""" WHAT: Image Generator — Stage 2 of the Multiverse AI pipeline. WHY: Converts the text prompt (from Stage 1) into a high-quality base visual scene. HOW: Reads its backend from profiles.py (get_stage_config). - "gemini" → Uses Gemini's native image generation (free quota, no HF credits) - "hf_inference" → Uses HuggingFace Inference Provider (FLUX.1-schnell, requires credits) - "local" → Loads a diffusion model locally via `diffusers` (requires GPU) - "mock" → Returns a procedurally generated PIL Image instantly """ import gc from io import BytesIO from typing import Dict, Union import torch from PIL import Image from .base import BaseModel from ..config import HF_TOKEN, GEMINI_API_KEY, DEVICE, MOCK_INFERENCE, get_stage_config class ImageGenerator(BaseModel): """ Wrapper for the Text-to-Image model. Routes to Gemini, HF Inference, local diffusers, or mock based on the active profile. """ def __init__(self): self.client = None self.pipeline = None # Used only by the "local" backend (diffusers) self.stage_config = get_stage_config("image_generation") self.backend = self.stage_config["backend"] self.model_id = self.stage_config["model"] def initialize(self) -> None: """ WHAT: Initializes the appropriate image generation client based on the active profile. WHY: Gemini, HF, and local diffusers all need different initialization steps. HOW: Reads self.backend and sets up the matching client/pipeline object. """ if MOCK_INFERENCE or self.backend == "mock": print("[ImageGenerator] Running in MOCK mode. Bypassing model download & load.") return # ── GEMINI backend ──────────────────────────────────────────────────── if self.backend == "gemini": if not GEMINI_API_KEY: # No key → degrade to free Pollinations image generation so the demo still works. print("[ImageGenerator] GEMINI_API_KEY missing — falling back to Pollinations (no key) for images.") self.backend = "pollinations" return from google import genai self.client = genai.Client(api_key=GEMINI_API_KEY) print(f"[ImageGenerator] Initialized Gemini client with model: {self.model_id}") return # ── HF_INFERENCE backend ────────────────────────────────────────────── if self.backend == "hf_inference": if not HF_TOKEN: raise ValueError("HF_TOKEN is missing. Cannot initialize HF ImageGenerator.") from huggingface_hub import InferenceClient # provider="hf-inference" routes through HF's own serverless tier self.client = InferenceClient(token=HF_TOKEN, provider="hf-inference") print(f"[ImageGenerator] Initialized HF InferenceClient with model: {self.model_id}") return # ── LOCAL backend ───────────────────────────────────────────────────── if self.backend == "local": # Import diffusers only when actually needed (avoids slow import on mock/cloud runs) from diffusers import StableDiffusionPipeline print(f"[ImageGenerator] Loading local diffusion model: {self.model_id}...") # torch_dtype=float16 halves VRAM usage — essential for 6GB GPUs self.pipeline = StableDiffusionPipeline.from_pretrained( self.model_id, torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32, ).to(DEVICE) print("[ImageGenerator] Local diffusion pipeline ready.") return def _mock_image(self, actual_prompt: str) -> Image.Image: """Procedurally generate a placeholder scene image (used by mock backend and as a fallback).""" from PIL import ImageDraw img = Image.new("RGB", (704, 512), color=(79, 70, 229)) draw = ImageDraw.Draw(img) draw.ellipse([277, 181, 427, 331], fill=(225, 29, 72)) label = actual_prompt[:40] if actual_prompt else "Mock Scene" draw.text((20, 20), f"Scene: {label}", fill=(255, 255, 255)) return img def generate(self, **kwargs) -> Image.Image: """ WHAT: Generates a PIL Image from the expanded prompt. WHY: Returns a real or mock image depending on the active backend. HOW: Routes to the correct backend generation method. """ prompt_data: Union[Dict[str, str], str] = kwargs.get("prompt", "") # Extract the image-specific sub-prompt if a full dict was passed in if isinstance(prompt_data, dict): actual_prompt = prompt_data.get("image_prompt", "") else: actual_prompt = str(prompt_data) # ── MOCK backend ────────────────────────────────────────────────────── if MOCK_INFERENCE or self.backend == "mock": return self._mock_image(actual_prompt) if not actual_prompt: raise ValueError("ImageGenerator requires a valid 'prompt' in kwargs.") # ── POLLINATIONS backend ────────────────────────────────────────────── # WHAT: Calls Pollinations.ai — a completely free, no-auth image generation service. # WHY: Requires zero API key, zero credits, zero quota. Just a plain HTTP GET request. # The URL encodes the prompt and the service returns a JPEG image directly. # HOW: We URL-encode the prompt string and make a GET request. The response body # IS the image bytes — we wrap them directly in a PIL Image object. if self.backend == "pollinations": try: import urllib.parse import requests as req encoded = urllib.parse.quote(actual_prompt) url = f"https://image.pollinations.ai/prompt/{encoded}?width=768&height=512&nologo=true&seed=42" print(f"[ImageGenerator] Calling Pollinations.ai for free image generation...") response = req.get(url, timeout=60) response.raise_for_status() return Image.open(BytesIO(response.content)).convert("RGB") except Exception as e: print(f"[ImageGenerator Warning] Pollinations failed ({e}). Using local mock image fallback.") return self._mock_image(actual_prompt) # ── GEMINI backend ──────────────────────────────────────────────────── if self.backend == "gemini": try: print(f"[ImageGenerator] Calling Gemini image generation model: {self.model_id}...") from google.genai import types response = self.client.models.generate_content( model=self.model_id, contents=actual_prompt, config=types.GenerateContentConfig( # Request both text AND image in the response response_modalities=["TEXT", "IMAGE"], ), ) # Gemini returns parts — find the first IMAGE part for part in response.candidates[0].content.parts: if part.inline_data is not None: # inline_data.data is raw image bytes — wrap in PIL return Image.open(BytesIO(part.inline_data.data)).convert("RGB") raise RuntimeError("Gemini returned no image in its response parts.") except Exception as e: raise RuntimeError(f"Gemini Image Generation failed: {e}") from e # ── HF_INFERENCE backend ────────────────────────────────────────────── if self.backend == "hf_inference": try: print(f"[ImageGenerator] Calling HF Inference API for model: {self.model_id}...") pil_image = self.client.text_to_image(actual_prompt, model=self.model_id) return pil_image except Exception as e: raise RuntimeError(f"HF Cloud Image Generation failed: {e}") from e # ── LOCAL backend ───────────────────────────────────────────────────── if self.backend == "local": try: print(f"[ImageGenerator] Running local diffusion inference on {DEVICE}...") with torch.inference_mode(): result = self.pipeline(actual_prompt, num_inference_steps=25) return result.images[0] except Exception as e: raise RuntimeError(f"Local Image Generation failed: {e}") from e raise RuntimeError(f"Unknown backend '{self.backend}' for ImageGenerator.") def cleanup(self) -> None: """ WHAT: Releases GPU VRAM after inference completes. WHY: The local diffusion pipeline holds ~4GB of VRAM. Releasing it after Stage 2 allows Stage 4 (MusicGen) and Stage 5 (Video) to load without OOM errors. HOW: Deletes the pipeline reference and calls gc.collect() + torch.cuda.empty_cache(). """ if self.pipeline is not None: del self.pipeline self.pipeline = None gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() print("[ImageGenerator] Released local pipeline from VRAM.")