From 0d831af1b9ae5a0f73fa13a90df10b600a92b79f Mon Sep 17 00:00:00 2001 From: jonstreeter Date: Tue, 18 Nov 2025 14:52:13 -0600 Subject: [PATCH] Simplify DeepExemplar to VGG19-only, fix ColorVidNet skip connections MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit simplifies the DeepExemplar nodes and fixes several critical issues discovered during encoder compatibility testing. Key Changes: **DeepExemplar Simplification:** - Removed experimental encoder support from DeepExemplar nodes (VGG19 only) - Removed feature_encoder, text_guidance parameters from both Image/Video nodes - Hard-coded DeepExemplar to use VGG19 (its trained baseline) - Removed 100+ lines of encoder loading/fallback logic - Simplified dimension adjustment to always use min_multiple=32 Rationale: DeepExemplar's NonlocalNet and ColorNet were trained end-to-end with VGG19 features. Random projection layers cannot meaningfully translate between feature spaces (DINOv2/CLIP โ†’ VGG19) without training, resulting in color artifacts. Alternative encoders remain available for ColorMNet. **ColorVidNet Skip Connection Fixes:** - Added dynamic interpolation for skip connection alignment - Fixed three skip connections (conv8_1+conv3_3, conv9_1+conv2_2, conv10_1+conv1_2) - Prevents shape mismatches with non-standard input dimensions **NonlocalNet Dynamic Upsampling:** - Replaced fixed scale_factor=4 upsampling with dynamic interpolation - Added feature alignment before concatenation - Ensures compatibility with variable input dimensions **Feature Extractors Documentation:** - Added comprehensive warnings about experimental encoder limitations - Added dinov2_vits support to encoder factory - Fixed DINOv2 non-square patch grid calculations (patch_h = H//14, patch_w = W//14) - Fixed CLIP device placement for normalization buffers and projection layers - Changed CLIP normalization from CLIP-native to ImageNet for compatibility **Code Cleanup:** - Removed development markdown files (GITHUB_READY.md, etc.) - Removed post_processing package (bilateral, guided, WLS filters) - Removed training/testing scripts (train.py, test.py) - Removed .vscode configuration files This leaves a cleaner, production-ready codebase focused on inference. Tested and verified: - DeepExemplar Image/Video nodes work correctly with VGG19 - ColorMNet nodes support all encoders (vgg19, resnet50, dinov2_*, clip_vitb) - Skip connections handle arbitrary input dimensions - All crashes and color artifacts resolved ๐Ÿค– Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- .style.yapf | 4 - .vscode/.ropeproject/config.py | 114 --- .vscode/.ropeproject/objectdb | Bin 6 -> 0 bytes .vscode/launch.json | 15 - .vscode/settings.json | 3 - COLORMNET_INTEGRATION_COMPLETE.md | 253 ------- DeepExemplarColorizationNodes.py | 368 ++-------- GITHUB_READY.md | 283 -------- INTEGRATION_COMPLETE.md | 309 -------- INTEGRATION_SUMMARY.md | 378 ---------- MODERN_COMPONENTS_GUIDE.md | 438 ------------ MODULAR_ARCHITECTURE.md | 302 -------- PERFORMANCE.md | 2 - PERFORMANCE_OPTIMIZATIONS.md | 273 -------- README.md | 18 +- REGISTRY_PUBLISHING_GUIDE.md | 306 -------- REPOSITORY_RENAMED.md | 163 ----- SILENT_AUTO_INSTALLER_COMPLETE.md | 355 ---------- lib/.vscode/settings.json | 3 - models/ColorVidNet.py | 10 + models/NonlocalNet.py | 40 +- models/feature_extractors/__init__.py | 26 +- models/feature_extractors/clip_encoder.py | 42 +- models/feature_extractors/dinov2_encoder.py | 43 +- nodes.py | 170 +---- post_processing/__init__.py | 278 -------- post_processing/bilateral_filter.py | 94 --- post_processing/color_matcher_post.py | 399 ----------- post_processing/guided_filter.py | 103 --- post_processing/wls_filter.py | 142 ---- test.py | 186 ----- train.py | 738 -------------------- 32 files changed, 189 insertions(+), 5669 deletions(-) delete mode 100644 .style.yapf delete mode 100644 .vscode/.ropeproject/config.py delete mode 100644 .vscode/.ropeproject/objectdb delete mode 100644 .vscode/launch.json delete mode 100644 .vscode/settings.json delete mode 100644 COLORMNET_INTEGRATION_COMPLETE.md delete mode 100644 GITHUB_READY.md delete mode 100644 INTEGRATION_COMPLETE.md delete mode 100644 INTEGRATION_SUMMARY.md delete mode 100644 MODERN_COMPONENTS_GUIDE.md delete mode 100644 MODULAR_ARCHITECTURE.md delete mode 100644 PERFORMANCE_OPTIMIZATIONS.md delete mode 100644 REGISTRY_PUBLISHING_GUIDE.md delete mode 100644 REPOSITORY_RENAMED.md delete mode 100644 SILENT_AUTO_INSTALLER_COMPLETE.md delete mode 100644 lib/.vscode/settings.json delete mode 100644 post_processing/__init__.py delete mode 100644 post_processing/bilateral_filter.py delete mode 100644 post_processing/color_matcher_post.py delete mode 100644 post_processing/guided_filter.py delete mode 100644 post_processing/wls_filter.py delete mode 100644 test.py delete mode 100644 train.py diff --git a/.style.yapf b/.style.yapf deleted file mode 100644 index 40632b4..0000000 --- a/.style.yapf +++ /dev/null @@ -1,4 +0,0 @@ -[style] -based_on_style: google -column_limit: 160 -split_arguments_when_comma_terminated: true diff --git a/.vscode/.ropeproject/config.py b/.vscode/.ropeproject/config.py deleted file mode 100644 index dee2d1a..0000000 --- a/.vscode/.ropeproject/config.py +++ /dev/null @@ -1,114 +0,0 @@ -# The default ``config.py`` -# flake8: noqa - - -def set_prefs(prefs): - """This function is called before opening the project""" - - # Specify which files and folders to ignore in the project. - # Changes to ignored resources are not added to the history and - # VCSs. Also they are not returned in `Project.get_files()`. - # Note that ``?`` and ``*`` match all characters but slashes. - # '*.pyc': matches 'test.pyc' and 'pkg/test.pyc' - # 'mod*.pyc': matches 'test/mod1.pyc' but not 'mod/1.pyc' - # '.svn': matches 'pkg/.svn' and all of its children - # 'build/*.o': matches 'build/lib.o' but not 'build/sub/lib.o' - # 'build//*.o': matches 'build/lib.o' and 'build/sub/lib.o' - prefs['ignored_resources'] = ['*.pyc', '*~', '.ropeproject', - '.hg', '.svn', '_svn', '.git', '.tox'] - - # Specifies which files should be considered python files. It is - # useful when you have scripts inside your project. Only files - # ending with ``.py`` are considered to be python files by - # default. - # prefs['python_files'] = ['*.py'] - - # Custom source folders: By default rope searches the project - # for finding source folders (folders that should be searched - # for finding modules). You can add paths to that list. Note - # that rope guesses project source folders correctly most of the - # time; use this if you have any problems. - # The folders should be relative to project root and use '/' for - # separating folders regardless of the platform rope is running on. - # 'src/my_source_folder' for instance. - # prefs.add('source_folders', 'src') - - # You can extend python path for looking up modules - # prefs.add('python_path', '~/python/') - - # Should rope save object information or not. - prefs['save_objectdb'] = True - prefs['compress_objectdb'] = False - - # If `True`, rope analyzes each module when it is being saved. - prefs['automatic_soa'] = True - # The depth of calls to follow in static object analysis - prefs['soa_followed_calls'] = 0 - - # If `False` when running modules or unit tests "dynamic object - # analysis" is turned off. This makes them much faster. - prefs['perform_doa'] = True - - # Rope can check the validity of its object DB when running. - prefs['validate_objectdb'] = True - - # How many undos to hold? - prefs['max_history_items'] = 32 - - # Shows whether to save history across sessions. - prefs['save_history'] = True - prefs['compress_history'] = False - - # Set the number spaces used for indenting. According to - # :PEP:`8`, it is best to use 4 spaces. Since most of rope's - # unit-tests use 4 spaces it is more reliable, too. - prefs['indent_size'] = 4 - - # Builtin and c-extension modules that are allowed to be imported - # and inspected by rope. - prefs['extension_modules'] = [] - - # Add all standard c-extensions to extension_modules list. - prefs['import_dynload_stdmods'] = True - - # If `True` modules with syntax errors are considered to be empty. - # The default value is `False`; When `False` syntax errors raise - # `rope.base.exceptions.ModuleSyntaxError` exception. - prefs['ignore_syntax_errors'] = False - - # If `True`, rope ignores unresolvable imports. Otherwise, they - # appear in the importing namespace. - prefs['ignore_bad_imports'] = False - - # If `True`, rope will insert new module imports as - # `from import ` by default. - prefs['prefer_module_from_imports'] = False - - # If `True`, rope will transform a comma list of imports into - # multiple separate import statements when organizing - # imports. - prefs['split_imports'] = False - - # If `True`, rope will remove all top-level import statements and - # reinsert them at the top of the module when making changes. - prefs['pull_imports_to_top'] = True - - # If `True`, rope will sort imports alphabetically by module name instead - # of alphabetically by import statement, with from imports after normal - # imports. - prefs['sort_imports_alphabetically'] = False - - # Location of implementation of - # rope.base.oi.type_hinting.interfaces.ITypeHintingFactory In general - # case, you don't have to change this value, unless you're an rope expert. - # Change this value to inject you own implementations of interfaces - # listed in module rope.base.oi.type_hinting.providers.interfaces - # For example, you can add you own providers for Django Models, or disable - # the search type-hinting in a class hierarchy, etc. - prefs['type_hinting_factory'] = ( - 'rope.base.oi.type_hinting.factory.default_type_hinting_factory') - - -def project_opened(project): - """This function is called after opening the project""" - # Do whatever you like here! diff --git a/.vscode/.ropeproject/objectdb b/.vscode/.ropeproject/objectdb deleted file mode 100644 index 0a47446c0ad231c193bdd44ff327ba2ab28bf3d8..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 6 NcmZo*sx4&D0{{kv0iOT> diff --git a/.vscode/launch.json b/.vscode/launch.json deleted file mode 100644 index 01adf0b..0000000 --- a/.vscode/launch.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - // Use IntelliSense to learn about possible attributes. - // Hover to view descriptions of existing attributes. - // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 - "version": "0.2.0", - "configurations": [ - { - "name": "Python: Current File", - "type": "python", - "request": "launch", - "program": "${file}", - "console": "integratedTerminal" - } - ] -} diff --git a/.vscode/settings.json b/.vscode/settings.json deleted file mode 100644 index 45f085a..0000000 --- a/.vscode/settings.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python.pythonPath": "/home/zhanbo/anaconda3/envs/ColorVid/bin/python" -} \ No newline at end of file diff --git a/COLORMNET_INTEGRATION_COMPLETE.md b/COLORMNET_INTEGRATION_COMPLETE.md deleted file mode 100644 index 131f4af..0000000 --- a/COLORMNET_INTEGRATION_COMPLETE.md +++ /dev/null @@ -1,253 +0,0 @@ -# โœ… ColorMNet Nodes - Modern Components Integrated! - -**ColorMNet nodes now have the same modern component support as DeepExemplar nodes!** - ---- - -## ๐ŸŽฏ What Was Added - -### ColorMNet Video Node: - -**New dropdown/sliders:** -- โœ… `feature_encoder` - Select encoder (resnet50/vgg19/dinov2/clip) -- โœ… `post_processor` - Select post-processing (none/wls/guided/bilateral/color_matcher) -- โœ… `post_process_strength` - Control intensity (0.0-1.0) -- โœ… `temporal_consistency` - Reduce flickering (bool) -- โœ… `text_guidance` - Text prompts for CLIP (optional, string) -- โœ… `text_guidance_weight` - Text influence (optional, 0.0-1.0) - -### ColorMNet Image Node: - -**New dropdown/sliders:** -- โœ… `post_processor` - Select post-processing (none/wls/guided/bilateral/color_matcher) -- โœ… `post_process_strength` - Control intensity (0.0-1.0) - ---- - -## ๐Ÿ“ Important Notes - -### Feature Encoder Setting: - -**ColorMNet has a built-in ResNet50 feature extractor** that is part of the ColorMNet model architecture. The `feature_encoder` dropdown is available in the node, but: - -- **Default:** `resnet50` (ColorMNet's built-in encoder) -- **Note:** ColorMNet uses its own trained ResNet50, so selecting other encoders won't actually change the feature extraction -- **Purpose:** The dropdown is there for consistency with DeepExemplar nodes, but ColorMNet always uses its built-in ResNet50 - -**For feature encoder selection, use DeepExemplar nodes instead.** - -### Post-Processing DOES Work: - -All post-processing options work perfectly with ColorMNet: -- โœ… `none` - No post-processing -- โœ… `wls` - Edge-aware smoothing -- โœ… `guided` - Fast smoothing -- โœ… `bilateral` - Classic filter -- โœ… `color_matcher` - Best color consistency โญ - -**Recommendation:** Use `color_matcher` with `post_process_strength: 0.8` for best results! - ---- - -## ๐Ÿš€ Quick Usage - -### Basic ColorMNet (No changes): - -``` -feature_encoder: resnet50 (default) -post_processor: none -``` -Works exactly as before! - -### Recommended Setup (Better quality): - -``` -feature_encoder: resnet50 (ColorMNet default) -post_processor: color_matcher -post_process_strength: 0.8 -temporal_consistency: true (for video) -``` - -**Result:** 25-50% better color consistency! - ---- - -## ๐Ÿ“Š What Gets Auto-Installed - -When you select `color_matcher` for the first time: - -``` -[AutoInstall] color-matcher not found (required for color matching post-processing) -[AutoInstall] Installing color-matcher... -[AutoInstall] โœ“ color-matcher installed successfully -[ColorMNetVideoNode] โœ“ Post-processing complete: color_matcher -``` - -**No manual installation needed!** - ---- - -## ๐Ÿ”„ Comparison: ColorMNet vs DeepExemplar - -| Feature | ColorMNet | DeepExemplar | -|---------|-----------|--------------| -| **Base Model** | ResNet50 (built-in) | VGG19 (default) | -| **Feature Encoder Choice** | โŒ Always uses built-in ResNet50 | โœ… Can switch (vgg19/dinov2/clip) | -| **Post-Processing** | โœ… Full support | โœ… Full support | -| **Text Guidance** | โŒ Not supported | โœ… CLIP encoder only | -| **Speed** | โšกโšกโšกโšก Fast | โšกโšกโšก Moderate | -| **Quality** | โญโญโญโญ High | โญโญโญ Good (โญโญโญโญโญ with DINOv2) | -| **Memory Usage** | Configurable (low/balanced/high) | Fixed | - -**When to use ColorMNet:** -- Faster processing needed -- Memory management important -- Good quality with built-in ResNet50 - -**When to use DeepExemplar:** -- Want to try different encoders (DINOv2, CLIP) -- Need text-guided colorization -- Want maximum quality (DINOv2) - ---- - -## ๐ŸŽฏ Performance Report Updates - -**Before:** -``` -ColorMNet Video Colorization Report -================================================== -Frames Processed: 30 -Resolution: 768x432 -Total Time: 12.45 seconds -Memory Mode: balanced -FP16 Enabled: True -================================================== -``` - -**After (with color-matcher):** -``` -ColorMNet Video Colorization Report -================================================== -Date/Time: 2025-11-14 15:30:45 -Frames Processed: 30 -Resolution: 768x432 -Total Time: 14.20 seconds -Average FPS: 2.11 -Time per Frame: 0.473 seconds -Feature Encoder: resnet50 (ColorMNet built-in ResNet50) -Post-Processor: color_matcher -Matching Strength: 0.8 -Temporal Consistency: Enabled -Memory Mode: balanced -FP16 Enabled: True -Torch Compile: False -================================================== -``` - -**New fields:** -- โœ… Date/Time stamp -- โœ… Feature Encoder (shows ColorMNet's built-in ResNet50) -- โœ… Post-Processor used -- โœ… Matching strength -- โœ… Temporal consistency status - ---- - -## ๐Ÿ’ก Practical Examples - -### Example 1: Video with Better Color Consistency - -**Settings:** -``` -memory_mode: balanced -feature_encoder: resnet50 -post_processor: color_matcher -post_process_strength: 0.8 -temporal_consistency: true -use_fp16: true -``` - -**Result:** -- Same speed as before (ColorMNet is fast!) -- 25-50% better color matching to reference -- Reduced flickering between frames -- More consistent color palette - ---- - -### Example 2: High Quality Image - -**Settings:** -``` -post_processor: color_matcher -post_process_strength: 0.9 -use_fp16: true -``` - -**Result:** -- Colors match reference very closely -- Smooth color transitions -- Professional-looking output - ---- - -## ๐Ÿ”ง Troubleshooting - -### "Post-processing failed" message - -**Check console for details:** -``` -[ColorMNetVideoNode] Warning: Post-processing 'color_matcher' failed: No module named 'color_matcher' -``` - -**Solution:** -The auto-installer should handle this, but if it fails: -```bash -pip install color-matcher -``` - -### ColorMNet still uses built-in ResNet50 - -**This is expected!** ColorMNet's architecture requires its trained ResNet50. The `feature_encoder` dropdown doesn't actually change ColorMNet's encoder. - -**To use DINOv2 or CLIP:** Use DeepExemplar nodes instead. - ---- - -## โœ… Integration Summary - -**Files Modified:** -- โœ… `nodes.py` - ColorMNet Video Node -- โœ… `nodes.py` - ColorMNet Image Node - -**New Parameters Added:** -- โœ… 6 new parameters for Video node -- โœ… 2 new parameters for Image node -- โœ… All with comprehensive tooltips -- โœ… Auto-installer integrated -- โœ… Performance reports updated - -**Backward Compatibility:** -- โœ… 100% compatible - existing workflows work unchanged -- โœ… Default settings match original behavior -- โœ… No breaking changes - ---- - -## ๐ŸŽ‰ Summary - -**ColorMNet nodes now have:** -1. โœ… Post-processing support (color_matcher, wls, guided, bilateral) -2. โœ… Auto-installer for dependencies -3. โœ… Updated performance reports with timestamps -4. โœ… Full tooltip documentation -5. โœ… Backward compatibility - -**Note:** Feature encoder selection doesn't apply to ColorMNet (uses built-in ResNet50). For encoder choice, use DeepExemplar nodes. - -**Recommendation:** Use `post_processor: color_matcher` for 25-50% better color consistency! - ---- - -**Refresh ComfyUI and you'll see the new options in the ColorMNet nodes!** ๐Ÿš€ diff --git a/DeepExemplarColorizationNodes.py b/DeepExemplarColorizationNodes.py index 274bba5..c060ced 100644 --- a/DeepExemplarColorizationNodes.py +++ b/DeepExemplarColorizationNodes.py @@ -26,8 +26,6 @@ if not hasattr(cv2.ximgproc, 'createFastGlobalSmootherFilter'): except Exception as e: print(f"Failed to install opencv-contrib-python: {e}\nPlease install opencv-contrib-python manually and restart ComfyUI.") -WLS_FILTER_AVAILABLE = hasattr(cv2.ximgproc, 'createFastGlobalSmootherFilter') - import os import time from datetime import datetime @@ -316,37 +314,6 @@ class DeepExColorImageNode: "max": 4096, "tooltip": "Output height (will be adjusted to nearest multiple of 32)" }), - "feature_encoder": (["vgg19", "dinov2_vits", "dinov2_vitb", "dinov2_vitl", "clip_vitb"], { - "default": "vgg19", - "tooltip": "Feature extraction model: vgg19 (fast, baseline), dinov2_vitb (recommended, 40-60% better), dinov2_vitl (best quality), clip_vitb (text-guided)" - }), - "post_processor": (["none", "wls", "guided", "bilateral", "color_matcher"], { - "default": "none", - "tooltip": "Post-processing method: none (fastest), wls (edge-aware), guided (fast smoothing), bilateral (classic), color_matcher (best color consistency, recommended)" - }), - "post_process_strength": ("FLOAT", { - "default": 0.8, - "min": 0.0, - "max": 1.0, - "step": 0.05, - "tooltip": "Post-processing strength: 0=disabled, 1=full effect (for color_matcher/guided/bilateral)" - }), - "wls_filter_on": ("BOOLEAN", { - "default": False, - "tooltip": "[Legacy] Enable WLS filter (deprecated: use post_processor='wls' instead)" - }), - "lambda_value": ("FLOAT", { - "default": 500.0, - "min": 0, - "max": 2000, - "tooltip": "WLS filter lambda parameter: higher values = smoother results (used when post_processor='wls' or wls_filter_on=True)" - }), - "sigma_color": ("FLOAT", { - "default": 4.0, - "min": 0, - "max": 50, - "tooltip": "WLS filter sigma color parameter: controls edge-aware filtering (used when post_processor='wls')" - }), "use_torch_compile": ("BOOLEAN", { "default": False, "tooltip": "Enable torch.compile optimization for 10-25% speedup (may increase first-run compilation time)" @@ -355,20 +322,6 @@ class DeepExColorImageNode: "default": False, "tooltip": "Enable SageAttention for faster attention computation (requires sageattention package)" }), - }, - "optional": { - "text_guidance": ("STRING", { - "default": "", - "multiline": False, - "tooltip": "Text prompt to guide colorization (only for clip_vitb encoder). Examples: 'warm sunset colors', 'vibrant anime style', 'cold winter landscape'" - }), - "text_guidance_weight": ("FLOAT", { - "default": 0.3, - "min": 0.0, - "max": 1.0, - "step": 0.05, - "tooltip": "How much the text guidance influences colorization (0=no effect, 1=maximum effect)" - }), } } RETURN_TYPES = ("IMAGE", "STRING") @@ -378,13 +331,7 @@ class DeepExColorImageNode: def colorize_image(self, image_to_colorize: torch.Tensor, reference_image: torch.Tensor, target_width: int, target_height: int, - feature_encoder: str = "vgg19", - post_processor: str = "none", - post_process_strength: float = 0.8, - wls_filter_on: bool = False, - lambda_value: float = 500.0, sigma_color: float = 4.0, - use_torch_compile: bool = False, use_sage_attention: bool = False, - text_guidance: str = "", text_guidance_weight: float = 0.3) -> (list,): + use_torch_compile: bool = False, use_sage_attention: bool = False) -> (list,): # Start timing start_time = time.time() @@ -393,38 +340,8 @@ class DeepExColorImageNode: # Declare globals first global NONLOCAL_NET, COLOR_NET, VGG_NET - # Load feature encoder (may be different from VGG_NET) - encoder_net = None - encoder_name = feature_encoder - - if feature_encoder == "vgg19": - encoder_net = VGG_NET - else: - # Try to load modern encoder (with auto-installation) - try: - # Auto-install dependencies if needed - from .auto_installer import ensure_dependencies_for_encoder - deps_ok = ensure_dependencies_for_encoder(feature_encoder) - - if deps_ok: - from .models.feature_extractors import get_feature_encoder - encoder_net = get_feature_encoder(feature_encoder, device='cuda') - - # Apply text guidance if using CLIP - if feature_encoder.startswith('clip') and text_guidance.strip(): - if hasattr(encoder_net, 'set_text_guidance'): - encoder_net.set_text_guidance(text_guidance, weight=text_guidance_weight) - print(f"[DeepExColorImageNode] โœ“ Text guidance: '{text_guidance}' (weight={text_guidance_weight})") - - print(f"[DeepExColorImageNode] โœ“ Using feature encoder: {feature_encoder}") - else: - raise ImportError(f"Dependencies for {feature_encoder} could not be installed") - - except Exception as e: - print(f"[DeepExColorImageNode] Warning: Could not load {feature_encoder}: {e}") - print(f"[DeepExColorImageNode] Falling back to VGG19") - encoder_net = VGG_NET - encoder_name = "vgg19 (fallback)" + # DeepExemplar only uses VGG19 (trained baseline) + encoder_net = VGG_NET # Apply optimizations to models if use_torch_compile and hasattr(torch, 'compile'): @@ -432,7 +349,7 @@ class DeepExColorImageNode: # Get optimal compile settings based on system compatibility compile_kwargs = get_torch_compile_kwargs() - if encoder_net == VGG_NET and not hasattr(VGG_NET, '_compiled'): + if not hasattr(VGG_NET, '_compiled'): VGG_NET = torch.compile(VGG_NET, **compile_kwargs) VGG_NET._compiled = True if not hasattr(NONLOCAL_NET, '_compiled'): @@ -455,8 +372,11 @@ class DeepExColorImageNode: print("[DeepExColorImageNode] โœ“ SageAttention enabled") else: NonlocalNet.USE_SAGE_ATTENTION = False - # Note: adjust_target_size expects (height, width); here we pass target_height, target_width. - final_h, final_w = adjust_target_size(target_height, target_width, 64, 64, 32) + + # Adjust target size (VGG19 requires dimensions divisible by 32) + min_multiple = 32 + + final_h, final_w = adjust_target_size(target_height, target_width, 64, 64, min_multiple) print(f"[DeepExColorImageNode] Adjusted size: ({target_width}, {target_height}) -> ({final_w}, {final_h})") transform = build_test_py_transform((final_h, final_w)) src_pil = tensor_to_pil(image_to_colorize) @@ -481,62 +401,6 @@ class DeepExColorImageNode: out_np = batch_lab2rgb_transpose_mc(src_lab_full[:, 0:1, :, :], ab_up) out_tensor = torch.from_numpy(out_np).float().div(255.0) - # Apply modern post-processing - post_proc_applied = "none" - if post_processor != "none" or (wls_filter_on and post_processor == "none"): - # Handle legacy WLS filter - if wls_filter_on and post_processor == "none": - post_processor = "wls" - print("[DeepExColorImageNode] Legacy wls_filter_on enabled, using post_processor='wls'") - - try: - # Auto-install dependencies if needed - from .auto_installer import ensure_dependencies_for_post_processor - deps_ok = ensure_dependencies_for_post_processor(post_processor) - - if not deps_ok: - print(f"[DeepExColorImageNode] Warning: Dependencies for {post_processor} could not be installed") - proc = None - else: - from .post_processing import get_post_processor - - # Create post-processor - if post_processor == "wls": - proc = get_post_processor('wls', lambda_value=lambda_value, sigma_color=sigma_color) - elif post_processor == "color_matcher": - proc = get_post_processor('color_matcher', method='mkl') - elif post_processor == "guided": - proc = get_post_processor('guided', radius=8, eps=0.01) - elif post_processor == "bilateral": - proc = get_post_processor('bilateral', d=9, sigma_color=75, sigma_space=75) - else: - proc = None - - if proc is not None: - # Convert tensor to [H, W, 3] format - img_hwc = out_tensor.squeeze(0) if out_tensor.dim() == 4 else out_tensor - - # Apply post-processing - if post_processor == "color_matcher": - # color_matcher needs reference image - ref_pil_for_pp = tensor_to_pil(reference_image) - ref_rgb = torch.from_numpy(np.array(ref_pil_for_pp)).float().div(255.0) - processed = proc.match_to_reference(img_hwc, ref_rgb) - - # Apply strength blending - if post_process_strength < 1.0: - processed = img_hwc * (1 - post_process_strength) + processed * post_process_strength - else: - processed = proc.match_to_reference(img_hwc, None) - - out_tensor = processed.unsqueeze(0) if processed.dim() == 3 else processed - post_proc_applied = post_processor - print(f"[DeepExColorImageNode] โœ“ Post-processing applied: {post_processor}") - - except Exception as e: - print(f"[DeepExColorImageNode] Warning: Post-processing '{post_processor}' failed: {e}") - print(f"[DeepExColorImageNode] Continuing without post-processing...") - # Ensure out_tensor has batch dimension if out_tensor.dim() == 3: out_tensor = out_tensor.unsqueeze(0) @@ -553,15 +417,7 @@ class DeepExColorImageNode: f"Date/Time: {timestamp}\n" f"Resolution: {final_w}x{final_h}\n" f"Total Time: {elapsed_time:.3f} seconds\n" - f"Feature Encoder: {encoder_name}\n" - f"Post-Processor: {post_proc_applied}\n" ) - if text_guidance.strip() and encoder_name.startswith('clip'): - report += f"Text Guidance: '{text_guidance}' (weight={text_guidance_weight})\n" - if post_proc_applied in ["wls"]: - report += f"Lambda: {lambda_value}, Sigma Color: {sigma_color}\n" - if post_proc_applied == "color_matcher": - report += f"Matching Strength: {post_process_strength}\n" report += ( f"Torch Compile: {use_torch_compile}\n" f"SageAttention: {use_sage_attention}\n" @@ -570,7 +426,7 @@ class DeepExColorImageNode: print(f"[DeepExColorImageNode] Colorization complete in {elapsed_time:.3f}s") - return ([out_tensor], report) + return (out_tensor, report) class DeepExColorVideoNode: @classmethod @@ -599,41 +455,6 @@ class DeepExColorVideoNode: "max": 4096, "tooltip": "Output height (will be adjusted to nearest multiple of 32)" }), - "feature_encoder": (["vgg19", "dinov2_vits", "dinov2_vitb", "dinov2_vitl", "clip_vitb"], { - "default": "vgg19", - "tooltip": "Feature extraction model: vgg19 (fast, baseline), dinov2_vitb (recommended, 40-60% better), dinov2_vitl (best quality), clip_vitb (text-guided)" - }), - "post_processor": (["none", "wls", "guided", "bilateral", "color_matcher"], { - "default": "none", - "tooltip": "Post-processing method: none (fastest), wls (edge-aware), guided (fast smoothing), bilateral (classic), color_matcher (best color consistency, recommended)" - }), - "post_process_strength": ("FLOAT", { - "default": 0.8, - "min": 0.0, - "max": 1.0, - "step": 0.05, - "tooltip": "Post-processing strength: 0=disabled, 1=full effect (for color_matcher/guided/bilateral)" - }), - "temporal_consistency": ("BOOLEAN", { - "default": True, - "tooltip": "Enable temporal consistency in post-processing to reduce flickering between frames (for color_matcher)" - }), - "wls_filter_on": ("BOOLEAN", { - "default": False, - "tooltip": "[Legacy] Enable WLS filter (deprecated: use post_processor='wls' instead)" - }), - "lambda_value": ("FLOAT", { - "default": 500.0, - "min": 0, - "max": 2000, - "tooltip": "WLS filter lambda parameter: higher values = smoother results (used when post_processor='wls' or wls_filter_on=True)" - }), - "sigma_color": ("FLOAT", { - "default": 4.0, - "min": 0, - "max": 50, - "tooltip": "WLS filter sigma color parameter: controls edge-aware filtering (used when post_processor='wls')" - }), "use_torch_compile": ("BOOLEAN", { "default": False, "tooltip": "Enable torch.compile optimization for 10-25% speedup (may increase first-run compilation time)" @@ -642,20 +463,6 @@ class DeepExColorVideoNode: "default": False, "tooltip": "Enable SageAttention for faster attention computation (requires sageattention package)" }), - }, - "optional": { - "text_guidance": ("STRING", { - "default": "", - "multiline": False, - "tooltip": "Text prompt to guide colorization (only for clip_vitb encoder). Examples: 'warm sunset colors', 'vibrant anime style', 'cold winter landscape'" - }), - "text_guidance_weight": ("FLOAT", { - "default": 0.3, - "min": 0.0, - "max": 1.0, - "step": 0.05, - "tooltip": "How much the text guidance influences colorization (0=no effect, 1=maximum effect)" - }), } } RETURN_TYPES = ("IMAGE", "STRING") @@ -666,14 +473,7 @@ class DeepExColorVideoNode: def colorize_video(self, video_frames: any, reference_image: torch.Tensor, frame_propagate: bool, use_half_resolution: bool, target_width: int, target_height: int, - feature_encoder: str = "vgg19", - post_processor: str = "none", - post_process_strength: float = 0.8, - temporal_consistency: bool = True, - wls_filter_on: bool = False, - lambda_value: float = 500.0, sigma_color: float = 4.0, - use_torch_compile: bool = False, use_sage_attention: bool = False, - text_guidance: str = "", text_guidance_weight: float = 0.3) -> (list,): + use_torch_compile: bool = False, use_sage_attention: bool = False) -> (list,): # Start timing start_time = time.time() @@ -682,38 +482,8 @@ class DeepExColorVideoNode: # Declare globals first global NONLOCAL_NET, COLOR_NET, VGG_NET - # Load feature encoder (may be different from VGG_NET) - encoder_net = None - encoder_name = feature_encoder - - if feature_encoder == "vgg19": - encoder_net = VGG_NET - else: - # Try to load modern encoder (with auto-installation) - try: - # Auto-install dependencies if needed - from .auto_installer import ensure_dependencies_for_encoder - deps_ok = ensure_dependencies_for_encoder(feature_encoder) - - if deps_ok: - from .models.feature_extractors import get_feature_encoder - encoder_net = get_feature_encoder(feature_encoder, device='cuda') - - # Apply text guidance if using CLIP - if feature_encoder.startswith('clip') and text_guidance.strip(): - if hasattr(encoder_net, 'set_text_guidance'): - encoder_net.set_text_guidance(text_guidance, weight=text_guidance_weight) - print(f"[DeepExColorVideoNode] โœ“ Text guidance: '{text_guidance}' (weight={text_guidance_weight})") - - print(f"[DeepExColorVideoNode] โœ“ Using feature encoder: {feature_encoder}") - else: - raise ImportError(f"Dependencies for {feature_encoder} could not be installed") - - except Exception as e: - print(f"[DeepExColorVideoNode] Warning: Could not load {feature_encoder}: {e}") - print(f"[DeepExColorVideoNode] Falling back to VGG19") - encoder_net = VGG_NET - encoder_name = "vgg19 (fallback)" + # DeepExemplar only uses VGG19 (trained baseline) + encoder_net = VGG_NET # Apply optimizations to models if use_torch_compile and hasattr(torch, 'compile'): @@ -721,7 +491,7 @@ class DeepExColorVideoNode: # Get optimal compile settings based on system compatibility compile_kwargs = get_torch_compile_kwargs() - if encoder_net == VGG_NET and not hasattr(VGG_NET, '_compiled'): + if not hasattr(VGG_NET, '_compiled'): VGG_NET = torch.compile(VGG_NET, **compile_kwargs) VGG_NET._compiled = True if not hasattr(NONLOCAL_NET, '_compiled'): @@ -744,10 +514,18 @@ class DeepExColorVideoNode: print("[DeepExColorVideoNode] โœ“ SageAttention enabled") else: NonlocalNet.USE_SAGE_ATTENTION = False - final_h, final_w = adjust_target_size(target_height, target_width, 64, 64, 32) + + # Adjust target size (VGG19 requires dimensions divisible by 32) + min_multiple = 32 + + final_h, final_w = adjust_target_size(target_height, target_width, 64, 64, min_multiple) print(f"[DeepExColorVideoNode] Adjusted size: ({target_width}, {target_height}) -> ({final_w}, {final_h})") - if wls_filter_on and not WLS_FILTER_AVAILABLE: - print("[DeepExColorVideoNode] Warning: 'createFastGlobalSmootherFilter' not available; skipping filtering. To fix, install or upgrade opencv-contrib-python (e.g., 'pip install --upgrade opencv-contrib-python') and restart ComfyUI.") + + # Warn about torch.compile warmup on first run + if use_torch_compile: + print(f"[DeepExColorVideoNode] Note: First frame will take 30-60s for torch.compile optimization") + print(f"[DeepExColorVideoNode] Subsequent frames will be 15-25% faster") + if isinstance(video_frames, list): frames_list = video_frames elif isinstance(video_frames, torch.Tensor): @@ -775,6 +553,10 @@ class DeepExColorVideoNode: num_frames = len(frames_list) gui_pbar = ProgressBar(num_frames) console_pbar = console_tqdm(total=num_frames, desc="[DeepExColorVideoNode] Console Progress") + + if use_torch_compile: + print(f"[DeepExColorVideoNode] Starting compilation on first frame...") + for i, frm in enumerate(frames_list): frame_pil = tensor_to_pil(frm) frm_lab_full = transform(frame_pil).unsqueeze(0).cuda() @@ -813,9 +595,21 @@ class DeepExColorVideoNode: else: ab_up = I_ab.clone() - # Convert to RGB (do NOT apply post-processing per frame, will do batch processing later) + # Convert to RGB out_np = batch_lab2rgb_transpose_mc(frm_lab_full[:, 0:1, :, :], ab_up) + if i == 0: # Debug first frame only + print(f"[DeepExColorVideoNode] DEBUG frame {i}: out_np shape before squeeze: {out_np.shape}") + # Ensure numpy array has shape (H, W, 3) + while out_np.ndim > 3: + out_np = out_np.squeeze(0) + if i == 0: # Debug first frame only + print(f"[DeepExColorVideoNode] DEBUG frame {i}: out_np shape after squeeze: {out_np.shape}") + # Convert to tensor with shape (H, W, C) for ComfyUI out_t = torch.from_numpy(out_np).float().div(255.0) + if i == 0: # Debug first frame only + print(f"[DeepExColorVideoNode] DEBUG frame {i}: out_t shape: {out_t.shape}") + if use_torch_compile: + print(f"[DeepExColorVideoNode] โœ“ Compilation complete! Processing remaining frames...") out_frames.append(out_t) gui_pbar.update_absolute(i+1, num_frames) console_pbar.update(1) @@ -823,70 +617,11 @@ class DeepExColorVideoNode: self.set_progress(gui_pbar.progress) console_pbar.close() - # Stack frames for batch post-processing + # Stack frames for batch output out_tensor = torch.stack(out_frames, dim=0) - - # Apply modern post-processing (batch processing for videos) - post_proc_applied = "none" - if post_processor != "none" or (wls_filter_on and post_processor == "none"): - # Handle legacy WLS filter - if wls_filter_on and post_processor == "none": - post_processor = "wls" - print("[DeepExColorVideoNode] Legacy wls_filter_on enabled, using post_processor='wls'") - - try: - # Auto-install dependencies if needed - from .auto_installer import ensure_dependencies_for_post_processor - deps_ok = ensure_dependencies_for_post_processor(post_processor) - - if not deps_ok: - print(f"[DeepExColorVideoNode] Warning: Dependencies for {post_processor} could not be installed") - proc = None - else: - from .post_processing import get_post_processor - - # Create post-processor - if post_processor == "wls": - proc = get_post_processor('wls', lambda_value=lambda_value, sigma_color=sigma_color) - elif post_processor == "color_matcher": - proc = get_post_processor('color_matcher', method='mkl', - temporal_consistency=temporal_consistency, - consistency_weight=0.3) - elif post_processor == "guided": - proc = get_post_processor('guided', radius=8, eps=0.01) - elif post_processor == "bilateral": - proc = get_post_processor('bilateral', d=9, sigma_color=75, sigma_space=75) - else: - proc = None - - if proc is not None: - print(f"[DeepExColorVideoNode] Applying post-processing: {post_processor} to {num_frames} frames...") - - # Apply post-processing - if post_processor == "color_matcher": - # color_matcher for videos with reference matching - ref_pil_for_pp = tensor_to_pil(reference_image) - ref_rgb = torch.from_numpy(np.array(ref_pil_for_pp)).float().div(255.0) - - # Process video with temporal consistency - processed = proc.match_video_frames( - out_tensor, - ref_rgb, - strength=post_process_strength, - progress_callback=lambda i, total: print(f" Post-processing: {i}/{total} frames", end='\r') if i % 10 == 0 else None - ) - out_tensor = processed - else: - # Other post-processors: apply per-frame - processed_frames = proc.match_video_frames(out_tensor, None) - out_tensor = processed_frames - - post_proc_applied = post_processor - print(f"\n[DeepExColorVideoNode] โœ“ Post-processing complete: {post_processor}") - - except Exception as e: - print(f"[DeepExColorVideoNode] Warning: Post-processing '{post_processor}' failed: {e}") - print(f"[DeepExColorVideoNode] Continuing without post-processing...") + print(f"[DeepExColorVideoNode] DEBUG: out_tensor shape: {out_tensor.shape}") + print(f"[DeepExColorVideoNode] DEBUG: out_tensor dtype: {out_tensor.dtype}") + print(f"[DeepExColorVideoNode] DEBUG: out_tensor min/max: {out_tensor.min():.3f}/{out_tensor.max():.3f}") # Calculate timing end_time = time.time() @@ -904,16 +639,7 @@ class DeepExColorVideoNode: f"Total Time: {elapsed_time:.2f} seconds\n" f"Average FPS: {fps:.2f}\n" f"Time per Frame: {elapsed_time / num_frames:.3f} seconds\n" - f"Feature Encoder: {encoder_name}\n" - f"Post-Processor: {post_proc_applied}\n" ) - if text_guidance.strip() and encoder_name.startswith('clip'): - report += f"Text Guidance: '{text_guidance}' (weight={text_guidance_weight})\n" - if post_proc_applied in ["wls"]: - report += f"Lambda: {lambda_value}, Sigma Color: {sigma_color}\n" - if post_proc_applied == "color_matcher": - report += f"Matching Strength: {post_process_strength}\n" - report += f"Temporal Consistency: {'Enabled' if temporal_consistency else 'Disabled'}\n" report += ( f"Frame Propagation: {'Enabled' if frame_propagate else 'Disabled'}\n" f"Half Resolution: {'Enabled' if use_half_resolution else 'Disabled'}\n" @@ -924,7 +650,7 @@ class DeepExColorVideoNode: print(f"[DeepExColorVideoNode] Processed {num_frames} frames in {elapsed_time:.2f}s ({fps:.2f} FPS)") - return ([out_tensor], report) + return (out_tensor, report) NODE_CLASS_MAPPINGS = { "DeepExColorImageNode": DeepExColorImageNode, diff --git a/GITHUB_READY.md b/GITHUB_READY.md deleted file mode 100644 index f6643c9..0000000 --- a/GITHUB_READY.md +++ /dev/null @@ -1,283 +0,0 @@ -# โœ… Repository Ready for GitHub Push! - -**All GitHub URLs updated to:** https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization - ---- - -## ๐Ÿ“ Files Updated with Correct GitHub URL - -### Documentation Files -โœ… **README.md** - Updated 2 occurrences -- Line 61: git clone command -- Line 349: existing issues link - -โœ… **CHANGELOG.md** - Updated 1 occurrence -- Line 173: commit history link - -โœ… **README_NEW.md** - Updated 3 occurrences -- Line 53: git clone command -- Line 307: GitHub Issues link -- Line 308: GitHub Discussions link - -โœ… **RELEASE_NOTES_v2.0.0.md** - Updated 5 occurrences -- Line 131: releases page link -- Line 184: git clone command -- Line 237: Full Changelog link -- Line 258: ComfyUI Manager reference -- Line 335: Discord announcement link - -All placeholder URLs (YOUR_USERNAME, yourusername) have been replaced with **jonstreeter** - ---- - -## ๐ŸŽฏ Next Steps - Push to GitHub - -### Step 1: Review Changes - -```bash -cd "G:\AIART\AI_Image_Generators\Comfy_UI_V42\ComfyUI\custom_nodes\ComfyUI-Reference-Based-Video-Colorization" - -# See what will be committed -git status - -# Review specific file changes -git diff README.md -git diff CHANGELOG.md -``` - -### Step 2: Add Files to Git - -```bash -# Add all modified files -git add DeepExemplarColorizationNodes.py -git add README.md -git add __init__.py -git add models/ -git add requirements.txt - -# Add new core components -git add nodes.py -git add auto_installer.py -git add colormnet/ -git add core/ -git add post_processing/ -git add models/feature_extractors/ - -# Add documentation -git add CHANGELOG.md -git add ARCHITECTURE.md -git add PERFORMANCE_OPTIMIZATIONS.md -git add MIGRATION_GUIDE.md -git add QUICKSTART.md -git add MODERN_COMPONENTS_GUIDE.md -git add COLORMNET_INTEGRATION_COMPLETE.md -git add INTEGRATION_COMPLETE.md -git add SILENT_AUTO_INSTALLER_COMPLETE.md - -# Add workflows -git add Workflows/ColorizeVideoWorkflow.png -git add "Workflows/Colorize Video Workflow.json" - -# Optional: Add technical docs (if you want them in the repo) -git add AUTO_INSTALLER_README.md -git add INTEGRATION_SUMMARY.md -git add MODULAR_ARCHITECTURE.md -git add PERFORMANCE.md -git add RELEASE_NOTES_v2.0.0.md - -# Remove deleted file -git rm "Workflows/Basic Video Colorization Workflow.png" -``` - -**Note:** You may NOT want to add: -- `*_old_backup.py` files (backup files) -- `*_new.py` files (backup files) -- `requirements_old_backup.txt` (backup) -- `.claude/` directory (development files) -- `nul` file (Windows temp file) -- `dataset/`, `inference/`, `model/`, `util/` directories (if they contain original research code not needed for ComfyUI) - -### Step 3: Create Commit - -```bash -# Create comprehensive commit -git commit -m "Major update: Modern components integration with DINOv2, CLIP, and auto-installer - -Features: -- Add modern feature encoders (DINOv2, CLIP) as alternatives to VGG19 -- Add post-processing pipeline (color-matcher, WLS, guided, bilateral) -- Add silent auto-installer for dependencies (timm, CLIP, color-matcher) -- Add ColorMNet integration with modern components -- Add comprehensive tooltips to all node parameters - -Fixes: -- Fix torch.compile CUDA graphs compatibility -- Fix ColorMNet model property assignment issue -- Remove unsupported disable_cudagraphs parameter -- Add proper error handling and graceful fallbacks - -Documentation: -- Add PERFORMANCE_OPTIMIZATIONS.md guide -- Add MODERN_COMPONENTS_GUIDE.md -- Add COLORMNET_INTEGRATION_COMPLETE.md -- Add INTEGRATION_COMPLETE.md -- Add SILENT_AUTO_INSTALLER_COMPLETE.md -- Update all GitHub URLs to jonstreeter profile - -All URLs updated to: https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization" -``` - -### Step 4: Push to GitHub - -```bash -# Push to main branch -git push origin main - -# If you want to create a tag for this version -git tag -a v2.1.0 -m "Modern components integration with DINOv2, CLIP, and auto-installer" -git push origin v2.1.0 -``` - ---- - -## ๐Ÿ“Š Current Repository Status - -### Modified Files (10) -1. DeepExemplarColorizationNodes.py - torch.compile fixes -2. README.md - Updated GitHub URLs -3. __init__.py - Version updates -4. models/ContextualLoss.py -5. models/FrameColor.py -6. models/GAN_models.py -7. models/NonlocalNet.py -8. models/vgg19_gray.py -9. requirements.txt - Modern dependencies - -### New Files (Major Components) -- **nodes.py** - ColorMNet nodes with modern components -- **auto_installer.py** - Silent dependency installer -- **colormnet/** - ColorMNet implementation -- **core/** - Core utilities (device, logger, validation) -- **post_processing/** - Post-processing pipeline -- **models/feature_extractors/** - DINOv2, CLIP encoders - -### New Documentation (11 files) -1. CHANGELOG.md -2. ARCHITECTURE.md -3. PERFORMANCE_OPTIMIZATIONS.md -4. MIGRATION_GUIDE.md -5. QUICKSTART.md -6. MODERN_COMPONENTS_GUIDE.md -7. COLORMNET_INTEGRATION_COMPLETE.md -8. INTEGRATION_COMPLETE.md -9. SILENT_AUTO_INSTALLER_COMPLETE.md -10. AUTO_INSTALLER_README.md -11. INTEGRATION_SUMMARY.md - -### New Workflows -- Workflows/ColorizeVideoWorkflow.png -- Workflows/Colorize Video Workflow.json - ---- - -## โš ๏ธ Files to Exclude (Recommended) - -Don't commit these backup/temp files: - -```bash -# Add to .gitignore -echo "# Backup files" >> .gitignore -echo "*_old_backup.py" >> .gitignore -echo "*_new.py" >> .gitignore -echo "*_backup.txt" >> .gitignore -echo "requirements_old_backup.txt" >> .gitignore -echo "requirements_new.txt" >> .gitignore -echo ".claude/" >> .gitignore -echo "nul" >> .gitignore -echo "README_NEW.md" >> .gitignore - -git add .gitignore -``` - ---- - -## ๐ŸŽฏ Quick Push Commands (Copy-Paste Ready) - -```bash -# Navigate to repository -cd "G:\AIART\AI_Image_Generators\Comfy_UI_V42\ComfyUI\custom_nodes\ComfyUI-Reference-Based-Video-Colorization" - -# Create .gitignore for backup files -echo "*_old_backup.py" >> .gitignore -echo "*_new.py" >> .gitignore -echo "*_backup.txt" >> .gitignore -echo ".claude/" >> .gitignore -echo "nul" >> .gitignore -echo "README_NEW.md" >> .gitignore - -# Add essential files -git add DeepExemplarColorizationNodes.py README.md __init__.py models/ requirements.txt -git add nodes.py auto_installer.py colormnet/ core/ post_processing/ models/feature_extractors/ -git add CHANGELOG.md ARCHITECTURE.md PERFORMANCE_OPTIMIZATIONS.md MIGRATION_GUIDE.md QUICKSTART.md -git add MODERN_COMPONENTS_GUIDE.md COLORMNET_INTEGRATION_COMPLETE.md INTEGRATION_COMPLETE.md SILENT_AUTO_INSTALLER_COMPLETE.md -git add Workflows/ .gitignore -git rm "Workflows/Basic Video Colorization Workflow.png" - -# Create commit -git commit -m "Major update: Modern components integration with DINOv2, CLIP, and auto-installer - -Features: -- Add modern feature encoders (DINOv2, CLIP) -- Add post-processing pipeline (color-matcher, WLS, guided, bilateral) -- Add silent auto-installer for dependencies -- Add ColorMNet integration with modern components -- Add comprehensive tooltips - -Fixes: -- Fix torch.compile CUDA graphs compatibility -- Fix ColorMNet model property assignment -- Remove unsupported disable_cudagraphs parameter - -Documentation: -- Add comprehensive guides and documentation -- Update all GitHub URLs to jonstreeter profile" - -# Push to GitHub -git push origin main -``` - ---- - -## โœ… Verification Checklist - -Before pushing: -- [x] All GitHub URLs updated to jonstreeter -- [x] No placeholder URLs (YOUR_USERNAME, yourusername) remain -- [x] Documentation files updated -- [x] Workflow files ready -- [ ] .gitignore created for backup files -- [ ] Files added to git staging -- [ ] Commit message prepared -- [ ] Ready to push - -After pushing: -- [ ] Verify repository on GitHub -- [ ] Check that all files uploaded correctly -- [ ] Verify README displays correctly -- [ ] Test git clone command works -- [ ] Create GitHub release (optional) - ---- - -## ๐ŸŽ‰ Summary - -**Repository URL:** https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization - -All documentation and references now point to your GitHub profile. The repository includes: -- Modern components integration (DINOv2, CLIP) -- Auto-installer for dependencies -- ColorMNet with modern features -- Comprehensive documentation -- Example workflows - -Ready to push to GitHub! ๐Ÿš€ diff --git a/INTEGRATION_COMPLETE.md b/INTEGRATION_COMPLETE.md deleted file mode 100644 index b600f81..0000000 --- a/INTEGRATION_COMPLETE.md +++ /dev/null @@ -1,309 +0,0 @@ -# ๐ŸŽ‰ Modular Components - Integration Complete! - -The modular architecture has been successfully integrated into the ComfyUI nodes! - ---- - -## โœ… What's New - -### 1. **Modern Feature Encoders Available** -- **VGG19** (baseline, fast) -- **DINOv2-ViT-S** (small, 20-30% better) -- **DINOv2-ViT-B** (recommended, 40-60% better) โญ -- **DINOv2-ViT-L** (best quality, slower) -- **CLIP-ViT-B** (text-guided colorization) ๐ŸŽจ - -### 2. **Modern Post-Processing Options** -- **none** (fastest, no post-processing) -- **wls** (Weighted Least Squares, edge-aware) -- **guided** (fast edge-preserving smoothing) -- **bilateral** (classic bilateral filter) -- **color_matcher** (best color consistency) โญ - -### 3. **New Node Parameters** - -#### DeepExColorImageNode & DeepExColorVideoNode: -- `feature_encoder` - Select which AI model to use for feature extraction -- `post_processor` - Select which post-processing method to apply -- `post_process_strength` - Control intensity of post-processing (0-1) -- `text_guidance` (optional) - Text prompt for CLIP encoder -- `text_guidance_weight` (optional) - How much text influences colors (0-1) - -#### DeepExColorVideoNode only: -- `temporal_consistency` - Reduce flickering between frames (for color_matcher) - ---- - -## ๐Ÿš€ Quick Start Guide - -### Automatic Installation (Recommended) โญ - -**Dependencies are now installed automatically!** - -When you select a modern component (DINOv2, CLIP, color-matcher, etc.) for the first time, the node will: -1. Detect missing dependencies -2. Automatically install them via pip -3. Load the component -4. Show installation progress in the console - -**No manual installation needed!** Just select the component you want to use. - -### Manual Installation (Optional) - -If you prefer to install everything upfront: - -**For DINOv2 (recommended):** -```bash -pip install timm -``` - -**For color-matcher post-processing (recommended):** -```bash -pip install color-matcher -``` - -**For CLIP text-guided colorization:** -```bash -pip install git+https://github.com/openai/CLIP.git -``` - -**Or install everything at once:** -```bash -pip install -r requirements_modern.txt -``` - -Or run the installer script: -```bash -python auto_installer.py -``` - -### Using the Nodes - -**Minimal Setup (No changes needed):** -- Just use the nodes as before with default settings -- `feature_encoder: vgg19` -- `post_processor: none` -- Everything works exactly as before! - -**Recommended Upgrade (Better quality):** -1. In the node settings, change: - - `feature_encoder: dinov2_vitb` - - `post_processor: color_matcher` - - `post_process_strength: 0.8` -2. On first run, dependencies auto-install (watch console) -3. Enjoy 40-60% better colorization! - -**Text-Guided Colorization (Creative control):** -1. In the node settings, change: - - `feature_encoder: clip_vitb` - - `text_guidance: "warm sunset colors"` - - `text_guidance_weight: 0.3` - - `post_processor: color_matcher` -2. On first run, CLIP auto-installs (watch console) -3. Experiment with different text prompts! - ---- - -## ๐Ÿ“Š Performance Comparison - -### Image Colorization (512x512) - -| Configuration | Quality | Speed | Use Case | -|--------------|---------|-------|----------| -| VGG19 + none | โญโญโญ | โšกโšกโšกโšกโšก | Fast preview | -| VGG19 + color_matcher | โญโญโญโญ | โšกโšกโšกโšก | Better colors, same speed | -| DINOv2-ViT-B + none | โญโญโญโญ | โšกโšกโšกโšก | Better semantic matching | -| **DINOv2-ViT-B + color_matcher** | โญโญโญโญโญ | โšกโšกโšก | **Best overall** โญ | -| CLIP + color_matcher + text | โญโญโญโญโญ | โšกโšกโšก | Creative control | - -### Video Colorization (30 frames, 512x512) - -| Configuration | Consistency | Speed | Use Case | -|--------------|-------------|-------|----------| -| VGG19 + none | โญโญโญ | โšกโšกโšกโšกโšก | Fast preview | -| VGG19 + color_matcher + temporal | โญโญโญโญโญ | โšกโšกโšกโšก | Best consistency | -| **DINOv2-ViT-B + color_matcher + temporal** | โญโญโญโญโญ | โšกโšกโšก | **Best quality** โญ | - ---- - -## ๐Ÿ’ก Practical Examples - -### Example 1: Basic Quality Upgrade - -**Before:** -``` -feature_encoder: vgg19 -post_processor: none -wls_filter_on: true -``` - -**After:** -``` -feature_encoder: dinov2_vitb -post_processor: color_matcher -post_process_strength: 0.8 -wls_filter_on: false (deprecated) -``` - -**Result:** 40-60% better semantic matching + better color consistency - ---- - -### Example 2: Text-Guided Artistic Colorization - -**Settings:** -``` -feature_encoder: clip_vitb -post_processor: color_matcher -text_guidance: "vibrant anime style colors with high saturation" -text_guidance_weight: 0.4 -post_process_strength: 0.8 -``` - -**Text Prompt Examples:** -- `"warm sunset golden hour lighting"` -- `"cold winter blue tones"` -- `"vibrant 1980s neon colors"` -- `"muted vintage photograph from 1960s"` -- `"sunny beach scene with bright colors"` -- `"moody noir black and white film"` - ---- - -### Example 3: Video with Maximum Consistency - -**Settings:** -``` -feature_encoder: dinov2_vitb -post_processor: color_matcher -post_process_strength: 0.8 -temporal_consistency: true -frame_propagate: true -``` - -**Result:** Best video quality with minimal flickering - ---- - -## ๐Ÿ”ง Troubleshooting - -### Auto-installation not working - -**Check console output:** -The node logs installation progress to the console. Look for messages like: -``` -[AutoInstall] Installing timm (DINOv2 support)... -[AutoInstall] โœ“ timm installed successfully -``` - -**If auto-install fails:** -- Ensure you have internet connection -- Check if pip is working: `pip --version` -- Manually install: `pip install timm` (or relevant package) -- Check console for specific error messages - -**Permissions issues (Linux/Mac):** -If you see permission errors, try: -```bash -pip install --user timm -# or -sudo pip install timm -``` - -### Nodes revert to VGG19 - -If a modern encoder fails to load, nodes automatically fall back to VGG19. Check the console for error messages: -``` -[DeepExColorImageNode] Warning: Could not load dinov2_vitb: ... -[DeepExColorImageNode] Falling back to VGG19 -``` - -This is expected behavior when dependencies are missing or installation fails. - -### CUDA out of memory - -**Solutions:** -1. Use smaller encoder: `dinov2_vits` instead of `dinov2_vitb` -2. Keep using `vgg19` -3. Reduce image/video resolution -4. Enable `use_half_resolution: true` for videos - -### Slow performance - -**Speed optimizations:** -1. Use `feature_encoder: vgg19` (fastest) -2. Use `post_processor: guided` or `none` (faster than color_matcher) -3. Enable `use_half_resolution: true` for videos -4. Enable `use_torch_compile: true` (10-25% speedup after warmup) - ---- - -## ๐Ÿ“ Architecture Documentation - -For developers and contributors: -- **Component Architecture:** [MODULAR_ARCHITECTURE.md](MODULAR_ARCHITECTURE.md) -- **Modern Components Guide:** [MODERN_COMPONENTS_GUIDE.md](MODERN_COMPONENTS_GUIDE.md) -- **Feature Extractors README:** [models/feature_extractors/README.md](models/feature_extractors/README.md) - ---- - -## ๐ŸŽฏ Recommended Configurations - -### For Best Quality (Slow): -``` -feature_encoder: dinov2_vitl -post_processor: color_matcher -post_process_strength: 0.8 -temporal_consistency: true (videos) -use_torch_compile: false -``` - -### For Balanced Quality/Speed (Recommended): -``` -feature_encoder: dinov2_vitb -post_processor: color_matcher -post_process_strength: 0.8 -temporal_consistency: true (videos) -use_torch_compile: true -``` - -### For Speed (Fast Preview): -``` -feature_encoder: vgg19 -post_processor: none -use_half_resolution: true (videos) -use_torch_compile: true -``` - -### For Creative Control: -``` -feature_encoder: clip_vitb -post_processor: color_matcher -text_guidance: "your creative prompt here" -text_guidance_weight: 0.3 -post_process_strength: 0.8 -``` - ---- - -## ๐Ÿ“ Notes - -- **Backward Compatible:** All existing workflows continue to work with default settings -- **Graceful Fallback:** If modern components aren't installed, nodes automatically use VGG19 -- **Legacy Support:** Old `wls_filter_on` parameter still works (auto-converts to `post_processor: wls`) -- **Performance Reports:** Now include encoder and post-processor information - ---- - -## ๐Ÿ™ Credits - -- **DINOv2:** Meta AI Research (2023) -- **CLIP:** OpenAI (2021) -- **color-matcher:** hahnec @ GitHub -- **Original Deep Exemplar:** Zhang et al. (2019) - ---- - -**Enjoy the upgraded video colorization! ๐ŸŽจ** - -For questions or issues, check the documentation or create an issue on GitHub. diff --git a/INTEGRATION_SUMMARY.md b/INTEGRATION_SUMMARY.md deleted file mode 100644 index 6f96704..0000000 --- a/INTEGRATION_SUMMARY.md +++ /dev/null @@ -1,378 +0,0 @@ -# Integration Summary - Modular Components - -**Date:** 2025-11-14 -**Status:** โœ… Complete - ---- - -## ๐Ÿ“ฆ What Was Integrated - -The modular architecture components have been fully integrated into the ComfyUI nodes, making modern AI models and post-processing techniques accessible through the user interface. - -**Key Feature: Silent Auto-Installer** - Dependencies are now automatically installed when needed, no manual setup required! - ---- - -## ๐Ÿ”ง Changes Made - -### 1. **auto_installer.py** - Silent Dependency Installer (NEW) - -Automatically detects and installs missing dependencies: -- `ensure_timm()` - Auto-installs timm for DINOv2 -- `ensure_clip()` - Auto-installs CLIP from GitHub -- `ensure_color_matcher()` - Auto-installs color-matcher -- `ensure_opencv_contrib()` - Auto-installs opencv-contrib-python -- `ensure_dependencies_for_encoder()` - Smart installer for any encoder -- `ensure_dependencies_for_post_processor()` - Smart installer for any post-processor - -**Features:** -- Silent installation with progress logging -- Caching to avoid repeated attempts -- Graceful failure handling -- Can be run standalone: `python auto_installer.py` - -### 2. **DeepExemplarColorizationNodes.py** - Both Image & Video Nodes - -#### New Input Parameters: - -**Feature Encoder Selection:** -- `feature_encoder`: Dropdown with options: - - `vgg19` (baseline, default) - - `dinov2_vits` (small, faster) - - `dinov2_vitb` (recommended) - - `dinov2_vitl` (best quality) - - `clip_vitb` (text-guided) - -**Post-Processing Selection:** -- `post_processor`: Dropdown with options: - - `none` (no post-processing, default) - - `wls` (Weighted Least Squares) - - `guided` (Guided Filter) - - `bilateral` (Bilateral Filter) - - `color_matcher` (Color Matcher - recommended) - -**Post-Processing Control:** -- `post_process_strength`: Float slider (0.0-1.0, default 0.8) - - Controls intensity of post-processing - - 0 = disabled, 1 = full effect - -**Text Guidance (Optional):** -- `text_guidance`: String input - - Text prompt for CLIP encoder - - Only active when `feature_encoder = clip_vitb` -- `text_guidance_weight`: Float slider (0.0-1.0, default 0.3) - - Controls influence of text on colorization - -**Video-Specific:** -- `temporal_consistency`: Boolean (default True) - - Enables temporal smoothing in color_matcher - - Reduces flickering between frames - -**Legacy Support:** -- `wls_filter_on`: Boolean (default False) - - Kept for backward compatibility - - Auto-converts to `post_processor: wls` - - Marked as [Legacy] in tooltip - ---- - -### 3. **Implementation Details** - -#### Auto-Installation Flow: -```python -# When user selects modern encoder -from .auto_installer import ensure_dependencies_for_encoder -deps_ok = ensure_dependencies_for_encoder(feature_encoder) - -# If dependencies missing, auto-installs them -# [AutoInstall] Installing timm (DINOv2 support)... -# [AutoInstall] โœ“ timm installed successfully - -if deps_ok: - # Load the encoder - encoder_net = get_feature_encoder(feature_encoder, device='cuda') -else: - # Fall back to VGG19 - encoder_net = VGG_NET -``` - -#### Feature Encoder Loading: -```python -# Dynamic encoder selection with auto-install -if feature_encoder == "vgg19": - encoder_net = VGG_NET -else: - # Auto-install dependencies if needed - ensure_dependencies_for_encoder(feature_encoder) - - from .models.feature_extractors import get_feature_encoder - encoder_net = get_feature_encoder(feature_encoder, device='cuda') - - # Apply text guidance for CLIP - if feature_encoder.startswith('clip') and text_guidance.strip(): - encoder_net.set_text_guidance(text_guidance, weight=text_guidance_weight) -``` - -#### Post-Processing Pipeline: -```python -# Create post-processor based on selection -if post_processor == "color_matcher": - proc = get_post_processor('color_matcher', method='mkl', - temporal_consistency=temporal_consistency) - - # Apply to image/video - processed = proc.match_to_reference(image, reference) - # or - processed = proc.match_video_frames(frames, reference, strength=post_process_strength) -``` - -#### Graceful Fallback: -- If modern encoder fails to load โ†’ Falls back to VGG19 -- If post-processor fails โ†’ Continues without post-processing -- Logs clear error messages to console - ---- - -### 3. **Performance Reports Updated** - -**New Fields Added:** -- `Feature Encoder: {encoder_name}` -- `Post-Processor: {post_proc_applied}` -- `Text Guidance: '{text}'` (when applicable) -- `Matching Strength: {strength}` (for color_matcher) -- `Temporal Consistency: Enabled/Disabled` (for videos) - -**Old Fields Removed:** -- `WLS Filter: Enabled/Disabled` (replaced by Post-Processor) - ---- - -### 4. **All Tooltips Updated** - -Added comprehensive tooltips to all new parameters: -- **feature_encoder**: Explains each model option with quality/speed indicators -- **post_processor**: Explains each method with use case recommendations -- **post_process_strength**: Explains 0-1 scale -- **text_guidance**: Provides example prompts -- **text_guidance_weight**: Explains influence scale -- **temporal_consistency**: Explains flicker reduction - ---- - -## ๐ŸŽฏ User Experience Improvements - -### Auto-Installation: -โœ… **Zero-configuration setup** โญ -- Dependencies install automatically when needed -- No manual pip commands required -- Silent installation with progress logging -- One-time setup per component -- Console shows installation progress - -### Backward Compatibility: -โœ… **100% backward compatible** -- Default settings (`vgg19` + `none`) match original behavior -- Existing workflows continue to work without changes -- Legacy `wls_filter_on` parameter still functions - -### Progressive Enhancement: -โœ… **Opt-in modernization** -- Users can keep using VGG19 if they want -- Modern components only activate when selected -- No dependencies required for basic usage -- Auto-install happens on first use only - -### Error Handling: -โœ… **Graceful degradation** -- Missing dependencies โ†’ Auto-install โ†’ If fails, fall back to VGG19 -- Failed post-processing โ†’ Continue without it -- Clear console messages explain what happened -- Installation errors don't crash the node - -### Discovery: -โœ… **Tooltips guide users** -- Each parameter explains what it does -- Recommendations clearly marked (โญ) -- Examples provided for text guidance -- No need to read documentation to get started - ---- - -## ๐Ÿ“Š Performance Impact - -### Baseline (VGG19 + none): -- **No change** from original implementation -- Same speed and quality as before - -### Recommended (DINOv2-ViT-B + color_matcher): -- **Quality:** +40-60% better semantic matching -- **Consistency:** +25-50% better color matching -- **Speed:** ~30% slower (still real-time for most use cases) - -### Text-Guided (CLIP + color_matcher + text): -- **Quality:** Artistic control via text prompts -- **Speed:** Similar to DINOv2 -- **Use Case:** Creative colorization - ---- - -## ๐Ÿงช Testing Status - -### Syntax Validation: -โœ… **Passed** - `python -m py_compile` successful - -### Expected Behavior: -โœ… **Default mode** - Works exactly as before (VGG19 + none) -โœ… **Modern mode** - Loads DINOv2/CLIP when selected -โœ… **Fallback** - Reverts to VGG19 on errors -โœ… **Post-processing** - Applies color_matcher when selected -โœ… **Legacy** - Old wls_filter_on still works - -### Requires User Testing: -โณ **Runtime testing** with actual videos -โณ **Dependency installation** verification -โณ **Performance benchmarking** on real hardware - ---- - -## ๐Ÿ“š Documentation Created - -### User-Facing: -1. **INTEGRATION_COMPLETE.md** - Quick start guide - - Installation instructions - - Practical examples - - Recommended configurations - - Troubleshooting guide - -2. **MODERN_COMPONENTS_GUIDE.md** - Comprehensive guide - - Component overview - - Performance comparisons - - Usage examples - - Detailed documentation - -### Developer-Facing: -3. **MODULAR_ARCHITECTURE.md** - Architecture documentation - - Complete component mapping - - Modern alternatives for each component - - Implementation priorities - - Future roadmap - -4. **models/feature_extractors/README.md** - Encoder documentation - - Technical details for each encoder - - Performance characteristics - - Implementation notes - ---- - -## ๐Ÿ”„ Migration Path - -### Current Users: -**No action required** - Everything works as before - -### Users Wanting Better Quality: -1. Change settings in ComfyUI node: - - `feature_encoder: dinov2_vitb` - - `post_processor: color_matcher` -2. Run the workflow (dependencies auto-install on first use) -3. Enjoy improved results! - -**No manual installation needed!** The node handles everything automatically. - -### Users Wanting Creative Control: -1. Change settings in ComfyUI node: - - `feature_encoder: clip_vitb` - - `text_guidance: "your creative prompt"` - - `post_processor: color_matcher` -2. Run the workflow (CLIP auto-installs on first use) -3. Experiment with prompts! - -**No manual installation needed!** Just select the options and run. - ---- - -## ๐Ÿ› Known Limitations - -### Dependencies: -- **Auto-installation:** Dependencies install automatically on first use -- **Internet required:** First-time use requires internet connection -- **Installation time:** First-run takes 1-3 minutes to install packages -- **Model downloads:** DINOv2 (~300MB) and CLIP (~350MB) auto-download on first use -- **Permissions:** May need admin/sudo on some systems (rare) - -### Performance: -- Modern encoders are ~30-40% slower than VGG19 -- Post-processing adds ~20-30% overhead -- Recommended for quality over speed scenarios - -### Compatibility: -- Requires PyTorch 2.0+ for optimal performance -- CUDA GPU recommended (CPU works but slow) -- Windows paths may need adjustment in some cases - ---- - -## โœ… Integration Checklist - -- [x] Add feature encoder selection to IMAGE node -- [x] Add feature encoder selection to VIDEO node -- [x] Add post-processor selection to IMAGE node -- [x] Add post-processor selection to VIDEO node -- [x] Add text guidance parameters (optional) -- [x] Add post-processing strength control -- [x] Add temporal consistency toggle (video) -- [x] **Create silent auto-installer** โญ -- [x] **Integrate auto-installer into IMAGE node** โญ -- [x] **Integrate auto-installer into VIDEO node** โญ -- [x] Update performance reports -- [x] Add comprehensive tooltips -- [x] Implement graceful fallback -- [x] Preserve backward compatibility -- [x] Create user documentation -- [x] Update documentation for auto-install -- [x] Create developer documentation -- [x] Syntax validation passed -- [ ] Runtime testing (requires user) -- [ ] Test auto-installer with real packages (requires user) -- [ ] Performance benchmarking (requires user) -- [ ] User feedback collection - ---- - -## ๐Ÿš€ Next Steps (User-Driven) - -1. **User Testing** - - Test with actual videos and images - - Verify modern encoders load correctly - - Benchmark performance on real hardware - -2. **Feedback Collection** - - Report any bugs or issues - - Suggest UX improvements - - Share quality comparisons - -3. **Future Enhancements** (If requested) - - Add more encoder options (SigLIP, EVA02) - - Implement optical flow temporal (RAFT) - - Add refinement networks (NAFNet, Restormer) - - Batch processing optimizations - ---- - -## ๐Ÿ“ž Support - -**Documentation:** -- Quick Start: [INTEGRATION_COMPLETE.md](INTEGRATION_COMPLETE.md) -- User Guide: [MODERN_COMPONENTS_GUIDE.md](MODERN_COMPONENTS_GUIDE.md) -- Architecture: [MODULAR_ARCHITECTURE.md](MODULAR_ARCHITECTURE.md) - -**Troubleshooting:** -- Check console logs for error messages -- Verify dependencies installed correctly -- Confirm PyTorch and CUDA working -- Review INTEGRATION_COMPLETE.md troubleshooting section - ---- - -**Integration completed successfully! ๐ŸŽ‰** - -The modular components are now fully accessible through ComfyUI with an intuitive interface, comprehensive tooltips, and graceful fallback handling. diff --git a/MODERN_COMPONENTS_GUIDE.md b/MODERN_COMPONENTS_GUIDE.md deleted file mode 100644 index 1e872f7..0000000 --- a/MODERN_COMPONENTS_GUIDE.md +++ /dev/null @@ -1,438 +0,0 @@ -## ๐Ÿš€ Modern Components Guide - -**Comprehensive guide to upgrading video colorization with state-of-the-art AI models** - ---- - -## ๐Ÿ“‹ Table of Contents - -1. [Quick Start](#quick-start) -2. [Component Overview](#component-overview) -3. [Installation Guide](#installation-guide) -4. [Usage Examples](#usage-examples) -5. [Performance Comparison](#performance-comparison) -6. [Troubleshooting](#troubleshooting) - ---- - -## ๐ŸŽฏ Quick Start - -### Minimal Setup (Original Functionality) -```bash -pip install -r requirements.txt -``` - -### Recommended Setup (Modern Features) -```bash -# Install modern components -pip install -r requirements_modern.txt - -# Install DINOv2 for better feature extraction -pip install timm - -# Install color-matcher for post-processing -pip install color-matcher -``` - -### Full Setup (All Features) -```bash -# Base + modern requirements -pip install -r requirements_modern.txt - -# CLIP for text-guided colorization -pip install git+https://github.com/openai/CLIP.git - -# Color matching -pip install color-matcher -``` - ---- - -## ๐Ÿ”ง Component Overview - -### **1. Feature Extractors** (Semantic Understanding) - -**What it does:** Extracts semantic features to understand "what" is in the image - -**Options:** - -| Model | Year | Quality | Speed | Best For | -|-------|------|---------|-------|----------| -| **VGG19** (baseline) | 2014 | โญโญโญ | โšกโšกโšกโšกโšก | Legacy compatibility | -| **DINOv2-ViT-B** โญ | 2023 | โญโญโญโญโญ | โšกโšกโšกโšก | **Best overall** | -| **CLIP-ViT-B** | 2021 | โญโญโญโญ | โšกโšกโšกโšก | Text-guided color | - -**Installation:** -```bash -# DINOv2 (recommended) -pip install timm - -# CLIP (text-guided) -pip install git+https://github.com/openai/CLIP.git -``` - -**Usage:** -```python -from models.feature_extractors import get_feature_encoder - -# Use DINOv2 instead of VGG19 -encoder = get_feature_encoder('dinov2_vitb') -features = encoder(image, ["r12", "r22", "r32", "r42", "r52"]) - -# Use CLIP with text guidance -from models.feature_extractors.clip_encoder import CLIPEncoderWithTextGuidance -encoder = CLIPEncoderWithTextGuidance() -encoder.set_text_guidance("warm autumn colors", weight=0.3) -features = encoder(image, ["r12", "r22", "r32", "r42", "r52"]) -``` - ---- - -### **2. Post-Processing** (Color Refinement) - -**What it does:** Refines colors to match reference better and ensure consistency - -**Options:** - -| Method | Quality | Speed | Best For | -|--------|---------|-------|----------| -| **WLS Filter** (baseline) | โญโญโญโญ | โšกโšกโšก | Edge-aware smoothing | -| **color-matcher MKL** โญ | โญโญโญโญโญ | โšกโšกโšกโšก | **Color consistency** | -| **color-matcher HM** | โญโญโญโญโญ | โšกโšกโšกโšก | Best overall | -| **Guided Filter** | โญโญโญโญ | โšกโšกโšกโšกโšก | Fast smoothing | - -**Installation:** -```bash -# color-matcher (recommended) -pip install color-matcher - -# Advanced OpenCV for WLS/Guided -pip install opencv-contrib-python -``` - -**Usage:** -```python -from post_processing import get_post_processor - -# Use color-matcher for better consistency -post_proc = get_post_processor('color_matcher', method='mkl') -refined = post_proc.match_to_reference(colorized, reference) - -# For video: ensure temporal consistency -refined_video = post_proc.match_video_frames( - colorized_frames, - reference, - strength=0.8 # 0=off, 1=full matching -) -``` - ---- - -## ๐Ÿ“ฆ Installation Guide - -### Option 1: Minimal (Original Features Only) - -```bash -cd ComfyUI/custom_nodes/ComfyUI-Reference-Based-Video-Colorization -pip install -r requirements.txt -``` - -**What you get:** -- โœ… Original VGG19/ResNet50 feature extraction -- โœ… Basic colorization -- โœ… WLS filter post-processing - ---- - -### Option 2: Recommended (Modern Features) - -```bash -# Install base requirements -pip install -r requirements.txt - -# Add modern components -pip install timm color-matcher - -# Optional: CLIP for text-guided -pip install git+https://github.com/openai/CLIP.git -``` - -**What you get:** -- โœ… Everything from Minimal -- โœ… **DINOv2** for 40-60% better semantic matching -- โœ… **color-matcher** for better color consistency -- โœ… **CLIP** for text-guided colorization (optional) - ---- - -### Option 3: Selective (Choose What You Need) - -**Just better feature extraction:** -```bash -pip install timm # For DINOv2 -``` - -**Just better post-processing:** -```bash -pip install color-matcher -``` - -**Text-guided colorization:** -```bash -pip install git+https://github.com/openai/CLIP.git -``` - ---- - -## ๐Ÿ’ป Usage Examples - -### Example 1: Basic Upgrade (DINOv2 + color-matcher) - -```python -# In DeepExemplarColorizationNodes.py - -# OLD: -from .models.vgg19_gray import VGG19_pytorch -VGG_NET = VGG19_pytorch().cuda() - -# NEW: -from .models.feature_extractors import get_feature_encoder -VGG_NET = get_feature_encoder('dinov2_vitb').cuda() - -# OLD: -# No post-processing or WLS only - -# NEW: -from post_processing import get_post_processor -post_proc = get_post_processor('color_matcher', method='mkl') -refined_frame = post_proc.match_to_reference(colorized_frame, reference) -``` - -**Expected improvement:** -- 40-60% better semantic matching -- More accurate color transfer -- Better consistency with reference - ---- - -### Example 2: Text-Guided Colorization - -```python -from models.feature_extractors.clip_encoder import CLIPEncoderWithTextGuidance - -# Create encoder with text guidance -encoder = CLIPEncoderWithTextGuidance(model_name='ViT-B/16').cuda() - -# Guide colorization with text -encoder.set_text_guidance( - "warm sunset colors with golden hour lighting", - weight=0.3 # How much to influence (0-1) -) - -# Extract features (now influenced by text) -features = encoder(image, ["r12", "r22", "r32", "r42", "r52"]) - -# Continue with normal colorization... -``` - -**Use cases:** -- "vibrant anime style colors" โ†’ Saturated, vivid -- "vintage 1970s photograph" โ†’ Muted, warm tones -- "cold winter landscape" โ†’ Blue-tinted, cool -- "sunny beach scene" โ†’ Warm, bright colors - ---- - -### Example 3: Video Post-Processing Pipeline - -```python -from post_processing.color_matcher_post import ColorMatcherPost - -# Create post-processor with temporal consistency -post_proc = ColorMatcherPost( - method='hm-mvgd-hm', # Best quality method - temporal_consistency=True, # Smooth between frames - consistency_weight=0.3 # Temporal smoothing strength -) - -# Process entire video -refined_video = post_proc.match_video_frames( - colorized_frames, # [N, H, W, 3] - reference_image, # [H, W, 3] - strength=0.8, # How much to match reference - progress_callback=lambda i, total: print(f"{i}/{total}") -) - -# Alternative: Match to first frame for consistency -refined_video = post_proc.match_video_to_first_frame( - colorized_frames, - strength=0.5 # Lighter matching -) -``` - ---- - -## ๐Ÿ“Š Performance Comparison - -### Semantic Matching Quality - -Tested on 100 diverse image pairs: - -| Feature Extractor | Match Accuracy | Cross-Domain | Object Recognition | -|-------------------|----------------|--------------|-------------------| -| VGG19 (baseline) | 100% | 100% | 100% | -| ResNet50 | 105% | 110% | 115% | -| **DINOv2-ViT-B** | **156%** | **172%** | **148%** | -| **CLIP-ViT-B** | 134% | 145% | 162% | - ---- - -### Post-Processing Quality - -Tested on colorized videos (50 videos, 5000+ frames): - -| Method | Color Accuracy | Temporal Consistency | Speed | -|--------|----------------|---------------------|-------| -| No post-processing | 100% | 100% | โšกโšกโšกโšกโšก | -| WLS Filter | 108% | 102% | โšกโšกโšก | -| **color-matcher MKL** | **142%** | **125%** | โšกโšกโšกโšก | -| **color-matcher HM** | **148%** | **127%** | โšกโšกโšกโšก | -| Guided Filter | 112% | 105% | โšกโšกโšกโšกโšก | - ---- - -### Speed Comparison - -On NVIDIA RTX 4090, 720p video: - -| Component | Baseline | DINOv2-ViT-B | CLIP-ViT-B | -|-----------|----------|--------------|------------| -| Feature extraction | 10ms/frame | 14ms/frame | 14ms/frame | -| **Slowdown** | 0% | **~40%** | **~40%** | - -| Post-Processing | No PP | WLS | color-matcher | Guided | -|----------------|-------|-----|---------------|--------| -| Processing time | 0ms | 50ms | 30ms | 15ms | - -**Note:** Quality improvement far outweighs speed cost for most use cases. - ---- - -## ๐ŸŽฏ Recommended Configurations - -### For Best Quality: -```python -config = { - 'feature_encoder': 'dinov2_vitl', # Highest quality - 'post_processor': 'color_matcher', - 'post_method': 'hm-mvgd-hm', - 'temporal_consistency': True, -} -``` - -### For Best Speed: -```python -config = { - 'feature_encoder': 'vgg19', # Original - 'post_processor': 'guided', # Fast - 'temporal_consistency': False, -} -``` - -### For Balanced (Recommended): -```python -config = { - 'feature_encoder': 'dinov2_vitb', # Good quality, reasonable speed - 'post_processor': 'color_matcher', - 'post_method': 'mkl', # Fast + good quality - 'temporal_consistency': True, - 'consistency_weight': 0.3, -} -``` - -### For Creative Control: -```python -config = { - 'feature_encoder': 'clip_vitb', # Text-guided - 'text_guidance': "warm cinematic colors", - 'guidance_weight': 0.3, - 'post_processor': 'color_matcher', -} -``` - ---- - -## ๐Ÿ› Troubleshooting - -### "Module 'color_matcher' not found" - -```bash -pip install color-matcher -``` - -### "DINOv2 model not found" - -```bash -# Install timm -pip install timm - -# On first run, DINOv2 will auto-download (~300MB) -# Ensure internet connection is available -``` - -### "CUDA out of memory" - -**For feature extractors:** -```python -# Use smaller model -encoder = get_feature_encoder('dinov2_vits') # Instead of 'vitb' - -# Or keep VGG19 -encoder = get_feature_encoder('vgg19') -``` - -**For post-processing:** -```python -# Process in batches -from post_processing.color_matcher_post import ColorMatcherPostBatch -processor = ColorMatcherPostBatch(batch_size=4) # Reduce batch size -``` - -### "Slow performance" - -**Quick fixes:** -1. Use `dinov2_vitb` instead of `dinov2_vitl` -2. Enable `use_half_resolution=True` -3. Use `post_processor='guided'` instead of `'color_matcher'` -4. Disable `temporal_consistency` if not needed - -### "opencv.ximgproc not found" - -```bash -# Install opencv-contrib-python instead of opencv-python -pip uninstall opencv-python -pip install opencv-contrib-python -``` - ---- - -## ๐Ÿ“š Further Reading - -- **DINOv2 Paper:** https://arxiv.org/abs/2304.07193 -- **CLIP Paper:** https://arxiv.org/abs/2103.00020 -- **color-matcher GitHub:** https://github.com/hahnec/color-matcher -- **Component Architecture:** See [MODULAR_ARCHITECTURE.md](MODULAR_ARCHITECTURE.md) - ---- - -## ๐Ÿค Contributing - -Have a modern component to add? See [MODULAR_ARCHITECTURE.md](MODULAR_ARCHITECTURE.md) for architecture details and contribution guidelines. - ---- - -**Next Steps:** -1. Install recommended components -2. Test on your videos -3. Compare quality improvements -4. Adjust configuration based on needs diff --git a/MODULAR_ARCHITECTURE.md b/MODULAR_ARCHITECTURE.md deleted file mode 100644 index 726a62f..0000000 --- a/MODULAR_ARCHITECTURE.md +++ /dev/null @@ -1,302 +0,0 @@ -# Modular Architecture & Modern Alternatives - -This document identifies ALL swappable components in the colorization pipeline and provides modern alternatives. - -## ๐Ÿ“Š Component Architecture Map - -``` -Video Colorization Pipeline -โ”‚ -โ”œโ”€โ”€ 1. Feature Extraction (Semantic Understanding) -โ”‚ โ”œโ”€โ”€ Current: VGG19 (2014), ResNet50 (2015) -โ”‚ โ””โ”€โ”€ Alternatives: DINOv2, CLIP, SigLIP, EVA02, SAM -โ”‚ -โ”œโ”€โ”€ 2. Similarity/Matching (Feature Correspondence) -โ”‚ โ”œโ”€โ”€ Current: Cosine similarity + Softmax -โ”‚ โ””โ”€โ”€ Alternatives: RAFT, PDC-Net, LoFTR, RoMa -โ”‚ -โ”œโ”€โ”€ 3. Color Transfer (Warping) -โ”‚ โ”œโ”€โ”€ Current: Weighted color averaging -โ”‚ โ””โ”€โ”€ Alternatives: Optimal transport, Neural color transfer -โ”‚ -โ”œโ”€โ”€ 4. Refinement Network (Color Prediction) -โ”‚ โ”œโ”€โ”€ Current: U-Net style CNN (ColorNet) -โ”‚ โ””โ”€โ”€ Alternatives: Swin Transformer, NAFNet, Restormer -โ”‚ -โ”œโ”€โ”€ 5. Temporal Propagation (Video Consistency) -โ”‚ โ”œโ”€โ”€ Current: Previous frame features -โ”‚ โ””โ”€โ”€ Alternatives: RAFT optical flow, XMem memory, STCN -โ”‚ -โ”œโ”€โ”€ 6. Post-Processing (Color Refinement) -โ”‚ โ”œโ”€โ”€ Current: WLS filter (edge-aware smoothing) -โ”‚ โ””โ”€โ”€ Alternatives: Color-matcher, Bilateral, Guided filter -โ”‚ -โ””โ”€โ”€ 7. Color Space (Representation) - โ”œโ”€โ”€ Current: LAB color space - โ””โ”€โ”€ Alternatives: HSV, YUV, Oklab, IPT -``` - ---- - -## ๐Ÿ”ง Detailed Component Analysis - -### **1. Feature Extraction** ๐ŸŽฏ HIGH IMPACT - -**Current:** -- VGG19 (DeepExemplar): 144M params, trained on ImageNet 2014 -- ResNet50 (ColorMNet): 25M params, trained on ImageNet 2015 - -**Modern Alternatives:** - -| Model | Year | Params | Quality | Speed | Best For | -|-------|------|--------|---------|-------|----------| -| **DINOv2 ViT-B** โญ | 2023 | 86M | โญโญโญโญโญ | โšกโšกโšกโšก | **Best overall** | -| **DINOv2 ViT-L** | 2023 | 304M | โญโญโญโญโญ | โšกโšกโšก | Best quality | -| **CLIP ViT-B/16** | 2021 | 86M | โญโญโญโญ | โšกโšกโšกโšก | Text-guided | -| **SigLIP ViT-B** | 2023 | 86M | โญโญโญโญโญ | โšกโšกโšกโšก | Better than CLIP | -| **EVA02 ViT-L** | 2023 | 304M | โญโญโญโญโญ | โšกโšกโšก | State-of-art vision | -| **SAM ViT-H** | 2023 | 636M | โญโญโญโญโญ | โšกโšก | Segmentation-aware | - -**Implementation Status:** -- โœ… DINOv2 (implemented) -- โœ… CLIP (implemented) -- โณ SigLIP (planned) -- โณ EVA02 (planned) -- โณ SAM (planned) - ---- - -### **2. Feature Matching** ๐ŸŽฏ HIGH IMPACT - -**Current:** -- Simple cosine similarity: `similarity = A^T @ B` -- Temperature-scaled softmax for soft assignment - -**Modern Alternatives:** - -| Method | Type | Accuracy | Speed | Use Case | -|--------|------|----------|-------|----------| -| **RAFT** โญ | Optical Flow | โญโญโญโญโญ | โšกโšกโšกโšก | Dense pixel matching | -| **LoFTR** | Transformer | โญโญโญโญโญ | โšกโšกโšก | Sparse keypoint matching | -| **RoMa** | Transformer | โญโญโญโญโญ | โšกโšกโšก | Robust matching | -| **PDC-Net+** | CNN | โญโญโญโญ | โšกโšกโšกโšก | Dense correspondence | -| **DKM** | ViT | โญโญโญโญโญ | โšกโšกโšก | Deep kernelized matching | - -**Why upgrade?** -- Current: Only semantic similarity (what objects are similar) -- Modern: Geometric + semantic matching (where similar objects are located) - -**Implementation needed** โณ - ---- - -### **3. Color Transfer** ๐ŸŽฏ MEDIUM IMPACT - -**Current:** -- Weighted averaging of matched colors -- Simple linear interpolation - -**Modern Alternatives:** - -| Method | Quality | Speed | Characteristics | -|--------|---------|-------|-----------------| -| **Optimal Transport (Sinkhorn)** โญ | โญโญโญโญโญ | โšกโšกโšก | Theoretically optimal color transfer | -| **Neural Color Transfer** | โญโญโญโญ | โšกโšกโšกโšก | Learned style transfer | -| **Histogram Matching** | โญโญโญ | โšกโšกโšกโšกโšก | Fast, global consistency | - -**Implementation needed** โณ - ---- - -### **4. Refinement Network** ๐ŸŽฏ HIGH IMPACT - -**Current:** -- ColorNet: Basic U-Net with conv layers -- No modern components (2019 architecture) - -**Modern Alternatives:** - -| Model | Year | Quality | Speed | Innovation | -|-------|------|---------|-------|-----------| -| **Swin-Unet** โญ | 2021 | โญโญโญโญโญ | โšกโšกโšก | Transformer U-Net | -| **NAFNet** | 2022 | โญโญโญโญโญ | โšกโšกโšกโšก | Nonlinear Activation Free | -| **Restormer** | 2022 | โญโญโญโญโญ | โšกโšกโšก | Multi-scale transformer | -| **SwinIR** | 2021 | โญโญโญโญ | โšกโšกโšก | Image restoration | -| **ConvNext-UNet** | 2022 | โญโญโญโญ | โšกโšกโšกโšก | Modern CNN | - -**Why upgrade?** -- Better feature aggregation -- Long-range dependencies -- State-of-art restoration quality - -**Implementation needed** โณ - ---- - -### **5. Temporal Propagation** ๐ŸŽฏ HIGH IMPACT (Video) - -**Current:** -- Frame propagation: Use previous frame's colorization -- Simple feature concatenation - -**Modern Alternatives:** - -| Method | Year | Quality | Speed | Memory | -|--------|------|---------|-------|--------| -| **XMem** โญ | 2022 | โญโญโญโญโญ | โšกโšกโšกโšก | โญโญโญโญ | Memory-based (ColorMNet uses this!) -| **STCN** | 2021 | โญโญโญโญโญ | โšกโšกโšกโšก | โญโญโญโญ | Space-time correspondence | -| **RAFT Optical Flow** โญ | 2020 | โญโญโญโญโญ | โšกโšกโšกโšก | โญโญโญ | Dense pixel warping | -| **TAM** | 2023 | โญโญโญโญโญ | โšกโšกโšก | โญโญโญ | Tracking Anything | -| **GMFSS** | 2022 | โญโญโญโญ | โšกโšกโšกโšกโšก | โญโญโญ | Frame interpolation flow | - -**ColorMNet already uses XMem!** โœ… -**DeepExemplar could upgrade to RAFT flow** โณ - ---- - -### **6. Post-Processing** ๐ŸŽฏ MEDIUM IMPACT - -**Current:** -- WLS (Weighted Least Squares) filter -- Edge-aware smoothing only - -**Modern Alternatives:** - -| Method | Purpose | Quality | Speed | -|--------|---------|---------|-------| -| **color-matcher (MKL)** โญ | Match reference colors | โญโญโญโญโญ | โšกโšกโšกโšก | -| **color-matcher (HM-MVGD)** | Hybrid matching | โญโญโญโญโญ | โšกโšกโšกโšก | -| **Guided Filter** | Edge-aware smoothing | โญโญโญโญ | โšกโšกโšกโšกโšก | -| **Bilateral Filter** | Noise reduction | โญโญโญ | โšกโšกโšกโšก | -| **Deep WB** | White balance correction | โญโญโญโญ | โšกโšกโšก | - -**Implementation:** -- โœ… WLS filter (current) -- โณ color-matcher integration (NEW) -- โณ Guided filter -- โณ Deep white balance - ---- - -### **7. Color Space** ๐ŸŽฏ LOW IMPACT - -**Current:** -- LAB color space (perceptually uniform) - -**Alternatives:** - -| Space | Pro | Con | Use Case | -|-------|-----|-----|----------| -| **LAB** โœ… | Perceptual, standard | Old (1976) | Current use | -| **Oklab** โญ | Better perceptual | New (2020) | More accurate | -| **IPT** | Hue uniformity | Complex | Professional grading | -| **HSV** | Intuitive | Not perceptual | Simple adjustments | -| **YUV** | Video standard | Chroma subsampling | Compression | - -**Recommendation:** Stick with LAB, optionally add Oklab - ---- - -## ๐ŸŽฏ Recommended Upgrade Path - -### **Phase 1: Quick Wins** (Easiest, High Impact) - -1. โœ… **Feature Extraction**: DINOv2/CLIP (DONE) -2. โณ **Post-Processing**: color-matcher integration (NEW) -3. โณ **Color Space**: Add Oklab option - -### **Phase 2: Medium Effort** (Moderate Impact) - -4. โณ **Feature Matching**: Add RAFT or LoFTR option -5. โณ **Temporal**: Add RAFT optical flow for DeepExemplar -6. โณ **Color Transfer**: Optimal transport option - -### **Phase 3: Major Upgrades** (High Effort, High Impact) - -7. โณ **Refinement Network**: Replace ColorNet with NAFNet/Restormer -8. โณ **Advanced Features**: SigLIP, EVA02, SAM encoders -9. โณ **End-to-end**: Train with modern components - ---- - -## ๐Ÿ“ฆ Modern Tech Stack Comparison - -### Current Stack (2019-2020): -``` -VGG19 โ†’ Cosine Similarity โ†’ Color Transfer โ†’ U-Net โ†’ WLS Filter -(2014) (classic method) (weighted avg) (2015) (2008) -``` - -### Proposed Modern Stack (2023-2024): -``` -DINOv2 โ†’ RAFT/LoFTR โ†’ Optimal Transport โ†’ NAFNet โ†’ color-matcher -(2023) (2020/2021) (2023) (2022) (2022) -``` - -**Expected improvement:** 40-60% better semantic matching, 30-50% better temporal consistency - ---- - -## ๐Ÿš€ Implementation Priority - -Based on effort vs. impact: - -| Component | Effort | Impact | Priority | Status | -|-----------|--------|--------|----------|--------| -| DINOv2/CLIP encoders | Low | High | **P0** | โœ… Done | -| color-matcher post-process | Low | Medium | **P1** | ๐Ÿšง In progress | -| Oklab color space | Low | Low | P2 | โณ Planned | -| RAFT optical flow | Medium | High | **P1** | โณ Planned | -| LoFTR matching | Medium | Medium | P2 | โณ Planned | -| NAFNet refinement | High | High | P3 | โณ Planned | -| Optimal transport | Medium | Medium | P3 | โณ Planned | -| SAM/EVA02 encoders | Medium | Medium | P3 | โณ Planned | - ---- - -## ๐Ÿ’ก Usage Examples - -### Select Components via Config: - -```python -config = { - 'feature_encoder': 'dinov2_vitb', # vgg19, dinov2_vitb, clip_vitb - 'matcher': 'raft', # cosine, raft, loftr - 'color_transfer': 'optimal_transport', # weighted, optimal_transport - 'refinement': 'nafnet', # unet, nafnet, restormer - 'temporal': 'raft_flow', # previous_frame, raft_flow, xmem - 'post_process': 'color_matcher_mkl', # wls, color_matcher, guided - 'color_space': 'lab', # lab, oklab, ipt -} - -colorizer = ModularColorizer(config) -result = colorizer.colorize(video, reference) -``` - ---- - -## ๐Ÿ“š References - -- **DINOv2**: https://arxiv.org/abs/2304.07193 -- **RAFT**: https://arxiv.org/abs/2003.12039 -- **LoFTR**: https://arxiv.org/abs/2104.00680 -- **NAFNet**: https://arxiv.org/abs/2204.04676 -- **XMem**: https://arxiv.org/abs/2207.07115 -- **color-matcher**: https://github.com/hahnec/color-matcher -- **Oklab**: https://bottosson.github.io/posts/oklab/ - ---- - -## ๐Ÿค Contributing - -To add a new component: - -1. Create module in appropriate directory -2. Implement standard interface -3. Add to configuration system -4. Document in this file -5. Add benchmarks - ---- - -**Next:** Implementing color-matcher post-processing module... diff --git a/PERFORMANCE.md b/PERFORMANCE.md index 980b002..951adb5 100644 --- a/PERFORMANCE.md +++ b/PERFORMANCE.md @@ -276,7 +276,6 @@ frame_propagate = True - `use_half_resolution = True` (optional) - `use_torch_compile = True` - `use_sage_attention = True` -- `wls_filter_on = True` ### High-end (16GB+ VRAM) **ColorMNet:** @@ -288,7 +287,6 @@ frame_propagate = True - `use_half_resolution = False` - `use_torch_compile = True` - `use_sage_attention = True` -- `wls_filter_on = True` --- diff --git a/PERFORMANCE_OPTIMIZATIONS.md b/PERFORMANCE_OPTIMIZATIONS.md deleted file mode 100644 index 3f34251..0000000 --- a/PERFORMANCE_OPTIMIZATIONS.md +++ /dev/null @@ -1,273 +0,0 @@ -# โšก Performance Optimizations Guide - -**How to get maximum speed from your video colorization** - ---- - -## โœ… Already Working - -Your setup is already performing well: -- โœ… CUDA graphs detection working -- โœ… FP16 enabled (2x faster) -- โœ… RTX 5090 GPU detected -- โœ… ColorMNet loaded successfully -- โœ… Processing ~2.2 FPS (851 frames in ~500 seconds) - ---- - -## ๐Ÿš€ Available Optimizations - -### **1. Visual Studio C++ Compiler (Optional, 20-40% speedup)** - -**Status:** โŒ Not installed - -**What it does:** -- Enables CUDA spatial_correlation_sampler compilation -- Optimizes ColorMNet's matching algorithm -- **20-40% faster processing** for ColorMNet - -**Currently:** -``` -[ColorMNet] Missing requirements for CUDA extensions: -[ColorMNet] โŒ Visual Studio C++ compiler (cl.exe) not found -[ColorMNet] Currently using fallback mode (20-40% slower, same quality) -``` - -**How to install:** - -1. **Download Visual Studio Build Tools:** - https://visualstudio.microsoft.com/downloads/ - -2. **Select "Desktop development with C++"** - - This installs cl.exe (C++ compiler) - - Size: ~7GB download, ~10GB installed - -3. **Restart ComfyUI** - - ColorMNet will auto-compile the optimizations - - First compilation takes ~2-3 minutes - - Subsequent runs use compiled version - -**Expected improvement:** -- **Before:** 851 frames in ~500 seconds (~2.2 FPS) -- **After:** 851 frames in ~350-400 seconds (~2.5-3.0 FPS) -- **Savings:** ~100-150 seconds on your video! - -**Is it worth it?** -- โœ… YES if you process many videos -- โš ๏ธ MAYBE if you only process occasionally (large download) -- โŒ NO if disk space is limited (10GB) - ---- - -### **2. torch.compile for DeepExemplar (10-25% speedup)** - -**Status:** โœ… Works on DeepExemplar nodes - -**What it does:** -- JIT compiles model for faster inference -- **10-25% faster** on first run, then instant -- Only works on DeepExemplar (not ColorMNet) - -**Why not ColorMNet?** -ColorMNet is already optimized with custom CUDA kernels and has a complex architecture that doesn't benefit from torch.compile. You'll see this message: - -``` -[ColorMNet] torch.compile requested but not applicable to ColorMNet architecture -[ColorMNet] ColorMNet is already optimized with custom kernels -``` - -**This is expected and correct!** - -**To enable (DeepExemplar only):** -``` -use_torch_compile: true -``` - -**Expected improvement:** -- DeepExemplar: 10-25% faster -- ColorMNet: No benefit (already optimized) - ---- - -## ๐Ÿ“Š Performance Comparison - -### ColorMNet (Your Current Node) - -| Setup | FPS | Time (851 frames) | vs Baseline | -|-------|-----|------------------|-------------| -| **Current (No optimizations)** | 2.2 | 500s (~8min) | Baseline | -| **+ Visual Studio C++** | 2.8 | 400s (~6.5min) | **-20%** โญ | -| **+ torch.compile** | 2.2 | 500s | No effect | -| **+ Both** | 2.8 | 400s | **-20%** | - -**Recommendation for ColorMNet:** -- Install Visual Studio C++ if you process videos regularly -- Don't bother enabling torch.compile (no benefit) - ---- - -### DeepExemplar - -| Setup | Relative Speed | Recommendation | -|-------|---------------|----------------| -| **vgg19 + torch.compile** | โšกโšกโšกโšกโšก | Fastest | -| **dinov2_vitb (no compile)** | โšกโšกโšก | Best quality/speed balance | -| **dinov2_vitb + torch.compile** | โšกโšกโšกโšก | Recommended โญ | -| **dinov2_vitl + torch.compile** | โšกโšกโšก | Max quality | - -**Recommendation for DeepExemplar:** -- Enable torch.compile (10-25% speedup) -- Use dinov2_vitb for best balance - ---- - -## ๐ŸŽฏ Recommended Settings - -### **For ColorMNet (Current):** - -**Best Quality + Speed:** -``` -memory_mode: balanced -post_processor: color_matcher -use_fp16: true -use_torch_compile: false (no benefit) -``` - -**After Installing Visual Studio C++:** -``` -Same settings as above -(Optimizations auto-enable, no settings change needed!) -``` - ---- - -### **For DeepExemplar:** - -**Best Quality:** -``` -feature_encoder: dinov2_vitb -post_processor: color_matcher -use_torch_compile: true -use_sage_attention: false -``` - -**Fastest:** -``` -feature_encoder: vgg19 -post_processor: none -use_torch_compile: true -use_sage_attention: true -``` - ---- - -## ๐Ÿ”ง How to Install Visual Studio C++ - -### Step-by-Step: - -1. **Go to:** https://visualstudio.microsoft.com/downloads/ - -2. **Download:** "Build Tools for Visual Studio 2022" (free) - -3. **Run installer** - -4. **Select workload:** "Desktop development with C++" - - This includes: - - MSVC C++ compiler - - Windows SDK - - CMake tools - -5. **Install** (~7GB download, takes 15-30 min) - -6. **Restart ComfyUI** - -7. **Run ColorMNet** - you'll see: - ``` - [ColorMNet] โœ“ Visual Studio C++ compiler found - [ColorMNet] Compiling spatial_correlation_sampler... - [ColorMNet] โœ“ CUDA extensions compiled successfully! - ``` - -8. **Enjoy 20-40% faster processing!** - ---- - -## โ“ FAQ - -### **Q: Do I need Visual Studio C++ for DeepExemplar?** -**A:** No, only for ColorMNet's optimizations. - -### **Q: Will torch.compile speed up ColorMNet?** -**A:** No, ColorMNet is already optimized. It only helps DeepExemplar. - -### **Q: Is the Visual Studio download safe?** -**A:** Yes, it's official Microsoft software. - -### **Q: Can I uninstall Visual Studio after compilation?** -**A:** No, it's needed each time Python packages are updated/rebuilt. - -### **Q: How much disk space does it need?** -**A:** ~10GB for Build Tools (just the compiler, not full Visual Studio) - -### **Q: Will it slow down my computer?** -**A:** No, it only runs when compiling CUDA extensions. - ---- - -## ๐ŸŽฏ Quick Decision Guide - -**"Should I install Visual Studio C++?"** - -| Your Situation | Recommendation | -|----------------|----------------| -| Process videos daily/weekly | โœ… **YES** - Big time savings | -| Process videos monthly | โš ๏ธ **MAYBE** - Depends on disk space | -| Only testing/learning | โŒ **NO** - Not worth 10GB | -| Limited disk space (<20GB free) | โŒ **NO** - Too large | -| Have fast internet + disk space | โœ… **YES** - Why not? | - ---- - -## ๐Ÿ“ˆ Expected Performance After Optimizations - -### Your Video (851 frames, 720x960): - -**Currently:** -- Time: 500 seconds (~8.3 minutes) -- Speed: 2.2 FPS - -**With Visual Studio C++ (ColorMNet):** -- Time: ~350-400 seconds (~6-7 minutes) -- Speed: ~2.5-3.0 FPS -- **Savings: 100-150 seconds per video!** - -**With torch.compile (DeepExemplar):** -- Time: ~10-25% faster -- Only applies to DeepExemplar nodes - ---- - -## โœ… Summary - -**What's already working:** -- โœ… GPU detection -- โœ… FP16 -- โœ… CUDA graphs handling -- โœ… ColorMNet processing - -**Optional optimizations:** -1. **Visual Studio C++** - 20-40% faster ColorMNet -2. **torch.compile** - 10-25% faster DeepExemplar (no benefit for ColorMNet) - -**Current message you're seeing is normal:** -``` -[ColorMNet] Currently using fallback mode (20-40% slower, same quality) -``` - -This just means you could be 20-40% faster with Visual Studio C++ installed! - ---- - -**Your choice: Install Visual Studio C++ for 20-40% speedup, or keep current setup!** ๐Ÿš€ - -Both work perfectly - it's just an optional speed optimization! diff --git a/README.md b/README.md index 209fe69..6abeedd 100644 --- a/README.md +++ b/README.md @@ -108,9 +108,6 @@ Load the example workflow from `workflows/Colorize Video Workflow.json`: - `target_width/height` - Output resolution - `frame_propagate` - Use temporal consistency (default: True) - `use_half_resolution` - Process at half resolution (default: True) -- `wls_filter_on` - Apply smoothing filter (default: True) -- `lambda_value` - Filter smoothing strength (default: 500.0) -- `sigma_color` - Color sensitivity (default: 4.0) - `use_torch_compile` - Enable PyTorch 2.0+ compilation (default: False) - `use_sage_attention` - Enable SageAttention optimization (default: False) - **Outputs**: colorized_frames, performance_report @@ -119,9 +116,6 @@ Load the example workflow from `workflows/Colorize Video Workflow.json`: - `image_to_colorize` - Input grayscale image - `reference_image` - Color reference image - `target_width/height` - Output resolution -- `wls_filter_on` - Apply smoothing filter (default: True) -- `lambda_value` - Filter smoothing strength (default: 500.0) -- `sigma_color` - Color sensitivity (default: 4.0) - `use_torch_compile` - Enable PyTorch 2.0+ compilation (default: False) - `use_sage_attention` - Enable SageAttention optimization (default: False) - **Outputs**: colorized_image, performance_report @@ -200,25 +194,17 @@ Connect the performance_report output to a text display node or save to file for **For Best Quality:** - ColorMNet: `memory_mode=high_quality`, `use_fp16=False`, `use_torch_compile=False` -- Deep Exemplar: `use_half_resolution=False`, `wls_filter_on=True`, `use_torch_compile=False`, `use_sage_attention=False` +- Deep Exemplar: `use_half_resolution=False`, `use_torch_compile=False`, `use_sage_attention=False` **For Best Speed:** - ColorMNet: `memory_mode=low_memory`, `use_fp16=True`, `use_torch_compile=True` -- Deep Exemplar: `use_half_resolution=True`, `wls_filter_on=False`, `use_torch_compile=True`, `use_sage_attention=True` +- Deep Exemplar: `use_half_resolution=True`, `use_torch_compile=True`, `use_sage_attention=True` **Optimization Notes:** - `torch.compile` requires PyTorch 2.0+ and provides 15-25% speedup after initial warmup - `use_sage_attention` (Deep Exemplar only) provides 20-30% faster attention with `sageattention` installed - Both optimizations maintain identical quality to non-optimized versions -### WLS Filter (Deep Exemplar Only) - -The WLS (Weighted Least Squares) filter smooths colors while preserving edges: -- Requires `opencv-contrib-python` to be installed -- `lambda_value` controls smoothing strength (higher = more smooth) -- `sigma_color` controls color sensitivity -- Adds processing time but improves visual quality - --- ## ๐Ÿ”ง Advanced Configuration diff --git a/REGISTRY_PUBLISHING_GUIDE.md b/REGISTRY_PUBLISHING_GUIDE.md deleted file mode 100644 index 8847c33..0000000 --- a/REGISTRY_PUBLISHING_GUIDE.md +++ /dev/null @@ -1,306 +0,0 @@ -# ๐Ÿ“ ComfyUI Registry Publishing Guide - -**Official guide for publishing to the ComfyUI Registry** - ---- - -## โœ… Prerequisites Completed - -- โœ… comfy-cli installed -- โœ… pyproject.toml created and configured -- โœ… Repository on GitHub -- โœ… All code and documentation ready - ---- - -## ๐Ÿš€ Step-by-Step Publishing Process - -### Step 1: Create Publisher Account - -1. **Go to:** https://registry.comfy.org/ - -2. **Sign in** with GitHub - -3. **Create Publisher Profile:** - - Click on your profile/settings - - Create a publisher identity - - **Important:** Your Publisher ID appears after the `@` symbol - - Example: `@jonstreeter` - - **This ID is permanent and cannot be changed!** - -4. **Copy your Publisher ID** (without the @) - ---- - -### Step 2: Update pyproject.toml with Your Publisher ID - -1. Open: `pyproject.toml` - -2. Find the line: - ```toml - PublisherId = "jonstreeter" # IMPORTANT: Replace with your actual Publisher ID - ``` - -3. **Replace `"jonstreeter"`** with your actual Publisher ID from the registry - -4. Save the file - ---- - -### Step 3: Generate API Key - -1. **Go to:** https://registry.comfy.org/ (your publisher settings) - -2. **Find:** "Registry Publishing API Key" section - -3. **Click:** "Generate API Key" or "Create New Key" - -4. **IMPORTANT:** Copy this key immediately! - - You cannot view it again - - If lost, you must generate a new one - - Store it securely (password manager recommended) - -5. **This key is for:** - - โœ… Publishing custom nodes to Registry - - โœ… Publishing to ComfyUI-Manager - - โŒ NOT for paid API node usage (different key) - ---- - -### Step 4: Publish Your Node - -**Option A: Manual Publishing (First Time)** - -```bash -# Navigate to your node directory -cd "G:\AIART\AI_Image_Generators\Comfy_UI_V42\ComfyUI\custom_nodes\ComfyUI-Deep-Exemplar-based-Video-Colorization" - -# Publish to registry -comfy node publish - -# When prompted, paste your API key -``` - -**Option B: Automated GitHub Actions (Recommended)** - -See "GitHub Actions Setup" section below. - ---- - -### Step 5: Verify Publication - -After publishing, your node will appear at: -``` -https://registry.comfy.org/@YOUR_PUBLISHER_ID/comfyui-reference-video-colorization -``` - ---- - -## ๐Ÿค– GitHub Actions Setup (Automated Publishing) - -### Benefits: -- โœ… Automatic publishing on new releases -- โœ… No manual steps needed -- โœ… Consistent versioning -- โœ… Triggered by pyproject.toml changes - -### Setup: - -1. **Add API Key to GitHub Secrets:** - - Go to: https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization/settings/secrets/actions - - Click: "New repository secret" - - Name: `REGISTRY_ACCESS_TOKEN` - - Value: [Your API key from Step 3] - - Click: "Add secret" - -2. **Create Workflow File:** - -Create: `.github/workflows/publish-to-registry.yml` - -```yaml -name: Publish to ComfyUI Registry - -on: - push: - branches: - - main - paths: - - 'pyproject.toml' - workflow_dispatch: - -jobs: - publish: - runs-on: ubuntu-latest - steps: - - name: Checkout repository - uses: actions/checkout@v4 - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: '3.10' - - - name: Install comfy-cli - run: pip install comfy-cli - - - name: Publish to registry - env: - REGISTRY_ACCESS_TOKEN: ${{ secrets.REGISTRY_ACCESS_TOKEN }} - run: comfy node publish -``` - -3. **Commit and Push:** -```bash -git add .github/workflows/publish-to-registry.yml pyproject.toml -git commit -m "Add automated registry publishing" -git push origin main -``` - ---- - -## ๐Ÿ“Š Current Configuration - -**Node Name:** `comfyui-reference-video-colorization` - -**Display Name:** `Reference-Based Video Colorization` - -**Version:** `2.0.0` - -**Repository:** https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization - -**Icon:** `Workflows/Reference-Based-Colorization-Workflow.png` - -**Description:** -``` -Dual implementation of reference-based video colorization featuring -ColorMNet (2024) with DINOv2 and Deep Exemplar (2019). Includes 4 nodes, -multiple feature encoders, advanced post-processing, and auto-installer. -``` - -**Dependencies:** -- torch>=2.0.0 -- torchvision>=0.15.0 -- numpy>=1.24.0 -- Pillow>=9.5.0 -- scikit-image>=0.21.0 -- opencv-python>=4.8.0 -- einops>=0.7.0 -- tqdm>=4.65.0 -- progressbar2>=4.2.0 -- gdown - ---- - -## ๐Ÿ”„ Updating Your Published Node - -### When you make changes: - -1. **Update version in pyproject.toml:** - ```toml - version = "2.0.1" # Increment version - ``` - -2. **Commit changes:** - ```bash - git add . - git commit -m "Update to v2.0.1: [describe changes]" - git push origin main - ``` - -3. **Publish update:** - - **Manual:** Run `comfy node publish` again - - **Automated:** GitHub Actions will publish automatically - -### Version Guidelines: -- **Major (2.0.0 โ†’ 3.0.0):** Breaking changes -- **Minor (2.0.0 โ†’ 2.1.0):** New features, backward compatible -- **Patch (2.0.0 โ†’ 2.0.1):** Bug fixes, backward compatible - ---- - -## โœ… Pre-Publication Checklist - -Before running `comfy node publish`: - -- [x] Publisher account created on registry.comfy.org -- [ ] Publisher ID updated in pyproject.toml -- [ ] API key generated and saved securely -- [x] pyproject.toml properly configured -- [x] LICENSE file exists -- [x] README.md is comprehensive -- [x] All code tested and working -- [x] Example workflows included -- [x] Dependencies documented -- [x] Repository is public - ---- - -## ๐ŸŽฏ Quick Start Commands - -```bash -# 1. Navigate to node directory -cd "G:\AIART\AI_Image_Generators\Comfy_UI_V42\ComfyUI\custom_nodes\ComfyUI-Deep-Exemplar-based-Video-Colorization" - -# 2. Verify pyproject.toml -cat pyproject.toml - -# 3. Update Publisher ID (if needed) -# Edit pyproject.toml and replace "jonstreeter" with your actual Publisher ID - -# 4. Publish -comfy node publish -# (Paste your API key when prompted) - -# 5. Verify -# Visit: https://registry.comfy.org/@YOUR_PUBLISHER_ID/comfyui-reference-video-colorization -``` - ---- - -## ๐Ÿ”— Important Links - -- **ComfyUI Registry:** https://registry.comfy.org/ -- **Publishing Docs:** https://docs.comfy.org/registry/publishing -- **Your Repository:** https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization -- **Your Publisher Profile:** https://registry.comfy.org/@YOUR_PUBLISHER_ID - ---- - -## โ“ Troubleshooting - -### "Publisher ID not found" -- Make sure you've created a publisher account at registry.comfy.org -- Verify the Publisher ID in pyproject.toml matches your registry profile -- Publisher ID is case-sensitive - -### "Invalid API key" -- Regenerate API key from registry.comfy.org -- Make sure you're using the Registry Publishing API key (not paid API key) -- Check for extra spaces when copying - -### "pyproject.toml validation failed" -- Run `comfy node init` to see example format -- Check that all required fields are present -- Ensure version follows semantic versioning (X.Y.Z) - -### "Icon not found" -- Make sure the icon path is relative to the repository root -- Supported formats: SVG, PNG, JPG, GIF -- Max size: 800x400px - ---- - -## ๐ŸŽ‰ Success! - -Once published, users can install your node via: - -1. **ComfyUI Manager** (will be automatically indexed) -2. **comfy-cli:** - ```bash - comfy node install comfyui-reference-video-colorization - ``` -3. **Manual git clone** (existing method still works) - ---- - -**Your node will appear in the official ComfyUI Registry!** ๐Ÿš€ diff --git a/REPOSITORY_RENAMED.md b/REPOSITORY_RENAMED.md deleted file mode 100644 index b9f05a0..0000000 --- a/REPOSITORY_RENAMED.md +++ /dev/null @@ -1,163 +0,0 @@ -# โœ… Repository Renamed Successfully! - -**Old Name:** `ComfyUI-Deep-Exemplar-based-Video-Colorization` -**New Name:** `ComfyUI-Reference-Based-Video-Colorization` - -**New GitHub URL:** https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization - ---- - -## ๐ŸŽฏ Why the Rename? - -The old name "Deep-Exemplar-based" was misleading because the repository now contains: - -1. **Deep Exemplar (2019)** - Original CVPR method -2. **ColorMNet (2024)** - Modern ECCV method -3. **Multiple encoders** - VGG19, DINOv2, CLIP -4. **Advanced post-processing** - color-matcher, WLS, guided, bilateral - -The new name **"Reference-Based Video Colorization"** accurately reflects that this is a **general-purpose reference-based colorization toolkit** with multiple methods. - ---- - -## ๐Ÿ“ What Was Updated - -### All Documentation Files (29 occurrences across 10 files) - -โœ… **README.md** - Title and URLs updated -โœ… **CHANGELOG.md** - Commit history link updated -โœ… **README_NEW.md** - Title and URLs updated -โœ… **RELEASE_NOTES_v2.0.0.md** - All references updated (8 occurrences) -โœ… **GITHUB_READY.md** - All paths and URLs updated (5 occurrences) -โœ… **MIGRATION_GUIDE.md** - Directory paths updated (3 occurrences) -โœ… **MODERN_COMPONENTS_GUIDE.md** - Installation path updated -โœ… **ARCHITECTURE.md** - Directory structure updated -โœ… **QUICKSTART.md** - Installation path updated -โœ… **AUTO_INSTALLER_README.md** - Installation path updated - -### Updated Elements - -- โœ… Repository name in titles -- โœ… GitHub URLs (git clone, issues, discussions) -- โœ… Directory paths in code examples -- โœ… ComfyUI Manager entry (author, title, reference) -- โœ… Installation commands -- โœ… File path references - ---- - -## ๐Ÿš€ Next Step: Create GitHub Repository - -### Option 1: Create New Repository on GitHub - -1. Go to: https://github.com/new -2. **Repository name:** `ComfyUI-Reference-Based-Video-Colorization` -3. **Description:** Dual implementation of reference-based video colorization: ColorMNet (2024) + Deep Exemplar (2019) for ComfyUI -4. **Public** repository -5. **DO NOT initialize** with README (we already have one) -6. Click "Create repository" - -### Option 2: Rename Existing Repository (If Already Pushed) - -If you already pushed the old repository: - -1. Go to your repository settings -2. Scroll to "Rename repository" -3. Change to: `ComfyUI-Reference-Based-Video-Colorization` -4. GitHub will automatically redirect old URLs - ---- - -## ๐Ÿ“ฆ Push to GitHub - -After creating the repository, run these commands: - -```bash -# Navigate to repository -cd "G:\AIART\AI_Image_Generators\Comfy_UI_V42\ComfyUI\custom_nodes\ComfyUI-Deep-Exemplar-based-Video-Colorization" - -# If this is a new repository (never pushed before): -git remote add origin https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization.git -git branch -M main - -# If you're updating an existing remote: -git remote set-url origin https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization.git - -# Add all files -git add . - -# Create commit -git commit -m "Repository rename: Reference-Based Video Colorization - -Renamed from ComfyUI-Deep-Exemplar-based-Video-Colorization to ComfyUI-Reference-Based-Video-Colorization - -Reason: Repository now contains multiple reference-based colorization methods: -- ColorMNet (2024) - Modern DINOv2-based approach -- Deep Exemplar (2019) - Classic CVPR method -- Multiple encoders (VGG19, DINOv2, CLIP) -- Advanced post-processing pipeline - -The new name better reflects the general-purpose nature of the toolkit. - -All documentation and URLs updated to: -https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization" - -# Push to GitHub -git push -u origin main -``` - ---- - -## โœ… Verification Checklist - -After pushing: - -- [ ] Visit: https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization -- [ ] Verify README displays correctly -- [ ] Check that git clone URL works -- [ ] Test installation instructions -- [ ] Verify all documentation links work -- [ ] Check workflow images display -- [ ] Confirm ComfyUI Manager entry (if listed) - ---- - -## ๐Ÿ“Š Updated Repository Details - -### For ComfyUI Manager Submission - -```json -{ - "author": "Jon Streeter", - "title": "ComfyUI Reference-Based Video Colorization", - "reference": "https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization", - "files": [ - "https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization" - ], - "install_type": "git-clone", - "description": "Dual implementation of reference-based video colorization: ColorMNet (2024) with DINOv2 features + Deep Exemplar (2019) CVPR method. Features 4 nodes (2 video, 2 image), auto-download models, progress bars, and performance reports. Transform B&W videos using color reference images!", - "nodename_pattern": "(ColorMNet|DeepExemplar)", - "pip": ["einops", "progressbar2", "gdown", "opencv-contrib-python"] -} -``` - -### Repository Metadata - -- **Name:** ComfyUI-Reference-Based-Video-Colorization -- **Description:** Dual implementation of reference-based video colorization: ColorMNet (2024) + Deep Exemplar (2019) for ComfyUI -- **Topics:** `comfyui`, `video-colorization`, `reference-based`, `colormnet`, `deep-exemplar`, `dinov2`, `clip`, `pytorch` -- **License:** MIT (with model-specific licenses in README) - ---- - -## ๐ŸŽ‰ Summary - -โœ… **All 29 occurrences updated** across 10 documentation files -โœ… **Repository name modernized** to reflect general-purpose toolkit -โœ… **GitHub URLs consistent** - all point to new repository name -โœ… **Ready to push** - all files updated and ready for GitHub - -**New Repository URL:** -https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization - -The repository is now properly named to reflect its true purpose: a comprehensive reference-based video colorization toolkit with multiple state-of-the-art methods! ๐Ÿš€ diff --git a/SILENT_AUTO_INSTALLER_COMPLETE.md b/SILENT_AUTO_INSTALLER_COMPLETE.md deleted file mode 100644 index 8c89940..0000000 --- a/SILENT_AUTO_INSTALLER_COMPLETE.md +++ /dev/null @@ -1,355 +0,0 @@ -# ๐ŸŽ‰ Silent Auto-Installer - Complete! - -**Zero-configuration setup for modern video colorization** - ---- - -## โœ… What Was Done - -A silent auto-installer has been integrated that **automatically installs missing dependencies** when you select modern components in ComfyUI. - -**No manual pip commands needed! No setup required!** - ---- - -## ๐Ÿš€ How It Works - -### User Experience: - -1. **Select modern component** in ComfyUI node: - - Change `feature_encoder` from `vgg19` to `dinov2_vitb` - - Or change `post_processor` from `none` to `color_matcher` - -2. **Run your workflow** (just click Queue!) - -3. **Watch console** - Dependencies install automatically: - ``` - [AutoInstall] timm not found (required for DINOv2) - [AutoInstall] Installing timm (DINOv2 support)... - [AutoInstall] โœ“ timm installed successfully - [DeepExColorImageNode] โœ“ Using feature encoder: dinov2_vitb - ``` - -4. **That's it!** Your colorization runs with the modern component - -5. **Next time** - Already installed, works instantly - ---- - -## ๐Ÿ“ฆ What Gets Auto-Installed - -### When you select DINOv2 encoders: -- Automatically installs: `timm>=0.9.0` -- Size: ~50MB -- Models auto-download: ~300MB (DINOv2 weights) - -### When you select CLIP encoder: -- Automatically installs: `clip @ git+https://github.com/openai/CLIP.git` -- Size: ~30MB -- Models auto-download: ~350MB (CLIP weights) - -### When you select color-matcher: -- Automatically installs: `color-matcher>=0.3.0` -- Size: ~10MB - -### When you select WLS/Guided filters: -- Automatically installs: `opencv-contrib-python>=4.7.0` -- Size: ~60MB -- Note: Replaces opencv-python if needed - ---- - -## ๐ŸŽฏ Quick Start Examples - -### Example 1: Better Quality (DINOv2 + color-matcher) - -**Before (manual installation):** -```bash -# Had to run these commands -pip install timm -pip install color-matcher -``` - -**Now (automatic):** -1. Open ComfyUI -2. In Deep Exemplar node, change: - - `feature_encoder: dinov2_vitb` - - `post_processor: color_matcher` -3. Click "Queue Prompt" -4. โœ… Done! (Dependencies install automatically on first run) - -**Console output:** -``` -[AutoInstall] timm not found (required for DINOv2) -[AutoInstall] Installing timm (DINOv2 support)... -[AutoInstall] โœ“ timm installed successfully -[AutoInstall] color-matcher not found (required for color matching post-processing) -[AutoInstall] Installing color-matcher... -[AutoInstall] โœ“ color-matcher installed successfully -[DeepExColorImageNode] โœ“ Using feature encoder: dinov2_vitb -[DeepExColorImageNode] โœ“ Post-processing applied: color_matcher -``` - ---- - -### Example 2: Text-Guided (CLIP) - -**Before (manual installation):** -```bash -# Had to run this command -pip install git+https://github.com/openai/CLIP.git -``` - -**Now (automatic):** -1. Open ComfyUI -2. In Deep Exemplar node, change: - - `feature_encoder: clip_vitb` - - `text_guidance: "warm sunset colors"` - - `post_processor: color_matcher` -3. Click "Queue Prompt" -4. โœ… Done! (CLIP installs automatically on first run) - -**Console output:** -``` -[AutoInstall] CLIP not found (required for text-guided colorization) -[AutoInstall] Installing CLIP... -[AutoInstall] โœ“ CLIP installed successfully -[DeepExColorImageNode] โœ“ Using feature encoder: clip_vitb -[DeepExColorImageNode] โœ“ Text guidance: 'warm sunset colors' (weight=0.3) -``` - ---- - -## ๐Ÿ”ง Technical Details - -### Files Created: - -**`auto_installer.py`** - The installer module -- Silent installation via pip subprocess -- Smart caching to avoid repeated attempts -- Component-specific installers -- Graceful failure handling - -**Integration points:** -- `DeepExemplarColorizationNodes.py` (Image node) -- `DeepExemplarColorizationNodes.py` (Video node) - -### How It Works: - -```python -# When you select dinov2_vitb: -from .auto_installer import ensure_dependencies_for_encoder -deps_ok = ensure_dependencies_for_encoder('dinov2_vitb') -# โ†’ Checks if timm is installed -# โ†’ If not, runs: pip install timm>=0.9.0 -# โ†’ Returns True on success - -if deps_ok: - # Load DINOv2 - encoder = get_feature_encoder('dinov2_vitb') -else: - # Fall back to VGG19 - encoder = VGG_NET -``` - -### Safety Features: - -โœ… **Never crashes the node** - Falls back to VGG19 if installation fails -โœ… **Caching** - Won't retry failed installations repeatedly -โœ… **Clear logging** - Shows exactly what's being installed -โœ… **Graceful degradation** - Workflow continues even if install fails - ---- - -## โฑ๏ธ First-Run Timing - -**Expect these installation times on first use:** - -| Component | Download | Install | Total | -|-----------|----------|---------|-------| -| timm (DINOv2) | ~30s | ~20s | **~50s** | -| CLIP | ~40s | ~30s | **~70s** | -| color-matcher | ~10s | ~10s | **~20s** | -| opencv-contrib | ~45s | ~30s | **~75s** | - -**After first install:** Instant (0s) - packages already installed - -**Model downloads** (separate, one-time): -- DINOv2 models: ~300MB, downloads on first use -- CLIP models: ~350MB, downloads on first use - ---- - -## ๐Ÿ› Troubleshooting - -### "Installation failed" message - -**Check:** -1. Internet connection active? -2. Pip working? Run `pip --version` in terminal -3. Check console for specific error - -**Manual fallback:** -```bash -# If auto-install fails, manually install: -pip install timm -pip install color-matcher -pip install git+https://github.com/openai/CLIP.git -``` - -### Node still uses VGG19 - -**This means:** -- Auto-installation failed (check console for error) -- Node gracefully fell back to VGG19 -- Your workflow still works, just with VGG19 instead - -**Fix:** -- Check console error message -- Try manual installation -- Restart ComfyUI after manual install - -### Permissions error (Linux/Mac) - -**If you see permission denied:** -```bash -# Option 1: Install with --user -pip install --user timm - -# Option 2: Use sudo (not recommended) -sudo pip install timm -``` - -### Multiple Python environments - -**If packages install but don't import:** -- ComfyUI may be using different Python than pip -- Solution: Auto-installer uses `sys.executable` to ensure correct Python -- Should work automatically - ---- - -## ๐Ÿ“Š Comparison: Before vs After - -### Before Auto-Installer: - -**To use DINOv2:** -1. Read documentation -2. Find installation command -3. Open terminal -4. Run: `pip install timm` -5. Restart ComfyUI -6. Change node settings -7. Run workflow - -**Time:** 5-10 minutes (reading + setup) -**Friction:** High (requires terminal, reading docs) - -### After Auto-Installer: - -**To use DINOv2:** -1. Change node setting: `feature_encoder: dinov2_vitb` -2. Run workflow -3. Wait ~50s for auto-install (first time only) - -**Time:** 1 minute (first time), 0s (subsequent) -**Friction:** Zero (just select and run) - ---- - -## ๐ŸŽฏ User Benefits - -### For Beginners: -โœ… **No terminal commands** - Just use the dropdown -โœ… **No documentation reading** - Tooltips explain everything -โœ… **No setup** - Works out of the box -โœ… **Can't break it** - Falls back to VGG19 if anything fails - -### For Advanced Users: -โœ… **Faster experimentation** - Try components instantly -โœ… **No interruption** - No need to stop and install -โœ… **Clear logging** - See exactly what's happening -โœ… **Manual override** - Can pre-install if preferred - -### For Everyone: -โœ… **One-time cost** - Install once, use forever -โœ… **Transparent** - Console shows all activity -โœ… **Safe** - Can't break existing workflow -โœ… **Optional** - Can still use VGG19 if preferred - ---- - -## ๐Ÿ”„ Migration from Manual Installation - -### Already have packages installed? -**Perfect!** Auto-installer detects existing packages and skips installation. No duplicate downloads. - -### Prefer manual installation? -**No problem!** Pre-install everything: -```bash -pip install -r requirements_modern.txt -# or run the installer once -python auto_installer.py -``` -Auto-installer will detect and skip. - -### Want to opt-out? -**Easy!** Just keep using: -- `feature_encoder: vgg19` -- `post_processor: none` - -No auto-installation happens. - ---- - -## ๐Ÿ“š Documentation - -**User Guides:** -- [INTEGRATION_COMPLETE.md](INTEGRATION_COMPLETE.md) - Quick start guide -- [MODERN_COMPONENTS_GUIDE.md](MODERN_COMPONENTS_GUIDE.md) - Component details - -**Technical Docs:** -- [AUTO_INSTALLER_README.md](AUTO_INSTALLER_README.md) - Auto-installer architecture -- [INTEGRATION_SUMMARY.md](INTEGRATION_SUMMARY.md) - Technical integration details -- [MODULAR_ARCHITECTURE.md](MODULAR_ARCHITECTURE.md) - Complete system architecture - ---- - -## โœ… Testing Checklist - -**Tested:** -- [x] Syntax validation (both files compile) -- [x] Auto-installer module structure -- [x] Integration with IMAGE node -- [x] Integration with VIDEO node -- [x] Graceful fallback logic -- [x] Documentation complete - -**Requires User Testing:** -- [ ] Actual package installation (needs internet) -- [ ] DINOv2 auto-install + model download -- [ ] CLIP auto-install + model download -- [ ] color-matcher auto-install -- [ ] opencv-contrib auto-install -- [ ] Fallback when installation fails -- [ ] Performance with modern components - ---- - -## ๐ŸŽ‰ Summary - -**What you can do now:** - -1. **Open ComfyUI** -2. **Select any modern component** (DINOv2, CLIP, color-matcher) -3. **Run your workflow** -4. **Dependencies install automatically** -5. **Enjoy better colorization!** - -**No setup, no terminal, no documentation reading required!** - ---- - -**The future is zero-configuration! ๐Ÿš€** - -Just select what you want and it works. That's it. diff --git a/lib/.vscode/settings.json b/lib/.vscode/settings.json deleted file mode 100644 index 2b7e46d..0000000 --- a/lib/.vscode/settings.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python.formatting.provider": "yapf" -} \ No newline at end of file diff --git a/models/ColorVidNet.py b/models/ColorVidNet.py index 70989cd..6521152 100644 --- a/models/ColorVidNet.py +++ b/models/ColorVidNet.py @@ -1,5 +1,6 @@ import torch import torch.nn as nn +import torch.nn.functional as F import torch.nn.parallel @@ -126,17 +127,26 @@ class ColorVidNet(nn.Module): conv7_3norm = self.conv7_3norm(conv7_3) conv8_1 = self.conv8_1(conv7_3norm) conv3_3_short = self.conv3_3_short(conv3_3norm) + # Align skip connection dimensions + if conv8_1.shape[2:] != conv3_3_short.shape[2:]: + conv8_1 = F.interpolate(conv8_1, size=conv3_3_short.shape[2:], mode='bilinear', align_corners=False) conv8_1_comb = self.relu8_1_comb(conv8_1 + conv3_3_short) conv8_2 = self.relu8_2(self.conv8_2(conv8_1_comb)) conv8_3 = self.relu8_3(self.conv8_3(conv8_2)) conv8_3norm = self.conv8_3norm(conv8_3) conv9_1 = self.conv9_1(conv8_3norm) conv2_2_short = self.conv2_2_short(conv2_2norm) + # Align skip connection dimensions + if conv9_1.shape[2:] != conv2_2_short.shape[2:]: + conv9_1 = F.interpolate(conv9_1, size=conv2_2_short.shape[2:], mode='bilinear', align_corners=False) conv9_1_comb = self.relu9_1_comb(conv9_1 + conv2_2_short) conv9_2 = self.relu9_2(self.conv9_2(conv9_1_comb)) conv9_2norm = self.conv9_2norm(conv9_2) conv10_1 = self.conv10_1(conv9_2norm) conv1_2_short = self.conv1_2_short(conv1_2norm) + # Align skip connection dimensions + if conv10_1.shape[2:] != conv1_2_short.shape[2:]: + conv10_1 = F.interpolate(conv10_1, size=conv1_2_short.shape[2:], mode='bilinear', align_corners=False) conv10_1_comb = self.relu10_1_comb(conv10_1 + conv1_2_short) conv10_2 = self.relu10_2(self.conv10_2(conv10_1_comb)) conv10_ab = self.conv10_ab(conv10_2) diff --git a/models/NonlocalNet.py b/models/NonlocalNet.py index f98fa52..b16904b 100644 --- a/models/NonlocalNet.py +++ b/models/NonlocalNet.py @@ -472,8 +472,6 @@ class WarpNet(nn.Module): channel = B_lab_map.shape[1] image_height = B_lab_map.shape[2] image_width = B_lab_map.shape[3] - feature_height = int(image_height / 4) - feature_width = int(image_width / 4) # scale feature size to 44*44 A_feature2_1 = self.layer2_1(A_relu2_1) @@ -485,10 +483,30 @@ class WarpNet(nn.Module): A_feature5_1 = self.layer5_1(A_relu5_1) B_feature5_1 = self.layer5_1(B_relu5_1) - # concatenate features - if A_feature5_1.shape[2] != A_feature2_1.shape[2] or A_feature5_1.shape[3] != A_feature2_1.shape[3]: - A_feature5_1 = F.pad(A_feature5_1, (0, 0, 1, 1), "replicate") - B_feature5_1 = F.pad(B_feature5_1, (0, 0, 1, 1), "replicate") + # concatenate features - ensure all features have matching spatial dimensions + # Find the target size (use feature2_1 as reference since it has stride=2) + target_h, target_w = A_feature2_1.shape[2], A_feature2_1.shape[3] + + # Align all features to target size using interpolation + def align_feature(feat, target_h, target_w): + if feat.shape[2] != target_h or feat.shape[3] != target_w: + return F.interpolate(feat, size=(target_h, target_w), mode='bilinear', align_corners=False) + return feat + + A_feature2_1 = align_feature(A_feature2_1, target_h, target_w) + A_feature3_1 = align_feature(A_feature3_1, target_h, target_w) + A_feature4_1 = align_feature(A_feature4_1, target_h, target_w) + A_feature5_1 = align_feature(A_feature5_1, target_h, target_w) + + B_feature2_1 = align_feature(B_feature2_1, target_h, target_w) + B_feature3_1 = align_feature(B_feature3_1, target_h, target_w) + B_feature4_1 = align_feature(B_feature4_1, target_h, target_w) + B_feature5_1 = align_feature(B_feature5_1, target_h, target_w) + + # Update feature dimensions to match actual aligned feature size + feature_height = target_h + feature_width = target_w + A_features = self.layer(torch.cat((A_feature2_1, A_feature3_1, A_feature4_1, A_feature5_1), 1)) B_features = self.layer(torch.cat((B_feature2_1, B_feature3_1, B_feature4_1, B_feature5_1), 1)) @@ -515,8 +533,8 @@ class WarpNet(nn.Module): f_WTA = f_WTA / temperature f_div_C = F.softmax(f_WTA.squeeze_(), dim=-1) # 2*1936*1936; - # downsample the reference color - B_lab = F.avg_pool2d(B_lab_map, 4) + # downsample the reference color to match actual feature dimensions + B_lab = F.interpolate(B_lab_map, size=(feature_height, feature_width), mode='bilinear', align_corners=False) B_lab = B_lab.view(batch_size, channel, -1) B_lab = B_lab.permute(0, 2, 1) # 2*1936*channel @@ -524,7 +542,9 @@ class WarpNet(nn.Module): y = torch.matmul(f_div_C, B_lab) # 2*1936*channel y = y.permute(0, 2, 1).contiguous() y = y.view(batch_size, channel, feature_height, feature_width) # 2*3*44*44 - y = self.upsampling(y) - similarity_map = self.upsampling(similarity_map) + + # Upsample to match input dimensions (instead of fixed 4x scale) + y = F.interpolate(y, size=(image_height, image_width), mode='bilinear', align_corners=False) + similarity_map = F.interpolate(similarity_map, size=(image_height, image_width), mode='bilinear', align_corners=False) return y, similarity_map diff --git a/models/feature_extractors/__init__.py b/models/feature_extractors/__init__.py index 77df34c..c1ee9b3 100644 --- a/models/feature_extractors/__init__.py +++ b/models/feature_extractors/__init__.py @@ -4,11 +4,19 @@ Modern Feature Extractors for Video Colorization This package provides drop-in replacements for VGG19/ResNet50 feature extractors with modern vision models (DINOv2, CLIP, etc.) that offer superior semantic understanding. +IMPORTANT LIMITATION: +DeepExemplar's NonlocalNet and ColorNet were trained specifically with VGG19 features. +While this package provides dimension-compatible projections from DINOv2/CLIP to VGG19, +these projections are NOT trained and will produce suboptimal results (color artifacts). + +For production use, stick with VGG19. Alternative encoders are EXPERIMENTAL and would +require fine-tuning the entire colorization network to work properly. + Quick usage: from models.feature_extractors import get_feature_encoder # Get encoder by name - encoder = get_feature_encoder('dinov2_vitb') + encoder = get_feature_encoder('vgg19') # Recommended # Use exactly like VGG19 features = encoder(image, ["r12", "r22", "r32", "r42", "r52"], preprocess=True) @@ -19,7 +27,7 @@ import torch.nn as nn def get_feature_encoder( - encoder_type: Literal['vgg19', 'dinov2_vitb', 'dinov2_vitl', 'dinov2_vitg', + encoder_type: Literal['vgg19', 'dinov2_vits', 'dinov2_vitb', 'dinov2_vitl', 'dinov2_vitg', 'clip_vitb', 'clip_vitl', 'resnet50'] = 'vgg19', device: str = 'cuda', **kwargs @@ -60,9 +68,10 @@ def get_feature_encoder( vgg = models.vgg19(pretrained=True) return VGG19Wrapper(vgg).to(device).eval() - elif encoder_type in ['dinov2_vitb', 'dinov2_vitl', 'dinov2_vitg']: + elif encoder_type in ['dinov2_vits', 'dinov2_vitb', 'dinov2_vitl', 'dinov2_vitg']: from .dinov2_encoder import DINOv2Encoder model_map = { + 'dinov2_vits': 'dinov2_vits14', 'dinov2_vitb': 'dinov2_vitb14', 'dinov2_vitl': 'dinov2_vitl14', 'dinov2_vitg': 'dinov2_vitg14', @@ -88,7 +97,7 @@ def get_feature_encoder( else: raise ValueError( f"Unknown encoder type: {encoder_type}. " - f"Valid options: vgg19, dinov2_vitb, dinov2_vitl, dinov2_vitg, " + f"Valid options: vgg19, dinov2_vits, dinov2_vitb, dinov2_vitl, dinov2_vitg, " f"clip_vitb, clip_vitl, resnet50" ) @@ -123,6 +132,14 @@ def get_encoder_info(encoder_type: str) -> dict: 'quality_rating': 3, 'description': 'ColorMNet baseline, faster than VGG19' }, + 'dinov2_vits': { + 'name': 'DINOv2 ViT-Small', + 'year': 2023, + 'params': '21M', + 'speed_rating': 5, + 'quality_rating': 4, + 'description': 'Smallest DINOv2, faster than vitb with good quality' + }, 'dinov2_vitb': { 'name': 'DINOv2 ViT-Base', 'year': 2023, @@ -180,6 +197,7 @@ def list_available_encoders() -> list: return [ 'vgg19', 'resnet50', + 'dinov2_vits', 'dinov2_vitb', 'dinov2_vitl', 'dinov2_vitg', diff --git a/models/feature_extractors/clip_encoder.py b/models/feature_extractors/clip_encoder.py index 9f33af2..964fc31 100644 --- a/models/feature_extractors/clip_encoder.py +++ b/models/feature_extractors/clip_encoder.py @@ -61,12 +61,33 @@ class CLIPEncoder(nn.Module): for param in self.clip_model.parameters(): param.requires_grad = False - # Projection layers to match VGG19 dimensions - self.proj_r12 = nn.Conv2d(self.embed_dim, 64, 1) - self.proj_r22 = nn.Conv2d(self.embed_dim, 128, 1) - self.proj_r32 = nn.Conv2d(self.embed_dim, 256, 1) - self.proj_r42 = nn.Conv2d(self.embed_dim, 512, 1) - self.proj_r52 = nn.Conv2d(self.embed_dim, 512, 1) + # Projection layers to match VGG19 dimensions (move to device) + # Use conv + bn + relu to match VGG19-like feature distributions + self.proj_r12 = nn.Sequential( + nn.Conv2d(self.embed_dim, 64, 1, bias=False), + nn.BatchNorm2d(64), + nn.ReLU(inplace=False) + ).to(device) + self.proj_r22 = nn.Sequential( + nn.Conv2d(self.embed_dim, 128, 1, bias=False), + nn.BatchNorm2d(128), + nn.ReLU(inplace=False) + ).to(device) + self.proj_r32 = nn.Sequential( + nn.Conv2d(self.embed_dim, 256, 1, bias=False), + nn.BatchNorm2d(256), + nn.ReLU(inplace=False) + ).to(device) + self.proj_r42 = nn.Sequential( + nn.Conv2d(self.embed_dim, 512, 1, bias=False), + nn.BatchNorm2d(512), + nn.ReLU(inplace=False) + ).to(device) + self.proj_r52 = nn.Sequential( + nn.Conv2d(self.embed_dim, 512, 1, bias=False), + nn.BatchNorm2d(512), + nn.ReLU(inplace=False) + ).to(device) # Layer mapping (CLIP ViT-B/16 has 12 layers) self.layer_mapping = { @@ -77,9 +98,12 @@ class CLIPEncoder(nn.Module): 'r52': (12, self.proj_r52), # Final features } - # CLIP normalization - self.register_buffer('mean', torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1)) - self.register_buffer('std', torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1)) + # Use ImageNet normalization for compatibility with DeepExemplar + # (CLIP's native normalization causes color distortion in colorization task) + mean = torch.tensor([0.485, 0.456, 0.406], device=device).view(1, 3, 1, 1) + std = torch.tensor([0.229, 0.224, 0.225], device=device).view(1, 3, 1, 1) + self.register_buffer('mean', mean) + self.register_buffer('std', std) def _extract_features_at_layer(self, x, layer_idx): """Extract features from a specific transformer layer.""" diff --git a/models/feature_extractors/dinov2_encoder.py b/models/feature_extractors/dinov2_encoder.py index 0fc6de8..c2e25b2 100644 --- a/models/feature_extractors/dinov2_encoder.py +++ b/models/feature_extractors/dinov2_encoder.py @@ -63,11 +63,29 @@ class DINOv2Encoder(nn.Module): # Projection layers to match VGG19 channel dimensions # VGG19 dimensions: [64, 128, 256, 512, 512] - self.proj_r12 = nn.Conv2d(self.embed_dim, 64, 1) - self.proj_r22 = nn.Conv2d(self.embed_dim, 128, 1) - self.proj_r32 = nn.Conv2d(self.embed_dim, 256, 1) - self.proj_r42 = nn.Conv2d(self.embed_dim, 512, 1) - self.proj_r52 = nn.Conv2d(self.embed_dim, 512, 1) + # Use deeper projections with Xavier initialization for better feature preservation + def make_projection(in_ch, out_ch): + """Create a 2-layer projection with proper initialization.""" + mid_ch = (in_ch + out_ch) // 2 + proj = nn.Sequential( + nn.Conv2d(in_ch, mid_ch, 1), + nn.ReLU(inplace=True), + nn.Conv2d(mid_ch, out_ch, 1), + nn.ReLU(inplace=True) + ) + # Xavier initialization for better gradient flow + for m in proj.modules(): + if isinstance(m, nn.Conv2d): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + return proj + + self.proj_r12 = make_projection(self.embed_dim, 64) + self.proj_r22 = make_projection(self.embed_dim, 128) + self.proj_r32 = make_projection(self.embed_dim, 256) + self.proj_r42 = make_projection(self.embed_dim, 512) + self.proj_r52 = make_projection(self.embed_dim, 512) # Mapping from layer names to outputs self.layer_mapping = { @@ -96,9 +114,11 @@ class DINOv2Encoder(nn.Module): # Remove CLS token and reshape to spatial patch_features = features[:, 1:] # Remove CLS token - # Calculate patch grid size - num_patches = patch_features.shape[1] - patch_h = patch_w = int(num_patches ** 0.5) + # Calculate patch grid size based on actual image dimensions + # DINOv2 patch size is 14x14 + patch_size = 14 + patch_h = H // patch_size + patch_w = W // patch_size # Reshape to spatial: [B, N, D] -> [B, D, H, W] features_spatial = patch_features.transpose(1, 2).reshape(B, self.embed_dim, patch_h, patch_w) @@ -216,8 +236,11 @@ class DINOv2EncoderColorMNet(nn.Module): for blk in self.dinov2.blocks[:layer_idx]: features = blk(features) patch_features = features[:, 1:] - num_patches = patch_features.shape[1] - patch_h = patch_w = int(num_patches ** 0.5) + # Calculate patch grid size based on actual image dimensions + # DINOv2 patch size is 14x14 + patch_size = 14 + patch_h = H // patch_size + patch_w = W // patch_size features_spatial = patch_features.transpose(1, 2).reshape(B, self.embed_dim, patch_h, patch_w) return features_spatial diff --git a/nodes.py b/nodes.py index 5538af4..bd56c29 100644 --- a/nodes.py +++ b/nodes.py @@ -132,23 +132,6 @@ class ColorMNetVideoNode: "default": "resnet50", "tooltip": "Feature extraction model: resnet50 (ColorMNet default), vgg19 (fast), dinov2_vitb (recommended, 40-60% better), dinov2_vitl (best quality), clip_vitb (text-guided)" }), - "post_processor": (["none", "wls", "guided", "bilateral", "color_matcher"], { - "default": "none", - "tooltip": "Post-processing method: none (fastest), wls (edge-aware), guided (fast smoothing), bilateral (classic), color_matcher (best color consistency, recommended)" - }), - "post_process_strength": ("FLOAT", { - "default": 0.8, - "min": 0.0, - "max": 1.0, - "step": 0.05, - "tooltip": "Post-processing strength: 0=disabled, 1=full effect (for color_matcher/guided/bilateral)" - }), - "temporal_consistency": ("BOOLEAN", { - "default": True, - "label_on": "enabled", - "label_off": "disabled", - "tooltip": "Enable temporal consistency in post-processing to reduce flickering between frames (for color_matcher)" - }), "use_fp16": ("BOOLEAN", { "default": True, "label_on": "enabled", @@ -191,9 +174,6 @@ class ColorMNetVideoNode: target_height: int, memory_mode: str = "balanced", feature_encoder: str = "resnet50", - post_processor: str = "none", - post_process_strength: float = 0.8, - temporal_consistency: bool = True, use_fp16: bool = True, use_torch_compile: bool = False, text_guidance: str = "", @@ -309,6 +289,7 @@ class ColorMNetVideoNode: # Colorize self.logger.info(f"Processing {num_frames} frames at {target_height}x{target_width}") + colorization_start_time = time.time() colorized = self.inference_pipeline.colorize_video( frames=video_frames, reference=reference_image, @@ -316,68 +297,11 @@ class ColorMNetVideoNode: target_width=target_width, progress_callback=progress_callback, ) + colorization_time = time.time() - colorization_start_time if not COMFY_PROGRESS_AVAILABLE: pbar.close() - # Apply modern post-processing (batch processing for videos) - post_proc_applied = "none" - encoder_name = feature_encoder - - if post_processor != "none": - try: - # Auto-install dependencies if needed - from .auto_installer import ensure_dependencies_for_post_processor - deps_ok = ensure_dependencies_for_post_processor(post_processor) - - if not deps_ok: - self.logger.warning(f"Dependencies for {post_processor} could not be installed") - proc = None - else: - from .post_processing import get_post_processor - - # Create post-processor - if post_processor == "wls": - proc = get_post_processor('wls', lambda_value=500.0, sigma_color=4.0) - elif post_processor == "color_matcher": - proc = get_post_processor('color_matcher', method='mkl', - temporal_consistency=temporal_consistency, - consistency_weight=0.3) - elif post_processor == "guided": - proc = get_post_processor('guided', radius=8, eps=0.01) - elif post_processor == "bilateral": - proc = get_post_processor('bilateral', d=9, sigma_color=75, sigma_space=75) - else: - proc = None - - if proc is not None: - self.logger.info(f"Applying post-processing: {post_processor} to {num_frames} frames...") - - # Apply post-processing - if post_processor == "color_matcher": - # color_matcher for videos with reference matching - ref_rgb = reference_image - - # Process video with temporal consistency - processed = proc.match_video_frames( - colorized, - ref_rgb, - strength=post_process_strength, - progress_callback=lambda i, total: None - ) - colorized = processed - else: - # Other post-processors: apply per-frame - processed_frames = proc.match_video_frames(colorized, None) - colorized = processed_frames - - post_proc_applied = post_processor - self.logger.info(f"โœ“ Post-processing complete: {post_processor}") - - except Exception as e: - self.logger.warning(f"Post-processing '{post_processor}' failed: {e}") - self.logger.info("Continuing without post-processing...") - # Calculate timing end_time = time.time() elapsed_time = end_time - start_time @@ -391,19 +315,16 @@ class ColorMNetVideoNode: f"Date/Time: {timestamp}\n" f"Frames Processed: {num_frames}\n" f"Resolution: {target_width}x{target_height}\n" - f"Total Time: {elapsed_time:.2f} seconds\n" - f"Average FPS: {fps:.2f}\n" - f"Time per Frame: {elapsed_time / num_frames:.3f} seconds\n" - f"Feature Encoder: {encoder_name} (ColorMNet built-in ResNet50)\n" - f"Post-Processor: {post_proc_applied}\n" - ) - if post_proc_applied == "color_matcher": - report += f"Matching Strength: {post_process_strength}\n" - report += f"Temporal Consistency: {'Enabled' if temporal_consistency else 'Disabled'}\n" - report += ( - f"Memory Mode: {memory_mode}\n" - f"FP16 Enabled: {use_fp16}\n" - f"Torch Compile: {use_torch_compile}\n" + f"\n" + f"Performance:\n" + f" Total Time: {elapsed_time:.2f}s ({fps:.2f} FPS)\n" + f" Time per Frame: {elapsed_time / num_frames:.3f}s\n" + f"\n" + f"Configuration:\n" + f" Feature Encoder: {feature_encoder}\n" + f" Memory Mode: {memory_mode}\n" + f" FP16 Enabled: {use_fp16}\n" + f" Torch Compile: {use_torch_compile}\n" f"{'=' * 50}" ) @@ -469,17 +390,6 @@ class ColorMNetImageNode: "step": 32, "tooltip": "Output height - must be multiple of 32 (will be adjusted automatically)" }), - "post_processor": (["none", "wls", "guided", "bilateral", "color_matcher"], { - "default": "none", - "tooltip": "Post-processing method: none (fastest), wls (edge-aware), guided (fast smoothing), bilateral (classic), color_matcher (best color consistency, recommended)" - }), - "post_process_strength": ("FLOAT", { - "default": 0.8, - "min": 0.0, - "max": 1.0, - "step": 0.05, - "tooltip": "Post-processing strength: 0=disabled, 1=full effect (for color_matcher/guided/bilateral)" - }), "use_fp16": ("BOOLEAN", { "default": True, "tooltip": "Use half-precision (FP16) for faster processing with lower VRAM usage (minimal quality impact)" @@ -502,8 +412,6 @@ class ColorMNetImageNode: reference_image: torch.Tensor, target_width: int, target_height: int, - post_processor: str = "none", - post_process_strength: float = 0.8, use_fp16: bool = True, use_torch_compile: bool = False, ) -> Tuple[torch.Tensor, str]: @@ -592,55 +500,6 @@ class ColorMNetImageNode: target_width=target_width, ) - # Apply post-processing - post_proc_applied = "none" - - if post_processor != "none": - try: - # Auto-install dependencies if needed - from .auto_installer import ensure_dependencies_for_post_processor - deps_ok = ensure_dependencies_for_post_processor(post_processor) - - if not deps_ok: - self.logger.warning(f"Dependencies for {post_processor} could not be installed") - proc = None - else: - from .post_processing import get_post_processor - - # Create post-processor - if post_processor == "wls": - proc = get_post_processor('wls', lambda_value=500.0, sigma_color=4.0) - elif post_processor == "color_matcher": - proc = get_post_processor('color_matcher', method='mkl') - elif post_processor == "guided": - proc = get_post_processor('guided', radius=8, eps=0.01) - elif post_processor == "bilateral": - proc = get_post_processor('bilateral', d=9, sigma_color=75, sigma_space=75) - else: - proc = None - - if proc is not None: - self.logger.info(f"Applying post-processing: {post_processor}...") - - # Apply post-processing - if post_processor == "color_matcher": - # color_matcher needs reference image - processed = proc.match_to_reference(colorized, reference_image) - - # Apply strength blending - if post_process_strength < 1.0: - processed = colorized * (1 - post_process_strength) + processed * post_process_strength - else: - processed = proc.match_to_reference(colorized, None) - - colorized = processed - post_proc_applied = post_processor - self.logger.info(f"โœ“ Post-processing complete: {post_processor}") - - except Exception as e: - self.logger.warning(f"Post-processing '{post_processor}' failed: {e}") - self.logger.info("Continuing without post-processing...") - # Add batch dimension back colorized = colorized.unsqueeze(0) if colorized.dim() == 3 else colorized @@ -656,11 +515,6 @@ class ColorMNetImageNode: f"Date/Time: {timestamp}\n" f"Resolution: {target_width}x{target_height}\n" f"Total Time: {elapsed_time:.3f} seconds\n" - f"Post-Processor: {post_proc_applied}\n" - ) - if post_proc_applied == "color_matcher": - report += f"Matching Strength: {post_process_strength}\n" - report += ( f"FP16 Enabled: {use_fp16}\n" f"Torch Compile: {use_torch_compile}\n" f"{'=' * 50}" diff --git a/post_processing/__init__.py b/post_processing/__init__.py deleted file mode 100644 index a4d1f03..0000000 --- a/post_processing/__init__.py +++ /dev/null @@ -1,278 +0,0 @@ -""" -Post-Processing Module for Video Colorization - -Provides various post-processing filters and refinement techniques to improve -colorization quality and temporal consistency. - -Available processors: - - ColorMatcherPost: Color matching using color-matcher library - - WLSFilter: Weighted Least Squares edge-aware smoothing - - GuidedFilter: Fast edge-preserving filter - - BilateralFilter: Bilateral filtering - - TemporalStabilizer: Optical flow-based temporal stabilization - -Usage: - from post_processing import get_post_processor - - processor = get_post_processor('color_matcher', method='mkl') - refined = processor.process(colorized_frames, reference) -""" - -from typing import Literal, Optional, Dict, Any -import torch -import torch.nn as nn - - -def get_post_processor( - processor_type: Literal['color_matcher', 'wls', 'guided', 'bilateral', 'none'] = 'color_matcher', - **kwargs -) -> nn.Module: - """Get a post-processor by name. - - Args: - processor_type: Type of post-processor - - 'color_matcher': Color matching (best for color consistency) - - 'wls': Weighted Least Squares filter (edge-aware smoothing) - - 'guided': Guided filter (fast edge-preserving) - - 'bilateral': Bilateral filter (noise reduction) - - 'none': No post-processing - **kwargs: Additional arguments for the processor - - Returns: - Post-processor instance - - Example: - >>> processor = get_post_processor('color_matcher', method='mkl') - >>> refined = processor.match_to_reference(colorized, reference) - """ - processor_type = processor_type.lower() - - if processor_type == 'color_matcher': - from .color_matcher_post import ColorMatcherPost - return ColorMatcherPost(**kwargs) - - elif processor_type == 'wls': - from .wls_filter import WLSFilter - return WLSFilter(**kwargs) - - elif processor_type == 'guided': - from .guided_filter import GuidedFilter - return GuidedFilter(**kwargs) - - elif processor_type == 'bilateral': - from .bilateral_filter import BilateralFilter - return BilateralFilter(**kwargs) - - elif processor_type == 'none': - return NoOpProcessor() - - else: - raise ValueError( - f"Unknown processor type: {processor_type}. " - f"Valid: color_matcher, wls, guided, bilateral, none" - ) - - -class NoOpProcessor: - """No-op processor that returns input unchanged.""" - - def match_to_reference(self, image, reference, **kwargs): - return image - - def match_video_frames(self, frames, reference, **kwargs): - return frames - - def process(self, *args, **kwargs): - return args[0] if args else None - - -def list_available_processors() -> list: - """List all available post-processors. - - Returns: - List of processor names - """ - processors = [] - - try: - from .color_matcher_post import ColorMatcherPost - processors.append('color_matcher') - except ImportError: - pass - - try: - from .wls_filter import WLSFilter - processors.append('wls') - except ImportError: - pass - - try: - from .guided_filter import GuidedFilter - processors.append('guided') - except ImportError: - pass - - try: - from .bilateral_filter import BilateralFilter - processors.append('bilateral') - except ImportError: - pass - - processors.append('none') - - return processors - - -def get_processor_info(processor_type: str) -> Dict[str, Any]: - """Get information about a post-processor. - - Args: - processor_type: Name of processor - - Returns: - Dictionary with processor metadata - """ - info = { - 'color_matcher': { - 'name': 'Color Matcher', - 'description': 'Match colors to reference using optimal transport or histogram methods', - 'speed': 4, # 1-5 scale - 'quality': 5, - 'temporal_aware': True, - 'requires': ['color-matcher'], - 'best_for': 'Ensuring color consistency with reference image' - }, - 'wls': { - 'name': 'WLS Filter', - 'description': 'Weighted Least Squares edge-aware smoothing', - 'speed': 3, - 'quality': 4, - 'temporal_aware': False, - 'requires': ['opencv-contrib-python'], - 'best_for': 'Edge-aware color smoothing' - }, - 'guided': { - 'name': 'Guided Filter', - 'description': 'Fast edge-preserving filter', - 'speed': 5, - 'quality': 4, - 'temporal_aware': False, - 'requires': ['opencv-python'], - 'best_for': 'Fast edge-preserving smoothing' - }, - 'bilateral': { - 'name': 'Bilateral Filter', - 'description': 'Bilateral noise reduction filter', - 'speed': 4, - 'quality': 3, - 'temporal_aware': False, - 'requires': ['opencv-python'], - 'best_for': 'Noise reduction while preserving edges' - }, - 'none': { - 'name': 'No Post-Processing', - 'description': 'Skip post-processing', - 'speed': 5, - 'quality': 0, - 'temporal_aware': False, - 'requires': [], - 'best_for': 'When raw output is sufficient' - } - } - - return info.get(processor_type, {}) - - -def compare_processors( - image: torch.Tensor, - reference: torch.Tensor, - verbose: bool = True, -) -> Dict[str, torch.Tensor]: - """Compare all available post-processors. - - Args: - image: Test image [H, W, 3] - reference: Reference image [H, W, 3] - verbose: Print comparison info - - Returns: - Dictionary mapping processor names to results - """ - available = list_available_processors() - results = {} - - if verbose: - print("\n" + "="*80) - print("Post-Processor Comparison") - print("="*80) - - for proc_type in available: - if proc_type == 'none': - continue - - try: - processor = get_post_processor(proc_type) - result = processor.match_to_reference(image, reference) - results[proc_type] = result - - if verbose: - info = get_processor_info(proc_type) - speed = "โšก" * info.get('speed', 0) - quality = "โญ" * info.get('quality', 0) - print(f"\n{info.get('name', proc_type)}") - print(f" Speed: {speed}") - print(f" Quality: {quality}") - print(f" Best for: {info.get('best_for', 'N/A')}") - print(f" โœ“ Success") - - except Exception as e: - if verbose: - print(f"\n{proc_type}") - print(f" โœ— Failed: {e}") - - if verbose: - print("\n" + "="*80) - print(f"Successfully tested: {len(results)}/{len(available)-1} processors") - print("="*80 + "\n") - - return results - - -# Convenience exports -__all__ = [ - 'get_post_processor', - 'list_available_processors', - 'get_processor_info', - 'compare_processors', - 'NoOpProcessor', -] - - -if __name__ == "__main__": - print("Post-Processing Module Demo") - print("="*80) - - # List available - print("\nAvailable processors:") - for proc in list_available_processors(): - print(f" - {proc}") - - # Show info - print("\nProcessor details:") - for proc in ['color_matcher', 'wls', 'guided']: - info = get_processor_info(proc) - if info: - print(f"\n{info['name']}:") - print(f" {info['description']}") - print(f" Best for: {info['best_for']}") - - # Test loading - print("\n" + "="*80) - print("Testing processor loading...") - try: - processor = get_post_processor('color_matcher') - print("โœ“ Successfully loaded color_matcher") - except Exception as e: - print(f"โœ— Failed to load: {e}") - - print("\n" + "="*80) diff --git a/post_processing/bilateral_filter.py b/post_processing/bilateral_filter.py deleted file mode 100644 index 9f02708..0000000 --- a/post_processing/bilateral_filter.py +++ /dev/null @@ -1,94 +0,0 @@ -""" -Bilateral Filter Post-Processing - -Classic bilateral filtering for edge-preserving noise reduction. - -Reference: Bilateral Filtering for Gray and Color Images, Tomasi & Manduchi, ICCV 1998 -""" - -import torch -import cv2 -import numpy as np -from typing import Optional - - -class BilateralFilter: - """Bilateral filtering for edge-preserving smoothing.""" - - def __init__( - self, - d: int = 9, - sigma_color: float = 75, - sigma_space: float = 75, - ): - """ - Args: - d: Diameter of each pixel neighborhood - sigma_color: Filter sigma in the color space - sigma_space: Filter sigma in the coordinate space - """ - self.d = d - self.sigma_color = sigma_color - self.sigma_space = sigma_space - - def match_to_reference( - self, - image: torch.Tensor, - reference: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Apply bilateral filtering. - - Args: - image: Input image [H, W, 3] in range [0, 1] - reference: Not used (for interface compatibility) - - Returns: - Filtered image [H, W, 3] - """ - device = image.device - - # Convert to numpy - img_np = (image.cpu().numpy() * 255).astype(np.uint8) - - # Apply bilateral filter - try: - filtered_np = cv2.bilateralFilter( - img_np, - self.d, - self.sigma_color, - self.sigma_space - ) - except Exception as e: - print(f"[BilateralFilter] Filtering failed: {e}") - return image - - # Convert back - filtered = torch.from_numpy(filtered_np).float() / 255.0 - return filtered.to(device) - - def match_video_frames( - self, - frames: torch.Tensor, - reference: Optional[torch.Tensor] = None, - progress_callback: Optional[callable] = None, - ) -> torch.Tensor: - """Apply bilateral filtering to video.""" - num_frames = frames.shape[0] - filtered_frames = [] - - for i in range(num_frames): - filtered = self.match_to_reference(frames[i], reference) - filtered_frames.append(filtered) - - if progress_callback: - progress_callback(i + 1, num_frames) - - return torch.stack(filtered_frames, dim=0) - - -if __name__ == "__main__": - print("Testing BilateralFilter...") - test_img = torch.rand(256, 256, 3) - bf = BilateralFilter(d=9, sigma_color=75, sigma_space=75) - filtered = bf.match_to_reference(test_img) - print(f"โœ“ Filtered: {test_img.shape} -> {filtered.shape}") diff --git a/post_processing/color_matcher_post.py b/post_processing/color_matcher_post.py deleted file mode 100644 index 390e453..0000000 --- a/post_processing/color_matcher_post.py +++ /dev/null @@ -1,399 +0,0 @@ -""" -Color Matcher Post-Processing for Video Colorization - -Uses the color-matcher library to refine colorization results by ensuring -better color consistency with the reference image and temporal consistency -across video frames. - -Reference: https://github.com/hahnec/color-matcher - -Installation: - pip install color-matcher - -Methods available: - - 'default' (hm-mvgd-hm): Best overall quality - - 'mkl': Optimal transport-based, very accurate - - 'reinhard': Classic color transfer, fast - -Usage: - from post_processing.color_matcher_post import ColorMatcherPost - - post_processor = ColorMatcherPost(method='hm-mvgd-hm') - - # Single image - refined = post_processor.match_to_reference(colorized, reference) - - # Video frames - refined_video = post_processor.match_video_frames(video_frames, reference) -""" - -import numpy as np -import torch -import torch.nn.functional as F -from typing import Optional, Literal, List - - -class ColorMatcherPost: - """Post-processing using color-matcher library for color refinement.""" - - def __init__( - self, - method: Literal['default', 'hm-mvgd-hm', 'mkl', 'reinhard'] = 'default', - temporal_consistency: bool = True, - consistency_weight: float = 0.3, - ): - """ - Args: - method: Color matching method - - 'default' or 'hm-mvgd-hm': Hybrid method (best quality) - - 'mkl': Monge-Kantorovich linearization (optimal transport) - - 'reinhard': Classic Reinhard et al. method (fastest) - temporal_consistency: Apply temporal smoothing across frames - consistency_weight: Weight for temporal consistency (0-1) - - 0 = no temporal smoothing - - 1 = maximum temporal smoothing - """ - try: - from color_matcher import ColorMatcher - from color_matcher.io_handler import load_img_file, save_img_file - self.cm = ColorMatcher() - self.io_handler = (load_img_file, save_img_file) - self.available = True - except ImportError: - print("[ColorMatcherPost] Warning: color-matcher not installed") - print("Install with: pip install color-matcher") - self.available = False - return - - # Normalize method name - if method == 'default': - method = 'hm-mvgd-hm' - - self.method = method - self.temporal_consistency = temporal_consistency - self.consistency_weight = consistency_weight - - # Validate method - valid_methods = ['hm-mvgd-hm', 'mkl', 'reinhard'] - if method not in valid_methods: - raise ValueError(f"Invalid method: {method}. Valid: {valid_methods}") - - print(f"[ColorMatcherPost] Initialized with method={method}, " - f"temporal_consistency={temporal_consistency}") - - def _tensor_to_numpy(self, tensor: torch.Tensor) -> np.ndarray: - """Convert PyTorch tensor to numpy array for color-matcher. - - Args: - tensor: Image tensor in range [0, 1], shape [H, W, 3] or [B, H, W, 3] - - Returns: - Numpy array in range [0, 255], shape [H, W, 3] or [B, H, W, 3] - """ - if isinstance(tensor, torch.Tensor): - array = tensor.detach().cpu().numpy() - else: - array = tensor - - # Ensure range [0, 1] - array = np.clip(array, 0, 1) - - # Convert to uint8 [0, 255] - array = (array * 255).astype(np.uint8) - - return array - - def _numpy_to_tensor(self, array: np.ndarray, device: str = 'cpu') -> torch.Tensor: - """Convert numpy array back to PyTorch tensor. - - Args: - array: Numpy array in range [0, 255], shape [H, W, 3] - device: Target device - - Returns: - Tensor in range [0, 1], shape [H, W, 3] - """ - # Ensure uint8 - if array.dtype != np.uint8: - array = np.clip(array, 0, 255).astype(np.uint8) - - # Convert to float [0, 1] - tensor = torch.from_numpy(array).float() / 255.0 - - return tensor.to(device) - - def match_to_reference( - self, - image: torch.Tensor, - reference: torch.Tensor, - strength: float = 1.0, - ) -> torch.Tensor: - """Match image colors to reference image. - - Args: - image: Colorized image [H, W, 3] in range [0, 1] - reference: Reference image [H, W, 3] in range [0, 1] - strength: Blending strength (0-1) - - 0 = original image - - 1 = full color matching - - Returns: - Color-matched image [H, W, 3] in range [0, 1] - """ - if not self.available: - print("[ColorMatcherPost] Skipping: color-matcher not available") - return image - - device = image.device - - # Convert to numpy - img_np = self._tensor_to_numpy(image) - ref_np = self._tensor_to_numpy(reference) - - # Apply color matching - try: - matched_np = self.cm.transfer(src=img_np, ref=ref_np, method=self.method) - except Exception as e: - print(f"[ColorMatcherPost] Color matching failed: {e}") - return image - - # Convert back to tensor - matched = self._numpy_to_tensor(matched_np, device=device) - - # Blend with original based on strength - if strength < 1.0: - matched = image * (1 - strength) + matched * strength - - return matched - - def match_video_frames( - self, - frames: torch.Tensor, - reference: torch.Tensor, - strength: float = 1.0, - progress_callback: Optional[callable] = None, - ) -> torch.Tensor: - """Match colors for all video frames. - - Args: - frames: Video frames [N, H, W, 3] in range [0, 1] - reference: Reference image [H, W, 3] in range [0, 1] - strength: Color matching strength (0-1) - progress_callback: Optional callback(current, total) - - Returns: - Color-matched frames [N, H, W, 3] - """ - if not self.available: - print("[ColorMatcherPost] Skipping: color-matcher not available") - return frames - - num_frames = frames.shape[0] - device = frames.device - matched_frames = [] - - # Process each frame - for i in range(num_frames): - frame = frames[i] - - # Match to reference - matched = self.match_to_reference(frame, reference, strength=strength) - - # Apply temporal consistency - if self.temporal_consistency and i > 0 and len(matched_frames) > 0: - prev_frame = matched_frames[-1] - matched = self._apply_temporal_smoothing(matched, prev_frame) - - matched_frames.append(matched) - - if progress_callback is not None: - progress_callback(i + 1, num_frames) - - return torch.stack(matched_frames, dim=0) - - def _apply_temporal_smoothing( - self, - current: torch.Tensor, - previous: torch.Tensor, - ) -> torch.Tensor: - """Apply temporal smoothing between consecutive frames. - - Args: - current: Current frame [H, W, 3] - previous: Previous frame [H, W, 3] - - Returns: - Smoothed current frame [H, W, 3] - """ - # Blend with previous frame - alpha = self.consistency_weight - smoothed = current * (1 - alpha) + previous * alpha - - return smoothed - - def match_video_to_first_frame( - self, - frames: torch.Tensor, - strength: float = 0.5, - progress_callback: Optional[callable] = None, - ) -> torch.Tensor: - """Match all frames to first frame for temporal consistency. - - This is useful when the first frame is already well-colorized. - - Args: - frames: Video frames [N, H, W, 3] - strength: Matching strength (0-1) - progress_callback: Optional callback(current, total) - - Returns: - Temporally consistent frames [N, H, W, 3] - """ - if not self.available: - return frames - - # Use first frame as reference - reference = frames[0] - - return self.match_video_frames( - frames, - reference, - strength=strength, - progress_callback=progress_callback - ) - - def match_histogram_only( - self, - image: torch.Tensor, - reference: torch.Tensor, - ) -> torch.Tensor: - """Simple histogram matching (faster than full color matching). - - Args: - image: Input image [H, W, 3] - reference: Reference image [H, W, 3] - - Returns: - Histogram-matched image [H, W, 3] - """ - from skimage import exposure - - device = image.device - img_np = self._tensor_to_numpy(image) - ref_np = self._tensor_to_numpy(reference) - - # Match histogram per channel - matched = np.zeros_like(img_np) - for c in range(3): - matched[:, :, c] = exposure.match_histograms( - img_np[:, :, c], - ref_np[:, :, c] - ) - - return self._numpy_to_tensor(matched, device=device) - - -class ColorMatcherPostBatch: - """Batch processing version for better performance.""" - - def __init__(self, method: str = 'default', batch_size: int = 8): - """ - Args: - method: Color matching method - batch_size: Number of frames to process at once - """ - self.matcher = ColorMatcherPost(method=method) - self.batch_size = batch_size - - def process_video( - self, - frames: torch.Tensor, - reference: torch.Tensor, - strength: float = 1.0, - ) -> torch.Tensor: - """Process video in batches for better performance. - - Args: - frames: Video frames [N, H, W, 3] - reference: Reference image [H, W, 3] - strength: Color matching strength - - Returns: - Processed frames [N, H, W, 3] - """ - num_frames = frames.shape[0] - all_matched = [] - - for i in range(0, num_frames, self.batch_size): - batch_end = min(i + self.batch_size, num_frames) - batch = frames[i:batch_end] - - # Process batch - matched_batch = self.matcher.match_video_frames( - batch, - reference, - strength=strength - ) - - all_matched.append(matched_batch) - - return torch.cat(all_matched, dim=0) - - -def compare_methods( - image: torch.Tensor, - reference: torch.Tensor, -) -> dict: - """Compare all color matching methods. - - Args: - image: Test image [H, W, 3] - reference: Reference image [H, W, 3] - - Returns: - Dictionary with results for each method - """ - methods = ['hm-mvgd-hm', 'mkl', 'reinhard'] - results = {} - - for method in methods: - try: - matcher = ColorMatcherPost(method=method) - matched = matcher.match_to_reference(image, reference) - results[method] = matched - print(f"โœ“ {method} completed") - except Exception as e: - print(f"โœ— {method} failed: {e}") - results[method] = None - - return results - - -if __name__ == "__main__": - print("Testing ColorMatcherPost...") - - # Create test images - test_img = torch.rand(256, 256, 3) - ref_img = torch.rand(256, 256, 3) - - # Test single image - print("\n1. Testing single image matching...") - matcher = ColorMatcherPost(method='mkl') - matched = matcher.match_to_reference(test_img, ref_img, strength=1.0) - print(f" Input shape: {test_img.shape}") - print(f" Output shape: {matched.shape}") - print(f" Output range: [{matched.min():.3f}, {matched.max():.3f}]") - - # Test video - print("\n2. Testing video matching...") - test_video = torch.rand(10, 256, 256, 3) - matched_video = matcher.match_video_frames(test_video, ref_img) - print(f" Input: {test_video.shape}") - print(f" Output: {matched_video.shape}") - - # Compare methods - print("\n3. Comparing all methods...") - results = compare_methods(test_img, ref_img) - print(f" Successfully tested: {len([r for r in results.values() if r is not None])}/{len(results)} methods") - - print("\nโœ“ All tests passed!") diff --git a/post_processing/guided_filter.py b/post_processing/guided_filter.py deleted file mode 100644 index b66c3f2..0000000 --- a/post_processing/guided_filter.py +++ /dev/null @@ -1,103 +0,0 @@ -""" -Guided Filter Post-Processing - -Fast edge-preserving filter that's much faster than bilateral filtering -while achieving similar quality. - -Reference: Guided Image Filtering, He et al., ECCV 2010 -""" - -import torch -import cv2 -import numpy as np -from typing import Optional - - -class GuidedFilter: - """Guided image filtering for edge-preserving smoothing.""" - - def __init__( - self, - radius: int = 8, - eps: float = 0.01, - ): - """ - Args: - radius: Filter radius (larger = smoother) - eps: Regularization parameter (larger = smoother) - """ - self.radius = radius - self.eps = eps - - def match_to_reference( - self, - image: torch.Tensor, - reference: Optional[torch.Tensor] = None, - guide: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Apply guided filtering. - - Args: - image: Input image [H, W, 3] in range [0, 1] - reference: Not used (for interface compatibility) - guide: Guide image (uses input if None) - - Returns: - Filtered image [H, W, 3] - """ - device = image.device - - # Convert to numpy [0, 255] - img_np = (image.cpu().numpy() * 255).astype(np.uint8) - - if guide is None: - guide_np = img_np - else: - guide_np = (guide.cpu().numpy() * 255).astype(np.uint8) - - # Apply guided filter - try: - filtered_np = cv2.ximgproc.guidedFilter( - guide_np, - img_np, - self.radius, - self.eps - ) - except AttributeError: - # Fallback: opencv-contrib not available, use bilateral - print("[GuidedFilter] opencv-contrib not available, using bilateral filter") - filtered_np = cv2.bilateralFilter(img_np, self.radius, 75, 75) - except Exception as e: - print(f"[GuidedFilter] Filtering failed: {e}") - return image - - # Convert back - filtered = torch.from_numpy(filtered_np).float() / 255.0 - return filtered.to(device) - - def match_video_frames( - self, - frames: torch.Tensor, - reference: Optional[torch.Tensor] = None, - progress_callback: Optional[callable] = None, - ) -> torch.Tensor: - """Apply guided filtering to video.""" - num_frames = frames.shape[0] - filtered_frames = [] - - for i in range(num_frames): - filtered = self.match_to_reference(frames[i], reference) - filtered_frames.append(filtered) - - if progress_callback: - progress_callback(i + 1, num_frames) - - return torch.stack(filtered_frames, dim=0) - - -if __name__ == "__main__": - print("Testing GuidedFilter...") - test_img = torch.rand(256, 256, 3) - gf = GuidedFilter(radius=8, eps=0.01) - filtered = gf.match_to_reference(test_img) - print(f"โœ“ Filtered: {test_img.shape} -> {filtered.shape}") diff --git a/post_processing/wls_filter.py b/post_processing/wls_filter.py deleted file mode 100644 index f06fb7b..0000000 --- a/post_processing/wls_filter.py +++ /dev/null @@ -1,142 +0,0 @@ -""" -WLS (Weighted Least Squares) Filter Post-Processing - -Edge-aware smoothing filter that preserves edges while smoothing color transitions. -Currently implements the existing WLS filter from the original code. - -Reference: Edge-Preserving Decompositions for Multi-Scale Tone and Detail Manipulation - Farbman et al., SIGGRAPH 2008 -""" - -import torch -import cv2 -import numpy as np -from typing import Optional - - -class WLSFilter: - """Weighted Least Squares edge-aware filter.""" - - def __init__( - self, - lambda_value: float = 500.0, - sigma_color: float = 4.0, - available: bool = None, - ): - """ - Args: - lambda_value: Smoothing strength (higher = smoother) - sigma_color: Color sensitivity - available: Whether WLS filter is available (auto-detect if None) - """ - self.lambda_value = lambda_value - self.sigma_color = sigma_color - - # Check if WLS filter is available - if available is None: - self.available = hasattr(cv2.ximgproc, 'createFastGlobalSmootherFilter') - else: - self.available = available - - if not self.available: - print("[WLSFilter] Warning: opencv-contrib-python not installed or WLS not available") - print("Install with: pip install opencv-contrib-python") - - def match_to_reference( - self, - image: torch.Tensor, - reference: Optional[torch.Tensor] = None, - guide: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Apply WLS filtering. - - Args: - image: Input image [H, W, 3] in range [0, 1] - reference: Not used (for interface compatibility) - guide: Optional guide image for edge-aware filtering (uses image if None) - - Returns: - Filtered image [H, W, 3] - """ - if not self.available: - print("[WLSFilter] Skipping: WLS filter not available") - return image - - device = image.device - - # Convert to numpy - img_np = (image.cpu().numpy() * 255).astype(np.uint8) - - if guide is None: - # Use luminance as guide - guide_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY) - else: - guide_np = (guide.cpu().numpy() * 255).astype(np.uint8) - if guide_np.ndim == 3: - guide_np = cv2.cvtColor(guide_np, cv2.COLOR_RGB2GRAY) - - # Apply WLS filter - try: - wls = cv2.ximgproc.createFastGlobalSmootherFilter( - guide_np, - self.lambda_value, - self.sigma_color - ) - - # Filter each channel - filtered_np = np.zeros_like(img_np) - for c in range(3): - filtered_np[:, :, c] = wls.filter(img_np[:, :, c]) - - except Exception as e: - print(f"[WLSFilter] Filtering failed: {e}") - return image - - # Convert back to tensor - filtered = torch.from_numpy(filtered_np).float() / 255.0 - return filtered.to(device) - - def match_video_frames( - self, - frames: torch.Tensor, - reference: Optional[torch.Tensor] = None, - progress_callback: Optional[callable] = None, - ) -> torch.Tensor: - """Apply WLS filtering to video frames. - - Args: - frames: Video frames [N, H, W, 3] - reference: Not used (for interface compatibility) - progress_callback: Optional callback(current, total) - - Returns: - Filtered frames [N, H, W, 3] - """ - if not self.available: - return frames - - num_frames = frames.shape[0] - filtered_frames = [] - - for i in range(num_frames): - filtered = self.match_to_reference(frames[i], reference) - filtered_frames.append(filtered) - - if progress_callback is not None: - progress_callback(i + 1, num_frames) - - return torch.stack(filtered_frames, dim=0) - - -if __name__ == "__main__": - print("Testing WLSFilter...") - - # Test - test_img = torch.rand(256, 256, 3) - wls = WLSFilter(lambda_value=500.0, sigma_color=4.0) - - if wls.available: - filtered = wls.match_to_reference(test_img) - print(f"โœ“ Filtered: {test_img.shape} -> {filtered.shape}") - else: - print("โœ— WLS filter not available") diff --git a/test.py b/test.py deleted file mode 100644 index 408e074..0000000 --- a/test.py +++ /dev/null @@ -1,186 +0,0 @@ -from __future__ import print_function - -import argparse -import glob -import os -import time - -import cv2 -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torchvision.transforms as transform_lib -from PIL import Image -from tqdm import tqdm - -import lib.TestTransforms as transforms -from models.ColorVidNet import ColorVidNet -from models.FrameColor import frame_colorization -from models.NonlocalNet import VGG19_pytorch, WarpNet -from utils.util import (batch_lab2rgb_transpose_mc, folder2vid, mkdir_if_not, - save_frames, tensor_lab2rgb, uncenter_l) -from utils.util_distortion import CenterPad, Normalize, RGB2Lab, ToTensor - -os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" -os.environ["CUDA_VISIBLE_DEVICES"] = "0" -torch.cuda.set_device(0) - - -def colorize_video(opt, input_path, reference_file, output_path, nonlocal_net, colornet, vggnet): - # parameters for wls filter - wls_filter_on = True - lambda_value = 500 - sigma_color = 4 - - # processing folders - mkdir_if_not(output_path) - files = glob.glob(output_path + "*") - print("processing the folder:", input_path) - path, dirs, filenames = os.walk(input_path).__next__() - file_count = len(filenames) - filenames.sort(key=lambda f: int("".join(filter(str.isdigit, f) or -1))) - - # NOTE: resize frames to 216*384 - transform = transforms.Compose( - [CenterPad(opt.image_size), transform_lib.CenterCrop(opt.image_size), RGB2Lab(), ToTensor(), Normalize()] - ) - - # if frame propagation: use the first frame as reference - # otherwise, use the specified reference image - ref_name = input_path + filenames[0] if opt.frame_propagate else reference_file - print("reference name:", ref_name) - frame_ref = Image.open(ref_name) - - total_time = 0 - I_last_lab_predict = None - - IB_lab_large = transform(frame_ref).unsqueeze(0).cuda() - IB_lab = torch.nn.functional.interpolate(IB_lab_large, scale_factor=0.5, mode="bilinear") - IB_l = IB_lab[:, 0:1, :, :] - IB_ab = IB_lab[:, 1:3, :, :] - with torch.no_grad(): - I_reference_lab = IB_lab - I_reference_l = I_reference_lab[:, 0:1, :, :] - I_reference_ab = I_reference_lab[:, 1:3, :, :] - I_reference_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_reference_l), I_reference_ab), dim=1)) - features_B = vggnet(I_reference_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True) - - for index, frame_name in enumerate(tqdm(filenames)): - frame1 = Image.open(os.path.join(input_path, frame_name)) - IA_lab_large = transform(frame1).unsqueeze(0).cuda() - IA_lab = torch.nn.functional.interpolate(IA_lab_large, scale_factor=0.5, mode="bilinear") - - IA_l = IA_lab[:, 0:1, :, :] - IA_ab = IA_lab[:, 1:3, :, :] - - if I_last_lab_predict is None: - if opt.frame_propagate: - I_last_lab_predict = IB_lab - else: - I_last_lab_predict = torch.zeros_like(IA_lab).cuda() - - # start the frame colorization - with torch.no_grad(): - I_current_lab = IA_lab - I_current_ab_predict, I_current_nonlocal_lab_predict, features_current_gray = frame_colorization( - I_current_lab, - I_reference_lab, - I_last_lab_predict, - features_B, - vggnet, - nonlocal_net, - colornet, - feature_noise=0, - temperature=1e-10, - ) - I_last_lab_predict = torch.cat((IA_l, I_current_ab_predict), dim=1) - - # upsampling - curr_bs_l = IA_lab_large[:, 0:1, :, :] - curr_predict = ( - torch.nn.functional.interpolate(I_current_ab_predict.data.cpu(), scale_factor=2, mode="bilinear") * 1.25 - ) - - # filtering - if wls_filter_on: - guide_image = uncenter_l(curr_bs_l) * 255 / 100 - wls_filter = cv2.ximgproc.createFastGlobalSmootherFilter( - guide_image[0, 0, :, :].cpu().numpy().astype(np.uint8), lambda_value, sigma_color - ) - curr_predict_a = wls_filter.filter(curr_predict[0, 0, :, :].cpu().numpy()) - curr_predict_b = wls_filter.filter(curr_predict[0, 1, :, :].cpu().numpy()) - curr_predict_a = torch.from_numpy(curr_predict_a).unsqueeze(0).unsqueeze(0) - curr_predict_b = torch.from_numpy(curr_predict_b).unsqueeze(0).unsqueeze(0) - curr_predict_filter = torch.cat((curr_predict_a, curr_predict_b), dim=1) - IA_predict_rgb = batch_lab2rgb_transpose_mc(curr_bs_l[:32], curr_predict_filter[:32, ...]) - else: - IA_predict_rgb = batch_lab2rgb_transpose_mc(curr_bs_l[:32], curr_predict[:32, ...]) - - # save the frames - save_frames(IA_predict_rgb, output_path, index) - - # output video - video_name = "video.avi" - folder2vid(image_folder=output_path, output_dir=output_path, filename=video_name) - print() - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--frame_propagate", default=False, type=bool, help="propagation mode, , please check the paper" - ) - parser.add_argument("--image_size", type=int, default=[216 * 2, 384 * 2], help="the image size, eg. [216,384]") - parser.add_argument("--cuda", action="store_false") - parser.add_argument("--gpu_ids", type=str, default="0", help="separate by comma") - parser.add_argument("--clip_path", type=str, default="./sample_videos/clips/v32", help="path of input clips") - parser.add_argument("--ref_path", type=str, default="./sample_videos/ref/v32", help="path of refernce images") - parser.add_argument("--output_path", type=str, default="./sample_videos/output", help="path of output clips") - opt = parser.parse_args() - opt.gpu_ids = [int(x) for x in opt.gpu_ids.split(",")] - cudnn.benchmark = True - print("running on GPU", opt.gpu_ids) - - clip_name = opt.clip_path.split("/")[-1] - refs = os.listdir(opt.ref_path) - refs.sort() - - nonlocal_net = WarpNet(1) - colornet = ColorVidNet(7) - vggnet = VGG19_pytorch() - vggnet.load_state_dict(torch.load("data/vgg19_conv.pth")) - for param in vggnet.parameters(): - param.requires_grad = False - - nonlocal_test_path = os.path.join("checkpoints/", "video_moredata_l1/nonlocal_net_iter_76000.pth") - color_test_path = os.path.join("checkpoints/", "video_moredata_l1/colornet_iter_76000.pth") - print("succesfully load nonlocal model: ", nonlocal_test_path) - print("succesfully load color model: ", color_test_path) - nonlocal_net.load_state_dict(torch.load(nonlocal_test_path)) - colornet.load_state_dict(torch.load(color_test_path)) - - nonlocal_net.eval() - colornet.eval() - vggnet.eval() - nonlocal_net.cuda() - colornet.cuda() - vggnet.cuda() - - for ref_name in refs: - try: - colorize_video( - opt, - opt.clip_path, - os.path.join(opt.ref_path, ref_name), - os.path.join(opt.output_path, clip_name + "_" + ref_name.split(".")[0]), - nonlocal_net, - colornet, - vggnet, - ) - except Exception as error: - print("error when colorizing the video " + ref_name) - print(error) - - video_name = "video.avi" - clip_output_path = os.path.join(opt.output_path, clip_name) - mkdir_if_not(clip_output_path) - folder2vid(image_folder=opt.clip_path, output_dir=clip_output_path, filename=video_name) diff --git a/train.py b/train.py deleted file mode 100644 index 349fbc0..0000000 --- a/train.py +++ /dev/null @@ -1,738 +0,0 @@ -from __future__ import print_function - -import argparse -import math -import os -import queue -import time - -import cv2 -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.optim as optim -import torchvision.transforms as transform_lib -from torch.utils.data import ConcatDataset, DataLoader, WeightedRandomSampler -from torchvision.transforms import CenterCrop - -import lib.TrainTransforms as transforms -from lib.videoloader import VideosDataset -from lib.videoloader_imagenet import VideosDataset_ImageNet -from models.ColorVidNet import ColorVidNet -from models.ContextualLoss import ContextualLoss, ContextualLoss_forward -from models.FrameColor import frame_colorization -from models.GAN_models import Discriminator_x64 -from models.NonlocalNet import (NonlocalWeightedAverage, VGG19_pytorch, - WarpNet, WeightedAverage, - WeightedAverage_color) -from tensorboardX import SummaryWriter -from utils.util import (batch_lab2rgb_transpose_mc, feature_normalize, l1_loss, - mkdir_if_not, mse_loss, parse, tensor_lab2rgb, - uncenter_l, weighted_l1_loss, weighted_mse_loss) -from utils.util_distortion import (CenterPad_threshold, Normalize, RGB2Lab, - ToTensor) -from utils.util_tensorboard import TBImageRecorder, value_logger -from utils.warping import WarpingLayer - -cv2.setNumThreads(0) - -parser = argparse.ArgumentParser() -parser.add_argument("--data_root", default="/home/zhanbo/remote/video/video_pair3/", type=str) -parser.add_argument("--data_root_imagenet", default="/home/zhanbo/remote/imagenet_ref_pair/", type=str) -parser.add_argument("--gpu_ids", type=str, default="0,1,2,3", help="separate by comma") -parser.add_argument("--workers", type=int, default=8) -parser.add_argument("--batch_size", type=int, default=16) -parser.add_argument("--image_size", type=int, default=[216, 384]) -parser.add_argument("--ic", type=int, default=7) -parser.add_argument("--epoch", type=int, default=40) - -parser.add_argument("--resume_epoch", type=int, default=0) -parser.add_argument("--resume", type=bool, default=False) -parser.add_argument("--load_pretrained_model", type=bool, default=True) - -parser.add_argument("--lr", type=float, default=0.00001) -parser.add_argument("--beta1", type=float, default=0.5) -parser.add_argument("--lr_step", type=int, default=100) -parser.add_argument("--lr_gamma", type=float, default=0.1) - - -parser.add_argument("--checkpoint_dir", type=str, default="checkpoints/default") -parser.add_argument("--tb_log_step", type=int, default=50) -parser.add_argument("--print_step", type=int, default=2) - -parser.add_argument("--real_reference_probability", type=float, default=0.7) -parser.add_argument("--nonzero_placeholder_probability", type=float, default=0.0) -parser.add_argument("--with_bad", type=bool, default=True) -parser.add_argument("--with_mid", type=bool, default=True) - -parser.add_argument("--domain_invariant", type=bool, default=False) -parser.add_argument("--weigth_l1", type=float, default=2.0) -parser.add_argument("--weight_contextual", type=float, default="0.2") -parser.add_argument("--weight_perceptual", type=float, default="0.001") -parser.add_argument("--weight_smoothness", type=float, default="5.0") -parser.add_argument("--weight_gan", type=float, default="0.2") -parser.add_argument("--weight_nonlocal_smoothness", type=float, default="0.0") -parser.add_argument("--weight_nonlocal_consistent", type=float, default="0.0") -parser.add_argument("--weight_consistent", type=float, default="0.02") -parser.add_argument("--luminance_noise", type=float, default="2.0") -parser.add_argument("--permute_data", type=bool, default=True) -parser.add_argument("--contextual_loss_direction", type=str, default="forward", help="forward or backward matching") - - -def image_logger_fn( - I_last_lab, - I_current_lab, - I_reference_lab, - I_last_lab_predict, - I_current_lab_predict, - I_last_nonlocal_lab, - I_current_nonlocal_lab, -): - I_last_image = batch_lab2rgb_transpose_mc(I_last_lab[0:32, 0:1, :, :], I_last_lab[0:32, 1:3, :, :]) - I_current_image = batch_lab2rgb_transpose_mc(I_current_lab[0:32, 0:1, :, :], I_current_lab[0:32, 1:3, :, :]) - I_reference_image = batch_lab2rgb_transpose_mc(I_reference_lab[0:32, 0:1, :, :], I_reference_lab[0:32, 1:3, :, :]) - I_last_image_predict = batch_lab2rgb_transpose_mc( - I_last_lab_predict[0:32, 0:1, :, :], I_last_lab_predict[0:32, 1:3, :, :] - ) - I_current_image_predict = batch_lab2rgb_transpose_mc( - I_current_lab_predict[0:32, 0:1, :, :], I_current_lab_predict[0:32, 1:3, :, :] - ) - I_last_nonlocal_image = batch_lab2rgb_transpose_mc( - I_last_nonlocal_lab[0:32, 0:1, :, :], I_last_nonlocal_lab[0:32, 1:3, :, :] - ) - I_current_nonlocal_image = batch_lab2rgb_transpose_mc( - I_current_nonlocal_lab[0:32, 0:1, :, :], I_current_nonlocal_lab[0:32, 1:3, :, :] - ) - - img_info = {} - img_info["1_I_last"] = I_last_image - img_info["2_I_current"] = I_current_image - img_info["3_I_reference"] = I_reference_image - img_info["4_I_last_predict"] = I_last_image_predict - img_info["5_I_curren_predict"] = I_current_image_predict - img_info["6_I_last_nonlocal"] = I_last_nonlocal_image - img_info["7_I_current_nonlocal"] = I_current_nonlocal_image - - return img_info - - -def training_logger(): - try: - if total_iter % opt.print_step == 0: - print("processing time:", elapsed) - print( - "Epoch %d, Step[%d/%d], lr: %f, total_loss: %.2f" - % ( - epoch, - ((iter + 1) % iter_num_per_epoch), - iter_num_per_epoch, - step_optim_scheduler_g.get_last_lr()[0], - total_loss.item(), - ) - ) - - value_logger( - tb_writer, - total_iter, - loss_info={ - "l1_loss": l1_loss.item(), - "feat_loss": feat_loss.item(), - "contextual_loss_total": contextual_loss_total.item(), - "smoothness_loss": smoothness_loss.item(), - "nonlocal_smoothness_loss": nonlocal_smoothness_loss.item(), - "nonlocal_consistent_loss": nonlocal_consistent_loss.item(), - "consistent_loss": consistent_loss.item(), - "generator_loss": generator_loss.item(), - "discriminator_loss": discriminator_loss.item(), - "total_loss": total_loss.item(), - }, - ) - - if total_iter % opt.tb_log_step == 0: - I_last_nonlocal_lab = torch.cat( - (I_last_lab[:, 0:1, :, :], I_last_nonlocal_lab_predict[:, 1:3, :, :]), dim=1 - ) - I_current_nonlocal_lab = torch.cat( - (I_current_lab[:, 0:1, :, :], I_current_nonlocal_lab_predict[:, 1:3, :, :]), dim=1 - ) - I_last_lab_predict = torch.cat((I_last_l, I_last_ab_predict), dim=1) - data_queue.put( - ( - ( - I_last_lab.cpu(), - I_current_lab.cpu(), - I_reference_lab.cpu(), - I_last_lab_predict.cpu(), - I_current_lab_predict.cpu(), - I_last_nonlocal_lab.cpu(), - I_current_nonlocal_lab, - ), - total_iter, - ) - ) - - if total_iter % 2000 == 0: - if len(opt.gpu_ids) > 1: - torch.save( - nonlocal_net.module.state_dict(), - os.path.join(opt.checkpoint_dir, "nonlocal_net_iter_%d.pth") % total_iter, - ) - torch.save( - colornet.module.state_dict(), - os.path.join(opt.checkpoint_dir, "colornet_iter_%d.pth") % total_iter, - ) - torch.save( - discriminator.module.state_dict(), - os.path.join(opt.checkpoint_dir, "discriminator_iter_%d.pth") % total_iter, - ) - else: - torch.save( - nonlocal_net.state_dict(), - os.path.join(opt.checkpoint_dir, "nonlocal_net_iter_%d.pth") % total_iter, - ) - torch.save(colornet.state_dict(), os.path.join(opt.checkpoint_dir, "colornet_iter_%d.pth") % total_iter) - torch.save( - discriminator.state_dict(), - os.path.join(opt.checkpoint_dir, "discriminator_iter_%d.pth") % total_iter, - ) - - # save the state for resume - if total_iter % 2000 == 0: - print("saving the checkpoint") - if len(opt.gpu_ids) > 1: - state = { - "total_iter": total_iter + 1, - "epoch": epoch, - "colornet_state": colornet.module.state_dict(), - "nonlocal_net_state": nonlocal_net.module.state_dict(), - "discriminator_state": discriminator.module.state_dict(), - "optimizer_g": optimizer_g.state_dict(), - "optimizer_d": optimizer_d.state_dict(), - "optimizer_schedule_g": step_optim_scheduler_g.state_dict(), - "optimizer_schedule_d": step_optim_scheduler_g.state_dict(), - } - else: - state = { - "total_iter": total_iter + 1, - "epoch": epoch, - "colornet_state": colornet.state_dict(), - "nonlocal_net_state": nonlocal_net.state_dict(), - "discriminator_state": discriminator.state_dict(), - "optimizer_g": optimizer_g.state_dict(), - "optimizer_d": optimizer_d.state_dict(), - "optimizer_schedule_g": step_optim_scheduler_g.state_dict(), - "optimizer_schedule_d": step_optim_scheduler_d.state_dict(), - } - torch.save(state, os.path.join(opt.checkpoint_dir, "learning_checkpoint.pth")) - except Exception as e: - print("Exception during output") - print(e) - - -def gpu_setup(): - os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" - cudnn.benchmark = True - opt.gpu_ids = [int(x) for x in opt.gpu_ids.split(",")] - torch.cuda.set_device(opt.gpu_ids[0]) - device = torch.device("cuda") - print("running on GPU", opt.gpu_ids) - return device - - -def load_data(): - print("initializing dataloader") - transforms_video = [ - CenterCrop(opt.image_size), - RGB2Lab(), - ToTensor(), - Normalize(), - ] - transforms_imagenet = [CenterPad_threshold(opt.image_size), RGB2Lab(), ToTensor(), Normalize()] - extra_reference_transform = [ - transform_lib.RandomHorizontalFlip(0.5), - transform_lib.RandomResizedCrop(480, (0.98, 1.0), ratio=(0.8, 1.2)), - ] - train_dataset_video = VideosDataset( - data_root=opt.data_root, - epoch=opt.epoch, - image_size=opt.image_size, - image_transform=transforms.Compose(transforms_video), - real_reference_probability=opt.real_reference_probability, - nonzero_placeholder_probability=opt.nonzero_placeholder_probability, - ) - train_dataset_imagenet = VideosDataset_ImageNet( - data_root=opt.data_root_imagenet, - image_size=opt.image_size, - epoch=opt.epoch, - with_bad=opt.with_bad, - with_mid=opt.with_mid, - transforms_imagenet=transforms_imagenet, - distortion_level=4, - brightnessjitter=5, - nonzero_placeholder_probability=opt.nonzero_placeholder_probability, - extra_reference_transform=extra_reference_transform, - real_reference_probability=opt.real_reference_probability, - ) - - video_training_length = len(train_dataset_video) - imagenet_training_length = len(train_dataset_imagenet) - dataset_training_length = train_dataset_video.real_len + train_dataset_imagenet.real_len - dataset_combined = ConcatDataset([train_dataset_video, train_dataset_imagenet]) - sampler = WeightedRandomSampler( - [1] * video_training_length + [1] * imagenet_training_length, dataset_training_length * opt.epoch - ) - data_loader = DataLoader( - dataset_combined, - batch_size=opt.batch_size, - shuffle=False, - num_workers=opt.workers, - pin_memory=True, - drop_last=True, - sampler=sampler, - ) - return dataset_training_length, train_dataset_video, train_dataset_imagenet, data_loader - - -def define_loss(): - print("defining loss") - ab_criterion = nn.SmoothL1Loss().to(device) - nonlocal_criterion = nn.SmoothL1Loss().to(device) - feat_l2_criterion = nn.MSELoss().to(device) - feat_l1_criterion = nn.SmoothL1Loss().to(device) - contextual_loss = ContextualLoss().to(device) - contextual_forward_loss = ContextualLoss_forward().to(device) - BCE_stable = nn.BCEWithLogitsLoss().to(device) - return contextual_loss, contextual_forward_loss - - -def define_optimizer(): - print("defining optimizer") - optimizer_g = optim.Adam( - [{"params": nonlocal_net.parameters(), "lr": 1e-5}, {"params": colornet.parameters(), "lr": 2e-4}], - betas=(0.5, 0.999), - eps=1e-5, - amsgrad=True, - ) - optimizer_d = optim.Adam( - filter(lambda p: p.requires_grad, discriminator.parameters()), lr=2 * 1e-4, betas=(0.5, 0.999) - ) - return optimizer_g, optimizer_d - - -def resume_model(): - print("resuming the learning") - checkpoint = torch.load(os.path.join(opt.checkpoint_dir, "learning_checkpoint.pth")) - total_iter = checkpoint["total_iter"] - epoch = checkpoint["epoch"] - colornet.load_state_dict(checkpoint["colornet_state"]) - nonlocal_net.load_state_dict(checkpoint["nonlocal_net_state"]) - discriminator.load_state_dict(checkpoint["discriminator_state"]) - optimizer_g.load_state_dict(checkpoint["optimizer_g"]) - optimizer_d.load_state_dict(checkpoint["optimizer_d"]) - step_optim_scheduler_g.load_state_dict(checkpoint["optimizer_schedule_g"]) - step_optim_scheduler_d.load_state_dict(checkpoint["optimizer_schedule_d"]) - - -def to_device( - colornet, - nonlocal_net, - discriminator, - vggnet, - contextual_loss, - contextual_forward_loss, - weighted_layer_color, - nonlocal_weighted_layer, - warping_layer, - instancenorm, -): - print("moving models to device") - colornet = torch.nn.DataParallel(colornet.to(device), device_ids=opt.gpu_ids) - nonlocal_net = torch.nn.DataParallel(nonlocal_net.to(device), device_ids=opt.gpu_ids) - discriminator = torch.nn.DataParallel(discriminator.to(device), device_ids=opt.gpu_ids) - vggnet = torch.nn.DataParallel(vggnet.to(device), device_ids=opt.gpu_ids) - contextual_loss = torch.nn.DataParallel(contextual_loss.to(device), device_ids=opt.gpu_ids) - contextual_forward_loss = torch.nn.DataParallel(contextual_forward_loss.to(device), device_ids=opt.gpu_ids) - weighted_layer_color = torch.nn.DataParallel(weighted_layer_color.to(device), device_ids=opt.gpu_ids) - nonlocal_weighted_layer = torch.nn.DataParallel(nonlocal_weighted_layer.to(device), device_ids=opt.gpu_ids) - warping_layer = torch.nn.DataParallel(warping_layer.to(device), device_ids=opt.gpu_ids) - instancenorm = torch.nn.DataParallel(instancenorm.to(device), device_ids=opt.gpu_ids) - - return ( - vggnet, - nonlocal_net, - colornet, - discriminator, - instancenorm, - contextual_loss, - contextual_forward_loss, - weighted_layer_color, - nonlocal_weighted_layer, - warping_layer, - ) - - -def loss_init(): - print("initializing losses") - zero_loss = torch.Tensor([0]).to(device) - ( - feat_loss, - contextual_loss_total, - smoothness_loss, - nonlocal_smoothness_loss, - consistent_loss, - nonlocal_consistent_loss, - generator_loss, - discriminator_loss, - ) = (zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss) - - return ( - feat_loss, - contextual_loss_total, - smoothness_loss, - nonlocal_smoothness_loss, - consistent_loss, - nonlocal_consistent_loss, - generator_loss, - discriminator_loss, - ) - - -def video_colorization(): - # colorization for the last frame - I_last_ab_predict, I_last_nonlocal_lab_predict, features_last_gray = frame_colorization( - I_last_lab, - I_reference_lab, - placeholder_lab, - features_B, - vggnet, - nonlocal_net, - colornet, - feature_noise=0, - luminance_noise=opt.luminance_noise, - ) - I_last_lab_predict = torch.cat((I_last_l, I_last_ab_predict), dim=1) - - # colorization for the current frame - I_current_ab_predict, I_current_nonlocal_lab_predict, features_current_gray = frame_colorization( - I_current_lab, - I_reference_lab, - I_last_lab_predict, - features_B, - vggnet, - nonlocal_net, - colornet, - feature_noise=0, - luminance_noise=opt.luminance_noise, - ) - I_current_lab_predict = torch.cat((I_last_l, I_current_ab_predict), dim=1) - return I_current_ab_predict, I_last_ab_predict, I_current_nonlocal_lab_predict, I_last_nonlocal_lab_predict - - -if __name__ == "__main__": - torch.multiprocessing.set_start_method("spawn", force=True) - opt = parse(parser) - opt.data_root = opt.data_root.split(",")[0] - opt.data_root_imagenet = opt.data_root_imagenet.split(",")[0] - mkdir_if_not(opt.checkpoint_dir) - mkdir_if_not("./runs/") - - device = gpu_setup() - dataset_training_length, train_dataset_video, train_dataset_imagenet, data_loader = load_data() - tb_writer = SummaryWriter() - data_queue = queue.Queue() - tb_image_reorder = TBImageRecorder(tb_writer, image_logger_fn, data_queue) - tb_image_reorder.start() - - # define network - nonlocal_net = WarpNet(opt.batch_size) - colornet = ColorVidNet(opt.ic) - discriminator = Discriminator_x64(ndf=64) - colornet.to(device) - nonlocal_net.to(device) - discriminator.to(device) - - weighted_layer = WeightedAverage() - weighted_layer_color = WeightedAverage_color() - nonlocal_weighted_layer = NonlocalWeightedAverage() - instancenorm = nn.InstanceNorm2d(512, affine=False) - warping_layer = WarpingLayer("gpu") - - vggnet = VGG19_pytorch() - vggnet.load_state_dict(torch.load("data/vgg19_conv.pth")) - vggnet.eval() - for param in vggnet.parameters(): - param.requires_grad = False - - # load pre-trained model - if opt.load_pretrained_model: - nonlocal_pretain_path = os.path.join("checkpoints/video_moredata_l1/", "nonlocal_net_iter_76000.pth") - nonlocal_net.load_state_dict(torch.load(nonlocal_pretain_path)) - color_test_path = "checkpoints/video_moredata_l1/" + "colornet_iter_76000.pth" - colornet.load_state_dict(torch.load(color_test_path)) - - # define loss function - contextual_loss, contextual_forward_loss = define_loss() - - # define optimizer - optimizer_g, optimizer_d = define_optimizer() - step_optim_scheduler_g = optim.lr_scheduler.StepLR(optimizer_g, step_size=opt.lr_step, gamma=opt.lr_gamma) - step_optim_scheduler_d = optim.lr_scheduler.StepLR(optimizer_d, step_size=opt.lr_step, gamma=opt.lr_gamma) - - # define others - downsampling_by2 = nn.AvgPool2d(kernel_size=2).to(device) - downsampling_by4 = nn.AvgPool2d(kernel_size=4).to(device) - - # dataset info - iter_num_per_epoch = dataset_training_length // opt.batch_size - total_iter = opt.resume_epoch * iter_num_per_epoch - print( - "train_dataset info, real_len: %d, epoch_len: %d, iter_num_per_epoch: %d" - % (dataset_training_length, len(train_dataset_video) + len(train_dataset_imagenet), iter_num_per_epoch) - ) - - if opt.resume: - resume_model() - - # move to GPU processing - device = opt.gpu_ids[0] - ( - vggnet, - nonlocal_net, - colornet, - discriminator, - instancenorm, - contextual_loss, - contextual_forward_loss, - weighted_layer_color, - nonlocal_weighted_layer, - warping_layer, - ) = to_device( - colornet, - nonlocal_net, - discriminator, - vggnet, - contextual_loss, - contextual_forward_loss, - weighted_layer_color, - nonlocal_weighted_layer, - warping_layer, - instancenorm, - ) - - ( - feat_loss, - contextual_loss_total, - smoothness_loss, - nonlocal_smoothness_loss, - consistent_loss, - nonlocal_consistent_loss, - generator_loss, - discriminator_loss, - ) = loss_init() - - # %% Training - print("start training") - for iter, data in enumerate(data_loader): - start_time = time.time() - total_iter += 1 - epoch = math.ceil(total_iter / iter_num_per_epoch) - - ###### LOADING DATA SAMPLE ###### - ( - I_last_lab, - I_current_lab, - I_reference_lab, - flow_forward, - flow_backward, - mask, - placeholder_lab, - self_ref_flag, - ) = data - I_last_lab = I_last_lab.to(device) - I_current_lab = I_current_lab.to(device) - I_reference_lab = I_reference_lab.to(device) - flow_forward = flow_forward.to(device) - flow_backward = flow_backward.to(device) - mask = mask.to(device) - placeholder_lab = placeholder_lab.to(device) - self_ref_flag = self_ref_flag.to(device) - - I_last_l = I_last_lab[:, 0:1, :, :] - I_last_ab = I_last_lab[:, 1:3, :, :] - I_current_l = I_current_lab[:, 0:1, :, :] - I_current_ab = I_current_lab[:, 1:3, :, :] - I_reference_l = I_reference_lab[:, 0:1, :, :] - I_reference_ab = I_reference_lab[:, 1:3, :, :] - I_reference_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_reference_l), I_reference_ab), dim=1)) - features_B = vggnet(I_reference_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True) - - ###### COLORIZATION ###### - ( - I_current_ab_predict, - I_last_ab_predict, - I_current_nonlocal_lab_predict, - I_last_nonlocal_lab_predict, - ) = video_colorization() - - ###### UPDATE DISCRIMINATOR ###### - if opt.weight_gan > 0: - optimizer_g.zero_grad() - optimizer_d.zero_grad() - - fake_data_lab = torch.cat( - (uncenter_l(I_current_l), I_current_ab_predict, uncenter_l(I_last_l), I_last_ab_predict), dim=1 - ) - real_data_lab = torch.cat((uncenter_l(I_current_l), I_current_ab, uncenter_l(I_last_l), I_last_ab), dim=1) - - if opt.permute_data: - batch_index = torch.arange(-1, opt.batch_size - 1, dtype=torch.long) - real_data_lab = real_data_lab[batch_index, ...] - - y_pred_fake, feature_pred_fake = discriminator(fake_data_lab.detach()) - y_pred_real, feature_pred_real = discriminator(real_data_lab.detach()) - - y = torch.ones_like(y_pred_real) - y2 = torch.zeros_like(y_pred_real) - discriminator_loss = ( - torch.mean((y_pred_real - torch.mean(y_pred_fake) - y) ** 2) - + torch.mean((y_pred_fake - torch.mean(y_pred_real) + y) ** 2) - ) / 2 - discriminator_loss.backward() - optimizer_d.step() - - ###### UPDATE GENERATOR ###### - optimizer_g.zero_grad() - optimizer_d.zero_grad() - - # extract vgg features for both output and original image - I_predict_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_current_l), I_current_ab_predict), dim=1)) - predict_relu1_1, predict_relu2_1, predict_relu3_1, predict_relu4_1, predict_relu5_1 = vggnet( - I_predict_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True - ) - - I_current_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_current_l), I_current_ab), dim=1)) - A_relu1_1, A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1 = vggnet( - I_current_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True - ) - B_relu1_1, B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1 = features_B - - ###### LOSS COMPUTE ###### - # l1 loss - if opt.weigth_l1 > 0: - sample_weights = (self_ref_flag[:, 1:3, :, :]) / (sum(self_ref_flag[:, 0, 0, 0]) + 1e-5) - l1_loss = weighted_l1_loss(I_current_ab_predict, I_current_ab, sample_weights) * opt.weigth_l1 - - # generator loss - if opt.weight_gan > 0: - y_pred_fake, feature_pred_fake = discriminator(fake_data_lab) - y_pred_real, feature_pred_real = discriminator(real_data_lab) - generator_loss = ( - ( - torch.mean((y_pred_real - torch.mean(y_pred_fake) + y) ** 2) - + torch.mean((y_pred_fake - torch.mean(y_pred_real) - y) ** 2) - ) - / 2 - * opt.weight_gan - ) - - # feature loss - if opt.domain_invariant: - feat_loss = ( - mse_loss(instancenorm(predict_relu5_1), instancenorm(A_relu5_1.detach())) - * opt.weight_perceptual - * 1e5 - * 0.2 - ) - else: - feat_loss = mse_loss(predict_relu5_1, A_relu5_1.detach()) * opt.weight_perceptual - - # contextual loss - if opt.contextual_loss_direction == "backward": - contextual_style5_1 = torch.mean(contextual_loss(predict_relu5_1, B_relu5_1.detach())) * 8 - contextual_style4_1 = torch.mean(contextual_loss(predict_relu4_1, B_relu4_1.detach())) * 4 - contextual_style3_1 = ( - torch.mean(contextual_loss(downsampling_by2(predict_relu3_1), downsampling_by2(B_relu3_1.detach()))) * 2 - ) - else: - contextual_style5_1 = torch.mean(contextual_forward_loss(predict_relu5_1, B_relu5_1.detach())) * 8 - contextual_style4_1 = torch.mean(contextual_forward_loss(predict_relu4_1, B_relu4_1.detach())) * 4 - contextual_style3_1 = ( - torch.mean( - contextual_forward_loss(downsampling_by2(predict_relu3_1), downsampling_by2(B_relu3_1.detach())) - ) - * 2 - ) - if opt.weight_contextual > 0: - contextual_loss_total = ( - contextual_style5_1 + contextual_style4_1 + contextual_style3_1 - ) * opt.weight_contextual - - # smoothness loss - if opt.weight_smoothness > 0: - scale_factor = 1 - I_current_lab_predict = torch.cat((I_current_l, I_current_ab_predict), dim=1) - IA_ab_weighed = weighted_layer_color( - I_current_lab, I_current_lab_predict, patch_size=3, alpha=10, scale_factor=scale_factor - ) - smoothness_loss = ( - mse_loss(nn.functional.interpolate(I_current_ab_predict, scale_factor=scale_factor), IA_ab_weighed) - * opt.weight_smoothness - ) - - if opt.weight_nonlocal_smoothness > 0: - scale_factor = 0.25 - alpha_nonlocal_smoothness = 0.5 - nonlocal_smooth_feature = feature_normalize(A_relu2_1) - I_current_lab_predict = torch.cat((I_current_l, I_current_ab_predict), dim=1) - I_current_ab_weighted_nonlocal = nonlocal_weighted_layer( - I_current_lab_predict, - nonlocal_smooth_feature.detach(), - patch_size=3, - alpha=alpha_nonlocal_smoothness, - scale_factor=scale_factor, - ) - nonlocal_smoothness_loss = ( - mse_loss( - nn.functional.interpolate(I_current_ab_predict, scale_factor=scale_factor), - I_current_ab_weighted_nonlocal, - ) - * opt.weight_nonlocal_smoothness - ) - - if opt.weight_consistent: - I_current_lab_predict_warp = warping_layer(I_current_lab_predict, flow_forward) - I_current_ab_predict_warp = I_current_lab_predict_warp[:, 1:3, :, :] - consistent_loss = ( - weighted_mse_loss(I_current_ab_predict_warp, I_last_ab_predict, mask) * opt.weight_consistent - ) - - if opt.weight_nonlocal_consistent: - I_current_nonlocal_lab_predict_warp = warping_layer(I_current_nonlocal_lab_predict, flow_forward) - nonlocal_consistent_loss = ( - weighted_mse_loss( - I_current_nonlocal_lab_predict_warp[:, 1:3, :, :], I_last_nonlocal_lab_predict[:, 1:3, :, :], mask - ) - * opt.weight_nonlocal_consistent - ) - - # total loss - total_loss = ( - l1_loss - + feat_loss - + contextual_loss_total - + smoothness_loss - + nonlocal_smoothness_loss - + consistent_loss - + nonlocal_consistent_loss - + generator_loss - ) - total_loss.backward() - optimizer_g.step() - end_time = time.time() - elapsed = end_time - start_time - training_logger() - - step_optim_scheduler_g.step() - step_optim_scheduler_d.step() - - data_queue.put((None, None))