Image url as input, evnironment variable for with and without tiling

handler.py

@@ -13,16 +13,20 @@ import boto3
 import uuid, io
 import torch
 import base64
+import requests
 
 
 class EndpointHandler:
     def __init__(self, path=""):
+        
+        self.tiling_size = int(os.environ["TILING_SIZE"])
 
         # Initialize the Real-ESRGAN model with specified parameters
         self.model = RealESRGANer(
             scale=4,  # Scale factor for the model
             # Path to the pre-trained model weights
-            model_path=f"/repository/weights/Real-ESRGAN-x4plus.pth",
+            # model_path=f"/repository/weights/Real-ESRGAN-x4plus.pth",
+            model_path=f"/workspace/real-esrgan/weights/Real-ESRGAN-x4plus.pth",
             # Initialize the RRDBNet model architecture with specified parameters
             model= RRDBNet(num_in_ch=3,
                            num_out_ch=3,
@@ -31,7 +35,7 @@ class EndpointHandler:
                            num_grow_ch=32,
                            scale=4
                            ),
-            tile=0,
+            tile=self.tiling_size,
             tile_pad=0,
             half=True,
         )
@@ -56,12 +60,13 @@ class EndpointHandler:
             outscale = float(inputs.pop("outscale", 3))
             
             # decode base64 image to PIL
-            image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
+            image = self.download_image_url(inputs['image_url'])
             in_size, in_mode = image.size, image.mode
             
             # check image size and mode and return dict
             assert in_mode in ["RGB", "RGBA", "L"], f"Unsupported image mode: {in_mode}"
-            assert in_size[0] * in_size[1] <  1400*1400, f"Image is too large: {in_size}: {in_size[0] * in_size[1]} is greater than {1400*1400}"
+            if self.tiling_size == 0:
+                assert in_size[0] * in_size[1] <  1400*1400, f"Image is too large: {in_size}: {in_size[0] * in_size[1]} is greater than {self.tiling_size*self.tiling_size}"
             assert outscale > 1 and outscale <=10, f"Outscale must be between 1 and 10: {outscale}"
             
             # debug
@@ -141,4 +146,11 @@ class EndpointHandler:
         image_url = "https://{0}.s3.amazonaws.com/{1}".format(self.bucket_name, key)
         
         # return the url and the key    
-        return image_url, key
+        return image_url, key
+    
+    def download_image_url(self, image_url):
+        "Download the image from the url and return the image."
+        
+        response = requests.get(image_url)
+        image = Image.open(BytesIO(response.content))
+        return image

test_handler.ipynb

@@ -15,6 +15,7 @@
     }
    ],
    "source": [
+    "import os\n",
     "from handler import EndpointHandler\n",
     "import base64\n",
     "from io import BytesIO\n",
@@ -28,6 +29,18 @@
    "execution_count": 2,
    "metadata": {},
    "outputs": [],
+   "source": [
+    "os.environ[\"AWS_ACCESS_KEY_ID\"] = \"\"\n",
+    "os.environ[\"AWS_SECRET_ACCESS_KEY\"] = \"\"\n",
+    "os.environ[\"S3_BUCKET_NAME\"] = \"\"\n",
+    "os.environ[\"TILING_SIZE\"] = \"1000\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "# init handler\n",
     "my_handler = EndpointHandler(path=\".\")"
@@ -42,45 +55,48 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "image.size: (1200, 517), image.mode: RGBA, outscale: 10.0\n"
+      "image.size: (1024, 1024), image.mode: RGB, outscale: 4.0\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "output.shape: (5170, 12000, 4)\n",
-      "https://upscale-process-results.s3.amazonaws.com/0da01291-5a39-40fd-b322-f5d83e6066f6.png 0da01291-5a39-40fd-b322-f5d83e6066f6.png\n",
-      "image.size: (1056, 1068), image.mode: RGB, outscale: 3.0\n",
-      "output.shape: (3204, 3168, 3)\n",
-      "https://upscale-process-results.s3.amazonaws.com/c1ba714d-50e2-45d0-ac9f-8a4a4f218320.png c1ba714d-50e2-45d0-ac9f-8a4a4f218320.png\n",
-      "image.size: (1056, 1068), image.mode: L, outscale: 5.49\n",
-      "output.shape: (5863, 5797, 3)\n",
-      "https://upscale-process-results.s3.amazonaws.com/e964c021-9169-49d8-9382-104f704a1d92.png e964c021-9169-49d8-9382-104f704a1d92.png\n"
+      "\tTile 1/4\n",
+      "\tTile 2/4\n",
+      "\tTile 3/4\n",
+      "\tTile 4/4\n",
+      "output.shape: (4096, 4096, 3)\n",
+      "https://jiffy-staging-upscaled-images.s3.amazonaws.com/d91323cb-0801-45b7-8109-9739212037ed.png d91323cb-0801-45b7-8109-9739212037ed.png\n"
      ]
     }
    ],
    "source": [
     "img_dir = \"test_data/\"\n",
-    "img_names = [\"4121783.png\", \"FB_IMG_1725931665635.jpg\", \"FB_IMG_1725931665635_gray.jpg\"]\n",
-    "out_scales = [10, 3, 5.49]\n",
-    "for img_name, outscale in zip(img_names, out_scales):\n",
-    "    image_path = img_dir + img_name\n",
+    "img_urls = [\"https://jiffy-transfers.imgix.net/2/attachments/r267odvvfmkp6c5lccj1y6f9trb0?ixlib=rb-0.3.5\",\n",
+    "            # \"https://jiffy-staging-transfers.imgix.net/2/development/attachments/zo31eau0ykhbwoddrjtlbyz6w9mp?ixlib=rb-0.3.5\", # larger than > 1.96M pixels\n",
+    "            # \"https://jiffy-staging-transfers.imgix.net/2/development/attachments/b8ecchms9rr9wk3g71kfpfprqg1v?ixlib=rb-0.3.5\" # larger than > 1.96M pixels\n",
+    "           ]\n",
+    "\n",
+    "out_scales = [4, 3, 2]\n",
+    "for img_url, outscale in zip(img_urls, out_scales):\n",
     "    # create payload\n",
-    "    with open(image_path, \"rb\") as i:\n",
-    "        b64 = base64.b64encode(i.read())\n",
-    "        b64 = b64.decode(\"utf-8\")\n",
-    "        payload = {\n",
-    "            \"inputs\": {\"image\": b64, \n",
+    "    payload = {\n",
+    "        \"inputs\": {\"image_url\": img_url, \n",
     "                    \"outscale\": outscale\n",
     "                    }\n",
     "            }\n",
-    "\n",
-    "\n",
+    "    \n",
     "    output_payload = my_handler(payload)\n",
-    "    print(output_payload[\"image_url\"], output_payload[\"image_key\"])\n",
-    "    "
+    "    print(output_payload[\"image_url\"], output_payload[\"image_key\"])\n"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {