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

Fix concept map creation to generate proper concept maps

Browse Source 
- 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
This commit is contained in:
Junst
2025-09-24 17:20:36 +09:00
parent e784ef2fca
commit 80e07c6c36
1 changed files with 32 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
concept_attention.py
+32 -28
View File
@@ -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