llama-joycaption-alpha-two-hf-llava

项目介绍：Llama JoyCaption Alpha Two

Llama JoyCaption Alpha Two 是一个新兴的视觉语言模型（Visual Language Model，简称VLM），旨在为图像生成模型提供免费的、开放的和不受限制的图像描述功能。该项目的目标是为社区提供一个可以用于训练扩散模型的工具。

项目背景

现代图像生成技术，如扩散模型，受益于图像的自动描述功能。然而，传统的文本描述工具要么价格高昂且有严格的筛选机制（例如ChatGPT），要么在特定领域表现不佳（如CogVLM在非成人领域表现较弱）。Llama JoyCaption 的诞生就是为了解决这些问题，它希望在图像描述的性能上追赶甚至达到GPT-4的水平，同时保持免费和开放。

项目特点

免费和开放：Llama JoyCaption 将作为一个免费和开放的工具发布，没有权重限制，任何人都可以获取其训练脚本和详细的构建信息，类似于 bigASP 项目。
不受限制：支持NSFW和SFW图像的平等覆盖，不对某些类型的内容进行过滤。
多样性支持：欢迎各类型的图像风格和文化背景内容，无论是数字艺术、照片现实、动漫还是其他风格。
最少过滤：JoyCaption 使用大量的图像进行训练，以最大程度地理解我们的世界。但会杜绝非法内容的使用。

如何开始使用

用户可以通过 Github 获取有关如何使用该模型的详细信息。以下是一个简洁的使用示例：

import torch
import torchvision.transforms.functional as TVF
from PIL import Image
from transformers import AutoTokenizer, LlavaForConditionalGeneration


IMAGE_PATH = "image.jpg"
PROMPT = "Write a long descriptive caption for this image in a formal tone."
MODEL_NAME = "fancyfeast/llama-joycaption-alpha-two-hf-llava"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
llava_model = LlavaForConditionalGeneration.from_pretrained(MODEL_NAME, torch_dtype="bfloat16", device_map=0)
llava_model.eval()

with torch.no_grad():
    image = Image.open(IMAGE_PATH)

    if image.size != (384, 384):
        image = image.resize((384, 384), Image.LANCZOS)

    image = image.convert("RGB")
    pixel_values = TVF.pil_to_tensor(image)
    pixel_values = pixel_values / 255.0
    pixel_values = TVF.normalize(pixel_values, [0.5], [0.5])
    pixel_values = pixel_values.to(torch.bfloat16).unsqueeze(0)

    convo = [
        {
            "role": "system",
            "content": "You are a helpful image captioner.",
        },
        {
            "role": "user",
            "content": PROMPT,
        },
    ]

    convo_string = tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
    convo_tokens = tokenizer.encode(convo_string, add_special_tokens=False, truncation=False)

    input_tokens = []
    for token in convo_tokens:
        if token == llava_model.config.image_token_index:
            input_tokens.extend([llava_model.config.image_token_index] * llava_model.config.image_seq_length)
        else:
            input_tokens.append(token)

    input_ids = torch.tensor(input_tokens, dtype=torch.long).unsqueeze(0)
    attention_mask = torch.ones_like(input_ids)

    generate_ids = llava_model.generate(input_ids=input_ids.to('cuda'), pixel_values=pixel_values.to('cuda'), attention_mask=attention_mask.to('cuda'), max_new_tokens=300, do_sample=True, suppress_tokens=None, use_cache=True)[0]

    generate_ids = generate_ids[input_ids.shape[1]:]

    caption = tokenizer.decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
    caption = caption.strip()
    print(caption)