diff --git a/README.md b/README.md index f897c3d1..d3f0ecc4 100644 --- a/README.md +++ b/README.md @@ -256,6 +256,7 @@ Documentation sources are also available in [`docs/`](docs/). | 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` |
@@ -282,7 +283,7 @@ Documentation sources are also available in [`docs/`](docs/). ## 📜 License -This project leverages [timm](https://github.com/huggingface/pytorch-image-models#licenses), [transformers](https://github.com/huggingface/transformers#license) and [diffusers](https://github.com/huggingface/diffusers#license) for converting pretrained weights from PyTorch to Keras. For licensing details, please refer to the respective repositories. Converted weights keep their upstream license (for example, the Stable Diffusion checkpoints are CreativeML OpenRAIL-M / OpenRAIL++-M, and SDXL-Turbo is non-commercial under the Stability AI Community License). +This project leverages [timm](https://github.com/huggingface/pytorch-image-models#licenses), [transformers](https://github.com/huggingface/transformers#license) and [diffusers](https://github.com/huggingface/diffusers#license) for converting pretrained weights from PyTorch to Keras. For licensing details, please refer to the respective repositories. Converted weights keep their upstream license (for example, the Stable Diffusion checkpoints are CreativeML OpenRAIL-M / OpenRAIL++-M, SDXL-Turbo is non-commercial under the Stability AI Community License, and Qwen-Image is Apache-2.0). - 🔖 **zeromodels Code**: This repository is licensed under the [Apache 2.0 License](https://www.apache.org/licenses/LICENSE-2.0). diff --git a/docs/getting_started.md b/docs/getting_started.md index cb4ac22d..044b99b0 100644 --- a/docs/getting_started.md +++ b/docs/getting_started.md @@ -51,10 +51,11 @@ results = processor.post_process_object_detection( - **Multimodal** - Vision-language generation and grounding. + 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) + [Kimi K2.5](kimi_k25.md) ┬╖ [LocateAnything](locateanything.md) ┬╖ + [Stable Diffusion](stable_diffusion.md) ┬╖ [Qwen-Image](qwen_image.md) - **Speech** diff --git a/docs/index.md b/docs/index.md index 54c7954e..c149f5d0 100644 --- a/docs/index.md +++ b/docs/index.md @@ -323,10 +323,11 @@ Encoders and decoder LLMs, dense and mixture-of-experts. ### Multimodal -Vision-language generation and grounding. +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) +[Kimi K2.5](kimi_k25.md) · [LocateAnything](locateanything.md) · +[Stable Diffusion](stable_diffusion.md) · [Qwen-Image](qwen_image.md) diff --git a/docs/loading_weights.md b/docs/loading_weights.md index 246d4e48..58cf6521 100644 --- a/docs/loading_weights.md +++ b/docs/loading_weights.md @@ -17,7 +17,7 @@ load, and how long it takes. | # | Way | What happens | Used by | |---|---|---|---| -| 1 | [**HuggingFace Hub**](#1-hub-keras-weights) | `zeromodels/`. The repo's `zm_config.json` rebuilds the model and `model.weights.h5` (or a sharded `.weights.json`) loads with no conversion. | Vision, detection, segmentation, depth, speech, text encoders, CLIP-family, classification backbones, diffusion (Stable Diffusion, hosted only: no way 2 / 3) | +| 1 | [**HuggingFace Hub**](#1-hub-keras-weights) | `zeromodels/`. The repo's `zm_config.json` rebuilds the model and `model.weights.h5` (or a sharded `.weights.json`) loads with no conversion. | Vision, detection, segmentation, depth, speech, text encoders, CLIP-family, classification backbones, diffusion (Stable Diffusion, Qwen-Image; hosted only: no way 2 / 3) | | 2 | [**On the fly**](#2-on-the-fly-conversion) | A bare variant whose entry carries an `hf_id`. Upstream safetensors are downloaded and converted in process. | The LLMs and VLMs: Qwen, Llama, Gemma, DeepSeek, GLM, Mistral, ... | | 3 | [**`hf:` prefix**](#3-the-hf-prefix) | Any Hub repo, named explicitly. Same conversion machinery as way 2, but you pick the repo. | Fine-tunes and community weights, for any architecture | diff --git a/docs/models.md b/docs/models.md index f5328f81..3b623407 100644 --- a/docs/models.md +++ b/docs/models.md @@ -110,3 +110,4 @@ Vision-language encoders, generative VLMs, grounding across detection, OCR, poin - [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) diff --git a/docs/qwen_image.md b/docs/qwen_image.md new file mode 100644 index 00000000..155e9d9e --- /dev/null +++ b/docs/qwen_image.md @@ -0,0 +1,356 @@ +# 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. The whole model is **one container**, +`QwenImageModel`, 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 +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), +[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 container: 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_`. + +### `QwenImageTextToImage` + +The text-to-image task: the `QwenImageModel` container plus `BaseDiffusion`'s +`generate`. It supplies the hooks 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", +) +``` + +| 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 space / empty prompt when unset | +| `negative_attention_mask` | `None` | mask for the negative ids | +| `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 | +|---|---|---| +| `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 | + +### `QwenImageModel` + +The container: one functional model whose 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. + +### `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 + +### `QwenImageTokenizer` + +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 +QwenImageTokenizer(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.models.qwen_image import ( + QwenImageTextToImage, + QwenImageTokenizer, +) + +model = QwenImageTextToImage.from_weights("zeromodels/qwen-image") +tokenizer = QwenImageTokenizer.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.models.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. + +### Container only + +```python +from zeromodels.models.qwen_image import QwenImageModel + +qi = QwenImageModel.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 = QwenImageTextToImage.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 container 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 `QwenImageModel`, 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` +(`transfer_qwen_image(repo)`, `pip install zeromodels[conversion]`) and host the +result. diff --git a/tests/base/model_test_registry.py b/tests/base/model_test_registry.py index 2eb0f352..c558b204 100644 --- a/tests/base/model_test_registry.py +++ b/tests/base/model_test_registry.py @@ -4791,6 +4791,61 @@ "expected_output_shape": dict(_sd3_outputs), } +# Qwen-Image: packed-latent double-stream DiT + Wan-derived VAE + Qwen2.5-VL text. +# sample_size=8 → packed_seq=(8/2)^2=16; VAE 32px → latent 4 with scale 8. +_qwen_image_tiny = { + "transformer_sample_size": 8, + "transformer_patch_size": 2, + "transformer_in_channels": 16, + "transformer_out_channels": 4, + "transformer_num_layers": 2, + "transformer_attention_head_dim": 8, + "transformer_num_attention_heads": 2, + "transformer_joint_attention_dim": 32, + "transformer_axes_dims_rope": (2, 2, 4), + "max_sequence_length": 16, + "vae_sample_size": 32, + "vae_base_dim": 16, + "vae_z_dim": 4, + "vae_dim_mult": (1, 1), + "vae_num_res_blocks": 1, + "vae_temperal_downsample": (False, True), + "vae_latents_mean": (0.0, 0.0, 0.0, 0.0), + "vae_latents_std": (1.0, 1.0, 1.0, 1.0), + "text_embed_dim": 32, + "text_mlp_dim": 64, + "text_num_layers": 2, + "text_num_heads": 4, + "text_num_kv_heads": 2, + "text_mrope_section": (1, 1, 2), + "max_seq_len": 32, + "vocab_size": 128, + "default_sample_size": 4, +} + +_qwen_image_outputs = { + "noise_pred": (2, 16, 16), + "moments": (2, 16, 16, 8), + "image": (2, 16, 16, 3), + "prompt_embeds": (2, 32, 32), +} +MODEL_TEST_CONFIGS["QwenImageModel"] = { + "module": "zeromodels.models.qwen_image", + "model_cls": "QwenImageModel", + "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), + "input_factory": "qwen_image_input", + "expected_output_shape": dict(_qwen_image_outputs), +} + def get_all_model_ids(): return list(MODEL_TEST_CONFIGS.keys()) diff --git a/tests/fixtures/cross_backend_parity.json b/tests/fixtures/cross_backend_parity.json index ead0d87f..ac35e568 100644 --- a/tests/fixtures/cross_backend_parity.json +++ b/tests/fixtures/cross_backend_parity.json @@ -20615,6 +20615,598 @@ ] } ], + "QwenImageModel": [ + { + "shape": [ + 2, + 16, + 16, + 3 + ], + "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 + ] + }, + { + "shape": [ + 2, + 16, + 16, + 8 + ], + "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 + ] + }, + { + "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 + ] + }, + { + "shape": [ + 2, + 32, + 32 + ], + "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 + ] + } + ], + "QwenImageTextToImage": [ + { + "shape": [ + 2, + 16, + 16, + 3 + ], + "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 + ] + }, + { + "shape": [ + 2, + 16, + 16, + 8 + ], + "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 + ] + }, + { + "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 + ] + }, + { + "shape": [ + 2, + 32, + 32 + ], + "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 + ] + } + ], "RFDETRDetect": [ { "shape": [ diff --git a/tests/fixtures/dummy_inputs.py b/tests/fixtures/dummy_inputs.py index 3d926da8..2467e0f9 100644 --- a/tests/fixtures/dummy_inputs.py +++ b/tests/fixtures/dummy_inputs.py @@ -236,6 +236,34 @@ def stable_diffusion_3_input( } +def qwen_image_input( + batch_size=2, + packed_seq=16, + in_channels=16, + image_size=32, + latent_size=8, + z_dim=4, + text_seq_len=16, + joint_attention_dim=32, + max_seq_len=32, +): + """Dummy inputs for the Qwen-Image container graph (packed latents).""" + return { + "sample": ops.ones((batch_size, packed_seq, in_channels)), + "timestep": ops.ones((batch_size,)), + "encoder_hidden_states": ops.ones( + (batch_size, text_seq_len, joint_attention_dim) + ), + "encoder_hidden_states_mask": ops.ones( + (batch_size, text_seq_len), dtype="int32" + ), + "image": ops.ones((batch_size, image_size, image_size, 3)), + "latent": ops.ones((batch_size, latent_size, latent_size, z_dim)), + "token_ids": ops.ones((batch_size, max_seq_len), dtype="int32"), + "padding_mask": ops.ones((batch_size, max_seq_len), dtype="int32"), + } + + def tips_v2_text_input(batch_size=2, max_seq_len=16): return { "token_ids": ops.ones((batch_size, max_seq_len), dtype="int32"), diff --git a/tests/integration/test_data_formats.py b/tests/integration/test_data_formats.py index 4ab49b48..765944c0 100644 --- a/tests/integration/test_data_formats.py +++ b/tests/integration/test_data_formats.py @@ -37,6 +37,9 @@ "Tipsv2DptDensePredict", "Tipsv2DptDepthEstimation", "Tipsv2DptSemanticSegment", + # Qwen-Image's Wan-derived VAE is built channels_last (NTHWC Conv3D). + "QwenImageModel", + "QwenImageTextToImage", # Qwen-VL inputs are pre-patchified (no spatial axes) -> layout-agnostic. "Qwen2VLModel", "Qwen2_5VLModel", diff --git a/website/mkdocs.yml b/website/mkdocs.yml index 9da072d7..0a838c93 100644 --- a/website/mkdocs.yml +++ b/website/mkdocs.yml @@ -237,6 +237,7 @@ nav: - 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 - TIPSv2: tipsv2.md - Loading Weights: loading_weights.md diff --git a/zeromodels/auto/auto_mapping_names.py b/zeromodels/auto/auto_mapping_names.py index 389be02e..92fe9034 100644 --- a/zeromodels/auto/auto_mapping_names.py +++ b/zeromodels/auto/auto_mapping_names.py @@ -281,6 +281,7 @@ "stable_diffusion_xl_refiner": "StableDiffusionXLRefinerModel", "stable_diffusion_3": "StableDiffusion3Model", "stable_diffusion_3_5": "StableDiffusion3_5Model", + "qwen_image": "QwenImageModel", "swin": "SwinModel", "swinv2": "SwinV2Model", "t5": "T5Model", @@ -418,6 +419,7 @@ "stable_diffusion_xl": "StableDiffusionXLTextToImage", "stable_diffusion_3": "StableDiffusion3TextToImage", "stable_diffusion_3_5": "StableDiffusion3_5TextToImage", + "qwen_image": "QwenImageTextToImage", }, "ImageToImage": { "stable_diffusion_xl_refiner": "StableDiffusionXLRefinerImageToImage", @@ -686,6 +688,7 @@ "stable_diffusion_3": "StableDiffusion3Config", "stable_diffusion_3_5": "StableDiffusion3_5Config", "stable_diffusion_3_t5_encoder": "StableDiffusion3T5EncoderConfig", + "qwen_image": "QwenImageConfig", "swin": "SwinConfig", "swinv2": "SwinV2Config", "t5": "T5Config", @@ -788,6 +791,7 @@ "stable_diffusion_xl_refiner": "StableDiffusionXLTokenizer", "stable_diffusion_3": "StableDiffusion3Tokenizer", "stable_diffusion_3_5": "StableDiffusion3_5Tokenizer", + "qwen_image": "QwenImageTokenizer", "t5": "T5Tokenizer", "tipsv2": "Tipsv2Tokenizer", "whisper": "WhisperTokenizer", diff --git a/zeromodels/base/base_scheduler.py b/zeromodels/base/base_scheduler.py index 9bebb6c4..42842c74 100644 --- a/zeromodels/base/base_scheduler.py +++ b/zeromodels/base/base_scheduler.py @@ -487,7 +487,7 @@ def step(self, model_output, timestep, sample, **kwargs): class FlowMatchEulerDiscreteScheduler(BaseScheduler): - """Euler sampler for rectified-flow models (Stable Diffusion 3 / 3.5, FLUX). + """Euler sampler for rectified-flow models (SD 3 / 3.5, FLUX, Qwen-Image). There is no beta schedule: the noise level is the flow time ``sigma`` in ``[0, 1]`` (``x_t = (1 - sigma) x_0 + sigma * noise``), the model predicts the @@ -496,20 +496,55 @@ class FlowMatchEulerDiscreteScheduler(BaseScheduler): Args: num_train_timesteps: The flow time resolution (1000). - shift: Timestep shift towards noisier levels, ``shift * s / (1 + (shift - 1) s)`` - (3.0 for SD3 / SD3.5). + shift: Fixed timestep shift ``shift * s / (1 + (shift - 1) s)`` (3.0 for + SD3). Ignored when ``use_dynamic_shifting`` is True. + use_dynamic_shifting: Resolution-dependent shift (Qwen-Image / FLUX); + ``set_timesteps(..., mu=...)`` required. + base_shift / max_shift / base_image_seq_len / max_image_seq_len: + Dynamic-shift parameters (stored for callers that compute ``mu``). + shift_terminal: Stretch the schedule so the last sigma equals this value. + time_shift_type: ``"exponential"`` or ``"linear"`` dynamic shift. """ - def __init__(self, num_train_timesteps=1000, shift=1.0, **kwargs): + def __init__( + self, + num_train_timesteps=1000, + shift=1.0, + use_dynamic_shifting=False, + base_shift=0.5, + max_shift=1.15, + base_image_seq_len=256, + max_image_seq_len=4096, + shift_terminal=None, + time_shift_type="exponential", + **kwargs, + ): # the base (beta) schedule is irrelevant here; keep the constructor # compatible with from_config (a repo's scheduler_config may carry extra keys) super().__init__(num_train_timesteps=num_train_timesteps) self.shift = shift + self.use_dynamic_shifting = bool(use_dynamic_shifting) + self.base_shift = base_shift + self.max_shift = max_shift + self.base_image_seq_len = base_image_seq_len + self.max_image_seq_len = max_image_seq_len + self.shift_terminal = shift_terminal + self.time_shift_type = time_shift_type + self.config_dict = { + "base_shift": base_shift, + "max_shift": max_shift, + "base_image_seq_len": base_image_seq_len, + "max_image_seq_len": max_image_seq_len, + "use_dynamic_shifting": self.use_dynamic_shifting, + "shift_terminal": shift_terminal, + "time_shift_type": time_shift_type, + } timesteps = np.linspace( 1, num_train_timesteps, num_train_timesteps, dtype=np.float32 )[::-1].copy() sigmas = timesteps / np.float32(num_train_timesteps) - sigmas = self.shift_sigmas(sigmas) + if not self.use_dynamic_shifting: + sigmas = self.shift_sigmas(sigmas) self.sigma_min = float(sigmas[-1]) self.sigma_max = float(sigmas[0]) self.sigmas = sigmas @@ -520,6 +555,26 @@ def shift_sigmas(self, sigmas): shift = np.float32(self.shift) return (shift * sigmas / (1 + (shift - 1) * sigmas)).astype(np.float32) + def time_shift(self, mu, sigma, t): + """Resolution-dependent sigma shift (Diffusers ``time_shift``).""" + t = np.asarray(t, dtype=np.float64) + if self.time_shift_type == "exponential": + return (np.exp(mu) / (np.exp(mu) + (1.0 / t - 1.0) ** sigma)).astype( + np.float32 + ) + if self.time_shift_type == "linear": + return (mu / (mu + (1.0 / t - 1.0) ** sigma)).astype(np.float32) + raise ValueError(f"Unknown time_shift_type {self.time_shift_type!r}") + + def stretch_shift_to_terminal(self, sigmas): + """Stretch sigmas so the last (pre-zero) value equals ``shift_terminal``.""" + if self.shift_terminal is None: + return sigmas + sigmas = np.asarray(sigmas, dtype=np.float64) + one_minus = 1.0 - sigmas + scale = one_minus[-1] / (1.0 - float(self.shift_terminal)) + return (1.0 - one_minus / scale).astype(np.float32) + @property def init_noise_sigma(self): return 1.0 @@ -527,15 +582,34 @@ def init_noise_sigma(self): def set_begin_index(self, index): self.step_index = int(index) - def set_timesteps(self, num_inference_steps): + def set_timesteps(self, num_inference_steps, sigmas=None, mu=None, **kwargs): + del kwargs self.num_inference_steps = num_inference_steps n_train = self.num_train_timesteps - timesteps = np.linspace( - self.sigma_max * n_train, self.sigma_min * n_train, num_inference_steps - ) - sigmas = timesteps / n_train - sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) - sigmas = sigmas.astype(np.float32) + if sigmas is None: + timesteps = np.linspace( + self.sigma_max * n_train, self.sigma_min * n_train, num_inference_steps + ) + sigmas = timesteps / n_train + else: + sigmas = np.asarray(sigmas, dtype=np.float32) + num_inference_steps = len(sigmas) + self.num_inference_steps = num_inference_steps + + if self.use_dynamic_shifting: + if mu is None: + raise ValueError( + "`mu` must be passed when use_dynamic_shifting is True" + ) + sigmas = self.time_shift(mu, 1.0, sigmas) + else: + sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) + sigmas = sigmas.astype(np.float32) + + if self.shift_terminal is not None: + sigmas = self.stretch_shift_to_terminal(sigmas) + + sigmas = np.asarray(sigmas, dtype=np.float32) self.sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) self.timesteps = sigmas * np.float32(n_train) self.step_index = 0 @@ -563,6 +637,13 @@ def to_config(self): "_class_name": type(self).__name__, "num_train_timesteps": self.num_train_timesteps, "shift": self.shift, + "use_dynamic_shifting": self.use_dynamic_shifting, + "base_shift": self.base_shift, + "max_shift": self.max_shift, + "base_image_seq_len": self.base_image_seq_len, + "max_image_seq_len": self.max_image_seq_len, + "shift_terminal": self.shift_terminal, + "time_shift_type": self.time_shift_type, } diff --git a/zeromodels/models/__init__.py b/zeromodels/models/__init__.py index 463e0fac..fb9e924e 100644 --- a/zeromodels/models/__init__.py +++ b/zeromodels/models/__init__.py @@ -104,6 +104,7 @@ qwen3_next, qwen3_vl, qwen3_vl_moe, + qwen_image, regnet, res2net, resmlp, diff --git a/zeromodels/models/qwen_image/__init__.py b/zeromodels/models/qwen_image/__init__.py new file mode 100644 index 00000000..d7fe9377 --- /dev/null +++ b/zeromodels/models/qwen_image/__init__.py @@ -0,0 +1,27 @@ +from .qwen_image_config import ( + QwenImageConfig, + QwenImageTextConfig, + QwenImageTransformerConfig, + QwenImageVAEConfig, +) +from .qwen_image_model import ( + AutoencoderKLQwenImage, + QwenImageModel, + QwenImageTextEncoderModel, + QwenImageTextToImage, + QwenImageTransformer2DModel, +) +from .qwen_image_tokenizer import QwenImageTokenizer + +__all__ = [ + "AutoencoderKLQwenImage", + "QwenImageConfig", + "QwenImageModel", + "QwenImageTextConfig", + "QwenImageTextEncoderModel", + "QwenImageTextToImage", + "QwenImageTokenizer", + "QwenImageTransformer2DModel", + "QwenImageTransformerConfig", + "QwenImageVAEConfig", +] diff --git a/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py b/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py new file mode 100644 index 00000000..07f52b09 --- /dev/null +++ b/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py @@ -0,0 +1,345 @@ +from typing import Dict + +import numpy as np +from tqdm import tqdm + +from zeromodels.conversion.exceptions import ( + WeightMappingError, + WeightShapeMismatchError, +) +from zeromodels.conversion.weight_split_util import split_model_weights +from zeromodels.conversion.weight_transfer_util import ( + compare_keras_torch_names, + transfer_weights, + zeros_init, +) + +QWEN_IMAGE_SOURCES = { + "qwen-image": "Qwen/Qwen-Image", +} + +WEIGHT_NAME_MAPPING: Dict[str, str] = { + "__": ".", + "/kernel": ".weight", + "/gamma": ".weight", + "/beta": ".bias", + "/scale": ".weight", + "/": ".", + # text encoder (Qwen2.5-VL language tower) + "token_embedding.embeddings": "model.embed_tokens.weight", + "language_model.final_norm.weight": "model.norm.weight", + "language_model.": "model.", + "decoder_layer_": "layers.", + "attention.query": "self_attn.q_proj", + "attention.key": "self_attn.k_proj", + "attention.value": "self_attn.v_proj", + "attention.output_proj": "self_attn.o_proj", + "attention_norm": "input_layernorm", + "mlp_norm": "post_attention_layernorm", + "mlp.gate": "mlp.gate_proj", + "mlp.up": "mlp.up_proj", + "mlp.down": "mlp.down_proj", + "gamma": "weight", + "beta": "bias", + "kernel": "weight", +} + + +def config_from_diffusers(repo, token=None): + import json + + from huggingface_hub import hf_hub_download + from diffusers import FlowMatchEulerDiscreteScheduler + from diffusers import QwenImageTransformer2DModel as DiffusersTransformer + from transformers import AutoConfig + + from zeromodels.models.qwen_image.qwen_image_config import ( + QwenImageConfig, + QwenImageVAEConfig, + ) + from zeromodels.models.qwen_image.qwen_image_model import ( + QwenImageTransformer2DModel, + ) + + transformer = dict( + DiffusersTransformer.load_config(repo, subfolder="transformer", token=token) + ) + vae = json.load( + open( + hf_hub_download(repo, "config.json", subfolder="vae", token=token), + encoding="utf-8", + ) + ) + text = AutoConfig.from_pretrained( + repo, subfolder="text_encoder", token=token + ).to_dict() + text_inner = text.get("text_config") or text + scheduler = { + k: v + for k, v in FlowMatchEulerDiscreteScheduler.load_config( + repo, subfolder="scheduler", token=token + ).items() + if k == "_class_name" or not k.startswith("_") + } + temperal = tuple(vae.get("temperal_downsample", (False, True, True))) + return QwenImageConfig( + transformer_config={ + **QwenImageTransformer2DModel.kwargs_from_diffusers_config(transformer), + "sample_size": 128, + "text_seq_len": 512, + }, + vae_config=QwenImageVAEConfig( + base_dim=vae.get("base_dim", 96), + z_dim=vae.get("z_dim", 16), + dim_mult=tuple(vae.get("dim_mult", (1, 2, 4, 4))), + num_res_blocks=vae.get("num_res_blocks", 2), + attn_scales=tuple(vae.get("attn_scales") or ()), + temperal_downsample=temperal, + dropout=vae.get("dropout", 0.0), + latents_mean=tuple(vae["latents_mean"]), + latents_std=tuple(vae["latents_std"]), + sample_size=1024, + ), + text_config={ + "vocab_size": text_inner.get("vocab_size", 152064), + "embed_dim": text_inner.get("hidden_size", 3584), + "mlp_dim": text_inner.get("intermediate_size", 18944), + "num_layers": text_inner.get("num_hidden_layers", 28), + "num_heads": text_inner.get("num_attention_heads", 28), + "num_kv_heads": text_inner.get("num_key_value_heads", 4), + "norm_eps": text_inner.get("rms_norm_eps", 1e-6), + "rope_theta": text_inner.get("rope_theta", 1000000.0), + "mrope_section": tuple( + (text_inner.get("rope_scaling") or {}).get( + "mrope_section", (16, 24, 24) + ) + ), + "tie_embeddings": text.get("tie_word_embeddings", False), + "max_seq_len": 1024, + }, + scheduler_config=scheduler, + bos_token_id=text.get("bos_token_id", 151643), + eos_token_id=text.get("eos_token_id", 151645), + pad_token_id=text.get("bos_token_id", 151643), + ) + + +def transfer_qwen_image( + repo, token=None, dtype="float16", build_sample_size=16, config=None +): + import gc + import json + + from huggingface_hub import hf_hub_download + from safetensors import safe_open + + from zeromodels.base.base_mixin import build_dtype_scope + from zeromodels.models.qwen_image.qwen_image_model import QwenImageModel + + 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) + with build_dtype_scope(dtype), zeros_init(): + model = QwenImageModel(**flat) + + # VAE RMSNorm uses /scale; drop /gamma so it cannot collide with other paths + vae_mapping = {k: v for k, v in WEIGHT_NAME_MAPPING.items() if "/gamma" not in k} + for step, (component, subfolder, mapping, index_name, filename) in enumerate( + ( + ( + model.transformer, + "transformer", + WEIGHT_NAME_MAPPING, + "diffusion_pytorch_model.safetensors.index.json", + None, + ), + ( + model.vae, + "vae", + vae_mapping, + None, + "diffusion_pytorch_model.safetensors", + ), + ), + start=2, + ): + print(f"[{step}/4] Transferring {subfolder}…", flush=True) + if index_name is not None: + index_path = hf_hub_download( + repo, index_name, subfolder=subfolder, token=token + ) + with open(index_path, encoding="utf-8") as f: + weight_map = json.load(f)["weight_map"] + shard_paths = { + shard: hf_hub_download(repo, shard, subfolder=subfolder, token=token) + for shard in sorted(set(weight_map.values())) + } + + class _State(dict): + def __contains__(self, key): + return key in weight_map + + def __getitem__(self, key): + with safe_open( + shard_paths[weight_map[key]], framework="np" + ) as shard: + return shard.get_tensor(key) + + def keys(self): + return weight_map.keys() + + def __iter__(self): + return iter(weight_map) + + def __len__(self): + return len(weight_map) + + state = _State() + else: + path = hf_hub_download(repo, filename, subfolder=subfolder, token=token) + state = {} + with safe_open(path, framework="np") as shard: + for key in shard.keys(): + state[key] = shard.get_tensor(key) + + consumed = set() + trainable, non_trainable = split_model_weights(component) + for keras_weight, _ in tqdm( + trainable + non_trainable, + desc=f"Transferring {subfolder} weights to Keras", + ): + if any(part in keras_weight.path for part in ("pos_freqs", "neg_freqs")): + continue + key = "/".join(keras_weight.path.split("/")[-2:]) + for old, new in mapping.items(): + key = key.replace(old, new) + if key not in state: + raise WeightMappingError(keras_weight.path, key) + consumed.add(key) + raw = state[key] + arr = np.asarray(raw) + kshape = tuple(keras_weight.shape) + if len(kshape) == 5 and arr.ndim == 5: + arr = np.transpose(arr, (2, 3, 4, 1, 0)) + elif len(kshape) == 4 and arr.ndim == 4: + arr = np.transpose(arr, (2, 3, 1, 0)) + elif ( + arr.ndim > 1 + and len(kshape) == 1 + and int(np.prod(arr.shape)) == kshape[0] + ): + arr = arr.reshape(kshape) + if len(keras_weight.shape) in (4, 5): + if tuple(keras_weight.shape) != arr.shape: + raise WeightShapeMismatchError( + keras_weight.path, keras_weight.shape, key, np.shape(raw) + ) + keras_weight.assign(arr) + continue + if tuple(keras_weight.shape) != tuple(arr.shape): + if len(keras_weight.shape) == 1 and np.asarray(raw).ndim > 1: + raise WeightShapeMismatchError( + keras_weight.path, keras_weight.shape, key, arr.shape + ) + if not compare_keras_torch_names( + keras_weight.path, keras_weight, key, raw + ): + raise WeightShapeMismatchError( + keras_weight.path, + keras_weight.shape, + key, + np.asarray(raw).shape, + ) + transfer_weights(keras_weight.path, keras_weight, arr) + unused = sorted(set(state) - consumed) + if unused: + raise ValueError( + f"{type(component).__name__}: {len(unused)} checkpoint tensors " + f"unused, e.g. {unused[:5]}." + ) + del state + gc.collect() + + print("[4/4] Transferring text encoder…", flush=True) + index_path = hf_hub_download( + repo, "model.safetensors.index.json", subfolder="text_encoder", token=token + ) + with open(index_path, encoding="utf-8") as f: + weight_map = json.load(f)["weight_map"] + shard_paths = { + shard: hf_hub_download(repo, shard, subfolder="text_encoder", token=token) + for shard in sorted(set(weight_map.values())) + } + hf_keys = {} + for key in weight_map: + if key.startswith("model.language_model."): + hf_keys["model." + key[len("model.language_model.") :]] = key + elif key.startswith("model.") and not key.startswith("model.visual."): + hf_keys[key] = key + consumed = set() + text_encoder = model.text_encoder + for weight in tqdm( + text_encoder.weights, desc="Transferring text_encoder weights to Keras" + ): + name = weight.path.removeprefix(f"{text_encoder.name}/") + for old, new in WEIGHT_NAME_MAPPING.items(): + name = name.replace(old, new) + if name not in hf_keys: + raise WeightMappingError(weight.path, name) + consumed.add(name) + shard_key = hf_keys[name] + with safe_open(shard_paths[weight_map[shard_key]], framework="np") as shard: + transfer_weights(weight.path, weight, shard.get_tensor(shard_key)) + unused = sorted(set(hf_keys) - consumed) + if unused: + raise ValueError( + f"text_encoder: {len(unused)} checkpoint tensors unused, e.g. {unused[:5]}." + ) + del weight_map, shard_paths, hf_keys + gc.collect() + return model, config + + +if __name__ == "__main__": + import gc + import os + + import keras + + OUT_DIR = os.environ.get( + "ZM_OUT_DIR", "C:/Users/gites/Desktop/code/qwen_image_weights" + ) + os.makedirs(OUT_DIR, exist_ok=True) + MAX_SHARD_GB = 5.0 + token = os.environ.get("HF_TOKEN") + dtype = os.environ.get("ZM_DTYPE", "bfloat16") + device = os.environ.get("ZM_DEVICE", "cpu") + selected = [v for v in os.environ.get("ZM_VARIANTS", "").split(",") if v] + sources = { + variant: source + for variant, source in QWEN_IMAGE_SOURCES.items() + if not selected or variant in selected + } + + for variant, source in sources.items(): + print(f"\n{'=' * 60}\nConverting: {variant} <- {source}\n{'=' * 60}") + with keras.device(device): + model, config = transfer_qwen_image(source, token=token, dtype=dtype) + + itemsize = 2 if "16" in dtype else 4 + n_bytes = sum(int(np.prod(w.shape)) * itemsize for w in model.weights) + stem = os.path.join(OUT_DIR, variant.replace("-", "_")) + if n_bytes > MAX_SHARD_GB * 1024**3: + out = f"{stem}.weights.json" + model.save_weights(out, max_shard_size=MAX_SHARD_GB) + else: + out = f"{stem}.weights.h5" + model.save_weights(out) + print(f" saved -> {out} ({n_bytes / 1024**3:.2f} GB {dtype})") + + del model + keras.backend.clear_session() + gc.collect() diff --git a/zeromodels/models/qwen_image/qwen_image_config.py b/zeromodels/models/qwen_image/qwen_image_config.py new file mode 100644 index 00000000..9adf5cb0 --- /dev/null +++ b/zeromodels/models/qwen_image/qwen_image_config.py @@ -0,0 +1,165 @@ +from zeromodels.base import BaseConfig +from zeromodels.models.qwen2_5_vl.qwen2_5_vl_config import Qwen2_5VLTextConfig + +DEFAULT_LATENTS_MEAN = ( + -0.7571, + -0.7089, + -0.9113, + 0.1075, + -0.1745, + 0.9653, + -0.1517, + 1.5508, + 0.4134, + -0.0715, + 0.5517, + -0.3632, + -0.1922, + -0.9497, + 0.2503, + -0.2921, +) +DEFAULT_LATENTS_STD = ( + 2.8184, + 1.4541, + 2.3275, + 2.6558, + 1.2196, + 1.7708, + 2.6052, + 2.0743, + 3.2687, + 2.1526, + 2.8652, + 1.5579, + 1.6382, + 1.1253, + 2.8251, + 1.9160, +) + + +class QwenImageVAEConfig(BaseConfig): + """Configuration for :class:`AutoencoderKLQwenImage`. + + Fields match Diffusers ``AutoencoderKLQwenImage``; ``sample_size`` is the + ZeroModels graph-build resolution (weights are resolution-independent). + """ + + model_type = "autoencoder_kl_qwen_image" + + base_dim: int = 96 + z_dim: int = 16 + dim_mult: tuple = (1, 2, 4, 4) + num_res_blocks: int = 2 + attn_scales: tuple = () + temperal_downsample: tuple = (False, True, True) + dropout: float = 0.0 + input_channels: int = 3 + latents_mean: tuple = DEFAULT_LATENTS_MEAN + latents_std: tuple = DEFAULT_LATENTS_STD + sample_size: int = 1024 + + +class QwenImageTransformerConfig(BaseConfig): + r"""Configuration for [`QwenImageTransformer2DModel`]. + + Defaults match ``Qwen/Qwen-Image`` ``transformer/config.json`` (60-layer + double-stream DiT, 24 heads × 128, joint text width 3584). + + Args: + patch_size: Latent patch side (2); packed tokens use ``in_channels``. + in_channels: Packed latent width (64 = 16 × 2 × 2). + out_channels: Unpacked latent channels (16). + num_layers: Dual-stream DiT blocks (60). + attention_head_dim / num_attention_heads: Head geometry (inner dim = + heads × head_dim = 3072). + joint_attention_dim: Text feature width from Qwen2.5-VL (3584). + axes_dims_rope: MS-RoPE axis splits ``(16, 56, 56)``. + sample_size: Latent spatial side the graph is built for (image / 8). + text_seq_len: Static text sequence length for the graph. + """ + + model_type = "qwen_image_transformer_2d" + + patch_size: int = 2 + in_channels: int = 64 + out_channels: int = 16 + num_layers: int = 60 + attention_head_dim: int = 128 + num_attention_heads: int = 24 + joint_attention_dim: int = 3584 + axes_dims_rope: tuple = (16, 56, 56) + guidance_embeds: bool = False + sample_size: int = 128 + text_seq_len: int = 512 + + +class QwenImageTextConfig(Qwen2_5VLTextConfig): + r"""Qwen2.5-VL text tower used as the Qwen-Image prompt encoder. + + Defaults match ``Qwen/Qwen-Image`` ``text_encoder/config.json`` (7B Instruct + text half: 28 layers, 3584-d, 28 heads / 4 KV). + """ + + model_type = "qwen_image_text" + + vocab_size: int = 152064 + embed_dim: int = 3584 + mlp_dim: int = 18944 + num_layers: int = 28 + num_heads: int = 28 + num_kv_heads: int = 4 + norm_eps: float = 1e-06 + rope_theta: float = 1000000.0 + mrope_section: tuple = (16, 24, 24) + tie_embeddings: bool = False + max_seq_len: int = 1024 + + +class QwenImageConfig(BaseConfig): + r"""Configuration for [`QwenImageModel`] / [`QwenImageTextToImage`]. + + One hosted container: 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. + + Args: + transformer_config: The double-stream DiT. + vae_config: The Wan-derived 16-channel VAE. + text_config: The Qwen2.5-VL text tower. + scheduler_config: Diffusers ``scheduler_config.json`` dict + (``FlowMatchEulerDiscreteScheduler`` with dynamic shifting). + prompt_template_encode_start_idx: Tokens dropped from the ChatML + template prefix when building prompt embeds (34 in Diffusers). + max_sequence_length: Prompt embed length after the template drop (512). + default_sample_size: Latent grid side used when height/width omitted + (128 → 1024px at VAE scale 8). + bos_token_id / eos_token_id / pad_token_id: Qwen2 special token ids. + """ + + model_type = "qwen_image" + + sub_configs = { + "transformer_config": QwenImageTransformerConfig, + "vae_config": QwenImageVAEConfig, + "text_config": QwenImageTextConfig, + } + sub_config_prefixes = { + "transformer_config": "transformer_", + "vae_config": "vae_", + "text_config": "text_", + } + group_extras = {"text_config": ("vocab_size", "max_seq_len")} + + transformer_config: QwenImageTransformerConfig | dict | None = None + vae_config: QwenImageVAEConfig | dict | None = None + text_config: QwenImageTextConfig | dict | None = None + scheduler_config: dict | None = None + prompt_template_encode_start_idx: int = 34 + max_sequence_length: int = 512 + default_sample_size: int = 128 + bos_token_id: int = 151643 + eos_token_id: int = 151645 + pad_token_id: int = 151643 diff --git a/zeromodels/models/qwen_image/qwen_image_layers.py b/zeromodels/models/qwen_image/qwen_image_layers.py new file mode 100644 index 00000000..c3e98b0a --- /dev/null +++ b/zeromodels/models/qwen_image/qwen_image_layers.py @@ -0,0 +1,708 @@ +from __future__ import annotations + +import keras +import numpy as np +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, + timestep_embedding, +) +from zeromodels.models.stable_diffusion_3.stable_diffusion_3_layers import ( + StableDiffusion3AdaLayerNorm, + StableDiffusion3GELUFeedForward, +) + +NORM_EPS = 1e-6 +MASK_NEG = -1e4 +ROPE_MAX_INDEX = 4096 + + +def apply_rotary_emb_qwen(x, freqs_cis, use_real=True, use_real_unbind_dim=-1): + """Apply Qwen RoPE to ``x`` of shape ``[B, S, H, D]``. + + ``freqs_cis`` is either rotation angles ``[S, D // 2]`` (preferred; matches + the Diffusers Neuron path) or a ``(cos, sin)`` pair each ``[S, D]`` with + adjacent-pair angles repeated. Equivalent to complex multiplication of + adjacent feature pairs ``(x[..., 2k], x[..., 2k+1])``. + + Args: + x: Query or key tensor ``[B, S, H, D]``. + freqs_cis: Angles ``[S, D // 2]`` or ``(cos, sin)`` each ``[S, D]``. + use_real: Kept for Diffusers API parity; only the real path is used. + use_real_unbind_dim: ``-1`` (adjacent pairs, Qwen/Flux) or ``-2``. + """ + del use_real # real path only (portable across Keras backends) + if isinstance(freqs_cis, (tuple, list)): + cos, sin = freqs_cis + else: + cos = ops.repeat(ops.cos(freqs_cis), 2, axis=-1) + sin = ops.repeat(ops.sin(freqs_cis), 2, axis=-1) + cos = ops.expand_dims(ops.expand_dims(cos, 0), 2) + sin = ops.expand_dims(ops.expand_dims(sin, 0), 2) + + x_dtype = x.dtype + x_f = ops.cast(x, "float32") + if use_real_unbind_dim == -1: + pair = ops.reshape(x_f, ops.shape(x)[:-1] + (-1, 2)) + x_real, x_imag = pair[..., 0], pair[..., 1] + x_rotated = ops.reshape(ops.stack([-x_imag, x_real], axis=-1), ops.shape(x)) + elif use_real_unbind_dim == -2: + pair = ops.reshape(x_f, ops.shape(x)[:-1] + (2, -1)) + x_real, x_imag = pair[..., 0, :], pair[..., 1, :] + x_rotated = ops.concatenate([-x_imag, x_real], axis=-1) + else: + raise ValueError( + f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2." + ) + out = x_f * ops.cast(cos, "float32") + x_rotated * ops.cast(sin, "float32") + return ops.cast(out, x_dtype) + + +def _rope_angles(index, dim, theta): + """Outer product of positions with ``1 / theta^(i/dim)`` frequencies.""" + assert dim % 2 == 0 + freqs = np.outer( + np.asarray(index, dtype=np.float64), + 1.0 + / np.power( + float(theta), + np.arange(0, dim, 2, dtype=np.float64) / dim, + ), + ) + return freqs.astype(np.float32) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageRMSNorm(layers.Layer): + """Diffusers ``RMSNorm`` (learned weight, ones init, ``eps=1e-6``). + + Normalizes over the last axis. Used for ``txt_norm`` and per-head Q/K norms + (``attn.norm_q``, ``norm_k``, ``norm_added_q``, ``norm_added_k``). + """ + + def __init__(self, eps=NORM_EPS, module_path=None, **kwargs): + if module_path is not None: + kwargs.setdefault("name", safe_name(module_path)) + super().__init__(**kwargs) + self.eps = eps + self.module_path = module_path + + def build(self, input_shape): + self.weight = self.add_weight( + name="weight", + shape=(int(input_shape[-1]),), + initializer="ones", + trainable=True, + ) + self.built = True + + def call(self, x): + dtype = x.dtype + x32 = ops.cast(x, "float32") + variance = ops.mean(ops.square(x32), axis=-1, keepdims=True) + x32 = x32 * ops.rsqrt(variance + self.eps) + return ops.cast(x32, dtype) * self.weight + + def compute_output_shape(self, input_shape): + return tuple(input_shape) + + def get_config(self): + config = super().get_config() + config.update({"eps": self.eps, "module_path": self.module_path}) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageTimestepProjEmbeddings(layers.Layer): + """Diffusers ``QwenTimestepProjEmbeddings``. + + ``Timesteps(256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000)`` + followed by ``TimestepEmbedding`` (``linear_1`` / SiLU / ``linear_2``) to + ``embedding_dim``. The Diffusers ``scale=1000`` is applied by multiplying + timesteps before :func:`timestep_embedding` (equivalent to scaling the + sinusoidal arguments). + + Weight path: ``{module_path}.timestep_embedder.linear_{1,2}``. + """ + + def __init__( + self, + embedding_dim, + module_path="time_text_embed", + time_freq_dim=256, + scale=1000.0, + **kwargs, + ): + kwargs.setdefault("name", safe_name(module_path)) + super().__init__(**kwargs) + self.embedding_dim = embedding_dim + self.module_path = module_path + self.time_freq_dim = time_freq_dim + self.scale = scale + embedder = f"{module_path}.timestep_embedder" + self.linear_1 = layers.Dense( + embedding_dim, name=safe_name(f"{embedder}.linear_1") + ) + self.linear_2 = layers.Dense( + embedding_dim, name=safe_name(f"{embedder}.linear_2") + ) + + def build(self, timestep_shape): + freq_shape = (timestep_shape[0], self.time_freq_dim) + self.linear_1.build(freq_shape) + self.linear_2.build((timestep_shape[0], self.embedding_dim)) + self.built = True + + def call(self, timestep): + t = ops.cast(timestep, "float32") * self.scale + emb = timestep_embedding( + t, + self.time_freq_dim, + flip_sin_to_cos=True, + downscale_freq_shift=0, + ) + emb = self.linear_1(emb) + emb = ops.silu(emb) + return self.linear_2(emb) + + def compute_output_shape(self, timestep_shape): + return (timestep_shape[0], self.embedding_dim) + + def get_config(self): + config = super().get_config() + config.update( + { + "embedding_dim": self.embedding_dim, + "module_path": self.module_path, + "time_freq_dim": self.time_freq_dim, + "scale": self.scale, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageEmbedRope(layers.Layer): + """Diffusers ``QwenEmbedRope`` with real-valued rotation angles. + + Precomputes positive/negative position frequencies as angles (not complex) + for backend portability. ``call(img_h, img_w, txt_seq_len, frame=1)`` returns + ``(vid_freqs, txt_freqs)`` each ``[S, sum(axes_dim) // 2]`` for + :func:`apply_rotary_emb_qwen`. + + Defaults match Qwen-Image T2I: ``theta=10000``, ``axes_dim=(16, 56, 56)``, + ``scale_rope=True``. + """ + + def __init__( + self, + theta=10000, + axes_dim=(16, 56, 56), + scale_rope=True, + module_path="pos_embed", + **kwargs, + ): + kwargs.setdefault("name", safe_name(module_path)) + super().__init__(**kwargs) + self.theta = int(theta) + self.axes_dim = tuple(int(d) for d in axes_dim) + self.scale_rope = bool(scale_rope) + self.module_path = module_path + self.rope_dim = sum(self.axes_dim) // 2 + self._axis_halves = [d // 2 for d in self.axes_dim] + + pos_index = np.arange(ROPE_MAX_INDEX) + neg_index = np.flip(pos_index) * -1 - 1 + pos_freqs = np.concatenate( + [_rope_angles(pos_index, d, self.theta) for d in self.axes_dim], + axis=1, + ) + neg_freqs = np.concatenate( + [_rope_angles(neg_index, d, self.theta) for d in self.axes_dim], + axis=1, + ) + self._pos_freqs_np = pos_freqs + self._neg_freqs_np = neg_freqs + + def build(self, input_shape=None): + self.pos_freqs = self.add_weight( + name="pos_freqs", + shape=(ROPE_MAX_INDEX, self.rope_dim), + initializer=keras.initializers.Constant(self._pos_freqs_np), + trainable=False, + dtype="float32", + autocast=False, + ) + self.neg_freqs = self.add_weight( + name="neg_freqs", + shape=(ROPE_MAX_INDEX, self.rope_dim), + initializer=keras.initializers.Constant(self._neg_freqs_np), + trainable=False, + dtype="float32", + autocast=False, + ) + self.built = True + + def _split_axes(self, freqs): + h0, h1, h2 = self._axis_halves + return ( + freqs[:, :h0], + freqs[:, h0 : h0 + h1], + freqs[:, h0 + h1 :], + ) + + def _video_freqs(self, frame, height, width, idx=0): + frame, height, width, idx = int(frame), int(height), int(width), int(idx) + seq_lens = frame * height * width + freqs_pos = self._split_axes(self.pos_freqs) + freqs_neg = self._split_axes(self.neg_freqs) + + freqs_frame = freqs_pos[0][idx : idx + frame] + freqs_frame = ops.reshape(freqs_frame, (frame, 1, 1, self._axis_halves[0])) + freqs_frame = ops.broadcast_to( + freqs_frame, (frame, height, width, self._axis_halves[0]) + ) + + if self.scale_rope: + h_neg = height - height // 2 + w_neg = width - width // 2 + freqs_height = ops.concatenate( + [freqs_neg[1][-h_neg:], freqs_pos[1][: height // 2]], axis=0 + ) + freqs_width = ops.concatenate( + [freqs_neg[2][-w_neg:], freqs_pos[2][: width // 2]], axis=0 + ) + else: + freqs_height = freqs_pos[1][:height] + freqs_width = freqs_pos[2][:width] + + freqs_height = ops.reshape(freqs_height, (1, height, 1, self._axis_halves[1])) + freqs_height = ops.broadcast_to( + freqs_height, (frame, height, width, self._axis_halves[1]) + ) + freqs_width = ops.reshape(freqs_width, (1, 1, width, self._axis_halves[2])) + freqs_width = ops.broadcast_to( + freqs_width, (frame, height, width, self._axis_halves[2]) + ) + + freqs = ops.concatenate([freqs_frame, freqs_height, freqs_width], axis=-1) + return ops.reshape(freqs, (seq_lens, self.rope_dim)) + + def call(self, img_h, img_w, txt_seq_len, frame=1): + """Return ``(vid_freqs, txt_freqs)`` angle tables for T2I RoPE. + + Args: + img_h / img_w: Latent patch grid height / width (Python ints preferred). + txt_seq_len: Text token count (matches encoder sequence length). + frame: Temporal size (``1`` for images). + """ + if not self.built: + self.build(None) + height, width = int(img_h), int(img_w) + txt_seq_len = int(txt_seq_len) + frame = int(frame) + + vid_freqs = self._video_freqs(frame, height, width, idx=0) + if self.scale_rope: + max_vid_index = max(height // 2, width // 2) + else: + max_vid_index = max(height, width) + txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + txt_seq_len] + return vid_freqs, txt_freqs + + def get_config(self): + config = super().get_config() + config.update( + { + "theta": self.theta, + "axes_dim": list(self.axes_dim), + "scale_rope": self.scale_rope, + "module_path": self.module_path, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageDoubleStreamAttention(layers.Layer): + """Joint text/image attention (Diffusers ``Attention`` + ``QwenDoubleStreamAttnProcessor2_0``). + + Separate QKV for the image stream (``to_q`` / ``to_k`` / ``to_v``) and text + stream (``add_q_proj`` / ``add_k_proj`` / ``add_v_proj``), RMSNorm on Q/K for + both, RoPE on both, then joint attention over ``[text, image]`` (text-first, + unlike SD3's image-first concat). Outputs project via ``to_out.0`` and + ``to_add_out``. + + Typical Qwen-Image sizes: ``dim=3072``, ``heads=24``, ``head_dim=128``. + """ + + def __init__( + self, + dim, + num_attention_heads, + attention_head_dim=None, + module_path="attn", + eps=NORM_EPS, + attn_implementation="fused", + **kwargs, + ): + kwargs.setdefault("name", safe_name(module_path)) + super().__init__(**kwargs) + self._convert_input_args = False + self.dim = dim + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim or (dim // num_attention_heads) + self.module_path = module_path + self.eps = eps + self.attn_implementation = attn_implementation + self.scale = self.attention_head_dim**-0.5 + + self.to_q = self._dense("to_q") + self.to_k = self._dense("to_k") + self.to_v = self._dense("to_v") + self.to_out = self._dense("to_out.0") + self.add_q_proj = self._dense("add_q_proj") + self.add_k_proj = self._dense("add_k_proj") + self.add_v_proj = self._dense("add_v_proj") + self.to_add_out = self._dense("to_add_out") + + self.norm_q = QwenImageRMSNorm(eps, module_path=f"{module_path}.norm_q") + self.norm_k = QwenImageRMSNorm(eps, module_path=f"{module_path}.norm_k") + self.norm_added_q = QwenImageRMSNorm( + eps, module_path=f"{module_path}.norm_added_q" + ) + self.norm_added_k = QwenImageRMSNorm( + eps, module_path=f"{module_path}.norm_added_k" + ) + + def _dense(self, leaf): + return layers.Dense(self.dim, name=safe_name(f"{self.module_path}.{leaf}")) + + def build( + self, + hidden_states_shape, + encoder_hidden_states_shape=None, + image_rotary_emb_shape=None, + encoder_hidden_states_mask_shape=None, + ): + del image_rotary_emb_shape, encoder_hidden_states_mask_shape + if encoder_hidden_states_shape is None: + encoder_hidden_states_shape = hidden_states_shape + for layer in (self.to_q, self.to_k, self.to_v, self.to_out): + layer.build(hidden_states_shape) + for layer in ( + self.add_q_proj, + self.add_k_proj, + self.add_v_proj, + self.to_add_out, + ): + layer.build(encoder_hidden_states_shape) + head_shape = ( + None, + None, + self.num_attention_heads, + self.attention_head_dim, + ) + for norm in (self.norm_q, self.norm_k, self.norm_added_q, self.norm_added_k): + norm.build(head_shape) + self.built = True + + def _to_heads(self, t): + # [B, S, dim] -> [B, S, H, D] + t = ops.reshape( + t, + (-1, ops.shape(t)[1], self.num_attention_heads, self.attention_head_dim), + ) + return t + + def call( + self, + hidden_states, + encoder_hidden_states, + image_rotary_emb=None, + encoder_hidden_states_mask=None, + ): + seq_txt = ops.shape(encoder_hidden_states)[1] + + img_q = self._to_heads(self.to_q(hidden_states)) + img_k = self._to_heads(self.to_k(hidden_states)) + img_v = self._to_heads(self.to_v(hidden_states)) + txt_q = self._to_heads(self.add_q_proj(encoder_hidden_states)) + txt_k = self._to_heads(self.add_k_proj(encoder_hidden_states)) + txt_v = self._to_heads(self.add_v_proj(encoder_hidden_states)) + + img_q, img_k = self.norm_q(img_q), self.norm_k(img_k) + txt_q, txt_k = self.norm_added_q(txt_q), self.norm_added_k(txt_k) + + if image_rotary_emb is not None: + img_freqs, txt_freqs = image_rotary_emb + img_q = apply_rotary_emb_qwen(img_q, img_freqs) + img_k = apply_rotary_emb_qwen(img_k, img_freqs) + txt_q = apply_rotary_emb_qwen(txt_q, txt_freqs) + txt_k = apply_rotary_emb_qwen(txt_k, txt_freqs) + + query = ops.concatenate([txt_q, img_q], axis=1) + key = ops.concatenate([txt_k, img_k], axis=1) + value = ops.concatenate([txt_v, img_v], axis=1) + + query = ops.transpose(query, (0, 2, 1, 3)) + key = ops.transpose(key, (0, 2, 1, 3)) + value = ops.transpose(value, (0, 2, 1, 3)) + + attention_mask = None + if encoder_hidden_states_mask is not None: + batch = ops.shape(hidden_states)[0] + seq_img = ops.shape(hidden_states)[1] + img_mask = ops.ones( + (batch, seq_img), dtype=encoder_hidden_states_mask.dtype + ) + joint_mask = ops.concatenate([encoder_hidden_states_mask, img_mask], axis=1) + keep = ops.cast(joint_mask, "float32") + attention_mask = (1.0 - keep) * MASK_NEG + attention_mask = attention_mask[:, None, None, :] + + out = fused_attention( + query, + key, + value, + self.scale, + attention_mask=attention_mask, + attn_implementation=active_attn_implementation() + or self.attn_implementation, + ) + out = ops.transpose(out, (0, 2, 1, 3)) + out = ops.reshape(out, (-1, ops.shape(out)[1], self.dim)) + + txt_out = self.to_add_out(out[:, :seq_txt]) + img_out = self.to_out(out[:, seq_txt:]) + return img_out, txt_out + + def compute_output_shape(self, input_shape, context_shape=None): + if context_shape is None and isinstance(input_shape, (list, tuple)): + if len(input_shape) >= 2: + input_shape, context_shape = input_shape[0], input_shape[1] + img_out = tuple(input_shape[:-1]) + (self.dim,) + txt_out = tuple(context_shape[:-1]) + (self.dim,) + return img_out, txt_out + + def get_config(self): + config = super().get_config() + config.update( + { + "dim": self.dim, + "num_attention_heads": self.num_attention_heads, + "attention_head_dim": self.attention_head_dim, + "module_path": self.module_path, + "eps": self.eps, + "attn_implementation": self.attn_implementation, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageTransformerBlock(layers.Layer): + """One double-stream DiT block (Diffusers ``QwenImageTransformerBlock``). + + Image and text streams each get SiLU+Dense(``6 * dim``) modulation + (``img_mod.1`` / ``txt_mod.1``), affine-free LayerNorm, joint attention, + and a gated ``gelu-approximate`` MLP. ``zero_cond_t`` is not supported in + this v1 port (always ``False``). + + ``call`` returns ``(encoder_hidden_states, hidden_states)`` — text then + image — matching Diffusers. + """ + + def __init__( + self, + dim, + num_attention_heads, + attention_head_dim=None, + module_path="transformer_blocks.0", + eps=NORM_EPS, + attn_implementation="fused", + **kwargs, + ): + kwargs.setdefault("name", safe_name(module_path)) + super().__init__(**kwargs) + self._convert_input_args = False + self.dim = dim + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim or (dim // num_attention_heads) + self.module_path = module_path + self.eps = eps + self.attn_implementation = attn_implementation + + self.img_mod = layers.Dense(6 * dim, name=safe_name(f"{module_path}.img_mod.1")) + self.txt_mod = layers.Dense(6 * dim, name=safe_name(f"{module_path}.txt_mod.1")) + self.img_norm1 = layers.LayerNormalization( + epsilon=eps, + center=False, + scale=False, + name=safe_name(f"{module_path}.img_norm1"), + ) + self.txt_norm1 = layers.LayerNormalization( + epsilon=eps, + center=False, + scale=False, + name=safe_name(f"{module_path}.txt_norm1"), + ) + self.img_norm2 = layers.LayerNormalization( + epsilon=eps, + center=False, + scale=False, + name=safe_name(f"{module_path}.img_norm2"), + ) + self.txt_norm2 = layers.LayerNormalization( + epsilon=eps, + center=False, + scale=False, + name=safe_name(f"{module_path}.txt_norm2"), + ) + self.attn = QwenImageDoubleStreamAttention( + dim=dim, + num_attention_heads=num_attention_heads, + attention_head_dim=self.attention_head_dim, + module_path=f"{module_path}.attn", + eps=eps, + attn_implementation=attn_implementation, + ) + self.img_mlp = StableDiffusion3GELUFeedForward( + dim, module_path=f"{module_path}.img_mlp" + ) + self.txt_mlp = StableDiffusion3GELUFeedForward( + dim, module_path=f"{module_path}.txt_mlp" + ) + + def build( + self, + hidden_states_shape, + encoder_hidden_states_shape=None, + temb_shape=None, + encoder_hidden_states_mask_shape=None, + image_rotary_emb_shape=None, + ): + del encoder_hidden_states_mask_shape, image_rotary_emb_shape + if isinstance(hidden_states_shape, (list, tuple)) and temb_shape is None: + if len(hidden_states_shape) >= 3 and hasattr( + hidden_states_shape[0], "__len__" + ): + ( + hidden_states_shape, + encoder_hidden_states_shape, + temb_shape, + ) = ( + hidden_states_shape[0], + hidden_states_shape[1], + hidden_states_shape[2], + ) + if encoder_hidden_states_shape is None or temb_shape is None: + raise ValueError( + "QwenImageTransformerBlock.build expects image, text, and temb shapes." + ) + self.img_mod.build(temb_shape) + self.txt_mod.build(temb_shape) + self.img_norm1.build(hidden_states_shape) + self.txt_norm1.build(encoder_hidden_states_shape) + self.img_norm2.build(hidden_states_shape) + self.txt_norm2.build(encoder_hidden_states_shape) + self.attn.build(hidden_states_shape, encoder_hidden_states_shape) + self.img_mlp.build(hidden_states_shape) + self.txt_mlp.build(encoder_hidden_states_shape) + self.built = True + + def _modulate(self, x, mod_params): + shift, scale, gate = ops.split(mod_params, 3, axis=-1) + modulated = x * (1.0 + scale[:, None, :]) + shift[:, None, :] + return modulated, gate + + def call( + self, + hidden_states, + encoder_hidden_states, + temb, + encoder_hidden_states_mask=None, + image_rotary_emb=None, + ): + img_mod_params = self.img_mod(ops.silu(temb)) + txt_mod_params = self.txt_mod(ops.silu(temb)) + img_mod1, img_mod2 = ops.split(img_mod_params, 2, axis=-1) + txt_mod1, txt_mod2 = ops.split(txt_mod_params, 2, axis=-1) + + img_modulated, img_gate1 = self._modulate( + self.img_norm1(hidden_states), img_mod1 + ) + txt_modulated, txt_gate1 = self._modulate( + self.txt_norm1(encoder_hidden_states), txt_mod1 + ) + + img_attn, txt_attn = self.attn( + img_modulated, + txt_modulated, + image_rotary_emb=image_rotary_emb, + encoder_hidden_states_mask=encoder_hidden_states_mask, + ) + hidden_states = hidden_states + img_gate1[:, None, :] * img_attn + encoder_hidden_states = encoder_hidden_states + txt_gate1[:, None, :] * txt_attn + + img_modulated2, img_gate2 = self._modulate( + self.img_norm2(hidden_states), img_mod2 + ) + hidden_states = hidden_states + img_gate2[:, None, :] * self.img_mlp( + img_modulated2 + ) + + txt_modulated2, txt_gate2 = self._modulate( + self.txt_norm2(encoder_hidden_states), txt_mod2 + ) + encoder_hidden_states = encoder_hidden_states + txt_gate2[ + :, None, : + ] * self.txt_mlp(txt_modulated2) + + if keras.backend.standardize_dtype(encoder_hidden_states.dtype) == "float16": + encoder_hidden_states = ops.clip(encoder_hidden_states, -65504.0, 65504.0) + if keras.backend.standardize_dtype(hidden_states.dtype) == "float16": + hidden_states = ops.clip(hidden_states, -65504.0, 65504.0) + + return encoder_hidden_states, hidden_states + + def compute_output_shape(self, input_shape): + if ( + isinstance(input_shape, (list, tuple)) + and len(input_shape) >= 2 + and hasattr(input_shape[0], "__len__") + ): + img_shape, txt_shape = input_shape[0], input_shape[1] + return tuple(txt_shape), tuple(img_shape) + return tuple(input_shape), tuple(input_shape) + + def get_config(self): + config = super().get_config() + config.update( + { + "dim": self.dim, + "num_attention_heads": self.num_attention_heads, + "attention_head_dim": self.attention_head_dim, + "module_path": self.module_path, + "eps": self.eps, + "attn_implementation": self.attn_implementation, + } + ) + return config + + +QwenImageAdaLayerNormContinuous = StableDiffusion3AdaLayerNorm + +__all__ = [ + "apply_rotary_emb_qwen", + "QwenImageRMSNorm", + "QwenImageTimestepProjEmbeddings", + "QwenImageEmbedRope", + "QwenImageDoubleStreamAttention", + "QwenImageTransformerBlock", + "QwenImageAdaLayerNormContinuous", + "StableDiffusion3AdaLayerNorm", + "StableDiffusion3GELUFeedForward", +] diff --git a/zeromodels/models/qwen_image/qwen_image_model.py b/zeromodels/models/qwen_image/qwen_image_model.py new file mode 100644 index 00000000..4ff2d36d --- /dev/null +++ b/zeromodels/models/qwen_image/qwen_image_model.py @@ -0,0 +1,811 @@ +from __future__ import annotations + +import keras +import numpy as np +from keras import layers, ops + +from zeromodels.base import BaseDiffusion, BaseModel +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 ( + DEFAULT_LATENTS_MEAN, + DEFAULT_LATENTS_STD, + QwenImageConfig, + QwenImageTextConfig, + QwenImageTransformerConfig, + QwenImageVAEConfig, +) +from zeromodels.models.qwen_image.qwen_image_layers import ( + QwenImageAdaLayerNormContinuous, + QwenImageEmbedRope, + QwenImageRMSNorm, + QwenImageTimestepProjEmbeddings, + QwenImageTransformerBlock, +) +from zeromodels.models.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" +) + + +def pack_latents(latents, height, width): + """Pack ``(B, H, W, C)`` latents into ``(B, H/2 * W/2, C*4)`` (Diffusers).""" + batch = ops.shape(latents)[0] + channels = ops.shape(latents)[-1] + latents = ops.reshape(latents, (batch, height // 2, 2, width // 2, 2, channels)) + latents = ops.transpose(latents, (0, 1, 3, 5, 2, 4)) + return ops.reshape(latents, (batch, (height // 2) * (width // 2), channels * 4)) + + +def unpack_latents(latents, height, width, channels): + """Unpack ``(B, seq, C*4)`` to ``(B, H, W, C)``.""" + batch = ops.shape(latents)[0] + latents = ops.reshape(latents, (batch, height // 2, width // 2, channels, 2, 2)) + latents = ops.transpose(latents, (0, 1, 4, 2, 5, 3)) + return ops.reshape(latents, (batch, height, width, channels)) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class AutoencoderKLQwenImage(BaseModel): + """Qwen-Image VAE (Diffusers ``AutoencoderKLQwenImage``), channels-last NDHWC. + + Encode / decode a single image frame (``T=1``). The functional graph is built + for ``sample_size`` but conv weights are resolution-independent. Public + helpers accept channels-last HWC images or Diffusers NCDHW 5D tensors. + + Latent normalisation with ``latents_mean`` / ``latents_std`` is left to the + pipeline (stored on the model for that purpose). + """ + + config_class = QwenImageVAEConfig + HF_MODEL_TYPE = None + + def __init__( + self, + base_dim=96, + z_dim=16, + dim_mult=(1, 2, 4, 4), + num_res_blocks=2, + attn_scales=(), + temperal_downsample=(False, True, True), + dropout=0.0, + input_channels=3, + latents_mean=DEFAULT_LATENTS_MEAN, + latents_std=DEFAULT_LATENTS_STD, + sample_size=1024, + apply_temporal=False, + name="AutoencoderKLQwenImage", + **kwargs, + ): + dim_mult = tuple(dim_mult) + temperal_downsample = tuple(temperal_downsample) + attn_scales = tuple(attn_scales) + latents_mean = tuple(latents_mean) + latents_std = tuple(latents_std) + temperal_upsample = tuple(reversed(temperal_downsample)) + + h_img, w_img = ( + sample_size + if isinstance(sample_size, (tuple, list)) + else (sample_size, sample_size) + ) + spatial_compression_ratio = 2 ** len(temperal_downsample) + h_lat, w_lat = ( + h_img // spatial_compression_ratio, + w_img // spatial_compression_ratio, + ) + + encoder = QwenImageEncoder3d( + dim=base_dim, + z_dim=z_dim * 2, + dim_mult=dim_mult, + num_res_blocks=num_res_blocks, + attn_scales=attn_scales, + temperal_downsample=temperal_downsample, + dropout=dropout, + input_channels=input_channels, + module_path="encoder", + apply_temporal=apply_temporal, + ) + decoder = QwenImageDecoder3d( + dim=base_dim, + z_dim=z_dim, + dim_mult=dim_mult, + num_res_blocks=num_res_blocks, + attn_scales=attn_scales, + temperal_upsample=temperal_upsample, + dropout=dropout, + input_channels=input_channels, + module_path="decoder", + apply_temporal=apply_temporal, + ) + quant_conv = QwenImageCausalConv3d( + z_dim * 2, kernel_size=1, padding=0, module_path="quant_conv" + ) + post_quant_conv = QwenImageCausalConv3d( + z_dim, kernel_size=1, padding=0, module_path="post_quant_conv" + ) + + image_in = layers.Input(shape=(h_img, w_img, input_channels), name="image") + latent_in = layers.Input(shape=(h_lat, w_lat, z_dim), name="latent") + + image_5d = ops.expand_dims(image_in, axis=1) + moments_5d = quant_conv(encoder(image_5d)) + moments = ops.squeeze(moments_5d, axis=1) + + latent_5d = ops.expand_dims(latent_in, axis=1) + decoded_5d = decoder(post_quant_conv(latent_5d)) + decoded = ops.squeeze(decoded_5d, axis=1) + decoded = ops.clip(decoded, -1.0, 1.0) + + super().__init__( + inputs={"image": image_in, "latent": latent_in}, + outputs={"moments": moments, "sample": decoded}, + name=name, + **kwargs, + ) + + self.base_dim = base_dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_upsample + self.dropout = dropout + self.input_channels = input_channels + self.latents_mean = latents_mean + self.latents_std = latents_std + self.sample_size = sample_size + self.apply_temporal = apply_temporal + self.spatial_compression_ratio = spatial_compression_ratio + self.vae_scale_factor = spatial_compression_ratio + self.encoder = encoder + self.decoder = decoder + self.quant_conv = quant_conv + self.post_quant_conv = post_quant_conv + + def _to_ndhwc(self, x, is_latent=False): + """Normalize inputs to NDHWC ``(B, T, H, W, C)``.""" + static_ndim = len(x.shape) + if static_ndim == 4: + return ops.expand_dims(x, axis=1) + if static_ndim == 5: + c1 = int(x.shape[1]) if x.shape[1] is not None else None + c_last = int(x.shape[-1]) if x.shape[-1] is not None else None + expect = self.z_dim if is_latent else self.input_channels + if c1 == expect and c_last != expect: + return ops.transpose(x, (0, 2, 3, 4, 1)) + return x + raise ValueError(f"Expected 4D or 5D tensor, got shape {x.shape}") + + def _maybe_squeeze_t(self, x, original_was_4d): + if original_was_4d: + return ops.squeeze(x, axis=1) + return x + + def encode(self, x, sample=False, seed=None, return_ncdhw=False): + """Encode image(s) to latents (mean, or reparameterized sample). + + Args: + x: ``(B, H, W, 3)`` HWC, ``(B, 1, H, W, 3)`` NDHWC, or + ``(B, 3, 1, H, W)`` Diffusers NCDHW. + sample: If True, draw ``z ~ N(mean, std)``; else return mean. + seed: RNG seed for sampling. + return_ncdhw: If True, return Diffusers layout ``(B, C, T, H, W)``. + """ + was_4d = len(x.shape) == 4 + x5 = self._to_ndhwc(x, is_latent=False) + moments = self.quant_conv(self.encoder(x5)) + mean, logvar = ops.split(moments, 2, axis=-1) + if sample: + logvar = ops.clip(logvar, -30.0, 20.0) + std = ops.exp(0.5 * logvar) + noise = keras.random.normal(ops.shape(mean), dtype=mean.dtype, seed=seed) + z = mean + std * noise + else: + z = mean + if return_ncdhw: + return ops.transpose(z, (0, 4, 1, 2, 3)) + return self._maybe_squeeze_t(z, was_4d) + + def decode(self, z, return_ncdhw=False): + """Decode latents to RGB in ``[-1, 1]``. + + Args: + z: ``(B, h, w, z_dim)``, ``(B, 1, h, w, z_dim)`` NDHWC, or + ``(B, z_dim, 1, h, w)`` NCDHW. + return_ncdhw: If True, return Diffusers layout. + """ + was_4d = len(z.shape) == 4 + z5 = self._to_ndhwc(z, is_latent=True) + x = self.decoder(self.post_quant_conv(z5)) + x = ops.clip(x, -1.0, 1.0) + if return_ncdhw: + return ops.transpose(x, (0, 4, 1, 2, 3)) + return self._maybe_squeeze_t(x, was_4d) + + def get_config(self): + config = super().get_config() + config.update(self.config.constructor_kwargs()) + config["apply_temporal"] = self.apply_temporal + return config + + @classmethod + def from_diffusers_config(cls, config, sample_size=1024, **kwargs): + return cls( + base_dim=config.get("base_dim", 96), + z_dim=config.get("z_dim", 16), + dim_mult=tuple(config.get("dim_mult", (1, 2, 4, 4))), + num_res_blocks=config.get("num_res_blocks", 2), + attn_scales=tuple(config.get("attn_scales", ())), + temperal_downsample=tuple( + config.get("temperal_downsample", (False, True, True)) + ), + dropout=config.get("dropout", 0.0), + input_channels=config.get("input_channels", 3), + latents_mean=tuple(config.get("latents_mean", DEFAULT_LATENTS_MEAN)), + latents_std=tuple(config.get("latents_std", DEFAULT_LATENTS_STD)), + sample_size=sample_size, + **kwargs, + ) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageTransformer2DModel(BaseModel): + """Qwen-Image denoiser (Diffusers ``QwenImageTransformer2DModel``). + + Double-stream MMDiT over **packed** latents ``(B, seq, in_channels)`` and + text features ``(B, text_seq, joint_attention_dim)``. Inputs dict keys: + ``sample``, ``timestep``, ``encoder_hidden_states``, optional + ``encoder_hidden_states_mask``. Output ``{"sample": packed velocity}``. + + Built for a fixed packed sequence length derived from ``sample_size`` + (latent side before 2×2 packing). + """ + + HF_MODEL_TYPE = None + config_class = QwenImageTransformerConfig + + def __init__( + self, + patch_size=2, + in_channels=64, + out_channels=16, + num_layers=60, + attention_head_dim=128, + num_attention_heads=24, + joint_attention_dim=3584, + axes_dims_rope=(16, 56, 56), + guidance_embeds=False, + sample_size=128, + text_seq_len=512, + name="QwenImageTransformer2DModel", + **kwargs, + ): + del guidance_embeds # base Qwen-Image is not guidance-distilled + keras_kwargs = {k: kwargs.pop(k) for k in ("trainable", "dtype") if k in kwargs} + axes_dims_rope = tuple(axes_dims_rope) + inner_dim = num_attention_heads * attention_head_dim + sample_h = ( + sample_size[0] if isinstance(sample_size, (tuple, list)) else sample_size + ) + pack_h = pack_w = sample_h // patch_size + packed_seq = pack_h * pack_w + + pos_embed = QwenImageEmbedRope( + theta=10000, + axes_dim=axes_dims_rope, + scale_rope=True, + module_path="pos_embed", + ) + time_text_embed = QwenImageTimestepProjEmbeddings( + embedding_dim=inner_dim, module_path="time_text_embed" + ) + txt_norm = QwenImageRMSNorm(eps=1e-6, module_path="txt_norm") + img_in = layers.Dense(inner_dim, name=safe_name("img_in")) + txt_in = layers.Dense(inner_dim, name=safe_name("txt_in")) + blocks = [ + QwenImageTransformerBlock( + dim=inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + module_path=f"transformer_blocks.{i}", + ) + for i in range(num_layers) + ] + norm_out = QwenImageAdaLayerNormContinuous( + inner_dim, num_chunks=2, module_path="norm_out" + ) + proj_out = layers.Dense( + patch_size * patch_size * out_channels, + name=safe_name("proj_out"), + ) + + sample_in = layers.Input(shape=(packed_seq, in_channels), name="sample") + timestep_in = layers.Input(shape=(), name="timestep") + enc_in = layers.Input( + shape=(text_seq_len, joint_attention_dim), name="encoder_hidden_states" + ) + enc_mask_in = layers.Input( + shape=(text_seq_len,), dtype="int32", name="encoder_hidden_states_mask" + ) + + hidden = img_in(sample_in) + encoder = txt_in(txt_norm(enc_in)) + temb = time_text_embed(timestep_in) + pos_embed.build(None) + vid_freqs, txt_freqs = pos_embed.call(pack_h, pack_w, text_seq_len, frame=1) + # Bake RoPE tables into the graph as constants (static T2I resolution). + image_rotary_emb = ( + ops.convert_to_tensor(ops.convert_to_numpy(vid_freqs)), + ops.convert_to_tensor(ops.convert_to_numpy(txt_freqs)), + ) + mask = ops.cast(enc_mask_in, "bool") + for block in blocks: + encoder, hidden = block( + hidden, + encoder, + temb, + encoder_hidden_states_mask=mask, + image_rotary_emb=image_rotary_emb, + ) + hidden = norm_out([hidden, temb]) + output = proj_out(hidden) + + super().__init__( + inputs={ + "sample": sample_in, + "timestep": timestep_in, + "encoder_hidden_states": enc_in, + "encoder_hidden_states_mask": enc_mask_in, + }, + outputs={"sample": output}, + name=name, + **keras_kwargs, + ) + + self.patch_size = patch_size + self.in_channels = in_channels + self.out_channels = out_channels + self.num_layers = num_layers + self.attention_head_dim = attention_head_dim + self.num_attention_heads = num_attention_heads + self.joint_attention_dim = joint_attention_dim + self.axes_dims_rope = axes_dims_rope + self.sample_size = sample_size + self.text_seq_len = text_seq_len + self.inner_dim = inner_dim + self.pack_h = pack_h + self.pack_w = pack_w + self.packed_seq = packed_seq + self.pos_embed = pos_embed + self.time_text_embed = time_text_embed + self.txt_norm = txt_norm + self.img_in = img_in + self.txt_in = txt_in + self.transformer_blocks = blocks + self.norm_out = norm_out + self.proj_out = proj_out + + @classmethod + def kwargs_from_diffusers_config(cls, cfg): + return { + "patch_size": cfg.get("patch_size", 2), + "in_channels": cfg.get("in_channels", 64), + "out_channels": cfg.get("out_channels", 16), + "num_layers": cfg.get("num_layers", 60), + "attention_head_dim": cfg.get("attention_head_dim", 128), + "num_attention_heads": cfg.get("num_attention_heads", 24), + "joint_attention_dim": cfg.get("joint_attention_dim", 3584), + "axes_dims_rope": tuple(cfg.get("axes_dims_rope", (16, 56, 56))), + "guidance_embeds": bool(cfg.get("guidance_embeds", False)), + } + + +@keras.saving.register_keras_serializable(package="zeromodels") +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`. + Inputs ``input_ids`` / ``attention_mask``; output ``last_hidden_state``. + """ + + HF_MODEL_TYPE = None + config_class = QwenImageTextConfig + + def __init__(self, max_seq_len=1024, name="text_encoder", **kwargs): + kwargs["build_vision"] = False + super().__init__(name=name, **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. + + 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. + """ + + 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) + + def build_components(self, config): + d, v = config.transformer_config, config.vae_config + transformer = QwenImageTransformer2DModel( + patch_size=d.patch_size, + in_channels=d.in_channels, + out_channels=d.out_channels, + num_layers=d.num_layers, + attention_head_dim=d.attention_head_dim, + num_attention_heads=d.num_attention_heads, + joint_attention_dim=d.joint_attention_dim, + axes_dims_rope=d.axes_dims_rope, + guidance_embeds=d.guidance_embeds, + sample_size=d.sample_size, + text_seq_len=config.max_sequence_length, + ) + vae = AutoencoderKLQwenImage( + base_dim=v.base_dim, + z_dim=v.z_dim, + dim_mult=v.dim_mult, + num_res_blocks=v.num_res_blocks, + attn_scales=v.attn_scales, + temperal_downsample=v.temperal_downsample, + dropout=v.dropout, + input_channels=v.input_channels, + latents_mean=v.latents_mean, + latents_std=v.latents_std, + sample_size=v.sample_size, + ) + text_encoder = QwenImageTextEncoderModel(config.text_config) + return {"transformer": transformer, "vae": vae, "text_encoder": text_encoder} + + def build_graph(self, config, components): + d, v, t = config.transformer_config, config.vae_config, config.text_config + transformer, vae, text_encoder = ( + components["transformer"], + components["vae"], + components["text_encoder"], + ) + text_seq = config.max_sequence_length + img_h = img_w = ( + v.sample_size + if not isinstance(v.sample_size, (tuple, list)) + else v.sample_size[0] + ) + lat_h = img_h // vae.vae_scale_factor + lat_w = img_w // vae.vae_scale_factor + + inputs = { + "sample": layers.Input( + shape=(transformer.packed_seq, d.in_channels), name="sample" + ), + "timestep": layers.Input(shape=(), name="timestep"), + "encoder_hidden_states": layers.Input( + shape=(text_seq, d.joint_attention_dim), name="encoder_hidden_states" + ), + "encoder_hidden_states_mask": layers.Input( + shape=(text_seq,), dtype="int32", name="encoder_hidden_states_mask" + ), + "image": layers.Input(shape=(img_h, img_w, 3), name="image"), + "latent": layers.Input(shape=(lat_h, lat_w, v.z_dim), name="latent"), + "token_ids": layers.Input( + shape=(t.max_seq_len,), dtype="int32", name="token_ids" + ), + "padding_mask": layers.Input( + shape=(t.max_seq_len,), dtype="int32", name="padding_mask" + ), + } + noise_pred = transformer( + { + "sample": inputs["sample"], + "timestep": inputs["timestep"], + "encoder_hidden_states": inputs["encoder_hidden_states"], + "encoder_hidden_states_mask": inputs["encoder_hidden_states_mask"], + } + )["sample"] + vae_out = vae({"image": inputs["image"], "latent": inputs["latent"]}) + text_out = text_encoder( + {"input_ids": inputs["token_ids"], "attention_mask": inputs["padding_mask"]} + ) + return inputs, { + "noise_pred": noise_pred, + "moments": vae_out["moments"], + "image": vae_out["sample"], + "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, + 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): + # 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× + if height is None or width is None: + side = self.config.default_sample_size * self.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 + + def unconditional_ids(self, batch): + # Empty / space negative prompt is encoded by the tokenizer template; + # here build a minimal pad row the task replaces via encode_negative_prompt. + length = self.config.text_config.max_seq_len + row = [self.config.pad_token_id] * length + return ops.convert_to_tensor([row] * batch, 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``. + """ + 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) + + # Gather non-padding tokens per row, drop template prefix, then pad. + # Implemented with numpy for clarity on host; convert back to tensors. + 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), + } + + 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") + 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) + + 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: + # 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) + else: + latents = ops.cast(ops.convert_to_tensor(latents), dtype) + return latents * self.scheduler.init_noise_sigma + + def denoise( + self, latents, embeddings, num_inference_steps, guidance_scale, timesteps=None + ): + """True CFG: separate cond/uncond forwards + norm renormalization.""" + do_cfg = guidance_scale > 1.0 + scheduler = self.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: + 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) + 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)) + latents = scheduler.step(noise_pred, t, latents) + return latents + + def generate( + 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, + ): + 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) + do_cfg = guidance_scale > 1.0 and ( + negative_input_ids is not None + or conditioning.get("negative_prompt") is not None + ) + # Diffusers enables true CFG when a negative prompt is provided. + if guidance_scale > 1.0 and negative_input_ids is None: + # Encode empty/space negative via unconditional_ids path. + neg_ids = self.unconditional_ids(batch) + neg_mask = ops.ones_like(neg_ids) + # Prefer caller-supplied negative when present. + uncond = self.encode_prompt(neg_ids, neg_mask) + do_cfg = True + elif negative_input_ids is not None: + uncond = self.encode_prompt(negative_input_ids, negative_attention_mask) + do_cfg = guidance_scale > 1.0 + else: + uncond = None + do_cfg = False + + # 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( + 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) diff --git a/zeromodels/models/qwen_image/qwen_image_tokenizer.py b/zeromodels/models/qwen_image/qwen_image_tokenizer.py new file mode 100644 index 00000000..ab798e9b --- /dev/null +++ b/zeromodels/models/qwen_image/qwen_image_tokenizer.py @@ -0,0 +1,58 @@ +"""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/qwen_image_vae.py b/zeromodels/models/qwen_image/qwen_image_vae.py new file mode 100644 index 00000000..c79dfbd9 --- /dev/null +++ b/zeromodels/models/qwen_image/qwen_image_vae.py @@ -0,0 +1,857 @@ +import math + +import keras +from keras import layers, ops + +from zeromodels.base.base_attention import fused_attention +from zeromodels.models.stable_diffusion.stable_diffusion_layers import safe_name + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageRMSNorm(layers.Layer): + """RMSNorm used by the Qwen-Image / Wan VAE (Diffusers ``QwenImageRMS_norm``).""" + + def __init__(self, dim, images=True, module_path=None, **kwargs): + if module_path is not None: + kwargs.setdefault("name", safe_name(module_path)) + super().__init__(**kwargs) + self.dim = int(dim) + self.images = bool(images) + self.module_path = module_path + self.scale = self.dim**0.5 + + def build(self, input_shape): + self.gamma = self.add_weight( + shape=(self.dim,), + initializer="ones", + trainable=True, + name="gamma", + ) + self.built = True + + def call(self, x): + dtype = x.dtype + x_f = ops.cast(x, "float32") + norm = ops.sqrt(ops.sum(ops.square(x_f), axis=-1, keepdims=True)) + x_f = x_f / ops.maximum(norm, 1e-12) + x_f = x_f * self.scale * ops.cast(self.gamma, "float32") + return ops.cast(x_f, dtype) + + def compute_output_shape(self, input_shape): + return tuple(input_shape) + + def get_config(self): + config = super().get_config() + config.update( + { + "dim": self.dim, + "images": self.images, + "module_path": self.module_path, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageCausalConv3d(layers.Layer): + """Causal 3D conv with asymmetric temporal pad (Diffusers ``QwenImageCausalConv3d``). + + Diffusers pads ``(W_l, W_r, H_l, H_r, T_l, T_r)`` with ``T_l = 2 * temporal_pad`` + and ``T_r = 0``. Keras uses NDHWC ``Conv3D`` + ``ops.pad``. + + Weight mapping from Diffusers ``Conv3d``: ``(O, I, T, H, W) -> (T, H, W, I, O)``. + """ + + def __init__( + self, + out_channels, + kernel_size=3, + stride=1, + padding=0, + module_path=None, + **kwargs, + ): + if module_path is not None: + kwargs.setdefault("name", safe_name(module_path)) + super().__init__(**kwargs) + self.out_channels = int(out_channels) + self.kernel_size = ( + (kernel_size, kernel_size, kernel_size) + if isinstance(kernel_size, int) + else tuple(kernel_size) + ) + self.stride = ( + (stride, stride, stride) if isinstance(stride, int) else tuple(stride) + ) + pad = ( + (padding, padding, padding) if isinstance(padding, int) else tuple(padding) + ) + self.padding_t, self.padding_h, self.padding_w = pad + self._pad_t_left = 2 * self.padding_t + self._pad_t_right = 0 + self._pad_h = self.padding_h + self._pad_w = self.padding_w + self.module_path = module_path + leaf_name = safe_name(module_path) if module_path else "conv" + self.conv = layers.Conv3D( + self.out_channels, + self.kernel_size, + strides=self.stride, + padding="valid", + data_format="channels_last", + name=leaf_name, + ) + + def _padded_shape(self, input_shape): + b, t, h, w, c = input_shape + t2 = None if t is None else t + self._pad_t_left + self._pad_t_right + h2 = None if h is None else h + 2 * self._pad_h + w2 = None if w is None else w + 2 * self._pad_w + return (b, t2, h2, w2, c) + + def build(self, input_shape): + self.conv.build(self._padded_shape(input_shape)) + self.built = True + + def call(self, x): + # NDHWC pad: [[B], [T], [H], [W], [C]] + x = ops.pad( + x, + ( + (0, 0), + (self._pad_t_left, self._pad_t_right), + (self._pad_h, self._pad_h), + (self._pad_w, self._pad_w), + (0, 0), + ), + ) + return self.conv(x) + + def compute_output_shape(self, input_shape): + b, t, h, w, _ = self._padded_shape(input_shape) + kt, kh, kw = self.kernel_size + st, sh, sw = self.stride + + def _out(size, k, s): + if size is None: + return None + return (size - k) // s + 1 + + return (b, _out(t, kt, st), _out(h, kh, sh), _out(w, kw, sw), self.out_channels) + + def get_config(self): + config = super().get_config() + config.update( + { + "out_channels": self.out_channels, + "kernel_size": self.kernel_size, + "stride": self.stride, + "padding": (self.padding_t, self.padding_h, self.padding_w), + "module_path": self.module_path, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageResample(layers.Layer): + """Spatial / spatio-temporal resample (Diffusers ``QwenImageResample``). + + For ``T=1`` image inference the temporal ``time_conv`` is created (weights) + but not applied, matching Diffusers' cold feat-cache first chunk. + """ + + def __init__(self, dim, mode, module_path, apply_temporal=False, **kwargs): + kwargs.setdefault("name", safe_name(module_path)) + super().__init__(**kwargs) + self.dim = int(dim) + self.mode = mode + self.module_path = module_path + self.apply_temporal = bool(apply_temporal) + + self.spatial_conv = None + self.time_conv = None + self._downsample = mode in ("downsample2d", "downsample3d") + self._upsample = mode in ("upsample2d", "upsample3d") + + if mode in ("upsample2d", "upsample3d"): + self.spatial_conv = layers.Conv2D( + dim // 2, + 3, + padding="same", + data_format="channels_last", + name=safe_name(f"{module_path}.resample.1"), + ) + if mode == "upsample3d": + self.time_conv = QwenImageCausalConv3d( + dim * 2, + kernel_size=(3, 1, 1), + padding=(1, 0, 0), + module_path=f"{module_path}.time_conv", + ) + elif mode in ("downsample2d", "downsample3d"): + self.spatial_conv = layers.Conv2D( + dim, + 3, + strides=2, + padding="valid", + data_format="channels_last", + name=safe_name(f"{module_path}.resample.1"), + ) + if mode == "downsample3d": + self.time_conv = QwenImageCausalConv3d( + dim, + kernel_size=(3, 1, 1), + stride=(2, 1, 1), + padding=(0, 0, 0), + module_path=f"{module_path}.time_conv", + ) + elif mode != "none": + raise ValueError( + f"Unknown resample mode {mode!r}; expected none/upsample2d/" + "upsample3d/downsample2d/downsample3d." + ) + + def build(self, input_shape): + # input_shape: (B, T, H, W, C) + b, t, h, w, c = input_shape + if self.spatial_conv is not None: + if self._downsample: + h_p = None if h is None else h + 1 + w_p = None if w is None else w + 1 + self.spatial_conv.build((b, h_p, w_p, c)) + else: + h_u = None if h is None else h * 2 + w_u = None if w is None else w * 2 + self.spatial_conv.build((b, h_u, w_u, c)) + if self.time_conv is not None: + b, t, h, w, c = input_shape + t_build = t if isinstance(t, int) and t >= 4 else 4 + self.time_conv.build((b, t_build, h, w, c)) + self.built = True + + def call(self, x): + if ( + self.apply_temporal + and self.time_conv is not None + and self.mode == "upsample3d" + ): + b = ops.shape(x)[0] + t = ops.shape(x)[1] + h = ops.shape(x)[2] + w = ops.shape(x)[3] + c = self.dim + x = self.time_conv(x) + x = ops.reshape(x, (b, t, h, w, 2, c)) + x = ops.transpose(x, (0, 1, 4, 2, 3, 5)) + x = ops.reshape(x, (b, t * 2, h, w, c)) + + b = ops.shape(x)[0] + t = ops.shape(x)[1] + h = ops.shape(x)[2] + w = ops.shape(x)[3] + c = ops.shape(x)[4] + x2 = ops.reshape(x, (b * t, h, w, c)) + + if self._upsample: + x2 = ops.image.resize(x2, (h * 2, w * 2), interpolation="nearest") + x2 = self.spatial_conv(x2) + out_c = self.dim // 2 + out_h, out_w = h * 2, w * 2 + elif self._downsample: + x2 = ops.pad(x2, ((0, 0), (0, 1), (0, 1), (0, 0))) + x2 = self.spatial_conv(x2) + out_c = self.dim + out_h = (h + 1 - 3) // 2 + 1 + out_w = (w + 1 - 3) // 2 + 1 + else: + out_c = c + out_h, out_w = h, w + + return ops.reshape(x2, (b, t, out_h, out_w, out_c)) + + def compute_output_shape(self, input_shape): + b, t, h, w, c = input_shape + if self._upsample: + return ( + b, + t, + None if h is None else h * 2, + None if w is None else w * 2, + self.dim // 2, + ) + if self._downsample: + return ( + b, + t, + None if h is None else (h + 1 - 3) // 2 + 1, + None if w is None else (w + 1 - 3) // 2 + 1, + self.dim, + ) + return tuple(input_shape) + + def get_config(self): + config = super().get_config() + config.update( + { + "dim": self.dim, + "mode": self.mode, + "module_path": self.module_path, + "apply_temporal": self.apply_temporal, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageResidualBlock(layers.Layer): + """Residual block with causal 3D convs (Diffusers ``QwenImageResidualBlock``).""" + + def __init__( + self, + in_dim, + out_dim, + module_path, + dropout=0.0, + **kwargs, + ): + kwargs.setdefault("name", safe_name(module_path)) + super().__init__(**kwargs) + self.in_dim = int(in_dim) + self.out_dim = int(out_dim) + self.module_path = module_path + self.dropout_rate = float(dropout) + + self.norm1 = QwenImageRMSNorm( + in_dim, images=False, module_path=f"{module_path}.norm1" + ) + self.conv1 = QwenImageCausalConv3d( + out_dim, 3, padding=1, module_path=f"{module_path}.conv1" + ) + self.norm2 = QwenImageRMSNorm( + out_dim, images=False, module_path=f"{module_path}.norm2" + ) + self.dropout = layers.Dropout( + self.dropout_rate, name=safe_name(f"{module_path}.dropout") + ) + self.conv2 = QwenImageCausalConv3d( + out_dim, 3, padding=1, module_path=f"{module_path}.conv2" + ) + self.conv_shortcut = None + if in_dim != out_dim: + self.conv_shortcut = QwenImageCausalConv3d( + out_dim, 1, padding=0, module_path=f"{module_path}.conv_shortcut" + ) + + def build(self, input_shape): + self.norm1.build(input_shape) + self.conv1.build(input_shape) + mid = list(input_shape) + mid[-1] = self.out_dim + mid = tuple(mid) + self.norm2.build(mid) + self.conv2.build(mid) + if self.conv_shortcut is not None: + self.conv_shortcut.build(input_shape) + self.built = True + + def call(self, x, training=None): + shortcut = x if self.conv_shortcut is None else self.conv_shortcut(x) + h = self.norm1(x) + h = ops.silu(h) + h = self.conv1(h) + h = self.norm2(h) + h = ops.silu(h) + h = self.dropout(h, training=training) + h = self.conv2(h) + return h + shortcut + + def compute_output_shape(self, input_shape): + out = list(input_shape) + out[-1] = self.out_dim + return tuple(out) + + def get_config(self): + config = super().get_config() + config.update( + { + "in_dim": self.in_dim, + "out_dim": self.out_dim, + "module_path": self.module_path, + "dropout": self.dropout_rate, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageAttentionBlock(layers.Layer): + """Single-head causal self-attention over spatial tokens (per time step).""" + + def __init__(self, dim, module_path, **kwargs): + kwargs.setdefault("name", safe_name(module_path)) + super().__init__(**kwargs) + self.dim = int(dim) + self.module_path = module_path + self.norm = QwenImageRMSNorm( + dim, images=True, module_path=f"{module_path}.norm" + ) + self.to_qkv = layers.Conv2D( + dim * 3, + 1, + data_format="channels_last", + name=safe_name(f"{module_path}.to_qkv"), + ) + self.proj = layers.Conv2D( + dim, + 1, + data_format="channels_last", + name=safe_name(f"{module_path}.proj"), + ) + + def build(self, input_shape): + # (B, T, H, W, C) + b, t, h, w, c = input_shape + self.norm.build((b, h, w, c)) + self.to_qkv.build((b, h, w, c)) + self.proj.build((b, h, w, c)) + self.built = True + + def call(self, x): + identity = x + b = ops.shape(x)[0] + t = ops.shape(x)[1] + h = ops.shape(x)[2] + w = ops.shape(x)[3] + c = self.dim + x2 = ops.reshape(x, (b * t, h, w, c)) + x2 = self.norm(x2) + qkv = self.to_qkv(x2) + # (B*T, H, W, 3C) -> (B*T, 1, HW, 3C) -> split to q,k,v (B*T, 1, HW, C) + qkv = ops.reshape(qkv, (b * t, 1, h * w, c * 3)) + q, k, v = ops.split(qkv, 3, axis=-1) + scale = 1.0 / math.sqrt(c) + attn = fused_attention(q, k, v, scale) + attn = ops.reshape(attn, (b * t, h, w, c)) + attn = self.proj(attn) + attn = ops.reshape(attn, (b, t, h, w, c)) + return attn + identity + + def compute_output_shape(self, input_shape): + return tuple(input_shape) + + def get_config(self): + config = super().get_config() + config.update({"dim": self.dim, "module_path": self.module_path}) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageMidBlock(layers.Layer): + """Mid block: resnet -> (attn -> resnet) * num_layers.""" + + def __init__(self, dim, module_path, dropout=0.0, num_layers=1, **kwargs): + kwargs.setdefault("name", safe_name(module_path)) + super().__init__(**kwargs) + self.dim = int(dim) + self.module_path = module_path + self.dropout_rate = float(dropout) + self.num_layers = int(num_layers) + + self.resnets = [ + QwenImageResidualBlock( + dim, dim, f"{module_path}.resnets.0", dropout=dropout + ) + ] + self.attentions = [] + for i in range(num_layers): + self.attentions.append( + QwenImageAttentionBlock(dim, f"{module_path}.attentions.{i}") + ) + self.resnets.append( + QwenImageResidualBlock( + dim, dim, f"{module_path}.resnets.{i + 1}", dropout=dropout + ) + ) + + def build(self, input_shape): + for layer in self.resnets: + layer.build(input_shape) + for layer in self.attentions: + layer.build(input_shape) + self.built = True + + def call(self, x, training=None): + x = self.resnets[0](x, training=training) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + x = attn(x) + x = resnet(x, training=training) + return x + + def compute_output_shape(self, input_shape): + return tuple(input_shape) + + def get_config(self): + config = super().get_config() + config.update( + { + "dim": self.dim, + "module_path": self.module_path, + "dropout": self.dropout_rate, + "num_layers": self.num_layers, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageEncoder3d(layers.Layer): + """3D VAE encoder (Diffusers ``QwenImageEncoder3d``).""" + + def __init__( + self, + dim=96, + z_dim=32, + dim_mult=(1, 2, 4, 4), + num_res_blocks=2, + attn_scales=(), + temperal_downsample=(False, True, True), + dropout=0.0, + input_channels=3, + module_path="encoder", + apply_temporal=False, + **kwargs, + ): + kwargs.setdefault("name", safe_name(module_path)) + super().__init__(**kwargs) + self.dim = int(dim) + self.z_dim = int(z_dim) + self.dim_mult = tuple(dim_mult) + self.num_res_blocks = int(num_res_blocks) + self.attn_scales = tuple(attn_scales) + self.temperal_downsample = tuple(temperal_downsample) + self.dropout_rate = float(dropout) + self.input_channels = int(input_channels) + self.module_path = module_path + self.apply_temporal = bool(apply_temporal) + + dims = [dim * u for u in [1] + list(self.dim_mult)] + scale = 1.0 + + self.conv_in = QwenImageCausalConv3d( + dims[0], 3, padding=1, module_path=f"{module_path}.conv_in" + ) + self.down_blocks = [] + idx = 0 + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + for _ in range(self.num_res_blocks): + self.down_blocks.append( + QwenImageResidualBlock( + in_dim, + out_dim, + f"{module_path}.down_blocks.{idx}", + dropout=dropout, + ) + ) + idx += 1 + if scale in self.attn_scales: + self.down_blocks.append( + QwenImageAttentionBlock( + out_dim, f"{module_path}.down_blocks.{idx}" + ) + ) + idx += 1 + in_dim = out_dim + if i != len(self.dim_mult) - 1: + mode = "downsample3d" if self.temperal_downsample[i] else "downsample2d" + self.down_blocks.append( + QwenImageResample( + out_dim, + mode, + f"{module_path}.down_blocks.{idx}", + apply_temporal=self.apply_temporal, + ) + ) + idx += 1 + scale /= 2.0 + + self.mid_block = QwenImageMidBlock( + dims[-1], f"{module_path}.mid_block", dropout=dropout, num_layers=1 + ) + self.norm_out = QwenImageRMSNorm( + dims[-1], images=False, module_path=f"{module_path}.norm_out" + ) + self.conv_out = QwenImageCausalConv3d( + z_dim, 3, padding=1, module_path=f"{module_path}.conv_out" + ) + self._out_dim = dims[-1] + + def build(self, input_shape): + self.conv_in.build(input_shape) + x_shape = list(input_shape) + x_shape[-1] = self.dim + x_shape = tuple(x_shape) + shape = x_shape + for layer in self.down_blocks: + layer.build(shape) + shape = layer.compute_output_shape(shape) + self.mid_block.build(shape) + self.norm_out.build(shape) + self.conv_out.build(shape) + self.built = True + + def call(self, x, training=None): + x = self.conv_in(x) + for layer in self.down_blocks: + if isinstance(layer, QwenImageResidualBlock): + x = layer(x, training=training) + else: + x = layer(x) + x = self.mid_block(x, training=training) + x = self.norm_out(x) + x = ops.silu(x) + x = self.conv_out(x) + return x + + def compute_output_shape(self, input_shape): + shape = list(input_shape) + shape[-1] = self.dim + shape = tuple(shape) + for layer in self.down_blocks: + shape = layer.compute_output_shape(shape) + out = list(shape) + out[-1] = self.z_dim + b, t, h, w, _ = input_shape + factor = 2 ** (len(self.dim_mult) - 1) + return ( + b, + t, + None if h is None else h // factor, + None if w is None else w // factor, + self.z_dim, + ) + + def get_config(self): + config = super().get_config() + config.update( + { + "dim": self.dim, + "z_dim": self.z_dim, + "dim_mult": self.dim_mult, + "num_res_blocks": self.num_res_blocks, + "attn_scales": self.attn_scales, + "temperal_downsample": self.temperal_downsample, + "dropout": self.dropout_rate, + "input_channels": self.input_channels, + "module_path": self.module_path, + "apply_temporal": self.apply_temporal, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageUpBlock(layers.Layer): + """Decoder up block (Diffusers ``QwenImageUpBlock``).""" + + def __init__( + self, + in_dim, + out_dim, + num_res_blocks, + module_path, + dropout=0.0, + upsample_mode=None, + apply_temporal=False, + **kwargs, + ): + kwargs.setdefault("name", safe_name(module_path)) + super().__init__(**kwargs) + self.in_dim = int(in_dim) + self.out_dim = int(out_dim) + self.num_res_blocks = int(num_res_blocks) + self.module_path = module_path + self.dropout_rate = float(dropout) + self.upsample_mode = upsample_mode + self.apply_temporal = bool(apply_temporal) + + self.resnets = [] + current = in_dim + for i in range(num_res_blocks + 1): + self.resnets.append( + QwenImageResidualBlock( + current, + out_dim, + f"{module_path}.resnets.{i}", + dropout=dropout, + ) + ) + current = out_dim + + self.upsampler = None + if upsample_mode is not None: + self.upsampler = QwenImageResample( + out_dim, + upsample_mode, + f"{module_path}.upsamplers.0", + apply_temporal=apply_temporal, + ) + + def build(self, input_shape): + shape = input_shape + for layer in self.resnets: + layer.build(shape) + shape = layer.compute_output_shape(shape) + if self.upsampler is not None: + self.upsampler.build(shape) + self.built = True + + def call(self, x, training=None): + for resnet in self.resnets: + x = resnet(x, training=training) + if self.upsampler is not None: + x = self.upsampler(x) + return x + + def compute_output_shape(self, input_shape): + shape = list(input_shape) + shape[-1] = self.out_dim + shape = tuple(shape) + if self.upsampler is not None: + shape = self.upsampler.compute_output_shape(shape) + return shape + + def get_config(self): + config = super().get_config() + config.update( + { + "in_dim": self.in_dim, + "out_dim": self.out_dim, + "num_res_blocks": self.num_res_blocks, + "module_path": self.module_path, + "dropout": self.dropout_rate, + "upsample_mode": self.upsample_mode, + "apply_temporal": self.apply_temporal, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class QwenImageDecoder3d(layers.Layer): + """3D VAE decoder (Diffusers ``QwenImageDecoder3d``).""" + + def __init__( + self, + dim=96, + z_dim=16, + dim_mult=(1, 2, 4, 4), + num_res_blocks=2, + attn_scales=(), + temperal_upsample=(True, True, False), + dropout=0.0, + input_channels=3, + module_path="decoder", + apply_temporal=False, + **kwargs, + ): + kwargs.setdefault("name", safe_name(module_path)) + super().__init__(**kwargs) + self.dim = int(dim) + self.z_dim = int(z_dim) + self.dim_mult = tuple(dim_mult) + self.num_res_blocks = int(num_res_blocks) + self.attn_scales = tuple(attn_scales) + self.temperal_upsample = tuple(temperal_upsample) + self.dropout_rate = float(dropout) + self.input_channels = int(input_channels) + self.module_path = module_path + self.apply_temporal = bool(apply_temporal) + + dims = [dim * u for u in [self.dim_mult[-1]] + list(self.dim_mult[::-1])] + + self.conv_in = QwenImageCausalConv3d( + dims[0], 3, padding=1, module_path=f"{module_path}.conv_in" + ) + self.mid_block = QwenImageMidBlock( + dims[0], f"{module_path}.mid_block", dropout=dropout, num_layers=1 + ) + self.up_blocks = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + if i > 0: + in_dim = in_dim // 2 + upsample_mode = None + if i != len(self.dim_mult) - 1: + upsample_mode = ( + "upsample3d" if self.temperal_upsample[i] else "upsample2d" + ) + self.up_blocks.append( + QwenImageUpBlock( + in_dim=in_dim, + out_dim=out_dim, + num_res_blocks=num_res_blocks, + module_path=f"{module_path}.up_blocks.{i}", + dropout=dropout, + upsample_mode=upsample_mode, + apply_temporal=apply_temporal, + ) + ) + self.norm_out = QwenImageRMSNorm( + dims[-1], images=False, module_path=f"{module_path}.norm_out" + ) + self.conv_out = QwenImageCausalConv3d( + input_channels, 3, padding=1, module_path=f"{module_path}.conv_out" + ) + self._out_channels = dims[-1] + + def build(self, input_shape): + self.conv_in.build(input_shape) + shape = list(input_shape) + shape[-1] = self.dim * self.dim_mult[-1] + shape = tuple(shape) + self.mid_block.build(shape) + for block in self.up_blocks: + block.build(shape) + shape = block.compute_output_shape(shape) + self.norm_out.build(shape) + self.conv_out.build(shape) + self.built = True + + def call(self, x, training=None): + x = self.conv_in(x) + x = self.mid_block(x, training=training) + for block in self.up_blocks: + x = block(x, training=training) + x = self.norm_out(x) + x = ops.silu(x) + x = self.conv_out(x) + return x + + def compute_output_shape(self, input_shape): + b, t, h, w, _ = input_shape + factor = 2 ** (len(self.dim_mult) - 1) + return ( + b, + t, + None if h is None else h * factor, + None if w is None else w * factor, + self.input_channels, + ) + + def get_config(self): + config = super().get_config() + config.update( + { + "dim": self.dim, + "z_dim": self.z_dim, + "dim_mult": self.dim_mult, + "num_res_blocks": self.num_res_blocks, + "attn_scales": self.attn_scales, + "temperal_upsample": self.temperal_upsample, + "dropout": self.dropout_rate, + "input_channels": self.input_channels, + "module_path": self.module_path, + "apply_temporal": self.apply_temporal, + } + ) + return config