Fix dtype mismatch and input shape issues for Flux model
1. Fix dtype mismatch between Float and BFloat16: - Detect model's actual dtype (BFloat16) and use it for all inputs - Ensure all tensors use the same dtype as the model - Fix 'mat1 and mat2 must have the same dtype' error 2. Fix input shape issues for different layer types: - Use correct shapes for query_norm and key_norm layers ([*, 128]) - Use appropriate shapes for QKV layers ([1, 1024, 256]) - Fix 'expected input with shape [*, 128], but got input of size[1, 1024, 256]' error 3. Improve model device and dtype handling: - Check model's device and dtype before creating inputs - Move model to correct device while preserving dtype - Use model's dtype for all input tensors 4. Add final hook triggering attempt: - Access model parameters directly to trigger hooks - Limit parameter access to avoid excessive computation - Multiple fallback strategies for hook activation This should resolve the dtype mismatch and shape errors while providing better hook triggering mechanisms for attention capture.
This commit is contained in:
Junst
+41
-16
@@ -234,23 +234,29 @@ class ConceptAttentionProcessor:
|
||||
if actual_model is not None:
|
||||
logger.info(f"Using actual model: {type(actual_model)}")
|
||||
|
||||
# Check model device and move to correct device if needed
|
||||
# Check model device and dtype, move to correct device if needed
|
||||
model_device = next(actual_model.parameters()).device
|
||||
model_dtype = next(actual_model.parameters()).dtype
|
||||
logger.info(f"Model device: {model_device}, Target device: {self.device}")
|
||||
logger.info(f"Model dtype: {model_dtype}")
|
||||
|
||||
if model_device != self.device:
|
||||
logger.info(f"Moving model from {model_device} to {self.device}")
|
||||
actual_model = actual_model.to(self.device)
|
||||
|
||||
# Update model dtype for consistency
|
||||
if model_dtype != torch.float32:
|
||||
logger.info(f"Model uses {model_dtype}, will use this dtype for inputs")
|
||||
|
||||
# For Flux models, we need the correct input format
|
||||
logger.info("Attempting Flux model forward pass with correct input format")
|
||||
|
||||
try:
|
||||
# Prepare inputs according to Flux model requirements - ensure all on same device as model
|
||||
x = torch.randn(1, 4, 64, 64, device=self.device, dtype=torch.float32) # Flux expects 4-channel input
|
||||
timestep = torch.tensor([0.0], device=self.device, dtype=torch.float32)
|
||||
context = torch.randn(1, 77, 2048, device=self.device, dtype=torch.float32) # CLIP context
|
||||
y = torch.randn(1, 512, device=self.device, dtype=torch.float32) # CLIP pooled
|
||||
# Prepare inputs according to Flux model requirements - use model's dtype
|
||||
x = torch.randn(1, 4, 64, 64, device=self.device, dtype=model_dtype) # Flux expects 4-channel input
|
||||
timestep = torch.tensor([0.0], device=self.device, dtype=model_dtype)
|
||||
context = torch.randn(1, 77, 2048, device=self.device, dtype=model_dtype) # CLIP context
|
||||
y = torch.randn(1, 512, device=self.device, dtype=model_dtype) # CLIP pooled
|
||||
|
||||
logger.info(f"Flux inputs - x: {x.shape}, timestep: {timestep.shape}, context: {context.shape}, y: {y.shape}")
|
||||
|
||||
@@ -289,12 +295,12 @@ class ConceptAttentionProcessor:
|
||||
first_block = blocks[0]
|
||||
logger.info(f"First block type: {type(first_block)}")
|
||||
|
||||
# Create minimal inputs for the block with correct parameters
|
||||
# Create minimal inputs for the block with correct parameters and dtype
|
||||
# DoubleStreamBlock needs: x, t, vec, pe
|
||||
block_x = torch.randn(1, 256, 64, 64, device=self.device, dtype=torch.float32)
|
||||
block_t = torch.randn(1, 256, device=self.device, dtype=torch.float32)
|
||||
block_vec = torch.randn(1, 77, 2048, device=self.device, dtype=torch.float32) # CLIP context
|
||||
block_pe = torch.randn(1, 256, device=self.device, dtype=torch.float32) # Positional encoding
|
||||
block_x = torch.randn(1, 256, 64, 64, device=self.device, dtype=model_dtype)
|
||||
block_t = torch.randn(1, 256, device=self.device, dtype=model_dtype)
|
||||
block_vec = torch.randn(1, 77, 2048, device=self.device, dtype=model_dtype) # CLIP context
|
||||
block_pe = torch.randn(1, 256, device=self.device, dtype=model_dtype) # Positional encoding
|
||||
|
||||
# Try to run the block
|
||||
try:
|
||||
@@ -330,16 +336,20 @@ class ConceptAttentionProcessor:
|
||||
|
||||
# Try to run a forward pass on the layer if possible
|
||||
if hasattr(layer, 'forward'):
|
||||
# Create appropriate input based on layer type
|
||||
# Create appropriate input based on layer type and expected shape
|
||||
if 'qkv' in name.lower():
|
||||
# QKV layer expects different input
|
||||
layer_input = torch.randn(1, 1024, 256, device=self.device, dtype=torch.float32)
|
||||
layer_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
||||
elif 'norm' in name.lower():
|
||||
# Norm layer
|
||||
layer_input = torch.randn(1, 1024, 256, device=self.device, dtype=torch.float32)
|
||||
# Norm layer - check if it's query_norm, key_norm, etc.
|
||||
if 'query_norm' in name.lower() or 'key_norm' in name.lower():
|
||||
# These expect shape [*, 128] based on error message
|
||||
layer_input = torch.randn(1, 128, device=self.device, dtype=model_dtype)
|
||||
else:
|
||||
layer_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
||||
else:
|
||||
# Default attention layer
|
||||
layer_input = torch.randn(1, 1024, 256, device=self.device, dtype=torch.float32)
|
||||
layer_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
||||
|
||||
try:
|
||||
layer_output = layer(layer_input)
|
||||
@@ -361,6 +371,21 @@ class ConceptAttentionProcessor:
|
||||
except Exception as alt_error:
|
||||
logger.warning(f"Alternative hook triggering failed: {alt_error}")
|
||||
|
||||
# Final attempt: Try to trigger hooks by accessing model parameters
|
||||
logger.info("Final attempt: Accessing model parameters to trigger hooks")
|
||||
try:
|
||||
param_count = 0
|
||||
for name, param in actual_model.named_parameters():
|
||||
if 'attn' in name.lower() or 'attention' in name.lower():
|
||||
# Access the parameter to potentially trigger computation
|
||||
_ = param.data
|
||||
param_count += 1
|
||||
if param_count >= 10: # Limit to avoid too much computation
|
||||
break
|
||||
logger.info(f"Accessed {param_count} attention-related parameters")
|
||||
except Exception as param_error:
|
||||
logger.warning(f"Parameter access failed: {param_error}")
|
||||
|
||||
elif hasattr(actual_model, 'forward'):
|
||||
result = actual_model.forward(x, timestep, context, y)
|
||||
logger.info(f"Flux forward result: {type(result)}, shape: {result.shape if result is not None else 'None'}")
|
||||
|
||||
Reference in New Issue
Block a user