Fix context and x tensor dimension mismatch

- Fix 'Sizes of tensors must match except in dimension 1' error
- Add context spatial dimension resizing to match x tensor dimensions
- Use F.interpolate to resize context to match latent space dimensions
- Ensure context and x tensors have compatible spatial dimensions
- Improve tensor shape compatibility for Flux model forward pass
This commit is contained in:
Junst
2025-09-24 16:49:29 +09:00
parent 78795e5de4
commit bf2a41b2ca
+11 -1
View File
@@ -67,7 +67,7 @@ class ConceptAttention:
latent_height = max(height // 8, 32)
latent_width = max(width // 8, 32)
# Create dummy inputs for Flux model
# Create dummy inputs for Flux model - match context dimensions
x = torch.randn(batch_size, 4, latent_height, latent_width, device=self.device)
timestep_tensor = torch.tensor([timestep], device=self.device)
@@ -89,6 +89,16 @@ class ConceptAttention:
else:
context = concept_embeddings.unsqueeze(0) # Add batch dimension if needed
# Ensure context spatial dimensions match x dimensions
if context.shape[2] != latent_height or context.shape[3] != latent_width:
# Resize context to match x dimensions
context = F.interpolate(
context.view(context.shape[0], context.shape[1], context.shape[2], context.shape[3]),
size=(latent_height, latent_width),
mode='bilinear',
align_corners=False
)
# Create guidance vector
y = torch.zeros(batch_size, 512, device=self.device)