Blame for image_generator.py

"""
TimeLapseForge — Image Generator Module v2.0
Supports BOTH local (diffusers) and API-based generation.
Unified interface for all providers.
"""

import gc
import torch
from typing import List, Dict, Optional, Callable, Tuple
from PIL import Image

from api_providers import get_provider, BaseProvider, PROVIDERS


def flush_memory():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


# ═══════════════════════════════════════════
# LOCAL GENERATOR (Same as before — diffusers)
# ═══════════════════════════════════════════

_local_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 _local_pipelines:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.float16 if device == "cuda" else torch.float32

        if pipeline_type == "t2i":
            try:
                pipe = AutoPipelineForText2Image.from_pretrained(
                    model_id, torch_dtype=dtype, variant="fp16"
                )
            except Exception:
                pipe = AutoPipelineForText2Image.from_pretrained(
                    model_id, torch_dtype=dtype
                )
            pipe = pipe.to(device)
            _local_pipelines[cache_key] = pipe

        elif pipeline_type == "i2i":
            t2i_key = f"{model_id}_t2i"
            if t2i_key in _local_pipelines:
                try:
                    pipe = AutoPipelineForImage2Image.from_pipe(_local_pipelines[t2i_key])
                except Exception:
                    try:
                        pipe = AutoPipelineForImage2Image.from_pretrained(
                            model_id, torch_dtype=dtype, variant="fp16"
                        )
                    except Exception:
                        pipe = AutoPipelineForImage2Image.from_pretrained(
                            model_id, torch_dtype=dtype
                        )
                    pipe = pipe.to(device)
            else:
                try:
                    pipe = AutoPipelineForImage2Image.from_pretrained(
                        model_id, torch_dtype=dtype, variant="fp16"
                    )
                except Exception:
                    pipe = AutoPipelineForImage2Image.from_pretrained(
                        model_id, torch_dtype=dtype
                    )
                pipe = pipe.to(device)
            _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),
    },
}


# ═══════════════════════════════════════════
# UNIFIED IMAGE GENERATOR
# ═══════════════════════════════════════════

class ImageGenerator:
    """
    Unified image generator supporting both local models and API providers.
    """

    def __init__(
        self,
        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.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
        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 = 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,
        )

    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,
        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 — local or API."""
        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 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_local_i2i(
                            full_prompt, images[-1], strength,
                            neg_prompt, seed, steps, guidance,
                        )
                else:
                    # 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)

            except Exception as e:
                print(f"Error generating panel {i + 1}: {e}")
                if images:
                    images.append(images[-1].copy())
                else:
                    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))

        flush_memory()
        return images

    def regenerate_single_panel(
        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]
        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:
            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:
            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 = existing_images.copy()
        updated[panel_index] = new_img
        return new_img, updated