File size: 1,183 Bytes
f102a76 a1195f3 e2cc4e8 f102a76 e2cc4e8 f102a76 e2cc4e8 f102a76 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
from fastapi import FastAPI, Response
from fastapi.middleware.cors import CORSMiddleware
import torch
from torch.cuda.amp import autocast
from diffusers import DiffusionPipeline
from diffusers import StableDiffusionPipeline
from io import BytesIO
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_credentials=True,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"]
)
# MODEL 1.4
model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionPipeline.from_pretrained(
model_id, torch_dtype=torch.float32, cache_dir="./cache"
)
# MOdel 1.0
# model_id = "stabilityai/stable-diffusion-xl-base-1.0"
# device = "cuda" if torch.cuda.is_available() else "cpu"
# pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32, use_safetensors=True, variant="fp16")
pipe = pipe.to(device)
@app.get("/")
def generate(prompt: str):
with autocast(device):
image = pipe(prompt, guidance_scale=8.5).images[0]
buffer = BytesIO()
image.save(buffer, format="PNG")
buffer.seek(0)
return Response(content=buffer.getvalue(), media_type="image/png")
|