| import unittest |
| from unittest import mock |
|
|
| from lora_utils import ensure_loras_loaded, parse_adapter_specs |
|
|
|
|
| class DummyAdapterHost: |
| def __init__(self): |
| self.loaded = [] |
| self.active_names = None |
| self.active_weights = None |
|
|
| def load_lora_adapter(self, state_dict, adapter_name, prefix=None): |
| self.loaded.append((state_dict, adapter_name, prefix)) |
|
|
| def set_adapters(self, adapter_names, adapter_weights=None): |
| self.active_names = adapter_names |
| self.active_weights = adapter_weights |
|
|
| def get_list_adapters(self): |
| return [adapter_name for _, adapter_name, _ in self.loaded] |
|
|
| def delete_adapters(self, adapter_name): |
| self.loaded = [entry for entry in self.loaded if entry[1] != adapter_name] |
|
|
|
|
| class DummyIdeogramPipeline: |
| def __init__(self): |
| self.transformer = DummyAdapterHost() |
| self.unconditional_transformer = DummyAdapterHost() |
|
|
|
|
| class DummyPeftLayer: |
| def __init__(self): |
| self.lora_A = {} |
| self.active_names = None |
| self.scales = [] |
| self.enabled = None |
| self.deleted = [] |
|
|
| def set_adapter(self, adapter_names): |
| self.active_names = adapter_names |
|
|
| def set_scale(self, adapter_name, scale): |
| self.scales.append((adapter_name, scale)) |
|
|
| def enable_adapters(self, enabled): |
| self.enabled = enabled |
|
|
| def delete_adapter(self, adapter_name): |
| self.deleted.append(adapter_name) |
| self.lora_A.pop(adapter_name, None) |
|
|
|
|
| class DummyPeftHost: |
| def __init__(self): |
| self.peft_config = {} |
| self.layer = DummyPeftLayer() |
|
|
| def modules(self): |
| return [self, self.layer] |
|
|
| def named_modules(self): |
| return [("", self), ("layer", self.layer)] |
|
|
|
|
| class DummyPeftOnlyIdeogramPipeline: |
| def __init__(self): |
| self.transformer = DummyPeftHost() |
| self.unconditional_transformer = DummyPeftHost() |
|
|
|
|
| class ParseAdapterSpecsTest(unittest.TestCase): |
| def test_blank_spec_returns_no_entries(self): |
| self.assertEqual(parse_adapter_specs("\n \n", 1.0), []) |
|
|
| def test_parses_repo_weight_and_default_scale(self): |
| entries = parse_adapter_specs("vladi/real-loras:ema4_flux2_klein_9b_000002000.safetensors", 1.0) |
|
|
| self.assertEqual(entries[0]["repo_id"], "vladi/real-loras") |
| self.assertEqual(entries[0]["weight_name"], "ema4_flux2_klein_9b_000002000.safetensors") |
| self.assertEqual(entries[0]["scale"], 1.0) |
|
|
| def test_parses_multiple_entries_and_inline_scale(self): |
| entries = parse_adapter_specs( |
| "vladi/real-loras:ema4_flux2_klein_9b_000002000.safetensors\n" |
| "vladi/loras:klein9b/klein_snofs_v1_4.safetensors@0.8", |
| 0.5, |
| ) |
|
|
| self.assertEqual(len(entries), 2) |
| self.assertEqual(entries[1]["repo_id"], "vladi/loras") |
| self.assertEqual(entries[1]["weight_name"], "klein9b/klein_snofs_v1_4.safetensors") |
| self.assertEqual(entries[1]["scale"], 0.4) |
|
|
| def test_parses_huggingface_url_with_subfolder(self): |
| entries = parse_adapter_specs( |
| "https://huggingface.co/vladi/loras/blob/dev/klein9b/klein_snofs_v1_4.safetensors@0.8", |
| 1.0, |
| ) |
|
|
| self.assertEqual(entries[0]["repo_id"], "vladi/loras") |
| self.assertEqual(entries[0]["weight_name"], "klein9b/klein_snofs_v1_4.safetensors") |
| self.assertEqual(entries[0]["revision"], "dev") |
| self.assertEqual(entries[0]["scale"], 0.8) |
|
|
| def test_rejects_missing_weight_name(self): |
| with self.assertRaises(ValueError): |
| parse_adapter_specs("vladi/real-loras", 1.0) |
|
|
| def test_rejects_duplicate_entries(self): |
| with self.assertRaises(ValueError): |
| parse_adapter_specs( |
| "vladi/real-loras:ema4_flux2_klein_9b_000002000.safetensors\n" |
| "vladi/real-loras:ema4_flux2_klein_9b_000002000.safetensors@0.8", |
| 1.0, |
| ) |
|
|
| def test_loads_lora_without_pipeline_loader(self): |
| pipe = DummyIdeogramPipeline() |
| active_by_key = {} |
| state_dict = {"transformer_blocks.0.attn.to_q.lora_A.weight": object()} |
|
|
| with mock.patch("lora_utils._download_lora_weight", return_value="/tmp/adapter.safetensors"), mock.patch( |
| "lora_utils._load_adapter_state_dict", return_value=state_dict |
| ): |
| entries = ensure_loras_loaded( |
| pipe, |
| "vladi/real-loras:ema4_flux2_klein_9b_000002000.safetensors", |
| 0.8, |
| active_by_key, |
| token=None, |
| ) |
|
|
| adapter_name = entries[0]["adapter_name"] |
| self.assertEqual(pipe.transformer.loaded[0], (state_dict, adapter_name, None)) |
| self.assertEqual(pipe.unconditional_transformer.loaded[0], (state_dict, adapter_name, None)) |
| self.assertEqual(pipe.transformer.active_names, [adapter_name]) |
| self.assertEqual(pipe.unconditional_transformer.active_names, [adapter_name]) |
| self.assertEqual(pipe.transformer.active_weights, [0.8]) |
| self.assertEqual(pipe.unconditional_transformer.active_weights, [0.8]) |
|
|
| def test_loads_lora_with_peft_fallback_when_transformers_have_no_loader(self): |
| pipe = DummyPeftOnlyIdeogramPipeline() |
| active_by_key = {} |
| state_dict = {"transformer.transformer_blocks.0.attn.to_q.lora_A.weight": object()} |
| loaded = [] |
|
|
| def fake_load_with_peft(host, host_state_dict, adapter_name): |
| loaded.append((host, host_state_dict, adapter_name)) |
| host.peft_config[adapter_name] = object() |
| host.layer.lora_A[adapter_name] = object() |
|
|
| with mock.patch("lora_utils._download_lora_weight", return_value="/tmp/adapter.safetensors"), mock.patch( |
| "lora_utils._load_adapter_state_dict", return_value=state_dict |
| ), mock.patch("lora_utils._load_lora_with_peft", side_effect=fake_load_with_peft): |
| entries = ensure_loras_loaded( |
| pipe, |
| "https://huggingface.co/vladi/loras/blob/dev/klein9b/klein_snofs_v1_4.safetensors", |
| 0.7, |
| active_by_key, |
| token="secret", |
| ) |
|
|
| adapter_name = entries[0]["adapter_name"] |
| expected_state_dict = {"transformer_blocks.0.attn.to_q.lora_A.weight": state_dict[next(iter(state_dict))]} |
| self.assertEqual(loaded[0], (pipe.transformer, expected_state_dict, adapter_name)) |
| self.assertEqual(loaded[1], (pipe.unconditional_transformer, expected_state_dict, adapter_name)) |
| self.assertEqual(pipe.transformer.layer.active_names, [adapter_name]) |
| self.assertEqual(pipe.unconditional_transformer.layer.active_names, [adapter_name]) |
| self.assertEqual(pipe.transformer.layer.scales, [(adapter_name, 0.7)]) |
| self.assertEqual(pipe.unconditional_transformer.layer.scales, [(adapter_name, 0.7)]) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
