[Core] Cleanup DiffusersPipelineLoader (#1932)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -1,19 +1,46 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+"""
+Tests for the DiffusersPipelineLoader.
+"""
+
 from types import SimpleNamespace
 
 import pytest
 import torch
 import torch.nn as nn
+from huggingface_hub import snapshot_download
+from vllm.config.load import LoadConfig
 
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
+from vllm_omni.diffusion.models.helios import HeliosPipeline
 from vllm_omni.diffusion.registry import initialize_model
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
+model_path = "hf-internal-testing/tiny-helios-modular-pipe"
+
+
+@pytest.fixture(scope="module")
+def prefetch_helios_model():
+    """Downloads the tiny helios model prior to running a test."""
+    snapshot_download(model_path)
+
+
+@pytest.fixture(scope="function")
+def mock_tp_group(mocker):
+    """Mocks the tensor parallel group; this is needed to initialize the Helios model."""
+    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
+    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
+    mock_group = mocker.MagicMock()
+    mock_group.world_size = 1
+    mock_group.rank_in_group = 0
+    mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
+
 
 class _DummyPipelineModel(nn.Module):
     def __init__(self, *, source_prefix: str):
@@ -42,7 +69,13 @@ class _DummyPipelineModel(nn.Module):
 
 
 def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader:
-    loader = object.__new__(DiffusersPipelineLoader)
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=False),
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+
     loader.counter_before_loading_weights = 0.0
     loader.counter_after_loading_weights = 0.0
 
@@ -150,16 +183,15 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
         quantization_config=None,
     )
 
-    loader = object.__new__(DiffusersPipelineLoader)
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
     loader.load_weights = lambda model: None  # type: ignore[assignment]
     loader._process_weights_after_loading = lambda model, target_device: None  # type: ignore[assignment]
-    loader._is_gguf_quantization = lambda _od_config: False  # type: ignore[assignment]
+    loader._is_gguf_quantization = lambda: False  # type: ignore[assignment]
 
     monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
     monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))
 
     model = loader.load_model(
-        od_config,
         load_device="cpu",
         load_format="custom_pipeline",
         custom_pipeline_name="tests.dummy.ConfigAwarePipeline",
@@ -168,3 +200,33 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
     assert model.captured_config is od_config
     assert model.seen_config_during_init is od_config
     assert get_current_diffusion_config_or_none() is None
+
+
+def test_get_all_weights(prefetch_helios_model, mock_tp_group):
+    """Ensure that get all weights on a tiny model resolves to nonempty weights."""
+    od_config = OmniDiffusionConfig(
+        model_class_name="HeliosPipeline",
+        model=model_path,
+    )
+    loader = DiffusersPipelineLoader(
+        load_config=LoadConfig(),
+        od_config=od_config,
+    )
+    pipeline = HeliosPipeline(od_config=od_config)
+
+    weights = list(loader.get_all_weights(pipeline))
+    assert len(weights) > 0
+
+
+def test_load_model(prefetch_helios_model, mock_tp_group):
+    """Ensure that load model creates an instance of the expected pipeline class."""
+    od_config = OmniDiffusionConfig(
+        model_class_name="HeliosPipeline",
+        model=model_path,
+    )
+    loader = DiffusersPipelineLoader(
+        load_config=LoadConfig(),
+        od_config=od_config,
+    )
+    model = loader.load_model(load_device="cpu")
+    assert isinstance(model, HeliosPipeline)

tests/diffusion/model_loader/test_diffusers_loader_gguf.py

@@ -62,6 +62,7 @@ def _make_loader() -> DiffusersPipelineLoader:
         ignore_patterns=["*.tmp"],
     )
     loader.od_config = None
+    loader.quant_config = None
     loader.counter_before_loading_weights = 0.0
     loader.counter_after_loading_weights = 0.0
     return loader
@@ -69,53 +70,47 @@ def _make_loader() -> DiffusersPipelineLoader:
 
 def test_is_gguf_quantization_accepts_dict_config():
     loader = _make_loader()
-    od_config = SimpleNamespace(
-        quantization_config={
-            "method": "gguf",
-            "gguf_model": "weights.gguf",
-        }
-    )
+    loader.quant_config = {
+        "method": "gguf",
+        "gguf_model": "weights.gguf",
+    }
 
-    assert loader._is_gguf_quantization(od_config) is True
+    assert loader._is_gguf_quantization() 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",
-        }
-    )
+    loader.quant_config = {
+        "method": "fp8",
+        "gguf_model": "weights.gguf",
+    }
 
-    assert loader._is_gguf_quantization(od_config) is False
+    assert loader._is_gguf_quantization() is False
 
 
 def test_is_gguf_quantization_requires_gguf_model_for_dict_config():
     loader = _make_loader()
-    od_config = SimpleNamespace(quantization_config={"method": "gguf"})
+    loader.quant_config = {"method": "gguf"}
 
     with pytest.raises(ValueError, match="gguf_model"):
-        loader._is_gguf_quantization(od_config)
+        loader._is_gguf_quantization()
 
 
 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",
-        )
+    loader.quant_config = SimpleNamespace(
+        get_name=lambda: "gguf",
+        gguf_model="weights.gguf",
     )
 
-    assert loader._is_gguf_quantization(od_config) is True
+    assert loader._is_gguf_quantization() 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"))
+    loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
 
-    assert loader._is_gguf_quantization(od_config) is True
+    assert loader._is_gguf_quantization() is True
 
 
 def test_get_model_loadable_names_collects_parameters_and_buffers():
@@ -171,16 +166,14 @@ def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.Mon
 
 def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.MonkeyPatch):
     loader = _make_loader()
+    loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
+    loader.od_config = SimpleNamespace(revision="main")
     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):
@@ -190,7 +183,7 @@ def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.Mon
     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))
+    weights = list(loader._get_gguf_weights_iterator(source, object()))
 
     assert weights[0][0] == "transformer.weight"
     assert torch.equal(weights[0][1], torch.ones((2, 2)))
@@ -201,13 +194,12 @@ def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.Mon
 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))]),
+        lambda source, model: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
     )
 
     def _hf_iter(source):
@@ -222,7 +214,7 @@ def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(
         lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
     )
 
-    loaded = loader._load_weights_with_gguf(model, od_config)
+    loaded = loader._load_weights_with_gguf(model)
 
     assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
     assert model.calls[0] == ["transformer.weight"]
@@ -233,13 +225,12 @@ def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(
 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(
+        lambda source, model: iter(
             [
                 ("transformer.weight", torch.ones((2, 2), dtype=torch.float32)),
                 ("transformer.bias", torch.zeros(2, dtype=torch.float32)),
@@ -262,7 +253,7 @@ def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete
         lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
     )
 
-    loaded = loader._load_weights_with_gguf(model, od_config)
+    loaded = loader._load_weights_with_gguf(model)
 
     assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
     assert hf_calls == ["vae"]