Implement direct attention module execution to trigger hooks
🔧 **Direct Attention Module Execution**:
1. **Individual Attention Module Access**:
- Directly access img_attn and txt_attn within DoubleStreamBlocks
- Target blocks 15-19 (like original ConceptAttention)
- Create appropriate inputs for SelfAttention modules
2. **Multiple Execution Patterns**:
- Try different calling patterns for SelfAttention (with/without pe)
- Use proper input shapes for image and text attention
- Handle both img_attn and txt_attn modules
3. **Hook Triggering**:
- Execute individual attention modules to trigger registered hooks
- Should capture actual attention outputs from SelfAttention
- More reliable than full model forward pass
This approach directly runs the attention modules that have hooks
registered, which should finally trigger the hooks and capture
real attention outputs for concept attention map generation.
This commit is contained in:
Junst
+84
-58
@@ -426,65 +426,91 @@ class ConceptAttentionProcessor:
|
||||
actual_model = actual_model.to(self.device)
|
||||
|
||||
try:
|
||||
# Try to run the actual Flux model forward pass to trigger hooks
|
||||
logger.info("🔍 Attempting to run actual Flux model forward pass")
|
||||
|
||||
# Create proper inputs for Flux model
|
||||
# Based on ComfyUI's Flux model requirements
|
||||
x = torch.randn(1, 4, 64, 64, device=self.device, dtype=model_dtype) # Latent 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}")
|
||||
|
||||
# Try to run the model using apply_model
|
||||
if hasattr(actual_model, 'apply_model'):
|
||||
logger.info("Using apply_model method for Flux")
|
||||
try:
|
||||
# Try with minimal parameters first
|
||||
result = actual_model.apply_model(x, timestep)
|
||||
logger.info(f"Flux apply_model result: {type(result)}, shape: {result.shape if result is not None else 'None'}")
|
||||
except Exception as e1:
|
||||
logger.warning(f"apply_model failed: {e1}")
|
||||
try:
|
||||
# Try with context
|
||||
result = actual_model.apply_model(x, timestep, c_crossattn=context)
|
||||
logger.info(f"Flux apply_model with context result: {type(result)}, shape: {result.shape if result is not None else 'None'}")
|
||||
except Exception as e2:
|
||||
logger.warning(f"apply_model with context failed: {e2}")
|
||||
|
||||
# Also try to run a single block to trigger hooks
|
||||
if hasattr(actual_model, 'diffusion_model'):
|
||||
diffusion_model = actual_model.diffusion_model
|
||||
if hasattr(diffusion_model, 'double_blocks'):
|
||||
double_blocks = diffusion_model.double_blocks
|
||||
logger.info(f"Found {len(double_blocks)} double blocks")
|
||||
|
||||
# Try to run the last few blocks (15-19 like original ConceptAttention)
|
||||
target_blocks = [15, 16, 17, 18, 19] if len(double_blocks) > 19 else list(range(max(0, len(double_blocks)-5), len(double_blocks)))
|
||||
|
||||
for block_idx in target_blocks:
|
||||
if block_idx < len(double_blocks):
|
||||
block = double_blocks[block_idx]
|
||||
logger.info(f"Attempting to run block {block_idx}: {type(block)}")
|
||||
|
||||
try:
|
||||
# Create inputs for the block
|
||||
# DoubleStreamBlock needs: img, txt, vec, pe
|
||||
block_img = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
||||
block_txt = torch.randn(1, 77, 256, device=self.device, dtype=model_dtype)
|
||||
block_vec = torch.randn(1, 256, device=self.device, dtype=model_dtype)
|
||||
block_pe = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
||||
|
||||
with torch.no_grad():
|
||||
block_output = block(block_img, block_txt, block_vec, block_pe)
|
||||
logger.info(f"🚀 Successfully ran block {block_idx}, output shape: {block_output.shape if hasattr(block_output, 'shape') else 'No shape'}")
|
||||
except Exception as block_error:
|
||||
logger.debug(f"Block {block_idx} execution failed: {block_error}")
|
||||
|
||||
# Try a different approach - directly access and run attention modules
|
||||
logger.info("🔍 Attempting direct attention module execution")
|
||||
|
||||
# Find the diffusion model
|
||||
diffusion_model = None
|
||||
if hasattr(actual_model, 'diffusion_model'):
|
||||
diffusion_model = actual_model.diffusion_model
|
||||
elif hasattr(actual_model, 'model') and hasattr(actual_model.model, 'diffusion_model'):
|
||||
diffusion_model = actual_model.model.diffusion_model
|
||||
|
||||
if diffusion_model is not None:
|
||||
logger.info(f"Found diffusion model: {type(diffusion_model)}")
|
||||
|
||||
# Try to access and run individual attention modules
|
||||
if hasattr(diffusion_model, 'double_blocks'):
|
||||
double_blocks = diffusion_model.double_blocks
|
||||
logger.info(f"Found {len(double_blocks)} double blocks")
|
||||
|
||||
# Target the last few blocks (15-19 like original ConceptAttention)
|
||||
target_blocks = [15, 16, 17, 18, 19] if len(double_blocks) > 19 else list(range(max(0, len(double_blocks)-5), len(double_blocks)))
|
||||
|
||||
for block_idx in target_blocks:
|
||||
if block_idx < len(double_blocks):
|
||||
block = double_blocks[block_idx]
|
||||
logger.info(f"Attempting to run block {block_idx}: {type(block)}")
|
||||
|
||||
try:
|
||||
# Try to run individual attention modules within the block
|
||||
if hasattr(block, 'img_attn'):
|
||||
img_attn = block.img_attn
|
||||
logger.info(f"Found img_attn in block {block_idx}: {type(img_attn)}")
|
||||
|
||||
# Create simple inputs for SelfAttention
|
||||
# SelfAttention typically takes (x, pe) where x is [batch, seq_len, dim]
|
||||
attn_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
||||
attn_pe = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
||||
|
||||
with torch.no_grad():
|
||||
try:
|
||||
# Try different calling patterns for SelfAttention
|
||||
attn_output = img_attn(attn_input, attn_pe)
|
||||
logger.info(f"🚀 Successfully ran img_attn in block {block_idx}, output shape: {attn_output.shape}")
|
||||
except Exception as e1:
|
||||
try:
|
||||
# Try with just input
|
||||
attn_output = img_attn(attn_input)
|
||||
logger.info(f"🚀 Successfully ran img_attn (input only) in block {block_idx}, output shape: {attn_output.shape}")
|
||||
except Exception as e2:
|
||||
logger.debug(f"img_attn execution failed in block {block_idx}: {e1}, {e2}")
|
||||
|
||||
if hasattr(block, 'txt_attn'):
|
||||
txt_attn = block.txt_attn
|
||||
logger.info(f"Found txt_attn in block {block_idx}: {type(txt_attn)}")
|
||||
|
||||
# Create inputs for text attention
|
||||
txt_input = torch.randn(1, 77, 256, device=self.device, dtype=model_dtype)
|
||||
txt_pe = torch.randn(1, 77, 256, device=self.device, dtype=model_dtype)
|
||||
|
||||
with torch.no_grad():
|
||||
try:
|
||||
txt_output = txt_attn(txt_input, txt_pe)
|
||||
logger.info(f"🚀 Successfully ran txt_attn in block {block_idx}, output shape: {txt_output.shape}")
|
||||
except Exception as e1:
|
||||
try:
|
||||
txt_output = txt_attn(txt_input)
|
||||
logger.info(f"🚀 Successfully ran txt_attn (input only) in block {block_idx}, output shape: {txt_output.shape}")
|
||||
except Exception as e2:
|
||||
logger.debug(f"txt_attn execution failed in block {block_idx}: {e1}, {e2}")
|
||||
|
||||
except Exception as block_error:
|
||||
logger.debug(f"Block {block_idx} access failed: {block_error}")
|
||||
|
||||
# Also try to trigger hooks by accessing model parameters
|
||||
logger.info("🔍 Attempting to trigger hooks by accessing model parameters")
|
||||
param_count = 0
|
||||
for name, param in actual_model.named_parameters():
|
||||
if any(keyword in name.lower() for keyword in ['attention', 'attn', 'qkv', 'proj']):
|
||||
_ = 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 e:
|
||||
logger.warning(f"Flux model forward pass failed: {e}")
|
||||
logger.warning(f"Direct attention execution failed: {e}")
|
||||
logger.info("Hooks may not have been triggered")
|
||||
else:
|
||||
logger.warning("No suitable model found, will use mock data")
|
||||
|
||||
Reference in New Issue
Block a user