Update app.py

app.py

@@ -6,27 +6,30 @@ from transformers import AutoTokenizer, Qwen2ForCausalLM
 device = "cpu"
 dtype = torch.float32
 
+# 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: str):
-    # 1) Nettoyage du prompt
+    # Nettoyage du prompt
     if prompt is None:
         prompt = ""
     prompt_clean = prompt.strip()
 
-    # 2) Si vide → on force un token valide
+    # Si vide -> token de secours
     if prompt_clean == "":
-        if tokenizer.eos_token is not None:
+        if tokenizer.eos_token:
             prompt_clean = tokenizer.eos_token
         else:
             prompt_clean = "."
 
+    # Tokenisation
     tokens = tokenizer(
         prompt_clean,
         return_tensors="pt",
@@ -35,17 +38,23 @@ def encode(prompt: str):
         max_length=512,
     )
 
-    with torch.inference_mode():
-        out = text_encoder(
-            **tokens,
-            output_hidden_states=True,
-            use_cache=False,
-        )
+    # Encodage Qwen (SANS inference_mode)
+    out = text_encoder(
+        **tokens,
+        output_hidden_states=True,
+        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]
 
-    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]
+    # pooled -> moyenne
+    pooled = embeds_2048.mean(dim=1)     # [1, 2048]
 
+    # Sauvegarde
     torch.save(embeds_2048, "embeds.pt")
     torch.save(pooled, "pooled.pt")
 
@@ -57,7 +66,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)"
 )