zeromodels/stable-diffusion-3.5-medium

See our collection for all Stable Diffusion 3.5 checkpoints.

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

Powered by Stability AI

zeromodels/stable-diffusion-3.5-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.5-medium 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 (3.37B parameters, 6.44 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, 384x384 position grid, RMS-normalized queries and keys and dual attention in blocks 0 to 12
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_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: omit to zero the T5 features
)
tokenizer = StableDiffusion3_5Tokenizer.from_weights("zeromodels/stable-diffusion-3.5-medium")

inputs = tokenizer("a photograph of an astronaut riding a horse")
images = model.generate(**inputs, num_inference_steps=40, guidance_scale=4.5, 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.5 checkpoint the same way with from_weights("zeromodels/<variant>"):

Variant	Hub	Training
stable-diffusion-3.5-large	zeromodels/stable-diffusion-3.5-large	1024px, rectified flow (shift 3), 28 steps at guidance 3.5: the 8B MMDiT with RMS-normalized queries and keys
stable-diffusion-3.5-large-turbo	zeromodels/stable-diffusion-3.5-large-turbo	1024px, 4 steps, no guidance: SD 3.5 large distilled with Adversarial Diffusion Distillation
stable-diffusion-3.5-medium	zeromodels/stable-diffusion-3.5-medium	1024px (up to 2 MP), rectified flow (shift 3), 40 steps at guidance 4.5: the 2.5B MMDiT-X with dual-attention blocks

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 6144px position grid).
• StableDiffusion3_5Model.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.5 docs.

License

The weights are redistributed under the Stability AI 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 Community License, Copyright © 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.5-medium 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.

Powered by Stability AI

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.