diff --git a/.gitignore b/.gitignore index 5627e45..3840acc 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 3d24f8d..74f18df 100644 --- a/examples/wan2.2/predict_animate.py +++ b/examples/wan2.2/predict_animate.py @@ -314,8 +314,8 @@ 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 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 @@ with torch.no_grad(): 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 6fd2a07..d883a57 100644 --- a/examples/wan2.2/predict_i2v.py +++ b/examples/wan2.2/predict_i2v.py @@ -136,13 +136,15 @@ transformer = Wan2_2Transformer3DModel.from_pretrained( 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 @@ if transformer_path is not None: 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 @@ pipeline = Wan2_2I2VPipeline( 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 ulysses_degree > 1 or ring_degree > 1: 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 @@ if coefficients is not None: 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 @@ with torch.no_grad(): 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 @@ with torch.no_grad(): 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 45c8209..d7dcf2d 100644 --- a/examples/wan2.2/predict_s2v.py +++ b/examples/wan2.2/predict_s2v.py @@ -318,8 +318,8 @@ 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 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 @@ with torch.no_grad(): 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 c64d73a..2c2d238 100755 --- a/examples/wan2.2/predict_t2v.py +++ b/examples/wan2.2/predict_t2v.py @@ -132,13 +132,15 @@ transformer = Wan2_2Transformer3DModel.from_pretrained( 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 @@ if transformer_path is not None: 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 @@ pipeline = Wan2_2Pipeline( 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 ulysses_degree > 1 or ring_degree > 1: 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 @@ if coefficients is not None: 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 @@ with torch.no_grad(): 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 @@ with torch.no_grad(): 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 36bdbfd..380d4d2 100755 --- a/examples/wan2.2/predict_ti2v.py +++ b/examples/wan2.2/predict_ti2v.py @@ -303,8 +303,8 @@ 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 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 @@ with torch.no_grad(): 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 161e083..d6bba17 100644 --- a/examples/wan2.2_fun/predict_i2v.py +++ b/examples/wan2.2_fun/predict_i2v.py @@ -305,8 +305,8 @@ 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 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 @@ with torch.no_grad(): 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 339565f..9922d46 100644 --- a/examples/wan2.2_fun/predict_i2v_5b.py +++ b/examples/wan2.2_fun/predict_i2v_5b.py @@ -307,8 +307,8 @@ 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 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 @@ with torch.no_grad(): 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 9f5efeb..c68e0f5 100644 --- a/examples/wan2.2_fun/predict_t2v.py +++ b/examples/wan2.2_fun/predict_t2v.py @@ -135,13 +135,15 @@ transformer = Wan2_2Transformer3DModel.from_pretrained( 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 @@ if transformer_path is not None: 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 @@ pipeline = Wan2_2FunInpaintPipeline( 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 ulysses_degree > 1 or ring_degree > 1: 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 @@ if coefficients is not None: 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 @@ with torch.no_grad(): 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 @@ with torch.no_grad(): 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 6f7290b..c4aa514 100644 --- a/examples/wan2.2_fun/predict_t2v_5b.py +++ b/examples/wan2.2_fun/predict_t2v_5b.py @@ -301,8 +301,8 @@ 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 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 @@ with torch.no_grad(): 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 63cd841..0d076b6 100644 --- a/examples/wan2.2_fun/predict_v2v_control.py +++ b/examples/wan2.2_fun/predict_v2v_control.py @@ -318,8 +318,8 @@ 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 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 @@ with torch.no_grad(): 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 8cf222f..7c78666 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 @@ 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 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 @@ with torch.no_grad(): 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 666e0a4..3ec8e52 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 @@ 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 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 @@ with torch.no_grad(): 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 a50173b..ee36ab0 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 @@ 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 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 @@ with torch.no_grad(): 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 eaf088d..43906df 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 @@ 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 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 @@ with torch.no_grad(): 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 5ecd1c7..7d1c2c3 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 @@ 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 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 @@ with torch.no_grad(): 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 1abf301..fa18f3c 100644 --- a/examples/wan2.2_vace_fun/predict_i2v.py +++ b/examples/wan2.2_vace_fun/predict_i2v.py @@ -322,8 +322,8 @@ 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 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 @@ with torch.no_grad(): 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 5d90aa3..9ff65a5 100644 --- a/examples/wan2.2_vace_fun/predict_s2v.py +++ b/examples/wan2.2_vace_fun/predict_s2v.py @@ -322,8 +322,8 @@ 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 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 @@ with torch.no_grad(): 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 a063e4f..5a22716 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 @@ 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 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 @@ with torch.no_grad(): 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 b57099b..e9a557f 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 @@ 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 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 @@ with torch.no_grad(): 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 dd28d98..b25eb5b 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 @@ 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 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 @@ with torch.no_grad(): 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 b995226..fd7f88b 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 babc601..849c681 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 0131aca..a67936c 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 e4a6c5f..f41dc79 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 3d6f772..67628d3 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 6855d4f..fcf9efa 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 eecc940..2d5d22c 100755 --- a/scripts/cogvideox_fun/train_reward_lora.py +++ b/scripts/cogvideox_fun/train_reward_lora.py @@ -38,7 +38,10 @@ from accelerate import Accelerator 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 9703a20..bcc0d0e 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 db12d30..4ebf0cb 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 8552f96..526dc9d 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 9afb731..0283d05 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 ab8cb60..2ef47d4 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 b880a64..4f717f2 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 08b966c..4f7efe6 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 0483742..c32cd9c 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 1d6c5ef..3498e51 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 99f5c0b..e2eb5f5 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 ff47968..fa47dde 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 1f4738c..14f19e3 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 514a8a7..f1f2ad4 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 c27d067..19e8c1d 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 cbaef25..64e7684 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 c29848d..74b3a3d 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 661f7a4..691539a 100644 --- a/scripts/get_rewards_1video.py +++ b/scripts/get_rewards_1video.py @@ -77,7 +77,10 @@ import sys 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 e31445e..3498bc3 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 c375c7e..fdf6dca 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 47393fa..c43393d 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 46b7ec0..fe816bb 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 3eaf964..225cda2 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 6d8e04a..3c996e1 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 279ca31..fdd9545 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 11bedcd..4d65008 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 6d517a7..816c01e 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 544b073..03514fe 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 cf1a7a3..4f3354e 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 8042567..7824a08 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 924c328..c68cc25 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 1fc3f05..2befaa4 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 6ffa6cf..d7ef5a1 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 965de39..8708ae6 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 1b5d7a4..915ae97 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 0d96fb3..80103b7 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/longcatvideo/train.py b/scripts/longcatvideo/train.py index 8fe9e03..a168d58 100644 --- a/scripts/longcatvideo/train.py +++ b/scripts/longcatvideo/train.py @@ -985,11 +985,6 @@ def main(): if args.gradient_checkpointing: transformer3d.enable_gradient_checkpointing() - elif args.selective_ac > 0: - from videox_fun.models.wan_transformer3d import WanAttentionBlock - from videox_fun.utils.ac_handle import apply_checkpointing, partial - apply_selective_ac = partial(apply_checkpointing, block=WanAttentionBlock) - apply_selective_ac(transformer3d, p=args.selective_ac) # Enable TF32 for faster training on Ampere GPUs, # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices diff --git a/scripts/longcatvideo/train_avatar.py b/scripts/longcatvideo/train_avatar.py index ae4011c..cf54481 100644 --- a/scripts/longcatvideo/train_avatar.py +++ b/scripts/longcatvideo/train_avatar.py @@ -1027,11 +1027,6 @@ def main(): if args.gradient_checkpointing: transformer3d.enable_gradient_checkpointing() - elif args.selective_ac > 0: - from videox_fun.models.wan_transformer3d import WanAttentionBlock - from videox_fun.utils.ac_handle import apply_checkpointing, partial - apply_selective_ac = partial(apply_checkpointing, block=WanAttentionBlock) - apply_selective_ac(transformer3d, p=args.selective_ac) # Enable TF32 for faster training on Ampere GPUs, # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices diff --git a/scripts/ltx2.3/README_TRAIN.md b/scripts/ltx2.3/README_TRAIN.md index 5b425d1..eb25bcd 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 dce5006..b372f0f 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 6b847af..0642098 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 d3d9d5e..12722fc 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 cbeab9c..8decc46 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 3701f98..6011a8d 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 4f4e076..bda37ff 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 5a870e5..83b9e81 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 a829082..477bc34 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 main(): 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 6e9cce0..17f5766 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 7c0fc15..97534e2 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 282fdcf..6a54ff5 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 2a5e0d4..47af20a 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 ac1674f..672c9ba 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 f274b76..9dd08d2 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 db7dc15..643519d 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 3414401..da3e176 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 65dfb0e..f915483 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 c78bccf..40620b3 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 e4dddab..2b394ac 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 c26d785..3d5f764 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 18816b8..2f72db5 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 ec6617b..1a2a0eb 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 cc2261e..5b2644a 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 ed7e1ec..b0c32b2 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 ec68899..e8d71aa 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 main(): 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 main(): 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 main(): 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 main(): ) 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 main(): 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 main(): 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 3941d3b..3a6289f 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 3450a07..eb87cd4 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 c8e89f9..911c2ba 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 ac608e5..d98fed6 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 4366b29..fa8815f 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 b167ead..664ccd9 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 148c94e..d80d0c9 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 5c98724..8458f38 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 ebff90e..fc18f98 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.py b/scripts/wan2.1/train.py index 7c705a3..d9eff84 100755 --- a/scripts/wan2.1/train.py +++ b/scripts/wan2.1/train.py @@ -1047,11 +1047,6 @@ def main(): if args.gradient_checkpointing: transformer3d.enable_gradient_checkpointing() - elif args.selective_ac > 0: - from videox_fun.models.wan_transformer3d import WanAttentionBlock - from videox_fun.utils.ac_handle import apply_checkpointing, partial - apply_selective_ac = partial(apply_checkpointing, block=WanAttentionBlock) - apply_selective_ac(transformer3d, p=args.selective_ac) # Enable TF32 for faster training on Ampere GPUs, # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices diff --git a/scripts/wan2.1/train_distill.py b/scripts/wan2.1/train_distill.py index 17ee068..ebe57c3 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 main(): 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 main(): 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 main(): 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 main(): ) 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 main(): 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 main(): 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 dbe4524..ffd3518 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 e4ae846..e48ab06 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 main(): 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 main(): ) 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 main(): 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 main(): 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 7801835..8322a04 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 af22e83..c801896 100755 --- a/scripts/wan2.1/train_reward_lora.py +++ b/scripts/wan2.1/train_reward_lora.py @@ -38,7 +38,10 @@ from accelerate import Accelerator 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 842ec1e..8dde7ed 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 06e35d6..4d4c9f2 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 050d5b4..41ff30b 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 bff96ce..bd2558b 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 4795f50..99aa1a4 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 c899cf5..b99ce17 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 d768410..42086e6 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 83f528a..3f20ad2 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 4bae97c..9e18bcf 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 import Accelerator 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 c8f6b29..d434b76 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 053be37..305a831 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 7828f1f..800b6b4 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 5e3970c..b1310c0 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 d56eca5..9ec8d31 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): ) 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): # 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 main(): 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 27e2d0e..b46f570 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 63df464..c3d049a 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 957fa2f..33aa63e 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 257c7c9..7add18b 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 6d81810..868cac3 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 34c5b01..1a96096 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 c9bdfe9..54fbf58 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 0494174..6c8b1d5 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 3a89981..ec875e2 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 02487aa..bc54457 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 047a555..478e1ac 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 d68ac9f..c000106 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 aaec812..0d4bb6d 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 ebf56ad..9a8370c 100644 --- a/scripts/wan2.2/train.py +++ b/scripts/wan2.2/train.py @@ -1775,12 +1775,12 @@ def main(): 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 367f994..5cb85f0 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): # 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 main(): 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 main(): 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 main(): ) 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): ) 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 8c91e4e..fb1c5c2 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 a7d5d93..3a2d4da 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 main(): 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 main(): 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 main(): 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 main(): # 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 main(): 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 main(): 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 main(): ) 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): 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 main(): ) 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 afba111..e18b178 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 2fb1015..32b5115 100755 --- a/scripts/wan2.2/train_lora.py +++ b/scripts/wan2.2/train_lora.py @@ -1821,12 +1821,12 @@ def main(): 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 db13c53..1477529 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 74fda38..6fed232 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 13150c0..7fa7abc 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 69f048a..619a3f8 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 7edd786..ecc3ab5 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 e8fba28..84274da 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 ad5f07e..ef7c9db 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 66120be..7ddb250 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 0000000..247c610 --- /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 0000000..b34abd7 --- /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 d184fe4..f53b31c 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 import Accelerator 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 4f5de63..e713ea0 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 a68f13d..2a84e8e 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 7364f38..4e404db 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 8016981..c15ee2e 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 63db974..703e63c 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 633816c..946ebd4 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 82a6090..ff09ad4 100644 --- a/scripts/z_image/train_distill.py +++ b/scripts/z_image/train_distill.py @@ -1033,8 +1033,8 @@ def main(): 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 main(): 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 main(): 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 main(): ) 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 6631c35..233c4b5 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 f05aedd..e7d5847 100644 --- a/scripts/z_image/train_distill_lora.py +++ b/scripts/z_image/train_distill_lora.py @@ -1043,7 +1043,7 @@ def main(): 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 main(): ) 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 7469e70..db3d4a4 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 77c8227..1937ac7 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 902e3ec..4665a52 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 0afeb72..d1c3386 100644 --- a/scripts/z_image_fun/train_control_distill.py +++ b/scripts/z_image_fun/train_control_distill.py @@ -1102,8 +1102,8 @@ def main(): 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 main(): 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 main(): 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 main(): ) 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 eb90723..3ee703c 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/__init__.py b/videox_fun/__init__.py index e69de29..2eac33a 100644 --- a/videox_fun/__init__.py +++ b/videox_fun/__init__.py @@ -0,0 +1,4 @@ +import importlib.util + +if importlib.util.find_spec("paifuser") is not None: + import paifuser diff --git a/videox_fun/data/dataset_image_video.py b/videox_fun/data/dataset_image_video.py index 3f6b116..c3bc3b6 100755 --- a/videox_fun/data/dataset_image_video.py +++ b/videox_fun/data/dataset_image_video.py @@ -14,7 +14,6 @@ import numpy as np 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 safetensors.torch import load_file 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 df4b103..9bd0193 100644 --- a/videox_fun/data/dataset_video.py +++ b/videox_fun/data/dataset_video.py @@ -9,12 +9,16 @@ import librosa 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 5891151..ab2d8d8 100644 --- a/videox_fun/data/utils.py +++ b/videox_fun/data/utils.py @@ -6,13 +6,81 @@ from contextlib import contextmanager 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/__init__.py b/videox_fun/dist/__init__.py index 79ad8a2..ac22716 100755 --- a/videox_fun/dist/__init__.py +++ b/videox_fun/dist/__init__.py @@ -1,5 +1,3 @@ -import importlib.util - from .cogvideox_xfuser import CogVideoXMultiGPUsAttnProcessor2_0 from .ernie_image_xfuser import ErnieImageMultiGPUsAttnProcessor from .flashhead_xfuser import usp_attn_flashhead_forward @@ -25,59 +23,3 @@ from .qwen_xfuser import QwenImageMultiGPUsAttnProcessor2_0 from .wan_xfuser import (usp_attn_forward, usp_attn_s2v_forward, usp_attn_self_forcing_forward) from .z_image_xfuser import ZMultiGPUsSingleStreamAttnProcessor - -# The pai_fuser is an internally developed acceleration package, which can be used on PAI. -if importlib.util.find_spec("paifuser") is not None: - # --------------------------------------------------------------- # - # The simple_wrapper is used to solve the problem - # about conflicts between cython and torch.compile - # --------------------------------------------------------------- # - def simple_wrapper(func): - def inner(*args, **kwargs): - return func(*args, **kwargs) - return inner - - # --------------------------------------------------------------- # - # Sparse Attention Kernel - # --------------------------------------------------------------- # - from paifuser.models import parallel_magvit_vae - from paifuser.ops import wan_usp_sparse_attention_wrapper - - from . import wan_xfuser - - # --------------------------------------------------------------- # - # Sparse Attention - # --------------------------------------------------------------- # - usp_sparse_attn_wrap_forward = simple_wrapper(wan_usp_sparse_attention_wrapper()(wan_xfuser.usp_attn_forward)) - wan_xfuser.usp_attn_forward = usp_sparse_attn_wrap_forward - usp_attn_forward = usp_sparse_attn_wrap_forward - print("Import PAI VAE Turbo and Sparse Attention") - - # --------------------------------------------------------------- # - # Fast Rope Kernel - # --------------------------------------------------------------- # - import types - - import torch - from paifuser.ops import (ENABLE_KERNEL, usp_fast_rope_apply_qk, - usp_rope_apply_real_qk) - - def deepcopy_function(f): - return types.FunctionType(f.__code__, f.__globals__, name=f.__name__, argdefs=f.__defaults__,closure=f.__closure__) - - local_rope_apply_qk = deepcopy_function(wan_xfuser.rope_apply_qk) - - if ENABLE_KERNEL: - def adaptive_fast_usp_rope_apply_qk(q, k, grid_sizes, freqs): - if torch.is_grad_enabled(): - return local_rope_apply_qk(q, k, grid_sizes, freqs) - else: - return usp_fast_rope_apply_qk(q, k, grid_sizes, freqs) - - else: - def adaptive_fast_usp_rope_apply_qk(q, k, grid_sizes, freqs): - return usp_rope_apply_real_qk(q, k, grid_sizes, freqs) - - wan_xfuser.rope_apply_qk = adaptive_fast_usp_rope_apply_qk - rope_apply_qk = adaptive_fast_usp_rope_apply_qk - print("Import PAI Fast rope") \ No newline at end of file diff --git a/videox_fun/dist/qwen_xfuser.py b/videox_fun/dist/qwen_xfuser.py index 55bb32a..9d02dca 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 +from typing import Optional, Tuple, Union -import numpy as np 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 6c2358c..dbf8ad8 100755 --- a/videox_fun/dist/wan_xfuser.py +++ b/videox_fun/dist/wan_xfuser.py @@ -308,16 +308,36 @@ def usp_attn_self_forcing_forward( 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 usp_attn_self_forcing_forward( 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 4ee220c..15d0bc0 100644 --- a/videox_fun/dist/z_image_xfuser.py +++ b/videox_fun/dist/z_image_xfuser.py @@ -11,10 +11,25 @@ from .fuser import (get_sequence_parallel_rank, 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 @@ class ZMultiGPUsSingleStreamAttnProcessor: 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 @@ class ZMultiGPUsSingleStreamAttnProcessor: 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 @@ class ZMultiGPUsSingleStreamAttnProcessor: 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 @@ class ZMultiGPUsSingleStreamAttnProcessor: 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/__init__.py b/videox_fun/models/__init__.py index 4826481..98e76a7 100755 --- a/videox_fun/models/__init__.py +++ b/videox_fun/models/__init__.py @@ -1,5 +1,3 @@ -import importlib.util - from diffusers import AutoencoderKL from transformers import (AutoProcessor, AutoTokenizer, CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, @@ -92,89 +90,3 @@ from .wan_vae import AutoencoderKLWan, AutoencoderKLWan_ from .wan_vae3_8 import AutoencoderKLWan2_2_, AutoencoderKLWan3_8 from .z_image_transformer2d import ZImageTransformer2DModel from .z_image_transformer2d_control import ZImageControlTransformer2DModel - -# The pai_fuser is an internally developed acceleration package, which can be used on PAI. -if importlib.util.find_spec("paifuser") is not None: - # --------------------------------------------------------------- # - # The simple_wrapper is used to solve the problem - # about conflicts between cython and torch.compile - # --------------------------------------------------------------- # - def simple_wrapper(func): - def inner(*args, **kwargs): - return func(*args, **kwargs) - return inner - - # --------------------------------------------------------------- # - # VAE Parallel Kernel - # --------------------------------------------------------------- # - from ..dist import parallel_magvit_vae - AutoencoderKLWan_.decode = simple_wrapper(parallel_magvit_vae(0.4, 8)(AutoencoderKLWan_.decode)) - AutoencoderKLWan2_2_.decode = simple_wrapper(parallel_magvit_vae(0.4, 16)(AutoencoderKLWan2_2_.decode)) - - # --------------------------------------------------------------- # - # Sparse Attention - # --------------------------------------------------------------- # - import torch - from paifuser.ops import wan_sparse_attention_wrapper - - WanSelfAttention.forward = simple_wrapper(wan_sparse_attention_wrapper()(WanSelfAttention.forward)) - print("Import Sparse Attention") - - WanTransformer3DModel.forward = simple_wrapper(WanTransformer3DModel.forward) - - # --------------------------------------------------------------- # - # CFG Skip Turbo - # --------------------------------------------------------------- # - import os - - if importlib.util.find_spec("paifuser.accelerator") is not None: - from paifuser.accelerator import (cfg_skip_turbo, disable_cfg_skip, - enable_cfg_skip, share_cfg_skip) - else: - from paifuser import (cfg_skip_turbo, disable_cfg_skip, - enable_cfg_skip, share_cfg_skip) - - WanTransformer3DModel.enable_cfg_skip = enable_cfg_skip()(WanTransformer3DModel.enable_cfg_skip) - WanTransformer3DModel.disable_cfg_skip = disable_cfg_skip()(WanTransformer3DModel.disable_cfg_skip) - WanTransformer3DModel.share_cfg_skip = share_cfg_skip()(WanTransformer3DModel.share_cfg_skip) - - QwenImageTransformer2DModel.enable_cfg_skip = enable_cfg_skip()(QwenImageTransformer2DModel.enable_cfg_skip) - QwenImageTransformer2DModel.disable_cfg_skip = disable_cfg_skip()(QwenImageTransformer2DModel.disable_cfg_skip) - print("Import CFG Skip Turbo") - - # --------------------------------------------------------------- # - # RMS Norm Kernel - # --------------------------------------------------------------- # - from paifuser.ops import rms_norm_forward - WanRMSNorm.forward = rms_norm_forward - print("Import PAI RMS Fuse") - - # --------------------------------------------------------------- # - # Fast Rope Kernel - # --------------------------------------------------------------- # - import types - - import torch - from paifuser.ops import (ENABLE_KERNEL, fast_rope_apply_qk, - rope_apply_real_qk) - - from . import wan_transformer3d - - def deepcopy_function(f): - return types.FunctionType(f.__code__, f.__globals__, name=f.__name__, argdefs=f.__defaults__,closure=f.__closure__) - - local_rope_apply_qk = deepcopy_function(wan_transformer3d.rope_apply_qk) - - if ENABLE_KERNEL: - def adaptive_fast_rope_apply_qk(q, k, grid_sizes, freqs): - if torch.is_grad_enabled(): - return local_rope_apply_qk(q, k, grid_sizes, freqs) - else: - return fast_rope_apply_qk(q, k, grid_sizes, freqs) - else: - def adaptive_fast_rope_apply_qk(q, k, grid_sizes, freqs): - return rope_apply_real_qk(q, k, grid_sizes, freqs) - - wan_transformer3d.rope_apply_qk = adaptive_fast_rope_apply_qk - rope_apply_qk = adaptive_fast_rope_apply_qk - print("Import PAI Fast rope") \ No newline at end of file diff --git a/videox_fun/models/attention_utils.py b/videox_fun/models/attention_utils.py index ad775c6..e10948a 100644 --- a/videox_fun/models/attention_utils.py +++ b/videox_fun/models/attention_utils.py @@ -4,12 +4,6 @@ import warnings import torch import torch.nn as nn -try: - import flash_attn.cute - FLASH_ATTN_4_AVAILABLE = True -except ModuleNotFoundError: - FLASH_ATTN_4_AVAILABLE = False - try: import flash_attn_interface FLASH_ATTN_3_AVAILABLE = True @@ -222,39 +216,13 @@ def flash_attention( if q_scale is not None: q = q * q_scale - if version is not None and version == 4 and not FLASH_ATTN_4_AVAILABLE: - warnings.warn( - 'Flash attention 4 is not available, use flash attention 3 or 2 instead.' - ) - if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE: warnings.warn( 'Flash attention 3 is not available, use flash attention 2 instead.' ) # apply attention - if (version is None or version == 4) and FLASH_ATTN_4_AVAILABLE: - # Note: dropout_p, window_size are not supported in FA3 now. - x = flash_attn.cute.flash_attn_varlen_func( - q=q, - k=k, - v=v, - cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum( - 0, dtype=torch.int32).to(q.device, non_blocking=True), - cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum( - 0, dtype=torch.int32).to(q.device, non_blocking=True), - seqused_q=None, - seqused_k=None, - max_seqlen_q=lq, - max_seqlen_k=lk, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic) - if isinstance(x, (tuple, list)): - x = x[0].unflatten(0, (b, lq)) - else: - x = x.unflatten(0, (b, lq)) - elif (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE: + if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE: # Note: dropout_p, window_size are not supported in FA3 now. x = flash_attn_interface.flash_attn_varlen_func( q=q, @@ -339,7 +307,7 @@ def attention( out = sageattn( q, k, v, attn_mask=attn_mask, tensor_layout="NHD", is_causal=causal, dropout_p=dropout_p) - elif attention_type == "FLASH_ATTENTION" and (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE or FLASH_ATTN_4_AVAILABLE): + elif attention_type == "FLASH_ATTENTION" and (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE): return flash_attention( q=q, k=k, diff --git a/videox_fun/models/wan_transformer3d_self_forcing.py b/videox_fun/models/wan_transformer3d_self_forcing.py index cee954a..38a2919 100644 --- a/videox_fun/models/wan_transformer3d_self_forcing.py +++ b/videox_fun/models/wan_transformer3d_self_forcing.py @@ -1061,9 +1061,16 @@ class WanTransformer3DModel_SelfForcing(WanTransformer3DModel): # 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 78f2290..fb488ab 100644 --- a/videox_fun/models/z_image_transformer2d.py +++ b/videox_fun/models/z_image_transformer2d.py @@ -863,6 +863,18 @@ class ZImageTransformer2DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): # 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 cb81fcd..340f8d1 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.pipelines.pipeline_utils import DiffusionPipeline 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 94ec8ea..ba7a627 100644 --- a/videox_fun/reward/hpsv3_predictor.py +++ b/videox_fun/reward/hpsv3_predictor.py @@ -918,6 +918,181 @@ class HPSv3RewardInferencer: 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 e1c0828..d8643bf 100644 --- a/videox_fun/reward/reward_fn.py +++ b/videox_fun/reward/reward_fn.py @@ -12,13 +12,35 @@ __all__ = ["AestheticReward", "HPSReward", "PickScoreReward", "MPSReward", "HPSv 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 @@ class AestheticReward(BaseReward): 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 @@ class HPSReward(BaseReward): 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 @@ class PickScoreReward(BaseReward): 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 @@ class MPSReward(BaseReward): 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 @@ class HPSv3Reward(BaseReward): dtype=torch.float16, max_reward=1, loss_scale=1, + differentiable=False, ): from .hpsv3_predictor import HPSv3RewardInferencer @@ -387,22 +420,58 @@ class HPSv3Reward(BaseReward): 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 @@ class HPSv3Reward(BaseReward): class VideoAlignReward(BaseReward): + is_image_reward = False + def __init__( self, model_path=None, @@ -442,6 +513,8 @@ class VideoAlignReward(BaseReward): 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 @@ class VideoAlignReward(BaseReward): 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 @@ class VideoAlignReward(BaseReward): 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 @@ class VideoAlignReward(BaseReward): 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 @@ class VideoAlignReward(BaseReward): # 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 @@ class VideoAlignReward(BaseReward): 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 @@ class VideoAlignReward(BaseReward): 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 @@ class VideoAlignReward(BaseReward): 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 42cfa71..f4021da 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 @@ class VideoVLMRewardInference: 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 @@ class VideoVLMRewardInference: 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/utils/__init__.py b/videox_fun/utils/__init__.py index 3bda0e4..d930511 100755 --- a/videox_fun/utils/__init__.py +++ b/videox_fun/utils/__init__.py @@ -1,5 +1,3 @@ -import importlib.util - from .cfg_optimization import cfg_skip from .discrete_sampler import DiscreteSampling from .fm_solvers import FlowDPMSolverMultistepScheduler @@ -18,33 +16,4 @@ from .trigflow_sampler import (RectifiedFlow_TrigFlowWrapper, from .utils import (calculate_dimensions, filter_kwargs, get_autocast_dtype, get_image_latent, get_image_to_video_latent, get_video_to_video_latent, save_videos_grid, - save_videos_with_audio_grid) - -# The pai_fuser is an internally developed acceleration package, which can be used on PAI. -if importlib.util.find_spec("paifuser") is not None: - # --------------------------------------------------------------- # - # FP8 Linear Kernel - # --------------------------------------------------------------- # - from paifuser.ops import (convert_model_weight_to_float8, - convert_weight_dtype_wrapper) - - from . import fp8_optimization - fp8_optimization.convert_model_weight_to_float8 = convert_model_weight_to_float8 - fp8_optimization.convert_weight_dtype_wrapper = convert_weight_dtype_wrapper - convert_model_weight_to_float8 = fp8_optimization.convert_model_weight_to_float8 - convert_weight_dtype_wrapper = fp8_optimization.convert_weight_dtype_wrapper - print("Import PAI Quantization Turbo") - - # --------------------------------------------------------------- # - # CFG Skip Turbo - # --------------------------------------------------------------- # - if importlib.util.find_spec("paifuser.accelerator") is not None: - from paifuser.accelerator import (cfg_skip_turbo, disable_cfg_skip, - enable_cfg_skip, share_cfg_skip) - else: - from paifuser import (cfg_skip_turbo, disable_cfg_skip, - enable_cfg_skip, share_cfg_skip) - from . import cfg_optimization - cfg_optimization.cfg_skip = cfg_skip_turbo - cfg_skip = cfg_skip_turbo - print("Import CFG Skip Turbo") \ No newline at end of file + save_videos_with_audio_grid) \ No newline at end of file diff --git a/videox_fun/utils/ac_handle.py b/videox_fun/utils/ac_handle.py deleted file mode 100644 index 91df98a..0000000 --- a/videox_fun/utils/ac_handle.py +++ /dev/null @@ -1,64 +0,0 @@ -from functools import partial - -from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( - CheckpointImpl, - apply_activation_checkpointing, - checkpoint_wrapper, -) - - -non_reentrant_wrapper = partial( - checkpoint_wrapper, - checkpoint_impl=CheckpointImpl.NO_REENTRANT, -) - - -def apply_checkpointing(model, block, p): - """ - Apply selective activation checkpointing. - - Selectivity is defined as a percentage p, which means we apply ac - on p of the total blocks. p is a floating number in the range of - [0, 1]. - - Some examples: - p = 0: no ac for all blocks. same as `fsdp_activation_checkpointing=False` - p = 1: apply ac on every block. i.e. "full ac". - p = 1/2: [ac, no-ac, ac, no-ac, ...] - p = 1/3: [no-ac, ac, no-ac, no-ac, ac, no-ac, ...] - p = 2/3: [ac, no-ac, ac, ac, no-ac, ac, ...] - Since blocks are homogeneous, we make ac blocks evenly spaced among - all blocks. - - Implementation: - For a given ac ratio p, we should essentially apply ac on every "1/p" - blocks. The first ac block can be as early as the 0th block, or as - late as the "1/p"th block, and we pick the middle one: (0.5p)th block. - Therefore, we are essentially to apply ac on: - (0.5/p)th block, (1.5/p)th block, (2.5/p)th block, etc., and of course, - with these values rounding to integers. - Since ac is applied recursively, we can simply use the following math - in the code to apply ac on corresponding blocks. - """ - block_idx = 0 - cut_off = 1 / 2 - # when passing p as a fraction number (e.g. 1/3), it will be interpreted - # as a string in argv, thus we need eval("1/3") here for fractions. - p = eval(p) if isinstance(p, str) else p - - def selective_checkpointing(submodule): - nonlocal block_idx - nonlocal cut_off - - if isinstance(submodule, block): - block_idx += 1 - if block_idx * p >= cut_off: - cut_off += 1 - return True - return False - - apply_activation_checkpointing( - model, - checkpoint_wrapper_fn=non_reentrant_wrapper, - check_fn=selective_checkpointing, - ) diff --git a/videox_fun/video_caption/utils/video_utils.py b/videox_fun/video_caption/utils/video_utils.py index bc9cc6c..5083555 100644 --- a/videox_fun/video_caption/utils/video_utils.py +++ b/videox_fun/video_caption/utils/video_utils.py @@ -6,7 +6,10 @@ from contextlib import contextmanager 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