75 lines
3.3 KiB
Markdown
75 lines
3.3 KiB
Markdown
# Documentation
|
||
- Class name: SAMLoader
|
||
- Category: ImpactPack
|
||
- Output node: False
|
||
- Repo Ref: https://github.com/ltdrdata/ComfyUI-Impact-Pack.git
|
||
|
||
SAMLoader节点旨在高效管理和加载SAM（Segment Anything Model）家族中的各种模型，包括ESAM。它提供了一个高级接口，用于根据用户定义的标准（如模型名称和设备模式）选择和使用模型，确保最佳性能和资源利用率。
|
||
|
||
# Input types
|
||
## Required
|
||
- model_name
|
||
- model_name参数对于确定要加载的SAM模型至关重要。它通过指导节点指向正确的模型文件来影响节点的执行，这对于后续的处理和预测任务至关重要。
|
||
- Comfy dtype: str
|
||
- Python dtype: str
|
||
## Optional
|
||
- device_mode
|
||
- device_mode参数允许用户指定模型执行的计算设备。它通过根据所选模式优化GPU或CPU资源的使用，在节点性能中扮演重要角色。
|
||
- Comfy dtype: str
|
||
- Python dtype: str
|
||
|
||
# Output types
|
||
- SAM_MODEL
|
||
- SAM_MODEL输出代表已加载的SAM模型，准备用于推理或进一步处理。它封装了模型的架构和学习到的参数，标志着节点主要功能的完成。
|
||
- Comfy dtype: torch.nn.Module
|
||
- Python dtype: torch.nn.Module
|
||
|
||
# Usage tips
|
||
- Infra type: GPU
|
||
|
||
# Source code
|
||
```
|
||
class SAMLoader:
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
models = [x for x in folder_paths.get_filename_list('sams') if 'hq' not in x]
|
||
return {'required': {'model_name': (models + ['ESAM'],), 'device_mode': (['AUTO', 'Prefer GPU', 'CPU'],)}}
|
||
RETURN_TYPES = ('SAM_MODEL',)
|
||
FUNCTION = 'load_model'
|
||
CATEGORY = 'ImpactPack'
|
||
|
||
def load_model(self, model_name, device_mode='auto'):
|
||
if model_name == 'ESAM':
|
||
if 'ESAM_ModelLoader_Zho' not in nodes.NODE_CLASS_MAPPINGS:
|
||
try_install_custom_node('https://github.com/ZHO-ZHO-ZHO/ComfyUI-YoloWorld-EfficientSAM', "To use 'ESAM' model, 'ComfyUI-YoloWorld-EfficientSAM' extension is required.")
|
||
raise Exception("'ComfyUI-YoloWorld-EfficientSAM' node isn't installed.")
|
||
esam_loader = nodes.NODE_CLASS_MAPPINGS['ESAM_ModelLoader_Zho']()
|
||
if device_mode == 'CPU':
|
||
esam = esam_loader.load_esam_model('CPU')[0]
|
||
else:
|
||
device_mode = 'CUDA'
|
||
esam = esam_loader.load_esam_model('CUDA')[0]
|
||
sam_obj = core.ESAMWrapper(esam, device_mode)
|
||
esam.sam_wrapper = sam_obj
|
||
print(f'Loads EfficientSAM model: (device:{device_mode})')
|
||
return (esam,)
|
||
modelname = folder_paths.get_full_path('sams', model_name)
|
||
if 'vit_h' in model_name:
|
||
model_kind = 'vit_h'
|
||
elif 'vit_l' in model_name:
|
||
model_kind = 'vit_l'
|
||
else:
|
||
model_kind = 'vit_b'
|
||
sam = sam_model_registry[model_kind](checkpoint=modelname)
|
||
size = os.path.getsize(modelname)
|
||
safe_to = core.SafeToGPU(size)
|
||
device = comfy.model_management.get_torch_device() if device_mode == 'Prefer GPU' else 'CPU'
|
||
if device_mode == 'Prefer GPU':
|
||
safe_to.to_device(sam, device)
|
||
is_auto_mode = device_mode == 'AUTO'
|
||
sam_obj = core.SAMWrapper(sam, is_auto_mode=is_auto_mode, safe_to_gpu=safe_to)
|
||
sam.sam_wrapper = sam_obj
|
||
print(f'Loads SAM model: {modelname} (device:{device_mode})')
|
||
return (sam,)
|
||
``` |