[Core] Cleanup DiffusersPipelineLoader (#1932)

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

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"]