src/feed/feed_generator_2.py

"""
몽글마을 피드 이미지 생성기 v2 (이미지 입력 파이프라인 전용)

feed_generator_1.py와의 차이점:
- 캐릭터 외형을 프롬프트 앞쪽에 배치 (SDXL 토큰 가중치 활용)
- appearance_str에 limbs 포함 (비전형 캐릭터 신체 반영)
- NEGATIVE에 human/person 추가
- char LoRA 0.9 / bg LoRA 1.1 (외형 강제력 상향)
- guidance_scale 8.5 (프롬프트 추종력 강화)
- 양쪽 인코더 모두 scene 정보 전달
"""
import os as _os; _os.chdir(_os.path.dirname(_os.path.dirname(_os.path.dirname(_os.path.abspath(__file__)))))
import os, gc
from pathlib import Path
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"

from diffusers import StableDiffusionXLPipeline, DPMSolverMultistepScheduler
import torch

CHARACTER_LORA = "models/lora_v3_32bit/pytorch_lora_weights.safetensors"
BG_LORA        = "models/lora_bg_v1/pytorch_lora_weights.safetensors"

NEGATIVE = (
    "realistic, 3d render, photograph, blurry, dark, gloomy, "
    "watercolor, sketch, painterly, smooth illustration, "
    "modern city, urban, scary, violence, text, watermark, "
    "multiple characters, tiling, repeated, "
    "human, person, girl, boy, man, woman, humanoid, human body, human face, "
    "skin, hair, clothes, dress, shirt"
)


def _to_appearance_str(appearance) -> str:
    if isinstance(appearance, str):
        return appearance.strip()
    if not isinstance(appearance, dict):
        return str(appearance).strip()

    animal    = appearance.get("animal_type", "bear")
    color     = appearance.get("body_color", "")
    shape     = appearance.get("body_shape", "")
    eye       = appearance.get("eye_style", "")
    nose      = appearance.get("nose_mouth") or appearance.get("nose_shape", "")
    cheeks    = appearance.get("cheeks", "")
    ears      = appearance.get("ear_shape", "")
    limbs     = appearance.get("limbs", "")
    texture   = appearance.get("texture", "")
    acc       = appearance.get("accessories", "")
    secondary = appearance.get("secondary_colors", [])

    parts = [f"a {color} pixel art stuffed {animal} character".strip()]
    if shape: parts.append(shape)
    if isinstance(secondary, list):
        parts.extend(s for s in secondary[:2] if s and len(str(s).split()) <= 3)
    if eye:   parts.append(eye)
    if cheeks and "no" not in cheeks.lower(): parts.append(cheeks)
    if ears and "not visible" not in ears.lower(): parts.append(ears)
    if nose and "not visible" not in nose.lower(): parts.append(nose)
    if limbs and "not visible" not in limbs.lower(): parts.append(limbs)
    if texture: parts.append(texture)
    if isinstance(acc, list):
        acc = ", ".join(str(a) for a in acc if a)
    if acc and str(acc).lower() not in ("none", ""): parts.append(acc)

    return ", ".join(p for p in parts if p)


def build_prompt(appearance, quest_en: str):
    """
    두 인코더 역할 분리:
    prompt   (CLIP-L)       → 장면/배경 담당: scene 앞에 배치
    prompt_2 (OpenCLIP-bigG) → 캐릭터 외형 담당: appearance 앞에 배치
    """
    appearance_str = _to_appearance_str(appearance)

    animal = "plush mascot"
    if isinstance(appearance, dict):
        animal = appearance.get("animal_type", "plush mascot")

    # CLIP-L: 장면 먼저 → 배경 생성 담당
    # 캐릭터 타입만 간단히 명시해 human 방지
    prompt = (
        f"monglestyle, {quest_en}, "
        f"stuffed {animal} plush toy mascot, "
        f"cozy pastel sky island village, "
        f"32-bit pixel art scene, detailed background, outdoor scenery, pastel colors"
    )
    # OpenCLIP-bigG: 외형 먼저 → 캐릭터 정체성 담당
    # 장면 힌트도 뒤에 포함해 배경 보조
    prompt_2 = (
        f"monglestyle, stuffed {animal} plush toy mascot, "
        f"{appearance_str}, "
        f"full body in action, {quest_en}, "
        f"pixel art background scene"
    )
    return prompt, prompt_2


def load_pipeline():
    print("SDXL 로드 중 (캐릭터 LoRA + 배경 LoRA)...")
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0",
        torch_dtype=torch.float16,
        use_safetensors=True,
    ).to("cuda")

    if not Path(CHARACTER_LORA).exists():
        raise FileNotFoundError(f"Character LoRA not found: {CHARACTER_LORA}")
    if not Path(BG_LORA).exists():
        raise FileNotFoundError(f"Background LoRA not found: {BG_LORA}")

    pipe.load_lora_weights(CHARACTER_LORA, adapter_name="character")
    pipe.load_lora_weights(BG_LORA,        adapter_name="bg")
    pipe.set_adapters(
        ["character", "bg"],
        adapter_weights=[0.9, 1.1],
    )
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(
        pipe.scheduler.config, use_karras_sigmas=True
    )
    pipe.enable_attention_slicing()
    print("로드 완료!\n")
    return pipe


def generate(appearance, quest_en: str, pipe, seed: int = 42):
    prompt, prompt_2 = build_prompt(appearance, quest_en)
    print(f"  prompt  : {prompt}")
    print(f"  prompt_2: {prompt_2}")

    return pipe(
        prompt=prompt,
        prompt_2=prompt_2,
        negative_prompt=NEGATIVE,
        num_inference_steps=25,
        guidance_scale=8.5,
        height=1024,
        width=1024,
        generator=torch.Generator("cuda").manual_seed(seed),
    ).images[0]


def unload_pipeline(pipe):
    del pipe
    gc.collect()
    torch.cuda.empty_cache()
    print("파이프라인 VRAM 해제 완료")