已合并
feat: add generative data utils #2924
feat: add generative data utils #2924
已合并
chenpeizhe创建于 8月4日
共 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
atomgit-bot
atomgit-botatomgit-bot8月4日

🟠 High Priority

变更行 L201:transform_size 默认值是字符串 "auto"。L210 和 L212 中,当 transform_size 不是 dict 时,直接使用原值(即 "auto")。

受影响行为:

失败模式:任何使用 ResizeVideo() 默认参数构造的实例,在 __call__ 时必然抛出异常。这是一个一触即发的运行时错误。

建议:将默认值改为 None,在 isinstance 检查为 False 时抛出明确的 ValueError 说明必须传入 dict 类型的 transform_size。

likedislike
不准确?
chenpeizhe
8月5日 评论:
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:
atomgit-bot
atomgit-botatomgit-bot8月4日

🟡 Medium Priority

变更行 L94–100:folder_anno = [i.strip().split(",") ...] 按逗号分割每行,然后 L100 for folder, anno in folder_anno: 硬编码解包为恰好两个变量。

受影响行为:如果数据文件某行包含的逗号分隔字段数不等于 2,会触发 ValueError: too many values to unpack(多于 2 个字段)或 ValueError: not enough values to unpack(只有 1 个字段)。此外 L99 的 [temp[1] for temp in folder_anno] 在字段数 <2 时也会触发 IndexError。

失败模式:数据加载因格式不合预期而崩溃。CSV/TSV 文件中字段数不一致是常见的脏数据问题。

建议:在解包前检查每行 split 后的长度是否为 2,否则跳过并给出警告。

likedislike
不准确?
chenpeizhe
8月5日 评论:
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)