Update app.py

app.py

@@ -10,6 +10,8 @@ import torch
 from accelerate import init_empty_weights
 # Load environment variables
 from transformers import AutoProcessor, AutoModelForPreTraining
+from transformers import MllamaForConditionalGeneration, AutoProcessor
+
 # Configure Gemini API
 # genai.configure(api_key=os.getenv("gkey2"))
 
@@ -25,7 +27,7 @@ prompt="<|image|><|begin_of_text|>You are a helpful assistant. Please respond to
 
 # Initialize the Llama model
 # model = Ollama(model="llama3.2")
-model_id = "meta-llama/Llama-3.2-11B-Vision"
+
 
 # model = MllamaForConditionalGeneration.from_pretrained(
 #     model_id,
@@ -33,8 +35,11 @@ model_id = "meta-llama/Llama-3.2-11B-Vision"
 #     device_map="auto",
 # )
 # processor = AutoProcessor.from_pretrained(model_id)
-processor = AutoProcessor.from_pretrained("meta-llama/Llama-3.2-11B-Vision")
-model = AutoModelForPreTraining.from_pretrained("meta-llama/Llama-3.2-11B-Vision")
+
+
+model_id = "meta-llama/Llama-3.2-11B-Vision"
+model = MllamaForConditionalGeneration.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16)
+processor = AutoProcessor.from_pretrained(model_id)
 
 # Define function to get response from the model
 def get_gemin_response(input_text, img):