Load with ZMDiffusion / ZMDiffusionTokenizer (zeromodels.diffusers)

README.md

@@ -22,7 +22,7 @@ tags:
 
 # 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)
+[![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/diffusers/#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
 
@@ -56,10 +56,10 @@ import os
 os.environ["KERAS_BACKEND"] = "torch"  # or "jax" / "tensorflow"
 
 from PIL import Image
-from zeromodels.models.stable_diffusion_2 import StableDiffusion2TextToImage, StableDiffusion2Tokenizer
+from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer
 
-model = StableDiffusion2TextToImage.from_weights("zeromodels/stable-diffusion-2-base")
-tokenizer = StableDiffusion2Tokenizer.from_weights("zeromodels/stable-diffusion-2-base")
+model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-2-base")
+tokenizer = ZMDiffusionTokenizer.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)
@@ -84,18 +84,20 @@ Load any Stable Diffusion 2.x checkpoint the same way with `from_weights("zeromo
 ## 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 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.
+- The trained components are attributes of the loaded `ZMDiffusion`
+  (UNet / VAE / text encoder as `.unet` / `.vae` / `.text_encoder`), each callable on its own.
 
 - 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/).
+- See the [Stable Diffusion 2.x docs](https://imvision12.github.io/ZeroModels/diffusers/#stable-diffusion-2).
 
 ## License
 

zm_config.json

@@ -1,8 +1,8 @@
 {
   "library_name": "zeromodels",
   "zeromodels_version": "1.3.3",
-  "model_module": "zeromodels.models.stable_diffusion_2",
-  "model_class": "StableDiffusion2Model",
+  "model_module": "zeromodels.diffusers",
+  "model_class": "ZMDiffusion",
   "variant": "stable-diffusion-2-base",
   "weights": "model.weights.h5",
   "schema_version": 2,
@@ -95,4 +95,4 @@
     "num_inference_steps": 50,
     "guidance_scale": 7.5
   }
-}
+}

zm_preprocessor.json

@@ -0,0 +1,10 @@
+{
+  "library_name": "zeromodels",
+  "zeromodels_version": "1.3.3",
+  "preprocessor_module": "zeromodels.diffusers",
+  "preprocessor_class": "ZMDiffusionTokenizer",
+  "variant": "stable-diffusion-2-base",
+  "model_type": "stable_diffusion_2",
+  "max_seq_len": 77,
+  "pad_token": "!"
+}