From 1c79f4ad10777468adf3974afa4627e932a7c6ba Mon Sep 17 00:00:00 2001 From: jonstreeter Date: Tue, 18 Nov 2025 14:52:31 -0600 Subject: [PATCH] Improve installer messages and error handling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- .gitignore | 8 ++++++ auto_installer.py | 28 ------------------ colormnet/downloader.py | 55 +++++++++++++++++++++++++----------- colormnet_model/attention.py | 29 ++++++++++++------- 4 files changed, 66 insertions(+), 54 deletions(-) diff --git a/.gitignore b/.gitignore index 0c9aa76..44932c6 100644 --- a/.gitignore +++ b/.gitignore @@ -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 diff --git a/auto_installer.py b/auto_installer.py index 94a6566..a343bda 100644 --- a/auto_installer.py +++ b/auto_installer.py @@ -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 diff --git a/colormnet/downloader.py b/colormnet/downloader.py index cfde529..6d1fcf6 100644 --- a/colormnet/downloader.py +++ b/colormnet/downloader.py @@ -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( diff --git a/colormnet_model/attention.py b/colormnet_model/attention.py index 7caa7a6..1e1c73d 100644 --- a/colormnet_model/attention.py +++ b/colormnet_model/attention.py @@ -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(