- Add LAB color transfer as new default method for superior perceptual color matching
- Add HSV hue-conditional saturation histogram matching for targeted oversaturation correction
- Add wavelet_adaptive hybrid method combining wavelet base with selective HSV correction
- Change default color correction from 'wavelet' to 'lab' for better color accuracy
- Implement full CIELAB color space conversion with D65 illuminant and histogram matching
- Add comprehensive documentation and method descriptions to color_fix.py
Major refactor introducing:
- GlobalModelCache for cross-node model sharing with dynamic config updates
* Models cached by node ID, reused across different upscaler instances
* Config changes (torch.compile, BlockSwap, tiling) handled dynamically
* Runner templates cached when both models present
- Complete alpha channel processing rewrite
* New edge-guided alpha upscaling in src/core/alpha_upscaling.py
* Remove broken VAE RGBA adapter approach
* Alpha properly extracted, upscaled, and merged with RGB
- BlockSwap enhancements
* Support I/O-only swapping without transformer block offloading
* Separate memory reporting for I/O components vs transformer blocks
* Clearer logging of GPU vs CPU placement
* New is_blockswap_enabled() utility
- Memory management improvements
* Add release_tensor_collection() for batch tensor cleanup
* Proper tensor memory release throughout pipeline
* Better cleanup in postprocess phase
- Code quality and refactoring
* Extensive docstrings with type hints across all modules
* Generic _update_model_config() reduces code duplication
* Better function signatures and parameter documentation
* Improved separation of concerns in model_manager.py
* Clearer debug logging categories and messages
Core Changes:
- Split monolithic comfyui_node.py into modular per-node files:
* video_upscaler.py - main upscaler node
* dit_model_loader.py - DiT model configuration
* vae_model_loader.py - VAE model configuration
* torch_compile_settings.py - torch.compile settings
* __init__.py - centralized node registry
BlockSwap Enhancements:
- Fixed swap_io_components to work independently of blocks_to_swap
- Both features can now be enabled separately or together
- Updated validation logic throughout (model_manager, blockswap, generation)
- Improved config description to reflect combined capabilities
Logging Improvements:
- Standardized indentation and improved debug messages clarity
- Added footer with support links (YouTube, GitHub) for community engagement
- Remove separate mask input/output - node now auto-detects RGB vs RGBA from pixels input
- Automatically invert alpha channel on RGBA input to match ComfyUI's mask convention
- Add full input→output resolution logging using actual transform pipeline
- Reuse video transform pipeline between dimension calculation and encoding
- Improve log readability: consistent format, concise color correction messages
- Update tooltips to clarify RGB/RGBA support
Major features:
- RGBA/Alpha channel support throughout upscaling pipeline
* New VAE RGBA adapter for dynamic 3→4 channel conversion
* Detect RGBA input and adapt VAE after weight loading
* Preserve alpha during encoding/decoding phases
* Split alpha for RGB-only color correction, then recombine
* ComfyUI node now accepts mask input and outputs mask
* Frame count validation between pixels and mask
Configuration management:
- Smart in-place config updates without model reload
- Compare BlockSwap and torch.compile configs for changes
- Cleanup old materialized models before creating new meta structures
- Preserve config attributes (_dit_compile_args, _vae_compile_args, _dit_block_swap_config) for change detection
- Better config change logging with descriptive before/after states
- Handle model reloads separately from config-only updates
Logging improvements:
- Consistent 2-space indentation for hierarchical log messages
- More detailed configuration change descriptions
- Better visibility of RGBA processing steps
- Improved batch optimization tips formatting
Memory and cleanup:
- Handle meta device properly in cleanup functions
- Consistent attribute cleanup across dit/vae
- Clear compile/blockswap configs only when models not cached
- Improved model movement checks for meta device
VAE torch.compile enhancements:
- Disable compilation for InflatedCausalConv3d (dynamic shapes)
- Apply to both encoder and decoder to prevent recompilation issues
Core Features:
- Add torch.compile support for DiT (20-40% speedup) and VAE (15-25% speedup)
- Ensure BlockSwap compatibility by applying BlockSwap before torch.compile
- Add SeedVR2TorchCompileSettings node for ComfyUI configuration
- Add CLI arguments: --compile_dit, --compile_vae, --compile_backend, --compile_mode, --compile_fullgraph, --compile_dynamic, --compile_dynamo_cache_size_limit, --compile_dynamo_recompile_limit
Optimizations:
- Optimize na.py and other backend files for torch.compile compatibility (replace .tolist() with tensor operations)
- Remove .item() calls from attention modules (max_seqlen_q/k now accept tensors)
- Add @torch._dynamo.disable decorators to timing/debug methods to prevent compilation warnings
- Replace torch.repeat with torch.repeat_interleave in modulation.py for better performance
Bug Fixes:
- Fix timing stack cleanup in memory_manager.py (end timer before early returns)
- Fix BlockSwap timing reporting (_get_swap_start_time and _log_swap_timing now properly excluded from compilation)
Code Quality:
- Add comprehensive type hints and docstrings across all modified modules
- Standardize BlockSwap timer names (blockswap_block_*, blockswap_io_*)
BREAKING CHANGE: Replaced BlockSwap and ExtraArgs nodes with DiT/VAE loader nodes.
Enables multi-GPU placement, independent caching, and granular memory control.
New Features:
- Add SeedVR2LoadDiTModel and SeedVR2LoadVAEModel loader nodes
- Support different devices for DiT and VAE (multi-GPU load balancing)
- Independent model caching (cache_model_dit/cache_model_vae)
- Separate encode/decode VAE tiling with independent tile configurations
- User-defined DiT and VAE model selection from registry or disk
Breaking Changes:
- Deprecated SeedVR2BlockSwap node (integrated into DiT loader)
- Deprecated SeedVR2ExtraArgs node (split into DiT/VAE loaders & Upscaler node)
- BlockSwap configuration now part of DiT loader node
Improvements:
- Follow ComfyUI conventions (pixels input, IMAGE output type)
- Update download_weight() to support user-defined DiT/VAE models
- Improve tooltips and default parameters across all nodes
- Set show_tensors=False in debug logging for performance
- Dual-device context tracking throughout pipeline
- CLI updated to match new dual-model API
- Add suppress_tensor_warnings() to constants.py to centralize tensor/numpy warning handling
- Fix type hints and List/Tuple imports
- Improve GGUFTensor.__torch_function__ to prevent recursion, add debug null checks
- Fix GGUFTensor.to() to properly preserve tensor_shape attribute
- Enhance GGUF error handling with proper exceptions instead of just warnings
- Simplify gguf_ops dequantize_weight by removing redundant fallback paths
- Add memory logging and proper CUDA sync cleanup in GGUF loading
- Translate French comments to English in euler.py
- Remove unnecessary GGUF special handling in memory_manager
- Fix _propagate_debug_to_modules to handle None debug instance
- Implement GGUF model loading with Q3_K_M through Q8_K_M quantization support
- Add GGUFTensor wrapper to preserve quantization and enable on-demand dequantization
- Maintain tensors in quantized format to reduce VRAM usage
- Add GGUF dequantization operations for inference
- Update model registry to include GGUF variants for 3B/7B models
- Fix wavelet blur radius limit to prevent OOM at high resolutions (max 1/8 of image dimension)
- Add safety clamp [-1,1] for SDR color range to prevent numerical errors
- This is a WIP commit as some additional cleaning/testing is needed
- Add type hints throughout for better code maintainability
- Implement model discovery across multiple paths via ComfyUI's folder_paths API
- Add centralized model file scanning in get_all_model_files() to avoid duplication
- Update model loading to search all registered paths before downloading
- Prevent redundant downloads when files exist in alternate paths
- Users can now specify custom model directories using 'seedvr2' key in extra_model_paths.yaml
Addresses request from https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/issues/128#issuecomment-3312407641
- Clean up all __init__.py files to remove commented code and unnecessary imports
- Root __init__.py now only exports NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS
- Remove unused dependencies: pytorch-extension (replaced by custom implementations) and peft
- Fix diffusers version syntax (remove 'v' prefix)
- Add pyproject.toml for ComfyUI registry compliance
- Renamed model directories for clarity: dit -> dit_7b, dit_v2 -> dit_3b
- Converted all absolute imports to relative imports throughout codebase
- Removed sys.path.append() manipulations that caused namespace conflicts
- Updated YAML configs to reference renamed model directories
- Simplified model variant detection logic using new directory names
- Standardized function calls with named arguments for better clarity
- Fixed generation context initialization and interrupt handling
This resolves import conflicts with other ComfyUI nodes (e.g., Basic data handling)
that use sys.path manipulation, making the module properly isolated and compatible.
Fixes#29, #114, #136
Major architectural change to minimize model swapping overhead by processing
all batches in three distinct phases instead of sequential per-batch processing:
- Phase 1: Encode all batches with VAE
- Phase 2: Upscale all latents with DiT
- Phase 3: Decode all latents with VAE
Core changes:
- Split monolithic generation_loop into modular functions:
- prepare_generation_context(): Shared state management
- setup_device_environment(): Device configuration
- prepare_runner(): Model loading with cache support
- encode_all_batches(): Batch VAE encoding
- upscale_all_batches(): Batch DiT upscaling
- decode_all_batches(): Batch VAE decoding
- Removed generation_step function (logic integrated into upscale phase)
- Added lazy precision initialization to avoid redundant setup
Performance improvements:
- Pre-allocated lists for memory efficiency
- Better cleanup of intermediate storage between phases
- Added unique timer names to clear_memory() to avoid naming conflicts
- Improved model state management with change detection and caching
UI/UX enhancements:
- Switched to ComfyUI's native ProgressBar with weighted phase progress
- Changed from per-batch FPS to overall average FPS (always visible)
- Improved log clarity with clear phase separators
- Added ASCII art logo to clearly identify SeedVR2 process start
- Better progress tracking with weighted percentages across three phases
Code cleanup:
- Removed deprecated timer_context from Debug class
- Removed unused time imports across multiple files
- Fixed LOCAL_RANK environment variable to handle string conversion properly
- Improved error handling with try/except/finally blocks in all phases
- Change log_memory_state() to default show_tensors=False (opt-in tensor counting), skipping expensive tensor counting during intermediate operations
- Enable detailed tensor analysis only at final cleanup
- Add early return pattern in log() method to avoid string operations when disabled
- Make reset_vram_peak() logging conditional on debug.enabled state
- Remove unused vram_info call in model_manager.py
- Use meta device initialization to avoid unnecessary memory allocation during model creation
- Reduces DiT and VAE initialization time by ~90% when loading to CPU
- Explicitly delete state dicts after loading to free memory immediately
- Refactor configure_model_inference() into focused helper functions for DRY code
- Improve debug logging with timestamps and cleaner section separators for readability
- Merge generation setup log steps: combine "Model Configuration" and "Input Preparation" into unified "Generation Setup" step
- Simplify dtype handling: remove redundant vae_dtype variable, VAE now uses configured dtype consistently
- Remove platform-specific VAE dtype logic (MPS special case)
- Unify VAE encode/decode: remove autocast wrapper and target_dtype parameter, both now use configured dtype
- Add consistent docstrings for vae_encode and vae_decode methods
- Improve precision logging: show both DiT and VAE dtypes, rename model_dtype to dit_dtype
- Optimize text embeddings movement: only move when preserve_vram is active (BlockSwap handles model layers separately)
- Load text embeddings directly to CPU when preserve_vram is enabled
- Update logging consistency: uppercase device names, unified category names, clearer messages
- Clarify FP8 RoPE conversion log message: specify "from FP8 to BFloat16"
- Merge configure_dit_model_inference and configure_vae_model_inference into single configure_model_inference function
- Simplify load_quantized_state_dict by removing unnecessary FP8 conversion logic
- Remove redundant gradient checkpointing call (training-only feature)
- Remove redundant model_weight variables
- Standardize float formatting to .2f across all debug logs for consistency
- Eliminate redundant device movement operations already handled during model creation
- Remove "Force move weights to device" in FP8CompatibleDiT forward() to avoid clash with blockswap - This was used for preserve_vram but will refactor preserve_vram in a separate commit
- Remove duplicate dtype detection in generation_step (now passed from generation_loop)
- Add device checks before CPU moves to avoid redundant operations
- Improve debug.log and debug.log_memory_state() formatting, content, and categories for better visibility
- Remove unused imports and excessive clear_memory() calls
- Clean up non-essential logs from always display
This is a work-in-progress commit that consolidates a series of changes to fix memory leaks, optimize VRAM/RAM usage, improve performance, and enhance code maintainability.
* **BlockSwap Pinned Memory:** Disabled `use_non_blocking=True` for CPU-to-GPU transfers to resolve a memory leak where pinned memory was not being released.
* **Logging-Induced Leaks:** Modified `log_memory_state()` to avoid holding references to tensors during analysis and added a history limit to the checkpoint system to prevent unbounded memory growth.
* **Incomplete Model Cleanup:** Ensured models are completely deleted and their tensor storage is released when `cache_model=False`.
* **Lingering Tensors:** Fixed an issue where a scalar tensor from sampling timesteps and text embeddings remained on the GPU between batches when `preserve_vram` is active.
* **Centralized Cleanup Functions:** Introduced `clear_memory()` to replace `clear_vram_cache()` and all manual `torch.cuda.empty_cache()` calls, providing consistent VRAM/RAM cleanup logic. The function features a `full` parameter to distinguish between a fast, GPU-only cache clear (~1-5ms) for frequent operations and a full cleanup with garbage collection (~10-50ms) for critical stages.
* **Direct-to-CPU Model Loading:** Modified DiT/VAE weight loading to load directly onto the CPU when `preserve_vram` or `BlockSwap` is active, avoiding unnecessary VRAM spikes during model preparation.
* **VAE Device Management:** Created the `manage_vae_device()` helper function to centralize the logic for moving the VAE between the CPU and GPU, reducing code duplication. This also fixed a bug that incorrectly kept the VAE on the GPU when `preserve_vram` was active.
* **CPU Offloading:** Implemented logic to move text embeddings and sampling timesteps to the CPU after each batch when `preserve_vram` is active, reducing idle VRAM usage.
* **VAE Decode Performance:** Replaced proactive, frequent memory clearing during VAE decode with a reactive Out-of-Memory (OOM) handling system. This fixed a significant performance regression and eliminated the need for the `keep_vae_loaded_during_decode` flag.
* **Reduced Overhead:** Removed redundant `gc.collect()` calls from multiple locations to decrease unnecessary processing overhead.
* **Interval-Based VRAM Tracking:** Modified the logging system to reset peak VRAM statistics after each `log_memory_state()` call, enabling accurate tracking of peak memory usage for specific processing intervals (e.g., encode, inference, decode).
* **Accurate RAM Monitoring:** Added the `get_ram_usage()` function for correct process-specific RAM tracking.
* **Efficient Log Refactoring:** Refactored `log_memory_state()` into modular helper methods, optimizing tensor analysis into a single-pass `gc` iteration to improve both performance and maintainability.
* **Log Clarity:** Refined memory state and debug logging to remove redundant snapshots and add new ones for critical operations like model loading, weight loading, VAE encoding, and decoding. Standardized log message conventions.
* **Per-Batch Timers:** Implemented timer namespacing to ensure that performance timers for each batch are logged correctly without overwriting one another.
* **Error Handling:** Added `try/except` blocks to key memory and device management functions to handle edge cases and improve robustness.
* **Code Cleanup:** Removed deprecated code and outdated comments throughout the related modules.
* **Documentation:** Updated comments and function docstrings to reflect the new memory management architecture.
- Move cache_model and enable_debug options from BlockSwap to main SeedVR2 node
- cache_model: Keep models in RAM between runs (skip reload for faster iterations)
- enable_debug: Show detailed memory/timing tracking with hierarchical display
- Replace hardcoded debug statements with unified Debug class
- Add visual icons for log categorization
- Implement parent-child timing relationships for operation breakdown
- Add VRAM/RAM usage tracking with before/after comparisons
Model caching now available for all workflows, not just BlockSwap users.
Debug system provides granular performance insights when enabled.
Simplified the code and used torchvision.datasets.utils.download_url instead of hf_hub_download to be consistent with ComfyUI Manager and show the progress bar
Download progress is now displaying in the logs, e.g.
📥 Downloading model: seedvr2_ema_3b_fp16.safetensors
seedvr2_ema_3b_fp16.safetensors: 6% @ 45.9MB/s
local_dir is using base_cache_dir if model_dir is None else model_dir (was hardcoded to base_cache_dir before)
Adrien Toupet