import torch
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)

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

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)