1import numpy as np
2import torch
3import torchvision.transforms as T
4from PIL import Image
5from torchvision.transforms.functional import InterpolationMode
6from transformers import AutoModel, AutoTokenizer
7import librosa
8from transformers.processing_utils import ProcessorMixin
9import torch
10
11class WhisperProcessor(ProcessorMixin):
12 attributes = ["feature_extractor"]
13 feature_extractor_class = "WhisperFeatureExtractor"
14 def __init__(self, feature_extractor):
15 super().__init__(feature_extractor)
16 self.current_processor = self.feature_extractor
17 self._in_target_context_manager = False
18
19 def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True):
20 return self.tokenizer.get_decoder_prompt_ids(task=task, language=language, no_timestamps=no_timestamps)
21
22 def get_T_after_cnn(self,L_in, dilation=1):
23 for (padding, kernel_size, stride) in eval("[(1,3,1)] + [(1,3,2)] "):
24 L_out = L_in + 2 * padding - dilation * (kernel_size - 1) - 1
25 L_out = 1 + L_out // stride
26 L_in = L_out
27 return L_out
28
29 def __call__(self, *args, **kwargs):
30 if self._in_target_context_manager:
31 return self.current_processor(*args, **kwargs)
32
33 audio = kwargs.pop("audio", None)
34 sampling_rate = kwargs.pop("sampling_rate", 16000)
35 text = kwargs.pop("text", None)
36 if len(args) > 0:
37 audio = args[0]
38 args = args[1:]
39
40 if audio is None and text is None:
41 raise ValueError("You need to specify either an `audio` or `text` input to process.")
42
43 if audio is not None:
44 L = (audio.shape[0] if audio.shape[0] <= 480000 else 480000) # max_length < 30s
45 mel_len = L // 160
46 audio_len_after_cnn = self.get_T_after_cnn(mel_len)
47 audio_token_num = (audio_len_after_cnn - 2) // 2 + 1
48 inputs = self.feature_extractor(audio, *args, sampling_rate=sampling_rate, **kwargs)
49 inputs['audio_len_after_cnn'] = torch.tensor(audio_len_after_cnn, dtype=torch.long)
50 inputs['audio_token_num'] = torch.tensor(audio_token_num, dtype=torch.long)
51 if text is not None:
52 encodings = self.tokenizer(text, **kwargs)
53
54 if text is None:
55 return inputs
56
57 elif audio is None:
58 return encodings
59 else:
60 inputs["labels"] = encodings["input_ids"]
61 return inputs
62
63 def batch_decode(self, *args, **kwargs):
64 return self.tokenizer.batch_decode(*args, **kwargs)
65
66 def decode(self, *args, **kwargs):
67 return self.tokenizer.decode(*args, **kwargs)
68
69 def get_prompt_ids(self, text: str, return_tensors="np"):
70 return self.tokenizer.get_prompt_ids(text, return_tensors=return_tensors)
71
72IMAGENET_MEAN = (0.485, 0.456, 0.406)
73IMAGENET_STD = (0.229, 0.224, 0.225)
74
75def build_transform(input_size):
76 MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
77 transform = T.Compose([
78 T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
79 T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
80 T.ToTensor(),
81 T.Normalize(mean=MEAN, std=STD)
82 ])
83 return transform
84
85def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
86 best_ratio_diff = float('inf')
87 best_ratio = (1, 1)
88 area = width * height
89 for ratio in target_ratios:
90 target_aspect_ratio = ratio[0] / ratio[1]
91 ratio_diff = abs(aspect_ratio - target_aspect_ratio)
92 if ratio_diff < best_ratio_diff:
93 best_ratio_diff = ratio_diff
94 best_ratio = ratio
95 elif ratio_diff == best_ratio_diff:
96 if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
97 best_ratio = ratio
98 return best_ratio
99
100def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
101 orig_width, orig_height = image.size
102 aspect_ratio = orig_width / orig_height
103
104 # calculate the existing image aspect ratio
105 target_ratios = set(
106 (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
107 i * j <= max_num and i * j >= min_num)
108 target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
109
110 # find the closest aspect ratio to the target
111 target_aspect_ratio = find_closest_aspect_ratio(
112 aspect_ratio, target_ratios, orig_width, orig_height, image_size)
113
114 # calculate the target width and height
115 target_width = image_size * target_aspect_ratio[0]
116 target_height = image_size * target_aspect_ratio[1]
117 blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
118
119 # resize the image
120 resized_img = image.resize((target_width, target_height))
121 processed_images = []
122 for i in range(blocks):
123 box = (
124 (i % (target_width // image_size)) * image_size,
125 (i // (target_width // image_size)) * image_size,
126 ((i % (target_width // image_size)) + 1) * image_size,
127 ((i // (target_width // image_size)) + 1) * image_size
128 )
129 # split the image
130 split_img = resized_img.crop(box)
131 processed_images.append(split_img)
132 assert len(processed_images) == blocks
133 if use_thumbnail and len(processed_images) != 1:
134 thumbnail_img = image.resize((image_size, image_size))
135 processed_images.append(thumbnail_img)
136 return processed_images
137
138def load_image(image_file, input_size=448, max_num=12):
139 image = Image.open(image_file).convert('RGB')
140 transform = build_transform(input_size=input_size)
141 images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
142 pixel_values = [transform(image) for image in images]
143 pixel_values = torch.stack(pixel_values)
144 return pixel_values
145
146def load_audio(audio_file, audio_processor):
147 audio_values, _ = librosa.load(audio_file, sr=16000) # sample rate should be 16000
148
149 audio_process_values = audio_processor(audio_values, sampling_rate=16000, return_tensors="pt")
150 input_features = audio_process_values['input_features']
151 audio_len_after_cnn = audio_process_values['audio_len_after_cnn']
152 audio_token_num = audio_process_values['audio_token_num']
153
154
155 audio_input = {'audio_values': input_features,
156 'audio_len_after_cnn': audio_len_after_cnn,
157 'audio_token_num': audio_token_num,
158 }
159 return audio_input
160
161path = 'OpenGVLab/InternOmni'
162model = AutoModel.from_pretrained(
163 path,
164 torch_dtype=torch.bfloat16,
165 low_cpu_mem_usage=True,
166 trust_remote_code=True).eval().cuda()
167tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True, use_fast=False)
168audio_processor = WhisperProcessor.from_pretrained(path)
169# set the max number of tiles in `max_num`
170pixel_values = load_image('./1.jpg', max_num=12).to(torch.bfloat16).cuda()
171audio = load_audio('./1.wav', audio_processor)
172generation_config = dict(max_new_tokens=1024, do_sample=True)
173
174# question = '请将这段语音识别成文字，并以文字形式展示出来。'
175response = model.Audio_chat(tokenizer=tokenizer, pixel_values=pixel_values,audio=audio, question=None, generation_config)
176print(f'Assistant: {response}')
177
