README.md

---
pipeline_tag: text-to-image
license: other
license_name: stabilityai-nc-research-community
license_link: https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers/blob/main/LICENSE.md
base_model: stabilityai/stable-diffusion-3-medium-diffusers
library_name: zeromodels
language:
- en
tags:
- keras
- zeromodels
- stable-diffusion
- stable-diffusion-3
- sd3
- mmdit
- text-to-image
- diffusion
- rectified-flow
- arxiv:2403.03206
- pytorch
- jax
- tf
---
*See [our collection](https://huggingface.co/collections/zeromodels/stable-diffusion-v3-6aa7941443a0c5ca085ebcc9) for all Stable Diffusion 3 checkpoints.*

# Run Stable Diffusion 3 with Keras 3: JAX, PyTorch, or TensorFlow

[![GitHub](https://img.shields.io/badge/GitHub-ZeroModels-181717?logo=github)](https://github.com/IMvision12/ZeroModels) [![Docs](https://img.shields.io/badge/Docs-Stable_Diffusion_3-1f6feb)](https://imvision12.github.io/ZeroModels/stable_diffusion_3/) [![HuggingFace](https://img.shields.io/badge/HuggingFace-Stable_Diffusion_3-ffd21e?logo=huggingface&logoColor=black)](https://huggingface.co/collections/zeromodels/stable-diffusion-v3-6aa7941443a0c5ca085ebcc9)

# zeromodels/stable-diffusion-3-medium

Paper: [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (arXiv:2403.03206)](https://arxiv.org/abs/2403.03206) | [HF Papers](https://huggingface.co/papers/2403.03206)

Pure-**Keras 3** conversion of [`stabilityai/stable-diffusion-3-medium-diffusers`](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers) for
[zeromodels](https://github.com/IMvision12/ZeroModels). One implementation runs unmodified on
**TensorFlow / Torch / JAX**. The container ships the MMDiT denoiser, the 16-channel VAE
and the two CLIP text encoders in `model.weights.json` shards (2.99B parameters,
5.72 GB), plus `zm_config.json` (the component configs, the checkpoint's
`FlowMatchEulerDiscreteScheduler` (shift 3.0) and the default
generation settings) and the two tokenizers (`tokenizer.json`, the CLIP BPE, and
`tokenizer_3.json`, the T5 SentencePiece). Weights are stored in **float16**, the
checkpoint's native precision (the VAE in float32), and load in float16 by default; pass
`load_dtype="float32"` to `from_weights` for a float32 model. This checkpoint generates
**1024x1024** images (a 128x128x16 latent).

The third text encoder, the 4.7B-parameter **T5-XXL**, is shared by every SD 3 / 3.5
checkpoint and hosted once at
[`zeromodels/t5-v1_1-xxl-encoder`](https://huggingface.co/zeromodels/t5-v1_1-xxl-encoder); attach it
with `text_encoder_3=` (below) or leave it out (the T5 features are zeroed, SD 3's
memory-saving mode).

For model details, intended use and limitations, see the upstream
[model card](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers).

## Architecture

| Component | zeromodels class | Details |
| --- | --- | --- |
| Denoiser | `SD3Transformer2DModel` | MMDiT: 24 joint blocks of 24 x 64-d heads (1536 wide), patch 2, 4096-d text features projected to 1536, 2048-d pooled conditioning, 192x192 position grid |
| Autoencoder | `AutoencoderKL` | (128, 256, 512, 512) channels, x8 spatial compression to 16 latent channels, `scaling_factor` 1.5305, `shift_factor` 0.0609, float32 (`force_upcast`) |
| Text encoder | functional CLIP text tower | CLIP ViT-L/14: 768-d, 12 layers, `quick_gelu`, 768-d projection; penultimate hidden state + projected pooled state |
| Text encoder 2 | functional CLIP text tower | OpenCLIP ViT-bigG/14: 1280-d, 32 layers, `gelu`, 1280-d projection; penultimate hidden state + projected pooled state |
| Text encoder 3 | `SD3T5EncoderModel` (separate repo) | T5 v1.1 XXL encoder: 4096-d, 24 layers, gated GELU, 256 tokens |
| Scheduler | `FlowMatchEulerDiscreteScheduler` | rectified flow over 1000 timesteps, shift 3.0 |

## Quick start

```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",  # optional: omit to zero the T5 features
)
tokenizer = StableDiffusion3Tokenizer.from_weights("zeromodels/stable-diffusion-3-medium")

inputs = tokenizer("a photograph of an astronaut riding a horse")
images = model.generate(**inputs, num_inference_steps=28, guidance_scale=7.0, seed=0)
Image.fromarray(images[0]).save("astronaut.png")  # (1024, 1024, 3) uint8
```

`generate` takes the tokenizer's `input_ids` / `attention_mask` / `input_ids_3` (batch them
for several prompts), an optional tokenized negative prompt (`negative_input_ids` /
`negative_input_ids_3`), `num_inference_steps`, `guidance_scale`, a `seed`, or explicit
`latents` of shape `(batch, 128, 128, 16)` for results that are identical across
backends; `image` / `strength` refine an image instead.

Load any Stable Diffusion 3 checkpoint the same way with `from_weights("zeromodels/<variant>")`:

| Variant | Hub | Training |
| --- | --- | --- |
| `stable-diffusion-3-medium` | [zeromodels/stable-diffusion-3-medium](https://huggingface.co/zeromodels/stable-diffusion-3-medium) | 1024px, rectified flow (shift 3), 28 steps at guidance 7.0: the 2B MMDiT of SD 3 medium |

## Tips

- Set `KERAS_BACKEND` **before** importing Keras / zeromodels.
- The graphs are built for 1024px. Pass `transformer_sample_size=<px / 8>, vae_sample_size=<px>` to
  `from_weights` to build for another multiple of 16px (the weights are resolution-independent
  up to the 3072px position grid).
- `StableDiffusion3Model.from_weights(...)` loads the same repo as the bare container
  (`.transformer` / `.vae` / `.text_encoder` / `.text_encoder_2`) without the generation loop.
- `model.text_encoder_3` can be any `SD3T5EncoderModel` (for example one loaded with
  `quantization="int8"`); it is not part of the container's weights.
- Both `channels_last` and `channels_first` are supported (`keras.config.set_image_data_format`
  before loading); `generate` always returns `(batch, H, W, 3)` uint8.
- On-the-fly `hf:` conversion is not supported for diffusion models; the checkpoints are
  hosted here, converted once.
- See the [Stable Diffusion 3 docs](https://imvision12.github.io/ZeroModels/stable_diffusion_3/).

## License

The weights are redistributed under the [Stability AI Non-Commercial Research Community License](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers/blob/main/LICENSE) of the upstream
checkpoint, including its use-based restrictions. By using them you agree to those terms.

## Notice

This Stability AI Model is licensed under the Stability AI Non-Commercial Research Community License, Copyright (c) Stability AI Ltd. All Rights Reserved.

Modifications by zeromodels (https://github.com/IMvision12/ZeroModels): the checkpoint
released at https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers was
converted to the Keras 3 weights layout of zeromodels (`model.weights.json,
model_00000.weights.h5, model_00001.weights.h5, zm_config.json, tokenizer.json,
tokenizer_3.json`), stored in float16, the upstream fp16 files, with the VAE in float32. The
model architecture and the parameter values are unchanged; the weight names and the file
format differ from the release.

## Special Thanks

Thank you to Stability AI for training and releasing Stable Diffusion, and to the
Hugging Face diffusers team, whose implementation this port was verified against.