Update app.py

app.py

@@ -41,10 +41,11 @@ class Phi3WithImage(torch.nn.Module):
         super().__init__()
         self.phi3 = AutoModelForCausalLM.from_pretrained(
             phi3_model_name,
-            torch_dtype=torch.bfloat16,
+            #torch_dtype=torch.bfloat16,
+            torch_dtype=torch.float32,
             device_map="auto",
             trust_remote_code=True,
-            quantization_config=bnb_config
+            #quantization_config=bnb_config
         )
         self.image_encoder = image_encoder
         self.image_embed_dim = image_embed_dim
@@ -114,12 +115,20 @@ image_transform = transforms.Compose([
 image_encoder = SigLIPImageEncoder(model_name=image_model_name, embed_dim=image_embed_dim, pretrained_path=siglip_pretrained_path).to(device)
 #model = Phi3WithImage(phi3_model_name, lora_model_path, image_encoder, image_embed_dim, bnb_config=bnb_config).to(device)
 model = Phi3WithImage(phi3_model_name,  image_encoder, image_embed_dim, bnb_config=bnb_config).to(device)
-model.phi3 = prepare_model_for_kbit_training(model.phi3)
+#model.phi3 = prepare_model_for_kbit_training(model.phi3)
 
 # Load LoRA model
 model.phi3 = PeftModel.from_pretrained(model.phi3, lora_model_path, device_map="auto", offload_dir='./offload')
 model.phi3 = model.phi3.merge_and_unload()
+model.phi3.save_pretrained("merged_model_fp16")
+
 #model.phi3 = PeftModel.from_pretrained(model.phi3, lora_model_path, device_map="auto")
+model.phi3 = AutoModelForCausalLM.from_pretrained(
+    "merged_model_fp16",
+    quantization_config=bnb_config,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
 model.eval() # Set to evaluation mode
 
 # 3. Inference Function