Files
CavinHuang-comfyui-nodes-docs/docs/BLIP Analyze Image.md
2024-05-26 20:35:16 +08:00

5.4 KiB
Raw Permalink Blame History

Documentation

WAS_BLIP_Analyze_Image节点旨在使用BLIP(Bootstrapped Language Image Pretraining)模型分析和解释图像内容。它提供了生成标题和用自然语言问题询问图像的功能,提供了对输入图像的视觉和上下文方面的洞察。

Input types

Required

  • image
    • 图像参数对节点的操作至关重要,因为它是模型分析以生成标题或回答问题的输入。它直接影响输出和分析质量。
    • Comfy dtype: IMAGE
    • Python dtype: PIL.Image or torch.Tensor
  • mode
    • 模式参数确定节点将对图像执行的分析类型。它可以是'caption'以生成描述,或者是'interrogate'以回答有关图像内容的问题。
    • Comfy dtype: COMBO['caption', 'interrogate']
    • Python dtype: str

Optional

  • question
    • 当模式设置为'interrogate'时,使用问题参数。它指定了模型将根据图像内容尝试回答的查询。问题的措辞可以影响答案的准确性和相关性。
    • Comfy dtype: STRING
    • Python dtype: str
  • blip_model
    • blip_model参数允许用户为节点提供预加载的BLIP模型,而不是下载新模型。这可以提高效率,并且在多次运行节点时特别有用。
    • Comfy dtype: BLIP_MODEL
    • Python dtype: Tuple[torch.nn.Module, str]

Output types

  • output
    • 输出参数代表节点分析的结果,可以是描述图像的标题或对提出的询问性问题的回答。它概括了节点对图像内容的理解。
    • Comfy dtype: TEXT
    • Python dtype: str

Usage tips

  • Infra type: GPU

Source code

class WAS_BLIP_Analyze_Image:

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {'required': {'image': ('IMAGE',), 'mode': (['caption', 'interrogate'],), 'question': ('STRING', {'default': 'What does the background consist of?', 'multiline': True})}, 'optional': {'blip_model': ('BLIP_MODEL',)}}
    RETURN_TYPES = (TEXT_TYPE,)
    FUNCTION = 'blip_caption_image'
    CATEGORY = 'WAS Suite/Text/AI'

    def blip_caption_image(self, image, mode, question, blip_model=None):

        def transformImage(input_image, image_size, device):
            raw_image = input_image.convert('RGB')
            raw_image = raw_image.resize((image_size, image_size))
            transform = transforms.Compose([transforms.Resize(raw_image.size, interpolation=InterpolationMode.BICUBIC), transforms.ToTensor(), transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))])
            image = transform(raw_image).unsqueeze(0).to(device)
            return image.view(1, -1, image_size, image_size)
        from torchvision import transforms
        from torchvision.transforms.functional import InterpolationMode
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        conf = getSuiteConfig()
        image = tensor2pil(image)
        size = 384
        tensor = transformImage(image, size, device)
        if blip_model:
            mode = blip_model[1]
        if mode == 'caption':
            if blip_model:
                model = blip_model[0].to(device)
            else:
                from .modules.BLIP.blip_module import blip_decoder
                blip_dir = os.path.join(MODELS_DIR, 'blip')
                if not os.path.exists(blip_dir):
                    os.makedirs(blip_dir, exist_ok=True)
                torch.hub.set_dir(blip_dir)
                if conf.__contains__('blip_model_url'):
                    model_url = conf['blip_model_url']
                else:
                    model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_capfilt_large.pth'
                model = blip_decoder(pretrained=model_url, image_size=size, vit='base')
                model.eval()
                model = model.to(device)
            with torch.no_grad():
                caption = model.generate(tensor, sample=False, num_beams=6, max_length=74, min_length=20)
                cstr(f'\x1b[33mBLIP Caption:\x1b[0m {caption[0]}').msg.print()
                return (caption[0],)
        elif mode == 'interrogate':
            if blip_model:
                model = blip_model[0].to(device)
            else:
                from .modules.BLIP.blip_module import blip_vqa
                blip_dir = os.path.join(MODELS_DIR, 'blip')
                if not os.path.exists(blip_dir):
                    os.makedirs(blip_dir, exist_ok=True)
                torch.hub.set_dir(blip_dir)
                if conf.__contains__('blip_model_vqa_url'):
                    model_url = conf['blip_model_vqa_url']
                else:
                    model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_vqa_capfilt_large.pth'
                model = blip_vqa(pretrained=model_url, image_size=size, vit='base')
                model.eval()
                model = model.to(device)
            with torch.no_grad():
                answer = model(tensor, question, train=False, inference='generate')
                cstr(f'\x1b[33m BLIP Answer:\x1b[0m {answer[0]}').msg.print()
                return (answer[0],)
        else:
            cstr(f'The selected mode `{mode}` is not a valid selection!').error.print()
            return ('Invalid BLIP mode!',)