Add application file

app.py

@@ -5,79 +5,69 @@ import onnxruntime as rt
 import os
 import spaces
 import torch
-# Use AutoImageProcessor for robust image preprocessing (ConvNeXt/ViT)
-from transformers import AutoImageProcessor 
-from scipy.special import softmax # Used for computing probabilities from model output
-import requests # Used for making HTTP API calls (Typhoon 2.5)
+from transformers import AutoImageProcessor, AutoModelForCausalLM, AutoTokenizer # NEW: LLM Imports
+from scipy.special import softmax 
 
-# 1. ONNX MODEL AND LLM CONFIGURATION
+# 1. MODEL CONFIGURATION AND LOADING
 # ----------------------------------------------------
+# 1.1 ONNX Model (Image Classifier) Configuration
 ONNX_MODEL_PATH = "model.onnx"
 CLASS_LABELS_FILE = "class_labels.txt"
-MODEL_ID = 'facebook/convnext-tiny-224' # The base model ID used for training
+MODEL_ID = 'facebook/convnext-tiny-224' 
 
-# **LLM CONFIGURATION (Actual API Call Setup)**
-MODEL_NAME = "scb10x/typhoon2.5-qwen3-4b" # Typhoon 2.5 Model ID
-# Fetch Hugging Face Token from Space Secrets
-HF_TOKEN = os.getenv("HF_TOKEN") 
-API_URL = f"https://router.huggingface.co/models/{MODEL_NAME}" # <--- แก้ไข URL เป็น router.huggingface.co
+# 1.2 LLM Configuration (Loaded Locally)
+LLM_MODEL_NAME = "scb10x/typhoon2.5-qwen3-4b" 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"LLM Device: {device}")
 
-
-# Load character classes from the file created during training
+# Load ONNX Runtime Session and LLM
 try:
-    with open(CLASS_LABELS_FILE, 'r', encoding='utf-8') as f:
-        CHARACTER_LABELS = [line.strip() for line in f.readlines()]
-except FileNotFoundError:
-    CHARACTER_LABELS = ['Luffy', 'Zoro', 'Nami', 'Sanji', 'Chopper', 'Franky', 'Brook', 'Usopp', 'Jinbei', 'Robin', 'Ace', 'Law', 'Shanks', 'Kurohige', 'Mihawk', 'Rayleigh']
-    print(f"⚠️ WARNING: {CLASS_LABELS_FILE} not found. Using default labels.")
-
-# Load ONNX Runtime Session
-try:
-    print(f"Attempting to load ONNX model from: {ONNX_MODEL_PATH}") # <-- NEW DEBUG LINE
+    # 1. Load ONNX Runtime (ConvNeXt)
+    if not os.path.exists(ONNX_MODEL_PATH):
+        raise FileNotFoundError(f"ONNX Model file not found at: {ONNX_MODEL_PATH}")
+        
+    print(f"Attempting to load ONNX model from: {ONNX_MODEL_PATH}")
     sess = rt.InferenceSession(ONNX_MODEL_PATH)
     onnx_input_name = sess.get_inputs()[0].name
     onnx_output_name = sess.get_outputs()[0].name
-    # Load Image Processor (essential for correct image preparation)
     processor = AutoImageProcessor.from_pretrained(MODEL_ID)
     print("ONNX model and Image Processor loaded successfully.")
+
+    # 2. Load LLM (Typhoon 2.5) Locally
+    print(f"Attempting to load LLM model: {LLM_MODEL_NAME} onto {device}...")
+    llm_tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_NAME, trust_remote_code=True)
+    llm_model = AutoModelForCausalLM.from_pretrained(
+        LLM_MODEL_NAME,
+        trust_remote_code=True,
+        torch_dtype=torch.float16, # Use float16 for efficiency
+        low_cpu_mem_usage=True,
+    )
+    llm_model.to(device)
+    print("Typhoon 2.5 LLM loaded successfully.")
+
 except Exception as e:
-    # <--- NEW DEBUG BLOCK: พิมพ์ Error จริงออกมา
-    print(f"FATAL ERROR LOADING ONNX MODEL: {e}")
-    print("Please ensure model.onnx is tracked by Git LFS and is uploaded correctly.")
+    print(f"FATAL ERROR LOADING MODELS: {e}")
+    print("Please ensure GPU is available and files are uploaded correctly (including model.onnx.data).")
     sess = None
+    llm_model = None
+    llm_tokenizer = None
 
-
-# 2. LLM API FUNCTION (Replaces Placeholder)
-# ----------------------------------------------------
-def query_typhoon_api(payload):
-    """Sends prompt to the Hugging Face Inference API of Typhoon 2.5."""
-    if not HF_TOKEN:
-        return "Error: HF_TOKEN is not set in Hugging Face Space secrets."
-    
-    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-    response = requests.post(API_URL, headers=headers, json=payload)
-    
-    # Check for non-successful status codes (e.g., 401 Unauthorized, 503 Service Unavailable)
-    if response.status_code != 200:
-        return f"Error {response.status_code}: API call failed. {response.text}"
-        
-    try:
-        # Extract the generated text from the response structure
-        result = response.json()[0]['generated_text']
-        # Remove the input prompt part from the output text
-        return result.split(payload['inputs'])[-1].strip()
-    except Exception as e:
-        return f"Error processing API response: {e}"
+# Load character classes 
+try:
+    with open(CLASS_LABELS_FILE, 'r', encoding='utf-8') as f:
+        CHARACTER_LABELS = [line.strip() for line in f.readlines()]
+except FileNotFoundError:
+    CHARACTER_LABELS = ['Luffy', 'Zoro', 'Nami', 'Sanji', 'Chopper', 'Franky', 'Brook', 'Usopp', 'Jinbei', 'Robin', 'Ace', 'Law', 'Shanks', 'Kurohige', 'Mihawk', 'Rayleigh']
 
 
-# 3. TYPHOON 2.5 LOGIC (Knowledge Base + Prompt Generation)
-# ---------------------------------------------------------
-# This dictionary serves as the LLM's knowledge base for the character's full name/role
+# 2. LLM GENERATION FUNCTION (Local Inference)
+# ----------------------------------------------------
+# ฐานข้อมูลข้อมูลเสริมตัวละคร
 CHARACTER_INFO = {
     "Ace": "โพโทกัส ดี เอส พี่ชายบุญธรรมของลูฟี่ ผู้ใช้พลังผลปีศาจเมระ เมระ",
     "Luffy": "มังกี้ ดี ลูฟี่ กัปตันกลุ่มโจรสลัดหมวกฟาง ผู้ใฝ่ฝันจะเป็นราชาโจรสลัด",
     "Zoro": "โรโรโนอา โซโล นักดาบสามเล่มแห่งกลุ่มหมวกฟาง ผู้มีเป้าหมายเป็นนักดาบอันดับหนึ่งของโลก",
-    "Nami": "นามิ นักเดินเรือสาวแห่งกลุ่มหมวกฟาง และเป็นนักทำแผนที่มือฉมัง",
+    "Nami": "นามิ นักเดินเรือสาวแห่งกลุ่มโจรสลัดหมวกฟาง และเป็นนักทำแผนที่มือฉมัง",
     "Sanji": "ซันจิ กุ๊กแห่งกลุ่มโจรสลัดหมวกฟาง และเป็นสุดยอดนักสู้ที่ใช้เท้าในการต่อสู้",
     "Chopper": "โทนี่ โทนี่ ช็อปเปอร์ หมอประจำเรือ ผู้มีใจรักเพื่อนและอ่อนไหวที่สุดในกลุ่ม",
     "Robin": "นิโค โรบิน นักโบราณคดี ผู้เดียวที่อ่านโพเนกลีฟได้",
@@ -92,11 +82,14 @@ CHARACTER_INFO = {
     "Rayleigh": "ซิลเวอร์ส เรย์ลี่ อดีตมือขวาของราชาโจรสลัด โกลด์ ดี. โรเจอร์",
 }
 
-
-def generate_thai_response(character_name, confidence):
+def generate_typhoon_response(character_name, confidence):
     """
-    Constructs a sophisticated prompt and queries the Typhoon 2.5 API.
+    ฟังก์ชัน LLM ที่ใช้ Local Inference ภายใน Space
     """
+    if llm_model is None:
+        return (f"❌ LLM ไม่พร้อมใช้งาน: ตัวละครคือ **{character_name}** "
+                f"[ความมั่นใจ: **{confidence*100:.2f}%**]")
+    
     info = CHARACTER_INFO.get(character_name, "ตัวละครวันพีซ")
     
     # 1. Build a clear, instructional prompt for the LLM
@@ -106,51 +99,54 @@ def generate_thai_response(character_name, confidence):
         f"กรุณาสร้างข้อความตอบกลับที่เป็นมิตรและเป็นภาษาไทย โดยขึ้นต้นด้วย 'ยืนยันผลการทำนาย!' "
         f"และรวมข้อมูลทั้งหมดนี้เข้าด้วยกันในประโยคเดียวโดยใช้ Markdown bold สำหรับชื่อตัวละครและความมั่นใจ (XX.XX%)."
     )
+
+    # 2. Generate text using the local LLM
+    messages = [{"role": "user", "content": prompt}]
+    input_ids = llm_tokenizer.apply_chat_template(
+        messages, add_generation_prompt=True, return_tensors="pt"
+    ).to(device)
+
+    output_ids = llm_model.generate(
+        input_ids,
+        max_new_tokens=100,
+        temperature=0.7,
+        do_sample=True,
+        pad_token_id=llm_tokenizer.eos_token_id,
+    )
     
-    # 2. Prepare Payload
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "max_new_tokens": 100,
-            "return_full_text": False
-        }
-    }
-    
-    # 3. Call the API and handle potential errors
-    llm_response = query_typhoon_api(payload)
-    
-    if llm_response.startswith("Error"):
-        # Fallback to a static, simple response if API fails
-        thai_name = info.split(' ')[0] 
-        return (f"⚠️ LLM API ไม่ตอบสนอง: ตัวละครคือ **{thai_name}** ({info}) "
-                f"[ความมั่นใจ: **{confidence*100:.2f}%**]")
+    # 3. Decode response
+    response = llm_tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    # Remove the input prompt from the response
+    response_text = response.split(prompt)[-1].strip()
+
+    return response_text
 
-    return llm_response
 
-# 4. ONNX INFERENCE FUNCTION
+# 3. ONNX INFERENCE FUNCTION
 # ----------------------------------------------------
-#@spaces.GPU
+# เราจะใช้ @spaces.GPU ตรงนี้เพื่อให้ LLM (ซึ่งอยู่ในฟังก์ชันที่ถูกเรียก) รันบน GPU ด้วย
+@spaces.GPU # ใช้ GPU สำหรับการรัน LLM
 def predict_one_piece_character(pil_image):
     if pil_image is None or sess is None:
         return "⚠️ โมเดลไม่พร้อมใช้งาน กรุณาตรวจสอบไฟล์ ONNX และการตั้งค่า"
     
     try:
-        # 4.1 Preprocessing (ConvNeXt standard input)
+        # 3.1 Preprocessing (ConvNeXt standard input)
         inputs = processor(images=pil_image, return_tensors="np")
         onnx_input = inputs['pixel_values'].astype(np.float32)
 
-        # 4.2 Run Inference
+        # 3.2 Run Inference (ConvNeXt ONNX)
         onnx_predictions = sess.run([onnx_output_name], {onnx_input_name: onnx_input})
         logits = onnx_predictions[0].squeeze() 
 
-        # 4.3 Post-processing (Softmax and Argmax)
+        # 3.3 Post-processing (Softmax and Argmax)
         probabilities = softmax(logits)
         predicted_index = np.argmax(probabilities)
         predicted_character = CHARACTER_LABELS[predicted_index]
         confidence = probabilities[predicted_index].item()
         
-        # 4.4 LLM Integration
-        final_response = generate_thai_response(predicted_character, confidence)
+        # 3.4 LLM Integration (Local Generation)
+        final_response = generate_typhoon_response(predicted_character, confidence)
         
         return final_response
 
@@ -158,14 +154,14 @@ def predict_one_piece_character(pil_image):
         print(f"RUNTIME ERROR: {e}") 
         return f"เกิดข้อผิดพลาดในการทำนาย: {e}"
 
-# 5. GRADIO INTERFACE
+# 4. GRADIO INTERFACE
 # ----------------------------------------------------
 interface = gr.Interface(
     fn=predict_one_piece_character, 
     inputs=gr.Image(type="pil", label="อัปโหลดรูปภาพตัวละครวันพีซ"), 
-    outputs=gr.Textbox(label="ผลการทำนายชื่อตัวละคร (Powered by Typhoon 2.5)"), 
-    title="🏴‍☠️ One Piece Classifier (ConvNeXt ONNX + Typhoon 2.5)",
-    description="อัปโหลดภาพตัวละครวันพีซ เพื่อให้ AI ทำนายชื่อตัวละครพร้อมสร้างข้อความตอบกลับที่น่าประทับใจ"
+    outputs=gr.Textbox(label="ผลการทำนายชื่อตัวละคร (Typhoon 2.5 Local)"), 
+    title="🏴‍☠️ One Piece Classifier (ConvNeXt ONNX + Typhoon 2.5 Local)",
+    description="แอปพลิเคชันจำแนกตัวละครวันพีซโดยรัน LLM ภายใน Space"
 )
 
 if __name__ == "__main__":