已合并
feat: add generative data utils #2924
chenpeizhe创建于 8月4日
feat: add generative data utils #2924
已合并
共 2 个文件变更+964-0
| @@ -0,0 +1,330 @@ | |||
| 1 | +# Modified from Latte (datasets/video_transforms.py) | ||
| 2 | + | ||
| 3 | +import random | ||
| 4 | + | ||
| 5 | +import torch | ||
| 6 | +import numpy as np | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +def _is_tensor_video_clip(clip): | ||
| 10 | + if not torch.is_tensor(clip): | ||
| 11 | + raise TypeError("Clip should be Tensor, but it is %s" % type(clip)) | ||
| 12 | + | ||
| 13 | + if not clip.ndimension() == 4: | ||
| 14 | + raise ValueError("Clip should be 4D, but it is %dD" % clip.dim()) | ||
| 15 | + | ||
| 16 | + return True | ||
| 17 | + | ||
| 18 | + | ||
| 19 | +def to_tensor(clip): | ||
| 20 | + """Convert clip tensor (T, C, H, W) from uint8 to float (divide by 255.0) and permute dims.""" | ||
| 21 | + _is_tensor_video_clip(clip) | ||
| 22 | + if not clip.dtype == torch.uint8: | ||
| 23 | + raise TypeError( | ||
| 24 | + "Clip tensor should have data type uint8, but it is %s" % str(clip.dtype) | ||
| 25 | + ) | ||
| 26 | + | ||
| 27 | + return clip.float() / 255.0 | ||
| 28 | + | ||
| 29 | + | ||
| 30 | +def crop(clip, i, j, h, w): | ||
| 31 | + """Crop video clip tensor (T, C, H, W).""" | ||
| 32 | + if len(clip.size()) != 4: | ||
| 33 | + raise ValueError("clip should be a 4D tensor") | ||
| 34 | + return clip[..., i: i + h, j: j + w] | ||
| 35 | + | ||
| 36 | + | ||
| 37 | +def resize(clip, target_size, interpolation_mode, align_corners=False, antialias=False): | ||
| 38 | + if len(target_size) != 2: | ||
| 39 | + raise ValueError( | ||
| 40 | + f"target size should be tuple (height, width), instead got {target_size}" | ||
| 41 | + ) | ||
| 42 | + return torch.nn.functional.interpolate( | ||
| 43 | + clip, | ||
| 44 | + size=target_size, | ||
| 45 | + mode=interpolation_mode, | ||
| 46 | + align_corners=align_corners, | ||
| 47 | + antialias=antialias, | ||
| 48 | + ) | ||
| 49 | + | ||
| 50 | + | ||
| 51 | +def center_crop_using_short_edge(clip): | ||
| 52 | + if not _is_tensor_video_clip(clip): | ||
| 53 | + raise ValueError("clip should be a 4D torch.tensor") | ||
| 54 | + h, w = clip.size(-2), clip.size(-1) | ||
| 55 | + if h < w: | ||
| 56 | + th, tw = h, h | ||
| 57 | + i = 0 | ||
| 58 | + j = int(round((w - tw) / 2.0)) | ||
| 59 | + else: | ||
| 60 | + th, tw = w, w | ||
| 61 | + i = int(round((h - th) / 2.0)) | ||
| 62 | + j = 0 | ||
| 63 | + return crop(clip, i, j, th, tw) | ||
| 64 | + | ||
| 65 | + | ||
| 66 | +def center_crop_th_tw(clip, th, tw, top_crop): | ||
| 67 | + if not _is_tensor_video_clip(clip): | ||
| 68 | + raise ValueError("clip should be a 4D torch.tensor") | ||
| 69 | + | ||
| 70 | + h, w = clip.size(-2), clip.size(-1) | ||
| 71 | + tr = th / tw | ||
| 72 | + if h / w > tr: | ||
| 73 | + new_h = int(w * tr) | ||
| 74 | + new_w = w | ||
| 75 | + else: | ||
| 76 | + new_h = h | ||
| 77 | + new_w = int(h / tr) | ||
| 78 | + | ||
| 79 | + i = 0 if top_crop else int(round((h - new_h) / 2.0)) | ||
| 80 | + j = int(round((w - new_w) / 2.0)) | ||
| 81 | + return crop(clip, i, j, new_h, new_w) | ||
| 82 | + | ||
| 83 | + | ||
| 84 | +def longsideresize(h, w, size, skip_low_resolution): | ||
| 85 | + if h <= size[0] and w <= size[1] and skip_low_resolution: | ||
| 86 | + return h, w | ||
| 87 | + | ||
| 88 | + if h / w > size[0] / size[1]: | ||
| 89 | + # hxw 720x1280 size 320x640 hw_raito 9/16 > size_ratio 8/16 neww=320/720*1280=568 newh=320 | ||
| 90 | + w = int(size[0] / h * w) | ||
| 91 | + h = size[0] | ||
| 92 | + else: | ||
| 93 | + # hxw 720x1280 size 480x640 hw_raito 9/16 < size_ratio 12/16 newh=640/1280*720=360 neww=640 | ||
| 94 | + # hxw 1080x1920 size 720x1280 hw_raito 9/16 = size_ratio 9/16 newh=1280/1920*1080=720 neww=1280 | ||
| 95 | + h = int(size[1] / w * h) | ||
| 96 | + w = size[1] | ||
| 97 | + return h, w | ||
| 98 | + | ||
| 99 | + | ||
| 100 | +def shortsideresize(h, w, size, skip_low_resolution): | ||
| 101 | + if h <= size[0] and w <= size[1] and skip_low_resolution: | ||
| 102 | + return h, w | ||
| 103 | + | ||
| 104 | + if h / w < size[0] / size[1]: | ||
| 105 | + w = int(size[0] / h * w) | ||
| 106 | + h = size[0] | ||
| 107 | + else: | ||
| 108 | + h = int(size[1] / w * h) | ||
| 109 | + w = size[1] | ||
| 110 | + return h, w | ||
| 111 | + | ||
| 112 | + | ||
| 113 | +def calculate_centered_alignment(h: int, w: int, stride: int) -> tuple: | ||
| 114 | + """Compute centered crop offsets and stride-aligned dimensions; returns (v_off, h_off, aligned_h, aligned_w).""" | ||
| 115 | + # Calculate aligned dimensions | ||
| 116 | + aligned_h = h // stride * stride | ||
| 117 | + aligned_w = w // stride * stride | ||
| 118 | + | ||
| 119 | + # Compute centering offsets | ||
| 120 | + vertical_offset = (h - aligned_h) // 2 | ||
| 121 | + horizontal_offset = (w - aligned_w) // 2 | ||
| 122 | + | ||
| 123 | + return (vertical_offset, horizontal_offset, aligned_h, aligned_w) | ||
| 124 | + | ||
| 125 | + | ||
| 126 | +def maxhwresize(ori_height, ori_width, max_hxw): | ||
| 127 | + if ori_height * ori_width > max_hxw: | ||
| 128 | + scale_factor = np.sqrt(max_hxw / (ori_height * ori_width)) | ||
| 129 | + new_height = int(ori_height * scale_factor) | ||
| 130 | + new_width = int(ori_width * scale_factor) | ||
| 131 | + else: | ||
| 132 | + new_height = ori_height | ||
| 133 | + new_width = ori_width | ||
| 134 | + return new_height, new_width | ||
| 135 | + | ||
| 136 | + | ||
| 137 | +class ToTensorVideo: | ||
| 138 | + """Convert clip from uint8 to float (divide by 255.0) and permute dims.""" | ||
| 139 | + | ||
| 140 | + def __init__(self): | ||
| 141 | + pass | ||
| 142 | + | ||
| 143 | + @staticmethod | ||
| 144 | + def __call__(clip): | ||
| 145 | + """Convert clip (T, C, H, W) uint8 -> float.""" | ||
| 146 | + return to_tensor(clip) | ||
| 147 | + | ||
| 148 | + def __repr__(self) -> str: | ||
| 149 | + return self.__class__.__name__ | ||
| 150 | + | ||
| 151 | + | ||
| 152 | +class CenterCropResizeVideo: | ||
| 153 | + """Center-crop the short side, then resize to the specified size.""" | ||
| 154 | + | ||
| 155 | + def __init__( | ||
| 156 | + self, | ||
| 157 | + transform_size=None, | ||
| 158 | + use_short_edge=False, | ||
| 159 | + top_crop=False, | ||
| 160 | + interpolation_mode="bilinear", | ||
| 161 | + align_corners=False, | ||
| 162 | + antialias=False, | ||
| 163 | + ): | ||
| 164 | + if transform_size is None or "max_height" not in transform_size or "max_width" not in transform_size: | ||
| 165 | + raise ValueError("Missing required param: max_height or max_width in data transform.") | ||
| 166 | + self.size = (transform_size["max_height"], transform_size["max_width"]) | ||
| 167 | + self.use_short_edge = use_short_edge | ||
| 168 | + self.top_crop = top_crop | ||
| 169 | + self.interpolation_mode = interpolation_mode | ||
| 170 | + self.align_corners = align_corners | ||
| 171 | + self.antialias = antialias | ||
| 172 | + | ||
| 173 | + def __call__(self, clip): | ||
| 174 | + """Return scale-resized / center-cropped clip (T, C, crop_size, crop_size).""" | ||
| 175 | + if self.use_short_edge: | ||
| 176 | + clip_center_crop = center_crop_using_short_edge(clip) | ||
| 177 | + else: | ||
| 178 | + clip_center_crop = center_crop_th_tw( | ||
| 179 | + clip, self.size[0], self.size[1], top_crop=self.top_crop | ||
| 180 | + ) | ||
| 181 | + | ||
| 182 | + clip_center_crop_resize = resize( | ||
| 183 | + clip_center_crop, | ||
| 184 | + target_size=self.size, | ||
| 185 | + interpolation_mode=self.interpolation_mode, | ||
| 186 | + align_corners=self.align_corners, | ||
| 187 | + antialias=self.antialias, | ||
| 188 | + ) | ||
| 189 | + return clip_center_crop_resize | ||
| 190 | + | ||
| 191 | + def __repr__(self) -> str: | ||
| 192 | + return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}" | ||
| 193 | + | ||
| 194 | + | ||
| 195 | +class ResizeVideo: | ||
| 196 | + def __init__( | ||
| 197 | + self, | ||
| 198 | + transform_size="auto", | ||
| 199 | + interpolation_mode="bilinear", | ||
| 200 | + skip_low_resolution=False, | ||
| 201 | + align_corners=False, | ||
| 202 | + antialias=False, | ||
| 203 | + mode="resize" # resize / longside / shortside / hxw | ||
| 204 | + ): | ||
| 205 | + self.mode = mode | ||
| 206 | + if mode == 'hxw': | ||
| 207 | + self.transform_size = transform_size["max_hxw"] if isinstance(transform_size, dict) else transform_size | ||
| 208 | + elif mode in ["resize", "longside", "shortside"]: | ||
| 209 | + self.transform_size = (transform_size["max_height"], transform_size["max_width"]) if isinstance(transform_size, dict) else transform_size | ||
| 210 | + else: | ||
| 211 | + raise NotImplementedError(f"ResizeVideo only support mode `resize` / `longside` / `shortside` / `hxw`, {mode} is not implemented.") | ||
| 212 | + | ||
| 213 | + self.interpolation_mode = interpolation_mode | ||
| 214 | + self.align_corners = align_corners | ||
| 215 | + self.antialias = antialias | ||
| 216 | + self.skip_low_resolution = skip_low_resolution | ||
| 217 | + if isinstance(self.transform_size, str) and self.transform_size == "auto": | ||
| 218 | + raise ValueError("transform_size='auto' is a placeholder and must be replaced via get_transforms.") | ||
| 219 | + | ||
| 220 | + def __call__(self, clip): | ||
| 221 | + """Return scale-resized clip (T, C, H, W).""" | ||
| 222 | + h, w = clip.shape[-2:] | ||
| 223 | + if self.mode == "hxw": | ||
| 224 | + tr_h, tr_w = maxhwresize(h, w, self.transform_size) | ||
| 225 | + elif self.mode == "resize": | ||
| 226 | + tr_h, tr_w = self.transform_size | ||
| 227 | + elif self.mode == "longside": | ||
| 228 | + tr_h, tr_w = longsideresize(h, w, self.transform_size, skip_low_resolution=self.skip_low_resolution) | ||
| 229 | + elif self.mode == "shortside": | ||
| 230 | + tr_h, tr_w = shortsideresize(h, w, self.transform_size, skip_low_resolution=self.skip_low_resolution) | ||
| 231 | + | ||
| 232 | + if h == tr_h and w == tr_w: | ||
| 233 | + return clip | ||
| 234 | + resize_clip = resize( | ||
| 235 | + clip, | ||
| 236 | + target_size=(tr_h, tr_w), | ||
| 237 | + interpolation_mode=self.interpolation_mode, | ||
| 238 | + align_corners=self.align_corners, | ||
| 239 | + antialias=self.antialias | ||
| 240 | + ) | ||
| 241 | + return resize_clip | ||
| 242 | + | ||
| 243 | + def __repr__(self) -> str: | ||
| 244 | + return f"{self.__class__.__name__}(size={self.transform_size}, interpolation_mode={self.interpolation_mode})" | ||
| 245 | + | ||
| 246 | + | ||
| 247 | +class TemporalRandomCrop: | ||
| 248 | + """Temporally crop frame indices at a random location to ``size`` frames.""" | ||
| 249 | + | ||
| 250 | + def __init__(self, size): | ||
| 251 | + self.size = size | ||
| 252 | + | ||
| 253 | + def __call__(self, total_frames): | ||
| 254 | + rand_end = max(0, total_frames - self.size) | ||
| 255 | + begin_index = random.randint(0, rand_end) | ||
| 256 | + end_index = min(begin_index + self.size, total_frames) | ||
| 257 | + return begin_index, end_index | ||
| 258 | + | ||
| 259 | + | ||
| 260 | +low_aesthetic_score_notices_video = [ | ||
| 261 | + "This video has a low aesthetic quality.", | ||
| 262 | + "The beauty of this video is minimal.", | ||
| 263 | + "This video scores low in aesthetic appeal.", | ||
| 264 | + "The aesthetic quality of this video is below average.", | ||
| 265 | + "This video ranks low for beauty.", | ||
| 266 | + "The artistic quality of this video is lacking.", | ||
| 267 | + "This video has a low score for aesthetic value.", | ||
| 268 | + "The visual appeal of this video is low.", | ||
| 269 | + "This video is rated low for beauty.", | ||
| 270 | + "The aesthetic quality of this video is poor.", | ||
| 271 | +] | ||
| 272 | + | ||
| 273 | +high_aesthetic_score_notices_video = [ | ||
| 274 | + "This video has a high aesthetic quality.", | ||
| 275 | + "The beauty of this video is exceptional.", | ||
| 276 | + "This video scores high in aesthetic value.", | ||
| 277 | + "With its harmonious colors and balanced composition.", | ||
| 278 | + "This video ranks highly for aesthetic quality", | ||
| 279 | + "The artistic quality of this video is excellent.", | ||
| 280 | + "This video is rated high for beauty.", | ||
| 281 | + "The aesthetic quality of this video is impressive.", | ||
| 282 | + "This video has a top aesthetic score.", | ||
| 283 | + "The visual appeal of this video is outstanding.", | ||
| 284 | +] | ||
| 285 | + | ||
| 286 | +low_aesthetic_score_notices_image = [ | ||
| 287 | + "This image has a low aesthetic quality.", | ||
| 288 | + "The beauty of this image is minimal.", | ||
| 289 | + "This image scores low in aesthetic appeal.", | ||
| 290 | + "The aesthetic quality of this image is below average.", | ||
| 291 | + "This image ranks low for beauty.", | ||
| 292 | + "The artistic quality of this image is lacking.", | ||
| 293 | + "This image has a low score for aesthetic value.", | ||
| 294 | + "The visual appeal of this image is low.", | ||
| 295 | + "This image is rated low for beauty.", | ||
| 296 | + "The aesthetic quality of this image is poor.", | ||
| 297 | +] | ||
| 298 | + | ||
| 299 | +high_aesthetic_score_notices_image = [ | ||
| 300 | + "This image has a high aesthetic quality.", | ||
| 301 | + "The beauty of this image is exceptional", | ||
| 302 | + "This photo scores high in aesthetic value.", | ||
| 303 | + "With its harmonious colors and balanced composition.", | ||
| 304 | + "This image ranks highly for aesthetic quality.", | ||
| 305 | + "The artistic quality of this photo is excellent.", | ||
| 306 | + "This image is rated high for beauty.", | ||
| 307 | + "The aesthetic quality of this image is impressive.", | ||
| 308 | + "This photo has a top aesthetic score.", | ||
| 309 | + "The visual appeal of this image is outstanding.", | ||
| 310 | +] | ||
| 311 | + | ||
| 312 | + | ||
| 313 | +def add_aesthetic_notice_video(caption, aesthetic_score): | ||
| 314 | + if aesthetic_score <= 4.25: | ||
| 315 | + notice = random.choice(low_aesthetic_score_notices_video) | ||
| 316 | + return random.choice([caption + ' ' + notice, notice + ' ' + caption]) | ||
| 317 | + if aesthetic_score >= 5.75: | ||
| 318 | + notice = random.choice(high_aesthetic_score_notices_video) | ||
| 319 | + return random.choice([caption + ' ' + notice, notice + ' ' + caption]) | ||
| 320 | + return caption | ||
| 321 | + | ||
| 322 | + | ||
| 323 | +def add_aesthetic_notice_image(caption, aesthetic_score): | ||
| 324 | + if aesthetic_score <= 4.25: | ||
| 325 | + notice = random.choice(low_aesthetic_score_notices_image) | ||
| 326 | + return random.choice([caption + ' ' + notice, notice + ' ' + caption]) | ||
| 327 | + if aesthetic_score >= 5.75: | ||
| 328 | + notice = random.choice(high_aesthetic_score_notices_image) | ||
| 329 | + return random.choice([caption + ' ' + notice, notice + ' ' + caption]) | ||
| 330 | + return caption | ||
| @@ -0,0 +1,634 @@ | |||
| 1 | +# -------------------------------------------------------- | ||
| 2 | +# Modified from huggingface diffusers repos | ||
| 3 | +# This source code is licensed under the notice found in the root directory of this source tree. | ||
| 4 | +# -------------------------------------------------------- | ||
| 5 | +# References: TextProcesser (diffusers deepfloyd_if pipeline) | ||
| 6 | + | ||
| 7 | +"""Data reader/media processing utilities for generative-model (Wan2.2) FSDP2 datasets. | ||
| 8 | +Moved out of ``data_utils/utils.py``; carries no megatron dependency.""" | ||
| 9 | + | ||
| 10 | +import os | ||
| 11 | +import re | ||
| 12 | +import html | ||
| 13 | +import random | ||
| 14 | +import math | ||
| 15 | +import urllib.parse as ul | ||
| 16 | +from collections import defaultdict | ||
| 17 | +from multiprocessing import shared_memory | ||
| 18 | +from concurrent.futures import ProcessPoolExecutor, as_completed | ||
| 19 | + | ||
| 20 | +try: | ||
| 21 | + import decord | ||
| 22 | +except Exception as e: | ||
| 23 | + print(f"Failed to import decord module. The reason of decord unavailable is {e}") | ||
| 24 | + | ||
| 25 | +import orjson | ||
| 26 | +import ftfy | ||
| 27 | +import torch | ||
| 28 | +import numpy as np | ||
| 29 | +import pandas as pd | ||
| 30 | +from PIL import Image | ||
| 31 | +from bs4 import BeautifulSoup | ||
| 32 | +from einops import rearrange | ||
| 33 | +from torchvision.datasets.folder import pil_loader | ||
| 34 | + | ||
| 35 | +from mindspeed_mm.fsdp.data.data_utils.generative_data_utils.transform_pipeline import get_transforms | ||
| 36 | + | ||
| 37 | + | ||
| 38 | +VID_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv") | ||
| 39 | +TENSOR_EXTENSIONS = (".pt", ".pth") | ||
| 40 | + | ||
| 41 | + | ||
| 42 | +class DataFileReader: | ||
| 43 | + """get the data from different types of files such as csv/json/parquet""" | ||
| 44 | + | ||
| 45 | + def __init__(self, data_storage_mode="standard", **kwargs): | ||
| 46 | + """``data_storage_mode`` controls loading; ``reserved_keys`` filters keys from data.json; ``use_multiprocess`` enables parallel file processing (not recommended for fewer than 4 files).""" | ||
| 47 | + self.data_storage_mode = data_storage_mode | ||
| 48 | + self.reserved_keys = kwargs.get("reserved_keys", None) | ||
| 49 | + self.use_multiprocess = kwargs.get("use_multiprocess", False) | ||
| 50 | + | ||
| 51 | + def __call__(self, data_path, return_type="list"): | ||
| 52 | + if self.data_storage_mode == "standard": | ||
| 53 | + return self.get_datasamples(data_path, return_type=return_type) | ||
| 54 | + elif self.data_storage_mode == "combine" or self.data_storage_mode == "sorafeatured": | ||
| 55 | + redirect_keys = ["path"] | ||
| 56 | + return self.get_cap_list(data_path, redirect_keys) | ||
| 57 | + elif self.data_storage_mode == "vace": | ||
| 58 | + redirect_keys = ["video", "src_video", "src_video_mask", "src_ref_images"] | ||
| 59 | + return self.get_cap_list(data_path, redirect_keys) | ||
| 60 | + else: | ||
| 61 | + raise NotImplementedError("Not support now.") | ||
| 62 | + | ||
| 63 | + @staticmethod | ||
| 64 | + def get_datasamples(data_path, return_type="list"): | ||
| 65 | + if data_path.endswith(".csv"): | ||
| 66 | + data_out = pd.read_csv(data_path) | ||
| 67 | + if return_type == "list": | ||
| 68 | + return data_out.to_dict("records") | ||
| 69 | + else: | ||
| 70 | + return data_out | ||
| 71 | + elif data_path.endswith(".json"): | ||
| 72 | + return orjson_load(data_path) | ||
| 73 | + elif data_path.endswith(".jsonl"): | ||
| 74 | + return orjson_load(data_path) | ||
| 75 | + elif data_path.endswith(".parquet"): | ||
| 76 | + data_out = pd.read_parquet(data_path) | ||
| 77 | + return data_out.to_dict("records") | ||
| 78 | + elif data_path.endswith(".txt"): | ||
| 79 | + with open(data_path, 'r') as f: | ||
| 80 | + data_out = f.readlines() | ||
| 81 | + data_out = [data.strip() for data in data_out] | ||
| 82 | + return data_out | ||
| 83 | + else: | ||
| 84 | + raise NotImplementedError(f"Unsupported file format: {data_path}") | ||
| 85 | + | ||
| 86 | + def get_cap_list(self, data_path, redirect_keys=None): | ||
| 87 | + with open(data_path, "r") as f: | ||
| 88 | + folder_anno = [] | ||
| 89 | + for line_no, line in enumerate(f.readlines(), 1): | ||
| 90 | + if len(line.strip()) == 0: | ||
| 91 | + continue | ||
| 92 | + fields = line.strip().split(",") | ||
| 93 | + if len(fields) != 2: | ||
| 94 | + print(f"Skip malformed line {line_no} in {data_path}: expected 2 fields, got {len(fields)}") | ||
| 95 | + continue | ||
| 96 | + folder_anno.append(fields) | ||
| 97 | + json_loader = JsonLoader([temp[1] for temp in folder_anno], use_multiprocess=self.use_multiprocess) | ||
| 98 | + for folder, anno in folder_anno: | ||
🟡 Medium Priority 变更行 L94–100: 受影响行为:如果数据文件某行包含的逗号分隔字段数不等于 2,会触发 失败模式:数据加载因格式不合预期而崩溃。CSV/TSV 文件中字段数不一致是常见的脏数据问题。 建议:在解包前检查每行 split 后的长度是否为 2,否则跳过并给出警告。 ![]() ![]() 不准确? | |||
| 99 | + json_loader.set_process_func(anno, self._change_path, redirect_keys, folder) | ||
| 100 | + json_loader.set_process_func("all", self._remove_unused_keys, self.reserved_keys) | ||


🟠 High Priority
变更行 L201:
transform_size默认值是字符串"auto"。L210 和 L212 中,当transform_size不是 dict 时,直接使用原值(即"auto")。受影响行为:
失败模式:任何使用
ResizeVideo()默认参数构造的实例,在__call__时必然抛出异常。这是一个一触即发的运行时错误。建议:将默认值改为 None,在 isinstance 检查为 False 时抛出明确的 ValueError 说明必须传入 dict 类型的 transform_size。