import diffusers
import torch
from fastapi import FastAPI, UploadFile, HTTPException
from PIL import Image

app = FastAPI()

# Inicializa el pipeline al arrancar el servidor
@app.on_event("startup")
async def startup_event():
    global pipe
    print("[DEBUG] Cargando modelo Marigold...")
    pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
        "prs-eth/marigold-depth-lcm-v1-0", variant="fp16", torch_dtype=torch.float16
    ).to("cuda")
    print("[DEBUG] Modelo Marigold cargado exitosamente.")

@app.post("/predict-depth/")
async def predict_depth(file: UploadFile):
    try:
        # Verifica si el archivo es una imagen válida
        if not file.content_type.startswith("image/"):
            raise HTTPException(status_code=400, detail="El archivo subido no es una imagen.")

        # Carga la imagen desde el archivo subido
        image = Image.open(file.file).convert("RGB")

        # Realiza la predicción de profundidad
        print("[DEBUG] Realizando predicción de profundidad...")
        depth = pipe(image)

        # Visualiza la profundidad
        vis = pipe.image_processor.visualize_depth(depth.prediction)
        output_path = "predicted_depth.png"
        vis[0].save(output_path)

        return {"message": "Predicción completada", "output_file": output_path}
    except Exception as e:
        print(f"[ERROR] {str(e)}")
        raise HTTPException(status_code=500, detail="Error procesando la imagen.")

    
@app.get("/")
async def root():
    return {"message": "API de generación de mapas de profundidad con Marigold"}