src/flux/modules/conditioner.py

from torch import Tensor, nn
from transformers import (CLIPTextModel, CLIPTokenizer, T5EncoderModel,
                          T5Tokenizer)
import os

class HFEmbedder(nn.Module):
    def __init__(self, version: str, max_length: int, is_clip, **hf_kwargs):
        super().__init__()
        self.is_clip = is_clip
        self.max_length = max_length
        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.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(
            text,
            truncation=True,
            max_length=self.max_length,
            return_length=False,
            return_overflowing_tokens=False,
            padding="max_length",
            return_tensors="pt",
        )

        outputs = self.hf_module(
            input_ids=batch_encoding["input_ids"].to(self.hf_module.device),
            attention_mask=None,
            output_hidden_states=False,
        )
        return outputs[self.output_key]