diff --git a/bridge/defaults.py b/bridge/defaults.py index f37ee55..54179c1 100644 --- a/bridge/defaults.py +++ b/bridge/defaults.py @@ -76,5 +76,4 @@ class LightX2VDefaultConfig: "audio_sr": 16000, "talk_objects": None, "boundary_step_index": 2, - "rope_type": "torch", } diff --git a/bridge/translator/inference.py b/bridge/translator/inference.py index 14dc1af..7bd3b1d 100644 --- a/bridge/translator/inference.py +++ b/bridge/translator/inference.py @@ -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 diff --git a/config_builder.py b/config_builder.py index 16142d2..06db53c 100644 --- a/config_builder.py +++ b/config_builder.py @@ -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), diff --git a/lightx2v b/lightx2v index 52e86de..94f0f19 160000 --- a/lightx2v +++ b/lightx2v @@ -1 +1 @@ -Subproject commit 52e86de7c252f58cab21ef3f4d533756a10607ac +Subproject commit 94f0f19459e66f61a723db7690f1707f308599ef diff --git a/model_utils.py b/model_utils.py index bd5eb66..c5e3ab9 100644 --- a/model_utils.py +++ b/model_utils.py @@ -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", ] diff --git a/nodes/inference.py b/nodes/inference.py index 6a42da3..d079ff8 100644 --- a/nodes/inference.py +++ b/nodes/inference.py @@ -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)