Skip to content

Commit

Permalink
Update.
Browse files Browse the repository at this point in the history
  • Loading branch information
robinjhuang committed Jun 26, 2024
1 parent 66d753b commit 1b68e8e
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 13 deletions.
4 changes: 2 additions & 2 deletions action.yml
Original file line number Diff line number Diff line change
Expand Up @@ -94,13 +94,13 @@ runs:
if: ${{ inputs.os != 'windows' }}
shell: bash -el {0}
run: |
pip install -r requirements.txt
pip3 install -r requirements.txt
- name: '[Unix] Install Comfy-CLI'
if: ${{ inputs.os != 'windows' }}
shell: bash -el {0}
run: |
pip install comfy-cli
pip3 install comfy-cli
comfy --version
comfy --skip-prompt --no-enable-telemetry env
Expand Down
36 changes: 25 additions & 11 deletions download-models.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,14 @@
import base64
import argparse


def ensure_directory_exists(path):
"""Ensure that a directory exists; if it doesn't, create it."""
if not os.path.exists(path):
os.makedirs(path)
print("Created directory: " + path)


def download_model(url, directory, model_name):
"""Download a model file from a URL into a specific directory.
Skip download if the file already exists."""
Expand All @@ -19,23 +21,24 @@ def download_model(url, directory, model_name):
if os.path.exists(file_path):
print(f"{model_name} already exists in {directory}. Skipping download.")
return

print(f"Downloading {model_name} from {url} to {directory}")
with requests.get(url, stream=True) as r:
r.raise_for_status()
with open(file_path, 'wb') as f:
for chunk in r.iter_content(chunk_size=8192):
with open(file_path, "wb") as f:
for chunk in r.iter_content(chunk_size=8192):
f.write(chunk)
print(f"Downloaded {model_name} to {directory}")


def main(args):
"""Main function to parse JSON and download models based on the input
mode."""
if args.mode == 'base64':
if args.mode == "base64":
models_json = base64.b64decode(args.input).decode("utf-8")
elif args.mode == 'raw':
elif args.mode == "raw":
models_json = args.input

print("Processing input:")
print(models_json)
print("Base directory: " + args.directory)
Expand All @@ -48,11 +51,22 @@ def main(args):
ensure_directory_exists(target_directory)
download_model(url, target_directory, model_name)


if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Download models based on a JSON input.")
parser.add_argument('mode', choices=['base64', 'raw'], help="Input mode: 'base64' for a base64 encoded string, 'raw' for a JSON content itself.")
parser.add_argument('input', help="Input string or json content, depending on the mode.")
parser.add_argument('directory', help="Base directory where models will be downloaded.")

parser = argparse.ArgumentParser(
description="Download models based on a JSON input."
)
parser.add_argument(
"mode",
choices=["base64", "raw"],
help="Input mode: 'base64' for a base64 encoded string, 'raw' for a JSON content itself.",
)
parser.add_argument(
"input", help="Input string or json content, depending on the mode."
)
parser.add_argument(
"directory", help="Base directory where models will be downloaded."
)

args = parser.parse_args()
main(args)

0 comments on commit 1b68e8e

Please sign in to comment.