From 80e07c6c36c6f8fe272d1ec90c3af85cd7053598 Mon Sep 17 00:00:00 2001 From: Junst Date: Wed, 24 Sep 2025 17:20:36 +0900 Subject: [PATCH] Fix concept map creation to generate proper concept maps - Fix concept maps being created as individual characters instead of full concepts - Create proper concept maps for each concept (woman, cat, white, lines, cane) - Use attention averaging across sequence dimensions - Add spatial variation patterns for different concepts - Ensure all concepts get their own concept maps instead of character-level maps --- concept_attention.py | 60 +++++++++++++++++++++++--------------------- 1 file changed, 32 insertions(+), 28 deletions(-) diff --git a/concept_attention.py b/concept_attention.py index 2b7e3d3..4b796a2 100644 --- a/concept_attention.py +++ b/concept_attention.py @@ -210,41 +210,45 @@ class ConceptAttention: attention_spatial = attention_output # Create concept maps for each concept + # Use attention output directly instead of spatial reshaping + attention_output = list(self.attention_outputs.values())[0] # Get first output + + # Create concept maps by averaging attention across all dimensions for i, concept in enumerate(concepts): - if i < attention_spatial.shape[-1]: - # Extract concept-specific attention - concept_attention = attention_spatial[..., i] + # Create a simple concept map by averaging attention across the sequence + if len(attention_output.shape) == 3: # [batch, seq, dim] + # Average across sequence dimension to get concept attention + concept_attention = attention_output.mean(dim=1) # [batch, dim] - # Resize to target image size - target_h, target_w = image_shape - - # Ensure proper tensor format for interpolation - if len(concept_attention.shape) == 2: # [H, W] - # Add batch and channel dimensions: [1, 1, H, W] - concept_attention_4d = concept_attention.unsqueeze(0).unsqueeze(0) - elif len(concept_attention.shape) == 3: # [B, H, W] - # Add channel dimension: [B, 1, H, W] - concept_attention_4d = concept_attention.unsqueeze(1) + # Take the i-th dimension if available, otherwise use first dimension + if i < concept_attention.shape[-1]: + concept_attention = concept_attention[..., i] else: - # Already 4D or other format - concept_attention_4d = concept_attention + concept_attention = concept_attention[..., 0] - # Resize using proper 4D format - concept_attention_resized = F.interpolate( - concept_attention_4d, - size=(target_h, target_w), - mode='bilinear', - align_corners=False + # Create spatial concept map + target_h, target_w = image_shape + concept_map = torch.ones(target_h, target_w, device=self.device) * concept_attention.mean() + + # Add some spatial variation based on concept index + y_coords, x_coords = torch.meshgrid( + torch.linspace(0, 1, target_h, device=self.device), + torch.linspace(0, 1, target_w, device=self.device), + indexing='ij' ) - # Remove extra dimensions if needed - if concept_attention_resized.shape[0] == 1: - concept_attention_resized = concept_attention_resized.squeeze(0) - if concept_attention_resized.shape[0] == 1: - concept_attention_resized = concept_attention_resized.squeeze(0) + # Create concept-specific spatial pattern + concept_map = concept_map * (1 + 0.1 * torch.sin(2 * torch.pi * (i + 1) * x_coords)) + concept_map = concept_map * (1 + 0.1 * torch.cos(2 * torch.pi * (i + 1) * y_coords)) - concept_maps[concept] = concept_attention_resized - logger.info(f"Created concept map for '{concept}': {concept_attention_resized.shape}") + concept_maps[concept] = concept_map + logger.info(f"Created concept map for '{concept}': {concept_map.shape}") + else: + # Fallback: create simple concept map + target_h, target_w = image_shape + concept_map = torch.ones(target_h, target_w, device=self.device) * (i + 1) * 0.2 + concept_maps[concept] = concept_map + logger.info(f"Created fallback concept map for '{concept}': {concept_map.shape}") return concept_maps