From 832af7212939c7e47ab17fe3202697dfcccfb504 Mon Sep 17 00:00:00 2001 From: Junst Date: Wed, 24 Sep 2025 14:09:01 +0900 Subject: [PATCH] Remove ALL fallback mechanisms - require real attention capture MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 🚫 **Complete Fallback Removal**: 1. **Removed All Fallback Mechanisms**: - No more fallback attention creation from model weights - No more simple fallback attention output - No more patches_replace fallback - No more dummy text embeddings 2. **Require Real Components**: - Real attention capture is mandatory - Real text encoder and tokenizer required - Real concept embeddings required - No mock data allowed 3. **Strict Error Handling**: - RuntimeError if no real attention captured - RuntimeError if text encoder/tokenizer missing - RuntimeError if text encoder methods not found - No graceful degradation This ensures only real attention capture and real text encoding are used, as explicitly requested by the user. --- concept_attention.py | 124 +++++-------------------------------------- 1 file changed, 12 insertions(+), 112 deletions(-) diff --git a/concept_attention.py b/concept_attention.py index 1da3d30..bdd4d3f 100644 --- a/concept_attention.py +++ b/concept_attention.py @@ -279,14 +279,8 @@ class ConceptAttention: for concept in concepts: # Check if text_encoder and tokenizer are valid if text_encoder is None or tokenizer is None: - logger.warning(f"Text encoder or tokenizer is None, creating fallback embedding for '{concept}'") - # Create fallback embedding with proper dtype - if hasattr(self.model, 'parameters'): - params = list(self.model.parameters()) - model_dtype = params[0].dtype if params else torch.float32 - else: - model_dtype = torch.float32 - embedding = torch.randn(1, 512, device=self.device, dtype=model_dtype) + logger.error(f"❌ Text encoder or tokenizer is None for concept '{concept}'! Real text encoding is required.") + raise RuntimeError(f"Text encoder and tokenizer are required for concept '{concept}'. No fallback allowed.") else: # Tokenize the concept tokens = tokenizer(concept, return_tensors="pt", padding=True, truncation=True) @@ -301,13 +295,8 @@ class ConceptAttention: elif hasattr(text_encoder, 'encode'): embedding = text_encoder.encode(tokens["input_ids"]) else: - # Fallback: create dummy embedding with proper dtype - if hasattr(self.model, 'parameters'): - params = list(self.model.parameters()) - model_dtype = params[0].dtype if params else torch.float32 - else: - model_dtype = torch.float32 - embedding = torch.randn(1, 512, device=self.device, dtype=model_dtype) + logger.error(f"❌ Text encoder has no recognized methods for concept '{concept}'!") + raise RuntimeError(f"Text encoder must have encode_text, forward, or encode method. No fallback allowed.") # Ensure embedding is on correct device and dtype embedding = embedding.to(device=self.device) @@ -316,15 +305,7 @@ class ConceptAttention: except Exception as e: logger.error(f"Error extracting concept embeddings: {e}") - # Create dummy embeddings as fallback with proper dtype - if hasattr(self.model, 'parameters'): - params = list(self.model.parameters()) - model_dtype = params[0].dtype if params else torch.float32 - else: - model_dtype = torch.float32 - - for concept in concepts: - concept_embeddings[concept] = torch.randn(1, 512, device=self.device, dtype=model_dtype) + raise RuntimeError(f"Failed to extract concept embeddings: {e}. No fallback allowed.") return concept_embeddings @@ -615,97 +596,16 @@ class ConceptAttention: except Exception as attn_error: logger.debug(f"Direct attention execution failed in block {block_idx}: {attn_error}") - # If still no outputs, try to create fallback attention outputs from model weights + # No fallback - require real attention capture if not self.attention_outputs: - logger.info("🔍 Creating fallback attention outputs from model weights") - - if hasattr(actual_model, 'double_blocks'): - double_blocks = actual_model.double_blocks - - for block_idx in [15, 16, 17, 18]: - if block_idx < len(double_blocks): - block = double_blocks[block_idx] - - if hasattr(block, 'img_attn'): - img_attn = block.img_attn - - # Create fallback attention output using model weights - if hasattr(img_attn, 'qkv') and hasattr(img_attn.qkv, 'weight'): - # Get the weight matrix - weight = img_attn.qkv.weight - logger.info(f"Creating fallback attention from weight shape: {weight.shape}") - - # Create a dummy input and compute output - dummy_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype) - - with torch.no_grad(): - # Compute output using the weight matrix - fallback_output = torch.matmul(dummy_input, weight.t()) - logger.info(f"🚀 Fallback attention output created, shape: {fallback_output.shape}") - - # Store the fallback output - hook_key = f"fallback_attention_block_{block_idx}_{id(img_attn)}" - self.attention_outputs[hook_key] = fallback_output - logger.info(f"📝 Stored fallback attention output: {hook_key}") - - # Break after first successful block - break - - # If still no outputs, create a simple fallback - if not self.attention_outputs: - logger.info("🔍 Creating simple fallback attention output") - fallback_output = torch.randn(1, 1024, 256, device=self.device, dtype=torch.float32) - hook_key = f"simple_fallback_attention" - self.attention_outputs[hook_key] = fallback_output - logger.info(f"📝 Stored simple fallback attention output: {hook_key}") + logger.error("❌ No attention outputs captured! Real attention capture is required.") + raise RuntimeError("Failed to capture real attention outputs. No fallback allowed.") - # If still no outputs, try to use ComfyUI's patches_replace system - if not self.attention_outputs: - logger.info("🔍 Trying ComfyUI patches_replace system as fallback") - - if hasattr(actual_model, 'apply_model'): - # Create proper inputs for Flux model - # Flux model expects: x (latent), t (timestep), c_crossattn (context), y (guidance) - x = torch.randn(1, 4, 64, 64, device=self.device, dtype=model_dtype) - t = torch.tensor([500], device=self.device, dtype=torch.long) - c_crossattn = torch.randn(1, 77, 2048, device=self.device, dtype=model_dtype) - y = torch.randn(1, 512, device=self.device, dtype=model_dtype) - - logger.info(f"Running apply_model with patches_replace system") - logger.info(f"Inputs: x={x.shape}, t={t.shape}, c_crossattn={c_crossattn.shape}, y={y.shape}") - - # Create a custom transformer_options with patches_replace - transformer_options = { - "patches_replace": { - "dit": { - ("double_block", 15): self._create_attention_capture_wrapper(15), - ("double_block", 16): self._create_attention_capture_wrapper(16), - ("double_block", 17): self._create_attention_capture_wrapper(17), - ("double_block", 18): self._create_attention_capture_wrapper(18), - } - } - } - - with torch.no_grad(): - output = actual_model.apply_model(x, t, c_crossattn=c_crossattn, y=y, transformer_options=transformer_options) - logger.info(f"🚀 apply_model with patches_replace completed, output shape: {output.shape}") - - # Store the output as attention output - hook_key = f"patches_replace_output_{id(actual_model)}" - self.attention_outputs[hook_key] = output - logger.info(f"📝 Stored patches_replace output: {hook_key}") - - except Exception as patches_error: - logger.debug(f"Attention capture failed: {patches_error}") - - # Log total attention outputs captured - logger.info(f"📊 Total attention outputs captured: {len(self.attention_outputs)}") - if self.attention_outputs: + # Log final results + logger.info(f"📊 Total attention outputs captured: {len(self.attention_outputs)}") logger.info(f"📊 Attention output keys: {list(self.attention_outputs.keys())}") - else: - logger.warning("⚠️ No attention outputs captured!") - - return self.attention_outputs + + return self.attention_outputs class ConceptAttentionProcessor: