From 505daa562c49516d54e8055fb2e69987ccbea863 Mon Sep 17 00:00:00 2001 From: Junst Date: Tue, 23 Sep 2025 19:39:49 +0900 Subject: [PATCH] Focus on actual attention outputs instead of Linear layers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 🎯 **Real Attention Outputs Focus**: 1. **Removed Linear Layer Capture**: - Removed 'linear' from hook keywords - Focus on actual attention modules (SelfAttention, etc.) - Exclude Linear layers from hook registration 2. **Target Actual Attention Modules**: - Hook registration now excludes 'linear' layers - Focus on 'attention' and 'attn' modules only - Should capture real attention outputs like original ConceptAttention 3. **Improved Hook Triggering**: - Try to trigger actual SelfAttention modules - Use proper inputs (x, pe) for attention modules - Should capture img_attn and txt_attn outputs 4. **Align with Original ConceptAttention**: - Original uses img_attn, concept_attn outputs - Not Linear layer outputs - Should get proper attention vectors for concept-image attention This should capture the actual attention outputs that the original ConceptAttention paper uses for concept attention map generation. --- concept_attention.py | 63 ++++++++++++++++++-------------------------- 1 file changed, 25 insertions(+), 38 deletions(-) diff --git a/concept_attention.py b/concept_attention.py index 7b7ea2e..5f58ed2 100644 --- a/concept_attention.py +++ b/concept_attention.py @@ -39,9 +39,10 @@ class ConceptAttention: module_name = getattr(module, '__class__', type(module)).__name__ # 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 [ 'selfattention', 'crossattention', 'doubleblock', - 'img_attn', 'txt_attn', 'attention', 'attn', 'linear' + 'img_attn', 'txt_attn', 'attention', 'attn' ]) or any(keyword in str(module).lower() for keyword in [ 'selfattention', 'crossattention', 'doubleblock', 'img_attn', 'txt_attn', 'attention', 'attn' @@ -58,7 +59,8 @@ class ConceptAttention: try: hook_count = 0 for name, module in actual_model.named_modules(): - if any(keyword in name.lower() for keyword in ['attention', 'attn', 'query', 'key', 'value', 'proj']): + # Focus on actual attention modules, not Linear layers + if any(keyword in name.lower() for keyword in ['attention', 'attn']) and 'linear' not in name.lower(): hook = module.register_forward_hook(hook_fn) self.attention_hooks.append(hook) hook_count += 1 @@ -477,48 +479,33 @@ class ConceptAttentionProcessor: img_attn = block.img_attn logger.info(f"Found img_attn in block {block_idx}: {type(img_attn)}") - # Try to access qkv layer - if hasattr(img_attn, 'qkv'): - qkv = img_attn.qkv - logger.info(f"Found qkv in block {block_idx}: {type(qkv)}") + # Try to trigger the actual attention module (not just qkv/proj) + try: + # Create appropriate inputs for the attention module + # SelfAttention expects: x, pe (positional encoding) + test_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype) + test_pe = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype) - # Try to trigger the qkv layer - try: - test_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype) - with torch.no_grad(): - qkv_output = qkv(test_input) - logger.info(f"🚀 Successfully triggered qkv in block {block_idx}, output shape: {qkv_output.shape}") - except Exception as qkv_error: - logger.debug(f"QKV forward failed in block {block_idx}: {qkv_error}") - - # Try to access proj layer - if hasattr(img_attn, 'proj'): - proj = img_attn.proj - logger.info(f"Found proj in block {block_idx}: {type(proj)}") - - # Try to trigger the proj layer - try: - test_input = torch.randn(1, 1024, 3072, device=self.device, dtype=model_dtype) - with torch.no_grad(): - proj_output = proj(test_input) - logger.info(f"🚀 Successfully triggered proj in block {block_idx}, output shape: {proj_output.shape}") - except Exception as proj_error: - logger.debug(f"Proj forward failed in block {block_idx}: {proj_error}") + with torch.no_grad(): + attn_output = img_attn(test_input, test_pe) + logger.info(f"🚀 Successfully triggered img_attn in block {block_idx}, output shape: {attn_output.shape}") + except Exception as attn_error: + logger.debug(f"Img attention forward failed in block {block_idx}: {attn_error}") if hasattr(block, 'txt_attn'): txt_attn = block.txt_attn logger.info(f"Found txt_attn in block {block_idx}: {type(txt_attn)}") - # Similar approach for txt_attn - if hasattr(txt_attn, 'qkv'): - qkv = txt_attn.qkv - try: - test_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype) - with torch.no_grad(): - qkv_output = qkv(test_input) - logger.info(f"🚀 Successfully triggered txt_attn qkv in block {block_idx}, output shape: {qkv_output.shape}") - except Exception as txt_qkv_error: - logger.debug(f"Txt QKV forward failed in block {block_idx}: {txt_qkv_error}") + # Try to trigger the actual attention module + try: + test_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype) + test_pe = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype) + + with torch.no_grad(): + attn_output = txt_attn(test_input, test_pe) + logger.info(f"🚀 Successfully triggered txt_attn in block {block_idx}, output shape: {attn_output.shape}") + except Exception as txt_attn_error: + logger.debug(f"Txt attention forward failed in block {block_idx}: {txt_attn_error}") except Exception as block_error: logger.debug(f"Block {block_idx} access failed: {block_error}")