Update app.py

app.py

@@ -1,66 +1,57 @@
-# app.py (Space ENCODER)
-
 import torch
 import gradio as gr
-from diffusers import Flux2Pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import tempfile, os
 
-REPO_ID = "black-forest-labs/FLUX.2-klein-4B"
+# Le vrai text encoder utilisé par FLUX.2 Klein 4B
+REPO_ID = "black-forest-labs/FLUX.2-klein-4B/text_encoder"
+
 device = "cpu"
 dtype = torch.float32
 
-# On NE CHARGE QUE la partie texte (comme tu voulais)
-pipe = Flux2Pipeline.from_pretrained(
+# Charger UNIQUEMENT le CausalLM Qwen utilisé par FLUX2 Klein
+tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
+text_encoder = AutoModelForCausalLM.from_pretrained(
     REPO_ID,
-    transformer=None,
-    vae=None,
-    scheduler=None,
     torch_dtype=dtype,
     low_cpu_mem_usage=True,
 )
-
-# On s'assure de ne garder que ce qui sert à l'encodage texte
-pipe.transformer = None
-pipe.vae = None
-pipe.scheduler = None
-
-pipe.to(device)
-
+text_encoder.to(device)
+text_encoder.eval()
 
 @torch.no_grad()
 def encode_text(prompt: str):
     if not prompt.strip():
         raise gr.Error("Prompt vide")
 
-    # encode_prompt renvoie EXACTEMENT ce que FLUX2 veut
-    prompt_embeds, pooled_prompt_embeds, text_ids = pipe.encode_prompt(
-        prompt=prompt,              # IMPORTANT : string simple, pas dict
-        device=device,
-        num_images_per_prompt=1,
-    )
+    # Tokenisation simple (pas de chat template)
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        truncation=True,
+        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
 
-    # On sauvegarde les trois tensors ensemble
-    data = {
-        "prompt_embeds": prompt_embeds.cpu(),
-        "pooled_prompt_embeds": pooled_prompt_embeds.cpu(),
-        "text_ids": text_ids.cpu(),
-    }
+    # Projection FLUX2 : 4096 → 7680
+    projected = text_encoder.model.project_out(hidden)
 
+    # Sauvegarde
     fd, path = tempfile.mkstemp(suffix=".pt")
     os.close(fd)
-    torch.save(data, path)
+    torch.save(projected.cpu(), path)
 
     return path
 
-
 demo = gr.Interface(
     fn=encode_text,
     inputs=gr.Textbox(label="Prompt"),
     outputs=gr.File(label="Embeddings FLUX2 (.pt)"),
-    title="FLUX.2 Klein — Text Encoder Officiel",
-    description="Renvoie les embeddings EXACTS que FLUX2 Klein attend.",
+    title="FLUX.2 Klein — Text Encoder Qwen3 Direct",
+    description="Encode le texte avec Qwen3 + projection FLUX2 (4096→7680).",
 )
 
-# api_name par défaut = "/predict", si tu veux explicitement :
-# demo.launch(api_name="/encode_text")
 demo.launch()