huggingface / huggingface/controlnet_aux

load_processor is not using local_files_only flag

Open
#96 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
497
Forks
104
PR merge metrics
No merged PRs in 30d

Description

### Describe the bug

There is no way to pass the local_files_only flag to the controlnet Processor `load_processor` function (see `controlnet_aux/processor.py`)

**Code:**

```python
def load_processor(self, processor_id: str) -> 'Processor':
"""Load controlnet aux processors

Args:
processor_id (str): processor name

Returns:
Processor: controlnet aux processor
"""
processor = MODELS[processor_id]['class']

# check if the proecssor is a checkpoint model
if MODELS[processor_id]['checkpoint']:
processor = processor.from_pretrained("lllyasviel/Annotators")
else:
processor = processor()
return processor
```

That function in-turn calls things such as `LeresDetector.from_pretrained` which looks like this

```python
class LeresDetector:
@classmethod
def from_pretrained(cls, pretrained_model_or_path, filename=None, pix2pix_filename=None, cache_dir=None):
filename = filename or "res101.pth"
pix2pix_filename = pix2pix_filename or "latest_net_G.pth"

if os.path.isdir(pretrained_model_or_path):
model_path = os.path.join(pretrained_model_or_path, filename)
else:
model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)

checkpoint = torch.load(model_path, map_location=torch.device('cpu'))

model = RelDepthModel(backbone='resnext101')
model.load_state_dict(strip_prefix_if_present(checkpoint['depth_model'], "module."), strict=True)
del checkpoint

if os.path.isdir(pretrained_model_or_path):
model_path = os.path.join(pretrained_model_or_path, pix2pix_filename)
else:
model_path = hf_hub_download(pretrained_model_or_path, pix2pix_filename, cache_dir=cache_dir)

opt = TestOptions().parse()
if not torch.cuda.is_available():
opt.gpu_ids = [] # cpu mode
pix2pixmodel = Pix2Pix4DepthModel(opt)
pix2pixmodel.save_dir = os.path.dirname(model_path)
pix2pixmodel.load_networks('latest')
pix2pixmodel.eval()

return cls(model, pix2pixmodel)
```

`LeresDetector.from_pretrained` in-turn calls `hf_hub_download` in order to download the model if the path is not a folder. `hf_hub_download` takes a `local_files_only` flag, however it is not being passed here.

Because of this, the hub attempts to access huggingface.co which causes the application to hang when offline.

## Proposed fix

Add a local_files_only flag to the Processor and each controlnet class

```python
def load_processor(self, processor_id: str, local_files_only: bool = False) -> 'Processor':
"""Load controlnet aux processors

Args:
processor_id (str): processor name

Returns:
Processor: controlnet aux processor
"""
processor = MODELS[processor_id]['class']

# check if the proecssor is a checkpoint model
if MODELS[processor_id]['checkpoint']:
processor = processor.from_pretrained("lllyasviel/Annotators", local_files_only=local_files_only)
else:
processor = processor()
return processor
```

```python
class LeresDetector:
@classmethod
def from_pretrained(cls, pretrained_model_or_path, filename=None, pix2pix_filename=None, cache_dir=None, local_files_only: bool = False):
filename = filename or "res101.pth"
pix2pix_filename = pix2pix_filename or "latest_net_G.pth"

if os.path.isdir(pretrained_model_or_path):
model_path = os.path.join(pretrained_model_or_path, filename)
else:
model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir, local_files_only=local_files_only)

checkpoint = torch.load(model_path, map_location=torch.device('cpu'))

model = RelDepthModel(backbone='resnext101')
model.load_state_dict(strip_prefix_if_present(checkpoint['depth_model'], "module."), strict=True)
del checkpoint

if os.path.isdir(pretrained_model_or_path):
model_path = os.path.join(pretrained_model_or_path, pix2pix_filename)
else:
model_path = hf_hub_download(pretrained_model_or_path, pix2pix_filename, cache_dir=cache_dir, local_files_only=local_files_only)

opt = TestOptions().parse()
if not torch.cuda.is_available():
opt.gpu_ids = [] # cpu mode
pix2pixmodel = Pix2Pix4DepthModel(opt)
pix2pixmodel.save_dir = os.path.dirname(model_path)
pix2pixmodel.load_networks('latest')
pix2pixmodel.eval()

return cls(model, pix2pixmodel)
```

### Reproduction

Attempt to initialize a processor using the `load_processor` function without an internet connection and see the script hang.

### Logs

_No response_

### System Info

controlnet_aux >=0.0.7

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.