Prevent duplicate downloading of models

Author: jhc13

taggui/auto_captioning/auto_captioning_model.py

@@ -37,6 +37,11 @@ def replace_template_variables(text: str, image: Image) -> str:
 
 class AutoCaptioningModel:
     dtype = torch.float16
+    # When loading a model, if the `use_safetensors` argument is not set and
+    # both a safetensors and a non-safetensors version of the model are
+    # available, both versions get downloaded. This should be set to `None` for
+    # models that do not have a safetensors version.
+    use_safetensors = True
     model_load_context_manager = nullcontext()
     transformers_model_class = AutoModelForVision2Seq
     image_mode = 'RGB'
@@ -90,7 +95,8 @@ def get_processor(self):
                                              trust_remote_code=True)
 
     def get_model_load_arguments(self) -> dict:
-        arguments = {'device_map': self.device, 'trust_remote_code': True}
+        arguments = {'device_map': self.device, 'trust_remote_code': True,
+                     'use_safetensors': self.use_safetensors}
         if self.load_in_4_bit:
             quantization_config = BitsAndBytesConfig(
                 load_in_4bit=True,

taggui/auto_captioning/models/florence_2.py

@@ -5,6 +5,7 @@
 
 
 class Florence2(AutoCaptioningModel):
+    use_safetensors = None
     transformers_model_class = AutoModelForCausalLM
     task_prompts = [
         '<CAPTION>',
@@ -30,6 +31,7 @@ def get_default_prompt(self) -> str:
 
 
 class Florence2Promptgen(Florence2):
+    use_safetensors = True
     task_prompts = [
         '<GENERATE_PROMPT>',
         '<CAPTION>',