src/amuseme/bg_generator.py

"""
bg_generator.py — Stable Diffusion-based background generator
Generates thematic images for lyric storyboard using stabilityai/sd-turbo.
"""
import torch
from diffusers import AutoPipelineForText2Image

from amuseme.logger import get_logger

logger = get_logger("bg_generator")

# Determine device and dtype
# In HF ZeroGPU, torch.cuda.is_available() returns True at root level
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32

logger.info(f"Initializing SD-Turbo on {device}...")
try:
    pipe = AutoPipelineForText2Image.from_pretrained(
        "stabilityai/sd-turbo",
        torch_dtype=torch_dtype,
        variant="fp16" if device == "cuda" else None
    )
    pipe.to(device)
except Exception as e:
    logger.warning(f"Could not pre-load SD-Turbo ({e}).")
    pipe = None


# Try to import spaces for ZeroGPU compatibility
try:
    import spaces
    HAS_SPACES = True
except ImportError:
    HAS_SPACES = False


# Baseline quality/style guidance applied to every generated background,
# independent of the per-line user prompt. Keeps the storyboard visually
# consistent (cinematic, no on-image text) regardless of what the lyric
# content or user's visual prompt asks for.
SYSTEM_PROMPT = (
    "cinematic lyric video background art, atmospheric, high detail, "
    "rich color grading, no text, no words, no letters, no captions, no watermark"
)


def _build_full_prompt(user_prompt: str) -> str:
    return f"{SYSTEM_PROMPT}, {user_prompt}" if user_prompt else SYSTEM_PROMPT


def _generate_batch(prompts: list[str]) -> list:
    """
    Generate a list of PIL Images corresponding to the input prompts.
    Uses 1-step inference for SD-Turbo (extremely fast).
    """
    if pipe is None:
        logger.error("Stable Diffusion pipeline is not initialized.")
        return []

    full_prompts = [_build_full_prompt(p) for p in prompts]
    logger.info(f"SD-Turbo Input Prompts (Total: {len(full_prompts)}):\n{full_prompts}")
    images = []
    for i, p in enumerate(full_prompts):
        logger.info(f"Generating storyboard image {i+1}/{len(full_prompts)}: '{p}'")
        try:
            # 1 step, guidance_scale=0.0 is optimal for sd-turbo
            result = pipe(p, num_inference_steps=1, guidance_scale=0.0)
            img = result.images[0]
            logger.info(f"SD-Turbo successfully generated image {i+1} with size {img.size}")
            images.append(img)
        except Exception as ex:
            logger.error(f"Failed to generate for prompt '{p}': {ex}")
            # Fallback to a plain dark grey image
            from PIL import Image
            fallback_img = Image.new("RGB", (512, 512), (30, 30, 30))
            logger.info(f"Fell back to placeholder image for prompt {i+1} with size {fallback_img.size}")
            images.append(fallback_img)
    return images


if HAS_SPACES:
    generate_storyboard = spaces.GPU(duration=60)(_generate_batch)
else:
    generate_storyboard = _generate_batch