Add stable-diffusion-2: zeromodels Keras 3 conversion of sd2-community/stable-diffusion-2

README.md

@@ -0,0 +1,109 @@
+---
+pipeline_tag: text-to-image
+license: openrail++
+base_model: sd2-community/stable-diffusion-2
+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
+
+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`](https://huggingface.co/sd2-community/stable-diffusion-2) 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 `v_prediction` objective
+and the default generation settings) and the tokenizer as `tokenizer.json`. Weights are stored in **float32**, exactly as released.
+This checkpoint generates **768x768** images (a 96x96 latent).
+
+For model details, intended use and limitations, see the upstream
+[model card](https://huggingface.co/sd2-community/stable-diffusion-2).
+
+## 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, 96x96x4 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, `v_prediction`; 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")
+tokenizer = StableDiffusion2Tokenizer.from_weights("zeromodels/stable-diffusion-2")
+
+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")  # (768, 768, 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, 96, 96, 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 (non-commercial research license) |
+
+## Tips
+
+- Set `KERAS_BACKEND` **before** importing Keras / zeromodels.
+- The graphs are built for 768px. 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](https://huggingface.co/sd2-community/stable-diffusion-2-1/blob/main/LICENSE-MODEL) of the upstream checkpoint,
+including its use-based restrictions. By using them you agree to those terms.
+
+## 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.

model.weights.h5

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

zm_config.json

@@ -0,0 +1,99 @@
+{
+  "library_name": "zeromodels",
+  "zeromodels_version": "1.3.3",
+  "model_module": "zeromodels.models.stable_diffusion_2",
+  "model_class": "StableDiffusion2Model",
+  "variant": "stable-diffusion-2",
+  "weights": "model.weights.h5",
+  "schema_version": 2,
+  "weight_dtype": "float32",
+  "model_type": "stable_diffusion_2",
+  "unet_config": {
+    "sample_size": 96,
+    "in_channels": 4,
+    "out_channels": 4,
+    "down_block_types": [
+      "CrossAttnDownBlock2D",
+      "CrossAttnDownBlock2D",
+      "CrossAttnDownBlock2D",
+      "DownBlock2D"
+    ],
+    "up_block_types": [
+      "UpBlock2D",
+      "CrossAttnUpBlock2D",
+      "CrossAttnUpBlock2D",
+      "CrossAttnUpBlock2D"
+    ],
+    "block_out_channels": [
+      320,
+      640,
+      1280,
+      1280
+    ],
+    "layers_per_block": 2,
+    "cross_attention_dim": 1024,
+    "num_attention_heads": [
+      5,
+      10,
+      20,
+      20
+    ],
+    "norm_num_groups": 32,
+    "use_linear_projection": true,
+    "transformer_layers_per_block": 1,
+    "addition_embed_type": null,
+    "addition_time_embed_dim": 256,
+    "projection_class_embeddings_input_dim": null,
+    "num_time_ids": 6,
+    "text_seq_len": 77
+  },
+  "vae_config": {
+    "in_channels": 3,
+    "out_channels": 3,
+    "latent_channels": 4,
+    "block_out_channels": [
+      128,
+      256,
+      512,
+      512
+    ],
+    "layers_per_block": 2,
+    "norm_num_groups": 32,
+    "sample_size": 768,
+    "scaling_factor": 0.18215,
+    "force_upcast": false,
+    "shift_factor": 0.0,
+    "use_quant_conv": true,
+    "use_post_quant_conv": true
+  },
+  "text_config": {
+    "hidden_dim": 1024,
+    "num_heads": 16,
+    "num_layers": 23,
+    "mlp_ratio": 4.0,
+    "vocab_size": 49408,
+    "max_seq_len": 77
+  },
+  "hidden_act": "gelu",
+  "layer_norm_eps": 1e-05,
+  "bos_token_id": 49406,
+  "eos_token_id": 49407,
+  "pad_token_id": 0,
+  "scheduler_config": {
+    "_class_name": "DDIMScheduler",
+    "beta_end": 0.012,
+    "beta_schedule": "scaled_linear",
+    "beta_start": 0.00085,
+    "clip_sample": false,
+    "num_train_timesteps": 1000,
+    "prediction_type": "v_prediction",
+    "set_alpha_to_one": false,
+    "skip_prk_steps": true,
+    "steps_offset": 1,
+    "trained_betas": null
+  },
+  "generate_args": {
+    "num_inference_steps": 50,
+    "guidance_scale": 7.5
+  }
+}