backup

app.py

@@ -2,7 +2,7 @@ import os
 import gradio as gr
 import torch
 from PIL import Image
-from transformers import MllamaForConditionalGeneration, AutoProcessor
+from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor
 from peft import PeftModel
 from huggingface_hub import login
 import spaces
@@ -31,13 +31,18 @@ login(token=os.environ["HF_TOKEN"], add_to_git_credential=True)
 base_model_path = "taesiri/BugsBunny-LLama-3.2-11B-Vision-Instruct-Medium-FullModel"
 # lora_weights_path = "taesiri/BugsBunny-LLama-3.2-11B-Vision-Base-Medium-LoRA"
 
-processor = AutoProcessor.from_pretrained(base_model_path)
-model = MllamaForConditionalGeneration.from_pretrained(
-    base_model_path,
-    torch_dtype=torch.bfloat16,
-    device_map="cuda",
-)
-model.tie_weights()
+try:
+    processor = AutoProcessor.from_pretrained(base_model_path, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model_path,
+        torch_dtype=torch.bfloat16,
+        device_map="cuda",
+        trust_remote_code=True,
+    )
+    model.tie_weights()
+except Exception as e:
+    print(f"Error loading model: {e}")
+    raise
 
 # model = PeftModel.from_pretrained(model, lora_weights_path)