Merge pull request #70 from jkoelker/fix/issue-69-lora-state-leakage

fix: prevent LoRA state leakage between generation requests

Author: jkoelker

src/oneiro/pipelines/flux2.py

@@ -64,6 +64,7 @@ def load(self, model_config: dict[str, Any], full_config: dict[str, Any] | None
         if loras:
             print(f"  Loading {len(loras)} LoRA(s)...")
             self.load_loras_sync(loras)
+            self.set_static_loras(loras)
 
         # Load embeddings if full_config provided
         if full_config:
@@ -131,3 +132,8 @@ def build_generation_kwargs(
                 "guidance_scale": guidance_scale,
                 "generator": generator,
             }
+
+    def post_generate(self, **kwargs: Any) -> None:
+        """Reset LoRA state after generation to prevent state leakage."""
+        super().post_generate(**kwargs)
+        self.restore_static_loras()

src/oneiro/pipelines/lora.py

@@ -605,8 +605,9 @@ class LoraLoaderMixin:
     def __init__(self) -> None:
         """Initialize LoRA state and continue MRO chain."""
         super().__init__()
-        self._lora_configs = []
-        self._loaded_adapters = []
+        self._lora_configs: list[LoraConfig] = []
+        self._loaded_adapters: list[str] = []
+        self._static_lora_configs: list[LoraConfig] = []
 
     def load_single_lora(
         self,
@@ -796,3 +797,42 @@ def active_loras(self) -> list[str]:
     def lora_count(self) -> int:
         """Get number of loaded LoRAs."""
         return len(self._loaded_adapters)
+
+    def set_static_loras(self, loras: list[LoraConfig]) -> None:
+        """Store the static LoRA baseline for post-generation restore.
+
+        Call this after loading LoRAs from config in load() to establish
+        which adapters should be preserved between generation requests.
+
+        Args:
+            loras: List of LoRA configurations that represent the static baseline
+        """
+        self._static_lora_configs = list(loras)
+
+    def restore_static_loras(self) -> None:
+        """Restore LoRA state to static adapters loaded from config.
+
+        Call this in post_generate() to prevent state leakage between requests.
+        If dynamic LoRAs were added during a request, this will unload them
+        and restore only the static config LoRAs.
+
+        Behavior:
+        - If no static LoRAs and no loaded adapters: no-op
+        - If adapters match static config: reset weights only
+        - If adapters differ: unload all and reload static LoRAs
+        """
+        static_names = [lora.adapter_name or lora.name for lora in self._static_lora_configs]
+        adapters_match = self._loaded_adapters == static_names
+
+        if adapters_match:
+            # Adapters match - just reset weights if there are static LoRAs
+            if self._static_lora_configs:
+                adapter_weights = [lora.weight for lora in self._static_lora_configs]
+                self.set_lora_adapters(static_names, adapter_weights)
+            return
+
+        # Adapters don't match - need to restore to static baseline
+        self.unload_loras()
+        if self._static_lora_configs:
+            self.load_loras_sync(self._static_lora_configs)
+            print(f"Restored {len(self._static_lora_configs)} static LoRA(s)")

src/oneiro/pipelines/qwen.py

@@ -146,6 +146,7 @@ def load(self, model_config: dict[str, Any], full_config: dict[str, Any] | None
         if loras:
             print(f"  Loading {len(loras)} LoRA(s)...")
             self.load_loras_sync(loras)
+            self.set_static_loras(loras)
 
         if full_config:
             embeddings = parse_embeddings_from_config(full_config, model_config)
@@ -247,3 +248,8 @@ def build_result(
             steps=steps,
             guidance_scale=guidance_scale,
         )
+
+    def post_generate(self, **kwargs: Any) -> None:
+        """Reset LoRA state after generation to prevent state leakage."""
+        super().post_generate(**kwargs)
+        self.restore_static_loras()

src/oneiro/pipelines/zimage.py

@@ -39,6 +39,7 @@ def load(self, model_config: dict[str, Any], full_config: dict[str, Any] | None
         if loras:
             print(f"  Loading {len(loras)} LoRA(s)...")
             self.load_loras_sync(loras)
+            self.set_static_loras(loras)
 
         # Load embeddings if full_config provided
         if full_config:
@@ -107,3 +108,8 @@ def build_result(
             steps=steps,
             guidance_scale=0.0,  # Always 0.0 for Turbo
         )
+
+    def post_generate(self, **kwargs: Any) -> None:
+        """Reset LoRA state after generation to prevent state leakage."""
+        super().post_generate(**kwargs)
+        self.restore_static_loras()

tests/test_lora.py

@@ -798,3 +798,153 @@ async def test_civitai_handles_empty_trained_words(self, tmp_path):
         )
 
         assert config.trigger_words == []
+
+
+class TestLoraLoaderMixinStaticLoras:
+    """Tests for LoraLoaderMixin static LoRA management methods."""
+
+    def test_set_static_loras_stores_configs(self):
+        """set_static_loras stores a copy of the configs."""
+        from oneiro.pipelines.lora import LoraLoaderMixin
+
+        class TestMixin(LoraLoaderMixin):
+            pass
+
+        mixin = TestMixin()
+        loras = [
+            LoraConfig(name="lora1", source=LoraSource.LOCAL, path="/path/1"),
+            LoraConfig(name="lora2", source=LoraSource.LOCAL, path="/path/2"),
+        ]
+
+        mixin.set_static_loras(loras)
+
+        assert mixin._static_lora_configs == loras
+        # Verify it's a copy, not the same list
+        assert mixin._static_lora_configs is not loras
+
+    def test_set_static_loras_empty_list(self):
+        """set_static_loras handles empty list."""
+        from oneiro.pipelines.lora import LoraLoaderMixin
+
+        class TestMixin(LoraLoaderMixin):
+            pass
+
+        mixin = TestMixin()
+        mixin.set_static_loras([])
+
+        assert mixin._static_lora_configs == []
+
+    def test_restore_static_loras_no_op_when_empty(self):
+        """restore_static_loras is no-op when no static loras and no loaded adapters."""
+        from oneiro.pipelines.lora import LoraLoaderMixin
+
+        class TestMixin(LoraLoaderMixin):
+            pass
+
+        mixin = TestMixin()
+        mixin.unload_loras = Mock()
+        mixin.load_loras_sync = Mock()
+        mixin.set_lora_adapters = Mock()
+
+        mixin.restore_static_loras()
+
+        mixin.unload_loras.assert_not_called()
+        mixin.load_loras_sync.assert_not_called()
+        mixin.set_lora_adapters.assert_not_called()
+
+    def test_restore_static_loras_resets_weights_when_adapters_match(self):
+        """restore_static_loras only resets weights when adapters match static config."""
+        from oneiro.pipelines.lora import LoraLoaderMixin
+
+        class TestMixin(LoraLoaderMixin):
+            pass
+
+        mixin = TestMixin()
+        loras = [
+            LoraConfig(name="lora1", source=LoraSource.LOCAL, path="/path/1", weight=0.8),
+            LoraConfig(name="lora2", source=LoraSource.LOCAL, path="/path/2", weight=0.5),
+        ]
+        mixin.set_static_loras(loras)
+        mixin._loaded_adapters = ["lora1", "lora2"]
+
+        mixin.unload_loras = Mock()
+        mixin.load_loras_sync = Mock()
+        mixin.set_lora_adapters = Mock()
+
+        mixin.restore_static_loras()
+
+        mixin.unload_loras.assert_not_called()
+        mixin.load_loras_sync.assert_not_called()
+        mixin.set_lora_adapters.assert_called_once_with(["lora1", "lora2"], [0.8, 0.5])
+
+    def test_restore_static_loras_reloads_when_adapters_differ(self):
+        """restore_static_loras reloads when loaded adapters differ from static."""
+        from oneiro.pipelines.lora import LoraLoaderMixin
+
+        class TestMixin(LoraLoaderMixin):
+            pass
+
+        mixin = TestMixin()
+        static_loras = [
+            LoraConfig(name="static-lora", source=LoraSource.LOCAL, path="/path/1"),
+        ]
+        mixin.set_static_loras(static_loras)
+        # Simulate dynamic lora was added
+        mixin._loaded_adapters = ["static-lora", "dynamic-lora"]
+
+        mixin.unload_loras = Mock()
+        mixin.load_loras_sync = Mock()
+        mixin.set_lora_adapters = Mock()
+
+        mixin.restore_static_loras()
+
+        mixin.unload_loras.assert_called_once()
+        mixin.load_loras_sync.assert_called_once_with(static_loras)
+        mixin.set_lora_adapters.assert_not_called()
+
+    def test_restore_static_loras_unloads_only_when_no_static(self):
+        """restore_static_loras unloads all when no static but has loaded adapters."""
+        from oneiro.pipelines.lora import LoraLoaderMixin
+
+        class TestMixin(LoraLoaderMixin):
+            pass
+
+        mixin = TestMixin()
+        mixin._loaded_adapters = ["dynamic-lora"]
+
+        mixin.unload_loras = Mock()
+        mixin.load_loras_sync = Mock()
+        mixin.set_lora_adapters = Mock()
+
+        mixin.restore_static_loras()
+
+        mixin.unload_loras.assert_called_once()
+        mixin.load_loras_sync.assert_not_called()
+
+    def test_restore_static_loras_uses_adapter_name_when_set(self):
+        """restore_static_loras uses adapter_name if specified, otherwise name."""
+        from oneiro.pipelines.lora import LoraLoaderMixin
+
+        class TestMixin(LoraLoaderMixin):
+            pass
+
+        mixin = TestMixin()
+        loras = [
+            LoraConfig(
+                name="lora1",
+                source=LoraSource.LOCAL,
+                path="/path/1",
+                adapter_name="custom_adapter",
+                weight=0.7,
+            ),
+        ]
+        mixin.set_static_loras(loras)
+        mixin._loaded_adapters = ["custom_adapter"]
+
+        mixin.unload_loras = Mock()
+        mixin.load_loras_sync = Mock()
+        mixin.set_lora_adapters = Mock()
+
+        mixin.restore_static_loras()
+
+        mixin.set_lora_adapters.assert_called_once_with(["custom_adapter"], [0.7])