Closes #96 add local files only flag

setup.py

@@ -180,7 +180,7 @@ def run(self):
 
 setup(
     name="controlnet_aux",
-    version="0.0.6",  # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)
+    version="0.0.8",  # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)
     description="Auxillary models for controlnet",
     long_description=open("README.md", "r", encoding="utf-8").read(),
     long_description_content_type="text/markdown",

src/controlnet_aux/processor.py

@@ -82,34 +82,36 @@
 
 
 class Processor:
-    def __init__(self, processor_id: str, params: Optional[Dict] = None) -> None:
+    def __init__(self, processor_id: str, params: Optional[Dict] = None, local_files_only: bool = False) -> None:
         """Processor that can be used to process images with controlnet aux processors
 
         Args:
             processor_id (str): processor name, options are 'hed, midas, mlsd, openpose,
                                 pidinet, normalbae, lineart, lineart_coarse, lineart_anime,
                                 canny, content_shuffle, zoe, mediapipe_face
             params (Optional[Dict]): parameters for the processor
+            local_files_only (bool): whether to load only local files
         """
         LOGGER.info(f"Loading {processor_id}")
 
         if processor_id not in MODELS:
             raise ValueError(f"{processor_id} is not a valid processor id. Please make sure to choose one of {', '.join(MODELS.keys())}")
 
         self.processor_id = processor_id
-        self.processor = self.load_processor(self.processor_id)
+        self.processor = self.load_processor(self.processor_id, local_files_only=local_files_only)
 
         # load default params
         self.params = MODEL_PARAMS[self.processor_id]
         # update with user params
         if params:
             self.params.update(params)
 
-    def load_processor(self, processor_id: str) -> 'Processor':
+    def load_processor(self, processor_id: str, local_files_only: bool = False) -> 'Processor':
         """Load controlnet aux processors
 
         Args:
             processor_id (str): processor name
+            local_files_only (bool): whether to load only local files
 
         Returns:
             Processor: controlnet aux processor
@@ -118,7 +120,7 @@ def load_processor(self, processor_id: str) -> 'Processor':
 
         # check if the proecssor is a checkpoint model
         if MODELS[processor_id]['checkpoint']:
-            processor = processor.from_pretrained("lllyasviel/Annotators")
+            processor = processor.from_pretrained("lllyasviel/Annotators", local_files_only=local_files_only)
         else:
             processor = processor()
         return processor