from fastapi import FastAPI
from pydantic import BaseModel
from diffusers import StableDiffusionPipeline
import torch

# Definir la API
app = FastAPI()

# Cargar el modelo en CPU (o en CUDA si tienes GPU disponible)
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to("cpu")

# Definir el esquema de entrada
class PromptInput(BaseModel):
    text: str

# Endpoint para generar imágenes
@app.post("/generate")
async def generate(prompt: PromptInput):
    image = pipe(prompt.text).images[0]
    image.save("output.png")
    return {"message": "Imagen generada", "path": "output.png"}

# Bloque para iniciar el servidor con uvicorn
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=7860, log_level="info")