V0.7

app.py

@@ -10,7 +10,7 @@ import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-from dd import encode_prompt
+from dd import encode_prompt   # <-- version Qwen compatible
 
 # ------------------------------------------------------
 # Logging
@@ -19,35 +19,35 @@ logging.basicConfig(
     level=os.getenv("LOG_LEVEL", "INFO"),
     format="%(asctime)s [%(levelname)s] %(name)s - %(message)s",
 )
-logger = logging.getLogger("mistral-text-encoding-gradio")
+logger = logging.getLogger("qwen-text-encoding-gradio")
 
 # ------------------------------------------------------
 # Config
 # ------------------------------------------------------
-TEXT_ENCODER_ID="Qwen/Qwen2.5-0.5B-Instruct"
+TEXT_ENCODER_ID = "Qwen/Qwen2.5-0.5B-Instruct"
 TOKENIZER_ID = "Qwen/Qwen2.5-0.5B-Instruct"
-DTYPE = torch.bfloat16
+DTYPE = torch.float32   # CPU safe
 
 # ------------------------------------------------------
-# Global model references
+# Load model
 # ------------------------------------------------------
-logger.info("Loading models...")
+logger.info("Loading Qwen 0.5B...")
 
 t0 = time.time()
 text_encoder = AutoModelForCausalLM.from_pretrained(
     TEXT_ENCODER_ID,
-    dtype=DTYPE,
+    torch_dtype=DTYPE,
+    low_cpu_mem_usage=True,
 ).to("cpu")
 
 t1 = time.time()
 tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_ID)
-logger.info("Loaded tokenizer in %.2fs", time.time() - t1)
 
+logger.info("Loaded Qwen in %.2fs", time.time() - t0)
 torch.set_grad_enabled(False)
 
 
 def get_vram_info():
-    """Get current VRAM usage info."""
     if torch.cuda.is_available():
         return {
             "vram_allocated_mb": round(torch.cuda.memory_allocated() / 1024 / 1024, 2),
@@ -56,6 +56,7 @@ def get_vram_info():
         }
     return {"vram": "CUDA not available"}
 
+
 @spaces.GPU()
 def encode_text(prompt: str):
     """Encode text and return a downloadable pytorch file."""
@@ -66,7 +67,7 @@ def encode_text(prompt: str):
 
     t0 = time.time()
 
-    # Handle multiple prompts (one per line)
+    # 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"
@@ -76,6 +77,7 @@ def encode_text(prompt: str):
 
     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,
@@ -90,11 +92,11 @@ def encode_text(prompt: str):
         tuple(text_ids.shape),
     )
 
-    # Save tensor to a temporary file
+    # Save tensor
     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".pt")
     torch.save(prompt_embeds.cpu(), temp_file.name)
 
-    # Clean up GPU tensors
+    # Cleanup
     del prompt_embeds, text_ids
     gc.collect()
     if torch.cuda.is_available():
@@ -108,13 +110,13 @@ def encode_text(prompt: str):
     )
 
     return temp_file.name, status
-    
+
 
 # ------------------------------------------------------
 # Gradio Interface
 # ------------------------------------------------------
-with gr.Blocks(title="Mistral Text Encoder") as demo:
-    gr.Markdown("# Mistral Text Encoder")
+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.")
 
     with gr.Row():