import os
import uuid
import torch
from PIL import Image
from diffusers import StableDiffusionXLImg2ImgPipeline
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 = "stabilityai/stable-diffusion-xl-base-1.0"
#MODEL_ID = "RunDiffusion/Juggernaut-XL-v9"
DEVICE = "cuda"
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
OUTPUT_DIR = os.path.join(SCRIPT_DIR, "output")
CANVAS_SIZE = 1024

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 = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True,
).to(DEVICE)

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

pipe.vae.to(dtype=torch.float32)
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(...), negative_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가 입력되지 않았습니다.")

    if not negative_prompt.strip():
        raise HTTPException(status_code=400, detail="negative_prompt가 입력되지 않았습니다.")

    input_image = Image.open(image_path).convert("RGB")
    input_image = input_image.resize((CANVAS_SIZE, CANVAS_SIZE))

    with torch.inference_mode():
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=input_image,
            strength=0.55,
            guidance_scale=10.0,
            num_inference_steps=50,
            generator=torch.Generator(device=DEVICE).manual_seed(42),
        ).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),
    }