Update app.py

app.py

@@ -1,3 +1,5 @@
+# app.py (Space ENCODER)
+
 import torch
 import gradio as gr
 from diffusers import Flux2Pipeline
@@ -7,6 +9,7 @@ REPO_ID = "black-forest-labs/FLUX.2-klein-4B"
 device = "cpu"
 dtype = torch.float32
 
+# On NE CHARGE QUE la partie texte (comme tu voulais)
 pipe = Flux2Pipeline.from_pretrained(
     REPO_ID,
     transformer=None,
@@ -16,34 +19,40 @@ pipe = Flux2Pipeline.from_pretrained(
     low_cpu_mem_usage=True,
 )
 
-# On supprime les parties inutiles
+# On s'assure de ne garder que ce qui sert à l'encodage texte
 pipe.transformer = None
 pipe.vae = None
 pipe.scheduler = None
 
 pipe.to(device)
 
+
 @torch.no_grad()
 def encode_text(prompt: str):
     if not prompt.strip():
         raise gr.Error("Prompt vide")
 
-    # FLUX2 Klein attend un format chat Qwen3
-    messages = [
-        {"role": "user", "content": prompt}
-    ]
-
-    prompt_embeds, _, _ = pipe.encode_prompt(
-        prompt=messages,
+    # encode_prompt renvoie EXACTEMENT ce que FLUX2 veut
+    prompt_embeds, pooled_prompt_embeds, text_ids = pipe.encode_prompt(
+        prompt=prompt,              # IMPORTANT : string simple, pas dict
         device=device,
         num_images_per_prompt=1,
     )
 
+    # On sauvegarde les trois tensors ensemble
+    data = {
+        "prompt_embeds": prompt_embeds.cpu(),
+        "pooled_prompt_embeds": pooled_prompt_embeds.cpu(),
+        "text_ids": text_ids.cpu(),
+    }
+
     fd, path = tempfile.mkstemp(suffix=".pt")
     os.close(fd)
-    torch.save(prompt_embeds.cpu(), path)
+    torch.save(data, path)
+
     return path
 
+
 demo = gr.Interface(
     fn=encode_text,
     inputs=gr.Textbox(label="Prompt"),
@@ -52,4 +61,6 @@ demo = gr.Interface(
     description="Renvoie les embeddings EXACTS que FLUX2 Klein attend.",
 )
 
+# api_name par défaut = "/predict", si tu veux explicitement :
+# demo.launch(api_name="/encode_text")
 demo.launch()