CavinHuang-comfyui-nodes-docs/docs/AV_SAMLoader.md

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

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,)