File size: 7,375 Bytes
45cf443 92c7321 45cf443 9a1a7d2 a4fb8db 9a1a7d2 c6bc767 9a1a7d2 45cf443 9a1a7d2 45cf443 9a1a7d2 45cf443 9a1a7d2 45cf443 9a1a7d2 c6bc767 45cf443 9a1a7d2 45cf443 a4fb8db 45cf443 c6bc767 45cf443 9a1a7d2 45cf443 a4fb8db 45cf443 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | import time
import argparse
import os
import warnings
import torch
import random
from PIL import Image
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
from models import VLM, VLMConfig
from utils import setup_seed, get_vlm_model_params
warnings.filterwarnings('ignore')
def init_model(args):
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, trust_remote_code=True)
if args.native:
ckp = f'{args.save_dir}/{args.weight}.pth'
if not os.path.exists(ckp):
moe_suffix = '_moe' if args.use_moe else ''
ckp = f'{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
state = torch.load(ckp, map_location=args.device)
n_layers = max(int(k.split('.')[2]) for k in state if k.startswith('model.layers.')) + 1
model = VLM(
VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=n_layers, use_moe=bool(args.use_moe)),
vision_model_path=args.vision_model_dir
)
model.load_state_dict({k: v for k, v in state.items() if 'mask' not in k}, strict=False)
processor = model.vision_encoder.processor if hasattr(model.vision_encoder, 'processor') else None
else:
model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
hf_vision, processor = VLM.get_vision_model(args.vision_model_dir)
if hf_vision is not None:
model.vision_encoder = hf_vision
get_vlm_model_params(model, model.config)
model = model.eval()
if "cuda" in args.device: model = model.half()
return model.to(args.device), tokenizer, processor
def main():
parser = argparse.ArgumentParser(description="MiniMind-V 视觉多模态推理")
parser.add_argument('--load_from', default='', type=str, help="模型加载路径（transformers格式，native模式不感知此参数）")
parser.add_argument('--tokenizer_path', default='checkpoint/omni/native_hf', type=str, help="tokenizer 路径")
parser.add_argument('--native', action='store_true', help="加载原生 torch checkpoint（由 save_dir/weight/hidden_size 定位）")
parser.add_argument('--save_dir', default='checkpoint', type=str, help="模型权重目录")
parser.add_argument('--weight', default='sft_vlm', type=str, help="权重名称前缀（pretrain_vlm, sft_vlm）")
parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
parser.add_argument('--vision_model_dir', default='checkpoint/siglip', type=str, help="视觉模型目录")
parser.add_argument('--max_new_tokens', default=512, type=int, help="最大生成长度")
parser.add_argument('--temperature', default=0.7, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
parser.add_argument('--top_p', default=0.85, type=float, help="nucleus采样阈值（0-1）")
parser.add_argument('--image_dir', default='./dataset/eval_images/', type=str, help="测试图像目录")
parser.add_argument('--show_speed', default=1, type=int, help="显示decode速度（tokens/s）")
parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备")
parser.add_argument('--open_thinking', default=0, type=int, help="是否开启自适应思考（0=否，1=是）")
args = parser.parse_args()
model, tokenizer, preprocess = init_model(args)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
mode = input('[0] 自动测试目录图片\n[1] 手动输入（图片路径 + 文本）\n')
if mode == '0':
prompt = "<image>\n请描述这张图中的主要物体和场景。"
for image_file in sorted(os.listdir(args.image_dir)):
if not image_file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
continue
setup_seed(random.randint(1, 31415926))
image_path = os.path.join(args.image_dir, image_file)
image = Image.open(image_path).convert('RGB')
pixel_values = {k: v.to(args.device) for k, v in VLM.image2tensor(image, preprocess).items()}
messages = [{"role": "user", "content": prompt.replace('<image>', getattr(model.config, 'image_special_token', '<|image_pad|>') * getattr(model.config, 'image_token_len', 64))}]
inputs_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, open_thinking=bool(args.open_thinking))
inputs = tokenizer(inputs_text, return_tensors="pt", truncation=True).to(args.device)
print(f'[图像]: {image_file}')
print(f"💬: {repr(prompt)}")
print('🤖: ', end='')
st = time.time()
generated_ids = model.generate(
inputs=inputs["input_ids"], attention_mask=inputs["attention_mask"],
max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer,
pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
top_p=args.top_p, temperature=args.temperature, pixel_values=pixel_values
)
gen_tokens = len(generated_ids[0]) - len(inputs["input_ids"][0])
print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n\n') if args.show_speed else print('\n\n')
else:
while True:
image_path = input('图片路径（留空跳过）: ').strip()
prompt = input('💬: ').strip()
if not prompt:
break
setup_seed(random.randint(1, 31415926))
pixel_values = None
if image_path and os.path.exists(image_path):
image = Image.open(image_path).convert('RGB')
pixel_values = {k: v.to(args.device) for k, v in VLM.image2tensor(image, preprocess).items()}
if '<image>' not in prompt:
prompt = '<image>\n' + prompt
content = prompt.replace('<image>', getattr(model.config, 'image_special_token', '<|image_pad|>') * getattr(model.config, 'image_token_len', 64))
messages = [{"role": "user", "content": content}]
inputs_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, open_thinking=bool(args.open_thinking))
inputs = tokenizer(inputs_text, return_tensors="pt", truncation=True).to(args.device)
print('🤖: ', end='')
st = time.time()
gen_kwargs = dict(
inputs=inputs["input_ids"], attention_mask=inputs["attention_mask"],
max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer,
pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
top_p=args.top_p, temperature=args.temperature,
)
if pixel_values is not None:
gen_kwargs['pixel_values'] = pixel_values
generated_ids = model.generate(**gen_kwargs)
gen_tokens = len(generated_ids[0]) - len(inputs["input_ids"][0])
print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n') if args.show_speed else print()
if __name__ == "__main__":
main()
|