fluxpublic / app.py
eztrz's picture
Update app.py
d89926a verified
raw
history blame
645 Bytes
# 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"}