diff --git a/docs/user_guide/diffusion/quantization/gguf.md b/docs/user_guide/diffusion/quantization/gguf.md index c3d5eb7eb..04e2d6aad 100644 --- a/docs/user_guide/diffusion/quantization/gguf.md +++ b/docs/user_guide/diffusion/quantization/gguf.md @@ -6,7 +6,7 @@ 3. Keep user-facing knobs minimal and consistent across offline and online flows. ## Scope -1. Models: Z-Image, and Flux2-klein. +1. Models: Z-Image, Flux2-klein, and Qwen-Image family. 2. Components: diffusion transformer weights, loader paths, and quantization configs. 3. Modes: native GGUF (transformer-only weights). @@ -124,19 +124,21 @@ x @ weight.T 1. `DiffusersPipelineLoader.load_model` detects `quantization_config.method == "gguf"`. 2. `gguf_model` is resolved as one of: local file, `repo/file.gguf`, or `repo:quant_type`. 3. GGUF weights are routed through adapters in `vllm_omni/diffusion/model_loader/gguf_adapters/`. -4. Name mapping is applied per-architecture (Z-Image, Flux2Klein). +4. Name mapping is applied per-architecture (Z-Image, Flux2Klein, Qwen-Image). 5. GGUF weights are loaded into transformer modules, remaining non-transformer weights come from the HF checkpoint. ## GGUF Adapter Design 1. `GGUFAdapter` is an abstract base class for model-specific adapters. 2. `Flux2KleinGGUFAdapter` implements Flux2-Klein remapping + qkv split + adaLN swap. 3. `ZImageGGUFAdapter` implements Z-Image qkv + ffn shard handling and linear qweight routing. -4. `get_gguf_adapter(...)` strictly selects by model class/config; unsupported models raise an error (no fallback adapter). +4. `QwenImageGGUFAdapter` implements Qwen-Image fused attention projection remapping. +5. `get_gguf_adapter(...)` strictly selects by model class/config; unsupported models raise an error (no fallback adapter). Adapter paths: - Base: `vllm_omni/diffusion/model_loader/gguf_adapters/base.py` - Z-Image: `vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py` - Flux2-Klein: `vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py` +- Qwen-Image: `vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py` ## User Usage (Offline) @@ -195,3 +197,31 @@ curl -X POST http://localhost:8000/v1/images/generations \ "num_inference_steps": 4 }' ``` + +## Qwen-Image Example + +### Offline +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model Qwen/Qwen-Image \ + --gguf-model QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf \ + --quantization gguf \ + --prompt "a red paper kite hanging from a pine tree in a winter courtyard" \ + --height 1024 \ + --width 1024 \ + --seed 42 \ + --num_inference_steps 20 \ + --output outputs/qwen_image_q4km.png +``` + +### Online +```bash +vllm serve Qwen/Qwen-Image \ + --omni \ + --port 8000 \ + --quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}' +``` + +Notes: +1. The GGUF repo provides transformer weights only; the base model repo still provides tokenizer, text encoder, scheduler, and VAE. +2. The same adapter path covers `QwenImagePipeline`, `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, and `QwenImageLayeredPipeline`. diff --git a/tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py b/tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py new file mode 100644 index 000000000..6b7392c80 --- /dev/null +++ b/tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter +from vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image import ( + QwenImageGGUFAdapter, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _make_od_config( + model_class_name: str = "QwenImagePipeline", + model_type: str = "qwen_image", +): + return SimpleNamespace( + model_class_name=model_class_name, + tf_model_config={"model_type": model_type}, + ) + + +def _make_source(prefix: str = "", subfolder: str = "transformer"): + return SimpleNamespace(prefix=prefix, subfolder=subfolder) + + +class _FakeTransformer: + def __init__(self) -> None: + self._params = { + "img_in.qweight": 1, + "img_in.qweight_type": 1, + "txt_in.qweight": 1, + "txt_in.qweight_type": 1, + "proj_out.qweight": 1, + "proj_out.qweight_type": 1, + "time_text_embed.timestep_embedder.linear_1.qweight": 1, + "time_text_embed.timestep_embedder.linear_1.qweight_type": 1, + "time_text_embed.timestep_embedder.linear_2.qweight": 1, + "time_text_embed.timestep_embedder.linear_2.qweight_type": 1, + "transformer_blocks.0.img_mod.1.qweight": 1, + "transformer_blocks.0.img_mod.1.qweight_type": 1, + "transformer_blocks.0.txt_mod.1.qweight": 1, + "transformer_blocks.0.txt_mod.1.qweight_type": 1, + "transformer_blocks.0.attn.to_qkv.qweight": 1, + "transformer_blocks.0.attn.to_qkv.qweight_type": 1, + "transformer_blocks.0.attn.add_kv_proj.qweight": 1, + "transformer_blocks.0.attn.add_kv_proj.qweight_type": 1, + "transformer_blocks.0.attn.to_out.weight": 1, + "transformer_blocks.0.img_mlp.net.0.proj.qweight": 1, + "transformer_blocks.0.img_mlp.net.0.proj.qweight_type": 1, + } + + def named_parameters(self): + return list(self._params.items()) + + def named_buffers(self): + return [] + + def get_submodule(self, _prefix: str): + return self + + +def test_qwen_adapter_selected_for_qwen_image_family(): + adapter = get_gguf_adapter( + "dummy.gguf", + _FakeTransformer(), + _make_source(), + _make_od_config(), + ) + assert isinstance(adapter, QwenImageGGUFAdapter) + + +def test_qwen_adapter_matches_multiple_pipeline_variants(): + for model_class_name in ( + "QwenImagePipeline", + "QwenImageEditPipeline", + "QwenImageEditPlusPipeline", + "QwenImageLayeredPipeline", + ): + assert QwenImageGGUFAdapter.is_compatible( + _make_od_config(model_class_name=model_class_name), + _FakeTransformer(), + _make_source(), + ) + + +def test_qwen_adapter_preserves_split_projection_names(monkeypatch: pytest.MonkeyPatch): + import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module + + monkeypatch.setattr( + qwen_image_module, + "gguf_quant_weights_iterator", + lambda _path: iter( + [ + ("transformer_blocks.0.attn.to_q.qweight_type", 1), + ("transformer_blocks.0.attn.to_q.qweight", 2), + ("transformer_blocks.0.attn.to_k.qweight_type", 3), + ("transformer_blocks.0.attn.to_k.qweight", 4), + ("transformer_blocks.0.attn.to_out.0.qweight_type", 5), + ("transformer_blocks.0.attn.to_out.0.qweight", 6), + ] + ), + ) + + adapter = QwenImageGGUFAdapter( + "dummy.gguf", + _FakeTransformer(), + _make_source(), + _make_od_config(), + ) + + weights = list(adapter.weights_iterator()) + + assert ("transformer_blocks.0.attn.to_q.qweight_type", 1) in weights + assert ("transformer_blocks.0.attn.to_q.qweight", 2) in weights + assert ("transformer_blocks.0.attn.to_k.qweight_type", 3) in weights + assert ("transformer_blocks.0.attn.to_k.qweight", 4) in weights + assert ("transformer_blocks.0.attn.to_out.0.qweight_type", 5) in weights + assert ("transformer_blocks.0.attn.to_out.0.qweight", 6) in weights + + +def test_qwen_adapter_keeps_top_level_quantized_weights(monkeypatch: pytest.MonkeyPatch): + import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module + + monkeypatch.setattr( + qwen_image_module, + "gguf_quant_weights_iterator", + lambda _path: iter( + [ + ("img_in.qweight_type", 1), + ("img_in.qweight", 2), + ("time_text_embed.timestep_embedder.linear_1.qweight_type", 3), + ("time_text_embed.timestep_embedder.linear_1.qweight", 4), + ] + ), + ) + + adapter = QwenImageGGUFAdapter( + "dummy.gguf", + _FakeTransformer(), + _make_source(), + _make_od_config(), + ) + + weights = list(adapter.weights_iterator()) + + assert ("img_in.qweight_type", 1) in weights + assert ("img_in.qweight", 2) in weights + assert ("time_text_embed.timestep_embedder.linear_1.qweight_type", 3) in weights + assert ("time_text_embed.timestep_embedder.linear_1.qweight", 4) in weights + + +def test_qwen_adapter_keeps_modulation_quantized_weights(monkeypatch: pytest.MonkeyPatch): + import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module + + monkeypatch.setattr( + qwen_image_module, + "gguf_quant_weights_iterator", + lambda _path: iter( + [ + ("transformer_blocks.0.img_mod.1.qweight_type", 1), + ("transformer_blocks.0.img_mod.1.qweight", 2), + ("transformer_blocks.0.txt_mod.1.qweight_type", 3), + ("transformer_blocks.0.txt_mod.1.qweight", 4), + ] + ), + ) + + adapter = QwenImageGGUFAdapter( + "dummy.gguf", + _FakeTransformer(), + _make_source(), + _make_od_config(), + ) + + weights = list(adapter.weights_iterator()) + + assert ("transformer_blocks.0.img_mod.1.qweight_type", 1) in weights + assert ("transformer_blocks.0.img_mod.1.qweight", 2) in weights + assert ("transformer_blocks.0.txt_mod.1.qweight_type", 3) in weights + assert ("transformer_blocks.0.txt_mod.1.qweight", 4) in weights diff --git a/tests/diffusion/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py similarity index 79% rename from tests/diffusion/test_diffusers_loader.py rename to tests/diffusion/model_loader/test_diffusers_loader.py index 3f6396027..1b13c400f 100644 --- a/tests/diffusion/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -6,6 +6,7 @@ import torch import torch.nn as nn from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -71,3 +72,24 @@ def test_empty_source_prefix_keeps_full_model_strict_check(): with pytest.raises(ValueError, match="vae.weight"): loader.load_weights(model) + + +def test_qwen_model_class_selects_qwen_gguf_adapter(): + od_config = type( + "Config", + (), + { + "model_class_name": "QwenImagePipeline", + "tf_model_config": {"model_type": "qwen_image"}, + }, + )() + source = DiffusersPipelineLoader.ComponentSource( + model_or_path="dummy", + subfolder="transformer", + revision=None, + prefix="transformer.", + ) + + adapter = get_gguf_adapter("dummy.gguf", object(), source, od_config) + + assert adapter.__class__.__name__ == "QwenImageGGUFAdapter" diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py b/vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py index 416ebc7a8..3266ba06c 100644 --- a/vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py +++ b/vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py @@ -7,6 +7,7 @@ import torch from .base import GGUFAdapter from .flux2_klein import Flux2KleinGGUFAdapter +from .qwen_image import QwenImageGGUFAdapter from .z_image import ZImageGGUFAdapter if TYPE_CHECKING: @@ -22,7 +23,7 @@ def get_gguf_adapter( source: DiffusersPipelineLoader.ComponentSource, od_config: OmniDiffusionConfig, ) -> GGUFAdapter: - adapter_classes = (ZImageGGUFAdapter, Flux2KleinGGUFAdapter) + adapter_classes = (QwenImageGGUFAdapter, ZImageGGUFAdapter, Flux2KleinGGUFAdapter) for adapter_cls in adapter_classes: if adapter_cls.is_compatible(od_config, model, source): return adapter_cls(gguf_file, model, source, od_config) @@ -40,6 +41,7 @@ def get_gguf_adapter( __all__ = [ "GGUFAdapter", "Flux2KleinGGUFAdapter", + "QwenImageGGUFAdapter", "ZImageGGUFAdapter", "get_gguf_adapter", ] diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py b/vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py new file mode 100644 index 000000000..cc8ccea5c --- /dev/null +++ b/vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from collections.abc import Iterable + +import torch + +from .base import GGUFAdapter, gguf_quant_weights_iterator + + +class QwenImageGGUFAdapter(GGUFAdapter): + """GGUF adapter for the Qwen-Image transformer family.""" + + @staticmethod + def is_compatible(od_config, model: torch.nn.Module, source) -> bool: + model_class = od_config.model_class_name or "" + if model_class.startswith("QwenImage"): + return True + cfg = od_config.tf_model_config + if cfg is not None: + model_type = str(cfg.get("model_type", "")).lower() + if model_type.startswith("qwen_image"): + return True + return False + + def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]: + yield from gguf_quant_weights_iterator(self.gguf_file) diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index 8da6ffdb2..3e9a0f0f3 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -21,6 +21,7 @@ from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, QKVParallelLinear, + ReplicatedLinear, RowParallelLinear, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader @@ -60,7 +61,7 @@ class ImageRopePrepare(nn.Module): Note: Our _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism). """ - def __init__(self, img_in: nn.Linear, pos_embed: nn.Module): + def __init__(self, img_in: nn.Module, pos_embed: nn.Module): super().__init__() self.img_in = img_in self.pos_embed = pos_embed @@ -157,11 +158,32 @@ class ModulateIndexPrepare(nn.Module): class QwenTimestepProjEmbeddings(nn.Module): - def __init__(self, embedding_dim, use_additional_t_cond=False): + def __init__( + self, + embedding_dim, + use_additional_t_cond: bool = False, + quant_config: QuantizationConfig | None = None, + ): super().__init__() self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000) self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + self.timestep_embedder.linear_1 = ReplicatedLinear( + 256, + embedding_dim, + bias=True, + return_bias=False, + quant_config=quant_config, + prefix="timestep_embedder.linear_1", + ) + self.timestep_embedder.linear_2 = ReplicatedLinear( + embedding_dim, + embedding_dim, + bias=True, + return_bias=False, + quant_config=quant_config, + prefix="timestep_embedder.linear_2", + ) self.use_additional_t_cond = use_additional_t_cond if use_additional_t_cond: self.addition_t_embedding = nn.Embedding(2, embedding_dim) @@ -681,7 +703,14 @@ class QwenImageTransformerBlock(nn.Module): # Image processing modules self.img_mod = nn.Sequential( nn.SiLU(), - nn.Linear(dim, 6 * dim, bias=True), # For scale, shift, gate for norm1 and norm2 + ReplicatedLinear( + dim, + 6 * dim, + bias=True, + return_bias=False, + quant_config=quant_config, + prefix="img_mod.1", + ), ) self.img_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) self.attn = QwenImageCrossAttention( @@ -698,7 +727,14 @@ class QwenImageTransformerBlock(nn.Module): # Text processing modules self.txt_mod = nn.Sequential( nn.SiLU(), - nn.Linear(dim, 6 * dim, bias=True), # For scale, shift, gate for norm1 and norm2 + ReplicatedLinear( + dim, + 6 * dim, + bias=True, + return_bias=False, + quant_config=quant_config, + prefix="txt_mod.1", + ), ) self.txt_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) # Text doesn't need separate attention - it's handled by img_attn joint computation @@ -912,13 +948,29 @@ class QwenImageTransformer2DModel(CachedTransformer): self.pos_embed = QwenEmbedLayer3DRope(theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True) self.time_text_embed = QwenTimestepProjEmbeddings( - embedding_dim=self.inner_dim, use_additional_t_cond=use_additional_t_cond + embedding_dim=self.inner_dim, + use_additional_t_cond=use_additional_t_cond, + quant_config=quant_config, ) self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6) - self.img_in = nn.Linear(in_channels, self.inner_dim) - self.txt_in = nn.Linear(joint_attention_dim, self.inner_dim) + self.img_in = ReplicatedLinear( + in_channels, + self.inner_dim, + bias=True, + return_bias=False, + quant_config=quant_config, + prefix="img_in", + ) + self.txt_in = ReplicatedLinear( + joint_attention_dim, + self.inner_dim, + bias=True, + return_bias=False, + quant_config=quant_config, + prefix="txt_in", + ) self.transformer_blocks = nn.ModuleList( [ @@ -934,7 +986,22 @@ class QwenImageTransformer2DModel(CachedTransformer): ) self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6) - self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True) + self.norm_out.linear = ReplicatedLinear( + self.inner_dim, + 2 * self.inner_dim, + bias=True, + return_bias=False, + quant_config=quant_config, + prefix="norm_out.linear", + ) + self.proj_out = ReplicatedLinear( + self.inner_dim, + patch_size * patch_size * self.out_channels, + bias=True, + return_bias=False, + quant_config=quant_config, + prefix="proj_out", + ) self.gradient_checkpointing = False self.zero_cond_t = zero_cond_t @@ -1097,7 +1164,7 @@ class QwenImageTransformer2DModel(CachedTransformer): original_name = name lookup_name = name for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in original_name: + if weight_name not in original_name or param_name in original_name: continue lookup_name = original_name.replace(weight_name, param_name) param = params_dict[lookup_name]