Compare commits

...
52 Commits
Author SHA1 Message Date
Adrien Toupet 58bc9e8bc9 Merge pull request #373 from AInVFX/main
v2.5.17: Proper bf16 detection for older GPUs #314
2025-12-05 21:07:21 -05:00
Adrien Toupet 3eec5847c0 Release v2.5.17: Older GPU compatibility fix 2025-12-05 21:05:30 -05:00
Adrien Toupet eae3aac60d Fix CUBLAS_STATUS_NOT_SUPPORTED on older GPUs via bf16 probe (again\!) (#314) 2025-12-05 20:10:36 -05:00
Adrien Toupet 0a660065f0 Merge pull request #371 from AInVFX/main
v2.5.16: Older GPU compatibility fix, quality regression fix, debug improvements
2025-12-05 15:52:34 -05:00
Adrien Toupet 11239eed13 Release v2.5.16: Quality regression fix, older GPU compatibility fix, system info debug 2025-12-05 15:50:08 -05:00
Adrien Toupet b4d7ab89eb Fix CUBLAS_STATUS_NOT_SUPPORTED on older GPUs (GTX 970) #314
Add automatic bfloat16 → float16 SDPA fallback for GPUs without native bf16 cuBLAS support
2025-12-05 15:39:22 -05:00
Adrien Toupet f061d97fe7 Revert bfloat16 detection - was causing quality regression / keep ensure_triton_compat() 2025-12-05 15:03:40 -05:00
Adrien Toupet 43c4f00e19 Revert bfloat16 detection - was causing quality regression / keep ensure_triton_compat() 2025-12-05 15:01:58 -05:00
Adrien Toupet f8998ebd75 Revert bfloat16 detection - was causing quality regression 2025-12-05 14:55:30 -05:00
Adrien Toupet aa968cf2c9 docs: simplify contribution workflow to main branch only 2025-12-05 13:54:39 -05:00
Adrien Toupet d78f6c268d Merge branch 'main' of https://github.com/ainvfx/ComfyUI-SeedVR2_VideoUpscaler 2025-12-05 11:12:32 -05:00
Adrien Toupet 18b44d66e1 feat: add environment info display in debug mode to help with issue reporting 2025-12-05 11:11:18 -05:00
Adrien Toupet f68fe920b8 Merge pull request #358 from AInVFX/main
v2.5.15: MPS compatibility fixes, autocast device type, VRAM tracking, triton 3.0+ compatibility
2025-12-03 13:16:46 -05:00
Adrien Toupet 65c1c1b6cd Release v2.5.15: MPS fixes, autocast device type, accurate VRAM tracking, triton 3.0 compatibility 2025-12-03 13:14:51 -05:00
Adrien Toupet b40f26167c Fix MPS compatibility: disable antialias for MPS tensors, fix bfloat16 arange (#354) 2025-12-03 13:09:59 -05:00
Adrien Toupet ed53581359 Fix triton.ops compatibility for bitsandbytes 0.45+ / triton 3.0+
Fixes #340 - Installation error with PyTorch 2.7+cu126 and triton_windows

Add compatibility shim for missing triton.ops.matmul_perf_model module.
Reverts local VAE types approach
2025-12-03 12:43:35 -05:00
Adrien Toupet 71ac9ffe54 fix: use max_memory_reserved for accurate VRAM peak tracking 2025-12-03 11:51:30 -05:00
Adrien Toupet ffba05907d Fix autocast device_type error by using .type attribute instead of str() #350 2025-12-03 11:32:39 -05:00
Adrien Toupet d4dd5e747d Merge pull request #344 from AInVFX/main
v2.5.14: MPS device fix, VRAM swap detection, enforce physical VRAM limit
2025-12-01 00:33:53 -05:00
Adrien Toupet e2faedaaa6 Release v2.5.14 - MPS device fix, VRAM swap detection, enforce physical VRAM limit 2025-12-01 00:30:49 -05:00
Adrien Toupet 5775ff0f99 Enforce VRAM limit to physical capacity - OOM instead of silent swap 2025-12-01 00:25:26 -05:00
Adrien Toupet ff937756a3 Add VRAM swap detection - show GPU+swap breakdown in peak stats, warn when swap detected 2025-11-30 23:47:20 -05:00
Adrien Toupet 5848cef05f fix(mps): normalize device strings to prevent unnecessary tensor movements
- Add _device_str() helper to normalize MPS variants (mps:0 → MPS)
- Fix device comparison: mps:0 and mps now correctly identified as same device
- Consistent MPS logging across all memory management functions
2025-11-30 21:03:26 -05:00
Adrien Toupet f5b902b8b0 Merge pull request #341 from AInVFX/main
v2.5.13: Fix triton import error, OOM on long video float32 conversion, macOS CLI watermark
2025-11-30 09:04:04 -05:00
Adrien Toupet 19825fa9fa Release v2.5.13: Fix triton import, OOM on long videos, macOS watermark 2025-11-30 09:01:53 -05:00
Adrien Toupet 021fc7b70f Fix triton.ops import error by using local VAE types
Fixes #340 - Installation error with PyTorch 2.7+cu126 and triton_windows

Replace diffusers.models.autoencoders.vae imports with local implementations
of DecoderOutput and DiagonalGaussianDistribution to avoid triggering the
bitsandbytes -> triton.ops import chain that fails on newer triton versions.
2025-11-30 08:27:01 -05:00
Adrien Toupet b0ac880a4f Fix OOM crash on float32 conversion for long videos. Gracefully fallback to native dtype if insufficient memory. Fixes #299 2025-11-30 08:03:07 -05:00
Adrien Toupet 6930e0f13a Fix CLI MPS watermark error on macOS (fixes #336) 2025-11-30 07:37:51 -05:00
Adrien Toupet 112275d4e5 Merge pull request #333 from AInVFX/main
release: v2.5.12 - Regression fix
2025-11-28 17:51:53 -05:00
Adrien Toupet 16508d353d release: v2.5.12 - fix color artifacts regression from in-place transform ops 2025-11-28 17:50:40 -05:00
Adrien Toupet 6f784b1324 Merge pull request #332 from AInVFX/main
v2.5.11: CUDNN attention, streaming VAE decode, LAB artifact fix, MPS/multi-GPU improvements
2025-11-28 17:09:48 -05:00
Adrien Toupet f791631495 release: v2.5.11 - CUDNN attention, long video memory fix, LAB artifacts fix, MPS/multi-GPU improvements 2025-11-28 16:49:08 -05:00
Adrien Toupet b15982b442 perf: in-place tensor ops to reduce memory allocation overhead
- Use in-place clamp/mul/add in postprocess normalization
- Use in-place clamp and Normalize in video transform pipeline
- Consolidate optimized channel permute functions in performance.py
- Use native permute() instead of einops rearrange in VAE encode/decode
- In-place wavelet decomposition and reconstruction operations
- Remove unused temporal_latent_blending function
2025-11-28 16:40:27 -05:00
Adrien Toupet 849f6b5997 Add peak RAM tracking alongside VRAM in debug summary 2025-11-28 15:32:07 -05:00
Adrien Toupet e34bec76d5 Auto-detect bfloat16 support to fix CUBLAS_STATUS_NOT_SUPPORTED on older GPUs (#314) 2025-11-28 14:30:41 -05:00
Adrien Toupet 70c7f21eb5 docs: clarify multi-GPU frame-level parallelism behavior in README 2025-11-28 14:00:59 -05:00
Adrien Toupet daa13fb4ca Fix BlockSwap logging confusion and CLI worker validation
- BlockSwap: Show effective/total blocks (e.g., 32/32) instead of raw requested value
- CLI: Skip CUDA device validation when CUDA_VISIBLE_DEVICES already set (worker process)
2025-11-28 13:40:07 -05:00
Adrien Toupet 7581e0014f Fix CLI: MPS subprocess allocator error (#290) and multi-GPU distribution (#309)
- Mac: Use direct processing instead of spawning subprocess (MPS allocator fails in child process)
- Multi-GPU: Set CUDA_VISIBLE_DEVICES before spawn so child inherits it before module-level torch import
- Remove redundant env setup in worker (now inherited from parent)
2025-11-28 12:11:42 -05:00
Adrien Toupet be2efd474f fix: use canonical torch.backends.mps.is_available() for reliable MPS detection
Replaces torch.mps.is_available() with torch.backends.mps.is_available()
across all files. The torch.backends API is the official PyTorch method
for MPS detection (since PyTorch 1.12) and works reliably on:
- macOS with Apple Silicon (returns True when MPS available)
- macOS without MPS support (returns False)
- Windows/Linux (hasattr guard prevents AttributeError)
2025-11-28 11:49:43 -05:00
Adrien Toupet bce62382db Fix color correction reference misalignment with temporal overlap 2025-11-28 09:43:46 -05:00
Adrien Toupet f792ca5937 fix: eliminate LAB color correction tile artifacts by preprocessing with wavelet reconstruction
Fixes #324
Fixes https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/discussions/271#discussioncomment-14997355
2025-11-28 00:10:17 -05:00
Adrien Toupet 3ab2c84091 fix: stream VAE decode directly to pre-allocated final_video tensor
Eliminates memory spike for long videos by:
- Pre-allocating final_video in Phase 3 before any batch processing
- Writing VAE-decoded RGB directly to final_video (no batch_samples accumulation)
- For RGBA: allocating 4 channels upfront, writing RGB in Phase 3, alpha in Phase 4
- Processing batches in-place through color correction and normalization
- Moving temporal overlap blending to Phase 3 decode phase

Fixes #130 (DefaultCPUAllocator: not enough memory for long video)
2025-11-27 20:15:42 -05:00
Adrien Toupet 8dd0c3061a feat: Add CUDNN attention backend support with PyTorch 2.3+ API
- Use torch.nn.attention.sdpa_kernel() (new API) with CUDNN_ATTENTION backend
- Fallback to torch.backends.cuda.sdp_kernel() when not available
Thanks to @eadwu for the original PR #317
2025-11-27 15:19:55 -05:00
Adrien Toupet 79e7f41216 Revert "Fix: MPS allocator error in model.to() transfer (#305)"
This reverts commit 9424020687.
2025-11-27 14:52:12 -05:00
Adrien Toupet 9424020687 Fix: MPS allocator error in model.to() transfer (#305)
Add fallback to individual parameter movement when bulk model.to(mps)
fails with allocator errors. Complements safetensors loading fix.
2025-11-14 12:06:13 -05:00
Adrien Toupet 912ab4a5da Merge pull request #298 from AInVFX/main
v2.5.10: Fix determinism, BlockSwap caching, and model path resolution
2025-11-13 12:04:00 -05:00
Adrien Toupet 65dd29a865 v2.5.10: Fix determinism, BlockSwap caching, and model path resolution
Core Fixes:
- Reset seed per batch to ensure deterministic generation across sessions and batch positions
- Fix temporal overlap logging when automatically reset to prevent incorrect frame counts
- Fix NoneType attribute error in VAE tiled encode/decode at maximum resolution (#296)

BlockSwap & Caching Architecture:
- Move BlockSwap state (_block_swap_config, _blockswap_bypass_protection) from runner to model
- Ensures state survives runner recreation during independent DiT/VAE caching scenarios
- Fix runner template caching to trigger when either DiT or VAE becomes cached (bidirectional)
- Resolves BlockSwap reload failures when only DiT was cached (#297)

Model Discovery:
- Implement case-insensitive YAML path resolution for extra_model_paths.yaml (#289-#295)
- Add debug logging for model discovery (searched paths, validation status, cache hits)
- Support any case variation (seedvr2, SEEDVR2, SeedVR2) in ComfyUI configuration
2025-11-13 12:02:37 -05:00
Adrien Toupet 9f78b30b3d Merge pull request #291 from AInVFX/main
Release v2.5.9: Bug fixes and enhancements
2025-11-12 14:36:32 -05:00
Adrien Toupet 96b26bb7f7 Update FUNDING.yml 2025-11-12 14:33:43 -05:00
Adrien Toupet 513ab67214 Update FUNDING.yml 2025-11-12 14:33:37 -05:00
Adrien Toupet 19c761522e Update FUNDING.yml 2025-11-12 14:30:39 -05:00
Adrien Toupet 40d37b7248 Release v2.5.9: Bug fixes and enhancements
- Fix: OpenCV memory layout error in tile debug visualization (#283)
- Fix: macOS MPS allocator fallback for safetensors loading (#290)
- Fix: Windows log buffering with flush=True (#278)
- Fix: ComfyUI registry icon URL (raw.githubusercontent)
- Feature: Version display in node name and CLI/ComfyUI header
- Feature: GitHub Sponsors support link
- License: Migrate from MIT to Apache 2.0 to match Bytedance Seed official repo
2025-11-12 14:27:40 -05:00
26 changed files with 1176 additions and 572 deletions
+1
View File
@@ -0,0 +1 @@
github: [adrientoupet, numz]
+2 -8
View File
@@ -16,15 +16,9 @@ We welcome code contributions from the community. Here's how to propose changes:
1. Fork this repository to your own GitHub account.
2. Create a new branch on your fork for your changes.
3. Make your changes in this branch.
4. When you are ready, submit a pull request to either:
- **`main` branch** - For stable features and bug fixes
- **`nightly` branch** - For experimental features or if you want faster integration (note: this branch changes more frequently and may be less stable)
4. When you are ready, submit a pull request to the **`main`** branch.
### Branch Guidelines
- **Main Branch**: Use this for production-ready changes, critical bug fixes, and well-tested features
- **Nightly Branch**: Use this for experimental features, cutting-edge improvements, or if you need your changes integrated quickly. Be aware that this branch is actively under development and may contain unstable features.
We use the GitHub Flow workflow. Choose the appropriate target branch based on the nature of your contribution.
We use the GitHub Flow workflow.
Before submitting a pull request, please make sure your code adheres to the project's coding conventions and it has passed all tests. If you are adding features, please also add appropriate tests.
+197 -17
View File
@@ -1,21 +1,201 @@
MIT License
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
Copyright (c) 2025 NumZ, AInVFX and Contributors
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
1. Definitions.
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "{}"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright 2025 seed
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
+97 -13
View File
@@ -35,6 +35,76 @@ We're actively working on improvements and new features. To stay informed:
- **🔮 Next Model Survey**: We're looking for community input on the next open-source super-powerful generic restoration model. Share your suggestions in [Issue #164](https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/issues/164)
## 🚀 Updates
**2025.12.05 - Version 2.5.17**
- **🔧 Fix: Older GPU compatibility (GTX 970, etc.)** - Runtime bf16 CUBLAS probe replaces compute capability heuristics, correctly detecting unsupported GPUs without affecting RTX 20XX
**2025.12.05 - Version 2.5.16**
- **🔧 Fix: Older GPU compatibility (GTX 970, etc.)** - Automatic fallback for GPUs without bfloat16 support
- **🐛 Fix: Quality regression** - Reverted bfloat16 detection that was causing artifact issues
- **📋 Debug: Environment info display** - Shows system info in debug mode to help with issue reporting
- **📚 Docs: Simplified contribution workflow** - Streamlined to main branch only
**2025.12.03 - Version 2.5.15**
- **🍎 Fix: MPS compatibility** - Disable antialias for MPS tensors and fix bfloat16 arange issues
- **⚡ Fix: Autocast device type** - Use proper device type attribute to prevent autocast errors
- **📊 Memory: Accurate VRAM tracking** - Use max_memory_reserved for more precise peak reporting
- **🔧 Fix: Triton compatibility** - Add shim for bitsandbytes 0.45+ / triton 3.0+ (fixes PyTorch 2.7 installation errors)
**2025.12.01 - Version 2.5.14**
- **🍎 Fix: MPS device comparison** - Normalize device strings to prevent unnecessary tensor movements
- **📊 Memory: VRAM swap detection** - Peak stats now show GPU+swap breakdown when overflow occurs, with warning when swap detected
- **🛡️ Memory: Enforce physical VRAM limit** - PyTorch now OOMs instead of silently swapping to shared memory (prevents extreme slowdowns on Windows)
**2025.11.30 - Version 2.5.13**
- **🔧 Fix: PyTorch 2.7+ triton import error** - Resolved installation crash caused by triton.ops import chain on newer triton versions
- **💾 Fix: OOM on float32 conversion for long videos** - Graceful fallback to native dtype when insufficient memory for float32 conversion
- **🍎 Fix: CLI watermark error on macOS** - Resolved MPS-related watermark processing crash on Apple Silicon
**2025.11.28 - Version 2.5.12**
- **🐛 Fix: Color artifacts regression** - Reverted in-place tensor operations in video transform pipeline that caused color artifacts on some images
**2025.11.28 - Version 2.5.11**
- **⚡ Feature: CUDNN attention backend** - Added support for PyTorch 2.3+ CUDNN_ATTENTION backend with automatic fallback for older versions (thanks @eadwu)
- **💾 Fix: Memory spike for long videos** - VAE decode now streams directly to pre-allocated tensor, eliminating OOM errors during long video processing
- **🎨 Fix: LAB color correction artifacts** - Resolved tile boundary artifacts using wavelet reconstruction preprocessing
- **🎨 Fix: Color reference misalignment** - Fixed color correction frame alignment with temporal overlap
- **🍎 Fix: MPS detection reliability** - Switched to canonical `torch.backends.mps.is_available()` API for consistent Apple Silicon detection
- **🖥️ Fix: Mac subprocess error** - CLI now uses direct processing on Mac to avoid MPS allocator failures in child processes
- **🖥️ Fix: Multi-GPU device assignment** - CUDA_VISIBLE_DEVICES now set before spawn for proper worker inheritance
- **📊 Fix: BlockSwap logging** - Now shows effective/total blocks (e.g., 32/32) instead of raw requested value
- **🔧 Feature: Auto bfloat16 detection** - Automatically detects bfloat16 support to prevent CUBLAS errors on older GPUs
- **📊 Feature: Peak RAM tracking** - Added RAM usage alongside VRAM in debug summary
- **⚡ Performance: In-place tensor ops** - Reduced memory allocation overhead with in-place operations throughout pipeline
- **📖 Docs: Multi-GPU clarification** - Clarified frame-level parallelism behavior expectations for multi-GPU setups
**2025.11.13 - Version 2.5.10**
- **🎯 Fix: Deterministic generation** - Identical images with the same seed now produce identical results across different sessions and batch positions
- **🔧 Fix: Model caching with BlockSwap** - Resolved issue where cached DiT models wouldn't properly reload when VAE caching state changed
- **💾 Fix: Runner caching optimization** - Runner templates now correctly cache whenever both DiT and VAE are cached, regardless of caching order
- **📁 Fix: Case-insensitive model paths** - Extra model paths in YAML config now work regardless of case (seedvr2, SEEDVR2, SeedVR2, etc.)
- **🐛 Fix: High resolution tile debug crash** - Fixed "NoneType has no attribute log" error when using maximum resolution with VAE tiling
- **📊 Fix: Temporal overlap logging** - Corrected frame count reporting when temporal overlap is automatically adjusted
- **🔍 Feature: Enhanced model path debugging** - Added detailed logging to help troubleshoot model loading issues (visible in debug mode)
**2025.11.12 - Version 2.5.9**
- **🐛 Fix: Tile debug visualization crash** - Fixed OpenCV error when using VAE tile debug mode on certain systems.
- **🍎 Fix: macOS MPS loading error** - Added automatic CPU fallback for MPS allocator issues on certain PyTorch/macOS versions.
- **🖥️ Fix: Windows log buffering** - Added flush to print statements for real-time log visibility in ComfyUI on Windows
- **📦 Fix: ComfyUI Registry logo** - Updated icon URL to display properly in ComfyUI node registry
- **ℹ️ Feature: Version display** - Added version number to node name and CLI/ComfyUI header for better tracking
- **💝 Feature: GitHub Sponsors** - Added sponsor button to support project development. Thank you everyone for your support!
- **📜 License: Apache 2.0** - Reverted License from MIT to Apache 2.0 to match ByteDance Seed project
**2025.11.10 - Version 2.5.8**
- **🐛 Fix (CLI): Windows batch processing duplicate files** - Fixed CLI batch mode processing each file twice on Windows due to case-insensitive filesystem. Improved directory scanning performance by 2-3x
@@ -808,26 +878,40 @@ python inference_cli.py media_folder/ \
### Multi-GPU Processing Explained
The CLI's multi-GPU mode automatically distributes the workload across multiple GPUs with intelligent temporal overlap handling:
The CLI's multi-GPU mode uses **frame-level parallelism**: the video is split into chunks and each GPU processes its chunk independently through all 4 phases (encode → upscale → decode → postprocess). This is ideal for long videos where you want to reduce total processing time by dividing the workload.
**How it works:**
1. Video is split into chunks, one per GPU
2. Each GPU processes its chunk independently
3. Chunks overlap by `--temporal_overlap` frames
4. Results are blended together seamlessly using the overlap region
1. Video frames are split evenly across GPUs (e.g., 100 frames on 2 GPUs → 50 frames each)
2. Each GPU loads its own copy of the models and processes its chunk independently
3. When `--temporal_overlap` is set, chunks include overlapping frames for seamless blending
4. Results are concatenated (and blended at overlap regions) into the final video
**Example for 2 GPUs with temporal_overlap=4:**
**Example for 100 frames on 2 GPUs with temporal_overlap=4:**
```
GPU 0: Frames 0-50 (includes 4 overlap frames at end)
GPU 1: Frames 46-100 (includes 4 overlap frames at beginning)
Result: Frames 0-100 with smooth transition at frame 48
GPU 0: Frames 0-53 (50 base + 4 overlap at end, processed as independent video)
GPU 1: Frames 50-99 (50 frames, 4 overlap at start, processed as independent video)
Result: Frames 0-99 with smooth blending at the transition point
```
**Important considerations:**
- Each GPU processes its chunk as a separate video with its own batch splitting
- `batch_size` controls batching *within* each GPU's chunk, not across GPUs
- For short videos (< 100 frames), single GPU is often more efficient due to model loading overhead
- Multi-GPU doubles VRAM usage (each GPU loads full models) but roughly halves processing time
**When to use multi-GPU:**
- Long videos (100+ frames) where splitting provides significant time savings
- When you have multiple GPUs with sufficient VRAM each
**When to use single GPU:**
- Short videos where model loading overhead outweighs parallel gains
- When you want all frames processed together for maximum temporal coherence
**Best practices:**
- Set `--temporal_overlap` to 2-8 frames for smooth blending
- Set `--temporal_overlap` to 2-4 frames for smooth blending between GPU chunks
- Higher overlap = smoother transitions but more redundant processing
- Use `--prepend_frames` to reduce artifacts at video start
- batch_size should divide evenly into chunk sizes for best results
- For optimal quality on short videos, use single GPU with `batch_size` matching your shot length
## ⚠️ Limitations
@@ -875,7 +959,7 @@ For detailed contribution guidelines, see [CONTRIBUTING.md](CONTRIBUTING.md).
2. Create your feature branch (`git checkout -b feature/AmazingFeature`)
3. Commit your changes (`git commit -m 'Add some AmazingFeature'`)
4. Push to the branch (`git push origin feature/AmazingFeature`)
5. Open a Pull Request to **main** branch for stable features or **nightly** branch for experimental features
5. Open a Pull Request to the **main** branch
**Get Help:**
- YouTube: [AInVFX Channel](https://www.youtube.com/@AInVFX)
@@ -891,4 +975,4 @@ Special thanks to our community contributors including [benjaminherb](https://gi
## 📜 License
The code in this repository is released under the MIT license as found in the [LICENSE](LICENSE) file.
The code in this repository is released under the Apache 2.0 license as found in the [LICENSE](LICENSE) file.
+1
View File
@@ -3,6 +3,7 @@ ComfyUI-SeedVR2_VideoUpscaler
Official SeedVR2 integration for ComfyUI
"""
from .src.optimization.compatibility import ensure_triton_compat # noqa: F401
from .src.interfaces import comfy_entrypoint, SeedVR2Extension
__all__ = ["comfy_entrypoint", "SeedVR2Extension"]
+41 -34
View File
@@ -64,8 +64,14 @@ os.environ['PYTHONPATH'] = script_dir + ':' + os.environ.get('PYTHONPATH', '')
if mp.get_start_method(allow_none=True) != 'spawn':
mp.set_start_method('spawn', force=True)
# Configure VRAM management and validate CUDA devices before heavy imports
if platform.system() != "Darwin":
# Configure platform-specific memory management before heavy imports
# Must be set BEFORE import torch
if platform.system() == "Darwin":
# MPS allocator requires: low_watermark <= high_watermark
# Setting both to 0.0 disables PyTorch memory limits, letting macOS manage memory
os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.0")
os.environ.setdefault("PYTORCH_MPS_LOW_WATERMARK_RATIO", "0.0")
else:
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")
# Pre-parse CUDA device argument for validation and environment setup
@@ -76,23 +82,25 @@ if platform.system() != "Darwin":
if _pre_args.cuda_device is not None:
device_list_env = [x.strip() for x in _pre_args.cuda_device.split(',') if x.strip()!='']
# Temporary torch import for CUDA device validation only
# Must happen before setting CUDA_VISIBLE_DEVICES and before main torch import
import torch as _torch_check
if _torch_check.cuda.is_available():
available_count = _torch_check.cuda.device_count()
invalid_devices = [d for d in device_list_env if not d.isdigit() or int(d) >= available_count]
if invalid_devices:
print(f"❌ [ERROR] Invalid CUDA device ID(s): {', '.join(invalid_devices)}. "
f"Available devices: 0-{available_count-1} (total: {available_count})")
# Skip validation if CUDA_VISIBLE_DEVICES is already set (worker process)
if os.environ.get("CUDA_VISIBLE_DEVICES") is None:
# Temporary torch import for CUDA device validation only
# Must happen before setting CUDA_VISIBLE_DEVICES and before main torch import
import torch as _torch_check
if _torch_check.cuda.is_available():
available_count = _torch_check.cuda.device_count()
invalid_devices = [d for d in device_list_env if not d.isdigit() or int(d) >= available_count]
if invalid_devices:
print(f"❌ [ERROR] Invalid CUDA device ID(s): {', '.join(invalid_devices)}. "
f"Available devices: 0-{available_count-1} (total: {available_count})")
sys.exit(1)
else:
print("❌ [ERROR] CUDA is not available on this system. Cannot use --cuda_device argument.")
sys.exit(1)
else:
print("❌ [ERROR] CUDA is not available on this system. Cannot use --cuda_device argument.")
sys.exit(1)
# Set CUDA_VISIBLE_DEVICES for single GPU after validation
if len(device_list_env) == 1:
os.environ["CUDA_VISIBLE_DEVICES"] = device_list_env[0]
# Set CUDA_VISIBLE_DEVICES for single GPU after validation
if len(device_list_env) == 1:
os.environ["CUDA_VISIBLE_DEVICES"] = device_list_env[0]
# Heavy dependency imports after environment configuration
import torch
@@ -394,9 +402,9 @@ def process_single_file(input_path: str, args: argparse.Namespace, device_list:
# Process frames
processing_start = time.time()
# Use direct processing if caching enabled
if runner_cache is not None:
# Direct single-GPU processing with model caching
# Use direct processing if caching enabled OR on Mac (MPS doesn't support multiprocessing well)
if runner_cache is not None or platform.system() == "Darwin":
# Direct single-GPU processing (required for Mac MPS, optional for caching)
result = _single_gpu_direct_processing(frames_tensor, args, device_list[0], runner_cache)
else:
# Multi-GPU or non-cached processing via worker processes
@@ -819,15 +827,11 @@ def _worker_process(
"""
Worker process for multi-GPU upscaling.
Sets up isolated CUDA environment and calls core processing logic.
Results returned via multiprocessing queue as numpy arrays.
CUDA_VISIBLE_DEVICES is set by parent before spawn, so this worker
only sees its assigned GPU. Results returned via queue as numpy arrays.
"""
if platform.system() != "Darwin":
# Limit CUDA visibility to the chosen GPU BEFORE importing torch
os.environ["CUDA_VISIBLE_DEVICES"] = device_id
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")
import torch
# Note: CUDA_VISIBLE_DEVICES and PYTORCH_CUDA_ALLOC_CONF are inherited
# from parent (set before spawn). torch is imported at module level.
# Create debug instance for this worker
worker_debug = Debug(enabled=shared_args["debug"])
@@ -842,7 +846,7 @@ def _worker_process(
result_tensor = _process_frames_core(
frames_tensor=frames_tensor,
args=args,
device_id="0", # Always "0" in worker (CUDA_VISIBLE_DEVICES set)
device_id="0", # Worker sees only 1 GPU (index 0) due to CUDA_VISIBLE_DEVICES
debug=worker_debug,
runner_cache=None # No caching in multiprocessing mode
)
@@ -936,6 +940,9 @@ def _gpu_processing(
# Start all workers
for idx, (device_id, chunk_tensor) in enumerate(zip(device_list, chunks)):
# Set CUDA_VISIBLE_DEVICES before spawning so child inherits it
os.environ["CUDA_VISIBLE_DEVICES"] = device_id
p = mp.Process(
target=_worker_process,
args=(idx, device_id, chunk_tensor.cpu().numpy(), shared_args, return_queue),
@@ -1215,14 +1222,14 @@ def main() -> None:
Raises:
SystemExit: On argument validation failure or processing error
"""
# print header
debug.print_header(cli=True)
# Parse arguments
args = parse_arguments()
# Update debug instance with --debug flag
debug.enabled = args.debug
# print header
debug.print_header(cli=True)
debug.log("Arguments:", category="setup")
for key, value in vars(args).items():
+2 -2
View File
@@ -1,7 +1,7 @@
[project]
name = "seedvr2_videoupscaler"
description = "SeedVR2 official ComfyUI integration: ByteDance-Seed's one-step diffusion-based video/image upscaling with memory-efficient inference"
version = "2.5.8"
version = "2.5.17"
authors = [
{name = "numz"},
{name = "adrientoupet"}
@@ -37,5 +37,5 @@ Forum = "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/discussions"
[tool.comfy]
PublisherId = "ainvfx"
DisplayName = "ComfyUI-SeedVR2_VideoUpscaler"
Icon = "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/blob/main/docs/seedvr_logo.png"
Icon = "https://raw.githubusercontent.com/numz/ComfyUI-SeedVR2_VideoUpscaler/refs/heads/main/docs/seedvr_logo.png"
includes = []
@@ -36,7 +36,7 @@ class UniformTrailingSamplingTimesteps(SamplingTimesteps):
dtype: torch.dtype = torch.float32,
):
# Create trailing timesteps with specified dtype
timesteps = torch.arange(1.0, 0.0, -1.0 / steps, device=device, dtype=dtype)
timesteps = torch.arange(1.0, 0.0, -1.0 / steps, device='cpu').to(device=device, dtype=dtype)
# Shift timesteps.
timesteps = shift * timesteps / (1 + (shift - 1) * timesteps)
+1 -1
View File
@@ -47,7 +47,7 @@ def get_device() -> torch.device:
"""
Get current rank device.
"""
if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cuda", get_local_rank())
+342 -276
View File
@@ -268,6 +268,10 @@ def encode_all_batches(
if step <= 0:
step = batch_size
temporal_overlap = 0
debug.log(f"temporal_overlap >= batch_size, resetting to 0", level="WARNING", category="setup", force=True)
# Store actual temporal overlap used (may differ from parameter if reset)
ctx['actual_temporal_overlap'] = temporal_overlap
# Calculate number of batches
num_encode_batches = 0
@@ -306,6 +310,14 @@ def encode_all_batches(
runner.vae, ctx['cache_context']['vae_model'], debug
)
ctx['cache_context']['vae_newly_cached'] = True
# If both models now cached, cache runner template
dit_is_cached = ctx['cache_context']['cached_dit'] or ctx['cache_context']['dit_newly_cached']
if dit_is_cached:
ctx['cache_context']['global_cache'].set_runner(
ctx['cache_context']['dit_id'], ctx['cache_context']['vae_id'],
runner, debug
)
# Set deterministic seed for VAE encoding (separate from diffusion noise)
# Uses seed + 1,000,000 to avoid collision with upscaling batch seeds
@@ -620,6 +632,7 @@ def upscale_all_batches(
runner.dit, ctx['cache_context']['dit_model'], debug
)
ctx['cache_context']['dit_newly_cached'] = True
# If both models now cached, cache runner template
vae_is_cached = ctx['cache_context']['cached_vae'] or ctx['cache_context']['vae_newly_cached']
if vae_is_cached:
@@ -628,11 +641,6 @@ def upscale_all_batches(
runner, debug
)
# Set base seed for DiT noise generation
# Ensures deterministic noise across all batches in this upscaling phase
set_seed(seed)
debug.log(f"Using seed: {seed}", category="dit")
# Move DiT to GPU for upscaling (no-op if already there)
manage_model_device(model=runner.dit, target_device=ctx['dit_device'],
model_name="DiT", debug=debug, runner=runner)
@@ -646,6 +654,11 @@ def upscale_all_batches(
check_interrupt(ctx)
debug.log(f"Upscaling batch {upscale_idx+1}/{num_valid_latents}", category="generation", force=True)
# Reset seed for each batch to ensure identical RNG state
# This ensures identical inputs produce identical outputs regardless of batch position
set_seed(seed)
debug.log(f"Using seed: {seed} for deterministic generation", category="dit")
debug.start_timer(f"upscale_batch_{upscale_idx+1}")
# Move to DiT device with correct dtype for upscaling (no-op if already there)
@@ -698,7 +711,7 @@ def upscale_all_batches(
debug.start_timer(f"dit_inference_{upscale_idx+1}")
with torch.no_grad():
if dit_dtype != ctx['compute_dtype']:
with torch.autocast(str(ctx['dit_device']), ctx['compute_dtype'], enabled=True):
with torch.autocast(ctx['dit_device'].type, ctx['compute_dtype'], enabled=True):
upscaled_latents = runner.inference(
noises=noises,
conditions=conditions,
@@ -796,7 +809,10 @@ def decode_all_batches(
"""
Phase 3: VAE Decoding.
Decodes all upscaled latents back to pixel space.
Decodes all upscaled latents back to pixel space and writes directly to
pre-allocated final_video tensor. This avoids memory duplication by not
storing intermediate batch_samples.
Requires context from upscale_all_batches with upscaled latents.
Args:
@@ -808,7 +824,8 @@ def decode_all_batches(
Returns:
dict: Updated context containing:
- batch_samples: List of decoded samples ready for post-processing
- final_video: Pre-allocated tensor with decoded samples (unnormalized, in [-1,1])
- decode_batch_info: List of (start_idx, end_idx, ori_length) for Phase 4 processing
- VAE cleanup completed
Raises:
@@ -831,11 +848,36 @@ def decode_all_batches(
# Count valid latents
num_valid_latents = len([l for l in ctx['all_upscaled_latents'] if l is not None])
# Pre-allocate to match original batches (use ori_lengths which is always available)
num_batches = len([l for l in ctx['all_ori_lengths'] if l is not None])
ctx['batch_samples'] = [None] * num_batches
# Get output dimensions from context (set during Phase 1)
if 'true_target_dims' not in ctx:
raise ValueError("true_target_dims not found in context. Run encode_all_batches first.")
true_h, true_w = ctx['true_target_dims']
total_frames = ctx.get('total_frames', 0)
C = 4 if ctx.get('is_rgba', False) else 3
# Pre-allocate final_video at the START of decode phase (before any batch processing)
# This ensures we only need memory for final_video + 1 batch, not final_video + all batch_samples
target_device = ctx['tensor_offload_device'] if ctx['tensor_offload_device'] is not None else 'cpu'
channels_str = "RGBA" if C == 4 else "RGB"
required_gb = (total_frames * true_h * true_w * C * 2) / (1024**3)
debug.log(f"Pre-allocating output tensor: {total_frames} frames, {true_w}x{true_h}px, {channels_str} ({required_gb:.2f}GB)",
category="setup", force=True)
ctx['final_video'] = torch.empty((total_frames, true_h, true_w, C), dtype=ctx['compute_dtype'], device=target_device)
# Track batch write positions for Phase 4 processing
# Each entry: (write_start, write_end, batch_idx, ori_length)
ctx['decode_batch_info'] = []
# Get temporal overlap from context (set during Phase 1)
temporal_overlap = ctx.get('actual_temporal_overlap', 0)
# Track padding removed for final summary
total_padding_removed = 0
current_write_idx = 0
decode_idx = 0
try:
@@ -884,41 +926,96 @@ def decode_all_batches(
samples = runner.vae_decode([upscaled_latent])
debug.end_timer("vae_decode", "VAE decode")
# Process samples
# Process samples - get the single decoded sample
debug.start_timer("optimized_video_rearrange")
samples = optimized_video_rearrange(samples)
debug.end_timer("optimized_video_rearrange", "Video rearrange")
# Convert from VAE dtype to compute dtype and offload to avoid VRAM accumulation
if ctx['tensor_offload_device'] is not None:
# samples is always a single-element list from vae_decode([upscaled_latent])]
if samples[0].is_cuda or samples[0].is_mps:
samples[0] = manage_tensor(
tensor=samples[0],
target_device=ctx['tensor_offload_device'],
tensor_name=f"sample_{decode_idx+1}",
dtype=ctx['compute_dtype'],
debug=debug,
reason="storing decoded samples for post-processing",
indent_level=1
)
else:
# No offload device, but still convert to compute dtype
samples[0] = manage_tensor(
tensor=samples[0],
target_device=samples[0].device,
tensor_name=f"sample_{decode_idx+1}",
dtype=ctx['compute_dtype'],
debug=debug,
reason="VAE dtype → compute dtype",
indent_level=1
)
ctx['batch_samples'][decode_idx] = samples
# Get the decoded sample (always single-element list)
sample = samples[0]
del samples
# Free the upscaled latent and GPU samples
# Get original length for this batch (before any padding was added)
ori_length = ctx['all_ori_lengths'][decode_idx] if decode_idx < len(ctx['all_ori_lengths']) else sample.shape[0]
# Trim temporal padding: sample is in [T, C, H, W] format after rearrange
if ori_length < sample.shape[0]:
padding_removed = sample.shape[0] - ori_length
debug.log(f"Trimming temporal padding: {padding_removed} frames removed ({sample.shape[0]} → {ori_length})",
category="video", indent_level=1)
sample = sample[:ori_length]
total_padding_removed += padding_removed
# Trim spatial padding to true target dimensions
current_h, current_w = sample.shape[-2:]
if current_h != true_h or current_w != true_w:
debug.log(f"Trimming spatial padding: {current_w}x{current_h} → {true_w}x{true_h}",
category="video", indent_level=1)
sample = sample[:, :, :true_h, :true_w]
# Convert to output format: [T, C, H, W] → [T, H, W, C]
# Note: We keep values in [-1, 1] range - normalization happens in Phase 4
sample = optimized_sample_to_image_format(sample) # T, C, H, W → T, H, W, C
# Calculate write position with temporal overlap handling
batch_frames = sample.shape[0]
if decode_idx == 0 or temporal_overlap == 0:
# First batch or no overlap: write all frames
write_start = current_write_idx
write_end = current_write_idx + batch_frames
else:
# Subsequent batches with overlap: blend overlapping region
if temporal_overlap < batch_frames and current_write_idx >= temporal_overlap:
# Blend overlapping region in-place on final_video
prev_tail = ctx['final_video'][current_write_idx - temporal_overlap:current_write_idx]
cur_head = sample[:temporal_overlap]
# Move to same device for blending if needed
if prev_tail.device != cur_head.device:
cur_head = cur_head.to(prev_tail.device)
blended = blend_overlapping_frames(prev_tail, cur_head, temporal_overlap)
ctx['final_video'][current_write_idx - temporal_overlap:current_write_idx] = blended
debug.log(f"Blended {temporal_overlap} overlapping frames at positions {current_write_idx - temporal_overlap}-{current_write_idx}",
category="video", indent_level=1)
# Write only non-overlapping part
sample = sample[temporal_overlap:]
batch_frames = sample.shape[0]
del prev_tail, cur_head, blended
write_start = current_write_idx
write_end = current_write_idx + batch_frames
# Move sample to target device and write directly to final_video
sample = manage_tensor(
tensor=sample,
target_device=target_device,
tensor_name=f"sample_{decode_idx+1}",
dtype=ctx['compute_dtype'],
debug=debug,
reason="writing to final_video",
indent_level=1
)
# Write to final_video - for RGBA, write only RGB channels (VAE outputs 3 channels)
if ctx.get('is_rgba', False):
ctx['final_video'][write_start:write_end, :, :, :3] = sample
else:
ctx['final_video'][write_start:write_end] = sample
# Store batch info for Phase 4 processing
ctx['decode_batch_info'].append((write_start, write_end, decode_idx, ori_length))
current_write_idx = write_end
debug.log(f"Wrote {batch_frames} frames to positions {write_start}-{write_end}",
category="video", indent_level=1)
# Free memory immediately - no batch_samples storage
release_tensor_memory(ctx['all_upscaled_latents'][batch_idx])
ctx['all_upscaled_latents'][batch_idx] = None
del upscaled_latent, samples
del upscaled_latent, sample
debug.end_timer(f"decode_batch_{decode_idx+1}", f"Decoded batch {decode_idx+1}")
@@ -927,6 +1024,9 @@ def decode_all_batches(
1, "Phase 3: Decoding")
decode_idx += 1
# Store padding stats for Phase 4 final summary
ctx['total_padding_removed'] = total_padding_removed
except Exception as e:
debug.log(f"Error in Phase 3 (Decoding): {e}", level="ERROR", category="error", force=True)
@@ -958,14 +1058,12 @@ def postprocess_all_batches(
"""
Phase 4: Post-processing and Final Assembly.
Applies color correction and assembles the final video from decoded batches.
Uses stream-based direct writing to minimize memory usage - processes each
batch and writes directly to pre-allocated output tensor without accumulating
intermediate results. Automatically applies temporal overlap blending and
removes prepended frames if specified.
Processes final_video slices in-place: applies alpha upscaling, color correction,
and normalization. Reads from and writes back to the same final_video tensor
to avoid memory duplication.
Args:
ctx: Context from decode_all_batches containing batch_samples (required)
ctx: Context from decode_all_batches containing final_video (required)
debug: Debug instance for logging (required)
progress_callback: Optional callback(current, total, frames, phase_name)
color_correction: Color correction method - "wavelet", "adain", or "none" (default: "wavelet")
@@ -979,7 +1077,7 @@ def postprocess_all_batches(
- All intermediate storage cleared for memory efficiency
Raises:
ValueError: If context is missing or has no batch samples
ValueError: If context is missing or has no final_video
"""
if debug is None:
raise ValueError("Debug instance must be provided to postprocess_all_batches")
@@ -987,9 +1085,13 @@ def postprocess_all_batches(
if ctx is None:
raise ValueError("Context is required for postprocess_all_batches. Run decode_all_batches first.")
# Validate we have batch samples
if 'batch_samples' not in ctx or not ctx['batch_samples']:
raise ValueError("No batch samples found. Run decode_all_batches first.")
# Validate we have final_video (pre-allocated in decode_all_batches)
if 'final_video' not in ctx or ctx['final_video'] is None:
raise ValueError("final_video not found. Run decode_all_batches first.")
# Validate we have batch info for processing
if 'decode_batch_info' not in ctx or not ctx['decode_batch_info']:
raise ValueError("decode_batch_info not found. Run decode_all_batches first.")
debug.log("", category="none", force=True)
debug.log("━━━━━━━━ Phase 4: Post-processing ━━━━━━━━", category="none", force=True)
@@ -1004,8 +1106,9 @@ def postprocess_all_batches(
debug.log("No frames to process", level="WARNING", category="generation", force=True)
return ctx
# Count valid samples for progress reporting
num_valid_samples = len([s for s in ctx['batch_samples'] if s is not None])
# Get batch info from decode phase
batch_info_list = ctx['decode_batch_info']
num_valid_samples = len(batch_info_list)
# Calculate total post-processing work units
# For RGBA: each batch needs 2 steps (alpha processing + color correction/assembly)
@@ -1022,15 +1125,12 @@ def postprocess_all_batches(
current_postprocessing_step = 0
# Pre-allocation will happen after processing first sample to get exact dimensions
ctx['final_video'] = None
current_frame_idx = 0
# Initialize padding tracking for final log message
total_padding_removed = 0
# Get padding stats from Phase 3
total_padding_removed = ctx.get('total_padding_removed', 0)
# Alpha processing - handle RGBA inputs with edge-guided upscaling
if ctx.get('is_rgba', False) and 'all_alpha_channels' in ctx and 'all_input_rgb' in ctx:
# Process alpha on final_video slices in-place
if has_alpha_processing:
debug.log("Processing Alpha channel with edge-guided upscaling...", category="alpha")
# Validate alpha channel data exists
@@ -1038,10 +1138,7 @@ def postprocess_all_batches(
debug.log("WARNING: Alpha channel data malformed, skipping alpha processing",
level="WARNING", category="alpha", force=True)
else:
for batch_idx in range(len(ctx['batch_samples'])):
if ctx['batch_samples'][batch_idx] is None:
continue
for write_start, write_end, batch_idx, ori_length in batch_info_list:
# Bounds checking for alpha channel lists
if batch_idx >= len(ctx['all_alpha_channels']) or ctx['all_alpha_channels'][batch_idx] is None:
continue
@@ -1055,9 +1152,14 @@ def postprocess_all_batches(
debug.log(f"Processing Alpha batch {batch_idx+1}/{num_valid_samples}", category="alpha", force=True)
debug.start_timer(f"alpha_batch_{batch_idx+1}")
# Get RGB slice from final_video for alpha processing
# final_video is [T, H, W, C], process_alpha_for_batch expects list of [T, C, H, W]
rgb_slice = ctx['final_video'][write_start:write_end, :, :, :3] # Only RGB
rgb_tchw = rgb_slice.permute(0, 3, 1, 2) # [T, H, W, 3] → [T, 3, H, W]
# Process Alpha and merge with RGB
ctx['batch_samples'][batch_idx] = process_alpha_for_batch(
rgb_samples=ctx['batch_samples'][batch_idx],
processed_samples = process_alpha_for_batch(
rgb_samples=[rgb_tchw],
alpha_original=ctx['all_alpha_channels'][batch_idx],
rgb_original=ctx['all_input_rgb'][batch_idx],
device=ctx['vae_device'],
@@ -1065,6 +1167,27 @@ def postprocess_all_batches(
debug=debug
)
# processed_samples[0] is [T, 4, H, W] (RGBA)
# Extract only the alpha channel and write to final_video's alpha slot
processed_rgba = processed_samples[0] # [T, 4, H, W]
alpha_channel = processed_rgba[:, 3:4, :, :] # [T, 1, H, W]
alpha_thwc = alpha_channel.permute(0, 2, 3, 1) # [T, 1, H, W] → [T, H, W, 1]
alpha_thwc = manage_tensor(
tensor=alpha_thwc,
target_device=ctx['final_video'].device,
tensor_name=f"alpha_channel_{batch_idx+1}",
dtype=ctx['compute_dtype'],
debug=debug,
reason="writing alpha channel to final_video",
indent_level=1
)
# Write only the alpha channel to the 4th channel slot
ctx['final_video'][write_start:write_end, :, :, 3:4] = alpha_thwc
del rgb_slice, rgb_tchw, processed_samples, processed_rgba, alpha_channel, alpha_thwc
# Free memory immediately
release_tensor_memory(ctx['all_alpha_channels'][batch_idx])
ctx['all_alpha_channels'][batch_idx] = None
@@ -1083,221 +1206,165 @@ def postprocess_all_batches(
debug.log("Alpha processing complete for all batches", category="alpha")
try:
# Stream-based processing: write directly to final_video without accumulation
for batch_idx, samples in enumerate(ctx['batch_samples']):
if samples is None:
continue
# Process each batch slice in final_video in-place
for info_idx, (write_start, write_end, batch_idx, ori_length) in enumerate(batch_info_list):
check_interrupt(ctx)
debug.log(f"Post-processing batch {batch_idx+1}/{num_valid_samples}", category="video", force=True)
debug.start_timer(f"postprocess_batch_{batch_idx+1}")
debug.log(f"Post-processing batch {info_idx+1}/{num_valid_samples}", category="video", force=True)
debug.start_timer(f"postprocess_batch_{info_idx+1}")
# Post-process each sample in the batch
for i, sample in enumerate(samples):
# Move to VAE device with correct dtype for processing (no-op if already there)
sample = manage_tensor(
tensor=sample,
target_device=ctx['vae_device'],
tensor_name=f"sample_{batch_idx+1}_{i}",
dtype=ctx['compute_dtype'],
debug=debug,
reason="post-processing",
indent_level=1
)
# Get original length for trimming (always available)
video_idx = min(batch_idx, len(ctx['all_ori_lengths']) - 1)
ori_length = ctx['all_ori_lengths'][video_idx] if 'all_ori_lengths' in ctx else sample.shape[0]
# Reconstruct transformed video on-demand for color correction
input_video = None
if color_correction != "none" and ctx.get('batch_metadata') is not None:
if video_idx < len(ctx['batch_metadata']) and ctx['batch_metadata'][video_idx] is not None:
# Reconstruct transformation
transformed_video = _reconstruct_and_transform_batch(ctx, video_idx, debug)
input_video = optimized_single_video_rearrange(transformed_video)
del transformed_video
# Trim both sample and input_video to original length if necessary (handles temporal padding)
if ori_length < sample.shape[0]:
padding_removed = sample.shape[0] - ori_length
debug.log(f"Removing padding: {padding_removed} frame{'s' if padding_removed != 1 else ''} trimmed ({sample.shape[0]} -> {ori_length})",
category="video", force=True, indent_level=1)
sample = sample[:ori_length]
if input_video is not None:
input_video = input_video[:ori_length]
# Get slice from final_video - currently in [T, H, W, C] format, values in [-1, 1]
sample_thwc = ctx['final_video'][write_start:write_end]
# For RGBA, we only process RGB channels for color correction
# Alpha was already written during alpha processing above
if ctx.get('is_rgba', False) and sample_thwc.shape[-1] == 4:
sample_thwc_rgb = sample_thwc[..., :3] # [T, H, W, 3]
sample = sample_thwc_rgb.permute(0, 3, 1, 2) # [T, H, W, 3] → [T, 3, H, W]
else:
sample = sample_thwc.permute(0, 3, 1, 2) # [T, H, W, C] → [T, C, H, W]
# Move to VAE device for processing
sample = manage_tensor(
tensor=sample,
target_device=ctx['vae_device'],
tensor_name=f"sample_{info_idx+1}",
dtype=ctx['compute_dtype'],
debug=debug,
reason="post-processing",
indent_level=1
)
# Reconstruct transformed video on-demand for color correction
input_video = None
if color_correction != "none" and ctx.get('batch_metadata') is not None:
if batch_idx < len(ctx['batch_metadata']) and ctx['batch_metadata'][batch_idx] is not None:
# Reconstruct transformation
transformed_video = _reconstruct_and_transform_batch(ctx, batch_idx, debug)
input_video = optimized_single_video_rearrange(transformed_video)
del transformed_video
# Accumulate padding removed across all batches
total_padding_removed += padding_removed
# Trim spatial dimensions to true target size (removes divisible padding)
if 'true_target_dims' in ctx:
true_h, true_w = ctx['true_target_dims']
current_h, current_w = sample.shape[-2:]
# For batches after the first with temporal overlap, the overlap frames
# were blended in Phase 3 and are not part of this slice. Skip them.
actual_overlap = ctx.get('actual_temporal_overlap', 0)
if info_idx > 0 and actual_overlap > 0:
input_video = input_video[actual_overlap:]
if current_h != true_h or current_w != true_w:
debug.log(f"Trimming spatial padding: {current_w}x{current_h} → {true_w}x{true_h}",
category="video", indent_level=1)
# Trim from bottom-right to match padding location
sample = sample[:, :, :true_h, :true_w]
if input_video is not None:
# Trim input_video to match sample length (handles padding differences)
if input_video.shape[0] > sample.shape[0]:
input_video = input_video[:sample.shape[0]]
# Trim spatial dimensions to true target size
if 'true_target_dims' in ctx:
true_h, true_w = ctx['true_target_dims']
if input_video.shape[-2] != true_h or input_video.shape[-1] != true_w:
input_video = input_video[:, :, :true_h, :true_w]
# Apply color correction if enabled (RGB only)
if color_correction != "none" and input_video is not None:
# Check if RGBA (samples are in T, C, H, W format at this point)
has_alpha = ctx.get('is_rgba', False)
alpha_channel = None
# Apply color correction if enabled (RGB only)
if color_correction != "none" and input_video is not None:
# Check if RGBA (samples are in T, C, H, W format at this point)
has_alpha = ctx.get('is_rgba', False)
alpha_channel = None
if has_alpha:
# Check actual channel count
if sample.shape[1] == 4:
# Extract and temporarily store alpha for reattachment after color correction
alpha_channel = sample[:, 3:4, :, :] # (T, 1, H, W)
sample = sample[:, :3, :, :] # Keep only RGB (T, 3, H, W)
# Ensure both tensors are on same device (GPU) for color correction
if input_video.device != sample.device:
input_video = manage_tensor(
tensor=input_video,
target_device=sample.device,
tensor_name=f"input_video_{batch_idx+1}",
debug=debug,
reason="color correction",
indent_level=1
)
# Apply selected color correction method
debug.start_timer(f"color_correction_{color_correction}")
if color_correction == "lab":
debug.log("Applying LAB perceptual color transfer", category="video", force=True, indent_level=1)
sample = lab_color_transfer(sample, input_video, debug, luminance_weight=0.8)
elif color_correction == "wavelet_adaptive":
debug.log("Applying wavelet with adaptive saturation correction", category="video", force=True, indent_level=1)
sample = wavelet_adaptive_color_correction(sample, input_video, debug)
elif color_correction == "wavelet":
debug.log("Applying wavelet color reconstruction", category="video", force=True, indent_level=1)
sample = wavelet_reconstruction(sample, input_video, debug)
elif color_correction == "hsv":
debug.log("Applying HSV hue-conditional saturation matching", category="video", force=True, indent_level=1)
sample = hsv_saturation_histogram_match(sample, input_video, debug)
elif color_correction == "adain":
debug.log("Applying AdaIN color correction", category="video", force=True, indent_level=1)
sample = adaptive_instance_normalization(sample, input_video)
else:
debug.log(f"Unknown color correction method: {color_correction}", level="WARNING", category="video", force=True, indent_level=1)
debug.end_timer(f"color_correction_{color_correction}", f"Color correction ({color_correction})")
# Free the reconstructed transformed video
del input_video
# Recombine with Alpha if it was present in input
if has_alpha and alpha_channel is not None:
# Concatenate in channels-first: (T, 3, H, W) + (T, 1, H, W) -> (T, 4, H, W)
sample = torch.cat([sample, alpha_channel], dim=1)
if has_alpha:
# Check actual channel count
if sample.shape[1] == 4:
# Extract and temporarily store alpha for reattachment after color correction
alpha_channel = sample[:, 3:4, :, :] # (T, 1, H, W)
sample = sample[:, :3, :, :] # Keep only RGB (T, 3, H, W)
else:
debug.log("Color correction disabled (set to none)", category="video", indent_level=1)
# Free the original length entry
if 'all_ori_lengths' in ctx and video_idx < len(ctx['all_ori_lengths']):
ctx['all_ori_lengths'][video_idx] = None
# Convert to final format (still on GPU at this point)
sample = optimized_sample_to_image_format(sample)
# Apply normalization only to RGB channels, preserve Alpha as-is
if ctx.get('is_rgba', False) and sample.shape[-1] == 4:
# Split RGBA: sample is (T, H, W, C) format after optimized_sample_to_image_format
rgb_channels = sample[..., :3] # (T, H, W, 3)
alpha_channel = sample[..., 3:4] # (T, H, W, 1)
# Normalize only RGB from [-1, 1] to [0, 1]
rgb_channels = rgb_channels.clip(-1, 1).mul_(0.5).add_(0.5)
# Merge back with unchanged Alpha
sample = torch.cat([rgb_channels, alpha_channel], dim=-1)
else:
# RGB only: apply normalization as usual
sample = sample.clip(-1, 1).mul_(0.5).add_(0.5)
# Draw tile boundaries for debugging (if tile info available)
for phase, attr in [('encode', 'encode_tile_boundaries'), ('decode', 'decode_tile_boundaries')]:
tiles = getattr(debug, attr, None)
if tiles:
sample = _draw_tile_boundaries(sample, debug, tiles, phase)
break
# Move to tensor_offload_device if specified
if ctx['tensor_offload_device'] is not None:
sample = manage_tensor(
tensor=sample,
target_device=ctx['tensor_offload_device'],
tensor_name=f"sample_{batch_idx+1}_{i}_final",
# Ensure both tensors are on same device (GPU) for color correction
if input_video.device != sample.device:
input_video = manage_tensor(
tensor=input_video,
target_device=sample.device,
tensor_name=f"input_video_{info_idx+1}",
debug=debug,
reason="storing final processed samples",
reason="color correction",
indent_level=1
)
# Get batch dimensions
batch_frames = sample.shape[0]
# Pre-allocate output tensor on first write
if ctx['final_video'] is None:
H, W = sample.shape[1], sample.shape[2]
# Use input format, not first sample's shape, for consistent allocation
C = 4 if ctx.get('is_rgba', False) else 3
channels_str = "RGBA" if C == 4 else "RGB"
debug.log(f"Pre-allocating output tensor: {total_frames} frames, {W}x{H}px, {channels_str}", category="setup")
target_device = ctx['tensor_offload_device'] if ctx['tensor_offload_device'] is not None else ctx['vae_device']
ctx['final_video'] = torch.empty((total_frames, H, W, C), dtype=ctx['compute_dtype'], device=target_device)
# Write with temporal overlap blending if enabled
if batch_idx == 0 or temporal_overlap == 0:
# First batch or no overlap: write all frames normally
ctx['final_video'][current_frame_idx:current_frame_idx + batch_frames] = sample
current_frame_idx += batch_frames
# Apply selected color correction method
debug.start_timer(f"color_correction_{color_correction}")
if color_correction == "lab":
debug.log("Applying LAB perceptual color transfer", category="video", force=True, indent_level=1)
sample = lab_color_transfer(sample, input_video, debug, luminance_weight=0.8)
elif color_correction == "wavelet_adaptive":
debug.log("Applying wavelet with adaptive saturation correction", category="video", force=True, indent_level=1)
sample = wavelet_adaptive_color_correction(sample, input_video, debug)
elif color_correction == "wavelet":
debug.log("Applying wavelet color reconstruction", category="video", force=True, indent_level=1)
sample = wavelet_reconstruction(sample, input_video, debug)
elif color_correction == "hsv":
debug.log("Applying HSV hue-conditional saturation matching", category="video", force=True, indent_level=1)
sample = hsv_saturation_histogram_match(sample, input_video, debug)
elif color_correction == "adain":
debug.log("Applying AdaIN color correction", category="video", force=True, indent_level=1)
sample = adaptive_instance_normalization(sample, input_video)
else:
# Subsequent batches with overlap: blend overlapping region in-place
if temporal_overlap < batch_frames and current_frame_idx >= temporal_overlap:
# Get overlapping regions
prev_tail = ctx['final_video'][current_frame_idx - temporal_overlap:current_frame_idx]
cur_head = sample[:temporal_overlap]
# Log blending operation
blend_range_start = current_frame_idx - temporal_overlap
blend_range_end = current_frame_idx
debug.log(f"Blending {temporal_overlap} overlapping frames (positions {blend_range_start}-{blend_range_end})",
category="video", indent_level=1)
# Blend and replace in-place
blended = blend_overlapping_frames(prev_tail, cur_head, temporal_overlap)
ctx['final_video'][current_frame_idx - temporal_overlap:current_frame_idx] = blended
# Write non-overlapping part
non_overlapping = sample[temporal_overlap:]
non_overlapping_count = non_overlapping.shape[0]
ctx['final_video'][current_frame_idx:current_frame_idx + non_overlapping_count] = non_overlapping
current_frame_idx += non_overlapping_count
else:
# Edge case: write normally if overlap >= batch_frames
ctx['final_video'][current_frame_idx:current_frame_idx + batch_frames] = sample
current_frame_idx += batch_frames
debug.log(f"Unknown color correction method: {color_correction}", level="WARNING", category="video", force=True, indent_level=1)
# Immediately release sample memory
del sample
debug.end_timer(f"color_correction_{color_correction}", f"Color correction ({color_correction})")
# Clear batch samples as we go to free memory progressively
release_tensor_memory(ctx['batch_samples'][batch_idx])
ctx['batch_samples'][batch_idx] = None
# Free the reconstructed transformed video
del input_video
# Recombine with Alpha if it was present in input
if has_alpha and alpha_channel is not None:
# Concatenate in channels-first: (T, 3, H, W) + (T, 1, H, W) -> (T, 4, H, W)
sample = torch.cat([sample, alpha_channel], dim=1)
debug.end_timer(f"postprocess_batch_{batch_idx+1}", f"Post-processed batch {batch_idx+1}")
else:
debug.log("Color correction disabled (set to none)", category="video", indent_level=1)
# Convert to final format: [T, C, H, W] → [T, H, W, C]
sample = optimized_sample_to_image_format(sample)
# Apply normalization only to RGB channels, preserve Alpha as-is
if ctx.get('is_rgba', False) and sample.shape[-1] == 4:
# Split RGBA: sample is (T, H, W, C) format after optimized_sample_to_image_format
rgb_channels = sample[..., :3] # (T, H, W, 3)
alpha_channel = sample[..., 3:4] # (T, H, W, 1)
# Normalize only RGB from [-1, 1] to [0, 1]
rgb_channels.clamp_(-1, 1).mul_(0.5).add_(0.5)
# Merge back with unchanged Alpha
sample = torch.cat([rgb_channels, alpha_channel], dim=-1)
else:
# RGB only: apply normalization as usual
sample.clamp_(-1, 1).mul_(0.5).add_(0.5)
# Draw tile boundaries for debugging (if tile info available)
for phase, attr in [('encode', 'encode_tile_boundaries'), ('decode', 'decode_tile_boundaries')]:
tiles = getattr(debug, attr, None)
if tiles:
sample = _draw_tile_boundaries(sample, debug, tiles, phase)
break
# Move to final_video device and write back in-place
sample = manage_tensor(
tensor=sample,
target_device=ctx['final_video'].device,
tensor_name=f"sample_{info_idx+1}_final",
dtype=ctx['compute_dtype'],
debug=debug,
reason="writing processed result to final_video",
indent_level=1
)
# Write back to final_video in-place
# For RGBA, write only RGB channels (alpha already written during alpha processing)
if ctx.get('is_rgba', False) and ctx['final_video'].shape[-1] == 4:
ctx['final_video'][write_start:write_end, :, :, :3] = sample
else:
ctx['final_video'][write_start:write_end] = sample
# Free sample memory
del sample, sample_thwc
debug.end_timer(f"postprocess_batch_{info_idx+1}", f"Post-processed batch {info_idx+1}")
# Update progress for main processing step
current_postprocessing_step += 1
@@ -1332,24 +1399,23 @@ def postprocess_all_batches(
if total_padding_removed > 0:
adjustments.append(f"{total_padding_removed} padding")
# Use actual temporal overlap from encoding (may have been reset)
actual_overlap = ctx.get('actual_temporal_overlap', temporal_overlap)
# Calculate and include temporal overlap blending info
if temporal_overlap > 0:
frames_blended = (num_valid_samples - 1) * temporal_overlap
if actual_overlap > 0:
frames_blended = (num_valid_samples - 1) * actual_overlap
adjustments.append(f"{frames_blended} overlap")
if adjustments:
# Add back all removed/blended frames to get true computed count
total_computed = frames_before_removal + total_padding_removed
if temporal_overlap > 0:
total_computed += (num_valid_samples - 1) * temporal_overlap
if actual_overlap > 0:
total_computed += (num_valid_samples - 1) * actual_overlap
frame_info += f" ({total_computed} computed with {' + '.join(adjustments)} removed)"
debug.log(f"Final output assembled: {frame_info}, Resolution: {Wf}x{Hf}px, Channels: {channels_str}",
category="generation", force=True)
if current_frame_idx != total_frames:
debug.log(f"WARNING: Frame count mismatch - expected {total_frames}, wrote {current_frame_idx}",
level="WARNING", category="generation", force=True)
else:
ctx['final_video'] = torch.empty((0, 0, 0, 0), dtype=ctx['compute_dtype'])
debug.log("No frames were processed", level="WARNING", category="generation", force=True)
@@ -1358,11 +1424,11 @@ def postprocess_all_batches(
debug.log(f"Error in Phase 4 (Post-processing): {e}", level="ERROR", category="generation", force=True)
raise
finally:
# 1. Clean up batch_samples from context (already mostly freed during processing)
if 'batch_samples' in ctx and ctx['batch_samples']:
release_tensor_collection(ctx['batch_samples'])
ctx['batch_samples'].clear()
del ctx['batch_samples']
# 1. Clean up decode_batch_info and padding stats
if 'decode_batch_info' in ctx:
del ctx['decode_batch_info']
if 'total_padding_removed' in ctx:
del ctx['total_padding_removed']
# 2. Clean up video transform caches
if 'video_transform' in ctx and ctx['video_transform'] is not None:
+9 -3
View File
@@ -36,6 +36,7 @@ from .model_configuration import configure_runner
from .infer import VideoDiffusionInfer
from ..data.image.transforms.divisible_crop import DivisiblePad
from ..data.image.transforms.na_resize import NaResize
from ..optimization.compatibility import COMPUTE_DTYPE, BFLOAT16_SUPPORTED
from ..optimization.memory_manager import manage_tensor
from ..utils.constants import get_script_directory
@@ -371,7 +372,7 @@ def setup_generation_context(
'dit_offload_device': dit_offload_device,
'vae_offload_device': vae_offload_device,
'tensor_offload_device': tensor_offload_device,
'compute_dtype': torch.bfloat16, # Hardcoded - gives the best compromise between memory & quality without artifacts
'compute_dtype': COMPUTE_DTYPE,
'interrupt_fn': interrupt_fn,
'video_transform': None,
'text_embeds': None,
@@ -401,7 +402,12 @@ def setup_generation_context(
f"LOCAL_RANK={os.environ['LOCAL_RANK']}",
category="setup"
)
reason = "quality" if ctx['compute_dtype'] == torch.float32 else "compatibility"
if ctx['compute_dtype'] == torch.float32:
reason = "quality"
elif not BFLOAT16_SUPPORTED:
reason = "compatibility (GPU lacks bfloat16 CUBLAS - 7B models unsupported, 3B may have artifacts)"
else:
reason = "performance"
debug.log(f"Unified compute dtype: {ctx['compute_dtype']} across entire pipeline for maximum {reason}", category="precision")
return ctx
@@ -722,7 +728,7 @@ def _draw_tile_boundaries(image: torch.Tensor, debug: 'Debug', tile_boundaries:
annotated_frames = []
for frame_idx in range(T):
# Convert frame to numpy (handle RGB and RGBA)
img = (image[frame_idx].float().cpu().numpy() * 255).astype(np.uint8) # [H, W, C]
img = np.ascontiguousarray((image[frame_idx].float().cpu().numpy() * 255).astype(np.uint8)) # [H, W, C]
# Draw boundary lines inside each tile
for idx, tile_info in enumerate(tile_boundaries):
+8 -36
View File
@@ -26,41 +26,13 @@ from ..common.diffusion import (
from ..common.distributed import (
get_device,
)
from ..optimization.performance import (
optimized_channels_to_last,
optimized_channels_to_second
)
from ..models.dit_3b import na
def optimized_channels_to_last(tensor):
"""🚀 Optimized replacement for rearrange(tensor, 'b c ... -> b ... c')
Moves channels from position 1 to last position using PyTorch native operations.
"""
if tensor.ndim == 3: # [batch, channels, spatial]
return tensor.permute(0, 2, 1)
elif tensor.ndim == 4: # [batch, channels, height, width]
return tensor.permute(0, 2, 3, 1)
elif tensor.ndim == 5: # [batch, channels, depth, height, width]
return tensor.permute(0, 2, 3, 4, 1)
else:
# Fallback for other dimensions - move channel (dim=1) to last
dims = list(range(tensor.ndim))
dims = [dims[0]] + dims[2:] + [dims[1]] # [0, 2, 3, ..., 1]
return tensor.permute(*dims)
def optimized_channels_to_second(tensor):
"""🚀 Optimized replacement for rearrange(tensor, 'b ... c -> b c ...')
Moves channels from last position to position 1 using PyTorch native operations.
"""
if tensor.ndim == 3: # [batch, spatial, channels]
return tensor.permute(0, 2, 1)
elif tensor.ndim == 4: # [batch, height, width, channels]
return tensor.permute(0, 3, 1, 2)
elif tensor.ndim == 5: # [batch, depth, height, width, channels]
return tensor.permute(0, 4, 1, 2, 3)
else:
# Fallback for other dimensions - move last dim to position 1
dims = list(range(tensor.ndim))
dims = [dims[0], dims[-1]] + dims[1:-1] # [0, -1, 1, 2, ..., -2]
return tensor.permute(*dims)
class VideoDiffusionInfer():
def __init__(self, config: DictConfig, debug: 'Debug',
encode_tiled: bool = False, encode_tile_size: Tuple[int, int] = (512, 512),
@@ -183,7 +155,7 @@ class VideoDiffusionInfer():
# Use autocast if VAE dtype differs from input dtype
if vae_dtype != sample.dtype:
with torch.autocast(str(device), sample.dtype, enabled=True):
with torch.autocast(device.type, sample.dtype, enabled=True):
if use_sample:
latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size,
tile_overlap=self.encode_tile_overlap).latent
@@ -201,7 +173,7 @@ class VideoDiffusionInfer():
tile_overlap=self.encode_tile_overlap).posterior.mode().squeeze(2)
latent = latent.unsqueeze(2) if latent.ndim == 4 else latent
latent = rearrange(latent, "b c ... -> b ... c")
latent = optimized_channels_to_last(latent)
latent = (latent - shift) * scale
latents.append(latent)
@@ -248,7 +220,7 @@ class VideoDiffusionInfer():
for i, latent in enumerate(latents):
latent = latent / scale + shift
latent = rearrange(latent, "b ... c -> b c ...")
latent = optimized_channels_to_second(latent)
latent = latent.squeeze(2)
# Detect VAE model dtype
@@ -259,7 +231,7 @@ class VideoDiffusionInfer():
# Use autocast if VAE dtype differs from latent dtype
if vae_dtype != latent.dtype:
with torch.autocast(str(device), latent.dtype, enabled=True):
with torch.autocast(device.type, latent.dtype, enabled=True):
sample = self.vae.decode(
latent,
tiled=self.decode_tiled, tile_size=self.decode_tile_size,
+22 -1
View File
@@ -115,7 +115,28 @@ def load_quantized_state_dict(checkpoint_path: str, device: torch.device = torch
debug.log("This is a one-time installation that will enable loading of .safetensors files",
level="INFO", category="info", force=True)
raise ImportError(error_msg)
state = load_safetensors_file(checkpoint_path, device=device_str)
# Try direct device loading first (optimal path)
try:
state = load_safetensors_file(checkpoint_path, device=device_str)
except RuntimeError as e:
# MPS allocator fallback: some PyTorch/macOS versions have issues with
# direct MPS loading (allocation failures, watermark errors, etc.)
error_msg = str(e).lower()
is_mps_alloc_error = device.type == "mps" and any(
keyword in error_msg for keyword in ["watermark", "allocat", "memory"]
)
if is_mps_alloc_error:
# Transparent fallback - only log if debug enabled
if debug:
debug.log("Using CPU intermediate loading for MPS compatibility",
category="info", indent_level=1)
state = load_safetensors_file(checkpoint_path, device="cpu")
# Tensors will be moved to MPS during model.load_state_dict()
else:
# Re-raise if it's a different error (file corruption, etc.)
raise
elif checkpoint_path.endswith('.gguf'):
validate_gguf_availability(f"load {os.path.basename(checkpoint_path)}", debug)
state = _load_gguf_state(
+3 -1
View File
@@ -31,7 +31,7 @@ class AreaResize:
self.max_area = max_area
self.downsample_only = downsample_only
self.interpolation = interpolation
if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
self.interpolation = InterpolationMode.BILINEAR
def __call__(self, image: Union[torch.Tensor, Image.Image]):
@@ -50,10 +50,12 @@ class AreaResize:
resized_height, resized_width = round(height * scale), round(width * scale)
antialias = not (isinstance(image, torch.Tensor) and image.device.type == 'mps')
return TVF.resize(
image,
size=(resized_height, resized_width),
interpolation=self.interpolation,
antialias=antialias,
)
+1 -1
View File
@@ -26,7 +26,7 @@ def NaResize(
max_resolution: int = 0,
interpolation: InterpolationMode = InterpolationMode.BICUBIC,
):
Interpolation = InterpolationMode.BILINEAR if (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()) else interpolation
Interpolation = InterpolationMode.BILINEAR if (hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()) else interpolation
if mode == "area":
return AreaResize(
max_area=resolution**2,
+5 -4
View File
@@ -30,7 +30,7 @@ class SideResize:
self.max_size = max_size
self.downsample_only = downsample_only
self.interpolation = interpolation
if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
self.interpolation = InterpolationMode.BILINEAR
def __call__(self, image: Union[torch.Tensor, Image.Image]):
@@ -56,8 +56,9 @@ class SideResize:
else:
size = self.size
# Resize to shortest edge
resized = TVF.resize(image, size, self.interpolation)
# Resize to shortest edge (disable antialias only for MPS tensors - not supported)
antialias = not (isinstance(image, torch.Tensor) and image.device.type == 'mps')
resized = TVF.resize(image, size, self.interpolation, antialias=antialias)
# Apply max_size constraint if specified
if self.max_size > 0:
@@ -69,6 +70,6 @@ class SideResize:
if max(h, w) > self.max_size:
scale = self.max_size / max(h, w)
new_h, new_w = round(h * scale), round(w * scale)
resized = TVF.resize(resized, (new_h, new_w), self.interpolation)
resized = TVF.resize(resized, (new_h, new_w), self.interpolation, antialias=antialias)
return resized
+12 -6
View File
@@ -6,7 +6,7 @@ Main ComfyUI node for high-quality video upscaling using diffusion models
import torch
from comfy_api.latest import io
from typing import Tuple, Dict, Any, Optional
from ..utils.constants import get_base_cache_dir
from ..utils.constants import get_base_cache_dir, __version__
from ..utils.downloads import download_weight
from ..utils.debug import Debug
from ..core.generation_phases import (
@@ -50,7 +50,7 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
def define_schema(cls) -> io.Schema:
return io.Schema(
node_id="SeedVR2VideoUpscaler",
display_name="SeedVR2 Video Upscaler",
display_name=f"SeedVR2 Video Upscaler (v{__version__})",
category="SEEDVR2",
description=(
"SeedVR2 main upscaling node: processes video frames using DiT and VAE models with diffusion-based enhancement. "
@@ -509,15 +509,21 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
)
sample = ctx['final_video']
debug.log("", category="none", force=True)
# Ensure CPU tensor in float32 for maximum ComfyUI compatibility
if torch.is_tensor(sample):
if sample.is_cuda or sample.is_mps:
sample = sample.cpu()
if sample.dtype != torch.float32:
sample = sample.to(torch.float32)
src_dtype = sample.dtype
try:
sample = sample.to(torch.float32)
debug.log(f"Converted output from {src_dtype} to float32", category="precision")
except Exception as e:
debug.log(f"Could not convert to float32: {e}. Output is {src_dtype}, compatibility with other nodes not guaranteed",
level="WARNING", category="precision", force=True)
debug.log("", category="none", force=True)
debug.log("Upscaling completed successfully!", category="success", force=True)
debug.end_timer("generation", "Video generation")
@@ -529,7 +535,7 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
debug.log_memory_state("After all phases complete", show_tensors=False, detailed_tensors=False)
# Final peak vram summary
debug.log_peak_vram_summary()
debug.log_peak_memory_summary()
# Final timing summary
debug.log("", category="none")
@@ -1308,7 +1308,8 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
if H <= tile_h and W <= tile_w:
return self.slicing_encode(x)
else:
self.debug.log(f"Using VAE tiled encoding (Tile: {tile_size}, Overlap: {tile_overlap})", category="vae", force=True, indent_level=1)
if self.debug:
self.debug.log(f"Using VAE tiled encoding (Tile: {tile_size}, Overlap: {tile_overlap})", category="vae", force=True, indent_level=1)
# Spatial scale factor (output/latent)
scale_factor = self.spatial_downsample_factor
@@ -1481,7 +1482,8 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
if H <= latent_tile_h and W <= latent_tile_w:
return self.slicing_decode(z)
else:
self.debug.log(f"Using VAE tiled decoding (Tile: {tile_size}, Overlap: {tile_overlap})", category="vae", force=True, indent_level=1)
if self.debug:
self.debug.log(f"Using VAE tiled decoding (Tile: {tile_size}, Overlap: {tile_overlap})", category="vae", force=True, indent_level=1)
latent_overlap_h = max(0, min((overlap_h // scale_factor), latent_tile_h - 1))
latent_overlap_w = max(0, min((overlap_w // scale_factor), latent_tile_w - 1))
+62 -40
View File
@@ -144,22 +144,21 @@ def apply_block_swap_to_dit(
total_blocks = len(model.blocks)
# Log configuration clearly based on what's enabled
block_text = "block" if blocks_to_swap <= 1 else "blocks"
if blocks_to_swap > 0 and swap_io_components:
debug.log(f"BlockSwap: {blocks_to_swap} transformer {block_text} + I/O components offloaded to {str(offload_device).upper()}", category="blockswap", force=True)
elif blocks_to_swap > 0:
debug.log(f"BlockSwap: {blocks_to_swap} transformer {block_text} offloaded to {str(offload_device).upper()}", category="blockswap", force=True)
elif swap_io_components:
debug.log(f"BlockSwap: I/O components offloaded to {str(offload_device).upper()} (blocks remain on GPU)", category="blockswap", force=True)
# Clamp blocks_to_swap to available blocks BEFORE logging
effective_blocks = min(blocks_to_swap, total_blocks) if blocks_to_swap > 0 else 0
debug.log(f"Model has {total_blocks} transformer blocks", category="blockswap")
# Log configuration clearly based on what's enabled
block_text = "block" if effective_blocks <= 1 else "blocks"
if effective_blocks > 0 and swap_io_components:
debug.log(f"BlockSwap: {effective_blocks}/{total_blocks} transformer {block_text} + I/O components offloaded to {str(offload_device).upper()}", category="blockswap", force=True)
elif effective_blocks > 0:
debug.log(f"BlockSwap: {effective_blocks}/{total_blocks} transformer {block_text} offloaded to {str(offload_device).upper()}", category="blockswap", force=True)
elif swap_io_components:
debug.log(f"BlockSwap: I/O components offloaded to {str(offload_device).upper()} (0/{total_blocks} blocks swapped)", category="blockswap", force=True)
# Configure model with blockswap attributes
if blocks_to_swap > 0:
blocks_to_swap = min(blocks_to_swap, total_blocks)
model.blocks_to_swap = blocks_to_swap - 1 # Convert to 0-indexed
debug.log(f"Transformer blocks to swap: {blocks_to_swap}/{total_blocks}", category="blockswap")
model.blocks_to_swap = effective_blocks - 1 # Convert to 0-indexed
else:
# No block swapping, set to -1 so no blocks match the swap condition
model.blocks_to_swap = -1
@@ -193,7 +192,7 @@ def apply_block_swap_to_dit(
runner._blockswap_active = True
# Store configuration for debugging and cleanup
runner._block_swap_config = {
model._block_swap_config = {
"blocks_swapped": blocks_to_swap,
"swap_io_components": swap_io_components,
"total_blocks": total_blocks,
@@ -650,11 +649,11 @@ def _protect_model_from_move(
Wraps model.to() method to prevent other code from accidentally moving
the entire model to GPU, which would defeat BlockSwap's memory savings.
Allows movement only when explicitly bypassed via runner flag.
Allows movement only when explicitly bypassed via model flag.
Args:
model: DiT model to protect
runner: VideoDiffusionInfer instance (stored as weak reference)
runner: VideoDiffusionInfer instance (for active status check)
debug: Debug instance for logging (required)
"""
if not hasattr(model, '_original_to'):
@@ -665,34 +664,46 @@ def _protect_model_from_move(
# Define the protected method without closures
def protected_model_to(self, device, *args, **kwargs):
# Check if protection is temporarily bypassed for offloading
# Flag is stored on model itself (not runner) to survive runner recreation
if getattr(self, "_blockswap_bypass_protection", False):
# Protection bypassed, allow movement
if hasattr(self, '_original_to'):
return self._original_to(device, *args, **kwargs)
# Get configured offload device directly from model
blockswap_offload_device = "cpu" # default
if hasattr(self, "_block_swap_config"):
blockswap_offload_device = self._block_swap_config.get("offload_device", "cpu")
# Check if BlockSwap is currently active via runner weak reference
runner_ref = getattr(self, '_blockswap_runner_ref', None)
blockswap_is_active = False
if runner_ref:
runner_obj = runner_ref()
if runner_obj and getattr(runner_obj, "_blockswap_bypass_protection", False):
# Protection bypassed, allow movement
if hasattr(self, '_original_to'):
return self._original_to(device, *args, **kwargs)
if runner_obj and hasattr(runner_obj, "_blockswap_active"):
blockswap_is_active = runner_obj._blockswap_active
# Check blockswap status using weak reference
# Get configured offload device from runner
blockswap_offload_device = "cpu" # default
if runner_ref:
runner_obj = runner_ref()
if runner_obj and hasattr(runner_obj, "_block_swap_config"):
blockswap_offload_device = runner_obj._block_swap_config.get("offload_device", "cpu")
# Block attempts to move model away from configured offload device when active
if blockswap_is_active and str(device) != str(blockswap_offload_device):
# Get debug instance from runner if available
debug_instance = None
if runner_ref:
runner_obj = runner_ref()
if runner_obj and hasattr(runner_obj, 'debug'):
debug_instance = runner_obj.debug
# Block attempts to move model away from configured offload device
if str(device) != str(blockswap_offload_device):
if runner_obj and hasattr(runner_obj, "_blockswap_active") and runner_obj._blockswap_active:
debug.log(f"Blocked attempt to move blockswapped model from {blockswap_offload_device} to {device}",
level="WARNING", category="blockswap", force=True)
return self
if debug_instance:
debug_instance.log(
f"Blocked attempt to move BlockSwap model from {blockswap_offload_device} to {device}",
level="WARNING", category="blockswap", force=True
)
return self
# Use original method stored as attribute
# Allow movement (either bypass is enabled or target is offload device)
if hasattr(self, '_original_to'):
return self._original_to(device, *args, **kwargs)
else:
# This shouldn't happen, but fallback to super().to()
# Fallback - shouldn't happen
return super(type(self), self).to(device, *args, **kwargs)
# Bind as a method to the model instance
@@ -712,7 +723,13 @@ def set_blockswap_bypass(runner, bypass: bool, debug):
if not hasattr(runner, "_blockswap_active") or not runner._blockswap_active:
return
runner._blockswap_bypass_protection = bypass
# Get the actual model (handle FP8CompatibleDiT wrapper)
model = runner.dit
if hasattr(model, "dit_model"):
model = model.dit_model
# Store on model so it survives runner recreation during caching
model._blockswap_bypass_protection = bypass
if bypass:
debug.log("BlockSwap protection disabled to allow model DiT offloading", category="success")
@@ -741,11 +758,16 @@ def cleanup_blockswap(runner, keep_state_for_cache=False):
debug = runner.debug
# Check if there's any BlockSwap state to clean up
# Get the actual model (handle FP8CompatibleDiT wrapper)
model = runner.dit
if hasattr(model, "dit_model"):
model = model.dit_model
# Check if there's any BlockSwap state to clean up (check both runner and model)
has_blockswap_state = (
hasattr(runner, "_blockswap_active") or
hasattr(runner, "_block_swap_config") or
hasattr(runner, "_blockswap_bypass_protection")
hasattr(model, "_block_swap_config") or
hasattr(model, "_blockswap_bypass_protection")
)
if not has_blockswap_state:
@@ -757,7 +779,7 @@ def cleanup_blockswap(runner, keep_state_for_cache=False):
# Minimal cleanup for caching - just mark as inactive and allow offloading
# Everything else stays intact for fast reactivation
if hasattr(runner, "_blockswap_active") and runner._blockswap_active:
if not getattr(runner, "_blockswap_bypass_protection", False):
if not getattr(model, "_blockswap_bypass_protection", False):
set_blockswap_bypass(runner=runner, bypass=True, debug=debug)
runner._blockswap_active = False
debug.log("BlockSwap deactivated for caching (configuration preserved)", category="success")
@@ -829,7 +851,7 @@ def cleanup_blockswap(runner, keep_state_for_cache=False):
# 5. Clean up BlockSwap-specific attributes
for attr in ['_blockswap_runner_ref', 'blocks_to_swap', 'main_device',
'offload_device', '_blockswap_configured']:
'offload_device']:
if hasattr(model, attr):
delattr(model, attr)
+63 -8
View File
@@ -5,10 +5,42 @@ Contains FP8/FP16 compatibility layers and wrappers for different model architec
Extracted from: seedvr2.py (lines 1045-1630)
"""
# Triton compatibility shim for bitsandbytes 0.45+ with triton 3.0+
# Must be called before any diffusers import
import sys
def ensure_triton_compat():
"""Create minimal triton.ops stubs only if missing, to allow bitsandbytes import."""
if 'triton.ops.matmul_perf_model' in sys.modules:
return
try:
from triton.ops.matmul_perf_model import early_config_prune # noqa: F401
return
except (ImportError, ModuleNotFoundError, AttributeError):
pass
import types
if 'triton.ops' not in sys.modules:
sys.modules['triton.ops'] = types.ModuleType('triton.ops')
matmul_perf = types.ModuleType('triton.ops.matmul_perf_model')
matmul_perf.early_config_prune = lambda configs, *a, **kw: configs
matmul_perf.estimate_matmul_time = lambda *a, **kw: 0.0
sys.modules['triton.ops'].matmul_perf_model = matmul_perf
sys.modules['triton.ops.matmul_perf_model'] = matmul_perf
# Run immediately on import
ensure_triton_compat()
import torch
import types
import os
# Flash Attention & Triton Compatibility Layer
# 1. Flash Attention - speedup for attention operations
try:
@@ -171,6 +203,24 @@ if not os.environ.get("SEEDVR2_OPTIMIZATIONS_LOGGED"):
print(f"🔧 Conv3d workaround active: PyTorch {torch_ver}, cuDNN {cudnn_ver} (fixing VAE 3x memory bug)")
# Bfloat16 CUBLAS support
def _probe_bfloat16_support() -> bool:
if not torch.cuda.is_available():
return True
try:
a = torch.randn(8, 8, dtype=torch.bfloat16, device='cuda:0')
_ = torch.matmul(a, a)
del a
return True
except RuntimeError as e:
if "CUBLAS_STATUS_NOT_SUPPORTED" in str(e):
return False
raise
BFLOAT16_SUPPORTED = _probe_bfloat16_support()
COMPUTE_DTYPE = torch.bfloat16 if BFLOAT16_SUPPORTED else torch.float16
def call_rope_with_stability(method, *args, **kwargs):
"""
Call RoPE method with stability fixes:
@@ -221,7 +271,7 @@ class FP8CompatibleDiT(torch.nn.Module):
self._convert_rope_freqs(target_dtype=self.compute_dtype)
self.debug.end_timer("_convert_rope_freqs", "RoPE freqs conversion")
if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
self.debug.log(f"Also converting NaDiT parameters/buffers for MPS backend", category="setup", force=True)
self.debug.start_timer("_force_nadit_precision")
self._force_nadit_precision(target_dtype=self.compute_dtype)
@@ -524,19 +574,24 @@ class FP8CompatibleDiT(torch.nn.Module):
k = k.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
v = v.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
attn_output = torch.nn.functional.scaled_dot_product_attention(
q, k, v,
dropout_p=0.0,
is_causal=False
)
else:
# Use optimized SDPA
with torch.backends.cuda.sdp_kernel(
enable_flash=True,
enable_math=True,
enable_mem_efficient=True
):
# Use optimized SDPA - PyTorch 2.3+ API with CUDNN support, fallback for older versions
if hasattr(torch.nn.attention, 'sdpa_kernel'):
ctx = torch.nn.attention.sdpa_kernel([
torch.nn.attention.SDPBackend.FLASH_ATTENTION,
torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION,
torch.nn.attention.SDPBackend.CUDNN_ATTENTION,
torch.nn.attention.SDPBackend.MATH])
else:
ctx = torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True)
with ctx:
attn_output = torch.nn.functional.scaled_dot_product_attention(
q, k, v,
dropout_p=0.0,
+54 -32
View File
@@ -13,6 +13,12 @@ import psutil
from typing import Tuple, Dict, Any, Optional, List, Union
def _device_str(device: Union[torch.device, str]) -> str:
"""Normalized uppercase device string for comparison and logging. MPS variants → 'MPS'."""
s = str(device).upper()
return 'MPS' if s.startswith('MPS') else s
def get_device_list(include_none: bool = False, include_cpu: bool = False) -> List[str]:
"""
Get list of available compute devices for SeedVR2
@@ -38,7 +44,7 @@ def get_device_list(include_none: bool = False, include_cpu: bool = False) -> Li
pass
try:
if hasattr(torch, "mps") and hasattr(torch.mps, "is_available") and torch.mps.is_available():
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
devs.append("mps") # MPS doesn't use device indices
has_mps = True
except Exception:
@@ -80,7 +86,7 @@ def get_basic_vram_info(device: Optional[torch.device] = None) -> Dict[str, Any]
elif not isinstance(device, torch.device):
device = torch.device(device)
free_memory, total_memory = torch.cuda.mem_get_info(device)
elif hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
# MPS doesn't support per-device queries or mem_get_info
# Use system memory as proxy
mem = psutil.virtual_memory()
@@ -100,12 +106,28 @@ def get_basic_vram_info(device: Optional[torch.device] = None) -> Dict[str, Any]
# Initial VRAM check at module load
vram_info = get_basic_vram_info(device=None)
if "error" not in vram_info:
backend = "MPS" if (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()) else "CUDA"
backend = "MPS" if (hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()) else "CUDA"
print(f"📊 Initial {backend} memory: {vram_info['free_gb']:.2f}GB free / {vram_info['total_gb']:.2f}GB total")
else:
print(f"⚠️ Memory check failed: {vram_info['error']} - No available backend!")
def _enforce_vram_limit() -> None:
"""
Enforce VRAM limit to physical capacity to prevent silent swap to system RAM.
Called once at module load. No-op on MPS or unsupported platforms.
"""
if not torch.cuda.is_available():
return
try:
for i in range(torch.cuda.device_count()):
torch.cuda.set_per_process_memory_fraction(1.0, i)
except Exception:
pass
_enforce_vram_limit()
def get_vram_usage(device: Optional[torch.device] = None, debug: Optional['Debug'] = None) -> Tuple[float, float, float]:
"""
Get current VRAM usage metrics for monitoring.
@@ -116,7 +138,7 @@ def get_vram_usage(device: Optional[torch.device] = None, debug: Optional['Debug
debug: Optional debug instance for logging
Returns:
tuple: (allocated_gb, reserved_gb, max_allocated_gb)
tuple: (allocated_gb, reserved_gb, max_reserved_gb)
Returns (0, 0, 0) if no GPU available
"""
try:
@@ -127,9 +149,9 @@ def get_vram_usage(device: Optional[torch.device] = None, debug: Optional['Debug
device = torch.device(device)
allocated = torch.cuda.memory_allocated(device) / (1024**3)
reserved = torch.cuda.memory_reserved(device) / (1024**3)
max_allocated = torch.cuda.max_memory_allocated(device) / (1024**3)
return allocated, reserved, max_allocated
elif hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
max_reserved = torch.cuda.max_memory_reserved(device) / (1024**3)
return allocated, reserved, max_reserved
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
# MPS doesn't support per-device queries - uses global memory tracking
allocated = torch.mps.current_allocated_memory() / (1024**3)
reserved = torch.mps.driver_allocated_memory() / (1024**3)
@@ -235,11 +257,11 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo
if free_ratio < 0.05:
should_clear = True
if debug:
backend = "MPS" if (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()) else "VRAM"
backend = "MPS" if (hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()) else "VRAM"
debug.log(f"{backend} pressure: {mem_info['free_gb']:.2f}GB free of {mem_info['total_gb']:.2f}GB", category="memory")
# For non-MPS systems, also check system RAM separately
if not should_clear and not (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()):
if not should_clear and not (hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()):
mem = psutil.virtual_memory()
if mem.available < mem.total * 0.05:
should_clear = True
@@ -265,7 +287,7 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
elif hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
torch.mps.empty_cache()
if debug:
@@ -302,7 +324,7 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo
handle = _os_memory_lib.GetCurrentProcess()
_os_memory_lib.SetProcessWorkingSetSize(handle, -1, -1)
elif hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
# macOS with MPS
import ctypes # Import only when needed
import ctypes.util
@@ -591,7 +613,7 @@ def manage_tensor(
target_dtype = dtype if dtype is not None else current_dtype
# Check if movement is actually needed
needs_device_move = current_device != target_device
needs_device_move = _device_str(current_device) != _device_str(target_device)
needs_dtype_change = dtype is not None and current_dtype != target_dtype
if not needs_device_move and not needs_dtype_change:
@@ -609,8 +631,8 @@ def manage_tensor(
# Log the movement
if debug:
current_device_str = str(current_device).upper()
target_device_str = str(target_device).upper()
current_device_str = _device_str(current_device)
target_device_str = _device_str(target_device)
dtype_info = ""
if needs_dtype_change:
@@ -681,8 +703,8 @@ def manage_model_device(model: torch.nn.Module, target_device: torch.device, mod
# Extract device type for comparison (both are torch.device objects)
target_type = target_device.type
current_device_upper = str(current_device).upper()
target_device_upper = str(target_device).upper()
current_device_upper = _device_str(current_device)
target_device_upper = _device_str(target_device)
# Compare normalized device types
if current_device_upper == target_device_upper and not is_blockswap_model:
@@ -737,10 +759,10 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
actual_source_device = param.device
break
source_device_desc = str(actual_source_device).upper() if actual_source_device else str(target_device).upper()
source_device_desc = _device_str(actual_source_device) if actual_source_device else _device_str(target_device)
if debug:
debug.log(f"Moving {model_name} from {source_device_desc} to {str(target_device).upper()} ({reason or 'model caching'})", category="general")
debug.log(f"Moving {model_name} from {source_device_desc} to {_device_str(target_device)} ({reason or 'model caching'})", category="general")
# Enable bypass to allow movement
set_blockswap_bypass(runner=runner, bypass=True, debug=debug)
@@ -755,14 +777,14 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
model.zero_grad(set_to_none=True)
if debug:
debug.end_timer(timer_name, f"BlockSwap model offloaded to {str(target_device).upper()}")
debug.end_timer(timer_name, f"BlockSwap model offloaded to {_device_str(target_device)}")
return True
else:
# Moving to GPU (reload)
# Check if we're in bypass mode (coming from offload)
if not getattr(runner, "_blockswap_bypass_protection", False):
if not getattr(model, "_blockswap_bypass_protection", False):
# Not in bypass mode, blocks are already configured
if debug:
debug.log(f"{model_name} with BlockSwap active - blocks already distributed across devices, skipping movement", category="general")
@@ -775,10 +797,10 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
actual_current_device = param.device
break
current_device_desc = str(actual_current_device).upper() if actual_current_device else "OFFLOAD"
current_device_desc = _device_str(actual_current_device) if actual_current_device else "OFFLOAD"
if debug:
debug.log(f"Moving {model_name} from {current_device_desc} to {str(target_device).upper()} ({reason or 'inference requirement'})", category="general")
debug.log(f"Moving {model_name} from {current_device_desc} to {_device_str(target_device)} ({reason or 'inference requirement'})", category="general")
timer_name = f"{model_name.lower()}_to_gpu"
if debug:
@@ -787,7 +809,7 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
# Restore blocks to their configured devices
if hasattr(model, "blocks") and hasattr(model, "blocks_to_swap"):
# Use configured offload_device from BlockSwap config
offload_device = runner._block_swap_config.get("offload_device")
offload_device = model._block_swap_config.get("offload_device")
if not offload_device:
raise ValueError("BlockSwap config missing offload_device")
@@ -801,7 +823,7 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
block.to(offload_device)
# Handle I/O components
if not runner._block_swap_config.get("swap_io_components", False):
if not model._block_swap_config.get("swap_io_components", False):
# I/O components should be on GPU if not offloaded
for name, module in model.named_children():
if name != "blocks":
@@ -814,11 +836,11 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
if debug:
# Get actual configuration from runner
if hasattr(runner, '_block_swap_config'):
blocks_on_gpu = runner._block_swap_config.get('total_blocks', 32) - runner._block_swap_config.get('blocks_swapped', 16)
total_blocks = runner._block_swap_config.get('total_blocks', 32)
main_device = runner._block_swap_config.get('main_device', 'GPU')
debug.log(f"BlockSwap blocks restored to configured devices ({blocks_on_gpu}/{total_blocks} blocks on {str(main_device).upper()})", category="success")
if hasattr(model, '_block_swap_config'):
blocks_on_gpu = model._block_swap_config.get('total_blocks', 32) - model._block_swap_config.get('blocks_swapped', 16)
total_blocks = model._block_swap_config.get('total_blocks', 32)
main_device = model._block_swap_config.get('main_device', 'GPU')
debug.log(f"BlockSwap blocks restored to configured devices ({blocks_on_gpu}/{total_blocks} blocks on {_device_str(main_device)})", category="success")
else:
debug.log("BlockSwap blocks restored to configured devices", category="success")
@@ -865,8 +887,8 @@ def _standard_model_movement(model: torch.nn.Module, current_device: torch.devic
# Log the movement with full device strings
if debug:
current_device_str = str(current_device).upper()
target_device_str = str(target_device).upper()
current_device_str = _device_str(current_device)
target_device_str = _device_str(target_device)
debug.log(f"Moving {model_name} from {current_device_str} to {target_device_str} ({reason})", category="general")
# Start timer based on direction
@@ -891,7 +913,7 @@ def _standard_model_movement(model: torch.nn.Module, current_device: torch.devic
# End timer
if debug:
debug.end_timer(timer_name, f"{model_name} moved to {str(target_device).upper()}")
debug.end_timer(timer_name, f"{model_name} moved to {_device_str(target_device)}")
return True
+34 -30
View File
@@ -9,6 +9,40 @@ import torch
from typing import List
def optimized_channels_to_last(tensor):
"""🚀 Optimized replacement for rearrange(tensor, 'b c ... -> b ... c')
Moves channels from position 1 to last position using PyTorch native operations.
"""
if tensor.ndim == 3: # [batch, channels, spatial]
return tensor.permute(0, 2, 1)
elif tensor.ndim == 4: # [batch, channels, height, width]
return tensor.permute(0, 2, 3, 1)
elif tensor.ndim == 5: # [batch, channels, depth, height, width]
return tensor.permute(0, 2, 3, 4, 1)
else:
# Fallback for other dimensions - move channel (dim=1) to last
dims = list(range(tensor.ndim))
dims = [dims[0]] + dims[2:] + [dims[1]] # [0, 2, 3, ..., 1]
return tensor.permute(*dims)
def optimized_channels_to_second(tensor):
"""🚀 Optimized replacement for rearrange(tensor, 'b ... c -> b c ...')
Moves channels from last position to position 1 using PyTorch native operations.
"""
if tensor.ndim == 3: # [batch, spatial, channels]
return tensor.permute(0, 2, 1)
elif tensor.ndim == 4: # [batch, height, width, channels]
return tensor.permute(0, 3, 1, 2)
elif tensor.ndim == 5: # [batch, depth, height, width, channels]
return tensor.permute(0, 4, 1, 2, 3)
else:
# Fallback for other dimensions - move last dim to position 1
dims = list(range(tensor.ndim))
dims = [dims[0], dims[-1]] + dims[1:-1] # [0, -1, 1, 2, ..., -2]
return tensor.permute(*dims)
def optimized_video_rearrange(video_tensors: List[torch.Tensor]) -> List[torch.Tensor]:
"""
🚀 OPTIMIZED version of video rearrangement
@@ -132,34 +166,4 @@ def optimized_sample_to_image_format(sample: torch.Tensor) -> torch.Tensor:
return sample.permute(0, 2, 3, 1)
def temporal_latent_blending(latents1: torch.Tensor, latents2: torch.Tensor, blend_frames: int) -> torch.Tensor:
"""
🎨 Temporal blending in latent space to avoid discontinuities
Args:
latents1: Latents from previous batch (end frames)
latents2: Latents from current batch (start frames)
blend_frames: Number of frames to blend
Returns:
Blended latents for smooth transition
"""
if latents1.shape[0] != latents2.shape[0]:
# Adjust dimensions if necessary
min_frames = min(latents1.shape[0], latents2.shape[0])
latents1 = latents1[:min_frames]
latents2 = latents2[:min_frames]
# Create linear blending weights
# Frame 0: 100% latents1, 0% latents2
# Frame n: 0% latents1, 100% latents2
weights1 = torch.linspace(1.0, 0.0, blend_frames).view(-1, 1, 1, 1).to(latents1.device)
weights2 = torch.linspace(0.0, 1.0, blend_frames).view(-1, 1, 1, 1).to(latents2.device)
# Apply blending
blended_latents = weights1 * latents1 + weights2 * latents2
return blended_latents
+25 -25
View File
@@ -178,7 +178,7 @@ def wavelet_decomposition(image: Tensor, levels: int = 5) -> tuple[Tensor, Tenso
for i in range(levels):
radius = 2 ** i
low_freq = wavelet_blur(image, radius)
high_freq += (image - low_freq)
high_freq.add_(image).sub_(low_freq)
image = low_freq
return high_freq, low_freq
@@ -237,13 +237,13 @@ def wavelet_reconstruction(content_feat: Tensor, style_feat: Tensor, debug: Opti
align_corners=False
)
# Reconstruct: content details + style color
result = content_high_freq + style_low_freq
# Reconstruct: content details + style color (in-place on content_high_freq)
content_high_freq.add_(style_low_freq)
# Safety clamp for normalized SDR range
# Safety clamp for normalized SDR range (in-place)
# This prevents numerical errors from propagating
# Note: For HDR support, this would need to be removed
return torch.clamp(result, -1.0, 1.0)
return content_high_freq.clamp_(-1.0, 1.0)
def lab_color_transfer(
@@ -255,17 +255,15 @@ def lab_color_transfer(
"""
Perceptually-accurate color transfer using CIELAB color space.
LAB provides superior perceptual uniformity compared to RGB/HSV, enabling
highly accurate color matching that preserves the original image's appearance.
This is the RECOMMENDED method for color correction.
Combines wavelet reconstruction (for spatial continuity) with LAB histogram
matching (for precise color matching). This eliminates tile artifacts while
providing superior color accuracy.
Algorithm:
1. Convert both images to LAB color space (D65 illuminant)
2. Apply histogram matching to all LAB channels:
- L* (luminance): Weighted blend to preserve detail
- a* (green-red): Full histogram matching
- b* (blue-yellow): Full histogram matching
3. Convert back to RGB
1. Apply wavelet reconstruction to get artifact-free base
2. Convert both wavelet result and style to LAB color space
3. Apply histogram matching to LAB channels
4. Convert back to RGB
Args:
content_feat: Target tensor [B, C, H, W] in [-1, 1] with upscaled details
@@ -278,7 +276,10 @@ def lab_color_transfer(
Returns:
Color-corrected tensor [B, C, H, W] in [-1, 1]
"""
# Handle spatial dimension mismatch
# Step 1: Apply wavelet to get artifact-free base with correct spatial structure
content_feat = wavelet_reconstruction(content_feat, style_feat, debug=None)
# Handle spatial dimension mismatch (should already match after wavelet)
if content_feat.shape != style_feat.shape:
debug.log(
f"LAB: Resizing style {style_feat.shape} to match content {content_feat.shape}",
@@ -291,14 +292,14 @@ def lab_color_transfer(
align_corners=False
)
# Store device
# Store device and convert to float32
device = content_feat.device
# Convert to float32 for accurate color space conversion
content_feat, original_dtype = ensure_float32_precision(content_feat)
style_feat, _ = ensure_float32_precision(style_feat)
# Precompute color space conversion matrices (once per batch, not per pixel)
# Precompute color space conversion matrices
rgb_to_xyz_matrix = torch.tensor([
[0.4124564, 0.3575761, 0.1804375],
[0.2126729, 0.7151522, 0.0721750],
@@ -315,17 +316,16 @@ def lab_color_transfer(
epsilon = 6.0 / 29.0
kappa = (29.0 / 3.0) ** 3
# Convert from [-1, 1] to [0, 1] range (in-place where possible)
content_rgb = content_feat.add(1.0).mul_(0.5).clamp_(0.0, 1.0)
style_rgb = style_feat.add(1.0).mul_(0.5).clamp_(0.0, 1.0)
del content_feat, style_feat
# Convert from [-1, 1] to [0, 1] range (in-place)
content_feat.add_(1.0).mul_(0.5).clamp_(0.0, 1.0)
style_feat.add_(1.0).mul_(0.5).clamp_(0.0, 1.0)
# Convert to LAB color space
content_lab = _rgb_to_lab_batch(content_rgb, device, rgb_to_xyz_matrix, epsilon, kappa)
del content_rgb
content_lab = _rgb_to_lab_batch(content_feat, device, rgb_to_xyz_matrix, epsilon, kappa)
del content_feat
style_lab = _rgb_to_lab_batch(style_rgb, device, rgb_to_xyz_matrix, epsilon, kappa)
del style_rgb, rgb_to_xyz_matrix
style_lab = _rgb_to_lab_batch(style_feat, device, rgb_to_xyz_matrix, epsilon, kappa)
del style_feat, rgb_to_xyz_matrix
# Match chrominance channels (a*, b*) for accurate color transfer
matched_a = _histogram_matching_channel(content_lab[:, 1], style_lab[:, 1], device)
+26 -4
View File
@@ -3,6 +3,9 @@ Shared constants and utilities for SeedVR2
Only includes constants actually used in the codebase
"""
# Version information
__version__ = "2.5.17"
import os
import warnings
import inspect
@@ -52,14 +55,33 @@ def get_base_cache_dir() -> str:
def get_all_model_paths() -> list:
"""Get all registered model paths including those from extra_model_paths.yaml"""
"""Get all registered model paths including those from extra_model_paths.yaml (case-insensitive)"""
try:
import folder_paths
# Ensure default path is registered first
get_base_cache_dir()
# Get all paths registered for seedvr2 model type
paths = folder_paths.get_folder_paths(SEEDVR2_MODEL_TYPE)
return paths if paths else [get_base_cache_dir()]
# Case-insensitive lookup: search through all registered folder types
# This handles any case variation users might use in extra_model_paths.yaml
all_paths = []
target_lower = SEEDVR2_MODEL_TYPE.lower()
# folder_paths.folder_names_and_paths is the underlying dict: {type: ([paths], extensions)}
if hasattr(folder_paths, 'folder_names_and_paths'):
for folder_type, (paths, _) in folder_paths.folder_names_and_paths.items():
if folder_type.lower() == target_lower:
all_paths.extend(paths)
# Remove duplicates while preserving order (os.path.normpath handles Windows/Linux path differences)
seen = set()
unique_paths = []
for path in all_paths:
normalized = os.path.normpath(path.lower())
if normalized not in seen:
seen.add(normalized)
unique_paths.append(path)
return unique_paths if unique_paths else [get_base_cache_dir()]
except:
return [get_base_cache_dir()]
+143 -27
View File
@@ -11,6 +11,15 @@ import gc
from typing import Optional, List, Dict, Any, Union
from datetime import datetime
from ..optimization.memory_manager import get_vram_usage, get_basic_vram_info, get_ram_usage, reset_vram_peak
from ..utils.constants import __version__
def _format_peak_with_swap(peak_gb: float, total_vram_gb: float) -> str:
"""Format peak memory, showing swap breakdown if overflow occurred."""
if total_vram_gb > 0 and peak_gb > total_vram_gb:
swap_gb = peak_gb - total_vram_gb
return f"{peak_gb:.2f}GB ({total_vram_gb:.0f}GB GPU + {swap_gb:.2f}GB swap)"
return f"{peak_gb:.2f}GB"
class Debug:
@@ -71,7 +80,8 @@ class Debug:
self.vram_history: List[float] = []
self.active_timer_stack: List[str] = []
self.timer_namespace: str = ""
self.phase_peaks: Dict[str, float] = {}
self.phase_vram_peaks: Dict[str, float] = {}
self.phase_ram_peaks: Dict[str, float] = {}
@torch._dynamo.disable # Skip tracing to avoid datetime.now() warnings
def log(self, message: str, level: str = "INFO", category: str = "general", force: bool = False, indent_level: int = 0) -> None:
@@ -111,7 +121,7 @@ class Debug:
# Add indentation
indent = " " * (indent_level * 2)
print(f"{prefix} {indent}{message}")
print(f"{prefix} {indent}{message}", flush=True)
def print_header(self, cli: bool = False) -> None:
"""Print the header with banner - always displayed"""
@@ -123,12 +133,84 @@ class Debug:
self.log(" ║ ███████ █████ █████ ██ ██ ██ ██ ██████ █████ ║", category="none", force=True)
self.log(" ║ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ║", category="none", force=True)
self.log(" ║ ███████ ███████ ███████ ██████ ████ ██ ██ ███████ ║", category="none", force=True)
if cli:
self.log(" ║ 💻 CLI mode © ByteDance Seed · NumZ · AInVFX ║", category="none", force=True)
else:
self.log(" ║ © ByteDance Seed · NumZ · AInVFX ║", category="none", force=True)
# Version number with dynamic padding to maintain visual alignment with any version length
version_text = f"v{__version__}"
prefix = " 💻 CLI mode · " if cli else " "
suffix = "© ByteDance Seed · NumZ · AInVFX "
emoji_compensation = 1 if cli else 0
padding_width = 59 - len(prefix) - len(version_text) - len(suffix) - 2 - emoji_compensation
padding = " " * max(1, padding_width)
self.log(f" ║{prefix}{version_text}{padding} {suffix}║", category="none", force=True)
self.log(" ╚══════════════════════════════════════════════════════════╝", category="none", force=True)
self.log("", category="none", force=True)
# Environment info - only in debug mode
if self.enabled:
self._print_environment_info(cli)
def _print_environment_info(self, cli: bool = False) -> None:
"""Print concise environment info for bug reports - zero cost when debug disabled"""
import platform
import sys
# OS
os_name = platform.system()
if os_name == "Windows":
os_str = f"Windows ({platform.version()})"
elif os_name == "Darwin":
os_str = f"macOS {platform.mac_ver()[0]}"
else:
try:
distro = platform.freedesktop_os_release()
os_str = f"{distro.get('NAME', 'Linux')} {distro.get('VERSION_ID', '')}"
except (OSError, AttributeError):
os_str = f"Linux {platform.release()}"
# Python & PyTorch & CUDA
py_ver = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
torch_ver = torch.__version__
cuda_ver = getattr(torch.version, 'cuda', None) or "N/A"
# GPU
if torch.cuda.is_available():
try:
props = torch.cuda.get_device_properties(0)
gpu_str = f"{props.name} ({round(props.total_memory / (1024**3))}GB)"
cudnn_ver = str(torch.backends.cudnn.version()) if torch.backends.cudnn.is_available() else "N/A"
except Exception:
gpu_str = "CUDA"
cudnn_ver = "N/A"
elif getattr(getattr(torch, 'mps', None), 'is_available', lambda: False)():
gpu_str = "Apple Silicon (MPS)"
cudnn_ver = "N/A"
else:
gpu_str = "CPU"
cudnn_ver = "N/A"
# Flash Attn & Triton - reuse existing module constants
try:
from ..optimization.compatibility import FLASH_ATTN_AVAILABLE, TRITON_AVAILABLE
flash_str, triton_str = ("✓" if FLASH_ATTN_AVAILABLE else "✗"), ("✓" if TRITON_AVAILABLE else "✗")
except ImportError:
flash_str = triton_str = "?"
# ComfyUI version
comfy_str = None
if not cli:
try:
from comfyui_version import __version__ as comfy_ver
comfy_str = comfy_ver
except ImportError:
pass
# Print
self.log(f"OS: {os_str} | GPU: {gpu_str}", category="info")
self.log(f"Python: {py_ver} | PyTorch: {torch_ver} | Flash Attn: {flash_str} | Triton: {triton_str}", category="info")
cuda_line = f"CUDA: {cuda_ver} | cuDNN: {cudnn_ver}"
self.log(f"{cuda_line} | ComfyUI: {comfy_str}" if comfy_str else cuda_line, category="info")
self.log("", category="none")
def print_footer(self) -> None:
"""Print the footer with links - always displayed"""
@@ -299,19 +381,30 @@ class Debug:
if show_diff and self.memory_checkpoints:
self._log_memory_diff(current_metrics=memory_info, force=force)
# Log detailed analysis if requested
# Warn if swap detected (peak > physical VRAM)
if memory_info['vram_total'] > 0 and memory_info['vram_peak_since_last'] > memory_info['vram_total']:
self.log("VRAM swap detected - severe slowdown expected. Consider optimizing (e.g., reduce resolution, batch_size, enable BlockSwap, VAE tiling...).",
level="WARNING", category="memory", force=True)
# Log detailed analysis if requested
if detailed_tensors and tensor_stats.get('details'):
self._log_detailed_tensor_analysis(details=tensor_stats['details'], force=force)
# Store checkpoint with memory limit
self._store_checkpoint(label, memory_info)
# Update phase peak if we're in an active phase
if self.current_phase and memory_info['vram_peak_since_last'] > 0:
self.phase_peaks[self.current_phase] = max(
self.phase_peaks.get(self.current_phase, 0),
memory_info['vram_peak_since_last']
)
# Update phase peaks if we're in an active phase
if self.current_phase:
if memory_info['vram_peak_since_last'] > 0:
self.phase_vram_peaks[self.current_phase] = max(
self.phase_vram_peaks.get(self.current_phase, 0),
memory_info['vram_peak_since_last']
)
if memory_info['ram_process'] > 0:
self.phase_ram_peaks[self.current_phase] = max(
self.phase_ram_peaks.get(self.current_phase, 0),
memory_info['ram_process']
)
# Reset PyTorch's peak memory stats for next interval
reset_vram_peak(device=None, debug=self)
@@ -333,7 +426,7 @@ class Debug:
}
# VRAM metrics
if torch.cuda.is_available() or (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()):
if torch.cuda.is_available() or (hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()):
metrics['vram_allocated'], metrics['vram_reserved'], current_global_peak = get_vram_usage(device=None, debug=self)
# Calculate peak since last log_memory_state
@@ -346,10 +439,11 @@ class Debug:
metrics['vram_free'] = vram_info["free_gb"]
metrics['vram_total'] = vram_info["total_gb"]
backend = "MPS" if (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()) else "VRAM"
backend = "MPS" if (hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()) else "VRAM"
peak_str = _format_peak_with_swap(metrics['vram_peak_since_last'], metrics['vram_total'])
metrics['summary_vram'] = (f" [{backend}] {metrics['vram_allocated']:.2f}GB allocated / "
f"{metrics['vram_reserved']:.2f}GB reserved / "
f"Peak: {metrics['vram_peak_since_last']:.2f}GB / "
f"Peak: {peak_str} / "
f"{metrics['vram_free']:.2f}GB free / "
f"{metrics['vram_total']:.2f}GB total")
else:
@@ -369,7 +463,7 @@ class Debug:
metrics['summary_ram'] = ""
# Update VRAM history for tracking
if torch.cuda.is_available() or (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()):
if torch.cuda.is_available() or (hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()):
self.vram_history.append(metrics['vram_allocated'])
return metrics
@@ -497,9 +591,9 @@ class Debug:
if diffs:
self.log(f"Memory changes: {', '.join(diffs)}", category="memory", force=force, indent_level=1)
def log_peak_vram_summary(self, force: bool = True) -> None:
"""Display peak VRAM usage across all phases"""
if not self.phase_peaks:
def log_peak_memory_summary(self, force: bool = True) -> None:
"""Display peak memory usage across all phases (VRAM and RAM combined)"""
if not self.phase_vram_peaks and not self.phase_ram_peaks:
return
phase_names = {
@@ -509,17 +603,38 @@ class Debug:
'phase4': 'Post-processing'
}
is_mps = hasattr(torch.backends, 'mps') and torch.backends.mps.is_available() and not torch.cuda.is_available()
# Get total VRAM for swap detection (reuse existing function)
total_vram_gb = 0.0
if not is_mps:
vram_info = get_basic_vram_info(device=None)
if "error" not in vram_info:
total_vram_gb = vram_info["total_gb"]
self.log("", category="none", force=force)
self.log("────────────────────────", category="none", force=force)
self.log("Peak VRAM by Phase:", category="memory", force=force)
self.log("Peak memory by phase:", category="memory", force=force)
for phase_key in sorted(self.phase_peaks.keys()):
all_phases = sorted(set(self.phase_vram_peaks.keys()) | set(self.phase_ram_peaks.keys()))
for phase_key in all_phases:
phase_num = phase_key[-1]
self.log(f" Phase {phase_num}: {phase_names[phase_key]}: {self.phase_peaks[phase_key]:.2f}GB",
category="memory", force=force)
phase_name = phase_names.get(phase_key, phase_key)
vram = self.phase_vram_peaks.get(phase_key, 0)
ram = self.phase_ram_peaks.get(phase_key, 0)
if is_mps:
self.log(f" Phase {phase_num} ({phase_name}): {vram:.2f}GB", category="memory", force=force)
else:
self.log(f" Phase {phase_num} ({phase_name}): {_format_peak_with_swap(vram, total_vram_gb)} | RAM {ram:.2f}GB", category="memory", force=force)
overall_max = max(self.phase_peaks.values())
self.log(f"Overall Peak VRAM: {overall_max:.2f}GB", category="memory", force=force)
if is_mps:
overall = max(self.phase_vram_peaks.values()) if self.phase_vram_peaks else 0
self.log(f"Overall Peak: {overall:.2f}GB", category="memory", force=force)
else:
overall_vram = max(self.phase_vram_peaks.values()) if self.phase_vram_peaks else 0
overall_ram = max(self.phase_ram_peaks.values()) if self.phase_ram_peaks else 0
self.log(f"Overall peak: {_format_peak_with_swap(overall_vram, total_vram_gb)} | RAM {overall_ram:.2f}GB", category="memory", force=force)
@torch._dynamo.disable # Skip tracing to avoid time.time() warnings
def _store_checkpoint(self, label: str, metrics: Dict[str, Any]) -> None:
@@ -629,5 +744,6 @@ class Debug:
self.timer_durations.clear()
self.timer_messages.clear()
self.active_timer_stack.clear()
self.phase_peaks.clear()
self.phase_vram_peaks.clear()
self.phase_ram_peaks.clear()
self.current_phase = None
+20
View File
@@ -14,6 +14,7 @@ import time
from .model_registry import MODEL_REGISTRY, DEFAULT_VAE
from .constants import (
get_base_cache_dir,
get_all_model_paths,
find_model_file,
get_validation_cache_path,
HUGGINGFACE_BASE_URL,
@@ -179,20 +180,36 @@ def download_weight(dit_model: str, vae_model: str, model_dir: Optional[str] = N
level="WARNING", category="setup")
continue
# Get model type for logging
model_type = "VAE" if model_info.category == "vae" else "DiT"
# Check if file exists in any registered path first
existing_filepath = find_model_file(filename, fallback_dir=cache_dir)
# Use existing file path if it exists, otherwise download to cache_dir
if os.path.exists(existing_filepath):
filepath = existing_filepath
# Debug log: Model found
if debug:
debug.log(f"{model_type} model found: {filepath}", category="setup")
else:
filepath = os.path.join(cache_dir, filename)
# Debug log: Model not found, will need to download
if debug:
searched_paths = get_all_model_paths()
debug.log(f"{model_type} model not found: {filename}", category="setup")
debug.log(f"Searched in {len(searched_paths)} location(s):", category="setup")
for i, path in enumerate(searched_paths, 1):
debug.log(f"[{i}] {path}", category="setup", indent_level=1)
expected_hash = model_info.sha256
repo = model_info.repo
## Quick cache check first
if is_file_validated_cached(filepath, cache_dir):
# Debug log: Model already validated (using cache)
if debug:
debug.log(f"{model_type} model already validated (cache): {filepath}", category="setup")
continue
# File exists - validate it
@@ -201,6 +218,9 @@ def download_weight(dit_model: str, vae_model: str, model_dir: Optional[str] = N
debug.log(f"Validating {filename}...", category="setup", force=True)
if validate_file(filepath, expected_hash, cache_dir):
# Debug log: Model validated successfully
if debug:
debug.log(f"{model_type} model validated successfully: {filepath}", category="setup")
continue
else:
# File is corrupted