Create image_generator.py

image_generator.py

@@ -0,0 +1,284 @@
+"""
+TimeLapseForge — Image Generator Module
+Uses SDXL Turbo / SDXL Base for consistent timelapse panel generation.
+Key technique: img2img chaining with low strength for natural transitions.
+"""
+
+import torch
+import gc
+from typing import List, Dict, Optional, Callable, Tuple
+from PIL import Image
+
+# Model cache
+_pipelines = {}
+
+
+def flush_memory():
+    """Clear GPU and CPU memory caches."""
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def get_pipeline(model_id: str = "stabilityai/sdxl-turbo", pipeline_type: str = "t2i"):
+    """
+    Load and cache the diffusion pipeline.
+    Uses from_pipe for img2img to share components and save memory.
+    """
+    global _pipelines
+    from diffusers import AutoPipelineForText2Image, AutoPipelineForImage2Image
+
+    cache_key = f"{model_id}_{pipeline_type}"
+
+    if cache_key not in _pipelines:
+        if pipeline_type == "t2i":
+            pipe = AutoPipelineForText2Image.from_pretrained(
+                model_id,
+                torch_dtype=torch.float16,
+                variant="fp16" if "turbo" in model_id or "sdxl" in model_id.lower() else None,
+            )
+            _pipelines[cache_key] = pipe
+
+        elif pipeline_type == "i2i":
+            # Try to reuse t2i components
+            t2i_key = f"{model_id}_t2i"
+            if t2i_key in _pipelines:
+                pipe = AutoPipelineForImage2Image.from_pipe(_pipelines[t2i_key])
+            else:
+                pipe = AutoPipelineForImage2Image.from_pretrained(
+                    model_id,
+                    torch_dtype=torch.float16,
+                    variant="fp16" if "turbo" in model_id or "sdxl" in model_id.lower() else None,
+                )
+            _pipelines[cache_key] = pipe
+
+    return _pipelines[cache_key]
+
+
+def get_model_config(model_id: str) -> Dict:
+    """Get optimal generation parameters for each model."""
+    configs = {
+        "stabilityai/sdxl-turbo": {
+            "num_inference_steps": 4,
+            "guidance_scale": 0.0,
+            "default_resolution": (512, 512),
+            "supports_guidance": False,
+        },
+        "stabilityai/stable-diffusion-xl-base-1.0": {
+            "num_inference_steps": 25,
+            "guidance_scale": 7.5,
+            "default_resolution": (1024, 1024),
+            "supports_guidance": True,
+        },
+        "runwayml/stable-diffusion-v1-5": {
+            "num_inference_steps": 25,
+            "guidance_scale": 7.5,
+            "default_resolution": (512, 512),
+            "supports_guidance": True,
+        },
+    }
+    # Default config for unknown models
+    return configs.get(model_id, {
+        "num_inference_steps": 20,
+        "guidance_scale": 7.0,
+        "default_resolution": (512, 512),
+        "supports_guidance": True,
+    })
+
+
+class ImageGenerator:
+    """
+    Generates timelapse panels using Stable Diffusion.
+    Uses img2img chaining for visual consistency between panels.
+    """
+
+    def __init__(self, model_id: str = "stabilityai/sdxl-turbo"):
+        self.model_id = model_id
+        self.config = get_model_config(model_id)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    def _move_to_device(self, pipe):
+        """Move pipeline to the appropriate device."""
+        if self.device == "cuda":
+            pipe.to("cuda")
+        return pipe
+
+    def generate_first_panel(
+        self,
+        prompt: str,
+        negative_prompt: str = "",
+        seed: int = 42,
+        width: Optional[int] = None,
+        height: Optional[int] = None,
+        steps: Optional[int] = None,
+        guidance: Optional[float] = None,
+    ) -> Image.Image:
+        """Generate the first panel using text-to-image."""
+        pipe = get_pipeline(self.model_id, "t2i")
+        pipe = self._move_to_device(pipe)
+
+        w, h = width or self.config["default_resolution"][0], height or self.config["default_resolution"][1]
+        n_steps = steps or self.config["num_inference_steps"]
+        cfg = guidance if guidance is not None else self.config["guidance_scale"]
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+
+        result = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt if negative_prompt else None,
+            num_inference_steps=n_steps,
+            guidance_scale=cfg,
+            width=w,
+            height=h,
+            generator=generator,
+        )
+        return result.images[0]
+
+    def generate_next_panel(
+        self,
+        prompt: str,
+        previous_image: Image.Image,
+        negative_prompt: str = "",
+        strength: float = 0.4,
+        seed: int = 42,
+        steps: Optional[int] = None,
+        guidance: Optional[float] = None,
+    ) -> Image.Image:
+        """Generate the next panel using img2img from the previous panel."""
+        pipe = get_pipeline(self.model_id, "i2i")
+        pipe = self._move_to_device(pipe)
+
+        n_steps = steps or self.config["num_inference_steps"]
+        cfg = guidance if guidance is not None else self.config["guidance_scale"]
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+
+        # Resize previous image to match model's expected resolution
+        target_w, target_h = self.config["default_resolution"]
+        prev_resized = previous_image.resize((target_w, target_h), Image.LANCZOS)
+
+        result = pipe(
+            prompt=prompt,
+            image=prev_resized,
+            negative_prompt=negative_prompt if negative_prompt else None,
+            num_inference_steps=n_steps,
+            guidance_scale=cfg,
+            strength=strength,
+            generator=generator,
+        )
+        return result.images[0]
+
+    def generate_all_panels(
+        self,
+        prompts: List[Dict[str, str]],
+        strength: float = 0.4,
+        base_seed: int = 42,
+        steps: Optional[int] = None,
+        guidance: Optional[float] = None,
+        width: Optional[int] = None,
+        height: Optional[int] = None,
+        progress_callback: Optional[Callable] = None,
+        reference_image: Optional[Image.Image] = None,
+    ) -> List[Image.Image]:
+        """
+        Generate all panels in sequence.
+        Panel 1: text-to-image (or img2img from reference if provided)
+        Panel 2+: img2img from previous panel with controlled strength
+        """
+        images = []
+
+        for i, prompt_data in enumerate(prompts):
+            main_prompt = prompt_data.get("main_prompt", "")
+            style = prompt_data.get("style_suffix", "")
+            full_prompt = f"{main_prompt}, {style}" if style else main_prompt
+            neg_prompt = prompt_data.get("negative_prompt", "")
+            seed = base_seed + i
+
+            try:
+                if i == 0:
+                    if reference_image is not None:
+                        # Use reference image as base for first panel
+                        img = self.generate_next_panel(
+                            prompt=full_prompt,
+                            previous_image=reference_image,
+                            negative_prompt=neg_prompt,
+                            strength=max(strength, 0.5),  # Slightly higher for first panel from ref
+                            seed=seed,
+                            steps=steps,
+                            guidance=guidance,
+                        )
+                    else:
+                        img = self.generate_first_panel(
+                            prompt=full_prompt,
+                            negative_prompt=neg_prompt,
+                            seed=seed,
+                            width=width,
+                            height=height,
+                            steps=steps,
+                            guidance=guidance,
+                        )
+                else:
+                    img = self.generate_next_panel(
+                        prompt=full_prompt,
+                        previous_image=images[-1],
+                        negative_prompt=neg_prompt,
+                        strength=strength,
+                        seed=seed,
+                        steps=steps,
+                        guidance=guidance,
+                    )
+
+                images.append(img)
+
+            except Exception as e:
+                print(f"Error generating panel {i + 1}: {e}")
+                # Use previous image as fallback (or a blank image for panel 1)
+                if images:
+                    images.append(images[-1].copy())
+                else:
+                    fallback = Image.new("RGB", self.config["default_resolution"], (50, 50, 50))
+                    images.append(fallback)
+
+            if progress_callback:
+                progress_callback(i + 1, len(prompts))
+
+        flush_memory()
+        return images
+
+    def regenerate_single_panel(
+        self,
+        panel_index: int,
+        prompts: List[Dict[str, str]],
+        existing_images: List[Image.Image],
+        strength: float = 0.4,
+        base_seed: int = 42,
+        steps: Optional[int] = None,
+        guidance: Optional[float] = None,
+    ) -> Tuple[Image.Image, List[Image.Image]]:
+        """
+        Regenerate a single panel and optionally cascade changes forward.
+        Returns the new image and the updated image list.
+        """
+        prompt_data = prompts[panel_index]
+        main_prompt = prompt_data.get("main_prompt", "")
+        style = prompt_data.get("style_suffix", "")
+        full_prompt = f"{main_prompt}, {style}" if style else main_prompt
+        neg_prompt = prompt_data.get("negative_prompt", "")
+        seed = base_seed + panel_index
+
+        if panel_index == 0:
+            new_img = self.generate_first_panel(
+                prompt=full_prompt, negative_prompt=neg_prompt, seed=seed, steps=steps, guidance=guidance
+            )
+        else:
+            new_img = self.generate_next_panel(
+                prompt=full_prompt,
+                previous_image=existing_images[panel_index - 1],
+                negative_prompt=neg_prompt,
+                strength=strength,
+                seed=seed,
+                steps=steps,
+                guidance=guidance,
+            )
+
+        updated_images = existing_images.copy()
+        updated_images[panel_index] = new_img
+        return new_img, updated_images