diff --git a/.gitignore b/.gitignore index 5627e450..3840acc8 100644 --- a/.gitignore +++ b/.gitignore @@ -1,17 +1,19 @@ -# Byte-compiled / optimized / DLL files -models* -output* -logs* -taming* -samples* -datasets* -asset* +# Used in VideoX-Fun _* -logs* +/models* +/output* +/logs* +/taming* +/samples* +/datasets* +/asset* +/repo* +/scripts_demo* + +# Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class -scripts_demo* # C extensions *.so diff --git a/examples/wan2.2/predict_animate.py b/examples/wan2.2/predict_animate.py index 3d24f8dd..74f18dfe 100644 --- a/examples/wan2.2/predict_animate.py +++ b/examples/wan2.2/predict_animate.py @@ -314,8 +314,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -363,8 +363,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2/predict_i2v.py b/examples/wan2.2/predict_i2v.py index 6fd2a073..d883a576 100644 --- a/examples/wan2.2/predict_i2v.py +++ b/examples/wan2.2/predict_i2v.py @@ -136,13 +136,15 @@ low_cpu_mem_usage=True, torch_dtype=weight_dtype, ) - -transformer_2 = Wan2_2Transformer3DModel.from_pretrained( - os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')), - transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), - low_cpu_mem_usage=True, - torch_dtype=weight_dtype, -) +if config['transformer_additional_kwargs'].get('transformer_combination_type', 'single') == "moe": + transformer_2 = Wan2_2Transformer3DModel.from_pretrained( + os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=weight_dtype, + ) +else: + transformer_2 = None if transformer_path is not None: print(f"From checkpoint: {transformer_path}") @@ -156,17 +158,18 @@ m, u = transformer.load_state_dict(state_dict, strict=False) print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") -if transformer_high_path is not None: - print(f"From checkpoint: {transformer_high_path}") - if transformer_high_path.endswith("safetensors"): - from safetensors.torch import load_file, safe_open - state_dict = load_file(transformer_high_path) - else: - state_dict = torch.load(transformer_high_path, map_location="cpu") - state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict +if transformer_2 is not None: + if transformer_high_path is not None: + print(f"From checkpoint: {transformer_high_path}") + if transformer_high_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(transformer_high_path) + else: + state_dict = torch.load(transformer_high_path, map_location="cpu") + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict - m, u = transformer_2.load_state_dict(state_dict, strict=False) - print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + m, u = transformer_2.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") # Get Vae Chosen_AutoencoderKL = { @@ -228,11 +231,13 @@ if ulysses_degree > 1 or ring_degree > 1: from functools import partial transformer.enable_multi_gpus_inference() - transformer_2.enable_multi_gpus_inference() + if transformer_2 is not None: + transformer_2.enable_multi_gpus_inference() if fsdp_dit: shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype) pipeline.transformer = shard_fn(pipeline.transformer) - pipeline.transformer_2 = shard_fn(pipeline.transformer_2) + if transformer_2 is not None: + pipeline.transformer_2 = shard_fn(pipeline.transformer_2) print("Add FSDP DIT") if fsdp_text_encoder: shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype) @@ -242,33 +247,38 @@ if compile_dit: for i in range(len(pipeline.transformer.blocks)): pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i]) - for i in range(len(pipeline.transformer_2.blocks)): - pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i]) + if transformer_2 is not None: + for i in range(len(pipeline.transformer_2.blocks)): + pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i]) print("Add Compile") if GPU_memory_mode == "sequential_cpu_offload": replace_parameters_by_name(transformer, ["modulation",], device=device) - replace_parameters_by_name(transformer_2, ["modulation",], device=device) transformer.freqs = transformer.freqs.to(device=device) - transformer_2.freqs = transformer_2.freqs.to(device=device) + if transformer_2 is not None: + replace_parameters_by_name(transformer_2, ["modulation",], device=device) + transformer_2.freqs = transformer_2.freqs.to(device=device) pipeline.enable_sequential_cpu_offload(device=device) elif GPU_memory_mode == "model_group_offload": register_auto_device_hook(pipeline.transformer) - register_auto_device_hook(pipeline.transformer_2) + if transformer_2 is not None: + register_auto_device_hook(pipeline.transformer_2) safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True) elif GPU_memory_mode == "model_cpu_offload_and_qfloat8": convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device) - convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device) convert_weight_dtype_wrapper(transformer, weight_dtype) - convert_weight_dtype_wrapper(transformer_2, weight_dtype) + if transformer_2 is not None: + convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device) + convert_weight_dtype_wrapper(transformer_2, weight_dtype) pipeline.enable_model_cpu_offload(device=device) elif GPU_memory_mode == "model_cpu_offload": pipeline.enable_model_cpu_offload(device=device) elif GPU_memory_mode == "model_full_load_and_qfloat8": convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device) - convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device) convert_weight_dtype_wrapper(transformer, weight_dtype) - convert_weight_dtype_wrapper(transformer_2, weight_dtype) + if transformer_2 is not None: + convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device) + convert_weight_dtype_wrapper(transformer_2, weight_dtype) pipeline.to(device=device) else: pipeline.to(device=device) @@ -279,17 +289,20 @@ pipeline.transformer.enable_teacache( coefficients, num_inference_steps, teacache_threshold, num_skip_start_steps=num_skip_start_steps, offload=teacache_offload ) - pipeline.transformer_2.share_teacache(transformer=pipeline.transformer) + if transformer_2 is not None: + pipeline.transformer_2.share_teacache(transformer=pipeline.transformer) if cfg_skip_ratio is not None: print(f"Enable cfg_skip_ratio {cfg_skip_ratio}.") pipeline.transformer.enable_cfg_skip(cfg_skip_ratio, num_inference_steps) - pipeline.transformer_2.share_cfg_skip(transformer=pipeline.transformer) + if transformer_2 is not None: + pipeline.transformer_2.share_cfg_skip(transformer=pipeline.transformer) generator = torch.Generator(device=device).manual_seed(seed) if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) +if lora_high_path is not None and transformer_2 is not None: pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): @@ -298,7 +311,8 @@ if enable_riflex: pipeline.transformer.enable_riflex(k = riflex_k, L_test = latent_frames) - pipeline.transformer_2.enable_riflex(k = riflex_k, L_test = latent_frames) + if transformer_2 is not None: + pipeline.transformer_2.enable_riflex(k = riflex_k, L_test = latent_frames) input_video, input_video_mask, clip_image = get_image_to_video_latent(validation_image_start, None, video_length=video_length, sample_size=sample_size) @@ -320,6 +334,7 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) +if lora_high_path is not None and transformer_2 is not None: pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): diff --git a/examples/wan2.2/predict_s2v.py b/examples/wan2.2/predict_s2v.py index 45c82095..d7dcf2d1 100644 --- a/examples/wan2.2/predict_s2v.py +++ b/examples/wan2.2/predict_s2v.py @@ -318,8 +318,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): segment_frame_length = segment_frame_length // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio if segment_frame_length != 1 else 1 @@ -356,8 +356,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2/predict_t2v.py b/examples/wan2.2/predict_t2v.py index c64d73a3..2c2d238f 100755 --- a/examples/wan2.2/predict_t2v.py +++ b/examples/wan2.2/predict_t2v.py @@ -132,13 +132,15 @@ low_cpu_mem_usage=True, torch_dtype=weight_dtype, ) - -transformer_2 = Wan2_2Transformer3DModel.from_pretrained( - os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')), - transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), - low_cpu_mem_usage=True, - torch_dtype=weight_dtype, -) +if config['transformer_additional_kwargs'].get('transformer_combination_type', 'single') == "moe": + transformer_2 = Wan2_2Transformer3DModel.from_pretrained( + os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=weight_dtype, + ) +else: + transformer_2 = None if transformer_path is not None: print(f"From checkpoint: {transformer_path}") @@ -152,17 +154,18 @@ m, u = transformer.load_state_dict(state_dict, strict=False) print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") -if transformer_high_path is not None: - print(f"From checkpoint: {transformer_high_path}") - if transformer_high_path.endswith("safetensors"): - from safetensors.torch import load_file, safe_open - state_dict = load_file(transformer_high_path) - else: - state_dict = torch.load(transformer_high_path, map_location="cpu") - state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict +if transformer_2 is not None: + if transformer_high_path is not None: + print(f"From checkpoint: {transformer_high_path}") + if transformer_high_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(transformer_high_path) + else: + state_dict = torch.load(transformer_high_path, map_location="cpu") + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict - m, u = transformer_2.load_state_dict(state_dict, strict=False) - print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + m, u = transformer_2.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") # Get Vae Chosen_AutoencoderKL = { @@ -223,11 +226,13 @@ if ulysses_degree > 1 or ring_degree > 1: from functools import partial transformer.enable_multi_gpus_inference() - transformer_2.enable_multi_gpus_inference() + if transformer_2 is not None: + transformer_2.enable_multi_gpus_inference() if fsdp_dit: shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype) pipeline.transformer = shard_fn(pipeline.transformer) - pipeline.transformer_2 = shard_fn(pipeline.transformer_2) + if transformer_2 is not None: + pipeline.transformer_2 = shard_fn(pipeline.transformer_2) print("Add FSDP DIT") if fsdp_text_encoder: shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype) @@ -237,33 +242,38 @@ if compile_dit: for i in range(len(pipeline.transformer.blocks)): pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i]) - for i in range(len(pipeline.transformer_2.blocks)): - pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i]) + if transformer_2 is not None: + for i in range(len(pipeline.transformer_2.blocks)): + pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i]) print("Add Compile") if GPU_memory_mode == "sequential_cpu_offload": replace_parameters_by_name(transformer, ["modulation",], device=device) - replace_parameters_by_name(transformer_2, ["modulation",], device=device) transformer.freqs = transformer.freqs.to(device=device) - transformer_2.freqs = transformer_2.freqs.to(device=device) + if transformer_2 is not None: + replace_parameters_by_name(transformer_2, ["modulation",], device=device) + transformer_2.freqs = transformer_2.freqs.to(device=device) pipeline.enable_sequential_cpu_offload(device=device) elif GPU_memory_mode == "model_group_offload": register_auto_device_hook(pipeline.transformer) - register_auto_device_hook(pipeline.transformer_2) + if transformer_2 is not None: + register_auto_device_hook(pipeline.transformer_2) safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True) elif GPU_memory_mode == "model_cpu_offload_and_qfloat8": convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device) - convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device) convert_weight_dtype_wrapper(transformer, weight_dtype) - convert_weight_dtype_wrapper(transformer_2, weight_dtype) + if transformer_2 is not None: + convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device) + convert_weight_dtype_wrapper(transformer_2, weight_dtype) pipeline.enable_model_cpu_offload(device=device) elif GPU_memory_mode == "model_cpu_offload": pipeline.enable_model_cpu_offload(device=device) elif GPU_memory_mode == "model_full_load_and_qfloat8": convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device) - convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device) convert_weight_dtype_wrapper(transformer, weight_dtype) - convert_weight_dtype_wrapper(transformer_2, weight_dtype) + if transformer_2 is not None: + convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device) + convert_weight_dtype_wrapper(transformer_2, weight_dtype) pipeline.to(device=device) else: pipeline.to(device=device) @@ -274,17 +284,20 @@ pipeline.transformer.enable_teacache( coefficients, num_inference_steps, teacache_threshold, num_skip_start_steps=num_skip_start_steps, offload=teacache_offload ) - pipeline.transformer_2.share_teacache(transformer=pipeline.transformer) + if transformer_2 is not None: + pipeline.transformer_2.share_teacache(transformer=pipeline.transformer) if cfg_skip_ratio is not None: print(f"Enable cfg_skip_ratio {cfg_skip_ratio}.") pipeline.transformer.enable_cfg_skip(cfg_skip_ratio, num_inference_steps) - pipeline.transformer_2.share_cfg_skip(transformer=pipeline.transformer) + if transformer_2 is not None: + pipeline.transformer_2.share_cfg_skip(transformer=pipeline.transformer) generator = torch.Generator(device=device).manual_seed(seed) if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) +if lora_high_path is not None and transformer_2 is not None: pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): @@ -293,7 +306,8 @@ if enable_riflex: pipeline.transformer.enable_riflex(k = riflex_k, L_test = latent_frames) - pipeline.transformer_2.enable_riflex(k = riflex_k, L_test = latent_frames) + if transformer_2 is not None: + pipeline.transformer_2.enable_riflex(k = riflex_k, L_test = latent_frames) sample = pipeline( prompt, @@ -310,6 +324,7 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) +if lora_high_path is not None and transformer_2 is not None: pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): diff --git a/examples/wan2.2/predict_ti2v.py b/examples/wan2.2/predict_ti2v.py index 36bdbfde..380d4d28 100755 --- a/examples/wan2.2/predict_ti2v.py +++ b/examples/wan2.2/predict_ti2v.py @@ -303,8 +303,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -338,8 +338,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2_fun/predict_i2v.py b/examples/wan2.2_fun/predict_i2v.py index 161e083c..d6bba172 100644 --- a/examples/wan2.2_fun/predict_i2v.py +++ b/examples/wan2.2_fun/predict_i2v.py @@ -305,8 +305,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -337,8 +337,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2_fun/predict_i2v_5b.py b/examples/wan2.2_fun/predict_i2v_5b.py index 339565f3..9922d46b 100644 --- a/examples/wan2.2_fun/predict_i2v_5b.py +++ b/examples/wan2.2_fun/predict_i2v_5b.py @@ -307,8 +307,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -339,8 +339,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2_fun/predict_t2v.py b/examples/wan2.2_fun/predict_t2v.py index 9f5efeb7..c68e0f50 100644 --- a/examples/wan2.2_fun/predict_t2v.py +++ b/examples/wan2.2_fun/predict_t2v.py @@ -135,13 +135,15 @@ low_cpu_mem_usage=True, torch_dtype=weight_dtype, ) - -transformer_2 = Wan2_2Transformer3DModel.from_pretrained( - os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')), - transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), - low_cpu_mem_usage=True, - torch_dtype=weight_dtype, -) +if config['transformer_additional_kwargs'].get('transformer_combination_type', 'single') == "moe": + transformer_2 = Wan2_2Transformer3DModel.from_pretrained( + os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=weight_dtype, + ) +else: + transformer_2 = None if transformer_path is not None: print(f"From checkpoint: {transformer_path}") @@ -155,17 +157,18 @@ m, u = transformer.load_state_dict(state_dict, strict=False) print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") -if transformer_high_path is not None: - print(f"From checkpoint: {transformer_high_path}") - if transformer_high_path.endswith("safetensors"): - from safetensors.torch import load_file, safe_open - state_dict = load_file(transformer_high_path) - else: - state_dict = torch.load(transformer_high_path, map_location="cpu") - state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict +if transformer_2 is not None: + if transformer_high_path is not None: + print(f"From checkpoint: {transformer_high_path}") + if transformer_high_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(transformer_high_path) + else: + state_dict = torch.load(transformer_high_path, map_location="cpu") + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict - m, u = transformer_2.load_state_dict(state_dict, strict=False) - print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + m, u = transformer_2.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") # Get Vae Chosen_AutoencoderKL = { @@ -227,11 +230,13 @@ if ulysses_degree > 1 or ring_degree > 1: from functools import partial transformer.enable_multi_gpus_inference() - transformer_2.enable_multi_gpus_inference() + if transformer_2 is not None: + transformer_2.enable_multi_gpus_inference() if fsdp_dit: shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype) pipeline.transformer = shard_fn(pipeline.transformer) - pipeline.transformer_2 = shard_fn(pipeline.transformer_2) + if transformer_2 is not None: + pipeline.transformer_2 = shard_fn(pipeline.transformer_2) print("Add FSDP DIT") if fsdp_text_encoder: shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype) @@ -241,33 +246,38 @@ if compile_dit: for i in range(len(pipeline.transformer.blocks)): pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i]) - for i in range(len(pipeline.transformer_2.blocks)): - pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i]) + if transformer_2 is not None: + for i in range(len(pipeline.transformer_2.blocks)): + pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i]) print("Add Compile") if GPU_memory_mode == "sequential_cpu_offload": replace_parameters_by_name(transformer, ["modulation",], device=device) - replace_parameters_by_name(transformer_2, ["modulation",], device=device) transformer.freqs = transformer.freqs.to(device=device) - transformer_2.freqs = transformer_2.freqs.to(device=device) + if transformer_2 is not None: + replace_parameters_by_name(transformer_2, ["modulation",], device=device) + transformer_2.freqs = transformer_2.freqs.to(device=device) pipeline.enable_sequential_cpu_offload(device=device) elif GPU_memory_mode == "model_group_offload": register_auto_device_hook(pipeline.transformer) - register_auto_device_hook(pipeline.transformer_2) + if transformer_2 is not None: + register_auto_device_hook(pipeline.transformer_2) safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True) elif GPU_memory_mode == "model_cpu_offload_and_qfloat8": convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device) - convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device) convert_weight_dtype_wrapper(transformer, weight_dtype) - convert_weight_dtype_wrapper(transformer_2, weight_dtype) + if transformer_2 is not None: + convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device) + convert_weight_dtype_wrapper(transformer_2, weight_dtype) pipeline.enable_model_cpu_offload(device=device) elif GPU_memory_mode == "model_cpu_offload": pipeline.enable_model_cpu_offload(device=device) elif GPU_memory_mode == "model_full_load_and_qfloat8": convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device) - convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device) convert_weight_dtype_wrapper(transformer, weight_dtype) - convert_weight_dtype_wrapper(transformer_2, weight_dtype) + if transformer_2 is not None: + convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device) + convert_weight_dtype_wrapper(transformer_2, weight_dtype) pipeline.to(device=device) else: pipeline.to(device=device) @@ -278,17 +288,20 @@ pipeline.transformer.enable_teacache( coefficients, num_inference_steps, teacache_threshold, num_skip_start_steps=num_skip_start_steps, offload=teacache_offload ) - pipeline.transformer_2.share_teacache(transformer=pipeline.transformer) + if transformer_2 is not None: + pipeline.transformer_2.share_teacache(transformer=pipeline.transformer) if cfg_skip_ratio is not None: print(f"Enable cfg_skip_ratio {cfg_skip_ratio}.") pipeline.transformer.enable_cfg_skip(cfg_skip_ratio, num_inference_steps) - pipeline.transformer_2.share_cfg_skip(transformer=pipeline.transformer) + if transformer_2 is not None: + pipeline.transformer_2.share_cfg_skip(transformer=pipeline.transformer) generator = torch.Generator(device=device).manual_seed(seed) if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) +if lora_high_path is not None and transformer_2 is not None: pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): @@ -297,7 +310,8 @@ if enable_riflex: pipeline.transformer.enable_riflex(k = riflex_k, L_test = latent_frames) - pipeline.transformer_2.enable_riflex(k = riflex_k, L_test = latent_frames) + if transformer_2 is not None: + pipeline.transformer_2.enable_riflex(k = riflex_k, L_test = latent_frames) input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=sample_size) @@ -319,6 +333,7 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) +if lora_high_path is not None and transformer_2 is not None: pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): diff --git a/examples/wan2.2_fun/predict_t2v_5b.py b/examples/wan2.2_fun/predict_t2v_5b.py index 6f7290b6..c4aa5144 100644 --- a/examples/wan2.2_fun/predict_t2v_5b.py +++ b/examples/wan2.2_fun/predict_t2v_5b.py @@ -301,8 +301,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -333,8 +333,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2_fun/predict_v2v_control.py b/examples/wan2.2_fun/predict_v2v_control.py index 63cd841f..0d076b6f 100644 --- a/examples/wan2.2_fun/predict_v2v_control.py +++ b/examples/wan2.2_fun/predict_v2v_control.py @@ -318,8 +318,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -364,8 +364,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2_fun/predict_v2v_control_5b.py b/examples/wan2.2_fun/predict_v2v_control_5b.py index 8cf222fc..7c78666d 100644 --- a/examples/wan2.2_fun/predict_v2v_control_5b.py +++ b/examples/wan2.2_fun/predict_v2v_control_5b.py @@ -318,8 +318,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -364,8 +364,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2_fun/predict_v2v_control_camera.py b/examples/wan2.2_fun/predict_v2v_control_camera.py index 666e0a48..3ec8e524 100644 --- a/examples/wan2.2_fun/predict_v2v_control_camera.py +++ b/examples/wan2.2_fun/predict_v2v_control_camera.py @@ -318,8 +318,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -364,8 +364,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2_fun/predict_v2v_control_camera_5b.py b/examples/wan2.2_fun/predict_v2v_control_camera_5b.py index a50173b6..ee36ab00 100644 --- a/examples/wan2.2_fun/predict_v2v_control_camera_5b.py +++ b/examples/wan2.2_fun/predict_v2v_control_camera_5b.py @@ -318,8 +318,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -364,8 +364,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2_fun/predict_v2v_control_ref.py b/examples/wan2.2_fun/predict_v2v_control_ref.py index eaf088d2..43906df8 100644 --- a/examples/wan2.2_fun/predict_v2v_control_ref.py +++ b/examples/wan2.2_fun/predict_v2v_control_ref.py @@ -318,8 +318,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -364,8 +364,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2_fun/predict_v2v_control_ref_5b.py b/examples/wan2.2_fun/predict_v2v_control_ref_5b.py index 5ecd1c7a..7d1c2c3b 100644 --- a/examples/wan2.2_fun/predict_v2v_control_ref_5b.py +++ b/examples/wan2.2_fun/predict_v2v_control_ref_5b.py @@ -318,8 +318,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -364,8 +364,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2_vace_fun/predict_i2v.py b/examples/wan2.2_vace_fun/predict_i2v.py index 1abf301e..fa18f3cb 100644 --- a/examples/wan2.2_vace_fun/predict_i2v.py +++ b/examples/wan2.2_vace_fun/predict_i2v.py @@ -322,8 +322,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -370,8 +370,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2_vace_fun/predict_s2v.py b/examples/wan2.2_vace_fun/predict_s2v.py index 5d90aa3f..9ff65a5f 100644 --- a/examples/wan2.2_vace_fun/predict_s2v.py +++ b/examples/wan2.2_vace_fun/predict_s2v.py @@ -322,8 +322,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -370,8 +370,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2_vace_fun/predict_v2v_control.py b/examples/wan2.2_vace_fun/predict_v2v_control.py index a063e4ff..5a227168 100644 --- a/examples/wan2.2_vace_fun/predict_v2v_control.py +++ b/examples/wan2.2_vace_fun/predict_v2v_control.py @@ -322,8 +322,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -370,8 +370,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2_vace_fun/predict_v2v_control_ref.py b/examples/wan2.2_vace_fun/predict_v2v_control_ref.py index b57099bf..e9a557f3 100644 --- a/examples/wan2.2_vace_fun/predict_v2v_control_ref.py +++ b/examples/wan2.2_vace_fun/predict_v2v_control_ref.py @@ -322,8 +322,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -370,8 +370,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/examples/wan2.2_vace_fun/predict_v2v_mask.py b/examples/wan2.2_vace_fun/predict_v2v_mask.py index dd28d98e..b25eb5ba 100644 --- a/examples/wan2.2_vace_fun/predict_v2v_mask.py +++ b/examples/wan2.2_vace_fun/predict_v2v_mask.py @@ -322,8 +322,8 @@ if lora_path is not None: pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") with torch.no_grad(): video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 @@ -370,8 +370,8 @@ if lora_path is not None: pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) - if transformer_2 is not None: - pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") +if lora_high_path is not None and transformer_2 is not None: + pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2") def save_results(): if not os.path.exists(save_path): diff --git a/scripts/cogvideox_fun/README_TRAIN.md b/scripts/cogvideox_fun/README_TRAIN.md index b9952263..fd7f88ba 100755 --- a/scripts/cogvideox_fun/README_TRAIN.md +++ b/scripts/cogvideox_fun/README_TRAIN.md @@ -45,9 +45,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -137,8 +138,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -147,7 +148,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -223,8 +224,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/CogVideoX-Fun-2b-InP` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size per GPU | 1 | | `--image_sample_size` | Maximum image training resolution | 512 | | `--video_sample_size` | Maximum video training resolution | 512 | diff --git a/scripts/cogvideox_fun/README_TRAIN_CONTROL.md b/scripts/cogvideox_fun/README_TRAIN_CONTROL.md index babc601e..849c681a 100755 --- a/scripts/cogvideox_fun/README_TRAIN_CONTROL.md +++ b/scripts/cogvideox_fun/README_TRAIN_CONTROL.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -145,8 +146,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Controls-Demo --local_dir ./datas If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -155,7 +156,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -231,8 +232,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/CogVideoX-Fun-V1.1-2b-Control` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Controls-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size per GPU | 4 | | `--image_sample_size` | Maximum image training resolution | 512 | | `--video_sample_size` | Maximum video training resolution | 512 | diff --git a/scripts/cogvideox_fun/README_TRAIN_CONTROL_zh-CN.md b/scripts/cogvideox_fun/README_TRAIN_CONTROL_zh-CN.md index 0131aca1..a67936ca 100755 --- a/scripts/cogvideox_fun/README_TRAIN_CONTROL_zh-CN.md +++ b/scripts/cogvideox_fun/README_TRAIN_CONTROL_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -145,8 +146,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Controls-Demo --local_dir ./datas 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -155,7 +156,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -231,8 +232,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/CogVideoX-Fun-V1.1-2b-Control` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Controls-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 4 | | `--image_sample_size` | 图像最大训练分辨率 | 512 | | `--video_sample_size` | 视频最大训练分辨率 | 512 | diff --git a/scripts/cogvideox_fun/README_TRAIN_LORA.md b/scripts/cogvideox_fun/README_TRAIN_LORA.md index e4a6c5f4..f41dc796 100755 --- a/scripts/cogvideox_fun/README_TRAIN_LORA.md +++ b/scripts/cogvideox_fun/README_TRAIN_LORA.md @@ -45,9 +45,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -137,8 +138,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -147,7 +148,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, relative paths are recommended. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, absolute paths are recommended. @@ -225,8 +226,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----------|-------------|---------------| | `--pretrained_model_name_or_path` | Pre-trained model path | `models/Diffusion_Transformer/CogVideoX-Fun-2b-InP` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size | 1 | | `--image_sample_size` | Maximum training resolution for images | 512 | | `--video_sample_size` | Maximum training resolution for videos | 512 | diff --git a/scripts/cogvideox_fun/README_TRAIN_LORA_zh-CN.md b/scripts/cogvideox_fun/README_TRAIN_LORA_zh-CN.md index 3d6f772a..67628d38 100644 --- a/scripts/cogvideox_fun/README_TRAIN_LORA_zh-CN.md +++ b/scripts/cogvideox_fun/README_TRAIN_LORA_zh-CN.md @@ -45,9 +45,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -137,8 +138,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -147,7 +148,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -225,8 +226,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/CogVideoX-Fun-2b-InP` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 512 | | `--video_sample_size` | 视频最大训练分辨率 | 512 | diff --git a/scripts/cogvideox_fun/README_TRAIN_zh-CN.md b/scripts/cogvideox_fun/README_TRAIN_zh-CN.md index 6855d4fb..fcf9efaa 100644 --- a/scripts/cogvideox_fun/README_TRAIN_zh-CN.md +++ b/scripts/cogvideox_fun/README_TRAIN_zh-CN.md @@ -45,9 +45,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -137,8 +138,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -147,7 +148,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -223,8 +224,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/CogVideoX-Fun-2b-InP` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 512 | | `--video_sample_size` | 视频最大训练分辨率 | 512 | diff --git a/scripts/cogvideox_fun/train_reward_lora.py b/scripts/cogvideox_fun/train_reward_lora.py index eecc9400..2d5d22c3 100755 --- a/scripts/cogvideox_fun/train_reward_lora.py +++ b/scripts/cogvideox_fun/train_reward_lora.py @@ -38,7 +38,10 @@ from accelerate.logging import get_logger from accelerate.state import AcceleratorState from accelerate.utils import ProjectConfiguration, set_seed -from decord import VideoReader +try: + from decord import VideoReader +except ImportError: + from videox_fun.data.utils import AVVideoReader as VideoReader from diffusers import CogVideoXDPMScheduler, DDIMScheduler from diffusers.optimization import get_scheduler from diffusers.utils import check_min_version, deprecate, is_wandb_available diff --git a/scripts/ernie_image/README_TRAIN.md b/scripts/ernie_image/README_TRAIN.md index 9703a20c..bcc0d0e6 100644 --- a/scripts/ernie_image/README_TRAIN.md +++ b/scripts/ernie_image/README_TRAIN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -139,7 +140,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -208,8 +209,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/ERNIE-Image` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--image_sample_size` | Maximum training resolution, auto bucketing | 1328 | | `--gradient_accumulation_steps` | Gradient accumulation steps (equivalent to larger batch) | 1 | @@ -267,8 +268,8 @@ You can configure validation parameters to periodically generate test images dur ```sh export MODEL_NAME="models/Diffusion_Transformer/ERNIE-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -306,8 +307,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/ERNIE-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/ernie_image/README_TRAIN_zh-CN.md b/scripts/ernie_image/README_TRAIN_zh-CN.md index db12d30b..4ebf0cb1 100644 --- a/scripts/ernie_image/README_TRAIN_zh-CN.md +++ b/scripts/ernie_image/README_TRAIN_zh-CN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -139,7 +140,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -208,8 +209,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/ERNIE-Image` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 最大训练分辨率,代码会自动分桶 | 1328 | | `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 | @@ -267,8 +268,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/ERNIE-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -306,8 +307,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/ERNIE-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/fantasytalking/README_TRAIN.md b/scripts/fantasytalking/README_TRAIN.md index 8552f962..526dc9db 100755 --- a/scripts/fantasytalking/README_TRAIN.md +++ b/scripts/fantasytalking/README_TRAIN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -149,8 +150,8 @@ The `metadata.json` for FantasyTalking is slightly different from the normal JSO If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -159,7 +160,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/path/to/your/metadata.json" +export DATASET_META_NAME="/path/to/your/metadata_add_width_height.json" ``` > 💡 **Tip**: If your dataset is small and stored locally, use relative paths. If your dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -249,8 +250,8 @@ Here is a detailed explanation of the key parameters in the training script: | `config_path` | Model configuration file path | `config/wan2.1/wan_civitai.yaml` | | `pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.1-I2V-14B-720P` | | `pretrained_audio_model_name_or_path` | Audio encoder path | `None` (automatically uses $MODEL_NAME/audio_encoder) | -| `train_data_dir` | Training dataset directory | `datasets/internal_datasets/` | -| `train_data_meta` | Training dataset metadata file | `datasets/internal_datasets/metadata.json` | +| `train_data_dir` | Training dataset directory | `datasets/X-Fun-Videos-Audios-Demo/` | +| `train_data_meta` | Training dataset metadata file | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `video_sample_size` | Video sample size (maximum resolution) | `512` | | `token_sample_size` | Token sample size | `512` | | `video_sample_stride` | Video sample stride | `1` | diff --git a/scripts/fantasytalking/README_TRAIN_zh-CN.md b/scripts/fantasytalking/README_TRAIN_zh-CN.md index 9afb7310..0283d05d 100755 --- a/scripts/fantasytalking/README_TRAIN_zh-CN.md +++ b/scripts/fantasytalking/README_TRAIN_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式三：使用 Docker** @@ -149,8 +150,8 @@ FantasyTalking 的 `metadata.json` 与 VideoX-Fun 的普通 JSON 格式略有不 如果你的数据使用的是相对路径，训练脚本中请这样配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -159,7 +160,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/path/to/your/metadata.json" +export DATASET_META_NAME="/path/to/your/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存放在本地，请使用相对路径。如果数据集存放在外部存储（如 NAS、OSS）或多机共享，请使用绝对路径。 @@ -249,8 +250,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `config_path` | 模型配置文件路径 | `config/wan2.1/wan_civitai.yaml` | | `pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-I2V-14B-720P` | | `pretrained_audio_model_name_or_path` | 音频编码器路径 | `None`（自动使用 $MODEL_NAME/audio_encoder） | -| `train_data_dir` | 训练数据集目录 | `datasets/internal_datasets/` | -| `train_data_meta` | 训练数据集元数据文件 | `datasets/internal_datasets/metadata.json` | +| `train_data_dir` | 训练数据集目录 | `datasets/X-Fun-Videos-Audios-Demo/` | +| `train_data_meta` | 训练数据集元数据文件 | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `video_sample_size` | 视频采样尺寸（最大分辨率） | `512` | | `token_sample_size` | Token 采样尺寸 | `512` | | `video_sample_stride` | 视频采样步幅 | `1` | diff --git a/scripts/flashhead/README_TRAIN_S2V.md b/scripts/flashhead/README_TRAIN_S2V.md index ab8cb609..2ef47d43 100644 --- a/scripts/flashhead/README_TRAIN_S2V.md +++ b/scripts/flashhead/README_TRAIN_S2V.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -164,7 +165,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -247,8 +248,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--config_path` | Model configuration file path | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/SoulX-FlashHead-1_3B` | | `--audio_encoder_path` | Audio encoder path (**FlashHead-S2V specific**) | `models/Diffusion_Transformer/wav2vec2-base-960h` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--video_sample_size` | Maximum video resolution for training | 512 | | `--token_sample_size` | Token length sampling size | 512 | diff --git a/scripts/flashhead/README_TRAIN_S2V_zh-CN.md b/scripts/flashhead/README_TRAIN_S2V_zh-CN.md index b880a647..4f717f2d 100644 --- a/scripts/flashhead/README_TRAIN_S2V_zh-CN.md +++ b/scripts/flashhead/README_TRAIN_S2V_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式三：使用 Docker** @@ -147,8 +148,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset 如果你的数据使用的是相对路径，训练脚本中请这样配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -157,7 +158,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存放在本地，请使用相对路径。如果数据集存放在外部存储（如 NAS、OSS）或多机共享，请使用绝对路径。 @@ -240,8 +241,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--config_path` | 模型配置文件路径 | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/SoulX-FlashHead-1_3B` | | `--audio_encoder_path` | 音频编码器路径（**FlashHead-S2V 特有**） | `models/Diffusion_Transformer/wav2vec2-base-960h` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元数据文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | 训练数据元数据文件 | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批训练的样本数 | 1 | | `--video_sample_size` | 视频最大训练分辨率 | 512 | | `--token_sample_size` | Token 长度采样大小 | 512 | diff --git a/scripts/flux/README_TRAIN.md b/scripts/flux/README_TRAIN.md index 08b966c4..4f7efe6f 100755 --- a/scripts/flux/README_TRAIN.md +++ b/scripts/flux/README_TRAIN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -139,7 +140,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -208,8 +209,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/FLUX.1-dev` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size per device | 1 | | `--image_sample_size` | Maximum training resolution (auto bucketing) | 1024 | | `--gradient_accumulation_steps` | Gradient accumulation steps (effective batch size increase) | 1 | @@ -266,8 +267,8 @@ You can configure validation parameters to periodically generate test images dur ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.1-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -305,8 +306,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.1-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/flux/README_TRAIN_LORA.md b/scripts/flux/README_TRAIN_LORA.md index 04837427..c32cd9c8 100755 --- a/scripts/flux/README_TRAIN_LORA.md +++ b/scripts/flux/README_TRAIN_LORA.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -139,7 +140,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -209,8 +210,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/FLUX.1-dev` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size per device | 1 | | `--image_sample_size` | Maximum training resolution (auto bucketing) | 1024 | | `--gradient_accumulation_steps` | Gradient accumulation steps (effective batch size increase) | 1 | @@ -267,8 +268,8 @@ You can configure validation parameters to periodically generate test images dur ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.1-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -307,8 +308,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.1-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/flux/README_TRAIN_LORA_zh-CN.md b/scripts/flux/README_TRAIN_LORA_zh-CN.md index 1d6c5ef1..3498e516 100755 --- a/scripts/flux/README_TRAIN_LORA_zh-CN.md +++ b/scripts/flux/README_TRAIN_LORA_zh-CN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun 如果你的数据使用相对路径，训练脚本中这样配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **使用绝对路径**： @@ -139,7 +140,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，使用相对路径。如果数据集存储在外部存储（如 NAS、OSS）或多机共享，使用绝对路径。 @@ -209,8 +210,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/FLUX.1-dev` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每张卡的批次大小 | 1 | | `--image_sample_size` | 最大训练分辨率（自动分桶） | 1024 | | `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch size） | 1 | @@ -267,8 +268,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.1-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -307,8 +308,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.1-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/flux/README_TRAIN_zh-CN.md b/scripts/flux/README_TRAIN_zh-CN.md index 99f5c0be..e2eb5f5c 100755 --- a/scripts/flux/README_TRAIN_zh-CN.md +++ b/scripts/flux/README_TRAIN_zh-CN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun 如果你的数据使用相对路径，训练脚本中这样配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **使用绝对路径**： @@ -139,7 +140,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，使用相对路径。如果数据集存储在外部存储（如 NAS、OSS）或多机共享，使用绝对路径。 @@ -208,8 +209,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/FLUX.1-dev` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每张卡的批次大小 | 1 | | `--image_sample_size` | 最大训练分辨率（自动分桶） | 1024 | | `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch size） | 1 | @@ -266,8 +267,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.1-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -305,8 +306,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.1-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/flux2/README_TRAIN.md b/scripts/flux2/README_TRAIN.md index ff479680..fa47dde6 100644 --- a/scripts/flux2/README_TRAIN.md +++ b/scripts/flux2/README_TRAIN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -139,7 +140,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -208,8 +209,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/FLUX.2-dev` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size per device | 1 | | `--image_sample_size` | Maximum training resolution (auto bucketing) | 1328 | | `--gradient_accumulation_steps` | Gradient accumulation steps (effective batch size increase) | 1 | @@ -266,8 +267,8 @@ You can configure validation parameters to periodically generate test images dur ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.2-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -305,8 +306,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.2-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/flux2/README_TRAIN_LORA.md b/scripts/flux2/README_TRAIN_LORA.md index 1f4738c9..14f19e33 100644 --- a/scripts/flux2/README_TRAIN_LORA.md +++ b/scripts/flux2/README_TRAIN_LORA.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -139,7 +140,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -209,8 +210,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/FLUX.2-dev` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size per device | 1 | | `--image_sample_size` | Maximum training resolution (auto bucketing) | 1328 | | `--gradient_accumulation_steps` | Gradient accumulation steps (effective batch size increase) | 1 | @@ -266,8 +267,8 @@ You can configure validation parameters to periodically generate test images dur ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.2-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -306,8 +307,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.2-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/flux2/README_TRAIN_LORA_zh-CN.md b/scripts/flux2/README_TRAIN_LORA_zh-CN.md index 514a8a7c..f1f2ad47 100644 --- a/scripts/flux2/README_TRAIN_LORA_zh-CN.md +++ b/scripts/flux2/README_TRAIN_LORA_zh-CN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -139,7 +140,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -209,8 +210,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/FLUX.2-dev` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，代码会自动分桶 | 1328 | | `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 | @@ -266,8 +267,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.2-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -306,8 +307,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.2-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/flux2/README_TRAIN_zh-CN.md b/scripts/flux2/README_TRAIN_zh-CN.md index c27d067d..19e8c1d6 100644 --- a/scripts/flux2/README_TRAIN_zh-CN.md +++ b/scripts/flux2/README_TRAIN_zh-CN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -139,7 +140,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -208,8 +209,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/FLUX.2-dev` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，代码会自动分桶 | 1328 | | `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 | @@ -266,8 +267,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.2-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -305,8 +306,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/FLUX.2-dev" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/flux2_fun/README_TRAIN.md b/scripts/flux2_fun/README_TRAIN.md index cbaef25d..64e76842 100644 --- a/scripts/flux2_fun/README_TRAIN.md +++ b/scripts/flux2_fun/README_TRAIN.md @@ -45,9 +45,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -144,45 +145,25 @@ It is recommended to use tools like [DWPose](https://github.com/IDEA-Research/DW ### 2.4 Relative vs Absolute Path Usage -**Option 1: Using Relative Paths (Recommended)** +**Relative Paths**: -When data paths are not fixed or need to be trained on different machines, relative paths are recommended. +If your data uses relative paths, set in the training script: -Configure relative paths in `metadata.json`, then specify the dataset root directory via `--train_data_dir` in the training script: - -```json -[ - { - "file_path": "train/image001.jpg", - "control_file_path": "control/image001.jpg", - "text": "A group of young men in suits and sunglasses are walking down a city street.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata.json" ``` -The training process will automatically search for files corresponding to the relative paths under `--train_data_dir`. +**Absolute Paths**: -**Option 2: Using Absolute Paths** +If your data uses absolute paths, set in the training script: -If the dataset path is fixed, you can directly configure absolute paths in `metadata.json`: - -```json -[ - { - "file_path": "/mnt/data/images/image001.jpg", - "control_file_path": "/mnt/data/controls/image001.jpg", - "text": "A group of young men in suits and sunglasses.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="" +export DATASET_META_NAME="/mnt/data/metadata.json" ``` -When using absolute paths, the `--train_data_dir` parameter serves only as a default path, and the absolute paths in the JSON will take priority. +> 💡 **Recommendation**: If the dataset is small and stored locally, relative paths are recommended; if the dataset is stored on external storage (such as NAS, OSS) or shared across multiple machines, absolute paths are recommended. --- diff --git a/scripts/flux2_fun/README_TRAIN_zh-CN.md b/scripts/flux2_fun/README_TRAIN_zh-CN.md index c29848d3..74b3a3d8 100644 --- a/scripts/flux2_fun/README_TRAIN_zh-CN.md +++ b/scripts/flux2_fun/README_TRAIN_zh-CN.md @@ -45,9 +45,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -144,45 +145,25 @@ Control 模式的 metadata.json 与普通 FLUX.2 的 json 略有不同，需要 ### 2.4 相对路径与绝对路径使用方案 -**方案 1：使用相对路径（推荐）** +**相对路径**： -当数据路径不固定，或需要在不同机器上训练时，推荐使用相对路径。 +如果数据的路径为相对路径，则在训练脚本中设置： -在 `metadata.json` 中配置相对路径，然后在训练脚本中通过 `--train_data_dir` 指定数据集根目录： - -```json -[ - { - "file_path": "train/image001.jpg", - "control_file_path": "control/image001.jpg", - "text": "A group of young men in suits and sunglasses are walking down a city street.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata.json" ``` -训练时会自动在 `--train_data_dir` 下寻找相对路径对应的文件。 +**绝对路径**： -**方案 2：使用绝对路径** +如果数据的路径为绝对路径，则在训练脚本中设置： -如果数据集路径固定，可以直接在 `metadata.json` 中配置绝对路径： - -```json -[ - { - "file_path": "/mnt/data/images/image001.jpg", - "control_file_path": "/mnt/data/controls/image001.jpg", - "text": "A group of young men in suits and sunglasses.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="" +export DATASET_META_NAME="/mnt/data/metadata.json" ``` -使用绝对路径时，`--train_data_dir` 参数仅作为默认路径，实际会优先使用 json 中的绝对路径。 +> 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 --- diff --git a/scripts/get_rewards_1video.py b/scripts/get_rewards_1video.py index 661f7a45..691539ac 100644 --- a/scripts/get_rewards_1video.py +++ b/scripts/get_rewards_1video.py @@ -77,7 +77,10 @@ import numpy as np import torch import torchvision.transforms as transforms -from decord import VideoReader +try: + from decord import VideoReader +except ImportError: + from videox_fun.data.utils import AVVideoReader as VideoReader from einops import rearrange project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) diff --git a/scripts/hunyuanvideo/README_TRAIN.md b/scripts/hunyuanvideo/README_TRAIN.md index e31445e4..3498bc3f 100644 --- a/scripts/hunyuanvideo/README_TRAIN.md +++ b/scripts/hunyuanvideo/README_TRAIN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -146,7 +147,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, relative paths are recommended. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, absolute paths are recommended. @@ -224,8 +225,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/HunyuanVideo` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Number of samples per batch | 16 | | `--image_sample_size` | Maximum training resolution for images | 512 | | `--video_sample_size` | Maximum training resolution for videos | 512 | diff --git a/scripts/hunyuanvideo/README_TRAIN_LORA.md b/scripts/hunyuanvideo/README_TRAIN_LORA.md index c375c7e4..fdf6dcaf 100644 --- a/scripts/hunyuanvideo/README_TRAIN_LORA.md +++ b/scripts/hunyuanvideo/README_TRAIN_LORA.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, set in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute paths**: @@ -146,7 +147,7 @@ If your data uses absolute paths, set in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, relative paths are recommended. If the dataset is stored on external storage (such as NAS, OSS) or shared across multiple machines, absolute paths are recommended. @@ -224,8 +225,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----------|-------------|---------------| | `--pretrained_model_name_or_path` | Pre-trained model path | `models/Diffusion_Transformer/HunyuanVideo` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Number of samples per batch | 1 | | `--image_sample_size` | Maximum training resolution for images | 640 | | `--video_sample_size` | Maximum training resolution for videos | 640 | diff --git a/scripts/hunyuanvideo/README_TRAIN_LORA_zh-CN.md b/scripts/hunyuanvideo/README_TRAIN_LORA_zh-CN.md index 47393fa7..c43393db 100644 --- a/scripts/hunyuanvideo/README_TRAIN_LORA_zh-CN.md +++ b/scripts/hunyuanvideo/README_TRAIN_LORA_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -146,7 +147,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -224,8 +225,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/HunyuanVideo` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | diff --git a/scripts/hunyuanvideo/README_TRAIN_zh-CN.md b/scripts/hunyuanvideo/README_TRAIN_zh-CN.md index 46b7ec03..fe816bba 100644 --- a/scripts/hunyuanvideo/README_TRAIN_zh-CN.md +++ b/scripts/hunyuanvideo/README_TRAIN_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -146,7 +147,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -224,8 +225,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/HunyuanVideo` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 16 | | `--image_sample_size` | 图像最大训练分辨率 | 512 | | `--video_sample_size` | 视频最大训练分辨率 | 512 | diff --git a/scripts/infinitetalk/README_TRAIN_S2V.md b/scripts/infinitetalk/README_TRAIN_S2V.md index 3eaf9646..225cda29 100644 --- a/scripts/infinitetalk/README_TRAIN_S2V.md +++ b/scripts/infinitetalk/README_TRAIN_S2V.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -164,7 +165,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -253,8 +254,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--config_path` | Model configuration file path | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/Wan2.1-I2V-14B-480P` | | `--audio_encoder_path` | Audio encoder path (**InfiniteTalk-S2V specific**) | `models/Diffusion_Transformer/chinese-wav2vec2-base/` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--video_sample_size` | Maximum video resolution for training | 640 | | `--token_sample_size` | Token length sampling size | 640 | diff --git a/scripts/infinitetalk/README_TRAIN_S2V_zh-CN.md b/scripts/infinitetalk/README_TRAIN_S2V_zh-CN.md index 6d8e04a7..3c996e1a 100644 --- a/scripts/infinitetalk/README_TRAIN_S2V_zh-CN.md +++ b/scripts/infinitetalk/README_TRAIN_S2V_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式三：使用 Docker** @@ -147,8 +148,8 @@ InfiniteTalk 的 `metadata.json` 与 VideoX-Fun 的普通 JSON 格式略有不 如果你的数据使用的是相对路径，训练脚本中请这样配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -157,7 +158,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存放在本地，请使用相对路径。如果数据集存放在外部存储（如 NAS、OSS）或多机共享，请使用绝对路径。 @@ -248,8 +249,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--config_path` | 模型配置文件路径 | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-I2V-14B-480P` | | `--audio_encoder_path` | 音频编码器路径（**InfiniteTalk 特有**） | `models/Diffusion_Transformer/chinese-wav2vec2-base/` | -| `--train_data_dir` | 训练数据集目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据集元数据文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据集目录 | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | 训练数据集元数据文件 | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--video_sample_size` | 视频采样尺寸（最大分辨率） | `640` | | `--token_sample_size` | Token 采样尺寸 | `640` | | `--video_sample_stride` | 视频采样步幅 | `1` | diff --git a/scripts/lens/README_TRAIN.md b/scripts/lens/README_TRAIN.md index 279ca317..fdd95455 100644 --- a/scripts/lens/README_TRAIN.md +++ b/scripts/lens/README_TRAIN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -139,7 +140,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -208,8 +209,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/Lens` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--image_sample_size` | Maximum training resolution, auto bucketing | 1328 | | `--gradient_accumulation_steps` | Gradient accumulation steps (equivalent to larger batch) | 1 | @@ -267,8 +268,8 @@ You can configure validation parameters to periodically generate test images dur ```sh export MODEL_NAME="models/Diffusion_Transformer/Lens" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -306,8 +307,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/Lens" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/lens/README_TRAIN_LORA.md b/scripts/lens/README_TRAIN_LORA.md index 11bedcd2..4d650088 100755 --- a/scripts/lens/README_TRAIN_LORA.md +++ b/scripts/lens/README_TRAIN_LORA.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -139,7 +140,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -209,8 +210,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/Lens` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--image_sample_size` | Maximum training resolution, auto bucketing | 1328 | | `--gradient_accumulation_steps` | Gradient accumulation steps (equivalent to larger batch) | 1 | @@ -269,8 +270,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/Lens" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -309,8 +310,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/Lens" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/lens/README_TRAIN_LORA_zh-CN.md b/scripts/lens/README_TRAIN_LORA_zh-CN.md index 6d517a71..816c01ed 100644 --- a/scripts/lens/README_TRAIN_LORA_zh-CN.md +++ b/scripts/lens/README_TRAIN_LORA_zh-CN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -139,7 +140,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -209,8 +210,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Lens` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，代码会自动分桶 | 1328 | | `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 | @@ -270,8 +271,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/Lens" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -310,8 +311,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/Lens" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/lens/README_TRAIN_zh-CN.md b/scripts/lens/README_TRAIN_zh-CN.md index 544b073f..03514fe1 100644 --- a/scripts/lens/README_TRAIN_zh-CN.md +++ b/scripts/lens/README_TRAIN_zh-CN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -139,7 +140,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -208,8 +209,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Lens` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 最大训练分辨率,代码会自动分桶 | 1328 | | `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 | @@ -267,8 +268,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/Lens" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -306,8 +307,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/Lens" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/longcatvideo/README_TRAIN.md b/scripts/longcatvideo/README_TRAIN.md index cf1a7a38..4f3354e2 100644 --- a/scripts/longcatvideo/README_TRAIN.md +++ b/scripts/longcatvideo/README_TRAIN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -146,7 +147,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -224,8 +225,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/LongCat-Video` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size per GPU | 1 | | `--image_sample_size` | Maximum image training resolution | 640 | | `--video_sample_size` | Maximum video training resolution | 640 | diff --git a/scripts/longcatvideo/README_TRAIN_AVATAR.md b/scripts/longcatvideo/README_TRAIN_AVATAR.md index 80425677..7824a087 100644 --- a/scripts/longcatvideo/README_TRAIN_AVATAR.md +++ b/scripts/longcatvideo/README_TRAIN_AVATAR.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -164,7 +165,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -247,8 +248,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con |-----|------|-------| | `--pretrained_model_name_or_path` | Path to pretrained base model | `models/Diffusion_Transformer/LongCat-Video` | | `--pretrained_avatar_model_name_or_path` | Path to pretrained avatar model | `models/Diffusion_Transformer/LongCat-Video-Avatar` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--video_sample_size` | Maximum video resolution for training | 640 | | `--token_sample_size` | Token length sampling size | 640 | diff --git a/scripts/longcatvideo/README_TRAIN_AVATAR_LORA.md b/scripts/longcatvideo/README_TRAIN_AVATAR_LORA.md index 924c328b..c68cc25d 100644 --- a/scripts/longcatvideo/README_TRAIN_AVATAR_LORA.md +++ b/scripts/longcatvideo/README_TRAIN_AVATAR_LORA.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -164,7 +165,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -256,8 +257,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con |-----|------|-------| | `--pretrained_model_name_or_path` | Path to pretrained base model | `models/Diffusion_Transformer/LongCat-Video` | | `--pretrained_avatar_model_name_or_path` | Path to pretrained avatar model | `models/Diffusion_Transformer/LongCat-Video-Avatar` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--image_sample_size` | Maximum training resolution, auto bucketing | 640 | | `--video_sample_size` | Maximum video resolution for training | 640 | diff --git a/scripts/longcatvideo/README_TRAIN_AVATAR_LORA_zh-CN.md b/scripts/longcatvideo/README_TRAIN_AVATAR_LORA_zh-CN.md index 1fc3f058..2befaa41 100644 --- a/scripts/longcatvideo/README_TRAIN_AVATAR_LORA_zh-CN.md +++ b/scripts/longcatvideo/README_TRAIN_AVATAR_LORA_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式三：使用 Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset 如果你的数据使用的是相对路径，训练脚本中请这样配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -164,7 +165,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存放在本地，请使用相对路径。如果数据集存放在外部存储（如 NAS、OSS）或多机共享，请使用绝对路径。 @@ -256,8 +257,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练基础模型路径 | `models/Diffusion_Transformer/LongCat-Video` | | `--pretrained_avatar_model_name_or_path` | 预训练 Avatar 模型路径 | `models/Diffusion_Transformer/LongCat-Video-Avatar` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元数据文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | 训练数据元数据文件 | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批训练的样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，自动 bucket | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | diff --git a/scripts/longcatvideo/README_TRAIN_AVATAR_zh-CN.md b/scripts/longcatvideo/README_TRAIN_AVATAR_zh-CN.md index 6ffa6cfc..d7ef5a16 100644 --- a/scripts/longcatvideo/README_TRAIN_AVATAR_zh-CN.md +++ b/scripts/longcatvideo/README_TRAIN_AVATAR_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式三：使用 Docker** @@ -147,8 +148,8 @@ LongCat-Video-Avatar 的 `metadata.json` 与 VideoX-Fun 的普通 JSON 格式略 如果你的数据使用的是相对路径，训练脚本中请这样配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -157,7 +158,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存放在本地，请使用相对路径。如果数据集存放在外部存储（如 NAS、OSS）或多机共享，请使用绝对路径。 @@ -242,8 +243,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con |------|------|--------| | `--pretrained_model_name_or_path` | 预训练基础模型路径 | `models/Diffusion_Transformer/LongCat-Video` | | `--pretrained_avatar_model_name_or_path` | 预训练 Avatar 模型路径 | `models/Diffusion_Transformer/LongCat-Video-Avatar` | -| `--train_data_dir` | 训练数据集目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据集元数据文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据集目录 | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | 训练数据集元数据文件 | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--video_sample_size` | 视频采样尺寸（最大分辨率） | `640` | | `--token_sample_size` | Token 采样尺寸 | `640` | | `--video_sample_stride` | 视频采样步幅 | `2` | diff --git a/scripts/longcatvideo/README_TRAIN_LORA.md b/scripts/longcatvideo/README_TRAIN_LORA.md index 965de392..8708ae63 100644 --- a/scripts/longcatvideo/README_TRAIN_LORA.md +++ b/scripts/longcatvideo/README_TRAIN_LORA.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -146,7 +147,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, relative paths are recommended. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, absolute paths are recommended. @@ -225,8 +226,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----------|-------------|---------------| | `--pretrained_model_name_or_path` | Pre-trained model path | `models/Diffusion_Transformer/LongCat-Video` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size | 1 | | `--image_sample_size` | Maximum training resolution for images | 640 | | `--video_sample_size` | Maximum training resolution for videos | 640 | diff --git a/scripts/longcatvideo/README_TRAIN_LORA_zh-CN.md b/scripts/longcatvideo/README_TRAIN_LORA_zh-CN.md index 1b5d7a41..915ae97b 100644 --- a/scripts/longcatvideo/README_TRAIN_LORA_zh-CN.md +++ b/scripts/longcatvideo/README_TRAIN_LORA_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -146,7 +147,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -225,8 +226,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/LongCat-Video` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | diff --git a/scripts/longcatvideo/README_TRAIN_zh-CN.md b/scripts/longcatvideo/README_TRAIN_zh-CN.md index 0d96fb3f..80103b79 100644 --- a/scripts/longcatvideo/README_TRAIN_zh-CN.md +++ b/scripts/longcatvideo/README_TRAIN_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -146,7 +147,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -224,8 +225,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/LongCat-Video` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | diff --git a/scripts/ltx2.3/README_TRAIN.md b/scripts/ltx2.3/README_TRAIN.md index 5b425d1a..eb25bcd2 100644 --- a/scripts/ltx2.3/README_TRAIN.md +++ b/scripts/ltx2.3/README_TRAIN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -164,7 +165,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -241,8 +242,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/LTX-2.3-Diffusers` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--image_sample_size` | Maximum training resolution, auto bucketing | 640 | | `--video_sample_size` | Maximum video resolution for training | 640 | @@ -307,8 +308,8 @@ You can configure validation parameters to periodically generate test videos dur ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2.3-Diffusers" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -354,8 +355,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2.3-Diffusers" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/ltx2.3/README_TRAIN_LORA.md b/scripts/ltx2.3/README_TRAIN_LORA.md index dce5006c..b372f0f6 100644 --- a/scripts/ltx2.3/README_TRAIN_LORA.md +++ b/scripts/ltx2.3/README_TRAIN_LORA.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -164,7 +165,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -251,8 +252,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/LTX-2.3-Diffusers` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--image_sample_size` | Maximum training resolution, auto bucketing | 640 | | `--video_sample_size` | Maximum video resolution for training | 640 | @@ -314,8 +315,8 @@ You can configure validation parameters to periodically generate test videos dur ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2.3-Diffusers" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -362,8 +363,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2.3-Diffusers" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/ltx2.3/README_TRAIN_LORA_zh-CN.md b/scripts/ltx2.3/README_TRAIN_LORA_zh-CN.md index 6b847af2..06420988 100644 --- a/scripts/ltx2.3/README_TRAIN_LORA_zh-CN.md +++ b/scripts/ltx2.3/README_TRAIN_LORA_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式三：使用 Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset 如果你的数据使用的是相对路径，训练脚本中请这样配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -164,7 +165,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存放在本地，请使用相对路径。如果数据集存放在外部存储（如 NAS、OSS）或多机共享，请使用绝对路径。 @@ -251,8 +252,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/LTX-2` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元数据文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | 训练数据元数据文件 | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批训练的样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，自动 bucket | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | @@ -314,8 +315,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -362,8 +363,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/ltx2.3/README_TRAIN_zh-CN.md b/scripts/ltx2.3/README_TRAIN_zh-CN.md index d3d9d5eb..12722fcc 100644 --- a/scripts/ltx2.3/README_TRAIN_zh-CN.md +++ b/scripts/ltx2.3/README_TRAIN_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式三：使用 Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset 如果你的数据使用的是相对路径，训练脚本中请这样配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -164,7 +165,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存放在本地，请使用相对路径。如果数据集存放在外部存储（如 NAS、OSS）或多机共享，请使用绝对路径。 @@ -241,8 +242,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/LTX-2.3-Diffusers` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元数据文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | 训练数据元数据文件 | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批训练的样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，自动 bucket | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | @@ -307,8 +308,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2.3-Diffusers" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -354,8 +355,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2.3-Diffusers" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/ltx2/README_TRAIN.md b/scripts/ltx2/README_TRAIN.md index cbeab9c1..8decc462 100644 --- a/scripts/ltx2/README_TRAIN.md +++ b/scripts/ltx2/README_TRAIN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -164,7 +165,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -244,8 +245,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/LTX-2` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--image_sample_size` | Maximum training resolution, auto bucketing | 640 | | `--video_sample_size` | Maximum video resolution for training | 640 | @@ -310,8 +311,8 @@ You can configure validation parameters to periodically generate test videos dur ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -357,8 +358,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/ltx2/README_TRAIN_LORA.md b/scripts/ltx2/README_TRAIN_LORA.md index 3701f981..6011a8d4 100644 --- a/scripts/ltx2/README_TRAIN_LORA.md +++ b/scripts/ltx2/README_TRAIN_LORA.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -164,7 +165,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -254,8 +255,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/LTX-2` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--image_sample_size` | Maximum training resolution, auto bucketing | 640 | | `--video_sample_size` | Maximum video resolution for training | 640 | @@ -317,8 +318,8 @@ You can configure validation parameters to periodically generate test videos dur ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -365,8 +366,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/ltx2/README_TRAIN_LORA_zh-CN.md b/scripts/ltx2/README_TRAIN_LORA_zh-CN.md index 4f4e0766..bda37ff7 100644 --- a/scripts/ltx2/README_TRAIN_LORA_zh-CN.md +++ b/scripts/ltx2/README_TRAIN_LORA_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式三：使用 Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset 如果你的数据使用的是相对路径，训练脚本中请这样配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -164,7 +165,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存放在本地，请使用相对路径。如果数据集存放在外部存储（如 NAS、OSS）或多机共享，请使用绝对路径。 @@ -254,8 +255,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/LTX-2` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元数据文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | 训练数据元数据文件 | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批训练的样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，自动 bucket | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | @@ -317,8 +318,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -365,8 +366,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/ltx2/README_TRAIN_zh-CN.md b/scripts/ltx2/README_TRAIN_zh-CN.md index 5a870e5d..83b9e819 100644 --- a/scripts/ltx2/README_TRAIN_zh-CN.md +++ b/scripts/ltx2/README_TRAIN_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式三：使用 Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset 如果你的数据使用的是相对路径，训练脚本中请这样配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -164,7 +165,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存放在本地，请使用相对路径。如果数据集存放在外部存储（如 NAS、OSS）或多机共享，请使用绝对路径。 @@ -244,8 +245,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/LTX-2` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元数据文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | 训练数据元数据文件 | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批训练的样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，自动 bucket | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | @@ -310,8 +311,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -357,8 +358,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/LTX-2" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/ltx2/train_upsampler.py b/scripts/ltx2/train_upsampler.py index a8290824..477bc345 100644 --- a/scripts/ltx2/train_upsampler.py +++ b/scripts/ltx2/train_upsampler.py @@ -124,7 +124,10 @@ def log_validation(vae, latent_upsampler, args, accelerator, weight_dtype, globa logger.info("No validation_paths provided, skipping validation.") return - from decord import VideoReader + try: + from decord import VideoReader + except ImportError: + from videox_fun.data.utils import AVVideoReader as VideoReader for i, video_path in enumerate(args.validation_paths): if not os.path.exists(video_path): @@ -159,8 +162,8 @@ def log_validation(vae, latent_upsampler, args, accelerator, weight_dtype, globa # Encode high-res gt_latents = vae.encode(pixel_values)[0].mode() - # Create low-res input - scale = args.spatial_scale + # Create low-res input (scale is fixed by upsampler architecture) + scale = float(accelerator.unwrap_model(latent_upsampler).config.rational_spatial_scale) low_h, low_w = int(h_target / scale), int(w_target / scale) # Downsample spatially: flatten batch and frames, interpolate, unflatten b, c, f, h, w = pixel_values.shape @@ -488,12 +491,6 @@ def parse_args(): default=None, help=("If you want to load the weight from other vaes, input its path."), ) - parser.add_argument( - "--spatial_scale", - type=float, - default=2.0, - help="Spatial upsampling scale factor (must match model config).", - ) parser.add_argument( '--trainable_modules', nargs='+', @@ -664,6 +661,11 @@ def main(): m, u = latent_upsampler.load_state_dict(state_dict, strict=False) print(f"Upsampler missing keys: {len(m)}, unexpected keys: {len(u)}") + # Spatial scale is fixed by the upsampler architecture; read it directly from the model config. + spatial_scale = float(latent_upsampler.config.rational_spatial_scale) + if accelerator.is_main_process: + logger.info(f"Using spatial_scale={spatial_scale} from latent_upsampler.config.rational_spatial_scale.") + # Set trainable parameters latent_upsampler.requires_grad_(False) latent_upsampler.train() @@ -857,7 +859,7 @@ def collate_fn(examples): batch_video_length = 1 # Compute low-res target size (aligned to spatial_compression_ratio) - scale = args.spatial_scale + scale = spatial_scale spatial_ratio = vae.config.spatial_compression_ratio closest_size_list = list(map(lambda x: int(x), closest_size)) low_h = int(closest_size_list[0] / scale / spatial_ratio) * spatial_ratio diff --git a/scripts/mova/README_TRAIN.md b/scripts/mova/README_TRAIN.md index 6e9cce0f..17f57661 100644 --- a/scripts/mova/README_TRAIN.md +++ b/scripts/mova/README_TRAIN.md @@ -47,20 +47,13 @@ pip install Pillow einops safetensors timm tomesd librosa transformers accelerat ### 2.1 Quick Test Dataset -For testing purposes, you can use the demo dataset from Hugging Face: +We provide a test dataset containing several video-audio training samples. ```bash -# Create dataset directory -mkdir -p datasets - -# Download demo dataset from Hugging Face -cd datasets -git clone https://huggingface.co/datasets/modelscope/X-Fun-Videos-Audios-Demo -cd .. +# Download demo dataset +modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./datasets/X-Fun-Videos-Audios-Demo ``` -This dataset contains 17 video-audio pairs for quick testing. - ### 2.2 Dataset Structure ``` @@ -133,8 +126,8 @@ This dataset contains 17 video-audio pairs for quick testing. If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -143,7 +136,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Tip**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -169,7 +162,7 @@ FSDP training is recommended as it can significantly save VRAM. ```bash export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi-node environments without RDMA # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -219,8 +212,8 @@ accelerate launch --mixed_precision="bf16" scripts/mova/train.py \ | Parameter | Description | Example Value | |-----------|-------------|---------------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/MOVA-360p` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Number of samples per batch | 1 | | `--image_sample_size` | Maximum training resolution for auto bucket | 480 | | `--video_sample_size` | Maximum video training resolution | 480 | @@ -269,8 +262,8 @@ Edit the script to load your trained checkpoint and generate test videos. ```bash export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -320,8 +313,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```bash export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/mova/README_TRAIN_LORA.md b/scripts/mova/README_TRAIN_LORA.md index 7c0fc151..97534e24 100644 --- a/scripts/mova/README_TRAIN_LORA.md +++ b/scripts/mova/README_TRAIN_LORA.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -164,7 +165,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Tip**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -192,7 +193,7 @@ The difference between DeepSpeed-Zero-2 and FSDP is whether model weights are sh ```bash export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi-node environments without RDMA # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -254,8 +255,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----------|-------------|---------------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/MOVA-360p` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Number of samples per batch | 1 | | `--image_sample_size` | Maximum training resolution for auto bucket | 360 | | `--video_sample_size` | Maximum video training resolution | 360 | @@ -337,8 +338,8 @@ You can configure validation parameters to regularly generate test videos during ```sh export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -394,8 +395,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -451,7 +452,7 @@ Assuming 2 machines, each with 8 GPUs: ```bash export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" export MASTER_ADDR="192.168.1.100" # Master machine IP export MASTER_PORT=10086 export WORLD_SIZE=2 # Total number of machines @@ -504,7 +505,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main ```bash export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" export MASTER_ADDR="192.168.1.100" # Same as Master export MASTER_PORT=10086 export WORLD_SIZE=2 diff --git a/scripts/mova/README_TRAIN_LORA_zh-CN.md b/scripts/mova/README_TRAIN_LORA_zh-CN.md index 282fdcf6..6a54ff50 100644 --- a/scripts/mova/README_TRAIN_LORA_zh-CN.md +++ b/scripts/mova/README_TRAIN_LORA_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式三：使用 Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset 如果你的数据使用的是相对路径，训练脚本中请这样配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -164,7 +165,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存放在本地，请使用相对路径。如果数据集存放在外部存储（如 NAS、OSS）或多机共享，请使用绝对路径。 @@ -192,7 +193,7 @@ DeepSpeed-Zero-2 与 FSDP 的区别在于模型权重是否分片。**如果多 ```bash export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -254,8 +255,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/MOVA-360p` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元数据文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | 训练数据元数据文件 | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批训练的样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，自动 bucket | 360 | | `--video_sample_size` | 视频最大训练分辨率 | 360 | @@ -338,8 +339,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -395,8 +396,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -452,7 +453,7 @@ accelerate launch --mixed_precision="bf16" scripts/mova/train_lora.py \ ```bash export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" export MASTER_ADDR="192.168.1.100" # Master 机器 IP export MASTER_PORT=10086 export WORLD_SIZE=2 # 总机器数 @@ -505,7 +506,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main ```bash export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" export MASTER_ADDR="192.168.1.100" # 与 Master 相同 export MASTER_PORT=10086 export WORLD_SIZE=2 diff --git a/scripts/mova/README_TRAIN_zh-CN.md b/scripts/mova/README_TRAIN_zh-CN.md index 2a5e0d41..47af20a2 100644 --- a/scripts/mova/README_TRAIN_zh-CN.md +++ b/scripts/mova/README_TRAIN_zh-CN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式三：使用 Docker** @@ -152,8 +153,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset 如果你的数据使用的是相对路径，训练脚本中请这样配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -162,7 +163,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存放在本地，请使用相对路径。如果数据集存放在外部存储（如 NAS、OSS）或多机共享，请使用绝对路径。 @@ -190,7 +191,7 @@ DeepSpeed-Zero-2 与 FSDP 的区别在于模型权重是否分片。**如果多 ```bash export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -240,8 +241,8 @@ accelerate launch --mixed_precision="bf16" scripts/mova/train.py \ | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/MOVA-360p` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元数据文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | 训练数据元数据文件 | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批训练的样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，自动 bucket | 480 | | `--video_sample_size` | 视频最大训练分辨率 | 480 | @@ -313,8 +314,8 @@ accelerate launch --mixed_precision="bf16" scripts/mova/train.py \ ```sh export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -363,8 +364,8 @@ accelerate launch --mixed_precision="bf16" scripts/mova/train.py \ ```sh export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 和 NCCL_P2P_DISABLE=1 用于无 RDMA 的多机环境 # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -422,7 +423,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```bash export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" export MASTER_ADDR="192.168.1.100" # Master 机器 IP export MASTER_PORT=10086 export WORLD_SIZE=2 # 总机器数 @@ -474,7 +475,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main ```bash export MODEL_NAME="models/Diffusion_Transformer/MOVA-360p" export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" export MASTER_ADDR="192.168.1.100" # 与 Master 相同 export MASTER_PORT=10086 export WORLD_SIZE=2 diff --git a/scripts/qwenimage/README_TRAIN.md b/scripts/qwenimage/README_TRAIN.md index ac1674f6..672c9ba8 100755 --- a/scripts/qwenimage/README_TRAIN.md +++ b/scripts/qwenimage/README_TRAIN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -139,7 +140,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -208,8 +209,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/Qwen-Image` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--image_sample_size` | Maximum training resolution, auto bucketing | 1328 | | `--gradient_accumulation_steps` | Gradient accumulation steps (equivalent to larger batch) | 1 | @@ -270,8 +271,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -320,8 +321,8 @@ python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/ Training shell command: ```sh export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -359,8 +360,8 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag ```sh export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/qwenimage/README_TRAIN_EDIT.md b/scripts/qwenimage/README_TRAIN_EDIT.md index f274b767..9dd08d23 100755 --- a/scripts/qwenimage/README_TRAIN_EDIT.md +++ b/scripts/qwenimage/README_TRAIN_EDIT.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -128,8 +129,8 @@ The metadata.json for Edit model is different from the normal version, requiring If your data uses relative paths, set in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Edit-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Edit-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -138,7 +139,7 @@ If your data uses absolute paths, set in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, relative paths are recommended; if the dataset is stored on external storage (such as NAS, OSS) or shared across multiple machines, absolute paths are recommended. @@ -216,8 +217,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Pre-trained model path | `models/Diffusion_Transformer/Qwen-Image-Edit` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Images-Edit-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Images-Edit-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Number of samples per batch | 1 | | `--image_sample_size` | Maximum training resolution, code will automatically bucket | 1328 | | `--gradient_accumulation_steps` | Gradient accumulation steps (equivalent to increasing batch) | 1 | diff --git a/scripts/qwenimage/README_TRAIN_EDIT_zh-CN.md b/scripts/qwenimage/README_TRAIN_EDIT_zh-CN.md index db7dc150..643519dc 100644 --- a/scripts/qwenimage/README_TRAIN_EDIT_zh-CN.md +++ b/scripts/qwenimage/README_TRAIN_EDIT_zh-CN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -128,8 +129,8 @@ Edit 模型的 metadata.json 与普通版本不同，需要添加 `source_file_p 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Edit-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Edit-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -138,7 +139,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -216,8 +217,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Qwen-Image-Edit` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Images-Edit-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Images-Edit-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，代码会自动分桶 | 1328 | | `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 | diff --git a/scripts/qwenimage/README_TRAIN_LORA.md b/scripts/qwenimage/README_TRAIN_LORA.md index 3414401b..da3e1762 100755 --- a/scripts/qwenimage/README_TRAIN_LORA.md +++ b/scripts/qwenimage/README_TRAIN_LORA.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -139,7 +140,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -209,8 +210,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/Qwen-Image` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--image_sample_size` | Maximum training resolution, auto bucketing | 1328 | | `--gradient_accumulation_steps` | Gradient accumulation steps (equivalent to larger batch) | 1 | @@ -269,8 +270,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -320,8 +321,8 @@ python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/ Training shell command: ```sh export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -360,8 +361,8 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag ```sh export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/qwenimage/README_TRAIN_LORA_zh-CN.md b/scripts/qwenimage/README_TRAIN_LORA_zh-CN.md index 65dfb0ec..f9154833 100644 --- a/scripts/qwenimage/README_TRAIN_LORA_zh-CN.md +++ b/scripts/qwenimage/README_TRAIN_LORA_zh-CN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -139,7 +140,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -209,8 +210,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Qwen-Image` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，代码会自动分桶 | 1328 | | `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 | @@ -270,8 +271,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -321,8 +322,8 @@ python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/ 执行命令为： ```sh export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -361,8 +362,8 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag ```sh export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/qwenimage/README_TRAIN_zh-CN.md b/scripts/qwenimage/README_TRAIN_zh-CN.md index c78bccf0..40620b34 100644 --- a/scripts/qwenimage/README_TRAIN_zh-CN.md +++ b/scripts/qwenimage/README_TRAIN_zh-CN.md @@ -41,9 +41,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -129,8 +130,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -139,7 +140,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -208,8 +209,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Qwen-Image` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，代码会自动分桶 | 1328 | | `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 | @@ -271,8 +272,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -321,8 +322,8 @@ python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/ 执行命令为： ```sh export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -360,8 +361,8 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag ```sh export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/qwenimage_fun/README_TRAIN.md b/scripts/qwenimage_fun/README_TRAIN.md index e4dddab1..2b394ac2 100644 --- a/scripts/qwenimage_fun/README_TRAIN.md +++ b/scripts/qwenimage_fun/README_TRAIN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -142,45 +143,25 @@ It is recommended to use tools like [DWPose](https://github.com/IDEA-Research/DW ### 2.4 Relative vs Absolute Path Usage -**Option 1: Using Relative Paths (Recommended)** +**Relative Paths**: -When data paths are not fixed or you need to train on different machines, relative paths are recommended. +If your data uses relative paths, set in the training script: -Configure relative paths in `metadata.json`, then specify the dataset root directory via `--train_data_dir` in the training script: - -```json -[ - { - "file_path": "train/image001.jpg", - "control_file_path": "control/image001.jpg", - "text": "A group of young men in suits and sunglasses are walking down a city street.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata.json" ``` -During training, the script will automatically search for files corresponding to relative paths under `--train_data_dir`. +**Absolute Paths**: -**Option 2: Using Absolute Paths** +If your data uses absolute paths, set in the training script: -If the dataset path is fixed, you can directly configure absolute paths in `metadata.json`: - -```json -[ - { - "file_path": "/mnt/data/images/image001.jpg", - "control_file_path": "/mnt/data/controls/image001.jpg", - "text": "A group of young men in suits and sunglasses.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="" +export DATASET_META_NAME="/mnt/data/metadata.json" ``` -When using absolute paths, the `--train_data_dir` parameter serves only as a default path, and the absolute paths in the JSON will take priority. +> 💡 **Recommendation**: If the dataset is small and stored locally, relative paths are recommended; if the dataset is stored on external storage (such as NAS, OSS) or shared across multiple machines, absolute paths are recommended. --- diff --git a/scripts/qwenimage_fun/README_TRAIN_zh-CN.md b/scripts/qwenimage_fun/README_TRAIN_zh-CN.md index c26d7854..3d5f764c 100644 --- a/scripts/qwenimage_fun/README_TRAIN_zh-CN.md +++ b/scripts/qwenimage_fun/README_TRAIN_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3:使用docker** @@ -142,45 +143,25 @@ Control 模式的 metadata.json 与普通 Qwen-Image 的 json 略有不同，需 ### 2.4 相对路径与绝对路径使用方案 -**方案 1:使用相对路径(推荐)** +**相对路径**： -当数据路径不固定,或需要在不同机器上训练时,推荐使用相对路径。 +如果数据的路径为相对路径，则在训练脚本中设置： -在 `metadata.json` 中配置相对路径,然后在训练脚本中通过 `--train_data_dir` 指定数据集根目录: - -```json -[ - { - "file_path": "train/image001.jpg", - "control_file_path": "control/image001.jpg", - "text": "A group of young men in suits and sunglasses are walking down a city street.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata.json" ``` -训练时会自动在 `--train_data_dir` 下寻找相对路径对应的文件。 +**绝对路径**： -**方案 2:使用绝对路径** +如果数据的路径为绝对路径，则在训练脚本中设置： -如果数据集路径固定,可以直接在 `metadata.json` 中配置绝对路径: - -```json -[ - { - "file_path": "/mnt/data/images/image001.jpg", - "control_file_path": "/mnt/data/controls/image001.jpg", - "text": "A group of young men in suits and sunglasses.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="" +export DATASET_META_NAME="/mnt/data/metadata.json" ``` -使用绝对路径时,`--train_data_dir` 参数仅作为默认路径,实际会优先使用 json 中的绝对路径。 +> 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 --- diff --git a/scripts/qwenimage_instantx/README_TRAIN.md b/scripts/qwenimage_instantx/README_TRAIN.md index 18816b8e..2f72db55 100644 --- a/scripts/qwenimage_instantx/README_TRAIN.md +++ b/scripts/qwenimage_instantx/README_TRAIN.md @@ -45,9 +45,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -144,45 +145,25 @@ It is recommended to use tools like [DWPose](https://github.com/IDEA-Research/DW ### 2.4 Relative vs Absolute Path Usage -**Option 1: Using Relative Paths (Recommended)** +**Relative Paths**: -When data paths are not fixed or need to be trained on different machines, relative paths are recommended. +If your data uses relative paths, set in the training script: -Configure relative paths in `metadata.json`, then specify the dataset root directory via `--train_data_dir` in the training script: - -```json -[ - { - "file_path": "train/image001.jpg", - "control_file_path": "control/image001.jpg", - "text": "A group of young men in suits and sunglasses are walking down a city street.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata.json" ``` -The training process will automatically search for files corresponding to the relative paths under `--train_data_dir`. +**Absolute Paths**: -**Option 2: Using Absolute Paths** +If your data uses absolute paths, set in the training script: -If the dataset path is fixed, you can directly configure absolute paths in `metadata.json`: - -```json -[ - { - "file_path": "/mnt/data/images/image001.jpg", - "control_file_path": "/mnt/data/controls/image001.jpg", - "text": "A group of young men in suits and sunglasses.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="" +export DATASET_META_NAME="/mnt/data/metadata.json" ``` -When using absolute paths, the `--train_data_dir` parameter serves only as a default path, and the absolute paths in the JSON will take priority. +> 💡 **Recommendation**: If the dataset is small and stored locally, relative paths are recommended; if the dataset is stored on external storage (such as NAS, OSS) or shared across multiple machines, absolute paths are recommended. --- diff --git a/scripts/qwenimage_instantx/README_TRAIN_zh-CN.md b/scripts/qwenimage_instantx/README_TRAIN_zh-CN.md index ec6617b6..1a2a0eb4 100644 --- a/scripts/qwenimage_instantx/README_TRAIN_zh-CN.md +++ b/scripts/qwenimage_instantx/README_TRAIN_zh-CN.md @@ -45,9 +45,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3:使用docker** @@ -144,45 +145,25 @@ Control 模式的 metadata.json 与普通 Qwen-Image 的 json 略有不同，需 ### 2.4 相对路径与绝对路径使用方案 -**方案 1:使用相对路径(推荐)** +**相对路径**： -当数据路径不固定,或需要在不同机器上训练时,推荐使用相对路径。 +如果数据的路径为相对路径，则在训练脚本中设置： -在 `metadata.json` 中配置相对路径,然后在训练脚本中通过 `--train_data_dir` 指定数据集根目录: - -```json -[ - { - "file_path": "train/image001.jpg", - "control_file_path": "control/image001.jpg", - "text": "A group of young men in suits and sunglasses are walking down a city street.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata.json" ``` -训练时会自动在 `--train_data_dir` 下寻找相对路径对应的文件。 +**绝对路径**： -**方案 2:使用绝对路径** +如果数据的路径为绝对路径，则在训练脚本中设置： -如果数据集路径固定,可以直接在 `metadata.json` 中配置绝对路径: - -```json -[ - { - "file_path": "/mnt/data/images/image001.jpg", - "control_file_path": "/mnt/data/controls/image001.jpg", - "text": "A group of young men in suits and sunglasses.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="" +export DATASET_META_NAME="/mnt/data/metadata.json" ``` -使用绝对路径时,`--train_data_dir` 参数仅作为默认路径,实际会优先使用 json 中的绝对路径。 +> 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 --- diff --git a/scripts/turbodiffusion/README_TRAIN_DISTILL.md b/scripts/turbodiffusion/README_TRAIN_DISTILL.md index cc2261e1..5b2644ae 100755 --- a/scripts/turbodiffusion/README_TRAIN_DISTILL.md +++ b/scripts/turbodiffusion/README_TRAIN_DISTILL.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -146,7 +147,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -199,7 +200,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -227,8 +228,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con |-----|------|-------| | `--config_path` | Config file path | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size per GPU | 1 | | `--image_sample_size` | Maximum image training resolution | 640 | | `--video_sample_size` | Maximum video training resolution | 640 | @@ -240,7 +241,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--num_train_epochs` | Number of training epochs | 100 | | `--checkpointing_steps` | Save checkpoint every N steps | 50 | | `--learning_rate` | Initial learning rate (generator) | 2e-06 | -| `--learning_rate_critic` | Initial learning rate (critic) | 2e-07 | +| `--learning_rate_critic` | Initial learning rate (critic) | 2e-06 | | `--lr_scheduler` | Learning rate scheduler | `constant_with_warmup` | | `--lr_warmup_steps` | Learning rate warmup steps | 100 | | `--seed` | Random seed | 42 | @@ -376,7 +377,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -434,7 +435,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -484,7 +485,7 @@ accelerate launch --mixed_precision="bf16" scripts/turbodiffusion/train_distill. --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -544,7 +545,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ diff --git a/scripts/turbodiffusion/README_TRAIN_DISTILL_zh-CN.md b/scripts/turbodiffusion/README_TRAIN_DISTILL_zh-CN.md index ed7e1ecd..b0c32b21 100755 --- a/scripts/turbodiffusion/README_TRAIN_DISTILL_zh-CN.md +++ b/scripts/turbodiffusion/README_TRAIN_DISTILL_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方法 3: 使用 Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果您的数据使用相对路径，请在训练脚本中配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -146,7 +147,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，请使用相对路径。如果数据集存储在外部存储（如 NAS、OSS）或在多台机器间共享，请使用绝对路径。 @@ -199,7 +200,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -227,8 +228,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con |-----|------|-------| | `--config_path` | 配置文件路径 | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元数据文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元数据文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每个 GPU 的批次大小 | 1 | | `--image_sample_size` | 最大图像训练分辨率 | 640 | | `--video_sample_size` | 最大视频训练分辨率 | 640 | @@ -240,7 +241,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--num_train_epochs` | 训练轮数 | 100 | | `--checkpointing_steps` | 每 N 步保存检查点 | 50 | | `--learning_rate` | 初始学习率（生成器） | 2e-06 | -| `--learning_rate_critic` | 初始学习率（判别器） | 2e-07 | +| `--learning_rate_critic` | 初始学习率（判别器） | 2e-06 | | `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` | | `--lr_warmup_steps` | 学习率预热步数 | 100 | | `--seed` | 随机种子 | 42 | @@ -376,7 +377,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -434,7 +435,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -484,7 +485,7 @@ accelerate launch --mixed_precision="bf16" scripts/turbodiffusion/train_distill. --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -544,7 +545,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ diff --git a/scripts/turbodiffusion/train_distill.py b/scripts/turbodiffusion/train_distill.py index ec688995..e8d71aab 100644 --- a/scripts/turbodiffusion/train_distill.py +++ b/scripts/turbodiffusion/train_distill.py @@ -730,6 +730,11 @@ def parse_args(): default=[1000, 750, 500, 250], help="The denoising step list.", ) + parser.add_argument( + "--randomize_step_indices", + action="store_true", + help="whether to use randomize timesteps indices in training.", + ) args = parser.parse_args() env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) @@ -857,6 +862,7 @@ def main(): weight_dtype = torch.bfloat16 args.mixed_precision = accelerator.mixed_precision + args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list] # Load scheduler, tokenizer and models. noise_scheduler = FlowMatchEulerDiscreteScheduler( **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs'])) @@ -1107,8 +1113,8 @@ def load_model_hook(models, input_dir): fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_transformer3d.parameters())) fake_trainable_params_optim = [ - {'params': [], 'lr': args.learning_rate}, - {'params': [], 'lr': args.learning_rate / 2}, + {'params': [], 'lr': args.learning_rate_critic}, + {'params': [], 'lr': args.learning_rate_critic / 2}, ] in_already = [] for name, param in fake_score_transformer3d.named_parameters(): @@ -1121,7 +1127,7 @@ def load_model_hook(models, input_dir): high_lr_flag = True fake_trainable_params_optim[0]['params'].append(param) if accelerator.is_main_process: - print(f"Set {name} to lr : {args.learning_rate}") + print(f"Set {name} to lr : {args.learning_rate_critic}") break if high_lr_flag: continue @@ -1130,7 +1136,7 @@ def load_model_hook(models, input_dir): in_already.append(name) fake_trainable_params_optim[1]['params'].append(param) if accelerator.is_main_process: - print(f"Set {name} to lr : {args.learning_rate / 2}") + print(f"Set {name} to lr : {args.learning_rate_critic / 2}") break if args.use_came: @@ -1501,7 +1507,7 @@ def collate_fn(examples): ) fake_score_lr_scheduler = get_scheduler( args.lr_scheduler, - optimizer=optimizer, + optimizer=critic_optimizer, num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, num_training_steps=args.max_train_steps * accelerator.num_processes, ) @@ -1625,7 +1631,55 @@ def unwrap_model(model): vae_stream_2 = None idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling) - denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)] + + def randomize_denoising_step_indices( + denoising_step_indices_list, + train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=0.3, + ): + indices = list(denoising_step_indices_list) + n = len(indices) + + if n <= 2: + low = indices[1] + high = indices[0] - 1 + random_tail = torch.randint(low, high + 1, (1,)).item() + + result = torch.tensor([indices[0], random_tail]) + else: + result = [0] * n + result[0] = indices[0] + result[-1] = indices[-1] + + for i in range(1, n - 1): + gap_upper = indices[i - 1] - indices[i] + gap_lower = indices[i] - indices[i + 1] + + max_jitter = int(min(gap_upper, gap_lower) * jitter_ratio) + + if max_jitter > 0: + jitter = torch.randint( + -max_jitter, max_jitter + 1, (1,) + ).item() + else: + jitter = 0 + + result[i] = indices[i] + jitter + + for i in range(1, n): + if result[i] >= result[i - 1]: + result[i] = result[i - 1] - 1 + + result = [max(1, min(train_sampling_steps, x)) for x in result] + result = torch.tensor(result) + + if dist.is_initialized(): + result = result.to(accelerator.device) + dist.broadcast(result, src=0) + result = result.cpu() + return result for epoch in range(first_epoch, args.num_train_epochs): train_dmd_loss = 0.0 @@ -1946,6 +2000,20 @@ def convert_flow_pred_to_x0( x0_pred = xt - sigma_t * flow_pred return x0_pred.to(original_dtype) + # Create discrete denoising steps (per-step, with optional randomization) + if getattr(args, 'randomize_step_indices', False): + random_indices = randomize_denoising_step_indices( + args.denoising_step_indices_list, + args.train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=getattr(args, 'index_jitter_ratio', 0.30), + ) + else: + random_indices = torch.tensor(args.denoising_step_indices_list) + + denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - random_indices] + # --- Main Training Logic --- bsz, channel, num_frames, height, width = target_shape if step % args.gen_update_interval == 0: diff --git a/scripts/turbodiffusion/train_distill.sh b/scripts/turbodiffusion/train_distill.sh index 3941d3be..3a6289f9 100644 --- a/scripts/turbodiffusion/train_distill.sh +++ b/scripts/turbodiffusion/train_distill.sh @@ -24,7 +24,7 @@ accelerate launch --mixed_precision="bf16" scripts/turbodiffusion/train_distill. --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ diff --git a/scripts/wan2.1/README_TRAIN.md b/scripts/wan2.1/README_TRAIN.md index 3450a07b..eb87cd4d 100755 --- a/scripts/wan2.1/README_TRAIN.md +++ b/scripts/wan2.1/README_TRAIN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -146,7 +147,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -228,8 +229,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----------|-------------|---------------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size | 1 | | `--image_sample_size` | Maximum training resolution for images | 640 | | `--video_sample_size` | Maximum training resolution for videos | 640 | diff --git a/scripts/wan2.1/README_TRAIN_DISTILL.md b/scripts/wan2.1/README_TRAIN_DISTILL.md index c8e89f9f..911c2ba2 100755 --- a/scripts/wan2.1/README_TRAIN_DISTILL.md +++ b/scripts/wan2.1/README_TRAIN_DISTILL.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -146,7 +147,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (such as NAS, OSS) or shared across multiple machines, use absolute paths. @@ -202,7 +203,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -229,8 +230,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size per GPU | 1 | | `--image_sample_size` | Maximum image training resolution | 640 | | `--video_sample_size` | Maximum video training resolution | 640 | @@ -242,7 +243,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--num_train_epochs` | Number of training epochs | 100 | | `--checkpointing_steps` | Save checkpoint every N steps | 50 | | `--learning_rate` | Initial learning rate (generator) | 2e-06 | -| `--learning_rate_critic` | Initial learning rate (critic) | 2e-07 | +| `--learning_rate_critic` | Initial learning rate (critic) | 2e-06 | | `--lr_scheduler` | Learning rate scheduler | `constant_with_warmup` | | `--lr_warmup_steps` | Learning rate warmup steps | 100 | | `--seed` | Random seed | 42 | @@ -360,7 +361,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -418,7 +419,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -468,7 +469,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1/train_distill.py \ --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -528,7 +529,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ diff --git a/scripts/wan2.1/README_TRAIN_DISTILL_LORA.md b/scripts/wan2.1/README_TRAIN_DISTILL_LORA.md index ac608e5a..d98fed6c 100755 --- a/scripts/wan2.1/README_TRAIN_DISTILL_LORA.md +++ b/scripts/wan2.1/README_TRAIN_DISTILL_LORA.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -146,7 +147,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -179,7 +180,7 @@ The difference between DeepSpeed-Zero-2 and FSDP in this repository is whether m ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -202,7 +203,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.1_distill_lora" \ --gradient_checkpointing \ @@ -248,8 +249,8 @@ In addition to distillation training, LoRA training adds the following specific | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Number of samples per batch | 1 | | `--image_sample_size` | Maximum training resolution for images | 640 | | `--video_sample_size` | Maximum training resolution for videos | 640 | @@ -261,7 +262,7 @@ In addition to distillation training, LoRA training adds the following specific | `--num_train_epochs` | Number of training epochs | 100 | | `--checkpointing_steps` | Save checkpoint every N steps | 50 | | `--learning_rate` | Initial learning rate (generator) | 1e-05 | -| `--learning_rate_critic` | Initial learning rate (discriminator) | 1e-06 | +| `--learning_rate_critic` | Initial learning rate (discriminator) | 1e-05 | | `--seed` | Random seed | 42 | | `--output_dir` | Output directory | `output_dir_wan2.1_distill_lora` | | `--gradient_checkpointing` | Activation recomputation | - | @@ -353,7 +354,7 @@ If VRAM is insufficient when using multiple GPUs with DeepSpeed-Zero-2, you can ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -376,7 +377,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.1_distill_lora" \ --gradient_checkpointing \ @@ -414,7 +415,7 @@ The training shell command is as follows: ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -437,7 +438,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.1_distill_lora" \ --gradient_checkpointing \ @@ -461,7 +462,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -484,7 +485,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1/train_distill_lora.py --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.1_distill_lora" \ --gradient_checkpointing \ @@ -517,7 +518,7 @@ Assuming 2 machines, each with 8 GPUs: ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" export MASTER_ADDR="192.168.1.100" # Master machine IP export MASTER_PORT=10086 export WORLD_SIZE=2 # Total number of machines @@ -545,7 +546,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.1_distill_lora" \ --gradient_checkpointing \ @@ -570,7 +571,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" export MASTER_ADDR="192.168.1.100" # Same as Master export MASTER_PORT=10086 export WORLD_SIZE=2 diff --git a/scripts/wan2.1/README_TRAIN_DISTILL_LORA_zh-CN.md b/scripts/wan2.1/README_TRAIN_DISTILL_LORA_zh-CN.md index 4366b297..fa8815f7 100755 --- a/scripts/wan2.1/README_TRAIN_DISTILL_LORA_zh-CN.md +++ b/scripts/wan2.1/README_TRAIN_DISTILL_LORA_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3:使用 docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径,则在训练脚本中设置: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**: @@ -146,7 +147,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**:如果数据集较小且存储在本地,推荐使用相对路径;如果数据集存储在外部存储(如 NAS、OSS)或多个机器共享存储,推荐使用绝对路径。 @@ -179,7 +180,7 @@ modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_ ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -202,7 +203,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.1_distill_lora" \ --gradient_checkpointing \ @@ -248,8 +249,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | @@ -261,7 +262,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--num_train_epochs` | 训练 epoch 数 | 100 | | `--checkpointing_steps` | 每 N 步保存 checkpoint | 50 | | `--learning_rate` | 初始学习率(生成器) | 1e-05 | -| `--learning_rate_critic` | 初始学习率(判别器) | 1e-06 | +| `--learning_rate_critic` | 初始学习率(判别器) | 1e-05 | | `--seed` | 随机种子 | 42 | | `--output_dir` | 输出目录 | `output_dir_wan2.1_distill_lora` | | `--gradient_checkpointing` | 激活重计算 | - | @@ -353,7 +354,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -376,7 +377,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.1_distill_lora" \ --gradient_checkpointing \ @@ -414,7 +415,7 @@ python scripts/zero_to_bf16.py output_dir/checkpoint-{your-num-steps} output_dir ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -437,7 +438,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.1_distill_lora" \ --gradient_checkpointing \ @@ -461,7 +462,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -484,7 +485,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1/train_distill_lora.py --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.1_distill_lora" \ --gradient_checkpointing \ @@ -517,7 +518,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1/train_distill_lora.py ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" export MASTER_ADDR="192.168.1.100" # Master 机器 IP export MASTER_PORT=10086 export WORLD_SIZE=2 # 机器总数 @@ -545,7 +546,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.1_distill_lora" \ --gradient_checkpointing \ @@ -570,7 +571,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" export MASTER_ADDR="192.168.1.100" # 与 Master 相同 export MASTER_PORT=10086 export WORLD_SIZE=2 diff --git a/scripts/wan2.1/README_TRAIN_DISTILL_zh-CN.md b/scripts/wan2.1/README_TRAIN_DISTILL_zh-CN.md index b167eada..664ccd96 100755 --- a/scripts/wan2.1/README_TRAIN_DISTILL_zh-CN.md +++ b/scripts/wan2.1/README_TRAIN_DISTILL_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -146,7 +147,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -202,7 +203,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -229,8 +230,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | @@ -242,7 +243,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--num_train_epochs` | 训练 epoch 数 | 100 | | `--checkpointing_steps` | 每 N 步保存 checkpoint | 50 | | `--learning_rate` | 初始学习率（生成器） | 2e-06 | -| `--learning_rate_critic` | 初始学习率（判别器） | 2e-07 | +| `--learning_rate_critic` | 初始学习率（判别器） | 2e-06 | | `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` | | `--lr_warmup_steps` | 学习率预热步数 | 100 | | `--seed` | 随机种子 | 42 | @@ -360,7 +361,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -418,7 +419,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -468,7 +469,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1/train_distill.py \ --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -528,7 +529,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ diff --git a/scripts/wan2.1/README_TRAIN_LORA.md b/scripts/wan2.1/README_TRAIN_LORA.md index 148c94ed..d80d0c94 100755 --- a/scripts/wan2.1/README_TRAIN_LORA.md +++ b/scripts/wan2.1/README_TRAIN_LORA.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, set in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -146,7 +147,7 @@ If your data uses absolute paths, set in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Suggestion**: If the dataset is small and stored locally, relative paths are recommended. If the dataset is stored on external storage (such as NAS, OSS) or shared across multiple machines, absolute paths are recommended. @@ -230,8 +231,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con |-----------|-------------|---------------| | `--config_path` | Configuration file path | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | Pre-trained model path | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Number of samples per batch | 1 | | `--image_sample_size` | Maximum training resolution for images | 640 | | `--video_sample_size` | Maximum training resolution for videos | 640 | diff --git a/scripts/wan2.1/README_TRAIN_LORA_zh-CN.md b/scripts/wan2.1/README_TRAIN_LORA_zh-CN.md index 5c987246..8458f384 100644 --- a/scripts/wan2.1/README_TRAIN_LORA_zh-CN.md +++ b/scripts/wan2.1/README_TRAIN_LORA_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -146,7 +147,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -230,8 +231,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con |-----|------|-------| | `--config_path` | 配置文件路径 | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | diff --git a/scripts/wan2.1/README_TRAIN_zh-CN.md b/scripts/wan2.1/README_TRAIN_zh-CN.md index ebff90ee..fc18f981 100644 --- a/scripts/wan2.1/README_TRAIN_zh-CN.md +++ b/scripts/wan2.1/README_TRAIN_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -146,7 +147,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -228,8 +229,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | diff --git a/scripts/wan2.1/train_distill.py b/scripts/wan2.1/train_distill.py index 17ee0689..ebe57c38 100644 --- a/scripts/wan2.1/train_distill.py +++ b/scripts/wan2.1/train_distill.py @@ -729,6 +729,11 @@ def parse_args(): default=[1000, 750, 500, 250], help="The denoising step list.", ) + parser.add_argument( + "--randomize_step_indices", + action="store_true", + help="whether to use randomize timesteps indices in training.", + ) args = parser.parse_args() env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) @@ -856,6 +861,7 @@ def main(): weight_dtype = torch.bfloat16 args.mixed_precision = accelerator.mixed_precision + args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list] # Load scheduler, tokenizer and models. noise_scheduler = FlowMatchEulerDiscreteScheduler( **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs'])) @@ -1106,8 +1112,8 @@ def load_model_hook(models, input_dir): fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_transformer3d.parameters())) fake_trainable_params_optim = [ - {'params': [], 'lr': args.learning_rate}, - {'params': [], 'lr': args.learning_rate / 2}, + {'params': [], 'lr': args.learning_rate_critic}, + {'params': [], 'lr': args.learning_rate_critic / 2}, ] in_already = [] for name, param in fake_score_transformer3d.named_parameters(): @@ -1120,7 +1126,7 @@ def load_model_hook(models, input_dir): high_lr_flag = True fake_trainable_params_optim[0]['params'].append(param) if accelerator.is_main_process: - print(f"Set {name} to lr : {args.learning_rate}") + print(f"Set {name} to lr : {args.learning_rate_critic}") break if high_lr_flag: continue @@ -1129,7 +1135,7 @@ def load_model_hook(models, input_dir): in_already.append(name) fake_trainable_params_optim[1]['params'].append(param) if accelerator.is_main_process: - print(f"Set {name} to lr : {args.learning_rate / 2}") + print(f"Set {name} to lr : {args.learning_rate_critic / 2}") break if args.use_came: @@ -1500,7 +1506,7 @@ def collate_fn(examples): ) fake_score_lr_scheduler = get_scheduler( args.lr_scheduler, - optimizer=optimizer, + optimizer=critic_optimizer, num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, num_training_steps=args.max_train_steps * accelerator.num_processes, ) @@ -1623,7 +1629,55 @@ def unwrap_model(model): vae_stream_2 = None idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling) - denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)] + + def randomize_denoising_step_indices( + denoising_step_indices_list, + train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=0.3, + ): + indices = list(denoising_step_indices_list) + n = len(indices) + + if n <= 2: + low = indices[1] + high = indices[0] - 1 + random_tail = torch.randint(low, high + 1, (1,)).item() + + result = torch.tensor([indices[0], random_tail]) + else: + result = [0] * n + result[0] = indices[0] + result[-1] = indices[-1] + + for i in range(1, n - 1): + gap_upper = indices[i - 1] - indices[i] + gap_lower = indices[i] - indices[i + 1] + + max_jitter = int(min(gap_upper, gap_lower) * jitter_ratio) + + if max_jitter > 0: + jitter = torch.randint( + -max_jitter, max_jitter + 1, (1,) + ).item() + else: + jitter = 0 + + result[i] = indices[i] + jitter + + for i in range(1, n): + if result[i] >= result[i - 1]: + result[i] = result[i - 1] - 1 + + result = [max(1, min(train_sampling_steps, x)) for x in result] + result = torch.tensor(result) + + if dist.is_initialized(): + result = result.to(accelerator.device) + dist.broadcast(result, src=0) + result = result.cpu() + return result for epoch in range(first_epoch, args.num_train_epochs): train_dmd_loss = 0.0 @@ -1944,6 +1998,20 @@ def convert_flow_pred_to_x0( x0_pred = xt - sigma_t * flow_pred return x0_pred.to(original_dtype) + # Create discrete denoising steps (per-step, with optional randomization) + if getattr(args, 'randomize_step_indices', False): + random_indices = randomize_denoising_step_indices( + args.denoising_step_indices_list, + args.train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=getattr(args, 'index_jitter_ratio', 0.30), + ) + else: + random_indices = torch.tensor(args.denoising_step_indices_list) + + denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - random_indices] + # --- Main Training Logic --- bsz, channel, num_frames, height, width = target_shape if step % args.gen_update_interval == 0: diff --git a/scripts/wan2.1/train_distill.sh b/scripts/wan2.1/train_distill.sh index dbe45245..ffd35186 100644 --- a/scripts/wan2.1/train_distill.sh +++ b/scripts/wan2.1/train_distill.sh @@ -23,7 +23,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1/train_distill.py \ --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ diff --git a/scripts/wan2.1/train_distill_lora.py b/scripts/wan2.1/train_distill_lora.py index e4ae8469..e48ab067 100644 --- a/scripts/wan2.1/train_distill_lora.py +++ b/scripts/wan2.1/train_distill_lora.py @@ -744,6 +744,11 @@ def parse_args(): default=[1000, 750, 500, 250], help="The denoising step list.", ) + parser.add_argument( + "--randomize_step_indices", + action="store_true", + help="whether to use randomize timesteps indices in training.", + ) parser.add_argument( "--lora_skip_name", type=str, @@ -883,6 +888,7 @@ def main(): weight_dtype = torch.bfloat16 args.mixed_precision = accelerator.mixed_precision + args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list] # Load scheduler, tokenizer and models. noise_scheduler = FlowMatchEulerDiscreteScheduler( **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs'])) @@ -1151,7 +1157,7 @@ def load_model_hook(models, input_dir): logging.info("Add fake_score_network parameters") fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_network.parameters())) - fake_trainable_params_optim = fake_score_network.prepare_optimizer_params(args.learning_rate / 2, args.learning_rate, args.learning_rate) + fake_trainable_params_optim = fake_score_network.prepare_optimizer_params(args.learning_rate_critic / 2, args.learning_rate_critic, args.learning_rate_critic) if args.use_came: optimizer = optimizer_cls( @@ -1521,7 +1527,7 @@ def collate_fn(examples): ) fake_score_lr_scheduler = get_scheduler( args.lr_scheduler, - optimizer=optimizer, + optimizer=critic_optimizer, num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, num_training_steps=args.max_train_steps * accelerator.num_processes, ) @@ -1664,7 +1670,55 @@ def save_model(ckpt_file, unwrapped_nw): vae_stream_2 = None idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling) - denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)] + + def randomize_denoising_step_indices( + denoising_step_indices_list, + train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=0.3, + ): + indices = list(denoising_step_indices_list) + n = len(indices) + + if n <= 2: + low = indices[1] + high = indices[0] - 1 + random_tail = torch.randint(low, high + 1, (1,)).item() + + result = torch.tensor([indices[0], random_tail]) + else: + result = [0] * n + result[0] = indices[0] + result[-1] = indices[-1] + + for i in range(1, n - 1): + gap_upper = indices[i - 1] - indices[i] + gap_lower = indices[i] - indices[i + 1] + + max_jitter = int(min(gap_upper, gap_lower) * jitter_ratio) + + if max_jitter > 0: + jitter = torch.randint( + -max_jitter, max_jitter + 1, (1,) + ).item() + else: + jitter = 0 + + result[i] = indices[i] + jitter + + for i in range(1, n): + if result[i] >= result[i - 1]: + result[i] = result[i - 1] - 1 + + result = [max(1, min(train_sampling_steps, x)) for x in result] + result = torch.tensor(result) + + if dist.is_initialized(): + result = result.to(accelerator.device) + dist.broadcast(result, src=0) + result = result.cpu() + return result for epoch in range(first_epoch, args.num_train_epochs): train_dmd_loss = 0.0 @@ -1985,6 +2039,20 @@ def convert_flow_pred_to_x0( x0_pred = xt - sigma_t * flow_pred return x0_pred.to(original_dtype) + # Create discrete denoising steps (per-step, with optional randomization) + if getattr(args, 'randomize_step_indices', False): + random_indices = randomize_denoising_step_indices( + args.denoising_step_indices_list, + args.train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=getattr(args, 'index_jitter_ratio', 0.30), + ) + else: + random_indices = torch.tensor(args.denoising_step_indices_list) + + denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - random_indices] + # --- Main Training Logic --- bsz, channel, num_frames, height, width = target_shape if step % args.gen_update_interval == 0: diff --git a/scripts/wan2.1/train_distill_lora.sh b/scripts/wan2.1/train_distill_lora.sh index 78018350..8322a04d 100644 --- a/scripts/wan2.1/train_distill_lora.sh +++ b/scripts/wan2.1/train_distill_lora.sh @@ -23,7 +23,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1/train_distill_lora.py --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.1_distill_lora" \ --gradient_checkpointing \ diff --git a/scripts/wan2.1/train_reward_lora.py b/scripts/wan2.1/train_reward_lora.py index af22e83b..c8018967 100755 --- a/scripts/wan2.1/train_reward_lora.py +++ b/scripts/wan2.1/train_reward_lora.py @@ -38,7 +38,10 @@ from accelerate.logging import get_logger from accelerate.state import AcceleratorState from accelerate.utils import ProjectConfiguration, set_seed -from decord import VideoReader +try: + from decord import VideoReader +except ImportError: + from videox_fun.data.utils import AVVideoReader as VideoReader from diffusers import DDIMScheduler, FlowMatchEulerDiscreteScheduler from diffusers.optimization import get_scheduler from diffusers.utils import check_min_version, is_wandb_available diff --git a/scripts/wan2.1_fun/README_TRAIN.md b/scripts/wan2.1_fun/README_TRAIN.md index 842ec1e0..8dde7ed1 100755 --- a/scripts/wan2.1_fun/README_TRAIN.md +++ b/scripts/wan2.1_fun/README_TRAIN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -146,7 +147,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -228,8 +229,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----------|-------------|---------------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.1-Fun-V1.1-1.3B-InP` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size | 1 | | `--image_sample_size` | Maximum training resolution for images | 640 | | `--video_sample_size` | Maximum training resolution for videos | 640 | diff --git a/scripts/wan2.1_fun/README_TRAIN_CONTROL.md b/scripts/wan2.1_fun/README_TRAIN_CONTROL.md index 06e35d6c..4d4c9f28 100755 --- a/scripts/wan2.1_fun/README_TRAIN_CONTROL.md +++ b/scripts/wan2.1_fun/README_TRAIN_CONTROL.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -145,8 +146,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Controls-Demo --local_dir ./datas If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata_control.json" +export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -155,7 +156,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata_control.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -239,8 +240,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con |-----|------|-------| | `--config_path` | Model config file path | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.1-Fun-V1.1-14B-Control` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata_control.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Controls-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size per GPU | 1 | | `--image_sample_size` | Maximum image training resolution | 640 | | `--video_sample_size` | Maximum video training resolution | 640 | diff --git a/scripts/wan2.1_fun/README_TRAIN_CONTROL_LORA.md b/scripts/wan2.1_fun/README_TRAIN_CONTROL_LORA.md index 050d5b4b..41ff30b9 100755 --- a/scripts/wan2.1_fun/README_TRAIN_CONTROL_LORA.md +++ b/scripts/wan2.1_fun/README_TRAIN_CONTROL_LORA.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -145,8 +146,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Controls-Demo --local_dir ./datas If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata_control.json" +export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -155,7 +156,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata_control.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -240,8 +241,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con |-----|------|-------| | `--config_path` | Model config file path | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.1-Fun-V1.1-14B-Control` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata_control.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Controls-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size per GPU | 1 | | `--image_sample_size` | Maximum image training resolution | 640 | | `--video_sample_size` | Maximum video training resolution | 640 | diff --git a/scripts/wan2.1_fun/README_TRAIN_CONTROL_LORA_zh-CN.md b/scripts/wan2.1_fun/README_TRAIN_CONTROL_LORA_zh-CN.md index bff96ce7..bd2558b8 100755 --- a/scripts/wan2.1_fun/README_TRAIN_CONTROL_LORA_zh-CN.md +++ b/scripts/wan2.1_fun/README_TRAIN_CONTROL_LORA_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -145,8 +146,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Controls-Demo --local_dir ./datas 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata_control.json" +export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -155,7 +156,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata_control.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata_control.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -240,8 +241,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con |-----|------|-------| | `--config_path` | 配置文件路径 | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-Fun-V1.1-14B-Control` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata_control.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Controls-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | diff --git a/scripts/wan2.1_fun/README_TRAIN_CONTROL_zh-CN.md b/scripts/wan2.1_fun/README_TRAIN_CONTROL_zh-CN.md index 4795f50b..99aa1a4e 100755 --- a/scripts/wan2.1_fun/README_TRAIN_CONTROL_zh-CN.md +++ b/scripts/wan2.1_fun/README_TRAIN_CONTROL_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -145,8 +146,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Controls-Demo --local_dir ./datas 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata_control.json" +export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -155,7 +156,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata_control.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata_control.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -239,8 +240,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con |-----|------|-------| | `--config_path` | 模型配置文件路径 | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-Fun-V1.1-14B-Control` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata_control.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Controls-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | diff --git a/scripts/wan2.1_fun/README_TRAIN_LORA.md b/scripts/wan2.1_fun/README_TRAIN_LORA.md index c899cf5f..b99ce17b 100755 --- a/scripts/wan2.1_fun/README_TRAIN_LORA.md +++ b/scripts/wan2.1_fun/README_TRAIN_LORA.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, set in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -146,7 +147,7 @@ If your data uses absolute paths, set in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Suggestion**: If the dataset is small and stored locally, relative paths are recommended. If the dataset is stored on external storage (such as NAS, OSS) or shared across multiple machines, absolute paths are recommended. @@ -230,8 +231,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con |-----------|-------------|---------------| | `--config_path` | Configuration file path | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | Pre-trained model path | `models/Diffusion_Transformer/Wan2.1-Fun-V1.1-1.3B-InP` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Number of samples per batch | 1 | | `--image_sample_size` | Maximum training resolution for images | 640 | | `--video_sample_size` | Maximum training resolution for videos | 640 | diff --git a/scripts/wan2.1_fun/README_TRAIN_LORA_zh-CN.md b/scripts/wan2.1_fun/README_TRAIN_LORA_zh-CN.md index d768410a..42086e65 100644 --- a/scripts/wan2.1_fun/README_TRAIN_LORA_zh-CN.md +++ b/scripts/wan2.1_fun/README_TRAIN_LORA_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -146,7 +147,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -230,8 +231,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con |-----|------|-------| | `--config_path` | 配置文件路径 | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-Fun-V1.1-1.3B-InP` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | diff --git a/scripts/wan2.1_fun/README_TRAIN_zh-CN.md b/scripts/wan2.1_fun/README_TRAIN_zh-CN.md index 83f528ab..3f20ad2e 100644 --- a/scripts/wan2.1_fun/README_TRAIN_zh-CN.md +++ b/scripts/wan2.1_fun/README_TRAIN_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -146,7 +147,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -228,8 +229,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-Fun-V1.1-1.3B-InP` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | diff --git a/scripts/wan2.1_fun/train_reward_lora.py b/scripts/wan2.1_fun/train_reward_lora.py index 4bae97c3..9e18bcfd 100755 --- a/scripts/wan2.1_fun/train_reward_lora.py +++ b/scripts/wan2.1_fun/train_reward_lora.py @@ -39,7 +39,10 @@ from accelerate.logging import get_logger from accelerate.state import AcceleratorState from accelerate.utils import ProjectConfiguration, set_seed -from decord import VideoReader +try: + from decord import VideoReader +except ImportError: + from videox_fun.data.utils import AVVideoReader as VideoReader from diffusers import DDIMScheduler, FlowMatchEulerDiscreteScheduler from diffusers.optimization import get_scheduler from diffusers.utils import check_min_version, is_wandb_available diff --git a/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md b/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md index c8f6b291..d434b765 100755 --- a/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md +++ b/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -135,8 +136,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -145,7 +146,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (such as NAS, OSS) or shared across multiple machines, use absolute paths. @@ -203,7 +204,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -233,8 +234,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size per GPU | 1 | | `--image_sample_size` | Maximum image training resolution | 640 | | `--video_sample_size` | Maximum video training resolution | 640 | @@ -246,7 +247,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--num_train_epochs` | Number of training epochs | 100 | | `--checkpointing_steps` | Save checkpoint every N steps | 50 | | `--learning_rate` | Initial learning rate (generator) | 2e-06 | -| `--learning_rate_critic` | Initial learning rate (critic) | 2e-07 | +| `--learning_rate_critic` | Initial learning rate (critic) | 2e-06 | | `--lr_scheduler` | Learning rate scheduler | `constant_with_warmup` | | `--lr_warmup_steps` | Learning rate warmup steps | 100 | | `--seed` | Random seed | 42 | @@ -366,7 +367,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -428,7 +429,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -482,7 +483,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -546,7 +547,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ diff --git a/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL_zh-CN.md b/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL_zh-CN.md index 053be37c..305a831c 100755 --- a/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL_zh-CN.md +++ b/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -135,8 +136,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -145,7 +146,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -204,7 +205,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -234,8 +235,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | @@ -247,7 +248,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--num_train_epochs` | 训练 epoch 数 | 100 | | `--checkpointing_steps` | 每 N 步保存 checkpoint | 50 | | `--learning_rate` | 初始学习率（生成器） | 2e-06 | -| `--learning_rate_critic` | 初始学习率（判别器） | 2e-07 | +| `--learning_rate_critic` | 初始学习率（判别器） | 2e-06 | | `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` | | `--lr_warmup_steps` | 学习率预热步数 | 100 | | `--seed` | 随机种子 | 42 | @@ -367,7 +368,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -429,7 +430,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -483,7 +484,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -547,7 +548,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ diff --git a/scripts/wan2.1_self_forcing/README_TRAIN_ODE.md b/scripts/wan2.1_self_forcing/README_TRAIN_ODE.md index 7828f1fa..800b6b4c 100755 --- a/scripts/wan2.1_self_forcing/README_TRAIN_ODE.md +++ b/scripts/wan2.1_self_forcing/README_TRAIN_ODE.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** diff --git a/scripts/wan2.1_self_forcing/README_TRAIN_ODE_zh-CN.md b/scripts/wan2.1_self_forcing/README_TRAIN_ODE_zh-CN.md index 5e3970c7..b1310c0c 100755 --- a/scripts/wan2.1_self_forcing/README_TRAIN_ODE_zh-CN.md +++ b/scripts/wan2.1_self_forcing/README_TRAIN_ODE_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用 docker** diff --git a/scripts/wan2.1_self_forcing/train_distill.py b/scripts/wan2.1_self_forcing/train_distill.py index d56eca52..9ec8d31a 100644 --- a/scripts/wan2.1_self_forcing/train_distill.py +++ b/scripts/wan2.1_self_forcing/train_distill.py @@ -117,6 +117,54 @@ def initialize_crossattn_cache_for_training(batch_size, text_len, num_layers, nu return crossattn_cache +def slice_last_n_latent_frames(tensor, n): + """Slice last n frames from [B, C, F, H, W] tensor.""" + if tensor.shape[2] <= n: + return tensor + return tensor[:, :, -n:] + + +def reencode_boundary_latent(vae, pred_latents, weight_dtype, score_num_frames=21): + """ + Re-encode the boundary frame to get a clean latent for the score window. + Follows Self-Forcing reference: decode all frames before the score window, take last pixel frame, re-encode. + Input: pred_latents [B, C, F, H, W] (all generated latent frames) + Output: boundary_latent [B, C, 1, H, W] + """ + with torch.no_grad(): + # Decode all frames except the last (score_num_frames - 1) to pixels + tail_len = score_num_frames - 1 + latent_to_decode = pred_latents[:, :, :-tail_len] + # VAE expects [B, C, F, H, W], decode returns [B, C, F, H, W] pixels + pixels = vae.decode(latent_to_decode.to(vae.dtype)).sample # [B, C, F, H, W] + # Take the last frame + frame = pixels[:, :, -1:, :, :] # [B, C, 1, H, W] + # Re-encode the last frame to get clean boundary latent + boundary_latent = vae.encode(frame)[0].sample().to(weight_dtype) # [B, C, 1, H, W] + return boundary_latent + + +def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_first_frame=False): + """ + Slice the last `score_num_frames` latent frames for score computation. + If pred has more than score_num_frames, re-encode boundary frame for clean context. + Returns: (pred_for_score, score_num_frames, need_gradient_mask) + """ + num_frames = pred.shape[2] + if num_frames <= score_num_frames: + return pred, num_frames, False + + # Re-encode boundary for cleaner score input + try: + boundary_latent = reencode_boundary_latent(vae, pred, weight_dtype, score_num_frames=score_num_frames) + pred_for_score = torch.cat([boundary_latent, pred[:, :, -(score_num_frames - 1):]], dim=2) + except Exception: + # Fallback: simple slice without boundary re-encoding + pred_for_score = pred[:, :, -score_num_frames:] + + return pred_for_score, score_num_frames, True + + def filter_kwargs(cls, kwargs): import inspect sig = inspect.signature(cls.__init__) @@ -740,6 +788,11 @@ def parse_args(): default=[1000, 750, 500, 250], help="The denoising step list.", ) + parser.add_argument( + "--randomize_step_indices", + action="store_true", + help="whether to use randomize timesteps indices in training.", + ) parser.add_argument( "--num_frame_per_block", type=int, @@ -756,6 +809,22 @@ def parse_args(): action="store_true", help="Use KV cache block-by-block training (matches original Self-Forcing)" ) + parser.add_argument( + "--score_num_frames", + type=int, + default=21, + help="Number of latent frames for score computation window (default: 21, matching base model). " + "fake_score/real_score always receive this many frames." + ) + parser.add_argument( + "--min_length_prob_bias", + type=float, + default=0.0, + help="Probability bias for sampling the minimum length (score_num_frames). " + "0.0 = uniform sampling (default), 0.5 = 50%% prob for min length, " + "remaining prob distributed equally among longer lengths. " + "Use this to increase 21-frame training ratio." + ) parser.add_argument( "--context_noise", type=int, @@ -900,6 +969,7 @@ def main(): weight_dtype = torch.bfloat16 args.mixed_precision = accelerator.mixed_precision + args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list] # Load scheduler, tokenizer and models. noise_scheduler = FlowMatchEulerDiscreteScheduler( **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs'])) @@ -1173,8 +1243,8 @@ def load_model_hook(models, input_dir): fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_transformer3d.parameters())) fake_trainable_params_optim = [ - {'params': [], 'lr': args.learning_rate}, - {'params': [], 'lr': args.learning_rate / 2}, + {'params': [], 'lr': args.learning_rate_critic}, + {'params': [], 'lr': args.learning_rate_critic / 2}, ] in_already = [] for name, param in fake_score_transformer3d.named_parameters(): @@ -1187,7 +1257,7 @@ def load_model_hook(models, input_dir): high_lr_flag = True fake_trainable_params_optim[0]['params'].append(param) if accelerator.is_main_process: - print(f"Set {name} to lr : {args.learning_rate}") + print(f"Set {name} to lr : {args.learning_rate_critic}") break if high_lr_flag: continue @@ -1196,7 +1266,7 @@ def load_model_hook(models, input_dir): in_already.append(name) fake_trainable_params_optim[1]['params'].append(param) if accelerator.is_main_process: - print(f"Set {name} to lr : {args.learning_rate / 2}") + print(f"Set {name} to lr : {args.learning_rate_critic / 2}") break if args.use_came: @@ -1461,7 +1531,7 @@ def collate_fn(examples): new_examples['encoder_hidden_states'] = prompt_embeds neg_txt = [ - "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" for text in batch['text'] + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" for text in new_examples['text'] ] neg_prompt_ids = tokenizer( neg_txt, @@ -1519,7 +1589,7 @@ def collate_fn(examples): new_examples['encoder_hidden_states'] = prompt_embeds neg_txt = [ - "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" for text in batch['text'] + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" for text in new_examples['text'] ] neg_prompt_ids = tokenizer( neg_txt, @@ -1578,7 +1648,7 @@ def collate_fn(examples): ) fake_score_lr_scheduler = get_scheduler( args.lr_scheduler, - optimizer=optimizer, + optimizer=critic_optimizer, num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, num_training_steps=args.max_train_steps * accelerator.num_processes, ) @@ -1701,7 +1771,55 @@ def unwrap_model(model): vae_stream_2 = None idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling) - denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)] + + def randomize_denoising_step_indices( + denoising_step_indices_list, + train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=0.3, + ): + indices = list(denoising_step_indices_list) + n = len(indices) + + if n <= 2: + low = indices[1] + high = indices[0] - 1 + random_tail = torch.randint(low, high + 1, (1,)).item() + + result = torch.tensor([indices[0], random_tail]) + else: + result = [0] * n + result[0] = indices[0] + result[-1] = indices[-1] + + for i in range(1, n - 1): + gap_upper = indices[i - 1] - indices[i] + gap_lower = indices[i] - indices[i + 1] + + max_jitter = int(min(gap_upper, gap_lower) * jitter_ratio) + + if max_jitter > 0: + jitter = torch.randint( + -max_jitter, max_jitter + 1, (1,) + ).item() + else: + jitter = 0 + + result[i] = indices[i] + jitter + + for i in range(1, n): + if result[i] >= result[i - 1]: + result[i] = result[i - 1] - 1 + + result = [max(1, min(train_sampling_steps, x)) for x in result] + result = torch.tensor(result) + + if dist.is_initialized(): + result = result.to(accelerator.device) + dist.broadcast(result, src=0) + result = result.cpu() + return result for epoch in range(first_epoch, args.num_train_epochs): train_dmd_loss = 0.0 @@ -1908,10 +2026,25 @@ def _batch_encode_vae(pixel_values): local_sample_size = aspect_ratio_sample_size[aspect_ratio_key] local_sample_size = [int(x / 16) * 16 for x in local_sample_size] + # Compute latent frame count + latent_num_frames = int((num_frames - 1) // vae.temporal_compression_ratio + 1) + + # Align latent_num_frames to num_frame_per_block for KV cache training + if args.use_kv_cache_training: + if args.independent_first_frame: + # latent_frames - 1 must be divisible by num_frame_per_block + k = latent_num_frames - 1 + k = (k // args.num_frame_per_block) * args.num_frame_per_block + latent_num_frames = k + 1 + else: + # latent_frames must be divisible by num_frame_per_block + latent_num_frames = (latent_num_frames // args.num_frame_per_block) * args.num_frame_per_block + latent_num_frames = max(latent_num_frames, args.num_frame_per_block) + target_shape = ( len(text), vae.latent_channels, - int((num_frames - 1) // vae.temporal_compression_ratio + 1), + latent_num_frames, int(local_sample_size[0] // vae.spatial_compression_ratio), int(local_sample_size[1] // vae.spatial_compression_ratio), ) @@ -2032,6 +2165,20 @@ def convert_flow_pred_to_x0( x0_pred = xt - sigma_t * flow_pred return x0_pred.to(original_dtype) + # Create discrete denoising steps (per-step, with optional randomization) + if getattr(args, 'randomize_step_indices', False): + random_indices = randomize_denoising_step_indices( + args.denoising_step_indices_list, + args.train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=getattr(args, 'index_jitter_ratio', 0.30), + ) + else: + random_indices = torch.tensor(args.denoising_step_indices_list) + + denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - random_indices] + # --- Main Training Logic --- bsz, channel, num_frames, height, width = target_shape if step % args.gen_update_interval == 0: @@ -2040,17 +2187,47 @@ def convert_flow_pred_to_x0( patch_h, patch_w = accelerator.unwrap_model(generator_transformer3d).config.patch_size[1:] frame_seq_length = (target_shape[3] * target_shape[4]) // (patch_h * patch_w) - # Determine block structure + # Determine block structure with variable-length support if not args.independent_first_frame: assert num_frames % args.num_frame_per_block == 0 - num_blocks = num_frames // args.num_frame_per_block + max_num_blocks = num_frames // args.num_frame_per_block + assert args.score_num_frames % args.num_frame_per_block == 0 + min_num_blocks = args.score_num_frames // args.num_frame_per_block else: assert (num_frames - 1) % args.num_frame_per_block == 0 - num_blocks = (num_frames - 1) // args.num_frame_per_block - - all_num_frames = [args.num_frame_per_block] * num_blocks + max_num_blocks = (num_frames - 1) // args.num_frame_per_block + if args.score_num_frames > 1: + assert (args.score_num_frames - 1) % args.num_frame_per_block == 0 + min_num_blocks = (args.score_num_frames - 1) // args.num_frame_per_block + else: + min_num_blocks = 0 + + # Random sample number of blocks (Self-Forcing variable-length training) + if args.min_length_prob_bias > 0.0 and max_num_blocks > min_num_blocks: + # Weighted sampling: give min_num_blocks a higher probability + num_options = max_num_blocks - min_num_blocks + 1 + bias = min(args.min_length_prob_bias, 0.99) + remaining_prob = (1.0 - bias) / (num_options - 1) + probs = [remaining_prob] * num_options + probs[0] = bias # min_num_blocks gets the bias + probs_tensor = torch.tensor(probs, device=accelerator.device) + block_indices = torch.multinomial(probs_tensor, 1, generator=torch_rng) + num_generated_blocks = (min_num_blocks + block_indices).item() + else: + num_generated_blocks = torch.randint( + min_num_blocks, max_num_blocks + 1, (1,), + generator=torch_rng, device=accelerator.device + ).item() + if dist.is_initialized(): + _sync = torch.tensor([num_generated_blocks], device=accelerator.device) + dist.broadcast(_sync, src=0) + num_generated_blocks = _sync.item() + + all_num_frames = [args.num_frame_per_block] * num_generated_blocks if args.independent_first_frame: all_num_frames = [1] + all_num_frames + + num_generated_frames = sum(all_num_frames) # Initialize KV cache num_layers = generator_transformer3d.config.num_layers @@ -2098,6 +2275,10 @@ def convert_flow_pred_to_x0( num_denoising_steps = len(denoising_step_list) final_step_index = generate_and_sync_list(num_denoising_steps, device=accelerator.device)[0] + # Only blocks in the last score_num_frames get gradient at exit step + # (matches Self-Forcing: start_gradient_frame_index = num_output_frames - 21) + start_gradient_frame_index = num_generated_frames - args.score_num_frames + for block_idx, current_num_frames in enumerate(all_num_frames): # Extract noise for current block start_idx = current_start_frame - num_input_frames @@ -2113,7 +2294,11 @@ def convert_flow_pred_to_x0( dtype=torch.int64 ) - context_manager = torch.no_grad() if not is_final_step else contextlib.nullcontext() + # Only enable gradient for final step AND block within score window + if not is_final_step or current_start_frame < start_gradient_frame_index: + context_manager = torch.no_grad() + else: + context_manager = contextlib.nullcontext() with context_manager: # Convert noisy_input to list format @@ -2194,13 +2379,54 @@ def convert_flow_pred_to_x0( current_start_frame += current_num_frames - # Final output - generator_pred = output_pred - seq_len = frame_seq_length * num_frames # For fake/real score computation + # Final output — slice generated frames (may be < num_frames for variable-length) + generator_pred_full = output_pred[:, :, :num_generated_frames] + + # Gradient mask: first block gets no gradient when generating > min frames + # (matches Self-Forcing reference: model/base.py L182-L190) + min_num_frames_score = args.score_num_frames + need_gradient_mask = (num_generated_frames != min_num_frames_score) + gradient_mask = None + if need_gradient_mask: + gradient_mask = torch.ones_like(generator_pred_full, dtype=torch.bool) + if args.independent_first_frame: + gradient_mask[:, :, :1] = False + else: + gradient_mask[:, :, :args.num_frame_per_block] = False + + # Slice for score computation: last score_num_frames frames + if num_generated_frames > args.score_num_frames: + # Re-encode boundary for cleaner score input + generator_pred_for_score, score_num_frames, _ = slice_for_score( + generator_pred_full, vae, weight_dtype, + score_num_frames=args.score_num_frames, + independent_first_frame=args.independent_first_frame, + ) + else: + generator_pred_for_score = generator_pred_full + score_num_frames = num_generated_frames + + # Compute score_mask for DMD loss (matches Self-Forcing: dmd.py L199-204) + score_mask = None + if gradient_mask is not None: + mask_offset = num_generated_frames - score_num_frames + score_mask = gradient_mask[:, :, mask_offset:mask_offset + score_num_frames] + + # generator_pred = the sliced version for DMD loss + generator_pred = generator_pred_for_score + seq_len = frame_seq_length * score_num_frames # Score always on fixed window else: # === Block mask training (flex attention, no KV cache) === # Block mask training: use flex attention to process entire video at once + # Note: for long videos, use KV cache mode instead + score_mask = None # Block mask mode: no gradient mask needed + if num_frames > args.score_num_frames: + raise ValueError( + f"Block mask mode does not support variable-length training " + f"(video produces {num_frames} latent frames > score_num_frames={args.score_num_frames}). " + f"Use --use_kv_cache_training for long video training." + ) patch_h_bm, patch_w_bm = accelerator.unwrap_model(generator_transformer3d).config.patch_size[1:] frame_seqlen_bm = (height * width) // (patch_h_bm * patch_w_bm) @@ -2379,11 +2605,20 @@ def convert_flow_pred_to_x0( fake_to_real_grad = fake_to_real_grad / normalizer fake_to_real_grad = torch.nan_to_num(fake_to_real_grad) - dmd_loss = 0.5 * F.mse_loss( - generator_pred.double(), - (generator_pred.double() - fake_to_real_grad.double()).detach(), - reduction="mean" - ) + # Apply gradient mask: only compute loss on unmasked elements + # (matches Self-Forcing dmd.py: F.mse_loss(x[mask], target[mask])) + if score_mask is not None: + dmd_loss = 0.5 * F.mse_loss( + generator_pred.double()[score_mask], + (generator_pred.double() - fake_to_real_grad.double()).detach()[score_mask], + reduction="mean" + ) + else: + dmd_loss = 0.5 * F.mse_loss( + generator_pred.double(), + (generator_pred.double() - fake_to_real_grad.double()).detach(), + reduction="mean" + ) avg_dmd_loss = accelerator.gather(dmd_loss.repeat(args.train_batch_size)).mean() train_dmd_loss += avg_dmd_loss.item() / args.gradient_accumulation_steps @@ -2414,15 +2649,41 @@ def convert_flow_pred_to_x0( # Calculate frame_seq_length frame_seq_length = (target_shape[3] * target_shape[4]) // (patch_h * patch_w) - # Determine block structure + # Determine block structure (variable-length, mirrors generator branch) if not args.independent_first_frame: - num_blocks = num_frames // args.num_frame_per_block + max_num_blocks_critic = num_frames // args.num_frame_per_block + min_num_blocks_critic = args.score_num_frames // args.num_frame_per_block else: - num_blocks = (num_frames - 1) // args.num_frame_per_block - - all_num_frames = [args.num_frame_per_block] * num_blocks + max_num_blocks_critic = (num_frames - 1) // args.num_frame_per_block + if args.score_num_frames > 1: + min_num_blocks_critic = (args.score_num_frames - 1) // args.num_frame_per_block + else: + min_num_blocks_critic = 0 + + # Random sample number of blocks (mirrors generator's variable-length training) + if args.min_length_prob_bias > 0.0 and max_num_blocks_critic > min_num_blocks_critic: + num_options = max_num_blocks_critic - min_num_blocks_critic + 1 + bias = min(args.min_length_prob_bias, 0.99) + remaining_prob = (1.0 - bias) / (num_options - 1) + probs = [remaining_prob] * num_options + probs[0] = bias # min_num_blocks_critic gets the bias + probs_tensor = torch.tensor(probs, device=accelerator.device) + block_indices = torch.multinomial(probs_tensor, 1, generator=torch_rng) + num_generated_blocks_critic = (min_num_blocks_critic + block_indices).item() + else: + num_generated_blocks_critic = torch.randint( + min_num_blocks_critic, max_num_blocks_critic + 1, (1,), + generator=torch_rng, device=accelerator.device + ).item() + if dist.is_initialized(): + _sync = torch.tensor([num_generated_blocks_critic], device=accelerator.device) + dist.broadcast(_sync, src=0) + num_generated_blocks_critic = _sync.item() + + all_num_frames = [args.num_frame_per_block] * num_generated_blocks_critic if args.independent_first_frame: all_num_frames = [1] + all_num_frames + num_generated_frames_critic = sum(all_num_frames) # Initialize KV cache num_layers = generator_transformer3d.config.num_layers @@ -2556,8 +2817,20 @@ def convert_flow_pred_to_x0( current_start_frame += current_num_frames - fake_score_denoised_pred = output_pred - seq_len = frame_seq_length * num_frames + fake_score_denoised_pred_full = output_pred[:, :, :num_generated_frames_critic] + + # Slice for critic score: last score_num_frames frames + if num_generated_frames_critic > args.score_num_frames: + fake_score_denoised_pred, critic_score_num_frames, _ = slice_for_score( + fake_score_denoised_pred_full, vae, weight_dtype, + score_num_frames=args.score_num_frames, + independent_first_frame=args.independent_first_frame, + ) + else: + fake_score_denoised_pred = fake_score_denoised_pred_full + critic_score_num_frames = num_generated_frames_critic + + seq_len = frame_seq_length * critic_score_num_frames else: with torch.no_grad(): diff --git a/scripts/wan2.1_self_forcing/train_distill.sh b/scripts/wan2.1_self_forcing/train_distill.sh index 27e2d0e8..b46f5700 100644 --- a/scripts/wan2.1_self_forcing/train_distill.sh +++ b/scripts/wan2.1_self_forcing/train_distill.sh @@ -24,7 +24,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ diff --git a/scripts/wan2.1_vace/README_TRAIN.md b/scripts/wan2.1_vace/README_TRAIN.md index 63df464a..c3d049a9 100644 --- a/scripts/wan2.1_vace/README_TRAIN.md +++ b/scripts/wan2.1_vace/README_TRAIN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -194,8 +195,8 @@ VACE training datasets require both original videos and corresponding control si If data paths are relative, set in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height_add_objects.json" ``` **Absolute paths**: @@ -204,7 +205,7 @@ If data paths are absolute, set in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height_add_objects.json" ``` > 💡 **Recommendation**: Use relative paths for small local datasets; use absolute paths for external storage (NAS, OSS) or shared storage across multiple machines. diff --git a/scripts/wan2.1_vace/README_TRAIN_zh-CN.md b/scripts/wan2.1_vace/README_TRAIN_zh-CN.md index 957fa2fb..33aa63e9 100755 --- a/scripts/wan2.1_vace/README_TRAIN_zh-CN.md +++ b/scripts/wan2.1_vace/README_TRAIN_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -194,8 +195,8 @@ VACE 训练数据集除了原始视频外，还需要提供一一对应的控制 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height_add_objects.json" ``` **绝对路径**： @@ -204,7 +205,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height_add_objects.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 diff --git a/scripts/wan2.2/README_TRAIN.md b/scripts/wan2.2/README_TRAIN.md index 257c7c9d..7add18b5 100755 --- a/scripts/wan2.2/README_TRAIN.md +++ b/scripts/wan2.2/README_TRAIN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -146,7 +147,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -339,8 +340,8 @@ Wan2.2 adopts an innovative dual-Transformer architecture: | Parameter | Description | Example Value | |-----------|-------------|---------------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size | 1 | | `--image_sample_size` | Maximum training resolution for images | 640 | | `--video_sample_size` | Maximum training resolution for videos | 640 | diff --git a/scripts/wan2.2/README_TRAIN_DISTILL.md b/scripts/wan2.2/README_TRAIN_DISTILL.md index 6d818101..868cac34 100755 --- a/scripts/wan2.2/README_TRAIN_DISTILL.md +++ b/scripts/wan2.2/README_TRAIN_DISTILL.md @@ -45,9 +45,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -137,8 +138,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -147,7 +148,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (such as NAS, OSS) or shared across multiple machines, use absolute paths. @@ -207,7 +208,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -233,7 +234,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-I2V-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -256,7 +257,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -282,7 +283,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -305,7 +306,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -343,8 +344,8 @@ Wan2.2 adopts an innovative dual-Transformer architecture: | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size per GPU | 1 | | `--image_sample_size` | Maximum image training resolution | 640 | | `--video_sample_size` | Maximum video training resolution | 640 | @@ -356,7 +357,7 @@ Wan2.2 adopts an innovative dual-Transformer architecture: | `--num_train_epochs` | Number of training epochs | 100 | | `--checkpointing_steps` | Save checkpoint every N steps | 50 | | `--learning_rate` | Initial learning rate (generator) | 2e-06 | -| `--learning_rate_critic` | Initial learning rate (critic) | 2e-07 | +| `--learning_rate_critic` | Initial learning rate (critic) | 2e-06 | | `--lr_scheduler` | Learning rate scheduler | `constant_with_warmup` | | `--lr_warmup_steps` | Learning rate warmup steps | 100 | | `--seed` | Random seed | 42 | @@ -492,7 +493,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -551,7 +552,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -602,7 +603,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_distill.py \ --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -663,7 +664,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ diff --git a/scripts/wan2.2/README_TRAIN_DISTILL_LORA.md b/scripts/wan2.2/README_TRAIN_DISTILL_LORA.md index 34c5b019..1a96096c 100755 --- a/scripts/wan2.2/README_TRAIN_DISTILL_LORA.md +++ b/scripts/wan2.2/README_TRAIN_DISTILL_LORA.md @@ -45,9 +45,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -137,8 +138,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -147,7 +148,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -217,7 +218,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.2_distill_lora" \ --gradient_checkpointing \ @@ -244,7 +245,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-I2V-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -267,7 +268,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.2_distill_lora" \ --gradient_checkpointing \ @@ -314,8 +315,8 @@ In addition to distillation training, LoRA training adds the following specific | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Number of samples per batch | 1 | | `--image_sample_size` | Maximum training resolution for images | 640 | | `--video_sample_size` | Maximum training resolution for videos | 640 | @@ -327,7 +328,7 @@ In addition to distillation training, LoRA training adds the following specific | `--num_train_epochs` | Number of training epochs | 100 | | `--checkpointing_steps` | Save checkpoint every N steps | 50 | | `--learning_rate` | Initial learning rate (generator) | 1e-05 | -| `--learning_rate_critic` | Initial learning rate (discriminator) | 1e-06 | +| `--learning_rate_critic` | Initial learning rate (discriminator) | 1e-05 | | `--seed` | Random seed | 42 | | `--output_dir` | Output directory | `output_dir_wan2.2_distill_lora` | | `--gradient_checkpointing` | Activation recomputation | - | @@ -442,7 +443,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.2_distill_lora" \ --gradient_checkpointing \ @@ -504,7 +505,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.2_distill_lora" \ --gradient_checkpointing \ @@ -552,7 +553,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_distill_lora.py --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.2_distill_lora" \ --gradient_checkpointing \ @@ -614,7 +615,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.2_distill_lora" \ --gradient_checkpointing \ diff --git a/scripts/wan2.2/README_TRAIN_DISTILL_LORA_zh-CN.md b/scripts/wan2.2/README_TRAIN_DISTILL_LORA_zh-CN.md index c9bdfe93..54fbf58a 100755 --- a/scripts/wan2.2/README_TRAIN_DISTILL_LORA_zh-CN.md +++ b/scripts/wan2.2/README_TRAIN_DISTILL_LORA_zh-CN.md @@ -45,9 +45,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3:使用 docker** @@ -137,8 +138,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径,则在训练脚本中设置: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**: @@ -147,7 +148,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**:如果数据集较小且存储在本地,推荐使用相对路径;如果数据集存储在外部存储(如 NAS、OSS)或多个机器共享存储,推荐使用绝对路径。 @@ -217,7 +218,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.2_distill_lora" \ --gradient_checkpointing \ @@ -244,7 +245,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-I2V-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -267,7 +268,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.2_distill_lora" \ --gradient_checkpointing \ @@ -314,8 +315,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | @@ -327,7 +328,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--num_train_epochs` | 训练 epoch 数 | 100 | | `--checkpointing_steps` | 每 N 步保存 checkpoint | 50 | | `--learning_rate` | 初始学习率(生成器) | 1e-05 | -| `--learning_rate_critic` | 初始学习率(判别器) | 1e-06 | +| `--learning_rate_critic` | 初始学习率(判别器) | 1e-05 | | `--seed` | 随机种子 | 42 | | `--output_dir` | 输出目录 | `output_dir_wan2.2_distill_lora` | | `--gradient_checkpointing` | 激活重计算 | - | @@ -442,7 +443,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.2_distill_lora" \ --gradient_checkpointing \ @@ -504,7 +505,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.2_distill_lora" \ --gradient_checkpointing \ @@ -556,7 +557,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_distill_lora.py --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.2_distill_lora" \ --gradient_checkpointing \ @@ -617,7 +618,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.2_distill_lora" \ --gradient_checkpointing \ diff --git a/scripts/wan2.2/README_TRAIN_DISTILL_zh-CN.md b/scripts/wan2.2/README_TRAIN_DISTILL_zh-CN.md index 04941741..6c8b1d50 100755 --- a/scripts/wan2.2/README_TRAIN_DISTILL_zh-CN.md +++ b/scripts/wan2.2/README_TRAIN_DISTILL_zh-CN.md @@ -45,9 +45,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -137,8 +138,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -147,7 +148,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -207,7 +208,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -233,7 +234,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-I2V-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -256,7 +257,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -282,7 +283,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -305,7 +306,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -343,8 +344,8 @@ Wan2.2采用了创新的双Transformer架构： | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | @@ -356,7 +357,7 @@ Wan2.2采用了创新的双Transformer架构： | `--num_train_epochs` | 训练 epoch 数 | 100 | | `--checkpointing_steps` | 每 N 步保存 checkpoint | 50 | | `--learning_rate` | 初始学习率（生成器） | 2e-06 | -| `--learning_rate_critic` | 初始学习率（判别器） | 2e-07 | +| `--learning_rate_critic` | 初始学习率（判别器） | 2e-06 | | `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` | | `--lr_warmup_steps` | 学习率预热步数 | 100 | | `--seed` | 随机种子 | 42 | @@ -492,7 +493,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -551,7 +552,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -602,7 +603,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_distill.py \ --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -663,7 +664,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ diff --git a/scripts/wan2.2/README_TRAIN_LORA.md b/scripts/wan2.2/README_TRAIN_LORA.md index 3a899811..ec875e2e 100755 --- a/scripts/wan2.2/README_TRAIN_LORA.md +++ b/scripts/wan2.2/README_TRAIN_LORA.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, set in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -146,7 +147,7 @@ If your data uses absolute paths, set in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Suggestion**: If the dataset is small and stored locally, relative paths are recommended. If the dataset is stored on external storage (such as NAS, OSS) or shared across multiple machines, absolute paths are recommended. @@ -343,8 +344,8 @@ Wan2.2 adopts an innovative dual-Transformer architecture: |-----------|-------------|---------------| | `--config_path` | Configuration file path | `config/wan2.2/wan_civitai_t2v.yaml` | | `--pretrained_model_name_or_path` | Pre-trained model path | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Number of samples per batch | 1 | | `--image_sample_size` | Maximum training resolution for images | 640 | | `--video_sample_size` | Maximum training resolution for videos | 640 | diff --git a/scripts/wan2.2/README_TRAIN_LORA_zh-CN.md b/scripts/wan2.2/README_TRAIN_LORA_zh-CN.md index 02487aa0..bc544574 100644 --- a/scripts/wan2.2/README_TRAIN_LORA_zh-CN.md +++ b/scripts/wan2.2/README_TRAIN_LORA_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -146,7 +147,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -343,8 +344,8 @@ Wan2.2采用了创新的双Transformer架构： |-----|------|-------| | `--config_path` | 配置文件路径 | `config/wan2.2/wan_civitai_t2v.yaml` | | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | diff --git a/scripts/wan2.2/README_TRAIN_S2V.md b/scripts/wan2.2/README_TRAIN_S2V.md index 047a5554..478e1ac3 100755 --- a/scripts/wan2.2/README_TRAIN_S2V.md +++ b/scripts/wan2.2/README_TRAIN_S2V.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -154,8 +155,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Audios-Demo --local_dir ./dataset If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata_control.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -164,7 +165,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata_control.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -244,8 +245,8 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_s2v.py \ |-----|------|-------| | `--config_path` | Model configuration file path | `config/wan2.2/wan_civitai_s2v.yaml` | | `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/Wan2.2-S2V-14B` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata_control.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--video_sample_size` | Maximum video resolution for training | 640 | | `--token_sample_size` | Token length sampling size | 640 | diff --git a/scripts/wan2.2/README_TRAIN_S2V_zh-CN.md b/scripts/wan2.2/README_TRAIN_S2V_zh-CN.md index d68ac9fa..c0001068 100755 --- a/scripts/wan2.2/README_TRAIN_S2V_zh-CN.md +++ b/scripts/wan2.2/README_TRAIN_S2V_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式三：使用 Docker** @@ -147,8 +148,8 @@ Wan2.2-S2V 的 `metadata.json` 与 VideoX-Fun 的普通 JSON 格式略有不同 如果你的数据使用的是相对路径，训练脚本中请这样配置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Audios-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -157,7 +158,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存放在本地，请使用相对路径。如果数据集存放在外部存储（如 NAS、OSS）或多机共享，请使用绝对路径。 @@ -239,8 +240,8 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_s2v.py \ |------|------|--------| | `--config_path` | 模型配置文件路径 | `config/wan2.2/wan_civitai_s2v.yaml` | | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.2-S2V-14B` | -| `--train_data_dir` | 训练数据集目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据集元数据文件 | `datasets/internal_datasets/metadata_control.json` | +| `--train_data_dir` | 训练数据集目录 | `datasets/X-Fun-Videos-Audios-Demo/` | +| `--train_data_meta` | 训练数据集元数据文件 | `datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json` | | `--video_sample_size` | 视频采样尺寸（最大分辨率） | `640` | | `--token_sample_size` | Token 采样尺寸 | `640` | | `--video_sample_stride` | 视频采样步幅 | `2` | diff --git a/scripts/wan2.2/README_TRAIN_zh-CN.md b/scripts/wan2.2/README_TRAIN_zh-CN.md index aaec812c..0d4bb6d9 100644 --- a/scripts/wan2.2/README_TRAIN_zh-CN.md +++ b/scripts/wan2.2/README_TRAIN_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -146,7 +147,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -339,8 +340,8 @@ Wan2.2采用了创新的双Transformer架构： | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | diff --git a/scripts/wan2.2/train.py b/scripts/wan2.2/train.py index ebf56ade..9a8370c2 100644 --- a/scripts/wan2.2/train.py +++ b/scripts/wan2.2/train.py @@ -1775,12 +1775,12 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): if args.train_mode == "ti2v": if rng is None: - t2v_in_ti2v = np.random.choice([0, 1], p = [0.50, 0.50]) + i2v_in_ti2v = np.random.choice([0, 1], p = [0.50, 0.50]) else: - t2v_in_ti2v = rng.choice([0, 1], p = [0.50, 0.50]) + i2v_in_ti2v = rng.choice([0, 1], p = [0.50, 0.50]) mask_bs = mask.size()[0] - if t2v_in_ti2v: + if i2v_in_ti2v: noisy_latents = (1 - mask) * inpaint_latents + mask * noisy_latents temp_ts = (mask[:, 0, :, ::2, ::2] * timesteps[:, None, None, None]).flatten(1) diff --git a/scripts/wan2.2/train_distill.py b/scripts/wan2.2/train_distill.py index 367f9942..5cb85f0c 100644 --- a/scripts/wan2.2/train_distill.py +++ b/scripts/wan2.2/train_distill.py @@ -764,6 +764,11 @@ def parse_args(): default=[1000, 750, 500, 250], help="The denoising step list.", ) + parser.add_argument( + "--randomize_step_indices", + action="store_true", + help="whether to use randomize timesteps indices in training.", + ) args = parser.parse_args() env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) @@ -891,6 +896,7 @@ def main(): weight_dtype = torch.bfloat16 args.mixed_precision = accelerator.mixed_precision + args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list] # Load scheduler, tokenizer and models. noise_scheduler = FlowMatchEulerDiscreteScheduler( **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs'])) @@ -977,7 +983,13 @@ def deepspeed_zero_init_disabled_context_manager(): state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict m, u = generator_transformer3d.load_state_dict(state_dict, strict=False) - print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + print(f"generator missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + m, u = fake_score_transformer3d.load_state_dict(state_dict, strict=False) + print(f"fake_score missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + m, u = real_score_transformer3d.load_state_dict(state_dict, strict=False) + print(f"real_score missing keys: {len(m)}, unexpected keys: {len(u)}") assert len(u) == 0 if args.vae_path is not None: @@ -1139,8 +1151,8 @@ def load_model_hook(models, input_dir): fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_transformer3d.parameters())) fake_trainable_params_optim = [ - {'params': [], 'lr': args.learning_rate}, - {'params': [], 'lr': args.learning_rate / 2}, + {'params': [], 'lr': args.learning_rate_critic}, + {'params': [], 'lr': args.learning_rate_critic / 2}, ] in_already = [] for name, param in fake_score_transformer3d.named_parameters(): @@ -1153,7 +1165,7 @@ def load_model_hook(models, input_dir): high_lr_flag = True fake_trainable_params_optim[0]['params'].append(param) if accelerator.is_main_process: - print(f"Set {name} to lr : {args.learning_rate}") + print(f"Set {name} to lr : {args.learning_rate_critic}") break if high_lr_flag: continue @@ -1162,7 +1174,7 @@ def load_model_hook(models, input_dir): in_already.append(name) fake_trainable_params_optim[1]['params'].append(param) if accelerator.is_main_process: - print(f"Set {name} to lr : {args.learning_rate / 2}") + print(f"Set {name} to lr : {args.learning_rate_critic / 2}") break if args.use_came: @@ -1198,6 +1210,7 @@ def load_model_hook(models, input_dir): # Get the training dataset sample_n_frames_bucket_interval = vae.config.temporal_compression_ratio + spatial_compression_ratio = vae.config.spatial_compression_ratio if args.fix_sample_size is not None and args.enable_bucket: args.video_sample_size = max(max(args.fix_sample_size), args.video_sample_size) @@ -1296,7 +1309,7 @@ def collate_fn(examples): aspect_ratio_random_crop_sample_size = {key : [x / 512 * args.video_sample_size / random_downsample_ratio for x in ASPECT_RATIO_RANDOM_CROP_512[key]] for key in ASPECT_RATIO_RANDOM_CROP_512.keys()} if args.fix_sample_size is not None: - fix_sample_size = [int(x / 16) * 16 for x in args.fix_sample_size] + fix_sample_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in args.fix_sample_size] elif args.random_ratio_crop: if rng is None: random_sample_size = aspect_ratio_random_crop_sample_size[ @@ -1306,10 +1319,10 @@ def collate_fn(examples): random_sample_size = aspect_ratio_random_crop_sample_size[ rng.choice(list(aspect_ratio_random_crop_sample_size.keys()), p = ASPECT_RATIO_RANDOM_CROP_PROB) ] - random_sample_size = [int(x / 16) * 16 for x in random_sample_size] + random_sample_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in random_sample_size] else: closest_size, closest_ratio = get_closest_ratio(h, w, ratios=aspect_ratio_sample_size) - closest_size = [int(x / 16) * 16 for x in closest_size] + closest_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in closest_size] min_example_length = min( [example["pixel_values"].shape[0] for example in examples] @@ -1533,7 +1546,7 @@ def collate_fn(examples): ) fake_score_lr_scheduler = get_scheduler( args.lr_scheduler, - optimizer=optimizer, + optimizer=critic_optimizer, num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, num_training_steps=args.max_train_steps * accelerator.num_processes, ) @@ -1671,7 +1684,55 @@ def unwrap_model(model): train_sampling_steps = args.train_sampling_steps idx_sampling = DiscreteSampling(train_sampling_steps, start_num_idx=start_num_idx, uniform_sampling=args.uniform_sampling) - denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)] + + def randomize_denoising_step_indices( + denoising_step_indices_list, + train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=0.3, + ): + indices = list(denoising_step_indices_list) + n = len(indices) + + if n <= 2: + low = indices[1] + high = indices[0] - 1 + random_tail = torch.randint(low, high + 1, (1,)).item() + + result = torch.tensor([indices[0], random_tail]) + else: + result = [0] * n + result[0] = indices[0] + result[-1] = indices[-1] + + for i in range(1, n - 1): + gap_upper = indices[i - 1] - indices[i] + gap_lower = indices[i] - indices[i + 1] + + max_jitter = int(min(gap_upper, gap_lower) * jitter_ratio) + + if max_jitter > 0: + jitter = torch.randint( + -max_jitter, max_jitter + 1, (1,) + ).item() + else: + jitter = 0 + + result[i] = indices[i] + jitter + + for i in range(1, n): + if result[i] >= result[i - 1]: + result[i] = result[i - 1] - 1 + + result = [max(1, min(train_sampling_steps, x)) for x in result] + result = torch.tensor(result) + + if dist.is_initialized(): + result = result.to(accelerator.device) + dist.broadcast(result, src=0) + result = result.cpu() + return result for epoch in range(first_epoch, args.num_train_epochs): train_dmd_loss = 0.0 @@ -1836,7 +1897,7 @@ def _batch_encode_vae(pixel_values): else: text = batch['text'] if args.fix_sample_size is not None: - local_sample_size = [int(x / 16) * 16 for x in args.fix_sample_size] + local_sample_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in args.fix_sample_size] num_frames = args.video_sample_n_frames else: if args.random_hw_adapt and args.training_with_video_token_length: @@ -1865,7 +1926,7 @@ def _batch_encode_vae(pixel_values): else: aspect_ratio_key = rng.choice(list(aspect_ratio_sample_size.keys())) local_sample_size = aspect_ratio_sample_size[aspect_ratio_key] - local_sample_size = [int(x / 16) * 16 for x in local_sample_size] + local_sample_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in local_sample_size] target_shape = ( len(text), @@ -1930,9 +1991,9 @@ def _batch_encode_vae(pixel_values): if args.train_mode == "ti2v": if rng is None: - t2v_in_ti2v = np.random.choice([0, 1], p = [0.50, 0.50]) + i2v_in_ti2v = np.random.choice([0, 1], p = [0.50, 0.50]) else: - t2v_in_ti2v = rng.choice([0, 1], p = [0.50, 0.50]) + i2v_in_ti2v = rng.choice([0, 1], p = [0.50, 0.50]) with accelerator.accumulate(generator_transformer3d): def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): @@ -1994,6 +2055,20 @@ def convert_flow_pred_to_x0( x0_pred = xt - sigma_t * flow_pred return x0_pred.to(original_dtype) + # Create discrete denoising steps (per-step, with optional randomization) + if getattr(args, 'randomize_step_indices', False): + random_indices = randomize_denoising_step_indices( + args.denoising_step_indices_list, + args.train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=getattr(args, 'index_jitter_ratio', 0.30), + ) + else: + random_indices = torch.tensor(args.denoising_step_indices_list) + + denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - random_indices] + # --- Main Training Logic --- bsz, channel, num_frames, height, width = target_shape if step % args.gen_update_interval == 0: @@ -2020,7 +2095,7 @@ def convert_flow_pred_to_x0( with context_manager: if args.train_mode == "ti2v": mask_bs = mask.size()[0] - if t2v_in_ti2v: + if i2v_in_ti2v: _generator_noise = (1 - mask) * inpaint_latents + mask * generator_noise temp_ts = (mask[:, 0, :, ::2, ::2] * timestep[:, None, None, None]).flatten(1) @@ -2041,7 +2116,7 @@ def convert_flow_pred_to_x0( generator_pred = convert_flow_pred_to_x0( scheduler=noise_scheduler, flow_pred=generator_pred, - xt=generator_noise, + xt=_generator_noise, timestep=timestep ) @@ -2069,8 +2144,8 @@ def convert_flow_pred_to_x0( with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device), torch.no_grad(): if args.train_mode == "ti2v": mask_bs = mask.size()[0] - if t2v_in_ti2v: - _generator_denoised_input = (1 - mask) * inpaint_latents + mask * generator_denoised_input + if i2v_in_ti2v: + _generator_denoised_input = (1 - mask) * generator_pred + mask * generator_denoised_input temp_ts = (mask[:, 0, :, ::2, ::2] * generator_timestep[:, None, None, None]).flatten(1) _generator_timestep = torch.cat([temp_ts, temp_ts.new_ones(mask_bs, seq_len - temp_ts.size(1)) * generator_timestep[:, None,]], dim = 1) @@ -2200,7 +2275,7 @@ def convert_flow_pred_to_x0( with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): if args.train_mode == "ti2v": mask_bs = mask.size()[0] - if t2v_in_ti2v: + if i2v_in_ti2v: _fake_score_critic_noise = (1 - mask) * inpaint_latents + mask * fake_score_critic_noise temp_ts = (mask[:, 0, :, ::2, ::2] * timestep[:, None, None, None]).flatten(1) @@ -2221,7 +2296,7 @@ def convert_flow_pred_to_x0( fake_score_denoised_pred = convert_flow_pred_to_x0( scheduler=noise_scheduler, flow_pred=fake_score_denoised_pred, - xt=fake_score_critic_noise, + xt=_fake_score_critic_noise, timestep=timestep ) @@ -2251,8 +2326,8 @@ def convert_flow_pred_to_x0( ) if args.train_mode == "ti2v": mask_bs = mask.size()[0] - if t2v_in_ti2v: - fake_score_denoised_input = (1 - mask) * inpaint_latents + mask * fake_score_denoised_input + if i2v_in_ti2v: + fake_score_denoised_input = (1 - mask) * fake_score_denoised_pred + mask * fake_score_denoised_input temp_ts = (mask[:, 0, :, ::2, ::2] * critic_timestep[:, None, None, None]).flatten(1) _critic_timestep = torch.cat([temp_ts, temp_ts.new_ones(mask_bs, seq_len - temp_ts.size(1)) * critic_timestep[:, None,]], dim = 1) diff --git a/scripts/wan2.2/train_distill.sh b/scripts/wan2.2/train_distill.sh index 8c91e4e8..fb1c5c24 100644 --- a/scripts/wan2.2/train_distill.sh +++ b/scripts/wan2.2/train_distill.sh @@ -23,7 +23,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_distill.py \ --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07\ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ diff --git a/scripts/wan2.2/train_distill_lora.py b/scripts/wan2.2/train_distill_lora.py index a7d5d938..3a2d4da8 100644 --- a/scripts/wan2.2/train_distill_lora.py +++ b/scripts/wan2.2/train_distill_lora.py @@ -778,6 +778,11 @@ def parse_args(): default=[1000, 750, 500, 250], help="The denoising step list.", ) + parser.add_argument( + "--randomize_step_indices", + action="store_true", + help="whether to use randomize timesteps indices in training.", + ) parser.add_argument( "--lora_skip_name", type=str, @@ -917,6 +922,7 @@ def main(): weight_dtype = torch.bfloat16 args.mixed_precision = accelerator.mixed_precision + args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list] # Load scheduler, tokenizer and models. noise_scheduler = FlowMatchEulerDiscreteScheduler( **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs'])) @@ -993,6 +999,9 @@ def deepspeed_zero_init_disabled_context_manager(): real_score_transformer3d.requires_grad_(False) fake_score_transformer3d.requires_grad_(False) + generator_transformer3d.train() + fake_score_transformer3d.train() + # Lora will work with this... if args.use_peft_lora: from peft import (LoraConfig, get_peft_model_state_dict, @@ -1040,7 +1049,13 @@ def deepspeed_zero_init_disabled_context_manager(): state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict m, u = generator_transformer3d.load_state_dict(state_dict, strict=False) - print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + print(f"generator missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + m, u = fake_score_transformer3d.load_state_dict(state_dict, strict=False) + print(f"fake_score missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + m, u = real_score_transformer3d.load_state_dict(state_dict, strict=False) + print(f"real_score missing keys: {len(m)}, unexpected keys: {len(u)}") assert len(u) == 0 if args.vae_path is not None: @@ -1180,7 +1195,7 @@ def load_model_hook(models, input_dir): logging.info("Add fake_score_network parameters") fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_network.parameters())) - fake_trainable_params_optim = fake_score_network.prepare_optimizer_params(args.learning_rate / 2, args.learning_rate, args.learning_rate) + fake_trainable_params_optim = fake_score_network.prepare_optimizer_params(args.learning_rate_critic / 2, args.learning_rate_critic, args.learning_rate_critic) if args.use_came: optimizer = optimizer_cls( @@ -1215,6 +1230,7 @@ def load_model_hook(models, input_dir): # Get the training dataset sample_n_frames_bucket_interval = vae.config.temporal_compression_ratio + spatial_compression_ratio = vae.config.spatial_compression_ratio if args.fix_sample_size is not None and args.enable_bucket: args.video_sample_size = max(max(args.fix_sample_size), args.video_sample_size) @@ -1313,7 +1329,7 @@ def collate_fn(examples): aspect_ratio_random_crop_sample_size = {key : [x / 512 * args.video_sample_size / random_downsample_ratio for x in ASPECT_RATIO_RANDOM_CROP_512[key]] for key in ASPECT_RATIO_RANDOM_CROP_512.keys()} if args.fix_sample_size is not None: - fix_sample_size = [int(x / 16) * 16 for x in args.fix_sample_size] + fix_sample_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in args.fix_sample_size] elif args.random_ratio_crop: if rng is None: random_sample_size = aspect_ratio_random_crop_sample_size[ @@ -1323,10 +1339,10 @@ def collate_fn(examples): random_sample_size = aspect_ratio_random_crop_sample_size[ rng.choice(list(aspect_ratio_random_crop_sample_size.keys()), p = ASPECT_RATIO_RANDOM_CROP_PROB) ] - random_sample_size = [int(x / 16) * 16 for x in random_sample_size] + random_sample_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in random_sample_size] else: closest_size, closest_ratio = get_closest_ratio(h, w, ratios=aspect_ratio_sample_size) - closest_size = [int(x / 16) * 16 for x in closest_size] + closest_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in closest_size] min_example_length = min( [example["pixel_values"].shape[0] for example in examples] @@ -1550,7 +1566,7 @@ def collate_fn(examples): ) fake_score_lr_scheduler = get_scheduler( args.lr_scheduler, - optimizer=optimizer, + optimizer=critic_optimizer, num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, num_training_steps=args.max_train_steps * accelerator.num_processes, ) @@ -1713,7 +1729,55 @@ def save_model(ckpt_file, unwrapped_nw): train_sampling_steps = args.train_sampling_steps idx_sampling = DiscreteSampling(train_sampling_steps, start_num_idx=start_num_idx, uniform_sampling=args.uniform_sampling) - denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)] + + def randomize_denoising_step_indices( + denoising_step_indices_list, + train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=0.3, + ): + indices = list(denoising_step_indices_list) + n = len(indices) + + if n <= 2: + low = indices[1] + high = indices[0] - 1 + random_tail = torch.randint(low, high + 1, (1,)).item() + + result = torch.tensor([indices[0], random_tail]) + else: + result = [0] * n + result[0] = indices[0] + result[-1] = indices[-1] + + for i in range(1, n - 1): + gap_upper = indices[i - 1] - indices[i] + gap_lower = indices[i] - indices[i + 1] + + max_jitter = int(min(gap_upper, gap_lower) * jitter_ratio) + + if max_jitter > 0: + jitter = torch.randint( + -max_jitter, max_jitter + 1, (1,) + ).item() + else: + jitter = 0 + + result[i] = indices[i] + jitter + + for i in range(1, n): + if result[i] >= result[i - 1]: + result[i] = result[i - 1] - 1 + + result = [max(1, min(train_sampling_steps, x)) for x in result] + result = torch.tensor(result) + + if dist.is_initialized(): + result = result.to(accelerator.device) + dist.broadcast(result, src=0) + result = result.cpu() + return result for epoch in range(first_epoch, args.num_train_epochs): train_dmd_loss = 0.0 @@ -1878,7 +1942,7 @@ def _batch_encode_vae(pixel_values): else: text = batch['text'] if args.fix_sample_size is not None: - local_sample_size = [int(x / 16) * 16 for x in args.fix_sample_size] + local_sample_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in args.fix_sample_size] num_frames = args.video_sample_n_frames else: if args.random_hw_adapt and args.training_with_video_token_length: @@ -1907,7 +1971,7 @@ def _batch_encode_vae(pixel_values): else: aspect_ratio_key = rng.choice(list(aspect_ratio_sample_size.keys())) local_sample_size = aspect_ratio_sample_size[aspect_ratio_key] - local_sample_size = [int(x / 16) * 16 for x in local_sample_size] + local_sample_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in local_sample_size] target_shape = ( len(text), @@ -1972,9 +2036,9 @@ def _batch_encode_vae(pixel_values): if args.train_mode == "ti2v": if rng is None: - t2v_in_ti2v = np.random.choice([0, 1], p = [0.50, 0.50]) + i2v_in_ti2v = np.random.choice([0, 1], p = [0.50, 0.50]) else: - t2v_in_ti2v = rng.choice([0, 1], p = [0.50, 0.50]) + i2v_in_ti2v = rng.choice([0, 1], p = [0.50, 0.50]) with accelerator.accumulate(generator_transformer3d): def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): @@ -2036,6 +2100,20 @@ def convert_flow_pred_to_x0( x0_pred = xt - sigma_t * flow_pred return x0_pred.to(original_dtype) + # Create discrete denoising steps (per-step, with optional randomization) + if getattr(args, 'randomize_step_indices', False): + random_indices = randomize_denoising_step_indices( + args.denoising_step_indices_list, + args.train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=getattr(args, 'index_jitter_ratio', 0.30), + ) + else: + random_indices = torch.tensor(args.denoising_step_indices_list) + + denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - random_indices] + # --- Main Training Logic --- bsz, channel, num_frames, height, width = target_shape if step % args.gen_update_interval == 0: @@ -2062,7 +2140,7 @@ def convert_flow_pred_to_x0( with context_manager: if args.train_mode == "ti2v": mask_bs = mask.size()[0] - if t2v_in_ti2v: + if i2v_in_ti2v: _generator_noise = (1 - mask) * inpaint_latents + mask * generator_noise temp_ts = (mask[:, 0, :, ::2, ::2] * timestep[:, None, None, None]).flatten(1) @@ -2083,7 +2161,7 @@ def convert_flow_pred_to_x0( generator_pred = convert_flow_pred_to_x0( scheduler=noise_scheduler, flow_pred=generator_pred, - xt=generator_noise, + xt=_generator_noise, timestep=timestep ) @@ -2111,8 +2189,8 @@ def convert_flow_pred_to_x0( with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device), torch.no_grad(): if args.train_mode == "ti2v": mask_bs = mask.size()[0] - if t2v_in_ti2v: - _generator_denoised_input = (1 - mask) * inpaint_latents + mask * generator_denoised_input + if i2v_in_ti2v: + _generator_denoised_input = (1 - mask) * generator_pred + mask * generator_denoised_input temp_ts = (mask[:, 0, :, ::2, ::2] * generator_timestep[:, None, None, None]).flatten(1) _generator_timestep = torch.cat([temp_ts, temp_ts.new_ones(mask_bs, seq_len - temp_ts.size(1)) * generator_timestep[:, None,]], dim = 1) @@ -2242,7 +2320,7 @@ def convert_flow_pred_to_x0( with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): if args.train_mode == "ti2v": mask_bs = mask.size()[0] - if t2v_in_ti2v: + if i2v_in_ti2v: _fake_score_critic_noise = (1 - mask) * inpaint_latents + mask * fake_score_critic_noise temp_ts = (mask[:, 0, :, ::2, ::2] * timestep[:, None, None, None]).flatten(1) @@ -2263,7 +2341,7 @@ def convert_flow_pred_to_x0( fake_score_denoised_pred = convert_flow_pred_to_x0( scheduler=noise_scheduler, flow_pred=fake_score_denoised_pred, - xt=fake_score_critic_noise, + xt=_fake_score_critic_noise, timestep=timestep ) @@ -2293,8 +2371,8 @@ def convert_flow_pred_to_x0( ) if args.train_mode == "ti2v": mask_bs = mask.size()[0] - if t2v_in_ti2v: - fake_score_denoised_input = (1 - mask) * inpaint_latents + mask * fake_score_denoised_input + if i2v_in_ti2v: + fake_score_denoised_input = (1 - mask) * fake_score_denoised_pred + mask * fake_score_denoised_input temp_ts = (mask[:, 0, :, ::2, ::2] * critic_timestep[:, None, None, None]).flatten(1) _critic_timestep = torch.cat([temp_ts, temp_ts.new_ones(mask_bs, seq_len - temp_ts.size(1)) * critic_timestep[:, None,]], dim = 1) diff --git a/scripts/wan2.2/train_distill_lora.sh b/scripts/wan2.2/train_distill_lora.sh index afba1116..e18b1783 100644 --- a/scripts/wan2.2/train_distill_lora.sh +++ b/scripts/wan2.2/train_distill_lora.sh @@ -23,7 +23,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_distill_lora.py --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-05 \ - --learning_rate_critic=1e-06 \ + --learning_rate_critic=1e-05 \ --seed=42 \ --output_dir="output_dir_wan2.2_distill_lora" \ --gradient_checkpointing \ diff --git a/scripts/wan2.2/train_lora.py b/scripts/wan2.2/train_lora.py index 2fb1015c..32b5115d 100755 --- a/scripts/wan2.2/train_lora.py +++ b/scripts/wan2.2/train_lora.py @@ -1821,12 +1821,12 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): if args.train_mode == "ti2v": if rng is None: - t2v_in_ti2v = np.random.choice([0, 1], p = [0.50, 0.50]) + i2v_in_ti2v = np.random.choice([0, 1], p = [0.50, 0.50]) else: - t2v_in_ti2v = rng.choice([0, 1], p = [0.50, 0.50]) + i2v_in_ti2v = rng.choice([0, 1], p = [0.50, 0.50]) mask_bs = mask.size()[0] - if t2v_in_ti2v: + if i2v_in_ti2v: noisy_latents = (1 - mask) * inpaint_latents + mask * noisy_latents temp_ts = (mask[:, 0, :, ::2, ::2] * timesteps[:, None, None, None]).flatten(1) diff --git a/scripts/wan2.2_fun/README_TRAIN.md b/scripts/wan2.2_fun/README_TRAIN.md index db13c532..14775299 100755 --- a/scripts/wan2.2_fun/README_TRAIN.md +++ b/scripts/wan2.2_fun/README_TRAIN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -146,7 +147,7 @@ If your data uses absolute paths, configure in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -337,8 +338,8 @@ Wan2.2 adopts an innovative dual-Transformer architecture: | Parameter | Description | Example Value | |-----------|-------------|---------------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.2-Fun-A14B-InP` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size | 1 | | `--image_sample_size` | Maximum training resolution for images | 640 | | `--video_sample_size` | Maximum training resolution for videos | 640 | diff --git a/scripts/wan2.2_fun/README_TRAIN_CONTROL.md b/scripts/wan2.2_fun/README_TRAIN_CONTROL.md index 74fda380..6fed2329 100755 --- a/scripts/wan2.2_fun/README_TRAIN_CONTROL.md +++ b/scripts/wan2.2_fun/README_TRAIN_CONTROL.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -147,8 +148,8 @@ Control training datasets require both original videos and corresponding control If your data uses relative paths, set in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -157,7 +158,7 @@ If your data uses absolute paths, set in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: Use relative paths for small local datasets; use absolute paths for external storage (NAS, OSS) or shared multi-machine storage. diff --git a/scripts/wan2.2_fun/README_TRAIN_CONTROL_LORA.md b/scripts/wan2.2_fun/README_TRAIN_CONTROL_LORA.md index 13150c0b..7fa7abc2 100755 --- a/scripts/wan2.2_fun/README_TRAIN_CONTROL_LORA.md +++ b/scripts/wan2.2_fun/README_TRAIN_CONTROL_LORA.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -147,8 +148,8 @@ In addition to the original videos, the Control training dataset also requires c If your data uses relative paths, set in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -157,7 +158,7 @@ If your data uses absolute paths, set in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Suggestion**: If the dataset is small and stored locally, relative paths are recommended. If the dataset is stored on external storage (such as NAS, OSS) or shared across multiple machines, absolute paths are recommended. diff --git a/scripts/wan2.2_fun/README_TRAIN_CONTROL_LORA_zh-CN.md b/scripts/wan2.2_fun/README_TRAIN_CONTROL_LORA_zh-CN.md index 69f048af..619a3f88 100755 --- a/scripts/wan2.2_fun/README_TRAIN_CONTROL_LORA_zh-CN.md +++ b/scripts/wan2.2_fun/README_TRAIN_CONTROL_LORA_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -147,8 +148,8 @@ Control 训练数据集除了原始视频外，还需要提供一一对应的控 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -157,7 +158,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 diff --git a/scripts/wan2.2_fun/README_TRAIN_CONTROL_zh-CN.md b/scripts/wan2.2_fun/README_TRAIN_CONTROL_zh-CN.md index 7edd7869..ecc3ab5a 100755 --- a/scripts/wan2.2_fun/README_TRAIN_CONTROL_zh-CN.md +++ b/scripts/wan2.2_fun/README_TRAIN_CONTROL_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -147,8 +148,8 @@ Control 训练数据集除了原始视频外，还需要提供一一对应的控 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -157,7 +158,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 diff --git a/scripts/wan2.2_fun/README_TRAIN_LORA.md b/scripts/wan2.2_fun/README_TRAIN_LORA.md index e8fba284..84274da1 100755 --- a/scripts/wan2.2_fun/README_TRAIN_LORA.md +++ b/scripts/wan2.2_fun/README_TRAIN_LORA.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -135,8 +136,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, set in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **Absolute Path**: @@ -145,7 +146,7 @@ If your data uses absolute paths, set in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Suggestion**: If the dataset is small and stored locally, relative paths are recommended. If the dataset is stored on external storage (such as NAS, OSS) or shared across multiple machines, absolute paths are recommended. diff --git a/scripts/wan2.2_fun/README_TRAIN_LORA_zh-CN.md b/scripts/wan2.2_fun/README_TRAIN_LORA_zh-CN.md index ad5f07ea..ef7c9dbe 100644 --- a/scripts/wan2.2_fun/README_TRAIN_LORA_zh-CN.md +++ b/scripts/wan2.2_fun/README_TRAIN_LORA_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -135,8 +136,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -145,7 +146,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 diff --git a/scripts/wan2.2_fun/README_TRAIN_zh-CN.md b/scripts/wan2.2_fun/README_TRAIN_zh-CN.md index 66120bec..7ddb250d 100644 --- a/scripts/wan2.2_fun/README_TRAIN_zh-CN.md +++ b/scripts/wan2.2_fun/README_TRAIN_zh-CN.md @@ -44,9 +44,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -136,8 +137,8 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -146,7 +147,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -343,8 +344,8 @@ Wan2.2 Fun采用了创新的双Transformer架构（A14B版本）： | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.2-Fun-A14B-InP` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | diff --git a/scripts/wan2.2_fun/train_distill.py b/scripts/wan2.2_fun/train_distill.py new file mode 100644 index 00000000..247c6104 --- /dev/null +++ b/scripts/wan2.2_fun/train_distill.py @@ -0,0 +1,2455 @@ +"""Modified from https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py +""" +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# 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 + +import argparse +import contextlib +import gc +import json +import logging +import math +import os +import pickle +import shutil +import sys + +import accelerate +import diffusers +import numpy as np +import torch +import torch.distributed as dist +import torch.nn.functional as F +import torch.utils.checkpoint +import torchvision.transforms.functional as TF +import transformers +from accelerate import Accelerator, FullyShardedDataParallelPlugin +from accelerate.logging import get_logger +from accelerate.state import AcceleratorState +from accelerate.utils import ProjectConfiguration, set_seed +from diffusers import DDIMScheduler, FlowMatchEulerDiscreteScheduler +from diffusers.optimization import get_scheduler +from diffusers.training_utils import (EMAModel, + compute_density_for_timestep_sampling, + compute_loss_weighting_for_sd3) +from diffusers.utils import check_min_version, deprecate, is_wandb_available +from diffusers.utils.torch_utils import is_compiled_module +from einops import rearrange +from omegaconf import OmegaConf +from packaging import version +from PIL import Image +from torch.distributed.fsdp.fully_sharded_data_parallel import ( + FullOptimStateDictConfig, FullStateDictConfig, ShardedOptimStateDictConfig, + ShardedStateDictConfig) +from torch.utils.data import Dataset, RandomSampler, BatchSampler +from torch.utils.tensorboard import SummaryWriter +from torchvision import transforms +from tqdm.auto import tqdm +from transformers import AutoTokenizer +from transformers.utils import ContextManagers + +import datasets + +current_file_path = os.path.abspath(__file__) +project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))] +for project_root in project_roots: + sys.path.insert(0, project_root) if project_root not in sys.path else None + +from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512, + ASPECT_RATIO_RANDOM_CROP_PROB, + AspectRatioBatchImageVideoSampler, + ImageVideoDataset, ImageVideoSampler, + RandomSampler, TextDataset, get_closest_ratio, + get_random_mask) +from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8, + CLIPModel, Wan2_2Transformer3DModel, + WanT5EncoderModel) +from videox_fun.pipeline import Wan2_2I2VPipeline, Wan2_2Pipeline +from videox_fun.utils.discrete_sampler import DiscreteSampling +from videox_fun.utils.utils import (calculate_dimensions, get_image_latent, + get_image_to_video_latent, + save_videos_grid) + +if is_wandb_available(): + import wandb + + +def filter_kwargs(cls, kwargs): + import inspect + sig = inspect.signature(cls.__init__) + valid_params = set(sig.parameters.keys()) - {'self', 'cls'} + filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params} + return filtered_kwargs + +def get_random_downsample_ratio(sample_size, image_ratio=[], + all_choices=False, rng=None): + def _create_special_list(length): + if length == 1: + return [1.0] + if length >= 2: + first_element = 0.75 + remaining_sum = 1.0 - first_element + other_elements_value = remaining_sum / (length - 1) + special_list = [first_element] + [other_elements_value] * (length - 1) + return special_list + + if sample_size >= 1536: + number_list = [1, 1.25, 1.5, 2, 2.5, 3] + image_ratio + elif sample_size >= 1024: + number_list = [1, 1.25, 1.5, 2] + image_ratio + elif sample_size >= 768: + number_list = [1, 1.25, 1.5] + image_ratio + elif sample_size >= 512: + number_list = [1] + image_ratio + else: + number_list = [1] + + if all_choices: + return number_list + + number_list_prob = np.array(_create_special_list(len(number_list))) + if rng is None: + return np.random.choice(number_list, p = number_list_prob) + else: + return rng.choice(number_list, p = number_list_prob) + +def resize_mask(mask, latent, process_first_frame_only=True): + latent_size = latent.size() + batch_size, channels, num_frames, height, width = mask.shape + + if process_first_frame_only: + target_size = list(latent_size[2:]) + target_size[0] = 1 + first_frame_resized = F.interpolate( + mask[:, :, 0:1, :, :], + size=target_size, + mode='trilinear', + align_corners=False + ) + + target_size = list(latent_size[2:]) + target_size[0] = target_size[0] - 1 + if target_size[0] != 0: + remaining_frames_resized = F.interpolate( + mask[:, :, 1:, :, :], + size=target_size, + mode='trilinear', + align_corners=False + ) + resized_mask = torch.cat([first_frame_resized, remaining_frames_resized], dim=2) + else: + resized_mask = first_frame_resized + else: + target_size = list(latent_size[2:]) + resized_mask = F.interpolate( + mask, + size=target_size, + mode='trilinear', + align_corners=False + ) + return resized_mask + +# Will error if the minimal version of diffusers is not installed. Remove at your own risks. +check_min_version("0.18.0.dev0") + +logger = get_logger(__name__, log_level="INFO") + +def log_validation(vae, text_encoder, tokenizer, transformer3d, args, config, accelerator, weight_dtype, global_step): + try: + is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine' + if is_deepspeed: + origin_config = transformer3d.config + transformer3d.config = accelerator.unwrap_model(transformer3d).config + with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + logger.info("Running validation... ") + scheduler = FlowMatchEulerDiscreteScheduler( + **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs'])) + ) + if args.boundary_type == "full": + transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d + transformer3d_2 = None + else: + if args.boundary_type == "low": + transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d + + sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer') + transformer3d_2 = Wan2_2Transformer3DModel.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, sub_path), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + ).to(weight_dtype) + else: + sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer') + transformer3d_1 = Wan2_2Transformer3DModel.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, sub_path), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + ).to(weight_dtype) + + transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d + + if args.train_mode != "normal": + pipeline = Wan2_2I2VPipeline( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer3d_1, + transformer_2=transformer3d_2, + scheduler=scheduler, + ) + else: + pipeline = Wan2_2Pipeline( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer3d_1, + transformer_2=transformer3d_2, + scheduler=scheduler, + ) + pipeline = pipeline.to(accelerator.device) + + if args.seed is None: + generator = None + else: + rank_seed = args.seed + accelerator.process_index + generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed) + logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}") + + for i in range(len(args.validation_prompts)): + if args.train_mode != "normal": + start_image = Image.open(args.validation_paths[i]) + width, height = start_image.width, start_image.height + width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height) + + video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1 + input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width]) + sample = pipeline( + args.validation_prompts[i], + num_frames = video_length, + negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + height = height, + width = width, + generator = generator, + + video = input_video, + mask_video = input_video_mask, + num_inference_steps = 4, + guidance_scale = 1.0, + boundary = config['transformer_additional_kwargs'].get('boundary', 0.900) + ).videos + + os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True) + save_videos_grid( + sample, + os.path.join( + args.output_dir, + f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.mp4" + ) + ) + else: + sample = pipeline( + args.validation_prompts[i], + num_frames = args.video_sample_n_frames, + negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + height = args.video_sample_size, + width = args.video_sample_size, + generator = generator, + num_inference_steps = 4, + guidance_scale = 1.0, + boundary = config['transformer_additional_kwargs'].get('boundary', 0.900) + ).videos + os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True) + save_videos_grid( + sample, + os.path.join( + args.output_dir, + f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.mp4" + ) + ) + + del pipeline + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if is_deepspeed: + transformer3d.config = origin_config + except Exception as e: + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + print(f"Eval error on rank {accelerator.process_index} with info {e}") + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + + +def linear_decay(initial_value, final_value, total_steps, current_step): + if current_step >= total_steps: + return final_value + current_step = max(0, current_step) + step_size = (final_value - initial_value) / total_steps + current_value = initial_value + step_size * current_step + return current_value + +def generate_timestep_with_lognorm(low, high, shape, device="cpu", generator=None): + u = torch.normal(mean=0.0, std=1.0, size=shape, device=device, generator=generator) + t = 1 / (1 + torch.exp(-u)) * (high - low) + low + return torch.clip(t.to(torch.int32), low, high - 1) + +def parse_args(): + parser = argparse.ArgumentParser(description="Simple example of a training script.") + parser.add_argument( + "--input_perturbation", type=float, default=0, help="The scale of input perturbation. Recommended 0.1." + ) + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + required=True, + help="Path to pretrained model or model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--revision", + type=str, + default=None, + required=False, + help="Revision of pretrained model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--variant", + type=str, + default=None, + help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16", + ) + parser.add_argument( + "--train_data_dir", + type=str, + default=None, + help=( + "A folder containing the training data. " + ), + ) + parser.add_argument( + "--train_data_meta", + type=str, + default=None, + help=( + "A csv containing the training data. " + ), + ) + parser.add_argument( + "--max_train_samples", + type=int, + default=None, + help=( + "For debugging purposes or quicker training, truncate the number of training examples to this " + "value if set." + ), + ) + parser.add_argument( + "--validation_prompts", + type=str, + default=None, + nargs="+", + help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."), + ) + parser.add_argument( + "--validation_paths", + type=str, + default=None, + nargs="+", + help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."), + ) + parser.add_argument( + "--negative_prompt", + type=str, + default="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + help=("The negative prompt of cfg distill"), + ) + parser.add_argument( + "--output_dir", + type=str, + default="sd-model-finetuned", + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument( + "--cache_dir", + type=str, + default=None, + help="The directory where the downloaded models and datasets will be stored.", + ) + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + parser.add_argument( + "--random_flip", + action="store_true", + help="whether to randomly flip images horizontally", + ) + parser.add_argument( + "--use_came", + action="store_true", + help="whether to use came", + ) + parser.add_argument( + "--multi_stream", + action="store_true", + help="whether to use cuda multi-stream", + ) + parser.add_argument( + "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader." + ) + parser.add_argument( + "--vae_mini_batch", type=int, default=32, help="mini batch size for vae." + ) + parser.add_argument("--num_train_epochs", type=int, default=100) + parser.add_argument( + "--max_train_steps", + type=int, + default=None, + help="Total number of training steps to perform. If provided, overrides num_train_epochs.", + ) + parser.add_argument( + "--gradient_accumulation_steps", + type=int, + default=1, + help="Number of updates steps to accumulate before performing a backward/update pass.", + ) + parser.add_argument( + "--gradient_checkpointing", + action="store_true", + help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + parser.add_argument( + "--learning_rate_critic", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + parser.add_argument( + "--scale_lr", + action="store_true", + default=False, + help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", + ) + parser.add_argument( + "--lr_scheduler", + type=str, + default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument( + "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." + ) + parser.add_argument( + "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes." + ) + parser.add_argument( + "--allow_tf32", + action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument( + "--non_ema_revision", + type=str, + default=None, + required=False, + help=( + "Revision of pretrained non-ema model identifier. Must be a branch, tag or git identifier of the local or" + " remote repository specified with --pretrained_model_name_or_path." + ), + ) + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=0, + help=( + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ), + ) + parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.") + parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.") + parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.") + parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer") + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") + parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") + parser.add_argument( + "--prediction_type", + type=str, + default=None, + help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.", + ) + parser.add_argument( + "--hub_model_id", + type=str, + default=None, + help="The name of the repository to keep in sync with the local `output_dir`.", + ) + parser.add_argument( + "--logging_dir", + type=str, + default="logs", + help=( + "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." + ), + ) + parser.add_argument( + "--report_model_info", action="store_true", help="Whether or not to report more info about model (such as norm, grad)." + ) + parser.add_argument( + "--mixed_precision", + type=str, + default=None, + choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + parser.add_argument( + "--report_to", + type=str, + default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument( + "--checkpointing_steps", + type=int, + default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints are only suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument( + "--checkpoints_total_limit", + type=int, + default=None, + help=("Max number of checkpoints to store."), + ) + parser.add_argument( + "--resume_from_checkpoint", + type=str, + default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument("--noise_offset", type=float, default=0, help="The scale of noise offset.") + parser.add_argument( + "--validation_epochs", + type=int, + default=5, + help="Run validation every X epochs.", + ) + parser.add_argument( + "--validation_steps", + type=int, + default=2000, + help="Run validation every X steps.", + ) + parser.add_argument( + "--tracker_project_name", + type=str, + default="text2image-fine-tune", + help=( + "The `project_name` argument passed to Accelerator.init_trackers for" + " more information see https://huggingface.co/docs/accelerate/v0.17.0/en/package_reference/accelerator#accelerate.Accelerator" + ), + ) + + parser.add_argument( + "--snr_loss", action="store_true", help="Whether or not to use snr_loss." + ) + parser.add_argument( + "--uniform_sampling", action="store_true", help="Whether or not to use uniform_sampling." + ) + parser.add_argument( + "--enable_text_encoder_in_dataloader", action="store_true", help="Whether or not to use text encoder in dataloader." + ) + parser.add_argument( + "--enable_bucket", action="store_true", help="Whether enable bucket sample in datasets." + ) + parser.add_argument( + "--random_ratio_crop", action="store_true", help="Whether enable random ratio crop sample in datasets." + ) + parser.add_argument( + "--random_frame_crop", action="store_true", help="Whether enable random frame crop sample in datasets." + ) + parser.add_argument( + "--random_hw_adapt", action="store_true", help="Whether enable random adapt height and width in datasets." + ) + parser.add_argument( + "--training_with_video_token_length", action="store_true", help="The training stage of the model in training.", + ) + parser.add_argument( + "--auto_tile_batch_size", action="store_true", help="Whether to auto tile batch size.", + ) + parser.add_argument( + "--motion_sub_loss", action="store_true", help="Whether enable motion sub loss." + ) + parser.add_argument( + "--motion_sub_loss_ratio", type=float, default=0.25, help="The ratio of motion sub loss." + ) + parser.add_argument( + "--train_sampling_steps", + type=int, + default=1000, + help="Run train_sampling_steps.", + ) + parser.add_argument( + "--keep_all_node_same_token_length", + action="store_true", + help="Reference of the length token.", + ) + parser.add_argument( + "--token_sample_size", + type=int, + default=512, + help="Sample size of the token.", + ) + parser.add_argument( + "--video_sample_size", + type=int, + default=512, + help="Sample size of the video.", + ) + parser.add_argument( + "--image_sample_size", + type=int, + default=512, + help="Sample size of the image.", + ) + parser.add_argument( + "--fix_sample_size", + nargs=2, type=int, default=None, + help="Fix Sample size [height, width] when using bucket and collate_fn." + ) + parser.add_argument( + "--video_sample_stride", + type=int, + default=4, + help="Sample stride of the video.", + ) + parser.add_argument( + "--video_sample_n_frames", + type=int, + default=17, + help="Num frame of video.", + ) + parser.add_argument( + "--video_repeat", + type=int, + default=0, + help="Num of repeat video.", + ) + parser.add_argument( + "--config_path", + type=str, + default=None, + help=( + "The config of the model in training." + ), + ) + parser.add_argument( + "--transformer_path", + type=str, + default=None, + help=("If you want to load the weight from other transformers, input its path."), + ) + parser.add_argument( + "--vae_path", + type=str, + default=None, + help=("If you want to load the weight from other vaes, input its path."), + ) + + parser.add_argument( + '--trainable_modules', + nargs='+', + help='Enter a list of trainable modules' + ) + parser.add_argument( + '--trainable_modules_low_learning_rate', + nargs='+', + default=[], + help='Enter a list of trainable modules with lower learning rate' + ) + parser.add_argument( + '--tokenizer_max_length', + type=int, + default=512, + help='Max length of tokenizer' + ) + parser.add_argument( + "--use_deepspeed", action="store_true", help="Whether or not to use deepspeed." + ) + parser.add_argument( + "--use_fsdp", action="store_true", help="Whether or not to use fsdp." + ) + parser.add_argument( + "--low_vram", action="store_true", help="Whether enable low_vram mode." + ) + parser.add_argument( + "--boundary_type", + type=str, + default="low", + help=( + 'The training boundary type for dual-Transformer architecture. ' + 'Support `"low"` (train low-noise model), `"high"` (train high-noise model), ' + 'and `"full"` (single model training like TI2V-5B).' + ), + ) + parser.add_argument( + "--train_mode", + type=str, + default="normal", + help=( + 'The training mode. Support `"normal"` (T2V, default) ' + 'and `"inpaint"` (Inpaint).' + ), + ) + parser.add_argument( + "--gen_update_interval", + type=int, + default=5, + help="The ratio to update transformer3d.", + ) + parser.add_argument( + "--fake_guidance_scale", + type=float, + default=0.0, + help="The cfg scale for fake iscore.", + ) + parser.add_argument( + "--real_guidance_scale", + type=float, + default=6.0, + help="The cfg scale for real score.", + ) + parser.add_argument( + '--denoising_step_indices_list', + nargs='+', + default=[1000, 750, 500, 250], + help="The denoising step list.", + ) + parser.add_argument( + "--randomize_step_indices", + action="store_true", + help="whether to use randomize timesteps indices in training.", + ) + + args = parser.parse_args() + env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) + if env_local_rank != -1 and env_local_rank != args.local_rank: + args.local_rank = env_local_rank + + # default to using the same revision for the non-ema model if not specified + if args.non_ema_revision is None: + args.non_ema_revision = args.revision + + return args + + +def main(): + args = parse_args() + + if args.report_to == "wandb" and args.hub_token is not None: + raise ValueError( + "You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token." + " Please use `huggingface-cli login` to authenticate with the Hub." + ) + + if args.non_ema_revision is not None: + deprecate( + "non_ema_revision!=None", + "0.15.0", + message=( + "Downloading 'non_ema' weights from revision branches of the Hub is deprecated. Please make sure to" + " use `--variant=non_ema` instead." + ), + ) + logging_dir = os.path.join(args.output_dir, args.logging_dir) + + config = OmegaConf.load(args.config_path) + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + ) + accelerator_fake_score_transformer3d = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + ) + + deepspeed_plugin = accelerator.state.deepspeed_plugin if hasattr(accelerator.state, "deepspeed_plugin") else None + fsdp_plugin = accelerator.state.fsdp_plugin if hasattr(accelerator.state, "fsdp_plugin") else None + if deepspeed_plugin is not None: + zero_stage = int(deepspeed_plugin.zero_stage) + fsdp_stage = 0 + print(f"Using DeepSpeed Zero stage: {zero_stage}") + + args.use_deepspeed = True + if zero_stage == 3: + print(f"Auto set save_state to True because zero_stage == 3") + args.save_state = True + elif fsdp_plugin is not None: + from torch.distributed.fsdp import ShardingStrategy + zero_stage = 0 + if fsdp_plugin.sharding_strategy is ShardingStrategy.FULL_SHARD: + fsdp_stage = 3 + elif fsdp_plugin.sharding_strategy is None: # The fsdp_plugin.sharding_strategy is None in FSDP 2. + fsdp_stage = 3 + elif fsdp_plugin.sharding_strategy is ShardingStrategy.SHARD_GRAD_OP: + fsdp_stage = 2 + else: + fsdp_stage = 0 + print(f"Using FSDP stage: {fsdp_stage}") + + args.use_fsdp = True + if fsdp_stage == 3: + print(f"Auto set save_state to True because fsdp_stage == 3") + args.save_state = True + else: + zero_stage = 0 + fsdp_stage = 0 + print("DeepSpeed is not enabled.") + + if accelerator.is_main_process: + writer = SummaryWriter(log_dir=logging_dir) + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + datasets.utils.logging.set_verbosity_warning() + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + datasets.utils.logging.set_verbosity_error() + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed) + rng = np.random.default_rng(np.random.PCG64(args.seed + accelerator.process_index)) + torch_rng = torch.Generator(accelerator.device).manual_seed(args.seed + accelerator.process_index) + else: + rng = None + torch_rng = None + index_rng = np.random.default_rng(np.random.PCG64(43)) + print(f"Init rng with seed {args.seed + accelerator.process_index}. Process_index is {accelerator.process_index}") + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + # For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora transformer3d) to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + args.mixed_precision = accelerator.mixed_precision + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + args.mixed_precision = accelerator.mixed_precision + + args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list] + # Load scheduler, tokenizer and models. + noise_scheduler = FlowMatchEulerDiscreteScheduler( + **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs'])) + ) + + # Get Tokenizer + tokenizer = AutoTokenizer.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')), + ) + + def deepspeed_zero_init_disabled_context_manager(): + """ + returns either a context list that includes one that will disable zero.Init or an empty context list + """ + deepspeed_plugin = AcceleratorState().deepspeed_plugin if accelerate.state.is_initialized() else None + if deepspeed_plugin is None: + return [] + + return [deepspeed_plugin.zero3_init_context_manager(enable=False)] + + # Currently Accelerate doesn't know how to handle multiple models under Deepspeed ZeRO stage 3. + # For this to work properly all models must be run through `accelerate.prepare`. But accelerate + # will try to assign the same optimizer with the same weights to all models during + # `deepspeed.initialize`, which of course doesn't work. + # + # For now the following workaround will partially support Deepspeed ZeRO-3, by excluding the 2 + # frozen models from being partitioned during `zero.Init` which gets called during + # `from_pretrained` So CLIPTextModel and AutoencoderKL will not enjoy the parameter sharding + # across multiple gpus and only UNet2DConditionModel will get ZeRO sharded. + with ContextManagers(deepspeed_zero_init_disabled_context_manager()): + # Get Text encoder + text_encoder = WanT5EncoderModel.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')), + additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=weight_dtype, + ) + text_encoder = text_encoder.eval() + # Get Vae + Chosen_AutoencoderKL = { + "AutoencoderKLWan": AutoencoderKLWan, + "AutoencoderKLWan3_8": AutoencoderKLWan3_8 + }[config['vae_kwargs'].get('vae_type', 'AutoencoderKLWan')] + vae = Chosen_AutoencoderKL.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['vae_kwargs'].get('vae_subpath', 'vae')), + additional_kwargs=OmegaConf.to_container(config['vae_kwargs']), + ) + vae.eval() + + # Get Transformer + if args.boundary_type == "low" or args.boundary_type == "full": + sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer') + else: + sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer') + + # Get Transformer + generator_transformer3d = Wan2_2Transformer3DModel.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, sub_path), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + ).to(weight_dtype) + real_score_transformer3d = Wan2_2Transformer3DModel.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, sub_path), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + ).to(weight_dtype) + fake_score_transformer3d = Wan2_2Transformer3DModel.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, sub_path), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + ).to(weight_dtype) + + # Freeze vae and text_encoder and set generator_transformer3d to trainable + vae.requires_grad_(False) + text_encoder.requires_grad_(False) + generator_transformer3d.requires_grad_(False) + real_score_transformer3d.requires_grad_(False) + fake_score_transformer3d.requires_grad_(False) + + if args.transformer_path is not None: + print(f"From checkpoint: {args.transformer_path}") + if args.transformer_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(args.transformer_path) + else: + state_dict = torch.load(args.transformer_path, map_location="cpu") + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict + + m, u = generator_transformer3d.load_state_dict(state_dict, strict=False) + print(f"generator missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + m, u = fake_score_transformer3d.load_state_dict(state_dict, strict=False) + print(f"fake_score missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + m, u = real_score_transformer3d.load_state_dict(state_dict, strict=False) + print(f"real_score missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + + if args.vae_path is not None: + print(f"From checkpoint: {args.vae_path}") + if args.vae_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(args.vae_path) + else: + state_dict = torch.load(args.vae_path, map_location="cpu") + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict + + m, u = vae.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + + # A good trainable modules is showed below now. + # For 3D Patch: trainable_modules = ['ff.net', 'pos_embed', 'attn2', 'proj_out', 'timepositionalencoding', 'h_position', 'w_position'] + # For 2D Patch: trainable_modules = ['ff.net', 'attn2', 'timepositionalencoding', 'h_position', 'w_position'] + generator_transformer3d.train() + fake_score_transformer3d.train() + if accelerator.is_main_process: + accelerator.print( + f"Trainable modules '{args.trainable_modules}'." + ) + for name, param in generator_transformer3d.named_parameters(): + for trainable_module_name in args.trainable_modules + args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + param.requires_grad = True + break + for name, param in fake_score_transformer3d.named_parameters(): + for trainable_module_name in args.trainable_modules + args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + param.requires_grad = True + break + + # `accelerate` 0.16.0 will have better support for customized saving + if version.parse(accelerate.__version__) >= version.parse("0.16.0"): + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + if fsdp_stage != 0 or zero_stage == 3: + def save_model_hook(models, weights, output_dir): + accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True) + if accelerator.is_main_process: + from safetensors.torch import save_file + + safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors") + accelerate_state_dict = {k: v.to(dtype=weight_dtype) for k, v in accelerate_state_dict.items()} + save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"}) + + with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file: + pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file) + + def load_model_hook(models, input_dir): + pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + loaded_number, _ = pickle.load(file) + batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0) + print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.") + else: + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + if accelerator.is_main_process: + models[0].save_pretrained(os.path.join(output_dir, "transformer")) + if not args.use_deepspeed: + weights.pop() + + with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file: + pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file) + + def load_model_hook(models, input_dir): + for i in range(len(models)): + # pop models so that they are not loaded again + model = models.pop() + + # load diffusers style into model + load_model = Wan2_2Transformer3DModel.from_pretrained( + input_dir, subfolder="transformer" + ) + model.register_to_config(**load_model.config) + + model.load_state_dict(load_model.state_dict()) + del load_model + + pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + loaded_number, _ = pickle.load(file) + batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0) + print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.") + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook) + accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook) + + if args.gradient_checkpointing: + generator_transformer3d.enable_gradient_checkpointing() + fake_score_transformer3d.enable_gradient_checkpointing() + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + # Initialize the optimizer + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`" + ) + + optimizer_cls = bnb.optim.AdamW8bit + elif args.use_came: + try: + from came_pytorch import CAME + except Exception: + raise ImportError( + "Please install came_pytorch to use CAME. You can do so by running `pip install came_pytorch`" + ) + + optimizer_cls = CAME + else: + optimizer_cls = torch.optim.AdamW + + trainable_params = list(filter(lambda p: p.requires_grad, generator_transformer3d.parameters())) + trainable_params_optim = [ + {'params': [], 'lr': args.learning_rate}, + {'params': [], 'lr': args.learning_rate / 2}, + ] + in_already = [] + for name, param in generator_transformer3d.named_parameters(): + high_lr_flag = False + if name in in_already: + continue + for trainable_module_name in args.trainable_modules: + if trainable_module_name in name: + in_already.append(name) + high_lr_flag = True + trainable_params_optim[0]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate}") + break + if high_lr_flag: + continue + for trainable_module_name in args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + in_already.append(name) + trainable_params_optim[1]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate / 2}") + break + + fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_transformer3d.parameters())) + fake_trainable_params_optim = [ + {'params': [], 'lr': args.learning_rate_critic}, + {'params': [], 'lr': args.learning_rate_critic / 2}, + ] + in_already = [] + for name, param in fake_score_transformer3d.named_parameters(): + high_lr_flag = False + if name in in_already: + continue + for trainable_module_name in args.trainable_modules: + if trainable_module_name in name: + in_already.append(name) + high_lr_flag = True + fake_trainable_params_optim[0]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate_critic}") + break + if high_lr_flag: + continue + for trainable_module_name in args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + in_already.append(name) + fake_trainable_params_optim[1]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate_critic / 2}") + break + + if args.use_came: + optimizer = optimizer_cls( + trainable_params_optim, + lr=args.learning_rate, + # weight_decay=args.adam_weight_decay, + betas=(0.9, 0.999, 0.9999), + eps=(1e-30, 1e-16) + ) + critic_optimizer = optimizer_cls( + fake_trainable_params_optim, + lr=args.learning_rate_critic, + # weight_decay=args.adam_weight_decay, + betas=(0.9, 0.999, 0.9999), + eps=(1e-30, 1e-16) + ) + else: + optimizer = optimizer_cls( + trainable_params_optim, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + critic_optimizer = optimizer_cls( + fake_trainable_params_optim, + lr=args.learning_rate_critic, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + # Get the training dataset + sample_n_frames_bucket_interval = vae.config.temporal_compression_ratio + spatial_compression_ratio = vae.config.spatial_compression_ratio + + if args.fix_sample_size is not None and args.enable_bucket: + args.video_sample_size = max(max(args.fix_sample_size), args.video_sample_size) + args.image_sample_size = max(max(args.fix_sample_size), args.image_sample_size) + args.training_with_video_token_length = False + args.random_hw_adapt = False + + # Get the dataset + if args.train_mode != "normal": + train_dataset = ImageVideoDataset( + args.train_data_meta, args.train_data_dir, + video_sample_size=args.video_sample_size, video_sample_stride=args.video_sample_stride, video_sample_n_frames=args.video_sample_n_frames, + video_repeat=args.video_repeat, + image_sample_size=args.image_sample_size, + enable_bucket=args.enable_bucket, enable_inpaint=True if args.train_mode != "normal" else False, + ) + else: + train_dataset = TextDataset( + args.train_data_meta + ) + + def get_length_to_frame_num(token_length): + if args.image_sample_size > args.video_sample_size: + sample_sizes = list(range(args.video_sample_size, args.image_sample_size + 1, 128)) + + if sample_sizes[-1] != args.image_sample_size: + sample_sizes.append(args.image_sample_size) + else: + sample_sizes = [args.image_sample_size] + + length_to_frame_num = { + sample_size: min(token_length / sample_size / sample_size, args.video_sample_n_frames) // sample_n_frames_bucket_interval * sample_n_frames_bucket_interval + 1 for sample_size in sample_sizes + } + + return length_to_frame_num + + if args.enable_bucket and args.train_mode != "normal": + aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + batch_sampler_generator = torch.Generator().manual_seed(args.seed) + batch_sampler = AspectRatioBatchImageVideoSampler( + sampler=RandomSampler(train_dataset, generator=batch_sampler_generator), dataset=train_dataset.dataset, + batch_size=args.train_batch_size, train_folder = args.train_data_dir, drop_last=True, + aspect_ratios=aspect_ratio_sample_size, + ) + + def collate_fn(examples): + # Get token length + target_token_length = args.video_sample_n_frames * args.token_sample_size * args.token_sample_size + length_to_frame_num = get_length_to_frame_num(target_token_length) + + # Create new output + new_examples = {} + new_examples["target_token_length"] = target_token_length + new_examples["pixel_values"] = [] + new_examples["text"] = [] + # Used in Inpaint mode + if args.train_mode != "normal": + new_examples["mask_pixel_values"] = [] + new_examples["mask"] = [] + new_examples["clip_pixel_values"] = [] + + # Get downsample ratio in image and videos + pixel_value = examples[0]["pixel_values"] + data_type = examples[0]["data_type"] + f, h, w, c = np.shape(pixel_value) + if data_type == 'image': + random_downsample_ratio = 1 if not args.random_hw_adapt else get_random_downsample_ratio(args.image_sample_size, image_ratio=[args.image_sample_size / args.video_sample_size], rng=rng) + + aspect_ratio_sample_size = {key : [x / 512 * args.image_sample_size / random_downsample_ratio for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + aspect_ratio_random_crop_sample_size = {key : [x / 512 * args.image_sample_size / random_downsample_ratio for x in ASPECT_RATIO_RANDOM_CROP_512[key]] for key in ASPECT_RATIO_RANDOM_CROP_512.keys()} + + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + else: + if args.random_hw_adapt: + if args.training_with_video_token_length: + local_min_size = np.min(np.array([np.mean(np.array([np.shape(example["pixel_values"])[1], np.shape(example["pixel_values"])[2]])) for example in examples])) + # The video will be resized to a lower resolution than its own. + choice_list = [length for length in list(length_to_frame_num.keys()) if length < local_min_size * 1.25] + if len(choice_list) == 0: + choice_list = list(length_to_frame_num.keys()) + if rng is None: + local_video_sample_size = np.random.choice(choice_list) + else: + local_video_sample_size = rng.choice(choice_list) + batch_video_length = length_to_frame_num[local_video_sample_size] + random_downsample_ratio = args.video_sample_size / local_video_sample_size + else: + random_downsample_ratio = get_random_downsample_ratio( + args.video_sample_size, rng=rng) + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + else: + random_downsample_ratio = 1 + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + + aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size / random_downsample_ratio for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + aspect_ratio_random_crop_sample_size = {key : [x / 512 * args.video_sample_size / random_downsample_ratio for x in ASPECT_RATIO_RANDOM_CROP_512[key]] for key in ASPECT_RATIO_RANDOM_CROP_512.keys()} + + if args.fix_sample_size is not None: + fix_sample_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in args.fix_sample_size] + elif args.random_ratio_crop: + if rng is None: + random_sample_size = aspect_ratio_random_crop_sample_size[ + np.random.choice(list(aspect_ratio_random_crop_sample_size.keys()), p = ASPECT_RATIO_RANDOM_CROP_PROB) + ] + else: + random_sample_size = aspect_ratio_random_crop_sample_size[ + rng.choice(list(aspect_ratio_random_crop_sample_size.keys()), p = ASPECT_RATIO_RANDOM_CROP_PROB) + ] + random_sample_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in random_sample_size] + else: + closest_size, closest_ratio = get_closest_ratio(h, w, ratios=aspect_ratio_sample_size) + closest_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in closest_size] + + min_example_length = min( + [example["pixel_values"].shape[0] for example in examples] + ) + batch_video_length = int(min(batch_video_length, min_example_length)) + + # Magvae needs the number of frames to be 4n + 1. + batch_video_length = (batch_video_length - 1) // sample_n_frames_bucket_interval * sample_n_frames_bucket_interval + 1 + + if batch_video_length <= 0: + batch_video_length = 1 + + for example in examples: + if args.fix_sample_size is not None: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + fix_sample_size = list(map(lambda x: int(x), fix_sample_size)) + transform = transforms.Compose([ + transforms.Resize(fix_sample_size, interpolation=transforms.InterpolationMode.BILINEAR), # Image.BICUBIC + transforms.CenterCrop(fix_sample_size), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + elif args.random_ratio_crop: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + b, c, h, w = pixel_values.size() + th, tw = random_sample_size + if th / tw > h / w: + nh = int(th) + nw = int(w / h * nh) + else: + nw = int(tw) + nh = int(h / w * nw) + + transform = transforms.Compose([ + transforms.Resize([nh, nw]), + transforms.CenterCrop([int(x) for x in random_sample_size]), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + else: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + closest_size = list(map(lambda x: int(x), closest_size)) + if closest_size[0] / h > closest_size[1] / w: + resize_size = closest_size[0], int(w * closest_size[0] / h) + else: + resize_size = int(h * closest_size[1] / w), closest_size[1] + + transform = transforms.Compose([ + transforms.Resize(resize_size, interpolation=transforms.InterpolationMode.BILINEAR), # Image.BICUBIC + transforms.CenterCrop(closest_size), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + + new_examples["pixel_values"].append(transform(pixel_values)[:batch_video_length]) + new_examples["text"].append(example["text"]) + + if args.train_mode != "normal": + mask = get_random_mask(new_examples["pixel_values"][-1].size()) + mask_pixel_values = new_examples["pixel_values"][-1] * (1 - mask) + # Wan 2.1 use 0 for masked pixels + # + torch.ones_like(new_examples["pixel_values"][-1]) * -1 * mask + new_examples["mask_pixel_values"].append(mask_pixel_values) + new_examples["mask"].append(mask) + + clip_pixel_values = new_examples["pixel_values"][-1][0].permute(1, 2, 0).contiguous() + clip_pixel_values = (clip_pixel_values * 0.5 + 0.5) * 255 + new_examples["clip_pixel_values"].append(clip_pixel_values) + + # Limit the number of frames to the same + new_examples["pixel_values"] = torch.stack([example for example in new_examples["pixel_values"]]) + if args.train_mode != "normal": + new_examples["mask_pixel_values"] = torch.stack([example for example in new_examples["mask_pixel_values"]]) + new_examples["mask"] = torch.stack([example for example in new_examples["mask"]]) + new_examples["clip_pixel_values"] = torch.stack([example for example in new_examples["clip_pixel_values"]]) + + # Encode prompts when enable_text_encoder_in_dataloader=True + if args.enable_text_encoder_in_dataloader: + prompt_ids = tokenizer( + new_examples['text'], + max_length=args.tokenizer_max_length, + padding="max_length", + add_special_tokens=True, + truncation=True, + return_tensors="pt" + ) + text_input_ids = prompt_ids.input_ids + prompt_attention_mask = prompt_ids.attention_mask + + seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long() + prompt_embeds = text_encoder(text_input_ids.to("cpu"), attention_mask=prompt_attention_mask.to("cpu"))[0] + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + + new_examples['encoder_attention_mask'] = prompt_ids.attention_mask + new_examples['encoder_hidden_states'] = prompt_embeds + + neg_txt = [ + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" for text in new_examples['text'] + ] + neg_prompt_ids = tokenizer( + neg_txt, + max_length=args.tokenizer_max_length, + padding="max_length", + add_special_tokens=True, + truncation=True, + return_tensors="pt" + ) + neg_text_input_ids = neg_prompt_ids.input_ids + neg_prompt_attention_mask = neg_prompt_ids.attention_mask + + neg_seq_lens = neg_prompt_attention_mask.gt(0).sum(dim=1).long() + neg_prompt_embeds = text_encoder(neg_text_input_ids.to("cpu"), attention_mask=neg_prompt_attention_mask.to("cpu"))[0] + neg_prompt_embeds = [u[:v] for u, v in zip(neg_prompt_embeds, neg_seq_lens)] + + new_examples['neg_encoder_attention_mask'] = neg_prompt_ids.attention_mask + new_examples['neg_encoder_hidden_states'] = neg_prompt_embeds + + return new_examples + + # DataLoaders creation: + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + collate_fn=collate_fn, + persistent_workers=True if args.dataloader_num_workers != 0 else False, + num_workers=args.dataloader_num_workers, + ) + elif args.train_mode == "normal": + def collate_fn(examples): + new_examples = {} + new_examples["text"] = [] + for example in examples: + new_examples["text"].append(example["text"]) + + # Encode prompts when enable_text_encoder_in_dataloader=True + if args.enable_text_encoder_in_dataloader: + prompt_ids = tokenizer( + new_examples['text'], + max_length=args.tokenizer_max_length, + padding="max_length", + add_special_tokens=True, + truncation=True, + return_tensors="pt" + ) + text_input_ids = prompt_ids.input_ids + prompt_attention_mask = prompt_ids.attention_mask + + seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long() + prompt_embeds = text_encoder(text_input_ids.to("cpu"), attention_mask=prompt_attention_mask.to("cpu"))[0] + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + + new_examples['encoder_attention_mask'] = prompt_ids.attention_mask + new_examples['encoder_hidden_states'] = prompt_embeds + + neg_txt = [ + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" for text in new_examples['text'] + ] + neg_prompt_ids = tokenizer( + neg_txt, + max_length=args.tokenizer_max_length, + padding="max_length", + add_special_tokens=True, + truncation=True, + return_tensors="pt" + ) + neg_text_input_ids = neg_prompt_ids.input_ids + neg_prompt_attention_mask = neg_prompt_ids.attention_mask + + neg_seq_lens = neg_prompt_attention_mask.gt(0).sum(dim=1).long() + neg_prompt_embeds = text_encoder(neg_text_input_ids.to("cpu"), attention_mask=neg_prompt_attention_mask.to("cpu"))[0] + neg_prompt_embeds = [u[:v] for u, v in zip(neg_prompt_embeds, neg_seq_lens)] + + new_examples['neg_encoder_attention_mask'] = neg_prompt_ids.attention_mask + new_examples['neg_encoder_hidden_states'] = neg_prompt_embeds + + return new_examples + + batch_sampler_generator = torch.Generator().manual_seed(args.seed) + batch_sampler = BatchSampler(RandomSampler(train_dataset, generator=batch_sampler_generator), batch_size=args.train_batch_size, drop_last=True) + + # DataLoaders creation: + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + collate_fn=collate_fn, + persistent_workers=True if args.dataloader_num_workers != 0 else False, + num_workers=args.dataloader_num_workers, + ) + else: + # DataLoaders creation: + batch_sampler_generator = torch.Generator().manual_seed(args.seed) + batch_sampler = ImageVideoSampler(RandomSampler(train_dataset, generator=batch_sampler_generator), train_dataset, args.train_batch_size) + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + persistent_workers=True if args.dataloader_num_workers != 0 else False, + num_workers=args.dataloader_num_workers, + ) + + # Scheduler and math around the number of training steps. + overrode_max_train_steps = False + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + overrode_max_train_steps = True + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, + num_training_steps=args.max_train_steps * accelerator.num_processes, + ) + fake_score_lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=critic_optimizer, + num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, + num_training_steps=args.max_train_steps * accelerator.num_processes, + ) + + # Prepare everything with our `accelerator`. + generator_transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + generator_transformer3d, optimizer, train_dataloader, lr_scheduler + ) + fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler= accelerator_fake_score_transformer3d.prepare( + fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler + ) + if fsdp_stage != 0 or zero_stage != 0: + from functools import partial + + from videox_fun.dist import set_multi_gpus_devices, shard_model + shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype) + real_score_transformer3d = shard_fn(real_score_transformer3d) + if fsdp_stage != 0 or zero_stage != 0: + from functools import partial + + from videox_fun.dist import set_multi_gpus_devices, shard_model + shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype) + text_encoder = shard_fn(text_encoder) + + # Move text_encode and vae to gpu and cast to weight_dtype + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + real_score_transformer3d.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu") + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if overrode_max_train_steps: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if accelerator.is_main_process: + tracker_config = dict(vars(args)) + keys_to_pop = [k for k, v in tracker_config.items() if isinstance(v, list)] + for k in keys_to_pop: + tracker_config.pop(k) + print(f"Removed tracker_config['{k}']") + accelerator.init_trackers(args.tracker_project_name, tracker_config) + + # Function for unwrapping if model was compiled with `torch.compile`. + def unwrap_model(model): + model = accelerator.unwrap_model(model) + model = model._orig_mod if is_compiled_module(model) else model + return model + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + + logger.info("***** Running training *****") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + global_step = 0 + first_epoch = 0 + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the most recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + + pkl_path = os.path.join(os.path.join(args.output_dir, path), "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + _, first_epoch = pickle.load(file) + else: + first_epoch = global_step // num_update_steps_per_epoch + print(f"Load pkl from {pkl_path}. Get first_epoch = {first_epoch}.") + + accelerator.print(f"Resuming from checkpoint {path}") + fake_score_path = os.path.join(path, "fake_score") + accelerator.load_state(os.path.join(args.output_dir, path)) + accelerator_fake_score_transformer3d.load_state(os.path.join(args.output_dir, fake_score_path)) + else: + initial_global_step = 0 + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + + if args.multi_stream and args.train_mode != "normal": + # create extra cuda streams to speedup inpaint vae computation + vae_stream_1 = torch.cuda.Stream() + vae_stream_2 = torch.cuda.Stream() + else: + vae_stream_1 = None + vae_stream_2 = None + + # Calculate the index we need + boundary = config['transformer_additional_kwargs'].get('boundary', 0.900) + split_timesteps = args.train_sampling_steps * boundary + differences = torch.abs(noise_scheduler.timesteps - split_timesteps) + closest_index = torch.argmin(differences).item() + if args.boundary_type == "high" or args.boundary_type == "low": + print(f"The boundary is {boundary} and the boundary_type is {args.boundary_type}. The closest_index we calculate is {closest_index}") + if args.boundary_type == "high": + start_num_idx = 0 + train_sampling_steps = closest_index + elif args.boundary_type == "low": + start_num_idx = closest_index + train_sampling_steps = args.train_sampling_steps - closest_index + else: + start_num_idx = 0 + train_sampling_steps = args.train_sampling_steps + + idx_sampling = DiscreteSampling(train_sampling_steps, start_num_idx=start_num_idx, uniform_sampling=args.uniform_sampling) + + def randomize_denoising_step_indices( + denoising_step_indices_list, + train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=0.3, + ): + indices = list(denoising_step_indices_list) + n = len(indices) + + if n <= 2: + low = indices[1] + high = indices[0] - 1 + random_tail = torch.randint(low, high + 1, (1,)).item() + + result = torch.tensor([indices[0], random_tail]) + else: + result = [0] * n + result[0] = indices[0] + result[-1] = indices[-1] + + for i in range(1, n - 1): + gap_upper = indices[i - 1] - indices[i] + gap_lower = indices[i] - indices[i + 1] + + max_jitter = int(min(gap_upper, gap_lower) * jitter_ratio) + + if max_jitter > 0: + jitter = torch.randint( + -max_jitter, max_jitter + 1, (1,) + ).item() + else: + jitter = 0 + + result[i] = indices[i] + jitter + + for i in range(1, n): + if result[i] >= result[i - 1]: + result[i] = result[i - 1] - 1 + + result = [max(1, min(train_sampling_steps, x)) for x in result] + result = torch.tensor(result) + + if dist.is_initialized(): + result = result.to(accelerator.device) + dist.broadcast(result, src=0) + result = result.cpu() + return result + + for epoch in range(first_epoch, args.num_train_epochs): + train_dmd_loss = 0.0 + train_denoising_loss = 0.0 + batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch) + for step, batch in enumerate(train_dataloader): + # Data batch sanity check + if args.train_mode != "normal" and epoch == first_epoch and step == 0: + pixel_values, texts = batch['pixel_values'].cpu(), batch['text'] + pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") + os.makedirs(os.path.join(args.output_dir, "sanity_check"), exist_ok=True) + for idx, (pixel_value, text) in enumerate(zip(pixel_values, texts)): + pixel_value = pixel_value[None, ...] + gif_name = '-'.join(text.replace('/', '').split()[:10]) if not text == '' else f'{global_step}-{idx}' + save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/{gif_name[:10]}.mp4", rescale=True) + + clip_pixel_values, mask_pixel_values, texts = batch['clip_pixel_values'].cpu(), batch['mask_pixel_values'].cpu(), batch['text'] + mask_pixel_values = rearrange(mask_pixel_values, "b f c h w -> b c f h w") + for idx, (clip_pixel_value, pixel_value, text) in enumerate(zip(clip_pixel_values, mask_pixel_values, texts)): + pixel_value = pixel_value[None, ...] + Image.fromarray(np.uint8(clip_pixel_value)).save(f"{args.output_dir}/sanity_check/clip_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.png") + save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/mask_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.mp4", rescale=True) + + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + if args.train_mode != "normal": + # Convert images to latent space + pixel_values = batch["pixel_values"].to(weight_dtype) + + # Increase the batch size when the length of the latent sequence of the current sample is small + if args.auto_tile_batch_size and args.training_with_video_token_length and zero_stage != 3: + if args.video_sample_n_frames * args.token_sample_size * args.token_sample_size // 16 >= pixel_values.size()[1] * pixel_values.size()[3] * pixel_values.size()[4]: + pixel_values = torch.tile(pixel_values, (4, 1, 1, 1, 1)) + if args.enable_text_encoder_in_dataloader: + batch['encoder_hidden_states'] = torch.tile(batch['encoder_hidden_states'], (4, 1, 1)) + batch['encoder_attention_mask'] = torch.tile(batch['encoder_attention_mask'], (4, 1)) + batch['neg_encoder_hidden_states'] = torch.tile(batch['neg_encoder_hidden_states'], (4, 1, 1)) + batch['neg_encoder_attention_mask'] = torch.tile(batch['neg_encoder_attention_mask'], (4, 1)) + else: + batch['text'] = batch['text'] * 4 + elif args.video_sample_n_frames * args.token_sample_size * args.token_sample_size // 4 >= pixel_values.size()[1] * pixel_values.size()[3] * pixel_values.size()[4]: + pixel_values = torch.tile(pixel_values, (2, 1, 1, 1, 1)) + if args.enable_text_encoder_in_dataloader: + batch['encoder_hidden_states'] = torch.tile(batch['encoder_hidden_states'], (2, 1, 1)) + batch['encoder_attention_mask'] = torch.tile(batch['encoder_attention_mask'], (2, 1)) + batch['neg_encoder_hidden_states'] = torch.tile(batch['neg_encoder_hidden_states'], (2, 1, 1)) + batch['neg_encoder_attention_mask'] = torch.tile(batch['neg_encoder_attention_mask'], (2, 1)) + else: + batch['text'] = batch['text'] * 2 + + clip_pixel_values = batch["clip_pixel_values"].to(weight_dtype) + mask_pixel_values = batch["mask_pixel_values"].to(weight_dtype) + mask = batch["mask"].to(weight_dtype) + # Increase the batch size when the length of the latent sequence of the current sample is small + if args.auto_tile_batch_size and args.training_with_video_token_length and zero_stage != 3: + if args.video_sample_n_frames * args.token_sample_size * args.token_sample_size // 16 >= pixel_values.size()[1] * pixel_values.size()[3] * pixel_values.size()[4]: + clip_pixel_values = torch.tile(clip_pixel_values, (4, 1, 1, 1)) + mask_pixel_values = torch.tile(mask_pixel_values, (4, 1, 1, 1, 1)) + mask = torch.tile(mask, (4, 1, 1, 1, 1)) + elif args.video_sample_n_frames * args.token_sample_size * args.token_sample_size // 4 >= pixel_values.size()[1] * pixel_values.size()[3] * pixel_values.size()[4]: + clip_pixel_values = torch.tile(clip_pixel_values, (2, 1, 1, 1)) + mask_pixel_values = torch.tile(mask_pixel_values, (2, 1, 1, 1, 1)) + mask = torch.tile(mask, (2, 1, 1, 1, 1)) + + if args.random_frame_crop: + def _create_special_list(length): + if length == 1: + return [1.0] + if length >= 2: + last_element = 0.90 + remaining_sum = 1.0 - last_element + other_elements_value = remaining_sum / (length - 1) + special_list = [other_elements_value] * (length - 1) + [last_element] + return special_list + select_frames = [_tmp for _tmp in list(range(sample_n_frames_bucket_interval + 1, args.video_sample_n_frames + sample_n_frames_bucket_interval, sample_n_frames_bucket_interval))] + select_frames_prob = np.array(_create_special_list(len(select_frames))) + + if len(select_frames) != 0: + if rng is None: + temp_n_frames = np.random.choice(select_frames, p = select_frames_prob) + else: + temp_n_frames = rng.choice(select_frames, p = select_frames_prob) + else: + temp_n_frames = 1 + + # Magvae needs the number of frames to be 4n + 1. + temp_n_frames = (temp_n_frames - 1) // sample_n_frames_bucket_interval + 1 + + pixel_values = pixel_values[:, :temp_n_frames, :, :] + mask_pixel_values = mask_pixel_values[:, :temp_n_frames, :, :] + mask = mask[:, :temp_n_frames, :, :] + + # Keep all node same token length to accelerate the traning when resolution grows. + if args.keep_all_node_same_token_length: + if args.token_sample_size > 256: + numbers_list = list(range(256, args.token_sample_size + 1, 128)) + + if numbers_list[-1] != args.token_sample_size: + numbers_list.append(args.token_sample_size) + else: + numbers_list = [256] + numbers_list = [_number * _number * args.video_sample_n_frames for _number in numbers_list] + + actual_token_length = index_rng.choice(numbers_list) + actual_video_length = (min( + actual_token_length / pixel_values.size()[-1] / pixel_values.size()[-2], args.video_sample_n_frames + ) - 1) // sample_n_frames_bucket_interval * sample_n_frames_bucket_interval + 1 + actual_video_length = int(max(actual_video_length, 1)) + + # Magvae needs the number of frames to be 4n + 1. + actual_video_length = (actual_video_length - 1) // sample_n_frames_bucket_interval + 1 + + pixel_values = pixel_values[:, :actual_video_length, :, :] + mask_pixel_values = mask_pixel_values[:, :actual_video_length, :, :] + mask = mask[:, :actual_video_length, :, :] + + if args.low_vram: + torch.cuda.empty_cache() + vae.to(accelerator.device) + real_score_transformer3d = real_score_transformer3d.to("cpu") + if not args.enable_text_encoder_in_dataloader: + text_encoder.to("cpu") + + with torch.no_grad(): + # This way is quicker when batch grows up + def _batch_encode_vae(pixel_values): + pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") + bs = args.vae_mini_batch + new_pixel_values = [] + for i in range(0, pixel_values.shape[0], bs): + pixel_values_bs = pixel_values[i : i + bs] + pixel_values_bs = vae.encode(pixel_values_bs)[0] + pixel_values_bs = pixel_values_bs.sample() + new_pixel_values.append(pixel_values_bs) + return torch.cat(new_pixel_values, dim = 0) + + # Encode inpaint latents. + mask_latents = _batch_encode_vae(mask_pixel_values) + if vae_stream_2 is not None: + torch.cuda.current_stream().wait_stream(vae_stream_2) + + # Compute t2v_flag for inpaint mode + if args.train_mode == "inpaint": + t2v_flag = [(_mask == 1).all() for _mask in mask] + new_t2v_flag = [] + for _mask in t2v_flag: + if _mask and np.random.rand() < 0.90: + new_t2v_flag.append(0) + else: + new_t2v_flag.append(1) + t2v_flag = torch.from_numpy(np.array(new_t2v_flag)).to(accelerator.device, dtype=weight_dtype) + + mask = rearrange(mask, "b f c h w -> b c f h w") + mask = torch.concat( + [ + torch.repeat_interleave(mask[:, :, 0:1], repeats=4, dim=2), + mask[:, :, 1:] + ], dim=2 + ) + mask = mask.view(mask.shape[0], mask.shape[2] // 4, 4, mask.shape[3], mask.shape[4]) + mask = mask.transpose(1, 2) + + mask_conditions = F.interpolate(mask[:, :1], size=mask_latents.size()[-3:], mode='trilinear', align_corners=True).to(accelerator.device, weight_dtype) + mask = resize_mask(1 - mask, mask_latents) + inpaint_latents = torch.concat([mask, mask_latents], dim=1) + inpaint_latents = t2v_flag[:, None, None, None, None] * inpaint_latents + + target_shape = mask_latents.size() + else: + text = batch['text'] + if args.fix_sample_size is not None: + local_sample_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in args.fix_sample_size] + num_frames = args.video_sample_n_frames + else: + if args.random_hw_adapt and args.training_with_video_token_length: + # Get token length + target_token_length = args.video_sample_n_frames * args.token_sample_size * args.token_sample_size + length_to_frame_num = get_length_to_frame_num(target_token_length) + + if rng is None: + local_length = np.random.choice(list(length_to_frame_num.keys())) + else: + local_length = rng.choice(list(length_to_frame_num.keys())) + num_frames = length_to_frame_num[local_length] + + aspect_ratio_sample_size = {key : [x / 512 * local_length for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + if rng is None: + aspect_ratio_key = np.random.choice(list(aspect_ratio_sample_size.keys())) + else: + aspect_ratio_key = rng.choice(list(aspect_ratio_sample_size.keys())) + local_sample_size = aspect_ratio_sample_size[aspect_ratio_key] + else: + num_frames = args.video_sample_n_frames + + aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + if rng is None: + aspect_ratio_key = np.random.choice(list(aspect_ratio_sample_size.keys())) + else: + aspect_ratio_key = rng.choice(list(aspect_ratio_sample_size.keys())) + local_sample_size = aspect_ratio_sample_size[aspect_ratio_key] + local_sample_size = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in local_sample_size] + + target_shape = ( + len(text), + vae.latent_channels, + int((num_frames - 1) // vae.temporal_compression_ratio + 1), + int(local_sample_size[0] // vae.spatial_compression_ratio), + int(local_sample_size[1] // vae.spatial_compression_ratio), + ) + + if args.low_vram: + vae.to('cpu') + real_score_transformer3d = real_score_transformer3d.to("cpu") + torch.cuda.empty_cache() + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device) + + if args.enable_text_encoder_in_dataloader: + prompt_embeds = batch['encoder_hidden_states'].to(device=accelerator.device) + neg_prompt_embeds = batch['neg_encoder_hidden_states'].to(device=accelerator.device) + else: + with torch.no_grad(): + prompt_ids = tokenizer( + batch['text'], + padding="max_length", + max_length=args.tokenizer_max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt" + ) + text_input_ids = prompt_ids.input_ids + prompt_attention_mask = prompt_ids.attention_mask + + seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long() + prompt_embeds = text_encoder(text_input_ids.to(accelerator.device), attention_mask=prompt_attention_mask.to(accelerator.device))[0] + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + + neg_txt = [ + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" for text in batch['text'] + ] + neg_prompt_ids = tokenizer( + neg_txt, + padding="max_length", + max_length=args.tokenizer_max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt" + ) + neg_text_input_ids = neg_prompt_ids.input_ids + neg_prompt_attention_mask = neg_prompt_ids.attention_mask + + neg_seq_lens = neg_prompt_attention_mask.gt(0).sum(dim=1).long() + neg_prompt_embeds = text_encoder(neg_text_input_ids.to(accelerator.device), attention_mask=neg_prompt_attention_mask.to(accelerator.device))[0] + neg_prompt_embeds = [u[:v] for u, v in zip(neg_prompt_embeds, neg_seq_lens)] + + if args.low_vram: + generator_transformer3d = generator_transformer3d.to(accelerator.device) + real_score_transformer3d = real_score_transformer3d.to(accelerator.device) + fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to('cpu') + torch.cuda.empty_cache() + + if args.train_mode == "inpaint" and spatial_compression_ratio >= 16: + mask_conditions[:, :, 1:, :, :] = 1 + _has_first_frame = not mask_conditions[:, :, 0, :, :].any() + + with accelerator.accumulate(generator_transformer3d): + def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): + sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype) + schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device) + timesteps = timesteps.to(accelerator.device) + + step_indices = [ + torch.argmin(torch.abs(schedule_timesteps - t)).item() + for t in timesteps + ] + step_indices = torch.tensor(step_indices, device=accelerator.device) + sigma = sigmas[step_indices].flatten() + + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + def add_noise(latents, noise, timesteps): + sigmas = get_sigmas(timesteps, n_dim=latents.ndim, dtype=latents.dtype) + return (1.0 - sigmas) * latents + sigmas * noise + + def generate_and_sync_list(num_denoising_steps, device): + indices = torch.randint(low=0, high=num_denoising_steps, size=(1,), generator=torch_rng, device=device) + if dist.is_initialized(): + dist.broadcast(indices, src=0) + return indices.tolist() + + def convert_flow_pred_to_x0( + scheduler, + flow_pred: torch.Tensor, + xt: torch.Tensor, + timestep: torch.Tensor + ) -> torch.Tensor: + """ + Convert flow matching's prediction to x0 prediction. + Supports both 4D [B, C, H, W] and 5D [B, C, F, H, W] inputs. + """ + original_dtype = flow_pred.dtype + device = flow_pred.device + + flow_pred = flow_pred.double() + xt = xt.double() + timesteps = scheduler.timesteps.to(device).double() + sigmas = scheduler.sigmas.to(device).double() + timestep = timestep.to(device).double() + + timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma_t = sigmas[timestep_id] + + ndim = flow_pred.ndim + if ndim == 4: + sigma_t = sigma_t.view(-1, 1, 1, 1) + elif ndim == 5: + sigma_t = sigma_t.view(-1, 1, 1, 1, 1) + else: + raise ValueError(f"Expected 4D or 5D input, got {ndim}D tensor.") + + x0_pred = xt - sigma_t * flow_pred + return x0_pred.to(original_dtype) + + # Create discrete denoising steps (per-step, with optional randomization) + if getattr(args, 'randomize_step_indices', False): + random_indices = randomize_denoising_step_indices( + args.denoising_step_indices_list, + args.train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=getattr(args, 'index_jitter_ratio', 0.30), + ) + else: + random_indices = torch.tensor(args.denoising_step_indices_list) + + denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - random_indices] + + # --- Main Training Logic --- + bsz, channel, num_frames, height, width = target_shape + if step % args.gen_update_interval == 0: + generator_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=weight_dtype) + num_denoising_steps = len(denoising_step_list) + final_step_index = generate_and_sync_list(num_denoising_steps, device=generator_noise.device)[0] + + # Precompute seq_len once (same for all steps) + patch_h, patch_w = accelerator.unwrap_model(generator_transformer3d).config.patch_size[1:] + seq_len = math.ceil((width * height) / (patch_h * patch_w) * num_frames) + + for index, current_timestep in enumerate(denoising_step_list): + is_final_step = (index == final_step_index) + timestep = torch.full( + generator_noise.shape[:1], + current_timestep, + device=generator_noise.device, + dtype=torch.int64 + ) + + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + context_manager = torch.no_grad() if not is_final_step else contextlib.nullcontext() + + with context_manager: + _timestep = timestep + _generator_noise = generator_noise + if args.train_mode == "inpaint" and spatial_compression_ratio >= 16: + _generator_noise = (1 - mask_conditions) * inpaint_latents[:, -vae.latent_channels:] + mask_conditions * _generator_noise + if _has_first_frame: + _temp_ts = (mask_conditions[:, 0, :, ::2, ::2] * _timestep[:, None, None, None]).flatten(1) + _timestep = torch.cat([_temp_ts, _temp_ts.new_ones(mask_conditions.size(0), seq_len - _temp_ts.size(1)) * _timestep[:, None,]], dim=1) + else: + _timestep = mask_conditions.new_ones(mask_conditions.size(0), seq_len) * _timestep[:, None,] + generator_pred = generator_transformer3d( + x=_generator_noise, + context=prompt_embeds, + t=_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + ) + generator_pred = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=generator_pred, + xt=_generator_noise, + timestep=timestep + ) + + if is_final_step: + break + + next_timestep = denoising_step_list[index + 1] * torch.ones( + generator_noise.shape[:1], dtype=torch.long, device=generator_noise.device + ) + generator_noise = add_noise( + generator_pred, + torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng), + next_timestep + ) + + indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu() + generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device) + generator_denoised_input = add_noise( + generator_pred, + torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng), + generator_timestep + ).detach().to(accelerator.device, dtype=weight_dtype) + + # Compute fake score + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device), torch.no_grad(): + _generator_timestep = generator_timestep + _generator_denoised_input = generator_denoised_input + if args.train_mode == "inpaint" and spatial_compression_ratio >= 16: + _generator_denoised_input = (1 - mask_conditions) * inpaint_latents[:, -vae.latent_channels:] + mask_conditions * _generator_denoised_input + if _has_first_frame: + _temp_ts = (mask_conditions[:, 0, :, ::2, ::2] * _generator_timestep[:, None, None, None]).flatten(1) + _generator_timestep = torch.cat([_temp_ts, _temp_ts.new_ones(mask_conditions.size(0), seq_len - _temp_ts.size(1)) * _generator_timestep[:, None,]], dim=1) + else: + _generator_timestep = mask_conditions.new_ones(mask_conditions.size(0), seq_len) * _generator_timestep[:, None,] + fake_score_main_cond = fake_score_transformer3d( + x=_generator_denoised_input, + context=prompt_embeds, + t=_generator_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + ) + fake_score_main_cond = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=fake_score_main_cond, + xt=generator_denoised_input, + timestep=generator_timestep + ) + + if args.fake_guidance_scale != 0.0: + fake_score_main_uncond = fake_score_transformer3d( + x=_generator_denoised_input, + context=neg_prompt_embeds, + t=_generator_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + ) + fake_score_main_uncond = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=fake_score_main_uncond, + xt=generator_denoised_input, + timestep=generator_timestep + ) + fake_score_main = fake_score_main_uncond + ( + fake_score_main_cond - fake_score_main_uncond + ) * args.fake_guidance_scale + else: + fake_score_main = fake_score_main_cond + + # Compute real score + real_score_main_cond = real_score_transformer3d( + x=_generator_denoised_input, + context=prompt_embeds, + t=_generator_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + ) + real_score_main_cond = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=real_score_main_cond, + xt=generator_denoised_input, + timestep=generator_timestep + ) + + real_score_main_uncond = real_score_transformer3d( + x=_generator_denoised_input, + context=neg_prompt_embeds, + t=_generator_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + ) + real_score_main_uncond = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=real_score_main_uncond, + xt=generator_denoised_input, + timestep=generator_timestep + ) + + real_score_main = real_score_main_uncond + ( + real_score_main_cond - real_score_main_uncond + ) * args.real_guidance_scale + + # DMD loss + fake_to_real_grad = fake_score_main - real_score_main + generator_to_real_norm = generator_pred - real_score_main + normalizer = torch.abs(generator_to_real_norm).mean(dim=[1, 2, 3, 4], keepdim=True) + fake_to_real_grad = fake_to_real_grad / normalizer + fake_to_real_grad = torch.nan_to_num(fake_to_real_grad) + + dmd_loss = 0.5 * F.mse_loss( + generator_pred.double(), + (generator_pred.double() - fake_to_real_grad.double()).detach(), + reduction="mean" + ) + avg_dmd_loss = accelerator.gather(dmd_loss.repeat(args.train_batch_size)).mean() + train_dmd_loss += avg_dmd_loss.item() / args.gradient_accumulation_steps + + if args.low_vram: + real_score_transformer3d = real_score_transformer3d.to("cpu") + fake_score_transformer3d = fake_score_transformer3d.to("cpu") + torch.cuda.empty_cache() + + accelerator.backward(dmd_loss) + if accelerator.sync_gradients: + accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + if args.low_vram: + fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device) + torch.cuda.empty_cache() + + with accelerator_fake_score_transformer3d.accumulate(fake_score_transformer3d): + # --- Fake Critic Denoising Loss --- + with torch.no_grad(): + fake_score_critic_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=weight_dtype) + num_denoising_steps = len(denoising_step_list) + final_step_index = generate_and_sync_list(num_denoising_steps, device=fake_score_critic_noise.device)[0] + + patch_h, patch_w = accelerator.unwrap_model(generator_transformer3d).config.patch_size[1:] + seq_len = math.ceil((width * height) / (patch_h * patch_w) * num_frames) + + for index, current_timestep in enumerate(denoising_step_list): + is_final_step = (index == final_step_index) + timestep = torch.full( + fake_score_critic_noise.shape[:1], + current_timestep, + device=fake_score_critic_noise.device, + dtype=torch.int64 + ) + + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + _timestep = timestep + _fake_score_critic_noise = fake_score_critic_noise + if args.train_mode == "inpaint" and spatial_compression_ratio >= 16: + _fake_score_critic_noise = (1 - mask_conditions) * inpaint_latents[:, -vae.latent_channels:] + mask_conditions * _fake_score_critic_noise + if _has_first_frame: + _temp_ts = (mask_conditions[:, 0, :, ::2, ::2] * _timestep[:, None, None, None]).flatten(1) + _timestep = torch.cat([_temp_ts, _temp_ts.new_ones(mask_conditions.size(0), seq_len - _temp_ts.size(1)) * _timestep[:, None,]], dim=1) + else: + _timestep = mask_conditions.new_ones(mask_conditions.size(0), seq_len) * _timestep[:, None,] + fake_score_denoised_pred = generator_transformer3d( + x=_fake_score_critic_noise, + context=prompt_embeds, + t=_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + ) + fake_score_denoised_pred = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=fake_score_denoised_pred, + xt=_fake_score_critic_noise, + timestep=timestep + ) + + if is_final_step: + break + + next_timestep = denoising_step_list[index + 1] * torch.ones( + fake_score_critic_noise.shape[:1], + dtype=torch.long, + device=fake_score_critic_noise.device + ) + + fake_score_critic_noise = add_noise( + fake_score_denoised_pred, + torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng), + next_timestep + ) + + indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu() + critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device) + critic_noise = torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng) + + fake_score_denoised_input = add_noise( + fake_score_denoised_pred, + critic_noise, + critic_timestep + ) + _critic_timestep = critic_timestep + if args.train_mode == "inpaint" and spatial_compression_ratio >= 16: + fake_score_denoised_input = (1 - mask_conditions) * inpaint_latents[:, -vae.latent_channels:] + mask_conditions * fake_score_denoised_input + if _has_first_frame: + _temp_ts = (mask_conditions[:, 0, :, ::2, ::2] * _critic_timestep[:, None, None, None]).flatten(1) + _critic_timestep = torch.cat([_temp_ts, _temp_ts.new_ones(mask_conditions.size(0), seq_len - _temp_ts.size(1)) * _critic_timestep[:, None,]], dim=1) + else: + _critic_timestep = mask_conditions.new_ones(mask_conditions.size(0), seq_len) * _critic_timestep[:, None,] + + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + fake_score_denoised_output = fake_score_transformer3d( + x=fake_score_denoised_input, + context=prompt_embeds, + t=_critic_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + ) + + def custom_mse_loss(noise_pred, target, weighting=None, threshold=50): + noise_pred = noise_pred.float() + target = target.float() + diff = noise_pred - target + mse_loss = F.mse_loss(noise_pred, target, reduction='none') + mask = (diff.abs() <= threshold).float() + masked_loss = mse_loss * mask + if weighting is not None: + masked_loss = masked_loss * weighting + final_loss = masked_loss.mean() + return final_loss + + denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred) + avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean() + train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps + + accelerator_fake_score_transformer3d.backward(denoising_loss) + if accelerator_fake_score_transformer3d.sync_gradients: + accelerator_fake_score_transformer3d.clip_grad_norm_(fake_trainable_params, args.max_grad_norm) + critic_optimizer.step() + fake_score_lr_scheduler.step() + critic_optimizer.zero_grad() + + if args.low_vram: + fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device) + generator_transformer3d = generator_transformer3d.to(accelerator.device) + + # Checks if the accelerator has performed an optimization step behind the scenes + if accelerator.sync_gradients: + + progress_bar.update(1) + global_step += 1 + accelerator.log({"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss}, step=global_step) + train_dmd_loss = 0.0 + train_denoising_loss = 0.0 + + if global_step % args.checkpointing_steps == 0: + if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + fake_score_save_path = os.path.join(save_path, "fake_score") + accelerator.save_state(save_path) + accelerator_fake_score_transformer3d.save_state(fake_score_save_path) + logger.info(f"Saved state to {save_path}") + + if args.validation_prompts is not None and global_step % args.validation_steps == 0: + log_validation( + vae, + text_encoder, + tokenizer, + generator_transformer3d, + args, + config, + accelerator, + weight_dtype, + global_step, + ) + + logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + + if global_step >= args.max_train_steps: + break + + if args.validation_prompts is not None and epoch % args.validation_epochs == 0: + log_validation( + vae, + text_encoder, + tokenizer, + generator_transformer3d, + args, + config, + accelerator, + weight_dtype, + global_step, + ) + + # Create the pipeline using the trained modules and save it. + accelerator.wait_for_everyone() + if accelerator.is_main_process: + generator_transformer3d = unwrap_model(generator_transformer3d) + + if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process: + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + fake_score_save_path = os.path.join(save_path, "fake_score") + accelerator.save_state(save_path) + accelerator_fake_score_transformer3d.save_state(fake_score_save_path) + logger.info(f"Saved state to {save_path}") + + accelerator.end_training() + + +if __name__ == "__main__": + main() diff --git a/scripts/wan2.2_fun/train_distill.sh b/scripts/wan2.2_fun/train_distill.sh new file mode 100644 index 00000000..b34abd7f --- /dev/null +++ b/scripts/wan2.2_fun/train_distill.sh @@ -0,0 +1,44 @@ +export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-Fun-A14B-InP" +export DATASET_NAME="datasets/internal_datasets/" +export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. +# export NCCL_IB_DISABLE=1 +# export NCCL_P2P_DISABLE=1 +NCCL_DEBUG=INFO + +accelerate launch --mixed_precision="bf16" scripts/wan2.2_fun/train_distill.py \ + --config_path="config/wan2.2/wan_civitai_i2v.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --video_repeat=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=50 \ + --learning_rate=2e-06 \ + --learning_rate_critic=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_wan2.2_fun_distill" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --uniform_sampling \ + --boundary_type="low" \ + --train_mode="inpaint" \ + --trainable_modules "." \ + --low_vram \ No newline at end of file diff --git a/scripts/wan2.2_fun/train_reward_lora.py b/scripts/wan2.2_fun/train_reward_lora.py index d184fe4e..f53b31c7 100644 --- a/scripts/wan2.2_fun/train_reward_lora.py +++ b/scripts/wan2.2_fun/train_reward_lora.py @@ -37,7 +37,10 @@ from accelerate.logging import get_logger from accelerate.state import AcceleratorState from accelerate.utils import ProjectConfiguration, set_seed -from decord import VideoReader +try: + from decord import VideoReader +except ImportError: + from videox_fun.data.utils import AVVideoReader as VideoReader from diffusers import FlowMatchEulerDiscreteScheduler from diffusers.optimization import get_scheduler from diffusers.utils import check_min_version, deprecate, is_wandb_available diff --git a/scripts/wan2.2_vace_fun/README_TRAIN.md b/scripts/wan2.2_vace_fun/README_TRAIN.md index 4f5de639..e713ea03 100644 --- a/scripts/wan2.2_vace_fun/README_TRAIN.md +++ b/scripts/wan2.2_vace_fun/README_TRAIN.md @@ -53,9 +53,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -204,8 +205,8 @@ VACE training datasets require original videos with corresponding control signal If data paths are relative, set in the training script: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" ``` **Absolute paths**: @@ -214,7 +215,7 @@ If data paths are absolute, set in the training script: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: Use relative paths for small local datasets; use absolute paths for external storage (NAS, OSS) or shared storage across multiple machines. diff --git a/scripts/wan2.2_vace_fun/README_TRAIN_zh-CN.md b/scripts/wan2.2_vace_fun/README_TRAIN_zh-CN.md index a68f13d6..2a84e8e9 100644 --- a/scripts/wan2.2_vace_fun/README_TRAIN_zh-CN.md +++ b/scripts/wan2.2_vace_fun/README_TRAIN_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -194,8 +195,8 @@ VACE 训练数据集除了原始视频外，还需要提供一一对应的控制 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height_add_objects.json" ``` **绝对路径**： @@ -204,7 +205,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -234,7 +235,7 @@ modelscope download --model PAI/Wan2.2-VACE-Fun-A14B --local_dir models/Diffusio ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-VACE-Fun-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height_add_objects.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -296,7 +297,7 @@ Wan2.2采用了创新的双Transformer架构： | `--config_path` | 配置文件路径 | `config/wan2.2/wan_civitai_t2v.yaml` | | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.2-VACE-Fun-A14B` | | `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Controls-Demo/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height_add_objects.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 图像最大训练分辨率 | 640 | | `--video_sample_size` | 视频最大训练分辨率 | 640 | @@ -388,7 +389,7 @@ Wan2.2采用了创新的双Transformer架构： ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-VACE-Fun-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height_add_objects.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -448,7 +449,7 @@ python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/ ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-VACE-Fun-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height_add_objects.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -498,7 +499,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-VACE-Fun-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height_add_objects.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -555,7 +556,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2_vace_fun/train.py \ ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-VACE-Fun-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height_add_objects.json" export MASTER_ADDR="192.168.1.100" # Master 机器 IP export MASTER_PORT=10086 export WORLD_SIZE=2 # 机器总数 @@ -607,7 +608,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main ```bash export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-VACE-Fun-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Controls-Demo/" -export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height.json" +export DATASET_META_NAME="datasets/X-Fun-Videos-Controls-Demo/metadata_add_width_height_add_objects.json" export MASTER_ADDR="192.168.1.100" # 与 Master 相同 export MASTER_PORT=10086 export WORLD_SIZE=2 diff --git a/scripts/z_image/README_TRAIN.md b/scripts/z_image/README_TRAIN.md index 7364f388..4e404dba 100644 --- a/scripts/z_image/README_TRAIN.md +++ b/scripts/z_image/README_TRAIN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -131,8 +132,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -141,7 +142,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -213,8 +214,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/Z-Image` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--image_sample_size` | Maximum training resolution, auto bucketing | 1328 | | `--gradient_accumulation_steps` | Gradient accumulation steps (equivalent to larger batch) | 1 | @@ -272,8 +273,8 @@ You can configure validation parameters to periodically generate test images dur ```sh export MODEL_NAME="models/Diffusion_Transformer/Z-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -311,8 +312,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/Z-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/z_image/README_TRAIN_LORA.md b/scripts/z_image/README_TRAIN_LORA.md index 8016981d..c15ee2ee 100644 --- a/scripts/z_image/README_TRAIN_LORA.md +++ b/scripts/z_image/README_TRAIN_LORA.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Method 3: Using Docker** @@ -131,8 +132,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun If your data uses relative paths, configure the training script as follows: ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **Absolute Paths**: @@ -141,7 +142,7 @@ If your data uses absolute paths, configure the training script as follows: ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. @@ -214,8 +215,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/Z-Image` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Samples per batch | 1 | | `--image_sample_size` | Maximum training resolution, auto bucketing | 1328 | | `--gradient_accumulation_steps` | Gradient accumulation steps (equivalent to larger batch) | 1 | @@ -271,8 +272,8 @@ You can configure validation parameters to periodically generate test images dur ```sh export MODEL_NAME="models/Diffusion_Transformer/Z-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -311,8 +312,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/Z-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/z_image/README_TRAIN_LORA_zh-CN.md b/scripts/z_image/README_TRAIN_LORA_zh-CN.md index 63db9743..703e63cf 100644 --- a/scripts/z_image/README_TRAIN_LORA_zh-CN.md +++ b/scripts/z_image/README_TRAIN_LORA_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -131,8 +132,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -141,7 +142,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -214,8 +215,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Z-Image` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，代码会自动分桶 | 1328 | | `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 | @@ -271,8 +272,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/Z-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -311,8 +312,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/Z-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/z_image/README_TRAIN_zh-CN.md b/scripts/z_image/README_TRAIN_zh-CN.md index 633816cc..946ebd47 100644 --- a/scripts/z_image/README_TRAIN_zh-CN.md +++ b/scripts/z_image/README_TRAIN_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3：使用docker** @@ -131,8 +132,8 @@ modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun 如果数据的路径为相对路径，则在训练脚本中设置： ```bash -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" ``` **绝对路径**： @@ -141,7 +142,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata.json" ```bash export DATASET_NAME="" -export DATASET_META_NAME="/mnt/data/metadata.json" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" ``` > 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 @@ -213,8 +214,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Z-Image` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Images-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 最大训练分辨率，代码会自动分桶 | 1328 | | `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 | @@ -272,8 +273,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```sh export MODEL_NAME="models/Diffusion_Transformer/Z-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -311,8 +312,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```sh export MODEL_NAME="models/Diffusion_Transformer/Z-Image" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/z_image/train_distill.py b/scripts/z_image/train_distill.py index 82a6090d..ff09ad49 100644 --- a/scripts/z_image/train_distill.py +++ b/scripts/z_image/train_distill.py @@ -1033,8 +1033,8 @@ def load_model_hook(models, input_dir): fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_transformer3d.parameters())) fake_trainable_params_optim = [ - {'params': [], 'lr': args.learning_rate}, - {'params': [], 'lr': args.learning_rate / 2}, + {'params': [], 'lr': args.learning_rate_critic}, + {'params': [], 'lr': args.learning_rate_critic / 2}, ] in_already = [] for name, param in fake_score_transformer3d.named_parameters(): @@ -1047,7 +1047,7 @@ def load_model_hook(models, input_dir): high_lr_flag = True fake_trainable_params_optim[0]['params'].append(param) if accelerator.is_main_process: - print(f"Set {name} to lr : {args.learning_rate}") + print(f"Set {name} to lr : {args.learning_rate_critic}") break if high_lr_flag: continue @@ -1056,7 +1056,7 @@ def load_model_hook(models, input_dir): in_already.append(name) fake_trainable_params_optim[1]['params'].append(param) if accelerator.is_main_process: - print(f"Set {name} to lr : {args.learning_rate / 2}") + print(f"Set {name} to lr : {args.learning_rate_critic / 2}") break if args.use_came: @@ -1173,7 +1173,7 @@ def collate_fn(examples): ) fake_score_lr_scheduler = get_scheduler( args.lr_scheduler, - optimizer=optimizer, + optimizer=critic_optimizer, num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, num_training_steps=args.max_train_steps * accelerator.num_processes, ) diff --git a/scripts/z_image/train_distill.sh b/scripts/z_image/train_distill.sh index 6631c35a..233c4b5e 100644 --- a/scripts/z_image/train_distill.sh +++ b/scripts/z_image/train_distill.sh @@ -17,7 +17,7 @@ accelerate launch --mixed_precision="bf16" scripts/z_image/train_distill.py \ --num_train_epochs=100 \ --checkpointing_steps=100 \ --learning_rate=2e-05 \ - --learning_rate_critic=2e-06 \ + --learning_rate_critic=2e-05 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ diff --git a/scripts/z_image/train_distill_lora.py b/scripts/z_image/train_distill_lora.py index f05aedda..e7d5847c 100644 --- a/scripts/z_image/train_distill_lora.py +++ b/scripts/z_image/train_distill_lora.py @@ -1043,7 +1043,7 @@ def load_model_hook(models, input_dir): logging.info("Add fake_score_network parameters") fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_network.parameters())) - fake_trainable_params_optim = fake_score_network.prepare_optimizer_params(args.learning_rate / 2, args.learning_rate, args.learning_rate) + fake_trainable_params_optim = fake_score_network.prepare_optimizer_params(args.learning_rate_critic / 2, args.learning_rate_critic, args.learning_rate_critic) if args.use_came: optimizer = optimizer_cls( @@ -1159,7 +1159,7 @@ def collate_fn(examples): ) fake_score_lr_scheduler = get_scheduler( args.lr_scheduler, - optimizer=optimizer, + optimizer=critic_optimizer, num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, num_training_steps=args.max_train_steps * accelerator.num_processes, ) diff --git a/scripts/z_image/train_distill_lora.sh b/scripts/z_image/train_distill_lora.sh index 7469e70c..db3d4a45 100644 --- a/scripts/z_image/train_distill_lora.sh +++ b/scripts/z_image/train_distill_lora.sh @@ -17,7 +17,7 @@ accelerate launch --mixed_precision="bf16" scripts/z_image/train_distill_lora.py --num_train_epochs=100 \ --checkpointing_steps=50 \ --learning_rate=1e-04 \ - --learning_rate_critic=1e-05 \ + --learning_rate_critic=1e-04 \ --seed=42 \ --output_dir="output_dir_z_image_distill_lora" \ --gradient_checkpointing \ diff --git a/scripts/z_image_fun/README_TRAIN.md b/scripts/z_image_fun/README_TRAIN.md index 77c82279..1937ac74 100644 --- a/scripts/z_image_fun/README_TRAIN.md +++ b/scripts/z_image_fun/README_TRAIN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **Option 3: Using Docker** @@ -142,45 +143,25 @@ It is recommended to use tools like [DWPose](https://github.com/IDEA-Research/DW ### 2.4 Relative vs Absolute Path Usage -**Option 1: Using Relative Paths (Recommended)** +**Relative Paths**: -When data paths are not fixed or you need to train on different machines, it is recommended to use relative paths. +If your data uses relative paths, set in the training script: -Configure relative paths in `metadata.json`, then specify the dataset root directory via `--train_data_dir` in the training script: - -```json -[ - { - "file_path": "train/image001.jpg", - "control_file_path": "control/image001.jpg", - "text": "A group of young men in suits and sunglasses are walking down a city street.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata.json" ``` -During training, the system will automatically search for files corresponding to relative paths under `--train_data_dir`. +**Absolute Paths**: -**Option 2: Using Absolute Paths** +If your data uses absolute paths, set in the training script: -If the dataset path is fixed, you can directly configure absolute paths in `metadata.json`: - -```json -[ - { - "file_path": "/mnt/data/images/image001.jpg", - "control_file_path": "/mnt/data/controls/image001.jpg", - "text": "A group of young men in suits and sunglasses.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="" +export DATASET_META_NAME="/mnt/data/metadata.json" ``` -When using absolute paths, the `--train_data_dir` parameter serves only as a default path, and the absolute paths in json will take precedence. +> 💡 **Recommendation**: If the dataset is small and stored locally, relative paths are recommended; if the dataset is stored on external storage (such as NAS, OSS) or shared across multiple machines, absolute paths are recommended. --- @@ -236,8 +217,8 @@ After downloading data according to **2.1 Quick Test Dataset** and weights accor ```bash export MODEL_NAME="models/Diffusion_Transformer/Z-Image-Turbo" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -279,8 +260,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----------|-------------|---------------| | `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Z-Image-Turbo` | -| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` | -| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Images-Controls-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json` | | `--train_batch_size` | Batch size per device | 1 | | `--image_sample_size` | Maximum training resolution, automatic bucketing | 1328 | | `--gradient_accumulation_steps` | Gradient accumulation steps (equivalent to larger batch) | 1 | @@ -331,8 +312,8 @@ If DeepSpeed-Zero-2 runs out of GPU memory, you can switch to FSDP for training: ```bash export MODEL_NAME="models/Diffusion_Transformer/Z-Image-Turbo" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -375,8 +356,8 @@ Training without DeepSpeed or FSDP may result in insufficient GPU memory. Only r ```bash export MODEL_NAME="models/Diffusion_Transformer/Z-Image-Turbo" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/z_image_fun/README_TRAIN_zh-CN.md b/scripts/z_image_fun/README_TRAIN_zh-CN.md index 902e3ec3..4665a52c 100644 --- a/scripts/z_image_fun/README_TRAIN_zh-CN.md +++ b/scripts/z_image_fun/README_TRAIN_zh-CN.md @@ -43,9 +43,10 @@ pip install -r requirements.txt pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" -pip install yunchang xfuser modelscope openpyxl deepspeed==0.17.0 numpy==1.26.4 +pip install yunchang xfuser modelscope openpyxl pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 ``` **方式 3:使用docker** @@ -142,45 +143,25 @@ Control 模式的 metadata.json 与普通 Z-Image 的 json 略有不同，需要 ### 2.4 相对路径与绝对路径使用方案 -**方案 1:使用相对路径(推荐)** +**相对路径**： -当数据路径不固定,或需要在不同机器上训练时,推荐使用相对路径。 +如果数据的路径为相对路径，则在训练脚本中设置： -在 `metadata.json` 中配置相对路径,然后在训练脚本中通过 `--train_data_dir` 指定数据集根目录: - -```json -[ - { - "file_path": "train/image001.jpg", - "control_file_path": "control/image001.jpg", - "text": "A group of young men in suits and sunglasses are walking down a city street.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata.json" ``` -训练时会自动在 `--train_data_dir` 下寻找相对路径对应的文件。 +**绝对路径**： -**方案 2:使用绝对路径** +如果数据的路径为绝对路径，则在训练脚本中设置： -如果数据集路径固定,可以直接在 `metadata.json` 中配置绝对路径: - -```json -[ - { - "file_path": "/mnt/data/images/image001.jpg", - "control_file_path": "/mnt/data/controls/image001.jpg", - "text": "A group of young men in suits and sunglasses.", - "width": 1024, - "height": 1024, - "type": "image" - } -] +```bash +export DATASET_NAME="" +export DATASET_META_NAME="/mnt/data/metadata.json" ``` -使用绝对路径时,`--train_data_dir` 参数仅作为默认路径,实际会优先使用 json 中的绝对路径。 +> 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 --- @@ -236,8 +217,8 @@ hf download alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.1 --local-dir model ```bash export MODEL_NAME="models/Diffusion_Transformer/Z-Image-Turbo" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -279,8 +260,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | 参数 | 说明 | 示例值 | |-----|------|-------| | `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Z-Image-Turbo` | -| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` | -| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` | +| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Images-Controls-Demo/` | +| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json` | | `--train_batch_size` | 每批次样本数 | 1 | | `--image_sample_size` | 最大训练分辨率,代码会自动分桶 | 1328 | | `--gradient_accumulation_steps` | 梯度累积步数(等效增大 batch) | 1 | @@ -331,8 +312,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con ```bash export MODEL_NAME="models/Diffusion_Transformer/Z-Image-Turbo" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -375,8 +356,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR ```bash export MODEL_NAME="models/Diffusion_Transformer/Z-Image-Turbo" -export DATASET_NAME="datasets/internal_datasets/" -export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 diff --git a/scripts/z_image_fun/train_control_distill.py b/scripts/z_image_fun/train_control_distill.py index 0afeb72a..d1c33866 100644 --- a/scripts/z_image_fun/train_control_distill.py +++ b/scripts/z_image_fun/train_control_distill.py @@ -1102,8 +1102,8 @@ def load_model_hook(models, input_dir): fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_transformer3d.parameters())) fake_trainable_params_optim = [ - {'params': [], 'lr': args.learning_rate}, - {'params': [], 'lr': args.learning_rate / 2}, + {'params': [], 'lr': args.learning_rate_critic}, + {'params': [], 'lr': args.learning_rate_critic / 2}, ] in_already = [] for name, param in fake_score_transformer3d.named_parameters(): @@ -1116,7 +1116,7 @@ def load_model_hook(models, input_dir): high_lr_flag = True fake_trainable_params_optim[0]['params'].append(param) if accelerator.is_main_process: - print(f"Set {name} to lr : {args.learning_rate}") + print(f"Set {name} to lr : {args.learning_rate_critic}") break if high_lr_flag: continue @@ -1125,7 +1125,7 @@ def load_model_hook(models, input_dir): in_already.append(name) fake_trainable_params_optim[1]['params'].append(param) if accelerator.is_main_process: - print(f"Set {name} to lr : {args.learning_rate / 2}") + print(f"Set {name} to lr : {args.learning_rate_critic / 2}") break if args.use_came: @@ -1387,7 +1387,7 @@ def _create_special_list(length): ) fake_score_lr_scheduler = get_scheduler( args.lr_scheduler, - optimizer=optimizer, + optimizer=critic_optimizer, num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, num_training_steps=args.max_train_steps * accelerator.num_processes, ) diff --git a/scripts/z_image_fun/train_turbo_control_distill.sh b/scripts/z_image_fun/train_turbo_control_distill.sh index eb90723d..3ee703cf 100644 --- a/scripts/z_image_fun/train_turbo_control_distill.sh +++ b/scripts/z_image_fun/train_turbo_control_distill.sh @@ -18,7 +18,7 @@ accelerate launch --mixed_precision="bf16" scripts/z_image_fun/train_control_dis --num_train_epochs=100 \ --checkpointing_steps=100 \ --learning_rate=2e-06 \ - --learning_rate_critic=2e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ diff --git a/videox_fun/data/dataset_image_video.py b/videox_fun/data/dataset_image_video.py index 3f6b116e..c3bc3b63 100755 --- a/videox_fun/data/dataset_image_video.py +++ b/videox_fun/data/dataset_image_video.py @@ -14,7 +14,6 @@ import torch import torch.nn.functional as F import torchvision.transforms as transforms -from decord import VideoReader from einops import rearrange from func_timeout import FunctionTimedOut, func_timeout from packaging import version as pver @@ -23,6 +22,11 @@ from torch.utils.data import BatchSampler, Sampler from torch.utils.data.dataset import Dataset +try: + from decord import VideoReader +except ImportError: + from .utils import AVVideoReader as VideoReader + from .utils import (VIDEO_READER_TIMEOUT, VideoReader_contextmanager, get_random_mask, get_video_reader_batch, padding_image, process_pose_file, process_pose_params, resize_frame, diff --git a/videox_fun/data/dataset_video.py b/videox_fun/data/dataset_video.py index df4b1038..9bd01930 100644 --- a/videox_fun/data/dataset_video.py +++ b/videox_fun/data/dataset_video.py @@ -9,12 +9,16 @@ import numpy as np import torch import torchvision.transforms as transforms -from decord import VideoReader from einops import rearrange from func_timeout import FunctionTimedOut, func_timeout from PIL import Image from torch.utils.data.dataset import Dataset +try: + from decord import VideoReader +except ImportError: + from .utils import AVVideoReader as VideoReader + from .utils import (VIDEO_READER_TIMEOUT, VideoReader_contextmanager, get_random_mask, get_video_reader_batch, resize_frame) diff --git a/videox_fun/data/utils.py b/videox_fun/data/utils.py index 58911511..ab2d8d8f 100644 --- a/videox_fun/data/utils.py +++ b/videox_fun/data/utils.py @@ -6,13 +6,81 @@ import cv2 import numpy as np import torch -from decord import VideoReader from einops import rearrange from packaging import version as pver from PIL import Image +try: + from decord import VideoReader + HAS_DECORD = True +except ImportError: + HAS_DECORD = False + print("Warning: decord is not installed. Falling back to PyAV for video reading. " + "Install decord for better performance: pip install decord") + VIDEO_READER_TIMEOUT = 20 + +class AVVideoReader: + """A VideoReader implementation using PyAV as a fallback when decord is unavailable. + + Provides the same interface as decord.VideoReader: + - len(reader) returns total frame count + - reader.get_batch(indices) returns a BatchFrames object with .asnumpy() + - reader.get_avg_fps() returns the average FPS + """ + def __init__(self, uri, num_threads=1, **kwargs): + import av + self._container = av.open(uri) + self._stream = self._container.streams.video[0] + self._stream.thread_type = 'AUTO' + self._num_frames = self._stream.frames + # Some videos may not report frame count; decode to count + if self._num_frames == 0: + for _ in self._container.decode(video=0): + self._num_frames += 1 + self._container.seek(0) + self._avg_fps = float(self._stream.average_rate) if self._stream.average_rate else 24.0 + + def __len__(self): + return self._num_frames + + def get_avg_fps(self): + return self._avg_fps + + def get_batch(self, indices): + """Read frames at specified indices. Returns an object with .asnumpy() method.""" + import av + indices_set = set(indices) + max_idx = max(indices) + frames_dict = {} + + self._container.seek(0) + frame_idx = 0 + for frame in self._container.decode(video=0): + if frame_idx in indices_set: + frames_dict[frame_idx] = frame.to_ndarray(format='rgb24') + if frame_idx >= max_idx: + break + frame_idx += 1 + + # Assemble frames in requested order + frames = [frames_dict[i] for i in indices] + return _AVBatchFrames(frames) + + def __del__(self): + if hasattr(self, '_container') and self._container is not None: + self._container.close() + + +class _AVBatchFrames: + """Wrapper to mimic decord's batch result with .asnumpy() interface.""" + def __init__(self, frames): + self._frames = frames + + def asnumpy(self): + return np.stack(self._frames) + def get_random_mask(shape, image_start_only=False): f, c, h, w = shape mask = torch.zeros((f, 1, h, w), dtype=torch.uint8) @@ -98,7 +166,10 @@ def get_random_mask(shape, image_start_only=False): @contextmanager def VideoReader_contextmanager(*args, **kwargs): - vr = VideoReader(*args, **kwargs) + if HAS_DECORD: + vr = VideoReader(*args, **kwargs) + else: + vr = AVVideoReader(*args, **kwargs) try: yield vr finally: diff --git a/videox_fun/dist/qwen_xfuser.py b/videox_fun/dist/qwen_xfuser.py index 55bb32af..9d02dcaa 100644 --- a/videox_fun/dist/qwen_xfuser.py +++ b/videox_fun/dist/qwen_xfuser.py @@ -1,34 +1,11 @@ -import functools -import glob -import json -import math -import os -import types -import warnings -from typing import Any, Dict, List, Optional, Tuple, Union - -import numpy as np +from typing import Optional, Tuple, Union + import torch -import torch.cuda.amp as amp -import torch.nn as nn import torch.nn.functional as F -from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin -from diffusers.loaders.single_file_model import FromOriginalModelMixin -from diffusers.models.attention import FeedForward from diffusers.models.attention_processor import Attention -from diffusers.models.embeddings import TimestepEmbedding, Timesteps -from diffusers.models.modeling_outputs import Transformer2DModelOutput -from diffusers.models.modeling_utils import ModelMixin -from diffusers.models.normalization import AdaLayerNormContinuous, RMSNorm -from diffusers.utils import (USE_PEFT_BACKEND, is_torch_version, logging, - scale_lora_layers, unscale_lora_layers) -from diffusers.utils.torch_utils import maybe_allow_in_graph -from torch import nn -from .fuser import (get_sequence_parallel_rank, - get_sequence_parallel_world_size, get_sp_group, - init_distributed_environment, initialize_model_parallel, - xFuserLongContextAttention) + +from .fuser import xFuserLongContextAttention + def apply_rotary_emb_qwen( x: torch.Tensor, diff --git a/videox_fun/dist/wan_xfuser.py b/videox_fun/dist/wan_xfuser.py index 6c2358c3..dbf8ad86 100755 --- a/videox_fun/dist/wan_xfuser.py +++ b/videox_fun/dist/wan_xfuser.py @@ -308,16 +308,36 @@ def qkv_fn(x): kv_cache["k"][:, local_start_index:local_end_index] = roped_key_full kv_cache["v"][:, local_start_index:local_end_index] = v_full - # Step 4: chunk back to SP distribution for attention computation + # Step 4: pad full Q seq dim to multiple of sp_size for even chunk under SP, + # then chunk back to SP distribution for attention computation. + unpadded_q_len = roped_query_full.shape[1] + q_pad = (sp_size - unpadded_q_len % sp_size) % sp_size + if q_pad > 0: + roped_query_full = torch.nn.functional.pad( + roped_query_full, (0, 0, 0, 0, 0, q_pad)) roped_query = torch.chunk(roped_query_full, sp_size, dim=1)[sp_rank] - + # Step 5: compute attention using xFuserLongContextAttention for sequence parallelism # Chunk KV cache window to match SP distribution - kv_k_full = kv_cache["k"][:, max(0, local_end_index - self.max_attention_size):local_end_index] - kv_v_full = kv_cache["v"][:, max(0, local_end_index - self.max_attention_size):local_end_index] + if self.local_attn_size == -1: + max_attention_size = local_end_index + else: + max_attention_size = self.local_attn_size * frame_seqlen + + kv_k_full = kv_cache["k"][:, max(0, local_end_index - max_attention_size):local_end_index] + kv_v_full = kv_cache["v"][:, max(0, local_end_index - max_attention_size):local_end_index] + + # Pad KV window seq dim to multiple of sp_size for even chunk under SP. + # Padded zeros yield zero K/V contribution; softmax normalization impact + # is bounded since kv_pad < sp_size << kv_len. + kv_pad = (sp_size - kv_k_full.shape[1] % sp_size) % sp_size + if kv_pad > 0: + kv_k_full = torch.nn.functional.pad(kv_k_full, (0, 0, 0, 0, 0, kv_pad)) + kv_v_full = torch.nn.functional.pad(kv_v_full, (0, 0, 0, 0, 0, kv_pad)) + kv_k = torch.chunk(kv_k_full, sp_size, dim=1)[sp_rank] kv_v = torch.chunk(kv_v_full, sp_size, dim=1)[sp_rank] - + x = xFuserLongContextAttention()( None, query=half(roped_query), @@ -325,10 +345,15 @@ def qkv_fn(x): value=kv_v, window_size=self.window_size ) - + + # Trim Q padding so the output length matches this rank's input length s + # (i.e. the per-rank chunk length produced by upstream torch.chunk(x, sp_size)). + if x.shape[1] > s: + x = x[:, :s] + kv_cache["global_end_index"].fill_(current_end) kv_cache["local_end_index"].fill_(local_end_index) - + # Output projection x = x.flatten(2) x = self.o(x) diff --git a/videox_fun/dist/z_image_xfuser.py b/videox_fun/dist/z_image_xfuser.py index 4ee220ca..15d0bc03 100644 --- a/videox_fun/dist/z_image_xfuser.py +++ b/videox_fun/dist/z_image_xfuser.py @@ -11,10 +11,25 @@ init_distributed_environment, initialize_model_parallel, xFuserLongContextAttention) + +def _pad_heads(tensor, sp_world_size): + """Pad head dimension (dim=2) to be divisible by sp_world_size. + Input shape: [B, S, H, D] -> pad H to next multiple of sp_world_size. + """ + n_heads = tensor.shape[2] + remainder = n_heads % sp_world_size + if remainder == 0: + return tensor, 0 + pad_amount = sp_world_size - remainder + # F.pad pads from last dim backward: (D_left, D_right, H_left, H_right) + tensor = F.pad(tensor, (0, 0, 0, pad_amount), value=0) + return tensor, pad_amount + + class ZMultiGPUsSingleStreamAttnProcessor: """ - Processor for Z-Image single stream attention that adapts the existing Attention class to match the behavior of the - original Z-ImageAttention module. + Processor for Z-Image single stream attention that adapts the existing + Attention class to match the behavior of the original Z-ImageAttention module. """ _attention_backend = None @@ -23,7 +38,8 @@ class ZMultiGPUsSingleStreamAttnProcessor: def __init__(self): if not hasattr(F, "scaled_dot_product_attention"): raise ImportError( - "ZSingleStreamAttnProcessor requires PyTorch 2.0. To use it, please upgrade PyTorch to version 2.0 or higher." + "ZSingleStreamAttnProcessor requires PyTorch 2.0. " + "To use it, please upgrade PyTorch to version 2.0 or higher." ) def __call__( @@ -54,7 +70,7 @@ def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tenso x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2)) freqs_cis = freqs_cis.unsqueeze(2) x_out = torch.view_as_real(x * freqs_cis).flatten(3) - return x_out.type_as(x_in) # todo + return x_out.type_as(x_in) if freqs_cis is not None: query = apply_rotary_emb(query, freqs_cis) @@ -64,7 +80,7 @@ def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tenso dtype = query.dtype query, key = query.to(dtype), key.to(dtype) - # From [batch, seq_len] to [batch, 1, 1, seq_len] -> broadcast to [batch, heads, seq_len, seq_len] + # From [batch, seq_len] to [batch, 1, 1, seq_len] if attention_mask is not None and attention_mask.ndim == 2: attention_mask = attention_mask[:, None, None, :] @@ -72,12 +88,20 @@ def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tenso def half(x): return x if x.dtype in half_dtypes else x.to(torch.bfloat16) + # Pad heads to be divisible by sp_world_size (e.g., 30 -> 32 for 8 GPUs) + sp_world_size = get_sequence_parallel_world_size() + query, head_pad = _pad_heads(query, sp_world_size) + key, _ = _pad_heads(key, sp_world_size) + value, _ = _pad_heads(value, sp_world_size) + hidden_states = xFuserLongContextAttention()( None, half(query), half(key), half(value), dropout_p=0.0, causal=False, ) - # Reshape back + # Trim padded heads back to original count, then reshape + if head_pad > 0: + hidden_states = hidden_states[:, :, :-head_pad, :] hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(dtype) @@ -85,4 +109,4 @@ def half(x): if len(attn.to_out) > 1: # dropout output = attn.to_out[1](output) - return output \ No newline at end of file + return output diff --git a/videox_fun/models/wan_transformer3d_self_forcing.py b/videox_fun/models/wan_transformer3d_self_forcing.py index cee954a3..38a29199 100644 --- a/videox_fun/models/wan_transformer3d_self_forcing.py +++ b/videox_fun/models/wan_transformer3d_self_forcing.py @@ -1061,9 +1061,16 @@ def forward( # Context Parallel: split input across GPUs if self.sp_world_size > 1: - x = torch.chunk(x, self.sp_world_size, dim=1)[self.sp_world_rank] if t.dim() != 1: + # Expand e0 from per-frame [B, F, 6, dim] to per-token + # [B, F*frame_seqlen, 6, dim] so it shares the same SP chunk + # granularity as x. Avoids mismatch when F % sp_size != 0. + F_curr = e0.shape[1] + assert x.shape[1] % F_curr == 0 + frame_seqlen_e0 = x.shape[1] // F_curr + e0 = e0.repeat_interleave(frame_seqlen_e0, dim=1) e0 = torch.chunk(e0, self.sp_world_size, dim=1)[self.sp_world_rank] + x = torch.chunk(x, self.sp_world_size, dim=1)[self.sp_world_rank] # Arguments kwargs = dict( diff --git a/videox_fun/models/z_image_transformer2d.py b/videox_fun/models/z_image_transformer2d.py index 78f22905..fb488abf 100644 --- a/videox_fun/models/z_image_transformer2d.py +++ b/videox_fun/models/z_image_transformer2d.py @@ -863,6 +863,18 @@ def custom_forward(*inputs): # Context Parallel if self.sp_world_size > 1: + # Ensure seq_len is divisible by SEQ_MULTI_OF * sp_world_size for chunking + required_multiple = SEQ_MULTI_OF * self.sp_world_size + current_seq_len = x.shape[1] + remainder = current_seq_len % required_multiple + if remainder != 0: + pad_amount = required_multiple - remainder + x = torch.nn.functional.pad(x, (0, 0, 0, pad_amount), value=0) + x_attn_mask = torch.nn.functional.pad(x_attn_mask, (0, pad_amount), value=False) + x_item_seqlens = [s + pad_amount for s in x_item_seqlens] + if x_freqs_cis is not None: + x_freqs_cis = torch.nn.functional.pad(x_freqs_cis, (0, 0, 0, pad_amount), value=0) + x = torch.chunk(x, self.sp_world_size, dim=1)[self.sp_world_rank] x_item_seqlens = [len(_) for _ in x] diff --git a/videox_fun/pipeline/pipeline_wan2_2_animate.py b/videox_fun/pipeline/pipeline_wan2_2_animate.py index cb81fcd8..340f8d19 100644 --- a/videox_fun/pipeline/pipeline_wan2_2_animate.py +++ b/videox_fun/pipeline/pipeline_wan2_2_animate.py @@ -17,7 +17,10 @@ from diffusers.utils import BaseOutput, logging, replace_example_docstring from diffusers.utils.torch_utils import randn_tensor from diffusers.video_processor import VideoProcessor -from decord import VideoReader +try: + from decord import VideoReader +except ImportError: + from ..data.utils import AVVideoReader as VideoReader from ..models import (AutoencoderKLWan, AutoTokenizer, CLIPModel, WanT5EncoderModel, Wan2_2Transformer3DModel_Animate) diff --git a/videox_fun/reward/hpsv3_predictor.py b/videox_fun/reward/hpsv3_predictor.py index 94ec8eab..ba7a627d 100644 --- a/videox_fun/reward/hpsv3_predictor.py +++ b/videox_fun/reward/hpsv3_predictor.py @@ -918,6 +918,181 @@ def reward(self, image_paths, prompts): return rewards + def prepare_batch_differentiable(self, image_tensors, prompts): + """Prepare batch with differentiable image preprocessing. + + Produces the same output as prepare_batch(), but the pixel_values + tensor retains its computation graph so that gradients can flow back + to the input image_tensors. + + Args: + image_tensors: List[torch.Tensor], each [C, H, W] in [0, 1]. + prompts: List[str], text prompts corresponding to each image. + Returns: + Batch dict with differentiable pixel_values. + """ + # --- Get processor vision config --- + img_proc = self.processor.image_processor + patch_size = img_proc.patch_size # typically 14 + temporal_patch_size = img_proc.temporal_patch_size # typically 2 + merge_size = img_proc.merge_size # typically 2 + image_mean = torch.tensor(img_proc.image_mean, dtype=torch.float32) + image_std = torch.tensor(img_proc.image_std, dtype=torch.float32) + + max_pixels = self.data_config.max_pixels + min_pixels = self.data_config.min_pixels + + # --- Step 1: Non-differentiable path for text tokens & metadata --- + with torch.no_grad(): + ref_pil_images = [] + ref_sizes = [] + for img in image_tensors: + img_det = img.detach().float() + C, H, W = img_det.shape + resized_height, resized_width = smart_resize( + H, W, factor=IMAGE_FACTOR, + min_pixels=min_pixels, max_pixels=max_pixels, + ) + ref_sizes.append((resized_height, resized_width)) + # Resize detached tensor for reference PIL + img_resized = torch.nn.functional.interpolate( + img_det.unsqueeze(0), size=[resized_height, resized_width], + mode="bicubic", antialias=True, align_corners=False, + ).squeeze(0) # [C, rH, rW] + # Convert to PIL for processor tokenization + frame_np = (img_resized.clamp(0, 1).cpu().numpy().transpose(1, 2, 0) * 255).astype('uint8') + ref_pil_images.append(Image.fromarray(frame_np)) + + # Build message list for text tokenization + message_list = [] + for text, pil_img in zip(prompts, ref_pil_images): + out_message = [ + { + "role": "user", + "content": [ + { + "type": "image", + "image": pil_img, + "min_pixels": min_pixels, + "max_pixels": max_pixels, + }, + { + "type": "text", + "text": ( + INSTRUCTION.format(text_prompt=text) + + prompt_with_special_token + if self.use_special_tokens + else prompt_without_special_token + ), + }, + ], + } + ] + message_list.append(out_message) + + image_inputs, _ = process_vision_info(message_list) + + batch_ref = self.processor( + text=self.processor.apply_chat_template(message_list, tokenize=False, add_generation_prompt=True), + images=image_inputs, + padding=True, + return_tensors="pt", + videos_kwargs={"do_rescale": True}, + ) + batch_ref = self._prepare_inputs(batch_ref) + + # --- Step 2: Differentiable image preprocessing --- + # Use the processor's image_grid_thw as the ground truth for resize + # dimensions. The processor's grid determines the number of image + # tokens in input_ids; pixel_values MUST produce the same patch count + # or masked_scatter will fail (too few patches) / rotary pos emb will + # mismatch (too many patches). + processor_grid = batch_ref["image_grid_thw"] # [B, 3] + + all_patches = [] + for idx, (img, (ref_h, ref_w)) in enumerate(zip(image_tensors, ref_sizes)): + # img: [C, H, W] in [0, 1] + img = img.float() + + # Compute target dimensions from the processor's grid to ensure + # patch count == image token count in input_ids. + proc_grid_t = processor_grid[idx, 0].item() + proc_grid_h = processor_grid[idx, 1].item() + proc_grid_w = processor_grid[idx, 2].item() + target_h = proc_grid_h * patch_size + target_w = proc_grid_w * patch_size + + # Differentiable resize to processor-aligned dimensions + img = torch.nn.functional.interpolate( + img.unsqueeze(0), size=[target_h, target_w], + mode="bicubic", antialias=True, align_corners=False, + ).squeeze(0) # [C, rH, rW] + + # Differentiable normalize: (x - mean) / std + mean = image_mean.to(img.device, img.dtype).view(3, 1, 1) + std = image_std.to(img.device, img.dtype).view(3, 1, 1) + img = (img - mean) / std # [C, rH, rW] + + # Temporal pad: repeat to temporal_patch_size frames + # [C, H, W] -> [temporal_patch_size, C, H, W] + img = img.unsqueeze(0).repeat(temporal_patch_size, 1, 1, 1) + + # Patch reshape (matches processor's internal logic) + T, C, rH, rW = img.shape + grid_t = T // temporal_patch_size # = 1 for images + grid_h = rH // patch_size + grid_w = rW // patch_size + + patches = img.reshape( + grid_t, temporal_patch_size, C, + grid_h // merge_size, merge_size, patch_size, + grid_w // merge_size, merge_size, patch_size, + ) + # Transpose to match processor: (0, 3, 6, 4, 7, 2, 1, 5, 8) + patches = patches.permute(0, 3, 6, 4, 7, 2, 1, 5, 8).contiguous() + patches = patches.reshape( + grid_t * grid_h * grid_w, + C * temporal_patch_size * patch_size * patch_size, + ) + all_patches.append(patches) + + # Concatenate patches from all images. .to() preserves grad_fn. + pixel_values = torch.cat(all_patches, dim=0).to(self.device) + + # --- Step 3: Assemble batch with differentiable pixel_values --- + # Use the processor's image_grid_thw directly since pixel_values + # have been resized to match it exactly. This also keeps + # position_ids and rope_deltas consistent with input_ids. + batch = { + "input_ids": batch_ref["input_ids"], + "attention_mask": batch_ref["attention_mask"], + "pixel_values": pixel_values, + "image_grid_thw": processor_grid, + } + # Preserve any extra keys from the processor (e.g. position_ids, rope_deltas) + for k, v in batch_ref.items(): + if k not in batch: + batch[k] = v + + return batch + + def reward_differentiable(self, image_tensors, prompts): + """Differentiable reward computation - returns raw logits with grad_fn. + + Unlike reward(), this does NOT use torch.inference_mode() or + torch.no_grad(). The returned tensor retains its computation graph + so gradients can flow back to the input image_tensors. + + Args: + image_tensors: List[torch.Tensor], each [C, H, W] in [0, 1]. + prompts: List[str], text prompts corresponding to each image. + Returns: + torch.Tensor: shape [B, output_dim] with grad_fn intact. + """ + batch = self.prepare_batch_differentiable(image_tensors, prompts) + logits = self.model(return_dict=True, **batch)["logits"] + return logits + # ========================= Main ========================= if __name__ == "__main__": diff --git a/videox_fun/reward/reward_fn.py b/videox_fun/reward/reward_fn.py index e1c08282..d8643bf0 100644 --- a/videox_fun/reward/reward_fn.py +++ b/videox_fun/reward/reward_fn.py @@ -12,13 +12,35 @@ class BaseReward(ABC): - """An base class for reward models. A custom Reward class must implement two functions below. + """An base class for reward models. A custom Reward Class must implement two functions below. """ + # Whether this reward operates on individual frames (image-level) rather than full videos. + # Image-based rewards (e.g. HPS, MPS, Aesthetic) need frame sampling before scoring. + is_image_reward = False + def __init__(self): """Define your reward model and image transformations (optional) here. """ pass + def to(self, device): + """Move the reward model to the specified device. + + Supports two common patterns: + - self.model (AestheticReward, HPSReward, PickScoreReward, MPSReward) + - self.inferencer.model (HPSv3Reward, VideoAlignReward) + + Subclasses with non-standard model storage should override this method. + """ + if hasattr(self, 'model') and isinstance(self.model, torch.nn.Module): + self.model.to(device) + elif hasattr(self, 'inferencer') and hasattr(self.inferencer, 'model'): + self.inferencer.model.to(device) + if hasattr(self.inferencer, 'device'): + self.inferencer.device = device + self.device = device + return self + @abstractmethod def __call__(self, batch_frames: torch.Tensor, batch_prompt: Optional[list[str]]=None) -> Tuple[torch.Tensor, torch.Tensor]: """Given batch frames with shape `[B, C, T, H, W]` extracted from a list of videos and a list of prompts @@ -36,6 +58,8 @@ class AestheticReward(BaseReward): """Aesthetic Predictor [V2](https://github.com/christophschuhmann/improved-aesthetic-predictor) and [V2.5](https://github.com/discus0434/aesthetic-predictor-v2-5) reward model. """ + is_image_reward = True + def __init__( self, encoder_path="openai/clip-vit-large-patch14", @@ -109,6 +133,8 @@ def get_reward(self, batch_frames: torch.Tensor, batch_prompt: Optional[list[str class HPSReward(BaseReward): """[HPS](https://github.com/tgxs002/HPSv2) v2 and v2.1 reward model. """ + is_image_reward = True + def __init__( self, model_path=None, @@ -205,6 +231,8 @@ def get_reward(self, batch_frames: torch.Tensor, batch_prompt: list[str]) -> tor class PickScoreReward(BaseReward): """[PickScore](https://github.com/yuvalkirstain/PickScore) reward model. """ + is_image_reward = True + def __init__( self, model_path="yuvalkirstain/PickScore_v1", @@ -272,6 +300,8 @@ def get_reward(self, batch_frames: torch.Tensor, batch_prompt: list[str]) -> tor class MPSReward(BaseReward): """[MPS](https://github.com/Kwai-Kolors/MPS) reward model. """ + is_image_reward = True + def __init__( self, model_path=None, @@ -371,6 +401,8 @@ def get_reward( class HPSv3Reward(BaseReward): """[HPSv3](https://github.com/tgxs002/HPSv2) v3 reward model based on Qwen2-VL. """ + is_image_reward = True + def __init__( self, checkpoint_path=None, @@ -379,6 +411,7 @@ def __init__( dtype=torch.float16, max_reward=1, loss_scale=1, + differentiable=False, ): from .hpsv3_predictor import HPSv3RewardInferencer @@ -387,22 +420,58 @@ def __init__( self.dtype = dtype self.max_reward = max_reward self.loss_scale = loss_scale + self.differentiable = differentiable self.inferencer = HPSv3RewardInferencer( checkpoint_path=self.checkpoint_path, device=self.device, + dtype=self.dtype, model_name_or_path=model_name_or_path, ) + # Freeze reward model parameters when using differentiable mode. + # The forward pass still builds the computation graph for input gradients. + if self.differentiable: + self.inferencer.model.requires_grad_(False) + def __call__(self, batch_frames: torch.Tensor, batch_prompt: list[str]) -> Tuple[torch.Tensor, torch.Tensor]: - rewards = self.get_reward(batch_frames, batch_prompt) - print(rewards) + if self.differentiable: + rewards = self.get_reward_differentiable(batch_frames, batch_prompt) + else: + rewards = self.get_reward(batch_frames, batch_prompt) if self.max_reward is None: loss_per_sample = (-1 * rewards) * self.loss_scale else: loss_per_sample = torch.abs(rewards - self.max_reward) * self.loss_scale return loss_per_sample.mean(), rewards.mean() + def get_reward_differentiable(self, batch_frames: torch.Tensor, batch_prompt: list[str]) -> torch.Tensor: + """Differentiable reward computation that preserves grad_fn. + + Gradients flow from the returned scalar rewards back through the reward + model to the input batch_frames. + + Args: + batch_frames: [B, C, T, H, W] tensor in [0, 1]. + batch_prompt: List of B text prompts. + Returns: + torch.Tensor: [B] scalar rewards (mu) with grad_fn. + """ + assert len(batch_frames) == len(batch_prompt) + batch_frames = rearrange(batch_frames, "b c t h w -> t b c h w") + total_rewards = [] + + for frames in batch_frames: + # frames: [B, C, H, W] in [0, 1] + image_tensors = [frame for frame in frames] + logits = self.inferencer.reward_differentiable(image_tensors, batch_prompt) + # logits: [B, output_dim], extract mu (index 0) + reward = logits[:, 0] + total_rewards.append(reward) + + rewards = torch.stack(total_rewards, dim=0).mean(dim=0) + return rewards + @torch.no_grad() def get_reward(self, batch_frames: torch.Tensor, batch_prompt: list[str]) -> torch.Tensor: assert len(batch_frames) == len(batch_prompt) @@ -429,6 +498,8 @@ def get_reward(self, batch_frames: torch.Tensor, batch_prompt: list[str]) -> tor class VideoAlignReward(BaseReward): + is_image_reward = False + def __init__( self, model_path=None, @@ -442,6 +513,8 @@ def __init__( num_frames=None, use_norm=True, return_all_dims=False, + use_legacy_video_io=True, + differentiable=False, ): from .video_align_predictor import VideoVLMRewardInference @@ -450,15 +523,19 @@ def __init__( self.dtype = dtype self.max_reward = max_reward self.loss_scale = loss_scale - self.reward_dim = reward_dim # Which dimension to extract as the scalar reward. + self.reward_dim = reward_dim # Which dimension(s) to extract as the scalar reward. # - "VQ" : Visual Quality (clearness, resolution, brightness, color) # - "MQ" : Motion Quality (consistency, smoothness, completeness) # - "TA" : Text-to-Video Alignment (prompt-content & motion match) # - "Overall": Overall Performance = VQ + MQ + TA (sum of the three) + # - Combinations like "VQ+MQ", "VQ+TA", "MQ+TA" are also supported, + # which sum the specified dimensions. self.fps = fps self.num_frames = num_frames self.use_norm = use_norm self.return_all_dims = return_all_dims # Return all dimensions instead of single reward_dim + self.use_legacy_video_io = use_legacy_video_io # If True, save to temp video then read back (old path) + self.differentiable = differentiable # If True, use differentiable path for backprop self.inferencer = VideoVLMRewardInference( load_from_pretrained=self.model_path, @@ -467,6 +544,11 @@ def __init__( model_name_or_path=model_name_or_path, ) + # Freeze reward model parameters when using differentiable mode. + # The forward pass still builds the computation graph for input gradients. + if self.differentiable: + self.inferencer.model.requires_grad_(False) + def _save_frames_to_temp_video(self, frames: torch.Tensor, fps: float = 8.0) -> str: """Save tensor frames to a temporary video file with lossless encoding. @@ -477,10 +559,11 @@ def _save_frames_to_temp_video(self, frames: torch.Tensor, fps: float = 8.0) -> Returns: Path to the temporary video file """ - import tempfile import os + import tempfile + import av - + # Use /dev/shm (tmpfs, RAM-based) to avoid disk IO, fallback to tempdir shm_dir = "/dev/shm" if os.path.exists(shm_dir) and os.access(shm_dir, os.W_OK): @@ -490,6 +573,7 @@ def _save_frames_to_temp_video(self, frames: torch.Tensor, fps: float = 8.0) -> # Generate unique filename based on frame content hash import hashlib + # Use multiple frames' bytes for robust hashing frame_data = frames.float().cpu().numpy().tobytes() frame_hash = hashlib.md5(frame_data[:10000] + frame_data[-10000:]).hexdigest()[:16] @@ -520,12 +604,18 @@ def _save_frames_to_temp_video(self, frames: torch.Tensor, fps: float = 8.0) -> return temp_video_path def __call__(self, batch_frames: torch.Tensor, batch_prompt: list[str]) -> Tuple[torch.Tensor, torch.Tensor]: - if self.return_all_dims: - rewards_dict = self.get_reward_all_dims(batch_frames, batch_prompt) - # Use Overall for loss computation - rewards = rewards_dict['Overall'] + if self.differentiable: + if self.return_all_dims: + rewards_dict = self.get_reward_all_dims_differentiable(batch_frames, batch_prompt) + rewards = rewards_dict['Overall'] + else: + rewards = self.get_reward_differentiable(batch_frames, batch_prompt) else: - rewards = self.get_reward(batch_frames, batch_prompt) + if self.return_all_dims: + rewards_dict = self.get_reward_all_dims(batch_frames, batch_prompt) + rewards = rewards_dict['Overall'] + else: + rewards = self.get_reward(batch_frames, batch_prompt) if self.max_reward is None: loss_per_sample = (-1 * rewards) * self.loss_scale @@ -533,39 +623,155 @@ def __call__(self, batch_frames: torch.Tensor, batch_prompt: list[str]) -> Tuple loss_per_sample = torch.abs(rewards - self.max_reward) * self.loss_scale return loss_per_sample.mean(), rewards.mean() - @torch.no_grad() - def get_reward(self, batch_frames: torch.Tensor, batch_prompt: list[str]) -> torch.Tensor: - assert len(batch_frames) == len(batch_prompt) - total_rewards = [] + def _get_rewards_legacy(self, batch_frames, batch_prompt): + """Legacy path: save tensors to temp video files, then read back via inferencer.reward().""" temp_video_paths = [] - try: for frames in batch_frames: - # Save frames to temp video frames = rearrange(frames, "c t h w -> t c h w") temp_video_path = self._save_frames_to_temp_video(frames, fps=self.fps) temp_video_paths.append(temp_video_path) - # Get rewards from VideoVLMRewardInference rewards_output = self.inferencer.reward( video_paths=temp_video_paths, prompts=batch_prompt, num_frames=self.num_frames, use_norm=self.use_norm, ) - - for reward_dict in rewards_output: - reward_value = reward_dict[self.reward_dim] - total_rewards.append(torch.tensor(reward_value, device=self.device, dtype=self.dtype)) - - rewards = torch.stack(total_rewards, dim=0) - finally: - # Clean up temporary video files for temp_path in temp_video_paths: if os.path.exists(temp_path): os.remove(temp_path) - + return rewards_output + + def _get_rewards_direct(self, batch_frames, batch_prompt): + """Direct tensor path: pass tensors to inferencer without file I/O.""" + video_tensors = [rearrange(frames, "c t h w -> t c h w") for frames in batch_frames] + rewards_output = self.inferencer.reward_from_tensors( + video_tensors=video_tensors, + prompts=batch_prompt, + num_frames=self.num_frames, + video_fps=self.fps, + use_norm=self.use_norm, + ) + return rewards_output + + def _norm_logits(self, logits): + """Apply per-dimension normalization to raw logits tensor (differentiable). + + Args: + logits: torch.Tensor of shape [B, 3] with columns [VQ, MQ, TA]. + Returns: + Normalized logits tensor of the same shape, with grad_fn preserved. + """ + if self.inferencer.inference_config is None: + return logits + # Cast to float32 for precision parity with the non-differentiable path, + # which normalizes in float64 after .item(). .float() preserves grad_fn. + logits = logits.float() + cfg = self.inferencer.inference_config + mean = torch.tensor( + [cfg['VQ_mean'], cfg['MQ_mean'], cfg['TA_mean']], + device=logits.device, dtype=logits.dtype, + ) + std = torch.tensor( + [cfg['VQ_std'], cfg['MQ_std'], cfg['TA_std']], + device=logits.device, dtype=logits.dtype, + ) + return (logits - mean) / std + + def get_reward_differentiable(self, batch_frames: torch.Tensor, batch_prompt: list[str]) -> torch.Tensor: + """Differentiable reward computation that preserves grad_fn. + + No torch.no_grad() context, no .item() calls, no torch.tensor() wrapping. + Gradients flow from the returned scalar rewards back through the reward + model to the input batch_frames. + + Args: + batch_frames: [B, C, T, H, W] tensor in [0, 1]. + batch_prompt: List of B text prompts. + Returns: + torch.Tensor: [B] scalar rewards with grad_fn. + """ + assert len(batch_frames) == len(batch_prompt) + video_tensors = [rearrange(frames, "c t h w -> t c h w") for frames in batch_frames] + logits = self.inferencer.reward_from_tensors_differentiable( + video_tensors=video_tensors, + prompts=batch_prompt, + num_frames=self.num_frames, + video_fps=self.fps, + ) # [B, 3] with columns [VQ, MQ, TA] + + if self.use_norm: + logits = self._norm_logits(logits) + + # Select the reward dimension(s). + # Supports single dim ("VQ"), "Overall", or combinations like "VQ+MQ". + dim_map = {"VQ": 0, "MQ": 1, "TA": 2} + if self.reward_dim == "Overall": + rewards = logits.sum(dim=-1) # [B] + elif "+" in self.reward_dim: + dims = [d.strip() for d in self.reward_dim.split("+")] + indices = [dim_map[d] for d in dims if d in dim_map] + if len(indices) == 0: + raise ValueError(f"Unknown reward_dim combination: {self.reward_dim}") + rewards = logits[:, indices].sum(dim=-1) # [B] + elif self.reward_dim in dim_map: + rewards = logits[:, dim_map[self.reward_dim]] # [B] + else: + raise ValueError(f"Unknown reward_dim: {self.reward_dim}") + + return rewards + + def get_reward_all_dims_differentiable(self, batch_frames: torch.Tensor, batch_prompt: list[str]) -> dict: + """Differentiable version of get_reward_all_dims. + + Returns: + dict: 'VQ', 'MQ', 'TA', 'Overall' keys, each a [B] tensor with grad_fn. + """ + assert len(batch_frames) == len(batch_prompt) + video_tensors = [rearrange(frames, "c t h w -> t c h w") for frames in batch_frames] + logits = self.inferencer.reward_from_tensors_differentiable( + video_tensors=video_tensors, + prompts=batch_prompt, + num_frames=self.num_frames, + video_fps=self.fps, + ) # [B, 3] + + if self.use_norm: + logits = self._norm_logits(logits) + + return { + 'VQ': logits[:, 0], + 'MQ': logits[:, 1], + 'TA': logits[:, 2], + 'Overall': logits.sum(dim=-1), + } + + @torch.no_grad() + def get_reward(self, batch_frames: torch.Tensor, batch_prompt: list[str]) -> torch.Tensor: + assert len(batch_frames) == len(batch_prompt) + total_rewards = [] + + if self.use_legacy_video_io: + rewards_output = self._get_rewards_legacy(batch_frames, batch_prompt) + else: + rewards_output = self._get_rewards_direct(batch_frames, batch_prompt) + + # Support single dim, "Overall", or combinations like "VQ+MQ" + if "+" in self.reward_dim: + dims = [d.strip() for d in self.reward_dim.split("+")] + else: + dims = [self.reward_dim] + + for reward_dict in rewards_output: + if "+" in self.reward_dim: + reward_value = sum(reward_dict[d] for d in dims) + else: + reward_value = reward_dict[self.reward_dim] + total_rewards.append(torch.tensor(reward_value, device=self.device, dtype=self.dtype)) + + rewards = torch.stack(total_rewards, dim=0) return rewards @torch.no_grad() @@ -576,45 +782,31 @@ def get_reward_all_dims(self, batch_frames: torch.Tensor, batch_prompt: list[str dict: Dictionary with keys 'VQ', 'MQ', 'TA', 'Overall', each containing a tensor of rewards. """ assert len(batch_frames) == len(batch_prompt) - temp_video_paths = [] all_rewards = {'VQ': [], 'MQ': [], 'TA': [], 'Overall': []} - - try: - for frames in batch_frames: - # Save frames to temp video - frames = rearrange(frames, "c t h w -> t c h w") - temp_video_path = self._save_frames_to_temp_video(frames, fps=self.fps) - temp_video_paths.append(temp_video_path) - # Get rewards from VideoVLMRewardInference - rewards_output = self.inferencer.reward( - video_paths=temp_video_paths, - prompts=batch_prompt, - num_frames=self.num_frames, - use_norm=self.use_norm, - ) + if self.use_legacy_video_io: + rewards_output = self._get_rewards_legacy(batch_frames, batch_prompt) + else: + rewards_output = self._get_rewards_direct(batch_frames, batch_prompt) - for reward_dict in rewards_output: - for dim in ['VQ', 'MQ', 'TA', 'Overall']: - reward_value = reward_dict[dim] - all_rewards[dim].append(torch.tensor(reward_value, device=self.device, dtype=self.dtype)) - - # Stack all dimensions - result = {} + for reward_dict in rewards_output: for dim in ['VQ', 'MQ', 'TA', 'Overall']: - result[dim] = torch.stack(all_rewards[dim], dim=0) - - finally: - # Clean up temporary video files - for temp_path in temp_video_paths: - if os.path.exists(temp_path): - os.remove(temp_path) - + reward_value = reward_dict[dim] + all_rewards[dim].append(torch.tensor(reward_value, device=self.device, dtype=self.dtype)) + + # Stack all dimensions + result = {} + for dim in ['VQ', 'MQ', 'TA', 'Overall']: + result[dim] = torch.stack(all_rewards[dim], dim=0) + return result if __name__ == "__main__": import numpy as np - from decord import VideoReader + try: + from decord import VideoReader + except ImportError: + from videox_fun.data.utils import AVVideoReader as VideoReader video_path_list = ["your_video_path_1.mp4", "your_video_path_2.mp4"] prompt_list = ["your_prompt_1", "your_prompt_2"] diff --git a/videox_fun/reward/video_align_predictor.py b/videox_fun/reward/video_align_predictor.py index 42cfa716..f4021da9 100644 --- a/videox_fun/reward/video_align_predictor.py +++ b/videox_fun/reward/video_align_predictor.py @@ -463,9 +463,9 @@ def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR) -> torch.Tensor | l resized_height, resized_width = smart_resize( height, width, factor=image_factor, min_pixels=min_pixels, max_pixels=max_pixels, ) - video = transforms.functional.resize( - video, [resized_height, resized_width], interpolation=InterpolationMode.BICUBIC, antialias=True, - ).float() + video = torch.nn.functional.interpolate( + video, size=[resized_height, resized_width], mode="bicubic", antialias=True, align_corners=False, + ) return video else: assert isinstance(ele["video"], (list, tuple)) @@ -956,7 +956,6 @@ def __init__(self, load_from_pretrained, load_from_pretrained_step=-1, device='c model, checkpoint_step = load_model_from_checkpoint(model, load_from_pretrained, load_from_pretrained_step) model.eval() - self.data_config = data_config self.inference_config = inference_config self.device = device @@ -1079,6 +1078,349 @@ def reward(self, video_paths, prompts, fps=None, num_frames=None, max_pixels=Non return rewards + def prepare_batch_from_tensors(self, video_tensors, prompts, fps=None, num_frames=None, max_pixels=None, video_fps=None): + """Prepare batch from video tensors directly, no file I/O. + Replicates the exact processing of prepare_batch() -> process_vision_info() -> fetch_video(). + + Args: + video_tensors: List[torch.Tensor], each [T, C, H, W] in [0, 1] + prompts: List[str] + fps: float, target sampling fps (for smart_nframes). If None, uses data_config.fps. + num_frames: int, if set, subsample to this many frames (after round_by_factor). + If None, uses fps-based sampling. + max_pixels: int, max pixels per frame for resize. If None, uses data_config.max_frame_pixels. + video_fps: float, intrinsic fps of the input video tensors. Required when num_frames is None. + """ + fps = self.data_config.fps if fps is None else fps + num_frames = self.data_config.num_frames if num_frames is None else num_frames + max_pixels = self.data_config.max_frame_pixels if max_pixels is None else max_pixels + + video_inputs = [] + for video in video_tensors: + # video: [T, C, H, W] in [0, 1] -> scale to [0, 255] + video = (video.float() * 255.0) + total_frames = video.shape[0] + + # Replicate smart_nframes logic exactly + if num_frames is not None: + nframes = round_by_factor(num_frames, FRAME_FACTOR) + else: + assert video_fps is not None, "video_fps is required when num_frames is None" + min_frames = ceil_by_factor(FPS_MIN_FRAMES, FRAME_FACTOR) + max_frames = floor_by_factor(min(FPS_MAX_FRAMES, total_frames), FRAME_FACTOR) + nframes = total_frames / video_fps * fps + nframes = min(max(nframes, min_frames), max_frames) + nframes = round_by_factor(nframes, FRAME_FACTOR) + + if nframes > total_frames: + nframes = total_frames + + # Always apply linspace subsampling (same as video reader backends) + idx = torch.linspace(0, total_frames - 1, nframes).round().long() + video = video[idx] + + # Resize: replicate fetch_video logic exactly + # When max_pixels is set from config, it overrides the computed value + # (matches: max_pixels = ele.get("max_pixels", computed) in fetch_video) + nframes, _, height, width = video.shape + min_pixels = VIDEO_MIN_PIXELS + if max_pixels is not None: + effective_max_pixels = max_pixels + else: + total_pixels = VIDEO_TOTAL_PIXELS + effective_max_pixels = max( + min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), + int(min_pixels * 1.05) + ) + resized_height, resized_width = smart_resize( + height, width, factor=IMAGE_FACTOR, + min_pixels=min_pixels, max_pixels=effective_max_pixels + ) + video = torch.nn.functional.interpolate( + video, size=[resized_height, resized_width], + mode="bicubic", antialias=True, align_corners=False, + ) + video_inputs.append(video) + + # Build text with chat template (same structure as prepare_batch) + chat_data = [ + [ + { + "role": "user", + "content": [ + {"type": "video", "video": "placeholder"}, + {"type": "text", "text": build_prompt(prompt, self.data_config.eval_dim, self.data_config.prompt_template_type)}, + ], + }, + ] + for prompt in prompts + ] + + batch = self.processor( + text=self.processor.apply_chat_template(chat_data, tokenize=False, add_generation_prompt=True), + images=None, + videos=video_inputs, + padding=True, + return_tensors="pt", + videos_kwargs={"do_rescale": True, "do_resize": False}, + ) + batch = self._prepare_inputs(batch) + return batch + + def reward_from_tensors(self, video_tensors, prompts, fps=None, num_frames=None, max_pixels=None, video_fps=None, use_norm=True): + """Compute rewards directly from video tensors without file I/O. + + Args: + video_tensors: List[torch.Tensor], each [T, C, H, W] in [0, 1] + prompts: List[str] + fps: float, target sampling fps. If None, uses data_config.fps. + num_frames: int, if set, subsample to this many frames. If None, uses fps-based sampling. + max_pixels: int, max pixels per frame for resize. + video_fps: float, intrinsic fps of the input video tensors. Required when num_frames is None. + use_norm: bool, whether to rescale the output rewards + Returns: + List[dict]: rewards with VQ, MQ, TA, Overall for each video. + """ + batch = self.prepare_batch_from_tensors(video_tensors, prompts, fps, num_frames, max_pixels, video_fps) + rewards = self.model( + return_dict=True, + **batch + )["logits"] + + rewards = [{'VQ': reward[0].item(), 'MQ': reward[1].item(), 'TA': reward[2].item()} for reward in rewards] + for i in range(len(rewards)): + if use_norm: + rewards[i] = self._norm(rewards[i]) + rewards[i]['Overall'] = rewards[i]['VQ'] + rewards[i]['MQ'] + rewards[i]['TA'] + + return rewards + + def prepare_batch_from_tensors_differentiable(self, video_tensors, prompts, fps=None, num_frames=None, max_pixels=None, video_fps=None): + """Prepare batch with differentiable video preprocessing. + + Produces the same output as prepare_batch_from_tensors(), but the + pixel_values_videos tensor retains its computation graph so that + gradients can flow back to the input video_tensors. + + The text tokenization (input_ids, attention_mask) and video_grid_thw + are computed via the standard non-differentiable path (they do not + require gradients). + + Args: + video_tensors: List[torch.Tensor], each [T, C, H, W] in [0, 1] + prompts: List[str] + fps: float, target sampling fps. If None, uses data_config.fps. + num_frames: int, if set, subsample to this many frames. + max_pixels: int, max pixels per frame for resize. + video_fps: float, intrinsic fps of the input video tensors. + """ + # --- Get processor vision config --- + img_proc = self.processor.image_processor + patch_size = img_proc.patch_size # typically 14 + temporal_patch_size = img_proc.temporal_patch_size # typically 2 + merge_size = img_proc.merge_size # typically 2 + image_mean = torch.tensor(img_proc.image_mean, dtype=torch.float32) + image_std = torch.tensor(img_proc.image_std, dtype=torch.float32) + + # --- Resolve default parameters --- + fps_val = self.data_config.fps if fps is None else fps + num_frames_val = self.data_config.num_frames if num_frames is None else num_frames + max_pixels_val = self.data_config.max_frame_pixels if max_pixels is None else max_pixels + + # --- Step 1: Non-differentiable path for text tokens & metadata --- + # We process detached & resized videos in [0, 1] range directly through the + # processor with do_rescale=False to obtain input_ids, attention_mask, and + # video_grid_thw. This avoids the *255 → processor resize incompatibility + # with newer transformers versions. + with torch.no_grad(): + ref_video_inputs = [] + for video in video_tensors: + video_det = video.detach().float() + total_frames = video_det.shape[0] + + # Frame sampling (same logic as Step 2) + if num_frames_val is not None: + nframes = round_by_factor(int(num_frames_val), FRAME_FACTOR) + else: + assert video_fps is not None, "video_fps is required when num_frames is None" + _min_frames = ceil_by_factor(FPS_MIN_FRAMES, FRAME_FACTOR) + _max_frames = floor_by_factor(min(FPS_MAX_FRAMES, total_frames), FRAME_FACTOR) + nframes = total_frames / video_fps * fps_val + nframes = min(max(nframes, _min_frames), _max_frames) + nframes = round_by_factor(int(nframes), FRAME_FACTOR) + # Ensure at least FRAME_FACTOR frames (temporal_patch_size minimum) + nframes = max(nframes, FRAME_FACTOR) + if nframes > total_frames: + nframes = total_frames + + idx = torch.linspace(0, total_frames - 1, nframes).round().long() + video_det = video_det[idx] + + # Resize + _nf, _C, _h, _w = video_det.shape + _min_pix = VIDEO_MIN_PIXELS + if max_pixels_val is not None: + _eff_max = max_pixels_val + else: + _eff_max = max( + min(VIDEO_MAX_PIXELS, VIDEO_TOTAL_PIXELS / _nf * FRAME_FACTOR), + int(_min_pix * 1.05) + ) + _rh, _rw = smart_resize(_h, _w, factor=IMAGE_FACTOR, min_pixels=_min_pix, max_pixels=_eff_max) + # Move to CPU for bicubic+antialias interpolation to avoid + # "CUDA error: invalid configuration argument" under low VRAM. + _device = video_det.device + video_det = torch.nn.functional.interpolate( + video_det.cpu(), size=[_rh, _rw], + mode="bicubic", antialias=True, align_corners=False, + ).to(_device) + ref_video_inputs.append(video_det) # [nframes, C, H, W] in [0, 1] + + # Build chat template for text tokenization + chat_data = [ + [ + { + "role": "user", + "content": [ + {"type": "video", "video": "placeholder"}, + {"type": "text", "text": build_prompt(prompt, self.data_config.eval_dim, self.data_config.prompt_template_type)}, + ], + }, + ] + for prompt in prompts + ] + + # Pass [0, 1] videos to processor with do_rescale=False (no resize needed) + batch_ref = self.processor( + text=self.processor.apply_chat_template(chat_data, tokenize=False, add_generation_prompt=True), + images=None, + videos=ref_video_inputs, + padding=True, + return_tensors="pt", + videos_kwargs={"do_rescale": False, "do_resize": False}, + ) + batch_ref = self._prepare_inputs(batch_ref) + + # --- Step 2: Differentiable video preprocessing --- + all_patches = [] + for video in video_tensors: + # video: [T, C, H, W] in [0, 1] + video = video.float() + total_frames = video.shape[0] + + # Frame sampling (index op preserves grad for selected frames) + if num_frames_val is not None: + nframes = round_by_factor(int(num_frames_val), FRAME_FACTOR) + else: + assert video_fps is not None, "video_fps is required when num_frames is None" + min_frames = ceil_by_factor(FPS_MIN_FRAMES, FRAME_FACTOR) + max_frames = floor_by_factor(min(FPS_MAX_FRAMES, total_frames), FRAME_FACTOR) + nframes = total_frames / video_fps * fps_val + nframes = min(max(nframes, min_frames), max_frames) + nframes = round_by_factor(int(nframes), FRAME_FACTOR) + + # Ensure at least FRAME_FACTOR frames (temporal_patch_size minimum) + nframes = max(nframes, FRAME_FACTOR) + if nframes > total_frames: + nframes = total_frames + + idx = torch.linspace(0, total_frames - 1, nframes).round().long() + video = video[idx] # [nframes, C, H, W] - grad preserved + + # Differentiable resize + nframes_now, C, height, width = video.shape + min_pixels = VIDEO_MIN_PIXELS + if max_pixels_val is not None: + effective_max_pixels = max_pixels_val + else: + total_pixels = VIDEO_TOTAL_PIXELS + effective_max_pixels = max( + min(VIDEO_MAX_PIXELS, total_pixels / nframes_now * FRAME_FACTOR), + int(min_pixels * 1.05) + ) + resized_height, resized_width = smart_resize( + height, width, factor=IMAGE_FACTOR, + min_pixels=min_pixels, max_pixels=effective_max_pixels + ) + video = torch.nn.functional.interpolate( + video, size=[resized_height, resized_width], + mode="bicubic", antialias=True, align_corners=False, + ) # differentiable + + # Differentiable normalize: (x - mean) / std + # Input is [0, 1] range. Processor would do rescale(1/255) on [0,255] input + # which gives [0, 1], then normalize. We skip the roundtrip. + mean = image_mean.to(video.device, video.dtype).view(1, 3, 1, 1) + std = image_std.to(video.device, video.dtype).view(1, 3, 1, 1) + video = (video - mean) / std # differentiable + + # Temporal padding (replicate last frame if not divisible) + T = video.shape[0] + if T % temporal_patch_size != 0: + pad_count = temporal_patch_size - (T % temporal_patch_size) + video = torch.cat([video, video[-1:].expand(pad_count, -1, -1, -1)], dim=0) + + # Differentiable patch reshape (matches processor's numpy reshape exactly) + T, C, H, W = video.shape + grid_t = T // temporal_patch_size + grid_h = H // patch_size + grid_w = W // patch_size + + patches = video.reshape( + grid_t, temporal_patch_size, C, + grid_h // merge_size, merge_size, patch_size, + grid_w // merge_size, merge_size, patch_size, + ) + # Transpose to match processor: (0, 3, 6, 4, 7, 2, 1, 5, 8) + patches = patches.permute(0, 3, 6, 4, 7, 2, 1, 5, 8).contiguous() + patches = patches.reshape( + grid_t * grid_h * grid_w, + C * temporal_patch_size * patch_size * patch_size, + ) + all_patches.append(patches) + + # Concatenate patches from all videos and move to model device. + # .to() is differentiable and preserves grad_fn. + pixel_values_videos = torch.cat(all_patches, dim=0).to(self.device) + + # --- Step 3: Assemble batch with differentiable pixel_values_videos --- + batch = { + "input_ids": batch_ref["input_ids"], + "attention_mask": batch_ref["attention_mask"], + "pixel_values_videos": pixel_values_videos, + "video_grid_thw": batch_ref["video_grid_thw"], + } + # Preserve any extra keys (e.g. position_ids, rope_deltas) + for k, v in batch_ref.items(): + if k not in batch: + batch[k] = v + + return batch + + def reward_from_tensors_differentiable(self, video_tensors, prompts, fps=None, num_frames=None, max_pixels=None, video_fps=None): + """Differentiable reward computation - returns raw logits tensor with grad_fn. + + Unlike reward_from_tensors(), this does NOT call .item() and does NOT + wrap results in Python dicts. The returned tensor retains its computation + graph so gradients can flow back to the input video_tensors. + + Args: + video_tensors: List[torch.Tensor], each [T, C, H, W] in [0, 1] + prompts: List[str] + fps: float, target sampling fps. + num_frames: int, if set, subsample to this many frames. + max_pixels: int, max pixels per frame for resize. + video_fps: float, intrinsic fps of the input video tensors. + Returns: + torch.Tensor: shape [B, output_dim] with grad_fn intact. + Columns are [VQ, MQ, TA] (or as configured by output_dim). + """ + batch = self.prepare_batch_from_tensors_differentiable( + video_tensors, prompts, fps, num_frames, max_pixels, video_fps, + ) + logits = self.model(return_dict=True, **batch)["logits"] + return logits + # ========================= Main ========================= diff --git a/videox_fun/video_caption/utils/video_utils.py b/videox_fun/video_caption/utils/video_utils.py index bc9cc6c2..50835550 100644 --- a/videox_fun/video_caption/utils/video_utils.py +++ b/videox_fun/video_caption/utils/video_utils.py @@ -6,7 +6,10 @@ from typing import List, Optional, Tuple import numpy as np -from decord import VideoReader +try: + from decord import VideoReader +except ImportError: + from ...data.utils import AVVideoReader as VideoReader from PIL import Image