From fce77b905c103d08213428e085442588b655efbd Mon Sep 17 00:00:00 2001
From: BennyKok <itechbenny@gmail.com>
Date: Thu, 11 Jan 2024 23:14:06 +0800
Subject: [PATCH] feat: add support for gif upload

---
 comfy-nodes/external_image.py       | 2 +-
 comfy-nodes/external_image_alpha.py | 2 +-
 custom_routes.py                    | 6 +++++-
 web/src/components/OutputRender.tsx | 9 +++++++--
 4 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/comfy-nodes/external_image.py b/comfy-nodes/external_image.py
index 0f13e21..599a054 100644
--- a/comfy-nodes/external_image.py
+++ b/comfy-nodes/external_image.py
@@ -34,7 +34,7 @@ class ComfyUIDeployExternalImage:
                 print("Fetching image from url: ", input_id)
                 response = requests.get(input_id)
                 image = Image.open(BytesIO(response.content))
-            elif input_id.startswith('data:image/png;base64,'):
+            elif input_id.startswith('data:image/png;base64,') or input_id.startswith('data:image/jpeg;base64,') or input_id.startswith('data:image/jpg;base64,'):
                 import base64
                 from io import BytesIO
                 print("Decoding base64 image")
diff --git a/comfy-nodes/external_image_alpha.py b/comfy-nodes/external_image_alpha.py
index 5b0aba7..c8099ed 100644
--- a/comfy-nodes/external_image_alpha.py
+++ b/comfy-nodes/external_image_alpha.py
@@ -34,7 +34,7 @@ class ComfyUIDeployExternalImageAlpha:
                 print("Fetching image from url: ", input_id)
                 response = requests.get(input_id)
                 image = Image.open(BytesIO(response.content))
-            elif input_id.startswith('data:image/png;base64,'):
+            elif input_id.startswith('data:image/png;base64,') or input_id.startswith('data:image/jpeg;base64,') or input_id.startswith('data:image/jpg;base64,'):
                 import base64
                 from io import BytesIO
                 print("Decoding base64 image")
diff --git a/custom_routes.py b/custom_routes.py
index 187d09f..33ba40a 100644
--- a/custom_routes.py
+++ b/custom_routes.py
@@ -376,7 +376,7 @@ async def update_run_with_output(prompt_id, data, node_id=None):
         }
 
         try:
-            have_upload = 'images' in data or 'files' in data
+            have_upload = 'images' in data or 'files' in data or 'gifs' in data
 
             print("\nhave_upload", have_upload, node_id)
 
@@ -390,6 +390,10 @@ async def update_run_with_output(prompt_id, data, node_id=None):
             files = data.get('files', [])
             for file in files:
                 await upload_file(prompt_id, file.get("filename"), subfolder=file.get("subfolder"), type=file.get("type"), content_type=file.get("content_type", "image/png"))
+
+            gifs = data.get('gifs', [])
+            for gif in gifs:
+                await upload_file(prompt_id, file.get("filename"), subfolder=file.get("subfolder"), type=file.get("type"), content_type=file.get("format", "image/gif"))
                 
             if have_upload:
                 await update_file_status(prompt_id, data, False, node_id=node_id)
diff --git a/web/src/components/OutputRender.tsx b/web/src/components/OutputRender.tsx
index 2977d8f..d88756c 100644
--- a/web/src/components/OutputRender.tsx
+++ b/web/src/components/OutputRender.tsx
@@ -5,7 +5,12 @@ export async function OutputRender(props: {
   run_id: string;
   filename: string;
 }) {
-  if (props.filename.endsWith(".png")) {
+  if (
+    props.filename.endsWith(".png") ||
+    props.filename.endsWith(".gif") ||
+    props.filename.endsWith(".jpg") ||
+    props.filename.endsWith(".jpeg")
+  ) {
     const url = await getFileDownloadUrl(
       `outputs/runs/${props.run_id}/${props.filename}`
     );
@@ -15,7 +20,7 @@ export async function OutputRender(props: {
     const url = await getFileDownloadUrl(
       `outputs/runs/${props.run_id}/${props.filename}`
     );
-    console.log(url);
+    // console.log(url);
 
     return <DownloadButton filename={props.filename} href={url} />;
   }