Update image_generator.py

image_generator.py

@@ -1,206 +1,196 @@
 """
-TimeLapseForge — Image Generator Module
-Uses SDXL Turbo / SDXL Base for consistent timelapse panel generation.
-Key technique: img2img chaining with low strength for natural transitions.
+TimeLapseForge — Image Generator Module v2.0
+Supports BOTH local (diffusers) and API-based generation.
+Unified interface for all providers.
 """
 
-import torch
 import gc
+import torch
 from typing import List, Dict, Optional, Callable, Tuple
 from PIL import Image
 
-# Model cache
-_pipelines = {}
+from api_providers import get_provider, BaseProvider, PROVIDERS
 
 
 def flush_memory():
-    """Clear GPU and CPU memory caches."""
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
 
 
-def get_device():
-    """Get the best available device."""
-    if torch.cuda.is_available():
-        return "cuda"
-    return "cpu"
+# ═══════════════════════════════════════════
+# LOCAL GENERATOR (Same as before — diffusers)
+# ═══════════════════════════════════════════
 
+_local_pipelines = {}
 
-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
+def get_local_pipeline(model_id, pipeline_type="t2i"):
+    global _local_pipelines
     from diffusers import AutoPipelineForText2Image, AutoPipelineForImage2Image
 
     cache_key = f"{model_id}_{pipeline_type}"
-
-    if cache_key not in _pipelines:
-        device = get_device()
+    if cache_key not in _local_pipelines:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
         dtype = torch.float16 if device == "cuda" else torch.float32
-        
-        # Determine variant
-        variant = None
-        if device == "cuda" and ("turbo" in model_id or "sdxl" in model_id.lower()):
-            variant = "fp16"
 
         if pipeline_type == "t2i":
             try:
                 pipe = AutoPipelineForText2Image.from_pretrained(
-                    model_id,
-                    torch_dtype=dtype,
-                    variant=variant,
+                    model_id, torch_dtype=dtype, variant="fp16"
                 )
             except Exception:
-                # Fallback without variant
                 pipe = AutoPipelineForText2Image.from_pretrained(
-                    model_id,
-                    torch_dtype=dtype,
+                    model_id, torch_dtype=dtype
                 )
             pipe = pipe.to(device)
-            _pipelines[cache_key] = pipe
+            _local_pipelines[cache_key] = pipe
 
         elif pipeline_type == "i2i":
-            # Try to reuse t2i components
             t2i_key = f"{model_id}_t2i"
-            if t2i_key in _pipelines:
+            if t2i_key in _local_pipelines:
                 try:
-                    pipe = AutoPipelineForImage2Image.from_pipe(_pipelines[t2i_key])
+                    pipe = AutoPipelineForImage2Image.from_pipe(_local_pipelines[t2i_key])
                 except Exception:
                     try:
                         pipe = AutoPipelineForImage2Image.from_pretrained(
-                            model_id, torch_dtype=dtype, variant=variant,
+                            model_id, torch_dtype=dtype, variant="fp16"
                         )
                     except Exception:
                         pipe = AutoPipelineForImage2Image.from_pretrained(
-                            model_id, torch_dtype=dtype,
+                            model_id, torch_dtype=dtype
                         )
                     pipe = pipe.to(device)
             else:
                 try:
                     pipe = AutoPipelineForImage2Image.from_pretrained(
-                        model_id, torch_dtype=dtype, variant=variant,
+                        model_id, torch_dtype=dtype, variant="fp16"
                     )
                 except Exception:
                     pipe = AutoPipelineForImage2Image.from_pretrained(
-                        model_id, torch_dtype=dtype,
+                        model_id, torch_dtype=dtype
                     )
                 pipe = pipe.to(device)
-            _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,
-        },
-    }
-    return configs.get(model_id, {
-        "num_inference_steps": 20,
-        "guidance_scale": 7.0,
+            _local_pipelines[cache_key] = pipe
+
+    return _local_pipelines[cache_key]
+
+
+LOCAL_MODELS = {
+    "stabilityai/sdxl-turbo": {
+        "num_inference_steps": 4, "guidance_scale": 0.0,
+        "default_resolution": (512, 512),
+    },
+    "stabilityai/stable-diffusion-xl-base-1.0": {
+        "num_inference_steps": 25, "guidance_scale": 7.5,
+        "default_resolution": (1024, 1024),
+    },
+    "runwayml/stable-diffusion-v1-5": {
+        "num_inference_steps": 25, "guidance_scale": 7.5,
         "default_resolution": (512, 512),
-        "supports_guidance": True,
-    })
+    },
+}
 
 
+# ═══════════════════════════════════════════
+# UNIFIED IMAGE GENERATOR
+# ═══════════════════════════════════════════
+
 class ImageGenerator:
     """
-    Generates timelapse panels using Stable Diffusion.
-    Uses img2img chaining for visual consistency between panels.
+    Unified image generator supporting both local models and API providers.
     """
 
-    def __init__(self, model_id: str = "stabilityai/sdxl-turbo"):
-        self.model_id = model_id
-        self.config = get_model_config(model_id)
-        self.device = get_device()
-
-    def generate_first_panel(
+    def __init__(
         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")
-
+        mode: str = "local",  # "local" or "api"
+        # Local params
+        local_model_id: str = "stabilityai/sdxl-turbo",
+        # API params
+        provider_name: str = "openai",
+        api_key: str = "",
+        api_model: str = "",
+        custom_base_url: str = "",
+        custom_endpoint_url: str = "",
+    ):
+        self.mode = mode
+
+        if mode == "local":
+            self.local_model_id = local_model_id
+            self.config = LOCAL_MODELS.get(local_model_id, {
+                "num_inference_steps": 20, "guidance_scale": 7.0,
+                "default_resolution": (512, 512),
+            })
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+            self.provider = None
+        else:
+            # API mode
+            kwargs = {}
+            if provider_name == "custom_openai":
+                kwargs["base_url"] = custom_base_url
+            elif provider_name == "direct_url":
+                kwargs["endpoint_url"] = custom_endpoint_url
+            self.provider = get_provider(provider_name, api_key, **kwargs)
+            self.api_model = api_model or self.provider.default_model
+            self.config = {"default_resolution": (1024, 1024)}
+            self.device = "cpu"
+
+    def _generate_local_t2i(self, prompt, negative_prompt="", seed=42,
+                            width=None, height=None, steps=None, guidance=None):
+        pipe = get_local_pipeline(self.local_model_id, "t2i")
         w = width or self.config["default_resolution"][0]
         h = 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)
-
-        kwargs = dict(
-            prompt=prompt,
-            num_inference_steps=n_steps,
-            guidance_scale=cfg,
-            width=w,
-            height=h,
-            generator=generator,
-        )
+        n_steps = steps or self.config.get("num_inference_steps", 20)
+        cfg = guidance if guidance is not None else self.config.get("guidance_scale", 7.0)
+        gen = torch.Generator(device=self.device).manual_seed(seed)
+
+        kwargs = dict(prompt=prompt, num_inference_steps=n_steps,
+                      guidance_scale=cfg, width=w, height=h, generator=gen)
         if negative_prompt:
             kwargs["negative_prompt"] = negative_prompt
-
-        result = pipe(**kwargs)
-        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")
-
-        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 resolution
+        return pipe(**kwargs).images[0]
+
+    def _generate_local_i2i(self, prompt, prev_image, strength=0.4,
+                            negative_prompt="", seed=42, steps=None, guidance=None):
+        pipe = get_local_pipeline(self.local_model_id, "i2i")
+        n_steps = steps or self.config.get("num_inference_steps", 20)
+        cfg = guidance if guidance is not None else self.config.get("guidance_scale", 7.0)
+        gen = torch.Generator(device=self.device).manual_seed(seed)
         target_w, target_h = self.config["default_resolution"]
-        prev_resized = previous_image.resize((target_w, target_h), Image.LANCZOS)
-
-        kwargs = dict(
-            prompt=prompt,
-            image=prev_resized,
-            num_inference_steps=n_steps,
-            guidance_scale=cfg,
-            strength=strength,
-            generator=generator,
-        )
+        prev_resized = prev_image.resize((target_w, target_h), Image.LANCZOS)
+
+        kwargs = dict(prompt=prompt, image=prev_resized, num_inference_steps=n_steps,
+                      guidance_scale=cfg, strength=strength, generator=gen)
         if negative_prompt:
             kwargs["negative_prompt"] = negative_prompt
+        return pipe(**kwargs).images[0]
+
+    def _generate_api_t2i(self, prompt, negative_prompt="", seed=None,
+                          width=1024, height=1024, **kwargs):
+        return self.provider.generate_image(
+            prompt=prompt, negative_prompt=negative_prompt,
+            width=width, height=height, seed=seed,
+            model=self.api_model, **kwargs,
+        )
 
-        result = pipe(**kwargs)
-        return result.images[0]
+    def _generate_api_i2i(self, prompt, prev_image, strength=0.4,
+                          negative_prompt="", seed=None, **kwargs):
+        if self.provider.supports_img2img:
+            return self.provider.img2img(
+                prompt=prompt, image=prev_image, strength=strength,
+                negative_prompt=negative_prompt, seed=seed,
+                model=self.api_model, **kwargs,
+            )
+        else:
+            # Fallback: include description of previous state in prompt
+            enhanced_prompt = (
+                f"{prompt}, maintaining visual consistency with the previous frame, "
+                f"same camera angle, same lighting, same environment, subtle changes only"
+            )
+            return self._generate_api_t2i(
+                prompt=enhanced_prompt, negative_prompt=negative_prompt,
+                seed=seed, width=prev_image.width, height=prev_image.height,
+                **kwargs,
+            )
 
     def generate_all_panels(
         self,
@@ -214,11 +204,7 @@ class ImageGenerator:
         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)
-        Panel 2+: img2img from previous panel
-        """
+        """Generate all panels in sequence — local or API."""
         images = []
 
         for i, prompt_data in enumerate(prompts):
@@ -229,37 +215,42 @@ class ImageGenerator:
             seed = base_seed + i
 
             try:
-                if i == 0:
-                    if reference_image is not None:
-                        img = self.generate_next_panel(
-                            prompt=full_prompt,
-                            previous_image=reference_image,
-                            negative_prompt=neg_prompt,
-                            strength=max(strength, 0.5),
-                            seed=seed,
-                            steps=steps,
-                            guidance=guidance,
-                        )
+                if self.mode == "local":
+                    if i == 0:
+                        if reference_image is not None:
+                            img = self._generate_local_i2i(
+                                full_prompt, reference_image,
+                                max(strength, 0.5), neg_prompt, seed, steps, guidance,
+                            )
+                        else:
+                            img = self._generate_local_t2i(
+                                full_prompt, neg_prompt, seed,
+                                width, height, steps, guidance,
+                            )
                     else:
-                        img = self.generate_first_panel(
-                            prompt=full_prompt,
-                            negative_prompt=neg_prompt,
-                            seed=seed,
-                            width=width,
-                            height=height,
-                            steps=steps,
-                            guidance=guidance,
+                        img = self._generate_local_i2i(
+                            full_prompt, images[-1], strength,
+                            neg_prompt, seed, steps, 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,
-                    )
+                    # API mode
+                    if i == 0:
+                        if reference_image is not None and self.provider.supports_img2img:
+                            img = self._generate_api_i2i(
+                                full_prompt, reference_image,
+                                max(strength, 0.5), neg_prompt, seed,
+                            )
+                        else:
+                            w = width or 1024
+                            h = height or 1024
+                            img = self._generate_api_t2i(
+                                full_prompt, neg_prompt, seed, w, h,
+                            )
+                    else:
+                        img = self._generate_api_i2i(
+                            full_prompt, images[-1], strength,
+                            neg_prompt, seed,
+                        )
 
                 images.append(img)
 
@@ -268,8 +259,8 @@ class ImageGenerator:
                 if images:
                     images.append(images[-1].copy())
                 else:
-                    fallback = Image.new("RGB", self.config["default_resolution"], (50, 50, 50))
-                    images.append(fallback)
+                    res = self.config.get("default_resolution", (1024, 1024))
+                    images.append(Image.new("RGB", res, (50, 50, 50)))
 
             if progress_callback:
                 progress_callback(i + 1, len(prompts))
@@ -278,14 +269,8 @@ class ImageGenerator:
         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,
+        self, panel_index, prompts, existing_images,
+        strength=0.4, base_seed=42, steps=None, guidance=None,
     ) -> Tuple[Image.Image, List[Image.Image]]:
         """Regenerate a single panel."""
         prompt_data = prompts[panel_index]
@@ -296,19 +281,23 @@ class ImageGenerator:
         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,
-            )
+            if self.mode == "local":
+                new_img = self._generate_local_t2i(
+                    full_prompt, neg_prompt, seed, steps=steps, guidance=guidance,
+                )
+            else:
+                new_img = self._generate_api_t2i(full_prompt, neg_prompt, seed)
         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,
-            )
+            prev = existing_images[panel_index - 1]
+            if self.mode == "local":
+                new_img = self._generate_local_i2i(
+                    full_prompt, prev, strength, neg_prompt, seed, steps, guidance,
+                )
+            else:
+                new_img = self._generate_api_i2i(
+                    full_prompt, prev, strength, neg_prompt, seed,
+                )
 
-        updated_images = existing_images.copy()
-        updated_images[panel_index] = new_img
-        return new_img, updated_images
+        updated = existing_images.copy()
+        updated[panel_index] = new_img
+        return new_img, updated