63 lines
1.7 KiB
Markdown
63 lines
1.7 KiB
Markdown
|
||
# Documentation
|
||
- Class name: AV_SAMLoader
|
||
- Category: Art Venture/Segmentation
|
||
- Output node: False
|
||
|
||
AV_SAMLoader节点专为加载Art Venture的分割任务定制的SAM（Segmentation-Aware Models，分割感知模型）而设计。它能根据模型名称动态加载SAM模型，确保在图像处理任务中使用适当的分割模型。
|
||
|
||
# Input types
|
||
## Required
|
||
- model_name
|
||
- 指定要加载的SAM模型的名称。此参数对于识别和加载正确的分割模型至关重要。
|
||
- Comfy dtype: COMBO[STRING]
|
||
- Python dtype: str
|
||
|
||
# Output types
|
||
- sam_model
|
||
- 返回加载完毕的SAM模型，可直接用于分割任务。
|
||
- Comfy dtype: AV_SAM_MODEL
|
||
- Python dtype: torch.nn.Module
|
||
|
||
|
||
## Usage tips
|
||
- Infra type: `GPU`
|
||
- Common nodes: unknown
|
||
|
||
|
||
## Source code
|
||
```python
|
||
class SAMLoader:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"model_name": (folder_paths.get_filename_list("sams"),),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("AV_SAM_MODEL",)
|
||
RETURN_NAMES = ("sam_model",)
|
||
FUNCTION = "load_model"
|
||
CATEGORY = "Art Venture/Segmentation"
|
||
|
||
def load_model(self, model_name):
|
||
modelname = folder_paths.get_full_path("sams", model_name)
|
||
|
||
state_dict = comfy.utils.load_torch_file(modelname)
|
||
encoder_size = state_dict["image_encoder.patch_embed.proj.bias"].shape[0]
|
||
|
||
if encoder_size == 1280:
|
||
model_kind = "vit_h"
|
||
elif encoder_size == 1024:
|
||
model_kind = "vit_l"
|
||
else:
|
||
model_kind = "vit_b"
|
||
|
||
sam = sam_model_registry[model_kind]()
|
||
sam.load_state_dict(state_dict)
|
||
|
||
return (sam,)
|
||
|
||
```
|