import os
import uuid
import torch
from PIL import Image
from diffusers import DiffusionPipeline
from fastapi import Body
from fastapi import FastAPI
from fastapi import HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles

# =========================
# 설정
# =========================
MODEL_ID = "black-forest-labs/FLUX.2-klein-9B"
DEVICE = "cuda"
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
CACHE_DIR = os.path.join(SCRIPT_DIR, "cache")
OUTPUT_DIR = os.path.join(SCRIPT_DIR, "output")
IMAGE_SIZE = 512

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["POST"],
    allow_headers=["Content-Type"],
)

os.makedirs(OUTPUT_DIR, exist_ok=True)
app.mount("/output", StaticFiles(directory=OUTPUT_DIR), name="output")

# =========================
# 모델 1회 로딩
# =========================
print("Loading model...")

pipe = DiffusionPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,  
    cache_dir=CACHE_DIR
).to(DEVICE)

#pipe = DiffusionPipeline.from_pretrained(
#    MODEL_ID,
#    torch_dtype=torch.float16,
#    low_cpu_mem_usage=True,
#    cache_dir=CACHE_DIR
#).to(DEVICE)

#pipe.enable_attention_slicing()
#pipe.enable_model_cpu_offload()
pipe.set_progress_bar_config(disable=True)

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.benchmark = True
print("Model loaded!")

# =========================
# API
# =========================
@app.post("/generate")
def generate(image: str = Body(...), prompt: str = Body(...)):
    image_path = image
    print(f"image: {image}", flush=True)

    if not os.path.isabs(image_path):
        image_path = os.path.join(SCRIPT_DIR, image_path)

    if not os.path.isfile(image_path):
        raise HTTPException(status_code=404, detail=f"입력 파일을 찾을 수 없습니다: {image}")

    if not prompt.strip():
        raise HTTPException(status_code=400, detail="prompt가 입력되지 않았습니다.")
        
    input_image = Image.open(image_path).convert("RGB")
    input_image = input_image.resize((IMAGE_SIZE, IMAGE_SIZE))

    with torch.inference_mode():
        result = pipe(
            image=input_image,
            prompt=prompt,
            height=IMAGE_SIZE,
            width=IMAGE_SIZE,
            num_inference_steps=15,
            generator=torch.Generator(device=DEVICE).manual_seed(1234)
        ).images[0]

    filename = f"{uuid.uuid4()}.png"
    save_path = os.path.join(OUTPUT_DIR, filename)
    result.save(save_path)

    return {
        "status": "success",
        "image": os.path.relpath(save_path, SCRIPT_DIR)
    }