File size: 6,053 Bytes
96d97a7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """
Performance optimization module for SeedVR2
Contains optimized tensor operations and video processing functions
Extracted from: seedvr2.py (lines 1633-1730)
"""
import torch
from typing import List
def optimized_channels_to_last(tensor):
"""🚀 Optimized replacement for rearrange(tensor, 'b c ... -> b ... c')
Moves channels from position 1 to last position using PyTorch native operations.
"""
if tensor.ndim == 3: # [batch, channels, spatial]
return tensor.permute(0, 2, 1)
elif tensor.ndim == 4: # [batch, channels, height, width]
return tensor.permute(0, 2, 3, 1)
elif tensor.ndim == 5: # [batch, channels, depth, height, width]
return tensor.permute(0, 2, 3, 4, 1)
else:
# Fallback for other dimensions - move channel (dim=1) to last
dims = list(range(tensor.ndim))
dims = [dims[0]] + dims[2:] + [dims[1]] # [0, 2, 3, ..., 1]
return tensor.permute(*dims)
def optimized_channels_to_second(tensor):
"""🚀 Optimized replacement for rearrange(tensor, 'b ... c -> b c ...')
Moves channels from last position to position 1 using PyTorch native operations.
"""
if tensor.ndim == 3: # [batch, spatial, channels]
return tensor.permute(0, 2, 1)
elif tensor.ndim == 4: # [batch, height, width, channels]
return tensor.permute(0, 3, 1, 2)
elif tensor.ndim == 5: # [batch, depth, height, width, channels]
return tensor.permute(0, 4, 1, 2, 3)
else:
# Fallback for other dimensions - move last dim to position 1
dims = list(range(tensor.ndim))
dims = [dims[0], dims[-1]] + dims[1:-1] # [0, -1, 1, 2, ..., -2]
return tensor.permute(*dims)
def optimized_video_rearrange(video_tensors: List[torch.Tensor]) -> List[torch.Tensor]:
"""
🚀 OPTIMIZED version of video rearrangement
Replaces slow loops with vectorized operations
Transforms:
- 3D: c h w -> t c h w (with t=1)
- 4D: c t h w -> t c h w
Expected gains: 5-10x faster than naive loops
Args:
video_tensors: List of video tensors to rearrange
Returns:
List of rearranged tensors in t c h w format
Raises:
ValueError: If video tensor has invalid dimensions (not 3D or 4D)
"""
if not video_tensors:
return []
# 🔍 Analyze dimensions to optimize processing
videos_3d = []
videos_4d = []
indices_3d = []
indices_4d = []
for i, video in enumerate(video_tensors):
if video.ndim == 3:
videos_3d.append(video)
indices_3d.append(i)
elif video.ndim == 4:
videos_4d.append(video)
indices_4d.append(i)
else:
raise ValueError(f"Video tensor at index {i} has invalid dimensions: {video.ndim}. Expected 3D or 4D.")
# 🎯 Prepare final result
samples = [None] * len(video_tensors)
# 🚀 BATCH PROCESSING for 3D videos (c h w -> 1 c h w)
if videos_3d:
# Stack + permute (faster than rearrange)
# c h w -> c 1 h w -> 1 c h w
batch_3d = torch.stack([v.unsqueeze(1) for v in videos_3d]) # [batch, c, 1, h, w]
batch_3d = batch_3d.permute(0, 2, 1, 3, 4) # [batch, 1, c, h, w]
for i, idx in enumerate(indices_3d):
samples[idx] = batch_3d[i] # [1, c, h, w]
# 🚀 BATCH PROCESSING for 4D videos (c t h w -> t c h w)
if videos_4d:
# Check if all 4D videos have the same shape for maximum optimization
shapes = [v.shape for v in videos_4d]
if len(set(shapes)) == 1:
# 🎯 MAXIMUM OPTIMIZATION: All shapes identical
# Stack + permute in single operation
batch_4d = torch.stack(videos_4d) # [batch, c, t, h, w]
batch_4d = batch_4d.permute(0, 2, 1, 3, 4) # [batch, t, c, h, w]
for i, idx in enumerate(indices_4d):
samples[idx] = batch_4d[i] # [t, c, h, w]
else:
# 🔄 FALLBACK: Different shapes, optimized individual processing
for i, idx in enumerate(indices_4d):
# Use permute instead of rearrange (faster)
samples[idx] = videos_4d[i].permute(1, 0, 2, 3) # c t h w -> t c h w
return samples
def optimized_single_video_rearrange(video: torch.Tensor) -> torch.Tensor:
"""
🚀 OPTIMIZED version for single video tensor
Replaces rearrange() with native PyTorch operations
Transforms:
- 3D: c h w -> 1 c h w (add temporal dimension)
- 4D: c t h w -> t c h w (permute dimensions)
Expected gains: 2-5x faster than rearrange()
Args:
video: Input video tensor
Returns:
Rearranged tensor with temporal dimension first
"""
if video.ndim == 3:
# c h w -> 1 c h w (add temporal dimension t=1)
return video.unsqueeze(0)
else: # ndim == 4
# c t h w -> t c h w (permute channels and temporal)
return video.permute(1, 0, 2, 3)
def optimized_sample_to_image_format(sample: torch.Tensor) -> torch.Tensor:
"""
🚀 OPTIMIZED version to convert sample to image format
Replaces rearrange() with native PyTorch operations
Transforms:
- 3D: c h w -> 1 h w c (add temporal dimension + permute to image format)
- 4D: t c h w -> t h w c (permute to image format)
Expected gains: 2-5x faster than rearrange()
Args:
sample: Input sample tensor
Returns:
Tensor in image format (channels last)
"""
if sample.ndim == 3:
# c h w -> 1 h w c (add temporal dimension then permute)
return sample.unsqueeze(0).permute(0, 2, 3, 1)
else: # ndim == 4
# t c h w -> t h w c (permute channels to last)
return sample.permute(0, 2, 3, 1)
|