visualnovel/painter.py

"""The Painter — backdrops + character sprites.

Design rules (see docs/ARCHITECTURE.md §5):
  - Compose prompts IN CODE from trusted fields so the locked anime style never drifts.
  - Cache every render by (kind, prompt, seed). A character's sprite is generated ONCE per
    mood and reused — never re-paint to "refresh" (that's what breaks consistency).
  - Pin seeds (character.sprite_seed / scene.backdrop_seed) for reproducibility.

`PainterBase` owns prompts + caching; subclasses implement `_render(prompt, seed, size)`.
MockPainter draws a labelled placeholder (zero ML deps) so the loop is visible immediately.
"""

from __future__ import annotations

import hashlib
from pathlib import Path
from typing import Protocol

from . import config
from .metrics import collector
from .schemas import Character, GameState
from .utils import _quiet_stderr


class Painter(Protocol):
    def backdrop(self, state: GameState) -> Path: ...
    def sprite(self, state: GameState, ch: Character) -> Path: ...


# --------------------------------------------------------------------------- #
#  Prompt composition (shared)
# --------------------------------------------------------------------------- #

# CLIP hard limit is 77 tokens (~4 chars/token).  Keep variable parts short so the
# fixed prefix + suffix stay within budget.  Rough budgets per prompt:
#   backdrop : 15 (prefix) + 18 (desc) + 5 (mood) + 12 (style) + 2 (suffix) ≈ 52
#   sprite   : 12 (prefix) + 10 (style) + 20 (appearance) + 5 (mood+suffix)  ≈ 47
_DESC_WORDS = 18  # scene description budget
_STYLE_WORDS = 12  # style_guide budget in prompts
_APP_WORDS = 20  # character appearance budget


def _w(text: str, n: int) -> str:
    """Return at most *n* words from *text* (space-split, no tokenizer needed)."""
    words = text.split()
    return " ".join(words[:n]) if len(words) > n else text


def _backdrop_style(style_guide: str) -> str:
    """Keep only the most impactful style tokens and strip character-specific ones."""
    skip = {"expressive eyes", "clean linework"}
    tokens = [t.strip() for t in style_guide.split(",") if t.strip() not in skip]
    return _w(", ".join(tokens), _STYLE_WORDS)


def backdrop_prompt(state: GameState) -> str:
    style = _backdrop_style(state.style_guide)
    desc = _w(state.scene.description, _DESC_WORDS)
    return (
        f"background art, empty scenic environment, wide establishing shot, no humans, "
        f"uninhabited location, {desc}, {state.scene.mood} atmosphere, "
        f"{style}, no text"
    )


def sprite_prompt(state: GameState, ch: Character) -> str:
    style = _w(state.style_guide, _STYLE_WORDS)
    app = _w(ch.appearance, _APP_WORDS)
    return (
        f"vtuber character design, solo, single character, full-body, centered, "
        f"{style}, {app}, {ch.mood} expression, "
        f"pure white background, simple background, no scenery, no text"
    )


# Negative prompt for backdrops: ban any human/character presence.
_BACKDROP_NEGATIVE = (
    f"{config.NEGATIVE_PROMPT}, "
    "person, people, human, character, figure, man, woman, boy, girl, face, body, "
    "portrait, anime character, silhouette, crowd, group"
)

# Negative prompt injected for every sprite render to enforce the white background.
# Combined with config.NEGATIVE_PROMPT (quality/artifact negatives).
_SPRITE_NEGATIVE = (
    f"{config.NEGATIVE_PROMPT}, "
    "background, scenery, landscape, environment, outdoors, indoors, room, sky, "
    "trees, grass, colored background, gradient background, detailed background, "
    "complex background, nature, architecture, buildings"
)


class PainterBase:
    """Caching + key derivation. Subclasses override `_render`."""

    _rembg_session = None

    def _remove_bg(self, img):
        """rembg with a reused ONNX session — rembg.remove() without one reloads
        the ~170MB u2net model on every call. Lazy so mock mode never imports it."""
        from rembg import new_session, remove  # noqa: PLC0415

        if self._rembg_session is None:
            self._rembg_session = new_session()
        return remove(img, session=self._rembg_session)

    def backdrop(self, state: GameState) -> Path:
        prompt = backdrop_prompt(state)
        return self._cached(
            kind="bg",
            key=state.scene.id,
            prompt=prompt,
            seed=state.scene.backdrop_seed,
            negative_prompt=_BACKDROP_NEGATIVE,
            guidance_scale=config.BACKDROP_GUIDANCE,
        )

    def sprite(self, state: GameState, ch: Character) -> Path:
        prompt = sprite_prompt(state, ch)
        return self._cached(
            kind="sprite",
            key=f"{ch.id}.{ch.mood}",
            prompt=prompt,
            seed=ch.sprite_seed,
            negative_prompt=_SPRITE_NEGATIVE,
        )

    def ending_backdrop(self, state: GameState, kind: str) -> Path:
        """Generate (and cache) a special ending illustration based on the ending kind."""
        style = _backdrop_style(state.style_guide)
        _ENDING_DESCS: dict[str, tuple[str, str]] = {
            "warm": (
                "cherry blossom park at golden sunset, petals falling softly, "
                "empty bench under blooming trees, warm amber light filtering through branches",
                "romantic warm",
            ),
            "bittersweet": (
                "misty autumn street at twilight, fallen leaves on cobblestones, "
                "distant lamplight, empty path fading into soft fog",
                "melancholic gentle",
            ),
            "strange": (
                "surreal moonlit garden, glowing silver motes drifting upward, "
                "impossible geometry, dreamlike luminous plants",
                "mysterious ethereal",
            ),
            "defeat": (
                "rain-soaked empty park bench at night, fallen leaves in puddles, "
                "single dim lamppost, deserted street receding into darkness",
                "desolate cold",
            ),
        }
        desc, mood = _ENDING_DESCS.get(kind, _ENDING_DESCS["warm"])
        prompt = (
            f"background art, empty scenic environment, wide establishing shot, no humans, "
            f"uninhabited location, {desc}, {mood} atmosphere, {style}, no text"
        )
        seed = (state.seed * 1234567 + sum(ord(c) for c in kind)) % (2**31)
        return self._cached(
            kind="ending",
            key=kind,
            prompt=prompt,
            seed=seed,
            negative_prompt=_BACKDROP_NEGATIVE,
            guidance_scale=config.BACKDROP_GUIDANCE,
        )

    # -- internals --
    def _cached(
        self,
        *,
        kind: str,
        key: str,
        prompt: str,
        seed: int,
        negative_prompt: str = "",
        guidance_scale: float = 0.0,
        transform=None,
    ) -> Path:
        h = hashlib.sha1(
            f"{kind}|{prompt}|{negative_prompt}|{seed}|{guidance_scale}".encode()
        ).hexdigest()[:12]
        path = config.CACHE_DIR / f"{kind}_{key}_{h}.png"
        if path.exists():
            collector.record_cache(hit=True)
            return path
        collector.record_cache(hit=False)
        img = self._render(
            prompt,
            seed,
            config.IMAGE_SIZE,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
        )
        if transform is not None:
            img = transform(img)
        img.save(path)
        return path

    def _render(
        self,
        prompt: str,
        seed: int,
        size: int,
        negative_prompt: str = "",
        guidance_scale: float = 0.0,
    ):
        raise NotImplementedError


# --------------------------------------------------------------------------- #
#  Mock — labelled placeholder, no models
# --------------------------------------------------------------------------- #
class MockPainter(PainterBase):
    def _render(
        self,
        prompt: str,
        seed: int,
        size: int,
        negative_prompt: str = "",
        guidance_scale: float = 0.0,
    ):
        from PIL import Image, ImageDraw  # noqa: PLC0415

        # deterministic colour from the seed so the same entity keeps the same placeholder
        r, g, b = (seed % 200 + 30, (seed // 7) % 200 + 30, (seed // 13) % 200 + 30)
        img = Image.new("RGB", (size, size), (r, g, b))
        d = ImageDraw.Draw(img)
        label = prompt[:60] + ("…" if len(prompt) > 60 else "")
        d.rectangle([8, 8, size - 8, size - 8], outline=(255, 255, 255), width=2)
        d.text((20, 20), "MOCK PAINTER", fill=(255, 255, 255))
        d.text((20, 44), label, fill=(235, 235, 235))
        d.text((20, size - 28), f"seed={seed}", fill=(220, 220, 220))
        return img


def _vae_kwargs(dtype, device: str) -> dict:
    """Pipeline kwargs swapping in the fp16-safe VAE (see config.IMAGE_VAE).

    Silences transformers' own logger first: pipeline loads pull in CLIP components
    whose advisory messages (e.g. 'requires torchvision') bypass stdlib logging config.
    """
    from transformers.utils import logging as hf_logging  # noqa: PLC0415

    hf_logging.set_verbosity_error()
    if not config.IMAGE_VAE or device == "cpu":
        return {}  # CPU runs fp32 anyway — no upcast cost to avoid
    from diffusers import AutoencoderKL  # noqa: PLC0415

    return {"vae": AutoencoderKL.from_pretrained(config.IMAGE_VAE, torch_dtype=dtype)}


def _parse_lora(lora: str) -> tuple[str, str | None]:
    """Return (repo_or_path, weight_name_or_None) for load_lora_weights.

    Accepts:
      - HF URL:   https://huggingface.co/owner/repo/resolve/main/sub/file.safetensors
      - HF repo:  owner/repo  (diffusers auto-detects the single .safetensors)
      - local:    /abs/path/to/lora.safetensors
    """
    import urllib.parse

    if lora.startswith("https://huggingface.co/"):
        path = urllib.parse.unquote(lora.removeprefix("https://huggingface.co/"))
        parts = path.split("/")
        # parts: [owner, repo, "resolve", branch, *file_parts]
        repo_id = "/".join(parts[:2])
        weight_name = "/".join(parts[4:])
        return (repo_id, weight_name)
    if "::" in lora:  # "owner/repo::weight_file.safetensors"
        repo_id, weight_name = lora.split("::", 1)
        return (repo_id, weight_name)
    return (lora, None)


# --------------------------------------------------------------------------- #
#  SDXL-Turbo (+ your fine-tuned anime LoRA). STUB — implement with Claude Code.
# --------------------------------------------------------------------------- #
class SdxlTurboPainter(PainterBase):
    def __init__(self) -> None:
        import torch  # noqa: PLC0415
        from diffusers.pipelines.auto_pipeline import AutoPipelineForText2Image  # noqa: PLC0415

        device = config.detect_device()  # "cuda" also covers AMD ROCm; "mps" = Apple Metal
        dtype = torch.float16 if device in ("cuda", "mps") else torch.float32
        print(f"[painter] Loading {config.IMAGE_MODEL} (downloading on first run)…", flush=True)
        self.pipe = AutoPipelineForText2Image.from_pretrained(
            config.IMAGE_MODEL,
            torch_dtype=dtype,
            variant="fp16" if device == "cuda" else None,
            **_vae_kwargs(dtype, device),
        ).to(device)
        if config.IMAGE_LORA:
            repo, weight_name = _parse_lora(config.IMAGE_LORA)
            kw = {"weight_name": weight_name} if weight_name else {}
            with _quiet_stderr():
                self.pipe.load_lora_weights(repo, **kw)
        self.torch = torch
        self.device = device

    def sprite(self, state: GameState, ch: Character) -> Path:
        prompt = sprite_prompt(state, ch)
        return self._cached(
            kind="sprite",
            key=f"{ch.id}.{ch.mood}",
            prompt=prompt,
            seed=ch.sprite_seed,
            negative_prompt=_SPRITE_NEGATIVE,
            transform=self._remove_bg,  # removes background → transparent PNG
        )

    def _render(
        self,
        prompt: str,
        seed: int,
        size: int,
        negative_prompt: str = "",
        guidance_scale: float = 0.0,
    ):
        gen = self.torch.Generator(device=self.device).manual_seed(seed)
        result = self.pipe(
            prompt=prompt,
            negative_prompt=negative_prompt or None,
            num_inference_steps=config.IMAGE_STEPS,
            guidance_scale=guidance_scale,
            height=size,
            width=size,
            generator=gen,
        )
        return result.images[0]


# --------------------------------------------------------------------------- #
#  SDXL-Lightning — 4-step adversarial distillation. Switch: VN_IMAGE_BACKEND=lightning
# --------------------------------------------------------------------------- #
class SdxlLightningPainter(PainterBase):
    def __init__(self) -> None:
        import torch  # noqa: PLC0415
        from diffusers import EulerDiscreteScheduler, StableDiffusionXLPipeline  # noqa: PLC0415

        device = config.detect_device()
        dtype = torch.float16 if device in ("cuda", "mps") else torch.float32
        print(f"[painter] Loading SDXL-Lightning ({config.IMAGE_MODEL})…", flush=True)
        self.pipe = StableDiffusionXLPipeline.from_pretrained(
            config.IMAGE_MODEL,
            torch_dtype=dtype,
            variant="fp16" if device == "cuda" else None,
            **_vae_kwargs(dtype, device),
        ).to(device)
        # Lightning requires EulerDiscrete with trailing timestep spacing
        self.pipe.scheduler = EulerDiscreteScheduler.from_config(
            self.pipe.scheduler.config, timestep_spacing="trailing"
        )
        if config.IMAGE_LORA:
            repo, weight_name = _parse_lora(config.IMAGE_LORA)
            kw = {"weight_name": weight_name} if weight_name else {}
            self.pipe.load_lora_weights(repo, **kw)
            self.pipe.fuse_lora()  # bake into weights for faster inference
        self.torch = torch
        self.device = device

    def sprite(self, state: GameState, ch: Character) -> Path:
        prompt = sprite_prompt(state, ch)
        return self._cached(
            kind="sprite",
            key=f"{ch.id}.{ch.mood}",
            prompt=prompt,
            seed=ch.sprite_seed,
            negative_prompt=_SPRITE_NEGATIVE,
            transform=self._remove_bg,  # removes background → transparent PNG
        )

    def _render(
        self,
        prompt: str,
        seed: int,
        size: int,
        negative_prompt: str = "",
        guidance_scale: float = 0.0,
    ):
        gen = self.torch.Generator(device=self.device).manual_seed(seed)
        result = self.pipe(
            prompt=prompt,
            negative_prompt=negative_prompt or None,
            num_inference_steps=config.IMAGE_STEPS,
            guidance_scale=guidance_scale,
            height=size,
            width=size,
            generator=gen,
        )
        return result.images[0]


# --------------------------------------------------------------------------- #
#  Modal cloud painter  (VN_IMAGE_BACKEND=modal)
# --------------------------------------------------------------------------- #
class ModalPainter(PainterBase):
    """Delegates _render to a Modal A10G container — no local GPU required."""

    def __init__(self) -> None:
        import modal  # noqa: PLC0415

        self._backend = modal.Cls.from_name("vn-app", "ModalPainterBackend")()

    def sprite(self, state: GameState, ch: Character) -> Path:
        """Override to request server-side background removal via rembg."""
        import io  # noqa: PLC0415

        from PIL import Image  # noqa: PLC0415

        prompt = sprite_prompt(state, ch)
        neg = _SPRITE_NEGATIVE
        seed = ch.sprite_seed
        kind, key = "sprite", f"{ch.id}.{ch.mood}"

        # Include remove_bg in the cache key so backdrop/sprite hashes never collide.
        h = hashlib.sha1(f"{kind}|{prompt}|{neg}|{seed}|rembg".encode()).hexdigest()[:12]
        path = config.CACHE_DIR / f"{kind}_{key}_{h}.png"
        if path.exists():
            collector.record_cache(hit=True)
            return path

        collector.record_cache(hit=False)
        png_bytes = self._backend.render.remote(
            prompt=prompt,
            negative_prompt=neg,
            seed=seed,
            size=config.IMAGE_SIZE,
            steps=config.IMAGE_STEPS,
            guidance_scale=0.0,
            remove_bg=True,
        )
        img = Image.open(io.BytesIO(png_bytes))
        img.save(path)
        return path

    def _render(
        self,
        prompt: str,
        seed: int,
        size: int,
        negative_prompt: str = "",
        guidance_scale: float = 0.0,
    ):
        import io  # noqa: PLC0415

        from PIL import Image  # noqa: PLC0415

        png_bytes = self._backend.render.remote(
            prompt=prompt,
            negative_prompt=negative_prompt or config.NEGATIVE_PROMPT,
            seed=seed,
            size=size,
            steps=config.IMAGE_STEPS,
            guidance_scale=guidance_scale,
            remove_bg=False,
        )
        return Image.open(io.BytesIO(png_bytes))


def get_painter() -> Painter:
    if config.USE_MOCK:
        return MockPainter()
    if config.IMAGE_BACKEND == "modal":
        return ModalPainter()
    if config.IMAGE_BACKEND == "lightning":
        return SdxlLightningPainter()
    return SdxlTurboPainter()