Update app.py

app.py

@@ -6,19 +6,29 @@ from transformers import AutoTokenizer, Qwen2ForCausalLM
 device = "cpu"
 dtype = torch.float32
 
-# Charger Qwen 0.5B (léger, CPU OK)
 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B")
 text_encoder = Qwen2ForCausalLM.from_pretrained(
     "Qwen/Qwen2-0.5B",
     torch_dtype=dtype,
 )
 
-# Projection 1536 → 2048 (pour FLUX.1-Schnell)
 proj = nn.Linear(1536, 2048)
 
-def encode(prompt):
+def encode(prompt: str):
+    # 1) Nettoyage du prompt
+    if prompt is None:
+        prompt = ""
+    prompt_clean = prompt.strip()
+
+    # 2) Si vide → on force un token valide
+    if prompt_clean == "":
+        if tokenizer.eos_token is not None:
+            prompt_clean = tokenizer.eos_token
+        else:
+            prompt_clean = "."
+
     tokens = tokenizer(
-        prompt,
+        prompt_clean,
         return_tensors="pt",
         padding=True,
         truncation=True,
@@ -32,16 +42,10 @@ def encode(prompt):
             use_cache=False,
         )
 
-    # Embeddings Qwen 1536 dims
-    embeds_1536 = out.hidden_states[-1]  # [1, L, 1536]
-
-    # Projection → 2048 dims
-    embeds_2048 = proj(embeds_1536)      # [1, L, 2048]
-
-    # pooled → moyenne
-    pooled = embeds_2048.mean(dim=1)     # [1, 2048]
+    embeds_1536 = out.hidden_states[-1]      # [1, L, 1536]
+    embeds_2048 = proj(embeds_1536)         # [1, L, 2048]
+    pooled = embeds_2048.mean(dim=1)        # [1, 2048]
 
-    # Sauvegarde
     torch.save(embeds_2048, "embeds.pt")
     torch.save(pooled, "pooled.pt")
 
@@ -53,7 +57,7 @@ demo = gr.Interface(
     outputs=[
         gr.Textbox(label="Shape"),
         gr.File(label="Embeddings 2048"),
-        gr.File(label="Pooled 2048")
+        gr.File(label="Pooled 2048"),
     ],
     title="External Text Encoder — 2048 dims (FLUX.1‑Schnell)"
 )