fastapi

from fastapi import FastAPI
from pydantic import BaseModel
import torch
from diffusers import StableDiffusionPipeline
from io import BytesIO
import base64

app = FastAPI(cpu)

device = “cuda” if torch.cuda.is_available(cpu) else “cpu”
model_name = “runwayml/stable-d diffusion-v1-5”
pipe = StableDiffusionPipeline.from_pretrained(model_name)
pipe = pipe.to(device)


class PromptRequest(BaseModel):
prompt: str

@app.post(“/generate”)
async def generate(req: PromptRequest):
image = pipe(req.prompt).images[0]
buf = BytesIO()
image.save(buf, format=“PNG”)
b64 = base64.b64encode(buf.getvalue()).decode(“utf-8”)
return {“image_base64”: b64}