prompt_generator.py

#!/usr/bin/env python3
"""Generate image prompts using Gemini image understanding."""

import argparse
import asyncio
import base64
from typing import Optional, Tuple

import aiohttp
import boto3

import config
import storage

GEMINI_MODEL = "gemini-3-pro-preview"
REPLICATE_MODEL = "google/gemini-3-pro"

# Global flag for using Replicate
USE_REPLICATE = False


async def generate_prompt_for_image_replicate_async(
    session: aiohttp.ClientSession,
    image_bytes: bytes,
    category: str = None,
    caption: str = None,
) -> str:
    """Use Gemini via Replicate to generate a prompt that describes the image (async)."""
    b64_data = base64.b64encode(image_bytes).decode("utf-8")
    data_uri = f"data:image/jpeg;base64,{b64_data}"

    system_prompt = """You are an expert at describing images in precise detail for image generation AI models.
Analyze this photograph and write a detailed prompt that could be used to recreate this exact image.

Focus on:
- Main subject and its characteristics
- Setting/environment
- Lighting conditions
- Camera angle and perspective
- Colors, textures, and materials
- Any distinctive features or details

Write a single detailed paragraph that captures all visual elements. Be specific and precise. Use 1024 characters or less. Respond only with the description, no other text."""

    if caption:
        system_prompt += f'\n\nFor reference, the original caption for this image is: "{caption}".\nUse this context to better understand the scene, but focus on describing what you see visually.'

    headers = {
        "Authorization": f"Bearer {config.REPLICATE_API_TOKEN}",
        "Content-Type": "application/json",
        "Prefer": "wait",
    }

    payload = {
        "input": {
            "prompt": system_prompt,
            "images": [data_uri],
            "temperature": 0.7,
            "max_output_tokens": 2048,
        }
    }

    api_base = "https://api.replicate.com/v1"

    async with session.post(
        f"{api_base}/models/{REPLICATE_MODEL}/predictions",
        json=payload,
        headers=headers,
        timeout=aiohttp.ClientTimeout(total=120),
    ) as response:
        if response.status not in (200, 201):
            text = await response.text()
            raise RuntimeError(f"Replicate API error {response.status}: {text}")
        result = await response.json()

    # Poll if not completed
    if result.get("status") not in ("succeeded", "failed", "canceled"):
        poll_url = result.get("urls", {}).get("get") or f"{api_base}/predictions/{result['id']}"
        for _ in range(60):
            await asyncio.sleep(1)
            async with session.get(poll_url, headers=headers) as poll_resp:
                result = await poll_resp.json()
                if result.get("status") in ("succeeded", "failed", "canceled"):
                    break

    if result.get("status") != "succeeded":
        raise RuntimeError(f"Replicate prediction failed: {result}")

    output = result.get("output")
    if isinstance(output, list):
        return "".join(output).strip()
    return str(output).strip() if output else ""


async def generate_prompt_for_image_async(
    session: aiohttp.ClientSession,
    image_bytes: bytes,
    category: str = None,
    caption: str = None,
) -> str:
    """Use Gemini to generate a detailed text prompt that describes the image (async).

    Args:
        session: aiohttp session
        image_bytes: Raw image bytes
        category: Optional category hint for better prompt generation
        caption: Optional original caption/description from the source

    Returns:
        Generated text prompt
    """
    # Encode image as base64
    b64_data = base64.b64encode(image_bytes).decode("utf-8")

    # Build the prompt based on category
    system_prompt = """You are an expert at describing images in precise detail for image generation AI models.
Analyze this photograph and write a detailed prompt that could be used to recreate this exact image.

Focus on:
- Main subject and its characteristics
- Setting/environment
- Lighting conditions
- Camera angle and perspective
- Colors, textures, and materials
- Any distinctive features or details

Write a single detailed paragraph that captures all visual elements. Be specific and precise. Use 1024 characters or less. Respond only with the description, no other text."""

    # If we have a caption, include it as context
    if caption:
        system_prompt += f'\n\nFor reference, the original caption for this image is: "{caption}".\nUse this context to better understand the scene, but focus on describing what you see visually.'

    # Call Gemini API
    url = f"https://generativelanguage.googleapis.com/v1beta/models/{GEMINI_MODEL}:generateContent"
    payload = {
        "contents": [
            {
                "parts": [
                    {"text": system_prompt},
                    {
                        "inline_data": {
                            "mime_type": "image/jpeg",
                            "data": b64_data,
                        }
                    },
                ]
            }
        ],
        "generationConfig": {
            "temperature": 0.7,
            "maxOutputTokens": 2048,
        },
    }

    async with session.post(
        url,
        json=payload,
        params={"key": config.GEMINI_API_KEY},
        timeout=aiohttp.ClientTimeout(total=60),
    ) as response:
        if response.status != 200:
            text = await response.text()
            raise RuntimeError(f"Gemini API error {response.status}: {text}")

        result = await response.json()

    # Extract text from response
    try:
        text = result["candidates"][0]["content"]["parts"][0]["text"]
        return text.strip()
    except (KeyError, IndexError) as e:
        raise RuntimeError(f"Unexpected Gemini response: {result}") from e


async def process_single_async(
    session: aiohttp.ClientSession,
    s3_client,
    img: dict,
    dry_run: bool = False,
    use_replicate: bool = False,
) -> Tuple[str, str, Optional[str]]:
    """Process a single image and generate a prompt (async).

    Returns:
        Tuple of (image_id, status, prompt or error message)
    """
    image_id = img["id"]

    try:
        # Download from S3
        s3_path = img["s3_processed_path"]
        if s3_path.startswith("s3://"):
            parts = s3_path.replace("s3://", "").split("/", 1)
            bucket = parts[0]
            key = parts[1]
        else:
            bucket = config.S3_BUCKET
            key = s3_path.lstrip("/")
            if key.startswith("data/"):
                key = key[5:]

        response = s3_client.get_object(Bucket=bucket, Key=key)
        image_bytes = response["Body"].read()

        # Generate prompt (include caption if available)
        caption = img.get("caption")
        if use_replicate or USE_REPLICATE:
            prompt = await generate_prompt_for_image_replicate_async(
                session, image_bytes, category=img["category"], caption=caption
            )
        else:
            prompt = await generate_prompt_for_image_async(
                session, image_bytes, category=img["category"], caption=caption
            )

        if not dry_run:
            storage.insert_image_prompt(image_id, prompt, model=GEMINI_MODEL)

        return (image_id, "SUCCESS", prompt[:100] + "...")

    except Exception as e:
        return (image_id, "FAILED", str(e))


async def process_batch_async(
    category: str = None,
    limit: int = None,
    dry_run: bool = False,
    concurrency: int = 5,
    group: int = None,
    use_replicate: bool = False,
):
    """Process all validated images without prompts (async with concurrency).

    Args:
        category: Filter by category
        limit: Max images to process
        dry_run: If True, don't save to database
        concurrency: Number of parallel requests
        group: Filter by group (0, 1, 2, etc.)
        use_replicate: If True, use Replicate API instead of direct Gemini
    """
    s3_client = boto3.client("s3", region_name=config.S3_REGION)

    # Get images without prompts
    images = storage.get_images_without_prompt(
        model=GEMINI_MODEL, category=category, group=group
    )

    if limit:
        images = images[:limit]

    backend = "Replicate" if use_replicate else "Gemini"
    print(f"Processing {len(images)} images via {backend} (concurrency={concurrency})")
    print("=" * 60)

    success_count = 0
    failed_count = 0

    async with aiohttp.ClientSession() as session:
        # Process in batches
        for i in range(0, len(images), concurrency):
            batch = images[i : i + concurrency]
            batch_num = i // concurrency + 1
            total_batches = (len(images) + concurrency - 1) // concurrency

            print(f"\nBatch {batch_num}/{total_batches} ({len(batch)} images)")

            # Create tasks for this batch
            tasks = [
                process_single_async(session, s3_client, img, dry_run, use_replicate)
                for img in batch
            ]

            # Run batch concurrently
            results = await asyncio.gather(*tasks)

            # Process results
            for image_id, status, message in results:
                if status == "SUCCESS":
                    success_count += 1
                    print(f"  [{image_id}] {status}: {message}")
                else:
                    failed_count += 1
                    print(f"  [{image_id}] {status}: {message}")

    print(f"\n{'=' * 60}")
    print(f"Done! Success: {success_count}, Failed: {failed_count}")


def process_batch(
    category: str = None,
    limit: int = None,
    dry_run: bool = False,
    concurrency: int = 5,
    group: int = None,
    use_replicate: bool = False,
):
    """Process all validated images without prompts."""
    asyncio.run(
        process_batch_async(
            category=category,
            limit=limit,
            dry_run=dry_run,
            concurrency=concurrency,
            group=group,
            use_replicate=use_replicate,
        )
    )


def process_single(image_id: str, dry_run: bool = False, use_replicate: bool = False) -> str:
    """Process a single image and generate a prompt.

    Args:
        image_id: Image ID from database
        dry_run: If True, don't save to database
        use_replicate: If True, use Replicate API instead of direct Gemini

    Returns:
        Generated prompt
    """
    # Get image from DB
    img = storage.get_image_by_id(image_id)
    if not img:
        raise ValueError(f"Image not found: {image_id}")

    s3_client = boto3.client("s3", region_name=config.S3_REGION)

    async def run():
        async with aiohttp.ClientSession() as session:
            return await process_single_async(session, s3_client, img, dry_run, use_replicate)

    image_id, status, message = asyncio.run(run())

    if status == "SUCCESS":
        print(f"Generated prompt: {message}")
    else:
        print(f"Failed: {message}")

    return message


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate image prompts with Gemini")
    subparsers = parser.add_subparsers(dest="command")

    # Single image
    single_parser = subparsers.add_parser("single", help="Process single image")
    single_parser.add_argument("image_id", help="Image ID")
    single_parser.add_argument("--dry-run", "-n", action="store_true")
    single_parser.add_argument("--replicate", "-r", action="store_true", help="Use Replicate API instead of direct Gemini")

    # Batch processing
    batch_parser = subparsers.add_parser("batch", help="Process batch of images")
    batch_parser.add_argument("--category", "-c", help="Filter by category")
    batch_parser.add_argument("--limit", "-l", type=int, help="Max images")
    batch_parser.add_argument("--dry-run", "-n", action="store_true")
    batch_parser.add_argument(
        "--concurrency",
        "-p",
        type=int,
        default=5,
        help="Number of parallel requests (default: 5)",
    )
    batch_parser.add_argument(
        "--group", "-g", type=int, help="Only process images in this group (0, 1, 2, etc.)"
    )
    batch_parser.add_argument("--replicate", "-r", action="store_true", help="Use Replicate API instead of direct Gemini")

    # List images with prompts
    list_parser = subparsers.add_parser("list", help="List images with prompts")
    list_parser.add_argument("--category", "-c", help="Filter by category")
    list_parser.add_argument("--limit", "-l", type=int, default=10)

    args = parser.parse_args()

    if args.command == "single":
        process_single(args.image_id, dry_run=args.dry_run, use_replicate=args.replicate)
    elif args.command == "batch":
        process_batch(
            category=args.category,
            limit=args.limit,
            dry_run=args.dry_run,
            concurrency=args.concurrency,
            group=args.group,
            use_replicate=args.replicate,
        )
    elif args.command == "list":
        images = storage.get_images_with_prompt(
            model=GEMINI_MODEL, category=args.category
        )
        for img in images[: args.limit]:
            print(f"\n{img['id']} ({img['category']}):")
            print(f"  {img['prompt'][:200]}...")
    else:
        parser.print_help()