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

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

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

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+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]
+
+
+class _DummyPipelineModel(nn.Module):
+    def __init__(self, *, source_prefix: str):
+        super().__init__()
+        self.transformer = nn.Linear(2, 2, bias=False)
+        self.vae = nn.Linear(2, 2, bias=False)
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path="dummy",
+                subfolder="transformer",
+                revision=None,
+                prefix=source_prefix,
+                fall_back_to_pt=True,
+            )
+        ]
+
+    def load_weights(self, weights):
+        params = dict(self.named_parameters())
+        loaded: set[str] = set()
+        for name, tensor in weights:
+            if name not in params:
+                continue
+            params[name].data.copy_(tensor.to(dtype=params[name].dtype))
+            loaded.add(name)
+        return loaded
+
+
+def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader:
+    loader = object.__new__(DiffusersPipelineLoader)
+    loader.counter_before_loading_weights = 0.0
+    loader.counter_after_loading_weights = 0.0
+
+    def _iter_weights(_model):
+        for name in weight_names:
+            yield name, torch.zeros((2, 2))
+
+    loader.get_all_weights = _iter_weights  # type: ignore[assignment]
+    return loader
+
+
+def test_strict_check_only_validates_source_prefix_parameters():
+    model = _DummyPipelineModel(source_prefix="transformer.")
+    loader = _make_loader_with_weights(["transformer.weight"])
+
+    # Should not require VAE parameters because they are outside weights_sources.
+    loader.load_weights(model)
+
+
+def test_strict_check_raises_when_source_parameters_are_missing():
+    model = _DummyPipelineModel(source_prefix="transformer.")
+    loader = _make_loader_with_weights([])
+
+    with pytest.raises(ValueError, match="transformer.weight"):
+        loader.load_weights(model)
+
+
+def test_empty_source_prefix_keeps_full_model_strict_check():
+    model = _DummyPipelineModel(source_prefix="")
+    loader = _make_loader_with_weights(["transformer.weight"])
+
+    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"