Update app.py

app.py

@@ -1,5 +1,17 @@
 # Copyright: Shayekh Bin Islam. KAIST, South Korea. 2026.
 
+try:
+    import spaces
+except ImportError:
+    class spaces:
+        @staticmethod
+        def GPU(*args, **kwargs):
+            def decorator(func):
+                return func
+            if len(args) == 1 and callable(args[0]) and not kwargs:
+                return args[0]
+            return decorator
+
 import gradio as gr
 import fitz  # PyMuPDF
 from PIL import Image
@@ -8,13 +20,12 @@ import json
 import base64
 import soundfile as sf
 import torch
-import spaces
 
 from supertonic import TTS
-from vllm import LLM, SamplingParams
+from transformers import AutoProcessor, AutoModelForImageTextToText
 
-llm = None
-sampling_params = None
+model = None
+processor = None
 tts = None
 voice_style = None
 
@@ -41,18 +52,26 @@ def get_base64_image(image):
 
 @spaces.GPU
 def extract_vocabulary(pdf_text, images, translit_lang, translit_format, target_lang):
-    """Use vLLM to extract vocabulary from text and images."""
-    global llm, sampling_params
+    """Use Transformers to extract vocabulary from text and images."""
+    global model, processor
     
     os.makedirs("log", exist_ok=True)
     
+    if pdf_text.strip() == "":
+        pdf_text = '''"No Text available, see provided images only."'''
+
+
+    non_english = ""
+    if translit_lang.upper() != "ENGLISH":
+        non_english = f" CRITICAL: You MUST use the native alphabet/script of {translit_lang.upper()}, do NOT use English letters unless requested."
+    
     prompt_text = f"""Extract 3 to 5 key Korean words or phrases from the following text and images.
 Return ONLY a valid JSON list of dictionaries, where each dictionary has four keys:
 - 'korean' (the Korean text)
-- 'transliteration' (the pronunciation transliterated into {translit_lang.upper()} script/characters, formatted as {translit_format}. CRITICAL: You MUST use the native alphabet/script of {translit_lang.upper()}, do NOT use English letters unless requested.)
+- 'transliteration' (the pronunciation transliterated into {translit_lang.upper()} script/characters, formatted as {translit_format}.{non_english})
 - 'translation' (the translation into {target_lang.upper()})
 - 'explanation' (a brief grammar or context note in {target_lang.upper()}).
-No markdown formatting, just raw JSON.
+No markdown formatting, just raw JSON with ```json and ``` markers.
 
 Text:
 {pdf_text[:1500]}
@@ -62,16 +81,19 @@ Text:
     with open("log/debug_vlm_prompt.txt", "w", encoding="utf-8") as f:
         f.write(prompt_text)
 
-    content = [{"type": "text", "text": prompt_text}]
+    content = []
+    pil_images = []
     
     for i, img in enumerate(images):
         # DEBUG: Log images
         img.save(f"log/debug_image_{i}.png", format="PNG")
+        pil_images.append(img)
         
         content.append({
-            "type": "image_url",
-            "image_url": {"url": get_base64_image(img)}
+            "type": "image",
         })
+        
+    content += [{"type": "text", "text": prompt_text}]
 
     messages = [
         {
@@ -81,26 +103,46 @@ Text:
     ]
 
     try:
-        outputs = llm.chat(messages=messages, sampling_params=sampling_params)
-        output_text = outputs[0].outputs[0].text
+        model.to("cuda")
+        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = processor(
+            text=[text],
+            images=pil_images if pil_images else None,
+            return_tensors="pt",
+            padding=True
+        ).to("cuda")
+
+        generated_ids = model.generate(
+            **inputs,
+            max_new_tokens=2048*4,
+            temperature=1.0,
+            top_p=0.95,
+            do_sample=True
+        )
+        
+        generated_ids = [
+            output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        output_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
         
         # DEBUG: Log raw output text
         with open("log/debug_vlm_output.txt", "w", encoding="utf-8") as f:
             f.write(output_text)
             
     except Exception as e:
-        print(f"Error during vLLM inference: {e}")
+        print(f"Error during Transformers inference: {e}")
         return []
 
     try:
-        clean_text = output_text.strip()
-        if clean_text.startswith("```json"):
-            clean_text = clean_text[7:]
-        if clean_text.startswith("```"):
-            clean_text = clean_text[3:]
-        if clean_text.endswith("```"):
-            clean_text = clean_text[:-3]
-        clean_text = clean_text.strip()
+        import re
+        # Extract JSON from markdown code fences or raw output
+        json_match = re.search(r'```(?:json)?\s*([\s\S]*?)```', output_text)
+        if json_match:
+            clean_text = json_match.group(1).strip()
+        else:
+            # Fallback: find first [ ... ] or { ... } block
+            json_match = re.search(r'(\[[\s\S]*\]|\{[\s\S]*\})', output_text)
+            clean_text = json_match.group(1).strip() if json_match else output_text.strip()
             
         data = json.loads(clean_text)
         if not isinstance(data, list):
@@ -458,26 +500,27 @@ def create_demo():
 
 @spaces.GPU
 def main():
-    print("Loading Qwen3.5-2B model via vLLM...")
-    llm = LLM(
-        model="Qwen/Qwen3.5-2B",
-        # model="Qwen/Qwen3.5-9B",
-        max_model_len=65536,  # Reduced from 262144 to fit on single GPU
-        tensor_parallel_size=1, # Kept at 1 since CUDA_VISIBLE_DEVICES=1
-        gpu_memory_utilization=0.5,
-        enable_prefix_caching=True,
-        trust_remote_code=True,
-        limit_mm_per_prompt={"image": 10} # Added image limits
-    )
-
-    sampling_params = SamplingParams(
-        temperature=1.0,
-        top_p=0.95,
-        top_k=20,
-        min_p=0.0,
-        presence_penalty=0.0,
-        repetition_penalty=1.0,
-        max_tokens=2048,
+    global model, processor, tts, voice_style
+    
+    
+    # model_id = "Qwen/Qwen3.5-9B"
+    model_id = "Qwen/Qwen3.5-2B"
+    
+    print(f"Loading {model_id} model via Transformers...")
+    
+    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+    
+    try:
+        with open("chat_template.jinja", "r", encoding="utf-8") as f:
+            processor.chat_template = f.read()
+    except Exception as e:
+        print("Could not load custom chat template:", e)
+        
+    model = AutoModelForImageTextToText.from_pretrained(
+        model_id, 
+        torch_dtype=torch.bfloat16, 
+        device_map="cpu", 
+        trust_remote_code=True
     )
 
     print("Loading Supertonic TTS...")
@@ -488,8 +531,8 @@ def main():
         voice_style = tts.get_voice_style(tts.voice_style_names[0])
         
     demo = create_demo()
-    demo.launch(server_name="0.0.0.0", server_port=7861)
-    
+    demo.launch(server_name="0.0.0.0", server_port=7865)
+
 
 if __name__ == "__main__":
     main()