Improve installer messages and error handling
- Enhanced .gitignore to exclude test scripts and temporary files - Removed post-processor dependency checks from auto_installer - Improved CUDA requirements feedback with clearer formatting - Added graceful handling for corrupted spatial_correlation_sampler - Fixed SyntaxError catching in addition to ImportError - Better user guidance for compilation prerequisites 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
jonstreeter
@@ -19,7 +19,13 @@ __pycache__/
|
||||
|
||||
# Checkpoints (too large for git)
|
||||
checkpoints/*.pth
|
||||
checkpoints/
|
||||
data/*.pth
|
||||
*.zip
|
||||
|
||||
# Test and helper scripts
|
||||
test_*.py
|
||||
install_*.py
|
||||
|
||||
# Virtual environment
|
||||
venv/
|
||||
@@ -28,9 +34,11 @@ ENV/
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
**/.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
.style.yapf
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
|
||||
@@ -138,32 +138,6 @@ def ensure_dependencies_for_encoder(encoder_type: str) -> bool:
|
||||
return True # Unknown encoder, assume no dependencies
|
||||
|
||||
|
||||
def ensure_dependencies_for_post_processor(processor_type: str) -> bool:
|
||||
"""
|
||||
Ensure all dependencies for a given post-processor are installed.
|
||||
|
||||
Args:
|
||||
processor_type: Post-processor name (e.g., 'color_matcher', 'wls')
|
||||
|
||||
Returns:
|
||||
True if all dependencies are available, False otherwise
|
||||
"""
|
||||
if processor_type == 'none':
|
||||
return True
|
||||
|
||||
elif processor_type == 'color_matcher':
|
||||
return ensure_color_matcher()
|
||||
|
||||
elif processor_type in ['wls', 'guided']:
|
||||
return ensure_opencv_contrib()
|
||||
|
||||
elif processor_type == 'bilateral':
|
||||
return True # opencv-python is already required
|
||||
|
||||
else:
|
||||
return True # Unknown processor, assume no dependencies
|
||||
|
||||
|
||||
def install_all_modern_components() -> dict:
|
||||
"""
|
||||
Install all modern components at once.
|
||||
@@ -177,8 +151,6 @@ def install_all_modern_components() -> dict:
|
||||
results = {
|
||||
'timm': ensure_timm(),
|
||||
'clip': ensure_clip(),
|
||||
'color_matcher': ensure_color_matcher(),
|
||||
'opencv_contrib': ensure_opencv_contrib(),
|
||||
}
|
||||
|
||||
# Summary
|
||||
|
||||
+39
-16
@@ -179,23 +179,31 @@ def check_cuda_requirements() -> Tuple[bool, str]:
|
||||
issues.append("❌ PyTorch not installed")
|
||||
|
||||
if issues:
|
||||
feedback = "\n[ColorMNet] Missing requirements for CUDA extensions:\n"
|
||||
feedback = "\n[ColorMNet] ═══════════════════════════════════════════════════════════════\n"
|
||||
feedback += "[ColorMNet] Missing requirements for CUDA-optimized compilation:\n"
|
||||
feedback += "[ColorMNet] ───────────────────────────────────────────────────────────────\n"
|
||||
for issue in issues:
|
||||
feedback += f"[ColorMNet] {issue}\n"
|
||||
feedback += "\n[ColorMNet] To enable optimized performance:\n"
|
||||
feedback += f"[ColorMNet] ✗ {issue}\n"
|
||||
|
||||
feedback += "\n[ColorMNet] To enable 2-5x faster performance, install:\n"
|
||||
feedback += "[ColorMNet] ───────────────────────────────────────────────────────────────\n"
|
||||
|
||||
if any("CUDA_HOME" in i for i in issues):
|
||||
feedback += "[ColorMNet] 1. Install CUDA Toolkit (match PyTorch version)\n"
|
||||
feedback += "[ColorMNet] Download: https://developer.nvidia.com/cuda-toolkit\n"
|
||||
feedback += "[ColorMNet] 2. Set CUDA_HOME environment variable\n"
|
||||
feedback += "[ColorMNet] Example: C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v12.8\n"
|
||||
feedback += "[ColorMNet] 1. CUDA Toolkit (match your PyTorch CUDA version)\n"
|
||||
feedback += "[ColorMNet] • Download: https://developer.nvidia.com/cuda-toolkit\n"
|
||||
feedback += "[ColorMNet] • Set CUDA_HOME environment variable after install\n"
|
||||
feedback += "[ColorMNet] Example: C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v12.8\n\n"
|
||||
|
||||
if any("cl.exe" in i for i in issues):
|
||||
feedback += "[ColorMNet] 3. Install Visual Studio Build Tools\n"
|
||||
feedback += "[ColorMNet] Download: https://visualstudio.microsoft.com/downloads/\n"
|
||||
feedback += "[ColorMNet] Select: Desktop development with C++\n"
|
||||
feedback += "[ColorMNet] 2. Visual Studio Build Tools (C++ compiler)\n"
|
||||
feedback += "[ColorMNet] • Download: https://visualstudio.microsoft.com/downloads/\n"
|
||||
feedback += "[ColorMNet] • Install: 'Desktop development with C++' workload\n"
|
||||
feedback += "[ColorMNet] • Restart terminal after installation\n\n"
|
||||
|
||||
feedback += "\n[ColorMNet] Currently using fallback mode (20-40% slower, same quality)\n"
|
||||
feedback += "[ColorMNet] ───────────────────────────────────────────────────────────────\n"
|
||||
feedback += "[ColorMNet] Currently using PyTorch fallback (20-40% slower, same quality)\n"
|
||||
feedback += "[ColorMNet] No action needed - colorization will work without compilation\n"
|
||||
feedback += "[ColorMNet] ═══════════════════════════════════════════════════════════════\n"
|
||||
return False, feedback
|
||||
|
||||
return True, "[ColorMNet] ✓ All CUDA requirements met"
|
||||
@@ -214,16 +222,31 @@ def try_install_correlation_sampler() -> Tuple[bool, str]:
|
||||
import spatial_correlation_sampler
|
||||
print("[ColorMNet] ✓ spatial_correlation_sampler already installed")
|
||||
return True, "Optimized CUDA extension available"
|
||||
except ImportError:
|
||||
except (ImportError, SyntaxError) as e:
|
||||
# SyntaxError can occur if the package is corrupted (null bytes in __init__.py)
|
||||
if isinstance(e, SyntaxError):
|
||||
print(f"[ColorMNet] ⚠ Corrupted spatial_correlation_sampler installation detected")
|
||||
print(f"[ColorMNet] Will use PyTorch fallback implementation")
|
||||
# Try to uninstall the corrupted package
|
||||
try:
|
||||
import subprocess
|
||||
subprocess.run([sys.executable, "-m", "pip", "uninstall", "spatial-correlation-sampler", "-y"],
|
||||
capture_output=True)
|
||||
except:
|
||||
pass
|
||||
pass
|
||||
|
||||
# Check requirements
|
||||
# Check requirements BEFORE attempting installation
|
||||
print("[ColorMNet] Checking compilation prerequisites...")
|
||||
requirements_met, feedback = check_cuda_requirements()
|
||||
if not requirements_met:
|
||||
print(feedback)
|
||||
return False, "Missing CUDA development tools"
|
||||
print("[ColorMNet] ⚠ Skipping compilation - prerequisites not met")
|
||||
print("[ColorMNet] ✓ Using PyTorch fallback implementation (fully functional)")
|
||||
return False, "Missing CUDA development tools - using fallback"
|
||||
|
||||
# Try to install
|
||||
print("[ColorMNet] ✓ All compilation prerequisites found")
|
||||
print("[ColorMNet] Attempting to install spatial_correlation_sampler...")
|
||||
print("[ColorMNet] This may take 2-5 minutes to compile...")
|
||||
|
||||
@@ -240,7 +263,7 @@ def try_install_correlation_sampler() -> Tuple[bool, str]:
|
||||
print("[ColorMNet] ✓ spatial_correlation_sampler installed successfully!")
|
||||
print("[ColorMNet] 🚀 Optimized CUDA mode enabled (2-5x faster)")
|
||||
return True, "Optimized CUDA extension installed"
|
||||
except ImportError:
|
||||
except (ImportError, SyntaxError):
|
||||
print("[ColorMNet] ⚠ Installation completed but import failed")
|
||||
return False, "Installation succeeded but module not importable"
|
||||
|
||||
@@ -286,7 +309,7 @@ def ensure_dependencies() -> bool:
|
||||
try:
|
||||
import spatial_correlation_sampler
|
||||
print("[ColorMNet] ✓ Pytorch-Correlation-extension available")
|
||||
except ImportError:
|
||||
except (ImportError, SyntaxError):
|
||||
print("[ColorMNet] Pytorch-Correlation-extension not found, installing...")
|
||||
print("[ColorMNet] Note: This may require C++ compiler on Windows")
|
||||
installed = install_git_dependency(
|
||||
|
||||
@@ -167,7 +167,7 @@ class MultiheadLocalAttentionV1(nn.Module):
|
||||
padding=0,
|
||||
dilation=1,
|
||||
dilation_patch=self.dilation)
|
||||
except ImportError:
|
||||
except (ImportError, SyntaxError):
|
||||
# Fallback to slower PyTorch-only implementation
|
||||
import warnings
|
||||
warnings.warn(
|
||||
@@ -292,14 +292,23 @@ class MultiheadLocalAttentionV2(nn.Module):
|
||||
self.enable_corr = enable_corr
|
||||
|
||||
if enable_corr:
|
||||
from spatial_correlation_sampler import SpatialCorrelationSampler
|
||||
self.correlation_sampler = SpatialCorrelationSampler(
|
||||
kernel_size=1,
|
||||
patch_size=self.window_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
dilation_patch=self.dilation)
|
||||
try:
|
||||
from spatial_correlation_sampler import SpatialCorrelationSampler
|
||||
self.correlation_sampler = SpatialCorrelationSampler(
|
||||
kernel_size=1,
|
||||
patch_size=self.window_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
dilation_patch=self.dilation)
|
||||
except (ImportError, SyntaxError):
|
||||
# Fallback to slower PyTorch-only implementation
|
||||
import warnings
|
||||
warnings.warn(
|
||||
"spatial_correlation_sampler not found. Using fallback implementation (20-40% slower). "
|
||||
"For better performance, install: pip install spatial-correlation-sampler",
|
||||
UserWarning
|
||||
)
|
||||
|
||||
self.projection = nn.Linear(d_model, d_model)
|
||||
|
||||
@@ -781,7 +790,7 @@ class LocalGatedPropagation(nn.Module):
|
||||
padding=0,
|
||||
dilation=1,
|
||||
dilation_patch=self.dilation)
|
||||
except ImportError:
|
||||
except (ImportError, SyntaxError):
|
||||
# Fallback to slower PyTorch-only implementation
|
||||
import warnings
|
||||
warnings.warn(
|
||||
|
||||
Reference in New Issue
Block a user