Merge pull request #64 from jkoelker/jk/refactor

refactor(pipelines): generate() templates

Author: jkoelker

.gitignore

@@ -95,3 +95,4 @@ uv.lock
 
 .beads
 .gitattributes
+.beads/artifacts/

src/oneiro/pipelines/base.py

@@ -44,7 +44,6 @@ def load(self, model_config: dict[str, Any], full_config: dict[str, Any] | None
             full_config: Full configuration dict (for accessing global sections like embeddings)
         """
 
-    @abstractmethod
     def generate(
         self,
         prompt: str,
@@ -56,7 +55,142 @@ def generate(
         guidance_scale: float = 0.0,
         **kwargs: Any,
     ) -> GenerationResult:
-        """Generate an image."""
+        """Generate an image using Template Method pattern.
+
+        Subclasses should override hooks (validate_pipeline, pre_generate,
+        build_generation_kwargs, run_inference, build_result, post_generate),
+        not this method.
+        """
+        self.validate_pipeline()
+        self.pre_generate(**kwargs)
+        try:
+            actual_seed, generator = self._prepare_seed(seed)
+            # Pop init_image and strength from kwargs to avoid passing twice
+            init_image = self._load_init_image(kwargs.pop("init_image", None))
+            strength = kwargs.pop("strength", 0.75)
+
+            gen_kwargs = self.build_generation_kwargs(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                width=width,
+                height=height,
+                steps=steps,
+                guidance_scale=guidance_scale,
+                generator=generator,
+                init_image=init_image,
+                strength=strength,
+                **kwargs,
+            )
+            is_img2img = init_image is not None
+            result = self.run_inference(gen_kwargs, is_img2img)
+
+            DevicePolicy.clear_cache()
+            return self.build_result(
+                result=result,
+                seed=actual_seed,
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                steps=steps,
+                guidance_scale=guidance_scale,
+            )
+        finally:
+            self.post_generate(**kwargs)
+
+    def validate_pipeline(self) -> None:
+        """Validate pipeline is ready for generation.
+
+        This is called before pre_generate() and before any kwargs are consumed.
+        Override for additional validation checks (e.g., config state validation).
+        """
+        if self.pipe is None:
+            raise RuntimeError("Pipeline not loaded")
+
+    def pre_generate(self, **kwargs: Any) -> None:  # noqa: B027
+        """Pre-generation hook called before building kwargs.
+
+        Override for scheduler/LoRA setup or other pre-processing.
+        This is an optional hook with a no-op default; it is intentionally
+        not abstract so subclasses can choose whether to implement it.
+
+        Note: This method may pop keys from kwargs to consume them before
+        build_generation_kwargs() is called. The modified kwargs are then
+        passed through to build_generation_kwargs() and post_generate().
+        """
+        pass
+
+    @abstractmethod
+    def build_generation_kwargs(
+        self,
+        prompt: str,
+        negative_prompt: str | None,
+        width: int,
+        height: int,
+        steps: int,
+        guidance_scale: float,
+        generator: torch.Generator,
+        init_image: Image.Image | None,
+        strength: float,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """Build pipeline-specific generation kwargs.
+
+        This is the REQUIRED hook that each subclass must implement.
+        Return a dict to be passed to self.pipe().
+        """
+
+    def run_inference(self, gen_kwargs: dict[str, Any], is_img2img: bool) -> Any:
+        """Run the diffusion pipeline.
+
+        Args:
+            gen_kwargs: Keyword arguments to pass to the underlying pipeline.
+            is_img2img: Whether this is an image-to-image generation. This flag
+                is not used by the base implementation but is provided for
+                subclasses that need to branch on img2img vs txt2img behavior.
+
+        Returns:
+            Pipeline output (typically has .images attribute).
+
+        Override if the pipeline call signature or behavior differs.
+        """
+        return self.pipe(**gen_kwargs)
+
+    def build_result(
+        self,
+        result: Any,
+        seed: int,
+        prompt: str,
+        negative_prompt: str | None,
+        steps: int,
+        guidance_scale: float,
+    ) -> GenerationResult:
+        """Build GenerationResult from pipeline output.
+
+        Override if result format differs.
+        """
+        output_image = result.images[0]
+        return GenerationResult(
+            image=output_image,
+            seed=seed,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            width=output_image.width,
+            height=output_image.height,
+            steps=steps,
+            guidance_scale=guidance_scale,
+        )
+
+    def post_generate(self, **kwargs: Any) -> None:  # noqa: B027
+        """Post-generation cleanup hook called after generation completes.
+
+        Override for LoRA restore or other cleanup. This is an optional hook
+        with a no-op default; it is intentionally not abstract so subclasses
+        can choose whether to implement it.
+
+        Note: The kwargs passed here have already had 'init_image' and 'strength'
+        removed by generate(). If a subclass needs access to these values,
+        it should save them in pre_generate() before they are consumed.
+        """
+        pass
 
     def unload(self) -> None:
         """Free GPU memory."""

src/oneiro/pipelines/civitai_checkpoint.py

@@ -10,6 +10,9 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
+import torch
+from PIL import Image
+
 from oneiro.device import DevicePolicy, OffloadMode
 from oneiro.pipelines.base import BasePipeline, GenerationResult
 from oneiro.pipelines.embedding import EmbeddingLoaderMixin, parse_embeddings_from_config
@@ -519,6 +522,7 @@ def __init__(self) -> None:
         self._current_scheduler: str | None = None
         self._static_lora_configs: list[LoraConfig] = []
         self._cpu_offload: bool = False
+        self._has_dynamic_loras: 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).
@@ -692,6 +696,12 @@ def configure_scheduler(self, scheduler_name: str | None) -> None:
         self._current_scheduler = scheduler_name
         print(f"  Scheduler: {scheduler_name}")
 
+    def validate_pipeline(self) -> None:
+        """Validate pipeline and config are ready for generation."""
+        super().validate_pipeline()
+        if self._pipeline_config is None:
+            raise RuntimeError("Pipeline config not initialized")
+
     def generate(
         self,
         prompt: str,
@@ -762,84 +772,70 @@ def generate(
             as the actual seed used, prompts, final image size, number of
             steps, and guidance scale.
         """
-        if self.pipe is None:
-            raise RuntimeError("Pipeline not loaded")
+        # Apply defaults from pipeline config (validation happens in super().generate())
+        # Note: We need to check _pipeline_config here before applying defaults,
+        # but full validation happens in validate_pipeline() called by super()
+        if self._pipeline_config is not None:
+            width = width if width is not None else self._pipeline_config.default_width
+            height = height if height is not None else self._pipeline_config.default_height
+            steps = steps if steps is not None else self._pipeline_config.default_steps
+            guidance_scale = (
+                guidance_scale
+                if guidance_scale is not None
+                else self._pipeline_config.default_guidance_scale
+            )
 
-        if self._pipeline_config is None:
-            raise RuntimeError("Pipeline config not initialized")
+        return super().generate(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            width=width or 1024,
+            height=height or 1024,
+            seed=seed,
+            steps=steps or 20,
+            guidance_scale=guidance_scale if guidance_scale is not None else 7.0,
+            **kwargs,
+        )
 
+    def pre_generate(self, **kwargs: Any) -> None:
+        """Pre-generation setup: scheduler override and dynamic LoRA loading."""
         scheduler_override = kwargs.pop("scheduler", None)
         if scheduler_override:
             self.configure_scheduler(scheduler_override)
 
         dynamic_loras = kwargs.pop("loras", None)
-        has_dynamic_loras = False
+        self._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
+            self._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
+                self._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(
+    def build_generation_kwargs(
         self,
         prompt: str,
         negative_prompt: str | None,
-        width: int | None,
-        height: int | None,
-        seed: int,
-        steps: int | None,
-        guidance_scale: float | None,
+        width: int,
+        height: int,
+        steps: int,
+        guidance_scale: float,
+        generator: "torch.Generator",
+        init_image: "Image.Image | None",
+        strength: float,
         **kwargs: Any,
-    ) -> GenerationResult:
+    ) -> dict[str, Any]:
+        """Build generation kwargs with embedding support."""
         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
-        guidance_scale = (
-            guidance_scale
-            if guidance_scale is not None
-            else self._pipeline_config.default_guidance_scale
-        )
-
-        actual_seed, generator = self._prepare_seed(seed)
 
-        # Handle img2img
-        init_image = self._load_init_image(kwargs.get("init_image"))
-        strength = kwargs.get("strength", 0.75)
-
-        # Build generation kwargs
         gen_kwargs: dict[str, Any] = {
             "num_inference_steps": steps,
             "guidance_scale": guidance_scale,
             "generator": generator,
         }
 
         # Use embedding-based prompt handling for pipelines that support it
-        # (SD 1.x, SD 2.x, SDXL, Flux, SD3) - enables weight syntax like (word:1.5)
         if self._supports_prompt_embeddings():
             gen_kwargs.update(self._encode_prompts_to_embeddings(prompt, negative_prompt))
         else:
@@ -849,29 +845,21 @@ def _run_generation(
                 gen_kwargs["negative_prompt"] = negative_prompt
 
         if init_image:
-            print(f"CivitAI img2img: '{prompt[:50]}...' seed={actual_seed} strength={strength}")
+            print(f"CivitAI img2img: '{prompt[:50]}...' strength={strength}")
             gen_kwargs["image"] = init_image
             gen_kwargs["strength"] = strength
         else:
-            print(f"CivitAI generating: '{prompt[:50]}...' seed={actual_seed}")
+            print(f"CivitAI generating: '{prompt[:50]}...'")
             gen_kwargs["height"] = height
             gen_kwargs["width"] = width
 
-        result = self.pipe(**gen_kwargs)
-
-        DevicePolicy.clear_cache()
+        return gen_kwargs
 
-        output_image = result.images[0]
-        return GenerationResult(
-            image=output_image,
-            seed=actual_seed,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            width=output_image.width,
-            height=output_image.height,
-            steps=steps,
-            guidance_scale=guidance_scale,
-        )
+    def post_generate(self, **kwargs: Any) -> None:
+        """Post-generation cleanup: restore static LoRAs if dynamic were used."""
+        if self._has_dynamic_loras:
+            self._restore_static_loras()
+            self._has_dynamic_loras = False
 
     def _load_dynamic_loras(self, loras: list[LoraConfig]) -> None:
         if self.pipe is None or not loras: