modified: src/flux/modules/conditioner.py

src/flux/modules/conditioner.py

@@ -11,25 +11,52 @@ class HFEmbedder(nn.Module):
         self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
 
         if self.is_clip:
-            #self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length)
-            self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained("/home/user/app/models/tokenizer", max_length=max_length, truncation=True)
+            self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length, truncation=True)
+            #self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained("/home/user/app/models/tokenizer", max_length=max_length, truncation=True)
             self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs)
 
-            # --- DEBUG ---
-            print(f"CLIP Tokenizer max length: {self.tokenizer.model_max_length}")
-            print(f"CLIP max position embeddings: {self.hf_module.config.max_position_embeddings}")
+            # --- DEBUG 信息 ---
+            print(f"--- CLIP Model Info ---")
+            print(f"  Requested version/path: {version}")
+            print(f"  Tokenizer loaded from: {getattr(self.tokenizer, 'name_or_path', 'Unknown')}")
+            print(f"  Model loaded from: {getattr(self.hf_module, 'name_or_path', 'Unknown')}")
+            print(f"  Tokenizer max length: {getattr(self.tokenizer, 'model_max_length', 'N/A')}")
+            print(f"  Model max position embeddings: {getattr(self.hf_module.config, 'max_position_embeddings', 'N/A')}")
+            
+            # 关键调试信息：词汇表大小
+            tokenizer_vocab_size = len(self.tokenizer.get_vocab()) if hasattr(self.tokenizer, 'get_vocab') else getattr(self.tokenizer, 'vocab_size', 'Unknown')
+            print(f"  Tokenizer vocab size (len(get_vocab())): {tokenizer_vocab_size}")
+            print(f"  Tokenizer vocab size (attribute): {getattr(self.tokenizer, 'vocab_size', 'N/A')}")
+            print(f"  Model config vocab size: {self.hf_module.config.vocab_size}")
+            print(f"  Actual model embedding weight shape: {self.hf_module.text_model.embeddings.token_embedding.weight.shape}")
+            print(f"-------------------------")
 
         else:
             self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length, truncation=True)
             self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs)
-            # --- DEBUG ---
-            '''print(f"T5 Tokenizer max length: {T5Tokenizer.model_max_length}")
-            print(f"T5 max position embeddings: {T5EncoderModel.config.max_position_embeddings}")'''
+            
+            # --- DEBUG 信息 ---
+            print(f"--- T5 Model Info ---")
+            print(f"  Requested version/path: {version}")
+            print(f"  Tokenizer loaded from: {getattr(self.tokenizer, 'name_or_path', 'Unknown')}")
+            print(f"  Model loaded from: {getattr(self.hf_module, 'name_or_path', 'Unknown')}")
+            print(f"  Tokenizer max length: {getattr(self.tokenizer, 'model_max_length', 'N/A')}")
+            print(f"  Model max position embeddings: {getattr(self.hf_module.config, 'd_model', 'N/A (T5 uses relative pos)')}") # T5 uses relative
+            tokenizer_vocab_size = len(self.tokenizer.get_vocab()) if hasattr(self.tokenizer, 'get_vocab') else getattr(self.tokenizer, 'vocab_size', 'Unknown')
+            print(f"  Tokenizer vocab size (len(get_vocab())): {tokenizer_vocab_size}")
+            print(f"  Tokenizer vocab size (attribute): {getattr(self.tokenizer, 'vocab_size', 'N/A')}")
+            print(f"  Model config vocab size: {self.hf_module.config.vocab_size}")
+            print(f"  Actual model embedding weight shape: {self.hf_module.encoder.embed_tokens.weight.shape}")
+            print(f"----------------------")
 
         self.hf_module = self.hf_module.eval().requires_grad_(False)
 
 
     def forward(self, text: list[str]) -> Tensor:
+         # Ensure text is a list
+        if isinstance(text, str):
+            text = [text]
+
         batch_encoding = self.tokenizer(
             text,
             truncation=True,
@@ -40,22 +67,53 @@ class HFEmbedder(nn.Module):
             return_tensors="pt",
         )
 
-        if self.is_clip:
-            flag = 'clip'
-        else:
-            flag = 't5'
-        print(f'foward {flag}')
+        encoder_type = 'clip' if self.is_clip else 't5'
+        print(f'Forward pass for {encoder_type}')
         input_ids = batch_encoding["input_ids"]
-        print(f"input_ids shape: {input_ids.shape}, max_length: {self.max_length}")  # Debug
-        assert input_ids.shape[1] == self.max_length, f"Sequence length {input_ids.shape[1]} does not match max_length {self.max_length}"
-        print(input_ids)
+        print(f"Input IDs shape: {input_ids.shape}, Max Length: {self.max_length}")
+        
+        # 更严格的断言
+        assert input_ids.shape == (len(text), self.max_length), f"Input IDs shape {input_ids.shape} does not match expected ({len(text)}, {self.max_length})"
+        print(f"Input IDs:\n{input_ids}")
 
-        print(f"{flag} self.tokenizer.vocab_size: {self.tokenizer.vocab_size}")  # Debug
-        print(f"{flag} self.hf_module.config.vocab_size: {self.hf_module.config.vocab_size}")  # Debug
+        # --- 关键调试：检查输入 ID 范围 ---
+        min_id, max_id = input_ids.min().item(), input_ids.max().item()
+        print(f"Input IDs range: [{min_id}, {max_id}]")
+
+        vocab_source = "tokenizer" if self.is_clip else "model_config"
+        vocab_size = len(self.tokenizer.get_vocab()) if self.is_clip and hasattr(self.tokenizer, 'get_vocab') else self.hf_module.config.vocab_size
+        print(f"Vocab size (from {vocab_source}): {vocab_size}")
+
+        if max_id >= vocab_size:
+            raise IndexError(f"Found input ID ({max_id}) >= vocab size ({vocab_size}). This will cause an embedding error.")
+
+        if min_id < 0:
+             raise IndexError(f"Found negative input ID ({min_id}). This is invalid.")
+
+        # 确保输入在正确的设备上
+        input_ids = input_ids.to(self.device)
+        attention_mask = batch_encoding["attention_mask"].to(self.device)
+
+        print(f"Input IDs device: {input_ids.device}")
+        print(f"Attention Mask device: {attention_mask.device}")
+        print(f"Model device: {next(self.hf_module.parameters()).device}")
+
+        try:
+            outputs = self.hf_module(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=False,
+            )
+        except IndexError as e:
+            # 捕获并提供更详细的错误上下文
+            print(f"*** IndexError caught during model forward pass ***")
+            print(f"Error: {e}")
+            print(f"Input IDs shape: {input_ids.shape}")
+            print(f"Input IDs range: [{input_ids.min().item()}, {input_ids.max().item()}]")
+            print(f"Model vocab size: {self.hf_module.config.vocab_size}")
+            if self.is_clip:
+                 print(f"Tokenizer vocab size: {len(self.tokenizer.get_vocab()) if hasattr(self.tokenizer, 'get_vocab') else 'N/A'}")
+                 print(f"Embedding layer weight shape: {self.hf_module.text_model.embeddings.token_embedding.weight.shape}")
+            raise # Re-raise the error after logging
 
-        outputs = self.hf_module(
-            input_ids=input_ids.to(self.hf_module.device),
-            attention_mask=batch_encoding["attention_mask"].to(self.hf_module.device),
-            output_hidden_states=False,
-        )
         return outputs[self.output_key]

src/flux/util.py

@@ -136,8 +136,8 @@ def load_t5(device: str | torch.device = "cuda", max_length: int = 77) -> HFEmbe
 
 
 def load_clip(device: str | torch.device = "cuda") -> HFEmbedder:
-    #return HFEmbedder("openai/clip-vit-base-patch32", max_length=77, is_clip=True, torch_dtype=torch.bfloat16).to(device)
-    return HFEmbedder("/home/user/app/models/text_encoder", max_length=77, is_clip=True, torch_dtype=torch.bfloat16).to(device)
+    return HFEmbedder("openai/clip-vit-base-patch32", max_length=77, is_clip=True, torch_dtype=torch.bfloat16).to(device)
+    #return HFEmbedder("/home/user/app/models/text_encoder", max_length=77, is_clip=True, torch_dtype=torch.bfloat16).to(device)
 
 
 def load_ae(name: str, device: str | torch.device = "cuda", hf_download: bool = True) -> AutoEncoder: