283 lines
9.2 KiB
Python
283 lines
9.2 KiB
Python
"""
|
|
Runtime debugging utilities for WanVideo Activation Editor
|
|
Can be imported and used anywhere for debugging
|
|
"""
|
|
|
|
import torch
|
|
import os
|
|
import json
|
|
from datetime import datetime
|
|
from typing import Any, Dict, Optional, List
|
|
import traceback
|
|
|
|
# Global debug flags
|
|
DEBUG = os.environ.get('WAN_ACTIVATION_DEBUG', '0') == '1'
|
|
VERBOSE = os.environ.get('WAN_ACTIVATION_VERBOSE', '0') == '1'
|
|
TRACE = os.environ.get('WAN_ACTIVATION_TRACE', '0') == '1'
|
|
|
|
# DEFAULT: Verbose logging is OFF
|
|
# To enable verbose logging, either:
|
|
# 1. Set log_level to "verbose" in WanVideoActivationEditor node
|
|
# 2. Set environment variable: export WAN_ACTIVATION_VERBOSE=1
|
|
# 3. Uncomment ONE line below:
|
|
# DEBUG = True # Basic debug output only
|
|
# VERBOSE = True # Verbose output (automatically enables DEBUG too)
|
|
|
|
# Runtime log level override
|
|
_LOG_LEVEL_OVERRIDE = None
|
|
|
|
def set_log_level(level: str):
|
|
"""Set the runtime log level. Levels: 'off', 'basic', 'verbose', 'trace'"""
|
|
global _LOG_LEVEL_OVERRIDE, DEBUG, VERBOSE, TRACE
|
|
_LOG_LEVEL_OVERRIDE = level
|
|
|
|
if level == "off":
|
|
DEBUG = False
|
|
VERBOSE = False
|
|
TRACE = False
|
|
elif level == "basic":
|
|
DEBUG = True
|
|
VERBOSE = False
|
|
TRACE = False
|
|
elif level == "verbose":
|
|
DEBUG = True
|
|
VERBOSE = True
|
|
TRACE = False
|
|
elif level == "trace":
|
|
DEBUG = True
|
|
VERBOSE = True
|
|
TRACE = True
|
|
|
|
# Immediate feedback
|
|
if level != "off":
|
|
print(f"[WanActivationEditor] Log level set to: {level}")
|
|
|
|
def get_log_level() -> str:
|
|
"""Get current log level"""
|
|
if _LOG_LEVEL_OVERRIDE:
|
|
return _LOG_LEVEL_OVERRIDE
|
|
elif TRACE:
|
|
return "trace"
|
|
elif VERBOSE:
|
|
return "verbose"
|
|
elif DEBUG:
|
|
return "basic"
|
|
else:
|
|
return "off"
|
|
|
|
def debug_print(message: str, level: str = "INFO"):
|
|
"""Print debug message if debugging is enabled"""
|
|
if DEBUG or VERBOSE:
|
|
timestamp = datetime.now().strftime('%H:%M:%S.%f')[:-3]
|
|
print(f"[WAN_DEBUG {timestamp}] [{level}] {message}")
|
|
|
|
def verbose_print(message: str):
|
|
"""Print verbose message if verbose mode is enabled"""
|
|
if VERBOSE:
|
|
debug_print(message, "VERBOSE")
|
|
|
|
def trace_print(message: str):
|
|
"""Print trace message if trace mode is enabled"""
|
|
if TRACE:
|
|
debug_print(message, "TRACE")
|
|
# Also print stack trace in trace mode
|
|
stack = traceback.extract_stack()[:-1]
|
|
for frame in stack[-3:]: # Last 3 frames
|
|
print(f" -> {frame.filename}:{frame.lineno} in {frame.name}")
|
|
|
|
def debug_tensor(tensor: torch.Tensor, name: str = "tensor", detailed: bool = False):
|
|
"""Debug print tensor information"""
|
|
if not (DEBUG or VERBOSE):
|
|
return
|
|
|
|
info = [
|
|
f"{name}:",
|
|
f" shape: {tensor.shape}",
|
|
f" device: {tensor.device}",
|
|
f" dtype: {tensor.dtype}",
|
|
]
|
|
|
|
if detailed or VERBOSE:
|
|
info.extend([
|
|
f" mean: {tensor.mean().item():.6f}",
|
|
f" std: {tensor.std().item():.6f}",
|
|
f" min: {tensor.min().item():.6f}",
|
|
f" max: {tensor.max().item():.6f}",
|
|
f" norm: {torch.norm(tensor).item():.6f}",
|
|
])
|
|
|
|
# Check for special values
|
|
if torch.isnan(tensor).any():
|
|
info.append(f" WARNING: Contains NaN values!")
|
|
if torch.isinf(tensor).any():
|
|
info.append(f" WARNING: Contains Inf values!")
|
|
if (tensor == 0).all():
|
|
info.append(f" WARNING: All zeros!")
|
|
|
|
debug_print("\n".join(info))
|
|
|
|
def debug_config(config: Dict[str, Any], name: str = "config"):
|
|
"""Debug print configuration dictionary"""
|
|
if not (DEBUG or VERBOSE):
|
|
return
|
|
|
|
debug_print(f"{name}: {json.dumps(config, indent=2, default=str)}")
|
|
|
|
def debug_injection_point(block_idx: int, context: torch.Tensor,
|
|
injection: Optional[torch.Tensor] = None,
|
|
strength: float = 0.0):
|
|
"""Debug print at injection point"""
|
|
if not (DEBUG or VERBOSE or TRACE):
|
|
return
|
|
|
|
debug_print(f"\n=== Injection Point: Block {block_idx} ===")
|
|
debug_tensor(context, f"context (block {block_idx})")
|
|
|
|
if injection is not None:
|
|
debug_tensor(injection, f"injection (block {block_idx})")
|
|
|
|
# Calculate difference
|
|
if context.shape == injection.shape:
|
|
diff = (context - injection).abs()
|
|
diff_percent = (diff > 0.01).float().mean().item() * 100
|
|
debug_print(f" difference: {diff.mean().item():.6f} ({diff_percent:.1f}% different)")
|
|
else:
|
|
debug_print(f" WARNING: Shape mismatch! {context.shape} vs {injection.shape}")
|
|
|
|
debug_print(f" strength: {strength}")
|
|
|
|
def debug_memory(prefix: str = ""):
|
|
"""Debug print memory usage"""
|
|
if not (DEBUG or VERBOSE):
|
|
return
|
|
|
|
if torch.cuda.is_available():
|
|
allocated = torch.cuda.memory_allocated() / 1024**3
|
|
reserved = torch.cuda.memory_reserved() / 1024**3
|
|
debug_print(f"{prefix}CUDA Memory - Allocated: {allocated:.2f}GB, Reserved: {reserved:.2f}GB")
|
|
|
|
def create_debug_checkpoint(name: str, data: Dict[str, Any]):
|
|
"""Create a debug checkpoint that can be loaded later"""
|
|
if not (DEBUG or VERBOSE):
|
|
return
|
|
|
|
checkpoint = {
|
|
'name': name,
|
|
'timestamp': datetime.now().isoformat(),
|
|
'data': data
|
|
}
|
|
|
|
# Save to temp file if TRACE mode
|
|
if TRACE:
|
|
filename = f"/tmp/wan_debug_{name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
|
|
try:
|
|
with open(filename, 'w') as f:
|
|
json.dump(checkpoint, f, indent=2, default=str)
|
|
trace_print(f"Debug checkpoint saved to: {filename}")
|
|
except Exception as e:
|
|
debug_print(f"Failed to save checkpoint: {e}", "ERROR")
|
|
|
|
def assert_shape(tensor: torch.Tensor, expected_shape: tuple, name: str = "tensor"):
|
|
"""Assert tensor has expected shape, with helpful debug info"""
|
|
if tensor.shape != expected_shape:
|
|
error_msg = f"{name} shape mismatch: expected {expected_shape}, got {tensor.shape}"
|
|
debug_print(error_msg, "ERROR")
|
|
debug_tensor(tensor, name, detailed=True)
|
|
raise ValueError(error_msg)
|
|
|
|
def compare_embeddings(embed1: torch.Tensor, embed2: torch.Tensor,
|
|
name1: str = "embed1", name2: str = "embed2") -> Dict[str, float]:
|
|
"""Compare two embeddings and return statistics"""
|
|
stats = {}
|
|
|
|
if embed1.shape != embed2.shape:
|
|
debug_print(f"Shape mismatch: {name1}={embed1.shape} vs {name2}={embed2.shape}", "WARNING")
|
|
return stats
|
|
|
|
# Move to same device for comparison
|
|
if embed1.device != embed2.device:
|
|
embed2 = embed2.to(embed1.device)
|
|
|
|
# Calculate statistics
|
|
diff = (embed1 - embed2).abs()
|
|
stats['mean_diff'] = diff.mean().item()
|
|
stats['max_diff'] = diff.max().item()
|
|
stats['percent_different'] = (diff > 0.01).float().mean().item() * 100
|
|
|
|
# Cosine similarity
|
|
embed1_flat = embed1.flatten()
|
|
embed2_flat = embed2.flatten()
|
|
cos_sim = torch.nn.functional.cosine_similarity(embed1_flat.unsqueeze(0),
|
|
embed2_flat.unsqueeze(0))
|
|
stats['cosine_similarity'] = cos_sim.item()
|
|
|
|
if DEBUG or VERBOSE:
|
|
debug_print(f"\nComparison {name1} vs {name2}:")
|
|
for key, value in stats.items():
|
|
debug_print(f" {key}: {value:.4f}")
|
|
|
|
return stats
|
|
|
|
# Decorators for function debugging
|
|
def debug_function(func):
|
|
"""Decorator to debug function calls"""
|
|
def wrapper(*args, **kwargs):
|
|
if DEBUG or VERBOSE:
|
|
debug_print(f"Calling {func.__name__}")
|
|
if VERBOSE:
|
|
debug_print(f" args: {args}")
|
|
debug_print(f" kwargs: {kwargs}")
|
|
|
|
result = func(*args, **kwargs)
|
|
|
|
if DEBUG or VERBOSE:
|
|
debug_print(f"Completed {func.__name__}")
|
|
|
|
return result
|
|
return wrapper
|
|
|
|
def trace_function(func):
|
|
"""Decorator for detailed function tracing"""
|
|
def wrapper(*args, **kwargs):
|
|
if TRACE:
|
|
trace_print(f"ENTER {func.__name__}")
|
|
start_time = datetime.now()
|
|
|
|
try:
|
|
result = func(*args, **kwargs)
|
|
|
|
if TRACE:
|
|
elapsed = (datetime.now() - start_time).total_seconds()
|
|
trace_print(f"EXIT {func.__name__} (took {elapsed:.3f}s)")
|
|
|
|
return result
|
|
except Exception as e:
|
|
if TRACE:
|
|
trace_print(f"ERROR in {func.__name__}: {e}")
|
|
raise
|
|
|
|
return wrapper
|
|
|
|
# Context manager for debugging blocks
|
|
class DebugBlock:
|
|
def __init__(self, name: str):
|
|
self.name = name
|
|
self.start_time = None
|
|
|
|
def __enter__(self):
|
|
if DEBUG or VERBOSE:
|
|
debug_print(f"=== BEGIN {self.name} ===")
|
|
self.start_time = datetime.now()
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
if DEBUG or VERBOSE:
|
|
elapsed = (datetime.now() - self.start_time).total_seconds() if self.start_time else 0
|
|
if exc_type is None:
|
|
debug_print(f"=== END {self.name} (took {elapsed:.3f}s) ===")
|
|
else:
|
|
debug_print(f"=== ERROR in {self.name}: {exc_val} ===", "ERROR")
|
|
|
|
if exc_type is not None and TRACE:
|
|
traceback.print_exc()
|