Andrea Pozzetti
2859c0a3e5
Split SECAD-Net nodes and utilities from the CADabra monorepo. Voxel-based sketch-extrude CAD reconstruction. Original source: https://github.com/PozzettiAndrea/ComfyUI-CADabra
108 lines
3.9 KiB
Python
108 lines
3.9 KiB
Python
# SPDX-License-Identifier: GPL-3.0-or-later
|
|
# Originally from ComfyUI-CADabra: https://github.com/PozzettiAndrea/ComfyUI-CADabra
|
|
# Copyright (C) 2025 ComfyUI-CADabra Contributors
|
|
|
|
"""
|
|
Model loader utilities for SECAD-Net
|
|
Downloads and caches pretrained models from Google Drive
|
|
"""
|
|
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
# SECAD-Net model URLs and metadata
|
|
SECADNET_MODELS = {
|
|
"secadnet_abc": {
|
|
"url": "https://drive.google.com/uc?export=download&id=1hLJ6JYI1eXoFG-Qc13Bx3o1s-nt9yHkH",
|
|
"filename": "ModelParameters/best.pth",
|
|
"description": "SECAD-Net pretrained on ABC dataset (4 primitives)",
|
|
"folder_id": "1hLJ6JYI1eXoFG-Qc13Bx3o1s-nt9yHkH",
|
|
},
|
|
}
|
|
|
|
|
|
def get_secadnet_models_dir() -> Path:
|
|
"""
|
|
Get the models directory for SECAD-Net models.
|
|
Creates ComfyUI/models/cadrecon/secadnet/ if it doesn't exist.
|
|
"""
|
|
current_dir = Path(__file__).parent.parent # ComfyUI-SECADNET/
|
|
comfyui_dir = current_dir.parent.parent # ComfyUI/
|
|
models_dir = comfyui_dir / "models" / "cadrecon" / "secadnet"
|
|
models_dir.mkdir(parents=True, exist_ok=True)
|
|
return models_dir
|
|
|
|
|
|
def get_secadnet_model_path(model_name: str) -> Optional[Path]:
|
|
"""Get the path to a specific SECAD-Net model."""
|
|
if model_name not in SECADNET_MODELS:
|
|
raise ValueError(f"Unknown model: {model_name}. Available: {list(SECADNET_MODELS.keys())}")
|
|
|
|
models_dir = get_secadnet_models_dir()
|
|
filename = SECADNET_MODELS[model_name]["filename"]
|
|
model_path = models_dir / filename
|
|
|
|
if model_path.exists():
|
|
return model_path
|
|
return None
|
|
|
|
|
|
def download_secadnet_model(model_name: str, force_download: bool = False) -> Optional[Path]:
|
|
"""
|
|
Download a SECAD-Net model from Google Drive if not already cached.
|
|
|
|
Note: Google Drive downloads require gdown for large files.
|
|
Install with: pip install gdown
|
|
"""
|
|
if model_name not in SECADNET_MODELS:
|
|
print(f"[ERROR] Unknown model: {model_name}")
|
|
print(f" Available models: {list(SECADNET_MODELS.keys())}")
|
|
return None
|
|
|
|
model_info = SECADNET_MODELS[model_name]
|
|
models_dir = get_secadnet_models_dir()
|
|
model_path = models_dir / model_info["filename"]
|
|
|
|
if model_path.exists() and not force_download:
|
|
print(f"[OK] Model already downloaded: {model_path}")
|
|
return model_path
|
|
|
|
print(f"[SECAD-Net] Model not found locally, downloading {model_name}...")
|
|
print(f" Description: {model_info['description']}")
|
|
|
|
try:
|
|
import gdown
|
|
folder_id = model_info.get("folder_id")
|
|
if folder_id:
|
|
print(f"[SECAD-Net] Downloading from Google Drive folder...")
|
|
output_dir = str(models_dir)
|
|
gdown.download_folder(id=folder_id, output=output_dir, quiet=False)
|
|
|
|
pth_files = list(models_dir.glob("**/*.pth"))
|
|
if pth_files:
|
|
best_pth = None
|
|
for pth in pth_files:
|
|
if pth.name == "best.pth":
|
|
best_pth = pth
|
|
break
|
|
if best_pth is None:
|
|
best_pth = pth_files[0]
|
|
|
|
model_path.parent.mkdir(parents=True, exist_ok=True)
|
|
if best_pth != model_path:
|
|
import shutil
|
|
shutil.copy2(best_pth, model_path)
|
|
print(f"[OK] Model downloaded: {model_path}")
|
|
return model_path
|
|
else:
|
|
print("[WARN] gdown completed but no .pth files found")
|
|
except ImportError:
|
|
print("[WARN] gdown not installed. Install with: pip install gdown")
|
|
print("[WARN] Falling back to direct download (may fail for large files)...")
|
|
|
|
print(f"\n[WARN] Automatic download failed. Please download manually:")
|
|
print(f" 1. Visit: https://drive.google.com/drive/folders/{model_info.get('folder_id', '')}")
|
|
print(f" 2. Download the .pth file")
|
|
print(f" 3. Save to: {model_path}")
|
|
return None
|