wehub-resource-sync
225 lines
11 KiB
Python
225 lines
11 KiB
Python
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
# coding: utf-8
|
||
import json
|
||
import torch
|
||
import random
|
||
import io
|
||
from einops import rearrange
|
||
from typing import List
|
||
import torch.nn.functional as F
|
||
import re
|
||
import numpy as np
|
||
|
||
RE_ZH = re.compile(r"[\u4e00-\u9fff]")
|
||
RE_EN = re.compile(r"[A-Za-z]")
|
||
|
||
def generate_system_prompt(system_prompt_type="caption", vision_type="video"):
|
||
if system_prompt_type == "caption":
|
||
str_list = [
|
||
f"Generate a detailed and accurate description of the {vision_type}, including all the key moments and visual details.",
|
||
f"Write an in-depth depiction of the {vision_type}, covering all its aspects.",
|
||
f"Write an exhaustive depiction of the given {vision_type}, capturing its essence and key moments.",
|
||
f"Describe the key features of the input {vision_type}, including color, shape, size, texture, objects, background.",
|
||
]
|
||
elif system_prompt_type == "t2v" or system_prompt_type == "i2v":
|
||
str_list = [f"Describe the {vision_type} by detailing the color, quantity, visible text, shape, size, texture, spatial relationships and motion/camera movements of the objects and background:"]
|
||
elif system_prompt_type == "t2i":
|
||
str_list = [f"Describe the {vision_type} by detailing the color, quantity, text, shape, size, texture, spatial relationships of the objects and background:"]
|
||
elif "edit" in system_prompt_type:
|
||
str_list = [f"Describe the key features of the input {vision_type} (color, shape, size, texture, objects, background), then explain how the user’s text instruction should alter or modify the {vision_type}. Generate a new {vision_type} that meets the user’s requirements while maintaining consistency with the original input where appropriate."]
|
||
elif "idip" in system_prompt_type:
|
||
str_list = [f"Describe the key features of the input image (color, shape, size, texture, objects, background, style), then incorporate the user’s text description to generate a new {vision_type} that satisfies the user’s requirements while preserving the essential identity and object or style information from the reference input."]
|
||
elif 'maze' in system_prompt_type:
|
||
str_list = [
|
||
"Describe the key elements of the input maze image (layout, white path, black walls, blue star, red flag, and overall background), then generate a 2D animation. The blue star should slide smoothly along the white path, stop exactly on the red flag, and then acquire a trophy. Ensure the blue star never crosses or enters the black maze walls. Keep the camera as a static top-down view showing the entire maze."
|
||
]
|
||
|
||
return random.choice(str_list)
|
||
|
||
|
||
def shift_position_ids(
|
||
position_ids: torch.Tensor,
|
||
pos_shift: any,
|
||
attn_modes: List[str],
|
||
split_lens: int,
|
||
shift_attn_mode=["full_noise", "full"],
|
||
pro_type=None,
|
||
i_sample_task=None,
|
||
i_sample_modality=None,
|
||
) -> torch.Tensor:
|
||
curr_split = 0
|
||
for i, attn_mode in enumerate(attn_modes):
|
||
if attn_mode in shift_attn_mode:
|
||
if pro_type == 10: # Related to sample_modality.
|
||
if position_ids[:, :, i_sample_modality == 4].sum() != 0:
|
||
pos_shift_type4 = 1000 - position_ids[:, :, i_sample_modality == 4][0, 0, 0]
|
||
position_ids[0, :, i_sample_modality == 4] += pos_shift_type4
|
||
if position_ids[:, :, i_sample_modality == 3].sum() != 0:
|
||
pos_shift_type3 = 2000 - position_ids[:, :, i_sample_modality == 3][0, 0, 0]
|
||
position_ids[0, :, i_sample_modality == 3] += pos_shift_type3
|
||
if position_ids[:, :, i_sample_modality == 2].sum() != 0 and sum(i_sample_modality == 2) == sum(i_sample_modality == 1):
|
||
position_ids[:, :, i_sample_modality == 1] = position_ids[:, :, i_sample_modality == 2]
|
||
|
||
curr_split += split_lens[i]
|
||
|
||
return position_ids
|
||
|
||
def detect_lang_simple(s: str) -> str:
|
||
"""
|
||
Fast heuristic: return 'zh' if Chinese is present, 'en' if English letters are present, otherwise 'other'.
|
||
Useful for quick routing. If both are present, this returns 'zh'; adjust if needed.
|
||
"""
|
||
# Remove digits before detection
|
||
s_without_digits = re.sub(r'\d+', '', s)
|
||
|
||
if RE_ZH.search(s_without_digits):
|
||
return "zh"
|
||
if RE_EN.search(s_without_digits):
|
||
return "en"
|
||
return "other"
|
||
|
||
def map_to_nearest_aspect_ratio(h, w, target_resolution=256):
|
||
"""
|
||
Map h and w to the closest preset aspect ratio and return adjusted h and w near the target resolution.
|
||
|
||
Preset ratios: ["21:9", "16:9", "4:3", "1:1", "3:4", "9:16"].
|
||
target_resolution: Base target resolution, default 256.
|
||
"""
|
||
# Precompute all preset aspect ratios as width / height
|
||
PRESET_RATIOS = [21 / 9, 16 / 9, 4 / 3, 1 / 1, 3 / 4, 9 / 16]
|
||
|
||
# Compute the original aspect ratio
|
||
original_ratio = w / h
|
||
|
||
# Find the closest preset ratio
|
||
min_index = min(range(len(PRESET_RATIOS)), key=lambda i: abs(original_ratio - PRESET_RATIOS[i]))
|
||
best_ratio = PRESET_RATIOS[min_index]
|
||
|
||
# Compute scale so the longer side is close to the target resolution
|
||
if best_ratio >= 1: # Landscape: width >= height
|
||
scale = target_resolution / best_ratio
|
||
adjusted_w = round(target_resolution)
|
||
adjusted_h = round(scale)
|
||
else: # Portrait: height > width
|
||
scale = target_resolution
|
||
adjusted_h = round(target_resolution)
|
||
adjusted_w = round(scale * best_ratio)
|
||
|
||
return adjusted_h, adjusted_w
|
||
|
||
|
||
def concat_resize_tensor_list(video_latents: List[torch.Tensor], dim: int = 0, is_offline: bool = False, max_num_frames: int = 121) -> torch.Tensor:
|
||
"""
|
||
Concatenate tensors along dim; resize H/W of tensors with different sizes to match target.
|
||
- tensors: Non-empty list; all tensors must have the same ndim.
|
||
- dim: Concatenation axis; negative values are supported.
|
||
- pad_value: Padding value, default 0.0.
|
||
Returns: Concatenated tensor.
|
||
"""
|
||
|
||
if is_offline:
|
||
H, W = video_latents[-1].shape[-3], video_latents[-1].shape[-2]
|
||
else:
|
||
H, W = video_latents[-1].shape[-2], video_latents[-1].shape[-1]
|
||
|
||
padded_video_latents = []
|
||
num_frames_target = video_latents[-1].shape[dim]
|
||
|
||
num_frames_all = num_frames_target
|
||
for index, video_latent in enumerate(video_latents):
|
||
if index != len(video_latents) - 1 and num_frames_all + video_latent.shape[dim] > max_num_frames: # Avoid producing videos longer than MAX_NUM_FRAMES
|
||
continue
|
||
num_frames_all += video_latent.shape[dim]
|
||
if is_offline:
|
||
# video_latent:[t,h,w,c] -> [t,c,h,w]
|
||
video_latent = rearrange(video_latent, "t h w c -> t c h w")
|
||
|
||
h, w = video_latent.shape[-2], video_latent.shape[-1]
|
||
if h != H or w != W:
|
||
video_latent = F.interpolate(video_latent, size=(H, W), mode="bilinear", align_corners=False)
|
||
padded_video_latents.append(video_latent)
|
||
|
||
padded_video_latents = torch.cat(padded_video_latents, dim=dim)
|
||
|
||
if is_offline:
|
||
# padded_video_latents: [t,c,h,w] -> [t,h,w,c]
|
||
padded_video_latents = rearrange(padded_video_latents, "t c h w -> t h w c")
|
||
|
||
return padded_video_latents
|
||
|
||
|
||
def concat_pad_tensor_list(video_latents: List[torch.Tensor], dim: int = 0, pad_value: float = 0.0, max_num_frames: int = 121) -> torch.Tensor:
|
||
"""
|
||
Concatenate tensors along dim; pad other axes to the maximum length on each axis with pad_value.
|
||
- tensors: Non-empty list; all tensors must have the same ndim.
|
||
- dim: Concatenation axis; negative values are supported.
|
||
- pad_value: Padding value, default 0.0.
|
||
Returns: Concatenated tensor.
|
||
"""
|
||
video_sizes = [item.shape for item in video_latents]
|
||
max_video_size = [max(item) for item in list(zip(*video_sizes))]
|
||
padded_video_latents = []
|
||
num_frames_target = video_latents[-1].shape[dim]
|
||
|
||
num_frames_all = num_frames_target
|
||
for index, video_latent in enumerate(video_latents):
|
||
if index != len(video_latents) - 1 and num_frames_all + video_latent.shape[dim] > max_num_frames: # Avoid producing videos longer than MAX_NUM_FRAMES
|
||
continue
|
||
num_frames_all += video_latent.shape[dim]
|
||
max_video_size[dim] = video_latent.shape[dim]
|
||
padded_video_latent = torch.zeros(max_video_size)
|
||
n1, n2, n3, n4 = video_latent.shape
|
||
padded_video_latent[:n1, :n2, :n3, :n4] = video_latent
|
||
padded_video_latents.append(padded_video_latent)
|
||
|
||
padded_video_latents = torch.cat(padded_video_latents, dim=dim)
|
||
|
||
return padded_video_latents
|
||
|
||
|
||
def parse_videochat2it_doubao_caption(row):
|
||
try:
|
||
IQA_i = "View the video attentively and provide a suitable answer to the posed question."
|
||
rewrite_VQA = json.loads(row['rewrite_VQA'])
|
||
|
||
IQA_q = rewrite_VQA['question'] if 'question' in rewrite_VQA.keys() else rewrite_VQA['Question']
|
||
IQA_a = rewrite_VQA['final_answer']
|
||
IQA_resoning = rewrite_VQA['reasoning']
|
||
# Combine reasoning and final_answer as the final answer
|
||
|
||
# if random.random() < 0.5: # 50% chance to include the reasoning process
|
||
# IQA_a = IQA_a + '\n' + IQA_resoning
|
||
# IQA_i = IQA_i + ' Please provide the reasoning process for selecting the correct answer.'
|
||
|
||
if 'options' not in IQA_q and 'Options' not in IQA_q: # When the question does not contain options
|
||
try:
|
||
options = rewrite_VQA['options']
|
||
except:
|
||
options = rewrite_VQA['Options']
|
||
if options == []:
|
||
return [IQA_i, IQA_q, IQA_a]
|
||
elif isinstance(options,list):
|
||
options = '\n'.join(options)
|
||
elif isinstance(options,dict):
|
||
options_str = [key + ' ' + value if value not in key else key for key,value in options.items()]
|
||
options = '\n'.join(options_str)
|
||
|
||
IQA_q = IQA_q + '\nOptions:\n' + options # Add options to the question
|
||
return [IQA_i, IQA_q, IQA_a]
|
||
except:
|
||
if 'rewrite_VQA' in row.keys():
|
||
raise ValueError(f"wrong rewrite_VQA in {row['rewrite_VQA']}")
|
||
else:
|
||
raise ValueError(f"wrong rewrite_VQA in {row}") |