Merge pull request #676 from roboflow/fix-lora

Fix lora
roboflow · Sep 25, 2024 · 2bcd35c · 2bcd35c
2 parents 8aa4c51 + 7d9f1e1
commit 2bcd35c
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/inference/models/transformers/transformers.py b/inference/models/transformers/transformers.py
@@ -4,6 +4,7 @@
 
 import numpy as np
 from peft import LoraConfig, get_peft_model
+from peft.peft_model import PeftModel
 from PIL import Image
 from transformers import AutoModel, AutoProcessor, PaliGemmaForConditionalGeneration
 
@@ -226,7 +227,11 @@ def initialize_model(self):
             cache_dir=cache_dir,
             token=token,
         ).to(self.dtype)
-        self.model = get_peft_model(self.base_model, lora_config).eval().to(self.dtype)
+        self.model = (
+            PeftModel.from_pretrained(self.base_model, self.cache_dir)
+            .eval()
+            .to(self.dtype)
+        )
 
         self.processor = self.processor_class.from_pretrained(
             self.cache_dir, revision=revision