Update my_model/captioner/image_captioning.py

my_model/captioner/image_captioning.py

@@ -22,6 +22,7 @@ class ImageCaptioningModel:
         self.device_map = config.DEVICE_MAP
         self.torch_dtype = config.TORCH_DTYPE
         self.load_in_8bit = config.LOAD_IN_8BIT
+        self.load_in_4bit = config.LOAD_IN_4BIT
         self.low_cpu_mem_usage = config.LOW_CPU_MEM_USAGE
         self.skip_secial_tokens = config.SKIP_SPECIAL_TOKENS
 
@@ -31,6 +32,7 @@ class ImageCaptioningModel:
         if self.model_type == 'i_blip':
             self.processor = InstructBlipProcessor.from_pretrained(self.model_path,
                                                                    load_in_8bit=self.load_in_8bit,
+                                                                   load_in_4bit=self.load_in_4bit
                                                                    torch_dtype=self.torch_dtype,
                                                                    device_map=self.device_map
                                                                    )