Update app.py

app.py

@@ -1,143 +1,55 @@
-import gc
-import os
-import io
-import time
-import tempfile
-import logging
-import spaces
-
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-from dd import encode_prompt   # <-- version Qwen compatible
-
-# ------------------------------------------------------
-# Logging
-# ------------------------------------------------------
-logging.basicConfig(
-    level=os.getenv("LOG_LEVEL", "INFO"),
-    format="%(asctime)s [%(levelname)s] %(name)s - %(message)s",
-)
-logger = logging.getLogger("qwen-text-encoding-gradio")
-
-# ------------------------------------------------------
-# Config
-# ------------------------------------------------------
-TEXT_ENCODER_ID = "Qwen/Qwen2.5-0.5B-Instruct"
-TOKENIZER_ID = "Qwen/Qwen2.5-0.5B-Instruct"
-DTYPE = torch.float32   # CPU safe
-
-# ------------------------------------------------------
-# Load model
-# ------------------------------------------------------
-logger.info("Loading Qwen 0.5B...")
-
-t0 = time.time()
-text_encoder = AutoModelForCausalLM.from_pretrained(
-    TEXT_ENCODER_ID,
-    torch_dtype=DTYPE,
-    low_cpu_mem_usage=True,
-).to("cpu")
-
-t1 = time.time()
-tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_ID)
-
-logger.info("Loaded Qwen in %.2fs", time.time() - t0)
-torch.set_grad_enabled(False)
-
-
-def get_vram_info():
-    if torch.cuda.is_available():
-        return {
-            "vram_allocated_mb": round(torch.cuda.memory_allocated() / 1024 / 1024, 2),
-            "vram_reserved_mb": round(torch.cuda.memory_reserved() / 1024 / 1024, 2),
-            "vram_max_allocated_mb": round(torch.cuda.max_memory_allocated() / 1024 / 1024, 2),
-        }
-    return {"vram": "CUDA not available"}
-
-
-@spaces.GPU()
-def encode_text(prompt: str):
-    """Encode text and return a downloadable pytorch file."""
-    global text_encoder, tokenizer
-
-    if text_encoder is None or tokenizer is None:
-        return None, "Model not loaded"
-
-    t0 = time.time()
-
-    # Multi‑prompts support
-    prompts = [p.strip() for p in prompt.strip().split("\n") if p.strip()]
-    if not prompts:
-        return None, "Please enter at least one prompt"
-
-    num_prompts = len(prompts)
-    prompt_input = prompts[0] if num_prompts == 1 else prompts
-
-    logger.info("Encoding %d prompt(s)", num_prompts)
-
-    # 🔥 Qwen‑compatible encode_prompt (dd.py)
-    prompt_embeds, text_ids = encode_prompt(
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        prompt=prompt_input,
+from diffusers import Flux2Pipeline
+
+# -----------------------------
+#  ENCODEUR FLUX MINIMAL
+# -----------------------------
+def encode_prompt(prompt: str):
+    # On charge uniquement le tokenizer + text_encoder
+    pipe = Flux2Pipeline.from_pretrained(
+        "black-forest-labs/FLUX.2-klein-4B",
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True,
+
+        # 🔥 On désactive tout ce qui n'est PAS utilisé par encode_prompt()
+        transformer=None,        # pas besoin du modèle de diffusion
+        vae=None,                # pas besoin pour encoder du texte
+        scheduler=None,          # inutile
+        safety_checker=None,     # inutile
+        feature_extractor=None,  # inutile
     )
 
-    duration = (time.time() - t0) * 1000.0
-    logger.info(
-        "Encoded in %.2f ms | prompt_embeds.shape=%s | text_ids.shape=%s",
-        duration,
-        tuple(prompt_embeds.shape),
-        tuple(text_ids.shape),
-    )
-
-    # Save tensor
-    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".pt")
-    torch.save(prompt_embeds.cpu(), temp_file.name)
-
-    # Cleanup
-    del prompt_embeds, text_ids
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
+    with torch.inference_mode():
+        prompt_embeds = pipe.encode_prompt(prompt)
 
-    vram = get_vram_info()
-    status = (
-        f"Encoded {num_prompts} prompt(s) in {duration:.2f}ms\n"
-        f"VRAM: {vram.get('vram_allocated_mb', 'N/A')} MB allocated, "
-        f"{vram.get('vram_max_allocated_mb', 'N/A')} MB peak"
-    )
+    return prompt_embeds
 
-    return temp_file.name, status
 
+# -----------------------------
+#  FONCTION GRADIO
+# -----------------------------
+def run_encoder(prompt):
+    embeds = encode_prompt(prompt)
 
-# ------------------------------------------------------
-# Gradio Interface
-# ------------------------------------------------------
-with gr.Blocks(title="Qwen Text Encoder") as demo:
-    gr.Markdown("# Qwen 2.5 Text Encoder")
-    gr.Markdown("Enter text to encode. For multiple prompts, put each on a new line.")
+    # On renvoie la forme pour vérifier
+    shape = tuple(embeds.shape)
 
-    with gr.Row():
-        with gr.Column():
-            prompt_input = gr.Textbox(
-                label="Prompt(s)",
-                placeholder="Enter your prompt here...\nOr multiple prompts, one per line",
-                lines=5,
-            )
-            encode_btn = gr.Button("Encode", variant="primary")
+    # On sauvegarde pour un autre Space si besoin
+    torch.save(embeds.cpu(), "embeds.pt")
 
-        with gr.Column():
-            output_file = gr.File(label="Download Embeddings (.pt)")
-            status_output = gr.Textbox(label="Status", interactive=False)
+    return f"Embeddings générés : shape = {shape}"
 
-    encode_btn.click(
-        fn=encode_text,
-        inputs=[prompt_input],
-        outputs=[output_file, status_output],
-    )
 
+# -----------------------------
+#  INTERFACE GRADIO
+# -----------------------------
+demo = gr.Interface(
+    fn=run_encoder,
+    inputs=gr.Textbox(label="Prompt", placeholder="Écris ton texte ici..."),
+    outputs=gr.Textbox(label="Résultat"),
+    title="Encodeur Texte FLUX.2 — Version Minimaliste",
+    description="Génère les embeddings texte officiels de FLUX.2 sans charger le VAE ou le modèle de diffusion.",
+)
 
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch()