backup

app.py

@@ -28,35 +28,19 @@ check_environment()
 login(token=os.environ["HF_TOKEN"], add_to_git_credential=True)
 
 # Load model and processor (do this outside the inference function to avoid reloading)
-base_model_path = "taesiri/BugsBunny-LLama-3.2-11B-Vision-Instruct-Medium-FullModel"
+base_model_path = (
+    "taesiri/BugsBunny-LLama-3.2-11B-Vision-BaseCaptioner-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",
-# )
-
-from transformers import AutoModelForCausalLM, AutoProcessor, LlamaTokenizer
-import torch
-
-model_path = "taesiri/BugsBunny-LLama-3.2-11B-Vision-Instruct-Medium-FullModel"
-
-# Load the processor (handles both text and vision inputs)
-processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
-
-# Load the model
-model = AutoModelForCausalLM.from_pretrained(
-    model_path, torch_dtype=torch.bfloat16, device_map="cuda", trust_remote_code=True
+processor = AutoProcessor.from_pretrained(base_model_path)
+model = MllamaForConditionalGeneration.from_pretrained(
+    base_model_path,
+    torch_dtype=torch.bfloat16,
+    device_map="cuda",
 )
-
-# If you specifically need the tokenizer
-tokenizer = LlamaTokenizer.from_pretrained(model_path)
-
-model.tie_weights()
-
 # model = PeftModel.from_pretrained(model, lora_weights_path)
+model.tie_weights()
 
 
 def describe_image_in_JSON(json_string):