Update app.py

app.py

@@ -1,22 +1,25 @@
 import torch
 import gradio as gr
 from transformers import AutoTokenizer, Qwen2ForCausalLM
-from diffusers import Flux2Pipeline
+from diffusers import Flux2Transformer2DModel
 
 device = "cpu"
 dtype = torch.float32
 
-# Charger FLUX pour récupérer les embedder internes
-pipe = Flux2Pipeline.from_pretrained(
+# Charger SEULEMENT le transformer FLUX (léger)
+transformer = Flux2Transformer2DModel.from_pretrained(
     "black-forest-labs/FLUX.2-klein-4B",
+    subfolder="transformer",
     torch_dtype=dtype,
-    text_encoder =None, 
-    vae=None, 
-    low_cpu_mem_usage=True,
 )
 
-pos_embedder = pipe.transformer.pos_embed
-extra_embedder = pipe.transformer.x_embedder
+# Extraire UNIQUEMENT les modules nécessaires
+pos_embedder = transformer.pos_embedder
+extra_embedder = transformer.x_embedder
+
+# Libérer le reste
+del transformer
+torch.cuda.empty_cache() if torch.cuda.is_available() else None
 
 # Charger Qwen (encodeur texte)
 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B")