README.md

---
pipeline_tag: text-to-image
license: openrail++
base_model: sd2-community/stable-diffusion-2-base
library_name: zeromodels
language:
- en
tags:
- keras
- zeromodels
- stable-diffusion
- stable-diffusion-diffusers
- text-to-image
- diffusion
- latent-diffusion
- arxiv:2112.10752
- pytorch
- jax
- tf
---
*See [our collection](https://huggingface.co/collections/zeromodels/stable-diffusion-v2-6aa7906bed026a6b11f4e7be) for all Stable Diffusion 2.x checkpoints.*

# Run Stable Diffusion 2.x 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_2.x-1f6feb)](https://imvision12.github.io/ZeroModels/stable_diffusion_2/) [![HuggingFace](https://img.shields.io/badge/HuggingFace-Stable_Diffusion_2.x-ffd21e?logo=huggingface&logoColor=black)](https://huggingface.co/collections/zeromodels/stable-diffusion-v2-6aa7906bed026a6b11f4e7be)

# zeromodels/stable-diffusion-2-base

Paper: [High-Resolution Image Synthesis with Latent Diffusion Models (arXiv:2112.10752)](https://arxiv.org/abs/2112.10752) | [HF Papers](https://huggingface.co/papers/2112.10752)

Pure-**Keras 3** conversion of [`sd2-community/stable-diffusion-2-base`](https://huggingface.co/sd2-community/stable-diffusion-2-base) for
[zeromodels](https://github.com/IMvision12/ZeroModels). One implementation runs unmodified on
**TensorFlow / Torch / JAX**. The whole text-to-image model ships as **one container**:
the UNet denoiser, the VAE and the OpenCLIP ViT-H/14 text encoder (penultimate layer) in `model.weights.h5`
(1.29B parameters, 4.81 GB), plus `zm_config.json` (the three
component configs, the checkpoint's `DDIMScheduler` schedule with its `epsilon` objective
and the default generation settings) and the tokenizer as `tokenizer.json`. Weights are stored in **float32**, exactly as released.
This checkpoint generates **512x512** images (a 64x64 latent).

For model details, intended use and limitations, see the upstream
[model card](https://huggingface.co/sd2-community/stable-diffusion-2-base).

## Architecture

| Component | zeromodels class | Details |
| --- | --- | --- |
| Denoiser | `UNet2DConditionModel` | (320, 640, 1280, 1280) channels, 2 ResNet blocks per level, (5, 10, 20, 20) attention heads on the 1024-d text context, linear token projection, 64x64x4 latent |
| Autoencoder | `AutoencoderKL` | (128, 256, 512, 512) channels, x8 spatial compression to 4 latent channels, `scaling_factor` 0.18215 |
| Text encoder | `CLIPTextModel` | OpenCLIP ViT-H/14 text encoder (penultimate layer): 1024-d, 23 layers, 16 heads, 77 tokens, `gelu` |
| Scheduler | `DDIMScheduler` | scaled_linear betas 0.00085 to 0.012 over 1000 steps, `epsilon`; DDIM / PNDM / Euler / Euler-ancestral are drop-in |

## Quick start

```python
import os
os.environ["KERAS_BACKEND"] = "torch"  # or "jax" / "tensorflow"

from PIL import Image
from zeromodels.models.stable_diffusion_2 import StableDiffusion2TextToImage, StableDiffusion2Tokenizer

model = StableDiffusion2TextToImage.from_weights("zeromodels/stable-diffusion-2-base")
tokenizer = StableDiffusion2Tokenizer.from_weights("zeromodels/stable-diffusion-2-base")

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

`generate` takes the tokenizer's `input_ids` (batch them for several prompts), an optional
`negative_input_ids` (tokenize the negative prompt), `num_inference_steps`, `guidance_scale`,
a `seed`, or explicit `latents` of shape `(batch, 64, 64, 4)` for results that are
identical across backends.

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

| Variant | Hub | Training |
| --- | --- | --- |
| `stable-diffusion-2-base` | [zeromodels/stable-diffusion-2-base](https://huggingface.co/zeromodels/stable-diffusion-2-base) | 512px, epsilon: from scratch, 550k steps at 256px on LAION-5B (aesthetics >= 4.5), then 850k steps at 512px |
| `stable-diffusion-2` | [zeromodels/stable-diffusion-2](https://huggingface.co/zeromodels/stable-diffusion-2) | 768px, v-prediction: 2-base + 150k steps at 768px |
| `stable-diffusion-2-1-base` | [zeromodels/stable-diffusion-2-1-base](https://huggingface.co/zeromodels/stable-diffusion-2-1-base) | 512px, epsilon: 2-base + 220k steps at 512px (punsafe 0.98) |
| `stable-diffusion-2-1` | [zeromodels/stable-diffusion-2-1](https://huggingface.co/zeromodels/stable-diffusion-2-1) | 768px, v-prediction: 2 + 55k steps (punsafe 0.1) + 155k steps (punsafe 0.98) at 768px |
| `sd-turbo` | [zeromodels/sd-turbo](https://huggingface.co/zeromodels/sd-turbo) | 512px, epsilon, Euler (trailing spacing), 1 to 4 steps, no guidance: SD 2.1 distilled with Adversarial Diffusion Distillation (Stability AI Community License) |

## Tips

- Set `KERAS_BACKEND` **before** importing Keras / zeromodels.
- The graphs are built for 512px. Pass `unet_sample_size=<px / 8>, vae_sample_size=<px>` to
  `from_weights` to build for another multiple of 64px (the weights are resolution-independent).
- Swap the sampler any time: `model.scheduler = EulerDiscreteScheduler.from_config(model.config.scheduler_config)`
  (`zeromodels.base.base_scheduler`).
- `StableDiffusion2Model.from_weights(...)` loads the same repo as the bare container
  (UNet / VAE / text encoder as `.unet` / `.vae` / `.text_encoder`) without the generation loop.

- 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 2.x docs](https://imvision12.github.io/ZeroModels/stable_diffusion_2/).

## License

The weights are redistributed under the [CreativeML Open RAIL++-M License](https://huggingface.co/sd2-community/stable-diffusion-2/blob/main/LICENSE-MODEL) of the upstream
checkpoint, including its use-based restrictions. By using them you agree to those terms.

## Notice

Modifications by zeromodels (https://github.com/IMvision12/ZeroModels): the checkpoint
released at https://huggingface.co/sd2-community/stable-diffusion-2-base was converted to
the Keras 3 weights layout of zeromodels (`model.weights.h5, zm_config.json,
tokenizer.json`), stored in float32 as released. 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 and the LAION / OpenCLIP teams for training and releasing Stable Diffusion, and to the
Hugging Face diffusers team, whose implementation this port was verified against.