dew.inputs

What a generative objective is fed: the sample field and its conditions.

InputSpec names the batch field the model learns to generate and, keyed by the model’s own keyword arguments, the conditions it is given. A Condition is an encoder, the batch field it reads, and the raw datum that stands for “no condition”, which classifier-free guidance and conditioning dropout substitute. Nothing here runs a model: the spec is a description, and the objective does the encoding.

Image and video batches arrive as uint8 pixels in [0, 255], the way the data workers write them. unit_range is the one conversion to the [-1, 1] range every diffusion loss, sample, artifact and image metric lives in; dew.artifacts.uint8_pixels is the one conversion back.

Name	Summary
CLIPText	The CLIP text tower, vendored in dew.nn.text_encoders, with the checkpoint’s tokenizer.
CharTable	Encodes text as a table lookup: one id per character, one fixed random vector per id.
Condition	Names one conditioning input: its encoder, the batch field holding its tokens, and the raw datum for the unconditional branch.
ConditionEncoder	Carries one modality from raw data to a conditioning value.
DiffusionConditioner	The text conditioning of a published latent diffusion checkpoint.
Field	A batch field and its per-example shape: Field("image", (128, 128, 3)).
InputSpec	Names the sample field and the conditions, keyed by the model keyword each is passed under: {"textcontext": Condition(...)}.
T5Text	The T5 encoder tower, vendored in dew.nn.text_encoders, with the checkpoint’s tokenizer.
pixel_field	The batch field carrying one image per row for a vision tower.
rebuild	The named encoder rebuilt from its JSON fields.
unit_range	uint8 pixels in [0, 255] as float32 in [-1, 1].

CLIPText

dataclass source

class CLIPText(
    checkpoint: str,
    transformer: CLIPTextTransformer,
    params: Variables,
    tokenizer: PreTrainedTokenizerBase,
    dtype: Dtype | None = None,
    revision: str | None = None,
    param_dtype: str = 'float32',
)

The CLIP text tower, vendored in dew.nn.text_encoders, with the checkpoint’s tokenizer.

Prompts are padded to the checkpoint’s own context length, which the tokenizer reports.

CLIPText.from_pretrained

def from_pretrained(
    checkpoint: str = DEFAULT_MODEL,
    *,
    dtype=None,
    revision: str | None = None,
    param_dtype: str = 'float32',
    params: Variables | None = None,
) -> CLIPText

CharTable

dataclass source

class CharTable(
    params: Variables,
    tokens: int = 8,
    features: int = 16,
    vocab: int = 130,
    seed: int = 0,
    dtype: Dtype | None = None,
    param_dtype: str = 'float32',
)

Encodes text as a table lookup: one id per character, one fixed random vector per id.

It costs nothing and downloads nothing, which makes it the text encoder of tests, benchmarks and smoke runs. It has the shape of a real one, a TextContext with a mask, so a model that takes CLIP’s output takes this one unchanged.

CharTable.from_pretrained

def from_pretrained(
    checkpoint: str = 'char_table',
    *,
    dtype=None,
    tokens: int = 8,
    features: int = 16,
    vocab: int = 130,
    seed: int = 0,
    param_dtype: str = 'float32',
    params: Variables | None = None,
)

The table seed draws, or the one params already holds.

There is nothing to load, so checkpoint goes unread here. It is on the signature because rebuild hands every encoder the name its to_json wrote, and this one writes the fixed "char_table".

CharTable.tokenize

def tokenize(texts: Sequence[str]) -> dict[str, np.ndarray]

CharTable.encode

def encode(params, tokens) -> TextContext

CharTable.captions

def captions(tokens) -> tuple[str, ...]

CharTable.to_json

def to_json() -> dict

Condition

dataclass source

class Condition(
    encoder: ConditionEncoder,
    field: str = 'text',
    unconditional: str | float | Mapping[str, object] = '',
)

Names one conditioning input: its encoder, the batch field holding its tokens, and the raw datum for the unconditional branch.

Condition.to_json

def to_json() -> dict

Condition.from_json

def from_json(record: Mapping, *, params: Variables | None = None) -> Condition

ConditionEncoder

class source

class ConditionEncoder(ABC, Generic[Raw, Encoded])

Carries one modality from raw data to a conditioning value.

parameter_collections: tuple[str, ...] | None

None declares a bare parameter tree; otherwise these collections own learned weights, including frozen ones. Other collections retain their dtype.

ConditionEncoder.from_pretrained

def from_pretrained(checkpoint: str, *, params: Variables | None = None) -> Self

Loads the tower named checkpoint, the one call that opens files.

Whatever else a checkpoint needs is a keyword field with a default, which is what to_json records and the registry rebuilds from. Supplied params are authoritative: the load reads metadata and never source weights, and keeps their values, dtypes and placement.

ConditionEncoder.tokenize

def tokenize(texts: Sequence[Raw]) -> Mapping[str, np.ndarray]

Raw data to the host arrays encode reads, one row per item.

ConditionEncoder.encode

def encode(params: Variables, tokens) -> Encoded

Tokens to the conditioning value, on device, under params.

ConditionEncoder.captions

def captions(tokens) -> tuple[str, ...]

What the tokens say, for a rendered artifact.

A modality that is not text has nothing to say and answers nothing.

ConditionEncoder.to_json

def to_json() -> dict

The keyword fields from_pretrained rebuilds this encoder from.

DiffusionConditioner

dataclass source

class DiffusionConditioner(
    towers: tuple[CLIPTextTransformer, ...],
    tokenizers: tuple[CLIPTokenizer, ...],
    names: tuple[str, ...],
    params: Variables,
    checkpoint: str,
    height: int,
    width: int,
    context_width: int,
    composition: Composition = 'clip',
    aesthetics: bool = False,
    t5: T5Segment | None = None,
    guidance: float | None = None,
    param_dtype: str = 'float32',
)

The text conditioning of a published latent diffusion checkpoint.

One encoder owns every family’s composition: which towers run, which of their states the model reads, and how the pooled vector is built. The towers themselves are the native CLIP and T5 towers, called the way their own source pipelines call them. The SD3 and Flux pipelines pass their T5 ids with no attention mask, which is what T5EncoderTransformer does with none, and no generic T5 default changes for it.

context_width: int

The width of the token sequence the denoiser reads: the UNet’s cross_attention_dim, the joint transformers’ joint_attention_dim. The SD3 composition pads its CLIP states out to it and writes the zero segment its pipeline substitutes for an absent third encoder at it.

stacked: bool

Whether the CLIP ids ride one array with a tower axis.

Every family but plain Stable Diffusion does. The XL refiner carries a single tower that way too, since its pipeline still writes a tower’s row rather than a bare batch.

DiffusionConditioner.from_pretrained

def from_pretrained(
    checkpoint: str,
    *,
    dtype: str | None = 'bfloat16',
    param_dtype: str = 'float32',
    revision: str | None = None,
    attention_impl: str = 'auto',
    params: Variables | None = None,
)

DiffusionConditioner.tokenize

def tokenize(texts: Sequence[str | Mapping[str, object]])

One row per item, with each text slot routed to the tower whose source pipeline reads it: text to the first CLIP tower, second to the second one, and the T5 tower’s own slot, which is third where a family has two CLIP towers beside it and second where it has one.

DiffusionConditioner.time_ids

def time_ids(
    count,
    dtype,
    *,
    original_size=None,
    crops_coords_top_left=(0, 0),
    target_size=None,
    aesthetic_score=6.0,
)

DiffusionConditioner.encode

def encode(params, tokens) -> DenoisingCondition

DiffusionConditioner.captions

def captions(tokens)

DiffusionConditioner.to_json

def to_json()

DiffusionConditioner.save_assets

def save_assets(destination: Path) -> None

Write the tokenizer files an exported directory carries beside the weights.

Field

dataclass source

class Field(key: str, shape: tuple[int, ...])

A batch field and its per-example shape: Field("image", (128, 128, 3)).

InputSpec

dataclass source

class InputSpec(
    sample: Field,
    conditions: Mapping[str, Condition] = dict(),
    mask: Field | None = None,
)

Names the sample field and the conditions, keyed by the model keyword each is passed under: {"textcontext": Condition(...)}.

tokenize is what a captioning dataset hands its text to. Every condition tokenizes the batch’s captions under its own field, so the encoder a run names decides the ids and the context length while the dataset carries the words alone.

mask: Field | None

A binary image mask for explicit masked-image latent conditioning.

InputSpec.tokenize

def tokenize(captions: Sequence[str]) -> dict[str, Mapping[str, np.ndarray]]

The batch fields this run’s conditions read out of captions.

Empty for a run that conditions on nothing, so the captions stop at the loader and no string array reaches a device.

InputSpec.to_json

def to_json() -> dict

InputSpec.from_json

def from_json(
    record: Mapping,
    *,
    params: Mapping[str, Variables] | None = None,
) -> InputSpec

Rebuilds the spec around supplied condition parameters, or loads each encoder’s own weights when none are given.

T5Text

dataclass source

class T5Text(
    checkpoint: str,
    transformer: T5EncoderTransformer,
    params: Variables,
    tokenizer: PreTrainedTokenizerBase,
    dtype: Dtype | None = None,
    revision: str | None = None,
    param_dtype: str = 'float32',
    max_length: int = 256,
)

The T5 encoder tower, vendored in dew.nn.text_encoders, with the checkpoint’s tokenizer.

It is the text half of an SD3.5/Flux-class run, whose MMDiT conditions on T5-XXL’s last hidden states. Prompts are padded to max_length, which the run’s record carries.

T5Text.from_pretrained

def from_pretrained(
    checkpoint: str = DEFAULT_T5_MODEL,
    *,
    dtype=None,
    revision: str | None = None,
    max_length: int = 256,
    param_dtype: str = 'float32',
    params: Variables | None = None,
) -> T5Text

pixel_field

function source

def pixel_field(height: int, width: int, channels: int = 3) -> Field

The batch field carrying one image per row for a vision tower.

It is float32 [channels, height, width], as the checkpoint’s processor emitted it, and rides beside the decoder’s token field.

rebuild

function source

def rebuild(
    name: str,
    fields: Mapping[str, object],
    *,
    params: Variables | None = None,
) -> ConditionEncoder

The named encoder rebuilt from its JSON fields.

A run’s record stores the registry name with the keyword fields to_json wrote. Those fields are unpacked here, so each encoder’s from_pretrained keeps its own concrete signature. The checkpoint is the one field every encoder takes and is read here; the rest are the encoder’s own and its signature checks them.

unit_range

function source

def unit_range(pixels: jax.typing.ArrayLike) -> jax.Array

uint8 pixels in [0, 255] as float32 in [-1, 1].