Update app.py

app.py

@@ -9,11 +9,9 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 torch_dtype = torch.float32
 
 # Load the fine-tuned base model
-base_model = AutoModelForCausalLM.from_pretrained('byh711/FLODA-deepfake', trust_remote_code=True, torch_dtype=torch_dtype).to(device)
+caption_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base-ft', trust_remote_code=True, revision='refs/pr/6', torch_dtype=torch_dtype).to(device)
+model = AutoModelForCausalLM.from_pretrained('byh711/FLODA-deepfake', trust_remote_code=True, torch_dtype=torch_dtype).to(device)
 processor = AutoProcessor.from_pretrained('byh711/FLODA-deepfake', trust_remote_code=True)
-
-# Load the LoRA weights
-model = PeftModel.from_pretrained(base_model, peft_model_path)
 model.eval()
 
 def caption_generate(task_prompt, text_input=None, image=None):
@@ -25,7 +23,7 @@ def caption_generate(task_prompt, text_input=None, image=None):
     else:
         prompt = task_prompt + text_input
     inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
-    generated_ids = model.generate(
+    generated_ids = caption_model.generate(
       input_ids=inputs["input_ids"],
       pixel_values=inputs["pixel_values"],
       max_new_tokens=1024,
@@ -52,7 +50,7 @@ def run_example(task_prompt, text_input=None, image=None):
     inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
     inputs = {k: v.to(torch_dtype) if v.is_floating_point() else v for k, v in inputs.items()}
 
-    generated_ids = base_model.generate(
+    generated_ids = model.generate(
         input_ids=inputs["input_ids"],
         pixel_values=inputs["pixel_values"],
         max_new_tokens=1024,