Update app.py

app.py

@@ -1,27 +1,10 @@
 import torch
 import gradio as gr
 from transformers import AutoTokenizer, Qwen2ForCausalLM
-from diffusers import Flux2Transformer2DModel
 
 device = "cpu"
 dtype = torch.float32
 
-# Charger uniquement le transformer FLUX (léger)
-transformer = Flux2Transformer2DModel.from_pretrained(
-    "black-forest-labs/FLUX.2-klein-4B",
-    subfolder="transformer",
-    torch_dtype=dtype,
-)
-
-# Modules internes
-pos_embed = transformer.pos_embed        # [1, 4096, 2560]
-x_embedder = transformer.x_embedder      # module → 2140 dims
-
-# Libérer le reste
-del transformer
-torch.cuda.empty_cache() if torch.cuda.is_available() else None
-
-# Charger Qwen (encodeur texte)
 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B")
 text_encoder = Qwen2ForCausalLM.from_pretrained(
     "Qwen/Qwen2-1.5B",
@@ -44,26 +27,23 @@ def encode(prompt):
             use_cache=False,
         )
 
-    text = out.hidden_states[-1]  # [1, L, 2560]
-    L = text.shape[1]
-
-    # 🔥 pos_embed n'est PAS un module → on slice
-    pos = pos_embed[:, :L, :]     # [1, L, 2560]
-
-    # extra embedder est un module → on l'appelle
-    extra = x_embedder(text)      # [1, L, 2140]
-
-    final = torch.cat([text, pos, extra], dim=-1)  # [1, L, 7260]
+    embeds = out.hidden_states[-1]          # [1, L, 2560]
+    pooled = embeds.mean(dim=1)            # [1, 2560]
 
-    torch.save(final, "embeds.pt")
+    torch.save(embeds, "embeds.pt")
+    torch.save(pooled, "pooled.pt")
 
-    return str(final.shape), "embeds.pt"
+    return str(embeds.shape), "embeds.pt", "pooled.pt"
 
 demo = gr.Interface(
     fn=encode,
     inputs=gr.Textbox(label="Prompt"),
-    outputs=[gr.Textbox(label="Shape"), gr.File(label="Embeddings")],
-    title="FLUX Klein — External Encoder (7260 dims)"
+    outputs=[
+        gr.Textbox(label="Shape"),
+        gr.File(label="Embeddings 2560"),
+        gr.File(label="Pooled 2560")
+    ],
+    title="External Text Encoder — 2560 dims"
 )
 
 demo.launch()