Update app.py

app.py

@@ -1,67 +1,58 @@
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModel
-import tempfile
-import os
-
-# ============================
-#  CONFIG
-# ============================
-
-# Tu peux changer ce modèle par un Mistral quand tu en trouves un adapté CPU
-# Exemple possible : "mistralai/Mistral-7B-v0.1" (très lourd pour 12 Go CPU)
-# Pour rester safe sur CPU, je mets un modèle plus léger par défaut :
-MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+from diffusers import Flux2Pipeline
+import tempfile, os
 
+REPO_ID = "black-forest-labs/FLUX.2-klein-4B"
 device = "cpu"
-dtype = torch.float32  # sur CPU, reste en float32 pour éviter les emmerdes
-
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModel.from_pretrained(MODEL_NAME, torch_dtype=dtype)
-model.to(device)
-model.eval()
-
-# ============================
-#  FONCTION D'ENCODAGE
-# ============================
+dtype = torch.float32   # CPU-safe
 
-def encode_text(prompt: str):
-    if not prompt or not prompt.strip():
-        raise gr.Error("Le prompt ne peut pas être vide.")
+# On charge la pipeline mais on supprime tout sauf le text encoder
+pipe = Flux2Pipeline.from_pretrained(
+    REPO_ID,
+    torch_dtype=dtype,
+    low_cpu_mem_usage=True,
+)
 
-    with torch.no_grad():
-        inputs = tokenizer(
-            prompt,
-            return_tensors="pt",
-            truncation=True,
-            max_length=256
-        ).to(device)
+# On supprime tout ce qui n'est PAS le text encoder
+pipe.transformer = None
+pipe.vae = None
+pipe.scheduler = None
 
-        outputs = model(**inputs)
+# On garde :
+# - pipe.tokenizer
+# - pipe.text_encoder
+# - pipe.text_encoder_2 (si présent)
+# - encode_prompt()
 
-        # pooling simple : moyenne sur la séquence
-        last_hidden = outputs.last_hidden_state  # [1, seq_len, hidden]
-        emb = last_hidden.mean(dim=1).squeeze(0)  # [hidden]
+pipe.to(device)
 
-        # sauvegarde dans un fichier temporaire
-        fd, path = tempfile.mkstemp(suffix=".pt")
-        os.close(fd)
-        torch.save(emb.cpu(), path)
+@torch.no_grad()
+def encode_text(prompt: str):
+    if not prompt.strip():
+        raise gr.Error("Prompt vide")
+
+    # encode_prompt = embeddings EXACTS attendus par FLUX2
+    prompt_embeds, _, _ = pipe.encode_prompt(
+        prompt=prompt,
+        device=device,
+        num_images_per_prompt=1,
+        do_classifier_free_guidance=False,
+    )
+
+    # Sauvegarde dans un fichier .pt
+    fd, path = tempfile.mkstemp(suffix=".pt")
+    os.close(fd)
+    torch.save(prompt_embeds.cpu(), path)
 
-    # IMPORTANT : on renvoie le chemin du fichier
     return path
 
-# ============================
-#  INTERFACE GRADIO
-# ============================
-
 demo = gr.Interface(
     fn=encode_text,
-    inputs=gr.Textbox(label="Prompt", placeholder="Texte à encoder..."),
-    outputs=gr.File(label="Fichier .pt des embeddings"),
-    title="Text Encoder — CPU",
-    description="Encode un texte en vecteur et renvoie un fichier .pt (PyTorch tensor).",
+    inputs=gr.Textbox(label="Prompt"),
+    outputs=gr.File(label="Embeddings FLUX2 (.pt)"),
+    title="FLUX.2 Klein — Text Encoder Officiel",
+    description="Renvoie les embeddings EXACTS que FLUX2 Klein attend.",
 )
 
-# API name pour gradio_client : /encode_text
 demo.launch()