""" Unified debugging system for SeedVR2 generation pipeline Provides structured logging, memory tracking, and performance monitoring for all pipeline stages, including BlockSwap operations. """ import time import torch import gc from typing import Optional, List, Dict, Any, Union from datetime import datetime 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_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: """ Unified debug logging for generation pipeline and BlockSwap monitoring Features: - Structured logging with categories - Memory tracking (VRAM/RAM) - Timing utilities - BlockSwap operation tracking - Minimal overhead when disabled - Timestamped logs for better troubleshooting - Force parameters for critical logs """ # Icon mapping for different categories CATEGORY_ICONS = { "general": "๐Ÿ”„", # General operations/processing "timing": "โšก", # Performance timing "memory": "๐Ÿ“Š", # Memory usage tracking "cache": "๐Ÿ’พ", # Cache operations "cleanup": "๐Ÿงน", # Cleanup operations "setup": "๐Ÿ”ง", # Configuration/setup "generation": "๐ŸŽฌ", # Generation process "dit": "๐Ÿš€", # Model loading/operations "blockswap": "๐Ÿ”€", # BlockSwap operations "download": "๐Ÿ“ฅ", # Download operations "success": "โœ…", # Successful completion "warning": "โš ๏ธ", # Warnings "error": "โŒ", # Errors "info": "โ„น๏ธ", # Statistics/info "tip" :"๐Ÿ’ก", # Tip/suggestion "video": "๐Ÿ“น", # Video/sequence info "reuse": "โ™ป๏ธ", # Reusing/recycling "runner": "๐Ÿƒ", # Runner operations "vae": "๐ŸŽจ", # VAE operations\ "precision": "๐ŸŽฏ", # Precision "device": "๐Ÿ–ฅ๏ธ", # Device info "file": "๐Ÿ“‚", # File operations "alpha": "๐Ÿ‘ป", # Alpha operations "starlove": "โญ๐Ÿ’", # Star + love "dialogue": "๐Ÿ’ฌ", # Dialogue "none" : "", } def __init__(self, enabled: bool = False, show_timestamps: bool = True): self.enabled = enabled self.show_timestamps = show_timestamps self.timers: Dict[str, float] = {} self.memory_checkpoints: List[Dict[str, Any]] = [] self.max_checkpoints = 100 self.timer_hierarchy: Dict[str, List[str]] = {} self.timer_durations: Dict[str, float] = {} self.timer_messages: Dict[str, str] = {} self.swap_times: List[Dict[str, Any]] = [] self.current_phase: Optional[str] = None self.vram_history: List[float] = [] self.active_timer_stack: List[str] = [] self.timer_namespace: str = "" 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 def log(self, message: str, level: str = "INFO", category: str = "general", force: bool = False, indent_level: int = 0) -> None: """ Log a categorized message with optional timestamp and indentation Args: message: Message to log level: Log level (INFO, WARN, ERROR) category: Category for the message force: If True, always log regardless of enabled state (for critical messages) indent_level: Indentation level (0=no indent, 1=2 spaces, 2=4 spaces, etc.) """ # Always log forced messages or if debugging is enabled - early return if not if not (self.enabled or force): return # Get icon for category, fallback to general icon icon = self.CATEGORY_ICONS.get(category, self.CATEGORY_ICONS["general"]) # Format prefix based on level if level == "WARNING": icon = self.CATEGORY_ICONS["warning"] elif level == "ERROR": icon = self.CATEGORY_ICONS["error"] # Build the log message with optional timestamp if self.show_timestamps: timestamp = datetime.now().strftime("%H:%M:%S.%f")[:-3] prefix = f"[{timestamp}] {icon}" else: prefix = f"{icon}" if level != "INFO": prefix += f" [{level}]" # Add indentation indent = " " * (indent_level * 2) print(f"{prefix} {indent}{message}", flush=True) def print_header(self, cli: bool = False) -> None: """Print the header with banner - always displayed""" # Temporarily disable timestamps for clean header display original_timestamps = self.show_timestamps self.show_timestamps = False # 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: self._print_environment_info(cli) def _print_environment_info(self, cli: bool = False) -> None: """Print concise environment info for bug reports - zero cost when debug disabled""" import platform import sys # OS os_name = platform.system() if os_name == "Windows": os_str = f"Windows ({platform.version()})" elif os_name == "Darwin": os_str = f"macOS {platform.mac_ver()[0]}" else: try: distro = platform.freedesktop_os_release() os_str = f"{distro.get('NAME', 'Linux')} {distro.get('VERSION_ID', '')}" except (OSError, AttributeError): os_str = f"Linux {platform.release()}" # Python & PyTorch & CUDA py_ver = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}" torch_ver = torch.__version__ cuda_ver = getattr(torch.version, 'cuda', None) or "N/A" # GPU if is_cuda_available(): try: props = torch.cuda.get_device_properties(0) gpu_str = f"{props.name} ({round(props.total_memory / (1024**3))}GB)" cudnn_ver = str(torch.backends.cudnn.version()) if torch.backends.cudnn.is_available() else "N/A" except Exception: gpu_str = "CUDA" cudnn_ver = "N/A" elif is_mps_available(): gpu_str = "Apple Silicon (MPS)" cudnn_ver = "N/A" else: gpu_str = "CPU" cudnn_ver = "N/A" # Flash Attn, SageAttn & Triton - reuse existing module constants try: 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 = sage_str = triton_str = "?" # ComfyUI version comfy_str = None if not cli: try: from comfyui_version import __version__ as comfy_ver comfy_str = comfy_ver except ImportError: pass # Print self.log(f"OS: {os_str} | GPU: {gpu_str}", category="info") self.log(f"Python: {py_ver} | PyTorch: {torch_ver} | 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") def print_footer(self) -> None: """Print the footer with links - always displayed""" self.log("", category="none", force=True) self.log("โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€", category="none", force=True) self.log("Questions? Updates? Watch, star & sponsor if you can!", category="dialogue", force=True) self.log("https://www.youtube.com/@AInVFX", category="generation", force=True) self.log("https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler", category="starlove", force=True) @torch._dynamo.disable # Skip tracing to avoid time.time() warnings def start_timer(self, name: str, force: bool = False) -> None: """ Start a named timer Args: name: Timer name force: If True, start timer even when debug is disabled """ if self.enabled or force: # Apply namespace if set if self.timer_namespace: name = f"{self.timer_namespace}_{name}" self.timers[name] = time.time() # Track phase for memory peak monitoring if name.startswith("phase") and name.endswith(("_encoding", "_upscaling", "_decoding", "_postprocessing")): # Extract phase number (e.g., "phase3_decoding" -> "3") phase_num = name.split("_")[0].replace("phase", "") self.current_phase = f"phase{phase_num}" # Auto-hierarchy: if there's an active timer, this is a child if self.active_timer_stack: parent = self.active_timer_stack[-1] if parent not in self.timer_hierarchy: self.timer_hierarchy[parent] = [] # Only add if not already a child (prevents duplicates) if name not in self.timer_hierarchy[parent]: self.timer_hierarchy[parent].append(name) # Push to stack self.active_timer_stack.append(name) @torch._dynamo.disable # Skip tracing to avoid time.time() warnings def end_timer(self, name: str, message: Optional[str] = None, force: bool = False, show_breakdown: bool = False, custom_children: Optional[Dict[str, float]] = None) -> float: """ End a timer and optionally log its duration Args: name: Timer name message: Optional message to log with the duration force: If True, log even when debug is disabled (for critical timings) show_breakdown: If True, show breakdown of child timers custom_children: Optional dict of child timer names and durations to override automatic hierarchy Returns: Duration in seconds (0.0 if timer not found) """ # Apply namespace if set if self.timer_namespace: name = f"{self.timer_namespace}_{name}" # Check if timer exists if name not in self.timers: return 0.0 duration = time.time() - self.timers[name] self.timer_durations[name] = duration # Store the message for later use in summary if message: self.timer_messages[name] = message del self.timers[name] # Pop from stack if this is the current active timer if self.active_timer_stack and self.active_timer_stack[-1] == name: self.active_timer_stack.pop() # If debug is disabled and not forcing, return early if not self.enabled and not force: return duration # ONLY log if show_breakdown is True - this means it's a major summary timer if message and show_breakdown: # Use custom children if provided, otherwise use automatic hierarchy if custom_children: children = custom_children child_total = sum(children.values()) unaccounted = duration - child_total self.log(f"{message}: {duration:.2f}s", category="timing", force=force) # Sort custom children by duration for better readability sorted_children = sorted(children.items(), key=lambda x: x[1], reverse=True) for child_name, child_duration in sorted_children: if child_duration >= 0.01: # Only show if >= 10ms self.log(f"โ””โ”€ {child_name}: {child_duration:.2f}s", category="timing", force=force, indent_level=1) else: # Use automatic hierarchy tracking children = self.timer_hierarchy.get(name, []) child_total = sum(self.timer_durations.get(child, 0) for child in children) unaccounted = duration - child_total self.log(f"{message}: {duration:.2f}s", category="timing", force=force) # Sort children by duration for better readability sorted_children = sorted(children, key=lambda c: self.timer_durations.get(c, 0), reverse=True) for child in sorted_children: child_duration = self.timer_durations.get(child, 0) if child_duration >= 0.01: # Only show if >= 10ms child_message = self.timer_messages.get(child, child) self.log(f"โ””โ”€ {child_message}: {child_duration:.2f}s", category="timing", force=force, indent_level=1) # Recursively show grandchildren if child in self.timer_hierarchy: grandchildren = self.timer_hierarchy[child] sorted_grandchildren = sorted(grandchildren, key=lambda c: self.timer_durations.get(c, 0), reverse=True) for grandchild in sorted_grandchildren: grandchild_duration = self.timer_durations.get(grandchild, 0) if grandchild_duration >= 0.01: # Only show if >= 10ms grandchild_message = self.timer_messages.get(grandchild, grandchild) self.log(f"โ””โ”€ {grandchild_message}: {grandchild_duration:.2f}s", category="timing", force=force, indent_level=2) if unaccounted > 0.01: # Show if more than 10ms unaccounted self.log(f"โ””โ”€ (other operations): {unaccounted:.2f}s", category="timing", force=force, indent_level=1) return duration def log_memory_state(self, label: str, show_diff: bool = True, show_tensors: bool = False, detailed_tensors: bool = False, force: bool = False) -> None: """ Log current memory state with minimal overhead. Args: label: Description for this checkpoint show_diff: Show change from last checkpoint show_tensors: Include tensor counts detailed_tensors: Show detailed tensor analysis (use sparingly) force: If True, always log regardless of enabled state """ if not (self.enabled or force): return # Collect memory metrics efficiently memory_info = self._collect_memory_metrics() # Show category self.log(f"{label}:", category="memory", force=force) # Show VRAM if memory_info['summary_vram']: self.log(f"{memory_info['summary_vram']}", category="memory", force=force) # Show RAM if memory_info['summary_ram']: self.log(f"{memory_info['summary_ram']}", category="memory", force=force) # Show tensors if show_tensors: tensor_stats = self._collect_tensor_stats(detailed=detailed_tensors) self.log(f"{tensor_stats['summary']}", category="memory", force=force) # Show diff from last checkpoint if show_diff and self.memory_checkpoints: self._log_memory_diff(current_metrics=memory_info, force=force) # 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 if detailed_tensors and tensor_stats.get('details'): self._log_detailed_tensor_analysis(details=tensor_stats['details'], force=force) # Store checkpoint with memory limit self._store_checkpoint(label, memory_info) # Update phase peaks if we're in an active phase if self.current_phase: 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( self.phase_ram_peaks.get(self.current_phase, 0), memory_info['ram_process'] ) # Reset PyTorch's peak memory stats for next interval reset_vram_peak(device=None, debug=self) def _collect_memory_metrics(self) -> Dict[str, Any]: """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_alloc': 0.0, 'vram_peak_rsv': 0.0, 'vram_overflow': 0.0, 'ram_process': 0.0, 'ram_available': 0.0, 'ram_total': 0.0, 'ram_others': 0.0, 'summary_vram': "", 'summary_ram': "" } 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 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 = "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 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" ) return metrics def _collect_tensor_stats(self, detailed: bool = False) -> Dict[str, Any]: """Collect tensor statistics with minimal overhead.""" stats = { 'gpu_count': 0, 'cpu_count': 0, 'total_count': 0, 'summary': "", 'details': None } if detailed: stats['details'] = { 'gpu_tensors': [], 'large_cpu_tensors': [], 'shape_patterns': {}, 'module_types': {} } # Single pass through gc objects for obj in gc.get_objects(): try: if torch.is_tensor(obj): stats['total_count'] += 1 is_gpu = obj.is_cuda or (hasattr(obj, 'is_mps') and obj.is_mps) if is_gpu: stats['gpu_count'] += 1 else: stats['cpu_count'] += 1 # Collect detailed info if requested if detailed and obj.numel() > 0: size_mb = obj.element_size() * obj.nelement() / (1024**2) if is_gpu or size_mb > 10: # Only track GPU tensors or large CPU tensors tensor_info = { 'shape': tuple(obj.shape), 'dtype': str(obj.dtype), 'size_mb': size_mb, 'requires_grad': obj.requires_grad } if is_gpu: stats['details']['gpu_tensors'].append(tensor_info) elif size_mb > 10: # Large CPU tensors (>10MB) stats['details']['large_cpu_tensors'].append(tensor_info) # Track shape patterns shape_key = str(tuple(obj.shape)) stats['details']['shape_patterns'][shape_key] = stats['details']['shape_patterns'].get(shape_key, 0) + 1 elif detailed and isinstance(obj, torch.nn.Module): module_type = type(obj).__name__ stats['details']['module_types'][module_type] = stats['details']['module_types'].get(module_type, 0) + 1 except (ReferenceError, AttributeError): # Object was deleted or doesn't have expected attributes pass stats['summary'] = f" [Tensors] {stats['gpu_count']} GPU / {stats['cpu_count']} CPU / {stats['total_count']} total" return stats def _log_detailed_tensor_analysis(self, details: Dict[str, Any], force: bool = False) -> None: """Log detailed tensor analysis when requested.""" # GPU tensors if details['gpu_tensors']: gpu_total_gb = sum(t['size_mb'] for t in details['gpu_tensors']) / 1024 self.log(f"GPU tensors: {len(details['gpu_tensors'])} using {gpu_total_gb:.2f}GB", category="memory", force=force, indent_level=1) # Show top 5 largest largest = sorted(details['gpu_tensors'], key=lambda x: x['size_mb'], reverse=True)[:5] for t in largest: self.log(f"{t['shape']}: {t['size_mb']:.2f}MB, {t['dtype']}", category="memory", force=force, indent_level=1) # Large CPU tensors if details['large_cpu_tensors']: cpu_large_gb = sum(t['size_mb'] for t in details['large_cpu_tensors']) / 1024 self.log(f"Large CPU tensors (>10MB):", category="memory", force=force, indent_level=1) self.log(f"{len(details['large_cpu_tensors'])} using {cpu_large_gb:.2f}GB", category="memory", force=force, indent_level=1) # Show top 3 largest largest = sorted(details['large_cpu_tensors'], key=lambda x: x['size_mb'], reverse=True)[:3] for t in largest: self.log(f"{t['shape']}: {t['size_mb']:.2f}MB, {t['dtype']}", category="memory", force=force, indent_level=1) # Common shape patterns if details['shape_patterns']: common_shapes = sorted(details['shape_patterns'].items(), key=lambda x: x[1], reverse=True)[:5] if len(common_shapes) > 0: self.log("Common tensor shapes:", category="memory", force=force, indent_level=1) for shape, count in common_shapes: if count > 1: self.log(f"{shape}: {count} instances", category="memory", force=force, indent_level=1) # Module instances if details['module_types']: multi_instance = [(k, v) for k, v in details['module_types'].items() if v > 1] if multi_instance: self.log("Multiple module instances:", category="memory", force=force, indent_level=1) for mtype, count in sorted(multi_instance, key=lambda x: x[1], reverse=True)[:5]: self.log(f"{mtype}: {count} instances", category="memory", force=force, indent_level=1) def _log_memory_diff(self, current_metrics: Dict[str, Any], force: bool = False) -> None: """Log memory changes from last checkpoint.""" last = self.memory_checkpoints[-1] vram_diff = current_metrics['vram_allocated'] - last.get('vram_allocated', 0) ram_diff = current_metrics['ram_process'] - last.get('ram_process', 0) diffs = [] if abs(vram_diff) > 0.01: sign = "+" if vram_diff > 0 else "" diffs.append(f"VRAM {sign}{vram_diff:.2f}GB") if abs(ram_diff) > 0.01: sign = "+" if ram_diff > 0 else "" diffs.append(f"RAM {sign}{ram_diff:.2f}GB") if diffs: 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.""" if not self.phase_vram_peaks_alloc and not self.phase_ram_peaks: return phase_names = { 'phase1': 'VAE encoding', 'phase2': 'DiT upscaling', 'phase3': 'VAE decoding', 'phase4': 'Post-processing' } is_mps = is_mps_available() # 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) if "error" not in vram_info: total_vram_gb = vram_info["total_gb"] self.log("", category="none", force=force) self.log("โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€", category="none", force=force) self.log("Peak memory by phase:", category="memory", force=force) 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) 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_num}. {phase_name}: {alloc:.2f}GB", category="memory", indent_level=1, force=force) else: 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: self.log(f"Overall peak: {overall_alloc:.2f}GB", category="memory", force=force) else: 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: """Store checkpoint with memory limit to prevent leaks.""" checkpoint = { 'label': label, 'timestamp': time.time(), 'vram_allocated': metrics['vram_allocated'], 'vram_reserved': metrics['vram_reserved'], 'vram_free': metrics['vram_free'], 'ram_process': metrics['ram_process'], 'ram_available': metrics['ram_available'], 'ram_others': metrics['ram_others'] } self.memory_checkpoints.append(checkpoint) # Prevent memory leak by limiting checkpoint history if len(self.memory_checkpoints) > self.max_checkpoints: # Keep first and last N/2 checkpoints for better history coverage mid = self.max_checkpoints // 2 self.memory_checkpoints = (self.memory_checkpoints[:mid] + self.memory_checkpoints[-mid:]) def log_swap_time(self, component_id: Union[int, str], duration: float, component_type: str = "block", force: bool = False) -> None: """ Log swap timing information for BlockSwap operations Args: component_id: Identifier for the component being swapped duration: Duration of the swap in seconds component_type: Type of component ('block' or other) force: If True, always log regardless of enabled state """ if self.enabled or force: # Store timing data self.swap_times.append({ 'component_id': component_id, 'component_type': component_type, 'duration': duration, }) # Format message based on component type if component_type == "block": message = f"Block {component_id} swap: {duration*1000:.2f}ms" else: message = f"{component_type} {component_id} swap: {duration*1000:.2f}ms" self.log(message, category="blockswap", force=force) def get_swap_summary(self) -> Dict[str, Any]: """Get summary of swap operations for analysis""" if not self.swap_times: return {} # Group by component type block_swaps = [s for s in self.swap_times if s['component_type'] == 'block'] io_swaps = [s for s in self.swap_times if s['component_type'] != 'block'] # Calculate statistics summary = { 'total_swaps': len(self.swap_times), 'block_swaps': len(block_swaps), 'io_swaps': len(io_swaps), } if block_swaps: block_times = [s['duration'] for s in block_swaps] summary['block_avg_ms'] = sum(block_times) * 1000 / len(block_times) summary['block_total_ms'] = sum(block_times) * 1000 summary['block_min_ms'] = min(block_times) * 1000 summary['block_max_ms'] = max(block_times) * 1000 # Track which blocks are swapped most frequently block_frequency = {} for swap in block_swaps: block_id = swap['component_id'] block_frequency[block_id] = block_frequency.get(block_id, 0) + 1 summary['most_swapped_block'] = max(block_frequency, key=block_frequency.get) summary['most_swapped_count'] = block_frequency[summary['most_swapped_block']] if io_swaps: io_times = [s['duration'] for s in io_swaps] summary['io_avg_ms'] = sum(io_times) * 1000 / len(io_times) summary['io_total_ms'] = sum(io_times) * 1000 # Track which I/O components are swapped io_components = list(set(s['component_id'] for s in io_swaps)) summary['io_components_swapped'] = io_components # VRAM efficiency metrics if self.vram_history: summary['peak_vram_gb'] = max(self.vram_history) summary['avg_vram_gb'] = sum(self.vram_history) / len(self.vram_history) summary['vram_variation_gb'] = max(self.vram_history) - min(self.vram_history) return summary def clear_history(self) -> None: """Clear all history tracking""" self.timers.clear() self.memory_checkpoints.clear() self.swap_times.clear() self.vram_history.clear() self.timer_hierarchy.clear() self.timer_durations.clear() self.timer_messages.clear() self.active_timer_stack.clear() self.phase_vram_peaks_alloc.clear() self.phase_vram_peaks_rsv.clear() self.phase_ram_peaks.clear() self.current_phase = None