# app.py

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

# Initialiser l'application FastAPI
app = FastAPI()

# Charger le modèle
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to(device)

# Définir le modèle de requête
class Prompt(BaseModel):
    text: str

@app.post("/generate")
def generate_image(prompt: Prompt):
    image = pipe(prompt.text).images[0]
    image.save("output.png")  # Enregistrer l'image générée
    return {"message": "Image generated", "image_path": "output.png"}