Commit Graph
49 Commits
Author SHA1 Message Date
Adrien Toupet e86d54ab3b fix: AdaIN color correction and AMD ROCm compatibility (v2.5.4)
- Fix AdaIN non-contiguous tensor error by using reshape() instead of view()
- Add cuDNN availability checks to prevent ROCm 'ATen not compiled with cuDNN' error
2025-11-08 09:11:55 -05:00
Adrien Toupet 9715d3e37a Fix: torch.mps AttributeError on Windows
Add defensive checks for torch.mps.is_available() to handle PyTorch versions where the method doesn't exist on non-Mac platforms. Resolves AttributeError: module 'torch.mps' has no attribute 'is_available'
2025-11-07 23:37:35 -05:00
Adrien Toupet afd5950d37 feat: torch.compile optimization and memory tracking improvements
- Eliminate graph breaks in na.py and attention.py for full torch.compile support
  * Replace cumsum-based tensor slicing with _tensor_split to avoid .item() calls
  * Use torch.tensor_split with .long().cpu() for PyTorch API requirements
  * Replace reshape-based averaging with split-stack-mean pattern
- Fix phase peak VRAM tracking to capture peaks during OOM retry cycles
- Standardize phase4 naming to 'postprocessing' for consistency
- Update RoPE docstrings for accuracy
- Update example workflows with icon
2025-11-06 17:48:39 -05:00
Adrien Toupet f8627298f7 Add peak VRAM tracking and summary display by phase 2025-11-05 21:21:23 -05:00
Adrien Toupet 9b79254c39 refactor(cli): improvements and bug fixes + 3b-Q8_0.gguf support
- Fix validation cache location to respect --model_dir parameter
- Fix output path handling for directories without extensions
- Remove spurious directory creation in get_base_cache_dir
- Enhanced CLI help with usage examples and argument grouping
- Auto-display help when script invoked without arguments
- Correct type hints (device_id: str, debug: Debug)
- Remove redundant type conversions and makedirs calls
- Reorganize imports to module top for clarity
- Improved docstrings & tooltip
- Change default batch_size from 1 to 5 to match ComfyUI integration
- Add support for seedvr2_ema_3b-Q8_0.gguf model
2025-11-05 00:28:19 -05:00
Adrien Toupet ad020d3803 feat(cli): improve UX with auto-format detection, FPS tracking, and consistent messaging with ComfyUI implementation
- Auto-detect output format per file type (mp4 for videos, png for images)
- Add visual separators between processed files for better readability
- Simplify FPS calculation to use wall-clock time for real-world throughput
- Consolidate banner/footer into shared Debug methods
- Update offload device args to support multi-GPU (cpu/cuda:N)
- Standardize terminology: 'upscaling' instead of 'video upscaling'
- Remove code duplication between CLI and ComfyUI implementations
- Consistent quote style (double quotes) throughout CLI
2025-11-04 11:54:57 -05:00
Adrien Toupet 9268346388 feat: CLI Add batch processing, fix multiprocessing issues, and unify model paths
Major Features:
- Renamed --video_path to --input supporting video files, images, and directories
- Added batch processing for directories (iterates all media files)
- Added single image upscaling with extract_frames_from_image()
- Auto-detect output format: images→PNG, videos→MP4 (overridable)
- Smart output path generation (single PNG vs frame sequences)

Critical Bug Fixes:
- Fixed 'str' object has no attribute 'type' by normalizing devices to torch.device
- Fixed 'Got unsupported ScalarType BFloat16' by converting ML dtypes to float32
- Fixed prepare_runner() signature mismatch (returned 2 values, claimed 3)
- Fixed KeyError 'cache_context' by storing cache_context in ctx
- Fixed duplicate optimization logging (3x imports) using environment variable

Performance Improvements:
- Removed mp.Manager() overhead
- Using direct mp.Queue(maxsize=0) for better throughput
- Improved multiprocessing reliability

Consistency & Quality:
- Unified model directory between CLI & ComfyUI to models/SEEDVR2 using constants
- Default CLI output folder to use ./output/
2025-11-04 00:12:13 -05:00
Adrien Toupet 3a4a4900df feat: optimize SeedVR2 memory management and color ops
- Optimize color operations for VRAM efficiency with safe precision handling
- Limit clear_memory calls to OOM retry and final cleanup only
- Fix model cleanup regression ensuring proper GPU memory release
2025-10-27 14:58:20 -04:00
Adrien Toupet c9dce827c0 feat: Add deterministic generation with seed control and CFG scale parameter
Core Changes:
- Implement deterministic generation with phase-specific seeding strategy
  * VAE encoding: seed+1M for deterministic sampling without quality loss
  * DiT upscaling: base seed for reproducible noise generation
- Add cfg_scale parameter for user control of upscaling strength (WIP)
- Fix ComfyUI V3 unique_id extraction using get_executing_context().node_id

Improvements:
- Standardize Optional['Debug'] type hints across codebase
- Make debug parameter required where it's essential (generate, infer)
- Remove legacy get_unique_id() stack inspection approach
- Add seed and cfg_scale logging for transparency
- Fix FP8CompatibleDiT parameter order consistency
- Refine input/latent noise scale steps (0.01 → 0.001 for finer control)
2025-10-21 13:31:54 -04:00
Adrien Toupet e735c2ad56 feat: V3 migration with GGUF fixes and attention optimizations
Major Changes:
- Migrate all nodes to ComfyUI V3 schema (stateless design, new IO types)
- Fix GGUF weight caching VRAM leak (non-persistent buffers + _apply override)
- Fix GGUF torch.compile compatibility (@torch._dynamo.disable on dequant)
- Centralize compatibility checks (Flash/Triton/GGUF/Conv3d in compatibility.py)
- Make flash_attn optional with graceful SDPA fallback
- Add attention_mode UI option (sdpa/flash_attn) to DiT loader
- Fix color correction batch padding error (trim input_video consistently)

Code Quality:
- Remove internal_execute for clarity (stateless node design)
- Add startup logging for optimization status
- Improve error messages with installation instructions
- Add get_unique_id() for V3 node-specific caching
- Standardize parameter names (dit_cache/vae_cache)
2025-10-21 00:10:37 -04:00
Adrien Toupet 9893a1d7b0 refactor: Add indent_level parameter to debug logging system
- Add indent_level param to debug.log() and manage_tensor()
- Replace all hardcoded spaces with indent_level (0/1/2)
- Fix encode_all_batches memory flow: encode before storage
- Minor whitespace cleanup
2025-10-19 10:13:17 -04:00
Adrien Toupet 97c7b9cd12 refactor: Overhaul device and memory management architecture
Major improvements to device handling, model offloading, and code clarity:

Device Management:
- Replace string-based devices with torch.device objects throughout
- Add explicit offload device parameters: dit_offload_device, vae_offload_device, tensor_offload_device
- Remove preserve_vram in favor of explicit offload control
- Improve get_device() to return torch.device objects consistently
- Enhance get_device_list() with smart MPS-only system handling

Context & Pipeline:
- Merge setup_device_environment and prepare_generation_context into single setup_generation_context
- Simplify LOCAL_RANK handling (set to '0' for single-GPU mode)
- Store device configuration on runner for submodule access

Tensor & Model Management:
- Add manage_tensor_device() for consistent tensor movement with logging
- Update manage_model_device() to use torch.device objects
- Add validation for BlockSwap and caching configurations
- Rename cache_in_ram to cache_model (more accurate naming)

CLI & Interface:
- Remove --preserve_vram flag (breaking change)
- Add --vae_offload_device and --tensor_offload_device flags
- Update ComfyUI node parameters with validation and better tooltips
- Improve device selection UI with offload_device options
2025-10-17 17:24:53 -04:00
Adrien Toupet a46e16d8cc refactor: 42, centralize float32 precision utilities and improve logging
- Move ensure_float32_precision() from alpha_upscaling to common/half_precision_fixes
- Update color_fix.py to use ensure_float32_precision
- Add type hints and docstrings to all half_precision_fixes functions
- Clarify batch size tips: emphasize avoiding padding waste over maximizing batch size
- Update default seed to 42 (the answer to life, the universe, and everything)
2025-10-16 14:30:47 -04:00
Adrien Toupet f1ac7d786e Add perceptually-accurate color correction methods
- 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
2025-10-15 13:14:48 -04:00
Adrien Toupet e48c2da776 feat: global model cache, alpha upscaling rewrite, and quality improvements
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
2025-10-15 00:51:56 -04:00
Adrien Toupet d6e04aa108 refactor: modularize node architecture and enhance BlockSwap flexibility
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
2025-10-11 17:37:58 -04:00
Adrien Toupet bf46dfd674 Simplify RGBA handling and improve logging consistency
- 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
2025-10-11 11:06:15 -04:00
Adrien Toupet 1b00665757 feat: Add RGBA support, improve config caching, and enhance logging
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
2025-10-11 02:26:59 -04:00
Adrien Toupet 9ee244d71b feat: implement torch.compile optimization with BlockSwap compatibility
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_*)
2025-10-10 12:34:49 -04:00
Adrien Toupet 5567f232af feat!: independent DiT/VAE management with separate loader nodes
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
2025-10-08 00:49:48 -04:00
Adrien Toupet 1fa603899f Fix GGUF implementation: centralize warning suppression, improve type hints, enhance error handling
- 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
2025-09-24 15:47:18 -04:00
Adrien Toupet 0b0c87ed4a Add GGUF quantized model support (based on PR #121 from @cmeka / @lihaoyun6)
- 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
2025-09-24 11:49:57 -04:00
Adrien Toupet 953524b13b feat: Add support for extra_model_paths.yaml configuration
- 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
2025-09-23 00:56:26 -04:00
Adrien Toupet efa851009f Fix corrupted model downloads with hash validation and resume support
- Add SHA256 hash validation for all models and VAE
- Implement download resume capability with progress bars
- Add validation cache to avoid redundant checks
- Replace unreliable torchvision downloads with urllib
- Fix safetensors header corruption detection
- Move SeedVR2 ASCII logo with credits before download operations for clearer logs

Addresses #128, #151
2025-09-22 23:45:43 -04:00
Adrien Toupet 2b00252a49 Simplify __init__.py files + fix dependencies + add pyproject.toml
- 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
2025-09-22 15:47:04 -04:00
Adrien Toupet d0d24a18cb Refactor: Replace dynamic relative imports with model class registry for cross-platform compatibility. Minor: Adjust ASCII logo width for better screen fit 2025-09-20 01:53:25 -04:00
Adrien Toupet 1518ecc1a5 refactor: Fix ComfyUI node conflicts via relative imports and clearer model structure
- 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
2025-09-16 16:45:31 -04:00
Adrien Toupet 477f57fd5a Refactor: Three-phase batch processing pipeline for improved performance
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
2025-09-16 14:38:57 -04:00
Adrien Toupet 66bcf436ab perf: optimize debug logging to reduce performance overhead
- 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
2025-08-28 23:50:49 -04:00
Adrien Toupet 89bde29cbe perf: optimize model initialization with meta device for 90% speedup
- 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
2025-08-26 13:23:30 -04:00
Adrien Toupet 5595d58597 refactor: streamline generation pipeline and improve dtype handling
- 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"
2025-08-25 17:11:29 -04:00
Adrien Toupet 47552c0ed5 refactor: unify model configuration and improve logging precision
- 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
2025-08-25 13:45:13 -04:00
Adrien Toupet 8333fb856e refactor(WIP): complete memory management overhaul with proper error handling
Memory Management:
- Add type hints to all memory functions for better IDE support and maintainability
- Replace silent exception handling with debug logging across all operations
- Removed unnecessary CPU transfers for GPU memory release
- Introduce unified manage_model_device() for consistent device management
- Remove torch._C._clear_cache() private API usage (incompatible across PyTorch versions)

Performance & Debugging:
- Add debug timers to critical operations (clear_memory, clear_runtime_caches, etc.)
- Consolidate all cleanup code into core functions: manage_*, release_*, clear_*, complete_cleanup
- Improved timer log messages for clarity

Code Quality:
- Remove unused imports
- Remove redundant RoPE cache clearing (clear_runtime_caches handles it)
- Simplify configure_runner() by eliminating duplicate code paths
- Update infer.py to use generic device management functions
- Add release_text_embeddings() helper to deduplicate embedding cleanup
2025-08-23 01:13:45 -04:00
Adrien Toupet 2053a80f38 refactor(WIP): Improve debug logging consistency and reduce redundancy
- 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
2025-08-22 15:38:00 -04:00
Adrien Toupet 9b7c113681 refactor(WIP): Memory Management Overhaul
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.
2025-08-22 09:20:41 -04:00
lihaoyun6 176562f829 Restored compatibility with FP8 safetensors for MPS backend; Changed the way to detect MPS device 2025-08-13 00:53:05 +08:00
Benjamin Herb fe67b2a042 feat: Add BlockSwap support for standalone CLI and update logging
- Add BlockSwap CLI arguments and activation
- Fix imports to work without ComfyUI dependencies
- Update to new debugging framework
2025-08-09 02:27:42 +02:00
Adrien Toupet 6d8c278a4e Fix QC bugs post updates 2025-08-08 00:53:19 +02:00
Adrien Toupet 1da2cdaf12 Fix VRAM memory leak when using cache_model, added log_memory_state detailed_tensors debug output for troubleshooting 2025-08-08 00:51:02 +02:00
Adrien Toupet d075dde799 feat: Extract model caching to main node + add hierarchical debug system
- 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.
2025-08-08 00:49:46 +02:00
Adrien Toupet 4118595968 Revert to 5c65c44 and fix integration issues
- Resolved merge conflicts

- Added missing constants.py and model_registry.py files

- Added graceful error handling for failed model downloads
2025-08-08 00:49:46 +02:00
Adrien Toupet 5682230ee9 Defaulting to latest diffusers version or higher to avoid huggingface_hub incompatibility 2025-07-23 12:29:33 -04:00
Adrien Toupet 9a42320977 Use torchvision.datasets.utils.download_url
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
2025-07-18 11:19:59 -04:00
Adrien Toupet afe33ddf0d Download progress bar + fixed local_dir
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)
2025-07-17 21:48:05 -04:00
NumZ 9461339bb9 Remove __init__ 2025-07-03 15:05:13 +02:00
NumZ d5ce75ad50 Add MultiGPU Standalone inference 2025-07-03 04:28:08 +02:00
NumZ 5c27397a40 Add MultiGPU Standalone inference 2025-07-03 04:22:28 +02:00
NumZ eb05b8ad26 fix import 2025-07-01 13:33:50 +02:00
NumZ 2786c05faa Speed Up 30 to 50%, fix Memory Leak, refacto) 2025-06-30 12:46:41 +02:00