nodes/Junst-ComfyUI-Concept-Diffusion
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add direct attention output capture and enhanced hook debugging

Browse Source 
🔧 **Direct Attention Capture & Enhanced Debugging**:

1. **Enhanced Hook Debugging**:
   - Added detailed logging to hook function to see if hooks are called
   - Track which modules trigger hooks and which don't
   - Monitor total attention outputs captured

2. **Direct Attention Output Capture**:
   - Store attention outputs directly without relying on hooks
   - Capture qkv, attention, and projection outputs separately
   - Use unique keys for each captured output

3. **Comprehensive Logging**:
   - Log total attention outputs captured
   - Show attention output keys for debugging
   - Track each step of attention computation

This approach bypasses hook issues by directly capturing
attention outputs during manual attention computation,
similar to how original ConceptAttention works.
This commit is contained in:
Junst
2025-09-23 20:36:24 +09:00
parent 1f8dc90cf8
commit 2ea5df78d5
1 changed files with 28 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
concept_attention.py
+28 -1
View File
@@ -35,9 +35,11 @@ class ConceptAttention:
self.attention_outputs = {}
def hook_fn(module, input, output):
# Store attention outputs for concept extraction
# Add debug logging to see if hook is being called at all
module_name = getattr(module, '__class__', type(module)).__name__
logger.info(f"🔍 HOOK CALLED on {module_name} - checking if attention-related")
# Store attention outputs for concept extraction
# Capture outputs from Flux DiT specific attention modules
# Focus on actual attention outputs, not Linear layers
if any(keyword in module_name.lower() for keyword in [
@@ -51,6 +53,9 @@ class ConceptAttention:
hook_key = f"{module_name}_{id(module)}"
self.attention_outputs[hook_key] = output
logger.info(f"🎯 HOOK TRIGGERED! Captured output from {module_name}: {output.shape}")
logger.info(f"Total attention outputs captured: {len(self.attention_outputs)}")
else:
logger.info(f"Hook called on {module_name} but not capturing (not attention-related)")
# Get the actual model from ModelPatcher
actual_model = getattr(self.model, 'model', self.model)
@@ -534,6 +539,11 @@ class ConceptAttentionProcessor:
qkv_output = img_attn.qkv(test_input)
logger.info(f"🚀 Successfully ran img_attn.qkv in block {block_idx}, output shape: {qkv_output.shape}")
# Store the qkv output directly as attention output
hook_key = f"img_attn_qkv_block_{block_idx}_{id(img_attn)}"
self.attention_outputs[hook_key] = qkv_output
logger.info(f"📝 Stored qkv output directly: {hook_key}")
# Try to run the full attention computation
if hasattr(img_attn, 'norm') and hasattr(img_attn, 'proj'):
# Split qkv output
@@ -547,9 +557,19 @@ class ConceptAttentionProcessor:
attn_output = torch.nn.functional.scaled_dot_product_attention(q, k, v)
logger.info(f"🚀 Successfully computed attention in block {block_idx}, output shape: {attn_output.shape}")
# Store the attention output directly
hook_key = f"img_attn_attention_block_{block_idx}_{id(img_attn)}"
self.attention_outputs[hook_key] = attn_output
logger.info(f"📝 Stored attention output directly: {hook_key}")
# Apply projection
proj_output = img_attn.proj(attn_output)
logger.info(f"🚀 Successfully ran img_attn.proj in block {block_idx}, output shape: {proj_output.shape}")
# Store the projection output directly
hook_key = f"img_attn_proj_block_{block_idx}_{id(img_attn)}"
self.attention_outputs[hook_key] = proj_output
logger.info(f"📝 Stored projection output directly: {hook_key}")
except Exception as attn_error:
logger.debug(f"Direct attention execution failed in block {block_idx}: {attn_error}")
@@ -557,6 +577,13 @@ class ConceptAttentionProcessor:
except Exception as block_error:
logger.debug(f"Block {block_idx} access failed: {block_error}")
# Log total attention outputs captured
logger.info(f"📊 Total attention outputs captured: {len(self.attention_outputs)}")
if self.attention_outputs:
logger.info(f"📊 Attention output keys: {list(self.attention_outputs.keys())}")
else:
logger.warning("⚠️ No attention outputs captured!")
except Exception as e:
logger.warning(f"DoubleStreamBlock execution failed: {e}")
logger.info("Hooks may not have been triggered")