From 4b76774a0468eb1e4b4aec0e1d39ed34bd67d2f1 Mon Sep 17 00:00:00 2001 From: "Jason.Stone" Date: Mon, 28 Apr 2025 09:16:37 +0000 Subject: [PATCH 1/4] Enhance model setup process: add TensorRT engine building and improve download handling --- .devcontainer/post-create.sh | 2 +- configs/models.yaml | 60 ++++- docker/entrypoint.sh | 48 +--- src/comfystream/scripts/setup_models.py | 326 +++++++++++++++++++----- 4 files changed, 322 insertions(+), 114 deletions(-) diff --git a/.devcontainer/post-create.sh b/.devcontainer/post-create.sh index 0afbc04e..9e72b404 100755 --- a/.devcontainer/post-create.sh +++ b/.devcontainer/post-create.sh @@ -22,7 +22,7 @@ if [ ! -d "/workspace/comfystream/nodes/web/static" ]; then fi # Create a symlink to the entrypoint script. -echo 'alias prepare_examples="/workspace/comfystream/docker/entrypoint.sh --download-models --build-engines"' >> ~/.bashrc +echo 'alias prepare_examples="/workspace/comfystream/docker/entrypoint.sh --download-models"' >> ~/.bashrc echo -e "\e[32mContainer ready! Run 'prepare_examples' to download models and build engines for example workflows.\e[0m" cd /workspace/comfystream diff --git a/configs/models.yaml b/configs/models.yaml index 4755789d..bdc1894f 100644 --- a/configs/models.yaml +++ b/configs/models.yaml @@ -15,16 +15,74 @@ models: extra_files: - url: "https://huggingface.co/aaronb/dreamshaper-8-dmd-1kstep/raw/main/config.json" path: "unet/dreamshaper-8-dmd-1kstep.json" - + tensorrt: + build: true + ## These engine files needs to be on the output directory for the TensorRT Loader to find them +## ------------------------------------------------------------------------------------------------------------------------------ ## + engines: + - script: "src/comfystream/scripts/build_trt.py" + engine_path: "output/tensorrt/static-dreamshaper8_SD15_$stat-b-1-h-512-w-512_00001_.engine" + args: + - "--model" + - "{model_path}" + - "--out-engine" + - "{engine_path}" + - script: "src/comfystream/scripts/build_trt.py" + engine_path: "output/tensorrt/dynamic-dreamshaper8_SD15_$dyn-b-1-4-2-h-448-704-512-w-448-704-512_00001_.engine" + args: + - "--model" + - "{model_path}" + - "--out-engine" + - "{engine_path}" + - "--width" + - "512" + - "--height" + - "512" + - "--min-width" + - "448" + - "--min-height" + - "448" + - "--max-width" + - "704" + - "--max-height" + - "704" +## ------------------------------------------------------------------------------------------------------------------------------ ## # Depth Anything V2 ONNX models depthanything-onnx: name: "DepthAnything ONNX" url: "https://huggingface.co/yuvraj108c/Depth-Anything-2-Onnx/resolve/main/depth_anything_v2_vitb.onnx?download=true" path: "tensorrt/depth-anything/depth_anything_vitl14.onnx" + tensorrt: + build: true + engines: + - script: "custom_nodes/ComfyUI-Depth-Anything-Tensorrt/export_trt.py" + engine_path: "tensorrt/depth-anything/depth_anything_vitl14-fp16.engine" + cwd: "custom_nodes/ComfyUI-Depth-Anything-Tensorrt" + args: + - "--trt-path" + - "{engine_path}" + - "--onnx-path" + - "{model_path}" + - "--use-fp32" + - "false" + depth-anything-v2-large-onnx: name: "DepthAnything V2 Large ONNX" url: "https://huggingface.co/yuvraj108c/Depth-Anything-2-Onnx/resolve/main/depth_anything_v2_vitl.onnx?download=true" path: "tensorrt/depth-anything/depth_anything_v2_vitl.onnx" + tensorrt: + build: true + engines: + - script: "custom_nodes/ComfyUI-Depth-Anything-Tensorrt/export_trt.py" + engine_path: "tensorrt/depth-anything/depth_anything_v2_vitl-fp16.engine" + cwd: "custom_nodes/ComfyUI-Depth-Anything-Tensorrt" + args: + - "--trt-path" + - "{engine_path}" + - "--onnx-path" + - "{model_path}" + - "--use-fp32" + - "false" # TAESD models taesd: diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index 009f4bda..6365f3ea 100755 --- a/docker/entrypoint.sh +++ b/docker/entrypoint.sh @@ -21,8 +21,7 @@ show_help() { echo "Usage: entrypoint.sh [OPTIONS]" echo "" echo "Options:" - echo " --download-models Download default models" - echo " --build-engines Build TensorRT engines for default models" + echo " --download-models Download default models and build required TensorRT engines" echo " --opencv-cuda Setup OpenCV with CUDA support" echo " --server Start the Comfystream server, UI and ComfyUI" echo " --help Show this help message" @@ -37,48 +36,9 @@ fi if [ "$1" = "--download-models" ]; then cd /workspace/comfystream conda activate comfystream - python src/comfystream/scripts/setup_models.py --workspace /workspace/ComfyUI - shift -fi - -DEPTH_ANYTHING_DIR="/workspace/ComfyUI/models/tensorrt/depth-anything" - -if [ "$1" = "--build-engines" ]; then - cd /workspace/comfystream - conda activate comfystream - - # Build Static Engine for Dreamshaper - python src/comfystream/scripts/build_trt.py --model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors --out-engine /workspace/ComfyUI/output/tensorrt/static-dreamshaper8_SD15_\$stat-b-1-h-512-w-512_00001_.engine - - # Build Dynamic Engine for Dreamshaper - python src/comfystream/scripts/build_trt.py \ - --model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors \ - --out-engine /workspace/ComfyUI/output/tensorrt/dynamic-dreamshaper8_SD15_\$dyn-b-1-4-2-h-448-704-512-w-448-704-512_00001_.engine \ - --width 512 \ - --height 512 \ - --min-width 448 \ - --min-height 448 \ - --max-width 704 \ - --max-height 704 - - # Build Engine for Depth Anything V2 - if [ ! -f "$DEPTH_ANYTHING_DIR/depth_anything_vitl14-fp16.engine" ]; then - if [ ! -d "$DEPTH_ANYTHING_DIR" ]; then - mkdir -p "$DEPTH_ANYTHING_DIR" - fi - cd "$DEPTH_ANYTHING_DIR" - python /workspace/ComfyUI/custom_nodes/ComfyUI-Depth-Anything-Tensorrt/export_trt.py - else - echo "Engine for DepthAnything2 already exists, skipping..." - fi - - # Build Engine for Depth Anything2 (large) - if [ ! -f "$DEPTH_ANYTHING_DIR/depth_anything_v2_vitl-fp16.engine" ]; then - cd "$DEPTH_ANYTHING_DIR" - python /workspace/ComfyUI/custom_nodes/ComfyUI-Depth-Anything-Tensorrt/export_trt.py --trt-path "${DEPTH_ANYTHING_DIR}/depth_anything_v2_vitl-fp16.engine" --onnx-path "${DEPTH_ANYTHING_DIR}/depth_anything_v2_vitl.onnx" - else - echo "Engine for DepthAnything2 (large) already exists, skipping..." - fi + + # Now also builds engines configured in models.yaml + python src/comfystream/scripts/setup_models.py --workspace /workspace/ComfyUI --build-engines shift fi diff --git a/src/comfystream/scripts/setup_models.py b/src/comfystream/scripts/setup_models.py index f99d7c13..123ac294 100644 --- a/src/comfystream/scripts/setup_models.py +++ b/src/comfystream/scripts/setup_models.py @@ -4,36 +4,165 @@ from tqdm import tqdm import yaml import argparse +import subprocess +import sys + +# Assuming utils is importable via standard mechanisms from utils import get_config_path, load_model_config +# --- Constants --- +COMFYSTREAM_ROOT = Path(__file__).resolve().parent.parent.parent.parent +DEFAULT_WORKSPACE = os.environ.get('COMFY_UI_WORKSPACE', os.path.expanduser('~/comfyui')) + def parse_args(): - parser = argparse.ArgumentParser(description='Setup ComfyUI models') + parser = argparse.ArgumentParser(description='Download ComfyUI models and optionally build TensorRT engines.') parser.add_argument('--workspace', - default=os.environ.get('COMFY_UI_WORKSPACE', os.path.expanduser('~/comfyui')), - help='ComfyUI workspace directory (default: ~/comfyui or $COMFY_UI_WORKSPACE)') + default=DEFAULT_WORKSPACE, + help=f'ComfyUI workspace directory (default: {DEFAULT_WORKSPACE} or $COMFY_UI_WORKSPACE)') + parser.add_argument('--build-engines', + action='store_true', + help='Build TensorRT engines for configured models after downloading.') + # Add argument for config path, default relative to script location + parser.add_argument('--config', + default=None, + help='Path to the models.yaml config file.') return parser.parse_args() def download_file(url, destination, description=None): """Download a file with progress bar""" - response = requests.get(url, stream=True) - total_size = int(response.headers.get('content-length', 0)) + print(f"Attempting to download from {url} to {destination}") + try: + response = requests.get(url, stream=True, timeout=30, allow_redirects=True) + response.raise_for_status() # Raise an exception for bad status codes + total_size = int(response.headers.get('content-length', 0)) + + desc = description or Path(destination).name + progress_bar = tqdm(total=total_size, unit='iB', unit_scale=True, desc=desc, leave=False) + + destination = Path(destination) + destination.parent.mkdir(parents=True, exist_ok=True) + + with open(destination, 'wb') as file: + for data in response.iter_content(chunk_size=8192): # Increased chunk size + size = file.write(data) + progress_bar.update(size) + progress_bar.close() + if total_size != 0 and progress_bar.n != total_size: + print(f"Error: Downloaded size ({progress_bar.n}) does not match expected size ({total_size}) for {desc}") + # Decide if this should be fatal? For now just print warning. + print(f"Successfully downloaded {desc}") + return True + except requests.exceptions.RequestException as e: + print(f"Error downloading {description or url}: {e}") + if destination and Path(destination).exists(): + # Clean up partial download + try: + Path(destination).unlink() + except OSError: + pass + return False + except Exception as e: + print(f"An unexpected error occurred during download of {description or url}: {e}") + if destination and Path(destination).exists(): + try: + Path(destination).unlink() + except OSError: + pass + return False + + +def build_tensorrt_engine(workspace_dir: Path, model_path: Path, engine_config: dict): + """Builds a single TensorRT engine based on the provided config.""" + script_path_str = engine_config.get('script') + engine_rel_path = engine_config.get('engine_path') + args_template = engine_config.get('args', []) + cwd_rel_path = engine_config.get('cwd') # Get optional CWD from config + + if not script_path_str or not engine_rel_path: + print(f"Error: Invalid engine config for model {model_path.name}. Missing 'script' or 'engine_path'. Config: {engine_config}") + return False + + # Let's try resolving relative to workspace first, then comfystream root. + script_path_ws = workspace_dir / script_path_str + script_path_cs = COMFYSTREAM_ROOT / script_path_str - desc = description or os.path.basename(destination) - progress_bar = tqdm(total=total_size, unit='iB', unit_scale=True, desc=desc) + if script_path_ws.exists(): + script_path = script_path_ws + elif script_path_cs.exists(): + script_path = script_path_cs + else: + print(f"Error: Build script not found at {script_path_ws} or {script_path_cs}") + return False - destination = Path(destination) - destination.parent.mkdir(parents=True, exist_ok=True) + # Resolve engine path relative to workspace/models + if '/' in engine_rel_path and Path(engine_rel_path).parts[0] == 'output': + engine_full_path = (workspace_dir / engine_rel_path).resolve() + else: + engine_full_path = (workspace_dir / "models" / engine_rel_path).resolve() + + # Ensure parent directory exists before creating engine file + engine_full_path.parent.mkdir(parents=True, exist_ok=True) + + # Resolve CWD if specified in config, relative to workspace + absolute_cwd = (workspace_dir / cwd_rel_path).resolve() if cwd_rel_path else None + if absolute_cwd and not absolute_cwd.is_dir(): + print(f"Warning: Specified CWD '{absolute_cwd}' does not exist or is not a directory. Ignoring CWD.") + absolute_cwd = None + + if engine_full_path.exists(): + print(f"Skipping build, engine already exists: {engine_full_path}") + return True + + # Format arguments + formatted_args = [] + try: + for arg in args_template: + formatted_args.append( + str(arg).format( + model_path=model_path.resolve(), + engine_path=engine_full_path, + workspace_dir=workspace_dir.resolve(), + # Add more placeholders if needed + ) + ) + except KeyError as e: + print(f"Error formatting arguments for {script_path.name}: Missing placeholder {e}") + return False + except Exception as e: + print(f"Error formatting arguments for {script_path.name}: {e}. Args: {args_template}") + return False + + # Construct command + command = [sys.executable, str(script_path)] + formatted_args + print(f"\nBuilding engine: {' '.join(command)}") + if absolute_cwd: + print(f"Running in directory: {absolute_cwd}") + + try: + # Use subprocess.run to execute the build script + # Inherit environment, capture output + # Let the build script's output stream directly to the console + result = subprocess.run(command, cwd=absolute_cwd, check=True, env=os.environ) + if not engine_full_path.exists(): + print(f"Error: Build command completed but engine file not found at {engine_full_path}") + return False + print(f"Successfully built engine: {engine_full_path}") + return True + except subprocess.CalledProcessError as e: + print(f"Error building engine with script {script_path.name}:") + print(f"Command: {' '.join(e.cmd)}") + print(f"Return Code: {e.returncode}") + return False + except FileNotFoundError: + print(f"Error: Python executable '{sys.executable}' or script '{script_path}' not found.") + return False + except Exception as e: + print(f"An unexpected error occurred during engine build: {e}") + return False - with open(destination, 'wb') as file: - for data in response.iter_content(chunk_size=1024): - size = file.write(data) - progress_bar.update(size) - progress_bar.close() -def setup_model_files(workspace_dir, config_path=None): - """Download and setup required model files based on configuration""" - if config_path is None: - config_path = get_config_path('models.yaml') +def setup_model_files(workspace_dir: Path, config_path: Path, build_engines_flag: bool): + """Download and setup required model files based on configuration, optionally build engines.""" try: config = load_model_config(config_path) except FileNotFoundError: @@ -44,60 +173,121 @@ def setup_model_files(workspace_dir, config_path=None): return models_path = workspace_dir / "models" - base_path = workspace_dir + base_path = workspace_dir # For resolving paths like custom_nodes/ - for _, model_info in config['models'].items(): - # Determine the full path based on whether it's in custom_nodes or models - if model_info['path'].startswith('custom_nodes/'): - full_path = base_path / model_info['path'] + download_success_count = 0 + download_fail_count = 0 + build_success_count = 0 + build_fail_count = 0 + + # Ensure base models directory exists + models_path.mkdir(parents=True, exist_ok=True) + + for model_key, model_info in config['models'].items(): + # Determine the full path for the main model file + path_str = model_info['path'] + if '/' in path_str and Path(path_str).parts[0] == 'custom_nodes': + # custom_nodes paths are relative to workspace_dir (base_path) + full_path = (base_path / path_str).resolve() else: - full_path = models_path / model_info['path'] + # Paths with subdirectories (e.g. controlnet/) are relative to models_path + # Also handles paths without '/' assuming they belong in models_path root (though less likely) + full_path = (models_path / path_str).resolve() + + model_name = model_info.get('name', model_key) + downloaded_main_file = False + + # Ensure parent directory exists before downloading + full_path.parent.mkdir(parents=True, exist_ok=True) if not full_path.exists(): - print(f"Downloading {model_info['name']}...") - download_file( - model_info['url'], - full_path, - f"Downloading {model_info['name']}" - ) - print(f"Downloaded {model_info['name']} to {full_path}") - - # Handle any extra files (like configs) - if 'extra_files' in model_info: - for extra in model_info['extra_files']: - extra_path = models_path / extra['path'] - if not extra_path.exists(): - download_file( - extra['url'], - extra_path, - f"Downloading {os.path.basename(extra['path'])}" - ) - print("Models download completed!") - -def setup_directories(workspace_dir): - """Create required directories in the workspace""" - # Create base directories - workspace_dir.mkdir(parents=True, exist_ok=True) - models_dir = workspace_dir / "models" - models_dir.mkdir(parents=True, exist_ok=True) - - # Create model subdirectories - model_dirs = [ - "checkpoints/SD1.5", - "controlnet", - "vae", - "tensorrt", - "unet", - "LLM", - ] - for dir_name in model_dirs: - (models_dir / dir_name).mkdir(parents=True, exist_ok=True) - -def setup_models(): + print(f"\nProcessing model: {model_name}") + if 'url' in model_info: + if download_file(model_info['url'], full_path, f"Downloading {model_name}"): + download_success_count += 1 + downloaded_main_file = True + else: + download_fail_count += 1 + else: + print(f"Warning: No URL specified for {model_name}, skipping download.") + else: + print(f"\nModel exists: {model_name} at {full_path}") + downloaded_main_file = True # Treat existing file as successfully "downloaded" for build logic + + # Handle any extra files (like configs) + if 'extra_files' in model_info: + for extra in model_info['extra_files']: + extra_path_str = extra['path'] + # Resolve extra file paths similarly + if '/' in extra_path_str and Path(extra_path_str).parts[0] == 'custom_nodes': + extra_full_path = (base_path / extra_path_str).resolve() + else: + extra_full_path = (models_path / extra_path_str).resolve() + + # Ensure parent directory exists before downloading extra file + extra_full_path.parent.mkdir(parents=True, exist_ok=True) + + if not extra_full_path.exists(): + if 'url' in extra: + print(f"Downloading extra file for {model_name}: {Path(extra_path_str).name}") + if download_file(extra['url'], extra_full_path, f"Downloading {Path(extra_path_str).name}"): + download_success_count += 1 + else: + download_fail_count += 1 + else: + print(f"Warning: No URL for extra file {extra_path_str}, skipping.") + else: + print(f"Extra file exists: {extra_full_path}") + + # --- Build TensorRT Engines --- + if build_engines_flag and downloaded_main_file and 'tensorrt' in model_info: + tensorrt_config = model_info['tensorrt'] + if tensorrt_config.get('build'): + print(f"\nAttempting to build TensorRT engines for {model_name}...") + engines_to_build = tensorrt_config.get('engines', []) + if not engines_to_build: + print(f"Warning: 'tensorrt: build: true' but no 'engines' listed for {model_name}") + continue + + for engine_conf in engines_to_build: + if build_tensorrt_engine(workspace_dir, full_path, engine_conf): + build_success_count += 1 + else: + build_fail_count += 1 + else: + print(f"Skipping TensorRT build for {model_name} (build flag is false in config)") + + print("\n--- Summary ---") + print(f"Model Downloads: {download_success_count} succeeded, {download_fail_count} failed.") + if build_engines_flag: + print(f"TensorRT Builds: {build_success_count} succeeded, {build_fail_count} failed.") + print("Setup process finished!") + + +def main(): args = parse_args() - workspace_dir = Path(args.workspace) + workspace_dir = Path(args.workspace).resolve() # Resolve workspace path + + # Create base workspace and output/tensorrt directories if they don't exist + # Model subdirs are created on demand in setup_model_files + workspace_dir.mkdir(parents=True, exist_ok=True) + (workspace_dir / "output" / "tensorrt").mkdir(parents=True, exist_ok=True) + + # Determine config path + if args.config: + config_path = Path(args.config).resolve() + else: + # Default to configs/models.yaml relative to COMFYSTREAM_ROOT + config_path = COMFYSTREAM_ROOT / 'configs' / 'models.yaml' + + print(f"Using workspace: {workspace_dir}") + print(f"Using model config: {config_path}") + if args.build_engines: + print("TensorRT engine building is ENABLED.") + else: + print("TensorRT engine building is DISABLED (use --build-engines to enable).") - setup_directories(workspace_dir) - setup_model_files(workspace_dir) + setup_model_files(workspace_dir, config_path, args.build_engines) -setup_models() +if __name__ == "__main__": + main() From f3b5d66c7fcad7e7c04f33bd9442808262315978 Mon Sep 17 00:00:00 2001 From: "Jason.Stone" Date: Mon, 28 Apr 2025 13:03:22 +0000 Subject: [PATCH 2/4] Remove --use-fp32 false flag from depth-anything model configs --- configs/models.yaml | 2 -- 1 file changed, 2 deletions(-) diff --git a/configs/models.yaml b/configs/models.yaml index bdc1894f..e18bcba3 100644 --- a/configs/models.yaml +++ b/configs/models.yaml @@ -64,7 +64,6 @@ models: - "--onnx-path" - "{model_path}" - "--use-fp32" - - "false" depth-anything-v2-large-onnx: name: "DepthAnything V2 Large ONNX" @@ -82,7 +81,6 @@ models: - "--onnx-path" - "{model_path}" - "--use-fp32" - - "false" # TAESD models taesd: From c62f3ad3358f2ca65fcdfd1128e8fd6cf17f6ab0 Mon Sep 17 00:00:00 2001 From: "Jason.Stone" Date: Mon, 28 Apr 2025 13:14:56 +0000 Subject: [PATCH 3/4] Remove --use-fp32 flag from depth-anything model conversion commands --- configs/models.yaml | 2 -- 1 file changed, 2 deletions(-) diff --git a/configs/models.yaml b/configs/models.yaml index e18bcba3..00b96a19 100644 --- a/configs/models.yaml +++ b/configs/models.yaml @@ -63,7 +63,6 @@ models: - "{engine_path}" - "--onnx-path" - "{model_path}" - - "--use-fp32" depth-anything-v2-large-onnx: name: "DepthAnything V2 Large ONNX" @@ -80,7 +79,6 @@ models: - "{engine_path}" - "--onnx-path" - "{model_path}" - - "--use-fp32" # TAESD models taesd: From ce3e5f0be14fa67d203fab1617ae10264cb3d247 Mon Sep 17 00:00:00 2001 From: "Jason.Stone" Date: Wed, 30 Apr 2025 08:11:51 +0000 Subject: [PATCH 4/4] Refactor TensorRT engine setup: remove unnecessary cwd specification and enhance logging with color-coded messages --- configs/models.yaml | 2 - src/comfystream/scripts/setup_models.py | 145 ++++++++++-------------- 2 files changed, 62 insertions(+), 85 deletions(-) diff --git a/configs/models.yaml b/configs/models.yaml index 00b96a19..4371b71f 100644 --- a/configs/models.yaml +++ b/configs/models.yaml @@ -57,7 +57,6 @@ models: engines: - script: "custom_nodes/ComfyUI-Depth-Anything-Tensorrt/export_trt.py" engine_path: "tensorrt/depth-anything/depth_anything_vitl14-fp16.engine" - cwd: "custom_nodes/ComfyUI-Depth-Anything-Tensorrt" args: - "--trt-path" - "{engine_path}" @@ -73,7 +72,6 @@ models: engines: - script: "custom_nodes/ComfyUI-Depth-Anything-Tensorrt/export_trt.py" engine_path: "tensorrt/depth-anything/depth_anything_v2_vitl-fp16.engine" - cwd: "custom_nodes/ComfyUI-Depth-Anything-Tensorrt" args: - "--trt-path" - "{engine_path}" diff --git a/src/comfystream/scripts/setup_models.py b/src/comfystream/scripts/setup_models.py index 123ac294..40100268 100644 --- a/src/comfystream/scripts/setup_models.py +++ b/src/comfystream/scripts/setup_models.py @@ -6,12 +6,13 @@ import argparse import subprocess import sys +from rich import print # Assuming utils is importable via standard mechanisms from utils import get_config_path, load_model_config # --- Constants --- -COMFYSTREAM_ROOT = Path(__file__).resolve().parent.parent.parent.parent +COMFYSTREAM_ROOT = Path(__file__).parents[3] DEFAULT_WORKSPACE = os.environ.get('COMFY_UI_WORKSPACE', os.path.expanduser('~/comfyui')) def parse_args(): @@ -22,15 +23,11 @@ def parse_args(): parser.add_argument('--build-engines', action='store_true', help='Build TensorRT engines for configured models after downloading.') - # Add argument for config path, default relative to script location - parser.add_argument('--config', - default=None, - help='Path to the models.yaml config file.') return parser.parse_args() def download_file(url, destination, description=None): """Download a file with progress bar""" - print(f"Attempting to download from {url} to {destination}") + print(f"[blue]Attempting to download from {url} to {destination}[/blue]") try: response = requests.get(url, stream=True, timeout=30, allow_redirects=True) response.raise_for_status() # Raise an exception for bad status codes @@ -43,17 +40,14 @@ def download_file(url, destination, description=None): destination.parent.mkdir(parents=True, exist_ok=True) with open(destination, 'wb') as file: - for data in response.iter_content(chunk_size=8192): # Increased chunk size + for data in response.iter_content(chunk_size=8192): size = file.write(data) progress_bar.update(size) progress_bar.close() - if total_size != 0 and progress_bar.n != total_size: - print(f"Error: Downloaded size ({progress_bar.n}) does not match expected size ({total_size}) for {desc}") - # Decide if this should be fatal? For now just print warning. - print(f"Successfully downloaded {desc}") + print(f"[blue]Successfully downloaded {desc}[/blue]") return True except requests.exceptions.RequestException as e: - print(f"Error downloading {description or url}: {e}") + print(f"[red]Error downloading {description or url}: {e}[/red]") if destination and Path(destination).exists(): # Clean up partial download try: @@ -62,7 +56,7 @@ def download_file(url, destination, description=None): pass return False except Exception as e: - print(f"An unexpected error occurred during download of {description or url}: {e}") + print(f"[red]An unexpected error occurred during download of {description or url}: {e}[/red]") if destination and Path(destination).exists(): try: Path(destination).unlink() @@ -74,43 +68,43 @@ def download_file(url, destination, description=None): def build_tensorrt_engine(workspace_dir: Path, model_path: Path, engine_config: dict): """Builds a single TensorRT engine based on the provided config.""" script_path_str = engine_config.get('script') - engine_rel_path = engine_config.get('engine_path') + engine_path_str = engine_config.get('engine_path') args_template = engine_config.get('args', []) - cwd_rel_path = engine_config.get('cwd') # Get optional CWD from config - if not script_path_str or not engine_rel_path: - print(f"Error: Invalid engine config for model {model_path.name}. Missing 'script' or 'engine_path'. Config: {engine_config}") - return False + assert script_path_str, f"[red]Error: Invalid engine config for model {model_path.name}. Missing 'script' field in model config.[/red]" + assert engine_path_str, f"[red]Error: Invalid engine config for model {model_path.name}. Missing 'engine_path' in model config.[/red]" # Let's try resolving relative to workspace first, then comfystream root. script_path_ws = workspace_dir / script_path_str script_path_cs = COMFYSTREAM_ROOT / script_path_str + script_path = None if script_path_ws.exists(): - script_path = script_path_ws + script_path = script_path_ws.resolve() elif script_path_cs.exists(): - script_path = script_path_cs + script_path = script_path_cs.resolve() else: - print(f"Error: Build script not found at {script_path_ws} or {script_path_cs}") + print(f"[red]Error: Build script not found at {script_path_ws} or {script_path_cs}[/red]") return False - # Resolve engine path relative to workspace/models - if '/' in engine_rel_path and Path(engine_rel_path).parts[0] == 'output': - engine_full_path = (workspace_dir / engine_rel_path).resolve() + # Determine the correct working directory from the script's location + absolute_cwd = script_path.parent + if not absolute_cwd.is_dir(): + # This should theoretically not happen if the script exists, but good practice + print(f"[red]Error: Determined CWD '{absolute_cwd}' is not a valid directory.[/red]") + return False + + # Resolve engine path relative to workspace/models or output/ + if '/' in engine_path_str and Path(engine_path_str).parts[0] == 'output': + engine_full_path = (workspace_dir / engine_path_str).resolve() else: - engine_full_path = (workspace_dir / "models" / engine_rel_path).resolve() + engine_full_path = (workspace_dir / "models" / engine_path_str).resolve() # Ensure parent directory exists before creating engine file engine_full_path.parent.mkdir(parents=True, exist_ok=True) - # Resolve CWD if specified in config, relative to workspace - absolute_cwd = (workspace_dir / cwd_rel_path).resolve() if cwd_rel_path else None - if absolute_cwd and not absolute_cwd.is_dir(): - print(f"Warning: Specified CWD '{absolute_cwd}' does not exist or is not a directory. Ignoring CWD.") - absolute_cwd = None - if engine_full_path.exists(): - print(f"Skipping build, engine already exists: {engine_full_path}") + print(f"[blue]Skipping build, engine already exists: {engine_full_path}[/blue]") return True # Format arguments @@ -122,42 +116,41 @@ def build_tensorrt_engine(workspace_dir: Path, model_path: Path, engine_config: model_path=model_path.resolve(), engine_path=engine_full_path, workspace_dir=workspace_dir.resolve(), - # Add more placeholders if needed ) ) except KeyError as e: - print(f"Error formatting arguments for {script_path.name}: Missing placeholder {e}") + print(f"[red]Error formatting arguments for {script_path.name}: Missing placeholder {e}[/red]") return False except Exception as e: - print(f"Error formatting arguments for {script_path.name}: {e}. Args: {args_template}") + print(f"[red]Error formatting arguments for {script_path.name}: {e}. Args: {args_template}[/red]") return False # Construct command command = [sys.executable, str(script_path)] + formatted_args - print(f"\nBuilding engine: {' '.join(command)}") - if absolute_cwd: - print(f"Running in directory: {absolute_cwd}") + print(f"\n[blue]Building engine: {' '.join(map(str, command))}[/blue]") + # Always print the CWD being used now + print(f"[blue]Running in directory: {absolute_cwd}[/blue]") try: # Use subprocess.run to execute the build script # Inherit environment, capture output - # Let the build script's output stream directly to the console - result = subprocess.run(command, cwd=absolute_cwd, check=True, env=os.environ) + # Ensure absolute_cwd is passed correctly + result = subprocess.run(command, cwd=str(absolute_cwd), check=True, env=os.environ) if not engine_full_path.exists(): - print(f"Error: Build command completed but engine file not found at {engine_full_path}") + print(f"[red]Error: Build command completed but engine file not found at {engine_full_path}[/red]") return False - print(f"Successfully built engine: {engine_full_path}") + print(f"[blue]Successfully built engine: {engine_full_path}[/blue]") return True except subprocess.CalledProcessError as e: - print(f"Error building engine with script {script_path.name}:") - print(f"Command: {' '.join(e.cmd)}") - print(f"Return Code: {e.returncode}") + print(f"[red]Error building engine with script {script_path.name}:[/red]") + print(f"[red]Command: {' '.join(map(str, e.cmd))}[/red]") + print(f"[red]Return Code: {e.returncode}[/red]") return False except FileNotFoundError: - print(f"Error: Python executable '{sys.executable}' or script '{script_path}' not found.") + print(f"[red]Error: Python executable '{sys.executable}' or script '{script_path}' not found.[/red]") return False except Exception as e: - print(f"An unexpected error occurred during engine build: {e}") + print(f"[red]An unexpected error occurred during engine build: {e}[/red]") return False @@ -166,10 +159,10 @@ def setup_model_files(workspace_dir: Path, config_path: Path, build_engines_flag try: config = load_model_config(config_path) except FileNotFoundError: - print(f"Error: Model config file not found at {config_path}") + print(f"[red]Error: Model config file not found at {config_path}[/red]") return except yaml.YAMLError as e: - print(f"Error parsing model config file: {e}") + print(f"[red]Error parsing model config file: {e}[/red]") return models_path = workspace_dir / "models" @@ -186,13 +179,7 @@ def setup_model_files(workspace_dir: Path, config_path: Path, build_engines_flag for model_key, model_info in config['models'].items(): # Determine the full path for the main model file path_str = model_info['path'] - if '/' in path_str and Path(path_str).parts[0] == 'custom_nodes': - # custom_nodes paths are relative to workspace_dir (base_path) - full_path = (base_path / path_str).resolve() - else: - # Paths with subdirectories (e.g. controlnet/) are relative to models_path - # Also handles paths without '/' assuming they belong in models_path root (though less likely) - full_path = (models_path / path_str).resolve() + full_path = (models_path / path_str).resolve() model_name = model_info.get('name', model_key) downloaded_main_file = False @@ -201,7 +188,7 @@ def setup_model_files(workspace_dir: Path, config_path: Path, build_engines_flag full_path.parent.mkdir(parents=True, exist_ok=True) if not full_path.exists(): - print(f"\nProcessing model: {model_name}") + print(f"\n[blue]Processing model: {model_name}[/blue]") if 'url' in model_info: if download_file(model_info['url'], full_path, f"Downloading {model_name}"): download_success_count += 1 @@ -209,44 +196,40 @@ def setup_model_files(workspace_dir: Path, config_path: Path, build_engines_flag else: download_fail_count += 1 else: - print(f"Warning: No URL specified for {model_name}, skipping download.") + print(f"[yellow]Warning: No URL specified for {model_name}, skipping download.[/yellow]") else: - print(f"\nModel exists: {model_name} at {full_path}") + print(f"\n[blue]Model exists: {model_name} at {full_path}[/blue]") downloaded_main_file = True # Treat existing file as successfully "downloaded" for build logic # Handle any extra files (like configs) if 'extra_files' in model_info: for extra in model_info['extra_files']: extra_path_str = extra['path'] - # Resolve extra file paths similarly - if '/' in extra_path_str and Path(extra_path_str).parts[0] == 'custom_nodes': - extra_full_path = (base_path / extra_path_str).resolve() - else: - extra_full_path = (models_path / extra_path_str).resolve() + extra_full_path = (models_path / extra_path_str).resolve() # Ensure parent directory exists before downloading extra file extra_full_path.parent.mkdir(parents=True, exist_ok=True) if not extra_full_path.exists(): if 'url' in extra: - print(f"Downloading extra file for {model_name}: {Path(extra_path_str).name}") + print(f"[blue]Downloading extra file for {model_name}: {Path(extra_path_str).name}[/blue]") if download_file(extra['url'], extra_full_path, f"Downloading {Path(extra_path_str).name}"): download_success_count += 1 else: download_fail_count += 1 else: - print(f"Warning: No URL for extra file {extra_path_str}, skipping.") + print(f"[yellow]Warning: No URL for extra file {extra_path_str}, skipping.[/yellow]") else: - print(f"Extra file exists: {extra_full_path}") + print(f"[blue]Extra file exists: {extra_full_path}[/blue]") # --- Build TensorRT Engines --- if build_engines_flag and downloaded_main_file and 'tensorrt' in model_info: tensorrt_config = model_info['tensorrt'] if tensorrt_config.get('build'): - print(f"\nAttempting to build TensorRT engines for {model_name}...") + print(f"\n[blue]Attempting to build TensorRT engines for {model_name}...[/blue]") engines_to_build = tensorrt_config.get('engines', []) if not engines_to_build: - print(f"Warning: 'tensorrt: build: true' but no 'engines' listed for {model_name}") + print(f"[yellow]Warning: 'tensorrt: build: true' but no 'engines' listed for {model_name}[/yellow]") continue for engine_conf in engines_to_build: @@ -255,13 +238,13 @@ def setup_model_files(workspace_dir: Path, config_path: Path, build_engines_flag else: build_fail_count += 1 else: - print(f"Skipping TensorRT build for {model_name} (build flag is false in config)") + print(f"[blue]Skipping TensorRT build for {model_name} (build flag is false in config)[/blue]") print("\n--- Summary ---") - print(f"Model Downloads: {download_success_count} succeeded, {download_fail_count} failed.") + print(f"Model Downloads: {'[green]' + str(download_success_count) + '[/green]'} succeeded, {'[red]' + str(download_fail_count) + '[/red]'} failed.") if build_engines_flag: - print(f"TensorRT Builds: {build_success_count} succeeded, {build_fail_count} failed.") - print("Setup process finished!") + print(f"TensorRT Builds: {'[green]' + str(build_success_count) + '[/green]'} succeeded, {'[red]' + str(build_fail_count) + '[/red]'} failed.") + print("[green]Setup process finished![/green]") def main(): @@ -269,23 +252,19 @@ def main(): workspace_dir = Path(args.workspace).resolve() # Resolve workspace path # Create base workspace and output/tensorrt directories if they don't exist - # Model subdirs are created on demand in setup_model_files workspace_dir.mkdir(parents=True, exist_ok=True) (workspace_dir / "output" / "tensorrt").mkdir(parents=True, exist_ok=True) # Determine config path - if args.config: - config_path = Path(args.config).resolve() - else: - # Default to configs/models.yaml relative to COMFYSTREAM_ROOT - config_path = COMFYSTREAM_ROOT / 'configs' / 'models.yaml' + config_path = COMFYSTREAM_ROOT / 'configs' / 'models.yaml' + assert config_path.exists(), f"[red]Model config file not found at {config_path}. Please check the path.[/red]" - print(f"Using workspace: {workspace_dir}") - print(f"Using model config: {config_path}") + print(f"[blue]Using workspace: {workspace_dir}[/blue]") + print(f"[blue]Using model config: {config_path}[/blue]") if args.build_engines: - print("TensorRT engine building is ENABLED.") + print("[blue]TensorRT engine building is ENABLED.[/blue]") else: - print("TensorRT engine building is DISABLED (use --build-engines to enable).") + print("[blue]TensorRT engine building is DISABLED (use --build-engines to enable).[/blue]") setup_model_files(workspace_dir, config_path, args.build_engines)