70 lines
2.6 KiB
Markdown
70 lines
2.6 KiB
Markdown
# Documentation
|
|
- Class name: MiDaS_Model_Loader
|
|
- Category: WAS Suite/Loaders
|
|
- Output node: False
|
|
- Repo Ref: https://github.com/WASasquatch/was-node-suite-comfyui
|
|
|
|
MiDaS_Model_Loader节点负责加载和准备MiDaS模型以供部署。它确保必要的依赖项已安装，并处理下载和加载指定的MiDaS模型变体。
|
|
|
|
# Input types
|
|
## Required
|
|
- midas_model
|
|
- 'midas_model'参数指定要加载的MiDaS模型类型。它对于确定用于深度估计任务的预训练模型权重和配置至关重要。
|
|
- Comfy dtype: str
|
|
- Python dtype: str
|
|
|
|
# Output types
|
|
- midas_model
|
|
- 'midas_model'输出提供了一个包含已加载的MiDaS模型及其相关转换函数的元组。此输出对于后续处理和分析深度信息至关重要。
|
|
- Comfy dtype: Tuple[torch.nn.Module, Callable]
|
|
- Python dtype: Tuple[torch.nn.Module, Callable[[torch.Tensor], torch.Tensor]]
|
|
|
|
# Usage tips
|
|
- Infra type: GPU
|
|
|
|
# Source code
|
|
```
|
|
class MiDaS_Model_Loader:
|
|
|
|
def __init__(self):
|
|
self.midas_dir = os.path.join(MODELS_DIR, 'midas')
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {'required': {'midas_model': (['DPT_Large', 'DPT_Hybrid'],)}}
|
|
RETURN_TYPES = ('MIDAS_MODEL',)
|
|
RETURN_NAMES = ('midas_model',)
|
|
FUNCTION = 'load_midas_model'
|
|
CATEGORY = 'WAS Suite/Loaders'
|
|
|
|
def load_midas_model(self, midas_model):
|
|
global MIDAS_INSTALLED
|
|
if not MIDAS_INSTALLED:
|
|
self.install_midas()
|
|
if midas_model == 'DPT_Large':
|
|
model_name = 'dpt_large_384.pt'
|
|
elif midas_model == 'DPT_Hybrid':
|
|
model_name = 'dpt_hybrid_384.pt'
|
|
else:
|
|
model_name = 'dpt_large_384.pt'
|
|
model_path = os.path.join(self.midas_dir, 'checkpoints' + os.sep + model_name)
|
|
torch.hub.set_dir(self.midas_dir)
|
|
if os.path.exists(model_path):
|
|
cstr(f'Loading MiDaS Model from `{model_path}`').msg.print()
|
|
midas_type = model_path
|
|
else:
|
|
cstr('Downloading and loading MiDaS Model...').msg.print()
|
|
midas = torch.hub.load('intel-isl/MiDaS', midas_model, trust_repo=True)
|
|
device = torch.device('cpu')
|
|
cstr(f'MiDaS is using passive device `{device}` until in use.').msg.print()
|
|
midas.to(device)
|
|
midas_transforms = torch.hub.load('intel-isl/MiDaS', 'transforms')
|
|
transform = midas_transforms.dpt_transform
|
|
return ((midas, transform),)
|
|
|
|
def install_midas(self):
|
|
global MIDAS_INSTALLED
|
|
if 'timm' not in packages():
|
|
install_package('timm')
|
|
MIDAS_INSTALLED = True
|
|
``` |