--- pipeline_tag: text-to-image license: other license_name: stabilityai-ai-community license_link: https://huggingface.co/stabilityai/sd-turbo/blob/main/LICENSE.md base_model: stabilityai/sd-turbo 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) **Powered by Stability AI** # zeromodels/sd-turbo 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 [`stabilityai/sd-turbo`](https://huggingface.co/stabilityai/sd-turbo) 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 `EulerDiscreteScheduler` 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/stabilityai/sd-turbo). ## 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, float32 (`force_upcast`) | | Text encoder | `CLIPTextModel` | OpenCLIP ViT-H/14 text encoder (penultimate layer): 1024-d, 23 layers, 16 heads, 77 tokens, `gelu` | | Scheduler | `EulerDiscreteScheduler` | scaled_linear betas 0.00085 to 0.012 over 1000 steps, `epsilon`, `trailing` timestep spacing; 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/sd-turbo") tokenizer = StableDiffusion2Tokenizer.from_weights("zeromodels/sd-turbo") inputs = tokenizer("a photograph of an astronaut riding a horse") images = model.generate(**inputs, num_inference_steps=1, guidance_scale=0.0, 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 | 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=, vae_sample_size=` 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 [Stability AI Community License](https://huggingface.co/stabilityai/sd-turbo/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/sd-turbo 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. **Powered by Stability AI** ## 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.