handler_character.py

"""
RunPod handler for Mongle character generation.

Inputs:
  Text branch:
    {
      "input": {
        "mode": "text",
        "character_ko": "...",
        "character_text": "optional English prompt",
        "name": "optional_job_name",
        "seed": 42,
        "character_steps": 30,
        "no_4bit_vlm": false
      }
    }

  Image branch:
    {
      "input": {
        "mode": "image",
        "character_image": "<base64 png/jpg or data URL>",
        "name": "optional_job_name",
        "seed": 42,
        "character_steps": 30,
        "no_4bit_vlm": false
      }
    }

Outputs:
  {
    "mode": "text" | "image",
    "name": "...",
    "character_image": "<base64 png>",
    "appearance": {...},
    "appearance_raw": "...",
    "character_text": "...",       # text branch only
    "input_nobg": "<base64 png>"   # image branch only
  }
"""

import base64
import io
import json
import os
import traceback
from pathlib import Path

import runpod
from PIL import Image

from test_image2feed_pipeline import (
    extract_appearance as extract_image_appearance,
    generate_character as generate_image_character,
    remove_background,
)
from test_text2feed_pipeline import (
    DEFAULT_CHARACTER_KO,
    extract_appearance as extract_text_appearance,
    generate_character as generate_text_character,
    translate_character,
)


OUTPUT_ROOT = Path("outputs/handler_character")


def _decode_image(data: str) -> Image.Image:
    if data.startswith("data:image"):
        data = data.split(",", 1)[1]
    return Image.open(io.BytesIO(base64.b64decode(data))).convert("RGB")


def _encode_png(image_or_path) -> str:
    if isinstance(image_or_path, (str, Path)):
        image = Image.open(image_or_path).convert("RGB")
    else:
        image = image_or_path.convert("RGB")
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


def _write_input_image(image_b64: str, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    image = _decode_image(image_b64)
    image.save(path)


def _read_text(path: Path) -> str:
    if not path.exists():
        return ""
    return path.read_text(encoding="utf-8")


def _text_branch(payload: dict, out_dir: Path) -> dict:
    character_ko = payload.get("character_ko") or DEFAULT_CHARACTER_KO
    character_text = payload.get("character_text") or ""
    seed = int(payload.get("seed", 42))
    steps = int(payload.get("character_steps", 30))
    use_4bit = not bool(payload.get("no_4bit_vlm", False))

    out_dir.mkdir(parents=True, exist_ok=True)
    character_path = out_dir / "character.png"

    if not character_text:
        character_text = translate_character(character_ko)

    generate_text_character(character_text, character_path, seed=seed, steps=steps)
    appearance = extract_text_appearance(character_path, out_dir, use_4bit=use_4bit)

    result = {
        "mode": "text",
        "character_ko": character_ko,
        "character_text": character_text,
        "character_image": _encode_png(character_path),
        "appearance": appearance,
        "appearance_raw": _read_text(out_dir / "appearance_raw.txt"),
    }
    (out_dir / "result.json").write_text(
        json.dumps({k: v for k, v in result.items() if k != "character_image"}, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    return result


def _image_branch(payload: dict, out_dir: Path) -> dict:
    image_b64 = payload.get("character_image") or payload.get("image")
    if not image_b64:
        raise ValueError("Missing input.character_image base64 PNG/JPEG data.")

    seed = int(payload.get("seed", 42))
    steps = int(payload.get("character_steps", 30))
    use_4bit = not bool(payload.get("no_4bit_vlm", False))

    out_dir.mkdir(parents=True, exist_ok=True)
    input_path = out_dir / "input.png"
    nobg_path = out_dir / "input_nobg.png"
    character_path = out_dir / "character.png"

    _write_input_image(image_b64, input_path)
    nobg_image = remove_background(input_path, nobg_path)
    appearance = extract_image_appearance(nobg_image, out_dir, use_4bit=use_4bit)
    generate_image_character(appearance, character_path, seed=seed, steps=steps)

    result = {
        "mode": "image",
        "input_nobg": _encode_png(nobg_path),
        "character_image": _encode_png(character_path),
        "appearance": appearance,
        "appearance_raw": _read_text(out_dir / "appearance_raw.txt"),
    }
    (out_dir / "result.json").write_text(
        json.dumps(
            {k: v for k, v in result.items() if k not in {"input_nobg", "character_image"}},
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )
    return result


def handler(event):
    payload = event.get("input", {}) if isinstance(event, dict) else {}
    mode = (payload.get("mode") or "").lower().strip()
    if not mode:
        mode = "image" if payload.get("character_image") or payload.get("image") else "text"

    name = payload.get("name") or f"{mode}_job"
    out_dir = OUTPUT_ROOT / name

    try:
        if mode == "text":
            result = _text_branch(payload, out_dir)
        elif mode == "image":
            result = _image_branch(payload, out_dir)
        else:
            raise ValueError("input.mode must be 'text' or 'image'.")

        result["name"] = name
        result["output_dir"] = str(out_dir).replace("\\", "/")
        return result
    except Exception:
        out_dir.mkdir(parents=True, exist_ok=True)
        error_text = traceback.format_exc()
        (out_dir / "error.log").write_text(error_text, encoding="utf-8")
        return {"error": error_text, "name": name, "output_dir": str(out_dir).replace("\\", "/")}


if os.getenv("RUNPOD_SERVERLESS", "1") == "1":
    runpod.serverless.start({"handler": handler})