File size: 4,085 Bytes
0d47fa5 7dc6a72 0d47fa5 7dc6a72 8eae2b8 7dc6a72 a0a2ed9 0d47fa5 7dc6a72 a0a2ed9 7dc6a72 a0a2ed9 0d47fa5 7dc6a72 a0a2ed9 0d47fa5 a0a2ed9 7dc6a72 |
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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
import sys
import os
import asyncio
import base64
import logging
from io import BytesIO
from pathlib import Path
import uvicorn
from config import Config
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from PIL import Image
from pydantic import BaseModel
sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
from utils.wrapper import StreamDiffusionWrapper
logger = logging.getLogger("uvicorn")
PROJECT_DIR = Path(__file__).parent.parent
class PredictInputModel(BaseModel):
"""
The input model for the /predict endpoint.
"""
prompt: str
class PredictResponseModel(BaseModel):
"""
The response model for the /predict endpoint.
"""
base64_image: str
class UpdatePromptResponseModel(BaseModel):
"""
The response model for the /update_prompt endpoint.
"""
prompt: str
class Api:
def __init__(self, config: Config) -> None:
"""
Initialize the API.
Parameters
----------
config : Config
The configuration.
"""
self.config = config
self.stream_diffusion = StreamDiffusionWrapper(
mode=config.mode,
model_id_or_path=config.model_id_or_path,
lcm_lora_id=config.lcm_lora_id,
vae_id=config.vae_id,
device=config.device,
dtype=config.dtype,
acceleration=config.acceleration,
t_index_list=config.t_index_list,
warmup=config.warmup,
use_safety_checker=config.use_safety_checker,
cfg_type="none",
)
self.app = FastAPI()
self.app.add_api_route(
"/api/predict",
self._predict,
methods=["POST"],
response_model=PredictResponseModel,
)
self.app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
self.app.mount(
"/", StaticFiles(directory="../view/build", html=True), name="public"
)
self._predict_lock = asyncio.Lock()
self._update_prompt_lock = asyncio.Lock()
async def _predict(self, inp: PredictInputModel) -> PredictResponseModel:
"""
Predict an image and return.
Parameters
----------
inp : PredictInputModel
The input.
Returns
-------
PredictResponseModel
The prediction result.
"""
async with self._predict_lock:
return PredictResponseModel(
base64_image=self._pil_to_base64(
self.stream_diffusion(prompt=inp.prompt)
)
)
def _pil_to_base64(self, image: Image.Image, format: str = "JPEG") -> bytes:
"""
Convert a PIL image to base64.
Parameters
----------
image : Image.Image
The PIL image.
format : str
The image format, by default "JPEG".
Returns
-------
bytes
The base64 image.
"""
buffered = BytesIO()
image.convert("RGB").save(buffered, format=format)
return base64.b64encode(buffered.getvalue()).decode("ascii")
def _base64_to_pil(self, base64_image: str) -> Image.Image:
"""
Convert a base64 image to PIL.
Parameters
----------
base64_image : str
The base64 image.
Returns
-------
Image.Image
The PIL image.
"""
if "base64," in base64_image:
base64_image = base64_image.split("base64,")[1]
return Image.open(BytesIO(base64.b64decode(base64_image))).convert("RGB")
if __name__ == "__main__":
from config import Config
config = Config()
uvicorn.run(
Api(config).app,
host=config.host,
port=config.port,
workers=config.workers,
)
|