Update app.py

app.py

@@ -1,6 +1,6 @@
 import torch
 from PIL import Image
-from transformers import AutoProcessor, AutoModelForCausalLM
+from transformers import AutoProcessor, AutoModelForPreTraining
 import gradio as gr
 import json
 import traceback
@@ -10,7 +10,7 @@ model_name = "meta-llama/Llama-3.2-11B-Vision-Instruct"
 token = os.getenv("HUGGINGFACE_TOKEN").strip()
 
 processor = AutoProcessor.from_pretrained(model_name, token=token)
-model = AutoModelForCausalLM.from_pretrained(
+model = AutoModelForPreTraining.from_pretrained(
     model_name,
     quantization_config={"load_in_4bit": True},
     token=token
@@ -36,15 +36,15 @@ def analyze_image(image, prompt):
             return_tensors="pt"
         ).to(model.device)
         
-        # Separate inputs for generate method
-        generate_inputs = {
-            k: v for k, v in inputs.items() 
-            if k not in ['pixel_values', 'aspect_ratio_ids', 'aspect_ratio_mask']
-        }
+        # # Separate inputs for generate method
+        # generate_inputs = {
+        #     k: v for k, v in inputs.items() 
+        #     if k not in ['pixel_values', 'aspect_ratio_ids', 'aspect_ratio_mask']
+        # }
         
         with torch.no_grad():
             output = model.generate(**generate_inputs, max_new_tokens=100)
-        
+        print(processor.decode(output[0]))
         result = processor.decode(output[0], skip_special_tokens=True)
         
         try: