From fa1785fd92ec0c16d6ebebf76ceefe21dff2335d Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Tue, 22 Sep 2026 16:09:45 -0700 Subject: [PATCH 1/8] Add Qwen Image Model --- docs/qwen_image.md | 61 + tests/base/model_test_registry.py | 55 + tests/fixtures/dummy_inputs.py | 28 + website/mkdocs.yml | 1 + zeromodels/auto/auto_mapping_names.py | 4 + zeromodels/base/base_scheduler.py | 105 +- zeromodels/models/__init__.py | 1 + zeromodels/models/qwen_image/__init__.py | 24 + .../convert_qwen_image_diffusers_to_keras.py | 387 ++++++ .../models/qwen_image/qwen_image_config.py | 109 ++ .../models/qwen_image/qwen_image_layers.py | 760 +++++++++++ .../models/qwen_image/qwen_image_model.py | 628 +++++++++ .../models/qwen_image/qwen_image_tokenizer.py | 62 + .../models/qwen_image/qwen_image_vae.py | 1194 +++++++++++++++++ 14 files changed, 3407 insertions(+), 12 deletions(-) create mode 100644 docs/qwen_image.md create mode 100644 zeromodels/models/qwen_image/__init__.py create mode 100644 zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py create mode 100644 zeromodels/models/qwen_image/qwen_image_config.py create mode 100644 zeromodels/models/qwen_image/qwen_image_layers.py create mode 100644 zeromodels/models/qwen_image/qwen_image_model.py create mode 100644 zeromodels/models/qwen_image/qwen_image_tokenizer.py create mode 100644 zeromodels/models/qwen_image/qwen_image_vae.py diff --git a/docs/qwen_image.md b/docs/qwen_image.md new file mode 100644 index 00000000..7b49adb6 --- /dev/null +++ b/docs/qwen_image.md @@ -0,0 +1,61 @@ +# Qwen-Image + +
+Weights: pretrained Keras weights will be hosted on Hugging Face under +zeromodels/qwen-image +(each repo carries zm_config.json + model.weights.h5 + +tokenizer.json). Load with from_weights("zeromodels/qwen-image"). +Conversion from Diffusers is offline only — on-the-fly hf: is not supported. +
+ +Qwen-Image text-to-image, ported to pure Keras 3 from +[Diffusers `QwenImagePipeline`](https://huggingface.co/docs/diffusers/api/pipelines/qwenimage) +(`/Qwen/Qwen-Image`). Latent flow-matching with: + +- **Denoiser**: 60-layer double-stream MMDiT (`QwenImageTransformer2DModel`) over + packed 2×2 latent patches (64-d tokens) with MS-RoPE +- **VAE**: Wan-derived `AutoencoderKLQwenImage` (16 latent channels, 8× spatial) +- **Text encoder**: Qwen2.5-VL-7B Instruct text tower (ChatML prompt template) +- **Scheduler**: `FlowMatchEulerDiscreteScheduler` with dynamic resolution shifting + +`QwenImageModel` is the hosted container (transformer + VAE + text encoder). +`QwenImageTextToImage` adds `generate` via `BaseDiffusion`, with **true CFG** +(separate cond/uncond forwards and prediction-norm renormalization, Diffusers +`true_cfg_scale`). + +## Status + +| Piece | Status | +|---|---| +| Configs / tokenizer / auto registry | Done | +| Transformer layers + model | Done (graph) | +| VAE | Done (T=1 image path) | +| Text-to-image `generate` + packed latents | Done | +| Flow-match dynamic shifting | Done | +| Diffusers → Keras weight transfer | In progress (`convert_qwen_image_diffusers_to_keras.py`) | +| Hosted `zeromodels/qwen-image` weights | Pending conversion | + +## API sketch + +```python +from zeromodels.models.qwen_image import QwenImageTextToImage, QwenImageTokenizer + +model = QwenImageTextToImage.from_weights("zeromodels/qwen-image") +tok = QwenImageTokenizer.from_weights("zeromodels/qwen-image") +image = model.generate( + **tok("a coffee shop entrance with a chalkboard sign"), + height=1024, + width=1024, + num_inference_steps=50, + guidance_scale=4.0, +) +``` + +## Variants + +| Variant | Hub (planned) | Source | +|---|---|---| +| `qwen-image` | `zeromodels/qwen-image` | [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) | + +Paper / model card: [Qwen-Image](https://huggingface.co/Qwen/Qwen-Image). +License: Apache-2.0. 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/dummy_inputs.py b/tests/fixtures/dummy_inputs.py index 3d926da8..b08e3312 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=16, + 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/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..8406db89 --- /dev/null +++ b/zeromodels/models/qwen_image/__init__.py @@ -0,0 +1,24 @@ +from .qwen_image_config import ( + QwenImageConfig, + QwenImageTextConfig, + QwenImageTransformerConfig, +) +from .qwen_image_model import ( + QwenImageModel, + QwenImageTextToImage, + QwenImageTransformer2DModel, +) +from .qwen_image_tokenizer import QwenImageTokenizer +from .qwen_image_vae import AutoencoderKLQwenImage, QwenImageVAEConfig + +__all__ = [ + "AutoencoderKLQwenImage", + "QwenImageConfig", + "QwenImageModel", + "QwenImageTextConfig", + "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..1f3457ab --- /dev/null +++ b/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py @@ -0,0 +1,387 @@ +"""Offline Diffusers ``Qwen/Qwen-Image`` → ZeroModels Keras weight conversion. + +Converts the transformer, VAE, and Qwen2.5-VL text encoder into a hosted +``zeromodels/qwen-image`` layout (``zm_config.json`` + sharded weights). + +Usage:: + + python -m zeromodels.models.qwen_image.convert_qwen_image_diffusers_to_keras + +Env: + ZM_OUT_DIR output directory (default ``./qwen_image_weights``) + HF_TOKEN optional Hub token + ZM_DTYPE ``float16`` / ``bfloat16`` / ``float32`` (default ``float16``) +""" + +from __future__ import annotations + +import gc +import json +import os + +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, +) +from zeromodels.models.qwen2_5_vl.convert_qwen2_5_vl_hf_to_keras import ( + transfer_qwen2_5_vl_weights, +) +from zeromodels.models.stable_diffusion.convert_stable_diffusion_diffusers_to_keras import ( + WEIGHT_NAME_MAPPING, +) + +QWEN_IMAGE_SOURCES = { + "qwen-image": "Qwen/Qwen-Image", +} + +# VAE Diffusers RMSNorm keeps the leaf name ``gamma`` (not ``weight``). +VAE_WEIGHT_NAME_MAPPING = { + k: v for k, v in WEIGHT_NAME_MAPPING.items() if "/gamma" not in k +} +VAE_WEIGHT_NAME_MAPPING["/"] = "." +VAE_WEIGHT_NAME_MAPPING["__"] = "." + + +def config_from_diffusers(repo, token=None): + """Build a :class:`QwenImageConfig` from a Diffusers Qwen-Image repo.""" + 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 + from zeromodels.models.qwen_image.qwen_image_model import ( + QwenImageTransformer2DModel, + ) + from zeromodels.models.qwen_image.qwen_image_vae import QwenImageVAEConfig + + 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 _keras_to_torch_key(keras_weight, mapping): + key = "/".join(keras_weight.path.split("/")[-2:]) + for old, new in mapping.items(): + key = key.replace(old, new) + return key + + +def _prepare_torch_array(keras_weight, torch_weight): + """Transpose Conv3D / squeeze VAE RMSNorm gamma to match Keras shapes.""" + arr = np.asarray(torch_weight) + kshape = tuple(keras_weight.shape) + if len(kshape) == 5 and arr.ndim == 5: + # Diffusers Conv3d NCDHW ``(O, I, T, H, W)`` → Keras NDHWC ``(T, H, W, I, O)``. + arr = np.transpose(arr, (2, 3, 4, 1, 0)) + elif arr.ndim > 1 and len(kshape) == 1 and int(np.prod(arr.shape)) == kshape[0]: + arr = arr.reshape(kshape) + return arr + + +def _load_safetensors_state(repo, subfolder, index_name, token=None): + """Lazy shard map: ``{key: (shard_path, key)}`` resolved on demand via safe_open.""" + from huggingface_hub import hf_hub_download + from safetensors import safe_open + + 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 _ShardDict(dict): + def __contains__(self, key): + return key in weight_map + + def __getitem__(self, key): + if key not in weight_map: + raise KeyError(key) + path = shard_paths[weight_map[key]] + with safe_open(path, 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) + + return _ShardDict() + + +def _load_single_safetensors(repo, subfolder, filename, token=None): + from huggingface_hub import hf_hub_download + from safetensors import safe_open + + 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) + return state + + +def transfer_component(component, state, mapping, desc, ignore=()): + """Copy ``state`` tensors into a Keras sub-model using ``mapping``.""" + consumed = set(ignore) + trainable, non_trainable = split_model_weights(component) + for keras_weight, _ in tqdm(trainable + non_trainable, desc=desc): + if any(part in keras_weight.path for part in ("pos_freqs", "neg_freqs")): + continue + key = _keras_to_torch_key(keras_weight, mapping) + if key in ignore: + continue + if key not in state: + raise WeightMappingError(keras_weight.path, key) + consumed.add(key) + raw = state[key] + torch_weight = _prepare_torch_array(keras_weight, raw) + if tuple(keras_weight.shape) != tuple(np.asarray(torch_weight).shape): + # Dense / Conv2D still need the shared transpose inside transfer_weights; + # only enforce exact match after _prepare (Conv3D / squeezed gamma). + if len(keras_weight.shape) == 5 or ( + len(keras_weight.shape) == 1 + and np.asarray(raw).ndim > 1 + ): + raise WeightShapeMismatchError( + keras_weight.path, + keras_weight.shape, + key, + torch_weight.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, torch_weight) + unused = sorted(set(state) - consumed) + # RoPE tables / unused buffers on the torch side are fine to ignore when empty. + if unused: + raise ValueError( + f"{type(component).__name__}: {len(unused)} checkpoint tensors " + f"unused, e.g. {unused[:5]}." + ) + + +def transfer_qwen_image(repo, token=None, dtype="float16", build_sample_size=16): + """Convert Diffusers Qwen-Image weights into a :class:`QwenImageModel`. + + Builds the Keras container at a small ``transformer_sample_size`` (weights are + resolution-independent for the DiT) to keep the functional graph smaller, + transfers transformer / VAE / text-encoder weights, and returns ``(model, config)``. + """ + from zeromodels.base.base_mixin import build_dtype_scope + from zeromodels.models.qwen_image.qwen_image_model import QwenImageModel + + config = config_from_diffusers(repo, token=token) + flat = config.constructor_kwargs() + # Smaller spatial graph for conversion; Dense DiT weights do not depend on it. + flat["transformer_sample_size"] = build_sample_size + flat["vae_sample_size"] = max(build_sample_size * 8, 64) + + print(f"[1/4] Building Keras QwenImageModel (dtype={dtype})…", flush=True) + import sys + import time + + t0 = time.time() + with build_dtype_scope(dtype), zeros_init(): + print(" building components…", flush=True) + model = QwenImageModel(**flat) + print( + f" build done in {time.time() - t0:.1f}s " + f"({len(model.weights)} weights)", + flush=True, + ) + sys.stdout.flush() + + print("[2/4] Transferring transformer…", flush=True) + transformer_state = _load_safetensors_state( + repo, + "transformer", + "diffusion_pytorch_model.safetensors.index.json", + token=token, + ) + transfer_component( + model.transformer, + transformer_state, + WEIGHT_NAME_MAPPING, + desc="transformer", + ) + del transformer_state + gc.collect() + + print("[3/4] Transferring VAE…") + vae_state = _load_single_safetensors( + repo, "vae", "diffusion_pytorch_model.safetensors", token=token + ) + transfer_component( + model.vae, + vae_state, + VAE_WEIGHT_NAME_MAPPING, + desc="vae", + ) + del vae_state + gc.collect() + + print("[4/4] Transferring text encoder (Qwen2.5-VL text tower)…") + text_state = _load_safetensors_state( + repo, "text_encoder", "model.safetensors.index.json", token=token + ) + # Drop vision + LM head: our text-only tower does not carry them. + text_only = { + k: text_state[k] + for k in text_state.keys() + if k.startswith("model.") and not k.startswith("model.visual.") + } + # Materialize into a plain dict for the shared VL transfer helper. + text_np = {k: np.asarray(text_only[k]) for k in tqdm(text_only, desc="load text")} + del text_state, text_only + gc.collect() + transfer_qwen2_5_vl_weights(model.text_encoder, text_np) + del text_np + gc.collect() + + return model, config + + +def save_converted(model, config, out_dir, variant="qwen-image", max_shard_gb=5.0): + """Write ``zm_config.json``, sharded weights, and copy the tokenizer files.""" + import shutil + + from huggingface_hub import hf_hub_download + + os.makedirs(out_dir, exist_ok=True) + stem = os.path.join(out_dir, variant.replace("-", "_")) + + config_path = os.path.join(out_dir, "zm_config.json") + with open(config_path, "w", encoding="utf-8") as f: + json.dump(config.to_dict(), f, indent=2) + print(f" wrote {config_path}") + + n_bytes = sum(int(np.prod(w.shape)) * 2 for w in model.weights) # ~fp16 + 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 weights -> {out} (~{n_bytes / 1024**3:.2f} GB fp16-est)") + + source = QWEN_IMAGE_SOURCES[variant] + for name in ( + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "vocab.json", + "merges.txt", + "added_tokens.json", + ): + try: + src = hf_hub_download(source, name, subfolder="tokenizer") + shutil.copy2(src, os.path.join(out_dir, name if name != "tokenizer.json" else "tokenizer.json")) + except Exception as exc: # noqa: BLE001 — best-effort tokenizer copy + print(f" skip tokenizer file {name}: {exc}") + + +if __name__ == "__main__": + import keras + + OUT_DIR = os.environ.get( + "ZM_OUT_DIR", + os.path.join(os.path.dirname(__file__), "..", "..", "..", "qwen_image_weights"), + ) + OUT_DIR = os.path.abspath(OUT_DIR) + os.makedirs(OUT_DIR, exist_ok=True) + token = os.environ.get("HF_TOKEN") + dtype = os.environ.get("ZM_DTYPE", "float16") + selected = [v for v in os.environ.get("ZM_VARIANTS", "qwen-image").split(",") if v] + + for variant in selected: + source = QWEN_IMAGE_SOURCES[variant] + print(f"\n{'=' * 60}\nConverting: {variant} <- {source}\n{'=' * 60}") + model, config = transfer_qwen_image(source, token=token, dtype=dtype) + save_converted(model, config, OUT_DIR, variant=variant) + del model + keras.backend.clear_session() + gc.collect() + print(f"Done: {variant} -> {OUT_DIR}") 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..2653fe1c --- /dev/null +++ b/zeromodels/models/qwen_image/qwen_image_config.py @@ -0,0 +1,109 @@ +"""Typed configs for Qwen-Image (Diffusers ``QwenImagePipeline``).""" + +from zeromodels.base import BaseConfig +from zeromodels.models.qwen2_5_vl.qwen2_5_vl_config import Qwen2_5VLTextConfig +from zeromodels.models.qwen_image.qwen_image_vae import QwenImageVAEConfig + + +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..de55d214 --- /dev/null +++ b/zeromodels/models/qwen_image/qwen_image_layers.py @@ -0,0 +1,760 @@ +"""Keras 3 layers for Qwen-Image's double-stream DiT. + +Ports the building blocks of Diffusers ``transformer_qwenimage.py`` (RMSNorm, +timestep projection, 3D RoPE, joint text/image attention, and the dual-stream +transformer block). The full ``QwenImageTransformer2DModel`` lives elsewhere. + +Module paths follow Diffusers so weight conversion can map leaves such as +``transformer_blocks.0.attn.to_q``, ``img_mod.1``, ``img_mlp.net.0.proj``, and +``time_text_embed.timestep_embedder.linear_1``. +""" + +from __future__ import annotations + +import math + +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 qwen_approximate_gelu(x): + """Tanh-approximate GELU (Diffusers ``gelu-approximate`` / BERT). + + ``0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))``. Prefer + :class:`StableDiffusion3GELUFeedForward` (``ops.gelu(..., approximate=True)``) + for the block MLPs; this helper is for call sites that need the formula + explicitly. + """ + return 0.5 * x * (1.0 + ops.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * ops.power(x, 3)))) + + +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) + # Broadcast over batch and heads: [S, D] -> [1, S, 1, D] + 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): + # Sinusoidal projection is fixed-width; MLP builds from (B, time_freq_dim). + 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): + # Diffusers Timesteps(scale=1000): emb = scale * (t[:, None] * freqs) + 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 # complex half-dim == angle width + self._axis_halves = [d // 2 for d in self.axes_dim] + + pos_index = np.arange(ROPE_MAX_INDEX) + # Diffusers: arange(4096).flip(0) * -1 - 1 -> [-4096, ..., -1] + 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, + ) + 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, + ) + self.built = True + + def _split_axes(self, freqs): + # Diffusers ``tensor.split(sizes, dim)``; Keras ``ops.split`` takes indices. + 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.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) + + # Joint sequence order: [text, image] (Diffusers Qwen; not SD3). + 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) + + # fused_attention expects [B, H, S, D] + 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 + ) + # Additive mask: keep=0, drop=MASK_NEG (Diffusers bool True=keep). + 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.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 + + # Diffusers Sequential(SiLU, Linear) -> weight leaf ``img_mod.1`` / ``txt_mod.1``. + 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: + # Explicit ``build((img, txt, temb))`` from a parent layer. + 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): + """``x * (1 + scale) + shift``, returning ``(modulated, gate)``.""" + 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) + ) + + # Diffusers clips fp16 dual-stream outputs to the fp16 finite range. + if str(encoder_hidden_states.dtype).endswith("float16"): + encoder_hidden_states = ops.clip( + encoder_hidden_states, -65504.0, 65504.0 + ) + if str(hidden_states.dtype).endswith("float16"): + hidden_states = ops.clip(hidden_states, -65504.0, 65504.0) + + return encoder_hidden_states, hidden_states + + def compute_output_shape(self, input_shape): + # Multi-input builds pass ``[img_shape, txt_shape, temb_shape]``; a lone + # tensor shape is ``(batch, seq, dim)``. + 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 + + +# Final-norm AdaLN: Diffusers ``AdaLayerNormContinuous`` uses (scale, shift) +# order with ``num_chunks=2`` — same as StableDiffusion3AdaLayerNorm. +QwenImageAdaLayerNormContinuous = StableDiffusion3AdaLayerNorm + + +__all__ = [ + "qwen_approximate_gelu", + "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..a7754a64 --- /dev/null +++ b/zeromodels/models/qwen_image/qwen_image_model.py @@ -0,0 +1,628 @@ +"""Qwen-Image models: transformer, container, and text-to-image task.""" + +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 ( + QwenImageConfig, + QwenImageTransformerConfig, +) +from zeromodels.models.qwen_image.qwen_image_layers import ( + QwenImageAdaLayerNormContinuous, + QwenImageEmbedRope, + QwenImageRMSNorm, + QwenImageTimestepProjEmbeddings, + QwenImageTransformerBlock, +) +from zeromodels.models.qwen_image.qwen_image_vae import AutoencoderKLQwenImage +from zeromodels.models.stable_diffusion.stable_diffusion_layers import safe_name + +QWEN_IMAGE_HUB_SIBLINGS = ("qwen-image",) + +PROMPT_TEMPLATE = ( + "<|im_start|>system\nDescribe the image by detailing the color, shape, size, " + "texture, quantity, text, spatial relationships of the objects and " + "background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" +) + + +def calculate_shift( + image_seq_len, + base_seq_len=256, + max_seq_len=4096, + base_shift=0.5, + max_shift=1.15, +): + """Resolution-dependent flow-match shift (Diffusers ``calculate_shift``).""" + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + return image_seq_len * m + b + + +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 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 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``. + 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, t = config.transformer_config, config.vae_config, config.text_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 = Qwen2_5VLModel( + vocab_size=t.vocab_size, + embed_dim=t.embed_dim, + mlp_dim=t.mlp_dim, + num_layers=t.num_layers, + num_heads=t.num_heads, + num_kv_heads=t.num_kv_heads, + norm_eps=t.norm_eps, + rope_theta=t.rope_theta, + mrope_section=t.mrope_section, + tie_embeddings=t.tie_embeddings, + build_vision=False, + name="text_encoder", + ) + 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 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 + + # Dynamic flow-match timesteps (mu from packed sequence length). + h, w = self._latent_side(height, width) + image_seq_len = (h // 2) * (w // 2) + sched_cfg = getattr(self.scheduler, "config_dict", None) or {} + mu = calculate_shift( + image_seq_len, + sched_cfg.get("base_image_seq_len", 256), + sched_cfg.get("max_image_seq_len", 4096), + sched_cfg.get("base_shift", 0.5), + sched_cfg.get("max_shift", 0.9), + ) + 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..c73f5dfb --- /dev/null +++ b/zeromodels/models/qwen_image/qwen_image_tokenizer.py @@ -0,0 +1,62 @@ +"""Qwen-Image tokenizer: Qwen2 ChatML template for prompt encoding.""" + +import keras + +from zeromodels.models.qwen2.qwen2_tokenizer import Qwen2Tokenizer + +# Diffusers ``QwenImagePipeline.prompt_template_encode`` +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, + ): + # Diffusers pads to tokenizer_max_length + drop_idx (= 1024 + 34). + 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] + # Pad/truncate to max_seq_len + template prefix budget (Diffusers). + 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) + # Cap length for the static text-encoder graph when shorter. + 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..987b9039 --- /dev/null +++ b/zeromodels/models/qwen_image/qwen_image_vae.py @@ -0,0 +1,1194 @@ +"""Keras 3 port of Diffusers ``AutoencoderKLQwenImage`` (Qwen-Image / Wan VAE). + +Internal activations use channels-last 5D ``(B, T, H, W, C)`` (NDHWC) so +``Conv3D`` is portable across Keras backends. Diffusers stores the same stacks +as NCDHW ``Conv3d``; converters must transpose kernels: + +* Conv3D: ``(O, I, T, H, W)`` -> ``(T, H, W, I, O)`` +* Conv2D (resample / attention): ``(O, I, H, W)`` -> ``(H, W, I, O)`` +* RMSNorm ``gamma``: squeeze Diffusers ``(C, 1, 1[, 1])`` -> ``(C,)`` + +v1 targets single-frame T2I (``T=1``): feat-cache video streaming is a no-op and +temporal ``time_conv`` paths inside ``Resample`` are skipped, matching Diffusers' +first-chunk behaviour when the cache is cold. Tiling / slicing are omitted. +Module paths mirror Diffusers for weight conversion. +""" + +from __future__ import annotations + +import math + +import keras +from keras import layers, ops + +from zeromodels.base import BaseModel +from zeromodels.base.base_attention import fused_attention +from zeromodels.base.base_config import BaseConfig +from zeromodels.models.stable_diffusion.stable_diffusion_layers import safe_name + +# --------------------------------------------------------------------------- +# Defaults (Qwen/Qwen-Image VAE config.json) +# --------------------------------------------------------------------------- + +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 + + +# --------------------------------------------------------------------------- +# Layout helpers +# --------------------------------------------------------------------------- + + +def _as_tuple3(value): + if isinstance(value, int): + return (value, value, value) + value = tuple(value) + if len(value) != 3: + raise ValueError(f"Expected int or length-3 tuple, got {value!r}") + return value + + +def ncdhw_to_ndhwc(x): + """``(B, C, T, H, W)`` -> ``(B, T, H, W, C)``.""" + return ops.transpose(x, (0, 2, 3, 4, 1)) + + +def ndhwc_to_ncdhw(x): + """``(B, T, H, W, C)`` -> ``(B, C, T, H, W)``.""" + return ops.transpose(x, (0, 4, 1, 2, 3)) + + +# --------------------------------------------------------------------------- +# Layers +# --------------------------------------------------------------------------- + + +@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): + # Normalize over channels (last axis for NDHWC / NHWC). + dtype = x.dtype + x_f = ops.cast(x, "float32") + variance = ops.mean(ops.square(x_f), axis=-1, keepdims=True) + x_f = x_f * ops.rsqrt(variance + 1e-6) + 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 = _as_tuple3(kernel_size) + self.stride = _as_tuple3(stride) + pad = _as_tuple3(padding) + self.padding_t, self.padding_h, self.padding_w = pad + # Causal: double temporal left pad, zero right 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 Conv3D keeps the Diffusers module path name so conversion's + # last-two-segment mapping yields ``encoder.conv_in.weight`` etc. + # (path ``.../encoder__conv_in/encoder__conv_in/kernel``). + 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: + # ZeroPad2d((0,1,0,1)) -> H+1, W+1 + 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: + # Downsample3d time_conv is stride-2 with k_t=3 and no pad; T=1 cannot + # run it. Build with a synthetic T so kernels exist for conversion; the + # T=1 image path never calls time_conv (cold-cache Diffusers behaviour). + 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): + # Optional temporal branch (video streaming); skipped for T=1 image path + # (matches Diffusers' cold feat-cache first chunk, which skips time_conv). + 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) + # (B, T, H, W, 2C) -> (B, T*2, H, W, C) interleaved like Diffusers + 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] + # Merge batch and time for 2D ops: (B*T, H, W, C) + 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) + # Walk blocks updating spatial dims approximately for build. + 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 + # Spatial / temporal sizes: 3 spatial downsamples -> /8; T unchanged (T=1 path). + 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 + + +# --------------------------------------------------------------------------- +# Model +# --------------------------------------------------------------------------- + + +@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" + ) + + # Functional graph: HWC image + HWC-ish latent (T squeezed to channels-last 4D + # with an explicit time axis of 1 in the 5D path). + 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) # (B, 1, H, W, C) + moments_5d = quant_conv(encoder(image_5d)) + moments = ops.squeeze(moments_5d, axis=1) # (B, h, w, 2*z) + + 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 + + # -- public encode / decode ------------------------------------------------ + + 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: + # (B, H, W, C) channels-last image or latent + return ops.expand_dims(x, axis=1) + if static_ndim == 5: + # Detect NCDHW (Diffusers): channel axis at 1 matches input_channels / z_dim. + 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 ncdhw_to_ndhwc(x) + 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 ndhwc_to_ncdhw(z) + 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 ndhwc_to_ncdhw(x) + 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, + ) From 024fec4656541e83dd531a1e4696e16a244e13b8 Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Tue, 22 Sep 2026 17:07:06 -0700 Subject: [PATCH 2/8] fix --- docs/qwen_image.md | 11 +- tests/base/model_test_registry.py | 1 - tests/fixtures/dummy_inputs.py | 3 - zeromodels/models/qwen_image/__init__.py | 2 + .../convert_qwen_image_diffusers_to_keras.py | 228 ++++++++++++------ .../models/qwen_image/qwen_image_model.py | 113 +++++---- 6 files changed, 234 insertions(+), 124 deletions(-) diff --git a/docs/qwen_image.md b/docs/qwen_image.md index 7b49adb6..9ac25880 100644 --- a/docs/qwen_image.md +++ b/docs/qwen_image.md @@ -18,7 +18,10 @@ Qwen-Image text-to-image, ported to pure Keras 3 from - **Text encoder**: Qwen2.5-VL-7B Instruct text tower (ChatML prompt template) - **Scheduler**: `FlowMatchEulerDiscreteScheduler` with dynamic resolution shifting -`QwenImageModel` is the hosted container (transformer + VAE + text encoder). +`QwenImageModel` is the hosted container (transformer + VAE). Like Stable +Diffusion 3's T5-XXL, the ~7B text tower is hosted separately as +`QwenImageTextEncoderModel` (`zeromodels/qwen-image-text-encoder`) and attached +with `from_weights(..., text_encoder=...)`. `QwenImageTextToImage` adds `generate` via `BaseDiffusion`, with **true CFG** (separate cond/uncond forwards and prediction-norm renormalization, Diffusers `true_cfg_scale`). @@ -40,7 +43,10 @@ Qwen-Image text-to-image, ported to pure Keras 3 from ```python from zeromodels.models.qwen_image import QwenImageTextToImage, QwenImageTokenizer -model = QwenImageTextToImage.from_weights("zeromodels/qwen-image") +model = QwenImageTextToImage.from_weights( + "zeromodels/qwen-image", + text_encoder="zeromodels/qwen-image-text-encoder", +) tok = QwenImageTokenizer.from_weights("zeromodels/qwen-image") image = model.generate( **tok("a coffee shop entrance with a chalkboard sign"), @@ -56,6 +62,7 @@ image = model.generate( | Variant | Hub (planned) | Source | |---|---|---| | `qwen-image` | `zeromodels/qwen-image` | [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) | +| `qwen-image-text-encoder` | `zeromodels/qwen-image-text-encoder` | [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) `text_encoder/` | Paper / model card: [Qwen-Image](https://huggingface.co/Qwen/Qwen-Image). License: Apache-2.0. diff --git a/tests/base/model_test_registry.py b/tests/base/model_test_registry.py index c558b204..3c4c42d7 100644 --- a/tests/base/model_test_registry.py +++ b/tests/base/model_test_registry.py @@ -4827,7 +4827,6 @@ "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", diff --git a/tests/fixtures/dummy_inputs.py b/tests/fixtures/dummy_inputs.py index b08e3312..f411ac98 100644 --- a/tests/fixtures/dummy_inputs.py +++ b/tests/fixtures/dummy_inputs.py @@ -245,7 +245,6 @@ def qwen_image_input( 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 { @@ -259,8 +258,6 @@ def qwen_image_input( ), "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"), } diff --git a/zeromodels/models/qwen_image/__init__.py b/zeromodels/models/qwen_image/__init__.py index 8406db89..73319cd9 100644 --- a/zeromodels/models/qwen_image/__init__.py +++ b/zeromodels/models/qwen_image/__init__.py @@ -5,6 +5,7 @@ ) from .qwen_image_model import ( QwenImageModel, + QwenImageTextEncoderModel, QwenImageTextToImage, QwenImageTransformer2DModel, ) @@ -16,6 +17,7 @@ "QwenImageConfig", "QwenImageModel", "QwenImageTextConfig", + "QwenImageTextEncoderModel", "QwenImageTextToImage", "QwenImageTokenizer", "QwenImageTransformer2DModel", 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 index 1f3457ab..cef4a37c 100644 --- a/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py +++ b/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py @@ -1,16 +1,27 @@ """Offline Diffusers ``Qwen/Qwen-Image`` → ZeroModels Keras weight conversion. -Converts the transformer, VAE, and Qwen2.5-VL text encoder into a hosted -``zeromodels/qwen-image`` layout (``zm_config.json`` + sharded weights). +Like SD3's T5-XXL, the ~7B Qwen2.5-VL text tower is converted in its own pass +(``clear_session`` in between), so the ~20B DiT and the text tower never share +memory: + +* ``/qwen-image/``: ``QwenImageModel`` (transformer + VAE) + + tokenizer files +* ``/qwen-image-text-encoder/``: ``QwenImageTextEncoderModel`` + +Each directory carries ``zm_config.json`` + ``model.weights.{h5,json}``. Usage:: python -m zeromodels.models.qwen_image.convert_qwen_image_diffusers_to_keras Env: - ZM_OUT_DIR output directory (default ``./qwen_image_weights``) - HF_TOKEN optional Hub token - ZM_DTYPE ``float16`` / ``bfloat16`` / ``float32`` (default ``float16``) + ZM_OUT_DIR output root (default ``./qwen_image_weights``) + HF_TOKEN optional Hub token + ZM_DTYPE ``float16`` / ``bfloat16`` / ``float32`` (default ``float16``) + ZM_DEVICE where weights are built (default ``cpu``; conversion is pure + copying, and the fp16 DiT alone is ~38 GiB) + ZM_VARIANTS comma list of ``qwen-image`` / ``qwen-image-text-encoder`` + (default: both) """ from __future__ import annotations @@ -32,8 +43,8 @@ transfer_weights, zeros_init, ) -from zeromodels.models.qwen2_5_vl.convert_qwen2_5_vl_hf_to_keras import ( - transfer_qwen2_5_vl_weights, +from zeromodels.models.qwen2_vl.convert_qwen2_vl_hf_to_keras import ( + WEIGHT_NAME_MAPPING as QWEN2_VL_WEIGHT_NAME_MAPPING, ) from zeromodels.models.stable_diffusion.convert_stable_diffusion_diffusers_to_keras import ( WEIGHT_NAME_MAPPING, @@ -42,6 +53,7 @@ QWEN_IMAGE_SOURCES = { "qwen-image": "Qwen/Qwen-Image", } +TEXT_ENCODER_VARIANT = "qwen-image-text-encoder" # VAE Diffusers RMSNorm keeps the leaf name ``gamma`` (not ``weight``). VAE_WEIGHT_NAME_MAPPING = { @@ -243,38 +255,29 @@ def transfer_component(component, state, mapping, desc, ignore=()): ) -def transfer_qwen_image(repo, token=None, dtype="float16", build_sample_size=16): - """Convert Diffusers Qwen-Image weights into a :class:`QwenImageModel`. +def transfer_qwen_image( + repo, token=None, dtype="float16", build_sample_size=16, config=None +): + """Convert the Diffusers transformer + VAE into a :class:`QwenImageModel`. - Builds the Keras container at a small ``transformer_sample_size`` (weights are - resolution-independent for the DiT) to keep the functional graph smaller, - transfers transformer / VAE / text-encoder weights, and returns ``(model, config)``. + Builds at a small ``transformer_sample_size`` (the DiT weights are + resolution-independent) to keep the functional graph small. The text tower + is converted separately by :func:`transfer_text_encoder`. + Returns ``(model, config)``. """ from zeromodels.base.base_mixin import build_dtype_scope from zeromodels.models.qwen_image.qwen_image_model import QwenImageModel - config = config_from_diffusers(repo, token=token) + config = config or config_from_diffusers(repo, token=token) flat = config.constructor_kwargs() - # Smaller spatial graph for conversion; Dense DiT weights do not depend on it. flat["transformer_sample_size"] = build_sample_size flat["vae_sample_size"] = max(build_sample_size * 8, 64) - print(f"[1/4] Building Keras QwenImageModel (dtype={dtype})…", flush=True) - import sys - import time - - t0 = time.time() + print(f"[1/3] Building QwenImageModel (dtype={dtype})…", flush=True) with build_dtype_scope(dtype), zeros_init(): - print(" building components…", flush=True) model = QwenImageModel(**flat) - print( - f" build done in {time.time() - t0:.1f}s " - f"({len(model.weights)} weights)", - flush=True, - ) - sys.stdout.flush() - print("[2/4] Transferring transformer…", flush=True) + print("[2/3] Transferring transformer…", flush=True) transformer_state = _load_safetensors_state( repo, "transformer", @@ -290,7 +293,7 @@ def transfer_qwen_image(repo, token=None, dtype="float16", build_sample_size=16) del transformer_state gc.collect() - print("[3/4] Transferring VAE…") + print("[3/3] Transferring VAE…", flush=True) vae_state = _load_single_safetensors( repo, "vae", "diffusion_pytorch_model.safetensors", token=token ) @@ -302,52 +305,97 @@ def transfer_qwen_image(repo, token=None, dtype="float16", build_sample_size=16) ) del vae_state gc.collect() + return model, config - print("[4/4] Transferring text encoder (Qwen2.5-VL text tower)…") - text_state = _load_safetensors_state( - repo, "text_encoder", "model.safetensors.index.json", token=token + +def transfer_text_encoder(repo, token=None, dtype="float16", config=None): + """Convert the Qwen2.5-VL text tower into a :class:`QwenImageTextEncoderModel`. + + Streams tensors from the safetensors shards one at a time (never a full + in-memory state dict); the vision tower and LM head are skipped. + Returns ``(model, text_config)``. + """ + from zeromodels.base.base_mixin import build_dtype_scope + from zeromodels.models.qwen_image.qwen_image_model import ( + QwenImageTextEncoderModel, ) - # Drop vision + LM head: our text-only tower does not carry them. - text_only = { - k: text_state[k] - for k in text_state.keys() - if k.startswith("model.") and not k.startswith("model.visual.") - } - # Materialize into a plain dict for the shared VL transfer helper. - text_np = {k: np.asarray(text_only[k]) for k in tqdm(text_only, desc="load text")} - del text_state, text_only - gc.collect() - transfer_qwen2_5_vl_weights(model.text_encoder, text_np) - del text_np - gc.collect() - return model, config + config = config or config_from_diffusers(repo, token=token) + text_config = config.text_config + print(f"[1/2] Building QwenImageTextEncoderModel (dtype={dtype})…", flush=True) + with build_dtype_scope(dtype), zeros_init(): + model = QwenImageTextEncoderModel(text_config) -def save_converted(model, config, out_dir, variant="qwen-image", max_shard_gb=5.0): - """Write ``zm_config.json``, sharded weights, and copy the tokenizer files.""" - import shutil + print("[2/2] Transferring text encoder…", flush=True) + state = _load_safetensors_state( + repo, "text_encoder", "model.safetensors.index.json", token=token + ) + # Both transformers layouts: ``model.layers.*`` and ``model.language_model.*``. + hf_keys = {} + for key in state.keys(): + 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() + for weight in tqdm(model.weights, desc="text_encoder"): + name = weight.path.replace("/", ".") + for old, new in QWEN2_VL_WEIGHT_NAME_MAPPING.items(): + name = name.replace(old, new) + if name not in hf_keys: + raise WeightMappingError(weight.path, name) + consumed.add(name) + transfer_weights(weight.path, weight, state[hf_keys[name]]) + unused = sorted(set(hf_keys) - consumed) + if unused: + raise ValueError( + f"text_encoder: {len(unused)} checkpoint tensors unused, e.g. {unused[:5]}." + ) + del state + gc.collect() + return model, text_config - from huggingface_hub import hf_hub_download + +def save_converted( + model, model_cls, config, out_dir, variant, dtype, max_shard_gb=5.0 +): + """Write ``model.weights.{h5,json}`` + ``zm_config.json`` into ``out_dir``.""" + from zeromodels.conversion.zm_config import write_zm_config os.makedirs(out_dir, exist_ok=True) - stem = os.path.join(out_dir, variant.replace("-", "_")) + import keras - config_path = os.path.join(out_dir, "zm_config.json") - with open(config_path, "w", encoding="utf-8") as f: - json.dump(config.to_dict(), f, indent=2) - print(f" wrote {config_path}") + def itemsize(dtype): + # np.dtype() rejects "bfloat16" + return 2 if "16" in keras.backend.standardize_dtype(dtype) else 4 - n_bytes = sum(int(np.prod(w.shape)) * 2 for w in model.weights) # ~fp16 + n_bytes = sum(int(np.prod(w.shape)) * itemsize(w.dtype) for w in model.weights) if n_bytes > max_shard_gb * 1024**3: - out = f"{stem}.weights.json" - model.save_weights(out, max_shard_size=max_shard_gb) + weights_filename = "model.weights.json" + model.save_weights( + os.path.join(out_dir, weights_filename), max_shard_size=max_shard_gb + ) else: - out = f"{stem}.weights.h5" - model.save_weights(out) - print(f" saved weights -> {out} (~{n_bytes / 1024**3:.2f} GB fp16-est)") + weights_filename = "model.weights.h5" + model.save_weights(os.path.join(out_dir, weights_filename)) + print(f" saved {weights_filename} ({n_bytes / 1024**3:.2f} GB)", flush=True) + write_zm_config( + out_dir, + model_cls, + variant, + config, + weights_filename=weights_filename, + weight_dtype=dtype, + ) + + +def copy_tokenizer_files(repo, out_dir, token=None): + import shutil + + from huggingface_hub import hf_hub_download - source = QWEN_IMAGE_SOURCES[variant] for name in ( "tokenizer.json", "tokenizer_config.json", @@ -357,31 +405,61 @@ def save_converted(model, config, out_dir, variant="qwen-image", max_shard_gb=5. "added_tokens.json", ): try: - src = hf_hub_download(source, name, subfolder="tokenizer") - shutil.copy2(src, os.path.join(out_dir, name if name != "tokenizer.json" else "tokenizer.json")) - except Exception as exc: # noqa: BLE001 — best-effort tokenizer copy + src = hf_hub_download(repo, name, subfolder="tokenizer", token=token) + except Exception as exc: # noqa: BLE001 — not every repo ships every file print(f" skip tokenizer file {name}: {exc}") + continue + shutil.copy2(src, os.path.join(out_dir, name)) if __name__ == "__main__": import keras - OUT_DIR = os.environ.get( - "ZM_OUT_DIR", - os.path.join(os.path.dirname(__file__), "..", "..", "..", "qwen_image_weights"), + from zeromodels.models.qwen_image.qwen_image_model import ( + QwenImageModel, + QwenImageTextEncoderModel, + ) + + OUT_DIR = os.path.abspath( + os.environ.get( + "ZM_OUT_DIR", + os.path.join( + os.path.dirname(__file__), "..", "..", "..", "qwen_image_weights" + ), + ) ) - OUT_DIR = os.path.abspath(OUT_DIR) - os.makedirs(OUT_DIR, exist_ok=True) token = os.environ.get("HF_TOKEN") dtype = os.environ.get("ZM_DTYPE", "float16") - selected = [v for v in os.environ.get("ZM_VARIANTS", "qwen-image").split(",") if v] + device = os.environ.get("ZM_DEVICE", "cpu") + default_variants = ",".join([*QWEN_IMAGE_SOURCES, TEXT_ENCODER_VARIANT]) + selected = [ + v for v in os.environ.get("ZM_VARIANTS", default_variants).split(",") if v + ] + source = QWEN_IMAGE_SOURCES["qwen-image"] + config = config_from_diffusers(source, token=token) for variant in selected: - source = QWEN_IMAGE_SOURCES[variant] - print(f"\n{'=' * 60}\nConverting: {variant} <- {source}\n{'=' * 60}") - model, config = transfer_qwen_image(source, token=token, dtype=dtype) - save_converted(model, config, OUT_DIR, variant=variant) + print( + f"\n{'=' * 60}\nConverting: {variant} <- {source} " + f"(device={device})\n{'=' * 60}", + flush=True, + ) + out_dir = os.path.join(OUT_DIR, variant) + with keras.device(device): + if variant == TEXT_ENCODER_VARIANT: + model, cfg = transfer_text_encoder( + source, token=token, dtype=dtype, config=config + ) + save_converted( + model, QwenImageTextEncoderModel, cfg, out_dir, variant, dtype + ) + else: + model, cfg = transfer_qwen_image( + QWEN_IMAGE_SOURCES[variant], token=token, dtype=dtype, config=config + ) + save_converted(model, QwenImageModel, cfg, out_dir, variant, dtype) + copy_tokenizer_files(QWEN_IMAGE_SOURCES[variant], out_dir, token=token) del model keras.backend.clear_session() gc.collect() - print(f"Done: {variant} -> {OUT_DIR}") + print(f"Done: {variant} -> {out_dir}", flush=True) diff --git a/zeromodels/models/qwen_image/qwen_image_model.py b/zeromodels/models/qwen_image/qwen_image_model.py index a7754a64..f6a1e62f 100644 --- a/zeromodels/models/qwen_image/qwen_image_model.py +++ b/zeromodels/models/qwen_image/qwen_image_model.py @@ -15,6 +15,7 @@ from zeromodels.models.qwen2_5_vl.qwen2_5_vl_model import Qwen2_5VLModel from zeromodels.models.qwen_image.qwen_image_config import ( QwenImageConfig, + QwenImageTextConfig, QwenImageTransformerConfig, ) from zeromodels.models.qwen_image.qwen_image_layers import ( @@ -27,7 +28,8 @@ from zeromodels.models.qwen_image.qwen_image_vae import AutoencoderKLQwenImage from zeromodels.models.stable_diffusion.stable_diffusion_layers import safe_name -QWEN_IMAGE_HUB_SIBLINGS = ("qwen-image",) +QWEN_IMAGE_HUB_SIBLINGS = frozenset({"QwenImageModel", "QwenImageTextToImage"}) +QWEN_IMAGE_TEXT_ENCODER_REPO = "zeromodels/qwen-image-text-encoder" PROMPT_TEMPLATE = ( "<|im_start|>system\nDescribe the image by detailing the color, shape, size, " @@ -226,12 +228,33 @@ def kwargs_from_diffusers_config(cls, cfg): } +@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. + + Hosted once, separately from the diffusion container + (``zeromodels/qwen-image-text-encoder``), like SD3's T5-XXL: the ~7B tower + plus the ~20B DiT do not fit one build on a 40 GB GPU. Attach it with + ``QwenImageTextToImage.from_weights(repo, text_encoder=...)``. + 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. + """Qwen-Image weights container: transformer + VAE. One functional ``keras.Model`` with disconnected paths (Diffusers - ``QwenImagePipeline`` components). Hosted as ``zeromodels/qwen-image``. + ``QwenImagePipeline`` components). Hosted as ``zeromodels/qwen-image``; the + Qwen2.5-VL text tower is a separate :class:`QwenImageTextEncoderModel` repo. On-the-fly ``hf:`` conversion is not supported. """ @@ -251,7 +274,7 @@ def __init__(self, name="QwenImageModel", **kwargs): setattr(self, attr, component) def build_components(self, config): - d, v, t = config.transformer_config, config.vae_config, config.text_config + d, v = config.transformer_config, config.vae_config transformer = QwenImageTransformer2DModel( patch_size=d.patch_size, in_channels=d.in_channels, @@ -278,33 +301,11 @@ def build_components(self, config): latents_std=v.latents_std, sample_size=v.sample_size, ) - text_encoder = Qwen2_5VLModel( - vocab_size=t.vocab_size, - embed_dim=t.embed_dim, - mlp_dim=t.mlp_dim, - num_layers=t.num_layers, - num_heads=t.num_heads, - num_kv_heads=t.num_kv_heads, - norm_eps=t.norm_eps, - rope_theta=t.rope_theta, - mrope_section=t.mrope_section, - tie_embeddings=t.tie_embeddings, - build_vision=False, - name="text_encoder", - ) - return { - "transformer": transformer, - "vae": vae, - "text_encoder": text_encoder, - } + return {"transformer": transformer, "vae": vae} 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"], - ) + d, v = config.transformer_config, config.vae_config + transformer, vae = components["transformer"], components["vae"] text_seq = config.max_sequence_length img_h = img_w = ( v.sample_size @@ -327,12 +328,6 @@ def build_graph(self, config, components): ), "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( { @@ -343,17 +338,10 @@ def build_graph(self, config, components): } )["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 from_hf(self, *args, **kwargs): @@ -370,10 +358,17 @@ class QwenImageTextToImage(QwenImageModel, BaseDiffusion): :: - model = QwenImageTextToImage.from_weights("zeromodels/qwen-image") + model = QwenImageTextToImage.from_weights( + "zeromodels/qwen-image", + text_encoder="zeromodels/qwen-image-text-encoder", + ) tok = QwenImageTokenizer.from_weights("zeromodels/qwen-image") image = model.generate(**tok("a cat"), height=1024, width=1024) + The text encoder lives outside the container's weights (like SD3's T5): + pass ``text_encoder=`` (a repo id or a built + :class:`QwenImageTextEncoderModel`) or assign ``model.text_encoder``. + 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. @@ -395,6 +390,32 @@ def __init__(self, scheduler=None, name="QwenImageTextToImage", **kwargs): ) self.scheduler = scheduler + @classmethod + def from_weights(cls, identifier, text_encoder=None, **kwargs): + """``BaseModel.from_weights`` plus ``text_encoder``: a hosted + :class:`QwenImageTextEncoderModel` repo id (loaded with the same + ``load_dtype``) or a built model to attach.""" + model = super().from_weights(identifier, **kwargs) + if text_encoder is not None: + if isinstance(text_encoder, str): + text_encoder = QwenImageTextEncoderModel.from_weights( + text_encoder, load_dtype=kwargs.get("load_dtype") + ) + model.text_encoder = text_encoder + return model + + @property + def text_encoder(self): + return self.__dict__.get("_text_encoder") + + def __setattr__(self, name, value): + # kept out of the tracked (saved / loaded / device-moved) sub-layers, so + # the container's weights stay transformer + VAE only + if name == "text_encoder": + self.__dict__["_text_encoder"] = value + return + super().__setattr__(name, value) + def default_scheduler(self): return FlowMatchEulerDiscreteScheduler( shift=1.0, @@ -444,6 +465,12 @@ def encode_prompt(self, input_ids, attention_mask=None, **conditioning): ``max_sequence_length``. """ del conditioning + if self.text_encoder is None: + raise ValueError( + "QwenImageTextToImage has no text encoder attached. Load one with " + f"from_weights(..., text_encoder={QWEN_IMAGE_TEXT_ENCODER_REPO!r}) " + "or assign model.text_encoder = QwenImageTextEncoderModel(...)." + ) input_ids = ops.cast(ops.convert_to_tensor(input_ids), "int32") if attention_mask is None: attention_mask = ops.ones_like(input_ids) From 0057ad59a69aa11646b3ed941c84eedaf9519aae Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Tue, 22 Sep 2026 17:43:34 -0700 Subject: [PATCH 3/8] fix --- .../convert_qwen_image_diffusers_to_keras.py | 24 +++++++++---------- .../models/qwen_image/qwen_image_vae.py | 7 +++--- 2 files changed, 15 insertions(+), 16 deletions(-) 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 index cef4a37c..250f5ffd 100644 --- a/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py +++ b/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py @@ -54,8 +54,6 @@ "qwen-image": "Qwen/Qwen-Image", } TEXT_ENCODER_VARIANT = "qwen-image-text-encoder" - -# VAE Diffusers RMSNorm keeps the leaf name ``gamma`` (not ``weight``). VAE_WEIGHT_NAME_MAPPING = { k: v for k, v in WEIGHT_NAME_MAPPING.items() if "/gamma" not in k } @@ -147,12 +145,13 @@ def _keras_to_torch_key(keras_weight, mapping): def _prepare_torch_array(keras_weight, torch_weight): - """Transpose Conv3D / squeeze VAE RMSNorm gamma to match Keras shapes.""" + """Transpose conv kernels / squeeze VAE RMSNorm gamma to match Keras shapes.""" arr = np.asarray(torch_weight) kshape = tuple(keras_weight.shape) if len(kshape) == 5 and arr.ndim == 5: - # Diffusers Conv3d NCDHW ``(O, I, T, H, W)`` → Keras NDHWC ``(T, H, W, I, O)``. 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) return arr @@ -223,13 +222,15 @@ def transfer_component(component, state, mapping, desc, ignore=()): consumed.add(key) raw = state[key] torch_weight = _prepare_torch_array(keras_weight, raw) + if len(keras_weight.shape) in (4, 5): + if tuple(keras_weight.shape) != torch_weight.shape: + raise WeightShapeMismatchError( + keras_weight.path, keras_weight.shape, key, np.shape(raw) + ) + keras_weight.assign(torch_weight) + continue if tuple(keras_weight.shape) != tuple(np.asarray(torch_weight).shape): - # Dense / Conv2D still need the shared transpose inside transfer_weights; - # only enforce exact match after _prepare (Conv3D / squeezed gamma). - if len(keras_weight.shape) == 5 or ( - len(keras_weight.shape) == 1 - and np.asarray(raw).ndim > 1 - ): + if len(keras_weight.shape) == 1 and np.asarray(raw).ndim > 1: raise WeightShapeMismatchError( keras_weight.path, keras_weight.shape, @@ -247,7 +248,6 @@ def transfer_component(component, state, mapping, desc, ignore=()): ) transfer_weights(keras_weight.path, keras_weight, torch_weight) unused = sorted(set(state) - consumed) - # RoPE tables / unused buffers on the torch side are fine to ignore when empty. if unused: raise ValueError( f"{type(component).__name__}: {len(unused)} checkpoint tensors " @@ -331,7 +331,6 @@ def transfer_text_encoder(repo, token=None, dtype="float16", config=None): state = _load_safetensors_state( repo, "text_encoder", "model.safetensors.index.json", token=token ) - # Both transformers layouts: ``model.layers.*`` and ``model.language_model.*``. hf_keys = {} for key in state.keys(): if key.startswith("model.language_model."): @@ -368,7 +367,6 @@ def save_converted( import keras def itemsize(dtype): - # np.dtype() rejects "bfloat16" return 2 if "16" in keras.backend.standardize_dtype(dtype) else 4 n_bytes = sum(int(np.prod(w.shape)) * itemsize(w.dtype) for w in model.weights) diff --git a/zeromodels/models/qwen_image/qwen_image_vae.py b/zeromodels/models/qwen_image/qwen_image_vae.py index 987b9039..c822fde6 100644 --- a/zeromodels/models/qwen_image/qwen_image_vae.py +++ b/zeromodels/models/qwen_image/qwen_image_vae.py @@ -142,11 +142,12 @@ def build(self, input_shape): self.built = True def call(self, x): - # Normalize over channels (last axis for NDHWC / NHWC). + # Diffusers: F.normalize(x, dim=channels) * sqrt(dim) * gamma, i.e. an L2 + # normalize (eps 1e-12 on the norm) over the last (channel) axis here. dtype = x.dtype x_f = ops.cast(x, "float32") - variance = ops.mean(ops.square(x_f), axis=-1, keepdims=True) - x_f = x_f * ops.rsqrt(variance + 1e-6) + 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) From cf4b5390226a8453d23605fb19115e9f03575ca8 Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Tue, 22 Sep 2026 18:15:25 -0700 Subject: [PATCH 4/8] fix --- docs/qwen_image.md | 13 +- tests/base/model_test_registry.py | 1 + tests/fixtures/dummy_inputs.py | 3 + .../convert_qwen_image_diffusers_to_keras.py | 191 ++++-------------- .../models/qwen_image/qwen_image_layers.py | 5 + .../models/qwen_image/qwen_image_model.py | 76 +++---- 6 files changed, 76 insertions(+), 213 deletions(-) diff --git a/docs/qwen_image.md b/docs/qwen_image.md index 9ac25880..b7908921 100644 --- a/docs/qwen_image.md +++ b/docs/qwen_image.md @@ -18,10 +18,8 @@ Qwen-Image text-to-image, ported to pure Keras 3 from - **Text encoder**: Qwen2.5-VL-7B Instruct text tower (ChatML prompt template) - **Scheduler**: `FlowMatchEulerDiscreteScheduler` with dynamic resolution shifting -`QwenImageModel` is the hosted container (transformer + VAE). Like Stable -Diffusion 3's T5-XXL, the ~7B text tower is hosted separately as -`QwenImageTextEncoderModel` (`zeromodels/qwen-image-text-encoder`) and attached -with `from_weights(..., text_encoder=...)`. +`QwenImageModel` is the hosted container (transformer + VAE + text tower), +stored as one set of sharded weights (`model.weights.json`, ~53 GiB at 16-bit). `QwenImageTextToImage` adds `generate` via `BaseDiffusion`, with **true CFG** (separate cond/uncond forwards and prediction-norm renormalization, Diffusers `true_cfg_scale`). @@ -43,10 +41,7 @@ with `from_weights(..., text_encoder=...)`. ```python from zeromodels.models.qwen_image import QwenImageTextToImage, QwenImageTokenizer -model = QwenImageTextToImage.from_weights( - "zeromodels/qwen-image", - text_encoder="zeromodels/qwen-image-text-encoder", -) +model = QwenImageTextToImage.from_weights("zeromodels/qwen-image") tok = QwenImageTokenizer.from_weights("zeromodels/qwen-image") image = model.generate( **tok("a coffee shop entrance with a chalkboard sign"), @@ -62,7 +57,5 @@ image = model.generate( | Variant | Hub (planned) | Source | |---|---|---| | `qwen-image` | `zeromodels/qwen-image` | [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) | -| `qwen-image-text-encoder` | `zeromodels/qwen-image-text-encoder` | [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) `text_encoder/` | - Paper / model card: [Qwen-Image](https://huggingface.co/Qwen/Qwen-Image). License: Apache-2.0. diff --git a/tests/base/model_test_registry.py b/tests/base/model_test_registry.py index 3c4c42d7..c558b204 100644 --- a/tests/base/model_test_registry.py +++ b/tests/base/model_test_registry.py @@ -4827,6 +4827,7 @@ "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", diff --git a/tests/fixtures/dummy_inputs.py b/tests/fixtures/dummy_inputs.py index f411ac98..b08e3312 100644 --- a/tests/fixtures/dummy_inputs.py +++ b/tests/fixtures/dummy_inputs.py @@ -245,6 +245,7 @@ def qwen_image_input( 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 { @@ -258,6 +259,8 @@ def qwen_image_input( ), "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"), } 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 index 250f5ffd..8a9622f8 100644 --- a/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py +++ b/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py @@ -1,31 +1,3 @@ -"""Offline Diffusers ``Qwen/Qwen-Image`` → ZeroModels Keras weight conversion. - -Like SD3's T5-XXL, the ~7B Qwen2.5-VL text tower is converted in its own pass -(``clear_session`` in between), so the ~20B DiT and the text tower never share -memory: - -* ``/qwen-image/``: ``QwenImageModel`` (transformer + VAE) + - tokenizer files -* ``/qwen-image-text-encoder/``: ``QwenImageTextEncoderModel`` - -Each directory carries ``zm_config.json`` + ``model.weights.{h5,json}``. - -Usage:: - - python -m zeromodels.models.qwen_image.convert_qwen_image_diffusers_to_keras - -Env: - ZM_OUT_DIR output root (default ``./qwen_image_weights``) - HF_TOKEN optional Hub token - ZM_DTYPE ``float16`` / ``bfloat16`` / ``float32`` (default ``float16``) - ZM_DEVICE where weights are built (default ``cpu``; conversion is pure - copying, and the fp16 DiT alone is ~38 GiB) - ZM_VARIANTS comma list of ``qwen-image`` / ``qwen-image-text-encoder`` - (default: both) -""" - -from __future__ import annotations - import gc import json import os @@ -53,7 +25,6 @@ QWEN_IMAGE_SOURCES = { "qwen-image": "Qwen/Qwen-Image", } -TEXT_ENCODER_VARIANT = "qwen-image-text-encoder" VAE_WEIGHT_NAME_MAPPING = { k: v for k, v in WEIGHT_NAME_MAPPING.items() if "/gamma" not in k } @@ -258,11 +229,10 @@ def transfer_component(component, state, mapping, desc, ignore=()): def transfer_qwen_image( repo, token=None, dtype="float16", build_sample_size=16, config=None ): - """Convert the Diffusers transformer + VAE into a :class:`QwenImageModel`. + """Convert the Diffusers pipeline into a :class:`QwenImageModel`. Builds at a small ``transformer_sample_size`` (the DiT weights are - resolution-independent) to keep the functional graph small. The text tower - is converted separately by :func:`transfer_text_encoder`. + resolution-independent) to keep the functional graph small. Returns ``(model, config)``. """ from zeromodels.base.base_mixin import build_dtype_scope @@ -273,11 +243,11 @@ def transfer_qwen_image( flat["transformer_sample_size"] = build_sample_size flat["vae_sample_size"] = max(build_sample_size * 8, 64) - print(f"[1/3] Building QwenImageModel (dtype={dtype})…", flush=True) + print(f"[1/4] Building QwenImageModel (dtype={dtype})…", flush=True) with build_dtype_scope(dtype), zeros_init(): model = QwenImageModel(**flat) - print("[2/3] Transferring transformer…", flush=True) + print("[2/4] Transferring transformer…", flush=True) transformer_state = _load_safetensors_state( repo, "transformer", @@ -293,7 +263,7 @@ def transfer_qwen_image( del transformer_state gc.collect() - print("[3/3] Transferring VAE…", flush=True) + print("[3/4] Transferring VAE…", flush=True) vae_state = _load_single_safetensors( repo, "vae", "diffusion_pytorch_model.safetensors", token=token ) @@ -305,29 +275,19 @@ def transfer_qwen_image( ) del vae_state gc.collect() + + print("[4/4] Transferring text encoder…", flush=True) + transfer_text_encoder(model.text_encoder, repo, token=token) + gc.collect() return model, config -def transfer_text_encoder(repo, token=None, dtype="float16", config=None): - """Convert the Qwen2.5-VL text tower into a :class:`QwenImageTextEncoderModel`. +def transfer_text_encoder(text_encoder, repo, token=None): + """Copy the Qwen2.5-VL text tower into ``text_encoder``. Streams tensors from the safetensors shards one at a time (never a full in-memory state dict); the vision tower and LM head are skipped. - Returns ``(model, text_config)``. """ - from zeromodels.base.base_mixin import build_dtype_scope - from zeromodels.models.qwen_image.qwen_image_model import ( - QwenImageTextEncoderModel, - ) - - config = config or config_from_diffusers(repo, token=token) - text_config = config.text_config - - print(f"[1/2] Building QwenImageTextEncoderModel (dtype={dtype})…", flush=True) - with build_dtype_scope(dtype), zeros_init(): - model = QwenImageTextEncoderModel(text_config) - - print("[2/2] Transferring text encoder…", flush=True) state = _load_safetensors_state( repo, "text_encoder", "model.safetensors.index.json", token=token ) @@ -339,8 +299,9 @@ def transfer_text_encoder(repo, token=None, dtype="float16", config=None): hf_keys[key] = key consumed = set() - for weight in tqdm(model.weights, desc="text_encoder"): - name = weight.path.replace("/", ".") + for weight in tqdm(text_encoder.weights, desc="text_encoder"): + path = weight.path.removeprefix(f"{text_encoder.name}/") + name = path.replace("/", ".") for old, new in QWEN2_VL_WEIGHT_NAME_MAPPING.items(): name = name.replace(old, new) if name not in hf_keys: @@ -353,111 +314,41 @@ def transfer_text_encoder(repo, token=None, dtype="float16", config=None): f"text_encoder: {len(unused)} checkpoint tensors unused, e.g. {unused[:5]}." ) del state - gc.collect() - return model, text_config - - -def save_converted( - model, model_cls, config, out_dir, variant, dtype, max_shard_gb=5.0 -): - """Write ``model.weights.{h5,json}`` + ``zm_config.json`` into ``out_dir``.""" - from zeromodels.conversion.zm_config import write_zm_config - - os.makedirs(out_dir, exist_ok=True) - import keras - - def itemsize(dtype): - return 2 if "16" in keras.backend.standardize_dtype(dtype) else 4 - - n_bytes = sum(int(np.prod(w.shape)) * itemsize(w.dtype) for w in model.weights) - if n_bytes > max_shard_gb * 1024**3: - weights_filename = "model.weights.json" - model.save_weights( - os.path.join(out_dir, weights_filename), max_shard_size=max_shard_gb - ) - else: - weights_filename = "model.weights.h5" - model.save_weights(os.path.join(out_dir, weights_filename)) - print(f" saved {weights_filename} ({n_bytes / 1024**3:.2f} GB)", flush=True) - write_zm_config( - out_dir, - model_cls, - variant, - config, - weights_filename=weights_filename, - weight_dtype=dtype, - ) - - -def copy_tokenizer_files(repo, out_dir, token=None): - import shutil - - from huggingface_hub import hf_hub_download - - for name in ( - "tokenizer.json", - "tokenizer_config.json", - "special_tokens_map.json", - "vocab.json", - "merges.txt", - "added_tokens.json", - ): - try: - src = hf_hub_download(repo, name, subfolder="tokenizer", token=token) - except Exception as exc: # noqa: BLE001 — not every repo ships every file - print(f" skip tokenizer file {name}: {exc}") - continue - shutil.copy2(src, os.path.join(out_dir, name)) if __name__ == "__main__": import keras - from zeromodels.models.qwen_image.qwen_image_model import ( - QwenImageModel, - QwenImageTextEncoderModel, - ) - - OUT_DIR = os.path.abspath( - os.environ.get( - "ZM_OUT_DIR", - os.path.join( - os.path.dirname(__file__), "..", "..", "..", "qwen_image_weights" - ), - ) - ) + 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", "float16") + dtype = os.environ.get("ZM_DTYPE", "bfloat16") device = os.environ.get("ZM_DEVICE", "cpu") - default_variants = ",".join([*QWEN_IMAGE_SOURCES, TEXT_ENCODER_VARIANT]) - selected = [ - v for v in os.environ.get("ZM_VARIANTS", default_variants).split(",") if v - ] - source = QWEN_IMAGE_SOURCES["qwen-image"] - config = config_from_diffusers(source, token=token) - - for variant in selected: - print( - f"\n{'=' * 60}\nConverting: {variant} <- {source} " - f"(device={device})\n{'=' * 60}", - flush=True, - ) - out_dir = os.path.join(OUT_DIR, variant) + 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): - if variant == TEXT_ENCODER_VARIANT: - model, cfg = transfer_text_encoder( - source, token=token, dtype=dtype, config=config - ) - save_converted( - model, QwenImageTextEncoderModel, cfg, out_dir, variant, dtype - ) - else: - model, cfg = transfer_qwen_image( - QWEN_IMAGE_SOURCES[variant], token=token, dtype=dtype, config=config - ) - save_converted(model, QwenImageModel, cfg, out_dir, variant, dtype) - copy_tokenizer_files(QWEN_IMAGE_SOURCES[variant], out_dir, token=token) + 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() - print(f"Done: {variant} -> {out_dir}", flush=True) + diff --git a/zeromodels/models/qwen_image/qwen_image_layers.py b/zeromodels/models/qwen_image/qwen_image_layers.py index de55d214..8ba50ae1 100644 --- a/zeromodels/models/qwen_image/qwen_image_layers.py +++ b/zeromodels/models/qwen_image/qwen_image_layers.py @@ -257,17 +257,22 @@ def __init__( self._neg_freqs_np = neg_freqs def build(self, input_shape=None): + # angles reach ~4096 rad: 16-bit storage would wreck the rotation 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 diff --git a/zeromodels/models/qwen_image/qwen_image_model.py b/zeromodels/models/qwen_image/qwen_image_model.py index f6a1e62f..b02bd61a 100644 --- a/zeromodels/models/qwen_image/qwen_image_model.py +++ b/zeromodels/models/qwen_image/qwen_image_model.py @@ -29,8 +29,6 @@ from zeromodels.models.stable_diffusion.stable_diffusion_layers import safe_name QWEN_IMAGE_HUB_SIBLINGS = frozenset({"QwenImageModel", "QwenImageTextToImage"}) -QWEN_IMAGE_TEXT_ENCODER_REPO = "zeromodels/qwen-image-text-encoder" - PROMPT_TEMPLATE = ( "<|im_start|>system\nDescribe the image by detailing the color, shape, size, " "texture, quantity, text, spatial relationships of the objects and " @@ -232,10 +230,7 @@ def kwargs_from_diffusers_config(cls, cfg): class QwenImageTextEncoderModel(Qwen2_5VLModel): """Qwen-Image prompt encoder: the Qwen2.5-VL-7B text tower, no vision / LM head. - Hosted once, separately from the diffusion container - (``zeromodels/qwen-image-text-encoder``), like SD3's T5-XXL: the ~7B tower - plus the ~20B DiT do not fit one build on a 40 GB GPU. Attach it with - ``QwenImageTextToImage.from_weights(repo, text_encoder=...)``. + The ``text_encoder`` component of :class:`QwenImageModel`. Inputs ``input_ids`` / ``attention_mask``; output ``last_hidden_state``. """ @@ -250,12 +245,11 @@ def __init__(self, max_seq_len=1024, name="text_encoder", **kwargs): @keras.saving.register_keras_serializable(package="zeromodels") class QwenImageModel(BaseModel): - """Qwen-Image weights container: transformer + VAE. + """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``; the - Qwen2.5-VL text tower is a separate :class:`QwenImageTextEncoderModel` repo. - On-the-fly ``hf:`` conversion is not supported. + ``QwenImagePipeline`` components). Hosted as ``zeromodels/qwen-image`` + (sharded weights). On-the-fly ``hf:`` conversion is not supported. """ config_class = QwenImageConfig @@ -301,11 +295,16 @@ def build_components(self, config): latents_std=v.latents_std, sample_size=v.sample_size, ) - return {"transformer": transformer, "vae": vae} + text_encoder = QwenImageTextEncoderModel(config.text_config) + return {"transformer": transformer, "vae": vae, "text_encoder": text_encoder} def build_graph(self, config, components): - d, v = config.transformer_config, config.vae_config - transformer, vae = components["transformer"], components["vae"] + 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 @@ -328,6 +327,12 @@ def build_graph(self, config, components): ), "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( { @@ -338,10 +343,14 @@ def build_graph(self, config, components): } )["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 from_hf(self, *args, **kwargs): @@ -358,17 +367,10 @@ class QwenImageTextToImage(QwenImageModel, BaseDiffusion): :: - model = QwenImageTextToImage.from_weights( - "zeromodels/qwen-image", - text_encoder="zeromodels/qwen-image-text-encoder", - ) + model = QwenImageTextToImage.from_weights("zeromodels/qwen-image") tok = QwenImageTokenizer.from_weights("zeromodels/qwen-image") image = model.generate(**tok("a cat"), height=1024, width=1024) - The text encoder lives outside the container's weights (like SD3's T5): - pass ``text_encoder=`` (a repo id or a built - :class:`QwenImageTextEncoderModel`) or assign ``model.text_encoder``. - 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. @@ -390,32 +392,6 @@ def __init__(self, scheduler=None, name="QwenImageTextToImage", **kwargs): ) self.scheduler = scheduler - @classmethod - def from_weights(cls, identifier, text_encoder=None, **kwargs): - """``BaseModel.from_weights`` plus ``text_encoder``: a hosted - :class:`QwenImageTextEncoderModel` repo id (loaded with the same - ``load_dtype``) or a built model to attach.""" - model = super().from_weights(identifier, **kwargs) - if text_encoder is not None: - if isinstance(text_encoder, str): - text_encoder = QwenImageTextEncoderModel.from_weights( - text_encoder, load_dtype=kwargs.get("load_dtype") - ) - model.text_encoder = text_encoder - return model - - @property - def text_encoder(self): - return self.__dict__.get("_text_encoder") - - def __setattr__(self, name, value): - # kept out of the tracked (saved / loaded / device-moved) sub-layers, so - # the container's weights stay transformer + VAE only - if name == "text_encoder": - self.__dict__["_text_encoder"] = value - return - super().__setattr__(name, value) - def default_scheduler(self): return FlowMatchEulerDiscreteScheduler( shift=1.0, @@ -465,12 +441,6 @@ def encode_prompt(self, input_ids, attention_mask=None, **conditioning): ``max_sequence_length``. """ del conditioning - if self.text_encoder is None: - raise ValueError( - "QwenImageTextToImage has no text encoder attached. Load one with " - f"from_weights(..., text_encoder={QWEN_IMAGE_TEXT_ENCODER_REPO!r}) " - "or assign model.text_encoder = QwenImageTextEncoderModel(...)." - ) input_ids = ops.cast(ops.convert_to_tensor(input_ids), "int32") if attention_mask is None: attention_mask = ops.ones_like(input_ids) From 385e61c880c495d47768e0032be070864a20b6c6 Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Tue, 22 Sep 2026 22:21:32 -0700 Subject: [PATCH 5/8] fix --- tests/fixtures/cross_backend_parity.json | 592 ++++++++++++++++++ tests/fixtures/dummy_inputs.py | 2 +- tests/integration/test_data_formats.py | 3 + .../models/qwen_image/qwen_image_layers.py | 7 +- .../models/qwen_image/qwen_image_model.py | 5 + 5 files changed, 605 insertions(+), 4 deletions(-) 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 b08e3312..2467e0f9 100644 --- a/tests/fixtures/dummy_inputs.py +++ b/tests/fixtures/dummy_inputs.py @@ -241,7 +241,7 @@ def qwen_image_input( packed_seq=16, in_channels=16, image_size=32, - latent_size=16, + latent_size=8, z_dim=4, text_seq_len=16, joint_attention_dim=32, 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/zeromodels/models/qwen_image/qwen_image_layers.py b/zeromodels/models/qwen_image/qwen_image_layers.py index 8ba50ae1..ff5a7b65 100644 --- a/zeromodels/models/qwen_image/qwen_image_layers.py +++ b/zeromodels/models/qwen_image/qwen_image_layers.py @@ -386,6 +386,7 @@ def __init__( ): 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) @@ -565,6 +566,7 @@ def __init__( ): 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) @@ -709,12 +711,11 @@ def call( + txt_gate2[:, None, :] * self.txt_mlp(txt_modulated2) ) - # Diffusers clips fp16 dual-stream outputs to the fp16 finite range. - if str(encoder_hidden_states.dtype).endswith("float16"): + if keras.backend.standardize_dtype(encoder_hidden_states.dtype) == "float16": encoder_hidden_states = ops.clip( encoder_hidden_states, -65504.0, 65504.0 ) - if str(hidden_states.dtype).endswith("float16"): + 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 diff --git a/zeromodels/models/qwen_image/qwen_image_model.py b/zeromodels/models/qwen_image/qwen_image_model.py index b02bd61a..bf1165e9 100644 --- a/zeromodels/models/qwen_image/qwen_image_model.py +++ b/zeromodels/models/qwen_image/qwen_image_model.py @@ -353,6 +353,11 @@ def build_graph(self, config, components): "prompt_embeds": text_out["last_hidden_state"], } + def get_config(self): + config = super().get_config() + config.update(self.config.constructor_kwargs()) + return config + def from_hf(self, *args, **kwargs): raise NotImplementedError( "On-the-fly hf: conversion is not supported for Qwen-Image; " From 120ceaeba3558be654ede1f1e346447c02cec856 Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Tue, 22 Sep 2026 23:07:00 -0700 Subject: [PATCH 6/8] refactor-code --- zeromodels/models/qwen_image/__init__.py | 3 +- .../convert_qwen_image_diffusers_to_keras.py | 376 +++++++++--------- .../models/qwen_image/qwen_image_config.py | 62 ++- .../models/qwen_image/qwen_image_layers.py | 45 +-- .../models/qwen_image/qwen_image_model.py | 245 ++++++++++-- .../models/qwen_image/qwen_image_tokenizer.py | 4 - .../models/qwen_image/qwen_image_vae.py | 358 +---------------- 7 files changed, 481 insertions(+), 612 deletions(-) diff --git a/zeromodels/models/qwen_image/__init__.py b/zeromodels/models/qwen_image/__init__.py index 73319cd9..d7fe9377 100644 --- a/zeromodels/models/qwen_image/__init__.py +++ b/zeromodels/models/qwen_image/__init__.py @@ -2,15 +2,16 @@ QwenImageConfig, QwenImageTextConfig, QwenImageTransformerConfig, + QwenImageVAEConfig, ) from .qwen_image_model import ( + AutoencoderKLQwenImage, QwenImageModel, QwenImageTextEncoderModel, QwenImageTextToImage, QwenImageTransformer2DModel, ) from .qwen_image_tokenizer import QwenImageTokenizer -from .qwen_image_vae import AutoencoderKLQwenImage, QwenImageVAEConfig __all__ = [ "AutoencoderKLQwenImage", 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 index 8a9622f8..5d7a7d73 100644 --- a/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py +++ b/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py @@ -1,6 +1,4 @@ -import gc -import json -import os +from typing import Dict import numpy as np from tqdm import tqdm @@ -15,35 +13,53 @@ transfer_weights, zeros_init, ) -from zeromodels.models.qwen2_vl.convert_qwen2_vl_hf_to_keras import ( - WEIGHT_NAME_MAPPING as QWEN2_VL_WEIGHT_NAME_MAPPING, -) -from zeromodels.models.stable_diffusion.convert_stable_diffusion_diffusers_to_keras import ( - WEIGHT_NAME_MAPPING, -) QWEN_IMAGE_SOURCES = { "qwen-image": "Qwen/Qwen-Image", } -VAE_WEIGHT_NAME_MAPPING = { - k: v for k, v in WEIGHT_NAME_MAPPING.items() if "/gamma" not in k + +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", } -VAE_WEIGHT_NAME_MAPPING["/"] = "." -VAE_WEIGHT_NAME_MAPPING["__"] = "." def config_from_diffusers(repo, token=None): - """Build a :class:`QwenImageConfig` from a Diffusers Qwen-Image repo.""" + 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 + from zeromodels.models.qwen_image.qwen_image_config import ( + QwenImageConfig, + QwenImageVAEConfig, + ) from zeromodels.models.qwen_image.qwen_image_model import ( QwenImageTransformer2DModel, ) - from zeromodels.models.qwen_image.qwen_image_vae import QwenImageVAEConfig transformer = dict( DiffusersTransformer.load_config(repo, subfolder="transformer", token=token) @@ -108,133 +124,15 @@ def config_from_diffusers(repo, token=None): ) -def _keras_to_torch_key(keras_weight, mapping): - key = "/".join(keras_weight.path.split("/")[-2:]) - for old, new in mapping.items(): - key = key.replace(old, new) - return key - - -def _prepare_torch_array(keras_weight, torch_weight): - """Transpose conv kernels / squeeze VAE RMSNorm gamma to match Keras shapes.""" - arr = np.asarray(torch_weight) - 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) - return arr - - -def _load_safetensors_state(repo, subfolder, index_name, token=None): - """Lazy shard map: ``{key: (shard_path, key)}`` resolved on demand via safe_open.""" - from huggingface_hub import hf_hub_download - from safetensors import safe_open - - 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 _ShardDict(dict): - def __contains__(self, key): - return key in weight_map - - def __getitem__(self, key): - if key not in weight_map: - raise KeyError(key) - path = shard_paths[weight_map[key]] - with safe_open(path, 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) - - return _ShardDict() - - -def _load_single_safetensors(repo, subfolder, filename, token=None): - from huggingface_hub import hf_hub_download - from safetensors import safe_open - - 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) - return state - - -def transfer_component(component, state, mapping, desc, ignore=()): - """Copy ``state`` tensors into a Keras sub-model using ``mapping``.""" - consumed = set(ignore) - trainable, non_trainable = split_model_weights(component) - for keras_weight, _ in tqdm(trainable + non_trainable, desc=desc): - if any(part in keras_weight.path for part in ("pos_freqs", "neg_freqs")): - continue - key = _keras_to_torch_key(keras_weight, mapping) - if key in ignore: - continue - if key not in state: - raise WeightMappingError(keras_weight.path, key) - consumed.add(key) - raw = state[key] - torch_weight = _prepare_torch_array(keras_weight, raw) - if len(keras_weight.shape) in (4, 5): - if tuple(keras_weight.shape) != torch_weight.shape: - raise WeightShapeMismatchError( - keras_weight.path, keras_weight.shape, key, np.shape(raw) - ) - keras_weight.assign(torch_weight) - continue - if tuple(keras_weight.shape) != tuple(np.asarray(torch_weight).shape): - if len(keras_weight.shape) == 1 and np.asarray(raw).ndim > 1: - raise WeightShapeMismatchError( - keras_weight.path, - keras_weight.shape, - key, - torch_weight.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, torch_weight) - unused = sorted(set(state) - consumed) - if unused: - raise ValueError( - f"{type(component).__name__}: {len(unused)} checkpoint tensors " - f"unused, e.g. {unused[:5]}." - ) - - def transfer_qwen_image( repo, token=None, dtype="float16", build_sample_size=16, config=None ): - """Convert the Diffusers pipeline into a :class:`QwenImageModel`. + import gc + import json + + from huggingface_hub import hf_hub_download + from safetensors import safe_open - Builds at a small ``transformer_sample_size`` (the DiT weights are - resolution-independent) to keep the functional graph small. - Returns ``(model, config)``. - """ from zeromodels.base.base_mixin import build_dtype_scope from zeromodels.models.qwen_image.qwen_image_model import QwenImageModel @@ -247,79 +145,176 @@ def transfer_qwen_image( with build_dtype_scope(dtype), zeros_init(): model = QwenImageModel(**flat) - print("[2/4] Transferring transformer…", flush=True) - transformer_state = _load_safetensors_state( - repo, - "transformer", - "diffusion_pytorch_model.safetensors.index.json", - token=token, - ) - transfer_component( - model.transformer, - transformer_state, - WEIGHT_NAME_MAPPING, - desc="transformer", - ) - del transformer_state - gc.collect() - - print("[3/4] Transferring VAE…", flush=True) - vae_state = _load_single_safetensors( - repo, "vae", "diffusion_pytorch_model.safetensors", token=token - ) - transfer_component( - model.vae, - vae_state, - VAE_WEIGHT_NAME_MAPPING, - desc="vae", - ) - del vae_state - gc.collect() + # 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) - transfer_text_encoder(model.text_encoder, repo, token=token) - gc.collect() - return model, config - - -def transfer_text_encoder(text_encoder, repo, token=None): - """Copy the Qwen2.5-VL text tower into ``text_encoder``. - - Streams tensors from the safetensors shards one at a time (never a full - in-memory state dict); the vision tower and LM head are skipped. - """ - state = _load_safetensors_state( - repo, "text_encoder", "model.safetensors.index.json", token=token + 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 state.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() - for weight in tqdm(text_encoder.weights, desc="text_encoder"): - path = weight.path.removeprefix(f"{text_encoder.name}/") - name = path.replace("/", ".") - for old, new in QWEN2_VL_WEIGHT_NAME_MAPPING.items(): + 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) - transfer_weights(weight.path, weight, state[hf_keys[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]}." + f"text_encoder: {len(unused)} checkpoint tensors unused, " + f"e.g. {unused[:5]}." ) - del state + 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") + 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") @@ -351,4 +346,3 @@ def transfer_text_encoder(text_encoder, repo, token=None): 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 index 2653fe1c..9adf5cb0 100644 --- a/zeromodels/models/qwen_image/qwen_image_config.py +++ b/zeromodels/models/qwen_image/qwen_image_config.py @@ -1,8 +1,64 @@ -"""Typed configs for Qwen-Image (Diffusers ``QwenImagePipeline``).""" - from zeromodels.base import BaseConfig from zeromodels.models.qwen2_5_vl.qwen2_5_vl_config import Qwen2_5VLTextConfig -from zeromodels.models.qwen_image.qwen_image_vae import QwenImageVAEConfig + +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): diff --git a/zeromodels/models/qwen_image/qwen_image_layers.py b/zeromodels/models/qwen_image/qwen_image_layers.py index ff5a7b65..c0c6f3db 100644 --- a/zeromodels/models/qwen_image/qwen_image_layers.py +++ b/zeromodels/models/qwen_image/qwen_image_layers.py @@ -1,18 +1,5 @@ -"""Keras 3 layers for Qwen-Image's double-stream DiT. - -Ports the building blocks of Diffusers ``transformer_qwenimage.py`` (RMSNorm, -timestep projection, 3D RoPE, joint text/image attention, and the dual-stream -transformer block). The full ``QwenImageTransformer2DModel`` lives elsewhere. - -Module paths follow Diffusers so weight conversion can map leaves such as -``transformer_blocks.0.attn.to_q``, ``img_mod.1``, ``img_mlp.net.0.proj``, and -``time_text_embed.timestep_embedder.linear_1``. -""" - from __future__ import annotations -import math - import keras import numpy as np from keras import layers, ops @@ -32,17 +19,6 @@ ROPE_MAX_INDEX = 4096 -def qwen_approximate_gelu(x): - """Tanh-approximate GELU (Diffusers ``gelu-approximate`` / BERT). - - ``0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))``. Prefer - :class:`StableDiffusion3GELUFeedForward` (``ops.gelu(..., approximate=True)``) - for the block MLPs; this helper is for call sites that need the formula - explicitly. - """ - return 0.5 * x * (1.0 + ops.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * ops.power(x, 3)))) - - 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]``. @@ -63,7 +39,6 @@ def apply_rotary_emb_qwen(x, freqs_cis, use_real=True, use_real_unbind_dim=-1): else: cos = ops.repeat(ops.cos(freqs_cis), 2, axis=-1) sin = ops.repeat(ops.sin(freqs_cis), 2, axis=-1) - # Broadcast over batch and heads: [S, D] -> [1, S, 1, D] cos = ops.expand_dims(ops.expand_dims(cos, 0), 2) sin = ops.expand_dims(ops.expand_dims(sin, 0), 2) @@ -177,14 +152,12 @@ def __init__( ) def build(self, timestep_shape): - # Sinusoidal projection is fixed-width; MLP builds from (B, time_freq_dim). 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): - # Diffusers Timesteps(scale=1000): emb = scale * (t[:, None] * freqs) t = ops.cast(timestep, "float32") * self.scale emb = timestep_embedding( t, @@ -239,11 +212,10 @@ def __init__( 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 # complex half-dim == angle width + 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) - # Diffusers: arange(4096).flip(0) * -1 - 1 -> [-4096, ..., -1] 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], @@ -257,7 +229,6 @@ def __init__( self._neg_freqs_np = neg_freqs def build(self, input_shape=None): - # angles reach ~4096 rad: 16-bit storage would wreck the rotation self.pos_freqs = self.add_weight( name="pos_freqs", shape=(ROPE_MAX_INDEX, self.rope_dim), @@ -277,7 +248,6 @@ def build(self, input_shape=None): self.built = True def _split_axes(self, freqs): - # Diffusers ``tensor.split(sizes, dim)``; Keras ``ops.split`` takes indices. h0, h1, h2 = self._axis_halves return ( freqs[:, :h0], @@ -479,12 +449,10 @@ def call( txt_q = apply_rotary_emb_qwen(txt_q, txt_freqs) txt_k = apply_rotary_emb_qwen(txt_k, txt_freqs) - # Joint sequence order: [text, image] (Diffusers Qwen; not SD3). 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) - # fused_attention expects [B, H, S, D] query = ops.transpose(query, (0, 2, 1, 3)) key = ops.transpose(key, (0, 2, 1, 3)) value = ops.transpose(value, (0, 2, 1, 3)) @@ -497,7 +465,6 @@ def call( joint_mask = ops.concatenate( [encoder_hidden_states_mask, img_mask], axis=1 ) - # Additive mask: keep=0, drop=MASK_NEG (Diffusers bool True=keep). keep = ops.cast(joint_mask, "float32") attention_mask = (1.0 - keep) * MASK_NEG attention_mask = attention_mask[:, None, None, :] @@ -574,7 +541,6 @@ def __init__( self.eps = eps self.attn_implementation = attn_implementation - # Diffusers Sequential(SiLU, Linear) -> weight leaf ``img_mod.1`` / ``txt_mod.1``. self.img_mod = layers.Dense( 6 * dim, name=safe_name(f"{module_path}.img_mod.1") ) @@ -630,7 +596,6 @@ def build( ): del encoder_hidden_states_mask_shape, image_rotary_emb_shape if isinstance(hidden_states_shape, (list, tuple)) and temb_shape is None: - # Explicit ``build((img, txt, temb))`` from a parent layer. if ( len(hidden_states_shape) >= 3 and hasattr(hidden_states_shape[0], "__len__") @@ -660,7 +625,6 @@ def build( self.built = True def _modulate(self, x, mod_params): - """``x * (1 + scale) + shift``, returning ``(modulated, gate)``.""" shift, scale, gate = ops.split(mod_params, 3, axis=-1) modulated = x * (1.0 + scale[:, None, :]) + shift[:, None, :] return modulated, gate @@ -721,8 +685,6 @@ def call( return encoder_hidden_states, hidden_states def compute_output_shape(self, input_shape): - # Multi-input builds pass ``[img_shape, txt_shape, temb_shape]``; a lone - # tensor shape is ``(batch, seq, dim)``. if ( isinstance(input_shape, (list, tuple)) and len(input_shape) >= 2 @@ -746,14 +708,9 @@ def get_config(self): ) return config - -# Final-norm AdaLN: Diffusers ``AdaLayerNormContinuous`` uses (scale, shift) -# order with ``num_chunks=2`` — same as StableDiffusion3AdaLayerNorm. QwenImageAdaLayerNormContinuous = StableDiffusion3AdaLayerNorm - __all__ = [ - "qwen_approximate_gelu", "apply_rotary_emb_qwen", "QwenImageRMSNorm", "QwenImageTimestepProjEmbeddings", diff --git a/zeromodels/models/qwen_image/qwen_image_model.py b/zeromodels/models/qwen_image/qwen_image_model.py index bf1165e9..b4c75be1 100644 --- a/zeromodels/models/qwen_image/qwen_image_model.py +++ b/zeromodels/models/qwen_image/qwen_image_model.py @@ -1,5 +1,3 @@ -"""Qwen-Image models: transformer, container, and text-to-image task.""" - from __future__ import annotations import keras @@ -14,9 +12,12 @@ ) 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, @@ -25,7 +26,11 @@ QwenImageTimestepProjEmbeddings, QwenImageTransformerBlock, ) -from zeromodels.models.qwen_image.qwen_image_vae import AutoencoderKLQwenImage +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"}) @@ -36,19 +41,6 @@ ) -def calculate_shift( - image_seq_len, - base_seq_len=256, - max_seq_len=4096, - base_shift=0.5, - max_shift=1.15, -): - """Resolution-dependent flow-match shift (Diffusers ``calculate_shift``).""" - m = (max_shift - base_shift) / (max_seq_len - base_seq_len) - b = base_shift - m * base_seq_len - return image_seq_len * m + b - - 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] @@ -72,6 +64,212 @@ def unpack_latents(latents, height, width, channels): 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``). @@ -594,17 +792,16 @@ def generate( uncond = None do_cfg = False - # Dynamic flow-match timesteps (mu from packed sequence length). + # 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 {} - mu = calculate_shift( - image_seq_len, - sched_cfg.get("base_image_seq_len", 256), - sched_cfg.get("max_image_seq_len", 4096), - sched_cfg.get("base_shift", 0.5), - sched_cfg.get("max_shift", 0.9), - ) + 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: diff --git a/zeromodels/models/qwen_image/qwen_image_tokenizer.py b/zeromodels/models/qwen_image/qwen_image_tokenizer.py index c73f5dfb..ab798e9b 100644 --- a/zeromodels/models/qwen_image/qwen_image_tokenizer.py +++ b/zeromodels/models/qwen_image/qwen_image_tokenizer.py @@ -4,7 +4,6 @@ from zeromodels.models.qwen2.qwen2_tokenizer import Qwen2Tokenizer -# Diffusers ``QwenImagePipeline.prompt_template_encode`` PROMPT_TEMPLATE = ( "<|im_start|>system\nDescribe the image by detailing the color, shape, size, " "texture, quantity, text, spatial relationships of the objects and " @@ -38,7 +37,6 @@ def __init__( max_seq_len=1024, **kwargs, ): - # Diffusers pads to tokenizer_max_length + drop_idx (= 1024 + 34). 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) @@ -49,11 +47,9 @@ def format_prompt(self, text): def call(self, inputs): texts = self.normalize_texts(inputs) templated = [self.format_prompt(t) for t in texts] - # Pad/truncate to max_seq_len + template prefix budget (Diffusers). 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) - # Cap length for the static text-encoder graph when shorter. return {"input_ids": input_ids, "attention_mask": attention_mask} def get_config(self): diff --git a/zeromodels/models/qwen_image/qwen_image_vae.py b/zeromodels/models/qwen_image/qwen_image_vae.py index c822fde6..af227e50 100644 --- a/zeromodels/models/qwen_image/qwen_image_vae.py +++ b/zeromodels/models/qwen_image/qwen_image_vae.py @@ -1,123 +1,11 @@ -"""Keras 3 port of Diffusers ``AutoencoderKLQwenImage`` (Qwen-Image / Wan VAE). - -Internal activations use channels-last 5D ``(B, T, H, W, C)`` (NDHWC) so -``Conv3D`` is portable across Keras backends. Diffusers stores the same stacks -as NCDHW ``Conv3d``; converters must transpose kernels: - -* Conv3D: ``(O, I, T, H, W)`` -> ``(T, H, W, I, O)`` -* Conv2D (resample / attention): ``(O, I, H, W)`` -> ``(H, W, I, O)`` -* RMSNorm ``gamma``: squeeze Diffusers ``(C, 1, 1[, 1])`` -> ``(C,)`` - -v1 targets single-frame T2I (``T=1``): feat-cache video streaming is a no-op and -temporal ``time_conv`` paths inside ``Resample`` are skipped, matching Diffusers' -first-chunk behaviour when the cache is cold. Tiling / slicing are omitted. -Module paths mirror Diffusers for weight conversion. -""" - -from __future__ import annotations - import math import keras from keras import layers, ops -from zeromodels.base import BaseModel from zeromodels.base.base_attention import fused_attention -from zeromodels.base.base_config import BaseConfig from zeromodels.models.stable_diffusion.stable_diffusion_layers import safe_name -# --------------------------------------------------------------------------- -# Defaults (Qwen/Qwen-Image VAE config.json) -# --------------------------------------------------------------------------- - -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 - - -# --------------------------------------------------------------------------- -# Layout helpers -# --------------------------------------------------------------------------- - - -def _as_tuple3(value): - if isinstance(value, int): - return (value, value, value) - value = tuple(value) - if len(value) != 3: - raise ValueError(f"Expected int or length-3 tuple, got {value!r}") - return value - - -def ncdhw_to_ndhwc(x): - """``(B, C, T, H, W)`` -> ``(B, T, H, W, C)``.""" - return ops.transpose(x, (0, 2, 3, 4, 1)) - - -def ndhwc_to_ncdhw(x): - """``(B, T, H, W, C)`` -> ``(B, C, T, H, W)``.""" - return ops.transpose(x, (0, 4, 1, 2, 3)) - - -# --------------------------------------------------------------------------- -# Layers -# --------------------------------------------------------------------------- - @keras.saving.register_keras_serializable(package="zeromodels") class QwenImageRMSNorm(layers.Layer): @@ -142,8 +30,6 @@ def build(self, input_shape): self.built = True def call(self, x): - # Diffusers: F.normalize(x, dim=channels) * sqrt(dim) * gamma, i.e. an L2 - # normalize (eps 1e-12 on the norm) over the last (channel) axis here. dtype = x.dtype x_f = ops.cast(x, "float32") norm = ops.sqrt(ops.sum(ops.square(x_f), axis=-1, keepdims=True)) @@ -189,19 +75,25 @@ def __init__( kwargs.setdefault("name", safe_name(module_path)) super().__init__(**kwargs) self.out_channels = int(out_channels) - self.kernel_size = _as_tuple3(kernel_size) - self.stride = _as_tuple3(stride) - pad = _as_tuple3(padding) + 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 - # Causal: double temporal left pad, zero right 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 Conv3D keeps the Diffusers module path name so conversion's - # last-two-segment mapping yields ``encoder.conv_in.weight`` etc. - # (path ``.../encoder__conv_in/encoder__conv_in/kernel``). leaf_name = safe_name(module_path) if module_path else "conv" self.conv = layers.Conv3D( self.out_channels, @@ -331,7 +223,6 @@ def build(self, input_shape): b, t, h, w, c = input_shape if self.spatial_conv is not None: if self._downsample: - # ZeroPad2d((0,1,0,1)) -> H+1, W+1 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)) @@ -340,17 +231,12 @@ def build(self, input_shape): 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: - # Downsample3d time_conv is stride-2 with k_t=3 and no pad; T=1 cannot - # run it. Build with a synthetic T so kernels exist for conversion; the - # T=1 image path never calls time_conv (cold-cache Diffusers behaviour). 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): - # Optional temporal branch (video streaming); skipped for T=1 image path - # (matches Diffusers' cold feat-cache first chunk, which skips time_conv). if ( self.apply_temporal and self.time_conv is not None @@ -362,7 +248,6 @@ def call(self, x): w = ops.shape(x)[3] c = self.dim x = self.time_conv(x) - # (B, T, H, W, 2C) -> (B, T*2, H, W, C) interleaved like Diffusers 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)) @@ -372,7 +257,6 @@ def call(self, x): h = ops.shape(x)[2] w = ops.shape(x)[3] c = ops.shape(x)[4] - # Merge batch and time for 2D ops: (B*T, H, W, C) x2 = ops.reshape(x, (b * t, h, w, c)) if self._upsample: @@ -713,7 +597,6 @@ def build(self, input_shape): x_shape = list(input_shape) x_shape[-1] = self.dim x_shape = tuple(x_shape) - # Walk blocks updating spatial dims approximately for build. shape = x_shape for layer in self.down_blocks: layer.build(shape) @@ -744,7 +627,6 @@ def compute_output_shape(self, input_shape): shape = layer.compute_output_shape(shape) out = list(shape) out[-1] = self.z_dim - # Spatial / temporal sizes: 3 spatial downsamples -> /8; T unchanged (T=1 path). b, t, h, w, _ = input_shape factor = 2 ** (len(self.dim_mult) - 1) return ( @@ -979,217 +861,3 @@ def get_config(self): } ) return config - - -# --------------------------------------------------------------------------- -# Model -# --------------------------------------------------------------------------- - - -@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" - ) - - # Functional graph: HWC image + HWC-ish latent (T squeezed to channels-last 4D - # with an explicit time axis of 1 in the 5D path). - 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) # (B, 1, H, W, C) - moments_5d = quant_conv(encoder(image_5d)) - moments = ops.squeeze(moments_5d, axis=1) # (B, h, w, 2*z) - - 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 - - # -- public encode / decode ------------------------------------------------ - - 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: - # (B, H, W, C) channels-last image or latent - return ops.expand_dims(x, axis=1) - if static_ndim == 5: - # Detect NCDHW (Diffusers): channel axis at 1 matches input_channels / z_dim. - 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 ncdhw_to_ndhwc(x) - 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 ndhwc_to_ncdhw(z) - 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 ndhwc_to_ncdhw(x) - 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, - ) From 0a9d5b28df92ba3de4025a40ef87f6487a89722c Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Tue, 22 Sep 2026 23:13:32 -0700 Subject: [PATCH 7/8] Update docs --- README.md | 3 +- docs/getting_started.md | 5 +- docs/index.md | 5 +- docs/loading_weights.md | 2 +- docs/models.md | 1 + docs/qwen_image.md | 374 +++++++++++++++++++++++++++++++++++----- 6 files changed, 342 insertions(+), 48 deletions(-) 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 index b7908921..ac9147aa 100644 --- a/docs/qwen_image.md +++ b/docs/qwen_image.md @@ -1,61 +1,351 @@ # Qwen-Image
-Weights: pretrained Keras weights will be hosted on Hugging Face under -zeromodels/qwen-image -(each repo carries zm_config.json + model.weights.h5 + -tokenizer.json). Load with from_weights("zeromodels/qwen-image"). -Conversion from Diffusers is offline only — on-the-fly hf: is not supported. +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 text-to-image, ported to pure Keras 3 from -[Diffusers `QwenImagePipeline`](https://huggingface.co/docs/diffusers/api/pipelines/qwenimage) -(`/Qwen/Qwen-Image`). Latent flow-matching with: +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`. -- **Denoiser**: 60-layer double-stream MMDiT (`QwenImageTransformer2DModel`) over - packed 2×2 latent patches (64-d tokens) with MS-RoPE -- **VAE**: Wan-derived `AutoencoderKLQwenImage` (16 latent channels, 8× spatial) -- **Text encoder**: Qwen2.5-VL-7B Instruct text tower (ChatML prompt template) -- **Scheduler**: `FlowMatchEulerDiscreteScheduler` with dynamic resolution shifting +The weights are converted once, offline, and hosted: on-the-fly `hf:` conversion +is deliberately **not supported** for diffusion models. -`QwenImageModel` is the hosted container (transformer + VAE + text tower), -stored as one set of sharded weights (`model.weights.json`, ~53 GiB at 16-bit). -`QwenImageTextToImage` adds `generate` via `BaseDiffusion`, with **true CFG** -(separate cond/uncond forwards and prediction-norm renormalization, Diffusers -`true_cfg_scale`). +Key facts of the port: -## Status +- **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)). -| Piece | Status | +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 | |---|---| -| Configs / tokenizer / auto registry | Done | -| Transformer layers + model | Done (graph) | -| VAE | Done (T=1 image path) | -| Text-to-image `generate` + packed latents | Done | -| Flow-match dynamic shifting | Done | -| Diffusers → Keras weight transfer | In progress (`convert_qwen_image_diffusers_to_keras.py`) | -| Hosted `zeromodels/qwen-image` weights | Pending conversion | +| `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 -## API sketch +### `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 -from zeromodels.models.qwen_image import QwenImageTextToImage, QwenImageTokenizer +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") -tok = QwenImageTokenizer.from_weights("zeromodels/qwen-image") -image = model.generate( - **tok("a coffee shop entrance with a chalkboard sign"), - height=1024, - width=1024, - num_inference_steps=50, - guidance_scale=4.0, +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 ``` -## Variants +### Negative prompt and batching -| Variant | Hub (planned) | Source | -|---|---|---| -| `qwen-image` | `zeromodels/qwen-image` | [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) | -Paper / model card: [Qwen-Image](https://huggingface.co/Qwen/Qwen-Image). -License: Apache-2.0. +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. From c7f81ccef05d31ea4e4caef89ef4581063ae069b Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Tue, 22 Sep 2026 23:13:52 -0700 Subject: [PATCH 8/8] Format --- docs/qwen_image.md | 7 ++- .../convert_qwen_image_diffusers_to_keras.py | 7 +-- .../models/qwen_image/qwen_image_layers.py | 47 +++++++------------ .../models/qwen_image/qwen_image_model.py | 32 ++++--------- .../models/qwen_image/qwen_image_vae.py | 24 ++++------ 5 files changed, 41 insertions(+), 76 deletions(-) diff --git a/docs/qwen_image.md b/docs/qwen_image.md index ac9147aa..155e9d9e 100644 --- a/docs/qwen_image.md +++ b/docs/qwen_image.md @@ -224,7 +224,12 @@ 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 + **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 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 index 5d7a7d73..07f52b09 100644 --- a/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py +++ b/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py @@ -199,9 +199,7 @@ def __len__(self): state = _State() else: - path = hf_hub_download( - repo, filename, subfolder=subfolder, token=token - ) + path = hf_hub_download(repo, filename, subfolder=subfolder, token=token) state = {} with safe_open(path, framework="np") as shard: for key in shard.keys(): @@ -298,8 +296,7 @@ def __len__(self): unused = sorted(set(hf_keys) - consumed) if unused: raise ValueError( - f"text_encoder: {len(unused)} checkpoint tensors unused, " - f"e.g. {unused[:5]}." + f"text_encoder: {len(unused)} checkpoint tensors unused, e.g. {unused[:5]}." ) del weight_map, shard_paths, hf_keys gc.collect() diff --git a/zeromodels/models/qwen_image/qwen_image_layers.py b/zeromodels/models/qwen_image/qwen_image_layers.py index c0c6f3db..c3e98b0a 100644 --- a/zeromodels/models/qwen_image/qwen_image_layers.py +++ b/zeromodels/models/qwen_image/qwen_image_layers.py @@ -47,9 +47,7 @@ def apply_rotary_emb_qwen(x, freqs_cis, use_real=True, use_real_unbind_dim=-1): 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) - ) + 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, :] @@ -280,9 +278,7 @@ def _video_freqs(self, frame, height, width, idx=0): 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.reshape(freqs_height, (1, height, 1, self._axis_halves[1])) freqs_height = ops.broadcast_to( freqs_height, (frame, height, width, self._axis_halves[1]) ) @@ -291,9 +287,7 @@ def _video_freqs(self, frame, height, width, idx=0): freqs_width, (frame, height, width, self._axis_halves[2]) ) - freqs = ops.concatenate( - [freqs_frame, freqs_height, freqs_width], axis=-1 - ) + 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): @@ -461,10 +455,10 @@ def call( 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 + 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, :] @@ -541,12 +535,8 @@ def __init__( 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_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, @@ -596,9 +586,8 @@ def build( ): 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__") + if len(hidden_states_shape) >= 3 and hasattr( + hidden_states_shape[0], "__len__" ): ( hidden_states_shape, @@ -656,9 +645,7 @@ def call( 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 - ) + encoder_hidden_states = encoder_hidden_states + txt_gate1[:, None, :] * txt_attn img_modulated2, img_gate2 = self._modulate( self.img_norm2(hidden_states), img_mod2 @@ -670,15 +657,12 @@ def call( 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) - ) + 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 - ) + 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) @@ -708,6 +692,7 @@ def get_config(self): ) return config + QwenImageAdaLayerNormContinuous = StableDiffusion3AdaLayerNorm __all__ = [ diff --git a/zeromodels/models/qwen_image/qwen_image_model.py b/zeromodels/models/qwen_image/qwen_image_model.py index b4c75be1..4ff2d36d 100644 --- a/zeromodels/models/qwen_image/qwen_image_model.py +++ b/zeromodels/models/qwen_image/qwen_image_model.py @@ -45,21 +45,15 @@ 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.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) - ) + 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.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)) @@ -307,9 +301,7 @@ def __init__( 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 + 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 @@ -682,9 +674,7 @@ def predict_noise(self, latents, timesteps, embeddings): "sample": latents, "timestep": timesteps, "encoder_hidden_states": embeddings["encoder_hidden_states"], - "encoder_hidden_states_mask": embeddings[ - "encoder_hidden_states_mask" - ], + "encoder_hidden_states_mask": embeddings["encoder_hidden_states_mask"], } )["sample"] @@ -706,9 +696,7 @@ def prepare_latents(self, batch, seed=None, latents=None, dtype="float32"): 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 - ) + 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) @@ -739,9 +727,7 @@ def denoise( 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) - ) + 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 @@ -784,9 +770,7 @@ def generate( 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 - ) + uncond = self.encode_prompt(negative_input_ids, negative_attention_mask) do_cfg = guidance_scale > 1.0 else: uncond = None diff --git a/zeromodels/models/qwen_image/qwen_image_vae.py b/zeromodels/models/qwen_image/qwen_image_vae.py index af227e50..c79dfbd9 100644 --- a/zeromodels/models/qwen_image/qwen_image_vae.py +++ b/zeromodels/models/qwen_image/qwen_image_vae.py @@ -84,9 +84,7 @@ def __init__( (stride, stride, stride) if isinstance(stride, int) else tuple(stride) ) pad = ( - (padding, padding, padding) - if isinstance(padding, int) - else tuple(padding) + (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 @@ -106,11 +104,7 @@ def __init__( 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 - ) + 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) @@ -337,7 +331,9 @@ def __init__( 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.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" ) @@ -397,7 +393,9 @@ def __init__(self, dim, module_path, **kwargs): 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.norm = QwenImageRMSNorm( + dim, images=True, module_path=f"{module_path}.norm" + ) self.to_qkv = layers.Conv2D( dim * 3, 1, @@ -565,11 +563,7 @@ def __init__( idx += 1 in_dim = out_dim if i != len(self.dim_mult) - 1: - mode = ( - "downsample3d" - if self.temperal_downsample[i] - else "downsample2d" - ) + mode = "downsample3d" if self.temperal_downsample[i] else "downsample2d" self.down_blocks.append( QwenImageResample( out_dim,