Removed Inpaint

app.py

@@ -58,9 +58,9 @@ async def startup_event():
     model = Model(CONFIG['DEVICE'])
 
     # Initialize the stable-diffusion-inpainting model
-    pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting", safety_checker=None)
+    # pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting", safety_checker=None)
 
-    pipe.to(CONFIG['DEVICE'])
+    # pipe.to(CONFIG['DEVICE'])
 
     # Initialize the connection manager
     connectionManager = ConnectionManager()
@@ -69,7 +69,7 @@ async def startup_event():
     app.package = {
         "model": model,
         "connectionManager": connectionManager,
-        "pipe": pipe
+        # "pipe": pipe
     }
 
 @app.get("/ping")
@@ -98,66 +98,66 @@ def show_about():
         "nvidia-smi": bash('nvidia-smi')
     }
 
-def resize_image(img, height=512, width=512):
-    '''Resize image to `size`'''
+# def resize_image(img, height=512, width=512):
+#     '''Resize image to `size`'''
 
-    size = (width, height)
+#     size = (width, height)
 
-    img_resized = img.resize(size, Image.ANTIALIAS)
-    return img_resized
+#     img_resized = img.resize(size, Image.ANTIALIAS)
+#     return img_resized
 
-def crop_image(img, d=64):
-    '''Make dimensions divisible by `d`'''
+# 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)
+#     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),
-    ]
+#     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
+#     img_cropped = img.crop(bbox)
+#     return img_cropped
 
-class InpaintBody(BaseModel):
-    image: str
-    mask: str
-    prompt: str
+# class InpaintBody(BaseModel):
+#     image: str
+#     mask: str
+#     prompt: str
 
-@app.post("/inpaint")
-async def do_inpaint(body: InpaintBody):
-    """
-    Perform inpainting on input data
-    """
+# @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
+#     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
+#     # 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):
@@ -188,65 +188,65 @@ 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
+# @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')
+#           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)
+#           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 
+#           # 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):

requirements.txt

@@ -6,9 +6,11 @@ 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
+
+# Stable Diffusion Inpaint
+# diffusers==0.14.0
+# transformers==4.27.1
+# accelerate==0.17.1
 
 pycocotools==2.0.6
 fastapi==0.94.1