# =========================================
# FLUX.2 klein 9B
# 의류 전체 복원
# =========================================
import os
import sys
import torch
from PIL import Image
from diffusers import DiffusionPipeline

# =========================================
# 설정
# =========================================
MODEL_ID = "black-forest-labs/FLUX.2-klein-9B"
DEVICE = "cuda"

def main():
    try:
        if len(sys.argv) < 3:
            raise ValueError("입력 파일과 출력 파일을 파라메터로 전달하세요. (예: python flux2_klein_9b.py input.png output.png)")

        input_path = sys.argv[1]
        output_path = sys.argv[2]

        if not os.path.isfile(input_path):
            raise FileNotFoundError(f"입력 파일을 찾을 수 없습니다: {input_path}")

        # =========================================
        # 모델 로드
        # =========================================
        pipe = DiffusionPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True
        ).to(DEVICE)

        pipe.enable_attention_slicing()
        pipe.enable_model_cpu_offload()

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

        # =========================================
        # Prompt
        # =========================================
        prompt = (
            "A highly detailed, realistic fashion catalog product photo of a complete clothing item. "
            "Front-view, perfectly symmetrical, fully visible from collar to bottom hem with complete sleeves. "
            "Natural fabric folds, preserving the exact original color, texture, and pattern. "
            "Plain white studio background, professional lighting, uncropped."
        )

        # =========================================
        # 생성
        # =========================================
        result = pipe(
            image=image,
            prompt=prompt,
            height=1024,
            width=1024,
            guidance_scale=5.0,
            num_inference_steps=25, # 속도 개선을 위해 50에서 25로 감소
            generator=torch.Generator(device=DEVICE).manual_seed(1234)
        ).images[0]

        # =========================================
        # 저장
        # =========================================
        output_dir = os.path.dirname(output_path)
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)

        result.save(output_path)
        print("Saved:", output_path)
        return 0

    except Exception as exc:
        print(f"ERROR: {exc}")
        return 1


if __name__ == "__main__":
    sys.exit(main())