#!/usr/bin/env python from fastapi import FastAPI, WebSocket, WebSocketDisconnect from fastapi.middleware.cors import CORSMiddleware from fastapi.logger import logger # General from config import CONFIG from pydantic import BaseModel from PIL import Image # Connection Manager from connectionManager import ConnectionManager # CutLER Model from model import Model import base64 from io import BytesIO from predict import predict # Stable Diffusion Inpainting Model # from diffusers import StableDiffusionInpaintPipeline # About import torch import os import sys # Server API import uvicorn app = FastAPI( title="AdVisual Model Hosting", description="Description of the ML Model", version="0.0.1", terms_of_service=None, contact=None, license_info=None, docs_url="/", ) # Allow CORS for local debugging if CONFIG['ENV'] == 'development': app.add_middleware(CORSMiddleware, allow_origins=["*"]) else: app.add_middleware(CORSMiddleware, allow_origins=["https://advisual.io"]) @app.on_event("startup") async def startup_event(): """ Initialize FastAPI and add variables """ logger.info('Running envirnoment: {}'.format(CONFIG['ENV'])) logger.info('PyTorch using device: {}'.format(CONFIG['DEVICE'])) # Initialize the CutLER model model = Model(CONFIG['DEVICE']) # Initialize the stable-diffusion-inpainting model # pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting", safety_checker=None) # pipe.to(CONFIG['DEVICE']) # Initialize the connection manager connectionManager = ConnectionManager() # add model and other preprocess tools too app state app.package = { "model": model, "connectionManager": connectionManager, # "pipe": pipe } @app.get("/ping") def ping(): return {"ok": True, "message": "Pong"} @app.get("/about") def show_about(): """ Get deployment information, for debugging """ logger.info('API /about called') def bash(command): output = os.popen(command).read() return output return { "sys.version": sys.version, "torch.__version__": torch.__version__, "torch.cuda.is_available()": torch.cuda.is_available(), "torch.version.cuda": torch.version.cuda, "torch.backends.cudnn.version()": torch.backends.cudnn.version(), "torch.backends.cudnn.enabled": torch.backends.cudnn.enabled, "nvidia-smi": bash('nvidia-smi') } # def resize_image(img, height=512, width=512): # '''Resize image to `size`''' # size = (width, height) # img_resized = img.resize(size, Image.ANTIALIAS) # return img_resized # def crop_image(img, d=64): # '''Make dimensions divisible by `d`''' # new_size = (img.size[0] - img.size[0] % d, # img.size[1] - img.size[1] % d) # bbox = [ # int((img.size[0] - new_size[0])/2), # int((img.size[1] - new_size[1])/2), # int((img.size[0] + new_size[0])/2), # int((img.size[1] + new_size[1])/2), # ] # img_cropped = img.crop(bbox) # return img_cropped # class InpaintBody(BaseModel): # image: str # mask: str # prompt: str # @app.post("/inpaint") # async def do_inpaint(body: InpaintBody): # """ # Perform inpainting on input data # """ # logger.info('API inpaint called') # image_data = body.image # mask_data = body.mask # prompt = body.prompt # # Extract base64 from mask and convert to PIL.Image # if (',' in image_data): # image = Image.open(BytesIO(base64.b64decode(image_data.split(',')[1]))) # else: # image = Image.open(BytesIO(base64.b64decode(image_data))) # # Extract base64 from mask and convert to PIL.Image # if (',' in mask_data): # mask = Image.open(BytesIO(base64.b64decode(mask_data.split(',')[1]))) # else: # mask = Image.open(BytesIO(base64.b64decode(mask_data))) # # Resize image and mask to 512x512 # image = crop_image(resize_image(image, 512, 512)) # mask = crop_image(resize_image(image, 512, 512)) # pipe = app.package.get('pipe') # result = pipe(prompt=prompt, image=image, mask_image=mask, num_inference_steps=10, num_images_per_prompt=1) # images = result['images'] # return images class ImageBody(BaseModel): image: str threshold: float = 0.15 num_objects: int = 1 @app.post("/predict") async def do_predict(body: ImageBody): """ Perform prediction on input data """ logger.info('API predict called') image: str = body.image threshold: float = body.threshold num_objects: int = body.num_objects # Run the algorithm result = predict(app.package, image, threshold, num_objects) # Convert the result to base64 and send the json back buffered = BytesIO() result.save(buffered, format="JPEG") img_str = 'data:image/jpeg;base64,' + base64.b64encode(buffered.getvalue()).decode("utf-8") return {"ok": True, "status": "FINISHED", "result": img_str} # @app.websocket("/ws-inpaint") # async def inpaint_websocket_endpoint(websocket: WebSocket): # connectionManager = app.package.get('connectionManager') # await connectionManager.connect(websocket) # await connectionManager.send_json({"ok": True, "status": "CONNECTED"}, websocket) # while True: # try: # data: ImageBody = await connectionManager.receive_json(websocket) # if (data is None): # # Wait for data # if not connectionManager.isConnected(websocket): # break # if connectionManager.shouldDisconnect(websocket): # await websocket.close() # connectionManager.disconnect(websocket) # break # continue # image_data: str = data.get('image') # mask_data: str = data.get('mask') # prompt: str = data.get('prompt') # await connectionManager.send_json({"ok": True, "status": "STARTED"}, websocket) # # Extract base64 from mask and convert to PIL.Image # if (',' in image_data): # image = Image.open(BytesIO(base64.b64decode(image_data.split(',')[1]))) # else: # image = Image.open(BytesIO(base64.b64decode(image_data))) # # Extract base64 from mask and convert to PIL.Image # if (',' in mask_data): # mask = Image.open(BytesIO(base64.b64decode(mask_data.split(',')[1]))) # else: # mask = Image.open(BytesIO(base64.b64decode(mask_data))) # # Resize image and mask to 512x512 # image = crop_image(resize_image(image, 512, 512)) # mask = crop_image(resize_image(image, 512, 512)) # pipe = app.package.get('pipe') # result = pipe(prompt=prompt, image=image, mask_image=mask, num_inference_steps=20, num_images_per_prompt=1) # images = result['images'] # # Convert the result to base64 and send the json back # result_array = [] # for image in images: # buffered = BytesIO() # image.save(buffered, format="JPEG") # img_str = 'data:image/jpeg;base64,' + base64.b64encode(buffered.getvalue()).decode("utf-8") # result_array.append(img_str) # await connectionManager.send_json({"ok": True, "status": "FINISHED", "result": result_array}, websocket) # await websocket.close() # connectionManager.disconnect(websocket) # except WebSocketDisconnect: # connectionManager.disconnect(websocket) # break @app.websocket("/ws") async def websocket_endpoint(websocket: WebSocket): connectionManager = app.package.get('connectionManager') await connectionManager.connect(websocket) await connectionManager.send_json({"ok": True, "status": "CONNECTED"}, websocket) while True: try: data: ImageBody = await connectionManager.receive_json(websocket) if (data is None): # Wait for data if not connectionManager.isConnected(websocket): break if connectionManager.shouldDisconnect(websocket): await websocket.close() connectionManager.disconnect(websocket) break continue image: str = data.get('image') threshold: float = data.get('threshold') or 0.15 num_objects: int = data.get('num_objects') or 1 await connectionManager.send_json({"ok": True, "status": "STARTED"}, websocket) # Run the algorithm result = predict(app.package, image, threshold, num_objects) # Convert the result to base64 and send the json back buffered = BytesIO() result.save(buffered, format="JPEG") img_str = 'data:image/jpeg;base64,' + base64.b64encode(buffered.getvalue()).decode("utf-8") await connectionManager.send_json({"ok": True, "status": "FINISHED", "result": img_str}, websocket) await websocket.close() connectionManager.disconnect(websocket) except WebSocketDisconnect: connectionManager.disconnect(websocket) break if __name__ == '__main__': # server api uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)