Merge branch 'numz:main' into main
This commit is contained in:
HB2k
@@ -4,7 +4,7 @@
|
||||
|
||||
Official release of [SeedVR2](https://github.com/ByteDance-Seed/SeedVR) for ComfyUI that enables high-quality video and image upscaling.
|
||||
|
||||
Can run as **Multi-GPU standalone CLI** too, see [🖥️ Run as Standalone](#️-run-as-standalone-cli) section.
|
||||
Can run as **Multi-GPU standalone CLI** too, see [🖥️ Run as Standalone](#-run-as-standalone-cli) section.
|
||||
|
||||
[](https://youtu.be/MBtWYXq_r60)
|
||||
|
||||
@@ -14,8 +14,8 @@ Can run as **Multi-GPU standalone CLI** too, see [🖥️ Run as Standalone](#
|
||||
|
||||
## 📋 Quick Access
|
||||
|
||||
- [🆙 Future Releases](#-future-releases)
|
||||
- [🚀 Updates](#-updates)
|
||||
- [🆙 Future Work](#-future-work)
|
||||
- [🚀 Release Notes](#-release-notes)
|
||||
- [🎯 Features](#-features)
|
||||
- [🔧 Requirements](#-requirements)
|
||||
- [📦 Installation](#-installation)
|
||||
@@ -26,7 +26,7 @@ Can run as **Multi-GPU standalone CLI** too, see [🖥️ Run as Standalone](#
|
||||
- [🙏 Credits](#-credits)
|
||||
- [📜 License](#-license)
|
||||
|
||||
## 🆙 Future Releases
|
||||
## 🆙 Future Work
|
||||
|
||||
We're actively working on improvements and new features. To stay informed:
|
||||
|
||||
@@ -34,13 +34,31 @@ We're actively working on improvements and new features. To stay informed:
|
||||
- **💬 Join the Community**: Learn from others, share your workflows, and get help in the [Discussions](https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/discussions)
|
||||
- **🔮 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
|
||||
## 🚀 Release Notes
|
||||
|
||||
**2025.12.08 - Version 2.5.18**
|
||||
**2025.12.12 - Version 2.5.20**
|
||||
|
||||
- **🚀 Feature: SageAttention Support** - Added `sd2` and `sd3` attention modes for experimental support of SageAttention v2 and v3 kernels.
|
||||
- **⚡ Feature: Precision Control** - Added `precision` selector (auto/fp16/bf16/bf32) to main upscaler node for fine-grained control over computation dtype.
|
||||
- **🔧 Optimization: High-end GPU Tuning** - Improved auto-detection defaults for Blackwell (50-series) and high-VRAM GPUs.
|
||||
- **⚡ Expanded attention backends** - Full support for Flash Attention 2 (Ampere+), Flash Attention 3 (Hopper+), SageAttention 2, and SageAttention 3 (Blackwell/RTX 50xx), with automatic fallback chains to PyTorch SDPA when unavailable *(based on PR by [@naxci1](https://github.com/naxci1) - thank you!)*
|
||||
- **🍎 macOS/Apple Silicon compatibility** - Replaced MPS autocast with explicit dtype conversion throughout VAE and DiT pipelines, resolving hangs and crashes on M-series Macs. BlockSwap now auto-disables with warning (unified memory makes it meaningless)
|
||||
- **🛡️ Flash Attention graceful fallback** - Added compatibility shims for corrupted or partially installed flash_attn/xformers DLLs, preventing startup crashes
|
||||
- **🛡️ AMD ROCm: bitsandbytes conflict fix** - Prevent kernel registration errors when diffusers attempts to re-import broken bitsandbytes installations
|
||||
- **📦 ComfyUI Manager: macOS classifier fix** - Removed NVIDIA CUDA classifier causing false "GPU not supported" warnings on macOS
|
||||
- **📚 Documentation updates** - Updated README with attention backend details, BlockSwap macOS notes, and clarified model caching descriptions
|
||||
|
||||
**2025.12.10 - Version 2.5.19**
|
||||
|
||||
- **🎨 New header logo design** - Refreshed ASCII art banner *(thanks [@naxci1](https://github.com/naxci1))*
|
||||
- **🧹 Remove dead flash attention wrapper** - Removed legacy code from FP8CompatibleDiT; FlashAttentionVarlen already handles backend switching via its `attention_mode` attribute
|
||||
- **🛡️ Fix graceful fallback from flash-attn** - Add compatibility shims for corrupted flash_attn/xformers DLLs, preventing startup crashes when CUDA extensions are broken
|
||||
- **📊 Improved VRAM tracking** - Separate allocated vs reserved memory tracking, Windows-only overflow detection (WDDM paging behavior)
|
||||
- **♻️ Centralize backend detection** - Unified `is_mps_available()`, `is_cuda_available()`, `get_gpu_backend()` helpers across codebase
|
||||
- **🔄 Revert 2.5.14 VRAM limit enforcement** - Removed `set_per_process_memory_fraction` call; Overflow detection and warnings remain.
|
||||
|
||||
**2025.12.09 - Version 2.5.18**
|
||||
|
||||
- **🚀 CLI: Streaming mode for long videos** - New `--chunk_size` flag processes videos in memory-bounded chunks, enabling arbitrarily long videos without RAM limits. Works with model caching (`--cache_dit`/`--cache_vae`) for chunk-to-chunk reuse *(inspired by [disk02](https://github.com/disk02) PR contribution)*
|
||||
- **⚡ CLI: Multi-GPU streaming** - Each GPU now streams its segment internally with independent model caching, improving memory efficiency and enabling `--temporal_overlap` blending at GPU boundaries
|
||||
- **🔧 CLI: Fix large video MemoryError** - Shared memory transfer replaces numpy pickling, preventing crashes on high-resolution/long video outputs *(inspired by [FurkanGozukara](https://github.com/FurkanGozukara) PR contribution)*
|
||||
|
||||
**2025.12.05 - Version 2.5.17**
|
||||
|
||||
@@ -223,7 +241,7 @@ We're actively working on improvements and new features. To stay informed:
|
||||
|
||||
**2025.07.03**
|
||||
|
||||
- 🛠️ Can run as **standalone mode** with **Multi GPU** see [🖥️ Run as Standalone](#️-run-as-standalone-cli)
|
||||
- 🛠️ Can run as **standalone mode** with **Multi GPU** see [🖥️ Run as Standalone](#run-as-standalone-cli)
|
||||
|
||||
**2025.06.30**
|
||||
|
||||
@@ -270,8 +288,8 @@ We're actively working on improvements and new features. To stay informed:
|
||||
### Performance Features
|
||||
- **torch.compile Integration**: Optional 20-40% DiT speedup and 15-25% VAE speedup with PyTorch 2.0+ compilation
|
||||
- **Multi-GPU CLI**: Distribute workload across multiple GPUs with automatic temporal overlap blending
|
||||
- **Model Caching**: Keep models loaded in memory for faster batch processing
|
||||
- **Flexible Attention Backends**: Choose between PyTorch SDPA (stable, always available) or Flash Attention 2 (faster on supported hardware)
|
||||
- **Model Caching**: Keep models loaded between generations for single-GPU directory processing or multi-GPU streaming
|
||||
- **Flexible Attention Backends**: Choose between PyTorch SDPA (stable, always available), Flash Attention 2/3, or SageAttention 2/3 for faster computation on supported hardware
|
||||
|
||||
### Quality Control
|
||||
- **Advanced Color Correction**: Five methods including LAB (recommended for highest fidelity), wavelet, wavelet adaptive, HSV, and AdaIN
|
||||
@@ -300,7 +318,7 @@ With the current optimizations (tiling, BlockSwap, GGUF quantization), SeedVR2 c
|
||||
- **Python**: 3.12+ (Python 3.12 and 3.13 tested and recommended)
|
||||
- **PyTorch**: 2.0+ for torch.compile support (optional but recommended)
|
||||
- **Triton**: Required for torch.compile with inductor backend (optional)
|
||||
- **Flash Attention 2**: Provides faster attention computation on supported hardware (optional, falls back to PyTorch SDPA)
|
||||
- **Flash Attention / SageAttention**: Flash Attention 2 (Ampere+), Flash Attention 3 (Hopper+), SageAttention 2 or SageAttention 3 (Blackwell) provide faster attention computation on supported hardware (optional, falls back to PyTorch SDPA)
|
||||
|
||||
## 📦 Installation
|
||||
|
||||
@@ -415,14 +433,21 @@ Configure the DiT (Diffusion Transformer) model for video upscaling.
|
||||
- Requires offload_device to be set and different from device
|
||||
|
||||
- **attention_mode**: Attention computation backend
|
||||
- `sdpa`: PyTorch scaled_dot_product_attention (default, stable, always available)
|
||||
- `flash_attn`: Flash Attention 2 (faster on supported hardware, requires flash-attn package)
|
||||
- `sdpa`: PyTorch scaled_dot_product_attention (default, always available)
|
||||
- `flash_attn_2`: Flash Attention 2 (Ampere+, requires flash-attn package)
|
||||
- `flash_attn_3`: Flash Attention 3 (Hopper+, requires flash-attn with FA3 support)
|
||||
- `sageattn_2`: SageAttention 2 (requires sageattention package)
|
||||
- `sageattn_3`: SageAttention 3 (Blackwell/RTX 50xx, requires sageattn3 package)
|
||||
|
||||
- **torch_compile_args**: Connect to SeedVR2 Torch Compile Settings node for 20-40% speedup
|
||||
|
||||
**BlockSwap Explained:**
|
||||
|
||||
BlockSwap enables running large models on GPUs with limited VRAM by dynamically swapping transformer blocks between GPU and CPU memory during inference. Here's how it works:
|
||||
BlockSwap enables running large models on GPUs with limited VRAM by dynamically swapping transformer blocks between GPU and CPU memory during inference.
|
||||
|
||||
> **Note:** BlockSwap is not available on macOS. Apple Silicon Macs use unified memory architecture where GPU and CPU share the same memory pool, making BlockSwap meaningless. The option will be automatically disabled with a warning if requested on macOS.
|
||||
|
||||
Here's how it works:
|
||||
|
||||
- **What it does**: Keeps only the currently-needed transformer blocks on the GPU, while storing the rest on CPU or another device
|
||||
- **When to use it**: When you get OOM (Out of Memory) errors during the upscaling phase
|
||||
@@ -777,9 +802,18 @@ The CLI provides comprehensive options for single-GPU, multi-GPU, and batch proc
|
||||
# Basic image upscaling
|
||||
python inference_cli.py image.jpg
|
||||
|
||||
# Basic video video upscaling with temporal consistency
|
||||
# Basic video upscaling with temporal consistency
|
||||
python inference_cli.py video.mp4 --resolution 720 --batch_size 33
|
||||
|
||||
# Streaming mode for long videos (memory-efficient)
|
||||
# Processes video in chunks of 330 frames to avoid loading entire video into RAM
|
||||
# Use --temporal_overlap to ensure smooth transitions between chunks
|
||||
python inference_cli.py long_video.mp4 \
|
||||
--resolution 1080 \
|
||||
--batch_size 33 \
|
||||
--chunk_size 330 \
|
||||
--temporal_overlap 3
|
||||
|
||||
# Multi-GPU processing with temporal overlap
|
||||
python inference_cli.py video.mp4 \
|
||||
--cuda_device 0,1 \
|
||||
@@ -836,7 +870,8 @@ python inference_cli.py media_folder/ \
|
||||
- `--batch_size`: Frames per batch (must follow 4n+1: 1, 5, 9, 13, 17, 21...). Ideally matches shot length for best temporal consistency (default: 5)
|
||||
- `--seed`: Random seed for reproducibility (default: 42)
|
||||
- `--skip_first_frames`: Skip N initial frames (default: 0)
|
||||
- `--load_cap`: Load maximum N frames from video. 0 = load all (default: 0)
|
||||
- `--load_cap`: Maximum total frames to load from video. 0 = load all (default: 0)
|
||||
- `--chunk_size`: Frames per chunk for streaming mode. When > 0, processes video in memory-bounded chunks of N frames, writing each chunk before loading the next. Essential for long videos that would otherwise exceed RAM. Use with `--temporal_overlap` for seamless chunk transitions. 0 = load all frames at once (default: 0)
|
||||
- `--prepend_frames`: Prepend N reversed frames to reduce start artifacts (auto-removed) (default: 0)
|
||||
- `--temporal_overlap`: Frames to overlap between batches/GPUs for smooth blending (default: 0)
|
||||
|
||||
@@ -848,9 +883,8 @@ python inference_cli.py media_folder/ \
|
||||
**Memory Management:**
|
||||
- `--dit_offload_device`: Device to offload DiT model: 'none' (keep on GPU), 'cpu', or 'cuda:X' (default: none)
|
||||
- `--vae_offload_device`: Device to offload VAE model: 'none', 'cpu', or 'cuda:X' (default: none)
|
||||
- `--blocks_to_swap`: Number of transformer blocks to swap (0=disabled, 3B: 0-32, 7B: 0-36). Requires dit_offload_device (default: 0)
|
||||
- `--swap_io_components`: Offload I/O components for additional VRAM savings. Requires dit_offload_device
|
||||
- `--use_non_blocking`: Use non-blocking memory transfers for BlockSwap (recommended)
|
||||
- `--blocks_to_swap`: Number of transformer blocks to swap (0=disabled, 3B: 0-32, 7B: 0-36). Requires dit_offload_device (default: 0). Not available on macOS.
|
||||
- `--swap_io_components`: Offload I/O components for additional VRAM savings. Requires dit_offload_device. Not available on macOS.
|
||||
|
||||
**VAE Tiling:**
|
||||
- `--vae_encode_tiled`: Enable VAE encode tiling to reduce VRAM during encoding
|
||||
@@ -862,7 +896,8 @@ python inference_cli.py media_folder/ \
|
||||
- `--tile_debug`: Visualize tiles: 'false' (default), 'encode', or 'decode'
|
||||
|
||||
**Performance Optimization:**
|
||||
- `--attention_mode`: Attention backend: 'sdpa' (default, stable) or 'flash_attn' (faster, requires package)
|
||||
- `--allow_vram_overflow`: Allow VRAM overflow to system RAM. Prevents OOM but may cause severe slowdown
|
||||
- `--attention_mode`: Attention backend: 'sdpa' (default), 'flash_attn_2' (Ampere+), 'flash_attn_3' (Hopper+), 'sageattn_2', or 'sageattn_3' (Blackwell)
|
||||
- `--compile_dit`: Enable torch.compile for DiT model (20-40% speedup, requires PyTorch 2.0+ and Triton)
|
||||
- `--compile_vae`: Enable torch.compile for VAE model (15-25% speedup, requires PyTorch 2.0+ and Triton)
|
||||
- `--compile_backend`: Compilation backend: 'inductor' (full optimization) or 'cudagraphs' (lightweight) (default: inductor)
|
||||
@@ -873,8 +908,8 @@ python inference_cli.py media_folder/ \
|
||||
- `--compile_dynamo_recompile_limit`: Max recompilation attempts before fallback (default: 128)
|
||||
|
||||
**Model Caching (batch processing):**
|
||||
- `--cache_dit`: Cache DiT model between files (single GPU only, speeds up directory processing)
|
||||
- `--cache_vae`: Cache VAE model between files (single GPU only, speeds up directory processing)
|
||||
- `--cache_dit`: Keep DiT model in memory between generations. Works with single-GPU directory processing or multi-GPU streaming (`--chunk_size`). Requires `--dit_offload_device`
|
||||
- `--cache_vae`: Keep VAE model in memory between generations. Works with single-GPU directory processing or multi-GPU streaming (`--chunk_size`). Requires `--vae_offload_device`
|
||||
|
||||
**Multi-GPU:**
|
||||
- `--cuda_device`: CUDA device id(s). Single id (e.g., '0') or comma-separated list '0,1' for multi-GPU
|
||||
@@ -977,7 +1012,7 @@ For detailed contribution guidelines, see [CONTRIBUTING.md](CONTRIBUTING.md).
|
||||
|
||||
This ComfyUI implementation is a collaborative project by **[NumZ](https://github.com/numz)** and **[AInVFX](https://www.youtube.com/@AInVFX)** (Adrien Toupet), based on the original [SeedVR2](https://github.com/ByteDance-Seed/SeedVR) by ByteDance Seed Team.
|
||||
|
||||
Special thanks to our community contributors including [benjaminherb](https://github.com/benjaminherb), [cmeka](https://github.com/cmeka), [FurkanGozukara](https://github.com/FurkanGozukara), [JohnAlcatraz](https://github.com/JohnAlcatraz), [lihaoyun6](https://github.com/lihaoyun6), [Luchuanzhao](https://github.com/Luchuanzhao), [Luke2642](https://github.com/Luke2642), [naxci1](https://github.com/naxci1), [q5sys](https://github.com/q5sys), and many others for their improvements, bug fixes, and testing.
|
||||
Special thanks to our community contributors including [naxci1](https://github.com/naxci1), [benjaminherb](https://github.com/benjaminherb), [cmeka](https://github.com/cmeka), [FurkanGozukara](https://github.com/FurkanGozukara), [JohnAlcatraz](https://github.com/JohnAlcatraz), [lihaoyun6](https://github.com/lihaoyun6), [Luchuanzhao](https://github.com/Luchuanzhao), [Luke2642](https://github.com/Luke2642), [proxyid](https://github.com/proxyid), [q5sys](https://github.com/q5sys), and many others for their improvements, bug fixes, and testing.
|
||||
|
||||
## 📜 License
|
||||
|
||||
|
||||
+504
-325
File diff suppressed because it is too large
Load Diff
+2
-3
@@ -1,15 +1,14 @@
|
||||
[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.17"
|
||||
version = "2.5.20"
|
||||
authors = [
|
||||
{name = "numz"},
|
||||
{name = "adrientoupet"}
|
||||
]
|
||||
license = {file = "LICENSE"}
|
||||
classifiers = [
|
||||
"Operating System :: OS Independent",
|
||||
"Environment :: GPU :: NVIDIA CUDA"
|
||||
"Operating System :: OS Independent"
|
||||
]
|
||||
dependencies = [
|
||||
"torch",
|
||||
|
||||
@@ -21,6 +21,7 @@ from datetime import timedelta
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from ...optimization.memory_manager import is_mps_available
|
||||
|
||||
def get_global_rank() -> int:
|
||||
"""
|
||||
@@ -47,7 +48,7 @@ def get_device() -> torch.device:
|
||||
"""
|
||||
Get current rank device.
|
||||
"""
|
||||
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
||||
if is_mps_available():
|
||||
return torch.device("mps")
|
||||
return torch.device("cuda", get_local_rank())
|
||||
|
||||
|
||||
@@ -700,7 +700,7 @@ def upscale_all_batches(
|
||||
)
|
||||
conditions = [condition]
|
||||
|
||||
# Detect DiT model dtype (handle FP8CompatibleDiT wrapper)
|
||||
# Detect DiT model dtype (handle CompatibleDiT wrapper)
|
||||
dit_model = runner.dit.dit_model if hasattr(runner.dit, 'dit_model') else runner.dit
|
||||
try:
|
||||
dit_dtype = next(dit_model.parameters()).dtype
|
||||
@@ -708,9 +708,10 @@ def upscale_all_batches(
|
||||
dit_dtype = ctx['compute_dtype'] # Fallback for meta device or empty model
|
||||
|
||||
# Use autocast if DiT dtype differs from compute dtype
|
||||
# Skip autocast on MPS (CompatibleDiT already handles dtype conversion)
|
||||
debug.start_timer(f"dit_inference_{upscale_idx+1}")
|
||||
with torch.no_grad():
|
||||
if dit_dtype != ctx['compute_dtype']:
|
||||
if dit_dtype != ctx['compute_dtype'] and ctx['dit_device'].type != 'mps':
|
||||
with torch.autocast(ctx['dit_device'].type, ctx['compute_dtype'], enabled=True):
|
||||
upscaled_latents = runner.inference(
|
||||
noises=noises,
|
||||
@@ -1414,7 +1415,7 @@ def postprocess_all_batches(
|
||||
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}",
|
||||
debug.log(f"Output assembled: {frame_info}, Resolution: {Wf}x{Hf}px, Channels: {channels_str}",
|
||||
category="generation", force=True)
|
||||
else:
|
||||
ctx['final_video'] = torch.empty((0, 0, 0, 0), dtype=ctx['compute_dtype'])
|
||||
|
||||
@@ -365,10 +365,6 @@ def setup_generation_context(
|
||||
interrupt_fn = None
|
||||
comfyui_available = False
|
||||
|
||||
# Determine compute dtype (allow override in context setup if needed, but standard flow uses COMPUTE_DTYPE)
|
||||
# The precision selection in UI overrides this later in model_configuration.
|
||||
compute_dtype = COMPUTE_DTYPE
|
||||
|
||||
# Create generation context
|
||||
ctx = {
|
||||
'dit_device': dit_device,
|
||||
@@ -376,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': compute_dtype,
|
||||
'compute_dtype': COMPUTE_DTYPE,
|
||||
'interrupt_fn': interrupt_fn,
|
||||
'video_transform': None,
|
||||
'text_embeds': None,
|
||||
@@ -436,7 +432,6 @@ def prepare_runner(
|
||||
decode_tile_overlap: Optional[Tuple[int, int]] = None,
|
||||
tile_debug: str = "false",
|
||||
attention_mode: str = 'sdpa',
|
||||
precision: str = 'auto',
|
||||
torch_compile_args_dit: Optional[Dict[str, Any]] = None,
|
||||
torch_compile_args_vae: Optional[Dict[str, Any]] = None
|
||||
) -> Tuple['VideoDiffusionInfer', Dict[str, Any]]:
|
||||
@@ -462,7 +457,7 @@ def prepare_runner(
|
||||
decode_tile_size: Tile size for decoding (height, width)
|
||||
decode_tile_overlap: Tile overlap for decoding (height, width)
|
||||
tile_debug: Tile visualization mode (false/encode/decode)
|
||||
attention_mode: Attention computation backend ('sdpa' or 'flash_attn')
|
||||
attention_mode: Attention computation backend ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
|
||||
torch_compile_args_dit: Optional torch.compile configuration for DiT model
|
||||
torch_compile_args_vae: Optional torch.compile configuration for VAE model
|
||||
|
||||
@@ -507,7 +502,6 @@ def prepare_runner(
|
||||
decode_tile_overlap=decode_tile_overlap,
|
||||
tile_debug=tile_debug,
|
||||
attention_mode=attention_mode,
|
||||
precision=precision,
|
||||
torch_compile_args_dit=torch_compile_args_dit,
|
||||
torch_compile_args_vae=torch_compile_args_vae
|
||||
)
|
||||
@@ -535,8 +529,8 @@ def load_text_embeddings(script_directory: str, device: torch.device,
|
||||
- Memory-efficient embedding preparation
|
||||
- Consistent movement logging
|
||||
"""
|
||||
text_pos_embeds = torch.load(os.path.join(script_directory, 'pos_emb.pt'))
|
||||
text_neg_embeds = torch.load(os.path.join(script_directory, 'neg_emb.pt'))
|
||||
text_pos_embeds = torch.load(os.path.join(script_directory, 'pos_emb.pt'), weights_only=True)
|
||||
text_neg_embeds = torch.load(os.path.join(script_directory, 'neg_emb.pt'), weights_only=True)
|
||||
|
||||
text_pos_embeds = manage_tensor(
|
||||
tensor=text_pos_embeds,
|
||||
@@ -825,4 +819,4 @@ def ensure_precision_initialized(
|
||||
debug.log(f"Model precision: {', '.join(parts)}", category="precision")
|
||||
|
||||
except Exception as e:
|
||||
debug.log(f"Could not log model dtypes: {e}", level="WARNING", category="precision", force=True)
|
||||
debug.log(f"Could not log model dtypes: {e}", level="WARNING", category="precision", force=True)
|
||||
+24
-3
@@ -154,15 +154,26 @@ class VideoDiffusionInfer():
|
||||
vae_dtype = dtype # Fallback
|
||||
|
||||
# Use autocast if VAE dtype differs from input dtype
|
||||
# Skip autocast on MPS (only supports bf16, unified memory = no benefit)
|
||||
# Instead, explicitly convert input to model dtype
|
||||
if vae_dtype != sample.dtype:
|
||||
with torch.autocast(device.type, sample.dtype, enabled=True):
|
||||
if device.type == 'mps':
|
||||
# MPS: explicit dtype conversion instead of autocast
|
||||
sample = sample.to(vae_dtype)
|
||||
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
|
||||
else:
|
||||
# Deterministic vae encode, only used for i2v inference (optionally)
|
||||
latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size,
|
||||
tile_overlap=self.encode_tile_overlap).posterior.mode().squeeze(2)
|
||||
else:
|
||||
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
|
||||
else:
|
||||
latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size,
|
||||
tile_overlap=self.encode_tile_overlap).posterior.mode().squeeze(2)
|
||||
else:
|
||||
if use_sample:
|
||||
latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size,
|
||||
@@ -230,13 +241,23 @@ class VideoDiffusionInfer():
|
||||
vae_dtype = dtype # Fallback
|
||||
|
||||
# Use autocast if VAE dtype differs from latent dtype
|
||||
# Skip autocast on MPS (only supports bf16, unified memory = no benefit)
|
||||
if vae_dtype != latent.dtype:
|
||||
with torch.autocast(device.type, latent.dtype, enabled=True):
|
||||
if device.type == 'mps':
|
||||
# MPS: explicit dtype conversion instead of autocast
|
||||
latent = latent.to(vae_dtype)
|
||||
sample = self.vae.decode(
|
||||
latent,
|
||||
tiled=self.decode_tiled, tile_size=self.decode_tile_size,
|
||||
tile_overlap=self.decode_tile_overlap
|
||||
).sample
|
||||
else:
|
||||
with torch.autocast(device.type, latent.dtype, enabled=True):
|
||||
sample = self.vae.decode(
|
||||
latent,
|
||||
tiled=self.decode_tiled, tile_size=self.decode_tile_size,
|
||||
tile_overlap=self.decode_tile_overlap
|
||||
).sample
|
||||
else:
|
||||
sample = self.vae.decode(
|
||||
latent,
|
||||
|
||||
@@ -70,13 +70,11 @@ from .model_cache import get_global_cache
|
||||
from ..common.config import load_config
|
||||
from ..models.video_vae_v3.modules.causal_inflation_lib import InflatedCausalConv3d
|
||||
from ..optimization.compatibility import (
|
||||
FP8CompatibleDiT,
|
||||
CompatibleDiT,
|
||||
TRITON_AVAILABLE,
|
||||
validate_flash_attention_availability,
|
||||
detect_high_end_system,
|
||||
log_system_capabilities
|
||||
validate_attention_mode
|
||||
)
|
||||
from ..optimization.blockswap import is_blockswap_enabled, apply_block_swap_to_dit, cleanup_blockswap
|
||||
from ..optimization.blockswap import is_blockswap_enabled, validate_blockswap_config, apply_block_swap_to_dit, cleanup_blockswap
|
||||
from ..optimization.memory_manager import cleanup_dit, cleanup_vae
|
||||
from ..utils.constants import find_model_file
|
||||
|
||||
@@ -173,7 +171,7 @@ def _describe_attention_mode(attention_mode: Optional[str]) -> str:
|
||||
Generate human-readable description of attention mode configuration.
|
||||
|
||||
Args:
|
||||
attention_mode: Attention mode string ('sdpa' or 'flash_attn' or 'sa2' or 'sa3')
|
||||
attention_mode: Attention mode string ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
|
||||
|
||||
Returns:
|
||||
Human-readable description string
|
||||
@@ -183,9 +181,10 @@ def _describe_attention_mode(attention_mode: Optional[str]) -> str:
|
||||
|
||||
mode_descriptions = {
|
||||
'sdpa': 'PyTorch SDPA',
|
||||
'flash_attn': 'Flash Attention 2',
|
||||
'sa2': 'SageAttention v2',
|
||||
'sa3': 'SageAttention v3'
|
||||
'flash_attn_2': 'Flash Attention 2',
|
||||
'flash_attn_3': 'Flash Attention 3',
|
||||
'sageattn_2': 'SageAttention 2',
|
||||
'sageattn_3': 'SageAttention 3 (Blackwell)'
|
||||
}
|
||||
|
||||
return mode_descriptions.get(attention_mode, attention_mode)
|
||||
@@ -440,7 +439,7 @@ def _update_dit_config(
|
||||
- dynamic: bool - Enable dynamic shapes
|
||||
- dynamo_cache_size_limit: int - Cache size limit
|
||||
- dynamo_recompile_limit: int - Recompilation limit
|
||||
attention_mode: Attention computation backend ('sdpa' or 'flash_attn')
|
||||
attention_mode: Attention computation backend ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
|
||||
debug: Debug instance for logging
|
||||
|
||||
Returns:
|
||||
@@ -749,7 +748,6 @@ def configure_runner(
|
||||
decode_tile_overlap: Optional[Tuple[int, int]] = None,
|
||||
tile_debug: str = "false",
|
||||
attention_mode: str = 'sdpa',
|
||||
precision: str = 'auto',
|
||||
torch_compile_args_dit: Optional[Dict[str, Any]] = None,
|
||||
torch_compile_args_vae: Optional[Dict[str, Any]] = None
|
||||
) -> Tuple[VideoDiffusionInfer, Dict[str, Any]]:
|
||||
@@ -776,7 +774,7 @@ def configure_runner(
|
||||
decode_tile_size: Tile size for decoding (height, width)
|
||||
decode_tile_overlap: Tile overlap for decoding (height, width)
|
||||
tile_debug: Tile visualization mode (false/encode/decode)
|
||||
attention_mode: Attention computation backend ('sdpa' or 'flash_attn')
|
||||
attention_mode: Attention computation backend ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
|
||||
torch_compile_args_dit: Optional torch.compile configuration for DiT model
|
||||
torch_compile_args_vae: Optional torch.compile configuration for VAE model
|
||||
|
||||
@@ -797,9 +795,14 @@ def configure_runner(
|
||||
if debug is None:
|
||||
raise ValueError("Debug instance must be provided to configure_runner")
|
||||
|
||||
# Log installed attention backends and versions
|
||||
log_system_capabilities(debug)
|
||||
|
||||
# Validate BlockSwap configuration early (before any model loading)
|
||||
block_swap_config = validate_blockswap_config(
|
||||
block_swap_config=block_swap_config,
|
||||
dit_device=ctx['dit_device'],
|
||||
dit_offload_device=ctx.get('dit_offload_device'),
|
||||
debug=debug
|
||||
)
|
||||
|
||||
# Phase 1: Initialize cache and get cached models
|
||||
cache_context = _initialize_cache_context(
|
||||
dit_cache, vae_cache, dit_id, vae_id,
|
||||
@@ -822,9 +825,6 @@ def configure_runner(
|
||||
block_swap_config, debug
|
||||
)
|
||||
|
||||
# Store precision setting
|
||||
runner._precision = precision
|
||||
|
||||
# Phase 4: Setup models (load from cache or create new)
|
||||
_setup_models(
|
||||
runner, cache_context, dit_model, vae_model,
|
||||
@@ -868,7 +868,7 @@ def _configure_runner_settings(
|
||||
decode_tile_size: Tile dimensions (height, width) for decoding in pixels
|
||||
decode_tile_overlap: Overlap dimensions (height, width) between decoding tiles
|
||||
tile_debug: Tile visualization mode (false/encode/decode)
|
||||
attention_mode: Attention computation backend ('sdpa' or 'flash_attn')
|
||||
attention_mode: Attention computation backend ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
|
||||
torch_compile_args_dit: torch.compile configuration for DiT model or None
|
||||
torch_compile_args_vae: torch.compile configuration for VAE model or None
|
||||
block_swap_config: BlockSwap configuration for DiT model or None
|
||||
@@ -906,16 +906,6 @@ def _configure_runner_settings(
|
||||
runner._tensor_offload_device = ctx['tensor_offload_device']
|
||||
runner._compute_dtype = ctx['compute_dtype']
|
||||
|
||||
# Auto-detection for 5070ti/similar hardware
|
||||
system_opts = detect_high_end_system()
|
||||
if system_opts.get('high_vram', False):
|
||||
if debug:
|
||||
debug.log(f"Detected high-end system optimizations: {system_opts}", category="setup")
|
||||
# Apply recommended settings if not overridden
|
||||
# For example, we might favor speed/quality trade-offs differently
|
||||
# Here we just log it as the user has control via UI, but we could set defaults if they were None
|
||||
pass
|
||||
|
||||
runner.debug = debug
|
||||
|
||||
|
||||
@@ -1186,43 +1176,27 @@ def apply_model_specific_config(model: torch.nn.Module, runner: VideoDiffusionIn
|
||||
"""
|
||||
if is_dit:
|
||||
# DiT-specific
|
||||
# Apply FP8 compatibility wrapper with compute_dtype
|
||||
if not isinstance(model, FP8CompatibleDiT):
|
||||
debug.log("Applying FP8/RoPE compatibility wrapper to DiT model", category="setup")
|
||||
debug.start_timer("FP8CompatibleDiT")
|
||||
# Apply compatibility wrapper with compute_dtype
|
||||
if not isinstance(model, CompatibleDiT):
|
||||
debug.log("Applying DiT compatibility wrapper", category="setup")
|
||||
debug.start_timer("CompatibleDiT")
|
||||
# Get compute_dtype from runner if available, fallback to bfloat16
|
||||
compute_dtype = getattr(runner, '_compute_dtype', torch.bfloat16)
|
||||
model = FP8CompatibleDiT(model, debug, compute_dtype=compute_dtype, skip_conversion=False)
|
||||
debug.end_timer("FP8CompatibleDiT", "FP8/RoPE compatibility wrapper application")
|
||||
model = CompatibleDiT(model, debug, compute_dtype=compute_dtype, skip_conversion=False)
|
||||
debug.end_timer("CompatibleDiT", "Compatibility wrapper application")
|
||||
else:
|
||||
debug.log("Reusing existing FP8/RoPE compatibility wrapper", category="reuse")
|
||||
debug.log("Reusing existing DiT compatibility wrapper", category="reuse")
|
||||
|
||||
# Apply attention mode and compute_dtype to all FlashAttentionVarlen modules
|
||||
if hasattr(runner, '_dit_attention_mode'):
|
||||
requested_attention_mode = runner._dit_attention_mode or 'sdpa'
|
||||
|
||||
# Validate and get final attention_mode (with warning if fallback needed)
|
||||
attention_mode = validate_flash_attention_availability(requested_attention_mode, debug)
|
||||
attention_mode = validate_attention_mode(requested_attention_mode, debug)
|
||||
|
||||
# Get compute_dtype from runner, override if precision set
|
||||
compute_dtype = getattr(runner, '_compute_dtype', torch.bfloat16)
|
||||
|
||||
# Apply precision override if set
|
||||
precision_override = getattr(runner, '_precision', 'auto')
|
||||
if precision_override == 'fp16':
|
||||
compute_dtype = torch.float16
|
||||
elif precision_override == 'bf16':
|
||||
compute_dtype = torch.bfloat16
|
||||
elif precision_override == 'bf32':
|
||||
# TF32 is an attribute of the context, not dtype, but we can respect it here or in setup
|
||||
# For compute dtype, usually bf32 implies tf32 or float32 with tf32
|
||||
compute_dtype = torch.float32
|
||||
# We will handle TF32 setting globally elsewhere or here if needed
|
||||
|
||||
# Log final decision prominently
|
||||
mode_desc = _describe_attention_mode(attention_mode)
|
||||
debug.log(f"Using Attention Mode: {mode_desc}", category="info", force=True)
|
||||
debug.log(f"Using Compute Dtype: {compute_dtype}", category="info", force=True)
|
||||
# Get compute_dtype from runner
|
||||
compute_dtype = getattr(runner, '_compute_dtype', torch.bfloat16)
|
||||
debug.log(f"Applying {attention_mode} attention mode and {compute_dtype} compute dtype to model", category="setup")
|
||||
|
||||
# Get the actual model (unwrap if needed)
|
||||
actual_model = model.dit_model if hasattr(model, 'dit_model') else model
|
||||
@@ -1503,4 +1477,4 @@ def _propagate_debug_to_modules(module: torch.nn.Module, debug: 'Debug') -> None
|
||||
for name, submodule in module.named_modules():
|
||||
if submodule.__class__.__name__ in target_modules:
|
||||
if not hasattr(submodule, 'debug'): # Only set if not already present
|
||||
submodule.debug = debug
|
||||
submodule.debug = debug
|
||||
@@ -19,6 +19,7 @@ import torch
|
||||
from PIL import Image
|
||||
from torchvision.transforms import functional as TVF
|
||||
from torchvision.transforms.functional import InterpolationMode
|
||||
from ....optimization.memory_manager import is_mps_available
|
||||
|
||||
|
||||
class AreaResize:
|
||||
@@ -31,7 +32,7 @@ class AreaResize:
|
||||
self.max_area = max_area
|
||||
self.downsample_only = downsample_only
|
||||
self.interpolation = interpolation
|
||||
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
||||
if is_mps_available():
|
||||
self.interpolation = InterpolationMode.BILINEAR
|
||||
|
||||
def __call__(self, image: Union[torch.Tensor, Image.Image]):
|
||||
|
||||
@@ -18,6 +18,7 @@ from torchvision.transforms import CenterCrop, Compose, InterpolationMode, Resiz
|
||||
|
||||
from .area_resize import AreaResize
|
||||
from .side_resize import SideResize
|
||||
from ....optimization.memory_manager import is_mps_available
|
||||
|
||||
def NaResize(
|
||||
resolution: int,
|
||||
@@ -26,7 +27,7 @@ def NaResize(
|
||||
max_resolution: int = 0,
|
||||
interpolation: InterpolationMode = InterpolationMode.BICUBIC,
|
||||
):
|
||||
Interpolation = InterpolationMode.BILINEAR if (hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()) else interpolation
|
||||
Interpolation = InterpolationMode.BILINEAR if is_mps_available() else interpolation
|
||||
if mode == "area":
|
||||
return AreaResize(
|
||||
max_area=resolution**2,
|
||||
|
||||
@@ -17,6 +17,7 @@ import torch
|
||||
from PIL import Image
|
||||
from torchvision.transforms import InterpolationMode
|
||||
from torchvision.transforms import functional as TVF
|
||||
from ....optimization.memory_manager import is_mps_available
|
||||
|
||||
class SideResize:
|
||||
def __init__(
|
||||
@@ -30,7 +31,7 @@ class SideResize:
|
||||
self.max_size = max_size
|
||||
self.downsample_only = downsample_only
|
||||
self.interpolation = interpolation
|
||||
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
||||
if is_mps_available():
|
||||
self.interpolation = InterpolationMode.BILINEAR
|
||||
|
||||
def __call__(self, image: Union[torch.Tensor, Image.Image]):
|
||||
|
||||
@@ -66,7 +66,8 @@ class SeedVR2LoadDiTModel(io.ComfyNode):
|
||||
"• 3B model: 0-32 blocks\n"
|
||||
"• 7B model: 0-36 blocks\n"
|
||||
"\n"
|
||||
"Requires offload_device to be set and different from device."
|
||||
"Requires offload_device to be set and different from device.\n"
|
||||
"Not available on macOS (unified memory architecture)."
|
||||
)
|
||||
),
|
||||
io.Boolean.Input("swap_io_components",
|
||||
@@ -74,7 +75,8 @@ class SeedVR2LoadDiTModel(io.ComfyNode):
|
||||
optional=True,
|
||||
tooltip=(
|
||||
"Offload input/output embeddings and normalization layers to reduce VRAM.\n"
|
||||
"Requires offload_device to be set and different from device."
|
||||
"Requires offload_device to be set and different from device.\n"
|
||||
"Not available on macOS (unified memory architecture)."
|
||||
)
|
||||
),
|
||||
io.Combo.Input("offload_device",
|
||||
@@ -100,15 +102,16 @@ class SeedVR2LoadDiTModel(io.ComfyNode):
|
||||
)
|
||||
),
|
||||
io.Combo.Input("attention_mode",
|
||||
options=["sdpa", "flash_attn", "sa2", "sa3"],
|
||||
options=["sdpa", "flash_attn_2", "flash_attn_3", "sageattn_2", "sageattn_3"],
|
||||
default="sdpa",
|
||||
optional=True,
|
||||
tooltip=(
|
||||
"Attention computation backend:\n"
|
||||
"• sdpa: PyTorch scaled_dot_product_attention (default, stable, always available)\n"
|
||||
"• flash_attn: Flash Attention 2 (faster on supported hardware, requires flash-attn package)\n"
|
||||
"• sa2: SageAttention v2 (requires sageattention package)\n"
|
||||
"• sa3: SageAttention v3 (requires sageattention package)\n"
|
||||
"• flash_attn_2: Flash Attention 2 (Ampere+, requires flash-attn package)\n"
|
||||
"• flash_attn_3: Flash Attention 3 (Hopper+, requires flash-attn with FA3 support)\n"
|
||||
"• sageattn_2: SageAttention 2 (requires sageattention package)\n"
|
||||
"• sageattn_3: SageAttention 3 (Blackwell/RTX 50xx only, requires sageattn3 package)\n"
|
||||
"\n"
|
||||
"SDPA is recommended - stable and works everywhere.\n"
|
||||
"Flash Attention and SageAttention provide speedup through optimized CUDA kernels on compatible GPUs."
|
||||
@@ -144,23 +147,15 @@ class SeedVR2LoadDiTModel(io.ComfyNode):
|
||||
cache_model: Whether to keep model loaded between runs
|
||||
blocks_to_swap: Number of transformer blocks to swap (requires offload_device != device)
|
||||
swap_io_components: Whether to offload I/O components (requires offload_device != device)
|
||||
attention_mode: Attention computation backend ('sdpa' or 'flash_attn')
|
||||
attention_mode: Attention computation backend ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
|
||||
torch_compile_args: Optional torch.compile configuration from settings node
|
||||
|
||||
Returns:
|
||||
NodeOutput containing configuration dictionary for SeedVR2 main node
|
||||
|
||||
Raises:
|
||||
ValueError: If BlockSwap is enabled but offload_device is invalid
|
||||
ValueError: If cache_model is enabled but offload_device is not set
|
||||
"""
|
||||
# Validate BlockSwap configuration
|
||||
if (blocks_to_swap > 0 or swap_io_components) and (offload_device == "none" or offload_device == device):
|
||||
raise ValueError(
|
||||
"BlockSwap requires offload_device to be set and different from device. "
|
||||
f"Current: device='{device}', offload_device='{offload_device}'. "
|
||||
"Please set offload_device to a different device (e.g., 'cpu' or another GPU)."
|
||||
)
|
||||
|
||||
# Validate cache_model configuration
|
||||
if cache_model and offload_device == "none":
|
||||
raise ValueError(
|
||||
|
||||
@@ -204,18 +204,6 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
|
||||
"• 'cuda:X': Offload to another GPU (good balance if available, faster than CPU)"
|
||||
)
|
||||
),
|
||||
io.Combo.Input("precision",
|
||||
options=["auto", "fp16", "bf16", "bf32"],
|
||||
default="auto",
|
||||
optional=True,
|
||||
tooltip=(
|
||||
"Precision for main generation process (default: auto).\n"
|
||||
"• auto: Automatically select based on device capabilities\n"
|
||||
"• fp16: Half precision (fastest, standard)\n"
|
||||
"• bf16: BFloat16 (better dynamic range, requires Ampere+ GPU)\n"
|
||||
"• bf32: Float32 with TF32 enabled (Ampere+ GPU)"
|
||||
)
|
||||
),
|
||||
io.Boolean.Input("enable_debug",
|
||||
default=False,
|
||||
optional=True,
|
||||
@@ -239,7 +227,7 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
|
||||
uniform_batch_size: bool = False, temporal_overlap: int = 0, prepend_frames: int = 0,
|
||||
color_correction: str = "wavelet", input_noise_scale: float = 0.0,
|
||||
latent_noise_scale: float = 0.0, offload_device: str = "none",
|
||||
precision: str = "auto", enable_debug: bool = False) -> io.NodeOutput:
|
||||
enable_debug: bool = False) -> io.NodeOutput:
|
||||
"""
|
||||
Execute SeedVR2 video upscaling with progress reporting
|
||||
|
||||
@@ -354,10 +342,6 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
|
||||
attention_mode = dit.get("attention_mode", "sdpa")
|
||||
vae_cache = vae.get("cache_model", False)
|
||||
|
||||
# Override attention mode if specified in dit config but allow validation later
|
||||
if "attention_mode" in dit:
|
||||
attention_mode = dit["attention_mode"]
|
||||
|
||||
# BlockSwap configuration - construct from individual values
|
||||
blocks_to_swap = dit.get("blocks_to_swap", 0)
|
||||
swap_io_components = dit.get("swap_io_components", False)
|
||||
@@ -365,13 +349,12 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
|
||||
|
||||
block_swap_config = None
|
||||
if blocks_to_swap > 0 or swap_io_components:
|
||||
# Convert offload device string to torch.device for BlockSwap
|
||||
block_swap_config = {
|
||||
"blocks_to_swap": blocks_to_swap,
|
||||
"swap_io_components": swap_io_components,
|
||||
}
|
||||
if dit_offload_str != "none":
|
||||
block_swap_config = {
|
||||
"blocks_to_swap": blocks_to_swap,
|
||||
"swap_io_components": swap_io_components,
|
||||
"offload_device": torch.device(dit_offload_str)
|
||||
}
|
||||
block_swap_config["offload_device"] = torch.device(dit_offload_str)
|
||||
|
||||
# Device configuration for offloading - convert "none" to None, else torch.device
|
||||
vae_offload_str = vae.get("offload_device", "none")
|
||||
@@ -447,7 +430,6 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
|
||||
decode_tile_overlap=(decode_tile_overlap, decode_tile_overlap),
|
||||
tile_debug=tile_debug,
|
||||
attention_mode=attention_mode,
|
||||
precision=precision,
|
||||
torch_compile_args_dit=dit_torch_compile_args,
|
||||
torch_compile_args_vae=vae_torch_compile_args
|
||||
)
|
||||
|
||||
@@ -15,16 +15,14 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
# Import flash_attn with automatic fallback from compatibility layer
|
||||
from ...optimization.compatibility import flash_attn_varlen_func, FLASH_ATTN_AVAILABLE, SAGE_ATTN_AVAILABLE
|
||||
# Import flash/sage attn with automatic fallback from compatibility layer
|
||||
from ...optimization.compatibility import (
|
||||
call_flash_attn_2_varlen, call_flash_attn_3_varlen,
|
||||
call_sage_attn_2_varlen, call_sage_attn_3_varlen
|
||||
)
|
||||
|
||||
from torch import nn
|
||||
|
||||
# Safe import for SageAttention
|
||||
try:
|
||||
import sageattention
|
||||
except ImportError:
|
||||
sageattention = None
|
||||
|
||||
def pytorch_varlen_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q=None, max_seqlen_k=None, dropout_p=0.0, softmax_scale=None, causal=False, deterministic=False):
|
||||
"""
|
||||
@@ -66,81 +64,6 @@ def pytorch_varlen_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q=N
|
||||
return torch.cat(output_splits, dim=0)
|
||||
|
||||
|
||||
@torch._dynamo.disable
|
||||
def _call_flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
|
||||
"""
|
||||
Wrapper for flash_attn_varlen_func that handles tensor-to-scalar conversion.
|
||||
|
||||
This function is excluded from torch.compile because:
|
||||
1. flash_attn is a C++ extension that can't be compiled anyway
|
||||
2. It requires Python int scalars for max_seqlen parameters
|
||||
3. Disabling compilation here keeps the rest of the model compilable
|
||||
"""
|
||||
if not FLASH_ATTN_AVAILABLE:
|
||||
raise ImportError("flash_attn is not available")
|
||||
|
||||
# Convert tensor max_seqlen to Python int if needed
|
||||
if torch.is_tensor(max_seqlen_q):
|
||||
max_seqlen_q = int(max_seqlen_q.item())
|
||||
if torch.is_tensor(max_seqlen_k):
|
||||
max_seqlen_k = int(max_seqlen_k.item())
|
||||
|
||||
return flash_attn_varlen_func(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_k=cu_seqlens_k,
|
||||
max_seqlen_q=max_seqlen_q,
|
||||
max_seqlen_k=max_seqlen_k,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
@torch._dynamo.disable
|
||||
def _call_sage_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, is_causal=False, implementation="sa2"):
|
||||
"""
|
||||
Wrapper for SageAttention variable length function.
|
||||
|
||||
Args:
|
||||
implementation: "sa2" (SageAttention v2) or "sa3" (SageAttention v3)
|
||||
"""
|
||||
if not SAGE_ATTN_AVAILABLE:
|
||||
raise ImportError("SageAttention is not available")
|
||||
|
||||
# SageAttention expects q, k, v as (total_tokens, heads, head_dim)
|
||||
# The input q, k, v here are (total_tokens, heads, head_dim)
|
||||
|
||||
# Convert tensor max_seqlen to Python int if needed
|
||||
if torch.is_tensor(max_seqlen_q):
|
||||
max_seqlen_q = int(max_seqlen_q.item())
|
||||
if torch.is_tensor(max_seqlen_k):
|
||||
max_seqlen_k = int(max_seqlen_k.item())
|
||||
|
||||
# Ensure tensors are contiguous
|
||||
q = q.contiguous()
|
||||
k = k.contiguous()
|
||||
v = v.contiguous()
|
||||
|
||||
# SageAttention API usage
|
||||
# sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, is_causal, sm_scale)
|
||||
try:
|
||||
from sageattention import sageattn_varlen
|
||||
except ImportError:
|
||||
# Fallback or error
|
||||
raise ImportError("sageattn_varlen not found in sageattention package")
|
||||
|
||||
# Check if sm_scale is needed (usually 1/sqrt(head_dim))
|
||||
sm_scale = 1.0 / (q.shape[-1] ** 0.5)
|
||||
|
||||
# Calling sageattn_varlen
|
||||
# Signature assumptions: q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, is_causal, sm_scale
|
||||
if not hasattr(_call_sage_attn_varlen_func, "_logged"):
|
||||
print(f"🚀 Executing SageAttention ({implementation}) kernel for the first time")
|
||||
_call_sage_attn_varlen_func._logged = True
|
||||
|
||||
return sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, is_causal, sm_scale)
|
||||
|
||||
|
||||
class TorchAttention(nn.Module):
|
||||
def tflops(self, args, kwargs, output) -> float:
|
||||
assert len(args) == 0 or len(args) > 2, "query, key should both provided by args / kwargs"
|
||||
@@ -156,12 +79,16 @@ class TorchAttention(nn.Module):
|
||||
|
||||
class FlashAttentionVarlen(nn.Module):
|
||||
"""
|
||||
Variable-length attention with configurable backend (Flash Attention or PyTorch SDPA).
|
||||
Variable-length attention with configurable backend.
|
||||
|
||||
Backend selection is validated during model configuration.
|
||||
Compilation behavior:
|
||||
- SDPA: Fully compilable, optimal performance
|
||||
- Flash Attention: Uses @torch._dynamo.disable wrapper (C++ extension)
|
||||
Supported backends:
|
||||
- sdpa: PyTorch SDPA (fully compilable, always available)
|
||||
- flash_attn_2: Flash Attention 2 (Ampere+)
|
||||
- flash_attn_3: Flash Attention 3 (Hopper+)
|
||||
- sageattn_2: SageAttention 2
|
||||
- sageattn_3: SageAttention 3 (Blackwell/RTX 50xx)
|
||||
|
||||
All non-SDPA backends use @torch._dynamo.disable wrapper (C++ extensions).
|
||||
"""
|
||||
|
||||
def __init__(self, attention_mode: str = 'sdpa', compute_dtype: torch.dtype = None):
|
||||
@@ -169,7 +96,7 @@ class FlashAttentionVarlen(nn.Module):
|
||||
Initialize with specified attention backend.
|
||||
|
||||
Args:
|
||||
attention_mode: 'flash_attn' or 'sdpa' (validated externally by validate_flash_attention_availability)
|
||||
attention_mode: 'sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3'
|
||||
compute_dtype: Compute dtype for attention (set by pipeline, defaults to None for auto-detection)
|
||||
"""
|
||||
super().__init__()
|
||||
@@ -193,20 +120,25 @@ class FlashAttentionVarlen(nn.Module):
|
||||
k = k.to(self.compute_dtype)
|
||||
v = v.to(self.compute_dtype)
|
||||
|
||||
if self.attention_mode == 'flash_attn':
|
||||
return _call_flash_attn_varlen_func(
|
||||
if self.attention_mode == 'flash_attn_3':
|
||||
return call_flash_attn_3_varlen(
|
||||
q, k, v, cu_seqlens_q, cu_seqlens_k,
|
||||
max_seqlen_q, max_seqlen_k, **kwargs
|
||||
)
|
||||
elif self.attention_mode in ['sa2', 'sa3']:
|
||||
# Use SageAttention
|
||||
# Extract causal flag if present in kwargs, default to False
|
||||
is_causal = kwargs.get('causal', False)
|
||||
return _call_sage_attn_varlen_func(
|
||||
elif self.attention_mode == 'flash_attn_2':
|
||||
return call_flash_attn_2_varlen(
|
||||
q, k, v, cu_seqlens_q, cu_seqlens_k,
|
||||
max_seqlen_q, max_seqlen_k, **kwargs
|
||||
)
|
||||
elif self.attention_mode == 'sageattn_3':
|
||||
return call_sage_attn_3_varlen(
|
||||
q, k, v, cu_seqlens_q, cu_seqlens_k,
|
||||
max_seqlen_q, max_seqlen_k,
|
||||
is_causal=is_causal,
|
||||
implementation=self.attention_mode
|
||||
max_seqlen_q, max_seqlen_k, **kwargs
|
||||
)
|
||||
elif self.attention_mode == 'sageattn_2':
|
||||
return call_sage_attn_2_varlen(
|
||||
q, k, v, cu_seqlens_q, cu_seqlens_k,
|
||||
max_seqlen_q, max_seqlen_k, **kwargs
|
||||
)
|
||||
else:
|
||||
# PyTorch SDPA
|
||||
|
||||
@@ -15,16 +15,14 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
# Import flash_attn with automatic fallback from compatibility layer
|
||||
from ...optimization.compatibility import flash_attn_varlen_func, FLASH_ATTN_AVAILABLE, SAGE_ATTN_AVAILABLE
|
||||
# Import flash/sage attn with automatic fallback from compatibility layer
|
||||
from ...optimization.compatibility import (
|
||||
call_flash_attn_2_varlen, call_flash_attn_3_varlen,
|
||||
call_sage_attn_2_varlen, call_sage_attn_3_varlen
|
||||
)
|
||||
|
||||
from torch import nn
|
||||
|
||||
# Safe import for SageAttention
|
||||
try:
|
||||
import sageattention
|
||||
except ImportError:
|
||||
sageattention = None
|
||||
|
||||
def pytorch_varlen_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q=None, max_seqlen_k=None, dropout_p=0.0, softmax_scale=None, causal=False, deterministic=False):
|
||||
"""
|
||||
@@ -66,81 +64,6 @@ def pytorch_varlen_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q=N
|
||||
return torch.cat(output_splits, dim=0)
|
||||
|
||||
|
||||
@torch._dynamo.disable
|
||||
def _call_flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
|
||||
"""
|
||||
Wrapper for flash_attn_varlen_func that handles tensor-to-scalar conversion.
|
||||
|
||||
This function is excluded from torch.compile because:
|
||||
1. flash_attn is a C++ extension that can't be compiled anyway
|
||||
2. It requires Python int scalars for max_seqlen parameters
|
||||
3. Disabling compilation here keeps the rest of the model compilable
|
||||
"""
|
||||
if not FLASH_ATTN_AVAILABLE:
|
||||
raise ImportError("flash_attn is not available")
|
||||
|
||||
# Convert tensor max_seqlen to Python int if needed
|
||||
if torch.is_tensor(max_seqlen_q):
|
||||
max_seqlen_q = int(max_seqlen_q.item())
|
||||
if torch.is_tensor(max_seqlen_k):
|
||||
max_seqlen_k = int(max_seqlen_k.item())
|
||||
|
||||
return flash_attn_varlen_func(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_k=cu_seqlens_k,
|
||||
max_seqlen_q=max_seqlen_q,
|
||||
max_seqlen_k=max_seqlen_k,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
@torch._dynamo.disable
|
||||
def _call_sage_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, is_causal=False, implementation="sa2"):
|
||||
"""
|
||||
Wrapper for SageAttention variable length function.
|
||||
|
||||
Args:
|
||||
implementation: "sa2" (SageAttention v2) or "sa3" (SageAttention v3)
|
||||
"""
|
||||
if not SAGE_ATTN_AVAILABLE:
|
||||
raise ImportError("SageAttention is not available")
|
||||
|
||||
# SageAttention expects q, k, v as (total_tokens, heads, head_dim)
|
||||
# The input q, k, v here are (total_tokens, heads, head_dim)
|
||||
|
||||
# Convert tensor max_seqlen to Python int if needed
|
||||
if torch.is_tensor(max_seqlen_q):
|
||||
max_seqlen_q = int(max_seqlen_q.item())
|
||||
if torch.is_tensor(max_seqlen_k):
|
||||
max_seqlen_k = int(max_seqlen_k.item())
|
||||
|
||||
# Ensure tensors are contiguous
|
||||
q = q.contiguous()
|
||||
k = k.contiguous()
|
||||
v = v.contiguous()
|
||||
|
||||
# SageAttention API usage
|
||||
# sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, is_causal, sm_scale)
|
||||
try:
|
||||
from sageattention import sageattn_varlen
|
||||
except ImportError:
|
||||
# Fallback or error
|
||||
raise ImportError("sageattn_varlen not found in sageattention package")
|
||||
|
||||
# Check if sm_scale is needed (usually 1/sqrt(head_dim))
|
||||
sm_scale = 1.0 / (q.shape[-1] ** 0.5)
|
||||
|
||||
# Calling sageattn_varlen
|
||||
# Signature assumptions: q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, is_causal, sm_scale
|
||||
if not hasattr(_call_sage_attn_varlen_func, "_logged"):
|
||||
print(f"🚀 Executing SageAttention ({implementation}) kernel for the first time")
|
||||
_call_sage_attn_varlen_func._logged = True
|
||||
|
||||
return sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, is_causal, sm_scale)
|
||||
|
||||
|
||||
class TorchAttention(nn.Module):
|
||||
def tflops(self, args, kwargs, output) -> float:
|
||||
assert len(args) == 0 or len(args) > 2, "query, key should both provided by args / kwargs"
|
||||
@@ -156,12 +79,16 @@ class TorchAttention(nn.Module):
|
||||
|
||||
class FlashAttentionVarlen(nn.Module):
|
||||
"""
|
||||
Variable-length attention with configurable backend (Flash Attention or PyTorch SDPA).
|
||||
Variable-length attention with configurable backend.
|
||||
|
||||
Backend selection is validated during model configuration.
|
||||
Compilation behavior:
|
||||
- SDPA: Fully compilable, optimal performance
|
||||
- Flash Attention: Uses @torch._dynamo.disable wrapper (C++ extension)
|
||||
Supported backends:
|
||||
- sdpa: PyTorch SDPA (fully compilable, always available)
|
||||
- flash_attn_2: Flash Attention 2 (Ampere+)
|
||||
- flash_attn_3: Flash Attention 3 (Hopper+)
|
||||
- sageattn_2: SageAttention 2
|
||||
- sageattn_3: SageAttention 3 (Blackwell/RTX 50xx)
|
||||
|
||||
All non-SDPA backends use @torch._dynamo.disable wrapper (C++ extensions).
|
||||
"""
|
||||
|
||||
def __init__(self, attention_mode: str = 'sdpa', compute_dtype: torch.dtype = None):
|
||||
@@ -169,7 +96,7 @@ class FlashAttentionVarlen(nn.Module):
|
||||
Initialize with specified attention backend.
|
||||
|
||||
Args:
|
||||
attention_mode: 'flash_attn' or 'sdpa' (validated externally by validate_flash_attention_availability)
|
||||
attention_mode: 'sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3'
|
||||
compute_dtype: Compute dtype for attention (set by pipeline, defaults to None for auto-detection)
|
||||
"""
|
||||
super().__init__()
|
||||
@@ -193,20 +120,25 @@ class FlashAttentionVarlen(nn.Module):
|
||||
k = k.to(self.compute_dtype)
|
||||
v = v.to(self.compute_dtype)
|
||||
|
||||
if self.attention_mode == 'flash_attn':
|
||||
return _call_flash_attn_varlen_func(
|
||||
if self.attention_mode == 'flash_attn_3':
|
||||
return call_flash_attn_3_varlen(
|
||||
q, k, v, cu_seqlens_q, cu_seqlens_k,
|
||||
max_seqlen_q, max_seqlen_k, **kwargs
|
||||
)
|
||||
elif self.attention_mode in ['sa2', 'sa3']:
|
||||
# Use SageAttention
|
||||
# Extract causal flag if present in kwargs, default to False
|
||||
is_causal = kwargs.get('causal', False)
|
||||
return _call_sage_attn_varlen_func(
|
||||
elif self.attention_mode == 'flash_attn_2':
|
||||
return call_flash_attn_2_varlen(
|
||||
q, k, v, cu_seqlens_q, cu_seqlens_k,
|
||||
max_seqlen_q, max_seqlen_k, **kwargs
|
||||
)
|
||||
elif self.attention_mode == 'sageattn_3':
|
||||
return call_sage_attn_3_varlen(
|
||||
q, k, v, cu_seqlens_q, cu_seqlens_k,
|
||||
max_seqlen_q, max_seqlen_k,
|
||||
is_causal=is_causal,
|
||||
implementation=self.attention_mode
|
||||
max_seqlen_q, max_seqlen_k, **kwargs
|
||||
)
|
||||
elif self.attention_mode == 'sageattn_2':
|
||||
return call_sage_attn_2_varlen(
|
||||
q, k, v, cu_seqlens_q, cu_seqlens_k,
|
||||
max_seqlen_q, max_seqlen_k, **kwargs
|
||||
)
|
||||
else:
|
||||
# PyTorch SDPA
|
||||
|
||||
@@ -47,6 +47,78 @@ def is_blockswap_enabled(config: Optional[Dict[str, Any]]) -> bool:
|
||||
return blocks_to_swap > 0 or swap_io_components
|
||||
|
||||
|
||||
def validate_blockswap_config(
|
||||
block_swap_config: Optional[Dict[str, Any]],
|
||||
dit_device: 'torch.device',
|
||||
dit_offload_device: Optional['torch.device'],
|
||||
debug: 'Debug'
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Validate and potentially modify BlockSwap configuration.
|
||||
|
||||
Performs platform-specific validation and configuration adjustment:
|
||||
- On macOS (MPS): Auto-disables BlockSwap since unified memory makes it meaningless
|
||||
- On other platforms: Validates that offload_device is properly configured
|
||||
|
||||
This is the single authoritative validation point for BlockSwap configuration,
|
||||
called early in configure_runner() before any model loading.
|
||||
|
||||
Args:
|
||||
block_swap_config: BlockSwap configuration dictionary (may be None)
|
||||
dit_device: Target device for DiT model inference
|
||||
dit_offload_device: Device for offloading DiT blocks (may be None)
|
||||
debug: Debug instance for logging warnings/errors
|
||||
|
||||
Returns:
|
||||
Validated/modified block_swap_config (may be None or modified copy)
|
||||
|
||||
Raises:
|
||||
ValueError: If BlockSwap is enabled but offload_device is invalid (non-MPS only)
|
||||
"""
|
||||
if not is_blockswap_enabled(block_swap_config):
|
||||
return block_swap_config
|
||||
|
||||
blocks_to_swap = block_swap_config.get("blocks_to_swap", 0)
|
||||
swap_io_components = block_swap_config.get("swap_io_components", False)
|
||||
|
||||
# Check for macOS unified memory - BlockSwap is meaningless there
|
||||
if dit_device.type == "mps":
|
||||
debug.log(
|
||||
f"BlockSwap disabled: macOS uses unified memory (no separate VRAM/RAM). "
|
||||
f"Ignoring blocks_to_swap={blocks_to_swap}, swap_io_components={swap_io_components}",
|
||||
level="WARNING", category="blockswap", force=True
|
||||
)
|
||||
# Return disabled config
|
||||
return {
|
||||
**block_swap_config,
|
||||
"blocks_to_swap": 0,
|
||||
"swap_io_components": False
|
||||
}
|
||||
|
||||
# Validate offload_device is set and different from dit_device
|
||||
offload_device_valid = (
|
||||
dit_offload_device is not None and
|
||||
str(dit_offload_device) != str(dit_device)
|
||||
)
|
||||
|
||||
if not offload_device_valid:
|
||||
config_details = []
|
||||
if blocks_to_swap > 0:
|
||||
config_details.append(f"blocks_to_swap={blocks_to_swap}")
|
||||
if swap_io_components:
|
||||
config_details.append("swap_io_components=True")
|
||||
|
||||
offload_str = str(dit_offload_device) if dit_offload_device else "none"
|
||||
raise ValueError(
|
||||
f"BlockSwap enabled ({', '.join(config_details)}) but dit_offload_device is invalid. "
|
||||
f"Current: device='{dit_device}', dit_offload_device='{offload_str}'. "
|
||||
f"BlockSwap requires offload_device on the DiT Model to be set and different from device. "
|
||||
f"Set --dit_offload_device cpu or disable BlockSwap."
|
||||
)
|
||||
|
||||
return block_swap_config
|
||||
|
||||
|
||||
# Timing helpers marked to skip torch.compile tracing
|
||||
# These functions are excluded from Dynamo's graph tracing to avoid warnings
|
||||
# about non-traceable builtins like time.time(), but they still execute normally
|
||||
@@ -125,7 +197,7 @@ def apply_block_swap_to_dit(
|
||||
|
||||
debug.start_timer("apply_blockswap")
|
||||
|
||||
# Get the actual model (handle FP8CompatibleDiT wrapper)
|
||||
# Get the actual model (handle CompatibleDiT wrapper)
|
||||
model = runner.dit
|
||||
if hasattr(model, "dit_model"):
|
||||
model = model.dit_model
|
||||
@@ -723,7 +795,7 @@ def set_blockswap_bypass(runner, bypass: bool, debug):
|
||||
if not hasattr(runner, "_blockswap_active") or not runner._blockswap_active:
|
||||
return
|
||||
|
||||
# Get the actual model (handle FP8CompatibleDiT wrapper)
|
||||
# Get the actual model (handle CompatibleDiT wrapper)
|
||||
model = runner.dit
|
||||
if hasattr(model, "dit_model"):
|
||||
model = model.dit_model
|
||||
@@ -758,7 +830,7 @@ def cleanup_blockswap(runner, keep_state_for_cache=False):
|
||||
|
||||
debug = runner.debug
|
||||
|
||||
# Get the actual model (handle FP8CompatibleDiT wrapper)
|
||||
# Get the actual model (handle CompatibleDiT wrapper)
|
||||
model = runner.dit
|
||||
if hasattr(model, "dit_model"):
|
||||
model = model.dit_model
|
||||
@@ -786,7 +858,7 @@ def cleanup_blockswap(runner, keep_state_for_cache=False):
|
||||
return
|
||||
|
||||
# Full cleanup when not caching
|
||||
# Get the actual model (handle FP8CompatibleDiT wrapper)
|
||||
# Get the actual model (handle CompatibleDiT wrapper)
|
||||
model = runner.dit
|
||||
if hasattr(model, "dit_model"):
|
||||
model = model.dit_model
|
||||
|
||||
+531
-358
@@ -5,9 +5,11 @@ 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
|
||||
# Compatibility shims - Must run before any torch/diffusers import
|
||||
import sys
|
||||
import types
|
||||
import importlib.machinery
|
||||
|
||||
|
||||
def ensure_triton_compat():
|
||||
"""Create minimal triton.ops stubs only if missing, to allow bitsandbytes import."""
|
||||
@@ -20,8 +22,6 @@ def ensure_triton_compat():
|
||||
except (ImportError, ModuleNotFoundError, AttributeError):
|
||||
pass
|
||||
|
||||
import types
|
||||
|
||||
if 'triton.ops' not in sys.modules:
|
||||
sys.modules['triton.ops'] = types.ModuleType('triton.ops')
|
||||
|
||||
@@ -32,102 +32,519 @@ def ensure_triton_compat():
|
||||
sys.modules['triton.ops'].matmul_perf_model = matmul_perf
|
||||
sys.modules['triton.ops.matmul_perf_model'] = matmul_perf
|
||||
|
||||
# Run immediately on import
|
||||
|
||||
def ensure_flash_attn_safe():
|
||||
"""
|
||||
Pre-test flash_attn package; stub if DLL is broken.
|
||||
Prevents diffusers from crashing when flash_attn has broken DLLs.
|
||||
"""
|
||||
if 'flash_attn' in sys.modules:
|
||||
return # Already loaded
|
||||
|
||||
try:
|
||||
import flash_attn
|
||||
except (ImportError, OSError):
|
||||
# DLL broken or not installed - create stub with proper __spec__
|
||||
stub = types.ModuleType('flash_attn')
|
||||
stub.__spec__ = importlib.machinery.ModuleSpec('flash_attn', None)
|
||||
stub.__file__ = None
|
||||
stub.__path__ = []
|
||||
stub.__loader__ = None
|
||||
# Provide attributes that diffusers/transformers import
|
||||
stub.flash_attn_func = None
|
||||
stub.flash_attn_varlen_func = None
|
||||
sys.modules['flash_attn'] = stub
|
||||
|
||||
|
||||
def ensure_xformers_flash_compat():
|
||||
"""
|
||||
Pre-test xformers._C_flashattention; stub if DLL is broken.
|
||||
Prevents xformers.ops.fmha.flash from crashing on import.
|
||||
"""
|
||||
if 'xformers._C_flashattention' in sys.modules:
|
||||
return # Already loaded
|
||||
|
||||
try:
|
||||
from xformers import _C_flashattention # noqa: F401
|
||||
except (ImportError, OSError):
|
||||
# DLL broken or not installed - create stub with proper __spec__
|
||||
class _FailingStub(types.ModuleType):
|
||||
"""Stub that lets xformers gracefully disable its flash backend."""
|
||||
def __getattr__(self, name):
|
||||
raise ImportError("_C_flashattention unavailable")
|
||||
|
||||
stub = _FailingStub('xformers._C_flashattention')
|
||||
stub.__spec__ = importlib.machinery.ModuleSpec('xformers._C_flashattention', None)
|
||||
stub.__file__ = None
|
||||
stub.__path__ = []
|
||||
stub.__loader__ = None
|
||||
sys.modules['xformers._C_flashattention'] = stub
|
||||
|
||||
|
||||
def ensure_bitsandbytes_safe():
|
||||
"""
|
||||
Pre-test bitsandbytes; stub if broken to prevent import conflicts.
|
||||
|
||||
On some systems (e.g., ROCm without proper binaries), bitsandbytes registers
|
||||
PyTorch kernels during import then fails. If another node already triggered
|
||||
this partial load, re-importing causes kernel registration conflicts.
|
||||
|
||||
This shim catches such failures and stubs the module so diffusers can load
|
||||
gracefully without bitsandbytes quantization support.
|
||||
"""
|
||||
if 'bitsandbytes' in sys.modules:
|
||||
return # Already loaded or stubbed
|
||||
|
||||
try:
|
||||
import bitsandbytes
|
||||
# Success - bitsandbytes works, other nodes can use it
|
||||
except (ImportError, OSError, RuntimeError):
|
||||
# Installation broken or not present - create stub
|
||||
stub = types.ModuleType('bitsandbytes')
|
||||
stub.__spec__ = importlib.machinery.ModuleSpec('bitsandbytes', None)
|
||||
stub.__file__ = None
|
||||
stub.__path__ = []
|
||||
stub.__version__ = "0.0.0"
|
||||
sys.modules['bitsandbytes'] = stub
|
||||
|
||||
|
||||
# Run all shims immediately on import, before torch/diffusers
|
||||
ensure_triton_compat()
|
||||
ensure_flash_attn_safe()
|
||||
ensure_xformers_flash_compat()
|
||||
ensure_bitsandbytes_safe()
|
||||
|
||||
|
||||
import torch
|
||||
import types
|
||||
import os
|
||||
from typing import Dict, Any, Optional
|
||||
|
||||
|
||||
# Flash Attention & Triton Compatibility Layer
|
||||
# 1. Flash Attention - speedup for attention operations
|
||||
# Flash/Sage Attention & Triton Compatibility Layer
|
||||
|
||||
# 1. Flash Attention 3 (Hopper+, faster, no dropout/window support)
|
||||
flash_attn_3_varlen_func = None
|
||||
FLASH_ATTN_3_AVAILABLE = False
|
||||
try:
|
||||
from flash_attn import flash_attn_varlen_func
|
||||
FLASH_ATTN_AVAILABLE = True
|
||||
except ImportError:
|
||||
flash_attn_varlen_func = None
|
||||
FLASH_ATTN_AVAILABLE = False
|
||||
import flash_attn_interface
|
||||
flash_attn_3_varlen_func = flash_attn_interface.flash_attn_varlen_func
|
||||
FLASH_ATTN_3_AVAILABLE = True
|
||||
except (ImportError, AttributeError, OSError):
|
||||
pass
|
||||
|
||||
# 1.1 SageAttention - speedup for attention operations
|
||||
# 2. Flash Attention 2 (wider compatibility, supports dropout/window)
|
||||
flash_attn_2_varlen_func = None
|
||||
FLASH_ATTN_2_AVAILABLE = False
|
||||
try:
|
||||
import sageattention
|
||||
SAGE_ATTN_AVAILABLE = True
|
||||
from flash_attn import flash_attn_varlen_func as _fa2_varlen
|
||||
import flash_attn_2_cuda # noqa: F401
|
||||
flash_attn_2_varlen_func = _fa2_varlen
|
||||
FLASH_ATTN_2_AVAILABLE = True
|
||||
except (ImportError, AttributeError, OSError):
|
||||
pass
|
||||
|
||||
FLASH_ATTN_AVAILABLE = FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE
|
||||
|
||||
# 3. SageAttention 2 (varlen support)
|
||||
sageattn_varlen = None
|
||||
SAGE_ATTN_2_AVAILABLE = False
|
||||
try:
|
||||
from sageattention import sageattn_varlen as _sa2_varlen
|
||||
sageattn_varlen = _sa2_varlen
|
||||
SAGE_ATTN_2_AVAILABLE = True
|
||||
except (ImportError, AttributeError, OSError):
|
||||
pass
|
||||
|
||||
# 4. SageAttention 3 / Blackwell (RTX 50xx only, batched attention)
|
||||
sageattn_blackwell = None
|
||||
SAGE_ATTN_3_AVAILABLE = False
|
||||
try:
|
||||
from sageattn3 import sageattn3_blackwell as _sa3_blackwell
|
||||
sageattn_blackwell = _sa3_blackwell
|
||||
SAGE_ATTN_3_AVAILABLE = True
|
||||
except (ImportError, AttributeError, OSError):
|
||||
try:
|
||||
from sageattention import sageattn_varlen
|
||||
# Basic check to see if it's functional or mock
|
||||
SAGE_ATTN_VARLEN_AVAILABLE = True
|
||||
except ImportError:
|
||||
SAGE_ATTN_VARLEN_AVAILABLE = False
|
||||
except ImportError:
|
||||
SAGE_ATTN_AVAILABLE = False
|
||||
SAGE_ATTN_VARLEN_AVAILABLE = False
|
||||
from sageattention import sageattn_blackwell as _sa3_blackwell
|
||||
sageattn_blackwell = _sa3_blackwell
|
||||
SAGE_ATTN_3_AVAILABLE = True
|
||||
except (ImportError, AttributeError, OSError):
|
||||
pass
|
||||
|
||||
SAGE_ATTN_AVAILABLE = SAGE_ATTN_2_AVAILABLE or SAGE_ATTN_3_AVAILABLE
|
||||
|
||||
|
||||
def validate_flash_attention_availability(requested_mode: str, debug=None) -> str:
|
||||
def validate_attention_mode(requested_mode: str, debug=None) -> str:
|
||||
"""
|
||||
Validate attention mode availability and warn if fallback needed.
|
||||
Validate attention mode availability with automatic fallback.
|
||||
|
||||
Args:
|
||||
requested_mode: 'flash_attn', 'sdpa', 'sd2', or 'sd3'
|
||||
requested_mode: 'sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3'
|
||||
debug: Optional debug instance for logging
|
||||
|
||||
Returns:
|
||||
Validated mode
|
||||
Validated mode that is available
|
||||
"""
|
||||
if requested_mode == 'flash_attn' and not FLASH_ATTN_AVAILABLE:
|
||||
# Flash Attention 3
|
||||
if requested_mode == 'flash_attn_3':
|
||||
if FLASH_ATTN_3_AVAILABLE:
|
||||
return requested_mode
|
||||
if FLASH_ATTN_2_AVAILABLE:
|
||||
if debug:
|
||||
debug.log(
|
||||
"Flash Attention 3 not available (requires Hopper+ GPU and flash-attn with FA3 support).\n"
|
||||
"Falling back to Flash Attention 2.",
|
||||
level="WARNING", category="setup", force=True
|
||||
)
|
||||
return 'flash_attn_2'
|
||||
error_msg = (
|
||||
f"Cannot use 'flash_attn' attention mode: Flash Attention is not installed.\n"
|
||||
f"\n"
|
||||
f"Flash Attention provides speedup on some hardware through optimized CUDA kernels.\n"
|
||||
f"Falling back to PyTorch SDPA (scaled dot-product attention).\n"
|
||||
f"\n"
|
||||
f"To fix this issue:\n"
|
||||
f" 1. Install Flash Attention: pip install flash-attn\n"
|
||||
f" 2. OR change attention_mode to 'sdpa' (default, always available)\n"
|
||||
f"\n"
|
||||
f"For more info: https://github.com/Dao-AILab/flash-attention"
|
||||
"Cannot use 'flash_attn_3' attention mode: Flash Attention is not installed.\n"
|
||||
"\n"
|
||||
"Flash Attention 3 provides maximum speedup on Hopper+ GPUs through optimized CUDA kernels.\n"
|
||||
"Falling back to PyTorch SDPA (scaled dot-product attention).\n"
|
||||
"\n"
|
||||
"To fix this issue:\n"
|
||||
" 1. Install Flash Attention: pip install flash-attn\n"
|
||||
" 2. OR change attention_mode to 'sdpa' (default, always available)\n"
|
||||
"\n"
|
||||
"For more info: https://github.com/Dao-AILab/flash-attention"
|
||||
)
|
||||
if debug:
|
||||
debug.log(error_msg, level="WARNING", category="setup", force=True)
|
||||
|
||||
return 'sdpa'
|
||||
|
||||
if requested_mode in ['sa2', 'sa3']:
|
||||
if not SAGE_ATTN_AVAILABLE:
|
||||
|
||||
# Flash Attention 2
|
||||
if requested_mode == 'flash_attn_2':
|
||||
if FLASH_ATTN_2_AVAILABLE:
|
||||
return requested_mode
|
||||
error_msg = (
|
||||
"Cannot use 'flash_attn_2' attention mode: Flash Attention 2 is not installed.\n"
|
||||
"\n"
|
||||
"Flash Attention 2 provides speedup on Ampere+ GPUs through optimized CUDA kernels.\n"
|
||||
"Falling back to PyTorch SDPA (scaled dot-product attention).\n"
|
||||
"\n"
|
||||
"To fix this issue:\n"
|
||||
" 1. Install Flash Attention: pip install flash-attn\n"
|
||||
" 2. OR change attention_mode to 'sdpa' (default, always available)\n"
|
||||
"\n"
|
||||
"For more info: https://github.com/Dao-AILab/flash-attention"
|
||||
)
|
||||
if debug:
|
||||
debug.log(error_msg, level="WARNING", category="setup", force=True)
|
||||
return 'sdpa'
|
||||
|
||||
# SageAttention 3 (Blackwell)
|
||||
if requested_mode == 'sageattn_3':
|
||||
if SAGE_ATTN_3_AVAILABLE:
|
||||
return requested_mode
|
||||
if SAGE_ATTN_2_AVAILABLE:
|
||||
if debug:
|
||||
debug.log(f"SageAttention not installed. Falling back from '{requested_mode}' to Flash Attention 2...", level="WARNING", category="setup", force=True)
|
||||
# Fallback to check FA2
|
||||
return validate_flash_attention_availability('flash_attn', debug)
|
||||
|
||||
elif not SAGE_ATTN_VARLEN_AVAILABLE:
|
||||
if debug:
|
||||
debug.log(f"SageAttention installed but 'sageattn_varlen' not found. Falling back from '{requested_mode}' to Flash Attention 2...", level="WARNING", category="setup", force=True)
|
||||
# Fallback to check FA2
|
||||
return validate_flash_attention_availability('flash_attn', debug)
|
||||
|
||||
# If the user explicitly requested sa3, we check for version compatibility.
|
||||
# If version is unknown or insufficient, we fallback to sa2.
|
||||
if requested_mode == 'sa3':
|
||||
try:
|
||||
version = sageattention.__version__
|
||||
# Assuming sa3 requires at least a certain version or just presence of version string.
|
||||
# If we can read version, we assume it's compliant enough or user knows what they are doing.
|
||||
if debug:
|
||||
debug.log(f"SageAttention version {version} detected. Using installed kernel for 'sa3' mode.", category="setup", force=True)
|
||||
except AttributeError:
|
||||
# Version unknown -> Assume it's an older version (sa2) and fallback
|
||||
if debug:
|
||||
debug.log("SageAttention version unknown (likely v2 or older). Falling back from 'sa3' to 'sa2'...", level="WARNING", category="setup", force=True)
|
||||
return validate_flash_attention_availability('sa2', debug)
|
||||
|
||||
pass
|
||||
debug.log(
|
||||
"SageAttention 3 (Blackwell) not available (requires RTX 50xx GPU and sageattn3 package).\n"
|
||||
"Falling back to SageAttention 2.",
|
||||
level="WARNING", category="setup", force=True
|
||||
)
|
||||
return 'sageattn_2'
|
||||
error_msg = (
|
||||
"Cannot use 'sageattn_3' attention mode: SageAttention is not installed.\n"
|
||||
"\n"
|
||||
"SageAttention 3 provides maximum speedup on Blackwell (RTX 50xx) GPUs.\n"
|
||||
"Falling back to PyTorch SDPA (scaled dot-product attention).\n"
|
||||
"\n"
|
||||
"To fix this issue:\n"
|
||||
" 1. Install SageAttention: pip install sageattention\n"
|
||||
" 2. For SA3 Blackwell support: pip install sageattn3\n"
|
||||
" 3. OR change attention_mode to 'flash_attn_2' or 'sdpa'\n"
|
||||
"\n"
|
||||
"For more info: https://github.com/thu-ml/SageAttention"
|
||||
)
|
||||
if debug:
|
||||
debug.log(error_msg, level="WARNING", category="setup", force=True)
|
||||
return 'sdpa'
|
||||
|
||||
# SageAttention 2
|
||||
if requested_mode == 'sageattn_2':
|
||||
if SAGE_ATTN_2_AVAILABLE:
|
||||
return requested_mode
|
||||
error_msg = (
|
||||
"Cannot use 'sageattn_2' attention mode: SageAttention is not installed.\n"
|
||||
"\n"
|
||||
"SageAttention provides speedup on NVIDIA GPUs through optimized CUDA kernels.\n"
|
||||
"Falling back to PyTorch SDPA (scaled dot-product attention).\n"
|
||||
"\n"
|
||||
"To fix this issue:\n"
|
||||
" 1. Install SageAttention: pip install sageattention\n"
|
||||
" 2. OR change attention_mode to 'flash_attn_2' or 'sdpa'\n"
|
||||
"\n"
|
||||
"For more info: https://github.com/thu-ml/SageAttention"
|
||||
)
|
||||
if debug:
|
||||
debug.log(error_msg, level="WARNING", category="setup", force=True)
|
||||
return 'sdpa'
|
||||
|
||||
return requested_mode
|
||||
|
||||
|
||||
@torch._dynamo.disable
|
||||
def call_flash_attn_2_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
|
||||
"""
|
||||
Wrapper for Flash Attention 2 flash_attn_varlen_func that handles tensor-to-scalar conversion.
|
||||
|
||||
Flash Attention 2 supports dropout_p and window_size parameters.
|
||||
Works on Ampere+ GPUs (RTX 30xx, 40xx, A100, etc.).
|
||||
|
||||
This function is excluded from torch.compile because:
|
||||
1. flash_attn is a C++ extension that can't be compiled anyway
|
||||
2. It requires Python int scalars for max_seqlen parameters
|
||||
3. Disabling compilation here keeps the rest of the model compilable
|
||||
|
||||
Args:
|
||||
q: Query tensor (total_seq, heads, head_dim)
|
||||
k: Key tensor (total_seq, heads, head_dim)
|
||||
v: Value tensor (total_seq, heads, head_dim)
|
||||
cu_seqlens_q: Cumulative sequence lengths for queries
|
||||
cu_seqlens_k: Cumulative sequence lengths for keys
|
||||
max_seqlen_q: Maximum query sequence length (can be tensor or int)
|
||||
max_seqlen_k: Maximum key sequence length (can be tensor or int)
|
||||
**kwargs: Additional arguments (dropout_p, softmax_scale, causal, window_size, deterministic)
|
||||
|
||||
Returns:
|
||||
Attention output tensor (total_seq, heads, head_dim)
|
||||
"""
|
||||
if not FLASH_ATTN_2_AVAILABLE:
|
||||
raise ImportError("Flash Attention 2 is not available")
|
||||
|
||||
# Convert tensor max_seqlen to Python int if needed
|
||||
if torch.is_tensor(max_seqlen_q):
|
||||
max_seqlen_q = int(max_seqlen_q.item())
|
||||
if torch.is_tensor(max_seqlen_k):
|
||||
max_seqlen_k = int(max_seqlen_k.item())
|
||||
|
||||
return flash_attn_2_varlen_func(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_k=cu_seqlens_k,
|
||||
max_seqlen_q=max_seqlen_q,
|
||||
max_seqlen_k=max_seqlen_k,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
|
||||
@torch._dynamo.disable
|
||||
def call_flash_attn_3_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
|
||||
"""
|
||||
Wrapper for Flash Attention 3 flash_attn_varlen_func that handles tensor-to-scalar conversion.
|
||||
|
||||
Flash Attention 3 is faster than FA2 but does NOT support dropout_p and window_size.
|
||||
Works on Hopper+ GPUs (H100, etc.) - requires flash_attn_interface package.
|
||||
|
||||
This function is excluded from torch.compile because:
|
||||
1. flash_attn is a C++ extension that can't be compiled anyway
|
||||
2. It requires Python int scalars for max_seqlen parameters
|
||||
3. Disabling compilation here keeps the rest of the model compilable
|
||||
|
||||
Args:
|
||||
q: Query tensor (total_seq, heads, head_dim)
|
||||
k: Key tensor (total_seq, heads, head_dim)
|
||||
v: Value tensor (total_seq, heads, head_dim)
|
||||
cu_seqlens_q: Cumulative sequence lengths for queries
|
||||
cu_seqlens_k: Cumulative sequence lengths for keys
|
||||
max_seqlen_q: Maximum query sequence length (can be tensor or int)
|
||||
max_seqlen_k: Maximum key sequence length (can be tensor or int)
|
||||
**kwargs: Additional arguments (softmax_scale, causal, deterministic)
|
||||
Note: dropout_p and window_size are ignored (not supported by FA3)
|
||||
|
||||
Returns:
|
||||
Attention output tensor (total_seq, heads, head_dim)
|
||||
"""
|
||||
if not FLASH_ATTN_3_AVAILABLE:
|
||||
raise ImportError("Flash Attention 3 is not available")
|
||||
|
||||
# Convert tensor max_seqlen to Python int if needed
|
||||
if torch.is_tensor(max_seqlen_q):
|
||||
max_seqlen_q = int(max_seqlen_q.item())
|
||||
if torch.is_tensor(max_seqlen_k):
|
||||
max_seqlen_k = int(max_seqlen_k.item())
|
||||
|
||||
# FA3 doesn't support dropout_p and window_size - filter them out
|
||||
fa3_kwargs = {key: val for key, val in kwargs.items() if key not in ('dropout_p', 'window_size')}
|
||||
|
||||
# FA3 returns a tuple (output, softmax_lse), we only need output
|
||||
return flash_attn_3_varlen_func(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_k=cu_seqlens_k,
|
||||
max_seqlen_q=max_seqlen_q,
|
||||
max_seqlen_k=max_seqlen_k,
|
||||
seqused_q=None,
|
||||
seqused_k=None,
|
||||
**fa3_kwargs
|
||||
)[0]
|
||||
|
||||
|
||||
@torch._dynamo.disable
|
||||
def call_sage_attn_2_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
|
||||
"""
|
||||
Wrapper for SageAttention 2 sageattn_varlen that handles tensor-to-scalar conversion.
|
||||
|
||||
SageAttention 2 provides optimized attention for NVIDIA GPUs with native varlen support.
|
||||
Works on most modern NVIDIA GPUs.
|
||||
|
||||
This function is excluded from torch.compile because:
|
||||
1. SageAttention is a C++ extension that can't be compiled anyway
|
||||
2. It requires Python int scalars for max_seqlen parameters
|
||||
3. Disabling compilation here keeps the rest of the model compilable
|
||||
|
||||
Args:
|
||||
q: Query tensor (total_seq, heads, head_dim)
|
||||
k: Key tensor (total_seq, heads, head_dim)
|
||||
v: Value tensor (total_seq, heads, head_dim)
|
||||
cu_seqlens_q: Cumulative sequence lengths for queries
|
||||
cu_seqlens_k: Cumulative sequence lengths for keys
|
||||
max_seqlen_q: Maximum query sequence length (can be tensor or int)
|
||||
max_seqlen_k: Maximum key sequence length (can be tensor or int)
|
||||
**kwargs: Additional arguments (causal supported, others ignored)
|
||||
|
||||
Returns:
|
||||
Attention output tensor (total_seq, heads, head_dim)
|
||||
"""
|
||||
if not SAGE_ATTN_2_AVAILABLE:
|
||||
raise ImportError("SageAttention 2 is not available")
|
||||
|
||||
# Convert tensor max_seqlen to Python int if needed
|
||||
if torch.is_tensor(max_seqlen_q):
|
||||
max_seqlen_q = int(max_seqlen_q.item())
|
||||
if torch.is_tensor(max_seqlen_k):
|
||||
max_seqlen_k = int(max_seqlen_k.item())
|
||||
|
||||
# SageAttention requires half precision (fp16/bf16)
|
||||
out_dtype = q.dtype
|
||||
half_dtypes = (torch.float16, torch.bfloat16)
|
||||
|
||||
if not (q.dtype == k.dtype == v.dtype):
|
||||
k = k.to(q.dtype)
|
||||
v = v.to(q.dtype)
|
||||
|
||||
if q.dtype not in half_dtypes:
|
||||
q = q.to(torch.bfloat16)
|
||||
k = k.to(torch.bfloat16)
|
||||
v = v.to(torch.bfloat16)
|
||||
|
||||
is_causal = kwargs.get('causal', False)
|
||||
sm_scale = 1.0 / (q.shape[-1] ** 0.5)
|
||||
|
||||
out = sageattn_varlen(
|
||||
q, k, v,
|
||||
cu_seqlens_q, cu_seqlens_k,
|
||||
max_seqlen_q, max_seqlen_k,
|
||||
is_causal, sm_scale
|
||||
)
|
||||
|
||||
return out.to(out_dtype) if out.dtype != out_dtype else out
|
||||
|
||||
|
||||
@torch._dynamo.disable
|
||||
def call_sage_attn_3_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
|
||||
"""
|
||||
Wrapper for SageAttention 3 (Blackwell) that converts varlen format to batched format.
|
||||
|
||||
SageAttention 3 / Blackwell provides maximum performance on RTX 50xx series GPUs.
|
||||
However, it only supports batched attention (uniform sequence lengths), not varlen.
|
||||
|
||||
This wrapper detects uniform-length batches and reshapes accordingly.
|
||||
For variable-length sequences, it automatically falls back to SageAttention 2.
|
||||
|
||||
This function is excluded from torch.compile because:
|
||||
1. SageAttention is a C++ extension that can't be compiled anyway
|
||||
2. It requires Python int scalars for max_seqlen parameters
|
||||
3. The varlen-to-batched conversion involves dynamic shapes
|
||||
4. Disabling compilation here keeps the rest of the model compilable
|
||||
|
||||
Args:
|
||||
q: Query tensor (total_seq, heads, head_dim)
|
||||
k: Key tensor (total_seq, heads, head_dim)
|
||||
v: Value tensor (total_seq, heads, head_dim)
|
||||
cu_seqlens_q: Cumulative sequence lengths for queries
|
||||
cu_seqlens_k: Cumulative sequence lengths for keys
|
||||
max_seqlen_q: Maximum query sequence length (can be tensor or int)
|
||||
max_seqlen_k: Maximum key sequence length (can be tensor or int)
|
||||
**kwargs: Additional arguments (passed to SA2 fallback if needed)
|
||||
|
||||
Returns:
|
||||
Attention output tensor (total_seq, heads, head_dim)
|
||||
"""
|
||||
if not SAGE_ATTN_3_AVAILABLE:
|
||||
raise ImportError("SageAttention 3 (Blackwell) is not available")
|
||||
|
||||
# Convert tensor max_seqlen to Python int if needed
|
||||
if torch.is_tensor(max_seqlen_q):
|
||||
max_seqlen_q = int(max_seqlen_q.item())
|
||||
if torch.is_tensor(max_seqlen_k):
|
||||
max_seqlen_k = int(max_seqlen_k.item())
|
||||
|
||||
# Check if all sequences have uniform length (required for SA3 batched API)
|
||||
# SA3/Blackwell uses batched attention, not varlen, so we need uniform lengths
|
||||
seq_lens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1]
|
||||
seq_lens_k = cu_seqlens_k[1:] - cu_seqlens_k[:-1]
|
||||
|
||||
uniform_q = (seq_lens_q == seq_lens_q[0]).all()
|
||||
uniform_k = (seq_lens_k == seq_lens_k[0]).all()
|
||||
|
||||
if not (uniform_q and uniform_k):
|
||||
# Fall back to SA2 for variable-length sequences
|
||||
# This is expected behavior - SA3 Blackwell doesn't support varlen natively
|
||||
if SAGE_ATTN_2_AVAILABLE:
|
||||
return call_sage_attn_2_varlen(
|
||||
q, k, v, cu_seqlens_q, cu_seqlens_k,
|
||||
max_seqlen_q, max_seqlen_k, **kwargs
|
||||
)
|
||||
raise RuntimeError(
|
||||
"SageAttention 3 (Blackwell) requires uniform sequence lengths, "
|
||||
"and SageAttention 2 is not available as fallback. "
|
||||
"Please install sageattention package or use flash_attn/sdpa instead."
|
||||
)
|
||||
|
||||
# Extract batch dimensions
|
||||
batch_size = len(cu_seqlens_q) - 1
|
||||
seq_len_q = int(seq_lens_q[0].item())
|
||||
seq_len_k = int(seq_lens_k[0].item())
|
||||
heads = q.shape[1]
|
||||
dim = q.shape[2]
|
||||
|
||||
# SageAttention requires half precision (fp16/bf16)
|
||||
out_dtype = q.dtype
|
||||
half_dtypes = (torch.float16, torch.bfloat16)
|
||||
|
||||
if not (q.dtype == k.dtype == v.dtype):
|
||||
k = k.to(q.dtype)
|
||||
v = v.to(q.dtype)
|
||||
|
||||
if q.dtype not in half_dtypes:
|
||||
q = q.to(torch.bfloat16)
|
||||
k = k.to(torch.bfloat16)
|
||||
v = v.to(torch.bfloat16)
|
||||
|
||||
# Reshape varlen (total_seq, heads, dim) -> batched (batch, seq, heads, dim)
|
||||
q_batched = q.view(batch_size, seq_len_q, heads, dim)
|
||||
k_batched = k.view(batch_size, seq_len_k, heads, dim)
|
||||
v_batched = v.view(batch_size, seq_len_k, heads, dim)
|
||||
|
||||
# SA3/Blackwell expects (batch, heads, seq, dim) layout
|
||||
q_batched = q_batched.transpose(1, 2) # (batch, heads, seq, dim)
|
||||
k_batched = k_batched.transpose(1, 2)
|
||||
v_batched = v_batched.transpose(1, 2)
|
||||
|
||||
# Call SA3 Blackwell
|
||||
out = sageattn_blackwell(q_batched, k_batched, v_batched, per_block_mean=False)
|
||||
|
||||
# Reshape back to varlen format (total_seq, heads, dim)
|
||||
out = out.transpose(1, 2).reshape(-1, heads, dim).contiguous()
|
||||
|
||||
return out.to(out_dtype) if out.dtype != out_dtype else out
|
||||
|
||||
|
||||
# 2. Triton - Required for torch.compile with inductor backend
|
||||
try:
|
||||
import triton
|
||||
@@ -226,21 +643,34 @@ NVIDIA_CONV3D_MEMORY_BUG_WORKAROUND = _check_conv3d_memory_bug()
|
||||
if not os.environ.get("SEEDVR2_OPTIMIZATIONS_LOGGED"):
|
||||
os.environ["SEEDVR2_OPTIMIZATIONS_LOGGED"] = "1"
|
||||
|
||||
# Flash Attention & Triton status
|
||||
has_both = FLASH_ATTN_AVAILABLE and TRITON_AVAILABLE
|
||||
has_neither = not FLASH_ATTN_AVAILABLE and not TRITON_AVAILABLE
|
||||
# Build status strings
|
||||
sage_status = "✅" if SAGE_ATTN_AVAILABLE else "❌"
|
||||
flash_status = "✅" if FLASH_ATTN_AVAILABLE else "❌"
|
||||
triton_status = "✅" if TRITON_AVAILABLE else "❌"
|
||||
|
||||
if has_both:
|
||||
print("⚡ SeedVR2 optimizations check: Flash Attention ✅ | Triton ✅")
|
||||
elif has_neither:
|
||||
print("⚠️ SeedVR2 optimizations check: Flash Attention ❌ | Triton ❌")
|
||||
print("💡 For best performance: pip install flash-attn triton")
|
||||
elif FLASH_ATTN_AVAILABLE:
|
||||
print("⚡ SeedVR2 optimizations check: Flash Attention ✅ | Triton ❌")
|
||||
print("💡 Install Triton for torch.compile: pip install triton")
|
||||
else: # TRITON_AVAILABLE only
|
||||
print("⚠️ SeedVR2 optimizations check: Flash Attention ❌ | Triton ✅")
|
||||
print("💡 Install Flash Attention for faster inference: pip install flash-attn")
|
||||
# Count available optimizations
|
||||
available = [SAGE_ATTN_AVAILABLE, FLASH_ATTN_AVAILABLE, TRITON_AVAILABLE]
|
||||
num_available = sum(available)
|
||||
|
||||
if num_available == 3:
|
||||
print(f"⚡ SeedVR2 optimizations check: SageAttention {sage_status} | Flash Attention {flash_status} | Triton {triton_status}")
|
||||
elif num_available == 0:
|
||||
print(f"⚠️ SeedVR2 optimizations check: SageAttention {sage_status} | Flash Attention {flash_status} | Triton {triton_status}")
|
||||
print("💡 For best performance: pip install sageattention flash-attn triton")
|
||||
else:
|
||||
icon = "⚡" if num_available >= 2 else "⚠️ "
|
||||
print(f"{icon} SeedVR2 optimizations check: SageAttention {sage_status} | Flash Attention {flash_status} | Triton {triton_status}")
|
||||
|
||||
# Build install suggestions for missing packages
|
||||
missing = []
|
||||
if not SAGE_ATTN_AVAILABLE:
|
||||
missing.append("sageattention")
|
||||
if not FLASH_ATTN_AVAILABLE:
|
||||
missing.append("flash-attn")
|
||||
if not TRITON_AVAILABLE:
|
||||
missing.append("triton")
|
||||
if missing:
|
||||
print(f"💡 Optional: pip install {' '.join(missing)}")
|
||||
|
||||
# Conv3d workaround status (if applicable)
|
||||
if NVIDIA_CONV3D_MEMORY_BUG_WORKAROUND:
|
||||
@@ -266,83 +696,38 @@ def _probe_bfloat16_support() -> bool:
|
||||
BFLOAT16_SUPPORTED = _probe_bfloat16_support()
|
||||
COMPUTE_DTYPE = torch.bfloat16 if BFLOAT16_SUPPORTED else torch.float16
|
||||
|
||||
def log_system_capabilities(debug=None):
|
||||
"""Log installed attention backends and versions at startup."""
|
||||
if not debug:
|
||||
return
|
||||
|
||||
# SageAttention
|
||||
sa_status = "Available" if SAGE_ATTN_AVAILABLE else "Not Installed"
|
||||
if SAGE_ATTN_AVAILABLE:
|
||||
try:
|
||||
sa_version = sageattention.__version__
|
||||
sa_status += f" (v{sa_version})"
|
||||
except AttributeError:
|
||||
sa_status += " (Version Unknown)"
|
||||
|
||||
# FlashAttention
|
||||
fa_status = "Available" if FLASH_ATTN_AVAILABLE else "Not Installed"
|
||||
|
||||
# Triton
|
||||
triton_status = "Available" if TRITON_AVAILABLE else "Not Installed"
|
||||
|
||||
debug.log(f"Attention Backends: SageAttention={sa_status} | FlashAttention={fa_status} | Triton={triton_status}", category="info", force=True)
|
||||
|
||||
def detect_high_end_system() -> Dict[str, Any]:
|
||||
"""
|
||||
Detect high-end systems (16GB+ VRAM, etc.) and return optimized defaults.
|
||||
|
||||
Returns:
|
||||
Dict with recommended settings or empty if no specific optimizations found.
|
||||
"""
|
||||
optimizations = {}
|
||||
try:
|
||||
# Basic VRAM check
|
||||
if torch.cuda.is_available():
|
||||
device = torch.device("cuda:0")
|
||||
props = torch.cuda.get_device_properties(device)
|
||||
total_vram_gb = props.total_memory / (1024**3)
|
||||
|
||||
# High-end GPU check (e.g., 5070ti/4080/4090/etc with >15GB VRAM)
|
||||
if total_vram_gb >= 15.5:
|
||||
optimizations['high_vram'] = True
|
||||
optimizations['recommended_dtype'] = 'bf16' if BFLOAT16_SUPPORTED else 'fp16'
|
||||
# For 16GB cards, BlockSwap might still be useful for 7B models but maybe less aggressive
|
||||
optimizations['block_swap_recommendation'] = 'moderate'
|
||||
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return optimizations
|
||||
|
||||
|
||||
def call_rope_with_stability(method, *args, **kwargs):
|
||||
"""
|
||||
Call RoPE method with stability fixes:
|
||||
1. Clear cache if available
|
||||
2. Disable autocast to prevent numerical issues
|
||||
2. Disable autocast to prevent numerical issues (CUDA only)
|
||||
This prevents artifacts in FP8/mixed precision models.
|
||||
"""
|
||||
if hasattr(method, 'cache_clear'):
|
||||
method.cache_clear()
|
||||
|
||||
with torch.cuda.amp.autocast(enabled=False):
|
||||
# Only use CUDA autocast context on CUDA devices
|
||||
# MPS has no CUDA autocast to disable
|
||||
if torch.cuda.is_available():
|
||||
with torch.cuda.amp.autocast(enabled=False):
|
||||
return method(*args, **kwargs)
|
||||
else:
|
||||
return method(*args, **kwargs)
|
||||
|
||||
|
||||
class FP8CompatibleDiT(torch.nn.Module):
|
||||
class CompatibleDiT(torch.nn.Module):
|
||||
"""
|
||||
Wrapper for DiT models with automatic compatibility management + advanced optimizations
|
||||
|
||||
Precision Handling:
|
||||
- FP8: Keeps native FP8 parameters (memory efficient), converts inputs/outputs to compute_dtype for arithmetic
|
||||
- FP16: Uses native FP16 precision throughout
|
||||
- BFloat16: Uses native BFloat16 precision throughout
|
||||
- Float32: Uses full precision for maximum quality
|
||||
- FP16/BFloat16/Float32: Uses native precision throughout
|
||||
- GGUF: On-the-fly dequantization to compute_dtype
|
||||
- MPS: Forces all parameters to compute_dtype (unified memory requires dtype consistency)
|
||||
- RoPE: Converted from FP8 to compute_dtype for numerical consistency
|
||||
|
||||
Optimizations:
|
||||
- Flash Attention: Automatic optimization of attention layers
|
||||
- RoPE Stabilization: Error handling for numerical stability in mixed precision
|
||||
- MPS Compatibility: Unified dtype conversion for Apple Silicon backends
|
||||
"""
|
||||
@@ -365,9 +750,12 @@ class FP8CompatibleDiT(torch.nn.Module):
|
||||
self.debug.start_timer("_convert_rope_freqs")
|
||||
self._convert_rope_freqs(target_dtype=self.compute_dtype)
|
||||
self.debug.end_timer("_convert_rope_freqs", "RoPE freqs conversion")
|
||||
|
||||
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)
|
||||
|
||||
# MPS requires unified dtype for all parameters/buffers (no autocast fallback)
|
||||
# Apply to ALL model types (FP8, FP16, GGUF) when dtype differs from compute_dtype
|
||||
if not skip_conversion and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
||||
if self.model_dtype != self.compute_dtype:
|
||||
self.debug.log(f"Converting NaDiT parameters/buffers to {self.compute_dtype} for MPS backend", category="setup", force=True)
|
||||
self.debug.start_timer("_force_nadit_precision")
|
||||
self._force_nadit_precision(target_dtype=self.compute_dtype)
|
||||
self.debug.end_timer("_force_nadit_precision", "NaDiT parameters/buffers conversion")
|
||||
@@ -377,11 +765,6 @@ class FP8CompatibleDiT(torch.nn.Module):
|
||||
self.debug.start_timer("_stabilize_rope_computations")
|
||||
self._stabilize_rope_computations()
|
||||
self.debug.end_timer("_stabilize_rope_computations", "RoPE stabilization")
|
||||
|
||||
# 🚀 FLASH ATTENTION OPTIMIZATION (Phase 2)
|
||||
self.debug.start_timer("_apply_flash_attention_optimization")
|
||||
self._apply_flash_attention_optimization()
|
||||
self.debug.end_timer("_apply_flash_attention_optimization", "Flash Attention application")
|
||||
|
||||
def _detect_model_dtype(self) -> torch.dtype:
|
||||
"""Detect main model dtype"""
|
||||
@@ -504,216 +887,6 @@ class FP8CompatibleDiT(torch.nn.Module):
|
||||
|
||||
if rope_count > 0:
|
||||
self.debug.log(f"Stabilized {rope_count} RoPE modules", category="success")
|
||||
|
||||
def _apply_flash_attention_optimization(self) -> None:
|
||||
"""🚀 FLASH ATTENTION OPTIMIZATION - 30-50% speedup of attention layers"""
|
||||
attention_layers_optimized = 0
|
||||
flash_attention_available = self._check_flash_attention_support()
|
||||
|
||||
for name, module in self.dit_model.named_modules():
|
||||
# Identify all attention layers
|
||||
if self._is_attention_layer(name, module):
|
||||
# Apply optimization based on availability
|
||||
if self._optimize_attention_layer(name, module, flash_attention_available):
|
||||
attention_layers_optimized += 1
|
||||
|
||||
if not flash_attention_available:
|
||||
self.debug.log("Flash Attention not available, using PyTorch SDPA as fallback", category="info", force=True)
|
||||
|
||||
def _check_flash_attention_support(self) -> bool:
|
||||
"""Check if Flash Attention is available"""
|
||||
# Check PyTorch SDPA (includes Flash Attention on H100/A100)
|
||||
if hasattr(torch.nn.functional, 'scaled_dot_product_attention'):
|
||||
return True
|
||||
|
||||
# Check flash-attn package (uses module-level check from top of file)
|
||||
return FLASH_ATTN_AVAILABLE
|
||||
|
||||
def _is_attention_layer(self, name: str, module: torch.nn.Module) -> bool:
|
||||
"""Identify if a module is an attention layer"""
|
||||
module_type_name = type(module).__name__
|
||||
|
||||
# Skip FlashAttentionVarlen modules - they manage their own attention backends
|
||||
if module_type_name == 'FlashAttentionVarlen':
|
||||
return False
|
||||
|
||||
attention_keywords = [
|
||||
'attention', 'attn', 'self_attn', 'cross_attn', 'mhattn', 'multihead',
|
||||
'transformer_block', 'dit_block'
|
||||
]
|
||||
|
||||
# Check by name
|
||||
if any(keyword in name.lower() for keyword in attention_keywords):
|
||||
return True
|
||||
|
||||
# Check by module type
|
||||
module_type = module_type_name.lower()
|
||||
if any(keyword in module_type for keyword in attention_keywords):
|
||||
return True
|
||||
|
||||
# Check by attributes (modules with q, k, v projections)
|
||||
if hasattr(module, 'q_proj') or hasattr(module, 'qkv') or hasattr(module, 'to_q'):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def _optimize_attention_layer(self, name: str, module: torch.nn.Module, flash_attention_available: bool) -> bool:
|
||||
"""Optimize a specific attention layer"""
|
||||
try:
|
||||
# Save original forward method
|
||||
if not hasattr(module, '_original_forward'):
|
||||
module._original_forward = module.forward
|
||||
|
||||
# Create new optimized forward method
|
||||
if flash_attention_available:
|
||||
optimized_forward = self._create_flash_attention_forward(module, name)
|
||||
else:
|
||||
optimized_forward = self._create_sdpa_forward(module, name)
|
||||
|
||||
# Replace forward method
|
||||
module.forward = optimized_forward
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
self.debug.log(f"Failed to optimize attention layer '{name}': {e}", level="WARNING", category="dit", force=True)
|
||||
return False
|
||||
|
||||
def _create_flash_attention_forward(self, module: torch.nn.Module, layer_name: str):
|
||||
"""Create optimized forward with Flash Attention"""
|
||||
original_forward = module._original_forward
|
||||
|
||||
def flash_attention_forward(*args, **kwargs):
|
||||
try:
|
||||
# Try to use Flash Attention via SDPA
|
||||
return self._sdpa_attention_forward(original_forward, module, *args, **kwargs)
|
||||
except Exception as e:
|
||||
# Fallback to original implementation
|
||||
self.debug.log(f"Flash Attention failed for {layer_name}, using original: {e}", level="WARNING", category="dit", force=True)
|
||||
return original_forward(*args, **kwargs)
|
||||
|
||||
return flash_attention_forward
|
||||
|
||||
def _create_sdpa_forward(self, module: torch.nn.Module, layer_name: str):
|
||||
"""Create optimized forward with PyTorch SDPA"""
|
||||
original_forward = module._original_forward
|
||||
|
||||
def sdpa_forward(*args, **kwargs):
|
||||
try:
|
||||
return self._sdpa_attention_forward(original_forward, module, *args, **kwargs)
|
||||
except Exception as e:
|
||||
# Fallback to original implementation
|
||||
return original_forward(*args, **kwargs)
|
||||
|
||||
return sdpa_forward
|
||||
|
||||
def _sdpa_attention_forward(self, original_forward, module: torch.nn.Module, *args, **kwargs):
|
||||
"""Optimized forward pass using SDPA (Scaled Dot Product Attention)"""
|
||||
# Detect if we can intercept and optimize this layer
|
||||
if len(args) >= 1 and isinstance(args[0], torch.Tensor):
|
||||
input_tensor = args[0]
|
||||
|
||||
# Check dimensions to ensure it's standard attention
|
||||
if len(input_tensor.shape) >= 3: # [batch, seq_len, hidden_dim] or similar
|
||||
try:
|
||||
return self._optimized_attention_computation(module, input_tensor, *args[1:], **kwargs)
|
||||
except:
|
||||
pass
|
||||
|
||||
# Fallback to original implementation
|
||||
return original_forward(*args, **kwargs)
|
||||
|
||||
def _optimized_attention_computation(self, module: torch.nn.Module, input_tensor: torch.Tensor, *args, **kwargs):
|
||||
"""Optimized attention computation with SDPA"""
|
||||
# Try to detect standard attention format
|
||||
batch_size, seq_len = input_tensor.shape[:2]
|
||||
|
||||
# Check if module has standard Q, K, V projections
|
||||
if hasattr(module, 'qkv') or (hasattr(module, 'q_proj') and hasattr(module, 'k_proj') and hasattr(module, 'v_proj')):
|
||||
return self._compute_sdpa_attention(module, input_tensor, *args, **kwargs)
|
||||
|
||||
# If no standard format detected, use original
|
||||
return module._original_forward(input_tensor, *args, **kwargs)
|
||||
|
||||
def _compute_sdpa_attention(self, module: torch.nn.Module, x: torch.Tensor, *args, **kwargs):
|
||||
"""Optimized SDPA computation for standard attention modules"""
|
||||
try:
|
||||
# Case 1: Module with combined QKV projection
|
||||
if hasattr(module, 'qkv'):
|
||||
qkv = module.qkv(x)
|
||||
# Reshape to separate Q, K, V
|
||||
batch_size, seq_len, _ = qkv.shape
|
||||
qkv = qkv.reshape(batch_size, seq_len, 3, -1)
|
||||
q, k, v = qkv.unbind(dim=2)
|
||||
|
||||
# Case 2: Separate Q, K, V projections
|
||||
elif hasattr(module, 'q_proj') and hasattr(module, 'k_proj') and hasattr(module, 'v_proj'):
|
||||
q = module.q_proj(x)
|
||||
k = module.k_proj(x)
|
||||
v = module.v_proj(x)
|
||||
else:
|
||||
# Unsupported format, use original
|
||||
return module._original_forward(x, *args, **kwargs)
|
||||
|
||||
# Detect number of heads
|
||||
head_dim = getattr(module, 'head_dim', None)
|
||||
num_heads = getattr(module, 'num_heads', None)
|
||||
|
||||
if head_dim is None or num_heads is None:
|
||||
# Try to guess from dimensions
|
||||
hidden_dim = q.shape[-1]
|
||||
if hasattr(module, 'num_heads'):
|
||||
num_heads = module.num_heads
|
||||
head_dim = hidden_dim // num_heads
|
||||
else:
|
||||
# Reasonable defaults
|
||||
head_dim = 64
|
||||
num_heads = hidden_dim // head_dim
|
||||
|
||||
# Reshape for multi-head attention
|
||||
batch_size, seq_len = q.shape[:2]
|
||||
q = q.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
|
||||
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.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 - 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,
|
||||
is_causal=False
|
||||
)
|
||||
|
||||
# Reshape back
|
||||
attn_output = attn_output.transpose(1, 2).contiguous().view(
|
||||
batch_size, seq_len, num_heads * head_dim
|
||||
)
|
||||
|
||||
# Output projection if it exists
|
||||
if hasattr(module, 'out_proj') or hasattr(module, 'o_proj'):
|
||||
proj = getattr(module, 'out_proj', None) or getattr(module, 'o_proj', None)
|
||||
attn_output = proj(attn_output)
|
||||
|
||||
return attn_output
|
||||
|
||||
except Exception as e:
|
||||
# In case of error, use original implementation
|
||||
return module._original_forward(x, *args, **kwargs)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
"""
|
||||
|
||||
@@ -10,8 +10,9 @@ import gc
|
||||
import sys
|
||||
import time
|
||||
import psutil
|
||||
import platform
|
||||
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'."""
|
||||
@@ -19,6 +20,31 @@ def _device_str(device: Union[torch.device, str]) -> str:
|
||||
return 'MPS' if s.startswith('MPS') else s
|
||||
|
||||
|
||||
def is_mps_available() -> bool:
|
||||
"""Check if MPS (Apple Metal) backend is available."""
|
||||
return hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()
|
||||
|
||||
|
||||
def is_cuda_available() -> bool:
|
||||
"""Check if CUDA backend is available."""
|
||||
return torch.cuda.is_available()
|
||||
|
||||
|
||||
def get_gpu_backend() -> str:
|
||||
"""Get the active GPU backend type.
|
||||
|
||||
Returns:
|
||||
'cuda': NVIDIA CUDA
|
||||
'mps': Apple Metal Performance Shaders
|
||||
'cpu': No GPU backend available
|
||||
"""
|
||||
if is_cuda_available():
|
||||
return 'cuda'
|
||||
if is_mps_available():
|
||||
return 'mps'
|
||||
return 'cpu'
|
||||
|
||||
|
||||
def get_device_list(include_none: bool = False, include_cpu: bool = False) -> List[str]:
|
||||
"""
|
||||
Get list of available compute devices for SeedVR2
|
||||
@@ -37,14 +63,14 @@ def get_device_list(include_none: bool = False, include_cpu: bool = False) -> Li
|
||||
has_mps = False
|
||||
|
||||
try:
|
||||
if hasattr(torch, "cuda") and hasattr(torch.cuda, "is_available") and torch.cuda.is_available():
|
||||
if is_cuda_available():
|
||||
devs += [f"cuda:{i}" for i in range(torch.cuda.device_count())]
|
||||
has_cuda = True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
if is_mps_available():
|
||||
devs.append("mps") # MPS doesn't use device indices
|
||||
has_mps = True
|
||||
except Exception:
|
||||
@@ -66,7 +92,7 @@ def get_device_list(include_none: bool = False, include_cpu: bool = False) -> Li
|
||||
result.extend(devs)
|
||||
|
||||
return result if result else []
|
||||
|
||||
|
||||
|
||||
def get_basic_vram_info(device: Optional[torch.device] = None) -> Dict[str, Any]:
|
||||
"""
|
||||
@@ -80,13 +106,13 @@ def get_basic_vram_info(device: Optional[torch.device] = None) -> Dict[str, Any]
|
||||
dict: {"free_gb": float, "total_gb": float} or {"error": str}
|
||||
"""
|
||||
try:
|
||||
if torch.cuda.is_available():
|
||||
if is_cuda_available():
|
||||
if device is None:
|
||||
device = torch.device("cuda:0")
|
||||
elif not isinstance(device, torch.device):
|
||||
device = torch.device(device)
|
||||
free_memory, total_memory = torch.cuda.mem_get_info(device)
|
||||
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
||||
elif is_mps_available():
|
||||
# MPS doesn't support per-device queries or mem_get_info
|
||||
# Use system memory as proxy
|
||||
mem = psutil.virtual_memory()
|
||||
@@ -106,29 +132,13 @@ 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.backends, 'mps') and torch.backends.mps.is_available()) else "CUDA"
|
||||
backend = "MPS" if is_mps_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]:
|
||||
def get_vram_usage(device: Optional[torch.device] = None, debug: Optional['Debug'] = None) -> Tuple[float, float, float, float]:
|
||||
"""
|
||||
Get current VRAM usage metrics for monitoring.
|
||||
Used for tracking memory consumption during processing.
|
||||
@@ -138,29 +148,30 @@ 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_reserved_gb)
|
||||
Returns (0, 0, 0) if no GPU available
|
||||
tuple: (allocated_gb, reserved_gb, peak_allocated_gb, peak_reserved_gb)
|
||||
Returns (0, 0, 0, 0) if no GPU available
|
||||
"""
|
||||
try:
|
||||
if torch.cuda.is_available():
|
||||
if is_cuda_available():
|
||||
if device is None:
|
||||
device = torch.device("cuda:0")
|
||||
elif not isinstance(device, torch.device):
|
||||
device = torch.device(device)
|
||||
allocated = torch.cuda.memory_allocated(device) / (1024**3)
|
||||
reserved = torch.cuda.memory_reserved(device) / (1024**3)
|
||||
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():
|
||||
peak_allocated = torch.cuda.max_memory_allocated(device) / (1024**3)
|
||||
peak_reserved = torch.cuda.max_memory_reserved(device) / (1024**3)
|
||||
return allocated, reserved, peak_allocated, peak_reserved
|
||||
elif is_mps_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)
|
||||
max_allocated = allocated # MPS doesn't track peak separately
|
||||
return allocated, reserved, max_allocated
|
||||
# MPS doesn't track peak separately
|
||||
return allocated, reserved, allocated, reserved
|
||||
except Exception as e:
|
||||
if debug:
|
||||
debug.log(f"Failed to get VRAM usage: {e}", level="WARNING", category="memory", force=True)
|
||||
return 0.0, 0.0, 0.0
|
||||
return 0.0, 0.0, 0.0, 0.0
|
||||
|
||||
|
||||
def get_ram_usage(debug: Optional['Debug'] = None) -> Tuple[float, float, float, float]:
|
||||
@@ -251,17 +262,17 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo
|
||||
# Use existing function for memory info
|
||||
mem_info = get_basic_vram_info(device=None)
|
||||
|
||||
if "error" not in mem_info:
|
||||
if "error" not in mem_info and mem_info["total_gb"] > 0:
|
||||
# Check VRAM/MPS memory pressure (5% free threshold)
|
||||
free_ratio = mem_info["free_gb"] / mem_info["total_gb"]
|
||||
if free_ratio < 0.05:
|
||||
should_clear = True
|
||||
if debug:
|
||||
backend = "MPS" if (hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()) else "VRAM"
|
||||
backend = "Unified Memory" if is_mps_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.backends, 'mps') and torch.backends.mps.is_available()):
|
||||
if not should_clear and not is_mps_available():
|
||||
mem = psutil.virtual_memory()
|
||||
if mem.available < mem.total * 0.05:
|
||||
should_clear = True
|
||||
@@ -284,10 +295,10 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo
|
||||
if debug:
|
||||
debug.start_timer(gpu_timer)
|
||||
|
||||
if torch.cuda.is_available():
|
||||
if is_cuda_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
||||
elif is_mps_available():
|
||||
torch.mps.empty_cache()
|
||||
|
||||
if debug:
|
||||
@@ -324,7 +335,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.backends, 'mps') and torch.backends.mps.is_available():
|
||||
elif is_mps_available():
|
||||
# macOS with MPS
|
||||
import ctypes # Import only when needed
|
||||
import ctypes.util
|
||||
@@ -401,7 +412,7 @@ def reset_vram_peak(device: Optional[torch.device] = None, debug: Optional['Debu
|
||||
if debug and debug.enabled:
|
||||
debug.log("Resetting VRAM peak memory statistics", category="memory")
|
||||
try:
|
||||
if torch.cuda.is_available():
|
||||
if is_cuda_available():
|
||||
if device is None:
|
||||
device = torch.device("cuda:0")
|
||||
elif not isinstance(device, torch.device):
|
||||
@@ -691,7 +702,7 @@ def manage_model_device(model: torch.nn.Module, target_device: torch.device, mod
|
||||
|
||||
if has_blockswap_config:
|
||||
is_blockswap_model = True
|
||||
# Get the actual model (handle FP8CompatibleDiT wrapper)
|
||||
# Get the actual model (handle CompatibleDiT wrapper)
|
||||
if hasattr(model, "dit_model"):
|
||||
actual_model = model.dit_model
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ Only includes constants actually used in the codebase
|
||||
"""
|
||||
|
||||
# Version information
|
||||
__version__ = "2.5.17"
|
||||
__version__ = "2.5.20"
|
||||
|
||||
import os
|
||||
import warnings
|
||||
|
||||
+138
-87
@@ -10,16 +10,33 @@ import torch
|
||||
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
|
||||
import platform
|
||||
from ..optimization.memory_manager import (
|
||||
get_vram_usage,
|
||||
get_basic_vram_info,
|
||||
get_ram_usage,
|
||||
reset_vram_peak,
|
||||
is_mps_available,
|
||||
is_cuda_available
|
||||
)
|
||||
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"
|
||||
def _format_peak_with_overflow(peak_gb: float, total_vram_gb: float) -> str:
|
||||
"""Format peak reserved memory, showing overflow breakdown on Windows.
|
||||
|
||||
Args:
|
||||
peak_gb: Peak reserved memory from PyTorch
|
||||
total_vram_gb: Physical GPU VRAM capacity
|
||||
"""
|
||||
if total_vram_gb <= 0:
|
||||
return f"{peak_gb:.2f}GB reserved"
|
||||
|
||||
overflow_gb = peak_gb - total_vram_gb
|
||||
if overflow_gb <= 0 or platform.system() != 'Windows':
|
||||
return f"{peak_gb:.2f}GB reserved"
|
||||
|
||||
return f"{peak_gb:.2f}GB reserved ({total_vram_gb:.0f}GB GPU + {overflow_gb:.2f}GB overflow)"
|
||||
|
||||
|
||||
class Debug:
|
||||
@@ -80,7 +97,8 @@ class Debug:
|
||||
self.vram_history: List[float] = []
|
||||
self.active_timer_stack: List[str] = []
|
||||
self.timer_namespace: str = ""
|
||||
self.phase_vram_peaks: Dict[str, float] = {}
|
||||
self.phase_vram_peaks_alloc: Dict[str, float] = {}
|
||||
self.phase_vram_peaks_rsv: Dict[str, float] = {}
|
||||
self.phase_ram_peaks: Dict[str, float] = {}
|
||||
|
||||
@torch._dynamo.disable # Skip tracing to avoid datetime.now() warnings
|
||||
@@ -125,26 +143,33 @@ class Debug:
|
||||
|
||||
def print_header(self, cli: bool = False) -> None:
|
||||
"""Print the header with banner - always displayed"""
|
||||
# Intro logo
|
||||
self.log("", category="none", force=True)
|
||||
self.log(" ╔══════════════════════════════════════════════════════════╗", category="none", force=True)
|
||||
self.log(" ║ ███████ ███████ ███████ ██████ ██ ██ ██████ ███████ ║", category="none", force=True)
|
||||
self.log(" ║ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ║", category="none", force=True)
|
||||
self.log(" ║ ███████ █████ █████ ██ ██ ██ ██ ██████ █████ ║", category="none", force=True)
|
||||
self.log(" ║ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ║", category="none", force=True)
|
||||
self.log(" ║ ███████ ███████ ███████ ██████ ████ ██ ██ ███████ ║", category="none", force=True)
|
||||
# Temporarily disable timestamps for clean header display
|
||||
original_timestamps = self.show_timestamps
|
||||
self.show_timestamps = False
|
||||
|
||||
# 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)
|
||||
# ASCII art logo
|
||||
self.log("", category="none", force=True)
|
||||
self.log("", category="none", force=True)
|
||||
self.log("███████╗███████╗███████╗██████╗ ██╗ ██╗██████╗ ██████╗ ███████╗", category="none", force=True, indent_level=1)
|
||||
self.log("██╔════╝██╔════╝██╔════╝██╔══██╗██║ ██║██╔══██╗ ╚════██╗ ██╔════╝", category="none", force=True, indent_level=1)
|
||||
self.log("███████╗█████╗ █████╗ ██║ ██║██║ ██║██████╔╝ █████╔╝ ███████╗", category="none", force=True, indent_level=1)
|
||||
self.log("╚════██║██╔══╝ ██╔══╝ ██║ ██║╚██╗ ██╔╝██╔══██╗ ██╔═══╝ ╚════██║", category="none", force=True, indent_level=1)
|
||||
self.log("███████║███████╗███████╗██████╔╝ ╚████╔╝ ██║ ██║ ███████╗ ██╗ ███████║", category="none", force=True, indent_level=1)
|
||||
self.log("╚══════╝╚══════╝╚══════╝╚═════╝ ╚═══╝ ╚═╝ ╚═╝ ╚══════╝ ╚═╝ ╚══════╝", category="none", force=True, indent_level=1)
|
||||
# Version and credits - left/right aligned to logo width
|
||||
version_text = f"v{__version__}"
|
||||
cli_indicator = "💻 CLI · " if cli else ""
|
||||
left_part = f"{cli_indicator}{version_text}"
|
||||
right_part = "© ByteDance Seed · NumZ · AInVFX"
|
||||
logo_width = 75
|
||||
emoji_compensation = 1 if cli else 0
|
||||
padding = logo_width - len(left_part) - len(right_part) - emoji_compensation
|
||||
self.log(f"{left_part}{' ' * max(1, padding)}{right_part}", category="none", force=True, indent_level=1)
|
||||
self.log("━" * logo_width, category="none", force=True, indent_level=1)
|
||||
self.log("", category="none", force=True)
|
||||
|
||||
# Restore timestamps setting
|
||||
self.show_timestamps = original_timestamps
|
||||
|
||||
# Environment info - only in debug mode
|
||||
if self.enabled:
|
||||
@@ -174,7 +199,7 @@ class Debug:
|
||||
cuda_ver = getattr(torch.version, 'cuda', None) or "N/A"
|
||||
|
||||
# GPU
|
||||
if torch.cuda.is_available():
|
||||
if is_cuda_available():
|
||||
try:
|
||||
props = torch.cuda.get_device_properties(0)
|
||||
gpu_str = f"{props.name} ({round(props.total_memory / (1024**3))}GB)"
|
||||
@@ -182,19 +207,37 @@ class Debug:
|
||||
except Exception:
|
||||
gpu_str = "CUDA"
|
||||
cudnn_ver = "N/A"
|
||||
elif getattr(getattr(torch, 'mps', None), 'is_available', lambda: False)():
|
||||
elif is_mps_available():
|
||||
gpu_str = "Apple Silicon (MPS)"
|
||||
cudnn_ver = "N/A"
|
||||
else:
|
||||
gpu_str = "CPU"
|
||||
cudnn_ver = "N/A"
|
||||
|
||||
# Flash Attn & Triton - reuse existing module constants
|
||||
# Flash Attn, SageAttn & 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 "✗")
|
||||
from ..optimization.compatibility import (
|
||||
FLASH_ATTN_2_AVAILABLE, FLASH_ATTN_3_AVAILABLE,
|
||||
SAGE_ATTN_2_AVAILABLE, SAGE_ATTN_3_AVAILABLE,
|
||||
TRITON_AVAILABLE
|
||||
)
|
||||
fa_parts = []
|
||||
if FLASH_ATTN_3_AVAILABLE:
|
||||
fa_parts.append("3")
|
||||
if FLASH_ATTN_2_AVAILABLE:
|
||||
fa_parts.append("2")
|
||||
flash_str = f"v{','.join(fa_parts)} ✓" if fa_parts else "✗"
|
||||
|
||||
sa_parts = []
|
||||
if SAGE_ATTN_3_AVAILABLE:
|
||||
sa_parts.append("3")
|
||||
if SAGE_ATTN_2_AVAILABLE:
|
||||
sa_parts.append("2")
|
||||
sage_str = f"v{','.join(sa_parts)} ✓" if sa_parts else "✗"
|
||||
|
||||
triton_str = "✓" if TRITON_AVAILABLE else "✗"
|
||||
except ImportError:
|
||||
flash_str = triton_str = "?"
|
||||
flash_str = sage_str = triton_str = "?"
|
||||
|
||||
# ComfyUI version
|
||||
comfy_str = None
|
||||
@@ -207,7 +250,7 @@ class Debug:
|
||||
|
||||
# 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")
|
||||
self.log(f"Python: {py_ver} | PyTorch: {torch_ver} | FlashAttn: {flash_str} | SageAttn: {sage_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")
|
||||
@@ -381,9 +424,12 @@ class Debug:
|
||||
if show_diff and self.memory_checkpoints:
|
||||
self._log_memory_diff(current_metrics=memory_info, force=force)
|
||||
|
||||
# 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...).",
|
||||
# Overflow warning (Windows only - WDDM can page to system RAM)
|
||||
overflow = memory_info.get('vram_overflow', 0.0)
|
||||
|
||||
if overflow > 0 and platform.system() == 'Windows':
|
||||
self.log(f"VRAM overflow: {overflow:.2f}GB paged to system RAM - 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
|
||||
@@ -395,10 +441,15 @@ class Debug:
|
||||
|
||||
# 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['vram_peak_alloc'] > 0:
|
||||
self.phase_vram_peaks_alloc[self.current_phase] = max(
|
||||
self.phase_vram_peaks_alloc.get(self.current_phase, 0),
|
||||
memory_info['vram_peak_alloc']
|
||||
)
|
||||
if memory_info['vram_peak_rsv'] > 0:
|
||||
self.phase_vram_peaks_rsv[self.current_phase] = max(
|
||||
self.phase_vram_peaks_rsv.get(self.current_phase, 0),
|
||||
memory_info['vram_peak_rsv']
|
||||
)
|
||||
if memory_info['ram_process'] > 0:
|
||||
self.phase_ram_peaks[self.current_phase] = max(
|
||||
@@ -410,13 +461,18 @@ class Debug:
|
||||
reset_vram_peak(device=None, debug=self)
|
||||
|
||||
def _collect_memory_metrics(self) -> Dict[str, Any]:
|
||||
"""Collect current memory metrics efficiently."""
|
||||
"""Collect current memory metrics."""
|
||||
is_mps = is_mps_available()
|
||||
has_gpu = is_mps or is_cuda_available()
|
||||
|
||||
metrics = {
|
||||
'vram_allocated': 0.0,
|
||||
'vram_reserved': 0.0,
|
||||
'vram_free': 0.0,
|
||||
'vram_total': 0.0,
|
||||
'vram_peak_since_last': 0.0,
|
||||
'vram_peak_alloc': 0.0,
|
||||
'vram_peak_rsv': 0.0,
|
||||
'vram_overflow': 0.0,
|
||||
'ram_process': 0.0,
|
||||
'ram_available': 0.0,
|
||||
'ram_total': 0.0,
|
||||
@@ -425,46 +481,36 @@ class Debug:
|
||||
'summary_ram': ""
|
||||
}
|
||||
|
||||
# VRAM metrics
|
||||
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
|
||||
# This captures the actual peak that occurred between calls
|
||||
metrics['vram_peak_since_last'] = current_global_peak
|
||||
|
||||
if has_gpu:
|
||||
metrics['vram_allocated'], metrics['vram_reserved'], metrics['vram_peak_alloc'], metrics['vram_peak_rsv'] = get_vram_usage(device=None, debug=self)
|
||||
vram_info = get_basic_vram_info(device=None)
|
||||
|
||||
if "error" not in vram_info:
|
||||
if "error" not in vram_info and vram_info["total_gb"] > 0:
|
||||
metrics['vram_free'] = vram_info["free_gb"]
|
||||
metrics['vram_total'] = vram_info["total_gb"]
|
||||
metrics['vram_overflow'] = max(0.0, metrics['vram_peak_rsv'] - metrics['vram_total'])
|
||||
|
||||
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: {peak_str} / "
|
||||
f"{metrics['vram_free']:.2f}GB free / "
|
||||
f"{metrics['vram_total']:.2f}GB total")
|
||||
else:
|
||||
metrics['summary_vram'] = ""
|
||||
else:
|
||||
metrics['summary_vram'] = ""
|
||||
backend = "Unified Memory" if is_mps else "VRAM"
|
||||
metrics['summary_vram'] = (
|
||||
f" [{backend}] {metrics['vram_allocated']:.2f}GB allocated / "
|
||||
f"{metrics['vram_reserved']:.2f}GB reserved / "
|
||||
f"Peak: {metrics['vram_peak_alloc']:.2f}GB / "
|
||||
f"{metrics['vram_free']:.2f}GB free / "
|
||||
f"{metrics['vram_total']:.2f}GB total"
|
||||
)
|
||||
|
||||
self.vram_history.append(metrics['vram_reserved'])
|
||||
|
||||
# RAM metrics using new function
|
||||
# RAM metrics
|
||||
metrics['ram_process'], metrics['ram_available'], metrics['ram_total'], metrics['ram_others'] = get_ram_usage(debug=self)
|
||||
|
||||
if metrics['ram_total'] > 0:
|
||||
metrics['summary_ram'] = (f" [RAM] {metrics['ram_process']:.2f}GB process / "
|
||||
f"{metrics['ram_others']:.2f}GB others / "
|
||||
f"{metrics['ram_available']:.2f}GB free / "
|
||||
f"{metrics['ram_total']:.2f}GB total")
|
||||
else:
|
||||
metrics['summary_ram'] = ""
|
||||
|
||||
# Update VRAM history for tracking
|
||||
if torch.cuda.is_available() or (hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()):
|
||||
self.vram_history.append(metrics['vram_allocated'])
|
||||
metrics['summary_ram'] = (
|
||||
f" [RAM] {metrics['ram_process']:.2f}GB process / "
|
||||
f"{metrics['ram_others']:.2f}GB others / "
|
||||
f"{metrics['ram_available']:.2f}GB free / "
|
||||
f"{metrics['ram_total']:.2f}GB total"
|
||||
)
|
||||
|
||||
return metrics
|
||||
|
||||
@@ -592,8 +638,8 @@ class Debug:
|
||||
self.log(f"Memory changes: {', '.join(diffs)}", category="memory", force=force, indent_level=1)
|
||||
|
||||
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:
|
||||
"""Display peak memory usage across all phases."""
|
||||
if not self.phase_vram_peaks_alloc and not self.phase_ram_peaks:
|
||||
return
|
||||
|
||||
phase_names = {
|
||||
@@ -603,9 +649,9 @@ class Debug:
|
||||
'phase4': 'Post-processing'
|
||||
}
|
||||
|
||||
is_mps = hasattr(torch.backends, 'mps') and torch.backends.mps.is_available() and not torch.cuda.is_available()
|
||||
is_mps = is_mps_available()
|
||||
|
||||
# Get total VRAM for swap detection (reuse existing function)
|
||||
# Get total VRAM for overflow formatting (Windows only)
|
||||
total_vram_gb = 0.0
|
||||
if not is_mps:
|
||||
vram_info = get_basic_vram_info(device=None)
|
||||
@@ -616,25 +662,29 @@ class Debug:
|
||||
self.log("────────────────────────", category="none", force=force)
|
||||
self.log("Peak memory by phase:", category="memory", force=force)
|
||||
|
||||
all_phases = sorted(set(self.phase_vram_peaks.keys()) | set(self.phase_ram_peaks.keys()))
|
||||
all_phases = sorted(set(self.phase_vram_peaks_alloc.keys()) | set(self.phase_ram_peaks.keys()))
|
||||
for phase_key in all_phases:
|
||||
phase_num = phase_key[-1]
|
||||
phase_name = phase_names.get(phase_key, phase_key)
|
||||
vram = self.phase_vram_peaks.get(phase_key, 0)
|
||||
alloc = self.phase_vram_peaks_alloc.get(phase_key, 0)
|
||||
rsv = self.phase_vram_peaks_rsv.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)
|
||||
self.log(f"{phase_num}. {phase_name}: {alloc:.2f}GB", category="memory", indent_level=1, 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)
|
||||
rsv_str = _format_peak_with_overflow(rsv, total_vram_gb)
|
||||
self.log(f"{phase_num}. {phase_name}: VRAM {alloc:.2f}GB allocated, {rsv_str} | RAM {ram:.2f}GB", category="memory", indent_level=1, force=force)
|
||||
|
||||
overall_alloc = max(self.phase_vram_peaks_alloc.values()) if self.phase_vram_peaks_alloc else 0
|
||||
overall_rsv = max(self.phase_vram_peaks_rsv.values()) if self.phase_vram_peaks_rsv else 0
|
||||
overall_ram = max(self.phase_ram_peaks.values()) if self.phase_ram_peaks else 0
|
||||
|
||||
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)
|
||||
self.log(f"Overall peak: {overall_alloc:.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)
|
||||
overall_rsv_str = _format_peak_with_overflow(overall_rsv, total_vram_gb)
|
||||
self.log(f"Overall peak: VRAM {overall_alloc:.2f}GB allocated, {overall_rsv_str} | 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:
|
||||
@@ -744,6 +794,7 @@ class Debug:
|
||||
self.timer_durations.clear()
|
||||
self.timer_messages.clear()
|
||||
self.active_timer_stack.clear()
|
||||
self.phase_vram_peaks.clear()
|
||||
self.phase_vram_peaks_alloc.clear()
|
||||
self.phase_vram_peaks_rsv.clear()
|
||||
self.phase_ram_peaks.clear()
|
||||
self.current_phase = None
|
||||
Reference in New Issue
Block a user