Update app.py

app.py

@@ -4,20 +4,36 @@ from diffusers import Flux2Pipeline
 
 pipe = Flux2Pipeline.from_pretrained(
     "black-forest-labs/FLUX.2-klein-4B",
-    torch_dtype=torch.float32,
-    low_cpu_mem_usage=True,
-    vae =None, 
-    transformer=None, 
-    scheduler=None, 
-    
+    transformer=None,
+    vae=None,
+    scheduler=None,
+    safety_checker=None,
+    feature_extractor=None,
 )
 
+tokenizer = pipe.tokenizer
+text_encoder = pipe.text_encoder
+
 def encode_text(prompt: str):
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=512,
+    )
+
     with torch.inference_mode():
-        prompt_embeds = pipe.encode_prompt(prompt)  # [B, L, 2560] dans ton cas
+        outputs = text_encoder(
+            **inputs,
+            output_hidden_states=True,
+            use_cache=False,
+        )
+
+    embeds = outputs.hidden_states[-1]  # [B, L, 4096]
 
-    torch.save(prompt_embeds.cpu(), "embeds.pt")
-    return f"shape={tuple(prompt_embeds.shape)}", "embeds.pt"
+    torch.save(embeds, "embeds.pt")
+    return f"shape={tuple(embeds.shape)}", "embeds.pt"
 
 demo = gr.Interface(
     fn=encode_text,