Update app.py

app.py

@@ -4,36 +4,16 @@ from diffusers import Flux2Pipeline
 
 pipe = Flux2Pipeline.from_pretrained(
     "black-forest-labs/FLUX.2-klein-4B",
-    transformer=None,
-    vae=None,
-    scheduler=None,
-    safety_checker=None,
-    feature_extractor=None,
+    torch_dtype=torch.float32,
+    low_cpu_mem_usage=True,
 )
 
-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():
-        outputs = text_encoder(
-            **inputs,
-            output_hidden_states=True,
-            use_cache=False,
-        )
-
-    embeds = outputs.hidden_states[-1]  # [B, L, 4096]
+        prompt_embeds = pipe.encode_prompt(prompt)  # [B, L, 2560] dans ton cas
 
-    torch.save(embeds, "embeds.pt")
-    return f"shape={tuple(embeds.shape)}", "embeds.pt"
+    torch.save(prompt_embeds.cpu(), "embeds.pt")
+    return f"shape={tuple(prompt_embeds.shape)}", "embeds.pt"
 
 demo = gr.Interface(
     fn=encode_text,