Update app.py

app.py

@@ -6,16 +6,16 @@ from diffusers import Flux2Transformer2DModel
 device = "cpu"
 dtype = torch.float32
 
-# Charger SEULEMENT le transformer FLUX (léger)
+# Charger uniquement le transformer FLUX (léger)
 transformer = Flux2Transformer2DModel.from_pretrained(
     "black-forest-labs/FLUX.2-klein-4B",
     subfolder="transformer",
     torch_dtype=dtype,
 )
 
-# Extraire UNIQUEMENT les modules nécessaires
-pos_embedder = transformer.pos_embed
-extra_embedder = transformer.x_embedder
+# Modules internes
+pos_embed = transformer.pos_embed        # [1, 4096, 2560]
+x_embedder = transformer.x_embedder      # module → 2140 dims
 
 # Libérer le reste
 del transformer
@@ -45,9 +45,13 @@ def encode(prompt):
         )
 
     text = out.hidden_states[-1]  # [1, L, 2560]
+    L = text.shape[1]
 
-    pos = pos_embedder(text)      # [1, L, 2560]
-    extra = extra_embedder(text)  # [1, L, 2140]
+    # 🔥 pos_embed n'est PAS un module → on slice
+    pos = pos_embed[:, :L, :]     # [1, L, 2560]
+
+    # extra embedder est un module → on l'appelle
+    extra = x_embedder(text)      # [1, L, 2140]
 
     final = torch.cat([text, pos, extra], dim=-1)  # [1, L, 7260]