fix app.py

app.py

@@ -1,97 +1,89 @@
+import os
+from huggingface_hub import login
+from transformers import AutoProcessor, AutoModelForVision2Seq
 import torch
-from transformers import AutoModelForVision2Seq, AutoProcessor, BitsAndBytesConfig
 from PIL import Image
 import gradio as gr
 
-# Global variables for model and processor
+# Login to Hugging Face Hub
+if 'HUGGING_FACE_HUB_TOKEN' in os.environ:
+    print("Logging in to Hugging Face Hub...")
+    login(token=os.environ['HUGGING_FACE_HUB_TOKEN'])
+else:
+    print("Warning: HUGGING_FACE_HUB_TOKEN not found")
+
+# Global variables
 model = None
 processor = None
 
-def load_model_and_processor():
+def load_model():
     global model, processor
     try:
         model_path = "Aekanun/thai-handwriting-llm"
-        base_model_path = "meta-llama/Llama-3.2-11B-Vision-Instruct"
-        
-        print("Loading processor...")
-        processor = AutoProcessor.from_pretrained(base_model_path)
+        print(f"Loading model and processor from {model_path}...")
         
-        print("Loading model...")
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.bfloat16
-        )
+        processor = AutoProcessor.from_pretrained(model_path)
+        model = AutoModelForVision2Seq.from_pretrained(model_path)
         
-        model = AutoModelForVision2Seq.from_pretrained(
-            model_path,
-            device_map="auto",
-            torch_dtype=torch.bfloat16,
-            quantization_config=bnb_config
-        )
+        if torch.cuda.is_available():
+            model = model.to("cuda")
+            
         return True
     except Exception as e:
         print(f"Error loading model: {str(e)}")
         return False
 
-def process_handwriting(image):
-    global model, processor
-    
+def process_image(image):
     if image is None:
         return "กรุณาอัพโหลดรูปภาพ"
-    
+        
     try:
+        # Ensure image is in PIL format
         if not isinstance(image, Image.Image):
             image = Image.fromarray(image)
         
-        prompt = """Transcribe the Thai handwritten text from the provided image.
-Only return the transcription in Thai language."""
-        
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": prompt},
-                    {"type": "image", "image": image}
-                ],
-            }
-        ]
-        
-        text = processor.apply_chat_template(messages, tokenize=False)
-        inputs = processor(text=text, images=image, return_tensors="pt")
-        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        # Convert to RGB if needed
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+            
+        # Process image
+        inputs = processor(images=image, return_tensors="pt")
         
+        # Move to GPU if available
+        if torch.cuda.is_available():
+            inputs = {k: v.to("cuda") for k, v in inputs.items()}
+            
+        # Generate text
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=256,
-                do_sample=False,
-                pad_token_id=processor.tokenizer.pad_token_id
+                max_new_tokens=100,
+                num_beams=4,
+                pad_token_id=processor.tokenizer.pad_token_id,
+                eos_token_id=processor.tokenizer.eos_token_id
             )
-        
-        transcription = processor.decode(outputs[0], skip_special_tokens=True)
-        return transcription
+            
+        # Decode output
+        predicted_text = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+        return predicted_text.strip()
         
     except Exception as e:
         return f"เกิดข้อผิดพลาด: {str(e)}"
 
-# Initialize application
+# Initialize
 print("Initializing application...")
-model_loaded = load_model_and_processor()
-
-if model_loaded:
-    print("Creating Gradio interface...")
+if load_model():
+    # Create Gradio interface
     demo = gr.Interface(
-        fn=process_handwriting,
+        fn=process_image,
         inputs=gr.Image(type="pil", label="อัพโหลดรูปลายมือเขียนภาษาไทย"),
         outputs=gr.Textbox(label="ข้อความที่แปลงได้"),
-        title="Thai Handwriting to Text",
-        description="อัพโหลดรูปภาพลายมือเขียนภาษาไทยเพื่อแปลงเป็นข้อความ"
+        title="Thai Handwriting Recognition",
+        description="อัพโหลดรูปภาพลายมือเขียนภาษาไทยเพื่อแปลงเป็นข้อความ",
+        examples=[["example1.jpg"], ["example2.jpg"]]
     )
-    
+
     if __name__ == "__main__":
-        print("Launching application...")
         demo.launch()
 else:
-    print("Failed to load model and processor. Please check the logs.")
+    print("Failed to initialize the application")