test_image2feed_pipeline.py

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

Flow:
  Real character image (input)
    -> rembg background removal
    -> Qwen2.5-VL appearance JSON extraction (pixel_prompt_core)
    -> SDXL + lora_v3_32bit pixel art character image
    -> Korean quest to English scene text
    -> SDXL + lora_v3_32bit + lora_bg_v1 feed image

Outputs:
  outputs/image2feed/<name>/
    input_nobg.png
    appearance_raw.txt
    appearance.json
    character.png
    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/image2feed")
CHAR_LORA_DIR = "models/lora_v3_32bit"

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
"""


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--name", default="sample_01")
    parser.add_argument("--character-image", required=True, help="Path to the input character image.")
    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("--character-steps", type=int, default=30)
    parser.add_argument("--feed-seed", type=int, default=123, help="Feed generation seed.")
    parser.add_argument("--force", action="store_true")
    parser.add_argument("--force-char", 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")
    return parser.parse_args()


def remove_background(image_path: Path, out_path: Path):
    from PIL import Image
    from rembg import remove

    print("Removing background...", flush=True)
    image = Image.open(image_path).convert("RGBA")
    removed = remove(image)
    white_bg = Image.new("RGBA", removed.size, (255, 255, 255, 255))
    white_bg.paste(removed, mask=removed.split()[3])
    result = white_bg.convert("RGB")
    result.save(out_path)
    print(f"Background removed: {out_path}", flush=True)
    return result


def extract_appearance(nobg_image, out_dir: Path, use_4bit: bool):
    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)
    raw = run_extraction(nobg_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)
    print(f"  pixel_prompt_core: {info.get('pixel_prompt_core', '')}", flush=True)

    del model, processor
    gc.collect()
    try:
        import torch
        torch.cuda.empty_cache()
    except Exception:
        pass
    return info


def build_character_prompt(appearance: dict):
    from src.feed.feed_generator_1 import _to_appearance_str

    appearance_str = _to_appearance_str(appearance)
    body_shape = appearance.get("body_shape", "") if isinstance(appearance, dict) else ""
    pose = appearance.get("pose", "") if isinstance(appearance, dict) else ""

    shape_hint = f"{body_shape}, " if body_shape else ""
    pose_hint = f"{pose}, " if pose else ""

    # CLIP-L: 스타일 + 원본의 체형/자세 (비율 강제 없음)
    prompt = (
        "monglestyle, "
        f"{shape_hint}{pose_hint}"
        "single stuffed animal toy mascot character, full body, centered, "
        "32-bit pixel art sprite, "
        "soft pixel shading, clean silhouette, soft brown outline, "
        "pure white background"
    )
    # OpenCLIP-bigG: 캐릭터 외형 전체
    prompt_2 = f"monglestyle, {appearance_str}"
    return prompt, prompt_2


def generate_character(appearance: dict, 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, prompt_2 = build_character_prompt(appearance)
    print(f"Character prompt  : {prompt}", flush=True)
    print(f"Character prompt_2: {prompt_2}", flush=True)
    image = pipe(
        prompt=prompt,
        prompt_2=prompt_2,
        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 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 generate_feed(quest_en: str, appearance, out_path: Path, seed: int):
    from src.feed.feed_generator_2 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("Image-to-feed pipeline started")
    status(f"Output dir: {out_dir}")

    src_image = Path(args.character_image).resolve()
    if not src_image.exists():
        raise FileNotFoundError(f"--character-image not found: {src_image}")

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

    # STEP 1: background removal
    if args.force or not nobg_path.exists():
        status("STEP 1: removing background")
        nobg_image = remove_background(src_image, nobg_path)
    else:
        from PIL import Image
        nobg_image = Image.open(nobg_path).convert("RGB")
        status(f"STEP 1: using cached nobg image: {nobg_path}")

    # STEP 2: extract appearance with Qwen2.5-VL
    if args.force or not appearance_path.exists():
        status("STEP 2: extracting appearance JSON with Qwen2.5-VL")
        appearance = extract_appearance(nobg_image, 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}")

    # STEP 3: generate pixel art character from full appearance description
    if args.force or args.force_char or not character_path.exists():
        status("STEP 3: generating pixel art character with SDXL + LoRA")
        generate_character(appearance, character_path, args.seed, args.character_steps)
    else:
        status(f"STEP 3: using cached character: {character_path}")

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

    # STEP 4: translate Korean quest to English scene
    if args.quest_en:
        quest_en = args.quest_en
        status("STEP 4: using provided English quest scene")
    elif args.force or args.force_quest or not quest_path.exists():
        status("STEP 4: 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 4: 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}")

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

    results = {
        "name": args.name,
        "character_image_input": str(src_image),
        "input_nobg": str(nobg_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("Image-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