Update app.py

app.py

@@ -2,15 +2,14 @@ import torch
 import gradio as gr
 from diffusers import Flux2Pipeline
 
-# -----------------------------
-#  CHARGEMENT MINIMAL DU TEXT ENCODER FLUX
-# -----------------------------
+device = "cpu"
+dtype = torch.float32
+
+# On charge uniquement tokenizer + text_encoder
 pipe = Flux2Pipeline.from_pretrained(
     "black-forest-labs/FLUX.2-klein-4B",
-    torch_dtype=torch.float32,
+    torch_dtype=dtype,
     low_cpu_mem_usage=True,
-
-    # On désactive tout ce qui n'est pas utile pour encode_prompt()
     transformer=None,
     vae=None,
     scheduler=None,
@@ -18,35 +17,49 @@ pipe = Flux2Pipeline.from_pretrained(
     feature_extractor=None,
 )
 
-# -----------------------------
-#  ENCODEUR
-# -----------------------------
+tokenizer = pipe.tokenizer
+text_encoder = pipe.text_encoder.to(device)
+
+
 def encode_text(prompt: str):
     if not prompt.strip():
         return "Prompt vide", None
 
+    # Tokenisation simple, sans chat template
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=512,
+    ).to(device)
+
     with torch.inference_mode():
-        embeds = pipe.encode_prompt(prompt)
+        outputs = text_encoder(
+            **inputs,
+            output_hidden_states=True,
+            use_cache=False,
+        )
+
+    # On prend la dernière couche cachée : [B, L, D]
+    embeds = outputs.hidden_states[-1].to("cpu")
 
-    # Sauvegarde dans un fichier temporaire
     file_path = "embeds.pt"
-    torch.save(embeds.cpu(), file_path)
+    torch.save(embeds, file_path)
 
-    return f"Embeddings générés : {tuple(embeds.shape)}", file_path
+    return f"Embeddings shape: {tuple(embeds.shape)}", file_path
 
 
-# -----------------------------
-#  INTERFACE GRADIO
-# -----------------------------
 demo = gr.Interface(
     fn=encode_text,
     inputs=gr.Textbox(label="Prompt", placeholder="Écris ton texte ici..."),
     outputs=[
         gr.Textbox(label="Infos"),
-        gr.File(label="Fichier .pt des embeddings")
+        gr.File(label="Fichier .pt des embeddings"),
     ],
-    title="Encodeur Texte FLUX.2 — Minimal",
-    description="Encodeur officiel FLUX.2 (Mistral-3-Small). Génère des embeddings compatibles avec Flux2Pipeline.",
+    title="Encodeur Texte FLUX.2 (Mistral-3) — Minimal",
+    description="Encode le prompt avec le text encoder FLUX.2 sans chat template.",
 )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()