prototype seedvr

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/seedvr/src/optimization/memory_manager.py

@@ -0,0 +1,131 @@
+"""
+Memory management module for SeedVR2
+Handles VRAM usage, cache management, and memory optimization
+
+Extracted from: seedvr2.py (lines 373-405, 607-626, 1016-1044)
+"""
+
+import os
+import torch
+import gc
+from typing import Tuple, Optional
+from src.common.cache import Cache
+from src.models.dit_v2.rope import RotaryEmbeddingBase
+
+
+def preinitialize_rope_cache(runner) -> None:
+    """
+    🚀 Pre-initialize RoPE cache to avoid OOM at first launch
+    
+    Args:
+        runner: The model runner containing DiT and VAE models
+    """
+
+    try:
+        # Create dummy tensors to simulate common shapes
+        # Format: [batch, channels, frames, height, width] for vid_shape
+        # Format: [batch, seq_len] for txt_shape
+        common_shapes = [
+            # Common video resolutions
+            (torch.tensor([[1, 3, 3]], dtype=torch.long), torch.tensor([[77]], dtype=torch.long)),    # 1 frame, 77 tokens
+            (torch.tensor([[4, 3, 3]], dtype=torch.long), torch.tensor([[77]], dtype=torch.long)),    # 4 frames
+            (torch.tensor([[5, 3, 3]], dtype=torch.long), torch.tensor([[77]], dtype=torch.long)),    # 5 frames (4n+1 format)
+            (torch.tensor([[1, 4, 4]], dtype=torch.long), torch.tensor([[77]], dtype=torch.long)),    # Higher resolution
+        ]
+
+        # Create mock cache for pre-initialization
+
+        temp_cache = Cache()
+
+        # Access RoPE modules in DiT (recursive search)
+        def find_rope_modules(module):
+            rope_modules = []
+            for name, child in module.named_modules():
+                if hasattr(child, 'get_freqs') and callable(child.get_freqs):
+                    rope_modules.append((name, child))
+            return rope_modules
+
+        rope_modules = find_rope_modules(runner.dit)
+
+        # Pre-calculate for each RoPE module found
+        for name, rope_module in rope_modules:
+            # Temporarily move module to CPU if necessary
+            original_device = next(rope_module.parameters()).device if list(rope_module.parameters()) else torch.device('cpu')
+            rope_module.to('cpu')
+
+            try:
+                for vid_shape, txt_shape in common_shapes:
+                    cache_key = f"720pswin_by_size_bysize_{tuple(vid_shape[0].tolist())}_sd3.mmrope_freqs_3d"
+
+                    def compute_freqs():
+                        # Calculate with reduced dimensions to avoid OOM
+                        with torch.no_grad():
+                            # Detect RoPE module type
+                            module_type = type(rope_module).__name__
+
+                            if module_type == 'NaRotaryEmbedding3d':
+                                # NaRotaryEmbedding3d: only takes shape (vid_shape)
+                                return rope_module.get_freqs(vid_shape.cpu())
+                            else:
+                                # Standard RoPE: takes vid_shape and txt_shape
+                                return rope_module.get_freqs(vid_shape.cpu(), txt_shape.cpu())
+
+                    # Store in cache
+                    temp_cache(cache_key, compute_freqs)
+
+            except Exception as e:
+                print(f"    ❌ Error in module {name}: {e}")
+            finally:
+                # Restore to original device
+                rope_module.to(original_device)
+
+        # Copy temporary cache to runner cache
+        if hasattr(runner, 'cache'):
+            runner.cache.cache.update(temp_cache.cache)
+        else:
+            runner.cache = temp_cache
+
+    except Exception as e:
+        print(f"  ⚠️ Error during RoPE pre-init: {e}")
+        print("  🔄 Model will work but could have OOM at first launch")
+
+
+def clear_rope_cache(runner) -> None:
+    """
+    🧹 Clear RoPE cache to free VRAM
+    
+    Args:
+        runner: The model runner containing the cache
+    """
+    print("🧹 Cleaning RoPE cache...")
+
+    if hasattr(runner, 'cache') and hasattr(runner.cache, 'cache'):
+        # Count entries before cleanup
+        cache_size = len(runner.cache.cache)
+
+        # Free all tensors from cache
+        for key, value in runner.cache.cache.items():
+            if isinstance(value, (tuple, list)):
+                for item in value:
+                    if hasattr(item, 'cpu'):
+                        item.cpu()
+                        del item
+            elif hasattr(value, 'cpu'):
+                value.cpu()
+                del value
+
+        # Clear the cache
+        runner.cache.cache.clear()
+        print(f"  ✅ RoPE cache cleared ({cache_size} entries removed)")
+
+    if hasattr(runner, 'dit'):
+        cleared_lru_count = 0
+        for module in runner.dit.modules():
+            if isinstance(module, RotaryEmbeddingBase):
+                if hasattr(module.get_axial_freqs, 'cache_clear'):
+                    module.get_axial_freqs.cache_clear()
+                    cleared_lru_count += 1
+        if cleared_lru_count > 0:
+            print(f"  ✅ Cleared {cleared_lru_count} LRU caches from RoPE modules.")
+
+    print("🎯 RoPE cache cleanup completed!")

modules/seedvr/src/optimization/performance.py

@@ -0,0 +1,160 @@
+"""
+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, Union
+
+
+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
+    """
+    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)
+        else:  # ndim == 4
+            videos_4d.append(video)
+            indices_4d.append(i)
+
+    # 🎯 Prepare final result
+    samples = [None] * len(video_tensors)
+
+    # 🚀 BATCH PROCESSING for 3D videos (c h w -> 1 c h w)
+    if videos_3d:
+        # Method 1: 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)
+
+
+def temporal_latent_blending(latents1: torch.Tensor, latents2: torch.Tensor, blend_frames: int) -> torch.Tensor:
+    """
+    🎨 Temporal blending in latent space to avoid discontinuities
+    
+    Args:
+        latents1: Latents from previous batch (end frames)
+        latents2: Latents from current batch (start frames)
+        blend_frames: Number of frames to blend
+    
+    Returns:
+        Blended latents for smooth transition
+    """
+    if latents1.shape[0] != latents2.shape[0]:
+        # Adjust dimensions if necessary
+        min_frames = min(latents1.shape[0], latents2.shape[0])
+        latents1 = latents1[:min_frames]
+        latents2 = latents2[:min_frames]
+
+    # Create linear blending weights
+    # Frame 0: 100% latents1, 0% latents2
+    # Frame n: 0% latents1, 100% latents2
+    weights1 = torch.linspace(1.0, 0.0, blend_frames).view(-1, 1, 1, 1).to(latents1.device)
+    weights2 = torch.linspace(0.0, 1.0, blend_frames).view(-1, 1, 1, 1).to(latents2.device)
+
+    # Apply blending
+    blended_latents = weights1 * latents1 + weights2 * latents2
+
+    return blended_latents
+
+
+