5.4 KiB
5.4 KiB
Documentation
- Class name: WAS_BLIP_Analyze_Image
- Category: WAS Suite/Text/AI
- Output node: False
- Repo Ref: https://github.com/WASasquatch/was-node-suite-comfyui
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!',)