Add full block execution to trigger hooks
🔧 **Full Block Execution**:
1. **Individual Attention Module Execution**:
- Try to run img_attn and txt_attn modules directly
- Use proper input shapes for SelfAttention
- Handle different calling patterns
2. **Full Block Execution**:
- Run entire DoubleStreamBlock to trigger hooks
- Create proper inputs for block execution
- Should trigger all hooks within the block
3. **Comprehensive Hook Triggering**:
- Try both individual attention and full block execution
- Multiple approaches to ensure hooks are triggered
- Better logging for debugging
This should finally trigger the registered hooks and capture
actual attention outputs for concept attention map generation.
This commit is contained in:
@@ -497,6 +497,21 @@ class ConceptAttentionProcessor:
|
||||
|
||||
except Exception as block_error:
|
||||
logger.debug(f"Block {block_idx} access failed: {block_error}")
|
||||
|
||||
# Try to run the entire block to trigger hooks
|
||||
try:
|
||||
logger.info(f"Attempting to run entire block {block_idx}")
|
||||
# Create inputs for the entire block
|
||||
block_img = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
||||
block_txt = torch.randn(1, 77, 256, device=self.device, dtype=model_dtype)
|
||||
block_vec = torch.randn(1, 256, device=self.device, dtype=model_dtype)
|
||||
block_pe = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
||||
|
||||
with torch.no_grad():
|
||||
block_output = block(block_img, block_txt, block_vec, block_pe)
|
||||
logger.info(f"🚀 Successfully ran entire block {block_idx}, output shape: {block_output.shape if hasattr(block_output, 'shape') else 'No shape'}")
|
||||
except Exception as full_block_error:
|
||||
logger.debug(f"Full block {block_idx} execution failed: {full_block_error}")
|
||||
|
||||
# Also try to trigger hooks by accessing model parameters
|
||||
logger.info("🔍 Attempting to trigger hooks by accessing model parameters")
|
||||
|
||||
Reference in New Issue
Block a user