`diffsynth.core.data`: 数据处理算子与通用数据集

数据处理算子

可用数据处理算子

diffsynth.core.data 提供了一系列数据处理算子，用于进行数据处理，包括：

数据格式转换算子
- ToInt: 转换为 int 格式
- ToFloat: 转换为 float 格式
- ToStr: 转换为 str 格式
- ToList: 转换为列表格式，以列表包裹此数据
- ToAbsolutePath: 将相对路径转换为绝对路径
文件加载算子
- LoadImage: 读取图片文件
- LoadVideo: 读取视频文件
- LoadAudio: 读取音频文件
- LoadGIF: 读取 GIF 文件
- LoadTorchPickle: 读取由 torch.save 保存的二进制文件【该算子可能导致二进制文件中的代码注入攻击，请谨慎使用！】
媒体文件处理算子
- ImageCropAndResize: 对图像进行裁剪和拉伸
Meta 算子
- SequencialProcess: 将序列中的每个数据路由到一个算子
- RouteByExtensionName: 按照文件扩展名路由到特定算子
- RouteByType: 按照数据类型路由到特定算子

算子使用

数据算子之间以 >> 符号连接形成数据处理流水线，例如：

from diffsynth.core.data.operators import *

data = "image.jpg"
data_pipeline = ToAbsolutePath(base_path="/data") >> LoadImage() >> ImageCropAndResize(max_pixels=512*512)
data = data_pipeline(data)

在经过每个算子后，数据被依次处理

ToAbsolutePath(base_path="/data"): "/data/image.jpg"
LoadImage(): <PIL.Image.Image image mode=RGB size=1024x1024 at 0x7F8E7AAEFC10>
ImageCropAndResize(max_pixels=512*512): <PIL.Image.Image image mode=RGB size=512x512 at 0x7F8E7A936F20>

我们可以组合出功能完备的数据流水线，例如通用数据集的默认视频数据算子为

RouteByType(operator_map=[
    (str, ToAbsolutePath(base_path) >> RouteByExtensionName(operator_map=[
        (("jpg", "jpeg", "png", "webp"), LoadImage() >> ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor) >> ToList()),
        (("gif",), LoadGIF(
            num_frames, time_division_factor, time_division_remainder,
            frame_processor=ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor),
        )),
        (("mp4", "avi", "mov", "wmv", "mkv", "flv", "webm"), LoadVideo(
            num_frames, time_division_factor, time_division_remainder,
            frame_processor=ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor),
        )),
    ])),
])

它包含如下逻辑：

如果是 str 类型的数据
- 如果是 "jpg", "jpeg", "png", "webp" 类型文件
  - 加载这张图片
  - 裁剪并缩放到特定分辨率
  - 打包进列表，视为单帧视频
- 如果是 "gif" 类型文件
  - 加载 gif 文件内容
  - 将每一帧裁剪和缩放到特定分辨率
- 如果是 "mp4", "avi", "mov", "wmv", "mkv", "flv", "webm" 类型文件
  - 加载 gif 文件内容
  - 将每一帧裁剪和缩放到特定分辨率
如果不是 str 类型的数据，报错

通用数据集

diffsynth.core.data 提供了统一的数据集实现，数据集需输入以下参数：

base_path: 根目录，若数据集中包含图片文件的相对路径，则需填入此字段用于加载这些路径指向的文件
metadata_path: 元数据目录，记录所有元数据的文件路径，支持 csv、json、jsonl 格式
repeat: 数据重复次数，默认为 1，该参数影响一个 epoch 的训练步数
data_file_keys: 需进行加载的数据字段名，例如 (image, edit_image)
main_data_operator: 主加载算子，需通过数据处理算子组装好数据处理流水线
special_operator_map: 特殊算子映射，对需要特殊处理的字段构建的算子映射

元数据

数据集的 metadata_path 指向元数据文件，支持 csv、json、jsonl 格式，以下提供了样例

csv 格式：可读性高、不支持列表数据、内存占用小

image,prompt
image_1.jpg,"a dog"
image_2.jpg,"a cat"

json 格式：可读性高、支持列表数据、内存占用大

[
    {
        "image": "image_1.jpg",
        "prompt": "a dog"
    },
    {
        "image": "image_2.jpg",
        "prompt": "a cat"
    }
]

jsonl 格式：可读性低、支持列表数据、内存占用小

{"image": "image_1.jpg", "prompt": "a dog"}
{"image": "image_2.jpg", "prompt": "a cat"}

如何选择最佳的元数据格式？

如果数据量大，达到千万级的数据量，由于 json 文件解析时需要额外内存，此时不可用，请使用 csv 或 jsonl 格式
如果数据集中包含列表数据，例如编辑模型需输入多张图，由于 csv 格式无法存储列表格式数据，此时不可用，请使用 json 或 jsonl 格式

数据加载逻辑

在没有进行额外设置时，数据集默认输出元数据集中的数据，图片和视频文件的路径会以字符串的格式输出，若要加载这些文件，则需要设置 data_file_keys、main_data_operator、special_operator_map。

在数据处理流程中，按如下逻辑进行处理：

如果字段位于 special_operator_map，则调用 special_operator_map 中的对应算子进行处理
如果字段不位于 special_operator_map
- 如果字段位于 data_file_keys，则调用 main_data_operator 算子进行处理
- 如果字段不位于 data_file_keys，则不进行处理

special_operator_map 可用于实现特殊的数据处理，例如模型 Wan-AI/Wan2.2-Animate-14B 中输入的人物面部视频 animate_face_video 是以固定分辨率处理的，与输出视频不一致，因此这一字段由专门的算子处理：

special_operator_map={
    "animate_face_video": ToAbsolutePath(args.dataset_base_path) >> LoadVideo(args.num_frames, 4, 1, frame_processor=ImageCropAndResize(512, 512, None, 16, 16)),
}

其他注意事项

当数据量过少时，可适当增加 repeat，延长单个 epoch 的训练时间，避免频繁保存模型产生较多耗时。

当数据量 * repeat 超过 10^9 时，我们观测到数据集的速度明显变慢，这似乎是 PyTorch 的 bug，我们尚不确定新版本的 PyTorch 是否已经修复了这一问题。

6.5 KiB Raw Blame History Unescape Escape

diffsynth.core.data: 数据处理算子与通用数据集