license: gpl-3.0

模型简介

这是一个实验性视觉模型，基于超长复杂结构生成输入图像的描述/提示词。它融合了图库式标签（逗号分隔关键词）和长文本描述两种形式。

示例：

瀑布, 无人物, 户外, 风景, 树木, 湖泊, 岩石, 河流, 水体, 自然, 植物, 天空, 草地, 白天, 岛屿, 蓝天, 独立场景, 山脉, 森林, 这幅运用数字艺术手法创作的宁静自然景观描绘了丛林环抱中的瀑布水潭。葱郁树木的鲜绿枝叶、盛放的粉红花朵与潋滟湖光共同营造出难以言喻的和谐静谧。高耸的瀑布在繁茂植被映衬下更显壮美，如同大自然馈赠的巨型礼物。整个场景散发着平和安宁的气息，呈现出一幅拥有壮观瀑布的热带景观，四周岩石嶙峋树木葱茏。水面上漂浮着零星树叶，点缀其间的花朵为环境增添色彩与质感。这些色彩明艳、花瓣精致的花朵总能提升场景美感，其巧妙布局引导观者视线，突显这件精妙作品的天然魅力。

开发初衷

这是对复杂长文本描述的实验性探索。目标是创建融合关键词标签与描述文本的混合体，使两者均可用于提示生成，并确保提示词的高质量。当前模型尚未完全实现该目标，需要进一步训练优化。

安装指南

需安装CUDA版PyTorch及相关依赖库。

简易使用脚本

pip install git+https://github.com/huggingface/transformers

from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
from PIL import Image
import requests
import torch

model_id = "mnemic/paligemma-longprompt-v1-safetensors"

url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true"
image = Image.open(requests.get(url, stream=True).raw)

model = PaliGemmaForConditionalGeneration.from_pretrained(model_id).to('cuda').eval()
processor = AutoProcessor.from_pretrained(model_id)

## 前缀设置
prompt = "caption en"
model_inputs = processor(text=prompt, images=image, return_tensors="pt").to('cuda')
input_len = model_inputs["input_ids"].shape[-1]

with torch.inference_mode():
    generation = model.generate(**model_inputs, max_new_tokens=256, do_sample=False)
    generation = generation[0][input_len:]
    decoded = processor.decode(generation, skip_special_tokens=True)
    print(decoded)

支持4/8位量化的批量处理脚本

from transformers import AutoProcessor, PaliGemmaForConditionalGeneration, BitsAndBytesConfig
from PIL import Image
import torch
import os
import glob
from colorama import init, Fore, Style
from datetime import datetime
import time
import re
from huggingface_hub import snapshot_download

# 初始化颜色输出
init(autoreset=True)

# 参数配置
quantization_bits = 8  # 可选None(全精度)/4(4位量化)/8(8位量化)
generation_token_length = 256
min_tokens = 20  # 生成结果的最小token数
max_word_character_length = 30  # 单词最大字符长度阈值
prune_end = True  # 修剪尾部不完整文本直至遇到句号或逗号
output_format = ".txt"  # 输出文件格式

# 生成质量优化
repetition_penalty = 1.15  # 重复惩罚系数
retry_words = ["no_parallel"]  # 触发重新生成的关键词
max_retries = 10
remove_words = ["#", "/", "、", "@", "__", "|", "  ", ";", "~", "\"", "*", "^", ",,", "ON DISPLAY:"]  # 需剔除的特殊字符
strip_contents_inside = ["(", "[", "{"]  # 需连内容删除的括号类型
remove_underscore_tags = True  # 是否移除含下划线的标签

# 模型路径配置
model_name = "mnemic/paligemma-longprompt-v1-safetensors"
models_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models')
model_path = os.path.join(models_dir, model_name.split('/')[-1])

# 目录设置
script_dir = os.path.dirname(os.path.abspath(__file__))
local_model_path = model_path
input_dir = os.path.join(script_dir, 'input')
output_in_input_dir = True  # 设为False可指定独立输出目录
output_dir = input_dir if output_in_input_dir else os.path.join(script_dir, 'output')

# 创建输出目录
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 模型下载函数
def download_model(model_name, model_path):
    if not os.path.exists(model_path):
        print(Fore.YELLOW + f"正在下载模型 {model_name} 至 {model_path}...")
        snapshot_download(repo_id=model_name, local_dir=model_path, local_dir_use_symlinks=False, local_files_only=False)
        print(Fore.GREEN + "模型下载完成")
    else:
        print(Fore.GREEN + f"模型已存在: {model_path}")

# 执行模型下载
download_model(model_name, model_path)

# 校验模型文件
required_files = ["config.json", "tokenizer_config.json"]
missing_files = [f for f in required_files if not os.path.exists(os.path.join(local_model_path, f))]
safetensor_files = [f for f in os.listdir(local_model_path) if f.endswith(".safetensors")]
if missing_files:
    raise FileNotFoundError(f"{local_model_path}中缺失必要文件: {', '.join(missing_files)}")
if not safetensor_files:
    raise FileNotFoundError(f"{local_model_path}中未找到safetensors文件")

# 加载模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(Fore.YELLOW + "正在加载模型与处理器...")
try:
    if quantization_bits == 4:
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
        )
        model = PaliGemmaForConditionalGeneration.from_pretrained(
            local_model_path,
            quantization_config=bnb_config,
            device_map={"": 0},
        ).eval()
    elif quantization_bits == 8:
        bnb_config = BitsAndBytesConfig(
            load_in_8bit=True,
        )
        model = PaliGemmaForConditionalGeneration.from_pretrained(
            local_model_path,
            quantization_config=bnb_config,
            device_map={"": 0},
        ).eval()
    elif quantization_bits is None:
        model = PaliGemmaForConditionalGeneration.from_pretrained(
            local_model_path
        ).eval()
        model.to(device)
    else:
        raise ValueError("quantization_bits参数值错误，仅支持None/4/8")

    processor = AutoProcessor.from_pretrained(local_model_path, local_files_only=True)
    print(Fore.GREEN + "模型加载成功")
except OSError as e:
    print(Fore.RED + f"加载错误: {e}")
    raise

# 递归处理输入目录图像
image_extensions = ['jpg', 'jpeg', 'png', 'webp']
image_paths = []
for ext in image_extensions:
    image_paths.extend(glob.glob(os.path.join(input_dir, '**', f'*.{ext}'), recursive=True))

print(Fore.YELLOW + f"发现待处理图像: {len(image_paths)}张\n")

def prune_text(text):
    if not prune_end:
        return text
    last_period_index = text.rfind('.')
    last_comma_index = text.rfind(',')
    prune_index = max(last_period_index, last_comma_index)
    if prune_index != -1:
        return text[:prune_index].strip()
    return text

def contains_retry_word(text, retry_words):
    return any(word in text for word in retry_words)

def remove_unwanted_words(text, remove_words):
    for word in remove_words:
        text = text.replace(word, ' ')
    return text

def strip_contents(text, chars):
    for char in chars:
        if char == "(":
            text = re.sub(r'\([^)]*\)', ' ', text)
        elif char == "[":
            text = re.sub(r'\[[^\]]*\]', ' ', text)
        elif char == "{":
            text = re.sub(r'\{[^}]*\}', ' ', text)
    text = re.sub(r'\s{2,}', ' ', text)
    text = re.sub(r'\s([,.!?;])', r'\1', text)
    text = re.sub(r'([,.!?;])\s', r'\1 ', text)
    return text.strip()

def remove_long_words(text, max_word_length):
    words = text.split()
    for i, word in enumerate(words):
        if len(word) > max_word_length:
            last_period_index = text.rfind('.', 0, text.find(word))
            last_comma_index = text.rfind(',', 0, text.find(word))
            prune_index = max(last_period_index, last_comma_index)
            if prune_index != -1:
                return text[:prune_index].strip()
            else:
                return text[:text.find(word)].strip()
    return text

def clean_text(text):
    text = remove_unwanted_words(text, remove_words)
    text = strip_contents(text, strip_contents_inside)
    text = remove_long_words(text, max_word_character_length)
    text = re.sub(r'[^\x00-\x7F]+', '', text)
    text = re.sub(r'\s+', ' ', text).strip()
    if remove_underscore_tags:
        text = ' '.join([word for word in text.split() if '_' not in word])
    return text

for image_path in image_paths:
    output_file_path = os.path.splitext(image_path)[0] + output_format if output_in_input_dir else os.path.join(output_dir, os.path.splitext(os.path.relpath(image_path, input_dir))[0] + output_format)
    
    if os.path.exists(output_file_path):
        continue

    try:
        start_time = datetime.now()
        print(Fore.CYAN + f"[{start_time.strftime('%Y-%m-%d %H:%M:%S')}] 开始处理 {image_path}")
        
        image = Image.open(image_path).convert('RGB')
        prompt = "caption en"
        model_inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
        input_len = model_inputs["input_ids"].shape[-1]

        retries = 0
        success = False
        while retries < max_retries:
            with torch.inference_mode():
                generation_start_time = time.time()
                generation = model.generate(
                    **model_inputs,
                    max_new_tokens=generation_token_length,
                    do_sample=True,
                    temperature=0.7,
                    top_k=50,
                    top_p=0.9,
                    no_repeat_ngram_size=2,
                    repetition_penalty=repetition_penalty
                )
                generation_end_time = time.time()
                generation = generation[0][input_len:]
                decoded = processor.decode(generation, skip_special_tokens=True)
                pruned_text = prune_text(decoded)
                
                if not contains_retry_word(pruned_text, retry_words) and len(pruned_text.split()) >= min_tokens:
                    success = True
                    break
                retries += 1
                print(Fore.YELLOW + f"{image_path} 生成结果包含重试词或token不足，正在第{retries}次重试...")
            
            if retries == max_retries:
                print(Fore.RED + f"{image_path} 已达最大重试次数，将保存当前结果")

        cleaned_text = clean_text(pruned_text)

        os.makedirs(os.path.dirname(output_file_path), exist_ok=True)
        with open(output_file_path, 'w', encoding='utf-8') as f:
            f.write(cleaned_text)
        
        end_time = datetime.now()
        duration = generation_end_time - generation_start_time
        
        print(Fore.GREEN + f"[{end_time.strftime('%Y-%m-%d %H:%M:%S')}] 已完成 {image_path} -> {output_file_path}")
        print(Fore.LIGHTBLACK_EX + f"生成结果: {cleaned_text}")
        print(Fore.LIGHTBLACK_EX + f"生成耗时: {duration:.2f}秒\n")
        
        del model_inputs
        torch.cuda.empty_cache()
    except Exception as e:
        print(Fore.RED + f"处理错误 {image_path}: {e}\n")

该脚本也适用于其他Paligemma模型。推荐使用：https://huggingface.co/gokaygokay/paligemma-rich-captions

特别感谢Gökay Aydoğan在模型微调与转换过程中的鼎力相助！