Update app.py

app.py

@@ -1,26 +1,30 @@
 import torch
 import gradio as gr
-from transformers import Qwen3ForCausalLM 
+from transformers import AutoTokenizer, Qwen2ForCausalLM
 from diffusers import Flux2Pipeline
 
-# Charger FLUX.2 Klein COMPLET pour récupérer le vrai text encoder (Qwen3)
+device = "cpu"
+dtype = torch.float32
+
+# Charger FLUX pour récupérer les embedder internes
 pipe = Flux2Pipeline.from_pretrained(
     "black-forest-labs/FLUX.2-klein-4B",
-    torch_dtype=torch.float32,
-    transformer=None,
-    vae=None,
-    scheduler=None, 
-    feature_extractor =None, 
+    torch_dtype=dtype,
     low_cpu_mem_usage=True,
 )
 
-# Récupération du tokenizer + text_encoder (Qwen3ForCausalLM)
-tokenizer = pipe.tokenizer
-text_encoder = Qwen3ForCausalLM.from_pretained(pipe.text_encoder)
+pos_embedder = pipe.transformer.pos_embedder
+extra_embedder = pipe.transformer.extra_embedder
+
+# Charger Qwen (encodeur texte)
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B")
+text_encoder = Qwen2ForCausalLM.from_pretrained(
+    "Qwen/Qwen2-1.5B",
+    torch_dtype=dtype,
+)
 
-def encode_text(prompt: str):
-    # Tokenisation simple
-    inputs = tokenizer(
+def encode(prompt):
+    tokens = tokenizer(
         prompt,
         return_tensors="pt",
         padding=True,
@@ -28,29 +32,29 @@ def encode_text(prompt: str):
         max_length=512,
     )
 
-    # Encodage texte → embeddings 2560 dims
     with torch.inference_mode():
-        outputs = text_encoder(
-            **inputs,
+        out = text_encoder(
+            **tokens,
             output_hidden_states=True,
             use_cache=False,
         )
 
-    # Dernière couche cachée = embeddings texte
-    embeds = outputs.hidden_states[-1]  # [B, L, 2560]
+    text = out.hidden_states[-1]  # [1, L, 2560]
+
+    pos = pos_embedder(text)      # [1, L, 2560]
+    extra = extra_embedder(text)  # [1, L, 2140]
+
+    final = torch.cat([text, pos, extra], dim=-1)  # [1, L, 7260]
 
-    # Sauvegarde dans un fichier .pt
-    torch.save(embeds, "embeds.pt")
+    torch.save(final, "embeds.pt")
 
-    return f"shape={tuple(embeds.shape)}", "embeds.pt"
+    return str(final.shape), "embeds.pt"
 
-# Interface Gradio
 demo = gr.Interface(
-    fn=encode_text,
+    fn=encode,
     inputs=gr.Textbox(label="Prompt"),
-    outputs=[gr.Textbox(label="Shape"), gr.File(label="Embeddings (.pt)")],
-    title="FLUX.2 Klein — Text Embedder (Qwen3 2560 dims)",
-    description="Encodeur texte officiel de FLUX.2 Klein (Qwen3ForCausalLM).",
+    outputs=[gr.Textbox(label="Shape"), gr.File(label="Embeddings")],
+    title="FLUX Klein — External Encoder (7260 dims)"
 )
 
 demo.launch()