File size: 5,851 Bytes
2869a12 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 | """
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})
|