fix

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from PIL import Image, ImageDraw
-from transformers import pipeline, AutoTokenizer
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from optimum.onnxruntime import ORTModelForCausalLM
 import torch
 
@@ -74,11 +74,8 @@ def get_llm(model_name, preferred_file: str | None = None):
 
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_repo)
     # Ensure pad token exists (common for decoder-only models)
-    if tokenizer.pad_token_id is None:
-        try:
-            tokenizer.pad_token = tokenizer.eos_token
-        except Exception:
-            pass
+    if tokenizer.pad_token_id is None and getattr(tokenizer, "eos_token_id", None) is not None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
 
     # Try a few common ONNX filenames found in community repos to avoid the
     # "Too many ONNX model files were found" ambiguity.
@@ -135,6 +132,7 @@ def get_llm(model_name, preferred_file: str | None = None):
     llm_cache[cache_key] = (model, tokenizer)
     return model, tokenizer
 
+
 def update_user_prompt(detected_objects, current_prompt):
     if "No objects detected" in detected_objects:
         return current_prompt
@@ -149,8 +147,7 @@ def update_user_prompt(detected_objects, current_prompt):
 def generate_text(
     model_name,
     onnx_file_choice,
-    system_prompt,
-    user_prompt,
+    messages,
     do_sample,
     temperature,
     top_p,
@@ -160,10 +157,6 @@ def generate_text(
 ):
     model, tokenizer = get_llm(model_name, preferred_file=None if onnx_file_choice == "auto" else onnx_file_choice)
 
-    messages = [
-        {"role": "system", "content": system_prompt},
-        {"role": "user", "content": user_prompt},
-    ]
 
     chat_template_kwargs = {
         "tokenize": False,
@@ -182,8 +175,6 @@ def generate_text(
     # Ensure attention_mask is present and pad_token is defined
     if "attention_mask" not in inputs:
         inputs = tokenizer([text], return_tensors="pt", padding=True)
-    if getattr(tokenizer, "pad_token_id", None) is None and getattr(tokenizer, "eos_token_id", None) is not None:
-        tokenizer.pad_token_id = tokenizer.eos_token_id
 
     gen_kwargs = {
         "max_new_tokens": int(max_new_tokens),
@@ -198,37 +189,16 @@ def generate_text(
         gen_kwargs["eos_token_id"] = tokenizer.eos_token_id
 
     with torch.inference_mode():
-        try:
-            gen_ids = model.generate(
-                **inputs,
-                **gen_kwargs,
-            )
-        except Exception as e:
-            msg = str(e)
-            # Retry with int8 if KV cache shape mismatch and user didn't pick int8
-            if (
-                "past_key_values" in msg or "INVALID_ARGUMENT" in msg
-            ) and onnx_file_choice != "model_int8.onnx":
-                # Reload as int8 and retry once
-                model, tokenizer = get_llm(model_name, preferred_file="model_int8.onnx")
-                gen_kwargs["use_cache"] = False
-                gen_ids = model.generate(
-                    **inputs,
-                    **gen_kwargs,
-                )
-                # Mark that we switched variant
-                switched_variant_note = "\n[Note] Switched to model_int8.onnx due to KV-cache shape mismatch."
-            else:
-                raise
-
-    # Decode only the newly generated tokens beyond the input length
+        gen_ids = model.generate(
+            **inputs,
+            **gen_kwargs,
+        )
+
     trimmed = [
         output_ids[len(input_ids):]
         for input_ids, output_ids in zip(inputs.input_ids, gen_ids)
     ]
     response = tokenizer.batch_decode(trimmed, skip_special_tokens=True)[0]
-    if 'switched_variant_note' in locals():
-        response = response + switched_variant_note
     return response
 
 def chat_respond(
@@ -258,38 +228,17 @@ def chat_respond(
             messages.append({"role": "assistant", "content": a})
     messages.append({"role": "user", "content": message})
 
-    # Generate using the same path as generate_text, but inline to avoid extra serialization
-    model, tokenizer = get_llm(model_name, preferred_file=None if onnx_file_choice == "auto" else onnx_file_choice)
-
-    chat_template_kwargs = {
-        "tokenize": False,
-        "add_generation_prompt": True,
-    }
-    if "qwen" in model_name.lower():
-        chat_template_kwargs["enable_thinking"] = False
-
-    text = tokenizer.apply_chat_template(messages, **chat_template_kwargs)
-    inputs = tokenizer([text], return_tensors="pt")
-
-    with torch.inference_mode():
-        gen_ids = model.generate(
-            **inputs,
-            max_new_tokens=int(max_new_tokens),
-            do_sample=bool(do_sample),
-            temperature=float(temperature),
-            top_p=float(top_p),
-            top_k=int(top_k),
-            repetition_penalty=float(repetition_penalty),
-        )
-
-    trimmed = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, gen_ids)]
-    try:
-        reply = tokenizer.batch_decode(trimmed, skip_special_tokens=True)[0]
-    except Exception as e:
-        # Gracefully surface decoding issues
-        reply = f"[Error] Failed to decode model output: {e}"
-    if 'switched_variant_note' in locals():
-        reply = reply + switched_variant_note
+    reply = generate_text(
+        model_name=model_name,
+        onnx_file_choice=onnx_file_choice,
+        messages=messages,
+        do_sample=do_sample,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        repetition_penalty=repetition_penalty,
+        max_new_tokens=max_new_tokens,
+    )
 
     new_history = (history or []) + [(message, reply)]
     return new_history, gr.update(value="")