Update app.py

app.py

@@ -1,16 +1,21 @@
 import torch
+import torch.nn as nn
 import gradio as gr
 from transformers import AutoTokenizer, Qwen2ForCausalLM
 
 device = "cpu"
 dtype = torch.float32
 
-tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B")
+# Charger Qwen 0.5B (léger, CPU OK)
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B")
 text_encoder = Qwen2ForCausalLM.from_pretrained(
-    "Qwen/Qwen2-3B",
+    "Qwen/Qwen2-0.5B",
     torch_dtype=dtype,
 )
 
+# Projection 1536 → 2048 (pour FLUX.1-Schnell)
+proj = nn.Linear(1536, 2048)
+
 def encode(prompt):
     tokens = tokenizer(
         prompt,
@@ -27,23 +32,30 @@ def encode(prompt):
             use_cache=False,
         )
 
-    embeds = out.hidden_states[-1]          # [1, L, 2560]
-    pooled = embeds.mean(dim=1)            # [1, 2560]
+    # 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]
 
-    torch.save(embeds, "embeds.pt")
+    # Sauvegarde
+    torch.save(embeds_2048, "embeds.pt")
     torch.save(pooled, "pooled.pt")
 
-    return str(embeds.shape), "embeds.pt", "pooled.pt"
+    return str(embeds_2048.shape), "embeds.pt", "pooled.pt"
 
 demo = gr.Interface(
     fn=encode,
     inputs=gr.Textbox(label="Prompt"),
     outputs=[
         gr.Textbox(label="Shape"),
-        gr.File(label="Embeddings 2560"),
-        gr.File(label="Pooled 2560")
+        gr.File(label="Embeddings 2048"),
+        gr.File(label="Pooled 2048")
     ],
-    title="External Text Encoder — 2560 dims"
+    title="External Text Encoder — 2048 dims (FLUX.1‑Schnell)"
 )
 
 demo.launch()