Update app.py

app.py

@@ -3,12 +3,13 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import tempfile, os
 
-# Le vrai CausalLM utilisé par FLUX2 Klein
+# Qwen 2.5 1.5B Instruct
 REPO_ID = "Qwen/Qwen2.5-1.5B-Instruct"
 
 device = "cpu"
 dtype = torch.float32
 
+# Charger UNIQUEMENT le CausalLM
 tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
 model = AutoModelForCausalLM.from_pretrained(
     REPO_ID,
@@ -26,6 +27,7 @@ def encode_text(prompt: str):
     if not prompt.strip():
         raise gr.Error("Prompt vide")
 
+    # Tokenisation simple
     inputs = tokenizer(
         prompt,
         return_tensors="pt",
@@ -33,12 +35,14 @@ def encode_text(prompt: str):
         max_length=256
     ).to(device)
 
+    # Sortie Qwen : hidden_states = [1, seq_len, 4096]
     outputs = model.model(**inputs, output_hidden_states=True)
-    hidden = outputs.hidden_states[-1]  # [1, seq_len, 4096]
+    hidden = outputs.hidden_states[-1]
 
     # Projection FLUX2 Klein
-    projected = project_out(hidden)     # [1, seq_len, 7680]
+    projected = project_out(hidden)  # [1, seq_len, 7680]
 
+    # Sauvegarde
     fd, path = tempfile.mkstemp(suffix=".pt")
     os.close(fd)
     torch.save(projected.cpu(), path)
@@ -49,8 +53,8 @@ demo = gr.Interface(
     fn=encode_text,
     inputs=gr.Textbox(label="Prompt"),
     outputs=gr.File(label="Embeddings FLUX2 (.pt)"),
-    title="FLUX.2 Klein — Text Encoder Qwen Direct",
-    description="Encode le texte avec Qwen2.5 + projection FLUX2 (4096→7680).",
+    title="FLUX2 Klein — Encoder Qwen2.5 1.5B",
+    description="Encode le texte avec Qwen2.5 1.5B + projection FLUX2 (4096→7680).",
 )
 
 demo.launch()