Create main.py

main.py

@@ -0,0 +1,162 @@
+import asyncio
+import base64
+import logging
+import os
+import sys
+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,
+            lora_dict=config.lora_dict,
+            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="./frontend/dist", 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,
+    )