Added stable diffusion inpaint endpoint

app.py

@@ -4,18 +4,23 @@ 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
 
-# Model
+# CutLER Model
 from model import Model
 import base64
 from io import BytesIO
-from pydantic import BaseModel
-from config import CONFIG
-
 from predict import predict
 
+# Stable Diffusion Inpainting Model
+from diffusers import StableDiffusionInpaintPipeline
+
 # About
 import torch
 import os
@@ -25,7 +30,7 @@ import sys
 import uvicorn
 
 app = FastAPI(
-    title="AdVisual MaskCut Model",
+    title="AdVisual Model Hosting",
     description="Description of the ML Model",
     version="0.0.1",
     terms_of_service=None,
@@ -35,7 +40,10 @@ app = FastAPI(
 )
 
 # Allow CORS for local debugging
-app.add_middleware(CORSMiddleware, allow_origins=["*"])
+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():
@@ -46,14 +54,22 @@ async def startup_event():
     logger.info('Running envirnoment: {}'.format(CONFIG['ENV']))
     logger.info('PyTorch using device: {}'.format(CONFIG['DEVICE']))
 
-    # Initialize the pytorch model
-    model = Model()
+    # 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
+        "connectionManager": connectionManager,
+        "pipe": pipe
     }
 
 @app.get("/ping")
@@ -82,6 +98,68 @@ def show_about():
         "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
@@ -110,6 +188,66 @@ async def do_predict(body: ImageBody):
     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=10, 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')
@@ -117,7 +255,7 @@ async def websocket_endpoint(websocket: WebSocket):
     await connectionManager.send_json({"ok": True, "status": "CONNECTED"}, websocket)
     while True:
         try: 
-          data = await connectionManager.receive_json(websocket)
+          data: ImageBody = await connectionManager.receive_json(websocket)
           if (data is None):
               # Wait for data
               if not connectionManager.isConnected(websocket):
@@ -127,6 +265,7 @@ async def websocket_endpoint(websocket: WebSocket):
                 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
@@ -146,7 +285,6 @@ async def websocket_endpoint(websocket: WebSocket):
           await websocket.close()
           connectionManager.disconnect(websocket)
         except WebSocketDisconnect:
-          print('websocket disconnect')
           connectionManager.disconnect(websocket)
           break 
 

config.py

@@ -8,8 +8,6 @@ import torch
 
 # Config that serves all environment
 GLOBAL_CONFIG = {
-    "MODEL_PATH": "../model/model.pt",
-    "SCALAR_PATH": "../model/scaler.joblib",
     "USE_CUDE_IF_AVAILABLE": True,
     "ROUND_DIGIT": 6
 }

connectionManager.py

@@ -1,8 +1,7 @@
 from fastapi import WebSocket
 
 from datetime import datetime
-import json
-from typing import Dict, List
+from typing import List
 
 class Connection: 
     websocket: WebSocket
@@ -19,6 +18,7 @@ class ConnectionManager:
         self.active_connections: List[Connection] = []
 
     async def connect(self, websocket: WebSocket):
+        print('Connecting')
         await websocket.accept()
         # Add connection time and websocket to active connections
         self.active_connections.append(Connection(websocket=websocket, connection_time=datetime.now()))
@@ -33,16 +33,20 @@ class ConnectionManager:
         for connection in self.active_connections:
             if connection.websocket == websocket:
                 if (datetime.now() - connection.connection_time).total_seconds() > self.timeout:
+                    print('Disconnecting...')
                     return True
         return False
 
     async def receive_json(self, websocket: WebSocket):
         if not self.isConnected(websocket):
           return None
+        print('Receiving...')
         data = await websocket.receive_json()
+        print('Received')
         return data
 
     def disconnect(self, websocket: WebSocket):
+        print('Disconnecting...')
         for connection in self.active_connections:
             if connection.websocket == websocket:
                 self.active_connections.remove(connection)
@@ -50,6 +54,7 @@ class ConnectionManager:
         return False
 
     async def send_json(self, json, websocket: WebSocket):
+        print('Sending JSON...')
         # Only send the message if the connection is still active
         if self.isConnected(websocket):
           await websocket.send_json(json)

model.py

@@ -18,9 +18,8 @@ from CutLER.third_party.TokenCut.unsupervised_saliency_detection import metric
 
 
 class Model:
-    def __init__(self):
-        self.device = torch.device(
-            'cuda:0' if torch.cuda.is_available() else 'cpu')
+    def __init__(self, device: str):
+        self.device = torch.device(device)
         self.backbone = self.load_backbone()
 
     def load_backbone(self):

requirements.txt

@@ -6,6 +6,9 @@ scikit-image==0.19.2
 torch==1.13.1
 torchvision==0.14.1
 tqdm==4.64.1
+diffusers==0.14.0
+transformers==4.27.1
+accelerate==0.17.1
 
 pycocotools==2.0.6
 fastapi==0.94.1