Update app.py

app.py

@@ -6,15 +6,15 @@ from transformers import AutoTokenizer, Qwen2ForCausalLM
 device = "cpu"
 dtype = torch.float32
 
-# Qwen 0.5B (léger, CPU OK)
+# Qwen 0.5B = hidden_size 896
 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)
+# Projection 896 -> 2048 pour FLUX.1-Schnell
+proj = nn.Linear(896, 2048)
 
 def encode(prompt: str):
     # Nettoyage du prompt
@@ -45,11 +45,11 @@ def encode(prompt: str):
         use_cache=False,
     )
 
-    # Embeddings Qwen (1536 dims)
-    embeds_1536 = out.hidden_states[-1]  # [1, L, 1536]
+    # Embeddings Qwen (896 dims)
+    embeds_896 = out.hidden_states[-1]  # [1, L, 896]
 
     # Projection -> 2048 dims
-    embeds_2048 = proj(embeds_1536)      # [1, L, 2048]
+    embeds_2048 = proj(embeds_896)      # [1, L, 2048]
 
     # pooled -> moyenne
     pooled = embeds_2048.mean(dim=1)     # [1, 2048]