| import base64 |
| import logging |
| import math |
| import os |
| from io import BytesIO |
| from tqdm.contrib.concurrent import thread_map |
|
|
| import numpy as np |
|
|
| import requests |
| import torch |
|
|
| from PIL import Image |
| import torchaudio |
| from typing import Union, Tuple, List |
|
|
| VIDEO_FETCH_VERSION = os.environ.get("VIDEO_FETCH_VERSION", "v1") |
| if VIDEO_FETCH_VERSION == "v1": |
| from .bailingmm_utils_video import v1_fetch_video as fetch_video |
| else: |
| from .bailingmm_utils_video import v2_fetch_video as fetch_video |
| from .bailingmm_utils_video import VideoInput |
|
|
| logger = logging.getLogger(__name__) |
|
|
| IMAGE_FACTOR = 32 |
| MIN_PIXELS = 4 * 32 * 32 |
| MAX_PIXELS = 16384 * 32 * 32 |
| MAX_RATIO = 200 |
|
|
| VideoInput = Union[ |
| List["Image.Image"], |
| "np.ndarray", |
| "torch.Tensor", |
| List["np.ndarray"], |
| List["torch.Tensor"], |
| List[List["Image.Image"]], |
| List[List["np.ndarrray"]], |
| List[List["torch.Tensor"]], |
| ] |
|
|
|
|
| def round_by_factor(number: int, factor: int) -> int: |
| """Returns the closest integer to 'number' that is divisible by 'factor'.""" |
| return round(number / factor) * factor |
|
|
| def ceil_by_factor(number: int, factor: int) -> int: |
| """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'.""" |
| return math.ceil(number / factor) * factor |
|
|
| def floor_by_factor(number: int, factor: int) -> int: |
| """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'.""" |
| return math.floor(number / factor) * factor |
|
|
| def is_image(image_file): |
| if isinstance(image_file, str) and (image_file.startswith("base64,") or image_file.lower().endswith( |
| ('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff'))): |
| return True |
| elif isinstance(image_file, Image.Image): |
| return True |
| else: |
| return False |
|
|
| def is_video(video_file): |
| if isinstance(video_file, str) and video_file.lower().endswith( |
| ('.mp4', '.mkv', '.avi', '.wmv', '.iso', ".webm")): |
| return True |
| else: |
| return False |
|
|
| def is_audio(audio_file): |
| if isinstance(audio_file, str) and audio_file.lower().endswith( |
| (".wav", ".mp3", ".aac", ".flac", ".alac", ".m4a", ".ogg", ".wma", ".aiff", ".amr", ".au")): |
| return True |
| else: |
| return False |
|
|
| def smart_resize( |
| height: int, width: int, factor: int = IMAGE_FACTOR, min_pixels: int = MIN_PIXELS, max_pixels: int = MAX_PIXELS |
| ) -> tuple[int, int]: |
| """ |
| Rescales the image so that the following conditions are met: |
| |
| 1. Both dimensions (height and width) are divisible by 'factor'. |
| |
| 2. The total number of pixels is within the range ['min_pixels', 'max_pixels']. |
| |
| 3. The aspect ratio of the image is maintained as closely as possible. |
| """ |
| if max(height, width) / min(height, width) > MAX_RATIO: |
| raise ValueError( |
| f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}" |
| ) |
| h_bar = max(factor, round_by_factor(height, factor)) |
| w_bar = max(factor, round_by_factor(width, factor)) |
| if h_bar * w_bar > max_pixels: |
| beta = math.sqrt((height * width) / max_pixels) |
| h_bar = floor_by_factor(height / beta, factor) |
| w_bar = floor_by_factor(width / beta, factor) |
| elif h_bar * w_bar < min_pixels: |
| beta = math.sqrt(min_pixels / (height * width)) |
| h_bar = ceil_by_factor(height * beta, factor) |
| w_bar = ceil_by_factor(width * beta, factor) |
| return h_bar, w_bar |
|
|
| def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image: |
| if "image" in ele: |
| image = ele["image"] |
| else: |
| image = ele["image_url"] |
| image_obj = None |
| if isinstance(image, Image.Image): |
| image_obj = image |
| elif image.startswith("http://") or image.startswith("https://"): |
| image_obj = Image.open(requests.get(image, stream=True).raw) |
| elif image.startswith("file://"): |
| image_obj = Image.open(image[7:]) |
| elif image.startswith("data:image"): |
| if "base64," in image: |
| _, base64_data = image.split("base64,", 1) |
| data = base64.b64decode(base64_data) |
| image_obj = Image.open(BytesIO(data)) |
| else: |
| image_obj = Image.open(image) |
| if image_obj is None: |
| raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}") |
| image = image_obj.convert("RGB") |
| |
| if "resized_height" in ele and "resized_width" in ele: |
| resized_height, resized_width = smart_resize( |
| ele["resized_height"], |
| ele["resized_width"], |
| factor=size_factor, |
| ) |
| else: |
| width, height = image.size |
| min_pixels = ele.get("min_pixels", MIN_PIXELS) |
| max_pixels = ele.get("max_pixels", MAX_PIXELS) |
| resized_height, resized_width = smart_resize( |
| height, |
| width, |
| factor=size_factor, |
| min_pixels=min_pixels, |
| max_pixels=max_pixels, |
| ) |
| image = image.resize((resized_width, resized_height)) |
|
|
| return image |
|
|
| def fetch_image_wo_resize(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image: |
| if "image" in ele: |
| image = ele["image"] |
| else: |
| image = ele["image_url"] |
| image_obj = None |
| if isinstance(image, Image.Image): |
| image_obj = image |
| elif image.startswith("http://") or image.startswith("https://"): |
| image_obj = Image.open(requests.get(image, stream=True).raw) |
| elif image.startswith("file://"): |
| image_obj = Image.open(image[7:]) |
| elif image.startswith("data:image"): |
| if "base64," in image: |
| _, base64_data = image.split("base64,", 1) |
| data = base64.b64decode(base64_data) |
| image_obj = Image.open(BytesIO(data)) |
| else: |
| image_obj = Image.open(image) |
| if image_obj is None: |
| raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}") |
| image = image_obj.convert("RGB") |
|
|
| return image |
|
|
| def fetch_audio(ele: dict[str, str | torch.Tensor], return_tensor="pt") -> Tuple[Union[torch.Tensor, np.ndarray], int]: |
| if "audio" in ele: |
| audio = ele["audio"] |
| else: |
| audio = ele["audio_url"] |
|
|
| if isinstance(audio, torch.Tensor): |
| waveform = audio |
| sample_rate: int = ele.get("sample_rate", 16000) |
| elif audio.startswith("http://") or audio.startswith("https://"): |
| audio_file = BytesIO(requests.get(audio, stream=True).content) |
| waveform, sample_rate = torchaudio.load(audio_file) |
| elif audio.startswith("file://"): |
| waveform, sample_rate = torchaudio.load(audio[7:]) |
| else: |
| waveform, sample_rate = torchaudio.load(audio) |
| if return_tensor == "pt": |
| return waveform, sample_rate |
| else: |
| return waveform.numpy(), sample_rate |
|
|
| def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]: |
| vision_infos = [] |
| if isinstance(conversations[0], dict): |
| conversations = [conversations] |
| for conversation in conversations: |
| for message in conversation: |
| if isinstance(message["content"], list): |
| for ele in message["content"]: |
| if ( |
| "image" in ele |
| or "image_url" in ele |
| or "video" in ele |
| or "video_url" in ele |
| or "audio" in ele |
| or "audio_url" in ele |
| or ele["type"] in ["image", "image_url", "video", "video_url", "audio", "audio_url"] |
| ): |
| vision_infos.append(ele) |
| return vision_infos |
|
|
|
|
| def process_reference_vision_info( |
| conversations: list[dict] | list[list[dict]], |
| ) -> list[Image.Image] | None: |
| vision_infos = extract_vision_info(conversations) |
| |
| image_inputs = [] |
|
|
| def inner_process_func(vision_info): |
| if "image" in vision_info or "image_url" in vision_info: |
| res_list = [] |
| if "image" in vision_info and isinstance(vision_info["image"], (tuple, list)): |
| for i in range(len(vision_info["image"])): |
| res_list.append(fetch_image_wo_resize({"type": "image", "image": vision_info["image"][i]})) |
| elif "image_url" in vision_info and vision_info["image_url"].get("url", None) is not None: |
| vision_info["image_url"] = vision_info["image_url"].get("url") |
| res_list.extend([fetch_image_wo_resize(vision_info)]) |
| else: |
| res_list.extend([fetch_image_wo_resize(vision_info)]) |
| return {'image_inputs':res_list} |
| else: |
| return None |
|
|
| vision_infos_reslist = thread_map(inner_process_func, vision_infos, disable=True) |
| for res in vision_infos_reslist: |
| if res is None: |
| raise ValueError("image, image_url, video, video_url, audio or audio_url should in content.") |
| elif 'image_inputs' in res: |
| image_inputs.extend(res['image_inputs']) |
|
|
| if len(image_inputs) > 1: |
| image_inputs = [image_inputs[0]] |
|
|
| if len(image_inputs) == 0: |
| image_inputs = None |
|
|
| return image_inputs |
|
|
|
|
| def process_vision_info( |
| conversations: list[dict] | list[list[dict]], |
| ) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None, list[ |
| torch.Tensor | list[np.ndarray]] | None]: |
| vision_infos = extract_vision_info(conversations) |
| |
| image_inputs = [] |
| video_inputs = [] |
| audio_inputs = [] |
|
|
| def inner_process_func(vision_info): |
| if "image" in vision_info or "image_url" in vision_info: |
| res_list = [] |
| if "image" in vision_info and isinstance(vision_info["image"], (tuple, list)): |
| for i in range(len(vision_info["image"])): |
| res_list.append(fetch_image({"type": "image", "image": vision_info["image"][i]})) |
| elif "image_url" in vision_info and vision_info["image_url"].get("url", None) is not None: |
| vision_info["image_url"] = vision_info["image_url"].get("url") |
| res_list.extend([fetch_image(vision_info)]) |
| else: |
| res_list.extend([fetch_image(vision_info)]) |
| return {'image_inputs':res_list} |
|
|
| elif "video" in vision_info or "video_url" in vision_info: |
| if "video_url" in vision_info and vision_info["video_url"].get("url", None) is not None: |
| data_value = vision_info["video_url"].get("url") |
| elif "video" in vision_info and not os.path.isdir(vision_info['video']): |
| data_value = vision_info['video'] |
| else: |
| data_value = [os.path.join(vision_info['video'], frame) for frame in sorted(os.listdir(vision_info['video']))] |
| vision_info['video']=data_value |
| return {"video_inputs": [fetch_video(vision_info, return_metadata=True)]} |
|
|
| elif "audio" in vision_info or "audio_url" in vision_info: |
| if "audio" in vision_info and isinstance(vision_info["audio"], (tuple, list)): |
| return {"audio_inputs":[fetch_audio(info) for info in vision_info["audio"]]} |
| elif "audio_url" in vision_info and vision_info["audio_url"].get("url", None) is not None: |
| vision_info["audio_url"] = vision_info["audio_url"].get("url") |
| return {"audio_inputs":[fetch_audio(vision_info)]} |
| else: |
| return {"audio_inputs":[fetch_audio(vision_info)]} |
| else: |
| return None |
|
|
| vision_infos_reslist = thread_map(inner_process_func, vision_infos, disable=True) |
| for res in vision_infos_reslist: |
| if res is None: |
| raise ValueError("image, image_url, video, video_url, audio or audio_url should in content.") |
| elif 'image_inputs' in res: |
| image_inputs.extend(res['image_inputs']) |
| elif 'video_inputs' in res: |
| video_inputs.extend(res['video_inputs']) |
| elif 'audio_inputs' in res: |
| audio_inputs.extend(res['audio_inputs']) |
|
|
| if len(image_inputs) == 0: |
| image_inputs = None |
| if len(video_inputs) == 0: |
| video_inputs = None |
| if len(audio_inputs) == 0: |
| audio_inputs = None |
| return image_inputs, video_inputs, audio_inputs |
|
|
|
|
| def get_closest_ratio(height: float, width: float, aspect_ratios: dict): |
| aspect_ratio = height / width |
| closest_ratio = min(aspect_ratios.keys(), key=lambda ratio: abs(float(ratio) - aspect_ratio)) |
| return aspect_ratios[closest_ratio], float(closest_ratio) |
|
|
| def process_ratio(ori_h, ori_w, highres=512): |
| ASPECT_RATIO_512 = { |
| "0.25": [256, 1024], |
| "0.26": [256, 992], |
| "0.27": [256, 960], |
| "0.28": [256, 928], |
| "0.32": [288, 896], |
| "0.33": [288, 864], |
| "0.35": [288, 832], |
| "0.4": [320, 800], |
| "0.42": [320, 768], |
| "0.48": [352, 736], |
| "0.5": [352, 704], |
| "0.52": [352, 672], |
| "0.57": [384, 672], |
| "0.6": [384, 640], |
| "0.68": [416, 608], |
| "0.72": [416, 576], |
| "0.78": [448, 576], |
| "0.82": [448, 544], |
| "0.88": [480, 544], |
| "0.94": [480, 512], |
| "1.0": [512, 512], |
| "1.07": [512, 480], |
| "1.13": [544, 480], |
| "1.21": [544, 448], |
| "1.29": [576, 448], |
| "1.38": [576, 416], |
| "1.46": [608, 416], |
| "1.67": [640, 384], |
| "1.75": [672, 384], |
| "2.0": [704, 352], |
| "2.09": [736, 352], |
| "2.4": [768, 320], |
| "2.5": [800, 320], |
| "2.89": [832, 288], |
| "3.0": [864, 288], |
| "3.11": [896, 288], |
| "3.62": [928, 256], |
| "3.75": [960, 256], |
| "3.88": [992, 256], |
| "4.0": [1024, 256], |
| } |
| ASPECT_RATIO_1024 = { |
| '0.25': [512, 2048], '0.26': [512, 1984], '0.27': [512, 1920], '0.28': [512, 1856], |
| '0.32': [576, 1792], '0.33': [576, 1728], '0.35': [576, 1664], '0.4': [640, 1600], |
| '0.42': [640, 1536], '0.48': [704, 1472], '0.5': [704, 1408], '0.52': [704, 1344], |
| '0.57': [768, 1344], '0.6': [768, 1280], '0.68': [832, 1216], '0.72': [832, 1152], |
| '0.78': [896, 1152], '0.82': [896, 1088], '0.88': [960, 1088], '0.94': [960, 1024], |
| '1.0': [1024, 1024], '1.07': [1024, 960], '1.13': [1088, 960], '1.21': [1088, 896], |
| '1.29': [1152, 896], '1.38': [1152, 832], '1.46': [1216, 832], '1.67': [1280, 768], |
| '1.75': [1344, 768], '2.0': [1408, 704], '2.09': [1472, 704], '2.4': [1536, 640], |
| '2.5': [1600, 640], '2.89': [1664, 576], '3.0': [1728, 576], '3.11': [1792, 576], |
| '3.62': [1856, 512], '3.75': [1920, 512], '3.88': [1984, 512], '4.0': [2048, 512], |
| } |
|
|
| ASPECT_RATIO_672 = { |
| '0.28': [352, 1280], |
| '0.32': [384, 1184], |
| '0.38': [416, 1088], |
| '0.44': [448, 1024], |
| '0.52': [480, 928], |
| '0.57': [512, 896], |
| '0.65': [544, 832], |
| '0.75': [576, 768], |
| '0.83': [608, 736], |
| '0.91': [640, 704], |
| '1.00': [672, 672], |
| '1.10': [704, 640], |
| '1.21': [736, 608], |
| '1.33': [768, 576], |
| '1.39': [800, 576], |
| '1.53': [832, 544], |
| '1.69': [864, 512], |
| '1.75': [896, 512], |
| '1.93': [928, 480], |
| '2.00': [960, 480], |
| '2.21': [992, 448], |
| '2.29': [1024, 448], |
| '2.54': [1056, 416], |
| '2.62': [1088, 416], |
| '2.69': [1120, 416], |
| '3.00': [1152, 384], |
| '3.08': [1184, 384], |
| '3.17': [1216, 384], |
| '3.55': [1248, 352], |
| '3.64': [1280, 352], |
| } |
|
|
| assert len(ASPECT_RATIO_512) == len(ASPECT_RATIO_1024) |
|
|
| aspect_ratio_dict = { |
| 512 : ASPECT_RATIO_512, |
| 672 : ASPECT_RATIO_672, |
| 1024 : ASPECT_RATIO_1024, |
| } |
|
|
| if highres is None or highres is False: |
| highres = 512 |
| elif highres is True: |
| highres = 1024 |
|
|
| aspect_ratio = aspect_ratio_dict[min([i for i in aspect_ratio_dict], key=lambda x: abs(x - highres))] |
|
|
| closest_size, _ = get_closest_ratio(ori_h, ori_w, aspect_ratios=aspect_ratio) |
| closest_size = list(map(lambda x: int(x), closest_size)) |
| if closest_size[0] / ori_h > closest_size[1] / ori_w: |
| resize_size = closest_size[0], int(ori_w * closest_size[0] / ori_h) |
| else: |
| resize_size = int(ori_h * closest_size[1] / ori_w), closest_size[1] |
| return closest_size, resize_size |
|
|
|
|
| def find_first_index_of_consecutive_ones(lst): |
| """ |
| 输入一个由0和1组成的列表，返回每个连续1片段的第一个1的索引。 |
| |
| 参数: |
| lst (list): 元素为0或1的列表 |
| |
| 返回: |
| list: 每个连续1片段的首个1的索引列表 |
| """ |
| result = [] |
| i = 0 |
| n = len(lst) |
| |
| while i < n: |
| if lst[i] == 1: |
| |
| result.append(i) |
| |
| while i < n and lst[i] == 1: |
| i += 1 |
| else: |
| i += 1 |
| |
| return result |
|
|
| def merge_consecutive_ones(lst, n): |
| """ |
| 输入一个由0和1组成的列表，将每个连续的1片段（长度 >= 1）中每n个1合并为一个1， |
| 要求每个连续1片段的长度必须能被n整除。 |
| 保持0和1的相对顺序。 |
| |
| 参数: |
| lst: list, 元素为0或1 |
| n: int, 合并的单位大小（正整数） |
| |
| 返回: |
| list: 合并后的列表 |
| """ |
| assert isinstance(lst, list), "输入必须是列表" |
| assert isinstance(n, int) and n > 0, "n必须是正整数" |
| |
| |
| i = 0 |
| while i < len(lst): |
| if lst[i] == 1: |
| count = 0 |
| start = i |
| |
| while i < len(lst) and lst[i] == 1: |
| count += 1 |
| i += 1 |
| |
| assert count % n == 0, f"连续1的片段从索引{start}开始，长度为{count}，不能被n={n}整除" |
| else: |
| i += 1 |
|
|
| |
| result = [] |
| i = 0 |
| while i < len(lst): |
| if lst[i] == 0: |
| result.append(0) |
| i += 1 |
| else: |
| |
| count = 0 |
| while i < len(lst) and lst[i] == 1: |
| count += 1 |
| i += 1 |
| |
| result.extend([1] * (count // n)) |
| |
| return result |
|
|
| def get_default_image_gen_hw(image_gen_highres, image_gen_aspect_ratio): |
| if image_gen_aspect_ratio is None: |
| image_gen_aspect_ratio = 1.0 |
|
|
| closest_size, _ = process_ratio(ori_h=512, ori_w=int(512.0 * image_gen_aspect_ratio), highres=image_gen_highres) |
| h, w = closest_size |
| return h, w |
|
|
