modified: src/flux/modules/conditioner.py

src/flux/modules/conditioner.py

@@ -1,6 +1,6 @@
 from torch import Tensor, nn
 from transformers import (CLIPTextModel, CLIPTokenizer, T5EncoderModel,
-                          T5Tokenizer)
+                          T5Tokenizer, AutoConfig)
 import os
 
 class HFEmbedder(nn.Module):
@@ -12,35 +12,15 @@ class HFEmbedder(nn.Module):
 
         if self.is_clip:
             self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length)
-            self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs)
+            config = AutoConfig.from_pretrained(version, max_position_embeddings=77)
+            self.hf_module = CLIPTextModel.from_pretrained(version, config=config, **hf_kwargs)
+            #self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs)
         else:
             self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length)
             self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs)
 
         self.hf_module = self.hf_module.eval().requires_grad_(False)
 
-        # In your HFEmbedder class's __init__ method, or after loading the clip model
-        # Assuming 'self.hf_module' is the CLIPTextModel instance for your CLIP embedder
-        if self.is_clip: # Add this check if HFEmbedder is used for both T5 and CLIP
-            clip_model_config = self.hf_module.config
-            print(f"CLIP Model Version: {self.hf_module.name_or_path}") # or 'version' if that's the attribute
-            print(f"CLIP Config max_position_embeddings: {clip_model_config.max_position_embeddings}")
-
-            # Accessing the CLIPTextEmbeddings module
-            # Path is typically model.embeddings for CLIPTextModel
-            text_embeddings_module = self.hf_module.embeddings
-
-            print(f"CLIP Position Embedding Layer num_embeddings: {text_embeddings_module.position_embedding.num_embeddings}")
-            print(f"CLIP Position IDs buffer 'position_ids' (from CLIPTextEmbeddings) shape: {text_embeddings_module.position_ids.shape}")
-
-            if text_embeddings_module.position_ids.shape[1] != text_embeddings_module.position_embedding.num_embeddings:
-                print("CRITICAL WARNING: Mismatch between position_ids buffer length and actual embedding layer size!")
-            if clip_model_config.max_position_embeddings != text_embeddings_module.position_embedding.num_embeddings:
-                print("CRITICAL WARNING: Mismatch between config.max_position_embeddings and actual embedding layer size!")
-            if self.max_length != text_embeddings_module.position_embedding.num_embeddings:
-                print(f"WARNING: Tokenizer max_length ({self.max_length}) "
-                    f"does not match position embedding size ({text_embeddings_module.position_embedding.num_embeddings})")
-
 
     def forward(self, text: list[str]) -> Tensor:
         batch_encoding = self.tokenizer(