better lora loading

lora_utils.py

@@ -135,6 +135,40 @@ def _sorted_lora_entries(entries):
     return sorted(entries, key=lambda entry: entry["adapter_name"])
 
 
+def _download_lora_weight(repo_id: str, weight_name: str, token=HF_TOKEN):
+    from huggingface_hub import hf_hub_download
+
+    kwargs = {}
+    if token:
+        kwargs["token"] = token
+    return hf_hub_download(repo_id, filename=weight_name, **kwargs)
+
+
+def _load_adapter_state_dict(local_path: str):
+    if local_path.endswith(".safetensors"):
+        from safetensors.torch import load_file as safetensors_load_file
+
+        return safetensors_load_file(local_path)
+
+    import torch
+
+    return torch.load(local_path, map_location="cpu")
+
+
+def _ensure_pipeline_lora_prefix(state_dict):
+    if any(key.startswith("transformer.") for key in state_dict.keys()):
+        return state_dict
+
+    if all(
+        key.startswith("single_transformer_blocks.")
+        or key.startswith("transformer_blocks.")
+        for key in state_dict.keys()
+    ):
+        return {f"transformer.{key}": value for key, value in state_dict.items()}
+
+    return state_dict
+
+
 def safe_unload_lora_adapters(pipe):
     deleted = False
     for host in _iter_adapter_hosts(pipe):
@@ -212,9 +246,13 @@ def apply_lora_adapters(pipe, lora_entries):
     adapter_names = [entry["adapter_name"] for entry in sorted_entries]
     adapter_weights = [entry["scale"] for entry in sorted_entries]
 
+    activated = False
     for host in _iter_adapter_hosts(pipe):
         if _set_adapters_on_host(host, adapter_names, adapter_weights):
-            return
+            activated = True
+
+    if activated:
+        return
 
     if len(adapter_names) == 1 and hasattr(pipe, "set_lora_scale"):
         pipe.set_lora_scale(adapter_weights[0])
@@ -231,13 +269,58 @@ def load_lora_adapter(pipe, entry, token=HF_TOKEN):
     if token:
         load_kwargs["token"] = token
 
-    try:
-        pipe.load_lora_weights(entry["repo_id"], **load_kwargs)
-    except TypeError:
-        if not token:
+    native_error = None
+    if hasattr(pipe, "load_lora_weights"):
+        try:
+            pipe.load_lora_weights(entry["repo_id"], **load_kwargs)
+            return
+        except TypeError as exc:
+            native_error = exc
+            if token:
+                load_kwargs.pop("token", None)
+                try:
+                    pipe.load_lora_weights(entry["repo_id"], **load_kwargs)
+                    return
+                except Exception as retry_exc:
+                    native_error = retry_exc
+        except Exception as exc:
+            native_error = exc
+
+    local_path = _download_lora_weight(entry["repo_id"], entry["weight_name"], token=token)
+    state_dict = _load_adapter_state_dict(local_path)
+
+    loaded = False
+    for host in _iter_adapter_hosts(pipe):
+        if not hasattr(host, "load_lora_adapter"):
+            continue
+        try:
+            host.load_lora_adapter(dict(state_dict), adapter_name=entry["adapter_name"], prefix=None)
+            loaded = True
+            continue
+        except TypeError:
+            try:
+                host.load_lora_adapter(dict(state_dict), adapter_name=entry["adapter_name"])
+                loaded = True
+                continue
+            except Exception as exc:
+                native_error = native_error or exc
+        except Exception as exc:
+            native_error = native_error or exc
+
+    if loaded:
+        return
+
+    if hasattr(pipe, "load_lora_weights"):
+        try:
+            fallback_kwargs = {"adapter_name": entry["adapter_name"]}
+            pipe.load_lora_weights(_ensure_pipeline_lora_prefix(state_dict), **fallback_kwargs)
+            return
+        except Exception as exc:
+            if native_error is not None:
+                raise ValueError(f"{native_error}; fallback failed with {exc}") from exc
             raise
-        load_kwargs.pop("token", None)
-        pipe.load_lora_weights(entry["repo_id"], **load_kwargs)
+
+    raise ValueError("This pipeline does not expose a LoRA loader on itself or its transformers.")
 
 
 def ensure_loras_loaded(pipe, spec_text: str, global_scale: float, active_by_key: dict, token=HF_TOKEN):

tests/test_lora_utils.py

@@ -1,6 +1,33 @@
 import unittest
+from unittest import mock
 
-from lora_utils import parse_adapter_specs
+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 ParseAdapterSpecsTest(unittest.TestCase):
@@ -48,6 +75,30 @@ class ParseAdapterSpecsTest(unittest.TestCase):
                 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])
+
 
 if __name__ == "__main__":
     unittest.main()