wehub-resource-sync
175 lines
6.5 KiB
Python
175 lines
6.5 KiB
Python
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License").
|
|
# You may not use this file except in compliance with the License.
|
|
# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from omegaconf import OmegaConf
|
|
|
|
|
|
DEFAULT_NEGATIVE_PROMPT = (
|
|
"色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，"
|
|
"整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，"
|
|
"画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，"
|
|
"手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
|
|
)
|
|
|
|
|
|
wan_default_config = {
|
|
"Wan2.2-TI2V-5B": {
|
|
"resolution": [1280, 704],
|
|
"temporal_compression_ratio": 4,
|
|
"spatial_compression_ratio": 16,
|
|
"num_heads": 24,
|
|
"head_dim": 128,
|
|
"num_transformer_blocks": 30,
|
|
"fps": 24,
|
|
}
|
|
}
|
|
|
|
|
|
SECTION_KEYS = (
|
|
"infra",
|
|
"algorithm",
|
|
"training",
|
|
"data",
|
|
"evaluation",
|
|
"inference",
|
|
"logging",
|
|
"checkpoints",
|
|
)
|
|
|
|
|
|
def _set_once(config, key, value, source):
|
|
if value is None:
|
|
return
|
|
if key in config and config[key] != value:
|
|
raise ValueError(
|
|
f"{key} is defined more than once with different values: "
|
|
f"{config[key]} vs {value} from {source}."
|
|
)
|
|
config[key] = value
|
|
|
|
|
|
def section_get(config, section_key, key, default=None, aliases=()):
|
|
"""Read a grouped config value, falling back to legacy flat names."""
|
|
section = config.get(section_key, None)
|
|
candidate_keys = (key, *aliases)
|
|
if section is not None:
|
|
for candidate in candidate_keys:
|
|
if candidate in section:
|
|
return section[candidate]
|
|
for candidate in candidate_keys:
|
|
if candidate in config:
|
|
return config[candidate]
|
|
return default
|
|
|
|
|
|
def normalize_config(config):
|
|
"""Expand grouped release configs into the flat runtime schema.
|
|
|
|
The training and inference code historically reads fields such as
|
|
``config.batch_size`` and ``config.model_kwargs`` directly. Release
|
|
configs can group those fields for readability, then call this function at
|
|
the entry point to preserve the existing runtime contract.
|
|
"""
|
|
for section_key in SECTION_KEYS:
|
|
section = config.get(section_key, None)
|
|
if section is None:
|
|
continue
|
|
for key, value in section.items():
|
|
config[key] = value
|
|
|
|
evaluation = config.get("evaluation", None)
|
|
if evaluation is not None:
|
|
if "interval" in evaluation:
|
|
_set_once(config, "generate_interval", evaluation.interval, "evaluation.interval")
|
|
_set_once(config, "vis_interval", evaluation.interval, "evaluation.interval")
|
|
if "num_frames" in evaluation:
|
|
num_frames = evaluation.num_frames
|
|
if isinstance(num_frames, (list, tuple)):
|
|
vis_lengths = list(num_frames)
|
|
inference_num_frames = vis_lengths[0] if vis_lengths else 0
|
|
else:
|
|
inference_num_frames = int(num_frames)
|
|
vis_lengths = [inference_num_frames]
|
|
_set_once(config, "inference_num_frames", inference_num_frames, "evaluation.num_frames")
|
|
_set_once(config, "vis_video_lengths", vis_lengths, "evaluation.num_frames")
|
|
if "use_ema" in evaluation:
|
|
_set_once(config, "vis_ema", evaluation.use_ema, "evaluation.use_ema")
|
|
|
|
model_section = config.get("model", None)
|
|
base_model_kwargs = config.get("model_kwargs", None)
|
|
model_kwargs = OmegaConf.create({})
|
|
if base_model_kwargs is not None:
|
|
model_kwargs = OmegaConf.merge(model_kwargs, base_model_kwargs)
|
|
|
|
if model_section is not None:
|
|
section_kwargs = model_section.get("kwargs", None)
|
|
if section_kwargs is not None:
|
|
model_kwargs = OmegaConf.merge(model_kwargs, section_kwargs)
|
|
|
|
model_name = model_section.get("name", None)
|
|
if model_name is not None:
|
|
model_kwargs.model_name = model_name
|
|
config.model_name = model_name
|
|
|
|
_set_once(
|
|
config,
|
|
"num_frame_per_block",
|
|
model_section.get("num_frame_per_block", None),
|
|
"model.num_frame_per_block",
|
|
)
|
|
|
|
if "model_name" in config and "model_name" not in model_kwargs:
|
|
model_kwargs.model_name = config.model_name
|
|
if "timestep_shift" in config and "timestep_shift" not in model_kwargs:
|
|
model_kwargs.timestep_shift = config.timestep_shift
|
|
if "timestep_shift" in model_kwargs:
|
|
_set_once(config, "timestep_shift", model_kwargs.timestep_shift, "model_kwargs.timestep_shift")
|
|
|
|
model_num_frame_per_block = model_kwargs.get("num_frame_per_block", None)
|
|
if model_num_frame_per_block is not None:
|
|
_set_once(config, "num_frame_per_block", model_num_frame_per_block, "model_kwargs.num_frame_per_block")
|
|
|
|
if len(model_kwargs) > 0:
|
|
config.model_kwargs = model_kwargs
|
|
|
|
if "wandb_host" not in config:
|
|
config.wandb_host = "https://api.wandb.ai"
|
|
|
|
if "negative_prompt" not in config:
|
|
config.negative_prompt = DEFAULT_NEGATIVE_PROMPT
|
|
|
|
if config.get("trainer", None) == "score_distillation":
|
|
dmd_defaults = {
|
|
"i2v": False,
|
|
"teacher_forcing": False,
|
|
"backward_simulation": True,
|
|
"independent_first_frame": False,
|
|
"num_train_timestep": 1000,
|
|
"denoising_loss_type": "flow",
|
|
"real_guidance_scale": 3.0,
|
|
"fake_guidance_scale": 0.0,
|
|
}
|
|
for key, value in dmd_defaults.items():
|
|
if key not in config:
|
|
config[key] = value
|
|
if "causal" not in config:
|
|
config.causal = bool(config.get("all_causal", True))
|
|
|
|
# Causal DMD uses the same Wan backbone for generator/teacher/critic unless
|
|
# a role-specific override is explicitly provided.
|
|
if getattr(config, "all_causal", False) and "model_kwargs" in config:
|
|
for role_key in ("real_model_kwargs", "fake_model_kwargs"):
|
|
if config.get(role_key, None) is None:
|
|
config[role_key] = OmegaConf.create(
|
|
OmegaConf.to_container(config.model_kwargs, resolve=True)
|
|
)
|
|
|
|
return config
|