126 lines
5.0 KiB
Python
126 lines
5.0 KiB
Python
"""Translate ``LightX2VInferenceConfig`` widget values into lightx2v config keys.
|
|
|
|
The wrapper's widget naming follows ComfyUI conventions (``height``, ``width``,
|
|
``video_length``, ``cfg_scale`` …). lightx2v's internal naming is different
|
|
(``target_height``, ``target_width``, ``target_video_length``,
|
|
``sample_guide_scale`` …). The single source of truth for that translation
|
|
is the ``WRAPPER_TO_LIGHTX2V_FIELDS`` table below — when adding a new field,
|
|
add a row there rather than burying the rename inside the function body.
|
|
"""
|
|
|
|
import os
|
|
from typing import Any, Dict
|
|
|
|
from ..defaults import LightX2VDefaultConfig
|
|
|
|
# Direct rename map: wrapper-side key -> lightx2v-side key.
|
|
# A row of ("foo", "foo") means the name matches but we still want to forward
|
|
# the value explicitly (rather than relying on the default config).
|
|
WRAPPER_TO_LIGHTX2V_FIELDS: Dict[str, str] = {
|
|
# Model selection — names match.
|
|
"model_cls": "model_cls",
|
|
"model_path": "model_path",
|
|
"task": "task",
|
|
# Inference loop.
|
|
"infer_steps": "infer_steps",
|
|
"seed": "seed",
|
|
"sample_shift": "sample_shift",
|
|
# Output shape — wrapper uses bare names, lightx2v prefixes with target_.
|
|
"height": "target_height",
|
|
"width": "target_width",
|
|
"video_length": "target_video_length",
|
|
"fps": "target_fps",
|
|
"video_duration": "video_duration",
|
|
# Image preprocessing.
|
|
"resize_mode": "resize_mode",
|
|
"fixed_area": "fixed_area",
|
|
# Sekotalk-specific.
|
|
"prev_frame_length": "prev_frame_length",
|
|
# Distillation.
|
|
"denoising_step_list": "denoising_step_list",
|
|
"use_31_block": "use_31_block",
|
|
}
|
|
|
|
# Attention type fans out to three internal slots in lightx2v.
|
|
_ATTN_TYPE_SLOTS = ("self_attn_1_type", "cross_attn_1_type", "cross_attn_2_type")
|
|
|
|
# LightX2V unified the dedicated distilled Wan runners into the regular Wan
|
|
# runners in 23fc3d45. Keep old ComfyUI workflows loadable by translating the
|
|
# retired model class names to the new model class plus scheduler selector.
|
|
_LEGACY_MODEL_CLASS_ALIASES = {
|
|
"wan2.1_distill": ("wan2.1", "dmd2"),
|
|
"wan2.1_mean_flow_distill": ("wan2.1", "mean_flow"),
|
|
"wan2.2_moe_distill": ("wan2.2_moe", "dmd2"),
|
|
}
|
|
|
|
|
|
def _add_wan22_branch_paths(model_path: str, updates: Dict[str, Any]) -> None:
|
|
"""Expose Wan2.2's two transformer branches in the runner config.
|
|
|
|
Hugging Face model packages commonly keep the branches below
|
|
``distill_models/`` while the unified runner accepts explicit branch
|
|
directories. Supplying both path variants is harmless: the runner picks
|
|
the original or quantized key according to ``dit_quantized``.
|
|
"""
|
|
if not model_path:
|
|
return
|
|
|
|
for branch in ("high_noise_model", "low_noise_model"):
|
|
candidates = (
|
|
os.path.join(model_path, branch),
|
|
os.path.join(model_path, "distill_models", branch),
|
|
)
|
|
branch_path = next((path for path in candidates if os.path.isdir(path)), None)
|
|
if branch_path is None:
|
|
continue
|
|
config_key = branch.removesuffix("_model")
|
|
updates.setdefault(f"{config_key}_original_ckpt", branch_path)
|
|
updates.setdefault(f"{config_key}_quantized_ckpt", branch_path)
|
|
|
|
|
|
def apply_inference_config(config: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Translate inference widget values to a partial lightx2v config dict."""
|
|
updates: Dict[str, Any] = {}
|
|
|
|
# Bulk rename via the table.
|
|
for wrapper_key, lightx2v_key in WRAPPER_TO_LIGHTX2V_FIELDS.items():
|
|
if wrapper_key not in config:
|
|
continue
|
|
# seed=-1 means "use lightx2v's default / random"; leave it out.
|
|
if wrapper_key == "seed" and config[wrapper_key] == -1:
|
|
continue
|
|
updates[lightx2v_key] = config[wrapper_key]
|
|
|
|
# cfg_scale -> sample_guide_scale (and toggle enable_cfg).
|
|
if "cfg_scale" in config:
|
|
updates["sample_guide_scale"] = config["cfg_scale"]
|
|
updates["enable_cfg"] = config["cfg_scale"] != 1.0
|
|
|
|
# Wan2.2 MoE has two CFG scales (high/low noise) and a boundary param.
|
|
model_cls = config.get("model_cls", "")
|
|
legacy_model = _LEGACY_MODEL_CLASS_ALIASES.get(model_cls)
|
|
if legacy_model is not None:
|
|
model_cls, distill_method = legacy_model
|
|
updates["model_cls"] = model_cls
|
|
updates["distill_method"] = distill_method
|
|
|
|
if "wan2.2_moe" in model_cls:
|
|
updates["boundary"] = 0.9
|
|
updates["sample_guide_scale"] = [config.get("cfg_scale"), config.get("cfg_scale2")]
|
|
_add_wan22_branch_paths(config.get("model_path", ""), updates)
|
|
if "wan2.2" in model_cls:
|
|
updates["use_image_encoder"] = False
|
|
|
|
# One widget value drives three lightx2v attention slots.
|
|
attention_type = config.get("attention_type", LightX2VDefaultConfig.DEFAULT_ATTENTION_TYPE)
|
|
for slot in _ATTN_TYPE_SLOTS:
|
|
updates[slot] = attention_type
|
|
|
|
# TAEW2.1 lightweight VAE lives next to the model.
|
|
if config.get("use_tiny_vae", False):
|
|
updates["use_tiny_vae"] = True
|
|
updates["tiny_vae"] = True
|
|
updates["tiny_vae_path"] = os.path.join(config["model_path"], "taew2_1.pth")
|
|
|
|
return updates
|