Add stable-diffusion-3.5-medium: zeromodels Keras 3 conversion of stabilityai/stable-diffusion-3.5-medium

README.md

@@ -0,0 +1,124 @@
+---
+pipeline_tag: text-to-image
+license: other
+base_model: stabilityai/stable-diffusion-3.5-medium
+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/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)
+
+# zeromodels/stable-diffusion-3.5-medium
+
+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-medium`](https://huggingface.co/stabilityai/stable-diffusion-3.5-medium) 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 (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`](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-medium).
+
+## 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
+
+```python
+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](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.
+- 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](https://imvision12.github.io/ZeroModels/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.
+
+## 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.

model_00000.weights.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b975072fd41d76c5c294926ce36678cdab6d8f6842dd1faecf73cd9c1d4c716f
+size 5360876896

model_00001.weights.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:227820b84981e8b140355154e955d4e61bc2d27aec4b05cac6b8146f5b320d2e
+size 1554114672

zm_config.json

@@ -0,0 +1,102 @@
+{
+  "library_name": "zeromodels",
+  "zeromodels_version": "1.3.3",
+  "model_module": "zeromodels.models.stable_diffusion_3_5",
+  "model_class": "StableDiffusion3_5Model",
+  "variant": "stable-diffusion-3.5-medium",
+  "weights": "model.weights.json",
+  "schema_version": 2,
+  "weight_dtype": "float16",
+  "weight_dtype_overrides": {
+    "vae_encoder": "float32",
+    "vae_decoder": "float32"
+  },
+  "model_type": "stable_diffusion_3_5",
+  "transformer_config": {
+    "sample_size": 128,
+    "patch_size": 2,
+    "in_channels": 16,
+    "out_channels": 16,
+    "num_layers": 24,
+    "attention_head_dim": 64,
+    "num_attention_heads": 24,
+    "joint_attention_dim": 4096,
+    "caption_projection_dim": 1536,
+    "pooled_projection_dim": 2048,
+    "pos_embed_max_size": 384,
+    "qk_norm": "rms_norm",
+    "dual_attention_layers": [
+      0,
+      1,
+      2,
+      3,
+      4,
+      5,
+      6,
+      7,
+      8,
+      9,
+      10,
+      11,
+      12
+    ],
+    "text_seq_len": 333
+  },
+  "vae_config": {
+    "in_channels": 3,
+    "out_channels": 3,
+    "latent_channels": 16,
+    "block_out_channels": [
+      128,
+      256,
+      512,
+      512
+    ],
+    "layers_per_block": 2,
+    "norm_num_groups": 32,
+    "sample_size": 1024,
+    "scaling_factor": 1.5305,
+    "force_upcast": true,
+    "shift_factor": 0.0609,
+    "use_quant_conv": false,
+    "use_post_quant_conv": false
+  },
+  "text_config": {
+    "hidden_dim": 768,
+    "num_heads": 12,
+    "num_layers": 12,
+    "mlp_ratio": 4.0,
+    "vocab_size": 49408,
+    "max_seq_len": 77,
+    "projection_dim": 768,
+    "hidden_act": "quick_gelu"
+  },
+  "text_config_2": {
+    "hidden_dim": 1280,
+    "num_heads": 20,
+    "num_layers": 32,
+    "mlp_ratio": 4.0,
+    "vocab_size": 49408,
+    "max_seq_len": 77,
+    "projection_dim": 1280,
+    "hidden_act": "gelu"
+  },
+  "hidden_act": "quick_gelu",
+  "layer_norm_eps": 1e-05,
+  "bos_token_id": 49406,
+  "eos_token_id": 49407,
+  "pad_token_id": 49407,
+  "pad_token_id_2": 0,
+  "max_sequence_length": 256,
+  "eos_token_id_3": 1,
+  "pad_token_id_3": 0,
+  "scheduler_config": {
+    "_class_name": "FlowMatchEulerDiscreteScheduler",
+    "num_train_timesteps": 1000,
+    "shift": 3.0
+  },
+  "generate_args": {
+    "num_inference_steps": 40,
+    "guidance_scale": 4.5
+  }
+}