File size: 3,068 Bytes
672243b
bd94aa2
 
 
 
 
 
 
 
 
 
 
 
 
 
3fe0b34
d56c8ea
 
bd94aa2
 
 
4fd0935
bd94aa2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
845eacd
8bb6490
 
 
845eacd
bd94aa2
845eacd
bd94aa2
 
 
 
672243b
760e14f
bd94aa2
 
 
 
845eacd
16500ab
25ff54e
16500ab
8bb6490
bd94aa2
25ff54e
8bb6490
bd94aa2
 
4fd0935
bd94aa2
 
845eacd
 
bd94aa2
 
 
 
 
672243b
760e14f
bd94aa2
 
 
 
 
16500ab
25ff54e
 
8bb6490
 
bd94aa2
25ff54e
8bb6490
bd94aa2
25ff54e
8bb6490
bd94aa2
 
4fd0935
bd94aa2
 
 
845eacd
 
bd94aa2
 
 
 
 
672243b
760e14f
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
113
114
115
116
117
118
119
120
from fastapi import FastAPI, Request, Response, Form, File, UploadFile
from contextlib import asynccontextmanager
from starlette.middleware.cors import CORSMiddleware

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/sd-turbo").to(
        "cpu"
    )

    img2img = AutoPipelineForImage2Image.from_pipe(text2img).to("cpu")

    inpaint = AutoPipelineForInpainting.from_pipe(img2img).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(...),
    num_inference_steps: int = Form(1),
):
    image = request.state.text2img(
        prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=0.0
    ).images[0]

    bytes = BytesIO()
    image.save(bytes, "PNG")
    # bytes.seek(0)
    return Response(content=bytes, media_type="image/png")


@app.post("/image-to-image/")
async def image_to_image(
    request: Request,
    prompt: str = Form(...),
    init_image: UploadFile = File(...),
    num_inference_steps: int = Form(2),
    strength: float = Form(1.0),
):
    init_bytes = await init_image.read()
    init_image = Image.open(BytesIO(init_bytes))
    init_image = init_image.convert("RGB").resize((512, 512))

    image = request.state.img2img(
        prompt,
        image=init_image,
        num_inference_steps=num_inference_steps,
        strength=strength,
        guidance_scale=0.0,
    ).images[0]

    bytes = BytesIO()
    image.save(bytes, "PNG")
    # bytes.seek(0)
    return Response(content=bytes, media_type="image/png")


@app.post("/inpainting/")
async def inpainting(
    request: Request,
    prompt: str = Form(...),
    init_image: UploadFile = File(...),
    mask_image: UploadFile = File(...),
    num_inference_steps: int = Form(2),
    strength: float = Form(1.0),
):
    init_bytes = await init_image.read()
    init_image = Image.open(BytesIO(init_bytes))
    init_image = init_image.convert("RGB").resize((512, 512))
    mask_bytes = await mask_image.read()
    mask_image = Image.open(BytesIO(mask_bytes))
    mask_image = mask_image.convert("RGB").resize((512, 512))

    image = request.state.inpaint(
        prompt,
        image=init_image,
        mask_image=mask_image,
        num_inference_steps=num_inference_steps,
        strength=strength,
        guidance_scale=0.0,
    ).images[0]

    bytes = BytesIO()
    image.save(bytes, "PNG")
    # bytes.seek(0)
    return Response(content=bytes, media_type="image/png")