[Quantization] Migrate GGUF diffusion model support to OOT plugin (#4769)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: Isotr0py <Isotr0py@outlook.com>
This commit is contained in:
Isotr0py authored and GitHub committed
1 parent
ff6f0da271
commit
b95905deb1
35 files changed
+272
-1614
No files matched your search
@@ -1,112 +0,0 @@
|
||||
# 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]))
|
||||
@@ -1,184 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
|
||||
from vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image import (
|
||||
QwenImageGGUFAdapter,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_od_config(
|
||||
model_class_name: str = "QwenImagePipeline",
|
||||
model_type: str = "qwen_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)
|
||||
|
||||
|
||||
class _FakeTransformer:
|
||||
def __init__(self) -> None:
|
||||
self._params = {
|
||||
"img_in.qweight": 1,
|
||||
"img_in.qweight_type": 1,
|
||||
"txt_in.qweight": 1,
|
||||
"txt_in.qweight_type": 1,
|
||||
"proj_out.qweight": 1,
|
||||
"proj_out.qweight_type": 1,
|
||||
"time_text_embed.timestep_embedder.linear_1.qweight": 1,
|
||||
"time_text_embed.timestep_embedder.linear_1.qweight_type": 1,
|
||||
"time_text_embed.timestep_embedder.linear_2.qweight": 1,
|
||||
"time_text_embed.timestep_embedder.linear_2.qweight_type": 1,
|
||||
"transformer_blocks.0.img_mod.1.qweight": 1,
|
||||
"transformer_blocks.0.img_mod.1.qweight_type": 1,
|
||||
"transformer_blocks.0.txt_mod.1.qweight": 1,
|
||||
"transformer_blocks.0.txt_mod.1.qweight_type": 1,
|
||||
"transformer_blocks.0.attn.to_qkv.qweight": 1,
|
||||
"transformer_blocks.0.attn.to_qkv.qweight_type": 1,
|
||||
"transformer_blocks.0.attn.add_kv_proj.qweight": 1,
|
||||
"transformer_blocks.0.attn.add_kv_proj.qweight_type": 1,
|
||||
"transformer_blocks.0.attn.to_out.weight": 1,
|
||||
"transformer_blocks.0.img_mlp.net.0.proj.qweight": 1,
|
||||
"transformer_blocks.0.img_mlp.net.0.proj.qweight_type": 1,
|
||||
}
|
||||
|
||||
def named_parameters(self):
|
||||
return list(self._params.items())
|
||||
|
||||
def named_buffers(self):
|
||||
return []
|
||||
|
||||
def get_submodule(self, _prefix: str):
|
||||
return self
|
||||
|
||||
|
||||
def test_qwen_adapter_selected_for_qwen_image_family():
|
||||
adapter = get_gguf_adapter(
|
||||
"dummy.gguf",
|
||||
_FakeTransformer(),
|
||||
_make_source(),
|
||||
_make_od_config(),
|
||||
)
|
||||
assert isinstance(adapter, QwenImageGGUFAdapter)
|
||||
|
||||
|
||||
def test_qwen_adapter_matches_multiple_pipeline_variants():
|
||||
for model_class_name in (
|
||||
"QwenImagePipeline",
|
||||
"QwenImageEditPipeline",
|
||||
"QwenImageEditPlusPipeline",
|
||||
"QwenImageLayeredPipeline",
|
||||
):
|
||||
assert QwenImageGGUFAdapter.is_compatible(
|
||||
_make_od_config(model_class_name=model_class_name),
|
||||
_FakeTransformer(),
|
||||
_make_source(),
|
||||
)
|
||||
|
||||
|
||||
def test_qwen_adapter_preserves_split_projection_names(monkeypatch: pytest.MonkeyPatch):
|
||||
import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
qwen_image_module,
|
||||
"gguf_quant_weights_iterator",
|
||||
lambda _path: iter(
|
||||
[
|
||||
("transformer_blocks.0.attn.to_q.qweight_type", 1),
|
||||
("transformer_blocks.0.attn.to_q.qweight", 2),
|
||||
("transformer_blocks.0.attn.to_k.qweight_type", 3),
|
||||
("transformer_blocks.0.attn.to_k.qweight", 4),
|
||||
("transformer_blocks.0.attn.to_out.0.qweight_type", 5),
|
||||
("transformer_blocks.0.attn.to_out.0.qweight", 6),
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
adapter = QwenImageGGUFAdapter(
|
||||
"dummy.gguf",
|
||||
_FakeTransformer(),
|
||||
_make_source(),
|
||||
_make_od_config(),
|
||||
)
|
||||
|
||||
weights = list(adapter.weights_iterator())
|
||||
|
||||
assert ("transformer_blocks.0.attn.to_q.qweight_type", 1) in weights
|
||||
assert ("transformer_blocks.0.attn.to_q.qweight", 2) in weights
|
||||
assert ("transformer_blocks.0.attn.to_k.qweight_type", 3) in weights
|
||||
assert ("transformer_blocks.0.attn.to_k.qweight", 4) in weights
|
||||
assert ("transformer_blocks.0.attn.to_out.0.qweight_type", 5) in weights
|
||||
assert ("transformer_blocks.0.attn.to_out.0.qweight", 6) in weights
|
||||
|
||||
|
||||
def test_qwen_adapter_keeps_top_level_quantized_weights(monkeypatch: pytest.MonkeyPatch):
|
||||
import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
qwen_image_module,
|
||||
"gguf_quant_weights_iterator",
|
||||
lambda _path: iter(
|
||||
[
|
||||
("img_in.qweight_type", 1),
|
||||
("img_in.qweight", 2),
|
||||
("time_text_embed.timestep_embedder.linear_1.qweight_type", 3),
|
||||
("time_text_embed.timestep_embedder.linear_1.qweight", 4),
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
adapter = QwenImageGGUFAdapter(
|
||||
"dummy.gguf",
|
||||
_FakeTransformer(),
|
||||
_make_source(),
|
||||
_make_od_config(),
|
||||
)
|
||||
|
||||
weights = list(adapter.weights_iterator())
|
||||
|
||||
assert ("img_in.qweight_type", 1) in weights
|
||||
assert ("img_in.qweight", 2) in weights
|
||||
assert ("time_text_embed.timestep_embedder.linear_1.qweight_type", 3) in weights
|
||||
assert ("time_text_embed.timestep_embedder.linear_1.qweight", 4) in weights
|
||||
|
||||
|
||||
def test_qwen_adapter_keeps_modulation_quantized_weights(monkeypatch: pytest.MonkeyPatch):
|
||||
import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
qwen_image_module,
|
||||
"gguf_quant_weights_iterator",
|
||||
lambda _path: iter(
|
||||
[
|
||||
("transformer_blocks.0.img_mod.1.qweight_type", 1),
|
||||
("transformer_blocks.0.img_mod.1.qweight", 2),
|
||||
("transformer_blocks.0.txt_mod.1.qweight_type", 3),
|
||||
("transformer_blocks.0.txt_mod.1.qweight", 4),
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
adapter = QwenImageGGUFAdapter(
|
||||
"dummy.gguf",
|
||||
_FakeTransformer(),
|
||||
_make_source(),
|
||||
_make_od_config(),
|
||||
)
|
||||
|
||||
weights = list(adapter.weights_iterator())
|
||||
|
||||
assert ("transformer_blocks.0.img_mod.1.qweight_type", 1) in weights
|
||||
assert ("transformer_blocks.0.img_mod.1.qweight", 2) in weights
|
||||
assert ("transformer_blocks.0.txt_mod.1.qweight_type", 3) in weights
|
||||
assert ("transformer_blocks.0.txt_mod.1.qweight", 4) in weights
|
||||
@@ -1,83 +0,0 @@
|
||||
# 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
|
||||
@@ -16,7 +16,6 @@ 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
|
||||
|
||||
@@ -111,27 +110,6 @@ def test_empty_source_prefix_keeps_full_model_strict_check():
|
||||
loader.load_weights(model)
|
||||
|
||||
|
||||
def test_qwen_model_class_selects_qwen_gguf_adapter():
|
||||
od_config = type(
|
||||
"Config",
|
||||
(),
|
||||
{
|
||||
"model_class_name": "QwenImagePipeline",
|
||||
"tf_model_config": {"model_type": "qwen_image"},
|
||||
},
|
||||
)()
|
||||
source = DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path="dummy",
|
||||
subfolder="transformer",
|
||||
revision=None,
|
||||
prefix="transformer.",
|
||||
)
|
||||
|
||||
adapter = get_gguf_adapter("dummy.gguf", object(), source, od_config)
|
||||
|
||||
assert adapter.__class__.__name__ == "QwenImageGGUFAdapter"
|
||||
|
||||
|
||||
class _ConfigAwareModel(nn.Module):
|
||||
def __init__(self, *, od_config):
|
||||
super().__init__()
|
||||
@@ -186,7 +164,6 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
|
||||
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: 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"))
|
||||
|
||||
@@ -1,259 +0,0 @@
|
||||
# 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.quant_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()
|
||||
loader.quant_config = {
|
||||
"method": "gguf",
|
||||
"gguf_model": "weights.gguf",
|
||||
}
|
||||
|
||||
assert loader._is_gguf_quantization() is True
|
||||
|
||||
|
||||
def test_is_gguf_quantization_rejects_non_gguf_dict_config():
|
||||
loader = _make_loader()
|
||||
loader.quant_config = {
|
||||
"method": "fp8",
|
||||
"gguf_model": "weights.gguf",
|
||||
}
|
||||
|
||||
assert loader._is_gguf_quantization() is False
|
||||
|
||||
|
||||
def test_is_gguf_quantization_requires_gguf_model_for_dict_config():
|
||||
loader = _make_loader()
|
||||
loader.quant_config = {"method": "gguf"}
|
||||
|
||||
with pytest.raises(ValueError, match="gguf_model"):
|
||||
loader._is_gguf_quantization()
|
||||
|
||||
|
||||
def test_is_gguf_quantization_accepts_object_config():
|
||||
loader = _make_loader()
|
||||
loader.quant_config = SimpleNamespace(
|
||||
get_name=lambda: "gguf",
|
||||
gguf_model="weights.gguf",
|
||||
)
|
||||
|
||||
assert loader._is_gguf_quantization() is True
|
||||
|
||||
|
||||
def test_is_gguf_quantization_uses_fallback_object_without_get_name():
|
||||
loader = _make_loader()
|
||||
loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
|
||||
|
||||
assert loader._is_gguf_quantization() 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()
|
||||
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.",
|
||||
)
|
||||
|
||||
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()))
|
||||
|
||||
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()
|
||||
|
||||
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
|
||||
monkeypatch.setattr(
|
||||
loader,
|
||||
"_get_gguf_weights_iterator",
|
||||
lambda source, model: 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)
|
||||
|
||||
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()
|
||||
|
||||
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
|
||||
monkeypatch.setattr(
|
||||
loader,
|
||||
"_get_gguf_weights_iterator",
|
||||
lambda source, model: 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)
|
||||
|
||||
assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
|
||||
assert hf_calls == ["vae"]
|
||||
Reference in new issue
Block a user