From deadb1caebeab549f4582c13dbbe6675ede11753 Mon Sep 17 00:00:00 2001 From: Junst Date: Wed, 17 Sep 2025 14:59:34 +0900 Subject: [PATCH] Enable actual model forward pass to capture real attention 1. Remove forward pass skip: - Hook registration was successful, so we can try to capture real attention - Attempt actual model forward pass with proper error handling - Log detailed information about model execution 2. Improve hook function: - Capture outputs from more attention-related modules (query, key, value, proj) - Add detailed logging for captured outputs - Better debugging information for attention capture 3. Enhanced error handling: - Try multiple model call patterns (apply_model, forward, direct call) - Graceful fallback to mock data if real extraction fails - Clear logging of what's happening at each step This should now attempt to capture real attention from the DiT model while maintaining fallback to mock data for robustness. --- concept_attention.py | 42 +++++++++++++++++++++++++++++++++++------- 1 file changed, 35 insertions(+), 7 deletions(-) diff --git a/concept_attention.py b/concept_attention.py index b0a031f..8f6768c 100644 --- a/concept_attention.py +++ b/concept_attention.py @@ -33,9 +33,11 @@ class ConceptAttention: def hook_fn(module, input, output): # Store attention outputs for concept extraction module_name = getattr(module, '__class__', type(module)).__name__ - if 'attention' in module_name.lower() or 'attn' in module_name.lower(): + # Capture outputs from attention-related modules + if any(keyword in module_name.lower() for keyword in ['attention', 'attn', 'query', 'key', 'value', 'proj']): self.attention_outputs[module_name] = output - logger.info(f"Captured attention output from {module_name}, shape: {output.shape}") + logger.info(f"Captured output from {module_name}, shape: {output.shape}") + logger.info(f"Output type: {type(output)}, device: {output.device if hasattr(output, 'device') else 'N/A'}") # Get the actual model from ModelPatcher actual_model = getattr(self.model, 'model', self.model) @@ -178,11 +180,37 @@ class ConceptAttentionProcessor: # Run model forward pass to capture attention with torch.no_grad(): - # For now, skip the actual model forward pass since ModelPatcher is complex - # and focus on getting the attention extraction working with mock data - logger.info("Skipping model forward pass for now, using mock data") - logger.info("This is expected behavior - real attention extraction will be implemented") - logger.info("when we have proper access to the DiT model's attention layers") + logger.info("Attempting to run model forward pass to capture attention") + + # Prepare input for the diffusion model + batch_size = image.shape[0] + height, width = image.shape[1], image.shape[2] + + # Create dummy timestep and noise for DiT + timestep = torch.tensor([0], device=self.device) + noise = torch.randn_like(image) + + try: + # Get the actual model from ModelPatcher + actual_model = getattr(self.model, 'model', self.model) + logger.info(f"Using actual model: {type(actual_model)}") + + # Try to run the model to trigger hooks + if hasattr(actual_model, 'apply_model'): + logger.info("Using apply_model method") + _ = actual_model.apply_model(noise, timestep) + elif hasattr(actual_model, 'forward'): + logger.info("Using forward method") + _ = actual_model(noise, timestep) + else: + logger.info("Trying direct model call") + _ = actual_model(noise) + + logger.info("Model forward pass completed successfully") + + except Exception as e: + logger.warning(f"Model forward pass failed: {e}") + logger.info("Will use mock data as fallback") # Process attention outputs to create concept maps concept_maps = self._create_concept_maps_from_attention(concept_embeddings, image.shape)