Merge pull request #58 from jkoelker/jk/lora-fix

fix: resolve device mismatch when loading auto-detected LoRAs

Author: jkoelker

src/oneiro/pipelines/__init__.py

@@ -152,6 +152,30 @@ async def generate(
         if self.pipeline is None:
             raise RuntimeError("No pipeline loaded")
 
+        loras: list[LoraConfig] | None = kwargs.pop("loras", None)
+        if loras:
+            pipeline_type = None
+            if self.current_model:
+                model_config = self.config.get("models", self.current_model)
+                if model_config:
+                    pipeline_type = model_config.get("type")
+
+            resolved_loras: list[LoraConfig] = []
+            for lora in loras:
+                try:
+                    await resolve_lora_path(
+                        lora,
+                        civitai_client=self._civitai_client,
+                        pipeline_type=pipeline_type,
+                        validate_compatibility=True,
+                    )
+                    resolved_loras.append(lora)
+                except Exception as e:
+                    print(f"Warning: Failed to resolve LoRA {lora.name}: {e}")
+
+            if resolved_loras:
+                kwargs["loras"] = resolved_loras
+
         return await asyncio.to_thread(
             self.pipeline.generate,
             prompt,

src/oneiro/pipelines/civitai_checkpoint.py

@@ -20,7 +20,7 @@
     get_weighted_text_embeddings_sd15,
     get_weighted_text_embeddings_sdxl,
 )
-from oneiro.pipelines.lora import LoraLoaderMixin, parse_loras_from_model_config
+from oneiro.pipelines.lora import LoraConfig, LoraLoaderMixin, parse_loras_from_model_config
 
 if TYPE_CHECKING:
     from oneiro.civitai import CivitaiClient, ModelVersion
@@ -518,6 +518,8 @@ def __init__(self) -> None:
         self._base_model: str | None = None
         self._full_config: dict[str, Any] | None = None
         self._current_scheduler: str | None = None
+        self._static_lora_configs: list[LoraConfig] = []
+        self._cpu_offload: bool = False
 
     def load(self, model_config: dict[str, Any], full_config: dict[str, Any] | None = None) -> None:
         """Load checkpoint from config (synchronous, requires checkpoint_path).
@@ -631,7 +633,8 @@ def _load_from_path(self, checkpoint_path: Path, model_config: dict[str, Any]) -
 
         # Apply optimizations
         cpu_offload = model_config.get("cpu_offload", True)
-        if cpu_offload and self._device == "cuda":
+        self._cpu_offload = cpu_offload and self._device == "cuda"
+        if self._cpu_offload:
             self.pipe.enable_model_cpu_offload()
         elif self._device == "cuda":
             self.pipe.to("cuda")
@@ -643,11 +646,13 @@ def _load_from_path(self, checkpoint_path: Path, model_config: dict[str, Any]) -
             if hasattr(self.pipe.vae, "enable_slicing"):
                 self.pipe.vae.enable_slicing()
 
-        # Load LoRAs if configured
         loras = parse_loras_from_model_config(model_config)
         if loras:
             print(f"  Loading {len(loras)} LoRA(s)...")
             self.load_loras_sync(loras)
+            self._static_lora_configs = list(loras)
+        else:
+            self._static_lora_configs = []
 
         # Load embeddings if full_config provided
         if self._full_config:
@@ -743,6 +748,8 @@ def generate(
               pass the resulting image to the pipeline.
             - ``strength``: Strength parameter for img2img generation, used
               when ``init_image`` is provided. Defaults to ``0.75``.
+            - ``loras``: List of LoraConfig objects for dynamic LoRA loading.
+              These are loaded before generation and unloaded after.
 
             Any other keyword arguments are passed through unchanged to the
             underlying diffusers pipeline call and may be used to control
@@ -765,7 +772,50 @@ def generate(
         if scheduler_override:
             self.configure_scheduler(scheduler_override)
 
-        # Use defaults from pipeline config
+        dynamic_loras = kwargs.pop("loras", None)
+        has_dynamic_loras = False
+        if dynamic_loras:
+            # Mark that we are entering a dynamic LoRA context before loading,
+            # so that failures during loading can be properly rolled back.
+            has_dynamic_loras = True
+            try:
+                self._load_dynamic_loras(dynamic_loras)
+            except Exception:
+                # If loading dynamic LoRAs fails after modifying the pipeline
+                # state (for example, after unloading static LoRAs), attempt
+                # to restore the original static LoRAs before propagating
+                # the error.
+                self._restore_static_loras()
+                has_dynamic_loras = False
+                raise
+
+        try:
+            return self._run_generation(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                width=width,
+                height=height,
+                seed=seed,
+                steps=steps,
+                guidance_scale=guidance_scale,
+                **kwargs,
+            )
+        finally:
+            if has_dynamic_loras:
+                self._restore_static_loras()
+
+    def _run_generation(
+        self,
+        prompt: str,
+        negative_prompt: str | None,
+        width: int | None,
+        height: int | None,
+        seed: int,
+        steps: int | None,
+        guidance_scale: float | None,
+        **kwargs: Any,
+    ) -> GenerationResult:
+        assert self._pipeline_config is not None
         width = width or self._pipeline_config.default_width
         height = height or self._pipeline_config.default_height
         steps = steps or self._pipeline_config.default_steps
@@ -824,6 +874,47 @@ def generate(
             guidance_scale=guidance_scale,
         )
 
+    def _load_dynamic_loras(self, loras: list[LoraConfig]) -> None:
+        if self.pipe is None or not loras:
+            return
+
+        lora_configs: list[LoraConfig] = [cfg for cfg in loras if isinstance(cfg, LoraConfig)]
+        if not lora_configs:
+            return
+
+        self.unload_loras()
+
+        # Only move pipeline to device manually when CPU offload is not enabled.
+        # With CPU offload, diffusers manages device placement automatically.
+        if not self._cpu_offload:
+            self.pipe.to(self._device)
+
+        loaded_names: list[str] = []
+        loaded_weights: list[float] = []
+
+        for lora in lora_configs:
+            try:
+                name = self.load_single_lora(lora)
+                loaded_names.append(name)
+                loaded_weights.append(lora.weight)
+                print(f"Loaded dynamic LoRA: {lora.name} (weight={lora.weight})")
+            except Exception as e:
+                print(f"Warning: Failed to load LoRA {lora.name}: {e}")
+
+        if loaded_names:
+            self.set_lora_adapters(loaded_names, loaded_weights)
+
+    def _restore_static_loras(self) -> None:
+        self.unload_loras()
+
+        if not self._static_lora_configs:
+            return
+
+        if not self._cpu_offload:
+            self.pipe.to(self._device)
+        self.load_loras_sync(self._static_lora_configs)
+        print(f"Restored {len(self._static_lora_configs)} static LoRA(s)")
+
     @property
     def pipeline_config(self) -> PipelineConfig | None:
         """Get the current pipeline configuration."""

tests/test_civitai_checkpoint.py

@@ -1,5 +1,6 @@
 """Tests for CivitAI checkpoint pipeline."""
 
+from pathlib import Path
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -14,6 +15,7 @@
     get_diffusers_pipeline_class,
     get_pipeline_config_for_base_model,
 )
+from oneiro.pipelines.lora import LoraConfig, LoraSource
 
 
 class TestPipelineConfig:
@@ -1156,3 +1158,166 @@ def test_falls_back_to_sd15_for_unknown_pipeline(self):
         mock_func.assert_called_once()
         assert result["prompt_embeds"] is mock_prompt
         assert result["negative_prompt_embeds"] is mock_neg_prompt
+
+
+class TestDynamicLoraGeneration:
+    """Tests for dynamic LoRA loading during generation."""
+
+    def _create_pipeline_with_mocks(self):
+        """Create a pipeline with common mocks for dynamic LoRA tests."""
+        pipeline = CivitaiCheckpointPipeline()
+        pipeline._pipeline_config = PipelineConfig(
+            pipeline_class="StableDiffusionXLPipeline",
+            default_steps=25,
+            default_guidance_scale=7.0,
+            default_width=1024,
+            default_height=1024,
+        )
+        mock_pipe = MagicMock()
+        mock_image = MagicMock()
+        mock_image.width = 1024
+        mock_image.height = 1024
+        mock_pipe.return_value.images = [mock_image]
+        pipeline.pipe = mock_pipe
+        pipeline._cpu_offload = False
+        return pipeline
+
+    def test_generate_with_dynamic_loras(self):
+        """generate() loads dynamic LoRAs passed via kwargs."""
+        pipeline = self._create_pipeline_with_mocks()
+
+        lora = LoraConfig(name="test-lora", source=LoraSource.LOCAL, path="/fake/path.safetensors")
+        lora._resolved_path = Path("/fake/path.safetensors")
+
+        with (
+            patch("oneiro.pipelines.civitai_checkpoint.torch"),
+            patch.object(pipeline, "_encode_prompts_to_embeddings"),
+            patch.object(pipeline, "_load_dynamic_loras") as mock_load,
+            patch.object(pipeline, "_restore_static_loras") as mock_restore,
+        ):
+            pipeline.generate("test prompt", loras=[lora])
+
+        mock_load.assert_called_once_with([lora])
+        mock_restore.assert_called_once()
+
+    def test_generate_restores_static_loras_after_dynamic(self):
+        """generate() restores static LoRAs after using dynamic ones."""
+        pipeline = self._create_pipeline_with_mocks()
+
+        static_lora = LoraConfig(
+            name="static-lora", source=LoraSource.LOCAL, path="/static.safetensors"
+        )
+        pipeline._static_lora_configs = [static_lora]
+
+        dynamic_lora = LoraConfig(
+            name="dynamic-lora", source=LoraSource.LOCAL, path="/dynamic.safetensors"
+        )
+        dynamic_lora._resolved_path = Path("/dynamic.safetensors")
+
+        with (
+            patch("oneiro.pipelines.civitai_checkpoint.torch"),
+            patch.object(pipeline, "_encode_prompts_to_embeddings"),
+            patch.object(pipeline, "unload_loras") as mock_unload,
+            patch.object(pipeline, "load_single_lora", return_value="dynamic-lora"),
+            patch.object(pipeline, "set_lora_adapters"),
+            patch.object(pipeline, "load_loras_sync") as mock_load_sync,
+        ):
+            pipeline.generate("test prompt", loras=[dynamic_lora])
+
+        assert mock_unload.call_count == 2
+        mock_load_sync.assert_called_once_with([static_lora])
+
+    def test_generate_handles_dynamic_lora_loading_failure(self):
+        """generate() restores static LoRAs when dynamic loading fails."""
+        pipeline = self._create_pipeline_with_mocks()
+
+        static_lora = LoraConfig(
+            name="static-lora", source=LoraSource.LOCAL, path="/static.safetensors"
+        )
+        pipeline._static_lora_configs = [static_lora]
+
+        dynamic_lora = LoraConfig(name="bad-lora", source=LoraSource.LOCAL, path="/bad.safetensors")
+        dynamic_lora._resolved_path = Path("/bad.safetensors")
+
+        with (
+            patch("oneiro.pipelines.civitai_checkpoint.torch"),
+            patch.object(pipeline, "_encode_prompts_to_embeddings"),
+            patch.object(pipeline, "_load_dynamic_loras", side_effect=RuntimeError("Load failed")),
+            patch.object(pipeline, "_restore_static_loras") as mock_restore,
+        ):
+            with pytest.raises(RuntimeError, match="Load failed"):
+                pipeline.generate("test prompt", loras=[dynamic_lora])
+
+        mock_restore.assert_called_once()
+
+    def test_generate_cleanup_on_generation_failure(self):
+        """generate() cleans up dynamic LoRAs even if generation fails."""
+        pipeline = self._create_pipeline_with_mocks()
+
+        dynamic_lora = LoraConfig(
+            name="dynamic-lora", source=LoraSource.LOCAL, path="/dynamic.safetensors"
+        )
+        dynamic_lora._resolved_path = Path("/dynamic.safetensors")
+
+        pipeline.pipe.side_effect = RuntimeError("Generation failed")
+
+        with (
+            patch("oneiro.pipelines.civitai_checkpoint.torch"),
+            patch.object(pipeline, "_encode_prompts_to_embeddings"),
+            patch.object(pipeline, "_load_dynamic_loras"),
+            patch.object(pipeline, "_restore_static_loras") as mock_restore,
+        ):
+            with pytest.raises(RuntimeError, match="Generation failed"):
+                pipeline.generate("test prompt", loras=[dynamic_lora])
+
+        mock_restore.assert_called_once()
+
+    def test_generate_without_dynamic_loras_skips_lora_handling(self):
+        """generate() skips LoRA handling when no dynamic LoRAs provided."""
+        pipeline = self._create_pipeline_with_mocks()
+
+        with (
+            patch("oneiro.pipelines.civitai_checkpoint.torch"),
+            patch.object(pipeline, "_encode_prompts_to_embeddings"),
+            patch.object(pipeline, "_load_dynamic_loras") as mock_load,
+            patch.object(pipeline, "_restore_static_loras") as mock_restore,
+        ):
+            pipeline.generate("test prompt")
+
+        mock_load.assert_not_called()
+        mock_restore.assert_not_called()
+
+    def test_load_dynamic_loras_respects_cpu_offload(self):
+        """_load_dynamic_loras() skips .to(device) when cpu_offload enabled."""
+        pipeline = self._create_pipeline_with_mocks()
+        pipeline._cpu_offload = True
+
+        lora = LoraConfig(name="test-lora", source=LoraSource.LOCAL, path="/fake.safetensors")
+        lora._resolved_path = Path("/fake.safetensors")
+
+        with (
+            patch.object(pipeline, "unload_loras"),
+            patch.object(pipeline, "load_single_lora", return_value="test-lora"),
+            patch.object(pipeline, "set_lora_adapters"),
+        ):
+            pipeline._load_dynamic_loras([lora])
+
+        pipeline.pipe.to.assert_not_called()
+
+    def test_load_dynamic_loras_moves_to_device_without_cpu_offload(self):
+        """_load_dynamic_loras() calls .to(device) when cpu_offload disabled."""
+        pipeline = self._create_pipeline_with_mocks()
+        pipeline._cpu_offload = False
+        pipeline._device = "cuda"
+
+        lora = LoraConfig(name="test-lora", source=LoraSource.LOCAL, path="/fake.safetensors")
+        lora._resolved_path = Path("/fake.safetensors")
+
+        with (
+            patch.object(pipeline, "unload_loras"),
+            patch.object(pipeline, "load_single_lora", return_value="test-lora"),
+            patch.object(pipeline, "set_lora_adapters"),
+        ):
+            pipeline._load_dynamic_loras([lora])
+
+        pipeline.pipe.to.assert_called_once_with("cuda")