[Quantization] feat: add qwen-image gguf adapter (#1755)

Signed-off-by: David Chen <530634352@qq.com>

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`.

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

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"

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",
 ]

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)

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]