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

Fix spatial dimension mismatch in concept map creation

Browse Source 
- Fix 'Input and output must have the same number of spatial dimensions' error
- Add proper tensor format handling for F.interpolate
- Support 2D, 3D, and 4D tensor formats for interpolation
- Ensure correct dimension handling for concept attention tensors
- Add dimension squeezing to remove extra batch/channel dimensions
This commit is contained in:
Junst
2025-09-24 17:15:03 +09:00
parent 3499ca8a0a
commit e784ef2fca
1 changed files with 21 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
concept_attention.py
+21 -2
View File
@@ -217,12 +217,31 @@ class ConceptAttention:
# 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)
else:
# Already 4D or other format
concept_attention_4d = concept_attention
# Resize using proper 4D format
concept_attention_resized = F.interpolate(
concept_attention.unsqueeze(0).unsqueeze(0),
concept_attention_4d,
size=(target_h, target_w),
mode='bilinear',
align_corners=False
).squeeze()
)
# 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)
concept_maps[concept] = concept_attention_resized
logger.info(f"Created concept map for '{concept}': {concept_attention_resized.shape}")