--- 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 | 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=, vae_sample_size=` 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.