imagegen.py

"""Illustration generation: Z-Image-Turbo (6B, Apache 2.0) via diffusers.

Character consistency strategy (validated in the local prototype):
fixed style prefix + per-character fixed ENGLISH appearance description +
one fixed seed per story session.
"""

from __future__ import annotations

import hashlib
import os

from PIL import Image

MOCK = os.environ.get("STORY_MOCK") == "1"

IMAGE_MODEL_ID = os.environ.get("STORY_IMAGE_MODEL", "Tongyi-MAI/Z-Image-Turbo")

STYLE_PREFIX = (
    "children's picture book illustration, soft watercolor style, gentle pastel colors, "
    "warm cozy bedtime mood, cute rounded shapes, storybook art, friendly and safe, no text"
)

WIDTH, HEIGHT = 1024, 768
STEPS = 8
GUIDANCE = 1.0  # Z-Image-Turbo is distilled: 8 steps, no CFG

_pipe = {"pipe": None}


def build_illustration_prompt(cast: list, illustration: str) -> str:
    """style prefix + appearances of characters present in the scene (fallback: whole cast)."""
    present = [c for c in cast if c.name and c.name in illustration]
    if not present:
        present = list(cast)
    cast_desc = ", ".join(f"{c.name} ({c.appearance.strip() or f'a cute {c.kind}'})" for c in present)
    return ", ".join(part for part in (STYLE_PREFIX, cast_desc, illustration.strip()) if part)


def preload():
    """Load at Space startup, not inside @spaces.GPU functions (see story.preload)."""
    if MOCK or _pipe["pipe"] is not None:
        return
    import torch
    from diffusers import ZImagePipeline

    _pipe["pipe"] = ZImagePipeline.from_pretrained(IMAGE_MODEL_ID, torch_dtype=torch.bfloat16).to("cuda")


def _load_pipe():
    preload()
    return _pipe["pipe"]


def _mock_image(prompt: str) -> Image.Image:
    h = int(hashlib.sha1(prompt.encode()).hexdigest()[:6], 16)
    base = ((h >> 16) & 255, (h >> 8) & 255, h & 255)
    img = Image.new("RGB", (WIDTH, HEIGHT))
    for y in range(HEIGHT):
        t = y / HEIGHT
        row = tuple(int(c * (1 - t) + 240 * t) for c in base)
        for x in range(0, WIDTH, 4):
            img.paste(Image.new("RGB", (4, 1), row), (x, y))
    return img


def generate_illustration(prompt: str, seed: int) -> Image.Image:
    if MOCK:
        return _mock_image(prompt)
    import torch

    pipe = _load_pipe()
    generator = torch.Generator(device="cuda").manual_seed(seed)
    result = pipe(
        prompt=prompt,
        num_inference_steps=STEPS,
        guidance_scale=GUIDANCE,
        width=WIDTH,
        height=HEIGHT,
        generator=generator,
    )
    return result.images[0]