69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
import folder_paths
|
|
from PIL import Image, ImageOps
|
|
import numpy as np
|
|
import torch
|
|
import folder_paths
|
|
from tqdm import tqdm
|
|
|
|
class ComfyUIDeployExternalCheckpoint:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"input_id": (
|
|
"STRING",
|
|
{"multiline": False, "default": "input_checkpoint"},
|
|
),
|
|
},
|
|
"optional": {
|
|
"default_checkpoint_name": (folder_paths.get_filename_list("checkpoints"), ),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (folder_paths.get_filename_list("checkpoints"),)
|
|
RETURN_NAMES = ("path",)
|
|
|
|
FUNCTION = "run"
|
|
|
|
CATEGORY = "deploy"
|
|
|
|
def run(self, input_id, default_checkpoint_name=None):
|
|
import requests
|
|
import os
|
|
import uuid
|
|
|
|
if input_id and input_id.startswith('http'):
|
|
unique_filename = str(uuid.uuid4()) + ".safetensors"
|
|
print(unique_filename)
|
|
print(folder_paths.folder_names_and_paths["checkpoints"][0][0])
|
|
destination_path = os.path.join(
|
|
folder_paths.folder_names_and_paths["checkpoints"][0][0], unique_filename)
|
|
print(destination_path)
|
|
print("Downloading external checkpoint - " +
|
|
input_id + " to " + destination_path)
|
|
response = requests.get(
|
|
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:
|
|
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 = {
|
|
"ComfyUIDeployExternalCheckpoint": ComfyUIDeployExternalCheckpoint}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ComfyUIDeployExternalCheckpoint": "External Checkpoint (ComfyUI Deploy)"}
|