nodes/maxious-comfyui-dap
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add native ComfyUI progress bar support for model downloads

Browse Source
This commit is contained in:
Alex Sadleir
2025-12-29 12:44:42 +11:00
parent 6ffcc5c6c5
commit e7483499c6
3 changed files with 65 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dap_modules/dap_loader.py
+35 -19
View File
@@ -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":
tests/conftest.py
+22 -17
View File
@@ -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]"""
tests/unit/test_dap_nodes.py
+8 -4
View File
@@ -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():