自学内容网 自学内容网
自学内容网 > 正文

类 WAS_CLIPSeg_Model_Loade

🕗 发布于 2024-07-22 03:33 pycharm  
# CLIPSeg Model Loader

class WAS_CLIPSeg_Model_Loader:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("STRING", {"default": "CIDAS/clipseg-rd64-refined", "multiline": False}),
            },
        }

    RETURN_TYPES = ("CLIPSEG_MODEL",)
    RETURN_NAMES = ("clipseg_model",)
    FUNCTION = "clipseg_model"

    CATEGORY = "WAS Suite/Loaders"

    def clipseg_model(self, model):
        from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation

        cache = os.path.join(MODELS_DIR, 'clipseg')

        inputs = CLIPSegProcessor.from_pretrained(model, cache_dir=cache)
        model = CLIPSegForImageSegmentation.from_pretrained(model, cache_dir=cache)

        return ( (inputs, model), )

# CLIPSeg Node

class WAS_CLIPSeg:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "text": ("STRING", {"default":"", "multiline": False}),
            },
            "optional": {
                "clipseg_model": ("CLIPSEG_MODEL",),
            }
        }

    RETURN_TYPES = ("MASK", "IMAGE")
    RETURN_NAMES = ("MASK", "MASK_IMAGE")
    FUNCTION = "CLIPSeg_image"

    CATEGORY = "WAS Suite/Image/Masking"

    def CLIPSeg_image(self, image, text=None, clipseg_model=None):
        from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation

        image = tensor2pil(image)
        cache = os.path.join(MODELS_DIR, 'clipseg')

        if clipseg_model:
            inputs = clipseg_model[0]
            model = clipseg_model[1]
        else:
            inputs = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined", cache_dir=cache)
            model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined", cache_dir=cache)

        with torch.no_grad():
            result = model(**inputs(text=text, images=image, padding=True, return_tensors="pt"))

        tensor = torch.sigmoid(result[0])
        mask = 1. - (tensor - tensor.min()) / tensor.max()
        mask = mask.unsqueeze(0)
        mask = tensor2pil(mask).convert("L")
        mask = mask.resize(image.size)

        return (pil2mask(mask), pil2tensor(ImageOps.invert(mask.convert("RGB"))))

# CLIPSeg Node

这段代码定义了一个名为 WAS_CLIPSeg_Model_Loader 的类，用于加载 CLIPSeg 模型并将其作为一个节点在 ComfyUI 中使用。以下是详细解释：

类 WAS_CLIPSeg_Model_Loader
这个类负责加载 CLIPSeg 模型，具体包括模型处理器和模型本身，并将其返回以供后续使用。

构造函数 init

python

复制

def __init__(self):
    pass

这个构造函数没有执行任何操作，只是一个占位符。类的实例化时不会进行任何特殊初始化。
类方法 INPUT_TYPES
python
复制

@classmethod
def INPUT_TYPES(cls):
    return {
        "required": {
            "model": ("STRING", {"default": "CIDAS/clipseg-rd64-refined", "multiline": False}),
        },
    }

这是一个类方法，返回一个字典，用于定义节点所需的输入类型。
required 字段定义了模型名称 (model) 是一个字符串类型，默认值为 “CIDAS/clipseg-rd64-refined”，且不支持多行输入。
类属性 RETURN_TYPES 和 RETURN_NAMES
python
复制

RETURN_TYPES = ("CLIPSEG_MODEL",)
RETURN_NAMES = ("clipseg_model",)
RETURN_TYPES 定义了该节点返回的输出类型，输出类型为 CLIPSEG_MODEL。
RETURN_NAMES 定义了该节点返回的输出名称，输出名称为 clipseg_model。

类属性 FUNCTION 和 CATEGORY
python
复制
FUNCTION = “clipseg_model”
CATEGORY = “WAS Suite/Loaders”
FUNCTION 定义了在类中执行的主要函数名称，即 clipseg_model。
CATEGORY 指定了该节点在 ComfyUI 中的分类，便于用户在 UI 中找到该节点。
方法 clipseg_model
python
复制

def clipseg_model(self, model):
    from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation

    cache = os.path.join(MODELS_DIR, 'clipseg')

    inputs = CLIPSegProcessor.from_pretrained(model, cache_dir=cache)
    model = CLIPSegForImageSegmentation.from_pretrained(model, cache_dir=cache)

    return ((inputs, model),)

该方法负责加载 CLIPSeg 模型。
具体步骤如下：

导入必要的库：

python
复制

from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation

定义缓存目录：

python
复制

cache = os.path.join(MODELS_DIR, 'clipseg')

这里将缓存目录定义为 MODELS_DIR 目录下的 clipseg 子目录，用于存放下载的模型文件。
加载处理器和模型：

python
复制

inputs = CLIPSegProcessor.from_pretrained(model, cache_dir=cache)
model = CLIPSegForImageSegmentation.from_pretrained(model, cache_dir=cache)

CLIPSegProcessor.from_pretrained(model, cache_dir=cache)：从预训练模型中加载处理器，如果本地缓存目录中存在模型文件，则从本地加载，否则从远程服务器下载并缓存到本地。
CLIPSegForImageSegmentation.from_pretrained(model, cache_dir=cache)：从预训练模型中加载图像分割模型，遵循与处理器相同的加载逻辑。
返回结果：

python
复制

return ((inputs, model),)

返回一个包含处理器和模型的元组 ((inputs, model),)，其中 inputs 是加载的处理器，model 是加载的图像分割模型。
主要工作流程
模型加载：通过 WAS_CLIPSeg_Model_Loader 类，加载 CLIPSeg 模型处理器和图像分割模型，并将它们作为节点返回。
缓存机制：如果本地缓存目录中存在模型文件，则从本地加载；如果不存在，则从远程服务器下载并缓存到本地，确保下次使用时可以直接从本地加载，减少网络请求和下载时间。
在 ComfyUI 中使用
这个类被定义为一个自定义节点，可以在 ComfyUI 中使用，用于加载 CLIPSeg 模型。
用户可以通过 UI 输入模型名称，然后该节点会自动加载相应的处理器和模型，供后续图像处理任务使用。
总结
WAS_CLIPSeg_Model_Loader 类通过指定模型名称，加载 CLIPSeg 模型处理器和图像分割模型，并将它们返回作为节点输出，便于在 ComfyUI 中进行图像分割任务。


原文地址：https://blog.csdn.net/vv1025/article/details/140586206

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！

相关文章

自学内容网
Copyright © 2015-2024