diff --git a/setup.py b/setup.py index d011f43..9c5f1ac 100644 --- a/setup.py +++ b/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", diff --git a/src/controlnet_aux/processor.py b/src/controlnet_aux/processor.py index 12cb6b0..4ceb421 100644 --- a/src/controlnet_aux/processor.py +++ b/src/controlnet_aux/processor.py @@ -82,7 +82,7 @@ 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: @@ -90,6 +90,7 @@ def __init__(self, processor_id: str, params: Optional[Dict] = None) -> None: 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}") @@ -97,7 +98,7 @@ def __init__(self, processor_id: str, params: Optional[Dict] = None) -> None: 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] @@ -105,11 +106,12 @@ def __init__(self, processor_id: str, params: Optional[Dict] = None) -> None: 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