import torch
import numpy as np
from PIL import Image
from diffusers import StableDiffusionInpaintPipeline

# =========================
# 설정
# =========================

IMAGE_PATH = "test1.png"
MASK_PATH = "top_mask.png"
OUTPUT_PATH = "inpaint_result.png"

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# =========================
# 이미지 로드
# =========================

image = Image.open(IMAGE_PATH).convert("RGB")
mask = Image.open(MASK_PATH).convert("L")

# 원본과 마스크 크기는 반드시 같아야 함
if image.size != mask.size:
    mask = mask.resize(image.size)

# =========================
# top_mask.png 기준으로 상의 영역만 추출
# =========================
image_np = np.array(image)
mask_np = np.array(mask)
mask_bin = mask_np > 127
mask = Image.fromarray((mask_bin.astype(np.uint8) * 255), mode="L")
mask_3ch = np.stack([mask_bin] * 3, axis=-1)
white_bg = np.ones_like(image_np, dtype=np.uint8) * 255
cloth_only = np.where(mask_3ch, image_np, white_bg).astype(np.uint8)
image = Image.fromarray(cloth_only)

# =========================
# 파이프라인 로드
# =========================

pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-inpainting",
    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32
)

pipe = pipe.to(DEVICE)

# =========================
# Inpainting 실행
# =========================

result = pipe(
    prompt=(
        "a clean white t-shirt, realistic fabric texture, "
        "natural folds, realistic lighting"
    ),
    negative_prompt=(
        "distorted clothes, bad anatomy, blurry, artifacts, "
        "extra arms, text, logo"
    ),
    image=image,
    mask_image=mask,
    strength=0.85,
    guidance_scale=7.5,
    num_inference_steps=30
).images[0]

# =========================
# 저장
# =========================

result.save(OUTPUT_PATH)

print("완료:", OUTPUT_PATH)