test_text2feed_pipeline.py

"""
Text-to-feed end-to-end test pipeline.

Flow:
  Korean character text
    -> English visual character prompt
    -> SDXL + lora_v3_32bit character image
    -> Qwen2.5-VL appearance JSON extraction
    -> Korean quest to English scene text
    -> SDXL + lora_v3_32bit + lora_bg_v1 feed image

Outputs:
  outputs/text2feed/<name>/
    character.png
    appearance_raw.txt
    appearance.json
    quest.json
    feed.png
    results.json
"""
import argparse
import gc
import json
import os
import sys
import traceback
from pathlib import Path

os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"

ROOT = Path(__file__).resolve().parent
os.chdir(ROOT)
sys.path.insert(0, str(ROOT))

OUTPUT_ROOT = Path("outputs/text2feed")
CHAR_LORA_DIR = "models/lora_v3_32bit"

DEFAULT_CHARACTER_KO = (
    "이 친구는 부드러운 체리핑크색 곰 인형이에요. "
    "크림색 배를 가지고 있고, 작은 둥근 귀와 분홍색 코, "
    "짧고 통통한 팔다리와 포근한 표정을 가지고 있어요."
)
DEFAULT_CHARACTER_EN = (
    "soft cherry pink bear plush mascot with a cream white belly, "
    "small rounded ears, tiny pink oval nose, simple smiling mouth, "
    "short rounded limbs, smooth plush texture"
)
DEFAULT_QUEST_KO = "공원에서 30분 달리기를 완료했어요!"

CHARACTER_NEGATIVE = (
    "realistic, photograph, 3d render, human, person, anime, scary, "
    "complex background, multiple characters, text, watermark, logo, "
    "extra limbs, long limbs, harsh black outline, low quality, blurry"
)

QUEST_SYSTEM = """You are an action-pose scene prompt writer for Mongle Village, a cozy sky island pixel art village.

Convert a Korean quest completion message into a short English scene description for image generation.

Rules:
- The FIRST words must describe the character's visible action pose.
- Keep the activity from the quest very explicit.
- Include body movement cues such as arms swinging, legs moving, holding a book, stirring a pot, walking steps.
- Describe one simple matching environment after the action.
- Set it in a cozy pastel sky island village world.
- Mention only one character.
- Do not write only a landscape/background description.
- Do not make the character standing still unless the quest is resting.
- 16-28 words max.
- Output ONLY the English scene description.

Examples:
Input: 공원에서 30분 달리기를 완료했어요!
Output: running with arms swinging and legs in motion along a fluffy cloud meadow path

Input: 오늘 책 한 권을 다 읽었어요
Output: sitting and holding an open storybook under a blossoming cloud tree beside a cozy cottage

Input: 직접 요리해서 건강한 밥을 먹었어요
Output: cooking with both paws stirring a pot in a cozy cottage kitchen with fresh vegetables
"""

CHARACTER_SYSTEM = """You are a visual prompt writer for a plush-to-pixel-art character generation pipeline.

Convert a Korean character description into a concise English visual prompt.

Rules:
- Focus only on visible appearance: animal type, body color, face, ears, body shape, limbs, texture, accessories.
- Convert personality words into visible expression only.
- If animal type is not specified, choose a soft plush animal that fits the description.
- If color is not specified, choose one pastel color.
- Do not include background, scene, action, story, name, or relationship.
- 25-45 words.
- Output ONLY the English visual prompt.

Example:
Input: 이 친구는 용감하고 사랑스러운 하얀 토끼 인형이에요. 볼이 발그레하고 포근해요.
Output: white bunny plush mascot, rosy cheeks, confident bright eyes, warm gentle smile, soft round chubby body, short rounded limbs, fluffy plush texture
"""


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--name", default="sample_01")
    parser.add_argument("--character-ko", default=DEFAULT_CHARACTER_KO)
    parser.add_argument(
        "--character-text",
        default="",
        help="English visual character prompt. If omitted, --character-ko is translated first.",
    )
    parser.add_argument("--quest-ko", default=DEFAULT_QUEST_KO)
    parser.add_argument("--quest-en", default="", help="If provided, skips Qwen text translation.")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--force", action="store_true")
    parser.add_argument("--force-quest", action="store_true")
    parser.add_argument("--force-feed", action="store_true")
    parser.add_argument("--no-4bit-vlm", action="store_true")
    parser.add_argument("--character-steps", type=int, default=30)
    parser.add_argument("--feed-seed", type=int, default=123)
    return parser.parse_args()


def character_prompt(character_text: str) -> str:
    return (
        "monglestyle, "
        f"{character_text}, "
        "single stuffed animal toy mascot character, full body, centered, "
        "front view, cute chibi proportions, 32-bit pixel art sprite, "
        "soft pixel shading, clean silhouette, soft brown outline, "
        "pure white background"
    )


def translate_character(character_ko: str):
    print("[character translation] importing load_qwen...", flush=True)
    from src.pipeline.persona2prompt import load_qwen, unload_qwen
    print("[character translation] importing torch...", flush=True)
    import torch

    print("Loading Qwen text model for character translation...", flush=True)
    model, tokenizer = load_qwen()
    messages = [
        {"role": "system", "content": CHARACTER_SYSTEM},
        {"role": "user", "content": character_ko},
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt").to("cuda")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=120,
            do_sample=False,
            temperature=None,
            top_p=None,
            pad_token_id=tokenizer.eos_token_id,
        )
    generated = outputs[0][inputs.input_ids.shape[1] :]
    character_en = tokenizer.decode(generated, skip_special_tokens=True).strip()
    unload_qwen(model, tokenizer)
    return character_en


def generate_character(character_text: str, out_path: Path, seed: int, steps: int):
    import torch
    from diffusers import DPMSolverMultistepScheduler, StableDiffusionXLPipeline

    print("Loading SDXL character pipeline...", flush=True)
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0",
        torch_dtype=torch.float16,
        use_safetensors=True,
    )
    pipe.load_lora_weights(CHAR_LORA_DIR)
    pipe.fuse_lora(lora_scale=0.9)
    pipe.unload_lora_weights()
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(
        pipe.scheduler.config, use_karras_sigmas=True
    )
    pipe.to("cuda")
    pipe.enable_attention_slicing()

    prompt = character_prompt(character_text)
    print(f"Character prompt: {prompt}", flush=True)
    image = pipe(
        prompt=prompt,
        negative_prompt=CHARACTER_NEGATIVE,
        num_inference_steps=steps,
        guidance_scale=7.5,
        height=1024,
        width=1024,
        generator=torch.Generator("cuda").manual_seed(seed),
    ).images[0]
    image.save(out_path)
    print(f"Character saved: {out_path}", flush=True)

    del pipe
    gc.collect()
    torch.cuda.empty_cache()


def extract_appearance(character_path: Path, out_dir: Path, use_4bit: bool):
    from PIL import Image
    from test_qwen25_vl_extract import (
        extract_json_from_text,
        load_model,
        normalize_info,
        run_extraction,
    )

    print("Loading Qwen2.5-VL appearance extractor...", flush=True)
    model, processor = load_model("Qwen/Qwen2.5-VL-7B-Instruct", use_4bit=use_4bit)
    image = Image.open(character_path).convert("RGB")
    raw = run_extraction(image, model, processor, max_new_tokens=900)

    raw_path = out_dir / "appearance_raw.txt"
    raw_path.write_text(raw, encoding="utf-8")

    info = normalize_info(extract_json_from_text(raw))
    appearance_path = out_dir / "appearance.json"
    appearance_path.write_text(json.dumps(info, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"Appearance saved: {appearance_path}", flush=True)

    del model, processor
    gc.collect()
    try:
        import torch

        torch.cuda.empty_cache()
    except Exception:
        pass
    return info


def translate_quest(quest_ko: str):
    print("[quest translation] importing load_qwen...", flush=True)
    from src.pipeline.persona2prompt import load_qwen, unload_qwen
    print("[quest translation] importing torch...", flush=True)
    import torch

    print("Loading Qwen text model for quest translation...", flush=True)
    model, tokenizer = load_qwen()
    messages = [
        {"role": "system", "content": QUEST_SYSTEM},
        {"role": "user", "content": quest_ko},
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt").to("cuda")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=80,
            do_sample=False,
            temperature=None,
            top_p=None,
            pad_token_id=tokenizer.eos_token_id,
        )
    generated = outputs[0][inputs.input_ids.shape[1] :]
    quest_en = tokenizer.decode(generated, skip_special_tokens=True).strip()
    unload_qwen(model, tokenizer)
    return quest_en


def remove_character_bg(character_path: Path, out_path: Path):
    from PIL import Image
    from rembg import remove

    print("Removing character background...", flush=True)
    image = Image.open(character_path).convert("RGBA")
    result = remove(image)
    result.save(out_path)
    print(f"Character nobg saved: {out_path}", flush=True)


def generate_feed(quest_en: str, appearance, out_path: Path, seed: int):
    from src.feed.feed_generator_1 import generate, load_pipeline, unload_pipeline

    print("Loading feed generation pipeline...", flush=True)
    pipe = load_pipeline()
    image = generate(appearance, quest_en, pipe, seed=seed)
    image.save(out_path)
    print(f"Feed saved: {out_path}", flush=True)
    unload_pipeline(pipe)


def main():
    args = parse_args()
    out_dir = OUTPUT_ROOT / args.name
    out_dir.mkdir(parents=True, exist_ok=True)
    status_path = out_dir / "status.log"

    def status(message: str):
        print(message, flush=True)
        with status_path.open("a", encoding="utf-8") as f:
            f.write(message + "\n")

    status("=" * 60)
    status("Text-to-feed pipeline started")
    status(f"Output dir: {out_dir}")

    character_path = out_dir / "character.png"
    character_nobg_path = out_dir / "character_nobg.png"
    character_prompt_path = out_dir / "character_prompt.json"
    appearance_path = out_dir / "appearance.json"
    quest_path = out_dir / "quest.json"
    feed_path = out_dir / "feed.png"
    results_path = out_dir / "results.json"

    if args.character_text:
        character_text = args.character_text
        status("STEP 0: using provided English character prompt")
    elif args.force or not character_prompt_path.exists():
        status("STEP 0: translating Korean character description to English prompt")
        character_text = translate_character(args.character_ko)
    else:
        character_text = json.loads(character_prompt_path.read_text(encoding="utf-8"))["character_text"]
        status(f"STEP 0: using cached character prompt: {character_prompt_path}")

    character_prompt_data = {
        "character_ko": args.character_ko,
        "character_text": character_text,
    }
    character_prompt_path.write_text(
        json.dumps(character_prompt_data, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    status(f"Character EN: {character_text}")

    if args.force or not character_path.exists():
        status("STEP 1: generating character image")
        generate_character(character_text, character_path, args.seed, args.character_steps)
    else:
        status(f"STEP 1: using cached character: {character_path}")

    if args.force or not character_nobg_path.exists():
        status("STEP 1.5: removing character background")
        remove_character_bg(character_path, character_nobg_path)
    else:
        status(f"STEP 1.5: using cached character nobg: {character_nobg_path}")

    if args.force or not appearance_path.exists():
        status("STEP 2: extracting appearance JSON with Qwen2.5-VL")
        appearance = extract_appearance(character_path, out_dir, use_4bit=not args.no_4bit_vlm)
    else:
        appearance = json.loads(appearance_path.read_text(encoding="utf-8"))
        status(f"STEP 2: using cached appearance: {appearance_path}")

    if args.quest_en:
        quest_en = args.quest_en
        status("STEP 3: using provided English quest scene")
    elif args.force or args.force_quest or not quest_path.exists():
        status("STEP 3: translating Korean quest to English scene")
        quest_en = translate_quest(args.quest_ko)
    else:
        quest_en = json.loads(quest_path.read_text(encoding="utf-8"))["quest_en"]
        status(f"STEP 3: using cached quest translation: {quest_path}")

    quest_data = {
        "quest_ko": args.quest_ko,
        "quest_en": quest_en,
    }
    quest_path.write_text(json.dumps(quest_data, ensure_ascii=False, indent=2), encoding="utf-8")
    status(f"Quest EN: {quest_en}")

    if args.force or args.force_feed or not feed_path.exists():
        status("STEP 4: generating feed image")
        generate_feed(quest_en, appearance, feed_path, args.feed_seed)
    else:
        status(f"STEP 4: using cached feed: {feed_path}")

    results = {
        "name": args.name,
        "character_ko": args.character_ko,
        "character_text": character_text,
        "character_prompt_json": str(character_prompt_path).replace("\\", "/"),
        "character_image": str(character_path).replace("\\", "/"),
        "character_nobg": str(character_nobg_path).replace("\\", "/"),
        "appearance_json": str(appearance_path).replace("\\", "/"),
        "appearance": appearance,
        "quest": quest_data,
        "feed_image": str(feed_path).replace("\\", "/"),
    }
    results_path.write_text(json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")
    status(f"Results saved: {results_path}")
    status("Text-to-feed pipeline finished")


if __name__ == "__main__":
    try:
        main()
    except Exception:
        fallback_dir = OUTPUT_ROOT / "sample_01"
        fallback_dir.mkdir(parents=True, exist_ok=True)
        error_path = fallback_dir / "error.log"
        error_text = traceback.format_exc()
        error_path.write_text(error_text, encoding="utf-8")
        print(error_text, flush=True)
        print(f"Error saved: {error_path}", flush=True)
        raise