Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
58bc9e8bc9 | ||
|
|
3eec5847c0 | ||
|
|
eae3aac60d | ||
|
|
0a660065f0 | ||
|
|
11239eed13 | ||
|
|
b4d7ab89eb | ||
|
|
f061d97fe7 | ||
|
|
43c4f00e19 | ||
|
|
f8998ebd75 | ||
|
|
aa968cf2c9 | ||
|
|
d78f6c268d | ||
|
|
18b44d66e1 | ||
|
|
f68fe920b8 | ||
|
|
65c1c1b6cd | ||
|
|
b40f26167c | ||
|
|
ed53581359 | ||
|
|
71ac9ffe54 | ||
|
|
ffba05907d | ||
|
|
d4dd5e747d | ||
|
|
e2faedaaa6 | ||
|
|
5775ff0f99 | ||
|
|
ff937756a3 | ||
|
|
5848cef05f | ||
|
|
f5b902b8b0 | ||
|
|
19825fa9fa | ||
|
|
021fc7b70f | ||
|
|
b0ac880a4f | ||
|
|
6930e0f13a | ||
|
|
112275d4e5 | ||
|
|
16508d353d | ||
|
|
6f784b1324 | ||
|
|
f791631495 | ||
|
|
b15982b442 | ||
|
|
849f6b5997 | ||
|
|
e34bec76d5 | ||
|
|
70c7f21eb5 | ||
|
|
daa13fb4ca | ||
|
|
7581e0014f | ||
|
|
be2efd474f | ||
|
|
bce62382db | ||
|
|
f792ca5937 | ||
|
|
3ab2c84091 | ||
|
|
8dd0c3061a | ||
|
|
79e7f41216 | ||
|
|
9424020687 | ||
|
|
912ab4a5da | ||
|
|
65dd29a865 | ||
|
|
9f78b30b3d | ||
|
|
96b26bb7f7 | ||
|
|
513ab67214 | ||
|
|
19c761522e | ||
|
|
40d37b7248 |
@@ -0,0 +1 @@
|
||||
github: [adrientoupet, numz]
|
||||
+2
-8
@@ -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.
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
@@ -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
@@ -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
@@ -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)
|
||||
|
||||
@@ -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
@@ -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:
|
||||
|
||||
@@ -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
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user