fix: adapt wrapper to LightX2V 0.6 runners

This commit is contained in:
gaclove
2026-10-22 23:44:02 +08:00
parent 78897cfe2c
commit 9cc24e44c3
6 changed files with 195 additions and 36 deletions
-1
View File
@@ -76,5 +76,4 @@ class LightX2VDefaultConfig:
"audio_sr": 16000,
"talk_objects": None,
"boundary_step_index": 2,
"rope_type": "torch",
}
+40
View File
@@ -44,6 +44,39 @@ WRAPPER_TO_LIGHTX2V_FIELDS: Dict[str, str] = {
# 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."""
@@ -65,9 +98,16 @@ def apply_inference_config(config: Dict[str, Any]) -> Dict[str, Any]:
# 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
+1
View File
@@ -313,6 +313,7 @@ class ConfigBuilder:
("model_cls", None),
("model_path", None),
("task", None),
("distill_method", None),
("t5_quantized", False),
("clip_quantized", False),
("lora_configs", None),
+13 -5
View File
@@ -30,19 +30,27 @@ def scan_models() -> List[str]:
def support_model_cls_list() -> List[str]:
"""Return model classes understood by the bundled LightX2V revision.
The retired ``*_distill`` entries are kept as compatibility aliases for saved
ComfyUI workflows. The current LightX2V runner registry uses the base Wan
class plus ``distill_method`` instead; ``bridge.translator.inference``
canonicalizes those aliases before the runner is built.
"""
return [
"wan2.1",
"wan2.1_distill",
"wan2.1_vace",
"cogvideox",
"seko_talk",
"seko_talk_ar",
"wan2.2_moe",
"wan2.2",
"wan2.2_moe_audio",
"wan2.2_audio",
"wan2.2_moe_distill",
"wan2.2_s2v",
"qwen_image",
"seedvr2",
# Legacy workflow values; normalized before build_runner().
"wan2.1_distill",
"wan2.1_mean_flow_distill",
"wan2.2_moe_distill",
]
+140 -29
View File
@@ -7,9 +7,36 @@ import torch
from comfy.utils import ProgressBar
from ..config_builder import ConfigBuilder
from ..lightx2v.lightx2v.infer import init_runner
from ..lightx2v.lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict
from ..lightx2v.lightx2v.utils.set_config import auto_calc_config, set_args2config
# LightX2V 0.6 replaced the resident ``init_runner`` entrypoint and the
# ``set_args2config``/``auto_calc_config`` pair with a startup/request split.
# Keep the imports optional so the wrapper can select the API at runtime while
# still supporting the previous submodule revision.
try:
from ..lightx2v.lightx2v.infer import init_runner as _legacy_init_runner
except ImportError:
_legacy_init_runner = None
try:
from ..lightx2v.lightx2v.utils.input_info import init_empty_input_info as _legacy_init_empty_input_info
from ..lightx2v.lightx2v.utils.input_info import update_input_info_from_dict as _legacy_update_input_info_from_dict
except ImportError:
_legacy_init_empty_input_info = None
_legacy_update_input_info_from_dict = None
try:
from ..lightx2v.lightx2v.utils.set_config import auto_calc_config as _legacy_auto_calc_config
from ..lightx2v.lightx2v.utils.set_config import set_args2config as _legacy_set_args2config
except ImportError:
_legacy_auto_calc_config = None
_legacy_set_args2config = None
try:
from ..lightx2v.lightx2v.models.runners.runner_factory import build_runner as _modern_build_runner
from ..lightx2v.lightx2v.utils.set_config import build_startup_config as _modern_build_startup_config
except ImportError:
_modern_build_runner = None
_modern_build_startup_config = None
class LightX2VModularInferenceV2:
@@ -62,16 +89,33 @@ class LightX2VModularInferenceV2:
# never get merged — rs2v_infer then KeyErrors on config["vae_stride"].
if "task" not in config:
config["task"] = "rs2v"
# set_config = set_args2config + auto_calc_config. set_args2config strips
# any key that's part of an InputInfo dataclass (target_video_length,
# infer_steps, seed, ...). For CLI runs auto_calc_config recovers them
# by merging --config_json, but we have no external JSON, so
# auto_calc_config's `config["target_video_length"]` access KeyErrors.
# Inject the bridge between set_args2config and auto_calc_config.
target_video_length = config.get("target_video_length", config.get("segment_length", config.get("video_length", 81)))
formatted = set_args2config(config)
formatted["target_video_length"] = target_video_length
formatted = auto_calc_config(formatted)
# LightX2V 0.6 uses the same startup builder as the regular runner;
# older revisions need the set_args2config + auto_calc_config pair.
# The shot runner's load_clip_configs() calls build_startup_config
# again, so passing the modern startup shape here is safe and keeps
# model config/vae metadata loading in LightX2V itself.
if _modern_build_startup_config is not None:
startup_data = dict(config)
startup_data["size"] = [
int(config.get("target_height", config.get("height", 480))),
int(config.get("target_width", config.get("width", 832))),
]
startup_data["num_frames"] = int(
config.get("target_video_length", config.get("video_length", config.get("num_frames", 81)))
)
formatted = _modern_build_startup_config(startup_data)
else:
# set_args2config strips any key that's part of an InputInfo
# dataclass (target_video_length, infer_steps, seed, ...). For
# CLI runs auto_calc_config recovers them by merging
# --config_json; our in-memory flow has no external JSON, so
# bridge target_video_length back before auto_calc_config.
target_video_length = config.get("target_video_length", config.get("segment_length", config.get("video_length", 81)))
if _legacy_set_args2config is None or _legacy_auto_calc_config is None:
raise RuntimeError("No compatible LightX2V shot-runner config API found")
formatted = _legacy_set_args2config(config)
formatted["target_video_length"] = target_video_length
formatted = _legacy_auto_calc_config(formatted)
main_cfg = {
"lightx2v_path": "",
"clip_configs": [
@@ -88,6 +132,53 @@ class LightX2VModularInferenceV2:
return load_clip_configs(main_cfg)
@staticmethod
def _build_modern_startup_config(config):
"""Translate wrapper shape fields into LightX2V 0.6 startup fields."""
if _modern_build_startup_config is None or _modern_build_runner is None:
raise RuntimeError("The installed LightX2V revision does not provide the modern startup API")
startup_data = dict(config)
startup_data["size"] = [
int(config.get("target_height", config.get("height", 480))),
int(config.get("target_width", config.get("width", 832))),
]
startup_data["num_frames"] = int(
config.get("target_video_length", config.get("video_length", config.get("num_frames", 81)))
)
return _modern_build_startup_config(startup_data)
@staticmethod
def _build_modern_request(config, runner):
"""Build only fields accepted by the current runner's request schema."""
task = config.get("task")
supported_fields = runner.get_supported_request_fields(task)
request = {}
for field in supported_fields:
if field == "size":
# SekoTalk's fixed-area/min-side modes derive the output size
# from the model's resize policy. The 0.6 request validator
# rejects an explicit size unless the request selected
# fixed_shape; keep the wrapper's target dimensions in startup
# config but do not send them as a per-request override.
if config.get("model_cls", "").startswith("seko") and config.get("resize_mode") != "fixed_shape":
continue
request[field] = [
int(config.get("target_height", config.get("height", 480))),
int(config.get("target_width", config.get("width", 832))),
]
elif field == "num_frames":
request[field] = int(
config.get("target_video_length", config.get("video_length", config.get("num_frames", 81)))
)
elif field in config:
request[field] = config[field]
request["task"] = task
request["return_result_tensor"] = True
request["save_result_path"] = ""
return request
def generate(self, prepared_config):
"""Run inference with prepared configuration."""
@@ -124,19 +215,22 @@ class LightX2VModularInferenceV2:
shot_cfg = self._build_rs2v_shot_config(config)
self.__class__._current_runner = ShotRS2VPipeline(shot_cfg)
else:
# set_args2config strips InputInfo-dataclass keys (target_video_length,
# infer_steps, ...). CLI flows recover them via --config_json merging
# inside auto_calc_config; our in-memory flow has no external JSON, so
# we bridge target_video_length manually between the two halves so
# auto_calc_config:194's modulo check on s2v/i2v doesn't KeyError.
target_video_length = config.get(
"target_video_length",
config.get("segment_length", config.get("video_length", 81)),
)
formatted_config = set_args2config(config)
formatted_config["target_video_length"] = target_video_length
formatted_config = auto_calc_config(formatted_config)
self.__class__._current_runner = init_runner(formatted_config)
if _modern_build_runner is not None and _modern_build_startup_config is not None:
startup_config = self._build_modern_startup_config(config)
self.__class__._current_runner = _modern_build_runner(startup_config)
else:
if _legacy_set_args2config is None or _legacy_auto_calc_config is None or _legacy_init_runner is None:
raise RuntimeError("No compatible LightX2V runner API found in the installed LightX2V submodule")
# The legacy API strips InputInfo-dataclass keys. Bridge
# target_video_length back before auto_calc_config.
target_video_length = config.get(
"target_video_length",
config.get("segment_length", config.get("video_length", 81)),
)
formatted_config = _legacy_set_args2config(config)
formatted_config["target_video_length"] = target_video_length
formatted_config = _legacy_auto_calc_config(formatted_config)
self.__class__._current_runner = _legacy_init_runner(formatted_config)
self.__class__._current_config_hash = config_hash
progress = ProgressBar(100)
@@ -153,10 +247,27 @@ class LightX2VModularInferenceV2:
config["save_result_path"] = ""
config["negative_prompt"] = config.get("negative_prompt", "")
if config.get("task") == "rs2v":
result_dict = current_runner.run_pipeline(config)
# ShotRS2VPipeline delegates request validation to its clip
# runner. Passing the prepared config wholesale worked with
# the old InputInfo path, but LightX2V 0.6 rejects startup
# fields (model paths, quantization flags, etc.) as request
# fields. Filter the request against the modern runner schema
# before handing it to the shot pipeline.
if _modern_build_runner is not None and _modern_build_startup_config is not None:
shot_runner = current_runner.clip_generators["rs2v_clip"]
request_data = self._build_modern_request(config, shot_runner)
result_dict = current_runner.run_pipeline(request_data)
else:
result_dict = current_runner.run_pipeline(config)
elif _modern_build_runner is not None and _modern_build_startup_config is not None:
request_data = self._build_modern_request(config, current_runner)
input_info = current_runner.prepare_request(request_data)
result_dict = current_runner.run_request(input_info)
else:
input_data = init_empty_input_info(config.task)
update_input_info_from_dict(input_data, config)
if _legacy_init_empty_input_info is None or _legacy_update_input_info_from_dict is None:
raise RuntimeError("No compatible LightX2V input-info API found in the installed LightX2V submodule")
input_data = _legacy_init_empty_input_info(config.task)
_legacy_update_input_info_from_dict(input_data, config)
current_runner.set_config(config)
result_dict = current_runner.run_pipeline(input_data)