Spaces:
Sleeping
Sleeping
File size: 2,872 Bytes
bd94aa2 |
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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
from fastapi import FastAPI, Request, Form, File, UploadFile
from fastapi.responses import StreamingResponse
from contextlib import asynccontextmanager
from starlette.middleware.cors import CORSMiddleware
import torch
from PIL import Image
from io import BytesIO
from diffusers import (
AutoPipelineForText2Image,
AutoPipelineForImage2Image,
AutoPipelineForInpainting,
)
@asynccontextmanager
async def lifespan(app: FastAPI):
text2img = AutoPipelineForText2Image.from_pretrained(
"stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16"
).to("cpu")
img2img = AutoPipelineForImage2Image.from_pipe(text2img).to("cpu")
inpaint = AutoPipelineForInpainting.from_pipe(text2img).to("cpu")
yield {"text2img": text2img, "img2img": img2img, "inpaint": inpaint}
del text2img
del img2img
del inpaint
app = FastAPI(lifespan=lifespan)
origins = ["*"]
app.add_middleware(
CORSMiddleware,
allow_origins=origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.get("/")
async def root():
return {"Hello": "World"}
@app.post("/text-to-image/")
async def text_to_image(request: Request, prompt: str = Form(...)):
image = request.state.text2img(
prompt=prompt, num_inference_steps=1, guidance_scale=0.0
).images[0]
bytes = BytesIO()
image.save(bytes, "PNG")
bytes.seek(0)
return StreamingResponse(bytes, media_type="image/png")
@app.post("/image-to-image/")
async def image_to_image(
request: Request, prompt: str = Form(...), init_image: UploadFile = File(...)
):
bytes = await init_image.read()
init_image = Image.open(BytesIO(bytes))
init_image = init_image.convert("RGB").resize((512, 512))
image = request.state.img2img.pipe(
prompt,
image=init_image,
num_inference_steps=2,
strength=0.5,
guidance_scale=0.0,
).images[0]
bytes = BytesIO()
image.save(bytes, "PNG")
bytes.seek(0)
return StreamingResponse(bytes, media_type="image/png")
@app.post("/inpainting/")
async def inpainting(
request: Request,
prompt: str = Form(...),
init_image: UploadFile = File(...),
mask_image: UploadFile = File(...),
):
bytes = await init_image.read()
init_image = Image.open(BytesIO(bytes))
init_image = init_image.convert("RGB").resize((512, 512))
bytes = await mask_image.read()
mask_image = Image.open(BytesIO(bytes))
mask_image = mask_image.convert("RGB").resize((512, 512))
image = request.state.inpaint.pipe(
prompt,
image=init_image,
mask_image=mask_image,
num_inference_steps=3,
strength=0.5,
guidance_scale=0.0,
).images[0]
bytes = BytesIO()
image.save(bytes, "PNG")
bytes.seek(0)
return StreamingResponse(bytes, media_type="image/png")
|