Update app.py

app.py

@@ -3,28 +3,29 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import tempfile, os
 
-# Le vrai text encoder utilisé par FLUX.2 Klein 4B
-REPO_ID = "black-forest-labs/FLUX.2-klein-4B/text_encoder"
+# Le vrai CausalLM utilisé par FLUX2 Klein
+REPO_ID = "Qwen/Qwen2.5-7B-Instruct"
 
 device = "cpu"
 dtype = torch.float32
 
-# Charger UNIQUEMENT le CausalLM Qwen utilisé par FLUX2 Klein
 tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
-text_encoder = AutoModelForCausalLM.from_pretrained(
+model = AutoModelForCausalLM.from_pretrained(
     REPO_ID,
     torch_dtype=dtype,
     low_cpu_mem_usage=True,
 )
-text_encoder.to(device)
-text_encoder.eval()
+model.to(device)
+model.eval()
+
+# Projection FLUX2 Klein : 4096 → 7680
+project_out = torch.nn.Linear(4096, 7680, bias=False)
 
 @torch.no_grad()
 def encode_text(prompt: str):
     if not prompt.strip():
         raise gr.Error("Prompt vide")
 
-    # Tokenisation simple (pas de chat template)
     inputs = tokenizer(
         prompt,
         return_tensors="pt",
@@ -32,14 +33,12 @@ def encode_text(prompt: str):
         max_length=256
     ).to(device)
 
-    # Sortie Qwen3 : hidden_states = [1, seq_len, 4096]
-    outputs = text_encoder.model(**inputs, output_hidden_states=True)
-    hidden = outputs.hidden_states[-1]  # dernière couche
+    outputs = model.model(**inputs, output_hidden_states=True)
+    hidden = outputs.hidden_states[-1]  # [1, seq_len, 4096]
 
-    # Projection FLUX2 : 4096 → 7680
-    projected = text_encoder.model.project_out(hidden)
+    # Projection FLUX2 Klein
+    projected = project_out(hidden)     # [1, seq_len, 7680]
 
-    # Sauvegarde
     fd, path = tempfile.mkstemp(suffix=".pt")
     os.close(fd)
     torch.save(projected.cpu(), path)
@@ -50,8 +49,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 Qwen3 Direct",
-    description="Encode le texte avec Qwen3 + projection FLUX2 (4096→7680).",
+    title="FLUX.2 Klein — Text Encoder Qwen Direct",
+    description="Encode le texte avec Qwen2.5 + projection FLUX2 (4096→7680).",
 )
 
 demo.launch()