Update app.py

app.py

@@ -2,14 +2,8 @@ import torch
 import gradio as gr
 from diffusers import Flux2Pipeline
 
-device = "cpu"
-dtype = torch.float32
-
-# On charge uniquement tokenizer + text_encoder
 pipe = Flux2Pipeline.from_pretrained(
     "black-forest-labs/FLUX.2-klein-4B",
-    torch_dtype=dtype,
-    low_cpu_mem_usage=True,
     transformer=None,
     vae=None,
     scheduler=None,
@@ -18,21 +12,16 @@ pipe = Flux2Pipeline.from_pretrained(
 )
 
 tokenizer = pipe.tokenizer
-text_encoder = pipe.text_encoder.to(device)
-
+text_encoder = pipe.text_encoder
 
 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():
         outputs = text_encoder(
@@ -41,25 +30,15 @@ def encode_text(prompt: str):
             use_cache=False,
         )
 
-    # On prend la dernière couche cachée : [B, L, D]
-    embeds = outputs.hidden_states[-1].to("cpu")
-
-    file_path = "embeds.pt"
-    torch.save(embeds, file_path)
-
-    return f"Embeddings shape: {tuple(embeds.shape)}", file_path
+    embeds = outputs.hidden_states[-1]  # [B, L, 4096]
 
+    torch.save(embeds, "embeds.pt")
+    return f"shape={tuple(embeds.shape)}", "embeds.pt"
 
 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"),
-    ],
-    title="Encodeur Texte FLUX.2 (Mistral-3) — Minimal",
-    description="Encode le prompt avec le text encoder FLUX.2 sans chat template.",
+    inputs=gr.Textbox(),
+    outputs=[gr.Textbox(), gr.File()],
 )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()