677 lines
33 KiB
Python
677 lines
33 KiB
Python
"""
|
||
Runtime patching for WanVideo transformer blocks.
|
||
Enables block-level activation editing without modifying WanVideoWrapper.
|
||
"""
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from typing import Dict, List, Optional, Any
|
||
import weakref
|
||
import os
|
||
import math
|
||
|
||
# Import debug utilities
|
||
from .debug_utils import (
|
||
debug_print, verbose_print, trace_print, debug_tensor,
|
||
debug_injection_point, debug_memory, DebugBlock,
|
||
debug_function, trace_function, DEBUG, VERBOSE
|
||
)
|
||
|
||
|
||
class ActivationPatcher:
|
||
"""Manages runtime patching of transformer blocks for activation editing."""
|
||
|
||
def __init__(self):
|
||
self.active_patches = weakref.WeakKeyDictionary()
|
||
self._text_embedding_layer = None
|
||
self._should_inject_hidden_states = False # Flag to enable hidden state injection
|
||
|
||
def patch_model(self, model_patcher, activation_config: Dict[str, Any]):
|
||
"""
|
||
Patch transformer blocks to enable activation editing.
|
||
|
||
Args:
|
||
model_patcher: ComfyUI ModelPatcher instance
|
||
activation_config: Configuration dict with:
|
||
- active_blocks: List of block indices to modify
|
||
- injection_strength: Blend strength (0-1)
|
||
- injection_embeds: Alternative embeddings to inject
|
||
"""
|
||
try:
|
||
# Access the actual model inside ComfyUI's ModelPatcher
|
||
actual_model = self._find_model(model_patcher)
|
||
if actual_model is None:
|
||
raise RuntimeError("Could not access model from ModelPatcher")
|
||
|
||
# Find transformer blocks
|
||
transformer_blocks = self._find_transformer_blocks(actual_model)
|
||
if transformer_blocks is None:
|
||
debug_print("ERROR: Failed to locate transformer blocks!")
|
||
raise RuntimeError("Failed to locate transformer blocks for patching")
|
||
|
||
debug_print(f"Found {len(transformer_blocks)} transformer blocks")
|
||
debug_print(f"Transformer blocks type: {type(transformer_blocks)}")
|
||
debug_print(f"First block type: {type(transformer_blocks[0]) if len(transformer_blocks) > 0 else 'empty'}")
|
||
|
||
# Find text_embedding layer for dimension projection
|
||
self._find_text_embedding_layer(actual_model)
|
||
|
||
# Store original forward methods
|
||
if model_patcher not in self.active_patches:
|
||
self.active_patches[model_patcher] = {
|
||
'original_forwards': {},
|
||
'config': activation_config
|
||
}
|
||
debug_print("Created new patch entry for model")
|
||
else:
|
||
debug_print("Model already has patch entry")
|
||
|
||
active_blocks = activation_config.get('active_blocks', [])
|
||
injection_strength = activation_config.get('injection_strength', 0.5)
|
||
injection_embeds = activation_config.get('injection_embeds')
|
||
injection_mode = activation_config.get('injection_mode', 'context')
|
||
|
||
# Set injection mode flags
|
||
self._should_inject_hidden_states = injection_mode in ['hidden_states', 'both']
|
||
self._should_inject_context = injection_mode in ['context', 'both']
|
||
|
||
debug_print(f"Config - active_blocks: {active_blocks}")
|
||
debug_print(f"Config - injection_strength: {injection_strength}")
|
||
debug_print(f"Config - injection_mode: {injection_mode}")
|
||
debug_print(f"Config - inject hidden states: {self._should_inject_hidden_states}")
|
||
debug_print(f"Config - inject context: {self._should_inject_context}")
|
||
debug_print(f"Config - injection_embeds type: {type(injection_embeds)}")
|
||
if injection_embeds is not None:
|
||
if hasattr(injection_embeds, 'shape'):
|
||
debug_print(f"Config - injection_embeds shape: {injection_embeds.shape}")
|
||
elif isinstance(injection_embeds, dict):
|
||
debug_print(f"Config - injection_embeds keys: {list(injection_embeds.keys())}")
|
||
|
||
# Patch each active block
|
||
patched_count = 0
|
||
debug_print(f"Starting to patch {len(active_blocks)} blocks...")
|
||
for block_idx in active_blocks:
|
||
if block_idx >= len(transformer_blocks):
|
||
debug_print(f"Warning: Block {block_idx} exceeds available blocks ({len(transformer_blocks)})")
|
||
continue
|
||
|
||
debug_print(f"Patching block {block_idx}...")
|
||
|
||
block = transformer_blocks[block_idx]
|
||
|
||
# Store original if not already stored
|
||
if block_idx not in self.active_patches[model_patcher]['original_forwards']:
|
||
self.active_patches[model_patcher]['original_forwards'][block_idx] = block.forward
|
||
|
||
# Create patched forward method
|
||
original_forward = self.active_patches[model_patcher]['original_forwards'][block_idx]
|
||
patched_forward = self._create_patched_forward(block_idx, original_forward, injection_strength, injection_embeds)
|
||
|
||
# Store reference to help with debugging
|
||
patched_forward._block_idx = block_idx
|
||
patched_forward._is_wan_patch = True
|
||
|
||
block.forward = patched_forward
|
||
|
||
# Verify patch was applied
|
||
if block.forward == patched_forward:
|
||
debug_print(f"✓ Block {block_idx} patched successfully")
|
||
debug_print(f" Forward method type: {type(block.forward)}")
|
||
debug_print(f" Has injection embeds: {injection_embeds is not None}")
|
||
if injection_embeds is not None:
|
||
if hasattr(injection_embeds, 'shape'):
|
||
debug_print(f" Injection shape: {injection_embeds.shape}")
|
||
elif isinstance(injection_embeds, dict):
|
||
debug_print(f" Injection is dict with keys: {list(injection_embeds.keys())}")
|
||
patched_count += 1
|
||
else:
|
||
debug_print(f"✗ Block {block_idx} patch failed!")
|
||
|
||
debug_print(f"Successfully patched {patched_count}/{len(active_blocks)} blocks")
|
||
|
||
# Add a test to ensure patches will work
|
||
if patched_count > 0 and injection_embeds is not None:
|
||
debug_print(f"Injection embeddings ready: shape={injection_embeds.shape if hasattr(injection_embeds, 'shape') else 'unknown'}")
|
||
|
||
# Return success status
|
||
return patched_count > 0
|
||
|
||
except Exception as e:
|
||
debug_print(f"ERROR in patch_model: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return False
|
||
|
||
def unpatch_model(self, model_patcher):
|
||
"""Remove patches from a model."""
|
||
if model_patcher not in self.active_patches:
|
||
return
|
||
|
||
try:
|
||
# Access the model
|
||
actual_model = self._find_model(model_patcher)
|
||
if actual_model is None:
|
||
return
|
||
|
||
# Find transformer blocks
|
||
transformer_blocks = self._find_transformer_blocks(actual_model)
|
||
if transformer_blocks is None:
|
||
return
|
||
|
||
# Restore original forwards
|
||
patch_info = self.active_patches[model_patcher]
|
||
for block_idx, original_forward in patch_info['original_forwards'].items():
|
||
if block_idx < len(transformer_blocks):
|
||
transformer_blocks[block_idx].forward = original_forward
|
||
verbose_print(f"Restored block {block_idx}")
|
||
|
||
# Clean up
|
||
del self.active_patches[model_patcher]
|
||
debug_print("Model unpatched")
|
||
|
||
except Exception as e:
|
||
debug_print(f"Error during unpatching: {e}")
|
||
|
||
def _find_model(self, model_patcher):
|
||
"""Find the actual model inside the ModelPatcher."""
|
||
# Try multiple ways to access the model
|
||
# WanVideoWrapper uses model.model.diffusion_model
|
||
if hasattr(model_patcher, 'model'):
|
||
model = model_patcher.model
|
||
if hasattr(model, 'diffusion_model'):
|
||
debug_print("Found model at model_patcher.model.diffusion_model")
|
||
return model.diffusion_model
|
||
elif hasattr(model, 'model') and hasattr(model.model, 'diffusion_model'):
|
||
debug_print("Found model at model_patcher.model.model.diffusion_model")
|
||
return model.model.diffusion_model
|
||
else:
|
||
debug_print("Found model at model_patcher.model")
|
||
return model
|
||
elif hasattr(model_patcher, 'transformer_blocks') or hasattr(model_patcher, 'blocks'):
|
||
debug_print("Model patcher appears to be the model itself")
|
||
return model_patcher
|
||
return None
|
||
|
||
def _find_transformer_blocks(self, model):
|
||
"""Find transformer blocks in the model."""
|
||
# Check various possible paths to transformer blocks
|
||
# WanVideoWrapper uses model.blocks
|
||
paths_to_try = [
|
||
lambda m: m.blocks if hasattr(m, 'blocks') else None,
|
||
lambda m: m.transformer_blocks if hasattr(m, 'transformer_blocks') else None,
|
||
lambda m: m.transformer.transformer_blocks if hasattr(m, 'transformer') and hasattr(m.transformer, 'transformer_blocks') else None,
|
||
lambda m: m.transformer.blocks if hasattr(m, 'transformer') and hasattr(m.transformer, 'blocks') else None,
|
||
lambda m: m.model.transformer_blocks if hasattr(m, 'model') and hasattr(m.model, 'transformer_blocks') else None,
|
||
lambda m: m.model.blocks if hasattr(m, 'model') and hasattr(m.model, 'blocks') else None,
|
||
lambda m: m._modules.get('transformer_blocks') if hasattr(m, '_modules') else None,
|
||
lambda m: m._modules.get('blocks') if hasattr(m, '_modules') else None,
|
||
]
|
||
|
||
for path_func in paths_to_try:
|
||
try:
|
||
result = path_func(model)
|
||
if result is not None and hasattr(result, '__len__'):
|
||
debug_print(f"Found blocks at: {path_func.__code__.co_consts[0]}")
|
||
return result
|
||
except (AttributeError, TypeError):
|
||
continue
|
||
|
||
return None
|
||
|
||
def _find_text_embedding_layer(self, model):
|
||
"""Find and store the text embedding layer."""
|
||
self._text_embedding_layer = None
|
||
self._text_len = 512 # Default value
|
||
|
||
# Search for text_embedding layer in various locations
|
||
if hasattr(model, 'text_embedding'):
|
||
self._text_embedding_layer = model.text_embedding
|
||
if hasattr(model, 'text_len'):
|
||
self._text_len = model.text_len
|
||
verbose_print(f"Found text_embedding layer (text_len={self._text_len})")
|
||
elif hasattr(model, 'transformer') and hasattr(model.transformer, 'text_embedding'):
|
||
self._text_embedding_layer = model.transformer.text_embedding
|
||
if hasattr(model.transformer, 'text_len'):
|
||
self._text_len = model.transformer.text_len
|
||
verbose_print(f"Found text_embedding layer in transformer (text_len={self._text_len})")
|
||
|
||
def _create_patched_forward(self, block_idx, original_forward, injection_strength, injection_embeds):
|
||
"""Create a patched forward function for a transformer block."""
|
||
patcher_instance = self
|
||
|
||
@torch.compiler.disable
|
||
def patched_forward(x, *args, **kwargs):
|
||
# Always print for first few calls to verify
|
||
if not hasattr(patched_forward, '_call_count'):
|
||
patched_forward._call_count = 0
|
||
patched_forward._x_shapes = [] # Track x shapes
|
||
patched_forward._call_count += 1
|
||
|
||
# Log first few calls for debugging
|
||
if patched_forward._call_count <= 3:
|
||
verbose_print(f"\n[ACTIVATION PATCH] Block {block_idx} forward called!")
|
||
verbose_print(f" Call #{patched_forward._call_count}")
|
||
verbose_print(f" x.shape: {x.shape if hasattr(x, 'shape') else 'unknown'}")
|
||
verbose_print(f" args count: {len(args)}")
|
||
if len(args) > 5:
|
||
verbose_print(f" context (arg 5) shape: {args[5].shape if hasattr(args[5], 'shape') else type(args[5])}")
|
||
verbose_print(f" kwargs keys: {list(kwargs.keys())}")
|
||
verbose_print(f" injection_embeds is None: {injection_embeds is None}")
|
||
verbose_print(f" Injection mode - Hidden states: {patcher_instance._should_inject_hidden_states}, Context: {patcher_instance._should_inject_context}")
|
||
|
||
# Track x shapes to understand data flow
|
||
if hasattr(x, 'shape') and x.shape not in patched_forward._x_shapes:
|
||
patched_forward._x_shapes.append(x.shape)
|
||
verbose_print(f" NEW x shape detected: {x.shape}")
|
||
|
||
if patched_forward._call_count == 1:
|
||
debug_print(f"Block {block_idx}: Forward method called (first time)")
|
||
debug_print(f" x shape: {x.shape if hasattr(x, 'shape') else 'unknown'}")
|
||
debug_print(f" args count: {len(args)}")
|
||
debug_print(f" kwargs keys: {list(kwargs.keys())}")
|
||
|
||
# EXPERIMENTAL: Try hidden state injection
|
||
modified_x = x
|
||
if injection_embeds is not None and hasattr(x, 'shape'):
|
||
# Check if we should inject into hidden states
|
||
if patcher_instance._should_inject_hidden_states:
|
||
try:
|
||
# Process hidden state injection
|
||
modified_x = patcher_instance._process_hidden_state_injection(
|
||
block_idx, x, injection_embeds, injection_strength
|
||
)
|
||
|
||
if modified_x is not None and modified_x is not x:
|
||
debug_print(f"Block {block_idx}: Applied hidden state injection")
|
||
verbose_print(f" Original x norm: {x.norm().item():.4f}")
|
||
verbose_print(f" Modified x norm: {modified_x.norm().item():.4f}")
|
||
|
||
# Use modified x for forward pass
|
||
x = modified_x
|
||
except Exception as e:
|
||
debug_print(f"Block {block_idx}: Error in hidden state injection: {e}")
|
||
|
||
# Original context injection logic
|
||
context = None
|
||
|
||
# Check if context is in kwargs (this is the typical case)
|
||
if 'context' in kwargs:
|
||
context = kwargs['context']
|
||
if patched_forward._call_count <= 3:
|
||
debug_print(f"Block {block_idx}: Found context in kwargs")
|
||
if hasattr(context, 'shape'):
|
||
debug_print(f"Block {block_idx}: Context shape: {context.shape}")
|
||
|
||
# Only process if we have both context and injection embeddings
|
||
if context is not None and injection_embeds is not None and patcher_instance._should_inject_context:
|
||
try:
|
||
# Process injection
|
||
blended_context = patcher_instance._process_injection(
|
||
block_idx, context, injection_embeds, injection_strength
|
||
)
|
||
|
||
if blended_context is not None:
|
||
# Update context in kwargs
|
||
kwargs['context'] = blended_context
|
||
|
||
# Call original forward with modified context
|
||
output = original_forward(x, *args, **kwargs)
|
||
|
||
# Log first successful injection
|
||
if not hasattr(patched_forward, '_logged'):
|
||
debug_print(f"Block {block_idx}: Successfully applied injection")
|
||
patched_forward._logged = True
|
||
|
||
return output
|
||
|
||
except Exception as e:
|
||
debug_print(f"Block {block_idx}: Error during injection: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
# Call original forward with potentially modified x
|
||
return original_forward(x, *args, **kwargs)
|
||
|
||
return patched_forward
|
||
|
||
def _process_injection(self, block_idx, context, injection_embeds, injection_strength):
|
||
"""Process the injection and blend with original context."""
|
||
# Log initial state
|
||
debug_print(f"\n=== Block {block_idx} Injection Processing ===")
|
||
verbose_print(f"Context type: {type(context)}, Injection type: {type(injection_embeds)}")
|
||
|
||
# Check if we should apply projection boost
|
||
if hasattr(self, '_boost_factor') and self._boost_factor > 1.0:
|
||
debug_print(f"Applying projection boost: {self._boost_factor}x")
|
||
|
||
# Extract tensors from dicts if needed
|
||
if isinstance(context, dict):
|
||
context = context.get('prompt_embeds', context)
|
||
verbose_print("Extracted context from dict")
|
||
if isinstance(injection_embeds, dict):
|
||
# Check if this is a latent embedding dict
|
||
if 'latent' in injection_embeds:
|
||
injection_embeds = injection_embeds.get('latent')
|
||
verbose_print("Extracted latent tensor from injection_embeds")
|
||
else:
|
||
injection_embeds = injection_embeds.get('prompt_embeds', injection_embeds)
|
||
verbose_print("Extracted injection_embeds from dict")
|
||
|
||
# Handle list format
|
||
if isinstance(context, list):
|
||
context = context[0] if len(context) > 0 else None
|
||
verbose_print("Extracted context from list")
|
||
if isinstance(injection_embeds, list):
|
||
injection_embeds = injection_embeds[0] if len(injection_embeds) > 0 else None
|
||
verbose_print("Extracted injection_embeds from list")
|
||
|
||
if not torch.is_tensor(context) or not torch.is_tensor(injection_embeds):
|
||
debug_print(f"Block {block_idx}: Invalid tensors - context is tensor: {torch.is_tensor(context)}, injection is tensor: {torch.is_tensor(injection_embeds)}")
|
||
return None
|
||
|
||
# Log tensor info
|
||
debug_print(f"Block {block_idx} shapes - Context: {context.shape}, Injection: {injection_embeds.shape}")
|
||
debug_print(f"Block {block_idx} devices - Context: {context.device}, Injection: {injection_embeds.device}")
|
||
debug_print(f"Block {block_idx} dtypes - Context: {context.dtype}, Injection: {injection_embeds.dtype}")
|
||
|
||
# Ensure same device and dtype
|
||
if injection_embeds.device != context.device or injection_embeds.dtype != context.dtype:
|
||
injection_embeds = injection_embeds.to(device=context.device, dtype=context.dtype)
|
||
verbose_print(f"Moved injection to device={context.device}, dtype={context.dtype}")
|
||
|
||
# Handle dimension mismatch
|
||
if context.shape[-1] != injection_embeds.shape[-1]:
|
||
debug_print(f"Block {block_idx}: Dimension mismatch {injection_embeds.shape[-1]} vs {context.shape[-1]}, projecting...")
|
||
injection_embeds = self._project_embeddings(injection_embeds, context)
|
||
if injection_embeds is None:
|
||
debug_print(f"Block {block_idx}: Projection failed!")
|
||
return None
|
||
debug_print(f"Block {block_idx}: After projection shape: {injection_embeds.shape}")
|
||
|
||
# Align shapes
|
||
injection_embeds = self._align_shapes(injection_embeds, context)
|
||
|
||
# Blend contexts
|
||
blended = (1 - injection_strength) * context + injection_strength * injection_embeds
|
||
|
||
# Calculate and log the actual change
|
||
diff = (blended - context).abs()
|
||
diff_percent = (diff > 0.01).float().mean().item() * 100
|
||
max_diff = diff.max().item()
|
||
mean_diff = diff.mean().item()
|
||
|
||
debug_print(f"Block {block_idx} Blending Results:")
|
||
debug_print(f" - Strength applied: {injection_strength}")
|
||
debug_print(f" - Mean difference: {mean_diff:.6f}")
|
||
debug_print(f" - Max difference: {max_diff:.6f}")
|
||
debug_print(f" - Percent changed: {diff_percent:.1f}%")
|
||
|
||
# Enhanced debugging at injection point
|
||
if VERBOSE:
|
||
debug_injection_point(block_idx, context, injection_embeds, injection_strength)
|
||
|
||
# Log statistics for first block only
|
||
if block_idx == 0 or (block_idx % 10 == 0):
|
||
self._log_blend_statistics(context, injection_embeds, blended, injection_strength)
|
||
|
||
return blended
|
||
|
||
def _project_embeddings(self, injection_embeds, context):
|
||
"""Project embeddings through text_embedding layer if needed."""
|
||
if self._text_embedding_layer is None:
|
||
verbose_print(f"No text_embedding layer available for projection")
|
||
return None
|
||
|
||
try:
|
||
# Pad to text_len if needed
|
||
if injection_embeds.shape[1] < self._text_len:
|
||
padding = torch.zeros(
|
||
injection_embeds.shape[0],
|
||
self._text_len - injection_embeds.shape[1],
|
||
injection_embeds.shape[2],
|
||
device=injection_embeds.device,
|
||
dtype=injection_embeds.dtype
|
||
)
|
||
injection_embeds = torch.cat([injection_embeds, padding], dim=1)
|
||
elif injection_embeds.shape[1] > self._text_len:
|
||
injection_embeds = injection_embeds[:, :self._text_len, :]
|
||
|
||
# Get text embedding layer's device and dtype
|
||
text_emb_params = next(self._text_embedding_layer.parameters())
|
||
|
||
# Project through text_embedding layer
|
||
injection_embeds = injection_embeds.to(device=text_emb_params.device, dtype=text_emb_params.dtype)
|
||
projected = self._text_embedding_layer(injection_embeds)
|
||
|
||
# Move back to context's device/dtype
|
||
return projected.to(device=context.device, dtype=context.dtype)
|
||
|
||
except Exception as e:
|
||
verbose_print(f"Failed to project embeddings: {e}")
|
||
return None
|
||
|
||
def _align_shapes(self, injection_embeds, context):
|
||
"""Align injection embeddings shape to match context."""
|
||
# Align sequence length
|
||
if context.shape[1] != injection_embeds.shape[1]:
|
||
if injection_embeds.shape[1] > context.shape[1]:
|
||
injection_embeds = injection_embeds[:, :context.shape[1], :]
|
||
else:
|
||
padding = torch.zeros(
|
||
injection_embeds.shape[0],
|
||
context.shape[1] - injection_embeds.shape[1],
|
||
injection_embeds.shape[2],
|
||
device=injection_embeds.device,
|
||
dtype=injection_embeds.dtype
|
||
)
|
||
injection_embeds = torch.cat([injection_embeds, padding], dim=1)
|
||
|
||
# Align batch dimension
|
||
if context.shape[0] != injection_embeds.shape[0]:
|
||
injection_embeds = injection_embeds.expand(context.shape[0], -1, -1)
|
||
|
||
return injection_embeds
|
||
|
||
def _log_blend_statistics(self, context, injection_embeds, blended, injection_strength):
|
||
"""Log statistics about the blending operation."""
|
||
with DebugBlock("blend_statistics"):
|
||
context_norm = torch.norm(context).item()
|
||
injection_norm = torch.norm(injection_embeds).item()
|
||
blended_norm = torch.norm(blended).item()
|
||
diff_norm = torch.norm(blended - context).item()
|
||
|
||
debug_print(f"Injection statistics:")
|
||
debug_print(f" - Strength: {injection_strength}")
|
||
debug_print(f" - Context norm: {context_norm:.3f}")
|
||
debug_print(f" - Injection norm: {injection_norm:.3f}")
|
||
debug_print(f" - Blended norm: {blended_norm:.3f}")
|
||
debug_print(f" - Change: {(diff_norm/context_norm)*100:.1f}%")
|
||
|
||
# More detailed tensor analysis
|
||
debug_tensor(context, "context", detailed=True)
|
||
debug_tensor(injection_embeds, "injection", detailed=True)
|
||
debug_tensor(blended, "blended", detailed=True)
|
||
|
||
# Analyze sequence composition if large
|
||
if context.shape[1] > 1000 and VERBOSE:
|
||
seq_len = context.shape[1]
|
||
text_tokens = 512
|
||
video_tokens = seq_len - text_tokens
|
||
debug_print(f"\nSequence analysis:")
|
||
debug_print(f" - Total length: {seq_len}")
|
||
debug_print(f" - Text tokens: {text_tokens}")
|
||
debug_print(f" - Video tokens: {video_tokens}")
|
||
|
||
# Try to factor video tokens
|
||
for frames in [16, 24, 32, 41, 48, 64]:
|
||
if video_tokens % frames == 0:
|
||
tokens_per_frame = video_tokens // frames
|
||
sqrt_val = math.sqrt(tokens_per_frame)
|
||
if sqrt_val == int(sqrt_val):
|
||
debug_print(f" - {frames} frames × {int(sqrt_val)}×{int(sqrt_val)} patches")
|
||
|
||
def _prepare_hidden_states_for_injection(self, injection_embeds, target_shape):
|
||
"""
|
||
Prepare injection embeddings to be compatible with hidden states.
|
||
This might involve projecting raw T5 embeddings through initial layers.
|
||
"""
|
||
debug_print("\n=== Preparing hidden states for injection ===")
|
||
debug_print(f"Target shape: {target_shape}")
|
||
debug_print(f"Injection embeds shape: {injection_embeds.shape if hasattr(injection_embeds, 'shape') else 'unknown'}")
|
||
|
||
# For now, we can only inject if dimensions already match
|
||
# In the future, we could project T5 embeddings through the initial layers
|
||
# to get them into hidden state space
|
||
|
||
if hasattr(injection_embeds, 'shape') and len(target_shape) >= 3:
|
||
if injection_embeds.shape[-1] == target_shape[-1]:
|
||
debug_print("Dimensions match! Can inject directly.")
|
||
return injection_embeds
|
||
else:
|
||
debug_print(f"Dimension mismatch: {injection_embeds.shape[-1]} vs {target_shape[-1]}")
|
||
debug_print("Would need to project through model layers first (not implemented)")
|
||
return None
|
||
|
||
return None
|
||
|
||
def _process_hidden_state_injection(self, block_idx, x, injection_embeds, injection_strength):
|
||
"""Process injection into hidden states instead of cross-attention context."""
|
||
try:
|
||
debug_print(f"\n=== Block {block_idx} Hidden State Injection ===")
|
||
debug_print(f"Hidden state (x) shape: {x.shape}")
|
||
debug_print(f"Injection embeds type: {type(injection_embeds)}")
|
||
|
||
# Extract injection tensor from dict if needed
|
||
if isinstance(injection_embeds, dict):
|
||
injection_embeds = injection_embeds.get('prompt_embeds', injection_embeds)
|
||
if isinstance(injection_embeds, list):
|
||
injection_embeds = injection_embeds[0] if len(injection_embeds) > 0 else None
|
||
|
||
if injection_embeds is None or not torch.is_tensor(injection_embeds):
|
||
debug_print(f"Block {block_idx}: No valid injection tensor for hidden states")
|
||
return x
|
||
|
||
debug_print(f"Injection tensor shape: {injection_embeds.shape}")
|
||
|
||
# Hidden states typically have shape [batch, seq_len, hidden_dim]
|
||
# Injection embeds might be [seq_len, embed_dim] or [batch, seq_len, embed_dim]
|
||
|
||
# Ensure injection embeds have batch dimension
|
||
if injection_embeds.dim() == 2:
|
||
injection_embeds = injection_embeds.unsqueeze(0)
|
||
|
||
# Ensure same device and dtype
|
||
if injection_embeds.device != x.device or injection_embeds.dtype != x.dtype:
|
||
injection_embeds = injection_embeds.to(device=x.device, dtype=x.dtype)
|
||
|
||
# Check if dimensions match
|
||
if x.shape[-1] != injection_embeds.shape[-1]:
|
||
debug_print(f"Block {block_idx}: Dimension mismatch for hidden states {x.shape[-1]} vs {injection_embeds.shape[-1]}")
|
||
|
||
# If injection is raw T5 embeddings (4096) and x is hidden states (5120), we can't directly blend
|
||
# We would need to project through the model's layers first
|
||
if injection_embeds.shape[-1] == 4096 and x.shape[-1] == 5120:
|
||
debug_print(f"Block {block_idx}: Cannot directly inject T5 embeddings into hidden states - dimensions incompatible")
|
||
return x
|
||
|
||
# For other mismatches, we can't proceed
|
||
return x
|
||
|
||
# Align sequence lengths
|
||
if x.shape[1] != injection_embeds.shape[1]:
|
||
if injection_embeds.shape[1] > x.shape[1]:
|
||
# Truncate injection to match x
|
||
injection_embeds = injection_embeds[:, :x.shape[1], :]
|
||
else:
|
||
# Pad injection to match x
|
||
padding = torch.zeros(
|
||
injection_embeds.shape[0],
|
||
x.shape[1] - injection_embeds.shape[1],
|
||
injection_embeds.shape[2],
|
||
device=injection_embeds.device,
|
||
dtype=injection_embeds.dtype
|
||
)
|
||
injection_embeds = torch.cat([injection_embeds, padding], dim=1)
|
||
|
||
# Align batch dimension
|
||
if x.shape[0] != injection_embeds.shape[0]:
|
||
injection_embeds = injection_embeds.expand(x.shape[0], -1, -1)
|
||
|
||
# Blend hidden states
|
||
blended_x = (1 - injection_strength) * x + injection_strength * injection_embeds
|
||
|
||
# Calculate and log the change
|
||
diff = (blended_x - x).abs()
|
||
diff_percent = (diff > 0.01).float().mean().item() * 100
|
||
max_diff = diff.max().item()
|
||
mean_diff = diff.mean().item()
|
||
|
||
debug_print(f"Block {block_idx} Hidden State Blending Results:")
|
||
debug_print(f" - Strength applied: {injection_strength}")
|
||
debug_print(f" - Mean difference: {mean_diff:.6f}")
|
||
debug_print(f" - Max difference: {max_diff:.6f}")
|
||
debug_print(f" - Percent changed: {diff_percent:.1f}%")
|
||
|
||
# Additional validation metrics
|
||
if VERBOSE:
|
||
# Check if the blending preserved important properties
|
||
x_norm = x.norm(dim=-1, keepdim=True)
|
||
blended_norm = blended_x.norm(dim=-1, keepdim=True)
|
||
norm_ratio = (blended_norm / (x_norm + 1e-6)).mean().item()
|
||
|
||
debug_print(f" - Norm preservation ratio: {norm_ratio:.4f} (1.0 = perfect preservation)")
|
||
|
||
# Check if we're actually changing the hidden states meaningfully
|
||
if diff_percent < 1.0:
|
||
debug_print(f" - WARNING: Very low change detected! Hidden states barely modified.")
|
||
elif diff_percent > 90.0:
|
||
debug_print(f" - WARNING: Extreme change detected! May destabilize generation.")
|
||
|
||
return blended_x
|
||
|
||
except Exception as e:
|
||
debug_print(f"Block {block_idx}: Error in hidden state injection: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return x
|
||
|
||
|
||
class ContextPreprocessor:
|
||
"""Handles preprocessing of injection embeddings to match transformer expectations."""
|
||
|
||
@staticmethod
|
||
def prepare_injection_context(injection_embeds, reference_context=None):
|
||
"""
|
||
Prepare injection embeddings to match the format expected by transformer blocks.
|
||
|
||
Args:
|
||
injection_embeds: Raw injection embeddings (from text encoder)
|
||
reference_context: Optional reference for shape matching
|
||
|
||
Returns:
|
||
Processed context ready for injection
|
||
"""
|
||
# Extract tensor from dict if needed
|
||
if isinstance(injection_embeds, dict):
|
||
injection_tensor = injection_embeds.get('prompt_embeds')
|
||
if isinstance(injection_tensor, list):
|
||
injection_tensor = injection_tensor[0] if len(injection_tensor) > 0 else None
|
||
else:
|
||
injection_tensor = injection_embeds
|
||
|
||
if injection_tensor is None:
|
||
return None
|
||
|
||
# Ensure 2D tensor [seq_len, dim]
|
||
if injection_tensor.dim() == 3 and injection_tensor.shape[0] == 1:
|
||
injection_tensor = injection_tensor.squeeze(0)
|
||
|
||
return injection_tensor
|
||
|
||
|
||
# Global patcher instance
|
||
_global_patcher = ActivationPatcher()
|
||
|
||
def get_patcher() -> ActivationPatcher:
|
||
"""Get the global activation patcher instance."""
|
||
return _global_patcher |