[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