| """Native Python defaults shared by generator training entries.""" |
| import json |
| from model_resources import wan_base_root, wan_model_paths, wan_tokenizer_path |
|
|
|
|
| def set_training_defaults(parser, variant): |
| model = wan_base_root() |
| data = "data/vace_train_v2" |
| caption_source = "teacher" if variant == "reward" else "grpo" |
| defaults = dict( |
| dataset_base_path=data, |
| dataset_metadata_path=f"{data}/metadata.csv" if variant == "hoi" else f"{data}/metadata_caption_branch_{caption_source}.csv", |
| data_file_keys="video,vace_video,vace_reference_image", |
| extra_inputs="vace_video,vace_reference_image", |
| height=480, width=480, num_frames=13, dataset_repeat=1, num_epochs=5, |
| model_paths=json.dumps(wan_model_paths(model)), |
| tokenizer_path=wan_tokenizer_path(model), |
| use_gradient_checkpointing_offload=True, |
| output_path=f"outputs/train_{variant}", |
| learning_rate=1e-4 if variant == "hoi" else 3e-4, |
| ) |
| if variant == "hoi": |
| defaults.update(lora_base_model="vace", lora_target_modules="q,k,v,o,ffn.0,ffn.2", lora_rank=128) |
| else: |
| defaults.update(preset_lora_path="checkpoints/generator_hoi/ckpt.safetensors", preset_lora_model="vace", find_unused_parameters=True) |
| parser.set_defaults(**defaults) |
| |
| for action in parser._actions: |
| if action.dest in defaults: |
| action.required = False |
| return parser |
|
|
