host/host_contract.md

Host contract — RGBA-Image-2.1 (Core AI port of Qwen-Image-2.1)

Everything a host computes around the three graphs, for text-to-image, batch 1, no CFG. The Python reference of every step below is in the zoo's conversion/qwenimage21/: qi21_tokenize.py (tokenizer), qi21_host.py (RoPE, pack/unpack), qi21_sched.py (sampler), pipeline_engine.py (the whole loop on the three bundles). The first three match the diffusers-main pipeline (Qwen/Qwen-Image-2.1 @ 790c926) exactly: the same token ids, bit-identical RoPE tables, bit-identical sigmas and Euler steps. pipeline_engine.py is scored against the pipeline's fp32 images.

1. The graphs

Every graph has one function, main. All tensors crossing a graph boundary are fp32, except input_ids (int32).

bundle	inputs	output	axes
qi21_encoder_dynL_w16a32_ids_iofp32.aimodel	input_ids [1,Lfull] int32	hidden [1,Lfull,4096]	Lfull 16..512
qi21_dit_full_bf16_dyn_iofp32.aimodel	img_tokens [1,N,64], txt_feats [1,L,4096], timestep [1], txt_cos [1,L,64], txt_sin [1,L,64], img_cos [1,N,64], img_sin [1,N,64]	vel [1,N,64]	L 8..512, N 64..4096
qi21_vae_{256,512,1024}_fp32.aimodel	latents_packed [1,N,64]	image [1,4,S,S]	fixed: N = (S/16)²

• Encoder: the Qwen3-VL-8B text stack (36 layers). bf16 weights, fp32 compute. embed_tokens is inside the graph. The output is the residual stream after the last layer, before the final RMSNorm. Do not apply a norm on the host.
• DiT: 32 blocks, bf16 weights and compute, fp32 boundary.
• VAE: decoder only, fp32. The graph unpacks the tokens and applies latents * std + mean itself.

2. Tokenize

Tokenizer: tokenizer/tokenizer.json (Qwen2 BPE, byte-level). No BOS token. An empty prompt is replaced by a single space " ".

Template (text-to-image):

<|im_start|>system\nComprehend and analyze the provided prompt.<|im_end|>\n<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n

\n is a newline character. Encode the whole string once, with no padding and no truncation, to get ids (length Lfull).

drop_idx is the number of tokens of the system part alone, <|im_start|>system\nComprehend and analyze the provided prompt.<|im_end|>\n. Compute it by encoding that string. With this tokenizer it is 14: [151644, 8948, 198, 1092, 30782, 408, 323, 23643, 279, 3897, 9934, 13, 151645, 198]. Check that ids[:drop_idx] equals those tokens.

Lfull must be 16..512 (the encoder's axis), and the DiT needs L = Lfull − drop_idx in 8..512. The template around the empty-prompt substitute " " is 23 tokens (L = 9).

3. Encode

hidden        = encoder(input_ids = ids as int32 [1, Lfull])      # [1, Lfull, 4096]
prompt_embeds = hidden[:, drop_idx : Lfull]                       # [1, L, 4096], L = Lfull - drop_idx

Pass exactly Lfull tokens; the axis is dynamic, so no padding is needed. (Attention is causal, so pad tokens never reach the first Lfull outputs, but a longer input moves the GPU result by about 4e-6 relative.) Run the encoder once per image.

4. Sequence layout and RoPE

The DiT's sequence is [text L | image N]: the L prompt tokens, then the image tokens in raster order. One image token covers one 16×16 px tile of the output: for an S×S image, h = w = S/16 and N = h·w (256² → 256 tokens, 512² → 1024, 1024² → 4096). There is no 2×2 packing.

Each token has a 3-axis position (frame, height, width):

• text token i (0-based) → (i, i, i);
• image token at row y, column x (0-based, token index y·w + x) → (L, y − (h − h//2), x − (w − w//2)). The grid is centred on zero. For h = 16, rows go −8..7.

host/rope_axis{a}_{cos,sin}.f32 hold one row per position −1024..8191: fp32, little-endian, row-major [9216, P_a] with P = (8, 28, 28). The row for position p is p + 1024. A token's 64 values are the three rows concatenated in axis order:

cos[token] = axis0_cos[f + 1024] ++ axis1_cos[hpos + 1024] ++ axis2_cos[wpos + 1024]   # 8 + 28 + 28
sin[token] = the same with the _sin tables
txt_cos, txt_sin = rows of the L text tokens     # [1, L, 64]
img_cos, img_sin = rows of the N image tokens    # [1, N, 64]

The tables are torch.polar(1, outer(pos, 1 / 10000^(arange(0, d, 2)/d))) for axis widths d = 16, 56, 56, the same numbers as QwenImage21Rope.freqs of diffusers main (bit-exact). Axes 1 and 2 have the same width, so their files are byte-identical. The four tensors depend only on L, h and w: build them once per image, not per step.

5. Noise and pack

The initial latent is standard normal noise: 64 channels × h rows × w columns, channel-major. The reference pipeline draws it as randn((1, 1, 64, h, w)) and packs it into DiT tokens:

x[0, y·w + x_, c] = z[c, y, x_]          # = z.view(1, 64, h·w).transpose(1, 2)

Any N(0, 1) noise generates an image. Matching the Python engine for a given seed needs its exact draw: torch.randn((1, 1, 64, h, w), generator=torch.Generator("cpu").manual_seed(seed)).

6. Sampler (FlowMatch Euler, 40 steps, no CFG)

Constants in host/scheduler.json (from the checkpoint's scheduler_config.json). All arrays fp32.

sigmas   = linspace(1, 1/steps, steps)                           # steps = 40
mu       = N · m + b,  m = (max_shift − base_shift) / (max_image_seq_len − base_image_seq_len),
                       b = base_shift − m · base_image_seq_len    # N = image tokens
sigmas   = e^mu / (e^mu + (1/sigmas − 1))                        # time_shift_type "exponential"
sigmas   = 1 − (1 − sigmas) / ((1 − sigmas[-1]) / (1 − shift_terminal))
timesteps = sigmas · num_train_timesteps                          # fp32
sigmas   = sigmas ++ [0]                                         # 41 values
for i in 0 ..< steps:
    t   = timesteps[i] / 1000                                    # fp32; this is the DiT `timestep` input
    vel = dit(img_tokens = x, txt_feats = prompt_embeds, timestep = [t], txt_cos, txt_sin, img_cos, img_sin)
    x   = x + (sigmas[i+1] − sigmas[i]) · vel                    # fp32

mu is 0.5 at 256² (N = 256), 0.5387… at 512², 0.6935… at 1024². Keep t = timesteps[i] / 1000 as written (multiply by 1000, then divide) to match the reference bit for bit. qi21_sched.py reproduces the reference sigmas, timesteps and every Euler step bit-exactly at 256² and 512².

7. Decode

image = vae_S(latents_packed = x)          # x after the last step, [1, N, 64], exactly as the sampler holds it
rgba8 = round(clip(image · 0.5 + 0.5, 0, 1) · 255)   # [1, 4, S, S] -> channels R, G, B, A

Feed the sampler's latent unchanged: no unpacking and no * std + mean on the host (the graph does both). The four channels are what the reference pipeline saves as an RGBA PNG.

8. Compile before running on the GPU

On macOS 27.0 (26A428) the Python runtime crashes on the 32-block DiT when the .aimodel is compiled just-in-time: the MPSGraph delegate forms a Neural Engine region inside the graph and the ANE inference fails (ANERegion.mm:414 failed assertion … Code=-19). A plain ahead-of-time compile fails the same way. What runs is an ahead-of-time compile with --expect-frequent-reshapes, loaded with SpecializationOptions.default():

xcrun coreai-build compile qi21_dit_full_bf16_dyn_iofp32.aimodel \
    --output aot/qi21_dit_full_bf16_dyn_iofp32 \
    --platform macOS --architecture h16c --preferred-compute gpu --expect-frequent-reshapes
# -> aot/qi21_dit_full_bf16_dyn_iofp32/qi21_dit_full_bf16_dyn_iofp32.h16c.aimodelc

Do the same for the encoder and the VAE bundle you use; every gate of this port ran all three this way, on an M4 Max with --architecture h16c (without --architecture, coreai-build compiles for every supported architecture). The compiled DiT is about 1.9× its .aimodel (27 GB); the encoder's stays at 14.10 GiB. Whether a Swift host using GraphModel(computeUnits: .gpu) hits the same Neural Engine region has not been tested.