import torch
from PIL import Image
from diffusers import StableDiffusionXLImg2ImgPipeline

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

INPUT_PATH = "input/test3.png"
OUTPUT_PATH = "output/output1.png"

DEVICE = "cuda"

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

image = Image.open(INPUT_PATH).convert("RGB")

# SDXL 권장 해상도
image = image.resize((1024, 1024))

# =========================================
# Base Img2Img Pipeline
# =========================================

base = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True
).to(DEVICE)

# 메모리 최적화
base.enable_attention_slicing()

# =========================================
# Refiner Pipeline
# =========================================

refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    text_encoder_2=base.text_encoder_2,
    vae=base.vae,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True
).to(DEVICE)

refiner.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.

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

Professional ecommerce catalog photography.
Official online store product shot.
Premium high-resolution apparel photography.

Preserve the original garment exactly:
original color,
original fabric texture,
original material,
original stitching,
existing logos,
existing prints,
existing patterns.

Smooth wrinkle-free fabric.
Luxury shopping mall catalog style.
"""

negative_prompt = """
human,
mannequin,
hanger,
wrinkles,
fabric creases,
crumpled fabric,
tilted garment,
asymmetrical clothing,
warped shape,
cropped clothing,
extra logos,
extra text,
blurry,
low quality
"""

# =========================================
# Base Generate
# =========================================

latent = base(
    prompt=prompt,
    negative_prompt=negative_prompt,
    image=image,

    strength=0.4,
    guidance_scale=7,
    num_inference_steps=40,

    denoising_end=0.8,
    output_type="latent"
).images

# =========================================
# Refiner
# =========================================

final_image = refiner(
    prompt=prompt,
    negative_prompt=negative_prompt,
    image=latent,

    strength=0.4,
    guidance_scale=7,
    num_inference_steps=40,

    denoising_start=0.8
).images[0]

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

final_image.save(OUTPUT_PATH)

print("Saved:", OUTPUT_PATH)