Views
No views yet
1
2import os
3from PIL import Image
4import numpy as np
5import torchaudio
6import torch
7from decord import VideoReader, cpu
8import whisper
9# fix seed
10torch.manual_seed(0)
11
12from intersuit.model.builder import load_pretrained_model
13from intersuit.mm_utils import tokenizer_image_speech_tokens, process_images
14from intersuit.constants import IMAGE_TOKEN_INDEX, SPEECH_TOKEN_INDEX
15
16import warnings
17warnings.filterwarnings("ignore")
18
19model_path = "ColorfulAI/LongVA-7B-Qwen2-Audio"
20video_path = "local_demo/assets/water.mp4"
21audio_path = "local_demo/wav/infer.wav"
22max_frames_num = 16 # you can change this to several thousands so long you GPU memory can handle it :)
23gen_kwargs = {"do_sample": True, "temperature": 0.5, "top_p": None, "num_beams": 1, "use_cache": True, "max_new_tokens": 1024}
24tokenizer, model, image_processor, _ = load_pretrained_model(model_path, None, "llava_qwen", device_map="cuda:0")
25
26query = "Give a detailed caption of the video as if I am blind."
27query = None # comment this to use ChatTTS to convert the query to audio
28
29#video input
30prompt = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<image><|im_end|>\n<|im_start|>user\n<speech>\n<|im_end|>\n<|im_start|>assistant\n"
31input_ids = tokenizer_image_speech_tokens(prompt, tokenizer, IMAGE_TOKEN_INDEX, SPEECH_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to(model.device)
32vr = VideoReader(video_path, ctx=cpu(0))
33total_frame_num = len(vr)
34uniform_sampled_frames = np.linspace(0, total_frame_num - 1, max_frames_num, dtype=int)
35frame_idx = uniform_sampled_frames.tolist()
36frames = vr.get_batch(frame_idx).asnumpy()
37video_tensor = image_processor.preprocess(frames, return_tensors="pt")["pixel_values"].to(model.device, dtype=torch.float16)
38
39#audio input
40# process speech for input question
41if query is not None:
42 import ChatTTS
43 chat = ChatTTS.Chat()
44 chat.load(source='local', compile=True)
45 audio_path = "./local_demo/wav/" + "infer.wav"
46 if os.path.exists(audio_path): os.remove(audio_path) # refresh
47 if not os.path.exists(audio_path):
48 wav = chat.infer(query)
49 try:
50 torchaudio.save(audio_path, torch.from_numpy(wav).unsqueeze(0), 24000)
51 except:
52 torchaudio.save(audio_path, torch.from_numpy(wav), 24000)
53 print(f"Human: {query}")
54
55else:
56 print("Human: <audio>")
57
58speech = whisper.load_audio(audio_path)
59speech = whisper.pad_or_trim(speech)
60speech = whisper.log_mel_spectrogram(speech, n_mels=128).permute(1, 0).to(device=model.device, dtype=torch.float16)
61speech_length = torch.LongTensor([speech.shape[0]]).to(model.device)
62
63with torch.inference_mode():
64 output_ids = model.generate(input_ids, images=[video_tensor], modalities=["video"], speeches=speech.unsqueeze(0), speech_lengths=speech_length, **gen_kwargs)
65outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
66print(f"Agent: {outputs}")
67