From e7483499c62a3a7100e6caed5c4bb48784c2371a Mon Sep 17 00:00:00 2001 From: Alex Sadleir Date: Mon, 29 Dec 2025 12:44:42 +1100 Subject: [PATCH] Add native ComfyUI progress bar support for model downloads --- dap_modules/dap_loader.py | 54 +++++++++++++++++++++++------------- tests/conftest.py | 39 ++++++++++++++------------ tests/unit/test_dap_nodes.py | 12 +++++--- 3 files changed, 65 insertions(+), 40 deletions(-) diff --git a/dap_modules/dap_loader.py b/dap_modules/dap_loader.py index 3ce702e..16b306c 100644 --- a/dap_modules/dap_loader.py +++ b/dap_modules/dap_loader.py @@ -1,9 +1,11 @@ import os import sys import torch +import requests import folder_paths import comfy.model_management as mm -from huggingface_hub import hf_hub_download +from comfy.utils import ProgressBar +from huggingface_hub import hf_hub_download, hf_hub_url from argparse import Namespace # Add dap_core to sys.path so imports work @@ -21,6 +23,28 @@ except ImportError: DAP = None +def download_with_progress(url, dest_path, filename): + print(f"Downloading {filename} from {url}...") + response = requests.get(url, stream=True) + response.raise_for_status() + + total_size = int(response.headers.get("content-length", 0)) + pbar = ProgressBar(total_size) + + # Ensure directory exists + os.makedirs(os.path.dirname(dest_path), exist_ok=True) + + # Download with progress updates + with open(dest_path, "wb") as f: + downloaded = 0 + for chunk in response.iter_content(chunk_size=1024 * 1024): # 1MB chunks + if chunk: + f.write(chunk) + downloaded += len(chunk) + pbar.update_absolute(downloaded, total_size) + print(f"Finished downloading {filename}") + + class DAP_Loader: @classmethod def INPUT_TYPES(s): @@ -40,38 +64,34 @@ class DAP_Loader: device = mm.get_torch_device() # 1. Define Model Files - # Mapping size to HF filename. Currently only 'model.pth' is available. - # size mapping is kept for future expansion. model_filename = "model.pth" + repo_id = "Insta360-Research/DAP-weights" # 2. Download/Locate Model - # Check standard ComfyUI locations first ckpt_path = folder_paths.get_full_path("checkpoints", model_filename) if not ckpt_path: - # Check dedicated folder dap_models_dir = os.path.join(folder_paths.models_dir, "dap") - if not os.path.exists(dap_models_dir): - os.makedirs(dap_models_dir) - ckpt_path = os.path.join(dap_models_dir, model_filename) if not os.path.exists(ckpt_path): - print(f"Downloading {model_filename} to {ckpt_path}...") + url = hf_hub_url(repo_id=repo_id, filename=model_filename) try: + download_with_progress(url, ckpt_path, model_filename) + except Exception as e: + # Fallback to standard hf_hub_download if manual fails + print( + f"Manual download failed: {e}. Falling back to hf_hub_download..." + ) ckpt_path = hf_hub_download( - repo_id="Insta360-Research/DAP-weights", + repo_id=repo_id, filename=model_filename, local_dir=dap_models_dir, ) - except Exception as e: - raise RuntimeError(f"Failed to download model: {e}") # 3. Initialize Model with Patching args = Namespace() - args.midas_model_type = ( - model_size # Fixed to 'vitl' for now based on hardcoded arch in dap.py - ) + args.midas_model_type = model_size args.fine_tune_type = "none" args.min_depth = 0.001 args.max_depth = 1.0 @@ -80,12 +100,9 @@ class DAP_Loader: original_cwd = os.getcwd() try: os.chdir(dap_core_path) - print(f"Changed CWD to {dap_core_path} for initialization") - model = DAP(args) state_dict = torch.load(ckpt_path, map_location="cpu") - # Handle possible "module." prefix if saved from DDP if state_dict and list(state_dict.keys())[0].startswith("module."): state_dict = { k.replace("module.", ""): v for k, v in state_dict.items() @@ -95,7 +112,6 @@ class DAP_Loader: finally: os.chdir(original_cwd) - print(f"Restored CWD to {original_cwd}") # 4. Quantization / Precision if precision == "fp16": diff --git a/tests/conftest.py b/tests/conftest.py index 4b8435b..f16f2b3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,29 +11,34 @@ sys.path.insert(0, str(plugin_dir)) sys.path.insert(0, str(plugin_dir / "dap_core")) # 2. Mock ComfyUI modules -# We use simple objects instead of MagicMock for the base modules to avoid weird recursion -mock_folder_paths = type("folder_paths", (), {})() +mock_folder_paths = MagicMock() mock_folder_paths.models_dir = str(plugin_dir / "models_mock") -mock_folder_paths.get_full_path = lambda *args, **kwargs: None -mock_folder_paths.get_filename_list = lambda *args, **kwargs: [] +mock_folder_paths.get_full_path.return_value = None sys.modules["folder_paths"] = mock_folder_paths -mock_mm = type("model_management", (), {})() -mock_mm.get_torch_device = lambda: torch.device("cpu") -mock_mm.unet_offload_device = lambda: torch.device("cpu") -mock_mm.load_model_gpu = lambda x: None -mock_mm.soft_empty_cache = lambda: None -mock_mm.is_device_mps = lambda x: False -mock_mm.get_autocast_device = lambda x: "cpu" - -mock_comfy = type("comfy", (), {})() -mock_comfy.model_management = mock_mm -sys.modules["comfy"] = mock_comfy +mock_mm = MagicMock() +mock_mm.get_torch_device.return_value = torch.device("cpu") +mock_mm.unet_offload_device.return_value = torch.device("cpu") sys.modules["comfy.model_management"] = mock_mm +mock_comfy = MagicMock() +mock_comfy.model_management = mock_mm +sys.modules["comfy"] = mock_comfy + +mock_utils = MagicMock() +sys.modules["comfy.utils"] = mock_utils + +# Mock huggingface_hub +mock_hf = MagicMock() +mock_hf.hf_hub_url.return_value = "http://mock-url" +sys.modules["huggingface_hub"] = mock_hf + +# Mock requests +mock_requests = MagicMock() +sys.modules["requests"] = mock_requests + + # 3. Fixtures - - @pytest.fixture def dummy_image(): """Create a dummy 2:1 panorama-style image tensor [1, 512, 1024, 3]""" diff --git a/tests/unit/test_dap_nodes.py b/tests/unit/test_dap_nodes.py index 0152006..6aae912 100644 --- a/tests/unit/test_dap_nodes.py +++ b/tests/unit/test_dap_nodes.py @@ -7,12 +7,16 @@ from dap_modules.dap_conversions import DAP_ERP_to_Cubemap from dap_modules.dap_normal import DAP_Normal_Map +from unittest.mock import MagicMock, patch + + def test_dap_loader_init(): """Verify DAP_Loader can be instantiated and has correct INPUT_TYPES""" - loader = DAP_Loader() - input_types = loader.INPUT_TYPES() - assert "model_size" in input_types["required"] - assert "precision" in input_types["required"] + with patch("huggingface_hub.hf_hub_download"), patch("requests.get"): + loader = DAP_Loader() + input_types = loader.INPUT_TYPES() + assert "model_size" in input_types["required"] + assert "precision" in input_types["required"] def test_dap_inference_init():