new file: .gitignore

.gitignore

@@ -0,0 +1,5 @@
+# Python
+__pycache__/
+*.pyc
+heatmap/
+src/gradio_utils/gradio_outputs/

app_test.py

@@ -1,37 +0,0 @@
-import gradio as gr
-
-def calculator(num1, operation, num2):
-    if operation == "add":
-        return num1 + num2
-    elif operation == "subtract":
-        return num1 - num2
-    elif operation == "multiply":
-        return num1 * num2
-    elif operation == "divide":
-        return num1 / num2
-
-with gr.Blocks() as demo:
-    with gr.Row():
-        with gr.Column():
-            num_1 = gr.Number(value=4)
-            operation = gr.Radio(["add", "subtract", "multiply", "divide"])
-            num_2 = gr.Number(value=0)
-            submit_btn = gr.Button(value="Calculate")
-        with gr.Column():
-            result = gr.Number()
-
-    submit_btn.click(
-        calculator, inputs=[num_1, operation, num_2], outputs=[result], api_name=False
-    )
-    examples = gr.Examples(
-        examples=[
-            [5, "add", 3],
-            [4, "divide", 2],
-            [-4, "multiply", 2.5],
-            [0, "subtract", 1.2],
-        ],
-        inputs=[num_1, operation, num_2],
-    )
-
-if __name__ == "__main__":
-    demo.launch(show_api=False)

src/flux/modules/conditioner.py

@@ -4,6 +4,7 @@ from transformers import (CLIPTextModel, CLIPTokenizer, T5EncoderModel,
 import os
 import torch
 
+'''
 class HFEmbedder(nn.Module):
     def __init__(self, version: str, max_length: int, is_clip, **hf_kwargs):
         super().__init__()
@@ -133,3 +134,38 @@ class HFEmbedder(nn.Module):
             raise # Re-raise the error after logging
 
         return outputs[self.output_key]
+'''
+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.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained("black-forest-labs/FLUX.1-dev/tokenizer_2", max_length=max_length)
+            self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs)
+
+        self.hf_module = self.hf_module.eval().requires_grad_(False)
+
+    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]