File size: 1,494 Bytes
c28ffff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | """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)
# Upstream required options now have concrete package defaults.
for action in parser._actions:
if action.dest in defaults:
action.required = False
return parser
|