[Bugfix] fix gguf TypeError: GGUFConfig.get_name() missing 1 required positional argument: 'self', test: add diffusion gguf unit coverage (#1865)

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

tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the Flux2-Klein GGUF adapter."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
+from vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein import (
+    Flux2KleinGGUFAdapter,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _make_od_config(
+    model_class_name: str = "Flux2Pipeline",
+    model_type: str = "flux",
+):
+    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)
+
+
+def test_flux2_klein_adapter_selected_for_flux_family():
+    adapter = get_gguf_adapter(
+        "dummy.gguf",
+        object(),
+        _make_source(),
+        _make_od_config(),
+    )
+
+    assert isinstance(adapter, Flux2KleinGGUFAdapter)
+
+
+def test_flux2_klein_adapter_matches_flux_model_type():
+    assert Flux2KleinGGUFAdapter.is_compatible(
+        _make_od_config(model_class_name="OtherPipeline", model_type="flux-dev"),
+        object(),
+        _make_source(),
+    )
+
+
+def test_flux2_klein_adapter_renames_core_projection_paths(monkeypatch: pytest.MonkeyPatch):
+    import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module
+
+    monkeypatch.setattr(
+        flux_module,
+        "gguf_quant_weights_iterator",
+        lambda _path: iter(
+            [
+                ("img_in.qweight", torch.ones((1, 1))),
+                ("time_in.in_layer.qweight_type", torch.tensor(1)),
+                ("double_blocks.0.img_attn.norm.query_norm.weight", torch.full((1, 1), 2.0)),
+                ("final_layer.linear.qweight", torch.full((1, 1), 3.0)),
+            ]
+        ),
+    )
+
+    adapter = Flux2KleinGGUFAdapter(
+        "dummy.gguf",
+        object(),
+        _make_source(),
+        _make_od_config(),
+    )
+
+    weights = list(adapter.weights_iterator())
+    names = [name for name, _ in weights]
+
+    assert "x_embedder.qweight" in names
+    assert "time_guidance_embed.timestep_embedder.linear_1.qweight_type" in names
+    assert "transformer_blocks.0.attn.norm_q.weight" in names
+    assert "proj_out.qweight" in names
+
+
+def test_flux2_klein_adapter_swaps_final_adaln_shift_and_scale(monkeypatch: pytest.MonkeyPatch):
+    import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module
+
+    monkeypatch.setattr(
+        flux_module,
+        "gguf_quant_weights_iterator",
+        lambda _path: iter(
+            [
+                (
+                    "final_layer.adaLN_modulation.1.scale",
+                    torch.tensor([1.0, 2.0, 3.0, 4.0]),
+                )
+            ]
+        ),
+    )
+
+    adapter = Flux2KleinGGUFAdapter(
+        "dummy.gguf",
+        object(),
+        _make_source(),
+        _make_od_config(),
+    )
+
+    weights = list(adapter.weights_iterator())
+
+    assert len(weights) == 1
+    assert weights[0][0] == "norm_out.linear.weight"
+    assert torch.equal(weights[0][1], torch.tensor([3.0, 4.0, 1.0, 2.0]))

tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py

@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the Z-Image GGUF adapter."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
+from vllm_omni.diffusion.model_loader.gguf_adapters.z_image import (
+    ZImageGGUFAdapter,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _make_od_config(
+    model_class_name: str = "ZImagePipeline",
+    model_type: str = "z_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)
+
+
+def test_z_image_adapter_selected_for_z_image_family():
+    adapter = get_gguf_adapter(
+        "dummy.gguf",
+        object(),
+        _make_source(),
+        _make_od_config(),
+    )
+
+    assert isinstance(adapter, ZImageGGUFAdapter)
+
+
+def test_z_image_adapter_matches_model_type_variants():
+    for model_type in ("z_image", "zimage", "z-image"):
+        assert ZImageGGUFAdapter.is_compatible(
+            _make_od_config(model_class_name="OtherPipeline", model_type=model_type),
+            object(),
+            _make_source(),
+        )
+
+
+def test_z_image_adapter_renames_known_gguf_tensor_paths(monkeypatch: pytest.MonkeyPatch):
+    import vllm_omni.diffusion.model_loader.gguf_adapters.z_image as z_image_module
+
+    monkeypatch.setattr(
+        z_image_module,
+        "gguf_quant_weights_iterator",
+        lambda _path: iter(
+            [
+                ("model.diffusion_model.final_layer.qweight", torch.ones((1, 1))),
+                ("model.diffusion_model.x_embedder.qweight_type", torch.tensor(1)),
+                ("transformer_blocks.0.attention.out.weight", torch.full((1, 1), 2.0)),
+                ("transformer_blocks.0.attention.qkv.qweight", torch.full((1, 1), 3.0)),
+            ]
+        ),
+    )
+
+    adapter = ZImageGGUFAdapter(
+        "dummy.gguf",
+        object(),
+        _make_source(),
+        _make_od_config(),
+    )
+
+    weights = list(adapter.weights_iterator())
+    names = [name for name, _ in weights]
+
+    assert "all_final_layer.2-1.qweight" in names
+    assert "all_x_embedder.2-1.qweight_type" in names
+    assert "transformer_blocks.0.attention.to_out.0.weight" in names
+    assert "transformer_blocks.0.attention.to_qkv.qweight" in names

tests/diffusion/model_loader/test_diffusers_loader_gguf.py

@@ -0,0 +1,268 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for GGUF-specific DiffusersPipelineLoader behavior."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+import vllm_omni.diffusion.model_loader.diffusers_loader as loader_module
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _DummyModel(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.transformer = nn.Linear(2, 2, bias=True)
+        self.vae = nn.Linear(2, 2, bias=False)
+        self.register_buffer("transformer_buffer", torch.ones(1))
+        self.calls: list[list[str]] = []
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path="dummy",
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            ),
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path="dummy",
+                subfolder="vae",
+                revision=None,
+                prefix="vae.",
+                fall_back_to_pt=True,
+            ),
+        ]
+
+    def load_weights(self, weights):
+        loadable = dict(self.named_parameters())
+        loadable.update(dict(self.named_buffers()))
+        seen: list[str] = []
+        loaded: set[str] = set()
+        for name, tensor in weights:
+            seen.append(name)
+            if name in loadable:
+                target = loadable[name]
+                target.data.copy_(tensor.to(dtype=target.dtype))
+                loaded.add(name)
+        self.calls.append(seen)
+        return loaded
+
+
+def _make_loader() -> DiffusersPipelineLoader:
+    loader = object.__new__(DiffusersPipelineLoader)
+    loader.load_config = SimpleNamespace(
+        download_dir="cache-dir",
+        ignore_patterns=["*.tmp"],
+    )
+    loader.od_config = None
+    loader.counter_before_loading_weights = 0.0
+    loader.counter_after_loading_weights = 0.0
+    return loader
+
+
+def test_is_gguf_quantization_accepts_dict_config():
+    loader = _make_loader()
+    od_config = SimpleNamespace(
+        quantization_config={
+            "method": "gguf",
+            "gguf_model": "weights.gguf",
+        }
+    )
+
+    assert loader._is_gguf_quantization(od_config) is True
+
+
+def test_is_gguf_quantization_rejects_non_gguf_dict_config():
+    loader = _make_loader()
+    od_config = SimpleNamespace(
+        quantization_config={
+            "method": "fp8",
+            "gguf_model": "weights.gguf",
+        }
+    )
+
+    assert loader._is_gguf_quantization(od_config) is False
+
+
+def test_is_gguf_quantization_requires_gguf_model_for_dict_config():
+    loader = _make_loader()
+    od_config = SimpleNamespace(quantization_config={"method": "gguf"})
+
+    with pytest.raises(ValueError, match="gguf_model"):
+        loader._is_gguf_quantization(od_config)
+
+
+def test_is_gguf_quantization_accepts_object_config():
+    loader = _make_loader()
+    od_config = SimpleNamespace(
+        quantization_config=SimpleNamespace(
+            get_name=lambda: "gguf",
+            gguf_model="weights.gguf",
+        )
+    )
+
+    assert loader._is_gguf_quantization(od_config) is True
+
+
+def test_is_gguf_quantization_uses_fallback_object_without_get_name():
+    loader = _make_loader()
+    od_config = SimpleNamespace(quantization_config=SimpleNamespace(gguf_model="weights.gguf"))
+
+    assert loader._is_gguf_quantization(od_config) is True
+
+
+def test_get_model_loadable_names_collects_parameters_and_buffers():
+    loader = _make_loader()
+    model = _DummyModel()
+
+    names = loader._get_model_loadable_names(model)
+
+    assert "transformer.weight" in names
+    assert "transformer.bias" in names
+    assert "transformer_buffer" in names
+
+
+def test_resolve_gguf_model_path_returns_local_file(tmp_path):
+    loader = _make_loader()
+    gguf_file = tmp_path / "model.gguf"
+    gguf_file.write_bytes(b"gguf")
+
+    resolved = loader._resolve_gguf_model_path(str(gguf_file), revision=None)
+
+    assert resolved == str(gguf_file)
+
+
+def test_resolve_gguf_model_path_downloads_explicit_gguf_filename(monkeypatch: pytest.MonkeyPatch):
+    loader = _make_loader()
+
+    monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False)
+    monkeypatch.setattr(
+        loader_module,
+        "hf_hub_download",
+        lambda repo_id, filename, revision, cache_dir: f"{cache_dir}/{repo_id}/{filename}",
+    )
+
+    resolved = loader._resolve_gguf_model_path("owner/repo/model-Q4_0.gguf", revision="main")
+
+    assert resolved == "cache-dir/owner/repo/model-Q4_0.gguf"
+
+
+def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.MonkeyPatch):
+    loader = _make_loader()
+
+    monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False)
+    monkeypatch.setattr(
+        loader_module,
+        "download_gguf",
+        lambda repo_id, quant_type, cache_dir, revision, ignore_patterns: f"{cache_dir}/{repo_id}/{quant_type}.gguf",
+    )
+
+    resolved = loader._resolve_gguf_model_path("owner/repo:Q4_0", revision="main")
+
+    assert resolved == "cache-dir/owner/repo/Q4_0.gguf"
+
+
+def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.MonkeyPatch):
+    loader = _make_loader()
+    source = DiffusersPipelineLoader.ComponentSource(
+        model_or_path="dummy",
+        subfolder="transformer",
+        revision=None,
+        prefix="transformer.",
+    )
+    od_config = SimpleNamespace(
+        revision="main",
+        quantization_config=SimpleNamespace(gguf_model="weights.gguf"),
+    )
+
+    class _Adapter:
+        def weights_iterator(self):
+            yield "weight", torch.ones((2, 2))
+            yield "bias", torch.zeros(2)
+
+    monkeypatch.setattr(loader, "_resolve_gguf_model_path", lambda gguf_model, revision: "resolved.gguf")
+    monkeypatch.setattr(loader_module, "get_gguf_adapter", lambda gguf_file, model, source, od_config: _Adapter())
+
+    weights = list(loader._get_gguf_weights_iterator(source, object(), od_config))
+
+    assert weights[0][0] == "transformer.weight"
+    assert torch.equal(weights[0][1], torch.ones((2, 2)))
+    assert weights[1][0] == "transformer.bias"
+    assert torch.equal(weights[1][1], torch.zeros(2))
+
+
+def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(monkeypatch: pytest.MonkeyPatch):
+    loader = _make_loader()
+    model = _DummyModel()
+    od_config = SimpleNamespace()
+
+    monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
+    monkeypatch.setattr(
+        loader,
+        "_get_gguf_weights_iterator",
+        lambda source, model, od_config: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
+    )
+
+    def _hf_iter(source):
+        if source.subfolder == "transformer":
+            return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))])
+        return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))])
+
+    monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter)
+    monkeypatch.setattr(
+        loader,
+        "_get_expected_parameter_names",
+        lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
+    )
+
+    loaded = loader._load_weights_with_gguf(model, od_config)
+
+    assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
+    assert model.calls[0] == ["transformer.weight"]
+    assert model.calls[1] == ["transformer.bias"]
+    assert model.calls[2] == ["vae.weight"]
+
+
+def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete(monkeypatch: pytest.MonkeyPatch):
+    loader = _make_loader()
+    model = _DummyModel()
+    od_config = SimpleNamespace()
+
+    monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
+    monkeypatch.setattr(
+        loader,
+        "_get_gguf_weights_iterator",
+        lambda source, model, od_config: iter(
+            [
+                ("transformer.weight", torch.ones((2, 2), dtype=torch.float32)),
+                ("transformer.bias", torch.zeros(2, dtype=torch.float32)),
+            ]
+        ),
+    )
+
+    hf_calls: list[str] = []
+
+    def _hf_iter(source):
+        hf_calls.append(source.subfolder or "")
+        if source.subfolder == "transformer":
+            return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))])
+        return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))])
+
+    monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter)
+    monkeypatch.setattr(
+        loader,
+        "_get_expected_parameter_names",
+        lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
+    )
+
+    loaded = loader._load_weights_with_gguf(model, od_config)
+
+    assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
+    assert hf_calls == ["vae"]

tests/diffusion/quantization/test_gguf_config.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for diffusion GGUF quantization config and linear method."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from vllm.model_executor.layers.linear import LinearBase
+from vllm.model_executor.layers.quantization.gguf import UNQUANTIZED_TYPES
+
+from vllm_omni.diffusion.quantization import (
+    DiffusionGgufConfig,
+    get_diffusion_quant_config,
+)
+from vllm_omni.diffusion.quantization.gguf import (
+    DiffusionGGUFLinearMethod,
+    UnquantizedLinearMethod,
+    _GGUFConfig,
+    dequant_gemm_gguf,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_gguf_config_creation_and_delegation():
+    config = DiffusionGgufConfig(
+        gguf_model="weights.gguf",
+        unquantized_modules=["proj_out"],
+    )
+
+    assert config.gguf_model == "weights.gguf"
+    assert config.unquantized_modules == ["proj_out"]
+    assert config.get_name() == "gguf"
+    assert isinstance(config.get_vllm_quant_config(), _GGUFConfig)
+
+
+def test_get_diffusion_quant_config_builds_gguf_config():
+    config = get_diffusion_quant_config(
+        "gguf",
+        gguf_model="repo/model-Q4_0.gguf",
+        unquantized_modules=["norm_out"],
+    )
+
+    assert isinstance(config, DiffusionGgufConfig)
+    assert config.gguf_model == "repo/model-Q4_0.gguf"
+    assert config.unquantized_modules == ["norm_out"]
+
+
+def test_gguf_vllm_config_returns_diffusion_linear_method_for_linear_layers():
+    linear = object.__new__(LinearBase)
+    method = _GGUFConfig(unquantized_modules=[]).get_quant_method(linear, "transformer.img_in")
+
+    assert isinstance(method, DiffusionGGUFLinearMethod)
+
+
+def test_gguf_vllm_config_respects_unquantized_modules():
+    linear = object.__new__(LinearBase)
+    method = _GGUFConfig(unquantized_modules=["proj_out"]).get_quant_method(linear, "transformer.proj_out")
+
+    assert isinstance(method, UnquantizedLinearMethod)
+
+
+def test_gguf_vllm_config_returns_none_for_non_linear_layers():
+    method = _GGUFConfig(unquantized_modules=[]).get_quant_method(torch.nn.LayerNorm(4), "norm")
+
+    assert method is None
+
+
+def test_dequant_gemm_gguf_uses_plain_matmul_for_unquantized_types():
+    qweight_type = next(iter(UNQUANTIZED_TYPES))
+    x = torch.tensor([[1.0, 2.0]], dtype=torch.float32)
+    qweight = torch.tensor([[3.0, 4.0], [5.0, 6.0]], dtype=torch.float32)
+
+    out = dequant_gemm_gguf(x, qweight, qweight_type)
+
+    assert torch.allclose(out, x @ qweight.T)
+
+
+def test_diffusion_gguf_linear_method_applies_bias_on_unquantized_weight():
+    qweight_type = next(iter(UNQUANTIZED_TYPES))
+    method = DiffusionGGUFLinearMethod(_GGUFConfig(unquantized_modules=[]))
+    layer = SimpleNamespace(
+        qweight=torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32),
+        qweight_type=SimpleNamespace(weight_type=qweight_type),
+    )
+    x = torch.tensor([[[1.0, 0.5]]], dtype=torch.float32)
+    bias = torch.tensor([0.25, -0.5], dtype=torch.float32)
+
+    out = method.apply(layer, x, bias)
+
+    expected = x @ layer.qweight.T + bias
+    assert torch.allclose(out, expected)
+
+
+def test_diffusion_gguf_linear_method_concatenates_sharded_outputs():
+    qweight_type = next(iter(UNQUANTIZED_TYPES))
+    method = DiffusionGGUFLinearMethod(_GGUFConfig(unquantized_modules=[]))
+    qweight = torch.nn.Parameter(
+        torch.tensor(
+            [
+                [1.0, 0.0],
+                [0.0, 1.0],
+                [2.0, 0.0],
+                [0.0, 2.0],
+            ],
+            dtype=torch.float32,
+        ),
+        requires_grad=False,
+    )
+    qweight.shard_id = ["left", "right"]
+    qweight.shard_offset_map = {
+        "left": (0, 2, 2),
+        "right": (2, 4, 2),
+    }
+    layer = SimpleNamespace(
+        qweight=qweight,
+        qweight_type=SimpleNamespace(
+            shard_weight_type={
+                "left": qweight_type,
+                "right": qweight_type,
+            }
+        ),
+    )
+    x = torch.tensor([[1.5, 2.0]], dtype=torch.float32)
+
+    out = method.apply(layer, x)
+
+    expected_q = x @ qweight[:2].T
+    expected_k = x @ qweight[2:4].T
+    assert torch.allclose(out, torch.cat([expected_q, expected_k], dim=-1))

vllm_omni/diffusion/quantization/base.py

@@ -38,7 +38,17 @@ class DiffusionQuantizationConfig(ABC):
         Delegates to the underlying vLLM config class's get_name().
         """
         if cls.quant_config_cls is not None:
-            return cls.quant_config_cls.get_name()
+            try:
+                return cls.quant_config_cls.get_name()
+            except TypeError as exc:
+                # Upstream quantization configs are not fully consistent:
+                # some expose get_name as a classmethod (e.g. FP8), while
+                # others expose it as an instance method (e.g. GGUF in
+                # vLLM 0.18.0). Fall back to a default-constructed instance
+                # so diffusion wrappers work with both styles.
+                if "required positional argument" not in str(exc):
+                    raise
+                return cls.quant_config_cls().get_name()
         raise NotImplementedError("Subclass must set quant_config_cls or override get_name().")
 
     def get_vllm_quant_config(self) -> "QuantizationConfig | None":