gradio/mme.py

import soundfile as sf
import os
###os.environ["CUDA_VISIBLE_DEVICES"] = "0"
from io import BytesIO
from urllib.request import urlopen
from qwen_vl_utils import process_vision_info
from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniModel, AutoTokenizer, AutoProcessor, Qwen2_5OmniThinkerModel, BitsAndBytesConfig
from src.llamafactory.model.loader import patch_tokenizer, patch_processor
from src.llamafactory.data.template import get_template_and_fix_tokenizer
from argparse import Namespace, ArgumentParser
import torch
import json
from tqdm import tqdm
import time
import math
from transformers.cache_utils import DynamicCache
import glob
from safetensors import safe_open

def endswith_tensor(seq, pattern):
    if isinstance(seq, torch.Tensor):
        return torch.equal(seq[-len(pattern):], pattern.to(seq.device))
    else:  # list
        return seq[-len(pattern):] == pattern
    
def truncate_kv_cache(past_key_values, n_tokens_to_remove=4):
    new_cache = DynamicCache()
    for layer_idx, (k, v) in enumerate(past_key_values):
        # k/v shape: [batch, num_heads, seq_len, head_dim]
        new_k = k[:, :, :-n_tokens_to_remove, :].contiguous()
        new_v = v[:, :, :-n_tokens_to_remove, :].contiguous()
        new_cache.update(new_k, new_v, layer_idx)
    return new_cache

# 添加命令行参数解析
parser = ArgumentParser(description="在指定CUDA设备上运行Qwen2_5Omni模型进行视频问答。")
parser.add_argument("--model_path", type=str, default="whole_model/model", help="模型文件夹的路径。")
parser.add_argument("--test_data_path", type=str, default="egoschema/Subset/test_egoschema.jsonl", help="测试数据的JSONL文件路径。")
parser.add_argument("--output_path", type=str, default="eval/basic/egoschema/test_result.jsonl", help="结果输出的JSONL文件路径。")

args = parser.parse_args()

# 不再使用 os.environ，直接创建 torch.device 对象
#device = torch.device(f"cuda:{args.cuda}" if torch.cuda.is_available() else "cpu")
#print(f"正在使用设备: {device}")

print(f"从 '{args.model_path}' 加载模型")
# Hardware-agnostic, env-driven loader. Defaults are Turing-safe (Quadro RTX 6000, sm_75):
#   ROMA_DTYPE=float16  (bf16 is not accelerated on Turing)
#   ROMA_ATTN=sdpa      (FlashAttention-2 is unsupported on Turing; use eager if sdpa fails)
#   ROMA_LOAD_8BIT=0    (fp16 sharded across all visible GPUs via device_map=auto; set 1 for 1-GPU)
_DTYPE = {"float16": torch.float16, "bfloat16": torch.bfloat16, "auto": "auto"}[os.environ.get("ROMA_DTYPE", "float16")]
_load_kwargs = dict(
    torch_dtype=_DTYPE,
    device_map="auto",
    attn_implementation=os.environ.get("ROMA_ATTN", "sdpa"),
    trust_remote_code=True,
)
if os.environ.get("ROMA_LOAD_8BIT", "0") == "1":
    _load_kwargs["quantization_config"] = BitsAndBytesConfig(
        load_in_8bit=True,
        llm_int8_skip_modules=["talker", "token2wav", "visual", "audio_tower",
                               "gate_head", "gate_mixer", "gate_head_pro_fc1", "gate_head_pro_fc2"],
    )
model = Qwen2_5OmniModel.from_pretrained(args.model_path, **_load_kwargs)
# if not hasattr(model.thinker, "gate_head"):
#     H = model.thinker.config.get_text_config().hidden_size
#     device = next(model.parameters()).device
#     dtype = next(model.parameters()).dtype

#     class GateMixer(torch.nn.Module):
#         def __init__(self, K, device, dtype):
#             super().__init__()
#             self.logits = torch.nn.Parameter(torch.zeros(K, device=device, dtype=dtype))
#         def weights(self): return torch.softmax(self.logits, dim=0)

#     model.thinker.gate_head = torch.nn.Linear(H, 1, bias=True).to(device=device, dtype=dtype)
#     model.thinker.gate_mixer = GateMixer(K=4, device=device, dtype=dtype)  # K按你训练时的值
#     model.thinker.gate_layer_ids = [-4, -3, -2, -1] 

# state = {}
# for shard in sorted(glob.glob(os.path.join(args.model_path, "model-*.safetensors"))):
#     with safe_open(shard, framework="pt", device="cpu") as f:
#         for k in f.keys():
#             if k.startswith("thinker.gate_head.") or k.startswith("thinker.gate_mixer."):
#                 state[k] = f.get_tensor(k)

# head_sd  = {k.split("thinker.gate_head.", 1)[1]:  v for k,v in state.items() if k.startswith("thinker.gate_head.")}
# mixer_sd = {k.split("thinker.gate_mixer.",1)[1]:  v for k,v in state.items() if k.startswith("thinker.gate_mixer.")}

# model.thinker.gate_head.load_state_dict(head_sd, strict=True)
# model.thinker.gate_mixer.load_state_dict(mixer_sd, strict=True)
# print("✅ gate_head / gate_mixer 权重已恢复")

tokenizer = AutoTokenizer.from_pretrained(
            args.model_path,
            use_fast=True,
            split_special_tokens = False,
            padding_side="left",
            trust_remote_code = True,
            cache_dir = None,
            revision = 'main',
            token = None
        )
newline_token_id = tokenizer.encode("\n", add_special_tokens=False)
processor_args_dict = {
    "image_max_pixels": 262144,
    "image_min_pixels": 1024,
    "image_do_pan_and_scan": False,
    "crop_to_patches": False,
    "video_max_pixels": 65536, #147456,
    "video_min_pixels": 256,
    "video_fps": 2.0,
    "video_maxlen": 14400, #2h
    "audio_sampling_rate": 16000,
    "use_audio_in_video": True
}

processor_args = Namespace(**processor_args_dict)
processor = AutoProcessor.from_pretrained(args.model_path, trust_remote_code = True, cache_dir = None, revision = 'main', token = None)

patch_processor(processor,tokenizer, processor_args)

args_dict = {
    "template": "streaming_turn",
    "train_on_prompt": False,
    "tool_format": None
}
template_args= Namespace(**args_dict)
template = get_template_and_fix_tokenizer(tokenizer, template_args)
###################
def transform_example_format(example: dict[str, any]) -> dict[str, any]:
    if not isinstance(example, dict):
        raise ValueError("输入 'example' 必须是一个字典。")

    output: dict[str, any] = {
        "_prompt": example.get("query", []), 
        "_response": example.get("ans", []),  
        "_system": "",                       
        "_tools": example.get("tools", "") if example.get("tools") else "",                                                                 
        "_images": example.get("images") if len(example.get("images"))!=0 else None,
        "_videos": example.get("videos") if len(example.get("videos"))!=0 else None,
        "_audios": []
    }
    return output
##################
def get_multimodal_input_ids(prompt,response,system,tools,images,videos,audios):
    messages = template.mm_plugin.process_messages(
        [[prompt,response]], images, videos, audios, processor,mode = "infer"
    )
    encoded_pairs = template.encode_multiturn(tokenizer, messages, system, tools)
    inputs_list = []
    for input_multimodal, _ in encoded_pairs:
        inputs_list.append(input_multimodal)
    return inputs_list
#################
# 开始处理数据
if args.test_data_path.endswith(".jsonl"):
    with open(args.test_data_path, "r") as f:
        conversations = [json.loads(line) for line in f]
else:
    with open(args.test_data_path, "r") as f:
        conversations = json.load(f)

processed_ids = set()
if os.path.exists(args.output_path):
    print(f"发现已存在的输出文件: {args.output_path}。正在读取已处理的ID...")
    try:
        with open(args.output_path, "r", encoding="utf-8") as f_out:
            for line in f_out:
                # 跳过空行
                if not line.strip():
                    continue
                try:
                    # 解析每一行JSON
                    processed_data = json.loads(line)
                    # 确保'id'键存在
                    if 'id' in processed_data:
                        processed_ids.add(processed_data['id'])
                except json.JSONDecodeError:
                    print(f"警告：跳过无法解析的行: {line.strip()}")
    except Exception as e:
        print(f"错误：读取输出文件时发生错误: {e}")

original_total = len(conversations)
tasks_to_process = [item for item in conversations if item.get('id') not in processed_ids]
remaining_total = len(tasks_to_process)

print("-" * 50)
print(f"总任务数: {original_total}")
print(f"已完成任务数: {len(processed_ids)}")
print(f"待处理任务数: {remaining_total}")
print("-" * 50)

all_probs = []
import numpy as np

# 如果所有任务都已完成，则直接退出
if remaining_total == 0:
    print("所有任务均已完成！程序退出。")
    exit()
with torch.no_grad():
    for data in tqdm(tasks_to_process):
        data_formated = transform_example_format(data)
        multimodal_input_id_list = get_multimodal_input_ids(
            prompt=data_formated["_prompt"],
            response=data_formated["_response"],
            system="",
            tools="",
            images=data_formated["_images"] or [],
            videos=data_formated["_videos"] or [],
            audios=[],
        )

        batch_images = []
        batch_videos = []
        batch_videos.append(data['videos'][0])
        batch_audios = []

        batch_imglens = []
        batch_imglens.append(0)
        batch_vidlens = []
        batch_audlens = []
        batch_audlens.append(1)

        batch_input_ids = []
        batch_input_ids.append(multimodal_input_id_list[0])
        messages = []
        messages.append([data_formated["_prompt"], data_formated["_response"]])

        mm_inputs = template.mm_plugin.get_mm_inputs(
                batch_images,
                batch_videos,
                batch_audios,
                batch_imglens,
                batch_vidlens,
                batch_audlens,
                batch_input_ids,
                processor,
                messages = messages,
            )
        features = {} 
        ### 在这里进行每个chunk的拼接
        input_ids = []
        sum_video_token = 0
        sum_audio_token = 0
        final_answer = {}

        # 1. 初始化状态变量
        past_key_values = None
        generated_ids_since_last_chunk = torch.tensor([], dtype=torch.long, device=model.device)
        final_answers = []
        final_answer_text = []
        final_answer_time = None
        last_rope_delta = None
    ######################
        flag = False
        ask_time = messages[0][0][0]['time'] + math.ceil(messages[0][0][0]['duration'])
        # import random
        # ask_time = random.randint(1, ask_time-1)
        for i, chunk in enumerate(multimodal_input_id_list):     
            if i >= ask_time:
                flag = True
                break
            if endswith_tensor(chunk,[151645,198,151644,77091,198]):
                chunk = chunk[:-5]
            # 确保 chunk 的类型与操作安全
            if isinstance(chunk, torch.Tensor):
                input_ids.extend(chunk.tolist())
                sum_video_token += int((chunk == 151656).sum().item()) # <|VIDEO|>
                sum_audio_token += int((chunk == 151646).sum().item()) # <|AUDIO|>
            else:
                input_ids.extend(chunk)
                sum_video_token += chunk.count(151656) #<|VIDEO|>
                sum_audio_token += chunk.count(151646) #<|AUDIO|>
            num_video_features = sum_video_token * 4
            num_audio_features = sum_audio_token        
        
        assistant_prefix = [151645, 198, 151644, 77091, 198]    #<|im_end|>\n<|im_start|>assistant\n
        user_prefix = [198, 151644, 872, 198]    #\n<|im_start|>user\n
        # 关键改动：不要把 Tensor 直接 extend 到 list
        input_ids.extend(assistant_prefix)
        
        if not flag and (i == len(multimodal_input_id_list)-1):
            i = i+1
        features['input_ids'] = torch.tensor([input_ids]).to(model.device) 
        features['attention_mask'] = torch.ones([1,len(input_ids)],dtype=torch.int64).to(model.device) 
        # 将 mm_inputs 中的所有 tensor 移动到指定设备
        features['video_grid_thw'] = mm_inputs['video_grid_thw'].clone().to(model.device)
        features['video_grid_thw'][0, 0] = i
        features['pixel_values_videos'] = mm_inputs['pixel_values_videos'][:num_video_features,:].to(model.dtype).to(model.device) 
        features['input_features'] = mm_inputs['input_features'][:, : , :(i)*100].to(model.dtype).to(model.device)
        features['feature_attention_mask'] = mm_inputs['feature_attention_mask'][:, :(i)*100].to(model.device)
        features['video_second_per_grid'] = mm_inputs['video_second_per_grid'].to(model.dtype).to(model.device)
        
########################### probing

#         probe_inputs = {
#             **features,
#             "input_ids":          features["input_ids"],
#             "attention_mask":     features["attention_mask"],
#             "use_cache":          True,
#             "output_hidden_states": True,
#             "return_dict":        True,
#             "past_key_values":    past_key_values,
#             "rope_deltas":        last_rope_delta,
#         }

#         out = model.thinker(**probe_inputs)
#         hs_all = out.hidden_states[1] if (isinstance(out.hidden_states, tuple) and isinstance(out.hidden_states[1], (list, tuple))) else out.hidden_states
#         B, T, H = hs_all[-1].shape

#         anchor_pos = features['input_ids'].size(1) - 1
#         anchor_idx = torch.tensor([[anchor_pos]], device=features['input_ids'].device)
#         idx = anchor_idx.unsqueeze(-1).expand(B, 1, H)

#         # 层混合
#         layer_ids = getattr(model.thinker, "gate_layer_ids", [-4, -3, -2, -1])
#         mix_w = model.thinker.gate_mixer.weights() 
#         print("mix_w:",mix_w)
#         h_mix = 0.0
#         L = len(hs_all)
#         for w, lid in zip(mix_w, layer_ids):
#             lid = lid if lid >= 0 else L + lid
#             lid = int(max(0, min(L - 1, lid)))
#             h_l = hs_all[lid]                 # [B,T,H]
#             h_anchor = torch.gather(h_l, 1, idx)    # [B,1,H]
#             h_mix = h_mix + w * h_anchor

#         logit = model.thinker.gate_head(h_mix).squeeze(-1).squeeze(-1)   # [B]
#         prob  = torch.sigmoid(logit).item()
#         gate_decision = (prob >= 0.1)
#         all_probs.append(prob)
#         print(prob)

# print("平均概率:", np.mean(all_probs), "概率方差:", np.var(all_probs))

######################################## end probing
        output = model.generate(
            **features, 
            thinker_max_new_tokens=25, 
            use_audio_in_video=True, 
            return_audio=False, 
            streaming=True,
            past_key_values=past_key_values,
            output_scores=True,
            mode="infer",
            rope_deltas = last_rope_delta
        )
        last_rope_delta = output["rope_deltas"]
        past_key_values = output.past_key_values

        newly_generated_ids = output.sequences[0,len(input_ids):]
        # 关键改动：空输出保护 + list[int] 保持
        if newly_generated_ids.numel() > 0:
            ids_list = newly_generated_ids.tolist()
            input_ids.extend(ids_list)
            if ids_list[-1] != 151645: 
                input_ids.extend([151643,151645])
            newly_generated_text = processor.decode(newly_generated_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)
        else:
            newly_generated_text = ""
        # 关键改动：不要把 Tensor 直接 extend
        input_ids.extend(user_prefix)

        print("newly_generated:", [newly_generated_text])

        stop_token_str = "<|im_end|>"
        is_answer_finished = False
        chunk_to_process = newly_generated_text

        if stop_token_str in newly_generated_text:
            stop_index = newly_generated_text.index(stop_token_str)
            chunk_to_process = newly_generated_text[:stop_index]
            # if chunk_to_process == "<|silence|>":
            #     chunk_to_process = ""
            is_answer_finished = True
        if chunk_to_process:
            if final_answer_time is None:
                final_answer_time = i+1
            # 关键改动：append 而不是 extend（避免按字符拆）
            final_answer_text.append(chunk_to_process)
        if is_answer_finished and final_answer_text:
            final_answer = {
                'time': final_answer_time,
                'text': "".join(final_answer_text)
            }
            final_answers.append(final_answer)
            
            with open(args.output_path, "a") as f:
                final_answer = {
                    'id': data['id'],
                    'question' : data['query'][0]['text'],
                    'prediction': final_answers,
                    "gt": data['answer']
                }
                f.write(json.dumps(final_answer, ensure_ascii=False) + "\n")   
            del data_formated, multimodal_input_id_list, mm_inputs, features, output, final_answers
            torch.cuda.empty_cache()
            continue
    #######
        for i, chunk in enumerate(multimodal_input_id_list):
            #print("进入到第二个循环")
            if i <= ask_time:
                continue
            # 将 chunk 移动到指定设备（关键改动：类型安全）
            if isinstance(chunk, torch.Tensor):
                input_ids.extend(chunk.tolist())
                prev_sum_video_token = sum_video_token
                sum_video_token += int((chunk == 151656).sum().item())
                sum_audio_token += int((chunk == 151646).sum().item())
            else:
                input_ids.extend(chunk) 
                prev_sum_video_token = sum_video_token
                sum_video_token += chunk.count(151656)
                sum_audio_token += chunk.count(151646)
            num_video_features = sum_video_token * 4
            num_audio_features = sum_audio_token
            video_features_before_this_chunk = prev_sum_video_token * 4
            
            features = {}
            if i != len(multimodal_input_id_list)-1:
                # 关键改动：不要 extend Tensor
                input_ids.extend(assistant_prefix)
            features['input_ids'] = torch.tensor([input_ids]).to(model.device)
            features['attention_mask'] = torch.ones([1,len(input_ids)],dtype=torch.int64).to(model.device)
            # 将 mm_inputs 中的所有 tensor 移动到指定设备
            features['video_grid_thw'] = mm_inputs['video_grid_thw'].clone().to(model.device)
            features['video_grid_thw'][0, 0] = 1
            features['pixel_values_videos'] = mm_inputs['pixel_values_videos'][video_features_before_this_chunk:num_video_features,:].to(model.dtype).to(model.device)
            features['input_features'] = mm_inputs['input_features'][:, : , (i)*100:(i+1)*100].to(model.dtype).to(model.device)
            features['feature_attention_mask'] = mm_inputs['feature_attention_mask'][:, (i)*100:(i+1)*100].to(model.device)
            features['video_second_per_grid'] = mm_inputs['video_second_per_grid'].to(model.dtype).to(model.device)
            
            start_time = time.time()
            
            output = model.generate(
                **features, 
                thinker_max_new_tokens=25, 
                use_audio_in_video=True, 
                return_audio=False, 
                streaming=True,
                past_key_values=past_key_values,
                output_scores=True,
                mode="infer",
                rope_deltas = last_rope_delta
            )
            last_rope_delta = output["rope_deltas"]
            end_time = time.time()
            print("完成一个chunk的耗时：", end_time - start_time)
            
            past_key_values = output.past_key_values
            newly_generated_ids = output.sequences[0,len(input_ids):]
        
            # 关键改动：空输出保护 + list[int] 保持
            if newly_generated_ids.numel() > 0:
                ids_list = newly_generated_ids.tolist()
                input_ids.extend(ids_list)
                if ids_list[-1] != 151645:
                    input_ids.extend([151643,151645])
                newly_generated_text = processor.decode(newly_generated_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)
            else:
                newly_generated_text = ""
            # 关键改动：不要把 Tensor 直接 extend
            input_ids.extend(user_prefix)

            print("newly_generated:", [newly_generated_text])

            stop_token_str = "<|im_end|>"
            is_answer_finished = False
            chunk_to_process = newly_generated_text

            if stop_token_str in newly_generated_text:
                stop_index = newly_generated_text.index(stop_token_str)
                chunk_to_process = newly_generated_text[:stop_index]
                is_answer_finished = True
            if chunk_to_process:
                if final_answer_time is None:
                    final_answer_time = i+1
                # 关键改动：append 而不是 extend
                final_answer_text.append(chunk_to_process)
            if is_answer_finished and final_answer_text:
                final_answer = {
                    'time': final_answer_time,
                    'text': "".join(final_answer_text)
                }
                final_answers.append(final_answer)
                break

        # 关键改动：写盘前兜底提交（即使没有 <|im_end|> 也写当前累积）
        if (not final_answers) and final_answer_text:
            final_answers.append({
                'time': final_answer_time or len(multimodal_input_id_list),
                'text': "".join(final_answer_text)
            })

        with open(args.output_path, "a") as f:
            final_answer = {
                'id': data['id'],
                'question' : data['query'][0]['text'],
                'prediction': final_answers,
                "gt": data['answer']
            }
            f.write(json.dumps(final_answer, ensure_ascii=False) + "\n")   
        del data_formated, multimodal_input_id_list, mm_inputs, features, output, final_answers
        torch.cuda.empty_cache()