+
+#### Negative prompt and batching {#stable-diffusion-negative-prompt-and-batching}
+
+Prompts batch through the tokenizer; a negative prompt is tokenized the same way and
+passed as `negative_input_ids` (one per prompt, or one broadcast row).
+
+```python
+prompts = [
+ "a watercolor painting of a lighthouse at sunset",
+ "a red vintage car on a coastal road",
+]
+inputs = tokenizer(prompts) # {"input_ids": (2, 77), "attention_mask": (2, 77)}
+negative = tokenizer(["blurry, low quality"] * len(prompts))["input_ids"]
+
+images = model.generate(**inputs, negative_input_ids=negative, num_inference_steps=30)
+for i, image in enumerate(images):
+ Image.fromarray(image).save(f"out_{i}.png")
+```
+
+#### Reproducible across backends {#stable-diffusion-reproducible-across-backends}
+
+`seed` reproduces a run on a given backend, but `keras.random` differs between torch,
+jax and tensorflow. For an image that is identical everywhere, pass the initial latent
+yourself (`(batch, 64, 64, 4)` at 512px):
+
+```python
+import numpy as np
+
+latents = np.random.default_rng(0).standard_normal((1, 64, 64, 4)).astype("float32")
+images = model.generate(**tokenizer("a bowl of ramen"), latents=latents)
+```
+
+#### Verified against diffusers {#stable-diffusion-verified-against-diffusers}
+
+The fox prompt above with the same initial latent, PNDM 50 steps and guidance 7.5
+through `ZMDiffusion` (left) and diffusers' `StableDiffusionPipeline`
+(right), both fp32:
+
+
+
+```
+tokenizer input_ids identical
+text_encoder max|d|=9.5e-06
+unet noise_pred (t=981) max|d|=9.5e-07
+vae decode max|d|=3.5e-06
+final image (uint8) max|d|=1 mean|d|=0.0022 PSNR=74.7 dB identical_pixels=99.34%
+```
+
+#### Components on their own {#stable-diffusion-components-on-their-own}
+
+```python
+from zeromodels.diffusers import ZMDiffusion
+
+sd = ZMDiffusion.from_weights("zeromodels/stable-diffusion-v1-5")
+latent = sd.vae.encode(image) # (1, 512, 512, 3) in [-1, 1] -> (1, 64, 64, 4)
+text = sd.text_encoder({"token_ids": ids, "padding_mask": mask})["last_hidden_state"]
+```
+
+#### Other resolutions {#stable-diffusion-other-resolutions}
+
+The graphs are built for a fixed size, the weights are not. Rebuild at another size with
+the constructor overrides (multiples of 64px; SD 1.x was trained at 512px):
+
+```python
+model = ZMDiffusion.from_weights(
+ "zeromodels/stable-diffusion-v1-5", unet_sample_size=96, vae_sample_size=768
+)
+images = model.generate(**tokenizer("a mountain lake at dawn")) # (1, 768, 768, 3)
+```
+
+### Data Format {#stable-diffusion-data-format}
+
+**Both `channels_last` and `channels_first` are supported.** The models read
+`keras.config.image_data_format()` when they are constructed (there is no
+`data_format` argument), so set it once, before `from_weights`, and every tensor that
+crosses the model boundary follows it:
+
+| | `channels_last` (default) | `channels_first` |
+|---|---|---|
+| `latents` passed to `generate` | `(batch, 64, 64, 4)` | `(batch, 4, 64, 64)` |
+| `AutoencoderKL.encode` input / `decode` output | `(batch, 512, 512, 3)` | `(batch, 3, 512, 512)` |
+| UNet `sample` | `(batch, 64, 64, 4)` | `(batch, 4, 64, 64)` |
+
+`generate` itself always hands back `(batch, height, width, 3)` uint8 images, whatever
+the layout, so the PIL step is the same in both. The same hosted `model.weights.h5`
+loads under either layout (conv kernels are `(kh, kw, in, out)` in both), and the two
+runs agree to within 1 uint8 level (99.7% of pixels identical, the rest being the
+NCHW / NHWC convolution kernels' rounding).
+
+```python
+import keras
+
+keras.config.set_image_data_format("channels_first")
+
+model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-v1-5")
+images = model.generate(
+ **tokenizer("a bowl of ramen"), seed=0
+) # still (1, 512, 512, 3)
+```
+
+Note that `keras.config.set_image_data_format` is global state. Set it once at the top
+of a script rather than toggling it between calls, since already-built models keep the
+layout they were constructed with.
+
+### Memory and speed {#stable-diffusion-memory-and-speed}
+
+The fp32 container is 3.97 GB; a 512px, batch-1, guided run peaks at 6.5 GB of GPU
+memory (the guided UNet call runs batch 2). On an RTX 4060 laptop the 50-step run takes
+~25 s in eager Torch, about 1.5x diffusers' eager time. On JAX the compiled denoiser step
+measured 2-3x faster than an op-by-op run; TensorFlow uses the same mechanism.
+
+### Loading Fine-tuned Weights {#stable-diffusion-loading-fine-tuned-weights}
+
+The five hosted checkpoints are the supported weights; any repo laid out like them
+(`zm_config.json` declaring `ZMDiffusion` with `model_type: "stable_diffusion"`,
+`model.weights.h5`, `tokenizer.json`) loads with `from_weights("
+
+#### 768px, v-prediction {#stable-diffusion-2-768px-v-prediction}
+
+The 768px checkpoints need no extra arguments; the v-prediction DDIM scheduler and the
+96x96 latent come from the repo:
+
+```python
+model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-2-1")
+tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-2-1")
+images = model.generate(
+ **tokenizer("a lighthouse on a cliff at dusk, oil painting")
+) # (1, 768, 768, 3)
+```
+
+#### SD-Turbo {#stable-diffusion-2-sd-turbo}
+
+One step, no guidance (the repo's defaults; up to 4 steps sharpen a little):
+
+```python
+model = ZMDiffusion.from_weights("zeromodels/sd-turbo")
+tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/sd-turbo")
+images = model.generate(
+ **tokenizer(
+ "a cinematic shot of a baby raccoon wearing an intricate italian priest robe"
+ )
+)
+```
+
+
+
+The same prompt and latent through diffusers' `StableDiffusionPipeline` (fp32) match to
+1 uint8 level: 99.9% of pixels identical at 1 step (PSNR 82 dB), 99.8% at 4 steps.
+
+Negative prompts, batching, explicit `latents` for cross-backend reproducibility, other
+resolutions and the components on their own work exactly as on the
+[Stable Diffusion](#stable-diffusion-end-to-end-example) page. Image-to-image
+(`generate(..., image=..., strength=...)`) works as described for
+[`BaseDiffusion`](main_classes.md#basediffusion).
+
+#### Verified against diffusers {#stable-diffusion-2-verified-against-diffusers}
+
+The ramen prompt above with the same initial latent, PNDM 50 steps and guidance 7.5
+through `ZMDiffusion` (left) and diffusers' `StableDiffusionPipeline`
+(right) on `stable-diffusion-2-1-base`, both fp32:
+
+
+
+```
+tokenizer input_ids identical
+text_encoder max|d|=1.5e-05
+unet noise_pred (t=981) max|d|=7.2e-07
+vae decode max|d|=2.6e-06
+final image (uint8) max|d|=1 mean|d|=0.0095 PSNR=68.3 dB identical_pixels=97.22%
+```
+
+The 768px v-prediction checkpoint (`stable-diffusion-2-1`, DDIM 20 steps, same latent)
+matches the same way: UNet 8.3e-6, VAE 7.9e-5, final image max 1 uint8 level, 99.3% of
+pixels identical, PSNR 74.6 dB.
+
+### Data Format {#stable-diffusion-2-data-format}
+
+Both `channels_last` and `channels_first` are supported, as for
+[Stable Diffusion](#stable-diffusion-data-format); `generate` always returns
+`(batch, H, W, 3)` uint8.
+
+### Memory and speed {#stable-diffusion-2-memory-and-speed}
+
+The fp32 container is 5.16 GB. At 512px a guided batch-1 run needs a little over 7 GB of
+GPU memory in eager Torch (about 1.5x diffusers' eager time); at 768px the self-attention
+runs over 9216 tokens, so the guided run needs roughly 12 GB, or the CPU.
+
+### Loading Fine-tuned Weights {#stable-diffusion-2-loading-fine-tuned-weights}
+
+Any repo laid out like the hosted ones (`zm_config.json` declaring
+`ZMDiffusion` with `model_type: "stable_diffusion_2"`, the weights, `tokenizer.json`)
+loads with `from_weights("
+
+#### Base + refiner {#stable-diffusion-xl-base-refiner}
+
+The SDXL "ensemble of experts": the base runs the first 80% of the schedule and hands its
+latent to the refiner, which finishes it. Same prompt and initial latent as above:
+
+```python
+refiner = ZMDiffusion.from_weights("zeromodels/stable-diffusion-xl-refiner-1.0")
+latent = model.generate(
+ **inputs, num_inference_steps=30, denoising_end=0.8, output_type="latent", seed=3
+)
+images = refiner.generate(
+ **inputs, latents=latent, num_inference_steps=30, denoising_start=0.8
+)
+```
+
+
+
+#### SDXL-Turbo {#stable-diffusion-xl-sdxl-turbo}
+
+One step, no guidance, 512px (the repo's defaults):
+
+```python
+model = ZMDiffusion.from_weights("zeromodels/sdxl-turbo")
+tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/sdxl-turbo")
+images = model.generate(
+ **tokenizer(
+ "a cinematic shot of a baby raccoon wearing an intricate italian priest robe"
+ )
+)
+```
+
+
+
+#### Negative prompts, micro-conditioning {#stable-diffusion-xl-negative-prompts-micro-conditioning}
+
+```python
+inputs = tokenizer(["a red bicycle", "a blue bicycle"])
+negative = tokenizer("blurry, low quality") # one row, shared by the batch
+images = model.generate(
+ **inputs,
+ negative_input_ids=negative["input_ids"],
+ original_size=(2048, 2048), # "a crop of a larger image": sharper detail
+ crops_coords_top_left=(0, 0),
+ target_size=(1024, 1024),
+ seed=0,
+)
+```
+
+#### Other resolutions {#stable-diffusion-xl-other-resolutions}
+
+The graphs are built for the repo's resolution; pass overrides to build for another
+(the weights are resolution-independent; multiples of 64px):
+
+```python
+model = ZMDiffusion.from_weights(
+ "zeromodels/stable-diffusion-xl-base-1.0",
+ unet_sample_size=(96, 128),
+ vae_sample_size=(768, 1024),
+)
+images = model.generate(**tokenizer("a wide mountain valley")) # (1, 768, 1024, 3)
+```
+
+#### float32 {#stable-diffusion-xl-float32}
+
+```python
+model = ZMDiffusion.from_weights(
+ "zeromodels/stable-diffusion-xl-base-1.0", load_dtype="float32"
+)
+```
+
+#### Verified against diffusers {#stable-diffusion-xl-verified-against-diffusers}
+
+The components of `stable-diffusion-xl-base-1.0` and `stable-diffusion-xl-refiner-1.0`
+were checked one by one against diffusers 0.39 / transformers in float32 on the same
+inputs (the converted float16 weights computed in float32):
+
+```
+base text_encoder penultimate max|d|=6.1e-05 (values up to 854)
+base text_encoder_2 penultimate max|d|=3.1e-05
+base text_encoder_2 pooled max|d|=3.6e-06
+base unet noise_pred (t=981) max|d|=9.3e-06
+base vae decode (512px) max|d|=3.5e-05
+refiner text_encoder_2 penultimate max|d|=3.1e-05
+refiner unet noise_pred (t=181) max|d|=1.2e-05
+```
+
+The pipelines (Euler `leading` spacing, classifier-free guidance with the zeroed or
+encoded negative branch, the micro-conditioning values, negative prompts, batches,
+`strength`, `denoising_end` / `denoising_start`) were checked against
+`StableDiffusionXLPipeline` and `StableDiffusionXLImg2ImgPipeline` on small
+random-weight models, where the generated images are pixel-identical on torch and jax.
+At full size in float16 (the default load; diffusers in float16 with its float32 VAE),
+the lighthouse above matches diffusers to PSNR 49.7 dB (98.5% of pixels within 2 uint8
+levels), the base + refiner ensemble to PSNR 51.9 dB (98.8% within 2) and SDXL-Turbo
+(1 step, 512px) to PSNR 44.9 dB (92% within 2; a single step amplifies float16 rounding).
+The lighthouse, zeromodels (left) and diffusers (right):
+
+
+
+### Data Format {#stable-diffusion-xl-data-format}
+
+Both `channels_last` and `channels_first` are supported, as for
+[Stable Diffusion](#stable-diffusion-data-format); `generate` always returns
+`(batch, H, W, 3)` uint8.
+
+### Memory and speed {#stable-diffusion-xl-memory-and-speed}
+
+Measured on the torch backend with the default float16 load (weights resident: 6.8 GB for
+a base model, 6.3 GB for the refiner, the VAE in float32):
+
+| Resolution | Denoising step (guided, batch 1) | VAE decode (float32) |
+|---|---|---|
+| 512px | peak 7.1 GB, 1.2 s/step on an RTX 4060 Laptop | peak 8.5 GB |
+| 768px | peak 8.0 GB | peak 10.5 GB |
+| 1024px | peak 10.4 GB | about 14 GB |
+
+Keras' functional executor keeps every layer output of a graph alive until the graph
+returns, so the decoder's feature maps at 1024px add several GB on top of the weights: a
+16 GB GPU runs 1024px comfortably, a 12 GB GPU 768px; on 8 GB, generate at 512px, or
+denoise on the GPU (`output_type="latent"`) and decode the latent on the CPU (about 100 s
+at 1024px) with a second, CPU-only process. In float32 the weights alone take 13.9 GB.
+The UNet's attention runs the portable `"sdpa"` math by default, which materializes the
+4096 x 4096 float32 logits of the 64 x 64 level at 1024px; loading with
+`attn_implementation="fused"` (the backend's fused kernel, see
+[`fused_attention`](main_classes.md#fused_attention)) cuts the 1024px step from a 10.3 GB
+peak to 7.6 GB on the same GPU. On JAX and TensorFlow the denoiser step is compiled once
+per run (see [`BaseDiffusion`](main_classes.md#basediffusion)).
+
+### Loading Fine-tuned Weights {#stable-diffusion-xl-loading-fine-tuned-weights}
+
+Any repo laid out like the hosted ones (`zm_config.json` declaring
+`ZMDiffusion` with `model_type: "stable_diffusion_xl"` or
+`"stable_diffusion_xl_refiner"`, the weights, `tokenizer.json`) loads with
+`from_weights("
+
+#### The T5 encoder on the CPU {#stable-diffusion-3-the-t5-encoder-on-the-cpu}
+
+The 9.5 GB T5-XXL rarely fits next to the MMDiT on a consumer GPU. Encode the prompt
+separately (an `SD3T5EncoderModel` in another process, or int8) and hand the ids to a
+model without T5 for the CLIP part only, or attach it in full on a large machine. The
+`encode_prompt` hook returns the conditioning dict the transformer takes, so any
+precomputed T5 features can be substituted there.
+
+#### Verified against diffusers {#stable-diffusion-3-verified-against-diffusers}
+
+The components of `stable-diffusion-3-medium` were checked one by one against diffusers
+0.39 / transformers 5 in float32 on the same inputs (the converted float16 weights
+computed in float32):
+
+```
+text_encoder penultimate max|d|=6.1e-05 (values up to 854)
+text_encoder_2 penultimate max|d|=3.1e-05
+pooled (L + G, projected) max|d|=3.8e-06
+t5-v1_1-xxl encoder (256 tok) max|d|=1.1e-04 (values up to 6.3)
+mmdit velocity (t=1000) max|d|=9.9e-06
+vae decode (512px) max|d|=6.3e-05
+vae encode max|d|=1.4e-05
+```
+
+The pipeline (flow-match Euler with shift 3, classifier-free guidance with the encoded
+empty prompt, the three tokenizers, batched negative prompts, the no-T5 mode) was checked
+against `StableDiffusion3Pipeline` on a small random-weight SD 3, where the generated images
+are pixel-identical on torch and jax. End to end, `stable-diffusion-3-medium` without the
+T5 encoder (28 steps, guidance 7, the same initial latent) reproduces the diffusers
+pipeline in float32 at 512px: 92.8% identical pixels, max |d| 2 / 255, PSNR 63.9 dB (the
+final latent differs by at most 2.5e-2 with values up to 5.1). In float16 on the GPU the
+two implementations round differently through the 28 guided steps: the pictures agree in
+composition but not pixel for pixel (PSNR 26 dB at 1024px, 29 dB at 768px, float16 MMDiT
+on both sides), the usual float16 sampling noise.
+
+### Data Format {#stable-diffusion-3-data-format}
+
+Both `channels_last` and `channels_first` are supported; `generate` always returns
+`(batch, H, W, 3)` uint8.
+
+### Memory and speed {#stable-diffusion-3-memory-and-speed}
+
+The checkpoint is 5.8 GB in float16 (the T5 encoder adds 9.5 GB). Measured on the torch
+backend with the default float16 load (RTX 4060 Laptop, 8 GB), T5 features precomputed:
+
+| Resolution | Denoising step (guided, batch 1) | VAE decode (float32) |
+|---|---|---|
+| 512px | peak 6.1 GB, 0.5 s/step | peak 7.6 GB, 4 s |
+| 768px | peak 6.4 GB, 0.8 s/step | peak 9.5 GB |
+| 1024px | peak 6.9 GB, 1.0 s/step | peak 12.4 GB |
+
+The MMDiT itself fits an 8 GB GPU at the native 1024px thanks to the fused attention
+kernel (the `"sdpa"` math needs 3.8 GB per block for the 4429 x 4429 float32 logits and
+runs out of memory). The VAE decoder is the SDXL one (16 latent channels aside) and Keras'
+functional executor keeps its feature maps alive, see the
+[SDXL notes](#stable-diffusion-xl-memory-and-speed): on 8 GB, denoise on the GPU
+(`output_type="latent"`) and decode on the CPU (75 to 90 s at 1024px). The T5-XXL encoder
+encodes a prompt and its negative on the CPU in about 10 to 25 s (float32 compute over
+the float16 weights); a 24 GB GPU holds everything. On JAX and TensorFlow the denoiser
+step is compiled once per run.
+
+### Loading Fine-tuned Weights {#stable-diffusion-3-loading-fine-tuned-weights}
+
+Any repo laid out like the hosted ones (`zm_config.json` declaring
+`ZMDiffusion` with `model_type: "stable_diffusion_3"`, the weights, the two tokenizer
+files) loads with
+`from_weights("
+
+#### Large Turbo {#stable-diffusion-3-5-large-turbo}
+
+```python
+model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-3.5-large-turbo")
+images = model.generate(
+ **tokenizer("a red panda reading a book")
+) # 4 steps, guidance 0 (the repo defaults)
+```
+
+#### Verified against diffusers {#stable-diffusion-3-5-verified-against-diffusers}
+
+The components of `stable-diffusion-3.5-medium` (the MMDiT-X with dual attention and RMS
+qk norms) were checked against diffusers 0.39 in float32 on the same inputs (the converted
+float16 weights computed in float32): MMDiT velocity max|d| 1.6e-5 (values up to 5.2), the
+text encoders and the VAE as for [SD 3](#stable-diffusion-3-verified-against-diffusers).
+The SD 3.5 block variants (dual attention, `qk_norm`) are also covered by the small
+random-weight pipeline test, pixel-identical to `StableDiffusion3Pipeline` on torch and
+jax, and by the SD 3 float32 end-to-end check (92.8% identical pixels, PSNR 63.9 dB at
+512px), which the two families share down to the sampler. The picture above is the
+1024px float16 output of `stable-diffusion-3.5-medium` with the T5 encoder.
+
+### Memory and speed {#stable-diffusion-3-5-memory-and-speed}
+
+The medium checkpoint is 6.8 GB in float16, the large ones 17.8 GB (a 24 GB GPU); the
+T5-XXL encoder adds 9.5 GB. See [Stable Diffusion 3](#stable-diffusion-3-memory-and-speed):
+the medium MMDiT-X at 1024px peaks at 7.5 GB during the guided denoising step on the torch
+backend (fused attention, float16), which an 8 GB GPU runs with the VAE decode on the CPU.
+
+### Loading Fine-tuned Weights {#stable-diffusion-3-5-loading-fine-tuned-weights}
+
+Any repo laid out like the hosted ones (`zm_config.json` declaring
+`ZMDiffusion` with `model_type: "stable_diffusion_3_5"`, the weights, the two tokenizer
+files) loads with
+`from_weights("
-Weights: pretrained Keras weights live on Hugging Face under
-zeromodels/<variant>
-(each repo carries
-
-Qwen-Image, ported to pure Keras 3: latent text-to-image flow-matching with a
-60-layer double-stream MMDiT denoiser, a Wan-derived KL autoencoder and the
-Qwen2.5-VL-7B Instruct text tower. The whole model is **one container**,
-`QwenImageModel`, a single functional graph with the three towers as disconnected
-sub-graphs (the way `CLIPModel` holds its two), so a hosted repo is one set of
-sharded weights and one `zm_config.json`. `QwenImageTextToImage` is that same
-graph plus `generate`: the true-CFG denoising loop of `BaseDiffusion`, the
-diffusion counterpart of `BaseGeneration`.
-
-The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion
-is deliberately **not supported** for diffusion models.
-
-Key facts of the port:
-
-- **Packed latents**: the denoiser sees `(B, H/2 Β· W/2, 64)` tokens; pack / unpack
- match Diffusers' 2Γ2 patching around the 16-channel VAE latent grid
- (`AutoencoderKLQwenImage`, 8Γ spatial, channels-last NDHWC with `T=1`).
-- **True CFG**: when `guidance_scale > 1`, two transformer forwards and Diffusers'
- prediction-norm renormalization (`true_cfg_scale`); a space negative prompt is
- the default when none is passed.
-- **bf16-safe dual-stream clip**: Diffusers clips dual-stream activations to the
- fp16 finite range only under float16; the port matches that, so bf16 runs keep
- the ~1e7 magnitudes Qwen-Image produces instead of collapsing to noise.
-- **Schedulers match Diffusers**: `FlowMatchEulerDiscreteScheduler` with dynamic
- resolution shifting (`mu` from packed sequence length) so timesteps agree with
- the reference within float rounding.
-- **Verified against Diffusers**: transformer, text encoder, VAE and scheduler
- outputs agree with the reference within float noise (see
- [Verified against Diffusers](#verified-against-diffusers)).
-
-Links:
-
-- Paper: [Qwen-Image Technical Report (arXiv:2508.02324)](https://arxiv.org/abs/2508.02324)
-- Reference implementation: [diffusers `QwenImagePipeline`](https://huggingface.co/docs/diffusers/api/pipelines/qwenimage)
-- License: [Apache-2.0](https://huggingface.co/Qwen/Qwen-Image/blob/main/LICENSE)
-
-See also [qwen2_5_vl.md](qwen2_5_vl.md) (the text tower),
-[qwen_image_21.md](qwen_image_21.md) (2.1 single-stream / unpatched),
-[stable_diffusion_3.md](stable_diffusion_3.md) (flow-match / MMDiT-style diffusion
-in ZeroModels).
-
-## Variants
-
-Preconverted, bfloat16 weights are hosted under `zeromodels/`. Load with
-`from_weights("zeromodels/zm_config.json + sharded
-*.weights.json / *.weights.h5 +
-tokenizer.json). Load with from_weights("zeromodels/<variant>").
-
-Weights: pretrained Keras weights live on Hugging Face under
-zeromodels/<variant>
-(each repo carries
-
-Qwen-Image-2.1, ported to pure Keras 3: latent text-to-image flow-matching with a
-32-layer **single-stream** block-causal DiT, a residual 64-channel KL autoencoder
-(16Γ spatial, RGBA), and the Qwen3-VL text tower. The whole model is **one
-container**, `QwenImage21Model`. `QwenImage21TextToImage` adds `generate`.
-
-The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion
-is deliberately **not supported** for diffusion models.
-
-Key facts of the port:
-
-- **Unpatched latents**: the denoiser sees `(B, H Β· W, 64)` tokens (VAE scale 16;
- no 2Γ2 packing). At 1024px that is a `(B, 4096, 64)` sequence.
-- **Block-causal attention**: text is causal; the target image block is
- bidirectional and can attend to all preceding text.
-- **`causal_condition`**: text tokens modulate from `t = 0` (timestep-independent),
- matching Diffusers' KV-cache-ready conditioning.
-- **True CFG optional**: Diffusers defaults to `true_cfg_scale=1.0` (no guidance).
- Pass `guidance_scale > 1` with a negative prompt to enable dual forwards.
-- **Pre-norm text features**: the text tower returns decoder outputs *before* the
- final RMSNorm, matching Diffusers' forward hook on the language-model norm.
-- **Schedulers match Diffusers**: `FlowMatchEulerDiscreteScheduler` with dynamic
- resolution shifting (`mu` from image sequence length) and `shift_terminal`.
-
-Links:
-
-- Source: [`Qwen/Qwen-Image-2.1`](https://huggingface.co/Qwen/Qwen-Image-2.1)
-- Reference: [diffusers `QwenImage21Pipeline`](https://huggingface.co/docs/diffusers/api/pipelines/qwenimage)
-- License: [Qwen Research License](https://huggingface.co/Qwen/Qwen-Image-2.1)
-- See also [qwen_image.md](qwen_image.md) (1.0 double-stream / packed),
- [qwen3_vl.md](qwen3_vl.md) (text tower)
-
-## Variants
-
-Preconverted, bfloat16 weights are hosted under `zeromodels/`. Load with
-`from_weights("zeromodels/zm_config.json + sharded
-*.weights.json / *.weights.h5 +
-tokenizer.json). Load with from_weights("zeromodels/<variant>").
-
Weights: Pre-converted Keras weights for all 13 checkpoints are hosted under the
@@ -10,21 +10,51 @@ with the
-ZeroModels loads 13 original dense sentence-embedding checkpoints. See the
-[Sentence Transformer models](sentence_transformers.md) page for the available
-ZeroModels-hosted checkpoint IDs and Hub links.
-It builds a Keras 3 encoder and applies the checkpoint's declared pooling,
-optional Dense projection, and optional L2 normalization. Parity has been
+ZeroModels loads 13 original dense sentence-embedding checkpoints (listed under
+[Checkpoints](#checkpoints)). It builds a Keras 3 encoder and applies the
+checkpoint's declared pooling, optional Dense projection, and optional L2
+normalization. Parity has been
verified on the Keras Torch backend; TensorFlow and JAX parity have not yet
been checked for all 13 checkpoints.
The public API is `zeromodels.sentence_transformers`, separate from the upstream
-`sentence_transformers` reference library. The former `zeromodels.st_models`
-and `zeromodels.models.sentence_transformer` paths remain compatibility aliases
-for existing code and saved metadata; new examples and exports use the new namespace.
-Use `ZMSentenceTransformer` and `ZMSentenceTransformerTokenizer`; the former
-`SentenceTransformer`, `SentenceTransformerModel`, and `SentenceTransformerTokenizer`
-names remain compatibility aliases.
+`sentence_transformers` reference library, and it is exactly two classes:
+`ZMSentenceTransformer` and `ZMSentenceTransformerTokenizer`. Hub repos whose
+`zm_config.json` or `zm_preprocessor.json` still name the former classes
+(`SentenceTransformerModel`, `SentenceTransformer`, `SentenceTransformerTextEmbed`,
+`SentenceTransformerTokenizer`) load into these two unchanged.
+
+| | |
+|---|---|
+| Class | `ZMSentenceTransformer` |
+| Tokenizer | `ZMSentenceTransformerTokenizer` |
+| Selector | `backbone_type` (one architecture, pluggable encoder) |
+| Recipes | `SENTENCE_TRANSFORMER_BACKBONES` maps each `backbone_type` to its encoder builder: the `zeromodels.models` MPNet / BERT / RoBERTa / XLM-R encoders, and DistilBERT / ALBERT builders in `sentence_transformer_backbones.py` |
+| Hub files | `zm_config.json` + `zm_preprocessor.json` + `tokenizer.json` |
+
+Supported `backbone_type` values: `mpnet`, `bert`, `roberta`, `xlm-roberta`,
+`distilbert`, `albert`.
+
+## Checkpoints
+
+| Checkpoint | Backbone | Embedding recipe | Dimensions |
+|---|---|---|---:|
+| [`all-mpnet-base-v2`](https://huggingface.co/zeromodels/all-mpnet-base-v2) | MPNet | Mean pooling β L2 normalization | 768 |
+| [`multi-qa-mpnet-base-dot-v1`](https://huggingface.co/zeromodels/multi-qa-mpnet-base-dot-v1) | MPNet | CLS pooling; dot-product similarity | 768 |
+| [`all-distilroberta-v1`](https://huggingface.co/zeromodels/all-distilroberta-v1) | RoBERTa | Mean pooling β L2 normalization | 768 |
+| [`all-MiniLM-L12-v2`](https://huggingface.co/zeromodels/all-MiniLM-L12-v2) | BERT | Mean pooling β L2 normalization | 384 |
+| [`multi-qa-distilbert-cos-v1`](https://huggingface.co/zeromodels/multi-qa-distilbert-cos-v1) | DistilBERT | Mean pooling β L2 normalization | 768 |
+| [`all-MiniLM-L6-v2`](https://huggingface.co/zeromodels/all-MiniLM-L6-v2) | BERT | Mean pooling β L2 normalization | 384 |
+| [`multi-qa-MiniLM-L6-cos-v1`](https://huggingface.co/zeromodels/multi-qa-MiniLM-L6-cos-v1) | BERT | Mean pooling β L2 normalization | 384 |
+| [`paraphrase-multilingual-mpnet-base-v2`](https://huggingface.co/zeromodels/paraphrase-multilingual-mpnet-base-v2) | XLM-RoBERTa | Mean pooling | 768 |
+| [`paraphrase-albert-small-v2`](https://huggingface.co/zeromodels/paraphrase-albert-small-v2) | ALBERT | Mean pooling | 768 |
+| [`paraphrase-multilingual-MiniLM-L12-v2`](https://huggingface.co/zeromodels/paraphrase-multilingual-MiniLM-L12-v2) | BERT | Mean pooling | 384 |
+| [`paraphrase-MiniLM-L3-v2`](https://huggingface.co/zeromodels/paraphrase-MiniLM-L3-v2) | BERT | Mean pooling | 384 |
+| [`distiluse-base-multilingual-cased-v1`](https://huggingface.co/zeromodels/distiluse-base-multilingual-cased-v1) | DistilBERT | Mean pooling β Tanh Dense | 512 |
+| [`distiluse-base-multilingual-cased-v2`](https://huggingface.co/zeromodels/distiluse-base-multilingual-cased-v2) | DistilBERT | Mean pooling β Tanh Dense | 512 |
+
+The `paraphrase-multilingual-mpnet-base-v2` name is historical; the checkpoint
+actually uses an XLM-RoBERTa backbone.
## Usage
@@ -135,8 +165,8 @@ tokenizer = ZMSentenceTransformerTokenizer.from_weights(
An explicit commit SHA can be supplied as `revision=` to pin the upstream
snapshot.
-`AutoZMTextEmbed` resolves ZeroModels `sentence_transformer` configs to the
-`SentenceTransformerTextEmbed` task view. For upstream `hf:` repos, load the
+`AutoZMTextEmbed` resolves ZeroModels `sentence_transformer` configs to
+`ZMSentenceTransformer` as well. For upstream `hf:` repos, load the
concrete class as shown above: their root `model_type` describes the backbone,
not the sentence-embedding task.
@@ -162,6 +192,6 @@ actual = model.encode(texts, tokenizer=tokenizer)
print("Maximum absolute difference:", np.max(np.abs(actual - expected)))
```
-Use matching checkpoint names from the models table, compare in float32, and
+Use matching checkpoint names from the [Checkpoints](#checkpoints) table, compare in float32, and
record Hub commit SHAs for reproducible results. This downloads original
weights for comparison; it does not convert or publish a new checkpoint.
diff --git a/docs/sentence_transformers.md b/docs/sentence_transformers.md
deleted file mode 100644
index 5a37ee8b..00000000
--- a/docs/sentence_transformers.md
+++ /dev/null
@@ -1,49 +0,0 @@
-# Sentence Transformer models
-
-ZeroModels provides 13 converted dense sentence-embedding checkpoints. Each
-checkpoint combines a text encoder with its original pooling, optional Dense
-projection, and optional L2 normalization. Load them through the
-[`zeromodels.sentence_transformers` API](sentence_transformer.md); no conversion is needed
-for the ZeroModels-hosted weights.
-
-| Checkpoint | Backbone | Embedding recipe | Dimensions |
-|---|---|---|---:|
-| [`all-mpnet-base-v2`](https://huggingface.co/zeromodels/all-mpnet-base-v2) | MPNet | Mean pooling β L2 normalization | 768 |
-| [`multi-qa-mpnet-base-dot-v1`](https://huggingface.co/zeromodels/multi-qa-mpnet-base-dot-v1) | MPNet | CLS pooling; dot-product similarity | 768 |
-| [`all-distilroberta-v1`](https://huggingface.co/zeromodels/all-distilroberta-v1) | RoBERTa | Mean pooling β L2 normalization | 768 |
-| [`all-MiniLM-L12-v2`](https://huggingface.co/zeromodels/all-MiniLM-L12-v2) | BERT | Mean pooling β L2 normalization | 384 |
-| [`multi-qa-distilbert-cos-v1`](https://huggingface.co/zeromodels/multi-qa-distilbert-cos-v1) | DistilBERT | Mean pooling β L2 normalization | 768 |
-| [`all-MiniLM-L6-v2`](https://huggingface.co/zeromodels/all-MiniLM-L6-v2) | BERT | Mean pooling β L2 normalization | 384 |
-| [`multi-qa-MiniLM-L6-cos-v1`](https://huggingface.co/zeromodels/multi-qa-MiniLM-L6-cos-v1) | BERT | Mean pooling β L2 normalization | 384 |
-| [`paraphrase-multilingual-mpnet-base-v2`](https://huggingface.co/zeromodels/paraphrase-multilingual-mpnet-base-v2) | XLM-RoBERTa | Mean pooling | 768 |
-| [`paraphrase-albert-small-v2`](https://huggingface.co/zeromodels/paraphrase-albert-small-v2) | ALBERT | Mean pooling | 768 |
-| [`paraphrase-multilingual-MiniLM-L12-v2`](https://huggingface.co/zeromodels/paraphrase-multilingual-MiniLM-L12-v2) | BERT | Mean pooling | 384 |
-| [`paraphrase-MiniLM-L3-v2`](https://huggingface.co/zeromodels/paraphrase-MiniLM-L3-v2) | BERT | Mean pooling | 384 |
-| [`distiluse-base-multilingual-cased-v1`](https://huggingface.co/zeromodels/distiluse-base-multilingual-cased-v1) | DistilBERT | Mean pooling β Tanh Dense | 512 |
-| [`distiluse-base-multilingual-cased-v2`](https://huggingface.co/zeromodels/distiluse-base-multilingual-cased-v2) | DistilBERT | Mean pooling β Tanh Dense | 512 |
-
-The `paraphrase-multilingual-mpnet-base-v2` name is historical; the checkpoint
-actually uses an XLM-RoBERTa backbone.
-
-## Quick start
-
-```python
-import os
-
-os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow"
-
-from zeromodels.sentence_transformers import (
- ZMSentenceTransformer,
- ZMSentenceTransformerTokenizer,
-)
-
-repo = "zeromodels/all-MiniLM-L6-v2"
-model = ZMSentenceTransformer.from_weights(repo)
-tokenizer = ZMSentenceTransformerTokenizer.from_weights(repo)
-
-embeddings = model.encode(["A cat on a mat.", "A dog on a mat."], tokenizer=tokenizer)
-print(embeddings.shape) # (2, 384)
-```
-
-See [Sentence Transformer](sentence_transformer.md) for the complete API,
-upstream `hf:` loading, similarity scoring, and parity-check instructions.
diff --git a/docs/stable_diffusion.md b/docs/stable_diffusion.md
deleted file mode 100644
index 2993961d..00000000
--- a/docs/stable_diffusion.md
+++ /dev/null
@@ -1,351 +0,0 @@
-# Stable Diffusion
-
-hf: prefix, e.g.
The backbone and embedding recipe are read from the checkpoint metadata.
-Weights: pretrained Keras weights live on Hugging Face under
-zeromodels/<variant>
-(each repo carries
-
-Stable Diffusion 1.x, ported to pure Keras 3: latent text-to-image diffusion with a
-UNet denoiser, a KL-regularized autoencoder and the CLIP ViT-L/14 text encoder. The
-whole model is **one container**, `StableDiffusionModel`, a single functional graph with
-the three towers as disconnected sub-graphs (the way `CLIPModel` holds its two), so a
-hosted repo is one `model.weights.h5` and one `zm_config.json`.
-`StableDiffusionTextToImage` is that same graph plus `generate`: the classifier-free
-guidance denoising loop of `BaseDiffusion`, the diffusion counterpart of
-`BaseGeneration`.
-
-All five SD 1.x checkpoints (v1-1 to v1-5) share this architecture; only the weights
-differ. The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion
-is deliberately **not supported** for diffusion models.
-
-Key facts of the port:
-
-- **Both data formats**: the UNet and VAE follow `keras.config.image_data_format()`
- (every conv takes it, GroupNorm normalizes the layout's channel axis, and the two
- token <-> grid boundaries transpose around their flatten under `channels_first`).
- Weights are layout-independent, so one hosted checkpoint serves both, and the
- converter is a `(O, I, H, W) -> (H, W, I, O)` kernel transpose and nothing else.
-- **Block-level layers**: the UNet and VAE are built from composite Keras layers
- (`StableDiffusionResnetBlock2D`, `StableDiffusionTransformer2DModel`,
- `StableDiffusionCrossAttention`, ...). A functional graph keeps every node's output
- alive until the forward ends, and at 512px the UNet's per-op intermediates (4096x4096
- attention maps above all) would need several GB; a layer's internals are freed when
- its call returns.
-- **Schedulers match diffusers to the bit**: the training noise schedule is built in
- float32 the way torch does it, so `alphas_cumprod` and the Euler sigmas are identical
- and a 50-step run stays within float rounding of the reference.
-- **Verified end to end**: at 512px, fp32, PNDM 50 steps and the same initial latent, the
- generated image matches diffusers' to within **1 uint8 level** (99.3% of pixels
- identical, PSNR 75 dB); the UNet, VAE and text encoder outputs agree to ~1e-6.
-
-Links:
-
-- Paper: [High-Resolution Image Synthesis with Latent Diffusion Models (arXiv:2112.10752)](https://arxiv.org/abs/2112.10752)
-- Reference implementation: [diffusers `StableDiffusionPipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img)
-- License: [CreativeML OpenRAIL-M](https://huggingface.co/spaces/CompVis/stable-diffusion-license)
-
-See also [clip.md](clip.md) (the text encoder is `CLIPTextModel`).
-
-## Variants
-
-Preconverted, float32 weights are hosted under `zeromodels/`. Load with
-`from_weights("zeromodels/zm_config.json + model.weights.h5 +
-tokenizer.json). Load with from_weights("zeromodels/<variant>").
-
-
-### Negative prompt and batching
-
-Prompts batch through the tokenizer; a negative prompt is tokenized the same way and
-passed as `negative_input_ids` (one per prompt, or one broadcast row).
-
-```python
-prompts = [
- "a watercolor painting of a lighthouse at sunset",
- "a red vintage car on a coastal road",
-]
-inputs = tokenizer(prompts) # {"input_ids": (2, 77), "attention_mask": (2, 77)}
-negative = tokenizer(["blurry, low quality"] * len(prompts))["input_ids"]
-
-images = model.generate(**inputs, negative_input_ids=negative, num_inference_steps=30)
-for i, image in enumerate(images):
- Image.fromarray(image).save(f"out_{i}.png")
-```
-
-### Reproducible across backends
-
-`seed` reproduces a run on a given backend, but `keras.random` differs between torch,
-jax and tensorflow. For an image that is identical everywhere, pass the initial latent
-yourself (`(batch, 64, 64, 4)` at 512px):
-
-```python
-import numpy as np
-
-latents = np.random.default_rng(0).standard_normal((1, 64, 64, 4)).astype("float32")
-images = model.generate(**tokenizer("a bowl of ramen"), latents=latents)
-```
-
-### Verified against diffusers
-
-The fox prompt above with the same initial latent, PNDM 50 steps and guidance 7.5
-through `StableDiffusionTextToImage` (left) and diffusers' `StableDiffusionPipeline`
-(right), both fp32:
-
-
-
-```
-tokenizer input_ids identical
-text_encoder max|d|=9.5e-06
-unet noise_pred (t=981) max|d|=9.5e-07
-vae decode max|d|=3.5e-06
-final image (uint8) max|d|=1 mean|d|=0.0022 PSNR=74.7 dB identical_pixels=99.34%
-```
-
-### Container only
-
-```python
-from zeromodels.models.stable_diffusion import StableDiffusionModel
-
-sd = StableDiffusionModel.from_weights("zeromodels/stable-diffusion-v1-5")
-latent = sd.vae.encode(image) # (1, 512, 512, 3) in [-1, 1] -> (1, 64, 64, 4)
-text = sd.text_encoder({"token_ids": ids, "padding_mask": mask})["last_hidden_state"]
-```
-
-### Other resolutions
-
-The graphs are built for a fixed size, the weights are not. Rebuild at another size with
-the constructor overrides (multiples of 64px; SD 1.x was trained at 512px):
-
-```python
-model = StableDiffusionTextToImage.from_weights(
- "zeromodels/stable-diffusion-v1-5", unet_sample_size=96, vae_sample_size=768
-)
-images = model.generate(**tokenizer("a mountain lake at dawn")) # (1, 768, 768, 3)
-```
-
-## Data Format
-
-**Both `channels_last` and `channels_first` are supported.** The models read
-`keras.config.image_data_format()` when they are constructed (there is no
-`data_format` argument), so set it once, before `from_weights`, and every tensor that
-crosses the model boundary follows it:
-
-| | `channels_last` (default) | `channels_first` |
-|---|---|---|
-| `latents` passed to `generate` | `(batch, 64, 64, 4)` | `(batch, 4, 64, 64)` |
-| `AutoencoderKL.encode` input / `decode` output | `(batch, 512, 512, 3)` | `(batch, 3, 512, 512)` |
-| UNet `sample` | `(batch, 64, 64, 4)` | `(batch, 4, 64, 64)` |
-
-`generate` itself always hands back `(batch, height, width, 3)` uint8 images, whatever
-the layout, so the PIL step is the same in both. The same hosted `model.weights.h5`
-loads under either layout (conv kernels are `(kh, kw, in, out)` in both), and the two
-runs agree to within 1 uint8 level (99.7% of pixels identical, the rest being the
-NCHW / NHWC convolution kernels' rounding).
-
-```python
-import keras
-
-keras.config.set_image_data_format("channels_first")
-
-model = StableDiffusionTextToImage.from_weights("zeromodels/stable-diffusion-v1-5")
-images = model.generate(
- **tokenizer("a bowl of ramen"), seed=0
-) # still (1, 512, 512, 3)
-```
-
-Note that `keras.config.set_image_data_format` is global state. Set it once at the top
-of a script rather than toggling it between calls, since already-built models keep the
-layout they were constructed with.
-
-## Memory and speed
-
-The fp32 container is 3.97 GB; a 512px, batch-1, guided run peaks at 6.5 GB of GPU
-memory (the guided UNet call runs batch 2). On an RTX 4060 laptop the 50-step run takes
-~25 s in eager Torch, about 1.5x diffusers' eager time. On JAX the compiled denoiser step
-measured 2-3x faster than an op-by-op run; TensorFlow uses the same mechanism.
-
-## Loading Fine-tuned Weights
-
-The five hosted checkpoints are the supported weights; any repo laid out like them
-(`zm_config.json` declaring `StableDiffusionModel`, `model.weights.h5`, `tokenizer.json`)
-loads with `from_weights("
-Weights: pretrained Keras weights live on Hugging Face under
-zeromodels/<variant>
-(each repo carries
-
-Stable Diffusion 2.x, ported to pure Keras 3: the [Stable Diffusion](stable_diffusion.md)
-latent-diffusion architecture with the second-generation conditioning. It reuses the
-SD 1.x family end to end (`StableDiffusion2Model` is `StableDiffusionModel`,
-`StableDiffusion2TextToImage` is `StableDiffusionTextToImage`, each with the SD 2
-configuration), so everything on that page applies: one container per repo, `generate`
-from `BaseDiffusion`, the schedulers, both data formats. What changes is the config:
-
-- **Text encoder**: OpenCLIP ViT-H/14's text tower, 1024 wide, 16 heads, `gelu`, shipped
- truncated to its 23rd (penultimate) layer, which SD 2 conditions on.
-- **UNet**: 1024-d cross-attention, one head count per level (`(5, 10, 20, 20)`, a 64-wide
- head everywhere, where SD 1.x used 8 heads at every level) and a linear token
- projection in the Transformer2D blocks instead of the 1x1 conv.
-- **Tokenizer**: the same CLIP BPE, padded with `!` (id 0) the OpenCLIP way rather than
- `<|endoftext|>`, so the empty prompt used for classifier-free guidance is
- `[<|startoftext|>, <|endoftext|>, !, !, ...]`.
-- **768px checkpoints** (`stable-diffusion-2`, `stable-diffusion-2-1`): built at a 96x96
- latent and trained with the **v-prediction** objective; their repos carry a
- v-prediction DDIM scheduler, which `generate` picks up from `scheduler_config`.
-- **SD-Turbo** (`sd-turbo`): SD 2.1 distilled with Adversarial Diffusion Distillation to
- generate in 1 to 4 steps **without guidance**; its repo carries an Euler scheduler with
- `trailing` timestep spacing and `generate_args` of 1 step, `guidance_scale=0.0`.
-
-The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion is
-deliberately **not supported** for diffusion models. The original `stabilityai/*` repos are
-no longer on the Hub; the conversion sources are the `sd2-community` mirrors.
-
-Links:
-
-- Paper: [High-Resolution Image Synthesis with Latent Diffusion Models (arXiv:2112.10752)](https://arxiv.org/abs/2112.10752)
-- Reference implementation: [diffusers `StableDiffusionPipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img)
-- License: [CreativeML Open RAIL++-M](https://huggingface.co/sd2-community/stable-diffusion-2-1/blob/main/LICENSE-MODEL)
-
-See also [stable_diffusion.md](stable_diffusion.md), [clip.md](clip.md).
-
-## Variants
-
-Preconverted, float32 weights are hosted under `zeromodels/`. Load with
-`from_weights("zeromodels/zm_config.json + model.weights.h5 +
-tokenizer.json). Load with from_weights("zeromodels/<variant>").
-
-
-### 768px, v-prediction
-
-The 768px checkpoints need no extra arguments; the v-prediction DDIM scheduler and the
-96x96 latent come from the repo:
-
-```python
-model = StableDiffusion2TextToImage.from_weights("zeromodels/stable-diffusion-2-1")
-tokenizer = StableDiffusion2Tokenizer.from_weights("zeromodels/stable-diffusion-2-1")
-images = model.generate(
- **tokenizer("a lighthouse on a cliff at dusk, oil painting")
-) # (1, 768, 768, 3)
-```
-
-### SD-Turbo
-
-One step, no guidance (the repo's defaults; up to 4 steps sharpen a little):
-
-```python
-model = StableDiffusion2TextToImage.from_weights("zeromodels/sd-turbo")
-tokenizer = StableDiffusion2Tokenizer.from_weights("zeromodels/sd-turbo")
-images = model.generate(
- **tokenizer(
- "a cinematic shot of a baby raccoon wearing an intricate italian priest robe"
- )
-)
-```
-
-
-
-The same prompt and latent through diffusers' `StableDiffusionPipeline` (fp32) match to
-1 uint8 level: 99.9% of pixels identical at 1 step (PSNR 82 dB), 99.8% at 4 steps.
-
-Negative prompts, batching, explicit `latents` for cross-backend reproducibility, other
-resolutions and the container-only use work exactly as on the
-[Stable Diffusion](stable_diffusion.md#end-to-end-example) page. Image-to-image
-(`generate(..., image=..., strength=...)`) works as described for
-[`BaseDiffusion`](main_classes.md#basediffusion).
-
-### Verified against diffusers
-
-The ramen prompt above with the same initial latent, PNDM 50 steps and guidance 7.5
-through `StableDiffusion2TextToImage` (left) and diffusers' `StableDiffusionPipeline`
-(right) on `stable-diffusion-2-1-base`, both fp32:
-
-
-
-```
-tokenizer input_ids identical
-text_encoder max|d|=1.5e-05
-unet noise_pred (t=981) max|d|=7.2e-07
-vae decode max|d|=2.6e-06
-final image (uint8) max|d|=1 mean|d|=0.0095 PSNR=68.3 dB identical_pixels=97.22%
-```
-
-The 768px v-prediction checkpoint (`stable-diffusion-2-1`, DDIM 20 steps, same latent)
-matches the same way: UNet 8.3e-6, VAE 7.9e-5, final image max 1 uint8 level, 99.3% of
-pixels identical, PSNR 74.6 dB.
-
-## Data Format
-
-Both `channels_last` and `channels_first` are supported, as for
-[Stable Diffusion](stable_diffusion.md#data-format); `generate` always returns
-`(batch, H, W, 3)` uint8.
-
-## Memory and speed
-
-The fp32 container is 5.16 GB. At 512px a guided batch-1 run needs a little over 7 GB of
-GPU memory in eager Torch (about 1.5x diffusers' eager time); at 768px the self-attention
-runs over 9216 tokens, so the guided run needs roughly 12 GB, or the CPU.
-
-## Loading Fine-tuned Weights
-
-Any repo laid out like the hosted ones (`zm_config.json` declaring
-`StableDiffusion2Model`, the weights, `tokenizer.json`) loads with
-`from_weights("
-Weights: pretrained Keras weights live on Hugging Face under
-zeromodels/<variant>
-(each repo carries
-
-Stable Diffusion 3, ported to pure Keras 3: a rectified-flow model whose denoiser is the
-**MMDiT**, a multimodal diffusion transformer in which the latent tokens and the text tokens
-attend to each other jointly, conditioned on the timestep and the pooled text embeddings
-through adaptive layer norms. It keeps the SD family's shape (one container per repo,
-`generate` from `BaseDiffusion`, the schedulers, both data formats) with new parts:
-
-- **MMDiT** (`SD3Transformer2DModel`): the 16-channel latent is patchified (2x2) into 1536-d
- tokens with a sinusoidal position table cropped from a 192x192 grid; 24 joint blocks of
- 24 heads, each with AdaLN-Zero modulation of both streams, joint attention over the
- concatenated latent + text tokens, and gated GELU feed-forwards; a final adaptive norm
- and projection unpatchify the velocity. 2B parameters.
-- **Three text encoders**: CLIP ViT-L/14 (768-d) and OpenCLIP ViT-bigG/14 (1280-d), each
- with a projection, and **T5-XXL** (4.7B, 4096-d, 256 tokens). The two CLIP penultimate
- states are concatenated (2048-d), zero-padded to 4096 and followed along the sequence by
- the T5 states (77 + 256 tokens); the projected pooled CLIP states (2048-d) feed the
- conditioning embedding.
-- **T5 is optional and separate**: the container holds the MMDiT, the VAE and the two CLIP
- towers (2.9B, 5.8 GB float16); the T5-XXL encoder is hosted once for every SD 3 / 3.5
- checkpoint (`zeromodels/t5-v1_1-xxl-encoder`, an
- [`SD3T5EncoderModel`](#sd3t5encodermodel), 9.5 GB float16) and attached with
- `text_encoder_3=`. Without it the T5 features are zeros, SD 3's documented
- memory-saving mode (prompt adherence drops, the images stay good).
-- **VAE**: 16 latent channels, no quant convolutions, `scaling_factor` 1.5305 and
- `shift_factor` 0.0609 (`z = (x - shift) * scale`), built in float32 (`force_upcast`).
-- **Sampler**: `FlowMatchEulerDiscreteScheduler` with `shift` 3.0 (rectified flow: the
- timesteps are `sigma * 1000`, a step is `x + (sigma_next - sigma) * v`); 28 steps at
- guidance 7.0 by default. No negative prompt means the encoded empty prompt.
-- **One tokenizer call**: `StableDiffusion3Tokenizer` runs the CLIP BPE (the second
- tower's `!`-padded ids are derived from the mask, as SDXL) and the T5 SentencePiece
- tokenizer, returning `{"input_ids", "attention_mask", "input_ids_3"}`.
-- **float16 weights**, the release precision; `from_weights` builds the model in float16
- by default (`load_dtype="float32"` for float32).
-
-The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion is
-deliberately **not supported** for diffusion models. The upstream repo is gated behind the
-Stability AI Non-Commercial Research Community License.
-
-Links:
-
-- Paper: [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (arXiv:2403.03206)](https://arxiv.org/abs/2403.03206)
-- Reference implementation: [diffusers `StableDiffusion3Pipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_3)
-- License: [Stability AI Non-Commercial Research Community License](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers/blob/main/LICENSE.md)
-
-See also [stable_diffusion_3_5.md](stable_diffusion_3_5.md), [stable_diffusion_xl.md](stable_diffusion_xl.md), [clip.md](clip.md).
-
-## Variants
-
-| Variant | Hub | Resolution | Sampler | License |
-|---|---|---|---|---|
-| `stable-diffusion-3-medium` | [`zeromodels/stable-diffusion-3-medium`](https://huggingface.co/zeromodels/stable-diffusion-3-medium) | 1024 | flow-match Euler (shift 3), 28 steps, guidance 7.0 | Stability AI Non-Commercial Research Community |
-| `t5-v1_1-xxl-encoder` | [`zeromodels/t5-v1_1-xxl-encoder`](https://huggingface.co/zeromodels/t5-v1_1-xxl-encoder) | text encoder (shared with SD 3.5) | | Apache 2.0 (Google T5 v1.1) |
-
-## API
-
-Configs are typed: `StableDiffusion3Config` (composite, `model_type` `"stable_diffusion_3"`)
-over `StableDiffusion3TransformerConfig`, `StableDiffusion3VAEConfig`,
-`StableDiffusion3TextConfig` (CLIP ViT-L/14 with its 768-d projection) and
-`StableDiffusionXLTextConfig2` (OpenCLIP ViT-bigG/14 with its 1280-d projection), plus
-the scheduler config, the CLIP and T5 token ids and `max_sequence_length` (256). Flat
-constructor with `transformer_` / `vae_` / `text_` / `text_2_` prefixes.
-
-### `StableDiffusion3TextToImage`
-
-`StableDiffusionTextToImage`'s `generate` over the SD 3 graph:
-
-```python
-generate(
- input_ids,
- attention_mask=None,
- negative_input_ids=None,
- num_inference_steps=None,
- guidance_scale=None,
- seed=None,
- latents=None,
- image=None,
- strength=None,
- denoising_start=None,
- denoising_end=None,
- output_type="image",
- input_ids_3=None,
- negative_input_ids_3=None,
-)
-```
-
-| Argument | Description |
-|---|---|
-| `input_ids`, `attention_mask`, `input_ids_3` | `**tokenizer(prompts)`: the CLIP ids and mask (the second tower's ids are derived from the mask) and the T5 ids (`(batch, 256)`). Without `input_ids_3`, or without an attached T5, the T5 features are zeros. |
-| `negative_input_ids`, `negative_input_ids_3` | The tokenized negative prompt (`neg = tokenizer("..."); negative_input_ids=neg["input_ids"], negative_input_ids_3=neg["input_ids_3"]`); the empty prompt when omitted. |
-| `num_inference_steps`, `guidance_scale`, `seed`, `latents`, `image`, `strength`, ... | As for [`BaseDiffusion.generate`](main_classes.md#basediffusion); defaults from the repo's `generate_args` (28 steps, guidance 7.0). `latents` is `(batch, 128, 128, 16)` at 1024px. |
-
-Returns `(batch, H, W, 3)` uint8 images.
-
-`from_weights` takes one extra argument, `text_encoder_3`: a hosted `SD3T5EncoderModel`
-repo id (loaded at the same `load_dtype`) or a built `SD3T5EncoderModel`;
-`sd.text_encoder_3` can also be assigned later (or set to `None`). It lives outside the
-container's weights.
-
-### `StableDiffusion3Model`
-
-The container: four disconnected paths in one functional model, `.transformer` / `.vae` /
-`.text_encoder` / `.text_encoder_2`.
-
-| Inputs | Outputs |
-|---|---|
-| `sample` (B, 128, 128, 16), `timestep` (B,), `encoder_hidden_states` (B, 333, 4096), `pooled_projections` (B, 2048) | `noise_pred` (B, 128, 128, 16) |
-| `image` (B, 1024, 1024, 3), `latent` (B, 128, 128, 16) | `moments` (B, 128, 128, 32), `image` (B, 1024, 1024, 3) |
-| `token_ids` (B, 77), `token_ids_2` (B, 77), `padding_mask` (B, 77) | `prompt_embeds` (B, 77, 2048), `pooled_prompt_embeds` (B, 2048) |
-
-### `SD3Transformer2DModel`
-
-The MMDiT on its own (`{"sample", "timestep", "encoder_hidden_states",
-"pooled_projections"} -> {"sample"}`), built from `StableDiffusion3TransformerConfig`
-(`sample_size` 128, `patch_size` 2, `num_layers` 24, `num_attention_heads` 24 x
-`attention_head_dim` 64, `joint_attention_dim` 4096, `caption_projection_dim` 1536,
-`pooled_projection_dim` 2048, `pos_embed_max_size` 192, `qk_norm`,
-`dual_attention_layers`). Its layers (`StableDiffusion3PatchEmbed`,
-`StableDiffusion3AdaLayerNorm`, `StableDiffusion3JointAttention`,
-`StableDiffusion3JointTransformerBlock`, `StableDiffusion3GELUFeedForward`) are named
-with the diffusers module paths.
-
-The joint attention runs over 4429 tokens at 1024px, so `StableDiffusion3JointAttention`
-defaults to the `"fused"` attention implementation (`keras.ops.dot_product_attention`,
-the backend's own kernel: torch's flash / memory-efficient kernels never materialize the
-4429 x 4429 logits; see [`fused_attention`](main_classes.md#fused_attention)) instead of
-the library-wide `"sdpa"` math. `from_weights(..., attn_implementation="sdpa")` switches
-it back.
-
-### `SD3T5EncoderModel`
-
-The third text encoder, the encoder half of T5 v1.1 XXL (transformers `T5EncoderModel` on
-`google/t5-v1_1-xxl`), as its own model in this family: `{"input_ids", "attention_mask"}`
-(`(batch, 256)` int32) `-> {"last_hidden_state": (batch, 256, 4096)}`, built from
-`StableDiffusion3T5EncoderConfig` (`vocab_size` 32128, `embed_dim` 4096, `key_value_dim`
-64, `mlp_dim` 10240, `num_layers` 24, `num_heads` 64, `relative_attention_num_buckets` 32,
-`relative_attention_max_distance` 128, `layer_norm_eps` 1e-6; 4.7B parameters). The blocks
-reuse the [T5](t5.md) family's attention, RMSNorm and relative position bias, and add the
-gated-GELU feed-forward of T5 v1.1 (`wo(gelu_tanh(wi_0(x)) * wi_1(x))`,
-`StableDiffusion3T5GatedFeedForward`), which the original T5 does not have. Loads on its own with
-`SD3T5EncoderModel.from_weights("zeromodels/t5-v1_1-xxl-encoder")` (`load_dtype`,
-`quantization="int8"` as needed); its inputs are the `input_ids_3` of the SD 3 tokenizer.
-
-## Preprocessing
-
-### `StableDiffusion3Tokenizer`
-
-```python
-StableDiffusion3Tokenizer(
- hf_id=None,
- tokenizer_file=None,
- tokenizer_file_3=None,
- max_seq_len=77,
- max_sequence_length=256,
-)
-```
-
-Returns `{"input_ids", "attention_mask", "input_ids_3"}`; pass all three to `generate`.
-
-## End-to-end example
-
-```python
-import os
-
-os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow"
-
-from PIL import Image
-from zeromodels.models.stable_diffusion_3 import (
- StableDiffusion3TextToImage,
- StableDiffusion3Tokenizer,
-)
-
-model = StableDiffusion3TextToImage.from_weights(
- "zeromodels/stable-diffusion-3-medium",
- text_encoder_3="zeromodels/t5-v1_1-xxl-encoder", # omit on small machines: T5 features zeroed
-)
-tokenizer = StableDiffusion3Tokenizer.from_weights(
- "zeromodels/stable-diffusion-3-medium"
-)
-
-inputs = tokenizer(
- "a fluffy corgi wearing round sunglasses, sitting on a surfboard at a sunny beach, photo"
-)
-images = model.generate(**inputs, num_inference_steps=28, guidance_scale=7.0, seed=4)
-
-Image.fromarray(images[0]).save("corgi.png") # (1024, 1024, 3) uint8
-```
-
-zm_config.json + the float16 weights +
-tokenizer.json / tokenizer_3.json). Load with
-from_weights("zeromodels/<variant>").
-
-
-### The T5 encoder on the CPU
-
-The 9.5 GB T5-XXL rarely fits next to the MMDiT on a consumer GPU. Encode the prompt
-separately (an `SD3T5EncoderModel` in another process, or int8) and hand the ids to a
-model without T5 for the CLIP part only, or attach it in full on a large machine. The
-`encode_prompt` hook returns the conditioning dict the transformer takes, so any
-precomputed T5 features can be substituted there.
-
-### Verified against diffusers
-
-The components of `stable-diffusion-3-medium` were checked one by one against diffusers
-0.39 / transformers 5 in float32 on the same inputs (the converted float16 weights
-computed in float32):
-
-```
-text_encoder penultimate max|d|=6.1e-05 (values up to 854)
-text_encoder_2 penultimate max|d|=3.1e-05
-pooled (L + G, projected) max|d|=3.8e-06
-t5-v1_1-xxl encoder (256 tok) max|d|=1.1e-04 (values up to 6.3)
-mmdit velocity (t=1000) max|d|=9.9e-06
-vae decode (512px) max|d|=6.3e-05
-vae encode max|d|=1.4e-05
-```
-
-The pipeline (flow-match Euler with shift 3, classifier-free guidance with the encoded
-empty prompt, the three tokenizers, batched negative prompts, the no-T5 mode) was checked
-against `StableDiffusion3Pipeline` on a small random-weight SD 3, where the generated images
-are pixel-identical on torch and jax. End to end, `stable-diffusion-3-medium` without the
-T5 encoder (28 steps, guidance 7, the same initial latent) reproduces the diffusers
-pipeline in float32 at 512px: 92.8% identical pixels, max |d| 2 / 255, PSNR 63.9 dB (the
-final latent differs by at most 2.5e-2 with values up to 5.1). In float16 on the GPU the
-two implementations round differently through the 28 guided steps: the pictures agree in
-composition but not pixel for pixel (PSNR 26 dB at 1024px, 29 dB at 768px, float16 MMDiT
-on both sides), the usual float16 sampling noise.
-
-## Data Format
-
-Both `channels_last` and `channels_first` are supported; `generate` always returns
-`(batch, H, W, 3)` uint8.
-
-## Memory and speed
-
-The container is 5.8 GB in float16 (the T5 encoder adds 9.5 GB). Measured on the torch
-backend with the default float16 load (RTX 4060 Laptop, 8 GB), T5 features precomputed:
-
-| Resolution | Denoising step (guided, batch 1) | VAE decode (float32) |
-|---|---|---|
-| 512px | peak 6.1 GB, 0.5 s/step | peak 7.6 GB, 4 s |
-| 768px | peak 6.4 GB, 0.8 s/step | peak 9.5 GB |
-| 1024px | peak 6.9 GB, 1.0 s/step | peak 12.4 GB |
-
-The MMDiT itself fits an 8 GB GPU at the native 1024px thanks to the fused attention
-kernel (the `"sdpa"` math needs 3.8 GB per block for the 4429 x 4429 float32 logits and
-runs out of memory). The VAE decoder is the SDXL one (16 latent channels aside) and Keras'
-functional executor keeps its feature maps alive, see the
-[SDXL notes](stable_diffusion_xl.md#memory-and-speed): on 8 GB, denoise on the GPU
-(`output_type="latent"`) and decode on the CPU (75 to 90 s at 1024px). The T5-XXL encoder
-encodes a prompt and its negative on the CPU in about 10 to 25 s (float32 compute over
-the float16 weights); a 24 GB GPU holds everything. On JAX and TensorFlow the denoiser
-step is compiled once per run.
-
-## Loading Fine-tuned Weights
-
-Any repo laid out like the hosted ones (`zm_config.json` declaring
-`StableDiffusion3Model`, the weights, the two tokenizer files) loads with
-`from_weights("
-Weights: pretrained Keras weights live on Hugging Face under
-zeromodels/<variant>
-(each repo carries
-
-Stable Diffusion 3.5, ported to pure Keras 3: the [Stable Diffusion 3](stable_diffusion_3.md)
-architecture, retrained and scaled. It reuses the SD 3 family end to end
-(`StableDiffusion3_5Model` is `StableDiffusion3Model`, `StableDiffusion3_5TextToImage` is
-`StableDiffusion3TextToImage`, each with the SD 3.5 configuration), so everything on that
-page applies: the MMDiT, the three text encoders with the separately hosted T5-XXL, the
-16-channel VAE, the flow-match sampler, one tokenizer call. What changes is the config:
-
-- **QK normalization**: every attention block RMS-normalizes its queries and keys per head
- (`qk_norm="rms_norm"`), the training stabilizer of SD 3.5.
-- **Large** (8B): 38 blocks of 38 heads (2432-d tokens), 28 steps at guidance 3.5.
-- **Large Turbo**: the large model distilled with Adversarial Diffusion Distillation to 4
- steps without guidance (`guidance_scale=0.0`).
-- **Medium** (2.5B, MMDiT-X): 24 blocks of 24 heads with **dual attention** in blocks 0 to
- 12 (a second, latent-only attention with its own modulation next to the joint one) and a
- 384x384 position grid (up to 2 MP); 40 steps at guidance 4.5.
-
-The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion is
-deliberately **not supported** for diffusion models. The upstream repos are gated behind
-the Stability AI Community License.
-
-Links:
-
-- Paper: [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (arXiv:2403.03206)](https://arxiv.org/abs/2403.03206)
-- Reference implementation: [diffusers `StableDiffusion3Pipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_3)
-- License: [Stability AI Community License](https://huggingface.co/stabilityai/stable-diffusion-3.5-large/blob/main/LICENSE.md)
-
-See also [stable_diffusion_3.md](stable_diffusion_3.md).
-
-## Variants
-
-Each repo is one container: MMDiT + VAE (84M) + CLIP ViT-L/14 (123M) + OpenCLIP
-ViT-bigG/14 (695M); the T5-XXL encoder (`zeromodels/t5-v1_1-xxl-encoder`, the SD 3
-family's `SD3T5EncoderModel`, 9.5 GB) is shared and optional.
-
-| Variant | Hub | MMDiT | Sampler | Size |
-|---|---|---|---|---|
-| `stable-diffusion-3.5-large` | [`zeromodels/stable-diffusion-3.5-large`](https://huggingface.co/zeromodels/stable-diffusion-3.5-large) | 38 x 38 heads, 8B | flow-match Euler (shift 3), 28 steps, guidance 3.5 | 17.8 GB float16 |
-| `stable-diffusion-3.5-large-turbo` | [`zeromodels/stable-diffusion-3.5-large-turbo`](https://huggingface.co/zeromodels/stable-diffusion-3.5-large-turbo) | 38 x 38 heads, 8B | 4 steps, no guidance | 17.8 GB float16 |
-| `stable-diffusion-3.5-medium` | [`zeromodels/stable-diffusion-3.5-medium`](https://huggingface.co/zeromodels/stable-diffusion-3.5-medium) | 24 x 24 heads + dual attention, 2.5B | flow-match Euler (shift 3), 40 steps, guidance 4.5 | 6.8 GB float16 |
-
-## API
-
-`StableDiffusion3_5Config` (`model_type` `"stable_diffusion_3_5"`) is
-`StableDiffusion3Config` over `StableDiffusion3_5TransformerConfig` (the SD 3.5 large
-defaults: `num_layers` 38, `num_attention_heads` 38, `caption_projection_dim` 2432,
-`qk_norm="rms_norm"`; medium's repo carries `dual_attention_layers=(0, ..., 12)` and
-`pos_embed_max_size=384`). `StableDiffusion3_5TextToImage.generate` and
-`StableDiffusion3_5Tokenizer` are the SD 3 ones; see
-[Stable Diffusion 3](stable_diffusion_3.md#api).
-
-## End-to-end example
-
-```python
-import os
-
-os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow"
-
-from PIL import Image
-from zeromodels.models.stable_diffusion_3_5 import (
- StableDiffusion3_5TextToImage,
- StableDiffusion3_5Tokenizer,
-)
-
-model = StableDiffusion3_5TextToImage.from_weights(
- "zeromodels/stable-diffusion-3.5-medium",
- text_encoder_3="zeromodels/t5-v1_1-xxl-encoder", # optional
-)
-tokenizer = StableDiffusion3_5Tokenizer.from_weights(
- "zeromodels/stable-diffusion-3.5-medium"
-)
-
-inputs = tokenizer(
- "a whimsical treehouse village at twilight, lanterns glowing, watercolor illustration"
-)
-images = model.generate(**inputs, num_inference_steps=40, guidance_scale=4.5, seed=5)
-
-Image.fromarray(images[0]).save("treehouse.png") # (1024, 1024, 3) uint8
-```
-
-zm_config.json + the float16 weights +
-tokenizer.json / tokenizer_3.json). Load with
-from_weights("zeromodels/<variant>").
-
-
-### Large Turbo
-
-```python
-model = StableDiffusion3_5TextToImage.from_weights(
- "zeromodels/stable-diffusion-3.5-large-turbo"
-)
-images = model.generate(
- **tokenizer("a red panda reading a book")
-) # 4 steps, guidance 0 (the repo defaults)
-```
-
-### Verified against diffusers
-
-The components of `stable-diffusion-3.5-medium` (the MMDiT-X with dual attention and RMS
-qk norms) were checked against diffusers 0.39 in float32 on the same inputs (the converted
-float16 weights computed in float32): MMDiT velocity max|d| 1.6e-5 (values up to 5.2), the
-text encoders and the VAE as for [SD 3](stable_diffusion_3.md#verified-against-diffusers).
-The SD 3.5 block variants (dual attention, `qk_norm`) are also covered by the small
-random-weight pipeline test, pixel-identical to `StableDiffusion3Pipeline` on torch and
-jax, and by the SD 3 float32 end-to-end check (92.8% identical pixels, PSNR 63.9 dB at
-512px), which the two families share down to the sampler. The picture above is the
-1024px float16 output of `stable-diffusion-3.5-medium` with the T5 encoder.
-
-## Memory and speed
-
-The medium container is 6.8 GB in float16, the large ones 17.8 GB (a 24 GB GPU); the
-T5-XXL encoder adds 9.5 GB. See [Stable Diffusion 3](stable_diffusion_3.md#memory-and-speed):
-the medium MMDiT-X at 1024px peaks at 7.5 GB during the guided denoising step on the torch
-backend (fused attention, float16), which an 8 GB GPU runs with the VAE decode on the CPU.
-
-## Loading Fine-tuned Weights
-
-Any repo laid out like the hosted ones (`zm_config.json` declaring
-`StableDiffusion3_5Model`, the weights, the two tokenizer files) loads with
-`from_weights("
-Weights: pretrained Keras weights live on Hugging Face under
-zeromodels/<variant>
-(each repo carries
-
-Stable Diffusion XL, ported to pure Keras 3: the [Stable Diffusion](stable_diffusion.md)
-latent-diffusion recipe scaled up to 1024px with a 2.6B-parameter UNet and two text
-encoders, plus the **refiner**, an image-to-image expert that finishes the base model's
-latents, and **SDXL-Turbo**, the base distilled to a few steps. It subclasses the SD 1.x
-family (`StableDiffusionXLModel` is a `StableDiffusionModel` with a fourth component,
-`StableDiffusionXLTextToImage` is `StableDiffusionTextToImage` with the SDXL conditioning
-hooks; the refiner classes subclass those), so everything on that page applies: one
-container per repo, `generate` from `BaseDiffusion`, the schedulers, both data formats.
-What changes:
-
-- **Two text encoders**: CLIP ViT-L/14's text tower (768-d, the SD 1.x one) and OpenCLIP
- ViT-bigG/14's (1280-d, 32 layers, `gelu`, with a 1280-d `text_projection`). The prompt
- goes through both; their **penultimate** hidden states (no final LayerNorm) are
- concatenated into the UNet's 2048-d cross-attention context, and the second tower's
- projected pooled (EOT) state is the `text_embeds` micro-conditioning.
-- **One tokenizer call**: both towers use the same CLIP BPE and differ only in the pad
- token (`<|endoftext|>` vs `!`), so `StableDiffusionXLTokenizer` is the first one and the
- model derives the second tower's ids from the returned `attention_mask`.
-- **Micro-conditioning**: the UNet's timestep embedding also receives the pooled text
- embedding and six size / crop values (`original_size`, `crops_coords_top_left`,
- `target_size`, each `(height, width)`), sinusoidally embedded and projected by the
- `add_embedding` MLP (`addition_embed_type="text_time"`).
-- **UNet**: three levels (320, 640, 1280) with no attention at the first, (1, 2, 10)
- transformer blocks per Transformer2D, (5, 10, 20) heads (64 wide), linear token
- projection, built at a 128x128 latent (1024px).
-- **VAE**: the SDXL autoencoder (`scaling_factor` 0.13025), built in **float32 whatever
- the load dtype** (`force_upcast`: it overflows in float16, as in diffusers).
-- **Classifier-free guidance** with no negative prompt uses **zero** embeddings
- (`force_zeros_for_empty_prompt`), not the encoded empty prompt. Default guidance is 5.0.
-- **Scheduler**: Euler with `leading` timestep spacing (`steps_offset` 1) for the base
- models and the refiners, ancestral Euler with `trailing` spacing for SDXL-Turbo; all
- come from the repo's `scheduler_config`.
-- **Refiner**: a second model (`StableDiffusionXLRefinerModel` /
- `StableDiffusionXLRefinerImageToImage`) with the OpenCLIP tower alone (1280-d
- context), a four-level UNet ((384, 768, 1536, 1536), attention at the middle levels, 4
- blocks per Transformer2D, 2.3B parameters) and an **aesthetic score** as the fifth
- micro-conditioning id. It refines a latent the base left partially denoised
- (`denoising_start`) or any image (`strength`); its empty negative prompt is encoded, not
- zeroed.
-- **float16 weights**: SDXL was released in float16, so the repos store that (the VAE in
- float32) and `from_weights` builds the model in float16 by default (6.6 GB of weights).
- Pass `load_dtype="float32"` for a float32 model (13.9 GB).
-
-The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion is
-deliberately **not supported** for diffusion models.
-
-Links:
-
-- Paper: [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis (arXiv:2307.01952)](https://arxiv.org/abs/2307.01952)
-- Reference implementation: [diffusers `StableDiffusionXLPipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl)
-- Licenses: [CreativeML Open RAIL++-M](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/LICENSE.md)
- (1.0 base and refiner), [Stability AI Non-Commercial Research Community License](https://huggingface.co/stabilityai/sdxl-turbo/blob/main/LICENSE.md)
- (Turbo); the 0.9 research previews are not hosted
-
-See also [stable_diffusion.md](stable_diffusion.md), [stable_diffusion_2.md](stable_diffusion_2.md), [clip.md](clip.md).
-
-## Variants
-
-Preconverted weights are hosted under `zeromodels/`. Load with
-`from_weights("zeromodels/zm_config.json + the float16 weights +
-tokenizer.json). Load with from_weights("zeromodels/<variant>").
-
-
-### Base + refiner
-
-The SDXL "ensemble of experts": the base runs the first 80% of the schedule and hands its
-latent to the refiner, which finishes it. Same prompt and initial latent as above:
-
-```python
-refiner = StableDiffusionXLRefinerImageToImage.from_weights(
- "zeromodels/stable-diffusion-xl-refiner-1.0"
-)
-latent = model.generate(
- **inputs, num_inference_steps=30, denoising_end=0.8, output_type="latent", seed=3
-)
-images = refiner.generate(
- **inputs, latents=latent, num_inference_steps=30, denoising_start=0.8
-)
-```
-
-
-
-### SDXL-Turbo
-
-One step, no guidance, 512px (the repo's defaults):
-
-```python
-model = StableDiffusionXLTextToImage.from_weights("zeromodels/sdxl-turbo")
-tokenizer = StableDiffusionXLTokenizer.from_weights("zeromodels/sdxl-turbo")
-images = model.generate(
- **tokenizer(
- "a cinematic shot of a baby raccoon wearing an intricate italian priest robe"
- )
-)
-```
-
-
-
-### Negative prompts, micro-conditioning
-
-```python
-inputs = tokenizer(["a red bicycle", "a blue bicycle"])
-negative = tokenizer("blurry, low quality") # one row, shared by the batch
-images = model.generate(
- **inputs,
- negative_input_ids=negative["input_ids"],
- original_size=(2048, 2048), # "a crop of a larger image": sharper detail
- crops_coords_top_left=(0, 0),
- target_size=(1024, 1024),
- seed=0,
-)
-```
-
-### Other resolutions
-
-The graphs are built for the repo's resolution; pass overrides to build for another
-(the weights are resolution-independent; multiples of 64px):
-
-```python
-model = StableDiffusionXLTextToImage.from_weights(
- "zeromodels/stable-diffusion-xl-base-1.0",
- unet_sample_size=(96, 128),
- vae_sample_size=(768, 1024),
-)
-images = model.generate(**tokenizer("a wide mountain valley")) # (1, 768, 1024, 3)
-```
-
-### float32
-
-```python
-model = StableDiffusionXLTextToImage.from_weights(
- "zeromodels/stable-diffusion-xl-base-1.0", load_dtype="float32"
-)
-```
-
-### Verified against diffusers
-
-The components of `stable-diffusion-xl-base-1.0` and `stable-diffusion-xl-refiner-1.0`
-were checked one by one against diffusers 0.39 / transformers in float32 on the same
-inputs (the converted float16 weights computed in float32):
-
-```
-base text_encoder penultimate max|d|=6.1e-05 (values up to 854)
-base text_encoder_2 penultimate max|d|=3.1e-05
-base text_encoder_2 pooled max|d|=3.6e-06
-base unet noise_pred (t=981) max|d|=9.3e-06
-base vae decode (512px) max|d|=3.5e-05
-refiner text_encoder_2 penultimate max|d|=3.1e-05
-refiner unet noise_pred (t=181) max|d|=1.2e-05
-```
-
-The pipelines (Euler `leading` spacing, classifier-free guidance with the zeroed or
-encoded negative branch, the micro-conditioning values, negative prompts, batches,
-`strength`, `denoising_end` / `denoising_start`) were checked against
-`StableDiffusionXLPipeline` and `StableDiffusionXLImg2ImgPipeline` on small
-random-weight models, where the generated images are pixel-identical on torch and jax.
-At full size in float16 (the default load; diffusers in float16 with its float32 VAE),
-the lighthouse above matches diffusers to PSNR 49.7 dB (98.5% of pixels within 2 uint8
-levels), the base + refiner ensemble to PSNR 51.9 dB (98.8% within 2) and SDXL-Turbo
-(1 step, 512px) to PSNR 44.9 dB (92% within 2; a single step amplifies float16 rounding).
-The lighthouse, zeromodels (left) and diffusers (right):
-
-
-
-## Data Format
-
-Both `channels_last` and `channels_first` are supported, as for
-[Stable Diffusion](stable_diffusion.md#data-format); `generate` always returns
-`(batch, H, W, 3)` uint8.
-
-## Memory and speed
-
-Measured on the torch backend with the default float16 load (weights resident: 6.8 GB for
-a base model, 6.3 GB for the refiner, the VAE in float32):
-
-| Resolution | Denoising step (guided, batch 1) | VAE decode (float32) |
-|---|---|---|
-| 512px | peak 7.1 GB, 1.2 s/step on an RTX 4060 Laptop | peak 8.5 GB |
-| 768px | peak 8.0 GB | peak 10.5 GB |
-| 1024px | peak 10.4 GB | about 14 GB |
-
-Keras' functional executor keeps every layer output of a graph alive until the graph
-returns, so the decoder's feature maps at 1024px add several GB on top of the weights: a
-16 GB GPU runs 1024px comfortably, a 12 GB GPU 768px; on 8 GB, generate at 512px, or
-denoise on the GPU (`output_type="latent"`) and decode the latent on the CPU (about 100 s
-at 1024px) with a second, CPU-only process. In float32 the weights alone take 13.9 GB.
-The UNet's attention runs the portable `"sdpa"` math by default, which materializes the
-4096 x 4096 float32 logits of the 64 x 64 level at 1024px; loading with
-`attn_implementation="fused"` (the backend's fused kernel, see
-[`fused_attention`](main_classes.md#fused_attention)) cuts the 1024px step from a 10.3 GB
-peak to 7.6 GB on the same GPU. On JAX and TensorFlow the denoiser step is compiled once
-per run (see [`BaseDiffusion`](main_classes.md#basediffusion)).
-
-## Loading Fine-tuned Weights
-
-Any repo laid out like the hosted ones (`zm_config.json` declaring
-`StableDiffusionXLModel`, the weights, `tokenizer.json`) loads with
-`from_weights("