Update app.py

app.py

@@ -3,53 +3,50 @@ import gradio as gr
 from diffusers import Flux2Pipeline
 
 # -----------------------------
-#  ENCODEUR FLUX MINIMAL
+#  CHARGEMENT MINIMAL DU TEXT ENCODER FLUX
 # -----------------------------
-def encode_prompt(prompt: str):
-    # On charge uniquement le tokenizer + text_encoder
-    pipe = Flux2Pipeline.from_pretrained(
-        "black-forest-labs/FLUX.2-klein-4B",
-        torch_dtype=torch.float32,
-        low_cpu_mem_usage=True,
-
-        # 🔥 On désactive tout ce qui n'est PAS utilisé par encode_prompt()
-        transformer=None,        # pas besoin du modèle de diffusion
-        vae=None,                # pas besoin pour encoder du texte
-        scheduler=None,          # inutile
-        safety_checker=None,     # inutile
-        feature_extractor=None,  # inutile
-    ).to("cpu")
-
-    with torch.inference_mode():
-        prompt_embeds = pipe.encode_prompt(prompt)
-
-    return prompt_embeds
-
+pipe = Flux2Pipeline.from_pretrained(
+    "black-forest-labs/FLUX.2-klein-4B",
+    torch_dtype=torch.float32,
+    low_cpu_mem_usage=True,
+
+    # On désactive tout ce qui n'est pas utile pour encode_prompt()
+    transformer=None,
+    vae=None,
+    scheduler=None,
+    safety_checker=None,
+    feature_extractor=None,
+)
 
 # -----------------------------
-#  FONCTION GRADIO
+#  ENCODEUR
 # -----------------------------
-def run_encoder(prompt):
-    embeds = encode_prompt(prompt)
+def encode_text(prompt: str):
+    if not prompt.strip():
+        return "Prompt vide", None
 
-    # On renvoie la forme pour vérifier
-    shape = tuple(embeds.shape)
+    with torch.inference_mode():
+        embeds = pipe.encode_prompt(prompt)
 
-    # On sauvegarde pour un autre Space si besoin
-    torch.save(embeds.cpu(), "embeds.pt")
+    # Sauvegarde dans un fichier temporaire
+    file_path = "embeds.pt"
+    torch.save(embeds.cpu(), file_path)
 
-    return f"Embeddings générés : shape = {shape}"
+    return f"Embeddings générés : {tuple(embeds.shape)}", file_path
 
 
 # -----------------------------
 #  INTERFACE GRADIO
 # -----------------------------
 demo = gr.Interface(
-    fn=run_encoder,
+    fn=encode,
     inputs=gr.Textbox(label="Prompt", placeholder="Écris ton texte ici..."),
-    outputs=gr.Textbox(label="Résultat"),
-    title="Encodeur Texte FLUX.2 — Version Minimaliste",
-    description="Génère les embeddings texte officiels de FLUX.2 sans charger le VAE ou le modèle de diffusion.",
+    outputs=[
+        gr.Textbox(label="Infos"),
+        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.",
 )
 
 demo.launch()