[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)