src/feed/feed_generator_1.py

"""
몽글마을 피드 이미지 생성기
입력: 캐릭터 외형(Florence-2 추출) + 퀘스트(영어)
출력: 피드 이미지 (캐릭터가 퀘스트를 수행하는 장면)
"""
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"
)


def _to_appearance_str(appearance) -> str:
    """dict 또는 string → 외형 텍스트 변환"""
    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", "")
    texture  = appearance.get("texture", "")
    acc      = appearance.get("accessories", "")
    features = appearance.get("distinctive_features", [])
    secondary = appearance.get("secondary_colors", [])

    parts = [f"a {color} pixel art stuffed {animal} character".strip()]
    if shape:    parts.append(shape)
    if isinstance(secondary, list):
        # 3단어 이하인 짧은 색상 값만 포함 (문장형 묘사 제외)
        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:     parts.append(ears)
    if nose:     parts.append(nose)
    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() != "none": parts.append(acc)
    # distinctive_features는 eye_style, cheeks 등과 중복되므로 제외

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


def build_prompt(appearance, quest_en: str):
    """
    prompt   → 장면/스타일 (CLIP-L)
    prompt_2 → 장면 + 캐릭터 외형 (OpenCLIP-bigG, 더 강한 인코더에 둘 다 전달)
    """
    appearance_str = _to_appearance_str(appearance)
    prompt = (
        f"monglestyle, {quest_en}, "
        f"cute chibi stuffed animal character in action, "
        f"full body pose, 32-bit pixel art scene, pastel colors"
    )
    prompt_2 = f"monglestyle, {appearance_str}"
    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.8, 0.8],
    )
    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=20,
        guidance_scale=7.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 해제 완료")