From 33b57aea348e7bb07943093c8b421c569c212654 Mon Sep 17 00:00:00 2001
From: BennyKok <itechbenny@gmail.com>
Date: Tue, 23 Jan 2024 11:27:54 +0800
Subject: [PATCH] fix: display download progress for external checkpoint

---
 comfy-nodes/external_checkpoints.py     | 32 +++++++++++++++++--------
 web/src/components/customInputNodes.tsx | 15 ++++++------
 2 files changed, 29 insertions(+), 18 deletions(-)

diff --git a/comfy-nodes/external_checkpoints.py b/comfy-nodes/external_checkpoints.py
index fc09f8a..b5cc25b 100644
--- a/comfy-nodes/external_checkpoints.py
+++ b/comfy-nodes/external_checkpoints.py
@@ -3,20 +3,20 @@ from PIL import Image, ImageOps
 import numpy as np
 import torch
 import folder_paths
+from tqdm import tqdm
 
-
-class ComfyUIDeployExternalCheckpoints:
+class ComfyUIDeployExternalCheckpoint:
     @classmethod
     def INPUT_TYPES(s):
         return {
             "required": {
                 "input_id": (
                     "STRING",
-                    {"multiline": False, "default": "input_checkpoints"},
+                    {"multiline": False, "default": "input_checkpoint"},
                 ),
             },
             "optional": {
-                "default_checkpoints_name": (folder_paths.get_filename_list("checkpoints"), ),
+                "default_checkpoint_name": (folder_paths.get_filename_list("checkpoints"), ),
             }
         }
 
@@ -27,7 +27,7 @@ class ComfyUIDeployExternalCheckpoints:
 
     CATEGORY = "deploy"
 
-    def run(self, input_id, default_checkpoints_name=None):
+    def run(self, input_id, default_checkpoint_name=None):
         import requests
         import os
         import uuid
@@ -39,18 +39,30 @@ class ComfyUIDeployExternalCheckpoints:
             destination_path = os.path.join(
                 folder_paths.folder_names_and_paths["checkpoints"][0][0], unique_filename)
             print(destination_path)
-            print("Downloading external checkpoints - " +
+            print("Downloading external checkpoint - " +
                   input_id + " to " + destination_path)
             response = requests.get(
-                input_id, headers={'User-Agent': 'Mozilla/5.0'}, allow_redirects=True)
+                input_id, headers={'User-Agent': 'Mozilla/5.0'}, allow_redirects=True, stream=True)
+            file_size = int(response.headers.get('Content-Length', 0))
+            chunk = 1
+            chunk_size = 1024
+            num_bars = int(file_size / chunk_size)
+
             with open(destination_path, 'wb') as out_file:
-                out_file.write(response.content)
+                for chunk in tqdm(
+                    response.iter_content(chunk_size=chunk_size),
+                    total=num_bars,
+                    unit='KB',
+                    desc="Downloading",
+                    leave=True  # leave=True to keep progress bars
+                ):
+                    out_file.write(chunk)
             return (unique_filename,)
         else:
             return (default_checkpoints_name,)
 
 
 NODE_CLASS_MAPPINGS = {
-    "ComfyUIDeployExternalCheckpoints": ComfyUIDeployExternalCheckpoints}
+    "ComfyUIDeployExternalCheckpoint": ComfyUIDeployExternalCheckpoint}
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "ComfyUIDeployExternalCheckpoints": "External Checkpoints (ComfyUI Deploy)"}
+    "ComfyUIDeployExternalCheckpoint": "External Checkpoint (ComfyUI Deploy)"}
diff --git a/web/src/components/customInputNodes.tsx b/web/src/components/customInputNodes.tsx
index 2825169..d6f9d68 100644
--- a/web/src/components/customInputNodes.tsx
+++ b/web/src/components/customInputNodes.tsx
@@ -1,10 +1,9 @@
 export const customInputNodes: Record<string, string> = {
-	ComfyUIDeployExternalText: "string",
-	ComfyUIDeployExternalImage: "string - (public image url)",
-	ComfyUIDeployExternalImageAlpha: "string - (public image url)",
-	ComfyUIDeployExternalNumber: "float",
-	ComfyUIDeployExternalNumberInt: "integer",
-	ComfyUIDeployExternalLora: "string - (public lora download url)",
-	ComfyUIDeployExternalCheckpoints:
-		"string - (public checkpoints download url)",
+  ComfyUIDeployExternalText: "string",
+  ComfyUIDeployExternalImage: "string - (public image url)",
+  ComfyUIDeployExternalImageAlpha: "string - (public image url)",
+  ComfyUIDeployExternalNumber: "float",
+  ComfyUIDeployExternalNumberInt: "integer",
+  ComfyUIDeployExternalLora: "string - (public lora download url)",
+  ComfyUIDeployExternalCheckpoint: "string - (public checkpoints download url)",
 };