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 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)"}

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

@ -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)",
};