Update app.py

app.py

@@ -1,13 +1,14 @@
 import gradio as gr
 import torch
 from PIL import Image
-from transformers import MllamaForConditionalGeneration, AutoProcessor
+from transformers import AutoModelForVision2Seq, AutoProcessor
 
 # โหลดโมเดลและตัวประมวลผล
 model_id = "0llheaven/Llama-3.2-11B-Vision-Radiology-mini"
 
-model = MllamaForConditionalGeneration.from_pretrained(
+model = AutoModelForVision2Seq.from_pretrained(
     model_id,
+    load_in_4bit = True,
     torch_dtype=torch.bfloat16,
     device_map="auto",
 )
@@ -35,20 +36,22 @@ def generate_caption(image):
     ).to(model.device)
 
     # สร้างข้อความตอบกลับจากโมเดล
-    output = model.generate(**inputs, max_new_tokens=256, use_cache=True, temperature=1.5, min_p=0.1)
-    return processor.decode(output[0])
+    output = model.generate(**inputs, 
+                            max_new_tokens=256, 
+                            use_cache=True, 
+                            temperature=1.5, 
+                            min_p=0.1)
+    
+    return processor.tokenizer.decode(output[0], skip_special_tokens=True).strip()
 
 # สร้างอินเตอร์เฟซด้วย Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("# Radiology Image Captioning")
-    with gr.Row():
-        image_input = gr.Image(type="pil", label="Upload Image")
-        output_text = gr.Textbox(label="Generated Caption")
-        title="Medical Vision Analysis"
-    generate_button = gr.Button("Generate Caption")
-    
-    # กำหนดการทำงานเมื่อกดปุ่ม
-    generate_button.click(fn=generate_caption, inputs=image_input, outputs=output_text)
+demo = gr.Interface(
+    fn=generate_caption,
+    image_input = gr.Image(type="pil", label="Upload Image"),
+    output_text = gr.Textbox(label="Generated Caption"),
+    title="Medical Vision Analysis"
+)
 
 # รันแอป
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()