Update app.py

app.py

@@ -2,17 +2,23 @@ import torch
 import gradio as gr
 from diffusers import Flux2Pipeline
 
-# Charger FLUX.2 COMPLET
+# Charger FLUX.2 Klein COMPLET pour récupérer le vrai text encoder (Qwen3)
 pipe = Flux2Pipeline.from_pretrained(
     "black-forest-labs/FLUX.2-klein-4B",
     torch_dtype=torch.float32,
+    transformer=None,
+    vae=None,
+    scheduler=None, 
+    feature_extractor =None, 
     low_cpu_mem_usage=True,
 )
 
+# Récupération du tokenizer + text_encoder (Qwen3ForCausalLM)
 tokenizer = pipe.tokenizer
 text_encoder = pipe.text_encoder
 
 def encode_text(prompt: str):
+    # Tokenisation simple
     inputs = tokenizer(
         prompt,
         return_tensors="pt",
@@ -21,6 +27,7 @@ def encode_text(prompt: str):
         max_length=512,
     )
 
+    # Encodage texte → embeddings 2560 dims
     with torch.inference_mode():
         outputs = text_encoder(
             **inputs,
@@ -28,15 +35,21 @@ def encode_text(prompt: str):
             use_cache=False,
         )
 
-    embeds = outputs.hidden_states[-1]  # [B, L, 4096]
+    # Dernière couche cachée = embeddings texte
+    embeds = outputs.hidden_states[-1]  # [B, L, 2560]
 
+    # Sauvegarde dans un fichier .pt
     torch.save(embeds, "embeds.pt")
+
     return f"shape={tuple(embeds.shape)}", "embeds.pt"
 
+# Interface Gradio
 demo = gr.Interface(
     fn=encode_text,
-    inputs=gr.Textbox(),
-    outputs=[gr.Textbox(), gr.File()],
+    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).",
 )
 
 demo.launch()