zeromodels/stable-diffusion-3-medium

See our collection for all Stable Diffusion 3 checkpoints.

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

zeromodels/stable-diffusion-3-medium

Paper: Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (arXiv:2403.03206) | HF Papers

Pure-Keras 3 conversion of stabilityai/stable-diffusion-3-medium-diffusers for 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; 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.

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

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	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.

License

The weights are redistributed under the Stability AI Non-Commercial Research Community 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.