From c8d6d4ec5b2df091453c76bd3841eeef276caaff Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Sun, 18 Oct 2026 16:04:45 -0700 Subject: [PATCH 1/4] Refactor Diffusion and Sentence Transformer --- README.md | 2 + docs/auto.md | 23 +- docs/diffusers.md | 244 + docs/getting_started.md | 2 +- docs/index.md | 2 +- docs/main_classes.md | 11 +- docs/models.md | 21 +- docs/qwen_image.md | 56 +- docs/qwen_image_21.md | 45 +- docs/sentence_transformer.md | 26 +- docs/stable_diffusion.md | 84 +- docs/stable_diffusion_2.md | 65 +- docs/stable_diffusion_3.md | 42 +- docs/stable_diffusion_3_5.md | 31 +- docs/stable_diffusion_xl.md | 90 +- pyproject.toml | 1 + tests/base/model_test_registry.py | 161 +- tests/fixtures/cross_backend_parity.json | 19035 +++++++--------- tests/fixtures/dummy_inputs.py | 6 +- tests/integration/test_auto_registry.py | 57 +- tests/integration/test_data_formats.py | 6 +- .../integration/test_sentence_transformer.py | 58 +- website/mkdocs.yml | 16 +- zeromodels/__init__.py | 2 + zeromodels/auto/__init__.py | 4 +- zeromodels/auto/auto_factory.py | 3 +- zeromodels/auto/auto_mapping_names.py | 67 +- zeromodels/base/base_diffusion.py | 8 +- zeromodels/diffusers/__init__.py | 7 + zeromodels/diffusers/diffusion_model.py | 325 + zeromodels/diffusers/diffusion_tokenizer.py | 247 + .../qwen_image/__init__.py | 8 +- .../convert_qwen_image_diffusers_to_keras.py | 10 +- .../qwen_image/qwen_image_config.py | 4 +- .../qwen_image/qwen_image_layers.py | 4 +- .../qwen_image/qwen_image_model.py | 380 +- .../qwen_image/qwen_image_vae.py | 2 +- .../qwen_image_21/__init__.py | 8 +- ...onvert_qwen_image_21_diffusers_to_keras.py | 10 +- .../qwen_image_21/qwen_image_21_config.py | 4 +- .../qwen_image_21/qwen_image_21_layers.py | 4 +- .../qwen_image_21/qwen_image_21_model.py | 315 +- .../qwen_image_21/qwen_image_21_vae.py | 4 +- .../stable_diffusion/__init__.py | 12 +- ...ert_stable_diffusion_diffusers_to_keras.py | 13 +- .../stable_diffusion_config.py | 13 +- .../stable_diffusion_layers.py | 0 .../stable_diffusion_model.py | 193 +- .../diffusers/stable_diffusion_2/__init__.py | 13 + ...t_stable_diffusion_2_diffusers_to_keras.py | 16 +- .../stable_diffusion_2_config.py | 11 +- .../stable_diffusion_2_model.py | 22 + .../stable_diffusion_3/__init__.py | 8 +- ...t_stable_diffusion_3_diffusers_to_keras.py | 21 +- .../stable_diffusion_3_config.py | 17 +- .../stable_diffusion_3_layers.py | 2 +- .../stable_diffusion_3_model.py | 202 +- .../stable_diffusion_3_5/__init__.py | 11 + ...stable_diffusion_3_5_diffusers_to_keras.py | 17 +- .../stable_diffusion_3_5_config.py | 9 +- .../stable_diffusion_3_5_model.py | 22 + .../stable_diffusion_xl/__init__.py | 14 +- ..._stable_diffusion_xl_diffusers_to_keras.py | 23 +- .../stable_diffusion_xl_config.py | 20 +- .../stable_diffusion_xl_model.py | 365 + zeromodels/models/__init__.py | 7 - .../models/qwen_image/qwen_image_tokenizer.py | 58 - .../qwen_image_21/qwen_image_21_tokenizer.py | 57 - .../stable_diffusion_tokenizer.py | 70 - .../models/stable_diffusion_2/__init__.py | 19 - .../stable_diffusion_2_model.py | 69 - .../stable_diffusion_2_tokenizer.py | 34 - .../stable_diffusion_3_tokenizer.py | 89 - .../models/stable_diffusion_3_5/__init__.py | 17 - .../stable_diffusion_3_5_model.py | 69 - .../stable_diffusion_3_5_tokenizer.py | 39 - .../stable_diffusion_xl_model.py | 496 - .../stable_diffusion_xl_tokenizer.py | 42 - zeromodels/sentence_transformers/__init__.py | 32 +- ...onvert_sentence_transformer_hf_to_keras.py | 84 +- .../sentence_transformer_backbones.py | 115 + .../sentence_transformer_config.py | 4 + .../sentence_transformer_model.py | 168 +- .../sentence_transformer_tokenizer.py | 8 +- zeromodels/st_models/__init__.py | 40 - 85 files changed, 10493 insertions(+), 13548 deletions(-) create mode 100644 docs/diffusers.md create mode 100644 zeromodels/diffusers/__init__.py create mode 100644 zeromodels/diffusers/diffusion_model.py create mode 100644 zeromodels/diffusers/diffusion_tokenizer.py rename zeromodels/{models => diffusers}/qwen_image/__init__.py (73%) rename zeromodels/{models => diffusers}/qwen_image/convert_qwen_image_diffusers_to_keras.py (97%) rename zeromodels/{models => diffusers}/qwen_image/qwen_image_config.py (96%) rename zeromodels/{models => diffusers}/qwen_image/qwen_image_layers.py (99%) rename zeromodels/{models => diffusers}/qwen_image/qwen_image_model.py (69%) rename zeromodels/{models => diffusers}/qwen_image/qwen_image_vae.py (99%) rename zeromodels/{models => diffusers}/qwen_image_21/__init__.py (73%) rename zeromodels/{models => diffusers}/qwen_image_21/convert_qwen_image_21_diffusers_to_keras.py (97%) rename zeromodels/{models => diffusers}/qwen_image_21/qwen_image_21_config.py (96%) rename zeromodels/{models => diffusers}/qwen_image_21/qwen_image_21_layers.py (99%) rename zeromodels/{models => diffusers}/qwen_image_21/qwen_image_21_model.py (69%) rename zeromodels/{models => diffusers}/qwen_image_21/qwen_image_21_vae.py (99%) rename zeromodels/{models => diffusers}/stable_diffusion/__init__.py (65%) rename zeromodels/{models => diffusers}/stable_diffusion/convert_stable_diffusion_diffusers_to_keras.py (95%) rename zeromodels/{models => diffusers}/stable_diffusion/stable_diffusion_config.py (94%) rename zeromodels/{models => diffusers}/stable_diffusion/stable_diffusion_layers.py (100%) rename zeromodels/{models => diffusers}/stable_diffusion/stable_diffusion_model.py (80%) create mode 100644 zeromodels/diffusers/stable_diffusion_2/__init__.py rename zeromodels/{models => diffusers}/stable_diffusion_2/convert_stable_diffusion_2_diffusers_to_keras.py (82%) rename zeromodels/{models => diffusers}/stable_diffusion_2/stable_diffusion_2_config.py (88%) create mode 100644 zeromodels/diffusers/stable_diffusion_2/stable_diffusion_2_model.py rename zeromodels/{models => diffusers}/stable_diffusion_3/__init__.py (70%) rename zeromodels/{models => diffusers}/stable_diffusion_3/convert_stable_diffusion_3_diffusers_to_keras.py (94%) rename zeromodels/{models => diffusers}/stable_diffusion_3/stable_diffusion_3_config.py (92%) rename zeromodels/{models => diffusers}/stable_diffusion_3/stable_diffusion_3_layers.py (99%) rename zeromodels/{models => diffusers}/stable_diffusion_3/stable_diffusion_3_model.py (76%) create mode 100644 zeromodels/diffusers/stable_diffusion_3_5/__init__.py rename zeromodels/{models => diffusers}/stable_diffusion_3_5/convert_stable_diffusion_3_5_diffusers_to_keras.py (76%) rename zeromodels/{models => diffusers}/stable_diffusion_3_5/stable_diffusion_3_5_config.py (81%) create mode 100644 zeromodels/diffusers/stable_diffusion_3_5/stable_diffusion_3_5_model.py rename zeromodels/{models => diffusers}/stable_diffusion_xl/__init__.py (57%) rename zeromodels/{models => diffusers}/stable_diffusion_xl/convert_stable_diffusion_xl_diffusers_to_keras.py (94%) rename zeromodels/{models => diffusers}/stable_diffusion_xl/stable_diffusion_xl_config.py (92%) create mode 100644 zeromodels/diffusers/stable_diffusion_xl/stable_diffusion_xl_model.py delete mode 100644 zeromodels/models/qwen_image/qwen_image_tokenizer.py delete mode 100644 zeromodels/models/qwen_image_21/qwen_image_21_tokenizer.py delete mode 100644 zeromodels/models/stable_diffusion/stable_diffusion_tokenizer.py delete mode 100644 zeromodels/models/stable_diffusion_2/__init__.py delete mode 100644 zeromodels/models/stable_diffusion_2/stable_diffusion_2_model.py delete mode 100644 zeromodels/models/stable_diffusion_2/stable_diffusion_2_tokenizer.py delete mode 100644 zeromodels/models/stable_diffusion_3/stable_diffusion_3_tokenizer.py delete mode 100644 zeromodels/models/stable_diffusion_3_5/__init__.py delete mode 100644 zeromodels/models/stable_diffusion_3_5/stable_diffusion_3_5_model.py delete mode 100644 zeromodels/models/stable_diffusion_3_5/stable_diffusion_3_5_tokenizer.py delete mode 100644 zeromodels/models/stable_diffusion_xl/stable_diffusion_xl_model.py delete mode 100644 zeromodels/models/stable_diffusion_xl/stable_diffusion_xl_tokenizer.py delete mode 100644 zeromodels/st_models/__init__.py diff --git a/README.md b/README.md index ae038a46..add0da12 100644 --- a/README.md +++ b/README.md @@ -271,6 +271,8 @@ All 13 checkpoints below have pre-converted Keras weights under the `zeromodels` - Text-to-Image Diffusion (text → image) + Every diffusion checkpoint loads with `ZMDiffusion` and `ZMDiffusionTokenizer` from `zeromodels.diffusers` ([documentation](https://imvision12.github.io/ZeroModels/diffusers/)). + | 🏷️ Model Name | 📜 Reference Paper | 📦 Source of Weights | |---------------|-------------------|---------------------| | Stable Diffusion 1.x (v1-1 to v1-5) | [High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) | `diffusers` | diff --git a/docs/auto.md b/docs/auto.md index 94b7304f..c066c74e 100644 --- a/docs/auto.md +++ b/docs/auto.md @@ -67,6 +67,23 @@ The tokenizer preserves the checkpoint's own preprocessing recipe. See upstream `hf:sentence-transformers/...` configs describe their backbone type; use the concrete classes to load their sentence-embedding recipe. +Diffusion repos resolve the same way: every diffusion `model_type` +(`stable_diffusion`, `stable_diffusion_xl`, `stable_diffusion_3_5`, `qwen_image`, ...) +maps to the one pipeline class and the one tokenizer of +[`zeromodels.diffusers`](diffusers.md): + +```python +from zeromodels import AutoZModel, AutoZMTokenizer + +repo = "zeromodels/stable-diffusion-v1-5" +pipe = AutoZModel.from_weights(repo) # -> ZMDiffusion (also AutoZMTextToImage) +tokenizer = AutoZMTokenizer.from_weights(repo) # -> ZMDiffusionTokenizer +images = pipe.generate(**tokenizer("a photo of a cat")) +``` + +The SDXL refiner (`stable_diffusion_xl_refiner`) is also `ZMDiffusion`, listed under +`AutoZMImageToImage`. + ## Task classes ### Pipeline components @@ -75,9 +92,9 @@ Exported component classes are not necessarily standalone Auto tasks. For Qwen-I and Qwen-Image 2.1, the VAE (`AutoencoderKLQwenImage`, `AutoencoderKLQwenImage21`), text encoder (`QwenImageTextEncoderModel`, `QwenImage21TextEncoderModel`), and transformer (`QwenImageTransformer2DModel`, `QwenImage21Transformer2DModel`) are -pipeline components. Load the complete model with `AutoZModel` or the generation -pipeline with `AutoZMTextToImage`, then access its `.vae`, `.text_encoder`, or -`.transformer` attribute. Direct component construction and serialization remain +pipeline components. Load the complete pipeline, a `ZMDiffusion`, with `AutoZModel` +or `AutoZMTextToImage`, then access its `.vae`, `.text_encoder`, or `.transformer` +attribute. Direct component construction and serialization remain available; component config types do not imply standalone Auto weight loading. These six classes are explicitly recorded with their pipeline owners in diff --git a/docs/diffusers.md b/docs/diffusers.md new file mode 100644 index 00000000..6876aa32 --- /dev/null +++ b/docs/diffusers.md @@ -0,0 +1,244 @@ +# Diffusers + +
+Weights: every diffusion checkpoint is hosted under the +zeromodels organization as one +zm_config.json + model.weights.h5 (or sharded +model.weights.json) + tokenizer.json. Load any of them with +ZMDiffusion.from_weights("zeromodels/<variant>"). On-the-fly +hf: conversion is not supported for diffusion models. +
+ +`zeromodels.diffusers` holds every text-to-image model of ZeroModels behind **one +pipeline class and one tokenizer**, the way `zeromodels.sentence_transformers` holds the +sentence-embedding models: + +- `ZMDiffusion`: the model and its `generate` (diffusers' `DiffusionPipeline`). +- `ZMDiffusionTokenizer`: the prompt tokenizer. + +These two are the whole public API (`zeromodels.diffusers.__all__`); everything else in the +package (the family recipes, configs and components) is internal. + +| | | +|---|---| +| Class | `ZMDiffusion` | +| Tokenizer | `ZMDiffusionTokenizer` | +| Selector | `model_type` (the architecture differs per family) | +| Recipes | Family classes in `diffusers//`, registered in `DIFFUSION_FAMILIES` | +| Hub files | `zm_config.json` + `zm_preprocessor.json` + `tokenizer.json` | + +A checkpoint's `model_type` (recorded in its `zm_config.json`) selects the family; the +family decides which components are built and how a prompt becomes an image, so loading +and generating look the same for Stable Diffusion 1.x and Qwen-Image alike. The public API +is `zeromodels.diffusers`, separate from the upstream `diffusers` library (which is only +needed to convert checkpoints). + +## Usage + +```python +import os + +os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" + +from PIL import Image +from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer + +repo = "zeromodels/stable-diffusion-xl-base-1.0" +pipe = ZMDiffusion.from_weights(repo) +tokenizer = ZMDiffusionTokenizer.from_weights(repo) + +images = pipe.generate(**tokenizer("a lighthouse on a rocky coast at dawn"), seed=0) +Image.fromarray(images[0]).save("lighthouse.png") # (1024, 1024, 3) uint8 +``` + +Use the same repo for the model and the tokenizer. Swapping the repo id is all it takes to +switch family; the defaults (`num_inference_steps`, `guidance_scale`, the sampler) come +from each repo. + +## Families + +| `model_type` | Hosted checkpoints | Page | +|---|---|---| +| `stable_diffusion` | `stable-diffusion-v1-1` to `stable-diffusion-v1-5` | [Stable Diffusion](stable_diffusion.md) | +| `stable_diffusion_2` | `stable-diffusion-2`, `stable-diffusion-2-base`, `stable-diffusion-2-1`, `stable-diffusion-2-1-base`, `sd-turbo` | [Stable Diffusion 2](stable_diffusion_2.md) | +| `stable_diffusion_xl` | `stable-diffusion-xl-base-1.0`, `sdxl-turbo` | [Stable Diffusion XL](stable_diffusion_xl.md) | +| `stable_diffusion_xl_refiner` | `stable-diffusion-xl-refiner-1.0` (image-to-image) | [Stable Diffusion XL](stable_diffusion_xl.md#the-refiner) | +| `stable_diffusion_3` | `stable-diffusion-3-medium` | [Stable Diffusion 3](stable_diffusion_3.md) | +| `stable_diffusion_3_5` | `stable-diffusion-3.5-large`, `stable-diffusion-3.5-large-turbo`, `stable-diffusion-3.5-medium` | [Stable Diffusion 3.5](stable_diffusion_3_5.md) | +| `qwen_image` | `qwen-image` | [Qwen-Image](qwen_image.md) | +| `qwen_image_21` | `qwen-image-2.1` | [Qwen-Image-2.1](qwen_image_21.md) | + +The family pages cover each architecture, its generation arguments, the verification +against diffusers and the memory needed. + +## `ZMDiffusion` + +```python +ZMDiffusion(model_type, scheduler=None, **config_fields) +ZMDiffusion(config) # or a family config positionally (zeromodels.diffusers.) +ZMDiffusion.from_weights("zeromodels/", **kwargs) +``` + +One functional `keras.Model` holds every trained component as disconnected paths: the +denoiser (`{"sample", "timestep", "encoder_hidden_states", ...} -> "noise_pred"`), the VAE +(`{"image", "latent"} -> "moments" / "image"`) and the text towers (token ids to prompt +embeddings). The components are attributes that `generate` runs step by step and that can +be called on their own: + +| Family | Denoiser | VAE | Text towers | +|---|---|---|---| +| SD 1.x / 2.x | `.unet` | `.vae` | `.text_encoder` | +| SDXL | `.unet` | `.vae` | `.text_encoder`, `.text_encoder_2` | +| SDXL refiner | `.unet` | `.vae` | `.text_encoder_2` | +| SD 3 / 3.5 | `.transformer` | `.vae` | `.text_encoder`, `.text_encoder_2` (+ optional `.text_encoder_3`) | +| Qwen-Image / 2.1 | `.transformer` | `.vae` | `.text_encoder` | + +`model.config` is the typed family config (`StableDiffusionXLConfig` for an SDXL +checkpoint), `model.model_type` the family, `model.scheduler` the sampler (built from the +repo's `scheduler_config`; assign another at any time) and `model.generate_args` the repo's +default generation settings. + +`from_weights` accepts everything [`BaseModel.from_weights`](loading_weights.md) does +(`load_dtype`, `quantization`, `attn_implementation`, `cache_converted`, size overrides such +as `unet_sample_size=96, vae_sample_size=768`) plus `text_encoder_3` for SD 3 / 3.5: a +hosted `SD3T5EncoderModel` repo id (`"zeromodels/t5-v1_1-xxl-encoder"`) or a built +encoder, attached outside the checkpoint's weights. + +### `generate` + +```python +generate( + input_ids, + attention_mask=None, + negative_input_ids=None, + num_inference_steps=None, + guidance_scale=None, + seed=None, + latents=None, + image=None, + strength=None, + denoising_start=None, + denoising_end=None, + output_type="image", + **family_arguments, +) +``` + +The shared loop is [`BaseDiffusion.generate`](main_classes.md#basediffusion): tokenized +prompts in, `(batch, H, W, 3)` uint8 images (or the final latent) out, with classifier-free +guidance, image-to-image (`image` + `strength`) and the base / refiner split +(`denoising_end` / `denoising_start`). Family-specific arguments: + +| Family | Arguments | +|---|---| +| SDXL / refiner | `original_size`, `crops_coords_top_left`, `target_size` (`(height, width)` pairs), `aesthetic_score` / `negative_aesthetic_score` on the refiner, `negative_*` variants | +| SD 3 / 3.5 | `input_ids_3` / `negative_input_ids_3` (the T5 tokens the tokenizer returns) | +| Qwen-Image / 2.1 | `height`, `width`, `negative_attention_mask` | + +## `ZMDiffusionTokenizer` + +```python +ZMDiffusionTokenizer( + model_type, + hf_id=None, + tokenizer_file=None, + tokenizer_file_3=None, + max_seq_len=None, + max_sequence_length=None, + pad_token=None, +) +ZMDiffusionTokenizer.from_weights("zeromodels/") +``` + +Calling it on a string or a list of strings returns what `generate` takes: + +| `model_type` | Encoding | Returns | +|---|---|---| +| `stable_diffusion`, `stable_diffusion_xl`, `stable_diffusion_xl_refiner` | CLIP BPE, `<|endoftext|>`-padded to 77 | `input_ids`, `attention_mask` | +| `stable_diffusion_2` | CLIP BPE, `!`-padded to 77 | `input_ids`, `attention_mask` | +| `stable_diffusion_3`, `stable_diffusion_3_5` | CLIP BPE (77) + T5 SentencePiece (`tokenizer_3.json`, 256) | `input_ids`, `attention_mask`, `input_ids_3` | +| `qwen_image`, `qwen_image_21` | the pipeline's ChatML template, Qwen2 / Qwen3 BPE | `input_ids`, `attention_mask` | + +`decode(ids)` turns ids back into text. `from_weights` reads the family and its settings +from the repo's `zm_preprocessor.json`, or from the `model_type` of its `zm_config.json` +when the repo has no `zm_preprocessor.json`, and the vocabulary from the repo's +`tokenizer.json`. + +## Auto loading + +The same repos work with the Auto API: + +```python +from zeromodels import AutoZModel, AutoZMTokenizer + +repo = "zeromodels/qwen-image" +pipe = AutoZModel.from_weights(repo) # ZMDiffusion +tokenizer = AutoZMTokenizer.from_weights(repo) # ZMDiffusionTokenizer +``` + +`AutoZMTextToImage` (and `AutoZMImageToImage` for the SDXL refiner) resolve to +`ZMDiffusion` as well. See [Auto](auto.md). + +## Hub metadata + +A repo's `zm_config.json` names the class and the family: + +```json +{ + "library_name": "zeromodels", + "model_module": "zeromodels.diffusers", + "model_class": "ZMDiffusion", + "variant": "stable-diffusion-v1-5", + "model_type": "stable_diffusion", + "unet_config": {"...": "..."}, + "vae_config": {"...": "..."}, + "text_config": {"...": "..."}, + "scheduler_config": {"_class_name": "PNDMScheduler", "...": "..."}, + "generate_args": {"num_inference_steps": 50, "guidance_scale": 7.5} +} +``` + +and `zm_preprocessor.json` the tokenizer recipe (the files themselves are read from the +same repo): + +```json +{ + "library_name": "zeromodels", + "preprocessor_module": "zeromodels.diffusers", + "preprocessor_class": "ZMDiffusionTokenizer", + "variant": "stable-diffusion-v1-5", + "model_type": "stable_diffusion", + "max_seq_len": 77, + "pad_token": "<|endoftext|>" +} +``` + +## Migrating from the per-family classes + +The per-family container and pipeline classes were replaced by `ZMDiffusion`, and the +per-family tokenizers by `ZMDiffusionTokenizer`, the only two public names. The +components and configs moved from `zeromodels.models.` into the internal +`zeromodels.diffusers.` packages. +The hosted weights are unchanged, and `zm_config.json` files that still name an old class +load into `ZMDiffusion`. + +| Before | Now | +|---|---| +| `StableDiffusionTextToImage`, `StableDiffusionModel` (and the `2`, `XL`, `3`, `3_5` variants) | `ZMDiffusion` | +| `StableDiffusionXLRefinerImageToImage`, `StableDiffusionXLRefinerModel` | `ZMDiffusion` | +| `QwenImageTextToImage`, `QwenImageModel`, `QwenImage21TextToImage`, `QwenImage21Model` | `ZMDiffusion` | +| `StableDiffusionTokenizer` (and every family tokenizer) | `ZMDiffusionTokenizer` | +| `from zeromodels.models.stable_diffusion_3 import SD3T5EncoderModel` | `ZMDiffusion.from_weights(repo, text_encoder_3="zeromodels/t5-v1_1-xxl-encoder")` | + +## Adding a family + +A family is a plain recipe class (for example `StableDiffusionFamily` in +`zeromodels/diffusers/stable_diffusion/stable_diffusion_model.py`) registered in +`DIFFUSION_FAMILIES` (`diffusion_model.py`) under its config's `model_type`. It builds the +components and the graph (`build_components`, `build_graph`) and implements the +`BaseDiffusion` hooks (`encode_prompt`, `unconditional_ids`, `predict_noise`, +`decode_latents`, `latent_shape`, `default_scheduler`, and `generate` / `denoise` / +`prepare_latents` where it differs from the shared loop), each taking the `ZMDiffusion` +instance; a hook the family does not define falls back to `BaseDiffusion`'s. Its config +is a plain `BaseConfig` whose `model_type` names the family. Add the tokenizer settings to +`DIFFUSION_TOKENIZERS` (`diffusion_tokenizer.py`) and the `model_type` to the Auto tables. diff --git a/docs/getting_started.md b/docs/getting_started.md index 65a7e29e..dfb27e3e 100644 --- a/docs/getting_started.md +++ b/docs/getting_started.md @@ -55,7 +55,7 @@ results = processor.post_process_object_detection( [Qwen3-VL](qwen3_vl.md) ┬╖ [InternVL](internvl.md) ┬╖ [Kimi K2.5](kimi_k25.md) ┬╖ [LocateAnything](locateanything.md) ┬╖ - [Stable Diffusion](stable_diffusion.md) ┬╖ [Qwen-Image](qwen_image.md) ┬╖ + [Diffusers](diffusers.md) ┬╖ [Stable Diffusion](stable_diffusion.md) ┬╖ [Qwen-Image](qwen_image.md) ┬╖ [Qwen-Image-2.1](qwen_image_21.md) - **Speech** diff --git a/docs/index.md b/docs/index.md index f5af4bdc..148cc9a8 100644 --- a/docs/index.md +++ b/docs/index.md @@ -327,7 +327,7 @@ Vision-language generation, grounding, and text-to-image diffusion. [Qwen3-VL](qwen3_vl.md) · [InternVL](internvl.md) · [Kimi K2.5](kimi_k25.md) · [LocateAnything](locateanything.md) · -[Stable Diffusion](stable_diffusion.md) · [Qwen-Image](qwen_image.md) · +[Diffusers](diffusers.md) · [Stable Diffusion](stable_diffusion.md) · [Qwen-Image](qwen_image.md) · [Qwen-Image-2.1](qwen_image_21.md) diff --git a/docs/main_classes.md b/docs/main_classes.md index 64bad0b9..787e971d 100644 --- a/docs/main_classes.md +++ b/docs/main_classes.md @@ -257,9 +257,12 @@ model.generate( ) ``` -The diffusion flavor, used by [Stable Diffusion](stable_diffusion.md), -[Stable Diffusion 2](stable_diffusion_2.md), [Stable Diffusion XL](stable_diffusion_xl.md) -and [Stable Diffusion 3 / 3.5](stable_diffusion_3.md). Where the LM +The diffusion flavor, mixed into [`ZMDiffusion`](diffusers.md), the one pipeline class of +every diffusion family ([Stable Diffusion](stable_diffusion.md), +[Stable Diffusion 2](stable_diffusion_2.md), [Stable Diffusion XL](stable_diffusion_xl.md), +[Stable Diffusion 3 / 3.5](stable_diffusion_3.md), [Qwen-Image](qwen_image.md), +[Qwen-Image-2.1](qwen_image_21.md)); each family supplies the hooks through its +family recipe, selected by the checkpoint's `model_type`. Where the LM mixins decode tokens, this one runs a scheduler's denoising loop with classifier-free guidance over a latent and decodes it to `(batch, H, W, 3)` uint8 images. A model supplies five hooks (`encode_prompt`, `unconditional_ids`, `predict_noise`, @@ -274,7 +277,7 @@ instance, like the LM decode loop. **Parameters** -- **input_ids**: `(batch, seq)` token ids from the family's tokenizer. +- **input_ids**: `(batch, seq)` token ids from `ZMDiffusionTokenizer`. - **negative_input_ids** (*optional*): the tokenized negative prompt, one row per prompt or a single row for the batch; the empty prompt when omitted. - **num_inference_steps** / **guidance_scale** (*optional*): default to the repo's diff --git a/docs/models.md b/docs/models.md index 5c06b6ab..0c0308a3 100644 --- a/docs/models.md +++ b/docs/models.md @@ -46,6 +46,20 @@ Dense text embeddings with checkpoint-specific pooling and normalization. - [All 13 Sentence Transformer checkpoints](sentence_transformers.md) - [Sentence Transformer API and usage](sentence_transformer.md) +## Diffusion models + +Text-to-image (and image-to-image) diffusion, every family loaded through `ZMDiffusion` +and `ZMDiffusionTokenizer`. + +- [Diffusers API and usage](diffusers.md) +- [Stable Diffusion](stable_diffusion.md) +- [Stable Diffusion 2](stable_diffusion_2.md) +- [Stable Diffusion XL](stable_diffusion_xl.md) +- [Stable Diffusion 3](stable_diffusion_3.md) +- [Stable Diffusion 3.5](stable_diffusion_3_5.md) +- [Qwen-Image](qwen_image.md) +- [Qwen-Image-2.1](qwen_image_21.md) + ## Vision models Detection, segmentation, monocular depth, and self-supervised backbones. @@ -112,10 +126,3 @@ Vision-language encoders, generative VLMs, grounding across detection, OCR, poin - [SAM 3](sam3.md) - [SigLIP](siglip.md) - [SigLIP 2](siglip2.md) -- [Stable Diffusion](stable_diffusion.md) -- [Stable Diffusion 2](stable_diffusion_2.md) -- [Stable Diffusion XL](stable_diffusion_xl.md) -- [Stable Diffusion 3](stable_diffusion_3.md) -- [Stable Diffusion 3.5](stable_diffusion_3_5.md) -- [Qwen-Image](qwen_image.md) -- [Qwen-Image-2.1](qwen_image_21.md) diff --git a/docs/qwen_image.md b/docs/qwen_image.md index dc76ed7a..a324bfd2 100644 --- a/docs/qwen_image.md +++ b/docs/qwen_image.md @@ -10,11 +10,11 @@ Qwen-Image, ported to pure Keras 3: latent text-to-image flow-matching with a 60-layer double-stream MMDiT denoiser, a Wan-derived KL autoencoder and the -Qwen2.5-VL-7B Instruct text tower. The whole model is **one container**, -`QwenImageModel`, a single functional graph with the three towers as disconnected +Qwen2.5-VL-7B Instruct text tower. It loads through [`ZMDiffusion`](diffusers.md), the one diffusion +pipeline class of zeromodels, with `model_type="qwen_image"`: a single functional graph with the three towers as disconnected sub-graphs (the way `CLIPModel` holds its two), so a hosted repo is one set of -sharded weights and one `zm_config.json`. `QwenImageTextToImage` is that same -graph plus `generate`: the true-CFG denoising loop of `BaseDiffusion`, the +sharded weights and one `zm_config.json`. plus +`generate`: the true-CFG denoising loop of `BaseDiffusion`, the diffusion counterpart of `BaseGeneration`. The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion @@ -52,7 +52,7 @@ in ZeroModels). ## Variants Preconverted, bfloat16 weights are hosted under `zeromodels/`. Load with -`from_weights("zeromodels/")`. Each repo is one container: MMDiT +`from_weights("zeromodels/")`. Each repo is one checkpoint: MMDiT transformer + VAE + Qwen2.5-VL text tower, ~53 GiB at 16-bit. Released under Apache-2.0. @@ -70,10 +70,11 @@ the checkpoint's `scheduler_config` and the Qwen special token ids. Each repo's `zm_config.json` parses through it; the constructor stays flat, with the sub-config fields prefixed `transformer_` / `vae_` / `text_`. -### `QwenImageTextToImage` +### `ZMDiffusion` -The text-to-image task: the `QwenImageModel` container plus `BaseDiffusion`'s -`generate`. It supplies the hooks the mixin needs (`encode_prompt` on the text +The pipeline class of every diffusion family (see [Diffusers](diffusers.md)); the +repo's `model_type` (`"qwen_image"` here) selects this family, which supplies the hooks +`BaseDiffusion`'s `generate` needs the mixin needs (`encode_prompt` on the text tower with ChatML template drop, `unconditional_ids`, `predict_noise` on the transformer, `decode_latents` on the VAE with mean/std un-normalization) and a `scheduler`, built from the config's `scheduler_config` @@ -114,6 +115,7 @@ Returns `(batch, height, width, 3)` uint8 numpy images. | Constructor arg | Default | Meaning | |---|---|---| +| `model_type` | required | `"qwen_image"` (read from the repo's `zm_config.json` on `from_weights`) | | `scheduler` | `None` | a `BaseScheduler`; built from `scheduler_config` when unset | | `transformer_sample_size` | `128` | latent side the transformer graph is built for (image / 8) | | `vae_sample_size` | `1024` | image size the VAE graphs are built for | @@ -123,15 +125,12 @@ Returns `(batch, height, width, 3)` uint8 numpy images. | `prompt_template_encode_start_idx` | `34` | ChatML template tokens dropped after the text encoder | | `max_sequence_length` | `512` | prompt embed length after the template drop | -### `QwenImageModel` - -The container: one functional model whose graph is three disconnected paths, one per -component. Inputs cover the transformer (`sample`, `timestep`, +The graph is three disconnected paths, one per component. Inputs cover the transformer (`sample`, `timestep`, `encoder_hidden_states`, `encoder_hidden_states_mask`), the VAE (`image`, `latent`) and the text tower (`token_ids`, `padding_mask`); outputs include the packed velocity, VAE moments / sample, and prompt embeds. The components are exposed as -`.transformer`, `.vae` and `.text_encoder` and share the container's weights, so -loading the container loads all three. It loads the same repo as the task class. +`.transformer`, `.vae` and `.text_encoder` and share the model's weights, so one load +fills all three, and each can be called on its own. ### `QwenImageTransformer2DModel` @@ -190,7 +189,7 @@ matching Diffusers. ## Preprocessing -### `QwenImageTokenizer` +### `ZMDiffusionTokenizer` The Qwen2 BPE tokenizer with the Diffusers ChatML prompt template wrapped around each string. Calling it returns `{"input_ids", "attention_mask"}` for a string or @@ -199,7 +198,7 @@ a list of strings. `generate` / `encode_prompt` drop the template prefix `max_sequence_length` (512) tokens. ```python -QwenImageTokenizer(hf_id=None, tokenizer_file=None, max_seq_len=1024) +ZMDiffusionTokenizer(model_type="qwen_image", hf_id=None, tokenizer_file=None, max_seq_len=1024) ``` There is no image processor and no processor class: text-to-image takes token ids @@ -215,13 +214,10 @@ import os os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" from PIL import Image -from zeromodels.models.qwen_image import ( - QwenImageTextToImage, - QwenImageTokenizer, -) +from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer -model = QwenImageTextToImage.from_weights("zeromodels/qwen-image") -tokenizer = QwenImageTokenizer.from_weights("zeromodels/qwen-image") +model = ZMDiffusion.from_weights("zeromodels/qwen-image") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/qwen-image") inputs = tokenizer( "a coffee shop entrance with a chalkboard sign reading 'Open', warm afternoon light" @@ -272,7 +268,7 @@ initial packed latent yourself. At 1024px the spatial noise before packing is ```python import numpy as np -from zeromodels.models.qwen_image.qwen_image_model import pack_latents +from zeromodels.diffusers.qwen_image.qwen_image_model import pack_latents h, w, channels = 128, 128, 16 noise = np.random.default_rng(0).standard_normal((1, h, w, channels)).astype("float32") @@ -297,12 +293,12 @@ End-to-end image parity depends on dtype and resolution; prefer bf16 (or fp32) over fp16 — Qwen-Image activations overflow the fp16 range without the Diffusers clip path. -### Container only +### Components on their own ```python -from zeromodels.models.qwen_image import QwenImageModel +from zeromodels.diffusers import ZMDiffusion -qi = QwenImageModel.from_weights("zeromodels/qwen-image") +qi = ZMDiffusion.from_weights("zeromodels/qwen-image") latent = qi.vae.encode(image) # (1, H, W, 3) in [-1, 1] -> (1, H/8, W/8, 16) text = qi.text_encoder({"input_ids": ids, "attention_mask": mask})["last_hidden_state"] ``` @@ -314,7 +310,7 @@ size with the constructor overrides (multiples of 16px after VAE×pack; the checkpoint targets 1024px): ```python -model = QwenImageTextToImage.from_weights( +model = ZMDiffusion.from_weights( "zeromodels/qwen-image", transformer_sample_size=64, vae_sample_size=512, @@ -340,7 +336,7 @@ tokens and has no spatial layout. `generate` always hands back ## Memory and speed -The bf16 / float16 container is about 53 GiB. Building the full graph on a ~40 GB +The bf16 / float16 checkpoint is about 53 GiB. Building the full graph on a ~40 GB GPU OOMs; build on CPU (or load with component offload) and move one tower at a time. A guided 1024px run needs substantial activation headroom on top of the ~38 GiB transformer; 512px is the practical default on 40 GB cards. Prefer the @@ -350,10 +346,10 @@ plain matmul softmax default. ## Loading Fine-tuned Weights The hosted checkpoint is the supported weight; any repo laid out like it -(`zm_config.json` declaring `QwenImageModel`, sharded `*.weights.json` / +(`zm_config.json` declaring `ZMDiffusion` with `model_type: "qwen_image"`, sharded `*.weights.json` / `*.weights.h5`, `tokenizer.json`) loads with `from_weights("/")`. The `hf:` prefix raises for diffusion models: convert a Diffusers-format checkpoint once with -`zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py` +`zeromodels/diffusers/qwen_image/convert_qwen_image_diffusers_to_keras.py` (`transfer_qwen_image(repo)`, `pip install zeromodels[conversion]`) and host the result. diff --git a/docs/qwen_image_21.md b/docs/qwen_image_21.md index 47657360..86dcb241 100644 --- a/docs/qwen_image_21.md +++ b/docs/qwen_image_21.md @@ -10,8 +10,10 @@ Qwen-Image-2.1, ported to pure Keras 3: latent text-to-image flow-matching with a 32-layer **single-stream** block-causal DiT, a residual 64-channel KL autoencoder -(16× spatial, RGBA), and the Qwen3-VL text tower. The whole model is **one -container**, `QwenImage21Model`. `QwenImage21TextToImage` adds `generate`. +(16× spatial, RGBA), and the Qwen3-VL text tower. It loads through +[`ZMDiffusion`](diffusers.md), the one diffusion pipeline class of zeromodels, with +`model_type="qwen_image_21"`: one functional graph holding the three towers, plus +`generate`. Prompts are tokenized by `ZMDiffusionTokenizer`. The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion is deliberately **not supported** for diffusion models. @@ -42,7 +44,7 @@ Links: ## Variants Preconverted, bfloat16 weights are hosted under `zeromodels/`. Load with -`from_weights("zeromodels/")`. Each repo is one container: DiT + VAE + +`from_weights("zeromodels/")`. Each repo is one checkpoint: DiT + VAE + Qwen3-VL text tower, ~28 GiB at 16-bit. | Variant | Hub | Source | @@ -59,10 +61,11 @@ over `QwenImage21TransformerConfig`, `QwenImage21VAEConfig` and token ids. Each repo's `zm_config.json` parses through it; the constructor stays flat, with the sub-config fields prefixed `transformer_` / `vae_` / `text_`. -### `QwenImage21TextToImage` +### `ZMDiffusion` -The text-to-image task: the `QwenImage21Model` container plus `BaseDiffusion`'s -`generate`. It supplies `encode_prompt` (ChatML template drop after the text +The pipeline class of every diffusion family (see [Diffusers](diffusers.md)); the +repo's `model_type` (`"qwen_image_21"` here) selects this family, which supplies the hooks +`BaseDiffusion`'s `generate` needs: `encode_prompt` (ChatML template drop after the text encoder), `predict_noise` on the transformer, `decode_latents` on the VAE with mean/std un-normalization, and a `FlowMatchEulerDiscreteScheduler` from `scheduler_config`. @@ -100,6 +103,7 @@ Returns `(batch, height, width, 3)` uint8 numpy images (VAE RGBA is cropped to R | Constructor arg | Default | Meaning | |---|---|---| +| `model_type` | required | `"qwen_image_21"` (read from the repo's `zm_config.json` on `from_weights`) | | `scheduler` | `None` | a `BaseScheduler`; built from `scheduler_config` when unset | | `transformer_sample_size` | `64` | latent side the DiT graph is built for (image / 16) | | `vae_sample_size` | `1024` | image size the VAE graphs are built for | @@ -109,9 +113,7 @@ Returns `(batch, height, width, 3)` uint8 numpy images (VAE RGBA is cropped to R | `prompt_template_encode_start_idx` | `14` | ChatML system-prefix tokens dropped after the text encoder | | `max_sequence_length` | `512` | prompt embed length after the template drop | -### `QwenImage21Model` - -The container: one functional model whose graph is three disconnected paths. +The graph is three disconnected paths, one per component. Inputs cover the transformer (`sample`, `timestep`, `encoder_hidden_states`, `encoder_hidden_states_mask`), the VAE (`image`, `latent`) and the text tower (`token_ids`, `padding_mask`). Components are exposed as `.transformer`, `.vae` @@ -171,7 +173,7 @@ linspace sigma schedule, matching Diffusers. ## Preprocessing -### `QwenImage21Tokenizer` +### `ZMDiffusionTokenizer` Qwen3 BPE with Diffusers' T2I ChatML template (`Comprehend and analyze the provided prompt.`). Returns `{"input_ids", "attention_mask"}`. `encode_prompt` @@ -179,7 +181,7 @@ drops `prompt_template_encode_start_idx` (14) tokens and keeps at most `max_sequence_length` (512). ```python -QwenImage21Tokenizer(hf_id=None, tokenizer_file=None, max_seq_len=1024) +ZMDiffusionTokenizer(model_type="qwen_image_21", hf_id=None, tokenizer_file=None, max_seq_len=1024) ``` ## End-to-end example @@ -192,13 +194,10 @@ import os os.environ["KERAS_BACKEND"] = "torch" from PIL import Image -from zeromodels.models.qwen_image_21 import ( - QwenImage21TextToImage, - QwenImage21Tokenizer, -) +from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer -model = QwenImage21TextToImage.from_weights("zeromodels/qwen-image-2.1") -tokenizer = QwenImage21Tokenizer.from_weights("zeromodels/qwen-image-2.1") +model = ZMDiffusion.from_weights("zeromodels/qwen-image-2.1") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/qwen-image-2.1") inputs = tokenizer("a photo of a capybara wearing a wizard hat, soft window light") images = model.generate( @@ -234,7 +233,7 @@ At 1024px the unpatched noise is `(batch, 64, 64, 64)` → `(batch, 4096, 64)`: ```python import numpy as np -from zeromodels.models.qwen_image_21.qwen_image_21_model import pack_latents +from zeromodels.diffusers.qwen_image_21.qwen_image_21_model import pack_latents h, w, channels = 64, 64, 64 noise = np.random.default_rng(0).standard_normal((1, h, w, channels)).astype("float32") @@ -248,7 +247,7 @@ Graphs are built for a fixed size; weights are not. Rebuild with constructor overrides (multiples of 32px; checkpoint targets 1024): ```python -model = QwenImage21TextToImage.from_weights( +model = ZMDiffusion.from_weights( "zeromodels/qwen-image-2.1", transformer_sample_size=32, vae_sample_size=512, @@ -261,10 +260,10 @@ images = model.generate(**tokenizer("a mountain lake at dawn"), height=512, widt ```bash pip install zeromodels[conversion] # KERAS_BACKEND=torch; prefer CPU for the full ~28 GiB bf16 build -python -m zeromodels.models.qwen_image_21.convert_qwen_image_21_diffusers_to_keras +python -m zeromodels.diffusers.qwen_image_21.convert_qwen_image_21_diffusers_to_keras ``` -`transfer_qwen_image_21(repo)` streams Diffusers shards into a Keras container and +`transfer_qwen_image_21(repo)` streams Diffusers shards into a `ZMDiffusion` and writes sharded `*.weights.json`. Building the full text tower on GPU can OOM during Functional shape tracing; convert on CPU, then load for CUDA inference. @@ -282,7 +281,7 @@ during Functional shape tracing; convert on CPU, then load for CUDA inference. ## Memory and speed -The bf16 container is about 28 GiB (DiT ~14 GiB + text ~14 GiB + VAE ~1.4 GiB). +The bf16 checkpoint is about 28 GiB (DiT ~14 GiB + text ~14 GiB + VAE ~1.4 GiB). Building the full Functional graph at 1024px on GPU can OOM during attention shape tracing (layers use `compute_output_spec` to avoid materializing seq² temps). Prefer CPU convert / build, then CUDA generate. A 1024px run needs substantial activation @@ -292,7 +291,7 @@ headroom on an 80 GB card; prefer fused attention ## Loading Fine-tuned Weights The hosted checkpoint is the supported weight; any repo laid out like it -(`zm_config.json` declaring `QwenImage21Model`, sharded `*.weights.json` / +(`zm_config.json` declaring `ZMDiffusion` with `model_type: "qwen_image_21"`, sharded `*.weights.json` / `*.weights.h5`, `tokenizer.json`) loads with `from_weights("/")`. The `hf:` prefix raises for diffusion models: convert Diffusers format once with `convert_qwen_image_21_diffusers_to_keras.py` and host the result. diff --git a/docs/sentence_transformer.md b/docs/sentence_transformer.md index 073442df..4962651c 100644 --- a/docs/sentence_transformer.md +++ b/docs/sentence_transformer.md @@ -19,12 +19,22 @@ verified on the Keras Torch backend; TensorFlow and JAX parity have not yet been checked for all 13 checkpoints. The public API is `zeromodels.sentence_transformers`, separate from the upstream -`sentence_transformers` reference library. The former `zeromodels.st_models` -and `zeromodels.models.sentence_transformer` paths remain compatibility aliases -for existing code and saved metadata; new examples and exports use the new namespace. -Use `ZMSentenceTransformer` and `ZMSentenceTransformerTokenizer`; the former -`SentenceTransformer`, `SentenceTransformerModel`, and `SentenceTransformerTokenizer` -names remain compatibility aliases. +`sentence_transformers` reference library, and it is exactly two classes: +`ZMSentenceTransformer` and `ZMSentenceTransformerTokenizer`. Hub repos whose +`zm_config.json` or `zm_preprocessor.json` still name the former classes +(`SentenceTransformerModel`, `SentenceTransformer`, `SentenceTransformerTextEmbed`, +`SentenceTransformerTokenizer`) load into these two unchanged. + +| | | +|---|---| +| Class | `ZMSentenceTransformer` | +| Tokenizer | `ZMSentenceTransformerTokenizer` | +| Selector | `backbone_type` (one architecture, pluggable encoder) | +| Recipes | `SENTENCE_TRANSFORMER_BACKBONES` maps each `backbone_type` to its encoder builder: the `zeromodels.models` MPNet / BERT / RoBERTa / XLM-R encoders, and DistilBERT / ALBERT builders in `sentence_transformer_backbones.py` | +| Hub files | `zm_config.json` + `zm_preprocessor.json` + `tokenizer.json` | + +Supported `backbone_type` values: `mpnet`, `bert`, `roberta`, `xlm-roberta`, +`distilbert`, `albert`. ## Usage @@ -135,8 +145,8 @@ tokenizer = ZMSentenceTransformerTokenizer.from_weights( An explicit commit SHA can be supplied as `revision=` to pin the upstream snapshot. -`AutoZMTextEmbed` resolves ZeroModels `sentence_transformer` configs to the -`SentenceTransformerTextEmbed` task view. For upstream `hf:` repos, load the +`AutoZMTextEmbed` resolves ZeroModels `sentence_transformer` configs to +`ZMSentenceTransformer` as well. For upstream `hf:` repos, load the concrete class as shown above: their root `model_type` describes the backbone, not the sentence-embedding task. diff --git a/docs/stable_diffusion.md b/docs/stable_diffusion.md index 2993961d..c56a5709 100644 --- a/docs/stable_diffusion.md +++ b/docs/stable_diffusion.md @@ -8,13 +8,13 @@ Stable Diffusion 1.x, ported to pure Keras 3: latent text-to-image diffusion with a -UNet denoiser, a KL-regularized autoencoder and the CLIP ViT-L/14 text encoder. The -whole model is **one container**, `StableDiffusionModel`, a single functional graph with -the three towers as disconnected sub-graphs (the way `CLIPModel` holds its two), so a -hosted repo is one `model.weights.h5` and one `zm_config.json`. -`StableDiffusionTextToImage` is that same graph plus `generate`: the classifier-free -guidance denoising loop of `BaseDiffusion`, the diffusion counterpart of -`BaseGeneration`. +UNet denoiser, a KL-regularized autoencoder and the CLIP ViT-L/14 text encoder. It loads +through [`ZMDiffusion`](diffusers.md), the one diffusion pipeline class of zeromodels, +with `model_type="stable_diffusion"`: a single functional graph with the three towers as +disconnected sub-graphs (the way `CLIPModel` holds its two), so a hosted repo is one +`model.weights.h5` and one `zm_config.json`, plus `generate`, the classifier-free +guidance denoising loop of `BaseDiffusion` (the diffusion counterpart of +`BaseGeneration`). Prompts are tokenized by `ZMDiffusionTokenizer`. All five SD 1.x checkpoints (v1-1 to v1-5) share this architecture; only the weights differ. The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion @@ -74,13 +74,14 @@ ViT-L/14 defaults), plus the checkpoint's `scheduler_config` and the CLIP specia ids. Each repo's `zm_config.json` parses through it; the constructor stays flat, with the sub-config fields prefixed `unet_` / `vae_` / `text_`. -### `StableDiffusionTextToImage` +### `ZMDiffusion` -The text-to-image task: the `StableDiffusionModel` container plus `BaseDiffusion`'s -`generate`. It supplies the hooks the mixin needs (`encode_prompt` on the CLIP tower, -the empty prompt's `unconditional_ids`, `predict_noise` on the UNet, `decode_latents` on -the VAE) and a `scheduler`, built from the config's `scheduler_config` (PNDM for every -SD 1.x checkpoint). +The pipeline class of every diffusion family (see [Diffusers](diffusers.md)); a repo's +`model_type` (`"stable_diffusion"` here) selects the SD 1.x family, which supplies the +hooks `BaseDiffusion`'s `generate` needs (`encode_prompt` on the CLIP tower, the empty +prompt's `unconditional_ids`, `predict_noise` on the UNet, `decode_latents` on the VAE) +and a `scheduler`, built from the config's `scheduler_config` (PNDM for every SD 1.x +checkpoint). ```python generate( @@ -111,6 +112,7 @@ shapes are the same at every step, so one compile serves the whole loop. | Constructor arg | Default | Meaning | |---|---|---| +| `model_type` | required | `"stable_diffusion"` (read from the repo's `zm_config.json` on `from_weights`) | | `scheduler` | `None` | a `BaseScheduler`; built from `scheduler_config` when unset | | `unet_sample_size` | `64` | latent size the UNet graph is built for (image size / 8) | | `vae_sample_size` | `512` | image size the VAE graphs are built for | @@ -118,14 +120,12 @@ shapes are the same at every step, so one compile serves the whole loop. | `bos_token_id` / `eos_token_id` / `pad_token_id` | `49406` / `49407` / `49407` | CLIP's special tokens, used to build the empty prompt | | `scheduler_config` | `None` | the diffusers scheduler dict (`_class_name`, betas, `steps_offset`, ...) | -### `StableDiffusionModel` - -The container: one functional model whose graph is three disconnected paths, one per -component. Inputs `{"sample", "timestep", "encoder_hidden_states"}` (UNet), -`{"image", "latent"}` (VAE encode / decode) and `{"token_ids", "padding_mask"}` (text); -outputs `{"noise_pred", "moments", "image", "text_embeds"}`. The components are exposed as -`.unet`, `.vae` and `.text_encoder` and share the container's weights, so loading the -container loads all three. It loads the same repo as the task class. +The graph is three disconnected paths, one per component. Inputs +`{"sample", "timestep", "encoder_hidden_states"}` (UNet), `{"image", "latent"}` (VAE +encode / decode) and `{"token_ids", "padding_mask"}` (text); outputs +`{"noise_pred", "moments", "image", "text_embeds"}`. The components are exposed as +`.unet`, `.vae` and `.text_encoder` and share the model's weights, so one load fills all +three, and each can be called on its own. ### `UNet2DConditionModel` @@ -189,15 +189,16 @@ model.scheduler = EulerDiscreteScheduler.from_config(model.config.scheduler_conf ## Preprocessing -### `StableDiffusionTokenizer` +### `ZMDiffusionTokenizer` -The CLIP ViT-L/14 BPE tokenizer: `<|startoftext|>` / `<|endoftext|>` framing, truncated -and `<|endoftext|>`-padded to 77 tokens, the fixed context length the UNet's -cross-attention is built for. Calling it returns `{"input_ids", "attention_mask"}` -for a string or a list of strings. +The tokenizer of every diffusion family; for `model_type="stable_diffusion"` it is the +CLIP ViT-L/14 BPE tokenizer: `<|startoftext|>` / `<|endoftext|>` framing, truncated and +`<|endoftext|>`-padded to 77 tokens, the fixed context length the UNet's cross-attention +is built for. Calling it returns `{"input_ids", "attention_mask"}` for a string or a list +of strings. `from_weights("zeromodels/")` reads the family from the repo. ```python -StableDiffusionTokenizer(hf_id=None, tokenizer_file=None, max_seq_len=77) +ZMDiffusionTokenizer(model_type="stable_diffusion", hf_id=None, tokenizer_file=None, max_seq_len=77) ``` There is no image processor and no processor class: text-to-image takes token ids in @@ -213,13 +214,10 @@ import os os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" from PIL import Image -from zeromodels.models.stable_diffusion import ( - StableDiffusionTextToImage, - StableDiffusionTokenizer, -) +from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer -model = StableDiffusionTextToImage.from_weights("zeromodels/stable-diffusion-v1-5") -tokenizer = StableDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-v1-5") +model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-v1-5") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-v1-5") inputs = tokenizer( "a red fox sitting in a field of lavender, golden hour, wildlife photography, sharp focus" @@ -265,7 +263,7 @@ images = model.generate(**tokenizer("a bowl of ramen"), latents=latents) ### Verified against diffusers The fox prompt above with the same initial latent, PNDM 50 steps and guidance 7.5 -through `StableDiffusionTextToImage` (left) and diffusers' `StableDiffusionPipeline` +through `ZMDiffusion` (left) and diffusers' `StableDiffusionPipeline` (right), both fp32: zeromodels (left) and diffusers (right) generations of a red fox in lavender from the same latent @@ -278,12 +276,12 @@ vae decode max|d|=3.5e-06 final image (uint8) max|d|=1 mean|d|=0.0022 PSNR=74.7 dB identical_pixels=99.34% ``` -### Container only +### Components on their own ```python -from zeromodels.models.stable_diffusion import StableDiffusionModel +from zeromodels.diffusers import ZMDiffusion -sd = StableDiffusionModel.from_weights("zeromodels/stable-diffusion-v1-5") +sd = ZMDiffusion.from_weights("zeromodels/stable-diffusion-v1-5") latent = sd.vae.encode(image) # (1, 512, 512, 3) in [-1, 1] -> (1, 64, 64, 4) text = sd.text_encoder({"token_ids": ids, "padding_mask": mask})["last_hidden_state"] ``` @@ -294,7 +292,7 @@ The graphs are built for a fixed size, the weights are not. Rebuild at another s the constructor overrides (multiples of 64px; SD 1.x was trained at 512px): ```python -model = StableDiffusionTextToImage.from_weights( +model = ZMDiffusion.from_weights( "zeromodels/stable-diffusion-v1-5", unet_sample_size=96, vae_sample_size=768 ) images = model.generate(**tokenizer("a mountain lake at dawn")) # (1, 768, 768, 3) @@ -324,7 +322,7 @@ import keras keras.config.set_image_data_format("channels_first") -model = StableDiffusionTextToImage.from_weights("zeromodels/stable-diffusion-v1-5") +model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-v1-5") images = model.generate( **tokenizer("a bowl of ramen"), seed=0 ) # still (1, 512, 512, 3) @@ -344,8 +342,10 @@ measured 2-3x faster than an op-by-op run; TensorFlow uses the same mechanism. ## Loading Fine-tuned Weights The five hosted checkpoints are the supported weights; any repo laid out like them -(`zm_config.json` declaring `StableDiffusionModel`, `model.weights.h5`, `tokenizer.json`) -loads with `from_weights("/")`. The `hf:` prefix raises for diffusion models: +(`zm_config.json` declaring `ZMDiffusion` with `model_type: "stable_diffusion"`, +`model.weights.h5`, `tokenizer.json`) loads with `from_weights("/")`; configs +written before `ZMDiffusion` (declaring `StableDiffusionModel`) load the same way. The +`hf:` prefix raises for diffusion models: convert a diffusers-format checkpoint once with -`zeromodels/models/stable_diffusion/convert_stable_diffusion_diffusers_to_keras.py` +`zeromodels/diffusers/stable_diffusion/convert_stable_diffusion_diffusers_to_keras.py` (`transfer_stable_diffusion(repo)`, `pip install zeromodels[conversion]`) and host the result. diff --git a/docs/stable_diffusion_2.md b/docs/stable_diffusion_2.md index f08fb024..84a2072c 100644 --- a/docs/stable_diffusion_2.md +++ b/docs/stable_diffusion_2.md @@ -8,11 +8,11 @@ Stable Diffusion 2.x, ported to pure Keras 3: the [Stable Diffusion](stable_diffusion.md) -latent-diffusion architecture with the second-generation conditioning. It reuses the -SD 1.x family end to end (`StableDiffusion2Model` is `StableDiffusionModel`, -`StableDiffusion2TextToImage` is `StableDiffusionTextToImage`, each with the SD 2 -configuration), so everything on that page applies: one container per repo, `generate` -from `BaseDiffusion`, the schedulers, both data formats. What changes is the config: +latent-diffusion architecture with the second-generation conditioning. It loads +through [`ZMDiffusion`](diffusers.md) with `model_type="stable_diffusion_2"`, the SD 1.x +family with the SD 2 configuration, so everything on that page applies: one functional +model and one weights file per repo, `generate` from `BaseDiffusion`, the schedulers, +both data formats. What changes is the config: - **Text encoder**: OpenCLIP ViT-H/14's text tower, 1024 wide, 16 heads, `gelu`, shipped truncated to its 23rd (penultimate) layer, which SD 2 conditions on. @@ -68,9 +68,10 @@ with the SD 2 widths), `AutoencoderKLConfig` and `StableDiffusion2TextConfig` (a `CLIPTextConfig` with the ViT-H/14 sizes), plus the scheduler config and the token ids (`pad_token_id` 0). Flat constructor, `unet_` / `vae_` / `text_` prefixes, like SD 1.x. -### `StableDiffusion2TextToImage` +### `ZMDiffusion` -`StableDiffusionTextToImage` with the SD 2 configuration; `generate` is unchanged: +`ZMDiffusion` with `model_type="stable_diffusion_2"` (read from the repo); `generate` is +the SD 1.x one: ```python generate( @@ -87,14 +88,11 @@ generate( Returns `(batch, H, W, 3)` uint8 images at the checkpoint's resolution; `latents` is `(batch, 64, 64, 4)` for the 512px checkpoints and `(batch, 96, 96, 4)` for the 768px ones (`channels_first`: channels second). Defaults come from the repo's `generate_args` -(50 steps, guidance 7.5). See [Stable Diffusion](stable_diffusion.md#stablediffusiontexttoimage) +(50 steps, guidance 7.5). See [Stable Diffusion](stable_diffusion.md#zmdiffusion) for the argument table. -### `StableDiffusion2Model` - -The container, `StableDiffusionModel` with the SD 2 configuration: the same three -disconnected paths (`.unet`, `.vae`, `.text_encoder`) and the same inputs / outputs, with -`encoder_hidden_states` 1024 wide. It loads the same repo as the task class. +The graph is SD 1.x's three disconnected paths (`.unet`, `.vae`, `.text_encoder`) with +the same inputs / outputs, `encoder_hidden_states` 1024 wide. The components are the SD 1.x classes, `UNet2DConditionModel` (built with `num_attention_heads=(5, 10, 20, 20)`, `use_linear_projection=True`, @@ -103,16 +101,16 @@ The components are the SD 1.x classes, `UNet2DConditionModel` (built with ## Preprocessing -### `StableDiffusion2Tokenizer` +### `ZMDiffusionTokenizer` -`StableDiffusionTokenizer` with `!` as the pad token: CLIP BPE, `<|startoftext|>` / -`<|endoftext|>` framing, truncated and `!`-padded to 77 tokens. Returns -`{"input_ids", "attention_mask"}`. +`ZMDiffusionTokenizer` with `model_type="stable_diffusion_2"`, which pads with `!`: +CLIP BPE, `<|startoftext|>` / `<|endoftext|>` framing, truncated and `!`-padded to 77 +tokens. Returns `{"input_ids", "attention_mask"}`. ```python -StableDiffusion2Tokenizer( - hf_id=None, tokenizer_file=None, max_seq_len=77, pad_token="!" -) +ZMDiffusionTokenizer( + model_type="stable_diffusion_2", hf_id=None, tokenizer_file=None, max_seq_len=77 +) # pad_token defaults to "!" ``` ## End-to-end example @@ -123,13 +121,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-1-base") -tokenizer = StableDiffusion2Tokenizer.from_weights( +model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-2-1-base") +tokenizer = ZMDiffusionTokenizer.from_weights( "zeromodels/stable-diffusion-2-1-base" ) @@ -149,8 +144,8 @@ The 768px checkpoints need no extra arguments; the v-prediction DDIM scheduler a 96x96 latent come from the repo: ```python -model = StableDiffusion2TextToImage.from_weights("zeromodels/stable-diffusion-2-1") -tokenizer = StableDiffusion2Tokenizer.from_weights("zeromodels/stable-diffusion-2-1") +model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-2-1") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-2-1") images = model.generate( **tokenizer("a lighthouse on a cliff at dusk, oil painting") ) # (1, 768, 768, 3) @@ -161,8 +156,8 @@ images = model.generate( One step, no guidance (the repo's defaults; up to 4 steps sharpen a little): ```python -model = StableDiffusion2TextToImage.from_weights("zeromodels/sd-turbo") -tokenizer = StableDiffusion2Tokenizer.from_weights("zeromodels/sd-turbo") +model = ZMDiffusion.from_weights("zeromodels/sd-turbo") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/sd-turbo") images = model.generate( **tokenizer( "a cinematic shot of a baby raccoon wearing an intricate italian priest robe" @@ -176,7 +171,7 @@ The same prompt and latent through diffusers' `StableDiffusionPipeline` (fp32) m 1 uint8 level: 99.9% of pixels identical at 1 step (PSNR 82 dB), 99.8% at 4 steps. Negative prompts, batching, explicit `latents` for cross-backend reproducibility, other -resolutions and the container-only use work exactly as on the +resolutions and the components on their own work exactly as on the [Stable Diffusion](stable_diffusion.md#end-to-end-example) page. Image-to-image (`generate(..., image=..., strength=...)`) works as described for [`BaseDiffusion`](main_classes.md#basediffusion). @@ -184,7 +179,7 @@ resolutions and the container-only use work exactly as on the ### Verified against diffusers The ramen prompt above with the same initial latent, PNDM 50 steps and guidance 7.5 -through `StableDiffusion2TextToImage` (left) and diffusers' `StableDiffusionPipeline` +through `ZMDiffusion` (left) and diffusers' `StableDiffusionPipeline` (right) on `stable-diffusion-2-1-base`, both fp32: zeromodels (left) and diffusers (right) generations of a bowl of ramen from stable-diffusion-2-1-base with the same latent @@ -216,8 +211,8 @@ runs over 9216 tokens, so the guided run needs roughly 12 GB, or the CPU. ## Loading Fine-tuned Weights Any repo laid out like the hosted ones (`zm_config.json` declaring -`StableDiffusion2Model`, the weights, `tokenizer.json`) loads with -`from_weights("/")`. The `hf:` prefix raises for diffusion models: convert a +`ZMDiffusion` with `model_type: "stable_diffusion_2"`, the weights, `tokenizer.json`) +loads with `from_weights("/")`. The `hf:` prefix raises for diffusion models: convert a diffusers-format SD 2 checkpoint once with -`zeromodels/models/stable_diffusion_2/convert_stable_diffusion_2_diffusers_to_keras.py` +`zeromodels/diffusers/stable_diffusion_2/convert_stable_diffusion_2_diffusers_to_keras.py` (`transfer_stable_diffusion_2(repo)`, `pip install zeromodels[conversion]`) and host the result. diff --git a/docs/stable_diffusion_3.md b/docs/stable_diffusion_3.md index bf699bbe..676698b0 100644 --- a/docs/stable_diffusion_3.md +++ b/docs/stable_diffusion_3.md @@ -11,8 +11,10 @@ Stable Diffusion 3, ported to pure Keras 3: a rectified-flow model whose denoiser is the **MMDiT**, a multimodal diffusion transformer in which the latent tokens and the text tokens attend to each other jointly, conditioned on the timestep and the pooled text embeddings -through adaptive layer norms. It keeps the SD family's shape (one container per repo, -`generate` from `BaseDiffusion`, the schedulers, both data formats) with new parts: +through adaptive layer norms. It loads through [`ZMDiffusion`](diffusers.md) with +`model_type="stable_diffusion_3"` and keeps the SD family's shape (one functional model +and one set of weights per repo, `generate` from `BaseDiffusion`, the schedulers, both +data formats) with new parts: - **MMDiT** (`SD3Transformer2DModel`): the 16-channel latent is patchified (2x2) into 1536-d tokens with a sinusoidal position table cropped from a 192x192 grid; 24 joint blocks of @@ -24,7 +26,7 @@ through adaptive layer norms. It keeps the SD family's shape (one container per states are concatenated (2048-d), zero-padded to 4096 and followed along the sequence by the T5 states (77 + 256 tokens); the projected pooled CLIP states (2048-d) feed the conditioning embedding. -- **T5 is optional and separate**: the container holds the MMDiT, the VAE and the two CLIP +- **T5 is optional and separate**: the checkpoint holds the MMDiT, the VAE and the two CLIP towers (2.9B, 5.8 GB float16); the T5-XXL encoder is hosted once for every SD 3 / 3.5 checkpoint (`zeromodels/t5-v1_1-xxl-encoder`, an [`SD3T5EncoderModel`](#sd3t5encodermodel), 9.5 GB float16) and attached with @@ -35,7 +37,7 @@ through adaptive layer norms. It keeps the SD family's shape (one container per - **Sampler**: `FlowMatchEulerDiscreteScheduler` with `shift` 3.0 (rectified flow: the timesteps are `sigma * 1000`, a step is `x + (sigma_next - sigma) * v`); 28 steps at guidance 7.0 by default. No negative prompt means the encoded empty prompt. -- **One tokenizer call**: `StableDiffusion3Tokenizer` runs the CLIP BPE (the second +- **One tokenizer call**: `ZMDiffusionTokenizer` runs the CLIP BPE (the second tower's `!`-padded ids are derived from the mask, as SDXL) and the T5 SentencePiece tokenizer, returning `{"input_ids", "attention_mask", "input_ids_3"}`. - **float16 weights**, the release precision; `from_weights` builds the model in float16 @@ -69,9 +71,10 @@ over `StableDiffusion3TransformerConfig`, `StableDiffusion3VAEConfig`, the scheduler config, the CLIP and T5 token ids and `max_sequence_length` (256). Flat constructor with `transformer_` / `vae_` / `text_` / `text_2_` prefixes. -### `StableDiffusion3TextToImage` +### `ZMDiffusion` -`StableDiffusionTextToImage`'s `generate` over the SD 3 graph: +`ZMDiffusion` with `model_type="stable_diffusion_3"` (read from the repo); `generate` +over the SD 3 graph: ```python generate( @@ -103,11 +106,9 @@ Returns `(batch, H, W, 3)` uint8 images. `from_weights` takes one extra argument, `text_encoder_3`: a hosted `SD3T5EncoderModel` repo id (loaded at the same `load_dtype`) or a built `SD3T5EncoderModel`; `sd.text_encoder_3` can also be assigned later (or set to `None`). It lives outside the -container's weights. +checkpoint's weights; passing `text_encoder_3` for a family other than SD 3 / 3.5 raises. -### `StableDiffusion3Model` - -The container: four disconnected paths in one functional model, `.transformer` / `.vae` / +The graph is four disconnected paths in one functional model, `.transformer` / `.vae` / `.text_encoder` / `.text_encoder_2`. | Inputs | Outputs | @@ -151,10 +152,11 @@ gated-GELU feed-forward of T5 v1.1 (`wo(gelu_tanh(wi_0(x)) * wi_1(x))`, ## Preprocessing -### `StableDiffusion3Tokenizer` +### `ZMDiffusionTokenizer` ```python -StableDiffusion3Tokenizer( +ZMDiffusionTokenizer( + model_type="stable_diffusion_3", hf_id=None, tokenizer_file=None, tokenizer_file_3=None, @@ -173,16 +175,13 @@ import os os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" from PIL import Image -from zeromodels.models.stable_diffusion_3 import ( - StableDiffusion3TextToImage, - StableDiffusion3Tokenizer, -) +from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer -model = StableDiffusion3TextToImage.from_weights( +model = ZMDiffusion.from_weights( "zeromodels/stable-diffusion-3-medium", text_encoder_3="zeromodels/t5-v1_1-xxl-encoder", # omit on small machines: T5 features zeroed ) -tokenizer = StableDiffusion3Tokenizer.from_weights( +tokenizer = ZMDiffusionTokenizer.from_weights( "zeromodels/stable-diffusion-3-medium" ) @@ -238,7 +237,7 @@ Both `channels_last` and `channels_first` are supported; `generate` always retur ## Memory and speed -The container is 5.8 GB in float16 (the T5 encoder adds 9.5 GB). Measured on the torch +The checkpoint is 5.8 GB in float16 (the T5 encoder adds 9.5 GB). Measured on the torch backend with the default float16 load (RTX 4060 Laptop, 8 GB), T5 features precomputed: | Resolution | Denoising step (guided, batch 1) | VAE decode (float32) | @@ -260,9 +259,10 @@ step is compiled once per run. ## Loading Fine-tuned Weights Any repo laid out like the hosted ones (`zm_config.json` declaring -`StableDiffusion3Model`, the weights, the two tokenizer files) loads with +`ZMDiffusion` with `model_type: "stable_diffusion_3"`, the weights, the two tokenizer +files) loads with `from_weights("/")`. The `hf:` prefix raises for diffusion models: convert a diffusers-format SD 3 checkpoint once with -`zeromodels/models/stable_diffusion_3/convert_stable_diffusion_3_diffusers_to_keras.py` +`zeromodels/diffusers/stable_diffusion_3/convert_stable_diffusion_3_diffusers_to_keras.py` (`transfer_stable_diffusion_3(repo)`; `transfer_t5_encoder(repo)` for the T5 tower) and host the result. diff --git a/docs/stable_diffusion_3_5.md b/docs/stable_diffusion_3_5.md index 5a2e843e..b2f2958b 100644 --- a/docs/stable_diffusion_3_5.md +++ b/docs/stable_diffusion_3_5.md @@ -9,10 +9,9 @@ Stable Diffusion 3.5, ported to pure Keras 3: the [Stable Diffusion 3](stable_diffusion_3.md) -architecture, retrained and scaled. It reuses the SD 3 family end to end -(`StableDiffusion3_5Model` is `StableDiffusion3Model`, `StableDiffusion3_5TextToImage` is -`StableDiffusion3TextToImage`, each with the SD 3.5 configuration), so everything on that -page applies: the MMDiT, the three text encoders with the separately hosted T5-XXL, the +architecture, retrained and scaled. It loads through +[`ZMDiffusion`](diffusers.md) with `model_type="stable_diffusion_3_5"`, the SD 3 family +with the SD 3.5 configuration, so everything on that page applies: the MMDiT, the three text encoders with the separately hosted T5-XXL, the 16-channel VAE, the flow-match sampler, one tokenizer call. What changes is the config: - **QK normalization**: every attention block RMS-normalizes its queries and keys per head @@ -38,7 +37,7 @@ See also [stable_diffusion_3.md](stable_diffusion_3.md). ## Variants -Each repo is one container: MMDiT + VAE (84M) + CLIP ViT-L/14 (123M) + OpenCLIP +Each repo is one checkpoint: MMDiT + VAE (84M) + CLIP ViT-L/14 (123M) + OpenCLIP ViT-bigG/14 (695M); the T5-XXL encoder (`zeromodels/t5-v1_1-xxl-encoder`, the SD 3 family's `SD3T5EncoderModel`, 9.5 GB) is shared and optional. @@ -54,8 +53,8 @@ family's `SD3T5EncoderModel`, 9.5 GB) is shared and optional. `StableDiffusion3Config` over `StableDiffusion3_5TransformerConfig` (the SD 3.5 large defaults: `num_layers` 38, `num_attention_heads` 38, `caption_projection_dim` 2432, `qk_norm="rms_norm"`; medium's repo carries `dual_attention_layers=(0, ..., 12)` and -`pos_embed_max_size=384`). `StableDiffusion3_5TextToImage.generate` and -`StableDiffusion3_5Tokenizer` are the SD 3 ones; see +`pos_embed_max_size=384`). `ZMDiffusion.generate` and +`ZMDiffusionTokenizer` (`model_type="stable_diffusion_3_5"`) behave as for SD 3; see [Stable Diffusion 3](stable_diffusion_3.md#api). ## End-to-end example @@ -66,16 +65,13 @@ 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, -) +from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer -model = StableDiffusion3_5TextToImage.from_weights( +model = ZMDiffusion.from_weights( "zeromodels/stable-diffusion-3.5-medium", text_encoder_3="zeromodels/t5-v1_1-xxl-encoder", # optional ) -tokenizer = StableDiffusion3_5Tokenizer.from_weights( +tokenizer = ZMDiffusionTokenizer.from_weights( "zeromodels/stable-diffusion-3.5-medium" ) @@ -92,7 +88,7 @@ Image.fromarray(images[0]).save("treehouse.png") # (1024, 1024, 3) uint8 ### Large Turbo ```python -model = StableDiffusion3_5TextToImage.from_weights( +model = ZMDiffusion.from_weights( "zeromodels/stable-diffusion-3.5-large-turbo" ) images = model.generate( @@ -114,7 +110,7 @@ jax, and by the SD 3 float32 end-to-end check (92.8% identical pixels, PSNR 63.9 ## Memory and speed -The medium container is 6.8 GB in float16, the large ones 17.8 GB (a 24 GB GPU); the +The medium checkpoint is 6.8 GB in float16, the large ones 17.8 GB (a 24 GB GPU); the T5-XXL encoder adds 9.5 GB. See [Stable Diffusion 3](stable_diffusion_3.md#memory-and-speed): the medium MMDiT-X at 1024px peaks at 7.5 GB during the guided denoising step on the torch backend (fused attention, float16), which an 8 GB GPU runs with the VAE decode on the CPU. @@ -122,7 +118,8 @@ backend (fused attention, float16), which an 8 GB GPU runs with the VAE decode o ## Loading Fine-tuned Weights Any repo laid out like the hosted ones (`zm_config.json` declaring -`StableDiffusion3_5Model`, the weights, the two tokenizer files) loads with +`ZMDiffusion` with `model_type: "stable_diffusion_3_5"`, the weights, the two tokenizer +files) loads with `from_weights("/")`. Convert a diffusers-format checkpoint once with -`zeromodels/models/stable_diffusion_3_5/convert_stable_diffusion_3_5_diffusers_to_keras.py` +`zeromodels/diffusers/stable_diffusion_3_5/convert_stable_diffusion_3_5_diffusers_to_keras.py` and host the result. diff --git a/docs/stable_diffusion_xl.md b/docs/stable_diffusion_xl.md index c62ad479..ebb4f440 100644 --- a/docs/stable_diffusion_xl.md +++ b/docs/stable_diffusion_xl.md @@ -10,11 +10,12 @@ Stable Diffusion XL, ported to pure Keras 3: the [Stable Diffusion](stable_diffusion.md) latent-diffusion recipe scaled up to 1024px with a 2.6B-parameter UNet and two text encoders, plus the **refiner**, an image-to-image expert that finishes the base model's -latents, and **SDXL-Turbo**, the base distilled to a few steps. It subclasses the SD 1.x -family (`StableDiffusionXLModel` is a `StableDiffusionModel` with a fourth component, -`StableDiffusionXLTextToImage` is `StableDiffusionTextToImage` with the SDXL conditioning -hooks; the refiner classes subclass those), so everything on that page applies: one -container per repo, `generate` from `BaseDiffusion`, the schedulers, both data formats. +latents, and **SDXL-Turbo**, the base distilled to a few steps. It loads +through [`ZMDiffusion`](diffusers.md) with `model_type="stable_diffusion_xl"` (the +refiner: `"stable_diffusion_xl_refiner"`), a family built on SD 1.x's with a fourth +component and the SDXL conditioning hooks, so everything on that page applies: one +functional model and one set of weights per repo, `generate` from `BaseDiffusion`, the +schedulers, both data formats. What changes: - **Two text encoders**: CLIP ViT-L/14's text tower (768-d, the SD 1.x one) and OpenCLIP @@ -23,8 +24,8 @@ What changes: concatenated into the UNet's 2048-d cross-attention context, and the second tower's projected pooled (EOT) state is the `text_embeds` micro-conditioning. - **One tokenizer call**: both towers use the same CLIP BPE and differ only in the pad - token (`<|endoftext|>` vs `!`), so `StableDiffusionXLTokenizer` is the first one and the - model derives the second tower's ids from the returned `attention_mask`. + token (`<|endoftext|>` vs `!`), so `ZMDiffusionTokenizer` returns the first one's ids + and the model derives the second tower's ids from the returned `attention_mask`. - **Micro-conditioning**: the UNet's timestep embedding also receives the pooled text embedding and six size / crop values (`original_size`, `crops_coords_top_left`, `target_size`, each `(height, width)`), sinusoidally embedded and projected by the @@ -39,8 +40,8 @@ What changes: - **Scheduler**: Euler with `leading` timestep spacing (`steps_offset` 1) for the base models and the refiners, ancestral Euler with `trailing` spacing for SDXL-Turbo; all come from the repo's `scheduler_config`. -- **Refiner**: a second model (`StableDiffusionXLRefinerModel` / - `StableDiffusionXLRefinerImageToImage`) with the OpenCLIP tower alone (1280-d +- **Refiner**: a second checkpoint (a `ZMDiffusion` with + `model_type="stable_diffusion_xl_refiner"`) with the OpenCLIP tower alone (1280-d context), a four-level UNet ((384, 768, 1536, 1536), attention at the middle levels, 4 blocks per Transformer2D, 2.3B parameters) and an **aesthetic score** as the fifth micro-conditioning id. It refines a latent the base left partially denoised @@ -71,16 +72,16 @@ Preconverted weights are hosted under `zeromodels/`. Load with 3.47B parameters, 6.6 GB in float16 (two `model.weights.json` shards); a refiner repo is UNet (2.26B) + VAE + OpenCLIP text encoder, 3.04B parameters, 5.8 GB. -| Variant | Hub | Class | Resolution | Sampler | License | +| Variant | Hub | `model_type` | Resolution | Sampler | License | |---|---|---|---|---|---| -| `stable-diffusion-xl-base-1.0` | [`zeromodels/stable-diffusion-xl-base-1.0`](https://huggingface.co/zeromodels/stable-diffusion-xl-base-1.0) | `StableDiffusionXLTextToImage` | 1024 | Euler, 50 steps, guidance 5.0 | CreativeML Open RAIL++-M | -| `stable-diffusion-xl-refiner-1.0` | [`zeromodels/stable-diffusion-xl-refiner-1.0`](https://huggingface.co/zeromodels/stable-diffusion-xl-refiner-1.0) | `StableDiffusionXLRefinerImageToImage` | 1024 | Euler, 50 steps, guidance 5.0, strength 0.3 | CreativeML Open RAIL++-M | -| `sdxl-turbo` | [`zeromodels/sdxl-turbo`](https://huggingface.co/zeromodels/sdxl-turbo) | `StableDiffusionXLTextToImage` | 512 | Euler ancestral, 1 step, no guidance | Stability AI Non-Commercial Research Community | +| `stable-diffusion-xl-base-1.0` | [`zeromodels/stable-diffusion-xl-base-1.0`](https://huggingface.co/zeromodels/stable-diffusion-xl-base-1.0) | `stable_diffusion_xl` | 1024 | Euler, 50 steps, guidance 5.0 | CreativeML Open RAIL++-M | +| `stable-diffusion-xl-refiner-1.0` | [`zeromodels/stable-diffusion-xl-refiner-1.0`](https://huggingface.co/zeromodels/stable-diffusion-xl-refiner-1.0) | `stable_diffusion_xl_refiner` | 1024 | Euler, 50 steps, guidance 5.0, strength 0.3 | CreativeML Open RAIL++-M | +| `sdxl-turbo` | [`zeromodels/sdxl-turbo`](https://huggingface.co/zeromodels/sdxl-turbo) | `stable_diffusion_xl` | 512 | Euler ancestral, 1 step, no guidance | Stability AI Non-Commercial Research Community | The SDXL 0.9 research previews (`stabilityai/stable-diffusion-xl-base-0.9` and `-refiner-0.9`, the same architecture with different weights, gated behind the research license) are not hosted; with an authorized token the converter turns them into the same -containers (`transfer_stable_diffusion_xl(repo, token=...)`, see +checkpoints (`transfer_stable_diffusion_xl(repo, token=...)`, see [Loading Fine-tuned Weights](#loading-fine-tuned-weights)). SDXL-Turbo is a research-only, non-commercial release; its repo defaults (`generate_args`) are 1 step and `guidance_scale=0.0`, which is how it was trained (do not add guidance). @@ -98,10 +99,10 @@ Flat constructor with `unet_` / `vae_` / `text_` / `text_2_` prefixes. same with `text_config=None`, the `StableDiffusionXLRefinerUNetConfig` denoiser, `requires_aesthetics_score=True` and `force_zeros_for_empty_prompt=False`. -### `StableDiffusionXLTextToImage` +### `ZMDiffusion` -`StableDiffusionTextToImage` over the SDXL graph; `generate` gains the micro-conditioning -arguments: +`ZMDiffusion` with `model_type="stable_diffusion_xl"` (read from the repo); `generate` +gains the micro-conditioning arguments: ```python generate( @@ -132,7 +133,7 @@ generate( |---|---| | `input_ids`, `attention_mask` | `**tokenizer(prompts)`; the mask tells the model which positions are padding, so the second tower's `!`-padded ids can be derived (without it, everything after the first `<|endoftext|>` counts as padding). | | `negative_input_ids` | Tokenized negative prompt, one row per prompt or a single row shared by the batch. Without it the unconditional branch is zero embeddings (`force_zeros_for_empty_prompt`), as in diffusers. | -| `num_inference_steps`, `guidance_scale`, `seed`, `latents` | As for [Stable Diffusion](stable_diffusion.md#stablediffusiontexttoimage); defaults from the repo's `generate_args` (50 steps, guidance 5.0 for the base model). | +| `num_inference_steps`, `guidance_scale`, `seed`, `latents` | As for [Stable Diffusion](stable_diffusion.md#zmdiffusion); defaults from the repo's `generate_args` (50 steps, guidance 5.0 for the base model). | | `image`, `strength`, `denoising_start`, `denoising_end`, `output_type` | The image-to-image and ensemble controls of [`BaseDiffusion.generate`](main_classes.md#basediffusion): refine an `image` (noised to `strength`), stop early (`denoising_end`) and hand the latent over (`output_type="latent"`), or resume one (`denoising_start`). | | `original_size`, `crops_coords_top_left`, `target_size` | SDXL's size / crop conditioning, `(height, width)` in pixels; the sizes default to the generated image's size, the crop to `(0, 0)`. | | `aesthetic_score` | The refiner's fifth conditioning value (6.0), in place of `target_size`; ignored by the base models. | @@ -147,21 +148,21 @@ original_size, crops_coords_top_left, target_size)` returns the dict the UNet ta (batch, 6)); `encode_negative_prompt` zeroes the first two when no negative ids are given; `predict_noise(latents, timesteps, embeddings)` runs the UNet. -### `StableDiffusionXLRefinerImageToImage` +#### The refiner -The refiner task: the same class over the refiner container (`text_encoder_2` alone, +The same class with `model_type="stable_diffusion_xl_refiner"` (`text_encoder_2` alone, `aesthetic_score` in the conditioning, `generate_args` with `strength` 0.3). Its `generate` is the one above; with no `image` / `latents` it runs text-to-image, which the refiner was not trained for. The two intended uses: ```python -base = StableDiffusionXLTextToImage.from_weights( +base = ZMDiffusion.from_weights( "zeromodels/stable-diffusion-xl-base-1.0" ) -refiner = StableDiffusionXLRefinerImageToImage.from_weights( +refiner = ZMDiffusion.from_weights( "zeromodels/stable-diffusion-xl-refiner-1.0" ) -tokenizer = StableDiffusionXLTokenizer.from_weights( +tokenizer = ZMDiffusionTokenizer.from_weights( "zeromodels/stable-diffusion-xl-base-1.0" ) inputs = tokenizer("a lighthouse on a rocky coast at dawn, dramatic clouds, cinematic") @@ -178,9 +179,9 @@ images = refiner.generate( images = refiner.generate(**inputs, image=images, strength=0.3, seed=0) ``` -### `StableDiffusionXLModel` +#### The graph -The container: four disconnected paths in one functional model, `.unet` / `.vae` / +Four disconnected paths in one functional model, `.unet` / `.vae` / `.text_encoder` / `.text_encoder_2`. | Inputs | Outputs | @@ -193,22 +194,23 @@ Each text tower is a functional CLIP text model (the `clip` family's layers) who are `penultimate_hidden_state`, `last_hidden_state`, `pooler_output` and, for the second, `text_embeds`. The UNet is the SD 1.x `UNet2DConditionModel` with the SDXL config (`transformer_layers_per_block=(1, 2, 10)`, `addition_embed_type="text_time"`); the VAE -is `AutoencoderKL` with `force_upcast=True`. `StableDiffusionXLRefinerModel` is the same -container without the first tower (`token_ids_2` + `padding_mask` in, `prompt_embeds` +is `AutoencoderKL` with `force_upcast=True`. The refiner is the same +graph without the first tower (`token_ids_2` + `padding_mask` in, `prompt_embeds` (B, 77, 1280) out) and with `time_ids` (B, 5). ## Preprocessing -### `StableDiffusionXLTokenizer` +### `ZMDiffusionTokenizer` -`StableDiffusionTokenizer` (CLIP BPE, `<|startoftext|>` / `<|endoftext|>` framing, +`ZMDiffusionTokenizer` with `model_type="stable_diffusion_xl"` (CLIP BPE, +`<|startoftext|>` / `<|endoftext|>` framing, `<|endoftext|>`-padded to 77 tokens). Returns `{"input_ids", "attention_mask"}`; pass both to `generate`. ```python -StableDiffusionXLTokenizer( - hf_id=None, tokenizer_file=None, max_seq_len=77, pad_token="<|endoftext|>" -) +ZMDiffusionTokenizer( + model_type="stable_diffusion_xl", hf_id=None, tokenizer_file=None, max_seq_len=77 +) # pad_token defaults to "<|endoftext|>"; the refiner repos use the same settings ``` A literal `!` in a prompt is encoded as OpenCLIP encodes it (the BPE token `!`, id @@ -223,15 +225,12 @@ import os os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" from PIL import Image -from zeromodels.models.stable_diffusion_xl import ( - StableDiffusionXLTextToImage, - StableDiffusionXLTokenizer, -) +from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer -model = StableDiffusionXLTextToImage.from_weights( +model = ZMDiffusion.from_weights( "zeromodels/stable-diffusion-xl-base-1.0" ) -tokenizer = StableDiffusionXLTokenizer.from_weights( +tokenizer = ZMDiffusionTokenizer.from_weights( "zeromodels/stable-diffusion-xl-base-1.0" ) @@ -251,7 +250,7 @@ The SDXL "ensemble of experts": the base runs the first 80% of the schedule and latent to the refiner, which finishes it. Same prompt and initial latent as above: ```python -refiner = StableDiffusionXLRefinerImageToImage.from_weights( +refiner = ZMDiffusion.from_weights( "zeromodels/stable-diffusion-xl-refiner-1.0" ) latent = model.generate( @@ -269,8 +268,8 @@ images = refiner.generate( One step, no guidance, 512px (the repo's defaults): ```python -model = StableDiffusionXLTextToImage.from_weights("zeromodels/sdxl-turbo") -tokenizer = StableDiffusionXLTokenizer.from_weights("zeromodels/sdxl-turbo") +model = ZMDiffusion.from_weights("zeromodels/sdxl-turbo") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/sdxl-turbo") images = model.generate( **tokenizer( "a cinematic shot of a baby raccoon wearing an intricate italian priest robe" @@ -301,7 +300,7 @@ The graphs are built for the repo's resolution; pass overrides to build for anot (the weights are resolution-independent; multiples of 64px): ```python -model = StableDiffusionXLTextToImage.from_weights( +model = ZMDiffusion.from_weights( "zeromodels/stable-diffusion-xl-base-1.0", unet_sample_size=(96, 128), vae_sample_size=(768, 1024), @@ -312,7 +311,7 @@ images = model.generate(**tokenizer("a wide mountain valley")) # (1, 768, 1024, ### float32 ```python -model = StableDiffusionXLTextToImage.from_weights( +model = ZMDiffusion.from_weights( "zeromodels/stable-diffusion-xl-base-1.0", load_dtype="float32" ) ``` @@ -378,8 +377,9 @@ per run (see [`BaseDiffusion`](main_classes.md#basediffusion)). ## Loading Fine-tuned Weights Any repo laid out like the hosted ones (`zm_config.json` declaring -`StableDiffusionXLModel`, the weights, `tokenizer.json`) loads with +`ZMDiffusion` with `model_type: "stable_diffusion_xl"` or +`"stable_diffusion_xl_refiner"`, the weights, `tokenizer.json`) loads with `from_weights("/")`. The `hf:` prefix raises for diffusion models: convert a diffusers-format SDXL checkpoint once with -`zeromodels/models/stable_diffusion_xl/convert_stable_diffusion_xl_diffusers_to_keras.py` +`zeromodels/diffusers/stable_diffusion_xl/convert_stable_diffusion_xl_diffusers_to_keras.py` (`transfer_stable_diffusion_xl(repo)`, `pip install zeromodels[conversion]`) and host the result. diff --git a/pyproject.toml b/pyproject.toml index 019f301d..3180f8cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -105,6 +105,7 @@ source = ["zeromodels", "tests"] omit = [ "zeromodels/models/*/convert_*.py", "zeromodels/sentence_transformers/convert_*.py", + "zeromodels/diffusers/*/convert_*.py", ] [tool.coverage.report] diff --git a/tests/base/model_test_registry.py b/tests/base/model_test_registry.py index 673c0490..af93d631 100644 --- a/tests/base/model_test_registry.py +++ b/tests/base/model_test_registry.py @@ -3021,30 +3021,6 @@ "token_embeddings": (2, 16, 32), }, }, - "SentenceTransformerTextEmbed": { - "module": "zeromodels.sentence_transformers", - "model_cls": "SentenceTransformerTextEmbed", - "model_type": "llm", - "init_kwargs": { - "backbone_type": "mpnet", - "vocab_size": 128, - "embed_dim": 32, - "num_layers": 2, - "num_heads": 2, - "mlp_dim": 64, - "max_position_embeddings": 64, - "relative_attention_num_buckets": 8, - "pooling": "mean", - "normalize_embeddings": True, - "max_seq_len": 32, - }, - "input_factory": "modernbert_input", - "input_factory_kwargs": {"seq_len": 16}, - "expected_output_shape": { - "embeddings": (2, 32), - "token_embeddings": (2, 16, 32), - }, - }, # ---- Text encoders (XLM-RoBERTa) ---- "XLMRobertaModel": { "module": "zeromodels.models.xlm_roberta", @@ -4629,9 +4605,8 @@ } -# Stable Diffusion 1.x: the hosted container (UNet + VAE + CLIP text, three -# disconnected paths in one functional graph) and the text-to-image task, which is -# the same graph plus the denoising loop. Tiny 2-level UNet at an 8x8 latent, a +# Diffusion: one ZMDiffusion entry per family (model_type). Stable Diffusion 1.x: +# UNet + VAE + CLIP text, three disconnected paths in one functional graph. Tiny 2-level UNet at an 8x8 latent, a # 16px two-level VAE (x2), a 1-layer text tower padded to 16 tokens. _sd_tiny = { "unet_sample_size": 8, @@ -4659,41 +4634,25 @@ "image": (2, 16, 16, 3), "text_embeds": (2, 16, 32), } -MODEL_TEST_CONFIGS["StableDiffusionModel"] = { - "module": "zeromodels.models.stable_diffusion", - "model_cls": "StableDiffusionModel", - "model_type": "diffusion", - "init_kwargs": dict(_sd_tiny), - "input_factory": "stable_diffusion_input", - "expected_output_shape": dict(_sd_outputs), -} -MODEL_TEST_CONFIGS["StableDiffusionTextToImage"] = { - "module": "zeromodels.models.stable_diffusion", - "model_cls": "StableDiffusionTextToImage", +MODEL_TEST_CONFIGS["ZMDiffusion_stable_diffusion"] = { + "module": "zeromodels.diffusers", + "model_cls": "ZMDiffusion", "model_type": "diffusion", - "init_kwargs": dict(_sd_tiny), + "init_kwargs": dict(_sd_tiny, model_type="stable_diffusion"), "input_factory": "stable_diffusion_input", "expected_output_shape": dict(_sd_outputs), } -# Stable Diffusion 2.x: the same container / task over the SD 2 configuration (per-level +# Stable Diffusion 2.x: the SD 1.x family over the SD 2 configuration (per-level # attention heads, linear token projection, gelu OpenCLIP-style text tower, "!" padding). _sd2_tiny = dict( _sd_tiny, unet_num_attention_heads=(2, 4), unet_use_linear_projection=True ) -MODEL_TEST_CONFIGS["StableDiffusion2Model"] = { - "module": "zeromodels.models.stable_diffusion_2", - "model_cls": "StableDiffusion2Model", - "model_type": "diffusion", - "init_kwargs": dict(_sd2_tiny), - "input_factory": "stable_diffusion_input", - "expected_output_shape": dict(_sd_outputs), -} -MODEL_TEST_CONFIGS["StableDiffusion2TextToImage"] = { - "module": "zeromodels.models.stable_diffusion_2", - "model_cls": "StableDiffusion2TextToImage", +MODEL_TEST_CONFIGS["ZMDiffusion_stable_diffusion_2"] = { + "module": "zeromodels.diffusers", + "model_cls": "ZMDiffusion", "model_type": "diffusion", - "init_kwargs": dict(_sd2_tiny), + "init_kwargs": dict(_sd2_tiny, model_type="stable_diffusion_2"), "input_factory": "stable_diffusion_input", "expected_output_shape": dict(_sd_outputs), } @@ -4740,19 +4699,11 @@ "prompt_embeds": (2, 16, 64), "pooled_prompt_embeds": (2, 16), } -MODEL_TEST_CONFIGS["StableDiffusionXLModel"] = { - "module": "zeromodels.models.stable_diffusion_xl", - "model_cls": "StableDiffusionXLModel", - "model_type": "diffusion", - "init_kwargs": dict(_sdxl_tiny), - "input_factory": "stable_diffusion_xl_input", - "expected_output_shape": dict(_sdxl_outputs), -} -MODEL_TEST_CONFIGS["StableDiffusionXLTextToImage"] = { - "module": "zeromodels.models.stable_diffusion_xl", - "model_cls": "StableDiffusionXLTextToImage", +MODEL_TEST_CONFIGS["ZMDiffusion_stable_diffusion_xl"] = { + "module": "zeromodels.diffusers", + "model_cls": "ZMDiffusion", "model_type": "diffusion", - "init_kwargs": dict(_sdxl_tiny), + "init_kwargs": dict(_sdxl_tiny, model_type="stable_diffusion_xl"), "input_factory": "stable_diffusion_xl_input", "expected_output_shape": dict(_sdxl_outputs), } @@ -4790,19 +4741,11 @@ "text_2_vocab_size": 128, } _sdxl_refiner_outputs = dict(_sdxl_outputs, prompt_embeds=(2, 16, 32)) -MODEL_TEST_CONFIGS["StableDiffusionXLRefinerModel"] = { - "module": "zeromodels.models.stable_diffusion_xl", - "model_cls": "StableDiffusionXLRefinerModel", +MODEL_TEST_CONFIGS["ZMDiffusion_stable_diffusion_xl_refiner"] = { + "module": "zeromodels.diffusers", + "model_cls": "ZMDiffusion", "model_type": "diffusion", - "init_kwargs": dict(_sdxl_refiner_tiny), - "input_factory": "stable_diffusion_xl_refiner_input", - "expected_output_shape": dict(_sdxl_refiner_outputs), -} -MODEL_TEST_CONFIGS["StableDiffusionXLRefinerImageToImage"] = { - "module": "zeromodels.models.stable_diffusion_xl", - "model_cls": "StableDiffusionXLRefinerImageToImage", - "model_type": "diffusion", - "init_kwargs": dict(_sdxl_refiner_tiny), + "init_kwargs": dict(_sdxl_refiner_tiny, model_type="stable_diffusion_xl_refiner"), "input_factory": "stable_diffusion_xl_refiner_input", "expected_output_shape": dict(_sdxl_refiner_outputs), } @@ -4850,25 +4793,17 @@ "prompt_embeds": (2, 16, 24), "pooled_prompt_embeds": (2, 24), } -MODEL_TEST_CONFIGS["StableDiffusion3Model"] = { - "module": "zeromodels.models.stable_diffusion_3", - "model_cls": "StableDiffusion3Model", +MODEL_TEST_CONFIGS["ZMDiffusion_stable_diffusion_3"] = { + "module": "zeromodels.diffusers", + "model_cls": "ZMDiffusion", "model_type": "diffusion", - "init_kwargs": dict(_sd3_tiny), - "input_factory": "stable_diffusion_3_input", - "expected_output_shape": dict(_sd3_outputs), -} -MODEL_TEST_CONFIGS["StableDiffusion3TextToImage"] = { - "module": "zeromodels.models.stable_diffusion_3", - "model_cls": "StableDiffusion3TextToImage", - "model_type": "diffusion", - "init_kwargs": dict(_sd3_tiny), + "init_kwargs": dict(_sd3_tiny, model_type="stable_diffusion_3"), "input_factory": "stable_diffusion_3_input", "expected_output_shape": dict(_sd3_outputs), } # the separately hosted third text encoder (T5 v1.1 XXL layout: gated-GELU blocks) MODEL_TEST_CONFIGS["SD3T5EncoderModel"] = { - "module": "zeromodels.models.stable_diffusion_3", + "module": "zeromodels.diffusers.stable_diffusion_3", "model_cls": "SD3T5EncoderModel", "model_type": "llm", "init_kwargs": { @@ -4890,19 +4825,11 @@ transformer_qk_norm="rms_norm", transformer_dual_attention_layers=(0,), ) -MODEL_TEST_CONFIGS["StableDiffusion3_5Model"] = { - "module": "zeromodels.models.stable_diffusion_3_5", - "model_cls": "StableDiffusion3_5Model", +MODEL_TEST_CONFIGS["ZMDiffusion_stable_diffusion_3_5"] = { + "module": "zeromodels.diffusers", + "model_cls": "ZMDiffusion", "model_type": "diffusion", - "init_kwargs": dict(_sd3_5_tiny), - "input_factory": "stable_diffusion_3_input", - "expected_output_shape": dict(_sd3_outputs), -} -MODEL_TEST_CONFIGS["StableDiffusion3_5TextToImage"] = { - "module": "zeromodels.models.stable_diffusion_3_5", - "model_cls": "StableDiffusion3_5TextToImage", - "model_type": "diffusion", - "init_kwargs": dict(_sd3_5_tiny), + "init_kwargs": dict(_sd3_5_tiny, model_type="stable_diffusion_3_5"), "input_factory": "stable_diffusion_3_input", "expected_output_shape": dict(_sd3_outputs), } @@ -4945,19 +4872,11 @@ "image": (2, 16, 16, 3), "prompt_embeds": (2, 32, 32), } -MODEL_TEST_CONFIGS["QwenImageModel"] = { - "module": "zeromodels.models.qwen_image", - "model_cls": "QwenImageModel", +MODEL_TEST_CONFIGS["ZMDiffusion_qwen_image"] = { + "module": "zeromodels.diffusers", + "model_cls": "ZMDiffusion", "model_type": "diffusion", - "init_kwargs": dict(_qwen_image_tiny), - "input_factory": "qwen_image_input", - "expected_output_shape": dict(_qwen_image_outputs), -} -MODEL_TEST_CONFIGS["QwenImageTextToImage"] = { - "module": "zeromodels.models.qwen_image", - "model_cls": "QwenImageTextToImage", - "model_type": "diffusion", - "init_kwargs": dict(_qwen_image_tiny), + "init_kwargs": dict(_qwen_image_tiny, model_type="qwen_image"), "input_factory": "qwen_image_input", "expected_output_shape": dict(_qwen_image_outputs), } @@ -5002,19 +4921,11 @@ "image": (2, 64, 64, 4), "prompt_embeds": (2, 16, 32), } -MODEL_TEST_CONFIGS["QwenImage21Model"] = { - "module": "zeromodels.models.qwen_image_21", - "model_cls": "QwenImage21Model", - "model_type": "diffusion", - "init_kwargs": dict(_qwen_image_21_tiny), - "input_factory": "qwen_image_21_input", - "expected_output_shape": dict(_qwen_image_21_outputs), -} -MODEL_TEST_CONFIGS["QwenImage21TextToImage"] = { - "module": "zeromodels.models.qwen_image_21", - "model_cls": "QwenImage21TextToImage", +MODEL_TEST_CONFIGS["ZMDiffusion_qwen_image_21"] = { + "module": "zeromodels.diffusers", + "model_cls": "ZMDiffusion", "model_type": "diffusion", - "init_kwargs": dict(_qwen_image_21_tiny), + "init_kwargs": dict(_qwen_image_21_tiny, model_type="qwen_image_21"), "input_factory": "qwen_image_21_input", "expected_output_shape": dict(_qwen_image_21_outputs), } diff --git a/tests/fixtures/cross_backend_parity.json b/tests/fixtures/cross_backend_parity.json index 5f98aa08..63910547 100644 --- a/tests/fixtures/cross_backend_parity.json +++ b/tests/fixtures/cross_backend_parity.json @@ -20615,5187 +20615,2246 @@ ] } ], - "QwenImage21Model": [ - { - "shape": [ - 2, - 64, - 64, - 4 - ], - "sample": [ - 0.000958, - -0.002757, - -0.002676, - -0.002686, - -0.002682, - -0.002686, - -0.002666, - -0.002711, - -0.00273, - -0.0208, - -0.020851, - -0.02082, - -0.020843, - -0.02082, - -0.020862, - -0.020826, - -0.020863, - -0.020826, - 0.015441, - 0.015452, - 0.015469, - 0.015462, - 0.015427, - 0.015388, - 0.015356, - 0.015404, - 0.015399, - 0.033911, - 0.033944, - 0.03398, - 0.033954, - 0.03374, - 0.034017, - 0.033818, - 0.033958, - 0.033937, - -0.002685, - -0.002678, - -0.002668, - -0.002742, - -0.002671, - -0.002652, - -0.00272, - -0.002738, - -0.002786, - -0.020785, - -0.020731, - -0.020771, - -0.020838, - -0.020827, - -0.020796, - -0.020871, - -0.020885, - -0.020835, - 0.015476, - 0.015424, - 0.015454, - 0.015447, - 0.015401, - 0.015431, - 0.015359, - 0.015402, - 0.015371, - 0.032569 - ] - }, + "RFDETRDetect": [ { "shape": [ 2, - 16, - 16, - 16 + 10, + 91 ], "sample": [ - 0.043742, - 0.006526, - -0.00716, - 0.002203, - 0.007001, - -0.003265, - 0.041306, - 0.048222, - 0.044051, - 0.006727, - -0.007272, - 0.002304, - 0.007318, - -0.003362, - 0.041373, - 0.048279, - 0.043725, - 0.006633, - -0.007264, - 0.00234, - 0.007361, - -0.003407, - 0.041377, - 0.048261, - 0.043969, - 0.006626, - -0.007279, - 0.002379, - 0.007225, - -0.003334, - 0.041278, - 0.048426, - 0.043925, - 0.006591, - -0.007236, - 0.002311, - 0.007166, - -0.003305, - 0.04118, - 0.048263, - 0.043971, - 0.006603, - -0.007174, - 0.002288, - 0.007199, - -0.003511, - 0.041368, - 0.048239, - 0.04406, - 0.006657, - -0.007153, - 0.002314, - 0.007224, - -0.003533, - 0.041315, - 0.048327, - 0.043871, - 0.006593, - -0.007152, - 0.002406, - 0.007264, - -0.003716, - 0.041199, - -0.03484 + -0.00643, + 0.021504, + -0.019894, + 0.008899, + -0.008518, + -0.000696, + 0.011243, + 0.015434, + -0.026406, + -0.000564, + -0.002507, + 0.007414, + 0.053563, + 0.006487, + 0.018004, + -0.001236, + 0.010106, + -0.031125, + 0.012363, + 0.009641, + -0.004421, + 0.020297, + 0.033772, + 0.014356, + 0.001726, + -0.006439, + -0.002284, + -0.011271, + -0.031966, + -0.001856, + -0.030547, + -0.006431, + -0.021222, + -0.017534, + -0.037648, + 0.058764, + 0.014301, + -0.00258, + -0.009698, + 0.019725, + 0.001287, + -0.006274, + -0.022411, + -0.027076, + -0.001648, + -0.011207, + 0.00236, + 0.029138, + 0.015287, + 0.021535, + 0.035749, + -0.020141, + 0.015779, + 0.007711, + 0.006392, + -0.001759, + -0.001101, + 0.011934, + -0.001055, + -0.004871, + -0.008768, + 0.011803, + 0.008267, + -0.023788 ] }, { "shape": [ 2, - 16, - 16 + 10, + 4 ], "sample": [ - 0.042214, - 0.105714, - 0.056851, - 0.076787, - 0.029758, - 0.00917, - 0.095684, - 0.106859, - 0.032329, - -0.067297, - 0.057183, - -0.051838, - -0.013893, - 0.035384, - -0.029907, - 0.154545, - -0.02961, - 0.08136, - 0.013182, - -0.054818, - -0.108304, - -0.122532, - -0.010875, - -0.090636, - -0.027775, - 0.009631, - -0.066816, - -0.096933, - -0.096951, - 0.185838, - -0.088463, - -0.070053, - -0.012736, - -0.047557, - 0.10501, - 0.094757, - 0.087138, - -0.094607, - -0.200976, - 0.063318, - 0.071323, - -0.137504, - 0.045576, - -0.086316, - -0.025822, - -0.020573, - 0.134116, - 0.039678, - -0.002558, - -0.075023, - 0.19758, - 0.001701, - 0.030064, - -0.005096, - -0.038978, - -0.093928, - -0.026556, - -0.030843, - -0.01305, - -0.017115, - 0.036709, - 0.005811, - 0.139922, - 0.009673 + 0.217749, + 0.218695, + 0.052857, + 0.051293, + 0.469524, + 0.051192, + 0.049849, + 0.654773, + 0.051088, + 0.047972, + 0.217339, + 0.407809, + 0.048604, + 0.717899, + 0.779579, + 0.051846, + 0.342028, + 0.780168, + 0.050525, + 0.051154, + 0.093879, + 0.05139, + 0.048411, + 0.590571, + 0.051932, + 0.048688, + 0.655617, + 0.846453, + 0.049342, + 0.844251, + 0.093517, + 0.052198, + 0.09275, + 0.781195, + 0.052856, + 0.051292, + 0.094523, + 0.051191, + 0.049848, + 0.029771, + 0.051088, + 0.047974, + 0.717339, + 0.782808, + 0.048602, + 0.030401, + 0.029579, + 0.051846, + 0.279528, + 0.967667, + 0.050525, + 0.051153, + 0.093877, + 0.051389, + 0.048407, + 0.46557, + 0.05193, + 0.048687, + 0.155618, + 0.533953, + 0.049341, + 0.594253, + 0.90602, + 0.049467 ] - }, + } + ], + "RFDETRInstanceSegment": [ { "shape": [ 2, - 16, - 32 + 10, + 91 ], "sample": [ - 0.002967, - -0.002949, - 0.002967, - -0.002949, - 0.002967, - 0.026536, - -0.02298, - 0.026536, - -0.02298, - 0.023317, - 0.011212, - 0.023317, - 0.011212, - 0.000229, - 0.004552, - 0.000229, - 0.004552, - 0.012545, - -0.008191, - 0.012545, - -0.008191, - 0.043657, - -0.018314, - 0.043657, - -0.018314, - 0.043657, - 0.022612, - 0.015819, - 0.022612, - 0.015819, - -0.001567, - -0.012379, - -0.001567, - -0.012379, - 0.008989, - 0.015044, - 0.008989, - 0.015044, - 0.005536, - -0.028382, - 0.005536, - -0.028382, - 0.01766, - 0.026689, - 0.01766, - 0.026689, - 0.01766, - -0.036786, - 0.017117, - -0.036786, - 0.017117, - -0.033625, - -0.005554, - -0.033625, - -0.005554, - -0.017666, - -0.012415, - -0.017666, - -0.012415, - 0.012932, - -0.040328, - 0.012932, - -0.040328, - -0.023516 + -0.00643, + 0.021504, + -0.019894, + 0.008899, + -0.008518, + -0.000696, + 0.011243, + 0.015434, + -0.026406, + -0.000564, + -0.002507, + 0.007414, + 0.053563, + 0.006487, + 0.018004, + -0.001236, + 0.010106, + -0.031125, + 0.012363, + 0.009641, + -0.004421, + 0.020297, + 0.033772, + 0.014357, + 0.001726, + -0.006439, + -0.002284, + -0.011271, + -0.031966, + -0.001856, + -0.030547, + -0.006431, + -0.021222, + -0.017534, + -0.037648, + 0.058764, + 0.014301, + -0.002579, + -0.009698, + 0.019725, + 0.001287, + -0.006274, + -0.022411, + -0.027076, + -0.001648, + -0.011207, + 0.002361, + 0.029138, + 0.015287, + 0.021535, + 0.035749, + -0.020141, + 0.015779, + 0.007711, + 0.006392, + -0.001759, + -0.001101, + 0.011934, + -0.001055, + -0.004871, + -0.008768, + 0.011803, + 0.008267, + -0.023788 ] - } - ], - "QwenImage21TextToImage": [ + }, { "shape": [ 2, - 64, - 64, + 10, 4 ], "sample": [ - 0.000958, - -0.002757, - -0.002676, - -0.002686, - -0.002682, - -0.002686, - -0.002666, - -0.002711, - -0.00273, - -0.0208, - -0.020851, - -0.02082, - -0.020843, - -0.02082, - -0.020862, - -0.020826, - -0.020863, - -0.020826, - 0.015441, - 0.015452, - 0.015469, - 0.015462, - 0.015427, - 0.015388, - 0.015356, - 0.015404, - 0.015399, - 0.033911, - 0.033944, - 0.03398, - 0.033954, - 0.03374, - 0.034017, - 0.033818, - 0.033958, - 0.033937, - -0.002685, - -0.002678, - -0.002668, - -0.002742, - -0.002671, - -0.002652, - -0.00272, - -0.002738, - -0.002786, - -0.020785, - -0.020731, - -0.020771, - -0.020838, - -0.020827, - -0.020796, - -0.020871, - -0.020885, - -0.020835, - 0.015476, - 0.015424, - 0.015454, - 0.015447, - 0.015401, - 0.015431, - 0.015359, - 0.015402, - 0.015371, - 0.032569 + 0.217749, + 0.218695, + 0.052857, + 0.051293, + 0.469524, + 0.051192, + 0.049849, + 0.654773, + 0.051088, + 0.047972, + 0.217339, + 0.407809, + 0.048604, + 0.717899, + 0.779579, + 0.051846, + 0.342028, + 0.780168, + 0.050525, + 0.051154, + 0.093879, + 0.05139, + 0.048411, + 0.590571, + 0.051932, + 0.048688, + 0.655617, + 0.846453, + 0.049342, + 0.844251, + 0.093517, + 0.052198, + 0.09275, + 0.781195, + 0.052856, + 0.051292, + 0.094523, + 0.051191, + 0.049848, + 0.029771, + 0.051088, + 0.047974, + 0.717339, + 0.782808, + 0.048602, + 0.030401, + 0.029579, + 0.051846, + 0.279528, + 0.967667, + 0.050525, + 0.051153, + 0.093877, + 0.051389, + 0.048407, + 0.46557, + 0.05193, + 0.048687, + 0.155618, + 0.533953, + 0.049341, + 0.594253, + 0.90602, + 0.049467 ] }, { "shape": [ 2, - 16, - 16, - 16 + 10, + 64, + 64 ], "sample": [ - 0.043742, - 0.006526, - -0.00716, - 0.002203, - 0.007001, - -0.003265, - 0.041306, - 0.048222, - 0.044051, - 0.006727, - -0.007272, - 0.002304, - 0.007318, - -0.003362, - 0.041373, - 0.048279, - 0.043725, - 0.006633, - -0.007264, - 0.00234, - 0.007361, - -0.003407, - 0.041377, - 0.048261, - 0.043969, - 0.006626, - -0.007279, - 0.002379, - 0.007225, - -0.003334, - 0.041278, - 0.048426, - 0.043925, - 0.006591, - -0.007236, - 0.002311, - 0.007166, - -0.003305, - 0.04118, - 0.048263, - 0.043971, - 0.006603, - -0.007174, - 0.002288, - 0.007199, - -0.003511, - 0.041368, - 0.048239, - 0.04406, - 0.006657, - -0.007153, - 0.002314, - 0.007224, - -0.003533, - 0.041315, - 0.048327, - 0.043871, - 0.006593, - -0.007152, - 0.002406, - 0.007264, - -0.003716, - 0.041199, - -0.03484 + -0.037459, + -0.037359, + -0.037447, + -0.036733, + -0.037721, + -0.037173, + -0.037194, + -0.037453, + -0.038214, + -0.037606, + -0.037757, + -0.037462, + -0.03753, + -0.038017, + -0.037157, + -0.037597, + -0.037569, + -0.037373, + -0.037004, + -0.037311, + -0.037633, + -0.037419, + -0.036852, + -0.037469, + -0.037402, + -0.036891, + -0.037606, + -0.037552, + -0.037339, + -0.037585, + -0.038214, + -0.037619, + -0.037505, + -0.037502, + -0.03783, + -0.037997, + -0.036738, + -0.037258, + -0.036703, + -0.037353, + -0.037827, + -0.036519, + -0.037207, + -0.037396, + -0.037232, + -0.037389, + -0.037559, + -0.037585, + -0.037921, + -0.037502, + -0.037347, + -0.037732, + -0.037177, + -0.037621, + -0.037888, + -0.037393, + -0.037585, + -0.037577, + -0.036716, + -0.037512, + -0.036471, + -0.037538, + -0.038004, + -0.037385 ] - }, + } + ], + "RTDETRDetect": [ { "shape": [ 2, - 16, - 16 + 10, + 80 ], "sample": [ - 0.042214, - 0.105714, - 0.056851, - 0.076787, - 0.029758, - 0.00917, - 0.095684, - 0.106859, - 0.032329, - -0.067297, - 0.057183, - -0.051838, - -0.013893, - 0.035384, - -0.029907, - 0.154545, - -0.02961, - 0.08136, - 0.013182, - -0.054818, - -0.108304, - -0.122532, - -0.010875, - -0.090636, - -0.027775, - 0.009631, - -0.066816, - -0.096933, - -0.096951, - 0.185838, - -0.088463, - -0.070053, - -0.012736, - -0.047557, - 0.10501, - 0.094757, - 0.087138, - -0.094607, - -0.200976, - 0.063318, - 0.071323, - -0.137504, - 0.045576, - -0.086316, - -0.025822, - -0.020573, - 0.134116, - 0.039678, - -0.002558, - -0.075023, - 0.19758, - 0.001701, - 0.030064, - -0.005096, - -0.038978, - -0.093928, - -0.026556, - -0.030843, - -0.01305, - -0.017115, - 0.036709, - 0.005811, - 0.139922, - 0.009673 + 0.04334, + -0.0334, + 0.029813, + 0.01065, + -0.003517, + -0.009205, + -0.029861, + -0.010261, + 0.0305, + -0.033181, + 0.005936, + 0.005659, + -0.011381, + -0.013536, + -0.006085, + -0.012837, + 0.00948, + 0.023415, + 0.009691, + 0.005063, + -0.000476, + -0.004171, + -0.027532, + -0.02021, + 0.035149, + 0.021277, + -0.029911, + 0.008936, + 0.020678, + -0.028142, + -0.014304, + -0.017786, + 0.015731, + 0.001191, + 0.017765, + -0.000786, + 0.010243, + -0.025679, + 0.039507, + -0.001914, + -0.00643, + 0.04333, + 0.002698, + 0.05316, + 0.010647, + 0.031321, + -0.022752, + -0.029854, + -0.028153, + 0.030504, + 0.033069, + -0.022454, + 0.005659, + -0.027778, + -0.023754, + -0.006086, + -0.031945, + 0.009491, + -0.001744, + -0.02274, + 0.005059, + -0.012732, + -0.004173, + -0.001074 ] }, { "shape": [ 2, - 16, - 32 + 10, + 4 ], "sample": [ - 0.002967, - -0.002949, - 0.002967, - -0.002949, - 0.002967, - 0.026536, - -0.02298, - 0.026536, - -0.02298, - 0.023317, - 0.011212, - 0.023317, - 0.011212, - 0.000229, - 0.004552, - 0.000229, - 0.004552, - 0.012545, - -0.008191, - 0.012545, - -0.008191, - 0.043657, - -0.018314, - 0.043657, - -0.018314, - 0.043657, - 0.022612, - 0.015819, - 0.022612, - 0.015819, - -0.001567, - -0.012379, - -0.001567, - -0.012379, - 0.008989, - 0.015044, - 0.008989, - 0.015044, - 0.005536, - -0.028382, - 0.005536, - -0.028382, - 0.01766, - 0.026689, - 0.01766, - 0.026689, - 0.01766, - -0.036786, - 0.017117, - -0.036786, - 0.017117, - -0.033625, - -0.005554, - -0.033625, - -0.005554, - -0.017666, - -0.012415, - -0.017666, - -0.012415, - 0.012932, - -0.040328, - 0.012932, - -0.040328, - -0.023516 + 0.636184, + 0.11791, + 0.197334, + 0.188859, + 0.11791, + 0.197334, + 0.188859, + 0.880158, + 0.197334, + 0.188859, + 0.880158, + 0.867548, + 0.188859, + 0.636184, + 0.359558, + 0.197334, + 0.386319, + 0.359558, + 0.197334, + 0.188859, + 0.359558, + 0.197334, + 0.188859, + 0.130347, + 0.197334, + 0.188859, + 0.130347, + 0.11791, + 0.188859, + 0.130347, + 0.359558, + 0.197334, + 0.636184, + 0.11791, + 0.197334, + 0.188859, + 0.11791, + 0.197334, + 0.188859, + 0.880158, + 0.197334, + 0.188859, + 0.880158, + 0.867548, + 0.188859, + 0.636184, + 0.359558, + 0.197334, + 0.386319, + 0.359558, + 0.197334, + 0.188859, + 0.359558, + 0.197334, + 0.188859, + 0.130347, + 0.197334, + 0.188859, + 0.130347, + 0.11791, + 0.188859, + 0.130347, + 0.359558, + 0.188859 ] } ], - "QwenImageModel": [ + "RTDETRV2Detect": [ { "shape": [ 2, - 16, - 16, - 3 + 10, + 80 ], "sample": [ - 0.003081, - 0.002639, - 0.002623, - -0.019443, - -0.020466, - -0.019435, - 0.010895, - 0.010283, - 0.010896, - 0.001854, - 0.001801, - -0.019436, - -0.019393, - -0.019437, - 0.010202, - 0.010282, - 0.010205, - 0.001856, - 0.001854, - 0.001859, - -0.019438, - -0.019436, - 0.010283, - 0.010284, - 0.010281, - 0.001854, - 0.001852, - 0.001851, - -0.01949, - -0.019487, - -0.019641, - 0.011624, - 0.01238, - 0.002639, - 0.001921, - 0.001924, - -0.019436, - -0.019431, - -0.019437, - 0.010286, - 0.010284, - 0.010288, - 0.001856, - 0.001857, - -0.019436, - -0.019428, - -0.019437, - 0.01028, - 0.010282, - 0.010282, - 0.001855, - 0.00184, - 0.001853, - -0.019504, - -0.019438, - 0.010197, - 0.010282, - 0.010199, - 0.001852, - 0.001637, - 0.001844, - -0.019552, - -0.019643, - 0.01312 + 0.039924, + -0.020892, + 0.014591, + 0.013451, + -0.016797, + -0.006289, + -0.032525, + -0.008799, + 0.030717, + -0.035124, + -0.010183, + 0.004655, + -0.003199, + -0.021864, + 0.007556, + -0.022853, + 0.010917, + 0.03682, + 0.007265, + -0.020871, + -0.009341, + 0.008063, + -0.028614, + -0.019212, + 0.041065, + 0.028744, + -0.034333, + -0.00611, + 0.00848, + -0.015636, + -0.010104, + -0.011154, + 0.028242, + -0.002491, + 0.021537, + 0.006164, + 0.014673, + -0.015604, + 0.044478, + -0.033233, + -0.00841, + 0.039918, + -0.007275, + 0.059041, + 0.013446, + 0.011633, + -0.003943, + -0.032518, + -0.02095, + 0.030726, + 0.035223, + -0.009655, + 0.004655, + -0.028753, + -0.047164, + 0.007561, + -0.029219, + 0.010919, + -0.003063, + 0.007837, + -0.020869, + -0.005034, + 0.008064, + -0.017304 ] }, { "shape": [ 2, - 16, - 16, - 8 + 10, + 4 ], "sample": [ - 0.04349, - 0.00292, - 0.00799, - -0.023711, - -0.006685, - 0.01876, - 0.003434, - -0.019355, - 0.043433, - 0.003021, - 0.007969, - -0.023718, - -0.006645, - 0.018855, - 0.003458, - -0.019362, - 0.043484, - 0.002997, - 0.008004, - -0.023699, - -0.006694, - 0.018847, - 0.003461, - -0.019371, - 0.043468, - 0.003043, - 0.008022, - -0.023695, - -0.006721, - 0.018852, - 0.003411, - -0.019382, - 0.043425, - 0.002957, - 0.008039, - -0.023618, - -0.006733, - 0.018822, - 0.003465, - -0.019322, - 0.043453, - 0.003008, - 0.008019, - -0.023752, - -0.006704, - 0.018895, - 0.003493, - -0.019344, - 0.043495, - 0.002949, - 0.007995, - -0.023701, - -0.00665, - 0.018842, - 0.003426, - -0.019335, - 0.043465, - 0.002983, - 0.008051, - -0.023634, - -0.006767, - 0.018829, - 0.003383, - -0.019354 + 0.636122, + 0.118101, + 0.197377, + 0.188947, + 0.118101, + 0.197377, + 0.188947, + 0.880129, + 0.197377, + 0.188947, + 0.880129, + 0.867759, + 0.188947, + 0.636122, + 0.35998, + 0.197377, + 0.386255, + 0.35998, + 0.197377, + 0.188947, + 0.35998, + 0.197377, + 0.188947, + 0.130316, + 0.197377, + 0.188948, + 0.130316, + 0.118101, + 0.188948, + 0.130316, + 0.35998, + 0.197377, + 0.636122, + 0.118101, + 0.197377, + 0.188947, + 0.118101, + 0.197377, + 0.188947, + 0.880129, + 0.197377, + 0.188947, + 0.880129, + 0.867759, + 0.188947, + 0.636122, + 0.35998, + 0.197377, + 0.386255, + 0.35998, + 0.197377, + 0.188947, + 0.35998, + 0.197377, + 0.188947, + 0.130316, + 0.197377, + 0.188948, + 0.130316, + 0.118101, + 0.188948, + 0.130316, + 0.35998, + 0.188948 ] - }, - { - "shape": [ - 2, - 16, - 16 - ], - "sample": [ - 0.025382, - 0.047502, - 0.021523, - 0.003751, - 0.020512, - -0.02819, - 0.07473, - 0.055641, - 0.015141, - -0.116319, - 0.040794, - -0.103205, - -0.024959, - -0.02472, - -0.037541, - 0.083859, - -0.040947, - 0.018693, - 0.035366, - -0.078992, - -0.042715, - -0.145621, - 0.029375, - -0.135195, - 0.013235, - -0.028536, - -0.030166, - -0.07373, - -0.03337, - 0.222781, - -0.015033, - -0.032118, - 0.061251, - -0.035176, - 0.139108, - 0.098706, - 0.064707, - -0.011321, - -0.187153, - 0.136417, - 0.052554, - -0.03679, - 0.032921, - -0.003105, - -0.036564, - -0.050207, - 0.065246, - 0.004824, - -0.068774, - -0.089157, - 0.068509, - -0.045322, - -0.101066, - -0.053207, - -0.043108, - -0.15871, - -0.029696, - -0.088673, - -0.017017, - -0.090879, - 0.021776, - -0.060155, - 0.129198, - -0.07387 - ] - }, + } + ], + "RegNetImageClassify": [ { "shape": [ 2, - 32, - 32 + 1000 ], "sample": [ - 0.005047, - 0.005047, - 0.005047, - -0.015945, - -0.015945, - 0.002267, - 0.002267, - 0.002829, - 0.002829, - -0.003136, - -0.003136, - 0.018783, - 0.018783, - -0.055568, - -0.055568, - 0.001576, - 0.001576, - 0.001848, - 0.001848, - 0.006598, - 0.006598, - -0.014055, - -0.014055, - 0.004142, - 0.004142, - 0.007703, - 0.007703, - -0.023382, - -0.023382, - 0.058163, - 0.058163, - -0.024143, - -0.024143, - 0.001445, - 0.001445, - 0.001554, - 0.001554, - -0.045016, - -0.045016, - 0.003507, - 0.003507, - 0.000497, - 0.000497, - -0.056717, - -0.056717, - 0.018945, - 0.018945, - 0.003564, - 0.003564, - 0.018921, - 0.018921, - 0.017499, - 0.017499, - -0.000682, - -0.000682, - 0.002824, - 0.002824, - 0.045148, - 0.045148, - 0.001521, - 0.001521, - -0.000113, - -0.000113, - 0.002808 + -0.009858, + 0.000647, + -0.004881, + 0.004978, + 0.039314, + 0.006037, + -0.028278, + 0.036571, + 0.028744, + -0.006683, + 0.013662, + 0.013965, + 0.018444, + 0.005598, + 0.004153, + 0.004897, + 0.00091, + -0.01713, + -0.037554, + -0.017242, + -0.017878, + -0.023708, + 0.029172, + -0.004541, + 0.022359, + 0.008551, + -0.019578, + -0.005603, + 0.007752, + 0.02809, + 0.004369, + -0.016323, + 0.023316, + 0.008502, + 0.030392, + -0.019346, + 0.038713, + -0.033542, + -0.039762, + -0.042688, + -0.007713, + 0.009866, + 0.020308, + 0.026513, + -0.022959, + -0.005805, + 0.019264, + -0.019209, + 0.018523, + 0.014165, + 0.003753, + 0.010536, + 0.015847, + 0.001725, + 0.006576, + 0.036405, + 0.023064, + 0.023107, + -0.012569, + 0.015555, + -0.019397, + 0.006557, + -0.034946, + 0.032827 ] } ], - "QwenImageTextToImage": [ + "Res2NetImageClassify": [ { "shape": [ 2, - 16, - 16, - 3 + 1000 ], "sample": [ - 0.003081, - 0.002639, - 0.002623, - -0.019443, - -0.020466, - -0.019435, - 0.010895, - 0.010283, - 0.010896, - 0.001854, - 0.001801, - -0.019436, - -0.019393, - -0.019437, - 0.010202, - 0.010282, - 0.010205, - 0.001856, - 0.001854, - 0.001859, - -0.019438, - -0.019436, - 0.010283, - 0.010284, - 0.010281, - 0.001854, - 0.001852, - 0.001851, - -0.01949, - -0.019487, - -0.019641, - 0.011624, - 0.01238, - 0.002639, - 0.001921, - 0.001924, - -0.019436, - -0.019431, - -0.019437, - 0.010286, - 0.010284, - 0.010288, - 0.001856, - 0.001857, - -0.019436, - -0.019428, - -0.019437, - 0.01028, - 0.010282, - 0.010282, - 0.001855, - 0.00184, - 0.001853, - -0.019504, - -0.019438, - 0.010197, - 0.010282, - 0.010199, - 0.001852, - 0.001637, - 0.001844, - -0.019552, - -0.019643, - 0.01312 + 0.033209, + 0.043605, + 0.007504, + -0.038013, + -0.005667, + 0.02878, + -0.012648, + 0.042833, + 0.007406, + -0.037209, + -0.007657, + 0.002883, + 0.038668, + 0.013185, + -0.007089, + 0.034501, + -0.000307, + 0.006438, + -0.015565, + -0.002151, + 0.049214, + 0.002254, + 0.033224, + -0.016511, + -0.027819, + 0.015059, + 0.033423, + -0.030005, + -0.007784, + 0.003483, + 0.004279, + -0.039813, + 0.036724, + -0.02716, + -0.006941, + -0.013, + -0.012659, + 0.029234, + -0.013319, + 0.003257, + 0.013238, + -0.023613, + -0.069263, + -0.041046, + 0.055614, + -0.026655, + -0.024231, + -0.007439, + -0.006986, + 0.004351, + 0.012456, + -0.003304, + -0.021892, + -0.051173, + -0.047338, + -0.009765, + 0.047473, + 0.035008, + 0.018161, + -0.030247, + -0.019736, + 0.002833, + -0.002, + -0.023537 ] - }, + } + ], + "ResMLPImageClassify": [ { "shape": [ 2, - 16, - 16, - 8 + 1000 ], "sample": [ - 0.04349, - 0.00292, - 0.00799, - -0.023711, - -0.006685, - 0.01876, - 0.003434, - -0.019355, - 0.043433, - 0.003021, - 0.007969, - -0.023718, - -0.006645, - 0.018855, - 0.003458, - -0.019362, - 0.043484, - 0.002997, - 0.008004, - -0.023699, - -0.006694, - 0.018847, - 0.003461, - -0.019371, - 0.043468, - 0.003043, - 0.008022, - -0.023695, - -0.006721, - 0.018852, - 0.003411, - -0.019382, - 0.043425, - 0.002957, - 0.008039, - -0.023618, - -0.006733, - 0.018822, - 0.003465, - -0.019322, - 0.043453, - 0.003008, - 0.008019, - -0.023752, - -0.006704, - 0.018895, - 0.003493, - -0.019344, - 0.043495, - 0.002949, - 0.007995, - -0.023701, - -0.00665, - 0.018842, - 0.003426, - -0.019335, - 0.043465, - 0.002983, - 0.008051, - -0.023634, - -0.006767, - 0.018829, - 0.003383, - -0.019354 + 0.020725, + 0.012823, + 0.003965, + -0.002472, + -0.016795, + 0.034916, + -0.028363, + 0.045582, + -0.002198, + -0.064605, + 0.013911, + -0.030648, + 0.030263, + 0.012549, + -0.016125, + 0.02806, + 0.008489, + 0.004956, + -0.005301, + -0.000166, + 0.003682, + -0.022652, + 0.031509, + 0.034854, + -0.024944, + 0.008796, + -0.003534, + -0.045192, + 0.012783, + 0.010212, + -0.007217, + -0.018173, + 0.025437, + 0.018432, + 0.015618, + -0.022768, + -0.011606, + -0.004398, + -0.012869, + 0.023885, + 0.038428, + -0.013238, + -0.045832, + -0.023672, + 0.015933, + -0.015671, + -0.005256, + -0.003881, + -0.000209, + -0.009331, + 0.042818, + 6.9e-05, + -0.022178, + -0.017999, + -0.024031, + -0.003404, + 0.032538, + 0.0165, + -0.021765, + -0.015351, + -0.009104, + -0.016671, + -0.007859, + -0.03041 ] - }, + } + ], + "ResNeXtImageClassify": [ { "shape": [ 2, - 16, - 16 + 1000 ], "sample": [ - 0.025382, - 0.047502, - 0.021523, - 0.003751, - 0.020512, - -0.02819, - 0.07473, - 0.055641, - 0.015141, - -0.116319, - 0.040794, - -0.103205, - -0.024959, - -0.02472, - -0.037541, - 0.083859, - -0.040947, - 0.018693, - 0.035366, - -0.078992, - -0.042715, - -0.145621, - 0.029375, - -0.135195, - 0.013235, - -0.028536, - -0.030166, - -0.07373, - -0.03337, - 0.222781, - -0.015033, - -0.032118, - 0.061251, - -0.035176, - 0.139108, - 0.098706, - 0.064707, - -0.011321, - -0.187153, - 0.136417, - 0.052554, - -0.03679, - 0.032921, - -0.003105, - -0.036564, - -0.050207, - 0.065246, - 0.004824, - -0.068774, - -0.089157, - 0.068509, - -0.045322, - -0.101066, - -0.053207, - -0.043108, - -0.15871, - -0.029696, - -0.088673, - -0.017017, - -0.090879, - 0.021776, - -0.060155, - 0.129198, - -0.07387 + 0.03366, + 0.027927, + 0.012954, + -0.024166, + -0.026871, + 0.045026, + -0.054242, + 0.049006, + 0.008981, + -0.045463, + -0.004104, + -0.00363, + 0.021577, + 0.022146, + -0.019124, + 0.02334, + 0.001306, + 0.036907, + 0.008641, + -0.002627, + 0.049313, + -0.019951, + 0.036502, + 0.023923, + -0.024004, + 0.017446, + 0.00852, + -0.051455, + -0.010452, + 0.004527, + -0.008951, + -0.024954, + 0.053898, + 0.003465, + -0.013708, + -0.047808, + 0.029421, + 0.009419, + -0.020658, + 0.018452, + 0.010167, + -0.029803, + -0.053733, + -0.038825, + 0.047533, + -0.05375, + -0.028515, + -0.029548, + 0.002434, + -0.034321, + -0.003224, + 0.003799, + -0.041061, + -0.060688, + -0.01578, + -0.004413, + 0.034116, + 0.038024, + 0.014058, + -0.039699, + -0.00568, + 0.001708, + 0.018758, + -0.031432 ] - }, + } + ], + "ResNetImageClassify": [ { "shape": [ 2, - 32, - 32 + 1000 ], "sample": [ - 0.005047, - 0.005047, - 0.005047, - -0.015945, - -0.015945, - 0.002267, - 0.002267, - 0.002829, - 0.002829, - -0.003136, - -0.003136, - 0.018783, - 0.018783, - -0.055568, - -0.055568, - 0.001576, - 0.001576, - 0.001848, - 0.001848, - 0.006598, - 0.006598, - -0.014055, - -0.014055, - 0.004142, - 0.004142, - 0.007703, - 0.007703, - -0.023382, - -0.023382, - 0.058163, - 0.058163, - -0.024143, - -0.024143, - 0.001445, - 0.001445, - 0.001554, - 0.001554, - -0.045016, - -0.045016, - 0.003507, - 0.003507, - 0.000497, - 0.000497, - -0.056717, - -0.056717, - 0.018945, - 0.018945, - 0.003564, - 0.003564, - 0.018921, - 0.018921, - 0.017499, - 0.017499, - -0.000682, - -0.000682, - 0.002824, - 0.002824, - 0.045148, - 0.045148, - 0.001521, - 0.001521, - -0.000113, - -0.000113, - 0.002808 + 0.033792, + 0.027935, + 0.012911, + -0.024139, + -0.026869, + 0.045084, + -0.054348, + 0.049086, + 0.008984, + -0.04558, + -0.004053, + -0.003616, + 0.021623, + 0.022099, + -0.019185, + 0.023365, + 0.001349, + 0.037089, + 0.008827, + -0.002739, + 0.049434, + -0.020054, + 0.036641, + 0.023838, + -0.024198, + 0.017397, + 0.008487, + -0.051323, + -0.010468, + 0.004666, + -0.008796, + -0.024942, + 0.053697, + 0.003274, + -0.013526, + -0.047853, + 0.029485, + 0.009487, + -0.020872, + 0.018396, + 0.010107, + -0.029635, + -0.053517, + -0.038769, + 0.047618, + -0.053648, + -0.028573, + -0.029465, + 0.002624, + -0.034464, + -0.003052, + 0.003632, + -0.04115, + -0.060654, + -0.015837, + -0.004219, + 0.034162, + 0.038158, + 0.014172, + -0.03983, + -0.005739, + 0.001556, + 0.018753, + -0.031437 ] } ], - "RFDETRDetect": [ + "ResNetV2ImageClassify": [ { "shape": [ 2, - 10, - 91 + 1000 ], "sample": [ - -0.00643, - 0.021504, - -0.019894, - 0.008899, - -0.008518, - -0.000696, - 0.011243, - 0.015434, - -0.026406, - -0.000564, - -0.002507, - 0.007414, - 0.053563, - 0.006487, - 0.018004, - -0.001236, - 0.010106, - -0.031125, - 0.012363, - 0.009641, - -0.004421, - 0.020297, - 0.033772, - 0.014356, - 0.001726, - -0.006439, - -0.002284, - -0.011271, - -0.031966, - -0.001856, - -0.030547, - -0.006431, - -0.021222, - -0.017534, - -0.037648, - 0.058764, - 0.014301, - -0.00258, - -0.009698, - 0.019725, - 0.001287, - -0.006274, - -0.022411, - -0.027076, - -0.001648, - -0.011207, - 0.00236, - 0.029138, - 0.015287, - 0.021535, - 0.035749, - -0.020141, - 0.015779, - 0.007711, - 0.006392, - -0.001759, - -0.001101, - 0.011934, - -0.001055, - -0.004871, - -0.008768, - 0.011803, - 0.008267, - -0.023788 + 0.022231, + 0.008351, + -0.031703, + -0.020745, + 0.007132, + -0.010516, + 0.001989, + 0.046573, + 0.049065, + -0.047551, + 0.000799, + -0.011843, + -0.003589, + 0.042244, + 0.019193, + 0.018882, + 0.010995, + 0.012241, + -0.013287, + -0.005206, + 0.022194, + -0.000138, + 0.003359, + 0.003543, + -0.027023, + 0.009793, + 0.009512, + 0.002733, + -0.004958, + -0.004093, + 0.042661, + 0.026909, + 0.056814, + 0.001709, + 0.005711, + -0.037617, + 0.018223, + 0.012137, + -0.048989, + 0.020044, + 0.01376, + -0.026024, + -0.031223, + -0.037329, + 0.045899, + -0.016684, + -0.008029, + -0.022839, + 0.013064, + -0.04703, + 0.030794, + 0.014406, + -0.054365, + -0.051362, + -0.047638, + -0.016485, + 0.048612, + 0.029496, + 0.008585, + -0.043712, + -0.010822, + -0.001486, + 0.003031, + -0.057003 ] - }, + } + ], + "RobertaMaskedLM": [ { "shape": [ 2, - 10, - 4 + 16, + 128 ], "sample": [ - 0.217749, - 0.218695, - 0.052857, - 0.051293, - 0.469524, - 0.051192, - 0.049849, - 0.654773, - 0.051088, - 0.047972, - 0.217339, - 0.407809, - 0.048604, - 0.717899, - 0.779579, - 0.051846, - 0.342028, - 0.780168, - 0.050525, - 0.051154, - 0.093879, - 0.05139, - 0.048411, - 0.590571, - 0.051932, - 0.048688, - 0.655617, - 0.846453, - 0.049342, - 0.844251, - 0.093517, - 0.052198, - 0.09275, - 0.781195, - 0.052856, - 0.051292, - 0.094523, - 0.051191, - 0.049848, - 0.029771, - 0.051088, - 0.047974, - 0.717339, - 0.782808, - 0.048602, - 0.030401, - 0.029579, - 0.051846, - 0.279528, - 0.967667, - 0.050525, - 0.051153, - 0.093877, - 0.051389, - 0.048407, - 0.46557, - 0.05193, - 0.048687, - 0.155618, - 0.533953, - 0.049341, - 0.594253, - 0.90602, - 0.049467 + 0.013965, + 0.019017, + -0.003735, + -0.001879, + 0.010813, + 0.002046, + -0.009915, + -0.008338, + 0.019872, + 0.000889, + 0.04867, + 0.014064, + 0.0143, + 0.004872, + -0.004134, + 0.004227, + -0.055303, + 0.023829, + -0.035749, + -0.000273, + -0.019216, + -0.004134, + -0.017732, + -0.02736, + 0.016447, + 0.010392, + -0.032707, + -0.01479, + 0.009025, + -0.004914, + -0.008511, + 0.011608, + 0.015601, + -0.016102, + 0.003131, + 0.032635, + 0.021671, + 0.00695, + 0.003794, + -0.008977, + 0.008033, + -0.002371, + 0.003972, + -0.017868, + 0.026751, + 0.03861, + 0.016864, + -0.015283, + 0.004662, + -0.002961, + -0.041212, + -0.022634, + -0.044025, + 0.001969, + 0.009821, + -0.048096, + -0.034329, + 0.002701, + 0.014121, + -0.009398, + -0.017948, + -0.010799, + 0.017219, + -0.021088 ] } ], - "RFDETRInstanceSegment": [ + "RobertaModel": [ { "shape": [ 2, - 10, - 91 + 16, + 32 ], "sample": [ - -0.00643, - 0.021504, - -0.019894, - 0.008899, - -0.008518, - -0.000696, - 0.011243, - 0.015434, - -0.026406, - -0.000564, - -0.002507, - 0.007414, - 0.053563, - 0.006487, - 0.018004, - -0.001236, - 0.010106, - -0.031125, - 0.012363, - 0.009641, - -0.004421, - 0.020297, - 0.033772, - 0.014357, - 0.001726, - -0.006439, - -0.002284, - -0.011271, - -0.031966, - -0.001856, - -0.030547, - -0.006431, - -0.021222, - -0.017534, - -0.037648, - 0.058764, - 0.014301, - -0.002579, - -0.009698, - 0.019725, - 0.001287, - -0.006274, - -0.022411, - -0.027076, - -0.001648, - -0.011207, - 0.002361, - 0.029138, - 0.015287, - 0.021535, - 0.035749, + -0.012946, + 0.015214, + -0.013106, + 0.015577, + -0.013179, + 0.069574, + 0.011613, + 0.070116, + 0.010722, + 0.024324, + 0.037582, + 0.024304, + 0.037166, + -0.035889, + -0.007053, + -0.036229, + -0.007861, + -0.073318, + -0.001794, + -0.073078, + -0.001493, + 0.031959, + -0.024997, + 0.033061, + -0.024845, + 0.032108, + 0.002208, + 0.016397, + 0.000681, + 0.016726, + 0.000822, + -0.005828, + 0.00068, + -0.00826, + 0.108422, + -0.049974, + 0.105383, + -0.04979, + -0.046333, + 0.025502, + -0.046555, + 0.026064, + 0.018798, + -0.003703, + 0.019448, + -0.00295, + 0.016, + 0.004491, + -0.02026, + 0.004788, -0.020141, - 0.015779, - 0.007711, - 0.006392, - -0.001759, - -0.001101, - 0.011934, - -0.001055, - -0.004871, - -0.008768, - 0.011803, - 0.008267, - -0.023788 + 0.024671, + -0.017721, + 0.024861, + -0.017792, + 0.02249, + -0.024283, + 0.022482, + -0.024532, + -0.107456, + 0.000908, + -0.110674, + 0.000475, + -0.016094 ] }, { "shape": [ 2, - 10, - 4 + 32 ], "sample": [ - 0.217749, - 0.218695, - 0.052857, - 0.051293, - 0.469524, - 0.051192, - 0.049849, - 0.654773, - 0.051088, - 0.047972, - 0.217339, - 0.407809, - 0.048604, - 0.717899, - 0.779579, - 0.051846, - 0.342028, - 0.780168, - 0.050525, - 0.051154, - 0.093879, - 0.05139, - 0.048411, - 0.590571, - 0.051932, - 0.048688, - 0.655617, - 0.846453, - 0.049342, - 0.844251, - 0.093517, - 0.052198, - 0.09275, - 0.781195, - 0.052856, - 0.051292, - 0.094523, - 0.051191, - 0.049848, - 0.029771, - 0.051088, - 0.047974, - 0.717339, - 0.782808, - 0.048602, - 0.030401, - 0.029579, - 0.051846, - 0.279528, - 0.967667, - 0.050525, - 0.051153, - 0.093877, - 0.051389, - 0.048407, - 0.46557, - 0.05193, - 0.048687, - 0.155618, - 0.533953, - 0.049341, - 0.594253, - 0.90602, - 0.049467 + 9.5e-05, + 0.011957, + 0.006713, + 0.004725, + 0.011335, + -0.006035, + 0.012561, + 0.002068, + -0.0103, + 0.035051, + -0.010026, + 0.006029, + 0.002999, + 0.004093, + -0.02845, + 0.024154, + 0.007347, + 0.007143, + 0.001298, + -0.013618, + -0.011444, + -0.030955, + -0.005628, + -0.02546, + 0.018342, + 0.023768, + -0.011254, + -0.004289, + -0.019031, + 0.040206, + -0.025693, + -0.011637, + 8e-05, + 0.011829, + 0.006999, + 0.004674, + 0.011376, + -0.006081, + 0.012517, + 0.001985, + -0.010364, + 0.034915, + -0.010216, + 0.006072, + 0.003048, + 0.004378, + -0.028556, + 0.02404, + 0.007438, + 0.007219, + 0.001509, + -0.013589, + -0.011347, + -0.031104, + -0.00554, + -0.025427, + 0.018239, + 0.023746, + -0.01139, + -0.004372, + -0.018953, + 0.040002, + -0.025732, + -0.011578 ] - }, - { + } + ], + "RobertaMultipleChoice": [ + { "shape": [ 2, - 10, - 64, - 64 + 3 ], "sample": [ - -0.037459, - -0.037359, - -0.037447, - -0.036733, - -0.037721, - -0.037173, - -0.037194, - -0.037453, - -0.038214, - -0.037606, - -0.037757, - -0.037462, - -0.03753, - -0.038017, - -0.037157, - -0.037597, - -0.037569, - -0.037373, - -0.037004, - -0.037311, - -0.037633, - -0.037419, - -0.036852, - -0.037469, - -0.037402, - -0.036891, - -0.037606, - -0.037552, - -0.037339, - -0.037585, - -0.038214, - -0.037619, - -0.037505, - -0.037502, - -0.03783, - -0.037997, - -0.036738, - -0.037258, - -0.036703, - -0.037353, - -0.037827, - -0.036519, - -0.037207, - -0.037396, - -0.037232, - -0.037389, - -0.037559, - -0.037585, - -0.037921, - -0.037502, - -0.037347, - -0.037732, - -0.037177, - -0.037621, - -0.037888, - -0.037393, - -0.037585, - -0.037577, - -0.036716, - -0.037512, - -0.036471, - -0.037538, - -0.038004, - -0.037385 + -0.015686, + -0.015691, + -0.01569, + -0.015684, + -0.015687, + -0.015699 ] } ], - "RTDETRDetect": [ + "RobertaQnA": [ { "shape": [ 2, - 10, - 80 + 16 ], "sample": [ - 0.04334, - -0.0334, - 0.029813, - 0.01065, - -0.003517, - -0.009205, - -0.029861, - -0.010261, - 0.0305, - -0.033181, - 0.005936, - 0.005659, - -0.011381, - -0.013536, - -0.006085, - -0.012837, - 0.00948, - 0.023415, - 0.009691, - 0.005063, - -0.000476, - -0.004171, - -0.027532, - -0.02021, - 0.035149, - 0.021277, - -0.029911, - 0.008936, - 0.020678, - -0.028142, - -0.014304, - -0.017786, - 0.015731, - 0.001191, - 0.017765, - -0.000786, - 0.010243, - -0.025679, - 0.039507, - -0.001914, - -0.00643, - 0.04333, - 0.002698, - 0.05316, - 0.010647, - 0.031321, - -0.022752, - -0.029854, - -0.028153, - 0.030504, - 0.033069, - -0.022454, - 0.005659, - -0.027778, - -0.023754, - -0.006086, - -0.031945, - 0.009491, - -0.001744, - -0.02274, - 0.005059, - -0.012732, - -0.004173, - -0.001074 + -0.040833, + -0.040412, + -0.040958, + -0.041051, + -0.040626, + -0.041051, + -0.040365, + -0.04072, + -0.040683, + -0.041041, + -0.040674, + -0.041258, + -0.040843, + -0.040957, + -0.040952, + -0.040941, + -0.040662, + -0.04115, + -0.040876, + -0.040881, + -0.040948, + -0.040716, + -0.041096, + -0.040604, + -0.041002, + -0.040509, + -0.040938, + -0.040975, + -0.041082, + -0.040653, + -0.040324, + -0.041059 ] }, { "shape": [ 2, - 10, + 16 + ], + "sample": [ + 0.019887, + 0.020086, + 0.019876, + 0.020019, + 0.020148, + 0.020007, + 0.020131, + 0.019951, + 0.020173, + 0.020156, + 0.019967, + 0.019955, + 0.020004, + 0.019911, + 0.019939, + 0.020037, + 0.020038, + 0.019925, + 0.020181, + 0.02007, + 0.020088, + 0.020004, + 0.019789, + 0.020143, + 0.020097, + 0.019977, + 0.020205, + 0.019934, + 0.02004, + 0.019997, + 0.020336, + 0.01988 + ] + } + ], + "RobertaSequenceClassify": [ + { + "shape": [ + 2, 4 ], "sample": [ - 0.636184, - 0.11791, - 0.197334, - 0.188859, - 0.11791, - 0.197334, - 0.188859, - 0.880158, - 0.197334, - 0.188859, - 0.880158, - 0.867548, - 0.188859, - 0.636184, - 0.359558, - 0.197334, - 0.386319, - 0.359558, - 0.197334, - 0.188859, - 0.359558, - 0.197334, - 0.188859, - 0.130347, - 0.197334, - 0.188859, - 0.130347, - 0.11791, - 0.188859, - 0.130347, - 0.359558, - 0.197334, - 0.636184, - 0.11791, - 0.197334, - 0.188859, - 0.11791, - 0.197334, - 0.188859, - 0.880158, - 0.197334, - 0.188859, - 0.880158, - 0.867548, - 0.188859, - 0.636184, - 0.359558, - 0.197334, - 0.386319, - 0.359558, - 0.197334, - 0.188859, - 0.359558, - 0.197334, - 0.188859, - 0.130347, - 0.197334, - 0.188859, - 0.130347, - 0.11791, - 0.188859, - 0.130347, - 0.359558, - 0.188859 + 0.000177, + 0.014164, + -0.036711, + -0.005497, + 0.000159, + 0.01417, + -0.036695, + -0.005501 ] } ], - "RTDETRV2Detect": [ - { - "shape": [ - 2, - 10, - 80 - ], - "sample": [ - 0.039924, - -0.020892, - 0.014591, - 0.013451, - -0.016797, - -0.006289, - -0.032525, - -0.008799, - 0.030717, - -0.035124, - -0.010183, - 0.004655, - -0.003199, - -0.021864, - 0.007556, - -0.022853, - 0.010917, - 0.03682, - 0.007265, - -0.020871, - -0.009341, - 0.008063, - -0.028614, - -0.019212, - 0.041065, - 0.028744, - -0.034333, - -0.00611, - 0.00848, - -0.015636, - -0.010104, - -0.011154, - 0.028242, - -0.002491, - 0.021537, - 0.006164, - 0.014673, - -0.015604, - 0.044478, - -0.033233, - -0.00841, - 0.039918, - -0.007275, - 0.059041, - 0.013446, - 0.011633, - -0.003943, - -0.032518, - -0.02095, - 0.030726, - 0.035223, - -0.009655, - 0.004655, - -0.028753, - -0.047164, - 0.007561, - -0.029219, - 0.010919, - -0.003063, - 0.007837, - -0.020869, - -0.005034, - 0.008064, - -0.017304 - ] - }, + "RobertaTokenClassify": [ { "shape": [ 2, - 10, - 4 + 16, + 5 ], "sample": [ - 0.636122, - 0.118101, - 0.197377, - 0.188947, - 0.118101, - 0.197377, - 0.188947, - 0.880129, - 0.197377, - 0.188947, - 0.880129, - 0.867759, - 0.188947, - 0.636122, - 0.35998, - 0.197377, - 0.386255, - 0.35998, - 0.197377, - 0.188947, - 0.35998, - 0.197377, - 0.188947, - 0.130316, - 0.197377, - 0.188948, - 0.130316, - 0.118101, - 0.188948, - 0.130316, - 0.35998, - 0.197377, - 0.636122, - 0.118101, - 0.197377, - 0.188947, - 0.118101, - 0.197377, - 0.188947, - 0.880129, - 0.197377, - 0.188947, - 0.880129, - 0.867759, - 0.188947, - 0.636122, - 0.35998, - 0.197377, - 0.386255, - 0.35998, - 0.197377, - 0.188947, - 0.35998, - 0.197377, - 0.188947, - 0.130316, - 0.197377, - 0.188948, - 0.130316, - 0.118101, - 0.188948, - 0.130316, - 0.35998, - 0.188948 + -0.011828, + -0.052942, + -0.012009, + -0.052773, + -0.012009, + -0.05311, + -0.011534, + -0.052935, + -0.011814, + -0.052717, + -0.011641, + -0.052989, + -0.011957, + -0.052632, + -0.011815, + -0.052956, + -0.011606, + -0.052607, + -0.011437, + -0.052679, + -0.011876, + 0.026477, + -0.011723, + 0.026765, + -0.011958, + 0.026455, + -0.011909, + 0.026545, + -0.011798, + 0.026639, + -0.011807, + 0.026565, + -0.011974, + 0.026476, + -0.011909, + 0.02662, + -0.011605, + 0.026641, + -0.011655, + 0.026516, + -0.011716, + 0.026667, + 0.000539, + 0.026477, + 0.000706, + 0.02649, + 0.000529, + 0.026434, + 0.000589, + 0.026668, + 0.000565, + 0.026388, + 0.000571, + 0.026631, + 0.000628, + 0.026565, + 0.00061, + 0.026627, + 0.000546, + 0.026415, + 0.000415, + 0.026348, + 0.000654, + -0.05073 ] } ], - "RegNetImageClassify": [ + "SAM2Model": [ { "shape": [ 2, - 1000 + 16, + 16, + 256 ], "sample": [ - -0.009858, - 0.000647, - -0.004881, - 0.004978, - 0.039314, - 0.006037, - -0.028278, - 0.036571, - 0.028744, - -0.006683, - 0.013662, - 0.013965, - 0.018444, - 0.005598, - 0.004153, - 0.004897, - 0.00091, - -0.01713, - -0.037554, - -0.017242, - -0.017878, - -0.023708, - 0.029172, - -0.004541, - 0.022359, - 0.008551, - -0.019578, - -0.005603, - 0.007752, - 0.02809, - 0.004369, - -0.016323, - 0.023316, - 0.008502, - 0.030392, - -0.019346, - 0.038713, - -0.033542, - -0.039762, - -0.042688, - -0.007713, - 0.009866, - 0.020308, - 0.026513, - -0.022959, - -0.005805, - 0.019264, - -0.019209, - 0.018523, - 0.014165, - 0.003753, - 0.010536, - 0.015847, - 0.001725, - 0.006576, - 0.036405, - 0.023064, - 0.023107, - -0.012569, - 0.015555, - -0.019397, - 0.006557, - -0.034946, - 0.032827 - ] - } - ], - "Res2NetImageClassify": [ - { - "shape": [ - 2, - 1000 - ], - "sample": [ - 0.033209, - 0.043605, - 0.007504, - -0.038013, - -0.005667, - 0.02878, - -0.012648, - 0.042833, - 0.007406, - -0.037209, - -0.007657, - 0.002883, - 0.038668, - 0.013185, - -0.007089, - 0.034501, - -0.000307, - 0.006438, - -0.015565, - -0.002151, - 0.049214, - 0.002254, - 0.033224, - -0.016511, - -0.027819, - 0.015059, - 0.033423, - -0.030005, - -0.007784, - 0.003483, - 0.004279, - -0.039813, - 0.036724, - -0.02716, - -0.006941, - -0.013, - -0.012659, - 0.029234, - -0.013319, - 0.003257, - 0.013238, - -0.023613, - -0.069263, - -0.041046, - 0.055614, - -0.026655, - -0.024231, - -0.007439, - -0.006986, - 0.004351, - 0.012456, - -0.003304, - -0.021892, - -0.051173, - -0.047338, - -0.009765, - 0.047473, - 0.035008, - 0.018161, - -0.030247, - -0.019736, - 0.002833, - -0.002, - -0.023537 - ] - } - ], - "ResMLPImageClassify": [ - { - "shape": [ - 2, - 1000 - ], - "sample": [ - 0.020725, - 0.012823, - 0.003965, - -0.002472, - -0.016795, - 0.034916, - -0.028363, - 0.045582, - -0.002198, - -0.064605, - 0.013911, - -0.030648, - 0.030263, - 0.012549, - -0.016125, - 0.02806, - 0.008489, - 0.004956, - -0.005301, - -0.000166, - 0.003682, - -0.022652, - 0.031509, - 0.034854, - -0.024944, - 0.008796, - -0.003534, - -0.045192, - 0.012783, - 0.010212, - -0.007217, - -0.018173, - 0.025437, - 0.018432, - 0.015618, - -0.022768, - -0.011606, - -0.004398, - -0.012869, - 0.023885, - 0.038428, - -0.013238, - -0.045832, - -0.023672, - 0.015933, - -0.015671, - -0.005256, - -0.003881, - -0.000209, - -0.009331, - 0.042818, - 6.9e-05, - -0.022178, - -0.017999, - -0.024031, - -0.003404, - 0.032538, - 0.0165, - -0.021765, - -0.015351, - -0.009104, - -0.016671, - -0.007859, - -0.03041 - ] - } - ], - "ResNeXtImageClassify": [ - { - "shape": [ - 2, - 1000 - ], - "sample": [ - 0.03366, - 0.027927, - 0.012954, - -0.024166, - -0.026871, - 0.045026, - -0.054242, - 0.049006, - 0.008981, - -0.045463, - -0.004104, - -0.00363, - 0.021577, - 0.022146, - -0.019124, - 0.02334, - 0.001306, - 0.036907, - 0.008641, - -0.002627, - 0.049313, - -0.019951, - 0.036502, - 0.023923, - -0.024004, - 0.017446, - 0.00852, - -0.051455, - -0.010452, - 0.004527, - -0.008951, - -0.024954, - 0.053898, - 0.003465, - -0.013708, - -0.047808, - 0.029421, - 0.009419, - -0.020658, - 0.018452, - 0.010167, - -0.029803, - -0.053733, - -0.038825, - 0.047533, - -0.05375, - -0.028515, - -0.029548, - 0.002434, - -0.034321, - -0.003224, - 0.003799, - -0.041061, - -0.060688, - -0.01578, - -0.004413, - 0.034116, - 0.038024, - 0.014058, - -0.039699, - -0.00568, - 0.001708, - 0.018758, - -0.031432 - ] - } - ], - "ResNetImageClassify": [ - { - "shape": [ - 2, - 1000 - ], - "sample": [ - 0.033792, - 0.027935, - 0.012911, - -0.024139, - -0.026869, - 0.045084, - -0.054348, - 0.049086, - 0.008984, - -0.04558, - -0.004053, - -0.003616, - 0.021623, - 0.022099, - -0.019185, - 0.023365, - 0.001349, - 0.037089, - 0.008827, - -0.002739, - 0.049434, - -0.020054, - 0.036641, - 0.023838, - -0.024198, - 0.017397, - 0.008487, - -0.051323, - -0.010468, - 0.004666, - -0.008796, - -0.024942, - 0.053697, - 0.003274, - -0.013526, - -0.047853, - 0.029485, - 0.009487, - -0.020872, - 0.018396, - 0.010107, - -0.029635, - -0.053517, - -0.038769, - 0.047618, - -0.053648, - -0.028573, - -0.029465, - 0.002624, - -0.034464, - -0.003052, - 0.003632, - -0.04115, - -0.060654, - -0.015837, - -0.004219, - 0.034162, - 0.038158, - 0.014172, - -0.03983, - -0.005739, - 0.001556, - 0.018753, - -0.031437 - ] - } - ], - "ResNetV2ImageClassify": [ - { - "shape": [ - 2, - 1000 - ], - "sample": [ - 0.022231, - 0.008351, - -0.031703, - -0.020745, - 0.007132, - -0.010516, - 0.001989, - 0.046573, - 0.049065, - -0.047551, - 0.000799, - -0.011843, - -0.003589, - 0.042244, - 0.019193, - 0.018882, - 0.010995, - 0.012241, - -0.013287, - -0.005206, - 0.022194, - -0.000138, - 0.003359, - 0.003543, - -0.027023, - 0.009793, - 0.009512, - 0.002733, - -0.004958, - -0.004093, - 0.042661, - 0.026909, - 0.056814, - 0.001709, - 0.005711, - -0.037617, - 0.018223, - 0.012137, - -0.048989, - 0.020044, - 0.01376, - -0.026024, - -0.031223, - -0.037329, - 0.045899, - -0.016684, - -0.008029, - -0.022839, - 0.013064, - -0.04703, - 0.030794, - 0.014406, - -0.054365, - -0.051362, - -0.047638, - -0.016485, - 0.048612, - 0.029496, - 0.008585, - -0.043712, - -0.010822, - -0.001486, - 0.003031, - -0.057003 - ] - } - ], - "RobertaMaskedLM": [ - { - "shape": [ - 2, - 16, - 128 - ], - "sample": [ - 0.013965, - 0.019017, - -0.003735, - -0.001879, - 0.010813, - 0.002046, - -0.009915, - -0.008338, - 0.019872, - 0.000889, - 0.04867, - 0.014064, - 0.0143, - 0.004872, - -0.004134, - 0.004227, - -0.055303, - 0.023829, - -0.035749, - -0.000273, - -0.019216, - -0.004134, - -0.017732, - -0.02736, - 0.016447, - 0.010392, - -0.032707, - -0.01479, - 0.009025, - -0.004914, - -0.008511, - 0.011608, - 0.015601, - -0.016102, - 0.003131, - 0.032635, - 0.021671, - 0.00695, - 0.003794, - -0.008977, - 0.008033, - -0.002371, - 0.003972, - -0.017868, - 0.026751, - 0.03861, - 0.016864, - -0.015283, - 0.004662, - -0.002961, - -0.041212, - -0.022634, - -0.044025, - 0.001969, - 0.009821, - -0.048096, - -0.034329, - 0.002701, - 0.014121, - -0.009398, - -0.017948, - -0.010799, - 0.017219, - -0.021088 - ] - } - ], - "RobertaModel": [ - { - "shape": [ - 2, - 16, - 32 - ], - "sample": [ - -0.012946, - 0.015214, - -0.013106, - 0.015577, - -0.013179, - 0.069574, - 0.011613, - 0.070116, - 0.010722, - 0.024324, - 0.037582, - 0.024304, - 0.037166, - -0.035889, - -0.007053, - -0.036229, - -0.007861, - -0.073318, - -0.001794, - -0.073078, - -0.001493, - 0.031959, - -0.024997, - 0.033061, - -0.024845, - 0.032108, - 0.002208, - 0.016397, - 0.000681, - 0.016726, - 0.000822, - -0.005828, - 0.00068, - -0.00826, - 0.108422, - -0.049974, - 0.105383, - -0.04979, - -0.046333, - 0.025502, - -0.046555, - 0.026064, - 0.018798, - -0.003703, - 0.019448, - -0.00295, - 0.016, - 0.004491, - -0.02026, - 0.004788, - -0.020141, - 0.024671, - -0.017721, - 0.024861, - -0.017792, - 0.02249, - -0.024283, - 0.022482, - -0.024532, - -0.107456, - 0.000908, - -0.110674, - 0.000475, - -0.016094 - ] - }, - { - "shape": [ - 2, - 32 - ], - "sample": [ - 9.5e-05, - 0.011957, - 0.006713, - 0.004725, - 0.011335, - -0.006035, - 0.012561, - 0.002068, - -0.0103, - 0.035051, - -0.010026, - 0.006029, - 0.002999, - 0.004093, - -0.02845, - 0.024154, - 0.007347, - 0.007143, - 0.001298, - -0.013618, - -0.011444, - -0.030955, - -0.005628, - -0.02546, - 0.018342, - 0.023768, - -0.011254, - -0.004289, - -0.019031, - 0.040206, - -0.025693, - -0.011637, - 8e-05, - 0.011829, - 0.006999, - 0.004674, - 0.011376, - -0.006081, - 0.012517, - 0.001985, - -0.010364, - 0.034915, - -0.010216, - 0.006072, - 0.003048, - 0.004378, - -0.028556, - 0.02404, - 0.007438, - 0.007219, - 0.001509, - -0.013589, - -0.011347, - -0.031104, - -0.00554, - -0.025427, - 0.018239, - 0.023746, - -0.01139, - -0.004372, - -0.018953, - 0.040002, - -0.025732, - -0.011578 - ] - } - ], - "RobertaMultipleChoice": [ - { - "shape": [ - 2, - 3 - ], - "sample": [ - -0.015686, - -0.015691, - -0.01569, - -0.015684, - -0.015687, - -0.015699 - ] - } - ], - "RobertaQnA": [ - { - "shape": [ - 2, - 16 - ], - "sample": [ - -0.040833, - -0.040412, - -0.040958, - -0.041051, - -0.040626, - -0.041051, - -0.040365, - -0.04072, - -0.040683, - -0.041041, - -0.040674, - -0.041258, - -0.040843, - -0.040957, - -0.040952, - -0.040941, - -0.040662, - -0.04115, - -0.040876, - -0.040881, - -0.040948, - -0.040716, - -0.041096, - -0.040604, - -0.041002, - -0.040509, - -0.040938, - -0.040975, - -0.041082, - -0.040653, - -0.040324, - -0.041059 - ] - }, - { - "shape": [ - 2, - 16 - ], - "sample": [ - 0.019887, - 0.020086, - 0.019876, - 0.020019, - 0.020148, - 0.020007, - 0.020131, - 0.019951, - 0.020173, - 0.020156, - 0.019967, - 0.019955, - 0.020004, - 0.019911, - 0.019939, - 0.020037, - 0.020038, - 0.019925, - 0.020181, - 0.02007, - 0.020088, - 0.020004, - 0.019789, - 0.020143, - 0.020097, - 0.019977, - 0.020205, - 0.019934, - 0.02004, - 0.019997, - 0.020336, - 0.01988 - ] - } - ], - "RobertaSequenceClassify": [ - { - "shape": [ - 2, - 4 - ], - "sample": [ - 0.000177, - 0.014164, - -0.036711, - -0.005497, - 0.000159, - 0.01417, - -0.036695, - -0.005501 - ] - } - ], - "RobertaTokenClassify": [ - { - "shape": [ - 2, - 16, - 5 - ], - "sample": [ - -0.011828, - -0.052942, - -0.012009, - -0.052773, - -0.012009, - -0.05311, - -0.011534, - -0.052935, - -0.011814, - -0.052717, - -0.011641, - -0.052989, - -0.011957, - -0.052632, - -0.011815, - -0.052956, - -0.011606, - -0.052607, - -0.011437, - -0.052679, - -0.011876, - 0.026477, - -0.011723, - 0.026765, - -0.011958, - 0.026455, - -0.011909, - 0.026545, - -0.011798, - 0.026639, - -0.011807, - 0.026565, - -0.011974, - 0.026476, - -0.011909, - 0.02662, - -0.011605, - 0.026641, - -0.011655, - 0.026516, - -0.011716, - 0.026667, - 0.000539, - 0.026477, - 0.000706, - 0.02649, - 0.000529, - 0.026434, - 0.000589, - 0.026668, - 0.000565, - 0.026388, - 0.000571, - 0.026631, - 0.000628, - 0.026565, - 0.00061, - 0.026627, - 0.000546, - 0.026415, - 0.000415, - 0.026348, - 0.000654, - -0.05073 - ] - } - ], - "SAM2Model": [ - { - "shape": [ - 2, - 16, - 16, - 256 - ], - "sample": [ - -0.015336, - -0.013034, - 0.006987, - 0.020153, - 0.049287, - 0.01717, - 0.044167, - -0.03045, - 0.032655, - 0.059362, - -0.011628, - 0.017312, - -0.02234, - -0.00998, - 0.018686, - 0.041445, - -0.02269, - 0.018144, - -0.032625, - -0.017222, - -0.009087, - 0.017343, - 0.014006, - -0.00969, - 0.003306, - 0.026315, - 0.005598, - -0.020814, - -0.035375, - -0.007683, - -0.000294, - 0.014428, - -0.048157, - 0.006775, - 0.045402, - 0.047048, - 0.008475, - 0.020926, - 0.001399, - -0.014787, - -0.0152, - -0.013629, - -0.019056, - -0.027307, - -0.047092, - -0.022155, - -0.003388, - 0.033364, - 0.024901, - 0.012694, - -0.007007, - 0.037517, - 0.002602, - 0.016988, - -0.031581, - 0.02448, - 0.031745, - 0.007232, - 0.002769, - 0.023481, - -0.01057, - 0.049235, - -0.043309, - -0.016593 - ] - }, - { - "shape": [ - 2, - 8, - 8, - 256 - ], - "sample": [ - 0.001858, - 0.005013, - -0.03215, - -0.023843, - 0.023074, - -0.009762, - 0.007063, - 0.007979, - -0.023957, - 0.041274, - -0.039285, - -0.017954, - 0.020519, - -0.03382, - -0.00841, - -0.00636, - 0.027404, - 0.019387, - -0.024081, - -0.010251, - -0.031152, - -0.001085, - -0.01668, - 0.017862, - -0.004157, - 0.002491, - 0.001924, - -0.004861, - -0.002372, - 0.03471, - -0.029995, - -0.01142, - 0.001389, - -0.026705, - 0.003308, - 0.010559, - -0.002314, - -0.00535, - -0.001025, - 0.028143, - -0.019863, - 0.015856, - -0.014986, - -0.011457, - -0.017505, - 0.020718, - -0.002761, - -0.027695, - -0.00788, - -0.002352, - 0.036122, - -0.014375, - -0.022155, - -0.022174, - 0.008941, - 0.003909, - 0.018927, - -0.014251, - -0.035766, - -0.011259, - -0.011355, - -0.017476, - 0.009668, - -0.038129 - ] - }, - { - "shape": [ - 2, - 4, - 4, - 256 - ], - "sample": [ - 0.030796, - 0.040807, - -0.004929, - -0.010892, - -0.019576, - -0.057813, - -0.069933, - -0.025368, - -0.008011, - 0.019821, - 0.021221, - 0.004241, - 0.02296, - 0.01008, - -0.018551, - -0.03447, - -0.003466, - -0.005969, - -0.062868, - 0.043891, - 0.066126, - -0.002958, - 0.044128, - -0.038288, - -0.038014, - 0.008986, - 0.027282, - -0.036212, - 0.068866, - -0.036626, - 0.007519, - 0.009738, - -0.035756, - 0.026206, - -0.000522, - -0.027787, - -0.042496, - -0.021335, - 0.032815, - -0.064257, - 0.009692, - 0.019132, - 0.006331, - -0.026821, - 0.017907, - 0.025561, - -0.034405, - -0.039337, - -0.054798, - -0.01864, - -0.035654, - -0.042077, - 0.030411, - 0.030888, - 0.021825, - -0.009091, - -0.063324, - 0.044919, - 0.035721, - 0.014497, - 0.016162, - -0.019659, - 0.017876, - -0.025821 - ] - } - ], - "SAM2PromptableSegment": [ - { - "shape": [ - 2, - 1, - 3 - ], - "sample": [ - 0.498468, - 0.503674, - 0.500753, - 0.498468, - 0.503674, - 0.500753 - ] - }, - { - "shape": [ - 2, - 1, - 1 - ], - "sample": [ - -0.00015, - -0.00015 - ] - }, - { - "shape": [ - 2, - 1, - 3, - 16, - 16 - ], - "sample": [ - 0.001723, - 0.001358, - 0.001621, - 0.001894, - 0.001981, - 0.001704, - 0.001766, - 0.00185, - 0.002127, - 0.002068, - 0.001415, - -3.9e-05, - 0.000403, - 0.000159, - -3.1e-05, - -0.00026, - -0.000519, - 0.000495, - 0.000265, - -0.000139, - 1.4e-05, - 0.000188, - -0.001361, - -0.001924, - -0.0008, - -0.001101, - -0.00219, - -0.001732, - -0.00103, - -0.001181, - -0.001395, - -0.001796, - 0.001847, - 0.001937, - 0.001716, - 0.002225, - 0.001588, - 0.001684, - 0.001877, - 0.002109, - 0.001345, - 0.002057, - 0.002675, - 0.000631, - 0.000172, - -9.5e-05, - 0.000222, - -2.5e-05, - 8.1e-05, - -0.000149, - -0.00035, - 0.000488, - 6.9e-05, - -0.001721, - -0.001244, - -0.001717, - -0.000922, - -0.001359, - -0.001446, - -0.001118, - -0.002134, - -0.001088, - -0.001499, - -0.00176 - ] - } - ], - "SAMModel": [ - { - "shape": [ - 2, - 4, - 4, - 256 - ], - "sample": [ - -0.025676, - -0.004309, - 0.0244, - 0.001673, - 0.030505, - -0.0386, - 0.01161, - 0.005268, - 0.010626, - 0.018627, - -0.025338, - -0.016413, - 0.025446, - -0.015198, - -0.008434, - -0.008578, - -0.003512, - 0.033251, - -0.070632, - -0.005308, - 0.000303, - 0.00923, - 0.026017, - 0.026024, - -0.019712, - -0.014215, - 0.0426, - -0.04387, - -0.028514, - 0.042532, - 0.02437, - 0.080674, - -0.039657, - 0.029955, - 0.05102, - -0.00787, - 0.007323, - -0.013503, - -0.027775, - 0.013422, - -0.014877, - 0.008596, - -0.027227, - -0.028858, - 0.013134, - -0.026335, - 0.023304, - -0.028756, - -0.027388, - 0.016623, - -0.037225, - -0.014127, - 0.010916, - 0.017536, - 0.015706, - 0.013338, - 0.02569, - 0.002734, - 0.061468, - 0.036984, - -0.011109, - 0.009696, - 0.073214, - -0.006962 - ] - } - ], - "SAMPromptableSegment": [ - { - "shape": [ - 2, - 1, - 3 - ], - "sample": [ - -0.026417, - 0.00425, - 0.015243, - -0.026417, - 0.00425, - 0.015242 - ] - }, - { - "shape": [ - 2, - 1, - 3, - 16, - 16 - ], - "sample": [ - -0.001283, - -0.001259, - -0.001282, - -0.001481, - -0.001512, - -0.001624, - -0.001298, - -0.001273, - -0.001249, - -0.001602, - -0.001545, - 0.000605, - 0.000614, - 0.000463, - 0.000507, - 0.000717, - 0.000802, - 0.000456, - 0.000451, - 0.000556, - 0.000791, - 0.000508, - 0.000305, - 0.000347, - 0.000117, - 0.000161, - 5.2e-05, - 2.2e-05, - 6e-05, - 9.3e-05, - 0.000286, - -1e-06, - -0.001527, - -0.001283, - -0.00129, - -0.001262, - -0.001508, - -0.001483, - -0.001583, - -0.001311, - -0.001245, - -0.001255, - -0.001556, - 0.000541, - 0.000472, - 0.000614, - 0.000621, - 0.000521, - 0.000505, - 0.000712, - 0.000555, - 0.000453, - 0.00044, - 0.000194, - 0.000103, - 0.000345, - 0.000305, - 0.000136, - 0.000109, - 2.5e-05, - 5.6e-05, - 0.000104, - 5.9e-05, - -3e-05 - ] - } - ], - "SD3T5EncoderModel": [ - { - "shape": [ - 2, - 16, - 32 - ], - "sample": [ - 0.00707, - 0.006957, - 0.009331, - 0.004635, - -0.015012, - -0.007533, - -0.008983, - 0.000647, - 0.005648, - -0.0074, - 0.003372, - -0.011393, - 0.001414, - 0.027524, - -0.005903, - -0.011301, - 0.012131, - 0.013565, - 0.020675, - 0.002784, - 0.028628, - -0.004386, - 0.000328, - -0.004168, - 0.012268, - -0.01554, - -2.8e-05, - -0.012873, - 2.4e-05, - 0.054132, - -0.010922, - 0.005437, - 0.015127, - -0.027579, - -0.002549, - 0.003275, - 0.001096, - 0.010374, - -0.008043, - -0.030602, - 0.000448, - -0.032054, - 0.014059, - -0.003486, - -0.030945, - -0.000766, - -0.015376, - -0.001499, - -0.00409, - 0.001824, - 0.002497, - 0.015728, - 0.018065, - -0.01547, - -0.037632, - -0.002506, - -0.001476, - -0.004758, - 0.0006, - -0.01241, - 0.000607, - -0.003328, - 5.1e-05, - -0.014076 - ] - } - ], - "SENetImageClassify": [ - { - "shape": [ - 2, - 1000 - ], - "sample": [ - 0.033375, - 0.030584, - 0.00939, - -0.023106, - -0.020252, - 0.044979, - -0.03994, - 0.037784, - 0.009241, - -0.05507, - -0.00591, - -0.007533, - 0.023775, - 0.014904, - -0.016802, - 0.021809, - 0.008639, - 0.026153, - -0.000224, - -0.003933, - 0.040378, - -0.011918, - 0.034627, - 0.018249, - -0.025581, - 0.020486, - 0.004803, - -0.052036, - -0.009788, - 0.009076, - -0.001145, - -0.021569, - 0.055696, - -0.000346, - -0.008035, - -0.037688, - 0.011819, - 0.013224, - -0.032714, - 0.018822, - 0.008501, - -0.026391, - -0.053209, - -0.031158, - 0.038006, - -0.038696, - -0.019482, - -0.018734, - 0.004919, - -0.031237, - 0.007237, - 0.009181, - -0.040009, - -0.048588, - -0.015589, - 0.005695, - 0.033936, - 0.033833, - 0.011182, - -0.033774, - -0.009193, - -0.003455, - 0.004113, - -0.034696 - ] - } - ], - "SegFormerSemanticSegment": [ - { - "shape": [ - 2, - 32, - 32, - 150 - ], - "sample": [ - -0.003387, - 0.000225, - 0.021734, - -0.01509, - -0.028037, - 0.006117, - 0.019701, - -0.003491, - 0.013617, - -0.010724, - -0.01798, - 9.9e-05, - -0.016505, - 0.017597, - 0.012508, - -0.005059, - -0.020864, - -0.019498, - -0.006232, - -0.018937, - -0.01325, - -0.021692, - -0.007152, - 0.00235, - -0.003654, - -0.00454, - -0.00058, - -0.027338, - 0.008141, - -4.1e-05, - -0.002294, - 0.037869, - -0.008589, - 0.007304, - 0.009095, - -0.029071, - 0.027942, - -0.005207, - 0.004451, - 0.008656, - 0.020167, - 0.025304, - 0.041152, - 0.013345, - -0.009861, - 0.013261, - 0.004396, - 0.016357, - -0.000784, - -0.023537, - 0.016361, - -0.014649, - -0.021318, - -0.02706, - -0.006946, - -0.024722, - 0.006379, - -0.014328, - 0.020379, - -0.008039, - 0.034676, - -0.004492, - -0.01603, - 0.018704 - ] - } - ], - "SentenceTransformerModel": [ - { - "shape": [ - 2, - 32 - ], - "sample": [ - -0.138292, - 0.107383, - 0.227017, - -0.157133, - 0.043625, - -0.147396, - 0.077896, - 0.145626, - 0.28461, - -0.071275, - 0.122478, - -0.137024, - -0.125705, - -0.16926, - -0.036518, - 0.066965, - 0.107061, - -0.087836, - 0.160158, - -0.163558, - -0.408743, - 0.039037, - 0.045357, - -0.055598, - -0.408058, - 0.181733, - -0.026227, - 0.110355, - 0.224607, - 0.15685, - -0.357952, - -0.10573, - -0.136634, - 0.107497, - 0.226582, - -0.15824, - 0.043457, - -0.147285, - 0.080659, - 0.144457, - 0.282996, - -0.068748, - 0.124357, - -0.137633, - -0.122886, - -0.167857, - -0.037022, - 0.067155, - 0.108349, - -0.086218, - 0.160484, - -0.164189, - -0.406857, - 0.036848, - 0.046163, - -0.051938, - -0.41074, - 0.180494, - -0.025904, - 0.110552, - 0.223734, - 0.157155, - -0.360933, - -0.106077 - ] - }, - { - "shape": [ - 2, - 16, - 32 - ], - "sample": [ - -0.020572, - 0.015957, - -0.019703, - 0.015892, - -0.020483, - -0.01368, - 0.016936, - -0.015326, - 0.014478, - 0.023795, - 0.033209, - 0.02408, - 0.033903, - -0.024844, - -0.023571, - -0.024457, - -0.02343, - -0.059766, - 0.006253, - -0.062119, - 0.006618, - 0.005499, - -0.020619, - 0.005742, - -0.023827, - 0.006142, - 0.009706, - 0.006107, - 0.008017, - 0.00602, - 0.021343, - -0.013266, - 0.015743, - -0.010255, - 0.043917, - -0.061277, - 0.042983, - -0.062477, - -0.010429, - 0.026378, - -0.009833, - 0.02688, - 0.019112, - -0.004345, - 0.019161, - -0.004257, - 0.018308, - 0.015562, - -0.020819, - 0.01682, - -0.02066, - 0.032244, - -0.018348, - 0.033819, - -0.017939, - 0.02347, - -0.025788, - 0.023409, - -0.025221, - -0.053642, - -0.004411, - -0.055167, - -0.005387, - -0.015949 - ] - } - ], - "SentenceTransformerTextEmbed": [ - { - "shape": [ - 2, - 32 - ], - "sample": [ - -0.138292, - 0.107383, - 0.227017, - -0.157133, - 0.043625, - -0.147396, - 0.077896, - 0.145626, - 0.28461, - -0.071275, - 0.122478, - -0.137024, - -0.125705, - -0.16926, - -0.036518, - 0.066965, - 0.107061, - -0.087836, - 0.160158, - -0.163558, - -0.408743, - 0.039037, - 0.045357, - -0.055598, - -0.408058, - 0.181733, - -0.026227, - 0.110355, - 0.224607, - 0.15685, - -0.357952, - -0.10573, - -0.136634, - 0.107497, - 0.226582, - -0.15824, - 0.043457, - -0.147285, - 0.080659, - 0.144457, - 0.282996, - -0.068748, - 0.124357, - -0.137633, - -0.122886, - -0.167857, - -0.037022, - 0.067155, - 0.108349, - -0.086218, - 0.160484, - -0.164189, - -0.406857, - 0.036848, - 0.046163, - -0.051938, - -0.41074, - 0.180494, - -0.025904, - 0.110552, - 0.223734, - 0.157155, - -0.360933, - -0.106077 - ] - }, - { - "shape": [ - 2, - 16, - 32 - ], - "sample": [ - -0.020572, - 0.015957, - -0.019703, - 0.015892, - -0.020483, - -0.01368, - 0.016936, - -0.015326, - 0.014478, - 0.023795, - 0.033209, - 0.02408, - 0.033903, - -0.024844, - -0.023571, - -0.024457, - -0.02343, - -0.059766, - 0.006253, - -0.062119, - 0.006618, - 0.005499, - -0.020619, - 0.005742, - -0.023827, - 0.006142, - 0.009706, - 0.006107, - 0.008017, - 0.00602, - 0.021343, - -0.013266, - 0.015743, - -0.010255, - 0.043917, - -0.061277, - 0.042983, - -0.062477, - -0.010429, - 0.026378, - -0.009833, - 0.02688, - 0.019112, - -0.004345, - 0.019161, - -0.004257, - 0.018308, - 0.015562, - -0.020819, - 0.01682, - -0.02066, - 0.032244, - -0.018348, - 0.033819, - -0.017939, - 0.02347, - -0.025788, - 0.023409, - -0.025221, - -0.053642, - -0.004411, - -0.055167, - -0.005387, - -0.015949 - ] - } - ], - "SentenceTransformer_all_MiniLM_L12_v2": [ - { - "shape": [ - 2, - 16 - ], - "sample": [ - -0.006283, - -0.28829, - -0.510357, - -0.181786, - 0.267572, - -0.204303, - 0.065134, - -0.036725, - 0.221611, - -0.276, - 0.044667, - -0.319701, - -0.356479, - 0.162738, - 0.068047, - -0.341659, - -0.007111, - -0.286077, - -0.509385, - -0.181498, - 0.266028, - -0.204473, - 0.059981, - -0.032958, - 0.218527, - -0.278198, - 0.045874, - -0.318675, - -0.36582, - 0.169479, - 0.064059, - -0.336116 - ] - }, - { - "shape": [ - 2, - 16, - 16 - ], - "sample": [ - -0.001388, - 0.02555, - 0.001283, - 0.024534, - 0.000669, - 0.0236, - -0.000574, - 0.024204, - -9.2e-05, - -0.031526, - -0.032391, - -0.030826, - -0.032146, - -0.031856, - -0.032201, - -0.030922, - -0.032239, - -0.032286, - -0.05965, - 0.004568, - -0.064821, - 0.002639, - -0.055885, - 0.007288, - -0.060506, - 0.002031, - -0.053601, - -0.035833, - -0.021838, - -0.035347, - -0.020152, - -0.036121, - -0.020518, - -0.036075, - -0.020536, - -0.036254, - 0.030103, - -0.039094, - 0.030121, - -0.035787, - 0.030136, - -0.047276, - 0.030123, - -0.041288, - 0.030123, - 0.014357, - -0.020862, - 0.01455, - -0.023323, - 0.019685, - -0.02363, - 0.00853, - -0.020957, - 0.031693, - 0.003902, - 0.005723, - 0.007174, - 0.006236, - 0.003883, - 0.006472, - 0.015475, - 0.010045, - -0.001431, - -0.036287 - ] - } - ], - "SentenceTransformer_all_MiniLM_L6_v2": [ - { - "shape": [ - 2, - 16 - ], - "sample": [ - -0.006283, - -0.28829, - -0.510357, - -0.181786, - 0.267572, - -0.204303, - 0.065134, - -0.036725, - 0.221611, - -0.276, - 0.044667, - -0.319701, - -0.356479, - 0.162738, - 0.068047, - -0.341659, - -0.007111, - -0.286077, - -0.509385, - -0.181498, - 0.266028, - -0.204473, - 0.059981, - -0.032958, - 0.218527, - -0.278198, - 0.045874, - -0.318675, - -0.36582, - 0.169479, - 0.064059, - -0.336116 - ] - }, - { - "shape": [ - 2, - 16, - 16 - ], - "sample": [ - -0.001388, - 0.02555, - 0.001283, - 0.024534, - 0.000669, - 0.0236, - -0.000574, - 0.024204, - -9.2e-05, - -0.031526, - -0.032391, - -0.030826, - -0.032146, - -0.031856, - -0.032201, - -0.030922, - -0.032239, - -0.032286, - -0.05965, - 0.004568, - -0.064821, - 0.002639, - -0.055885, - 0.007288, - -0.060506, - 0.002031, - -0.053601, - -0.035833, - -0.021838, - -0.035347, - -0.020152, - -0.036121, - -0.020518, - -0.036075, - -0.020536, - -0.036254, - 0.030103, - -0.039094, - 0.030121, - -0.035787, - 0.030136, - -0.047276, - 0.030123, - -0.041288, - 0.030123, - 0.014357, - -0.020862, - 0.01455, - -0.023323, - 0.019685, - -0.02363, - 0.00853, - -0.020957, - 0.031693, - 0.003902, - 0.005723, - 0.007174, - 0.006236, - 0.003883, - 0.006472, - 0.015475, - 0.010045, - -0.001431, - -0.036287 - ] - } - ], - "SentenceTransformer_all_distilroberta_v1": [ - { - "shape": [ - 2, - 16 - ], - "sample": [ - -0.006428, - -0.288801, - -0.508576, - -0.181893, - 0.268083, - -0.204161, - 0.058898, - -0.038747, - 0.221928, - -0.27647, - 0.045056, - -0.320287, - -0.360028, - 0.155845, - 0.068021, - -0.342725, - -0.003668, - -0.288044, - -0.509014, - -0.183166, - 0.268015, - -0.205943, - 0.054317, - -0.032845, - 0.218373, - -0.281978, - 0.046054, - -0.321448, - -0.358621, - 0.167182, - 0.063274, - -0.335987 - ] - }, - { - "shape": [ - 2, - 16, - 16 - ], - "sample": [ - 0.00054, - 0.024421, - 0.00022, - 0.024225, - 0.000921, - 0.023616, - -0.00275, - 0.026056, - 8.3e-05, - -0.030825, - -0.032213, - -0.031122, - -0.032074, - -0.0312, - -0.032084, - -0.033028, - -0.032695, - -0.030283, - -0.064057, - 0.005521, - -0.059847, - 0.003993, - -0.058114, - 0.005001, - -0.053613, - 0.000932, - -0.055822, - -0.035447, - -0.021108, - -0.036471, - -0.020006, - -0.035757, - -0.020632, - -0.036374, - -0.020734, - -0.03649, - 0.030099, - -0.045783, - 0.030133, - -0.03234, - 0.030137, - -0.041763, - 0.030125, - -0.045943, - 0.030106, - 0.020362, - -0.021803, - 0.011972, - -0.022738, - 0.021537, - -0.023277, - 0.009515, - -0.022329, - 0.040238, - 0.014673, - 0.006805, - 0.01218, - 0.008585, - -0.000855, - 0.006342, - 0.002824, - 0.008836, - -0.000799, - -0.036171 - ] - } - ], - "SentenceTransformer_all_mpnet_base_v2": [ - { - "shape": [ - 2, - 16 - ], - "sample": [ - 0.131721, - -0.261829, - -0.385196, - -0.173886, - 0.281246, - 0.001545, - 0.245415, - -0.072511, - -0.050302, - -0.224904, - 0.388975, - -0.262667, - -0.358547, - 0.329789, - 0.2014, - -0.199019, - 0.130397, - -0.259231, - -0.389688, - -0.17159, - 0.279488, - 0.000753, - 0.247705, - -0.083166, - -0.051451, - -0.230933, - 0.382909, - -0.258658, - -0.349881, - 0.325492, - 0.221085, - -0.203389 - ] - }, - { - "shape": [ - 2, - 16, - 16 - ], - "sample": [ - 0.012509, - -0.007654, - 0.012005, - -0.005402, - 0.012474, - -0.005269, - 0.013393, - -0.00231, - 0.01668, - -0.023257, - -0.027662, - -0.027021, - -0.028156, - -0.025185, - -0.027928, - -0.026351, - -0.02774, - -0.028348, - -0.031596, - 0.042527, - -0.038723, - 0.039582, - -0.045372, - 0.041086, - -0.044355, - 0.041554, - -0.03575, - -0.027495, - -0.018154, - -0.024365, - -0.019006, - -0.026141, - -0.018336, - -0.038645, - -0.018094, - -0.020396, - 0.030104, - -0.045994, - 0.030238, - -0.01341, - 0.030433, - -0.053719, - 0.03042, - -0.077768, - 0.03017, - 0.034797, - -0.002746, - 0.033777, - -0.006548, - 0.037556, - 0.012106, - 0.030717, - -0.003646, - 0.039079, - 0.038353, - 0.012978, - 0.040298, - 0.028337, - 0.011519, - 0.025773, - 0.027593, - 0.011926, - 0.014531, - -0.022591 - ] - } - ], - "SentenceTransformer_distiluse_base_multilingual_cased_v1": [ - { - "shape": [ - 2, - 8 - ], - "sample": [ - -0.005696, - 0.011375, - -0.000598, - -0.005064, - -0.008866, - 0.021272, - 0.007137, - -0.015385, - -0.005726, - 0.011304, - -0.000321, - -0.005126, - -0.008742, - 0.021266, - 0.00725, - -0.015222 - ] - }, - { - "shape": [ - 2, - 16, - 16 - ], - "sample": [ - 0.00223, - -0.007174, - 0.002873, - -0.009289, - 0.003751, - -0.007166, - 0.003752, - -0.007555, - 0.001998, - 0.011331, - -0.02211, - -0.001803, - -0.023278, - 0.00989, - -0.0181, - -0.00481, - -0.031739, - 0.000239, - 0.015857, - 0.06128, - 0.015085, - 0.06705, - 0.017578, - 0.05266, - 0.013209, - 0.060971, - 0.015753, - 0.049514, - -0.001151, - 0.049367, - -0.001122, - 0.033595, - -0.000648, - 0.041646, - -0.000901, - 0.048977, - -0.010922, - -0.039426, - -0.004766, - -0.042949, - -0.007115, - -0.037974, - -0.003507, - -0.042364, - -0.006889, - -0.005795, - -0.003808, - -0.007997, - -0.003384, - 0.000805, - -0.004321, - -0.008732, - -0.005165, - -0.006883, - -0.002073, - 0.050016, - -0.005411, - 0.052228, - -0.006485, - 0.053975, - -0.00512, - 0.0522, - -0.004789, - 0.017832 - ] - } - ], - "SentenceTransformer_distiluse_base_multilingual_cased_v2": [ - { - "shape": [ - 2, - 8 - ], - "sample": [ - -0.005696, - 0.011375, - -0.000598, - -0.005064, - -0.008866, - 0.021272, - 0.007137, - -0.015385, - -0.005726, - 0.011304, - -0.000321, - -0.005126, - -0.008742, - 0.021266, - 0.00725, - -0.015222 - ] - }, - { - "shape": [ - 2, - 16, - 16 - ], - "sample": [ - 0.00223, - -0.007174, - 0.002873, - -0.009289, - 0.003751, - -0.007166, - 0.003752, - -0.007555, - 0.001998, - 0.011331, - -0.02211, - -0.001803, - -0.023278, - 0.00989, - -0.0181, - -0.00481, - -0.031739, - 0.000239, - 0.015857, - 0.06128, - 0.015085, - 0.06705, - 0.017578, - 0.05266, - 0.013209, - 0.060971, - 0.015753, - 0.049514, - -0.001151, - 0.049367, - -0.001122, - 0.033595, - -0.000648, - 0.041646, - -0.000901, - 0.048977, - -0.010922, - -0.039426, - -0.004766, - -0.042949, - -0.007115, - -0.037974, - -0.003507, - -0.042364, - -0.006889, - -0.005795, - -0.003808, - -0.007997, - -0.003384, - 0.000805, - -0.004321, - -0.008732, - -0.005165, - -0.006883, - -0.002073, - 0.050016, - -0.005411, - 0.052228, - -0.006485, - 0.053975, - -0.00512, - 0.0522, - -0.004789, - 0.017832 - ] - } - ], - "SentenceTransformer_multi_qa_MiniLM_L6_cos_v1": [ - { - "shape": [ - 2, - 16 - ], - "sample": [ - -0.006283, - -0.28829, - -0.510357, - -0.181786, - 0.267572, - -0.204303, - 0.065134, - -0.036725, - 0.221611, - -0.276, - 0.044667, - -0.319701, - -0.356479, - 0.162738, - 0.068047, - -0.341659, - -0.007111, - -0.286077, - -0.509385, - -0.181498, - 0.266028, - -0.204473, - 0.059981, - -0.032958, - 0.218527, - -0.278198, - 0.045874, - -0.318675, - -0.36582, - 0.169479, - 0.064059, - -0.336116 - ] - }, - { - "shape": [ - 2, - 16, - 16 - ], - "sample": [ - -0.001388, - 0.02555, - 0.001283, - 0.024534, - 0.000669, - 0.0236, - -0.000574, - 0.024204, - -9.2e-05, - -0.031526, - -0.032391, - -0.030826, - -0.032146, - -0.031856, - -0.032201, - -0.030922, - -0.032239, - -0.032286, - -0.05965, - 0.004568, - -0.064821, - 0.002639, - -0.055885, - 0.007288, - -0.060506, - 0.002031, - -0.053601, - -0.035833, - -0.021838, - -0.035347, - -0.020152, - -0.036121, - -0.020518, - -0.036075, - -0.020536, - -0.036254, - 0.030103, - -0.039094, - 0.030121, - -0.035787, - 0.030136, - -0.047276, - 0.030123, - -0.041288, - 0.030123, - 0.014357, - -0.020862, - 0.01455, - -0.023323, - 0.019685, - -0.02363, - 0.00853, - -0.020957, - 0.031693, - 0.003902, - 0.005723, - 0.007174, - 0.006236, - 0.003883, - 0.006472, - 0.015475, - 0.010045, - -0.001431, - -0.036287 - ] - } - ], - "SentenceTransformer_multi_qa_distilbert_cos_v1": [ - { - "shape": [ - 2, - 16 - ], - "sample": [ - 0.025973, - -0.206354, - 0.137715, - -0.007853, - -0.064003, - -0.047599, - -0.033647, - 0.035637, - -0.075466, - 0.007383, - 0.548441, - 0.416146, - -0.376682, - -0.05115, - 0.495402, - 0.243534, - 0.025587, - -0.19556, - 0.143738, - -0.007995, - -0.066312, - -0.045355, - -0.03309, - 0.023793, - -0.075691, - -0.024764, - 0.571854, - 0.411226, - -0.383017, - -0.0612, - 0.497094, - 0.179573 - ] - }, - { - "shape": [ - 2, - 16, - 16 - ], - "sample": [ - 0.00223, - -0.007174, - 0.002873, - -0.009289, - 0.003751, - -0.007166, - 0.003752, - -0.007555, - 0.001998, - 0.011331, - -0.02211, - -0.001803, - -0.023278, - 0.00989, - -0.0181, - -0.00481, - -0.031739, - 0.000239, - 0.015857, - 0.06128, - 0.015085, - 0.06705, - 0.017578, - 0.05266, - 0.013209, - 0.060971, - 0.015753, - 0.049514, - -0.001151, - 0.049367, - -0.001122, - 0.033595, - -0.000648, - 0.041646, - -0.000901, - 0.048977, - -0.010922, - -0.039426, - -0.004766, - -0.042949, - -0.007115, - -0.037974, - -0.003507, - -0.042364, - -0.006889, - -0.005795, - -0.003808, - -0.007997, - -0.003384, - 0.000805, - -0.004321, - -0.008732, - -0.005165, - -0.006883, - -0.002073, - 0.050016, - -0.005411, - 0.052228, - -0.006485, - 0.053975, - -0.00512, - 0.0522, - -0.004789, - 0.017832 - ] - } - ], - "SentenceTransformer_multi_qa_mpnet_base_dot_v1": [ - { - "shape": [ - 2, - 16 - ], - "sample": [ - 0.012509, - -0.027962, - -0.041582, - -0.018233, - 0.030212, - -0.002694, - 0.010153, - 0.004949, - -0.007654, - -0.025851, - 0.042763, - -0.02836, - -0.03605, - 0.034484, - 0.022787, - -0.014465, - 0.007949, - -0.027739, - -0.047459, - -0.018336, - 0.030239, - -0.001236, - 0.020397, - -0.011969, - -0.007566, - -0.027303, - 0.041396, - -0.038645, - -0.02535, - 0.034228, - 0.041315, - -0.024879 - ] - }, - { - "shape": [ - 2, - 16, - 16 - ], - "sample": [ - 0.012509, - -0.007654, - 0.012005, - -0.005402, - 0.012474, - -0.005269, - 0.013393, - -0.00231, - 0.01668, - -0.023257, - -0.027662, - -0.027021, - -0.028156, - -0.025185, - -0.027928, - -0.026351, - -0.02774, - -0.028348, - -0.031596, - 0.042527, - -0.038723, - 0.039582, - -0.045372, - 0.041086, - -0.044355, - 0.041554, - -0.03575, - -0.027495, - -0.018154, - -0.024365, - -0.019006, - -0.026141, - -0.018336, - -0.038645, - -0.018094, - -0.020396, - 0.030104, - -0.045994, - 0.030238, - -0.01341, - 0.030433, - -0.053719, - 0.03042, - -0.077768, - 0.03017, - 0.034797, - -0.002746, - 0.033777, - -0.006548, - 0.037556, - 0.012106, - 0.030717, - -0.003646, - 0.039079, - 0.038353, - 0.012978, - 0.040298, - 0.028337, - 0.011519, - 0.025773, - 0.027593, - 0.011926, - 0.014531, - -0.022591 + -0.015336, + -0.013034, + 0.006987, + 0.020153, + 0.049287, + 0.01717, + 0.044167, + -0.03045, + 0.032655, + 0.059362, + -0.011628, + 0.017312, + -0.02234, + -0.00998, + 0.018686, + 0.041445, + -0.02269, + 0.018144, + -0.032625, + -0.017222, + -0.009087, + 0.017343, + 0.014006, + -0.00969, + 0.003306, + 0.026315, + 0.005598, + -0.020814, + -0.035375, + -0.007683, + -0.000294, + 0.014428, + -0.048157, + 0.006775, + 0.045402, + 0.047048, + 0.008475, + 0.020926, + 0.001399, + -0.014787, + -0.0152, + -0.013629, + -0.019056, + -0.027307, + -0.047092, + -0.022155, + -0.003388, + 0.033364, + 0.024901, + 0.012694, + -0.007007, + 0.037517, + 0.002602, + 0.016988, + -0.031581, + 0.02448, + 0.031745, + 0.007232, + 0.002769, + 0.023481, + -0.01057, + 0.049235, + -0.043309, + -0.016593 ] - } - ], - "SentenceTransformer_paraphrase_MiniLM_L3_v2": [ + }, { "shape": [ 2, - 16 + 8, + 8, + 256 ], "sample": [ - -0.000708, - -0.032463, - -0.05747, - -0.02047, - 0.03013, - -0.023006, - 0.007335, - -0.004135, - 0.024955, - -0.03108, - 0.00503, - -0.036001, - -0.040142, - 0.018325, - 0.007663, - -0.038473, - -0.000805, - -0.032396, - -0.057684, - -0.020553, - 0.030126, - -0.023155, - 0.006792, - -0.003732, - 0.024747, - -0.031504, - 0.005195, - -0.036088, - -0.041427, - 0.019192, - 0.007254, - -0.038063 + 0.001858, + 0.005013, + -0.03215, + -0.023843, + 0.023074, + -0.009762, + 0.007063, + 0.007979, + -0.023957, + 0.041274, + -0.039285, + -0.017954, + 0.020519, + -0.03382, + -0.00841, + -0.00636, + 0.027404, + 0.019387, + -0.024081, + -0.010251, + -0.031152, + -0.001085, + -0.01668, + 0.017862, + -0.004157, + 0.002491, + 0.001924, + -0.004861, + -0.002372, + 0.03471, + -0.029995, + -0.01142, + 0.001389, + -0.026705, + 0.003308, + 0.010559, + -0.002314, + -0.00535, + -0.001025, + 0.028143, + -0.019863, + 0.015856, + -0.014986, + -0.011457, + -0.017505, + 0.020718, + -0.002761, + -0.027695, + -0.00788, + -0.002352, + 0.036122, + -0.014375, + -0.022155, + -0.022174, + 0.008941, + 0.003909, + 0.018927, + -0.014251, + -0.035766, + -0.011259, + -0.011355, + -0.017476, + 0.009668, + -0.038129 ] }, { "shape": [ 2, - 16, - 16 + 4, + 4, + 256 ], "sample": [ - -0.001388, - 0.02555, - 0.001283, - 0.024534, - 0.000669, - 0.0236, - -0.000574, - 0.024204, - -9.2e-05, - -0.031526, - -0.032391, - -0.030826, - -0.032146, - -0.031856, - -0.032201, - -0.030922, - -0.032239, - -0.032286, - -0.05965, - 0.004568, - -0.064821, - 0.002639, - -0.055885, - 0.007288, - -0.060506, - 0.002031, - -0.053601, - -0.035833, - -0.021838, - -0.035347, - -0.020152, - -0.036121, - -0.020518, - -0.036075, - -0.020536, - -0.036254, - 0.030103, - -0.039094, - 0.030121, - -0.035787, - 0.030136, - -0.047276, - 0.030123, - -0.041288, - 0.030123, - 0.014357, - -0.020862, - 0.01455, - -0.023323, - 0.019685, - -0.02363, - 0.00853, - -0.020957, - 0.031693, - 0.003902, - 0.005723, - 0.007174, - 0.006236, - 0.003883, - 0.006472, - 0.015475, - 0.010045, - -0.001431, - -0.036287 + 0.030796, + 0.040807, + -0.004929, + -0.010892, + -0.019576, + -0.057813, + -0.069933, + -0.025368, + -0.008011, + 0.019821, + 0.021221, + 0.004241, + 0.02296, + 0.01008, + -0.018551, + -0.03447, + -0.003466, + -0.005969, + -0.062868, + 0.043891, + 0.066126, + -0.002958, + 0.044128, + -0.038288, + -0.038014, + 0.008986, + 0.027282, + -0.036212, + 0.068866, + -0.036626, + 0.007519, + 0.009738, + -0.035756, + 0.026206, + -0.000522, + -0.027787, + -0.042496, + -0.021335, + 0.032815, + -0.064257, + 0.009692, + 0.019132, + 0.006331, + -0.026821, + 0.017907, + 0.025561, + -0.034405, + -0.039337, + -0.054798, + -0.01864, + -0.035654, + -0.042077, + 0.030411, + 0.030888, + 0.021825, + -0.009091, + -0.063324, + 0.044919, + 0.035721, + 0.014497, + 0.016162, + -0.019659, + 0.017876, + -0.025821 ] } ], - "SentenceTransformer_paraphrase_albert_small_v2": [ + "SAM2PromptableSegment": [ { "shape": [ 2, - 16 + 1, + 3 ], "sample": [ - 0.019625, - -0.000452, - -0.065652, - 0.038684, - -0.016475, - -0.077698, - 0.066542, - 0.033501, - -0.000845, - 0.026137, - -0.01556, - 0.010658, - -0.019583, - -0.014868, - -0.019154, - 0.056296, - 0.019622, - -0.0004, - -0.065506, - 0.038478, - -0.016457, - -0.077675, - 0.066516, - 0.033927, - -0.000842, - 0.026165, - -0.015646, - 0.010543, - -0.019583, - -0.014792, - -0.018992, - 0.055857 + 0.498468, + 0.503674, + 0.500753, + 0.498468, + 0.503674, + 0.500753 ] }, { "shape": [ 2, + 1, + 1 + ], + "sample": [ + -0.00015, + -0.00015 + ] + }, + { + "shape": [ + 2, + 1, + 3, 16, 16 ], "sample": [ - 0.019678, - -0.000842, - 0.019711, - -0.000838, - 0.01957, - -0.000845, - 0.019591, - -0.000842, - 0.019633, - 0.026074, - -0.000336, - 0.026192, - -0.000561, - 0.026158, - -0.000532, - 0.026085, - -0.000782, - 0.026046, - -0.066363, - -0.015558, - -0.065007, - -0.015725, - -0.066967, - -0.015305, - -0.065715, - -0.015685, - -0.064612, - 0.0103, - 0.038897, - 0.010929, - 0.039194, - 0.010795, - 0.038024, - 0.010261, - 0.03857, - 0.010324, - -0.016849, - -0.019532, - -0.016213, - -0.019658, - -0.016232, - -0.019634, - -0.017388, - -0.019531, - -0.015906, - -0.015022, - -0.077507, - -0.014948, - -0.077572, - -0.014514, - -0.07773, - -0.014773, - -0.077588, - -0.014602, - 0.066499, - -0.019658, - 0.066418, - -0.018982, - 0.066457, - -0.018579, - 0.066452, - -0.018653, - 0.066585, - 0.057405 + 0.001723, + 0.001358, + 0.001621, + 0.001894, + 0.001981, + 0.001704, + 0.001766, + 0.00185, + 0.002127, + 0.002068, + 0.001415, + -3.9e-05, + 0.000403, + 0.000159, + -3.1e-05, + -0.00026, + -0.000519, + 0.000495, + 0.000265, + -0.000139, + 1.4e-05, + 0.000188, + -0.001361, + -0.001924, + -0.0008, + -0.001101, + -0.00219, + -0.001732, + -0.00103, + -0.001181, + -0.001395, + -0.001796, + 0.001847, + 0.001937, + 0.001716, + 0.002225, + 0.001588, + 0.001684, + 0.001877, + 0.002109, + 0.001345, + 0.002057, + 0.002675, + 0.000631, + 0.000172, + -9.5e-05, + 0.000222, + -2.5e-05, + 8.1e-05, + -0.000149, + -0.00035, + 0.000488, + 6.9e-05, + -0.001721, + -0.001244, + -0.001717, + -0.000922, + -0.001359, + -0.001446, + -0.001118, + -0.002134, + -0.001088, + -0.001499, + -0.00176 + ] + } + ], + "SAMModel": [ + { + "shape": [ + 2, + 4, + 4, + 256 + ], + "sample": [ + -0.025676, + -0.004309, + 0.0244, + 0.001673, + 0.030505, + -0.0386, + 0.01161, + 0.005268, + 0.010626, + 0.018627, + -0.025338, + -0.016413, + 0.025446, + -0.015198, + -0.008434, + -0.008578, + -0.003512, + 0.033251, + -0.070632, + -0.005308, + 0.000303, + 0.00923, + 0.026017, + 0.026024, + -0.019712, + -0.014215, + 0.0426, + -0.04387, + -0.028514, + 0.042532, + 0.02437, + 0.080674, + -0.039657, + 0.029955, + 0.05102, + -0.00787, + 0.007323, + -0.013503, + -0.027775, + 0.013422, + -0.014877, + 0.008596, + -0.027227, + -0.028858, + 0.013134, + -0.026335, + 0.023304, + -0.028756, + -0.027388, + 0.016623, + -0.037225, + -0.014127, + 0.010916, + 0.017536, + 0.015706, + 0.013338, + 0.02569, + 0.002734, + 0.061468, + 0.036984, + -0.011109, + 0.009696, + 0.073214, + -0.006962 ] } ], - "SentenceTransformer_paraphrase_multilingual_MiniLM_L12_v2": [ + "SAMPromptableSegment": [ { "shape": [ 2, - 16 + 1, + 3 ], "sample": [ - -0.000708, - -0.032463, - -0.05747, - -0.02047, - 0.03013, - -0.023006, - 0.007335, - -0.004135, - 0.024955, - -0.03108, - 0.00503, - -0.036001, - -0.040142, - 0.018325, - 0.007663, - -0.038473, - -0.000805, - -0.032396, - -0.057684, - -0.020553, - 0.030126, - -0.023155, - 0.006792, - -0.003732, - 0.024747, - -0.031504, - 0.005195, - -0.036088, - -0.041427, - 0.019192, - 0.007254, - -0.038063 + -0.026417, + 0.00425, + 0.015243, + -0.026417, + 0.00425, + 0.015242 ] }, { "shape": [ 2, + 1, + 3, 16, 16 ], "sample": [ - -0.001388, - 0.02555, - 0.001283, - 0.024534, - 0.000669, - 0.0236, - -0.000574, - 0.024204, - -9.2e-05, - -0.031526, - -0.032391, - -0.030826, - -0.032146, - -0.031856, - -0.032201, - -0.030922, - -0.032239, - -0.032286, - -0.05965, - 0.004568, - -0.064821, - 0.002639, - -0.055885, - 0.007288, - -0.060506, - 0.002031, - -0.053601, - -0.035833, - -0.021838, - -0.035347, - -0.020152, - -0.036121, - -0.020518, - -0.036075, - -0.020536, - -0.036254, - 0.030103, - -0.039094, - 0.030121, - -0.035787, - 0.030136, - -0.047276, - 0.030123, - -0.041288, - 0.030123, - 0.014357, - -0.020862, - 0.01455, - -0.023323, - 0.019685, - -0.02363, - 0.00853, - -0.020957, - 0.031693, - 0.003902, - 0.005723, - 0.007174, - 0.006236, - 0.003883, - 0.006472, - 0.015475, - 0.010045, - -0.001431, - -0.036287 + -0.001283, + -0.001259, + -0.001282, + -0.001481, + -0.001512, + -0.001624, + -0.001298, + -0.001273, + -0.001249, + -0.001602, + -0.001545, + 0.000605, + 0.000614, + 0.000463, + 0.000507, + 0.000717, + 0.000802, + 0.000456, + 0.000451, + 0.000556, + 0.000791, + 0.000508, + 0.000305, + 0.000347, + 0.000117, + 0.000161, + 5.2e-05, + 2.2e-05, + 6e-05, + 9.3e-05, + 0.000286, + -1e-06, + -0.001527, + -0.001283, + -0.00129, + -0.001262, + -0.001508, + -0.001483, + -0.001583, + -0.001311, + -0.001245, + -0.001255, + -0.001556, + 0.000541, + 0.000472, + 0.000614, + 0.000621, + 0.000521, + 0.000505, + 0.000712, + 0.000555, + 0.000453, + 0.00044, + 0.000194, + 0.000103, + 0.000345, + 0.000305, + 0.000136, + 0.000109, + 2.5e-05, + 5.6e-05, + 0.000104, + 5.9e-05, + -3e-05 + ] + } + ], + "SD3T5EncoderModel": [ + { + "shape": [ + 2, + 16, + 32 + ], + "sample": [ + 0.00707, + 0.006957, + 0.009331, + 0.004635, + -0.015012, + -0.007533, + -0.008983, + 0.000647, + 0.005648, + -0.0074, + 0.003372, + -0.011393, + 0.001414, + 0.027524, + -0.005903, + -0.011301, + 0.012131, + 0.013565, + 0.020675, + 0.002784, + 0.028628, + -0.004386, + 0.000328, + -0.004168, + 0.012268, + -0.01554, + -2.8e-05, + -0.012873, + 2.4e-05, + 0.054132, + -0.010922, + 0.005437, + 0.015127, + -0.027579, + -0.002549, + 0.003275, + 0.001096, + 0.010374, + -0.008043, + -0.030602, + 0.000448, + -0.032054, + 0.014059, + -0.003486, + -0.030945, + -0.000766, + -0.015376, + -0.001499, + -0.00409, + 0.001824, + 0.002497, + 0.015728, + 0.018065, + -0.01547, + -0.037632, + -0.002506, + -0.001476, + -0.004758, + 0.0006, + -0.01241, + 0.000607, + -0.003328, + 5.1e-05, + -0.014076 ] } ], - "SentenceTransformer_paraphrase_multilingual_mpnet_base_v2": [ + "SENetImageClassify": [ { "shape": [ 2, - 16 + 1000 ], "sample": [ - -0.000722, - -0.032459, - -0.05716, - -0.020443, - 0.03013, - -0.022946, - 0.00662, - -0.004355, - 0.024943, - -0.031073, - 0.005064, - -0.035998, - -0.040464, - 0.017516, - 0.007645, - -0.038519, - -0.000412, - -0.032375, - -0.057211, - -0.020587, - 0.030124, - -0.023147, - 0.006105, - -0.003692, - 0.024544, - -0.031693, - 0.005176, - -0.036129, - -0.040307, - 0.01879, - 0.007112, - -0.037763 + 0.033375, + 0.030584, + 0.00939, + -0.023106, + -0.020252, + 0.044979, + -0.03994, + 0.037784, + 0.009241, + -0.05507, + -0.00591, + -0.007533, + 0.023775, + 0.014904, + -0.016802, + 0.021809, + 0.008639, + 0.026153, + -0.000224, + -0.003933, + 0.040378, + -0.011918, + 0.034627, + 0.018249, + -0.025581, + 0.020486, + 0.004803, + -0.052036, + -0.009788, + 0.009076, + -0.001145, + -0.021569, + 0.055696, + -0.000346, + -0.008035, + -0.037688, + 0.011819, + 0.013224, + -0.032714, + 0.018822, + 0.008501, + -0.026391, + -0.053209, + -0.031158, + 0.038006, + -0.038696, + -0.019482, + -0.018734, + 0.004919, + -0.031237, + 0.007237, + 0.009181, + -0.040009, + -0.048588, + -0.015589, + 0.005695, + 0.033936, + 0.033833, + 0.011182, + -0.033774, + -0.009193, + -0.003455, + 0.004113, + -0.034696 ] - }, + } + ], + "SegFormerSemanticSegment": [ { "shape": [ 2, - 16, - 16 + 32, + 32, + 150 ], "sample": [ - 0.00054, - 0.024421, - 0.00022, - 0.024225, - 0.000921, - 0.023616, - -0.00275, - 0.026056, - 8.3e-05, - -0.030825, - -0.032213, - -0.031122, - -0.032074, - -0.0312, - -0.032084, - -0.033028, - -0.032695, - -0.030283, - -0.064057, - 0.005521, - -0.059847, - 0.003993, - -0.058114, - 0.005001, - -0.053613, - 0.000932, - -0.055822, - -0.035447, - -0.021108, - -0.036471, - -0.020006, - -0.035757, - -0.020632, - -0.036374, - -0.020734, - -0.03649, - 0.030099, - -0.045783, - 0.030133, - -0.03234, - 0.030137, - -0.041763, - 0.030125, - -0.045943, - 0.030106, - 0.020362, - -0.021803, - 0.011972, - -0.022738, - 0.021537, - -0.023277, - 0.009515, - -0.022329, - 0.040238, - 0.014673, - 0.006805, - 0.01218, - 0.008585, - -0.000855, - 0.006342, - 0.002824, - 0.008836, - -0.000799, - -0.036171 + -0.003387, + 0.000225, + 0.021734, + -0.01509, + -0.028037, + 0.006117, + 0.019701, + -0.003491, + 0.013617, + -0.010724, + -0.01798, + 9.9e-05, + -0.016505, + 0.017597, + 0.012508, + -0.005059, + -0.020864, + -0.019498, + -0.006232, + -0.018937, + -0.01325, + -0.021692, + -0.007152, + 0.00235, + -0.003654, + -0.00454, + -0.00058, + -0.027338, + 0.008141, + -4.1e-05, + -0.002294, + 0.037869, + -0.008589, + 0.007304, + 0.009095, + -0.029071, + 0.027942, + -0.005207, + 0.004451, + 0.008656, + 0.020167, + 0.025304, + 0.041152, + 0.013345, + -0.009861, + 0.013261, + 0.004396, + 0.016357, + -0.000784, + -0.023537, + 0.016361, + -0.014649, + -0.021318, + -0.02706, + -0.006946, + -0.024722, + 0.006379, + -0.014328, + 0.020379, + -0.008039, + 0.034676, + -0.004492, + -0.01603, + 0.018704 ] } ], @@ -26499,3204 +23558,3158 @@ ] } ], - "StableDiffusion2Model": [ + "SwinImageClassify": [ { "shape": [ 2, - 16, - 16, - 3 + 1000 ], "sample": [ - 0.015141, - 0.014566, - 0.01507, - -0.013042, - -0.014793, - -0.013933, - 0.01303, - 0.014152, - 0.013037, - 0.014178, - 0.014338, - -0.013882, - -0.013874, - -0.013862, - 0.01476, - 0.014129, - 0.014743, - 0.014246, - 0.014257, - 0.014225, - -0.014007, - -0.013831, - 0.014013, - 0.014092, - 0.014013, - 0.014246, - 0.014301, - 0.014293, - -0.013722, - -0.013693, - -0.013301, - 0.014484, - 0.010029, - 0.014481, - 0.015246, - 0.015245, - -0.013927, - -0.013932, - -0.01385, - 0.014118, - 0.014105, - 0.014152, - 0.014243, - 0.01403, - -0.013841, - -0.013774, - -0.013812, - 0.014105, - 0.014097, - 0.014094, - 0.014242, - 0.014448, - 0.014266, - -0.014007, - -0.013837, - 0.014, - 0.014121, - 0.014003, - 0.014259, - 0.011565, - 0.014413, + 0.03727, + 0.010736, + 0.00379, -0.011314, - -0.01329, - 0.013126 + -0.001801, + 0.048002, + -0.037877, + 0.035164, + 0.011903, + -0.06182, + -0.001327, + -0.021618, + 0.03845, + 0.023628, + -0.012034, + 0.030154, + 0.013825, + 0.008501, + 0.0072, + 0.02061, + 0.008837, + -0.024512, + 0.017009, + 0.004001, + -0.040273, + 0.004263, + -0.000877, + -0.035575, + 0.030585, + -0.026599, + 0.030435, + 0.002476, + 0.025173, + 0.005227, + 0.011432, + -0.012005, + -0.002249, + -0.007996, + 0.011941, + 0.001431, + 0.020885, + -0.018125, + -0.041154, + -0.025066, + 0.029893, + -0.007118, + -0.012232, + -0.001271, + 0.004738, + -0.010213, + 0.047078, + -0.012023, + -0.002808, + -0.047233, + -0.017357, + -0.020568, + 0.021827, + 0.024365, + -0.028548, + -0.02213, + 0.001802, + -0.002804, + -0.002764, + -0.030708 ] - }, + } + ], + "SwinV2ImageClassify": [ { "shape": [ 2, - 8, - 8, - 8 + 1000 ], "sample": [ - 0.051589, - 0.051341, - 0.051341, - 0.05132, - 0.051681, - -0.006162, - -0.006151, - -0.006155, - -0.006046, - 0.006829, - 0.006821, - 0.006841, - 0.006748, - 0.010876, - 0.010889, - 0.010892, - 0.010883, - -0.021278, - -0.02128, - -0.021273, - -0.021104, - -0.005376, - -0.005379, - -0.005368, - -0.005528, - -0.005407, - -0.000975, - -0.000971, - -0.00076, - -0.000734, - -0.043369, - -0.043349, - -0.043402, - -0.043551, - 0.051337, - 0.051297, - 0.05151, - 0.051494, - -0.006147, - -0.006312, - -0.006151, - -0.006165, - 0.006816, - 0.006766, - 0.006782, - 0.006824, - 0.006837, - 0.010968, - 0.010946, - 0.0109, - 0.010906, - -0.021234, - -0.021271, - -0.021282, - -0.021255, - -0.005589, - -0.005439, - -0.005423, - -0.005422, - -0.000794, - -0.00072, - -0.000731, - -0.000732, - -0.043442 + 0.028752, + -0.006332, + 0.000363, + 0.002035, + -0.006701, + 0.047487, + -0.039371, + 0.023125, + 0.016489, + -0.060217, + -0.003939, + -0.024132, + 0.057357, + 0.01048, + -0.038023, + 0.010662, + 0.032048, + -0.001057, + 0.010816, + 0.01114, + 0.005381, + -0.038666, + 0.016187, + 0.005265, + -0.040001, + 0.002051, + 0.002523, + -0.020339, + 0.033311, + -0.027095, + 0.016853, + -0.019808, + 0.010958, + -0.014414, + 0.002567, + -0.027859, + -0.005412, + -0.005288, + 0.006264, + 0.010331, + 0.029314, + -0.024317, + -0.059341, + -0.047102, + 0.015536, + -0.005734, + -0.020719, + 0.00087, + -0.006053, + -0.031878, + 0.041451, + 0.015935, + 0.005199, + -0.049402, + -0.009516, + 0.011538, + 0.025188, + 0.020908, + -0.015169, + -0.022819, + 0.003664, + -0.009767, + 0.008631, + -0.043226 ] - }, + } + ], + "T5ConditionalGenerate": [ { "shape": [ 2, - 8, - 8, - 4 + 16, + 32 ], "sample": [ - 0.006219, - 0.002704, - 0.003533, - 0.003418, - 0.008001, - 0.005587, - 0.005602, - 0.005686, - 0.008691, - 0.002119, - -0.000714, - 0.001933, - -0.002422, + 0.007004, + 0.007011, + 0.009395, + 0.004649, + -0.014974, + -0.007318, + -0.008844, + 0.000876, + 0.005626, + -0.007449, + 0.003385, + -0.011547, + 0.00152, + 0.027466, + -0.006198, + -0.011448, + 0.011551, + 0.013829, + 0.020342, + 0.00263, + 0.029595, + -0.004371, + 0.000389, + -0.004182, + 0.012, + -0.015576, + -2.8e-05, + -0.012556, + 2.5e-05, + 0.054732, + -0.01089, + 0.00547, + 0.015315, + -0.027271, + -0.002516, + 0.003293, + 0.001099, + 0.010529, + -0.00809, + -0.031443, + 0.000377, + -0.031829, + 0.014496, + -0.003451, + -0.031001, + -0.00064, + -0.015198, + -0.001283, + -0.004123, 0.002237, - 0.001836, - 0.002578, - -0.001956, - 0.00071, - -0.003084, - -0.002698, - -0.005522, - -0.001268, - -0.001686, - -1.2e-05, - -0.004803, - 0.000822, - 0.000538, - 0.013731, - 0.01433, - 0.014584, - 0.015002, - 0.015691, - 0.013323, - 0.015859, - 0.016175, - 0.015331, - 0.004414, - 0.004255, - 0.006079, - 0.002555, - 0.005861, - 0.004837, - 0.005748, - 0.002341, - 0.005368, - -9.1e-05, - 0.001814, - 0.003786, - 0.001644, - 0.001444, - 0.00205, - 0.00168, - 8.1e-05, - 0.000167, - -0.001057, - -0.002688, - -0.000987, - -0.000541, - -0.000336, - -0.002546, - -0.005307, - -0.003004, - -0.006414, - 0.018894 + 0.002431, + 0.015693, + 0.018091, + -0.015281, + -0.037994, + -0.002521, + -0.001512, + -0.004742, + 0.000613, + -0.012441, + 0.000619, + -0.003316, + 8.3e-05, + -0.013941 ] }, { "shape": [ 2, - 16, + 8, 32 ], "sample": [ - -0.012036, - 0.027224, - -0.017103, - 0.027472, - -0.020236, - 0.000332, - -0.016442, - 0.015828, - -0.018295, - 0.014603, - -0.005828, - 0.016616, - -0.00709, - -0.005843, - 0.009211, - 0.000567, - 0.049181, - -0.012505, - 0.000328, - -0.009735, - -0.003474, - -0.007774, - -0.024105, - -0.00845, - -0.020531, - -0.008289, - -0.052391, - 0.003785, - -0.052397, - 0.003749, - -0.019434, - 0.013164, - -0.022706, - 0.014415, - 0.028551, - -0.025012, - 0.017438, - -0.033488, - 0.051644, - -0.027285, - 0.047169, - -0.013412, - 0.049211, - 0.009689, - 0.039967, - 0.02768, - 0.015537, - -0.001814, - -0.01621, - -0.001629, - -0.021368, - 0.015966, - -0.035636, - 0.013777, - -0.038478, - 0.024581, - 0.014664, - 0.015239, - 0.014259, - -0.002774, - 0.006665, - -0.001414, - -0.025644, - 0.009934 + -0.012666, + 0.008346, + -0.012532, + 0.025897, + -0.012666, + 0.008346, + -0.012532, + 0.025897, + -0.012666, + -0.008003, + 0.011512, + 0.020149, + -0.022597, + -0.008003, + 0.011512, + 0.020149, + -0.022597, + -0.008003, + 0.000519, + 0.005653, + -0.008304, + -8.3e-05, + 0.000519, + 0.005653, + -0.008304, + -8.3e-05, + 0.000519, + -0.005227, + -0.026162, + 0.014672, + -0.008533, + -0.005227, + -0.026154, + 0.014576, + -0.008315, + -0.005249, + -0.015154, + -0.005074, + -0.000946, + 0.010883, + -0.015154, + -0.005074, + -0.000946, + 0.010883, + -0.015154, + -0.024924, + 0.01542, + -0.004209, + -0.007258, + -0.024924, + 0.01542, + -0.004209, + -0.007258, + -0.024924, + -0.002469, + 0.002147, + -0.007502, + -0.002374, + -0.002469, + 0.002147, + -0.007502, + -0.002374, + -0.002469, + 0.002716 ] - } - ], - "StableDiffusion2TextToImage": [ + }, { "shape": [ 2, - 16, - 16, - 3 + 8, + 128 ], "sample": [ - 0.015141, - 0.014566, - 0.01507, - -0.013042, - -0.014793, - -0.013933, - 0.01303, - 0.014152, - 0.013037, - 0.014178, - 0.014338, - -0.013882, - -0.013874, - -0.013862, - 0.01476, - 0.014129, - 0.014743, - 0.014246, - 0.014257, - 0.014225, - -0.014007, - -0.013831, - 0.014013, - 0.014092, - 0.014013, - 0.014246, - 0.014301, - 0.014293, - -0.013722, - -0.013693, - -0.013301, - 0.014484, - 0.010029, - 0.014481, - 0.015246, - 0.015245, - -0.013927, - -0.013932, - -0.01385, - 0.014118, - 0.014105, - 0.014152, - 0.014243, - 0.01403, - -0.013841, - -0.013774, - -0.013812, - 0.014105, - 0.014097, - 0.014094, - 0.014242, - 0.014448, - 0.014266, - -0.014007, - -0.013837, - 0.014, - 0.014121, - 0.014003, - 0.014259, - 0.011565, - 0.014413, - -0.011314, - -0.01329, - 0.013126 + -1.3e-05, + -8.4e-05, + 6.1e-05, + -0.000177, + -0.000162, + -1.9e-05, + 0.000194, + 0.000133, + -0.000201, + 0.000181, + -0.000129, + 0.000109, + 2e-05, + 0.000283, + -0.000111, + -6e-06, + -9e-05, + -0.00027, + 6.1e-05, + -0.000117, + 4e-06, + -0.000192, + -0.000222, + 3.4e-05, + 3e-06, + -6.8e-05, + 0.000154, + -0.000511, + 0.000117, + -0.000262, + 0.000157, + -0.000222, + 0.000341, + -0.00019, + 0.000101, + 0.000131, + 0.000283, + -0.0, + 0.000358, + 3.3e-05, + 6.7e-05, + 0.000137, + -0.000217, + 0.000333, + 0.000191, + 0.000265, + 8.4e-05, + 5.3e-05, + -0.000328, + -0.000244, + 2.3e-05, + 0.000154, + 0.000141, + 8.2e-05, + 4e-06, + -0.00014, + -8.6e-05, + -0.00017, + 0.000444, + 0.000237, + 4e-05, + 1.8e-05, + -0.00063, + 0.000147 ] - }, + } + ], + "T5EncoderModel": [ { "shape": [ 2, - 8, - 8, - 8 + 16, + 32 ], "sample": [ - 0.051589, - 0.051341, - 0.051341, - 0.05132, - 0.051681, - -0.006162, - -0.006151, - -0.006155, - -0.006046, - 0.006829, - 0.006821, - 0.006841, - 0.006748, - 0.010876, - 0.010889, - 0.010892, - 0.010883, - -0.021278, - -0.02128, - -0.021273, - -0.021104, - -0.005376, - -0.005379, - -0.005368, - -0.005528, - -0.005407, - -0.000975, - -0.000971, - -0.00076, - -0.000734, - -0.043369, - -0.043349, - -0.043402, - -0.043551, - 0.051337, - 0.051297, - 0.05151, - 0.051494, - -0.006147, - -0.006312, - -0.006151, - -0.006165, - 0.006816, - 0.006766, - 0.006782, - 0.006824, - 0.006837, - 0.010968, - 0.010946, - 0.0109, - 0.010906, - -0.021234, - -0.021271, - -0.021282, - -0.021255, - -0.005589, - -0.005439, - -0.005423, - -0.005422, - -0.000794, - -0.00072, - -0.000731, - -0.000732, - -0.043442 + 0.007004, + 0.007011, + 0.009395, + 0.004649, + -0.014974, + -0.007318, + -0.008844, + 0.000876, + 0.005626, + -0.007449, + 0.003385, + -0.011547, + 0.00152, + 0.027466, + -0.006198, + -0.011448, + 0.011551, + 0.013829, + 0.020342, + 0.00263, + 0.029595, + -0.004371, + 0.000389, + -0.004182, + 0.012, + -0.015576, + -2.8e-05, + -0.012556, + 2.5e-05, + 0.054732, + -0.01089, + 0.00547, + 0.015315, + -0.027271, + -0.002516, + 0.003293, + 0.001099, + 0.010529, + -0.00809, + -0.031443, + 0.000377, + -0.031829, + 0.014496, + -0.003451, + -0.031001, + -0.00064, + -0.015198, + -0.001283, + -0.004123, + 0.002237, + 0.002431, + 0.015693, + 0.018091, + -0.015281, + -0.037994, + -0.002521, + -0.001512, + -0.004742, + 0.000613, + -0.012441, + 0.000619, + -0.003316, + 8.3e-05, + -0.013941 ] - }, + } + ], + "T5Model": [ { "shape": [ 2, - 8, - 8, - 4 + 16, + 32 ], "sample": [ - 0.006219, - 0.002704, - 0.003533, - 0.003418, - 0.008001, - 0.005587, - 0.005602, - 0.005686, - 0.008691, - 0.002119, - -0.000714, - 0.001933, - -0.002422, + 0.007004, + 0.007011, + 0.009395, + 0.004649, + -0.014974, + -0.007318, + -0.008844, + 0.000876, + 0.005626, + -0.007449, + 0.003385, + -0.011547, + 0.00152, + 0.027466, + -0.006198, + -0.011448, + 0.011551, + 0.013829, + 0.020342, + 0.00263, + 0.029595, + -0.004371, + 0.000389, + -0.004182, + 0.012, + -0.015576, + -2.8e-05, + -0.012556, + 2.5e-05, + 0.054732, + -0.01089, + 0.00547, + 0.015315, + -0.027271, + -0.002516, + 0.003293, + 0.001099, + 0.010529, + -0.00809, + -0.031443, + 0.000377, + -0.031829, + 0.014496, + -0.003451, + -0.031001, + -0.00064, + -0.015198, + -0.001283, + -0.004123, 0.002237, - 0.001836, - 0.002578, - -0.001956, - 0.00071, - -0.003084, - -0.002698, - -0.005522, - -0.001268, - -0.001686, - -1.2e-05, - -0.004803, - 0.000822, - 0.000538, - 0.013731, - 0.01433, - 0.014584, - 0.015002, - 0.015691, - 0.013323, - 0.015859, - 0.016175, - 0.015331, - 0.004414, - 0.004255, - 0.006079, - 0.002555, - 0.005861, - 0.004837, - 0.005748, - 0.002341, - 0.005368, - -9.1e-05, - 0.001814, - 0.003786, - 0.001644, - 0.001444, - 0.00205, - 0.00168, - 8.1e-05, - 0.000167, - -0.001057, - -0.002688, - -0.000987, - -0.000541, - -0.000336, - -0.002546, - -0.005307, - -0.003004, - -0.006414, - 0.018894 + 0.002431, + 0.015693, + 0.018091, + -0.015281, + -0.037994, + -0.002521, + -0.001512, + -0.004742, + 0.000613, + -0.012441, + 0.000619, + -0.003316, + 8.3e-05, + -0.013941 ] }, { "shape": [ 2, - 16, + 8, 32 ], "sample": [ - -0.012036, - 0.027224, - -0.017103, - 0.027472, - -0.020236, - 0.000332, - -0.016442, - 0.015828, - -0.018295, - 0.014603, - -0.005828, - 0.016616, - -0.00709, - -0.005843, - 0.009211, - 0.000567, - 0.049181, - -0.012505, - 0.000328, - -0.009735, - -0.003474, - -0.007774, - -0.024105, - -0.00845, - -0.020531, - -0.008289, - -0.052391, - 0.003785, - -0.052397, - 0.003749, - -0.019434, - 0.013164, - -0.022706, - 0.014415, - 0.028551, - -0.025012, - 0.017438, - -0.033488, - 0.051644, - -0.027285, - 0.047169, - -0.013412, - 0.049211, - 0.009689, - 0.039967, - 0.02768, - 0.015537, - -0.001814, - -0.01621, - -0.001629, - -0.021368, - 0.015966, - -0.035636, - 0.013777, - -0.038478, - 0.024581, - 0.014664, - 0.015239, - 0.014259, - -0.002774, - 0.006665, - -0.001414, - -0.025644, - 0.009934 + -0.012666, + 0.008346, + -0.012532, + 0.025897, + -0.012666, + 0.008346, + -0.012532, + 0.025897, + -0.012666, + -0.008003, + 0.011512, + 0.020149, + -0.022597, + -0.008003, + 0.011512, + 0.020149, + -0.022597, + -0.008003, + 0.000519, + 0.005653, + -0.008304, + -8.3e-05, + 0.000519, + 0.005653, + -0.008304, + -8.3e-05, + 0.000519, + -0.005227, + -0.026162, + 0.014672, + -0.008533, + -0.005227, + -0.026154, + 0.014576, + -0.008315, + -0.005249, + -0.015154, + -0.005074, + -0.000946, + 0.010883, + -0.015154, + -0.005074, + -0.000946, + 0.010883, + -0.015154, + -0.024924, + 0.01542, + -0.004209, + -0.007258, + -0.024924, + 0.01542, + -0.004209, + -0.007258, + -0.024924, + -0.002469, + 0.002147, + -0.007502, + -0.002374, + -0.002469, + 0.002147, + -0.007502, + -0.002374, + -0.002469, + 0.002716 ] } ], - "StableDiffusion3Model": [ + "T5QnA": [ { "shape": [ 2, - 16, - 16, - 3 + 16 ], "sample": [ - 0.014955, - 0.014535, - 0.015039, - -0.01294, - -0.015078, - -0.01354, - 0.012603, - 0.014234, - 0.013028, - 0.014104, - 0.014657, - -0.013961, - -0.01397, - -0.013705, - 0.014085, - 0.014107, - 0.01454, - 0.014684, - 0.014214, - 0.01473, - -0.014538, - -0.013919, - 0.013891, - 0.013894, - 0.013345, - 0.014624, - 0.015308, - 0.015109, - -0.013916, - -0.013549, - -0.013454, - 0.015281, - 0.010778, - 0.014866, - 0.015276, - 0.014992, - -0.014008, - -0.013972, - -0.013274, - 0.014006, - 0.014443, - 0.014339, - 0.013793, - 0.014472, - -0.013407, - -0.013649, - -0.013639, - 0.014388, - 0.014662, - 0.014543, - 0.015311, - 0.014633, - 0.013674, - -0.01387, - -0.01422, - 0.013761, - 0.014448, - 0.013939, - 0.01379, - 0.011211, - 0.013861, - -0.010878, - -0.012983, - 0.013007 + -0.036581, + -0.036534, + -0.033987, + -0.033718, + -0.032503, + -0.032634, + -0.036825, + -0.033032, + -0.033208, + -0.036711, + -0.035561, + -0.033981, + -0.037899, + -0.030949, + -0.038119, + -0.034084, + -0.036592, + -0.032261, + -0.03053, + -0.029801, + -0.031892, + -0.03281, + -0.035399, + -0.036351, + -0.034096, + -0.02928, + -0.035885, + -0.036726, + -0.032882, + -0.03446, + -0.03706, + -0.033631 ] }, { "shape": [ 2, - 8, - 8, - 8 + 16 ], "sample": [ - 0.017274, - 0.019562, - 0.019478, - 0.019517, - 0.016443, - 0.035483, - 0.035683, - 0.035216, - 0.030619, - -0.011854, - -0.012384, - -0.011986, - -0.008521, - 0.008372, - 0.008306, - 0.008393, - 0.006177, - 0.009252, - 0.009079, - 0.008854, - 0.009121, - -0.050123, - -0.050035, - -0.050299, - -0.042257, - -0.050059, - 0.019482, - 0.019644, - 0.015013, + 0.013309, + 0.016126, + 0.015854, + 0.015916, + 0.014135, + 0.016762, + 0.015658, + 0.016548, + 0.016202, + 0.014686, + 0.014325, + 0.011945, + 0.013924, + 0.014791, 0.013524, - 0.016211, - 0.016645, - 0.0115, - 0.018172, - 0.019602, - 0.017579, - 0.015945, - 0.01631, - 0.035517, - 0.037192, - 0.036485, - 0.035644, - -0.012204, - -0.015612, - -0.012235, - -0.012165, - -0.012408, - 0.012303, - 0.009171, - 0.008228, - 0.008402, - 0.00942, - 0.008764, - 0.00924, - 0.009194, - -0.047951, - -0.049781, - -0.049926, - -0.050071, - 0.018187, - 0.013454, - 0.013745, - 0.013502, - 0.018325 + 0.01577, + 0.013307, + 0.016444, + 0.014633, + 0.013717, + 0.013196, + 0.014541, + 0.01449, + 0.014696, + 0.013105, + 0.015376, + 0.015856, + 0.014193, + 0.013698, + 0.015209, + 0.013511, + 0.016833 ] - }, + } + ], + "T5SequenceClassify": [ { "shape": [ 2, - 8, - 8, 4 ], "sample": [ - 0.012091, - 0.097546, - 0.1286, - 0.065201, - -0.030095, - 0.150227, - 0.080477, - -0.046941, - 0.100516, - -0.039675, - -0.109624, - -0.035173, - -0.018982, - 0.017097, - 0.0496, - -0.016586, - -0.028893, - -0.045285, - 0.017648, - -0.007875, - -0.051069, - 0.043383, - 0.113692, - -0.142047, - -0.113984, - -0.090202, - -0.088588, - 0.136368, - 0.105976, - 0.16715, - 0.126743, - 0.109739, - -0.003245, - 0.006413, - 0.211325, - 0.010501, - 0.100809, - -0.026821, - -0.011584, - -0.063026, - 0.151531, - 0.07651, - 0.149342, - 0.198445, - -0.054413, - 0.071169, - 0.004813, - -0.014978, - 0.154012, - 0.049127, - -0.056784, - 0.082541, - 0.070801, - 0.150369, - 0.045853, - -0.040565, - 0.095416, - -0.002005, - 0.055735, - 0.052691, - -0.016378, - 0.060451, - 0.060556, - 0.000441 + 0.000502, + 0.015776, + -0.036908, + -0.005068, + 0.000508, + 0.015436, + -0.03677, + -0.005324 ] - }, + } + ], + "T5TokenClassify": [ { "shape": [ 2, - 24 + 16, + 4 ], "sample": [ - 0.000567, - -0.002695, - -0.000344, - 0.001872, - -0.000507, - -0.002267, - -0.003681, - 0.002676, - 0.004716, - -0.00181, - 0.001028, - 0.001284, - -0.000525, - -0.000201, - 0.001415, - 0.002244, - -0.000562, - 0.001038, - 0.00074, - 0.002248, - 0.000935, - 0.002127, - -0.000101, - 0.000829, - 0.000567, - -0.002695, - -0.000344, - 0.001872, - -0.000507, - -0.002267, - -0.003681, - 0.002676, - 0.004716, - -0.00181, - 0.001028, - 0.001284, - -0.000525, - -0.000201, - 0.001415, - 0.002244, - -0.000562, - 0.001038, - 0.00074, - 0.002248, - 0.000935, - 0.002127, - -0.000101, - 0.000829 + -0.019692, + -0.057709, + -0.023285, + -0.059521, + -0.022501, + -0.062189, + -0.01735, + -0.058669, + -0.01329, + -0.055911, + -0.021385, + -0.060795, + -0.01553, + -0.057615, + -0.021567, + -0.053714, + -0.016597, + -0.056771, + -0.017363, + -0.054833, + -0.019004, + -0.054143, + -0.019093, + -0.055566, + -0.020616, + -0.052903, + -0.020262, + -0.051335, + -0.01721, + -0.049579, + -0.016186, + -0.051587, + -0.01902, + -0.057753, + -0.016269, + -0.053339, + -0.019398, + -0.056683, + -0.020809, + -0.05843, + -0.017729, + -0.055149, + -0.020295, + -0.053454, + -0.015878, + -0.054542, + -0.017298, + -0.053031, + -0.016035, + -0.06077, + -0.01906, + -0.056622, + -0.017711, + -0.054803, + -0.019422, + -0.058055, + -0.015465, + -0.055705, + -0.019656, + -0.053127, + -0.016712, + -0.055788, + -0.015164, + 0.027239 + ] + } + ], + "TableTransformerDetect": [ + { + "shape": [ + 2, + 10, + 3 + ], + "sample": [ + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897, + -0.035959, + -0.025007, + 0.029897 ] }, { "shape": [ 2, - 16, - 24 + 10, + 4 ], "sample": [ - 0.01991, - 0.004829, - 0.01725, - 0.038287, - -0.006734, - 0.004269, - 0.045626, - 0.02678, - 0.05008, - 0.019808, - 0.082032, - 0.008582, - 0.014972, - 0.010277, - 0.034913, - 0.022746, - 0.015038, - 0.01132, - -0.021046, - -0.056212, - -0.001057, - -0.067862, - -0.024714, - -0.057371, - -0.035649, - -0.058719, - -0.052335, - -0.043667, - -0.032667, - -0.017269, - 0.023394, - -0.025872, - 0.037565, - -0.039031, - 0.032609, - -0.011299, - 0.029473, - 0.035091, - 0.006203, - 0.019498, - 0.028072, - 0.015388, - -0.055906, - 0.0522, - -0.031006, - 0.028746, - -0.001313, - -0.05912, - -0.015437, - -0.072604, - 0.00309, - -0.052105, - -0.060652, - 0.002501, - -0.005301, - -0.016556, - 0.004076, - -0.030116, - 0.015446, - -0.040764, - 0.025618, - -0.082329, - 0.017716, - 0.019545 + 0.498714, + 0.503207, + 0.503151, + 0.506022, + 0.503207, + 0.503151, + 0.506022, + 0.498714, + 0.503151, + 0.506022, + 0.498714, + 0.503207, + 0.506022, + 0.498714, + 0.503207, + 0.503151, + 0.498714, + 0.503207, + 0.503151, + 0.506022, + 0.503207, + 0.503151, + 0.506022, + 0.498714, + 0.503151, + 0.506022, + 0.498714, + 0.503207, + 0.506022, + 0.498714, + 0.503207, + 0.503151, + 0.498714, + 0.503207, + 0.503151, + 0.506022, + 0.503207, + 0.503151, + 0.506022, + 0.498714, + 0.503151, + 0.506022, + 0.498714, + 0.503207, + 0.506022, + 0.498714, + 0.503207, + 0.503151, + 0.498714, + 0.503207, + 0.503151, + 0.506022, + 0.503207, + 0.503151, + 0.506022, + 0.498714, + 0.503151, + 0.506022, + 0.498714, + 0.503207, + 0.506022, + 0.498714, + 0.503207, + 0.506022 ] } ], - "StableDiffusion3TextToImage": [ + "Tipsv2DptDensePredict": [ { "shape": [ 2, 16, - 16, - 3 + 16 ], "sample": [ - 0.014955, - 0.014535, - 0.015039, - -0.01294, - -0.015078, - -0.01354, - 0.012603, - 0.014234, - 0.013028, - 0.014104, - 0.014657, - -0.013961, - -0.01397, - -0.013705, - 0.014085, - 0.014107, - 0.01454, - 0.014684, - 0.014214, - 0.01473, - -0.014538, - -0.013919, - 0.013891, - 0.013894, - 0.013345, - 0.014624, - 0.015308, - 0.015109, - -0.013916, - -0.013549, - -0.013454, - 0.015281, - 0.010778, - 0.014866, - 0.015276, - 0.014992, - -0.014008, - -0.013972, - -0.013274, - 0.014006, - 0.014443, - 0.014339, - 0.013793, - 0.014472, - -0.013407, - -0.013649, - -0.013639, - 0.014388, - 0.014662, - 0.014543, - 0.015311, - 0.014633, - 0.013674, - -0.01387, - -0.01422, - 0.013761, - 0.014448, - 0.013939, - 0.01379, - 0.011211, - 0.013861, - -0.010878, - -0.012983, - 0.013007 + 0.01559, + 0.015544, + 0.015598, + 0.015541, + 0.015599, + 0.015541, + 0.015599, + 0.015542, + 0.015599, + 0.015542, + 0.015541, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015544, + 0.015544, + 0.015544, + 0.015544, + 0.015541, + 0.015541, + 0.015541, + 0.015541, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015544, + 0.015526 ] }, { "shape": [ 2, - 8, - 8, - 8 + 16, + 16, + 5 ], "sample": [ - 0.017274, - 0.019562, - 0.019478, - 0.019517, - 0.016443, - 0.035483, - 0.035683, - 0.035216, - 0.030619, - -0.011854, - -0.012384, - -0.011986, - -0.008521, - 0.008372, - 0.008306, - 0.008393, - 0.006177, - 0.009252, - 0.009079, - 0.008854, - 0.009121, - -0.050123, - -0.050035, - -0.050299, - -0.042257, - -0.050059, - 0.019482, - 0.019644, - 0.015013, - 0.013524, - 0.016211, - 0.016645, - 0.0115, - 0.018172, - 0.019602, - 0.017579, - 0.015945, - 0.01631, - 0.035517, - 0.037192, - 0.036485, - 0.035644, - -0.012204, - -0.015612, - -0.012235, - -0.012165, - -0.012408, - 0.012303, - 0.009171, - 0.008228, - 0.008402, - 0.00942, - 0.008764, - 0.00924, - 0.009194, - -0.047951, - -0.049781, - -0.049926, - -0.050071, - 0.018187, - 0.013454, - 0.013745, - 0.013502, - 0.018325 + -0.010536, + -0.010721, + 0.01322, + 0.013027, + -0.009794, + -0.039424, + -0.039745, + -0.023638, + -0.023823, + -0.010883, + 0.01303, + 0.013027, + -0.009837, + -0.039424, + -0.039421, + -0.023638, + -0.023642, + -0.010883, + 0.013028, + 0.013027, + -0.009836, + -0.039424, + -0.039423, + -0.023638, + -0.023639, + -0.010883, + 0.013028, + 0.013028, + -0.009836, + -0.009836, + -0.039418, + -0.023689, + -0.023834, + -0.010721, + 0.013027, + 0.013027, + -0.009836, + -0.009836, + -0.039424, + -0.023638, + -0.023638, + -0.010883, + 0.013027, + 0.013026, + -0.009836, + -0.009836, + -0.039424, + -0.023638, + -0.023638, + -0.010884, + -0.010883, + 0.013024, + -0.009836, + -0.009835, + -0.039424, + -0.023637, + -0.023638, + -0.010755, + -0.010882, + 0.012814, + -0.009836, + -0.009669, + -0.039418, + -0.023813 ] - }, + } + ], + "Tipsv2DptDepthEstimation": [ { "shape": [ 2, - 8, - 8, - 4 + 16, + 16 ], "sample": [ - 0.012091, - 0.097546, - 0.1286, - 0.065201, - -0.030095, - 0.150227, - 0.080477, - -0.046941, - 0.100516, - -0.039675, - -0.109624, - -0.035173, - -0.018982, - 0.017097, - 0.0496, - -0.016586, - -0.028893, - -0.045285, - 0.017648, - -0.007875, - -0.051069, - 0.043383, - 0.113692, - -0.142047, - -0.113984, - -0.090202, - -0.088588, - 0.136368, - 0.105976, - 0.16715, - 0.126743, - 0.109739, - -0.003245, - 0.006413, - 0.211325, - 0.010501, - 0.100809, - -0.026821, - -0.011584, - -0.063026, - 0.151531, - 0.07651, - 0.149342, - 0.198445, - -0.054413, - 0.071169, - 0.004813, - -0.014978, - 0.154012, - 0.049127, - -0.056784, - 0.082541, - 0.070801, - 0.150369, - 0.045853, - -0.040565, - 0.095416, - -0.002005, - 0.055735, - 0.052691, - -0.016378, - 0.060451, - 0.060556, - 0.000441 + 0.01559, + 0.015544, + 0.015598, + 0.015541, + 0.015599, + 0.015541, + 0.015599, + 0.015542, + 0.015599, + 0.015542, + 0.015541, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015544, + 0.015544, + 0.015544, + 0.015544, + 0.015541, + 0.015541, + 0.015541, + 0.015541, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015542, + 0.015544, + 0.015526 ] - }, + } + ], + "Tipsv2DptSemanticSegment": [ { "shape": [ 2, - 24 + 16, + 16, + 5 ], "sample": [ - 0.000567, - -0.002695, - -0.000344, - 0.001872, - -0.000507, - -0.002267, - -0.003681, - 0.002676, - 0.004716, - -0.00181, - 0.001028, - 0.001284, - -0.000525, - -0.000201, - 0.001415, - 0.002244, - -0.000562, - 0.001038, - 0.00074, - 0.002248, - 0.000935, - 0.002127, - -0.000101, - 0.000829, - 0.000567, - -0.002695, - -0.000344, - 0.001872, - -0.000507, - -0.002267, - -0.003681, - 0.002676, - 0.004716, - -0.00181, - 0.001028, - 0.001284, - -0.000525, - -0.000201, - 0.001415, - 0.002244, - -0.000562, - 0.001038, - 0.00074, - 0.002248, - 0.000935, - 0.002127, - -0.000101, - 0.000829 + -0.010536, + -0.010721, + 0.01322, + 0.013027, + -0.009794, + -0.039424, + -0.039745, + -0.023638, + -0.023823, + -0.010883, + 0.01303, + 0.013027, + -0.009837, + -0.039424, + -0.039421, + -0.023638, + -0.023642, + -0.010883, + 0.013028, + 0.013027, + -0.009836, + -0.039424, + -0.039423, + -0.023638, + -0.023639, + -0.010883, + 0.013028, + 0.013028, + -0.009836, + -0.009836, + -0.039418, + -0.023689, + -0.023834, + -0.010721, + 0.013027, + 0.013027, + -0.009836, + -0.009836, + -0.039424, + -0.023638, + -0.023638, + -0.010883, + 0.013027, + 0.013026, + -0.009836, + -0.009836, + -0.039424, + -0.023638, + -0.023638, + -0.010884, + -0.010883, + 0.013024, + -0.009836, + -0.009835, + -0.039424, + -0.023637, + -0.023638, + -0.010755, + -0.010882, + 0.012814, + -0.009836, + -0.009669, + -0.039418, + -0.023813 ] - }, + } + ], + "Tipsv2Model": [ { "shape": [ 2, - 16, - 24 + 32 ], "sample": [ - 0.01991, - 0.004829, - 0.01725, - 0.038287, - -0.006734, - 0.004269, - 0.045626, - 0.02678, - 0.05008, - 0.019808, - 0.082032, - 0.008582, - 0.014972, - 0.010277, - 0.034913, - 0.022746, - 0.015038, - 0.01132, - -0.021046, - -0.056212, - -0.001057, - -0.067862, - -0.024714, - -0.057371, - -0.035649, - -0.058719, - -0.052335, - -0.043667, - -0.032667, - -0.017269, - 0.023394, - -0.025872, - 0.037565, - -0.039031, - 0.032609, - -0.011299, - 0.029473, - 0.035091, - 0.006203, - 0.019498, - 0.028072, - 0.015388, - -0.055906, - 0.0522, - -0.031006, - 0.028746, - -0.001313, - -0.05912, - -0.015437, - -0.072604, - 0.00309, - -0.052105, - -0.060652, - 0.002501, - -0.005301, - -0.016556, - 0.004076, - -0.030116, - 0.015446, - -0.040764, - 0.025618, - -0.082329, - 0.017716, - 0.019545 + -0.204265, + -0.059741, + 0.026511, + -0.164984, + 0.046244, + -0.079875, + -0.577878, + 0.032957, + -0.457186, + -0.031679, + 0.134449, + -0.152016, + -0.073435, + 0.033245, + -0.147961, + -0.169231, + -0.114642, + -0.173875, + -0.086014, + 0.194842, + -0.035061, + -0.001086, + 0.069644, + 0.302139, + -0.014048, + -0.114194, + 0.130663, + -0.098736, + -0.06882, + 0.182741, + 0.03171, + -0.098478, + -0.204267, + -0.059737, + 0.026501, + -0.164981, + 0.046248, + -0.079876, + -0.577906, + 0.032955, + -0.457162, + -0.031677, + 0.134442, + -0.152013, + -0.073436, + 0.033242, + -0.147959, + -0.169224, + -0.11463, + -0.173872, + -0.086019, + 0.194818, + -0.035056, + -0.001096, + 0.069653, + 0.302146, + -0.014045, + -0.114194, + 0.13069, + -0.098734, + -0.068817, + 0.182733, + 0.031709, + -0.098476 ] - } - ], - "StableDiffusion3_5Model": [ + }, { "shape": [ 2, - 16, - 16, - 3 + 2 ], "sample": [ - 0.014955, - 0.014535, - 0.015039, - -0.01294, - -0.015078, - -0.01354, - 0.012603, - 0.014234, - 0.013028, - 0.014104, - 0.014657, - -0.013961, - -0.01397, - -0.013705, - 0.014085, - 0.014107, - 0.01454, - 0.014684, - 0.014214, - 0.01473, - -0.014538, - -0.013919, - 0.013891, - 0.013894, - 0.013345, - 0.014624, - 0.015308, - 0.015109, - -0.013916, - -0.013549, - -0.013454, - 0.015281, - 0.010778, - 0.014866, - 0.015276, - 0.014992, - -0.014008, - -0.013972, - -0.013274, - 0.014006, - 0.014443, - 0.014339, - 0.013793, - 0.014472, - -0.013407, - -0.013649, - -0.013639, - 0.014388, - 0.014662, - 0.014543, - 0.015311, - 0.014633, - 0.013674, - -0.01387, - -0.01422, - 0.013761, - 0.014448, - 0.013939, - 0.01379, - 0.011211, - 0.013861, - -0.010878, - -0.012983, - 0.013007 + 33.656609, + 33.656609, + 33.657059, + 33.657059 ] }, { "shape": [ 2, - 8, - 8, - 8 + 2 ], "sample": [ - 0.017274, - 0.019562, - 0.019478, - 0.019517, - 0.016443, - 0.035483, - 0.035683, - 0.035216, - 0.030619, - -0.011854, - -0.012384, - -0.011986, - -0.008521, - 0.008372, - 0.008306, - 0.008393, - 0.006177, - 0.009252, - 0.009079, - 0.008854, - 0.009121, - -0.050123, - -0.050035, - -0.050299, - -0.042257, - -0.050059, - 0.019482, - 0.019644, - 0.015013, - 0.013524, - 0.016211, - 0.016645, - 0.0115, - 0.018172, - 0.019602, - 0.017579, - 0.015945, - 0.01631, - 0.035517, - 0.037192, - 0.036485, - 0.035644, - -0.012204, - -0.015612, - -0.012235, - -0.012165, - -0.012408, - 0.012303, - 0.009171, - 0.008228, - 0.008402, - 0.00942, - 0.008764, - 0.00924, - 0.009194, - -0.047951, - -0.049781, - -0.049926, - -0.050071, - 0.018187, - 0.013454, - 0.013745, - 0.013502, - 0.018325 + 33.656609, + 33.657059, + 33.656609, + 33.657059 ] }, { "shape": [ 2, - 8, - 8, - 4 + 32 ], "sample": [ - 0.01323, - 0.098254, - 0.129669, - 0.066239, - -0.028639, - 0.150828, - 0.081944, - -0.045938, - 0.101357, - -0.039674, - -0.10952, - -0.035309, - -0.019402, - 0.016394, - 0.048952, - -0.017228, - -0.029025, - -0.045206, - 0.018156, - -0.007403, - -0.051539, - 0.042579, - 0.113596, - -0.142042, - -0.113807, - -0.089958, - -0.088259, - 0.136318, - 0.105069, - 0.166323, - 0.125724, - 0.108803, - -0.003717, - 0.005988, - 0.211244, - 0.010028, - 0.100874, - -0.026819, - -0.011773, - -0.063357, - 0.151896, - 0.077693, - 0.151372, - 0.199062, - -0.054222, - 0.072171, - 0.006164, - -0.014138, - 0.153723, - 0.049434, - -0.056111, - 0.083136, - 0.071386, - 0.150738, - 0.046609, - -0.039163, - 0.09659, - -0.00176, - 0.055724, - 0.053063, - -0.015067, - 0.061407, - 0.061014, - 0.002211 + 0.231413, + -0.231878, + 0.036957, + -0.112705, + 0.224442, + -0.284503, + -0.020773, + -0.191486, + -0.101979, + -0.30212, + -0.231222, + 0.084227, + 0.094925, + 0.060053, + -0.129362, + 0.210374, + -0.067982, + -0.058104, + 0.185936, + -0.040592, + 0.028479, + 0.224797, + 0.177719, + 0.404719, + 0.221314, + -0.168479, + -0.003695, + -0.149697, + 0.026421, + -0.147367, + 0.080145, + -0.234062, + 0.231413, + -0.231878, + 0.036957, + -0.112705, + 0.224442, + -0.284503, + -0.020773, + -0.191486, + -0.101979, + -0.30212, + -0.231222, + 0.084227, + 0.094925, + 0.060053, + -0.129362, + 0.210374, + -0.067982, + -0.058104, + 0.185936, + -0.040592, + 0.028479, + 0.224797, + 0.177719, + 0.404719, + 0.221314, + -0.168479, + -0.003695, + -0.149697, + 0.026421, + -0.147367, + 0.080145, + -0.234062 ] - }, + } + ], + "Tipsv2TextModel": [ { "shape": [ 2, - 24 + 16, + 32 ], "sample": [ - 0.000567, - -0.002695, - -0.000344, - 0.001872, - -0.000507, - -0.002267, - -0.003681, - 0.002676, - 0.004716, - -0.00181, - 0.001028, - 0.001284, - -0.000525, - -0.000201, - 0.001415, - 0.002244, - -0.000562, - 0.001038, - 0.00074, - 0.002248, - 0.000935, - 0.002127, - -0.000101, - 0.000829, - 0.000567, - -0.002695, - -0.000344, - 0.001872, - -0.000507, - -0.002267, - -0.003681, - 0.002676, - 0.004716, - -0.00181, - 0.001028, - 0.001284, - -0.000525, - -0.000201, - 0.001415, - 0.002244, - -0.000562, - 0.001038, - 0.00074, - 0.002248, - 0.000935, - 0.002127, - -0.000101, - 0.000829 + 0.033123, + -0.010784, + 0.0338, + -0.013429, + 0.033636, + -0.008718, + -0.032206, + -0.008677, + -0.033497, + 0.026862, + 0.007951, + 0.026524, + 0.005263, + -0.008706, + -0.016543, + -0.003263, + -0.018358, + 0.004191, + 0.032402, + 0.00429, + 0.031997, + 0.030361, + -0.041748, + 0.032885, + -0.040408, + 0.042542, + -0.003329, + 0.025781, + -0.002476, + 0.028777, + -0.019681, + 0.058398, + -0.022705, + 0.059073, + -0.015275, + 0.029098, + -0.013964, + 0.033865, + -0.050568, + -0.025017, + -0.046109, + -0.0233, + -0.03752, + 0.003248, + -0.035767, + -0.002517, + -0.033605, + -0.019767, + 0.0116, + -0.024619, + 0.012422, + 0.000619, + 0.016388, + 0.002959, + 0.003817, + -0.020142, + 0.008776, + -0.01486, + 0.008877, + 0.012328, + -0.026626, + 0.009385, + -0.017078, + -0.035824 ] }, { "shape": [ 2, - 16, - 24 + 32 ], "sample": [ - 0.01991, - 0.004829, - 0.01725, - 0.038287, - -0.006734, - 0.004269, - 0.045626, - 0.02678, - 0.05008, - 0.019808, - 0.082032, - 0.008582, - 0.014972, - 0.010277, - 0.034913, - 0.022746, - 0.015038, - 0.01132, - -0.021046, - -0.056212, - -0.001057, - -0.067862, - -0.024714, - -0.057371, - -0.035649, - -0.058719, - -0.052335, - -0.043667, - -0.032667, - -0.017269, - 0.023394, - -0.025872, - 0.037565, - -0.039031, - 0.032609, - -0.011299, - 0.029473, - 0.035091, - 0.006203, - 0.019498, - 0.028072, - 0.015388, - -0.055906, - 0.0522, - -0.031006, - 0.028746, - -0.001313, - -0.05912, - -0.015437, - -0.072604, - 0.00309, - -0.052105, - -0.060652, - 0.002501, - -0.005301, - -0.016556, - 0.004076, - -0.030116, - 0.015446, - -0.040764, - 0.025618, - -0.082329, - 0.017716, - 0.019545 + 0.033625, + -0.033692, + 0.00537, + -0.016376, + 0.032612, + -0.041339, + -0.003018, + -0.027823, + -0.014818, + -0.043898, + -0.033597, + 0.012238, + 0.013793, + 0.008726, + -0.018796, + 0.030568, + -0.009878, + -0.008443, + 0.027017, + -0.005898, + 0.004138, + 0.032663, + 0.025823, + 0.058806, + 0.032157, + -0.02448, + -0.000537, + -0.021751, + 0.003839, + -0.021413, + 0.011645, + -0.03401, + 0.033625, + -0.033692, + 0.00537, + -0.016376, + 0.032612, + -0.041339, + -0.003018, + -0.027823, + -0.014818, + -0.043898, + -0.033597, + 0.012238, + 0.013793, + 0.008726, + -0.018796, + 0.030568, + -0.009878, + -0.008443, + 0.027017, + -0.005898, + 0.004138, + 0.032663, + 0.025823, + 0.058806, + 0.032157, + -0.02448, + -0.000537, + -0.021751, + 0.003839, + -0.021413, + 0.011645, + -0.03401 ] } ], - "StableDiffusion3_5TextToImage": [ + "Tipsv2VisionModel": [ { "shape": [ 2, - 16, - 16, - 3 + 6, + 32 ], "sample": [ - 0.014955, - 0.014535, - 0.015039, - -0.01294, - -0.015078, - -0.01354, - 0.012603, - 0.014234, - 0.013028, - 0.014104, - 0.014657, - -0.013961, - -0.01397, - -0.013705, - 0.014085, - 0.014107, - 0.01454, - 0.014684, - 0.014214, - 0.01473, - -0.014538, - -0.013919, - 0.013891, - 0.013894, - 0.013345, - 0.014624, - 0.015308, - 0.015109, - -0.013916, - -0.013549, - -0.013454, - 0.015281, - 0.010778, - 0.014866, - 0.015276, - 0.014992, - -0.014008, - -0.013972, - -0.013274, - 0.014006, - 0.014443, - 0.014339, - 0.013793, - 0.014472, - -0.013407, - -0.013649, - -0.013639, - 0.014388, - 0.014662, - 0.014543, - 0.015311, - 0.014633, - 0.013674, - -0.01387, - -0.01422, - 0.013761, - 0.014448, - 0.013939, - 0.01379, - 0.011211, - 0.013861, - -0.010878, - -0.012983, - 0.013007 + -0.035782, + -0.10123, + -0.012864, + -0.015068, + -0.002461, + 0.005555, + 0.010684, + 0.039316, + -0.048693, + 0.008908, + -0.024619, + 0.011059, + -0.073175, + -0.01496, + -0.068544, + -0.016729, + -0.007399, + 0.005127, + -0.013034, + 0.007938, + 0.025508, + -0.018023, + -0.012591, + -0.02752, + -0.031995, + 0.021785, + 0.006093, + 0.01765, + 0.056828, + -0.042498, + -0.018075, + -0.056264, + 0.004642, + -0.080085, + -0.025919, + -0.006141, + 0.022894, + -0.013294, + 0.009124, + 0.007434, + 0.098242, + 0.033887, + -0.015847, + -0.028111, + -0.027522, + -0.030784, + 0.029121, + 0.101468, + -0.011032, + 0.003873, + -0.030811, + 0.011487, + -0.043628, + -0.090914, + 0.007013, + -0.036377, + -0.030575, + -0.010353, + -0.037672, + 0.03382, + -0.001482, + -0.017995, + 0.02622, + -0.015573 ] }, { "shape": [ 2, - 8, - 8, - 8 + 32 ], "sample": [ - 0.017274, - 0.019562, - 0.019478, - 0.019517, - 0.016443, - 0.035483, - 0.035683, - 0.035216, - 0.030619, - -0.011854, - -0.012384, - -0.011986, - -0.008521, - 0.008372, - 0.008306, - 0.008393, - 0.006177, - 0.009252, - 0.009079, - 0.008854, - 0.009121, - -0.050123, - -0.050035, - -0.050299, - -0.042257, - -0.050059, - 0.019482, - 0.019644, - 0.015013, - 0.013524, - 0.016211, - 0.016645, - 0.0115, - 0.018172, - 0.019602, - 0.017579, - 0.015945, - 0.01631, - 0.035517, - 0.037192, - 0.036485, - 0.035644, - -0.012204, - -0.015612, - -0.012235, - -0.012165, - -0.012408, - 0.012303, - 0.009171, - 0.008228, - 0.008402, - 0.00942, - 0.008764, - 0.00924, - 0.009194, - -0.047951, - -0.049781, - -0.049926, - -0.050071, - 0.018187, - 0.013454, - 0.013745, - 0.013502, - 0.018325 + -0.035782, + -0.010465, + 0.004644, + -0.028901, + 0.008101, + -0.013992, + -0.10123, + 0.005773, + -0.080088, + -0.005549, + 0.023552, + -0.026629, + -0.012864, + 0.005824, + -0.025919, + -0.029645, + -0.020083, + -0.030459, + -0.015068, + 0.034131, + -0.006142, + -0.00019, + 0.0122, + 0.052927, + -0.002461, + -0.020004, + 0.022889, + -0.017296, + -0.012056, + 0.032012, + 0.005555, + -0.017251, + -0.035783, + -0.010465, + 0.004642, + -0.028901, + 0.008102, + -0.013993, + -0.101237, + 0.005773, + -0.080085, + -0.005549, + 0.023551, + -0.026629, + -0.012864, + 0.005823, + -0.025919, + -0.029644, + -0.020081, + -0.030459, + -0.015069, + 0.034128, + -0.006141, + -0.000192, + 0.012202, + 0.052929, + -0.00246, + -0.020004, + 0.022894, + -0.017296, + -0.012055, + 0.032011, + 0.005555, + -0.017251 ] - }, + } + ], + "VGGImageClassify": [ { "shape": [ 2, - 8, - 8, - 4 + 1000 ], "sample": [ - 0.01323, - 0.098254, - 0.129669, - 0.066239, - -0.028639, - 0.150828, - 0.081944, - -0.045938, - 0.101357, - -0.039674, - -0.10952, - -0.035309, - -0.019402, - 0.016394, - 0.048952, - -0.017228, - -0.029025, - -0.045206, - 0.018156, - -0.007403, - -0.051539, - 0.042579, - 0.113596, - -0.142042, - -0.113807, - -0.089958, - -0.088259, - 0.136318, - 0.105069, - 0.166323, - 0.125724, - 0.108803, - -0.003717, - 0.005988, - 0.211244, - 0.010028, - 0.100874, - -0.026819, - -0.011773, - -0.063357, - 0.151896, - 0.077693, - 0.151372, - 0.199062, - -0.054222, - 0.072171, - 0.006164, - -0.014138, - 0.153723, - 0.049434, - -0.056111, - 0.083136, - 0.071386, - 0.150738, - 0.046609, - -0.039163, - 0.09659, - -0.00176, - 0.055724, - 0.053063, - -0.015067, - 0.061407, - 0.061014, - 0.002211 + -0.075346, + 0.09173, + 0.003965, + -0.047696, + 0.005099, + 0.042046, + 0.039386, + -0.008242, + -0.100712, + -0.033049, + -0.077039, + 0.001321, + 0.02788, + 0.014689, + 0.005051, + -0.009848, + 0.001221, + -0.024884, + 0.038048, + -0.040728, + 0.039735, + -0.022384, + 0.008776, + 0.05443, + 0.013342, + -0.089412, + 0.101408, + -0.101833, + -0.098672, + -0.066069, + 0.071206, + 0.016617, + 0.046086, + -0.015095, + -0.079544, + -0.071482, + 0.027208, + 0.058843, + -0.1108, + 0.03061, + 0.098512, + 0.012451, + -0.084589, + 0.013578, + 0.000479, + 0.023356, + 0.059713, + -0.013926, + 0.049146, + -0.121879, + 0.023654, + 0.099224, + -0.010519, + 0.009046, + -0.055002, + 0.007435, + 0.050897, + 0.028916, + -0.005484, + 0.010513, + -0.058174, + -0.022928, + -0.0136, + -0.00607 ] - }, + } + ], + "ViTImageClassify": [ { "shape": [ 2, - 24 + 1000 ], "sample": [ - 0.000567, - -0.002695, - -0.000344, - 0.001872, - -0.000507, - -0.002267, - -0.003681, - 0.002676, - 0.004716, - -0.00181, - 0.001028, - 0.001284, - -0.000525, - -0.000201, - 0.001415, - 0.002244, - -0.000562, - 0.001038, - 0.00074, - 0.002248, - 0.000935, - 0.002127, - -0.000101, - 0.000829, - 0.000567, - -0.002695, - -0.000344, - 0.001872, - -0.000507, - -0.002267, - -0.003681, - 0.002676, - 0.004716, - -0.00181, - 0.001028, - 0.001284, - -0.000525, - -0.000201, - 0.001415, - 0.002244, - -0.000562, - 0.001038, - 0.00074, - 0.002248, - 0.000935, - 0.002127, - -0.000101, - 0.000829 + 0.028365, + 0.019775, + 0.01387, + -0.005907, + -0.00608, + 0.030517, + -0.024835, + 0.047901, + 0.014553, + -0.057438, + -0.006771, + -0.025818, + 0.026576, + 0.013906, + -0.00355, + 0.03239, + 0.024765, + 0.001365, + 0.000628, + 0.010015, + 0.012252, + -0.018734, + 0.010178, + 0.043254, + -0.013845, + 0.013418, + 0.001323, + -0.045827, + 0.009029, + -0.001661, + 0.002305, + -0.004402, + 0.006255, + 0.019019, + 0.020271, + -0.026835, + -0.004927, + 0.003636, + -0.015582, + -0.002634, + 0.03241, + -0.011949, + -0.063028, + -0.027115, + 0.029731, + -0.010011, + -0.018461, + -0.00617, + 0.011066, + -0.01429, + 0.019819, + -0.007293, + -0.001256, + -0.033223, + -0.027069, + 0.000452, + 0.037526, + 0.018432, + -0.015729, + -0.020209, + 0.007031, + -0.018771, + -0.016291, + -0.033565 ] - }, + } + ], + "WhisperAudioClassify": [ { "shape": [ 2, - 16, - 24 + 10 ], "sample": [ - 0.01991, - 0.004829, - 0.01725, - 0.038287, - -0.006734, - 0.004269, - 0.045626, - 0.02678, - 0.05008, - 0.019808, - 0.082032, - 0.008582, - 0.014972, - 0.010277, - 0.034913, - 0.022746, - 0.015038, - 0.01132, - -0.021046, - -0.056212, - -0.001057, - -0.067862, - -0.024714, - -0.057371, - -0.035649, - -0.058719, - -0.052335, - -0.043667, - -0.032667, - -0.017269, - 0.023394, - -0.025872, - 0.037565, - -0.039031, - 0.032609, - -0.011299, - 0.029473, - 0.035091, - 0.006203, - 0.019498, - 0.028072, - 0.015388, - -0.055906, - 0.0522, - -0.031006, - 0.028746, - -0.001313, - -0.05912, - -0.015437, - -0.072604, - 0.00309, - -0.052105, - -0.060652, - 0.002501, - -0.005301, - -0.016556, - 0.004076, - -0.030116, - 0.015446, - -0.040764, - 0.025618, - -0.082329, - 0.017716, - 0.019545 + -0.019674, + -0.00448, + -0.05384, + 0.026557, + -0.047316, + -0.004436, + 0.012442, + 0.001493, + -0.031542, + -0.013457, + -0.019705, + -0.004451, + -0.053866, + 0.026555, + -0.047313, + -0.004437, + 0.012436, + 0.001537, + -0.031587, + -0.01345 ] } ], - "StableDiffusionModel": [ + "WhisperConditionalGenerate": [ { "shape": [ 2, - 16, - 16, - 3 + 20, + 64 ], "sample": [ - 0.015141, - 0.014566, - 0.01507, - -0.013042, - -0.014793, - -0.013933, - 0.01303, - 0.014152, - 0.013037, - 0.014178, - 0.014338, - -0.013882, - -0.013874, - -0.013862, - 0.01476, - 0.014129, - 0.014743, - 0.014246, - 0.014257, - 0.014225, - -0.014007, - -0.013831, - 0.014013, - 0.014092, - 0.014013, - 0.014246, - 0.014301, - 0.014293, - -0.013722, - -0.013693, - -0.013301, - 0.014484, - 0.010029, - 0.014481, - 0.015246, - 0.015245, - -0.013927, - -0.013932, - -0.01385, - 0.014118, - 0.014105, - 0.014152, - 0.014243, - 0.01403, - -0.013841, - -0.013774, - -0.013812, - 0.014105, - 0.014097, - 0.014094, - 0.014242, - 0.014448, - 0.014266, - -0.014007, - -0.013837, - 0.014, - 0.014121, - 0.014003, - 0.014259, - 0.011565, - 0.014413, - -0.011314, - -0.01329, - 0.013126 + -0.016339, + 0.014459, + 0.027236, + -0.012696, + 0.003826, + 0.004806, + 0.00172, + 0.018539, + -0.003298, + -0.020976, + 0.037257, + -0.035281, + -0.004034, + 0.007173, + 0.031324, + -0.022888, + -0.013037, + -0.005674, + 0.017987, + 0.011197, + -0.007481, + 0.023746, + -0.024699, + 0.009173, + 0.0112, + -0.031971, + -0.023946, + -0.021795, + 0.052292, + -0.002748, + 0.00709, + 0.000376, + -0.006427, + -0.0098, + -0.010043, + -0.034723, + -0.013803, + 0.017374, + 0.028874, + -0.014603, + -0.010544, + 0.037403, + 0.020164, + -0.035462, + -0.004364, + 0.017992, + 0.101977, + -0.021944, + 0.019578, + -0.026028, + 0.050127, + -0.044843, + -0.005331, + 0.008791, + 0.024862, + 0.00735, + 0.006498, + 0.005705, + -0.001536, + 0.021862, + 0.012463, + -0.021342, + 0.037456, + -0.03329 ] }, { "shape": [ 2, - 8, - 8, - 8 + 5, + 128 ], "sample": [ - 0.051589, - 0.051341, - 0.051341, - 0.05132, - 0.051681, - -0.006162, - -0.006151, - -0.006155, - -0.006046, - 0.006829, - 0.006821, - 0.006841, - 0.006748, - 0.010876, - 0.010889, - 0.010892, - 0.010883, - -0.021278, - -0.02128, - -0.021273, - -0.021104, - -0.005376, - -0.005379, - -0.005368, - -0.005528, - -0.005407, - -0.000975, - -0.000971, - -0.00076, - -0.000734, - -0.043369, - -0.043349, - -0.043402, - -0.043551, - 0.051337, - 0.051297, - 0.05151, - 0.051494, - -0.006147, - -0.006312, - -0.006151, - -0.006165, - 0.006816, - 0.006766, - 0.006782, - 0.006824, - 0.006837, - 0.010968, - 0.010946, - 0.0109, - 0.010906, - -0.021234, - -0.021271, - -0.021282, - -0.021255, - -0.005589, - -0.005439, - -0.005423, - -0.005422, - -0.000794, - -0.00072, - -0.000731, - -0.000732, - -0.043442 + 0.004865, + -0.003695, + 0.002367, + -0.000958, + -0.000965, + -0.000914, + 0.004061, + 5.2e-05, + -0.004466, + 0.000528, + -0.00554, + 0.002936, + 0.010828, + -0.003573, + -0.003714, + -0.004223, + -0.01237, + -0.00467, + -0.001861, + -0.004012, + -0.009534, + -0.005811, + -0.004497, + -0.0022, + 0.003601, + -0.002138, + 0.004155, + 0.009307, + -0.003024, + -0.002535, + -3.1e-05, + 0.004786, + -0.00753, + -0.003172, + 0.002225, + -0.001803, + -9.7e-05, + 0.003097, + 0.001154, + 0.008286, + -0.001478, + 0.002645, + 0.000562, + 0.008003, + -0.000749, + 0.00237, + -0.005276, + -0.004898, + -0.00232, + 0.003043, + -0.001344, + 0.001505, + 0.003906, + -0.001533, + 0.004861, + -0.002786, + 0.003926, + 0.001462, + 0.008712, + 0.0024, + 0.003343, + 0.001272, + 0.000233, + -0.000715 ] - }, + } + ], + "WhisperModel": [ { "shape": [ 2, - 8, - 8, - 4 + 20, + 64 ], "sample": [ - 0.006219, - 0.002704, - 0.003533, - 0.003418, - 0.008001, - 0.005587, - 0.005602, - 0.005686, - 0.008691, - 0.002119, - -0.000714, - 0.001933, - -0.002422, - 0.002237, - 0.001836, - 0.002578, - -0.001956, - 0.00071, - -0.003084, - -0.002698, - -0.005522, - -0.001268, - -0.001686, - -1.2e-05, - -0.004803, - 0.000822, - 0.000538, - 0.013731, - 0.01433, - 0.014584, - 0.015002, - 0.015691, - 0.013323, - 0.015859, - 0.016175, - 0.015331, - 0.004414, - 0.004255, - 0.006079, - 0.002555, - 0.005861, - 0.004837, - 0.005748, - 0.002341, - 0.005368, - -9.1e-05, - 0.001814, - 0.003786, - 0.001644, - 0.001444, - 0.00205, - 0.00168, - 8.1e-05, - 0.000167, - -0.001057, - -0.002688, - -0.000987, - -0.000541, - -0.000336, - -0.002546, - -0.005307, - -0.003004, - -0.006414, - 0.018894 + -0.016339, + 0.014459, + 0.027236, + -0.012696, + 0.003826, + 0.004806, + 0.00172, + 0.018539, + -0.003298, + -0.020976, + 0.037257, + -0.035281, + -0.004034, + 0.007173, + 0.031324, + -0.022888, + -0.013037, + -0.005674, + 0.017987, + 0.011197, + -0.007481, + 0.023746, + -0.024699, + 0.009173, + 0.0112, + -0.031971, + -0.023946, + -0.021795, + 0.052292, + -0.002748, + 0.00709, + 0.000376, + -0.006427, + -0.0098, + -0.010043, + -0.034723, + -0.013803, + 0.017374, + 0.028874, + -0.014603, + -0.010544, + 0.037403, + 0.020164, + -0.035462, + -0.004364, + 0.017992, + 0.101977, + -0.021944, + 0.019578, + -0.026028, + 0.050127, + -0.044843, + -0.005331, + 0.008791, + 0.024862, + 0.00735, + 0.006498, + 0.005705, + -0.001536, + 0.021862, + 0.012463, + -0.021342, + 0.037456, + -0.03329 ] }, { "shape": [ 2, - 16, - 32 + 5, + 128 ], "sample": [ - -0.012037, - 0.027224, - -0.017103, - 0.027472, - -0.020236, - 0.000331, - -0.01644, - 0.015825, - -0.018293, - 0.014603, - -0.005827, - 0.016616, - -0.007089, - -0.005842, - 0.009213, - 0.000568, - 0.049181, - -0.012505, - 0.000328, - -0.009735, - -0.003473, - -0.007774, - -0.024105, - -0.00845, - -0.020531, - -0.008289, - -0.052391, - 0.003785, - -0.052397, - 0.003749, - -0.019436, - 0.013164, - -0.022708, - 0.014415, - 0.028549, - -0.025015, - 0.017437, - -0.03349, - 0.051644, - -0.027284, - 0.04717, - -0.013412, - 0.04921, - 0.009694, - 0.039967, - 0.027683, - 0.015538, - -0.001814, - -0.016217, - -0.001629, - -0.021373, - 0.015966, - -0.035635, - 0.013777, - -0.038477, - 0.024585, - 0.014664, - 0.015243, - 0.014259, - -0.002774, - 0.006661, - -0.001414, - -0.025646, - 0.00993 + 0.004865, + -0.003695, + 0.002367, + -0.000958, + -0.000965, + -0.000914, + 0.004061, + 5.2e-05, + -0.004466, + 0.000528, + -0.00554, + 0.002936, + 0.010828, + -0.003573, + -0.003714, + -0.004223, + -0.01237, + -0.00467, + -0.001861, + -0.004012, + -0.009534, + -0.005811, + -0.004497, + -0.0022, + 0.003601, + -0.002138, + 0.004155, + 0.009307, + -0.003024, + -0.002535, + -3.1e-05, + 0.004786, + -0.00753, + -0.003172, + 0.002225, + -0.001803, + -9.7e-05, + 0.003097, + 0.001154, + 0.008286, + -0.001478, + 0.002645, + 0.000562, + 0.008003, + -0.000749, + 0.00237, + -0.005276, + -0.004898, + -0.00232, + 0.003043, + -0.001344, + 0.001505, + 0.003906, + -0.001533, + 0.004861, + -0.002786, + 0.003926, + 0.001462, + 0.008712, + 0.0024, + 0.003343, + 0.001272, + 0.000233, + -0.000715 ] } ], - "StableDiffusionTextToImage": [ + "XLMRobertaMaskedLM": [ { "shape": [ 2, 16, - 16, - 3 + 128 ], "sample": [ - 0.015141, - 0.014566, - 0.01507, - -0.013042, - -0.014793, - -0.013933, - 0.01303, - 0.014152, - 0.013037, - 0.014178, - 0.014338, - -0.013882, - -0.013874, - -0.013862, - 0.01476, - 0.014129, - 0.014743, - 0.014246, - 0.014257, - 0.014225, - -0.014007, - -0.013831, - 0.014013, - 0.014092, - 0.014013, - 0.014246, - 0.014301, - 0.014293, - -0.013722, - -0.013693, - -0.013301, - 0.014484, - 0.010029, - 0.014481, - 0.015246, - 0.015245, - -0.013927, - -0.013932, - -0.01385, - 0.014118, - 0.014105, - 0.014152, - 0.014243, - 0.01403, - -0.013841, - -0.013774, - -0.013812, - 0.014105, - 0.014097, - 0.014094, - 0.014242, - 0.014448, - 0.014266, - -0.014007, - -0.013837, - 0.014, + 0.013965, + 0.019017, + -0.003735, + -0.001879, + 0.010813, + 0.002046, + -0.009915, + -0.008338, + 0.019872, + 0.000889, + 0.04867, + 0.014064, + 0.0143, + 0.004872, + -0.004134, + 0.004227, + -0.055303, + 0.023829, + -0.035749, + -0.000273, + -0.019216, + -0.004134, + -0.017732, + -0.02736, + 0.016447, + 0.010392, + -0.032707, + -0.01479, + 0.009025, + -0.004914, + -0.008511, + 0.011608, + 0.015601, + -0.016102, + 0.003131, + 0.032635, + 0.021671, + 0.00695, + 0.003794, + -0.008977, + 0.008033, + -0.002371, + 0.003972, + -0.017868, + 0.026751, + 0.03861, + 0.016864, + -0.015283, + 0.004662, + -0.002961, + -0.041212, + -0.022634, + -0.044025, + 0.001969, + 0.009821, + -0.048096, + -0.034329, + 0.002701, 0.014121, - 0.014003, - 0.014259, - 0.011565, - 0.014413, - -0.011314, - -0.01329, - 0.013126 + -0.009398, + -0.017948, + -0.010799, + 0.017219, + -0.021088 + ] + } + ], + "XLMRobertaModel": [ + { + "shape": [ + 2, + 16, + 32 + ], + "sample": [ + -0.012946, + 0.015214, + -0.013106, + 0.015577, + -0.013179, + 0.069574, + 0.011613, + 0.070116, + 0.010722, + 0.024324, + 0.037582, + 0.024304, + 0.037166, + -0.035889, + -0.007053, + -0.036229, + -0.007861, + -0.073318, + -0.001794, + -0.073078, + -0.001493, + 0.031959, + -0.024997, + 0.033061, + -0.024845, + 0.032108, + 0.002208, + 0.016397, + 0.000681, + 0.016726, + 0.000822, + -0.005828, + 0.00068, + -0.00826, + 0.108422, + -0.049974, + 0.105383, + -0.04979, + -0.046333, + 0.025502, + -0.046555, + 0.026064, + 0.018798, + -0.003703, + 0.019448, + -0.00295, + 0.016, + 0.004491, + -0.02026, + 0.004788, + -0.020141, + 0.024671, + -0.017721, + 0.024861, + -0.017792, + 0.02249, + -0.024283, + 0.022482, + -0.024532, + -0.107456, + 0.000908, + -0.110674, + 0.000475, + -0.016094 ] }, { "shape": [ 2, - 8, - 8, - 8 + 32 ], "sample": [ - 0.051589, - 0.051341, - 0.051341, - 0.05132, - 0.051681, - -0.006162, - -0.006151, - -0.006155, - -0.006046, - 0.006829, - 0.006821, - 0.006841, - 0.006748, - 0.010876, - 0.010889, - 0.010892, - 0.010883, - -0.021278, - -0.02128, - -0.021273, - -0.021104, - -0.005376, - -0.005379, - -0.005368, - -0.005528, - -0.005407, - -0.000975, - -0.000971, - -0.00076, - -0.000734, - -0.043369, - -0.043349, - -0.043402, - -0.043551, - 0.051337, - 0.051297, - 0.05151, - 0.051494, - -0.006147, - -0.006312, - -0.006151, - -0.006165, - 0.006816, - 0.006766, - 0.006782, - 0.006824, - 0.006837, - 0.010968, - 0.010946, - 0.0109, - 0.010906, - -0.021234, - -0.021271, - -0.021282, - -0.021255, - -0.005589, - -0.005439, - -0.005423, - -0.005422, - -0.000794, - -0.00072, - -0.000731, - -0.000732, - -0.043442 + 9.5e-05, + 0.011957, + 0.006713, + 0.004725, + 0.011335, + -0.006035, + 0.012561, + 0.002068, + -0.0103, + 0.035051, + -0.010026, + 0.006029, + 0.002999, + 0.004093, + -0.02845, + 0.024154, + 0.007347, + 0.007143, + 0.001298, + -0.013618, + -0.011444, + -0.030955, + -0.005628, + -0.02546, + 0.018342, + 0.023768, + -0.011254, + -0.004289, + -0.019031, + 0.040206, + -0.025693, + -0.011637, + 8e-05, + 0.011829, + 0.006999, + 0.004674, + 0.011376, + -0.006081, + 0.012517, + 0.001985, + -0.010364, + 0.034915, + -0.010216, + 0.006072, + 0.003048, + 0.004378, + -0.028556, + 0.02404, + 0.007438, + 0.007219, + 0.001509, + -0.013589, + -0.011347, + -0.031104, + -0.00554, + -0.025427, + 0.018239, + 0.023746, + -0.01139, + -0.004372, + -0.018953, + 0.040002, + -0.025732, + -0.011578 ] - }, + } + ], + "XLMRobertaMultipleChoice": [ { "shape": [ 2, - 8, - 8, - 4 + 3 ], "sample": [ - 0.006219, - 0.002704, - 0.003533, - 0.003418, - 0.008001, - 0.005587, - 0.005602, - 0.005686, - 0.008691, - 0.002119, - -0.000714, - 0.001933, - -0.002422, - 0.002237, - 0.001836, - 0.002578, - -0.001956, - 0.00071, - -0.003084, - -0.002698, - -0.005522, - -0.001268, - -0.001686, - -1.2e-05, - -0.004803, - 0.000822, - 0.000538, - 0.013731, - 0.01433, - 0.014584, - 0.015002, - 0.015691, - 0.013323, - 0.015859, - 0.016175, - 0.015331, - 0.004414, - 0.004255, - 0.006079, - 0.002555, - 0.005861, - 0.004837, - 0.005748, - 0.002341, - 0.005368, - -9.1e-05, - 0.001814, - 0.003786, - 0.001644, - 0.001444, - 0.00205, - 0.00168, - 8.1e-05, - 0.000167, - -0.001057, - -0.002688, - -0.000987, - -0.000541, - -0.000336, - -0.002546, - -0.005307, - -0.003004, - -0.006414, - 0.018894 + -0.015686, + -0.015691, + -0.01569, + -0.015684, + -0.015687, + -0.015699 ] - }, + } + ], + "XLMRobertaQnA": [ { "shape": [ 2, - 16, - 32 + 16 ], "sample": [ - -0.012037, - 0.027224, - -0.017103, - 0.027472, - -0.020236, - 0.000331, - -0.01644, - 0.015825, - -0.018293, - 0.014603, - -0.005827, - 0.016616, - -0.007089, - -0.005842, - 0.009213, - 0.000568, - 0.049181, - -0.012505, - 0.000328, - -0.009735, - -0.003473, - -0.007774, - -0.024105, - -0.00845, - -0.020531, - -0.008289, - -0.052391, - 0.003785, - -0.052397, - 0.003749, - -0.019436, - 0.013164, - -0.022708, - 0.014415, - 0.028549, - -0.025015, - 0.017437, - -0.03349, - 0.051644, - -0.027284, - 0.04717, - -0.013412, - 0.04921, - 0.009694, - 0.039967, - 0.027683, - 0.015538, - -0.001814, - -0.016217, - -0.001629, - -0.021373, - 0.015966, - -0.035635, - 0.013777, - -0.038477, - 0.024585, - 0.014664, - 0.015243, - 0.014259, - -0.002774, - 0.006661, - -0.001414, - -0.025646, - 0.00993 + -0.040833, + -0.040412, + -0.040958, + -0.041051, + -0.040626, + -0.041051, + -0.040365, + -0.04072, + -0.040683, + -0.041041, + -0.040674, + -0.041258, + -0.040843, + -0.040957, + -0.040952, + -0.040941, + -0.040662, + -0.04115, + -0.040876, + -0.040881, + -0.040948, + -0.040716, + -0.041096, + -0.040604, + -0.041002, + -0.040509, + -0.040938, + -0.040975, + -0.041082, + -0.040653, + -0.040324, + -0.041059 ] - } - ], - "StableDiffusionXLModel": [ + }, { "shape": [ 2, - 16, - 16, - 3 + 16 ], "sample": [ - 0.015117, - 0.014538, - 0.015099, - -0.013055, - -0.014773, - -0.013969, - 0.013046, - 0.014114, - 0.013047, - 0.014206, - 0.01431, - -0.013879, - -0.013879, - -0.013813, - 0.014727, - 0.014137, - 0.014728, - 0.014265, - 0.014223, - 0.014277, - -0.013994, - -0.013862, - 0.013943, - 0.014071, - 0.01396, - 0.0142, - 0.014354, - 0.014264, - -0.01367, - -0.013662, - -0.013273, - 0.014495, - 0.009997, - 0.014473, - 0.015255, - 0.015218, - -0.013923, - -0.01392, - -0.013851, - 0.014064, - 0.014162, - 0.014085, - 0.014243, - 0.014099, - -0.013891, - -0.013708, - -0.013863, - 0.014129, - 0.01412, - 0.014114, - 0.014283, - 0.014466, - 0.014317, - -0.014031, - -0.013894, - 0.013986, - 0.01413, - 0.014016, - 0.014216, - 0.011546, - 0.014328, - -0.011294, - -0.013285, - 0.013122 + 0.019887, + 0.020086, + 0.019876, + 0.020019, + 0.020148, + 0.020007, + 0.020131, + 0.019951, + 0.020173, + 0.020156, + 0.019967, + 0.019955, + 0.020004, + 0.019911, + 0.019939, + 0.020037, + 0.020038, + 0.019925, + 0.020181, + 0.02007, + 0.020088, + 0.020004, + 0.019789, + 0.020143, + 0.020097, + 0.019977, + 0.020205, + 0.019934, + 0.02004, + 0.019997, + 0.020336, + 0.01988 ] - }, + } + ], + "XLMRobertaSequenceClassify": [ { "shape": [ 2, - 8, - 8, - 8 + 4 ], "sample": [ - 0.051605, - 0.051336, - 0.051333, - 0.051329, - 0.051681, - -0.006148, - -0.006164, - -0.006143, - -0.006038, - 0.006828, - 0.006835, - 0.006854, - 0.006746, - 0.01086, - 0.010873, - 0.010902, - 0.010893, - -0.021265, - -0.021276, - -0.021282, - -0.021111, - -0.005397, - -0.005385, - -0.005377, - -0.005534, - -0.005419, - -0.000964, - -0.000972, - -0.000765, - -0.000732, - -0.043374, - -0.043356, - -0.043409, - -0.043543, - 0.051326, - 0.051296, - 0.051495, - 0.051486, - -0.006144, - -0.006299, - -0.006143, - -0.006163, - 0.006826, - 0.006775, - 0.006793, - 0.006821, - 0.006818, - 0.010973, - 0.010951, - 0.010897, - 0.010886, - -0.02123, - -0.021273, - -0.021282, - -0.021264, - -0.005587, - -0.005444, - -0.005419, - -0.005421, - -0.000796, - -0.000722, - -0.000731, - -0.00073, - -0.043448 + 0.000177, + 0.014164, + -0.036711, + -0.005497, + 0.000159, + 0.01417, + -0.036695, + -0.005501 ] - }, + } + ], + "XLMRobertaTokenClassify": [ { "shape": [ 2, - 8, - 8, - 4 + 16, + 5 ], "sample": [ - 0.005391, - 0.001505, - 0.002643, - 0.002754, - 0.005352, - 0.003054, - 0.002996, - 0.002889, - 0.006612, - 0.001375, - -0.00157, - 0.001698, - -0.003237, - 0.002067, - 0.001719, - 0.002174, - -0.003029, - -0.000492, - -0.005456, - -0.005545, - -0.007489, - -0.002657, - -0.003191, - -0.000723, - -0.006884, - -0.000489, - -0.000329, - 0.013296, - 0.014399, - 0.012816, - 0.013579, - 0.014752, - 0.014613, - 0.016026, - 0.016497, - 0.015358, - 0.001693, - 0.000713, - 0.00396, - 0.001166, - 0.003711, - 0.002246, - 0.002561, - 0.000427, - 0.00251, - -0.001125, - 0.002253, - 0.004757, - 0.001045, - 0.000851, - 0.001967, - 0.001514, - -0.000799, - -5.2e-05, - -0.002928, - -0.003612, - -0.003489, - -0.002121, - -0.001713, - -0.003564, - -0.005562, - -0.003289, - -0.007317, - 0.016749 + -0.011828, + -0.052942, + -0.012009, + -0.052773, + -0.012009, + -0.05311, + -0.011534, + -0.052935, + -0.011814, + -0.052717, + -0.011641, + -0.052989, + -0.011957, + -0.052632, + -0.011815, + -0.052956, + -0.011606, + -0.052607, + -0.011437, + -0.052679, + -0.011876, + 0.026477, + -0.011723, + 0.026765, + -0.011958, + 0.026455, + -0.011909, + 0.026545, + -0.011798, + 0.026639, + -0.011807, + 0.026565, + -0.011974, + 0.026476, + -0.011909, + 0.02662, + -0.011605, + 0.026641, + -0.011655, + 0.026516, + -0.011716, + 0.026667, + 0.000539, + 0.026477, + 0.000706, + 0.02649, + 0.000529, + 0.026434, + 0.000589, + 0.026668, + 0.000565, + 0.026388, + 0.000571, + 0.026631, + 0.000628, + 0.026565, + 0.00061, + 0.026627, + 0.000546, + 0.026415, + 0.000415, + 0.026348, + 0.000654, + -0.05073 + ] + } + ], + "XceptionImageClassify": [ + { + "shape": [ + 2, + 1000 + ], + "sample": [ + 0.029016, + 0.032873, + -0.008282, + -0.01587, + 0.0041, + 0.033549, + -0.04755, + 0.04435, + 0.008561, + -0.067786, + 0.00482, + 0.00107, + 0.032935, + 0.019229, + 0.020901, + 0.014184, + 0.006072, + 0.011499, + 0.001939, + -0.001087, + 0.013775, + -0.011099, + 0.030425, + 0.030052, + -0.014473, + 0.007834, + 0.020641, + -0.035, + -0.002316, + 0.003105, + -0.000788, + -0.031078, + 0.034641, + 0.002396, + 0.009609, + -0.031307, + -0.005042, + 0.00558, + 0.00853, + 0.021877, + 0.036303, + -0.00843, + -0.054513, + -0.029453, + 0.032808, + -0.008157, + -0.008873, + -0.013846, + 0.004452, + 0.001395, + 0.027834, + 0.00415, + -0.018065, + -0.05072, + -0.03487, + -0.002989, + 0.052267, + 0.016791, + -0.008023, + -0.014243, + -0.007622, + -0.034641, + -0.014824, + -0.03292 ] - }, + } + ], + "ZMDiffusion_qwen_image": [ { "shape": [ 2, - 16 + 16, + 16, + 3 ], "sample": [ - -0.001651, - 0.004384, - 0.000777, - 0.001588, - 0.001558, - 0.002165, - -0.002537, - 0.002124, - -0.002929, - 0.004796, - 0.002518, - -0.000757, - -0.006795, - -0.003231, - -0.005023, - 0.002541, - -0.001651, - 0.004384, - 0.000777, - 0.001588, - 0.001558, - 0.002165, - -0.002537, - 0.002124, - -0.002929, - 0.004796, - 0.002518, - -0.000757, - -0.006795, - -0.003231, - -0.005023, - 0.002541 + 0.003081, + 0.002639, + 0.002623, + -0.019443, + -0.020466, + -0.019435, + 0.010895, + 0.010283, + 0.010896, + 0.001854, + 0.001801, + -0.019436, + -0.019393, + -0.019437, + 0.010202, + 0.010282, + 0.010205, + 0.001856, + 0.001854, + 0.001859, + -0.019438, + -0.019436, + 0.010283, + 0.010284, + 0.010281, + 0.001854, + 0.001852, + 0.001851, + -0.01949, + -0.019487, + -0.019641, + 0.011624, + 0.01238, + 0.002639, + 0.001921, + 0.001924, + -0.019436, + -0.019431, + -0.019437, + 0.010286, + 0.010284, + 0.010288, + 0.001856, + 0.001857, + -0.019436, + -0.019428, + -0.019437, + 0.01028, + 0.010282, + 0.010282, + 0.001855, + 0.00184, + 0.001853, + -0.019504, + -0.019438, + 0.010197, + 0.010282, + 0.010199, + 0.001852, + 0.001637, + 0.001844, + -0.019552, + -0.019643, + 0.01312 ] }, { "shape": [ 2, 16, - 64 + 16, + 8 ], "sample": [ - 0.022025, - -0.090905, - -0.005076, - 0.007606, - 0.008709, - -0.020836, - -0.006304, - 0.056435, - -0.066661, - -0.024251, - -0.024757, - -0.038454, - 0.083587, - -0.02392, - 0.07096, - -0.037477, - -0.053752, - 0.016964, - -0.025098, - -0.000555, - 0.048412, - -0.056334, - -0.012582, - 0.013893, - -0.012235, - 0.048392, - 0.008929, - -0.023006, - -0.018066, - 0.025023, - -0.021386, - -0.020507, - -0.118211, - -0.040016, - -0.012872, - -0.005505, - 0.041438, - -0.029135, - 0.015047, - 0.003094, - 0.003135, - -0.052091, - -0.026491, - 0.025553, - 0.032734, - -0.047294, - -0.003716, - -0.016993, - 0.002541, - -0.002171, - -0.05959, - -0.001694, - 0.016903, - -0.071585, - 0.046993, - 0.011406, - -0.007823, - -0.061899, - 0.043031, - -0.024788, - -0.030914, - -0.014168, - -0.01589, - -0.052181 + 0.04349, + 0.00292, + 0.00799, + -0.023711, + -0.006685, + 0.01876, + 0.003434, + -0.019355, + 0.043433, + 0.003021, + 0.007969, + -0.023718, + -0.006645, + 0.018855, + 0.003458, + -0.019362, + 0.043484, + 0.002997, + 0.008004, + -0.023699, + -0.006694, + 0.018847, + 0.003461, + -0.019371, + 0.043468, + 0.003043, + 0.008022, + -0.023695, + -0.006721, + 0.018852, + 0.003411, + -0.019382, + 0.043425, + 0.002957, + 0.008039, + -0.023618, + -0.006733, + 0.018822, + 0.003465, + -0.019322, + 0.043453, + 0.003008, + 0.008019, + -0.023752, + -0.006704, + 0.018895, + 0.003493, + -0.019344, + 0.043495, + 0.002949, + 0.007995, + -0.023701, + -0.00665, + 0.018842, + 0.003426, + -0.019335, + 0.043465, + 0.002983, + 0.008051, + -0.023634, + -0.006767, + 0.018829, + 0.003383, + -0.019354 ] - } - ], - "StableDiffusionXLRefinerImageToImage": [ + }, { "shape": [ 2, 16, - 16, - 3 + 16 ], "sample": [ - 0.015128, - 0.014555, - 0.015074, - -0.013027, - -0.014804, - -0.013913, - 0.013123, - 0.014134, - 0.013115, - 0.014175, - 0.01428, - -0.01386, - -0.013889, - -0.013869, - 0.014721, - 0.014142, - 0.014713, - 0.014238, - 0.014288, - 0.014254, - -0.013996, - -0.013801, - 0.013976, - 0.014124, - 0.013997, - 0.014224, - 0.01436, - 0.014271, - -0.013665, - -0.013693, - -0.013272, - 0.014464, - 0.010032, - 0.014493, - 0.015215, - 0.015217, - -0.013947, - -0.013935, - -0.013853, - 0.014106, - 0.014111, - 0.014094, - 0.014294, - 0.014129, - -0.01385, - -0.013772, - -0.01386, - 0.014143, - 0.014081, - 0.014134, - 0.014289, - 0.014426, - 0.014263, - -0.014039, - -0.013828, - 0.013967, - 0.014145, - 0.013975, - 0.014264, - 0.011534, - 0.014323, - -0.011284, - -0.013253, - 0.013137 + 0.025382, + 0.047502, + 0.021523, + 0.003751, + 0.020512, + -0.02819, + 0.07473, + 0.055641, + 0.015141, + -0.116319, + 0.040794, + -0.103205, + -0.024959, + -0.02472, + -0.037541, + 0.083859, + -0.040947, + 0.018693, + 0.035366, + -0.078992, + -0.042715, + -0.145621, + 0.029375, + -0.135195, + 0.013235, + -0.028536, + -0.030166, + -0.07373, + -0.03337, + 0.222781, + -0.015033, + -0.032118, + 0.061251, + -0.035176, + 0.139108, + 0.098706, + 0.064707, + -0.011321, + -0.187153, + 0.136417, + 0.052554, + -0.03679, + 0.032921, + -0.003105, + -0.036564, + -0.050207, + 0.065246, + 0.004824, + -0.068774, + -0.089157, + 0.068509, + -0.045322, + -0.101066, + -0.053207, + -0.043108, + -0.15871, + -0.029696, + -0.088673, + -0.017017, + -0.090879, + 0.021776, + -0.060155, + 0.129198, + -0.07387 ] }, { "shape": [ 2, - 8, - 8, - 8 + 32, + 32 ], "sample": [ - 0.051589, - 0.051339, - 0.05134, - 0.051333, - 0.051681, - -0.006151, - -0.00616, - -0.006148, - -0.006045, - 0.006832, - 0.006832, - 0.006855, - 0.00675, - 0.0109, - 0.010893, - 0.010888, - 0.010893, - -0.02127, - -0.021268, - -0.021271, - -0.021117, - -0.005368, - -0.00538, - -0.005377, - -0.005524, - -0.005409, - -0.000983, - -0.000983, - -0.000756, - -0.000733, - -0.043366, - -0.043358, - -0.043399, - -0.043558, - 0.051339, - 0.051289, - 0.051522, - 0.051488, - -0.006163, - -0.006294, - -0.006157, - -0.006156, - 0.006832, - 0.006773, - 0.006796, - 0.006826, - 0.006832, - 0.010985, - 0.010941, - 0.01088, - 0.010879, - -0.021234, - -0.021277, - -0.021264, - -0.02128, - -0.005601, - -0.005452, - -0.005426, - -0.00543, - -0.000791, - -0.00072, - -0.000728, - -0.000727, - -0.043441 + 0.005047, + 0.005047, + 0.005047, + -0.015945, + -0.015945, + 0.002267, + 0.002267, + 0.002829, + 0.002829, + -0.003136, + -0.003136, + 0.018783, + 0.018783, + -0.055568, + -0.055568, + 0.001576, + 0.001576, + 0.001848, + 0.001848, + 0.006598, + 0.006598, + -0.014055, + -0.014055, + 0.004142, + 0.004142, + 0.007703, + 0.007703, + -0.023382, + -0.023382, + 0.058163, + 0.058163, + -0.024143, + -0.024143, + 0.001445, + 0.001445, + 0.001554, + 0.001554, + -0.045016, + -0.045016, + 0.003507, + 0.003507, + 0.000497, + 0.000497, + -0.056717, + -0.056717, + 0.018945, + 0.018945, + 0.003564, + 0.003564, + 0.018921, + 0.018921, + 0.017499, + 0.017499, + -0.000682, + -0.000682, + 0.002824, + 0.002824, + 0.045148, + 0.045148, + 0.001521, + 0.001521, + -0.000113, + -0.000113, + 0.002808 ] - }, + } + ], + "ZMDiffusion_qwen_image_21": [ { "shape": [ 2, - 8, - 8, + 64, + 64, 4 ], "sample": [ - 0.000481, - -0.002919, - -0.00179, - -0.000417, - 0.003282, - 0.003562, - 0.004124, - 0.003136, - 0.003245, - -0.009335, - -0.007721, - -0.008315, - -0.007362, - -0.010133, - -0.009758, - -0.009103, - -0.008528, - -0.010434, - -0.010003, - -0.007856, - -0.008889, - -0.008763, - -0.01003, - -0.008096, - -0.009952, - -0.006777, - -0.007405, - 0.009624, - 0.010534, - 0.013684, - 0.014548, - 0.014422, - 0.008618, - 0.009694, - 0.008467, - 0.007777, - 0.00285, - 0.003307, - 0.004301, - 0.00407, - 0.003844, - 0.0041, - 0.002652, - 0.003409, - 0.001842, - -0.00776, - -0.008313, - -0.003254, - -0.007042, - -0.009972, - -0.009185, - -0.0035, - -0.008157, - -0.009825, - -0.008416, - -0.006144, - -0.007326, - -0.007724, - -0.008343, - -0.006462, - -0.009628, - -0.010498, - -0.010122, - 0.016935 + 0.000958, + -0.002757, + -0.002676, + -0.002686, + -0.002682, + -0.002686, + -0.002666, + -0.002711, + -0.00273, + -0.0208, + -0.020851, + -0.02082, + -0.020843, + -0.02082, + -0.020862, + -0.020826, + -0.020863, + -0.020826, + 0.015441, + 0.015452, + 0.015469, + 0.015462, + 0.015427, + 0.015388, + 0.015356, + 0.015404, + 0.015399, + 0.033911, + 0.033944, + 0.03398, + 0.033954, + 0.03374, + 0.034017, + 0.033818, + 0.033958, + 0.033937, + -0.002685, + -0.002678, + -0.002668, + -0.002742, + -0.002671, + -0.002652, + -0.00272, + -0.002738, + -0.002786, + -0.020785, + -0.020731, + -0.020771, + -0.020838, + -0.020827, + -0.020796, + -0.020871, + -0.020885, + -0.020835, + 0.015476, + 0.015424, + 0.015454, + 0.015447, + 0.015401, + 0.015431, + 0.015359, + 0.015402, + 0.015371, + 0.032569 + ] + }, + { + "shape": [ + 2, + 16, + 16, + 16 + ], + "sample": [ + 0.043742, + 0.006526, + -0.00716, + 0.002203, + 0.007001, + -0.003265, + 0.041306, + 0.048222, + 0.044051, + 0.006727, + -0.007272, + 0.002304, + 0.007318, + -0.003362, + 0.041373, + 0.048279, + 0.043725, + 0.006633, + -0.007264, + 0.00234, + 0.007361, + -0.003407, + 0.041377, + 0.048261, + 0.043969, + 0.006626, + -0.007279, + 0.002379, + 0.007225, + -0.003334, + 0.041278, + 0.048426, + 0.043925, + 0.006591, + -0.007236, + 0.002311, + 0.007166, + -0.003305, + 0.04118, + 0.048263, + 0.043971, + 0.006603, + -0.007174, + 0.002288, + 0.007199, + -0.003511, + 0.041368, + 0.048239, + 0.04406, + 0.006657, + -0.007153, + 0.002314, + 0.007224, + -0.003533, + 0.041315, + 0.048327, + 0.043871, + 0.006593, + -0.007152, + 0.002406, + 0.007264, + -0.003716, + 0.041199, + -0.03484 ] }, { "shape": [ 2, + 16, 16 ], "sample": [ - -0.004386, - 0.005251, - -0.001644, - 0.002351, - 0.002163, - 0.003551, - -0.002766, - 0.001111, - -0.001758, - 0.005881, - -0.000112, - 0.002642, - -0.006228, - -0.002458, - -0.002431, - -0.000691, - -0.004386, - 0.005251, - -0.001644, - 0.002351, - 0.002163, - 0.003551, - -0.002766, - 0.001111, - -0.001758, - 0.005881, - -0.000112, - 0.002642, - -0.006228, - -0.002458, - -0.002431, - -0.000691 + 0.042214, + 0.105714, + 0.056851, + 0.076787, + 0.029758, + 0.00917, + 0.095684, + 0.106859, + 0.032329, + -0.067297, + 0.057183, + -0.051838, + -0.013893, + 0.035384, + -0.029907, + 0.154545, + -0.02961, + 0.08136, + 0.013182, + -0.054818, + -0.108304, + -0.122532, + -0.010875, + -0.090636, + -0.027775, + 0.009631, + -0.066816, + -0.096933, + -0.096951, + 0.185838, + -0.088463, + -0.070053, + -0.012736, + -0.047557, + 0.10501, + 0.094757, + 0.087138, + -0.094607, + -0.200976, + 0.063318, + 0.071323, + -0.137504, + 0.045576, + -0.086316, + -0.025822, + -0.020573, + 0.134116, + 0.039678, + -0.002558, + -0.075023, + 0.19758, + 0.001701, + 0.030064, + -0.005096, + -0.038978, + -0.093928, + -0.026556, + -0.030843, + -0.01305, + -0.017115, + 0.036709, + 0.005811, + 0.139922, + 0.009673 ] }, { @@ -29706,220 +26719,146 @@ 32 ], "sample": [ - -0.07549, - -0.036797, - -0.06729, - -0.010316, - -0.087311, - 0.026806, - 0.046183, - 0.080015, - 0.028766, - 0.010841, - -0.001441, - 0.025617, - -0.022712, - 0.002993, - 0.03557, - 0.032339, - 0.039893, - -0.003567, - -0.040687, - -0.018836, - -0.055347, - -0.010683, - 0.006929, - 0.011367, - 0.034529, - -0.008013, - 0.033182, - -0.019324, - 0.047424, - -0.047148, - -0.093025, - 0.054083, - -0.060365, - 0.062953, - 0.030306, - 0.02033, - 0.052184, - 0.016348, - -0.005382, - -0.010774, - 0.03376, - -0.011427, - -0.024654, - -0.077269, - -0.040009, - -0.079375, - -0.008919, - 0.042451, - -0.07994, - 0.030593, - -0.049702, - -0.078214, - 0.049456, - -0.093707, - 0.04243, - -0.034284, - -0.032833, - -0.028629, - -0.038116, - 0.025034, - 0.0324, - -0.014014, - 0.029427, - -0.024193 - ] - } - ], - "StableDiffusionXLRefinerModel": [ - { - "shape": [ - 2, - 16, - 16, - 3 - ], - "sample": [ - 0.015128, - 0.014555, - 0.015074, - -0.013027, - -0.014804, - -0.013913, - 0.013123, - 0.014134, - 0.013115, - 0.014175, - 0.01428, - -0.01386, - -0.013889, - -0.013869, - 0.014721, - 0.014142, - 0.014713, - 0.014238, - 0.014288, - 0.014254, - -0.013996, - -0.013801, - 0.013976, - 0.014124, - 0.013997, - 0.014224, - 0.01436, - 0.014271, - -0.013665, - -0.013693, - -0.013272, - 0.014464, - 0.010032, - 0.014493, - 0.015215, - 0.015217, - -0.013947, - -0.013935, - -0.013853, - 0.014106, - 0.014111, - 0.014094, - 0.014294, - 0.014129, - -0.01385, - -0.013772, - -0.01386, - 0.014143, - 0.014081, - 0.014134, - 0.014289, - 0.014426, - 0.014263, - -0.014039, - -0.013828, - 0.013967, - 0.014145, - 0.013975, - 0.014264, - 0.011534, - 0.014323, - -0.011284, - -0.013253, - 0.013137 + 0.002967, + -0.002949, + 0.002967, + -0.002949, + 0.002967, + 0.026536, + -0.02298, + 0.026536, + -0.02298, + 0.023317, + 0.011212, + 0.023317, + 0.011212, + 0.000229, + 0.004552, + 0.000229, + 0.004552, + 0.012545, + -0.008191, + 0.012545, + -0.008191, + 0.043657, + -0.018314, + 0.043657, + -0.018314, + 0.043657, + 0.022612, + 0.015819, + 0.022612, + 0.015819, + -0.001567, + -0.012379, + -0.001567, + -0.012379, + 0.008989, + 0.015044, + 0.008989, + 0.015044, + 0.005536, + -0.028382, + 0.005536, + -0.028382, + 0.01766, + 0.026689, + 0.01766, + 0.026689, + 0.01766, + -0.036786, + 0.017117, + -0.036786, + 0.017117, + -0.033625, + -0.005554, + -0.033625, + -0.005554, + -0.017666, + -0.012415, + -0.017666, + -0.012415, + 0.012932, + -0.040328, + 0.012932, + -0.040328, + -0.023516 ] - }, + } + ], + "ZMDiffusion_stable_diffusion": [ { "shape": [ 2, - 8, - 8, - 8 + 16, + 16, + 3 ], "sample": [ - 0.051589, - 0.051339, - 0.05134, - 0.051333, - 0.051681, - -0.006151, - -0.00616, - -0.006148, - -0.006045, - 0.006832, - 0.006832, - 0.006855, - 0.00675, - 0.0109, - 0.010893, - 0.010888, - 0.010893, - -0.02127, - -0.021268, - -0.021271, - -0.021117, - -0.005368, - -0.00538, - -0.005377, - -0.005524, - -0.005409, - -0.000983, - -0.000983, - -0.000756, - -0.000733, - -0.043366, - -0.043358, - -0.043399, - -0.043558, - 0.051339, - 0.051289, - 0.051522, - 0.051488, - -0.006163, - -0.006294, - -0.006157, - -0.006156, - 0.006832, - 0.006773, - 0.006796, - 0.006826, - 0.006832, - 0.010985, - 0.010941, - 0.01088, - 0.010879, - -0.021234, - -0.021277, - -0.021264, - -0.02128, - -0.005601, - -0.005452, - -0.005426, - -0.00543, - -0.000791, - -0.00072, - -0.000728, - -0.000727, - -0.043441 + 0.015141, + 0.014566, + 0.01507, + -0.013042, + -0.014793, + -0.013933, + 0.01303, + 0.014152, + 0.013037, + 0.014178, + 0.014338, + -0.013882, + -0.013874, + -0.013862, + 0.01476, + 0.014129, + 0.014743, + 0.014246, + 0.014257, + 0.014225, + -0.014007, + -0.013831, + 0.014013, + 0.014092, + 0.014013, + 0.014246, + 0.014301, + 0.014293, + -0.013722, + -0.013693, + -0.013301, + 0.014484, + 0.010029, + 0.014481, + 0.015246, + 0.015245, + -0.013927, + -0.013932, + -0.01385, + 0.014118, + 0.014105, + 0.014152, + 0.014243, + 0.01403, + -0.013841, + -0.013774, + -0.013812, + 0.014105, + 0.014097, + 0.014094, + 0.014242, + 0.014448, + 0.014266, + -0.014007, + -0.013837, + 0.014, + 0.014121, + 0.014003, + 0.014259, + 0.011565, + 0.014413, + -0.011314, + -0.01329, + 0.013126 ] }, { @@ -29927,113 +26866,147 @@ 2, 8, 8, - 4 + 8 ], "sample": [ - 0.000481, - -0.002919, - -0.00179, - -0.000417, - 0.003282, - 0.003562, - 0.004124, - 0.003136, - 0.003245, - -0.009335, - -0.007721, - -0.008315, - -0.007362, - -0.010133, - -0.009758, - -0.009103, - -0.008528, - -0.010434, - -0.010003, - -0.007856, - -0.008889, - -0.008763, - -0.01003, - -0.008096, - -0.009952, - -0.006777, - -0.007405, - 0.009624, - 0.010534, - 0.013684, - 0.014548, - 0.014422, - 0.008618, - 0.009694, - 0.008467, - 0.007777, - 0.00285, - 0.003307, - 0.004301, - 0.00407, - 0.003844, - 0.0041, - 0.002652, - 0.003409, - 0.001842, - -0.00776, - -0.008313, - -0.003254, - -0.007042, - -0.009972, - -0.009185, - -0.0035, - -0.008157, - -0.009825, - -0.008416, - -0.006144, - -0.007326, - -0.007724, - -0.008343, - -0.006462, - -0.009628, - -0.010498, - -0.010122, - 0.016935 + 0.051589, + 0.051341, + 0.051341, + 0.05132, + 0.051681, + -0.006162, + -0.006151, + -0.006155, + -0.006046, + 0.006829, + 0.006821, + 0.006841, + 0.006748, + 0.010876, + 0.010889, + 0.010892, + 0.010883, + -0.021278, + -0.02128, + -0.021273, + -0.021104, + -0.005376, + -0.005379, + -0.005368, + -0.005528, + -0.005407, + -0.000975, + -0.000971, + -0.00076, + -0.000734, + -0.043369, + -0.043349, + -0.043402, + -0.043551, + 0.051337, + 0.051297, + 0.05151, + 0.051494, + -0.006147, + -0.006312, + -0.006151, + -0.006165, + 0.006816, + 0.006766, + 0.006782, + 0.006824, + 0.006837, + 0.010968, + 0.010946, + 0.0109, + 0.010906, + -0.021234, + -0.021271, + -0.021282, + -0.021255, + -0.005589, + -0.005439, + -0.005423, + -0.005422, + -0.000794, + -0.00072, + -0.000731, + -0.000732, + -0.043442 ] }, { "shape": [ 2, - 16 + 8, + 8, + 4 ], "sample": [ - -0.004386, - 0.005251, - -0.001644, - 0.002351, - 0.002163, - 0.003551, - -0.002766, - 0.001111, - -0.001758, - 0.005881, - -0.000112, - 0.002642, - -0.006228, - -0.002458, - -0.002431, - -0.000691, - -0.004386, - 0.005251, - -0.001644, - 0.002351, - 0.002163, - 0.003551, - -0.002766, - 0.001111, - -0.001758, - 0.005881, - -0.000112, - 0.002642, - -0.006228, - -0.002458, - -0.002431, - -0.000691 + 0.006219, + 0.002704, + 0.003533, + 0.003418, + 0.008001, + 0.005587, + 0.005602, + 0.005686, + 0.008691, + 0.002119, + -0.000714, + 0.001933, + -0.002422, + 0.002237, + 0.001836, + 0.002578, + -0.001956, + 0.00071, + -0.003084, + -0.002698, + -0.005522, + -0.001268, + -0.001686, + -1.2e-05, + -0.004803, + 0.000822, + 0.000538, + 0.013731, + 0.01433, + 0.014584, + 0.015002, + 0.015691, + 0.013323, + 0.015859, + 0.016175, + 0.015331, + 0.004414, + 0.004255, + 0.006079, + 0.002555, + 0.005861, + 0.004837, + 0.005748, + 0.002341, + 0.005368, + -9.1e-05, + 0.001814, + 0.003786, + 0.001644, + 0.001444, + 0.00205, + 0.00168, + 8.1e-05, + 0.000167, + -0.001057, + -0.002688, + -0.000987, + -0.000541, + -0.000336, + -0.002546, + -0.005307, + -0.003004, + -0.006414, + 0.018894 ] }, { @@ -30043,74 +27016,74 @@ 32 ], "sample": [ - -0.07549, - -0.036797, - -0.06729, - -0.010316, - -0.087311, - 0.026806, - 0.046183, - 0.080015, - 0.028766, - 0.010841, - -0.001441, - 0.025617, - -0.022712, - 0.002993, - 0.03557, - 0.032339, - 0.039893, - -0.003567, - -0.040687, - -0.018836, - -0.055347, - -0.010683, - 0.006929, - 0.011367, - 0.034529, - -0.008013, - 0.033182, - -0.019324, - 0.047424, - -0.047148, - -0.093025, - 0.054083, - -0.060365, - 0.062953, - 0.030306, - 0.02033, - 0.052184, - 0.016348, - -0.005382, - -0.010774, - 0.03376, - -0.011427, - -0.024654, - -0.077269, - -0.040009, - -0.079375, - -0.008919, - 0.042451, - -0.07994, - 0.030593, - -0.049702, - -0.078214, - 0.049456, - -0.093707, - 0.04243, - -0.034284, - -0.032833, - -0.028629, - -0.038116, - 0.025034, - 0.0324, - -0.014014, - 0.029427, - -0.024193 + -0.012037, + 0.027224, + -0.017103, + 0.027472, + -0.020236, + 0.000331, + -0.01644, + 0.015825, + -0.018293, + 0.014603, + -0.005827, + 0.016616, + -0.007089, + -0.005842, + 0.009213, + 0.000568, + 0.049181, + -0.012505, + 0.000328, + -0.009735, + -0.003473, + -0.007774, + -0.024105, + -0.00845, + -0.020531, + -0.008289, + -0.052391, + 0.003785, + -0.052397, + 0.003749, + -0.019436, + 0.013164, + -0.022708, + 0.014415, + 0.028549, + -0.025015, + 0.017437, + -0.03349, + 0.051644, + -0.027284, + 0.04717, + -0.013412, + 0.04921, + 0.009694, + 0.039967, + 0.027683, + 0.015538, + -0.001814, + -0.016217, + -0.001629, + -0.021373, + 0.015966, + -0.035635, + 0.013777, + -0.038477, + 0.024585, + 0.014664, + 0.015243, + 0.014259, + -0.002774, + 0.006661, + -0.001414, + -0.025646, + 0.00993 ] } ], - "StableDiffusionXLTextToImage": [ + "ZMDiffusion_stable_diffusion_2": [ { "shape": [ 2, @@ -30119,2967 +27092,3281 @@ 3 ], "sample": [ - 0.015117, - 0.014538, - 0.015099, - -0.013055, - -0.014773, - -0.013969, - 0.013046, - 0.014114, - 0.013047, - 0.014206, - 0.01431, - -0.013879, - -0.013879, - -0.013813, - 0.014727, - 0.014137, - 0.014728, - 0.014265, - 0.014223, - 0.014277, - -0.013994, + 0.015141, + 0.014566, + 0.01507, + -0.013042, + -0.014793, + -0.013933, + 0.01303, + 0.014152, + 0.013037, + 0.014178, + 0.014338, + -0.013882, + -0.013874, -0.013862, - 0.013943, - 0.014071, - 0.01396, - 0.0142, - 0.014354, - 0.014264, - -0.01367, - -0.013662, - -0.013273, - 0.014495, - 0.009997, - 0.014473, - 0.015255, - 0.015218, - -0.013923, - -0.01392, - -0.013851, - 0.014064, - 0.014162, - 0.014085, - 0.014243, - 0.014099, - -0.013891, - -0.013708, - -0.013863, + 0.01476, 0.014129, - 0.01412, - 0.014114, - 0.014283, - 0.014466, - 0.014317, - -0.014031, - -0.013894, - 0.013986, - 0.01413, - 0.014016, - 0.014216, - 0.011546, - 0.014328, - -0.011294, - -0.013285, - 0.013122 + 0.014743, + 0.014246, + 0.014257, + 0.014225, + -0.014007, + -0.013831, + 0.014013, + 0.014092, + 0.014013, + 0.014246, + 0.014301, + 0.014293, + -0.013722, + -0.013693, + -0.013301, + 0.014484, + 0.010029, + 0.014481, + 0.015246, + 0.015245, + -0.013927, + -0.013932, + -0.01385, + 0.014118, + 0.014105, + 0.014152, + 0.014243, + 0.01403, + -0.013841, + -0.013774, + -0.013812, + 0.014105, + 0.014097, + 0.014094, + 0.014242, + 0.014448, + 0.014266, + -0.014007, + -0.013837, + 0.014, + 0.014121, + 0.014003, + 0.014259, + 0.011565, + 0.014413, + -0.011314, + -0.01329, + 0.013126 + ] + }, + { + "shape": [ + 2, + 8, + 8, + 8 + ], + "sample": [ + 0.051589, + 0.051341, + 0.051341, + 0.05132, + 0.051681, + -0.006162, + -0.006151, + -0.006155, + -0.006046, + 0.006829, + 0.006821, + 0.006841, + 0.006748, + 0.010876, + 0.010889, + 0.010892, + 0.010883, + -0.021278, + -0.02128, + -0.021273, + -0.021104, + -0.005376, + -0.005379, + -0.005368, + -0.005528, + -0.005407, + -0.000975, + -0.000971, + -0.00076, + -0.000734, + -0.043369, + -0.043349, + -0.043402, + -0.043551, + 0.051337, + 0.051297, + 0.05151, + 0.051494, + -0.006147, + -0.006312, + -0.006151, + -0.006165, + 0.006816, + 0.006766, + 0.006782, + 0.006824, + 0.006837, + 0.010968, + 0.010946, + 0.0109, + 0.010906, + -0.021234, + -0.021271, + -0.021282, + -0.021255, + -0.005589, + -0.005439, + -0.005423, + -0.005422, + -0.000794, + -0.00072, + -0.000731, + -0.000732, + -0.043442 + ] + }, + { + "shape": [ + 2, + 8, + 8, + 4 + ], + "sample": [ + 0.006219, + 0.002704, + 0.003533, + 0.003418, + 0.008001, + 0.005587, + 0.005602, + 0.005686, + 0.008691, + 0.002119, + -0.000714, + 0.001933, + -0.002422, + 0.002237, + 0.001836, + 0.002578, + -0.001956, + 0.00071, + -0.003084, + -0.002698, + -0.005522, + -0.001268, + -0.001686, + -1.2e-05, + -0.004803, + 0.000822, + 0.000538, + 0.013731, + 0.01433, + 0.014584, + 0.015002, + 0.015691, + 0.013323, + 0.015859, + 0.016175, + 0.015331, + 0.004414, + 0.004255, + 0.006079, + 0.002555, + 0.005861, + 0.004837, + 0.005748, + 0.002341, + 0.005368, + -9.1e-05, + 0.001814, + 0.003786, + 0.001644, + 0.001444, + 0.00205, + 0.00168, + 8.1e-05, + 0.000167, + -0.001057, + -0.002688, + -0.000987, + -0.000541, + -0.000336, + -0.002546, + -0.005307, + -0.003004, + -0.006414, + 0.018894 ] }, { "shape": [ 2, - 8, - 8, - 8 + 16, + 32 ], "sample": [ - 0.051605, - 0.051336, - 0.051333, - 0.051329, - 0.051681, - -0.006148, - -0.006164, - -0.006143, - -0.006038, - 0.006828, - 0.006835, - 0.006854, - 0.006746, - 0.01086, - 0.010873, - 0.010902, - 0.010893, - -0.021265, - -0.021276, - -0.021282, - -0.021111, - -0.005397, - -0.005385, - -0.005377, - -0.005534, - -0.005419, - -0.000964, - -0.000972, - -0.000765, - -0.000732, - -0.043374, - -0.043356, - -0.043409, - -0.043543, - 0.051326, - 0.051296, - 0.051495, - 0.051486, - -0.006144, - -0.006299, - -0.006143, - -0.006163, - 0.006826, - 0.006775, - 0.006793, - 0.006821, - 0.006818, - 0.010973, - 0.010951, - 0.010897, - 0.010886, - -0.02123, - -0.021273, - -0.021282, - -0.021264, - -0.005587, - -0.005444, - -0.005419, - -0.005421, - -0.000796, - -0.000722, - -0.000731, - -0.00073, - -0.043448 + -0.012036, + 0.027224, + -0.017103, + 0.027472, + -0.020236, + 0.000332, + -0.016442, + 0.015828, + -0.018295, + 0.014603, + -0.005828, + 0.016616, + -0.00709, + -0.005843, + 0.009211, + 0.000567, + 0.049181, + -0.012505, + 0.000328, + -0.009735, + -0.003474, + -0.007774, + -0.024105, + -0.00845, + -0.020531, + -0.008289, + -0.052391, + 0.003785, + -0.052397, + 0.003749, + -0.019434, + 0.013164, + -0.022706, + 0.014415, + 0.028551, + -0.025012, + 0.017438, + -0.033488, + 0.051644, + -0.027285, + 0.047169, + -0.013412, + 0.049211, + 0.009689, + 0.039967, + 0.02768, + 0.015537, + -0.001814, + -0.01621, + -0.001629, + -0.021368, + 0.015966, + -0.035636, + 0.013777, + -0.038478, + 0.024581, + 0.014664, + 0.015239, + 0.014259, + -0.002774, + 0.006665, + -0.001414, + -0.025644, + 0.009934 ] - }, + } + ], + "ZMDiffusion_stable_diffusion_3": [ { "shape": [ 2, - 8, - 8, - 4 + 16, + 16, + 3 ], "sample": [ - 0.005391, - 0.001505, - 0.002643, - 0.002754, - 0.005352, - 0.003054, - 0.002996, - 0.002889, - 0.006612, - 0.001375, - -0.00157, - 0.001698, - -0.003237, - 0.002067, - 0.001719, - 0.002174, - -0.003029, - -0.000492, - -0.005456, - -0.005545, - -0.007489, - -0.002657, - -0.003191, - -0.000723, - -0.006884, - -0.000489, - -0.000329, - 0.013296, - 0.014399, - 0.012816, - 0.013579, - 0.014752, - 0.014613, - 0.016026, - 0.016497, - 0.015358, - 0.001693, - 0.000713, - 0.00396, - 0.001166, - 0.003711, - 0.002246, - 0.002561, - 0.000427, - 0.00251, - -0.001125, - 0.002253, - 0.004757, - 0.001045, - 0.000851, - 0.001967, - 0.001514, - -0.000799, - -5.2e-05, - -0.002928, - -0.003612, - -0.003489, - -0.002121, - -0.001713, - -0.003564, - -0.005562, - -0.003289, - -0.007317, - 0.016749 + 0.014955, + 0.014535, + 0.015039, + -0.01294, + -0.015078, + -0.01354, + 0.012603, + 0.014234, + 0.013028, + 0.014104, + 0.014657, + -0.013961, + -0.01397, + -0.013705, + 0.014085, + 0.014107, + 0.01454, + 0.014684, + 0.014214, + 0.01473, + -0.014538, + -0.013919, + 0.013891, + 0.013894, + 0.013345, + 0.014624, + 0.015308, + 0.015109, + -0.013916, + -0.013549, + -0.013454, + 0.015281, + 0.010778, + 0.014866, + 0.015276, + 0.014992, + -0.014008, + -0.013972, + -0.013274, + 0.014006, + 0.014443, + 0.014339, + 0.013793, + 0.014472, + -0.013407, + -0.013649, + -0.013639, + 0.014388, + 0.014662, + 0.014543, + 0.015311, + 0.014633, + 0.013674, + -0.01387, + -0.01422, + 0.013761, + 0.014448, + 0.013939, + 0.01379, + 0.011211, + 0.013861, + -0.010878, + -0.012983, + 0.013007 ] }, { "shape": [ 2, - 16 + 8, + 8, + 8 ], "sample": [ - -0.001651, - 0.004384, - 0.000777, - 0.001588, - 0.001558, - 0.002165, - -0.002537, - 0.002124, - -0.002929, - 0.004796, - 0.002518, - -0.000757, - -0.006795, - -0.003231, - -0.005023, - 0.002541, - -0.001651, - 0.004384, - 0.000777, - 0.001588, - 0.001558, - 0.002165, - -0.002537, - 0.002124, - -0.002929, - 0.004796, - 0.002518, - -0.000757, - -0.006795, - -0.003231, - -0.005023, - 0.002541 + 0.017274, + 0.019562, + 0.019478, + 0.019517, + 0.016443, + 0.035483, + 0.035683, + 0.035216, + 0.030619, + -0.011854, + -0.012384, + -0.011986, + -0.008521, + 0.008372, + 0.008306, + 0.008393, + 0.006177, + 0.009252, + 0.009079, + 0.008854, + 0.009121, + -0.050123, + -0.050035, + -0.050299, + -0.042257, + -0.050059, + 0.019482, + 0.019644, + 0.015013, + 0.013524, + 0.016211, + 0.016645, + 0.0115, + 0.018172, + 0.019602, + 0.017579, + 0.015945, + 0.01631, + 0.035517, + 0.037192, + 0.036485, + 0.035644, + -0.012204, + -0.015612, + -0.012235, + -0.012165, + -0.012408, + 0.012303, + 0.009171, + 0.008228, + 0.008402, + 0.00942, + 0.008764, + 0.00924, + 0.009194, + -0.047951, + -0.049781, + -0.049926, + -0.050071, + 0.018187, + 0.013454, + 0.013745, + 0.013502, + 0.018325 ] }, { "shape": [ 2, - 16, - 64 + 8, + 8, + 4 ], "sample": [ - 0.022025, - -0.090905, - -0.005076, - 0.007606, - 0.008709, - -0.020836, - -0.006304, - 0.056435, - -0.066661, - -0.024251, - -0.024757, - -0.038454, - 0.083587, - -0.02392, - 0.07096, - -0.037477, - -0.053752, - 0.016964, - -0.025098, - -0.000555, - 0.048412, - -0.056334, - -0.012582, - 0.013893, - -0.012235, - 0.048392, - 0.008929, - -0.023006, - -0.018066, - 0.025023, - -0.021386, - -0.020507, - -0.118211, - -0.040016, - -0.012872, - -0.005505, - 0.041438, - -0.029135, - 0.015047, - 0.003094, - 0.003135, - -0.052091, - -0.026491, - 0.025553, - 0.032734, - -0.047294, - -0.003716, - -0.016993, - 0.002541, - -0.002171, - -0.05959, - -0.001694, - 0.016903, - -0.071585, - 0.046993, - 0.011406, - -0.007823, - -0.061899, - 0.043031, - -0.024788, - -0.030914, - -0.014168, - -0.01589, - -0.052181 + 0.012091, + 0.097546, + 0.1286, + 0.065201, + -0.030095, + 0.150227, + 0.080477, + -0.046941, + 0.100516, + -0.039675, + -0.109624, + -0.035173, + -0.018982, + 0.017097, + 0.0496, + -0.016586, + -0.028893, + -0.045285, + 0.017648, + -0.007875, + -0.051069, + 0.043383, + 0.113692, + -0.142047, + -0.113984, + -0.090202, + -0.088588, + 0.136368, + 0.105976, + 0.16715, + 0.126743, + 0.109739, + -0.003245, + 0.006413, + 0.211325, + 0.010501, + 0.100809, + -0.026821, + -0.011584, + -0.063026, + 0.151531, + 0.07651, + 0.149342, + 0.198445, + -0.054413, + 0.071169, + 0.004813, + -0.014978, + 0.154012, + 0.049127, + -0.056784, + 0.082541, + 0.070801, + 0.150369, + 0.045853, + -0.040565, + 0.095416, + -0.002005, + 0.055735, + 0.052691, + -0.016378, + 0.060451, + 0.060556, + 0.000441 ] - } - ], - "SwinImageClassify": [ + }, { "shape": [ 2, - 1000 + 24 ], "sample": [ - 0.03727, - 0.010736, - 0.00379, - -0.011314, - -0.001801, - 0.048002, - -0.037877, - 0.035164, - 0.011903, - -0.06182, - -0.001327, - -0.021618, - 0.03845, - 0.023628, - -0.012034, - 0.030154, - 0.013825, - 0.008501, - 0.0072, - 0.02061, - 0.008837, - -0.024512, - 0.017009, - 0.004001, - -0.040273, - 0.004263, - -0.000877, - -0.035575, - 0.030585, - -0.026599, - 0.030435, - 0.002476, - 0.025173, - 0.005227, - 0.011432, - -0.012005, - -0.002249, - -0.007996, - 0.011941, - 0.001431, - 0.020885, - -0.018125, - -0.041154, - -0.025066, - 0.029893, - -0.007118, - -0.012232, - -0.001271, - 0.004738, - -0.010213, - 0.047078, - -0.012023, - -0.002808, - -0.047233, - -0.017357, - -0.020568, - 0.021827, - 0.024365, - -0.028548, - -0.02213, - 0.001802, - -0.002804, - -0.002764, - -0.030708 + 0.000567, + -0.002695, + -0.000344, + 0.001872, + -0.000507, + -0.002267, + -0.003681, + 0.002676, + 0.004716, + -0.00181, + 0.001028, + 0.001284, + -0.000525, + -0.000201, + 0.001415, + 0.002244, + -0.000562, + 0.001038, + 0.00074, + 0.002248, + 0.000935, + 0.002127, + -0.000101, + 0.000829, + 0.000567, + -0.002695, + -0.000344, + 0.001872, + -0.000507, + -0.002267, + -0.003681, + 0.002676, + 0.004716, + -0.00181, + 0.001028, + 0.001284, + -0.000525, + -0.000201, + 0.001415, + 0.002244, + -0.000562, + 0.001038, + 0.00074, + 0.002248, + 0.000935, + 0.002127, + -0.000101, + 0.000829 ] - } - ], - "SwinV2ImageClassify": [ + }, { "shape": [ 2, - 1000 + 16, + 24 ], "sample": [ - 0.028752, - -0.006332, - 0.000363, - 0.002035, - -0.006701, - 0.047487, - -0.039371, - 0.023125, - 0.016489, - -0.060217, - -0.003939, - -0.024132, - 0.057357, - 0.01048, - -0.038023, - 0.010662, - 0.032048, + 0.01991, + 0.004829, + 0.01725, + 0.038287, + -0.006734, + 0.004269, + 0.045626, + 0.02678, + 0.05008, + 0.019808, + 0.082032, + 0.008582, + 0.014972, + 0.010277, + 0.034913, + 0.022746, + 0.015038, + 0.01132, + -0.021046, + -0.056212, -0.001057, - 0.010816, - 0.01114, - 0.005381, - -0.038666, - 0.016187, - 0.005265, - -0.040001, - 0.002051, - 0.002523, - -0.020339, - 0.033311, - -0.027095, - 0.016853, - -0.019808, - 0.010958, - -0.014414, - 0.002567, - -0.027859, - -0.005412, - -0.005288, - 0.006264, - 0.010331, - 0.029314, - -0.024317, - -0.059341, - -0.047102, - 0.015536, - -0.005734, - -0.020719, - 0.00087, - -0.006053, - -0.031878, - 0.041451, - 0.015935, - 0.005199, - -0.049402, - -0.009516, - 0.011538, - 0.025188, - 0.020908, - -0.015169, - -0.022819, - 0.003664, - -0.009767, - 0.008631, - -0.043226 + -0.067862, + -0.024714, + -0.057371, + -0.035649, + -0.058719, + -0.052335, + -0.043667, + -0.032667, + -0.017269, + 0.023394, + -0.025872, + 0.037565, + -0.039031, + 0.032609, + -0.011299, + 0.029473, + 0.035091, + 0.006203, + 0.019498, + 0.028072, + 0.015388, + -0.055906, + 0.0522, + -0.031006, + 0.028746, + -0.001313, + -0.05912, + -0.015437, + -0.072604, + 0.00309, + -0.052105, + -0.060652, + 0.002501, + -0.005301, + -0.016556, + 0.004076, + -0.030116, + 0.015446, + -0.040764, + 0.025618, + -0.082329, + 0.017716, + 0.019545 ] } ], - "T5ConditionalGenerate": [ + "ZMDiffusion_stable_diffusion_3_5": [ { "shape": [ 2, 16, - 32 + 16, + 3 ], "sample": [ - 0.007004, - 0.007011, - 0.009395, - 0.004649, - -0.014974, - -0.007318, - -0.008844, - 0.000876, - 0.005626, - -0.007449, - 0.003385, - -0.011547, - 0.00152, - 0.027466, - -0.006198, - -0.011448, - 0.011551, - 0.013829, - 0.020342, - 0.00263, - 0.029595, - -0.004371, - 0.000389, - -0.004182, - 0.012, - -0.015576, - -2.8e-05, - -0.012556, - 2.5e-05, - 0.054732, - -0.01089, - 0.00547, - 0.015315, - -0.027271, - -0.002516, - 0.003293, - 0.001099, - 0.010529, - -0.00809, - -0.031443, - 0.000377, - -0.031829, - 0.014496, - -0.003451, - -0.031001, - -0.00064, - -0.015198, - -0.001283, - -0.004123, - 0.002237, - 0.002431, - 0.015693, - 0.018091, - -0.015281, - -0.037994, - -0.002521, - -0.001512, - -0.004742, - 0.000613, - -0.012441, - 0.000619, - -0.003316, - 8.3e-05, - -0.013941 + 0.014955, + 0.014535, + 0.015039, + -0.01294, + -0.015078, + -0.01354, + 0.012603, + 0.014234, + 0.013028, + 0.014104, + 0.014657, + -0.013961, + -0.01397, + -0.013705, + 0.014085, + 0.014107, + 0.01454, + 0.014684, + 0.014214, + 0.01473, + -0.014538, + -0.013919, + 0.013891, + 0.013894, + 0.013345, + 0.014624, + 0.015308, + 0.015109, + -0.013916, + -0.013549, + -0.013454, + 0.015281, + 0.010778, + 0.014866, + 0.015276, + 0.014992, + -0.014008, + -0.013972, + -0.013274, + 0.014006, + 0.014443, + 0.014339, + 0.013793, + 0.014472, + -0.013407, + -0.013649, + -0.013639, + 0.014388, + 0.014662, + 0.014543, + 0.015311, + 0.014633, + 0.013674, + -0.01387, + -0.01422, + 0.013761, + 0.014448, + 0.013939, + 0.01379, + 0.011211, + 0.013861, + -0.010878, + -0.012983, + 0.013007 ] }, { "shape": [ 2, 8, - 32 + 8, + 8 ], "sample": [ - -0.012666, - 0.008346, - -0.012532, - 0.025897, - -0.012666, - 0.008346, - -0.012532, - 0.025897, - -0.012666, - -0.008003, - 0.011512, - 0.020149, - -0.022597, - -0.008003, - 0.011512, - 0.020149, - -0.022597, - -0.008003, - 0.000519, - 0.005653, - -0.008304, - -8.3e-05, - 0.000519, - 0.005653, - -0.008304, - -8.3e-05, - 0.000519, - -0.005227, - -0.026162, - 0.014672, - -0.008533, - -0.005227, - -0.026154, - 0.014576, - -0.008315, - -0.005249, - -0.015154, - -0.005074, - -0.000946, - 0.010883, - -0.015154, - -0.005074, - -0.000946, - 0.010883, - -0.015154, - -0.024924, - 0.01542, - -0.004209, - -0.007258, - -0.024924, - 0.01542, - -0.004209, - -0.007258, - -0.024924, - -0.002469, - 0.002147, - -0.007502, - -0.002374, - -0.002469, - 0.002147, - -0.007502, - -0.002374, - -0.002469, - 0.002716 + 0.017274, + 0.019562, + 0.019478, + 0.019517, + 0.016443, + 0.035483, + 0.035683, + 0.035216, + 0.030619, + -0.011854, + -0.012384, + -0.011986, + -0.008521, + 0.008372, + 0.008306, + 0.008393, + 0.006177, + 0.009252, + 0.009079, + 0.008854, + 0.009121, + -0.050123, + -0.050035, + -0.050299, + -0.042257, + -0.050059, + 0.019482, + 0.019644, + 0.015013, + 0.013524, + 0.016211, + 0.016645, + 0.0115, + 0.018172, + 0.019602, + 0.017579, + 0.015945, + 0.01631, + 0.035517, + 0.037192, + 0.036485, + 0.035644, + -0.012204, + -0.015612, + -0.012235, + -0.012165, + -0.012408, + 0.012303, + 0.009171, + 0.008228, + 0.008402, + 0.00942, + 0.008764, + 0.00924, + 0.009194, + -0.047951, + -0.049781, + -0.049926, + -0.050071, + 0.018187, + 0.013454, + 0.013745, + 0.013502, + 0.018325 ] }, { "shape": [ 2, 8, - 128 + 8, + 4 ], "sample": [ - -1.3e-05, - -8.4e-05, - 6.1e-05, - -0.000177, - -0.000162, - -1.9e-05, - 0.000194, - 0.000133, - -0.000201, - 0.000181, - -0.000129, - 0.000109, - 2e-05, - 0.000283, - -0.000111, - -6e-06, - -9e-05, - -0.00027, - 6.1e-05, - -0.000117, - 4e-06, - -0.000192, - -0.000222, - 3.4e-05, - 3e-06, - -6.8e-05, - 0.000154, - -0.000511, - 0.000117, - -0.000262, - 0.000157, - -0.000222, - 0.000341, - -0.00019, - 0.000101, - 0.000131, - 0.000283, - -0.0, - 0.000358, - 3.3e-05, - 6.7e-05, - 0.000137, - -0.000217, - 0.000333, - 0.000191, - 0.000265, - 8.4e-05, - 5.3e-05, - -0.000328, - -0.000244, - 2.3e-05, - 0.000154, - 0.000141, - 8.2e-05, - 4e-06, - -0.00014, - -8.6e-05, - -0.00017, - 0.000444, - 0.000237, - 4e-05, - 1.8e-05, - -0.00063, - 0.000147 + 0.01323, + 0.098254, + 0.129669, + 0.066239, + -0.028639, + 0.150828, + 0.081944, + -0.045938, + 0.101357, + -0.039674, + -0.10952, + -0.035309, + -0.019402, + 0.016394, + 0.048952, + -0.017228, + -0.029025, + -0.045206, + 0.018156, + -0.007403, + -0.051539, + 0.042579, + 0.113596, + -0.142042, + -0.113807, + -0.089958, + -0.088259, + 0.136318, + 0.105069, + 0.166323, + 0.125724, + 0.108803, + -0.003717, + 0.005988, + 0.211244, + 0.010028, + 0.100874, + -0.026819, + -0.011773, + -0.063357, + 0.151896, + 0.077693, + 0.151372, + 0.199062, + -0.054222, + 0.072171, + 0.006164, + -0.014138, + 0.153723, + 0.049434, + -0.056111, + 0.083136, + 0.071386, + 0.150738, + 0.046609, + -0.039163, + 0.09659, + -0.00176, + 0.055724, + 0.053063, + -0.015067, + 0.061407, + 0.061014, + 0.002211 ] - } - ], - "T5EncoderModel": [ + }, { "shape": [ 2, - 16, - 32 + 24 ], "sample": [ - 0.007004, - 0.007011, - 0.009395, - 0.004649, - -0.014974, - -0.007318, - -0.008844, - 0.000876, - 0.005626, - -0.007449, - 0.003385, - -0.011547, - 0.00152, - 0.027466, - -0.006198, - -0.011448, - 0.011551, - 0.013829, - 0.020342, - 0.00263, - 0.029595, - -0.004371, - 0.000389, - -0.004182, - 0.012, - -0.015576, - -2.8e-05, - -0.012556, - 2.5e-05, - 0.054732, - -0.01089, - 0.00547, - 0.015315, - -0.027271, - -0.002516, - 0.003293, - 0.001099, - 0.010529, - -0.00809, - -0.031443, - 0.000377, - -0.031829, - 0.014496, - -0.003451, - -0.031001, - -0.00064, - -0.015198, - -0.001283, - -0.004123, - 0.002237, - 0.002431, - 0.015693, - 0.018091, - -0.015281, - -0.037994, - -0.002521, - -0.001512, - -0.004742, - 0.000613, - -0.012441, - 0.000619, - -0.003316, - 8.3e-05, - -0.013941 + 0.000567, + -0.002695, + -0.000344, + 0.001872, + -0.000507, + -0.002267, + -0.003681, + 0.002676, + 0.004716, + -0.00181, + 0.001028, + 0.001284, + -0.000525, + -0.000201, + 0.001415, + 0.002244, + -0.000562, + 0.001038, + 0.00074, + 0.002248, + 0.000935, + 0.002127, + -0.000101, + 0.000829, + 0.000567, + -0.002695, + -0.000344, + 0.001872, + -0.000507, + -0.002267, + -0.003681, + 0.002676, + 0.004716, + -0.00181, + 0.001028, + 0.001284, + -0.000525, + -0.000201, + 0.001415, + 0.002244, + -0.000562, + 0.001038, + 0.00074, + 0.002248, + 0.000935, + 0.002127, + -0.000101, + 0.000829 ] - } - ], - "T5Model": [ + }, { "shape": [ 2, 16, - 32 + 24 ], "sample": [ - 0.007004, - 0.007011, - 0.009395, - 0.004649, - -0.014974, - -0.007318, - -0.008844, - 0.000876, - 0.005626, - -0.007449, - 0.003385, - -0.011547, - 0.00152, - 0.027466, - -0.006198, - -0.011448, - 0.011551, - 0.013829, - 0.020342, - 0.00263, - 0.029595, - -0.004371, - 0.000389, - -0.004182, - 0.012, - -0.015576, - -2.8e-05, - -0.012556, - 2.5e-05, - 0.054732, - -0.01089, - 0.00547, - 0.015315, - -0.027271, - -0.002516, - 0.003293, - 0.001099, - 0.010529, - -0.00809, - -0.031443, - 0.000377, - -0.031829, - 0.014496, - -0.003451, - -0.031001, - -0.00064, - -0.015198, - -0.001283, - -0.004123, - 0.002237, - 0.002431, - 0.015693, - 0.018091, - -0.015281, - -0.037994, - -0.002521, - -0.001512, - -0.004742, - 0.000613, - -0.012441, - 0.000619, - -0.003316, - 8.3e-05, - -0.013941 + 0.01991, + 0.004829, + 0.01725, + 0.038287, + -0.006734, + 0.004269, + 0.045626, + 0.02678, + 0.05008, + 0.019808, + 0.082032, + 0.008582, + 0.014972, + 0.010277, + 0.034913, + 0.022746, + 0.015038, + 0.01132, + -0.021046, + -0.056212, + -0.001057, + -0.067862, + -0.024714, + -0.057371, + -0.035649, + -0.058719, + -0.052335, + -0.043667, + -0.032667, + -0.017269, + 0.023394, + -0.025872, + 0.037565, + -0.039031, + 0.032609, + -0.011299, + 0.029473, + 0.035091, + 0.006203, + 0.019498, + 0.028072, + 0.015388, + -0.055906, + 0.0522, + -0.031006, + 0.028746, + -0.001313, + -0.05912, + -0.015437, + -0.072604, + 0.00309, + -0.052105, + -0.060652, + 0.002501, + -0.005301, + -0.016556, + 0.004076, + -0.030116, + 0.015446, + -0.040764, + 0.025618, + -0.082329, + 0.017716, + 0.019545 ] - }, + } + ], + "ZMDiffusion_stable_diffusion_xl": [ { "shape": [ 2, - 8, - 32 + 16, + 16, + 3 ], "sample": [ - -0.012666, - 0.008346, - -0.012532, - 0.025897, - -0.012666, - 0.008346, - -0.012532, - 0.025897, - -0.012666, - -0.008003, - 0.011512, - 0.020149, - -0.022597, - -0.008003, - 0.011512, - 0.020149, - -0.022597, - -0.008003, - 0.000519, - 0.005653, - -0.008304, - -8.3e-05, - 0.000519, - 0.005653, - -0.008304, - -8.3e-05, - 0.000519, - -0.005227, - -0.026162, - 0.014672, - -0.008533, - -0.005227, - -0.026154, - 0.014576, - -0.008315, - -0.005249, - -0.015154, - -0.005074, - -0.000946, - 0.010883, - -0.015154, - -0.005074, - -0.000946, - 0.010883, - -0.015154, - -0.024924, - 0.01542, - -0.004209, - -0.007258, - -0.024924, - 0.01542, - -0.004209, - -0.007258, - -0.024924, - -0.002469, - 0.002147, - -0.007502, - -0.002374, - -0.002469, - 0.002147, - -0.007502, - -0.002374, - -0.002469, - 0.002716 + 0.015117, + 0.014538, + 0.015099, + -0.013055, + -0.014773, + -0.013969, + 0.013046, + 0.014114, + 0.013047, + 0.014206, + 0.01431, + -0.013879, + -0.013879, + -0.013813, + 0.014727, + 0.014137, + 0.014728, + 0.014265, + 0.014223, + 0.014277, + -0.013994, + -0.013862, + 0.013943, + 0.014071, + 0.01396, + 0.0142, + 0.014354, + 0.014264, + -0.01367, + -0.013662, + -0.013273, + 0.014495, + 0.009997, + 0.014473, + 0.015255, + 0.015218, + -0.013923, + -0.01392, + -0.013851, + 0.014064, + 0.014162, + 0.014085, + 0.014243, + 0.014099, + -0.013891, + -0.013708, + -0.013863, + 0.014129, + 0.01412, + 0.014114, + 0.014283, + 0.014466, + 0.014317, + -0.014031, + -0.013894, + 0.013986, + 0.01413, + 0.014016, + 0.014216, + 0.011546, + 0.014328, + -0.011294, + -0.013285, + 0.013122 ] - } - ], - "T5QnA": [ + }, { "shape": [ 2, - 16 + 8, + 8, + 8 ], "sample": [ - -0.036581, - -0.036534, - -0.033987, - -0.033718, - -0.032503, - -0.032634, - -0.036825, - -0.033032, - -0.033208, - -0.036711, - -0.035561, - -0.033981, - -0.037899, - -0.030949, - -0.038119, - -0.034084, - -0.036592, - -0.032261, - -0.03053, - -0.029801, - -0.031892, - -0.03281, - -0.035399, - -0.036351, - -0.034096, - -0.02928, - -0.035885, - -0.036726, - -0.032882, - -0.03446, - -0.03706, - -0.033631 + 0.051605, + 0.051336, + 0.051333, + 0.051329, + 0.051681, + -0.006148, + -0.006164, + -0.006143, + -0.006038, + 0.006828, + 0.006835, + 0.006854, + 0.006746, + 0.01086, + 0.010873, + 0.010902, + 0.010893, + -0.021265, + -0.021276, + -0.021282, + -0.021111, + -0.005397, + -0.005385, + -0.005377, + -0.005534, + -0.005419, + -0.000964, + -0.000972, + -0.000765, + -0.000732, + -0.043374, + -0.043356, + -0.043409, + -0.043543, + 0.051326, + 0.051296, + 0.051495, + 0.051486, + -0.006144, + -0.006299, + -0.006143, + -0.006163, + 0.006826, + 0.006775, + 0.006793, + 0.006821, + 0.006818, + 0.010973, + 0.010951, + 0.010897, + 0.010886, + -0.02123, + -0.021273, + -0.021282, + -0.021264, + -0.005587, + -0.005444, + -0.005419, + -0.005421, + -0.000796, + -0.000722, + -0.000731, + -0.00073, + -0.043448 ] }, { "shape": [ 2, - 16 + 8, + 8, + 4 ], "sample": [ - 0.013309, - 0.016126, - 0.015854, - 0.015916, - 0.014135, - 0.016762, - 0.015658, - 0.016548, - 0.016202, - 0.014686, - 0.014325, - 0.011945, - 0.013924, - 0.014791, - 0.013524, - 0.01577, - 0.013307, - 0.016444, - 0.014633, - 0.013717, - 0.013196, - 0.014541, - 0.01449, - 0.014696, - 0.013105, - 0.015376, - 0.015856, - 0.014193, - 0.013698, - 0.015209, - 0.013511, - 0.016833 - ] - } - ], - "T5SequenceClassify": [ + 0.005391, + 0.001505, + 0.002643, + 0.002754, + 0.005352, + 0.003054, + 0.002996, + 0.002889, + 0.006612, + 0.001375, + -0.00157, + 0.001698, + -0.003237, + 0.002067, + 0.001719, + 0.002174, + -0.003029, + -0.000492, + -0.005456, + -0.005545, + -0.007489, + -0.002657, + -0.003191, + -0.000723, + -0.006884, + -0.000489, + -0.000329, + 0.013296, + 0.014399, + 0.012816, + 0.013579, + 0.014752, + 0.014613, + 0.016026, + 0.016497, + 0.015358, + 0.001693, + 0.000713, + 0.00396, + 0.001166, + 0.003711, + 0.002246, + 0.002561, + 0.000427, + 0.00251, + -0.001125, + 0.002253, + 0.004757, + 0.001045, + 0.000851, + 0.001967, + 0.001514, + -0.000799, + -5.2e-05, + -0.002928, + -0.003612, + -0.003489, + -0.002121, + -0.001713, + -0.003564, + -0.005562, + -0.003289, + -0.007317, + 0.016749 + ] + }, { "shape": [ 2, - 4 + 16 ], "sample": [ - 0.000502, - 0.015776, - -0.036908, - -0.005068, - 0.000508, - 0.015436, - -0.03677, - -0.005324 + -0.001651, + 0.004384, + 0.000777, + 0.001588, + 0.001558, + 0.002165, + -0.002537, + 0.002124, + -0.002929, + 0.004796, + 0.002518, + -0.000757, + -0.006795, + -0.003231, + -0.005023, + 0.002541, + -0.001651, + 0.004384, + 0.000777, + 0.001588, + 0.001558, + 0.002165, + -0.002537, + 0.002124, + -0.002929, + 0.004796, + 0.002518, + -0.000757, + -0.006795, + -0.003231, + -0.005023, + 0.002541 ] - } - ], - "T5TokenClassify": [ + }, { "shape": [ 2, 16, - 4 + 64 ], "sample": [ - -0.019692, - -0.057709, - -0.023285, - -0.059521, - -0.022501, - -0.062189, - -0.01735, - -0.058669, - -0.01329, - -0.055911, - -0.021385, - -0.060795, - -0.01553, - -0.057615, - -0.021567, - -0.053714, - -0.016597, - -0.056771, - -0.017363, - -0.054833, - -0.019004, - -0.054143, - -0.019093, - -0.055566, - -0.020616, - -0.052903, - -0.020262, - -0.051335, - -0.01721, - -0.049579, - -0.016186, - -0.051587, - -0.01902, - -0.057753, - -0.016269, - -0.053339, - -0.019398, - -0.056683, - -0.020809, - -0.05843, - -0.017729, - -0.055149, - -0.020295, - -0.053454, - -0.015878, - -0.054542, - -0.017298, - -0.053031, - -0.016035, - -0.06077, - -0.01906, - -0.056622, - -0.017711, - -0.054803, - -0.019422, - -0.058055, - -0.015465, - -0.055705, - -0.019656, - -0.053127, - -0.016712, - -0.055788, - -0.015164, - 0.027239 + 0.022025, + -0.090905, + -0.005076, + 0.007606, + 0.008709, + -0.020836, + -0.006304, + 0.056435, + -0.066661, + -0.024251, + -0.024757, + -0.038454, + 0.083587, + -0.02392, + 0.07096, + -0.037477, + -0.053752, + 0.016964, + -0.025098, + -0.000555, + 0.048412, + -0.056334, + -0.012582, + 0.013893, + -0.012235, + 0.048392, + 0.008929, + -0.023006, + -0.018066, + 0.025023, + -0.021386, + -0.020507, + -0.118211, + -0.040016, + -0.012872, + -0.005505, + 0.041438, + -0.029135, + 0.015047, + 0.003094, + 0.003135, + -0.052091, + -0.026491, + 0.025553, + 0.032734, + -0.047294, + -0.003716, + -0.016993, + 0.002541, + -0.002171, + -0.05959, + -0.001694, + 0.016903, + -0.071585, + 0.046993, + 0.011406, + -0.007823, + -0.061899, + 0.043031, + -0.024788, + -0.030914, + -0.014168, + -0.01589, + -0.052181 ] } ], - "TableTransformerDetect": [ + "ZMDiffusion_stable_diffusion_xl_refiner": [ { "shape": [ 2, - 10, + 16, + 16, 3 ], "sample": [ - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897, - -0.035959, - -0.025007, - 0.029897 + 0.015128, + 0.014555, + 0.015074, + -0.013027, + -0.014804, + -0.013913, + 0.013123, + 0.014134, + 0.013115, + 0.014175, + 0.01428, + -0.01386, + -0.013889, + -0.013869, + 0.014721, + 0.014142, + 0.014713, + 0.014238, + 0.014288, + 0.014254, + -0.013996, + -0.013801, + 0.013976, + 0.014124, + 0.013997, + 0.014224, + 0.01436, + 0.014271, + -0.013665, + -0.013693, + -0.013272, + 0.014464, + 0.010032, + 0.014493, + 0.015215, + 0.015217, + -0.013947, + -0.013935, + -0.013853, + 0.014106, + 0.014111, + 0.014094, + 0.014294, + 0.014129, + -0.01385, + -0.013772, + -0.01386, + 0.014143, + 0.014081, + 0.014134, + 0.014289, + 0.014426, + 0.014263, + -0.014039, + -0.013828, + 0.013967, + 0.014145, + 0.013975, + 0.014264, + 0.011534, + 0.014323, + -0.011284, + -0.013253, + 0.013137 + ] + }, + { + "shape": [ + 2, + 8, + 8, + 8 + ], + "sample": [ + 0.051589, + 0.051339, + 0.05134, + 0.051333, + 0.051681, + -0.006151, + -0.00616, + -0.006148, + -0.006045, + 0.006832, + 0.006832, + 0.006855, + 0.00675, + 0.0109, + 0.010893, + 0.010888, + 0.010893, + -0.02127, + -0.021268, + -0.021271, + -0.021117, + -0.005368, + -0.00538, + -0.005377, + -0.005524, + -0.005409, + -0.000983, + -0.000983, + -0.000756, + -0.000733, + -0.043366, + -0.043358, + -0.043399, + -0.043558, + 0.051339, + 0.051289, + 0.051522, + 0.051488, + -0.006163, + -0.006294, + -0.006157, + -0.006156, + 0.006832, + 0.006773, + 0.006796, + 0.006826, + 0.006832, + 0.010985, + 0.010941, + 0.01088, + 0.010879, + -0.021234, + -0.021277, + -0.021264, + -0.02128, + -0.005601, + -0.005452, + -0.005426, + -0.00543, + -0.000791, + -0.00072, + -0.000728, + -0.000727, + -0.043441 ] }, { "shape": [ 2, - 10, + 8, + 8, 4 ], "sample": [ - 0.498714, - 0.503207, - 0.503151, - 0.506022, - 0.503207, - 0.503151, - 0.506022, - 0.498714, - 0.503151, - 0.506022, - 0.498714, - 0.503207, - 0.506022, - 0.498714, - 0.503207, - 0.503151, - 0.498714, - 0.503207, - 0.503151, - 0.506022, - 0.503207, - 0.503151, - 0.506022, - 0.498714, - 0.503151, - 0.506022, - 0.498714, - 0.503207, - 0.506022, - 0.498714, - 0.503207, - 0.503151, - 0.498714, - 0.503207, - 0.503151, - 0.506022, - 0.503207, - 0.503151, - 0.506022, - 0.498714, - 0.503151, - 0.506022, - 0.498714, - 0.503207, - 0.506022, - 0.498714, - 0.503207, - 0.503151, - 0.498714, - 0.503207, - 0.503151, - 0.506022, - 0.503207, - 0.503151, - 0.506022, - 0.498714, - 0.503151, - 0.506022, - 0.498714, - 0.503207, - 0.506022, - 0.498714, - 0.503207, - 0.506022 + 0.000481, + -0.002919, + -0.00179, + -0.000417, + 0.003282, + 0.003562, + 0.004124, + 0.003136, + 0.003245, + -0.009335, + -0.007721, + -0.008315, + -0.007362, + -0.010133, + -0.009758, + -0.009103, + -0.008528, + -0.010434, + -0.010003, + -0.007856, + -0.008889, + -0.008763, + -0.01003, + -0.008096, + -0.009952, + -0.006777, + -0.007405, + 0.009624, + 0.010534, + 0.013684, + 0.014548, + 0.014422, + 0.008618, + 0.009694, + 0.008467, + 0.007777, + 0.00285, + 0.003307, + 0.004301, + 0.00407, + 0.003844, + 0.0041, + 0.002652, + 0.003409, + 0.001842, + -0.00776, + -0.008313, + -0.003254, + -0.007042, + -0.009972, + -0.009185, + -0.0035, + -0.008157, + -0.009825, + -0.008416, + -0.006144, + -0.007326, + -0.007724, + -0.008343, + -0.006462, + -0.009628, + -0.010498, + -0.010122, + 0.016935 ] - } - ], - "Tipsv2DptDensePredict": [ + }, { "shape": [ 2, - 16, 16 ], "sample": [ - 0.01559, - 0.015544, - 0.015598, - 0.015541, - 0.015599, - 0.015541, - 0.015599, - 0.015542, - 0.015599, - 0.015542, - 0.015541, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015544, - 0.015544, - 0.015544, - 0.015544, - 0.015541, - 0.015541, - 0.015541, - 0.015541, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015544, - 0.015526 + -0.004386, + 0.005251, + -0.001644, + 0.002351, + 0.002163, + 0.003551, + -0.002766, + 0.001111, + -0.001758, + 0.005881, + -0.000112, + 0.002642, + -0.006228, + -0.002458, + -0.002431, + -0.000691, + -0.004386, + 0.005251, + -0.001644, + 0.002351, + 0.002163, + 0.003551, + -0.002766, + 0.001111, + -0.001758, + 0.005881, + -0.000112, + 0.002642, + -0.006228, + -0.002458, + -0.002431, + -0.000691 ] }, { "shape": [ 2, 16, - 16, - 5 - ], - "sample": [ - -0.010536, - -0.010721, - 0.01322, - 0.013027, - -0.009794, - -0.039424, - -0.039745, - -0.023638, - -0.023823, - -0.010883, - 0.01303, - 0.013027, - -0.009837, - -0.039424, - -0.039421, - -0.023638, - -0.023642, - -0.010883, - 0.013028, - 0.013027, - -0.009836, - -0.039424, - -0.039423, - -0.023638, - -0.023639, - -0.010883, - 0.013028, - 0.013028, - -0.009836, - -0.009836, - -0.039418, - -0.023689, - -0.023834, - -0.010721, - 0.013027, - 0.013027, - -0.009836, - -0.009836, - -0.039424, - -0.023638, - -0.023638, - -0.010883, - 0.013027, - 0.013026, - -0.009836, - -0.009836, - -0.039424, - -0.023638, - -0.023638, - -0.010884, - -0.010883, - 0.013024, - -0.009836, - -0.009835, - -0.039424, - -0.023637, - -0.023638, - -0.010755, - -0.010882, - 0.012814, - -0.009836, - -0.009669, - -0.039418, - -0.023813 - ] - } - ], - "Tipsv2DptDepthEstimation": [ - { - "shape": [ - 2, - 16, - 16 + 32 ], "sample": [ - 0.01559, - 0.015544, - 0.015598, - 0.015541, - 0.015599, - 0.015541, - 0.015599, - 0.015542, - 0.015599, - 0.015542, - 0.015541, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015544, - 0.015544, - 0.015544, - 0.015544, - 0.015541, - 0.015541, - 0.015541, - 0.015541, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015542, - 0.015544, - 0.015526 + -0.07549, + -0.036797, + -0.06729, + -0.010316, + -0.087311, + 0.026806, + 0.046183, + 0.080015, + 0.028766, + 0.010841, + -0.001441, + 0.025617, + -0.022712, + 0.002993, + 0.03557, + 0.032339, + 0.039893, + -0.003567, + -0.040687, + -0.018836, + -0.055347, + -0.010683, + 0.006929, + 0.011367, + 0.034529, + -0.008013, + 0.033182, + -0.019324, + 0.047424, + -0.047148, + -0.093025, + 0.054083, + -0.060365, + 0.062953, + 0.030306, + 0.02033, + 0.052184, + 0.016348, + -0.005382, + -0.010774, + 0.03376, + -0.011427, + -0.024654, + -0.077269, + -0.040009, + -0.079375, + -0.008919, + 0.042451, + -0.07994, + 0.030593, + -0.049702, + -0.078214, + 0.049456, + -0.093707, + 0.04243, + -0.034284, + -0.032833, + -0.028629, + -0.038116, + 0.025034, + 0.0324, + -0.014014, + 0.029427, + -0.024193 ] } ], - "Tipsv2DptSemanticSegment": [ + "ZMSentenceTransformer": [ { "shape": [ 2, - 16, - 16, - 5 + 32 ], "sample": [ - -0.010536, - -0.010721, - 0.01322, - 0.013027, - -0.009794, - -0.039424, - -0.039745, - -0.023638, - -0.023823, - -0.010883, - 0.01303, - 0.013027, - -0.009837, - -0.039424, - -0.039421, - -0.023638, - -0.023642, - -0.010883, - 0.013028, - 0.013027, - -0.009836, - -0.039424, - -0.039423, - -0.023638, - -0.023639, - -0.010883, - 0.013028, - 0.013028, - -0.009836, - -0.009836, - -0.039418, - -0.023689, - -0.023834, - -0.010721, - 0.013027, - 0.013027, - -0.009836, - -0.009836, - -0.039424, - -0.023638, - -0.023638, - -0.010883, - 0.013027, - 0.013026, - -0.009836, - -0.009836, - -0.039424, - -0.023638, - -0.023638, - -0.010884, - -0.010883, - 0.013024, - -0.009836, - -0.009835, - -0.039424, - -0.023637, - -0.023638, - -0.010755, - -0.010882, - 0.012814, - -0.009836, - -0.009669, - -0.039418, - -0.023813 + -0.138292, + 0.107383, + 0.227017, + -0.157133, + 0.043625, + -0.147396, + 0.077896, + 0.145626, + 0.28461, + -0.071275, + 0.122478, + -0.137024, + -0.125705, + -0.16926, + -0.036518, + 0.066965, + 0.107061, + -0.087836, + 0.160158, + -0.163558, + -0.408743, + 0.039037, + 0.045357, + -0.055598, + -0.408058, + 0.181733, + -0.026227, + 0.110355, + 0.224607, + 0.15685, + -0.357952, + -0.10573, + -0.136634, + 0.107497, + 0.226582, + -0.15824, + 0.043457, + -0.147285, + 0.080659, + 0.144457, + 0.282996, + -0.068748, + 0.124357, + -0.137633, + -0.122886, + -0.167857, + -0.037022, + 0.067155, + 0.108349, + -0.086218, + 0.160484, + -0.164189, + -0.406857, + 0.036848, + 0.046163, + -0.051938, + -0.41074, + 0.180494, + -0.025904, + 0.110552, + 0.223734, + 0.157155, + -0.360933, + -0.106077 ] - } - ], - "Tipsv2Model": [ + }, { "shape": [ - 2, - 32 - ], - "sample": [ - -0.204265, - -0.059741, - 0.026511, - -0.164984, - 0.046244, - -0.079875, - -0.577878, - 0.032957, - -0.457186, - -0.031679, - 0.134449, - -0.152016, - -0.073435, - 0.033245, - -0.147961, - -0.169231, - -0.114642, - -0.173875, - -0.086014, - 0.194842, - -0.035061, - -0.001086, - 0.069644, - 0.302139, - -0.014048, - -0.114194, - 0.130663, - -0.098736, - -0.06882, - 0.182741, - 0.03171, - -0.098478, - -0.204267, - -0.059737, - 0.026501, - -0.164981, - 0.046248, - -0.079876, - -0.577906, - 0.032955, - -0.457162, - -0.031677, - 0.134442, - -0.152013, - -0.073436, - 0.033242, - -0.147959, - -0.169224, - -0.11463, - -0.173872, - -0.086019, - 0.194818, - -0.035056, - -0.001096, - 0.069653, - 0.302146, - -0.014045, - -0.114194, - 0.13069, - -0.098734, - -0.068817, - 0.182733, - 0.031709, - -0.098476 + 2, + 16, + 32 + ], + "sample": [ + -0.020572, + 0.015957, + -0.019703, + 0.015892, + -0.020483, + -0.01368, + 0.016936, + -0.015326, + 0.014478, + 0.023795, + 0.033209, + 0.02408, + 0.033903, + -0.024844, + -0.023571, + -0.024457, + -0.02343, + -0.059766, + 0.006253, + -0.062119, + 0.006618, + 0.005499, + -0.020619, + 0.005742, + -0.023827, + 0.006142, + 0.009706, + 0.006107, + 0.008017, + 0.00602, + 0.021343, + -0.013266, + 0.015743, + -0.010255, + 0.043917, + -0.061277, + 0.042983, + -0.062477, + -0.010429, + 0.026378, + -0.009833, + 0.02688, + 0.019112, + -0.004345, + 0.019161, + -0.004257, + 0.018308, + 0.015562, + -0.020819, + 0.01682, + -0.02066, + 0.032244, + -0.018348, + 0.033819, + -0.017939, + 0.02347, + -0.025788, + 0.023409, + -0.025221, + -0.053642, + -0.004411, + -0.055167, + -0.005387, + -0.015949 ] - }, + } + ], + "ZMSentenceTransformer_all_MiniLM_L12_v2": [ { "shape": [ 2, - 2 + 16 ], "sample": [ - 33.656609, - 33.656609, - 33.657059, - 33.657059 + -0.006283, + -0.28829, + -0.510357, + -0.181786, + 0.267572, + -0.204303, + 0.065134, + -0.036725, + 0.221611, + -0.276, + 0.044667, + -0.319701, + -0.356479, + 0.162738, + 0.068047, + -0.341659, + -0.007111, + -0.286077, + -0.509385, + -0.181498, + 0.266028, + -0.204473, + 0.059981, + -0.032958, + 0.218527, + -0.278198, + 0.045874, + -0.318675, + -0.36582, + 0.169479, + 0.064059, + -0.336116 ] }, { "shape": [ 2, - 2 + 16, + 16 ], "sample": [ - 33.656609, - 33.657059, - 33.656609, - 33.657059 + -0.001388, + 0.02555, + 0.001283, + 0.024534, + 0.000669, + 0.0236, + -0.000574, + 0.024204, + -9.2e-05, + -0.031526, + -0.032391, + -0.030826, + -0.032146, + -0.031856, + -0.032201, + -0.030922, + -0.032239, + -0.032286, + -0.05965, + 0.004568, + -0.064821, + 0.002639, + -0.055885, + 0.007288, + -0.060506, + 0.002031, + -0.053601, + -0.035833, + -0.021838, + -0.035347, + -0.020152, + -0.036121, + -0.020518, + -0.036075, + -0.020536, + -0.036254, + 0.030103, + -0.039094, + 0.030121, + -0.035787, + 0.030136, + -0.047276, + 0.030123, + -0.041288, + 0.030123, + 0.014357, + -0.020862, + 0.01455, + -0.023323, + 0.019685, + -0.02363, + 0.00853, + -0.020957, + 0.031693, + 0.003902, + 0.005723, + 0.007174, + 0.006236, + 0.003883, + 0.006472, + 0.015475, + 0.010045, + -0.001431, + -0.036287 ] - }, + } + ], + "ZMSentenceTransformer_all_MiniLM_L6_v2": [ { "shape": [ 2, - 32 + 16 ], "sample": [ - 0.231413, - -0.231878, - 0.036957, - -0.112705, - 0.224442, - -0.284503, - -0.020773, - -0.191486, - -0.101979, - -0.30212, - -0.231222, - 0.084227, - 0.094925, - 0.060053, - -0.129362, - 0.210374, - -0.067982, - -0.058104, - 0.185936, - -0.040592, - 0.028479, - 0.224797, - 0.177719, - 0.404719, - 0.221314, - -0.168479, - -0.003695, - -0.149697, - 0.026421, - -0.147367, - 0.080145, - -0.234062, - 0.231413, - -0.231878, - 0.036957, - -0.112705, - 0.224442, - -0.284503, - -0.020773, - -0.191486, - -0.101979, - -0.30212, - -0.231222, - 0.084227, - 0.094925, - 0.060053, - -0.129362, - 0.210374, - -0.067982, - -0.058104, - 0.185936, - -0.040592, - 0.028479, - 0.224797, - 0.177719, - 0.404719, - 0.221314, - -0.168479, - -0.003695, - -0.149697, - 0.026421, - -0.147367, - 0.080145, - -0.234062 + -0.006283, + -0.28829, + -0.510357, + -0.181786, + 0.267572, + -0.204303, + 0.065134, + -0.036725, + 0.221611, + -0.276, + 0.044667, + -0.319701, + -0.356479, + 0.162738, + 0.068047, + -0.341659, + -0.007111, + -0.286077, + -0.509385, + -0.181498, + 0.266028, + -0.204473, + 0.059981, + -0.032958, + 0.218527, + -0.278198, + 0.045874, + -0.318675, + -0.36582, + 0.169479, + 0.064059, + -0.336116 ] - } - ], - "Tipsv2TextModel": [ + }, { "shape": [ 2, 16, - 32 + 16 ], "sample": [ - 0.033123, - -0.010784, - 0.0338, - -0.013429, - 0.033636, - -0.008718, - -0.032206, - -0.008677, - -0.033497, - 0.026862, - 0.007951, - 0.026524, - 0.005263, - -0.008706, - -0.016543, - -0.003263, - -0.018358, - 0.004191, - 0.032402, - 0.00429, - 0.031997, - 0.030361, - -0.041748, - 0.032885, - -0.040408, - 0.042542, - -0.003329, - 0.025781, - -0.002476, - 0.028777, - -0.019681, - 0.058398, - -0.022705, - 0.059073, - -0.015275, - 0.029098, - -0.013964, - 0.033865, - -0.050568, - -0.025017, - -0.046109, - -0.0233, - -0.03752, - 0.003248, - -0.035767, - -0.002517, - -0.033605, - -0.019767, - 0.0116, - -0.024619, - 0.012422, - 0.000619, - 0.016388, - 0.002959, - 0.003817, - -0.020142, - 0.008776, - -0.01486, - 0.008877, - 0.012328, - -0.026626, - 0.009385, - -0.017078, - -0.035824 + -0.001388, + 0.02555, + 0.001283, + 0.024534, + 0.000669, + 0.0236, + -0.000574, + 0.024204, + -9.2e-05, + -0.031526, + -0.032391, + -0.030826, + -0.032146, + -0.031856, + -0.032201, + -0.030922, + -0.032239, + -0.032286, + -0.05965, + 0.004568, + -0.064821, + 0.002639, + -0.055885, + 0.007288, + -0.060506, + 0.002031, + -0.053601, + -0.035833, + -0.021838, + -0.035347, + -0.020152, + -0.036121, + -0.020518, + -0.036075, + -0.020536, + -0.036254, + 0.030103, + -0.039094, + 0.030121, + -0.035787, + 0.030136, + -0.047276, + 0.030123, + -0.041288, + 0.030123, + 0.014357, + -0.020862, + 0.01455, + -0.023323, + 0.019685, + -0.02363, + 0.00853, + -0.020957, + 0.031693, + 0.003902, + 0.005723, + 0.007174, + 0.006236, + 0.003883, + 0.006472, + 0.015475, + 0.010045, + -0.001431, + -0.036287 ] - }, + } + ], + "ZMSentenceTransformer_all_distilroberta_v1": [ { "shape": [ 2, - 32 + 16 ], "sample": [ - 0.033625, - -0.033692, - 0.00537, - -0.016376, - 0.032612, - -0.041339, - -0.003018, - -0.027823, - -0.014818, - -0.043898, - -0.033597, - 0.012238, - 0.013793, - 0.008726, - -0.018796, - 0.030568, - -0.009878, - -0.008443, - 0.027017, - -0.005898, - 0.004138, - 0.032663, - 0.025823, - 0.058806, - 0.032157, - -0.02448, - -0.000537, - -0.021751, - 0.003839, - -0.021413, - 0.011645, - -0.03401, - 0.033625, - -0.033692, - 0.00537, - -0.016376, - 0.032612, - -0.041339, - -0.003018, - -0.027823, - -0.014818, - -0.043898, - -0.033597, - 0.012238, - 0.013793, - 0.008726, - -0.018796, - 0.030568, - -0.009878, - -0.008443, - 0.027017, - -0.005898, - 0.004138, - 0.032663, - 0.025823, - 0.058806, - 0.032157, - -0.02448, - -0.000537, - -0.021751, - 0.003839, - -0.021413, - 0.011645, - -0.03401 + -0.006428, + -0.288801, + -0.508576, + -0.181893, + 0.268083, + -0.204161, + 0.058898, + -0.038747, + 0.221928, + -0.27647, + 0.045056, + -0.320287, + -0.360028, + 0.155845, + 0.068021, + -0.342725, + -0.003668, + -0.288044, + -0.509014, + -0.183166, + 0.268015, + -0.205943, + 0.054317, + -0.032845, + 0.218373, + -0.281978, + 0.046054, + -0.321448, + -0.358621, + 0.167182, + 0.063274, + -0.335987 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + 0.00054, + 0.024421, + 0.00022, + 0.024225, + 0.000921, + 0.023616, + -0.00275, + 0.026056, + 8.3e-05, + -0.030825, + -0.032213, + -0.031122, + -0.032074, + -0.0312, + -0.032084, + -0.033028, + -0.032695, + -0.030283, + -0.064057, + 0.005521, + -0.059847, + 0.003993, + -0.058114, + 0.005001, + -0.053613, + 0.000932, + -0.055822, + -0.035447, + -0.021108, + -0.036471, + -0.020006, + -0.035757, + -0.020632, + -0.036374, + -0.020734, + -0.03649, + 0.030099, + -0.045783, + 0.030133, + -0.03234, + 0.030137, + -0.041763, + 0.030125, + -0.045943, + 0.030106, + 0.020362, + -0.021803, + 0.011972, + -0.022738, + 0.021537, + -0.023277, + 0.009515, + -0.022329, + 0.040238, + 0.014673, + 0.006805, + 0.01218, + 0.008585, + -0.000855, + 0.006342, + 0.002824, + 0.008836, + -0.000799, + -0.036171 ] } ], - "Tipsv2VisionModel": [ + "ZMSentenceTransformer_all_mpnet_base_v2": [ { "shape": [ 2, - 6, - 32 + 16 ], "sample": [ - -0.035782, - -0.10123, - -0.012864, - -0.015068, - -0.002461, - 0.005555, - 0.010684, - 0.039316, - -0.048693, - 0.008908, - -0.024619, - 0.011059, - -0.073175, - -0.01496, - -0.068544, - -0.016729, - -0.007399, - 0.005127, - -0.013034, - 0.007938, - 0.025508, - -0.018023, - -0.012591, - -0.02752, - -0.031995, - 0.021785, - 0.006093, - 0.01765, - 0.056828, - -0.042498, - -0.018075, - -0.056264, - 0.004642, - -0.080085, - -0.025919, - -0.006141, - 0.022894, - -0.013294, - 0.009124, - 0.007434, - 0.098242, - 0.033887, - -0.015847, - -0.028111, - -0.027522, - -0.030784, - 0.029121, - 0.101468, - -0.011032, - 0.003873, - -0.030811, - 0.011487, - -0.043628, - -0.090914, - 0.007013, - -0.036377, - -0.030575, - -0.010353, - -0.037672, - 0.03382, - -0.001482, - -0.017995, - 0.02622, - -0.015573 + 0.131721, + -0.261829, + -0.385196, + -0.173886, + 0.281246, + 0.001545, + 0.245415, + -0.072511, + -0.050302, + -0.224904, + 0.388975, + -0.262667, + -0.358547, + 0.329789, + 0.2014, + -0.199019, + 0.130397, + -0.259231, + -0.389688, + -0.17159, + 0.279488, + 0.000753, + 0.247705, + -0.083166, + -0.051451, + -0.230933, + 0.382909, + -0.258658, + -0.349881, + 0.325492, + 0.221085, + -0.203389 ] }, { "shape": [ 2, - 32 + 16, + 16 ], "sample": [ - -0.035782, - -0.010465, - 0.004644, - -0.028901, - 0.008101, - -0.013992, - -0.10123, - 0.005773, - -0.080088, - -0.005549, - 0.023552, - -0.026629, - -0.012864, - 0.005824, - -0.025919, - -0.029645, - -0.020083, - -0.030459, - -0.015068, - 0.034131, - -0.006142, - -0.00019, - 0.0122, - 0.052927, - -0.002461, - -0.020004, - 0.022889, - -0.017296, - -0.012056, - 0.032012, - 0.005555, - -0.017251, - -0.035783, - -0.010465, - 0.004642, - -0.028901, - 0.008102, - -0.013993, - -0.101237, - 0.005773, - -0.080085, - -0.005549, - 0.023551, - -0.026629, - -0.012864, - 0.005823, - -0.025919, - -0.029644, - -0.020081, - -0.030459, - -0.015069, - 0.034128, - -0.006141, - -0.000192, - 0.012202, - 0.052929, - -0.00246, - -0.020004, - 0.022894, - -0.017296, - -0.012055, - 0.032011, - 0.005555, - -0.017251 + 0.012509, + -0.007654, + 0.012005, + -0.005402, + 0.012474, + -0.005269, + 0.013393, + -0.00231, + 0.01668, + -0.023257, + -0.027662, + -0.027021, + -0.028156, + -0.025185, + -0.027928, + -0.026351, + -0.02774, + -0.028348, + -0.031596, + 0.042527, + -0.038723, + 0.039582, + -0.045372, + 0.041086, + -0.044355, + 0.041554, + -0.03575, + -0.027495, + -0.018154, + -0.024365, + -0.019006, + -0.026141, + -0.018336, + -0.038645, + -0.018094, + -0.020396, + 0.030104, + -0.045994, + 0.030238, + -0.01341, + 0.030433, + -0.053719, + 0.03042, + -0.077768, + 0.03017, + 0.034797, + -0.002746, + 0.033777, + -0.006548, + 0.037556, + 0.012106, + 0.030717, + -0.003646, + 0.039079, + 0.038353, + 0.012978, + 0.040298, + 0.028337, + 0.011519, + 0.025773, + 0.027593, + 0.011926, + 0.014531, + -0.022591 ] } ], - "VGGImageClassify": [ + "ZMSentenceTransformer_distiluse_base_multilingual_cased_v1": [ { "shape": [ 2, - 1000 + 8 ], "sample": [ - -0.075346, - 0.09173, - 0.003965, - -0.047696, - 0.005099, - 0.042046, - 0.039386, - -0.008242, - -0.100712, - -0.033049, - -0.077039, - 0.001321, - 0.02788, - 0.014689, - 0.005051, - -0.009848, - 0.001221, - -0.024884, - 0.038048, - -0.040728, - 0.039735, - -0.022384, - 0.008776, - 0.05443, - 0.013342, - -0.089412, - 0.101408, - -0.101833, - -0.098672, - -0.066069, - 0.071206, - 0.016617, - 0.046086, - -0.015095, - -0.079544, - -0.071482, - 0.027208, - 0.058843, - -0.1108, - 0.03061, - 0.098512, - 0.012451, - -0.084589, - 0.013578, - 0.000479, - 0.023356, - 0.059713, - -0.013926, - 0.049146, - -0.121879, - 0.023654, - 0.099224, - -0.010519, - 0.009046, - -0.055002, - 0.007435, - 0.050897, - 0.028916, - -0.005484, - 0.010513, - -0.058174, - -0.022928, - -0.0136, - -0.00607 + -0.005696, + 0.011375, + -0.000598, + -0.005064, + -0.008866, + 0.021272, + 0.007137, + -0.015385, + -0.005726, + 0.011304, + -0.000321, + -0.005126, + -0.008742, + 0.021266, + 0.00725, + -0.015222 ] - } - ], - "ViTImageClassify": [ + }, { "shape": [ 2, - 1000 + 16, + 16 ], "sample": [ - 0.028365, - 0.019775, - 0.01387, - -0.005907, - -0.00608, - 0.030517, - -0.024835, - 0.047901, - 0.014553, - -0.057438, - -0.006771, - -0.025818, - 0.026576, - 0.013906, - -0.00355, - 0.03239, - 0.024765, - 0.001365, - 0.000628, - 0.010015, - 0.012252, - -0.018734, - 0.010178, - 0.043254, - -0.013845, - 0.013418, - 0.001323, - -0.045827, - 0.009029, - -0.001661, - 0.002305, - -0.004402, - 0.006255, - 0.019019, - 0.020271, - -0.026835, - -0.004927, - 0.003636, - -0.015582, - -0.002634, - 0.03241, - -0.011949, - -0.063028, - -0.027115, - 0.029731, - -0.010011, - -0.018461, - -0.00617, - 0.011066, - -0.01429, - 0.019819, - -0.007293, - -0.001256, - -0.033223, - -0.027069, - 0.000452, - 0.037526, - 0.018432, - -0.015729, - -0.020209, - 0.007031, - -0.018771, - -0.016291, - -0.033565 + 0.00223, + -0.007174, + 0.002873, + -0.009289, + 0.003751, + -0.007166, + 0.003752, + -0.007555, + 0.001998, + 0.011331, + -0.02211, + -0.001803, + -0.023278, + 0.00989, + -0.0181, + -0.00481, + -0.031739, + 0.000239, + 0.015857, + 0.06128, + 0.015085, + 0.06705, + 0.017578, + 0.05266, + 0.013209, + 0.060971, + 0.015753, + 0.049514, + -0.001151, + 0.049367, + -0.001122, + 0.033595, + -0.000648, + 0.041646, + -0.000901, + 0.048977, + -0.010922, + -0.039426, + -0.004766, + -0.042949, + -0.007115, + -0.037974, + -0.003507, + -0.042364, + -0.006889, + -0.005795, + -0.003808, + -0.007997, + -0.003384, + 0.000805, + -0.004321, + -0.008732, + -0.005165, + -0.006883, + -0.002073, + 0.050016, + -0.005411, + 0.052228, + -0.006485, + 0.053975, + -0.00512, + 0.0522, + -0.004789, + 0.017832 ] } ], - "WhisperAudioClassify": [ + "ZMSentenceTransformer_distiluse_base_multilingual_cased_v2": [ { "shape": [ 2, - 10 + 8 ], "sample": [ - -0.019674, - -0.00448, - -0.05384, - 0.026557, - -0.047316, - -0.004436, - 0.012442, - 0.001493, - -0.031542, - -0.013457, - -0.019705, - -0.004451, - -0.053866, - 0.026555, - -0.047313, - -0.004437, - 0.012436, - 0.001537, - -0.031587, - -0.01345 + -0.005696, + 0.011375, + -0.000598, + -0.005064, + -0.008866, + 0.021272, + 0.007137, + -0.015385, + -0.005726, + 0.011304, + -0.000321, + -0.005126, + -0.008742, + 0.021266, + 0.00725, + -0.015222 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + 0.00223, + -0.007174, + 0.002873, + -0.009289, + 0.003751, + -0.007166, + 0.003752, + -0.007555, + 0.001998, + 0.011331, + -0.02211, + -0.001803, + -0.023278, + 0.00989, + -0.0181, + -0.00481, + -0.031739, + 0.000239, + 0.015857, + 0.06128, + 0.015085, + 0.06705, + 0.017578, + 0.05266, + 0.013209, + 0.060971, + 0.015753, + 0.049514, + -0.001151, + 0.049367, + -0.001122, + 0.033595, + -0.000648, + 0.041646, + -0.000901, + 0.048977, + -0.010922, + -0.039426, + -0.004766, + -0.042949, + -0.007115, + -0.037974, + -0.003507, + -0.042364, + -0.006889, + -0.005795, + -0.003808, + -0.007997, + -0.003384, + 0.000805, + -0.004321, + -0.008732, + -0.005165, + -0.006883, + -0.002073, + 0.050016, + -0.005411, + 0.052228, + -0.006485, + 0.053975, + -0.00512, + 0.0522, + -0.004789, + 0.017832 ] } ], - "WhisperConditionalGenerate": [ + "ZMSentenceTransformer_multi_qa_MiniLM_L6_cos_v1": [ { "shape": [ 2, - 20, - 64 + 16 ], "sample": [ - -0.016339, - 0.014459, - 0.027236, - -0.012696, - 0.003826, - 0.004806, - 0.00172, - 0.018539, - -0.003298, - -0.020976, - 0.037257, - -0.035281, - -0.004034, - 0.007173, - 0.031324, - -0.022888, - -0.013037, - -0.005674, - 0.017987, - 0.011197, - -0.007481, - 0.023746, - -0.024699, - 0.009173, - 0.0112, - -0.031971, - -0.023946, - -0.021795, - 0.052292, - -0.002748, - 0.00709, - 0.000376, - -0.006427, - -0.0098, - -0.010043, - -0.034723, - -0.013803, - 0.017374, - 0.028874, - -0.014603, - -0.010544, - 0.037403, - 0.020164, - -0.035462, - -0.004364, - 0.017992, - 0.101977, - -0.021944, - 0.019578, - -0.026028, - 0.050127, - -0.044843, - -0.005331, - 0.008791, - 0.024862, - 0.00735, - 0.006498, - 0.005705, - -0.001536, - 0.021862, - 0.012463, - -0.021342, - 0.037456, - -0.03329 + -0.006283, + -0.28829, + -0.510357, + -0.181786, + 0.267572, + -0.204303, + 0.065134, + -0.036725, + 0.221611, + -0.276, + 0.044667, + -0.319701, + -0.356479, + 0.162738, + 0.068047, + -0.341659, + -0.007111, + -0.286077, + -0.509385, + -0.181498, + 0.266028, + -0.204473, + 0.059981, + -0.032958, + 0.218527, + -0.278198, + 0.045874, + -0.318675, + -0.36582, + 0.169479, + 0.064059, + -0.336116 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + -0.001388, + 0.02555, + 0.001283, + 0.024534, + 0.000669, + 0.0236, + -0.000574, + 0.024204, + -9.2e-05, + -0.031526, + -0.032391, + -0.030826, + -0.032146, + -0.031856, + -0.032201, + -0.030922, + -0.032239, + -0.032286, + -0.05965, + 0.004568, + -0.064821, + 0.002639, + -0.055885, + 0.007288, + -0.060506, + 0.002031, + -0.053601, + -0.035833, + -0.021838, + -0.035347, + -0.020152, + -0.036121, + -0.020518, + -0.036075, + -0.020536, + -0.036254, + 0.030103, + -0.039094, + 0.030121, + -0.035787, + 0.030136, + -0.047276, + 0.030123, + -0.041288, + 0.030123, + 0.014357, + -0.020862, + 0.01455, + -0.023323, + 0.019685, + -0.02363, + 0.00853, + -0.020957, + 0.031693, + 0.003902, + 0.005723, + 0.007174, + 0.006236, + 0.003883, + 0.006472, + 0.015475, + 0.010045, + -0.001431, + -0.036287 ] - }, + } + ], + "ZMSentenceTransformer_multi_qa_distilbert_cos_v1": [ { "shape": [ 2, - 5, - 128 + 16 ], "sample": [ - 0.004865, - -0.003695, - 0.002367, - -0.000958, - -0.000965, - -0.000914, - 0.004061, - 5.2e-05, - -0.004466, - 0.000528, - -0.00554, - 0.002936, - 0.010828, - -0.003573, - -0.003714, - -0.004223, - -0.01237, - -0.00467, - -0.001861, - -0.004012, - -0.009534, - -0.005811, - -0.004497, - -0.0022, - 0.003601, - -0.002138, - 0.004155, - 0.009307, - -0.003024, - -0.002535, - -3.1e-05, - 0.004786, - -0.00753, - -0.003172, - 0.002225, - -0.001803, - -9.7e-05, - 0.003097, - 0.001154, - 0.008286, - -0.001478, - 0.002645, - 0.000562, - 0.008003, - -0.000749, - 0.00237, - -0.005276, - -0.004898, - -0.00232, - 0.003043, - -0.001344, - 0.001505, - 0.003906, - -0.001533, - 0.004861, - -0.002786, - 0.003926, - 0.001462, - 0.008712, - 0.0024, - 0.003343, - 0.001272, - 0.000233, - -0.000715 + 0.025973, + -0.206354, + 0.137715, + -0.007853, + -0.064003, + -0.047599, + -0.033647, + 0.035637, + -0.075466, + 0.007383, + 0.548441, + 0.416146, + -0.376682, + -0.05115, + 0.495402, + 0.243534, + 0.025587, + -0.19556, + 0.143738, + -0.007995, + -0.066312, + -0.045355, + -0.03309, + 0.023793, + -0.075691, + -0.024764, + 0.571854, + 0.411226, + -0.383017, + -0.0612, + 0.497094, + 0.179573 ] - } - ], - "WhisperModel": [ + }, { "shape": [ 2, - 20, - 64 + 16, + 16 ], "sample": [ - -0.016339, - 0.014459, - 0.027236, - -0.012696, - 0.003826, - 0.004806, - 0.00172, - 0.018539, - -0.003298, - -0.020976, - 0.037257, - -0.035281, - -0.004034, - 0.007173, - 0.031324, - -0.022888, - -0.013037, - -0.005674, - 0.017987, - 0.011197, - -0.007481, - 0.023746, - -0.024699, - 0.009173, - 0.0112, - -0.031971, - -0.023946, - -0.021795, - 0.052292, - -0.002748, - 0.00709, - 0.000376, - -0.006427, - -0.0098, - -0.010043, - -0.034723, - -0.013803, - 0.017374, - 0.028874, - -0.014603, - -0.010544, - 0.037403, - 0.020164, - -0.035462, - -0.004364, - 0.017992, - 0.101977, - -0.021944, - 0.019578, - -0.026028, - 0.050127, - -0.044843, - -0.005331, - 0.008791, - 0.024862, - 0.00735, - 0.006498, - 0.005705, - -0.001536, - 0.021862, - 0.012463, - -0.021342, - 0.037456, - -0.03329 + 0.00223, + -0.007174, + 0.002873, + -0.009289, + 0.003751, + -0.007166, + 0.003752, + -0.007555, + 0.001998, + 0.011331, + -0.02211, + -0.001803, + -0.023278, + 0.00989, + -0.0181, + -0.00481, + -0.031739, + 0.000239, + 0.015857, + 0.06128, + 0.015085, + 0.06705, + 0.017578, + 0.05266, + 0.013209, + 0.060971, + 0.015753, + 0.049514, + -0.001151, + 0.049367, + -0.001122, + 0.033595, + -0.000648, + 0.041646, + -0.000901, + 0.048977, + -0.010922, + -0.039426, + -0.004766, + -0.042949, + -0.007115, + -0.037974, + -0.003507, + -0.042364, + -0.006889, + -0.005795, + -0.003808, + -0.007997, + -0.003384, + 0.000805, + -0.004321, + -0.008732, + -0.005165, + -0.006883, + -0.002073, + 0.050016, + -0.005411, + 0.052228, + -0.006485, + 0.053975, + -0.00512, + 0.0522, + -0.004789, + 0.017832 ] - }, + } + ], + "ZMSentenceTransformer_multi_qa_mpnet_base_dot_v1": [ { "shape": [ 2, - 5, - 128 + 16 ], "sample": [ - 0.004865, - -0.003695, - 0.002367, - -0.000958, - -0.000965, - -0.000914, - 0.004061, - 5.2e-05, - -0.004466, - 0.000528, - -0.00554, - 0.002936, - 0.010828, - -0.003573, - -0.003714, - -0.004223, - -0.01237, - -0.00467, - -0.001861, - -0.004012, - -0.009534, - -0.005811, - -0.004497, - -0.0022, - 0.003601, - -0.002138, - 0.004155, - 0.009307, - -0.003024, - -0.002535, - -3.1e-05, - 0.004786, - -0.00753, - -0.003172, - 0.002225, - -0.001803, - -9.7e-05, - 0.003097, - 0.001154, - 0.008286, - -0.001478, - 0.002645, - 0.000562, - 0.008003, - -0.000749, - 0.00237, - -0.005276, - -0.004898, - -0.00232, - 0.003043, - -0.001344, - 0.001505, - 0.003906, - -0.001533, - 0.004861, - -0.002786, - 0.003926, - 0.001462, - 0.008712, - 0.0024, - 0.003343, - 0.001272, - 0.000233, - -0.000715 + 0.012509, + -0.027962, + -0.041582, + -0.018233, + 0.030212, + -0.002694, + 0.010153, + 0.004949, + -0.007654, + -0.025851, + 0.042763, + -0.02836, + -0.03605, + 0.034484, + 0.022787, + -0.014465, + 0.007949, + -0.027739, + -0.047459, + -0.018336, + 0.030239, + -0.001236, + 0.020397, + -0.011969, + -0.007566, + -0.027303, + 0.041396, + -0.038645, + -0.02535, + 0.034228, + 0.041315, + -0.024879 ] - } - ], - "XLMRobertaMaskedLM": [ + }, { "shape": [ 2, 16, - 128 + 16 ], "sample": [ - 0.013965, - 0.019017, - -0.003735, - -0.001879, - 0.010813, - 0.002046, - -0.009915, - -0.008338, - 0.019872, - 0.000889, - 0.04867, - 0.014064, - 0.0143, - 0.004872, - -0.004134, - 0.004227, - -0.055303, - 0.023829, - -0.035749, - -0.000273, - -0.019216, - -0.004134, - -0.017732, - -0.02736, - 0.016447, - 0.010392, - -0.032707, - -0.01479, - 0.009025, - -0.004914, - -0.008511, - 0.011608, - 0.015601, - -0.016102, - 0.003131, - 0.032635, - 0.021671, - 0.00695, - 0.003794, - -0.008977, - 0.008033, - -0.002371, - 0.003972, - -0.017868, - 0.026751, - 0.03861, - 0.016864, - -0.015283, - 0.004662, - -0.002961, - -0.041212, - -0.022634, - -0.044025, - 0.001969, - 0.009821, - -0.048096, - -0.034329, - 0.002701, - 0.014121, - -0.009398, - -0.017948, - -0.010799, - 0.017219, - -0.021088 + 0.012509, + -0.007654, + 0.012005, + -0.005402, + 0.012474, + -0.005269, + 0.013393, + -0.00231, + 0.01668, + -0.023257, + -0.027662, + -0.027021, + -0.028156, + -0.025185, + -0.027928, + -0.026351, + -0.02774, + -0.028348, + -0.031596, + 0.042527, + -0.038723, + 0.039582, + -0.045372, + 0.041086, + -0.044355, + 0.041554, + -0.03575, + -0.027495, + -0.018154, + -0.024365, + -0.019006, + -0.026141, + -0.018336, + -0.038645, + -0.018094, + -0.020396, + 0.030104, + -0.045994, + 0.030238, + -0.01341, + 0.030433, + -0.053719, + 0.03042, + -0.077768, + 0.03017, + 0.034797, + -0.002746, + 0.033777, + -0.006548, + 0.037556, + 0.012106, + 0.030717, + -0.003646, + 0.039079, + 0.038353, + 0.012978, + 0.040298, + 0.028337, + 0.011519, + 0.025773, + 0.027593, + 0.011926, + 0.014531, + -0.022591 ] } ], - "XLMRobertaModel": [ + "ZMSentenceTransformer_paraphrase_MiniLM_L3_v2": [ { "shape": [ 2, - 16, - 32 + 16 ], "sample": [ - -0.012946, - 0.015214, - -0.013106, - 0.015577, - -0.013179, - 0.069574, - 0.011613, - 0.070116, - 0.010722, - 0.024324, - 0.037582, - 0.024304, - 0.037166, - -0.035889, - -0.007053, - -0.036229, - -0.007861, - -0.073318, - -0.001794, - -0.073078, - -0.001493, - 0.031959, - -0.024997, - 0.033061, - -0.024845, - 0.032108, - 0.002208, - 0.016397, - 0.000681, - 0.016726, - 0.000822, - -0.005828, - 0.00068, - -0.00826, - 0.108422, - -0.049974, - 0.105383, - -0.04979, - -0.046333, - 0.025502, - -0.046555, - 0.026064, - 0.018798, - -0.003703, - 0.019448, - -0.00295, - 0.016, - 0.004491, - -0.02026, - 0.004788, - -0.020141, - 0.024671, - -0.017721, - 0.024861, - -0.017792, - 0.02249, - -0.024283, - 0.022482, - -0.024532, - -0.107456, - 0.000908, - -0.110674, - 0.000475, - -0.016094 + -0.000708, + -0.032463, + -0.05747, + -0.02047, + 0.03013, + -0.023006, + 0.007335, + -0.004135, + 0.024955, + -0.03108, + 0.00503, + -0.036001, + -0.040142, + 0.018325, + 0.007663, + -0.038473, + -0.000805, + -0.032396, + -0.057684, + -0.020553, + 0.030126, + -0.023155, + 0.006792, + -0.003732, + 0.024747, + -0.031504, + 0.005195, + -0.036088, + -0.041427, + 0.019192, + 0.007254, + -0.038063 ] }, { "shape": [ 2, - 32 + 16, + 16 ], "sample": [ - 9.5e-05, - 0.011957, - 0.006713, - 0.004725, - 0.011335, - -0.006035, - 0.012561, - 0.002068, - -0.0103, - 0.035051, - -0.010026, - 0.006029, - 0.002999, - 0.004093, - -0.02845, - 0.024154, - 0.007347, - 0.007143, - 0.001298, - -0.013618, - -0.011444, - -0.030955, - -0.005628, - -0.02546, - 0.018342, - 0.023768, - -0.011254, - -0.004289, - -0.019031, - 0.040206, - -0.025693, - -0.011637, - 8e-05, - 0.011829, - 0.006999, - 0.004674, - 0.011376, - -0.006081, - 0.012517, - 0.001985, - -0.010364, - 0.034915, - -0.010216, - 0.006072, - 0.003048, - 0.004378, - -0.028556, - 0.02404, - 0.007438, - 0.007219, - 0.001509, - -0.013589, - -0.011347, - -0.031104, - -0.00554, - -0.025427, - 0.018239, - 0.023746, - -0.01139, - -0.004372, - -0.018953, - 0.040002, - -0.025732, - -0.011578 + -0.001388, + 0.02555, + 0.001283, + 0.024534, + 0.000669, + 0.0236, + -0.000574, + 0.024204, + -9.2e-05, + -0.031526, + -0.032391, + -0.030826, + -0.032146, + -0.031856, + -0.032201, + -0.030922, + -0.032239, + -0.032286, + -0.05965, + 0.004568, + -0.064821, + 0.002639, + -0.055885, + 0.007288, + -0.060506, + 0.002031, + -0.053601, + -0.035833, + -0.021838, + -0.035347, + -0.020152, + -0.036121, + -0.020518, + -0.036075, + -0.020536, + -0.036254, + 0.030103, + -0.039094, + 0.030121, + -0.035787, + 0.030136, + -0.047276, + 0.030123, + -0.041288, + 0.030123, + 0.014357, + -0.020862, + 0.01455, + -0.023323, + 0.019685, + -0.02363, + 0.00853, + -0.020957, + 0.031693, + 0.003902, + 0.005723, + 0.007174, + 0.006236, + 0.003883, + 0.006472, + 0.015475, + 0.010045, + -0.001431, + -0.036287 + ] + } + ], + "ZMSentenceTransformer_paraphrase_albert_small_v2": [ + { + "shape": [ + 2, + 16 + ], + "sample": [ + 0.019625, + -0.000452, + -0.065652, + 0.038684, + -0.016475, + -0.077698, + 0.066542, + 0.033501, + -0.000845, + 0.026137, + -0.01556, + 0.010658, + -0.019583, + -0.014868, + -0.019154, + 0.056296, + 0.019622, + -0.0004, + -0.065506, + 0.038478, + -0.016457, + -0.077675, + 0.066516, + 0.033927, + -0.000842, + 0.026165, + -0.015646, + 0.010543, + -0.019583, + -0.014792, + -0.018992, + 0.055857 ] - } - ], - "XLMRobertaMultipleChoice": [ + }, { "shape": [ 2, - 3 + 16, + 16 ], "sample": [ - -0.015686, - -0.015691, - -0.01569, - -0.015684, - -0.015687, - -0.015699 + 0.019678, + -0.000842, + 0.019711, + -0.000838, + 0.01957, + -0.000845, + 0.019591, + -0.000842, + 0.019633, + 0.026074, + -0.000336, + 0.026192, + -0.000561, + 0.026158, + -0.000532, + 0.026085, + -0.000782, + 0.026046, + -0.066363, + -0.015558, + -0.065007, + -0.015725, + -0.066967, + -0.015305, + -0.065715, + -0.015685, + -0.064612, + 0.0103, + 0.038897, + 0.010929, + 0.039194, + 0.010795, + 0.038024, + 0.010261, + 0.03857, + 0.010324, + -0.016849, + -0.019532, + -0.016213, + -0.019658, + -0.016232, + -0.019634, + -0.017388, + -0.019531, + -0.015906, + -0.015022, + -0.077507, + -0.014948, + -0.077572, + -0.014514, + -0.07773, + -0.014773, + -0.077588, + -0.014602, + 0.066499, + -0.019658, + 0.066418, + -0.018982, + 0.066457, + -0.018579, + 0.066452, + -0.018653, + 0.066585, + 0.057405 ] } ], - "XLMRobertaQnA": [ + "ZMSentenceTransformer_paraphrase_multilingual_MiniLM_L12_v2": [ { "shape": [ 2, 16 ], "sample": [ - -0.040833, - -0.040412, - -0.040958, - -0.041051, - -0.040626, - -0.041051, - -0.040365, - -0.04072, - -0.040683, - -0.041041, - -0.040674, - -0.041258, - -0.040843, - -0.040957, - -0.040952, - -0.040941, - -0.040662, - -0.04115, - -0.040876, - -0.040881, - -0.040948, - -0.040716, - -0.041096, - -0.040604, - -0.041002, - -0.040509, - -0.040938, - -0.040975, - -0.041082, - -0.040653, - -0.040324, - -0.041059 + -0.000708, + -0.032463, + -0.05747, + -0.02047, + 0.03013, + -0.023006, + 0.007335, + -0.004135, + 0.024955, + -0.03108, + 0.00503, + -0.036001, + -0.040142, + 0.018325, + 0.007663, + -0.038473, + -0.000805, + -0.032396, + -0.057684, + -0.020553, + 0.030126, + -0.023155, + 0.006792, + -0.003732, + 0.024747, + -0.031504, + 0.005195, + -0.036088, + -0.041427, + 0.019192, + 0.007254, + -0.038063 ] }, { "shape": [ 2, + 16, 16 ], "sample": [ - 0.019887, - 0.020086, - 0.019876, - 0.020019, - 0.020148, - 0.020007, - 0.020131, - 0.019951, - 0.020173, - 0.020156, - 0.019967, - 0.019955, - 0.020004, - 0.019911, - 0.019939, - 0.020037, - 0.020038, - 0.019925, - 0.020181, - 0.02007, - 0.020088, - 0.020004, - 0.019789, - 0.020143, - 0.020097, - 0.019977, - 0.020205, - 0.019934, - 0.02004, - 0.019997, - 0.020336, - 0.01988 + -0.001388, + 0.02555, + 0.001283, + 0.024534, + 0.000669, + 0.0236, + -0.000574, + 0.024204, + -9.2e-05, + -0.031526, + -0.032391, + -0.030826, + -0.032146, + -0.031856, + -0.032201, + -0.030922, + -0.032239, + -0.032286, + -0.05965, + 0.004568, + -0.064821, + 0.002639, + -0.055885, + 0.007288, + -0.060506, + 0.002031, + -0.053601, + -0.035833, + -0.021838, + -0.035347, + -0.020152, + -0.036121, + -0.020518, + -0.036075, + -0.020536, + -0.036254, + 0.030103, + -0.039094, + 0.030121, + -0.035787, + 0.030136, + -0.047276, + 0.030123, + -0.041288, + 0.030123, + 0.014357, + -0.020862, + 0.01455, + -0.023323, + 0.019685, + -0.02363, + 0.00853, + -0.020957, + 0.031693, + 0.003902, + 0.005723, + 0.007174, + 0.006236, + 0.003883, + 0.006472, + 0.015475, + 0.010045, + -0.001431, + -0.036287 ] } ], - "XLMRobertaSequenceClassify": [ + "ZMSentenceTransformer_paraphrase_multilingual_mpnet_base_v2": [ { "shape": [ 2, - 4 + 16 ], "sample": [ - 0.000177, - 0.014164, - -0.036711, - -0.005497, - 0.000159, - 0.01417, - -0.036695, - -0.005501 + -0.000722, + -0.032459, + -0.05716, + -0.020443, + 0.03013, + -0.022946, + 0.00662, + -0.004355, + 0.024943, + -0.031073, + 0.005064, + -0.035998, + -0.040464, + 0.017516, + 0.007645, + -0.038519, + -0.000412, + -0.032375, + -0.057211, + -0.020587, + 0.030124, + -0.023147, + 0.006105, + -0.003692, + 0.024544, + -0.031693, + 0.005176, + -0.036129, + -0.040307, + 0.01879, + 0.007112, + -0.037763 ] - } - ], - "XLMRobertaTokenClassify": [ + }, { "shape": [ 2, 16, - 5 - ], - "sample": [ - -0.011828, - -0.052942, - -0.012009, - -0.052773, - -0.012009, - -0.05311, - -0.011534, - -0.052935, - -0.011814, - -0.052717, - -0.011641, - -0.052989, - -0.011957, - -0.052632, - -0.011815, - -0.052956, - -0.011606, - -0.052607, - -0.011437, - -0.052679, - -0.011876, - 0.026477, - -0.011723, - 0.026765, - -0.011958, - 0.026455, - -0.011909, - 0.026545, - -0.011798, - 0.026639, - -0.011807, - 0.026565, - -0.011974, - 0.026476, - -0.011909, - 0.02662, - -0.011605, - 0.026641, - -0.011655, - 0.026516, - -0.011716, - 0.026667, - 0.000539, - 0.026477, - 0.000706, - 0.02649, - 0.000529, - 0.026434, - 0.000589, - 0.026668, - 0.000565, - 0.026388, - 0.000571, - 0.026631, - 0.000628, - 0.026565, - 0.00061, - 0.026627, - 0.000546, - 0.026415, - 0.000415, - 0.026348, - 0.000654, - -0.05073 - ] - } - ], - "XceptionImageClassify": [ - { - "shape": [ - 2, - 1000 + 16 ], "sample": [ - 0.029016, - 0.032873, - -0.008282, - -0.01587, - 0.0041, - 0.033549, - -0.04755, - 0.04435, - 0.008561, - -0.067786, - 0.00482, - 0.00107, - 0.032935, - 0.019229, - 0.020901, - 0.014184, - 0.006072, - 0.011499, - 0.001939, - -0.001087, - 0.013775, - -0.011099, - 0.030425, - 0.030052, - -0.014473, - 0.007834, - 0.020641, - -0.035, - -0.002316, - 0.003105, - -0.000788, - -0.031078, - 0.034641, - 0.002396, - 0.009609, - -0.031307, - -0.005042, - 0.00558, - 0.00853, - 0.021877, - 0.036303, - -0.00843, - -0.054513, - -0.029453, - 0.032808, - -0.008157, - -0.008873, - -0.013846, - 0.004452, - 0.001395, - 0.027834, - 0.00415, - -0.018065, - -0.05072, - -0.03487, - -0.002989, - 0.052267, - 0.016791, - -0.008023, - -0.014243, - -0.007622, - -0.034641, - -0.014824, - -0.03292 + 0.00054, + 0.024421, + 0.00022, + 0.024225, + 0.000921, + 0.023616, + -0.00275, + 0.026056, + 8.3e-05, + -0.030825, + -0.032213, + -0.031122, + -0.032074, + -0.0312, + -0.032084, + -0.033028, + -0.032695, + -0.030283, + -0.064057, + 0.005521, + -0.059847, + 0.003993, + -0.058114, + 0.005001, + -0.053613, + 0.000932, + -0.055822, + -0.035447, + -0.021108, + -0.036471, + -0.020006, + -0.035757, + -0.020632, + -0.036374, + -0.020734, + -0.03649, + 0.030099, + -0.045783, + 0.030133, + -0.03234, + 0.030137, + -0.041763, + 0.030125, + -0.045943, + 0.030106, + 0.020362, + -0.021803, + 0.011972, + -0.022738, + 0.021537, + -0.023277, + 0.009515, + -0.022329, + 0.040238, + 0.014673, + 0.006805, + 0.01218, + 0.008585, + -0.000855, + 0.006342, + 0.002824, + 0.008836, + -0.000799, + -0.036171 ] } ] diff --git a/tests/fixtures/dummy_inputs.py b/tests/fixtures/dummy_inputs.py index 8a312a50..198dc36a 100644 --- a/tests/fixtures/dummy_inputs.py +++ b/tests/fixtures/dummy_inputs.py @@ -145,7 +145,7 @@ def siglip_input(batch_size=2, image_size=64, max_seq_len=64): def stable_diffusion_input( batch_size=2, sample_size=8, image_size=16, max_seq_len=16, cross_attention_dim=32 ): - # the container's three disconnected paths: UNet (sample / timestep / text + # ZMDiffusion's three disconnected paths: UNet (sample / timestep / text # context), VAE (image to encode, latent to decode) and the CLIP text tower return { "sample": ops.ones((batch_size, sample_size, sample_size, 4)), @@ -247,7 +247,7 @@ def qwen_image_input( joint_attention_dim=32, max_seq_len=32, ): - """Dummy inputs for the Qwen-Image container graph (packed latents).""" + """Dummy inputs for the Qwen-Image ZMDiffusion graph (packed latents).""" return { "sample": ops.ones((batch_size, packed_seq, in_channels)), "timestep": ops.ones((batch_size,)), @@ -275,7 +275,7 @@ def qwen_image_21_input( context_in_dim=32, max_seq_len=16, ): - """Dummy inputs for the Qwen-Image-2.1 container graph (unpatched latents).""" + """Dummy inputs for the Qwen-Image-2.1 ZMDiffusion graph (unpatched latents).""" return { "sample": ops.ones((batch_size, img_seq, in_channels)), "timestep": ops.ones((batch_size,)), diff --git a/tests/integration/test_auto_registry.py b/tests/integration/test_auto_registry.py index 43c41e6c..561b9ef0 100644 --- a/tests/integration/test_auto_registry.py +++ b/tests/integration/test_auto_registry.py @@ -71,6 +71,15 @@ ) +# One class serving several task tables: ZMDiffusion is the backbone, the text-to-image +# and the image-to-image class of every diffusion family (the family is the model_type); +# ZMSentenceTransformer is both the backbone and the text-embedding class. +_MULTI_TASK_CLASSES = { + "ZMDiffusion": {"Model", "TextToImage", "ImageToImage"}, + "ZMSentenceTransformer": {"Model", "TextEmbed"}, +} + + def _task_of(class_name): for suffix in _TASK_SUFFIXES: if class_name.endswith(suffix): @@ -86,12 +95,13 @@ def _hf_types(cls): def _iter_model_classes(): + import zeromodels.diffusers as diffusers import zeromodels.models as models import zeromodels.sentence_transformers as sentence_transformers from zeromodels.base.base_model import BaseModel seen = set() - for namespace in (models, sentence_transformers): + for namespace in (models, sentence_transformers, diffusers): packages = [namespace] packages.extend( getattr(namespace, family) @@ -268,6 +278,9 @@ def test_every_table_value_resolves_and_matches_its_task(): for task, table in names.MODEL_TASK_MAPPING_NAMES.items(): for model_type, class_name in table.items(): A._resolve_class(class_name) # raises if the class was renamed/removed + if class_name in _MULTI_TASK_CLASSES: + assert task in _MULTI_TASK_CLASSES[class_name], (class_name, task) + continue assert _task_of(class_name) == task, ( f"{class_name} is in the {task!r} table but its task is " f"{_task_of(class_name)!r}" @@ -294,9 +307,8 @@ def test_every_table_value_resolves_and_matches_its_task(): # tower and vice versa), so neither is a table default; load them via the concrete class. "Gemma4MultimodalModel", "Qwen3_5VLModel", - # Components of the Stable Diffusion containers (loaded through the family's - # XModel / XTextToImage), never hosted repos of their own; keep diffusers' names, - # which carry no task suffix. + # Components of a Stable Diffusion ZMDiffusion (loaded through it), never hosted + # repos of their own; keep diffusers' names, which carry no task suffix. "AutoencoderKL", "SD3Transformer2DModel", } @@ -350,20 +362,37 @@ def test_component_only_contract_has_real_unmapped_components_and_mapped_owners( assert component != owner +def test_diffusers_public_api_is_the_pipeline_and_the_tokenizer(): + import zeromodels.diffusers as diffusers + + assert diffusers.__all__ == ["ZMDiffusion", "ZMDiffusionTokenizer"] + + @pytest.mark.parametrize( - "model_type,model_name,pipeline_name", + "model_type", [ - ("qwen_image", "QwenImageModel", "QwenImageTextToImage"), - ("qwen_image_21", "QwenImage21Model", "QwenImage21TextToImage"), + "stable_diffusion", + "stable_diffusion_2", + "stable_diffusion_xl", + "stable_diffusion_xl_refiner", + "stable_diffusion_3", + "stable_diffusion_3_5", + "qwen_image", + "qwen_image_21", ], ) -def test_qwen_image_pipeline_auto_routes_remain_registered( - model_type, model_name, pipeline_name -): - assert names.MODEL_TASK_MAPPING_NAMES["Model"][model_type] == model_name - assert names.MODEL_TASK_MAPPING_NAMES["TextToImage"][model_type] == pipeline_name - assert model_name not in names.COMPONENT_ONLY_MODEL_NAMES - assert pipeline_name not in names.COMPONENT_ONLY_MODEL_NAMES +def test_diffusion_auto_routes_resolve_to_zm_diffusion(model_type): + """Every diffusion family loads through the one pipeline class and tokenizer: + the backbone and task autos name ZMDiffusion, the tokenizer auto names + ZMDiffusionTokenizer, and ZMDiffusion knows the model_type.""" + from zeromodels.diffusers.diffusion_model import DIFFUSION_FAMILIES + + task = "ImageToImage" if model_type.endswith("_refiner") else "TextToImage" + assert names.MODEL_TASK_MAPPING_NAMES["Model"][model_type] == "ZMDiffusion" + assert names.MODEL_TASK_MAPPING_NAMES[task][model_type] == "ZMDiffusion" + assert names.TOKENIZER_MAPPING_NAMES[model_type] == "ZMDiffusionTokenizer" + assert model_type in DIFFUSION_FAMILIES + assert "ZMDiffusion" not in names.COMPONENT_ONLY_MODEL_NAMES # Classes whose zm config model_type DELIBERATELY differs from the HF checkpoint's diff --git a/tests/integration/test_data_formats.py b/tests/integration/test_data_formats.py index 517db8b1..00610565 100644 --- a/tests/integration/test_data_formats.py +++ b/tests/integration/test_data_formats.py @@ -38,11 +38,9 @@ "Tipsv2DptDepthEstimation", "Tipsv2DptSemanticSegment", # Qwen-Image's Wan-derived VAE is built channels_last (NTHWC Conv3D). - "QwenImageModel", - "QwenImageTextToImage", + "ZMDiffusion_qwen_image", # Qwen-Image-2.1 residual VAE is channels_last (spatial Conv2d, RGBA). - "QwenImage21Model", - "QwenImage21TextToImage", + "ZMDiffusion_qwen_image_21", # Qwen-VL inputs are pre-patchified (no spatial axes) -> layout-agnostic. "Qwen2VLModel", "Qwen2_5VLModel", diff --git a/tests/integration/test_sentence_transformer.py b/tests/integration/test_sentence_transformer.py index 0edd9c53..a095a635 100644 --- a/tests/integration/test_sentence_transformer.py +++ b/tests/integration/test_sentence_transformer.py @@ -1,6 +1,5 @@ """Offline contract tests for the multi-backbone sentence-embedding API.""" -import importlib import json import numpy as np @@ -9,36 +8,36 @@ from tests.base.model_test_registry import MODEL_TEST_CONFIGS from zeromodels.sentence_transformers import ( - SentenceTransformerTextEmbed, ZMSentenceTransformer, ZMSentenceTransformerTokenizer, ) +from zeromodels.sentence_transformers.sentence_transformer_backbones import ( + SENTENCE_TRANSFORMER_BACKBONES, +) -@pytest.mark.parametrize( - "module", - ( - "zeromodels.models.sentence_transformer", - "zeromodels.st_models", - "zeromodels.st_models.sentence_transformer", - "zeromodels.sentence_transformers.sentence_transformer", - ), -) -def test_top_level_st_api_and_legacy_imports_resolve_to_same_classes(module): - legacy = importlib.import_module(module) - package = ( - module + ".sentence_transformer" if module.endswith("st_models") else module - ) - legacy_model_module = importlib.import_module( - package + ".sentence_transformer_model" - ) - assert ZMSentenceTransformer.__name__ == "ZMSentenceTransformer" - assert ZMSentenceTransformerTokenizer.__name__ == "ZMSentenceTransformerTokenizer" - assert legacy.ZMSentenceTransformer is ZMSentenceTransformer - assert legacy_model_module.ZMSentenceTransformer is ZMSentenceTransformer - assert legacy.SentenceTransformerModel is ZMSentenceTransformer - assert legacy_model_module.SentenceTransformerModel is ZMSentenceTransformer - assert legacy.SentenceTransformerTokenizer is ZMSentenceTransformerTokenizer +def test_public_api_is_the_model_and_the_tokenizer(): + import zeromodels.sentence_transformers as sentence_transformers + + assert sentence_transformers.__all__ == [ + "ZMSentenceTransformer", + "ZMSentenceTransformerTokenizer", + ] + for removed in ( + "SentenceTransformer", + "SentenceTransformerModel", + "SentenceTransformerTokenizer", + "SentenceTransformerTextEmbed", + "SentenceTransformerConfig", + ): + assert not hasattr(sentence_transformers, removed), removed + + +def test_backbone_registry_matches_accepted_hf_model_types(): + assert set(ZMSentenceTransformer.HF_MODEL_TYPE) == { + "sentence_transformer", + *SENTENCE_TRANSFORMER_BACKBONES, + } @pytest.mark.parametrize( @@ -136,7 +135,7 @@ def test_backbone_builds_and_returns_sentence_embeddings(backbone_type): "input_ids": np.array([[2, 3, 4, 1]], dtype="int32"), "attention_mask": np.array([[1, 1, 1, 0]], dtype="int32"), } - if backbone_type in ("bert", "roberta", "xlm-roberta", "albert"): + if SENTENCE_TRANSFORMER_BACKBONES[backbone_type].token_type_ids: inputs["token_type_ids"] = np.zeros((1, 4), dtype="int32") actual = model(inputs, training=False) assert actual["embeddings"].shape == (1, 16) @@ -179,10 +178,7 @@ def test_pooling_parser_rejects_ambiguous_modes(monkeypatch, tmp_path): def test_sentence_transformer_registered_for_text_embed_task(): from zeromodels.auto import AutoZMTextEmbed - assert ( - AutoZMTextEmbed.mapping()["sentence_transformer"] - is SentenceTransformerTextEmbed - ) + assert AutoZMTextEmbed.mapping()["sentence_transformer"] is ZMSentenceTransformer def test_dot_similarity_is_not_cosine_normalized(): diff --git a/website/mkdocs.yml b/website/mkdocs.yml index 365e80aa..7dd90248 100644 --- a/website/mkdocs.yml +++ b/website/mkdocs.yml @@ -172,6 +172,15 @@ nav: - Sentence Transformer models: - sentence_transformers.md - Sentence Transformer: sentence_transformer.md + - Diffusion models: + - diffusers.md + - Stable Diffusion: stable_diffusion.md + - Stable Diffusion 2: stable_diffusion_2.md + - Stable Diffusion XL: stable_diffusion_xl.md + - Stable Diffusion 3: stable_diffusion_3.md + - Stable Diffusion 3.5: stable_diffusion_3_5.md + - Qwen-Image: qwen_image.md + - Qwen-Image-2.1: qwen_image_21.md - Vision models: - BEiT: beit.md - Classification backbones: classification_backbones.md @@ -235,13 +244,6 @@ nav: - SAM 3: sam3.md - SigLIP: siglip.md - SigLIP 2: siglip2.md - - Stable Diffusion: stable_diffusion.md - - Stable Diffusion 2: stable_diffusion_2.md - - Stable Diffusion XL: stable_diffusion_xl.md - - Stable Diffusion 3: stable_diffusion_3.md - - Stable Diffusion 3.5: stable_diffusion_3_5.md - - Qwen-Image: qwen_image.md - - Qwen-Image-2.1: qwen_image_21.md - TIPSv2: tipsv2.md - Loading Weights: loading_weights.md diff --git a/zeromodels/__init__.py b/zeromodels/__init__.py index d85bd976..cf3de14d 100644 --- a/zeromodels/__init__.py +++ b/zeromodels/__init__.py @@ -1,6 +1,7 @@ from zeromodels import ( auto, conversion, + diffusers, models, samplers, sentence_transformers, @@ -20,6 +21,7 @@ __all__ = [ "auto", + "diffusers", "models", "samplers", "sentence_transformers", diff --git a/zeromodels/auto/__init__.py b/zeromodels/auto/__init__.py index f57401fa..30e994e7 100644 --- a/zeromodels/auto/__init__.py +++ b/zeromodels/auto/__init__.py @@ -6,8 +6,8 @@ classes (``AutoZModel``, ``AutoZMDetect``, ``AutoZMImageClassify``, ``AutoZMTextGenerate``, ...) plus ``AutoZMConfig`` / ``AutoZMTokenizer`` / ``AutoZMProcessor`` / ``AutoZMImageProcessor`` are generated from the committed task tables. Class discovery -includes :mod:`zeromodels.models` and :mod:`zeromodels.sentence_transformers`; this -package re-exports whatever the factory built. +includes :mod:`zeromodels.models`, :mod:`zeromodels.sentence_transformers` and +:mod:`zeromodels.diffusers`; this package re-exports whatever the factory built. """ from zeromodels.auto import auto_factory diff --git a/zeromodels/auto/auto_factory.py b/zeromodels/auto/auto_factory.py index 45f71086..4005f289 100644 --- a/zeromodels/auto/auto_factory.py +++ b/zeromodels/auto/auto_factory.py @@ -103,11 +103,12 @@ def _load_repo_config(identifier): def _class_index(): global _CLASS_INDEX if _CLASS_INDEX is None: + import zeromodels.diffusers as diffusers import zeromodels.models as models import zeromodels.sentence_transformers as sentence_transformers index = {} - for namespace in (models, sentence_transformers): + for namespace in (models, sentence_transformers, diffusers): packages = [namespace] packages.extend( getattr(namespace, family) diff --git a/zeromodels/auto/auto_mapping_names.py b/zeromodels/auto/auto_mapping_names.py index edfee13d..4c7da3b6 100644 --- a/zeromodels/auto/auto_mapping_names.py +++ b/zeromodels/auto/auto_mapping_names.py @@ -11,17 +11,18 @@ the zeromodels-repo and ``hf:`` load paths share one table. """ -# Public building blocks of full diffusion pipelines, not standalone Auto tasks. -# Values name the Auto-loadable owner exposing vae/text_encoder/transformer. +# Public building blocks of the diffusion pipeline, not standalone Auto tasks. +# Values name the Auto-loadable owner exposing vae/text_encoder/transformer +# (ZMDiffusion, which builds them for the Qwen-Image families). # Direct construction and serialization remain supported; a config model_type # alone does not promise an independent Auto checkpoint-loading route. COMPONENT_ONLY_MODEL_NAMES = { - "AutoencoderKLQwenImage": "QwenImageModel", - "QwenImageTextEncoderModel": "QwenImageModel", - "QwenImageTransformer2DModel": "QwenImageModel", - "AutoencoderKLQwenImage21": "QwenImage21Model", - "QwenImage21TextEncoderModel": "QwenImage21Model", - "QwenImage21Transformer2DModel": "QwenImage21Model", + "AutoencoderKLQwenImage": "ZMDiffusion", + "QwenImageTextEncoderModel": "ZMDiffusion", + "QwenImageTransformer2DModel": "ZMDiffusion", + "AutoencoderKLQwenImage21": "ZMDiffusion", + "QwenImage21TextEncoderModel": "ZMDiffusion", + "QwenImage21Transformer2DModel": "ZMDiffusion", } # task suffix -> {model_type: ClassName}. AutoZModel is task 'Model'; @@ -289,14 +290,14 @@ "sentence_transformer": "ZMSentenceTransformer", "siglip": "SigLIPModel", "siglip2": "SigLIP2Model", - "stable_diffusion": "StableDiffusionModel", - "stable_diffusion_2": "StableDiffusion2Model", - "stable_diffusion_xl": "StableDiffusionXLModel", - "stable_diffusion_xl_refiner": "StableDiffusionXLRefinerModel", - "stable_diffusion_3": "StableDiffusion3Model", - "stable_diffusion_3_5": "StableDiffusion3_5Model", - "qwen_image": "QwenImageModel", - "qwen_image_21": "QwenImage21Model", + "stable_diffusion": "ZMDiffusion", + "stable_diffusion_2": "ZMDiffusion", + "stable_diffusion_xl": "ZMDiffusion", + "stable_diffusion_xl_refiner": "ZMDiffusion", + "stable_diffusion_3": "ZMDiffusion", + "stable_diffusion_3_5": "ZMDiffusion", + "qwen_image": "ZMDiffusion", + "qwen_image_21": "ZMDiffusion", "swin": "SwinModel", "swinv2": "SwinV2Model", "t5": "T5Model", @@ -370,7 +371,7 @@ }, "TextEmbed": { "clip": "CLIPTextEmbed", - "sentence_transformer": "SentenceTransformerTextEmbed", + "sentence_transformer": "ZMSentenceTransformer", }, "TextGenerate": { "deepseek_v2": "DeepseekV2TextGenerate", @@ -430,16 +431,16 @@ "tipsv2": "Tipsv2TextModel", }, "TextToImage": { - "stable_diffusion": "StableDiffusionTextToImage", - "stable_diffusion_2": "StableDiffusion2TextToImage", - "stable_diffusion_xl": "StableDiffusionXLTextToImage", - "stable_diffusion_3": "StableDiffusion3TextToImage", - "stable_diffusion_3_5": "StableDiffusion3_5TextToImage", - "qwen_image": "QwenImageTextToImage", - "qwen_image_21": "QwenImage21TextToImage", + "stable_diffusion": "ZMDiffusion", + "stable_diffusion_2": "ZMDiffusion", + "stable_diffusion_xl": "ZMDiffusion", + "stable_diffusion_3": "ZMDiffusion", + "stable_diffusion_3_5": "ZMDiffusion", + "qwen_image": "ZMDiffusion", + "qwen_image_21": "ZMDiffusion", }, "ImageToImage": { - "stable_diffusion_xl_refiner": "StableDiffusionXLRefinerImageToImage", + "stable_diffusion_xl_refiner": "ZMDiffusion", }, "TokenClassify": { "bert": "BertTokenClassify", @@ -808,14 +809,14 @@ "siglip": "SigLIPTokenizer", "siglip2": "SigLIP2Tokenizer", "speech_to_text": "Speech2TextTokenizer", - "stable_diffusion": "StableDiffusionTokenizer", - "stable_diffusion_2": "StableDiffusion2Tokenizer", - "stable_diffusion_xl": "StableDiffusionXLTokenizer", - "stable_diffusion_xl_refiner": "StableDiffusionXLTokenizer", - "stable_diffusion_3": "StableDiffusion3Tokenizer", - "stable_diffusion_3_5": "StableDiffusion3_5Tokenizer", - "qwen_image": "QwenImageTokenizer", - "qwen_image_21": "QwenImage21Tokenizer", + "stable_diffusion": "ZMDiffusionTokenizer", + "stable_diffusion_2": "ZMDiffusionTokenizer", + "stable_diffusion_xl": "ZMDiffusionTokenizer", + "stable_diffusion_xl_refiner": "ZMDiffusionTokenizer", + "stable_diffusion_3": "ZMDiffusionTokenizer", + "stable_diffusion_3_5": "ZMDiffusionTokenizer", + "qwen_image": "ZMDiffusionTokenizer", + "qwen_image_21": "ZMDiffusionTokenizer", "t5": "T5Tokenizer", "tipsv2": "Tipsv2Tokenizer", "whisper": "WhisperTokenizer", diff --git a/zeromodels/base/base_diffusion.py b/zeromodels/base/base_diffusion.py index 65515713..100b27fb 100644 --- a/zeromodels/base/base_diffusion.py +++ b/zeromodels/base/base_diffusion.py @@ -9,10 +9,10 @@ class BaseDiffusion: """Backend-agnostic latent-diffusion generation, the peer of :class:`BaseGeneration` (mirrors diffusers' ``DiffusionPipeline.__call__``). - A mixin added to a diffusion container (the family's ``XModel``, which builds - its denoiser / VAE / text-encoder towers internally) to give it ``generate``, - exactly as ``T5ConditionalGenerate(T5Model, BaseSeq2SeqGeneration)`` adds an - imperative decode loop to a functional backbone. It bundles the machinery every + A mixin added to a functional diffusion model (:class:`ZMDiffusion`, which + builds its denoiser / VAE / text-encoder towers internally) to give it + ``generate``, exactly as ``T5ConditionalGenerate(T5Model, BaseSeq2SeqGeneration)`` + adds an imperative decode loop to a functional backbone. It bundles the machinery every diffusion task shares: resolving the generation settings, batching the conditioning for classifier-free guidance, sampling the initial latent, the scheduler's denoising loop, decoding the latent and the uint8 postprocess. The diff --git a/zeromodels/diffusers/__init__.py b/zeromodels/diffusers/__init__.py new file mode 100644 index 00000000..e5196624 --- /dev/null +++ b/zeromodels/diffusers/__init__.py @@ -0,0 +1,7 @@ +from .diffusion_model import ZMDiffusion +from .diffusion_tokenizer import ZMDiffusionTokenizer + +__all__ = [ + "ZMDiffusion", + "ZMDiffusionTokenizer", +] diff --git a/zeromodels/diffusers/diffusion_model.py b/zeromodels/diffusers/diffusion_model.py new file mode 100644 index 00000000..108458de --- /dev/null +++ b/zeromodels/diffusers/diffusion_model.py @@ -0,0 +1,325 @@ +import functools + +import keras + +from zeromodels.base import BaseConfig, BaseModel +from zeromodels.base.base_diffusion import BaseDiffusion +from zeromodels.base.base_model import strip_functional_graph_keys +from zeromodels.base.base_scheduler import get_scheduler + +from .qwen_image.qwen_image_model import QwenImageFamily +from .qwen_image_21.qwen_image_21_model import QwenImage21Family +from .stable_diffusion.stable_diffusion_model import StableDiffusionFamily +from .stable_diffusion_2.stable_diffusion_2_model import StableDiffusion2Family +from .stable_diffusion_3.stable_diffusion_3_model import ( + SD3T5EncoderModel, + StableDiffusion3Family, +) +from .stable_diffusion_3_5.stable_diffusion_3_5_model import StableDiffusion3_5Family +from .stable_diffusion_xl.stable_diffusion_xl_model import ( + StableDiffusionXLFamily, + StableDiffusionXLRefinerFamily, +) + +DIFFUSION_FAMILIES = { + family.model_type: family() + for family in ( + StableDiffusionFamily, + StableDiffusion2Family, + StableDiffusionXLFamily, + StableDiffusionXLRefinerFamily, + StableDiffusion3Family, + StableDiffusion3_5Family, + QwenImageFamily, + QwenImage21Family, + ) +} + +# Classes that hosted zm_config.json files named before ZMDiffusion replaced the +# per-family container and pipeline classes; their repos load unchanged (same +# model_type, same weights layout). +LEGACY_DIFFUSION_CLASSES = frozenset( + { + "ZMDiffusion", + "StableDiffusionModel", + "StableDiffusionTextToImage", + "StableDiffusion2Model", + "StableDiffusion2TextToImage", + "StableDiffusionXLModel", + "StableDiffusionXLTextToImage", + "StableDiffusionXLRefinerModel", + "StableDiffusionXLRefinerImageToImage", + "StableDiffusion3Model", + "StableDiffusion3TextToImage", + "StableDiffusion3_5Model", + "StableDiffusion3_5TextToImage", + "QwenImageModel", + "QwenImageTextToImage", + "QwenImage21Model", + "QwenImage21TextToImage", + } +) + + +def get_diffusion_family(model_type): + family = DIFFUSION_FAMILIES.get(model_type) + if family is None: + raise ValueError( + f"Unknown diffusion model_type {model_type!r}. ZMDiffusion supports: " + f"{', '.join(sorted(DIFFUSION_FAMILIES))}." + ) + return family + + +class ZMDiffusionMeta(type(BaseModel)): + """Carries a family config's ``model_type`` into the constructor kwargs. + + ``ZMDiffusion(StableDiffusionXLConfig())`` is expanded by the base metaclass + into the config's flat fields, which do not include ``model_type`` (a class + attribute of the config); this adds it so the family can be selected, and keeps + ``model.config is config`` when nothing overrides it, as for every BaseModel. + """ + + def __call__(cls, *args, **kwargs): + if not (args and isinstance(args[0], BaseConfig)): + return super().__call__(*args, **kwargs) + config, overridden = args[0], bool(kwargs) + kwargs.setdefault("model_type", config.model_type) + model = super().__call__(*args, **kwargs) + if not overridden: + model.config = config + return model + + +@keras.saving.register_keras_serializable(package="zeromodels") +class ZMDiffusion(BaseModel, BaseDiffusion, metaclass=ZMDiffusionMeta): + """Text-to-image and image-to-image diffusion, pure Keras 3 and cross-backend. + + The one diffusion pipeline class of zeromodels (diffusers' ``DiffusionPipeline`` + in this library's form), for every hosted family: Stable Diffusion 1.x / 2.x / + XL (base, refiner, Turbo), Stable Diffusion 3 / 3.5, Qwen-Image and + Qwen-Image-2.1. A checkpoint's ``model_type`` (recorded in its + ``zm_config.json``) selects the family; the family decides which components + are built and how a prompt becomes an image, so loading and generating look the + same for all of them:: + + from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer + + repo = "zeromodels/stable-diffusion-v1-5" + pipe = ZMDiffusion.from_weights(repo) + tokenizer = ZMDiffusionTokenizer.from_weights(repo) + image = pipe.generate(**tokenizer("a photo of a cat")) # (1, 512, 512, 3) uint8 + + One functional ``keras.Model`` holds every trained component as disconnected + paths (``{"sample", "timestep", "encoder_hidden_states", ...} -> "noise_pred"`` + for the denoiser, ``{"image", "latent"} -> "moments" / "image"`` for the VAE, + the token ids to the prompt embeddings for the text towers), and exposes them as + attributes (``.unet`` or ``.transformer``, ``.vae``, ``.text_encoder``, + ``.text_encoder_2``) that ``generate`` runs step by step. The weights are + converted once from diffusers checkpoints and hosted under + ``zeromodels/``; on-the-fly ``hf:`` conversion is not supported. + + ``generate`` (:class:`BaseDiffusion`) takes the tokenizer's output plus + ``negative_input_ids``, ``num_inference_steps``, ``guidance_scale``, ``seed``, + ``latents`` (explicit initial noise, for results identical across backends) and + ``output_type`` (``"image"`` uint8 or ``"latent"``). Image-to-image takes an + ``image`` and a ``strength``; ``denoising_end`` / ``denoising_start`` split a run + between the SDXL base and refiner. Family-specific arguments: SDXL's + ``original_size`` / ``crops_coords_top_left`` / ``target_size`` (and + ``aesthetic_score`` on the refiner), SD 3's ``input_ids_3`` (the T5 tokens), and + Qwen-Image's ``height`` / ``width`` / ``negative_attention_mask``. Defaults come + from the repo's ``generate_args``. + + SD 3 / 3.5 can attach the separately hosted T5-XXL encoder: + ``ZMDiffusion.from_weights(repo, text_encoder_3="zeromodels/t5-v1_1-xxl-encoder")`` + or ``pipe.text_encoder_3 = SD3T5EncoderModel.from_weights(...)``; it lives + outside the checkpoint's weights. Without it the T5 features are zeros. + + Args: + model_type: The family, e.g. ``"stable_diffusion"``, + ``"stable_diffusion_xl"``, ``"stable_diffusion_3_5"``, ``"qwen_image"``. + Read from the repo on ``from_weights``. + scheduler: A :class:`BaseScheduler`. Defaults to the one the config's + ``scheduler_config`` describes (the checkpoint's own sampler), else the + family's default. Swap any time with ``pipe.scheduler = ...``. + **kwargs: The family config's flat fields (``unet_`` / ``transformer_`` / + ``vae_`` / ``text_`` prefixed), or pass a family config positionally: + ``ZMDiffusion(StableDiffusionXLConfig())``. Pass ``unet_sample_size`` / + ``vae_sample_size`` to build for another resolution. + """ + + # one class serves every family, so the typed config is the family's own + # (resolved from model_type in __init__ and _config_from_zm_spec), not a class + # attribute the base metaclass would apply + config_class = None + HF_MODEL_TYPE = None + HUB_REPO_SIBLINGS = LEGACY_DIFFUSION_CLASSES + + def __init__(self, model_type=None, scheduler=None, name="ZMDiffusion", **kwargs): + family = get_diffusion_family(model_type) + keras_kwargs = {k: kwargs.pop(k) for k in ("trainable", "dtype") if k in kwargs} + config = family.config_class.from_dict(kwargs) + + data_format = keras.config.image_data_format() + channels_axis = -1 if data_format == "channels_last" else 1 + + components = family.build_components(config, data_format, channels_axis) + inputs, outputs = family.build_graph(config, components, data_format) + super().__init__(inputs=inputs, outputs=outputs, name=name, **keras_kwargs) + + self.model_type = model_type + self.family = family + self.config = config + self.data_format = data_format + self.channels_axis = channels_axis + for attr, component in components.items(): + setattr(self, attr, component) + self.generate_args = dict(getattr(family, "generate_args", None) or {}) + if scheduler is None: + scheduler = ( + get_scheduler(config.scheduler_config) + if getattr(config, "scheduler_config", None) + else family.default_scheduler() + ) + self.scheduler = scheduler + + @classmethod + def from_weights(cls, identifier, text_encoder_3=None, **kwargs): + model = super().from_weights(identifier, **kwargs) + if text_encoder_3 is not None: + if not getattr(model.family, "supports_text_encoder_3", False): + raise ValueError( + f"text_encoder_3 is only used by Stable Diffusion 3 / 3.5, not " + f"by a {model.model_type!r} checkpoint." + ) + if isinstance(text_encoder_3, str): + text_encoder_3 = SD3T5EncoderModel.from_weights( + text_encoder_3, load_dtype=kwargs.get("load_dtype") + ) + model.text_encoder_3 = text_encoder_3 + return model + + @classmethod + def from_variant(cls, variant, /, quantization=None, **kwargs): + # a bare name is the hosted repo's basename; `quantization` is applied by + # from_weights after the load, exactly as on the org/repo branch + return cls.from_hub_repo(f"zeromodels/{variant}", **kwargs) + + @classmethod + def from_hf(cls, repo, **kwargs): + raise NotImplementedError( + "On-the-fly `hf:` conversion is not supported for diffusion models. Load " + "the preconverted weights from a zeromodels repo, e.g. " + "ZMDiffusion.from_weights('zeromodels/stable-diffusion-v1-5'), or run the " + f"family's convert_*_diffusers_to_keras.py to produce them from {repo!r}." + ) + + @classmethod + def _config_from_zm_spec(cls, spec, variant): + # parse the repo's zm_config.json with its family's typed config (refusing + # keys this version does not know, as the base loader does for typed + # configs) and keep the model_type the flat kwargs would otherwise lose + if spec is None or isinstance(spec.get("config"), dict): + return super()._config_from_zm_spec(spec, variant) + from zeromodels.conversion.zm_config import ZM_METADATA_KEYS + + fields = {k: v for k, v in spec.items() if k not in ZM_METADATA_KEYS} + model_type = fields.get("model_type") + config_cls = get_diffusion_family(model_type).config_class + unknown = config_cls.unknown_keys(fields) + if unknown: + raise ValueError( + f"ZMDiffusion.from_weights: '{variant}' zm_config.json has config " + f"key(s) {sorted(unknown)} that this zeromodels version does not " + f"recognize. The repo is likely newer than the installed zeromodels: " + f"upgrade it, or remove the key(s) if you are hand-editing the config." + ) + kwargs = config_cls.from_dict(fields).constructor_kwargs() + return {"model_type": model_type, **kwargs} + + @staticmethod + def _apply_generate_args(model, spec): + # merged over the family's defaults (an instance attribute here, since one + # class serves every family) + if spec is not None and spec.get("generate_args"): + model.generate_args = {**model.generate_args, **spec["generate_args"]} + + def get_config(self): + config = strip_functional_graph_keys(super().get_config()) + config.update(self.config.constructor_kwargs()) + config["model_type"] = self.model_type + return config + + @property + def text_encoder_3(self): + return self.__dict__.get("text_encoder_3") + + def __setattr__(self, name, value): + # the T5 encoder is kept out of the tracked (saved / loaded / device-moved) + # sub-layers: Keras (and torch's Module on that backend) would otherwise + # register any model assigned as an attribute + if name == "text_encoder_3": + self.__dict__["text_encoder_3"] = value + return + super().__setattr__(name, value) + + @property + def latent_shape(self): + return self.family.latent_shape(self) + + def default_scheduler(self): + return self.family.default_scheduler() + + def unconditional_ids(self, batch): + return self.family.unconditional_ids(self, batch) + + def encode_prompt(self, input_ids, attention_mask=None, **conditioning): + return self.family.encode_prompt( + self, input_ids, attention_mask, **conditioning + ) + + def family_hook(self, name): + # the family's own version of a BaseDiffusion hook, bound to this model, or + # None when the family keeps the shared implementation + hook = getattr(self.family, name, None) + return None if hook is None else functools.partial(hook, self) + + def encode_negative_prompt(self, negative_input_ids, batch, **conditioning): + hook = ( + self.family_hook("encode_negative_prompt") or super().encode_negative_prompt + ) + return hook(negative_input_ids, batch, **conditioning) + + def predict_noise(self, latents, timesteps, embeddings): + return self.family.predict_noise(self, latents, timesteps, embeddings) + + def decode_latents(self, latents, **kwargs): + return self.family.decode_latents(self, latents, **kwargs) + + def encode_latents(self, image): + hook = self.family_hook("encode_latents") + if hook is None: + raise NotImplementedError( + f"{self.model_type!r} checkpoints do not support image-to-image." + ) + return hook(image) + + def prepare_latents(self, batch, seed=None, latents=None, dtype="float32"): + hook = self.family_hook("prepare_latents") or super().prepare_latents + return hook(batch, seed=seed, latents=latents, dtype=dtype) + + def denoise( + self, latents, embeddings, num_inference_steps, guidance_scale, timesteps=None + ): + hook = self.family_hook("denoise") or super().denoise + return hook( + latents, + embeddings, + num_inference_steps, + guidance_scale, + timesteps=timesteps, + ) + + def generate(self, input_ids, **kwargs): + hook = self.family_hook("generate") or super().generate + return hook(input_ids, **kwargs) diff --git a/zeromodels/diffusers/diffusion_tokenizer.py b/zeromodels/diffusers/diffusion_tokenizer.py new file mode 100644 index 00000000..1ff6a3af --- /dev/null +++ b/zeromodels/diffusers/diffusion_tokenizer.py @@ -0,0 +1,247 @@ +import keras +from keras import ops + +from zeromodels.base import BaseTokenizer +from zeromodels.models.clip.clip_tokenizer import CLIPTokenizer +from zeromodels.models.qwen2.qwen2_tokenizer import Qwen2Tokenizer +from zeromodels.models.qwen3.qwen3_tokenizer import Qwen3Tokenizer +from zeromodels.models.t5.t5_tokenizer import T5Tokenizer + +QWEN_IMAGE_PROMPT_TEMPLATE = ( + "<|im_start|>system\nDescribe the image by detailing the color, shape, size, " + "texture, quantity, text, spatial relationships of the objects and " + "background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" +) +QWEN_IMAGE_21_PROMPT_TEMPLATE = ( + "<|im_start|>system\nComprehend and analyze the provided prompt.<|im_end|>\n" + "<|im_start|>user\n{}<|im_end|>\n" + "<|im_start|>assistant\n" +) + +# Per model_type: the tokenizer backend and its defaults. "clip": the CLIP BPE +# tokenizer padded to a fixed length (SDXL and SD 3 derive their second tower's +# "!"-padded ids from its mask; SD 3 adds the T5 tokenizer from tokenizer_3.json). +# "qwen": a Qwen BPE tokenizer wrapping each prompt in the pipeline's ChatML +# template (the template prefix is dropped after the text encoder). +DIFFUSION_TOKENIZERS = { + "stable_diffusion": {"kind": "clip", "pad_token": "<|endoftext|>"}, + "stable_diffusion_2": {"kind": "clip", "pad_token": "!"}, + "stable_diffusion_xl": {"kind": "clip", "pad_token": "<|endoftext|>"}, + "stable_diffusion_xl_refiner": {"kind": "clip", "pad_token": "<|endoftext|>"}, + "stable_diffusion_3": {"kind": "clip", "pad_token": "<|endoftext|>", "t5": True}, + "stable_diffusion_3_5": { + "kind": "clip", + "pad_token": "<|endoftext|>", + "t5": True, + }, + "qwen_image": { + "kind": "qwen", + "tokenizer_class": Qwen2Tokenizer, + "prompt_template": QWEN_IMAGE_PROMPT_TEMPLATE, + "prompt_template_start_idx": 34, + "empty_prompt": "", + }, + "qwen_image_21": { + "kind": "qwen", + "tokenizer_class": Qwen3Tokenizer, + "prompt_template": QWEN_IMAGE_21_PROMPT_TEMPLATE, + "prompt_template_start_idx": 14, + "empty_prompt": " ", + }, +} +DEFAULT_MAX_SEQ_LEN = {"clip": 77, "qwen": 1024} + + +@keras.saving.register_keras_serializable(package="zeromodels") +class ZMDiffusionTokenizer(BaseTokenizer): + """The prompt tokenizer of every :class:`ZMDiffusion` checkpoint. + + One class for all diffusion families, selected by ``model_type`` like the model: + loading by repo id reads it from the repo (``zm_preprocessor.json`` when present, + else ``zm_config.json``) together with the repo's ``tokenizer.json``:: + + tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-xl-base-1.0") + image = pipe.generate(**tokenizer("a photo of a cat")) + + - Stable Diffusion 1.x / XL (base and refiner): CLIP ViT-L/14's byte-pair + encoding with ``<|startoftext|>`` / ``<|endoftext|>`` framing, truncated and + ``<|endoftext|>``-padded to 77 tokens. SDXL's second text tower pads with + ``!`` instead; :class:`ZMDiffusion` derives those ids from the returned + ``attention_mask``, so one call feeds ``generate``. + - Stable Diffusion 2.x: the same encoding padded with ``!`` (id 0), the way + OpenCLIP ViT-H/14 was trained. + - Stable Diffusion 3 / 3.5: the CLIP tokens plus the T5 SentencePiece tokens of + the repo's ``tokenizer_3.json`` (````-terminated, ````-padded to + ``max_sequence_length``, 256) as ``input_ids_3``. + - Qwen-Image / Qwen-Image-2.1: each prompt wrapped in the pipeline's ChatML + template, then Qwen2 / Qwen3 byte-pair encoding, right-padded to the longest + prompt (``input_ids`` / ``attention_mask``). + + Args: + model_type: The diffusion family (see :class:`ZMDiffusion`). + hf_id: Hosted repo to read ``tokenizer.json`` (and ``tokenizer_3.json``) from. + Required unless ``tokenizer_file`` is given (no default repo). + tokenizer_file: Explicit ``tokenizer.json`` path (overrides ``hf_id``). + tokenizer_file_3: Explicit T5 ``tokenizer.json`` path (SD 3 / 3.5). + max_seq_len: Padded / truncated length (77 for CLIP; the prompt length + limit for Qwen, 1024, not counting the template prefix). + max_sequence_length: T5 padded / truncated length (SD 3 / 3.5, 256). + pad_token: CLIP pad token string (``<|endoftext|>``, ``!`` for SD 2.x). + """ + + def __init__( + self, + model_type=None, + hf_id=None, + tokenizer_file=None, + tokenizer_file_3=None, + max_seq_len=None, + max_sequence_length=None, + pad_token=None, + **kwargs, + ): + spec = DIFFUSION_TOKENIZERS.get(model_type) + if spec is None: + raise ValueError( + f"Unknown diffusion model_type {model_type!r}. ZMDiffusionTokenizer " + f"supports: {', '.join(sorted(DIFFUSION_TOKENIZERS))}." + ) + if tokenizer_file is None and hf_id is None: + raise ValueError( + "ZMDiffusionTokenizer() needs an hf_id (a hosted repo, read from its " + "tokenizer.json) or an explicit tokenizer_file: there is no default " + "repo. Load it by repo id, e.g. " + "from_weights('zeromodels/stable-diffusion-v1-5')." + ) + super().__init__(**kwargs) + self.model_type = model_type + self.kind = spec["kind"] + self.hf_id = hf_id + self.tokenizer_file = self.resolve_tokenizer_json_from_hf(hf_id, tokenizer_file) + if max_seq_len is None: + max_seq_len = DEFAULT_MAX_SEQ_LEN[self.kind] + self.max_seq_len = max_seq_len + self.pad_token = None + if self.kind == "clip": + self.pad_token = pad_token or spec["pad_token"] + self.tokenizer = CLIPTokenizer( + tokenizer_file=self.tokenizer_file, + max_seq_len=max_seq_len, + pad_token=self.pad_token, + ) + else: + self.prompt_template = spec["prompt_template"] + self.prompt_template_start_idx = spec["prompt_template_start_idx"] + self.empty_prompt = spec["empty_prompt"] + self.tokenizer = spec["tokenizer_class"](tokenizer_file=self.tokenizer_file) + self.tokenizer_file_3 = None + self.max_sequence_length = None + self.tokenizer_3 = None + if spec.get("t5"): + self.tokenizer_file_3 = self.resolve_tokenizer_json_from_hf( + hf_id, tokenizer_file_3, filename="tokenizer_3.json" + ) + self.max_sequence_length = max_sequence_length or 256 + self.tokenizer_3 = T5Tokenizer( + tokenizer_file=self.tokenizer_file_3, + max_seq_len=self.max_sequence_length, + ) + + @classmethod + def from_hub_repo(cls, repo_id, **kwargs): + from zeromodels.conversion.zm_config import ( + ZM_METADATA_KEYS, + load_zm_config, + load_zm_preprocessor, + ) + + repo_id = repo_id.rstrip("/") + spec = load_zm_preprocessor(repo_id) + if spec is not None: + config = {k: v for k, v in spec.items() if k not in ZM_METADATA_KEYS} + else: + # repos without zm_preprocessor.json: the family comes from the model's + # zm_config.json, the rest from the family defaults + model_spec = load_zm_config(repo_id) + if not model_spec or "model_type" not in model_spec: + raise OSError( + f"{repo_id} has neither zm_preprocessor.json nor a zm_config.json " + f"with a model_type, so the diffusion family is unknown. Pass " + f"model_type= explicitly." + ) + config = {"model_type": model_spec["model_type"]} + config.update(kwargs) + config.setdefault("hf_id", repo_id) + return cls(**config) + + @classmethod + def from_variant(cls, variant, /, **kwargs): + return cls.from_hub_repo(f"zeromodels/{variant}", **kwargs) + + @classmethod + def from_hf(cls, repo, **kwargs): + raise NotImplementedError( + "On-the-fly `hf:` loading is not supported for diffusion models: use " + "ZMDiffusionTokenizer.from_weights('zeromodels/')." + ) + + @property + def vocab_size(self): + return self.tokenizer.vocab_size + + def format_prompt(self, text): + return self.prompt_template.format(text if text else self.empty_prompt) + + def call(self, inputs): + if self.kind == "qwen": + texts = self.tokenizer.normalize_texts(inputs) + max_length = self.max_seq_len + self.prompt_template_start_idx + encoded = [ + self.tokenizer.encode(self.format_prompt(t))[:max_length] for t in texts + ] + input_ids, attention_mask = self.pad_batch(encoded) + return {"input_ids": input_ids, "attention_mask": attention_mask} + out = self.tokenizer(inputs) + if self.tokenizer_3 is not None: + # the T5 tokenizer truncates to max_sequence_length and pads to the + # longest prompt; the transformer takes the fixed length, so pad the rest + ids = self.tokenizer_3(inputs)["input_ids"] + short = self.max_sequence_length - int(ids.shape[1]) + if short > 0: + ids = ops.pad( + ids, + ((0, 0), (0, short)), + constant_values=self.tokenizer_3.pad_token_id, + ) + out["input_ids_3"] = ids + return out + + def decode(self, ids, skip_special_tokens=True): + if self.kind == "qwen": + return self.tokenizer.decode(ids, skip_special_tokens=skip_special_tokens) + return self.tokenizer.detokenize(ids) + + def get_config(self): + config = super().get_config() + config.update( + { + "model_type": self.model_type, + "hf_id": self.hf_id, + "tokenizer_file": self.tokenizer_file, + "tokenizer_file_3": self.tokenizer_file_3, + "max_seq_len": self.max_seq_len, + "max_sequence_length": self.max_sequence_length, + "pad_token": self.pad_token, + } + ) + return config + + def get_hub_config(self): + # the portable zm_preprocessor.json recipe: the family and its settings, + # with the tokenizer files read from the same repo (no local paths) + config = {"model_type": self.model_type, "max_seq_len": self.max_seq_len} + if self.pad_token is not None: + config["pad_token"] = self.pad_token + if self.max_sequence_length is not None: + config["max_sequence_length"] = self.max_sequence_length + return config diff --git a/zeromodels/models/qwen_image/__init__.py b/zeromodels/diffusers/qwen_image/__init__.py similarity index 73% rename from zeromodels/models/qwen_image/__init__.py rename to zeromodels/diffusers/qwen_image/__init__.py index d7fe9377..d4807983 100644 --- a/zeromodels/models/qwen_image/__init__.py +++ b/zeromodels/diffusers/qwen_image/__init__.py @@ -6,21 +6,17 @@ ) from .qwen_image_model import ( AutoencoderKLQwenImage, - QwenImageModel, + QwenImageFamily, QwenImageTextEncoderModel, - QwenImageTextToImage, QwenImageTransformer2DModel, ) -from .qwen_image_tokenizer import QwenImageTokenizer __all__ = [ "AutoencoderKLQwenImage", "QwenImageConfig", - "QwenImageModel", + "QwenImageFamily", "QwenImageTextConfig", "QwenImageTextEncoderModel", - "QwenImageTextToImage", - "QwenImageTokenizer", "QwenImageTransformer2DModel", "QwenImageTransformerConfig", "QwenImageVAEConfig", diff --git a/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py b/zeromodels/diffusers/qwen_image/convert_qwen_image_diffusers_to_keras.py similarity index 97% rename from zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py rename to zeromodels/diffusers/qwen_image/convert_qwen_image_diffusers_to_keras.py index d7ee8d5b..3d1a0620 100644 --- a/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py +++ b/zeromodels/diffusers/qwen_image/convert_qwen_image_diffusers_to_keras.py @@ -51,11 +51,11 @@ def config_from_diffusers(repo, token=None): from huggingface_hub import hf_hub_download - from zeromodels.models.qwen_image.qwen_image_config import ( + from zeromodels.diffusers.qwen_image.qwen_image_config import ( QwenImageConfig, QwenImageVAEConfig, ) - from zeromodels.models.qwen_image.qwen_image_model import ( + from zeromodels.diffusers.qwen_image.qwen_image_model import ( QwenImageTransformer2DModel, ) @@ -127,16 +127,16 @@ def transfer_qwen_image( from safetensors import safe_open from zeromodels.base.base_mixin import build_dtype_scope - from zeromodels.models.qwen_image.qwen_image_model import QwenImageModel + from zeromodels.diffusers.diffusion_model import ZMDiffusion config = config or config_from_diffusers(repo, token=token) flat = config.constructor_kwargs() flat["transformer_sample_size"] = build_sample_size flat["vae_sample_size"] = max(build_sample_size * 8, 64) - print(f"[1/4] Building QwenImageModel (dtype={dtype})…", flush=True) + print(f"[1/4] Building ZMDiffusion (qwen_image, dtype={dtype})…", flush=True) with build_dtype_scope(dtype), zeros_init(): - model = QwenImageModel(**flat) + model = ZMDiffusion(model_type=config.model_type, **flat) vae_mapping = { k: v for k, v in WEIGHT_NAME_MAPPING.items() if k not in ("/gamma", "gamma") diff --git a/zeromodels/models/qwen_image/qwen_image_config.py b/zeromodels/diffusers/qwen_image/qwen_image_config.py similarity index 96% rename from zeromodels/models/qwen_image/qwen_image_config.py rename to zeromodels/diffusers/qwen_image/qwen_image_config.py index 9adf5cb0..b58c1c59 100644 --- a/zeromodels/models/qwen_image/qwen_image_config.py +++ b/zeromodels/diffusers/qwen_image/qwen_image_config.py @@ -118,9 +118,9 @@ class QwenImageTextConfig(Qwen2_5VLTextConfig): class QwenImageConfig(BaseConfig): - r"""Configuration for [`QwenImageModel`] / [`QwenImageTextToImage`]. + r"""Configuration for a Qwen-Image [`ZMDiffusion`] (``model_type="qwen_image"``). - One hosted container: MMDiT transformer + Qwen-Image VAE + Qwen2.5-VL text + One hosted checkpoint: MMDiT transformer + Qwen-Image VAE + Qwen2.5-VL text encoder. Nested serialize (``transformer_config`` / ``vae_config`` / ``text_config``); flat constructor with ``transformer_`` / ``vae_`` / ``text_`` prefixes. diff --git a/zeromodels/models/qwen_image/qwen_image_layers.py b/zeromodels/diffusers/qwen_image/qwen_image_layers.py similarity index 99% rename from zeromodels/models/qwen_image/qwen_image_layers.py rename to zeromodels/diffusers/qwen_image/qwen_image_layers.py index c3e98b0a..a6506125 100644 --- a/zeromodels/models/qwen_image/qwen_image_layers.py +++ b/zeromodels/diffusers/qwen_image/qwen_image_layers.py @@ -5,11 +5,11 @@ from keras import layers, ops from zeromodels.base.base_attention import active_attn_implementation, fused_attention -from zeromodels.models.stable_diffusion.stable_diffusion_layers import ( +from zeromodels.diffusers.stable_diffusion.stable_diffusion_layers import ( safe_name, timestep_embedding, ) -from zeromodels.models.stable_diffusion_3.stable_diffusion_3_layers import ( +from zeromodels.diffusers.stable_diffusion_3.stable_diffusion_3_layers import ( StableDiffusion3AdaLayerNorm, StableDiffusion3GELUFeedForward, ) diff --git a/zeromodels/models/qwen_image/qwen_image_model.py b/zeromodels/diffusers/qwen_image/qwen_image_model.py similarity index 69% rename from zeromodels/models/qwen_image/qwen_image_model.py rename to zeromodels/diffusers/qwen_image/qwen_image_model.py index ff7bda67..894c760c 100644 --- a/zeromodels/models/qwen_image/qwen_image_model.py +++ b/zeromodels/diffusers/qwen_image/qwen_image_model.py @@ -4,15 +4,11 @@ import numpy as np from keras import layers, ops -from zeromodels.base import BaseDiffusion, BaseModel -from zeromodels.base.base_attention import with_model_attn_implementation +from zeromodels.base import BaseModel +from zeromodels.base.base_attention import use_attn_implementation from zeromodels.base.base_mixin import inference_scope -from zeromodels.base.base_scheduler import ( - FlowMatchEulerDiscreteScheduler, - get_scheduler, -) -from zeromodels.models.qwen2_5_vl.qwen2_5_vl_model import Qwen2_5VLModel -from zeromodels.models.qwen_image.qwen_image_config import ( +from zeromodels.base.base_scheduler import FlowMatchEulerDiscreteScheduler +from zeromodels.diffusers.qwen_image.qwen_image_config import ( DEFAULT_LATENTS_MEAN, DEFAULT_LATENTS_STD, QwenImageConfig, @@ -20,26 +16,20 @@ QwenImageTransformerConfig, QwenImageVAEConfig, ) -from zeromodels.models.qwen_image.qwen_image_layers import ( +from zeromodels.diffusers.qwen_image.qwen_image_layers import ( QwenImageAdaLayerNormContinuous, QwenImageEmbedRope, QwenImageRMSNorm, QwenImageTimestepProjEmbeddings, QwenImageTransformerBlock, ) -from zeromodels.models.qwen_image.qwen_image_vae import ( +from zeromodels.diffusers.qwen_image.qwen_image_vae import ( QwenImageCausalConv3d, QwenImageDecoder3d, QwenImageEncoder3d, ) -from zeromodels.models.stable_diffusion.stable_diffusion_layers import safe_name - -QWEN_IMAGE_HUB_SIBLINGS = frozenset({"QwenImageModel", "QwenImageTextToImage"}) -PROMPT_TEMPLATE = ( - "<|im_start|>system\nDescribe the image by detailing the color, shape, size, " - "texture, quantity, text, spatial relationships of the objects and " - "background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" -) +from zeromodels.diffusers.stable_diffusion.stable_diffusion_layers import safe_name +from zeromodels.models.qwen2_5_vl.qwen2_5_vl_model import Qwen2_5VLModel def pack_latents(latents, height, width): @@ -421,7 +411,7 @@ def kwargs_from_diffusers_config(cls, cfg): class QwenImageTextEncoderModel(Qwen2_5VLModel): """Qwen-Image prompt encoder: the Qwen2.5-VL-7B text tower, no vision / LM head. - The ``text_encoder`` component of :class:`QwenImageModel`. + The ``text_encoder`` component of a Qwen-Image :class:`ZMDiffusion`. Inputs ``input_ids`` / ``attention_mask``; output ``last_hidden_state``. """ @@ -434,31 +424,29 @@ def __init__(self, max_seq_len=1024, name="text_encoder", **kwargs): self.max_seq_len = max_seq_len -@keras.saving.register_keras_serializable(package="zeromodels") -class QwenImageModel(BaseModel): - """Qwen-Image weights container: transformer + VAE + Qwen2.5-VL text tower. +class QwenImageFamily: + """Qwen-Image inside :class:`ZMDiffusion` (``model_type="qwen_image"``). - One functional ``keras.Model`` with disconnected paths (Diffusers - ``QwenImagePipeline`` components). Hosted as ``zeromodels/qwen-image`` - (sharded weights). On-the-fly ``hf:`` conversion is not supported. + Three components, ``transformer`` (the double-stream MMDiT, + :class:`QwenImageTransformer2DModel`), ``vae`` (the Wan-derived + :class:`AutoencoderKLQwenImage`) and ``text_encoder`` (the Qwen2.5-VL-7B text + tower), wired into three disconnected paths (diffusers' ``QwenImagePipeline`` + components). Hosted as ``zeromodels/qwen-image`` (sharded weights). + + Generation follows the reference: the ChatML-templated prompt through the text + tower with the template prefix dropped, packed 2x2 latents, the flow-match Euler + sampler with resolution-dependent shifting, and true classifier-free guidance + (``guidance_scale``, diffusers' ``true_cfg_scale``: separate conditional and + unconditional passes, the combination renormalized to the conditional + prediction's norm). Without a negative prompt the unconditional branch is the + templated blank prompt. ``height`` / ``width`` choose the output size. """ + model_type = "qwen_image" config_class = QwenImageConfig - HF_MODEL_TYPE = None - HUB_REPO_SIBLINGS = QWEN_IMAGE_HUB_SIBLINGS - - def __init__(self, name="QwenImageModel", **kwargs): - keras_kwargs = {k: kwargs.pop(k) for k in ("trainable", "dtype") if k in kwargs} - config = self.config_class.from_dict(kwargs) - - components = self.build_components(config) - inputs, outputs = self.build_graph(config, components) - super().__init__(inputs=inputs, outputs=outputs, name=name, **keras_kwargs) - - for attr, component in components.items(): - setattr(self, attr, component) + generate_args = {"num_inference_steps": 50, "guidance_scale": 4.0} - def build_components(self, config): + def build_components(self, config, data_format, channels_axis): d, v = config.transformer_config, config.vae_config transformer = QwenImageTransformer2DModel( patch_size=d.patch_size, @@ -489,7 +477,7 @@ def build_components(self, config): text_encoder = QwenImageTextEncoderModel(config.text_config) return {"transformer": transformer, "vae": vae, "text_encoder": text_encoder} - def build_graph(self, config, components): + def build_graph(self, config, components, data_format): d, v, t = config.transformer_config, config.vae_config, config.text_config transformer, vae, text_encoder = ( components["transformer"], @@ -544,50 +532,6 @@ def build_graph(self, config, components): "prompt_embeds": text_out["last_hidden_state"], } - def get_config(self): - config = super().get_config() - config.update(self.config.constructor_kwargs()) - return config - - def from_hf(self, *args, **kwargs): - raise NotImplementedError( - "On-the-fly hf: conversion is not supported for Qwen-Image; " - "use convert_qwen_image_diffusers_to_keras.py and " - "from_weights('zeromodels/qwen-image')." - ) - - -@keras.saving.register_keras_serializable(package="zeromodels") -class QwenImageTextToImage(QwenImageModel, BaseDiffusion): - """Text-to-image Qwen-Image (Diffusers ``QwenImagePipeline``). - - :: - - model = QwenImageTextToImage.from_weights("zeromodels/qwen-image") - tok = QwenImageTokenizer.from_weights("zeromodels/qwen-image") - image = model.generate(**tok("a cat"), height=1024, width=1024) - - Uses true CFG (``guidance_scale`` / Diffusers ``true_cfg_scale``) with - dual forward passes and prediction-norm renormalization. Packed latents + - flow-match Euler with dynamic resolution shifting. - """ - - config_class = QwenImageConfig - HUB_REPO_SIBLINGS = QWEN_IMAGE_HUB_SIBLINGS - generate_args = {"num_inference_steps": 50, "guidance_scale": 4.0} - DEFAULT_GUIDANCE_SCALE = 4.0 - - def __init__(self, scheduler=None, name="QwenImageTextToImage", **kwargs): - scheduler_config = kwargs.get("scheduler_config") - super().__init__(name=name, **kwargs) - if scheduler is None: - scheduler = ( - get_scheduler(scheduler_config) - if scheduler_config - else self.default_scheduler() - ) - self.scheduler = scheduler - def default_scheduler(self): return FlowMatchEulerDiscreteScheduler( shift=1.0, @@ -600,30 +544,37 @@ def default_scheduler(self): time_shift_type="exponential", ) - @property - def vae_scale_factor(self): - return self.vae.vae_scale_factor - - @property - def latent_shape(self): - # Packed shape used by prepare_latents after packing. - h, w = self._latent_side() - channels = self.vae.z_dim - return ((h // 2) * (w // 2), channels * 4) - - def _latent_side(self, height=None, width=None): - scale = self.vae_scale_factor * 2 # VAE 8× and pack 2× + def latent_side(self, model, height=None, width=None): + # diffusers: round the size down to a multiple of vae_scale_factor * 2 (the + # VAE downsampling times the 2x2 packing), in latent pixels + scale = model.vae.vae_scale_factor * 2 if height is None or width is None: - side = self.config.default_sample_size * self.vae_scale_factor + side = model.config.default_sample_size * model.vae.vae_scale_factor height = width = side - # Match Diffusers: round down to multiple of vae_scale_factor*2, then *2 latent. - h = 2 * (int(height) // scale) - w = 2 * (int(width) // scale) - return h, w + return 2 * (int(height) // scale), 2 * (int(width) // scale) + + def generation_side(self, model): + height, width = model.__dict__.get("generation_size", (None, None)) + return self.latent_side(model, height, width) - def _negative_prompt_inputs(self, batch, prompt=None): - """Tokenize the templated negative prompt used for true CFG.""" - from .qwen_image_tokenizer import QwenImageTokenizer + def image_seq_len(self, height, width): + return (height // 2) * (width // 2) + + def latent_shape(self, model): + # the packed shape, at the default resolution + h, w = self.latent_side(model) + return (self.image_seq_len(h, w), model.vae.z_dim * 4) + + def pack_latents(self, latents, height, width): + return pack_latents(latents, height, width) + + def unpack_latents(self, latents, height, width, channels): + return unpack_latents(latents, height, width, channels) + + def negative_prompt_inputs(self, model, batch, prompt=None): + # the templated negative prompt of true CFG (a blank prompt by default), + # tokenized with the hosted Qwen-Image tokenizer + from zeromodels.diffusers.diffusion_tokenizer import ZMDiffusionTokenizer if prompt is None: prompts = [" "] * batch @@ -633,26 +584,22 @@ def _negative_prompt_inputs(self, batch, prompt=None): prompts = list(prompt) if len(prompts) != batch: raise ValueError("negative_prompt must have one entry per image") - tokenizer = getattr(self, "_negative_prompt_tokenizer", None) + tokenizer = model.__dict__.get("negative_prompt_tokenizer") if tokenizer is None: - tokenizer = QwenImageTokenizer.from_weights("zeromodels/qwen-image") - self._negative_prompt_tokenizer = tokenizer + tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/qwen-image") + model.__dict__["negative_prompt_tokenizer"] = tokenizer return tokenizer(prompts) - def unconditional_ids(self, batch): - """Token IDs for a templated space prompt, matching Diffusers CFG.""" + def unconditional_ids(self, model, batch): return ops.convert_to_tensor( - self._negative_prompt_inputs(batch)["input_ids"], dtype="int32" + self.negative_prompt_inputs(model, batch)["input_ids"], dtype="int32" ) - def encode_prompt(self, input_ids, attention_mask=None, **conditioning): - """Encode ChatML-templated token ids → truncated prompt embeds + mask. - - Expects ``input_ids`` already wrapped with the Diffusers prompt template - (see :class:`QwenImageTokenizer`). Drops the template prefix - (``prompt_template_encode_start_idx``) and pads/truncates to - ``max_sequence_length``. - """ + def encode_prompt(self, model, input_ids, attention_mask=None, **conditioning): + # expects ids already wrapped in the ChatML prompt template (as + # ZMDiffusionTokenizer returns them); drops the template prefix + # (prompt_template_encode_start_idx) and pads / truncates to + # max_sequence_length del conditioning input_ids = ops.cast(ops.convert_to_tensor(input_ids), "int32") if attention_mask is None: @@ -660,15 +607,14 @@ def encode_prompt(self, input_ids, attention_mask=None, **conditioning): else: attention_mask = ops.cast(ops.convert_to_tensor(attention_mask), "int32") - out = self.text_encoder( + out = model.text_encoder( {"input_ids": input_ids, "attention_mask": attention_mask} ) hidden = out["last_hidden_state"] - drop = int(self.config.prompt_template_encode_start_idx) - max_len = int(self.config.max_sequence_length) + drop = int(model.config.prompt_template_encode_start_idx) + max_len = int(model.config.max_sequence_length) - # Gather non-padding tokens per row, drop template prefix, then pad. - # Implemented with numpy for clarity on host; convert back to tensors. + # gather each row's real tokens on the host, drop the template prefix, pad hidden_np = ops.convert_to_numpy(hidden) mask_np = ops.convert_to_numpy(attention_mask) batch = hidden_np.shape[0] @@ -686,51 +632,64 @@ def encode_prompt(self, input_ids, attention_mask=None, **conditioning): "encoder_hidden_states_mask": ops.convert_to_tensor(out_mask), } - @with_model_attn_implementation - def predict_noise(self, latents, timesteps, embeddings): - return self.transformer( - { - "sample": latents, - "timestep": timesteps, - "encoder_hidden_states": embeddings["encoder_hidden_states"], - "encoder_hidden_states_mask": embeddings["encoder_hidden_states_mask"], - } - )["sample"] - - def decode_latents(self, latents, height=None, width=None): - h, w = self._latent_side(height, width) - latents = unpack_latents(latents, h, w, self.vae.z_dim) - # Diffusers: latents / (1/std) + mean == latents * std + mean - mean = ops.convert_to_tensor(self.vae.latents_mean, dtype="float32") - std = ops.convert_to_tensor(self.vae.latents_std, dtype="float32") + def predict_noise(self, model, latents, timesteps, embeddings): + inputs = { + "sample": latents, + "timestep": timesteps, + "encoder_hidden_states": embeddings["encoder_hidden_states"], + "encoder_hidden_states_mask": embeddings["encoder_hidden_states_mask"], + } + impl = getattr(model, "_attn_implementation", None) + if impl is None: + return model.transformer(inputs)["sample"] + with use_attn_implementation(impl): + return model.transformer(inputs)["sample"] + + def decode_latents(self, model, latents, height=None, width=None): + h, w = self.latent_side(model, height, width) + latents = self.unpack_latents(latents, h, w, model.vae.z_dim) + # diffusers: latents / (1 / std) + mean == latents * std + mean + mean = ops.convert_to_tensor(model.vae.latents_mean, dtype="float32") + std = ops.convert_to_tensor(model.vae.latents_std, dtype="float32") mean = ops.reshape(mean, (1, 1, 1, -1)) std = ops.reshape(std, (1, 1, 1, -1)) latents = latents * std + mean - return self.vae.decode(latents) + return model.vae.decode(latents) - def prepare_latents(self, batch, seed=None, latents=None, dtype="float32"): - h, w = self._latent_side( - getattr(self, "_gen_height", None), getattr(self, "_gen_width", None) - ) - channels = self.vae.z_dim + def prepare_latents(self, model, batch, seed=None, latents=None, dtype="float32"): + h, w = self.generation_side(model) if latents is None: - # Spatial noise then pack (matches Diffusers prepare_latents). - noise = keras.random.normal((batch, h, w, channels), seed=seed, dtype=dtype) - latents = pack_latents(noise, h, w) + # spatial noise, then packed (diffusers' prepare_latents) + noise = keras.random.normal( + (batch, h, w, model.vae.z_dim), seed=seed, dtype=dtype + ) + latents = self.pack_latents(noise, h, w) else: latents = ops.cast(ops.convert_to_tensor(latents), dtype) - return latents * self.scheduler.init_noise_sigma + return latents * model.scheduler.init_noise_sigma + + def combine_guidance(self, noise_pred, neg_pred, guidance_scale): + comb = neg_pred + guidance_scale * (noise_pred - neg_pred) + cond_norm = ops.sqrt(ops.sum(ops.square(noise_pred), axis=-1, keepdims=True)) + comb_norm = ops.sqrt(ops.sum(ops.square(comb), axis=-1, keepdims=True)) + return comb * (cond_norm / (comb_norm + 1e-8)) def denoise( - self, latents, embeddings, num_inference_steps, guidance_scale, timesteps=None + self, + model, + latents, + embeddings, + num_inference_steps, + guidance_scale, + timesteps=None, ): - """True CFG: separate cond/uncond forwards + norm renormalization.""" + # true CFG: separate conditional / unconditional passes; embeddings is the + # (uncond, cond) pair under guidance do_cfg = guidance_scale > 1.0 - scheduler = self.scheduler + scheduler = model.scheduler if timesteps is None: scheduler.set_timesteps(num_inference_steps) timesteps = scheduler.timesteps - # embeddings may be a pair (uncond, cond) under CFG from generate(). if do_cfg and isinstance(embeddings, (tuple, list)): uncond_emb, cond_emb = embeddings else: @@ -739,20 +698,55 @@ def denoise( for t in timesteps: t_batch = ops.full((ops.shape(latents)[0],), float(t) / 1000.0) - noise_pred = self.predict_noise(latents, t_batch, cond_emb) + noise_pred = model.predict_noise(latents, t_batch, cond_emb) if do_cfg and uncond_emb is not None: - neg_pred = self.predict_noise(latents, t_batch, uncond_emb) - comb = neg_pred + guidance_scale * (noise_pred - neg_pred) - cond_norm = ops.sqrt( - ops.sum(ops.square(noise_pred), axis=-1, keepdims=True) - ) - comb_norm = ops.sqrt(ops.sum(ops.square(comb), axis=-1, keepdims=True)) - noise_pred = comb * (cond_norm / (comb_norm + 1e-8)) + neg_pred = model.predict_noise(latents, t_batch, uncond_emb) + noise_pred = self.combine_guidance(noise_pred, neg_pred, guidance_scale) latents = scheduler.step(noise_pred, t, latents) return latents + def negative_embeddings( + self, + model, + batch, + negative_input_ids, + negative_attention_mask, + guidance_scale, + conditioning, + ): + # the unconditional branch, or None to run without guidance + if guidance_scale <= 1.0: + return None + if negative_input_ids is None: + negative_inputs = self.negative_prompt_inputs( + model, batch, conditioning.get("negative_prompt") + ) + negative_input_ids = negative_inputs["input_ids"] + negative_attention_mask = negative_inputs["attention_mask"] + return model.encode_prompt(negative_input_ids, negative_attention_mask) + + def set_timesteps(self, model, num_inference_steps, height, width): + # diffusers' calculate_shift: the resolution-dependent flow-match mu + h, w = self.latent_side(model, height, width) + image_seq_len = self.image_seq_len(h, w) + sched_cfg = getattr(model.scheduler, "config_dict", None) or {} + base_seq_len = sched_cfg.get("base_image_seq_len", 256) + max_seq_len = sched_cfg.get("max_image_seq_len", 4096) + base_shift = sched_cfg.get("base_shift", 0.5) + max_shift = sched_cfg.get("max_shift", 0.9) + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + mu = image_seq_len * m + (base_shift - m * base_seq_len) + sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps) + if hasattr(model.scheduler, "set_timesteps"): + try: + model.scheduler.set_timesteps(num_inference_steps, sigmas=sigmas, mu=mu) + except TypeError: + model.scheduler.set_timesteps(num_inference_steps) + return model.scheduler.timesteps + def generate( self, + model, input_ids, attention_mask=None, negative_input_ids=None, @@ -766,58 +760,34 @@ def generate( output_type="image", **conditioning, ): - num_inference_steps, guidance_scale, seed = self.resolve_generation_args( + num_inference_steps, guidance_scale, seed = model.resolve_generation_args( num_inference_steps, guidance_scale, seed ) - self._gen_height = height - self._gen_width = width + model.__dict__["generation_size"] = (height, width) input_ids = ops.cast(ops.convert_to_tensor(input_ids), "int32") batch = int(input_ids.shape[0]) with inference_scope(): - embeddings = self.encode_prompt(input_ids, attention_mask, **conditioning) - do_cfg = guidance_scale > 1.0 - if do_cfg: - if negative_input_ids is None: - negative_inputs = self._negative_prompt_inputs( - batch, conditioning.get("negative_prompt") - ) - negative_input_ids = negative_inputs["input_ids"] - negative_attention_mask = negative_inputs["attention_mask"] - uncond = self.encode_prompt(negative_input_ids, negative_attention_mask) - else: - uncond = None - - # Diffusers calculate_shift: resolution-dependent flow-match mu - h, w = self._latent_side(height, width) - image_seq_len = (h // 2) * (w // 2) - sched_cfg = getattr(self.scheduler, "config_dict", None) or {} - base_seq_len = sched_cfg.get("base_image_seq_len", 256) - max_seq_len = sched_cfg.get("max_image_seq_len", 4096) - base_shift = sched_cfg.get("base_shift", 0.5) - max_shift = sched_cfg.get("max_shift", 0.9) - m = (max_shift - base_shift) / (max_seq_len - base_seq_len) - mu = image_seq_len * m + (base_shift - m * base_seq_len) - sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps) - if hasattr(self.scheduler, "set_timesteps"): - try: - self.scheduler.set_timesteps( - num_inference_steps, sigmas=sigmas, mu=mu - ) - except TypeError: - self.scheduler.set_timesteps(num_inference_steps) - timesteps = self.scheduler.timesteps - - latents = self.prepare_latents(batch, seed=seed, latents=latents) - emb = (uncond, embeddings) if do_cfg else embeddings - latents = self.denoise( + embeddings = model.encode_prompt(input_ids, attention_mask, **conditioning) + uncond = self.negative_embeddings( + model, + batch, + negative_input_ids, + negative_attention_mask, + guidance_scale, + conditioning, + ) + do_cfg = uncond is not None + timesteps = self.set_timesteps(model, num_inference_steps, height, width) + latents = model.prepare_latents(batch, seed=seed, latents=latents) + latents = model.denoise( latents, - emb, + (uncond, embeddings) if do_cfg else embeddings, num_inference_steps, guidance_scale if do_cfg else 1.0, timesteps=timesteps, ) if output_type == "latent": return latents - image = self.decode_latents(latents, height=height, width=width) - return self.postprocess_image(image) + image = model.decode_latents(latents, height=height, width=width) + return model.postprocess_image(image) diff --git a/zeromodels/models/qwen_image/qwen_image_vae.py b/zeromodels/diffusers/qwen_image/qwen_image_vae.py similarity index 99% rename from zeromodels/models/qwen_image/qwen_image_vae.py rename to zeromodels/diffusers/qwen_image/qwen_image_vae.py index c79dfbd9..e5e8eb80 100644 --- a/zeromodels/models/qwen_image/qwen_image_vae.py +++ b/zeromodels/diffusers/qwen_image/qwen_image_vae.py @@ -4,7 +4,7 @@ from keras import layers, ops from zeromodels.base.base_attention import fused_attention -from zeromodels.models.stable_diffusion.stable_diffusion_layers import safe_name +from zeromodels.diffusers.stable_diffusion.stable_diffusion_layers import safe_name @keras.saving.register_keras_serializable(package="zeromodels") diff --git a/zeromodels/models/qwen_image_21/__init__.py b/zeromodels/diffusers/qwen_image_21/__init__.py similarity index 73% rename from zeromodels/models/qwen_image_21/__init__.py rename to zeromodels/diffusers/qwen_image_21/__init__.py index cace1f2d..55bd21cb 100644 --- a/zeromodels/models/qwen_image_21/__init__.py +++ b/zeromodels/diffusers/qwen_image_21/__init__.py @@ -6,21 +6,17 @@ ) from .qwen_image_21_model import ( AutoencoderKLQwenImage21, - QwenImage21Model, + QwenImage21Family, QwenImage21TextEncoderModel, - QwenImage21TextToImage, QwenImage21Transformer2DModel, ) -from .qwen_image_21_tokenizer import QwenImage21Tokenizer __all__ = [ "AutoencoderKLQwenImage21", "QwenImage21Config", - "QwenImage21Model", + "QwenImage21Family", "QwenImage21TextConfig", "QwenImage21TextEncoderModel", - "QwenImage21TextToImage", - "QwenImage21Tokenizer", "QwenImage21Transformer2DModel", "QwenImage21TransformerConfig", "QwenImage21VAEConfig", diff --git a/zeromodels/models/qwen_image_21/convert_qwen_image_21_diffusers_to_keras.py b/zeromodels/diffusers/qwen_image_21/convert_qwen_image_21_diffusers_to_keras.py similarity index 97% rename from zeromodels/models/qwen_image_21/convert_qwen_image_21_diffusers_to_keras.py rename to zeromodels/diffusers/qwen_image_21/convert_qwen_image_21_diffusers_to_keras.py index 532fb032..2e994fbc 100644 --- a/zeromodels/models/qwen_image_21/convert_qwen_image_21_diffusers_to_keras.py +++ b/zeromodels/diffusers/qwen_image_21/convert_qwen_image_21_diffusers_to_keras.py @@ -61,11 +61,11 @@ def config_from_diffusers(repo, token=None): from huggingface_hub import hf_hub_download - from zeromodels.models.qwen_image_21.qwen_image_21_config import ( + from zeromodels.diffusers.qwen_image_21.qwen_image_21_config import ( QwenImage21Config, QwenImage21VAEConfig, ) - from zeromodels.models.qwen_image_21.qwen_image_21_model import ( + from zeromodels.diffusers.qwen_image_21.qwen_image_21_model import ( QwenImage21Transformer2DModel, ) @@ -143,7 +143,7 @@ def transfer_qwen_image_21( from safetensors import safe_open from zeromodels.base.base_mixin import build_dtype_scope - from zeromodels.models.qwen_image_21.qwen_image_21_model import QwenImage21Model + from zeromodels.diffusers.diffusion_model import ZMDiffusion config = config or config_from_diffusers(repo, token=token) flat = config.constructor_kwargs() @@ -153,9 +153,9 @@ def transfer_qwen_image_21( flat["transformer_text_seq_len"] = flat["max_sequence_length"] flat["max_seq_len"] = min(int(flat.get("max_seq_len", 1024)), 64) - print(f"[1/4] Building QwenImage21Model (dtype={dtype})…", flush=True) + print(f"[1/4] Building ZMDiffusion (qwen_image_21, dtype={dtype})…", flush=True) with build_dtype_scope(dtype), zeros_init(): - model = QwenImage21Model(**flat) + model = ZMDiffusion(model_type=config.model_type, **flat) # VAE RMSNorm checkpoints keep the name ``gamma`` (not ``weight``). vae_mapping = { diff --git a/zeromodels/models/qwen_image_21/qwen_image_21_config.py b/zeromodels/diffusers/qwen_image_21/qwen_image_21_config.py similarity index 96% rename from zeromodels/models/qwen_image_21/qwen_image_21_config.py rename to zeromodels/diffusers/qwen_image_21/qwen_image_21_config.py index 4d2e289b..094297ef 100644 --- a/zeromodels/models/qwen_image_21/qwen_image_21_config.py +++ b/zeromodels/diffusers/qwen_image_21/qwen_image_21_config.py @@ -210,9 +210,9 @@ class QwenImage21TextConfig(Qwen3VLTextConfig): class QwenImage21Config(BaseConfig): - """Configuration for :class:`QwenImage21Model` / :class:`QwenImage21TextToImage`. + """Configuration for a Qwen-Image-2.1 :class:`ZMDiffusion` (``model_type="qwen_image_21"``). - One hosted container: single-stream DiT + Qwen-Image-2.1 VAE + Qwen3-VL text + One hosted checkpoint: single-stream DiT + Qwen-Image-2.1 VAE + Qwen3-VL text encoder. Nested serialize; flat constructor with ``transformer_`` / ``vae_`` / ``text_`` prefixes. """ diff --git a/zeromodels/models/qwen_image_21/qwen_image_21_layers.py b/zeromodels/diffusers/qwen_image_21/qwen_image_21_layers.py similarity index 99% rename from zeromodels/models/qwen_image_21/qwen_image_21_layers.py rename to zeromodels/diffusers/qwen_image_21/qwen_image_21_layers.py index d1c1e290..654c8242 100644 --- a/zeromodels/models/qwen_image_21/qwen_image_21_layers.py +++ b/zeromodels/diffusers/qwen_image_21/qwen_image_21_layers.py @@ -7,11 +7,11 @@ from keras import layers, ops from zeromodels.base.base_attention import fused_attention -from zeromodels.models.qwen_image.qwen_image_layers import ( +from zeromodels.diffusers.qwen_image.qwen_image_layers import ( QwenImageRMSNorm, apply_rotary_emb_qwen, ) -from zeromodels.models.stable_diffusion.stable_diffusion_layers import ( +from zeromodels.diffusers.stable_diffusion.stable_diffusion_layers import ( safe_name, timestep_embedding, ) diff --git a/zeromodels/models/qwen_image_21/qwen_image_21_model.py b/zeromodels/diffusers/qwen_image_21/qwen_image_21_model.py similarity index 69% rename from zeromodels/models/qwen_image_21/qwen_image_21_model.py rename to zeromodels/diffusers/qwen_image_21/qwen_image_21_model.py index d919310e..66009bc9 100644 --- a/zeromodels/models/qwen_image_21/qwen_image_21_model.py +++ b/zeromodels/diffusers/qwen_image_21/qwen_image_21_model.py @@ -4,18 +4,9 @@ import numpy as np from keras import layers, ops -from zeromodels.base import BaseDiffusion, BaseModel, CausalMask -from zeromodels.base.base_attention import with_model_attn_implementation -from zeromodels.base.base_mixin import inference_scope -from zeromodels.base.base_scheduler import ( - FlowMatchEulerDiscreteScheduler, - get_scheduler, -) -from zeromodels.models.qwen3_vl.qwen3_vl_model import ( - Qwen3VLTextModel, - qwen3_text_cos_sin, -) -from zeromodels.models.qwen_image_21.qwen_image_21_config import ( +from zeromodels.base import BaseModel, CausalMask +from zeromodels.diffusers.qwen_image.qwen_image_model import QwenImageFamily +from zeromodels.diffusers.qwen_image_21.qwen_image_21_config import ( DEFAULT_LATENTS_MEAN, DEFAULT_LATENTS_STD, QwenImage21Config, @@ -23,7 +14,7 @@ QwenImage21TransformerConfig, QwenImage21VAEConfig, ) -from zeromodels.models.qwen_image_21.qwen_image_21_layers import ( +from zeromodels.diffusers.qwen_image_21.qwen_image_21_layers import ( IMG_TOKENS_PER_SLOT, MASK_NEG, QwenImage21AdaLayerNormContinuous, @@ -34,14 +25,16 @@ build_qwenimage21_rope_angles, build_token_metadata, ) -from zeromodels.models.qwen_image_21.qwen_image_21_vae import ( +from zeromodels.diffusers.qwen_image_21.qwen_image_21_vae import ( QwenImage21CausalConv, QwenImage21Decoder3d, QwenImage21Encoder3d, ) -from zeromodels.models.stable_diffusion.stable_diffusion_layers import safe_name - -QWEN_IMAGE_21_HUB_SIBLINGS = frozenset({"QwenImage21Model", "QwenImage21TextToImage"}) +from zeromodels.diffusers.stable_diffusion.stable_diffusion_layers import safe_name +from zeromodels.models.qwen3_vl.qwen3_vl_model import ( + Qwen3VLTextModel, + qwen3_text_cos_sin, +) def pack_latents(latents, height, width): @@ -534,24 +527,27 @@ def __init__(self, name="text_encoder", **kwargs): self.tie_embeddings = config.tie_embeddings -@keras.saving.register_keras_serializable(package="zeromodels") -class QwenImage21Model(BaseModel): - """Qwen-Image-2.1 weights container: transformer + VAE + Qwen3-VL text tower.""" +class QwenImage21Family(QwenImageFamily): + """Qwen-Image-2.1 inside :class:`ZMDiffusion` (``model_type="qwen_image_21"``). + + Three components, ``transformer`` (the single-stream DiT, + :class:`QwenImage21Transformer2DModel`), ``vae`` (the residual RGBA + :class:`AutoencoderKLQwenImage21`) and ``text_encoder`` (the Qwen3-VL text tower, + pre-final-norm states), wired like Qwen-Image's (diffusers' + ``QwenImage21Pipeline`` components). Hosted as ``zeromodels/qwen-image-2.1``. + Generation is Qwen-Image's with the 2.1 differences: unpatched latents (one + token per latent pixel), a plain guidance combination, the alpha channel dropped + from the decoded image, and the reference defaults of 40 steps without guidance. + Classifier-free guidance runs only with ``guidance_scale > 1`` and an explicit + negative prompt. + """ + + model_type = "qwen_image_21" config_class = QwenImage21Config - HF_MODEL_TYPE = None - HUB_REPO_SIBLINGS = QWEN_IMAGE_21_HUB_SIBLINGS + generate_args = {"num_inference_steps": 40, "guidance_scale": 1.0} - def __init__(self, name="QwenImage21Model", **kwargs): - keras_kwargs = {k: kwargs.pop(k) for k in ("trainable", "dtype") if k in kwargs} - config = self.config_class.from_dict(kwargs) - components = self.build_components(config) - inputs, outputs = self.build_graph(config, components) - super().__init__(inputs=inputs, outputs=outputs, name=name, **keras_kwargs) - for attr, component in components.items(): - setattr(self, attr, component) - - def build_components(self, config): + def build_components(self, config, data_format, channels_axis): d, v = config.transformer_config, config.vae_config transformer = QwenImage21Transformer2DModel( patch_size=d.patch_size, @@ -593,7 +589,7 @@ def build_components(self, config): text_encoder = QwenImage21TextEncoderModel(**text_kw) return {"transformer": transformer, "vae": vae, "text_encoder": text_encoder} - def build_graph(self, config, components): + def build_graph(self, config, components, data_format): d, v, t = config.transformer_config, config.vae_config, config.text_config transformer, vae, text_encoder = ( components["transformer"], @@ -648,239 +644,44 @@ def build_graph(self, config, components): "prompt_embeds": text_out["last_hidden_state"], } - def get_config(self): - config = super().get_config() - config.update(self.config.constructor_kwargs()) - return config + def image_seq_len(self, height, width): + return height * width - def from_hf(self, *args, **kwargs): - raise NotImplementedError( - "On-the-fly hf: conversion is not supported for Qwen-Image-2.1; " - "use convert_qwen_image_21_diffusers_to_keras.py and " - "from_weights('zeromodels/qwen-image-2.1')." - ) + def latent_shape(self, model): + h, w = self.latent_side(model) + return (self.image_seq_len(h, w), model.vae.z_dim) + def pack_latents(self, latents, height, width): + return pack_latents(latents, height, width) -@keras.saving.register_keras_serializable(package="zeromodels") -class QwenImage21TextToImage(QwenImage21Model, BaseDiffusion): - """Text-to-image Qwen-Image-2.1 (Diffusers ``QwenImage21Pipeline``). - - :: + def unpack_latents(self, latents, height, width, channels): + return unpack_latents(latents, height, width, channels) - model = QwenImage21TextToImage.from_weights("zeromodels/qwen-image-2.1") - tok = QwenImage21Tokenizer.from_weights("zeromodels/qwen-image-2.1") - image = model.generate(**tok("a cat"), height=1024, width=1024) - - Defaults match the reference: 40 flow-match steps, ``guidance_scale=1.0`` - (no CFG). Pass ``guidance_scale>1`` with a negative prompt for true CFG. - """ - - config_class = QwenImage21Config - HUB_REPO_SIBLINGS = QWEN_IMAGE_21_HUB_SIBLINGS - generate_args = {"num_inference_steps": 40, "guidance_scale": 1.0} - DEFAULT_GUIDANCE_SCALE = 1.0 - - def __init__(self, scheduler=None, name="QwenImage21TextToImage", **kwargs): - scheduler_config = kwargs.get("scheduler_config") - super().__init__(name=name, **kwargs) - if scheduler is None: - scheduler = ( - get_scheduler(scheduler_config) - if scheduler_config - else self.default_scheduler() - ) - self.scheduler = scheduler - - def default_scheduler(self): - return FlowMatchEulerDiscreteScheduler( - shift=1.0, - use_dynamic_shifting=True, - base_shift=0.5, - max_shift=0.9, - base_image_seq_len=256, - max_image_seq_len=8192, - shift_terminal=0.02, - time_shift_type="exponential", - ) - - @property - def vae_scale_factor(self): - return self.vae.vae_scale_factor - - @property - def latent_shape(self): - h, w = self.latent_side() - return (h * w, self.vae.z_dim) - - def latent_side(self, height=None, width=None): - scale = self.vae_scale_factor * 2 - if height is None or width is None: - side = self.config.default_sample_size * self.vae_scale_factor - height = width = side - h = 2 * (int(height) // scale) - w = 2 * (int(width) // scale) - return h, w - - def unconditional_ids(self, batch): - length = self.config.text_config.max_seq_len - row = [self.config.pad_token_id] * length + def unconditional_ids(self, model, batch): + length = model.config.text_config.max_seq_len + row = [model.config.pad_token_id] * length return ops.convert_to_tensor([row] * batch, dtype="int32") - def encode_prompt(self, input_ids, attention_mask=None, **conditioning): - del conditioning - input_ids = ops.cast(ops.convert_to_tensor(input_ids), "int32") - if attention_mask is None: - attention_mask = ops.ones_like(input_ids) - else: - attention_mask = ops.cast(ops.convert_to_tensor(attention_mask), "int32") - - out = self.text_encoder( - {"input_ids": input_ids, "attention_mask": attention_mask} - ) - hidden = out["last_hidden_state"] - drop = int(self.config.prompt_template_encode_start_idx) - max_len = int(self.config.max_sequence_length) - - hidden_np = ops.convert_to_numpy(hidden) - mask_np = ops.convert_to_numpy(attention_mask) - batch = hidden_np.shape[0] - dim = hidden_np.shape[-1] - embeds = np.zeros((batch, max_len, dim), dtype=hidden_np.dtype) - out_mask = np.zeros((batch, max_len), dtype=np.int32) - for i in range(batch): - valid = hidden_np[i][mask_np[i].astype(bool)] - valid = valid[drop:][:max_len] - n = valid.shape[0] - embeds[i, :n] = valid - out_mask[i, :n] = 1 - return { - "encoder_hidden_states": ops.convert_to_tensor(embeds), - "encoder_hidden_states_mask": ops.convert_to_tensor(out_mask), - } - - @with_model_attn_implementation - def predict_noise(self, latents, timesteps, embeddings): - return self.transformer( - { - "sample": latents, - "timestep": timesteps, - "encoder_hidden_states": embeddings["encoder_hidden_states"], - "encoder_hidden_states_mask": embeddings["encoder_hidden_states_mask"], - } - )["sample"] - - def decode_latents(self, latents, height=None, width=None): - h, w = self.latent_side(height, width) - latents = unpack_latents(latents, h, w, self.vae.z_dim) - mean = ops.convert_to_tensor(self.vae.latents_mean, dtype="float32") - std = ops.convert_to_tensor(self.vae.latents_std, dtype="float32") - mean = ops.reshape(mean, (1, 1, 1, -1)) - std = ops.reshape(std, (1, 1, 1, -1)) - latents = latents * std + mean - image = self.vae.decode(latents) + def decode_latents(self, model, latents, height=None, width=None): + image = super().decode_latents(model, latents, height=height, width=width) if int(image.shape[-1]) == 4: - image = image[..., :3] + image = image[..., :3] # RGBA VAE: drop the alpha channel return image - def prepare_latents(self, batch, seed=None, latents=None, dtype="float32"): - h, w = self.latent_side( - getattr(self, "_gen_height", None), getattr(self, "_gen_width", None) - ) - channels = self.vae.z_dim - if latents is None: - noise = keras.random.normal((batch, h, w, channels), seed=seed, dtype=dtype) - latents = pack_latents(noise, h, w) - else: - latents = ops.cast(ops.convert_to_tensor(latents), dtype) - return latents * self.scheduler.init_noise_sigma + def combine_guidance(self, noise_pred, neg_pred, guidance_scale): + return neg_pred + guidance_scale * (noise_pred - neg_pred) - def denoise( - self, latents, embeddings, num_inference_steps, guidance_scale, timesteps=None - ): - do_cfg = guidance_scale > 1.0 - scheduler = self.scheduler - if timesteps is None: - scheduler.set_timesteps(num_inference_steps) - timesteps = scheduler.timesteps - if do_cfg and isinstance(embeddings, (tuple, list)): - uncond_emb, cond_emb = embeddings - else: - cond_emb = embeddings - uncond_emb = None - - for t in timesteps: - t_batch = ops.full((ops.shape(latents)[0],), float(t) / 1000.0) - noise_pred = self.predict_noise(latents, t_batch, cond_emb) - if do_cfg and uncond_emb is not None: - neg_pred = self.predict_noise(latents, t_batch, uncond_emb) - noise_pred = neg_pred + guidance_scale * (noise_pred - neg_pred) - latents = scheduler.step(noise_pred, t, latents) - return latents - - def generate( + def negative_embeddings( self, - input_ids, - attention_mask=None, - negative_input_ids=None, - negative_attention_mask=None, - num_inference_steps=None, - guidance_scale=None, - seed=None, - latents=None, - height=None, - width=None, - output_type="image", - **conditioning, + model, + batch, + negative_input_ids, + negative_attention_mask, + guidance_scale, + conditioning, ): - num_inference_steps, guidance_scale, seed = self.resolve_generation_args( - num_inference_steps, guidance_scale, seed - ) - self._gen_height = height - self._gen_width = width - input_ids = ops.cast(ops.convert_to_tensor(input_ids), "int32") - batch = int(input_ids.shape[0]) - - with inference_scope(): - embeddings = self.encode_prompt(input_ids, attention_mask, **conditioning) - if guidance_scale > 1.0 and negative_input_ids is None: - do_cfg = False - uncond = None - elif negative_input_ids is not None and guidance_scale > 1.0: - uncond = self.encode_prompt(negative_input_ids, negative_attention_mask) - do_cfg = True - else: - uncond = None - do_cfg = False - - h, w = self.latent_side(height, width) - image_seq_len = h * w - sched_cfg = getattr(self.scheduler, "config_dict", None) or {} - base_seq_len = sched_cfg.get("base_image_seq_len", 256) - max_seq_len = sched_cfg.get("max_image_seq_len", 4096) - base_shift = sched_cfg.get("base_shift", 0.5) - max_shift = sched_cfg.get("max_shift", 0.9) - m = (max_shift - base_shift) / (max_seq_len - base_seq_len) - mu = image_seq_len * m + (base_shift - m * base_seq_len) - sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps) - if hasattr(self.scheduler, "set_timesteps"): - try: - self.scheduler.set_timesteps( - num_inference_steps, sigmas=sigmas, mu=mu - ) - except TypeError: - self.scheduler.set_timesteps(num_inference_steps) - timesteps = self.scheduler.timesteps - - latents = self.prepare_latents(batch, seed=seed, latents=latents) - emb = (uncond, embeddings) if do_cfg else embeddings - latents = self.denoise( - latents, - emb, - num_inference_steps, - guidance_scale if do_cfg else 1.0, - timesteps=timesteps, - ) - if output_type == "latent": - return latents - image = self.decode_latents(latents, height=height, width=width) - return self.postprocess_image(image) + # guidance only with an explicit negative prompt (the reference's default + # is no guidance at all) + if guidance_scale <= 1.0 or negative_input_ids is None: + return None + return model.encode_prompt(negative_input_ids, negative_attention_mask) diff --git a/zeromodels/models/qwen_image_21/qwen_image_21_vae.py b/zeromodels/diffusers/qwen_image_21/qwen_image_21_vae.py similarity index 99% rename from zeromodels/models/qwen_image_21/qwen_image_21_vae.py rename to zeromodels/diffusers/qwen_image_21/qwen_image_21_vae.py index 7bbecaad..2827ea35 100644 --- a/zeromodels/models/qwen_image_21/qwen_image_21_vae.py +++ b/zeromodels/diffusers/qwen_image_21/qwen_image_21_vae.py @@ -4,8 +4,8 @@ from keras import layers, ops from zeromodels.base.base_attention import fused_attention -from zeromodels.models.qwen_image.qwen_image_vae import QwenImageRMSNorm -from zeromodels.models.stable_diffusion.stable_diffusion_layers import safe_name +from zeromodels.diffusers.qwen_image.qwen_image_vae import QwenImageRMSNorm +from zeromodels.diffusers.stable_diffusion.stable_diffusion_layers import safe_name @keras.saving.register_keras_serializable(package="zeromodels") diff --git a/zeromodels/models/stable_diffusion/__init__.py b/zeromodels/diffusers/stable_diffusion/__init__.py similarity index 65% rename from zeromodels/models/stable_diffusion/__init__.py rename to zeromodels/diffusers/stable_diffusion/__init__.py index bbf02ef5..9431c1c2 100644 --- a/zeromodels/models/stable_diffusion/__init__.py +++ b/zeromodels/diffusers/stable_diffusion/__init__.py @@ -6,20 +6,16 @@ ) from .stable_diffusion_model import ( AutoencoderKL, - StableDiffusionModel, - StableDiffusionTextToImage, + StableDiffusionFamily, UNet2DConditionModel, ) -from .stable_diffusion_tokenizer import StableDiffusionTokenizer __all__ = [ - "StableDiffusionModel", - "StableDiffusionTextToImage", - "UNet2DConditionModel", "AutoencoderKL", + "AutoencoderKLConfig", "StableDiffusionConfig", + "StableDiffusionFamily", "StableDiffusionTextConfig", "UNet2DConditionConfig", - "AutoencoderKLConfig", - "StableDiffusionTokenizer", + "UNet2DConditionModel", ] diff --git a/zeromodels/models/stable_diffusion/convert_stable_diffusion_diffusers_to_keras.py b/zeromodels/diffusers/stable_diffusion/convert_stable_diffusion_diffusers_to_keras.py similarity index 95% rename from zeromodels/models/stable_diffusion/convert_stable_diffusion_diffusers_to_keras.py rename to zeromodels/diffusers/stable_diffusion/convert_stable_diffusion_diffusers_to_keras.py index 84a68e64..673c93d1 100644 --- a/zeromodels/models/stable_diffusion/convert_stable_diffusion_diffusers_to_keras.py +++ b/zeromodels/diffusers/stable_diffusion/convert_stable_diffusion_diffusers_to_keras.py @@ -35,10 +35,10 @@ def config_from_diffusers(repo, token=None, config_cls=None): from diffusers import AutoencoderKL, PNDMScheduler, UNet2DConditionModel from transformers import CLIPTextConfig, CLIPTokenizerFast - from zeromodels.models.stable_diffusion.stable_diffusion_config import ( + from zeromodels.diffusers.stable_diffusion.stable_diffusion_config import ( StableDiffusionConfig, ) - from zeromodels.models.stable_diffusion.stable_diffusion_model import ( + from zeromodels.diffusers.stable_diffusion.stable_diffusion_model import ( UNet2DConditionModel as KerasUNet2DConditionModel, ) @@ -96,7 +96,7 @@ def config_from_diffusers(repo, token=None, config_cls=None): ) -def transfer_stable_diffusion(repo, token=None, model_cls=None, config_cls=None): +def transfer_stable_diffusion(repo, token=None, config_cls=None): import gc import torch @@ -104,14 +104,11 @@ def transfer_stable_diffusion(repo, token=None, model_cls=None, config_cls=None) from diffusers import UNet2DConditionModel as DiffusersUNet2DConditionModel from transformers import CLIPTextModel as HFCLIPTextModel + from zeromodels.diffusers.diffusion_model import ZMDiffusion from zeromodels.models.clip import CLIPTextModel - from zeromodels.models.stable_diffusion.stable_diffusion_model import ( - StableDiffusionModel, - ) - model_cls = model_cls or StableDiffusionModel config = config_from_diffusers(repo, token=token, config_cls=config_cls) - model = model_cls(config) + model = ZMDiffusion(config) load = {"torch_dtype": torch.float32, "token": token} # the SD 1.x era VAE files spell the mid-block attention query / key / value / diff --git a/zeromodels/models/stable_diffusion/stable_diffusion_config.py b/zeromodels/diffusers/stable_diffusion/stable_diffusion_config.py similarity index 94% rename from zeromodels/models/stable_diffusion/stable_diffusion_config.py rename to zeromodels/diffusers/stable_diffusion/stable_diffusion_config.py index ad5c6128..68b00cc9 100644 --- a/zeromodels/models/stable_diffusion/stable_diffusion_config.py +++ b/zeromodels/diffusers/stable_diffusion/stable_diffusion_config.py @@ -54,7 +54,7 @@ class UNet2DConditionConfig(BaseConfig): Examples: ```python - >>> from zeromodels.models.stable_diffusion import UNet2DConditionConfig, UNet2DConditionModel + >>> from zeromodels.diffusers.stable_diffusion import UNet2DConditionConfig, UNet2DConditionModel >>> config = UNet2DConditionConfig() >>> model = UNet2DConditionModel(config) @@ -125,7 +125,7 @@ class AutoencoderKLConfig(BaseConfig): Examples: ```python - >>> from zeromodels.models.stable_diffusion import AutoencoderKLConfig, AutoencoderKL + >>> from zeromodels.diffusers.stable_diffusion import AutoencoderKLConfig, AutoencoderKL >>> config = AutoencoderKLConfig() >>> vae = AutoencoderKL.from_config(config) @@ -161,7 +161,7 @@ class StableDiffusionTextConfig(CLIPTextConfig): class StableDiffusionConfig(BaseConfig): - r"""Configuration for [`StableDiffusionModel`], the hosted SD 1.x weights container. + r"""Configuration for a Stable Diffusion 1.x [`ZMDiffusion`] (``model_type="stable_diffusion"``). One config for the three trained components, so a repo carries a single ``zm_config.json`` (nested ``unet_config`` / ``vae_config`` / ``text_config``) @@ -189,16 +189,17 @@ class StableDiffusionConfig(BaseConfig): diffusers ``scheduler_config.json`` form (``_class_name``, ``beta_schedule``, ``beta_start`` / ``beta_end``, ``steps_offset``, ...). Stored inside ``zm_config.json`` rather than as a separate file; - ``StableDiffusionTextToImage`` builds its scheduler from it (PNDM with + ``ZMDiffusion`` builds its scheduler from it (PNDM with the SD 1.x schedule when absent). Examples: ```python - >>> from zeromodels.models.stable_diffusion import StableDiffusionConfig, StableDiffusionModel + >>> from zeromodels.diffusers import ZMDiffusion + >>> from zeromodels.diffusers.stable_diffusion import StableDiffusionConfig >>> config = StableDiffusionConfig() - >>> model = StableDiffusionModel(config) # random weights, SD 1.x shapes + >>> model = ZMDiffusion(config) # random weights, SD 1.x shapes >>> model.config.unet_config.block_out_channels (320, 640, 1280, 1280) ```""" diff --git a/zeromodels/models/stable_diffusion/stable_diffusion_layers.py b/zeromodels/diffusers/stable_diffusion/stable_diffusion_layers.py similarity index 100% rename from zeromodels/models/stable_diffusion/stable_diffusion_layers.py rename to zeromodels/diffusers/stable_diffusion/stable_diffusion_layers.py diff --git a/zeromodels/models/stable_diffusion/stable_diffusion_model.py b/zeromodels/diffusers/stable_diffusion/stable_diffusion_model.py similarity index 80% rename from zeromodels/models/stable_diffusion/stable_diffusion_model.py rename to zeromodels/diffusers/stable_diffusion/stable_diffusion_model.py index 7956b43e..2f3a935f 100644 --- a/zeromodels/models/stable_diffusion/stable_diffusion_model.py +++ b/zeromodels/diffusers/stable_diffusion/stable_diffusion_model.py @@ -2,9 +2,8 @@ from keras import layers, ops from zeromodels.base import BaseModel -from zeromodels.base.base_diffusion import BaseDiffusion from zeromodels.base.base_mixin import build_dtype_scope -from zeromodels.base.base_scheduler import PNDMScheduler, get_scheduler +from zeromodels.base.base_scheduler import PNDMScheduler from zeromodels.models.clip import CLIPTextModel from .stable_diffusion_config import ( @@ -25,12 +24,6 @@ timestep_embedding, ) -# The container and the task build the identical graph, so both load the one -# hosted repo (whose zm_config.json names the container), like CLIP's heads. -STABLE_DIFFUSION_HUB_SIBLINGS = frozenset( - {"StableDiffusionModel", "StableDiffusionTextToImage"} -) - CROSS_ATTN_DOWN = "CrossAttnDownBlock2D" DOWN = "DownBlock2D" CROSS_ATTN_UP = "CrossAttnUpBlock2D" @@ -570,7 +563,7 @@ class AutoencoderKL(BaseModel): (``image -> encoder -> quant_conv -> moments`` and ``latent -> post_quant_conv -> decoder -> sample``). :meth:`encode` and :meth:`decode` run one path on its own. Being a single functional model is what - lets its weights ride inside the hosted :class:`StableDiffusionModel`'s one + lets its weights ride inside a hosted :class:`ZMDiffusion` checkpoint's one ``model.weights.h5``. Args mirror the diffusers config: ``in_channels`` / ``out_channels`` (3), @@ -748,60 +741,26 @@ def from_diffusers_config(cls, config, sample_size=512, **kwargs): ) -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusionModel(BaseModel): - """Stable Diffusion 1.x weights container: the UNet, the VAE and the CLIP text - encoder as one functional ``keras.Model``. - - This is the hosted unit: a ``zeromodels/`` repo carries one - ``zm_config.json`` (a :class:`StableDiffusionConfig`) and one - ``model.weights.h5``, and the standard ``from_weights("zeromodels/")`` - loader builds this container and loads every component's weights in one go. It - is the ``XModel`` of the family; :class:`StableDiffusionTextToImage` is the task - class, this same graph plus the denoising loop (``BaseDiffusion``) and a - scheduler, the way ``T5ConditionalGenerate`` extends ``T5Model``. - - Like ``CLIPModel`` (text + vision towers in one graph), it is a single - functional model whose graph is three disconnected paths, one per component: - ``{"sample", "timestep", "encoder_hidden_states"} -> "noise_pred"`` (UNet), - ``{"image", "latent"} -> "moments" / "image"`` (VAE) and - ``{"token_ids", "padding_mask"} -> "text_embeds"`` (CLIP). The components are - exposed as ``.unet`` / ``.vae`` / ``.text_encoder`` (they share this model's - weights) and are what the task class runs step by step; calling the container - end to end is rarely useful. On-the-fly ``hf:`` conversion is deliberately - unsupported for diffusion models; the weights are converted once by - ``convert_stable_diffusion_diffusers_to_keras.py`` and hosted. +class StableDiffusionFamily: + """Stable Diffusion 1.x inside :class:`ZMDiffusion` (``model_type="stable_diffusion"``). - The constructor is flat, like every zeromodels model: the sub-config fields - arrive prefixed ``unet_`` / ``vae_`` / ``text_`` (see - :class:`StableDiffusionConfig`), or pass a config positionally; ``.config`` is - always the typed config. Pass ``unet_sample_size`` / ``vae_sample_size`` to - build for another resolution (the weights are resolution-independent). + Three components, ``unet`` (:class:`UNet2DConditionModel`), ``vae`` + (:class:`AutoencoderKL`) and ``text_encoder`` (CLIP ViT-L/14's text tower), wired + into three disconnected paths: ``{"sample", "timestep", + "encoder_hidden_states"} -> "noise_pred"``, ``{"image", "latent"} -> + "moments" / "image"`` and ``{"token_ids", "padding_mask"} -> "text_embeds"``. + Generation follows diffusers' ``StableDiffusionPipeline``: the prompt through the + CLIP tower (padding attended, as the reference does), the empty prompt as the + unconditional branch, the UNet as the denoiser and the scaled VAE decoder. The + sampler defaults to PNDM with the SD 1.x schedule when the config carries no + ``scheduler_config``. """ + model_type = "stable_diffusion" config_class = StableDiffusionConfig - HF_MODEL_TYPE = None - HUB_REPO_SIBLINGS = STABLE_DIFFUSION_HUB_SIBLINGS - - def __init__(self, name="StableDiffusionModel", **kwargs): - keras_kwargs = {k: kwargs.pop(k) for k in ("trainable", "dtype") if k in kwargs} - config = self.config_class.from_dict(kwargs) # regroup the flat kwargs - - data_format = keras.config.image_data_format() - channels_axis = -1 if data_format == "channels_last" else 1 - - components = self.build_components(config, data_format, channels_axis) - inputs, outputs = self.build_graph(config, components, data_format) - super().__init__(inputs=inputs, outputs=outputs, name=name, **keras_kwargs) - - self.data_format = data_format - self.channels_axis = channels_axis - for attr, component in components.items(): - setattr(self, attr, component) + generate_args = {"num_inference_steps": 50, "guidance_scale": 7.5} def build_components(self, config, data_format, channels_axis): - """The trained components, ``{attribute name: model}``: the UNet, the VAE - and the CLIP text tower (subclasses swap or add towers).""" u, v, t = config.unet_config, config.vae_config, config.text_config unet = UNet2DConditionModel( u, data_format=data_format, channels_axis=channels_axis @@ -824,7 +783,6 @@ def build_components(self, config, data_format, channels_axis): return {"unet": unet, "vae": vae, "text_encoder": text_encoder} def unet_vae_inputs(self, config, components, data_format): - """The UNet's and the VAE's ``keras.Input`` dict (shared by every SD family).""" u, v = config.unet_config, config.vae_config vae = components["vae"] sample_h, sample_w = ( @@ -838,7 +796,7 @@ def unet_vae_inputs(self, config, components, data_format): else (v.sample_size, v.sample_size) ) lat_h, lat_w = img_h // vae.vae_scale_factor, img_w // vae.vae_scale_factor - inputs = { + return { "sample": layers.Input( shape=(u.in_channels, sample_h, sample_w) if data_format == "channels_first" @@ -863,11 +821,8 @@ def unet_vae_inputs(self, config, components, data_format): name="latent", ), } - return inputs def build_graph(self, config, components, data_format): - """Wire the components into the container's ``(inputs, outputs)`` dicts: - three disconnected paths, one per component.""" t = config.text_config unet, vae, text_encoder = ( components["unet"], @@ -901,108 +856,32 @@ def build_graph(self, config, components, data_format): } return inputs, outputs - @classmethod - def from_hf(cls, repo, **kwargs): - raise NotImplementedError( - "On-the-fly `hf:` conversion is not supported for diffusion models. Load " - "the preconverted weights from a zeromodels repo, e.g. " - f"{cls.__name__}.from_weights('zeromodels/stable-diffusion-v1-5'), or run " - "convert_stable_diffusion_diffusers_to_keras.py to produce them from " - f"{repo!r}." - ) - - @classmethod - def from_variant(cls, variant, /, quantization=None, **kwargs): - # A bare name is the hosted repo's basename. `quantization` is applied by - # from_weights after the load, exactly as on the org/repo branch. - return cls.from_hub_repo(f"zeromodels/{variant}", **kwargs) - - def get_config(self): - config = super().get_config() - config.update(self.config.constructor_kwargs()) - return config - - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusionTextToImage(StableDiffusionModel, BaseDiffusion): - """Text-to-image Stable Diffusion 1.x, pure Keras 3 and cross-backend (the - diffusers ``StableDiffusionPipeline`` in this library's task-class form). - - :class:`StableDiffusionModel` (the UNet, VAE and CLIP text towers, built - internally like ``CLIPModel``) plus :class:`BaseDiffusion`'s ``generate``, the - way ``T5ConditionalGenerate`` is ``T5Model`` plus a generation mixin: this class - only supplies the hooks (``encode_prompt`` on the CLIP tower, the empty prompt's - ``unconditional_ids``, ``predict_noise`` on the UNet, ``decode_latents`` on the - VAE) and the scheduler. It is the same graph and weights as the container, so - both classes load the same hosted repo. - - Inference is the library's usual two-step, tokenizer then ``generate``:: - - sd = StableDiffusionTextToImage.from_weights("zeromodels/stable-diffusion-v1-5") - tokenizer = StableDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-v1-5") - image = sd.generate(**tokenizer("a photo of a cat")) # (1, 512, 512, 3) uint8 - - The graphs are built for the repo's resolution (512px); pass - ``unet_sample_size`` / ``vae_sample_size`` overrides to ``from_weights`` to - build for another size (the weights are resolution-independent). On-the-fly - ``hf:`` conversion is not supported for diffusion models. - - Args: - scheduler: A :class:`BaseScheduler`. Defaults to the one the config's - ``scheduler_config`` describes (the checkpoint's own sampler and noise - schedule, carried in the repo's ``zm_config.json``), or PNDM with the SD - 1.x schedule when the config has none. Swap any time with - ``sd.scheduler = EulerDiscreteScheduler()``. - **kwargs: The flat :class:`StableDiffusionConfig` fields (see - :class:`StableDiffusionModel`). - """ - - HUB_REPO_SIBLINGS = STABLE_DIFFUSION_HUB_SIBLINGS - # Default generation settings (diffusers' StableDiffusionPipeline defaults); - # a repo's zm_config.json generate_args override them on load. - generate_args = {"num_inference_steps": 50, "guidance_scale": 7.5} - - def __init__(self, scheduler=None, name="StableDiffusionTextToImage", **kwargs): - # read from the flat kwargs: the metaclass attaches self.config only after - # __init__ returns - scheduler_config = kwargs.get("scheduler_config") - super().__init__(name=name, **kwargs) - if scheduler is None: - scheduler = ( - get_scheduler(scheduler_config) - if scheduler_config - else self.default_scheduler() - ) - self.scheduler = scheduler - def default_scheduler(self): - """The family's sampler when the config carries no ``scheduler_config``.""" return PNDMScheduler(steps_offset=1) - @property - def latent_shape(self): - size = self.unet.sample_size + def latent_shape(self, model): + size = model.unet.sample_size height, width = size if isinstance(size, (tuple, list)) else (size, size) - if self.data_format == "channels_first": - return (self.unet.in_channels, height, width) - return (height, width, self.unet.in_channels) + if model.data_format == "channels_first": + return (model.unet.in_channels, height, width) + return (height, width, model.unet.in_channels) - def unconditional_ids(self, batch): - cfg = self.config + def unconditional_ids(self, model, batch): + cfg = model.config length = cfg.text_config.max_seq_len row = [cfg.bos_token_id, cfg.eos_token_id] + [cfg.pad_token_id] * (length - 2) return ops.convert_to_tensor([row] * batch, dtype="int32") - def encode_prompt(self, input_ids, attention_mask=None): + def encode_prompt(self, model, input_ids, attention_mask=None): # SD attends the padding tokens (the reference passes no attention mask), # so the text tower gets an all-ones mask regardless of the tokenizer's input_ids = ops.cast(ops.convert_to_tensor(input_ids), "int32") mask = ops.ones_like(input_ids) - out = self.text_encoder({"token_ids": input_ids, "padding_mask": mask}) + out = model.text_encoder({"token_ids": input_ids, "padding_mask": mask}) return out["last_hidden_state"] - def predict_noise(self, latents, timesteps, embeddings): - return self.unet( + def predict_noise(self, model, latents, timesteps, embeddings): + return model.unet( { "sample": latents, "timestep": timesteps, @@ -1010,20 +889,20 @@ def predict_noise(self, latents, timesteps, embeddings): } )["sample"] - def decode_latents(self, latents): - image = self.vae.decode( - latents / self.vae.scaling_factor + self.vae.shift_factor + def decode_latents(self, model, latents): + image = model.vae.decode( + latents / model.vae.scaling_factor + model.vae.shift_factor ) - if self.data_format == "channels_first": + if model.data_format == "channels_first": image = ops.transpose(image, (0, 2, 3, 1)) # generate() hands out HWC return image - def encode_latents(self, image): + def encode_latents(self, model, image): # the posterior mean, scaled: the image-to-image starting point (the # reference samples the posterior; its noise is negligible next to the # noise added for the strength) - if self.data_format == "channels_first": + if model.data_format == "channels_first": image = ops.transpose(image, (0, 3, 1, 2)) # generate() hands in HWC return ( - self.vae.encode(image) - self.vae.shift_factor - ) * self.vae.scaling_factor + model.vae.encode(image) - model.vae.shift_factor + ) * model.vae.scaling_factor diff --git a/zeromodels/diffusers/stable_diffusion_2/__init__.py b/zeromodels/diffusers/stable_diffusion_2/__init__.py new file mode 100644 index 00000000..5938060f --- /dev/null +++ b/zeromodels/diffusers/stable_diffusion_2/__init__.py @@ -0,0 +1,13 @@ +from .stable_diffusion_2_config import ( + StableDiffusion2Config, + StableDiffusion2TextConfig, + StableDiffusion2UNetConfig, +) +from .stable_diffusion_2_model import StableDiffusion2Family + +__all__ = [ + "StableDiffusion2Config", + "StableDiffusion2Family", + "StableDiffusion2TextConfig", + "StableDiffusion2UNetConfig", +] diff --git a/zeromodels/models/stable_diffusion_2/convert_stable_diffusion_2_diffusers_to_keras.py b/zeromodels/diffusers/stable_diffusion_2/convert_stable_diffusion_2_diffusers_to_keras.py similarity index 82% rename from zeromodels/models/stable_diffusion_2/convert_stable_diffusion_2_diffusers_to_keras.py rename to zeromodels/diffusers/stable_diffusion_2/convert_stable_diffusion_2_diffusers_to_keras.py index 97ff389b..8351499f 100644 --- a/zeromodels/models/stable_diffusion_2/convert_stable_diffusion_2_diffusers_to_keras.py +++ b/zeromodels/diffusers/stable_diffusion_2/convert_stable_diffusion_2_diffusers_to_keras.py @@ -1,9 +1,9 @@ import numpy as np -from zeromodels.models.stable_diffusion.convert_stable_diffusion_diffusers_to_keras import ( +from zeromodels.diffusers.stable_diffusion.convert_stable_diffusion_diffusers_to_keras import ( config_from_diffusers as stable_diffusion_config_from_diffusers, ) -from zeromodels.models.stable_diffusion.convert_stable_diffusion_diffusers_to_keras import ( +from zeromodels.diffusers.stable_diffusion.convert_stable_diffusion_diffusers_to_keras import ( transfer_stable_diffusion, ) @@ -24,7 +24,7 @@ def config_from_diffusers(repo, token=None): - from zeromodels.models.stable_diffusion_2.stable_diffusion_2_config import ( + from zeromodels.diffusers.stable_diffusion_2.stable_diffusion_2_config import ( StableDiffusion2Config, ) @@ -34,18 +34,12 @@ def config_from_diffusers(repo, token=None): def transfer_stable_diffusion_2(repo, token=None): - from zeromodels.models.stable_diffusion_2.stable_diffusion_2_config import ( + from zeromodels.diffusers.stable_diffusion_2.stable_diffusion_2_config import ( StableDiffusion2Config, ) - from zeromodels.models.stable_diffusion_2.stable_diffusion_2_model import ( - StableDiffusion2Model, - ) return transfer_stable_diffusion( - repo, - token=token, - model_cls=StableDiffusion2Model, - config_cls=StableDiffusion2Config, + repo, token=token, config_cls=StableDiffusion2Config ) diff --git a/zeromodels/models/stable_diffusion_2/stable_diffusion_2_config.py b/zeromodels/diffusers/stable_diffusion_2/stable_diffusion_2_config.py similarity index 88% rename from zeromodels/models/stable_diffusion_2/stable_diffusion_2_config.py rename to zeromodels/diffusers/stable_diffusion_2/stable_diffusion_2_config.py index edf18cba..a7e31807 100644 --- a/zeromodels/models/stable_diffusion_2/stable_diffusion_2_config.py +++ b/zeromodels/diffusers/stable_diffusion_2/stable_diffusion_2_config.py @@ -1,8 +1,8 @@ -from zeromodels.models.clip.clip_config import CLIPTextConfig -from zeromodels.models.stable_diffusion.stable_diffusion_config import ( +from zeromodels.diffusers.stable_diffusion.stable_diffusion_config import ( StableDiffusionConfig, UNet2DConditionConfig, ) +from zeromodels.models.clip.clip_config import CLIPTextConfig class StableDiffusion2UNetConfig(UNet2DConditionConfig): @@ -34,7 +34,7 @@ class StableDiffusion2TextConfig(CLIPTextConfig): class StableDiffusion2Config(StableDiffusionConfig): - r"""Configuration for [`StableDiffusion2Model`], the hosted SD 2.x weights container. + r"""Configuration for a Stable Diffusion 2.x [`ZMDiffusion`] (``model_type="stable_diffusion_2"``). Same shape as [`StableDiffusionConfig`] (one nested ``unet_config`` / ``vae_config`` / ``text_config`` plus the scheduler and token ids; flat @@ -57,10 +57,11 @@ class StableDiffusion2Config(StableDiffusionConfig): Examples: ```python - >>> from zeromodels.models.stable_diffusion_2 import StableDiffusion2Config, StableDiffusion2Model + >>> from zeromodels.diffusers import ZMDiffusion + >>> from zeromodels.diffusers.stable_diffusion_2 import StableDiffusion2Config >>> config = StableDiffusion2Config() - >>> model = StableDiffusion2Model(config) # random weights, SD 2.x shapes + >>> model = ZMDiffusion(config) # random weights, SD 2.x shapes >>> model.config.unet_config.num_attention_heads (5, 10, 20, 20) ```""" diff --git a/zeromodels/diffusers/stable_diffusion_2/stable_diffusion_2_model.py b/zeromodels/diffusers/stable_diffusion_2/stable_diffusion_2_model.py new file mode 100644 index 00000000..16f41593 --- /dev/null +++ b/zeromodels/diffusers/stable_diffusion_2/stable_diffusion_2_model.py @@ -0,0 +1,22 @@ +from zeromodels.diffusers.stable_diffusion.stable_diffusion_model import ( + StableDiffusionFamily, +) + +from .stable_diffusion_2_config import StableDiffusion2Config + + +class StableDiffusion2Family(StableDiffusionFamily): + """Stable Diffusion 2.x inside :class:`ZMDiffusion` (``model_type="stable_diffusion_2"``). + + The SD 1.x family with the SD 2.x configuration (:class:`StableDiffusion2Config`): + the same three components and paths, with a 1024-d text context, one attention + head count per UNet level, a linear token projection in the Transformer2D blocks + and the 23-layer ``gelu`` OpenCLIP ViT-H/14 text tower. The checkpoint's + ``!``-padded empty prompt drives classifier-free guidance and its sampler (PNDM + for ``stable-diffusion-2-1-base``, DDIM for the others, v-prediction for the + 768px checkpoints) comes from the repo's ``scheduler_config``. The 768px repos + already carry ``unet_sample_size=96, vae_sample_size=768``. + """ + + model_type = "stable_diffusion_2" + config_class = StableDiffusion2Config diff --git a/zeromodels/models/stable_diffusion_3/__init__.py b/zeromodels/diffusers/stable_diffusion_3/__init__.py similarity index 70% rename from zeromodels/models/stable_diffusion_3/__init__.py rename to zeromodels/diffusers/stable_diffusion_3/__init__.py index df0f4cbe..87c680ce 100644 --- a/zeromodels/models/stable_diffusion_3/__init__.py +++ b/zeromodels/diffusers/stable_diffusion_3/__init__.py @@ -8,20 +8,16 @@ from .stable_diffusion_3_model import ( SD3T5EncoderModel, SD3Transformer2DModel, - StableDiffusion3Model, - StableDiffusion3TextToImage, + StableDiffusion3Family, ) -from .stable_diffusion_3_tokenizer import StableDiffusion3Tokenizer __all__ = [ "SD3T5EncoderModel", "SD3Transformer2DModel", - "StableDiffusion3Model", - "StableDiffusion3TextToImage", "StableDiffusion3Config", + "StableDiffusion3Family", "StableDiffusion3T5EncoderConfig", "StableDiffusion3TextConfig", "StableDiffusion3TransformerConfig", "StableDiffusion3VAEConfig", - "StableDiffusion3Tokenizer", ] diff --git a/zeromodels/models/stable_diffusion_3/convert_stable_diffusion_3_diffusers_to_keras.py b/zeromodels/diffusers/stable_diffusion_3/convert_stable_diffusion_3_diffusers_to_keras.py similarity index 94% rename from zeromodels/models/stable_diffusion_3/convert_stable_diffusion_3_diffusers_to_keras.py rename to zeromodels/diffusers/stable_diffusion_3/convert_stable_diffusion_3_diffusers_to_keras.py index 73f0958d..7fc9f9d0 100644 --- a/zeromodels/models/stable_diffusion_3/convert_stable_diffusion_3_diffusers_to_keras.py +++ b/zeromodels/diffusers/stable_diffusion_3/convert_stable_diffusion_3_diffusers_to_keras.py @@ -12,10 +12,10 @@ compare_keras_torch_names, transfer_weights, ) -from zeromodels.models.stable_diffusion.convert_stable_diffusion_diffusers_to_keras import ( +from zeromodels.diffusers.stable_diffusion.convert_stable_diffusion_diffusers_to_keras import ( WEIGHT_NAME_MAPPING, ) -from zeromodels.models.stable_diffusion_xl.convert_stable_diffusion_xl_diffusers_to_keras import ( +from zeromodels.diffusers.stable_diffusion_xl.convert_stable_diffusion_xl_diffusers_to_keras import ( LazyNumpyStateDict, text_config_kwargs, transfer_text_encoder, @@ -39,10 +39,10 @@ def config_from_diffusers(repo, token=None, config_cls=None): from diffusers import SD3Transformer2DModel as DiffusersSD3Transformer2DModel from transformers import CLIPTextConfig, CLIPTokenizerFast, T5TokenizerFast - from zeromodels.models.stable_diffusion_3.stable_diffusion_3_config import ( + from zeromodels.diffusers.stable_diffusion_3.stable_diffusion_3_config import ( StableDiffusion3Config, ) - from zeromodels.models.stable_diffusion_3.stable_diffusion_3_model import ( + from zeromodels.diffusers.stable_diffusion_3.stable_diffusion_3_model import ( SD3Transformer2DModel, ) @@ -133,9 +133,7 @@ def from_pretrained_fp16(cls, repo, subfolder, token=None): ) -def transfer_stable_diffusion_3( - repo, token=None, dtype="float16", model_cls=None, config_cls=None -): +def transfer_stable_diffusion_3(repo, token=None, dtype="float16", config_cls=None): import gc import torch @@ -144,14 +142,11 @@ def transfer_stable_diffusion_3( from transformers import CLIPTextModelWithProjection from zeromodels.base.base_mixin import build_dtype_scope - from zeromodels.models.stable_diffusion_3.stable_diffusion_3_model import ( - StableDiffusion3Model, - ) + from zeromodels.diffusers.diffusion_model import ZMDiffusion - model_cls = model_cls or StableDiffusion3Model config = config_from_diffusers(repo, token=token, config_cls=config_cls) with build_dtype_scope(dtype): - model = model_cls(config) + model = ZMDiffusion(config) legacy = { ".query.": ".to_q.", @@ -234,7 +229,7 @@ def transfer_t5_encoder(repo, token=None, dtype="float16"): from transformers import T5Config from zeromodels.base.base_mixin import build_dtype_scope - from zeromodels.models.stable_diffusion_3.stable_diffusion_3_model import ( + from zeromodels.diffusers.stable_diffusion_3.stable_diffusion_3_model import ( SD3T5EncoderModel, ) diff --git a/zeromodels/models/stable_diffusion_3/stable_diffusion_3_config.py b/zeromodels/diffusers/stable_diffusion_3/stable_diffusion_3_config.py similarity index 92% rename from zeromodels/models/stable_diffusion_3/stable_diffusion_3_config.py rename to zeromodels/diffusers/stable_diffusion_3/stable_diffusion_3_config.py index 3da5514f..121f6292 100644 --- a/zeromodels/models/stable_diffusion_3/stable_diffusion_3_config.py +++ b/zeromodels/diffusers/stable_diffusion_3/stable_diffusion_3_config.py @@ -1,8 +1,8 @@ from zeromodels.base import BaseConfig -from zeromodels.models.stable_diffusion.stable_diffusion_config import ( +from zeromodels.diffusers.stable_diffusion.stable_diffusion_config import ( AutoencoderKLConfig, ) -from zeromodels.models.stable_diffusion_xl.stable_diffusion_xl_config import ( +from zeromodels.diffusers.stable_diffusion_xl.stable_diffusion_xl_config import ( StableDiffusionXLTextConfig2, ) @@ -47,7 +47,7 @@ class StableDiffusion3TransformerConfig(BaseConfig): Examples: ```python - >>> from zeromodels.models.stable_diffusion_3 import StableDiffusion3TransformerConfig, SD3Transformer2DModel + >>> from zeromodels.diffusers.stable_diffusion_3 import StableDiffusion3TransformerConfig, SD3Transformer2DModel >>> config = StableDiffusion3TransformerConfig() >>> model = SD3Transformer2DModel(config) @@ -107,7 +107,7 @@ class StableDiffusion3T5EncoderConfig(BaseConfig): Examples: ```python - >>> from zeromodels.models.stable_diffusion_3 import StableDiffusion3T5EncoderConfig, SD3T5EncoderModel + >>> from zeromodels.diffusers.stable_diffusion_3 import StableDiffusion3T5EncoderConfig, SD3T5EncoderModel >>> config = StableDiffusion3T5EncoderConfig(num_layers=2, embed_dim=64, mlp_dim=128, num_heads=4) >>> model = SD3T5EncoderModel(config) @@ -156,9 +156,9 @@ class StableDiffusion3TextConfig(StableDiffusionXLTextConfig2): class StableDiffusion3Config(BaseConfig): - r"""Configuration for [`StableDiffusion3Model`], the hosted SD 3 weights container. + r"""Configuration for a Stable Diffusion 3 [`ZMDiffusion`] (``model_type="stable_diffusion_3"``). - One config for the container's components, so a repo carries a single + One config for the checkpoint's components, so a repo carries a single ``zm_config.json`` (nested ``transformer_config`` / ``vae_config`` / ``text_config`` / ``text_config_2``) plus the weights, and loads with the standard ``from_weights("zeromodels/")``. The constructor stays flat: @@ -196,10 +196,11 @@ class StableDiffusion3Config(BaseConfig): Examples: ```python - >>> from zeromodels.models.stable_diffusion_3 import StableDiffusion3Config, StableDiffusion3Model + >>> from zeromodels.diffusers import ZMDiffusion + >>> from zeromodels.diffusers.stable_diffusion_3 import StableDiffusion3Config >>> config = StableDiffusion3Config() - >>> model = StableDiffusion3Model(config) # random weights, SD 3 medium shapes + >>> model = ZMDiffusion(config) # random weights, SD 3 medium shapes >>> model.config.transformer_config.num_layers 24 ```""" diff --git a/zeromodels/models/stable_diffusion_3/stable_diffusion_3_layers.py b/zeromodels/diffusers/stable_diffusion_3/stable_diffusion_3_layers.py similarity index 99% rename from zeromodels/models/stable_diffusion_3/stable_diffusion_3_layers.py rename to zeromodels/diffusers/stable_diffusion_3/stable_diffusion_3_layers.py index 9b373c0d..f7ef3b0b 100644 --- a/zeromodels/models/stable_diffusion_3/stable_diffusion_3_layers.py +++ b/zeromodels/diffusers/stable_diffusion_3/stable_diffusion_3_layers.py @@ -3,7 +3,7 @@ from keras import layers, ops from zeromodels.base.base_attention import active_attn_implementation, fused_attention -from zeromodels.models.stable_diffusion.stable_diffusion_layers import safe_name +from zeromodels.diffusers.stable_diffusion.stable_diffusion_layers import safe_name from zeromodels.models.t5.t5_layers import T5LayerNorm, T5SelfAttentionLayer NORM_EPS = 1e-6 diff --git a/zeromodels/models/stable_diffusion_3/stable_diffusion_3_model.py b/zeromodels/diffusers/stable_diffusion_3/stable_diffusion_3_model.py similarity index 76% rename from zeromodels/models/stable_diffusion_3/stable_diffusion_3_model.py rename to zeromodels/diffusers/stable_diffusion_3/stable_diffusion_3_model.py index 9ea4a7fe..5805f0ac 100644 --- a/zeromodels/models/stable_diffusion_3/stable_diffusion_3_model.py +++ b/zeromodels/diffusers/stable_diffusion_3/stable_diffusion_3_model.py @@ -2,18 +2,18 @@ from keras import layers, ops from zeromodels.base import BaseModel +from zeromodels.base.base_diffusion import BaseDiffusion from zeromodels.base.base_mixin import inference_scope from zeromodels.base.base_scheduler import FlowMatchEulerDiscreteScheduler -from zeromodels.models.stable_diffusion.stable_diffusion_layers import ( +from zeromodels.diffusers.stable_diffusion.stable_diffusion_layers import ( time_embedding_mlp, timestep_embedding, ) -from zeromodels.models.stable_diffusion.stable_diffusion_model import ( +from zeromodels.diffusers.stable_diffusion.stable_diffusion_model import ( AutoencoderKL, - StableDiffusionModel, - StableDiffusionTextToImage, + StableDiffusionFamily, ) -from zeromodels.models.stable_diffusion_xl.stable_diffusion_xl_model import ( +from zeromodels.diffusers.stable_diffusion_xl.stable_diffusion_xl_model import ( build_text_encoder, ) from zeromodels.models.t5.t5_layers import T5LayerNorm @@ -31,12 +31,6 @@ StableDiffusion3T5GatedEncoderBlock, ) -# The container and the task build the identical graph, so both load the one -# hosted repo (whose zm_config.json names the container). -STABLE_DIFFUSION_3_HUB_SIBLINGS = frozenset( - {"StableDiffusion3Model", "StableDiffusion3TextToImage"} -) - @keras.saving.register_keras_serializable(package="zeromodels") class SD3Transformer2DModel(BaseModel): @@ -245,14 +239,14 @@ class SD3T5EncoderModel(BaseModel): blocks, the gated-GELU feed-forward of T5 v1.1) and a final RMSNorm. The attention, norm and position-bias layers are the t5 family's; the gated feed-forward, which the original T5 does not have, is - :class:`~zeromodels.models.stable_diffusion_3.stable_diffusion_3_layers.StableDiffusion3T5GatedFeedForward`. + :class:`~zeromodels.diffusers.stable_diffusion_3.stable_diffusion_3_layers.StableDiffusion3T5GatedFeedForward`. Inputs ``{"input_ids", "attention_mask"}`` (``(B, T)`` int32), output ``{"last_hidden_state": (B, T, embed_dim)}``. Every SD 3 / 3.5 checkpoint conditions on the same 4.7B-parameter encoder, so it is hosted once (``zeromodels/t5-v1_1-xxl-encoder``, float16 weights) and attached - to a task model on demand: - ``StableDiffusion3TextToImage.from_weights(repo, text_encoder_3="zeromodels/t5-v1_1-xxl-encoder")``. + to a :class:`ZMDiffusion` on demand: + ``ZMDiffusion.from_weights(repo, text_encoder_3="zeromodels/t5-v1_1-xxl-encoder")``. Args: vocab_size: SentencePiece vocabulary (32128). @@ -400,40 +394,36 @@ def get_config(self): return config -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusion3Model(StableDiffusionModel): - """Stable Diffusion 3 weights container: the MMDiT, the 16-channel VAE and the - two CLIP text encoders (ViT-L/14 and OpenCLIP ViT-bigG/14, each with its - projection) as one functional ``keras.Model``. - - The SD 1.x container (:class:`StableDiffusionModel`) with the SD 3 parts: four - disconnected sub-graphs, ``{"sample", "timestep", "encoder_hidden_states", - "pooled_projections"} -> "noise_pred"`` (the transformer), ``{"image", - "latent"} -> "moments" / "image"`` (the VAE) and ``{"token_ids", "token_ids_2", - "padding_mask"} -> "prompt_embeds" / "pooled_prompt_embeds"`` (the CLIP towers: - penultimate states concatenated to 2048-d, projected pooled states concatenated - to 2048-d). Components: ``.transformer`` / ``.vae`` / ``.text_encoder`` / - ``.text_encoder_2``. - - The third text encoder, the 4.7B-parameter T5-XXL, is not part of the - container: it is shared by every SD 3 / 3.5 checkpoint and hosted once - (``zeromodels/t5-v1_1-xxl-encoder``, an :class:`SD3T5EncoderModel`), and the - task class attaches it on demand (``text_encoder_3``); without it the - T5 features are zeros, SD 3's documented memory-saving mode. Hosted as one - ``zm_config.json`` + float16 weights per variant; on-the-fly ``hf:`` conversion - is not supported for diffusion models. - - Args: - **kwargs: The flat :class:`StableDiffusion3Config` fields - (``transformer_`` / ``vae_`` / ``text_`` / ``text_2_`` prefixed), or a - config positionally. +class StableDiffusion3Family(StableDiffusionFamily): + """Stable Diffusion 3 inside :class:`ZMDiffusion` (``model_type="stable_diffusion_3"``). + + Four components, ``transformer`` (the MMDiT, :class:`SD3Transformer2DModel`), + ``vae`` (16 channels) and the two CLIP towers ``text_encoder`` (ViT-L/14) / + ``text_encoder_2`` (OpenCLIP ViT-bigG/14), each with its projection, wired into + four disconnected paths: ``{"sample", "timestep", "encoder_hidden_states", + "pooled_projections"} -> "noise_pred"``, ``{"image", "latent"} -> "moments" / + "image"`` and ``{"token_ids", "token_ids_2", "padding_mask"} -> "prompt_embeds" / + "pooled_prompt_embeds"`` (penultimate states concatenated to 2048-d, projected + pooled states concatenated to 2048-d). + + The third text encoder, the 4.7B-parameter T5-XXL, is not part of the checkpoint: + every SD 3 / 3.5 checkpoint shares it, so it is hosted once + (``zeromodels/t5-v1_1-xxl-encoder``, an :class:`SD3T5EncoderModel`) and attached + with ``ZMDiffusion.from_weights(repo, text_encoder_3=...)`` or by assigning + ``model.text_encoder_3``. Without it the T5 features are zeros, SD 3's documented + memory-saving mode. + + Generation follows diffusers' ``StableDiffusion3Pipeline`` with the rectified-flow + sampler (:class:`FlowMatchEulerDiscreteScheduler`): the CLIP states, zero-padded + to the T5 width, followed by the T5 tokens (``input_ids_3``, 256) along the + sequence, the pooled CLIP states along the features. With no negative prompt the + empty prompt is encoded (no zeroing). """ + model_type = "stable_diffusion_3" config_class = StableDiffusion3Config - HUB_REPO_SIBLINGS = STABLE_DIFFUSION_3_HUB_SIBLINGS - - def __init__(self, name="StableDiffusion3Model", **kwargs): - super().__init__(name=name, **kwargs) + generate_args = {"num_inference_steps": 28, "guidance_scale": 7.0} + supports_text_encoder_3 = True def build_components(self, config, data_format, channels_axis): d, v, t, t2 = ( @@ -567,118 +557,44 @@ def build_graph(self, config, components, data_format): } return inputs, outputs - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusion3TextToImage(StableDiffusionTextToImage, StableDiffusion3Model): - """Text-to-image Stable Diffusion 3, pure Keras 3 and cross-backend (the - diffusers ``StableDiffusion3Pipeline`` in this library's task-class form). - - :class:`StableDiffusionTextToImage`'s ``generate`` (``BaseDiffusion``) over the - :class:`StableDiffusion3Model` graph with the rectified-flow sampler - (:class:`FlowMatchEulerDiscreteScheduler`). The prompt goes through the two - CLIP towers (penultimate states concatenated, zero-padded to the T5 width) and, - when a T5-XXL encoder is attached, through it (256 tokens); the text tokens are - concatenated along the sequence, the pooled CLIP states along the features. With - no negative prompt the empty prompt is encoded (no zeroing). The T5 encoder is - loaded separately and optionally:: - - sd = StableDiffusion3TextToImage.from_weights( - "zeromodels/stable-diffusion-3-medium", - text_encoder_3="zeromodels/t5-v1_1-xxl-encoder", # or omit: T5 features zeroed - ) - tokenizer = StableDiffusion3Tokenizer.from_weights("zeromodels/stable-diffusion-3-medium") - image = sd.generate(**tokenizer("a photo of a cat")) # (1, 1024, 1024, 3) uint8 - - ``sd.text_encoder_3`` can also be assigned directly (an - :class:`SD3T5EncoderModel`, e.g. one loaded with ``quantization="int8"``); it - lives outside the container's weights. - - Args: - scheduler: A :class:`BaseScheduler`; defaults to the config's - ``scheduler_config``, else the flow-match Euler sampler with shift 3. - **kwargs: The flat :class:`StableDiffusion3Config` fields. - """ - - config_class = StableDiffusion3Config - HUB_REPO_SIBLINGS = STABLE_DIFFUSION_3_HUB_SIBLINGS - # diffusers' StableDiffusion3Pipeline defaults; a repo's generate_args override - generate_args = {"num_inference_steps": 28, "guidance_scale": 7.0} - - def __init__(self, scheduler=None, name="StableDiffusion3TextToImage", **kwargs): - super().__init__(scheduler=scheduler, name=name, **kwargs) - - @classmethod - def from_weights(cls, identifier, text_encoder_3=None, **kwargs): - """``BaseModel.from_weights`` plus ``text_encoder_3``: a hosted - :class:`SD3T5EncoderModel` repo id (loaded with the same ``load_dtype``) or - a built model to attach as the third text encoder.""" - model = super().from_weights(identifier, **kwargs) - if text_encoder_3 is not None: - if isinstance(text_encoder_3, str): - text_encoder_3 = SD3T5EncoderModel.from_weights( - text_encoder_3, load_dtype=kwargs.get("load_dtype") - ) - model.text_encoder_3 = text_encoder_3 - return model - - @property - def text_encoder_3(self): - return self.__dict__.get("_text_encoder_3") - - def __setattr__(self, name, value): - # the T5 encoder is kept out of the tracked (saved / loaded / device-moved) - # sub-layers, like the LM mixins' caches: Keras (and torch's Module on that - # backend) would otherwise register any model assigned as an attribute - if name == "text_encoder_3": - self.__dict__["_text_encoder_3"] = value - return - super().__setattr__(name, value) - def default_scheduler(self): return FlowMatchEulerDiscreteScheduler(shift=3.0) - @property - def latent_shape(self): - size = self.transformer.sample_size + def latent_shape(self, model): + size = model.transformer.sample_size height, width = size if isinstance(size, (tuple, list)) else (size, size) - if self.data_format == "channels_first": - return (self.transformer.in_channels, height, width) - return (height, width, self.transformer.in_channels) - - def unconditional_ids(self, batch): - cfg = self.config - length = cfg.text_config.max_seq_len - row = [cfg.bos_token_id, cfg.eos_token_id] + [cfg.pad_token_id] * (length - 2) - return ops.convert_to_tensor([row] * batch, dtype="int32") + if model.data_format == "channels_first": + return (model.transformer.in_channels, height, width) + return (height, width, model.transformer.in_channels) - def unconditional_ids_3(self, batch): - cfg = self.config + def unconditional_ids_3(self, model, batch): + cfg = model.config row = [cfg.eos_token_id_3] + [cfg.pad_token_id_3] * ( cfg.max_sequence_length - 1 ) return ops.convert_to_tensor([row] * batch, dtype="int32") - def prompt_mask(self, input_ids): + def prompt_mask(self, model, input_ids): # the tokens up to and including the first <|endoftext|>; what follows is # padding (the first tokenizer pads with <|endoftext|> itself) - is_eos = ops.cast(ops.equal(input_ids, self.config.eos_token_id), "int32") + is_eos = ops.cast(ops.equal(input_ids, model.config.eos_token_id), "int32") return ops.equal(ops.cumsum(is_eos, axis=1) - is_eos, 0) - def encode_prompt(self, input_ids, attention_mask=None, input_ids_3=None): + def encode_prompt(self, model, input_ids, attention_mask=None, input_ids_3=None): input_ids = ops.cast(ops.convert_to_tensor(input_ids), "int32") batch = int(input_ids.shape[0]) if attention_mask is None: - real = self.prompt_mask(input_ids) + real = self.prompt_mask(model, input_ids) else: real = ops.greater( ops.cast(ops.convert_to_tensor(attention_mask), "int32"), 0 ) # the second tokenizer pads with "!" (id 0) where the first pads with # <|endoftext|>; the towers attend the padding, as the reference does - input_ids_2 = ops.where(real, input_ids, self.config.pad_token_id_2) + input_ids_2 = ops.where(real, input_ids, model.config.pad_token_id_2) ones = ops.ones_like(input_ids) - out = self.text_encoder({"token_ids": input_ids, "padding_mask": ones}) - out_2 = self.text_encoder_2({"token_ids": input_ids_2, "padding_mask": ones}) + out = model.text_encoder({"token_ids": input_ids, "padding_mask": ones}) + out_2 = model.text_encoder_2({"token_ids": input_ids_2, "padding_mask": ones}) clip_embeds = ops.concatenate( [out["penultimate_hidden_state"], out_2["penultimate_hidden_state"]], axis=-1, @@ -686,20 +602,20 @@ def encode_prompt(self, input_ids, attention_mask=None, input_ids_3=None): pooled = ops.concatenate([out["text_embeds"], out_2["text_embeds"]], axis=-1) # the CLIP features are zero-padded to the T5 width, then the T5 tokens # (or zeros without a T5 encoder) follow along the sequence - t5_dim = self.transformer.joint_attention_dim + t5_dim = model.transformer.joint_attention_dim clip_embeds = ops.pad( clip_embeds, ((0, 0), (0, 0), (0, t5_dim - clip_embeds.shape[-1])) ) - length = self.config.max_sequence_length - if self.text_encoder_3 is None: + length = model.config.max_sequence_length + if model.text_encoder_3 is None: t5_embeds = ops.zeros((batch, length, t5_dim), dtype=clip_embeds.dtype) else: if input_ids_3 is None: - input_ids_3 = self.unconditional_ids_3(batch) + input_ids_3 = self.unconditional_ids_3(model, batch) input_ids_3 = ops.cast(ops.convert_to_tensor(input_ids_3), "int32") if int(input_ids_3.shape[0]) == 1 and batch > 1: input_ids_3 = ops.repeat(input_ids_3, batch, axis=0) # one for all - t5_embeds = self.text_encoder_3( + t5_embeds = model.text_encoder_3( {"input_ids": input_ids_3, "attention_mask": ops.ones_like(input_ids_3)} )["last_hidden_state"] t5_embeds = ops.cast(t5_embeds, clip_embeds.dtype) @@ -708,13 +624,15 @@ def encode_prompt(self, input_ids, attention_mask=None, input_ids_3=None): "pooled_projections": pooled, } - def encode_negative_prompt(self, negative_input_ids, batch, **conditioning): + def encode_negative_prompt(self, model, negative_input_ids, batch, **conditioning): if negative_input_ids is None: # the empty prompt for the T5 tower as well, not the positive prompt's ids conditioning = dict(conditioning, input_ids_3=None) - return super().encode_negative_prompt(negative_input_ids, batch, **conditioning) + return BaseDiffusion.encode_negative_prompt( + model, negative_input_ids, batch, **conditioning + ) - def predict_noise(self, latents, timesteps, embeddings): - return self.transformer( + def predict_noise(self, model, latents, timesteps, embeddings): + return model.transformer( {"sample": latents, "timestep": timesteps, **embeddings} )["sample"] diff --git a/zeromodels/diffusers/stable_diffusion_3_5/__init__.py b/zeromodels/diffusers/stable_diffusion_3_5/__init__.py new file mode 100644 index 00000000..ff0e8ded --- /dev/null +++ b/zeromodels/diffusers/stable_diffusion_3_5/__init__.py @@ -0,0 +1,11 @@ +from .stable_diffusion_3_5_config import ( + StableDiffusion3_5Config, + StableDiffusion3_5TransformerConfig, +) +from .stable_diffusion_3_5_model import StableDiffusion3_5Family + +__all__ = [ + "StableDiffusion3_5Config", + "StableDiffusion3_5Family", + "StableDiffusion3_5TransformerConfig", +] diff --git a/zeromodels/models/stable_diffusion_3_5/convert_stable_diffusion_3_5_diffusers_to_keras.py b/zeromodels/diffusers/stable_diffusion_3_5/convert_stable_diffusion_3_5_diffusers_to_keras.py similarity index 76% rename from zeromodels/models/stable_diffusion_3_5/convert_stable_diffusion_3_5_diffusers_to_keras.py rename to zeromodels/diffusers/stable_diffusion_3_5/convert_stable_diffusion_3_5_diffusers_to_keras.py index 52c919c3..7170b8ae 100644 --- a/zeromodels/models/stable_diffusion_3_5/convert_stable_diffusion_3_5_diffusers_to_keras.py +++ b/zeromodels/diffusers/stable_diffusion_3_5/convert_stable_diffusion_3_5_diffusers_to_keras.py @@ -1,9 +1,9 @@ import numpy as np -from zeromodels.models.stable_diffusion_3.convert_stable_diffusion_3_diffusers_to_keras import ( +from zeromodels.diffusers.stable_diffusion_3.convert_stable_diffusion_3_diffusers_to_keras import ( config_from_diffusers as stable_diffusion_3_config_from_diffusers, ) -from zeromodels.models.stable_diffusion_3.convert_stable_diffusion_3_diffusers_to_keras import ( +from zeromodels.diffusers.stable_diffusion_3.convert_stable_diffusion_3_diffusers_to_keras import ( transfer_stable_diffusion_3, ) @@ -15,7 +15,7 @@ def config_from_diffusers(repo, token=None): - from zeromodels.models.stable_diffusion_3_5.stable_diffusion_3_5_config import ( + from zeromodels.diffusers.stable_diffusion_3_5.stable_diffusion_3_5_config import ( StableDiffusion3_5Config, ) @@ -25,19 +25,12 @@ def config_from_diffusers(repo, token=None): def transfer_stable_diffusion_3_5(repo, token=None, dtype="float16"): - from zeromodels.models.stable_diffusion_3_5.stable_diffusion_3_5_config import ( + from zeromodels.diffusers.stable_diffusion_3_5.stable_diffusion_3_5_config import ( StableDiffusion3_5Config, ) - from zeromodels.models.stable_diffusion_3_5.stable_diffusion_3_5_model import ( - StableDiffusion3_5Model, - ) return transfer_stable_diffusion_3( - repo, - token=token, - dtype=dtype, - model_cls=StableDiffusion3_5Model, - config_cls=StableDiffusion3_5Config, + repo, token=token, dtype=dtype, config_cls=StableDiffusion3_5Config ) diff --git a/zeromodels/models/stable_diffusion_3_5/stable_diffusion_3_5_config.py b/zeromodels/diffusers/stable_diffusion_3_5/stable_diffusion_3_5_config.py similarity index 81% rename from zeromodels/models/stable_diffusion_3_5/stable_diffusion_3_5_config.py rename to zeromodels/diffusers/stable_diffusion_3_5/stable_diffusion_3_5_config.py index f71459a8..30a8010e 100644 --- a/zeromodels/models/stable_diffusion_3_5/stable_diffusion_3_5_config.py +++ b/zeromodels/diffusers/stable_diffusion_3_5/stable_diffusion_3_5_config.py @@ -1,4 +1,4 @@ -from zeromodels.models.stable_diffusion_3.stable_diffusion_3_config import ( +from zeromodels.diffusers.stable_diffusion_3.stable_diffusion_3_config import ( StableDiffusion3Config, StableDiffusion3TransformerConfig, ) @@ -21,7 +21,7 @@ class StableDiffusion3_5TransformerConfig(StableDiffusion3TransformerConfig): class StableDiffusion3_5Config(StableDiffusion3Config): - r"""Configuration for [`StableDiffusion3_5Model`], the hosted SD 3.5 weights container. + r"""Configuration for a Stable Diffusion 3.5 [`ZMDiffusion`] (``model_type="stable_diffusion_3_5"``). Same shape as [`StableDiffusion3Config`] (nested ``transformer_config`` / ``vae_config`` / ``text_config`` / ``text_config_2`` plus the scheduler, the @@ -32,10 +32,11 @@ class StableDiffusion3_5Config(StableDiffusion3Config): Examples: ```python - >>> from zeromodels.models.stable_diffusion_3_5 import StableDiffusion3_5Config, StableDiffusion3_5Model + >>> from zeromodels.diffusers import ZMDiffusion + >>> from zeromodels.diffusers.stable_diffusion_3_5 import StableDiffusion3_5Config >>> config = StableDiffusion3_5Config() - >>> model = StableDiffusion3_5Model(config) # random weights, SD 3.5 large shapes + >>> model = ZMDiffusion(config) # random weights, SD 3.5 large shapes >>> model.config.transformer_config.qk_norm 'rms_norm' ```""" diff --git a/zeromodels/diffusers/stable_diffusion_3_5/stable_diffusion_3_5_model.py b/zeromodels/diffusers/stable_diffusion_3_5/stable_diffusion_3_5_model.py new file mode 100644 index 00000000..111a3657 --- /dev/null +++ b/zeromodels/diffusers/stable_diffusion_3_5/stable_diffusion_3_5_model.py @@ -0,0 +1,22 @@ +from zeromodels.diffusers.stable_diffusion_3.stable_diffusion_3_model import ( + StableDiffusion3Family, +) + +from .stable_diffusion_3_5_config import StableDiffusion3_5Config + + +class StableDiffusion3_5Family(StableDiffusion3Family): + """Stable Diffusion 3.5 inside :class:`ZMDiffusion` (``model_type="stable_diffusion_3_5"``). + + The SD 3 family with the SD 3.5 configuration (:class:`StableDiffusion3_5Config`): + RMS-normalized queries and keys in every MMDiT block and, for the medium + checkpoint, the dual-attention (MMDiT-X) blocks. Same components + (``transformer`` / ``vae`` / ``text_encoder`` / ``text_encoder_2``), same + separately hosted T5-XXL (``text_encoder_3``). Each repo's ``generate_args`` hold + the checkpoint's defaults: 28 steps at guidance 3.5 for large, 40 at 4.5 for + medium, 4 without guidance for large-turbo. + """ + + model_type = "stable_diffusion_3_5" + config_class = StableDiffusion3_5Config + generate_args = {"num_inference_steps": 28, "guidance_scale": 3.5} diff --git a/zeromodels/models/stable_diffusion_xl/__init__.py b/zeromodels/diffusers/stable_diffusion_xl/__init__.py similarity index 57% rename from zeromodels/models/stable_diffusion_xl/__init__.py rename to zeromodels/diffusers/stable_diffusion_xl/__init__.py index 0943c298..1c217fd0 100644 --- a/zeromodels/models/stable_diffusion_xl/__init__.py +++ b/zeromodels/diffusers/stable_diffusion_xl/__init__.py @@ -7,23 +7,17 @@ StableDiffusionXLVAEConfig, ) from .stable_diffusion_xl_model import ( - StableDiffusionXLModel, - StableDiffusionXLRefinerImageToImage, - StableDiffusionXLRefinerModel, - StableDiffusionXLTextToImage, + StableDiffusionXLFamily, + StableDiffusionXLRefinerFamily, ) -from .stable_diffusion_xl_tokenizer import StableDiffusionXLTokenizer __all__ = [ - "StableDiffusionXLModel", - "StableDiffusionXLTextToImage", - "StableDiffusionXLRefinerModel", - "StableDiffusionXLRefinerImageToImage", "StableDiffusionXLConfig", + "StableDiffusionXLFamily", "StableDiffusionXLRefinerConfig", + "StableDiffusionXLRefinerFamily", "StableDiffusionXLRefinerUNetConfig", "StableDiffusionXLTextConfig2", "StableDiffusionXLUNetConfig", "StableDiffusionXLVAEConfig", - "StableDiffusionXLTokenizer", ] diff --git a/zeromodels/models/stable_diffusion_xl/convert_stable_diffusion_xl_diffusers_to_keras.py b/zeromodels/diffusers/stable_diffusion_xl/convert_stable_diffusion_xl_diffusers_to_keras.py similarity index 94% rename from zeromodels/models/stable_diffusion_xl/convert_stable_diffusion_xl_diffusers_to_keras.py rename to zeromodels/diffusers/stable_diffusion_xl/convert_stable_diffusion_xl_diffusers_to_keras.py index b29628d7..08d6a924 100644 --- a/zeromodels/models/stable_diffusion_xl/convert_stable_diffusion_xl_diffusers_to_keras.py +++ b/zeromodels/diffusers/stable_diffusion_xl/convert_stable_diffusion_xl_diffusers_to_keras.py @@ -12,10 +12,10 @@ compare_keras_torch_names, transfer_weights, ) -from zeromodels.models.clip.convert_clip_hf_to_keras import transfer_clip_weights -from zeromodels.models.stable_diffusion.convert_stable_diffusion_diffusers_to_keras import ( +from zeromodels.diffusers.stable_diffusion.convert_stable_diffusion_diffusers_to_keras import ( WEIGHT_NAME_MAPPING, ) +from zeromodels.models.clip.convert_clip_hf_to_keras import transfer_clip_weights # SDXL ships in float16 (the release checkpoints' precision; the diffusers fp32 # files are upcasts of the fp16 variant), so the container is built and saved in @@ -23,7 +23,8 @@ # rebuilds at that precision by default; load_dtype="float32" upcasts. The 0.9 # repos are gated (research license): set HF_TOKEN to convert them. The refiner # repos (an image-to-image model: one text tower, a four-level UNet) convert to -# StableDiffusionXLRefinerModel, everything else to StableDiffusionXLModel. +# a ZMDiffusion of model_type stable_diffusion_xl_refiner, everything else to +# stable_diffusion_xl. STABLE_DIFFUSION_XL_SOURCES = { "stable-diffusion-xl-base-0.9": "stabilityai/stable-diffusion-xl-base-0.9", "stable-diffusion-xl-refiner-0.9": "stabilityai/stable-diffusion-xl-refiner-0.9", @@ -67,10 +68,10 @@ def config_from_diffusers(repo, token=None): from diffusers.pipelines.pipeline_utils import DiffusionPipeline from transformers import CLIPTextConfig, CLIPTokenizerFast - from zeromodels.models.stable_diffusion.stable_diffusion_model import ( + from zeromodels.diffusers.stable_diffusion.stable_diffusion_model import ( UNet2DConditionModel as KerasUNet2DConditionModel, ) - from zeromodels.models.stable_diffusion_xl.stable_diffusion_xl_config import ( + from zeromodels.diffusers.stable_diffusion_xl.stable_diffusion_xl_config import ( StableDiffusionXLConfig, StableDiffusionXLRefinerConfig, ) @@ -166,19 +167,11 @@ def transfer_stable_diffusion_xl(repo, token=None, dtype="float16"): from transformers import CLIPTextModel, CLIPTextModelWithProjection from zeromodels.base.base_mixin import build_dtype_scope - from zeromodels.models.stable_diffusion_xl.stable_diffusion_xl_model import ( - StableDiffusionXLModel, - StableDiffusionXLRefinerModel, - ) + from zeromodels.diffusers.diffusion_model import ZMDiffusion config = config_from_diffusers(repo, token=token) - model_cls = ( - StableDiffusionXLRefinerModel - if config.text_config is None - else StableDiffusionXLModel - ) with build_dtype_scope(dtype): - model = model_cls(config) + model = ZMDiffusion(config) fp16 = {"torch_dtype": torch.float16, "variant": "fp16", "token": token} fp32 = {"torch_dtype": torch.float32, "token": token} diff --git a/zeromodels/models/stable_diffusion_xl/stable_diffusion_xl_config.py b/zeromodels/diffusers/stable_diffusion_xl/stable_diffusion_xl_config.py similarity index 92% rename from zeromodels/models/stable_diffusion_xl/stable_diffusion_xl_config.py rename to zeromodels/diffusers/stable_diffusion_xl/stable_diffusion_xl_config.py index 0ed30070..ae35aba1 100644 --- a/zeromodels/models/stable_diffusion_xl/stable_diffusion_xl_config.py +++ b/zeromodels/diffusers/stable_diffusion_xl/stable_diffusion_xl_config.py @@ -1,10 +1,10 @@ -from zeromodels.models.clip.clip_config import CLIPTextConfig -from zeromodels.models.stable_diffusion.stable_diffusion_config import ( +from zeromodels.diffusers.stable_diffusion.stable_diffusion_config import ( AutoencoderKLConfig, StableDiffusionConfig, StableDiffusionTextConfig, UNet2DConditionConfig, ) +from zeromodels.models.clip.clip_config import CLIPTextConfig class StableDiffusionXLUNetConfig(UNet2DConditionConfig): @@ -70,7 +70,7 @@ class StableDiffusionXLTextConfig2(CLIPTextConfig): class StableDiffusionXLConfig(StableDiffusionConfig): - r"""Configuration for [`StableDiffusionXLModel`], the hosted SDXL weights container. + r"""Configuration for a Stable Diffusion XL [`ZMDiffusion`] (``model_type="stable_diffusion_xl"``). The [`StableDiffusionConfig`] shape (nested sub-configs plus the scheduler and token ids; flat constructor) with a fourth component: SDXL conditions on two @@ -103,10 +103,11 @@ class StableDiffusionXLConfig(StableDiffusionConfig): Examples: ```python - >>> from zeromodels.models.stable_diffusion_xl import StableDiffusionXLConfig, StableDiffusionXLModel + >>> from zeromodels.diffusers import ZMDiffusion + >>> from zeromodels.diffusers.stable_diffusion_xl import StableDiffusionXLConfig >>> config = StableDiffusionXLConfig() - >>> model = StableDiffusionXLModel(config) # random weights, SDXL shapes + >>> model = ZMDiffusion(config) # random weights, SDXL shapes >>> model.config.unet_config.transformer_layers_per_block (1, 2, 10) ```""" @@ -165,8 +166,8 @@ class StableDiffusionXLRefinerUNetConfig(StableDiffusionXLUNetConfig): class StableDiffusionXLRefinerConfig(StableDiffusionXLConfig): - r"""Configuration for [`StableDiffusionXLRefinerModel`], the hosted SDXL refiner - weights container. + r"""Configuration for a Stable Diffusion XL refiner [`ZMDiffusion`] + (``model_type="stable_diffusion_xl_refiner"``). The [`StableDiffusionXLConfig`] shape with the refiner's parts: no first text tower (``text_config`` is ``None``; the OpenCLIP ViT-bigG/14 ``text_config_2`` @@ -191,10 +192,11 @@ class StableDiffusionXLRefinerConfig(StableDiffusionXLConfig): Examples: ```python - >>> from zeromodels.models.stable_diffusion_xl import StableDiffusionXLRefinerConfig, StableDiffusionXLRefinerModel + >>> from zeromodels.diffusers import ZMDiffusion + >>> from zeromodels.diffusers.stable_diffusion_xl import StableDiffusionXLRefinerConfig >>> config = StableDiffusionXLRefinerConfig() - >>> model = StableDiffusionXLRefinerModel(config) # random weights, refiner shapes + >>> model = ZMDiffusion(config) # random weights, refiner shapes >>> model.config.unet_config.block_out_channels (384, 768, 1536, 1536) ```""" diff --git a/zeromodels/diffusers/stable_diffusion_xl/stable_diffusion_xl_model.py b/zeromodels/diffusers/stable_diffusion_xl/stable_diffusion_xl_model.py new file mode 100644 index 00000000..63d2b741 --- /dev/null +++ b/zeromodels/diffusers/stable_diffusion_xl/stable_diffusion_xl_model.py @@ -0,0 +1,365 @@ +import keras +from keras import layers, ops + +from zeromodels.base.base_diffusion import BaseDiffusion +from zeromodels.base.base_scheduler import EulerDiscreteScheduler +from zeromodels.diffusers.stable_diffusion.stable_diffusion_model import ( + AutoencoderKL, + StableDiffusionFamily, + UNet2DConditionModel, +) +from zeromodels.models.clip.clip_layers import CLIPTextModelEmbedding +from zeromodels.models.clip.clip_model import residual_attention_block + +from .stable_diffusion_xl_config import ( + StableDiffusionXLConfig, + StableDiffusionXLRefinerConfig, +) + + +def build_text_encoder( + config, hidden_act, layer_norm_eps, projection_dim=None, name="text_encoder" +): + """Functional CLIP text tower (the ``clip`` family's layers and leaf names) + that also exposes the penultimate hidden state SDXL conditions on. + + Outputs ``penultimate_hidden_state`` (the input of the last transformer block, + no final LayerNorm: diffusers' ``hidden_states[-2]``), the usual + ``last_hidden_state`` / ``pooler_output`` (EOT position, argmax of the ids) and, + when ``projection_dim`` is set, ``text_embeds``, the pooled state through the + bias-free ``text_projection`` (the second SDXL tower). Built under a + ``keras.name_scope(name)`` so the two towers' identically named layers keep + distinct weight paths inside the container. + """ + length = config.max_seq_len + with keras.name_scope(name): + # int32 inputs: under a float16 load policy a float input would round the + # ids (49407 is not representable in float16) + token_ids = layers.Input(shape=(length,), dtype="int32", name="token_ids") + padding_mask = layers.Input(shape=(length,), dtype="int32", name="padding_mask") + x = CLIPTextModelEmbedding( + vocab_size=config.vocab_size, + max_seq_len=length, + embed_dim=config.hidden_dim, + name="text_model_embedding", + )(token_ids) + + causal_mask = ops.cast(ops.triu(ops.ones((length, length)), k=1), "float32") + causal_mask = causal_mask * (-1e8) + key_mask = ops.reshape(ops.cast(padding_mask, "float32"), (-1, 1, 1, length)) + key_mask = (1.0 - ops.repeat(key_mask, length, axis=2)) * (-1e8) + + penultimate = None + for i in range(config.num_layers): + if i == config.num_layers - 1: + penultimate = x + x = residual_attention_block( + x, + proj_dim=config.hidden_dim, + num_heads=config.num_heads, + layer_name_prefix="text_model_encoder", + layer_idx=i, + causal_attention_mask=causal_mask, + attention_mask=key_mask, + mlp_ratio=config.mlp_ratio, + hidden_act=hidden_act, + layer_norm_eps=layer_norm_eps, + ) + last_hidden_state = layers.LayerNormalization( + epsilon=layer_norm_eps, name="text_model_layernorm" + )(x) + one_hot = ops.one_hot( + ops.argmax(token_ids, axis=-1), length, dtype=last_hidden_state.dtype + ) + pooler_output = ops.einsum("bi,bij->bj", one_hot, last_hidden_state) + + outputs = { + "penultimate_hidden_state": penultimate, + "last_hidden_state": last_hidden_state, + "pooler_output": pooler_output, + } + if projection_dim is not None: + # the Dense runs on a (B, 1, D) view, as CLIPTextEmbed does, so the + # kernel maps to the checkpoint's text_projection.weight unchanged + projected = layers.Dense( + projection_dim, use_bias=False, name="text_projection" + )(ops.expand_dims(pooler_output, axis=1)) + outputs["text_embeds"] = ops.squeeze(projected, axis=1) + return keras.Model( + inputs={"token_ids": token_ids, "padding_mask": padding_mask}, + outputs=outputs, + name=name, + ) + + +class StableDiffusionXLFamily(StableDiffusionFamily): + """Stable Diffusion XL inside :class:`ZMDiffusion` (``model_type="stable_diffusion_xl"``). + + Four components, ``unet`` / ``vae`` / ``text_encoder`` (CLIP ViT-L/14) / + ``text_encoder_2`` (OpenCLIP ViT-bigG/14), wired into four disconnected paths: + ``{"sample", "timestep", "encoder_hidden_states", "text_embeds", "time_ids"} -> + "noise_pred"`` (the UNet with its ``text_time`` micro-conditioning), ``{"image", + "latent"} -> "moments" / "image"`` (the VAE, float32 whatever the load dtype, + ``force_upcast``) and ``{"token_ids", "token_ids_2", "padding_mask"} -> + "prompt_embeds" / "pooled_prompt_embeds"`` (the two towers' penultimate states + concatenated, and the second tower's projected pooled state). + + Generation follows diffusers' ``StableDiffusionXLPipeline``: the second tower's + ``!``-padded ids are derived from the first tokenizer's output, the pooled state + becomes the ``text_embeds`` micro-conditioning next to the ``time_ids`` built from + ``original_size`` / ``crops_coords_top_left`` / ``target_size`` (``(height, + width)`` pairs, defaulting to the generated size and ``(0, 0)``; ``negative_`` + variants condition the negative branch). With no negative prompt the + unconditional branch is zero embeddings (``force_zeros_for_empty_prompt``). The + sampler is the repo's ``scheduler_config`` (Euler with ``leading`` spacing for the + base model, ancestral Euler with ``trailing`` spacing for SDXL-Turbo). For the + ensemble of experts, run the base with ``denoising_end=0.8, + output_type="latent"`` and hand the latent to the refiner. + """ + + model_type = "stable_diffusion_xl" + config_class = StableDiffusionXLConfig + generate_args = {"num_inference_steps": 50, "guidance_scale": 5.0} + + def build_components(self, config, data_format, channels_axis): + u, v, t, t2 = ( + config.unet_config, + config.vae_config, + config.text_config, + config.text_config_2, + ) + unet = UNet2DConditionModel( + u, data_format=data_format, channels_axis=channels_axis + ) + vae = AutoencoderKL( + **v.constructor_kwargs(), + data_format=data_format, + channels_axis=channels_axis, + ) + components = {"unet": unet, "vae": vae} + if t is not None: # the refiner has no CLIP ViT-L/14 tower + components["text_encoder"] = build_text_encoder( + t, config.hidden_act, config.layer_norm_eps, name="text_encoder" + ) + components["text_encoder_2"] = build_text_encoder( + t2, + t2.hidden_act, + config.layer_norm_eps, + projection_dim=t2.projection_dim, + name="text_encoder_2", + ) + return components + + def build_graph(self, config, components, data_format): + u, length = config.unet_config, config.text_config_2.max_seq_len + unet, vae = components["unet"], components["vae"] + text_encoder = components.get("text_encoder") + text_encoder_2 = components["text_encoder_2"] + inputs = self.unet_vae_inputs(config, components, data_format) + pooled_dim = ( + u.projection_class_embeddings_input_dim + - u.num_time_ids * u.addition_time_embed_dim + ) + inputs["text_embeds"] = layers.Input(shape=(pooled_dim,), name="text_embeds") + inputs["time_ids"] = layers.Input(shape=(u.num_time_ids,), name="time_ids") + if text_encoder is not None: + inputs["token_ids"] = layers.Input( + shape=(length,), dtype="int32", name="token_ids" + ) + inputs["token_ids_2"] = layers.Input( + shape=(length,), dtype="int32", name="token_ids_2" + ) + inputs["padding_mask"] = layers.Input( + shape=(length,), dtype="int32", name="padding_mask" + ) + + noise_pred = unet( + { + key: inputs[key] + for key in ( + "sample", + "timestep", + "encoder_hidden_states", + "text_embeds", + "time_ids", + ) + } + )["sample"] + vae_out = vae({"image": inputs["image"], "latent": inputs["latent"]}) + text_out_2 = text_encoder_2( + { + "token_ids": inputs["token_ids_2"], + "padding_mask": inputs["padding_mask"], + } + ) + hidden_states = [text_out_2["penultimate_hidden_state"]] + if text_encoder is not None: + text_out = text_encoder( + { + "token_ids": inputs["token_ids"], + "padding_mask": inputs["padding_mask"], + } + ) + hidden_states.insert(0, text_out["penultimate_hidden_state"]) + prompt_embeds = ( + ops.concatenate(hidden_states, axis=-1) + if len(hidden_states) > 1 + else hidden_states[0] + ) + outputs = { + "noise_pred": noise_pred, + "moments": vae_out["moments"], + "image": vae_out["sample"], + "prompt_embeds": prompt_embeds, + "pooled_prompt_embeds": text_out_2["text_embeds"], + } + return inputs, outputs + + def default_scheduler(self): + return EulerDiscreteScheduler(timestep_spacing="leading", steps_offset=1) + + def generate( + self, + model, + input_ids, + crops_coords_top_left=(0, 0), + aesthetic_score=6.0, + negative_aesthetic_score=2.5, + **kwargs, + ): + # the micro-conditioning defaults of diffusers' SDXL pipelines; aesthetic + # scores replace the target size on the refiner (requires_aesthetics_score) + return BaseDiffusion.generate( + model, + input_ids, + crops_coords_top_left=crops_coords_top_left, + aesthetic_score=aesthetic_score, + negative_aesthetic_score=negative_aesthetic_score, + **kwargs, + ) + + def image_size(self, model): + size = model.vae.sample_size + return tuple(size) if isinstance(size, (tuple, list)) else (size, size) + + def time_ids( + self, + model, + batch, + original_size=None, + crops_coords_top_left=(0, 0), + target_size=None, + aesthetic_score=6.0, + ): + # the (batch, 6) float32 micro-conditioning row: original size, crop offset + # and target size, each (height, width); (batch, 5) with the aesthetic score + # in place of the target size when requires_aesthetics_score (the refiner) + original_size = tuple(original_size or self.image_size(model)) + row = list(original_size) + list(crops_coords_top_left) + if model.config.requires_aesthetics_score: + row = row + [aesthetic_score] + else: + row = row + list(target_size or self.image_size(model)) + return ops.convert_to_tensor([row] * batch, dtype="float32") + + def unconditional_ids(self, model, batch): + cfg = model.config + length = cfg.text_config_2.max_seq_len + row = [cfg.bos_token_id, cfg.eos_token_id] + [cfg.pad_token_id] * (length - 2) + return ops.convert_to_tensor([row] * batch, dtype="int32") + + def prompt_mask(self, model, input_ids): + # the tokens up to and including the first <|endoftext|>; what follows is + # padding (the first tokenizer pads with <|endoftext|> itself) + is_eos = ops.cast(ops.equal(input_ids, model.config.eos_token_id), "int32") + return ops.equal(ops.cumsum(is_eos, axis=1) - is_eos, 0) + + def encode_prompt( + self, + model, + input_ids, + attention_mask=None, + original_size=None, + crops_coords_top_left=(0, 0), + target_size=None, + aesthetic_score=6.0, + ): + input_ids = ops.cast(ops.convert_to_tensor(input_ids), "int32") + if attention_mask is None: + real = self.prompt_mask(model, input_ids) + else: + real = ops.greater( + ops.cast(ops.convert_to_tensor(attention_mask), "int32"), 0 + ) + # the second tokenizer pads with "!" (id 0) where the first pads with + # <|endoftext|>; the towers attend the padding (no attention mask), as the + # reference does + input_ids_2 = ops.where(real, input_ids, model.config.pad_token_id_2) + ones = ops.ones_like(input_ids) + out_2 = model.text_encoder_2({"token_ids": input_ids_2, "padding_mask": ones}) + hidden = out_2["penultimate_hidden_state"] + if model.config.text_config is not None: + out = model.text_encoder({"token_ids": input_ids, "padding_mask": ones}) + hidden = ops.concatenate([out["penultimate_hidden_state"], hidden], -1) + batch = int(input_ids.shape[0]) + return { + "encoder_hidden_states": hidden, + "text_embeds": out_2["text_embeds"], + "time_ids": self.time_ids( + model, + batch, + original_size, + crops_coords_top_left, + target_size, + aesthetic_score, + ), + } + + def encode_negative_prompt(self, model, negative_input_ids, batch, **conditioning): + if ( + negative_input_ids is not None + or not model.config.force_zeros_for_empty_prompt + ): + return BaseDiffusion.encode_negative_prompt( + model, negative_input_ids, batch, **conditioning + ) + # force_zeros_for_empty_prompt: zero embeddings stand in for the empty + # prompt; the micro-conditioning keeps its (negative) values + u, length = model.config.unet_config, model.config.text_config_2.max_seq_len + pooled_dim = ( + u.projection_class_embeddings_input_dim + - u.num_time_ids * u.addition_time_embed_dim + ) + return { + "encoder_hidden_states": ops.zeros( + (batch, length, u.cross_attention_dim), dtype="float32" + ), + "text_embeds": ops.zeros((batch, pooled_dim), dtype="float32"), + "time_ids": self.time_ids(model, batch, **conditioning), + } + + def predict_noise(self, model, latents, timesteps, embeddings): + return model.unet({"sample": latents, "timestep": timesteps, **embeddings})[ + "sample" + ] + + +class StableDiffusionXLRefinerFamily(StableDiffusionXLFamily): + """The Stable Diffusion XL refiner inside :class:`ZMDiffusion` + (``model_type="stable_diffusion_xl_refiner"``). + + The SDXL family with the refiner configuration + (:class:`StableDiffusionXLRefinerConfig`): no CLIP ViT-L/14 tower (the UNet's + context is the OpenCLIP ViT-bigG/14 tower's 1280-d penultimate state alone), the + four-level refiner UNet and five micro-conditioning ids, the aesthetic score + (``aesthetic_score`` 6.0 / ``negative_aesthetic_score`` 2.5) in place of the + target size. It is an image-to-image checkpoint (diffusers' + ``StableDiffusionXLImg2ImgPipeline``): ``generate`` refines an ``image`` (noised + to ``strength``, 0.3 by default) or, as the second half of the ensemble of + experts, a latent the base left partially denoised (``denoising_start`` at the + base's ``denoising_end``). + """ + + model_type = "stable_diffusion_xl_refiner" + config_class = StableDiffusionXLRefinerConfig + generate_args = {"num_inference_steps": 50, "guidance_scale": 5.0, "strength": 0.3} diff --git a/zeromodels/models/__init__.py b/zeromodels/models/__init__.py index 74301652..15b61d82 100644 --- a/zeromodels/models/__init__.py +++ b/zeromodels/models/__init__.py @@ -104,8 +104,6 @@ qwen3_next, qwen3_vl, qwen3_vl_moe, - qwen_image, - qwen_image_21, regnet, res2net, resmlp, @@ -124,11 +122,6 @@ siglip, siglip2, speech2text, - stable_diffusion, - stable_diffusion_2, - stable_diffusion_3, - stable_diffusion_3_5, - stable_diffusion_xl, swin, swinv2, t5, diff --git a/zeromodels/models/qwen_image/qwen_image_tokenizer.py b/zeromodels/models/qwen_image/qwen_image_tokenizer.py deleted file mode 100644 index ab798e9b..00000000 --- a/zeromodels/models/qwen_image/qwen_image_tokenizer.py +++ /dev/null @@ -1,58 +0,0 @@ -"""Qwen-Image tokenizer: Qwen2 ChatML template for prompt encoding.""" - -import keras - -from zeromodels.models.qwen2.qwen2_tokenizer import Qwen2Tokenizer - -PROMPT_TEMPLATE = ( - "<|im_start|>system\nDescribe the image by detailing the color, shape, size, " - "texture, quantity, text, spatial relationships of the objects and " - "background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" -) -PROMPT_TEMPLATE_START_IDX = 34 - - -@keras.saving.register_keras_serializable(package="zeromodels") -class QwenImageTokenizer(Qwen2Tokenizer): - """Tokenizer for Qwen-Image text-to-image. - - Wraps each prompt in the Diffusers ChatML template, then BPE-encodes with - the Qwen2 tokenizer. Returns ``input_ids`` / ``attention_mask`` for - :meth:`QwenImageTextToImage.generate` (the task drops the template prefix - after the text encoder). - - Load:: - - tok = QwenImageTokenizer.from_weights("zeromodels/qwen-image") - model.generate(**tok("a photo of a cat")) - """ - - prompt_template = PROMPT_TEMPLATE - prompt_template_start_idx = PROMPT_TEMPLATE_START_IDX - - def __init__( - self, - hf_id=None, - tokenizer_file=None, - max_seq_len=1024, - **kwargs, - ): - self.max_seq_len = max_seq_len - self.tokenizer_max_length = max_seq_len - super().__init__(hf_id=hf_id, tokenizer_file=tokenizer_file, **kwargs) - - def format_prompt(self, text): - return self.prompt_template.format(text) - - def call(self, inputs): - texts = self.normalize_texts(inputs) - templated = [self.format_prompt(t) for t in texts] - max_length = self.tokenizer_max_length + self.prompt_template_start_idx - encoded = [self.encode(t)[:max_length] for t in templated] - input_ids, attention_mask = self.pad_batch(encoded) - return {"input_ids": input_ids, "attention_mask": attention_mask} - - def get_config(self): - config = super().get_config() - config.update({"max_seq_len": self.max_seq_len}) - return config diff --git a/zeromodels/models/qwen_image_21/qwen_image_21_tokenizer.py b/zeromodels/models/qwen_image_21/qwen_image_21_tokenizer.py deleted file mode 100644 index f0975ba3..00000000 --- a/zeromodels/models/qwen_image_21/qwen_image_21_tokenizer.py +++ /dev/null @@ -1,57 +0,0 @@ -import keras - -from zeromodels.models.qwen3.qwen3_tokenizer import Qwen3Tokenizer - -SYS_PROMPT = "Comprehend and analyze the provided prompt." -PROMPT_TEMPLATE = ( - f"<|im_start|>system\n{SYS_PROMPT}<|im_end|>\n" - f"<|im_start|>user\n{{}}<|im_end|>\n" - f"<|im_start|>assistant\n" -) -PROMPT_TEMPLATE_START_IDX = 14 - - -@keras.saving.register_keras_serializable(package="zeromodels") -class QwenImage21Tokenizer(Qwen3Tokenizer): - """Tokenizer for Qwen-Image-2.1 text-to-image. - - Wraps each prompt in the Diffusers ChatML template (system prompt - ``Comprehend and analyze the provided prompt.``), then BPE-encodes. - Returns ``input_ids`` / ``attention_mask`` for - :meth:`QwenImage21TextToImage.generate`. - - Load:: - - tok = QwenImage21Tokenizer.from_weights("zeromodels/qwen-image-2.1") - model.generate(**tok("a photo of a cat")) - """ - - prompt_template = PROMPT_TEMPLATE - prompt_template_start_idx = PROMPT_TEMPLATE_START_IDX - - def __init__( - self, - hf_id=None, - tokenizer_file=None, - max_seq_len=1024, - **kwargs, - ): - self.max_seq_len = max_seq_len - self.tokenizer_max_length = max_seq_len - super().__init__(hf_id=hf_id, tokenizer_file=tokenizer_file, **kwargs) - - def format_prompt(self, text): - return self.prompt_template.format(text if text else " ") - - def call(self, inputs): - texts = self.normalize_texts(inputs) - templated = [self.format_prompt(t) for t in texts] - max_length = self.tokenizer_max_length + self.prompt_template_start_idx - encoded = [self.encode(t)[:max_length] for t in templated] - input_ids, attention_mask = self.pad_batch(encoded) - return {"input_ids": input_ids, "attention_mask": attention_mask} - - def get_config(self): - config = super().get_config() - config.update({"max_seq_len": self.max_seq_len}) - return config diff --git a/zeromodels/models/stable_diffusion/stable_diffusion_tokenizer.py b/zeromodels/models/stable_diffusion/stable_diffusion_tokenizer.py deleted file mode 100644 index 6a20d1a0..00000000 --- a/zeromodels/models/stable_diffusion/stable_diffusion_tokenizer.py +++ /dev/null @@ -1,70 +0,0 @@ -import keras - -from zeromodels.models.clip.clip_tokenizer import CLIPTokenizer - - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusionTokenizer(CLIPTokenizer): - """Stable Diffusion 1.x text tokenizer: the CLIP ViT-L/14 BPE tokenizer. - - SD 1.x conditions on OpenAI's CLIP ViT-L/14 text encoder, so this is that - tokenizer byte for byte: byte-pair encoding with ``<|startoftext|>`` / - ``<|endoftext|>`` framing, truncated and ``<|endoftext|>``-padded to - ``max_seq_len`` (77), the fixed text length the UNet's cross-attention is built - for. Loads by repo id like the model, reading the hosted repo's root - ``tokenizer.json``: ``from_weights("zeromodels/stable-diffusion-v1-5")``. - - Args: - hf_id: Hosted repo to read ``tokenizer.json`` from. Required unless - ``tokenizer_file`` is given (no default repo). - tokenizer_file: Explicit ``tokenizer.json`` path (overrides ``hf_id``). - max_seq_len: Padded / truncated length (default 77). - pad_token: Pad token string (``<|endoftext|>``). - """ - - def __init__( - self, - hf_id=None, - tokenizer_file=None, - max_seq_len=77, - pad_token="<|endoftext|>", - **kwargs, - ): - if tokenizer_file is None and hf_id is None: - raise ValueError( - f"{type(self).__name__}() needs an hf_id (a hosted repo, read from its " - "tokenizer.json) or an explicit tokenizer_file: there is no default " - "repo. Load it by repo id, e.g. " - "from_weights('zeromodels/stable-diffusion-v1-5')." - ) - if tokenizer_file is None: - tokenizer_file = self.download_tokenizer_json(hf_id) - super().__init__( - tokenizer_file=tokenizer_file, - max_seq_len=max_seq_len, - pad_token=pad_token, - **kwargs, - ) - self.hf_id = hf_id - - @staticmethod - def download_tokenizer_json(hf_id): - import os - - from huggingface_hub import hf_hub_download - - return hf_hub_download( - hf_id, "tokenizer.json", token=os.environ.get("HF_TOKEN") - ) - - @classmethod - def from_hf(cls, repo, **kwargs): - raise NotImplementedError( - "On-the-fly `hf:` loading is not supported for diffusion models: use " - f"{cls.__name__}.from_weights('zeromodels/')." - ) - - def get_config(self): - config = super().get_config() - config["hf_id"] = self.hf_id - return config diff --git a/zeromodels/models/stable_diffusion_2/__init__.py b/zeromodels/models/stable_diffusion_2/__init__.py deleted file mode 100644 index 45b64e47..00000000 --- a/zeromodels/models/stable_diffusion_2/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -from .stable_diffusion_2_config import ( - StableDiffusion2Config, - StableDiffusion2TextConfig, - StableDiffusion2UNetConfig, -) -from .stable_diffusion_2_model import ( - StableDiffusion2Model, - StableDiffusion2TextToImage, -) -from .stable_diffusion_2_tokenizer import StableDiffusion2Tokenizer - -__all__ = [ - "StableDiffusion2Model", - "StableDiffusion2TextToImage", - "StableDiffusion2Config", - "StableDiffusion2TextConfig", - "StableDiffusion2UNetConfig", - "StableDiffusion2Tokenizer", -] diff --git a/zeromodels/models/stable_diffusion_2/stable_diffusion_2_model.py b/zeromodels/models/stable_diffusion_2/stable_diffusion_2_model.py deleted file mode 100644 index e27ec14c..00000000 --- a/zeromodels/models/stable_diffusion_2/stable_diffusion_2_model.py +++ /dev/null @@ -1,69 +0,0 @@ -import keras - -from zeromodels.models.stable_diffusion.stable_diffusion_model import ( - StableDiffusionModel, - StableDiffusionTextToImage, -) - -from .stable_diffusion_2_config import StableDiffusion2Config - -# The container and the task build the identical graph, so both load the one -# hosted repo (whose zm_config.json names the container). -STABLE_DIFFUSION_2_HUB_SIBLINGS = frozenset( - {"StableDiffusion2Model", "StableDiffusion2TextToImage"} -) - - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusion2Model(StableDiffusionModel): - """Stable Diffusion 2.x weights container: the UNet, the VAE and the OpenCLIP - ViT-H/14 text encoder as one functional ``keras.Model``. - - The SD 1.x container with the SD 2.x configuration - (:class:`StableDiffusion2Config`): the same three disconnected sub-graphs and - the same components (``.unet`` / ``.vae`` / ``.text_encoder``), with a 1024-d - text context, one attention head count per UNet level, a linear token - projection in the Transformer2D blocks and the 23-layer ``gelu`` text tower. - Hosted as one ``zm_config.json`` + ``model.weights.h5`` per checkpoint under - ``zeromodels/stable-diffusion-2*``; on-the-fly ``hf:`` conversion is not - supported for diffusion models. - - Args: - **kwargs: The flat :class:`StableDiffusion2Config` fields (``unet_`` / - ``vae_`` / ``text_`` prefixed), or a config positionally. Pass - ``unet_sample_size=96, vae_sample_size=768`` for the 768px checkpoints - (their repos already carry it). - """ - - config_class = StableDiffusion2Config - HUB_REPO_SIBLINGS = STABLE_DIFFUSION_2_HUB_SIBLINGS - - def __init__(self, name="StableDiffusion2Model", **kwargs): - super().__init__(name=name, **kwargs) - - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusion2TextToImage(StableDiffusionTextToImage): - """Text-to-image Stable Diffusion 2.x, pure Keras 3 and cross-backend. - - :class:`StableDiffusionTextToImage` with the SD 2.x configuration: the same - ``generate`` (``BaseDiffusion``) over the SD 2 graph, so the checkpoint's - ``!``-padded empty prompt drives classifier-free guidance and its scheduler - (PNDM for ``stable-diffusion-2-1-base``, DDIM for the others, v-prediction for - the 768px checkpoints) comes from the repo's ``scheduler_config``:: - - sd = StableDiffusion2TextToImage.from_weights("zeromodels/stable-diffusion-2-1") - tokenizer = StableDiffusion2Tokenizer.from_weights("zeromodels/stable-diffusion-2-1") - image = sd.generate(**tokenizer("a photo of a cat")) # (1, 768, 768, 3) uint8 - - Args: - scheduler: A :class:`BaseScheduler`; defaults to the config's - ``scheduler_config``. - **kwargs: The flat :class:`StableDiffusion2Config` fields. - """ - - config_class = StableDiffusion2Config - HUB_REPO_SIBLINGS = STABLE_DIFFUSION_2_HUB_SIBLINGS - - def __init__(self, scheduler=None, name="StableDiffusion2TextToImage", **kwargs): - super().__init__(scheduler=scheduler, name=name, **kwargs) diff --git a/zeromodels/models/stable_diffusion_2/stable_diffusion_2_tokenizer.py b/zeromodels/models/stable_diffusion_2/stable_diffusion_2_tokenizer.py deleted file mode 100644 index c9d33103..00000000 --- a/zeromodels/models/stable_diffusion_2/stable_diffusion_2_tokenizer.py +++ /dev/null @@ -1,34 +0,0 @@ -import keras - -from zeromodels.models.stable_diffusion.stable_diffusion_tokenizer import ( - StableDiffusionTokenizer, -) - - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusion2Tokenizer(StableDiffusionTokenizer): - """Stable Diffusion 2.x text tokenizer: the OpenCLIP ViT-H/14 BPE tokenizer. - - The same byte-pair encoding and ``<|startoftext|>`` / ``<|endoftext|>`` framing - as CLIP ViT-L/14's, truncated to 77 tokens, but padded with ``!`` (id 0) the - way OpenCLIP does, which is what the SD 2 text encoder was trained on. Loads by - repo id like the model: ``from_weights("zeromodels/stable-diffusion-2-1")``. - - Args: - hf_id: Hosted repo to read ``tokenizer.json`` from. Required unless - ``tokenizer_file`` is given (no default repo). - tokenizer_file: Explicit ``tokenizer.json`` path (overrides ``hf_id``). - max_seq_len: Padded / truncated length (default 77). - pad_token: Pad token string (``!``). - """ - - def __init__( - self, hf_id=None, tokenizer_file=None, max_seq_len=77, pad_token="!", **kwargs - ): - super().__init__( - hf_id=hf_id, - tokenizer_file=tokenizer_file, - max_seq_len=max_seq_len, - pad_token=pad_token, - **kwargs, - ) diff --git a/zeromodels/models/stable_diffusion_3/stable_diffusion_3_tokenizer.py b/zeromodels/models/stable_diffusion_3/stable_diffusion_3_tokenizer.py deleted file mode 100644 index ee4ebb65..00000000 --- a/zeromodels/models/stable_diffusion_3/stable_diffusion_3_tokenizer.py +++ /dev/null @@ -1,89 +0,0 @@ -import keras -from keras import ops - -from zeromodels.models.stable_diffusion.stable_diffusion_tokenizer import ( - StableDiffusionTokenizer, -) -from zeromodels.models.t5.t5_tokenizer import T5Tokenizer - - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusion3Tokenizer(StableDiffusionTokenizer): - """Stable Diffusion 3 text tokenizer: the CLIP ViT-L/14 BPE tokenizer plus the - T5 SentencePiece tokenizer, in one call. - - SD 3 tokenizes the prompt three times: with the two CLIP tokenizers (the same - BPE, padded with ``<|endoftext|>`` and ``!`` respectively; this class is the - first, and :class:`StableDiffusion3TextToImage` derives the second tower's ids - from the returned ``attention_mask``, as SDXL does) and with the T5 tokenizer - (````-terminated, ````-padded to ``max_sequence_length``, 256). The - returned dict is ``{"input_ids", "attention_mask", "input_ids_3"}``, what - ``generate`` takes. Loads by repo id like the model, reading the hosted repo's - ``tokenizer.json`` (CLIP) and ``tokenizer_3.json`` (T5): - ``from_weights("zeromodels/stable-diffusion-3-medium")``. - - Args: - hf_id: Hosted repo to read the tokenizer files from. Required unless the - files are given (no default repo). - tokenizer_file: Explicit CLIP ``tokenizer.json`` path (overrides ``hf_id``). - tokenizer_file_3: Explicit T5 ``tokenizer.json`` path (overrides ``hf_id``). - max_seq_len: CLIP padded / truncated length (default 77). - max_sequence_length: T5 padded / truncated length (default 256). - pad_token: CLIP pad token string (``<|endoftext|>``). - """ - - def __init__( - self, - hf_id=None, - tokenizer_file=None, - tokenizer_file_3=None, - max_seq_len=77, - max_sequence_length=256, - pad_token="<|endoftext|>", - **kwargs, - ): - super().__init__( - hf_id=hf_id, - tokenizer_file=tokenizer_file, - max_seq_len=max_seq_len, - pad_token=pad_token, - **kwargs, - ) - if tokenizer_file_3 is None: - tokenizer_file_3 = self.download_tokenizer_json(hf_id, "tokenizer_3.json") - self.tokenizer_file_3 = tokenizer_file_3 - self.max_sequence_length = max_sequence_length - self.tokenizer_3 = T5Tokenizer( - tokenizer_file=tokenizer_file_3, max_seq_len=max_sequence_length - ) - - @staticmethod - def download_tokenizer_json(hf_id, filename="tokenizer.json"): - import os - - from huggingface_hub import hf_hub_download - - return hf_hub_download(hf_id, filename, token=os.environ.get("HF_TOKEN")) - - def call(self, inputs): - out = super().call(inputs) - # the T5 tokenizer truncates to max_sequence_length and pads to the longest - # prompt; the transformer takes the fixed 256-token length, so pad the rest - ids = self.tokenizer_3(inputs)["input_ids"] - short = self.max_sequence_length - int(ids.shape[1]) - if short > 0: - ids = ops.pad( - ids, ((0, 0), (0, short)), constant_values=self.tokenizer_3.pad_token_id - ) - out["input_ids_3"] = ids - return out - - def get_config(self): - config = super().get_config() - config.update( - { - "tokenizer_file_3": self.tokenizer_file_3, - "max_sequence_length": self.max_sequence_length, - } - ) - return config diff --git a/zeromodels/models/stable_diffusion_3_5/__init__.py b/zeromodels/models/stable_diffusion_3_5/__init__.py deleted file mode 100644 index 29ea0ff8..00000000 --- a/zeromodels/models/stable_diffusion_3_5/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -from .stable_diffusion_3_5_config import ( - StableDiffusion3_5Config, - StableDiffusion3_5TransformerConfig, -) -from .stable_diffusion_3_5_model import ( - StableDiffusion3_5Model, - StableDiffusion3_5TextToImage, -) -from .stable_diffusion_3_5_tokenizer import StableDiffusion3_5Tokenizer - -__all__ = [ - "StableDiffusion3_5Model", - "StableDiffusion3_5TextToImage", - "StableDiffusion3_5Config", - "StableDiffusion3_5TransformerConfig", - "StableDiffusion3_5Tokenizer", -] diff --git a/zeromodels/models/stable_diffusion_3_5/stable_diffusion_3_5_model.py b/zeromodels/models/stable_diffusion_3_5/stable_diffusion_3_5_model.py deleted file mode 100644 index 12c56582..00000000 --- a/zeromodels/models/stable_diffusion_3_5/stable_diffusion_3_5_model.py +++ /dev/null @@ -1,69 +0,0 @@ -import keras - -from zeromodels.models.stable_diffusion_3.stable_diffusion_3_model import ( - StableDiffusion3Model, - StableDiffusion3TextToImage, -) - -from .stable_diffusion_3_5_config import StableDiffusion3_5Config - -# The container and the task build the identical graph, so both load the one -# hosted repo (whose zm_config.json names the container). -STABLE_DIFFUSION_3_5_HUB_SIBLINGS = frozenset( - {"StableDiffusion3_5Model", "StableDiffusion3_5TextToImage"} -) - - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusion3_5Model(StableDiffusion3Model): - """Stable Diffusion 3.5 weights container: the MMDiT (large) or MMDiT-X (medium), - the 16-channel VAE and the two CLIP text encoders as one functional - ``keras.Model``. - - The SD 3 container (:class:`StableDiffusion3Model`) with the SD 3.5 - configuration (:class:`StableDiffusion3_5Config`): RMS-normalized queries and - keys in every block and, for the medium checkpoint, the dual-attention blocks. - Same components (``.transformer`` / ``.vae`` / ``.text_encoder`` / - ``.text_encoder_2``), same separately hosted T5-XXL. Hosted under - ``zeromodels/stable-diffusion-3.5-*`` in float16. - - Args: - **kwargs: The flat :class:`StableDiffusion3_5Config` fields, or a config - positionally. - """ - - config_class = StableDiffusion3_5Config - HUB_REPO_SIBLINGS = STABLE_DIFFUSION_3_5_HUB_SIBLINGS - - def __init__(self, name="StableDiffusion3_5Model", **kwargs): - super().__init__(name=name, **kwargs) - - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusion3_5TextToImage(StableDiffusion3TextToImage): - """Text-to-image Stable Diffusion 3.5, pure Keras 3 and cross-backend. - - :class:`StableDiffusion3TextToImage` with the SD 3.5 configuration: the same - ``generate`` over the SD 3.5 graph, the repo's flow-match sampler and - defaults (28 steps at guidance 3.5 for large, 40 at 4.5 for medium, 4 without - guidance for large-turbo):: - - sd = StableDiffusion3_5TextToImage.from_weights( - "zeromodels/stable-diffusion-3.5-large", text_encoder_3="zeromodels/t5-v1_1-xxl-encoder" - ) - tokenizer = StableDiffusion3_5Tokenizer.from_weights("zeromodels/stable-diffusion-3.5-large") - image = sd.generate(**tokenizer("a photo of a cat")) # (1, 1024, 1024, 3) uint8 - - Args: - scheduler: A :class:`BaseScheduler`; defaults to the config's - ``scheduler_config``. - **kwargs: The flat :class:`StableDiffusion3_5Config` fields. - """ - - config_class = StableDiffusion3_5Config - HUB_REPO_SIBLINGS = STABLE_DIFFUSION_3_5_HUB_SIBLINGS - # diffusers' defaults for SD 3.5 large; each repo's generate_args override - generate_args = {"num_inference_steps": 28, "guidance_scale": 3.5} - - def __init__(self, scheduler=None, name="StableDiffusion3_5TextToImage", **kwargs): - super().__init__(scheduler=scheduler, name=name, **kwargs) diff --git a/zeromodels/models/stable_diffusion_3_5/stable_diffusion_3_5_tokenizer.py b/zeromodels/models/stable_diffusion_3_5/stable_diffusion_3_5_tokenizer.py deleted file mode 100644 index b3259457..00000000 --- a/zeromodels/models/stable_diffusion_3_5/stable_diffusion_3_5_tokenizer.py +++ /dev/null @@ -1,39 +0,0 @@ -import keras - -from zeromodels.models.stable_diffusion_3.stable_diffusion_3_tokenizer import ( - StableDiffusion3Tokenizer, -) - - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusion3_5Tokenizer(StableDiffusion3Tokenizer): - """Stable Diffusion 3.5 text tokenizer: the SD 3 tokenizer (CLIP BPE plus the T5 - SentencePiece tokenizer, ``{"input_ids", "attention_mask", "input_ids_3"}``); - the checkpoints share their tokenizers with SD 3. Loads by repo id like the - model: ``from_weights("zeromodels/stable-diffusion-3.5-large")``. - - Args: - hf_id: Hosted repo to read the tokenizer files from. Required unless the - files are given (no default repo). - tokenizer_file / tokenizer_file_3: Explicit CLIP / T5 ``tokenizer.json`` paths. - max_seq_len: CLIP padded / truncated length (default 77). - max_sequence_length: T5 padded / truncated length (default 256). - """ - - def __init__( - self, - hf_id=None, - tokenizer_file=None, - tokenizer_file_3=None, - max_seq_len=77, - max_sequence_length=256, - **kwargs, - ): - super().__init__( - hf_id=hf_id, - tokenizer_file=tokenizer_file, - tokenizer_file_3=tokenizer_file_3, - max_seq_len=max_seq_len, - max_sequence_length=max_sequence_length, - **kwargs, - ) diff --git a/zeromodels/models/stable_diffusion_xl/stable_diffusion_xl_model.py b/zeromodels/models/stable_diffusion_xl/stable_diffusion_xl_model.py deleted file mode 100644 index ccb94a9d..00000000 --- a/zeromodels/models/stable_diffusion_xl/stable_diffusion_xl_model.py +++ /dev/null @@ -1,496 +0,0 @@ -import keras -from keras import layers, ops - -from zeromodels.base.base_scheduler import EulerDiscreteScheduler -from zeromodels.models.clip.clip_layers import CLIPTextModelEmbedding -from zeromodels.models.clip.clip_model import residual_attention_block -from zeromodels.models.stable_diffusion.stable_diffusion_model import ( - AutoencoderKL, - StableDiffusionModel, - StableDiffusionTextToImage, - UNet2DConditionModel, -) - -from .stable_diffusion_xl_config import ( - StableDiffusionXLConfig, - StableDiffusionXLRefinerConfig, -) - -# The container and the task build the identical graph, so both load the one -# hosted repo (whose zm_config.json names the container). -STABLE_DIFFUSION_XL_HUB_SIBLINGS = frozenset( - {"StableDiffusionXLModel", "StableDiffusionXLTextToImage"} -) -STABLE_DIFFUSION_XL_REFINER_HUB_SIBLINGS = frozenset( - {"StableDiffusionXLRefinerModel", "StableDiffusionXLRefinerImageToImage"} -) - - -def build_text_encoder( - config, hidden_act, layer_norm_eps, projection_dim=None, name="text_encoder" -): - """Functional CLIP text tower (the ``clip`` family's layers and leaf names) - that also exposes the penultimate hidden state SDXL conditions on. - - Outputs ``penultimate_hidden_state`` (the input of the last transformer block, - no final LayerNorm: diffusers' ``hidden_states[-2]``), the usual - ``last_hidden_state`` / ``pooler_output`` (EOT position, argmax of the ids) and, - when ``projection_dim`` is set, ``text_embeds``, the pooled state through the - bias-free ``text_projection`` (the second SDXL tower). Built under a - ``keras.name_scope(name)`` so the two towers' identically named layers keep - distinct weight paths inside the container. - """ - length = config.max_seq_len - with keras.name_scope(name): - # int32 inputs: under a float16 load policy a float input would round the - # ids (49407 is not representable in float16) - token_ids = layers.Input(shape=(length,), dtype="int32", name="token_ids") - padding_mask = layers.Input(shape=(length,), dtype="int32", name="padding_mask") - x = CLIPTextModelEmbedding( - vocab_size=config.vocab_size, - max_seq_len=length, - embed_dim=config.hidden_dim, - name="text_model_embedding", - )(token_ids) - - causal_mask = ops.cast(ops.triu(ops.ones((length, length)), k=1), "float32") - causal_mask = causal_mask * (-1e8) - key_mask = ops.reshape(ops.cast(padding_mask, "float32"), (-1, 1, 1, length)) - key_mask = (1.0 - ops.repeat(key_mask, length, axis=2)) * (-1e8) - - penultimate = None - for i in range(config.num_layers): - if i == config.num_layers - 1: - penultimate = x - x = residual_attention_block( - x, - proj_dim=config.hidden_dim, - num_heads=config.num_heads, - layer_name_prefix="text_model_encoder", - layer_idx=i, - causal_attention_mask=causal_mask, - attention_mask=key_mask, - mlp_ratio=config.mlp_ratio, - hidden_act=hidden_act, - layer_norm_eps=layer_norm_eps, - ) - last_hidden_state = layers.LayerNormalization( - epsilon=layer_norm_eps, name="text_model_layernorm" - )(x) - one_hot = ops.one_hot( - ops.argmax(token_ids, axis=-1), length, dtype=last_hidden_state.dtype - ) - pooler_output = ops.einsum("bi,bij->bj", one_hot, last_hidden_state) - - outputs = { - "penultimate_hidden_state": penultimate, - "last_hidden_state": last_hidden_state, - "pooler_output": pooler_output, - } - if projection_dim is not None: - # the Dense runs on a (B, 1, D) view, as CLIPTextEmbed does, so the - # kernel maps to the checkpoint's text_projection.weight unchanged - projected = layers.Dense( - projection_dim, use_bias=False, name="text_projection" - )(ops.expand_dims(pooler_output, axis=1)) - outputs["text_embeds"] = ops.squeeze(projected, axis=1) - return keras.Model( - inputs={"token_ids": token_ids, "padding_mask": padding_mask}, - outputs=outputs, - name=name, - ) - - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusionXLModel(StableDiffusionModel): - """Stable Diffusion XL weights container: the UNet, the VAE and the two text - encoders (CLIP ViT-L/14 and OpenCLIP ViT-bigG/14) as one functional - ``keras.Model``. - - The SD 1.x container (:class:`StableDiffusionModel`) with the SDXL - configuration and a fourth component: four disconnected sub-graphs, - ``{"sample", "timestep", "encoder_hidden_states", "text_embeds", "time_ids"} - -> "noise_pred"`` (the UNet with its ``text_time`` micro-conditioning), - ``{"image", "latent"} -> "moments" / "image"`` (the VAE, built in float32 - whatever the load dtype, ``force_upcast``) and ``{"token_ids", "token_ids_2", - "padding_mask"} -> "prompt_embeds" / "pooled_prompt_embeds"`` (the two text - towers: their penultimate hidden states concatenated to the UNet's 2048-d - context, and the second tower's projected pooled state). The components are - ``.unet`` / ``.vae`` / ``.text_encoder`` / ``.text_encoder_2``. Hosted as one - ``zm_config.json`` + float16 weights (the checkpoints' native precision) per - variant under ``zeromodels/stable-diffusion-xl-base-1.0`` and - ``zeromodels/sdxl-turbo``; on-the-fly ``hf:`` conversion is not supported for - diffusion models. - - Args: - **kwargs: The flat :class:`StableDiffusionXLConfig` fields (``unet_`` / - ``vae_`` / ``text_`` / ``text_2_`` prefixed), or a config positionally. - Pass ``unet_sample_size`` / ``vae_sample_size`` to build for another - resolution (the weights are resolution-independent). - """ - - config_class = StableDiffusionXLConfig - HUB_REPO_SIBLINGS = STABLE_DIFFUSION_XL_HUB_SIBLINGS - - def __init__(self, name="StableDiffusionXLModel", **kwargs): - super().__init__(name=name, **kwargs) - - def build_components(self, config, data_format, channels_axis): - u, v, t, t2 = ( - config.unet_config, - config.vae_config, - config.text_config, - config.text_config_2, - ) - unet = UNet2DConditionModel( - u, data_format=data_format, channels_axis=channels_axis - ) - vae = AutoencoderKL( - **v.constructor_kwargs(), - data_format=data_format, - channels_axis=channels_axis, - ) - components = {"unet": unet, "vae": vae} - if t is not None: # the refiner has no CLIP ViT-L/14 tower - components["text_encoder"] = build_text_encoder( - t, config.hidden_act, config.layer_norm_eps, name="text_encoder" - ) - components["text_encoder_2"] = build_text_encoder( - t2, - t2.hidden_act, - config.layer_norm_eps, - projection_dim=t2.projection_dim, - name="text_encoder_2", - ) - return components - - def build_graph(self, config, components, data_format): - u, length = config.unet_config, config.text_config_2.max_seq_len - unet, vae = components["unet"], components["vae"] - text_encoder = components.get("text_encoder") - text_encoder_2 = components["text_encoder_2"] - inputs = self.unet_vae_inputs(config, components, data_format) - pooled_dim = ( - u.projection_class_embeddings_input_dim - - u.num_time_ids * u.addition_time_embed_dim - ) - inputs["text_embeds"] = layers.Input(shape=(pooled_dim,), name="text_embeds") - inputs["time_ids"] = layers.Input(shape=(u.num_time_ids,), name="time_ids") - if text_encoder is not None: - inputs["token_ids"] = layers.Input( - shape=(length,), dtype="int32", name="token_ids" - ) - inputs["token_ids_2"] = layers.Input( - shape=(length,), dtype="int32", name="token_ids_2" - ) - inputs["padding_mask"] = layers.Input( - shape=(length,), dtype="int32", name="padding_mask" - ) - - noise_pred = unet( - { - key: inputs[key] - for key in ( - "sample", - "timestep", - "encoder_hidden_states", - "text_embeds", - "time_ids", - ) - } - )["sample"] - vae_out = vae({"image": inputs["image"], "latent": inputs["latent"]}) - text_out_2 = text_encoder_2( - { - "token_ids": inputs["token_ids_2"], - "padding_mask": inputs["padding_mask"], - } - ) - hidden_states = [text_out_2["penultimate_hidden_state"]] - if text_encoder is not None: - text_out = text_encoder( - { - "token_ids": inputs["token_ids"], - "padding_mask": inputs["padding_mask"], - } - ) - hidden_states.insert(0, text_out["penultimate_hidden_state"]) - prompt_embeds = ( - ops.concatenate(hidden_states, axis=-1) - if len(hidden_states) > 1 - else hidden_states[0] - ) - outputs = { - "noise_pred": noise_pred, - "moments": vae_out["moments"], - "image": vae_out["sample"], - "prompt_embeds": prompt_embeds, - "pooled_prompt_embeds": text_out_2["text_embeds"], - } - return inputs, outputs - - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusionXLTextToImage(StableDiffusionTextToImage, StableDiffusionXLModel): - """Text-to-image Stable Diffusion XL, pure Keras 3 and cross-backend (the - diffusers ``StableDiffusionXLPipeline`` in this library's task-class form). - - :class:`StableDiffusionTextToImage`'s ``generate`` (``BaseDiffusion``) over the - :class:`StableDiffusionXLModel` graph. The hooks follow the SDXL recipe: the - prompt goes through both text towers (the second tower's ``!``-padded ids are - derived from the first tokenizer's output), their penultimate hidden states - are concatenated into the UNet context and the second tower's projected pooled - state becomes the ``text_embeds`` micro-conditioning, next to the ``time_ids`` - built from ``original_size`` / ``crops_coords_top_left`` / ``target_size``. - With no negative prompt the unconditional branch of classifier-free guidance - is zero embeddings (``force_zeros_for_empty_prompt``), as in the reference. - The sampler is the repo's ``scheduler_config`` (Euler with ``leading`` spacing - for the base model, ancestral Euler with ``trailing`` spacing for SDXL-Turbo):: - - sd = StableDiffusionXLTextToImage.from_weights("zeromodels/stable-diffusion-xl-base-1.0") - tokenizer = StableDiffusionXLTokenizer.from_weights("zeromodels/stable-diffusion-xl-base-1.0") - image = sd.generate(**tokenizer("a photo of a cat")) # (1, 1024, 1024, 3) uint8 - - Args: - scheduler: A :class:`BaseScheduler`; defaults to the config's - ``scheduler_config``, else Euler with ``leading`` spacing. - **kwargs: The flat :class:`StableDiffusionXLConfig` fields. - """ - - config_class = StableDiffusionXLConfig - HUB_REPO_SIBLINGS = STABLE_DIFFUSION_XL_HUB_SIBLINGS - # diffusers' StableDiffusionXLPipeline defaults; a repo's generate_args override - generate_args = {"num_inference_steps": 50, "guidance_scale": 5.0} - - def __init__(self, scheduler=None, name="StableDiffusionXLTextToImage", **kwargs): - super().__init__(scheduler=scheduler, name=name, **kwargs) - - def default_scheduler(self): - return EulerDiscreteScheduler(timestep_spacing="leading", steps_offset=1) - - def generate( - self, - input_ids, - attention_mask=None, - negative_input_ids=None, - num_inference_steps=None, - guidance_scale=None, - seed=None, - latents=None, - image=None, - strength=None, - denoising_start=None, - denoising_end=None, - output_type="image", - original_size=None, - crops_coords_top_left=(0, 0), - target_size=None, - aesthetic_score=6.0, - negative_original_size=None, - negative_crops_coords_top_left=None, - negative_target_size=None, - negative_aesthetic_score=2.5, - ): - """Generate images from tokenized prompts, or refine an ``image`` / latent - (see ``BaseDiffusion.generate`` for the shared arguments). - - The SDXL micro-conditioning arguments are ``(height, width)`` pairs in - pixels: ``original_size`` and ``target_size`` default to the generated - image's size, ``crops_coords_top_left`` to ``(0, 0)``; the ``negative_`` - variants condition the negative branch and default to the positive values. - ``aesthetic_score`` / ``negative_aesthetic_score`` (6.0 / 2.5) replace the - target size on the refiner (``requires_aesthetics_score``). For the SDXL - ensemble, run the base with ``denoising_end=0.8, output_type="latent"`` and - hand the latent to the refiner's ``generate(..., latents=..., - denoising_start=0.8)``. - """ - return super().generate( - input_ids, - attention_mask=attention_mask, - negative_input_ids=negative_input_ids, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - seed=seed, - latents=latents, - image=image, - strength=strength, - denoising_start=denoising_start, - denoising_end=denoising_end, - output_type=output_type, - original_size=original_size, - crops_coords_top_left=crops_coords_top_left, - target_size=target_size, - aesthetic_score=aesthetic_score, - negative_original_size=negative_original_size, - negative_crops_coords_top_left=negative_crops_coords_top_left, - negative_target_size=negative_target_size, - negative_aesthetic_score=negative_aesthetic_score, - ) - - def image_size(self): - size = self.vae.sample_size - return tuple(size) if isinstance(size, (tuple, list)) else (size, size) - - def time_ids( - self, - batch, - original_size=None, - crops_coords_top_left=(0, 0), - target_size=None, - aesthetic_score=6.0, - ): - """The ``(batch, 6)`` float32 micro-conditioning row: original size, crop - offset and target size, each ``(height, width)``; ``(batch, 5)`` with the - aesthetic score in place of the target size when - ``requires_aesthetics_score`` (the refiner).""" - original_size = tuple(original_size or self.image_size()) - row = list(original_size) + list(crops_coords_top_left) - if self.config.requires_aesthetics_score: - row = row + [aesthetic_score] - else: - row = row + list(target_size or self.image_size()) - return ops.convert_to_tensor([row] * batch, dtype="float32") - - def unconditional_ids(self, batch): - cfg = self.config - length = cfg.text_config_2.max_seq_len - row = [cfg.bos_token_id, cfg.eos_token_id] + [cfg.pad_token_id] * (length - 2) - return ops.convert_to_tensor([row] * batch, dtype="int32") - - def prompt_mask(self, input_ids): - # the tokens up to and including the first <|endoftext|>; what follows is - # padding (the first tokenizer pads with <|endoftext|> itself) - is_eos = ops.cast(ops.equal(input_ids, self.config.eos_token_id), "int32") - return ops.equal(ops.cumsum(is_eos, axis=1) - is_eos, 0) - - def encode_prompt( - self, - input_ids, - attention_mask=None, - original_size=None, - crops_coords_top_left=(0, 0), - target_size=None, - aesthetic_score=6.0, - ): - input_ids = ops.cast(ops.convert_to_tensor(input_ids), "int32") - if attention_mask is None: - real = self.prompt_mask(input_ids) - else: - real = ops.greater( - ops.cast(ops.convert_to_tensor(attention_mask), "int32"), 0 - ) - # the second tokenizer pads with "!" (id 0) where the first pads with - # <|endoftext|>; the towers attend the padding (no attention mask), as the - # reference does - input_ids_2 = ops.where(real, input_ids, self.config.pad_token_id_2) - ones = ops.ones_like(input_ids) - out_2 = self.text_encoder_2({"token_ids": input_ids_2, "padding_mask": ones}) - hidden = out_2["penultimate_hidden_state"] - if self.config.text_config is not None: - out = self.text_encoder({"token_ids": input_ids, "padding_mask": ones}) - hidden = ops.concatenate([out["penultimate_hidden_state"], hidden], -1) - batch = int(input_ids.shape[0]) - return { - "encoder_hidden_states": hidden, - "text_embeds": out_2["text_embeds"], - "time_ids": self.time_ids( - batch, - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - ), - } - - def encode_negative_prompt(self, negative_input_ids, batch, **conditioning): - if ( - negative_input_ids is not None - or not self.config.force_zeros_for_empty_prompt - ): - return super().encode_negative_prompt( - negative_input_ids, batch, **conditioning - ) - # force_zeros_for_empty_prompt: zero embeddings stand in for the empty - # prompt; the micro-conditioning keeps its (negative) values - u, length = self.config.unet_config, self.config.text_config_2.max_seq_len - pooled_dim = ( - u.projection_class_embeddings_input_dim - - u.num_time_ids * u.addition_time_embed_dim - ) - return { - "encoder_hidden_states": ops.zeros( - (batch, length, u.cross_attention_dim), dtype="float32" - ), - "text_embeds": ops.zeros((batch, pooled_dim), dtype="float32"), - "time_ids": self.time_ids(batch, **conditioning), - } - - def predict_noise(self, latents, timesteps, embeddings): - return self.unet({"sample": latents, "timestep": timesteps, **embeddings})[ - "sample" - ] - - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusionXLRefinerModel(StableDiffusionXLModel): - """Stable Diffusion XL refiner weights container: the refiner UNet, the VAE and - the OpenCLIP ViT-bigG/14 text encoder as one functional ``keras.Model``. - - :class:`StableDiffusionXLModel` with the refiner configuration - (:class:`StableDiffusionXLRefinerConfig`): three disconnected sub-graphs (no - CLIP ViT-L/14 tower, so the UNet's context is the second tower's 1280-d - penultimate state alone), the four-level refiner UNet and five - micro-conditioning ids (size, crop, aesthetic score). Components: - ``.unet`` / ``.vae`` / ``.text_encoder_2``. Hosted under - ``zeromodels/stable-diffusion-xl-refiner-1.0`` (float16, the VAE float32). - - Args: - **kwargs: The flat :class:`StableDiffusionXLRefinerConfig` fields, or a - config positionally. - """ - - config_class = StableDiffusionXLRefinerConfig - HUB_REPO_SIBLINGS = STABLE_DIFFUSION_XL_REFINER_HUB_SIBLINGS - - def __init__(self, name="StableDiffusionXLRefinerModel", **kwargs): - super().__init__(name=name, **kwargs) - - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusionXLRefinerImageToImage( - StableDiffusionXLTextToImage, StableDiffusionXLRefinerModel -): - """Image-to-image Stable Diffusion XL refiner, pure Keras 3 and cross-backend - (diffusers' ``StableDiffusionXLImg2ImgPipeline`` with the refiner weights). - - :class:`StableDiffusionXLTextToImage`'s hooks over the - :class:`StableDiffusionXLRefinerModel` graph, the prompt encoded by the - OpenCLIP ViT-bigG/14 tower alone and the aesthetic score - (``aesthetic_score`` 6.0 / ``negative_aesthetic_score`` 2.5) as the fifth - micro-conditioning id. ``generate`` refines an ``image`` (noised to - ``strength``, 0.3 by default) or, as the second half of the SDXL ensemble of - experts, a latent the base model left partially denoised - (``denoising_start`` at the base's ``denoising_end``):: - - base = StableDiffusionXLTextToImage.from_weights("zeromodels/stable-diffusion-xl-base-1.0") - refiner = StableDiffusionXLRefinerImageToImage.from_weights("zeromodels/stable-diffusion-xl-refiner-1.0") - tokenizer = StableDiffusionXLTokenizer.from_weights("zeromodels/stable-diffusion-xl-base-1.0") - inputs = tokenizer("a photo of a cat") - latent = base.generate(**inputs, denoising_end=0.8, output_type="latent") - image = refiner.generate(**inputs, latents=latent, denoising_start=0.8) - - Args: - scheduler: A :class:`BaseScheduler`; defaults to the config's - ``scheduler_config``, else Euler with ``leading`` spacing. - **kwargs: The flat :class:`StableDiffusionXLRefinerConfig` fields. - """ - - config_class = StableDiffusionXLRefinerConfig - HUB_REPO_SIBLINGS = STABLE_DIFFUSION_XL_REFINER_HUB_SIBLINGS - # diffusers' StableDiffusionXLImg2ImgPipeline defaults - generate_args = {"num_inference_steps": 50, "guidance_scale": 5.0, "strength": 0.3} - - def __init__( - self, scheduler=None, name="StableDiffusionXLRefinerImageToImage", **kwargs - ): - super().__init__(scheduler=scheduler, name=name, **kwargs) diff --git a/zeromodels/models/stable_diffusion_xl/stable_diffusion_xl_tokenizer.py b/zeromodels/models/stable_diffusion_xl/stable_diffusion_xl_tokenizer.py deleted file mode 100644 index 88046a65..00000000 --- a/zeromodels/models/stable_diffusion_xl/stable_diffusion_xl_tokenizer.py +++ /dev/null @@ -1,42 +0,0 @@ -import keras - -from zeromodels.models.stable_diffusion.stable_diffusion_tokenizer import ( - StableDiffusionTokenizer, -) - - -@keras.saving.register_keras_serializable(package="zeromodels") -class StableDiffusionXLTokenizer(StableDiffusionTokenizer): - """Stable Diffusion XL text tokenizer: the CLIP ViT-L/14 BPE tokenizer. - - SDXL tokenizes the prompt twice, once per text encoder, with the same - byte-pair encoding and ``<|startoftext|>`` / ``<|endoftext|>`` framing: the - two tokenizers differ only in the pad token (``<|endoftext|>`` for the CLIP - ViT-L/14 tower, ``!`` for the OpenCLIP ViT-bigG/14 one). This is the first; - :class:`StableDiffusionXLTextToImage` derives the second tower's ids from the - returned ``attention_mask``, so one call feeds ``generate``. Loads by repo id - like the model: ``from_weights("zeromodels/stable-diffusion-xl-base-1.0")``. - - Args: - hf_id: Hosted repo to read ``tokenizer.json`` from. Required unless - ``tokenizer_file`` is given (no default repo). - tokenizer_file: Explicit ``tokenizer.json`` path (overrides ``hf_id``). - max_seq_len: Padded / truncated length (default 77). - pad_token: Pad token string (``<|endoftext|>``). - """ - - def __init__( - self, - hf_id=None, - tokenizer_file=None, - max_seq_len=77, - pad_token="<|endoftext|>", - **kwargs, - ): - super().__init__( - hf_id=hf_id, - tokenizer_file=tokenizer_file, - max_seq_len=max_seq_len, - pad_token=pad_token, - **kwargs, - ) diff --git a/zeromodels/sentence_transformers/__init__.py b/zeromodels/sentence_transformers/__init__.py index d55870ec..9fa98f2b 100644 --- a/zeromodels/sentence_transformers/__init__.py +++ b/zeromodels/sentence_transformers/__init__.py @@ -1,37 +1,7 @@ -"""Sentence Transformer models and tokenizers for the ZeroModels embedding API.""" - -import importlib -import sys - -from .sentence_transformer_config import SentenceTransformerConfig -from .sentence_transformer_model import ( - SentenceTransformerTextEmbed, - ZMSentenceTransformer, -) +from .sentence_transformer_model import ZMSentenceTransformer from .sentence_transformer_tokenizer import ZMSentenceTransformerTokenizer -SentenceTransformer = ZMSentenceTransformer -SentenceTransformerModel = ZMSentenceTransformer -SentenceTransformerTokenizer = ZMSentenceTransformerTokenizer - -# Keep the former nested import paths usable without a nested implementation. -sentence_transformer = sys.modules[__name__] -sys.modules[f"{__name__}.sentence_transformer"] = sentence_transformer -for _name in ( - "convert_sentence_transformer_hf_to_keras", - "sentence_transformer_backbones", - "sentence_transformer_config", - "sentence_transformer_layers", - "sentence_transformer_model", - "sentence_transformer_tokenizer", -): - sys.modules[f"{__name__}.sentence_transformer.{_name}"] = importlib.import_module( - f"{__name__}.{_name}" - ) - __all__ = [ - "SentenceTransformerConfig", - "SentenceTransformerTextEmbed", "ZMSentenceTransformer", "ZMSentenceTransformerTokenizer", ] diff --git a/zeromodels/sentence_transformers/convert_sentence_transformer_hf_to_keras.py b/zeromodels/sentence_transformers/convert_sentence_transformer_hf_to_keras.py index 4a41e88f..e5b82600 100644 --- a/zeromodels/sentence_transformers/convert_sentence_transformer_hf_to_keras.py +++ b/zeromodels/sentence_transformers/convert_sentence_transformer_hf_to_keras.py @@ -76,44 +76,56 @@ def transfer_sentence_transformer_weights( break hf[normalize_hf_key(key)] = value kind = keras_model.backbone_type - if kind == "mpnet": - transfer_mpnet_weights(keras_model, hf) - elif kind == "bert": - transfer_bert_weights(keras_model, hf) - elif kind == "roberta": - transfer_roberta_weights(keras_model, hf) - elif kind == "xlm-roberta": - transfer_xlm_roberta_weights(keras_model, hf) - elif kind in ("distilbert", "albert"): - prefix = f"{kind}." - names = {key.removeprefix(prefix): key for key in hf} - mapping = WEIGHT_NAME_MAPPING[kind] - for weight in keras_model.weights: - if "/" not in weight.path: - continue - layer_path, suffix = weight.path.rsplit("/", 1) - index = None - if kind == "distilbert" and layer_path.startswith("distil_layer_"): - block, component = layer_path.split("/", 1) - index = int(block.removeprefix("distil_layer_")) - layer_path = f"distil_layer_{{index}}/{component}" - hf_name = mapping.get(layer_path) - if hf_name is None: - if layer_path.startswith(("distil_", "albert_")): - raise WeightMappingError(weight.path, "WEIGHT_NAME_MAPPING entry") - continue - suffix = WEIGHT_SUFFIX_MAPPING.get(suffix, suffix) - hf_name = f"{hf_name.format(index=index)}.{suffix}" - if hf_name not in names: - raise WeightMappingError(weight.path, hf_name) - source = hf[names[hf_name]] - if weight.path == "albert_embedding_projection/kernel": - weight.assign(np.asarray(source).T) - else: - transfer_weights(weight.path, weight, source) - else: + if kind not in BACKBONE_TRANSFERS: raise ValueError(f"Unsupported sentence embedding backbone: {kind}") + BACKBONE_TRANSFERS[kind](keras_model, hf) + transfer_projection(keras_model, hf_state_dict) + +def transfer_local_backbone(keras_model, hf): + """Copy a DistilBERT / ALBERT encoder built by sentence_transformer_backbones.""" + kind = keras_model.backbone_type + prefix = f"{kind}." + names = {key.removeprefix(prefix): key for key in hf} + mapping = WEIGHT_NAME_MAPPING[kind] + for weight in keras_model.weights: + if "/" not in weight.path: + continue + layer_path, suffix = weight.path.rsplit("/", 1) + index = None + if kind == "distilbert" and layer_path.startswith("distil_layer_"): + block, component = layer_path.split("/", 1) + index = int(block.removeprefix("distil_layer_")) + layer_path = f"distil_layer_{{index}}/{component}" + hf_name = mapping.get(layer_path) + if hf_name is None: + if layer_path.startswith(("distil_", "albert_")): + raise WeightMappingError(weight.path, "WEIGHT_NAME_MAPPING entry") + continue + suffix = WEIGHT_SUFFIX_MAPPING.get(suffix, suffix) + hf_name = f"{hf_name.format(index=index)}.{suffix}" + if hf_name not in names: + raise WeightMappingError(weight.path, hf_name) + source = hf[names[hf_name]] + if weight.path == "albert_embedding_projection/kernel": + weight.assign(np.asarray(source).T) + else: + transfer_weights(weight.path, weight, source) + + +# backbone_type -> encoder weight transfer (mirrors SENTENCE_TRANSFORMER_BACKBONES) +BACKBONE_TRANSFERS = { + "mpnet": transfer_mpnet_weights, + "bert": transfer_bert_weights, + "roberta": transfer_roberta_weights, + "xlm-roberta": transfer_xlm_roberta_weights, + "distilbert": transfer_local_backbone, + "albert": transfer_local_backbone, +} + + +def transfer_projection(keras_model, hf_state_dict): + """Copy the optional Dense projection (a separate ST module) when present.""" needs_proj = any(w.path.startswith("projection/") for w in keras_model.weights) if not needs_proj: return diff --git a/zeromodels/sentence_transformers/sentence_transformer_backbones.py b/zeromodels/sentence_transformers/sentence_transformer_backbones.py index 6bd10e5e..532670c9 100644 --- a/zeromodels/sentence_transformers/sentence_transformer_backbones.py +++ b/zeromodels/sentence_transformers/sentence_transformer_backbones.py @@ -1,8 +1,13 @@ import math +from typing import Callable, NamedTuple import keras from keras import layers, ops +from zeromodels.models.bert.bert_model import bert_backbone +from zeromodels.models.mpnet.mpnet_model import mpnet_backbone +from zeromodels.models.roberta.roberta_model import roberta_backbone + def _attention(query, key, value, attention_mask, num_heads): batch = ops.shape(query)[0] @@ -223,3 +228,113 @@ def albert_backbone( for _ in range(num_layers): x = shared([x, attention_mask]) return x + + +def encoder_kwargs(config): + """The encoder hyperparameters every backbone builder takes.""" + return { + "vocab_size": config.vocab_size, + "embed_dim": config.embed_dim, + "num_layers": config.num_layers, + "num_heads": config.num_heads, + "mlp_dim": config.mlp_dim, + "max_position_embeddings": config.max_position_embeddings, + "hidden_act": config.hidden_act, + "layer_norm_eps": config.layer_norm_eps, + } + + +def build_mpnet(inputs, config): + sequence_output, _ = mpnet_backbone( + inputs["input_ids"], + inputs["attention_mask"], + **encoder_kwargs(config), + relative_attention_num_buckets=config.relative_attention_num_buckets, + pad_token_id=config.pad_token_id, + dropout=config.dropout, + attention_dropout=config.attention_dropout, + add_pooler=False, + ) + return sequence_output + + +def build_bert(inputs, config): + sequence_output, _ = bert_backbone( + inputs["input_ids"], + inputs["attention_mask"], + inputs["token_type_ids"], + **encoder_kwargs(config), + type_vocab_size=config.type_vocab_size, + dropout=config.dropout, + attention_dropout=config.attention_dropout, + add_pooler=False, + ) + return sequence_output + + +def build_roberta(inputs, config): + sequence_output, _ = roberta_backbone( + inputs["input_ids"], + inputs["attention_mask"], + inputs["token_type_ids"], + **encoder_kwargs(config), + type_vocab_size=config.type_vocab_size, + pad_token_id=config.pad_token_id, + dropout=config.dropout, + attention_dropout=config.attention_dropout, + add_pooler=False, + ) + return sequence_output + + +def build_distilbert(inputs, config): + return distilbert_backbone( + inputs["input_ids"], inputs["attention_mask"], **encoder_kwargs(config) + ) + + +def build_albert(inputs, config): + return albert_backbone( + inputs["input_ids"], + inputs["attention_mask"], + inputs["token_type_ids"], + **encoder_kwargs(config), + embedding_size=config.embedding_size or config.embed_dim, + type_vocab_size=config.type_vocab_size, + num_hidden_groups=config.num_hidden_groups, + inner_group_num=config.inner_group_num, + ) + + +class SentenceTransformerBackbone(NamedTuple): + """One encoder ``ZMSentenceTransformer`` can be built on. + + ``build(inputs, config)`` wires the encoder into the functional graph and + returns the token embeddings; ``token_type_ids`` says whether the encoder (and + so the model and the tokenizer) takes segment ids. + """ + + build: Callable + token_type_ids: bool + + +# backbone_type -> encoder. MPNet / BERT / RoBERTa / XLM-R reuse the zeromodels.models +# encoders; DistilBERT and ALBERT are built above. +SENTENCE_TRANSFORMER_BACKBONES = { + "mpnet": SentenceTransformerBackbone(build_mpnet, token_type_ids=False), + "bert": SentenceTransformerBackbone(build_bert, token_type_ids=True), + "roberta": SentenceTransformerBackbone(build_roberta, token_type_ids=True), + "xlm-roberta": SentenceTransformerBackbone(build_roberta, token_type_ids=True), + "distilbert": SentenceTransformerBackbone(build_distilbert, token_type_ids=False), + "albert": SentenceTransformerBackbone(build_albert, token_type_ids=True), +} + + +def get_backbone(backbone_type): + backbone = SENTENCE_TRANSFORMER_BACKBONES.get(backbone_type) + if backbone is None: + raise ValueError( + f"Unsupported sentence embedding backbone: {backbone_type!r}; expected " + f"one of {sorted(SENTENCE_TRANSFORMER_BACKBONES)}." + ) + return backbone diff --git a/zeromodels/sentence_transformers/sentence_transformer_config.py b/zeromodels/sentence_transformers/sentence_transformer_config.py index 89be4331..53177ca9 100644 --- a/zeromodels/sentence_transformers/sentence_transformer_config.py +++ b/zeromodels/sentence_transformers/sentence_transformer_config.py @@ -39,3 +39,7 @@ class SentenceTransformerConfig(BaseConfig): projection_activation: str = "linear" similarity_fn_name: str = "cosine" max_seq_len: int = 384 + + +# resolvable by the Auto loaders (AutoZMConfig); not part of the package API +__all__ = ["SentenceTransformerConfig"] diff --git a/zeromodels/sentence_transformers/sentence_transformer_model.py b/zeromodels/sentence_transformers/sentence_transformer_model.py index 66998dd9..e63977a8 100644 --- a/zeromodels/sentence_transformers/sentence_transformer_model.py +++ b/zeromodels/sentence_transformers/sentence_transformer_model.py @@ -14,15 +14,22 @@ from zeromodels.conversion.hf_download_utils import download_hf_state_dict from zeromodels.conversion.weight_transfer_util import skip_mismatched_weights -from .sentence_transformer_backbones import albert_backbone, distilbert_backbone +from .sentence_transformer_backbones import SENTENCE_TRANSFORMER_BACKBONES, get_backbone from .sentence_transformer_config import SentenceTransformerConfig from .sentence_transformer_layers import ( SentenceTransformerNormalize, SentenceTransformerPooling, ) +# Class names hosted zm_config.json files declared before ZMSentenceTransformer; their +# repos load unchanged. SENTENCE_TRANSFORMER_HUB_SIBLINGS = frozenset( - {"ZMSentenceTransformer", "SentenceTransformerModel", "SentenceTransformer"} + { + "ZMSentenceTransformer", + "SentenceTransformerModel", + "SentenceTransformer", + "SentenceTransformerTextEmbed", + } ) @@ -30,6 +37,9 @@ @keras.saving.register_keras_serializable( package="zeromodels", name="SentenceTransformerModel" ) +@keras.saving.register_keras_serializable( + package="zeromodels", name="SentenceTransformerTextEmbed" +) class ZMSentenceTransformer(BaseModel): """Sentence embedding model: encoder + pooling (+ optional Dense) + normalize. @@ -64,15 +74,7 @@ class ZMSentenceTransformer(BaseModel): # whose root ``config.json`` is the backbone transformers config # (``model_type: mpnet`` for all-mpnet-base-v2). Pooling / normalize come # from sibling ST files read in :meth:`from_hf`. - HF_MODEL_TYPE = ( - "sentence_transformer", - "mpnet", - "bert", - "roberta", - "xlm-roberta", - "distilbert", - "albert", - ) + HF_MODEL_TYPE = ("sentence_transformer", *SENTENCE_TRANSFORMER_BACKBONES) HUB_REPO_SIBLINGS = SENTENCE_TRANSFORMER_HUB_SIBLINGS def __init__( @@ -106,17 +108,7 @@ def __init__( for k in ("model", "hf_id", "url", "num_classes"): kwargs.pop(k, None) - if backbone_type not in ( - "mpnet", - "bert", - "roberta", - "xlm-roberta", - "distilbert", - "albert", - ): - raise ValueError( - f"Unsupported sentence embedding backbone: {backbone_type!r}" - ) + backbone = get_backbone(backbone_type) if pooling not in ("mean", "cls", "max"): raise ValueError(f"pooling must be mean|cls|max, got {pooling!r}") if projection_activation not in ("linear", "tanh"): @@ -126,83 +118,43 @@ def __init__( if similarity_fn_name not in ("cosine", "dot"): raise ValueError(f"Unsupported similarity function: {similarity_fn_name!r}") + config = SentenceTransformerConfig( + backbone_type=backbone_type, + vocab_size=vocab_size, + embed_dim=embed_dim, + num_layers=num_layers, + num_heads=num_heads, + mlp_dim=mlp_dim, + max_position_embeddings=max_position_embeddings, + relative_attention_num_buckets=relative_attention_num_buckets, + hidden_act=hidden_act, + layer_norm_eps=layer_norm_eps, + pad_token_id=pad_token_id, + type_vocab_size=type_vocab_size, + embedding_size=embedding_size, + num_hidden_groups=num_hidden_groups, + inner_group_num=inner_group_num, + dropout=dropout, + attention_dropout=attention_dropout, + pooling=pooling, + normalize_embeddings=normalize_embeddings, + projection_dim=projection_dim, + projection_activation=projection_activation, + similarity_fn_name=similarity_fn_name, + max_seq_len=max_seq_len, + ) + inputs = { "input_ids": layers.Input(shape=(None,), dtype="int32", name="input_ids"), "attention_mask": layers.Input( shape=(None,), dtype="int32", name="attention_mask" ), } - if backbone_type in ("bert", "roberta", "xlm-roberta", "albert"): + if backbone.token_type_ids: inputs["token_type_ids"] = layers.Input( shape=(None,), dtype="int32", name="token_type_ids" ) - common = { - "vocab_size": vocab_size, - "embed_dim": embed_dim, - "num_layers": num_layers, - "num_heads": num_heads, - "mlp_dim": mlp_dim, - "max_position_embeddings": max_position_embeddings, - "hidden_act": hidden_act, - "layer_norm_eps": layer_norm_eps, - } - if backbone_type == "mpnet": - from zeromodels.models.mpnet.mpnet_model import mpnet_backbone - - sequence_output, _ = mpnet_backbone( - inputs["input_ids"], - inputs["attention_mask"], - **common, - relative_attention_num_buckets=relative_attention_num_buckets, - pad_token_id=pad_token_id, - dropout=dropout, - attention_dropout=attention_dropout, - add_pooler=False, - ) - elif backbone_type in ("roberta", "xlm-roberta"): - from zeromodels.models.roberta.roberta_model import roberta_backbone - - sequence_output, _ = roberta_backbone( - inputs["input_ids"], - inputs["attention_mask"], - inputs["token_type_ids"], - **common, - type_vocab_size=type_vocab_size, - pad_token_id=pad_token_id, - dropout=dropout, - attention_dropout=attention_dropout, - add_pooler=False, - ) - elif backbone_type == "bert": - from zeromodels.models.bert.bert_model import bert_backbone - - sequence_output, _ = bert_backbone( - inputs["input_ids"], - inputs["attention_mask"], - inputs["token_type_ids"], - **common, - type_vocab_size=type_vocab_size, - dropout=dropout, - attention_dropout=attention_dropout, - add_pooler=False, - ) - elif backbone_type == "distilbert": - sequence_output = distilbert_backbone( - inputs["input_ids"], - inputs["attention_mask"], - **common, - ) - else: - sequence_output = albert_backbone( - inputs["input_ids"], - inputs["attention_mask"], - inputs["token_type_ids"], - **common, - embedding_size=embedding_size or embed_dim, - type_vocab_size=type_vocab_size, - num_hidden_groups=num_hidden_groups, - inner_group_num=inner_group_num, - ) + sequence_output = backbone.build(inputs, config) pooled = SentenceTransformerPooling(pooling=pooling, name="pooling")( [sequence_output, inputs["attention_mask"]] @@ -254,31 +206,7 @@ def __init__( self.similarity_fn_name = similarity_fn_name self.max_seq_len = max_seq_len self.output_embedding_dim = out_dim - self.config = SentenceTransformerConfig( - backbone_type=backbone_type, - vocab_size=vocab_size, - embed_dim=embed_dim, - num_layers=num_layers, - num_heads=num_heads, - mlp_dim=mlp_dim, - max_position_embeddings=max_position_embeddings, - relative_attention_num_buckets=relative_attention_num_buckets, - hidden_act=hidden_act, - layer_norm_eps=layer_norm_eps, - pad_token_id=pad_token_id, - type_vocab_size=type_vocab_size, - embedding_size=embedding_size, - num_hidden_groups=num_hidden_groups, - inner_group_num=inner_group_num, - dropout=dropout, - attention_dropout=attention_dropout, - pooling=pooling, - normalize_embeddings=normalize_embeddings, - projection_dim=projection_dim, - projection_activation=projection_activation, - similarity_fn_name=similarity_fn_name, - max_seq_len=max_seq_len, - ) + self.config = config @classmethod def config_from_hf(cls, hf_config): @@ -571,13 +499,3 @@ def get_config(self): config = super().get_config() config.update(self.config.constructor_kwargs()) return config - - -@keras.saving.register_keras_serializable(package="zeromodels") -class SentenceTransformerTextEmbed(ZMSentenceTransformer): - """Text-embedding task view of :class:`ZMSentenceTransformer`.""" - - -# Preserve old deep imports and serialized objects without duplicating classes. -SentenceTransformerModel = ZMSentenceTransformer -SentenceTransformer = ZMSentenceTransformer diff --git a/zeromodels/sentence_transformers/sentence_transformer_tokenizer.py b/zeromodels/sentence_transformers/sentence_transformer_tokenizer.py index de792fc3..85a92dfd 100644 --- a/zeromodels/sentence_transformers/sentence_transformer_tokenizer.py +++ b/zeromodels/sentence_transformers/sentence_transformer_tokenizer.py @@ -9,6 +9,8 @@ from zeromodels.base import BaseTokenizer +from .sentence_transformer_backbones import get_backbone + @keras.saving.register_keras_serializable(package="zeromodels") @keras.saving.register_keras_serializable( @@ -249,8 +251,7 @@ def call(self, inputs: Union[str, List[str]], text_pair=None): return self.encode_batch_to_inputs( inputs, text_pair, - token_type_ids=self.backbone_type - in ("bert", "roberta", "xlm-roberta", "albert"), + token_type_ids=get_backbone(self.backbone_type).token_type_ids, ) def get_config(self): @@ -344,6 +345,3 @@ def _has_lowercase(config): for item in config.get("normalizers", []) ) ) - - -SentenceTransformerTokenizer = ZMSentenceTransformerTokenizer diff --git a/zeromodels/st_models/__init__.py b/zeromodels/st_models/__init__.py deleted file mode 100644 index 90b6eba0..00000000 --- a/zeromodels/st_models/__init__.py +++ /dev/null @@ -1,40 +0,0 @@ -"""Compatibility aliases for the renamed ``zeromodels.sentence_transformers`` API.""" - -import importlib -import sys - -from zeromodels.sentence_transformers import ( - SentenceTransformer, - SentenceTransformerConfig, - SentenceTransformerModel, - SentenceTransformerTextEmbed, - SentenceTransformerTokenizer, - ZMSentenceTransformer, - ZMSentenceTransformerTokenizer, - sentence_transformer, -) - -# Saved configs and deep imports must resolve to the canonical classes, not -# load a second copy with a different module name or Keras registration. -sys.modules[f"{__name__}.sentence_transformer"] = sentence_transformer -for _name in ( - "convert_sentence_transformer_hf_to_keras", - "sentence_transformer_backbones", - "sentence_transformer_config", - "sentence_transformer_layers", - "sentence_transformer_model", - "sentence_transformer_tokenizer", -): - sys.modules[f"{__name__}.sentence_transformer.{_name}"] = importlib.import_module( - f"{sentence_transformer.__name__}.{_name}" - ) - -__all__ = [ - "SentenceTransformer", - "SentenceTransformerConfig", - "SentenceTransformerModel", - "SentenceTransformerTextEmbed", - "SentenceTransformerTokenizer", - "ZMSentenceTransformer", - "ZMSentenceTransformerTokenizer", -] From dc0e0aaecf2472bd66799e2eff9276cf2bba2ae1 Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Sun, 18 Oct 2026 16:12:28 -0700 Subject: [PATCH 2/4] fix-stmodels --- docs/qwen_image.md | 4 +++- docs/qwen_image_21.md | 4 +++- docs/stable_diffusion.md | 4 +++- docs/stable_diffusion_2.md | 4 +--- docs/stable_diffusion_3.md | 4 +--- docs/stable_diffusion_3_5.md | 8 ++----- docs/stable_diffusion_xl.md | 24 +++++-------------- .../integration/test_sentence_transformer.py | 1 - zeromodels/__init__.py | 2 -- 9 files changed, 19 insertions(+), 36 deletions(-) diff --git a/docs/qwen_image.md b/docs/qwen_image.md index a324bfd2..feca55ce 100644 --- a/docs/qwen_image.md +++ b/docs/qwen_image.md @@ -198,7 +198,9 @@ a list of strings. `generate` / `encode_prompt` drop the template prefix `max_sequence_length` (512) tokens. ```python -ZMDiffusionTokenizer(model_type="qwen_image", hf_id=None, tokenizer_file=None, max_seq_len=1024) +ZMDiffusionTokenizer( + model_type="qwen_image", hf_id=None, tokenizer_file=None, max_seq_len=1024 +) ``` There is no image processor and no processor class: text-to-image takes token ids diff --git a/docs/qwen_image_21.md b/docs/qwen_image_21.md index 86dcb241..88e95de6 100644 --- a/docs/qwen_image_21.md +++ b/docs/qwen_image_21.md @@ -181,7 +181,9 @@ drops `prompt_template_encode_start_idx` (14) tokens and keeps at most `max_sequence_length` (512). ```python -ZMDiffusionTokenizer(model_type="qwen_image_21", hf_id=None, tokenizer_file=None, max_seq_len=1024) +ZMDiffusionTokenizer( + model_type="qwen_image_21", hf_id=None, tokenizer_file=None, max_seq_len=1024 +) ``` ## End-to-end example diff --git a/docs/stable_diffusion.md b/docs/stable_diffusion.md index c56a5709..86e221aa 100644 --- a/docs/stable_diffusion.md +++ b/docs/stable_diffusion.md @@ -198,7 +198,9 @@ is built for. Calling it returns `{"input_ids", "attention_mask"}` for a string of strings. `from_weights("zeromodels/")` reads the family from the repo. ```python -ZMDiffusionTokenizer(model_type="stable_diffusion", hf_id=None, tokenizer_file=None, max_seq_len=77) +ZMDiffusionTokenizer( + model_type="stable_diffusion", hf_id=None, tokenizer_file=None, max_seq_len=77 +) ``` There is no image processor and no processor class: text-to-image takes token ids in diff --git a/docs/stable_diffusion_2.md b/docs/stable_diffusion_2.md index 84a2072c..5f25701d 100644 --- a/docs/stable_diffusion_2.md +++ b/docs/stable_diffusion_2.md @@ -124,9 +124,7 @@ from PIL import Image from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-2-1-base") -tokenizer = ZMDiffusionTokenizer.from_weights( - "zeromodels/stable-diffusion-2-1-base" -) +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-2-1-base") inputs = tokenizer( "a steaming bowl of ramen on a wooden table, food photography, shallow depth of field" diff --git a/docs/stable_diffusion_3.md b/docs/stable_diffusion_3.md index 676698b0..26e67feb 100644 --- a/docs/stable_diffusion_3.md +++ b/docs/stable_diffusion_3.md @@ -181,9 +181,7 @@ model = ZMDiffusion.from_weights( "zeromodels/stable-diffusion-3-medium", text_encoder_3="zeromodels/t5-v1_1-xxl-encoder", # omit on small machines: T5 features zeroed ) -tokenizer = ZMDiffusionTokenizer.from_weights( - "zeromodels/stable-diffusion-3-medium" -) +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-3-medium") inputs = tokenizer( "a fluffy corgi wearing round sunglasses, sitting on a surfboard at a sunny beach, photo" diff --git a/docs/stable_diffusion_3_5.md b/docs/stable_diffusion_3_5.md index b2f2958b..dc246c7b 100644 --- a/docs/stable_diffusion_3_5.md +++ b/docs/stable_diffusion_3_5.md @@ -71,9 +71,7 @@ model = ZMDiffusion.from_weights( "zeromodels/stable-diffusion-3.5-medium", text_encoder_3="zeromodels/t5-v1_1-xxl-encoder", # optional ) -tokenizer = ZMDiffusionTokenizer.from_weights( - "zeromodels/stable-diffusion-3.5-medium" -) +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-3.5-medium") inputs = tokenizer( "a whimsical treehouse village at twilight, lanterns glowing, watercolor illustration" @@ -88,9 +86,7 @@ Image.fromarray(images[0]).save("treehouse.png") # (1024, 1024, 3) uint8 ### Large Turbo ```python -model = ZMDiffusion.from_weights( - "zeromodels/stable-diffusion-3.5-large-turbo" -) +model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-3.5-large-turbo") images = model.generate( **tokenizer("a red panda reading a book") ) # 4 steps, guidance 0 (the repo defaults) diff --git a/docs/stable_diffusion_xl.md b/docs/stable_diffusion_xl.md index ebb4f440..4d31a17b 100644 --- a/docs/stable_diffusion_xl.md +++ b/docs/stable_diffusion_xl.md @@ -156,15 +156,9 @@ The same class with `model_type="stable_diffusion_xl_refiner"` (`text_encoder_2` refiner was not trained for. The two intended uses: ```python -base = ZMDiffusion.from_weights( - "zeromodels/stable-diffusion-xl-base-1.0" -) -refiner = ZMDiffusion.from_weights( - "zeromodels/stable-diffusion-xl-refiner-1.0" -) -tokenizer = ZMDiffusionTokenizer.from_weights( - "zeromodels/stable-diffusion-xl-base-1.0" -) +base = ZMDiffusion.from_weights("zeromodels/stable-diffusion-xl-base-1.0") +refiner = ZMDiffusion.from_weights("zeromodels/stable-diffusion-xl-refiner-1.0") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-xl-base-1.0") inputs = tokenizer("a lighthouse on a rocky coast at dawn, dramatic clouds, cinematic") # ensemble of experts: the base denoises 80% of the schedule, the refiner the rest @@ -227,12 +221,8 @@ os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" from PIL import Image from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer -model = ZMDiffusion.from_weights( - "zeromodels/stable-diffusion-xl-base-1.0" -) -tokenizer = ZMDiffusionTokenizer.from_weights( - "zeromodels/stable-diffusion-xl-base-1.0" -) +model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-xl-base-1.0") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-xl-base-1.0") inputs = tokenizer( "a lighthouse on a rocky coast at dawn, dramatic clouds, cinematic, highly detailed" @@ -250,9 +240,7 @@ The SDXL "ensemble of experts": the base runs the first 80% of the schedule and latent to the refiner, which finishes it. Same prompt and initial latent as above: ```python -refiner = ZMDiffusion.from_weights( - "zeromodels/stable-diffusion-xl-refiner-1.0" -) +refiner = ZMDiffusion.from_weights("zeromodels/stable-diffusion-xl-refiner-1.0") latent = model.generate( **inputs, num_inference_steps=30, denoising_end=0.8, output_type="latent", seed=3 ) diff --git a/tests/integration/test_sentence_transformer.py b/tests/integration/test_sentence_transformer.py index a095a635..2df1c673 100644 --- a/tests/integration/test_sentence_transformer.py +++ b/tests/integration/test_sentence_transformer.py @@ -44,7 +44,6 @@ def test_backbone_registry_matches_accepted_hf_model_types(): "module", ( "zeromodels.models.sentence_transformer", - "zeromodels.st_models.sentence_transformer", "zeromodels.sentence_transformers.sentence_transformer", "zeromodels.sentence_transformers", ), diff --git a/zeromodels/__init__.py b/zeromodels/__init__.py index cf3de14d..7e0bfa61 100644 --- a/zeromodels/__init__.py +++ b/zeromodels/__init__.py @@ -5,7 +5,6 @@ models, samplers, sentence_transformers, - st_models, utils, ) from zeromodels.auto import ( @@ -25,7 +24,6 @@ "models", "samplers", "sentence_transformers", - "st_models", "utils", "conversion", "version", From 6db67c1816ef026d358b8c5f9e47573ae75c3d1f Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Sun, 18 Oct 2026 16:26:23 -0700 Subject: [PATCH 3/4] docs --- README.md | 32 +- docs/auto.md | 2 +- docs/diffusers.md | 1961 ++++++++++++++++++++++++++++++++- docs/getting_started.md | 3 +- docs/index.md | 3 +- docs/main_classes.md | 8 +- docs/models.md | 26 +- docs/qwen_image.md | 357 ------ docs/qwen_image_21.md | 299 ----- docs/sentence_transformer.md | 34 +- docs/sentence_transformers.md | 49 - docs/stable_diffusion.md | 353 ------ docs/stable_diffusion_2.md | 216 ---- docs/stable_diffusion_3.md | 266 ----- docs/stable_diffusion_3_5.md | 121 -- docs/stable_diffusion_xl.md | 373 ------- website/mkdocs.yml | 17 +- 17 files changed, 2015 insertions(+), 2105 deletions(-) delete mode 100644 docs/qwen_image.md delete mode 100644 docs/qwen_image_21.md delete mode 100644 docs/sentence_transformers.md delete mode 100644 docs/stable_diffusion.md delete mode 100644 docs/stable_diffusion_2.md delete mode 100644 docs/stable_diffusion_3.md delete mode 100644 docs/stable_diffusion_3_5.md delete mode 100644 docs/stable_diffusion_xl.md diff --git a/README.md b/README.md index add0da12..96b0675e 100644 --- a/README.md +++ b/README.md @@ -42,8 +42,6 @@ Documentation sources are also available in [`docs/`](docs/). ## 💬 Sentence Transformers Models (sentence_transformers) -All 13 checkpoints below have pre-converted Keras weights under the `zeromodels` Hugging Face organization. Use `ZMSentenceTransformer` and `ZMSentenceTransformerTokenizer` from `zeromodels.sentence_transformers` for sentence similarity, semantic search, or clustering. - | 🏷️ Model Name | 📜 Reference Paper | 📦 Source of Weights | |---------------|-------------------|---------------------| | [all-mpnet-base-v2](https://huggingface.co/zeromodels/all-mpnet-base-v2) | [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://arxiv.org/abs/1908.10084) | `sentence-transformers` | @@ -62,6 +60,20 @@ All 13 checkpoints below have pre-converted Keras weights under the `zeromodels`
+## 🎨 Diffusion Models (diffusers) + +| 🏷️ Model Name | 📜 Reference Paper | 📦 Source of Weights | +|---------------|-------------------|---------------------| +| Stable Diffusion 1.x | [High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) | `diffusers` | +| Stable Diffusion 2.x | [High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) | `diffusers` | +| Stable Diffusion XL | [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis](https://arxiv.org/abs/2307.01952) | `diffusers` | +| Stable Diffusion 3 | [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://arxiv.org/abs/2403.03206) | `diffusers` | +| Stable Diffusion 3.5 | [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://arxiv.org/abs/2403.03206) | `diffusers` | +| Qwen-Image | [Qwen-Image Technical Report](https://arxiv.org/abs/2508.02324) | `diffusers` | +| Qwen-Image-2.1 | [Qwen/Qwen-Image-2.1](https://huggingface.co/Qwen/Qwen-Image-2.1) | `diffusers` | + +
+ ## 📑 Models ### 📝 Text Models @@ -269,22 +281,6 @@ All 13 checkpoints below have pre-converted Keras weights under the `zeromodels`
-- Text-to-Image Diffusion (text → image) - - Every diffusion checkpoint loads with `ZMDiffusion` and `ZMDiffusionTokenizer` from `zeromodels.diffusers` ([documentation](https://imvision12.github.io/ZeroModels/diffusers/)). - - | 🏷️ Model Name | 📜 Reference Paper | 📦 Source of Weights | - |---------------|-------------------|---------------------| - | Stable Diffusion 1.x (v1-1 to v1-5) | [High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) | `diffusers` | - | Stable Diffusion 2.x (2, 2-base, 2-1, 2-1-base, SD-Turbo) | [High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) | `diffusers` | - | Stable Diffusion XL (base + refiner 1.0, SDXL-Turbo) | [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis](https://arxiv.org/abs/2307.01952) | `diffusers` | - | Stable Diffusion 3 (medium) | [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://arxiv.org/abs/2403.03206) | `diffusers` | - | Stable Diffusion 3.5 (large, large-turbo, medium) | [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://arxiv.org/abs/2403.03206) | `diffusers` | - | Qwen-Image | [Qwen-Image Technical Report](https://arxiv.org/abs/2508.02324) | `diffusers` | - | Qwen-Image-2.1 | [Qwen/Qwen-Image-2.1](https://huggingface.co/Qwen/Qwen-Image-2.1) (single-stream DiT + Qwen3-VL) | `diffusers` | - -
- ### 🔊 Audio Models - Speech (speech → text) diff --git a/docs/auto.md b/docs/auto.md index c066c74e..113310bf 100644 --- a/docs/auto.md +++ b/docs/auto.md @@ -63,7 +63,7 @@ print(embeddings.shape) # (2, 384) Here `AutoZModel` loads the complete sentence-embedding stack, including its pooling, optional Dense projection, and normalization—not just the text encoder. The tokenizer preserves the checkpoint's own preprocessing recipe. See -[Sentence Transformer](sentence_transformer.md) for the concrete API. Original +[Sentence Transformers](sentence_transformer.md) for the concrete API. Original upstream `hf:sentence-transformers/...` configs describe their backbone type; use the concrete classes to load their sentence-embedding recipe. diff --git a/docs/diffusers.md b/docs/diffusers.md index 6876aa32..1aa27996 100644 --- a/docs/diffusers.md +++ b/docs/diffusers.md @@ -1,4 +1,4 @@ -# Diffusers +# Diffusion Models
Weights: every diffusion checkpoint is hosted under the @@ -59,17 +59,17 @@ from each repo. | `model_type` | Hosted checkpoints | Page | |---|---|---| -| `stable_diffusion` | `stable-diffusion-v1-1` to `stable-diffusion-v1-5` | [Stable Diffusion](stable_diffusion.md) | -| `stable_diffusion_2` | `stable-diffusion-2`, `stable-diffusion-2-base`, `stable-diffusion-2-1`, `stable-diffusion-2-1-base`, `sd-turbo` | [Stable Diffusion 2](stable_diffusion_2.md) | -| `stable_diffusion_xl` | `stable-diffusion-xl-base-1.0`, `sdxl-turbo` | [Stable Diffusion XL](stable_diffusion_xl.md) | -| `stable_diffusion_xl_refiner` | `stable-diffusion-xl-refiner-1.0` (image-to-image) | [Stable Diffusion XL](stable_diffusion_xl.md#the-refiner) | -| `stable_diffusion_3` | `stable-diffusion-3-medium` | [Stable Diffusion 3](stable_diffusion_3.md) | -| `stable_diffusion_3_5` | `stable-diffusion-3.5-large`, `stable-diffusion-3.5-large-turbo`, `stable-diffusion-3.5-medium` | [Stable Diffusion 3.5](stable_diffusion_3_5.md) | -| `qwen_image` | `qwen-image` | [Qwen-Image](qwen_image.md) | -| `qwen_image_21` | `qwen-image-2.1` | [Qwen-Image-2.1](qwen_image_21.md) | +| `stable_diffusion` | `stable-diffusion-v1-1` to `stable-diffusion-v1-5` | [Stable Diffusion](#stable-diffusion) | +| `stable_diffusion_2` | `stable-diffusion-2`, `stable-diffusion-2-base`, `stable-diffusion-2-1`, `stable-diffusion-2-1-base`, `sd-turbo` | [Stable Diffusion 2](#stable-diffusion-2) | +| `stable_diffusion_xl` | `stable-diffusion-xl-base-1.0`, `sdxl-turbo` | [Stable Diffusion XL](#stable-diffusion-xl) | +| `stable_diffusion_xl_refiner` | `stable-diffusion-xl-refiner-1.0` (image-to-image) | [Stable Diffusion XL](#stable-diffusion-xl-the-refiner) | +| `stable_diffusion_3` | `stable-diffusion-3-medium` | [Stable Diffusion 3](#stable-diffusion-3) | +| `stable_diffusion_3_5` | `stable-diffusion-3.5-large`, `stable-diffusion-3.5-large-turbo`, `stable-diffusion-3.5-medium` | [Stable Diffusion 3.5](#stable-diffusion-3-5) | +| `qwen_image` | `qwen-image` | [Qwen-Image](#qwen-image) | +| `qwen_image_21` | `qwen-image-2.1` | [Qwen-Image-2.1](#qwen-image-2-1) | -The family pages cover each architecture, its generation arguments, the verification -against diffusers and the memory needed. +The family sections below cover each architecture, its generation arguments, the +verification against diffusers and the memory needed. ## `ZMDiffusion` @@ -213,6 +213,1945 @@ same repo): } ``` +## Stable Diffusion 1.x {#stable-diffusion} + +Stable Diffusion 1.x, ported to pure Keras 3: latent text-to-image diffusion with a +UNet denoiser, a KL-regularized autoencoder and the CLIP ViT-L/14 text encoder. It loads +through [`ZMDiffusion`](#zmdiffusion), the one diffusion pipeline class of zeromodels, +with `model_type="stable_diffusion"`: a single functional graph with the three towers as +disconnected sub-graphs (the way `CLIPModel` holds its two), so a hosted repo is one +`model.weights.h5` and one `zm_config.json`, plus `generate`, the classifier-free +guidance denoising loop of `BaseDiffusion` (the diffusion counterpart of +`BaseGeneration`). Prompts are tokenized by `ZMDiffusionTokenizer`. + +All five SD 1.x checkpoints (v1-1 to v1-5) share this architecture; only the weights +differ. The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion +is deliberately **not supported** for diffusion models. + +Key facts of the port: + +- **Both data formats**: the UNet and VAE follow `keras.config.image_data_format()` + (every conv takes it, GroupNorm normalizes the layout's channel axis, and the two + token <-> grid boundaries transpose around their flatten under `channels_first`). + Weights are layout-independent, so one hosted checkpoint serves both, and the + converter is a `(O, I, H, W) -> (H, W, I, O)` kernel transpose and nothing else. +- **Block-level layers**: the UNet and VAE are built from composite Keras layers + (`StableDiffusionResnetBlock2D`, `StableDiffusionTransformer2DModel`, + `StableDiffusionCrossAttention`, ...). A functional graph keeps every node's output + alive until the forward ends, and at 512px the UNet's per-op intermediates (4096x4096 + attention maps above all) would need several GB; a layer's internals are freed when + its call returns. +- **Schedulers match diffusers to the bit**: the training noise schedule is built in + float32 the way torch does it, so `alphas_cumprod` and the Euler sigmas are identical + and a 50-step run stays within float rounding of the reference. +- **Verified end to end**: at 512px, fp32, PNDM 50 steps and the same initial latent, the + generated image matches diffusers' to within **1 uint8 level** (99.3% of pixels + identical, PSNR 75 dB); the UNet, VAE and text encoder outputs agree to ~1e-6. + +Links: + +- Paper: [High-Resolution Image Synthesis with Latent Diffusion Models (arXiv:2112.10752)](https://arxiv.org/abs/2112.10752) +- Reference implementation: [diffusers `StableDiffusionPipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img) +- License: [CreativeML OpenRAIL-M](https://huggingface.co/spaces/CompVis/stable-diffusion-license) + +See also [clip.md](clip.md) (the text encoder is `CLIPTextModel`). + +### Variants {#stable-diffusion-variants} + +Preconverted, float32 weights are hosted under `zeromodels/`. Load with +`from_weights("zeromodels/")`. Each repo is one container: UNet (860M) + VAE +(84M) + CLIP text encoder (123M), 1.07B parameters, 3.97 GB. All checkpoints are released +under the CreativeML OpenRAIL-M license, which carries use-based restrictions. + +| Variant | Hub | Training | +|---|---|---| +| `stable-diffusion-v1-1` | [`zeromodels/stable-diffusion-v1-1`](https://huggingface.co/zeromodels/stable-diffusion-v1-1) | 237k steps at 256px on laion2B-en, then 194k steps at 512px on laion-high-resolution | +| `stable-diffusion-v1-2` | [`zeromodels/stable-diffusion-v1-2`](https://huggingface.co/zeromodels/stable-diffusion-v1-2) | v1-1 + 515k steps at 512px on laion-aesthetics v2 5+ | +| `stable-diffusion-v1-3` | [`zeromodels/stable-diffusion-v1-3`](https://huggingface.co/zeromodels/stable-diffusion-v1-3) | v1-2 + 195k steps at 512px, 10% text-conditioning dropout | +| `stable-diffusion-v1-4` | [`zeromodels/stable-diffusion-v1-4`](https://huggingface.co/zeromodels/stable-diffusion-v1-4) | v1-2 + 225k steps at 512px, 10% text-conditioning dropout | +| `stable-diffusion-v1-5` | [`zeromodels/stable-diffusion-v1-5`](https://huggingface.co/zeromodels/stable-diffusion-v1-5) | v1-2 + 595k steps at 512px, 10% text-conditioning dropout | + +The 10% text-conditioning dropout of v1-3 onwards is what makes classifier-free guidance +work well; v1-5 is the usual default. + +### API {#stable-diffusion-api} + +Configs are typed: `StableDiffusionConfig` (composite) over `UNet2DConditionConfig`, +`AutoencoderKLConfig` and `StableDiffusionTextConfig` (a `CLIPTextConfig` with the +ViT-L/14 defaults), plus the checkpoint's `scheduler_config` and the CLIP special token +ids. Each repo's `zm_config.json` parses through it; the constructor stays flat, with the +sub-config fields prefixed `unet_` / `vae_` / `text_`. + +#### `ZMDiffusion` {#stable-diffusion-zmdiffusion} + +The pipeline class of every diffusion family (see [Diffusers](#zmdiffusion)); a repo's +`model_type` (`"stable_diffusion"` here) selects the SD 1.x family, which supplies the +hooks `BaseDiffusion`'s `generate` needs (`encode_prompt` on the CLIP tower, the empty +prompt's `unconditional_ids`, `predict_noise` on the UNet, `decode_latents` on the VAE) +and a `scheduler`, built from the config's `scheduler_config` (PNDM for every SD 1.x +checkpoint). + +```python +generate( + input_ids, + attention_mask=None, + negative_input_ids=None, + num_inference_steps=None, + guidance_scale=None, + seed=None, + latents=None, +) +``` + +| Arg | Default | Meaning | +|---|---|---| +| `input_ids` | required | `(batch, 77)` token ids, `**tokenizer(prompts)` | +| `attention_mask` | `None` | accepted alongside the tokenizer's output; SD attends its padding tokens, so it is unused | +| `negative_input_ids` | `None` | `(batch, 77)` tokenized negative prompt for guidance; the empty prompt when unset | +| `num_inference_steps` | `None` | scheduler steps; the repo's `generate_args` (50) when unset | +| `guidance_scale` | `None` | classifier-free guidance strength; `generate_args` (7.5) when unset, `<= 1` disables it | +| `seed` | `None` | seed for the initial latent, reproducible per backend | +| `latents` | `None` | explicit `(batch, 64, 64, 4)` initial latent; identical results across backends | + +Returns `(batch, 512, 512, 3)` uint8 numpy images. The denoiser call (the UNet plus the +guidance combine) runs as a per-backend compiled function cached on the instance, +`jax.jit` on JAX and `tf.function(jit_compile=True)` on TensorFlow, eager on Torch; its +shapes are the same at every step, so one compile serves the whole loop. + +| Constructor arg | Default | Meaning | +|---|---|---| +| `model_type` | required | `"stable_diffusion"` (read from the repo's `zm_config.json` on `from_weights`) | +| `scheduler` | `None` | a `BaseScheduler`; built from `scheduler_config` when unset | +| `unet_sample_size` | `64` | latent size the UNet graph is built for (image size / 8) | +| `vae_sample_size` | `512` | image size the VAE graphs are built for | +| `unet_*` / `vae_*` / `text_*` | SD 1.x | the flat sub-config fields (see [Configuration](configuration.md)) | +| `bos_token_id` / `eos_token_id` / `pad_token_id` | `49406` / `49407` / `49407` | CLIP's special tokens, used to build the empty prompt | +| `scheduler_config` | `None` | the diffusers scheduler dict (`_class_name`, betas, `steps_offset`, ...) | + +The graph is three disconnected paths, one per component. Inputs +`{"sample", "timestep", "encoder_hidden_states"}` (UNet), `{"image", "latent"}` (VAE +encode / decode) and `{"token_ids", "padding_mask"}` (text); outputs +`{"noise_pred", "moments", "image", "text_embeds"}`. The components are exposed as +`.unet`, `.vae` and `.text_encoder` and share the model's weights, so one load fills all +three, and each can be called on its own. + +#### `UNet2DConditionModel` {#stable-diffusion-unet2dconditionmodel} + +The denoiser (diffusers' `UNet2DConditionModel`): a symmetric down / mid / up UNet whose +ResNet blocks are conditioned on the sinusoidal timestep embedding and whose +`CrossAttn` blocks attend the latent to the text context. Inputs +`{"sample": (B, h, w, 4), "timestep": (B,), "encoder_hidden_states": (B, 77, 768)}`, +output `{"sample": (B, h, w, 4)}`. + +| Arg | Default | Meaning | +|---|---|---| +| `sample_size` | `64` | latent spatial size the graph is built for | +| `in_channels` / `out_channels` | `4` | latent channels | +| `down_block_types` / `up_block_types` | SD 1.x | `CrossAttnDownBlock2D` x3 + `DownBlock2D`, mirrored | +| `block_out_channels` | `(320, 640, 1280, 1280)` | width per level | +| `layers_per_block` | `2` | ResNet blocks per level | +| `cross_attention_dim` | `768` | width of the text context | +| `num_attention_heads` | `8` | heads in the cross-attention blocks | +| `norm_num_groups` | `32` | GroupNorm groups | +| `text_seq_len` | `77` | static text length | + +#### `AutoencoderKL` {#stable-diffusion-autoencoderkl} + +The VAE (diffusers' `AutoencoderKL`): `encode(image, sample=False)` returns the posterior +mean (or a sample) of a `(B, H/8, W/8, 4)` latent and `decode(latent)` the `(B, H, W, 3)` +image in `[-1, 1]`; multiply / divide by `scaling_factor` (0.18215) around the UNet, as +`generate` does. One functional graph with the encoder and decoder as two disconnected +paths. + +| Arg | Default | Meaning | +|---|---|---| +| `in_channels` / `out_channels` | `3` | image channels | +| `latent_channels` | `4` | latent channels | +| `block_out_channels` | `(128, 256, 512, 512)` | width per level, x8 spatial compression | +| `layers_per_block` | `2` | ResNet blocks per encoder level | +| `norm_num_groups` | `32` | GroupNorm groups | +| `sample_size` | `512` | image size the graphs are built for | +| `scaling_factor` | `0.18215` | latent scaling | + +#### Schedulers {#stable-diffusion-schedulers} + +`zeromodels.base.base_scheduler` has the SD samplers as weightless classes with the +diffusers interface (`set_timesteps`, `scale_model_input`, `step`, +`init_noise_sigma`), configured the diffusers way: + +| Scheduler | Notes | +|---|---| +| `PNDMScheduler` | the SD 1.x default (PLMS, `skip_prk_steps=True`) | +| `DDIMScheduler` | deterministic, `eta=0` | +| `EulerDiscreteScheduler` | Karras-style sigma sampler, `linspace` spacing | +| `EulerAncestralDiscreteScheduler` | stochastic Euler (`seed` for reproducibility) | + +`get_scheduler(config)` builds the one named by a diffusers scheduler config dict, which +is what the model does with its `scheduler_config`; swap at any time: + +```python +from zeromodels.base.base_scheduler import EulerDiscreteScheduler + +model.scheduler = EulerDiscreteScheduler.from_config(model.config.scheduler_config) +``` + +### Preprocessing {#stable-diffusion-preprocessing} + +#### `ZMDiffusionTokenizer` {#stable-diffusion-zmdiffusiontokenizer} + +The tokenizer of every diffusion family; for `model_type="stable_diffusion"` it is the +CLIP ViT-L/14 BPE tokenizer: `<|startoftext|>` / `<|endoftext|>` framing, truncated and +`<|endoftext|>`-padded to 77 tokens, the fixed context length the UNet's cross-attention +is built for. Calling it returns `{"input_ids", "attention_mask"}` for a string or a list +of strings. `from_weights("zeromodels/")` reads the family from the repo. + +```python +ZMDiffusionTokenizer( + model_type="stable_diffusion", hf_id=None, tokenizer_file=None, max_seq_len=77 +) +``` + +There is no image processor and no processor class: text-to-image takes token ids in +and hands uint8 images out. + +### End-to-end example {#stable-diffusion-end-to-end-example} + +#### Single prompt {#stable-diffusion-single-prompt} + +```python +import os + +os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" + +from PIL import Image +from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer + +model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-v1-5") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-v1-5") + +inputs = tokenizer( + "a red fox sitting in a field of lavender, golden hour, wildlife photography, sharp focus" +) +images = model.generate(**inputs, num_inference_steps=50, guidance_scale=7.5, seed=1) + +Image.fromarray(images[0]).save("fox.png") # (512, 512, 3) uint8 +``` + +Stable Diffusion v1-5: a red fox sitting in a field of lavender at golden hour, 512px + +#### Negative prompt and batching {#stable-diffusion-negative-prompt-and-batching} + +Prompts batch through the tokenizer; a negative prompt is tokenized the same way and +passed as `negative_input_ids` (one per prompt, or one broadcast row). + +```python +prompts = [ + "a watercolor painting of a lighthouse at sunset", + "a red vintage car on a coastal road", +] +inputs = tokenizer(prompts) # {"input_ids": (2, 77), "attention_mask": (2, 77)} +negative = tokenizer(["blurry, low quality"] * len(prompts))["input_ids"] + +images = model.generate(**inputs, negative_input_ids=negative, num_inference_steps=30) +for i, image in enumerate(images): + Image.fromarray(image).save(f"out_{i}.png") +``` + +#### Reproducible across backends {#stable-diffusion-reproducible-across-backends} + +`seed` reproduces a run on a given backend, but `keras.random` differs between torch, +jax and tensorflow. For an image that is identical everywhere, pass the initial latent +yourself (`(batch, 64, 64, 4)` at 512px): + +```python +import numpy as np + +latents = np.random.default_rng(0).standard_normal((1, 64, 64, 4)).astype("float32") +images = model.generate(**tokenizer("a bowl of ramen"), latents=latents) +``` + +#### Verified against diffusers {#stable-diffusion-verified-against-diffusers} + +The fox prompt above with the same initial latent, PNDM 50 steps and guidance 7.5 +through `ZMDiffusion` (left) and diffusers' `StableDiffusionPipeline` +(right), both fp32: + +zeromodels (left) and diffusers (right) generations of a red fox in lavender from the same latent + +``` +tokenizer input_ids identical +text_encoder max|d|=9.5e-06 +unet noise_pred (t=981) max|d|=9.5e-07 +vae decode max|d|=3.5e-06 +final image (uint8) max|d|=1 mean|d|=0.0022 PSNR=74.7 dB identical_pixels=99.34% +``` + +#### Components on their own {#stable-diffusion-components-on-their-own} + +```python +from zeromodels.diffusers import ZMDiffusion + +sd = ZMDiffusion.from_weights("zeromodels/stable-diffusion-v1-5") +latent = sd.vae.encode(image) # (1, 512, 512, 3) in [-1, 1] -> (1, 64, 64, 4) +text = sd.text_encoder({"token_ids": ids, "padding_mask": mask})["last_hidden_state"] +``` + +#### Other resolutions {#stable-diffusion-other-resolutions} + +The graphs are built for a fixed size, the weights are not. Rebuild at another size with +the constructor overrides (multiples of 64px; SD 1.x was trained at 512px): + +```python +model = ZMDiffusion.from_weights( + "zeromodels/stable-diffusion-v1-5", unet_sample_size=96, vae_sample_size=768 +) +images = model.generate(**tokenizer("a mountain lake at dawn")) # (1, 768, 768, 3) +``` + +### Data Format {#stable-diffusion-data-format} + +**Both `channels_last` and `channels_first` are supported.** The models read +`keras.config.image_data_format()` when they are constructed (there is no +`data_format` argument), so set it once, before `from_weights`, and every tensor that +crosses the model boundary follows it: + +| | `channels_last` (default) | `channels_first` | +|---|---|---| +| `latents` passed to `generate` | `(batch, 64, 64, 4)` | `(batch, 4, 64, 64)` | +| `AutoencoderKL.encode` input / `decode` output | `(batch, 512, 512, 3)` | `(batch, 3, 512, 512)` | +| UNet `sample` | `(batch, 64, 64, 4)` | `(batch, 4, 64, 64)` | + +`generate` itself always hands back `(batch, height, width, 3)` uint8 images, whatever +the layout, so the PIL step is the same in both. The same hosted `model.weights.h5` +loads under either layout (conv kernels are `(kh, kw, in, out)` in both), and the two +runs agree to within 1 uint8 level (99.7% of pixels identical, the rest being the +NCHW / NHWC convolution kernels' rounding). + +```python +import keras + +keras.config.set_image_data_format("channels_first") + +model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-v1-5") +images = model.generate( + **tokenizer("a bowl of ramen"), seed=0 +) # still (1, 512, 512, 3) +``` + +Note that `keras.config.set_image_data_format` is global state. Set it once at the top +of a script rather than toggling it between calls, since already-built models keep the +layout they were constructed with. + +### Memory and speed {#stable-diffusion-memory-and-speed} + +The fp32 container is 3.97 GB; a 512px, batch-1, guided run peaks at 6.5 GB of GPU +memory (the guided UNet call runs batch 2). On an RTX 4060 laptop the 50-step run takes +~25 s in eager Torch, about 1.5x diffusers' eager time. On JAX the compiled denoiser step +measured 2-3x faster than an op-by-op run; TensorFlow uses the same mechanism. + +### Loading Fine-tuned Weights {#stable-diffusion-loading-fine-tuned-weights} + +The five hosted checkpoints are the supported weights; any repo laid out like them +(`zm_config.json` declaring `ZMDiffusion` with `model_type: "stable_diffusion"`, +`model.weights.h5`, `tokenizer.json`) loads with `from_weights("/")`; configs +written before `ZMDiffusion` (declaring `StableDiffusionModel`) load the same way. The +`hf:` prefix raises for diffusion models: +convert a diffusers-format checkpoint once with +`zeromodels/diffusers/stable_diffusion/convert_stable_diffusion_diffusers_to_keras.py` +(`transfer_stable_diffusion(repo)`, `pip install zeromodels[conversion]`) and host the result. + +## Stable Diffusion 2.x {#stable-diffusion-2} + +Stable Diffusion 2.x, ported to pure Keras 3: the [Stable Diffusion](#stable-diffusion) +latent-diffusion architecture with the second-generation conditioning. It loads +through [`ZMDiffusion`](#zmdiffusion) with `model_type="stable_diffusion_2"`, the SD 1.x +family with the SD 2 configuration, so everything on that page applies: one functional +model and one weights file per repo, `generate` from `BaseDiffusion`, the schedulers, +both data formats. What changes is the config: + +- **Text encoder**: OpenCLIP ViT-H/14's text tower, 1024 wide, 16 heads, `gelu`, shipped + truncated to its 23rd (penultimate) layer, which SD 2 conditions on. +- **UNet**: 1024-d cross-attention, one head count per level (`(5, 10, 20, 20)`, a 64-wide + head everywhere, where SD 1.x used 8 heads at every level) and a linear token + projection in the Transformer2D blocks instead of the 1x1 conv. +- **Tokenizer**: the same CLIP BPE, padded with `!` (id 0) the OpenCLIP way rather than + `<|endoftext|>`, so the empty prompt used for classifier-free guidance is + `[<|startoftext|>, <|endoftext|>, !, !, ...]`. +- **768px checkpoints** (`stable-diffusion-2`, `stable-diffusion-2-1`): built at a 96x96 + latent and trained with the **v-prediction** objective; their repos carry a + v-prediction DDIM scheduler, which `generate` picks up from `scheduler_config`. +- **SD-Turbo** (`sd-turbo`): SD 2.1 distilled with Adversarial Diffusion Distillation to + generate in 1 to 4 steps **without guidance**; its repo carries an Euler scheduler with + `trailing` timestep spacing and `generate_args` of 1 step, `guidance_scale=0.0`. + +The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion is +deliberately **not supported** for diffusion models. The original `stabilityai/*` repos are +no longer on the Hub; the conversion sources are the `sd2-community` mirrors. + +Links: + +- Paper: [High-Resolution Image Synthesis with Latent Diffusion Models (arXiv:2112.10752)](https://arxiv.org/abs/2112.10752) +- Reference implementation: [diffusers `StableDiffusionPipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img) +- License: [CreativeML Open RAIL++-M](https://huggingface.co/sd2-community/stable-diffusion-2-1/blob/main/LICENSE-MODEL) + +See also [Stable Diffusion 1.x](#stable-diffusion), [clip.md](clip.md). + +### Variants {#stable-diffusion-2-variants} + +Preconverted, float32 weights are hosted under `zeromodels/`. Load with +`from_weights("zeromodels/")`. Each repo is one container: UNet (866M) + VAE +(84M) + OpenCLIP text encoder (340M), 1.29B parameters, 5.16 GB (4.81 GiB, one +`model.weights.h5`). The SD 2 checkpoints are released under the CreativeML Open RAIL++-M +license; SD-Turbo under the Stability AI Non-Commercial Research Community License. + +| Variant | Hub | Resolution | Objective | Training | +|---|---|---|---|---| +| `stable-diffusion-2-base` | [`zeromodels/stable-diffusion-2-base`](https://huggingface.co/zeromodels/stable-diffusion-2-base) | 512 | 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) | 768 | 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) | 512 | 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) | 768 | 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) | 512 | epsilon, 1 to 4 steps, no guidance | SD 2.1 distilled with Adversarial Diffusion Distillation (non-commercial) | + +Use the `-base` checkpoints for 512px images and the others for 768px; each repo's +`zm_config.json` builds the graph at its native size. + +### API {#stable-diffusion-2-api} + +Configs are typed: `StableDiffusion2Config` (composite, `model_type` +`"stable_diffusion_2"`) over `StableDiffusion2UNetConfig` (a `UNet2DConditionConfig` +with the SD 2 widths), `AutoencoderKLConfig` and `StableDiffusion2TextConfig` (a +`CLIPTextConfig` with the ViT-H/14 sizes), plus the scheduler config and the token ids +(`pad_token_id` 0). Flat constructor, `unet_` / `vae_` / `text_` prefixes, like SD 1.x. + +#### `ZMDiffusion` {#stable-diffusion-2-zmdiffusion} + +`ZMDiffusion` with `model_type="stable_diffusion_2"` (read from the repo); `generate` is +the SD 1.x one: + +```python +generate( + input_ids, + attention_mask=None, + negative_input_ids=None, + num_inference_steps=None, + guidance_scale=None, + seed=None, + latents=None, +) +``` + +Returns `(batch, H, W, 3)` uint8 images at the checkpoint's resolution; `latents` is +`(batch, 64, 64, 4)` for the 512px checkpoints and `(batch, 96, 96, 4)` for the 768px +ones (`channels_first`: channels second). Defaults come from the repo's `generate_args` +(50 steps, guidance 7.5). See [Stable Diffusion](#stable-diffusion-zmdiffusion) +for the argument table. + +The graph is SD 1.x's three disconnected paths (`.unet`, `.vae`, `.text_encoder`) with +the same inputs / outputs, `encoder_hidden_states` 1024 wide. + +The components are the SD 1.x classes, `UNet2DConditionModel` (built with +`num_attention_heads=(5, 10, 20, 20)`, `use_linear_projection=True`, +`cross_attention_dim=1024`) and `AutoencoderKL`; see their tables on the +[Stable Diffusion](#stable-diffusion-api) page. + +### Preprocessing {#stable-diffusion-2-preprocessing} + +#### `ZMDiffusionTokenizer` {#stable-diffusion-2-zmdiffusiontokenizer} + +`ZMDiffusionTokenizer` with `model_type="stable_diffusion_2"`, which pads with `!`: +CLIP BPE, `<|startoftext|>` / `<|endoftext|>` framing, truncated and `!`-padded to 77 +tokens. Returns `{"input_ids", "attention_mask"}`. + +```python +ZMDiffusionTokenizer( + model_type="stable_diffusion_2", hf_id=None, tokenizer_file=None, max_seq_len=77 +) # pad_token defaults to "!" +``` + +### End-to-end example {#stable-diffusion-2-end-to-end-example} + +```python +import os + +os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" + +from PIL import Image +from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer + +model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-2-1-base") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-2-1-base") + +inputs = tokenizer( + "a steaming bowl of ramen on a wooden table, food photography, shallow depth of field" +) +images = model.generate(**inputs, num_inference_steps=50, guidance_scale=7.5, seed=2) + +Image.fromarray(images[0]).save("ramen.png") # (512, 512, 3) uint8 +``` + +Stable Diffusion 2.1-base: a steaming bowl of ramen on a wooden table, 512px + +#### 768px, v-prediction {#stable-diffusion-2-768px-v-prediction} + +The 768px checkpoints need no extra arguments; the v-prediction DDIM scheduler and the +96x96 latent come from the repo: + +```python +model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-2-1") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-2-1") +images = model.generate( + **tokenizer("a lighthouse on a cliff at dusk, oil painting") +) # (1, 768, 768, 3) +``` + +#### SD-Turbo {#stable-diffusion-2-sd-turbo} + +One step, no guidance (the repo's defaults; up to 4 steps sharpen a little): + +```python +model = ZMDiffusion.from_weights("zeromodels/sd-turbo") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/sd-turbo") +images = model.generate( + **tokenizer( + "a cinematic shot of a baby raccoon wearing an intricate italian priest robe" + ) +) +``` + +SD-Turbo, one step: a baby raccoon in an italian priest robe, 512px + +The same prompt and latent through diffusers' `StableDiffusionPipeline` (fp32) match to +1 uint8 level: 99.9% of pixels identical at 1 step (PSNR 82 dB), 99.8% at 4 steps. + +Negative prompts, batching, explicit `latents` for cross-backend reproducibility, other +resolutions and the components on their own work exactly as on the +[Stable Diffusion](#stable-diffusion-end-to-end-example) page. Image-to-image +(`generate(..., image=..., strength=...)`) works as described for +[`BaseDiffusion`](main_classes.md#basediffusion). + +#### Verified against diffusers {#stable-diffusion-2-verified-against-diffusers} + +The ramen prompt above with the same initial latent, PNDM 50 steps and guidance 7.5 +through `ZMDiffusion` (left) and diffusers' `StableDiffusionPipeline` +(right) on `stable-diffusion-2-1-base`, both fp32: + +zeromodels (left) and diffusers (right) generations of a bowl of ramen from stable-diffusion-2-1-base with the same latent + +``` +tokenizer input_ids identical +text_encoder max|d|=1.5e-05 +unet noise_pred (t=981) max|d|=7.2e-07 +vae decode max|d|=2.6e-06 +final image (uint8) max|d|=1 mean|d|=0.0095 PSNR=68.3 dB identical_pixels=97.22% +``` + +The 768px v-prediction checkpoint (`stable-diffusion-2-1`, DDIM 20 steps, same latent) +matches the same way: UNet 8.3e-6, VAE 7.9e-5, final image max 1 uint8 level, 99.3% of +pixels identical, PSNR 74.6 dB. + +### Data Format {#stable-diffusion-2-data-format} + +Both `channels_last` and `channels_first` are supported, as for +[Stable Diffusion](#stable-diffusion-data-format); `generate` always returns +`(batch, H, W, 3)` uint8. + +### Memory and speed {#stable-diffusion-2-memory-and-speed} + +The fp32 container is 5.16 GB. At 512px a guided batch-1 run needs a little over 7 GB of +GPU memory in eager Torch (about 1.5x diffusers' eager time); at 768px the self-attention +runs over 9216 tokens, so the guided run needs roughly 12 GB, or the CPU. + +### Loading Fine-tuned Weights {#stable-diffusion-2-loading-fine-tuned-weights} + +Any repo laid out like the hosted ones (`zm_config.json` declaring +`ZMDiffusion` with `model_type: "stable_diffusion_2"`, the weights, `tokenizer.json`) +loads with `from_weights("/")`. The `hf:` prefix raises for diffusion models: convert a +diffusers-format SD 2 checkpoint once with +`zeromodels/diffusers/stable_diffusion_2/convert_stable_diffusion_2_diffusers_to_keras.py` +(`transfer_stable_diffusion_2(repo)`, `pip install zeromodels[conversion]`) and host the result. + +## Stable Diffusion XL {#stable-diffusion-xl} + +Stable Diffusion XL, ported to pure Keras 3: the [Stable Diffusion](#stable-diffusion) +latent-diffusion recipe scaled up to 1024px with a 2.6B-parameter UNet and two text +encoders, plus the **refiner**, an image-to-image expert that finishes the base model's +latents, and **SDXL-Turbo**, the base distilled to a few steps. It loads +through [`ZMDiffusion`](#zmdiffusion) with `model_type="stable_diffusion_xl"` (the +refiner: `"stable_diffusion_xl_refiner"`), a family built on SD 1.x's with a fourth +component and the SDXL conditioning hooks, so everything on that page applies: one +functional model and one set of weights per repo, `generate` from `BaseDiffusion`, the +schedulers, both data formats. +What changes: + +- **Two text encoders**: CLIP ViT-L/14's text tower (768-d, the SD 1.x one) and OpenCLIP + ViT-bigG/14's (1280-d, 32 layers, `gelu`, with a 1280-d `text_projection`). The prompt + goes through both; their **penultimate** hidden states (no final LayerNorm) are + concatenated into the UNet's 2048-d cross-attention context, and the second tower's + projected pooled (EOT) state is the `text_embeds` micro-conditioning. +- **One tokenizer call**: both towers use the same CLIP BPE and differ only in the pad + token (`<|endoftext|>` vs `!`), so `ZMDiffusionTokenizer` returns the first one's ids + and the model derives the second tower's ids from the returned `attention_mask`. +- **Micro-conditioning**: the UNet's timestep embedding also receives the pooled text + embedding and six size / crop values (`original_size`, `crops_coords_top_left`, + `target_size`, each `(height, width)`), sinusoidally embedded and projected by the + `add_embedding` MLP (`addition_embed_type="text_time"`). +- **UNet**: three levels (320, 640, 1280) with no attention at the first, (1, 2, 10) + transformer blocks per Transformer2D, (5, 10, 20) heads (64 wide), linear token + projection, built at a 128x128 latent (1024px). +- **VAE**: the SDXL autoencoder (`scaling_factor` 0.13025), built in **float32 whatever + the load dtype** (`force_upcast`: it overflows in float16, as in diffusers). +- **Classifier-free guidance** with no negative prompt uses **zero** embeddings + (`force_zeros_for_empty_prompt`), not the encoded empty prompt. Default guidance is 5.0. +- **Scheduler**: Euler with `leading` timestep spacing (`steps_offset` 1) for the base + models and the refiners, ancestral Euler with `trailing` spacing for SDXL-Turbo; all + come from the repo's `scheduler_config`. +- **Refiner**: a second checkpoint (a `ZMDiffusion` with + `model_type="stable_diffusion_xl_refiner"`) with the OpenCLIP tower alone (1280-d + context), a four-level UNet ((384, 768, 1536, 1536), attention at the middle levels, 4 + blocks per Transformer2D, 2.3B parameters) and an **aesthetic score** as the fifth + micro-conditioning id. It refines a latent the base left partially denoised + (`denoising_start`) or any image (`strength`); its empty negative prompt is encoded, not + zeroed. +- **float16 weights**: SDXL was released in float16, so the repos store that (the VAE in + float32) and `from_weights` builds the model in float16 by default (6.6 GB of weights). + Pass `load_dtype="float32"` for a float32 model (13.9 GB). + +The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion is +deliberately **not supported** for diffusion models. + +Links: + +- Paper: [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis (arXiv:2307.01952)](https://arxiv.org/abs/2307.01952) +- Reference implementation: [diffusers `StableDiffusionXLPipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl) +- Licenses: [CreativeML Open RAIL++-M](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/LICENSE.md) + (1.0 base and refiner), [Stability AI Non-Commercial Research Community License](https://huggingface.co/stabilityai/sdxl-turbo/blob/main/LICENSE.md) + (Turbo); the 0.9 research previews are not hosted + +See also [Stable Diffusion 1.x](#stable-diffusion), [Stable Diffusion 2.x](#stable-diffusion-2), [clip.md](clip.md). + +### Variants {#stable-diffusion-xl-variants} + +Preconverted weights are hosted under `zeromodels/`. Load with +`from_weights("zeromodels/")`. A base repo is one container: UNet (2.57B) + VAE +(84M) + CLIP ViT-L/14 text encoder (123M) + OpenCLIP ViT-bigG/14 text encoder (695M), +3.47B parameters, 6.6 GB in float16 (two `model.weights.json` shards); a refiner repo is +UNet (2.26B) + VAE + OpenCLIP text encoder, 3.04B parameters, 5.8 GB. + +| Variant | Hub | `model_type` | Resolution | Sampler | License | +|---|---|---|---|---|---| +| `stable-diffusion-xl-base-1.0` | [`zeromodels/stable-diffusion-xl-base-1.0`](https://huggingface.co/zeromodels/stable-diffusion-xl-base-1.0) | `stable_diffusion_xl` | 1024 | Euler, 50 steps, guidance 5.0 | CreativeML Open RAIL++-M | +| `stable-diffusion-xl-refiner-1.0` | [`zeromodels/stable-diffusion-xl-refiner-1.0`](https://huggingface.co/zeromodels/stable-diffusion-xl-refiner-1.0) | `stable_diffusion_xl_refiner` | 1024 | Euler, 50 steps, guidance 5.0, strength 0.3 | CreativeML Open RAIL++-M | +| `sdxl-turbo` | [`zeromodels/sdxl-turbo`](https://huggingface.co/zeromodels/sdxl-turbo) | `stable_diffusion_xl` | 512 | Euler ancestral, 1 step, no guidance | Stability AI Non-Commercial Research Community | + +The SDXL 0.9 research previews (`stabilityai/stable-diffusion-xl-base-0.9` and +`-refiner-0.9`, the same architecture with different weights, gated behind the research +license) are not hosted; with an authorized token the converter turns them into the same +checkpoints (`transfer_stable_diffusion_xl(repo, token=...)`, see +[Loading Fine-tuned Weights](#stable-diffusion-xl-loading-fine-tuned-weights)). SDXL-Turbo is a +research-only, non-commercial release; its repo defaults (`generate_args`) are 1 step and +`guidance_scale=0.0`, which is how it was trained (do not add guidance). + +### API {#stable-diffusion-xl-api} + +Configs are typed: `StableDiffusionXLConfig` (composite, `model_type` +`"stable_diffusion_xl"`) over `StableDiffusionXLUNetConfig`, `StableDiffusionXLVAEConfig`, +`StableDiffusionTextConfig` (the CLIP ViT-L/14 tower, as SD 1.x) and +`StableDiffusionXLTextConfig2` (the OpenCLIP ViT-bigG/14 tower, with `projection_dim` and +its own `hidden_act`), plus the scheduler config, the token ids (`pad_token_id_2` for the +second tokenizer's `!`), `force_zeros_for_empty_prompt` and `requires_aesthetics_score`. +Flat constructor with `unet_` / `vae_` / `text_` / `text_2_` prefixes. +`StableDiffusionXLRefinerConfig` (`model_type` `"stable_diffusion_xl_refiner"`) is the +same with `text_config=None`, the `StableDiffusionXLRefinerUNetConfig` denoiser, +`requires_aesthetics_score=True` and `force_zeros_for_empty_prompt=False`. + +#### `ZMDiffusion` {#stable-diffusion-xl-zmdiffusion} + +`ZMDiffusion` with `model_type="stable_diffusion_xl"` (read from the repo); `generate` +gains the micro-conditioning arguments: + +```python +generate( + input_ids, + attention_mask=None, + negative_input_ids=None, + num_inference_steps=None, + guidance_scale=None, + seed=None, + latents=None, + image=None, + strength=None, + denoising_start=None, + denoising_end=None, + output_type="image", + original_size=None, + crops_coords_top_left=(0, 0), + target_size=None, + aesthetic_score=6.0, + negative_original_size=None, + negative_crops_coords_top_left=None, + negative_target_size=None, + negative_aesthetic_score=2.5, +) +``` + +| Argument | Description | +|---|---| +| `input_ids`, `attention_mask` | `**tokenizer(prompts)`; the mask tells the model which positions are padding, so the second tower's `!`-padded ids can be derived (without it, everything after the first `<|endoftext|>` counts as padding). | +| `negative_input_ids` | Tokenized negative prompt, one row per prompt or a single row shared by the batch. Without it the unconditional branch is zero embeddings (`force_zeros_for_empty_prompt`), as in diffusers. | +| `num_inference_steps`, `guidance_scale`, `seed`, `latents` | As for [Stable Diffusion](#stable-diffusion-zmdiffusion); defaults from the repo's `generate_args` (50 steps, guidance 5.0 for the base model). | +| `image`, `strength`, `denoising_start`, `denoising_end`, `output_type` | The image-to-image and ensemble controls of [`BaseDiffusion.generate`](main_classes.md#basediffusion): refine an `image` (noised to `strength`), stop early (`denoising_end`) and hand the latent over (`output_type="latent"`), or resume one (`denoising_start`). | +| `original_size`, `crops_coords_top_left`, `target_size` | SDXL's size / crop conditioning, `(height, width)` in pixels; the sizes default to the generated image's size, the crop to `(0, 0)`. | +| `aesthetic_score` | The refiner's fifth conditioning value (6.0), in place of `target_size`; ignored by the base models. | +| `negative_*` | The same for the negative branch; each defaults to its positive value, except `negative_aesthetic_score` (2.5). | + +Returns `(batch, H, W, 3)` uint8 images (or the latent); `latents` is +`(batch, 128, 128, 4)` at 1024px (`channels_first`: channels second). + +The hooks, for anyone composing their own loop: `encode_prompt(input_ids, attention_mask, +original_size, crops_coords_top_left, target_size)` returns the dict the UNet takes +(`encoder_hidden_states` (batch, 77, 2048), `text_embeds` (batch, 1280), `time_ids` +(batch, 6)); `encode_negative_prompt` zeroes the first two when no negative ids are given; +`predict_noise(latents, timesteps, embeddings)` runs the UNet. + +##### The refiner {#stable-diffusion-xl-the-refiner} + +The same class with `model_type="stable_diffusion_xl_refiner"` (`text_encoder_2` alone, +`aesthetic_score` in the conditioning, `generate_args` with `strength` 0.3). Its +`generate` is the one above; with no `image` / `latents` it runs text-to-image, which the +refiner was not trained for. The two intended uses: + +```python +base = ZMDiffusion.from_weights("zeromodels/stable-diffusion-xl-base-1.0") +refiner = ZMDiffusion.from_weights("zeromodels/stable-diffusion-xl-refiner-1.0") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-xl-base-1.0") +inputs = tokenizer("a lighthouse on a rocky coast at dawn, dramatic clouds, cinematic") + +# ensemble of experts: the base denoises 80% of the schedule, the refiner the rest +latent = base.generate( + **inputs, num_inference_steps=50, denoising_end=0.8, output_type="latent" +) +images = refiner.generate( + **inputs, latents=latent, num_inference_steps=50, denoising_start=0.8 +) + +# image-to-image: refine any image (uint8 or [0, 1] float, (batch, H, W, 3)) +images = refiner.generate(**inputs, image=images, strength=0.3, seed=0) +``` + +##### The graph {#stable-diffusion-xl-the-graph} + +Four disconnected paths in one functional model, `.unet` / `.vae` / +`.text_encoder` / `.text_encoder_2`. + +| Inputs | Outputs | +|---|---| +| `sample` (B, 128, 128, 4), `timestep` (B,), `encoder_hidden_states` (B, 77, 2048), `text_embeds` (B, 1280), `time_ids` (B, 6) | `noise_pred` (B, 128, 128, 4) | +| `image` (B, 1024, 1024, 3), `latent` (B, 128, 128, 4) | `moments` (B, 128, 128, 8), `image` (B, 1024, 1024, 3) | +| `token_ids` (B, 77), `token_ids_2` (B, 77), `padding_mask` (B, 77) | `prompt_embeds` (B, 77, 2048), `pooled_prompt_embeds` (B, 1280) | + +Each text tower is a functional CLIP text model (the `clip` family's layers) whose outputs +are `penultimate_hidden_state`, `last_hidden_state`, `pooler_output` and, for the second, +`text_embeds`. The UNet is the SD 1.x `UNet2DConditionModel` with the SDXL config +(`transformer_layers_per_block=(1, 2, 10)`, `addition_embed_type="text_time"`); the VAE +is `AutoencoderKL` with `force_upcast=True`. The refiner is the same +graph without the first tower (`token_ids_2` + `padding_mask` in, `prompt_embeds` +(B, 77, 1280) out) and with `time_ids` (B, 5). + +### Preprocessing {#stable-diffusion-xl-preprocessing} + +#### `ZMDiffusionTokenizer` {#stable-diffusion-xl-zmdiffusiontokenizer} + +`ZMDiffusionTokenizer` with `model_type="stable_diffusion_xl"` (CLIP BPE, +`<|startoftext|>` / `<|endoftext|>` framing, +`<|endoftext|>`-padded to 77 tokens). Returns `{"input_ids", "attention_mask"}`; pass both +to `generate`. + +```python +ZMDiffusionTokenizer( + model_type="stable_diffusion_xl", hf_id=None, tokenizer_file=None, max_seq_len=77 +) # pad_token defaults to "<|endoftext|>"; the refiner repos use the same settings +``` + +A literal `!` in a prompt is encoded as OpenCLIP encodes it (the BPE token `!`, id +256) in the second tower's derived ids; diffusers' `tokenizer_2`, whose pad token is `!`, +turns it into the pad id 0 instead. + +### End-to-end example {#stable-diffusion-xl-end-to-end-example} + +```python +import os + +os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" + +from PIL import Image +from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer + +model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-xl-base-1.0") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-xl-base-1.0") + +inputs = tokenizer( + "a lighthouse on a rocky coast at dawn, dramatic clouds, cinematic, highly detailed" +) +images = model.generate(**inputs, num_inference_steps=30, guidance_scale=5.0, seed=3) + +Image.fromarray(images[0]).save("lighthouse.png") # (1024, 1024, 3) uint8 +``` + +Stable Diffusion XL base 1.0: a lighthouse on a rocky coast at dawn, 1024px + +#### Base + refiner {#stable-diffusion-xl-base-refiner} + +The SDXL "ensemble of experts": the base runs the first 80% of the schedule and hands its +latent to the refiner, which finishes it. Same prompt and initial latent as above: + +```python +refiner = ZMDiffusion.from_weights("zeromodels/stable-diffusion-xl-refiner-1.0") +latent = model.generate( + **inputs, num_inference_steps=30, denoising_end=0.8, output_type="latent", seed=3 +) +images = refiner.generate( + **inputs, latents=latent, num_inference_steps=30, denoising_start=0.8 +) +``` + +Stable Diffusion XL base + refiner: the lighthouse, refined + +#### SDXL-Turbo {#stable-diffusion-xl-sdxl-turbo} + +One step, no guidance, 512px (the repo's defaults): + +```python +model = ZMDiffusion.from_weights("zeromodels/sdxl-turbo") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/sdxl-turbo") +images = model.generate( + **tokenizer( + "a cinematic shot of a baby raccoon wearing an intricate italian priest robe" + ) +) +``` + +SDXL-Turbo, one step: a baby raccoon in an italian priest robe, 512px + +#### Negative prompts, micro-conditioning {#stable-diffusion-xl-negative-prompts-micro-conditioning} + +```python +inputs = tokenizer(["a red bicycle", "a blue bicycle"]) +negative = tokenizer("blurry, low quality") # one row, shared by the batch +images = model.generate( + **inputs, + negative_input_ids=negative["input_ids"], + original_size=(2048, 2048), # "a crop of a larger image": sharper detail + crops_coords_top_left=(0, 0), + target_size=(1024, 1024), + seed=0, +) +``` + +#### Other resolutions {#stable-diffusion-xl-other-resolutions} + +The graphs are built for the repo's resolution; pass overrides to build for another +(the weights are resolution-independent; multiples of 64px): + +```python +model = ZMDiffusion.from_weights( + "zeromodels/stable-diffusion-xl-base-1.0", + unet_sample_size=(96, 128), + vae_sample_size=(768, 1024), +) +images = model.generate(**tokenizer("a wide mountain valley")) # (1, 768, 1024, 3) +``` + +#### float32 {#stable-diffusion-xl-float32} + +```python +model = ZMDiffusion.from_weights( + "zeromodels/stable-diffusion-xl-base-1.0", load_dtype="float32" +) +``` + +#### Verified against diffusers {#stable-diffusion-xl-verified-against-diffusers} + +The components of `stable-diffusion-xl-base-1.0` and `stable-diffusion-xl-refiner-1.0` +were checked one by one against diffusers 0.39 / transformers in float32 on the same +inputs (the converted float16 weights computed in float32): + +``` +base text_encoder penultimate max|d|=6.1e-05 (values up to 854) +base text_encoder_2 penultimate max|d|=3.1e-05 +base text_encoder_2 pooled max|d|=3.6e-06 +base unet noise_pred (t=981) max|d|=9.3e-06 +base vae decode (512px) max|d|=3.5e-05 +refiner text_encoder_2 penultimate max|d|=3.1e-05 +refiner unet noise_pred (t=181) max|d|=1.2e-05 +``` + +The pipelines (Euler `leading` spacing, classifier-free guidance with the zeroed or +encoded negative branch, the micro-conditioning values, negative prompts, batches, +`strength`, `denoising_end` / `denoising_start`) were checked against +`StableDiffusionXLPipeline` and `StableDiffusionXLImg2ImgPipeline` on small +random-weight models, where the generated images are pixel-identical on torch and jax. +At full size in float16 (the default load; diffusers in float16 with its float32 VAE), +the lighthouse above matches diffusers to PSNR 49.7 dB (98.5% of pixels within 2 uint8 +levels), the base + refiner ensemble to PSNR 51.9 dB (98.8% within 2) and SDXL-Turbo +(1 step, 512px) to PSNR 44.9 dB (92% within 2; a single step amplifies float16 rounding). +The lighthouse, zeromodels (left) and diffusers (right): + +zeromodels (left) and diffusers (right) generations of the lighthouse from stable-diffusion-xl-base-1.0 with the same latent + +### Data Format {#stable-diffusion-xl-data-format} + +Both `channels_last` and `channels_first` are supported, as for +[Stable Diffusion](#stable-diffusion-data-format); `generate` always returns +`(batch, H, W, 3)` uint8. + +### Memory and speed {#stable-diffusion-xl-memory-and-speed} + +Measured on the torch backend with the default float16 load (weights resident: 6.8 GB for +a base model, 6.3 GB for the refiner, the VAE in float32): + +| Resolution | Denoising step (guided, batch 1) | VAE decode (float32) | +|---|---|---| +| 512px | peak 7.1 GB, 1.2 s/step on an RTX 4060 Laptop | peak 8.5 GB | +| 768px | peak 8.0 GB | peak 10.5 GB | +| 1024px | peak 10.4 GB | about 14 GB | + +Keras' functional executor keeps every layer output of a graph alive until the graph +returns, so the decoder's feature maps at 1024px add several GB on top of the weights: a +16 GB GPU runs 1024px comfortably, a 12 GB GPU 768px; on 8 GB, generate at 512px, or +denoise on the GPU (`output_type="latent"`) and decode the latent on the CPU (about 100 s +at 1024px) with a second, CPU-only process. In float32 the weights alone take 13.9 GB. +The UNet's attention runs the portable `"sdpa"` math by default, which materializes the +4096 x 4096 float32 logits of the 64 x 64 level at 1024px; loading with +`attn_implementation="fused"` (the backend's fused kernel, see +[`fused_attention`](main_classes.md#fused_attention)) cuts the 1024px step from a 10.3 GB +peak to 7.6 GB on the same GPU. On JAX and TensorFlow the denoiser step is compiled once +per run (see [`BaseDiffusion`](main_classes.md#basediffusion)). + +### Loading Fine-tuned Weights {#stable-diffusion-xl-loading-fine-tuned-weights} + +Any repo laid out like the hosted ones (`zm_config.json` declaring +`ZMDiffusion` with `model_type: "stable_diffusion_xl"` or +`"stable_diffusion_xl_refiner"`, the weights, `tokenizer.json`) loads with +`from_weights("/")`. The `hf:` prefix raises for diffusion models: convert a +diffusers-format SDXL checkpoint once with +`zeromodels/diffusers/stable_diffusion_xl/convert_stable_diffusion_xl_diffusers_to_keras.py` +(`transfer_stable_diffusion_xl(repo)`, `pip install zeromodels[conversion]`) and host the result. + +## Stable Diffusion 3 {#stable-diffusion-3} + +Stable Diffusion 3, ported to pure Keras 3: a rectified-flow model whose denoiser is the +**MMDiT**, a multimodal diffusion transformer in which the latent tokens and the text tokens +attend to each other jointly, conditioned on the timestep and the pooled text embeddings +through adaptive layer norms. It loads through [`ZMDiffusion`](#zmdiffusion) with +`model_type="stable_diffusion_3"` and keeps the SD family's shape (one functional model +and one set of weights per repo, `generate` from `BaseDiffusion`, the schedulers, both +data formats) with new parts: + +- **MMDiT** (`SD3Transformer2DModel`): the 16-channel latent is patchified (2x2) into 1536-d + tokens with a sinusoidal position table cropped from a 192x192 grid; 24 joint blocks of + 24 heads, each with AdaLN-Zero modulation of both streams, joint attention over the + concatenated latent + text tokens, and gated GELU feed-forwards; a final adaptive norm + and projection unpatchify the velocity. 2B parameters. +- **Three text encoders**: CLIP ViT-L/14 (768-d) and OpenCLIP ViT-bigG/14 (1280-d), each + with a projection, and **T5-XXL** (4.7B, 4096-d, 256 tokens). The two CLIP penultimate + states are concatenated (2048-d), zero-padded to 4096 and followed along the sequence by + the T5 states (77 + 256 tokens); the projected pooled CLIP states (2048-d) feed the + conditioning embedding. +- **T5 is optional and separate**: the checkpoint holds the MMDiT, the VAE and the two CLIP + towers (2.9B, 5.8 GB float16); the T5-XXL encoder is hosted once for every SD 3 / 3.5 + checkpoint (`zeromodels/t5-v1_1-xxl-encoder`, an + [`SD3T5EncoderModel`](#stable-diffusion-3-sd3t5encodermodel), 9.5 GB float16) and attached with + `text_encoder_3=`. Without it the T5 features are zeros, SD 3's documented + memory-saving mode (prompt adherence drops, the images stay good). +- **VAE**: 16 latent channels, no quant convolutions, `scaling_factor` 1.5305 and + `shift_factor` 0.0609 (`z = (x - shift) * scale`), built in float32 (`force_upcast`). +- **Sampler**: `FlowMatchEulerDiscreteScheduler` with `shift` 3.0 (rectified flow: the + timesteps are `sigma * 1000`, a step is `x + (sigma_next - sigma) * v`); 28 steps at + guidance 7.0 by default. No negative prompt means the encoded empty prompt. +- **One tokenizer call**: `ZMDiffusionTokenizer` runs the CLIP BPE (the second + tower's `!`-padded ids are derived from the mask, as SDXL) and the T5 SentencePiece + tokenizer, returning `{"input_ids", "attention_mask", "input_ids_3"}`. +- **float16 weights**, the release precision; `from_weights` builds the model in float16 + by default (`load_dtype="float32"` for float32). + +The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion is +deliberately **not supported** for diffusion models. The upstream repo is gated behind the +Stability AI Non-Commercial Research Community License. + +Links: + +- Paper: [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (arXiv:2403.03206)](https://arxiv.org/abs/2403.03206) +- Reference implementation: [diffusers `StableDiffusion3Pipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_3) +- License: [Stability AI Non-Commercial Research Community License](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers/blob/main/LICENSE.md) + +See also [Stable Diffusion 3.5](#stable-diffusion-3-5), [Stable Diffusion XL](#stable-diffusion-xl), [clip.md](clip.md). + +### Variants {#stable-diffusion-3-variants} + +| Variant | Hub | Resolution | Sampler | License | +|---|---|---|---|---| +| `stable-diffusion-3-medium` | [`zeromodels/stable-diffusion-3-medium`](https://huggingface.co/zeromodels/stable-diffusion-3-medium) | 1024 | flow-match Euler (shift 3), 28 steps, guidance 7.0 | Stability AI Non-Commercial Research Community | +| `t5-v1_1-xxl-encoder` | [`zeromodels/t5-v1_1-xxl-encoder`](https://huggingface.co/zeromodels/t5-v1_1-xxl-encoder) | text encoder (shared with SD 3.5) | | Apache 2.0 (Google T5 v1.1) | + +### API {#stable-diffusion-3-api} + +Configs are typed: `StableDiffusion3Config` (composite, `model_type` `"stable_diffusion_3"`) +over `StableDiffusion3TransformerConfig`, `StableDiffusion3VAEConfig`, +`StableDiffusion3TextConfig` (CLIP ViT-L/14 with its 768-d projection) and +`StableDiffusionXLTextConfig2` (OpenCLIP ViT-bigG/14 with its 1280-d projection), plus +the scheduler config, the CLIP and T5 token ids and `max_sequence_length` (256). Flat +constructor with `transformer_` / `vae_` / `text_` / `text_2_` prefixes. + +#### `ZMDiffusion` {#stable-diffusion-3-zmdiffusion} + +`ZMDiffusion` with `model_type="stable_diffusion_3"` (read from the repo); `generate` +over the SD 3 graph: + +```python +generate( + input_ids, + attention_mask=None, + negative_input_ids=None, + num_inference_steps=None, + guidance_scale=None, + seed=None, + latents=None, + image=None, + strength=None, + denoising_start=None, + denoising_end=None, + output_type="image", + input_ids_3=None, + negative_input_ids_3=None, +) +``` + +| Argument | Description | +|---|---| +| `input_ids`, `attention_mask`, `input_ids_3` | `**tokenizer(prompts)`: the CLIP ids and mask (the second tower's ids are derived from the mask) and the T5 ids (`(batch, 256)`). Without `input_ids_3`, or without an attached T5, the T5 features are zeros. | +| `negative_input_ids`, `negative_input_ids_3` | The tokenized negative prompt (`neg = tokenizer("..."); negative_input_ids=neg["input_ids"], negative_input_ids_3=neg["input_ids_3"]`); the empty prompt when omitted. | +| `num_inference_steps`, `guidance_scale`, `seed`, `latents`, `image`, `strength`, ... | As for [`BaseDiffusion.generate`](main_classes.md#basediffusion); defaults from the repo's `generate_args` (28 steps, guidance 7.0). `latents` is `(batch, 128, 128, 16)` at 1024px. | + +Returns `(batch, H, W, 3)` uint8 images. + +`from_weights` takes one extra argument, `text_encoder_3`: a hosted `SD3T5EncoderModel` +repo id (loaded at the same `load_dtype`) or a built `SD3T5EncoderModel`; +`sd.text_encoder_3` can also be assigned later (or set to `None`). It lives outside the +checkpoint's weights; passing `text_encoder_3` for a family other than SD 3 / 3.5 raises. + +The graph is four disconnected paths in one functional model, `.transformer` / `.vae` / +`.text_encoder` / `.text_encoder_2`. + +| Inputs | Outputs | +|---|---| +| `sample` (B, 128, 128, 16), `timestep` (B,), `encoder_hidden_states` (B, 333, 4096), `pooled_projections` (B, 2048) | `noise_pred` (B, 128, 128, 16) | +| `image` (B, 1024, 1024, 3), `latent` (B, 128, 128, 16) | `moments` (B, 128, 128, 32), `image` (B, 1024, 1024, 3) | +| `token_ids` (B, 77), `token_ids_2` (B, 77), `padding_mask` (B, 77) | `prompt_embeds` (B, 77, 2048), `pooled_prompt_embeds` (B, 2048) | + +#### `SD3Transformer2DModel` {#stable-diffusion-3-sd3transformer2dmodel} + +The MMDiT on its own (`{"sample", "timestep", "encoder_hidden_states", +"pooled_projections"} -> {"sample"}`), built from `StableDiffusion3TransformerConfig` +(`sample_size` 128, `patch_size` 2, `num_layers` 24, `num_attention_heads` 24 x +`attention_head_dim` 64, `joint_attention_dim` 4096, `caption_projection_dim` 1536, +`pooled_projection_dim` 2048, `pos_embed_max_size` 192, `qk_norm`, +`dual_attention_layers`). Its layers (`StableDiffusion3PatchEmbed`, +`StableDiffusion3AdaLayerNorm`, `StableDiffusion3JointAttention`, +`StableDiffusion3JointTransformerBlock`, `StableDiffusion3GELUFeedForward`) are named +with the diffusers module paths. + +The joint attention runs over 4429 tokens at 1024px, so `StableDiffusion3JointAttention` +defaults to the `"fused"` attention implementation (`keras.ops.dot_product_attention`, +the backend's own kernel: torch's flash / memory-efficient kernels never materialize the +4429 x 4429 logits; see [`fused_attention`](main_classes.md#fused_attention)) instead of +the library-wide `"sdpa"` math. `from_weights(..., attn_implementation="sdpa")` switches +it back. + +#### `SD3T5EncoderModel` {#stable-diffusion-3-sd3t5encodermodel} + +The third text encoder, the encoder half of T5 v1.1 XXL (transformers `T5EncoderModel` on +`google/t5-v1_1-xxl`), as its own model in this family: `{"input_ids", "attention_mask"}` +(`(batch, 256)` int32) `-> {"last_hidden_state": (batch, 256, 4096)}`, built from +`StableDiffusion3T5EncoderConfig` (`vocab_size` 32128, `embed_dim` 4096, `key_value_dim` +64, `mlp_dim` 10240, `num_layers` 24, `num_heads` 64, `relative_attention_num_buckets` 32, +`relative_attention_max_distance` 128, `layer_norm_eps` 1e-6; 4.7B parameters). The blocks +reuse the [T5](t5.md) family's attention, RMSNorm and relative position bias, and add the +gated-GELU feed-forward of T5 v1.1 (`wo(gelu_tanh(wi_0(x)) * wi_1(x))`, +`StableDiffusion3T5GatedFeedForward`), which the original T5 does not have. Loads on its own with +`SD3T5EncoderModel.from_weights("zeromodels/t5-v1_1-xxl-encoder")` (`load_dtype`, +`quantization="int8"` as needed); its inputs are the `input_ids_3` of the SD 3 tokenizer. + +### Preprocessing {#stable-diffusion-3-preprocessing} + +#### `ZMDiffusionTokenizer` {#stable-diffusion-3-zmdiffusiontokenizer} + +```python +ZMDiffusionTokenizer( + model_type="stable_diffusion_3", + hf_id=None, + tokenizer_file=None, + tokenizer_file_3=None, + max_seq_len=77, + max_sequence_length=256, +) +``` + +Returns `{"input_ids", "attention_mask", "input_ids_3"}`; pass all three to `generate`. + +### End-to-end example {#stable-diffusion-3-end-to-end-example} + +```python +import os + +os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" + +from PIL import Image +from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer + +model = ZMDiffusion.from_weights( + "zeromodels/stable-diffusion-3-medium", + text_encoder_3="zeromodels/t5-v1_1-xxl-encoder", # omit on small machines: T5 features zeroed +) +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-3-medium") + +inputs = tokenizer( + "a fluffy corgi wearing round sunglasses, sitting on a surfboard at a sunny beach, photo" +) +images = model.generate(**inputs, num_inference_steps=28, guidance_scale=7.0, seed=4) + +Image.fromarray(images[0]).save("corgi.png") # (1024, 1024, 3) uint8 +``` + +Stable Diffusion 3 medium: a corgi with sunglasses on a surfboard, 1024px + +#### The T5 encoder on the CPU {#stable-diffusion-3-the-t5-encoder-on-the-cpu} + +The 9.5 GB T5-XXL rarely fits next to the MMDiT on a consumer GPU. Encode the prompt +separately (an `SD3T5EncoderModel` in another process, or int8) and hand the ids to a +model without T5 for the CLIP part only, or attach it in full on a large machine. The +`encode_prompt` hook returns the conditioning dict the transformer takes, so any +precomputed T5 features can be substituted there. + +#### Verified against diffusers {#stable-diffusion-3-verified-against-diffusers} + +The components of `stable-diffusion-3-medium` were checked one by one against diffusers +0.39 / transformers 5 in float32 on the same inputs (the converted float16 weights +computed in float32): + +``` +text_encoder penultimate max|d|=6.1e-05 (values up to 854) +text_encoder_2 penultimate max|d|=3.1e-05 +pooled (L + G, projected) max|d|=3.8e-06 +t5-v1_1-xxl encoder (256 tok) max|d|=1.1e-04 (values up to 6.3) +mmdit velocity (t=1000) max|d|=9.9e-06 +vae decode (512px) max|d|=6.3e-05 +vae encode max|d|=1.4e-05 +``` + +The pipeline (flow-match Euler with shift 3, classifier-free guidance with the encoded +empty prompt, the three tokenizers, batched negative prompts, the no-T5 mode) was checked +against `StableDiffusion3Pipeline` on a small random-weight SD 3, where the generated images +are pixel-identical on torch and jax. End to end, `stable-diffusion-3-medium` without the +T5 encoder (28 steps, guidance 7, the same initial latent) reproduces the diffusers +pipeline in float32 at 512px: 92.8% identical pixels, max |d| 2 / 255, PSNR 63.9 dB (the +final latent differs by at most 2.5e-2 with values up to 5.1). In float16 on the GPU the +two implementations round differently through the 28 guided steps: the pictures agree in +composition but not pixel for pixel (PSNR 26 dB at 1024px, 29 dB at 768px, float16 MMDiT +on both sides), the usual float16 sampling noise. + +### Data Format {#stable-diffusion-3-data-format} + +Both `channels_last` and `channels_first` are supported; `generate` always returns +`(batch, H, W, 3)` uint8. + +### Memory and speed {#stable-diffusion-3-memory-and-speed} + +The checkpoint is 5.8 GB in float16 (the T5 encoder adds 9.5 GB). Measured on the torch +backend with the default float16 load (RTX 4060 Laptop, 8 GB), T5 features precomputed: + +| Resolution | Denoising step (guided, batch 1) | VAE decode (float32) | +|---|---|---| +| 512px | peak 6.1 GB, 0.5 s/step | peak 7.6 GB, 4 s | +| 768px | peak 6.4 GB, 0.8 s/step | peak 9.5 GB | +| 1024px | peak 6.9 GB, 1.0 s/step | peak 12.4 GB | + +The MMDiT itself fits an 8 GB GPU at the native 1024px thanks to the fused attention +kernel (the `"sdpa"` math needs 3.8 GB per block for the 4429 x 4429 float32 logits and +runs out of memory). The VAE decoder is the SDXL one (16 latent channels aside) and Keras' +functional executor keeps its feature maps alive, see the +[SDXL notes](#stable-diffusion-xl-memory-and-speed): on 8 GB, denoise on the GPU +(`output_type="latent"`) and decode on the CPU (75 to 90 s at 1024px). The T5-XXL encoder +encodes a prompt and its negative on the CPU in about 10 to 25 s (float32 compute over +the float16 weights); a 24 GB GPU holds everything. On JAX and TensorFlow the denoiser +step is compiled once per run. + +### Loading Fine-tuned Weights {#stable-diffusion-3-loading-fine-tuned-weights} + +Any repo laid out like the hosted ones (`zm_config.json` declaring +`ZMDiffusion` with `model_type: "stable_diffusion_3"`, the weights, the two tokenizer +files) loads with +`from_weights("/")`. The `hf:` prefix raises for diffusion models: convert a +diffusers-format SD 3 checkpoint once with +`zeromodels/diffusers/stable_diffusion_3/convert_stable_diffusion_3_diffusers_to_keras.py` +(`transfer_stable_diffusion_3(repo)`; `transfer_t5_encoder(repo)` for the T5 tower) and host the +result. + +## Stable Diffusion 3.5 {#stable-diffusion-3-5} + +Stable Diffusion 3.5, ported to pure Keras 3: the [Stable Diffusion 3](#stable-diffusion-3) +architecture, retrained and scaled. It loads through +[`ZMDiffusion`](#zmdiffusion) with `model_type="stable_diffusion_3_5"`, the SD 3 family +with the SD 3.5 configuration, so everything on that page applies: the MMDiT, the three text encoders with the separately hosted T5-XXL, the +16-channel VAE, the flow-match sampler, one tokenizer call. What changes is the config: + +- **QK normalization**: every attention block RMS-normalizes its queries and keys per head + (`qk_norm="rms_norm"`), the training stabilizer of SD 3.5. +- **Large** (8B): 38 blocks of 38 heads (2432-d tokens), 28 steps at guidance 3.5. +- **Large Turbo**: the large model distilled with Adversarial Diffusion Distillation to 4 + steps without guidance (`guidance_scale=0.0`). +- **Medium** (2.5B, MMDiT-X): 24 blocks of 24 heads with **dual attention** in blocks 0 to + 12 (a second, latent-only attention with its own modulation next to the joint one) and a + 384x384 position grid (up to 2 MP); 40 steps at guidance 4.5. + +The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion is +deliberately **not supported** for diffusion models. The upstream repos are gated behind +the Stability AI Community License. + +Links: + +- Paper: [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (arXiv:2403.03206)](https://arxiv.org/abs/2403.03206) +- Reference implementation: [diffusers `StableDiffusion3Pipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_3) +- License: [Stability AI Community License](https://huggingface.co/stabilityai/stable-diffusion-3.5-large/blob/main/LICENSE.md) + +See also [Stable Diffusion 3](#stable-diffusion-3). + +### Variants {#stable-diffusion-3-5-variants} + +Each repo is one checkpoint: MMDiT + VAE (84M) + CLIP ViT-L/14 (123M) + OpenCLIP +ViT-bigG/14 (695M); the T5-XXL encoder (`zeromodels/t5-v1_1-xxl-encoder`, the SD 3 +family's `SD3T5EncoderModel`, 9.5 GB) is shared and optional. + +| Variant | Hub | MMDiT | Sampler | Size | +|---|---|---|---|---| +| `stable-diffusion-3.5-large` | [`zeromodels/stable-diffusion-3.5-large`](https://huggingface.co/zeromodels/stable-diffusion-3.5-large) | 38 x 38 heads, 8B | flow-match Euler (shift 3), 28 steps, guidance 3.5 | 17.8 GB float16 | +| `stable-diffusion-3.5-large-turbo` | [`zeromodels/stable-diffusion-3.5-large-turbo`](https://huggingface.co/zeromodels/stable-diffusion-3.5-large-turbo) | 38 x 38 heads, 8B | 4 steps, no guidance | 17.8 GB float16 | +| `stable-diffusion-3.5-medium` | [`zeromodels/stable-diffusion-3.5-medium`](https://huggingface.co/zeromodels/stable-diffusion-3.5-medium) | 24 x 24 heads + dual attention, 2.5B | flow-match Euler (shift 3), 40 steps, guidance 4.5 | 6.8 GB float16 | + +### API {#stable-diffusion-3-5-api} + +`StableDiffusion3_5Config` (`model_type` `"stable_diffusion_3_5"`) is +`StableDiffusion3Config` over `StableDiffusion3_5TransformerConfig` (the SD 3.5 large +defaults: `num_layers` 38, `num_attention_heads` 38, `caption_projection_dim` 2432, +`qk_norm="rms_norm"`; medium's repo carries `dual_attention_layers=(0, ..., 12)` and +`pos_embed_max_size=384`). `ZMDiffusion.generate` and +`ZMDiffusionTokenizer` (`model_type="stable_diffusion_3_5"`) behave as for SD 3; see +[Stable Diffusion 3](#stable-diffusion-3-api). + +### End-to-end example {#stable-diffusion-3-5-end-to-end-example} + +```python +import os + +os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" + +from PIL import Image +from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer + +model = ZMDiffusion.from_weights( + "zeromodels/stable-diffusion-3.5-medium", + text_encoder_3="zeromodels/t5-v1_1-xxl-encoder", # optional +) +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-3.5-medium") + +inputs = tokenizer( + "a whimsical treehouse village at twilight, lanterns glowing, watercolor illustration" +) +images = model.generate(**inputs, num_inference_steps=40, guidance_scale=4.5, seed=5) + +Image.fromarray(images[0]).save("treehouse.png") # (1024, 1024, 3) uint8 +``` + +Stable Diffusion 3.5 medium: a treehouse village at twilight, 1024px + +#### Large Turbo {#stable-diffusion-3-5-large-turbo} + +```python +model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-3.5-large-turbo") +images = model.generate( + **tokenizer("a red panda reading a book") +) # 4 steps, guidance 0 (the repo defaults) +``` + +#### Verified against diffusers {#stable-diffusion-3-5-verified-against-diffusers} + +The components of `stable-diffusion-3.5-medium` (the MMDiT-X with dual attention and RMS +qk norms) were checked against diffusers 0.39 in float32 on the same inputs (the converted +float16 weights computed in float32): MMDiT velocity max|d| 1.6e-5 (values up to 5.2), the +text encoders and the VAE as for [SD 3](#stable-diffusion-3-verified-against-diffusers). +The SD 3.5 block variants (dual attention, `qk_norm`) are also covered by the small +random-weight pipeline test, pixel-identical to `StableDiffusion3Pipeline` on torch and +jax, and by the SD 3 float32 end-to-end check (92.8% identical pixels, PSNR 63.9 dB at +512px), which the two families share down to the sampler. The picture above is the +1024px float16 output of `stable-diffusion-3.5-medium` with the T5 encoder. + +### Memory and speed {#stable-diffusion-3-5-memory-and-speed} + +The medium checkpoint is 6.8 GB in float16, the large ones 17.8 GB (a 24 GB GPU); the +T5-XXL encoder adds 9.5 GB. See [Stable Diffusion 3](#stable-diffusion-3-memory-and-speed): +the medium MMDiT-X at 1024px peaks at 7.5 GB during the guided denoising step on the torch +backend (fused attention, float16), which an 8 GB GPU runs with the VAE decode on the CPU. + +### Loading Fine-tuned Weights {#stable-diffusion-3-5-loading-fine-tuned-weights} + +Any repo laid out like the hosted ones (`zm_config.json` declaring +`ZMDiffusion` with `model_type: "stable_diffusion_3_5"`, the weights, the two tokenizer +files) loads with +`from_weights("/")`. Convert a diffusers-format checkpoint once with +`zeromodels/diffusers/stable_diffusion_3_5/convert_stable_diffusion_3_5_diffusers_to_keras.py` +and host the result. + +## Qwen-Image {#qwen-image} + +Qwen-Image, ported to pure Keras 3: latent text-to-image flow-matching with a +60-layer double-stream MMDiT denoiser, a Wan-derived KL autoencoder and the +Qwen2.5-VL-7B Instruct text tower. It loads through [`ZMDiffusion`](#zmdiffusion), the one diffusion +pipeline class of zeromodels, with `model_type="qwen_image"`: a single functional graph with the three towers as disconnected +sub-graphs (the way `CLIPModel` holds its two), so a hosted repo is one set of +sharded weights and one `zm_config.json`. plus +`generate`: the true-CFG denoising loop of `BaseDiffusion`, the +diffusion counterpart of `BaseGeneration`. + +The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion +is deliberately **not supported** for diffusion models. + +Key facts of the port: + +- **Packed latents**: the denoiser sees `(B, H/2 · W/2, 64)` tokens; pack / unpack + match Diffusers' 2×2 patching around the 16-channel VAE latent grid + (`AutoencoderKLQwenImage`, 8× spatial, channels-last NDHWC with `T=1`). +- **True CFG**: when `guidance_scale > 1`, two transformer forwards and Diffusers' + prediction-norm renormalization (`true_cfg_scale`); a space negative prompt is + the default when none is passed. +- **bf16-safe dual-stream clip**: Diffusers clips dual-stream activations to the + fp16 finite range only under float16; the port matches that, so bf16 runs keep + the ~1e7 magnitudes Qwen-Image produces instead of collapsing to noise. +- **Schedulers match Diffusers**: `FlowMatchEulerDiscreteScheduler` with dynamic + resolution shifting (`mu` from packed sequence length) so timesteps agree with + the reference within float rounding. +- **Verified against Diffusers**: transformer, text encoder, VAE and scheduler + outputs agree with the reference within float noise (see + [Verified against Diffusers](#qwen-image-verified-against-diffusers)). + +Links: + +- Paper: [Qwen-Image Technical Report (arXiv:2508.02324)](https://arxiv.org/abs/2508.02324) +- Reference implementation: [diffusers `QwenImagePipeline`](https://huggingface.co/docs/diffusers/api/pipelines/qwenimage) +- License: [Apache-2.0](https://huggingface.co/Qwen/Qwen-Image/blob/main/LICENSE) + +See also [qwen2_5_vl.md](qwen2_5_vl.md) (the text tower), +[Qwen-Image-2.1](#qwen-image-2-1) (2.1 single-stream / unpatched), +[Stable Diffusion 3](#stable-diffusion-3) (flow-match / MMDiT-style diffusion +in ZeroModels). + +### Variants {#qwen-image-variants} + +Preconverted, bfloat16 weights are hosted under `zeromodels/`. Load with +`from_weights("zeromodels/")`. Each repo is one checkpoint: MMDiT +transformer + VAE + Qwen2.5-VL text tower, ~53 GiB at 16-bit. Released under +Apache-2.0. + +| Variant | Hub | Source | +|---|---|---| +| `qwen-image` | [`zeromodels/qwen-image`](https://huggingface.co/zeromodels/qwen-image) | [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) | + +Default `generate_args`: 50 flow-match steps, `guidance_scale=4.0`, 1024×1024. + +### API {#qwen-image-api} + +Configs are typed: `QwenImageConfig` (composite, `model_type` `"qwen_image"`) over +`QwenImageTransformerConfig`, `QwenImageVAEConfig` and `QwenImageTextConfig`, plus +the checkpoint's `scheduler_config` and the Qwen special token ids. Each repo's +`zm_config.json` parses through it; the constructor stays flat, with the +sub-config fields prefixed `transformer_` / `vae_` / `text_`. + +#### `ZMDiffusion` {#qwen-image-zmdiffusion} + +The pipeline class of every diffusion family (see [Diffusers](#zmdiffusion)); the +repo's `model_type` (`"qwen_image"` here) selects this family, which supplies the hooks +`BaseDiffusion`'s `generate` needs the mixin needs (`encode_prompt` on the text +tower with ChatML template drop, `unconditional_ids`, `predict_noise` on the +transformer, `decode_latents` on the VAE with mean/std un-normalization) and a +`scheduler`, built from the config's `scheduler_config` +(`FlowMatchEulerDiscreteScheduler` with dynamic shifting). + +```python +generate( + input_ids, + attention_mask=None, + negative_input_ids=None, + negative_attention_mask=None, + num_inference_steps=None, + guidance_scale=None, + seed=None, + latents=None, + height=None, + width=None, + output_type="image", + **conditioning, +) +``` + +| Arg | Default | Meaning | +|---|---|---| +| `input_ids` | required | ChatML-templated token ids, `**tokenizer(prompts)` | +| `attention_mask` | `None` | padding mask from the tokenizer | +| `negative_input_ids` | `None` | tokenized negative prompt for true CFG; a templated space prompt when unset | +| `negative_attention_mask` | `None` | mask for the negative ids | +| `negative_prompt` | `None` | optional text (or one string per image) passed via `**conditioning`; tokenized when negative ids are unset | +| `num_inference_steps` | `None` | scheduler steps; the repo's `generate_args` (50) when unset | +| `guidance_scale` | `None` | true CFG strength; `generate_args` (4.0) when unset, `<= 1` disables it | +| `seed` | `None` | seed for the initial latent, reproducible per backend | +| `latents` | `None` | explicit packed initial latent; identical results across backends | +| `height` / `width` | `None` | output pixel size; `default_sample_size * 8` (1024) when unset; rounded to VAE×pack multiples | +| `output_type` | `"image"` | `"image"` for uint8 RGB, `"latent"` for packed latents | + +Returns `(batch, height, width, 3)` uint8 numpy images. + +| Constructor arg | Default | Meaning | +|---|---|---| +| `model_type` | required | `"qwen_image"` (read from the repo's `zm_config.json` on `from_weights`) | +| `scheduler` | `None` | a `BaseScheduler`; built from `scheduler_config` when unset | +| `transformer_sample_size` | `128` | latent side the transformer graph is built for (image / 8) | +| `vae_sample_size` | `1024` | image size the VAE graphs are built for | +| `transformer_*` / `vae_*` / `text_*` | Qwen-Image | the flat sub-config fields (see [Configuration](configuration.md)) | +| `bos_token_id` / `eos_token_id` / `pad_token_id` | `151643` / `151645` / `151643` | Qwen special tokens | +| `scheduler_config` | `None` | the Diffusers scheduler dict (`_class_name`, shift, dynamic shifting, ...) | +| `prompt_template_encode_start_idx` | `34` | ChatML template tokens dropped after the text encoder | +| `max_sequence_length` | `512` | prompt embed length after the template drop | + +The graph is three disconnected paths, one per component. Inputs cover the transformer (`sample`, `timestep`, +`encoder_hidden_states`, `encoder_hidden_states_mask`), the VAE (`image`, `latent`) +and the text tower (`token_ids`, `padding_mask`); outputs include the packed +velocity, VAE moments / sample, and prompt embeds. The components are exposed as +`.transformer`, `.vae` and `.text_encoder` and share the model's weights, so one load +fills all three, and each can be called on its own. + +#### `QwenImageTransformer2DModel` {#qwen-image-qwenimagetransformer2dmodel} + +The denoiser (Diffusers' `QwenImageTransformer2DModel`): a 60-layer double-stream +MMDiT over **packed** latents and text features, with MS-RoPE +(`axes_dims_rope=(16, 56, 56)`). Inputs +`{"sample": (B, seq, 64), "timestep": (B,), "encoder_hidden_states": (B, text_seq, 3584)}` +(optional `encoder_hidden_states_mask`), output `{"sample": (B, seq, 64)}` packed +velocity. + +| Arg | Default | Meaning | +|---|---|---| +| `patch_size` | `2` | latent pack side | +| `in_channels` | `64` | packed token width (`16 × 2 × 2`) | +| `out_channels` | `16` | unpacked latent channels | +| `num_layers` | `60` | dual-stream DiT blocks | +| `attention_head_dim` / `num_attention_heads` | `128` / `24` | head geometry (inner dim 3072) | +| `joint_attention_dim` | `3584` | text feature width from Qwen2.5-VL | +| `axes_dims_rope` | `(16, 56, 56)` | MS-RoPE axis splits | +| `sample_size` | `128` | latent spatial side the graph is built for | +| `text_seq_len` | `512` | static text length after the template drop | + +#### `AutoencoderKLQwenImage` {#qwen-image-autoencoderklqwenimage} + +The VAE (Diffusers' `AutoencoderKLQwenImage`): `encode(image, sample=False)` returns +the posterior mean (or a sample) of a `(B, H/8, W/8, 16)` latent and +`decode(latent)` the `(B, H, W, 3)` image in `[-1, 1]`. Latent normalisation uses +`latents_mean` / `latents_std` around the denoiser, as `generate` does. One +functional graph with the encoder and decoder as two disconnected paths +(channels-last NDHWC, `T=1` for still images). + +| Arg | Default | Meaning | +|---|---|---| +| `base_dim` | `96` | base channel width | +| `z_dim` | `16` | latent channels | +| `dim_mult` | `(1, 2, 4, 4)` | width multipliers per level | +| `num_res_blocks` | `2` | residual blocks per level | +| `temperal_downsample` | `(False, True, True)` | which levels downsample time (unused at `T=1`) | +| `sample_size` | `1024` | image size the graphs are built for | +| `latents_mean` / `latents_std` | Qwen-Image | per-channel latent normalisation | + +#### Schedulers {#qwen-image-schedulers} + +`zeromodels.base.base_scheduler` has the flow-match sampler with the Diffusers +interface (`set_timesteps`, `step`, `init_noise_sigma`), configured the Diffusers +way: + +| Scheduler | Notes | +|---|---| +| `FlowMatchEulerDiscreteScheduler` | Qwen-Image default; `use_dynamic_shifting=True`, exponential time shift | + +`get_scheduler(config)` builds the one named by a Diffusers scheduler config dict, +which is what the model does with its `scheduler_config`. At generate time the +task sets timesteps with resolution-dependent `mu` and a linspace sigma schedule, +matching Diffusers. + +### Preprocessing {#qwen-image-preprocessing} + +#### `ZMDiffusionTokenizer` {#qwen-image-zmdiffusiontokenizer} + +The Qwen2 BPE tokenizer with the Diffusers ChatML prompt template wrapped around +each string. Calling it returns `{"input_ids", "attention_mask"}` for a string or +a list of strings. `generate` / `encode_prompt` drop the template prefix +(`prompt_template_encode_start_idx=34`) after the text encoder and keep at most +`max_sequence_length` (512) tokens. + +```python +ZMDiffusionTokenizer( + model_type="qwen_image", hf_id=None, tokenizer_file=None, max_seq_len=1024 +) +``` + +There is no image processor and no processor class: text-to-image takes token ids +in and hands uint8 images out. + +### End-to-end example {#qwen-image-end-to-end-example} + +#### Single prompt {#qwen-image-single-prompt} + +```python +import os + +os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" + +from PIL import Image +from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer + +model = ZMDiffusion.from_weights("zeromodels/qwen-image") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/qwen-image") + +inputs = tokenizer( + "a coffee shop entrance with a chalkboard sign reading 'Open', warm afternoon light" +) +images = model.generate( + **inputs, + height=1024, + width=1024, + num_inference_steps=50, + guidance_scale=4.0, + seed=1, +) + +Image.fromarray(images[0]).save("coffee.png") # (1024, 1024, 3) uint8 +``` + +#### Negative prompt and batching {#qwen-image-negative-prompt-and-batching} + +Prompts batch through the tokenizer; a negative prompt is tokenized the same way +and passed as `negative_input_ids` (one per prompt, or one broadcast row). + +```python +prompts = [ + "a watercolor painting of a lighthouse at sunset", + "a red vintage car on a coastal road", +] +inputs = tokenizer(prompts) +negative = tokenizer(["blurry, low quality"] * len(prompts))["input_ids"] + +images = model.generate( + **inputs, + negative_input_ids=negative, + height=512, + width=512, + num_inference_steps=30, +) +for i, image in enumerate(images): + Image.fromarray(image).save(f"out_{i}.png") +``` + +#### Reproducible across backends {#qwen-image-reproducible-across-backends} + +`seed` reproduces a run on a given backend, but `keras.random` differs between +torch, jax and tensorflow. For an image that is identical everywhere, pass the +initial packed latent yourself. At 1024px the spatial noise before packing is +`(batch, 128, 128, 16)`; after packing it is `(batch, 4096, 64)`: + +```python +import numpy as np + +from zeromodels.diffusers.qwen_image.qwen_image_model import pack_latents + +h, w, channels = 128, 128, 16 +noise = np.random.default_rng(0).standard_normal((1, h, w, channels)).astype("float32") +latents = pack_latents(noise, h, w) +images = model.generate(**tokenizer("a bowl of ramen"), latents=latents) +``` + +#### Verified against Diffusers {#qwen-image-verified-against-diffusers} + +Component checks against Diffusers / Transformers with the same weights (fp32 +unless noted): + +``` +scheduler timesteps / mu within 1e-4 / identical +transformer (fp32, N blocks) relative error ~8e-6 +transformer (bf16, N blocks) cosine similarity 0.99998 +text encoder (fp32) relative error ~1.7e-6 +VAE encode / decode matches Diffusers (NCDHW <-> NDHWC transpose) +``` + +End-to-end image parity depends on dtype and resolution; prefer bf16 (or fp32) +over fp16 — Qwen-Image activations overflow the fp16 range without the Diffusers +clip path. + +#### Components on their own {#qwen-image-components-on-their-own} + +```python +from zeromodels.diffusers import ZMDiffusion + +qi = ZMDiffusion.from_weights("zeromodels/qwen-image") +latent = qi.vae.encode(image) # (1, H, W, 3) in [-1, 1] -> (1, H/8, W/8, 16) +text = qi.text_encoder({"input_ids": ids, "attention_mask": mask})["last_hidden_state"] +``` + +#### Other resolutions {#qwen-image-other-resolutions} + +The graphs are built for a fixed size, the weights are not. Rebuild at another +size with the constructor overrides (multiples of 16px after VAE×pack; the +checkpoint targets 1024px): + +```python +model = ZMDiffusion.from_weights( + "zeromodels/qwen-image", + transformer_sample_size=64, + vae_sample_size=512, +) +images = model.generate( + **tokenizer("a mountain lake at dawn"), height=512, width=512 +) # (1, 512, 512, 3) +``` + +### Data Format {#qwen-image-data-format} + +**Channels-last only for the VAE.** `AutoencoderKLQwenImage` is built as NDHWC +`Conv3D` (with `T=1` for still images); the transformer works on packed sequence +tokens and has no spatial layout. `generate` always hands back +`(batch, height, width, 3)` uint8 images. + +| | Shape | +|---|---| +| `latents` passed to `generate` (packed) | `(batch, (H/16)·(W/16), 64)` | +| VAE encode input / decode output | `(batch, H, W, 3)` | +| VAE latent | `(batch, H/8, W/8, 16)` | +| Transformer `sample` | `(batch, seq, 64)` | + +### Memory and speed {#qwen-image-memory-and-speed} + +The bf16 / float16 checkpoint is about 53 GiB. Building the full graph on a ~40 GB +GPU OOMs; build on CPU (or load with component offload) and move one tower at a +time. A guided 1024px run needs substantial activation headroom on top of the +~38 GiB transformer; 512px is the practical default on 40 GB cards. Prefer the +fused attention path (`keras.ops.dot_product_attention` / torch SDPA) over the +plain matmul softmax default. + +### Loading Fine-tuned Weights {#qwen-image-loading-fine-tuned-weights} + +The hosted checkpoint is the supported weight; any repo laid out like it +(`zm_config.json` declaring `ZMDiffusion` with `model_type: "qwen_image"`, sharded `*.weights.json` / +`*.weights.h5`, `tokenizer.json`) loads with `from_weights("/")`. The +`hf:` prefix raises for diffusion models: convert a Diffusers-format checkpoint +once with +`zeromodels/diffusers/qwen_image/convert_qwen_image_diffusers_to_keras.py` +(`transfer_qwen_image(repo)`, `pip install zeromodels[conversion]`) and host the +result. + +## Qwen-Image-2.1 {#qwen-image-2-1} + +Qwen-Image-2.1, ported to pure Keras 3: latent text-to-image flow-matching with a +32-layer **single-stream** block-causal DiT, a residual 64-channel KL autoencoder +(16× spatial, RGBA), and the Qwen3-VL text tower. It loads through +[`ZMDiffusion`](#zmdiffusion), the one diffusion pipeline class of zeromodels, with +`model_type="qwen_image_21"`: one functional graph holding the three towers, plus +`generate`. Prompts are tokenized by `ZMDiffusionTokenizer`. + +The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion +is deliberately **not supported** for diffusion models. + +Key facts of the port: + +- **Unpatched latents**: the denoiser sees `(B, H · W, 64)` tokens (VAE scale 16; + no 2×2 packing). At 1024px that is a `(B, 4096, 64)` sequence. +- **Block-causal attention**: text is causal; the target image block is + bidirectional and can attend to all preceding text. +- **`causal_condition`**: text tokens modulate from `t = 0` (timestep-independent), + matching Diffusers' KV-cache-ready conditioning. +- **True CFG optional**: Diffusers defaults to `true_cfg_scale=1.0` (no guidance). + Pass `guidance_scale > 1` with a negative prompt to enable dual forwards. +- **Pre-norm text features**: the text tower returns decoder outputs *before* the + final RMSNorm, matching Diffusers' forward hook on the language-model norm. +- **Schedulers match Diffusers**: `FlowMatchEulerDiscreteScheduler` with dynamic + resolution shifting (`mu` from image sequence length) and `shift_terminal`. + +Links: + +- Source: [`Qwen/Qwen-Image-2.1`](https://huggingface.co/Qwen/Qwen-Image-2.1) +- Reference: [diffusers `QwenImage21Pipeline`](https://huggingface.co/docs/diffusers/api/pipelines/qwenimage) +- License: [Qwen Research License](https://huggingface.co/Qwen/Qwen-Image-2.1) +- See also [Qwen-Image](#qwen-image) (1.0 double-stream / packed), + [qwen3_vl.md](qwen3_vl.md) (text tower) + +### Variants {#qwen-image-2-1-variants} + +Preconverted, bfloat16 weights are hosted under `zeromodels/`. Load with +`from_weights("zeromodels/")`. Each repo is one checkpoint: DiT + VAE + +Qwen3-VL text tower, ~28 GiB at 16-bit. + +| Variant | Hub | Source | +|---|---|---| +| `qwen-image-2.1` | [`zeromodels/qwen-image-2.1`](https://huggingface.co/zeromodels/qwen-image-2.1) | [`Qwen/Qwen-Image-2.1`](https://huggingface.co/Qwen/Qwen-Image-2.1) | + +Default `generate_args`: 40 flow-match steps, `guidance_scale=1.0`, 1024×1024. + +### API {#qwen-image-2-1-api} + +Configs are typed: `QwenImage21Config` (composite, `model_type` `"qwen_image_21"`) +over `QwenImage21TransformerConfig`, `QwenImage21VAEConfig` and +`QwenImage21TextConfig`, plus the checkpoint's `scheduler_config` and Qwen special +token ids. Each repo's `zm_config.json` parses through it; the constructor stays +flat, with the sub-config fields prefixed `transformer_` / `vae_` / `text_`. + +#### `ZMDiffusion` {#qwen-image-2-1-zmdiffusion} + +The pipeline class of every diffusion family (see [Diffusers](#zmdiffusion)); the +repo's `model_type` (`"qwen_image_21"` here) selects this family, which supplies the hooks +`BaseDiffusion`'s `generate` needs: `encode_prompt` (ChatML template drop after the text +encoder), `predict_noise` on the transformer, `decode_latents` on the VAE with +mean/std un-normalization, and a `FlowMatchEulerDiscreteScheduler` from +`scheduler_config`. + +```python +generate( + input_ids, + attention_mask=None, + negative_input_ids=None, + negative_attention_mask=None, + num_inference_steps=None, + guidance_scale=None, + seed=None, + latents=None, + height=None, + width=None, + output_type="image", +) +``` + +| Arg | Default | Meaning | +|---|---|---| +| `input_ids` | required | ChatML-templated token ids, `**tokenizer(prompts)` | +| `attention_mask` | `None` | padding mask from the tokenizer | +| `negative_input_ids` | `None` | tokenized negative prompt; CFG only when set **and** `guidance_scale > 1` | +| `negative_attention_mask` | `None` | mask for the negative ids | +| `num_inference_steps` | `None` | scheduler steps; `generate_args` (40) when unset | +| `guidance_scale` | `None` | true CFG strength; `generate_args` (1.0) when unset | +| `seed` | `None` | seed for the initial latent | +| `latents` | `None` | explicit unpatched initial latent `(B, H·W, 64)` | +| `height` / `width` | `None` | output pixel size; `default_sample_size * 16` (1024) when unset | +| `output_type` | `"image"` | `"image"` for uint8 RGB, `"latent"` for sequence latents | + +Returns `(batch, height, width, 3)` uint8 numpy images (VAE RGBA is cropped to RGB). + +| Constructor arg | Default | Meaning | +|---|---|---| +| `model_type` | required | `"qwen_image_21"` (read from the repo's `zm_config.json` on `from_weights`) | +| `scheduler` | `None` | a `BaseScheduler`; built from `scheduler_config` when unset | +| `transformer_sample_size` | `64` | latent side the DiT graph is built for (image / 16) | +| `vae_sample_size` | `1024` | image size the VAE graphs are built for | +| `transformer_*` / `vae_*` / `text_*` | Qwen-Image-2.1 | flat sub-config fields (see [Configuration](configuration.md)) | +| `bos_token_id` / `eos_token_id` / `pad_token_id` | `151643` / `151645` / `151643` | Qwen special tokens | +| `scheduler_config` | `None` | Diffusers scheduler dict (dynamic shifting, `max_shift=0.9`, ...) | +| `prompt_template_encode_start_idx` | `14` | ChatML system-prefix tokens dropped after the text encoder | +| `max_sequence_length` | `512` | prompt embed length after the template drop | + +The graph is three disconnected paths, one per component. +Inputs cover the transformer (`sample`, `timestep`, `encoder_hidden_states`, +`encoder_hidden_states_mask`), the VAE (`image`, `latent`) and the text tower +(`token_ids`, `padding_mask`). Components are exposed as `.transformer`, `.vae` +and `.text_encoder`. + +#### `QwenImage21Transformer2DModel` {#qwen-image-2-1-qwenimage21transformer2dmodel} + +The denoiser (Diffusers' `QwenImage21Transformer2DModel`): a 32-layer +single-stream DiT over **unpatched** latents and Qwen3-VL text features, with +3-axis RoPE (`axes_dims_rope=(16, 56, 56)`). Inputs +`{"sample": (B, seq, 64), "timestep": (B,), "encoder_hidden_states": (B, text_seq, 4096)}` +(plus mask), output `{"sample": (B, seq, 64)}` target velocity. + +| Arg | Default | Meaning | +|---|---|---| +| `patch_size` | `1` | unpatched (asserted) | +| `in_channels` / `out_channels` | `64` / `64` | latent token width | +| `num_layers` | `32` | single-stream blocks | +| `attention_head_dim` / `num_attention_heads` | `128` / `32` | head geometry (inner dim 4096) | +| `context_in_dim` | `4096` | text feature width from Qwen3-VL | +| `mlp_ratio` | `3` | SwiGLU expansion | +| `causal_condition` | `True` | text modulates from `t=0` | +| `sample_size` | `64` | latent spatial side the graph is built for | +| `text_seq_len` | `512` | static text length after the template drop | + +#### `AutoencoderKLQwenImage21` {#qwen-image-2-1-autoencoderklqwenimage21} + +The VAE (Diffusers' `AutoencoderKLQwenImage21`): residual encoder/decoder, +`encode(image)` → `(B, H/16, W/16, 64)` latents, `decode(latent)` → +`(B, H, W, 4)` RGBA in `[-1, 1]`. Latent normalisation uses `latents_mean` / +`latents_std` around the denoiser. + +| Arg | Default | Meaning | +|---|---|---| +| `base_dim` / `decoder_base_dim` | `96` / `144` | channel widths | +| `z_dim` | `64` | latent channels | +| `dim_mult` | `(1, 2, 4, 8, 8)` | width multipliers per level | +| `scale_factor_spatial` | `16` | pixel → latent downscale | +| `input_channels` / `out_channels` | `4` / `4` | RGBA | +| `is_residual` | `True` | residual blocks | +| `sample_size` | `1024` | image size the graphs are built for | +| `latents_mean` / `latents_std` | Qwen-Image-2.1 | per-channel latent normalisation | + +#### `QwenImage21TextEncoderModel` {#qwen-image-2-1-qwenimage21textencodermodel} + +Qwen3-VL text tower only (no vision / LM head). Returns +`{"last_hidden_state"}` as **pre-final-norm** hidden states. + +#### Schedulers {#qwen-image-2-1-schedulers} + +| Scheduler | Notes | +|---|---| +| `FlowMatchEulerDiscreteScheduler` | `use_dynamic_shifting=True`, exponential time shift, `shift_terminal=0.02`, `max_shift=0.9` | + +At generate time the task sets timesteps with resolution-dependent `mu` and a +linspace sigma schedule, matching Diffusers. + +### Preprocessing {#qwen-image-2-1-preprocessing} + +#### `ZMDiffusionTokenizer` {#qwen-image-2-1-zmdiffusiontokenizer} + +Qwen3 BPE with Diffusers' T2I ChatML template (`Comprehend and analyze the +provided prompt.`). Returns `{"input_ids", "attention_mask"}`. `encode_prompt` +drops `prompt_template_encode_start_idx` (14) tokens and keeps at most +`max_sequence_length` (512). + +```python +ZMDiffusionTokenizer( + model_type="qwen_image_21", hf_id=None, tokenizer_file=None, max_seq_len=1024 +) +``` + +### End-to-end example {#qwen-image-2-1-end-to-end-example} + +#### Single prompt {#qwen-image-2-1-single-prompt} + +```python +import os + +os.environ["KERAS_BACKEND"] = "torch" + +from PIL import Image +from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer + +model = ZMDiffusion.from_weights("zeromodels/qwen-image-2.1") +tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/qwen-image-2.1") + +inputs = tokenizer("a photo of a capybara wearing a wizard hat, soft window light") +images = model.generate( + **inputs, + height=1024, + width=1024, + num_inference_steps=40, + guidance_scale=1.0, + seed=0, +) + +Image.fromarray(images[0]).save("capybara.png") # (1024, 1024, 3) uint8 +``` + +#### Optional true CFG {#qwen-image-2-1-optional-true-cfg} + +```python +inputs = tokenizer("a watercolor lighthouse at sunset") +negative = tokenizer("blurry, low quality")["input_ids"] + +images = model.generate( + **inputs, + negative_input_ids=negative, + guidance_scale=4.0, + height=1024, + width=1024, +) +``` + +#### Reproducible latents {#qwen-image-2-1-reproducible-latents} + +At 1024px the unpatched noise is `(batch, 64, 64, 64)` → `(batch, 4096, 64)`: + +```python +import numpy as np +from zeromodels.diffusers.qwen_image_21.qwen_image_21_model import pack_latents + +h, w, channels = 64, 64, 64 +noise = np.random.default_rng(0).standard_normal((1, h, w, channels)).astype("float32") +latents = pack_latents(noise, h, w) +images = model.generate(**tokenizer("a bowl of ramen"), latents=latents) +``` + +#### Other resolutions {#qwen-image-2-1-other-resolutions} + +Graphs are built for a fixed size; weights are not. Rebuild with constructor +overrides (multiples of 32px; checkpoint targets 1024): + +```python +model = ZMDiffusion.from_weights( + "zeromodels/qwen-image-2.1", + transformer_sample_size=32, + vae_sample_size=512, +) +images = model.generate(**tokenizer("a mountain lake at dawn"), height=512, width=512) +``` + +#### Offline conversion {#qwen-image-2-1-offline-conversion} + +```bash +pip install zeromodels[conversion] +# KERAS_BACKEND=torch; prefer CPU for the full ~28 GiB bf16 build +python -m zeromodels.diffusers.qwen_image_21.convert_qwen_image_21_diffusers_to_keras +``` + +`transfer_qwen_image_21(repo)` streams Diffusers shards into a `ZMDiffusion` and +writes sharded `*.weights.json`. Building the full text tower on GPU can OOM +during Functional shape tracing; convert on CPU, then load for CUDA inference. + +### Data Format {#qwen-image-2-1-data-format} + +**Channels-last for the VAE.** Transformer works on sequence tokens. + +| | Shape | +|---|---| +| `latents` passed to `generate` (unpatched) | `(batch, (H/16)·(W/16), 64)` | +| VAE encode input | `(batch, H, W, 4)` RGBA in `[-1, 1]` | +| VAE decode output | `(batch, H, W, 4)`; `generate` returns RGB uint8 | +| VAE latent | `(batch, H/16, W/16, 64)` | +| Transformer `sample` | `(batch, seq, 64)` | + +### Memory and speed {#qwen-image-2-1-memory-and-speed} + +The bf16 checkpoint is about 28 GiB (DiT ~14 GiB + text ~14 GiB + VAE ~1.4 GiB). +Building the full Functional graph at 1024px on GPU can OOM during attention shape +tracing (layers use `compute_output_spec` to avoid materializing seq² temps). Prefer +CPU convert / build, then CUDA generate. A 1024px run needs substantial activation +headroom on an 80 GB card; prefer fused attention +(`keras.ops.dot_product_attention` / torch SDPA). + +### Loading Fine-tuned Weights {#qwen-image-2-1-loading-fine-tuned-weights} + +The hosted checkpoint is the supported weight; any repo laid out like it +(`zm_config.json` declaring `ZMDiffusion` with `model_type: "qwen_image_21"`, sharded `*.weights.json` / +`*.weights.h5`, `tokenizer.json`) loads with `from_weights("/")`. The +`hf:` prefix raises for diffusion models: convert Diffusers format once with +`convert_qwen_image_21_diffusers_to_keras.py` and host the result. + ## Migrating from the per-family classes The per-family container and pipeline classes were replaced by `ZMDiffusion`, and the diff --git a/docs/getting_started.md b/docs/getting_started.md index dfb27e3e..1e5dac2c 100644 --- a/docs/getting_started.md +++ b/docs/getting_started.md @@ -55,8 +55,7 @@ results = processor.post_process_object_detection( [Qwen3-VL](qwen3_vl.md) ┬╖ [InternVL](internvl.md) ┬╖ [Kimi K2.5](kimi_k25.md) ┬╖ [LocateAnything](locateanything.md) ┬╖ - [Diffusers](diffusers.md) ┬╖ [Stable Diffusion](stable_diffusion.md) ┬╖ [Qwen-Image](qwen_image.md) ┬╖ - [Qwen-Image-2.1](qwen_image_21.md) + [Diffusion Models](diffusers.md) - **Speech** diff --git a/docs/index.md b/docs/index.md index 148cc9a8..9aa032d6 100644 --- a/docs/index.md +++ b/docs/index.md @@ -327,8 +327,7 @@ Vision-language generation, grounding, and text-to-image diffusion. [Qwen3-VL](qwen3_vl.md) · [InternVL](internvl.md) · [Kimi K2.5](kimi_k25.md) · [LocateAnything](locateanything.md) · -[Diffusers](diffusers.md) · [Stable Diffusion](stable_diffusion.md) · [Qwen-Image](qwen_image.md) · -[Qwen-Image-2.1](qwen_image_21.md) +[Diffusion Models](diffusers.md)
diff --git a/docs/main_classes.md b/docs/main_classes.md index 787e971d..ff25e4d1 100644 --- a/docs/main_classes.md +++ b/docs/main_classes.md @@ -258,10 +258,10 @@ model.generate( ``` The diffusion flavor, mixed into [`ZMDiffusion`](diffusers.md), the one pipeline class of -every diffusion family ([Stable Diffusion](stable_diffusion.md), -[Stable Diffusion 2](stable_diffusion_2.md), [Stable Diffusion XL](stable_diffusion_xl.md), -[Stable Diffusion 3 / 3.5](stable_diffusion_3.md), [Qwen-Image](qwen_image.md), -[Qwen-Image-2.1](qwen_image_21.md)); each family supplies the hooks through its +every diffusion family ([Stable Diffusion](diffusers.md#stable-diffusion), +[Stable Diffusion 2](diffusers.md#stable-diffusion-2), [Stable Diffusion XL](diffusers.md#stable-diffusion-xl), +[Stable Diffusion 3 / 3.5](diffusers.md#stable-diffusion-3), [Qwen-Image](diffusers.md#qwen-image), +[Qwen-Image-2.1](diffusers.md#qwen-image-2-1)); each family supplies the hooks through its family recipe, selected by the checkpoint's `model_type`. Where the LM mixins decode tokens, this one runs a scheduler's denoising loop with classifier-free guidance over a latent and decodes it to `(batch, H, W, 3)` uint8 images. A model diff --git a/docs/models.md b/docs/models.md index 0c0308a3..2d2106cb 100644 --- a/docs/models.md +++ b/docs/models.md @@ -41,24 +41,22 @@ Encoders for embeddings and masked language modelling, and decoder LLMs, dense a ## Sentence Transformer models -Dense text embeddings with checkpoint-specific pooling and normalization. +Dense sentence embeddings for similarity, semantic search and clustering. Every checkpoint +loads through two classes, `ZMSentenceTransformer` and `ZMSentenceTransformerTokenizer` +from `zeromodels.sentence_transformers`: a Keras encoder (MPNet, BERT, RoBERTa, XLM-R, +DistilBERT or ALBERT, chosen by the checkpoint's `backbone_type`) followed by its declared +pooling, optional Dense projection and optional L2 normalization. -- [All 13 Sentence Transformer checkpoints](sentence_transformers.md) -- [Sentence Transformer API and usage](sentence_transformer.md) +- [Sentence Transformers](sentence_transformer.md): API, usage and all 13 checkpoints ## Diffusion models -Text-to-image (and image-to-image) diffusion, every family loaded through `ZMDiffusion` -and `ZMDiffusionTokenizer`. - -- [Diffusers API and usage](diffusers.md) -- [Stable Diffusion](stable_diffusion.md) -- [Stable Diffusion 2](stable_diffusion_2.md) -- [Stable Diffusion XL](stable_diffusion_xl.md) -- [Stable Diffusion 3](stable_diffusion_3.md) -- [Stable Diffusion 3.5](stable_diffusion_3_5.md) -- [Qwen-Image](qwen_image.md) -- [Qwen-Image-2.1](qwen_image_21.md) +Text-to-image generation (and image-to-image with the SDXL refiner). Every family loads +through two classes, `ZMDiffusion` and `ZMDiffusionTokenizer` from `zeromodels.diffusers`; +the checkpoint's `model_type` selects the family. The page has its own +**Diffusion Models** entry in the navigation. + +- [Diffusion Models](diffusers.md): API, usage and every family (Stable Diffusion 1.x / 2.x / XL / 3 / 3.5, Qwen-Image, Qwen-Image-2.1) ## Vision models diff --git a/docs/qwen_image.md b/docs/qwen_image.md deleted file mode 100644 index feca55ce..00000000 --- a/docs/qwen_image.md +++ /dev/null @@ -1,357 +0,0 @@ -# Qwen-Image - -
-Weights: pretrained Keras weights live on Hugging Face under -zeromodels/<variant> -(each repo carries zm_config.json + sharded -*.weights.json / *.weights.h5 + -tokenizer.json). Load with from_weights("zeromodels/<variant>"). -
- -Qwen-Image, ported to pure Keras 3: latent text-to-image flow-matching with a -60-layer double-stream MMDiT denoiser, a Wan-derived KL autoencoder and the -Qwen2.5-VL-7B Instruct text tower. It loads through [`ZMDiffusion`](diffusers.md), the one diffusion -pipeline class of zeromodels, with `model_type="qwen_image"`: a single functional graph with the three towers as disconnected -sub-graphs (the way `CLIPModel` holds its two), so a hosted repo is one set of -sharded weights and one `zm_config.json`. plus -`generate`: the true-CFG denoising loop of `BaseDiffusion`, the -diffusion counterpart of `BaseGeneration`. - -The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion -is deliberately **not supported** for diffusion models. - -Key facts of the port: - -- **Packed latents**: the denoiser sees `(B, H/2 · W/2, 64)` tokens; pack / unpack - match Diffusers' 2×2 patching around the 16-channel VAE latent grid - (`AutoencoderKLQwenImage`, 8× spatial, channels-last NDHWC with `T=1`). -- **True CFG**: when `guidance_scale > 1`, two transformer forwards and Diffusers' - prediction-norm renormalization (`true_cfg_scale`); a space negative prompt is - the default when none is passed. -- **bf16-safe dual-stream clip**: Diffusers clips dual-stream activations to the - fp16 finite range only under float16; the port matches that, so bf16 runs keep - the ~1e7 magnitudes Qwen-Image produces instead of collapsing to noise. -- **Schedulers match Diffusers**: `FlowMatchEulerDiscreteScheduler` with dynamic - resolution shifting (`mu` from packed sequence length) so timesteps agree with - the reference within float rounding. -- **Verified against Diffusers**: transformer, text encoder, VAE and scheduler - outputs agree with the reference within float noise (see - [Verified against Diffusers](#verified-against-diffusers)). - -Links: - -- Paper: [Qwen-Image Technical Report (arXiv:2508.02324)](https://arxiv.org/abs/2508.02324) -- Reference implementation: [diffusers `QwenImagePipeline`](https://huggingface.co/docs/diffusers/api/pipelines/qwenimage) -- License: [Apache-2.0](https://huggingface.co/Qwen/Qwen-Image/blob/main/LICENSE) - -See also [qwen2_5_vl.md](qwen2_5_vl.md) (the text tower), -[qwen_image_21.md](qwen_image_21.md) (2.1 single-stream / unpatched), -[stable_diffusion_3.md](stable_diffusion_3.md) (flow-match / MMDiT-style diffusion -in ZeroModels). - -## Variants - -Preconverted, bfloat16 weights are hosted under `zeromodels/`. Load with -`from_weights("zeromodels/")`. Each repo is one checkpoint: MMDiT -transformer + VAE + Qwen2.5-VL text tower, ~53 GiB at 16-bit. Released under -Apache-2.0. - -| Variant | Hub | Source | -|---|---|---| -| `qwen-image` | [`zeromodels/qwen-image`](https://huggingface.co/zeromodels/qwen-image) | [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) | - -Default `generate_args`: 50 flow-match steps, `guidance_scale=4.0`, 1024×1024. - -## API - -Configs are typed: `QwenImageConfig` (composite, `model_type` `"qwen_image"`) over -`QwenImageTransformerConfig`, `QwenImageVAEConfig` and `QwenImageTextConfig`, plus -the checkpoint's `scheduler_config` and the Qwen special token ids. Each repo's -`zm_config.json` parses through it; the constructor stays flat, with the -sub-config fields prefixed `transformer_` / `vae_` / `text_`. - -### `ZMDiffusion` - -The pipeline class of every diffusion family (see [Diffusers](diffusers.md)); the -repo's `model_type` (`"qwen_image"` here) selects this family, which supplies the hooks -`BaseDiffusion`'s `generate` needs the mixin needs (`encode_prompt` on the text -tower with ChatML template drop, `unconditional_ids`, `predict_noise` on the -transformer, `decode_latents` on the VAE with mean/std un-normalization) and a -`scheduler`, built from the config's `scheduler_config` -(`FlowMatchEulerDiscreteScheduler` with dynamic shifting). - -```python -generate( - input_ids, - attention_mask=None, - negative_input_ids=None, - negative_attention_mask=None, - num_inference_steps=None, - guidance_scale=None, - seed=None, - latents=None, - height=None, - width=None, - output_type="image", - **conditioning, -) -``` - -| Arg | Default | Meaning | -|---|---|---| -| `input_ids` | required | ChatML-templated token ids, `**tokenizer(prompts)` | -| `attention_mask` | `None` | padding mask from the tokenizer | -| `negative_input_ids` | `None` | tokenized negative prompt for true CFG; a templated space prompt when unset | -| `negative_attention_mask` | `None` | mask for the negative ids | -| `negative_prompt` | `None` | optional text (or one string per image) passed via `**conditioning`; tokenized when negative ids are unset | -| `num_inference_steps` | `None` | scheduler steps; the repo's `generate_args` (50) when unset | -| `guidance_scale` | `None` | true CFG strength; `generate_args` (4.0) when unset, `<= 1` disables it | -| `seed` | `None` | seed for the initial latent, reproducible per backend | -| `latents` | `None` | explicit packed initial latent; identical results across backends | -| `height` / `width` | `None` | output pixel size; `default_sample_size * 8` (1024) when unset; rounded to VAE×pack multiples | -| `output_type` | `"image"` | `"image"` for uint8 RGB, `"latent"` for packed latents | - -Returns `(batch, height, width, 3)` uint8 numpy images. - -| Constructor arg | Default | Meaning | -|---|---|---| -| `model_type` | required | `"qwen_image"` (read from the repo's `zm_config.json` on `from_weights`) | -| `scheduler` | `None` | a `BaseScheduler`; built from `scheduler_config` when unset | -| `transformer_sample_size` | `128` | latent side the transformer graph is built for (image / 8) | -| `vae_sample_size` | `1024` | image size the VAE graphs are built for | -| `transformer_*` / `vae_*` / `text_*` | Qwen-Image | the flat sub-config fields (see [Configuration](configuration.md)) | -| `bos_token_id` / `eos_token_id` / `pad_token_id` | `151643` / `151645` / `151643` | Qwen special tokens | -| `scheduler_config` | `None` | the Diffusers scheduler dict (`_class_name`, shift, dynamic shifting, ...) | -| `prompt_template_encode_start_idx` | `34` | ChatML template tokens dropped after the text encoder | -| `max_sequence_length` | `512` | prompt embed length after the template drop | - -The graph is three disconnected paths, one per component. Inputs cover the transformer (`sample`, `timestep`, -`encoder_hidden_states`, `encoder_hidden_states_mask`), the VAE (`image`, `latent`) -and the text tower (`token_ids`, `padding_mask`); outputs include the packed -velocity, VAE moments / sample, and prompt embeds. The components are exposed as -`.transformer`, `.vae` and `.text_encoder` and share the model's weights, so one load -fills all three, and each can be called on its own. - -### `QwenImageTransformer2DModel` - -The denoiser (Diffusers' `QwenImageTransformer2DModel`): a 60-layer double-stream -MMDiT over **packed** latents and text features, with MS-RoPE -(`axes_dims_rope=(16, 56, 56)`). Inputs -`{"sample": (B, seq, 64), "timestep": (B,), "encoder_hidden_states": (B, text_seq, 3584)}` -(optional `encoder_hidden_states_mask`), output `{"sample": (B, seq, 64)}` packed -velocity. - -| Arg | Default | Meaning | -|---|---|---| -| `patch_size` | `2` | latent pack side | -| `in_channels` | `64` | packed token width (`16 × 2 × 2`) | -| `out_channels` | `16` | unpacked latent channels | -| `num_layers` | `60` | dual-stream DiT blocks | -| `attention_head_dim` / `num_attention_heads` | `128` / `24` | head geometry (inner dim 3072) | -| `joint_attention_dim` | `3584` | text feature width from Qwen2.5-VL | -| `axes_dims_rope` | `(16, 56, 56)` | MS-RoPE axis splits | -| `sample_size` | `128` | latent spatial side the graph is built for | -| `text_seq_len` | `512` | static text length after the template drop | - -### `AutoencoderKLQwenImage` - -The VAE (Diffusers' `AutoencoderKLQwenImage`): `encode(image, sample=False)` returns -the posterior mean (or a sample) of a `(B, H/8, W/8, 16)` latent and -`decode(latent)` the `(B, H, W, 3)` image in `[-1, 1]`. Latent normalisation uses -`latents_mean` / `latents_std` around the denoiser, as `generate` does. One -functional graph with the encoder and decoder as two disconnected paths -(channels-last NDHWC, `T=1` for still images). - -| Arg | Default | Meaning | -|---|---|---| -| `base_dim` | `96` | base channel width | -| `z_dim` | `16` | latent channels | -| `dim_mult` | `(1, 2, 4, 4)` | width multipliers per level | -| `num_res_blocks` | `2` | residual blocks per level | -| `temperal_downsample` | `(False, True, True)` | which levels downsample time (unused at `T=1`) | -| `sample_size` | `1024` | image size the graphs are built for | -| `latents_mean` / `latents_std` | Qwen-Image | per-channel latent normalisation | - -### Schedulers - -`zeromodels.base.base_scheduler` has the flow-match sampler with the Diffusers -interface (`set_timesteps`, `step`, `init_noise_sigma`), configured the Diffusers -way: - -| Scheduler | Notes | -|---|---| -| `FlowMatchEulerDiscreteScheduler` | Qwen-Image default; `use_dynamic_shifting=True`, exponential time shift | - -`get_scheduler(config)` builds the one named by a Diffusers scheduler config dict, -which is what the model does with its `scheduler_config`. At generate time the -task sets timesteps with resolution-dependent `mu` and a linspace sigma schedule, -matching Diffusers. - -## Preprocessing - -### `ZMDiffusionTokenizer` - -The Qwen2 BPE tokenizer with the Diffusers ChatML prompt template wrapped around -each string. Calling it returns `{"input_ids", "attention_mask"}` for a string or -a list of strings. `generate` / `encode_prompt` drop the template prefix -(`prompt_template_encode_start_idx=34`) after the text encoder and keep at most -`max_sequence_length` (512) tokens. - -```python -ZMDiffusionTokenizer( - model_type="qwen_image", hf_id=None, tokenizer_file=None, max_seq_len=1024 -) -``` - -There is no image processor and no processor class: text-to-image takes token ids -in and hands uint8 images out. - -## End-to-end example - -### Single prompt - -```python -import os - -os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" - -from PIL import Image -from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer - -model = ZMDiffusion.from_weights("zeromodels/qwen-image") -tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/qwen-image") - -inputs = tokenizer( - "a coffee shop entrance with a chalkboard sign reading 'Open', warm afternoon light" -) -images = model.generate( - **inputs, - height=1024, - width=1024, - num_inference_steps=50, - guidance_scale=4.0, - seed=1, -) - -Image.fromarray(images[0]).save("coffee.png") # (1024, 1024, 3) uint8 -``` - -### Negative prompt and batching - -Prompts batch through the tokenizer; a negative prompt is tokenized the same way -and passed as `negative_input_ids` (one per prompt, or one broadcast row). - -```python -prompts = [ - "a watercolor painting of a lighthouse at sunset", - "a red vintage car on a coastal road", -] -inputs = tokenizer(prompts) -negative = tokenizer(["blurry, low quality"] * len(prompts))["input_ids"] - -images = model.generate( - **inputs, - negative_input_ids=negative, - height=512, - width=512, - num_inference_steps=30, -) -for i, image in enumerate(images): - Image.fromarray(image).save(f"out_{i}.png") -``` - -### Reproducible across backends - -`seed` reproduces a run on a given backend, but `keras.random` differs between -torch, jax and tensorflow. For an image that is identical everywhere, pass the -initial packed latent yourself. At 1024px the spatial noise before packing is -`(batch, 128, 128, 16)`; after packing it is `(batch, 4096, 64)`: - -```python -import numpy as np - -from zeromodels.diffusers.qwen_image.qwen_image_model import pack_latents - -h, w, channels = 128, 128, 16 -noise = np.random.default_rng(0).standard_normal((1, h, w, channels)).astype("float32") -latents = pack_latents(noise, h, w) -images = model.generate(**tokenizer("a bowl of ramen"), latents=latents) -``` - -### Verified against Diffusers - -Component checks against Diffusers / Transformers with the same weights (fp32 -unless noted): - -``` -scheduler timesteps / mu within 1e-4 / identical -transformer (fp32, N blocks) relative error ~8e-6 -transformer (bf16, N blocks) cosine similarity 0.99998 -text encoder (fp32) relative error ~1.7e-6 -VAE encode / decode matches Diffusers (NCDHW <-> NDHWC transpose) -``` - -End-to-end image parity depends on dtype and resolution; prefer bf16 (or fp32) -over fp16 — Qwen-Image activations overflow the fp16 range without the Diffusers -clip path. - -### Components on their own - -```python -from zeromodels.diffusers import ZMDiffusion - -qi = ZMDiffusion.from_weights("zeromodels/qwen-image") -latent = qi.vae.encode(image) # (1, H, W, 3) in [-1, 1] -> (1, H/8, W/8, 16) -text = qi.text_encoder({"input_ids": ids, "attention_mask": mask})["last_hidden_state"] -``` - -### Other resolutions - -The graphs are built for a fixed size, the weights are not. Rebuild at another -size with the constructor overrides (multiples of 16px after VAE×pack; the -checkpoint targets 1024px): - -```python -model = ZMDiffusion.from_weights( - "zeromodels/qwen-image", - transformer_sample_size=64, - vae_sample_size=512, -) -images = model.generate( - **tokenizer("a mountain lake at dawn"), height=512, width=512 -) # (1, 512, 512, 3) -``` - -## Data Format - -**Channels-last only for the VAE.** `AutoencoderKLQwenImage` is built as NDHWC -`Conv3D` (with `T=1` for still images); the transformer works on packed sequence -tokens and has no spatial layout. `generate` always hands back -`(batch, height, width, 3)` uint8 images. - -| | Shape | -|---|---| -| `latents` passed to `generate` (packed) | `(batch, (H/16)·(W/16), 64)` | -| VAE encode input / decode output | `(batch, H, W, 3)` | -| VAE latent | `(batch, H/8, W/8, 16)` | -| Transformer `sample` | `(batch, seq, 64)` | - -## Memory and speed - -The bf16 / float16 checkpoint is about 53 GiB. Building the full graph on a ~40 GB -GPU OOMs; build on CPU (or load with component offload) and move one tower at a -time. A guided 1024px run needs substantial activation headroom on top of the -~38 GiB transformer; 512px is the practical default on 40 GB cards. Prefer the -fused attention path (`keras.ops.dot_product_attention` / torch SDPA) over the -plain matmul softmax default. - -## Loading Fine-tuned Weights - -The hosted checkpoint is the supported weight; any repo laid out like it -(`zm_config.json` declaring `ZMDiffusion` with `model_type: "qwen_image"`, sharded `*.weights.json` / -`*.weights.h5`, `tokenizer.json`) loads with `from_weights("/")`. The -`hf:` prefix raises for diffusion models: convert a Diffusers-format checkpoint -once with -`zeromodels/diffusers/qwen_image/convert_qwen_image_diffusers_to_keras.py` -(`transfer_qwen_image(repo)`, `pip install zeromodels[conversion]`) and host the -result. diff --git a/docs/qwen_image_21.md b/docs/qwen_image_21.md deleted file mode 100644 index 88e95de6..00000000 --- a/docs/qwen_image_21.md +++ /dev/null @@ -1,299 +0,0 @@ -# Qwen-Image-2.1 - -
-Weights: pretrained Keras weights live on Hugging Face under -zeromodels/<variant> -(each repo carries zm_config.json + sharded -*.weights.json / *.weights.h5 + -tokenizer.json). Load with from_weights("zeromodels/<variant>"). -
- -Qwen-Image-2.1, ported to pure Keras 3: latent text-to-image flow-matching with a -32-layer **single-stream** block-causal DiT, a residual 64-channel KL autoencoder -(16× spatial, RGBA), and the Qwen3-VL text tower. It loads through -[`ZMDiffusion`](diffusers.md), the one diffusion pipeline class of zeromodels, with -`model_type="qwen_image_21"`: one functional graph holding the three towers, plus -`generate`. Prompts are tokenized by `ZMDiffusionTokenizer`. - -The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion -is deliberately **not supported** for diffusion models. - -Key facts of the port: - -- **Unpatched latents**: the denoiser sees `(B, H · W, 64)` tokens (VAE scale 16; - no 2×2 packing). At 1024px that is a `(B, 4096, 64)` sequence. -- **Block-causal attention**: text is causal; the target image block is - bidirectional and can attend to all preceding text. -- **`causal_condition`**: text tokens modulate from `t = 0` (timestep-independent), - matching Diffusers' KV-cache-ready conditioning. -- **True CFG optional**: Diffusers defaults to `true_cfg_scale=1.0` (no guidance). - Pass `guidance_scale > 1` with a negative prompt to enable dual forwards. -- **Pre-norm text features**: the text tower returns decoder outputs *before* the - final RMSNorm, matching Diffusers' forward hook on the language-model norm. -- **Schedulers match Diffusers**: `FlowMatchEulerDiscreteScheduler` with dynamic - resolution shifting (`mu` from image sequence length) and `shift_terminal`. - -Links: - -- Source: [`Qwen/Qwen-Image-2.1`](https://huggingface.co/Qwen/Qwen-Image-2.1) -- Reference: [diffusers `QwenImage21Pipeline`](https://huggingface.co/docs/diffusers/api/pipelines/qwenimage) -- License: [Qwen Research License](https://huggingface.co/Qwen/Qwen-Image-2.1) -- See also [qwen_image.md](qwen_image.md) (1.0 double-stream / packed), - [qwen3_vl.md](qwen3_vl.md) (text tower) - -## Variants - -Preconverted, bfloat16 weights are hosted under `zeromodels/`. Load with -`from_weights("zeromodels/")`. Each repo is one checkpoint: DiT + VAE + -Qwen3-VL text tower, ~28 GiB at 16-bit. - -| Variant | Hub | Source | -|---|---|---| -| `qwen-image-2.1` | [`zeromodels/qwen-image-2.1`](https://huggingface.co/zeromodels/qwen-image-2.1) | [`Qwen/Qwen-Image-2.1`](https://huggingface.co/Qwen/Qwen-Image-2.1) | - -Default `generate_args`: 40 flow-match steps, `guidance_scale=1.0`, 1024×1024. - -## API - -Configs are typed: `QwenImage21Config` (composite, `model_type` `"qwen_image_21"`) -over `QwenImage21TransformerConfig`, `QwenImage21VAEConfig` and -`QwenImage21TextConfig`, plus the checkpoint's `scheduler_config` and Qwen special -token ids. Each repo's `zm_config.json` parses through it; the constructor stays -flat, with the sub-config fields prefixed `transformer_` / `vae_` / `text_`. - -### `ZMDiffusion` - -The pipeline class of every diffusion family (see [Diffusers](diffusers.md)); the -repo's `model_type` (`"qwen_image_21"` here) selects this family, which supplies the hooks -`BaseDiffusion`'s `generate` needs: `encode_prompt` (ChatML template drop after the text -encoder), `predict_noise` on the transformer, `decode_latents` on the VAE with -mean/std un-normalization, and a `FlowMatchEulerDiscreteScheduler` from -`scheduler_config`. - -```python -generate( - input_ids, - attention_mask=None, - negative_input_ids=None, - negative_attention_mask=None, - num_inference_steps=None, - guidance_scale=None, - seed=None, - latents=None, - height=None, - width=None, - output_type="image", -) -``` - -| Arg | Default | Meaning | -|---|---|---| -| `input_ids` | required | ChatML-templated token ids, `**tokenizer(prompts)` | -| `attention_mask` | `None` | padding mask from the tokenizer | -| `negative_input_ids` | `None` | tokenized negative prompt; CFG only when set **and** `guidance_scale > 1` | -| `negative_attention_mask` | `None` | mask for the negative ids | -| `num_inference_steps` | `None` | scheduler steps; `generate_args` (40) when unset | -| `guidance_scale` | `None` | true CFG strength; `generate_args` (1.0) when unset | -| `seed` | `None` | seed for the initial latent | -| `latents` | `None` | explicit unpatched initial latent `(B, H·W, 64)` | -| `height` / `width` | `None` | output pixel size; `default_sample_size * 16` (1024) when unset | -| `output_type` | `"image"` | `"image"` for uint8 RGB, `"latent"` for sequence latents | - -Returns `(batch, height, width, 3)` uint8 numpy images (VAE RGBA is cropped to RGB). - -| Constructor arg | Default | Meaning | -|---|---|---| -| `model_type` | required | `"qwen_image_21"` (read from the repo's `zm_config.json` on `from_weights`) | -| `scheduler` | `None` | a `BaseScheduler`; built from `scheduler_config` when unset | -| `transformer_sample_size` | `64` | latent side the DiT graph is built for (image / 16) | -| `vae_sample_size` | `1024` | image size the VAE graphs are built for | -| `transformer_*` / `vae_*` / `text_*` | Qwen-Image-2.1 | flat sub-config fields (see [Configuration](configuration.md)) | -| `bos_token_id` / `eos_token_id` / `pad_token_id` | `151643` / `151645` / `151643` | Qwen special tokens | -| `scheduler_config` | `None` | Diffusers scheduler dict (dynamic shifting, `max_shift=0.9`, ...) | -| `prompt_template_encode_start_idx` | `14` | ChatML system-prefix tokens dropped after the text encoder | -| `max_sequence_length` | `512` | prompt embed length after the template drop | - -The graph is three disconnected paths, one per component. -Inputs cover the transformer (`sample`, `timestep`, `encoder_hidden_states`, -`encoder_hidden_states_mask`), the VAE (`image`, `latent`) and the text tower -(`token_ids`, `padding_mask`). Components are exposed as `.transformer`, `.vae` -and `.text_encoder`. - -### `QwenImage21Transformer2DModel` - -The denoiser (Diffusers' `QwenImage21Transformer2DModel`): a 32-layer -single-stream DiT over **unpatched** latents and Qwen3-VL text features, with -3-axis RoPE (`axes_dims_rope=(16, 56, 56)`). Inputs -`{"sample": (B, seq, 64), "timestep": (B,), "encoder_hidden_states": (B, text_seq, 4096)}` -(plus mask), output `{"sample": (B, seq, 64)}` target velocity. - -| Arg | Default | Meaning | -|---|---|---| -| `patch_size` | `1` | unpatched (asserted) | -| `in_channels` / `out_channels` | `64` / `64` | latent token width | -| `num_layers` | `32` | single-stream blocks | -| `attention_head_dim` / `num_attention_heads` | `128` / `32` | head geometry (inner dim 4096) | -| `context_in_dim` | `4096` | text feature width from Qwen3-VL | -| `mlp_ratio` | `3` | SwiGLU expansion | -| `causal_condition` | `True` | text modulates from `t=0` | -| `sample_size` | `64` | latent spatial side the graph is built for | -| `text_seq_len` | `512` | static text length after the template drop | - -### `AutoencoderKLQwenImage21` - -The VAE (Diffusers' `AutoencoderKLQwenImage21`): residual encoder/decoder, -`encode(image)` → `(B, H/16, W/16, 64)` latents, `decode(latent)` → -`(B, H, W, 4)` RGBA in `[-1, 1]`. Latent normalisation uses `latents_mean` / -`latents_std` around the denoiser. - -| Arg | Default | Meaning | -|---|---|---| -| `base_dim` / `decoder_base_dim` | `96` / `144` | channel widths | -| `z_dim` | `64` | latent channels | -| `dim_mult` | `(1, 2, 4, 8, 8)` | width multipliers per level | -| `scale_factor_spatial` | `16` | pixel → latent downscale | -| `input_channels` / `out_channels` | `4` / `4` | RGBA | -| `is_residual` | `True` | residual blocks | -| `sample_size` | `1024` | image size the graphs are built for | -| `latents_mean` / `latents_std` | Qwen-Image-2.1 | per-channel latent normalisation | - -### `QwenImage21TextEncoderModel` - -Qwen3-VL text tower only (no vision / LM head). Returns -`{"last_hidden_state"}` as **pre-final-norm** hidden states. - -### Schedulers - -| Scheduler | Notes | -|---|---| -| `FlowMatchEulerDiscreteScheduler` | `use_dynamic_shifting=True`, exponential time shift, `shift_terminal=0.02`, `max_shift=0.9` | - -At generate time the task sets timesteps with resolution-dependent `mu` and a -linspace sigma schedule, matching Diffusers. - -## Preprocessing - -### `ZMDiffusionTokenizer` - -Qwen3 BPE with Diffusers' T2I ChatML template (`Comprehend and analyze the -provided prompt.`). Returns `{"input_ids", "attention_mask"}`. `encode_prompt` -drops `prompt_template_encode_start_idx` (14) tokens and keeps at most -`max_sequence_length` (512). - -```python -ZMDiffusionTokenizer( - model_type="qwen_image_21", hf_id=None, tokenizer_file=None, max_seq_len=1024 -) -``` - -## End-to-end example - -### Single prompt - -```python -import os - -os.environ["KERAS_BACKEND"] = "torch" - -from PIL import Image -from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer - -model = ZMDiffusion.from_weights("zeromodels/qwen-image-2.1") -tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/qwen-image-2.1") - -inputs = tokenizer("a photo of a capybara wearing a wizard hat, soft window light") -images = model.generate( - **inputs, - height=1024, - width=1024, - num_inference_steps=40, - guidance_scale=1.0, - seed=0, -) - -Image.fromarray(images[0]).save("capybara.png") # (1024, 1024, 3) uint8 -``` - -### Optional true CFG - -```python -inputs = tokenizer("a watercolor lighthouse at sunset") -negative = tokenizer("blurry, low quality")["input_ids"] - -images = model.generate( - **inputs, - negative_input_ids=negative, - guidance_scale=4.0, - height=1024, - width=1024, -) -``` - -### Reproducible latents - -At 1024px the unpatched noise is `(batch, 64, 64, 64)` → `(batch, 4096, 64)`: - -```python -import numpy as np -from zeromodels.diffusers.qwen_image_21.qwen_image_21_model import pack_latents - -h, w, channels = 64, 64, 64 -noise = np.random.default_rng(0).standard_normal((1, h, w, channels)).astype("float32") -latents = pack_latents(noise, h, w) -images = model.generate(**tokenizer("a bowl of ramen"), latents=latents) -``` - -### Other resolutions - -Graphs are built for a fixed size; weights are not. Rebuild with constructor -overrides (multiples of 32px; checkpoint targets 1024): - -```python -model = ZMDiffusion.from_weights( - "zeromodels/qwen-image-2.1", - transformer_sample_size=32, - vae_sample_size=512, -) -images = model.generate(**tokenizer("a mountain lake at dawn"), height=512, width=512) -``` - -### Offline conversion - -```bash -pip install zeromodels[conversion] -# KERAS_BACKEND=torch; prefer CPU for the full ~28 GiB bf16 build -python -m zeromodels.diffusers.qwen_image_21.convert_qwen_image_21_diffusers_to_keras -``` - -`transfer_qwen_image_21(repo)` streams Diffusers shards into a `ZMDiffusion` and -writes sharded `*.weights.json`. Building the full text tower on GPU can OOM -during Functional shape tracing; convert on CPU, then load for CUDA inference. - -## Data Format - -**Channels-last for the VAE.** Transformer works on sequence tokens. - -| | Shape | -|---|---| -| `latents` passed to `generate` (unpatched) | `(batch, (H/16)·(W/16), 64)` | -| VAE encode input | `(batch, H, W, 4)` RGBA in `[-1, 1]` | -| VAE decode output | `(batch, H, W, 4)`; `generate` returns RGB uint8 | -| VAE latent | `(batch, H/16, W/16, 64)` | -| Transformer `sample` | `(batch, seq, 64)` | - -## Memory and speed - -The bf16 checkpoint is about 28 GiB (DiT ~14 GiB + text ~14 GiB + VAE ~1.4 GiB). -Building the full Functional graph at 1024px on GPU can OOM during attention shape -tracing (layers use `compute_output_spec` to avoid materializing seq² temps). Prefer -CPU convert / build, then CUDA generate. A 1024px run needs substantial activation -headroom on an 80 GB card; prefer fused attention -(`keras.ops.dot_product_attention` / torch SDPA). - -## Loading Fine-tuned Weights - -The hosted checkpoint is the supported weight; any repo laid out like it -(`zm_config.json` declaring `ZMDiffusion` with `model_type: "qwen_image_21"`, sharded `*.weights.json` / -`*.weights.h5`, `tokenizer.json`) loads with `from_weights("/")`. The -`hf:` prefix raises for diffusion models: convert Diffusers format once with -`convert_qwen_image_21_diffusers_to_keras.py` and host the result. diff --git a/docs/sentence_transformer.md b/docs/sentence_transformer.md index 4962651c..5531d49f 100644 --- a/docs/sentence_transformer.md +++ b/docs/sentence_transformer.md @@ -1,4 +1,4 @@ -# Sentence Transformer +# Sentence Transformers
Weights: Pre-converted Keras weights for all 13 checkpoints are hosted under the @@ -10,11 +10,10 @@ with the hf: prefix, e.g. The backbone and embedding recipe are read from the checkpoint metadata.
-ZeroModels loads 13 original dense sentence-embedding checkpoints. See the -[Sentence Transformer models](sentence_transformers.md) page for the available -ZeroModels-hosted checkpoint IDs and Hub links. -It builds a Keras 3 encoder and applies the checkpoint's declared pooling, -optional Dense projection, and optional L2 normalization. Parity has been +ZeroModels loads 13 original dense sentence-embedding checkpoints (listed under +[Checkpoints](#checkpoints)). It builds a Keras 3 encoder and applies the +checkpoint's declared pooling, optional Dense projection, and optional L2 +normalization. Parity has been verified on the Keras Torch backend; TensorFlow and JAX parity have not yet been checked for all 13 checkpoints. @@ -36,6 +35,27 @@ The public API is `zeromodels.sentence_transformers`, separate from the upstream Supported `backbone_type` values: `mpnet`, `bert`, `roberta`, `xlm-roberta`, `distilbert`, `albert`. +## Checkpoints + +| Checkpoint | Backbone | Embedding recipe | Dimensions | +|---|---|---|---:| +| [`all-mpnet-base-v2`](https://huggingface.co/zeromodels/all-mpnet-base-v2) | MPNet | Mean pooling → L2 normalization | 768 | +| [`multi-qa-mpnet-base-dot-v1`](https://huggingface.co/zeromodels/multi-qa-mpnet-base-dot-v1) | MPNet | CLS pooling; dot-product similarity | 768 | +| [`all-distilroberta-v1`](https://huggingface.co/zeromodels/all-distilroberta-v1) | RoBERTa | Mean pooling → L2 normalization | 768 | +| [`all-MiniLM-L12-v2`](https://huggingface.co/zeromodels/all-MiniLM-L12-v2) | BERT | Mean pooling → L2 normalization | 384 | +| [`multi-qa-distilbert-cos-v1`](https://huggingface.co/zeromodels/multi-qa-distilbert-cos-v1) | DistilBERT | Mean pooling → L2 normalization | 768 | +| [`all-MiniLM-L6-v2`](https://huggingface.co/zeromodels/all-MiniLM-L6-v2) | BERT | Mean pooling → L2 normalization | 384 | +| [`multi-qa-MiniLM-L6-cos-v1`](https://huggingface.co/zeromodels/multi-qa-MiniLM-L6-cos-v1) | BERT | Mean pooling → L2 normalization | 384 | +| [`paraphrase-multilingual-mpnet-base-v2`](https://huggingface.co/zeromodels/paraphrase-multilingual-mpnet-base-v2) | XLM-RoBERTa | Mean pooling | 768 | +| [`paraphrase-albert-small-v2`](https://huggingface.co/zeromodels/paraphrase-albert-small-v2) | ALBERT | Mean pooling | 768 | +| [`paraphrase-multilingual-MiniLM-L12-v2`](https://huggingface.co/zeromodels/paraphrase-multilingual-MiniLM-L12-v2) | BERT | Mean pooling | 384 | +| [`paraphrase-MiniLM-L3-v2`](https://huggingface.co/zeromodels/paraphrase-MiniLM-L3-v2) | BERT | Mean pooling | 384 | +| [`distiluse-base-multilingual-cased-v1`](https://huggingface.co/zeromodels/distiluse-base-multilingual-cased-v1) | DistilBERT | Mean pooling → Tanh Dense | 512 | +| [`distiluse-base-multilingual-cased-v2`](https://huggingface.co/zeromodels/distiluse-base-multilingual-cased-v2) | DistilBERT | Mean pooling → Tanh Dense | 512 | + +The `paraphrase-multilingual-mpnet-base-v2` name is historical; the checkpoint +actually uses an XLM-RoBERTa backbone. + ## Usage ```python @@ -172,6 +192,6 @@ actual = model.encode(texts, tokenizer=tokenizer) print("Maximum absolute difference:", np.max(np.abs(actual - expected))) ``` -Use matching checkpoint names from the models table, compare in float32, and +Use matching checkpoint names from the [Checkpoints](#checkpoints) table, compare in float32, and record Hub commit SHAs for reproducible results. This downloads original weights for comparison; it does not convert or publish a new checkpoint. diff --git a/docs/sentence_transformers.md b/docs/sentence_transformers.md deleted file mode 100644 index 5a37ee8b..00000000 --- a/docs/sentence_transformers.md +++ /dev/null @@ -1,49 +0,0 @@ -# Sentence Transformer models - -ZeroModels provides 13 converted dense sentence-embedding checkpoints. Each -checkpoint combines a text encoder with its original pooling, optional Dense -projection, and optional L2 normalization. Load them through the -[`zeromodels.sentence_transformers` API](sentence_transformer.md); no conversion is needed -for the ZeroModels-hosted weights. - -| Checkpoint | Backbone | Embedding recipe | Dimensions | -|---|---|---|---:| -| [`all-mpnet-base-v2`](https://huggingface.co/zeromodels/all-mpnet-base-v2) | MPNet | Mean pooling → L2 normalization | 768 | -| [`multi-qa-mpnet-base-dot-v1`](https://huggingface.co/zeromodels/multi-qa-mpnet-base-dot-v1) | MPNet | CLS pooling; dot-product similarity | 768 | -| [`all-distilroberta-v1`](https://huggingface.co/zeromodels/all-distilroberta-v1) | RoBERTa | Mean pooling → L2 normalization | 768 | -| [`all-MiniLM-L12-v2`](https://huggingface.co/zeromodels/all-MiniLM-L12-v2) | BERT | Mean pooling → L2 normalization | 384 | -| [`multi-qa-distilbert-cos-v1`](https://huggingface.co/zeromodels/multi-qa-distilbert-cos-v1) | DistilBERT | Mean pooling → L2 normalization | 768 | -| [`all-MiniLM-L6-v2`](https://huggingface.co/zeromodels/all-MiniLM-L6-v2) | BERT | Mean pooling → L2 normalization | 384 | -| [`multi-qa-MiniLM-L6-cos-v1`](https://huggingface.co/zeromodels/multi-qa-MiniLM-L6-cos-v1) | BERT | Mean pooling → L2 normalization | 384 | -| [`paraphrase-multilingual-mpnet-base-v2`](https://huggingface.co/zeromodels/paraphrase-multilingual-mpnet-base-v2) | XLM-RoBERTa | Mean pooling | 768 | -| [`paraphrase-albert-small-v2`](https://huggingface.co/zeromodels/paraphrase-albert-small-v2) | ALBERT | Mean pooling | 768 | -| [`paraphrase-multilingual-MiniLM-L12-v2`](https://huggingface.co/zeromodels/paraphrase-multilingual-MiniLM-L12-v2) | BERT | Mean pooling | 384 | -| [`paraphrase-MiniLM-L3-v2`](https://huggingface.co/zeromodels/paraphrase-MiniLM-L3-v2) | BERT | Mean pooling | 384 | -| [`distiluse-base-multilingual-cased-v1`](https://huggingface.co/zeromodels/distiluse-base-multilingual-cased-v1) | DistilBERT | Mean pooling → Tanh Dense | 512 | -| [`distiluse-base-multilingual-cased-v2`](https://huggingface.co/zeromodels/distiluse-base-multilingual-cased-v2) | DistilBERT | Mean pooling → Tanh Dense | 512 | - -The `paraphrase-multilingual-mpnet-base-v2` name is historical; the checkpoint -actually uses an XLM-RoBERTa backbone. - -## Quick start - -```python -import os - -os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" - -from zeromodels.sentence_transformers import ( - ZMSentenceTransformer, - ZMSentenceTransformerTokenizer, -) - -repo = "zeromodels/all-MiniLM-L6-v2" -model = ZMSentenceTransformer.from_weights(repo) -tokenizer = ZMSentenceTransformerTokenizer.from_weights(repo) - -embeddings = model.encode(["A cat on a mat.", "A dog on a mat."], tokenizer=tokenizer) -print(embeddings.shape) # (2, 384) -``` - -See [Sentence Transformer](sentence_transformer.md) for the complete API, -upstream `hf:` loading, similarity scoring, and parity-check instructions. diff --git a/docs/stable_diffusion.md b/docs/stable_diffusion.md deleted file mode 100644 index 86e221aa..00000000 --- a/docs/stable_diffusion.md +++ /dev/null @@ -1,353 +0,0 @@ -# Stable Diffusion - -
-Weights: pretrained Keras weights live on Hugging Face under -zeromodels/<variant> -(each repo carries zm_config.json + model.weights.h5 + -tokenizer.json). Load with from_weights("zeromodels/<variant>"). -
- -Stable Diffusion 1.x, ported to pure Keras 3: latent text-to-image diffusion with a -UNet denoiser, a KL-regularized autoencoder and the CLIP ViT-L/14 text encoder. It loads -through [`ZMDiffusion`](diffusers.md), the one diffusion pipeline class of zeromodels, -with `model_type="stable_diffusion"`: a single functional graph with the three towers as -disconnected sub-graphs (the way `CLIPModel` holds its two), so a hosted repo is one -`model.weights.h5` and one `zm_config.json`, plus `generate`, the classifier-free -guidance denoising loop of `BaseDiffusion` (the diffusion counterpart of -`BaseGeneration`). Prompts are tokenized by `ZMDiffusionTokenizer`. - -All five SD 1.x checkpoints (v1-1 to v1-5) share this architecture; only the weights -differ. The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion -is deliberately **not supported** for diffusion models. - -Key facts of the port: - -- **Both data formats**: the UNet and VAE follow `keras.config.image_data_format()` - (every conv takes it, GroupNorm normalizes the layout's channel axis, and the two - token <-> grid boundaries transpose around their flatten under `channels_first`). - Weights are layout-independent, so one hosted checkpoint serves both, and the - converter is a `(O, I, H, W) -> (H, W, I, O)` kernel transpose and nothing else. -- **Block-level layers**: the UNet and VAE are built from composite Keras layers - (`StableDiffusionResnetBlock2D`, `StableDiffusionTransformer2DModel`, - `StableDiffusionCrossAttention`, ...). A functional graph keeps every node's output - alive until the forward ends, and at 512px the UNet's per-op intermediates (4096x4096 - attention maps above all) would need several GB; a layer's internals are freed when - its call returns. -- **Schedulers match diffusers to the bit**: the training noise schedule is built in - float32 the way torch does it, so `alphas_cumprod` and the Euler sigmas are identical - and a 50-step run stays within float rounding of the reference. -- **Verified end to end**: at 512px, fp32, PNDM 50 steps and the same initial latent, the - generated image matches diffusers' to within **1 uint8 level** (99.3% of pixels - identical, PSNR 75 dB); the UNet, VAE and text encoder outputs agree to ~1e-6. - -Links: - -- Paper: [High-Resolution Image Synthesis with Latent Diffusion Models (arXiv:2112.10752)](https://arxiv.org/abs/2112.10752) -- Reference implementation: [diffusers `StableDiffusionPipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img) -- License: [CreativeML OpenRAIL-M](https://huggingface.co/spaces/CompVis/stable-diffusion-license) - -See also [clip.md](clip.md) (the text encoder is `CLIPTextModel`). - -## Variants - -Preconverted, float32 weights are hosted under `zeromodels/`. Load with -`from_weights("zeromodels/")`. Each repo is one container: UNet (860M) + VAE -(84M) + CLIP text encoder (123M), 1.07B parameters, 3.97 GB. All checkpoints are released -under the CreativeML OpenRAIL-M license, which carries use-based restrictions. - -| Variant | Hub | Training | -|---|---|---| -| `stable-diffusion-v1-1` | [`zeromodels/stable-diffusion-v1-1`](https://huggingface.co/zeromodels/stable-diffusion-v1-1) | 237k steps at 256px on laion2B-en, then 194k steps at 512px on laion-high-resolution | -| `stable-diffusion-v1-2` | [`zeromodels/stable-diffusion-v1-2`](https://huggingface.co/zeromodels/stable-diffusion-v1-2) | v1-1 + 515k steps at 512px on laion-aesthetics v2 5+ | -| `stable-diffusion-v1-3` | [`zeromodels/stable-diffusion-v1-3`](https://huggingface.co/zeromodels/stable-diffusion-v1-3) | v1-2 + 195k steps at 512px, 10% text-conditioning dropout | -| `stable-diffusion-v1-4` | [`zeromodels/stable-diffusion-v1-4`](https://huggingface.co/zeromodels/stable-diffusion-v1-4) | v1-2 + 225k steps at 512px, 10% text-conditioning dropout | -| `stable-diffusion-v1-5` | [`zeromodels/stable-diffusion-v1-5`](https://huggingface.co/zeromodels/stable-diffusion-v1-5) | v1-2 + 595k steps at 512px, 10% text-conditioning dropout | - -The 10% text-conditioning dropout of v1-3 onwards is what makes classifier-free guidance -work well; v1-5 is the usual default. - -## API - -Configs are typed: `StableDiffusionConfig` (composite) over `UNet2DConditionConfig`, -`AutoencoderKLConfig` and `StableDiffusionTextConfig` (a `CLIPTextConfig` with the -ViT-L/14 defaults), plus the checkpoint's `scheduler_config` and the CLIP special token -ids. Each repo's `zm_config.json` parses through it; the constructor stays flat, with the -sub-config fields prefixed `unet_` / `vae_` / `text_`. - -### `ZMDiffusion` - -The pipeline class of every diffusion family (see [Diffusers](diffusers.md)); a repo's -`model_type` (`"stable_diffusion"` here) selects the SD 1.x family, which supplies the -hooks `BaseDiffusion`'s `generate` needs (`encode_prompt` on the CLIP tower, the empty -prompt's `unconditional_ids`, `predict_noise` on the UNet, `decode_latents` on the VAE) -and a `scheduler`, built from the config's `scheduler_config` (PNDM for every SD 1.x -checkpoint). - -```python -generate( - input_ids, - attention_mask=None, - negative_input_ids=None, - num_inference_steps=None, - guidance_scale=None, - seed=None, - latents=None, -) -``` - -| Arg | Default | Meaning | -|---|---|---| -| `input_ids` | required | `(batch, 77)` token ids, `**tokenizer(prompts)` | -| `attention_mask` | `None` | accepted alongside the tokenizer's output; SD attends its padding tokens, so it is unused | -| `negative_input_ids` | `None` | `(batch, 77)` tokenized negative prompt for guidance; the empty prompt when unset | -| `num_inference_steps` | `None` | scheduler steps; the repo's `generate_args` (50) when unset | -| `guidance_scale` | `None` | classifier-free guidance strength; `generate_args` (7.5) when unset, `<= 1` disables it | -| `seed` | `None` | seed for the initial latent, reproducible per backend | -| `latents` | `None` | explicit `(batch, 64, 64, 4)` initial latent; identical results across backends | - -Returns `(batch, 512, 512, 3)` uint8 numpy images. The denoiser call (the UNet plus the -guidance combine) runs as a per-backend compiled function cached on the instance, -`jax.jit` on JAX and `tf.function(jit_compile=True)` on TensorFlow, eager on Torch; its -shapes are the same at every step, so one compile serves the whole loop. - -| Constructor arg | Default | Meaning | -|---|---|---| -| `model_type` | required | `"stable_diffusion"` (read from the repo's `zm_config.json` on `from_weights`) | -| `scheduler` | `None` | a `BaseScheduler`; built from `scheduler_config` when unset | -| `unet_sample_size` | `64` | latent size the UNet graph is built for (image size / 8) | -| `vae_sample_size` | `512` | image size the VAE graphs are built for | -| `unet_*` / `vae_*` / `text_*` | SD 1.x | the flat sub-config fields (see [Configuration](configuration.md)) | -| `bos_token_id` / `eos_token_id` / `pad_token_id` | `49406` / `49407` / `49407` | CLIP's special tokens, used to build the empty prompt | -| `scheduler_config` | `None` | the diffusers scheduler dict (`_class_name`, betas, `steps_offset`, ...) | - -The graph is three disconnected paths, one per component. Inputs -`{"sample", "timestep", "encoder_hidden_states"}` (UNet), `{"image", "latent"}` (VAE -encode / decode) and `{"token_ids", "padding_mask"}` (text); outputs -`{"noise_pred", "moments", "image", "text_embeds"}`. The components are exposed as -`.unet`, `.vae` and `.text_encoder` and share the model's weights, so one load fills all -three, and each can be called on its own. - -### `UNet2DConditionModel` - -The denoiser (diffusers' `UNet2DConditionModel`): a symmetric down / mid / up UNet whose -ResNet blocks are conditioned on the sinusoidal timestep embedding and whose -`CrossAttn` blocks attend the latent to the text context. Inputs -`{"sample": (B, h, w, 4), "timestep": (B,), "encoder_hidden_states": (B, 77, 768)}`, -output `{"sample": (B, h, w, 4)}`. - -| Arg | Default | Meaning | -|---|---|---| -| `sample_size` | `64` | latent spatial size the graph is built for | -| `in_channels` / `out_channels` | `4` | latent channels | -| `down_block_types` / `up_block_types` | SD 1.x | `CrossAttnDownBlock2D` x3 + `DownBlock2D`, mirrored | -| `block_out_channels` | `(320, 640, 1280, 1280)` | width per level | -| `layers_per_block` | `2` | ResNet blocks per level | -| `cross_attention_dim` | `768` | width of the text context | -| `num_attention_heads` | `8` | heads in the cross-attention blocks | -| `norm_num_groups` | `32` | GroupNorm groups | -| `text_seq_len` | `77` | static text length | - -### `AutoencoderKL` - -The VAE (diffusers' `AutoencoderKL`): `encode(image, sample=False)` returns the posterior -mean (or a sample) of a `(B, H/8, W/8, 4)` latent and `decode(latent)` the `(B, H, W, 3)` -image in `[-1, 1]`; multiply / divide by `scaling_factor` (0.18215) around the UNet, as -`generate` does. One functional graph with the encoder and decoder as two disconnected -paths. - -| Arg | Default | Meaning | -|---|---|---| -| `in_channels` / `out_channels` | `3` | image channels | -| `latent_channels` | `4` | latent channels | -| `block_out_channels` | `(128, 256, 512, 512)` | width per level, x8 spatial compression | -| `layers_per_block` | `2` | ResNet blocks per encoder level | -| `norm_num_groups` | `32` | GroupNorm groups | -| `sample_size` | `512` | image size the graphs are built for | -| `scaling_factor` | `0.18215` | latent scaling | - -### Schedulers - -`zeromodels.base.base_scheduler` has the SD samplers as weightless classes with the -diffusers interface (`set_timesteps`, `scale_model_input`, `step`, -`init_noise_sigma`), configured the diffusers way: - -| Scheduler | Notes | -|---|---| -| `PNDMScheduler` | the SD 1.x default (PLMS, `skip_prk_steps=True`) | -| `DDIMScheduler` | deterministic, `eta=0` | -| `EulerDiscreteScheduler` | Karras-style sigma sampler, `linspace` spacing | -| `EulerAncestralDiscreteScheduler` | stochastic Euler (`seed` for reproducibility) | - -`get_scheduler(config)` builds the one named by a diffusers scheduler config dict, which -is what the model does with its `scheduler_config`; swap at any time: - -```python -from zeromodels.base.base_scheduler import EulerDiscreteScheduler - -model.scheduler = EulerDiscreteScheduler.from_config(model.config.scheduler_config) -``` - -## Preprocessing - -### `ZMDiffusionTokenizer` - -The tokenizer of every diffusion family; for `model_type="stable_diffusion"` it is the -CLIP ViT-L/14 BPE tokenizer: `<|startoftext|>` / `<|endoftext|>` framing, truncated and -`<|endoftext|>`-padded to 77 tokens, the fixed context length the UNet's cross-attention -is built for. Calling it returns `{"input_ids", "attention_mask"}` for a string or a list -of strings. `from_weights("zeromodels/")` reads the family from the repo. - -```python -ZMDiffusionTokenizer( - model_type="stable_diffusion", hf_id=None, tokenizer_file=None, max_seq_len=77 -) -``` - -There is no image processor and no processor class: text-to-image takes token ids in -and hands uint8 images out. - -## End-to-end example - -### Single prompt - -```python -import os - -os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" - -from PIL import Image -from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer - -model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-v1-5") -tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-v1-5") - -inputs = tokenizer( - "a red fox sitting in a field of lavender, golden hour, wildlife photography, sharp focus" -) -images = model.generate(**inputs, num_inference_steps=50, guidance_scale=7.5, seed=1) - -Image.fromarray(images[0]).save("fox.png") # (512, 512, 3) uint8 -``` - -Stable Diffusion v1-5: a red fox sitting in a field of lavender at golden hour, 512px - -### Negative prompt and batching - -Prompts batch through the tokenizer; a negative prompt is tokenized the same way and -passed as `negative_input_ids` (one per prompt, or one broadcast row). - -```python -prompts = [ - "a watercolor painting of a lighthouse at sunset", - "a red vintage car on a coastal road", -] -inputs = tokenizer(prompts) # {"input_ids": (2, 77), "attention_mask": (2, 77)} -negative = tokenizer(["blurry, low quality"] * len(prompts))["input_ids"] - -images = model.generate(**inputs, negative_input_ids=negative, num_inference_steps=30) -for i, image in enumerate(images): - Image.fromarray(image).save(f"out_{i}.png") -``` - -### Reproducible across backends - -`seed` reproduces a run on a given backend, but `keras.random` differs between torch, -jax and tensorflow. For an image that is identical everywhere, pass the initial latent -yourself (`(batch, 64, 64, 4)` at 512px): - -```python -import numpy as np - -latents = np.random.default_rng(0).standard_normal((1, 64, 64, 4)).astype("float32") -images = model.generate(**tokenizer("a bowl of ramen"), latents=latents) -``` - -### Verified against diffusers - -The fox prompt above with the same initial latent, PNDM 50 steps and guidance 7.5 -through `ZMDiffusion` (left) and diffusers' `StableDiffusionPipeline` -(right), both fp32: - -zeromodels (left) and diffusers (right) generations of a red fox in lavender from the same latent - -``` -tokenizer input_ids identical -text_encoder max|d|=9.5e-06 -unet noise_pred (t=981) max|d|=9.5e-07 -vae decode max|d|=3.5e-06 -final image (uint8) max|d|=1 mean|d|=0.0022 PSNR=74.7 dB identical_pixels=99.34% -``` - -### Components on their own - -```python -from zeromodels.diffusers import ZMDiffusion - -sd = ZMDiffusion.from_weights("zeromodels/stable-diffusion-v1-5") -latent = sd.vae.encode(image) # (1, 512, 512, 3) in [-1, 1] -> (1, 64, 64, 4) -text = sd.text_encoder({"token_ids": ids, "padding_mask": mask})["last_hidden_state"] -``` - -### Other resolutions - -The graphs are built for a fixed size, the weights are not. Rebuild at another size with -the constructor overrides (multiples of 64px; SD 1.x was trained at 512px): - -```python -model = ZMDiffusion.from_weights( - "zeromodels/stable-diffusion-v1-5", unet_sample_size=96, vae_sample_size=768 -) -images = model.generate(**tokenizer("a mountain lake at dawn")) # (1, 768, 768, 3) -``` - -## Data Format - -**Both `channels_last` and `channels_first` are supported.** The models read -`keras.config.image_data_format()` when they are constructed (there is no -`data_format` argument), so set it once, before `from_weights`, and every tensor that -crosses the model boundary follows it: - -| | `channels_last` (default) | `channels_first` | -|---|---|---| -| `latents` passed to `generate` | `(batch, 64, 64, 4)` | `(batch, 4, 64, 64)` | -| `AutoencoderKL.encode` input / `decode` output | `(batch, 512, 512, 3)` | `(batch, 3, 512, 512)` | -| UNet `sample` | `(batch, 64, 64, 4)` | `(batch, 4, 64, 64)` | - -`generate` itself always hands back `(batch, height, width, 3)` uint8 images, whatever -the layout, so the PIL step is the same in both. The same hosted `model.weights.h5` -loads under either layout (conv kernels are `(kh, kw, in, out)` in both), and the two -runs agree to within 1 uint8 level (99.7% of pixels identical, the rest being the -NCHW / NHWC convolution kernels' rounding). - -```python -import keras - -keras.config.set_image_data_format("channels_first") - -model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-v1-5") -images = model.generate( - **tokenizer("a bowl of ramen"), seed=0 -) # still (1, 512, 512, 3) -``` - -Note that `keras.config.set_image_data_format` is global state. Set it once at the top -of a script rather than toggling it between calls, since already-built models keep the -layout they were constructed with. - -## Memory and speed - -The fp32 container is 3.97 GB; a 512px, batch-1, guided run peaks at 6.5 GB of GPU -memory (the guided UNet call runs batch 2). On an RTX 4060 laptop the 50-step run takes -~25 s in eager Torch, about 1.5x diffusers' eager time. On JAX the compiled denoiser step -measured 2-3x faster than an op-by-op run; TensorFlow uses the same mechanism. - -## Loading Fine-tuned Weights - -The five hosted checkpoints are the supported weights; any repo laid out like them -(`zm_config.json` declaring `ZMDiffusion` with `model_type: "stable_diffusion"`, -`model.weights.h5`, `tokenizer.json`) loads with `from_weights("/")`; configs -written before `ZMDiffusion` (declaring `StableDiffusionModel`) load the same way. The -`hf:` prefix raises for diffusion models: -convert a diffusers-format checkpoint once with -`zeromodels/diffusers/stable_diffusion/convert_stable_diffusion_diffusers_to_keras.py` -(`transfer_stable_diffusion(repo)`, `pip install zeromodels[conversion]`) and host the result. diff --git a/docs/stable_diffusion_2.md b/docs/stable_diffusion_2.md deleted file mode 100644 index 5f25701d..00000000 --- a/docs/stable_diffusion_2.md +++ /dev/null @@ -1,216 +0,0 @@ -# Stable Diffusion 2 - -
-Weights: pretrained Keras weights live on Hugging Face under -zeromodels/<variant> -(each repo carries zm_config.json + model.weights.h5 + -tokenizer.json). Load with from_weights("zeromodels/<variant>"). -
- -Stable Diffusion 2.x, ported to pure Keras 3: the [Stable Diffusion](stable_diffusion.md) -latent-diffusion architecture with the second-generation conditioning. It loads -through [`ZMDiffusion`](diffusers.md) with `model_type="stable_diffusion_2"`, the SD 1.x -family with the SD 2 configuration, so everything on that page applies: one functional -model and one weights file per repo, `generate` from `BaseDiffusion`, the schedulers, -both data formats. What changes is the config: - -- **Text encoder**: OpenCLIP ViT-H/14's text tower, 1024 wide, 16 heads, `gelu`, shipped - truncated to its 23rd (penultimate) layer, which SD 2 conditions on. -- **UNet**: 1024-d cross-attention, one head count per level (`(5, 10, 20, 20)`, a 64-wide - head everywhere, where SD 1.x used 8 heads at every level) and a linear token - projection in the Transformer2D blocks instead of the 1x1 conv. -- **Tokenizer**: the same CLIP BPE, padded with `!` (id 0) the OpenCLIP way rather than - `<|endoftext|>`, so the empty prompt used for classifier-free guidance is - `[<|startoftext|>, <|endoftext|>, !, !, ...]`. -- **768px checkpoints** (`stable-diffusion-2`, `stable-diffusion-2-1`): built at a 96x96 - latent and trained with the **v-prediction** objective; their repos carry a - v-prediction DDIM scheduler, which `generate` picks up from `scheduler_config`. -- **SD-Turbo** (`sd-turbo`): SD 2.1 distilled with Adversarial Diffusion Distillation to - generate in 1 to 4 steps **without guidance**; its repo carries an Euler scheduler with - `trailing` timestep spacing and `generate_args` of 1 step, `guidance_scale=0.0`. - -The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion is -deliberately **not supported** for diffusion models. The original `stabilityai/*` repos are -no longer on the Hub; the conversion sources are the `sd2-community` mirrors. - -Links: - -- Paper: [High-Resolution Image Synthesis with Latent Diffusion Models (arXiv:2112.10752)](https://arxiv.org/abs/2112.10752) -- Reference implementation: [diffusers `StableDiffusionPipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img) -- License: [CreativeML Open RAIL++-M](https://huggingface.co/sd2-community/stable-diffusion-2-1/blob/main/LICENSE-MODEL) - -See also [stable_diffusion.md](stable_diffusion.md), [clip.md](clip.md). - -## Variants - -Preconverted, float32 weights are hosted under `zeromodels/`. Load with -`from_weights("zeromodels/")`. Each repo is one container: UNet (866M) + VAE -(84M) + OpenCLIP text encoder (340M), 1.29B parameters, 5.16 GB (4.81 GiB, one -`model.weights.h5`). The SD 2 checkpoints are released under the CreativeML Open RAIL++-M -license; SD-Turbo under the Stability AI Non-Commercial Research Community License. - -| Variant | Hub | Resolution | Objective | Training | -|---|---|---|---|---| -| `stable-diffusion-2-base` | [`zeromodels/stable-diffusion-2-base`](https://huggingface.co/zeromodels/stable-diffusion-2-base) | 512 | 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) | 768 | 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) | 512 | 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) | 768 | 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) | 512 | epsilon, 1 to 4 steps, no guidance | SD 2.1 distilled with Adversarial Diffusion Distillation (non-commercial) | - -Use the `-base` checkpoints for 512px images and the others for 768px; each repo's -`zm_config.json` builds the graph at its native size. - -## API - -Configs are typed: `StableDiffusion2Config` (composite, `model_type` -`"stable_diffusion_2"`) over `StableDiffusion2UNetConfig` (a `UNet2DConditionConfig` -with the SD 2 widths), `AutoencoderKLConfig` and `StableDiffusion2TextConfig` (a -`CLIPTextConfig` with the ViT-H/14 sizes), plus the scheduler config and the token ids -(`pad_token_id` 0). Flat constructor, `unet_` / `vae_` / `text_` prefixes, like SD 1.x. - -### `ZMDiffusion` - -`ZMDiffusion` with `model_type="stable_diffusion_2"` (read from the repo); `generate` is -the SD 1.x one: - -```python -generate( - input_ids, - attention_mask=None, - negative_input_ids=None, - num_inference_steps=None, - guidance_scale=None, - seed=None, - latents=None, -) -``` - -Returns `(batch, H, W, 3)` uint8 images at the checkpoint's resolution; `latents` is -`(batch, 64, 64, 4)` for the 512px checkpoints and `(batch, 96, 96, 4)` for the 768px -ones (`channels_first`: channels second). Defaults come from the repo's `generate_args` -(50 steps, guidance 7.5). See [Stable Diffusion](stable_diffusion.md#zmdiffusion) -for the argument table. - -The graph is SD 1.x's three disconnected paths (`.unet`, `.vae`, `.text_encoder`) with -the same inputs / outputs, `encoder_hidden_states` 1024 wide. - -The components are the SD 1.x classes, `UNet2DConditionModel` (built with -`num_attention_heads=(5, 10, 20, 20)`, `use_linear_projection=True`, -`cross_attention_dim=1024`) and `AutoencoderKL`; see their tables on the -[Stable Diffusion](stable_diffusion.md#api) page. - -## Preprocessing - -### `ZMDiffusionTokenizer` - -`ZMDiffusionTokenizer` with `model_type="stable_diffusion_2"`, which pads with `!`: -CLIP BPE, `<|startoftext|>` / `<|endoftext|>` framing, truncated and `!`-padded to 77 -tokens. Returns `{"input_ids", "attention_mask"}`. - -```python -ZMDiffusionTokenizer( - model_type="stable_diffusion_2", hf_id=None, tokenizer_file=None, max_seq_len=77 -) # pad_token defaults to "!" -``` - -## End-to-end example - -```python -import os - -os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" - -from PIL import Image -from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer - -model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-2-1-base") -tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-2-1-base") - -inputs = tokenizer( - "a steaming bowl of ramen on a wooden table, food photography, shallow depth of field" -) -images = model.generate(**inputs, num_inference_steps=50, guidance_scale=7.5, seed=2) - -Image.fromarray(images[0]).save("ramen.png") # (512, 512, 3) uint8 -``` - -Stable Diffusion 2.1-base: a steaming bowl of ramen on a wooden table, 512px - -### 768px, v-prediction - -The 768px checkpoints need no extra arguments; the v-prediction DDIM scheduler and the -96x96 latent come from the repo: - -```python -model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-2-1") -tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-2-1") -images = model.generate( - **tokenizer("a lighthouse on a cliff at dusk, oil painting") -) # (1, 768, 768, 3) -``` - -### SD-Turbo - -One step, no guidance (the repo's defaults; up to 4 steps sharpen a little): - -```python -model = ZMDiffusion.from_weights("zeromodels/sd-turbo") -tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/sd-turbo") -images = model.generate( - **tokenizer( - "a cinematic shot of a baby raccoon wearing an intricate italian priest robe" - ) -) -``` - -SD-Turbo, one step: a baby raccoon in an italian priest robe, 512px - -The same prompt and latent through diffusers' `StableDiffusionPipeline` (fp32) match to -1 uint8 level: 99.9% of pixels identical at 1 step (PSNR 82 dB), 99.8% at 4 steps. - -Negative prompts, batching, explicit `latents` for cross-backend reproducibility, other -resolutions and the components on their own work exactly as on the -[Stable Diffusion](stable_diffusion.md#end-to-end-example) page. Image-to-image -(`generate(..., image=..., strength=...)`) works as described for -[`BaseDiffusion`](main_classes.md#basediffusion). - -### Verified against diffusers - -The ramen prompt above with the same initial latent, PNDM 50 steps and guidance 7.5 -through `ZMDiffusion` (left) and diffusers' `StableDiffusionPipeline` -(right) on `stable-diffusion-2-1-base`, both fp32: - -zeromodels (left) and diffusers (right) generations of a bowl of ramen from stable-diffusion-2-1-base with the same latent - -``` -tokenizer input_ids identical -text_encoder max|d|=1.5e-05 -unet noise_pred (t=981) max|d|=7.2e-07 -vae decode max|d|=2.6e-06 -final image (uint8) max|d|=1 mean|d|=0.0095 PSNR=68.3 dB identical_pixels=97.22% -``` - -The 768px v-prediction checkpoint (`stable-diffusion-2-1`, DDIM 20 steps, same latent) -matches the same way: UNet 8.3e-6, VAE 7.9e-5, final image max 1 uint8 level, 99.3% of -pixels identical, PSNR 74.6 dB. - -## Data Format - -Both `channels_last` and `channels_first` are supported, as for -[Stable Diffusion](stable_diffusion.md#data-format); `generate` always returns -`(batch, H, W, 3)` uint8. - -## Memory and speed - -The fp32 container is 5.16 GB. At 512px a guided batch-1 run needs a little over 7 GB of -GPU memory in eager Torch (about 1.5x diffusers' eager time); at 768px the self-attention -runs over 9216 tokens, so the guided run needs roughly 12 GB, or the CPU. - -## Loading Fine-tuned Weights - -Any repo laid out like the hosted ones (`zm_config.json` declaring -`ZMDiffusion` with `model_type: "stable_diffusion_2"`, the weights, `tokenizer.json`) -loads with `from_weights("/")`. The `hf:` prefix raises for diffusion models: convert a -diffusers-format SD 2 checkpoint once with -`zeromodels/diffusers/stable_diffusion_2/convert_stable_diffusion_2_diffusers_to_keras.py` -(`transfer_stable_diffusion_2(repo)`, `pip install zeromodels[conversion]`) and host the result. diff --git a/docs/stable_diffusion_3.md b/docs/stable_diffusion_3.md deleted file mode 100644 index 26e67feb..00000000 --- a/docs/stable_diffusion_3.md +++ /dev/null @@ -1,266 +0,0 @@ -# Stable Diffusion 3 - -
-Weights: pretrained Keras weights live on Hugging Face under -zeromodels/<variant> -(each repo carries zm_config.json + the float16 weights + -tokenizer.json / tokenizer_3.json). Load with -from_weights("zeromodels/<variant>"). -
- -Stable Diffusion 3, ported to pure Keras 3: a rectified-flow model whose denoiser is the -**MMDiT**, a multimodal diffusion transformer in which the latent tokens and the text tokens -attend to each other jointly, conditioned on the timestep and the pooled text embeddings -through adaptive layer norms. It loads through [`ZMDiffusion`](diffusers.md) with -`model_type="stable_diffusion_3"` and keeps the SD family's shape (one functional model -and one set of weights per repo, `generate` from `BaseDiffusion`, the schedulers, both -data formats) with new parts: - -- **MMDiT** (`SD3Transformer2DModel`): the 16-channel latent is patchified (2x2) into 1536-d - tokens with a sinusoidal position table cropped from a 192x192 grid; 24 joint blocks of - 24 heads, each with AdaLN-Zero modulation of both streams, joint attention over the - concatenated latent + text tokens, and gated GELU feed-forwards; a final adaptive norm - and projection unpatchify the velocity. 2B parameters. -- **Three text encoders**: CLIP ViT-L/14 (768-d) and OpenCLIP ViT-bigG/14 (1280-d), each - with a projection, and **T5-XXL** (4.7B, 4096-d, 256 tokens). The two CLIP penultimate - states are concatenated (2048-d), zero-padded to 4096 and followed along the sequence by - the T5 states (77 + 256 tokens); the projected pooled CLIP states (2048-d) feed the - conditioning embedding. -- **T5 is optional and separate**: the checkpoint holds the MMDiT, the VAE and the two CLIP - towers (2.9B, 5.8 GB float16); the T5-XXL encoder is hosted once for every SD 3 / 3.5 - checkpoint (`zeromodels/t5-v1_1-xxl-encoder`, an - [`SD3T5EncoderModel`](#sd3t5encodermodel), 9.5 GB float16) and attached with - `text_encoder_3=`. Without it the T5 features are zeros, SD 3's documented - memory-saving mode (prompt adherence drops, the images stay good). -- **VAE**: 16 latent channels, no quant convolutions, `scaling_factor` 1.5305 and - `shift_factor` 0.0609 (`z = (x - shift) * scale`), built in float32 (`force_upcast`). -- **Sampler**: `FlowMatchEulerDiscreteScheduler` with `shift` 3.0 (rectified flow: the - timesteps are `sigma * 1000`, a step is `x + (sigma_next - sigma) * v`); 28 steps at - guidance 7.0 by default. No negative prompt means the encoded empty prompt. -- **One tokenizer call**: `ZMDiffusionTokenizer` runs the CLIP BPE (the second - tower's `!`-padded ids are derived from the mask, as SDXL) and the T5 SentencePiece - tokenizer, returning `{"input_ids", "attention_mask", "input_ids_3"}`. -- **float16 weights**, the release precision; `from_weights` builds the model in float16 - by default (`load_dtype="float32"` for float32). - -The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion is -deliberately **not supported** for diffusion models. The upstream repo is gated behind the -Stability AI Non-Commercial Research Community License. - -Links: - -- Paper: [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (arXiv:2403.03206)](https://arxiv.org/abs/2403.03206) -- Reference implementation: [diffusers `StableDiffusion3Pipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_3) -- License: [Stability AI Non-Commercial Research Community License](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers/blob/main/LICENSE.md) - -See also [stable_diffusion_3_5.md](stable_diffusion_3_5.md), [stable_diffusion_xl.md](stable_diffusion_xl.md), [clip.md](clip.md). - -## Variants - -| Variant | Hub | Resolution | Sampler | License | -|---|---|---|---|---| -| `stable-diffusion-3-medium` | [`zeromodels/stable-diffusion-3-medium`](https://huggingface.co/zeromodels/stable-diffusion-3-medium) | 1024 | flow-match Euler (shift 3), 28 steps, guidance 7.0 | Stability AI Non-Commercial Research Community | -| `t5-v1_1-xxl-encoder` | [`zeromodels/t5-v1_1-xxl-encoder`](https://huggingface.co/zeromodels/t5-v1_1-xxl-encoder) | text encoder (shared with SD 3.5) | | Apache 2.0 (Google T5 v1.1) | - -## API - -Configs are typed: `StableDiffusion3Config` (composite, `model_type` `"stable_diffusion_3"`) -over `StableDiffusion3TransformerConfig`, `StableDiffusion3VAEConfig`, -`StableDiffusion3TextConfig` (CLIP ViT-L/14 with its 768-d projection) and -`StableDiffusionXLTextConfig2` (OpenCLIP ViT-bigG/14 with its 1280-d projection), plus -the scheduler config, the CLIP and T5 token ids and `max_sequence_length` (256). Flat -constructor with `transformer_` / `vae_` / `text_` / `text_2_` prefixes. - -### `ZMDiffusion` - -`ZMDiffusion` with `model_type="stable_diffusion_3"` (read from the repo); `generate` -over the SD 3 graph: - -```python -generate( - input_ids, - attention_mask=None, - negative_input_ids=None, - num_inference_steps=None, - guidance_scale=None, - seed=None, - latents=None, - image=None, - strength=None, - denoising_start=None, - denoising_end=None, - output_type="image", - input_ids_3=None, - negative_input_ids_3=None, -) -``` - -| Argument | Description | -|---|---| -| `input_ids`, `attention_mask`, `input_ids_3` | `**tokenizer(prompts)`: the CLIP ids and mask (the second tower's ids are derived from the mask) and the T5 ids (`(batch, 256)`). Without `input_ids_3`, or without an attached T5, the T5 features are zeros. | -| `negative_input_ids`, `negative_input_ids_3` | The tokenized negative prompt (`neg = tokenizer("..."); negative_input_ids=neg["input_ids"], negative_input_ids_3=neg["input_ids_3"]`); the empty prompt when omitted. | -| `num_inference_steps`, `guidance_scale`, `seed`, `latents`, `image`, `strength`, ... | As for [`BaseDiffusion.generate`](main_classes.md#basediffusion); defaults from the repo's `generate_args` (28 steps, guidance 7.0). `latents` is `(batch, 128, 128, 16)` at 1024px. | - -Returns `(batch, H, W, 3)` uint8 images. - -`from_weights` takes one extra argument, `text_encoder_3`: a hosted `SD3T5EncoderModel` -repo id (loaded at the same `load_dtype`) or a built `SD3T5EncoderModel`; -`sd.text_encoder_3` can also be assigned later (or set to `None`). It lives outside the -checkpoint's weights; passing `text_encoder_3` for a family other than SD 3 / 3.5 raises. - -The graph is four disconnected paths in one functional model, `.transformer` / `.vae` / -`.text_encoder` / `.text_encoder_2`. - -| Inputs | Outputs | -|---|---| -| `sample` (B, 128, 128, 16), `timestep` (B,), `encoder_hidden_states` (B, 333, 4096), `pooled_projections` (B, 2048) | `noise_pred` (B, 128, 128, 16) | -| `image` (B, 1024, 1024, 3), `latent` (B, 128, 128, 16) | `moments` (B, 128, 128, 32), `image` (B, 1024, 1024, 3) | -| `token_ids` (B, 77), `token_ids_2` (B, 77), `padding_mask` (B, 77) | `prompt_embeds` (B, 77, 2048), `pooled_prompt_embeds` (B, 2048) | - -### `SD3Transformer2DModel` - -The MMDiT on its own (`{"sample", "timestep", "encoder_hidden_states", -"pooled_projections"} -> {"sample"}`), built from `StableDiffusion3TransformerConfig` -(`sample_size` 128, `patch_size` 2, `num_layers` 24, `num_attention_heads` 24 x -`attention_head_dim` 64, `joint_attention_dim` 4096, `caption_projection_dim` 1536, -`pooled_projection_dim` 2048, `pos_embed_max_size` 192, `qk_norm`, -`dual_attention_layers`). Its layers (`StableDiffusion3PatchEmbed`, -`StableDiffusion3AdaLayerNorm`, `StableDiffusion3JointAttention`, -`StableDiffusion3JointTransformerBlock`, `StableDiffusion3GELUFeedForward`) are named -with the diffusers module paths. - -The joint attention runs over 4429 tokens at 1024px, so `StableDiffusion3JointAttention` -defaults to the `"fused"` attention implementation (`keras.ops.dot_product_attention`, -the backend's own kernel: torch's flash / memory-efficient kernels never materialize the -4429 x 4429 logits; see [`fused_attention`](main_classes.md#fused_attention)) instead of -the library-wide `"sdpa"` math. `from_weights(..., attn_implementation="sdpa")` switches -it back. - -### `SD3T5EncoderModel` - -The third text encoder, the encoder half of T5 v1.1 XXL (transformers `T5EncoderModel` on -`google/t5-v1_1-xxl`), as its own model in this family: `{"input_ids", "attention_mask"}` -(`(batch, 256)` int32) `-> {"last_hidden_state": (batch, 256, 4096)}`, built from -`StableDiffusion3T5EncoderConfig` (`vocab_size` 32128, `embed_dim` 4096, `key_value_dim` -64, `mlp_dim` 10240, `num_layers` 24, `num_heads` 64, `relative_attention_num_buckets` 32, -`relative_attention_max_distance` 128, `layer_norm_eps` 1e-6; 4.7B parameters). The blocks -reuse the [T5](t5.md) family's attention, RMSNorm and relative position bias, and add the -gated-GELU feed-forward of T5 v1.1 (`wo(gelu_tanh(wi_0(x)) * wi_1(x))`, -`StableDiffusion3T5GatedFeedForward`), which the original T5 does not have. Loads on its own with -`SD3T5EncoderModel.from_weights("zeromodels/t5-v1_1-xxl-encoder")` (`load_dtype`, -`quantization="int8"` as needed); its inputs are the `input_ids_3` of the SD 3 tokenizer. - -## Preprocessing - -### `ZMDiffusionTokenizer` - -```python -ZMDiffusionTokenizer( - model_type="stable_diffusion_3", - hf_id=None, - tokenizer_file=None, - tokenizer_file_3=None, - max_seq_len=77, - max_sequence_length=256, -) -``` - -Returns `{"input_ids", "attention_mask", "input_ids_3"}`; pass all three to `generate`. - -## End-to-end example - -```python -import os - -os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" - -from PIL import Image -from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer - -model = ZMDiffusion.from_weights( - "zeromodels/stable-diffusion-3-medium", - text_encoder_3="zeromodels/t5-v1_1-xxl-encoder", # omit on small machines: T5 features zeroed -) -tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-3-medium") - -inputs = tokenizer( - "a fluffy corgi wearing round sunglasses, sitting on a surfboard at a sunny beach, photo" -) -images = model.generate(**inputs, num_inference_steps=28, guidance_scale=7.0, seed=4) - -Image.fromarray(images[0]).save("corgi.png") # (1024, 1024, 3) uint8 -``` - -Stable Diffusion 3 medium: a corgi with sunglasses on a surfboard, 1024px - -### The T5 encoder on the CPU - -The 9.5 GB T5-XXL rarely fits next to the MMDiT on a consumer GPU. Encode the prompt -separately (an `SD3T5EncoderModel` in another process, or int8) and hand the ids to a -model without T5 for the CLIP part only, or attach it in full on a large machine. The -`encode_prompt` hook returns the conditioning dict the transformer takes, so any -precomputed T5 features can be substituted there. - -### Verified against diffusers - -The components of `stable-diffusion-3-medium` were checked one by one against diffusers -0.39 / transformers 5 in float32 on the same inputs (the converted float16 weights -computed in float32): - -``` -text_encoder penultimate max|d|=6.1e-05 (values up to 854) -text_encoder_2 penultimate max|d|=3.1e-05 -pooled (L + G, projected) max|d|=3.8e-06 -t5-v1_1-xxl encoder (256 tok) max|d|=1.1e-04 (values up to 6.3) -mmdit velocity (t=1000) max|d|=9.9e-06 -vae decode (512px) max|d|=6.3e-05 -vae encode max|d|=1.4e-05 -``` - -The pipeline (flow-match Euler with shift 3, classifier-free guidance with the encoded -empty prompt, the three tokenizers, batched negative prompts, the no-T5 mode) was checked -against `StableDiffusion3Pipeline` on a small random-weight SD 3, where the generated images -are pixel-identical on torch and jax. End to end, `stable-diffusion-3-medium` without the -T5 encoder (28 steps, guidance 7, the same initial latent) reproduces the diffusers -pipeline in float32 at 512px: 92.8% identical pixels, max |d| 2 / 255, PSNR 63.9 dB (the -final latent differs by at most 2.5e-2 with values up to 5.1). In float16 on the GPU the -two implementations round differently through the 28 guided steps: the pictures agree in -composition but not pixel for pixel (PSNR 26 dB at 1024px, 29 dB at 768px, float16 MMDiT -on both sides), the usual float16 sampling noise. - -## Data Format - -Both `channels_last` and `channels_first` are supported; `generate` always returns -`(batch, H, W, 3)` uint8. - -## Memory and speed - -The checkpoint is 5.8 GB in float16 (the T5 encoder adds 9.5 GB). Measured on the torch -backend with the default float16 load (RTX 4060 Laptop, 8 GB), T5 features precomputed: - -| Resolution | Denoising step (guided, batch 1) | VAE decode (float32) | -|---|---|---| -| 512px | peak 6.1 GB, 0.5 s/step | peak 7.6 GB, 4 s | -| 768px | peak 6.4 GB, 0.8 s/step | peak 9.5 GB | -| 1024px | peak 6.9 GB, 1.0 s/step | peak 12.4 GB | - -The MMDiT itself fits an 8 GB GPU at the native 1024px thanks to the fused attention -kernel (the `"sdpa"` math needs 3.8 GB per block for the 4429 x 4429 float32 logits and -runs out of memory). The VAE decoder is the SDXL one (16 latent channels aside) and Keras' -functional executor keeps its feature maps alive, see the -[SDXL notes](stable_diffusion_xl.md#memory-and-speed): on 8 GB, denoise on the GPU -(`output_type="latent"`) and decode on the CPU (75 to 90 s at 1024px). The T5-XXL encoder -encodes a prompt and its negative on the CPU in about 10 to 25 s (float32 compute over -the float16 weights); a 24 GB GPU holds everything. On JAX and TensorFlow the denoiser -step is compiled once per run. - -## Loading Fine-tuned Weights - -Any repo laid out like the hosted ones (`zm_config.json` declaring -`ZMDiffusion` with `model_type: "stable_diffusion_3"`, the weights, the two tokenizer -files) loads with -`from_weights("/")`. The `hf:` prefix raises for diffusion models: convert a -diffusers-format SD 3 checkpoint once with -`zeromodels/diffusers/stable_diffusion_3/convert_stable_diffusion_3_diffusers_to_keras.py` -(`transfer_stable_diffusion_3(repo)`; `transfer_t5_encoder(repo)` for the T5 tower) and host the -result. diff --git a/docs/stable_diffusion_3_5.md b/docs/stable_diffusion_3_5.md deleted file mode 100644 index dc246c7b..00000000 --- a/docs/stable_diffusion_3_5.md +++ /dev/null @@ -1,121 +0,0 @@ -# Stable Diffusion 3.5 - -
-Weights: pretrained Keras weights live on Hugging Face under -zeromodels/<variant> -(each repo carries zm_config.json + the float16 weights + -tokenizer.json / tokenizer_3.json). Load with -from_weights("zeromodels/<variant>"). -
- -Stable Diffusion 3.5, ported to pure Keras 3: the [Stable Diffusion 3](stable_diffusion_3.md) -architecture, retrained and scaled. It loads through -[`ZMDiffusion`](diffusers.md) with `model_type="stable_diffusion_3_5"`, the SD 3 family -with the SD 3.5 configuration, so everything on that page applies: the MMDiT, the three text encoders with the separately hosted T5-XXL, the -16-channel VAE, the flow-match sampler, one tokenizer call. What changes is the config: - -- **QK normalization**: every attention block RMS-normalizes its queries and keys per head - (`qk_norm="rms_norm"`), the training stabilizer of SD 3.5. -- **Large** (8B): 38 blocks of 38 heads (2432-d tokens), 28 steps at guidance 3.5. -- **Large Turbo**: the large model distilled with Adversarial Diffusion Distillation to 4 - steps without guidance (`guidance_scale=0.0`). -- **Medium** (2.5B, MMDiT-X): 24 blocks of 24 heads with **dual attention** in blocks 0 to - 12 (a second, latent-only attention with its own modulation next to the joint one) and a - 384x384 position grid (up to 2 MP); 40 steps at guidance 4.5. - -The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion is -deliberately **not supported** for diffusion models. The upstream repos are gated behind -the Stability AI Community License. - -Links: - -- Paper: [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (arXiv:2403.03206)](https://arxiv.org/abs/2403.03206) -- Reference implementation: [diffusers `StableDiffusion3Pipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_3) -- License: [Stability AI Community License](https://huggingface.co/stabilityai/stable-diffusion-3.5-large/blob/main/LICENSE.md) - -See also [stable_diffusion_3.md](stable_diffusion_3.md). - -## Variants - -Each repo is one checkpoint: MMDiT + VAE (84M) + CLIP ViT-L/14 (123M) + OpenCLIP -ViT-bigG/14 (695M); the T5-XXL encoder (`zeromodels/t5-v1_1-xxl-encoder`, the SD 3 -family's `SD3T5EncoderModel`, 9.5 GB) is shared and optional. - -| Variant | Hub | MMDiT | Sampler | Size | -|---|---|---|---|---| -| `stable-diffusion-3.5-large` | [`zeromodels/stable-diffusion-3.5-large`](https://huggingface.co/zeromodels/stable-diffusion-3.5-large) | 38 x 38 heads, 8B | flow-match Euler (shift 3), 28 steps, guidance 3.5 | 17.8 GB float16 | -| `stable-diffusion-3.5-large-turbo` | [`zeromodels/stable-diffusion-3.5-large-turbo`](https://huggingface.co/zeromodels/stable-diffusion-3.5-large-turbo) | 38 x 38 heads, 8B | 4 steps, no guidance | 17.8 GB float16 | -| `stable-diffusion-3.5-medium` | [`zeromodels/stable-diffusion-3.5-medium`](https://huggingface.co/zeromodels/stable-diffusion-3.5-medium) | 24 x 24 heads + dual attention, 2.5B | flow-match Euler (shift 3), 40 steps, guidance 4.5 | 6.8 GB float16 | - -## API - -`StableDiffusion3_5Config` (`model_type` `"stable_diffusion_3_5"`) is -`StableDiffusion3Config` over `StableDiffusion3_5TransformerConfig` (the SD 3.5 large -defaults: `num_layers` 38, `num_attention_heads` 38, `caption_projection_dim` 2432, -`qk_norm="rms_norm"`; medium's repo carries `dual_attention_layers=(0, ..., 12)` and -`pos_embed_max_size=384`). `ZMDiffusion.generate` and -`ZMDiffusionTokenizer` (`model_type="stable_diffusion_3_5"`) behave as for SD 3; see -[Stable Diffusion 3](stable_diffusion_3.md#api). - -## End-to-end example - -```python -import os - -os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" - -from PIL import Image -from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer - -model = ZMDiffusion.from_weights( - "zeromodels/stable-diffusion-3.5-medium", - text_encoder_3="zeromodels/t5-v1_1-xxl-encoder", # optional -) -tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-3.5-medium") - -inputs = tokenizer( - "a whimsical treehouse village at twilight, lanterns glowing, watercolor illustration" -) -images = model.generate(**inputs, num_inference_steps=40, guidance_scale=4.5, seed=5) - -Image.fromarray(images[0]).save("treehouse.png") # (1024, 1024, 3) uint8 -``` - -Stable Diffusion 3.5 medium: a treehouse village at twilight, 1024px - -### Large Turbo - -```python -model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-3.5-large-turbo") -images = model.generate( - **tokenizer("a red panda reading a book") -) # 4 steps, guidance 0 (the repo defaults) -``` - -### Verified against diffusers - -The components of `stable-diffusion-3.5-medium` (the MMDiT-X with dual attention and RMS -qk norms) were checked against diffusers 0.39 in float32 on the same inputs (the converted -float16 weights computed in float32): MMDiT velocity max|d| 1.6e-5 (values up to 5.2), the -text encoders and the VAE as for [SD 3](stable_diffusion_3.md#verified-against-diffusers). -The SD 3.5 block variants (dual attention, `qk_norm`) are also covered by the small -random-weight pipeline test, pixel-identical to `StableDiffusion3Pipeline` on torch and -jax, and by the SD 3 float32 end-to-end check (92.8% identical pixels, PSNR 63.9 dB at -512px), which the two families share down to the sampler. The picture above is the -1024px float16 output of `stable-diffusion-3.5-medium` with the T5 encoder. - -## Memory and speed - -The medium checkpoint is 6.8 GB in float16, the large ones 17.8 GB (a 24 GB GPU); the -T5-XXL encoder adds 9.5 GB. See [Stable Diffusion 3](stable_diffusion_3.md#memory-and-speed): -the medium MMDiT-X at 1024px peaks at 7.5 GB during the guided denoising step on the torch -backend (fused attention, float16), which an 8 GB GPU runs with the VAE decode on the CPU. - -## Loading Fine-tuned Weights - -Any repo laid out like the hosted ones (`zm_config.json` declaring -`ZMDiffusion` with `model_type: "stable_diffusion_3_5"`, the weights, the two tokenizer -files) loads with -`from_weights("/")`. Convert a diffusers-format checkpoint once with -`zeromodels/diffusers/stable_diffusion_3_5/convert_stable_diffusion_3_5_diffusers_to_keras.py` -and host the result. diff --git a/docs/stable_diffusion_xl.md b/docs/stable_diffusion_xl.md deleted file mode 100644 index 4d31a17b..00000000 --- a/docs/stable_diffusion_xl.md +++ /dev/null @@ -1,373 +0,0 @@ -# Stable Diffusion XL - -
-Weights: pretrained Keras weights live on Hugging Face under -zeromodels/<variant> -(each repo carries zm_config.json + the float16 weights + -tokenizer.json). Load with from_weights("zeromodels/<variant>"). -
- -Stable Diffusion XL, ported to pure Keras 3: the [Stable Diffusion](stable_diffusion.md) -latent-diffusion recipe scaled up to 1024px with a 2.6B-parameter UNet and two text -encoders, plus the **refiner**, an image-to-image expert that finishes the base model's -latents, and **SDXL-Turbo**, the base distilled to a few steps. It loads -through [`ZMDiffusion`](diffusers.md) with `model_type="stable_diffusion_xl"` (the -refiner: `"stable_diffusion_xl_refiner"`), a family built on SD 1.x's with a fourth -component and the SDXL conditioning hooks, so everything on that page applies: one -functional model and one set of weights per repo, `generate` from `BaseDiffusion`, the -schedulers, both data formats. -What changes: - -- **Two text encoders**: CLIP ViT-L/14's text tower (768-d, the SD 1.x one) and OpenCLIP - ViT-bigG/14's (1280-d, 32 layers, `gelu`, with a 1280-d `text_projection`). The prompt - goes through both; their **penultimate** hidden states (no final LayerNorm) are - concatenated into the UNet's 2048-d cross-attention context, and the second tower's - projected pooled (EOT) state is the `text_embeds` micro-conditioning. -- **One tokenizer call**: both towers use the same CLIP BPE and differ only in the pad - token (`<|endoftext|>` vs `!`), so `ZMDiffusionTokenizer` returns the first one's ids - and the model derives the second tower's ids from the returned `attention_mask`. -- **Micro-conditioning**: the UNet's timestep embedding also receives the pooled text - embedding and six size / crop values (`original_size`, `crops_coords_top_left`, - `target_size`, each `(height, width)`), sinusoidally embedded and projected by the - `add_embedding` MLP (`addition_embed_type="text_time"`). -- **UNet**: three levels (320, 640, 1280) with no attention at the first, (1, 2, 10) - transformer blocks per Transformer2D, (5, 10, 20) heads (64 wide), linear token - projection, built at a 128x128 latent (1024px). -- **VAE**: the SDXL autoencoder (`scaling_factor` 0.13025), built in **float32 whatever - the load dtype** (`force_upcast`: it overflows in float16, as in diffusers). -- **Classifier-free guidance** with no negative prompt uses **zero** embeddings - (`force_zeros_for_empty_prompt`), not the encoded empty prompt. Default guidance is 5.0. -- **Scheduler**: Euler with `leading` timestep spacing (`steps_offset` 1) for the base - models and the refiners, ancestral Euler with `trailing` spacing for SDXL-Turbo; all - come from the repo's `scheduler_config`. -- **Refiner**: a second checkpoint (a `ZMDiffusion` with - `model_type="stable_diffusion_xl_refiner"`) with the OpenCLIP tower alone (1280-d - context), a four-level UNet ((384, 768, 1536, 1536), attention at the middle levels, 4 - blocks per Transformer2D, 2.3B parameters) and an **aesthetic score** as the fifth - micro-conditioning id. It refines a latent the base left partially denoised - (`denoising_start`) or any image (`strength`); its empty negative prompt is encoded, not - zeroed. -- **float16 weights**: SDXL was released in float16, so the repos store that (the VAE in - float32) and `from_weights` builds the model in float16 by default (6.6 GB of weights). - Pass `load_dtype="float32"` for a float32 model (13.9 GB). - -The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion is -deliberately **not supported** for diffusion models. - -Links: - -- Paper: [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis (arXiv:2307.01952)](https://arxiv.org/abs/2307.01952) -- Reference implementation: [diffusers `StableDiffusionXLPipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl) -- Licenses: [CreativeML Open RAIL++-M](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/LICENSE.md) - (1.0 base and refiner), [Stability AI Non-Commercial Research Community License](https://huggingface.co/stabilityai/sdxl-turbo/blob/main/LICENSE.md) - (Turbo); the 0.9 research previews are not hosted - -See also [stable_diffusion.md](stable_diffusion.md), [stable_diffusion_2.md](stable_diffusion_2.md), [clip.md](clip.md). - -## Variants - -Preconverted weights are hosted under `zeromodels/`. Load with -`from_weights("zeromodels/")`. A base repo is one container: UNet (2.57B) + VAE -(84M) + CLIP ViT-L/14 text encoder (123M) + OpenCLIP ViT-bigG/14 text encoder (695M), -3.47B parameters, 6.6 GB in float16 (two `model.weights.json` shards); a refiner repo is -UNet (2.26B) + VAE + OpenCLIP text encoder, 3.04B parameters, 5.8 GB. - -| Variant | Hub | `model_type` | Resolution | Sampler | License | -|---|---|---|---|---|---| -| `stable-diffusion-xl-base-1.0` | [`zeromodels/stable-diffusion-xl-base-1.0`](https://huggingface.co/zeromodels/stable-diffusion-xl-base-1.0) | `stable_diffusion_xl` | 1024 | Euler, 50 steps, guidance 5.0 | CreativeML Open RAIL++-M | -| `stable-diffusion-xl-refiner-1.0` | [`zeromodels/stable-diffusion-xl-refiner-1.0`](https://huggingface.co/zeromodels/stable-diffusion-xl-refiner-1.0) | `stable_diffusion_xl_refiner` | 1024 | Euler, 50 steps, guidance 5.0, strength 0.3 | CreativeML Open RAIL++-M | -| `sdxl-turbo` | [`zeromodels/sdxl-turbo`](https://huggingface.co/zeromodels/sdxl-turbo) | `stable_diffusion_xl` | 512 | Euler ancestral, 1 step, no guidance | Stability AI Non-Commercial Research Community | - -The SDXL 0.9 research previews (`stabilityai/stable-diffusion-xl-base-0.9` and -`-refiner-0.9`, the same architecture with different weights, gated behind the research -license) are not hosted; with an authorized token the converter turns them into the same -checkpoints (`transfer_stable_diffusion_xl(repo, token=...)`, see -[Loading Fine-tuned Weights](#loading-fine-tuned-weights)). SDXL-Turbo is a -research-only, non-commercial release; its repo defaults (`generate_args`) are 1 step and -`guidance_scale=0.0`, which is how it was trained (do not add guidance). - -## API - -Configs are typed: `StableDiffusionXLConfig` (composite, `model_type` -`"stable_diffusion_xl"`) over `StableDiffusionXLUNetConfig`, `StableDiffusionXLVAEConfig`, -`StableDiffusionTextConfig` (the CLIP ViT-L/14 tower, as SD 1.x) and -`StableDiffusionXLTextConfig2` (the OpenCLIP ViT-bigG/14 tower, with `projection_dim` and -its own `hidden_act`), plus the scheduler config, the token ids (`pad_token_id_2` for the -second tokenizer's `!`), `force_zeros_for_empty_prompt` and `requires_aesthetics_score`. -Flat constructor with `unet_` / `vae_` / `text_` / `text_2_` prefixes. -`StableDiffusionXLRefinerConfig` (`model_type` `"stable_diffusion_xl_refiner"`) is the -same with `text_config=None`, the `StableDiffusionXLRefinerUNetConfig` denoiser, -`requires_aesthetics_score=True` and `force_zeros_for_empty_prompt=False`. - -### `ZMDiffusion` - -`ZMDiffusion` with `model_type="stable_diffusion_xl"` (read from the repo); `generate` -gains the micro-conditioning arguments: - -```python -generate( - input_ids, - attention_mask=None, - negative_input_ids=None, - num_inference_steps=None, - guidance_scale=None, - seed=None, - latents=None, - image=None, - strength=None, - denoising_start=None, - denoising_end=None, - output_type="image", - original_size=None, - crops_coords_top_left=(0, 0), - target_size=None, - aesthetic_score=6.0, - negative_original_size=None, - negative_crops_coords_top_left=None, - negative_target_size=None, - negative_aesthetic_score=2.5, -) -``` - -| Argument | Description | -|---|---| -| `input_ids`, `attention_mask` | `**tokenizer(prompts)`; the mask tells the model which positions are padding, so the second tower's `!`-padded ids can be derived (without it, everything after the first `<|endoftext|>` counts as padding). | -| `negative_input_ids` | Tokenized negative prompt, one row per prompt or a single row shared by the batch. Without it the unconditional branch is zero embeddings (`force_zeros_for_empty_prompt`), as in diffusers. | -| `num_inference_steps`, `guidance_scale`, `seed`, `latents` | As for [Stable Diffusion](stable_diffusion.md#zmdiffusion); defaults from the repo's `generate_args` (50 steps, guidance 5.0 for the base model). | -| `image`, `strength`, `denoising_start`, `denoising_end`, `output_type` | The image-to-image and ensemble controls of [`BaseDiffusion.generate`](main_classes.md#basediffusion): refine an `image` (noised to `strength`), stop early (`denoising_end`) and hand the latent over (`output_type="latent"`), or resume one (`denoising_start`). | -| `original_size`, `crops_coords_top_left`, `target_size` | SDXL's size / crop conditioning, `(height, width)` in pixels; the sizes default to the generated image's size, the crop to `(0, 0)`. | -| `aesthetic_score` | The refiner's fifth conditioning value (6.0), in place of `target_size`; ignored by the base models. | -| `negative_*` | The same for the negative branch; each defaults to its positive value, except `negative_aesthetic_score` (2.5). | - -Returns `(batch, H, W, 3)` uint8 images (or the latent); `latents` is -`(batch, 128, 128, 4)` at 1024px (`channels_first`: channels second). - -The hooks, for anyone composing their own loop: `encode_prompt(input_ids, attention_mask, -original_size, crops_coords_top_left, target_size)` returns the dict the UNet takes -(`encoder_hidden_states` (batch, 77, 2048), `text_embeds` (batch, 1280), `time_ids` -(batch, 6)); `encode_negative_prompt` zeroes the first two when no negative ids are given; -`predict_noise(latents, timesteps, embeddings)` runs the UNet. - -#### The refiner - -The same class with `model_type="stable_diffusion_xl_refiner"` (`text_encoder_2` alone, -`aesthetic_score` in the conditioning, `generate_args` with `strength` 0.3). Its -`generate` is the one above; with no `image` / `latents` it runs text-to-image, which the -refiner was not trained for. The two intended uses: - -```python -base = ZMDiffusion.from_weights("zeromodels/stable-diffusion-xl-base-1.0") -refiner = ZMDiffusion.from_weights("zeromodels/stable-diffusion-xl-refiner-1.0") -tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-xl-base-1.0") -inputs = tokenizer("a lighthouse on a rocky coast at dawn, dramatic clouds, cinematic") - -# ensemble of experts: the base denoises 80% of the schedule, the refiner the rest -latent = base.generate( - **inputs, num_inference_steps=50, denoising_end=0.8, output_type="latent" -) -images = refiner.generate( - **inputs, latents=latent, num_inference_steps=50, denoising_start=0.8 -) - -# image-to-image: refine any image (uint8 or [0, 1] float, (batch, H, W, 3)) -images = refiner.generate(**inputs, image=images, strength=0.3, seed=0) -``` - -#### The graph - -Four disconnected paths in one functional model, `.unet` / `.vae` / -`.text_encoder` / `.text_encoder_2`. - -| Inputs | Outputs | -|---|---| -| `sample` (B, 128, 128, 4), `timestep` (B,), `encoder_hidden_states` (B, 77, 2048), `text_embeds` (B, 1280), `time_ids` (B, 6) | `noise_pred` (B, 128, 128, 4) | -| `image` (B, 1024, 1024, 3), `latent` (B, 128, 128, 4) | `moments` (B, 128, 128, 8), `image` (B, 1024, 1024, 3) | -| `token_ids` (B, 77), `token_ids_2` (B, 77), `padding_mask` (B, 77) | `prompt_embeds` (B, 77, 2048), `pooled_prompt_embeds` (B, 1280) | - -Each text tower is a functional CLIP text model (the `clip` family's layers) whose outputs -are `penultimate_hidden_state`, `last_hidden_state`, `pooler_output` and, for the second, -`text_embeds`. The UNet is the SD 1.x `UNet2DConditionModel` with the SDXL config -(`transformer_layers_per_block=(1, 2, 10)`, `addition_embed_type="text_time"`); the VAE -is `AutoencoderKL` with `force_upcast=True`. The refiner is the same -graph without the first tower (`token_ids_2` + `padding_mask` in, `prompt_embeds` -(B, 77, 1280) out) and with `time_ids` (B, 5). - -## Preprocessing - -### `ZMDiffusionTokenizer` - -`ZMDiffusionTokenizer` with `model_type="stable_diffusion_xl"` (CLIP BPE, -`<|startoftext|>` / `<|endoftext|>` framing, -`<|endoftext|>`-padded to 77 tokens). Returns `{"input_ids", "attention_mask"}`; pass both -to `generate`. - -```python -ZMDiffusionTokenizer( - model_type="stable_diffusion_xl", hf_id=None, tokenizer_file=None, max_seq_len=77 -) # pad_token defaults to "<|endoftext|>"; the refiner repos use the same settings -``` - -A literal `!` in a prompt is encoded as OpenCLIP encodes it (the BPE token `!`, id -256) in the second tower's derived ids; diffusers' `tokenizer_2`, whose pad token is `!`, -turns it into the pad id 0 instead. - -## End-to-end example - -```python -import os - -os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" - -from PIL import Image -from zeromodels.diffusers import ZMDiffusion, ZMDiffusionTokenizer - -model = ZMDiffusion.from_weights("zeromodels/stable-diffusion-xl-base-1.0") -tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/stable-diffusion-xl-base-1.0") - -inputs = tokenizer( - "a lighthouse on a rocky coast at dawn, dramatic clouds, cinematic, highly detailed" -) -images = model.generate(**inputs, num_inference_steps=30, guidance_scale=5.0, seed=3) - -Image.fromarray(images[0]).save("lighthouse.png") # (1024, 1024, 3) uint8 -``` - -Stable Diffusion XL base 1.0: a lighthouse on a rocky coast at dawn, 1024px - -### Base + refiner - -The SDXL "ensemble of experts": the base runs the first 80% of the schedule and hands its -latent to the refiner, which finishes it. Same prompt and initial latent as above: - -```python -refiner = ZMDiffusion.from_weights("zeromodels/stable-diffusion-xl-refiner-1.0") -latent = model.generate( - **inputs, num_inference_steps=30, denoising_end=0.8, output_type="latent", seed=3 -) -images = refiner.generate( - **inputs, latents=latent, num_inference_steps=30, denoising_start=0.8 -) -``` - -Stable Diffusion XL base + refiner: the lighthouse, refined - -### SDXL-Turbo - -One step, no guidance, 512px (the repo's defaults): - -```python -model = ZMDiffusion.from_weights("zeromodels/sdxl-turbo") -tokenizer = ZMDiffusionTokenizer.from_weights("zeromodels/sdxl-turbo") -images = model.generate( - **tokenizer( - "a cinematic shot of a baby raccoon wearing an intricate italian priest robe" - ) -) -``` - -SDXL-Turbo, one step: a baby raccoon in an italian priest robe, 512px - -### Negative prompts, micro-conditioning - -```python -inputs = tokenizer(["a red bicycle", "a blue bicycle"]) -negative = tokenizer("blurry, low quality") # one row, shared by the batch -images = model.generate( - **inputs, - negative_input_ids=negative["input_ids"], - original_size=(2048, 2048), # "a crop of a larger image": sharper detail - crops_coords_top_left=(0, 0), - target_size=(1024, 1024), - seed=0, -) -``` - -### Other resolutions - -The graphs are built for the repo's resolution; pass overrides to build for another -(the weights are resolution-independent; multiples of 64px): - -```python -model = ZMDiffusion.from_weights( - "zeromodels/stable-diffusion-xl-base-1.0", - unet_sample_size=(96, 128), - vae_sample_size=(768, 1024), -) -images = model.generate(**tokenizer("a wide mountain valley")) # (1, 768, 1024, 3) -``` - -### float32 - -```python -model = ZMDiffusion.from_weights( - "zeromodels/stable-diffusion-xl-base-1.0", load_dtype="float32" -) -``` - -### Verified against diffusers - -The components of `stable-diffusion-xl-base-1.0` and `stable-diffusion-xl-refiner-1.0` -were checked one by one against diffusers 0.39 / transformers in float32 on the same -inputs (the converted float16 weights computed in float32): - -``` -base text_encoder penultimate max|d|=6.1e-05 (values up to 854) -base text_encoder_2 penultimate max|d|=3.1e-05 -base text_encoder_2 pooled max|d|=3.6e-06 -base unet noise_pred (t=981) max|d|=9.3e-06 -base vae decode (512px) max|d|=3.5e-05 -refiner text_encoder_2 penultimate max|d|=3.1e-05 -refiner unet noise_pred (t=181) max|d|=1.2e-05 -``` - -The pipelines (Euler `leading` spacing, classifier-free guidance with the zeroed or -encoded negative branch, the micro-conditioning values, negative prompts, batches, -`strength`, `denoising_end` / `denoising_start`) were checked against -`StableDiffusionXLPipeline` and `StableDiffusionXLImg2ImgPipeline` on small -random-weight models, where the generated images are pixel-identical on torch and jax. -At full size in float16 (the default load; diffusers in float16 with its float32 VAE), -the lighthouse above matches diffusers to PSNR 49.7 dB (98.5% of pixels within 2 uint8 -levels), the base + refiner ensemble to PSNR 51.9 dB (98.8% within 2) and SDXL-Turbo -(1 step, 512px) to PSNR 44.9 dB (92% within 2; a single step amplifies float16 rounding). -The lighthouse, zeromodels (left) and diffusers (right): - -zeromodels (left) and diffusers (right) generations of the lighthouse from stable-diffusion-xl-base-1.0 with the same latent - -## Data Format - -Both `channels_last` and `channels_first` are supported, as for -[Stable Diffusion](stable_diffusion.md#data-format); `generate` always returns -`(batch, H, W, 3)` uint8. - -## Memory and speed - -Measured on the torch backend with the default float16 load (weights resident: 6.8 GB for -a base model, 6.3 GB for the refiner, the VAE in float32): - -| Resolution | Denoising step (guided, batch 1) | VAE decode (float32) | -|---|---|---| -| 512px | peak 7.1 GB, 1.2 s/step on an RTX 4060 Laptop | peak 8.5 GB | -| 768px | peak 8.0 GB | peak 10.5 GB | -| 1024px | peak 10.4 GB | about 14 GB | - -Keras' functional executor keeps every layer output of a graph alive until the graph -returns, so the decoder's feature maps at 1024px add several GB on top of the weights: a -16 GB GPU runs 1024px comfortably, a 12 GB GPU 768px; on 8 GB, generate at 512px, or -denoise on the GPU (`output_type="latent"`) and decode the latent on the CPU (about 100 s -at 1024px) with a second, CPU-only process. In float32 the weights alone take 13.9 GB. -The UNet's attention runs the portable `"sdpa"` math by default, which materializes the -4096 x 4096 float32 logits of the 64 x 64 level at 1024px; loading with -`attn_implementation="fused"` (the backend's fused kernel, see -[`fused_attention`](main_classes.md#fused_attention)) cuts the 1024px step from a 10.3 GB -peak to 7.6 GB on the same GPU. On JAX and TensorFlow the denoiser step is compiled once -per run (see [`BaseDiffusion`](main_classes.md#basediffusion)). - -## Loading Fine-tuned Weights - -Any repo laid out like the hosted ones (`zm_config.json` declaring -`ZMDiffusion` with `model_type: "stable_diffusion_xl"` or -`"stable_diffusion_xl_refiner"`, the weights, `tokenizer.json`) loads with -`from_weights("/")`. The `hf:` prefix raises for diffusion models: convert a -diffusers-format SDXL checkpoint once with -`zeromodels/diffusers/stable_diffusion_xl/convert_stable_diffusion_xl_diffusers_to_keras.py` -(`transfer_stable_diffusion_xl(repo)`, `pip install zeromodels[conversion]`) and host the result. diff --git a/website/mkdocs.yml b/website/mkdocs.yml index 7dd90248..264a0cdb 100644 --- a/website/mkdocs.yml +++ b/website/mkdocs.yml @@ -169,18 +169,7 @@ nav: - RoBERTa: roberta.md - T5: t5.md - XLM-RoBERTa: xlm_roberta.md - - Sentence Transformer models: - - sentence_transformers.md - - Sentence Transformer: sentence_transformer.md - - Diffusion models: - - diffusers.md - - Stable Diffusion: stable_diffusion.md - - Stable Diffusion 2: stable_diffusion_2.md - - Stable Diffusion XL: stable_diffusion_xl.md - - Stable Diffusion 3: stable_diffusion_3.md - - Stable Diffusion 3.5: stable_diffusion_3_5.md - - Qwen-Image: qwen_image.md - - Qwen-Image-2.1: qwen_image_21.md + - Sentence Transformers: sentence_transformer.md - Vision models: - BEiT: beit.md - Classification backbones: classification_backbones.md @@ -246,6 +235,10 @@ nav: - SigLIP 2: siglip2.md - TIPSv2: tipsv2.md + # Diffusion models are their own entry: one page, since every family loads through one + # pipeline class (ZMDiffusion) and one tokenizer (ZMDiffusionTokenizer). + - Diffusion Models: diffusers.md + - Loading Weights: loading_weights.md # The AutoZM* loaders sit next to Loading Weights: they are the shortcut that reads a From f96fc07f7585843505967f72ab8dba650ca6d940 Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Sun, 18 Oct 2026 16:50:04 -0700 Subject: [PATCH 4/4] Update-workflow --- .github/TESTING.md | 80 ++++++----- .github/workflows/docs.yml | 6 +- .github/workflows/model_tests.yml | 178 ++++++++---------------- .github/workflows/run_tests_command.yml | 87 ++++++++++++ 4 files changed, 195 insertions(+), 156 deletions(-) create mode 100644 .github/workflows/run_tests_command.yml diff --git a/.github/TESTING.md b/.github/TESTING.md index a7e5b4f9..7b89fc50 100644 --- a/.github/TESTING.md +++ b/.github/TESTING.md @@ -1,41 +1,53 @@ -# Pull request test commands +# Pull request checks -The automatic PR workflow (`test_code.yml`) contains only **Lint & Format**. -Model, tokenizer, image-processor, and processor jobs live in a separate workflow -(`model_tests.yml`) that does not subscribe to pull request events. They do not -appear as queued or skipped PR checks before a command is posted. +## Automatic, on every pull request -Those tests run on a pull request only when **IMvision12** posts a new -conversation comment containing exactly: +Every pull request (from anyone) runs, with no command needed: + +- **Tests / Lint & Format** (`test_code.yml`): `ruff check` and `ruff format --check`. +- **Docs / build** (`docs.yml`): the strict MkDocs build and the asset check. Deploying + the site stays `main`-only. + +## Model, tokenizer and processor tests, on request + +The heavy suite (`model_tests.yml`: 9 model shards across Torch, TensorFlow and JAX, +1 tokenizer job, 3 image-processor jobs, 3 processor jobs, and one summary job) runs on +a pull request only when **IMvision12** posts a new comment containing exactly: ```text /run-tests ``` -Other users' commands, issue comments, edited comments, and other text do not -authorize the test suite. Re-running a comment-triggered test job also requires -IMvision12 to be the triggering account. - -The command starts the existing 16 test jobs: 9 model shards (Torch, TensorFlow, -and JAX), 1 tokenizer job, 3 image-processor jobs, and 3 processor jobs. They test -the open PR's head commit captured when the request is authorized. A later push -does not start another test run; post `/run-tests` again to test the new commit. - -Results appear on the tested PR commit as **Requested model, tokenizer & processor -tests**, with a link to the workflow's individual jobs. If these tests must block -merging, require this aggregate check in the repository's branch protection or -ruleset. The test check is created only after an authorized command. - -The updated workflow must be present on the repository's default branch before -GitHub delivers comment-triggered runs. No additional personal access token or -repository secret is needed. - -Release events start the same lint and test jobs automatically in -`model_tests.yml`, with the existing commands and matrices unchanged. The release, -documentation, and standalone cross-backend parity workflow files are unchanged. - -Authorization and result-reporting jobs have check-writing permission but never -check out PR code. Test jobs have read-only repository access and do not persist -checkout credentials. Comment-triggered tests do not use shared Actions caches -or publish Codecov coverage from the default-branch event context; release runs -retain their existing cache and coverage steps. +Other users' comments, edited comments and other text are ignored. + +How it works: `run_tests_command.yml` (triggered by the comment, never checking out PR +code) validates the request and dispatches `model_tests.yml` **on the pull request's +branch**, pinned to the PR head commit at that moment. A dispatched run is attached to +that commit, so every test job appears in the PR's checks list as +**Model tests / ...**, alongside the automatic checks. The comment gets a 🚀 reaction and +a reply with the tested commit. + +- A later push is not tested automatically: comment `/run-tests` again. A new request on + the same PR cancels the previous run. +- The branch must contain the current `model_tests.yml` (with its `workflow_dispatch` + trigger); otherwise the reply asks to merge or rebase `main`. +- **Pull requests from forks**: a fork's branch cannot be dispatched, so the tests run + from `main` on the PR head commit. They are not attached to the PR; the reply links + the run. +- `run_tests_command.yml` only takes effect once it is on `main` (comment events always + use the default branch's workflow files). + +**Model tests / Requested model, tokenizer & processor tests** passes only when every +test job passes. To make the suite block merging, require that check in the branch +protection rules or ruleset. + +## Releases + +Release events run lint and the whole suite in `model_tests.yml` automatically, with +pip caching and Codecov upload (both disabled for pull request runs). + +## Security + +The comment workflow has `actions: write` (to dispatch) and comment permissions, and never +checks out or executes PR code. The test jobs that run PR code have read-only repository +access, do not persist checkout credentials, and use no secrets or shared caches. diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index ed2821db..7cedbbdd 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -8,12 +8,8 @@ on: - "assets/**" - "website/**" - ".github/workflows/docs.yml" + # every pull request builds the site (deploying stays main-only, below) pull_request: - paths: - - "docs/**" - - "assets/**" - - "website/**" - - ".github/workflows/docs.yml" workflow_dispatch: permissions: diff --git a/.github/workflows/model_tests.yml b/.github/workflows/model_tests.yml index 38ae3748..2a928306 100644 --- a/.github/workflows/model_tests.yml +++ b/.github/workflows/model_tests.yml @@ -1,9 +1,33 @@ name: Model tests +run-name: >- + ${{ github.event_name == 'workflow_dispatch' + && format('Model tests for PR #{0}', inputs.pr_number) + || format('Model tests ({0})', github.event_name) }} + +# Runs on a release, and on a pull request when IMvision12 comments /run-tests: +# run_tests_command.yml dispatches this workflow on the PR branch, so the run is +# attached to the PR head commit and every job below is listed in the PR's checks. on: release: types: [created] - issue_comment: - types: [created] + workflow_dispatch: + inputs: + pr_number: + description: "Pull request number (set by /run-tests)" + required: false + default: "" + pr_sha: + description: "Pull request head commit to test (set by /run-tests)" + required: false + default: "" + +concurrency: + # a new /run-tests on the same PR replaces the previous run + group: model-tests-${{ inputs.pr_number || github.ref }} + cancel-in-progress: true + +permissions: + contents: read jobs: lint-and-format: @@ -19,69 +43,13 @@ jobs: - run: ruff check . - run: ruff format --check . - authorize-tests: - name: Authorize PR tests - if: >- - github.event_name == 'issue_comment' && - github.event.issue.pull_request && - github.event.comment.user.login == 'IMvision12' && - github.triggering_actor == 'IMvision12' && - github.event.comment.body == '/run-tests' - runs-on: ubuntu-latest - permissions: - contents: read - pull-requests: read - checks: write - outputs: - run_tests: ${{ steps.authorize.outputs.run_tests }} - head_sha: ${{ steps.authorize.outputs.head_sha }} - check_run_id: ${{ steps.authorize.outputs.check_run_id }} - steps: - # This privileged job never checks out or executes PR-controlled files. - - name: Validate owner command and pin the PR commit - id: authorize - uses: actions/github-script@v9 - with: - script: | - const { issue, comment, action } = context.payload; - if (action !== 'created' || !issue?.pull_request || - comment?.user?.login?.toLowerCase() !== 'imvision12' || - comment.body !== '/run-tests') { - throw new Error('Only IMvision12 can request PR tests with /run-tests.'); - } - const { data: pr } = await github.rest.pulls.get({ - ...context.repo, - pull_number: issue.number, - }); - if (pr.state !== 'open') { - throw new Error('Tests can only be requested on an open pull request.'); - } - if (!/^[0-9a-f]{40}$/.test(pr.head.sha)) { - throw new Error('GitHub returned an invalid PR head commit.'); - } - const { data: check } = await github.rest.checks.create({ - ...context.repo, - name: 'Requested model, tokenizer & processor tests', - head_sha: pr.head.sha, - status: 'in_progress', - details_url: `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}`, - output: { - title: 'PR tests requested by IMvision12', - summary: `Testing PR #${issue.number} at ${pr.head.sha}. New pushes require a new /run-tests comment.`, - }, - }); - core.setOutput('head_sha', pr.head.sha); - core.setOutput('check_run_id', String(check.id)); - core.setOutput('run_tests', 'true'); - test-code: name: Test the code (${{ matrix.backend }} ${{ matrix.shard }}) - needs: [lint-and-format, authorize-tests] + needs: [lint-and-format] if: >- always() && !cancelled() && - ((github.event_name == 'release' && needs.lint-and-format.result == 'success') || - (github.event_name == 'issue_comment' && github.triggering_actor == 'IMvision12' && - needs.authorize-tests.outputs.run_tests == 'true')) + (github.event_name == 'workflow_dispatch' || + (github.event_name == 'release' && needs.lint-and-format.result == 'success')) runs-on: ubuntu-latest permissions: contents: read @@ -107,13 +75,13 @@ jobs: steps: - uses: actions/checkout@v4 with: - ref: ${{ needs.authorize-tests.outputs.head_sha }} - persist-credentials: ${{ github.event_name != 'issue_comment' }} + ref: ${{ inputs.pr_sha }} + persist-credentials: false - uses: actions/setup-python@v5 with: python-version: '3.11' - name: Cache pip packages - if: github.event_name != 'issue_comment' + if: github.event_name == 'release' uses: actions/cache@v4 with: path: ~/.cache/pip @@ -147,7 +115,7 @@ jobs: - name: Generate coverage run: coverage xml -o coverage.xml - name: Upload coverage - if: github.event_name != 'issue_comment' + if: github.event_name == 'release' uses: codecov/codecov-action@v5 with: files: coverage.xml @@ -156,12 +124,11 @@ jobs: test-tokenizers: name: Test tokenizers (vs HF) - needs: [lint-and-format, authorize-tests] + needs: [lint-and-format] if: >- always() && !cancelled() && - ((github.event_name == 'release' && needs.lint-and-format.result == 'success') || - (github.event_name == 'issue_comment' && github.triggering_actor == 'IMvision12' && - needs.authorize-tests.outputs.run_tests == 'true')) + (github.event_name == 'workflow_dispatch' || + (github.event_name == 'release' && needs.lint-and-format.result == 'success')) runs-on: ubuntu-latest permissions: contents: read @@ -170,13 +137,13 @@ jobs: steps: - uses: actions/checkout@v4 with: - ref: ${{ needs.authorize-tests.outputs.head_sha }} - persist-credentials: ${{ github.event_name != 'issue_comment' }} + ref: ${{ inputs.pr_sha }} + persist-credentials: false - uses: actions/setup-python@v5 with: python-version: '3.11' - name: Cache pip packages - if: github.event_name != 'issue_comment' + if: github.event_name == 'release' uses: actions/cache@v4 with: path: ~/.cache/pip @@ -197,12 +164,11 @@ jobs: test-image-processors: name: Test image processors (vs HF, ${{ matrix.backend }}) - needs: [lint-and-format, authorize-tests] + needs: [lint-and-format] if: >- always() && !cancelled() && - ((github.event_name == 'release' && needs.lint-and-format.result == 'success') || - (github.event_name == 'issue_comment' && github.triggering_actor == 'IMvision12' && - needs.authorize-tests.outputs.run_tests == 'true')) + (github.event_name == 'workflow_dispatch' || + (github.event_name == 'release' && needs.lint-and-format.result == 'success')) runs-on: ubuntu-latest permissions: contents: read @@ -215,13 +181,13 @@ jobs: steps: - uses: actions/checkout@v4 with: - ref: ${{ needs.authorize-tests.outputs.head_sha }} - persist-credentials: ${{ github.event_name != 'issue_comment' }} + ref: ${{ inputs.pr_sha }} + persist-credentials: false - uses: actions/setup-python@v5 with: python-version: '3.11' - name: Cache pip packages - if: github.event_name != 'issue_comment' + if: github.event_name == 'release' uses: actions/cache@v4 with: path: ~/.cache/pip @@ -245,12 +211,11 @@ jobs: test-processors: name: Test processors (vs HF, ${{ matrix.backend }}) - needs: [lint-and-format, authorize-tests] + needs: [lint-and-format] if: >- always() && !cancelled() && - ((github.event_name == 'release' && needs.lint-and-format.result == 'success') || - (github.event_name == 'issue_comment' && github.triggering_actor == 'IMvision12' && - needs.authorize-tests.outputs.run_tests == 'true')) + (github.event_name == 'workflow_dispatch' || + (github.event_name == 'release' && needs.lint-and-format.result == 'success')) runs-on: ubuntu-latest permissions: contents: read @@ -263,13 +228,13 @@ jobs: steps: - uses: actions/checkout@v4 with: - ref: ${{ needs.authorize-tests.outputs.head_sha }} - persist-credentials: ${{ github.event_name != 'issue_comment' }} + ref: ${{ inputs.pr_sha }} + persist-credentials: false - uses: actions/setup-python@v5 with: python-version: '3.11' - name: Cache pip packages - if: github.event_name != 'issue_comment' + if: github.event_name == 'release' uses: actions/cache@v4 with: path: ~/.cache/pip @@ -291,37 +256,16 @@ jobs: pytest tests/processors/test_keras_processors_vs_hf.py \ -v --durations=20 - report-pr-tests: - name: Report requested PR tests - needs: [authorize-tests, test-code, test-tokenizers, test-image-processors, test-processors] - if: >- - always() && github.event_name == 'issue_comment' && - needs.authorize-tests.outputs.run_tests == 'true' + all-tests: + # one result for branch protection; same name as the former summary check + name: Requested model, tokenizer & processor tests + needs: [test-code, test-tokenizers, test-image-processors, test-processors] + if: always() && !cancelled() runs-on: ubuntu-latest - permissions: - checks: write steps: - # Keep the write token isolated from the jobs that execute PR code. - - name: Publish the result on the tested PR commit - uses: actions/github-script@v9 + - name: Check that every test job passed env: - TEST_RESULTS: ${{ toJSON(needs) }} - CHECK_RUN_ID: ${{ needs.authorize-tests.outputs.check_run_id }} - with: - script: | - const needs = JSON.parse(process.env.TEST_RESULTS); - const groups = ['test-code', 'test-tokenizers', 'test-image-processors', 'test-processors']; - const results = groups.map(name => needs[name].result); - const conclusion = results.every(result => result === 'success') - ? 'success' - : results.includes('cancelled') ? 'cancelled' : 'failure'; - await github.rest.checks.update({ - ...context.repo, - check_run_id: Number(process.env.CHECK_RUN_ID), - status: 'completed', - conclusion, - output: { - title: `Requested PR tests: ${conclusion}`, - summary: groups.map(name => `- ${name}: ${needs[name].result}`).join('\n'), - }, - }); + RESULTS: ${{ toJSON(needs) }} + run: | + echo "$RESULTS" | jq -r 'to_entries[] | "\(.key): \(.value.result)"' + echo "$RESULTS" | jq -e 'all(.[]; .result == "success")' > /dev/null diff --git a/.github/workflows/run_tests_command.yml b/.github/workflows/run_tests_command.yml new file mode 100644 index 00000000..29c4f22c --- /dev/null +++ b/.github/workflows/run_tests_command.yml @@ -0,0 +1,87 @@ +name: Run tests command + +# `/run-tests` from IMvision12 on a pull request starts the model, tokenizer and processor +# tests (model_tests.yml). A comment-triggered run is attached to the default branch, so +# its jobs would never appear in the PR's checks; instead this workflow dispatches +# model_tests.yml ON THE PR BRANCH, whose run is attached to the PR head commit and so +# lists every test job in the PR's checks. This job never checks out or runs PR code. +on: + issue_comment: + types: [created] + +permissions: {} + +jobs: + dispatch: + name: Start requested PR tests + if: >- + github.event.issue.pull_request && + github.event.comment.user.login == 'IMvision12' && + github.triggering_actor == 'IMvision12' && + github.event.comment.body == '/run-tests' + runs-on: ubuntu-latest + permissions: + actions: write + contents: read + issues: write + pull-requests: write + steps: + - name: Validate the command and dispatch model_tests.yml + uses: actions/github-script@v9 + with: + script: | + const { issue, comment, action } = context.payload; + if (action !== 'created' || !issue?.pull_request || + comment?.user?.login?.toLowerCase() !== 'imvision12' || + comment.body.trim() !== '/run-tests') { + throw new Error('Only IMvision12 can request PR tests with /run-tests.'); + } + const { data: pr } = await github.rest.pulls.get({ + ...context.repo, + pull_number: issue.number, + }); + if (pr.state !== 'open') { + throw new Error('Tests can only be requested on an open pull request.'); + } + if (!/^[0-9a-f]{40}$/.test(pr.head.sha)) { + throw new Error('GitHub returned an invalid PR head commit.'); + } + const repoName = `${context.repo.owner}/${context.repo.repo}`; + const sameRepo = pr.head.repo?.full_name === repoName; + // A branch of this repository is dispatched directly (its run is attached to + // the PR head and shows in the PR checks). A fork's branch cannot be + // dispatched, so its tests run from the default branch on the PR commit. + const ref = sameRepo ? pr.head.ref : pr.base.repo.default_branch; + const runsUrl = `${context.serverUrl}/${repoName}/actions/workflows/model_tests.yml`; + try { + await github.rest.actions.createWorkflowDispatch({ + ...context.repo, + workflow_id: 'model_tests.yml', + ref, + inputs: { pr_number: String(issue.number), pr_sha: pr.head.sha }, + }); + } catch (error) { + await github.rest.issues.createComment({ + ...context.repo, + issue_number: issue.number, + body: `Could not start the tests on \`${ref}\`: ${error.message}\n\n` + + 'The branch needs the current `.github/workflows/model_tests.yml` ' + + '(with its `workflow_dispatch` trigger); merge or rebase `main` and ' + + 'comment `/run-tests` again.', + }); + throw error; + } + await github.rest.reactions.createForIssueComment({ + ...context.repo, + comment_id: comment.id, + content: 'rocket', + }); + const where = sameRepo + ? 'Every test job is listed in this PR\'s checks.' + : `This PR comes from a fork, so the run is not attached to the PR; follow it at ${runsUrl}.`; + await github.rest.issues.createComment({ + ...context.repo, + issue_number: issue.number, + body: `Started the model, tokenizer and processor tests on ${pr.head.sha}. ${where}\n\n` + + 'A later push is not tested automatically: comment `/run-tests` again.', + });