import torch
from PIL import Image
from diffusers import StableDiffusionXLImg2ImgPipeline

# =========================
# 설정
# =========================
MODEL_ID = "RunDiffusion/Juggernaut-XL-v9"
INPUT_PATH = "input/test2.png"
OUTPUT_PATH = "output/output.png"

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

# =========================
# 이미지 로드
# =========================
image = Image.open(INPUT_PATH).convert("RGB")
image = image.resize((1024, 1024))

# =========================
# Pipeline 로드
# =========================
pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=DTYPE,
    variant="fp16",
    use_safetensors=True
).to(DEVICE)

pipe.enable_attention_slicing()

# =========================
# Prompt
# =========================
prompt = """
Ultra realistic high-end fashion ecommerce catalog product photo of a single apparel item.

Front view, perfectly centered composition.
Perfectly laid flat on a pure white studio background.
Perfectly flattened garment.
Ironed flat appearance.
Professionally pressed fabric.
Freshly ironed clothing surface.
Crisp and straight garment shape.
Evenly spread fabric.
Perfect left-right symmetry.
Balanced garment proportions.
Aligned sleeves and hem.

Full garment completely visible from top to bottom, uncropped.
Clean garment outline with tidy edges.

Professional ecommerce catalog photography.
Official online store product shot.
Clean softbox studio lighting, soft natural shadow, sharp focus.
Highly detailed textile texture, seams, and stitching.

Preserve the original garment color, fabric texture, material, stitching, seams,
existing logos, existing prints, existing patterns, and proportions.

Do not redesign or modify the garment.
No new logos, no new text, no new typography, no new branding.
Smooth wrinkle-free fabric.
Luxury shopping mall catalog style.
"""

negative_prompt = """
human, person, mannequin, ghost mannequin, body, hands, hanger,
cropped, cut off, incomplete clothing,
tilted, twisted, asymmetrical,
distorted, deformed, warped,
uneven shape, misaligned sleeves, unbalanced proportions,
wrinkles, fabric creases, crumpled fabric, messy folds,
overlapping fabric, double layers,
extra logo, extra text, typography, new branding, symbols, letters,
low quality, blurry, noisy, bad lighting, background objects
"""

# =========================
# Img2Img 실행
# =========================
result = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    image=image,
    strength=0.6,           # 너무 높으면 원본 이미지에서 멀어질 수 있음
    guidance_scale=10,     # 너무 높으면 왜곡 가능
    num_inference_steps=35
).images[0]

# =========================
# 저장
# =========================
result.save(OUTPUT_PATH)
print("Saved:", OUTPUT_PATH)