File size: 8,247 Bytes
63e89bd
 
 
3b6a55a
 
63e89bd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
af3da36
63e89bd
c2e272b
 
63e89bd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
af3da36
63e89bd
af3da36
63e89bd
 
 
af3da36
63e89bd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
af3da36
 
63e89bd
 
 
af3da36
 
63e89bd
af3da36
63e89bd
 
 
 
af3da36
63e89bd
 
 
c2e272b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
63e89bd
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
---
pipeline_tag: text-to-image
license: other
license_name: stabilityai-ai-community
license_link: https://huggingface.co/stabilityai/stable-diffusion-3.5-large/blob/main/LICENSE.md
base_model: stabilityai/stable-diffusion-3.5-large
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-v35-6aa7963880e76656e5646ffc) for all Stable Diffusion 3.5 checkpoints.*

# Run Stable Diffusion 3.5 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.5-1f6feb)](https://imvision12.github.io/ZeroModels/diffusers/#stable-diffusion-3-5) [![HuggingFace](https://img.shields.io/badge/HuggingFace-Stable_Diffusion_3.5-ffd21e?logo=huggingface&logoColor=black)](https://huggingface.co/collections/zeromodels/stable-diffusion-v35-6aa7963880e76656e5646ffc)

**Powered by Stability AI**

# zeromodels/stable-diffusion-3.5-large

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.5-large`](https://huggingface.co/stabilityai/stable-diffusion-3.5-large) 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 (9.05B parameters,
17.01 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.5-large).

## Architecture

| Component | zeromodels class | Details |
| --- | --- | --- |
| Denoiser | `SD3Transformer2DModel` | MMDiT: 38 joint blocks of 38 x 64-d heads (2432 wide), patch 2, 4096-d text features projected to 2432, 2048-d pooled conditioning, 192x192 position grid, RMS-normalized queries and keys |
| 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.diffusers import ZMDiffusion, ZMDiffusionTokenizer

model = ZMDiffusion.from_weights(
    "zeromodels/stable-diffusion-3.5-large",
    text_encoder_3="zeromodels/t5-v1_1-xxl-encoder",  # optional: omit to zero the T5 features
)
tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-3.5-large")

inputs = tokenizer("a photograph of an astronaut riding a horse")
images = model.generate(**inputs, num_inference_steps=28, guidance_scale=3.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](https://huggingface.co/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](https://huggingface.co/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](https://huggingface.co/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.
- `AutoZModel` / `AutoZMTokenizer` (`from zeromodels import ...`) load this repo too;
  they resolve to `ZMDiffusion` / `ZMDiffusionTokenizer` from its `model_type`.
- 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).
- The trained components are attributes of the loaded `ZMDiffusion`
  (`.transformer` / `.vae` / `.text_encoder` / `.text_encoder_2`), each callable on its own.
- `model.text_encoder_3` can be any `SD3T5EncoderModel` (for example one loaded with
  `quantization="int8"`); it is not part of this checkpoint'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](https://imvision12.github.io/ZeroModels/diffusers/#stable-diffusion-3-5).

## License

The weights are redistributed under the [Stability AI Community License](https://huggingface.co/stabilityai/stable-diffusion-3.5-large/blob/main/LICENSE.md) 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-large was converted to
the Keras 3 weights layout of zeromodels (`model.weights.json, model_00000.weights.h5,
model_00001.weights.h5, model_00002.weights.h5, model_00003.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.