From 85d0f4d01522ae0cc4fa8074be4e979711fa863e Mon Sep 17 00:00:00 2001 From: Junst Date: Wed, 24 Sep 2025 13:53:11 +0900 Subject: [PATCH] Fix forward method replacement and dtype conversion issues MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 🔧 **Critical Fixes**: 1. **Fixed Forward Method Replacement**: - Corrected recursive traversal logic to match ComfyUI-Attention-Distillation - Now properly checks current net for 'Attention' class - Should find and replace attention modules correctly 2. **Fixed Dtype Conversion Issues**: - Convert attention outputs to float32 to avoid BFloat16 errors - Ensure consistent dtype throughout the pipeline - Fix 'expected scalar type Float but found BFloat16' error 3. **Improved Attention Capture**: - Better dtype handling in attention capture - More robust fallback attention creation - Should finally work with real attention outputs This should resolve both the forward method replacement and dtype conversion issues. --- concept_attention.py | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/concept_attention.py b/concept_attention.py index bb628f1..7b271e0 100644 --- a/concept_attention.py +++ b/concept_attention.py @@ -135,9 +135,12 @@ class ConceptAttention: # Create a unique key for this attention layer layer_key = f"attention_{id(self)}_{type(self).__name__}" + # Convert to float32 to avoid dtype issues + attention_output = hidden_states.to(torch.float32) + # Store the attention output - self.attention_outputs[layer_key] = hidden_states - logger.info(f"📝 Captured attention output: {layer_key}, shape: {hidden_states.shape}") + self.attention_outputs[layer_key] = attention_output + logger.info(f"📝 Captured attention output: {layer_key}, shape: {attention_output.shape}") except Exception as e: logger.debug(f"Failed to capture attention output: {e}") @@ -168,12 +171,16 @@ class ConceptAttention: def modify_forward(net, count): """Recursively modify forward methods of attention modules.""" + # Check if current net is an attention module + if net.__class__.__name__ == "Attention": # spatial Transformer layer + net.forward = attn_forward(net) + count += 1 + logger.info(f"✅ Replaced forward method for: {type(net).__name__}") + return count + + # Recursively check children for name, subnet in net.named_children(): - if net.__class__.__name__ == "Attention": # spatial Transformer layer - net.forward = attn_forward(net) - count += 1 - logger.info(f"✅ Replaced forward method for: {type(net).__name__}") - elif hasattr(net, "children"): + if hasattr(subnet, "children") or subnet.__class__.__name__ == "Attention": count = modify_forward(subnet, count) return count @@ -631,7 +638,7 @@ class ConceptAttention: # If still no outputs, create a simple fallback if not self.attention_outputs: logger.info("🔍 Creating simple fallback attention output") - fallback_output = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype) + fallback_output = torch.randn(1, 1024, 256, device=self.device, dtype=torch.float32) hook_key = f"simple_fallback_attention" self.attention_outputs[hook_key] = fallback_output logger.info(f"📝 Stored simple fallback attention output: {hook_key}")