compose.py

"""Compose reference photos into generated scenes using Nano Banana 2 Edit."""

import os
import time

import requests
from dotenv import load_dotenv

load_dotenv()

RUNPOD_API_KEY = os.environ.get("RUNPOD_API_KEY", "")
NANO_BANANA_URL = "https://api.runpod.ai/v2/google-nano-banana-2-edit"

# Nano Banana 2 recommends 1-3 reference images for best stability.
MAX_REFERENCE_IMAGES = 3


def compose_image(
    scene_image_url: str,
    reference_urls: list[str],
    prompt: str,
    resolution: str = "1k",
    timeout: int = 120,
) -> dict | None:
    """Compose reference photos into a generated scene using Nano Banana 2.

    Args:
        scene_image_url: URL of the generated scene image (e.g. SDXL output).
        reference_urls: URLs of reference photos (your photos).
        prompt: Composition instruction prompt.
        resolution: Output resolution ('1k', '2k', '4k').
        timeout: Max seconds to wait for completion.

    Returns:
        Dict with 'image_url' and 'cost', or None on failure.
    """
    if not RUNPOD_API_KEY:
        print("[compose] No API key, skipping composition")
        return None

    # Nano Banana 2 takes all images in one array.
    # Put scene image first, then reference photos.
    # Limit reference images to MAX_REFERENCE_IMAGES for stability.
    images = [scene_image_url] + reference_urls[:MAX_REFERENCE_IMAGES]

    print(f"[compose] Sending {len(images)} images to Nano Banana 2 ({resolution})...")
    print(f"[compose] Prompt: {prompt}...")

    try:
        # Use async /run + poll since composition can take 30-90s
        r = requests.post(
            f"{NANO_BANANA_URL}/run",
            headers={
                "Authorization": f"Bearer {RUNPOD_API_KEY}",
                "Content-Type": "application/json",
            },
            json={
                "input": {
                    "images": images,
                    "prompt": prompt,
                    "resolution": resolution,
                    "enable_safety_checker": False,
                }
            },
            timeout=30,
        )
        r.raise_for_status()
        job = r.json()
        job_id = job["id"]
        print(f"[compose] Job submitted: {job_id}")

        # Poll for completion — only log on state changes
        start = time.time()
        last_status = None
        while time.time() - start < timeout:
            sr = requests.get(
                f"{NANO_BANANA_URL}/status/{job_id}",
                headers={"Authorization": f"Bearer {RUNPOD_API_KEY}"},
                timeout=15,
            )
            status = sr.json()
            current_status = status["status"]

            if current_status != last_status:
                elapsed_so_far = time.time() - start
                print(
                    f"[compose] {current_status} ({elapsed_so_far:.0f}s)",
                    flush=True,
                )
                last_status = current_status

            if current_status == "COMPLETED":
                output = status["output"]
                image_url = output.get("result") or output.get("image_url")
                cost = output.get("cost", 0)
                elapsed = time.time() - start
                print(f"[compose] Done in {elapsed:.1f}s (cost: ${cost:.4f})")
                return {"image_url": image_url, "cost": cost}

            elif current_status == "FAILED":
                error = status.get("error", "Unknown error")
                print(f"[compose] FAILED: {error}")
                return None

            else:
                time.sleep(3)

        print(f"[compose] Timed out after {timeout}s")
        return None

    except Exception as e:
        print(f"[compose] Error: {e}")
        return None


def build_compose_prompt() -> str:
    """Build a composition prompt for Nano Banana 2."""
    return (
        "Place the subjects from the reference photos into the scene. "
        "Match the exact art style, color palette, and rendering technique of the input scene image. "
        "Maintain the subjects' identity and likeness while blending them seamlessly into the scene's style."
    )