chunzhi/comfyui-deploy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: display download progress for external checkpoint

Browse Source
This commit is contained in:
BennyKok 2024-01-23 11:27:54 +08:00
parent 71dcbe0d11
commit 33b57aea34
2 changed files with 29 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
comfy-nodes/external_checkpoints.py
View File

@ -3,20 +3,20 @@ from PIL import Image, ImageOps
import numpy as np import numpy as np
import torch import torch
import folder_paths import folder_paths
from tqdm import tqdm
class ComfyUIDeployExternalCheckpoint:
class ComfyUIDeployExternalCheckpoints:
@classmethod @classmethod
def INPUT_TYPES(s): def INPUT_TYPES(s):
return { return {
"required": { "required": {
"input_id": ( "input_id": (
"STRING", "STRING",
{"multiline": False, "default": "input_checkpoints"}, {"multiline": False, "default": "input_checkpoint"},
), ),
}, },
"optional": { "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" CATEGORY = "deploy"
def run(self, input_id, default_checkpoints_name=None): def run(self, input_id, default_checkpoint_name=None):
import requests import requests
import os import os
import uuid import uuid
@ -39,18 +39,30 @@ class ComfyUIDeployExternalCheckpoints:
destination_path = os.path.join( destination_path = os.path.join(
folder_paths.folder_names_and_paths["checkpoints"][0][0], unique_filename) folder_paths.folder_names_and_paths["checkpoints"][0][0], unique_filename)
print(destination_path) print(destination_path)
print("Downloading external checkpoints - " + print("Downloading external checkpoint - " +
input_id + " to " + destination_path) input_id + " to " + destination_path)
response = requests.get( 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: 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,) return (unique_filename,)
else: else:
return (default_checkpoints_name,) return (default_checkpoints_name,)
NODE_CLASS_MAPPINGS = { NODE_CLASS_MAPPINGS = {
"ComfyUIDeployExternalCheckpoints": ComfyUIDeployExternalCheckpoints} "ComfyUIDeployExternalCheckpoint": ComfyUIDeployExternalCheckpoint}
NODE_DISPLAY_NAME_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = {
"ComfyUIDeployExternalCheckpoints": "External Checkpoints (ComfyUI Deploy)"} "ComfyUIDeployExternalCheckpoint": "External Checkpoint (ComfyUI Deploy)"}

3
web/src/components/customInputNodes.tsx
View File

@ -5,6 +5,5 @@ export const customInputNodes: Record<string, string> = {
ComfyUIDeployExternalNumber: "float", ComfyUIDeployExternalNumber: "float",
ComfyUIDeployExternalNumberInt: "integer", ComfyUIDeployExternalNumberInt: "integer",
ComfyUIDeployExternalLora: "string - (public lora download url)", ComfyUIDeployExternalLora: "string - (public lora download url)",
ComfyUIDeployExternalCheckpoints: ComfyUIDeployExternalCheckpoint: "string - (public checkpoints download url)",
"string - (public checkpoints download url)",
}; };