diff --git a/README.md b/README.md index cb79910b5..db79dc22c 100644 --- a/README.md +++ b/README.md @@ -176,6 +176,8 @@ For attention operators installation, please refer to our documentation: **[Engl ### Usage Example +See the [MiniMax-H3 guide](scripts/minimax_h3/README.md) for checkpoint layout, local LoRA paths, CLI presets, and server/POST examples. + ```python # examples/minimax_h3/minimax_h3_t2av_dmd.py """ @@ -194,7 +196,7 @@ pipe = LightX2VPipeline( # The DMD config uses the released 768p LoRA, 4 inference steps, # video_flow_shift=6, audio_flow_shift=3, and lora alpha=128. pipe.create_generator( - config_json="configs/minimax_h3/dmd/minimax_h3_bf16_4step_single_gpu_offload.json" + config_json="configs/minimax_h3/dmd/minimax_h3_bf16_4step.json" ) # Generation parameters diff --git a/README_zh.md b/README_zh.md index 09ef3708f..1d24e9099 100644 --- a/README_zh.md +++ b/README_zh.md @@ -176,6 +176,9 @@ uv pip install -v . # pip install -v . 注意力算子安装说明请参考我们的文档:**[英文文档](https://lightx2v-en.readthedocs.io/en/latest/getting_started/quickstart.html#step-4-install-attention-operators) | [中文文档](https://lightx2v-zhcn.readthedocs.io/zh-cn/latest/getting_started/quickstart.html#id9)** ### 使用示例 + +权重目录、本地 LoRA 路径、CLI 配置及服务与 POST 示例见 [MiniMax-H3 使用说明](scripts/minimax_h3/README_zh.md)。 + ```python # examples/minimax_h3/minimax_h3_t2av_dmd.py """ @@ -194,7 +197,7 @@ pipe = LightX2VPipeline( # DMD 配置使用已发布的 768p LoRA、4 步推理、 # video_flow_shift=6、audio_flow_shift=3 和 LoRA alpha=128。 pipe.create_generator( - config_json="configs/minimax_h3/dmd/minimax_h3_bf16_4step_single_gpu_offload.json" + config_json="configs/minimax_h3/dmd/minimax_h3_bf16_4step.json" ) # 生成参数 diff --git a/app/gradio_demo.py b/app/gradio_demo.py index 383a325b2..6744225f1 100644 --- a/app/gradio_demo.py +++ b/app/gradio_demo.py @@ -10,13 +10,12 @@ import os import warnings -import torch from loguru import logger from utils.i18n import DEFAULT_LANG, set_language from utils.model_utils import cleanup_memory, extract_op_name, get_model_configs from utils.ui_builder import build_ui, generate_unique_filename, get_auto_config_dict -from lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict +from lightx2v.models.runners.runner_factory import build_runner from lightx2v.utils.set_config import get_default_config warnings.filterwarnings("ignore", category=UserWarning, module="huggingface_hub") @@ -41,12 +40,8 @@ global_runner = None -current_config = None -cur_dit_path = None -cur_use_lora = None -cur_lora_path = None -cur_high_lora_path = None -cur_low_lora_path = None +current_startup_config = None +current_lora_configs = [] def run_inference( @@ -146,17 +141,10 @@ def run_inference( model_cls = model_config["model_cls"] model_path = model_config["model_path"] - global global_runner, current_config, cur_dit_path, cur_use_lora, cur_lora_path, cur_high_lora_path, cur_low_lora_path + global global_runner, current_startup_config, current_lora_configs logger.info(f"Auto-determined model_cls: {model_cls} (model type: {model_type_input})") - if model_cls.startswith("wan2.2"): - current_dit_path = f"{high_noise_path_input}|{low_noise_path_input}" if high_noise_path_input and low_noise_path_input else None - else: - current_dit_path = dit_path_input - - needs_reinit = lazy_load or unload_modules or global_runner is None or cur_dit_path != current_dit_path or cur_use_lora != use_lora - config_graio = { "infer_steps": infer_steps, "target_video_length": num_frames, @@ -207,22 +195,8 @@ def run_inference( "aspect_ratio": aspect_ratio, } - args = argparse.Namespace( - model_cls=model_cls, - seed=seed, - task=task, - model_path=model_path, - prompt=prompt, - negative_prompt=negative_prompt, - image_path=image_path, - save_result_path=save_result_path, - return_result_tensor=False, - aspect_ratio=aspect_ratio, - target_shape=[], - ) - input_info = init_empty_input_info(args.task) config = get_default_config() - config.update({k: v for k, v in vars(args).items()}) + config.update({"model_cls": model_cls, "task": task, "model_path": model_path}) config.update(config_graio) config.update(model_config) @@ -233,92 +207,62 @@ def run_inference( logger.info(f"Using model: {model_path}") logger.info(f"Inference config:\n{json.dumps(config, indent=4, ensure_ascii=False)}") + startup_config = {key: value for key, value in config.items() if key not in {"aspect_ratio", "target_video_length", "lora_configs"}} + lora_configs = config.get("lora_configs") or [] + + current_targets = {item.get("name") for item in current_lora_configs} + new_targets = {item.get("name") for item in lora_configs} + config_changed = current_startup_config != startup_config + # Wan2.2 switch_lora cannot remove an adapter from a branch. + lora_targets_changed = current_targets != new_targets + needs_reinit = global_runner is None or lazy_load or unload_modules or config_changed or lora_targets_changed + # 初始化或重用 runner runner = global_runner if needs_reinit: if runner is not None: + global_runner = None del runner - torch.cuda.empty_cache() - gc.collect() - - from lightx2v.infer import init_runner - - runner = init_runner(config) + # Inference freezes the model graph; unfreeze it before collection. + gc.unfreeze() + cleanup_memory() - data = args.__dict__ - update_input_info_from_dict(input_info, data) - - current_config = config - cur_dit_path = current_dit_path - cur_use_lora = use_lora - cur_lora_path = lora_path - - # 保存 Wan2.2 的 LoRA 路径 - if model_cls.startswith("wan2.2"): - lora_configs = config.get("lora_configs") - if lora_configs: - lora_name_to_info = {item["name"]: item for item in lora_configs} - cur_high_lora_path = lora_name_to_info.get("high_noise_model", {}).get("path") - cur_low_lora_path = lora_name_to_info.get("low_noise_model", {}).get("path") - else: - cur_high_lora_path = None - cur_low_lora_path = None + runner = build_runner(config) if not lazy_load: global_runner = runner - else: - runner.config = config - data = args.__dict__ - update_input_info_from_dict(input_info, data) - - # 如果 use_lora 为 True 且 lora_path 变化了,调用 switch_lora - if use_lora: - lora_configs = config.get("lora_configs") - if model_cls.startswith("wan2.2") and lora_configs: - # 对于 Wan2.2 模型,从 lora_configs 中获取 high_noise 和 low_noise 的 LoRA 路径 - lora_name_to_info = {item["name"]: item for item in lora_configs} - high_lora_path = None - high_lora_strength = 1.0 - low_lora_path = None - low_lora_strength = 1.0 - - if "high_noise_model" in lora_name_to_info: - high_lora_info = lora_name_to_info["high_noise_model"] - high_lora_path = high_lora_info["path"] - high_lora_strength = high_lora_info.get("strength", 1.0) - - if "low_noise_model" in lora_name_to_info: - low_lora_info = lora_name_to_info["low_noise_model"] - low_lora_path = low_lora_info["path"] - low_lora_strength = low_lora_info.get("strength", 1.0) - - # 检查 high_lora_path 和 low_lora_path 是否变化 - high_lora_changed = high_lora_path != cur_high_lora_path - low_lora_changed = low_lora_path != cur_low_lora_path - - if high_lora_changed or low_lora_changed: - if hasattr(runner, "switch_lora"): - runner.switch_lora( - high_lora_path=high_lora_path, - high_lora_strength=high_lora_strength, - low_lora_path=low_lora_path, - low_lora_strength=low_lora_strength, - ) - logger.info(f"Switched LoRA for Wan2.2: high={high_lora_path}, low={low_lora_path}") - cur_high_lora_path = high_lora_path - cur_low_lora_path = low_lora_path - else: - logger.warning("Runner does not support switch_lora method") - elif lora_path and lora_path != cur_lora_path: - lora_strength_val = float(lora_strength) if lora_strength is not None else 1.0 - if hasattr(runner, "switch_lora"): - runner.switch_lora(lora_path, lora_strength_val) - logger.info(f"Switched LoRA to: {lora_path} with strength={lora_strength_val}") - else: - logger.warning("Runner does not support switch_lora method") - cur_lora_path = lora_path - - runner.run_pipeline(input_info) + elif lora_configs != current_lora_configs: + if model_cls.startswith("wan2.2"): + lora_by_name = {item["name"]: item for item in lora_configs} + high_lora = lora_by_name.get("high_noise_model", {}) + low_lora = lora_by_name.get("low_noise_model", {}) + switched = runner.switch_lora( + high_lora_path=high_lora.get("path"), + high_lora_strength=high_lora.get("strength", 1.0), + low_lora_path=low_lora.get("path"), + low_lora_strength=low_lora.get("strength", 1.0), + ) + else: + switched = runner.switch_lora(lora_configs[0]["path"], lora_configs[0]["strength"]) + if not switched: + raise RuntimeError("Failed to switch LoRA") + + current_startup_config = startup_config + current_lora_configs = lora_configs + + form_data = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "image_path": image_path, + "seed": seed, + "save_result_path": save_result_path, + "return_result_tensor": False, + "aspect_ratio": aspect_ratio, + "target_video_length": num_frames, + } + supported_request_fields = runner.get_supported_request_fields(task) + input_info = runner.prepare_request({key: value for key, value in form_data.items() if key in supported_request_fields}) + runner.run_request(input_info) cleanup_memory() return save_result_path diff --git a/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld.json b/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld.json index 862c20804..b9ca51a5b 100644 --- a/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld.json +++ b/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld.json @@ -6,7 +6,7 @@ "target_width": 640, "target_video_length": 17, "target_fps": 10.0, - "enable_cfg": true, + "enable_cfg": false, "action_mode": "forward_dynamics", "domain_name": "agibotworld", "view_point": "concat_view", diff --git a/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld_multichunk.json b/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld_multichunk.json index 7bd4f812f..0c98cd329 100644 --- a/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld_multichunk.json +++ b/configs/cosmos3/cosmos3_nano_omni_action_fd_agibotworld_multichunk.json @@ -6,7 +6,7 @@ "target_width": 640, "target_video_length": 17, "target_fps": 10.0, - "enable_cfg": true, + "enable_cfg": false, "action_mode": "forward_dynamics", "domain_name": "agibotworld", "view_point": "concat_view", diff --git a/configs/cosmos3/cosmos3_nano_omni_action_id_av.json b/configs/cosmos3/cosmos3_nano_omni_action_id_av.json index 36bba2638..11825d95b 100644 --- a/configs/cosmos3/cosmos3_nano_omni_action_id_av.json +++ b/configs/cosmos3/cosmos3_nano_omni_action_id_av.json @@ -6,7 +6,7 @@ "target_width": 832, "target_video_length": 61, "target_fps": 10.0, - "enable_cfg": true, + "enable_cfg": false, "action_mode": "inverse_dynamics", "domain_name": "av", "view_point": "ego_view", diff --git a/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld.json b/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld.json index 862c20804..b9ca51a5b 100644 --- a/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld.json +++ b/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld.json @@ -6,7 +6,7 @@ "target_width": 640, "target_video_length": 17, "target_fps": 10.0, - "enable_cfg": true, + "enable_cfg": false, "action_mode": "forward_dynamics", "domain_name": "agibotworld", "view_point": "concat_view", diff --git a/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld_multichunk.json b/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld_multichunk.json index 7bd4f812f..0c98cd329 100644 --- a/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld_multichunk.json +++ b/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld_multichunk.json @@ -6,7 +6,7 @@ "target_width": 640, "target_video_length": 17, "target_fps": 10.0, - "enable_cfg": true, + "enable_cfg": false, "action_mode": "forward_dynamics", "domain_name": "agibotworld", "view_point": "concat_view", diff --git a/configs/cosmos3/cosmos3_super_omni_action_id_av.json b/configs/cosmos3/cosmos3_super_omni_action_id_av.json index 36bba2638..11825d95b 100644 --- a/configs/cosmos3/cosmos3_super_omni_action_id_av.json +++ b/configs/cosmos3/cosmos3_super_omni_action_id_av.json @@ -6,7 +6,7 @@ "target_width": 832, "target_video_length": 61, "target_fps": 10.0, - "enable_cfg": true, + "enable_cfg": false, "action_mode": "inverse_dynamics", "domain_name": "av", "view_point": "ego_view", diff --git a/configs/disagg/multi_node/wan22_i2v_distill_controller.json b/configs/disagg/multi_node/wan22_i2v_distill_controller.json index 5e7614428..c3ecb1bb1 100644 --- a/configs/disagg/multi_node/wan22_i2v_distill_controller.json +++ b/configs/disagg/multi_node/wan22_i2v_distill_controller.json @@ -43,7 +43,6 @@ "low_noise_quantized_ckpt": "/path/to/wan2.2_i2v_A14b_low_noise_int8_lightx2v_4step.safetensors", "high_noise_original_ckpt": "/path/to/wan2.2_i2v_A14b_high_noise_int8_lightx2v_4step.safetensors", "low_noise_original_ckpt": "/path/to/wan2.2_i2v_A14b_low_noise_int8_lightx2v_4step.safetensors", - "image_path": "/path/to/img_0.jpg", "disagg_mode": "controller", "disagg_config": { "bootstrap_addr": "192.168.0.166", diff --git a/configs/disagg/multi_node/wan22_i2v_distill_decoder.json b/configs/disagg/multi_node/wan22_i2v_distill_decoder.json index 20f136f9e..41a9a59fd 100644 --- a/configs/disagg/multi_node/wan22_i2v_distill_decoder.json +++ b/configs/disagg/multi_node/wan22_i2v_distill_decoder.json @@ -43,7 +43,6 @@ "low_noise_quantized_ckpt": "/path/to/wan2.2_i2v_A14b_low_noise_int8_lightx2v_4step.safetensors", "high_noise_original_ckpt": "/path/to/wan2.2_i2v_A14b_high_noise_int8_lightx2v_4step.safetensors", "low_noise_original_ckpt": "/path/to/wan2.2_i2v_A14b_low_noise_int8_lightx2v_4step.safetensors", - "image_path": "/path/to/img_0.jpg", "disagg_mode": "decoder", "disagg_config": { "bootstrap_addr": "192.168.0.166", diff --git a/configs/disagg/multi_node/wan22_i2v_distill_encoder.json b/configs/disagg/multi_node/wan22_i2v_distill_encoder.json index 775979410..b3bfe8394 100644 --- a/configs/disagg/multi_node/wan22_i2v_distill_encoder.json +++ b/configs/disagg/multi_node/wan22_i2v_distill_encoder.json @@ -43,7 +43,6 @@ "low_noise_quantized_ckpt": "/path/to/wan2.2_i2v_A14b_low_noise_int8_lightx2v_4step.safetensors", "high_noise_original_ckpt": "/path/to/wan2.2_i2v_A14b_high_noise_int8_lightx2v_4step.safetensors", "low_noise_original_ckpt": "/path/to/wan2.2_i2v_A14b_low_noise_int8_lightx2v_4step.safetensors", - "image_path": "/path/to/img_0.jpg", "disagg_mode": "encoder", "disagg_config": { "bootstrap_addr": "192.168.0.166", diff --git a/configs/disagg/multi_node/wan22_i2v_distill_transformer.json b/configs/disagg/multi_node/wan22_i2v_distill_transformer.json index a5ef6c90d..b08b35812 100644 --- a/configs/disagg/multi_node/wan22_i2v_distill_transformer.json +++ b/configs/disagg/multi_node/wan22_i2v_distill_transformer.json @@ -43,7 +43,6 @@ "low_noise_quantized_ckpt": "/path/to/wan2.2_i2v_A14b_low_noise_int8_lightx2v_4step.safetensors", "high_noise_original_ckpt": "/path/to/wan2.2_i2v_A14b_high_noise_int8_lightx2v_4step.safetensors", "low_noise_original_ckpt": "/path/to/wan2.2_i2v_A14b_low_noise_int8_lightx2v_4step.safetensors", - "image_path": "/path/to/img_0.jpg", "disagg_mode": "transformer", "disagg_config": { "bootstrap_addr": "192.168.0.166", diff --git a/configs/disagg/single_node/wan22_i2v_distill_controller.json b/configs/disagg/single_node/wan22_i2v_distill_controller.json index b32892f2a..8c6a742c5 100644 --- a/configs/disagg/single_node/wan22_i2v_distill_controller.json +++ b/configs/disagg/single_node/wan22_i2v_distill_controller.json @@ -43,7 +43,6 @@ "low_noise_quantized_ckpt": "/path/to/wan2.2_i2v_A14b_low_noise_int8_lightx2v_4step.safetensors", "high_noise_original_ckpt": "/path/to/wan2.2_i2v_A14b_high_noise_int8_lightx2v_4step.safetensors", "low_noise_original_ckpt": "/path/to/wan2.2_i2v_A14b_low_noise_int8_lightx2v_4step.safetensors", - "image_path": "/path/to/img_0.jpg", "disagg_mode": "controller", "disagg_config": { "bootstrap_addr": "127.0.0.1", diff --git a/configs/dreamzero/dreamzero_droid_i2va.json b/configs/dreamzero/dreamzero_droid_i2va.json index a87d065bf..222aae998 100644 --- a/configs/dreamzero/dreamzero_droid_i2va.json +++ b/configs/dreamzero/dreamzero_droid_i2va.json @@ -15,7 +15,6 @@ "cross_attn_1_type": "dreamzero_cross_fa", "cross_attn_2_type": "dreamzero_cross_fa", "rms_norm_type": "torch", - "negative_prompt": "Vibrant colors, overexposed, static, blurry details, text, subtitles, style, artwork, painting, image, still, grayscale, dull, worst quality, low quality, JPEG artifacts, ugly, mutilated, extra fingers, bad hands, bad face, deformed, disfigured, mutated limbs, fused fingers, stagnant image, cluttered background, three legs, many people in the background, walking backwards.", "num_chunks": 15, "dit_step_mask": [true, true, true, false, false, false, true, false, false, false, true, false, false, true, true, true], "obs_cam_keys": [ diff --git a/configs/dreamzero/dreamzero_droid_i2va_dist_cfg.json b/configs/dreamzero/dreamzero_droid_i2va_dist_cfg.json index dbe271729..3acc67acb 100644 --- a/configs/dreamzero/dreamzero_droid_i2va_dist_cfg.json +++ b/configs/dreamzero/dreamzero_droid_i2va_dist_cfg.json @@ -15,7 +15,6 @@ "cross_attn_1_type": "dreamzero_cross_fa", "cross_attn_2_type": "dreamzero_cross_fa", "rms_norm_type": "torch", - "negative_prompt": "Vibrant colors, overexposed, static, blurry details, text, subtitles, style, artwork, painting, image, still, grayscale, dull, worst quality, low quality, JPEG artifacts, ugly, mutilated, extra fingers, bad hands, bad face, deformed, disfigured, mutated limbs, fused fingers, stagnant image, cluttered background, three legs, many people in the background, walking backwards.", "num_chunks": 15, "dit_step_mask": [true, true, true, false, false, false, true, false, false, false, true, false, false, true, true, true], "obs_cam_keys": [ diff --git a/configs/hunyuan_video_15/vsr/hy15_i2v_480p.json b/configs/hunyuan_video_15/vsr/hy15_i2v_480p.json index 4f774fff6..a1b607c07 100755 --- a/configs/hunyuan_video_15/vsr/hy15_i2v_480p.json +++ b/configs/hunyuan_video_15/vsr/hy15_i2v_480p.json @@ -13,6 +13,7 @@ "flow_shift": 2.0, "base_resolution": "480p", "guidance_scale": 1.0, + "enable_cfg": false, "num_inference_steps": 6, "use_meanflow": true } diff --git a/configs/minimax_h3/dmd/minimax_h3_bf16_4step_single_gpu_offload.json b/configs/minimax_h3/dmd/minimax_h3_bf16_4step.json similarity index 77% rename from configs/minimax_h3/dmd/minimax_h3_bf16_4step_single_gpu_offload.json rename to configs/minimax_h3/dmd/minimax_h3_bf16_4step.json index 585021c12..b275ad7df 100644 --- a/configs/minimax_h3/dmd/minimax_h3_bf16_4step_single_gpu_offload.json +++ b/configs/minimax_h3/dmd/minimax_h3_bf16_4step.json @@ -1,5 +1,5 @@ { - "infer_steps": 5, + "infer_steps": 4, "target_video_length": 362, "target_height": 768, "target_width": 1344, @@ -22,13 +22,10 @@ "h3_step_update": "training_euler", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "lora_dynamic_apply": true, "lora_configs": [ { - "path": "lightx2v/Minimax-h3-Turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors", + "path": "/path/to/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors", "strength": 1.0, "alpha": 128 } diff --git a/configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol_attn_single_gpu_offload.json b/configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol.json similarity index 88% rename from configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol_attn_single_gpu_offload.json rename to configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol.json index ce8b2179f..4632e3cec 100644 --- a/configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol_attn_single_gpu_offload.json +++ b/configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol.json @@ -1,10 +1,9 @@ { - "infer_steps": 5, + "infer_steps": 4, "target_video_length": 362, "target_height": 768, "target_width": 1344, "fps": 24, - "target_fps": 24, "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", @@ -34,9 +33,6 @@ "h3_step_update": "training_euler", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "lora_dynamic_apply": true, "lora_configs": [ { diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step.json index 8a8ca0c62..f78e192dd 100755 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step.json @@ -1,5 +1,5 @@ { - "infer_steps": 5, + "infer_steps": 4, "target_video_length": 362, "target_height": 768, "target_width": 1344, @@ -22,16 +22,13 @@ "h3_step_update": "training_euler", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "dit_quantized": true, "dit_quant_scheme": "fp8-sgl", "dit_quantized_ckpt": "/path/to/minimax_h3_fp8.safetensors", "lora_dynamic_apply": true, "lora_configs": [ { - "path": "/path/to/pytorch_lora_weights.safetensors", + "path": "/path/to/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors", "strength": 1.0, "alpha": 128 } diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090.json index cc38a2be1..dea873de9 100644 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090.json @@ -1,5 +1,5 @@ { - "infer_steps": 5, + "infer_steps": 4, "target_video_length": 362, "target_height": 768, "target_width": 1344, @@ -18,21 +18,19 @@ "rope_type": "minimax_h3_triton_rope", "feature_caching": "NoCaching", "use_compile": true, + "warmup": true, "video_flow_shift": 6.0, "audio_flow_shift": 3.0, "h3_step_update": "training_euler", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "dit_quantized": true, "dit_quant_scheme": "fp8-sgl", "dit_quantized_ckpt": "/path/to/minimax_h3_fp8.safetensors", "lora_dynamic_apply": true, "lora_configs": [ { - "path": "lightx2v/Minimax-h3-Turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors", + "path": "/path/to/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors", "strength": 1.0, "alpha": 128 } diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8.json similarity index 74% rename from configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae.json rename to configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8.json index 4f0d182c6..0bd223eb2 100644 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8.json @@ -1,5 +1,5 @@ { - "infer_steps": 5, + "infer_steps": 4, "target_video_length": 362, "target_height": 768, "target_width": 1344, @@ -23,6 +23,7 @@ "rope_type": "torch_real_rope", "feature_caching": "NoCaching", "use_compile": true, + "warmup": true, "use_adaln_cache": true, "vae_use_compile": true, "vae_attn_type": "sage_attn2", @@ -31,19 +32,16 @@ "h3_step_update": "training_euler", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "dit_quantized": true, "dit_quant_scheme": "fp8-sgl", - "dit_quantized_ckpt": "/path/to/models/minimax_h3/h3_quantized/fp8/minimax_h3_fp8.safetensors", + "dit_quantized_ckpt": "/path/to/minimax_h3_fp8.safetensors", "video_vae_quantized": true, "video_vae_quant_scheme": "fp8-sgl", - "video_vae_quantized_ckpt": "/path/to/models/minimax_h3/h3_quantized/fp8/minimax_h3_video_vae_fp8_sgl_bias_fp16.safetensors", + "video_vae_quantized_ckpt": "/path/to/minimax_h3_video_vae_fp8_sgl_bias_fp16.safetensors", "lora_dynamic_apply": true, "lora_configs": [ { - "path": "lightx2v/Minimax-h3-Turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors", + "path": "/path/to/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors", "strength": 1.0, "alpha": 128 } diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae_sla.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sla.json similarity index 74% rename from configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae_sla.json rename to configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sla.json index 3676c0136..c51f93b18 100755 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae_sla.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sla.json @@ -1,10 +1,9 @@ { - "infer_steps": 5, + "infer_steps": 4, "target_video_length": 362, "target_height": 768, "target_width": 1344, "fps": 24, - "target_fps": 24, "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", @@ -25,6 +24,7 @@ "rope_type": "torch_real_rope", "feature_caching": "NoCaching", "use_compile": true, + "warmup": true, "vae_use_compile": true, "vae_attn_type": "sage_attn2", "video_flow_shift": 6.0, @@ -32,19 +32,16 @@ "h3_step_update": "training_euler", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "dit_quantized": true, "dit_quant_scheme": "fp8-sgl", - "dit_quantized_ckpt": "/path/to/models/minimax_h3/h3_quantized/fp8/minimax_h3_fp8.safetensors", + "dit_quantized_ckpt": "/path/to/minimax_h3_fp8.safetensors", "video_vae_quantized": true, "video_vae_quant_scheme": "fp8-sgl", - "video_vae_quantized_ckpt": "/path/to/models/minimax_h3/h3_quantized/fp8/minimax_h3_video_vae_fp8_sgl_bias_fp16.safetensors", + "video_vae_quantized_ckpt": "/path/to/minimax_h3_video_vae_fp8_sgl_bias_fp16.safetensors", "lora_dynamic_apply": true, "lora_configs": [ { - "path": "lightx2v/Minimax-h3-Turbo-SLA/minimax_h3_fl2v_turbo_4step_v0.1_768p_sla_bf16.safetensors", + "path": "/path/to/minimax_h3_fl2v_turbo_4step_v0.1_768p_sla_bf16.safetensors", "strength": 1.0, "alpha": 128 } diff --git a/configs/minimax_h3/dmd/minimax_h3_sp8_4step_5090_with_fp8_vae_sol.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sol.json similarity index 93% rename from configs/minimax_h3/dmd/minimax_h3_sp8_4step_5090_with_fp8_vae_sol.json rename to configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sol.json index 849dda762..5baccfcfd 100644 --- a/configs/minimax_h3/dmd/minimax_h3_sp8_4step_5090_with_fp8_vae_sol.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sol.json @@ -1,10 +1,9 @@ { - "infer_steps": 5, + "infer_steps": 4, "target_video_length": 362, "target_height": 768, "target_width": 1344, "fps": 24, - "target_fps": 24, "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", @@ -50,9 +49,6 @@ "h3_step_update": "training_euler", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "dit_quantized": true, "dit_quant_scheme": "fp8-sgl", "dit_quantized_ckpt": "/path/to/minimax_h3_fp8.safetensors", diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_8step.json b/configs/minimax_h3/dmd/minimax_h3_fp8_8step.json index 1f697993e..788dfd3c0 100644 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_8step.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_8step.json @@ -1,5 +1,5 @@ { - "infer_steps": 9, + "infer_steps": 8, "target_video_length": 362, "target_height": 768, "target_width": 1344, @@ -22,9 +22,6 @@ "h3_step_update": "training_euler", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "dit_quantized": true, "dit_quant_scheme": "fp8-sgl", "dit_quantized_ckpt": "/path/to/minimax_h3_fp8.safetensors", diff --git a/configs/minimax_h3/dmd/minimax_h3_int8_4step.json b/configs/minimax_h3/dmd/minimax_h3_int8_4step.json index a451bbbf1..9bd651067 100644 --- a/configs/minimax_h3/dmd/minimax_h3_int8_4step.json +++ b/configs/minimax_h3/dmd/minimax_h3_int8_4step.json @@ -1,5 +1,5 @@ { - "infer_steps": 5, + "infer_steps": 4, "target_video_length": 362, "target_height": 768, "target_width": 1344, @@ -22,16 +22,13 @@ "h3_step_update": "training_euler", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "dit_quantized": true, "dit_quant_scheme": "int8-sgl", "dit_quantized_ckpt": "/path/to/minimax_h3_int8.safetensors", "lora_dynamic_apply": true, "lora_configs": [ { - "path": "lightx2v/Minimax-h3-Turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors", + "path": "/path/to/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors", "strength": 1.0, "alpha": 128 } diff --git a/configs/minimax_h3/dmd/minimax_h3_int8_convrot_8step.json b/configs/minimax_h3/dmd/minimax_h3_int8_convrot_8step.json index 8f92d10ce..c79948621 100644 --- a/configs/minimax_h3/dmd/minimax_h3_int8_convrot_8step.json +++ b/configs/minimax_h3/dmd/minimax_h3_int8_convrot_8step.json @@ -1,5 +1,5 @@ { - "infer_steps": 9, + "infer_steps": 8, "target_video_length": 362, "target_height": 768, "target_width": 1344, @@ -22,9 +22,6 @@ "h3_step_update": "training_euler", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "dit_quantized": true, "dit_quant_scheme": "int8-convrot", "dit_quantized_ckpt": "/path/to/minimax_h3_int8_convrot.safetensors", diff --git a/configs/minimax_h3/dmd/minimax_h3_ref2av.json b/configs/minimax_h3/dmd/minimax_h3_ref2av_4step.json similarity index 88% rename from configs/minimax_h3/dmd/minimax_h3_ref2av.json rename to configs/minimax_h3/dmd/minimax_h3_ref2av_4step.json index 42e2250b8..de146a14d 100755 --- a/configs/minimax_h3/dmd/minimax_h3_ref2av.json +++ b/configs/minimax_h3/dmd/minimax_h3_ref2av_4step.json @@ -1,5 +1,5 @@ { - "infer_steps": 5, + "infer_steps": 4, "target_video_length": 362, "target_height": 768, "target_width": 1344, @@ -22,9 +22,6 @@ "h3_step_update": "training_euler", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "text_encoder_tensor_parallel": true, "lora_configs": [ { diff --git a/configs/minimax_h3/fp8/minimax_h3_t2av.json b/configs/minimax_h3/fp8/minimax_h3.json similarity index 84% rename from configs/minimax_h3/fp8/minimax_h3_t2av.json rename to configs/minimax_h3/fp8/minimax_h3.json index 8edfc1124..0d743f2f1 100644 --- a/configs/minimax_h3/fp8/minimax_h3_t2av.json +++ b/configs/minimax_h3/fp8/minimax_h3.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, @@ -20,9 +20,6 @@ "audio_flow_shift": 3.0, "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "dit_quantized": true, "dit_quant_scheme": "fp8-sgl", "dit_quantized_ckpt": "/path/to/minimax_h3_fp8.safetensors" diff --git a/configs/minimax_h3/fp8/minimax_h3_t2av_encoder_fp8.json b/configs/minimax_h3/fp8/minimax_h3_encoder_fp8.json similarity index 85% rename from configs/minimax_h3/fp8/minimax_h3_t2av_encoder_fp8.json rename to configs/minimax_h3/fp8/minimax_h3_encoder_fp8.json index e661d600d..29febcf8c 100644 --- a/configs/minimax_h3/fp8/minimax_h3_t2av_encoder_fp8.json +++ b/configs/minimax_h3/fp8/minimax_h3_encoder_fp8.json @@ -1,10 +1,9 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, "fps": 24, - "target_fps": 24, "enable_cfg": false, "cpu_offload": true, "offload_granularity": "model", @@ -24,9 +23,6 @@ "audio_flow_shift": 3.0, "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "dit_quantized": true, "dit_quant_scheme": "fp8-sgl", "dit_quantized_ckpt": "/path/to/minimax_h3_fp8.safetensors" diff --git a/configs/minimax_h3/fp8/minimax_h3_t2av_sp8_5090.json b/configs/minimax_h3/fp8/minimax_h3_sp_5090.json similarity index 80% rename from configs/minimax_h3/fp8/minimax_h3_t2av_sp8_5090.json rename to configs/minimax_h3/fp8/minimax_h3_sp_5090.json index 817c868ee..5757f4bfb 100644 --- a/configs/minimax_h3/fp8/minimax_h3_t2av_sp8_5090.json +++ b/configs/minimax_h3/fp8/minimax_h3_sp_5090.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 362, "target_height": 768, "target_width": 1344, @@ -18,16 +18,14 @@ "rope_type": "minimax_h3_triton_rope", "feature_caching": "NoCaching", "use_compile": true, + "warmup": true, "video_flow_shift": 12.0, "audio_flow_shift": 3.0, "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "dit_quantized": true, "dit_quant_scheme": "fp8-sgl", - "dit_quantized_ckpt": "/path/to/models/minimax_h3/h3_quantized/fp8/minimax_h3_fp8.safetensors", + "dit_quantized_ckpt": "/path/to/minimax_h3_fp8.safetensors", "parallel": { "seq_p_size": 8, "seq_p_attn_type": "ulysses", diff --git a/configs/minimax_h3/fp8/minimax_h3_t2av_vae_fp8.json b/configs/minimax_h3/fp8/minimax_h3_vae_fp8.json similarity index 70% rename from configs/minimax_h3/fp8/minimax_h3_t2av_vae_fp8.json rename to configs/minimax_h3/fp8/minimax_h3_vae_fp8.json index 56fda6881..ebdaa25ea 100644 --- a/configs/minimax_h3/fp8/minimax_h3_t2av_vae_fp8.json +++ b/configs/minimax_h3/fp8/minimax_h3_vae_fp8.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 768, "target_width": 1344, @@ -18,17 +18,15 @@ "rope_type": "minimax_h3_triton_rope", "feature_caching": "NoCaching", "use_compile": true, + "warmup": true, "video_flow_shift": 12.0, "audio_flow_shift": 3.0, "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "dit_quantized": true, "dit_quant_scheme": "fp8-sgl", - "dit_quantized_ckpt": "path/minimax_h3/h3_quantized/fp8/minimax_h3_fp8.safetensors", + "dit_quantized_ckpt": "/path/to/minimax_h3_fp8.safetensors", "video_vae_quantized": true, "video_vae_quant_scheme": "fp8-sgl", - "video_vae_quantized_ckpt": "path/minimax_h3/h3_quantized/fp8/minimax_h3_video_vae_fp8_sgl_bias_fp16.safetensors" + "video_vae_quantized_ckpt": "/path/to/minimax_h3_video_vae_fp8_sgl_bias_fp16.safetensors" } diff --git a/configs/minimax_h3/minimax_h3_fl2av.json b/configs/minimax_h3/minimax_h3.json similarity index 77% rename from configs/minimax_h3/minimax_h3_fl2av.json rename to configs/minimax_h3/minimax_h3.json index 62d0daaa1..660248780 100644 --- a/configs/minimax_h3/minimax_h3_fl2av.json +++ b/configs/minimax_h3/minimax_h3.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, @@ -19,8 +19,5 @@ "video_flow_shift": 12.0, "audio_flow_shift": 3.0, "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true + "audio_sampling_rate": 32000 } diff --git a/configs/minimax_h3/minimax_h3_t2av_block_offload.json b/configs/minimax_h3/minimax_h3_block_offload.json similarity index 79% rename from configs/minimax_h3/minimax_h3_t2av_block_offload.json rename to configs/minimax_h3/minimax_h3_block_offload.json index a0162c18a..8bd0ddcf2 100644 --- a/configs/minimax_h3/minimax_h3_t2av_block_offload.json +++ b/configs/minimax_h3/minimax_h3_block_offload.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 768, "target_width": 1344, @@ -20,8 +20,5 @@ "video_flow_shift": 12.0, "audio_flow_shift": 3.0, "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true + "audio_sampling_rate": 32000 } diff --git a/configs/minimax_h3/minimax_h3_t2av_compile.json b/configs/minimax_h3/minimax_h3_compile.json similarity index 78% rename from configs/minimax_h3/minimax_h3_t2av_compile.json rename to configs/minimax_h3/minimax_h3_compile.json index 8a3289a46..0ea702fcf 100644 --- a/configs/minimax_h3/minimax_h3_t2av_compile.json +++ b/configs/minimax_h3/minimax_h3_compile.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, @@ -20,8 +20,5 @@ "video_flow_shift": 12.0, "audio_flow_shift": 3.0, "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true + "audio_sampling_rate": 32000 } diff --git a/configs/minimax_h3/minimax_h3_fl2av_compile.json b/configs/minimax_h3/minimax_h3_fl2av_compile.json deleted file mode 100644 index 93f02caf9..000000000 --- a/configs/minimax_h3/minimax_h3_fl2av_compile.json +++ /dev/null @@ -1,26 +0,0 @@ -{ - "infer_steps": 30, - "target_video_length": 124, - "target_height": 544, - "target_width": 960, - "fps": 24, - "enable_cfg": false, - "cpu_offload": true, - "offload_granularity": "model", - "text_encoder_cpu_offload": true, - "vae_cpu_offload": true, - "lazy_load": false, - "unload_modules": false, - "attn_type": "sage_attn2", - "rms_type": "sgl-kernel", - "rope_type": "torch_real_rope", - "feature_caching": "NoCaching", - "use_compile": true, - "video_flow_shift": 12.0, - "audio_flow_shift": 3.0, - "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true -} diff --git a/configs/minimax_h3/minimax_h3_i2av.json b/configs/minimax_h3/minimax_h3_i2av.json deleted file mode 100644 index 62d0daaa1..000000000 --- a/configs/minimax_h3/minimax_h3_i2av.json +++ /dev/null @@ -1,26 +0,0 @@ -{ - "infer_steps": 30, - "target_video_length": 124, - "target_height": 544, - "target_width": 960, - "fps": 24, - "enable_cfg": false, - "cpu_offload": true, - "offload_granularity": "model", - "text_encoder_cpu_offload": true, - "vae_cpu_offload": true, - "lazy_load": false, - "unload_modules": false, - "attn_type": "sage_attn2", - "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", - "feature_caching": "NoCaching", - "use_compile": false, - "video_flow_shift": 12.0, - "audio_flow_shift": 3.0, - "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true -} diff --git a/configs/minimax_h3/minimax_h3_i2av_compile.json b/configs/minimax_h3/minimax_h3_i2av_compile.json deleted file mode 100644 index 93f02caf9..000000000 --- a/configs/minimax_h3/minimax_h3_i2av_compile.json +++ /dev/null @@ -1,26 +0,0 @@ -{ - "infer_steps": 30, - "target_video_length": 124, - "target_height": 544, - "target_width": 960, - "fps": 24, - "enable_cfg": false, - "cpu_offload": true, - "offload_granularity": "model", - "text_encoder_cpu_offload": true, - "vae_cpu_offload": true, - "lazy_load": false, - "unload_modules": false, - "attn_type": "sage_attn2", - "rms_type": "sgl-kernel", - "rope_type": "torch_real_rope", - "feature_caching": "NoCaching", - "use_compile": true, - "video_flow_shift": 12.0, - "audio_flow_shift": 3.0, - "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true -} diff --git a/configs/minimax_h3/minimax_h3_l2av.json b/configs/minimax_h3/minimax_h3_l2av.json deleted file mode 100644 index 62d0daaa1..000000000 --- a/configs/minimax_h3/minimax_h3_l2av.json +++ /dev/null @@ -1,26 +0,0 @@ -{ - "infer_steps": 30, - "target_video_length": 124, - "target_height": 544, - "target_width": 960, - "fps": 24, - "enable_cfg": false, - "cpu_offload": true, - "offload_granularity": "model", - "text_encoder_cpu_offload": true, - "vae_cpu_offload": true, - "lazy_load": false, - "unload_modules": false, - "attn_type": "sage_attn2", - "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", - "feature_caching": "NoCaching", - "use_compile": false, - "video_flow_shift": 12.0, - "audio_flow_shift": 3.0, - "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true -} diff --git a/configs/minimax_h3/minimax_h3_l2av_compile.json b/configs/minimax_h3/minimax_h3_l2av_compile.json deleted file mode 100644 index 93f02caf9..000000000 --- a/configs/minimax_h3/minimax_h3_l2av_compile.json +++ /dev/null @@ -1,26 +0,0 @@ -{ - "infer_steps": 30, - "target_video_length": 124, - "target_height": 544, - "target_width": 960, - "fps": 24, - "enable_cfg": false, - "cpu_offload": true, - "offload_granularity": "model", - "text_encoder_cpu_offload": true, - "vae_cpu_offload": true, - "lazy_load": false, - "unload_modules": false, - "attn_type": "sage_attn2", - "rms_type": "sgl-kernel", - "rope_type": "torch_real_rope", - "feature_caching": "NoCaching", - "use_compile": true, - "video_flow_shift": 12.0, - "audio_flow_shift": 3.0, - "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true -} diff --git a/configs/minimax_h3/minimax_h3_ref2av.json b/configs/minimax_h3/minimax_h3_ref2av.json deleted file mode 100644 index 62d0daaa1..000000000 --- a/configs/minimax_h3/minimax_h3_ref2av.json +++ /dev/null @@ -1,26 +0,0 @@ -{ - "infer_steps": 30, - "target_video_length": 124, - "target_height": 544, - "target_width": 960, - "fps": 24, - "enable_cfg": false, - "cpu_offload": true, - "offload_granularity": "model", - "text_encoder_cpu_offload": true, - "vae_cpu_offload": true, - "lazy_load": false, - "unload_modules": false, - "attn_type": "sage_attn2", - "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", - "feature_caching": "NoCaching", - "use_compile": false, - "video_flow_shift": 12.0, - "audio_flow_shift": 3.0, - "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true -} diff --git a/configs/minimax_h3/minimax_h3_ref2av_compile.json b/configs/minimax_h3/minimax_h3_ref2av_compile.json deleted file mode 100644 index 93f02caf9..000000000 --- a/configs/minimax_h3/minimax_h3_ref2av_compile.json +++ /dev/null @@ -1,26 +0,0 @@ -{ - "infer_steps": 30, - "target_video_length": 124, - "target_height": 544, - "target_width": 960, - "fps": 24, - "enable_cfg": false, - "cpu_offload": true, - "offload_granularity": "model", - "text_encoder_cpu_offload": true, - "vae_cpu_offload": true, - "lazy_load": false, - "unload_modules": false, - "attn_type": "sage_attn2", - "rms_type": "sgl-kernel", - "rope_type": "torch_real_rope", - "feature_caching": "NoCaching", - "use_compile": true, - "video_flow_shift": 12.0, - "audio_flow_shift": 3.0, - "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true -} diff --git a/configs/minimax_h3/minimax_h3_t2av_sol_attn_block_offload.json b/configs/minimax_h3/minimax_h3_sol_block_offload.json similarity index 82% rename from configs/minimax_h3/minimax_h3_t2av_sol_attn_block_offload.json rename to configs/minimax_h3/minimax_h3_sol_block_offload.json index 9ee04dc1d..70d92eb47 100644 --- a/configs/minimax_h3/minimax_h3_t2av_sol_attn_block_offload.json +++ b/configs/minimax_h3/minimax_h3_sol_block_offload.json @@ -1,10 +1,9 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 362, "target_height": 768, "target_width": 1344, "fps": 24, - "target_fps": 24, "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", @@ -32,8 +31,5 @@ "video_flow_shift": 12.0, "audio_flow_shift": 3.0, "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true + "audio_sampling_rate": 32000 } diff --git a/configs/minimax_h3/minimax_h3_t2av_sp.json b/configs/minimax_h3/minimax_h3_sp.json similarity index 84% rename from configs/minimax_h3/minimax_h3_t2av_sp.json rename to configs/minimax_h3/minimax_h3_sp.json index 6f4c2fa5f..71b57b82c 100644 --- a/configs/minimax_h3/minimax_h3_t2av_sp.json +++ b/configs/minimax_h3/minimax_h3_sp.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, @@ -21,9 +21,6 @@ "audio_flow_shift": 3.0, "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "parallel": { "seq_p_size": 4, "seq_p_attn_type": "ulysses" diff --git a/configs/minimax_h3/minimax_h3_t2av.json b/configs/minimax_h3/minimax_h3_t2av.json deleted file mode 100644 index 47aaaf461..000000000 --- a/configs/minimax_h3/minimax_h3_t2av.json +++ /dev/null @@ -1,30 +0,0 @@ -{ - "infer_steps": 30, - "target_video_length": 362, - "target_height": 1344, - "target_width": 768, - "fps": 24, - "enable_cfg": false, - "cpu_offload": true, - "offload_granularity": "model", - "text_encoder_cpu_offload": true, - "vae_cpu_offload": true, - "lazy_load": false, - "unload_modules": false, - "attn_type": "sage_attn2", - "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", - "feature_caching": "NoCaching", - "use_compile": false, - "video_flow_shift": 12.0, - "audio_flow_shift": 3.0, - "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, - "parallel": { - "seq_p_size": 4, - "seq_p_attn_type": "ulysses" - } -} diff --git a/configs/minimax_h3/minimax_h3_t2av_dmd_lora_4step.json b/configs/minimax_h3/minimax_h3_t2av_dmd_lora_4step.json deleted file mode 100644 index d045efafc..000000000 --- a/configs/minimax_h3/minimax_h3_t2av_dmd_lora_4step.json +++ /dev/null @@ -1,36 +0,0 @@ -{ - "infer_steps": 5, - "target_video_length": 124, - "target_height": 544, - "target_width": 960, - "fps": 24, - "enable_cfg": false, - "cpu_offload": true, - "offload_granularity": "block", - "text_encoder_cpu_offload": true, - "text_encoder_offload_granularity": "block", - "vae_cpu_offload": true, - "lazy_load": false, - "unload_modules": false, - "attn_type": "flash_attn3", - "rms_type": "torch_native", - "rope_type": "torch_real_rope", - "feature_caching": "NoCaching", - "use_compile": false, - "video_flow_shift": 6.0, - "audio_flow_shift": 3.0, - "h3_step_update": "training_euler", - "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, - "lora_dynamic_apply": false, - "lora_configs": [ - { - "path": "/path/to/pytorch_lora_weights.safetensors", - "strength": 1.0, - "alpha": 128 - } - ] -} diff --git a/configs/minimax_h3/minimax_h3_t2av_video_codec_options.json b/configs/minimax_h3/minimax_h3_t2av_video_codec_options.json deleted file mode 100644 index 726ab931d..000000000 --- a/configs/minimax_h3/minimax_h3_t2av_video_codec_options.json +++ /dev/null @@ -1,31 +0,0 @@ -{ - "infer_steps": 30, - "target_video_length": 362, - "target_height": 544, - "target_width": 960, - "fps": 24, - "target_fps": 24, - "enable_cfg": false, - "cpu_offload": true, - "offload_granularity": "model", - "text_encoder_cpu_offload": true, - "vae_cpu_offload": true, - "lazy_load": false, - "unload_modules": false, - "attn_type": "sage_attn2", - "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", - "feature_caching": "NoCaching", - "use_compile": false, - "video_flow_shift": 12.0, - "audio_flow_shift": 3.0, - "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "video_codec_options": { - "preset": "ultrafast", - "crf": "18" - }, - "keep_latents_dtype_in_scheduler": true -} diff --git a/configs/minimax_h3/minimax_h3_t2av_tp.json b/configs/minimax_h3/minimax_h3_tp.json similarity index 84% rename from configs/minimax_h3/minimax_h3_t2av_tp.json rename to configs/minimax_h3/minimax_h3_tp.json index 01317d9b8..88f0e533f 100644 --- a/configs/minimax_h3/minimax_h3_t2av_tp.json +++ b/configs/minimax_h3/minimax_h3_tp.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, @@ -22,9 +22,6 @@ "audio_flow_shift": 3.0, "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "parallel": { "tensor_p_size": 2 } diff --git a/configs/minimax_h3/minimax_h3_t2av_tp_sp.json b/configs/minimax_h3/minimax_h3_tp_sp.json similarity index 85% rename from configs/minimax_h3/minimax_h3_t2av_tp_sp.json rename to configs/minimax_h3/minimax_h3_tp_sp.json index 1f31748cf..c7c10025d 100644 --- a/configs/minimax_h3/minimax_h3_t2av_tp_sp.json +++ b/configs/minimax_h3/minimax_h3_tp_sp.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, @@ -22,9 +22,6 @@ "audio_flow_shift": 3.0, "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, "parallel": { "tensor_p_size": 2, "seq_p_size": 2, diff --git a/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_15s.json b/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_15s.json index e7a79c483..bd8d77607 100644 --- a/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_15s.json +++ b/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_15s.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 362, "target_height": 768, "target_width": 1344, diff --git a/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_5s.json b/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_5s.json index bb90a282f..623d98623 100644 --- a/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_5s.json +++ b/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_5s.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 120, "target_height": 768, "target_width": 1344, diff --git a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_dmd_lora_4step_sp_tp.json b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_dmd_lora_4step_sp_tp.json index b6d493086..4bf1738c9 100644 --- a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_dmd_lora_4step_sp_tp.json +++ b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_dmd_lora_4step_sp_tp.json @@ -1,5 +1,5 @@ { - "infer_steps": 5, + "infer_steps": 4, "target_video_length": 124, "target_height": 768, "target_width": 1344, diff --git a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_int8_sp_tp.json b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_int8_sp_tp.json index d8cfd96a2..42f1ce43d 100644 --- a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_int8_sp_tp.json +++ b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_int8_sp_tp.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp.json b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp.json index 2ab902462..837a18206 100644 --- a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp.json +++ b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp_cpu_offload.json b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp_cpu_offload.json index 03551258f..541897853 100644 --- a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp_cpu_offload.json +++ b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp_cpu_offload.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_tp.json b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_tp.json index 799e9755b..ae4b580f1 100644 --- a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_tp.json +++ b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_tp.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json b/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json index d01f864fe..ea60dc55b 100644 --- a/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json +++ b/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json @@ -1,5 +1,5 @@ { - "infer_steps": 5, + "infer_steps": 4, "target_video_length": 362, "target_height": 768, "target_width": 1344, diff --git a/configs/platforms/intel_xpu/minimax_h3_t2av.json b/configs/platforms/intel_xpu/minimax_h3_t2av.json index 8f2d83055..82c499443 100644 --- a/configs/platforms/intel_xpu/minimax_h3_t2av.json +++ b/configs/platforms/intel_xpu/minimax_h3_t2av.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/intel_xpu/minimax_h3_t2av_dmd_lora_4step.json b/configs/platforms/intel_xpu/minimax_h3_t2av_dmd_lora_4step.json index 754576a7c..72f53c852 100644 --- a/configs/platforms/intel_xpu/minimax_h3_t2av_dmd_lora_4step.json +++ b/configs/platforms/intel_xpu/minimax_h3_t2av_dmd_lora_4step.json @@ -1,5 +1,5 @@ { - "infer_steps": 5, + "infer_steps": 4, "target_video_length": 124, "target_height": 768, "target_width": 1344, diff --git a/configs/platforms/intel_xpu/minimax_h3_t2av_fp8.json b/configs/platforms/intel_xpu/minimax_h3_t2av_fp8.json index 13e94d921..910bc1090 100644 --- a/configs/platforms/intel_xpu/minimax_h3_t2av_fp8.json +++ b/configs/platforms/intel_xpu/minimax_h3_t2av_fp8.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/intel_xpu/minimax_h3_t2av_int8.json b/configs/platforms/intel_xpu/minimax_h3_t2av_int8.json index a7dac66bf..a9ad4848b 100644 --- a/configs/platforms/intel_xpu/minimax_h3_t2av_int8.json +++ b/configs/platforms/intel_xpu/minimax_h3_t2av_int8.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/metax/minimax_h3_t2av_tp1_block_offload.json b/configs/platforms/metax/minimax_h3_t2av_tp1_block_offload.json index 13450eff3..b723935ce 100644 --- a/configs/platforms/metax/minimax_h3_t2av_tp1_block_offload.json +++ b/configs/platforms/metax/minimax_h3_t2av_tp1_block_offload.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/metax/minimax_h3_t2av_tp_sp.json b/configs/platforms/metax/minimax_h3_t2av_tp_sp.json index dc2254d22..eac89a061 100644 --- a/configs/platforms/metax/minimax_h3_t2av_tp_sp.json +++ b/configs/platforms/metax/minimax_h3_t2av_tp_sp.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/mlu/minimax_h3_t2av_sp.json b/configs/platforms/mlu/minimax_h3_t2av_sp.json index e916eba12..dfe81eeef 100644 --- a/configs/platforms/mlu/minimax_h3_t2av_sp.json +++ b/configs/platforms/mlu/minimax_h3_t2av_sp.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/mlu/minimax_h3_t2av_tp.json b/configs/platforms/mlu/minimax_h3_t2av_tp.json index 8a2e35d07..ea9dfa860 100644 --- a/configs/platforms/mlu/minimax_h3_t2av_tp.json +++ b/configs/platforms/mlu/minimax_h3_t2av_tp.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/mlu/minimax_h3_t2av_tp_sp.json b/configs/platforms/mlu/minimax_h3_t2av_tp_sp.json index 906dff7b0..f3341f721 100644 --- a/configs/platforms/mlu/minimax_h3_t2av_tp_sp.json +++ b/configs/platforms/mlu/minimax_h3_t2av_tp_sp.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/mthreads_musa/minimax_h3_t2av.json b/configs/platforms/mthreads_musa/minimax_h3_t2av.json index 39dd7b37a..c40e57d2d 100644 --- a/configs/platforms/mthreads_musa/minimax_h3_t2av.json +++ b/configs/platforms/mthreads_musa/minimax_h3_t2av.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/mthreads_musa/minimax_h3_t2av_fp8.json b/configs/platforms/mthreads_musa/minimax_h3_t2av_fp8.json index e71b02944..d5a0b2f33 100644 --- a/configs/platforms/mthreads_musa/minimax_h3_t2av_fp8.json +++ b/configs/platforms/mthreads_musa/minimax_h3_t2av_fp8.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp.json b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp.json index 06ab5f4ed..9842643a2 100644 --- a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp.json +++ b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp4_sp2.json b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp4_sp2.json index e0301007f..a60275c76 100644 --- a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp4_sp2.json +++ b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp4_sp2.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp_fp8.json b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp_fp8.json index 596b4ba78..de92bce18 100644 --- a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp_fp8.json +++ b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp_fp8.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 544, "target_width": 960, diff --git a/configs/seko_talk/shot/rs2v/rs2v.json b/configs/seko_talk/shot/rs2v/rs2v.json index b03e7efd6..2392196a8 100644 --- a/configs/seko_talk/shot/rs2v/rs2v.json +++ b/configs/seko_talk/shot/rs2v/rs2v.json @@ -1,7 +1,10 @@ { "model_cls": "seko_talk", "task": "rs2v", - "model_path":"/path/to/SekoTalk-v2.7_beta2-bf16-step4", + "model_path": "/path/to/SekoTalk-v2.7_beta2-bf16-step4", + "infer_steps": 4, + "target_video_length": 81, + "resize_mode": "adaptive", "target_fps": 16, "audio_sr": 16000, "self_attn_1_type": "flash_attn3", @@ -11,17 +14,5 @@ "sample_shift": 5, "enable_cfg": false, "use_31_block": true, - "target_video_length": 81, - "prev_frame_length": 0, - - "default_input_info": - { - "infer_steps": 4, - "resize_mode": "adaptive", - "prompt": "The video features a male speaking to the camera with arms spread out, a slightly furrowed brow, and a focused gaze.", - "negative_prompt": "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走", - "image_path": "assets/inputs/audio/seko_input.png", - "audio_path": "assets/inputs/audio/seko_input.mp3", - "save_result_path": "save_results/output_seko_talk_shot_rs2v.mp4" - } + "prev_frame_length": 0 } diff --git a/configs/seko_talk/shot/stream/f2v.json b/configs/seko_talk/shot/stream/f2v.json index 5359149f4..a77a9077f 100644 --- a/configs/seko_talk/shot/stream/f2v.json +++ b/configs/seko_talk/shot/stream/f2v.json @@ -1,7 +1,10 @@ { "model_cls": "seko_talk", "task": "s2v", - "model_path":"/path/to/Wan2.1-i2V1202-Audio-14B-720P", + "model_path": "/path/to/Wan2.1-i2V1202-Audio-14B-720P", + "infer_steps": 4, + "target_video_length": 33, + "resize_mode": "adaptive", "target_fps": 16, "audio_sr": 16000, "self_attn_1_type": "flash_attn3", @@ -12,7 +15,6 @@ "enable_cfg": false, "use_31_block": true, "rope_type": "torch_complex_rope", - "target_video_length": 33, "prev_frame_length": 1, "f2v_process": true, "cpu_offload": true, @@ -29,15 +31,5 @@ "path": "/path/to/lightx2v_I2V_14B_480p_cfg_step_distill_rank32_bf16.safetensors", "strength": 1.0 } - ], - "default_input_info": - { - "infer_steps": 4, - "resize_mode": "adaptive", - "prompt": "The video features a male speaking to the camera with arms spread out, a slightly furrowed brow, and a focused gaze.", - "negative_prompt": "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走", - "image_path": "assets/inputs/audio/seko_input.png", - "audio_path": "assets/inputs/audio/seko_input.mp3", - "save_result_path": "save_results/output_seko_talk_shot_stream.mp4" - } + ] } diff --git a/configs/seko_talk/shot/stream/s2v.json b/configs/seko_talk/shot/stream/s2v.json index c4dfcd529..78b0f8911 100644 --- a/configs/seko_talk/shot/stream/s2v.json +++ b/configs/seko_talk/shot/stream/s2v.json @@ -1,7 +1,10 @@ { "model_cls": "seko_talk", "task": "s2v", - "model_path":"/path/to/Wan2.1-R2V721-Audio-14B-720P", + "model_path": "/path/to/Wan2.1-R2V721-Audio-14B-720P", + "infer_steps": 3, + "target_video_length": 33, + "resize_mode": "adaptive", "target_fps": 16, "audio_sr": 16000, "self_attn_1_type": "flash_attn3", @@ -12,7 +15,6 @@ "enable_cfg": false, "use_31_block": true, "rope_type": "torch_complex_rope", - "target_video_length": 33, "prev_frame_length": 5, "cpu_offload": true, "offload_granularity": "block", @@ -22,16 +24,5 @@ "offload_ratio": 1, "use_tiling_vae": true, "audio_encoder_cpu_offload": true, - "audio_adapter_cpu_offload": false, - - "default_input_info": - { - "infer_steps": 4, - "resize_mode": "adaptive", - "prompt": "The video features a male speaking to the camera with arms spread out, a slightly furrowed brow, and a focused gaze.", - "negative_prompt": "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走", - "image_path": "assets/inputs/audio/seko_input.png", - "audio_path": "assets/inputs/audio/seko_input.mp3", - "save_result_path": "save_results/output_seko_talk_shot_stream.mp4" - } + "audio_adapter_cpu_offload": false } diff --git a/docs/EN/source/getting_started/model_structure.md b/docs/EN/source/getting_started/model_structure.md index 6edfc288b..9c1fc6266 100644 --- a/docs/EN/source/getting_started/model_structure.md +++ b/docs/EN/source/getting_started/model_structure.md @@ -197,7 +197,7 @@ lightx2v_path=/path/to/LightX2V # Run script cd LightX2V/scripts -bash wan/run_wan_i2v_distill_4step_cfg.sh +bash wan/distill/run_wan_i2v_distill_4step_cfg.sh ``` > 💡 **Tip**: When there's only one model file in the directory, LightX2V will automatically load it. @@ -272,7 +272,7 @@ Edit configuration file (e.g., `configs/distill/wan_i2v_distill_4step_cfg.json`) ```bash cd LightX2V/scripts -bash wan/run_wan_i2v_distill_4step_cfg.sh +bash wan/distill/run_wan_i2v_distill_4step_cfg.sh ``` > 💡 **Tip**: Other components (T5, CLIP, VAE, tokenizer, etc.) need to be manually organized into the model directory @@ -315,10 +315,10 @@ lightx2v_path=/path/to/LightX2V # Run script cd LightX2V/scripts -bash wan22/run_wan22_moe_i2v_distill.sh +bash wan22/distill/run_wan22_moe_i2v_distill_fp8_4step.sh ``` -> 💡 **Tip**: When there's only one model file in each subdirectory, LightX2V will automatically load it. +> 💡 **Tip**: Set `high_noise_quantized_ckpt` and `low_noise_quantized_ckpt` in `configs/distill/wan22/wan_moe_i2v_distill_quant.json` to the two downloaded files before running the script. #### Scenario B: Multiple Model Files Per Directory diff --git a/docs/EN/source/method_tutorials/attention.md b/docs/EN/source/method_tutorials/attention.md index 585b3d445..718dcfcc2 100644 --- a/docs/EN/source/method_tutorials/attention.md +++ b/docs/EN/source/method_tutorials/attention.md @@ -50,9 +50,10 @@ The example config is `configs/attentions/wan_i2v_sol_attn.json`. It enables Mor The MiniMax-H3 example retains the 15-second, 768p, block CPU-offload setup: +Set `lightx2v_path` and `model_path` in `scripts/minimax_h3/run_minimax_h3_t2av.sh` to your local directories, and change its `--config_json` argument to `"${lightx2v_path}/configs/minimax_h3/minimax_h3_sol_block_offload.json"` before running. + ```bash -MODEL_PATH=/path/to/MiniMax-H3 \ - bash scripts/minimax_h3/run_minimax_h3_t2av_sol_attn_offload.sh +bash scripts/minimax_h3/run_minimax_h3_t2av.sh ``` -Its config is `configs/minimax_h3/minimax_h3_t2av_sol_attn_block_offload.json`. Sol-Attn is used only by the 50-layer main transformer; the short text refiner uses dense Torch SDPA. The first six denoising steps and transformer layer 0 use SageAttention2 through `dense_backend=sage_attn2`. H3 attention uses a mixed `[text | audio | video]` packed sequence rather than one 3D video grid, so this config uses `reorder=none`; Wan's Morton3D reorder must not be enabled directly. +Its config is `configs/minimax_h3/minimax_h3_sol_block_offload.json`. Sol-Attn is used only by the 50-layer main transformer; the short text refiner uses dense Torch SDPA. The first six denoising steps and transformer layer 0 use SageAttention2 through `dense_backend=sage_attn2`. H3 attention uses a mixed `[text | audio | video]` packed sequence rather than one 3D video grid, so this config uses `reorder=none`; Wan's Morton3D reorder must not be enabled directly. diff --git a/docs/ZH_CN/source/getting_started/model_structure.md b/docs/ZH_CN/source/getting_started/model_structure.md index 8ceb48a10..4f7370436 100644 --- a/docs/ZH_CN/source/getting_started/model_structure.md +++ b/docs/ZH_CN/source/getting_started/model_structure.md @@ -197,7 +197,7 @@ lightx2v_path=/path/to/LightX2V # 运行脚本 cd LightX2V/scripts -bash wan/run_wan_i2v_distill_4step_cfg.sh +bash wan/distill/run_wan_i2v_distill_4step_cfg.sh ``` > 💡 **提示**:当目录下只有一个模型文件时,LightX2V 会自动加载该文件。 @@ -272,7 +272,7 @@ wan2.1_i2v_720p_multi/ ```bash cd LightX2V/scripts -bash wan/run_wan_i2v_distill_4step_cfg.sh +bash wan/distill/run_wan_i2v_distill_4step_cfg.sh ``` ### Wan2.2 单文件模型 @@ -313,10 +313,10 @@ lightx2v_path=/path/to/LightX2V # 运行脚本 cd LightX2V/scripts -bash wan22/run_wan22_moe_i2v_distill.sh +bash wan22/distill/run_wan22_moe_i2v_distill_fp8_4step.sh ``` -> 💡 **提示**:当每个子目录下只有一个模型文件时,LightX2V 会自动加载。 +> 💡 **提示**:运行脚本前,需要将 `configs/distill/wan22/wan_moe_i2v_distill_quant.json` 中的 `high_noise_quantized_ckpt` 和 `low_noise_quantized_ckpt` 设置为上面下载的两个文件。 #### 场景 B:每个目录下有多个模型文件 diff --git a/docs/ZH_CN/source/method_tutorials/attention.md b/docs/ZH_CN/source/method_tutorials/attention.md index 44bd9bf4c..091d79469 100644 --- a/docs/ZH_CN/source/method_tutorials/attention.md +++ b/docs/ZH_CN/source/method_tutorials/attention.md @@ -50,9 +50,10 @@ MODEL_PATH=/path/to/Wan2.1-I2V-14B-480P \ MiniMax-H3 的示例沿用 15 秒、768p、block CPU offload 配置: +运行前,将 `scripts/minimax_h3/run_minimax_h3_t2av.sh` 中的 `lightx2v_path` 和 `model_path` 改为本地目录,并将其 `--config_json` 参数改为 `"${lightx2v_path}/configs/minimax_h3/minimax_h3_sol_block_offload.json"`。 + ```bash -MODEL_PATH=/path/to/MiniMax-H3 \ - bash scripts/minimax_h3/run_minimax_h3_t2av_sol_attn_offload.sh +bash scripts/minimax_h3/run_minimax_h3_t2av.sh ``` -对应配置为 `configs/minimax_h3/minimax_h3_t2av_sol_attn_block_offload.json`。Sol-Attn 仅用于 50 层主 Transformer,短文本 refiner 使用 dense Torch SDPA;前 6 个去噪步骤和第 0 层通过 `dense_backend=sage_attn2` 使用 SageAttention2。H3 的 attention 序列按 `[text | audio | video]` 混合打包,不是单一的三维视频网格,因此该配置使用 `reorder=none`,不能直接启用 Wan 的 Morton3D 重排。 +对应配置为 `configs/minimax_h3/minimax_h3_sol_block_offload.json`。Sol-Attn 仅用于 50 层主 Transformer,短文本 refiner 使用 dense Torch SDPA;前 6 个去噪步骤和第 0 层通过 `dense_backend=sage_attn2` 使用 SageAttention2。H3 的 attention 序列按 `[text | audio | video]` 混合打包,不是单一的三维视频网格,因此该配置使用 `reorder=none`,不能直接启用 Wan 的 Morton3D 重排。 diff --git a/examples/minimax_h3/minimax_h3_t2av_dmd.py b/examples/minimax_h3/minimax_h3_t2av_dmd.py index e91c0204b..c847980e1 100644 --- a/examples/minimax_h3/minimax_h3_t2av_dmd.py +++ b/examples/minimax_h3/minimax_h3_t2av_dmd.py @@ -1,6 +1,6 @@ """MiniMax-H3 4-step 768p T2AV inference with the released DMD LoRA. -Before running, set ``model_path`` to a local MiniMax-H3 model directory and +Before running, set ``MODEL_PATH`` to a local MiniMax-H3 model directory and ensure the LoRA path in the selected config resolves to a local checkpoint. """ @@ -11,8 +11,8 @@ from lightx2v import LightX2VPipeline -MODEL_PATH = "/data/nvme6/gushiqiao/models/MiniMax-H3" -CONFIG_PATH = "configs/minimax_h3/dmd/minimax_h3_bf16_4step_single_gpu_offload.json" +MODEL_PATH = "/path/to/MiniMax-H3" +CONFIG_PATH = "configs/minimax_h3/dmd/minimax_h3_bf16_4step.json" OUTPUT_PATH = "save_results/minimax_h3_t2av_dmd_768p.mp4" diff --git a/examples/neopp/README.md b/examples/neopp/README.md index 95f69706e..708f47528 100644 --- a/examples/neopp/README.md +++ b/examples/neopp/README.md @@ -14,10 +14,20 @@ uses these LightX2V interfaces: bytes, and wraps `_run_infer_step` to check cancellation. - `runner.set_kvcache(...)` injects conditioning; `set_inference_params(...)` supplies matching position offsets, CFG settings and output format. -- `generate(seed=None, save_result_path="", target_shape=[height, width])` +- `generate(task="t2i", seed=None, save_result_path="", target_shape=[height, width])` continues the RNG state restored by LightLLM for later images in a session. An explicit integer seed starts a seeded generation. +Pipelines initialized with only `support_tasks` require `task` on every generation +call. The linked LightLLM adapter currently omits it; add +`task="t2i" if is_t2i else "i2i"` to its +`self.pipe.generate(...)` call when upgrading LightX2V. Keep its KV injection and +session RNG handling unchanged. + +An explicit constructor `task`, such as `LightX2VPipeline(..., task="t2i")`, is the +default for calls that omit it. Passing `generate(task="i2i", ...)` selects a task +for that request without changing the default. + Default image editing also uses `set_kvcache`: LightLLM includes image conditioning in the two KV branches. The optional `image_guidance_scale != 1` mode uses a separate three-branch interface that is not implemented by this @@ -27,7 +37,8 @@ The Python files in this directory replay KV dumps for backend development and performance debugging. Replace their `/path/to/...` placeholders with the checkpoint, KV files and output directory for your model and request. The filenames and position offsets in the examples describe particular dumps; -update both to match your captured request. Each turn calls `load_kvcache(...)` +update both and select `task="t2i"` or `task="i2i"` to match your captured request. +Each turn calls `load_kvcache(...)` and `set_inference_params(...)` before `generate(...)`; `save_result_for_debug=True` saves the image to the requested file. @@ -105,8 +116,8 @@ bash /path/to/LightX2V/scripts/neopp/run_neopp_dense_t2i_1k.sh ``` For image editing, capture KV from a LightLLM request that includes the input -images. CFG settings come from the selected JSON. In `replay_kv.py`, an explicit -seed overrides the JSON seed; if both are omitted, the default is 42. +images. CFG settings come from the selected JSON. In `replay_kv.py`, the seed +comes from `--seed` and defaults to 42 when omitted. `run_neopp_dense_i2i_1k_cfg3.sh` has been removed because the current runner does not implement three-branch image guidance. Use the supported two-branch image diff --git a/examples/neopp/neopp_dense_1k.py b/examples/neopp/neopp_dense_1k.py index 1e5540d3c..d8b2eeb9a 100644 --- a/examples/neopp/neopp_dense_1k.py +++ b/examples/neopp/neopp_dense_1k.py @@ -35,6 +35,7 @@ ) pipe.generate( + task="t2i", seed=200, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_1k_0.png", target_shape=[1024, 1024], # Height, Width @@ -58,6 +59,7 @@ ) pipe.generate( + task="t2i", seed=201, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_1k_1.png", target_shape=[1024, 1024], # Height, Width @@ -81,6 +83,7 @@ ) pipe.generate( + task="t2i", seed=202, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_1k_2.png", target_shape=[1024, 1024], # Height, Width diff --git a/examples/neopp/neopp_dense_1k_fp8.py b/examples/neopp/neopp_dense_1k_fp8.py index 8ba5680b9..35770e5bd 100644 --- a/examples/neopp/neopp_dense_1k_fp8.py +++ b/examples/neopp/neopp_dense_1k_fp8.py @@ -35,6 +35,7 @@ ) pipe.generate( + task="t2i", seed=200, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_1k_fp8_0.png", target_shape=[1024, 1024], # Height, Width @@ -58,6 +59,7 @@ ) pipe.generate( + task="t2i", seed=201, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_1k_fp8_1.png", target_shape=[1024, 1024], # Height, Width @@ -81,6 +83,7 @@ ) pipe.generate( + task="t2i", seed=202, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_1k_fp8_2.png", target_shape=[1024, 1024], # Height, Width diff --git a/examples/neopp/neopp_dense_1k_parallel_cfg.py b/examples/neopp/neopp_dense_1k_parallel_cfg.py index 44dff8395..f9866a09a 100644 --- a/examples/neopp/neopp_dense_1k_parallel_cfg.py +++ b/examples/neopp/neopp_dense_1k_parallel_cfg.py @@ -37,6 +37,7 @@ ) pipe.generate( + task="t2i", seed=200, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_1k_0.png", target_shape=[1024, 1024], # Height, Width @@ -60,6 +61,7 @@ ) pipe.generate( + task="t2i", seed=201, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_1k_1.png", target_shape=[1024, 1024], # Height, Width @@ -83,6 +85,7 @@ ) pipe.generate( + task="t2i", seed=202, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_1k_2.png", target_shape=[1024, 1024], # Height, Width diff --git a/examples/neopp/neopp_dense_1k_parallel_cfg_seq.py b/examples/neopp/neopp_dense_1k_parallel_cfg_seq.py index c88053241..4ee3e2225 100644 --- a/examples/neopp/neopp_dense_1k_parallel_cfg_seq.py +++ b/examples/neopp/neopp_dense_1k_parallel_cfg_seq.py @@ -37,6 +37,7 @@ ) pipe.generate( + task="t2i", seed=200, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_1k_0.png", target_shape=[1024, 1024], # Height, Width @@ -60,6 +61,7 @@ ) pipe.generate( + task="t2i", seed=201, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_1k_1.png", target_shape=[1024, 1024], # Height, Width @@ -83,6 +85,7 @@ ) pipe.generate( + task="t2i", seed=202, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_1k_2.png", target_shape=[1024, 1024], # Height, Width diff --git a/examples/neopp/neopp_dense_2k.py b/examples/neopp/neopp_dense_2k.py index 9b5e5b41f..ff791e245 100644 --- a/examples/neopp/neopp_dense_2k.py +++ b/examples/neopp/neopp_dense_2k.py @@ -35,6 +35,7 @@ ) pipe.generate( + task="t2i", seed=200, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_2k_0.png", target_shape=[2048, 2048], # Height, Width @@ -58,6 +59,7 @@ ) pipe.generate( + task="t2i", seed=None, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_2k_1.png", target_shape=[2048, 2048], # Height, Width @@ -81,6 +83,7 @@ ) pipe.generate( + task="t2i", seed=None, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_2k_2.png", target_shape=[2048, 2048], # Height, Width diff --git a/examples/neopp/neopp_dense_2k_8steps.py b/examples/neopp/neopp_dense_2k_8steps.py index 1ceb35a93..e5265b38d 100644 --- a/examples/neopp/neopp_dense_2k_8steps.py +++ b/examples/neopp/neopp_dense_2k_8steps.py @@ -35,6 +35,7 @@ ) pipe.generate( + task="t2i", seed=200, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_2k_0.png", target_shape=[2048, 2048], # Height, Width @@ -58,6 +59,7 @@ # ) # pipe.generate( +# task="t2i", # seed=None, # save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_2k_1.png", # target_shape=[2048, 2048], # Height, Width @@ -81,6 +83,7 @@ # ) # pipe.generate( +# task="t2i", # seed=None, # save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_2k_2.png", # target_shape=[2048, 2048], # Height, Width diff --git a/examples/neopp/neopp_dense_2k_fp8.py b/examples/neopp/neopp_dense_2k_fp8.py index 88080d4c4..262f3c945 100644 --- a/examples/neopp/neopp_dense_2k_fp8.py +++ b/examples/neopp/neopp_dense_2k_fp8.py @@ -35,6 +35,7 @@ ) pipe.generate( + task="t2i", seed=200, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_2k_fp8_0.png", target_shape=[2048, 2048], # Height, Width @@ -58,6 +59,7 @@ ) pipe.generate( + task="t2i", seed=201, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_2k_fp8_1.png", target_shape=[2048, 2048], # Height, Width @@ -81,6 +83,7 @@ ) pipe.generate( + task="t2i", seed=202, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_2k_fp8_2.png", target_shape=[2048, 2048], # Height, Width diff --git a/examples/neopp/neopp_dense_2k_parallel_cfg_seq.py b/examples/neopp/neopp_dense_2k_parallel_cfg_seq.py index 3be81e95f..5714b2188 100644 --- a/examples/neopp/neopp_dense_2k_parallel_cfg_seq.py +++ b/examples/neopp/neopp_dense_2k_parallel_cfg_seq.py @@ -37,6 +37,7 @@ ) pipe.generate( + task="t2i", seed=200, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_2k_0.png", target_shape=[2048, 2048], # Height, Width @@ -60,6 +61,7 @@ ) pipe.generate( + task="t2i", seed=201, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_2k_1.png", target_shape=[2048, 2048], # Height, Width @@ -83,6 +85,7 @@ ) pipe.generate( + task="t2i", seed=202, save_result_path="/path/to/save_results/output_lightx2v_neopp_dense_2k_2.png", target_shape=[2048, 2048], # Height, Width diff --git a/examples/neopp/neopp_moe_1k.py b/examples/neopp/neopp_moe_1k.py index 1f31e2515..56c74d62f 100644 --- a/examples/neopp/neopp_moe_1k.py +++ b/examples/neopp/neopp_moe_1k.py @@ -35,6 +35,7 @@ ) pipe.generate( + task="t2i", seed=200, save_result_path="/path/to/save_results/output_lightx2v_neopp_moe_1k_0.png", target_shape=[1024, 1024], # Height, Width diff --git a/examples/neopp/neopp_moe_2k.py b/examples/neopp/neopp_moe_2k.py index bc2830628..beac034fa 100644 --- a/examples/neopp/neopp_moe_2k.py +++ b/examples/neopp/neopp_moe_2k.py @@ -35,6 +35,7 @@ ) pipe.generate( + task="t2i", seed=200, save_result_path="/path/to/save_results/output_lightx2v_neopp_moe_2k_0.png", target_shape=[2048, 2048], # Height, Width diff --git a/examples/neopp/replay_kv.py b/examples/neopp/replay_kv.py index 25af03c21..b3ff7962d 100644 --- a/examples/neopp/replay_kv.py +++ b/examples/neopp/replay_kv.py @@ -1,5 +1,4 @@ import argparse -import json import torch.distributed as dist @@ -13,14 +12,11 @@ parser.add_argument("--uncond_kv") parser.add_argument("--index_offset_cond", type=int, required=True) parser.add_argument("--index_offset_uncond", type=int) -parser.add_argument("--seed", type=int) +parser.add_argument("--seed", type=int, default=42) parser.add_argument("--target_shape", type=int, nargs=2, required=True, metavar=("HEIGHT", "WIDTH")) parser.add_argument("--save_result_path", required=True) args = parser.parse_args() -with open(args.config_json, "r") as f: - config = json.load(f) - pipe = LightX2VPipeline( model_path=args.model_path, model_cls="neopp", @@ -37,9 +33,8 @@ index_offset_uncond=args.index_offset_uncond, **inference_params, ) -seed = args.seed if args.seed is not None else config.get("seed") pipe.generate( - seed=42 if seed is None else seed, + seed=args.seed, target_shape=args.target_shape, save_result_path=args.save_result_path, ) diff --git a/examples/sensenova_vision/example_visualize.py b/examples/sensenova_vision/example_visualize.py index 503244bab..098e221e6 100755 --- a/examples/sensenova_vision/example_visualize.py +++ b/examples/sensenova_vision/example_visualize.py @@ -4,15 +4,13 @@ import argparse import sys -from argparse import Namespace from pathlib import Path import torch from PIL import Image from lightx2v.models.runners.bagel.sensenova_vision_runner import SenseNovaVisionRunner -from lightx2v.utils.input_info import SenseNovaVisionInputInfo -from lightx2v.utils.set_config import set_config +from lightx2v.utils.set_config import build_startup_config def parse_args(): @@ -67,14 +65,14 @@ def main(): visualize_panoptic_segmentation, ) - config_args = Namespace( - model_cls="sensenova_vision", - task="omni_vision_task", - model_path=args.model_path, - config_json=str(lightx2v_root / "configs/sensenova_vision/sensenova_vision.json"), - seed=args.seed, + config = build_startup_config( + { + "model_cls": "sensenova_vision", + "task": "omni_vision_task", + "model_path": args.model_path, + "config_json": str(lightx2v_root / "configs/sensenova_vision/sensenova_vision.json"), + } ) - config = set_config(config_args) config["sensenova_source_path"] = str(source_root) runner = SenseNovaVisionRunner(config) runner.init_modules() @@ -87,15 +85,16 @@ def source_file(relative_path): return str(source_root / relative_path) def run(subtask, image_paths, prompt, save_name="", seed=None, **kwargs): - info = SenseNovaVisionInputInfo( - seed=args.seed if seed is None else seed, - prompt=prompt, - image_path=",".join(source_file(path) for path in image_paths), - save_result_path=str(output_dir / save_name) if save_name else "", - omni_vision_subtask=subtask, + request_data = { + "seed": args.seed if seed is None else seed, + "prompt": prompt, + "image_path": ",".join(source_file(path) for path in image_paths), + "save_result_path": str(output_dir / save_name) if save_name else "", + "omni_vision_subtask": subtask, **kwargs, - ) - return runner.run_pipeline(info) + } + input_info = runner.prepare_request(request_data) + return runner.run_request(input_info) # 1. General understanding. if selected("01"): diff --git a/examples/wan/wan_animate.py b/examples/wan/wan_animate.py index 317f34d71..d17f41144 100755 --- a/examples/wan/wan_animate.py +++ b/examples/wan/wan_animate.py @@ -34,7 +34,7 @@ model_cls="wan2.2_animate", task="animate", ) -pipe.replace_flag = True # Set to True for replace mode, False for animate mode +pipe.startup_config["replace_flag"] = True # Set to True for replace mode, False for animate mode # Alternative: create generator from config JSON file # pipe.create_generator( @@ -55,7 +55,6 @@ seed = 42 prompt = "视频中的人在做动作" -negative_prompt = "镜头晃动,色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" src_pose_path = "../save_results/animate/process_results/src_pose.mp4" src_face_path = "../save_results/animate/process_results/src_face.mp4" src_ref_images = "../save_results/animate/process_results/src_ref.png" @@ -67,6 +66,5 @@ src_face_path=src_face_path, src_ref_images=src_ref_images, prompt=prompt, - negative_prompt=negative_prompt, save_result_path=save_result_path, ) diff --git a/examples/worldmirror/run_worldmirror.py b/examples/worldmirror/run_worldmirror.py index 1ee07faca..e37a4a3e4 100644 --- a/examples/worldmirror/run_worldmirror.py +++ b/examples/worldmirror/run_worldmirror.py @@ -12,15 +12,15 @@ Single-GPU, exported-model dir (default subfolder):: python examples/worldmirror/run_worldmirror.py \\ - --input_path /workspace/HY-World-2.0/examples/worldrecon/realistic/Workspace \\ - --pretrained_model_name_or_path /data/nvme1/models/HY-World-2.0 \\ + --input_path /path/to/HY-World-2.0/examples/worldrecon/realistic/Workspace \\ + --pretrained_model_name_or_path /path/to/HY-World-2.0 \\ --no_interactive Multi-GPU (sequence-parallel) + bf16:: torchrun --nproc_per_node=2 examples/worldmirror/run_worldmirror.py \\ --input_path /path/to/images \\ - --pretrained_model_name_or_path /data/nvme1/models/HY-World-2.0 \\ + --pretrained_model_name_or_path /path/to/HY-World-2.0 \\ --enable_bf16 --no_interactive Training-output format (separate yaml + ckpt):: @@ -54,7 +54,7 @@ def build_parser(): p.add_argument("--strict_output_path", type=str, default=None, help="If set, save results directly to this path (no subdir/timestamp).") # --- Model loading --- - p.add_argument("--pretrained_model_name_or_path", type=str, default="/data/nvme1/models/HY-World-2.0", help="Local directory containing HY-WorldMirror-2.0 weights.") + p.add_argument("--pretrained_model_name_or_path", type=str, default="/path/to/HY-World-2.0", help="Local directory containing HY-WorldMirror-2.0 weights.") p.add_argument("--subfolder", type=str, default="HY-WorldMirror-2.0", help="Subfolder inside the model directory.") p.add_argument("--config_path", type=str, default=None, help="Optional training YAML; used with --ckpt_path.") p.add_argument("--ckpt_path", type=str, default=None, help="Optional .ckpt/.safetensors; used with --config_path.") @@ -182,20 +182,16 @@ def build_config(args): return config_dict -def run_one(runner, input_path, strict_output_path, output_path, task_name="recon"): - from lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict - - input_info = init_empty_input_info(task_name) - update_input_info_from_dict( - input_info, +def run_one(runner, input_path, strict_output_path, output_path): + input_info = runner.prepare_request( { "input_path": input_path, "save_result_path": output_path, "strict_output_path": strict_output_path, "return_result_tensor": False, - }, + } ) - return runner.run_pipeline(input_info) + return runner.run_request(input_info) def main(): @@ -205,21 +201,15 @@ def main(): import torch import torch.distributed as dist - from lightx2v.models.runners.worldmirror.worldmirror_runner import ( - WorldMirrorRunner, # noqa: F401 — registers "worldmirror" - _broadcast_string, - ) - from lightx2v.utils.lockable_dict import LockableDict - from lightx2v.utils.registry_factory import RUNNER_REGISTER + from lightx2v.models.runners.runner_factory import build_runner + from lightx2v.models.runners.worldmirror.worldmirror_runner import _broadcast_string + from lightx2v.utils.set_config import build_startup_config config_dict = build_config(args) if args.strict_output_path is None: os.makedirs(args.output_path, exist_ok=True) - config = LockableDict(config_dict) - - runner = RUNNER_REGISTER[config["model_cls"]](config) - runner.init_modules() + runner = build_runner(build_startup_config(config_dict)) is_distributed = runner.is_distributed rank = runner.rank diff --git a/examples/worldmirror/test_worldmirror.py b/examples/worldmirror/test_worldmirror.py index fc34d4409..fe7efc806 100644 --- a/examples/worldmirror/test_worldmirror.py +++ b/examples/worldmirror/test_worldmirror.py @@ -2,8 +2,8 @@ Equivalent to the HY-World-2.0 CLI: python -m hyworld2.worldrecon.pipeline \ - --input_path examples/worldrecon/realistic/Workspace \ - --pretrained_model_name_or_path /data/nvme1/models/HY-World-2.0 \ + --input_path /path/to/HY-World-2.0/examples/worldrecon/realistic/Workspace \ + --pretrained_model_name_or_path /path/to/HY-World-2.0 \ --no_interactive Run: @@ -21,10 +21,10 @@ if _REPO_ROOT not in sys.path: sys.path.insert(0, _REPO_ROOT) -DEFAULT_CONFIG_PATH = "/workspace/LightX2V/configs/worldmirror/worldmirror_recon.json" -DEFAULT_MODEL_PATH = "/data/nvme1/models/HY-World-2.0" -DEFAULT_INPUT_PATH = "/workspace/HY-World-2.0/examples/worldrecon/realistic/Workspace" -DEFAULT_OUTPUT_PATH = "/workspace/LightX2V/save_results/HY-WorldMirror/" +DEFAULT_CONFIG_PATH = os.path.join(_REPO_ROOT, "configs/worldmirror/worldmirror_recon.json") +DEFAULT_MODEL_PATH = "/path/to/HY-World-2.0" +DEFAULT_INPUT_PATH = "/path/to/HY-World-2.0/examples/worldrecon/realistic/Workspace" +DEFAULT_OUTPUT_PATH = os.path.join(_REPO_ROOT, "save_results/HY-WorldMirror") def main(): @@ -37,13 +37,8 @@ def main(): parser.add_argument("--enable_bf16", action="store_true") args = parser.parse_args() - # Importing lightx2v no longer eagerly pulls in every runner — see - # lightx2v/__init__.py's PEP 562 __getattr__ — so we can just import - # the runner we need directly. - from lightx2v.models.runners.worldmirror.worldmirror_runner import WorldMirrorRunner # noqa: F401 - from lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict - from lightx2v.utils.lockable_dict import LockableDict - from lightx2v.utils.registry_factory import RUNNER_REGISTER + from lightx2v.models.runners.runner_factory import build_runner + from lightx2v.utils.set_config import build_startup_config with open(args.config_path, "r") as f: config_dict = json.load(f) @@ -54,23 +49,17 @@ def main(): os.makedirs(args.output_path, exist_ok=True) - config = LockableDict(config_dict) + runner = build_runner(build_startup_config(config_dict)) - runner = RUNNER_REGISTER[config["model_cls"]](config) - runner.init_modules() - - input_info = init_empty_input_info(config["task"]) - update_input_info_from_dict( - input_info, + input_info = runner.prepare_request( { "input_path": args.input_path, "save_result_path": args.output_path, "strict_output_path": args.strict_output_path, "return_result_tensor": True, - }, + } ) - - result = runner.run_pipeline(input_info) + result = runner.run_request(input_info) print(f"[test_worldmirror] output: {result}") return result diff --git a/examples/worldplay/test_worldplay_ar.py b/examples/worldplay/test_worldplay_ar.py index 6b070a9ea..55e024467 100644 --- a/examples/worldplay/test_worldplay_ar.py +++ b/examples/worldplay/test_worldplay_ar.py @@ -2,11 +2,12 @@ import os # Paths -CONFIG_PATH = "/workspace/LightX2V/configs/worldplay/worldplay_ar_i2v_480p.json" -MODEL_PATH = "/data/nvme1/models/hunyuan/HunyuanVideo-1.5" -ACTION_CKPT = "/data/nvme1/models/hunyuan/HY-WorldPlay/ar_model/diffusion_pytorch_model.safetensors" -IMAGE_PATH = "/workspace/HY-WorldPlay/assets/img/test.png" -OUTPUT_PATH = "/workspace/LightX2V/save_results/HY-WorldPlay/" +LIGHTX2V_PATH = "/path/to/LightX2V" +CONFIG_PATH = os.path.join(LIGHTX2V_PATH, "configs/worldplay/worldplay_ar_i2v_480p.json") +MODEL_PATH = "/path/to/HunyuanVideo-1.5" +ACTION_CKPT = "/path/to/HY-WorldPlay/ar_model/diffusion_pytorch_model.safetensors" +IMAGE_PATH = "/path/to/HY-WorldPlay/assets/img/test.png" +OUTPUT_PATH = os.path.join(LIGHTX2V_PATH, "save_results/HY-WorldPlay") # Input parameters PROMPT = "A paved pathway leads towards a stone arch bridge spanning a calm body of water. Lush green trees and foliage line the path and the far bank of the water. A traditional-style pavilion with a tiered, reddish-brown roof sits on the far shore. The water reflects the surrounding greenery and the sky. The scene is bathed in soft, natural light, creating a tranquil and serene atmosphere." @@ -17,9 +18,8 @@ def main(): - from lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict - from lightx2v.utils.lockable_dict import LockableDict - from lightx2v.utils.registry_factory import RUNNER_REGISTER + from lightx2v.models.runners.runner_factory import build_runner + from lightx2v.utils.set_config import build_startup_config # Load config from JSON with open(CONFIG_PATH, "r") as f: @@ -30,27 +30,20 @@ def main(): config_dict["action_ckpt"] = ACTION_CKPT config_dict["transformer_model_path"] = os.path.join(MODEL_PATH, "transformer/480p_i2v") - config = LockableDict(config_dict) - - runner = RUNNER_REGISTER[config["model_cls"]](config) - - runner.init_modules() + runner = build_runner(build_startup_config(config_dict)) # Prepare input info input_data = { "seed": SEED, "prompt": PROMPT, - "negative_prompt": "", "image_path": IMAGE_PATH, "save_result_path": os.path.join(OUTPUT_PATH, "worldplay_ar_test.mp4"), "return_result_tensor": False, "pose": POSE, } - input_info = init_empty_input_info("i2v") - update_input_info_from_dict(input_info, input_data) - - result = runner.run_pipeline(input_info) + input_info = runner.prepare_request(input_data) + result = runner.run_request(input_info) return result diff --git a/examples/worldplay/test_worldplay_bi.py b/examples/worldplay/test_worldplay_bi.py index 45e8d34cc..7fbd9dd9c 100644 --- a/examples/worldplay/test_worldplay_bi.py +++ b/examples/worldplay/test_worldplay_bi.py @@ -2,11 +2,12 @@ import os # Paths -CONFIG_PATH = "/workspace/LightX2V/configs/worldplay/worldplay_bi_i2v_480p.json" -MODEL_PATH = "/data/nvme1/models/hunyuan/hf_cache/hub/models--tencent--HunyuanVideo-1.5/snapshots/9b49404b3f5df2a8f0b31df27a0c7ab872e7b038" -ACTION_CKPT = "/data/nvme1/models/hunyuan/HY-WorldPlay/bidirectional_model/diffusion_pytorch_model.safetensors" -IMAGE_PATH = "/workspace/HY-WorldPlay/assets/img/test.png" -OUTPUT_PATH = "/workspace/LightX2V/save_results/HY-WorldPlay/" +LIGHTX2V_PATH = "/path/to/LightX2V" +CONFIG_PATH = os.path.join(LIGHTX2V_PATH, "configs/worldplay/worldplay_bi_i2v_480p.json") +MODEL_PATH = "/path/to/HunyuanVideo-1.5" +ACTION_CKPT = "/path/to/HY-WorldPlay/bidirectional_model/diffusion_pytorch_model.safetensors" +IMAGE_PATH = "/path/to/HY-WorldPlay/assets/img/test.png" +OUTPUT_PATH = os.path.join(LIGHTX2V_PATH, "save_results/HY-WorldPlay") # Input parameters PROMPT = "A paved pathway leads towards a stone arch bridge spanning a calm body of water. Lush green trees and foliage line the path and the far bank of the water. A traditional-style pavilion with a tiered, reddish-brown roof sits on the far shore. The water reflects the surrounding greenery and the sky. The scene is bathed in soft, natural light, creating a tranquil and serene atmosphere." @@ -17,9 +18,8 @@ def main(): - from lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict - from lightx2v.utils.lockable_dict import LockableDict - from lightx2v.utils.registry_factory import RUNNER_REGISTER + from lightx2v.models.runners.runner_factory import build_runner + from lightx2v.utils.set_config import build_startup_config # Load config from JSON with open(CONFIG_PATH, "r") as f: @@ -30,11 +30,7 @@ def main(): config_dict["action_ckpt"] = ACTION_CKPT config_dict["transformer_model_path"] = os.path.join(MODEL_PATH, "transformer/480p_i2v") - config = LockableDict(config_dict) - - runner = RUNNER_REGISTER[config["model_cls"]](config) - - runner.init_modules() + runner = build_runner(build_startup_config(config_dict)) # Prepare input info input_data = { @@ -47,10 +43,8 @@ def main(): "pose": POSE, } - input_info = init_empty_input_info("i2v") - update_input_info_from_dict(input_info, input_data) - - result = runner.run_pipeline(input_info) + input_info = runner.prepare_request(input_data) + result = runner.run_request(input_info) return result diff --git a/examples/worldplay/test_worldplay_distill.py b/examples/worldplay/test_worldplay_distill.py index ade86eab1..a90f10036 100644 --- a/examples/worldplay/test_worldplay_distill.py +++ b/examples/worldplay/test_worldplay_distill.py @@ -2,11 +2,12 @@ import os # Paths -CONFIG_PATH = "/workspace/LightX2V/configs/worldplay/worldplay_distill_i2v_480p.json" -MODEL_PATH = "/data/nvme1/models/hunyuan/hf_cache/hub/models--tencent--HunyuanVideo-1.5/snapshots/9b49404b3f5df2a8f0b31df27a0c7ab872e7b038" -ACTION_CKPT = "/data/nvme1/models/hunyuan/HY-WorldPlay/ar_distilled_action_model/diffusion_pytorch_model.safetensors" -IMAGE_PATH = "/workspace/HY-WorldPlay/assets/img/test.png" -OUTPUT_PATH = "/workspace/LightX2V/save_results/HY-WorldPlay/" +LIGHTX2V_PATH = "/path/to/LightX2V" +CONFIG_PATH = os.path.join(LIGHTX2V_PATH, "configs/worldplay/worldplay_distill_i2v_480p.json") +MODEL_PATH = "/path/to/HunyuanVideo-1.5" +ACTION_CKPT = "/path/to/HY-WorldPlay/ar_distilled_action_model/diffusion_pytorch_model.safetensors" +IMAGE_PATH = "/path/to/HY-WorldPlay/assets/img/test.png" +OUTPUT_PATH = os.path.join(LIGHTX2V_PATH, "save_results/HY-WorldPlay") # Input parameters PROMPT = "A paved pathway leads towards a stone arch bridge spanning a calm body of water. Lush green trees and foliage line the path and the far bank of the water. A traditional-style pavilion with a tiered, reddish-brown roof sits on the far shore. The water reflects the surrounding greenery and the sky. The scene is bathed in soft, natural light, creating a tranquil and serene atmosphere. The pathway is composed of large, rectangular stones, and the bridge is constructed of light gray stone. The overall composition emphasizes the peaceful and harmonious nature of the landscape." @@ -17,9 +18,8 @@ def main(): - from lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict - from lightx2v.utils.lockable_dict import LockableDict - from lightx2v.utils.registry_factory import RUNNER_REGISTER + from lightx2v.models.runners.runner_factory import build_runner + from lightx2v.utils.set_config import build_startup_config # Load config from JSON with open(CONFIG_PATH, "r") as f: @@ -30,27 +30,20 @@ def main(): config_dict["action_ckpt"] = ACTION_CKPT config_dict["transformer_model_path"] = os.path.join(MODEL_PATH, "transformer/480p_i2v") - config = LockableDict(config_dict) - - runner = RUNNER_REGISTER[config["model_cls"]](config) - - runner.init_modules() + runner = build_runner(build_startup_config(config_dict)) # Prepare input info input_data = { "seed": SEED, "prompt": PROMPT, - "negative_prompt": "", "image_path": IMAGE_PATH, "save_result_path": os.path.join(OUTPUT_PATH, "worldplay_distill_test.mp4"), "return_result_tensor": False, "pose": POSE, } - input_info = init_empty_input_info("i2v") - update_input_info_from_dict(input_info, input_data) - - result = runner.run_pipeline(input_info) + input_info = runner.prepare_request(input_data) + result = runner.run_request(input_info) return result diff --git a/lightx2v/__init__.py b/lightx2v/__init__.py index a2250870d..034b34002 100755 --- a/lightx2v/__init__.py +++ b/lightx2v/__init__.py @@ -1,4 +1,4 @@ -__version__ = "0.1.0" +__version__ = "0.5.0" __author__ = "LightX2V Contributors" __license__ = "Apache 2.0" diff --git a/lightx2v/disagg/README.md b/lightx2v/disagg/README.md index c7efd0862..35289c33a 100644 --- a/lightx2v/disagg/README.md +++ b/lightx2v/disagg/README.md @@ -106,7 +106,12 @@ bash scripts/disagg/run_dynamic.sh | `SEED` | 随机种子。 | `42` | | `PROMPT` | 文本提示词。 | 脚本内置示例 prompt | | `NEGATIVE_PROMPT` | 负向提示词。 | 脚本内置示例 negative prompt | -| `SAVE_RESULT_PATH` | 最终视频保存路径。 | `save_results/wan22_i2v_dynamic.mp4` | +| `SAVE_RESULT_PATH` | 自动批次的输出基础路径,也作为 user 压测输出前缀的默认值。 | `save_results/wan22_i2v_dynamic.mp4` | +| `DISAGG_WORKLOAD_SAVE_PREFIX` | user 压测输出基础路径;请求生成端追加阶段名称和请求编号。 | 脚本使用 `SAVE_RESULT_PATH` | + +独立 Disagg 服务通过文件交付结果,每个生成请求必须携带非空 `save_path`,缺失时会在派发前报错。自动批次缺少基础路径会直接失败;外部请求缺少路径会记录为该请求失败,并继续接收后续请求。只启动 Encoder、Transformer 或 Decoder 进程时无需提供输出路径,路径随请求传入。 + +`LOAD_FROM_USER=0` 时,Controller 将基础路径按实际 `room` 编号展开,例如 `output.mp4` 对应 `output0.mp4`、`output1.mp4`。外部请求的显式路径保持原样;`LOAD_FROM_USER=1` 的脚本复用 workload 的前缀功能生成不同文件名,并将 prompt、图片和 seed 传给请求生成端。 ## 推荐的常见组合 diff --git a/lightx2v/disagg/disagg_mixin.py b/lightx2v/disagg/disagg_mixin.py index b9a33ae53..53854eeb6 100644 --- a/lightx2v/disagg/disagg_mixin.py +++ b/lightx2v/disagg/disagg_mixin.py @@ -43,12 +43,24 @@ except ImportError: RDMAClient = None from lightx2v.utils.envs import GET_DTYPE +from lightx2v.utils.utils import seed_all from lightx2v_platform.base.global_var import AI_DEVICE logger = logging.getLogger(__name__) _DISAGG_PROFILING = os.environ.get("DISAGG_PROFILING", "0") == "1" +_DISAGG_REQUEST_FIELDS = ( + "prompt", + "negative_prompt", + "save_result_path", + "return_result_tensor", + "seed", + "target_video_length", + "aspect_ratio", + "i2i_denoise_strength", +) + def _prof_log(tag: str, elapsed: float): if _DISAGG_PROFILING: @@ -101,6 +113,26 @@ def _estimate_encoder_buffer_sizes(config) -> List[int]: return buffer_sizes +def validate_disagg_buffer_capacity(buffers: List[torch.Tensor], required_sizes: List[int], phase: str) -> None: + capacities = [buffer.numel() for buffer in buffers] + if len(capacities) != len(required_sizes) or any(required > capacity for required, capacity in zip(required_sizes, capacities)): + raise ValueError( + f"[Disagg] {phase} request exceeds the configured transfer buffer capacity: " + f"required={required_sizes}, capacity={capacities}. " + "Increase target_video_length, target_height, or target_width in every stage's startup config, then restart the services." + ) + + +def wait_for_disagg_transfer(transfer, description: str) -> None: + while True: + status = transfer.poll() + if status == DataPoll.Success: + return + if status == DataPoll.Failed: + raise RuntimeError(f"[Disagg] {description} failed") + time.sleep(0.01) + + def _buffer_view(buf: torch.Tensor, dtype: torch.dtype, shape: tuple) -> torch.Tensor: """Create a typed view over a raw uint8 buffer without copying.""" view = torch.empty(0, dtype=dtype, device=buf.device) @@ -166,6 +198,7 @@ def init_disagg(self, config): self._disagg_active_encoder_room: Optional[int] = None self._disagg_active_transformer_room: Optional[int] = None self._disagg_active_decoder_room: Optional[int] = None + self._disagg_request_config: Optional[Dict[str, Any]] = None if self._disagg_mode == "encoder": if self._disagg_decentralized: @@ -321,54 +354,71 @@ def _disagg_json_safe_value(self, obj: Any) -> Any: return {str(k): self._disagg_json_safe_value(v) for k, v in obj.items()} return str(obj) - def _disagg_build_request_config_snapshot(self) -> Dict[str, Any]: + def resolve_request_seed(self, request_data): + # CLI preparation can precede init_disagg; dispatched requests already carry the upstream seed. + request_config = getattr(self, "_disagg_request_config", None) + if request_config is None or "seed" not in request_config: + if self.config.get("disagg_mode") in ("transformer", "decode"): + # Static workers receive the resolved seed with the tensor metadata. + return None + return super().resolve_request_seed(request_data) + return request_config["seed"] + + def build_disagg_request_config(self, input_info, request_config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + """Build the generation and routing state for one Disagg request.""" + source_config = self._disagg_request_config if request_config is None else request_config + request_config = dict(source_config or {}) + disagg_cfg = self.config.get("disagg_config", {}) + + if not self._disagg_decentralized: + request_config.setdefault("data_bootstrap_room", int(disagg_cfg.get("bootstrap_room", self._disagg_bootstrap_room))) + + for key in _DISAGG_REQUEST_FIELDS: + value = getattr(input_info, key, None) + if value is not None: + request_config[key] = value + + target_shape = getattr(input_info, "target_shape", None) + if target_shape: + request_config["target_shape"] = list(target_shape) + request_config["target_height"], request_config["target_width"] = target_shape + + return request_config + + def _disagg_effective_config(self, request_config: Dict[str, Any]) -> Dict[str, Any]: + config = dict(self.config) + request_config = dict(request_config) + request_disagg_cfg = request_config.pop("disagg_config", None) + config.update(request_config) + if request_disagg_cfg is not None: + config["disagg_config"] = {**dict(self.config.get("disagg_config", {})), **dict(request_disagg_cfg)} + return config + + def _disagg_build_request_config_snapshot(self, request_config: Dict[str, Any]) -> Dict[str, Any]: """Payload for phase1/phase2 ring: per-request fields for workers.""" - disagg_cfg = dict(self.config.get("disagg_config", {}) or {}) - room = int(self.config.get("data_bootstrap_room", disagg_cfg.get("bootstrap_room", self._disagg_bootstrap_room))) - payload: Dict[str, Any] = { - "data_bootstrap_room": room, - "task": self.config.get("task"), - "model_cls": self.config.get("model_cls"), - } + config = self._disagg_effective_config(request_config) + disagg_cfg = dict(config.get("disagg_config", {}) or {}) + room = int(config.get("data_bootstrap_room", disagg_cfg.get("bootstrap_room", self._disagg_bootstrap_room))) + payload: Dict[str, Any] = {"data_bootstrap_room": room} for key in ( - "seed", - "infer_steps", - "aspect_ratio", - "enable_cfg", - "sample_guide_scale", + *_DISAGG_REQUEST_FIELDS, "target_height", "target_width", - "text_len", "controller_result_host", "controller_result_port", + "target_shape", ): - if key in self.config and self.config.get(key) is not None: - payload[key] = self._disagg_json_safe_value(self.config.get(key)) - - ii = getattr(self, "input_info", None) - if ii is not None: - if getattr(ii, "prompt", None): - payload["prompt"] = ii.prompt - if getattr(ii, "negative_prompt", None) is not None: - payload["negative_prompt"] = ii.negative_prompt - if getattr(ii, "save_result_path", None): - payload["save_result_path"] = ii.save_result_path - payload["save_path"] = ii.save_result_path - if getattr(ii, "target_shape", None) is not None: - payload["target_shape"] = self._disagg_json_safe_value(ii.target_shape) - if getattr(ii, "aspect_ratio", None) and "aspect_ratio" not in payload: - payload["aspect_ratio"] = ii.aspect_ratio - if getattr(ii, "seed", None) is not None and "seed" not in payload: - payload["seed"] = ii.seed - - dpr = self.config.get("disagg_phase1_receiver_engine_rank") - if dpr is not None: - try: - payload["disagg_phase1_receiver_engine_rank"] = int(dpr) - except (TypeError, ValueError): - pass + if key in config and config.get(key) is not None: + payload[key] = self._disagg_json_safe_value(config.get(key)) + + if payload.get("save_result_path"): + payload["save_path"] = payload["save_result_path"] - return {k: v for k, v in payload.items() if v is not None} + phase1_receiver_rank = config.get("disagg_phase1_receiver_engine_rank") + if phase1_receiver_rank is not None: + payload["disagg_phase1_receiver_engine_rank"] = int(phase1_receiver_rank) + + return payload def _disagg_connect_queue_client( self, @@ -479,29 +529,25 @@ def disagg_transformer_prepare_dispatch(self, packet: Dict[str, Any]) -> None: req = dict(packet.get("request_config") or {}) enc_addr = str(packet.get("encoder_node_address", "127.0.0.1")) - - with self.config.temporarily_unlocked(): - self.config.update(req) - - disagg_cfg = self.config.get("disagg_config", {}) - room = int(self.config.get("data_bootstrap_room", disagg_cfg.get("bootstrap_room", 0))) - self.disagg_transformer_teardown_session() + self._disagg_request_config = None + config = self._disagg_effective_config(req) + disagg_cfg = config.get("disagg_config", {}) + room = int(config.get("data_bootstrap_room", disagg_cfg.get("bootstrap_room", 0))) - self._disagg_sender_rank = int(disagg_cfg.get("sender_engine_rank", self._disagg_sender_rank)) + sender_rank = int(disagg_cfg.get("sender_engine_rank", self._disagg_sender_rank)) pkt_recv_rank = req.get("disagg_phase1_receiver_engine_rank") if pkt_recv_rank is not None: - self._disagg_receiver_rank = int(pkt_recv_rank) + receiver_rank = int(pkt_recv_rank) else: - self._disagg_receiver_rank = int(disagg_cfg.get("receiver_engine_rank", self._disagg_receiver_rank)) + receiver_rank = int(disagg_cfg.get("receiver_engine_rank", self._disagg_receiver_rank)) - buffer_sizes = _estimate_encoder_buffer_sizes(self.config) - self._disagg_alloc_buffers(buffer_sizes) + self._disagg_alloc_buffers(packet["buffer_sizes"]) data_ptrs = [buf.data_ptr() for buf in self._disagg_rdma_buffers] data_lens = [buf.numel() for buf in self._disagg_rdma_buffers] data_args = DataArgs( - sender_engine_rank=self._disagg_sender_rank, - receiver_engine_rank=self._disagg_receiver_rank, + sender_engine_rank=sender_rank, + receiver_engine_rank=receiver_rank, data_ptrs=data_ptrs, data_lens=data_lens, data_item_lens=data_lens, @@ -516,11 +562,11 @@ def disagg_transformer_prepare_dispatch(self, packet: Dict[str, Any]) -> None: if disagg_cfg.get("decoder_engine_rank") is None: raise RuntimeError("decentralized transformer requires decoder_engine_rank in disagg_config") - p2_transformer_rank = int(self._disagg_receiver_rank) + p2_transformer_rank = receiver_rank p2_decoder_rank = int(disagg_cfg.get("decoder_engine_rank", 2)) p2_bootstrap_addr = str(disagg_cfg.get("bootstrap_addr", "127.0.0.1")) - buffer_sizes_p2 = estimate_transformer_buffer_sizes(self.config) + buffer_sizes_p2 = estimate_transformer_buffer_sizes(config) self._disagg_alloc_p2_buffers(buffer_sizes_p2) p2_ptrs = [buf.data_ptr() for buf in self._disagg_p2_rdma_buffers] p2_lens = [buf.numel() for buf in self._disagg_p2_rdma_buffers] @@ -540,9 +586,9 @@ def disagg_transformer_prepare_dispatch(self, packet: Dict[str, Any]) -> None: if self._disagg_phase2_queue is None: raise RuntimeError("phase2 meta queue not connected; check Controller and rdma_phase2_* config") - merged_req = {**self._disagg_build_request_config_snapshot(), **req} + merged_req = self._disagg_build_request_config_snapshot(req) dc_out = {**dict(disagg_cfg)} - dc_out["sender_engine_rank"] = int(self._disagg_receiver_rank) + dc_out["sender_engine_rank"] = receiver_rank dc_out["receiver_engine_rank"] = int(disagg_cfg.get("decoder_engine_rank", 4)) dc_out["bootstrap_room"] = room merged_req["disagg_config"] = dc_out @@ -552,6 +598,7 @@ def disagg_transformer_prepare_dispatch(self, packet: Dict[str, Any]) -> None: "transformer_session_id": self._disagg_p2_data_mgr.get_session_id(), } self._disagg_phase2_queue.produce(phase2_meta) + self._disagg_request_config = merged_req self._disagg_active_transformer_room = room logger.info("[Disagg] Transformer dispatch prepared for room=%s", room) @@ -586,21 +633,18 @@ def disagg_decoder_prepare_dispatch(self, packet: Dict[str, Any]) -> None: req = dict(packet.get("request_config") or {}) trans_addr = str(packet.get("transformer_node_address", "127.0.0.1")) - - with self.config.temporarily_unlocked(): - self.config.update(req) - - disagg_cfg = self.config.get("disagg_config", {}) - room = int(self.config.get("data_bootstrap_room", disagg_cfg.get("bootstrap_room", 0))) - self.disagg_decoder_teardown_session() + self._disagg_request_config = None + config = self._disagg_effective_config(req) + disagg_cfg = config.get("disagg_config", {}) + room = int(config.get("data_bootstrap_room", disagg_cfg.get("bootstrap_room", 0))) p2_transformer_rank = int(disagg_cfg.get("sender_engine_rank", 1)) p2_decoder_rank = int(disagg_cfg.get("receiver_engine_rank", 2)) from lightx2v.disagg.utils import estimate_transformer_buffer_sizes - buffer_sizes = estimate_transformer_buffer_sizes(self.config) + buffer_sizes = estimate_transformer_buffer_sizes(config) self._disagg_alloc_p2_buffers(buffer_sizes) data_ptrs = [buf.data_ptr() for buf in self._disagg_p2_rdma_buffers] data_lens = [buf.numel() for buf in self._disagg_p2_rdma_buffers] @@ -615,6 +659,7 @@ def disagg_decoder_prepare_dispatch(self, packet: Dict[str, Any]) -> None: self._disagg_p2_data_mgr.init(data_args, room) self._disagg_p2_receiver = DataReceiver(self._disagg_p2_data_mgr, trans_addr, room) self._disagg_p2_receiver.init() + self._disagg_request_config = self._disagg_build_request_config_snapshot(req) self._disagg_active_decoder_room = room logger.info("[Disagg] Decoder dispatch prepared for room=%s", room) @@ -644,20 +689,20 @@ def _disagg_encoder_teardown_room(self, room: int) -> None: if self._disagg_active_encoder_room == room: self._disagg_active_encoder_room = None - def _disagg_encoder_setup_room(self, room: int) -> None: + def _disagg_encoder_setup_room(self, room: int, request_config: Dict[str, Any], buffer_sizes: List[int]) -> None: if self._disagg_active_encoder_room == room and self._disagg_sender is not None: return if self._disagg_active_encoder_room is not None and self._disagg_active_encoder_room != room: self._disagg_encoder_teardown_room(self._disagg_active_encoder_room) + config = self._disagg_effective_config(request_config) recv_rank = int( - self.config.get( + config.get( "disagg_phase1_receiver_engine_rank", - self.config.get("disagg_config", {}).get("receiver_engine_rank", self._disagg_receiver_rank), + config.get("disagg_config", {}).get("receiver_engine_rank", self._disagg_receiver_rank), ) ) - buffer_sizes = _estimate_encoder_buffer_sizes(self.config) self._disagg_alloc_buffers(buffer_sizes) data_ptrs = [buf.data_ptr() for buf in self._disagg_rdma_buffers] data_lens = [buf.numel() for buf in self._disagg_rdma_buffers] @@ -673,13 +718,14 @@ def _disagg_encoder_setup_room(self, room: int) -> None: self._disagg_sender = DataSender(self._disagg_data_mgr, self._disagg_bootstrap_addr, room) self._disagg_active_encoder_room = room - def _disagg_produce_phase1_for_encoder(self) -> None: + def _disagg_produce_phase1_for_encoder(self, request_config: Dict[str, Any]) -> None: if self._disagg_phase1_queue is None: raise RuntimeError("phase1 meta queue not connected") - req = self._disagg_build_request_config_snapshot() + req = self._disagg_build_request_config_snapshot(request_config) room = int(req.get("data_bootstrap_room", self._disagg_bootstrap_room)) phase1_meta = { "request_config": req, + "buffer_sizes": [buf.numel() for buf in self._disagg_rdma_buffers], "encoder_node_address": self._disagg_data_mgr.get_localhost(), "encoder_session_id": self._disagg_data_mgr.get_session_id(), } @@ -690,21 +736,12 @@ def _disagg_produce_phase1_for_encoder(self) -> None: # Encoder role: serialize and send # ------------------------------------------------------------------ # - def send_encoder_outputs(self, inputs: dict, latent_shape: list): + def send_encoder_outputs(self, inputs: dict, latent_shape: list, request_config: Optional[Dict[str, Any]] = None): """Serialize encoder outputs into RDMA buffers and send via Mooncake.""" _t_send_start = time.perf_counter() - config = self.config + config = self._disagg_effective_config(request_config or {}) disagg_cfg = config.get("disagg_config", {}) - if getattr(self, "_disagg_decentralized", False): - room = int(config.get("data_bootstrap_room", disagg_cfg.get("bootstrap_room", 0))) - self._ensure_disagg_phase1_queue_producer(disagg_cfg) - if self._disagg_phase1_queue is None: - raise RuntimeError("[Disagg] decentralized encoder could not connect phase1 queue") - _t0 = time.perf_counter() - self._disagg_encoder_setup_room(room) - self._disagg_produce_phase1_for_encoder() - _prof_log("send_enc/phase1_ring_produce", time.perf_counter() - _t0) text_encoder_output = inputs["text_encoder_output"] image_encoder_output = inputs.get("image_encoder_output") @@ -738,8 +775,6 @@ def send_encoder_outputs(self, inputs: dict, latent_shape: list): item = image_encoder_output[0] vae_encoder_out = item.get("image_latents", item) if isinstance(item, dict) else item - text_len = int(config.get("text_len", 512)) - text_dim = int(config.get("text_encoder_dim", 4096)) clip_dim = int(config.get("clip_embed_dim", 1024)) z_dim = int(config.get("vae_z_dim", 16)) @@ -757,6 +792,25 @@ def send_encoder_outputs(self, inputs: dict, latent_shape: list): enable_cfg = bool(config.get("enable_cfg", False)) use_image_encoder = bool(config.get("use_image_encoder", True)) + request_sizes = _estimate_encoder_buffer_sizes(config) + if vae_encoder_out is not None: + # Reference-image latents can be larger than the requested output image. + vae_buffer_index = 1 + int(enable_cfg) + int(use_image_encoder) + request_sizes[vae_buffer_index] = vae_encoder_out.numel() * torch.tensor([], dtype=GET_DTYPE()).element_size() + + if self._disagg_decentralized: + room = int(config.get("data_bootstrap_room", disagg_cfg.get("bootstrap_room", 0))) + self._ensure_disagg_phase1_queue_producer(disagg_cfg) + if self._disagg_phase1_queue is None: + raise RuntimeError("[Disagg] decentralized encoder could not connect phase1 queue") + _t0 = time.perf_counter() + self._disagg_encoder_setup_room(room, request_config or {}, request_sizes) + self._disagg_produce_phase1_for_encoder(request_config or {}) + _prof_log("send_enc/phase1_ring_produce", time.perf_counter() - _t0) + else: + validate_disagg_buffer_capacity(self._disagg_rdma_buffers, request_sizes, "Phase 1") + self._disagg_data_mgr.data_args[self._disagg_bootstrap_room].data_item_lens = request_sizes + buffer_index = 0 # context @@ -790,7 +844,7 @@ def send_encoder_outputs(self, inputs: dict, latent_shape: list): vae_buf = _buffer_view( self._disagg_rdma_buffers[buffer_index], GET_DTYPE(), - (z_dim + 4, t_prime, h_prime, w_prime), + tuple(vae_encoder_out.shape) if vae_encoder_out is not None else (z_dim + 4, t_prime, h_prime, w_prime), ) vae_buf.zero_() if vae_encoder_out is not None: @@ -809,6 +863,7 @@ def send_encoder_outputs(self, inputs: dict, latent_shape: list): # meta includes shapes, hashes, and image_info (for QwenImage) meta = { "version": 1, + "seed": self.input_info.seed, "task": task, "context_shape": list(context.shape), "context_hash": _sha256_tensor(context), @@ -847,15 +902,9 @@ def default(self, obj): buffer_ptrs = [buf.data_ptr() for buf in self._disagg_rdma_buffers] _t_mooncake_start = time.perf_counter() self._disagg_sender.send(buffer_ptrs) - - # Wait for transfer completion - while True: - status = self._disagg_sender.poll() - if status == DataPoll.Success: - _prof_log("send_enc/mooncake_transfer", time.perf_counter() - _t_mooncake_start) - logger.info("Disagg: encoder outputs sent successfully.") - break - time.sleep(0.01) + wait_for_disagg_transfer(self._disagg_sender, "Encoder to Transformer transfer") + _prof_log("send_enc/mooncake_transfer", time.perf_counter() - _t_mooncake_start) + logger.info("Disagg: encoder outputs sent successfully.") if getattr(self, "_disagg_decentralized", False): self._disagg_encoder_teardown_room(int(config.get("data_bootstrap_room", disagg_cfg.get("bootstrap_room", 0)))) @@ -865,19 +914,17 @@ def default(self, obj): # Transformer role: receive and deserialize # ------------------------------------------------------------------ # - def receive_encoder_outputs(self) -> dict: + def receive_encoder_outputs(self, request_config: Optional[Dict[str, Any]] = None) -> dict: """Poll for data from Encoder and reconstruct standard inputs dict.""" _t_recv_start = time.perf_counter() - config = self.config + config = self._disagg_effective_config(request_config or {}) - # Wait for data - while True: - status = self._disagg_receiver.poll() - if status == DataPoll.Success: - _prof_log("recv_enc/mooncake_poll_wait", time.perf_counter() - _t_recv_start) - logger.info("Disagg: encoder outputs received successfully.") - break - time.sleep(0.01) + if not getattr(self, "_disagg_decentralized", False): + validate_disagg_buffer_capacity(self._disagg_rdma_buffers, _estimate_encoder_buffer_sizes(config), "Phase 1") + + wait_for_disagg_transfer(self._disagg_receiver, "Encoder to Transformer transfer") + _prof_log("recv_enc/mooncake_poll_wait", time.perf_counter() - _t_recv_start) + logger.info("Disagg: encoder outputs received successfully.") # Immediately snapshot all RDMA destination buffers after poll() returns. # Without this, a concurrent Encoder send for the next request can overwrite @@ -1003,6 +1050,10 @@ def receive_encoder_outputs(self) -> dict: if meta: self._disagg_verify_integrity(meta, context, context_null, clip_encoder_out, vae_encoder_out, latent_shape, enable_cfg, task) + if self.input_info.seed is None: + self.input_info.seed = meta["seed"] + seed_all(self.input_info.seed) + _prof_log("recv_enc/deserialize_total", time.perf_counter() - _t_recv_start) return { @@ -1067,16 +1118,16 @@ def send_transformer_outputs(self, latents: torch.Tensor): import numpy as _np - # Include pixel-space dimensions so the Decoder can reconstruct auto_height/width - # correctly even when latents are in packed (sequence) format (e.g. QwenImage). _input_info = getattr(self, "input_info", None) + target_shape = getattr(_input_info, "target_shape", None) latents_meta = { "version": 1, + "seed": self.input_info.seed, "latents_shape": list(latents_to_send.shape), "latents_dtype": str(latents_to_send.dtype), "latents_hash": _sha256_tensor(latents_to_send), - "auto_height": getattr(_input_info, "auto_height", None), - "auto_width": getattr(_input_info, "auto_width", None), + "auto_height": target_shape[0] if target_shape else None, + "auto_width": target_shape[1] if target_shape else None, } meta_bytes = json.dumps(latents_meta, ensure_ascii=True).encode("utf-8") meta_buf = self._disagg_p2_rdma_buffers[1] @@ -1086,25 +1137,24 @@ def send_transformer_outputs(self, latents: torch.Tensor): meta_view.zero_() meta_view[: len(meta_bytes)].copy_(torch.from_numpy(_np.frombuffer(meta_bytes, dtype=_np.uint8).copy())) + room = self._disagg_p2_sender.bootstrap_room + self._disagg_p2_data_mgr.data_args[room].data_item_lens = [latents_nbytes, meta_buf.numel()] + torch.cuda.synchronize() _prof_log("send_trans/serialize_buffers", time.perf_counter() - _t_p2_send_start) buffer_ptrs = [buf.data_ptr() for buf in self._disagg_p2_rdma_buffers] _t_p2_mooncake = time.perf_counter() self._disagg_p2_sender.send(buffer_ptrs) - while True: - status = self._disagg_p2_sender.poll() - if status == DataPoll.Success: - _prof_log("send_trans/mooncake_transfer", time.perf_counter() - _t_p2_mooncake) - _prof_log("send_trans/total", time.perf_counter() - _t_p2_send_start) - logger.info("[Disagg] Transformer latents sent to Decoder successfully.") - break - time.sleep(0.01) + wait_for_disagg_transfer(self._disagg_p2_sender, "Transformer to Decoder transfer") + _prof_log("send_trans/mooncake_transfer", time.perf_counter() - _t_p2_mooncake) + _prof_log("send_trans/total", time.perf_counter() - _t_p2_send_start) + logger.info("[Disagg] Transformer latents sent to Decoder successfully.") # ------------------------------------------------------------------ # # Decoder role: receive latents from Transformer (Phase 2) # ------------------------------------------------------------------ # - def receive_transformer_outputs(self) -> torch.Tensor: + def receive_transformer_outputs(self, request_config: Optional[Dict[str, Any]] = None) -> torch.Tensor: """Poll Phase 2 and reconstruct latents tensor from RDMA buffer.""" _t_p2_recv_start = time.perf_counter() if self._disagg_p2_receiver is None: @@ -1112,13 +1162,15 @@ def receive_transformer_outputs(self) -> torch.Tensor: if len(self._disagg_p2_rdma_buffers) < 2: raise RuntimeError("[Disagg] Phase2 RDMA buffers require [latents, meta] entries.") - while True: - status = self._disagg_p2_receiver.poll() - if status == DataPoll.Success: - _prof_log("recv_trans/mooncake_poll_wait", time.perf_counter() - _t_p2_recv_start) - logger.info("[Disagg] Decoder received latents from Transformer successfully.") - break - time.sleep(0.01) + if not getattr(self, "_disagg_decentralized", False): + from lightx2v.disagg.utils import estimate_transformer_buffer_sizes + + config = self._disagg_effective_config(request_config or {}) + validate_disagg_buffer_capacity(self._disagg_p2_rdma_buffers, estimate_transformer_buffer_sizes(config), "Phase 2") + + wait_for_disagg_transfer(self._disagg_p2_receiver, "Transformer to Decoder transfer") + _prof_log("recv_trans/mooncake_poll_wait", time.perf_counter() - _t_p2_recv_start) + logger.info("[Disagg] Decoder received latents from Transformer successfully.") # Immediately snapshot all Phase2 RDMA destination buffers after poll() returns. # Without this, a concurrent Transformer send for the next request can overwrite @@ -1154,11 +1206,12 @@ def receive_transformer_outputs(self) -> torch.Tensor: latents = _buffer_view(received_p2_buffers[0], latents_dtype, latent_shape) if meta.get("latents_hash") is not None and _sha256_tensor(latents) != meta.get("latents_hash"): raise ValueError("[Disagg] Latents hash mismatch between transformer and decoder") + if self.input_info.seed is None: + self.input_info.seed = meta["seed"] + seed_all(self.input_info.seed) latents = latents.to(AI_DEVICE).contiguous() logger.info(f"[Disagg] Phase2 latents restored: shape={latent_shape}, dtype={latents_dtype}") - # Store the Phase 2 metadata so the caller (e.g. QwenImageRunner decode mode) can - # access pixel-space dimensions (auto_height/auto_width) that are not recoverable - # from the packed latent tensor shape alone. + # Packed latents do not retain their pixel-space dimensions. self._p2_receive_meta = meta _prof_log("recv_trans/deserialize_total", time.perf_counter() - _t_p2_recv_start) return latents diff --git a/lightx2v/disagg/examples/infer.py b/lightx2v/disagg/examples/infer.py index 183025778..1b1e88ea6 100644 --- a/lightx2v/disagg/examples/infer.py +++ b/lightx2v/disagg/examples/infer.py @@ -1,54 +1,20 @@ import argparse import os -import torch import torch.distributed as dist from loguru import logger -from lightx2v.common.ops import * -from lightx2v.models.runners.bagel.bagel_runner import BagelRunner # noqa: F401 -from lightx2v.models.runners.flux2.flux2_runner import Flux2Runner # noqa: F401 -from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner import HunyuanVideo15Runner # noqa: F401 -from lightx2v.models.runners.longcat_image.longcat_image_runner import LongCatImageRunner # noqa: F401 -from lightx2v.models.runners.ltx2.ltx2_runner import LTX2Runner # noqa: F401 -from lightx2v.models.runners.neopp.neopp_runner import NeoppRunner # noqa: F401 -from lightx2v.models.runners.qwen_image.qwen_image_runner import QwenImageRunner # noqa: F401 -from lightx2v.models.runners.seedvr.seedvr_runner import SeedVRRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_animate_runner import WanAnimateRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_audio_runner import Wan22AudioRunner, WanAudioRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_matrix_game2_runner import WanSFMtxg2Runner # noqa: F401 -from lightx2v.models.runners.wan.wan_matrix_game3_runner import WanMatrixGame3Runner # noqa: F401 -from lightx2v.models.runners.wan.wan_runner import Wan22MoeRunner, WanRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_sf_runner import WanSFRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_vace_runner import Wan22MoeVaceRunner, WanVaceRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_ar_runner import WorldPlayARRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_bi_runner import WorldPlayBIRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_distill_runner import WorldPlayDistillRunner # noqa: F401 -from lightx2v.models.runners.z_image.z_image_runner import ZImageRunner # noqa: F401 +from lightx2v.models.runners.runner_factory import build_runner from lightx2v.utils.envs import * -from lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict from lightx2v.utils.profiler import * -from lightx2v.utils.registry_factory import RUNNER_REGISTER -from lightx2v.utils.set_config import print_config, set_config, set_parallel_config -from lightx2v.utils.utils import seed_all, validate_config_paths +from lightx2v.utils.set_config import build_cli_inputs, init_parallel, print_config, print_request +from lightx2v.utils.utils import validate_config_paths from lightx2v_platform.registry_factory import PLATFORM_DEVICE_REGISTER -try: - from lightx2v.models.runners.worldmirror.worldmirror_runner import WorldMirrorRunner # noqa: F401 -except Exception as exc: # pragma: no cover - optional dependency guard - logger.warning("WorldMirrorRunner import skipped: {}", exc) - - -def init_runner(config): - torch.set_grad_enabled(False) - runner = RUNNER_REGISTER[config["model_cls"]](config) - runner.init_modules() - return runner - def main(): parser = argparse.ArgumentParser() - parser.add_argument("--seed", type=int, default=42, help="The seed for random generator") + parser.add_argument("--seed", type=int, default=None, help="The seed for random generator") parser.add_argument( "--model_cls", type=str, @@ -81,53 +47,52 @@ def main(): "lingbot_world_fast", "worldmirror", ], - default="wan2.1", ) - - parser.add_argument("--task", type=str, choices=["t2v", "i2v", "t2i", "i2i", "flf2v", "vace", "animate", "s2v", "rs2v", "t2av", "i2av", "ltx2_s2v", "sr", "recon"], default="t2v") - parser.add_argument("--support_tasks", type=str, nargs="+", default=[], help="Set supported tasks for the model") + parser.add_argument( + "--task", + type=str, + choices=["t2v", "i2v", "t2i", "i2i", "flf2v", "vace", "animate", "s2v", "rs2v", "t2av", "i2av", "ltx2_s2v", "sr", "recon"], + required=True, + ) parser.add_argument("--model_path", type=str, required=True) parser.add_argument("--sf_model_path", type=str, required=False) parser.add_argument("--config_json", type=str, required=True) - parser.add_argument("--prompt", type=str, default="", help="The input prompt for text-to-video generation") - parser.add_argument("--negative_prompt", type=str, default="") - + parser.add_argument("--prompt", type=str, default=None, help="The input prompt for text-to-video generation") + parser.add_argument("--negative_prompt", type=str, default=None) parser.add_argument( "--image_path", type=str, - default="", + default=None, help="The path to input image file(s) for image-to-video (i2v) or image-to-audio-video (i2av) task. Multiple paths should be comma-separated. Example: 'path1.jpg,path2.jpg'", ) - parser.add_argument("--last_frame_path", type=str, default="", help="The path to last frame file for first-last-frame-to-video (flf2v) task") + parser.add_argument("--last_frame_path", type=str, default=None, help="The path to last frame file for first-last-frame-to-video (flf2v) task") parser.add_argument( "--audio_path", type=str, - default="", + default=None, help="Input audio path: Wan s2v / rs2v, or required for LTX-2 task ltx2_s2v.", ) - parser.add_argument("--image_strength", type=str, default="1.0", help="i2av: single float, or comma-separated floats (one per image, or one value broadcast). Example: 1.0 or 1.0,0.85,0.9") - parser.add_argument( - "--image_frame_idx", type=str, default="", help="i2av: comma-separated pixel frame indices (one per image). Omit or empty to evenly space frames in [0, num_frames-1]. Example: 0,40,80" - ) - # [Warning] For vace task, need refactor. + parser.add_argument("--video_path", type=str, default=None, help="Input video path.") parser.add_argument( - "--src_ref_images", + "--image_strength", type=str, default=None, - help="The file list of the source reference images. Separated by ','. Default None.", + help="i2av: single float, or comma-separated floats (one per image, or one value broadcast). Example: 1.0 or 1.0,0.85,0.9", ) parser.add_argument( - "--src_video", + "--image_frame_idx", type=str, default=None, - help="The file of the source video. Default None.", + help="i2av: comma-separated pixel frame indices (one per image). Omit or empty to evenly space frames in [0, num_frames-1]. Example: 0,40,80", ) + # [Warning] For vace task, need refactor. parser.add_argument( - "--src_mask", + "--src_ref_images", type=str, default=None, - help="The file of the source mask. Default None.", + help="The file list of the source reference images. Separated by ','. Default None.", ) + parser.add_argument("--mask_path", type=str, default=None, help="Input mask path.") parser.add_argument( "--src_pose_path", type=str, @@ -146,12 +111,6 @@ def main(): default=None, help="The file of the source background. Default None.", ) - parser.add_argument( - "--src_mask_path", - type=str, - default=None, - help="The file of the source mask. Default None.", - ) parser.add_argument( "--pose", type=str, @@ -164,65 +123,37 @@ def main(): default=None, help="Directory path for lingbot camera/action control files (poses.npy, intrinsics.npy, optional action.npy).", ) - parser.add_argument( - "--action_ckpt", - type=str, - default=None, - help="Path to action model checkpoint for WorldPlay models.", - ) # WorldMirror (3D reconstruction) specific parser.add_argument("--input_path", type=str, default=None, help="(worldmirror/recon) Path to a directory of images, a video file, or a single image.") parser.add_argument("--strict_output_path", type=str, default=None, help="(worldmirror/recon) If set, write outputs directly here instead of under save_result_path///.") parser.add_argument("--prior_cam_path", type=str, default=None, help="(worldmirror/recon) Optional camera prior JSON (extrinsics + intrinsics).") parser.add_argument("--prior_depth_path", type=str, default=None, help="(worldmirror/recon) Optional depth prior directory (one .npy/.png per image).") - parser.add_argument("--subfolder", type=str, default=None, help="(worldmirror/recon) Subfolder inside model_path containing weights. Overrides config.") - parser.add_argument("--disable_heads", type=str, nargs="*", default=None, help="(worldmirror/recon) Heads to disable: any of camera depth normal points gs.") - parser.add_argument("--enable_bf16", action="store_true", default=False, help="(worldmirror/recon) Run the WorldMirror model in bf16.") - parser.add_argument("--save_rendered", action="store_true", default=False, help="(worldmirror/recon) Render an interpolated fly-through video from Gaussian splats.") + parser.add_argument("--save_rendered", action="store_true", default=None, help="(worldmirror/recon) Render an interpolated fly-through video from Gaussian splats.") parser.add_argument("--render_interp_per_pair", type=int, default=None, help="(worldmirror/recon) Interpolated frames per camera pair for --save_rendered.") - parser.add_argument("--render_depth", action="store_true", default=False, help="(worldmirror/recon) Also render a depth video with --save_rendered.") - parser.add_argument("--wm_config_path", type=str, default=None, help="(worldmirror/recon) Optional training YAML (pair with --wm_ckpt_path).") - parser.add_argument("--wm_ckpt_path", type=str, default=None, help="(worldmirror/recon) Optional .ckpt/.safetensors (pair with --wm_config_path).") + parser.add_argument("--render_depth", action="store_true", default=None, help="(worldmirror/recon) Also render a depth video with --save_rendered.") parser.add_argument("--save_result_path", type=str, default=None, help="The path to save video path/file") - parser.add_argument("--return_result_tensor", action="store_true", help="Whether to return result tensor. (Useful for comfyui)") - parser.add_argument("--target_shape", type=int, nargs="+", default=[], help="Set return video or image shape") - parser.add_argument("--aspect_ratio", type=str, default="") - parser.add_argument("--video_path", type=str, default=None, help="input video path(for sr/v2v task)") - parser.add_argument("--sr_ratio", type=float, default=2.0, help="super resolution ratio for sr task") - parser.add_argument( - "--num_iterations", - type=int, - default=None, - help="Override the number of Matrix-Game-3 generation segments. Final video length follows 57 + 40 * (num_iterations - 1).", - ) + parser.add_argument("--return_result_tensor", action="store_true", default=None, help="Whether to return result tensor. (Useful for comfyui)") + parser.add_argument("--target_shape", type=int, nargs="+", default=None, help="Set return video or image shape") + parser.add_argument("--aspect_ratio", type=str, default=None) + parser.add_argument("--sr_ratio", type=float, default=None, help="super resolution ratio for sr task") args = parser.parse_args() - # validate_task_arguments(args) - - seed_all(args.seed) - - # set config - config = set_config(args) - # init input_info - input_info = init_empty_input_info(args.task, args.support_tasks) - - if config["parallel"]: + startup_config, request_data = build_cli_inputs(args) + if startup_config["parallel"]: platform_device = PLATFORM_DEVICE_REGISTER.get(os.getenv("PLATFORM", "cuda"), None) platform_device.init_parallel_env() - set_parallel_config(config) + init_parallel(startup_config) - print_config(config) + print_config(startup_config, title="Startup config") - validate_config_paths(config) + validate_config_paths(startup_config) with ProfilingContext4DebugL1("Total Cost"): - # init runner - runner = init_runner(config) - # start to infer - data = args.__dict__ - update_input_info_from_dict(input_info, data) - runner.run_pipeline(input_info) + runner = build_runner(startup_config) + input_info = runner.prepare_request(request_data) + print_request(input_info, runner.get_supported_request_fields(input_info.task)) + runner.run_request(input_info) # Clean up distributed process group if dist.is_initialized(): diff --git a/lightx2v/disagg/examples/run_service.py b/lightx2v/disagg/examples/run_service.py index a8dd95892..df6b4bdbb 100644 --- a/lightx2v/disagg/examples/run_service.py +++ b/lightx2v/disagg/examples/run_service.py @@ -4,7 +4,7 @@ from loguru import logger -from lightx2v.disagg.utils import set_config +from lightx2v.utils.set_config import build_startup_config from lightx2v.utils.utils import seed_all logging.basicConfig(level=logging.INFO) @@ -17,7 +17,8 @@ def _build_parser() -> argparse.ArgumentParser: parser.add_argument("--model_path", type=str, required=True) parser.add_argument("--config_json", type=str, required=True) - parser.add_argument("--seed", type=int, default=None) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--image_path", type=str, default=None) parser.add_argument( "--prompt", type=str, @@ -32,11 +33,7 @@ def _build_parser() -> argparse.ArgumentParser: "畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" ), ) - parser.add_argument( - "--save_result_path", - type=str, - default="/root/zht/LightX2V/save_results/test_disagg.mp4", - ) + parser.add_argument("--save_result_path", type=str, default=None) parser.add_argument( "--service", @@ -90,23 +87,11 @@ def _resolve_service_mode(args: argparse.Namespace, raw_cfg: dict) -> str: def _build_runtime_config(args: argparse.Namespace) -> tuple[dict, dict]: raw_cfg = _load_raw_json(args.config_json) - config = set_config( - model_path=args.model_path, - task=args.task, - model_cls=args.model_cls, - config_path=args.config_json, - ) + config = build_startup_config({"model_path": args.model_path, "task": args.task, "model_cls": args.model_cls, "config_json": args.config_json}) config = _normalize_disagg_config(config) raw_cfg = _normalize_disagg_config(raw_cfg) - if args.seed is not None: - config["seed"] = args.seed - elif config.get("seed") is None: - config["seed"] = 42 - config["prompt"] = args.prompt - config["negative_prompt"] = args.negative_prompt - config["save_path"] = args.save_result_path return config, raw_cfg @@ -119,7 +104,7 @@ def main(): rank_key = f"{service_mode}_engine_rank" config[rank_key] = int(args.engine_rank) - seed_all(config["seed"]) + seed_all(args.seed) logger.info("Starting disagg service mode={}", service_mode) if service_mode == "encoder": @@ -137,7 +122,14 @@ def main(): elif service_mode == "controller": from lightx2v.disagg.services.controller import ControllerService - ControllerService().run(config) + request_data = { + "prompt": args.prompt, + "negative_prompt": args.negative_prompt, + "image_path": args.image_path, + "seed": args.seed, + "save_path": args.save_result_path, + } + ControllerService().run(config, request_data) else: raise ValueError(f"Unsupported service mode: {service_mode}") diff --git a/lightx2v/disagg/examples/run_user.py b/lightx2v/disagg/examples/run_user.py index f360bcd44..680867eea 100644 --- a/lightx2v/disagg/examples/run_user.py +++ b/lightx2v/disagg/examples/run_user.py @@ -19,6 +19,10 @@ def _build_parser() -> argparse.ArgumentParser: parser.add_argument("--controller_request_port", type=int, default=REQUEST_POLLING_PORT - 2) parser.add_argument("--max_requests", type=int, default=0, help="0 means no hard cap") parser.add_argument("--sleep_min_ms", type=float, default=5.0, help="minimum loop sleep in ms") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--prompt", type=str) + parser.add_argument("--negative_prompt", type=str) + parser.add_argument("--image_path", type=str) return parser @@ -28,6 +32,8 @@ def main(): req_mgr = ReqManager() stages = load_stage_specs() base_config = load_base_config() + request_data = {"seed": args.seed, "prompt": args.prompt, "negative_prompt": args.negative_prompt, "image_path": args.image_path} + base_config.update({key: value for key, value in request_data.items() if value is not None}) shape = DisaggLoadShape() start_workload_clock() diff --git a/lightx2v/disagg/examples/wan_i2v_service.py b/lightx2v/disagg/examples/wan_i2v_service.py index c5030d669..b0ba12a77 100644 --- a/lightx2v/disagg/examples/wan_i2v_service.py +++ b/lightx2v/disagg/examples/wan_i2v_service.py @@ -70,18 +70,12 @@ def main(): decoder_engine_rank=2, ) - config["image_path"] = image_path - config["prompt"] = prompt - config["negative_prompt"] = negative_prompt - config["save_path"] = save_result_path + request_data = {"seed": seed, "prompt": prompt, "negative_prompt": negative_prompt, "image_path": image_path, "save_path": save_result_path} logger.info(f"Config initialized for task: {task}") seed_all(seed) - # 2. Add seed to config so services use it - config["seed"] = seed - - # 3. Define service threads + # 2. Define service threads def run_encoder(): logger.info("Initializing Encoder Service...") encoder_service = EncoderService(config) @@ -107,13 +101,13 @@ def run_controller(): logger.info("Initializing Controller Service...") controller_service = ControllerService() logger.info("Dispatching request to services...") - controller_service.run(config) + controller_service.run(config, request_data) encoder_stop_event.set() transformer_stop_event.set() decoder_stop_event.set() logger.info("Controller Service completed.") - # 4. Start threads + # 3. Start threads encoder_thread = threading.Thread(target=run_encoder) transformer_thread = threading.Thread(target=run_transformer) decoder_thread = threading.Thread(target=run_decoder) @@ -125,7 +119,7 @@ def run_controller(): decoder_thread.start() controller_thread.start() - # 5. Wait for completion + # 4. Wait for completion encoder_thread.join() transformer_thread.join() decoder_thread.join() diff --git a/lightx2v/disagg/examples/wan_t2v_service.py b/lightx2v/disagg/examples/wan_t2v_service.py index 80833dd6e..06698e4ac 100644 --- a/lightx2v/disagg/examples/wan_t2v_service.py +++ b/lightx2v/disagg/examples/wan_t2v_service.py @@ -54,11 +54,7 @@ def main(): logger.info(f"Config initialized for task: {task}") seed_all(seed) - # Add seed into config so services can use it if needed - config["seed"] = seed - config["prompt"] = prompt - config["negative_prompt"] = negative_prompt - config["save_path"] = save_result_path + request_data = {"seed": seed, "prompt": prompt, "negative_prompt": negative_prompt, "save_path": save_result_path} encoder_stop_event = threading.Event() transformer_stop_event = threading.Event() @@ -90,7 +86,7 @@ def run_controller(): logger.info("Initializing Controller Service...") controller_service = ControllerService() logger.info("Dispatching request to services...") - controller_service.run(config) + controller_service.run(config, request_data) encoder_stop_event.set() transformer_stop_event.set() decoder_stop_event.set() diff --git a/lightx2v/disagg/services/controller.py b/lightx2v/disagg/services/controller.py index 6b52bcfa8..c6d65c5ba 100644 --- a/lightx2v/disagg/services/controller.py +++ b/lightx2v/disagg/services/controller.py @@ -164,12 +164,6 @@ def _build_service_command(self, instance_type: str, engine_rank: int, instance_ str(instance_cfg.get("model_path")), "--config_json", service_config_json, - "--prompt", - str(instance_cfg.get("prompt", "")), - "--negative_prompt", - str(instance_cfg.get("negative_prompt", "")), - "--save_result_path", - str(instance_cfg.get("save_path", "")), ] def _maybe_wrap_service_command_with_nsys( @@ -1921,8 +1915,8 @@ def _address_to_rank(instance_address: str) -> int: self.rdma_buffer_request.produce(config) self.logger.info("Request enqueued to encoder request RDMA buffer") - def run(self, config): - """Initialize controller buffers, stream request configs from workload, then wait for all callbacks.""" + def run(self, config, request_data=None): + """Dispatch incoming workload or repeat request_data for automatic requests.""" if config is None: raise ValueError("config cannot be None") @@ -2033,7 +2027,6 @@ def run(self, config): time.sleep(5.0) - base_save_path = config.get("save_path") expected_rooms: set[int] = set() received_rooms: set[int] = set() received_results: list[dict] = [] @@ -2059,7 +2052,7 @@ def run(self, config): self.logger.info("LOAD_FROM_USER enabled, waiting workload configs on port=%s", request_ingress_port) else: self.logger.info( - "LOAD_FROM_USER disabled, generating requests from config: count=%s", + "LOAD_FROM_USER disabled, repeating automatic request: count=%s", auto_request_count, ) @@ -2076,14 +2069,13 @@ def run(self, config): else: if generated_request_count >= auto_request_count: break - workload_config = {} + workload_config = dict(request_data or {}) generated_request_count += 1 request_config = dict(config) request_config.update(self._to_plain(workload_config)) - if request_config.get("seed") is None: - seed = config.get("seed") - request_config["seed"] = 42 if seed is None else seed + seed = workload_config.get("seed") + request_config["seed"] = 42 if seed is None else seed room = request_config.get("data_bootstrap_room", next_room) try: @@ -2099,7 +2091,6 @@ def run(self, config): request_config["data_bootstrap_room"] = room request_config["controller_result_host"] = bootstrap_addr request_config["controller_result_port"] = result_port - metrics = request_config.get("request_metrics") if not isinstance(metrics, dict): metrics = {} @@ -2109,9 +2100,22 @@ def run(self, config): metrics["stages"] = {} request_config["request_metrics"] = metrics - if base_save_path and not request_config.get("save_path"): - save_path = Path(base_save_path) - request_config["save_path"] = str(save_path.with_name(f"{save_path.stem}{room}{save_path.suffix}")) + save_path = workload_config.get("save_path") + if not save_path: + error = "save_path is required for disaggregated generation requests" + if not load_from_user: + raise ValueError(error) + expected_rooms.add(room) + self._handle_decoder_result( + {"ok": False, "data_bootstrap_room": room, "save_path": None, "error": error, "request_metrics": metrics}, + expected_rooms=expected_rooms, + received_rooms=received_rooms, + received_results=received_results, + ) + continue + if not load_from_user: + output_path = Path(save_path) + request_config["save_path"] = str(output_path.with_name(f"{output_path.stem}{room}{output_path.suffix}")) with self._lock: current_request = request_config diff --git a/lightx2v/disagg/services/decoder.py b/lightx2v/disagg/services/decoder.py index db473574a..a0e587944 100644 --- a/lightx2v/disagg/services/decoder.py +++ b/lightx2v/disagg/services/decoder.py @@ -189,6 +189,10 @@ def alloc_memory(self, request: AllocationRequest) -> MemoryHandle: return MemoryHandle(buffers=buffers) def process(self, config): + save_path = config.get("save_path") + if not save_path: + raise ValueError("save_path is required for disaggregated generation requests") + seed_all(config["seed"]) self.logger.info("Starting processing in DecoderService...") room = config.get("data_bootstrap_room", 0) @@ -301,10 +305,6 @@ def _infer_latents_shape_from_config() -> tuple[int, int, int, int]: gen_video_final = wan_vae_to_comfy(gen_video) decoder_metrics["compute_end_ts"] = time.time() - save_path = config.get("save_path") - if save_path is None: - raise ValueError("save_path is required in config.") - self.logger.info(f"Saving video to {save_path}...") save_to_video(gen_video_final, save_path, fps=config.get("fps", 16), method="ffmpeg") decoder_metrics["output_enqueued_ts"] = time.time() diff --git a/lightx2v/disagg/utils.py b/lightx2v/disagg/utils.py index d596fa5db..a5f538303 100644 --- a/lightx2v/disagg/utils.py +++ b/lightx2v/disagg/utils.py @@ -1,4 +1,3 @@ -import json import logging import math import os @@ -8,22 +7,14 @@ import torchvision.transforms.functional as TF from PIL import Image -from lightx2v.models.networks.lora_adapter import LoraAdapter from lightx2v.utils.envs import GET_DTYPE -from lightx2v.utils.set_config import set_config as set_config_base +from lightx2v.utils.set_config import build_startup_config from lightx2v.utils.utils import find_torch_model_path from lightx2v_platform.base.global_var import AI_DEVICE logger = logging.getLogger(__name__) -class ConfigObj: - """Helper class to convert dictionary to object with attributes""" - - def __init__(self, **kwargs): - self.__dict__.update(kwargs) - - def read_image_input(image_path): img_ori = Image.open(image_path).convert("RGB") img = TF.to_tensor(img_ori).sub_(0.5).div_(0.5).unsqueeze(0).to(AI_DEVICE) @@ -78,16 +69,7 @@ def set_config( "vae_cpu_offload": vae_offload, # Map to internal keys } - # Simulate logic from LightX2VPipeline.create_generator - # which calls set_infer_config / set_infer_config_json - # Here we directly populate args_dict with the required inference config - - if config_path is not None: - with open(config_path, "r") as f: - config_json_content = json.load(f) - args_dict.update(config_json_content) - else: - # Replicating set_infer_config logic + if config_path is None: if model_cls == "ltx2": args_dict["distilled_sigma_values"] = distilled_sigma_values args_dict["infer_steps"] = len(distilled_sigma_values) - 1 if distilled_sigma_values is not None else infer_steps @@ -125,42 +107,7 @@ def set_config( args_dict["norm_modulate_backend"] = norm_modulate_backend args_dict.update(kwargs) - - # Convert to object for set_config compatibility - args = ConfigObj(**args_dict) - - # Use existing set_config from utils - config = set_config_base(args) - - return config - - -def build_wan_model_with_lora(wan_module, config, model_kwargs, lora_configs, model_type="high_noise_model"): - lora_dynamic_apply = config.get("lora_dynamic_apply", False) - - if lora_dynamic_apply: - if model_type in ["high_noise_model", "low_noise_model"]: - # For wan2.2 - lora_name_to_info = {item["name"]: item for item in lora_configs} - lora_path = lora_name_to_info[model_type]["path"] - lora_strength = lora_name_to_info[model_type]["strength"] - else: - # For wan2.1 - lora_path = lora_configs[0]["path"] - lora_strength = lora_configs[0]["strength"] - - model_kwargs["lora_path"] = lora_path - model_kwargs["lora_strength"] = lora_strength - model = wan_module(**model_kwargs) - else: - assert not config.get("dit_quantized", False), "Online LoRA only for quantized models; merging LoRA is unsupported." - assert not config.get("lazy_load", False), "Lazy load mode does not support LoRA merging." - model = wan_module(**model_kwargs) - lora_wrapper = LoraAdapter(model) - if model_type in ["high_noise_model", "low_noise_model"]: - lora_configs = [lora_config for lora_config in lora_configs if lora_config["name"] == model_type] - lora_wrapper.apply_lora(lora_configs, model_type=model_type) - return model + return build_startup_config(args_dict) def load_wan_text_encoder(config: Dict[str, Any]): @@ -344,7 +291,7 @@ def load_wan_transformer(config: Dict[str, Any]): flush=True, ) from lightx2v.models.networks.wan.model import WanModel - from lightx2v.models.runners.wan.wan_runner import MultiModelStruct, get_wan_model_class + from lightx2v.models.runners.wan.wan_runner import MultiModelStruct, build_wan_model_with_lora, get_wan_model_class from lightx2v.models.schedulers.wan.scheduler_factory import get_wan_distill_method print( diff --git a/lightx2v/infer.py b/lightx2v/infer.py index f95818a3e..b9c137179 100755 --- a/lightx2v/infer.py +++ b/lightx2v/infer.py @@ -5,63 +5,15 @@ import torch.distributed as dist from loguru import logger -from lightx2v.common.ops import * from lightx2v.models.networks.bagel.sensenova_tasks import OMNI_VISION_SUBTASK_CHOICES -from lightx2v.models.networks.wan.animate2_identity import WAN_ANIMATE2_MODEL_ID -from lightx2v.models.runners.bagel.bagel_runner import BagelRunner # noqa: F401 -from lightx2v.models.runners.bagel.sensenova_vision_runner import SenseNovaVisionRunner # noqa: F401 -from lightx2v.models.runners.cosmos3.cosmos3_runner import Cosmos3Runner # noqa: F401 -from lightx2v.models.runners.ernie_image.ernie_image_runner import ErnieImageRunner # noqa: F401 -from lightx2v.models.runners.flux2.flux2_runner import Flux2Runner # noqa: F401 -from lightx2v.models.runners.hidream_o1_image.hidream_o1_image_runner import HidreamO1ImageRunner # noqa: F401 -from lightx2v.models.runners.hunyuan3d.hunyuan3d_shape_runner import Hunyuan3DShapeRunner # noqa: F401 -from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner # noqa: F401 -from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner import HunyuanVideo15Runner # noqa: F401 -from lightx2v.models.runners.lingbot_video.lingbot_video_runner import LingBotVideoRunner # noqa: F401 -from lightx2v.models.runners.longcat_image.longcat_image_runner import LongCatImageRunner # noqa: F401 -from lightx2v.models.runners.ltx2.ltx2_runner import LTX2ARRunner, LTX2Runner # noqa: F401 -from lightx2v.models.runners.ltx2.ltx25_runner import LTX25Runner # noqa: F401 -from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner # noqa: F401 -from lightx2v.models.runners.motus.motus_runner import MotusRunner # noqa: F401 -from lightx2v.models.runners.neopp.neopp_runner import NeoppRunner # noqa: F401 -from lightx2v.models.runners.qwen_image.qwen_image_runner import QwenImageRunner # noqa: F401 -from lightx2v.models.runners.seedvr.seedvr_runner import SeedVRRunner # noqa: F401 -from lightx2v.models.runners.swiftvr.swiftvr_runner import SwiftVRRunner # noqa: F401 -from lightx2v.models.runners.wan.fastwam_runner import FastWAMRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_animate2_runner import WanAnimate2Runner # noqa: F401 -from lightx2v.models.runners.wan.wan_animate_runner import WanAnimateRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_audio_runner import Wan22AudioRunner, WanAudioRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_dancer_runner import WanDancerRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_dreamzero_runner import WanDreamZeroRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_infinitetalk_runner import InfiniteTalkRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_lingbot_va_runner import LingbotVARunner # noqa: F401 -from lightx2v.models.runners.wan.wan_matrix_game2_runner import WanSFMtxg2Runner # noqa: F401 -from lightx2v.models.runners.wan.wan_matrix_game3_runner import WanMatrixGame3Runner # noqa: F401 -from lightx2v.models.runners.wan.wan_runner import Wan22MoeRunner, WanRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_s2v_runner import WanS2VRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_sf_runner import WanSFRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_vace_runner import Wan22MoeVaceRunner, WanVaceRunner # noqa: F401 -from lightx2v.models.runners.worldmirror.worldmirror_runner import WorldMirrorRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_ar_runner import WorldPlayARRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_bi_runner import WorldPlayBIRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_distill_runner import WorldPlayDistillRunner # noqa: F401 -from lightx2v.models.runners.z_image.z_image_runner import ZImageRunner # noqa: F401 +from lightx2v.models.runners.runner_factory import RUNNER_MODULES, build_runner from lightx2v.utils.envs import * -from lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict from lightx2v.utils.profiler import * -from lightx2v.utils.registry_factory import RUNNER_REGISTER -from lightx2v.utils.set_config import print_config, set_config, set_parallel_config -from lightx2v.utils.utils import seed_all, validate_config_paths +from lightx2v.utils.set_config import build_cli_inputs, init_parallel, print_config, print_request +from lightx2v.utils.utils import validate_config_paths from lightx2v_platform.registry_factory import PLATFORM_DEVICE_REGISTER -def init_runner(config): - torch.set_grad_enabled(False) - runner = RUNNER_REGISTER[config["model_cls"]](config) - runner.init_modules() - return runner - - def distributed_barrier(): import torch.distributed as dist @@ -84,60 +36,12 @@ def distributed_barrier(): def main(): parser = argparse.ArgumentParser() - parser.add_argument("--seed", type=int, default=42, help="The seed for random generator") + parser.add_argument("--seed", type=int, default=None, help="The seed for random generator") parser.add_argument( "--model_cls", type=str, required=True, - choices=[ - "wan2.1", - "wan_dancer", - "wan2.1_vace", - "wan2.1_sf", - "wan2.1_sf_mtxg2", - "seko_talk", - "seko_talk_ar", - "wan2.2_moe", - "lingbot_world", - "wan2.2", - "wan2.2_matrix_game3", - "wan2.2_audio", - "wan2.2_moe_vace", - "qwen_image", - "ernie_image", - "hidream_o1_image", - "longcat_image", - "cosmos3", - "wan2.2_animate", - WAN_ANIMATE2_MODEL_ID, - "wan2.2_s2v", - "hunyuan_video_1.5", - "hunyuan_image3", - "hunyuan3d", - "worldplay_distill", - "worldplay_ar", - "worldplay_bi", - "z_image", - "flux2", - "ltx2", - "ltx2_ar", - "ltx2_5", - "minimax_h3", - "bagel", - "sensenova_vision", - "seedvr2", - "swiftvr", - "neopp", - "motus", - "lingbot_world_fast", - "worldmirror", - "lingbot_va", - "dreamzero", - "infinitetalk", - "fastwam", - "lingbot_video", - ], - default="wan2.1", + choices=RUNNER_MODULES, ) parser.add_argument( @@ -169,42 +73,55 @@ def main(): "i23d", "omni_vision_task", ], - default="t2v", + required=True, ) - parser.add_argument("--support_tasks", type=str, nargs="+", default=[], help="Set supported tasks for the model") parser.add_argument( "--omni_vision_subtask", type=str, choices=OMNI_VISION_SUBTASK_CHOICES, default=None, - help="SenseNova-Vision subtask used with --task omni_vision_task.", + help="Subtask used with --task omni_vision_task.", ) parser.add_argument("--model_path", type=str, required=True) parser.add_argument("--config_json", type=str, required=True) - parser.add_argument("--prompt", type=str, default="", help="The input prompt for text-to-video generation") - parser.add_argument("--prompt_ref", type=str, default="人物动作的参考视频", help="Reference/driving-video prompt for Wan-Animate-2.") - parser.add_argument("--negative_prompt", type=str, default="") + parser.add_argument("--prompt", type=str, default=None, help="The input prompt for text-to-video generation") + parser.add_argument("--prompt_ref", type=str, default=None, help="Reference/driving-video prompt for Wan-Animate-2.") + parser.add_argument("--negative_prompt", type=str, default=None) + parser.add_argument("--bot_task", type=str, default=None, help="HunyuanImage3 text generation mode.") + parser.add_argument("--max_new_tokens", type=int, default=None, help="Maximum number of generated text tokens.") + parser.add_argument("--system_prompt", type=str, default=None, help="System prompt for text generation.") + parser.add_argument("--text_do_sample", action=argparse.BooleanOptionalAction, default=None, help="Enable sampling during text generation.") + parser.add_argument("--text_temperature", type=float, default=None, help="Text sampling temperature.") + parser.add_argument("--text_top_k", type=int, default=None, help="Top-k text sampling limit.") + parser.add_argument("--text_top_p", type=float, default=None, help="Top-p text sampling threshold.") parser.add_argument( "--image_path", type=str, - default="", + default=None, help="The path to input image file(s), including HunyuanImage3 ti2t/ti2i and MiniMax-H3 ref2av reference images. Multiple paths should be comma-separated. Example: 'path1.jpg,path2.jpg'", ) - parser.add_argument("--state_path", type=str, default="", help="The path to input robot state file for robot i2v/i2va inference.") - parser.add_argument("--last_frame_path", type=str, default="", help="The path to last frame file for first-last-frame-to-video (flf2v) task") + parser.add_argument("--state_path", type=str, default=None, help="The path to input robot state file for robot i2v/i2va inference.") + parser.add_argument("--last_frame_path", type=str, default=None, help="The path to last frame file for first-last-frame-to-video (flf2v) task") parser.add_argument( "--audio_path", type=str, - default="", + default=None, help="Input audio path: Wan s2v / rs2v, LTX-2 ltx2_s2v, or MiniMax-H3 ref2av reference audio. H3 accepts comma-separated paths.", ) - parser.add_argument("--image_strength", type=str, default="1.0", help="i2av: single float, or comma-separated floats (one per image, or one value broadcast). Example: 1.0 or 1.0,0.85,0.9") + parser.add_argument( + "--video_path", + type=str, + default=None, + help="Input source video path. Its role is determined by the selected task.", + ) + parser.add_argument("--video_duration", type=float, default=None, help="Requested output duration in seconds for audio-driven video generation.") + parser.add_argument("--image_strength", type=str, default=None, help="i2av: single float, or comma-separated floats (one per image, or one value broadcast). Example: 1.0 or 1.0,0.85,0.9") parser.add_argument( "--num_frames", dest="target_video_length", type=int, default=None, - help="LTX-2.5: explicit output frame count (must be 8k+1). Omit to use DurationHead when enabled.", + help="Requested output frame count. Model-specific length constraints apply.", ) parser.add_argument( "--i2i_denoise_strength", @@ -212,8 +129,10 @@ def main(): default=None, help="(i2i) Single-image edit denoising strength in [0.0, 1.0]. 0.0 preserves the source image most; 1.0 redraws most. Omit to keep the model's existing behavior.", ) + parser.add_argument("--inpaint_blur_sigma", type=float, default=None, help="Flux2 inpainting mask blur sigma.") + parser.add_argument("--inpaint_blur_size", type=int, default=None, help="Flux2 inpainting mask blur kernel size.") parser.add_argument( - "--image_frame_idx", type=str, default="", help="i2av: comma-separated pixel frame indices (one per image). Omit or empty to evenly space frames in [0, num_frames-1]. Example: 0,40,80" + "--image_frame_idx", type=str, default=None, help="i2av: comma-separated pixel frame indices (one per image). Omit or empty to evenly space frames in [0, num_frames-1]. Example: 0,40,80" ) # [Warning] For vace task, need refactor. parser.add_argument( @@ -222,22 +141,11 @@ def main(): default=None, help="The file list of the source reference images. Separated by ','. Default None.", ) - parser.add_argument( - "--src_video", - type=str, - default=None, - help="The file of the source video. Default None.", - ) - parser.add_argument( - "--src_mask", - type=str, - default=None, - help="The file of the source mask. Default None.", - ) + parser.add_argument("--mask_path", type=str, default=None, help="Input mask path.") parser.add_argument( "--src_pose_path", type=str, - default="", + default=None, help="Pose driving video for Wan s2v / animate (e.g. examples/pose.mp4).", ) parser.add_argument( @@ -252,12 +160,6 @@ def main(): default=None, help="The file of the source background. Default None.", ) - parser.add_argument( - "--src_mask_path", - type=str, - default=None, - help="The file of the source mask. Default None.", - ) parser.add_argument( "--pose", type=str, @@ -273,89 +175,59 @@ def main(): parser.add_argument("--action_mode", type=str, default=None, choices=["forward_dynamics", "inverse_dynamics", "policy"], help="Cosmos3 action mode.") parser.add_argument("--domain_name", type=str, default=None, help="Cosmos3 action embodiment domain name.") parser.add_argument("--view_point", type=str, default=None, help="Cosmos3 action viewpoint label.") - parser.add_argument("--action_chunk_size", type=int, default=None, help="Cosmos3 action chunk size.") - parser.add_argument("--action_chunk_index", type=int, default=None, help="Cosmos3 action chunk index when action_path contains action_chunks.") - parser.add_argument( - "--action_ckpt", - type=str, - default=None, - help="Path to action model checkpoint for WorldPlay models.", - ) # WorldMirror (3D reconstruction) specific parser.add_argument("--input_path", type=str, default=None, help="(worldmirror/recon) Path to a directory of images, a video file, or a single image.") parser.add_argument("--strict_output_path", type=str, default=None, help="(worldmirror/recon) If set, write outputs directly here instead of under save_result_path///.") parser.add_argument("--prior_cam_path", type=str, default=None, help="(worldmirror/recon) Optional camera prior JSON (extrinsics + intrinsics).") parser.add_argument("--prior_depth_path", type=str, default=None, help="(worldmirror/recon) Optional depth prior directory (one .npy/.png per image).") - parser.add_argument("--subfolder", type=str, default=None, help="(worldmirror/recon) Subfolder inside model_path containing weights. Overrides config.") - parser.add_argument("--disable_heads", type=str, nargs="*", default=None, help="(worldmirror/recon) Heads to disable: any of camera depth normal points gs.") - parser.add_argument("--enable_bf16", action="store_true", default=False, help="(worldmirror/recon) Run the WorldMirror model in bf16.") - parser.add_argument("--save_rendered", action="store_true", default=False, help="(worldmirror/recon) Render an interpolated fly-through video from Gaussian splats.") + parser.add_argument("--save_rendered", action=argparse.BooleanOptionalAction, default=None, help="(worldmirror/recon) Render an interpolated fly-through video from Gaussian splats.") parser.add_argument("--render_interp_per_pair", type=int, default=None, help="(worldmirror/recon) Interpolated frames per camera pair for --save_rendered.") - parser.add_argument("--render_depth", action="store_true", default=False, help="(worldmirror/recon) Also render a depth video with --save_rendered.") - parser.add_argument("--wm_config_path", type=str, default=None, help="(worldmirror/recon) Optional training YAML (pair with --wm_ckpt_path).") - parser.add_argument("--wm_ckpt_path", type=str, default=None, help="(worldmirror/recon) Optional .ckpt/.safetensors (pair with --wm_config_path).") + parser.add_argument("--render_depth", action=argparse.BooleanOptionalAction, default=None, help="(worldmirror/recon) Also render a depth video with --save_rendered.") parser.add_argument("--save_result_path", type=str, default=None, help="The path to save video path/file") + parser.add_argument("--return_result_tensor", action="store_true", default=None, help="Whether to return result tensor. (Useful for comfyui)") parser.add_argument("--save_action_path", type=str, default=None, help="The path to save action predictions for Motus, LingBot-VA, or DreamZero.") - parser.add_argument("--return_result_tensor", action="store_true", help="Whether to return result tensor. (Useful for comfyui)") - parser.add_argument("--target_shape", type=int, nargs="+", default=[], help="Set return video or image shape") - parser.add_argument("--aspect_ratio", type=str, default="") + parser.add_argument("--raw_output_path", type=str, default=None, help="Raw prediction output path for SenseNova-Vision.") + parser.add_argument("--glb_output_path", type=str, default=None, help="GLB scene output path for SenseNova-Vision.") + parser.add_argument("--postprocess_predictions", action=argparse.BooleanOptionalAction, default=None, help="Postprocess SenseNova-Vision predictions.") + parser.add_argument("--target_shape", type=int, nargs="+", default=None, help="Set return video or image shape") + parser.add_argument("--aspect_ratio", type=str, default=None) + parser.add_argument("--infer_align_image_size", action=argparse.BooleanOptionalAction, default=None, help="Align HunyuanImage3 reference image sizes during inference.") parser.add_argument( "--keep_original_aspect", - action="store_true", + action=argparse.BooleanOptionalAction, + default=None, help="(i2i) When exactly one reference image is provided, preserve its aspect ratio with max_size=2048.", ) parser.add_argument( "--layout_bboxes", type=str, - default="", - help="(i2i) Layout boxes as a JSON string or JSON file path for HiDream layout-conditioned editing.", - ) - parser.add_argument( - "--video_path", - type=str, - default=None, - help="Input video path for sr/v2v/v2av, or MiniMax-H3 ref2av reference video. H3 accepts comma-separated paths. For v2av this is the pre-processed control/reference video (pose / canny / depth / motion-track for motion-transfer, or the degraded source video for ICEdit).", - ) - parser.add_argument("--sr_ratio", type=float, default=2.0, help="super resolution ratio for sr task") - parser.add_argument( - "--num_iterations", - type=int, default=None, - help="Override the number of Matrix-Game-3 generation segments. Final video length follows 57 + 40 * (num_iterations - 1).", + help="(i2i) Layout boxes as a JSON string or JSON file path for HiDream layout-conditioned editing.", ) + parser.add_argument("--sr_ratio", type=float, default=None, help="super resolution ratio for sr task") parser.add_argument( - "--reference_video_strength", type=float, default=1.0, help="(v2av) IC-LoRA reference-video conditioning strength in [0.0, 1.0]. 1.0 = full adherence to the control signal, 0.0 = ignore it." + "--reference_video_strength", type=float, default=None, help="(v2av) IC-LoRA reference-video conditioning strength in [0.0, 1.0]. 1.0 = full adherence to the control signal, 0.0 = ignore it." ) parser.add_argument("--reference_video_frame_cap", type=int, default=None, help="(v2av) Maximum number of frames to read from the reference/control video. Defaults to the full clip.") parser.add_argument("--mux_audio_video_path", type=str, default=None, help="(v2av, optional) After saving, mux audio from this file into the output mp4 (ffmpeg). ") args = parser.parse_args() - if args.model_cls == WAN_ANIMATE2_MODEL_ID and args.seed < 0: - parser.error(f"{WAN_ANIMATE2_MODEL_ID} requires a non-negative --seed") - seed_all(args.seed) - - # set config - config = set_config(args) - # init input_info - input_info = init_empty_input_info(args.task, args.support_tasks) - - if config["parallel"]: + startup_config, request_data = build_cli_inputs(args) + if startup_config["parallel"]: platform_device = PLATFORM_DEVICE_REGISTER.get(os.getenv("PLATFORM", "cuda"), None) platform_device.init_parallel_env() - set_parallel_config(config) + init_parallel(startup_config) - print_config(config) + print_config(startup_config, title="Startup config") - validate_config_paths(config) + validate_config_paths(startup_config) with ProfilingContext4DebugL1("Total Cost"): - # init runner - runner = init_runner(config) - # start to infer - data = args.__dict__ - update_input_info_from_dict(input_info, data) - runner.run_pipeline(input_info) + runner = build_runner(startup_config) + input_info = runner.prepare_request(request_data) + print_request(input_info, runner.get_supported_request_fields(input_info.task)) + runner.run_request(input_info) # Clean up distributed process group if dist.is_initialized(): diff --git a/lightx2v/models/networks/bagel/sensenova_tasks.py b/lightx2v/models/networks/bagel/sensenova_tasks.py index 46a454f78..3e11e77ed 100644 --- a/lightx2v/models/networks/bagel/sensenova_tasks.py +++ b/lightx2v/models/networks/bagel/sensenova_tasks.py @@ -179,11 +179,6 @@ def normalize_omni_vision_subtask(subtask): return normalized -def get_omni_vision_task_spec(subtask): - normalized = normalize_omni_vision_subtask(subtask) - return normalized, OMNI_VISION_TASK_SPECS[normalized] - - TASK_TO_MODE = { "raw_query": "dense_perception", "depth": "dense_perception", diff --git a/lightx2v/models/networks/base_model.py b/lightx2v/models/networks/base_model.py index 35a8ceaf7..037b03150 100644 --- a/lightx2v/models/networks/base_model.py +++ b/lightx2v/models/networks/base_model.py @@ -330,7 +330,7 @@ def _apply_weights(self, weight_dict=None): self.post_weight.load(self.original_weight_dict) # Handle LoRA if needed - if self.config.get("lora_dynamic_apply", False): + if self.config.get("lora_dynamic_apply", False) and self.lora_path is not None: assert self.config.get("lora_configs", False) if hasattr(self, "_register_lora"): self._register_lora(self.lora_path, self.lora_strength) diff --git a/lightx2v/models/networks/cosmos3/model.py b/lightx2v/models/networks/cosmos3/model.py index e2c827e40..3f1212ed2 100644 --- a/lightx2v/models/networks/cosmos3/model.py +++ b/lightx2v/models/networks/cosmos3/model.py @@ -154,18 +154,19 @@ def infer(self, inputs): self.post_weight.to_cuda() text_encoder_output = inputs["text_encoder_output"] - do_cfg = self.config.get("enable_cfg", True) and self.scheduler.sample_guide_scale != 1.0 - if do_cfg and self.config.get("cfg_parallel", False): - cfg_p_group = self.config["device_mesh"].get_group(mesh_dim="cfg_p") - assert dist.get_world_size(cfg_p_group) == 2, "cfg_p_world_size must be equal to 2" - cfg_p_rank = dist.get_rank(cfg_p_group) - input_ids = text_encoder_output["cond_input_ids"] if cfg_p_rank == 0 else text_encoder_output["uncond_input_ids"] - output = self._infer_cond_uncond(input_ids) - cond, uncond = self._gather_cfg_parallel_output(output, cfg_p_group) - self._set_scheduler_noise_pred(self._combine_cfg_output(cond, uncond)) - elif do_cfg: - cond = self._infer_cond_uncond(text_encoder_output["cond_input_ids"]) - uncond = self._infer_cond_uncond(text_encoder_output["uncond_input_ids"]) + do_cfg = self.config.get("enable_cfg", True) + if do_cfg: + assert self.scheduler.sample_guide_scale != 1.0, "enable_cfg=true requires sample_guide_scale != 1" + if self.config.get("cfg_parallel", False): + cfg_p_group = self.config["device_mesh"].get_group(mesh_dim="cfg_p") + assert dist.get_world_size(cfg_p_group) == 2, "cfg_p_world_size must be equal to 2" + cfg_p_rank = dist.get_rank(cfg_p_group) + input_ids = text_encoder_output["cond_input_ids"] if cfg_p_rank == 0 else text_encoder_output["uncond_input_ids"] + output = self._infer_cond_uncond(input_ids) + cond, uncond = self._gather_cfg_parallel_output(output, cfg_p_group) + else: + cond = self._infer_cond_uncond(text_encoder_output["cond_input_ids"]) + uncond = self._infer_cond_uncond(text_encoder_output["uncond_input_ids"]) self._set_scheduler_noise_pred(self._combine_cfg_output(cond, uncond)) else: cond = self._infer_cond_uncond(text_encoder_output["cond_input_ids"]) diff --git a/lightx2v/models/networks/ernie_image/model.py b/lightx2v/models/networks/ernie_image/model.py index 1906b45a3..c77e63b3b 100644 --- a/lightx2v/models/networks/ernie_image/model.py +++ b/lightx2v/models/networks/ernie_image/model.py @@ -74,6 +74,7 @@ def infer(self, inputs): text_output = inputs["text_encoder_output"] if self.config.get("enable_cfg", False): + assert self.scheduler.sample_guide_scale != 1.0, "enable_cfg=true requires sample_guide_scale != 1" noise_pred_cond = self._infer_cond_uncond( latents, text_output["prompt_embeds"], diff --git a/lightx2v/models/networks/flux2/infer/utils.py b/lightx2v/models/networks/flux2/infer/utils.py deleted file mode 100644 index 4af3d510f..000000000 --- a/lightx2v/models/networks/flux2/infer/utils.py +++ /dev/null @@ -1 +0,0 @@ -"""RoPE is selected by Flux2 attention weights and applied directly during inference.""" diff --git a/lightx2v/models/networks/flux2/model.py b/lightx2v/models/networks/flux2/model.py index e94acbac6..283cbb850 100644 --- a/lightx2v/models/networks/flux2/model.py +++ b/lightx2v/models/networks/flux2/model.py @@ -437,9 +437,10 @@ def _init_infer_class(self): @torch.no_grad() def infer(self, inputs): latents = self.scheduler.latents - do_cfg = self.config.get("enable_cfg", True) and self.config.get("sample_guide_scale", 1.0) > 1.0 + do_cfg = self.config.get("enable_cfg", True) if do_cfg: + assert self.scheduler.sample_guide_scale != 1.0, "enable_cfg=true requires sample_guide_scale != 1" use_cfg_parallel = self.config.get("cfg_parallel", False) if use_cfg_parallel and hasattr(self.scheduler, "input_image_latents") and self.scheduler.input_image_latents is not None: if hasattr(self.scheduler, "image_rotary_emb") and hasattr(self.scheduler, "negative_image_rotary_emb"): @@ -482,7 +483,7 @@ def infer(self, inputs): noise_pred_cond = noise_pred_list[0] noise_pred_uncond = noise_pred_list[1] - guidance_scale = self.config.get("sample_guide_scale", 1.0) + guidance_scale = self.scheduler.sample_guide_scale noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) self.scheduler.noise_pred = noise_pred else: @@ -504,7 +505,7 @@ def infer(self, inputs): img_ids=img_ids, ) - guidance_scale = self.config.get("sample_guide_scale", 1.0) + guidance_scale = self.scheduler.sample_guide_scale noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) self.scheduler.noise_pred = noise_pred else: diff --git a/lightx2v/models/networks/hidream_o1_image/infer/rope.py b/lightx2v/models/networks/hidream_o1_image/infer/rope.py deleted file mode 100644 index 73d0b9748..000000000 --- a/lightx2v/models/networks/hidream_o1_image/infer/rope.py +++ /dev/null @@ -1 +0,0 @@ -"""HiDream RoPE is selected by decoder-block weights and applied directly during inference.""" diff --git a/lightx2v/models/networks/hunyuan_image3/model.py b/lightx2v/models/networks/hunyuan_image3/model.py index bbb983621..7de630628 100644 --- a/lightx2v/models/networks/hunyuan_image3/model.py +++ b/lightx2v/models/networks/hunyuan_image3/model.py @@ -578,7 +578,9 @@ def _set_cfg_scheduler_predictions(self, noise_pred_cond, noise_pred_uncond, noi self.scheduler.noise_pred = noise_pred_guided def combine_cfg_predictions(self, noise_pred_cond, noise_pred_uncond): - return noise_pred_uncond + self._guidance_scale() * (noise_pred_cond - noise_pred_uncond) + guidance_scale = self._guidance_scale() + assert guidance_scale != 1.0, "CFG requires guidance_scale != 1" + return noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) @torch.no_grad() def infer_branch(self, inputs, infer_condition=True): diff --git a/lightx2v/models/networks/hunyuan_video/model.py b/lightx2v/models/networks/hunyuan_video/model.py index 1840c6368..5c3a8620c 100755 --- a/lightx2v/models/networks/hunyuan_video/model.py +++ b/lightx2v/models/networks/hunyuan_video/model.py @@ -94,6 +94,7 @@ def infer(self, inputs): self.transformer_weights.non_block_weights_to_cuda() if self.config["enable_cfg"]: + assert self.scheduler.sample_guide_scale != 1.0, "enable_cfg=true requires sample_guide_scale != 1" if self.config["cfg_parallel"]: # ==================== CFG Parallel Processing ==================== cfg_p_group = self.config["device_mesh"].get_group(mesh_dim="cfg_p") diff --git a/lightx2v/models/networks/longcat_image/infer/utils.py b/lightx2v/models/networks/longcat_image/infer/utils.py deleted file mode 100644 index aa894f845..000000000 --- a/lightx2v/models/networks/longcat_image/infer/utils.py +++ /dev/null @@ -1 +0,0 @@ -"""RoPE is selected by LongCat attention weights and applied directly during inference.""" diff --git a/lightx2v/models/networks/longcat_image/model.py b/lightx2v/models/networks/longcat_image/model.py index b1073a8d4..d20a9a802 100755 --- a/lightx2v/models/networks/longcat_image/model.py +++ b/lightx2v/models/networks/longcat_image/model.py @@ -118,6 +118,7 @@ def infer(self, inputs): latents = self.scheduler.latents if self.config.get("enable_cfg", True): + assert self.scheduler.sample_guide_scale != 1.0, "enable_cfg=true requires sample_guide_scale != 1" # Check if CFG parallel should be used # Note: I2I task may have different sequence lengths for positive/negative prompts, # which is not yet supported in CFG parallel mode diff --git a/lightx2v/models/networks/ltx2/model.py b/lightx2v/models/networks/ltx2/model.py index 7c6c6f37a..64fdeb32d 100755 --- a/lightx2v/models/networks/ltx2/model.py +++ b/lightx2v/models/networks/ltx2/model.py @@ -679,31 +679,30 @@ def infer(self, inputs): if self.config["cfg_parallel"]: raise NotImplementedError("LTX2 mm_guider 与 cfg_parallel 同时使用尚未实现,请关闭其一。") self._infer_mm_guider_cfg(inputs) - elif self.config["cfg_parallel"]: - # ==================== CFG Parallel Processing ==================== - cfg_p_group = self.config["device_mesh"].get_group(mesh_dim="cfg_p") - assert dist.get_world_size(cfg_p_group) == 2, "cfg_p_world_size must be equal to 2" - cfg_p_rank = dist.get_rank(cfg_p_group) - if cfg_p_rank == 0: - v_noise_pred, a_noise_pred = self._infer_cond_uncond(inputs, infer_condition=True) - else: - v_noise_pred, a_noise_pred = self._infer_cond_uncond(inputs, infer_condition=False) - - v_noise_pred_list = [torch.zeros_like(v_noise_pred) for _ in range(2)] - a_noise_pred_list = [torch.zeros_like(a_noise_pred) for _ in range(2)] - dist.all_gather(v_noise_pred_list, v_noise_pred, group=cfg_p_group) - dist.all_gather(a_noise_pred_list, a_noise_pred, group=cfg_p_group) - v_noise_pred_cond = v_noise_pred_list[0] # cfg_p_rank == 0 - v_noise_pred_uncond = v_noise_pred_list[1] # cfg_p_rank == 1 - a_noise_pred_cond = a_noise_pred_list[0] # cfg_p_rank == 0 - a_noise_pred_uncond = a_noise_pred_list[1] # cfg_p_rank == 1 - - self.scheduler.v_noise_pred = v_noise_pred_uncond + self.scheduler.sample_guide_scale * (v_noise_pred_cond - v_noise_pred_uncond) - self.scheduler.a_noise_pred = a_noise_pred_uncond + self.scheduler.sample_guide_scale * (a_noise_pred_cond - a_noise_pred_uncond) else: - # ==================== CFG Processing ==================== - v_noise_pred_cond, a_noise_pred_cond = self._infer_cond_uncond(inputs, infer_condition=True) - v_noise_pred_uncond, a_noise_pred_uncond = self._infer_cond_uncond(inputs, infer_condition=False) + assert self.scheduler.sample_guide_scale != 1.0, "enable_cfg=true requires sample_guide_scale != 1" + if self.config["cfg_parallel"]: + # ==================== CFG Parallel Processing ==================== + cfg_p_group = self.config["device_mesh"].get_group(mesh_dim="cfg_p") + assert dist.get_world_size(cfg_p_group) == 2, "cfg_p_world_size must be equal to 2" + cfg_p_rank = dist.get_rank(cfg_p_group) + if cfg_p_rank == 0: + v_noise_pred, a_noise_pred = self._infer_cond_uncond(inputs, infer_condition=True) + else: + v_noise_pred, a_noise_pred = self._infer_cond_uncond(inputs, infer_condition=False) + + v_noise_pred_list = [torch.zeros_like(v_noise_pred) for _ in range(2)] + a_noise_pred_list = [torch.zeros_like(a_noise_pred) for _ in range(2)] + dist.all_gather(v_noise_pred_list, v_noise_pred, group=cfg_p_group) + dist.all_gather(a_noise_pred_list, a_noise_pred, group=cfg_p_group) + v_noise_pred_cond = v_noise_pred_list[0] # cfg_p_rank == 0 + v_noise_pred_uncond = v_noise_pred_list[1] # cfg_p_rank == 1 + a_noise_pred_cond = a_noise_pred_list[0] # cfg_p_rank == 0 + a_noise_pred_uncond = a_noise_pred_list[1] # cfg_p_rank == 1 + else: + # ==================== CFG Processing ==================== + v_noise_pred_cond, a_noise_pred_cond = self._infer_cond_uncond(inputs, infer_condition=True) + v_noise_pred_uncond, a_noise_pred_uncond = self._infer_cond_uncond(inputs, infer_condition=False) self.scheduler.v_noise_pred = v_noise_pred_uncond + self.scheduler.sample_guide_scale * (v_noise_pred_cond - v_noise_pred_uncond) self.scheduler.a_noise_pred = a_noise_pred_uncond + self.scheduler.sample_guide_scale * (a_noise_pred_cond - a_noise_pred_uncond) diff --git a/lightx2v/models/networks/qwen_image/model.py b/lightx2v/models/networks/qwen_image/model.py index c8619b1fd..df5718ada 100755 --- a/lightx2v/models/networks/qwen_image/model.py +++ b/lightx2v/models/networks/qwen_image/model.py @@ -108,6 +108,7 @@ def infer(self, inputs): latents_input = latents if self.config["enable_cfg"]: + assert self.scheduler.sample_guide_scale != 1.0, "enable_cfg=true requires sample_guide_scale != 1" if self.config["cfg_parallel"]: # ==================== CFG Parallel Processing ==================== cfg_p_group = self.config["device_mesh"].get_group(mesh_dim="cfg_p") diff --git a/lightx2v/models/networks/wan/animate2_identity.py b/lightx2v/models/networks/wan/animate2_identity.py deleted file mode 100644 index fd2b1815a..000000000 --- a/lightx2v/models/networks/wan/animate2_identity.py +++ /dev/null @@ -1 +0,0 @@ -WAN_ANIMATE2_MODEL_ID = "wan2.2_animate2_distilled" diff --git a/lightx2v/models/networks/wan/dreamzero_model.py b/lightx2v/models/networks/wan/dreamzero_model.py index 9a54e6b6c..d48c8e6ec 100644 --- a/lightx2v/models/networks/wan/dreamzero_model.py +++ b/lightx2v/models/networks/wan/dreamzero_model.py @@ -198,6 +198,7 @@ def infer(self, inputs): cache_name = inputs.get("cache_name", "pos") if enable_cfg: + assert guide_scale != 1.0, "enable_cfg=true requires guide_scale != 1" if self.config.get("cfg_parallel", False): cond_video, cond_action, uncond_video = self._infer_cfg_parallel( inputs, diff --git a/lightx2v/models/networks/wan/model.py b/lightx2v/models/networks/wan/model.py index 9da2efd4d..fa40cb806 100755 --- a/lightx2v/models/networks/wan/model.py +++ b/lightx2v/models/networks/wan/model.py @@ -305,6 +305,7 @@ def infer(self, inputs): self.transformer_weights.non_block_weights_to_cuda() if self.config["enable_cfg"]: + assert self.scheduler.sample_guide_scale != 1.0, "enable_cfg=true requires sample_guide_scale != 1" if self.config["cfg_parallel"]: # ==================== CFG Parallel Processing ==================== cfg_p_group = self.config["device_mesh"].get_group(mesh_dim="cfg_p") diff --git a/lightx2v/models/networks/wan/weights/pre_weights.py b/lightx2v/models/networks/wan/weights/pre_weights.py index 1c8c687a3..742fe15a4 100755 --- a/lightx2v/models/networks/wan/weights/pre_weights.py +++ b/lightx2v/models/networks/wan/weights/pre_weights.py @@ -1,5 +1,4 @@ from lightx2v.common.modules.weight_module import WeightModule -from lightx2v.models.networks.wan.animate2_identity import WAN_ANIMATE2_MODEL_ID from lightx2v.utils.registry_factory import CONV3D_WEIGHT_REGISTER, EMBEDDING_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, TENSOR_REGISTER @@ -101,7 +100,7 @@ def __init__(self, config): if config["task"] in ["i2v", "flf2v", "animate", "s2v", "rs2v"] and config.get("use_image_encoder", True): # Wan-Animate-2's MLPProj uses nn.LayerNorm's default epsilon. # Preserve the established epsilon for all other Wan variants. - image_proj_norm_eps = 1e-5 if config["model_cls"] == WAN_ANIMATE2_MODEL_ID else 1e-6 + image_proj_norm_eps = 1e-5 if config["model_cls"] == "wan2.2_animate2_distilled" else 1e-6 self.add_module( "proj_0", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]( @@ -170,7 +169,7 @@ def __init__(self, config): "emb_pos", TENSOR_REGISTER["Default"](f"img_emb.emb_pos"), ) - if config["task"] == "animate" and config.get("model_cls") != WAN_ANIMATE2_MODEL_ID: + if config["task"] == "animate" and config.get("model_cls") != "wan2.2_animate2_distilled": self.add_module( "pose_patch_embedding", CONV3D_WEIGHT_REGISTER["Default"]( diff --git a/lightx2v/models/networks/worldplay/pose_utils.py b/lightx2v/models/networks/worldplay/pose_utils.py index cb5c6ed23..2052df9dc 100644 --- a/lightx2v/models/networks/worldplay/pose_utils.py +++ b/lightx2v/models/networks/worldplay/pose_utils.py @@ -231,18 +231,16 @@ def pose_string_to_json(pose_string): return pose_json -def get_latent_num_from_pose(pose_data): - """Get the number of latent frames from pose data without full tensor conversion.""" +def load_pose(pose_data): + """Load pose data from a command string, JSON path, or dictionary.""" + if isinstance(pose_data, dict): + return pose_data if isinstance(pose_data, str): if pose_data.endswith(".json"): - pose_json = json.load(open(pose_data, "r")) - else: - pose_json = pose_string_to_json(pose_data) - elif isinstance(pose_data, dict): - pose_json = pose_data - else: - raise ValueError(f"Invalid pose_data type: {type(pose_data)}") - return len(pose_json) + with open(pose_data, "r") as f: + return json.load(f) + return pose_string_to_json(pose_data) + raise ValueError(f"Invalid pose_data type: {type(pose_data)}. Expected str or dict.") def pose_to_input(pose_data, latent_num, tps=False): @@ -263,19 +261,7 @@ def pose_to_input(pose_data, latent_num, tps=False): - intrinsic_list: torch.Tensor (batch, latent_num, 3, 3) - normalized intrinsics - action_one_label: torch.Tensor (batch, latent_num) - discrete action labels (0-80) """ - # Handle different input types - if isinstance(pose_data, str): - if pose_data.endswith(".json"): - # Load from JSON file - pose_json = json.load(open(pose_data, "r")) - else: - # Parse pose string - pose_json = pose_string_to_json(pose_data) - elif isinstance(pose_data, dict): - pose_json = pose_data - else: - raise ValueError(f"Invalid pose_data type: {type(pose_data)}. Expected str or dict.") - + pose_json = load_pose(pose_data) pose_keys = list(pose_json.keys()) latent_num_from_pose = len(pose_keys) if latent_num_from_pose != latent_num: diff --git a/lightx2v/models/networks/z_image/infer/post_infer.py b/lightx2v/models/networks/z_image/infer/post_infer.py index 997ceb511..7e8864561 100755 --- a/lightx2v/models/networks/z_image/infer/post_infer.py +++ b/lightx2v/models/networks/z_image/infer/post_infer.py @@ -56,9 +56,9 @@ def infer(self, weights, hidden_states, temb_img_silu, image_tokens_len=None): raise ValueError(f"out_dim mismatch: {out_dim} != {expected_out_dim} (transformer_out_channels={transformer_out_channels})") out_channels = transformer_out_channels - target_shape = self.scheduler.input_info.target_shape + latent_shape = self.scheduler.input_info.latent_shape - _, _, height, width = target_shape + _, _, height, width = latent_shape num_frames = 1 pH = pW = patch_size pF = f_patch_size @@ -68,7 +68,7 @@ def infer(self, weights, hidden_states, temb_img_silu, image_tokens_len=None): expected_T = F_tokens * H_tokens * W_tokens if T != expected_T: - raise ValueError(f"Token count mismatch: T={T} != expected_T={expected_T} (from target_shape={target_shape})") + raise ValueError(f"Token count mismatch: T={T} != expected_T={expected_T} (from latent_shape={latent_shape})") # Unpatchify: [T, out_dim] -> [C, H, W] # Reshape: [T, out_dim] -> [F_tokens, H_tokens, W_tokens, pF, pH, pW, out_channels] diff --git a/lightx2v/models/networks/z_image/infer/pre_infer.py b/lightx2v/models/networks/z_image/infer/pre_infer.py index 52e817824..7077c0baf 100755 --- a/lightx2v/models/networks/z_image/infer/pre_infer.py +++ b/lightx2v/models/networks/z_image/infer/pre_infer.py @@ -130,10 +130,10 @@ def infer(self, weights, hidden_states, encoder_hidden_states): num_tokens, patch_dim = hidden_states.shape - original_shape = self.scheduler.input_info.target_shape - if len(original_shape) >= 2: - original_height = original_shape[-2] - original_width = original_shape[-1] + latent_shape = self.scheduler.input_info.latent_shape + if len(latent_shape) >= 2: + original_height = latent_shape[-2] + original_width = latent_shape[-1] original_frames = 1 F_tokens = original_frames // f_patch_size diff --git a/lightx2v/models/networks/z_image/model.py b/lightx2v/models/networks/z_image/model.py index 03ed3752a..3ceb37a08 100755 --- a/lightx2v/models/networks/z_image/model.py +++ b/lightx2v/models/networks/z_image/model.py @@ -109,6 +109,7 @@ def infer(self, inputs): latents_input = latents if self.config["enable_cfg"]: + assert self.scheduler.sample_guide_scale != 1.0, "enable_cfg=true requires sample_guide_scale != 1" if self.config["cfg_parallel"]: # ==================== CFG Parallel Processing ==================== cfg_p_group = self.config["device_mesh"].get_group(mesh_dim="cfg_p") diff --git a/lightx2v/models/runners/bagel/bagel_runner.py b/lightx2v/models/runners/bagel/bagel_runner.py index 22632871d..a8417a52d 100644 --- a/lightx2v/models/runners/bagel/bagel_runner.py +++ b/lightx2v/models/runners/bagel/bagel_runner.py @@ -9,6 +9,7 @@ from lightx2v.models.runners.bagel.i2i_utils import load_bagel_i2i_input_image, resize_pil_to_shape, resolve_bagel_i2i_image_shape from lightx2v.models.runners.bagel.t2i_utils import get_bagel_latent_downsample, resolve_bagel_t2i_image_shape from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS from lightx2v.models.schedulers.bagel.scheduler import BagelScheduler from lightx2v.models.video_encoders.hf.bagel.vae import BagelVae from lightx2v.server.metrics import monitor_cli @@ -26,6 +27,11 @@ def _has_save_path(input_info): @RUNNER_REGISTER("bagel") class BagelRunner(DefaultRunner): + supported_request_fields_by_task = { + "t2i": COMMON_REQUEST_FIELDS | {"aspect_ratio", "prompt", "target_shape"}, + "i2i": COMMON_REQUEST_FIELDS | {"image_path", "prompt", "target_shape"}, + } + def __init__(self, config): super().__init__(config) @@ -177,9 +183,6 @@ def _finalize_pipeline_outputs(self, input_info, images, latents=None, generator return {"images": images} def run_pipeline(self, input_info): - if self.config["task"] not in ["t2i", "i2i"]: - raise NotImplementedError("BAGEL image generation in LightX2V currently supports task='t2i' and task='i2i'") - self.input_info = input_info logger.info(f"input_info: {self.input_info}") if getattr(self.input_info, "negative_prompt", ""): diff --git a/lightx2v/models/runners/bagel/sensenova_vision_runner.py b/lightx2v/models/runners/bagel/sensenova_vision_runner.py index 15f61659a..348a0b13f 100644 --- a/lightx2v/models/runners/bagel/sensenova_vision_runner.py +++ b/lightx2v/models/runners/bagel/sensenova_vision_runner.py @@ -4,7 +4,6 @@ import gc import json import os -import random from pathlib import Path import numpy as np @@ -14,17 +13,19 @@ from loguru import logger from lightx2v.models.networks.bagel.sensenova_tasks import ( + OMNI_VISION_TASK_SPECS, TEXT_OUTPUT_MODES, clean_text_output, ensure_image_placeholders, get_mode_profile, - get_omni_vision_task_spec, + normalize_omni_vision_subtask, resolve_prompt, ) from lightx2v.models.networks.bagel.sensenova_transforms import build_sensenova_transforms from lightx2v.models.networks.bagel.sensenova_vision_model import SenseNovaVisionModel from lightx2v.models.runners.bagel.bagel_runner import BagelRunner from lightx2v.models.runners.bagel.sensenova_postprocess import load_official_postprocess, resolve_pose_string +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS from lightx2v.models.video_encoders.hf.bagel.sensenova_vae import SenseNovaVisionVae from lightx2v.utils.registry_factory import RUNNER_REGISTER from lightx2v_platform.base.global_var import AI_DEVICE @@ -36,22 +37,12 @@ def _is_main_process(): return not dist.is_initialized() or dist.get_rank() == 0 -def _set_request_seed(seed): - if seed is None: - return - seed = int(seed) - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False - - @RUNNER_REGISTER("sensenova_vision") class SenseNovaVisionRunner(BagelRunner): + supported_request_fields_by_task = { + "omni_vision_task": COMMON_REQUEST_FIELDS | {"glb_output_path", "image_path", "omni_vision_subtask", "postprocess_predictions", "prompt", "raw_output_path"}, + } + def load_bagel_model(self): return SenseNovaVisionModel(self.config) @@ -62,6 +53,11 @@ def init_modules(self): super().init_modules() self.sensenova_transforms = build_sensenova_transforms() + def prepare_request(self, request_data): + input_info = super().prepare_request(request_data) + input_info.omni_vision_subtask = normalize_omni_vision_subtask(input_info.omni_vision_subtask) + return input_info + def _configure_mode(self, mode): profile = get_mode_profile(mode) if "num_timesteps" not in profile: @@ -189,11 +185,10 @@ def _save_images(self, images, input_info, log_prefix="SenseNova-Vision image sa def _postprocess_recon3d(self, pointmaps, prepared, input_info, actual_image_count): pointmaps = np.asarray(pointmaps[:actual_image_count], dtype=np.float32) raw_path = getattr(input_info, "raw_output_path", "") - if not raw_path: - save_path = getattr(input_info, "save_result_path", "") or "sensenova_recon3d.npy" - raw_path = str(self._derive_path(save_path, "_raw", ".npy")) - raw_path = Path(raw_path) - if _is_main_process(): + if not raw_path and input_info.save_result_path is not None: + raw_path = self._derive_path(input_info.save_result_path, "_raw", ".npy") + raw_path = Path(raw_path) if raw_path else None + if raw_path is not None and _is_main_process(): raw_path.parent.mkdir(parents=True, exist_ok=True) np.save(raw_path, pointmaps) logger.info(f"SenseNova-Vision raw point maps saved: {raw_path}") @@ -218,19 +213,18 @@ def _postprocess_recon3d(self, pointmaps, prepared, input_info, actual_image_cou mask_black_bg=False, mask_white_bg=False, ) - if not glb_path: + if not glb_path and raw_path is not None: glb_path = str(raw_path.with_name(f"{raw_path.stem}_scene.glb")) - if _is_main_process(): + if glb_path and _is_main_process(): Path(glb_path).parent.mkdir(parents=True, exist_ok=True) scene.export(file_obj=glb_path) logger.info(f"SenseNova-Vision reconstructed scene saved: {glb_path}") - return pointmaps, scene, str(raw_path), glb_path or None + return pointmaps, scene, str(raw_path) if raw_path is not None else None, glb_path or None def run_pipeline(self, input_info): self.input_info = input_info - _set_request_seed(getattr(input_info, "seed", 42)) - subtask, task_spec = get_omni_vision_task_spec(getattr(input_info, "omni_vision_subtask", "")) - input_info.omni_vision_subtask = subtask + subtask = input_info.omni_vision_subtask + task_spec = OMNI_VISION_TASK_SPECS[subtask] task = task_spec.runner_task mode = task_spec.mode self._configure_mode(mode) diff --git a/lightx2v/models/runners/base_runner.py b/lightx2v/models/runners/base_runner.py index 057090c9a..bcb0cf77b 100755 --- a/lightx2v/models/runners/base_runner.py +++ b/lightx2v/models/runners/base_runner.py @@ -7,6 +7,8 @@ import torch.distributed as dist from loguru import logger +from lightx2v.utils.input_info import INPUT_INFO_TYPES, UNSET, InputInfo +from lightx2v.utils.utils import seed_all from lightx2v_platform.base.global_var import AI_DEVICE @@ -16,8 +18,17 @@ class BaseRunner(ABC): Defines interface methods that all subclasses must implement """ + input_info_cls_by_task: dict[str, type[InputInfo]] = {} + supported_request_fields_by_task: dict[str, frozenset[str]] = {} + def __init__(self, config): self.config = config + task = config.get("task") + if not task: + raise ValueError("task must be set when the runner is created") + if task not in self.supported_request_fields_by_task: + raise ValueError(f"{type(self).__name__} does not support task {task!r}") + self.supported_tasks = self.get_supported_tasks() self.vae_encoder_need_img_original = False self.input_info = None self.enable_reuse = config.get("enable_reuse", False) @@ -68,6 +79,60 @@ def warmup(self): if self.config.get("warmup", False): raise NotImplementedError(f"Warmup is not supported for {type(self).__name__}") + def get_supported_tasks(self): + """Return tasks accepted by this initialized runner.""" + return (self.config["task"],) + + def create_input_info(self, request_data): + """Create the runtime context for one inference request.""" + task = request_data["task"] + input_info_cls = self.input_info_cls_by_task.get(task) or INPUT_INFO_TYPES[task] + input_info = input_info_cls() + input_info.update(self.config) + input_info.update(request_data) + + if "aspect_ratio" in request_data and "target_shape" not in request_data: + input_info.target_shape = [] + elif "target_shape" not in request_data and "target_shape" not in self.config and "target_height" in self.config and "target_width" in self.config: + input_info.update({"target_shape": [self.config["target_height"], self.config["target_width"]]}) + + input_info.seed = self.resolve_request_seed(request_data) + return input_info + + def resolve_request_seed(self, request_data): + seed = request_data.get("seed") + return 42 if seed is None else seed + + def get_supported_request_fields(self, task): + """Return supported request fields for the given task.""" + supported_request_fields = self.supported_request_fields_by_task[task] + if not self.config.get("enable_cfg", False): + supported_request_fields = supported_request_fields - {"negative_prompt"} + return supported_request_fields + + def prepare_request(self, request_data): + """Build and validate the runtime context for one request.""" + request_data = {key: value for key, value in request_data.items() if value is not UNSET and value is not None} + task = request_data.get("task") + if task is None: + if len(self.supported_tasks) > 1: + raise ValueError("task is required when the runner supports multiple tasks") + task = self.config["task"] + request_data["task"] = task + if task not in self.supported_tasks: + task_names = ", ".join(self.supported_tasks) + raise ValueError(f"Task {task!r} is not supported by this runner; expected one of: {task_names}") + unsupported_fields = set(request_data) - self.get_supported_request_fields(task) + if unsupported_fields: + raise ValueError(f"{type(self).__name__} ({task}) does not support request fields: {', '.join(sorted(unsupported_fields))}") + return self.create_input_info(request_data) + + def run_request(self, input_info): + """Run a request that has already passed request preparation.""" + if input_info.seed is not None: + seed_all(input_info.seed) + return self.run_pipeline(input_info) + def set_reuse(self, reuse, reuse_prefix_segments=0): if reuse and not self.enable_reuse: raise ValueError(f"This {type(self).__name__} service does not enable reuse") @@ -107,51 +172,6 @@ def _maybe_freeze_gc(self): self._gc_frozen = True logger.info(f"[GC] gc.collect() reclaimed {collected} objects; gc.freeze() moved ~{n} live tracked objects out of future GC walks") - def apply_disagg_request_overrides(self, config_modify): - """Mirror flat disagg request fields into ``disagg_config`` in disagg mode only.""" - if not isinstance(config_modify, dict): - return - if not self.config.get("disagg_mode"): - return - disagg_config = self.config.get("disagg_config") - if not isinstance(disagg_config, dict): - return - - def _safe_int(key): - value = config_modify.get(key) - if value is None: - return None - try: - return int(value) - except (TypeError, ValueError): - return None - - with self.config.temporarily_unlocked(): - data_bootstrap_room = _safe_int("data_bootstrap_room") - if data_bootstrap_room is not None: - self.config["data_bootstrap_room"] = data_bootstrap_room - - disagg_bootstrap_room = _safe_int("disagg_bootstrap_room") - if disagg_bootstrap_room is not None: - disagg_config["bootstrap_room"] = disagg_bootstrap_room - self.config["data_bootstrap_room"] = disagg_bootstrap_room - - decoder_bootstrap_room = _safe_int("disagg_decoder_bootstrap_room") - if decoder_bootstrap_room is not None: - disagg_config["decoder_bootstrap_room"] = decoder_bootstrap_room - - phase1_receiver_engine_rank = _safe_int("disagg_phase1_receiver_engine_rank") - if phase1_receiver_engine_rank is not None: - self.config["disagg_phase1_receiver_engine_rank"] = phase1_receiver_engine_rank - - for flat_key, disagg_key in ( - ("disagg_phase1_receiver_engine_rank", "receiver_engine_rank"), - ("disagg_phase2_sender_engine_rank", "receiver_engine_rank"), - ): - value = _safe_int(flat_key) - if value is not None: - disagg_config[disagg_key] = value - def load_transformer(self): """Load transformer model diff --git a/lightx2v/models/runners/cosmos3/cosmos3_runner.py b/lightx2v/models/runners/cosmos3/cosmos3_runner.py index 0c1b753eb..5b14fceb6 100644 --- a/lightx2v/models/runners/cosmos3/cosmos3_runner.py +++ b/lightx2v/models/runners/cosmos3/cosmos3_runner.py @@ -24,9 +24,11 @@ normalize_policy_prompt_format, ) from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, PROMPT_FIELDS, VIDEO_REQUEST_FIELDS from lightx2v.models.schedulers.cosmos3.scheduler import Cosmos3Scheduler from lightx2v.models.video_encoders.hf.cosmos3.vae import Cosmos3WanVAE from lightx2v.utils.envs import * +from lightx2v.utils.input_info import Cosmos3InputInfo from lightx2v.utils.profiler import * from lightx2v.utils.registry_factory import RUNNER_REGISTER from lightx2v.utils.utils import save_to_video @@ -128,9 +130,54 @@ def compose_droid_policy_image(images): @RUNNER_REGISTER("cosmos3") class Cosmos3Runner(DefaultRunner): + input_info_cls_by_task = {task: Cosmos3InputInfo for task in ("t2i", "t2v", "i2v", "t2av", "i2av", "i2va", "v2av")} + supported_request_fields_by_task = { + "t2i": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | {"target_shape"}, + "t2v": VIDEO_REQUEST_FIELDS, + "i2v": VIDEO_REQUEST_FIELDS | {"image_path"}, + "t2av": VIDEO_REQUEST_FIELDS, + "i2av": VIDEO_REQUEST_FIELDS | {"image_path"}, + "i2va": COMMON_REQUEST_FIELDS + | PROMPT_FIELDS + | { + "action_mode", + "action_path", + "domain_name", + "image_path", + "policy_image", + "policy_state", + "save_action_path", + "state_path", + "target_shape", + "video_path", + "view_point", + }, + "v2av": COMMON_REQUEST_FIELDS + | PROMPT_FIELDS + | { + "action_mode", + "action_path", + "domain_name", + "image_path", + "save_action_path", + "target_shape", + "video_path", + "view_point", + }, + } + model_cpu_offload_seq = "transformer->vae->sound_tokenizer" _callback_tensor_inputs = ["latents"] + def create_input_info(self, request_data): + input_info = super().create_input_info(request_data) + # Action metadata resolves explicit request values, then the action file, then config. + input_info.domain_name = request_data.get("domain_name", "") + input_info.view_point = request_data.get("view_point", "") + input_info.action_chunk_size = None + input_info.raw_action_dim = None + return input_info + @ProfilingContext4DebugL2("Load models") def load_model(self): self.model = self.load_transformer() @@ -169,8 +216,6 @@ def init_modules(self): assert self.config.get("cpu_offload", False) if hasattr(self, "model") and self.model is not None: self.model.set_scheduler(self.scheduler) - if self.config["task"] not in ("t2i", "t2v", "i2v", "t2av", "i2av", "i2va", "v2av"): - raise NotImplementedError(f"Cosmos3Runner currently supports tasks t2i/t2v/i2v/t2av/i2av/i2va/v2av, got {self.config['task']}") if (self.config.get("enable_sound", False) or self.config["task"] in ("t2av", "i2av")) and not self.config.get("sound_gen", False): raise ValueError("Cosmos3 sound generation requires a checkpoint with sound_gen=True.") if (self.config.get("action_mode", "") or self.config["task"] in ("i2va", "v2av")) and not self.config.get("action_gen", False): @@ -229,14 +274,6 @@ def _get_action_value(self, key, default=None): return spec[key] return self.config.get(key, default) - def _get_target_video_length(self): - if self._get_action_mode(): - return int(getattr(self.input_info, "target_video_length", 0) or self.config.get("target_video_length", 1)) - input_frames = int(getattr(self.input_info, "target_video_length", 0) or 0) - if input_frames and input_frames != 81: - return input_frames - return int(self.config.get("target_video_length", input_frames or 1)) - def _prepare_action_context(self): if hasattr(self, "_action_spec"): del self._action_spec @@ -248,8 +285,6 @@ def _prepare_action_context(self): chunk_size = int(self._get_action_value("action_chunk_size", self.config.get("action_chunk_size", 16))) self.input_info.action_chunk_size = chunk_size self.input_info.target_video_length = chunk_size + 1 - if spec.get("fps") and "target_fps" not in self.config: - self.input_info.target_fps = float(spec["fps"]) @staticmethod def _build_action_json_prompt(description, view_point, num_frames, fps, height, width, additional_view_description=None): @@ -312,9 +347,8 @@ def _tokenize_chat(self, text: str, is_image: bool, use_system_prompt=True): def tokenize_prompt(self, prompt, negative_prompt=None): prompt = self._resolve_prompt_text(prompt) negative_prompt = self._resolve_prompt_text(negative_prompt) if negative_prompt is not None else None - height = int(self.input_info.auto_height) - width = int(self.input_info.auto_width) - num_frames = self._get_target_video_length() + height, width = self.input_info.target_shape + num_frames = self.get_target_video_length() fps = float(self.config.get("target_fps", 24.0)) is_image = num_frames == 1 negative_prompt = "" if negative_prompt is None else negative_prompt @@ -387,7 +421,7 @@ def _run_input_encoder_local(self): "image_encoder_output": None, } - def set_target_shape(self): + def set_latent_shape(self): if len(self.input_info.target_shape) == 2: height, width = self.input_info.target_shape height, width = int(height), int(width) @@ -406,13 +440,12 @@ def set_target_shape(self): height, width = rounded_height, rounded_width latent_channels = int(self.config.get("latent_channel", 48)) - pixel_frames = self._get_target_video_length() + pixel_frames = self.get_target_video_length() latent_frames = (pixel_frames - 1) // temporal_scale + 1 - self.input_info.auto_height = height - self.input_info.auto_width = width - self.input_info.target_shape = (1, latent_channels, latent_frames, height // spatial_scale, width // spatial_scale) + self.input_info.target_shape = [height, width] + self.input_info.latent_shape = (1, latent_channels, latent_frames, height // spatial_scale, width // spatial_scale) self.input_info.image_shapes = [[(latent_frames, height // spatial_scale, width // spatial_scale)]] - logger.info(f"Cosmos3 Runner set target shape: {width}x{height}, latent: {self.input_info.target_shape}") + logger.info(f"Cosmos3 Runner set target shape: {width}x{height}, latent: {self.input_info.latent_shape}") def _load_i2v_condition_frame(self): image_path = getattr(self.input_info, "image_path", "") @@ -420,8 +453,7 @@ def _load_i2v_condition_frame(self): raise ValueError("Cosmos3 i2v requires --image_path.") if not os.path.isfile(image_path): raise FileNotFoundError(f"Cosmos3 i2v image_path does not exist: {image_path}") - height = int(self.input_info.auto_height) - width = int(self.input_info.auto_width) + height, width = self.input_info.target_shape resample = getattr(Image, "Resampling", Image).BILINEAR with Image.open(image_path) as image: image = image.convert("RGB").resize((width, height), resample=resample) @@ -433,13 +465,13 @@ def _load_i2v_condition_frame(self): def _prepare_i2v_condition_latents(self): if self.config["task"] not in ("i2v", "i2av"): return - if hasattr(self.input_info, "vision_condition_latents") and self.input_info.vision_condition_latents is not None: + if self.input_info.vision_condition_latents is not None: return loaded_vae_here = not hasattr(self, "vae") or self.vae is None if loaded_vae_here: self.vae = self.load_vae() frame = self._load_i2v_condition_frame() - num_frames = self._get_target_video_length() + num_frames = self.get_target_video_length() video = frame.unsqueeze(2).expand(-1, -1, num_frames, -1, -1).contiguous() condition_latents = self.vae.encode(video) self.input_info.vision_condition_latents = condition_latents @@ -614,7 +646,7 @@ def _prepare_action_condition_latents(self): action_mode = self._get_action_mode() if not action_mode: return - if hasattr(self.input_info, "action_latents") or hasattr(self.input_info, "action_latent_shape"): + if self.input_info.action_latents is not None or self.input_info.action_latent_shape is not None: return chunk_size = int(getattr(self.input_info, "action_chunk_size", 0) or self._get_action_value("action_chunk_size", 16)) action_dim = int(self.config.get("action_dim", self.config.get("max_action_dim", 64))) @@ -627,11 +659,10 @@ def _prepare_action_condition_latents(self): if raw_action_dim > action_dim: raise ValueError(f"Cosmos3 raw_action_dim={raw_action_dim} exceeds model action_dim={action_dim}") - height = int(self.input_info.auto_height) - width = int(self.input_info.auto_width) + height, width = self.input_info.target_shape num_frames = chunk_size + 1 image_path = getattr(self.input_info, "image_path", None) or self.config.get("image_path", "") - video_path = getattr(self.input_info, "video_path", None) or self.config.get("video_path", "") + video_path = self.input_info.video_path or self.config.get("video_path", "") policy_image = getattr(self.input_info, "policy_image", None) loaded_vae_here = not hasattr(self, "vae") or self.vae is None @@ -711,10 +742,9 @@ def _clear_action_condition_state(self): "action_condition_frame_indexes", "action_domain_id", "raw_action_dim", - "action_start_frame_offset", ): - if hasattr(self.input_info, name): - delattr(self.input_info, name) + setattr(self.input_info, name, None) + self.input_info.action_start_frame_offset = 1 @ProfilingContext4DebugL2("Run DiT") def _run_dit_local(self, total_steps=None): @@ -724,8 +754,7 @@ def _run_dit_local(self, total_steps=None): self._prepare_i2v_condition_latents() self._prepare_action_condition_latents() self.model.scheduler.prepare(self.input_info) - if hasattr(self.input_info, "vision_condition_latents"): - self.input_info.vision_condition_latents = None + self.input_info.vision_condition_latents = None return self.run(total_steps) @ProfilingContext4DebugL1( @@ -949,7 +978,7 @@ def _finalize_pipeline_outputs(self, input_info, images, latents=None, generator return outputs def _is_video_output(self): - return int(self.config.get("target_video_length", 1)) > 1 + return self.get_target_video_length() > 1 def end_run(self): if hasattr(self, "model") and self.model is not None: @@ -1019,7 +1048,7 @@ def _run_action_forward_multichunk_pipeline(self, input_info): def run_pipeline(self, input_info): self.input_info = input_info self._prepare_action_context() - self.set_target_shape() + self.set_latent_shape() self.inputs = self.run_input_encoder() logger.info(f"input_info: {self.input_info}") if self._is_action_forward_multichunk(): @@ -1054,15 +1083,12 @@ class Cosmos3Policy: """ def __init__(self, config, *, actions_per_plan=None, binarize_gripper=True): - from lightx2v.utils.input_info import init_empty_input_info - if str(config.get("action_mode", "")).strip().lower() != "policy": raise ValueError("Cosmos3Policy requires action_mode='policy'.") if str(config.get("domain_name", "")).strip().lower() != "droid_lerobot": raise ValueError("Cosmos3Policy requires domain_name='droid_lerobot'.") self.config = config - self._input_info_factory = lambda: init_empty_input_info("i2va") self.action_dim = int(config.get("raw_action_dim", 8)) self.action_chunk_size = int(config.get("action_chunk_size", 32)) requested = self.action_chunk_size if actions_per_plan is None else int(actions_per_plan) @@ -1084,21 +1110,22 @@ def _plan(self, images, state, task_description): if state.size != self.action_dim: raise ValueError(f"Cosmos3 Policy-DROID state length {state.size} != {self.action_dim}") - input_info = self._input_info_factory() - input_info.seed = self._seed_sequence.next_seed() - input_info.prompt = str(task_description) - input_info.negative_prompt = "" - input_info.action_mode = "policy" - input_info.domain_name = "droid_lerobot" - input_info.view_point = str(self.config.get("view_point", "concat_view")) - input_info.return_result_tensor = True - input_info.policy_image = compose_droid_policy_image(images) - input_info.policy_state = state + request_data = { + "seed": self._seed_sequence.next_seed(), + "prompt": str(task_description), + "action_mode": "policy", + "domain_name": "droid_lerobot", + "view_point": str(self.config.get("view_point", "concat_view")), + "return_result_tensor": True, + "policy_image": compose_droid_policy_image(images), + "policy_state": state, + } if not dist.is_initialized() or dist.get_rank() == 0: - logger.info(f"Cosmos3 policy plan: seed={input_info.seed}, prompt_format={self.prompt_format}") + logger.info(f"Cosmos3 policy plan: seed={request_data['seed']}, prompt_format={self.prompt_format}") - result = self.runner.run_pipeline(input_info) + input_info = self.runner.prepare_request(request_data) + result = self.runner.run_request(input_info) chunk = result.get("action") if isinstance(result, dict) else None if chunk is None: raise RuntimeError("Cosmos3 Policy-DROID inference returned no action chunk") diff --git a/lightx2v/models/runners/default_runner.py b/lightx2v/models/runners/default_runner.py index 2d5ff4c0f..3c6843b02 100755 --- a/lightx2v/models/runners/default_runner.py +++ b/lightx2v/models/runners/default_runner.py @@ -287,30 +287,15 @@ def load_model(self): self.vfi_model = self.load_vfi_model() if "video_frame_interpolation" in self.config else None self.vsr_model = self.load_vsr_model() if "video_super_resolution" in self.config else None - def set_inputs(self, inputs): - self.input_info.seed = inputs.get("seed", 42) - self.input_info.prompt = inputs.get("prompt", "") - if "prompt_ref" in self.input_info.__dataclass_fields__: - self.input_info.prompt_ref = inputs.get("prompt_ref", self.input_info.prompt_ref) - self.input_info.negative_prompt = inputs.get("negative_prompt", "") - if "image_path" in self.input_info.__dataclass_fields__: - self.input_info.image_path = inputs.get("image_path", "") - if "state_path" in self.input_info.__dataclass_fields__: - self.input_info.state_path = inputs.get("state_path", "") - if "audio_path" in self.input_info.__dataclass_fields__: - self.input_info.audio_path = inputs.get("audio_path", "") - if "video_path" in self.input_info.__dataclass_fields__: - self.input_info.video_path = inputs.get("video_path", "") - if "src_video" in self.input_info.__dataclass_fields__: - self.input_info.src_video = inputs.get("src_video", "") - self.input_info.save_result_path = inputs.get("save_result_path", "") - if "save_action_path" in self.input_info.__dataclass_fields__: - self.input_info.save_action_path = inputs.get("save_action_path", "") - - def set_config(self, config_modify): - logger.info(f"modify config: {config_modify}") - with self.config.temporarily_unlocked(): - self.config.update(config_modify) + def get_target_video_length(self): + value = getattr(self.input_info, "target_video_length", None) + return int(self.config["target_video_length"] if value is None else value) + + def get_target_size(self): + target_shape = getattr(self.input_info, "target_shape", None) + if target_shape: + return int(target_shape[0]), int(target_shape[1]) + return int(self.config["target_height"]), int(self.config["target_width"]) def set_progress_callback(self, callback): self.progress_callback = callback @@ -413,8 +398,6 @@ def read_image_input(self, img_path): self.input_info.original_size = img_ori.size resize_mode = self.config.get("resize_mode", None) - # Treat empty string the same as missing — InputInfo dataclasses default - # `resize_mode` to "", which used to silently skip the resize branch. if resize_mode: img, h, w = resize_image( img, @@ -479,14 +462,13 @@ def _run_input_encoder_local_flf2v(self): @ProfilingContext4DebugL2("Run Encoders") def _run_input_encoder_local_vace(self): - src_video = self.input_info.src_video - src_mask = self.input_info.src_mask - src_ref_images = self.input_info.src_ref_images + ref_images = self.input_info.src_ref_images + target_height, target_width = self.get_target_size() src_video, src_mask, src_ref_images = self.prepare_source( - [src_video], - [src_mask], - [None if src_ref_images is None else src_ref_images.split(",")], - (self.config["target_width"], self.config["target_height"]), + [self.input_info.video_path or None], + [self.input_info.mask_path or None], + [ref_images.split(",") if ref_images else None], + (target_width, target_height), ) self.src_ref_images = src_ref_images @@ -577,6 +559,9 @@ def run_vae_decoder_stream(self, latents): del self.vae_decoder self.maybe_empty_cache() + def get_output_fps(self): + return getattr(self.input_info, "output_fps", None) or self.config.get("fps", 16) + def process_images_after_vae_decoder(self): return_result_tensor = self.input_info.return_result_tensor save_result = self.input_info.save_result_path is not None @@ -593,10 +578,11 @@ def process_images_after_vae_decoder(self): if "video_frame_interpolation" in self.config: assert self.vfi_model is not None and self.config["video_frame_interpolation"].get("target_fps", None) is not None target_fps = self.config["video_frame_interpolation"]["target_fps"] - logger.info(f"Interpolating frames from {self.config.get('fps', 16)} to {target_fps}") + source_fps = self.get_output_fps() + logger.info(f"Interpolating frames from {source_fps} to {target_fps}") self.gen_video_final = self.vfi_model.interpolate_frames( self.gen_video_final, - source_fps=self.config.get("fps", 16), + source_fps=source_fps, target_fps=target_fps, ) @@ -609,7 +595,7 @@ def process_images_after_vae_decoder(self): if "video_frame_interpolation" in self.config and self.config["video_frame_interpolation"].get("target_fps"): fps = self.config["video_frame_interpolation"]["target_fps"] else: - fps = self.config.get("fps", 16) + fps = self.get_output_fps() out_path = self.input_info.save_result_path img_in = (getattr(self.input_info, "image_path", None) or "").strip() diff --git a/lightx2v/models/runners/ernie_image/ernie_image_runner.py b/lightx2v/models/runners/ernie_image/ernie_image_runner.py index 89fa6d276..4717b8dce 100644 --- a/lightx2v/models/runners/ernie_image/ernie_image_runner.py +++ b/lightx2v/models/runners/ernie_image/ernie_image_runner.py @@ -9,6 +9,7 @@ from lightx2v.models.input_encoders.hf.ernie_image.mistral3_model import ErnieImageTextEncoder from lightx2v.models.networks.ernie_image.model import ErnieImageTransformerModel from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import IMAGE_REQUEST_FIELDS from lightx2v.models.schedulers.ernie_image.scheduler import ErnieImageScheduler from lightx2v.models.video_encoders.hf.ernie_image.vae import ErnieImageVAE from lightx2v.utils.envs import GET_DTYPE @@ -31,6 +32,9 @@ def calculate_dimensions(target_area, ratio, multiple_of): class ErnieImageRunner(DefaultRunner): model_cpu_offload_seq = "pe->text_encoder->transformer->vae" _callback_tensor_inputs = ["latents", "prompt_embeds"] + supported_request_fields_by_task = { + "t2i": IMAGE_REQUEST_FIELDS, + } def __init__(self, config): super().__init__(config) @@ -69,8 +73,6 @@ def init_modules(self): elif self.config.get("lazy_load", False): assert self.config.get("cpu_offload", False) self.run_dit = self._run_dit_local - if self.config["task"] != "t2i": - raise NotImplementedError(f"ErnieImageRunner only supports t2i, got: {self.config['task']}") self.run_input_encoder = self._run_input_encoder_local_t2i @ProfilingContext4DebugL2("Run DiT") @@ -101,8 +103,7 @@ def _run_input_encoder_local_t2i(self): @ProfilingContext4DebugL1("Run Text Encoder") def run_text_encoder(self, text, neg_prompt=None): - width = getattr(self.input_info, "auto_width", self.config.get("target_width", 1024)) - height = getattr(self.input_info, "auto_height", self.config.get("target_height", 1024)) + height, width = self.get_target_size() prompt_embeds_list, revised_prompts = self.text_encoders[0].infer( [text], use_pe=self.config.get("use_pe", True), @@ -128,40 +129,34 @@ def run_text_encoder(self, text, neg_prompt=None): text_encoder_output["negative_prompt_embeds"] = negative_prompt_embeds return text_encoder_output - def set_target_shape(self): + def set_latent_shape(self): vae_scale_factor = self.config.get("vae_scale_factor", 16) if len(self.input_info.target_shape) == 2: height, width = [int(v) for v in self.input_info.target_shape] else: - target_height = self.config.get("target_height", None) - target_width = self.config.get("target_width", None) - if target_height and target_width: - height, width = int(target_height), int(target_width) + aspect_ratio = self.input_info.aspect_ratio or self.config.get("aspect_ratio", "1:1") + if ":" in aspect_ratio: + w_ratio, h_ratio = [float(item) for item in aspect_ratio.split(":", 1)] + ratio = w_ratio / h_ratio else: - aspect_ratio = self.input_info.aspect_ratio or self.config.get("aspect_ratio", "1:1") - if ":" in aspect_ratio: - w_ratio, h_ratio = [float(item) for item in aspect_ratio.split(":", 1)] - ratio = w_ratio / h_ratio - else: - ratio = float(aspect_ratio) - width, height = calculate_dimensions( - self.resolution * self.resolution, - ratio, - vae_scale_factor, - ) + ratio = float(aspect_ratio) + width, height = calculate_dimensions( + self.resolution * self.resolution, + ratio, + vae_scale_factor, + ) if height % vae_scale_factor != 0 or width % vae_scale_factor != 0: raise ValueError(f"Height and width must be divisible by {vae_scale_factor}, got {height}x{width}.") - self.input_info.auto_width = width - self.input_info.auto_height = height - self.input_info.target_shape = ( + self.input_info.target_shape = [height, width] + self.input_info.latent_shape = ( 1, self.config.get("in_channels", 128), height // vae_scale_factor, width // vae_scale_factor, ) - logger.info(f"ERNIE-Image target shape: {width}x{height}, latent shape: {self.input_info.target_shape}") + logger.info(f"ERNIE-Image target shape: {width}x{height}, latent shape: {self.input_info.latent_shape}") def run(self, total_steps=None): if total_steps is None: @@ -218,7 +213,7 @@ def _finalize_pipeline_outputs(self, input_info, images, revised_prompts, latent @ProfilingContext4DebugL1("RUN pipeline") def run_pipeline(self, input_info): self.input_info = input_info - self.set_target_shape() + self.set_latent_shape() self.inputs = self.run_input_encoder() logger.info(f"input_info: {self.input_info}") latents, generator = self.run_dit() diff --git a/lightx2v/models/runners/flux2/flux2_runner.py b/lightx2v/models/runners/flux2/flux2_runner.py index c0c92dc89..63c5c52cb 100644 --- a/lightx2v/models/runners/flux2/flux2_runner.py +++ b/lightx2v/models/runners/flux2/flux2_runner.py @@ -1,5 +1,4 @@ import gc -import math import os import numpy as np @@ -8,9 +7,11 @@ from lightx2v.models.networks.flux2.model import Flux2DevTransformerModel, Flux2KleinTransformerModel from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS from lightx2v.models.schedulers.flux2.feature_caching.scheduler import Flux2DevSchedulerCaching, Flux2SchedulerCaching from lightx2v.models.schedulers.flux2.scheduler import Flux2DevScheduler, Flux2Scheduler from lightx2v.models.video_encoders.hf.flux2.vae import Flux2VAE +from lightx2v.utils.input_info import Flux2I2IInputInfo from lightx2v.utils.profiler import ProfilingContext4DebugL1, ProfilingContext4DebugL2 from lightx2v.utils.registry_factory import RUNNER_REGISTER from lightx2v.utils.utils import is_main_process @@ -19,20 +20,21 @@ torch_device_module = getattr(torch, AI_DEVICE) -def calculate_dimensions(target_area, ratio): - width = math.sqrt(target_area * ratio) - height = width / ratio - - width = round(width / 32) * 32 - height = round(height / 32) * 32 - - return width, height, None - - @RUNNER_REGISTER("flux2") class Flux2Runner(DefaultRunner): model_cpu_offload_seq = "text_encoder->transformer->vae" _callback_tensor_inputs = ["latents", "prompt_embeds"] + input_info_cls_by_task = {"i2i": Flux2I2IInputInfo} + supported_request_fields_by_task = { + "t2i": COMMON_REQUEST_FIELDS | {"aspect_ratio", "prompt", "target_shape"}, + "i2i": COMMON_REQUEST_FIELDS | {"image_path", "prompt"}, + } + + def get_supported_request_fields(self, task): + supported_request_fields = super().get_supported_request_fields(task) + if task == "i2i" and self.config.get("inpaint_mask_enabled", False): + supported_request_fields |= {"inpaint_blur_sigma", "inpaint_blur_size"} + return supported_request_fields def __init__(self, config): self.model_variant = config["model_variant"] @@ -88,7 +90,7 @@ def run_text_encoder(self, text, image_list=None, neg_prompt=None): text_encoder_output = {"prompt_embeds": prompt_embeds, "text_ids": text_ids} - uses_cfg = self.config.get("enable_cfg", True) and self.config.get("sample_guide_scale", 1.0) > 1.0 + uses_cfg = self.config.get("enable_cfg", True) if self.model_variant == "klein" and uses_cfg: neg_prompt_embeds_list, _ = self.text_encoders[0].infer([""]) neg_prompt_embeds = neg_prompt_embeds_list[0].unsqueeze(0) @@ -176,7 +178,7 @@ def _run_input_encoder_local_i2i(self): main_img = input_image[0] image_processor.check_image_input(main_img) processed_img, target_shape = self._preprocess_condition_image(image_processor, main_img, max_image_area, vae_scale_factor) - self.input_info.target_shape = target_shape + self.input_info.target_shape = list(target_shape) processed_tensor = processed_img.to(AI_DEVICE) condition_images.extend([processed_tensor, processed_tensor]) @@ -189,7 +191,7 @@ def _run_input_encoder_local_i2i(self): processed_img, target_shape = self._preprocess_condition_image(image_processor, img, max_image_area, vae_scale_factor) condition_images.append(processed_img.to(AI_DEVICE)) if index == 0: - self.input_info.target_shape = target_shape + self.input_info.target_shape = list(target_shape) torch_device_module.empty_cache() gc.collect() @@ -243,8 +245,8 @@ def _prepare_inpaint_mask(self, mask): mask = mask.permute(2, 0, 1).unsqueeze(0) mask = mask.mean(dim=1, keepdim=True) - blur_size = getattr(self.input_info, "inpaint_blur_size", None) - blur_sigma = getattr(self.input_info, "inpaint_blur_sigma", None) + blur_size = self.input_info.inpaint_blur_size + blur_sigma = self.input_info.inpaint_blur_sigma if blur_size is not None and blur_sigma is not None: from torchvision.transforms import GaussianBlur @@ -357,20 +359,13 @@ def get_custom_shape(self): width, height = as_maps[self.config.get("aspect_ratio", "16:9")] return (width, height) - def set_target_shape(self): + def set_latent_shape(self): task = self.config.get("task", "t2i") if task == "i2i": height, width = self.input_info.target_shape else: - custom_shape = self.get_custom_shape() - if custom_shape is not None: - width, height = custom_shape - else: - calculated_width, calculated_height, _ = calculate_dimensions(self.resolution * self.resolution, 16 / 9) - multiple_of = self.config.get("vae_scale_factor", 8) * 2 - width = calculated_width // multiple_of * multiple_of - height = calculated_height // multiple_of * multiple_of - self.input_info.target_shape = (height, width) + width, height = self.get_custom_shape() + self.input_info.target_shape = [height, width] multiple_of = self.config.get("vae_scale_factor", 8) * 2 @@ -383,9 +378,6 @@ def set_target_shape(self): self.input_info.latent_shape = (packed_batch, packed_h * packed_w, packed_channels) self.input_info.latent_image_ids = self._prepare_latent_ids(packed_batch, packed_h, packed_w).to(AI_DEVICE) - def set_img_shapes(self): - pass - @ProfilingContext4DebugL1("Run VAE Decoder") def run_vae_decoder(self, latents): if self.config.get("lazy_load", False) or self.config.get("unload_modules", False): @@ -424,14 +416,13 @@ def run_pipeline(self, input_info): self.inputs = self.run_input_encoder() logger.info(f"input_info: {self.input_info}") - self.set_target_shape() - self.set_img_shapes() + self.set_latent_shape() latents, generator = self.run_dit() images = self.run_vae_decoder(latents) self.end_run() - if not input_info.return_result_tensor and is_main_process(): + if not input_info.return_result_tensor and input_info.save_result_path is not None and is_main_process(): image = images[0] image.save(input_info.save_result_path) logger.info(f"Image saved: {input_info.save_result_path}") diff --git a/lightx2v/models/runners/hidream_o1_image/hidream_o1_image_runner.py b/lightx2v/models/runners/hidream_o1_image/hidream_o1_image_runner.py index 77ac89936..539102f3b 100644 --- a/lightx2v/models/runners/hidream_o1_image/hidream_o1_image_runner.py +++ b/lightx2v/models/runners/hidream_o1_image/hidream_o1_image_runner.py @@ -8,9 +8,11 @@ from lightx2v.models.networks.hidream_o1_image.model import HidreamO1ImageModel from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS from lightx2v.models.schedulers.hidream_o1_image.scheduler import HidreamO1ImageScheduler from lightx2v.server.metrics import monitor_cli from lightx2v.utils.envs import * +from lightx2v.utils.input_info import HidreamI2IInputInfo from lightx2v.utils.profiler import * from lightx2v.utils.registry_factory import RUNNER_REGISTER @@ -35,6 +37,20 @@ def _add_special_tokens(tokenizer): class HidreamO1ImageRunner(DefaultRunner): """LightX2V runner for HiDream-O1-Image t2i / i2i.""" + input_info_cls_by_task = {"i2i": HidreamI2IInputInfo} + supported_request_fields_by_task = { + "t2i": COMMON_REQUEST_FIELDS | {"prompt", "target_shape"}, + "i2i": COMMON_REQUEST_FIELDS + | { + "i2i_denoise_strength", + "image_path", + "keep_original_aspect", + "layout_bboxes", + "prompt", + "target_shape", + }, + } + def __init__(self, config): super().__init__(config) self.processor = None @@ -54,9 +70,6 @@ def init_scheduler(self): def init_modules(self): task = self.config["task"] - if task not in ("t2i", "i2i"): - raise NotImplementedError(f"HidreamO1ImageRunner supports t2i and i2i, got: {task}") - logger.info(f"Initializing HiDream-O1-Image {task} runner...") self.load_model() self.run_input_encoder = self._run_input_encoder_local_t2i if task == "t2i" else self._run_input_encoder_local_i2i @@ -173,10 +186,6 @@ def _run_input_encoder_local_i2i(self): if not ref_image_paths: raise ValueError("HiDream i2i requires --image_path with one or more reference image paths.") - layout_bboxes = getattr(self.input_info, "layout_bboxes", "") or self.config.get("layout_bboxes") - if layout_bboxes in ("", None): - layout_bboxes = None - from lightx2v.models.networks.hidream_o1_image.i2i_utils import build_i2i_samples generation_config = self._resolve_generation_config() @@ -185,15 +194,15 @@ def _run_input_encoder_local_i2i(self): ref_image_paths=ref_image_paths, height=self._resolve_size("height", "target_height", 2048), width=self._resolve_size("width", "target_width", 2048), - keep_original_aspect=getattr(self.input_info, "keep_original_aspect", False) or self.config.get("keep_original_aspect", False), - layout_bboxes=layout_bboxes, + keep_original_aspect=self.input_info.keep_original_aspect, + layout_bboxes=self.input_info.layout_bboxes or None, tokenizer=self.tokenizer, processor=self.processor, model_config=self.model.model_config, device=self.model.device, dtype=self.dtype, enable_cfg=generation_config["enable_cfg"], - i2i_denoise_strength=getattr(self.input_info, "i2i_denoise_strength", None), + i2i_denoise_strength=self.input_info.i2i_denoise_strength, ) for sample in inputs["samples"]: sample["tgt_image_len"] = inputs["tgt_image_len"] @@ -202,7 +211,7 @@ def _run_input_encoder_local_i2i(self): "seed": self.input_info.seed, "save_result_path": self.input_info.save_result_path, "generation_config": generation_config, - "i2i_denoise_strength": getattr(self.input_info, "i2i_denoise_strength", None), + "i2i_denoise_strength": self.input_info.i2i_denoise_strength, } ) return inputs diff --git a/lightx2v/models/runners/hunyuan3d/hunyuan3d_shape_runner.py b/lightx2v/models/runners/hunyuan3d/hunyuan3d_shape_runner.py index 4268d1b0e..84922a898 100644 --- a/lightx2v/models/runners/hunyuan3d/hunyuan3d_shape_runner.py +++ b/lightx2v/models/runners/hunyuan3d/hunyuan3d_shape_runner.py @@ -13,6 +13,7 @@ from lightx2v.models.networks.hunyuan3d.utils import torchvision_fix from lightx2v.models.networks.hunyuan3d.utils.checkpoint import load_checkpoint_dict, resolve_ckpt_paths, resolve_model_dir from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS from lightx2v.models.schedulers.hunyuan3d.scheduler import Hunyuan3DShapeScheduler from lightx2v.models.video_encoders.hf.hunyuan3d.decoder import Hunyuan3DShapeVAEDecoder from lightx2v.server.metrics import monitor_cli @@ -28,6 +29,10 @@ class Hunyuan3DShapeRunner(DefaultRunner): """Image-to-3D-mesh runner for Hunyuan3D-2.1 shape pipeline.""" + supported_request_fields_by_task = { + "i23d": (COMMON_REQUEST_FIELDS - {"return_result_tensor"}) | {"image_path"}, + } + def __init__(self, config): super().__init__(config) self._ckpt = None @@ -166,7 +171,7 @@ def _run_infer_step(self, step_index: int, dit_inputs: dict) -> None: def run_main(self): latent_shape = (self.inputs["image_tensor"].shape[0], *self.vae_decoder.vae.latent_shape) self.scheduler.prepare( - seed=getattr(self.input_info, "seed", None), + seed=self.input_info.seed, batch_size=latent_shape[0], latent_shape=latent_shape, ) diff --git a/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py b/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py index 3e3de6e75..9c06c33a8 100644 --- a/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py +++ b/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py @@ -22,7 +22,9 @@ DistributedAutotuneContext, FlashInferAutotuneController, ) +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS from lightx2v.models.schedulers.hunyuan_image3.scheduler import HunyuanImage3Scheduler +from lightx2v.utils.input_info import TI2IInputInfo from lightx2v.utils.registry_factory import RUNNER_REGISTER @@ -36,6 +38,21 @@ class HunyuanImage3TextGenerationPlan: @RUNNER_REGISTER("hunyuan_image3") class HunyuanImage3Runner(DefaultRunner): model_cpu_offload_seq = "transformer" + TEXT_REQUEST_FIELDS = COMMON_REQUEST_FIELDS | {"bot_task", "max_new_tokens", "prompt", "stream_callback", "system_prompt", "text_do_sample", "text_temperature", "text_top_k", "text_top_p"} + input_info_cls_by_task = {"i2i": TI2IInputInfo} + supported_request_fields_by_task = { + "t2t": TEXT_REQUEST_FIELDS, + "ti2t": TEXT_REQUEST_FIELDS | {"image_path", "infer_align_image_size"}, + "t2i": COMMON_REQUEST_FIELDS | {"prompt", "target_shape"}, + "ti2i": COMMON_REQUEST_FIELDS | {"image_path", "infer_align_image_size", "prompt", "target_shape"}, + "i2i": COMMON_REQUEST_FIELDS | {"image_path", "infer_align_image_size", "prompt", "target_shape"}, + } + + def get_supported_request_fields(self, task): + supported_request_fields = super().get_supported_request_fields(task) + if self.config.get("image_size"): + supported_request_fields -= {"target_shape"} + return supported_request_fields def __init__(self, config): super().__init__(config) @@ -984,7 +1001,7 @@ def _generate_text_tokens( ): generation_options = generation_options or {} device = input_ids.device - seed = int(generation_options.get("text_seed", self.config.get("text_seed", self.config.get("seed", 42)))) + seed = int(generation_options.get("text_seed", self.config.get("text_seed", self.input_info.seed))) generator = torch.Generator(device=device).manual_seed(seed) max_new_tokens = int(generation_options.get("max_new_tokens", self.config.get("max_new_tokens", getattr(self.hunyuan_generation_config, "max_new_tokens", 2048)))) transition_map = {stop_id: list(append_ids) for stop_id, append_ids in plan.stage_transitions} @@ -1589,9 +1606,7 @@ def generate_t2t(self, input_info): @torch.no_grad() def generate_ti2t(self, input_info): self._ensure_pipeline_modules() - seed = getattr(input_info, "seed", None) - if seed is None: - seed = self.config.get("seed", 42) + seed = input_info.seed image_paths = self._split_image_paths(getattr(input_info, "image_path", None) or self.config.get("image_path")) infer_align_image_size = getattr(input_info, "infer_align_image_size", None) if infer_align_image_size is None: @@ -1606,7 +1621,7 @@ def generate_t2i(self, input_info): self._ensure_pipeline_modules() prompt = getattr(input_info, "prompt", "") image_size = self._resolve_image_size(input_info) - seed = getattr(input_info, "seed", None) or self.config.get("seed", 42) + seed = input_info.seed cot_text = self._generate_cot_text(prompt, image_size) prepared_inputs = self._prepare_text_to_image_inputs( prompt, @@ -1623,7 +1638,7 @@ def generate_t2i(self, input_info): def generate_ti2i(self, input_info): self._ensure_pipeline_modules() prompt = getattr(input_info, "prompt", "") - seed = getattr(input_info, "seed", None) or self.config.get("seed", 42) + seed = input_info.seed image_paths = self._split_image_paths(getattr(input_info, "image_path", None) or self.config.get("image_path")) infer_align_image_size = getattr(input_info, "infer_align_image_size", None) if infer_align_image_size is None: diff --git a/lightx2v/models/runners/hunyuan_video/hunyuan_video_15_runner.py b/lightx2v/models/runners/hunyuan_video/hunyuan_video_15_runner.py index e7116dfb7..3833547c2 100755 --- a/lightx2v/models/runners/hunyuan_video/hunyuan_video_15_runner.py +++ b/lightx2v/models/runners/hunyuan_video/hunyuan_video_15_runner.py @@ -13,6 +13,7 @@ from lightx2v.models.input_encoders.hf.hunyuan15.siglip.model import SiglipVisionEncoder from lightx2v.models.networks.hunyuan_video.model import HunyuanVideo15Model from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, PROMPT_FIELDS from lightx2v.models.schedulers.hunyuan_video.feature_caching.scheduler import HunyuanVideo15SchedulerCaching from lightx2v.models.schedulers.hunyuan_video.scheduler import HunyuanVideo15SRScheduler, HunyuanVideo15Scheduler from lightx2v.models.schedulers.hunyuan_video.step_distill.scheduler import HunyuanVideo15StepDistillScheduler @@ -29,6 +30,17 @@ @RUNNER_REGISTER("hunyuan_video_1.5") class HunyuanVideo15Runner(DefaultRunner): + supported_request_fields_by_task = { + "t2v": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | {"target_video_length"}, + "i2v": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | {"image_path", "target_video_length"}, + } + + def get_supported_request_fields(self, task): + supported_request_fields = super().get_supported_request_fields(task) + if self.config.get("video_super_resolution", {}).get("enable_cfg", False): + supported_request_fields |= {"negative_prompt"} + return supported_request_fields + def __init__(self, config): config["is_sr_running"] = False @@ -44,6 +56,7 @@ def __init__(self, config): self.config_sr["is_sr_running"] = False self.config_sr["sample_shift"] = config["video_super_resolution"]["flow_shift"] # for SR model self.config_sr["sample_guide_scale"] = config["video_super_resolution"]["guidance_scale"] # for SR model + self.config_sr["enable_cfg"] = config["video_super_resolution"].get("enable_cfg", config["enable_cfg"]) self.config_sr["infer_steps"] = config["video_super_resolution"]["num_inference_steps"] super().__init__(config) @@ -138,7 +151,7 @@ def get_latent_shape_with_target_hw(self, origin_size=None): target_height, target_width = self.get_closest_resolution_given_original_size((int(width), int(height)), target_size) latent_shape = [ self.config.get("in_channels", 32), - (self.config["target_video_length"] - 1) // self.config["vae_stride"][0] + 1, + (self.get_target_video_length() - 1) // self.config["vae_stride"][0] + 1, target_height // self.config["vae_stride"][1], target_width // self.config["vae_stride"][2], ] @@ -230,7 +243,7 @@ def get_sr_latent_shape_with_target_hw(self): target_width, target_height = hr_bucket_map((lr_video_width, lr_video_height)) latent_shape = [ self.config_sr.get("in_channels", 32), - (self.config_sr["target_video_length"] - 1) // self.config_sr["vae_stride"][0] + 1, + (self.get_target_video_length() - 1) // self.config_sr["vae_stride"][0] + 1, target_height // self.config_sr["vae_stride"][1], target_width // self.config_sr["vae_stride"][2], ] @@ -283,12 +296,13 @@ def get_closest_ratio(self, height: float, width: float, ratios: list, buckets: return closest_size, closest_ratio def run_text_encoder(self, input_info): + config = self.config_sr if self.sr_version and self.config_sr["is_sr_running"] else self.config prompt = input_info.prompt neg_prompt = input_info.negative_prompt # run qwen25vl - if self.config.get("enable_cfg", False) and self.config["cfg_parallel"]: - cfg_p_group = self.config["device_mesh"].get_group(mesh_dim="cfg_p") + if config.get("enable_cfg", False) and config["cfg_parallel"]: + cfg_p_group = config["device_mesh"].get_group(mesh_dim="cfg_p") cfg_p_rank = dist.get_rank(cfg_p_group) if cfg_p_rank == 0: context = self.text_encoders[0].infer([prompt]) @@ -298,7 +312,7 @@ def run_text_encoder(self, input_info): text_encoder_output = {"context_null": context_null} else: context = self.text_encoders[0].infer([prompt]) - context_null = self.text_encoders[0].infer([neg_prompt]) if self.config.get("enable_cfg", False) else None + context_null = self.text_encoders[0].infer([neg_prompt]) if config.get("enable_cfg", False) else None text_encoder_output = { "context": context, "context_null": context_null, diff --git a/lightx2v/models/runners/lingbot_video/lingbot_video_runner.py b/lightx2v/models/runners/lingbot_video/lingbot_video_runner.py index a6472c7a0..91c459d50 100644 --- a/lightx2v/models/runners/lingbot_video/lingbot_video_runner.py +++ b/lightx2v/models/runners/lingbot_video/lingbot_video_runner.py @@ -12,6 +12,7 @@ from lightx2v.models.networks.lingbot_video.model import LingBotVideoTransformerModel from lightx2v.models.networks.lora_adapter import LoraAdapter from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, PROMPT_FIELDS, VIDEO_REQUEST_FIELDS from lightx2v.models.schedulers.lingbot_video.scheduler import LingBotVideoScheduler from lightx2v.models.video_encoders.hf.lingbot_video.vae import LingBotVideoWanVAE from lightx2v.utils.input_info import I2VInputInfo, T2IInputInfo, T2VInputInfo @@ -76,12 +77,16 @@ def smart_resize(height, width, factor, min_pixels=None, max_pixels=None): @RUNNER_REGISTER("lingbot_video") class LingBotVideoRunner(DefaultRunner): + supported_request_fields_by_task = { + "t2i": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | {"target_shape"}, + "t2v": VIDEO_REQUEST_FIELDS, + "i2v": VIDEO_REQUEST_FIELDS | {"image_path"}, + } + model_cpu_offload_seq = "text_encoder->transformer->vae" _WARMUP_RESOLUTIONS = ((480, 480), (320, 832)) def __init__(self, config): - if config.get("task") not in {"t2i", "t2v", "i2v"}: - raise NotImplementedError("LingBot-Video LightX2V backend currently supports t2i, t2v, and i2v.") if config.get("lazy_load", False) or config.get("unload_modules", False): raise NotImplementedError("LingBot-Video lazy_load/unload_modules are not implemented yet.") super().__init__(config) @@ -132,7 +137,7 @@ def _prepare_warmup_inputs(self, height, width, text_encoder_output=None): "text_encoder_output": text_encoder_output, "image_encoder_output": None, } - self.set_target_shape() + self.set_latent_shape() return text_encoder_output def clear_warmup_state(self): @@ -186,16 +191,11 @@ def run_text_encoder(self, prompt, neg_prompt=None, images=None): prompt_output = self.text_encoders[0].infer(prompt, images=images) text_encoder_output["prompt_embeds"] = prompt_output["prompt_embeds"] text_encoder_output["prompt_mask"] = prompt_output["prompt_mask"] - if hasattr(self.input_info, "txt_seq_lens"): - self.input_info.txt_seq_lens = [prompt_output["prompt_embeds"].shape[1]] - if self.config.get("enable_cfg", True): neg_prompt = "" if neg_prompt is None else neg_prompt negative_output = self.text_encoders[0].infer(neg_prompt, images=images) text_encoder_output["negative_prompt_embeds"] = negative_output["prompt_embeds"] text_encoder_output["negative_prompt_mask"] = negative_output["prompt_mask"] - if hasattr(self.input_info, "txt_seq_lens"): - self.input_info.txt_seq_lens.append(negative_output["prompt_embeds"].shape[1]) return text_encoder_output @ProfilingContext4DebugL2("Run Encoders") @@ -246,7 +246,7 @@ def _ensure_scheduler_generator(self): @ProfilingContext4DebugL2("Run Encoders") def _run_input_encoder_local_i2v(self, image=None): - height, width = self._resolve_output_size() + height, width = self.get_target_size() if image is None: image_path = self.input_info.image_path.split(",")[0] if not image_path: @@ -264,15 +264,8 @@ def _run_input_encoder_local_i2v(self, image=None): "image_encoder_output": {"cond_latent": cond_latent}, } - def _resolve_output_size(self): - if len(self.input_info.target_shape) == 2: - height, width = int(self.input_info.target_shape[0]), int(self.input_info.target_shape[1]) - else: - height, width = int(self.config["target_height"]), int(self.config["target_width"]) - return height, width - - def set_target_shape(self): - height, width = self._resolve_output_size() + def set_latent_shape(self): + height, width = self.get_target_size() if height <= 0 or width <= 0: raise ValueError(f"LingBot-Video target shape must be positive, got {height}x{width}.") if height % 16 != 0 or width % 16 != 0: @@ -281,7 +274,7 @@ def set_target_shape(self): if self.config["task"] == "t2i": frames = 1 else: - frames = int(self.config["target_video_length"]) + frames = self.get_target_video_length() if frames != 1 and (frames - 1) % int(self.config.get("vae_scale_factor_temporal", 4)) != 0: raise ValueError(f"LingBot-Video target_video_length must be 1 or 4n+1, got {frames}.") latent_t = (frames - 1) // int(self.config.get("vae_scale_factor_temporal", 4)) + 1 @@ -289,13 +282,8 @@ def set_target_shape(self): latent_w = width // int(self.config.get("vae_scale_factor_spatial", 8)) latent_shape = (1, int(self.config.get("in_channels", 16)), latent_t, latent_h, latent_w) - self.input_info.auto_height = height - self.input_info.auto_width = width - self.input_info.target_shape = latent_shape + self.input_info.target_shape = [height, width] self.input_info.latent_shape = latent_shape - patch_h, patch_w = self.config.get("patch_size", [1, 2, 2])[1:] - if hasattr(self.input_info, "image_shapes"): - self.input_info.image_shapes = [[(latent_t, latent_h // patch_h, latent_w // patch_w)]] logger.info(f"LingBot-Video target shape: frames={frames}, image={height}x{width}, latent={latent_shape}") def _apply_condition_latent(self): @@ -363,7 +351,7 @@ def _finalize_pipeline_outputs(self, outputs, latents=None, generator=None): def run_pipeline(self, input_info): self.input_info = input_info self.inputs = self.run_input_encoder() - self.set_target_shape() + self.set_latent_shape() logger.info(f"input_info: {self.input_info}") latents, generator = self.run_dit() outputs = self.run_vae_decoder(latents) diff --git a/lightx2v/models/runners/longcat_image/longcat_image_runner.py b/lightx2v/models/runners/longcat_image/longcat_image_runner.py index fe218d279..56c020de0 100755 --- a/lightx2v/models/runners/longcat_image/longcat_image_runner.py +++ b/lightx2v/models/runners/longcat_image/longcat_image_runner.py @@ -8,6 +8,7 @@ from lightx2v.models.input_encoders.hf.longcat.longcat_text_encoder import LongCatImageTextEncoder from lightx2v.models.networks.longcat_image.model import LongCatImageTransformerModel from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, IMAGE_REQUEST_FIELDS, PROMPT_FIELDS from lightx2v.models.schedulers.longcat_image.scheduler import LongCatImageScheduler from lightx2v.models.video_encoders.hf.longcat_image.vae import LongCatImageVAE from lightx2v.server.metrics import monitor_cli @@ -26,16 +27,6 @@ torch_device_module = getattr(torch, AI_DEVICE) -def calculate_dimensions(target_area, ratio): - width = math.sqrt(target_area * ratio) - height = width / ratio - - width = round(width / 32) * 32 - height = round(height / 32) * 32 - - return width, height, None - - def calculate_target_dimensions_from_image(image_size, target_area=1024 * 1024, multiple_of=16): """Calculate target dimensions from image size while preserving aspect ratio. @@ -62,10 +53,10 @@ def calculate_target_dimensions_from_image(image_size, target_area=1024 * 1024, class LongCatImageRunner(DefaultRunner): model_cpu_offload_seq = "text_encoder->transformer->vae" _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__(self, config): - super().__init__(config) - self.resolution = self.config.get("resolution", 1024) + supported_request_fields_by_task = { + "t2i": IMAGE_REQUEST_FIELDS, + "i2i": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | {"image_path"}, + } def load_transformer(self): model = LongCatImageTransformerModel(os.path.join(self.config["model_path"], "transformer"), self.config, self.init_device) @@ -190,9 +181,7 @@ def _preprocess_image(self, image): image = image_processor.resize(image, height, width) image_tensor = image_processor.preprocess(image, height, width) - # Store dimensions for later use - self.input_info.auto_width = width - self.input_info.auto_height = height + self.input_info.target_shape = [height, width] return image_tensor.to(AI_DEVICE, dtype=GET_DTYPE()) @@ -334,26 +323,15 @@ def get_custom_shape(self): width, height = as_maps[self.config.get("aspect_ratio", "16:9")] return (width, height) - def set_target_shape(self): - # For I2I task, use dimensions from _preprocess_image() if already set + def set_latent_shape(self): task = self.config.get("task", "t2i") - if task == "i2i" and hasattr(self.input_info, "auto_width") and self.input_info.auto_width: - width = self.input_info.auto_width - height = self.input_info.auto_height + if task == "i2i": + height, width = self.input_info.target_shape else: - custom_shape = self.get_custom_shape() - if custom_shape is not None: - width, height = custom_shape - else: - calculated_width, calculated_height, _ = calculate_dimensions(self.resolution * self.resolution, 16 / 9) - multiple_of = self.config.get("vae_scale_factor", 8) * 2 - width = calculated_width // multiple_of * multiple_of - height = calculated_height // multiple_of * multiple_of - - self.input_info.auto_width = width - self.input_info.auto_height = height + width, height = self.get_custom_shape() logger.info(f"LongCat Image Runner set target shape: {width}x{height}") + self.input_info.target_shape = [height, width] # VAE applies 8x compression on images but we must also account for packing which requires # latent height and width to be divisible by 2. @@ -361,14 +339,7 @@ def set_target_shape(self): height = 2 * (int(height) // (vae_scale_factor * 2)) width = 2 * (int(width) // (vae_scale_factor * 2)) num_channels_latents = 16 # LongCat uses 16 latent channels - self.input_info.target_shape = (1, num_channels_latents, height, width) - - def set_img_shapes(self): - width, height = self.input_info.auto_width, self.input_info.auto_height - vae_scale_factor = self.config.get("vae_scale_factor", 8) - # For T2I task - image_shapes = [(1, height // vae_scale_factor // 2, width // vae_scale_factor // 2)] * 1 - self.input_info.image_shapes = image_shapes + self.input_info.latent_shape = (1, num_channels_latents, height, width) def init_scheduler(self): self.scheduler = LongCatImageScheduler(self.config) @@ -391,14 +362,13 @@ def run_pipeline(self, input_info): self.input_info = input_info self.inputs = self.run_input_encoder() - self.set_target_shape() - self.set_img_shapes() + self.set_latent_shape() logger.info(f"input_info: {self.input_info}") latents, generator = self.run_dit() images = self.run_vae_decoder(latents) self.end_run() - if not input_info.return_result_tensor and is_main_process(): + if not input_info.return_result_tensor and input_info.save_result_path is not None and is_main_process(): image = images[0] image.save(input_info.save_result_path) logger.info(f"Image saved: {input_info.save_result_path}") diff --git a/lightx2v/models/runners/ltx2/ltx25_runner.py b/lightx2v/models/runners/ltx2/ltx25_runner.py index a274dfa68..10a7d3e2f 100644 --- a/lightx2v/models/runners/ltx2/ltx25_runner.py +++ b/lightx2v/models/runners/ltx2/ltx25_runner.py @@ -10,6 +10,7 @@ from lightx2v.models.input_encoders.hf.ltx2.model import LTX25TextEncoder from lightx2v.models.networks.ltx2.ltx25_model import LTX25Model from lightx2v.models.runners.ltx2.ltx2_runner import LTX2Runner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, VIDEO_OUTPUT_FIELDS from lightx2v.models.schedulers.ltx2.ltx25_scheduler import LTX25Scheduler from lightx2v.models.video_encoders.hf.ltx2.model import LTX25AudioVAE, LTX25VideoVAE from lightx2v.utils.envs import GET_DTYPE @@ -36,11 +37,12 @@ class LTX25Runner(LTX2Runner): text_encoder_root_key = "text_encoder_original_ckpt" video_vae_checkpoint_key = "video_vae_original_ckpt" audio_vae_checkpoint_key = "audio_vae_original_ckpt" + supported_request_fields_by_task = { + "t2av": COMMON_REQUEST_FIELDS | VIDEO_OUTPUT_FIELDS | {"prompt"}, + "i2av": COMMON_REQUEST_FIELDS | VIDEO_OUTPUT_FIELDS | {"image_frame_idx", "image_path", "image_strength", "prompt"}, + } def __init__(self, config): - task = config.get("task") - if task not in ("t2av", "i2av"): - raise NotImplementedError(f"LTX-2.5 currently supports t2av and i2av, got {task!r}") if config.get("enable_cfg", False) or float(config.get("sample_guide_scale", 1.0)) != 1.0: raise ValueError("The LTX-2.5 distilled pipeline requires CFG=1 (enable_cfg=false)") if config.get("disagg_mode"): @@ -123,7 +125,7 @@ def _resolve_target_video_length(self, text_encoder_output) -> int: logger.info(f"LTX-2.5 target video length: {num_frames} frames ({source})") return num_frames - def _prepare_stage1_target_shape(self) -> None: + def prepare_stage1_target_shape(self) -> None: """Interpret request/config dimensions as final two-stage dimensions.""" if self.input_info.target_shape: if len(self.input_info.target_shape) != 2: @@ -134,7 +136,8 @@ def _prepare_stage1_target_shape(self) -> None: final_width = int(self.config["target_width"]) if final_height % 64 != 0 or final_width % 64 != 0: raise ValueError(f"LTX-2.5 distilled two-stage output height and width must be divisible by 64, got {final_height}x{final_width}") - self.input_info.target_shape = [final_height // 2, final_width // 2] + self.input_info.target_shape = [final_height, final_width] + super().prepare_stage1_target_shape() def _validate_sequence_parallel_shape(self, num_frames: int, guiding_keyframes: int = 0) -> None: """Reject SP layouts that would introduce unmasked video tokens. @@ -171,7 +174,7 @@ def _run_input_encoder_local_t2av(self): self._clear_ltx2_reference_video_state() self.video_denoise_mask = None self.initial_video_latent = None - self._prepare_stage1_target_shape() + self.prepare_stage1_target_shape() text_encoder_output = self.run_text_encoder(self.input_info) num_frames = self._resolve_target_video_length(text_encoder_output) self._validate_sequence_parallel_shape(num_frames) @@ -183,7 +186,7 @@ def _run_input_encoder_local_i2av(self): self._clear_ltx2_reference_audio_state() self._clear_ltx2_reference_video_state() self._normalize_i2av_input_fields() - self._prepare_stage1_target_shape() + self.prepare_stage1_target_shape() text_encoder_output = self.run_text_encoder(self.input_info) num_frames = self._resolve_target_video_length(text_encoder_output) image_paths = [path.strip() for path in (self.input_info.image_path or "").split(",") if path.strip()] diff --git a/lightx2v/models/runners/ltx2/ltx2_runner.py b/lightx2v/models/runners/ltx2/ltx2_runner.py index 75af1fec0..8bb43f642 100755 --- a/lightx2v/models/runners/ltx2/ltx2_runner.py +++ b/lightx2v/models/runners/ltx2/ltx2_runner.py @@ -10,6 +10,7 @@ from lightx2v.models.networks.lora_adapter import LoraAdapter from lightx2v.models.networks.ltx2.model import LTX2ARModel, LTX2Model from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import VIDEO_REQUEST_FIELDS from lightx2v.models.schedulers.ltx2.scheduler import LTX2ARScheduler, LTX2Scheduler, LatentState from lightx2v.models.video_encoders.hf.ltx2.audio_vae.audio_vae import encode_audio from lightx2v.models.video_encoders.hf.ltx2.audio_vae.ops import Audio @@ -89,6 +90,22 @@ def _ltx2_resize_video_denoise_mask_for_stage2(mask: torch.Tensor, target_h: int @RUNNER_REGISTER("ltx2") class LTX2Runner(DefaultRunner): + supported_request_fields_by_task = { + "t2av": VIDEO_REQUEST_FIELDS, + "i2av": VIDEO_REQUEST_FIELDS | {"image_frame_idx", "image_path", "image_strength"}, + "ltx2_s2v": VIDEO_REQUEST_FIELDS | {"audio_path", "image_frame_idx", "image_path", "image_strength"}, + "v2av": VIDEO_REQUEST_FIELDS + | { + "image_frame_idx", + "image_path", + "image_strength", + "mux_audio_video_path", + "reference_video_frame_cap", + "reference_video_strength", + "video_path", + }, + } + _WARMUP_RESOLUTIONS = ((480, 480), (512, 768)) _UPSAMPLER_WARMUP_RESOLUTIONS = ((480, 480), (1024, 1536)) transformer_model_class = LTX2Model @@ -100,6 +117,12 @@ class LTX2Runner(DefaultRunner): video_vae_checkpoint_key = None audio_vae_checkpoint_key = None + def create_input_info(self, request_data): + input_info = super().create_input_info(request_data) + if input_info.task == "v2av" and "target_shape" not in request_data and "target_shape" not in self.config: + input_info.target_shape = [] + return input_info + def __init__(self, config): super().__init__(config) @@ -377,6 +400,12 @@ def get_latent_shape_with_target_hw(self): return video_latent_shape, audio_latent_shape + def prepare_stage1_target_shape(self): + """Convert the requested final size to the first-stage size.""" + if self.config.get("use_upsampler", False): + height, width = self.input_info.target_shape + self.input_info.target_shape = [height // 2, width // 2] + def _clear_ltx2_reference_audio_state(self) -> None: """Avoid leaking ltx2_s2v audio conditioning into t2av/i2av runs on a reused runner.""" self.initial_audio_latent = None @@ -389,6 +418,7 @@ def _run_input_encoder_local_t2av(self): self._clear_ltx2_reference_video_state() self.video_denoise_mask = None self.initial_video_latent = None + self.prepare_stage1_target_shape() self.input_info.video_latent_shape, self.input_info.audio_latent_shape = self.get_latent_shape_with_target_hw() # Important: set latent_shape in input_info text_encoder_output = self.run_text_encoder(self.input_info) self.maybe_empty_cache() @@ -419,6 +449,7 @@ def _run_input_encoder_local_i2av(self): self._clear_ltx2_reference_audio_state() self._clear_ltx2_reference_video_state() self._normalize_i2av_input_fields() + self.prepare_stage1_target_shape() self.input_info.video_latent_shape, self.input_info.audio_latent_shape = self.get_latent_shape_with_target_hw() text_encoder_output = self.run_text_encoder(self.input_info) self.video_denoise_mask, self.initial_video_latent = self.run_vae_encoder() @@ -467,7 +498,7 @@ def _probe_video_hw(path: str) -> tuple[int, int] | None: return None def _override_target_hw_from_ref_video(self) -> None: - """v2av: set ``input_info.target_shape`` from probed ``video_path`` (control mp4). + """v2av: set ``input_info.target_shape`` from the source control video. Skip if ``target_shape`` already set. Base H/W = final//2 when upsampler else final; snap to VAE grid (spatial 32) vs ``ref_downscale_factor``. Probe/config miss → no-op. @@ -475,7 +506,7 @@ def _override_target_hw_from_ref_video(self) -> None: if self.input_info.target_shape: return - ref_path = (getattr(self.input_info, "video_path", None) or "").strip() + ref_path = (self.input_info.video_path or "").strip() hw = self._probe_video_hw(ref_path) if hw is None: return @@ -507,7 +538,7 @@ def _snap_nearest(x: int, d: int) -> int: f"base-gen {base_w}x{base_h}, base_div={base_div}, " f"ref_downscale_factor={ref_factor}, use_upsampler={use_upsampler})." ) - self.input_info.target_shape = [base_h, base_w] + self.input_info.target_shape = [eff_final_h, eff_final_w] @ProfilingContext4DebugL2("Run Encoders") def _run_input_encoder_local_v2av(self): @@ -518,22 +549,17 @@ def _run_input_encoder_local_v2av(self): self._normalize_i2av_input_fields() self._override_target_hw_from_ref_video() if not self.input_info.target_shape: - if self.config.get("use_upsampler", False): - self.input_info.target_shape = [ - self.config["target_height"] // 2, - self.config["target_width"] // 2, - ] - else: - self.input_info.target_shape = [ - self.config["target_height"], - self.config["target_width"], - ] + self.input_info.target_shape = [ + self.config["target_height"], + self.config["target_width"], + ] + self.prepare_stage1_target_shape() # Reference/control video → pixel tensor, then align temporal length with # the clip (official-style: decode up to ``num_frames`` cap, actual length # follows the shorter of cap vs. on-disk frames). Only then derive # ``target_video_length`` / latent shapes so audio and denoising match. - ref_path = (getattr(self.input_info, "video_path", None) or "").strip() + ref_path = (self.input_info.video_path or "").strip() if not ref_path: raise ValueError("v2av requires a non-empty video_path (pre-processed control / reference video).") @@ -623,6 +649,7 @@ def _run_input_encoder_local_ltx2_s2v(self): """Reference audio (frozen in latent) + optional reference images; mux original waveform when saving.""" self._clear_ltx2_reference_video_state() self._normalize_i2av_input_fields() + self.prepare_stage1_target_shape() self.input_info.video_latent_shape, self.input_info.audio_latent_shape = self.get_latent_shape_with_target_hw() ap = (getattr(self.input_info, "audio_path", None) or "").strip() @@ -1186,6 +1213,8 @@ def run_segment(self, segment_idx=0, stage_name=None, cleanup_inputs=None): class LTX2ARRunner(LTX2Runner): """Chunkwise autoregressive LTX2.3 runner for teacher-forcing checkpoints.""" + supported_request_fields_by_task = {"t2av": LTX2Runner.supported_request_fields_by_task["t2av"]} + def init_scheduler(self): self.scheduler = LTX2ARScheduler(self.config) @@ -1210,8 +1239,6 @@ def init_run(self): self._prepare_ar_states() def _validate_ar_config(self): - if self.config.get("task") != "t2av": - raise NotImplementedError("ltx2_ar currently supports task=t2av only.") if self.config.get("use_upsampler", False): raise NotImplementedError("ltx2_ar does not support the latent upsampler.") chunk = int(self.config.get("ar_config", {}).get("num_frame_per_chunk", 0)) diff --git a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py index cbc1b33c3..81a16a63b 100644 --- a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py +++ b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py @@ -38,12 +38,13 @@ trim_reference_num_frames, ) from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, VIDEO_OUTPUT_FIELDS from lightx2v.models.schedulers.minimax_h3 import MiniMaxH3Scheduler from lightx2v.models.video_encoders.hf.ltx2.audio_vae.ops import Audio from lightx2v.models.video_encoders.hf.minimax_h3 import MiniMaxH3VideoVAE from lightx2v.server.metrics import monitor_cli from lightx2v.utils.envs import DTYPE_MAP, GET_RECORDER_MODE -from lightx2v.utils.input_info import FL2AVInputInfo, I2AVInputInfo, L2AVInputInfo, Ref2AVInputInfo, T2AVInputInfo +from lightx2v.utils.input_info import INPUT_INFO_TYPES from lightx2v.utils.ltx2_media_io import encode_video from lightx2v.utils.profiler import ProfilingContext4DebugL1, ProfilingContext4DebugL2 from lightx2v.utils.registry_factory import RUNNER_REGISTER @@ -95,15 +96,25 @@ class MiniMaxH3Runner(DefaultRunner): (544, 960, 124), ) _WARMUP_STEP_COUNT = 2 - _WARMUP_TASKS = ("t2av", "fl2av", "i2av", "l2av", "ref2av") + supported_request_fields_by_task = { + "t2av": COMMON_REQUEST_FIELDS | VIDEO_OUTPUT_FIELDS | {"prompt"}, + "i2av": COMMON_REQUEST_FIELDS | VIDEO_OUTPUT_FIELDS | {"image_path", "prompt"}, + "l2av": COMMON_REQUEST_FIELDS | VIDEO_OUTPUT_FIELDS | {"last_frame_path", "prompt"}, + "fl2av": COMMON_REQUEST_FIELDS | VIDEO_OUTPUT_FIELDS | {"image_path", "last_frame_path", "prompt"}, + "ref2av": COMMON_REQUEST_FIELDS | VIDEO_OUTPUT_FIELDS | {"audio_path", "image_path", "prompt", "video_path"}, + } def __init__(self, config): - if config.get("task") not in {"t2av", "i2av", "l2av", "fl2av", "ref2av"}: - raise ValueError("MiniMax-H3 supports t2av/i2av/l2av/fl2av/ref2av") - self.loaded_transformer_partition = "transformer_ref" if config["task"] == "ref2av" else "transformer" if config.get("lazy_load", False) or config.get("unload_modules", False): raise NotImplementedError("MiniMax-H3 does not support lazy_load or unload_modules yet; use the released sharded checkpoint with model or block CPU offload.") super().__init__(config) + self.loaded_transformer_partition = "transformer_ref" if config["task"] == "ref2av" else "transformer" + + def get_supported_tasks(self): + """Return tasks supported by the loaded transformer weights.""" + if self.config["task"] == "ref2av": + return ("ref2av",) + return ("t2av", "i2av", "l2av", "fl2av") def init_modules(self): super().init_modules() @@ -116,8 +127,6 @@ def init_modules(self): @ProfilingContext4DebugL1("Warmup") def run_warmup(self): task = self.config["task"] - if task not in self._WARMUP_TASKS: - raise NotImplementedError(f"MiniMax-H3 warmup does not support task: {task}") if task == "ref2av" and self.config.get("vae_use_compile", False): height, width, _ = self._WARMUP_SHAPES[0] @@ -159,26 +168,21 @@ def run_warmup(self): def _prepare_warmup_inputs(self, height, width, num_frames): task = self.config["task"] - common = { - "seed": 0, - "prompt": "A sunrise over distant mountains reflected across a calm lake beneath drifting clouds." + self.input_info = INPUT_INFO_TYPES[task]( + task=task, + seed=0, + prompt="A sunrise over distant mountains reflected across a calm lake beneath drifting clouds." if (height, width, num_frames) == self._WARMUP_SHAPES[0] else "A cinematic fox walking through a snowy forest.", - "target_shape": [height, width], - "target_video_length": num_frames, - "return_result_tensor": True, - } + target_shape=[height, width], + target_video_length=num_frames, + return_result_tensor=True, + ) image = Image.new("RGB", (width, height), color=0) - if task == "t2av": - self.input_info = T2AVInputInfo(**common) - elif task == "i2av": - self.input_info = I2AVInputInfo(**common, image_path=image) - elif task == "l2av": - self.input_info = L2AVInputInfo(**common, last_frame_path=image) - elif task == "fl2av": - self.input_info = FL2AVInputInfo(**common, image_path=image, last_frame_path=image.copy()) - else: - self.input_info = Ref2AVInputInfo(**common, image_path=image) + if task in ("i2av", "fl2av", "ref2av"): + self.input_info.image_path = image + if task in ("l2av", "fl2av"): + self.input_info.last_frame_path = image.copy() if task == "fl2av" else image def clear_conditioning_state(self): self.condition_video_latents = [] @@ -339,21 +343,17 @@ def _load_rgb_image(value): return ImageOps.exif_transpose(image).convert("RGB") def _prepare_keyframes(self): - task = self.config["task"] - if task == "t2av": - if not isinstance(self.input_info, T2AVInputInfo): - raise TypeError(f"MiniMax-H3 t2av expects T2AVInputInfo, got {type(self.input_info).__name__}") - return [], () + task = self.input_info.task if task == "i2av": - if not isinstance(self.input_info, I2AVInputInfo) or not self.input_info.image_path: + if not self.input_info.image_path: raise ValueError("MiniMax-H3 i2av requires exactly one --image_path") values, anchors = [self.input_info.image_path], ("first",) elif task == "l2av": - if not isinstance(self.input_info, L2AVInputInfo) or not self.input_info.last_frame_path: + if not self.input_info.last_frame_path: raise ValueError("MiniMax-H3 l2av requires --last_frame_path") values, anchors = [self.input_info.last_frame_path], ("last",) elif task == "fl2av": - if not isinstance(self.input_info, FL2AVInputInfo) or not self.input_info.image_path or not self.input_info.last_frame_path: + if not self.input_info.image_path or not self.input_info.last_frame_path: raise ValueError("MiniMax-H3 fl2av requires --image_path and --last_frame_path") values, anchors = [self.input_info.image_path, self.input_info.last_frame_path], ("first", "last") else: @@ -386,8 +386,6 @@ def _split_reference_paths(value): return [value] def _prepare_references(self): - if not isinstance(self.input_info, Ref2AVInputInfo): - raise TypeError(f"MiniMax-H3 ref2av expects Ref2AVInputInfo, got {type(self.input_info).__name__}") entries = [] for kind, value in ( ("image", self.input_info.image_path), @@ -500,7 +498,7 @@ def _encode_references(self, references): @ProfilingContext4DebugL2("Run Input Encoder") def _run_input_encoder_local_h3(self): - task = self.config["task"] + task = self.input_info.task requested_partition = "transformer_ref" if task == "ref2av" else "transformer" if requested_partition != self.loaded_transformer_partition: raise ValueError( diff --git a/lightx2v/models/runners/motus/motus_runner.py b/lightx2v/models/runners/motus/motus_runner.py index 43c2b5e78..2b2af0bd1 100644 --- a/lightx2v/models/runners/motus/motus_runner.py +++ b/lightx2v/models/runners/motus/motus_runner.py @@ -8,12 +8,14 @@ from lightx2v.models.input_encoders.hf.wan.t5.model import T5EncoderModel from lightx2v.models.networks.motus.model import MotusModel +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS from lightx2v.models.runners.wan.wan_runner import Wan22DenseRunner from lightx2v.models.schedulers.motus.scheduler import MotusScheduler from lightx2v.models.video_encoders.hf.wan.vae_2_2 import Wan2_2_VAE from lightx2v.models.video_encoders.hf.wan.vae_tiny import Wan2_2_VAE_tiny from lightx2v.server.metrics import monitor_cli from lightx2v.utils.envs import * +from lightx2v.utils.input_info import MotusInputInfo from lightx2v.utils.profiler import * from lightx2v.utils.registry_factory import RUNNER_REGISTER from lightx2v.utils.utils import find_torch_model_path, save_to_video, wan_vae_to_comfy @@ -45,6 +47,11 @@ def _merge_wan_dense_defaults(config): @RUNNER_REGISTER("motus") class MotusRunner(Wan22DenseRunner): + input_info_cls_by_task = {"i2v": MotusInputInfo} + supported_request_fields_by_task = { + "i2v": (COMMON_REQUEST_FIELDS - {"return_result_tensor"}) | {"image_path", "prompt", "save_action_path", "state_path"}, + } + def __init__(self, config): _merge_wan_dense_defaults(config) super().__init__(config) diff --git a/lightx2v/models/runners/neopp/neopp_runner.py b/lightx2v/models/runners/neopp/neopp_runner.py index 56b8536d1..8c03aa632 100755 --- a/lightx2v/models/runners/neopp/neopp_runner.py +++ b/lightx2v/models/runners/neopp/neopp_runner.py @@ -9,8 +9,10 @@ from lightx2v.models.networks.lora_adapter import LoraAdapter from lightx2v.models.networks.neopp.model import NeoppModel from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS from lightx2v.models.schedulers.neopp.scheduler import NeoppMoeScheduler from lightx2v.utils.envs import * +from lightx2v.utils.input_info import NeoppInputInfo from lightx2v.utils.profiler import * from lightx2v.utils.registry_factory import RUNNER_REGISTER from lightx2v.utils.utils import * @@ -37,6 +39,23 @@ def build_neopp_model_with_lora(neopp_module, config, model_kwargs, lora_configs @RUNNER_REGISTER("neopp") class NeoppRunner(DefaultRunner): + input_info_cls_by_task = {"t2i": NeoppInputInfo, "i2i": NeoppInputInfo} + supported_request_fields_by_task = { + "t2i": (COMMON_REQUEST_FIELDS - {"return_result_tensor"}) | {"target_shape"}, + "i2i": (COMMON_REQUEST_FIELDS - {"return_result_tensor"}) | {"target_shape"}, + } + + def get_supported_tasks(self): + # LightLLM encodes both text and image conditioning into the injected KV. + return ("t2i", "i2i") + + def prepare_request(self, request_data): + input_info = super().prepare_request(request_data) + # LightLLM restores the session RNG before explicitly passing seed=None. + if "seed" in request_data and request_data["seed"] is None: + input_info.seed = None + return input_info + def __init__(self, config): super().__init__(config) self.patch_size = self.config.get("patch_size", 16) @@ -52,9 +71,6 @@ def __init__(self, config): self.enable_cfg = self.config.get("enable_cfg", True) self.past_key_values_cond = None self.past_key_values_uncond = None - self.past_key_values_text_uncond = None - self.past_key_values_img_uncond = None - self.num_input_images = config.get("num_input_images", 1) if self.config["seq_parallel"]: self.seq_p_group = self.config.get("device_mesh").get_group(mesh_dim="seq_p") else: @@ -162,64 +178,8 @@ def get_latent_shape_with_target_hw(self): latent_shape = [1, 3, target_height, target_width] return latent_shape - def multi_pipeline_run_debug(self, input_info): - self.input_info = input_info - if self.config.get("load_kv_cache_in_pipeline_for_debug", False): - self.load_kvcache( - "/data/nvme1/yongyang/FL/neo_9b_new/vlm_tensor/to_x2v_cond_kv_0_289.pt", - "/data/nvme1/yongyang/FL/neo_9b_new/vlm_tensor/to_x2v_uncond_kv_0_9.pt", - ) - self.set_inference_params( - index_offset_cond=289, - index_offset_uncond=9, - cfg_interval=(-1, 2), - cfg_scale=4.0, - cfg_norm="global", - timestep_shift=3.0, - ) - self.input_info.save_result_path = self.input_info.save_result_path.replace(".png", "_0.png") - - self.inputs = self.run_input_encoder() - gen_result = self.run_main() - self.clear_kvcache() - - self.input_info = self.input_info - if self.config.get("load_kv_cache_in_pipeline_for_debug", False): - self.load_kvcache( - "/data/nvme1/yongyang/FL/neo_9b_new/vlm_tensor/to_x2v_cond_kv_1_346.pt", - "/data/nvme1/yongyang/FL/neo_9b_new/vlm_tensor/to_x2v_uncond_kv_1_12.pt", - ) - self.set_inference_params( - index_offset_cond=346, - index_offset_uncond=12, - cfg_interval=(-1, 2), - cfg_scale=4.0, - cfg_norm="global", - timestep_shift=3.0, - ) - self.input_info.save_result_path = self.input_info.save_result_path.replace("_0.png", "_1.png") - - self.inputs = self.run_input_encoder() - gen_result = self.run_main() - self.clear_kvcache() - return gen_result - def run_pipeline(self, input_info): self.input_info = input_info - if self.config.get("load_kv_cache_in_pipeline_for_debug", False): - self.load_kvcache( - "/data/nvme1/yongyang/FL/neo_9b_new/vlm_tensor/to_x2v_cond_kv_0_289.pt", - "/data/nvme1/yongyang/FL/neo_9b_new/vlm_tensor/to_x2v_uncond_kv_0_9.pt", - ) - self.set_inference_params( - index_offset_cond=289, - index_offset_uncond=9, - cfg_interval=(-1, 2), - cfg_scale=4.0, - cfg_norm="global", - timestep_shift=3.0, - ) - try: self.inputs = self.run_input_encoder() return self.run_main() diff --git a/lightx2v/models/runners/qwen_image/qwen_image_runner.py b/lightx2v/models/runners/qwen_image/qwen_image_runner.py index a659b1564..7d3871dba 100755 --- a/lightx2v/models/runners/qwen_image/qwen_image_runner.py +++ b/lightx2v/models/runners/qwen_image/qwen_image_runner.py @@ -11,6 +11,7 @@ from lightx2v.models.networks.lora_adapter import LoraAdapter from lightx2v.models.networks.qwen_image.model import QwenImageTransformerModel from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import IMAGE_REQUEST_FIELDS from lightx2v.models.schedulers.qwen_image.scheduler import QwenImageScheduler from lightx2v.models.video_encoders.hf.qwen_image.vae import AutoencoderKLQwenImageVAE from lightx2v.server.metrics import monitor_cli @@ -59,6 +60,16 @@ class QwenImageRunner(DisaggMixin, DefaultRunner): _callback_tensor_inputs = ["latents", "prompt_embeds"] _WARMUP_RESOLUTIONS = ((480, 480), (832, 1248)) _WARMUP_TASKS = ("t2i", "i2i") + supported_request_fields_by_task = { + "t2i": IMAGE_REQUEST_FIELDS, + "i2i": IMAGE_REQUEST_FIELDS | {"i2i_denoise_strength", "image_path"}, + } + + def get_supported_request_fields(self, task): + supported_request_fields = super().get_supported_request_fields(task) + if task == "i2i" and self.config.get("layered", False): + supported_request_fields -= {"i2i_denoise_strength"} + return supported_request_fields def __init__(self, config): super().__init__(config) @@ -149,8 +160,7 @@ def _prepare_warmup_inputs(self, height, width, t2i_text_cache=None): "text_encoder_output": text_encoder_output, "image_encoder_output": image_encoder_output, } - self.set_target_shape() - self.set_img_shapes() + self.set_latent_shape() return t2i_text_cache def clear_warmup_state(self): @@ -172,11 +182,6 @@ def clean_lazy_load_warmup(self): model = None self.maybe_empty_cache(collect_garbage=True) - def set_config(self, config_modify): - """Apply per-request overrides and optionally sync disagg fields.""" - super().set_config(config_modify) - self.apply_disagg_request_overrides(config_modify) - @ProfilingContext4DebugL2("Load models") def load_model(self): disagg_mode = self.config.get("disagg_mode") @@ -292,10 +297,8 @@ def init_modules(self): return if self.config["task"] == "t2i": self.run_input_encoder = self._run_input_encoder_local_t2i - elif self.config["task"] == "i2i": - self.run_input_encoder = self._run_input_encoder_local_i2i else: - raise NotImplementedError(f"QwenImageRunner does not support task: {self.config['task']}") + self.run_input_encoder = self._run_input_encoder_local_i2i @ProfilingContext4DebugL2("Run DiT") def _run_dit_local(self, total_steps=None): @@ -462,11 +465,6 @@ def get_custom_shape(self): logger.info(f"Qwen Image Runner got custom shape: {width}x{height}") return (width, height) - target_height = self.config.get("target_height", None) - target_width = self.config.get("target_width", None) - if target_height and target_width: - return (target_width, target_height) - aspect_ratio = self.input_info.aspect_ratio if self.input_info.aspect_ratio else self.config.get("aspect_ratio", None) if aspect_ratio in as_maps: logger.info(f"Qwen Image Runner got aspect ratio: {aspect_ratio}") @@ -476,43 +474,38 @@ def get_custom_shape(self): return None - def set_target_shape(self): + def set_latent_shape(self): # In disagg transformer mode, use the shape transmitted from encoder if self.config.get("disagg_mode") == "transformer" and getattr(self, "inputs", {}).get("latent_shape"): latent_shape = self.inputs["latent_shape"] - self.input_info.target_shape = tuple(latent_shape) - # Reconstruct auto_height and auto_width + self.input_info.latent_shape = tuple(latent_shape) scale_factor = self.config["vae_scale_factor"] - self.input_info.auto_height = latent_shape[-2] * scale_factor - self.input_info.auto_width = latent_shape[-1] * scale_factor - logger.info(f"Qwen Image Runner restored target shape from disagg: {latent_shape}") - return - - custom_shape = self.get_custom_shape() - if custom_shape is not None: - width, height = custom_shape - else: - width, height = self.input_info.original_size[-1] - calculated_width, calculated_height, _ = calculate_dimensions(self.resolution * self.resolution, width / height) - multiple_of = self.config["vae_scale_factor"] * 2 - width = calculated_width // multiple_of * multiple_of - height = calculated_height // multiple_of * multiple_of - logger.info(f"Qwen Image Runner set target shape: {width}x{height}") - self.input_info.auto_width = width - self.input_info.auto_height = height - - # VAE applies 8x compression on images but we must also account for packing which requires - # latent height and width to be divisible by 2. - height = 2 * (int(height) // (self.config["vae_scale_factor"] * 2)) - width = 2 * (int(width) // (self.config["vae_scale_factor"] * 2)) - num_channels_latents = self.config["in_channels"] // 4 - if not self.is_layered: - self.input_info.target_shape = (1, 1, num_channels_latents, height, width) + self.input_info.target_shape = [latent_shape[-2] * scale_factor, latent_shape[-1] * scale_factor] + logger.info(f"Qwen Image Runner restored latent shape from disagg: {latent_shape}") else: - self.input_info.target_shape = (1, self.layers + 1, num_channels_latents, height, width) + custom_shape = self.get_custom_shape() + if custom_shape is not None: + width, height = custom_shape + else: + width, height = self.input_info.original_size[-1] + calculated_width, calculated_height, _ = calculate_dimensions(self.resolution * self.resolution, width / height) + multiple_of = self.config["vae_scale_factor"] * 2 + width = calculated_width // multiple_of * multiple_of + height = calculated_height // multiple_of * multiple_of + logger.info(f"Qwen Image Runner set target shape: {width}x{height}") + self.input_info.target_shape = [height, width] + + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (self.config["vae_scale_factor"] * 2)) + width = 2 * (int(width) // (self.config["vae_scale_factor"] * 2)) + num_channels_latents = self.config["in_channels"] // 4 + if not self.is_layered: + self.input_info.latent_shape = (1, 1, num_channels_latents, height, width) + else: + self.input_info.latent_shape = (1, self.layers + 1, num_channels_latents, height, width) - def set_img_shapes(self): - width, height = self.input_info.auto_width, self.input_info.auto_height + height, width = self.input_info.target_shape if self.config["task"] == "t2i": image_shapes = [[(1, height // self.config["vae_scale_factor"] // 2, width // self.config["vae_scale_factor"] // 2)]] elif self.config["task"] == "i2i": @@ -544,7 +537,7 @@ def run_image_encoder(self): def _save_images(self, images, input_info, log_prefix="Image saved"): if dist.is_initialized() and dist.get_rank() != 0: return - if input_info.return_result_tensor: + if input_info.return_result_tensor or input_info.save_result_path is None: return image_prefix = input_info.save_result_path.rsplit(".", 1)[0] @@ -592,8 +585,7 @@ def _run_pipeline_local(self, input_info): self.stage_reuse_cache() if self.config["task"] == "i2i" and "image_encoder_output" in self.inputs: self.input_info.image_encoder_output = self.inputs["image_encoder_output"] - self.set_target_shape() - self.set_img_shapes() + self.set_latent_shape() logger.info(f"input_info: {self.input_info}") latents, generator = self.run_dit() images = self.run_vae_decoder(latents) @@ -609,20 +601,20 @@ def _run_pipeline_local(self, input_info): if self.input_info is not None: self.end_run() - def _run_pipeline_disagg_encoder(self): + def _run_pipeline_disagg_encoder(self, request_config): self.inputs = self.run_input_encoder() - self.set_target_shape() - self.set_img_shapes() + self.set_latent_shape() logger.info(f"input_info: {self.input_info}") - latent_shape = list(self.input_info.target_shape) - self.send_encoder_outputs(self.inputs, latent_shape) + request_config = self.build_disagg_request_config(self.input_info, request_config) + latent_shape = list(self.input_info.latent_shape) + self.send_encoder_outputs(self.inputs, latent_shape, request_config) logger.info("[Disagg] Encoder role completed. Skipping DiT run_main.") if GET_RECORDER_MODE(): monitor_cli.lightx2v_worker_request_success.inc() return None - def _run_pipeline_disagg_transformer(self, input_info): - self.inputs = self.receive_encoder_outputs() + def _run_pipeline_disagg_transformer(self, input_info, request_config): + self.inputs = self.receive_encoder_outputs(request_config) if self.config["task"] == "i2i" and "image_encoder_output" in self.inputs: self.input_info.image_encoder_output = self.inputs["image_encoder_output"] prompt_embeds = self.inputs.get("text_encoder_output", {}).get("prompt_embeds") @@ -632,8 +624,7 @@ def _run_pipeline_disagg_transformer(self, input_info): if neg_embeds is not None: self.input_info.txt_seq_lens.append(neg_embeds.shape[1]) - self.set_target_shape() - self.set_img_shapes() + self.set_latent_shape() logger.info(f"input_info: {self.input_info}") latents, generator = self.run_dit() @@ -649,25 +640,24 @@ def _run_pipeline_disagg_transformer(self, input_info): self._save_images(images, input_info, log_prefix="Image saved") return self._finalize_pipeline_outputs(input_info, images, latents=latents, generator=generator) - def _run_pipeline_disagg_decode(self, input_info): + def _run_pipeline_disagg_decode(self, input_info, request_config): # Decoder role: receive DiT latents from Transformer, decode with VAE, save image - latents = self.receive_transformer_outputs() + latents = self.receive_transformer_outputs(request_config) scale_factor = self.config["vae_scale_factor"] p2_meta = getattr(self, "_p2_receive_meta", {}) - auto_height = p2_meta.get("auto_height") - auto_width = p2_meta.get("auto_width") - if auto_height is None or auto_width is None: + target_height = p2_meta.get("auto_height") + target_width = p2_meta.get("auto_width") + if target_height is None or target_width is None: # Fallback for spatial-format latents (non-packed models) latent_h = latents.shape[-2] latent_w = latents.shape[-1] - auto_height = latent_h * scale_factor * 2 - auto_width = latent_w * scale_factor * 2 - self.input_info.auto_height = int(auto_height) - self.input_info.auto_width = int(auto_width) + target_height = latent_h * scale_factor * 2 + target_width = latent_w * scale_factor * 2 + self.input_info.target_shape = [int(target_height), int(target_width)] # Compute image_shapes: number of spatial patches per image - h_patches = int(auto_height) // (scale_factor * 2) - w_patches = int(auto_width) // (scale_factor * 2) + h_patches = int(target_height) // (scale_factor * 2) + w_patches = int(target_width) // (scale_factor * 2) self.input_info.image_shapes = [[(1, h_patches, w_patches)]] images = self.run_vae_decoder(latents) self.end_run() @@ -682,13 +672,20 @@ def _run_pipeline_disagg_decode(self, input_info): @ProfilingContext4DebugL1("RUN pipeline") def run_pipeline(self, input_info): - self.input_info = input_info disagg_mode = self.config.get("disagg_mode") + if disagg_mode in ("transformer", "decode"): + input_info.update(self._disagg_request_config or {}) + self.input_info = input_info + request_config = self.build_disagg_request_config(input_info) if disagg_mode else None - if disagg_mode == "decode": - return self._run_pipeline_disagg_decode(input_info) - if disagg_mode == "encoder": - return self._run_pipeline_disagg_encoder() - if disagg_mode == "transformer": - return self._run_pipeline_disagg_transformer(input_info) - return self._run_pipeline_local(input_info) + try: + if disagg_mode == "decode": + return self._run_pipeline_disagg_decode(input_info, request_config) + if disagg_mode == "encoder": + return self._run_pipeline_disagg_encoder(request_config) + if disagg_mode == "transformer": + return self._run_pipeline_disagg_transformer(input_info, request_config) + return self._run_pipeline_local(input_info) + finally: + if disagg_mode: + self._disagg_request_config = None diff --git a/lightx2v/models/runners/request_fields.py b/lightx2v/models/runners/request_fields.py new file mode 100644 index 000000000..7a9a5af0a --- /dev/null +++ b/lightx2v/models/runners/request_fields.py @@ -0,0 +1,8 @@ +"""Reusable request field groups; each runner declares its task-specific fields.""" + +COMMON_REQUEST_FIELDS = frozenset({"return_result_tensor", "save_result_path", "seed", "task"}) +PROMPT_FIELDS = frozenset({"negative_prompt", "prompt"}) +VIDEO_OUTPUT_FIELDS = frozenset({"target_shape", "target_video_length"}) + +IMAGE_REQUEST_FIELDS = COMMON_REQUEST_FIELDS | PROMPT_FIELDS | {"aspect_ratio", "target_shape"} +VIDEO_REQUEST_FIELDS = COMMON_REQUEST_FIELDS | PROMPT_FIELDS | VIDEO_OUTPUT_FIELDS diff --git a/lightx2v/models/runners/runner_factory.py b/lightx2v/models/runners/runner_factory.py new file mode 100644 index 000000000..d4617254d --- /dev/null +++ b/lightx2v/models/runners/runner_factory.py @@ -0,0 +1,66 @@ +from importlib import import_module + +import torch + +from lightx2v.utils.registry_factory import RUNNER_REGISTER + +RUNNER_MODULES = { + "bagel": "lightx2v.models.runners.bagel.bagel_runner", + "cosmos3": "lightx2v.models.runners.cosmos3.cosmos3_runner", + "dreamzero": "lightx2v.models.runners.wan.wan_dreamzero_runner", + "ernie_image": "lightx2v.models.runners.ernie_image.ernie_image_runner", + "fastwam": "lightx2v.models.runners.wan.fastwam_runner", + "flux2": "lightx2v.models.runners.flux2.flux2_runner", + "hidream_o1_image": "lightx2v.models.runners.hidream_o1_image.hidream_o1_image_runner", + "hunyuan3d": "lightx2v.models.runners.hunyuan3d.hunyuan3d_shape_runner", + "hunyuan_image3": "lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner", + "hunyuan_video_1.5": "lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner", + "infinitetalk": "lightx2v.models.runners.wan.wan_infinitetalk_runner", + "lingbot_va": "lightx2v.models.runners.wan.wan_lingbot_va_runner", + "lingbot_video": "lightx2v.models.runners.lingbot_video.lingbot_video_runner", + "lingbot_world": "lightx2v.models.runners.wan.wan_runner", + "lingbot_world_fast": "lightx2v.models.runners.wan.wan_lingbot_fast_runner", + "longcat_image": "lightx2v.models.runners.longcat_image.longcat_image_runner", + "ltx2": "lightx2v.models.runners.ltx2.ltx2_runner", + "ltx2_5": "lightx2v.models.runners.ltx2.ltx25_runner", + "ltx2_ar": "lightx2v.models.runners.ltx2.ltx2_runner", + "minimax_h3": "lightx2v.models.runners.minimax_h3.minimax_h3_runner", + "motus": "lightx2v.models.runners.motus.motus_runner", + "neopp": "lightx2v.models.runners.neopp.neopp_runner", + "qwen_image": "lightx2v.models.runners.qwen_image.qwen_image_runner", + "seedvr2": "lightx2v.models.runners.seedvr.seedvr_runner", + "seko_talk": "lightx2v.models.runners.wan.wan_audio_runner", + "seko_talk_ar": "lightx2v.models.runners.wan.wan_audio_runner", + "sensenova_vision": "lightx2v.models.runners.bagel.sensenova_vision_runner", + "swiftvr": "lightx2v.models.runners.swiftvr.swiftvr_runner", + "wan2.1": "lightx2v.models.runners.wan.wan_runner", + "wan2.1_sf": "lightx2v.models.runners.wan.wan_sf_runner", + "wan2.1_sf_mtxg2": "lightx2v.models.runners.wan.wan_matrix_game2_runner", + "wan2.1_vace": "lightx2v.models.runners.wan.wan_vace_runner", + "wan2.2": "lightx2v.models.runners.wan.wan_runner", + "wan2.2_animate": "lightx2v.models.runners.wan.wan_animate_runner", + "wan2.2_animate2_distilled": "lightx2v.models.runners.wan.wan_animate2_runner", + "wan2.2_audio": "lightx2v.models.runners.wan.wan_audio_runner", + "wan2.2_matrix_game3": "lightx2v.models.runners.wan.wan_matrix_game3_runner", + "wan2.2_moe": "lightx2v.models.runners.wan.wan_runner", + "wan2.2_moe_vace": "lightx2v.models.runners.wan.wan_vace_runner", + "wan2.2_s2v": "lightx2v.models.runners.wan.wan_s2v_runner", + "wan_dancer": "lightx2v.models.runners.wan.wan_dancer_runner", + "worldmirror": "lightx2v.models.runners.worldmirror.worldmirror_runner", + "worldplay_ar": "lightx2v.models.runners.worldplay.worldplay_ar_runner", + "worldplay_bi": "lightx2v.models.runners.worldplay.worldplay_bi_runner", + "worldplay_distill": "lightx2v.models.runners.worldplay.worldplay_distill_runner", + "z_image": "lightx2v.models.runners.z_image.z_image_runner", +} + + +def build_runner(config): + """Instantiate a runner and initialize its model modules.""" + model_cls = config["model_cls"] + import_module("lightx2v.common.ops") + import_module(RUNNER_MODULES[model_cls]) + + torch.set_grad_enabled(False) + runner = RUNNER_REGISTER[model_cls](config) + runner.init_modules() + return runner diff --git a/lightx2v/models/runners/seedvr/seedvr_runner.py b/lightx2v/models/runners/seedvr/seedvr_runner.py index 78741ec95..2bfe9ba2b 100755 --- a/lightx2v/models/runners/seedvr/seedvr_runner.py +++ b/lightx2v/models/runners/seedvr/seedvr_runner.py @@ -23,6 +23,7 @@ from torch import Tensor from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS from lightx2v.models.schedulers.seedvr.scheduler import SeedVRScheduler from lightx2v.models.video_encoders.hf.seedvr import attn_video_vae_v3_s8_c16_t4_inflation_sd3_init from lightx2v.models.video_encoders.hf.seedvr.color_fix import wavelet_reconstruction @@ -141,6 +142,16 @@ def read_video(filename, start_pts=0, end_pts=None, pts_unit="pts", output_forma class SeedVRRunner(DefaultRunner): """Runner for SeedVR video super-resolution model.""" + supported_request_fields_by_task = { + "sr": COMMON_REQUEST_FIELDS | {"image_path", "sr_ratio", "video_path"}, + } + + def get_supported_request_fields(self, task): + supported_request_fields = super().get_supported_request_fields(task) + if self.config.get("seq_parallel", False): + supported_request_fields -= {"image_path"} + return supported_request_fields + def __init__(self, config): super().__init__(config) self.run_input_encoder = self._run_input_encoder_local_sr @@ -227,17 +238,16 @@ def _get_sr_segment_params(self): logger.warning(f"[SeedVRRunner] sr_overlap >= sr_segment_length, clamp to {overlap}") return seg_len, overlap - def _set_output_fps(self, fps): + def set_output_fps(self, fps): if fps is None: return try: fps = float(fps) - except Exception: + except (TypeError, ValueError): return if fps <= 0: return - with self.config.temporarily_unlocked(): - self.config["fps"] = fps + self.input_info.output_fps = fps def _probe_video_torchcodec(self, video_path): from torchcodec.decoders import VideoDecoder @@ -254,7 +264,7 @@ def _probe_video_torchcodec(self, video_path): fps = float(self.config.get("fps", 16)) else: fps = float(fps) - self._set_output_fps(fps) + self.set_output_fps(fps) return int(total_frames), fps, [] @@ -272,7 +282,7 @@ def _probe_video(self, video_path): if fps_for_seek is None or fps_for_seek == 0: fps_for_seek = float(self.config.get("fps", 16)) if fps is not None and fps != 0: - self._set_output_fps(fps) + self.set_output_fps(fps) return total_frames, fps_for_seek, pts def _build_sr_segments(self, total_frames, seg_len, overlap): @@ -331,7 +341,7 @@ def _read_video_segment(self, video_path, start_idx, end_idx): ) if info is not None and self._sr_fps in [None, 0]: self._sr_fps = info.get("video_fps", self._sr_fps) - self._set_output_fps(self._sr_fps) + self.set_output_fps(self._sr_fps) if video.shape[0] > total_len: video = video[:total_len] @@ -713,7 +723,7 @@ def get_condition(self, latent: Tensor, latent_blur: Tensor, task: str) -> Tenso def _run_input_encoder_local_sr(self): """Prepare the input video, VAE latents and diffusion condition.""" - if "video_path" in self.input_info.__dataclass_fields__ and self.input_info.video_path: + if self.input_info.video_path: video_path = self.input_info.video_path if getattr(self, "_sr_segment", None) is not None: @@ -732,7 +742,7 @@ def _run_input_encoder_local_sr(self): read_video = _get_read_video() video, _, info = read_video(video_path, output_format="TCHW") if info is not None: - self._set_output_fps(info.get("video_fps", None)) + self.set_output_fps(info.get("video_fps", None)) if video.numel() == 0: raise ValueError(f"Failed to read video from {video_path}") @@ -740,7 +750,7 @@ def _run_input_encoder_local_sr(self): input_dtype = torch.float32 if self._seedvr_sp_size > 1 else GET_DTYPE() img = video.to(device=input_device, dtype=input_dtype).div_(255.0) input_source = video_path - elif "image_path" in self.input_info.__dataclass_fields__ and self.input_info.image_path: + elif self.input_info.image_path: from PIL import Image img_path = self.input_info.image_path @@ -804,7 +814,7 @@ def _run_input_encoder_local_sr(self): def run_pipeline(self, input_info): self.input_info = input_info - video_path = getattr(self.input_info, "video_path", "") + video_path = self.input_info.video_path if self._seedvr_sp_size > 1 and not video_path: raise ValueError("SeedVR VAE sequence parallel currently supports video SR input only") seg_len, overlap = self._get_sr_segment_params() @@ -845,7 +855,7 @@ def run_pipeline(self, input_info): if stream_file_output: video_recorder = SeedVRVideoRecorder( livestream_url=original_save_path, - fps=float(self.config.get("fps", 16)), + fps=float(self.get_output_fps()), ) video_recorder.config_crf = int(self.config.get("video_crf", 16)) video_recorder.config_preset = str(self.config.get("video_preset", "medium")) @@ -870,7 +880,7 @@ def run_pipeline(self, input_info): self._stream_sr_segment_video(raw, video_recorder, idx, len(segments)) else: segment_path = os.path.join(tmp_dir, f"segment_{idx:05d}.mp4") - self._save_sr_segment_video(raw, segment_path, fps=self.config.get("fps", 16)) + self._save_sr_segment_video(raw, segment_path, fps=self.get_output_fps()) segment_paths.append(segment_path) if raw is not None: del raw @@ -897,7 +907,7 @@ def run_pipeline(self, input_info): if not segment_paths: raise RuntimeError("SeedVR produced no video segments to save.") self._concat_sr_segment_videos(segment_paths, original_save_path) - input_video_path = getattr(self.input_info, "video_path", "") + input_video_path = self.input_info.video_path if input_video_path: mux_audio_from_video(input_video_path, original_save_path) logger.info(f"✅ Video saved successfully to: {original_save_path} ✅") diff --git a/lightx2v/models/runners/swiftvr/swiftvr_runner.py b/lightx2v/models/runners/swiftvr/swiftvr_runner.py index 26f2500d4..092dca797 100644 --- a/lightx2v/models/runners/swiftvr/swiftvr_runner.py +++ b/lightx2v/models/runners/swiftvr/swiftvr_runner.py @@ -22,6 +22,7 @@ padded_frame_count, ) from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS from lightx2v.server.metrics import monitor_cli from lightx2v.utils.envs import GET_DTYPE, GET_RECORDER_MODE from lightx2v.utils.profiler import ProfilingContext4DebugL1 @@ -57,12 +58,14 @@ class PendingVideoWrite: class SwiftVRRunner(DefaultRunner): """Native LightX2V runner for SwiftVR image and video restoration.""" + supported_request_fields_by_task = { + "sr": COMMON_REQUEST_FIELDS | {"image_path", "sr_ratio", "target_shape", "video_path"}, + } + # Two spatial shapes trigger dynamic compilation before serving requests. WARMUP_RESOLUTIONS = ((720, 1280), (2048, 1536)) def __init__(self, config): - if config["task"] != "sr": - raise ValueError("SwiftVR only supports the `sr` task.") if config.get("parallel"): raise ValueError("SwiftVR currently supports single-GPU inference only.") if config.get("cpu_offload"): diff --git a/lightx2v/models/runners/wan/fastwam_runner.py b/lightx2v/models/runners/wan/fastwam_runner.py index 27b9c17dd..cb14b7bd7 100644 --- a/lightx2v/models/runners/wan/fastwam_runner.py +++ b/lightx2v/models/runners/wan/fastwam_runner.py @@ -12,6 +12,7 @@ from lightx2v.models.input_encoders.hf.wan.t5.model import T5EncoderModel from lightx2v.models.networks.wan.fastwam_model import FastWAMNativeModel from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS from lightx2v.models.video_encoders.hf.wan.vae_2_2 import Wan2_2_VAE from lightx2v.utils.envs import GET_DTYPE from lightx2v.utils.registry_factory import RUNNER_REGISTER @@ -377,6 +378,10 @@ def close(self): @RUNNER_REGISTER("fastwam") class FastWAMRunner(BaseRunner): + supported_request_fields_by_task = { + "i2va": COMMON_REQUEST_FIELDS | {"image_path", "prompt", "save_action_path", "state_path"}, + } + def init_modules(self): logger.info("Loading FastWAM policy...") self.policy = FastWAMPolicy.from_config(self.config) diff --git a/lightx2v/models/runners/wan/wan_animate2_runner.py b/lightx2v/models/runners/wan/wan_animate2_runner.py index 0c71075dd..61d7e31af 100644 --- a/lightx2v/models/runners/wan/wan_animate2_runner.py +++ b/lightx2v/models/runners/wan/wan_animate2_runner.py @@ -16,8 +16,8 @@ VideoReader = None from lightx2v.common.kvcache import KVCacheManager -from lightx2v.models.networks.wan.animate2_identity import WAN_ANIMATE2_MODEL_ID from lightx2v.models.networks.wan.animate2_model import WanAnimate2Model +from lightx2v.models.runners.request_fields import VIDEO_REQUEST_FIELDS from lightx2v.models.runners.wan.wan_runner import WanRunner, build_wan_model_with_lora from lightx2v.models.schedulers.wan.animate2 import WanAnimate2Scheduler from lightx2v.utils.envs import GET_DTYPE @@ -168,7 +168,7 @@ def _abort(self): return self.returncode -@RUNNER_REGISTER(WAN_ANIMATE2_MODEL_ID) +@RUNNER_REGISTER("wan2.2_animate2_distilled") class WanAnimate2Runner(WanRunner): """Native LightX2V runner for Wan-Animate-2. @@ -178,6 +178,10 @@ class WanAnimate2Runner(WanRunner): ``wan2.2_animate`` runner. """ + supported_request_fields_by_task = { + "animate": VIDEO_REQUEST_FIELDS | {"image_path", "prompt_ref", "src_pose_path", "src_ref_images", "video_path"}, + } + def __init__(self, config): super().__init__(config) if self.config.get("disagg_mode"): @@ -188,8 +192,6 @@ def __init__(self, config): raise NotImplementedError("Wan-Animate-2 supports model/block offload, not phase offload.") if self.config.get("enable_reuse", False): raise NotImplementedError("Wan-Animate-2 request reuse is not implemented for autoregressive inputs.") - if self.config["task"] != "animate": - raise ValueError(f"{WAN_ANIMATE2_MODEL_ID} requires task='animate'.") if self.config.get("use_stream_vae", False): raise NotImplementedError("Wan-Animate-2 must drop its leading latent before Wan VAE decode; use_stream_vae is not supported.") if self.config.get("feature_caching", "NoCaching") != "NoCaching": @@ -200,8 +202,6 @@ def __init__(self, config): raise ValueError("Wan-Animate-2 requires both use_image_encoder=true and use_img_emb=true.") if not self.config.get("use_31_block", True): raise ValueError("Wan-Animate-2 requires use_31_block=true for its CLIP image features.") - if self.config.get("enable_cfg", False) != (float(self.config["sample_guide_scale"]) > 1.0): - raise ValueError("Wan-Animate-2 enables CFG exactly when sample_guide_scale > 1.") def init_scheduler(self): if self.config.get("disagg_mode") == "decode": @@ -232,11 +232,6 @@ def get_vae_parallel(self): # guaranteed to be splittable by LightX2V's spatial VAE grid either. return False - def set_inputs(self, inputs): - """Keep the reference prompt request-scoped in service mode.""" - super().set_inputs(inputs) - self.input_info.prompt_ref = inputs.get("prompt_ref", "人物动作的参考视频") - @staticmethod def _padding_resize(image, height, width, return_padding_info=False): """Match the upstream black-padding resize, including integer rounding.""" @@ -398,15 +393,14 @@ def prepare_input(self): # DefaultRunner's audio mux reads video_path. Keep legacy src_pose_path # fallback inputs source-compatible by recording the resolved driver. self.input_info.video_path = video_path - if self.input_info.seed is None or int(self.input_info.seed) < 0: - raise ValueError("Wan-Animate-2 requires a non-negative --seed.") reference_bgr = cv2.imread(reference_path, cv2.IMREAD_COLOR) if reference_bgr is None: raise ValueError(f"Failed to decode reference image: {reference_path}") reference_rgb = reference_bgr[:, :, ::-1] - target_area = int(self.config["target_width"]) * int(self.config["target_height"]) + target_height, target_width = self.get_target_size() + target_area = target_width * target_height self.reference_image, self.output_crop = self._resize_by_area( reference_rgb, target_area, @@ -420,7 +414,7 @@ def prepare_input(self): driving_shape = driving_frames[0].shape[:2] self.real_frame_len = len(driving_frames) - clip_len = int(self.config["target_video_length"]) + clip_len = self.get_target_video_length() if clip_len <= 1 or (clip_len - 1) % 4: raise ValueError(f"target_video_length must be 4k+1 and greater than 1, got {clip_len}.") padded_len = self._padding_length(self.real_frame_len, clip_len, overlap=1) @@ -633,6 +627,7 @@ def _build_segment_inputs(self, segment_idx): if list(generation_y.shape[1:]) != latent_shape[1:]: raise RuntimeError(f"Generation conditioning shape {tuple(generation_y.shape)} does not match latent shape {latent_shape}.") + target_height, target_width = self.get_target_size() animate2 = { "reference_latents": reference_latents, "reference_y": reference_y, @@ -641,7 +636,7 @@ def _build_segment_inputs(self, segment_idx): "generation_clip": self.generation_clip, "reference_kv_cache": self._build_reference_cache(reference_latents), "origin_len": clip_len, - "origin_area": [int(self.config["target_width"]), int(self.config["target_height"])], + "origin_area": [target_width, target_height], "clip_len": clip_len, } self.input_info.latent_shape = latent_shape @@ -753,7 +748,7 @@ def _close_stream_video(self, *, wait, mux_audio): if not os.path.isfile(output_path) or os.path.getsize(output_path) == 0: raise RuntimeError(f"Wan-Animate-2 FFmpeg stream did not produce a valid output file: {output_path}") - input_video_path = getattr(self.input_info, "video_path", "") + input_video_path = self.input_info.video_path if input_video_path: muxed_path = mux_audio_from_video( input_video_path, diff --git a/lightx2v/models/runners/wan/wan_animate_runner.py b/lightx2v/models/runners/wan/wan_animate_runner.py index 6433cd7af..7a5060e03 100755 --- a/lightx2v/models/runners/wan/wan_animate_runner.py +++ b/lightx2v/models/runners/wan/wan_animate_runner.py @@ -17,6 +17,7 @@ from lightx2v.models.input_encoders.hf.animate.face_encoder import FaceEncoder from lightx2v.models.input_encoders.hf.animate.motion_encoder import Generator from lightx2v.models.networks.wan.animate_model import WanAnimateModel +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, PROMPT_FIELDS from lightx2v.models.runners.wan.wan_runner import WanRunner, build_wan_model_with_lora from lightx2v.server.metrics import monitor_cli from lightx2v.utils.envs import * @@ -28,6 +29,23 @@ @RUNNER_REGISTER("wan2.2_animate") class WanAnimateRunner(WanRunner): + supported_request_fields_by_task = { + "animate": COMMON_REQUEST_FIELDS + | PROMPT_FIELDS + | { + "src_face_path", + "src_pose_path", + "src_ref_images", + "target_video_length", + }, + } + + def get_supported_request_fields(self, task): + supported_request_fields = super().get_supported_request_fields(task) + if self.config.get("replace_flag", False): + return supported_request_fields | {"mask_path", "src_bg_path"} + return supported_request_fields + def __init__(self, config): super().__init__(config) assert self.config["task"] == "animate" @@ -260,7 +278,7 @@ def run_vae_encoder( size=(H, W), mode="bicubic", ), - torch.zeros(3, self.config["target_video_length"] - self.mask_reft_len, H, W, dtype=GET_DTYPE()), + torch.zeros(3, self.get_target_video_length() - self.mask_reft_len, H, W, dtype=GET_DTYPE()), ], dim=1, ) @@ -283,7 +301,7 @@ def run_vae_encoder( mask_pixel_values=mask_pixel_values.unsqueeze(0), ) else: - y_reft = self.vae_encoder.encode(torch.zeros(1, 3, self.config["target_video_length"] - self.mask_reft_len, H, W, dtype=GET_DTYPE(), device=AI_DEVICE)) + y_reft = self.vae_encoder.encode(torch.zeros(1, 3, self.get_target_video_length() - self.mask_reft_len, H, W, dtype=GET_DTYPE(), device=AI_DEVICE)) msk_reft = self.get_i2v_mask(self.latent_t, self.latent_h, self.latent_w, self.mask_reft_len) y_reft = torch.concat([msk_reft, y_reft]) @@ -297,14 +315,15 @@ def prepare_input(self): src_ref_path = self.input_info.src_ref_images self.cond_images, self.face_images, self.refer_images = self.prepare_source(src_pose_path, src_face_path, src_ref_path) self.refer_pixel_values = torch.tensor(self.refer_images / 127.5 - 1, dtype=GET_DTYPE(), device=AI_DEVICE).permute(2, 0, 1) # chw - self.latent_t = self.config["target_video_length"] // self.config["vae_stride"][0] + 1 + target_video_length = self.get_target_video_length() + self.latent_t = target_video_length // self.config["vae_stride"][0] + 1 self.latent_h = self.refer_pixel_values.shape[-2] // self.config["vae_stride"][1] self.latent_w = self.refer_pixel_values.shape[-1] // self.config["vae_stride"][2] self.input_info.latent_shape = [self.config.get("num_channels_latents", 16), self.latent_t + 1, self.latent_h, self.latent_w] self.real_frame_len = len(self.cond_images) target_len = self.get_valid_len( self.real_frame_len, - self.config["target_video_length"], + target_video_length, overlap=self.config["refert_num"] if "refert_num" in self.config else 1, ) logger.info("real frames: {} target frames: {}".format(self.real_frame_len, target_len)) @@ -313,18 +332,19 @@ def prepare_input(self): if self.config["replace_flag"] if "replace_flag" in self.config else False: src_bg_path = self.input_info.src_bg_path - src_mask_path = self.input_info.src_mask_path + src_mask_path = self.input_info.mask_path self.bg_images, self.mask_images = self.prepare_source_for_replace(src_bg_path, src_mask_path) self.bg_images = self.inputs_padding(self.bg_images, target_len) self.mask_images = self.inputs_padding(self.mask_images, target_len) def get_video_segment_num(self): total_frames = len(self.cond_images) - self.move_frames = self.config["target_video_length"] - self.config["refert_num"] - if total_frames <= self.config["target_video_length"]: + target_video_length = self.get_target_video_length() + self.move_frames = target_video_length - self.config["refert_num"] + if total_frames <= target_video_length: self.video_segment_num = 1 else: - self.video_segment_num = 1 + (total_frames - self.config["target_video_length"] + self.move_frames - 1) // self.move_frames + self.video_segment_num = 1 + (total_frames - target_video_length + self.move_frames - 1) // self.move_frames def init_run(self): self.all_out_frames = [] @@ -355,7 +375,7 @@ def run_vae_decoder(self, latents): ) def init_run_segment(self, segment_idx): start = segment_idx * self.move_frames - end = start + self.config["target_video_length"] + end = start + self.get_target_video_length() if start == 0: self.mask_reft_len = 0 else: @@ -422,7 +442,7 @@ def process_images_after_vae_decoder(self): self.gen_video_final = torch.cat(self.all_out_frames, dim=2)[:, :, : self.real_frame_len] del self.all_out_frames gc.collect() - super().process_images_after_vae_decoder() + return super().process_images_after_vae_decoder() @ProfilingContext4DebugL1( "Run Image Encoder", diff --git a/lightx2v/models/runners/wan/wan_audio_runner.py b/lightx2v/models/runners/wan/wan_audio_runner.py index d1fb5e228..7b375a718 100755 --- a/lightx2v/models/runners/wan/wan_audio_runner.py +++ b/lightx2v/models/runners/wan/wan_audio_runner.py @@ -20,6 +20,7 @@ from lightx2v.models.input_encoders.hf.seko_audio.audio_adapter import AudioAdapter, CausalAudioSlidingProcessor from lightx2v.models.input_encoders.hf.seko_audio.audio_encoder import SekoAudioEncoderModel from lightx2v.models.networks.wan.audio_model import WanAudioARModel, WanAudioModel +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, PROMPT_FIELDS from lightx2v.models.runners.wan.wan_runner import WanRunner, build_wan_model_with_lora from lightx2v.models.schedulers.wan.audio.scheduler import EulerScheduler, WanAudioARScheduler from lightx2v.models.video_encoders.hf.wan.vae_2_2 import Wan2_2_VAE @@ -27,7 +28,7 @@ from lightx2v.utils.async_vae import AsyncVAEChunkDecoder from lightx2v.utils.audio_io import load_audio_file from lightx2v.utils.envs import * -from lightx2v.utils.input_info import UNSET +from lightx2v.utils.input_info import UNSET, S2VInputInfo from lightx2v.utils.profiler import * from lightx2v.utils.registry_factory import RUNNER_REGISTER from lightx2v.utils.utils import find_torch_model_path, fixed_shape_resize, get_optimal_patched_size_with_sp, isotropic_crop_resize, load_weights, wan_vae_to_comfy @@ -284,6 +285,18 @@ def load_image(image: Union[str, Image.Image], to_rgb: bool = True) -> Image.Ima @RUNNER_REGISTER("seko_talk") class WanAudioRunner(WanRunner): # type:ignore + supported_request_fields_by_task = { + task: COMMON_REQUEST_FIELDS + | PROMPT_FIELDS + | { + "audio_path", + "image_path", + "target_video_length", + "video_duration", + } + for task in ("s2v", "rs2v") + } + def __init__(self, config): super().__init__(config) self.name = self.config.get("name", "WanAudioRunner") @@ -320,17 +333,16 @@ def read_audio_input(self, audio_path): if GET_RECORDER_MODE(): monitor_cli.lightx2v_input_audio_len.observe(audio_len) - expected_frames = min(max(1, int(self.video_duration * target_fps)), audio_len) - if expected_frames < int(self.video_duration * target_fps): - logger.warning(f"Input video duration is greater than actual audio duration, using audio duration instead: audio_duration={audio_len / target_fps}, video_duration={self.video_duration}") + video_duration = self.input_info.video_duration or self.video_duration + expected_frames = min(max(1, int(video_duration * target_fps)), audio_len) + if expected_frames < int(video_duration * target_fps): + logger.warning(f"Input video duration is greater than actual audio duration, using audio duration instead: audio_duration={audio_len / target_fps}, video_duration={video_duration}") # Segment audio (CLI / input_info wins over config_json; target_video_length is not merged into config) target_video_length = self.config.get("target_video_length", 81) - ii = getattr(self, "input_info", None) - if ii is not None and hasattr(ii, "target_video_length"): - tvl = ii.target_video_length - if tvl is not None and tvl is not UNSET and tvl > 0: - target_video_length = tvl + request_frames = self.input_info.target_video_length + if request_frames is not None and request_frames is not UNSET and request_frames > 0: + target_video_length = request_frames if self.config.get("model_cls") == "seko_talk_ar": audio_start, audio_end = self._audio_processor.get_audio_range(0, expected_frames) audio_segments = [AudioSegment(audio_array[:, audio_start:audio_end], 0, expected_frames)] @@ -366,7 +378,7 @@ def get_audio_files_from_audio_path(self, audio_path): def _get_image_resize_kwargs(self): input_info = getattr(self, "input_info", None) return { - "resize_mode": (getattr(input_info, "resize_mode", None) if input_info is not None else None) or self.config.get("resize_mode", "adaptive"), + "resize_mode": self.config.get("resize_mode", "adaptive"), "bucket_shape": self.config.get("bucket_shape", None), "fixed_area": (getattr(input_info, "fixed_area", None) if input_info is not None else None) or self.config.get("fixed_area", None), "fixed_shape": self.config.get("fixed_shape", None), @@ -408,11 +420,7 @@ def read_image_input(self, img_path): latent_h = patched_h * self.config["patch_size"][1] latent_w = patched_w * self.config["patch_size"][2] - if hasattr(self.input_info, "target_video_length") and self.input_info.target_video_length is not None and self.input_info.target_video_length > 0: - target_video_length = self.input_info.target_video_length - latent_shape = self.get_latent_shape_with_lat_hw(latent_h, latent_w, target_video_length) - else: - latent_shape = self.get_latent_shape_with_lat_hw(latent_h, latent_w) + latent_shape = self.get_latent_shape_with_lat_hw(latent_h, latent_w, self.input_info.target_video_length) logger.info(f"[wan_audio] target_h: {target_shape[0]}, target_w: {target_shape[1]}, latent_h: {latent_h}, latent_w: {latent_w}") @@ -525,9 +533,8 @@ def prepare_prev_latents(self, prev_video: Optional[torch.Tensor], prev_frame_le """Prepare previous latents for conditioning""" dtype = GET_DTYPE() tgt_h, tgt_w = self.input_info.target_shape[0], self.input_info.target_shape[1] - if hasattr(self.input_info, "target_video_length") and self.input_info.target_video_length is not None and self.input_info.target_video_length > 0: - target_video_length = self.input_info.target_video_length - else: + target_video_length = self.input_info.target_video_length + if target_video_length is None or target_video_length <= 0: target_video_length = self.config["target_video_length"] prev_frames = torch.zeros((1, 3, target_video_length, tgt_h, tgt_w), device=AI_DEVICE) @@ -556,7 +563,6 @@ def prepare_prev_latents(self, prev_video: Optional[torch.Tensor], prev_frame_le prev_latents = self.vae_encoder.encode(prev_frames.to(dtype)) else: prev_latents = None - prev_mask = self.model.scheduler.mask else: prev_latents = self.vae_encoder.encode(prev_frames.to(dtype)) @@ -897,7 +903,7 @@ def run_clip_main(self): self.scheduler.set_audio_adapter(self.audio_adapter) self.model.scheduler.prepare( - seed=self.input_info.seed, latent_shape=self.input_info.latent_shape, infer_steps=self.input_info.infer_steps, image_encoder_output=self.inputs["image_encoder_output"] + seed=self.input_info.seed, latent_shape=self.input_info.latent_shape, infer_steps=self.config["infer_steps"], image_encoder_output=self.inputs["image_encoder_output"] ) if self.config.get("model_cls") == "wan2.2" and self.config["task"] in ["i2v", "s2v", "rs2v"]: @@ -944,53 +950,39 @@ def run_clip_pipeline(self, input_info): @RUNNER_REGISTER("wan2.2_audio") class Wan22AudioRunner(WanAudioRunner): - def __init__(self, config): - super().__init__(config) - - def load_vae_decoder(self): - # offload config - vae_offload = self.config.get("vae_cpu_offload", self.config.get("cpu_offload")) - if vae_offload: - vae_device = torch.device("cpu") - else: - vae_device = torch.device(AI_DEVICE) - vae_config = { - "vae_path": find_torch_model_path(self.config, "vae_path", "Wan2.2_VAE.pth"), - "device": vae_device, - "cpu_offload": vae_offload, - "offload_cache": self.config.get("vae_offload_cache", False), - "dummy_model": self.config.get("dummy_model", False), - } - vae_decoder = Wan2_2_VAE(**vae_config) - return vae_decoder + supported_request_fields_by_task = {task: WanAudioRunner.supported_request_fields_by_task["s2v"] for task in ("i2v", "s2v")} + input_info_cls_by_task = {"i2v": S2VInputInfo} + # The legacy i2v task still needs both image and audio inputs. + _run_input_encoder_local_i2v = WanAudioRunner._run_input_encoder_local_s2v def load_vae_encoder(self): - # offload config vae_offload = self.config.get("vae_cpu_offload", self.config.get("cpu_offload")) - if vae_offload: - vae_device = torch.device("cpu") - else: - vae_device = torch.device(AI_DEVICE) - vae_config = { - "vae_path": find_torch_model_path(self.config, "vae_path", "Wan2.2_VAE.pth"), - "device": vae_device, - "cpu_offload": vae_offload, - "offload_cache": self.config.get("vae_offload_cache", False), - "dummy_model": self.config.get("dummy_model", False), - } - if self.config.task not in ["i2v", "s2v", "rs2v"]: - return None - else: - return Wan2_2_VAE(**vae_config) + return Wan2_2_VAE( + vae_path=find_torch_model_path(self.config, "vae_path", "Wan2.2_VAE.pth"), + device=torch.device("cpu" if vae_offload else AI_DEVICE), + cpu_offload=vae_offload, + offload_cache=self.config.get("vae_offload_cache", False), + dummy_model=self.config.get("dummy_model", False), + ) + + load_vae_decoder = load_vae_encoder def load_vae(self): - vae_encoder = self.load_vae_encoder() - vae_decoder = self.load_vae_decoder() - return vae_encoder, vae_decoder + return self.load_vae_encoder(), self.load_vae_decoder() @RUNNER_REGISTER("seko_talk_ar") class WanAudioARRunner(WanAudioRunner): + supported_request_fields_by_task = { + "rs2v": (WanAudioRunner.supported_request_fields_by_task["rs2v"] - {"target_video_length"}) | {"target_shape"}, + } + + def get_supported_request_fields(self, task): + supported_request_fields = super().get_supported_request_fields(task) + if self.config.get("prompt_travel"): + supported_request_fields -= {"prompt"} + return supported_request_fields + @dataclass(frozen=True) class PromptTravelSegment: start_frame: int diff --git a/lightx2v/models/runners/wan/wan_dancer_runner.py b/lightx2v/models/runners/wan/wan_dancer_runner.py index a15c23e03..9f716cb13 100644 --- a/lightx2v/models/runners/wan/wan_dancer_runner.py +++ b/lightx2v/models/runners/wan/wan_dancer_runner.py @@ -11,6 +11,7 @@ from lightx2v.models.input_encoders.hf.wan.wan_dancer.wan_dancer import extract_music_features, split_music_features from lightx2v.models.networks.wan.dancer_model import WanDancerModel +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, PROMPT_FIELDS from lightx2v.models.runners.wan.wan_runner import WanRunner from lightx2v.models.schedulers.wan.dancer import WanDancerScheduler, WanDancerStepDistillScheduler from lightx2v.models.video_encoders.hf.wan.dancer_vae import WanDancerVAE @@ -31,6 +32,16 @@ def _barrier(): @RUNNER_REGISTER("wan_dancer") class WanDancerRunner(WanRunner): + supported_request_fields_by_task = { + "s2v": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | {"audio_path", "image_path"}, + } + + def get_supported_request_fields(self, task): + supported_request_fields = super().get_supported_request_fields(task) + if self.config.get("dancer_stage") == "local": + return supported_request_fields | {"video_path"} + return supported_request_fields + def __init__(self, config): super().__init__(config) self.vae_cls = WanDancerVAE diff --git a/lightx2v/models/runners/wan/wan_dreamzero_runner.py b/lightx2v/models/runners/wan/wan_dreamzero_runner.py index 612ab13cc..c34391e30 100644 --- a/lightx2v/models/runners/wan/wan_dreamzero_runner.py +++ b/lightx2v/models/runners/wan/wan_dreamzero_runner.py @@ -15,6 +15,7 @@ from lightx2v.models.input_encoders.hf.wan.t5.model import T5EncoderModel from lightx2v.models.input_encoders.hf.wan.xlm_roberta.model import CLIPModel from lightx2v.models.networks.wan.dreamzero_model import DreamZeroModel +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, PROMPT_FIELDS from lightx2v.models.runners.wan.wan_runner import WanRunner from lightx2v.models.schedulers.wan.dreamzero.scheduler import DreamZeroFlowUniPCScheduler from lightx2v.models.video_encoders.hf.wan.vae import WanVAE @@ -31,6 +32,10 @@ @RUNNER_REGISTER("dreamzero") class WanDreamZeroRunner(WanRunner): + supported_request_fields_by_task = { + "i2va": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | {"image_path", "save_action_path", "state_path"}, + } + def __init__(self, config): config["enable_cfg"] = config.get("enable_cfg", config.get("sample_guide_scale", 1.0) > 1) super().__init__(config) @@ -209,13 +214,9 @@ def init_modules(self): @ProfilingContext4DebugL2("Run Encoders") def _run_input_encoder_local_i2va(self): original_prompt = self.input_info.prompt - original_negative_prompt = self.input_info.negative_prompt self.input_info.prompt = self._format_droid_prompt(original_prompt) - if not original_negative_prompt: - self.input_info.negative_prompt = self.config.get("negative_prompt", "") text_encoder_output = self.run_text_encoder(self.input_info) self.input_info.prompt = original_prompt - self.input_info.negative_prompt = original_negative_prompt torch_device_module.empty_cache() gc.collect() return {"text_encoder_output": text_encoder_output, "image_encoder_output": None} @@ -397,7 +398,7 @@ def _encode_first_frame_condition(self, videos): image_zeros = torch.zeros( videos.shape[0], 3, - self.config["target_video_length"] - 1, + self.get_target_video_length() - 1, height, width, dtype=videos.dtype, @@ -683,13 +684,15 @@ def run_main(self): def process_images_after_vae_decoder(self): self.gen_video_final = self.gen_video - video_path = getattr(self.input_info, "save_result_path", None) + video_path = self.input_info.save_result_path if not video_path: raise ValueError("DreamZero requires save_result_path from input_info.") video_path = str(video_path) - action_path = getattr(self.input_info, "save_action_path", "") or str(Path(video_path).with_suffix(".actions.npy")) - if not os.path.isabs(str(action_path)): - action_path = os.path.join(os.path.dirname(video_path) or ".", str(action_path)) + action_path = self.input_info.save_action_path + if not action_path: + action_path = str(Path(video_path).with_suffix(".actions.npy")) + elif not os.path.isabs(action_path): + action_path = os.path.join(os.path.dirname(video_path), action_path) save_to_video(self.gen_video_final, video_path, fps=self.config.get("target_fps", 10), method=self.config.get("save_video_method", "imageio")) os.makedirs(os.path.dirname(action_path) or ".", exist_ok=True) np.save(action_path, self.pred_action.numpy()) diff --git a/lightx2v/models/runners/wan/wan_infinitetalk_runner.py b/lightx2v/models/runners/wan/wan_infinitetalk_runner.py index 782557de9..c9f738814 100644 --- a/lightx2v/models/runners/wan/wan_infinitetalk_runner.py +++ b/lightx2v/models/runners/wan/wan_infinitetalk_runner.py @@ -16,6 +16,7 @@ from lightx2v.models.input_encoders.hf.infinitetalk.audio_encoder import InfiniteTalkAudioEncoder from lightx2v.models.networks.wan.infinitetalk_model import WanInfiniteTalkModel +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, PROMPT_FIELDS from lightx2v.models.runners.wan.wan_runner import WanRunner from lightx2v.models.schedulers.wan.infinitetalk.scheduler import InfiniteTalkScheduler from lightx2v.server.metrics import monitor_cli @@ -103,6 +104,18 @@ def _is_video(path): @RUNNER_REGISTER("infinitetalk") class InfiniteTalkRunner(WanRunner): + supported_request_fields_by_task = { + "s2v": COMMON_REQUEST_FIELDS + | PROMPT_FIELDS + | { + "audio_path", + "image_path", + "target_video_length", + "video_duration", + "video_path", + }, + } + def __init__(self, config): super().__init__(config) assert self.config["task"] == "s2v", "InfiniteTalk runner expects task=s2v" @@ -207,7 +220,7 @@ def _load_input_data(self): audio_paths = [item.strip() for item in str(audio_path).split(",") if item.strip()] cond_audio = {f"person{idx + 1}": path for idx, path in enumerate(audio_paths)} - cond_video = getattr(self.input_info, "src_video", "") or getattr(self.input_info, "image_path", "") or self.config.get("cond_video", "") or self.config.get("image_path", "") + cond_video = self.input_info.video_path or self.input_info.image_path or self.config.get("cond_video", "") or self.config.get("image_path", "") data = { "prompt": getattr(self.input_info, "prompt", "") or self.config.get("prompt", ""), "cond_video": cond_video, @@ -222,14 +235,14 @@ def _load_input_data(self): if self.config.get("bbox", None): data["bbox"] = self.config["bbox"] - input_cond_video = getattr(self.input_info, "src_video", "") or getattr(self.input_info, "image_path", "") + input_cond_video = self.input_info.video_path or self.input_info.image_path if input_cond_video: data["cond_video"] = input_cond_video if not data.get("prompt"): raise ValueError("InfiniteTalk requires prompt from --prompt or config infinitetalk_input/prompt.") if not data.get("cond_video"): - raise ValueError("InfiniteTalk requires cond_video from --src_video, --image_path, or config.") + raise ValueError("InfiniteTalk requires cond_video from --video_path, --image_path, or config.") if not data.get("cond_audio"): raise ValueError("InfiniteTalk requires cond_audio from --audio_path or config.") diff --git a/lightx2v/models/runners/wan/wan_lingbot_fast_runner.py b/lightx2v/models/runners/wan/wan_lingbot_fast_runner.py index 458b48f20..9e0dd2dfb 100755 --- a/lightx2v/models/runners/wan/wan_lingbot_fast_runner.py +++ b/lightx2v/models/runners/wan/wan_lingbot_fast_runner.py @@ -30,6 +30,10 @@ class LingbotFastRunner(LingbotRunner): Adds SF scheduling and segment-based inference. """ + supported_request_fields_by_task = { + "i2v": LingbotRunner.supported_request_fields_by_task["i2v"] - {"target_video_length"}, + } + def __init__(self, config): WanRunner.__init__(self, config) self.control_type = config.get("control_type", "cam") diff --git a/lightx2v/models/runners/wan/wan_lingbot_va_runner.py b/lightx2v/models/runners/wan/wan_lingbot_va_runner.py index 9031c010e..b0e908fce 100644 --- a/lightx2v/models/runners/wan/wan_lingbot_va_runner.py +++ b/lightx2v/models/runners/wan/wan_lingbot_va_runner.py @@ -11,6 +11,7 @@ from lightx2v.common.kvcache import KVCacheManager from lightx2v.models.networks.wan.lingbot_va_model import WanLingbotVAModel +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, PROMPT_FIELDS from lightx2v.models.runners.wan.wan_runner import Wan22DenseRunner from lightx2v.models.schedulers.wan.lingbot_va.scheduler import LingbotVAFlowMatchScheduler from lightx2v.models.video_encoders.hf.wan.vae_2_2 import count_conv3d, patchify @@ -64,6 +65,10 @@ def encode_chunk(self, video, stream_name): @RUNNER_REGISTER("lingbot_va") class LingbotVARunner(Wan22DenseRunner): + supported_request_fields_by_task = { + "i2va": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | {"image_path"}, + } + def __init__(self, config): config["enable_cfg"] = config.get("enable_cfg", config.get("sample_guide_scale", 1.0) > 1) config["enable_action_cfg"] = config.get("enable_action_cfg", config.get("action_sample_guide_scale", 1.0) > 1) @@ -487,7 +492,7 @@ def run_main(self): def process_images_after_vae_decoder(self): self.gen_video_final = self.gen_video - video_path = getattr(self.input_info, "save_result_path", None) + video_path = self.input_info.save_result_path if not video_path: raise ValueError("LingBot-VA requires save_result_path from input_info.") video_path = str(video_path) diff --git a/lightx2v/models/runners/wan/wan_matrix_game2_runner.py b/lightx2v/models/runners/wan/wan_matrix_game2_runner.py index f7dd15f42..6c73fcd36 100755 --- a/lightx2v/models/runners/wan/wan_matrix_game2_runner.py +++ b/lightx2v/models/runners/wan/wan_matrix_game2_runner.py @@ -7,6 +7,7 @@ from lightx2v.models.input_encoders.hf.wan.matrix_game2.clip import CLIPModel from lightx2v.models.input_encoders.hf.wan.matrix_game2.conditions import Bench_actions_gta_drive, Bench_actions_templerun, Bench_actions_universal from lightx2v.models.networks.wan.matrix_game2_model import WanSFMtxg2Model +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS from lightx2v.models.runners.wan.wan_sf_runner import WanSFRunner from lightx2v.models.video_encoders.hf.wan.vae_sf import WanMtxg2VAE from lightx2v.server.metrics import monitor_cli @@ -152,6 +153,10 @@ def get_current_action(mode="universal"): @RUNNER_REGISTER("wan2.1_sf_mtxg2") class WanSFMtxg2Runner(WanSFRunner): + supported_request_fields_by_task = { + "i2v": COMMON_REQUEST_FIELDS | {"image_path"}, + } + def __init__(self, config): super().__init__(config) self.frame_process = v2.Compose( diff --git a/lightx2v/models/runners/wan/wan_matrix_game3_runner.py b/lightx2v/models/runners/wan/wan_matrix_game3_runner.py index ccadad9e1..08f99bc96 100644 --- a/lightx2v/models/runners/wan/wan_matrix_game3_runner.py +++ b/lightx2v/models/runners/wan/wan_matrix_game3_runner.py @@ -23,10 +23,12 @@ Rotation = None Slerp = None +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, PROMPT_FIELDS from lightx2v.models.runners.wan.wan_runner import Wan22DenseRunner, build_wan_model_with_lora from lightx2v.models.schedulers.scheduler import BaseScheduler from lightx2v.server.metrics import monitor_cli from lightx2v.utils.envs import GET_DTYPE +from lightx2v.utils.input_info import ActionI2VInputInfo from lightx2v.utils.profiler import GET_RECORDER_MODE, ProfilingContext4DebugL1, ProfilingContext4DebugL2 from lightx2v.utils.registry_factory import RUNNER_REGISTER from lightx2v_platform.base.global_var import AI_DEVICE @@ -1027,12 +1029,16 @@ class WanMatrixGame3Runner(Wan22DenseRunner): - Roll latent history across overlapping segments, then trim duplicated decoded frames. """ + input_info_cls_by_task = {"i2v": ActionI2VInputInfo} + supported_request_fields_by_task = { + "i2v": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | {"action_path", "image_path", "pose", "target_shape"}, + } + def __init__(self, config): with config.temporarily_unlocked(): # The public pipeline still instantiates us as "wan2.2_matrix_game3", but # the shared Wan2.2 runner expects `model_cls == "wan2.2"` for common setup. original_model_cls = str(config.get("model_cls", "wan2.2_matrix_game3")) - config["runner_model_cls"] = original_model_cls config["model_cls"] = "wan2.2" config["mode"] = "matrix_game3" config["use_image_encoder"] = False @@ -1093,15 +1099,6 @@ def __init__(self, config): self._mg3_tail_latents: Optional[torch.Tensor] = None self._mg3_noise_generator: Optional[torch.Generator] = None - def set_inputs(self, inputs): - super().set_inputs(inputs) - # Some callers still use `pose`, others use `action_path`. Mirror both so the - # runner remains compatible with older LightX2V entry points. - if "action_path" in self.input_info.__dataclass_fields__: - self.input_info.action_path = inputs.get("action_path", inputs.get("pose", "")) - if "pose" in self.input_info.__dataclass_fields__: - self.input_info.pose = inputs.get("pose", inputs.get("action_path", "")) - def run_text_encoder(self, input_info): # Official Matrix-Game-3 base inference uses a non-empty default negative # prompt for CFG. If the caller leaves `--negative_prompt` empty, reuse the @@ -1260,8 +1257,7 @@ def _segment_latent_shape(self, lat_h: int, lat_w: int, frame_count: int) -> lis def run_vae_encoder(self, img): # Unlike the generic Wan2.2 i2v path, MG3 only encodes the first frame. The # remaining temporal slots are left zeroed and later mixed with scheduler noise. - target_h = int(self.config["target_height"]) - target_w = int(self.config["target_width"]) + target_h, target_w = self.get_target_size() target_ratio = target_h / target_w input_h, input_w = img.height, img.width if input_h / input_w > target_ratio: diff --git a/lightx2v/models/runners/wan/wan_runner.py b/lightx2v/models/runners/wan/wan_runner.py index 1f2527cb2..1c258bd48 100755 --- a/lightx2v/models/runners/wan/wan_runner.py +++ b/lightx2v/models/runners/wan/wan_runner.py @@ -23,13 +23,14 @@ from lightx2v.models.networks.wan.lingbot_model import WanLingbotModel from lightx2v.models.networks.wan.model import WanModel from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import VIDEO_REQUEST_FIELDS from lightx2v.models.schedulers.wan.scheduler_factory import create_wan_scheduler, get_wan_distill_method from lightx2v.models.video_encoders.hf.wan.vae import WanVAE from lightx2v.models.video_encoders.hf.wan.vae_2_2 import Wan2_2_VAE from lightx2v.models.video_encoders.hf.wan.vae_tiny import Wan2_2_VAE_tiny, WanVAE_tiny from lightx2v.server.metrics import monitor_cli from lightx2v.utils.envs import * -from lightx2v.utils.input_info import T2VInputInfo +from lightx2v.utils.input_info import ActionI2VInputInfo, T2VInputInfo, align_target_video_length from lightx2v.utils.profiler import * from lightx2v.utils.registry_factory import RUNNER_REGISTER from lightx2v.utils.utils import * @@ -37,29 +38,20 @@ def build_wan_model_with_lora(wan_module, config, model_kwargs, lora_configs, model_type="high_noise_model"): - lora_dynamic_apply = config.get("lora_dynamic_apply", False) - - if lora_dynamic_apply: - if model_type in ["high_noise_model", "low_noise_model"]: - # For wan2.2 - lora_name_to_info = {item["name"]: item for item in lora_configs} - lora_path = lora_name_to_info[model_type]["path"] - lora_strength = lora_name_to_info[model_type]["strength"] - else: - # For wan2.1 - lora_path = lora_configs[0]["path"] - lora_strength = lora_configs[0]["strength"] - - model_kwargs["lora_path"] = lora_path - model_kwargs["lora_strength"] = lora_strength + if model_type in ["high_noise_model", "low_noise_model"]: + lora_configs = [lora_config for lora_config in lora_configs if lora_config["name"] == model_type] + if not lora_configs: + return wan_module(**model_kwargs) + + if config.get("lora_dynamic_apply", False): + model_kwargs["lora_path"] = lora_configs[0]["path"] + model_kwargs["lora_strength"] = lora_configs[0]["strength"] model = wan_module(**model_kwargs) else: assert not config.get("dit_quantized", False), "Online LoRA only for quantized models; merging LoRA is unsupported." assert not config.get("lazy_load", False), "Lazy load mode does not support LoRA merging." model = wan_module(**model_kwargs) lora_adapter = LoraAdapter(model) - if model_type in ["high_noise_model", "low_noise_model"]: - lora_configs = [lora_config for lora_config in lora_configs if lora_config["name"] == model_type] lora_adapter.apply_lora(lora_configs, model_type=model_type) return model @@ -72,6 +64,18 @@ def get_wan_model_class(distill_method): @RUNNER_REGISTER("wan2.1") class WanRunner(DisaggMixin, DefaultRunner): + supported_request_fields_by_task = { + "t2v": VIDEO_REQUEST_FIELDS, + "i2v": VIDEO_REQUEST_FIELDS | {"image_path"}, + "flf2v": VIDEO_REQUEST_FIELDS | {"image_path", "last_frame_path"}, + } + FIXED_FRAME_ATTENTION_TYPES = { + "svg_attn", + "radial_attn", + "nbhd_attn", + "nbhd_attn_flashinfer", + "general_sparse_attn", + } _WARMUP_RESOLUTIONS = ((480, 480), (720, 1280)) _WARMUP_TASKS = ("t2v", "i2v", "flf2v") _SUPPORTS_GENERIC_WARMUP = True @@ -84,6 +88,44 @@ def __init__(self, config): self.vae_name = config.get("vae_name", "Wan2.1_VAE.pth") self.tiny_vae_name = "taew2_1.pth" + def get_supported_request_fields(self, task): + supported_request_fields = super().get_supported_request_fields(task) + if self.config.get("self_attn_1_type") in self.FIXED_FRAME_ATTENTION_TYPES: + supported_request_fields -= {"target_video_length"} + if task in ("i2v", "flf2v") and self.config.get("resize_mode") and type(self).read_image_input is DefaultRunner.read_image_input: + supported_request_fields -= {"target_shape"} + return supported_request_fields + + def prepare_request(self, request_data): + input_info = super().prepare_request(request_data) + num_frames = getattr(input_info, "target_video_length", None) + if "target_video_length" in self.get_supported_request_fields(input_info.task) and num_frames is not None: + if num_frames < 1: + raise ValueError(f"num_frames must be positive, got {num_frames}") + temporal_stride = int(self.config["vae_stride"][0]) + if (num_frames - 1) % temporal_stride != 0: + aligned_frames = align_target_video_length(num_frames, temporal_stride) + if not dist.is_initialized() or dist.get_rank() == 0: + logger.warning(f"Wan num_frames must satisfy {temporal_stride}n+1; using {aligned_frames} instead of {num_frames}.") + input_info.target_video_length = aligned_frames + + if self.config.get("disagg_mode") in ("transformer", "decode"): + return input_info + + task = self.config["task"] + if task == "i2v": + if not input_info.image_path: + raise ValueError("Wan i2v requires image_path") + elif task == "flf2v": + if not input_info.image_path: + raise ValueError("Wan flf2v requires image_path") + if not input_info.last_frame_path: + raise ValueError("Wan flf2v requires last_frame_path") + elif task == "vace": + if not (input_info.video_path or input_info.src_ref_images): + raise ValueError("Wan VACE requires video_path or src_ref_images") + return input_info + def check_reuse_support(self): model_cls = self.config["model_cls"] if model_cls not in ( @@ -479,27 +521,18 @@ def reuse_key(self): reuse_key = { "prompt": self.input_info.prompt, "negative_prompt": self.input_info.negative_prompt, - "target_video_length": self.config["target_video_length"], + "target_video_length": self.get_target_video_length(), } + if "target_shape" in self.get_supported_request_fields(self.config["task"]): + reuse_key["target_shape"] = list(self.get_target_size()) if self.config["task"] == "i2v": - reuse_key.update( - { - "image_path": self.input_info.image_path.split(","), - "resize_mode": self.config.get("resize_mode"), - } - ) - elif self.config["task"] == "t2v": - target_shape = self.input_info.target_shape or ( - self.config["target_height"], - self.config["target_width"], - ) - reuse_key["target_shape"] = list(target_shape) + reuse_key["image_path"] = self.input_info.image_path.split(",") return reuse_key def reuse_input_info(self): return { - "latent_shape": list(self.input_info.latent_shape), - "target_shape": list(self.input_info.target_shape), + "latent_shape": [int(dim) for dim in self.input_info.latent_shape], + "target_shape": [int(dim) for dim in self.input_info.target_shape], } def _run_pipeline_local(self): @@ -517,23 +550,24 @@ def _run_pipeline_local(self): if self.input_info is not None: self.end_run() - def _run_pipeline_disagg_encoder(self): + def _run_pipeline_disagg_encoder(self, request_config): self.inputs = self.run_input_encoder() + request_config = self.build_disagg_request_config(self.input_info, request_config) latent_shape = list(self.input_info.latent_shape) - self.send_encoder_outputs(self.inputs, latent_shape) + self.send_encoder_outputs(self.inputs, latent_shape, request_config) logger.info("[Disagg] Encoder role completed.") return None - def _run_pipeline_disagg_transformer(self): - self.inputs = self.receive_encoder_outputs() + def _run_pipeline_disagg_transformer(self, request_config): + self.inputs = self.receive_encoder_outputs(request_config) latent_shape = self.inputs.get("latent_shape") if latent_shape: self.input_info.latent_shape = latent_shape return self._run_transformer_role() - def _run_pipeline_disagg_decode(self): + def _run_pipeline_disagg_decode(self, request_config): # Decoder role: receive DiT latents, run VAE, save video - latents = self.receive_transformer_outputs() + latents = self.receive_transformer_outputs(request_config) self.gen_video = self.run_vae_decoder(latents) self.gen_video_final = self.gen_video return self.process_images_after_vae_decoder() @@ -542,18 +576,24 @@ def _run_pipeline_disagg_decode(self): def run_pipeline(self, input_info): if GET_RECORDER_MODE(): monitor_cli.lightx2v_worker_request_count.inc() - self.input_info = input_info disagg_mode = self.config.get("disagg_mode") + if disagg_mode in ("transformer", "decode"): + input_info.update(self._disagg_request_config or {}) + self.input_info = input_info + request_config = self.build_disagg_request_config(input_info) if disagg_mode else None - if disagg_mode == "encoder": - gen_video_final = self._run_pipeline_disagg_encoder() - elif disagg_mode == "transformer": - gen_video_final = self._run_pipeline_disagg_transformer() - elif disagg_mode == "decode": - gen_video_final = self._run_pipeline_disagg_decode() - else: - # Keep default runner pipeline behavior unchanged in local mode. - gen_video_final = self._run_pipeline_local() + try: + if disagg_mode == "encoder": + gen_video_final = self._run_pipeline_disagg_encoder(request_config) + elif disagg_mode == "transformer": + gen_video_final = self._run_pipeline_disagg_transformer(request_config) + elif disagg_mode == "decode": + gen_video_final = self._run_pipeline_disagg_decode(request_config) + else: + gen_video_final = self._run_pipeline_local() + finally: + if disagg_mode: + self._disagg_request_config = None if GET_RECORDER_MODE(): monitor_cli.lightx2v_worker_request_success.inc() @@ -761,7 +801,7 @@ def _build_vae_encoder_input(self, first_frame, last_frame, height, width, world h_start, h_end = 0, height w_start, w_end = 0, width - target_video_length = self.config["target_video_length"] + target_video_length = self.get_target_video_length() if target_video_length < 1: raise ValueError(f"target_video_length must be positive, got {target_video_length}") if last_frame is not None and target_video_length < 2: @@ -793,7 +833,8 @@ def run_vae_encoder(self, first_frame, last_frame=None): if self.config.get("resize_mode", None) is None: h, w = first_frame.shape[2:] aspect_ratio = h / w - max_area = self.config["target_height"] * self.config["target_width"] + target_height, target_width = self.get_target_size() + max_area = target_height * target_width # Calculate initial latent dimensions ori_latent_h = round(np.sqrt(max_area * aspect_ratio) // self.config["vae_stride"][1] // self.config["patch_size"][1] * self.config["patch_size"][1]) @@ -845,7 +886,7 @@ def get_vae_encoder_output(self, first_frame, lat_h, lat_w, last_frame=None): world_size_h, world_size_w = self._resolve_vae_encode_grid(lat_h, lat_w) msk = torch.ones( 1, - self.config["target_video_length"], + self.get_target_video_length(), lat_h, lat_w, device=torch.device(AI_DEVICE), @@ -893,19 +934,18 @@ def get_encoder_output_i2v(self, clip_encoder_out, vae_encoder_out, text_encoder def get_latent_shape_with_lat_hw(self, latent_h, latent_w): latent_shape = [ self.config.get("num_channels_latents", 16), - (self.config["target_video_length"] - 1) // self.config["vae_stride"][0] + 1, + (self.get_target_video_length() - 1) // self.config["vae_stride"][0] + 1, latent_h, latent_w, ] return latent_shape def get_latent_shape_with_target_hw(self): - target_height = self.input_info.target_shape[0] if self.input_info.target_shape and len(self.input_info.target_shape) == 2 else self.config["target_height"] - target_width = self.input_info.target_shape[1] if self.input_info.target_shape and len(self.input_info.target_shape) == 2 else self.config["target_width"] + target_height, target_width = self.get_target_size() latent_shape = [ self.config.get("num_channels_latents", 16), - (self.config["target_video_length"] - 1) // self.config["vae_stride"][0] + 1, + (self.get_target_video_length() - 1) // self.config["vae_stride"][0] + 1, int(target_height) // self.config["vae_stride"][1], int(target_width) // self.config["vae_stride"][2], ] @@ -1125,6 +1165,7 @@ def switch_lora(self, high_lora_path: str = None, high_lora_strength: float = 1. @RUNNER_REGISTER("wan2.2") class Wan22DenseRunner(WanRunner): + supported_request_fields_by_task = {task: WanRunner.supported_request_fields_by_task[task] for task in ("t2v", "i2v")} _SUPPORTS_GENERIC_WARMUP = True def __init__(self, config): @@ -1217,7 +1258,8 @@ def get_warmup_image_encoder_output(self, latent_shape): metrics_labels=["Wan22DenseRunner"], ) def run_vae_encoder(self, img): - max_area = self.config.target_height * self.config.target_width + target_height, target_width = self.get_target_size() + max_area = target_height * target_width ih, iw = img.height, img.width dh, dw = self.config.patch_size[1] * self.config.vae_stride[1], self.config.patch_size[2] * self.config.vae_stride[2] ow, oh = best_output_size(iw, ih, dw, dh, max_area) @@ -1252,15 +1294,15 @@ def get_vae_encoder_output(self, img): @RUNNER_REGISTER("lingbot_world") class LingbotRunner(Wan22MoeRunner): + input_info_cls_by_task = {"i2v": ActionI2VInputInfo} + supported_request_fields_by_task = { + "i2v": WanRunner.supported_request_fields_by_task["i2v"] | {"action_path", "pose"}, + } + def __init__(self, config): super().__init__(config) self.control_type = config.get("control_type", "cam") - def set_inputs(self, inputs): - super().set_inputs(inputs) - if "pose" in self.input_info.__dataclass_fields__: - self.input_info.pose = inputs.get("action_path", inputs.get("pose", "")) - def load_image_encoder(self): if self.config.get("use_image_encoder", True): return super().load_image_encoder() @@ -1421,7 +1463,7 @@ def _build_lingbot_dit_cond_dict(self, action_path: str) -> dict: logger.warning("unexpected poses.npy shape: {}", c2ws_np.shape) return {} len_c2ws = ((len(c2ws_np) - 1) // 4) * 4 + 1 - frame_num = min(int(self.config["target_video_length"]), len_c2ws) + frame_num = min(self.get_target_video_length(), len_c2ws) c2ws_np = c2ws_np[:frame_num] c2ws_np = self._interp_c2ws_to_latf(c2ws_np, lat_f) c2ws = torch.from_numpy(c2ws_np).to(torch.device(AI_DEVICE)) diff --git a/lightx2v/models/runners/wan/wan_s2v_runner.py b/lightx2v/models/runners/wan/wan_s2v_runner.py index 38e82410f..956630033 100644 --- a/lightx2v/models/runners/wan/wan_s2v_runner.py +++ b/lightx2v/models/runners/wan/wan_s2v_runner.py @@ -1,9 +1,7 @@ # Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. import gc import os -import random import subprocess -import sys import numpy as np import torch @@ -14,6 +12,7 @@ from lightx2v.models.input_encoders.hf.wan.s2v.audio_encoder import AudioEncoder from lightx2v.models.networks.wan.s2v_model import WanS2VModel from lightx2v.models.networks.wan.s2v_utils import get_size_less_than_area +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, PROMPT_FIELDS from lightx2v.models.runners.wan.wan_runner import WanRunner from lightx2v.models.schedulers.wan.s2v.s2v_scheduler import WanS2VScheduler from lightx2v.server.metrics import monitor_cli @@ -55,6 +54,10 @@ def merge_video_audio(video_path: str, audio_path: str): @RUNNER_REGISTER("wan2.2_s2v") class WanS2VRunner(WanRunner): + supported_request_fields_by_task = { + "s2v": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | {"audio_path", "image_path", "src_pose_path"}, + } + def __init__(self, config): self.vae_name = "Wan2.1_VAE.pth" super().__init__(config) @@ -286,7 +289,7 @@ def run_main(self): videos_last_frames = inputs["motion_latents"].detach() out_clips = [] - seed = inputs["seed"] if inputs["seed"] >= 0 else random.randint(0, sys.maxsize) + seed = inputs["seed"] num_repeat = inputs["num_repeat"] src_pose_path = getattr(self.input_info, "src_pose_path", None) or "" diff --git a/lightx2v/models/runners/wan/wan_sf_runner.py b/lightx2v/models/runners/wan/wan_sf_runner.py index 591b0207a..2be345656 100755 --- a/lightx2v/models/runners/wan/wan_sf_runner.py +++ b/lightx2v/models/runners/wan/wan_sf_runner.py @@ -21,6 +21,10 @@ @RUNNER_REGISTER("wan2.1_sf") class WanSFRunner(WanRunner): + supported_request_fields_by_task = { + "t2v": WanRunner.supported_request_fields_by_task["t2v"] - {"target_video_length"}, + } + def __init__(self, config): super().__init__(config) diff --git a/lightx2v/models/runners/wan/wan_vace_runner.py b/lightx2v/models/runners/wan/wan_vace_runner.py index 1f62b6836..b5dd0e0c5 100755 --- a/lightx2v/models/runners/wan/wan_vace_runner.py +++ b/lightx2v/models/runners/wan/wan_vace_runner.py @@ -7,6 +7,7 @@ from lightx2v.models.input_encoders.hf.vace.vace_processor import VaceVideoProcessor from lightx2v.models.networks.wan.vace_model import WanVaceModel +from lightx2v.models.runners.request_fields import VIDEO_REQUEST_FIELDS from lightx2v.models.runners.wan.wan_runner import MultiModelStruct, WanRunner, build_wan_model_with_lora from lightx2v.server.metrics import monitor_cli from lightx2v.utils.envs import * @@ -16,9 +17,12 @@ @RUNNER_REGISTER("wan2.1_vace") class WanVaceRunner(WanRunner): + supported_request_fields_by_task = { + "vace": VIDEO_REQUEST_FIELDS | {"mask_path", "src_ref_images", "video_path"}, + } + def __init__(self, config): super().__init__(config) - assert self.config["task"] == "vace" self.vid_proc = VaceVideoProcessor( downsample=tuple([x * y for x, y in zip(self.config["vae_stride"], self.config["patch_size"])]), min_area=720 * 1280, @@ -59,7 +63,7 @@ def prepare_source(self, src_video, src_mask, src_ref_images, image_size, device src_mask[i] = torch.clamp((src_mask[i][:1, :, :, :] + 1) / 2, min=0, max=1) image_sizes.append(src_video[i].shape[2:]) elif sub_src_video is None: - src_video[i] = torch.zeros((3, self.config["target_video_length"], image_size[0], image_size[1]), device=device) + src_video[i] = torch.zeros((3, self.get_target_video_length(), image_size[0], image_size[1]), device=device) src_mask[i] = torch.ones_like(src_video[i], device=device) image_sizes.append(image_size) else: diff --git a/lightx2v/models/runners/worldmirror/worldmirror_runner.py b/lightx2v/models/runners/worldmirror/worldmirror_runner.py index 18fc84323..6f34ae2bb 100644 --- a/lightx2v/models/runners/worldmirror/worldmirror_runner.py +++ b/lightx2v/models/runners/worldmirror/worldmirror_runner.py @@ -48,6 +48,7 @@ ) from lightx2v.models.networks.worldmirror.utils.render_utils import render_interpolated_video from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS from lightx2v.utils.registry_factory import RUNNER_REGISTER from lightx2v_platform.base.global_var import AI_DEVICE @@ -132,6 +133,19 @@ def _broadcast_string(s, rank, src=0): class WorldMirrorRunner(BaseRunner): """Runner for HY-WorldMirror-2.0 3D reconstruction model.""" + supported_request_fields_by_task = { + "recon": COMMON_REQUEST_FIELDS + | { + "input_path", + "prior_cam_path", + "prior_depth_path", + "render_depth", + "render_interp_per_pair", + "save_rendered", + "strict_output_path", + }, + } + def __init__(self, config): super().__init__(config) self.model = None @@ -427,8 +441,8 @@ def run_pipeline(self, input_info): raise ValueError("input_info.input_path must be set") cfg = self.config - output_path = input_info.save_result_path or cfg.get("output_path", "inference_output") - strict_output_path = input_info.strict_output_path or cfg.get("strict_output_path", None) + output_path = input_info.save_result_path + strict_output_path = input_info.strict_output_path target_size = cfg.get("target_size", 952) fps = cfg.get("fps", 1) @@ -460,12 +474,12 @@ def run_pipeline(self, input_info): max_resolution = cfg.get("max_resolution", 1920) compress_gs_max_points = cfg.get("compress_gs_max_points", 5_000_000) - save_rendered = cfg.get("save_rendered", False) - render_interp_per_pair = cfg.get("render_interp_per_pair", 15) - render_depth = cfg.get("render_depth", False) + save_rendered = input_info.save_rendered + render_interp_per_pair = input_info.render_interp_per_pair + render_depth = input_info.render_depth - prior_cam_path = input_info.prior_cam_path or cfg.get("prior_cam_path", None) - prior_depth_path = input_info.prior_depth_path or cfg.get("prior_depth_path", None) + prior_cam_path = input_info.prior_cam_path + prior_depth_path = input_info.prior_depth_path log_time = cfg.get("log_time", True) case_t0 = time.perf_counter() @@ -485,9 +499,10 @@ def run_pipeline(self, input_info): if log_time: timings["data_loading"] = time.perf_counter() - t0 + outdir = None if strict_output_path is not None: outdir = Path(strict_output_path) - else: + elif output_path is not None: outdir = Path(output_path) / subdir_name / timestamp # 2. Adaptive resolution @@ -531,7 +546,7 @@ def run_pipeline(self, input_info): timings["gpu_mem_peak_gb"] = peak # 4. Post-processing and saving — rank 0 only so files aren't duplicated. - if self.rank == 0: + if self.rank == 0 and outdir is not None: B, S, C, H, W = imgs.shape t0 = time.perf_counter() @@ -630,16 +645,16 @@ def run_pipeline(self, input_info): if log_time: timings["render_video"] = -1.0 - if not self.is_distributed: - del predictions - torch.cuda.empty_cache() - timings["case_total"] = time.perf_counter() - case_t0 if log_time: print_and_save_timings(timings, outdir) logger.info(f"[WorldMirror] Results saved to: {outdir}") + if not self.is_distributed: + del predictions + torch.cuda.empty_cache() + if self.is_distributed: # Free local tensors and resync state across ranks before the # next request comes in. @@ -650,9 +665,10 @@ def run_pipeline(self, input_info): torch.cuda.empty_cache() dist.barrier() + result = {"output_dir": str(outdir) if outdir is not None else None} if input_info.return_result_tensor: - return {"output_dir": str(outdir), "timings": timings if self.rank == 0 else None} - return {"output_dir": str(outdir)} + result["timings"] = timings if self.rank == 0 else None + return result def end_run(self): self.input_info = None diff --git a/lightx2v/models/runners/worldplay/worldplay_ar_runner.py b/lightx2v/models/runners/worldplay/worldplay_ar_runner.py index 46cf323bb..9553a7dbe 100644 --- a/lightx2v/models/runners/worldplay/worldplay_ar_runner.py +++ b/lightx2v/models/runners/worldplay/worldplay_ar_runner.py @@ -5,9 +5,10 @@ from loguru import logger from lightx2v.models.networks.worldplay.ar_model import WorldPlayARModel -from lightx2v.models.networks.worldplay.pose_utils import pose_to_input +from lightx2v.models.networks.worldplay.pose_utils import load_pose, pose_to_input from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner import HunyuanVideo15Runner from lightx2v.models.schedulers.worldplay.ar_scheduler import WorldPlayARScheduler +from lightx2v.utils.input_info import UNSET, WorldPlayI2VInputInfo, WorldPlayT2VInputInfo from lightx2v.utils.profiler import ProfilingContext4DebugL2 from lightx2v.utils.registry_factory import RUNNER_REGISTER from lightx2v_platform.base.global_var import AI_DEVICE @@ -33,6 +34,9 @@ class WorldPlayARRunner(HunyuanVideo15Runner): - Memory window selection for long videos """ + input_info_cls_by_task = {"t2v": WorldPlayT2VInputInfo, "i2v": WorldPlayI2VInputInfo} + supported_request_fields_by_task = {task: request_fields | {"pose"} for task, request_fields in HunyuanVideo15Runner.supported_request_fields_by_task.items()} + def __init__(self, config): # AR-specific parameters self.chunk_latent_frames = config.get("chunk_latent_frames", 4) @@ -54,6 +58,19 @@ def __init__(self, config): super().__init__(config) + def prepare_request(self, request_data): + input_info = super().prepare_request(request_data) + if input_info.pose is None: + return input_info + + input_info.pose = load_pose(input_info.pose) + num_frames = (len(input_info.pose) - 1) * self.config["vae_stride"][0] + 1 + requested_frames = request_data.get("target_video_length") + if requested_frames is not None and requested_frames is not UNSET and requested_frames != num_frames: + raise ValueError(f"pose corresponds to {num_frames} frames, but num_frames is {requested_frames}; they must match.") + input_info.target_video_length = num_frames + return input_info + def init_scheduler(self): """Initialize WorldPlay AR scheduler.""" self.scheduler = WorldPlayARScheduler(self.config) @@ -81,14 +98,6 @@ def load_transformer(self): def _run_input_encoder_local_i2v(self): """Run encoders with pose processing for i2v task.""" img_ori = self.read_image_input(self.input_info.image_path) - if hasattr(self.input_info, "pose") and self.input_info.pose is not None: - from lightx2v.models.networks.worldplay.pose_utils import get_latent_num_from_pose - - latent_num = get_latent_num_from_pose(self.input_info.pose) - vae_stride_t = self.config["vae_stride"][0] - with self.config.temporarily_unlocked(): - self.config["target_video_length"] = latent_num * vae_stride_t - (vae_stride_t - 1) - logger.info(f"Auto-set target_video_length={self.config['target_video_length']} from pose ({latent_num} latent frames)") if self.sr_version and self.config_sr["is_sr_running"]: self.latent_sr_shape = self.get_sr_latent_shape_with_target_hw() self.input_info.latent_shape = self.get_latent_shape_with_target_hw(origin_size=img_ori.size) @@ -99,7 +108,7 @@ def _run_input_encoder_local_i2v(self): # Process pose input if available pose_output = None - if hasattr(self.input_info, "pose") and self.input_info.pose is not None: + if self.input_info.pose is not None: pose_output = self._process_pose_input(self.input_info.pose, self.input_info.latent_shape[1]) torch_device_module.empty_cache() @@ -118,14 +127,6 @@ def _run_input_encoder_local_i2v(self): @ProfilingContext4DebugL2("Run Encoders") def _run_input_encoder_local_t2v(self): """Run encoders with pose processing for t2v task.""" - if hasattr(self.input_info, "pose") and self.input_info.pose is not None: - from lightx2v.models.networks.worldplay.pose_utils import get_latent_num_from_pose - - latent_num = get_latent_num_from_pose(self.input_info.pose) - vae_stride_t = self.config["vae_stride"][0] - with self.config.temporarily_unlocked(): - self.config["target_video_length"] = latent_num * vae_stride_t - (vae_stride_t - 1) - logger.info(f"Auto-set target_video_length={self.config['target_video_length']} from pose ({latent_num} latent frames)") self.input_info.latent_shape = self.get_latent_shape_with_target_hw() text_encoder_output = self.run_text_encoder(self.input_info) @@ -134,7 +135,7 @@ def _run_input_encoder_local_t2v(self): # Process pose input if available pose_output = None - if hasattr(self.input_info, "pose") and self.input_info.pose is not None: + if self.input_info.pose is not None: pose_output = self._process_pose_input(self.input_info.pose, self.input_info.latent_shape[1]) torch_device_module.empty_cache() diff --git a/lightx2v/models/runners/worldplay/worldplay_bi_runner.py b/lightx2v/models/runners/worldplay/worldplay_bi_runner.py index 2dd6df426..6d0604ecd 100644 --- a/lightx2v/models/runners/worldplay/worldplay_bi_runner.py +++ b/lightx2v/models/runners/worldplay/worldplay_bi_runner.py @@ -11,6 +11,7 @@ from lightx2v.models.networks.worldplay.pose_utils import pose_to_input from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner import HunyuanVideo15Runner from lightx2v.models.schedulers.worldplay.bi_scheduler import WorldPlayBIScheduler +from lightx2v.utils.input_info import WorldPlayI2VInputInfo, WorldPlayT2VInputInfo from lightx2v.utils.profiler import ProfilingContext4DebugL2 from lightx2v.utils.registry_factory import RUNNER_REGISTER from lightx2v_platform.base.global_var import AI_DEVICE @@ -111,6 +112,9 @@ class WorldPlayBIRunner(HunyuanVideo15Runner): - Chunk-based generation with context frame selection """ + input_info_cls_by_task = {"t2v": WorldPlayT2VInputInfo, "i2v": WorldPlayI2VInputInfo} + supported_request_fields_by_task = {task: request_fields | {"pose"} for task, request_fields in HunyuanVideo15Runner.supported_request_fields_by_task.items()} + def __init__(self, config): # BI-specific parameters self.chunk_latent_frames = config.get("chunk_latent_frames", 16) # BI uses 16 by default @@ -230,7 +234,7 @@ def _run_input_encoder_local_i2v(self): # Process pose input if available pose_output = None - if hasattr(self.input_info, "pose") and self.input_info.pose is not None: + if self.input_info.pose is not None: pose_output = self._process_pose_input(self.input_info.pose, self.input_info.latent_shape[1]) torch_device_module.empty_cache() @@ -257,7 +261,7 @@ def _run_input_encoder_local_t2v(self): # Process pose input if available pose_output = None - if hasattr(self.input_info, "pose") and self.input_info.pose is not None: + if self.input_info.pose is not None: pose_output = self._process_pose_input(self.input_info.pose, self.input_info.latent_shape[1]) torch_device_module.empty_cache() diff --git a/lightx2v/models/runners/worldplay/worldplay_distill_runner.py b/lightx2v/models/runners/worldplay/worldplay_distill_runner.py index 376397ae2..c735175a4 100644 --- a/lightx2v/models/runners/worldplay/worldplay_distill_runner.py +++ b/lightx2v/models/runners/worldplay/worldplay_distill_runner.py @@ -6,6 +6,7 @@ from lightx2v.models.networks.worldplay.pose_utils import pose_to_input from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner import HunyuanVideo15Runner from lightx2v.models.schedulers.worldplay.scheduler import WorldPlayDistillScheduler +from lightx2v.utils.input_info import WorldPlayI2VInputInfo, WorldPlayT2VInputInfo from lightx2v.utils.profiler import ProfilingContext4DebugL2 from lightx2v.utils.registry_factory import RUNNER_REGISTER from lightx2v_platform.base.global_var import AI_DEVICE @@ -25,6 +26,9 @@ class WorldPlayDistillRunner(HunyuanVideo15Runner): - Few-step inference (4 steps by default) """ + input_info_cls_by_task = {"t2v": WorldPlayT2VInputInfo, "i2v": WorldPlayI2VInputInfo} + supported_request_fields_by_task = {task: request_fields | {"pose"} for task, request_fields in HunyuanVideo15Runner.supported_request_fields_by_task.items()} + def __init__(self, config): # Set default distill parameters if "denoising_step_list" not in config: @@ -78,7 +82,7 @@ def _run_input_encoder_local_i2v(self): # Process pose input if available pose_output = None - if hasattr(self.input_info, "pose") and self.input_info.pose is not None: + if self.input_info.pose is not None: pose_output = self._process_pose_input( self.input_info.pose, self.input_info.latent_shape[1], # num latent frames @@ -109,7 +113,7 @@ def _run_input_encoder_local_t2v(self): # Process pose input if available pose_output = None - if hasattr(self.input_info, "pose") and self.input_info.pose is not None: + if self.input_info.pose is not None: pose_output = self._process_pose_input( self.input_info.pose, self.input_info.latent_shape[1], # num latent frames diff --git a/lightx2v/models/runners/z_image/z_image_runner.py b/lightx2v/models/runners/z_image/z_image_runner.py index 4c7eaf8ca..a0645ef4b 100755 --- a/lightx2v/models/runners/z_image/z_image_runner.py +++ b/lightx2v/models/runners/z_image/z_image_runner.py @@ -1,5 +1,4 @@ import gc -import math import torch import torchvision.transforms.functional as TF @@ -10,6 +9,7 @@ from lightx2v.models.networks.lora_adapter import LoraAdapter from lightx2v.models.networks.z_image.model import ZImageTransformerModel from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import IMAGE_REQUEST_FIELDS from lightx2v.models.schedulers.z_image.scheduler import ZImageScheduler from lightx2v.models.video_encoders.hf.z_image.vae import AutoencoderKLZImageVAE from lightx2v.server.metrics import monitor_cli @@ -22,16 +22,6 @@ torch_device_module = getattr(torch, AI_DEVICE) -def calculate_dimensions(target_area, ratio): - width = math.sqrt(target_area * ratio) - height = width / ratio - - width = round(width / 32) * 32 - height = round(height / 32) * 32 - - return width, height, None - - def build_z_image_model_with_lora(z_image_module, config, model_kwargs, lora_configs): lora_dynamic_apply = config.get("lora_dynamic_apply", False) @@ -54,15 +44,10 @@ def build_z_image_model_with_lora(z_image_module, config, model_kwargs, lora_con class ZImageRunner(DefaultRunner): model_cpu_offload_seq = "text_encoder->transformer->vae" _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__(self, config): - super().__init__(config) - - @ProfilingContext4DebugL2("Load models") - def load_model(self): - self.model = self.load_transformer() - self.text_encoders = self.load_text_encoder() - self.vae = self.load_vae() + supported_request_fields_by_task = { + "t2i": IMAGE_REQUEST_FIELDS, + "i2i": IMAGE_REQUEST_FIELDS | {"i2i_denoise_strength", "image_path"}, + } def load_transformer(self): z_image_model_kwargs = { @@ -324,8 +309,9 @@ def get_input_target_shape(self): raise NotImplementedError - def set_target_shape(self): + def set_latent_shape(self): width, height = self.get_input_target_shape() + self.input_info.target_shape = [height, width] # VAE applies 8x compression on images but we must also account for packing which requires # latent height and width to be divisible by 2. @@ -334,26 +320,10 @@ def set_target_shape(self): height = 2 * (int(height) // (vae_scale_factor * 2)) width = 2 * (int(width) // (vae_scale_factor * 2)) num_channels_latents = self.config.get("num_channels_latents", 16) - self.input_info.target_shape = (1, num_channels_latents, height, width) - - def set_img_shapes(self): - if hasattr(self.input_info, "target_shape") and self.input_info.target_shape is not None: - if len(self.input_info.target_shape) != 4: - raise ValueError(f"target_shape must be 4D [B, C, H, W], got {len(self.input_info.target_shape)}D: {self.input_info.target_shape}") - _, _, latent_height, latent_width = self.input_info.target_shape - else: - width, height = self.get_input_target_shape() - - vae_scale_factor = self.config["vae_scale_factor"] - latent_height = 2 * (int(height) // (vae_scale_factor * 2)) - latent_width = 2 * (int(width) // (vae_scale_factor * 2)) + self.input_info.latent_shape = (1, num_channels_latents, height, width) patch_size = self.config.get("patch_size", 2) - patch_height = latent_height // patch_size - patch_width = latent_width // patch_size - - image_shapes = [(1, patch_height, patch_width)] - self.input_info.image_shapes = image_shapes + self.input_info.image_shapes = [(1, height // patch_size, width // patch_size)] def init_scheduler(self): self.scheduler = ZImageScheduler(self.config) @@ -395,15 +365,14 @@ def run_pipeline(self, input_info): if self.config["task"] == "i2i" and "image_encoder_output" in self.inputs: self.input_info.image_encoder_output = self.inputs["image_encoder_output"] - self.set_target_shape() - self.set_img_shapes() + self.set_latent_shape() logger.info(f"input_info: {self.input_info}") latents, generator = self.run_dit() images = self.run_vae_decoder(latents) self.end_run() - if not input_info.return_result_tensor and is_main_process(): + if not input_info.return_result_tensor and input_info.save_result_path is not None and is_main_process(): image = images[0] image.save(input_info.save_result_path) logger.info(f"Image saved: {input_info.save_result_path}") diff --git a/lightx2v/models/schedulers/bagel/scheduler.py b/lightx2v/models/schedulers/bagel/scheduler.py index e2e9699ff..36ddb0b78 100644 --- a/lightx2v/models/schedulers/bagel/scheduler.py +++ b/lightx2v/models/schedulers/bagel/scheduler.py @@ -48,14 +48,13 @@ def set_timesteps(self): self.dts = timesteps[:-1] - timesteps[1:] self.timesteps = timesteps[:-1] - def prepare_vae_latent(self, curr_kvlens, curr_rope, image_sizes, new_token_ids, seed=None): + def prepare_vae_latent(self, curr_kvlens, curr_rope, image_sizes, new_token_ids, seed): packed_text_ids, packed_text_indexes = list(), list() packed_vae_position_ids, packed_vae_token_indexes, packed_init_noises = list(), list(), list() packed_position_ids, packed_seqlens, packed_indexes = list(), list(), list() packed_key_value_indexes = list() query_curr = curr = 0 - seed = int(seed if seed is not None else self.config.get("seed", 42)) # A CLI/request seed must describe this request, not merely the first # request that happened to initialize the scheduler. self.generator = torch.Generator(device="cpu").manual_seed(seed) diff --git a/lightx2v/models/schedulers/cosmos3/scheduler.py b/lightx2v/models/schedulers/cosmos3/scheduler.py index b7ae7beec..2ed9056e7 100644 --- a/lightx2v/models/schedulers/cosmos3/scheduler.py +++ b/lightx2v/models/schedulers/cosmos3/scheduler.py @@ -513,21 +513,23 @@ def _build_unipc(self): return Cosmos3UniPCMultistepScheduler(**kwargs) def prepare_latents(self, input_info): - shape = tuple(input_info.target_shape) + shape = tuple(input_info.latent_shape) + target_video_length = getattr(input_info, "target_video_length", None) + target_video_length = int(self.config.get("target_video_length", 1) if target_video_length is None else target_video_length) condition_latents = getattr(input_info, "vision_condition_latents", None) condition_frame_indexes = getattr(input_info, "vision_condition_frame_indexes", None) if condition_latents is not None: shape = tuple(condition_latents.shape) - input_info.target_shape = shape + input_info.latent_shape = shape if not shape: height = int(self.config.get("target_height", 1024)) width = int(self.config.get("target_width", 1024)) scale = int(self.config.get("vae_scale_factor_spatial", self.config.get("vae_scale_factor", 16))) channels = int(self.config.get("latent_channel", 48)) temporal_scale = int(self.config.get("vae_scale_factor_temporal", 4)) - frames = (int(self.config.get("target_video_length", 1)) - 1) // temporal_scale + 1 + frames = (target_video_length - 1) // temporal_scale + 1 shape = (1, channels, frames, height // scale, width // scale) - input_info.target_shape = shape + input_info.latent_shape = shape self.generator = torch.Generator(device=AI_DEVICE).manual_seed(int(input_info.seed)) noise = torch.randn(shape, generator=self.generator, device=AI_DEVICE, dtype=GET_DTYPE()) self.vision_condition_frame_indexes = None @@ -554,11 +556,10 @@ def prepare_latents(self, input_info): sound_dim = int(self.config.get("sound_dim", 64)) sound_len = int(self.config.get("sound_latent_length", 0)) if sound_len <= 0: - num_frames = int(self.config.get("target_video_length", 189)) fps = float(self.config.get("target_fps", 24.0)) sampling_rate = int(self.config.get("sound_sampling_rate", 48000)) hop_size = int(self.config.get("sound_hop_size", 1920)) - sound_len = (int(num_frames / fps * sampling_rate) + hop_size - 1) // hop_size + sound_len = (int(target_video_length / fps * sampling_rate) + hop_size - 1) // hop_size sound_shape = (sound_dim, sound_len) self.sound_latents = torch.randn(tuple(sound_shape), generator=self.generator, device=AI_DEVICE, dtype=GET_DTYPE()) diff --git a/lightx2v/models/schedulers/ernie_image/scheduler.py b/lightx2v/models/schedulers/ernie_image/scheduler.py index f5a23d2fd..a14eed463 100644 --- a/lightx2v/models/schedulers/ernie_image/scheduler.py +++ b/lightx2v/models/schedulers/ernie_image/scheduler.py @@ -21,7 +21,7 @@ def __init__(self, config): self.rope_request_id = 0 def prepare_latents(self, input_info): - shape = tuple(input_info.target_shape) + shape = tuple(input_info.latent_shape) self.latents = torch.randn( shape, generator=self.generator, diff --git a/lightx2v/models/schedulers/flux2/scheduler.py b/lightx2v/models/schedulers/flux2/scheduler.py index 32e007cde..292f679a5 100755 --- a/lightx2v/models/schedulers/flux2/scheduler.py +++ b/lightx2v/models/schedulers/flux2/scheduler.py @@ -88,15 +88,8 @@ def prepare(self, input_info): else: logger.info(f"Generator is not None, using existing generator for latents") - if hasattr(input_info, "latent_image_ids"): - self.latent_image_ids = input_info.latent_image_ids - else: - self.latent_image_ids = None - - if hasattr(input_info, "txt_ids"): - self.txt_ids = input_info.txt_ids - else: - self.txt_ids = None + self.latent_image_ids = input_info.latent_image_ids + self.txt_ids = input_info.txt_ids self.latents = randn_tensor(input_info.latent_shape, generator=self.generator, device=AI_DEVICE, dtype=self.dtype) diff --git a/lightx2v/models/schedulers/hunyuan3d/scheduler.py b/lightx2v/models/schedulers/hunyuan3d/scheduler.py index 6120a777a..b3b62c5ab 100644 --- a/lightx2v/models/schedulers/hunyuan3d/scheduler.py +++ b/lightx2v/models/schedulers/hunyuan3d/scheduler.py @@ -65,7 +65,7 @@ def __init__(self, config): self.dtype = GET_DTYPE() self.current_timestep = None - def prepare(self, seed=None, batch_size=1, latent_shape=None): + def prepare(self, seed, batch_size=1, latent_shape=None): infer_steps = int(self.config.get("infer_steps", 50)) self.infer_steps = infer_steps @@ -78,10 +78,7 @@ def prepare(self, seed=None, batch_size=1, latent_shape=None): ) self.timesteps = timesteps - if seed is not None: - self.generator = torch.Generator(device=self.device).manual_seed(int(seed)) - else: - self.generator = None + self.generator = torch.Generator(device=self.device).manual_seed(seed) if latent_shape is None: raise ValueError("latent_shape must be provided to Hunyuan3DShapeScheduler.prepare") diff --git a/lightx2v/models/schedulers/hunyuan_image3/scheduler.py b/lightx2v/models/schedulers/hunyuan_image3/scheduler.py index 36d5966f3..ac679f599 100644 --- a/lightx2v/models/schedulers/hunyuan_image3/scheduler.py +++ b/lightx2v/models/schedulers/hunyuan_image3/scheduler.py @@ -16,8 +16,7 @@ def __init__(self, config): self.noise_pred = None def prepare(self, input_info): - seed = getattr(input_info, "seed", None) or self.config.get("seed", 42) - self.generator = torch.Generator(device=AI_DEVICE).manual_seed(seed) + self.generator = torch.Generator(device=AI_DEVICE).manual_seed(input_info.seed) def set_timesteps(self, num_inference_steps=None, device=None): num_inference_steps = num_inference_steps or self.infer_steps diff --git a/lightx2v/models/schedulers/lingbot_video/scheduler.py b/lightx2v/models/schedulers/lingbot_video/scheduler.py index d09adc3f0..30da2d5af 100755 --- a/lightx2v/models/schedulers/lingbot_video/scheduler.py +++ b/lightx2v/models/schedulers/lingbot_video/scheduler.py @@ -54,7 +54,7 @@ def _get_distilled_sigma_values(self, config): return values def prepare(self, input_info): - super().prepare(int(input_info.seed), input_info.target_shape) + super().prepare(int(input_info.seed), input_info.latent_shape) def set_timesteps(self, infer_steps=None, device=None, sigmas=None, mu=None, shift=None): if self.distilled_sigma_values is None: diff --git a/lightx2v/models/schedulers/longcat_image/scheduler.py b/lightx2v/models/schedulers/longcat_image/scheduler.py index 4bdcd5ffa..9fdbaa6f3 100755 --- a/lightx2v/models/schedulers/longcat_image/scheduler.py +++ b/lightx2v/models/schedulers/longcat_image/scheduler.py @@ -309,9 +309,8 @@ def _pack_latents(latents, batch_size, num_channels, height, width): def prepare_latents(self, input_info): """Prepare random latents for denoising.""" self.input_info = input_info - shape = input_info.target_shape - # target_shape is already in latent space: (B, C, H, W) - # where C=16 (VAE latent channels), H and W are latent dimensions + shape = input_info.latent_shape + # C=16 (VAE latent channels); H and W are latent dimensions. vae_latent_channels = shape[1] # 16 latent_height = shape[-2] latent_width = shape[-1] diff --git a/lightx2v/models/schedulers/minimax_h3/scheduler.py b/lightx2v/models/schedulers/minimax_h3/scheduler.py index fe00f905a..4b1a0b0df 100644 --- a/lightx2v/models/schedulers/minimax_h3/scheduler.py +++ b/lightx2v/models/schedulers/minimax_h3/scheduler.py @@ -18,10 +18,10 @@ from lightx2v_platform.base.global_var import AI_DEVICE -def _make_schedule(num_grid_points: int, shift: float, device) -> tuple[torch.Tensor, torch.Tensor]: - if num_grid_points < 2: - raise ValueError(f"MiniMax-H3 infer_steps must be at least 2, got {num_grid_points}") - base = torch.linspace(1.0, 0.0, num_grid_points, dtype=torch.float32, device="cpu") +def _make_schedule(infer_steps: int, shift: float, device) -> tuple[torch.Tensor, torch.Tensor]: + if infer_steps < 1: + raise ValueError(f"MiniMax-H3 infer_steps must be at least 1, got {infer_steps}") + base = torch.linspace(1.0, 0.0, infer_steps + 1, dtype=torch.float32, device="cpu") sigmas = shift * base / (1.0 + (shift - 1.0) * base) sigmas = torch.unique_consecutive(sigmas).to(device) return sigmas, 1.0 - sigmas[:-1] @@ -43,7 +43,7 @@ class MiniMaxH3Scheduler(BaseScheduler): def __init__(self, config): super().__init__(config) - self.num_grid_points = int(config["infer_steps"]) + infer_steps = int(config["infer_steps"]) self.video_shift = float(config.get("video_flow_shift", 12.0)) self.audio_shift = float(config.get("audio_flow_shift", 3.0)) self.step_update = config.get("h3_step_update", "reference_blend") @@ -51,12 +51,10 @@ def __init__(self, config): raise ValueError(f"MiniMax-H3 h3_step_update must be 'reference_blend' or 'training_euler', got {self.step_update!r}") if self.video_shift <= 0 or self.audio_shift <= 0: raise ValueError("MiniMax-H3 flow shifts must be positive") - self.video_sigmas, self.video_timesteps = _make_schedule(self.num_grid_points, self.video_shift, AI_DEVICE) - self.audio_sigmas, self.audio_timesteps = _make_schedule(self.num_grid_points, self.audio_shift, AI_DEVICE) + self.video_sigmas, self.video_timesteps = _make_schedule(infer_steps, self.video_shift, AI_DEVICE) + self.audio_sigmas, self.audio_timesteps = _make_schedule(infer_steps, self.audio_shift, AI_DEVICE) if self.video_timesteps.numel() != self.audio_timesteps.numel(): raise ValueError("video and audio schedules collapsed to different step counts") - # The user-facing value counts sigma grid points including terminal 0. - # LightX2V's loop count is the number of model evaluations. self.infer_steps = int(self.video_timesteps.numel()) self.video_latents = None self.audio_latents = None diff --git a/lightx2v/models/schedulers/qwen_image/scheduler.py b/lightx2v/models/schedulers/qwen_image/scheduler.py index 984190c6d..4229691f7 100755 --- a/lightx2v/models/schedulers/qwen_image/scheduler.py +++ b/lightx2v/models/schedulers/qwen_image/scheduler.py @@ -581,7 +581,7 @@ def prepare_i2i_denoise_strength_latents(self, input_info): if self.latents.shape[0] != 1: raise ValueError(f"i2i_denoise_strength currently supports single-image single-output editing only, got output latent batch {self.latents.shape[0]}.") - shape = input_info.target_shape + shape = input_info.latent_shape target_height, target_width = shape[-2], shape[-1] num_channels_latents = self.latents.shape[-1] // 4 image_latents = self._resize_i2i_image_latents(image_latents, target_height, target_width, num_channels_latents) @@ -592,7 +592,7 @@ def prepare_i2i_denoise_strength_latents(self, input_info): def prepare_latents(self, input_info): self.input_info = input_info - shape = input_info.target_shape + shape = input_info.latent_shape # shape: [B, T, C, H, W] width, height = shape[-1], shape[-2] num_channels_latents = self.config.get("num_channels_latents", 16) diff --git a/lightx2v/models/schedulers/wan/audio/scheduler.py b/lightx2v/models/schedulers/wan/audio/scheduler.py index 171f7869f..40410c6fb 100755 --- a/lightx2v/models/schedulers/wan/audio/scheduler.py +++ b/lightx2v/models/schedulers/wan/audio/scheduler.py @@ -45,6 +45,11 @@ def __init__(self, config): self.prev_latents = None self.prev_len = 0 + def clear(self): + super().clear() + self.prev_latents = None + self.prev_len = 0 + def set_audio_adapter(self, audio_adapter): self.audio_adapter = audio_adapter self._audio_t_emb_cache.clear() diff --git a/lightx2v/models/schedulers/wan/infinitetalk/scheduler.py b/lightx2v/models/schedulers/wan/infinitetalk/scheduler.py index b977122fe..4484f9df6 100644 --- a/lightx2v/models/schedulers/wan/infinitetalk/scheduler.py +++ b/lightx2v/models/schedulers/wan/infinitetalk/scheduler.py @@ -28,7 +28,6 @@ def __init__(self, config): self.rope_request_id = 0 def seed_everything(self, seed): - seed = seed if seed >= 0 else random.randint(0, 99999999) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) np.random.seed(seed) diff --git a/lightx2v/models/schedulers/wan/scheduler.py b/lightx2v/models/schedulers/wan/scheduler.py index ae411193f..a0fca15d8 100755 --- a/lightx2v/models/schedulers/wan/scheduler.py +++ b/lightx2v/models/schedulers/wan/scheduler.py @@ -13,7 +13,6 @@ class WanScheduler(BaseScheduler): def __init__(self, config): super().__init__(config) self.infer_steps = self.config["infer_steps"] - self.target_video_length = self.config["target_video_length"] self.sample_shift = self.config["sample_shift"] if self.config["seq_parallel"]: self.seq_p_group = self.config.get("device_mesh").get_group(mesh_dim="seq_p") @@ -35,7 +34,6 @@ def __init__(self, config): def refresh_from_config(self, config): self.config = config self.infer_steps = int(self.config["infer_steps"]) - self.target_video_length = int(self.config["target_video_length"]) self.sample_shift = float(self.config["sample_shift"]) self.sample_guide_scale = self.config.get("sample_guide_scale") self.caching_records = [True] * self.infer_steps diff --git a/lightx2v/models/schedulers/z_image/scheduler.py b/lightx2v/models/schedulers/z_image/scheduler.py index 9c61eaf82..33a3123cd 100755 --- a/lightx2v/models/schedulers/z_image/scheduler.py +++ b/lightx2v/models/schedulers/z_image/scheduler.py @@ -373,10 +373,10 @@ def prepare_i2i_denoise_strength_latents(self, input_info): def prepare_latents(self, input_info): self.input_info = input_info - shape = input_info.target_shape + shape = input_info.latent_shape if len(shape) != 4: - raise ValueError(f"target_shape must be 4D [B, C, H, W], got {len(shape)}D: {shape}") + raise ValueError(f"latent_shape must be 4D [B, C, H, W], got {len(shape)}D: {shape}") latents = randn_tensor(shape, generator=self.generator, device=AI_DEVICE, dtype=self.dtype) diff --git a/lightx2v/models/video_encoders/hf/longcat_image/vae.py b/lightx2v/models/video_encoders/hf/longcat_image/vae.py index 7bbe9f508..02988168b 100755 --- a/lightx2v/models/video_encoders/hf/longcat_image/vae.py +++ b/lightx2v/models/video_encoders/hf/longcat_image/vae.py @@ -89,7 +89,7 @@ def decode(self, latents, input_info): if self.cpu_offload: self.model.to(AI_DEVICE) - width, height = input_info.auto_width, input_info.auto_height + height, width = input_info.target_shape # Full VAE latent dimensions full_latent_height = height // self.vae_scale_factor full_latent_width = width // self.vae_scale_factor diff --git a/lightx2v/models/video_encoders/hf/qwen_image/vae.py b/lightx2v/models/video_encoders/hf/qwen_image/vae.py index db76b1740..dbbac0ad8 100755 --- a/lightx2v/models/video_encoders/hf/qwen_image/vae.py +++ b/lightx2v/models/video_encoders/hf/qwen_image/vae.py @@ -130,7 +130,7 @@ def _decode_dist(self, latents): def decode(self, latents, input_info): if self.cpu_offload: self.model.to(AI_DEVICE) - width, height = input_info.auto_width, input_info.auto_height + height, width = input_info.target_shape if self.is_layered: latents = self._unpack_latents(latents, height, width, self.config["vae_scale_factor"], self.layers) else: diff --git a/lightx2v/models/video_encoders/trt/qwen_image/vae_trt.py b/lightx2v/models/video_encoders/trt/qwen_image/vae_trt.py index 84ceef2d8..575af84a4 100644 --- a/lightx2v/models/video_encoders/trt/qwen_image/vae_trt.py +++ b/lightx2v/models/video_encoders/trt/qwen_image/vae_trt.py @@ -355,7 +355,7 @@ def encode_vae_image(self, image): @torch.no_grad() def decode(self, latents, input_info): """Decode latents to image.""" - width, height = input_info.auto_width, input_info.auto_height + height, width = input_info.target_shape if self.is_layered: latents = self._unpack_latents(latents, height, width, self.config["vae_scale_factor"], self.layers) else: diff --git a/lightx2v/pipeline.py b/lightx2v/pipeline.py index ddb0eb46e..a8739b540 100755 --- a/lightx2v/pipeline.py +++ b/lightx2v/pipeline.py @@ -1,4 +1,3 @@ -import json import os os.environ.setdefault("PROFILING_DEBUG_LEVEL", "2") @@ -7,79 +6,22 @@ os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import torch -import torch.distributed as dist from loguru import logger -from lightx2v.common.ops import * -from lightx2v.models.networks.wan.animate2_identity import WAN_ANIMATE2_MODEL_ID -from lightx2v.models.runners.ernie_image.ernie_image_runner import ErnieImageRunner # noqa: F401 -from lightx2v.models.runners.flux2.flux2_runner import Flux2Runner # noqa: F401 -from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner import HunyuanVideo15Runner # noqa: F401 -from lightx2v.models.runners.lingbot_video.lingbot_video_runner import LingBotVideoRunner # noqa: F401 -from lightx2v.models.runners.longcat_image.longcat_image_runner import LongCatImageRunner # noqa: F401 -from lightx2v.models.runners.ltx2.ltx2_runner import LTX2Runner # noqa: F401 -from lightx2v.models.runners.ltx2.ltx25_runner import LTX25Runner # noqa: F401 -from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner # noqa: F401 -from lightx2v.models.runners.neopp.neopp_runner import NeoppRunner # noqa: F401 -from lightx2v.models.runners.qwen_image.qwen_image_runner import QwenImageRunner # noqa: F401 -from lightx2v.models.runners.seedvr.seedvr_runner import SeedVRRunner # noqa: F401 -from lightx2v.models.runners.swiftvr.swiftvr_runner import SwiftVRRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_animate2_runner import WanAnimate2Runner # noqa: F401 -from lightx2v.models.runners.wan.wan_animate_runner import WanAnimateRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_audio_runner import Wan22AudioRunner, WanAudioRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_lingbot_fast_runner import LingbotFastRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_lingbot_va_runner import LingbotVARunner # noqa: F401 -from lightx2v.models.runners.wan.wan_matrix_game2_runner import WanSFMtxg2Runner # noqa: F401 -from lightx2v.models.runners.wan.wan_matrix_game3_runner import WanMatrixGame3Runner # noqa: F401 -from lightx2v.models.runners.wan.wan_runner import Wan22MoeRunner, WanRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_sf_runner import WanSFRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_vace_runner import WanVaceRunner # noqa: F401 -from lightx2v.models.runners.worldmirror.worldmirror_runner import WorldMirrorRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_ar_runner import WorldPlayARRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_bi_runner import WorldPlayBIRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_distill_runner import WorldPlayDistillRunner # noqa: F401 -from lightx2v.models.runners.z_image.z_image_runner import ZImageRunner # noqa: F401 -from lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict -from lightx2v.utils.registry_factory import RUNNER_REGISTER -from lightx2v.utils.set_config import set_config, set_parallel_config -from lightx2v.utils.utils import seed_all, validate_config_paths +from lightx2v.models.runners.runner_factory import build_runner +from lightx2v.utils.input_info import UNSET +from lightx2v.utils.set_config import build_startup_config, init_parallel +from lightx2v.utils.utils import validate_config_paths from lightx2v_platform.registry_factory import PLATFORM_DEVICE_REGISTER -def dict_like(cls): - cls.__getitem__ = lambda self, key: getattr(self, key) - cls.__setitem__ = lambda self, key, value: setattr(self, key, value) - cls.__delitem__ = lambda self, key: delattr(self, key) - cls.__contains__ = lambda self, key: hasattr(self, key) - - def update(self, *args, **kwargs): - for arg in args: - if isinstance(arg, dict): - items = arg.items() - else: - items = arg - for k, v in items: - setattr(self, k, v) - for k, v in kwargs.items(): - setattr(self, k, v) - - def get(self, key, default=None): - return getattr(self, key, default) - - cls.get = get - cls.update = update - - return cls - - -@dict_like class LightX2VPipeline: def __init__( self, task="", model_path="", model_cls="", - support_tasks=[], + support_tasks=None, sf_model_path=None, dit_original_ckpt=None, low_noise_original_ckpt=None, @@ -88,19 +30,36 @@ def __init__( distill_method=None, model_variant=None, ): + requested_model_cls = model_cls + if model_cls in ["qwen-image", "qwen-image-2512", "qwen-image-edit", "qwen-image-edit-2509", "qwen-image-edit-2511"]: + model_cls = "qwen_image" + self.task = task - self.support_tasks = support_tasks + # Select startup components without making support_tasks a request default. + if not task and support_tasks: + task = support_tasks[0] self.model_path = model_path self.model_cls = model_cls self.model_variant = model_variant - self.sf_model_path = sf_model_path - self.dit_original_ckpt = dit_original_ckpt - self.low_noise_original_ckpt = low_noise_original_ckpt - self.high_noise_original_ckpt = high_noise_original_ckpt - self.distill_method = distill_method - self.transformer_model_name = transformer_model_name + self.runner = None + self.startup_config = { + key: value + for key, value in { + "task": task, + "model_path": model_path, + "model_cls": model_cls, + "model_variant": model_variant, + "sf_model_path": sf_model_path, + "dit_original_ckpt": dit_original_ckpt, + "low_noise_original_ckpt": low_noise_original_ckpt, + "high_noise_original_ckpt": high_noise_original_ckpt, + "distill_method": distill_method, + "transformer_model_name": transformer_model_name, + }.items() + if value is not None + } - if self.model_cls in [ + if model_cls in [ "wan2.1", "wan2.1_vace", "wan2.1_sf", @@ -108,60 +67,48 @@ def __init__( "seko_talk", "seko_talk_ar", "wan2.2_moe", - "wan2.2_audio", "wan2.2_animate", - WAN_ANIMATE2_MODEL_ID, + "wan2.2_animate2_distilled", "wan2.2_s2v", ]: - self.vae_stride = (4, 8, 8) - if self.model_cls.startswith("wan2.2") and self.model_cls != WAN_ANIMATE2_MODEL_ID: - self.use_image_encoder = False - elif self.model_cls in ["wan2.2", "wan2.2_matrix_game3"]: - self.vae_stride = (4, 16, 16) - self.num_channels_latents = 48 - if self.model_cls == "wan2.2_matrix_game3": - self.use_image_encoder = False - elif self.model_cls == "hunyuan_video_1.5": - self.vae_stride = (4, 16, 16) - self.num_channels_latents = 32 - elif self.model_cls in ["ltx2", "ltx2_5"]: - self.num_channels_latents = 128 - self.audio_mel_bins = 16 - elif self.model_cls in ["cosmos3"]: - self.vae_stride = (4, 16, 16) - self.num_channels_latents = 48 - elif self.model_cls in ["minimax_h3", "minimax-h3", "minimaxh3"]: - self.model_cls = "minimax_h3" - self.support_tasks = self.support_tasks or ["t2av", "i2av", "l2av", "fl2av", "ref2av"] - self.vae_spatial_scale_factor = 16 - self.vae_scale_factor = 16 - self.fps = 24 - self.audio_sampling_rate = 32000 - self.audio_flow_shift = 3.0 - elif self.model_cls in ["lingbot_video", "lingbot-video"]: - self.model_cls = "lingbot_video" - self.vae_stride = (4, 8, 8) - self.num_channels_latents = 16 - - if model_cls in ["qwen-image", "qwen-image-2512", "qwen-image-edit", "qwen-image-edit-2509", "qwen-image-edit-2511"]: - self.CONDITION_IMAGE_SIZE = 147456 - self.USE_IMAGE_ID_IN_PROMPT = True - if model_cls == "qwen-image-edit": - self.CONDITION_IMAGE_SIZE = 1048576 - self.USE_IMAGE_ID_IN_PROMPT = False - self.model_cls = "qwen_image" - if self.task in ["i2i"]: - self.prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" - self.prompt_template_encode_start_idx = 64 - elif self.task in ["t2i"]: - self.prompt_template_encode = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" - self.prompt_template_encode_start_idx = 34 - elif self.model_cls in ["z_image"]: - self.model_cls = "z_image" - elif model_cls in ["ernie_image", "ernie-image", "ERNIE-Image"]: - self.model_cls = "ernie_image" - elif model_cls in ["longcat_image", "longcat-image"]: - self.model_cls = "longcat_image" + self.startup_config["vae_stride"] = (4, 8, 8) + if model_cls.startswith("wan2.2") and model_cls != "wan2.2_animate2_distilled": + self.startup_config["use_image_encoder"] = False + elif model_cls in ["wan2.2", "wan2.2_matrix_game3", "wan2.2_audio"]: + self.startup_config.update(vae_stride=(4, 16, 16), num_channels_latents=48) + if model_cls in ["wan2.2_matrix_game3", "wan2.2_audio"]: + self.startup_config["use_image_encoder"] = False + elif model_cls == "hunyuan_video_1.5": + self.startup_config.update(vae_stride=(4, 16, 16), num_channels_latents=32) + elif model_cls in ["ltx2", "ltx2_5"]: + self.startup_config.update(num_channels_latents=128, audio_mel_bins=16) + elif model_cls == "cosmos3": + self.startup_config.update(vae_stride=(4, 16, 16), num_channels_latents=48) + elif model_cls == "minimax_h3": + self.startup_config.update( + vae_spatial_scale_factor=16, + vae_scale_factor=16, + fps=24, + audio_sampling_rate=32000, + audio_flow_shift=3.0, + ) + elif model_cls == "lingbot_video": + self.startup_config.update(vae_stride=(4, 8, 8), num_channels_latents=16) + + if requested_model_cls in ["qwen-image", "qwen-image-2512", "qwen-image-edit", "qwen-image-edit-2509", "qwen-image-edit-2511"]: + self.startup_config.update(CONDITION_IMAGE_SIZE=147456, USE_IMAGE_ID_IN_PROMPT=True) + if requested_model_cls == "qwen-image-edit": + self.startup_config.update(CONDITION_IMAGE_SIZE=1048576, USE_IMAGE_ID_IN_PROMPT=False) + if task == "i2i": + self.startup_config["prompt_template_encode"] = ( + "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + ) + self.startup_config["prompt_template_encode_start_idx"] = 64 + elif task == "t2i": + self.startup_config["prompt_template_encode"] = ( + "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + ) + self.startup_config["prompt_template_encode_start_idx"] = 34 def create_generator( self, @@ -176,7 +123,7 @@ def create_generator( aspect_ratio="16:9", boundary=0.900, boundary_step_index=2, - denoising_step_list=[1000, 750, 500, 250], + denoising_step_list=(1000, 750, 500, 250), config_json=None, rope_type="torch_complex_rope", resize_mode=None, @@ -185,10 +132,10 @@ def create_generator( norm_modulate_backend="torch", distilled_sigma_values=None, ): - self.resize_mode = resize_mode - if config_json is not None: - self.set_infer_config_json(config_json) - else: + if self.runner is not None: + raise RuntimeError("Generator has already been created for this pipeline") + + if config_json is None: self.set_infer_config( attn_mode, rope_type, @@ -209,20 +156,32 @@ def create_generator( distilled_sigma_values, ) - config = set_config(self) + startup_config = dict(self.startup_config, config_json=config_json) + if resize_mode is not None: + startup_config["resize_mode"] = resize_mode + + config = build_startup_config(startup_config) + self.model_cls = config["model_cls"] + self.model_path = config["model_path"] + self.model_variant = config.get("model_variant", self.model_variant) validate_config_paths(config) if config["parallel"]: platform_device = PLATFORM_DEVICE_REGISTER.get(os.getenv("PLATFORM", "cuda"), None) platform_device.init_parallel_env() - set_parallel_config(config) + init_parallel(config) - self.runner = self._init_runner(config) + self.runner = build_runner(config) print(self.runner.config) - logger.info(f"Initializing {self.model_cls} runner for {self.task} task...") + logger.info(f"Initializing {self.model_cls} runner for {config['task']} task...") logger.info(f"Model path: {self.model_path}") logger.info("LightGenerator initialized successfully!") + def modify_config(self, config_modify): + """Update runtime configuration for embedding callers such as LightLLM.""" + with self.runner.config.temporarily_unlocked(): + self.runner.config.update(config_modify) + def set_infer_config( self, attn_mode, @@ -243,54 +202,47 @@ def set_infer_config( norm_modulate_backend, distilled_sigma_values, ): + config = { + "infer_steps": infer_steps, + "target_width": width, + "target_height": height, + "sample_guide_scale": guidance_scale, + "sample_shift": sample_shift, + "enable_cfg": guidance_scale != 1 and not (self.model_cls == "z_image" and guidance_scale == 0), + "rope_type": rope_type, + "fps": fps, + "aspect_ratio": aspect_ratio, + "boundary": boundary, + "boundary_step_index": boundary_step_index, + "denoising_step_list": list(denoising_step_list), + "audio_fps": audio_fps, + "double_precision_rope": double_precision_rope, + "norm_modulate_backend": norm_modulate_backend, + } if self.model_cls in ["ltx2", "ltx2_5"]: - self.distilled_sigma_values = distilled_sigma_values - self.infer_steps = len(distilled_sigma_values) - 1 if distilled_sigma_values is not None else infer_steps - else: - self.infer_steps = infer_steps - self.target_width = width - self.target_height = height + if distilled_sigma_values is not None: + config["distilled_sigma_values"] = distilled_sigma_values + config["infer_steps"] = len(distilled_sigma_values) - 1 if num_frames is not None: - self.target_video_length = num_frames - elif self.model_cls == "ltx2_5" and getattr(self, "auto_duration", False): + config["target_video_length"] = num_frames + elif self.model_cls == "ltx2_5" and self.startup_config.get("auto_duration", False): # ``None`` is meaningful for LTX-2.5: ask DurationHead to select # the request length. Preserve fixed JSON lengths for LTX-2/2.3 # and for LTX-2.5 profiles with auto duration disabled. - self.target_video_length = None - self.sample_guide_scale = guidance_scale - self.sample_shift = sample_shift - if self.sample_guide_scale == 1 or (self.model_cls == "z_image" and self.sample_guide_scale == 0): - self.enable_cfg = False - else: - self.enable_cfg = True - self.rope_type = rope_type - self.fps = fps - self.aspect_ratio = aspect_ratio - self.boundary = boundary - self.boundary_step_index = boundary_step_index - self.denoising_step_list = denoising_step_list - self.audio_fps = audio_fps - self.double_precision_rope = double_precision_rope + config["target_video_length"] = None if self.model_cls.startswith("wan"): - self.self_attn_1_type = attn_mode - self.cross_attn_1_type = attn_mode - self.cross_attn_2_type = attn_mode + config.update(self_attn_1_type=attn_mode, cross_attn_1_type=attn_mode, cross_attn_2_type=attn_mode) elif self.model_cls in ["hunyuan_video_1.5", "qwen_image", "longcat_image", "ltx2", "ltx2_5", "z_image", "lingbot_video", "minimax_h3"]: - self.attn_type = attn_mode + config["attn_type"] = attn_mode if self.model_cls == "minimax_h3": - self.video_flow_shift = sample_shift - self.fps = fps - self.audio_sampling_rate = 32000 - self.audio_flow_shift = getattr(self, "audio_flow_shift", 3.0) - self.vae_spatial_scale_factor = 16 - self.vae_scale_factor = 16 - self.norm_modulate_backend = norm_modulate_backend - - def set_infer_config_json(self, config_json): - logger.info(f"Loading infer config from {config_json}") - with open(config_json, "r") as f: - config_json = json.load(f) - self.update(config_json) + config.update( + video_flow_shift=sample_shift, + audio_sampling_rate=32000, + audio_flow_shift=self.startup_config.get("audio_flow_shift", 3.0), + vae_spatial_scale_factor=16, + vae_scale_factor=16, + ) + self.startup_config.update(config) def enable_lightvae( self, @@ -300,12 +252,13 @@ def enable_lightvae( tae_path=None, ): assert self.model_cls not in ["qwen_image", "longcat_image"] - self.use_lightvae = use_lightvae - self.use_tae = use_tae - self.vae_path = vae_path - self.tae_path = tae_path - if self.use_tae and self.model_cls.startswith("wan") and "lighttae" in tae_path: - self.need_scaled = True + self.startup_config.update(use_lightvae=use_lightvae, use_tae=use_tae) + if vae_path is not None: + self.startup_config["vae_path"] = vae_path + if tae_path is not None: + self.startup_config["tae_path"] = tae_path + if use_tae and self.model_cls.startswith("wan") and "lighttae" in tae_path: + self.startup_config["need_scaled"] = True def enable_quantize( self, @@ -319,35 +272,44 @@ def enable_quantize( image_encoder_quantized_ckpt=False, quant_scheme="fp8-sgl", text_encoder_quant_scheme=None, - skip_fp8_block_index=[0, 43, 44, 45, 46, 47], + skip_fp8_block_index=(0, 43, 44, 45, 46, 47), ): - self.dit_quantized = dit_quantized - self.dit_quant_scheme = quant_scheme - self.dit_quantized_ckpt = dit_quantized_ckpt - self.low_noise_quantized_ckpt = low_noise_quantized_ckpt - self.high_noise_quantized_ckpt = high_noise_quantized_ckpt + self.startup_config.update( + dit_quantized=dit_quantized, + dit_quant_scheme=quant_scheme, + ) + quantized_checkpoints = { + "dit_quantized_ckpt": dit_quantized_ckpt, + "low_noise_quantized_ckpt": low_noise_quantized_ckpt, + "high_noise_quantized_ckpt": high_noise_quantized_ckpt, + } + self.startup_config.update({key: value for key, value in quantized_checkpoints.items() if value is not None}) if self.model_cls.startswith("wan"): - self.t5_quant_scheme = quant_scheme - self.t5_quantized = text_encoder_quantized - self.t5_quantized_ckpt = text_encoder_quantized_ckpt - self.clip_quant_scheme = quant_scheme - self.clip_quantized = image_encoder_quantized - self.clip_quantized_ckpt = image_encoder_quantized_ckpt + self.startup_config.update( + t5_quant_scheme=quant_scheme, + t5_quantized=text_encoder_quantized, + t5_quantized_ckpt=text_encoder_quantized_ckpt, + clip_quant_scheme=quant_scheme, + clip_quantized=image_encoder_quantized, + clip_quantized_ckpt=image_encoder_quantized_ckpt, + ) elif self.model_cls in ["hunyuan_video_1.5", "qwen_image"]: - self.qwen25vl_quantized = text_encoder_quantized - self.qwen25vl_quantized_ckpt = text_encoder_quantized_ckpt - self.qwen25vl_quant_scheme = text_encoder_quant_scheme + self.startup_config.update( + qwen25vl_quantized=text_encoder_quantized, + qwen25vl_quantized_ckpt=text_encoder_quantized_ckpt, + ) + if text_encoder_quant_scheme is not None: + self.startup_config["qwen25vl_quant_scheme"] = text_encoder_quant_scheme elif self.model_cls in ["ltx2", "ltx2_5"]: - self.skip_fp8_block_index = skip_fp8_block_index + self.startup_config["skip_fp8_block_index"] = list(skip_fp8_block_index) elif self.model_cls == "z_image": - self.qwen3_quantized = text_encoder_quantized - self.qwen3_quantized_ckpt = text_encoder_quantized_ckpt - self.qwen3_quant_scheme = text_encoder_quant_scheme - elif self.model_cls == "minimax_h3": - self.dit_quantized = dit_quantized - self.dit_quantized_ckpt = dit_quantized_ckpt - self.dit_quant_scheme = quant_scheme + self.startup_config.update( + qwen3_quantized=text_encoder_quantized, + qwen3_quantized_ckpt=text_encoder_quantized_ckpt, + ) + if text_encoder_quant_scheme is not None: + self.startup_config["qwen3_quant_scheme"] = text_encoder_quant_scheme def enable_offload( self, @@ -357,9 +319,7 @@ def enable_offload( image_encoder_offload=False, vae_offload=False, ): - self.cpu_offload = cpu_offload - self.offload_granularity = offload_granularity - self.vae_cpu_offload = vae_offload + self.startup_config.update(cpu_offload=cpu_offload, offload_granularity=offload_granularity, vae_cpu_offload=vae_offload) if self.model_cls in [ "wan2.1", "wan2.1_vace", @@ -372,35 +332,35 @@ def enable_offload( "wan2.2_matrix_game3", "wan2.2_audio", "wan2.2_animate", - WAN_ANIMATE2_MODEL_ID, + "wan2.2_animate2_distilled", "wan2.2_s2v", ]: - self.t5_cpu_offload = text_encoder_offload - self.clip_cpu_offload = image_encoder_offload + self.startup_config.update(t5_cpu_offload=text_encoder_offload, clip_cpu_offload=image_encoder_offload) elif self.model_cls == "hunyuan_video_1.5": - self.qwen25vl_cpu_offload = text_encoder_offload - self.siglip_cpu_offload = image_encoder_offload - self.byt5_cpu_offload = image_encoder_offload + self.startup_config.update( + qwen25vl_cpu_offload=text_encoder_offload, + siglip_cpu_offload=image_encoder_offload, + byt5_cpu_offload=image_encoder_offload, + ) elif self.model_cls in ["qwen_image", "longcat_image"]: - self.qwen25vl_cpu_offload = text_encoder_offload + self.startup_config["qwen25vl_cpu_offload"] = text_encoder_offload elif self.model_cls in ["ltx2", "ltx2_5"]: - self.gemma_cpu_offload = text_encoder_offload + self.startup_config["gemma_cpu_offload"] = text_encoder_offload elif self.model_cls == "z_image": - self.qwen3_cpu_offload = text_encoder_offload + self.startup_config["qwen3_cpu_offload"] = text_encoder_offload elif self.model_cls == "minimax_h3": - self.text_encoder_cpu_offload = text_encoder_offload + self.startup_config["text_encoder_cpu_offload"] = text_encoder_offload def enable_lora(self, lora_configs, lora_dynamic_apply=False): - self.lora_configs = lora_configs - self.lora_dynamic_apply = lora_dynamic_apply + self.startup_config.update(lora_configs=lora_configs, lora_dynamic_apply=lora_dynamic_apply) def switch_lora(self, lora_path: str, strength: float = 1.0): if lora_path == "": logger.info("Removing LoRA weights") else: logger.info(f"Switching LoRA to: {lora_path} with strength={strength}") - if not self.lora_dynamic_apply: + if not self.runner.config.get("lora_dynamic_apply", False): logger.error("LoRA dynamic apply is not enabled. Please enable it first.") return self.runner.switch_lora(lora_path, strength) @@ -408,29 +368,29 @@ def switch_lora(self, lora_path: str, strength: float = 1.0): def enable_cache( self, cache_method="Tea", - coefficients=[], + coefficients=(), teacache_thresh=0.15, use_ret_steps=False, magcache_calibration=False, magcache_K=6, magcache_thresh=0.24, magcache_retention_ratio=0.2, - magcache_ratios=[], + magcache_ratios=(), ): - self.feature_caching = cache_method + self.startup_config["feature_caching"] = cache_method if cache_method == "Tea": - self.coefficients = coefficients - self.teacache_thresh = teacache_thresh - self.use_ret_steps = use_ret_steps + self.startup_config.update(coefficients=list(coefficients), teacache_thresh=teacache_thresh, use_ret_steps=use_ret_steps) elif cache_method == "Mag": - self.magcache_calibration = magcache_calibration - self.magcache_K = magcache_K - self.magcache_thresh = magcache_thresh - self.magcache_retention_ratio = magcache_retention_ratio - self.magcache_ratios = magcache_ratios + self.startup_config.update( + magcache_calibration=magcache_calibration, + magcache_K=magcache_K, + magcache_thresh=magcache_thresh, + magcache_retention_ratio=magcache_retention_ratio, + magcache_ratios=list(magcache_ratios), + ) def enable_parallel(self, cfg_p_size=1, seq_p_size=1, seq_p_attn_type="ulysses"): - self.parallel = { + self.startup_config["parallel"] = { "cfg_p_size": cfg_p_size, "seq_p_size": seq_p_size, "seq_p_attn_type": seq_p_attn_type, @@ -439,92 +399,65 @@ def enable_parallel(self, cfg_p_size=1, seq_p_size=1, seq_p_attn_type="ulysses") @torch.no_grad() def generate( self, - seed=42, - prompt="", - negative_prompt="", - save_result_path="lightx2v_gen_result.png", + seed=UNSET, + prompt=None, + negative_prompt=None, + save_result_path=None, task=None, image_path=None, action_path=None, - video_path=None, # For SR task (video super-resolution) + video_path=None, image_strength=None, i2i_denoise_strength=None, image_frame_idx=None, last_frame_path=None, audio_path=None, src_ref_images=None, - src_video=None, - src_mask=None, - return_result_tensor=False, - target_shape=[], + mask_path=None, + return_result_tensor=None, + target_shape=None, num_frames=None, - sr_ratio=2.0, - prompt_ref="人物动作的参考视频", + sr_ratio=None, + prompt_ref=None, + **task_inputs, ): - # Run inference (following LightX2V pattern) - # Note: image_path supports comma-separated paths for multiple images - # image_strength can be a scalar (float/int) or a list matching the number of images - # i2i_denoise_strength controls single-image edit redraw strength when explicitly set - # image_frame_idx: optional list of pixel frame indices (one per image), or None to evenly space in [0, num_frames-1] - uses_default_output_path = save_result_path == "lightx2v_gen_result.png" - is_video_output = self.model_cls in {WAN_ANIMATE2_MODEL_ID, "ltx2", "ltx2_5"} or (self.model_cls == "swiftvr" and bool(video_path)) - if uses_default_output_path and is_video_output: - save_result_path = "lightx2v_gen_result.mp4" - self.seed = seed - self.image_path = image_path - self.action_path = action_path - self.video_path = video_path # For SR task - self.sr_ratio = sr_ratio - self.last_frame_path = last_frame_path - self.audio_path = audio_path - self.src_ref_images = src_ref_images - self.src_video = src_video - self.src_mask = src_mask - self.prompt = prompt - self.prompt_ref = prompt_ref - self.negative_prompt = negative_prompt - self.save_result_path = save_result_path - self.return_result_tensor = return_result_tensor - self.target_shape = target_shape - if num_frames is not None: - self.target_video_length = num_frames - elif self.model_cls == "ltx2_5" and getattr(self, "auto_duration", False): - # Let DurationHead choose a length only for an auto-duration - # LTX-2.5 request. Keep the profile/create_generator value for - # LTX-2/2.3 and fixed-length LTX-2.5 profiles. - self.target_video_length = None - self.image_strength = image_strength - self.i2i_denoise_strength = i2i_denoise_strength - self.image_frame_idx = image_frame_idx - if task is not None: - self.task = task - self.modify_config({"task": self.task}) + """Generate one result, validating task-specific inputs in the runner. + + Seed is a non-negative integer; omitted/None defaults to 42. + NeoPP preserves explicit None for LightLLM session RNG continuation. + An omitted output path is passed to the runner as None, skipping file saving. + An omitted task uses the task explicitly set when creating the pipeline. + When only support_tasks is provided, each call must select a task. + """ + if task is None: + task = self.task + if not task: + raise ValueError("task is required when the pipeline has no default task") + request_data = { + "task": task, + "seed": seed, + "prompt": prompt, + "prompt_ref": prompt_ref, + "negative_prompt": negative_prompt, + "save_result_path": save_result_path, + "image_path": image_path, + "action_path": action_path, + "video_path": video_path, + "last_frame_path": last_frame_path, + "audio_path": audio_path, + "src_ref_images": src_ref_images, + "mask_path": mask_path, + "return_result_tensor": return_result_tensor, + "target_shape": target_shape, + "target_video_length": num_frames, + "sr_ratio": sr_ratio, + "image_strength": image_strength, + "i2i_denoise_strength": i2i_denoise_strength, + "image_frame_idx": image_frame_idx, + } + request_data.update(task_inputs) - # MiniMax-H3 validates the task-specific input dataclass in its runner. - # Do not merge its T2AV/I2AV/L2AV/FL2AV/Ref2AV schemas here. - input_support_tasks = [] if self.model_cls == "minimax_h3" else self.support_tasks - input_info = init_empty_input_info(self.task, input_support_tasks) - if self.model_cls == WAN_ANIMATE2_MODEL_ID and (self.seed is None or self.seed < 0): - raise ValueError(f"{WAN_ANIMATE2_MODEL_ID} requires a non-negative seed") - if self.seed is not None: - seed_all(self.seed) - update_input_info_from_dict(input_info, self) - gen_result = self.runner.run_pipeline(input_info) + input_info = self.runner.prepare_request(request_data) + gen_result = self.runner.run_request(input_info) logger.info("Generated successfully!") - logger.info(f"Saved in {save_result_path}") return gen_result - - def _init_runner(self, config): - torch.set_grad_enabled(False) - runner = RUNNER_REGISTER[config["model_cls"]](config) - runner.init_modules() - return runner - - def _init_parallel(self): - dist.init_process_group(backend="nccl") - torch.cuda.set_device(dist.get_rank()) - - def modify_config(self, config_modify): - logger.info(f"modify config: {config_modify}") - with self.runner.config.temporarily_unlocked(): - self.runner.config.update(config_modify) diff --git a/lightx2v/server/README.md b/lightx2v/server/README.md index d7e3dd5b1..4cdad3f8c 100644 --- a/lightx2v/server/README.md +++ b/lightx2v/server/README.md @@ -216,12 +216,12 @@ sequenceDiagram end par Parallel Inference across all ranks - TIW0->>TIW0: runner.set_inputs(task_data) - TIW0->>TIW0: runner.run_pipeline() + TIW0->>TIW0: input_info = runner.prepare_request(task_data) + TIW0->>TIW0: runner.run_request(input_info) and Note over TIW1: If world_size > 1 - TIW1->>TIW1: runner.set_inputs(task_data) - TIW1->>TIW1: runner.run_pipeline() + TIW1->>TIW1: input_info = runner.prepare_request(task_data) + TIW1->>TIW1: runner.run_request(input_info) end Note over TIW0,TIW1: Synchronization @@ -304,10 +304,17 @@ stateDiagram-v2 ```python class VideoTaskRequest(BaseTaskRequest): - target_video_length: int = 81 + target_video_length: Optional[int] = Field( + None, + validation_alias=AliasChoices("num_frames", "target_video_length"), + ) + reuse_prefix_segments: int = Field(0, ge=0) + video_path: str = "" + sr_ratio: float = Field(2.0, gt=0) audio_path: str = "" video_duration: int = 5 talk_objects: Optional[list[TalkObject]] = None + src_ref_images: list[str] = Field(default_factory=list) ``` ### ImageTaskRequest @@ -325,11 +332,20 @@ class BaseTaskRequest(BaseModel): prompt: str = "" negative_prompt: str = "" image_path: str = "" # URL, base64, or local path - save_result_path: str = "" - infer_steps: int = 5 - seed: int # auto-generated + save_result_path: Optional[str] = None # omitted/null: do not save a file + seed: Optional[int] # non-negative; omitted/null: 42 ``` +To download an asynchronous task's file result, provide `save_result_path` in the request. Omitting it or passing null skips file saving. Synchronous image APIs return the image from memory. + +Single-task services use the task selected by `--task` at startup, so POST requests can omit `task`. For runners that support multiple tasks, every native JSON or form request must specify `task`; unsupported tasks are rejected. The OpenAI image generation and editing endpoints select `t2i` and `i2i`, respectively. + +Native image and video request schemas accept the inference fields supported by their runners, including animation conditions, action/state paths, LTX reference controls, and image-edit options. The runner validates whether a field is supported by the selected model, task, and startup configuration. Unknown JSON or form fields are rejected with HTTP 422; startup settings such as `infer_steps`, `resize_mode`, and `warmup` belong in the startup configuration. + +The `/form` endpoints accept the same named request fields as the JSON endpoints, alongside their existing file uploads. Encode structured values such as `target_shape`, `image_frame_idx`, `image_strength` lists, `src_ref_images`, `talk_objects`, and WorldPlay `pose` objects as JSON strings. For example, use `target_shape='[480,832]'`. Text fields, including prompts and `layout_bboxes`, retain their submitted text. + +Video, pose-video, mask-video, action, and state paths refer to files or directories on the server. Image and audio inputs retain their existing URL, Base64, and local-path handling. Python-only inputs such as callbacks and in-memory policy tensors are not exposed as JSON fields; synchronous image APIs select tensor output internally. + ## Configuration ### Environment Variables diff --git a/lightx2v/server/__main__.py b/lightx2v/server/__main__.py index 06f2d2d4c..d4947b7a1 100644 --- a/lightx2v/server/__main__.py +++ b/lightx2v/server/__main__.py @@ -8,22 +8,16 @@ def main(): parser.add_argument("--model_path", type=str, required=True, help="Path to model") parser.add_argument("--model_cls", type=str, required=True, help="Model class name") + parser.add_argument("--task", type=str, required=True, help="Inference task") + parser.add_argument("--config_json", type=str, required=True, help="Path to startup config") parser.add_argument("--lora_dir", type=str, default=None, help="Directory containing LoRA files (.safetensors)") parser.add_argument("--host", type=str, default="0.0.0.0", help="Server host") parser.add_argument("--port", type=int, default=8000, help="Server port") + parser.add_argument("--metric_port", type=int, default=None, help="Metrics server port") parser.add_argument("--max_queue_size", type=int, default=10, help="Maximum active tasks (pending + processing)") - args, unknown = parser.parse_known_args() - - for i in range(0, len(unknown), 2): - if unknown[i].startswith("--"): - key = unknown[i][2:] - if i + 1 < len(unknown) and not unknown[i + 1].startswith("--"): - value = unknown[i + 1] - setattr(args, key, value) - - run_server(args) + run_server(parser.parse_args()) if __name__ == "__main__": diff --git a/lightx2v/server/api/openai_images.py b/lightx2v/server/api/openai_images.py index 7c90932a6..c99c70380 100644 --- a/lightx2v/server/api/openai_images.py +++ b/lightx2v/server/api/openai_images.py @@ -179,19 +179,22 @@ def _build_openai_response(request: Request, task_id: str, image_bytes: bytes, r def _build_image_task_request( prompt: str, *, - negative_prompt: str = "", + task: str, + negative_prompt: Optional[str] = None, seed: Optional[int] = None, target_shape: Optional[list[int]] = None, image_path: str = "", image_mask_path: str = "", i2i_denoise_strength: Optional[float] = None, ) -> ImageTaskRequest: - payload = { - "prompt": prompt, - "negative_prompt": negative_prompt, - "image_path": image_path, + payload = {"task": task, "prompt": prompt} + optional_fields = { "image_mask_path": image_mask_path, + "image_path": image_path, } + payload.update({key: value for key, value in optional_fields.items() if value}) + if negative_prompt is not None: + payload["negative_prompt"] = negative_prompt if target_shape: payload["target_shape"] = target_shape if seed is not None: @@ -217,6 +220,7 @@ async def create_openai_image_generation(request: Request, body: OpenAIImageGene raise HTTPException(status_code=400, detail=str(e)) message = _build_image_task_request( + task="t2i", prompt=body.prompt, seed=body.seed, target_shape=target_shape, @@ -254,9 +258,9 @@ async def create_openai_image_edit( seed: int | None = Form(default=None), i2i_denoise_strength: float | None = Form(default=None), ): + form = await request.form() image_uploads = list(image or []) if not image_uploads: - form = await request.form() image_uploads = [upload for upload in form.getlist("image[]") if hasattr(upload, "filename") and hasattr(upload, "read")] _ = model, user @@ -291,8 +295,9 @@ async def create_openai_image_edit( image_mask_path = await _save_upload_file(mask, services.file_service) message = _build_image_task_request( + task="i2i", prompt=prompt, - negative_prompt=negative_prompt, + negative_prompt=form.get("negative_prompt"), seed=seed, target_shape=target_shape, image_path=image_path, diff --git a/lightx2v/server/api/tasks/common.py b/lightx2v/server/api/tasks/common.py index f3bb411ea..bd1b0fd6f 100644 --- a/lightx2v/server/api/tasks/common.py +++ b/lightx2v/server/api/tasks/common.py @@ -1,10 +1,13 @@ import gc +import json from pathlib import Path import torch from fastapi import APIRouter, HTTPException +from fastapi.exceptions import RequestValidationError from fastapi.responses import StreamingResponse from loguru import logger +from pydantic import ValidationError from ...schema import StopTaskResponse from ...task_manager import TaskStatus, task_manager @@ -14,6 +17,21 @@ router = APIRouter() +def parse_form_request(request_cls, request_data): + """Decode structured form fields and validate them with the JSON request schema.""" + for field in ("target_shape", "src_ref_images", "image_frame_idx", "image_strength", "pose", "talk_objects"): + value = request_data.get(field) + if isinstance(value, str) and value.lstrip().startswith(("[", "{")): + try: + request_data[field] = json.loads(value) + except ValueError as exc: + raise HTTPException(status_code=422, detail=f"{field} must contain valid JSON") from exc + try: + return request_cls(**request_data) + except ValidationError as exc: + raise RequestValidationError(exc.errors(include_input=False)) from exc + + def _stream_file_response(file_path: Path, filename: str | None = None) -> StreamingResponse: services = get_services() assert services.file_service is not None, "File service is not initialized" diff --git a/lightx2v/server/api/tasks/image.py b/lightx2v/server/api/tasks/image.py index 2b9eb7712..e41922f08 100644 --- a/lightx2v/server/api/tasks/image.py +++ b/lightx2v/server/api/tasks/image.py @@ -8,6 +8,7 @@ from ...schema import ImageTaskRequest, TaskResponse from ...task_manager import TaskStatus, task_manager from ..deps import get_services, validate_url_async +from .common import parse_form_request router = APIRouter() @@ -177,13 +178,14 @@ async def create_image_task_sync( @router.post("/form", response_model=TaskResponse) async def create_image_task_form( + request: Request, + task: str | None = Form(default=None), image_file: UploadFile = File(None), prompt: str = Form(default=""), save_result_path: str = Form(default=""), negative_prompt: str = Form(default=""), - infer_steps: int = Form(default=5), - seed: int = Form(default=42), - aspect_ratio: str = Form(default="16:9"), + seed: int | None = Form(default=None), + aspect_ratio: str | None = Form(default=None), ): services = get_services() assert services.file_service is not None, "File service is not initialized" @@ -193,14 +195,12 @@ async def create_image_task_form( content = await image_file.read() image_path = str(await asyncio.to_thread(services.file_service.save_uploaded_file, content, image_file.filename)) - message = ImageTaskRequest( - prompt=prompt, - negative_prompt=negative_prompt, - image_path=image_path, - save_result_path=save_result_path, - infer_steps=infer_steps, - seed=seed, - aspect_ratio=aspect_ratio, - ) + request_data = {"seed": seed} if seed is not None else {} + if image_path: + request_data["image_path"] = image_path + # FastAPI replaces empty form strings with defaults; preserve submitted text. + form = await request.form() + form_data = {key: value for key, value in form.items() if key != "image_file"} + message = parse_form_request(ImageTaskRequest, form_data | request_data) return await create_image_task(message) diff --git a/lightx2v/server/api/tasks/video.py b/lightx2v/server/api/tasks/video.py index 892b9ef5f..db85fd540 100644 --- a/lightx2v/server/api/tasks/video.py +++ b/lightx2v/server/api/tasks/video.py @@ -1,11 +1,12 @@ import asyncio -from fastapi import APIRouter, File, Form, HTTPException, UploadFile +from fastapi import APIRouter, File, Form, HTTPException, Request, UploadFile from loguru import logger from ...schema import TaskResponse, VideoTaskRequest from ...task_manager import task_manager from ..deps import get_services, validate_url_async +from .common import parse_form_request router = APIRouter() @@ -34,17 +35,17 @@ async def create_video_task(message: VideoTaskRequest): @router.post("/form", response_model=TaskResponse) async def create_video_task_form( - image_file: UploadFile = File(...), + request: Request, + task: str | None = Form(default=None), + image_file: UploadFile = File(None), last_frame_file: UploadFile = File(None), prompt: str = Form(default=""), save_result_path: str = Form(default=""), negative_prompt: str = Form(default=""), - infer_steps: int = Form(default=5), - target_video_length: int = Form(default=81), - seed: int = Form(default=42), + target_video_length: int | None = Form(default=None), + seed: int | None = Form(default=None), audio_file: UploadFile = File(None), - video_duration: int = Form(default=5), - target_fps: int = Form(default=16), + video_duration: float | None = Form(default=None), ): services = get_services() assert services.file_service is not None, "File service is not initialized" @@ -64,18 +65,18 @@ async def create_video_task_form( content = await audio_file.read() audio_path = str(await asyncio.to_thread(services.file_service.save_uploaded_file, content, audio_file.filename, services.file_service.input_audio_dir)) - message = VideoTaskRequest( - prompt=prompt, - negative_prompt=negative_prompt, - last_frame_path=last_frame_path, - image_path=image_path, - save_result_path=save_result_path, - infer_steps=infer_steps, - target_video_length=target_video_length, - seed=seed, - audio_path=audio_path, - video_duration=video_duration, - target_fps=target_fps, - ) + request_data = {"seed": seed} if seed is not None else {} + optional_fields = { + "audio_path": audio_path, + "image_path": image_path, + "last_frame_path": last_frame_path, + "target_video_length": target_video_length, + "video_duration": video_duration, + } + request_data.update({key: value for key, value in optional_fields.items() if value not in (None, "")}) + # FastAPI replaces empty form strings with defaults; preserve submitted text. + form = await request.form() + form_data = {key: value for key, value in form.items() if key not in {"image_file", "last_frame_file", "audio_file"}} + message = parse_form_request(VideoTaskRequest, form_data | request_data) return await create_video_task(message) diff --git a/lightx2v/server/schema.py b/lightx2v/server/schema.py index eb2cf6c52..c4129c179 100644 --- a/lightx2v/server/schema.py +++ b/lightx2v/server/schema.py @@ -1,7 +1,6 @@ -import random from typing import Any, Optional -from pydantic import BaseModel, Field +from pydantic import AliasChoices, BaseModel, ConfigDict, Field from ..utils.generate_task_id import generate_task_id @@ -24,38 +23,26 @@ class Usage(BaseModel): output_tokens_details: UsageOutputTokensDetails -def generate_random_seed() -> int: - return random.randint(0, 2**32 - 1) - - class TalkObject(BaseModel): + model_config = ConfigDict(extra="forbid") + audio: str = Field(..., description="Audio path") mask: str = Field(..., description="Mask path") -class DisaggOverrideRequest(BaseModel): - """Optional Mooncake / disagg overrides (merged into runner.config per request).""" - - data_bootstrap_room: Optional[int] = Field(None, description="Per-request Mooncake bootstrap room (disagg)") - disagg_phase1_receiver_engine_rank: Optional[int] = Field( - None, - description="Transformer receiver rank for phase1 send (decentralized / multi-transformer)", - ) - disagg_bootstrap_room: Optional[int] = Field(None, description="Alias for data_bootstrap_room in some clients") - disagg_decoder_bootstrap_room: Optional[int] = Field(None, description="Phase2 Mooncake room override") - +class BaseTaskRequest(BaseModel): + model_config = ConfigDict(extra="forbid") -class BaseTaskRequest(DisaggOverrideRequest): task_id: str = Field(default_factory=generate_task_id, description="Task ID (auto-generated)") + task: Optional[str] = Field(None, description="Required for multi-task runners; single-task services use their startup task") prompt: str = Field("", description="Generation prompt") negative_prompt: str = Field("", description="Negative prompt") image_path: str = Field("", description="Base64 encoded image or URL") last_frame_path: str = Field("", description="Last frame image path (base64, or local path)") image_mask_path: str = Field("", description="Mask image path (supports URL, base64, or local path)") - save_result_path: str = Field("", description="Save result path (optional, defaults to task_id, suffix auto-detected)") + save_result_path: Optional[str] = Field(None, description="Output path; omitted or null skips file saving") presigned_url: str = Field("", description="Optional presigned URL for uploading final sync result") - infer_steps: int = Field(5, description="Inference steps") - seed: int = Field(default_factory=generate_random_seed, description="Random seed (auto-generated if not set)") + seed: Optional[int] = Field(None, description="Non-negative seed; omitted or null defaults to 42") reuse: bool = Field(False, description="Reuse the previous successful request") target_shape: list[int] = Field([], description="Return video or image shape") lora_name: Optional[str] = Field(None, description="LoRA filename to load from lora_dir, None to disable LoRA") @@ -63,40 +50,62 @@ class BaseTaskRequest(DisaggOverrideRequest): # Internal switch: sync API sets this True to return image from memory only. prefer_memory_result: bool = Field(default=False, exclude=True) - def __init__(self, **data): - super().__init__(**data) - if not self.save_result_path: - self.save_result_path = f"{self.task_id}" - def get(self, key, default=None): return getattr(self, key, default) class VideoTaskRequest(BaseTaskRequest): - target_video_length: int = Field(81, description="Target video length") + target_video_length: Optional[int] = Field( + None, + validation_alias=AliasChoices("num_frames", "target_video_length"), + description="Number of output frames; defaults to the startup config", + ) reuse_prefix_segments: int = Field( 0, ge=0, description="Number of prefix segments to reuse from the previous successful request", ) - video_path: str = Field("", description="Input video path (for SR/V2V-like tasks)") + video_path: str = Field("", description="Server-local input video path; comma-separated paths for multiple references") sr_ratio: float = Field(2.0, gt=0, description="Super-resolution scale factor used when target_shape is not set") audio_path: str = Field("", description="Input audio path (Wan-Audio)") - video_duration: int = Field(5, description="Video duration (Wan-Audio)") + video_duration: float = Field(5, description="Video duration in seconds (Wan-Audio)") talk_objects: Optional[list[TalkObject]] = Field(None, description="Talk objects (Wan-Audio)") - target_fps: Optional[int] = Field(16, description="Target FPS for video frame interpolation (overrides config)") - resize_mode: Optional[str] = Field("adaptive", description="Resize mode (adaptive, keep_ratio_fixed_area, fixed_min_area, fixed_max_area, fixed_shape, fixed_min_side)") + src_ref_images: list[str] = Field(default_factory=list, description="VACE reference images as base64, URL, or server-local paths") + prompt_ref: Optional[str] = Field(None, description="Prompt describing the driving video for Wan-Animate-2") + src_pose_path: Optional[str] = Field(None, description="Server-local pose video path") + src_face_path: Optional[str] = Field(None, description="Server-local face video path") + src_bg_path: Optional[str] = Field(None, description="Server-local background video path") + mask_path: Optional[str] = Field(None, description="Server-local mask video path for VACE or animation replacement") + pose: str | dict[str, Any] | None = Field(None, description="Action string, server-local pose JSON path, or WorldPlay pose object") + action_path: Optional[str] = Field(None, description="Server-local action file or control directory") + state_path: Optional[str] = Field(None, description="Server-local robot state file") + action_mode: Optional[str] = Field(None, description="Cosmos3 action mode") + domain_name: Optional[str] = Field(None, description="Cosmos3 embodiment domain") + view_point: Optional[str] = Field(None, description="Cosmos3 viewpoint") + save_action_path: Optional[str] = Field(None, description="Action output path, resolved by the runner") + image_strength: float | list[float] | None = Field(None, description="LTX image conditioning strength, shared or per image") + image_frame_idx: Optional[list[int]] = Field(None, description="LTX pixel frame index for each conditioning image") + reference_video_strength: Optional[float] = Field(None, description="LTX reference-video conditioning strength") + reference_video_frame_cap: Optional[int] = Field(None, description="Maximum number of LTX reference-video frames") + mux_audio_video_path: Optional[str] = Field(None, description="Server-local media file whose audio is muxed into the saved LTX video") class ImageTaskRequest(BaseTaskRequest): aspect_ratio: str = Field("16:9", description="Output aspect ratio") i2i_denoise_strength: Optional[float] = Field(None, description="Single-image I2I edit denoising strength in [0.0, 1.0]; omit to keep existing behavior") + inpaint_blur_sigma: Optional[float] = Field(None, description="Flux2 inpainting mask blur sigma") + inpaint_blur_size: Optional[int] = Field(None, description="Flux2 inpainting mask blur kernel size") sr_ratio: float = Field(2.0, gt=0, description="Super-resolution scale factor used when target_shape is not set") + keep_original_aspect: Optional[bool] = Field(None, description="Preserve a single HiDream reference image's aspect ratio") + layout_bboxes: Optional[str] = Field(None, description="HiDream layout boxes as a JSON string or server-local JSON file path") + infer_align_image_size: Optional[bool] = Field(None, description="Align HunyuanImage3 reference image sizes during inference") class SenseNovaVisionTaskRequest(BaseModel): """One request for the multi-task SenseNova-Vision service.""" + model_config = ConfigDict(extra="forbid") + task_id: str = Field(default_factory=generate_task_id, description="Task ID (auto-generated)") task: str = Field(..., description="Public SenseNova-Vision task name") prompt: str = Field("", description="Task prompt or question") @@ -104,8 +113,7 @@ class SenseNovaVisionTaskRequest(BaseModel): default_factory=list, description="Input images as base64/data URLs, HTTP(S) URLs, or server-local paths", ) - seed: int = Field(default_factory=generate_random_seed, description="Random seed") - target_shape: list[int] = Field(default_factory=list, description="Optional output [height, width]") + seed: Optional[int] = Field(None, description="Non-negative seed; omitted or null defaults to 42") visualize: bool = Field(True, description="Generate an official-style visualization when supported") postprocess_3d: bool = Field(False, description="Generate a GLB scene for recon3d") @@ -148,12 +156,11 @@ class SenseNovaVisionGenerationResponse(BaseModel): class TaskRequest(BaseTaskRequest): - target_video_length: int = Field(81, description="Target video length (video only)") + target_video_length: Optional[int] = Field(None, description="Target video length; defaults to the startup config") audio_path: str = Field("", description="Input audio path (Wan-Audio)") video_duration: int = Field(5, description="Video duration (Wan-Audio)") talk_objects: Optional[list[TalkObject]] = Field(None, description="Talk objects (Wan-Audio)") aspect_ratio: str = Field("16:9", description="Output aspect ratio (T2I only)") - target_fps: Optional[int] = Field(16, description="Target FPS for video frame interpolation (overrides config)") class TaskStatusMessage(BaseModel): @@ -163,7 +170,7 @@ class TaskStatusMessage(BaseModel): class TaskResponse(BaseModel): task_id: str task_status: str - save_result_path: str + save_result_path: Optional[str] # Filled after image generation in-process; never serialized in JSON responses. result_png: Optional[bytes] = Field(default=None, exclude=True) usage: Optional[Usage] = Field(default=None, exclude=True) diff --git a/lightx2v/server/services/generation/base.py b/lightx2v/server/services/generation/base.py index 76331ab58..461716ec5 100644 --- a/lightx2v/server/services/generation/base.py +++ b/lightx2v/server/services/generation/base.py @@ -118,23 +118,28 @@ async def _process_talk_objects(self, talk_objects: list, task_data: Dict[str, A with open(config_path, "w") as f: json.dump({"talk_objects": task_data["talk_objects"]}, f) - def _prepare_output_path(self, save_result_path: str, task_data: Dict[str, Any]) -> None: - actual_save_path = self.file_service.get_output_path(save_result_path) - if not actual_save_path.suffix: - actual_save_path = actual_save_path.with_suffix(self.get_output_extension()) - task_data["save_result_path"] = str(actual_save_path) + def prepare_task_data(self, message): + task_data = {field: getattr(message, field) for field in message.model_fields_set} + task_data["task_id"] = message.task_id + output_path = task_data.get("save_result_path") + if output_path: + actual_save_path = self.file_service.get_output_path(output_path) + if not actual_save_path.suffix: + actual_save_path = actual_save_path.with_suffix(self.get_output_extension()) + task_data["save_result_path"] = str(actual_save_path) + else: + task_data["save_result_path"] = None + return task_data async def generate_with_stop_event(self, message: Any, stop_event) -> Optional[Any]: try: - task_data = {field: getattr(message, field) for field in message.model_fields_set if field != "task_id"} - task_data["task_id"] = message.task_id - task_data["target_shape"] = message.target_shape + task_data = self.prepare_task_data(message) if stop_event.is_set(): logger.info(f"Task {message.task_id} cancelled before processing") return None - if hasattr(message, "image_path") and message.image_path: + if message.image_path: task_data["image_path"] = await self._resolve_image_path(message.image_path) logger.info(f"Task {message.task_id} image path: {task_data.get('image_path')}") @@ -142,6 +147,12 @@ async def generate_with_stop_event(self, message: Any, stop_event) -> Optional[A task_data["last_frame_path"] = await self._resolve_image_path(message.last_frame_path) logger.info(f"Task {message.task_id} last frame path: {task_data.get('last_frame_path')}") + reference_images = getattr(message, "src_ref_images", None) + if reference_images: + reference_image_paths = [await self._resolve_image_path(image) for image in reference_images] + task_data["src_ref_images"] = ",".join(reference_image_paths) + logger.info(f"Task {message.task_id} reference image paths: {task_data['src_ref_images']}") + if hasattr(message, "image_mask_path") and message.image_mask_path: task_data["image_mask_path"] = await self._resolve_image_path(message.image_mask_path) logger.info(f"Task {message.task_id} image mask path: {task_data.get('image_mask_path')}") @@ -155,9 +166,10 @@ async def generate_with_stop_event(self, message: Any, stop_event) -> Optional[A if hasattr(message, "talk_objects") and message.talk_objects: await self._process_talk_objects(message.talk_objects, task_data) - self._prepare_output_path(message.save_result_path, task_data) - task_data["seed"] = message.seed - task_data["resize_mode"] = message.resize_mode + task_data.pop("image_mask_path", None) + task_data.pop("talk_objects", None) + task_data.pop("prefer_memory_result", None) + task_data.pop("presigned_url", None) result = await self.inference_service.submit_task_async(task_data) @@ -168,13 +180,11 @@ async def generate_with_stop_event(self, message: Any, stop_event) -> Optional[A raise RuntimeError("Task processing failed") if result.get("status") == "success": - actual_save_path = self.file_service.get_output_path(message.save_result_path) - if not actual_save_path.suffix: - actual_save_path = actual_save_path.with_suffix(self.get_output_extension()) + output_path = result["save_result_path"] return TaskResponse( task_id=message.task_id, task_status="completed", - save_result_path=actual_save_path.name, + save_result_path=str(Path(output_path).absolute()) if output_path is not None else None, ) else: error_msg = result.get("error", "Inference failed") diff --git a/lightx2v/server/services/generation/image.py b/lightx2v/server/services/generation/image.py index 0c0bad004..50f650774 100644 --- a/lightx2v/server/services/generation/image.py +++ b/lightx2v/server/services/generation/image.py @@ -1,3 +1,4 @@ +from pathlib import Path from typing import Any, Optional from loguru import logger @@ -15,14 +16,7 @@ def get_task_type(self) -> str: async def generate_with_stop_event(self, message: Any, stop_event) -> Optional[Any]: try: - task_data = {field: getattr(message, field) for field in message.model_fields_set if field != "task_id"} - task_data["task_id"] = message.task_id - task_data["target_shape"] = message.target_shape - - if hasattr(message, "aspect_ratio"): - task_data["aspect_ratio"] = message.aspect_ratio - if hasattr(message, "i2i_denoise_strength"): - task_data["i2i_denoise_strength"] = message.i2i_denoise_strength + task_data = self.prepare_task_data(message) if stop_event.is_set(): logger.info(f"Task {message.task_id} cancelled before processing") @@ -38,12 +32,12 @@ async def generate_with_stop_event(self, message: Any, stop_event) -> Optional[A self._pack_image_and_mask_as_dir(task_data) logger.info(f"Task {message.task_id} packed image+mask dir: {task_data.get('image_path')}") - self._prepare_output_path(message.save_result_path, task_data) - task_data["seed"] = message.seed + task_data.pop("image_mask_path", None) prefer_memory_result = bool(getattr(message, "prefer_memory_result", False)) task_data.pop("prefer_memory_result", None) task_data.pop("presigned_url", None) - task_data["return_result_tensor"] = prefer_memory_result + if prefer_memory_result: + task_data["return_result_tensor"] = True result = await self.inference_service.submit_task_async(task_data) @@ -54,9 +48,6 @@ async def generate_with_stop_event(self, message: Any, stop_event) -> Optional[A raise RuntimeError("Task processing failed") if result.get("status") == "success": - actual_save_path = self.file_service.get_output_path(message.save_result_path) - if not actual_save_path.suffix: - actual_save_path = actual_save_path.with_suffix(self.get_output_extension()) if prefer_memory_result: result_png = result.get("result_png") if not result_png: @@ -65,15 +56,16 @@ async def generate_with_stop_event(self, message: Any, stop_event) -> Optional[A return TaskResponse( task_id=message.task_id, task_status="completed", - save_result_path="", + save_result_path=None, result_png=result_png, usage=usage, ) + output_path = result["save_result_path"] return TaskResponse( task_id=message.task_id, task_status="completed", - save_result_path=actual_save_path.name, + save_result_path=str(Path(output_path).absolute()) if output_path is not None else None, ) else: error_msg = result.get("error", "Inference failed") diff --git a/lightx2v/server/services/generation/sensenova_vision.py b/lightx2v/server/services/generation/sensenova_vision.py index 508bf78a0..98c7c5b3c 100644 --- a/lightx2v/server/services/generation/sensenova_vision.py +++ b/lightx2v/server/services/generation/sensenova_vision.py @@ -103,9 +103,6 @@ def validate_sensenova_request(message: SenseNovaVisionTaskRequest) -> tuple[str raise ValueError(f"SenseNova-Vision task={task!r} requires {expected} input image(s), got {image_count}.") if spec.requires_prompt and not str(message.prompt or "").strip(): raise ValueError(f"SenseNova-Vision task={task!r} requires a non-empty prompt.") - if message.target_shape: - if len(message.target_shape) != 2 or any(int(value) <= 0 for value in message.target_shape): - raise ValueError("SenseNova-Vision target_shape must be [height, width] with two positive integers.") if message.postprocess_3d and task != "recon3d": raise ValueError("postprocess_3d is only valid for task='recon3d'.") mode = spec.mode @@ -290,8 +287,6 @@ async def generate_with_stop_event(self, message: Any, stop_event) -> Optional[A "task_id": message.task_id, "prompt": prompt, "image_path": ",".join(image_paths), - "seed": int(message.seed), - "target_shape": list(message.target_shape), "save_result_path": str(save_result_path), "omni_vision_subtask": task, "raw_output_path": str(raw_output_path) if raw_output_path else "", @@ -300,6 +295,8 @@ async def generate_with_stop_event(self, message: Any, stop_event) -> Optional[A "return_result_tensor": False, "_return_pipeline_result": True, } + if message.seed is not None: + task_data["seed"] = message.seed inference_result = await self.inference_service.submit_task_async(task_data) if inference_result is None: diff --git a/lightx2v/server/services/inference/worker.py b/lightx2v/server/services/inference/worker.py index 84c993403..e714bd4fd 100644 --- a/lightx2v/server/services/inference/worker.py +++ b/lightx2v/server/services/inference/worker.py @@ -7,9 +7,8 @@ import torch from loguru import logger -from lightx2v.infer import init_runner -from lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict -from lightx2v.utils.set_config import set_config, set_parallel_config +from lightx2v.models.runners.runner_factory import build_runner +from lightx2v.utils.set_config import build_startup_config, init_parallel from ..distributed_utils import DistributedManager from .pipeline_image_encode import encode_pipeline_return_to_png_bytes @@ -46,19 +45,24 @@ def init(self, args) -> bool: else: logger.info(f"LoRA directory set to: {self.lora_dir}") - config = set_config(args) + config = build_startup_config( + { + "config_json": args.config_json, + "model_cls": args.model_cls, + "model_path": args.model_path, + "task": args.task, + } + ) if config["parallel"]: - set_parallel_config(config) + init_parallel(config) if self.rank == 0: logger.info(f"Config:\n {config}") - self.runner = init_runner(config) + self.runner = build_runner(config) logger.info(f"Rank {self.rank}/{self.world_size - 1} initialization completed") - self.input_info = init_empty_input_info(args.task) - return True except Exception as e: @@ -71,11 +75,13 @@ async def process_request(self, task_data: Dict[str, Any]) -> Dict[str, Any]: error_type = "" pipeline_return = None return_pipeline_result = False + task_id = task_data.get("task_id", "unknown") try: if self.world_size > 1 and self.rank == 0: task_data = self.dist_manager.broadcast_task_data(task_data) + task_id = task_data.pop("task_id", "unknown") return_pipeline_result = bool(task_data.pop("_return_pipeline_result", False)) # Handle dynamic LoRA loading @@ -84,30 +90,12 @@ async def process_request(self, task_data: Dict[str, Any]) -> Dict[str, Any]: reuse = task_data.pop("reuse", False) reuse_prefix_segments = task_data.pop("reuse_prefix_segments", 0) + input_info = self.runner.prepare_request(task_data) if self.lora_dir: self.switch_lora(lora_name, lora_strength) - task_data["task"] = self.runner.config["task"] - task_data["return_result_tensor"] = bool(task_data.get("return_result_tensor", False)) - task_data["negative_prompt"] = task_data.get("negative_prompt", "") - - target_fps = task_data.pop("target_fps", None) - if target_fps is not None: - vfi_cfg = self.runner.config.get("video_frame_interpolation") - if vfi_cfg: - task_data["video_frame_interpolation"] = {**vfi_cfg, "target_fps": target_fps} - else: - logger.warning(f"Target FPS {target_fps} is set, but video frame interpolation is not configured") - - if self.runner.config.get("model_cls") in {"sensenova_vision", "swiftvr"}: - # These services accept different input modes through one runner. - # Recreate request state so fields cannot leak between requests. - self.input_info = init_empty_input_info(self.runner.config["task"]) - update_input_info_from_dict(self.input_info, task_data) - self.runner.set_reuse(reuse, reuse_prefix_segments) - self.runner.set_config(task_data) - pipeline_return = self.runner.run_pipeline(self.input_info) + pipeline_return = self.runner.run_request(input_info) await asyncio.sleep(0) @@ -123,7 +111,7 @@ async def process_request(self, task_data: Dict[str, Any]) -> Dict[str, Any]: if self.rank == 0: if has_error: return { - "task_id": task_data.get("task_id", "unknown"), + "task_id": task_id, "status": "failed", "error": error_msg, "error_type": error_type, @@ -131,7 +119,7 @@ async def process_request(self, task_data: Dict[str, Any]) -> Dict[str, Any]: } else: out: Dict[str, Any] = { - "task_id": task_data["task_id"], + "task_id": task_id, "status": "success", "save_result_path": task_data.get("save_result_path"), "message": "Inference completed", @@ -140,7 +128,7 @@ async def process_request(self, task_data: Dict[str, Any]) -> Dict[str, Any]: encode_start = time.perf_counter() png = encode_pipeline_return_to_png_bytes(pipeline_return) encode_elapsed_ms = (time.perf_counter() - encode_start) * 1000 - logger.info(f"Task {task_data.get('task_id')} encode result_png cost {encode_elapsed_ms:.2f} ms") + logger.info(f"Task {task_id} encode result_png cost {encode_elapsed_ms:.2f} ms") if png: out["result_png"] = png usage = self.runner.compute_usage( diff --git a/lightx2v/shot_runner/rs2v_infer.py b/lightx2v/shot_runner/rs2v_infer.py index 26a28e4cb..ad0bdea95 100755 --- a/lightx2v/shot_runner/rs2v_infer.py +++ b/lightx2v/shot_runner/rs2v_infer.py @@ -10,7 +10,7 @@ from lightx2v.shot_runner.shot_base import ShotPipeline, load_clip_configs from lightx2v.shot_runner.utils import RS2V_SlidingWindowReader, save_audio, save_to_video from lightx2v.utils.audio_io import load_audio_file -from lightx2v.utils.input_info import UNSET, calculate_target_video_length_from_duration, init_input_info_from_args +from lightx2v.utils.input_info import calculate_target_video_length_from_duration from lightx2v.utils.profiler import * from lightx2v.utils.utils import is_main_process, seed_all, vae_to_comfyui_image, vae_to_comfyui_image_inplace from lightx2v.utils.va_controller import VAController @@ -86,10 +86,10 @@ def _calc_total_clips(total_samples, audio_per_frame, target_video_length): @staticmethod def _update_latent_shape(clip_input_info, target_len, vae_stride): - if hasattr(clip_input_info, "latent_shape") and clip_input_info.latent_shape is not None: - s = clip_input_info.latent_shape + shape = clip_input_info.latent_shape + if shape: new_t = (target_len - 1) // vae_stride + 1 - clip_input_info.latent_shape = [s[0], new_t, s[2], s[3]] + clip_input_info.latent_shape = [shape[0], new_t, shape[2], shape[3]] def _compute_segment_params(self, idx, audio_clip, pad_len, target_video_length, target_fps, audio_per_frame, vae_stride, clip_input_info): """Compute per-segment parameters (target_video_length, latent_shape, trimmed audio_clip). @@ -160,16 +160,16 @@ def generate(self, args): audio_per_frame = audio_sr // target_fps vae_stride = rs2v.config["vae_stride"][0] - clip_input_info = init_input_info_from_args(rs2v.config["task"], args) - clip_input_info = self.check_input_info(clip_input_info, rs2v.config) + clip_input_info = self.prepare_input_info(args, rs2v.config) + clip_input_info.seed = rs2v.resolve_request_seed({"seed": clip_input_info.seed}) + seed_all(clip_input_info.seed) - if clip_input_info.target_video_length is None or clip_input_info.target_video_length == UNSET: - if clip_input_info.video_duration is not None and clip_input_info.video_duration != UNSET: - segment_duration = min(clip_input_info.video_duration, 5.0) - clip_input_info.target_video_length = calculate_target_video_length_from_duration(segment_duration, target_fps) - logger.info(f"Auto-calculated target_video_length={clip_input_info.target_video_length} from video_duration={clip_input_info.video_duration}s (segment={segment_duration}s)") - else: - clip_input_info.target_video_length = rs2v.config.get("target_video_length", 81) + if getattr(args, "target_video_length", None) is None and clip_input_info.video_duration is not None: + segment_duration = min(clip_input_info.video_duration, 5.0) + clip_input_info.target_video_length = calculate_target_video_length_from_duration(segment_duration, target_fps) + logger.info(f"Auto-calculated target_video_length={clip_input_info.target_video_length} from video_duration={clip_input_info.video_duration}s (segment={segment_duration}s)") + elif clip_input_info.target_video_length is None: + clip_input_info.target_video_length = rs2v.config.get("target_video_length", 81) target_video_length = clip_input_info.target_video_length base_seed = clip_input_info.seed @@ -178,7 +178,7 @@ def generate(self, args): clip_input_info.audio_num = len(audio_files) # should set before _load_mask_latents (rs2v.process_single_mask) - # it will resize mask image by input_info.resize_mode and input_info.fixed_area + # it will resize mask image using config.resize_mode and input_info.fixed_area # otherwise repeat generate will use the old input_info and cause error rs2v.input_info = clip_input_info audio_array = self._load_audio_array(audio_files, audio_sr, clip_input_info.video_duration) @@ -259,7 +259,7 @@ def run_pipeline(self, input_info): def main(): parser = argparse.ArgumentParser() - parser.add_argument("--seed", type=int, default=42, help="The seed for random generator") + parser.add_argument("--seed", type=int, default=None, help="The seed for random generator") parser.add_argument("--config_json", type=str, required=True) parser.add_argument("--prompt", type=str, default="", help="The input prompt for text-to-video generation") parser.add_argument("--negative_prompt", type=str, default="") @@ -267,14 +267,12 @@ def main(): parser.add_argument("--audio_path", type=str, default="", help="The path to input audio file or directory for audio-to-video (s2v) task") parser.add_argument("--save_result_path", type=str, default=None, help="The path to save video path/file") parser.add_argument("--return_result_tensor", action="store_true", help="Whether to return result tensor. (Useful for comfyui)") - parser.add_argument("--target_shape", nargs="+", default=[], help="Set return video or image shape") - parser.add_argument("--infer_steps", type=int, default=4, help="Number of inference steps") + parser.add_argument("--target_shape", type=int, nargs="+", default=None, help="Set return video or image shape") parser.add_argument("--video_duration", type=float, default=20, help="Video duration in seconds") parser.add_argument("--stream_save_video", action="store_true", help="Whether to save video by stream") args = parser.parse_args() - seed_all(args.seed) clip_configs = load_clip_configs(args.config_json) with ProfilingContext4DebugL1("Init Pipeline Cost Time"): diff --git a/lightx2v/shot_runner/shot_base.py b/lightx2v/shot_runner/shot_base.py index cc09496a7..b973ad07c 100755 --- a/lightx2v/shot_runner/shot_base.py +++ b/lightx2v/shot_runner/shot_base.py @@ -1,16 +1,15 @@ import json -from argparse import Namespace -from dataclasses import dataclass +from dataclasses import dataclass, fields from pathlib import Path from typing import Any import torch from loguru import logger -from lightx2v.utils.input_info import fill_input_info_from_defaults +from lightx2v.models.runners.runner_factory import build_runner +from lightx2v.utils.input_info import UNSET, SekoTalkInputs from lightx2v.utils.profiler import * -from lightx2v.utils.registry_factory import RUNNER_REGISTER -from lightx2v.utils.set_config import print_config, set_config, set_parallel_config +from lightx2v.utils.set_config import build_startup_config, init_parallel, print_config from lightx2v_platform.registry_factory import PLATFORM_DEVICE_REGISTER @@ -52,12 +51,11 @@ def load_clip_configs(main_json_path): config = item["config"] else: config_json = str(Path(lightx2v_path) / item["path"]) - config_json = {"config_json": config_json} - config = set_config(Namespace(**config_json)) + config = build_startup_config(get_config_json(config_json)) if "parallel" in cfg: # Add parallel config to clip json config["parallel"] = cfg["parallel"] - set_parallel_config(config) + init_parallel(config) clip_configs.append(ClipConfig(name=item["name"], config_json=config)) return clip_configs @@ -76,11 +74,18 @@ def __init__(self, clip_configs: list[ClipConfig]): name = clip_config.name self.clip_generators[name] = self.create_clip_generator(clip_config) - def check_input_info(self, user_input_info, clip_config): - default_input_info = clip_config.get("default_input_info", None) - if default_input_info is not None: - fill_input_info_from_defaults(user_input_info, default_input_info) - return user_input_info.normalize_unset_to_none() + def prepare_input_info(self, args, clip_config): + task = clip_config["task"] + if task not in ("s2v", "rs2v"): + raise ValueError(f"Unsupported task: {task}") + + input_info = SekoTalkInputs(prompt="", negative_prompt="", seed=None, save_result_path=None) + input_info.update(clip_config) + input_info.update({key: value for key, value in vars(args).items() if value is not None}) + for input_field in fields(input_info): + if getattr(input_info, input_field.name) is UNSET: + setattr(input_info, input_field.name, None) + return input_info def _input_data_to_dict(self, input_data): if isinstance(input_data, dict): @@ -100,7 +105,7 @@ def update_input_info(self, input_data): setattr(self.shot_cfg, key, data[key]) for clip_input in self.clip_inputs.values(): - update_input_info_from_dict(clip_input, data) + clip_input.update(data) if hasattr(clip_input, "overlap_frame"): clip_input.overlap_frame = None if hasattr(clip_input, "overlap_latent"): @@ -108,24 +113,13 @@ def update_input_info(self, input_data): if hasattr(clip_input, "audio_clip"): clip_input.audio_clip = None - def _init_runner(self, config): - torch.set_grad_enabled(False) - runner = RUNNER_REGISTER[config["model_cls"]](config) - runner.init_modules() - return runner - - def set_config(self, config_modify): - for runner in self.clip_generators.values(): - if hasattr(runner, "set_config"): - runner.set_config(config_modify) - def set_progress_callback(self, callback): self.progress_callback = callback def create_clip_generator(self, clip_config: ClipConfig): logger.info(f"Clip {clip_config.name} initializing ... ") print_config(clip_config.config_json) - runner = self._init_runner(clip_config.config_json) + runner = build_runner(clip_config.config_json) logger.info(f"Clip {clip_config.name} initialized successfully!") return runner diff --git a/lightx2v/shot_runner/stream_infer.py b/lightx2v/shot_runner/stream_infer.py index ce0ebff32..2cc19ce23 100755 --- a/lightx2v/shot_runner/stream_infer.py +++ b/lightx2v/shot_runner/stream_infer.py @@ -10,7 +10,6 @@ from lightx2v.shot_runner.shot_base import ShotPipeline, load_clip_configs from lightx2v.shot_runner.utils import SlidingWindowReader, save_audio, save_to_video from lightx2v.utils.audio_io import load_audio_file -from lightx2v.utils.input_info import init_input_info_from_args from lightx2v.utils.profiler import * from lightx2v.utils.utils import seed_all @@ -28,12 +27,11 @@ def generate(self, args): model_fps = s2v.config.get("target_fps", 16) model_sr = s2v.config.get("audio_sr", 16000) - # 获取用户输入信息 - s2v_input_info = init_input_info_from_args(s2v.config["task"], args, infer_steps=3) - f2v_input_info = init_input_info_from_args(f2v.config["task"], args) - # 从默认配置中补全输入信息 - s2v_input_info = self.check_input_info(s2v_input_info, s2v.config) - f2v_input_info = self.check_input_info(f2v_input_info, f2v.config) + s2v_input_info = self.prepare_input_info(args, s2v.config) + f2v_input_info = self.prepare_input_info(args, f2v.config) + s2v_input_info.seed = s2v.resolve_request_seed({"seed": s2v_input_info.seed}) + f2v_input_info.seed = s2v_input_info.seed + seed_all(s2v_input_info.seed) assert s2v_input_info.audio_path == f2v_input_info.audio_path, "s2v and f2v must use the same audio input" @@ -68,9 +66,6 @@ def generate(self, args): inputs.audio_clip = audio_clip i = i + 1 - # if i % 4 == 0: - # inputs.infer_steps = 2#s2v 一半时间用2步推理 - if self.global_tail_video is not None: # 根据当前 pipe 需要多少 overlap_len 来裁剪 tail inputs.overlap_frame = self.global_tail_video[:, :, -pipe.prev_frame_length :] gen_clip_video, audio_clip, _ = pipe.run_clip_pipeline(inputs) @@ -84,18 +79,19 @@ def generate(self, args): gen_lvideo = torch.cat(gen_video_list, dim=2).float() gen_lvideo = torch.clamp(gen_lvideo, -1, 1) merge_audio = np.concatenate(cut_audio_list, axis=0).astype(np.float32) - out_path = os.path.join("./", "video_merge.mp4") - audio_file = os.path.join("./", "audio_merge.wav") + if s2v_input_info.save_result_path is not None: + out_path = os.path.join("./", "video_merge.mp4") + audio_file = os.path.join("./", "audio_merge.wav") - save_to_video(gen_lvideo, out_path, model_fps) - save_audio(merge_audio, audio_file, out_path, output_path=args.save_result_path) - os.remove(out_path) - os.remove(audio_file) + save_to_video(gen_lvideo, out_path, model_fps) + save_audio(merge_audio, audio_file, out_path, output_path=s2v_input_info.save_result_path) + os.remove(out_path) + os.remove(audio_file) def main(): parser = argparse.ArgumentParser() - parser.add_argument("--seed", type=int, default=42, help="The seed for random generator") + parser.add_argument("--seed", type=int, default=None, help="The seed for random generator") parser.add_argument("--config_json", type=str, required=True) parser.add_argument("--prompt", type=str, default="", help="The input prompt for text-to-video generation") parser.add_argument("--negative_prompt", type=str, default="") @@ -103,10 +99,9 @@ def main(): parser.add_argument("--audio_path", type=str, default="", help="The path to input audio file or directory for audio-to-video (s2v) task") parser.add_argument("--save_result_path", type=str, default=None, help="The path to save video path/file") parser.add_argument("--return_result_tensor", action="store_true", help="Whether to return result tensor. (Useful for comfyui)") - parser.add_argument("--target_shape", nargs="+", default=[], help="Set return video or image shape") + parser.add_argument("--target_shape", type=int, nargs="+", default=None, help="Set return video or image shape") args = parser.parse_args() - seed_all(args.seed) clip_configs = load_clip_configs(args.config_json) with ProfilingContext4DebugL1("Init Pipeline Cost Time"): diff --git a/lightx2v/utils/input_info.py b/lightx2v/utils/input_info.py index 10f524e96..fed9602df 100755 --- a/lightx2v/utils/input_info.py +++ b/lightx2v/utils/input_info.py @@ -1,5 +1,5 @@ -import inspect -from dataclasses import MISSING, dataclass, field, fields, make_dataclass +from collections.abc import Mapping +from dataclasses import dataclass, field, fields from typing import Any, Optional import torch @@ -14,84 +14,98 @@ def __repr__(self): @dataclass -class T2VInputInfo: - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) +class InputInfo: + """Mutable context shared by the runner, models, and schedulers for one inference. + + Startup defaults and request values initialize the context. The inference + pipeline then adds derived state to the same object. + + ``target_shape`` is always ``[height, width]`` in pixels. Model tensor + dimensions belong in ``latent_shape``. + """ + + task: str = "" + seed: int = 0 + save_result_path: Optional[str] = None + return_result_tensor: bool = False + + def update(self, values: Mapping[str, Any]) -> None: + for input_field in fields(self): + if input_field.name in values: + setattr(self, input_field.name, values[input_field.name]) + + +@dataclass +class T2VInputInfo(InputInfo): + prompt: str = "" + negative_prompt: str = "" # shape related - resize_mode: str = field(default_factory=str) + target_video_length: Optional[int] = None latent_shape: list = field(default_factory=list) target_shape: list = field(default_factory=list) @dataclass -class I2VInputInfo: - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - image_path: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) +class I2VInputInfo(InputInfo): + prompt: str = "" + negative_prompt: str = "" + image_path: str = "" # shape related - resize_mode: str = field(default_factory=str) + target_video_length: Optional[int] = None original_shape: list = field(default_factory=list) resized_shape: list = field(default_factory=list) latent_shape: list = field(default_factory=list) target_shape: list = field(default_factory=list) - # WorldPlay-specific: pose/action conditioning (optional) - pose: str = field(default_factory=lambda: None) - # Lingbot i2v camera/action conditioning (optional) - action_path: str = field(default_factory=str) @dataclass -class SRInputInfo: - seed: int = field(default_factory=int) - image_path: str = field(default_factory=str) # Single image input - video_path: str = field(default_factory=str) # Video input for SR - sr_ratio: float = field(default_factory=lambda: 2.0) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) +class ActionI2VInputInfo(I2VInputInfo): + pose: Optional[str] = None + action_path: str = "" + + +@dataclass +class MotusInputInfo(I2VInputInfo): + state_path: str = "" + save_action_path: str = "" + + +@dataclass +class SRInputInfo(InputInfo): + image_path: str = "" # Single image input + video_path: str = "" # Video input for SR + sr_ratio: float = 2.0 # shape related - resize_mode: str = field(default_factory=str) original_shape: list = field(default_factory=list) resized_shape: list = field(default_factory=list) latent_shape: list = field(default_factory=list) target_shape: list = field(default_factory=list) + output_fps: Optional[float] = field(default=None, repr=False) @dataclass -class Flf2vInputInfo: - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - image_path: str = field(default_factory=str) - last_frame_path: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) +class Flf2vInputInfo(InputInfo): + prompt: str = "" + negative_prompt: str = "" + image_path: str = "" + last_frame_path: str = "" # shape related - resize_mode: str = field(default_factory=str) + target_video_length: Optional[int] = None original_shape: list = field(default_factory=list) resized_shape: list = field(default_factory=list) latent_shape: list = field(default_factory=list) target_shape: list = field(default_factory=list) -# Need Check @dataclass -class VaceInputInfo: - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - src_ref_images: str = field(default_factory=str) - src_video: str = field(default_factory=str) - src_mask: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) +class VaceInputInfo(InputInfo): + prompt: str = "" + negative_prompt: str = "" + src_ref_images: Optional[str] = None + video_path: Optional[str] = None + mask_path: Optional[str] = None # shape related - resize_mode: str = field(default_factory=str) + target_video_length: Optional[int] = None original_shape: list = field(default_factory=list) resized_shape: list = field(default_factory=list) latent_shape: list = field(default_factory=list) @@ -99,88 +113,75 @@ class VaceInputInfo: @dataclass -class S2VInputInfo: - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - image_path: str = field(default_factory=str) - video_path: str = field(default_factory=str) - src_video: str = field(default_factory=str) - audio_path: str = field(default_factory=str) - src_pose_path: str = field(default_factory=str) - audio_num: int = field(default_factory=int) - with_mask: bool = field(default_factory=lambda: False) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) +class S2VInputInfo(InputInfo): + prompt: str = "" + negative_prompt: str = "" + image_path: str = "" + video_path: str = "" + audio_path: str = "" + src_pose_path: str = "" + audio_num: int = 0 + with_mask: bool = False stream_config: dict = field(default_factory=dict) # shape related - resize_mode: str = field(default_factory=str) - fixed_area: str = field(default_factory=str) + fixed_area: str = "" original_shape: list = field(default_factory=list) resized_shape: list = field(default_factory=list) latent_shape: list = field(default_factory=list) target_shape: list = field(default_factory=list) - target_video_length: int = field(default_factory=int) + target_video_length: Optional[int] = None + video_duration: Optional[float] = None # prev info - overlap_frame: torch.Tensor = field(default_factory=lambda: None) - overlap_latent: torch.Tensor = field(default_factory=lambda: None) + overlap_frame: Optional[torch.Tensor] = None + overlap_latent: Optional[torch.Tensor] = None # input preprocess audio - audio_clip: torch.Tensor = field(default_factory=lambda: None) + audio_clip: Optional[torch.Tensor] = None @dataclass -class RS2VInputInfo: - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - image_path: str = field(default_factory=str) - audio_path: str = field(default_factory=str) - audio_num: int = field(default_factory=int) - with_mask: bool = field(default_factory=lambda: False) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) +class RS2VInputInfo(InputInfo): + prompt: str = "" + negative_prompt: str = "" + image_path: str = "" + audio_path: str = "" + audio_num: int = 0 + with_mask: bool = False stream_config: dict = field(default_factory=dict) # shape related - resize_mode: str = field(default_factory=str) original_shape: list = field(default_factory=list) resized_shape: list = field(default_factory=list) latent_shape: list = field(default_factory=list) target_shape: list = field(default_factory=list) - target_video_length: int = field(default_factory=int) + target_video_length: Optional[int] = None + video_duration: Optional[float] = None # prev info - overlap_frame: torch.Tensor = field(default_factory=lambda: None) - overlap_latent: torch.Tensor = field(default_factory=lambda: None) + overlap_frame: Optional[torch.Tensor] = None + overlap_latent: Optional[torch.Tensor] = None # input preprocess audio - audio_clip: torch.Tensor = field(default_factory=lambda: None) + audio_clip: Optional[torch.Tensor] = None # input reference state - ref_state: int = field(default_factory=int) + ref_state: int = 0 # flags for first and last clip - is_first: bool = field(default_factory=lambda: False) - is_last: bool = field(default_factory=lambda: False) - - -# Need Check -@dataclass -class AnimateInputInfo: - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - prompt_ref: str = field(default_factory=lambda: "人物动作的参考视频") - negative_prompt: str = field(default_factory=str) - image_path: str = field(default_factory=str) - src_pose_path: str = field(default_factory=str) - src_face_path: str = field(default_factory=str) - src_ref_images: str = field(default_factory=str) - video_path: str = field(default_factory=str) - src_bg_path: str = field(default_factory=str) - src_mask_path: str = field(default_factory=str) - # None: use config_json replace_flag; True/False: per-request (e.g. worker frontend) - replace_flag: Optional[bool] = None - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) + is_first: bool = False + is_last: bool = False + + +@dataclass +class AnimateInputInfo(InputInfo): + prompt: str = "" + prompt_ref: str = "人物动作的参考视频" + negative_prompt: str = "" + image_path: str = "" + src_pose_path: str = "" + src_face_path: str = "" + src_ref_images: str = "" + video_path: str = "" + src_bg_path: str = "" + mask_path: str = "" # shape related - resize_mode: str = field(default_factory=str) + target_video_length: Optional[int] = None original_shape: list = field(default_factory=list) resized_shape: list = field(default_factory=list) latent_shape: list = field(default_factory=list) @@ -188,26 +189,30 @@ class AnimateInputInfo: @dataclass -class T2IInputInfo: - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) +class T2IInputInfo(InputInfo): + prompt: str = "" + negative_prompt: str = "" # shape related - resize_mode: str = field(default_factory=str) target_shape: list = field(default_factory=list) + latent_shape: list = field(default_factory=list) image_shapes: list = field(default_factory=list) txt_seq_lens: list = field(default_factory=list) # [postive_txt_seq_len, negative_txt_seq_len] - aspect_ratio: str = field(default_factory=str) + aspect_ratio: str = "" + latent_image_ids: Any = field(default=None, repr=False) + txt_ids: Optional[torch.Tensor] = field(default=None, repr=False) + revised_prompts: Any = field(default=None, repr=False) + + +@dataclass +class NeoppInputInfo(InputInfo): + seed: Optional[int] = 0 + target_shape: list = field(default_factory=list) + latent_shape: list = field(default_factory=list) @dataclass -class T2TInputInfo: - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) +class T2TInputInfo(InputInfo): + prompt: str = "" max_new_tokens: Optional[int] = None text_do_sample: Optional[bool] = None text_temperature: Optional[float] = None @@ -220,31 +225,40 @@ class T2TInputInfo: @dataclass class TI2TInputInfo(T2TInputInfo): - image_path: str = field(default_factory=str) + image_path: str = "" infer_align_image_size: Optional[bool] = None @dataclass -class I2IInputInfo: - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - image_path: str = field(default_factory=str) +class I2IInputInfo(InputInfo): + prompt: str = "" + negative_prompt: str = "" + image_path: str = "" i2i_denoise_strength: Optional[float] = None - inpaint_blur_size: Optional[int] = None - inpaint_blur_sigma: Optional[float] = None - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) - keep_original_aspect: bool = field(default_factory=lambda: False) - layout_bboxes: str = field(default_factory=str) # shape related - resize_mode: str = field(default_factory=str) target_shape: list = field(default_factory=list) + latent_shape: list = field(default_factory=list) image_shapes: list = field(default_factory=list) txt_seq_lens: list = field(default_factory=list) # [postive_txt_seq_len, negative_txt_seq_len] processed_image_size: list = field(default_factory=list) original_size: list = field(default_factory=list) - aspect_ratio: str = field(default_factory=str) + aspect_ratio: str = "" + image_encoder_output: Any = field(default=None, repr=False) + input_image: Any = field(default=None, repr=False) + latent_image_ids: Any = field(default=None, repr=False) + txt_ids: Optional[torch.Tensor] = field(default=None, repr=False) + + +@dataclass +class Flux2I2IInputInfo(I2IInputInfo): + inpaint_blur_size: Optional[int] = None + inpaint_blur_sigma: Optional[float] = None + + +@dataclass +class HidreamI2IInputInfo(I2IInputInfo): + keep_original_aspect: bool = False + layout_bboxes: str = "" @dataclass @@ -253,81 +267,72 @@ class TI2IInputInfo(I2IInputInfo): @dataclass -class T2AVInputInfo: - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) +class T2AVInputInfo(InputInfo): + prompt: str = "" + negative_prompt: str = "" # shape related - resize_mode: str = field(default_factory=str) + video_latent_shape: list = field(default_factory=list) audio_latent_shape: list = field(default_factory=list) latent_shape: list = field(default_factory=list) target_shape: list = field(default_factory=list) - target_video_length: int = field(default_factory=int) + target_video_length: Optional[int] = None @dataclass -class I2AVInputInfo: - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - image_path: str = field(default_factory=str) - image_strength: float = field(default_factory=float) +class I2AVInputInfo(InputInfo): + prompt: str = "" + negative_prompt: str = "" + image_path: str = "" + image_strength: float = 1.0 image_frame_idx: Optional[list[int]] = None - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) # shape related - resize_mode: str = field(default_factory=str) original_shape: list = field(default_factory=list) resized_shape: list = field(default_factory=list) + video_latent_shape: list = field(default_factory=list) + audio_latent_shape: list = field(default_factory=list) latent_shape: list = field(default_factory=list) target_shape: list = field(default_factory=list) - target_video_length: int = field(default_factory=int) + target_video_length: Optional[int] = None @dataclass class L2AVInputInfo(T2AVInputInfo): - last_frame_path: str = field(default_factory=str) + last_frame_path: str = "" @dataclass class FL2AVInputInfo(T2AVInputInfo): - image_path: str = field(default_factory=str) - last_frame_path: str = field(default_factory=str) + image_path: str = "" + last_frame_path: str = "" @dataclass class Ref2AVInputInfo(T2AVInputInfo): # Reuse the repository-wide media CLI. Comma-separated strings and Python # sequences are normalized by MiniMaxH3Runner. - image_path: Any = field(default_factory=str) - video_path: Any = field(default_factory=str) - audio_path: Any = field(default_factory=str) - - -@dataclass -class I2VAInputInfo: - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - image_path: str = field(default_factory=str) - video_path: str = field(default_factory=str) - action_path: str = field(default_factory=str) - state_path: str = field(default_factory=str) - action_mode: str = field(default_factory=str) - domain_name: str = field(default_factory=str) - view_point: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) - save_action_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) + image_path: Any = "" + video_path: Any = "" + audio_path: Any = "" + + +@dataclass +class I2VAInputInfo(InputInfo): + prompt: str = "" + negative_prompt: str = "" + image_path: str = "" + video_path: str = "" + action_path: str = "" + state_path: str = "" + action_mode: str = "" + domain_name: str = "" + view_point: str = "" + save_action_path: str = "" # shape related - resize_mode: str = field(default_factory=str) original_shape: list = field(default_factory=list) resized_shape: list = field(default_factory=list) latent_shape: list = field(default_factory=list) target_shape: list = field(default_factory=list) - target_video_length: int = field(default_factory=int) + target_video_length: Optional[int] = None # Optional in-memory policy inputs. Offline/CLI inference continues to use # image_path/state_path; long-running integrations (for example ROS) can # avoid writing a PNG and NPY file for every control step. @@ -336,7 +341,23 @@ class I2VAInputInfo: @dataclass -class V2AVInputInfo: +class Cosmos3InputInfo(I2VAInputInfo): + image_shapes: list = field(default_factory=list) + txt_seq_lens: list = field(default_factory=list) + audio_latent_shape: list = field(default_factory=list) + action_chunk_size: Optional[int] = None + vision_condition_latents: Any = field(default=None, repr=False) + vision_condition_frame_indexes: Optional[list[int]] = None + action_latents: Any = field(default=None, repr=False) + action_latent_shape: Optional[tuple[int, ...]] = None + action_condition_frame_indexes: Optional[list[int]] = None + action_domain_id: Optional[int] = None + raw_action_dim: Optional[int] = None + action_start_frame_offset: int = 1 + + +@dataclass +class V2AVInputInfo(I2AVInputInfo): """LTX-2.3 IC-LoRA video-to-audio-video. Drives both motion-transfer (Union / Pose / Motion-Track-Control) and @@ -346,152 +367,87 @@ class V2AVInputInfo: ``image_path`` / ``image_strength`` / ``image_frame_idx`` fields. """ - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - # Optional character / keyframe image conditioning (motion transfer). - image_path: str = field(default_factory=str) - image_strength: float = field(default_factory=float) - image_frame_idx: Optional[list[int]] = None # Pre-processed reference / control video (pose / canny / depth / track for # motion transfer, or the degraded source video for ICEdit). - video_path: str = field(default_factory=str) - action_path: str = field(default_factory=str) - action_mode: str = field(default_factory=str) - domain_name: str = field(default_factory=str) - view_point: str = field(default_factory=str) - reference_video_strength: float = field(default_factory=lambda: 1.0) + video_path: str = "" + reference_video_strength: float = 1.0 reference_video_frame_cap: Optional[int] = None # Optional: mux audio from this file after save (e.g. original driving video). # ``video_path`` is often a silent pose/canny/depth control clip; DefaultRunner's # v2av mux path is not used because LTX2Runner overrides ``process_images_after_vae_decoder``. - mux_audio_video_path: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) - save_action_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) - # shape related - resize_mode: str = field(default_factory=str) - original_shape: list = field(default_factory=list) - resized_shape: list = field(default_factory=list) - latent_shape: list = field(default_factory=list) - target_shape: list = field(default_factory=list) - target_video_length: int = field(default_factory=int) + mux_audio_video_path: str = "" @dataclass -class LTX2S2VInputInfo: +class LTX2S2VInputInfo(I2AVInputInfo): """LTX-2 audio-conditioned video (reference audio + optional reference images).""" - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - image_path: str = field(default_factory=str) - image_strength: float = field(default_factory=float) - image_frame_idx: Optional[list[int]] = None - audio_path: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) - resize_mode: str = field(default_factory=str) - original_shape: list = field(default_factory=list) - resized_shape: list = field(default_factory=list) - latent_shape: list = field(default_factory=list) - target_shape: list = field(default_factory=list) - target_video_length: int = field(default_factory=int) + audio_path: str = "" @dataclass -class WorldPlayI2VInputInfo: +class WorldPlayI2VInputInfo(I2VInputInfo): """Input info for WorldPlay model (image-to-video with action/pose conditioning).""" - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - image_path: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) - # shape related - resize_mode: str = field(default_factory=str) - original_shape: list = field(default_factory=list) - resized_shape: list = field(default_factory=list) - latent_shape: list = field(default_factory=list) - target_shape: list = field(default_factory=list) - # WorldPlay-specific: pose/action conditioning - pose: str = field(default_factory=str) # Pose string (e.g., "w-3, right-0.5") or JSON path - model_type: str = field(default_factory=lambda: "ar") # "ar" (autoregressive) or "bi" (bidirectional) - chunk_latent_frames: int = field(default_factory=lambda: 4) + pose: str | dict | None = None + model_type: str = "ar" # "ar" (autoregressive) or "bi" (bidirectional) + chunk_latent_frames: int = 4 # Computed pose tensors (set during processing) - viewmats: torch.Tensor = field(default_factory=lambda: None) - Ks: torch.Tensor = field(default_factory=lambda: None) - action: torch.Tensor = field(default_factory=lambda: None) + viewmats: Optional[torch.Tensor] = None + Ks: Optional[torch.Tensor] = None + action: Optional[torch.Tensor] = None @dataclass -class Hunyuan3DShapeInputInfo: +class Hunyuan3DShapeInputInfo(InputInfo): """Input info for Hunyuan3D-2.1 image-to-3D-mesh shape generation.""" - seed: int = field(default_factory=int) - image_path: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) + image_path: str = "" @dataclass -class WorldMirrorReconInputInfo: +class WorldMirrorReconInputInfo(InputInfo): """Input info for HY-WorldMirror-2.0 3D reconstruction. Unlike the diffusion tasks, this task takes a directory / video / image and saves multi-view depth / normal / Gaussian-splat results to disk. """ - seed: int = field(default_factory=int) # Input may be a directory of images, a single image, or a video. - input_path: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) # output root dir - strict_output_path: str = field(default_factory=lambda: None) - return_result_tensor: bool = field(default_factory=lambda: False) + input_path: str = "" + strict_output_path: Optional[str] = None # Optional priors - prior_cam_path: str = field(default_factory=lambda: None) - prior_depth_path: str = field(default_factory=lambda: None) + prior_cam_path: Optional[str] = None + prior_depth_path: Optional[str] = None + save_rendered: bool = False + render_interp_per_pair: int = 15 + render_depth: bool = False @dataclass -class WorldPlayT2VInputInfo: +class WorldPlayT2VInputInfo(T2VInputInfo): """Input info for WorldPlay model (text-to-video with action/pose conditioning).""" - seed: int = field(default_factory=int) - prompt: str = field(default_factory=str) - negative_prompt: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) - # shape related - resize_mode: str = field(default_factory=str) - latent_shape: list = field(default_factory=list) - target_shape: list = field(default_factory=list) - # WorldPlay-specific: pose/action conditioning - pose: str = field(default_factory=str) # Pose string (e.g., "w-3, right-0.5") or JSON path - model_type: str = field(default_factory=lambda: "ar") # "ar" (autoregressive) or "bi" (bidirectional) - chunk_latent_frames: int = field(default_factory=lambda: 4) + pose: str | dict | None = None + model_type: str = "ar" # "ar" (autoregressive) or "bi" (bidirectional) + chunk_latent_frames: int = 4 # Computed pose tensors (set during processing) - viewmats: torch.Tensor = field(default_factory=lambda: None) - Ks: torch.Tensor = field(default_factory=lambda: None) - action: torch.Tensor = field(default_factory=lambda: None) + viewmats: Optional[torch.Tensor] = None + Ks: Optional[torch.Tensor] = None + action: Optional[torch.Tensor] = None @dataclass -class SenseNovaVisionInputInfo: - seed: int = field(default_factory=lambda: 42) - prompt: str = field(default_factory=str) - image_path: str = field(default_factory=str) - save_result_path: str = field(default_factory=str) - return_result_tensor: bool = field(default_factory=lambda: False) - target_shape: list = field(default_factory=list) - omni_vision_subtask: str = field(default_factory=str) - raw_output_path: str = field(default_factory=str) - glb_output_path: str = field(default_factory=str) +class SenseNovaVisionInputInfo(InputInfo): + prompt: str = "" + image_path: str = "" + omni_vision_subtask: str = "" + raw_output_path: str = "" + glb_output_path: str = "" postprocess_predictions: Optional[bool] = None -task_dict = { +INPUT_INFO_TYPES = { "t2v": T2VInputInfo, "i2v": I2VInputInfo, "sr": SRInputInfo, @@ -513,53 +469,12 @@ class SenseNovaVisionInputInfo: "i2va": I2VAInputInfo, "v2av": V2AVInputInfo, "ltx2_s2v": LTX2S2VInputInfo, - "worldplay_i2v": WorldPlayI2VInputInfo, - "worldplay_t2v": WorldPlayT2VInputInfo, "recon": WorldMirrorReconInputInfo, "i23d": Hunyuan3DShapeInputInfo, "omni_vision_task": SenseNovaVisionInputInfo, } -def init_empty_input_info(task, support_tasks=[]): - if len(support_tasks) == 0: - support_tasks = [task] - # assert task in support_tasks, f"Task {task} not in support tasks {support_tasks}" - - if len(support_tasks) == 1: - support_task = support_tasks[0] - if support_task not in task_dict: - raise ValueError(f"Unsupported task: {support_task}") - return task_dict[support_task]() - - merged_fields = [] - merged_field_names = set() - - for support_task in support_tasks: - if support_task not in task_dict: - raise ValueError(f"Unsupported task: {support_task}") - - support_input_info_cls = task_dict[support_task] - for support_field in fields(support_input_info_cls): - if support_field.name in merged_field_names: - continue - merged_field_names.add(support_field.name) - - if support_field.default_factory is not MISSING: - merged_fields.append((support_field.name, support_field.type, field(default_factory=support_field.default_factory))) - elif support_field.default is not MISSING: - merged_fields.append((support_field.name, support_field.type, field(default=support_field.default))) - else: - merged_fields.append((support_field.name, support_field.type, field(default=None))) - - if not merged_fields: - raise ValueError("support_tasks must not be empty") - - merged_cls_name = "Merged" + "".join(task.upper() for task in support_tasks) + "InputInfo" - merged_input_info_cls = make_dataclass(merged_cls_name, merged_fields) - return merged_input_info_cls() - - def calculate_target_video_length_from_duration(duration_seconds: float, fps: int = 16) -> int: """Calculate target_video_length from video duration using the formula: target_video_length = (fps * seconds + 3) // 4 * 4 + 1 @@ -578,13 +493,16 @@ def calculate_target_video_length_from_duration(duration_seconds: float, fps: in 3s: (16*3 + 3) // 4 * 4 + 1 = 49 frames 5s: (16*5 + 3) // 4 * 4 + 1 = 81 frames """ - target_video_length = (int(fps * duration_seconds) + 3) // 4 * 4 + 1 - return target_video_length + return align_target_video_length(int(fps * duration_seconds) + 3, 4) + + +def align_target_video_length(num_frames: int, temporal_stride: int) -> int: + """Align a frame count so that ``num_frames - 1`` is stride-divisible.""" + return num_frames // temporal_stride * temporal_stride + 1 @dataclass -class SekoTalkInputs: - infer_steps: int | Any = UNSET +class SekoTalkInputs(InputInfo): target_video_length: int | Any = UNSET seed: int | Any = UNSET prompt: str | Any = UNSET @@ -598,15 +516,16 @@ class SekoTalkInputs: return_result_tensor: bool | Any = UNSET stream_config: dict | Any = UNSET - resize_mode: str | Any = UNSET fixed_area: str | Any = UNSET target_shape: list | Any = UNSET + latent_shape: list | Any = UNSET # prev info overlap_frame: torch.Tensor | Any = UNSET overlap_latent: torch.Tensor | Any = UNSET # input preprocess audio audio_clip: torch.Tensor | Any = UNSET + person_mask_latens: torch.Tensor | Any = field(default=UNSET, repr=False) # input reference state ref_state: int | Any = UNSET @@ -615,66 +534,3 @@ class SekoTalkInputs: is_last: bool | Any = UNSET # if save video by stream stream_save_video: bool | Any = UNSET - - @classmethod - def from_args(cls, args, **overrides): - """ - Build InputInfo from argparse.Namespace (or any object with __dict__) - Priority: - args < overrides - """ - field_names = {f.name for f in fields(cls)} - data = {k: v for k, v in vars(args).items() if k in field_names} - data.update(overrides) - return cls(**data) - - def normalize_unset_to_none(self): - """ - Replace all UNSET fields with None. - Call this right before running / inference. - """ - for f in fields(self): - if getattr(self, f.name) is UNSET: - setattr(self, f.name, None) - return self - - -def init_input_info_from_args(task, args, **overrides): - if task in ["s2v", "rs2v"]: - return SekoTalkInputs.from_args(args, **overrides) - else: - raise ValueError(f"Unsupported task: {task}") - - -def fill_input_info_from_defaults(input_info, defaults): - for key in input_info.__dataclass_fields__: - if key in defaults and getattr(input_info, key) is UNSET: - setattr(input_info, key, defaults[key]) - - -def update_input_info_from_dict(input_info, data): - for key in input_info.__dataclass_fields__: - if key in data: - setattr(input_info, key, data[key]) - - -def update_input_info_from_object(input_info, obj): - for key in input_info.__dataclass_fields__: - if hasattr(obj, key): - setattr(input_info, key, getattr(obj, key)) - - -def get_all_input_info_keys(): - all_keys = set() - - current_module = inspect.currentframe().f_globals - - for name, obj in current_module.items(): - if inspect.isclass(obj) and name.endswith("InputInfo") and hasattr(obj, "__dataclass_fields__"): - all_keys.update(obj.__dataclass_fields__.keys()) - - return all_keys - - -# 创建包含所有InputInfo字段的集合 -ALL_INPUT_INFO_KEYS = get_all_input_info_keys() diff --git a/lightx2v/utils/set_config.py b/lightx2v/utils/set_config.py index 86b04fce0..46aeb34b6 100755 --- a/lightx2v/utils/set_config.py +++ b/lightx2v/utils/set_config.py @@ -1,93 +1,58 @@ import json import os +from dataclasses import fields import torch import torch.distributed as dist from loguru import logger from torch.distributed.tensor.device_mesh import init_device_mesh -from lightx2v.utils.input_info import ALL_INPUT_INFO_KEYS +from lightx2v.utils.input_info import align_target_video_length from lightx2v.utils.lockable_dict import LockableDict from lightx2v.utils.utils import find_torch_model_path, is_main_process from lightx2v_platform.base.global_var import AI_DEVICE def get_default_config(): - default_config = { - "do_mm_calib": False, - "cpu_offload": False, - "max_area": False, - "vae_stride": (4, 8, 8), - "patch_size": (1, 2, 2), - "feature_caching": "NoCaching", # ["NoCaching", "TaylorSeer", "Tea"] - "teacache_thresh": 0.26, - "use_ret_steps": False, - "use_bfloat16": True, - "lora_configs": None, # List of dicts with 'path' and 'strength' keys - "parallel": False, - "seq_parallel": False, - "cfg_parallel": False, - "enable_cfg": False, - "warmup": False, - "use_image_encoder": True, - } - default_config = LockableDict(default_config) - return default_config - - -def validate_model_task_args(args): - """Validate model/task combinations before assembling the runtime config.""" - task = getattr(args, "task", None) - model_cls = getattr(args, "model_cls", None) - has_omni_vision_subtask = hasattr(args, "omni_vision_subtask") - omni_vision_subtask = getattr(args, "omni_vision_subtask", None) - - if task == "omni_vision_task": - if model_cls != "sensenova_vision": - raise ValueError("--task omni_vision_task requires --model_cls sensenova_vision") - # Offline inference exposes this argument and must select one subtask. - # The resident server omits it because each request chooses a subtask - # after the complete model has been loaded. - if has_omni_vision_subtask and not omni_vision_subtask: - raise ValueError("--omni_vision_subtask is required when --task omni_vision_task") - else: - if model_cls == "sensenova_vision": - raise ValueError("--model_cls sensenova_vision requires --task omni_vision_task") - if omni_vision_subtask: - raise ValueError("--omni_vision_subtask is only valid with --task omni_vision_task") - - -def set_args2config(args): - config = get_default_config() - config.update({k: v for k, v in vars(args).items() if k not in ALL_INPUT_INFO_KEYS and v is not None}) - - # Snapshot worldmirror-specific CLI flags so they can win over JSON - # defaults after auto_calc_config merges the JSON in. See the - # ``worldmirror`` branch of ``auto_calc_config`` for the replay logic. - _wm_cli_keys = ( - "subfolder", - "disable_heads", - "enable_bf16", - "save_rendered", - "render_interp_per_pair", - "render_depth", - "wm_config_path", - "wm_ckpt_path", + return LockableDict( + { + "do_mm_calib": False, + "cpu_offload": False, + "max_area": False, + "vae_stride": (4, 8, 8), + "patch_size": (1, 2, 2), + "feature_caching": "NoCaching", # ["NoCaching", "TaylorSeer", "Tea"] + "teacache_thresh": 0.26, + "use_ret_steps": False, + "use_bfloat16": True, + "lora_configs": None, # List of dicts with 'path' and 'strength' keys + "parallel": False, + "seq_parallel": False, + "cfg_parallel": False, + "enable_cfg": False, + "warmup": False, + "use_image_encoder": True, + } ) - config["_wm_cli_snapshot"] = {k: getattr(args, k, None) for k in _wm_cli_keys if hasattr(args, k)} - return config -def auto_calc_config(config): - cli_num_iterations = config.get("num_iterations", None) - if config.get("config_json", None) is not None: +def build_startup_config(config_data): + """Assemble startup settings from deployment and model configs.""" + config = get_default_config() + config.update(config_data) + if config.get("config_json") is not None: logger.info(f"Loading some config from {config['config_json']}") with open(config["config_json"], "r") as f: config_json = json.load(f) config.update(config_json) - if cli_num_iterations is not None: - config["num_iterations"] = cli_num_iterations + config["task"] = config_data["task"] + + load_model_config(config) + return config + +def load_model_config(config): + """Load model settings and normalize the startup configuration.""" if config.get("model_cls") == "ltx2_5": # Match Wan's checkpoint lookup contract: an explicit component path # wins; otherwise find the released filename below --model_path. @@ -118,27 +83,19 @@ def auto_calc_config(config): transformer_config.pop("rope_type", None) config.update(transformer_config) config["ltx_model_version"] = metadata.get("model_version", "") - except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc: + except (OSError, ValueError, TypeError) as exc: raise ValueError(f"Failed to read LTX-2.5 transformer metadata from {transformer_path}: {exc}") from exc assert os.path.exists(config["model_path"]), f"Model path not found: {config['model_path']}" - if config["model_cls"] == "hunyuan_video_1.5": # Special config for hunyuan video 1.5 model folder structure - config["transformer_model_path"] = os.path.join(config["model_path"], "transformer", config["transformer_model_name"]) # transformer_model_name: [480p_t2v, 480p_i2v, 720p_t2v, 720p_i2v] - if os.path.exists(os.path.join(config["transformer_model_path"], "config.json")): - with open(os.path.join(config["transformer_model_path"], "config.json"), "r") as f: - model_config = json.load(f) - config.update(model_config) - elif config["model_cls"] in ["worldplay_distill", "worldplay_ar", "worldplay_bi"]: # Special config for WorldPlay models + if config["model_cls"] in {"hunyuan_video_1.5", "worldplay_distill", "worldplay_ar", "worldplay_bi"}: config["transformer_model_path"] = os.path.join(config["model_path"], "transformer", config["transformer_model_name"]) if os.path.exists(os.path.join(config["transformer_model_path"], "config.json")): with open(os.path.join(config["transformer_model_path"], "config.json"), "r") as f: model_config = json.load(f) config.update(model_config) - elif config["model_cls"] == "hunyuan3d": - # Hunyuan3D shape loads hunyuan3d-dit-v2-1/config.yaml + checkpoint in the runner. - pass - elif config["model_cls"] == "hidream_o1_image": + elif config["model_cls"] in {"hunyuan3d", "hidream_o1_image"}: + # These runners load their own model configuration. pass elif config["model_cls"] == "sensenova_vision": llm_config_path = os.path.join(config["model_path"], "llm_config.json") @@ -186,34 +143,6 @@ def auto_calc_config(config): config["transformer_model_path"] = candidate else: config["transformer_model_path"] = config["model_path"] - - # Re-apply worldmirror-specific CLI overrides. ``auto_calc_config`` just - # finished merging the JSON, which will have clobbered any store_true / - # default-bearing CLI flags that the user set on the command line. We - # pull those specific keys back from the original args snapshot so that - # ``--enable_bf16`` / ``--save_rendered`` / ``--disable_heads`` etc. - # always win over JSON defaults. - # - # Caveat: ``argparse``'s ``store_true`` flags look the same (value - # ``False``) whether the user omitted them or explicitly wants False, - # so we only treat non-False, non-None values as explicit overrides. - # This matches the documented "CLI overrides JSON-true is not possible - # without a dedicated --no_xxx flag" behaviour. - _cli_snapshot = config.pop("_wm_cli_snapshot", None) - if isinstance(_cli_snapshot, dict): - for k, v in _cli_snapshot.items(): - if v is None or v is False: - continue - config[k] = v - # Translate the ``--wm_config_path`` / ``--wm_ckpt_path`` CLI - # aliases to the runner-visible keys, and don't leave the - # aliases lying around in the printed/locked config. - wm_config = config.pop("wm_config_path", None) - wm_ckpt = config.pop("wm_ckpt_path", None) - if wm_config: - config["config_path"] = wm_config - if wm_ckpt: - config["ckpt_path"] = wm_ckpt elif config["model_cls"] == "dreamzero": config_path = os.path.join(config["model_path"], "config.json") if os.path.exists(config_path): @@ -228,43 +157,26 @@ def auto_calc_config(config): config["target_video_length"] = config["num_frames"] if "out_dim" in config: config["num_channels_latents"] = config["out_dim"] - elif config["model_cls"] == "longcat_image": # Special config for longcat_image: load both root and transformer config - if os.path.exists(os.path.join(config["model_path"], "config.json")): - with open(os.path.join(config["model_path"], "config.json"), "r") as f: - model_config = json.load(f) - config.update(model_config) - if os.path.exists(os.path.join(config["model_path"], "transformer", "config.json")): - with open(os.path.join(config["model_path"], "transformer", "config.json"), "r") as f: - model_config = json.load(f) - config.update(model_config) - elif config["model_cls"] == "cosmos3": - transformer_config_path = os.path.join(config["model_path"], "transformer", "config.json") - if os.path.exists(transformer_config_path): - with open(transformer_config_path, "r") as f: - model_config = json.load(f) - config.update(model_config) - config.setdefault("target_video_length", 1) - config.setdefault("target_fps", config.get("base_fps", 24)) - config.setdefault("enable_cfg", True) - elif config["model_cls"] == "lingbot_video": + elif config["model_cls"] == "longcat_image": + for subfolder in ("", "transformer"): + config_path = os.path.join(config["model_path"], subfolder, "config.json") + if os.path.exists(config_path): + with open(config_path, "r") as f: + config.update(json.load(f)) + elif config["model_cls"] in {"cosmos3", "lingbot_video"}: transformer_config_path = os.path.join(config["model_path"], "transformer", "config.json") if os.path.exists(transformer_config_path): with open(transformer_config_path, "r") as f: model_config = json.load(f) config.update(model_config) config.setdefault("target_video_length", 1) - config.setdefault("target_fps", 24) - config.setdefault("enable_cfg", True) - config.setdefault("vae_stride", (4, 8, 8)) - config.setdefault("vae_scale_factor_spatial", 8) - config.setdefault("vae_scale_factor_temporal", 4) - config.setdefault("vae_scale_factor", 8) + config.setdefault("target_fps", config.get("base_fps", 24) if config["model_cls"] == "cosmos3" else 24) + if config["model_cls"] == "lingbot_video": + config.setdefault("vae_scale_factor_spatial", 8) + config.setdefault("vae_scale_factor_temporal", 4) + config.setdefault("vae_scale_factor", 8) elif config["model_cls"] == "minimax_h3": - supported_tasks = {"t2av", "i2av", "l2av", "fl2av", "ref2av"} - task = config.get("task") - if task not in supported_tasks: - raise ValueError(f"MiniMax-H3 supports {sorted(supported_tasks)}, got {task!r}") - transformer_subfolder = "transformer_ref" if task == "ref2av" else "transformer" + transformer_subfolder = "transformer_ref" if config["task"] == "ref2av" else "transformer" transformer_path = os.path.join(config["model_path"], transformer_subfolder) transformer_config_path = os.path.join(transformer_path, "config.json") if not os.path.isfile(transformer_config_path): @@ -281,40 +193,21 @@ def auto_calc_config(config): }.get(config.get("dit_quant_scheme"), config.get("dit_quant_scheme", "Default")) config["enable_cfg"] = False config["fps"] = 24 - config["vae_spatial_scale_factor"] = 16 - config["vae_scale_factor"] = 16 config.setdefault("video_flow_shift", 12.0) config.setdefault("audio_flow_shift", 3.0) config.setdefault("audio_sampling_rate", 32000) else: - if os.path.exists(os.path.join(config["model_path"], "config.json")): - with open(os.path.join(config["model_path"], "config.json"), "r") as f: - model_config = json.load(f) - if config["model_cls"] in ["ltx2", "ltx2_ar", "ltx2_5"]: - # LTX uses rope_type for the layout ("split"), while LightX2V - # uses it to select a registered RoPE implementation. - model_config.pop("rope_type", None) - config.update(model_config) - elif os.path.exists(os.path.join(config["model_path"], "low_noise_model", "config.json")): # 需要一个更优雅的update方法 - with open(os.path.join(config["model_path"], "low_noise_model", "config.json"), "r") as f: - model_config = json.load(f) - config.update(model_config) - elif os.path.exists(os.path.join(config["model_path"], "distill_models", "low_noise_model", "config.json")): # 需要一个更优雅的update方法 - with open(os.path.join(config["model_path"], "distill_models", "low_noise_model", "config.json"), "r") as f: - model_config = json.load(f) - config.update(model_config) - elif os.path.exists(os.path.join(config["model_path"], "original", "config.json")): - with open(os.path.join(config["model_path"], "original", "config.json"), "r") as f: - model_config = json.load(f) - config.update(model_config) - elif os.path.exists(os.path.join(config["model_path"], "transformer", "config.json")): - with open(os.path.join(config["model_path"], "transformer", "config.json"), "r") as f: + for subfolder in ("", "low_noise_model", "distill_models/low_noise_model", "original", "transformer"): + config_path = os.path.join(config["model_path"], subfolder, "config.json") + if not os.path.exists(config_path): + continue + with open(config_path, "r") as f: model_config = json.load(f) - if config["model_cls"] in ["ltx2", "ltx2_ar", "ltx2_5"]: - # Upstream LTX2 uses rope_type for the layout name ("split"), - # while LightX2V uses it as the registered RoPE implementation. + if config["model_cls"] in {"ltx2", "ltx2_ar", "ltx2_5"} and subfolder in ("", "transformer"): + # LTX uses rope_type for the layout, while LightX2V uses it + # to select a registered RoPE implementation. model_config.pop("rope_type", None) - elif config["model_cls"] == "z_image": + elif config["model_cls"] == "z_image" and subfolder == "transformer": # https://huggingface.co/Tongyi-MAI/Z-Image-Turbo/blob/main/transformer/config.json z_image_patch_size = model_config.pop("all_patch_size", [2]) z_image_f_patch_size = model_config.pop("all_f_patch_size", [1]) @@ -330,14 +223,45 @@ def auto_calc_config(config): model_config["f_patch_size"] = z_image_f_patch_size[0] config.update(model_config) + break # load quantized config - if config.get("dit_quantized_ckpt", None) is not None: + if config.get("dit_quantized_ckpt") is not None: config_path = os.path.join(config["dit_quantized_ckpt"], "config.json") if os.path.exists(config_path): with open(config_path, "r") as f: model_config = json.load(f) config.update(model_config) + vae_config_path = os.path.join(config["model_path"], "vae", "config.json") + if os.path.exists(vae_config_path): + with open(vae_config_path, "r") as f: + vae_config = json.load(f) + if "temperal_downsample" in vae_config: + config["vae_scale_factor"] = 2 ** len(vae_config["temperal_downsample"]) + elif "block_out_channels" in vae_config: + config["vae_scale_factor"] = 2 ** (len(vae_config["block_out_channels"]) - 1) + if config["model_cls"] == "ernie_image": + config["vae_scale_factor"] = 2 ** len(vae_config["block_out_channels"]) + elif config["model_cls"] == "cosmos3": + config["vae_scale_factor_spatial"] = int(vae_config.get("scale_factor_spatial", 16)) + config["vae_scale_factor_temporal"] = int(vae_config.get("scale_factor_temporal", 4)) + config["vae_scale_factor"] = config["vae_scale_factor_spatial"] + + if config["model_cls"] == "lingbot_video": + config["vae_scale_factor_spatial"] = int(config.get("vae_scale_factor_spatial", 8)) + config["vae_scale_factor_temporal"] = int(config.get("vae_scale_factor_temporal", 4)) + config["vae_scale_factor"] = config["vae_scale_factor_spatial"] + if config["model_cls"] == "minimax_h3": + # The generic Diffusers-VAE heuristic above counts six encoder stages + # and would incorrectly derive 32. H3 downsamples space by exactly 16. + config["vae_spatial_scale_factor"] = 16 + config["vae_scale_factor"] = 16 + if config["model_cls"] == "cosmos3" and os.path.exists(os.path.join(config["model_path"], "sound_tokenizer", "config.json")): + with open(os.path.join(config["model_path"], "sound_tokenizer", "config.json"), "r") as f: + sound_config = json.load(f) + config["sound_sampling_rate"] = int(sound_config.get("sampling_rate", 48000)) + config["sound_hop_size"] = int(sound_config.get("hop_size", 1920)) + # Some upstream/offical configs use `num_inference_steps`, while the shared # LightX2V scheduler stack expects `infer_steps`. if "infer_steps" not in config and "num_inference_steps" in config: @@ -361,110 +285,73 @@ def auto_calc_config(config): config["target_video_length"] = (latent_frames - 1) * temporal_stride + 1 logger.info(f"Auto-set LingBot-VA target_video_length={config['target_video_length']} from {latent_frames} latent frames and temporal stride {temporal_stride}.") - if config["model_cls"] != "minimax_h3" and config["task"] in ["i2v", "t2av", "i2av", "i2va", "s2v", "rs2v", "ltx2_s2v", "v2av"] and "target_video_length" in config and "vae_stride" in config: - if config["target_video_length"] % config["vae_stride"][0] != 1: - logger.warning(f"`num_frames - 1` has to be divisible by {config['vae_stride'][0]}. Rounding to the nearest number.") - config["target_video_length"] = config["target_video_length"] // config["vae_stride"][0] * config["vae_stride"][0] + 1 - - # Load diffusers vae config - if os.path.exists(os.path.join(config["model_path"], "vae", "config.json")): - with open(os.path.join(config["model_path"], "vae", "config.json"), "r") as f: - vae_config = json.load(f) - if "temperal_downsample" in vae_config: - config["vae_scale_factor"] = 2 ** len(vae_config["temperal_downsample"]) - elif "block_out_channels" in vae_config: - config["vae_scale_factor"] = 2 ** (len(vae_config["block_out_channels"]) - 1) - if config["model_cls"] == "ernie_image": - config["vae_scale_factor"] = 2 ** len(vae_config["block_out_channels"]) - - if config["model_cls"] == "cosmos3" and os.path.exists(os.path.join(config["model_path"], "vae", "config.json")): - with open(os.path.join(config["model_path"], "vae", "config.json"), "r") as f: - vae_config = json.load(f) - config["vae_scale_factor_spatial"] = int(vae_config.get("scale_factor_spatial", 16)) - config["vae_scale_factor_temporal"] = int(vae_config.get("scale_factor_temporal", 4)) - config["vae_scale_factor"] = config["vae_scale_factor_spatial"] - if config["model_cls"] == "lingbot_video": - config["vae_scale_factor_spatial"] = int(config.get("vae_scale_factor_spatial", 8)) - config["vae_scale_factor_temporal"] = int(config.get("vae_scale_factor_temporal", 4)) - config["vae_scale_factor"] = config["vae_scale_factor_spatial"] - if config["model_cls"] == "minimax_h3": - # The generic Diffusers-VAE heuristic above counts six encoder stages - # and would incorrectly derive 32. H3 downsamples space by exactly 16. - config["vae_spatial_scale_factor"] = 16 - config["vae_scale_factor"] = 16 - if config["model_cls"] == "cosmos3" and os.path.exists(os.path.join(config["model_path"], "sound_tokenizer", "config.json")): - with open(os.path.join(config["model_path"], "sound_tokenizer", "config.json"), "r") as f: - sound_config = json.load(f) - config["sound_sampling_rate"] = int(sound_config.get("sampling_rate", 48000)) - config["sound_hop_size"] = int(sound_config.get("hop_size", 1920)) - - return config - - -def set_config(args): - validate_model_task_args(args) - config = set_args2config(args) - config = auto_calc_config(config) - return config - - -def set_parallel_config(config): - if config["parallel"]: - tensor_p_size = int(config["parallel"].get("tensor_p_size", 1)) - cfg_p_size = int(config["parallel"].get("cfg_p_size", 1)) - seq_p_size = int(config["parallel"].get("seq_p_size", 1)) - world_size = dist.get_world_size() - expected_world_size = tensor_p_size * cfg_p_size * seq_p_size - if expected_world_size != world_size: - raise ValueError( - f"Parallel sizes must match the distributed world size: tensor_p_size ({tensor_p_size}) * cfg_p_size ({cfg_p_size}) * seq_p_size ({seq_p_size}) != world_size ({world_size})." - ) - - phase_aware = bool(config.get("model_cls") == "hunyuan_image3" and config["parallel"].get("phase_aware", False)) - if phase_aware: - from lightx2v.models.networks.hunyuan_image3.parallel import initialize_hunyuan_image3_parallel_runtime - - initialize_hunyuan_image3_parallel_runtime(config) - elif tensor_p_size > 1: - # Tensor parallel is the innermost dimension. Optional CFG and - # sequence dimensions are prepended so ranks with the same - # non-TP coordinates form contiguous TP groups. For TP+SP+CFG: - # mesh shape/names = [cfg_p, seq_p, tensor_p]. - mesh_shape = [] - mesh_dim_names = [] - if cfg_p_size > 1: - mesh_shape.append(cfg_p_size) - mesh_dim_names.append("cfg_p") - if seq_p_size > 1: - mesh_shape.append(seq_p_size) - mesh_dim_names.append("seq_p") - mesh_shape.append(tensor_p_size) - mesh_dim_names.append("tensor_p") - config["device_mesh"] = init_device_mesh( - AI_DEVICE, - tuple(mesh_shape), - mesh_dim_names=tuple(mesh_dim_names), - ) - config["tensor_parallel"] = True - config["seq_parallel"] = seq_p_size > 1 - config["cfg_parallel"] = bool(config.get("enable_cfg", False) and cfg_p_size > 1) - else: - # Original 2D mesh for cfg_p and seq_p - config["device_mesh"] = init_device_mesh(AI_DEVICE, (cfg_p_size, seq_p_size), mesh_dim_names=("cfg_p", "seq_p")) - config["tensor_parallel"] = False - config["seq_parallel"] = seq_p_size > 1 - config["cfg_parallel"] = bool(config.get("enable_cfg", False) and cfg_p_size > 1) - - # warmup dist - if AI_DEVICE == "cuda": - warmup_device = f"{AI_DEVICE}:{torch.cuda.current_device()}" - else: - warmup_device = AI_DEVICE - _a = torch.zeros([1], device=warmup_device) - dist.all_reduce(_a) + if ( + config["model_cls"] != "minimax_h3" + and config["task"] in ["i2v", "t2av", "i2av", "i2va", "s2v", "rs2v", "ltx2_s2v", "v2av"] + and config.get("target_video_length") is not None + and "vae_stride" in config + ): + temporal_stride = int(config["vae_stride"][0]) + if (config["target_video_length"] - 1) % temporal_stride != 0: + original_length = config["target_video_length"] + config["target_video_length"] = align_target_video_length(original_length, temporal_stride) + logger.warning(f"`num_frames - 1` must be divisible by {temporal_stride}; using {config['target_video_length']} instead of {original_length}.") + + +def build_cli_inputs(args): + args_data = {key: value for key, value in vars(args).items() if value is not None} + startup_fields = {"config_json", "model_cls", "model_path", "sf_model_path", "task"} + startup_args = {key: value for key, value in args_data.items() if key in startup_fields} + request_data = {key: value for key, value in args_data.items() if key not in startup_fields} + request_data["task"] = args.task + startup_config = build_startup_config(startup_args) + return startup_config, request_data + + +def init_parallel(config): + """Create the model's parallel mesh and warm up its communication.""" + parallel = config["parallel"] + if not parallel: + return + + tensor_p_size = int(parallel.get("tensor_p_size", 1)) + cfg_p_size = int(parallel.get("cfg_p_size", 1)) + seq_p_size = int(parallel.get("seq_p_size", 1)) + world_size = dist.get_world_size() + expected_world_size = tensor_p_size * cfg_p_size * seq_p_size + if expected_world_size != world_size: + raise ValueError(f"Parallel sizes must match the distributed world size: tensor_p_size ({tensor_p_size}) * cfg_p_size ({cfg_p_size}) * seq_p_size ({seq_p_size}) != world_size ({world_size}).") + + if config.get("model_cls") == "hunyuan_image3" and parallel.get("phase_aware", False): + from lightx2v.models.networks.hunyuan_image3.parallel import initialize_hunyuan_image3_parallel_runtime + + initialize_hunyuan_image3_parallel_runtime(config) + else: + # Keep the original CFG/SP dimensions without TP. With TP, omit unit + # dimensions and place TP last so its ranks form contiguous groups. + mesh_dims = [("cfg_p", cfg_p_size), ("seq_p", seq_p_size)] + if tensor_p_size > 1: + mesh_dims = [(name, size) for name, size in mesh_dims if size > 1] + mesh_dims.append(("tensor_p", tensor_p_size)) + config["device_mesh"] = init_device_mesh( + AI_DEVICE, + tuple(size for _, size in mesh_dims), + mesh_dim_names=tuple(name for name, _ in mesh_dims), + ) + config["tensor_parallel"] = tensor_p_size > 1 + config["seq_parallel"] = seq_p_size > 1 + config["cfg_parallel"] = bool(config.get("enable_cfg", False) and cfg_p_size > 1) + + warmup_device = f"cuda:{torch.cuda.current_device()}" if AI_DEVICE == "cuda" else AI_DEVICE + warmup_tensor = torch.zeros([1], device=warmup_device) + dist.all_reduce(warmup_tensor) + + +def print_config(config, title="config"): + if is_main_process(): + logger.info(f"{title}:\n{json.dumps(config, ensure_ascii=False, indent=4, default=str)}") -def print_config(config): - config_to_print = config.copy() - if is_main_process(): - logger.info(f"config:\n{json.dumps(config_to_print, ensure_ascii=False, indent=4, default=str)}") +def print_request(input_info, supported_request_fields): + request = {input_field.name: getattr(input_info, input_field.name) for input_field in fields(input_info) if input_field.repr and input_field.name in supported_request_fields} + print_config(request, title="Effective request") diff --git a/lightx2v/utils/va_controller.py b/lightx2v/utils/va_controller.py index 27d948115..8155a7f99 100644 --- a/lightx2v/utils/va_controller.py +++ b/lightx2v/utils/va_controller.py @@ -36,9 +36,8 @@ def __init__(self, model_runner): self.init_reader(model_runner) def init_base(self, config, input_info, has_vfi_model, has_vsr_model): - if "stream_config" in input_info.__dataclass_fields__: - self.stream_config = input_info.stream_config - logger.info(f"VAController init base with stream config: {self.stream_config}") + self.stream_config = input_info.stream_config + logger.info(f"VAController init base with stream config: {self.stream_config}") self.audio_path = input_info.audio_path self.output_video_path = input_info.save_result_path if isinstance(self.output_video_path, dict): diff --git a/lightx2v_ros/src/inference/inference/cosmos3_node/main.py b/lightx2v_ros/src/inference/inference/cosmos3_node/main.py index 4cfc8ae9e..7d5fe6476 100644 --- a/lightx2v_ros/src/inference/inference/cosmos3_node/main.py +++ b/lightx2v_ros/src/inference/inference/cosmos3_node/main.py @@ -10,7 +10,7 @@ from std_msgs.msg import Bool, Float32MultiArray, Int32, String from lightx2v.models.runners.cosmos3.cosmos3_runner import Cosmos3Policy -from lightx2v.utils.set_config import auto_calc_config, get_default_config, set_parallel_config +from lightx2v.utils.set_config import build_startup_config, init_parallel from lightx2v.utils.utils import seed_all from lightx2v_platform.registry_factory import PLATFORM_DEVICE_REGISTER @@ -106,8 +106,7 @@ def _build_policy_config(self): if not model_path: raise ValueError("Cosmos3 ROS node requires `model_path`.") - config = get_default_config() - config.update( + config = build_startup_config( { "model_cls": "cosmos3", "task": "i2va", @@ -116,9 +115,7 @@ def _build_policy_config(self): "seed": int(self.get_parameter("seed").value), } ) - config = auto_calc_config(config) - # ROS parameters are explicit runtime overrides and must win over the - # values loaded from config_json by auto_calc_config(). + # The ROS prompt format overrides deployment defaults. config["policy_prompt_format"] = prompt_format if int(config.get("raw_action_dim", 8)) != self.contract.action_dim: raise ValueError(f"Cosmos3 raw_action_dim={config.get('raw_action_dim')} != RoboLab action_dim={self.contract.action_dim}") @@ -137,7 +134,7 @@ def _initialize_parallel(self): if platform is None: raise RuntimeError(f"unsupported LightX2V platform: {os.getenv('PLATFORM', 'cuda')}") platform.init_parallel_env() - set_parallel_config(self.policy_config) + init_parallel(self.policy_config) def _broadcast(self, payload): if not dist.is_initialized(): diff --git a/lightx2v_ros/src/inference/inference/fastwam_node/main.py b/lightx2v_ros/src/inference/inference/fastwam_node/main.py index 6676a2db7..52aba4f92 100644 --- a/lightx2v_ros/src/inference/inference/fastwam_node/main.py +++ b/lightx2v_ros/src/inference/inference/fastwam_node/main.py @@ -6,7 +6,7 @@ from std_msgs.msg import Bool, Float32MultiArray, Int32, String from lightx2v.models.runners.wan.fastwam_runner import FastWAMPolicy -from lightx2v.utils.set_config import auto_calc_config, get_default_config +from lightx2v.utils.set_config import build_startup_config class FastWAMNode(Node): @@ -59,8 +59,7 @@ def build_policy_config(self): model_path = str(self.get_parameter("model_path").value).strip() if not model_path: raise ValueError("FastWAM ROS node requires `model_path`.") - config = get_default_config() - config.update( + config = build_startup_config( { "model_cls": "fastwam", "task": "i2va", @@ -68,7 +67,6 @@ def build_policy_config(self): "config_json": config_json, } ) - config = auto_calc_config(config) # The config_json is authoritative for policy params; warn loudly on any # mismatch with the environment contract so dimension bugs surface early. diff --git a/lightx2v_ros/src/inference/inference/lingbot_va_node/main.py b/lightx2v_ros/src/inference/inference/lingbot_va_node/main.py index a40bd4a1d..c97e35717 100644 --- a/lightx2v_ros/src/inference/inference/lingbot_va_node/main.py +++ b/lightx2v_ros/src/inference/inference/lingbot_va_node/main.py @@ -6,7 +6,7 @@ from std_msgs.msg import Bool, Float32MultiArray, Int32, String from lightx2v.models.runners.wan.wan_lingbot_va_runner import LingbotVAPolicy -from lightx2v.utils.set_config import auto_calc_config, get_default_config +from lightx2v.utils.set_config import build_startup_config class LingbotVANode(Node): @@ -73,8 +73,7 @@ def build_policy_config(self): model_path = str(self.get_parameter("model_path").value).strip() if not model_path: raise ValueError("LingBot-VA ROS node requires `model_path`.") - config = get_default_config() - config.update( + return build_startup_config( { "model_cls": "lingbot_va", "task": "i2va", @@ -83,7 +82,6 @@ def build_policy_config(self): "seed": self.seed, } ) - return auto_calc_config(config) def _make_image_cb(self, camera): def _callback(msg): diff --git a/pyproject.toml b/pyproject.toml index 18794576e..0dcbdc739 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ build-backend = "setuptools.build_meta" [project] name = "lightx2v" -version = "0.1.0" +version = "0.5.0" authors = [ {name = "LightX2V Contributors"}, ] diff --git a/scripts/disagg/run_dynamic.sh b/scripts/disagg/run_dynamic.sh index 179c8d063..b6dbe989c 100644 --- a/scripts/disagg/run_dynamic.sh +++ b/scripts/disagg/run_dynamic.sh @@ -98,12 +98,10 @@ else user_max_requests=${DISAGG_AUTO_REQUEST_COUNT} fi -seed_args=() -if [[ -v SEED ]]; then - seed_args=(--seed "${SEED}") -fi +seed=${SEED:-42} prompt=${PROMPT:-"Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard."} negative_prompt=${NEGATIVE_PROMPT:-"镜头晃动,色调艳丽,过曝,静态"} +image_path=${IMAGE_PATH:-${lightx2v_path}/assets/inputs/imgs/img_0.jpg} save_result_path=${SAVE_RESULT_PATH:-${lightx2v_path}/save_results/wan22_i2v_dynamic.mp4} controller_log=${lightx2v_path}/save_results/disagg_wan22_i2v_dynamic_controller.log @@ -480,9 +478,10 @@ python -m lightx2v.disagg.examples.run_service \ --task i2v \ --model_path ${model_path} \ --config_json ${controller_cfg} \ - "${seed_args[@]}" \ + --seed "${seed}" \ --prompt "${prompt}" \ --negative_prompt "${negative_prompt}" \ + --image_path "${image_path}" \ --save_result_path ${save_result_path} \ > ${controller_log} 2>&1 & controller_pid=$! @@ -495,9 +494,14 @@ if [[ "${LOAD_FROM_USER}" != "0" ]]; then echo "waiting ${user_start_delay_s}s before run_user to let remote services warm up" sleep "${user_start_delay_s}" fi + export DISAGG_WORKLOAD_SAVE_PREFIX=${DISAGG_WORKLOAD_SAVE_PREFIX:-${save_result_path}} python -m lightx2v.disagg.examples.run_user \ --controller_host "${DISAGG_CONTROLLER_HOST}" \ --controller_request_port "${DISAGG_CONTROLLER_REQUEST_PORT}" \ + --seed "${seed}" \ + --prompt "${prompt}" \ + --negative_prompt "${negative_prompt}" \ + --image_path "${image_path}" \ --max_requests "${user_max_requests}" \ > ${user_log} 2>&1 & user_pid=$! diff --git a/scripts/dreamzero/run_dreamzero_droid_i2va.sh b/scripts/dreamzero/run_dreamzero_droid_i2va.sh index 1772f7477..2ceac419d 100755 --- a/scripts/dreamzero/run_dreamzero_droid_i2va.sh +++ b/scripts/dreamzero/run_dreamzero_droid_i2va.sh @@ -17,6 +17,7 @@ python -m lightx2v.infer \ --config_json ${lightx2v_path}/configs/dreamzero/dreamzero_droid_i2va.json \ --seed 1140 \ --prompt "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan" \ +--negative_prompt "Vibrant colors, overexposed, static, blurry details, text, subtitles, style, artwork, painting, image, still, grayscale, dull, worst quality, low quality, JPEG artifacts, ugly, mutilated, extra fingers, bad hands, bad face, deformed, disfigured, mutated limbs, fused fingers, stagnant image, cluttered background, three legs, many people in the background, walking backwards." \ --image_path $input_path \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_dreamzero_droid_i2va.mp4 \ --save_action_path ${lightx2v_path}/save_results/output_lightx2v_dreamzero_droid_i2va_actions.npy diff --git a/scripts/dreamzero/run_dreamzero_droid_i2va_dist_cfg.sh b/scripts/dreamzero/run_dreamzero_droid_i2va_dist_cfg.sh index e1e716f04..96f97a696 100755 --- a/scripts/dreamzero/run_dreamzero_droid_i2va_dist_cfg.sh +++ b/scripts/dreamzero/run_dreamzero_droid_i2va_dist_cfg.sh @@ -17,6 +17,7 @@ torchrun --nproc_per_node=2 -m lightx2v.infer \ --config_json ${lightx2v_path}/configs/dreamzero/dreamzero_droid_i2va_dist_cfg.json \ --seed 1140 \ --prompt "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan" \ +--negative_prompt "Vibrant colors, overexposed, static, blurry details, text, subtitles, style, artwork, painting, image, still, grayscale, dull, worst quality, low quality, JPEG artifacts, ugly, mutilated, extra fingers, bad hands, bad face, deformed, disfigured, mutated limbs, fused fingers, stagnant image, cluttered background, three legs, many people in the background, walking backwards." \ --image_path $input_path \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_dreamzero_droid_i2va_dist_cfg.mp4 \ --save_action_path ${lightx2v_path}/save_results/output_lightx2v_dreamzero_droid_i2va_dist_cfg_actions.npy diff --git a/scripts/infinitetalk/run_infinitetalk_single_video.sh b/scripts/infinitetalk/run_infinitetalk_single_video.sh index e6852eae8..0ddbf295a 100755 --- a/scripts/infinitetalk/run_infinitetalk_single_video.sh +++ b/scripts/infinitetalk/run_infinitetalk_single_video.sh @@ -16,8 +16,7 @@ python -m lightx2v.infer \ --model_path $model_path \ --config_json ${lightx2v_path}/configs/infinitetalk/infinitetalk_480p_single_distilled.json \ --prompt "A man is talking" \ ---negative_prompt "bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \ ---src_video /data/nvme4/gushiqiao/new/InfiniteTalk/examples/single/ref_video.mp4 \ +--video_path /data/nvme4/gushiqiao/new/InfiniteTalk/examples/single/ref_video.mp4 \ --audio_path /data/nvme4/gushiqiao/new/InfiniteTalk/examples/single/1.wav \ --save_result_path ${lightx2v_path}/save_results/infinitetalk_single_video_480p.mp4 \ --seed 42 diff --git a/scripts/ltx2/ltx2_3/run_ltx2_3_t2av_8_unsample_3.sh b/scripts/ltx2/ltx2_3/run_ltx2_3_t2av_8_unsample_3.sh index e8ff41f63..af50babff 100755 --- a/scripts/ltx2/ltx2_3/run_ltx2_3_t2av_8_unsample_3.sh +++ b/scripts/ltx2/ltx2_3/run_ltx2_3_t2av_8_unsample_3.sh @@ -16,5 +16,4 @@ python -m lightx2v.infer \ --model_path $model_path \ --config_json ${lightx2v_path}/configs/ltx2/ltx2_3_upsample_compile.json \ --prompt "A beautiful sunset over the ocean" \ ---negative_prompt "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of field, background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent lighting direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny valley effect, incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, mismatched lip sync, silent or muted audio, distorted voice, robotic voice, echo, background noise, off-sync audio, incorrect dialogue, added dialogue, repetitive speech, jittery movement, awkward pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts." \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_ltx2_3_t2av_8_upsample.mp4 diff --git a/scripts/minimax_h3/README.md b/scripts/minimax_h3/README.md new file mode 100644 index 000000000..6b5a54519 --- /dev/null +++ b/scripts/minimax_h3/README.md @@ -0,0 +1,189 @@ +# MiniMax-H3 + +[English](README.md) | [简体中文](README_zh.md) + +MiniMax-H3 generates video with synchronized stereo audio. Run the commands below from the LightX2V repository root. Set `lightx2v_path` and `model_path` in the selected shell script before running it. + +## Weights and tasks + +Use the released Diffusers component layout under `model_path`: + +```text +MiniMax-H3/ +├── transformer/ # t2av, i2av, l2av, fl2av +├── transformer_ref/ # ref2av +├── text_encoder/ +├── tokenizer/ +├── processor/ +├── vae/ +└── audio_vae/ +``` + +Each transformer directory needs its `config.json`, weight index, and checkpoint shards. Download the components for the task family you use; downloading only the original `FL2VA/` or `Ref2VA/` checkpoint directories does not provide this layout. FP8/INT8 presets still use the original component configs, tokenizer, and VAEs, and specify additional local quantized checkpoints in JSON. + +| Task | Request inputs | Transformer | +| --- | --- | --- | +| `t2av` | `prompt` | `transformer` | +| `i2av` | `prompt`, `image_path` (first frame) | `transformer` | +| `l2av` | `prompt`, `last_frame_path` | `transformer` | +| `fl2av` | `prompt`, `image_path`, `last_frame_path` | `transformer` | +| `ref2av` | `prompt`, reference images and/or videos, optional reference audio | `transformer_ref` | + +The base transformer can serve all four base tasks without reloading. Reference generation uses a separate transformer and service. These checkpoints are CFG-distilled: do not send `negative_prompt`, including an empty string. + +## Offline inference + +The five task scripts share `configs/minimax_h3/minimax_h3.json`: one GPU, BF16 weights, model CPU offload, and 124 frames at `[height, width] = [544, 960]`. The script's `--task` selects the transformer and input handling; the JSON filename does not select a task. + +```bash +bash scripts/minimax_h3/run_minimax_h3_t2av.sh +bash scripts/minimax_h3/run_minimax_h3_i2av.sh +bash scripts/minimax_h3/run_minimax_h3_l2av.sh +bash scripts/minimax_h3/run_minimax_h3_fl2av.sh +bash scripts/minimax_h3/run_minimax_h3_ref2av.sh +``` + +The ordinary config uses SageAttention2 and SGL kernels. Select a config whose attention, quantization, and parallel settings match your installed kernels and devices. Paths, prompt, seed, and output path are visible in each script; each JSON remains a complete startup config. + +To select another mode, change `--config_json` in the listed script to the corresponding file under `configs/minimax_h3/`. The single-GPU script also launches FP8, compile, and four-step LoRA modes; `run_minimax_h3_t2av_parallel.sh` provides a shared launcher for SP, TP, and mixed parallel configs. The table lists each preset's configured parallel sizes. In filenames, `encoder` refers to the text encoder and `vae` to the video VAE. + +| Configuration under `configs/minimax_h3/` | Launch script | Mode | +| --- | --- | --- | +| `minimax_h3.json` | Any of the five task scripts above | Single-GPU BF16 | +| `minimax_h3_compile.json` | Any of the five task scripts above | Compile with startup warmup | +| `minimax_h3_block_offload.json` | `run_minimax_h3_t2av.sh` | Single-GPU BF16 block offload | +| `minimax_h3_sp.json` | `run_minimax_h3_t2av_parallel.sh` | SP4 | +| `minimax_h3_tp.json` | `run_minimax_h3_t2av_parallel.sh` | TP2 | +| `minimax_h3_tp_sp.json` | `run_minimax_h3_t2av_parallel.sh` | TP2 × SP2 | +| `minimax_h3_sol_block_offload.json` | `run_minimax_h3_t2av.sh` | Single-GPU Sol-Attn | +| `fp8/minimax_h3.json` | `run_minimax_h3_t2av.sh` | Single-GPU DiT FP8 | +| `fp8/minimax_h3_encoder_fp8.json` | `run_minimax_h3_t2av.sh` | Single-GPU DiT + text-encoder FP8 | +| `fp8/minimax_h3_vae_fp8.json` | `run_minimax_h3_t2av.sh` | Single-GPU DiT + video-VAE FP8 | +| `fp8/minimax_h3_sp_5090.json` | `run_minimax_h3_t2av_parallel.sh` | SP8, 5090 FP8 config | +| `dmd/minimax_h3_bf16_4step.json` | `run_minimax_h3_t2av.sh` | Single-GPU BF16, 4-step LoRA | +| `dmd/minimax_h3_bf16_4step_sol.json` | `run_minimax_h3_t2av.sh` | Single-GPU 4-step Sol-Attn | +| `dmd/minimax_h3_fp8_4step.json` | `run_minimax_h3_t2av_parallel.sh` | SP4, FP8 + 4-step LoRA | +| `dmd/minimax_h3_int8_4step.json` | `run_minimax_h3_t2av_parallel.sh` | SP4, INT8 + 4-step LoRA | +| `dmd/minimax_h3_fp8_8step.json` | `run_minimax_h3_t2av_parallel.sh` | SP8, FP8 + 8-step LoRA | +| `dmd/minimax_h3_int8_convrot_8step.json` | `run_minimax_h3_t2av_parallel.sh` | SP8, INT8 ConvRot + 8-step LoRA | +| `dmd/minimax_h3_fp8_4step_5090.json` | `run_minimax_h3_t2av_parallel.sh` | SP8, 5090, 4-step LoRA | +| `dmd/minimax_h3_fp8_4step_5090_vae_fp8.json` | `run_minimax_h3_t2av_parallel.sh` | SP8, FP8 VAE + 4-step LoRA | +| `dmd/minimax_h3_fp8_4step_5090_vae_fp8_sla.json` | `run_minimax_h3_t2av_parallel.sh` | SP8, FP8 VAE + matching SLA LoRA | +| `dmd/minimax_h3_fp8_4step_5090_vae_fp8_sol.json` | `run_minimax_h3_t2av_parallel.sh` | SP8, FP8 text encoder/VAE + Sol-Attn + 4-step LoRA | +| `dmd/minimax_h3_ref2av_4step.json` | `run_minimax_h3_ref2av.sh`, with 8 processes as described below | Reference-task 4-step LoRA | + +`run_minimax_h3_t2av_parallel.sh` defaults to SP4. To select another parallel preset, change `--config_json` and keep `torchrun --nproc_per_node` and `CUDA_VISIBLE_DEVICES` consistent with the JSON. The process count is `tensor_p_size × seq_p_size` for these presets; an omitted parallel size is 1. + +| Parallel mode in the selected JSON | `CUDA_VISIBLE_DEVICES` | `--nproc_per_node` | +| --- | --- | --- | +| TP2 | `0,1` | `2` | +| SP4 or TP2 × SP2 | `0,1,2,3` | `4` | +| SP8, including the 5090 presets | `0,1,2,3,4,5,6,7` | `8` | + +For the 8-GPU reference LoRA config, set `CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7` in `run_minimax_h3_ref2av.sh` and replace `python -m lightx2v.infer` with `torchrun --standalone --nproc_per_node=8 -m lightx2v.infer`. Keep its `--task ref2av` and reference input arguments. + +Both single-GPU Sol configs default to 362 frames at `[768, 1344]`. Select the corresponding Sol JSON through `--config_json` in `run_minimax_h3_t2av.sh` to use these output defaults. + +All MiniMax-H3 configs that enable compilation also set `warmup: true`. The ordinary and compile configs share the same output defaults and can be used by both service launch scripts as well. + +### Video encoding + +To reproduce the 362-frame output encoder example, update these fields in a copy of `minimax_h3.json` and point the launch script at that complete JSON: + +```json +{ + "target_video_length": 362, + "video_codec_options": { + "preset": "ultrafast", + "crf": "18" + } +} +``` + +`video_codec_options` is a startup setting consumed by the existing MP4 encoder. Other dimensions and frame counts remain request overrides; no separate config is needed for each output shape. + +## LoRA and Python + +Replace `/path/to/...` in every selected quantized checkpoint and LoRA field with an existing local file. The loader does not download LoRAs from a repository ID. Use a LoRA for the correct task family and keep `alpha` consistent with that checkpoint's training configuration; do not interchange base, reference, SLA, or differently versioned LoRAs merely because their shapes match. + +The named base 4-step v1.0 LoRA uses `alpha=128`, `video_flow_shift=6`, and `audio_flow_shift=3`. Other presets retain their own alpha and shift settings. + +Choose `lora_dynamic_apply` in the selected JSON according to the DiT weights: + +| DiT weights | Supported setting | Behavior | +| --- | --- | --- | +| Original BF16 | `false` or `true` | `false` merges the adapter when loading weights; `true` applies it during inference. | +| Quantized FP8 or INT8, including ConvRot | `true` | Apply the adapter during inference; merging into quantized DiT weights is unsupported. | + +Quantizing only the encoder or VAE does not impose this restriction. Dynamic LoRA currently accepts one adapter and requires its `alpha` to be explicitly configured. The existing model loader rejects unsupported combinations. + +The shared `dmd/minimax_h3_bf16_4step.json` defaults to `lora_dynamic_apply: true`, with 362 frames at `[768, 1344]`. Set the same field to `false` to use merging. Select this JSON in `run_minimax_h3_t2av.sh`; for a shorter, smaller output, add `--target_shape 544 960` and `--num_frames 124` to its inference command. + +Set `infer_steps` in JSON to the number of model evaluations: `4` for 4-step inference and `8` for 8-step inference. The scheduler includes the terminal zero automatically. When migrating an older MiniMax-H3 config, subtract one from its `infer_steps` value (`5` → `4`, `9` → `8`, `30` → `29`) to preserve the original sampling schedule. The bundled configs already use this convention. + +For Python, set `MODEL_PATH` in [the example](../../examples/minimax_h3/minimax_h3_t2av_dmd.py) and the local LoRA path in its selected JSON, then run: + +```bash +python examples/minimax_h3/minimax_h3_t2av_dmd.py +``` + +## Server and POST + +Start the base service, then submit a request from another terminal: + +```bash +bash scripts/minimax_h3/server/start_server.sh +``` + +```bash +python scripts/minimax_h3/server/post_t2av.py +python scripts/minimax_h3/server/post_i2av.py +python scripts/minimax_h3/server/post_l2av.py +python scripts/minimax_h3/server/post_fl2av.py +``` + +Each base request must specify `task`, because the loaded transformer supports four tasks. The image examples encode client-local files as Base64. + +For reference generation, use the separate reference service instead: + +```bash +bash scripts/minimax_h3/server/start_server_ref2av.sh +``` + +```bash +python scripts/minimax_h3/server/post_ref2av.py +``` + +Both launch scripts default to port 8000. To run both services simultaneously, choose separate GPUs, `--port`, and `--metric_port` values, and update the POST URLs. The reference service accepts only `ref2av`, so `task` may be omitted there; the example includes it for clarity. + +Reference requests can supply `image_path`, `video_path`, and `audio_path` together. For multiple files of one kind, use comma-separated server-local paths, for example: + +```json +{ + "task": "ref2av", + "prompt": "Generate an audio-video scene following the references.", + "image_path": "/path/to/character.jpg,/path/to/scene.jpg", + "video_path": "/path/to/motion.mp4", + "audio_path": "/path/to/voice.wav", + "seed": 42, + "save_result_path": "./minimax_h3_references.mp4" +} +``` + +Audio must be accompanied by an image or video. The runner accepts at most 9 images, 3 videos, and 12 references in total; the audio-reference limit is 3, including video soundtracks. Single image inputs also accept Base64 and HTTP(S) URLs; input videos use server-local paths. The default reference-image preprocessing follows the released Diffusers sizing; `reference_image_resize_mode: "match"` in JSON is an explicit alternative that limits reference-image area to the output canvas. + +POST returns a task ID before inference completes. Query status, then download the completed result: + +```bash +curl http://localhost:8000/v1/tasks/TASK_ID/status +curl --fail http://localhost:8000/v1/tasks/TASK_ID/result -o minimax_h3.mp4 +``` + +Use a relative path such as `./minimax_h3_t2av.mp4` for `save_result_path`; the service resolves it relative to its output directory and can return it through the download endpoint. Omitting the path or sending `null` skips saving; that request has no downloadable result. + +## Request and startup settings + +- Startup: model paths, default task, weights/LoRA, kernels, offload, parallelism, compile, and warmup. Prompt, media paths, seed, and output path belong in the Python call, CLI command, or POST body. +- Output defaults: JSON sets `target_video_length` and `target_height`/`target_width`. Requests can override them with `num_frames` and `target_shape` (`[height, width]`). Dimensions must be multiples of 32. Frame counts align upward to `17*n+5`, with supported aligned counts from 124 to 362; for example, 125 becomes 141. +- Seed: omitted or `null` defaults to 42; an explicit non-negative value, including 0, is used as supplied. +- Saving: every CLI example explicitly provides an output path. Removing it skips file saving. The service follows the same rule. Output is MP4 with 24 FPS video and 32 kHz stereo audio. diff --git a/scripts/minimax_h3/README_zh.md b/scripts/minimax_h3/README_zh.md new file mode 100644 index 000000000..eebd86a14 --- /dev/null +++ b/scripts/minimax_h3/README_zh.md @@ -0,0 +1,189 @@ +# MiniMax-H3 + +[English](README.md) | [简体中文](README_zh.md) + +MiniMax-H3 可以生成带有同步立体声音频的视频。以下命令均在 LightX2V 仓库根目录执行。运行前,请先设置所选 shell 脚本中的 `lightx2v_path` 和 `model_path`。 + +## 模型权重与任务 + +`model_path` 下需要采用发布的 Diffusers 组件目录结构: + +```text +MiniMax-H3/ +├── transformer/ # t2av, i2av, l2av, fl2av +├── transformer_ref/ # ref2av +├── text_encoder/ +├── tokenizer/ +├── processor/ +├── vae/ +└── audio_vae/ +``` + +每个 transformer 目录都需要包含 `config.json`、权重索引和权重分片。请按所需任务下载对应组件;只下载原始的 `FL2VA/` 或 `Ref2VA/` 权重目录,并不能满足上述目录要求。FP8/INT8 配置仍需使用原始组件配置、tokenizer 和 VAE,并在 JSON 中另行指定本地量化权重。 + +| 任务 | 请求输入 | Transformer | +| --- | --- | --- | +| `t2av` | `prompt` | `transformer` | +| `i2av` | `prompt`、`image_path`(首帧) | `transformer` | +| `l2av` | `prompt`、`last_frame_path` | `transformer` | +| `fl2av` | `prompt`、`image_path`、`last_frame_path` | `transformer` | +| `ref2av` | `prompt`、参考图片和/或视频,以及可选的参考音频 | `transformer_ref` | + +基础 transformer 加载一次即可处理前四种任务。参考生成使用独立的 transformer 和服务。这些权重已经完成 CFG 蒸馏,请勿传入 `negative_prompt`,包括空字符串。 + +## 离线推理 + +五种任务脚本共用 `configs/minimax_h3/minimax_h3.json`:单 GPU、BF16 权重、模型级 CPU 卸载,默认输出 124 帧,`[高度, 宽度] = [544, 960]`。脚本中的 `--task` 决定加载哪组 transformer 以及如何处理输入,JSON 文件名不决定任务。 + +```bash +bash scripts/minimax_h3/run_minimax_h3_t2av.sh +bash scripts/minimax_h3/run_minimax_h3_i2av.sh +bash scripts/minimax_h3/run_minimax_h3_l2av.sh +bash scripts/minimax_h3/run_minimax_h3_fl2av.sh +bash scripts/minimax_h3/run_minimax_h3_ref2av.sh +``` + +普通配置使用 SageAttention2 和 SGL 算子。请根据已安装的算子和设备,选择合适的注意力、量化和并行配置。模型路径、prompt、seed 和输出路径直接写在脚本中;每个 JSON 都是一份完整的启动配置。 + +切换运行方式时,将下表中脚本的 `--config_json` 改为 `configs/minimax_h3/` 下对应的文件。单 GPU 脚本也可运行 FP8、compile 和四步 LoRA 配置;`run_minimax_h3_t2av_parallel.sh` 统一用于 SP、TP 和混合并行。下表列出各配置设置的并行规模。文件名中的 `encoder` 指文本编码器,`vae` 指视频 VAE。 + +| `configs/minimax_h3/` 下的配置 | 启动脚本 | 运行方式 | +| --- | --- | --- | +| `minimax_h3.json` | 上述五种任务脚本中的任意一个 | 单 GPU BF16 | +| `minimax_h3_compile.json` | 上述五种任务脚本中的任意一个 | 编译并在启动时预热 | +| `minimax_h3_block_offload.json` | `run_minimax_h3_t2av.sh` | 单 GPU BF16,分块卸载 | +| `minimax_h3_sp.json` | `run_minimax_h3_t2av_parallel.sh` | SP4 | +| `minimax_h3_tp.json` | `run_minimax_h3_t2av_parallel.sh` | TP2 | +| `minimax_h3_tp_sp.json` | `run_minimax_h3_t2av_parallel.sh` | TP2 × SP2 | +| `minimax_h3_sol_block_offload.json` | `run_minimax_h3_t2av.sh` | 单 GPU Sol-Attn | +| `fp8/minimax_h3.json` | `run_minimax_h3_t2av.sh` | 单 GPU DiT FP8 | +| `fp8/minimax_h3_encoder_fp8.json` | `run_minimax_h3_t2av.sh` | 单 GPU DiT + 文本编码器 FP8 | +| `fp8/minimax_h3_vae_fp8.json` | `run_minimax_h3_t2av.sh` | 单 GPU DiT + 视频 VAE FP8 | +| `fp8/minimax_h3_sp_5090.json` | `run_minimax_h3_t2av_parallel.sh` | SP8,5090 FP8 配置 | +| `dmd/minimax_h3_bf16_4step.json` | `run_minimax_h3_t2av.sh` | 单 GPU BF16,四步 LoRA | +| `dmd/minimax_h3_bf16_4step_sol.json` | `run_minimax_h3_t2av.sh` | 单 GPU 四步 Sol-Attn | +| `dmd/minimax_h3_fp8_4step.json` | `run_minimax_h3_t2av_parallel.sh` | SP4,FP8 + 四步 LoRA | +| `dmd/minimax_h3_int8_4step.json` | `run_minimax_h3_t2av_parallel.sh` | SP4,INT8 + 四步 LoRA | +| `dmd/minimax_h3_fp8_8step.json` | `run_minimax_h3_t2av_parallel.sh` | SP8,FP8 + 八步 LoRA | +| `dmd/minimax_h3_int8_convrot_8step.json` | `run_minimax_h3_t2av_parallel.sh` | SP8,INT8 ConvRot + 八步 LoRA | +| `dmd/minimax_h3_fp8_4step_5090.json` | `run_minimax_h3_t2av_parallel.sh` | SP8,5090,四步 LoRA | +| `dmd/minimax_h3_fp8_4step_5090_vae_fp8.json` | `run_minimax_h3_t2av_parallel.sh` | SP8,FP8 VAE + 四步 LoRA | +| `dmd/minimax_h3_fp8_4step_5090_vae_fp8_sla.json` | `run_minimax_h3_t2av_parallel.sh` | SP8,FP8 VAE + 配套的 SLA LoRA | +| `dmd/minimax_h3_fp8_4step_5090_vae_fp8_sol.json` | `run_minimax_h3_t2av_parallel.sh` | SP8,FP8 文本编码器/VAE + Sol-Attn + 四步 LoRA | +| `dmd/minimax_h3_ref2av_4step.json` | `run_minimax_h3_ref2av.sh`,按下文说明使用 8 个进程 | 参考任务四步 LoRA | + +`run_minimax_h3_t2av_parallel.sh` 默认使用 SP4。切换并行配置时,需要修改 `--config_json`,并让 `torchrun --nproc_per_node` 和 `CUDA_VISIBLE_DEVICES` 与 JSON 保持一致。这些配置所需的进程数为 `tensor_p_size × seq_p_size`;未设置的并行维度按 1 计算。 + +| 所选 JSON 中的并行方式 | `CUDA_VISIBLE_DEVICES` | `--nproc_per_node` | +| --- | --- | --- | +| TP2 | `0,1` | `2` | +| SP4 或 TP2 × SP2 | `0,1,2,3` | `4` | +| SP8,包括 5090 配置 | `0,1,2,3,4,5,6,7` | `8` | + +使用 8 GPU 参考任务 LoRA 配置时,在 `run_minimax_h3_ref2av.sh` 中设置 `CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7`,并将 `python -m lightx2v.infer` 改为 `torchrun --standalone --nproc_per_node=8 -m lightx2v.infer`。保留其中的 `--task ref2av` 和参考输入参数。 + +两个单 GPU Sol 配置均默认输出 362 帧,尺寸为 `[768, 1344]`。在 `run_minimax_h3_t2av.sh` 中通过 `--config_json` 选择对应的 Sol JSON,即可使用这些默认规格。 + +所有启用编译的 MiniMax-H3 配置都同时设置了 `warmup: true`。普通配置与编译配置的默认输出规格一致,两份服务启动脚本也都可以使用。 + +### 视频编码 + +如需使用原先的 362 帧视频编码示例,在 `minimax_h3.json` 的副本中修改以下字段,再让启动脚本指向这份完整的 JSON: + +```json +{ + "target_video_length": 362, + "video_codec_options": { + "preset": "ultrafast", + "crf": "18" + } +} +``` + +`video_codec_options` 属于启动配置,由现有的 MP4 编码器使用。尺寸和帧数仍可由请求覆盖,无需为每种输出规格单独保留一份配置。 + +## LoRA 与 Python 调用 + +将所选配置中量化权重和 LoRA 字段的 `/path/to/...` 替换为实际存在的本地文件。加载器不会根据仓库 ID 自动下载 LoRA。请使用对应任务的 LoRA,并让 `alpha` 与该权重的训练配置保持一致;基础任务、参考任务、SLA 和不同版本的 LoRA 不能仅因形状相同就相互替换。 + +基础任务四步 v1.0 LoRA 使用 `alpha=128`、`video_flow_shift=6` 和 `audio_flow_shift=3`。其他配置使用各自的 alpha 和 shift 设置。 + +根据 DiT 权重,在所选 JSON 中设置 `lora_dynamic_apply`: + +| DiT 权重 | 支持的设置 | 行为 | +| --- | --- | --- | +| 原始 BF16 | `false` 或 `true` | `false` 在加载权重时合并 LoRA;`true` 在推理时动态应用 LoRA。 | +| 量化 FP8 或 INT8,包括 ConvRot | `true` | 在推理时动态应用 LoRA,不支持将 LoRA 合并进量化 DiT 权重。 | + +只量化文本编码器或 VAE 不受上述合并限制。动态 LoRA 当前只接受一个适配器,并要求显式提供 `alpha`。现有模型加载器会拒绝不支持的组合。 + +共用配置 `dmd/minimax_h3_bf16_4step.json` 默认设置 `lora_dynamic_apply: true`,输出 362 帧,尺寸为 `[768, 1344]`。将该字段改为 `false` 即可使用加载时合并。在 `run_minimax_h3_t2av.sh` 中选择这份 JSON;如需更短、更小的输出,在推理命令中添加 `--target_shape 544 960` 和 `--num_frames 124`。 + +在 JSON 中将 `infer_steps` 设置为实际模型计算次数:四步推理填 `4`,八步推理填 `8`。Scheduler 会自动包含末尾零点。迁移旧版 MiniMax-H3 配置时,将原 `infer_steps` 减一(`5` → `4`、`9` → `8`、`30` → `29`),即可保持原有采样过程。仓库中的配置已完成同步。 + +使用 Python 时,先设置[示例](../../examples/minimax_h3/minimax_h3_t2av_dmd.py)中的 `MODEL_PATH`,以及所选 JSON 中的本地 LoRA 路径,再执行: + +```bash +python examples/minimax_h3/minimax_h3_t2av_dmd.py +``` + +## 服务与 POST 请求 + +先启动基础任务服务,再从另一个终端发送请求: + +```bash +bash scripts/minimax_h3/server/start_server.sh +``` + +```bash +python scripts/minimax_h3/server/post_t2av.py +python scripts/minimax_h3/server/post_i2av.py +python scripts/minimax_h3/server/post_l2av.py +python scripts/minimax_h3/server/post_fl2av.py +``` + +基础任务请求必须指定 `task`,因为服务加载的 transformer 支持四种任务。图片示例会将客户端本地文件编码为 Base64 后发送。 + +参考生成需要单独启动参考任务服务: + +```bash +bash scripts/minimax_h3/server/start_server_ref2av.sh +``` + +```bash +python scripts/minimax_h3/server/post_ref2av.py +``` + +两份启动脚本默认都使用 8000 端口。如需同时运行两个服务,请分别设置 GPU、`--port` 和 `--metric_port`,并修改 POST 示例中的 URL。参考任务服务只接受 `ref2av`,因此请求可以省略 `task`;示例为便于理解仍显式填写。 + +参考请求可以同时提供 `image_path`、`video_path` 和 `audio_path`。同一类型包含多个文件时,使用逗号分隔的服务端本地路径,例如: + +```json +{ + "task": "ref2av", + "prompt": "Generate an audio-video scene following the references.", + "image_path": "/path/to/character.jpg,/path/to/scene.jpg", + "video_path": "/path/to/motion.mp4", + "audio_path": "/path/to/voice.wav", + "seed": 42, + "save_result_path": "./minimax_h3_references.mp4" +} +``` + +音频必须与图片或视频一起提供。Runner 最多接受 9 张图片、3 个视频、共计 12 个参考素材;参考音频最多 3 个,视频中自带的音轨也计入此限制。单张图片还支持 Base64 和 HTTP(S) URL;输入视频使用服务端本地路径。参考图片默认按发布的 Diffusers 尺寸规则预处理;也可在 JSON 中显式设置 `reference_image_resize_mode: "match"`,将参考图片面积限制在输出画布面积以内。 + +POST 会在推理完成前返回任务 ID。随后可查询状态,并在任务完成后下载结果: + +```bash +curl http://localhost:8000/v1/tasks/TASK_ID/status +curl --fail http://localhost:8000/v1/tasks/TASK_ID/result -o minimax_h3.mp4 +``` + +`save_result_path` 建议使用 `./minimax_h3_t2av.mp4` 这样的相对路径。服务会相对于其输出目录解析该路径,并通过下载接口返回文件。省略路径或传入 `null` 都会跳过保存,该请求将没有可下载的结果。 + +## 请求参数与启动配置 + +- 启动配置:模型路径、默认任务、权重/LoRA、算子、卸载、并行、编译和预热。Prompt、媒体路径、seed 和输出路径应放在 Python 调用、CLI 命令或 POST 请求体中。 +- 输出默认值:JSON 设置 `target_video_length` 和 `target_height`/`target_width`。请求可通过 `num_frames` 和 `target_shape`(`[高度, 宽度]`)覆盖。宽高必须是 32 的倍数。帧数向上对齐到 `17*n+5`,支持的对齐后帧数范围为 124 到 362;例如,125 会调整为 141。 +- Seed:省略或传入 `null` 时默认使用 42;显式提供的非负整数按原值使用,包括 0。 +- 保存:每个 CLI 示例都显式提供输出路径,移除该参数即跳过文件保存,服务采用相同规则。输出为 MP4,包含 24 FPS 视频和 32 kHz 立体声音频。 diff --git a/scripts/minimax_h3/run_minimax_h3_fl2av.sh b/scripts/minimax_h3/run_minimax_h3_fl2av.sh index 2f4025978..1e5c9a54b 100755 --- a/scripts/minimax_h3/run_minimax_h3_fl2av.sh +++ b/scripts/minimax_h3/run_minimax_h3_fl2av.sh @@ -1,23 +1,23 @@ #!/bin/bash # set path firstly -lightx2v_path=/data/nvme6/gushiqiao/codes/LightX2V -model_path=/data/nvme6/gushiqiao/models/MiniMax-H3 +lightx2v_path=/path/to/LightX2V +model_path=/path/to/MiniMax-H3 export CUDA_VISIBLE_DEVICES=0 # set environment variables -source ${lightx2v_path}/scripts/base/base.sh +source "${lightx2v_path}/scripts/base/base.sh" export DTYPE=BF16 export SENSITIVE_LAYER_DTYPE=BF16 python -m lightx2v.infer \ ---model_cls minimax_h3 \ ---task fl2av \ ---model_path $model_path \ ---config_json ${lightx2v_path}/configs/minimax_h3/minimax_h3_fl2av.json \ ---prompt "Create a coherent transition with natural synchronized sound." \ ---image_path ${lightx2v_path}/assets/inputs/imgs/flf2v_input_first_frame-fs8.png \ ---last_frame_path ${lightx2v_path}/assets/inputs/imgs/flf2v_input_last_frame-fs8.png \ ---save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_fl2av.mp4 \ ---seed 42 + --model_cls minimax_h3 \ + --task fl2av \ + --model_path "${model_path}" \ + --config_json "${lightx2v_path}/configs/minimax_h3/minimax_h3.json" \ + --prompt "Create a coherent transition with natural synchronized sound." \ + --image_path "${lightx2v_path}/assets/inputs/imgs/flf2v_input_first_frame-fs8.png" \ + --last_frame_path "${lightx2v_path}/assets/inputs/imgs/flf2v_input_last_frame-fs8.png" \ + --save_result_path "${lightx2v_path}/save_results/output_lightx2v_minimax_h3_fl2av.mp4" \ + --seed 42 diff --git a/scripts/minimax_h3/run_minimax_h3_i2av.sh b/scripts/minimax_h3/run_minimax_h3_i2av.sh index 2489f163e..980522c72 100755 --- a/scripts/minimax_h3/run_minimax_h3_i2av.sh +++ b/scripts/minimax_h3/run_minimax_h3_i2av.sh @@ -1,22 +1,22 @@ #!/bin/bash # set path firstly -lightx2v_path=/data/nvme6/gushiqiao/codes/LightX2V -model_path=/data/nvme6/gushiqiao/models/MiniMax-H3 +lightx2v_path=/path/to/LightX2V +model_path=/path/to/MiniMax-H3 export CUDA_VISIBLE_DEVICES=0 # set environment variables -source ${lightx2v_path}/scripts/base/base.sh +source "${lightx2v_path}/scripts/base/base.sh" export DTYPE=BF16 export SENSITIVE_LAYER_DTYPE=BF16 python -m lightx2v.infer \ ---model_cls minimax_h3 \ ---task i2av \ ---model_path $model_path \ ---config_json ${lightx2v_path}/configs/minimax_h3/minimax_h3_i2av.json \ ---prompt "Animate this image with natural synchronized sound." \ ---image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ ---save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_i2av.mp4 \ ---seed 42 + --model_cls minimax_h3 \ + --task i2av \ + --model_path "${model_path}" \ + --config_json "${lightx2v_path}/configs/minimax_h3/minimax_h3.json" \ + --prompt "Animate this image with natural synchronized sound." \ + --image_path "${lightx2v_path}/assets/inputs/imgs/img_0.jpg" \ + --save_result_path "${lightx2v_path}/save_results/output_lightx2v_minimax_h3_i2av.mp4" \ + --seed 42 diff --git a/scripts/minimax_h3/run_minimax_h3_l2av.sh b/scripts/minimax_h3/run_minimax_h3_l2av.sh index ba9df6a3d..0e001236c 100755 --- a/scripts/minimax_h3/run_minimax_h3_l2av.sh +++ b/scripts/minimax_h3/run_minimax_h3_l2av.sh @@ -1,22 +1,22 @@ #!/bin/bash # set path firstly -lightx2v_path=/data/nvme6/gushiqiao/codes/LightX2V -model_path=/data/nvme6/gushiqiao/models/MiniMax-H3 +lightx2v_path=/path/to/LightX2V +model_path=/path/to/MiniMax-H3 export CUDA_VISIBLE_DEVICES=0 # set environment variables -source ${lightx2v_path}/scripts/base/base.sh +source "${lightx2v_path}/scripts/base/base.sh" export DTYPE=BF16 export SENSITIVE_LAYER_DTYPE=BF16 python -m lightx2v.infer \ ---model_cls minimax_h3 \ ---task l2av \ ---model_path $model_path \ ---config_json ${lightx2v_path}/configs/minimax_h3/minimax_h3_l2av.json \ ---prompt "Generate the preceding scene with natural synchronized sound." \ ---last_frame_path ${lightx2v_path}/assets/inputs/imgs/flf2v_input_last_frame-fs8.png \ ---save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_l2av.mp4 \ ---seed 42 + --model_cls minimax_h3 \ + --task l2av \ + --model_path "${model_path}" \ + --config_json "${lightx2v_path}/configs/minimax_h3/minimax_h3.json" \ + --prompt "Generate the preceding scene with natural synchronized sound." \ + --last_frame_path "${lightx2v_path}/assets/inputs/imgs/flf2v_input_last_frame-fs8.png" \ + --save_result_path "${lightx2v_path}/save_results/output_lightx2v_minimax_h3_l2av.mp4" \ + --seed 42 diff --git a/scripts/minimax_h3/run_minimax_h3_ref2av.sh b/scripts/minimax_h3/run_minimax_h3_ref2av.sh index 0bc111c51..5bbe48bc2 100755 --- a/scripts/minimax_h3/run_minimax_h3_ref2av.sh +++ b/scripts/minimax_h3/run_minimax_h3_ref2av.sh @@ -1,22 +1,26 @@ #!/bin/bash # set path firstly -lightx2v_path=/data/nvme6/gushiqiao/codes/LightX2V -model_path=/data/nvme6/gushiqiao/models/MiniMax-H3 +lightx2v_path=/path/to/LightX2V +model_path=/path/to/MiniMax-H3 export CUDA_VISIBLE_DEVICES=0 # set environment variables -source ${lightx2v_path}/scripts/base/base.sh +source "${lightx2v_path}/scripts/base/base.sh" export DTYPE=BF16 export SENSITIVE_LAYER_DTYPE=BF16 +# Ref2AV uses transformer_ref/ weights and supports image, video, and audio references. +# Pass multiple files of one type as comma-separated paths. +# Add --video_path "/path/to/reference.mp4" or --audio_path "/path/to/reference.wav" as needed. +# Audio requires at least one image or video reference. python -m lightx2v.infer \ ---model_cls minimax_h3 \ ---task ref2av \ ---model_path $model_path \ ---config_json ${lightx2v_path}/configs/minimax_h3/minimax_h3_ref2av.json \ ---prompt "Generate an audio-video scene following the references." \ ---image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ ---save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_ref2av.mp4 \ ---seed 42 + --model_cls minimax_h3 \ + --task ref2av \ + --model_path "${model_path}" \ + --config_json "${lightx2v_path}/configs/minimax_h3/minimax_h3.json" \ + --prompt "Generate an audio-video scene following the references." \ + --image_path "${lightx2v_path}/assets/inputs/imgs/img_0.jpg" \ + --save_result_path "${lightx2v_path}/save_results/output_lightx2v_minimax_h3_ref2av.mp4" \ + --seed 42 diff --git a/scripts/minimax_h3/run_minimax_h3_t2av.sh b/scripts/minimax_h3/run_minimax_h3_t2av.sh index c00b51901..904dee980 100755 --- a/scripts/minimax_h3/run_minimax_h3_t2av.sh +++ b/scripts/minimax_h3/run_minimax_h3_t2av.sh @@ -1,25 +1,21 @@ #!/bin/bash # set path firstly -lightx2v_path=/data/nvme6/gushiqiao/codes/LightX2V -model_path=/data/nvme6/gushiqiao/models/MiniMax-H3 +lightx2v_path=/path/to/LightX2V +model_path=/path/to/MiniMax-H3 export CUDA_VISIBLE_DEVICES=0 # set environment variables -source ${lightx2v_path}/scripts/base/base.sh +source "${lightx2v_path}/scripts/base/base.sh" export DTYPE=BF16 export SENSITIVE_LAYER_DTYPE=BF16 -prompt='integrated_multimodal_description: [Shot 1] Cinematic low-angle tracking shot following a stylish woman from behind as she strolls down a bustling post-rain Tokyo street. The asphalt is completely wet, acting as a black mirror that perfectly reflects the dense canopy of overhead neon signs—warm pink lanterns, icy cyan katakana signage, and giant animated billboards playing silently. She walks slightly to the left of frame, revealing the back of her sleek black leather jacket, which glistens with specular highlights, and the hem of a flowing crimson dress that swirls around her calves. Black heeled boots splash subtly in shallow puddles, and a black leather purse hangs from her shoulder. The camera slowly pushes forward and gently rises, while out-of-focus pedestrians in modern clothing cross the frame, adding life. [Shot 2] At 00:03.500, a sharp cut to a medium profile shot from her right side, camera dollying sideways in perfect sync. She comes into clear view: oversized black sunglasses perched on her nose reflect a giant LED screen across the street, with purple and blue animations gliding across the lenses. Her bold matte red lipstick stands out against fair skin, and a hint of a confident smile plays on her lips. The sharp tailoring of her jacket catches rim light, and her stride is poised and rhythmic. The background is a bokeh of neon blur, while the wet ground distorts the red dress’s reflection into abstract color streaks. A subtle handheld camera shake increases immediacy. [Shot 3] At 00:06.800, a stylized slow-motion frontal medium close-up as she walks directly toward the lens, which pulls back. Time stretches—she casually removes her sunglasses in one smooth motion, revealing sharp winged eyeliner and a piercing gaze that locks directly with the viewer. A shaft of hot pink neon light sweeps across her cheekbones, then she slides the glasses back on with a soft click. The camera then racks focus from her face to the endless corridor of neon-lit street behind her as she walks past, dissolving into a blur of vibrant city lights. -overall_soundscape: Rich city atmosphere on wet streets: a constant damp hiss of car tires rolling through water in the distance, the resonant electrical hum and faint crackle of neon transformers overhead, a muffled J-pop bassline leaking from a nearby record store, layers of pedestrian chatter and soft laughter in Japanese, and in the foreground, the crisp, wet footsteps of her heeled boots striking the mirrored asphalt, with occasional tiny splashes. When she removes her sunglasses, a delicate, intimate "click" of the frame folding is audible, momentarily cutting through the noise. -non_diegetic_music: A lo-fi electronic city-pop track with a relaxed breakbeat and dreamy analog synth pads, setting a confident, seductive mood. As she takes off her sunglasses in slow motion, a warm, soulful saxophone phrase sweeps in with reverb, then gently settles back into the groove as she walks on, gradually fading out with the ambient hum.' - python -m lightx2v.infer \ ---model_cls minimax_h3 \ ---task t2av \ ---model_path $model_path \ ---config_json ${lightx2v_path}/configs/minimax_h3/fp8/minimax_h3_t2av.json \ ---prompt "$prompt" \ ---save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av10.mp4 \ ---seed 0 + --model_cls minimax_h3 \ + --task t2av \ + --model_path "${model_path}" \ + --config_json "${lightx2v_path}/configs/minimax_h3/minimax_h3.json" \ + --prompt 'integrated_multimodal_description: A cinematic fox walks through a snowy pine forest at dawn. overall_soundscape: Soft wind, crunching snow, and distant birds. non_diegetic_music: Quiet warm strings.' \ + --save_result_path "${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av.mp4" \ + --seed 42 diff --git a/scripts/minimax_h3/run_minimax_h3_t2av_dmd_lora_4step.sh b/scripts/minimax_h3/run_minimax_h3_t2av_dmd_lora_4step.sh deleted file mode 100644 index 06095ec92..000000000 --- a/scripts/minimax_h3/run_minimax_h3_t2av_dmd_lora_4step.sh +++ /dev/null @@ -1,26 +0,0 @@ -#!/bin/bash -set -e - -lightx2v_path=${LIGHTX2V_PATH:-/data/nvme6/gushiqiao/codes/latest/LightX2V} -model_path=${MINIMAX_H3_MODEL_PATH:-/data/nvme6/gushiqiao/models/MiniMax-H3} -python_bin=${PYTHON_BIN:-python} - -if ! command -v "${python_bin}" >/dev/null 2>&1; then - echo "Python executable not found: ${python_bin}" >&2 - exit 1 -fi - -export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0} - -source "${lightx2v_path}/scripts/base/base.sh" -export DTYPE=BF16 -export SENSITIVE_LAYER_DTYPE=FP32 - -PYTHONPATH="${lightx2v_path}" "${python_bin}" -m lightx2v.infer \ - --model_cls minimax_h3 \ - --task t2av \ - --model_path "${model_path}" \ - --config_json "${lightx2v_path}/configs/minimax_h3/minimax_h3_t2av_dmd_lora_4step.json" \ - --prompt "integrated_multimodal_description: Shot 1 Live-action sports documentary style, a handheld ringside close shot frames two female boxers wearing headguards and gloves during controlled sparring in a small neighborhood gym. The camera tracks left at fast speed as the boxer in blue slips a jab, answers with a rapid body-head combination against the pads, and pivots away while sweat droplets catch the overhead light in brief slow motion. Her coach with a firm, energetic voice S1 calls: English Move, breathe, reset! overall_soundscape: Gloves thud against padded targets, shoes squeak on canvas, the round timer beeps once, and both athletes breathe sharply. A jump rope and muffled training activity continue in the background. non_diegetic_music: A tight electronic beat at a fast tempo with clipped bass pulses, fading under the final breath." \ - --save_result_path "${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av_dmd_lora_4step_5s.mp4" \ - --seed 42 diff --git a/scripts/minimax_h3/run_minimax_h3_t2av_parallel.sh b/scripts/minimax_h3/run_minimax_h3_t2av_parallel.sh new file mode 100644 index 000000000..feb5b1326 --- /dev/null +++ b/scripts/minimax_h3/run_minimax_h3_t2av_parallel.sh @@ -0,0 +1,21 @@ +#!/bin/bash + +# set path firstly +lightx2v_path=/path/to/LightX2V +model_path=/path/to/MiniMax-H3 + +export CUDA_VISIBLE_DEVICES=0,1,2,3 + +# set environment variables +source "${lightx2v_path}/scripts/base/base.sh" +export DTYPE=BF16 +export SENSITIVE_LAYER_DTYPE=BF16 + +torchrun --standalone --nproc_per_node=4 -m lightx2v.infer \ + --model_cls minimax_h3 \ + --task t2av \ + --model_path "${model_path}" \ + --config_json "${lightx2v_path}/configs/minimax_h3/minimax_h3_sp.json" \ + --prompt "A cinematic fox walking through a snowy forest" \ + --save_result_path "${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av_parallel.mp4" \ + --seed 42 diff --git a/scripts/minimax_h3/run_minimax_h3_t2av_sol_attn_offload.sh b/scripts/minimax_h3/run_minimax_h3_t2av_sol_attn_offload.sh deleted file mode 100755 index 6f99e4647..000000000 --- a/scripts/minimax_h3/run_minimax_h3_t2av_sol_attn_offload.sh +++ /dev/null @@ -1,29 +0,0 @@ -#!/bin/bash - -# set path firstly -lightx2v_path= -model_path= - -export CUDA_VISIBLE_DEVICES=0 - -# set environment variables -source ${lightx2v_path}/scripts/base/base.sh -export DTYPE=BF16 -export SENSITIVE_LAYER_DTYPE=BF16 -export SOL_ATTN_STRICT=1 - -prompt='integrated_multimodal_description: [Shot 1] Live-action wildlife cinematography, a low-angle medium-wide tracking shot follows a red fox moving purposefully through a dense, snow-covered pine forest at dawn. The camera tracks backward at moderate speed, keeping the fox’s face and amber eyes sharply focused as its paws plunge into fresh powder and scatter fine snow crystals toward the lens. Its thick red-and-white winter coat ripples naturally in the cold wind while visible breath streams from its muzzle. Pale golden sunbeams flicker rapidly across its body as it passes between dark tree trunks. The fox suddenly hears a distant cracking branch, raises its ears, turns sharply to the right, and accelerates into a sprint. -[Shot 2] At 00:05.200, the camera cuts to a fast lateral tracking shot moving parallel to the sprinting fox. It weaves between closely spaced pine trunks, bounds over exposed roots, and ducks beneath a snow-laden branch. Its paws strike the ground in a rapid rhythm, throwing broad sprays of powder behind it. The disturbed branch snaps upward and releases a cascading curtain of snow as the camera passes through the falling crystals. The fox races down a short slope, briefly loses its footing in deep powder, recovers immediately, and launches toward a fallen log. -[Shot 3] At 00:10.300, the shot cuts to a low frontal angle on the opposite side of the log as the fox leaps directly across the frame in brief slow motion, individual snow crystals suspended around its outstretched body. As it lands, the camera arcs left with large amplitude at fast speed and transitions back to normal motion, following the fox into an open forest clearing. A small flock of ravens bursts from the nearby trees and crosses the pale sky while wind drives loose snow through shafts of golden light. The fox slows near the center of the clearing, turns its head toward the distant mountains, then runs into the luminous morning mist as the camera rises rapidly above the treetops to reveal the vast frozen forest. -overall_soundscape: Rapid paws crunch through deep snow, branches scrape against fur, frozen wood cracks, and cascading powder lands in soft layered impacts. The fox breathes faster during the sprint while wind rush intensifies through the trees; raven wings beat overhead and several sharp calls echo across the clearing. -non_diegetic_music: Repeating low-string ostinatos and deep hand-drum pulses gradually accelerate during the chase. A rising French-horn phrase and rapid high strings peak as the fox leaps over the log, then expand into sustained orchestral chords as the camera rises above the forest.' - - -python -m lightx2v.infer \ ---model_cls minimax_h3 \ ---task t2av \ ---model_path $model_path \ ---config_json ${lightx2v_path}/configs/minimax_h3/minimax_h3_t2av_sol_attn_block_offload.json \ ---prompt "$prompt" \ ---save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av10.mp4 \ ---seed 0 diff --git a/scripts/minimax_h3/run_minimax_h3_t2av_sol_attn_offload_4step.sh b/scripts/minimax_h3/run_minimax_h3_t2av_sol_attn_offload_4step.sh deleted file mode 100755 index 0537e7dfd..000000000 --- a/scripts/minimax_h3/run_minimax_h3_t2av_sol_attn_offload_4step.sh +++ /dev/null @@ -1,29 +0,0 @@ -#!/bin/bash - -# set path firstly -lightx2v_path= -model_path= - -export CUDA_VISIBLE_DEVICES=0 - -# set environment variables -source ${lightx2v_path}/scripts/base/base.sh -export DTYPE=BF16 -export SENSITIVE_LAYER_DTYPE=FP32 -export SOL_ATTN_STRICT=1 - -prompt='integrated_multimodal_description: [Shot 1] Live-action wildlife cinematography, a low-angle medium-wide tracking shot follows a red fox moving purposefully through a dense, snow-covered pine forest at dawn. The camera tracks backward at moderate speed, keeping the fox’s face and amber eyes sharply focused as its paws plunge into fresh powder and scatter fine snow crystals toward the lens. Its thick red-and-white winter coat ripples naturally in the cold wind while visible breath streams from its muzzle. Pale golden sunbeams flicker rapidly across its body as it passes between dark tree trunks. The fox suddenly hears a distant cracking branch, raises its ears, turns sharply to the right, and accelerates into a sprint. -[Shot 2] At 00:05.200, the camera cuts to a fast lateral tracking shot moving parallel to the sprinting fox. It weaves between closely spaced pine trunks, bounds over exposed roots, and ducks beneath a snow-laden branch. Its paws strike the ground in a rapid rhythm, throwing broad sprays of powder behind it. The disturbed branch snaps upward and releases a cascading curtain of snow as the camera passes through the falling crystals. The fox races down a short slope, briefly loses its footing in deep powder, recovers immediately, and launches toward a fallen log. -[Shot 3] At 00:10.300, the shot cuts to a low frontal angle on the opposite side of the log as the fox leaps directly across the frame in brief slow motion, individual snow crystals suspended around its outstretched body. As it lands, the camera arcs left with large amplitude at fast speed and transitions back to normal motion, following the fox into an open forest clearing. A small flock of ravens bursts from the nearby trees and crosses the pale sky while wind drives loose snow through shafts of golden light. The fox slows near the center of the clearing, turns its head toward the distant mountains, then runs into the luminous morning mist as the camera rises rapidly above the treetops to reveal the vast frozen forest. -overall_soundscape: Rapid paws crunch through deep snow, branches scrape against fur, frozen wood cracks, and cascading powder lands in soft layered impacts. The fox breathes faster during the sprint while wind rush intensifies through the trees; raven wings beat overhead and several sharp calls echo across the clearing. -non_diegetic_music: Repeating low-string ostinatos and deep hand-drum pulses gradually accelerate during the chase. A rising French-horn phrase and rapid high strings peak as the fox leaps over the log, then expand into sustained orchestral chords as the camera rises above the forest.' - - -python -m lightx2v.infer \ ---model_cls minimax_h3 \ ---task t2av \ ---model_path $model_path \ ---config_json ${lightx2v_path}/configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol_attn_single_gpu_offload.json \ ---prompt "$prompt" \ ---save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av_sol_attn_4step.mp4 \ ---seed 0 diff --git a/scripts/minimax_h3/run_minimax_h3_t2av_sp.sh b/scripts/minimax_h3/run_minimax_h3_t2av_sp.sh deleted file mode 100755 index 8de457ca5..000000000 --- a/scripts/minimax_h3/run_minimax_h3_t2av_sp.sh +++ /dev/null @@ -1,21 +0,0 @@ -#!/bin/bash - -# set path firstly -lightx2v_path=/path/to/LightX2V -model_path=/path/to/models/minimax_h3/h3_hf_bf16 - -export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 - -# set environment variables -source ${lightx2v_path}/scripts/base/base.sh -export DTYPE=BF16 -export SENSITIVE_LAYER_DTYPE=BF16 - -torchrun --standalone --nproc_per_node=8 -m lightx2v.infer \ ---model_cls minimax_h3 \ ---task t2av \ ---model_path $model_path \ ---config_json ${lightx2v_path}/configs/minimax_h3/fp8/minimax_h3_t2av_sp8_5090.json \ ---prompt "A cinematic fox walking through a snowy forest" \ ---save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av_sp.mp4 \ ---seed 42 diff --git a/scripts/minimax_h3/run_minimax_h3_t2av_sp8_4step_5090_with_fp8_vae_sol.sh b/scripts/minimax_h3/run_minimax_h3_t2av_sp8_4step_5090_with_fp8_vae_sol.sh deleted file mode 100755 index 7854f9b13..000000000 --- a/scripts/minimax_h3/run_minimax_h3_t2av_sp8_4step_5090_with_fp8_vae_sol.sh +++ /dev/null @@ -1,27 +0,0 @@ -#!/bin/bash - -lightx2v_path= -model_path= - -export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 - -source ${lightx2v_path}/scripts/base/base.sh -export DTYPE=BF16 -export SENSITIVE_LAYER_DTYPE=BF16 - -torchrun --standalone --nproc_per_node=8 -m lightx2v.infer \ ---model_cls minimax_h3 \ ---task t2av \ ---model_path ${model_path} \ ---config_json ${lightx2v_path}/configs/minimax_h3/dmd/minimax_h3_sp8_4step_5090_with_fp8_vae_sol.json \ ---prompt "integrated_multimodal_description: [Shot 1] Live-action wildlife cinematography, a low-angle medium-wide tracking shot follows a red fox moving purposefully through a dense, snow-covered pine forest at dawn. The camera tracks backward at moderate speed, keeping the fox’s face and amber eyes sharply focused as its paws plunge into fresh powder and scatter fine snow crystals toward the lens. Its thick red-and-white winter coat ripples naturally in the cold wind while visible breath streams from its muzzle. Pale golden sunbeams flicker rapidly across its body as it passes between dark tree trunks. The fox suddenly hears a distant cracking branch, raises its ears, turns sharply to the right, and accelerates into a sprint. - -[Shot 2] At 00:05.200, the camera cuts to a fast lateral tracking shot moving parallel to the sprinting fox. It weaves between closely spaced pine trunks, bounds over exposed roots, and ducks beneath a snow-laden branch. Its paws strike the ground in a rapid rhythm, throwing broad sprays of powder behind it. The disturbed branch snaps upward and releases a cascading curtain of snow as the camera passes through the falling crystals. The fox races down a short slope, briefly loses its footing in deep powder, recovers immediately, and launches toward a fallen log. - -[Shot 3] At 00:10.300, the shot cuts to a low frontal angle on the opposite side of the log as the fox leaps directly across the frame in brief slow motion, individual snow crystals suspended around its outstretched body. As it lands, the camera arcs left with large amplitude at fast speed and transitions back to normal motion, following the fox into an open forest clearing. A small flock of ravens bursts from the nearby trees and crosses the pale sky while wind drives loose snow through shafts of golden light. The fox slows near the center of the clearing, turns its head toward the distant mountains, then runs into the luminous morning mist as the camera rises rapidly above the treetops to reveal the vast frozen forest. - -overall_soundscape: Rapid paws crunch through deep snow, branches scrape against fur, frozen wood cracks, and cascading powder lands in soft layered impacts. The fox breathes faster during the sprint while wind rush intensifies through the trees; raven wings beat overhead and several sharp calls echo across the clearing. - -non_diegetic_music: Repeating low-string ostinatos and deep hand-drum pulses gradually accelerate during the chase. A rising French-horn phrase and rapid high strings peak as the fox leaps over the log, then expand into sustained orchestral chords as the camera rises above the forest. " \ ---save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_sp8_4step_5090_with_fp8_vae_sol.mp4 \ ---seed 42 diff --git a/scripts/minimax_h3/run_minimax_h3_t2av_tp.sh b/scripts/minimax_h3/run_minimax_h3_t2av_tp.sh deleted file mode 100755 index 021e6b480..000000000 --- a/scripts/minimax_h3/run_minimax_h3_t2av_tp.sh +++ /dev/null @@ -1,19 +0,0 @@ -#!/bin/bash - -lightx2v_path=/data/nvme6/gushiqiao/codes/LightX2V -model_path=/data/nvme6/gushiqiao/models/MiniMax-H3 - -export CUDA_VISIBLE_DEVICES=0,1 - -source ${lightx2v_path}/scripts/base/base.sh -export DTYPE=BF16 -export SENSITIVE_LAYER_DTYPE=BF16 - -torchrun --standalone --nproc_per_node=2 -m lightx2v.infer \ ---model_cls minimax_h3 \ ---task t2av \ ---model_path ${model_path} \ ---config_json ${lightx2v_path}/configs/minimax_h3/minimax_h3_t2av_tp.json \ ---prompt "A cinematic fox walking through a snowy forest" \ ---save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av_tp.mp4 \ ---seed 42 diff --git a/scripts/minimax_h3/run_minimax_h3_t2av_tp_sp.sh b/scripts/minimax_h3/run_minimax_h3_t2av_tp_sp.sh deleted file mode 100755 index 0bcdb7acc..000000000 --- a/scripts/minimax_h3/run_minimax_h3_t2av_tp_sp.sh +++ /dev/null @@ -1,19 +0,0 @@ -#!/bin/bash - -lightx2v_path=/data/nvme6/gushiqiao/codes/LightX2V -model_path=/data/nvme6/gushiqiao/models/MiniMax-H3 - -export CUDA_VISIBLE_DEVICES=0,1,2,3 - -source ${lightx2v_path}/scripts/base/base.sh -export DTYPE=BF16 -export SENSITIVE_LAYER_DTYPE=BF16 - -torchrun --standalone --nproc_per_node=4 -m lightx2v.infer \ ---model_cls minimax_h3 \ ---task t2av \ ---model_path ${model_path} \ ---config_json ${lightx2v_path}/configs/minimax_h3/minimax_h3_t2av_tp_sp.json \ ---prompt "A cinematic fox walking through a snowy forest" \ ---save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av_tp_sp.mp4 \ ---seed 42 diff --git a/scripts/minimax_h3/run_minimax_h3_t2av_warmup.sh b/scripts/minimax_h3/run_minimax_h3_t2av_warmup.sh deleted file mode 100755 index b40c61a12..000000000 --- a/scripts/minimax_h3/run_minimax_h3_t2av_warmup.sh +++ /dev/null @@ -1,23 +0,0 @@ -#!/bin/bash - -# set path firstly -lightx2v_path= -model_path= - -export CUDA_VISIBLE_DEVICES=0 - -# set environment variables -source ${lightx2v_path}/scripts/base/base.sh -export DTYPE=BF16 -export SENSITIVE_LAYER_DTYPE=BF16 - -prompt='A cinematic fox walking through a snowy forest.' - -python -m lightx2v.infer \ ---model_cls minimax_h3 \ ---task t2av \ ---model_path $model_path \ ---config_json ${lightx2v_path}/configs/minimax_h3/minimax_h3_t2av_compile.json \ ---prompt "$prompt" \ ---save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av10.mp4 \ ---seed 0 diff --git a/scripts/minimax_h3/server/post_fl2av.py b/scripts/minimax_h3/server/post_fl2av.py new file mode 100644 index 000000000..b44aa4798 --- /dev/null +++ b/scripts/minimax_h3/server/post_fl2av.py @@ -0,0 +1,31 @@ +import base64 + +import requests +from loguru import logger + + +def image_to_base64(image_path): + """Convert an image file to base64 string""" + with open(image_path, "rb") as f: + image_data = f.read() + return base64.b64encode(image_data).decode("utf-8") + + +if __name__ == "__main__": + url = "http://localhost:8000/v1/tasks/video/" + + message = { + "task": "fl2av", + "prompt": "Create a coherent transition with natural synchronized sound.", + "image_path": image_to_base64("assets/inputs/imgs/flf2v_input_first_frame-fs8.png"), + "last_frame_path": image_to_base64("assets/inputs/imgs/flf2v_input_last_frame-fs8.png"), + "seed": 42, + "num_frames": 124, + "target_shape": [544, 960], + "save_result_path": "./minimax_h3_fl2av.mp4", + } + + logger.info(f"message: {message}") + response = requests.post(url, json=message) + response.raise_for_status() + logger.info(f"response: {response.json()}") diff --git a/scripts/minimax_h3/server/post_i2av.py b/scripts/minimax_h3/server/post_i2av.py new file mode 100644 index 000000000..6a971046d --- /dev/null +++ b/scripts/minimax_h3/server/post_i2av.py @@ -0,0 +1,30 @@ +import base64 + +import requests +from loguru import logger + + +def image_to_base64(image_path): + """Convert an image file to base64 string""" + with open(image_path, "rb") as f: + image_data = f.read() + return base64.b64encode(image_data).decode("utf-8") + + +if __name__ == "__main__": + url = "http://localhost:8000/v1/tasks/video/" + + message = { + "task": "i2av", + "prompt": "Animate this image with natural synchronized sound.", + "image_path": image_to_base64("assets/inputs/imgs/img_0.jpg"), + "seed": 42, + "num_frames": 124, + "target_shape": [544, 960], + "save_result_path": "./minimax_h3_i2av.mp4", + } + + logger.info(f"message: {message}") + response = requests.post(url, json=message) + response.raise_for_status() + logger.info(f"response: {response.json()}") diff --git a/scripts/minimax_h3/server/post_l2av.py b/scripts/minimax_h3/server/post_l2av.py new file mode 100644 index 000000000..575c230fd --- /dev/null +++ b/scripts/minimax_h3/server/post_l2av.py @@ -0,0 +1,30 @@ +import base64 + +import requests +from loguru import logger + + +def image_to_base64(image_path): + """Convert an image file to base64 string""" + with open(image_path, "rb") as f: + image_data = f.read() + return base64.b64encode(image_data).decode("utf-8") + + +if __name__ == "__main__": + url = "http://localhost:8000/v1/tasks/video/" + + message = { + "task": "l2av", + "prompt": "Generate the preceding scene with natural synchronized sound.", + "last_frame_path": image_to_base64("assets/inputs/imgs/flf2v_input_last_frame-fs8.png"), + "seed": 42, + "num_frames": 124, + "target_shape": [544, 960], + "save_result_path": "./minimax_h3_l2av.mp4", + } + + logger.info(f"message: {message}") + response = requests.post(url, json=message) + response.raise_for_status() + logger.info(f"response: {response.json()}") diff --git a/scripts/minimax_h3/server/post_ref2av.py b/scripts/minimax_h3/server/post_ref2av.py new file mode 100644 index 000000000..310959a02 --- /dev/null +++ b/scripts/minimax_h3/server/post_ref2av.py @@ -0,0 +1,33 @@ +import base64 + +import requests +from loguru import logger + + +def image_to_base64(image_path): + """Convert an image file to base64 string""" + with open(image_path, "rb") as f: + image_data = f.read() + return base64.b64encode(image_data).decode("utf-8") + + +if __name__ == "__main__": + url = "http://localhost:8000/v1/tasks/video/" + + message = { + "task": "ref2av", + "prompt": "Generate an audio-video scene following the references.", + "image_path": image_to_base64("assets/inputs/imgs/img_0.jpg"), + # Video/audio paths below are local to the server. Audio needs a visual reference. + # "video_path": "/path/to/reference.mp4", + # "audio_path": "/path/to/reference.wav", + "seed": 42, + "num_frames": 124, + "target_shape": [544, 960], + "save_result_path": "./minimax_h3_ref2av.mp4", + } + + logger.info(f"message: {message}") + response = requests.post(url, json=message) + response.raise_for_status() + logger.info(f"response: {response.json()}") diff --git a/scripts/minimax_h3/server/post_t2av.py b/scripts/minimax_h3/server/post_t2av.py new file mode 100644 index 000000000..1aa94c4fb --- /dev/null +++ b/scripts/minimax_h3/server/post_t2av.py @@ -0,0 +1,19 @@ +import requests +from loguru import logger + +if __name__ == "__main__": + url = "http://localhost:8000/v1/tasks/video/" + + message = { + "task": "t2av", + "prompt": "integrated_multimodal_description: A cinematic fox walks through a snowy pine forest at dawn. overall_soundscape: Soft wind, crunching snow, and distant birds. non_diegetic_music: Quiet warm strings.", + "seed": 42, + "num_frames": 124, + "target_shape": [544, 960], + "save_result_path": "./minimax_h3_t2av.mp4", + } + + logger.info(f"message: {message}") + response = requests.post(url, json=message) + response.raise_for_status() + logger.info(f"response: {response.json()}") diff --git a/scripts/minimax_h3/server/start_server.sh b/scripts/minimax_h3/server/start_server.sh new file mode 100644 index 000000000..25c309a42 --- /dev/null +++ b/scripts/minimax_h3/server/start_server.sh @@ -0,0 +1,19 @@ +#!/bin/bash + +lightx2v_path=/path/to/LightX2V +model_path=/path/to/MiniMax-H3 + +export CUDA_VISIBLE_DEVICES=0 +source "${lightx2v_path}/scripts/base/base.sh" + +export DTYPE=BF16 +export SENSITIVE_LAYER_DTYPE=BF16 + +# The base transformer serves t2av, i2av, l2av, and fl2av requests. +python -m lightx2v.server \ + --model_cls minimax_h3 \ + --task t2av \ + --model_path "${model_path}" \ + --config_json "${lightx2v_path}/configs/minimax_h3/minimax_h3.json" \ + --host 0.0.0.0 \ + --port 8000 diff --git a/scripts/minimax_h3/server/start_server_ref2av.sh b/scripts/minimax_h3/server/start_server_ref2av.sh new file mode 100644 index 000000000..33a0191ad --- /dev/null +++ b/scripts/minimax_h3/server/start_server_ref2av.sh @@ -0,0 +1,19 @@ +#!/bin/bash + +lightx2v_path=/path/to/LightX2V +model_path=/path/to/MiniMax-H3 + +export CUDA_VISIBLE_DEVICES=0 +source "${lightx2v_path}/scripts/base/base.sh" + +export DTYPE=BF16 +export SENSITIVE_LAYER_DTYPE=BF16 + +# Reference generation loads transformer_ref instead of the base transformer. +python -m lightx2v.server \ + --model_cls minimax_h3 \ + --task ref2av \ + --model_path "${model_path}" \ + --config_json "${lightx2v_path}/configs/minimax_h3/minimax_h3.json" \ + --host 0.0.0.0 \ + --port 8000 diff --git a/scripts/seko_talk/shot/run_rs2v.sh b/scripts/seko_talk/shot/run_rs2v.sh index 717587fd0..360361764 100755 --- a/scripts/seko_talk/shot/run_rs2v.sh +++ b/scripts/seko_talk/shot/run_rs2v.sh @@ -12,9 +12,7 @@ export CUDA_VISIBLE_DEVICES=0 python -m lightx2v.shot_runner.rs2v_infer \ --config_json ${lightx2v_path}/configs/seko_talk/shot/rs2v/main.json \ --prompt "The video features a male speaking to the camera with arms spread out, a slightly furrowed brow, and a focused gaze." \ ---negative_prompt 色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走 \ --image_path ${lightx2v_path}/assets/inputs/audio/seko_input.png \ --audio_path ${lightx2v_path}/assets/inputs/audio/seko_input.mp3 \ --save_result_path ${lightx2v_path}/save_results/output_seko_talk_shot_rs2v.mp4 \ ---infer_steps 4 \ --video_duration 10 diff --git a/scripts/seko_talk/shot/run_rs2v_dist.sh b/scripts/seko_talk/shot/run_rs2v_dist.sh index 9335a9396..f0e9dac65 100755 --- a/scripts/seko_talk/shot/run_rs2v_dist.sh +++ b/scripts/seko_talk/shot/run_rs2v_dist.sh @@ -12,7 +12,6 @@ export CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc-per-node 2 -m lightx2v.shot_runner.rs2v_infer \ --config_json ${lightx2v_path}/configs/seko_talk/shot/rs2v/main_dist.json \ --prompt "A cultured woman speaking passionately and eloquently, her expression alive with emotion, conveying resolve, dignity, and a deep sense of purpose." \ ---negative_prompt 色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走 \ --image_path ${lightx2v_path}/assets/inputs/audio/seko_input.png \ --audio_path ${lightx2v_path}/assets/inputs/audio/seko_input.mp3 \ --save_result_path ${lightx2v_path}/save_results/output_seko_talk_shot_rs2v.mp4 diff --git a/scripts/seko_talk/shot/run_rs2v_multi_person.sh b/scripts/seko_talk/shot/run_rs2v_multi_person.sh index 288f2908c..e403a1102 100755 --- a/scripts/seko_talk/shot/run_rs2v_multi_person.sh +++ b/scripts/seko_talk/shot/run_rs2v_multi_person.sh @@ -12,9 +12,7 @@ export CUDA_VISIBLE_DEVICES=0 python -m lightx2v.shot_runner.rs2v_infer \ --config_json ${lightx2v_path}/configs/seko_talk/shot/rs2v/main.json \ --prompt "The video features a man and a woman standing by a bench in the park, their expressions tense and voices raised as they argue. The man gestures with both hands, his arms swinging slightly as if to emphasize each heated word, while the woman stands with her hands on her waist, her brows furrowed in frustration. The background is a wide expanse of sunlit grass, the golden light contrasting with the sharp energy of their quarrel. Their voices seem to clash in the air, and the rhythm of their hand movements and body postures interweaves with the rising tension, creating a vivid scene of confrontation." \ ---negative_prompt 色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走 \ --image_path ${lightx2v_path}/assets/inputs/audio/multi_person/seko_input.png \ --audio_path ${lightx2v_path}/assets/inputs/audio/multi_person \ --save_result_path ${lightx2v_path}/save_results/output_seko_talk_shot_rs2v_multi_person.mp4 \ ---infer_steps 3 \ --video_duration 20 diff --git a/scripts/seko_talk/shot/run_stream.sh b/scripts/seko_talk/shot/run_stream.sh index fd6853af6..62eae7a59 100755 --- a/scripts/seko_talk/shot/run_stream.sh +++ b/scripts/seko_talk/shot/run_stream.sh @@ -12,7 +12,6 @@ export CUDA_VISIBLE_DEVICES=0 python -m lightx2v.shot_runner.stream_infer \ --config_json ${lightx2v_path}/configs/seko_talk/shot/stream/main.json \ --prompt "The video features a male speaking to the camera with arms spread out, a slightly furrowed brow, and a focused gaze." \ ---negative_prompt 色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走 \ --image_path ${lightx2v_path}/assets/inputs/audio/seko_input.png \ --audio_path ${lightx2v_path}/assets/inputs/audio/seko_input.mp3 \ --save_result_path ${lightx2v_path}/save_results/output_seko_talk_shot_stream.mp4 diff --git a/scripts/server/benchmark_sync_s3_latency.py b/scripts/server/benchmark_sync_s3_latency.py index 95ac1e19e..306875308 100644 --- a/scripts/server/benchmark_sync_s3_latency.py +++ b/scripts/server/benchmark_sync_s3_latency.py @@ -81,12 +81,12 @@ def build_sync_payload(args: argparse.Namespace, presigned_url: str = "") -> Dic payload: Dict[str, Any] = { "prompt": args.prompt, "negative_prompt": args.negative_prompt, - "infer_steps": args.infer_steps, "seed": args.seed, "aspect_ratio": args.aspect_ratio, "save_result_path": args.save_result_path, } - if args.target_shape: + payload = {key: value for key, value in payload.items() if value is not None} + if args.target_shape is not None: payload["target_shape"] = args.target_shape if presigned_url: payload["presigned_url"] = presigned_url @@ -176,12 +176,11 @@ def main() -> None: parser.add_argument("--order", type=str, default="alternate", choices=["alternate", "client_first", "server_first"]) parser.add_argument("--prompt", type=str, required=True, help="Prompt text") - parser.add_argument("--negative_prompt", type=str, default="", help="Negative prompt text") - parser.add_argument("--infer_steps", type=int, default=30, help="Inference steps") - parser.add_argument("--seed", type=int, default=42, help="Random seed") - parser.add_argument("--aspect_ratio", type=str, default="16:9", help="Aspect ratio") + parser.add_argument("--negative_prompt", type=str, default=None, help="Negative prompt text") + parser.add_argument("--seed", type=int, default=None, help="Random seed") + parser.add_argument("--aspect_ratio", type=str, default=None, help="Aspect ratio") parser.add_argument("--target_shape", type=int, nargs="+", default=None, help="Target shape, e.g. 1536 2752") - parser.add_argument("--save_result_path", type=str, default="", help="Server-side save_result_path") + parser.add_argument("--save_result_path", type=str, default=None, help="Server-side save_result_path") parser.add_argument("--timeout_seconds", type=int, default=600) parser.add_argument("--poll_interval_seconds", type=float, default=0.5) diff --git a/scripts/server/disagg/qwen/post_qwen_i2i.py b/scripts/server/disagg/qwen/post_qwen_i2i.py index 5059fdbdf..0c1a377fe 100644 --- a/scripts/server/disagg/qwen/post_qwen_i2i.py +++ b/scripts/server/disagg/qwen/post_qwen_i2i.py @@ -47,7 +47,6 @@ def poll_task(url, task_id, timeout=300, interval=5): payload = { "prompt": "Change the person to a standing position, bending over to hold the dog's front paws.", - "negative_prompt": "", "image_path": image_to_base64(IMAGE_PATH), "seed": 42, "save_result_path": "save_results/qwen_i2i_disagg_3way.png", diff --git a/scripts/server/post_async_t2i_and_wait.py b/scripts/server/post_async_t2i_and_wait.py index e0fa76a1c..a01541c1b 100644 --- a/scripts/server/post_async_t2i_and_wait.py +++ b/scripts/server/post_async_t2i_and_wait.py @@ -9,22 +9,21 @@ def submit_t2i_task( base_url: str, prompt: str, - negative_prompt: str, - infer_steps: int, - seed: int, - aspect_ratio: str, + negative_prompt: Optional[str], + seed: Optional[int], + aspect_ratio: Optional[str], target_shape: Optional[List[int]], - save_result_path: str, + save_result_path: Optional[str], ) -> str: payload = { "prompt": prompt, "negative_prompt": negative_prompt, - "infer_steps": infer_steps, "seed": seed, "aspect_ratio": aspect_ratio, "save_result_path": save_result_path, } - if target_shape: + payload = {key: value for key, value in payload.items() if value is not None} + if target_shape is not None: payload["target_shape"] = target_shape submit_url = f"{base_url.rstrip('/')}/v1/tasks/image/" @@ -78,10 +77,9 @@ def main(): parser = argparse.ArgumentParser(description="Submit T2I task to /v1/tasks/image/ and wait for final result.") parser.add_argument("--url", type=str, default="http://127.0.0.1:8000", help="Server base url") parser.add_argument("--prompt", type=str, required=True, help="Prompt text") - parser.add_argument("--negative_prompt", type=str, default="", help="Negative prompt text") - parser.add_argument("--infer_steps", type=int, default=30, help="Inference steps") - parser.add_argument("--seed", type=int, default=42, help="Random seed") - parser.add_argument("--aspect_ratio", type=str, default="16:9", help="Aspect ratio for image task") + parser.add_argument("--negative_prompt", type=str, default=None, help="Negative prompt text") + parser.add_argument("--seed", type=int, default=None, help="Random seed") + parser.add_argument("--aspect_ratio", type=str, default=None, help="Aspect ratio for image task") parser.add_argument( "--target_shape", type=int, @@ -89,7 +87,7 @@ def main(): default=None, help="Target output shape, e.g. --target_shape 1536 2752", ) - parser.add_argument("--save_result_path", type=str, default="", help="Server-side save_result_path") + parser.add_argument("--save_result_path", type=str, default=None, help="Server-side save_result_path") parser.add_argument("--timeout_seconds", type=int, default=600, help="Polling timeout in seconds") parser.add_argument("--poll_interval", type=float, default=2.0, help="Polling interval in seconds") parser.add_argument("--output", type=str, default="save_results/t2i_result.png", help="Local output image path") @@ -100,7 +98,6 @@ def main(): base_url=args.url, prompt=args.prompt, negative_prompt=args.negative_prompt, - infer_steps=args.infer_steps, seed=args.seed, aspect_ratio=args.aspect_ratio, target_shape=args.target_shape, @@ -116,8 +113,11 @@ def main(): ) print(f"Task completed: {final_status}") - output_path = download_result(args.url, task_id, args.output) - print(f"Result saved to: {output_path}") + if final_status["save_result_path"] is not None: + output_path = download_result(args.url, task_id, args.output) + print(f"Result saved to: {output_path}") + else: + print("No file was saved; provide --save_result_path to download a result.") if __name__ == "__main__": diff --git a/scripts/server/post_seko_talk_ar.py b/scripts/server/post_seko_talk_ar.py index 09b903609..fb723854d 100644 --- a/scripts/server/post_seko_talk_ar.py +++ b/scripts/server/post_seko_talk_ar.py @@ -12,12 +12,6 @@ "gestures. The camera is fixed in a static medium shot." ) -DEFAULT_NEGATIVE_PROMPT = ( - "low quality, blurry, pixelated, low resolution, noise, artifacts, poor lighting, " - "overexposed, underexposed, distorted, unnatural, deformed, watermark, logo, text, " - "bad hands, malformed hands, missing fingers, static" -) - def submit_task(args) -> str: url = f"{args.url.rstrip('/')}/v1/tasks/video/" @@ -28,11 +22,9 @@ def submit_task(args) -> str: "audio_path": args.audio_path, "save_result_path": args.save_result_path, "seed": args.seed, - "infer_steps": args.infer_steps, "video_duration": args.video_duration, - "target_fps": args.target_fps, - "resize_mode": args.resize_mode, } + message = {key: value for key, value in message.items() if value is not None} if args.target_shape: message["target_shape"] = args.target_shape @@ -82,17 +74,14 @@ def download_result(base_url: str, task_id: str, output: str) -> Path: def parse_args(): parser = argparse.ArgumentParser(description="Submit a Seko Talk AR rs2v task to LightX2V server.") parser.add_argument("--url", type=str, default="http://127.0.0.1:8000", help="Server base URL") - parser.add_argument("--image_path", type=str, default="/data/nvme4/gushiqiao/new/example/1_素材图.png", help="Reference image path, URL, or base64") - parser.add_argument("--audio_path", type=str, default="/data/nvme4/gushiqiao/new/example/1_素材图.mp3", help="Driving audio path, URL, or base64") + parser.add_argument("--image_path", type=str, default="/path/to/reference.png", help="Reference image path, URL, or base64") + parser.add_argument("--audio_path", type=str, default="/path/to/audio.mp3", help="Driving audio path, URL, or base64") parser.add_argument("--prompt", type=str, default=DEFAULT_PROMPT) - parser.add_argument("--negative_prompt", type=str, default=DEFAULT_NEGATIVE_PROMPT) - parser.add_argument("--save_result_path", type=str, default="seko_talk_ar_server_test.mp4", help="Server-side output filename/path") + parser.add_argument("--negative_prompt", type=str, default=None, help="Optional negative prompt for CFG-enabled deployments") + parser.add_argument("--save_result_path", type=str, default=None, help="Server-side output filename/path") parser.add_argument("--output", type=str, default="save_results/seko_talk_ar_server_test.mp4", help="Downloaded result path") - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--infer_steps", type=int, default=4) + parser.add_argument("--seed", type=int, default=None) parser.add_argument("--video_duration", type=int, default=5) - parser.add_argument("--target_fps", type=int, default=16) - parser.add_argument("--resize_mode", type=str, default="fixed_shape") parser.add_argument("--target_shape", type=int, nargs=2, default=None, help="Optional target shape: H W") parser.add_argument("--timeout_seconds", type=int, default=1800) parser.add_argument("--poll_interval", type=float, default=2.0) @@ -106,8 +95,11 @@ def main(): final_status = wait_task_done(args.url, task_id, args.timeout_seconds, args.poll_interval) logger.info(f"final status: {final_status}") if not args.no_download: - output_path = download_result(args.url, task_id, args.output) - logger.info(f"result saved to: {output_path}") + if final_status["save_result_path"] is not None: + output_path = download_result(args.url, task_id, args.output) + logger.info(f"result saved to: {output_path}") + else: + logger.info("No file was saved; provide --save_result_path to download a result.") if __name__ == "__main__": diff --git a/scripts/server/post_sync_i2i_base64.py b/scripts/server/post_sync_i2i_base64.py index cd34b6942..b982eba4c 100644 --- a/scripts/server/post_sync_i2i_base64.py +++ b/scripts/server/post_sync_i2i_base64.py @@ -14,13 +14,12 @@ def main(): parser = argparse.ArgumentParser(description="Call /v1/tasks/image/sync with base64 image inputs.") parser.add_argument("--url", type=str, default="http://127.0.0.1:8000", help="Server base url") parser.add_argument("--prompt", type=str, required=True, help="Prompt text") - parser.add_argument("--negative_prompt", type=str, default="", help="Negative prompt text") - parser.add_argument("--infer_steps", type=int, default=30, help="Inference steps") - parser.add_argument("--seed", type=int, default=42, help="Random seed") - parser.add_argument("--aspect_ratio", type=str, default="16:9", help="Aspect ratio for image task") + parser.add_argument("--negative_prompt", type=str, default=None, help="Negative prompt text") + parser.add_argument("--seed", type=int, default=None, help="Random seed") + parser.add_argument("--aspect_ratio", type=str, default=None, help="Aspect ratio for image task") parser.add_argument("--timeout_seconds", type=int, default=600, help="Sync API timeout_seconds") parser.add_argument("--poll_interval_seconds", type=float, default=0.5, help="Sync API poll_interval_seconds") - parser.add_argument("--save_result_path", type=str, default="", help="Server-side save_result_path") + parser.add_argument("--save_result_path", type=str, default=None, help="Server-side save_result_path") parser.add_argument("--output", type=str, default="sync_result.png", help="Local output image path") # Base64 inputs (preferred) @@ -47,11 +46,11 @@ def main(): "prompt": args.prompt, "negative_prompt": args.negative_prompt, "image_path": image_base64, - "infer_steps": args.infer_steps, "seed": args.seed, "aspect_ratio": args.aspect_ratio, "save_result_path": args.save_result_path, } + payload = {key: value for key, value in payload.items() if value is not None} if image_mask_base64: payload["image_mask_path"] = image_mask_base64 diff --git a/scripts/server/post_sync_i2i_presigned.py b/scripts/server/post_sync_i2i_presigned.py index d4dfa78c9..e516ba7a8 100644 --- a/scripts/server/post_sync_i2i_presigned.py +++ b/scripts/server/post_sync_i2i_presigned.py @@ -80,13 +80,14 @@ def build_payload(args: argparse.Namespace) -> Dict[str, Any]: payload: Dict[str, Any] = { "prompt": args.prompt, "negative_prompt": args.negative_prompt, - "infer_steps": args.infer_steps, "seed": args.seed, "aspect_ratio": args.aspect_ratio, "save_result_path": args.save_result_path, "presigned_url": args.presigned_url, } + payload = {key: value for key, value in payload.items() if value is not None} + image_base64 = args.image_base64 if not image_base64 and args.image_path: image_base64 = file_to_base64(args.image_path) @@ -111,13 +112,12 @@ def main() -> None: parser = argparse.ArgumentParser(description="Call /v1/tasks/image/sync with presigned_url upload.") parser.add_argument("--url", type=str, default="http://127.0.0.1:8000", help="Server base url") parser.add_argument("--prompt", type=str, required=True, help="Prompt text") - parser.add_argument("--negative_prompt", type=str, default="", help="Negative prompt text") - parser.add_argument("--infer_steps", type=int, default=30, help="Inference steps") - parser.add_argument("--seed", type=int, default=42, help="Random seed") - parser.add_argument("--aspect_ratio", type=str, default="16:9", help="Aspect ratio for image task") + parser.add_argument("--negative_prompt", type=str, default=None, help="Negative prompt text") + parser.add_argument("--seed", type=int, default=None, help="Random seed") + parser.add_argument("--aspect_ratio", type=str, default=None, help="Aspect ratio for image task") parser.add_argument("--timeout_seconds", type=int, default=600, help="Sync API timeout_seconds") parser.add_argument("--poll_interval_seconds", type=float, default=0.5, help="Sync API poll_interval_seconds") - parser.add_argument("--save_result_path", type=str, default="", help="Server-side save_result_path") + parser.add_argument("--save_result_path", type=str, default=None, help="Server-side save_result_path") parser.add_argument("--presigned_url", type=str, default="", help="Presigned URL used by server to upload final PNG") parser.add_argument("--s3_endpoint_url", type=str, default="", help="S3 compatible endpoint, e.g. https://s3.amazonaws.com") parser.add_argument("--s3_region", type=str, default="", help="S3 region, defaults to AWS_DEFAULT_REGION or us-east-1") diff --git a/scripts/server/post_sync_t2i_base64.py b/scripts/server/post_sync_t2i_base64.py index 0b8958293..c7cdec653 100644 --- a/scripts/server/post_sync_t2i_base64.py +++ b/scripts/server/post_sync_t2i_base64.py @@ -9,12 +9,12 @@ def build_payload(args: argparse.Namespace) -> dict: payload = { "prompt": args.prompt, "negative_prompt": args.negative_prompt, - "infer_steps": args.infer_steps, "seed": args.seed, "aspect_ratio": args.aspect_ratio, "save_result_path": args.save_result_path, } - if args.target_shape: + payload = {key: value for key, value in payload.items() if value is not None} + if args.target_shape is not None: payload["target_shape"] = args.target_shape return payload @@ -43,10 +43,9 @@ def main() -> None: parser = argparse.ArgumentParser(description="Call /v1/tasks/image/sync for T2I and save final image.") parser.add_argument("--url", type=str, default="http://127.0.0.1:8000", help="Server base url") parser.add_argument("--prompt", type=str, required=True, help="Prompt text") - parser.add_argument("--negative_prompt", type=str, default="", help="Negative prompt text") - parser.add_argument("--infer_steps", type=int, default=30, help="Inference steps") - parser.add_argument("--seed", type=int, default=42, help="Random seed") - parser.add_argument("--aspect_ratio", type=str, default="16:9", help="Aspect ratio for image task") + parser.add_argument("--negative_prompt", type=str, default=None, help="Negative prompt text") + parser.add_argument("--seed", type=int, default=None, help="Random seed") + parser.add_argument("--aspect_ratio", type=str, default=None, help="Aspect ratio for image task") parser.add_argument( "--target_shape", type=int, @@ -54,7 +53,7 @@ def main() -> None: default=None, help="Target output shape, e.g. --target_shape 1536 2752", ) - parser.add_argument("--save_result_path", type=str, default="", help="Server-side save_result_path") + parser.add_argument("--save_result_path", type=str, default=None, help="Server-side save_result_path") parser.add_argument("--timeout_seconds", type=int, default=600, help="Sync API timeout_seconds") parser.add_argument("--poll_interval_seconds", type=float, default=0.5, help="Sync API poll_interval_seconds") parser.add_argument("--output", type=str, default="save_results/t2i_sync_result.png", help="Local output image path") diff --git a/scripts/server/post_sync_t2i_presigned.py b/scripts/server/post_sync_t2i_presigned.py index 1784f8bb6..0ba5df6eb 100644 --- a/scripts/server/post_sync_t2i_presigned.py +++ b/scripts/server/post_sync_t2i_presigned.py @@ -81,13 +81,13 @@ def build_payload(args: argparse.Namespace) -> Dict[str, Any]: payload: Dict[str, Any] = { "prompt": args.prompt, "negative_prompt": args.negative_prompt, - "infer_steps": args.infer_steps, "seed": args.seed, "aspect_ratio": args.aspect_ratio, "save_result_path": args.save_result_path, "presigned_url": args.presigned_url, } - if args.target_shape: + payload = {key: value for key, value in payload.items() if value is not None} + if args.target_shape is not None: payload["target_shape"] = args.target_shape return payload @@ -111,10 +111,9 @@ def main() -> None: parser = argparse.ArgumentParser(description="Call /v1/tasks/image/sync for T2I with presigned_url upload.") parser.add_argument("--url", type=str, default="http://127.0.0.1:8000", help="Server base url") parser.add_argument("--prompt", type=str, required=True, help="Prompt text") - parser.add_argument("--negative_prompt", type=str, default="", help="Negative prompt text") - parser.add_argument("--infer_steps", type=int, default=30, help="Inference steps") - parser.add_argument("--seed", type=int, default=42, help="Random seed") - parser.add_argument("--aspect_ratio", type=str, default="16:9", help="Aspect ratio for image task") + parser.add_argument("--negative_prompt", type=str, default=None, help="Negative prompt text") + parser.add_argument("--seed", type=int, default=None, help="Random seed") + parser.add_argument("--aspect_ratio", type=str, default=None, help="Aspect ratio for image task") parser.add_argument( "--target_shape", type=int, @@ -122,7 +121,7 @@ def main() -> None: default=None, help="Target output shape, e.g. --target_shape 1536 2752", ) - parser.add_argument("--save_result_path", type=str, default="", help="Server-side save_result_path") + parser.add_argument("--save_result_path", type=str, default=None, help="Server-side save_result_path") parser.add_argument("--timeout_seconds", type=int, default=600, help="Sync API timeout_seconds") parser.add_argument("--poll_interval_seconds", type=float, default=0.5, help="Sync API poll_interval_seconds") parser.add_argument("--presigned_url", type=str, default="", help="Presigned URL used by server to upload final PNG") diff --git a/scripts/server/start_server_i2i.sh b/scripts/server/start_server_i2i.sh index 2c31d9157..f5f1046e5 100755 --- a/scripts/server/start_server_i2i.sh +++ b/scripts/server/start_server_i2i.sh @@ -21,6 +21,5 @@ echo "Service stopped" # { # "prompt": "turn the style of the photo to vintage comic book", -# "image_path": "assets/inputs/imgs/snake.png", -# "infer_steps": 50 +# "image_path": "assets/inputs/imgs/snake.png" # } diff --git a/scripts/server/start_server_t2i.sh b/scripts/server/start_server_t2i.sh index c31ef2812..5f2bea9fd 100755 --- a/scripts/server/start_server_t2i.sh +++ b/scripts/server/start_server_t2i.sh @@ -23,6 +23,5 @@ echo "Service stopped" # { # "prompt": "a beautiful sunset over the ocean", -# "aspect_ratio": "16:9", -# "infer_steps": 50 +# "aspect_ratio": "16:9" # } diff --git a/scripts/server/sync_api_usage.md b/scripts/server/sync_api_usage.md index af91b7bae..bfc83e4b6 100644 --- a/scripts/server/sync_api_usage.md +++ b/scripts/server/sync_api_usage.md @@ -2,6 +2,8 @@ 本文档说明如何调用 `POST /v1/tasks/image/sync` 接口。 +文本、输入媒体和 seed 由请求提供,seed 省略时使用 42。尺寸、帧数、宽高比等规格可继承部署 JSON 的默认值,并由请求覆盖。`negative_prompt` 省略时不发送该字段;显式值(包括空字符串)会按当前模型的能力校验,开启 CFG 的模型仍可能对空字符串应用默认模板。省略 `save_result_path` 或传 null 时不保存文件,同步图片接口继续从内存返回图片。 + ## 1. 接口说明 - **接口**:`POST /v1/tasks/image/sync` @@ -21,8 +23,6 @@ curl -X POST "http://127.0.0.1:8000/v1/tasks/image/sync?timeout_seconds=600&poll -H "Content-Type: application/json" \ -d '{ "prompt": "a cute cat, studio light", - "negative_prompt": "", - "infer_steps": 30, "seed": 42, "aspect_ratio": "16:9" }' \ @@ -38,8 +38,6 @@ url = "http://127.0.0.1:8000/v1/tasks/image/sync" params = {"timeout_seconds": 600, "poll_interval_seconds": 0.5} payload = { "prompt": "a cute cat, studio light", - "negative_prompt": "", - "infer_steps": 30, "seed": 42, "aspect_ratio": "16:9", } @@ -63,8 +61,6 @@ curl -X POST "http://127.0.0.1:8000/v1/tasks/image/sync?timeout_seconds=600&poll -H "Content-Type: application/json" \ -d '{ "prompt": "a cute cat, studio light", - "negative_prompt": "", - "infer_steps": 30, "seed": 42, "aspect_ratio": "16:9", "presigned_url": "https://your-presigned-put-url" @@ -91,8 +87,6 @@ url = "http://127.0.0.1:8000/v1/tasks/image/sync" params = {"timeout_seconds": 600, "poll_interval_seconds": 0.5} payload = { "prompt": "a cute cat, studio light", - "negative_prompt": "", - "infer_steps": 30, "seed": 42, "aspect_ratio": "16:9", "presigned_url": "https://your-presigned-put-url", diff --git a/scripts/wan/distill/run_wan_i2v_distill_4step_cfg.sh b/scripts/wan/distill/run_wan_i2v_distill_4step_cfg.sh index 1c5a09f5b..fa75ed2af 100755 --- a/scripts/wan/distill/run_wan_i2v_distill_4step_cfg.sh +++ b/scripts/wan/distill/run_wan_i2v_distill_4step_cfg.sh @@ -13,5 +13,7 @@ python -m lightx2v.infer \ --config_json ${lightx2v_path}/configs/distill/wan_i2v_distill_4step_cfg.json \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ + --num_frames 81 \ + --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_i2v_distill.mp4 diff --git a/scripts/wan/distill/run_wan_i2v_distill_fp8_4step_cfg.sh b/scripts/wan/distill/run_wan_i2v_distill_fp8_4step_cfg.sh index 99c346898..16b7f924c 100755 --- a/scripts/wan/distill/run_wan_i2v_distill_fp8_4step_cfg.sh +++ b/scripts/wan/distill/run_wan_i2v_distill_fp8_4step_cfg.sh @@ -13,5 +13,7 @@ python -m lightx2v.infer \ --config_json ${lightx2v_path}/configs/distill/wan21/wan_i2v_distill_fp8_4step_cfg.json \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ + --num_frames 81 \ + --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_i2v_distill_fp8_4step_cfg.mp4 diff --git a/scripts/wan/distill/run_wan_i2v_distill_fp8_4step_cfg_ulysses.sh b/scripts/wan/distill/run_wan_i2v_distill_fp8_4step_cfg_ulysses.sh index 916633523..1c825c665 100755 --- a/scripts/wan/distill/run_wan_i2v_distill_fp8_4step_cfg_ulysses.sh +++ b/scripts/wan/distill/run_wan_i2v_distill_fp8_4step_cfg_ulysses.sh @@ -13,5 +13,7 @@ torchrun --nproc_per_node=8 -m lightx2v.infer \ --config_json ${lightx2v_path}/configs/distill/wan21/wan_i2v_distill_fp8_4step_cfg_ulysses.json \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ + --num_frames 81 \ + --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_i2v_distill_fp8_4step_cfg_ulysses.mp4 diff --git a/scripts/wan/distill/run_wan_i2v_distill_int8_4step_cfg.sh b/scripts/wan/distill/run_wan_i2v_distill_int8_4step_cfg.sh index cdd4325f0..d70995e14 100755 --- a/scripts/wan/distill/run_wan_i2v_distill_int8_4step_cfg.sh +++ b/scripts/wan/distill/run_wan_i2v_distill_int8_4step_cfg.sh @@ -13,5 +13,7 @@ python -m lightx2v.infer \ --config_json ${lightx2v_path}/configs/distill/wan21/wan_i2v_distill_int8_4step_cfg.json \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ + --num_frames 81 \ + --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_i2v_distill_int8_4step_cfg.mp4 diff --git a/scripts/wan/distill/run_wan_i2v_distill_lora_4step_cfg.sh b/scripts/wan/distill/run_wan_i2v_distill_lora_4step_cfg.sh index 771f989fd..b07dff386 100755 --- a/scripts/wan/distill/run_wan_i2v_distill_lora_4step_cfg.sh +++ b/scripts/wan/distill/run_wan_i2v_distill_lora_4step_cfg.sh @@ -13,5 +13,7 @@ python -m lightx2v.infer \ --config_json ${lightx2v_path}/configs/distill/wan21/wan_i2v_distill_lora_4step_cfg.json \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ + --num_frames 81 \ + --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_i2v_distill_lora_4step_cfg.mp4 diff --git a/scripts/wan/distill/run_wan_i2v_distill_lora_4step_cfg_ulysses.sh b/scripts/wan/distill/run_wan_i2v_distill_lora_4step_cfg_ulysses.sh index bcc2392d2..ccd313b31 100755 --- a/scripts/wan/distill/run_wan_i2v_distill_lora_4step_cfg_ulysses.sh +++ b/scripts/wan/distill/run_wan_i2v_distill_lora_4step_cfg_ulysses.sh @@ -13,5 +13,7 @@ torchrun --nproc_per_node=8 -m lightx2v.infer \ --config_json ${lightx2v_path}/configs/distill/wan21/wan_i2v_distill_lora_4step_cfg_ulysses.json \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ + --num_frames 81 \ + --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_i2v_distill_lora_4step_cfg_ulysses.mp4 diff --git a/scripts/wan/distill/run_wan_i2v_distill_model_4step_cfg.sh b/scripts/wan/distill/run_wan_i2v_distill_model_4step_cfg.sh index 1d93330c1..76eaff20c 100755 --- a/scripts/wan/distill/run_wan_i2v_distill_model_4step_cfg.sh +++ b/scripts/wan/distill/run_wan_i2v_distill_model_4step_cfg.sh @@ -13,5 +13,7 @@ python -m lightx2v.infer \ --config_json ${lightx2v_path}/configs/distill/wan21/wan_i2v_distill_model_4step_cfg.json \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ + --num_frames 81 \ + --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_i2v_distill_model_4step_cfg.mp4 diff --git a/scripts/wan/distill/run_wan_i2v_distill_model_4step_cfg_ulysses.sh b/scripts/wan/distill/run_wan_i2v_distill_model_4step_cfg_ulysses.sh index ff75fb55f..fc6be55e4 100755 --- a/scripts/wan/distill/run_wan_i2v_distill_model_4step_cfg_ulysses.sh +++ b/scripts/wan/distill/run_wan_i2v_distill_model_4step_cfg_ulysses.sh @@ -13,5 +13,7 @@ torchrun --nproc_per_node=8 -m lightx2v.infer \ --config_json ${lightx2v_path}/configs/distill/wan21/wan_i2v_distill_model_4step_cfg_ulysses.json \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ + --num_frames 81 \ + --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_i2v_distill_model_4step_cfg_ulysses.mp4 diff --git a/scripts/wan/distill/run_wan_t2v_distill_4step_cfg_dynamic.sh b/scripts/wan/distill/run_wan_t2v_distill_4step_cfg_dynamic.sh index 5635704ea..4f3ab8af4 100755 --- a/scripts/wan/distill/run_wan_t2v_distill_4step_cfg_dynamic.sh +++ b/scripts/wan/distill/run_wan_t2v_distill_4step_cfg_dynamic.sh @@ -12,6 +12,7 @@ python -m lightx2v.infer \ --model_path $model_path \ --config_json ${lightx2v_path}/configs/distill/wan_t2v_distill_4step_cfg_dynamic.json \ --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ + --num_frames 81 \ --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_t2v_distill_4step_cfg_dynamic.mp4 diff --git a/scripts/wan/distill/run_wan_t2v_distill_fp8_4step_cfg.sh b/scripts/wan/distill/run_wan_t2v_distill_fp8_4step_cfg.sh index 128f9f437..e2822f66d 100755 --- a/scripts/wan/distill/run_wan_t2v_distill_fp8_4step_cfg.sh +++ b/scripts/wan/distill/run_wan_t2v_distill_fp8_4step_cfg.sh @@ -12,6 +12,7 @@ python -m lightx2v.infer \ --model_path $model_path \ --config_json ${lightx2v_path}/configs/distill/wan21/wan_t2v_distill_fp8_4step_cfg.json \ --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ + --num_frames 81 \ --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_t2v_distill_fp8_4step_cfg.mp4 diff --git a/scripts/wan/distill/run_wan_t2v_distill_fp8_4step_cfg_ulysses.sh b/scripts/wan/distill/run_wan_t2v_distill_fp8_4step_cfg_ulysses.sh index 0dbd41263..272300437 100755 --- a/scripts/wan/distill/run_wan_t2v_distill_fp8_4step_cfg_ulysses.sh +++ b/scripts/wan/distill/run_wan_t2v_distill_fp8_4step_cfg_ulysses.sh @@ -12,6 +12,7 @@ torchrun --nproc_per_node=8 -m lightx2v.infer \ --model_path $model_path \ --config_json ${lightx2v_path}/configs/distill/wan21/wan_t2v_distill_fp8_4step_cfg_ulysses.json \ --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ + --num_frames 81 \ --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_t2v_distill_fp8_4step_cfg_ulysses.mp4 diff --git a/scripts/wan/distill/run_wan_t2v_distill_lora_4step_cfg.sh b/scripts/wan/distill/run_wan_t2v_distill_lora_4step_cfg.sh index 9f4d383da..38769b67d 100755 --- a/scripts/wan/distill/run_wan_t2v_distill_lora_4step_cfg.sh +++ b/scripts/wan/distill/run_wan_t2v_distill_lora_4step_cfg.sh @@ -12,6 +12,7 @@ python -m lightx2v.infer \ --model_path $model_path \ --config_json ${lightx2v_path}/configs/distill/wan21/wan_t2v_distill_lora_4step_cfg.json \ --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ + --num_frames 81 \ --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_t2v_distill_lora_4step_cfg.mp4 diff --git a/scripts/wan/distill/run_wan_t2v_distill_lora_4step_cfg_ulysses.sh b/scripts/wan/distill/run_wan_t2v_distill_lora_4step_cfg_ulysses.sh index 93de1168b..d2ae9ee24 100755 --- a/scripts/wan/distill/run_wan_t2v_distill_lora_4step_cfg_ulysses.sh +++ b/scripts/wan/distill/run_wan_t2v_distill_lora_4step_cfg_ulysses.sh @@ -12,6 +12,7 @@ torchrun --nproc_per_node=8 -m lightx2v.infer \ --model_path $model_path \ --config_json ${lightx2v_path}/configs/distill/wan21/wan_t2v_distill_lora_4step_cfg_ulysses.json \ --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ + --num_frames 81 \ --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_t2v_distill_lora_4step_cfg_ulysses.mp4 diff --git a/scripts/wan/distill/run_wan_t2v_distill_model_4step_cfg.sh b/scripts/wan/distill/run_wan_t2v_distill_model_4step_cfg.sh index 122e99535..997bbd3d6 100755 --- a/scripts/wan/distill/run_wan_t2v_distill_model_4step_cfg.sh +++ b/scripts/wan/distill/run_wan_t2v_distill_model_4step_cfg.sh @@ -12,6 +12,7 @@ python -m lightx2v.infer \ --model_path $model_path \ --config_json ${lightx2v_path}/configs/distill/wan21/wan_t2v_distill_model_4step_cfg.json \ --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ + --num_frames 81 \ --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_t2v_distill_model_4step_cfg.mp4 diff --git a/scripts/wan/distill/run_wan_t2v_distill_model_4step_cfg_ulysses.sh b/scripts/wan/distill/run_wan_t2v_distill_model_4step_cfg_ulysses.sh index 7870294e9..d77c4d171 100755 --- a/scripts/wan/distill/run_wan_t2v_distill_model_4step_cfg_ulysses.sh +++ b/scripts/wan/distill/run_wan_t2v_distill_model_4step_cfg_ulysses.sh @@ -12,6 +12,7 @@ torchrun --nproc_per_node=8 -m lightx2v.infer \ --model_path $model_path \ --config_json ${lightx2v_path}/configs/distill/wan21/wan_t2v_distill_model_4step_cfg_ulysses.json \ --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ + --num_frames 81 \ --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_t2v_distill_model_4step_cfg_ulysses.mp4 diff --git a/scripts/wan/run_wan_flf2v.sh b/scripts/wan/run_wan_flf2v.sh index 3273294ac..48d7d200f 100755 --- a/scripts/wan/run_wan_flf2v.sh +++ b/scripts/wan/run_wan_flf2v.sh @@ -15,5 +15,7 @@ python -m lightx2v.infer \ --negative_prompt "镜头晃动,色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \ --image_path ${lightx2v_path}/assets/inputs/imgs/flf2v_input_first_frame-fs8.png \ --last_frame_path ${lightx2v_path}/assets/inputs/imgs/flf2v_input_last_frame-fs8.png \ + --num_frames 81 \ + --target_shape 720 1280 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_flf2v.mp4 diff --git a/scripts/wan/run_wan_i2v.sh b/scripts/wan/run_wan_i2v.sh index 576984334..8f8e46b75 100755 --- a/scripts/wan/run_wan_i2v.sh +++ b/scripts/wan/run_wan_i2v.sh @@ -14,5 +14,7 @@ python -m lightx2v.infer \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ --negative_prompt "镜头晃动,色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ + --num_frames 81 \ + --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_i2v.mp4 diff --git a/scripts/wan/run_wan_i2v_lazy_load.sh b/scripts/wan/run_wan_i2v_lazy_load.sh index cbce64e85..9ffc4a94b 100755 --- a/scripts/wan/run_wan_i2v_lazy_load.sh +++ b/scripts/wan/run_wan_i2v_lazy_load.sh @@ -15,5 +15,7 @@ python -m lightx2v.infer \ --config_json ${lightx2v_path}/configs/offload/disk/wan_i2v_phase_lazy_load_720p.json \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ + --num_frames 81 \ + --target_shape 1280 720 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_i2v_lazy_load.mp4 diff --git a/scripts/wan/run_wan_i2v_sol_attn.sh b/scripts/wan/run_wan_i2v_sol_attn.sh index bb6098c5d..8c2d4897c 100755 --- a/scripts/wan/run_wan_i2v_sol_attn.sh +++ b/scripts/wan/run_wan_i2v_sol_attn.sh @@ -1,21 +1,20 @@ #!/bin/bash -lightx2v_path=/mnt/miaohua/wangshankun/LightX2V -model_path=/tmp/data/Wan2.1-I2V-14B-720P -config_json=${CONFIG_JSON:-${lightx2v_path}/configs/attentions/wan_i2v_sol_attn.json} -save_result_path=${SAVE_RESULT_PATH:-${lightx2v_path}/save_results/output_lightx2v_wan_i2v_sol_attn.mp4} + +lightx2v_path=/path/to/LightX2V +model_path=/path/to/Wan2.1-I2V-14B-480P export CUDA_VISIBLE_DEVICES=0 -export SOL_ATTN_STRICT=1 -# set environment variables source ${lightx2v_path}/scripts/base/base.sh python -m lightx2v.infer \ ---model_cls wan2.1 \ ---task i2v \ ---model_path $model_path \ - --model_path "${model_path}" \ - --config_json "${config_json}" \ - --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ - --negative_prompt "镜头晃动,色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \ - --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ - --save_result_path "${save_result_path}" + --model_cls wan2.1 \ + --task i2v \ + --model_path $model_path \ + --config_json ${lightx2v_path}/configs/attentions/wan_i2v_sol_attn.json \ + --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ + --negative_prompt "镜头晃动,色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \ + --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ + --num_frames 81 \ + --target_shape 480 832 \ + --seed 42 \ + --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_i2v_sol_attn.mp4 diff --git a/scripts/wan/run_wan_t2v.sh b/scripts/wan/run_wan_t2v.sh index 617b46c90..57bae04fc 100755 --- a/scripts/wan/run_wan_t2v.sh +++ b/scripts/wan/run_wan_t2v.sh @@ -13,6 +13,7 @@ python -m lightx2v.infer \ --config_json ${lightx2v_path}/configs/wan/wan_t2v.json \ --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ --negative_prompt "镜头晃动,色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \ + --num_frames 81 \ --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_t2v.mp4 diff --git a/scripts/wan/run_wan_vace.sh b/scripts/wan/run_wan_vace.sh index d0c563999..af267c49e 100755 --- a/scripts/wan/run_wan_vace.sh +++ b/scripts/wan/run_wan_vace.sh @@ -14,5 +14,7 @@ python -m lightx2v.infer \ --prompt "在一个欢乐而充满节日气氛的场景中,穿着鲜艳红色春服的小女孩正与她的可爱卡通蛇嬉戏。她的春服上绣着金色吉祥图案,散发着喜庆的气息,脸上洋溢着灿烂的笑容。蛇身呈现出亮眼的绿色,形状圆润,宽大的眼睛让它显得既友善又幽默。小女孩欢快地用手轻轻抚摸着蛇的头部,共同享受着这温馨的时刻。周围五彩斑斓的灯笼和彩带装饰着环境,阳光透过洒在她们身上,营造出一个充满友爱与幸福的新年氛围。" \ --negative_prompt "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \ --src_ref_images ${lightx2v_path}/assets/inputs/imgs/girl.png,${lightx2v_path}/assets/inputs/imgs/snake.png \ + --num_frames 81 \ + --target_shape 480 832 \ --seed 42 \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan_vace.mp4 diff --git a/scripts/wan/server/post_flf2v.py b/scripts/wan/server/post_flf2v.py new file mode 100644 index 000000000..170a39b15 --- /dev/null +++ b/scripts/wan/server/post_flf2v.py @@ -0,0 +1,27 @@ +import base64 +from pathlib import Path + +import requests +from loguru import logger + +if __name__ == "__main__": + url = "http://localhost:8000/v1/tasks/video/" + lightx2v_path = Path(__file__).resolve().parents[3] + image_path = lightx2v_path / "assets/inputs/imgs/flf2v_input_first_frame-fs8.png" + last_frame_path = lightx2v_path / "assets/inputs/imgs/flf2v_input_last_frame-fs8.png" + + message = { + "prompt": "CG animation style, a small blue bird takes off from the ground, flapping its wings. The bird's feathers are delicate, with a unique pattern on its chest. The background shows a blue sky with white clouds under bright sunshine. The camera follows the bird upward, capturing its flight and the vastness of the sky from a close-up, low-angle perspective.", + "negative_prompt": "镜头晃动,色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走", + "image_path": base64.b64encode(image_path.read_bytes()).decode("utf-8"), + "last_frame_path": base64.b64encode(last_frame_path.read_bytes()).decode("utf-8"), + "seed": 42, + "num_frames": 81, + "target_shape": [720, 1280], + "save_result_path": "./output_lightx2v_wan_flf2v.mp4", + } + + logger.info(f"image_path: {image_path}") + logger.info(f"last_frame_path: {last_frame_path}") + response = requests.post(url, json=message) + logger.info(f"response: {response.json()}") diff --git a/scripts/wan/server/post_i2v.py b/scripts/wan/server/post_i2v.py new file mode 100644 index 000000000..d584a4444 --- /dev/null +++ b/scripts/wan/server/post_i2v.py @@ -0,0 +1,26 @@ +import base64 +from pathlib import Path + +import requests +from loguru import logger + +if __name__ == "__main__": + url = "http://localhost:8000/v1/tasks/video/" + lightx2v_path = Path(__file__).resolve().parents[3] + image_path = lightx2v_path / "assets/inputs/imgs/img_0.jpg" + + message = { + "prompt": "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression.", + "negative_prompt": "镜头晃动,色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走", + "image_path": base64.b64encode(image_path.read_bytes()).decode("utf-8"), + "seed": 42, + "num_frames": 81, + "target_shape": [480, 832], + # Set to True after one successful request when using start_server_i2v_reuse.sh. + "reuse": False, + "save_result_path": "./output_lightx2v_wan_i2v.mp4", + } + + logger.info(f"image_path: {image_path}") + response = requests.post(url, json=message) + logger.info(f"response: {response.json()}") diff --git a/scripts/wan/server/post_t2v.py b/scripts/wan/server/post_t2v.py new file mode 100644 index 000000000..f56e921d1 --- /dev/null +++ b/scripts/wan/server/post_t2v.py @@ -0,0 +1,18 @@ +import requests +from loguru import logger + +if __name__ == "__main__": + url = "http://localhost:8000/v1/tasks/video/" + + message = { + "prompt": "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.", + "negative_prompt": "镜头晃动,色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走", + "seed": 42, + "num_frames": 81, + "target_shape": [480, 832], + "save_result_path": "./output_lightx2v_wan_t2v.mp4", + } + + logger.info(f"message: {message}") + response = requests.post(url, json=message) + logger.info(f"response: {response.json()}") diff --git a/scripts/wan/server/post_vace.py b/scripts/wan/server/post_vace.py new file mode 100644 index 000000000..430a2d130 --- /dev/null +++ b/scripts/wan/server/post_vace.py @@ -0,0 +1,27 @@ +import base64 +from pathlib import Path + +import requests +from loguru import logger + +if __name__ == "__main__": + url = "http://localhost:8000/v1/tasks/video/" + lightx2v_path = Path(__file__).resolve().parents[3] + reference_image_paths = [ + lightx2v_path / "assets/inputs/imgs/girl.png", + lightx2v_path / "assets/inputs/imgs/snake.png", + ] + + message = { + "prompt": "在一个欢乐而充满节日气氛的场景中,穿着鲜艳红色春服的小女孩正与她的可爱卡通蛇嬉戏。她的春服上绣着金色吉祥图案,散发着喜庆的气息,脸上洋溢着灿烂的笑容。蛇身呈现出亮眼的绿色,形状圆润,宽大的眼睛让它显得既友善又幽默。小女孩欢快地用手轻轻抚摸着蛇的头部,共同享受着这温馨的时刻。周围五彩斑斓的灯笼和彩带装饰着环境,阳光透过洒在她们身上,营造出一个充满友爱与幸福的新年氛围。", + "negative_prompt": "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走", + "src_ref_images": [base64.b64encode(path.read_bytes()).decode("utf-8") for path in reference_image_paths], + "seed": 42, + "num_frames": 81, + "target_shape": [480, 832], + "save_result_path": "./output_lightx2v_wan_vace.mp4", + } + + logger.info(f"reference_image_paths: {reference_image_paths}") + response = requests.post(url, json=message) + logger.info(f"response: {response.json()}") diff --git a/scripts/wan/server/start_server_flf2v.sh b/scripts/wan/server/start_server_flf2v.sh new file mode 100755 index 000000000..4a12cf404 --- /dev/null +++ b/scripts/wan/server/start_server_flf2v.sh @@ -0,0 +1,15 @@ +#!/bin/bash + +lightx2v_path=/path/to/LightX2V +model_path=/path/to/Wan2.1-FLF2V-14B-720P + +export CUDA_VISIBLE_DEVICES=0 +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.server \ + --model_cls wan2.1 \ + --task flf2v \ + --model_path $model_path \ + --config_json ${lightx2v_path}/configs/wan/wan_flf2v.json \ + --host 0.0.0.0 \ + --port 8000 diff --git a/scripts/wan/server/start_server_i2v.sh b/scripts/wan/server/start_server_i2v.sh new file mode 100755 index 000000000..a80a05751 --- /dev/null +++ b/scripts/wan/server/start_server_i2v.sh @@ -0,0 +1,15 @@ +#!/bin/bash + +lightx2v_path=/path/to/LightX2V +model_path=/path/to/Wan2.1-I2V-14B-480P + +export CUDA_VISIBLE_DEVICES=0 +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.server \ + --model_cls wan2.1 \ + --task i2v \ + --model_path $model_path \ + --config_json ${lightx2v_path}/configs/wan/wan_i2v.json \ + --host 0.0.0.0 \ + --port 8000 diff --git a/scripts/wan/server/start_server_i2v_compile.sh b/scripts/wan/server/start_server_i2v_compile.sh new file mode 100755 index 000000000..71f977df7 --- /dev/null +++ b/scripts/wan/server/start_server_i2v_compile.sh @@ -0,0 +1,15 @@ +#!/bin/bash + +lightx2v_path=/path/to/LightX2V +model_path=/path/to/Wan2.1-I2V-14B-480P + +export CUDA_VISIBLE_DEVICES=0 +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.server \ + --model_cls wan2.1 \ + --task i2v \ + --model_path $model_path \ + --config_json ${lightx2v_path}/configs/wan/wan_i2v_compile.json \ + --host 0.0.0.0 \ + --port 8000 diff --git a/scripts/wan/server/start_server_i2v_reuse.sh b/scripts/wan/server/start_server_i2v_reuse.sh new file mode 100755 index 000000000..b772d7bed --- /dev/null +++ b/scripts/wan/server/start_server_i2v_reuse.sh @@ -0,0 +1,15 @@ +#!/bin/bash + +lightx2v_path=/path/to/LightX2V +model_path=/path/to/Wan2.1-I2V-14B-480P + +export CUDA_VISIBLE_DEVICES=0 +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.server \ + --model_cls wan2.1 \ + --task i2v \ + --model_path $model_path \ + --config_json ${lightx2v_path}/configs/wan/wan_i2v_reuse.json \ + --host 0.0.0.0 \ + --port 8000 diff --git a/scripts/wan/server/start_server_t2v.sh b/scripts/wan/server/start_server_t2v.sh new file mode 100755 index 000000000..cc77b30cf --- /dev/null +++ b/scripts/wan/server/start_server_t2v.sh @@ -0,0 +1,15 @@ +#!/bin/bash + +lightx2v_path=/path/to/LightX2V +model_path=/path/to/Wan2.1-T2V-1.3B + +export CUDA_VISIBLE_DEVICES=0 +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.server \ + --model_cls wan2.1 \ + --task t2v \ + --model_path $model_path \ + --config_json ${lightx2v_path}/configs/wan/wan_t2v.json \ + --host 0.0.0.0 \ + --port 8000 diff --git a/scripts/wan/server/start_server_vace.sh b/scripts/wan/server/start_server_vace.sh new file mode 100755 index 000000000..b93a8ac8f --- /dev/null +++ b/scripts/wan/server/start_server_vace.sh @@ -0,0 +1,15 @@ +#!/bin/bash + +lightx2v_path=/path/to/LightX2V +model_path=/path/to/Wan2.1-VACE-14B + +export CUDA_VISIBLE_DEVICES=0 +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.server \ + --model_cls wan2.1_vace \ + --task vace \ + --model_path $model_path \ + --config_json ${lightx2v_path}/configs/wan/wan_vace.json \ + --host 0.0.0.0 \ + --port 8000 diff --git a/scripts/wan22/extreme/run_wan22_moe_i2v_extreme.sh b/scripts/wan22/extreme/run_wan22_moe_i2v_extreme.sh index 038d79362..7f1f55aa9 100755 --- a/scripts/wan22/extreme/run_wan22_moe_i2v_extreme.sh +++ b/scripts/wan22/extreme/run_wan22_moe_i2v_extreme.sh @@ -15,6 +15,5 @@ python -m lightx2v.infer \ --model_path $model_path \ --config_json ${lightx2v_path}/configs/wan22/extreme/wan_moe_i2v_distill_nvfp4_sparse_attn.json \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ ---negative_prompt "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan22_moe_i2v_extreme.mp4 diff --git a/scripts/wan22/extreme/run_wan22_moe_i2v_extreme_sp_parallel.sh b/scripts/wan22/extreme/run_wan22_moe_i2v_extreme_sp_parallel.sh index f4d567194..89cf8eda0 100755 --- a/scripts/wan22/extreme/run_wan22_moe_i2v_extreme_sp_parallel.sh +++ b/scripts/wan22/extreme/run_wan22_moe_i2v_extreme_sp_parallel.sh @@ -15,6 +15,5 @@ torchrun --nproc_per_node=8 -m lightx2v.infer \ --model_path $model_path \ --config_json ${lightx2v_path}/configs/wan22/extreme/wan_moe_i2v_distill_nvfp4_sparse_attn_sp_parallel.json \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ ---negative_prompt "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan22_moe_i2v_extreme_sp_parallel.mp4 diff --git a/scripts/wan22/extreme/run_wan22_moe_t2v_extreme.sh b/scripts/wan22/extreme/run_wan22_moe_t2v_extreme.sh index 22a7e3495..6a837c4f0 100755 --- a/scripts/wan22/extreme/run_wan22_moe_t2v_extreme.sh +++ b/scripts/wan22/extreme/run_wan22_moe_t2v_extreme.sh @@ -15,5 +15,4 @@ python -m lightx2v.infer \ --model_path $model_path \ --config_json ${lightx2v_path}/configs/wan22/extreme/wan_moe_t2v_distill_nvfp4_sparse_attn.json \ --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ ---negative_prompt "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan22_moe_t2v_extreme.mp4 diff --git a/scripts/wan22/extreme/run_wan22_moe_t2v_extreme_sp_parallel.sh b/scripts/wan22/extreme/run_wan22_moe_t2v_extreme_sp_parallel.sh index 71f410b6e..e055054b6 100755 --- a/scripts/wan22/extreme/run_wan22_moe_t2v_extreme_sp_parallel.sh +++ b/scripts/wan22/extreme/run_wan22_moe_t2v_extreme_sp_parallel.sh @@ -15,5 +15,4 @@ torchrun --nproc_per_node=8 -m lightx2v.infer \ --model_path $model_path \ --config_json ${lightx2v_path}/configs/wan22/extreme/wan_moe_t2v_distill_nvfp4_sparse_attn_sp_parallel.json \ --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ ---negative_prompt "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan22_moe_t2v_extreme_sp_parallel.mp4 diff --git a/scripts/wan22/run_wan22_animate_replace.sh b/scripts/wan22/run_wan22_animate_replace.sh index c5345dd85..48ec149e0 100755 --- a/scripts/wan22/run_wan22_animate_replace.sh +++ b/scripts/wan22/run_wan22_animate_replace.sh @@ -33,7 +33,6 @@ python -m lightx2v.infer \ --src_face_path ${lightx2v_path}/save_results/animate/process_results/src_face.mp4 \ --src_ref_images ${lightx2v_path}/save_results/animate/process_results/src_ref.png \ --src_bg_path ${lightx2v_path}/save_results/animate/process_results/src_bg.mp4 \ ---src_mask_path ${lightx2v_path}/save_results/animate/process_results/src_mask.mp4 \ +--mask_path ${lightx2v_path}/save_results/animate/process_results/src_mask.mp4 \ --prompt "视频中的人在做动作" \ ---negative_prompt "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan22_replace.mp4 diff --git a/scripts/wan22/run_wan22_animate_replace_lora.sh b/scripts/wan22/run_wan22_animate_replace_lora.sh index 70a1d8b5a..e0cc643bf 100755 --- a/scripts/wan22/run_wan22_animate_replace_lora.sh +++ b/scripts/wan22/run_wan22_animate_replace_lora.sh @@ -34,7 +34,6 @@ python -m lightx2v.infer \ --src_face_path ${lightx2v_path}/save_results/animate/process_results/src_face.mp4 \ --src_ref_images ${lightx2v_path}/save_results/animate/process_results/src_ref.png \ --src_bg_path ${lightx2v_path}/save_results/animate/process_results/src_bg.mp4 \ ---src_mask_path ${lightx2v_path}/save_results/animate/process_results/src_mask.mp4 \ +--mask_path ${lightx2v_path}/save_results/animate/process_results/src_mask.mp4 \ --prompt "视频中的人在做动作" \ ---negative_prompt "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan22_replace.mp4 diff --git a/scripts/wan22_moe_vace/run_wan22_moe_vace.sh b/scripts/wan22_moe_vace/run_wan22_moe_vace.sh index 09d0a5f23..de60e8579 100755 --- a/scripts/wan22_moe_vace/run_wan22_moe_vace.sh +++ b/scripts/wan22_moe_vace/run_wan22_moe_vace.sh @@ -17,6 +17,6 @@ python -m lightx2v.infer \ --config_json ${lightx2v_path}/configs/wan22_vace/a800/bf16/wan22_moe_vace.json \ --prompt "图片的女人,穿着白色连衣裙,模仿视频的动作,翩翩起舞." \ --negative_prompt "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \ ---src_video /path/to/post+depth.mp4 \ +--video_path /path/to/post+depth.mp4 \ --src_ref_images /path/to/image.png \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan22_moe_vace.mp4\ diff --git a/scripts/wan22_moe_vace/run_wan22_moe_vace_dist.sh b/scripts/wan22_moe_vace/run_wan22_moe_vace_dist.sh index 458326c42..2ac4cda32 100755 --- a/scripts/wan22_moe_vace/run_wan22_moe_vace_dist.sh +++ b/scripts/wan22_moe_vace/run_wan22_moe_vace_dist.sh @@ -17,6 +17,6 @@ torchrun --nproc_per_node=2 -m lightx2v.infer \ --config_json ${lightx2v_path}/configs/wan22_vace/a800/bf16/wan22_moe_vace_cfg_parallel.json \ --prompt "图片的女人,穿着白色连衣裙,模仿视频的动作,翩翩起舞." \ --negative_prompt "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \ ---src_video /path/to/post+depth.mp4 \ +--video_path /path/to/post+depth.mp4 \ --src_ref_images /path/to/image.png \ --save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan22_moe_vace_cfg_parallel.mp4\ diff --git a/scripts/worldmirror/run_calibration.py b/scripts/worldmirror/run_calibration.py index 1f57970de..30c2c843d 100644 --- a/scripts/worldmirror/run_calibration.py +++ b/scripts/worldmirror/run_calibration.py @@ -95,7 +95,6 @@ def main(): # Build runner and force-init modules (loads model + installs adapters). from lightx2v.models.runners.worldmirror.worldmirror_runner import WorldMirrorRunner # noqa: E402 - from lightx2v.utils.input_info import init_empty_input_info # noqa: E402 logger.info("[calib] Building WorldMirrorRunner...") runner = WorldMirrorRunner(config) @@ -116,14 +115,16 @@ def main(): continue logger.info(f"[calib] ({i + 1}/{len(args.scenes)}) scene={scene}") - input_info = init_empty_input_info("recon") - input_info.input_path = scene_path - # Send output into a throwaway tmp dir — save_* are all off above - # but the runner still wants a writable path. - input_info.save_result_path = "/tmp/wm_calib_output" - t0 = time.perf_counter() - runner.run_pipeline(input_info) + input_info = runner.prepare_request( + { + "input_path": scene_path, + # Send output into a throwaway tmp dir — save_* are all off above + # but the runner still wants a writable path. + "save_result_path": "/tmp/wm_calib_output", + } + ) + runner.run_request(input_info) if torch.cuda.is_available(): torch.cuda.synchronize() logger.info(f"[calib] done in {time.perf_counter() - t0:.1f}s") diff --git a/scripts/worldplay/run_worldplay_ar_sp4.sh b/scripts/worldplay/run_worldplay_ar_sp4.sh index 334c94025..1122c8618 100755 --- a/scripts/worldplay/run_worldplay_ar_sp4.sh +++ b/scripts/worldplay/run_worldplay_ar_sp4.sh @@ -11,7 +11,7 @@ MODEL_PATH=/data/nvme1/models/hunyuan/HunyuanVideo-1.5 # Input parameters PROMPT='A paved pathway leads towards a stone arch bridge spanning a calm body of water. Lush green trees and foliage line the path and the far bank of the water.' IMAGE_PATH=/workspace/HY-WorldPlay/assets/img/test.png -POSE='w-15,s-15' # Forward 15 frames, then backward 15 frames. Auto target_video_length=121 +POSE='w-15,s-15' # Forward 15 frames, then backward 15 frames SEED=42 # Output @@ -27,6 +27,7 @@ torchrun --nproc_per_node=4 -m lightx2v.infer \ --prompt "$PROMPT" \ --image_path $IMAGE_PATH \ --pose "$POSE" \ + --num_frames 121 \ --seed $SEED \ --save_result_path $OUTPUT_PATH diff --git a/tests/test_cfg_contract.py b/tests/test_cfg_contract.py new file mode 100644 index 000000000..a5c6e3f7a --- /dev/null +++ b/tests/test_cfg_contract.py @@ -0,0 +1,242 @@ +import json +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler + +from lightx2v.models.networks.cosmos3.model import Cosmos3TransformerModel +from lightx2v.models.networks.ernie_image.model import ErnieImageTransformerModel +from lightx2v.models.networks.flux2.model import Flux2DevTransformerModel, Flux2KleinTransformerModel +from lightx2v.models.networks.hunyuan_image3.model import HunyuanImage3Model +from lightx2v.models.networks.hunyuan_video.model import HunyuanVideo15Model +from lightx2v.models.networks.longcat_image.model import LongCatImageTransformerModel +from lightx2v.models.networks.ltx2.model import LTX2Model +from lightx2v.models.networks.qwen_image.model import QwenImageTransformerModel +from lightx2v.models.networks.wan.dreamzero_model import DreamZeroModel +from lightx2v.models.networks.wan.model import WanModel +from lightx2v.models.networks.z_image.model import ZImageTransformerModel +from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner import HunyuanVideo15Runner +from lightx2v.models.runners.wan.wan_runner import MultiModelStruct +from lightx2v.models.schedulers.cosmos3.scheduler import Cosmos3Scheduler +from lightx2v.models.schedulers.flux2.scheduler import Flux2Scheduler +from lightx2v.models.schedulers.hunyuan_video.scheduler import HunyuanVideo15Scheduler +from lightx2v.models.schedulers.lingbot_video.scheduler import LingBotVideoScheduler +from lightx2v.models.schedulers.wan.animate2.scheduler import WanAnimate2Scheduler +from lightx2v.models.schedulers.wan.scheduler import WanScheduler + +STANDARD_MODELS = [ + WanModel, + Cosmos3TransformerModel, + Flux2KleinTransformerModel, + HunyuanVideo15Model, + ErnieImageTransformerModel, + LongCatImageTransformerModel, + QwenImageTransformerModel, + ZImageTransformerModel, + LTX2Model, +] + + +@pytest.fixture +def config(tmp_path): + FlowMatchEulerDiscreteScheduler().save_pretrained(tmp_path / "scheduler") + return { + "model_path": str(tmp_path), + "model_cls": "wan2.2_audio", + "infer_steps": 4, + "enable_cfg": True, + "sample_guide_scale": 4.0, + "sample_shift": 3.0, + "seq_parallel": False, + "parallel": None, + "dim": 16, + "num_heads": 2, + } + + +@pytest.fixture +def inputs(): + return { + "text_encoder_output": { + "cond_input_ids": "positive", + "uncond_input_ids": "negative", + "prompt_embeds": "positive", + "negative_prompt_embeds": "negative", + } + } + + +def make_model(model_cls, enabled=True, scale=4.0): + model = object.__new__(model_cls) + model.config = {"enable_cfg": enabled, "cfg_parallel": False, "task": "t2i"} + model.cpu_offload = False + model.scheduler = SimpleNamespace(latents=torch.zeros(1), sample_guide_scale=scale) + values = [torch.tensor([10.0]), torch.tensor([2.0])] + if model_cls is Cosmos3TransformerModel: + values = [SimpleNamespace(vision=value, sound=None, action=None) for value in values] + elif model_cls is LTX2Model: + values = [(value, value * 2) for value in values] + model._infer_cond_uncond = Mock(side_effect=values) + return model + + +@pytest.mark.parametrize("model_cls", STANDARD_MODELS) +def test_enabled_cfg_rejects_unit_scale_before_forward(model_cls, inputs): + model = make_model(model_cls, scale=1.0) + with pytest.raises(AssertionError, match="enable_cfg=true requires sample_guide_scale != 1"): + model.infer(inputs) + model._infer_cond_uncond.assert_not_called() + + +@pytest.mark.parametrize("model_cls", STANDARD_MODELS) +def test_disabled_cfg_accepts_unit_scale_and_runs_one_branch(model_cls, inputs): + model = make_model(model_cls, enabled=False, scale=1.0) + model.infer(inputs) + model._infer_cond_uncond.assert_called_once() + + +@pytest.mark.parametrize("scheduler_cls", [WanScheduler, Cosmos3Scheduler, Flux2Scheduler, HunyuanVideo15Scheduler]) +def test_scheduler_initialization_does_not_validate_cfg(config, scheduler_cls): + scheduler = scheduler_cls({**config, "sample_guide_scale": 1.0}) + assert scheduler.sample_guide_scale == 1.0 + + +@pytest.mark.parametrize("model_cls", [WanModel, Cosmos3TransformerModel, Flux2KleinTransformerModel, HunyuanVideo15Model, LTX2Model]) +@pytest.mark.parametrize(("enabled", "scale"), [(False, 4.0), (True, 0.5), (True, 4.0)]) +def test_transformer_cfg_uses_only_the_flag(model_cls, enabled, scale, inputs): + model = make_model(model_cls, enabled=enabled, scale=scale) + model.infer(inputs) + assert model._infer_cond_uncond.call_count == (2 if enabled else 1) + expected = torch.tensor([2.0 + scale * 8 if enabled else 10.0]) + if model_cls is LTX2Model: + torch.testing.assert_close(model.scheduler.v_noise_pred, expected) + torch.testing.assert_close(model.scheduler.a_noise_pred, expected * 2) + else: + torch.testing.assert_close(model.scheduler.noise_pred, expected) + + +@pytest.mark.parametrize("scheduler_cls", [WanScheduler, WanAnimate2Scheduler, LingBotVideoScheduler]) +def test_disagg_refresh_is_checked_when_model_runs(config, scheduler_cls, inputs): + model = make_model(WanModel) + model.scheduler = scheduler_cls(config) + model.scheduler.refresh_from_config({**config, "sample_guide_scale": 1.0}) + with pytest.raises(AssertionError, match="sample_guide_scale"): + model.infer(inputs) + model._infer_cond_uncond.assert_not_called() + + +def test_wan_moe_checks_the_selected_stage_scale(config, inputs): + wrapper = object.__new__(MultiModelStruct) + wrapper.config = {**config, "sample_guide_scale": [4.0, 1.0], "cpu_offload": False} + wrapper.scheduler = WanScheduler(wrapper.config) + wrapper.model = [make_model(WanModel), make_model(WanModel)] + for model in wrapper.model: + model.scheduler = wrapper.scheduler + wrapper.uses_high_noise_model = Mock(side_effect=[True, False]) + wrapper.infer(inputs) + wrapper.model[0]._infer_cond_uncond.assert_called() + with pytest.raises(AssertionError, match="sample_guide_scale"): + wrapper.infer(inputs) + wrapper.model[1]._infer_cond_uncond.assert_not_called() + + +@pytest.mark.parametrize("model_cls", [Cosmos3TransformerModel, LTX2Model]) +@pytest.mark.parametrize("rank", [0, 1]) +def test_cfg_parallel_keeps_guided_outputs(monkeypatch, model_cls, rank, inputs): + model = make_model(model_cls) + model.config.update(cfg_parallel=True, device_mesh=Mock()) + cond, uncond = torch.tensor([10.0]), torch.tensor([2.0]) + local = cond if rank == 0 else uncond + model._infer_cond_uncond = Mock(return_value=(local, local * 2) if model_cls is LTX2Model else SimpleNamespace(vision=local, sound=None, action=None)) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group: rank) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group: 2) + calls = [] + + def gather(outputs, value, group): + factor = 2 if calls else 1 + calls.append(value) + outputs[0].copy_(cond * factor) + outputs[1].copy_(uncond * factor) + + monkeypatch.setattr(torch.distributed, "all_gather", gather) + model.infer(inputs) + model._infer_cond_uncond.assert_called_once() + if model_cls is LTX2Model: + torch.testing.assert_close(model.scheduler.v_noise_pred, torch.tensor([34.0])) + torch.testing.assert_close(model.scheduler.a_noise_pred, torch.tensor([68.0])) + else: + torch.testing.assert_close(model.scheduler.noise_pred, torch.tensor([34.0])) + + +@pytest.mark.parametrize("model_cls", [Cosmos3TransformerModel, LTX2Model, WanModel]) +def test_unit_scale_fails_before_cfg_parallel_collectives(model_cls, inputs): + model = make_model(model_cls, scale=1.0) + model.config["cfg_parallel"] = True + with pytest.raises(AssertionError, match="sample_guide_scale"): + model.infer(inputs) + model._infer_cond_uncond.assert_not_called() + + +def test_ltx_multimodal_guidance_keeps_its_own_scale_rules(inputs): + model = make_model(LTX2Model, scale=1.0) + model.scheduler.mm_guider_enabled = True + model._infer_mm_guider_cfg = Mock() + model.infer(inputs) + model._infer_mm_guider_cfg.assert_called_once_with(inputs) + model._infer_cond_uncond.assert_not_called() + + +def test_flux_dev_embedding_guidance_keeps_unit_scale(inputs): + model = make_model(Flux2DevTransformerModel, scale=1.0) + model.infer(inputs) + model._infer_cond_uncond.assert_called_once() + + +@pytest.mark.parametrize("field", ["sample_guide_scale", "diff_guidance_scale"]) +def test_hunyuan_image_checks_the_scale_when_combining_cfg(field): + model = object.__new__(HunyuanImage3Model) + model.config = {field: 4.0} + cond, uncond = torch.tensor([10.0]), torch.tensor([2.0]) + torch.testing.assert_close(model.combine_cfg_predictions(cond, uncond), torch.tensor([34.0])) + model.config[field] = 1.0 + with pytest.raises(AssertionError, match="guidance_scale"): + model.combine_cfg_predictions(cond, uncond) + + +def test_dreamzero_checks_runtime_guide_scale(): + model = make_model(DreamZeroModel) + model.config["sample_guide_scale"] = 4.0 + with pytest.raises(AssertionError, match="guide_scale"): + model.infer({"guide_scale": 1.0}) + + +def test_sr_uses_its_own_cfg_flag_and_conditioning(monkeypatch, config): + root = Path(__file__).resolve().parents[1] + config.update(json.loads((root / "configs/hunyuan_video_15/vsr/hy15_i2v_480p.json").read_text())) + config["transformer_model_path"] = "/models/hunyuan/480p_i2v" + monkeypatch.setattr(DefaultRunner, "__init__", lambda self, config: setattr(self, "config", config)) + runner = HunyuanVideo15Runner(config) + assert runner.config["enable_cfg"] is True + assert runner.config_sr["enable_cfg"] is False + runner.config_sr.update(is_sr_running=True, cfg_parallel=True) + qwen = Mock() + byt5 = Mock() + byt5.infer.return_value = ("features", "masks") + runner.text_encoders = [qwen, byt5] + output = runner.run_text_encoder(SimpleNamespace(prompt="positive", negative_prompt="negative")) + qwen.infer.assert_called_once_with(["positive"]) + assert output["context_null"] is None + + +def test_cosmos_action_presets_explicitly_disable_cfg(): + root = Path(__file__).resolve().parents[1] + paths = list((root / "configs/cosmos3").glob("cosmos3_*_omni_action_*.json")) + assert len(paths) == 6 + for path in paths: + config = json.loads(path.read_text()) + assert config["sample_guide_scale"] == 1.0 + assert config["enable_cfg"] is False diff --git a/tests/test_cli_parser.py b/tests/test_cli_parser.py new file mode 100644 index 000000000..e1871d4e2 --- /dev/null +++ b/tests/test_cli_parser.py @@ -0,0 +1,238 @@ +import json +from pathlib import Path + +import pytest + +from lightx2v import infer +from lightx2v.disagg.examples import infer as disagg_infer +from lightx2v.models.runners.bagel.sensenova_vision_runner import SenseNovaVisionRunner +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.flux2.flux2_runner import Flux2Runner +from lightx2v.models.runners.hidream_o1_image.hidream_o1_image_runner import HidreamO1ImageRunner +from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner +from lightx2v.models.runners.wan.wan_audio_runner import Wan22AudioRunner +from lightx2v.models.runners.wan.wan_infinitetalk_runner import InfiniteTalkRunner +from lightx2v.models.runners.wan.wan_runner import WanRunner +from lightx2v.models.runners.worldmirror.worldmirror_runner import WorldMirrorRunner +from lightx2v.utils.set_config import build_cli_inputs + + +class CapturedInputs(Exception): + def __init__(self, config, request): + self.config = config + self.request = request + + +@pytest.fixture +def cli_inputs(tmp_path, monkeypatch): + def capture(args): + raise CapturedInputs(*build_cli_inputs(args)) + + def parse(config, options=(), model_cls="wan2.1", task="t2v", entrypoint=infer): + monkeypatch.setattr(entrypoint, "build_cli_inputs", capture) + model_path = tmp_path / "model" + model_path.mkdir(exist_ok=True) + if model_cls == "sensenova_vision": + for filename in ("llm_config.json", "vit_config.json"): + (model_path / filename).write_text("{}") + config_path = tmp_path / "deployment.json" + config_path.write_text(json.dumps(config)) + monkeypatch.setattr( + "sys.argv", + ["infer", "--model_cls", model_cls, "--task", task, "--model_path", str(model_path), "--config_json", str(config_path), *options], + ) + with pytest.raises(CapturedInputs) as captured: + entrypoint.main() + return captured.value.config, captured.value.request + + return parse + + +@pytest.mark.parametrize("entrypoint", [infer, disagg_infer]) +@pytest.mark.parametrize("options", [(), ("--return_result_tensor",)]) +def test_cli_return_result_tensor(cli_inputs, entrypoint, options): + config, request = cli_inputs({}, options, entrypoint=entrypoint) + + assert ("return_result_tensor" in request) == bool(options) + assert "return_result_tensor" not in config + runner = object.__new__(WanRunner) + BaseRunner.__init__(runner, config) + assert runner.prepare_request(request).return_result_tensor is bool(options) + + +def test_static_cli_omissions_resolve_request_content_and_output_specs(cli_inputs): + defaults = {"target_video_length": 81} + config, request = cli_inputs(defaults) + + assert request == {"task": "t2v"} + for field, value in defaults.items(): + assert config[field] == value + runner = object.__new__(WanRunner) + BaseRunner.__init__(runner, config) + input_info = runner.prepare_request({"task": runner.config["task"], **request}) + assert input_info.save_result_path is None + assert input_info.seed == 42 + assert input_info.prompt == input_info.negative_prompt == "" + assert input_info.target_video_length == 81 + + +def test_static_cli_preserves_zero_empty_string_and_frame_alias(cli_inputs): + defaults = {"target_video_length": 81} + config, request = cli_inputs(defaults, ["--seed", "0", "--negative_prompt", "", "--save_result_path", "", "--num_frames", "49", "--target_shape", "480", "832"]) + + assert request == {"task": "t2v", "seed": 0, "negative_prompt": "", "save_result_path": "", "target_video_length": 49, "target_shape": [480, 832]} + for field, value in defaults.items(): + assert config[field] == value + assert not {"seed", "prompt", "negative_prompt", "save_result_path"} & config.keys() + + +@pytest.mark.parametrize("task", ["i2v", "s2v"]) +def test_wan22_audio_cli_preserves_audio_request(cli_inputs, task): + preset = Path(__file__).resolve().parents[1] / "configs/seko_talk/seko_talk_08_5B_base.json" + config, request = cli_inputs( + json.loads(preset.read_text()), + ["--image_path", "portrait.png", "--audio_path", "speech.wav", "--seed", "0"], + model_cls="wan2.2_audio", + task=task, + ) + runner = object.__new__(Wan22AudioRunner) + BaseRunner.__init__(runner, config) + input_info = runner.prepare_request(request) + + assert input_info.task == task + assert input_info.image_path == "portrait.png" + assert input_info.audio_path == "speech.wav" + assert input_info.seed == 0 + assert input_info.save_result_path is None + assert config["vae_stride"] == [4, 16, 16] + assert config["use_image_encoder"] is False + assert runner.get_latent_shape_with_lat_hw(4, 4, 17) == [48, 5, 4, 4] + + +def test_sensenova_subtask_uses_json_until_explicitly_overridden(cli_inputs): + defaults = {"omni_vision_subtask": "depth"} + config, request = cli_inputs(defaults, model_cls="sensenova_vision", task="omni_vision_task") + assert config["omni_vision_subtask"] == "depth" + assert request == {"task": "omni_vision_task"} + + config, request = cli_inputs(defaults, ["--omni_vision_subtask", "normal"], model_cls="sensenova_vision", task="omni_vision_task") + assert config["omni_vision_subtask"] == "depth" + assert request == {"task": "omni_vision_task", "omni_vision_subtask": "normal"} + + +@pytest.mark.parametrize("defaults", [{}, {"use_compile": True}, {"use_compile": True, "warmup": True}, {"warmup": False}]) +def test_warmup_is_only_read_from_system_config(cli_inputs, defaults): + config, request = cli_inputs(defaults) + + assert config["warmup"] is defaults.get("warmup", False) + assert request == {"task": "t2v"} + + +def test_cli_task_is_explicit_and_cannot_be_replaced_by_json(cli_inputs): + config, request = cli_inputs({"task": "i2v"}, task="t2v") + + assert config["task"] == request["task"] == "t2v" + + +@pytest.mark.parametrize("option", ["--warmup", "--no-warmup"]) +def test_static_cli_does_not_accept_warmup(cli_inputs, capsys, option): + with pytest.raises(SystemExit) as exc: + cli_inputs({}, [option]) + + assert exc.value.code == 2 + assert f"unrecognized arguments: {option}" in capsys.readouterr().err + + +@pytest.mark.parametrize( + "model_cls,runner_cls,task,defaults,options,expected", + [ + ( + "flux2", + Flux2Runner, + "i2i", + {"inpaint_mask_enabled": True}, + ["--inpaint_blur_sigma", "0.5", "--inpaint_blur_size", "3"], + {"inpaint_blur_sigma": 0.5, "inpaint_blur_size": 3}, + ), + ( + "hunyuan_image3", + HunyuanImage3Runner, + "t2t", + {"enable_cfg": False, "bot_task": "auto", "moe_backend": "torch", "text_do_sample": True}, + [ + "--bot_task", + "think_recaption", + "--max_new_tokens", + "128", + "--system_prompt", + "Be concise.", + "--no-text_do_sample", + "--text_temperature", + "0.7", + "--text_top_k", + "0", + "--text_top_p", + "0.9", + ], + {"bot_task": "think_recaption", "max_new_tokens": 128, "system_prompt": "Be concise.", "text_do_sample": False, "text_temperature": 0.7, "text_top_k": 0, "text_top_p": 0.9}, + ), + ( + "hunyuan_image3", + HunyuanImage3Runner, + "t2t", + {"enable_cfg": False, "bot_task": "auto", "moe_backend": "torch", "text_do_sample": False}, + ["--text_do_sample"], + {"text_do_sample": True}, + ), + ( + "hunyuan_image3", + HunyuanImage3Runner, + "i2i", + {"moe_backend": "torch", "infer_align_image_size": True}, + ["--no-infer_align_image_size"], + {"infer_align_image_size": False}, + ), + ( + "infinitetalk", + InfiniteTalkRunner, + "s2v", + {}, + ["--video_duration", "2.5"], + {"video_duration": 2.5}, + ), + ( + "sensenova_vision", + SenseNovaVisionRunner, + "omni_vision_task", + {"omni_vision_subtask": "depth", "postprocess_predictions": True}, + ["--raw_output_path", "raw.npz", "--glb_output_path", "scene.glb", "--no-postprocess_predictions"], + {"raw_output_path": "raw.npz", "glb_output_path": "scene.glb", "postprocess_predictions": False}, + ), + ], +) +def test_cli_model_specific_options_reach_runner(cli_inputs, model_cls, runner_cls, task, defaults, options, expected): + config, request = cli_inputs(defaults, options, model_cls=model_cls, task=task) + runner = object.__new__(runner_cls) + BaseRunner.__init__(runner, config) + input_info = runner.prepare_request(request) + + for field, value in expected.items(): + assert request[field] == value + assert getattr(input_info, field) == value + + +@pytest.mark.parametrize( + "model_cls,runner_cls,task,field", + [ + ("hidream_o1_image", HidreamO1ImageRunner, "i2i", "keep_original_aspect"), + ("worldmirror", WorldMirrorRunner, "recon", "save_rendered"), + ("worldmirror", WorldMirrorRunner, "recon", "render_depth"), + ], +) +def test_cli_boolean_can_disable_json_default(cli_inputs, model_cls, runner_cls, task, field): + config, request = cli_inputs({field: True}, [f"--no-{field}"], model_cls=model_cls, task=task) + runner = object.__new__(runner_cls) + BaseRunner.__init__(runner, config) + assert config[field] is True + assert getattr(runner.prepare_request(request), field) is False + assert getattr(runner.prepare_request({}), field) is True diff --git a/tests/test_disagg_request_config.py b/tests/test_disagg_request_config.py new file mode 100644 index 000000000..91c3106b2 --- /dev/null +++ b/tests/test_disagg_request_config.py @@ -0,0 +1,402 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightx2v.disagg.conn import DataPoll +from lightx2v.disagg.disagg_mixin import ( + DisaggMixin, + _estimate_encoder_buffer_sizes, + validate_disagg_buffer_capacity, + wait_for_disagg_transfer, +) +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS +from lightx2v.server.schema import ImageTaskRequest, VideoTaskRequest +from lightx2v.utils.input_info import I2IInputInfo, T2VInputInfo +from lightx2v.utils.lockable_dict import LockableDict + + +def create_disagg_runner(mode="encoder", decentralized=True): + runner = object.__new__(DisaggMixin) + runner.config = LockableDict( + { + "task": "t2v", + "model_cls": "wan2.1", + "infer_steps": 40, + "target_video_length": 81, + "target_height": 480, + "target_width": 832, + "disagg_config": { + "bootstrap_room": 10, + "sender_engine_rank": 0, + "receiver_engine_rank": 3, + }, + } + ) + runner.config.lock() + runner._disagg_mode = mode + runner._disagg_decentralized = decentralized + runner._disagg_bootstrap_room = 10 + runner._disagg_sender_rank = 0 + runner._disagg_receiver_rank = 3 + runner._disagg_request_config = None + return runner + + +class SeedDisaggRunner(DisaggMixin, BaseRunner): + supported_request_fields_by_task = {"t2v": COMMON_REQUEST_FIELDS} + + def run_pipeline(self, input_info): + input_info.update(self._disagg_request_config or {}) + return input_info.seed + + +@pytest.mark.parametrize("selected_seed", [0, 73]) +def test_disagg_uses_upstream_seed_before_seeding_each_stage(selected_seed, monkeypatch): + seeded = [] + monkeypatch.setattr("lightx2v.models.runners.base_runner.seed_all", seeded.append) + encoder = SeedDisaggRunner({"task": "t2v", "seed": selected_seed}) + # CLI prepares its request before init_disagg creates the dispatch state. + encoder_input = encoder.prepare_request({"task": encoder.config["task"], "seed": selected_seed}) + assert encoder_input.seed == selected_seed + request_config = {"seed": encoder_input.seed} + + for mode in ("transformer", "decode"): + runner = SeedDisaggRunner({"task": "t2v", "seed": 999, "disagg_mode": mode}) + runner._gc_frozen = True + runner._disagg_request_config = request_config.copy() + input_info = runner.prepare_request({"task": runner.config["task"], "seed": 111}) + assert input_info.seed == selected_seed + assert runner.run_request(input_info) == selected_seed + assert runner.config["seed"] == 999 + request_config = runner._disagg_request_config.copy() + + assert seeded == [selected_seed, selected_seed] + + runner._disagg_request_config = None + assert runner.prepare_request({"task": runner.config["task"], "seed": 9}).seed is None + assert runner.prepare_request({"task": runner.config["task"]}).seed is None + assert encoder.prepare_request({"task": encoder.config["task"], "seed": 9}).seed == 9 + assert encoder.prepare_request({"task": encoder.config["task"]}).seed == 42 + + +@pytest.mark.parametrize("request_cls", [ImageTaskRequest, VideoTaskRequest]) +def test_public_requests_do_not_expose_disagg_routing(request_cls): + routing_fields = { + "data_bootstrap_room", + "disagg_bootstrap_room", + "disagg_decoder_bootstrap_room", + "disagg_phase1_receiver_engine_rank", + } + + assert routing_fields.isdisjoint(request_cls.model_fields) + + +def test_disagg_request_config_keeps_startup_config_unchanged(): + runner = create_disagg_runner() + first = T2VInputInfo( + prompt="first", + seed=1, + target_video_length=49, + target_shape=[480, 832], + return_result_tensor=True, + ) + second = T2VInputInfo( + prompt="second", + seed=2, + target_video_length=81, + target_shape=[720, 1280], + ) + + first_config = runner.build_disagg_request_config( + first, + { + "data_bootstrap_room": 101, + "disagg_phase1_receiver_engine_rank": 3, + }, + ) + second_config = runner.build_disagg_request_config( + second, + { + "data_bootstrap_room": 102, + "disagg_phase1_receiver_engine_rank": 4, + }, + ) + + assert first_config["data_bootstrap_room"] == 101 + assert second_config["data_bootstrap_room"] == 102 + assert first_config["disagg_phase1_receiver_engine_rank"] == 3 + assert second_config["disagg_phase1_receiver_engine_rank"] == 4 + assert first_config["target_video_length"] == 49 + assert first_config["target_height"] == 480 + assert first_config["target_width"] == 832 + assert first_config["return_result_tensor"] is True + assert second_config["target_video_length"] == 81 + assert second_config["target_height"] == 720 + assert second_config["target_width"] == 1280 + + assert "data_bootstrap_room" not in runner.config + assert runner.config["target_video_length"] == 81 + assert runner.config["target_height"] == 480 + assert runner.config["target_width"] == 832 + with pytest.raises(TypeError, match="Dictionary is locked"): + runner.config["target_video_length"] = 49 + + +def test_disagg_payload_carries_request_geometry(): + runner = create_disagg_runner() + request_config = runner.build_disagg_request_config( + T2VInputInfo( + prompt="request", + seed=7, + target_video_length=49, + target_shape=[480, 832], + ) + ) + + payload = runner._disagg_build_request_config_snapshot(request_config) + + assert payload["target_video_length"] == 49 + assert payload["target_shape"] == [480, 832] + assert payload["target_height"] == 480 + assert payload["target_width"] == 832 + assert payload["return_result_tensor"] is False + assert "infer_steps" not in payload + assert "model_cls" not in payload + + +def test_i2i_dispatch_preserves_reference_buffer_size_and_denoise_strength(monkeypatch): + class MemoryDataManager: + def __init__(self): + self.data_args = {} + + def init(self, args, room): + self.data_args[room] = args + + def get_localhost(self): + return "127.0.0.1" + + def get_session_id(self): + return "test-session" + + def allocate_buffers(runner, sizes): + runner._disagg_rdma_buffers = [torch.zeros(size, dtype=torch.uint8) for size in sizes] + + def allocate_p2_buffers(runner, sizes): + runner._disagg_p2_rdma_buffers = [torch.zeros(size, dtype=torch.uint8) for size in sizes] + + monkeypatch.setattr(DisaggMixin, "_disagg_alloc_buffers", allocate_buffers) + monkeypatch.setattr(DisaggMixin, "_disagg_alloc_p2_buffers", allocate_p2_buffers) + monkeypatch.setattr("lightx2v.disagg.disagg_mixin.DataSender", lambda mgr, host, room: SimpleNamespace(bootstrap_room=room)) + monkeypatch.setattr("lightx2v.disagg.disagg_mixin.DataReceiver", lambda mgr, host, room: SimpleNamespace(init=lambda: None)) + encoder = create_disagg_runner() + transformer = create_disagg_runner(mode="transformer") + for runner in (encoder, transformer): + runner.config = { + **runner.config, + "task": "i2i", + "target_video_length": 1, + "vae_stride": [1, 8, 8], + "text_len": 2, + "text_encoder_dim": 2, + "use_image_encoder": False, + "disagg_config": {**runner.config["disagg_config"], "decoder_engine_rank": 2}, + } + runner._disagg_data_mgr = MemoryDataManager() + runner._disagg_bootstrap_addr = "127.0.0.1" + encoder._disagg_active_encoder_room = None + transformer._disagg_active_transformer_room = None + transformer._disagg_p2_data_mgr = MemoryDataManager() + phase1_packets, phase2_packets = [], [] + encoder._disagg_phase1_queue = SimpleNamespace(produce=phase1_packets.append) + transformer._disagg_phase2_queue = SimpleNamespace(produce=phase2_packets.append) + + input_info = I2IInputInfo(seed=42, target_shape=[512, 512], i2i_denoise_strength=0.5) + request_config = encoder.build_disagg_request_config(input_info, {"data_bootstrap_room": 101}) + sizes = _estimate_encoder_buffer_sizes(encoder._disagg_effective_config(request_config)) + # A 1024x1024 reference image encoded into bfloat16 packed latents. + sizes[1] = 4096 * 64 * 2 + encoder._disagg_encoder_setup_room(101, request_config, sizes) + encoder._disagg_produce_phase1_for_encoder(request_config) + packet = phase1_packets[0] + transformer.disagg_transformer_prepare_dispatch(packet) + + assert packet["buffer_sizes"] == sizes + assert transformer._disagg_data_mgr.data_args[101].data_lens == sizes + assert transformer._disagg_request_config["i2i_denoise_strength"] == 0.5 + assert phase2_packets[0]["request_config"]["i2i_denoise_strength"] == 0.5 + assert transformer._disagg_request_config["target_shape"] == [512, 512] + + +def test_static_buffer_capacity_comes_from_startup_config(): + class Buffer: + def __init__(self, size): + self.size = size + + def numel(self): + return self.size + + runner = create_disagg_runner(decentralized=False) + capacity_config = dict(runner.config) + capacity_config["task"] = "i2v" + + assert capacity_config["target_video_length"] == 81 + assert capacity_config["target_height"] == 480 + assert capacity_config["target_width"] == 832 + + capacity_sizes = _estimate_encoder_buffer_sizes(capacity_config) + buffers = [Buffer(size) for size in capacity_sizes] + + request_config = dict(capacity_config) + request_config.update(target_video_length=49, target_height=480, target_width=832) + validate_disagg_buffer_capacity(buffers, _estimate_encoder_buffer_sizes(request_config), "Phase 1") + + request_config.update(target_video_length=81, target_height=720, target_width=1280) + with pytest.raises(ValueError, match="startup config"): + validate_disagg_buffer_capacity(buffers, _estimate_encoder_buffer_sizes(request_config), "Phase 1") + + +def test_disagg_transfer_failure_is_reported(monkeypatch): + class Transfer: + def __init__(self): + self.statuses = iter((DataPoll.WaitingForInput, DataPoll.Failed)) + + def poll(self): + return next(self.statuses) + + monkeypatch.setattr("lightx2v.disagg.disagg_mixin.time.sleep", lambda _: None) + + with pytest.raises(RuntimeError, match="Encoder to Transformer transfer failed"): + wait_for_disagg_transfer(Transfer(), "Encoder to Transformer transfer") + + +def test_downstream_request_context_is_restored_without_inheritance(): + runner = create_disagg_runner(mode="transformer") + runner._disagg_request_config = { + "data_bootstrap_room": 25, + "prompt": "upstream", + "seed": 9, + "target_video_length": 49, + "target_shape": [480, 832], + "target_height": 480, + "target_width": 832, + } + input_info = T2VInputInfo() + + input_info.update(runner._disagg_request_config or {}) + request_config = runner.build_disagg_request_config(input_info) + + assert input_info.prompt == "upstream" + assert input_info.target_video_length == 49 + assert input_info.target_shape == [480, 832] + assert request_config["data_bootstrap_room"] == 25 + + runner._disagg_request_config = None + next_config = runner.build_disagg_request_config( + T2VInputInfo( + prompt="next", + seed=10, + target_video_length=81, + target_shape=[720, 1280], + ), + {"data_bootstrap_room": 26}, + ) + + assert next_config["data_bootstrap_room"] == 26 + assert next_config["prompt"] == "next" + assert next_config["target_video_length"] == 81 + assert next_config["target_shape"] == [720, 1280] + + +def test_effective_disagg_config_uses_request_context_without_mutating_startup_config(): + runner = create_disagg_runner(mode="transformer") + request_config = { + "target_video_length": 49, + "disagg_config": { + "sender_engine_rank": 3, + "receiver_engine_rank": 4, + }, + } + + config = runner._disagg_effective_config(request_config) + payload = runner._disagg_build_request_config_snapshot(request_config) + + assert config["target_video_length"] == 49 + assert config["infer_steps"] == 40 + assert config["model_cls"] == "wan2.1" + assert config["disagg_config"]["sender_engine_rank"] == 3 + assert config["disagg_config"]["receiver_engine_rank"] == 4 + assert runner.config["target_video_length"] == 81 + assert runner.config["disagg_config"]["sender_engine_rank"] == 0 + assert runner.config["disagg_config"]["receiver_engine_rank"] == 3 + assert "infer_steps" not in payload + assert "model_cls" not in payload + + +def test_request_context_reaches_every_stage_without_leaking_to_the_next_request(): + encoder = create_disagg_runner(mode="encoder") + transformer = create_disagg_runner(mode="transformer") + decoder = create_disagg_runner(mode="decode") + + first_input = T2VInputInfo( + prompt="first", + negative_prompt="first negative", + save_result_path="first.mp4", + return_result_tensor=True, + seed=1, + target_video_length=49, + target_shape=[480, 832], + ) + first_config = encoder._disagg_build_request_config_snapshot( + encoder.build_disagg_request_config( + first_input, + { + "data_bootstrap_room": 101, + "disagg_phase1_receiver_engine_rank": 3, + }, + ) + ) + + transformer._disagg_request_config = first_config + transformer_input = T2VInputInfo() + transformer_input.update(transformer._disagg_request_config or {}) + transformer_config = transformer._disagg_build_request_config_snapshot(transformer.build_disagg_request_config(transformer_input)) + + decoder._disagg_request_config = transformer_config + decoder_input = T2VInputInfo() + decoder_input.update(decoder._disagg_request_config or {}) + + assert decoder_input.prompt == "first" + assert decoder_input.negative_prompt == "first negative" + assert decoder_input.save_result_path == "first.mp4" + assert decoder_input.return_result_tensor is True + assert decoder_input.seed == 1 + assert decoder_input.target_video_length == 49 + assert decoder_input.target_shape == [480, 832] + + encoder._disagg_request_config = None + transformer._disagg_request_config = None + decoder._disagg_request_config = None + + second_config = encoder.build_disagg_request_config( + T2VInputInfo( + prompt="second", + save_result_path="second.mp4", + seed=2, + target_video_length=81, + target_shape=[720, 1280], + ), + { + "data_bootstrap_room": 102, + "disagg_phase1_receiver_engine_rank": 4, + }, + ) + + assert second_config["data_bootstrap_room"] == 102 + assert second_config["prompt"] == "second" + assert second_config["save_result_path"] == "second.mp4" + assert second_config["target_video_length"] == 81 + assert second_config["target_shape"] == [720, 1280] + assert "negative_prompt" in second_config + assert second_config["negative_prompt"] == "" diff --git a/tests/test_disagg_service_seed.py b/tests/test_disagg_service_seed.py new file mode 100644 index 000000000..8ed5a827e --- /dev/null +++ b/tests/test_disagg_service_seed.py @@ -0,0 +1,290 @@ +import importlib.util +import json +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from lightx2v.disagg.examples import run_service, run_user +from lightx2v.disagg.workload import StageSpec +from lightx2v.models.schedulers.wan.scheduler import WanScheduler + + +@pytest.fixture +def service_modules(monkeypatch): + # The CPU tests replace only the unavailable pyverbs transport imports. + modules = {} + service_dir = Path(run_service.__file__).parents[1] / "services" + with monkeypatch.context() as transport_imports: + for name, class_name in (("rdma_server", "RDMAServer"), ("rdma_client", "RDMAClient")): + module = ModuleType(f"lightx2v.disagg.{name}") + setattr(module, class_name, object) + transport_imports.setitem(sys.modules, module.__name__, module) + for name in ("controller", "encoder", "transformer", "decoder"): + spec = importlib.util.spec_from_file_location(f"test_disagg_{name}", service_dir / f"{name}.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + modules[name] = module + return modules + + +@pytest.mark.parametrize("role", ["controller", "encoder", "transformer", "decoder"]) +@pytest.mark.parametrize("seed_args, expected_seed", [([], 42), (["--seed", "0"], 0), (["--seed", "12"], 12)]) +def test_service_cli_separates_startup_config_and_automatic_request(service_modules, monkeypatch, tmp_path, role, seed_args, expected_seed): + config_path = tmp_path / "deployment.json" + config_path.write_text(json.dumps({"cpu_offload": True})) + argv = ["run_service", "--service", role, "--model_path", str(tmp_path), "--config_json", str(config_path), "--prompt", "request", "--image_path", "input.png", *seed_args] + if role == "controller": + argv.extend(["--save_result_path", "./output.mp4"]) + monkeypatch.setattr(sys, "argv", argv) + module = service_modules[role] + service = Mock() + service_cls = Mock(return_value=service) + monkeypatch.setattr(module, f"{role.title()}Service", service_cls) + monkeypatch.setitem(sys.modules, f"lightx2v.disagg.services.{role}", module) + seed_all = Mock() + monkeypatch.setattr(run_service, "seed_all", seed_all) + + run_service.main() + + if role == "controller": + config, request = service.run.call_args.args + assert request["seed"] == expected_seed + assert request["prompt"] == "request" + assert request["image_path"] == "input.png" + assert request["save_path"] == "./output.mp4" + else: + (config,) = service_cls.call_args.args + service.run.assert_called_once_with() + assert not {"seed", "prompt", "negative_prompt", "image_path", "save_path"} & config.keys() + assert config["cpu_offload"] is True + seed_all.assert_called_once_with(expected_seed) + + +@pytest.mark.parametrize("seed", [None, 0, 73]) +def test_spawn_command_contains_only_startup_arguments(service_modules, tmp_path, seed): + config_path = tmp_path / "encoder.json" + config_path.write_text(json.dumps({"cpu_offload": True})) + controller = service_modules["controller"].ControllerService.__new__(service_modules["controller"].ControllerService) + command = controller._build_service_command("encoder", 0, {"model_path": str(tmp_path), "seed": seed}, str(config_path)) + args = run_service._build_parser().parse_args(command[3:]) + + assert args.seed == 42 + assert not {"--seed", "--prompt", "--negative_prompt", "--image_path", "--save_result_path"} & set(command) + + +@pytest.fixture +def controller_runtime(service_modules, monkeypatch): + module = service_modules["controller"] + controller = module.ControllerService() + packets = [] + monkeypatch.setattr(module.time, "sleep", lambda _: None) + monkeypatch.setenv("DISAGG_AUTO_REQUEST_COUNT", "2") + monkeypatch.setenv("IS_CENTRALIZED", "0") + monkeypatch.setenv("ENABLE_MONITOR", "0") + monkeypatch.setenv("DISAGG_INSTANCE_WARMUP_WAIT_S", "0") + for name in ("_init_gpu_pool", "_init_request_rdma_buffer", "create_instance", "_run_centralized_ok_server", "_dump_controller_metrics"): + monkeypatch.setattr(controller, name, lambda *args, **kwargs: None) + monkeypatch.setattr(controller, "_drain_decoder_results_non_block", lambda **kwargs: kwargs["received_rooms"].update(kwargs["expected_rooms"])) + controller.req_mgr = SimpleNamespace(send=lambda host, port, packet: packets.append(dict(packet))) + controller.rdma_buffer_request = SimpleNamespace(produce=lambda packet: packets.append(dict(packet))) + for offset, policy in enumerate((controller.encoder_policy, controller.transformer_policy, controller.decoder_policy)): + monkeypatch.setattr(policy, "schedule", lambda offset=offset: f"127.0.0.1:{module.REQUEST_POLLING_PORT + offset}") + return controller, packets + + +@pytest.mark.parametrize("centralized", [False, True]) +@pytest.mark.parametrize("startup_config", [{}, {"seed": None}, {"seed": 0}, {"seed": 73}]) +@pytest.mark.parametrize("load_from_user", [False, True]) +def test_controller_resolves_each_request_before_dispatch(controller_runtime, monkeypatch, centralized, startup_config, load_from_user): + controller, packets = controller_runtime + monkeypatch.setenv("LOAD_FROM_USER", str(int(load_from_user))) + monkeypatch.setenv("IS_CENTRALIZED", str(int(centralized))) + requests = [{"save_path": f"./external_{index}.mp4", **seed} for index, seed in enumerate([{}, {"seed": None}, {"seed": 0}, {"seed": 12}, {}])] + workload = iter([*requests, {"workload_end": True}]) + controller.req_mgr.receive = lambda _: next(workload) + config = dict(startup_config) + request_data = {"seed": 73, "prompt": "automatic prompt", "negative_prompt": "", "image_path": "input.png", "save_path": "auto.mp4"} + original_request = dict(request_data) + + controller.run(config, request_data) + + expected = [42, 42, 0, 12, 42] if load_from_user else [73, 73] + replicas = 3 if centralized else 1 + assert [packet["seed"] for packet in packets] == [seed for seed in expected for _ in range(replicas)] + paths = [request["save_path"] for request in requests] if load_from_user else ["auto0.mp4", "auto1.mp4"] + assert [packet["save_path"] for packet in packets] == [path for path in paths for _ in range(replicas)] + for packet in packets: + if load_from_user: + assert "prompt" not in packet + assert "image_path" not in packet + else: + assert all(packet[key] == value for key, value in request_data.items() if key != "save_path") + assert config == startup_config + assert request_data == original_request + + +@pytest.mark.parametrize("request_data", [{}, {"save_path": None}, {"save_path": ""}]) +@pytest.mark.parametrize("load_from_user", [False, True]) +def test_controller_rejects_missing_output_before_dispatch(controller_runtime, monkeypatch, request_data, load_from_user): + controller, packets = controller_runtime + monkeypatch.setenv("LOAD_FROM_USER", str(int(load_from_user))) + workload = iter([request_data, {"workload_end": True}]) + controller.req_mgr.receive = lambda _: next(workload) + metrics = Mock() + monkeypatch.setattr(controller, "_dump_controller_metrics", metrics) + if load_from_user: + controller.run({"save_path": "startup.mp4"}, request_data) + result = metrics.call_args.args[0][0] + assert result["ok"] is False + assert "save_path is required" in result["error"] + else: + with pytest.raises(ValueError, match="save_path is required"): + controller.run({"save_path": "startup.mp4"}, request_data) + assert packets == [] + + +def test_invalid_external_request_does_not_interrupt_the_batch(controller_runtime, monkeypatch): + controller, packets = controller_runtime + monkeypatch.setenv("LOAD_FROM_USER", "1") + workload = iter([{"save_path": "first.mp4"}, {"save_path": None}, {"save_path": "last.mp4"}, {"workload_end": True}]) + controller.req_mgr.receive = lambda _: next(workload) + controller.run({}) + assert [(packet["data_bootstrap_room"], packet["save_path"]) for packet in packets] == [(0, "first.mp4"), (2, "last.mp4")] + + +def test_automatic_output_names_use_resolved_rooms(controller_runtime, monkeypatch): + controller, packets = controller_runtime + monkeypatch.setenv("LOAD_FROM_USER", "0") + controller.run({"data_bootstrap_room": 7}, {"save_path": "outputs/video.test.mp4"}) + assert [(packet["data_bootstrap_room"], packet["save_path"]) for packet in packets] == [(7, "outputs/video.test7.mp4"), (8, "outputs/video.test8.mp4")] + + +@pytest.mark.parametrize("request_data", [{}, {"save_path": None}, {"save_path": ""}]) +def test_decoder_rejects_missing_output_before_computation(service_modules, monkeypatch, request_data): + module = service_modules["decoder"] + service = module.DecoderService.__new__(module.DecoderService) + seed_all = Mock() + monkeypatch.setattr(module, "seed_all", seed_all) + with pytest.raises(ValueError, match="save_path is required"): + service.process({"seed": 42, **request_data}) + seed_all.assert_not_called() + + +def test_decoder_saves_to_the_explicit_path(service_modules, monkeypatch): + module = service_modules["decoder"] + service = module.DecoderService.__new__(module.DecoderService) + service.logger = Mock() + latents = torch.zeros(1, 1, 2, 2) + meta = json.dumps({"latents_shape": [1, 1, 2, 2], "latents_dtype": "torch.float32"}).encode() + b"\x00" + service._rdma_buffers = {0: [latents, torch.tensor(list(meta), dtype=torch.uint8)]} + service.data_receiver = {0: object()} + service.vae_decoder = SimpleNamespace(decode=lambda tensor: tensor) + monkeypatch.setattr(module, "AI_DEVICE", "cpu") + monkeypatch.setattr(module, "GET_DTYPE", lambda: torch.float32) + monkeypatch.setattr(module, "seed_all", lambda _: None) + monkeypatch.setattr(module, "wan_vae_to_comfy", lambda video: video) + save_to_video = Mock() + monkeypatch.setattr(module, "save_to_video", save_to_video) + + assert service.process({"seed": 42, "save_path": "./output.mp4"}) == "./output.mp4" + save_to_video.assert_called_once() + assert save_to_video.call_args.args[1] == "./output.mp4" + + +@pytest.mark.parametrize("seed_args, expected_seed", [([], 42), (["--seed", "0"], 0), (["--seed", "73"], 73)]) +def test_user_workload_carries_request_inputs_and_unique_outputs(monkeypatch, tmp_path, seed_args, expected_seed): + argv = ["run_user", "--prompt", "request prompt", "--negative_prompt", "", "--image_path", "input.png", "--max_requests", "2", *seed_args] + monkeypatch.setattr(sys, "argv", argv) + monkeypatch.setenv("DISAGG_WORKLOAD_SAVE_PREFIX", str(tmp_path / "output.mp4")) + stage = StageSpec("test", duration_s=10, user_count=1, spawn_rate=1) + monkeypatch.setattr(run_user, "load_stage_specs", lambda: [stage]) + monkeypatch.setattr(run_user, "load_base_config", lambda: {"task": "i2v"}) + monkeypatch.setattr(run_user, "DisaggLoadShape", lambda: SimpleNamespace(tick=lambda: (1, 1))) + monkeypatch.setattr(run_user, "current_stage", lambda _: stage) + monkeypatch.setattr(run_user.time, "sleep", lambda _: None) + monkeypatch.setattr(run_user, "send_workload_end_signal", lambda: None) + packets = [] + monkeypatch.setattr(run_user, "ReqManager", lambda: SimpleNamespace(send=lambda host, port, packet: packets.append(packet))) + + run_user.main() + + assert [packet["save_path"] for packet in packets] == [str(tmp_path / "output_test_0.mp4"), str(tmp_path / "output_test_1.mp4")] + assert all(packet["seed"] == expected_seed and packet["prompt"] == "request prompt" and packet["image_path"] == "input.png" and packet["negative_prompt"] == "" for packet in packets) + + +@pytest.mark.parametrize("name", ["encoder", "transformer", "decoder"]) +def test_services_seed_when_computation_starts(service_modules, monkeypatch, name): + module = service_modules[name] + service_class = getattr(module, f"{name.title()}Service") + service = service_class.__new__(service_class) + seeded = [] + + class ComputationReached(Exception): + pass + + def seed_at_computation(seed): + seeded.append(seed) + raise ComputationReached + + monkeypatch.setattr(module, "seed_all", seed_at_computation) + for seed in (0, 73): + with pytest.raises(ComputationReached): + service.process({"seed": seed, "save_path": "./output.mp4"}) + assert seeded == [0, 73] + + +@pytest.mark.parametrize("seed", [0, 73]) +def test_encoder_phase1_packet_keeps_the_resolved_seed(service_modules, monkeypatch, seed): + module = service_modules["encoder"] + service = module.EncoderService.__new__(module.EncoderService) + service.logger = SimpleNamespace(info=lambda *args: None) + service.text_encoder = SimpleNamespace(infer=lambda prompts: [torch.ones(2, 2)]) + service._rdma_buffers = {1: [torch.empty(size, dtype=torch.uint8) for size in (16, 32, 4096)]} + service.data_sender = {1: SimpleNamespace(send=lambda ptrs: None)} + service.data_mgr = SimpleNamespace(get_localhost=lambda: "encoder", get_session_id=lambda: "session") + service._centralized_request_mode = False + service.sync_comm = False + packets = [] + seeded = [] + service._produce_phase1_request_with_retry = lambda room, packet: packets.append(packet) + monkeypatch.setattr(module, "seed_all", seeded.append) + monkeypatch.setattr(module, "AI_DEVICE", "cpu") + monkeypatch.setattr(module, "GET_DTYPE", lambda: torch.float32) + config = {"seed": seed, "task": "t2v", "prompt": "test", "data_bootstrap_room": 1, "text_len": 2, "target_height": 16, "target_width": 16, "target_video_length": 1, "vae_stride": [4, 8, 8]} + + service.process(config) + + assert seeded == [seed] + assert packets[0]["request_config"]["seed"] == seed + assert packets[0]["request_config"] is not config + + +def test_transformer_discards_previous_request_generator(service_modules, monkeypatch): + module = service_modules["transformer"] + service = module.TransformerService.__new__(module.TransformerService) + service.logger = SimpleNamespace(info=lambda *args: None) + service.scheduler = WanScheduler.__new__(WanScheduler) + service.scheduler.generator = torch.Generator().manual_seed(999) + service.rdma_buffer1 = {} + service.rdma_buffer2 = {} + service.data_receiver = {} + service.data_sender = {} + service._phase2_remote_rooms = set() + monkeypatch.setattr(module, "seed_all", lambda seed: None) + monkeypatch.setattr("lightx2v.models.schedulers.wan.scheduler.AI_DEVICE", "cpu") + + latents = [] + for seed in (73, 12, 73): + config = {"seed": seed, "infer_steps": 1, "sample_shift": 5, "task": "t2v"} + # Stop at the first hardware boundary, after real per-request scheduler preparation. + with pytest.raises(RuntimeError, match="phase1 RDMA buffers"): + service.process(config) + service.scheduler.prepare_latents(seed, (1, 1, 2, 2)) + latents.append(service.scheduler.latents.clone()) + + assert torch.equal(latents[0], latents[2]) + assert not torch.equal(latents[0], latents[1]) diff --git a/tests/test_disagg_static_seed.py b/tests/test_disagg_static_seed.py new file mode 100644 index 000000000..393e19370 --- /dev/null +++ b/tests/test_disagg_static_seed.py @@ -0,0 +1,212 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightx2v.disagg.conn import DataPoll +from lightx2v.disagg.disagg_mixin import DisaggMixin, _estimate_encoder_buffer_sizes +from lightx2v.disagg.utils import estimate_transformer_buffer_sizes +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.qwen_image.qwen_image_runner import QwenImageRunner +from lightx2v.models.runners.wan.wan_runner import WanRunner +from lightx2v.utils.input_info import I2IInputInfo + + +class MemoryTransfer: + def __init__(self, source, destination, data_args=None): + self.source = source + self.destination = destination + self.bootstrap_room = 0 + self.ready = False + self.rearms = 0 + self.data_args = data_args + + def send(self, buffer_ptrs): + for index, (source, destination) in enumerate(zip(self.source, self.destination)): + size = self.data_args.data_item_lens[index] if self.data_args is not None else source.numel() + destination[:size].copy_(source[:size]) + self.ready = True + + def poll(self): + assert self.ready, "The CPU test must send before receiving" + return DataPoll.Success + + def init(self): + self.ready = False + self.rearms += 1 + + +@pytest.mark.parametrize(("runner_cls", "model_cls", "task"), [(WanRunner, "wan2.1", "t2v"), (QwenImageRunner, "qwen_image", "t2i")]) +def test_static_stages_transfer_one_resolved_seed_across_requests(runner_cls, model_cls, task, monkeypatch): + monkeypatch.setattr("lightx2v.disagg.disagg_mixin.AI_DEVICE", "cpu") + monkeypatch.setattr("lightx2v.disagg.disagg_mixin.GET_DTYPE", lambda: torch.float32) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) + seeded = [] + + def seed_cpu(seed): + seeded.append(seed) + torch.manual_seed(seed) + + monkeypatch.setattr("lightx2v.models.runners.base_runner.seed_all", seed_cpu) + monkeypatch.setattr("lightx2v.disagg.disagg_mixin.seed_all", seed_cpu) + + config = { + "model_cls": model_cls, + "task": task, + "seed": 17, + "infer_steps": 1, + "target_video_length": 1, + "target_height": 8, + "target_width": 8, + "vae_stride": [1, 2, 2], + "vae_z_dim": 1, + "vae_scale_factor": 2, + "text_len": 2, + "text_encoder_dim": 2, + "use_image_encoder": False, + "enable_cfg": True, + "disagg_config": {"bootstrap_room": 0}, + } + runners = [] + noises = {role: [] for role in ("encoder", "transformer", "decode")} + latent_shape = [1, 1, 4, 4] + sent_latents, decoded_latents = [], [] + + for role in noises: + runner = runner_cls.__new__(runner_cls) + BaseRunner.__init__(runner, {**config, "disagg_mode": role}) + runner._gc_frozen = True + runner._disagg_mode = role + runner._disagg_decentralized = False + runner._disagg_bootstrap_room = 0 + runner._disagg_request_config = None + runner._disagg_rdma_buffers = [torch.zeros(size, dtype=torch.uint8) for size in _estimate_encoder_buffer_sizes(config)] + runner._disagg_p2_rdma_buffers = [torch.zeros(size, dtype=torch.uint8) for size in estimate_transformer_buffer_sizes(config)] + runner._disagg_data_mgr = SimpleNamespace(data_args={0: SimpleNamespace(data_item_lens=None)}) + runner._disagg_p2_data_mgr = SimpleNamespace(data_args={0: SimpleNamespace(data_item_lens=None)}) + runner.end_run = lambda: None + runners.append(runner) + + encoder, transformer, decoder = runners + phase1 = MemoryTransfer(encoder._disagg_rdma_buffers, transformer._disagg_rdma_buffers) + phase2 = MemoryTransfer(transformer._disagg_p2_rdma_buffers, decoder._disagg_p2_rdma_buffers) + encoder._disagg_sender = transformer._disagg_receiver = phase1 + transformer._disagg_p2_sender = decoder._disagg_p2_receiver = phase2 + + def encode(): + noises["encoder"].append(torch.randn(8)) + encoder.input_info.latent_shape = latent_shape.copy() + context = torch.arange(4, dtype=torch.float32).reshape(1, 2, 2) + if runner_cls is WanRunner: + return {"text_encoder_output": {"context": context, "context_null": -context}} + return {"text_encoder_output": {"prompt_embeds": context, "negative_prompt_embeds": -context}} + + def denoise(): + noises["transformer"].append(torch.randn(8)) + latents = torch.randn(latent_shape) + sent_latents.append(latents.clone()) + return latents + + def decode(latents): + noises["decode"].append(torch.randn(8)) + decoded_latents.append(latents.clone()) + return latents + + encoder.run_input_encoder = encode + decoder.run_vae_decoder = decode + if runner_cls is WanRunner: + transformer._run_transformer_role = lambda: transformer.send_transformer_outputs(denoise()) + decoder.process_images_after_vae_decoder = lambda: decoder.gen_video + else: + for runner in runners: + runner.set_latent_shape = lambda runner=runner: setattr(runner.input_info, "latent_shape", latent_shape.copy()) + transformer.run_dit = lambda: (denoise(), None) + decoder._save_images = lambda *args, **kwargs: None + + with torch.random.fork_rng(devices=[]): + for index, (request_data, expected_seed) in enumerate([({}, 42), ({"seed": 0}, 0), ({"seed": 29}, 29)]): + source_input = encoder.prepare_request({"task": encoder.config["task"], "return_result_tensor": True, **request_data}) + transformer_input = transformer.prepare_request({"task": transformer.config["task"], "seed": 999, "return_result_tensor": True}) + decoder_input = decoder.prepare_request({"task": decoder.config["task"], "seed": 999, "return_result_tensor": True}) + assert source_input.seed == expected_seed + assert transformer_input.seed is None + assert decoder_input.seed is None + + encoder.run_request(source_input) + transformer.run_request(transformer_input) + decoder.run_request(decoder_input) + + assert seeded[-3:] == [expected_seed] * 3 + assert len(seeded) == (index + 1) * 3 + assert all(runner.input_info.seed == expected_seed for runner in runners) + assert all(runner.config["seed"] == 17 for runner in runners) + assert torch.equal(noises["encoder"][-1], noises["transformer"][-1]) + assert torch.equal(noises["encoder"][-1], noises["decode"][-1]) + assert torch.equal(sent_latents[-1], decoded_latents[-1]) + + assert phase1.rearms == phase2.rearms == 3 + assert not torch.equal(noises["encoder"][0], noises["encoder"][1]) + assert not torch.equal(noises["encoder"][0], noises["encoder"][2]) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("enable_cfg", [False, True]) +def test_i2i_transfers_reference_latents_larger_than_output(dtype, enable_cfg, monkeypatch): + monkeypatch.setattr("lightx2v.disagg.disagg_mixin.AI_DEVICE", "cpu") + monkeypatch.setattr("lightx2v.disagg.disagg_mixin.GET_DTYPE", lambda: dtype) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + config = { + "model_cls": "qwen_image", + "task": "i2i", + "target_video_length": 1, + "target_height": 1664, + "target_width": 1664, + "vae_stride": [1, 8, 8], + "vae_z_dim": 16, + "text_len": 2, + "text_encoder_dim": 2, + "clip_embed_dim": 2, + "enable_cfg": enable_cfg, + "use_image_encoder": True, + } + runners = [] + for _ in range(2): + runner = object.__new__(DisaggMixin) + runner.config = config.copy() + runner._disagg_decentralized = False + runner._disagg_bootstrap_room = 0 + runner._disagg_request_config = None + runner._disagg_rdma_buffers = [torch.zeros(size, dtype=torch.uint8) for size in _estimate_encoder_buffer_sizes(config)] + runner._disagg_data_mgr = SimpleNamespace(data_args={0: SimpleNamespace(data_item_lens=None)}) + runners.append(runner) + encoder, transformer = runners + data_args = encoder._disagg_data_mgr.data_args[0] + transfer = MemoryTransfer(encoder._disagg_rdma_buffers, transformer._disagg_rdma_buffers, data_args) + encoder._disagg_sender = transformer._disagg_receiver = transfer + original_ptrs = [buffer.data_ptr() for buffer in encoder._disagg_rdma_buffers] + reference = torch.arange(4096 * 64, dtype=torch.float32).reshape(1, 4096, 64).to(dtype) + context = torch.ones(1, 2, 2, dtype=dtype) + inputs = { + "text_encoder_output": {"prompt_embeds": context, "negative_prompt_embeds": -context if enable_cfg else None}, + "image_encoder_output": [{"image_latents": reference}], + } + + for output_size in (512, 768): + encoder.input_info = I2IInputInfo(seed=42, target_shape=[output_size, output_size]) + transformer.input_info = I2IInputInfo(seed=42, target_shape=[output_size, output_size]) + request_config = encoder.build_disagg_request_config(encoder.input_info) + latent_shape = [1, 1, 16, output_size // 8, output_size // 8] + encoder.send_encoder_outputs(inputs, latent_shape, request_config) + received = transformer.receive_encoder_outputs(request_config) + + assert torch.equal(received["image_encoder_output"][0]["image_latents"], reference) + assert received["latent_shape"] == latent_shape + assert data_args.data_item_lens[2 + int(enable_cfg)] == reference.nbytes + assert [buffer.data_ptr() for buffer in encoder._disagg_rdma_buffers] == original_ptrs + + assert transfer.rearms == 2 + + encoder._disagg_rdma_buffers[2 + int(enable_cfg)] = torch.empty(reference.nbytes - 1, dtype=torch.uint8) + with pytest.raises(ValueError, match="Phase 1 request exceeds"): + encoder.send_encoder_outputs(inputs, latent_shape, request_config) diff --git a/tests/test_disagg_wan_model_selection.py b/tests/test_disagg_wan_model_selection.py new file mode 100644 index 000000000..07211e321 --- /dev/null +++ b/tests/test_disagg_wan_model_selection.py @@ -0,0 +1,58 @@ +import pytest + +from lightx2v.disagg.utils import load_wan_transformer + + +@pytest.mark.parametrize( + ("model_cls", "model_config", "model_count"), + [ + ("wan2.1", {}, 1), + ("wan2.2_moe", {"boundary_step_index": 2}, 2), + ], +) +def test_disagg_loader_uses_distill_model_class(monkeypatch, model_cls, model_config, model_count): + selected_methods = [] + + class RecordingModel: + def __init__(self, **kwargs): + self.model_type = kwargs.get("model_type") + + def select_model_class(distill_method): + selected_methods.append(distill_method) + return RecordingModel + + monkeypatch.setattr("lightx2v.models.runners.wan.wan_runner.get_wan_model_class", select_model_class) + + config = { + "model_cls": model_cls, + "model_path": "/path/to/model", + "distill_method": "dmd2", + "cpu_offload": True, + **model_config, + } + model = load_wan_transformer(config) + models = model.model if model_cls == "wan2.2_moe" else [model] + + assert selected_methods == ["dmd2"] + assert len(models) == model_count + assert all(isinstance(item, RecordingModel) for item in models) + + +@pytest.mark.parametrize("branch,index", [("high_noise_model", 0), ("low_noise_model", 1)]) +def test_disagg_lora_can_target_one_branch(monkeypatch, branch, index): + monkeypatch.setattr("lightx2v.models.networks.wan.model.WanModel", lambda **kwargs: kwargs) + config = { + "model_cls": "wan2.2_moe", + "model_path": "/path/to/model", + "cpu_offload": True, + "boundary": 0.9, + "lora_dynamic_apply": True, + "lora_configs": [{"name": branch, "path": "/path/to/adapter.safetensors", "strength": 0.0}], + } + + models = load_wan_transformer(config).model + + assert models[index]["lora_path"] == "/path/to/adapter.safetensors" + assert models[index]["lora_strength"] == 0.0 + assert "lora_path" not in models[1 - index] + assert "lora_strength" not in models[1 - index] diff --git a/tests/test_entrypoint_defaults.py b/tests/test_entrypoint_defaults.py new file mode 100644 index 000000000..e5ed64cd0 --- /dev/null +++ b/tests/test_entrypoint_defaults.py @@ -0,0 +1,483 @@ +import argparse +import ast +import gc +import importlib.util +import json +import os +import weakref +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +from loguru import logger + +from lightx2v.models.networks.base_model import BaseTransformerModel +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.cosmos3 import cosmos3_runner +from lightx2v.models.runners.wan import wan_runner +from lightx2v.models.runners.wan.wan_runner import Wan22MoeRunner, WanRunner +from lightx2v.utils.lockable_dict import LockableDict +from lightx2v.utils.set_config import get_default_config + +ROOT = Path(__file__).resolve().parents[1] +CLIENTS = ( + "post_async_t2i_and_wait", + "post_sync_t2i_base64", + "post_sync_t2i_presigned", + "post_sync_i2i_base64", + "post_sync_i2i_presigned", +) + + +def load_client(name): + spec = importlib.util.spec_from_file_location(name, ROOT / "scripts/server" / f"{name}.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class CapturedPayload(Exception): + def __init__(self, payload): + self.payload = payload + + +def capture_post(url, *, json, timeout): + raise CapturedPayload(json) + + +@pytest.mark.parametrize("name", CLIENTS) +@pytest.mark.parametrize("explicit", [False, True]) +def test_image_client_main_only_sends_explicit_request_fields(monkeypatch, name, explicit): + client = load_client(name) + argv = [name, "--prompt", "a cat"] + if "i2i" in name: + argv += ["--image_base64", "aW1hZ2U="] + if "presigned" in name: + argv += ["--presigned_url", "https://example.invalid/upload"] + if explicit: + argv += ["--negative_prompt", "", "--seed", "0", "--aspect_ratio", "1:1", "--save_result_path", ""] + monkeypatch.setattr("sys.argv", argv) + monkeypatch.setattr(client.requests, "post", capture_post) + + with pytest.raises(CapturedPayload) as captured: + client.main() + + payload = captured.value.payload + assert payload["prompt"] == "a cat" + expected = {"negative_prompt": "", "seed": 0, "aspect_ratio": "1:1", "save_result_path": ""} + assert {key: payload[key] for key in expected if key in payload} == (expected if explicit else {}) + assert "target_shape" not in payload + if "i2i" in name: + assert payload["image_path"] == "aW1hZ2U=" + + +@pytest.mark.parametrize("output_path", [None, "result.png"]) +def test_async_client_downloads_only_saved_results(monkeypatch, output_path): + client = load_client("post_async_t2i_and_wait") + monkeypatch.setattr("sys.argv", ["post_async_t2i_and_wait", "--prompt", "image"]) + monkeypatch.setattr(client, "submit_t2i_task", lambda **kwargs: "task1") + monkeypatch.setattr(client, "wait_task_done", lambda **kwargs: {"status": "completed", "save_result_path": output_path}) + downloads = [] + monkeypatch.setattr(client, "download_result", lambda *args: downloads.append(args) or Path(args[2])) + + client.main() + + assert len(downloads) == (0 if output_path is None else 1) + + +@pytest.mark.parametrize("explicit", [False, True]) +def test_benchmark_preserves_request_omission(monkeypatch, explicit): + client = load_client("benchmark_sync_s3_latency") + argv = ["benchmark", "--prompt", "a cat"] + if explicit: + argv += ["--negative_prompt", "", "--seed", "0", "--aspect_ratio", "1:1", "--save_result_path", ""] + monkeypatch.setattr("sys.argv", argv) + parse_args = argparse.ArgumentParser.parse_args + + def capture_args(parser): + args = parse_args(parser) + raise CapturedPayload(client.build_sync_payload(args)) + + monkeypatch.setattr(argparse.ArgumentParser, "parse_args", capture_args) + with pytest.raises(CapturedPayload) as captured: + client.main() + + expected = {"prompt": "a cat"} + if explicit: + expected.update(negative_prompt="", seed=0, aspect_ratio="1:1", save_result_path="") + assert captured.value.payload == expected + + +@pytest.mark.parametrize("explicit", [False, True]) +def test_seko_client_preserves_request_omission(monkeypatch, explicit): + client = load_client("post_seko_talk_ar") + argv = ["post_seko_talk_ar"] + if explicit: + argv += ["--negative_prompt", "", "--seed", "0", "--save_result_path", ""] + monkeypatch.setattr("sys.argv", argv) + monkeypatch.setattr(client.requests, "post", capture_post) + + with pytest.raises(CapturedPayload) as captured: + client.main() + + payload = captured.value.payload + expected = {"negative_prompt": "", "seed": 0, "save_result_path": ""} + assert {key: payload[key] for key in expected if key in payload} == (expected if explicit else {}) + + +class RecordingCosmosRunner(cosmos3_runner.Cosmos3Runner): + def __init__(self, config): + BaseRunner.__init__(self, config) + self.scheduler = SimpleNamespace(sample_guide_scale=float(config.get("sample_guide_scale", 4.0))) + self.requests = [] + + def init_modules(self): + pass + + def run_request(self, input_info): + self.requests.append(input_info) + return {"action": np.zeros((2, 8), dtype=np.float32)} + + +@pytest.mark.parametrize( + ("cfg_config", "accepts_negative_prompt"), + [ + ({}, False), + ({"enable_cfg": False}, False), + ({"enable_cfg": False, "sample_guide_scale": 1.0}, False), + ({"enable_cfg": True, "sample_guide_scale": 0.5}, True), + ({"enable_cfg": True, "sample_guide_scale": 6.0}, True), + ], + ids=["defaults", "disabled", "unit_scale", "fractional_scale", "enabled"], +) +def test_cosmos_policy_uses_runner_contract_without_injecting_negative_prompt(monkeypatch, cfg_config, accepts_negative_prompt): + monkeypatch.setattr(cosmos3_runner, "Cosmos3Runner", RecordingCosmosRunner) + config = {"task": "i2va", "action_mode": "policy", "domain_name": "droid_lerobot", **cfg_config} + policy = cosmos3_runner.Cosmos3Policy(config) + images = {name: np.zeros((8, 8, 3), dtype=np.uint8) for name in ("wrist_cam", "over_shoulder_left_camera", "over_shoulder_right_camera")} + + action = policy.next_action(images=images, state=np.zeros(8), task_description="pick up the cup") + + assert action.shape == (8,) + assert policy.runner.requests[0].negative_prompt == "" + assert policy.runner.requests[0].policy_image.shape == (540, 640, 3) + for negative_prompt in ("", "blur"): + if accepts_negative_prompt: + assert policy.runner.prepare_request({"task": policy.runner.config["task"], "negative_prompt": negative_prompt}).negative_prompt == negative_prompt + else: + with pytest.raises(ValueError, match="negative_prompt"): + policy.runner.prepare_request({"task": policy.runner.config["task"], "negative_prompt": negative_prompt}) + + policy.reset() + with pytest.raises(ValueError, match="state length"): + policy.next_action(images=images, state=np.zeros(7), task_description="pick up the cup") + + +@pytest.mark.parametrize(("field", "value"), [("action_mode", "other"), ("domain_name", "other")]) +def test_cosmos_policy_keeps_domain_boundary(field, value): + config = {"action_mode": "policy", "domain_name": "droid_lerobot", field: value} + with pytest.raises(ValueError, match=field): + cosmos3_runner.Cosmos3Policy(config) + + +@pytest.fixture +def gradio_harness(): + # Load only this function to avoid starting the Gradio UI and its file logger. + path = ROOT / "app/gradio_demo.py" + tree = ast.parse(path.read_text()) + function = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "run_inference") + requests = [] + configs = [] + + def build_runner(config): + assert isinstance(config, LockableDict) + configs.append(config) + runner_cls = Wan22MoeRunner if config["model_cls"] == "wan2.2_moe" else WanRunner + runner = object.__new__(runner_cls) + BaseRunner.__init__(runner, config) + runner.run_request = lambda input_info: requests.append(input_info) + runner.switch_lora = Mock(return_value=True) + return runner + + auto_fields = ( + "rope_chunk", + "rope_chunk_size", + "cpu_offload", + "offload_granularity", + "lazy_load", + "t5_cpu_offload", + "clip_cpu_offload", + "vae_cpu_offload", + "unload_modules", + "attention_type", + "quant_op", + "use_tiling_vae", + "clean_cuda_cache", + ) + namespace = { + "argparse": argparse, + "json": json, + "os": os, + "gc": gc, + "logger": logger, + "cleanup_memory": lambda: None, + "get_auto_config_dict": lambda **kwargs: {f"{name}_val": False for name in auto_fields}, + "extract_op_name": lambda value: value, + "generate_unique_filename": lambda *args, **kwargs: "result.mp4", + "get_default_config": get_default_config, + "build_runner": build_runner, + "output_dir": ".", + "global_runner": None, + "current_startup_config": None, + "current_lora_configs": [], + } + model_path = ROOT / "app/utils/model_utils.py" + model_tree = ast.parse(model_path.read_text()) + model_function = next(node for node in model_tree.body if isinstance(node, ast.FunctionDef) and node.name == "get_model_configs") + namespace["build_wan21"] = lambda model_path, dit_path, *args: {"model_cls": "wan2.1", "model_path": model_path, "dit_original_ckpt": dit_path} + namespace["build_wan22"] = lambda model_path, high_path, low_path, *args: { + "model_cls": "wan2.2_moe", + "model_path": model_path, + "high_noise_original_ckpt": high_path, + "low_noise_original_ckpt": low_path, + } + exec(compile(ast.Module(body=[model_function], type_ignores=[]), str(model_path), "exec"), namespace) + exec(compile(ast.Module(body=[function], type_ignores=[]), str(path), "exec"), namespace) + yield namespace, configs, requests + gc.unfreeze() + + +@pytest.mark.parametrize("cfg_scale", [1, 5]) +def test_gradio_config_reaches_runner_with_real_lockable_dict(gradio_harness, cfg_scale): + namespace, configs, requests = gradio_harness + + result = namespace["run_inference"](model_type_input="Wan2.1", model_path_input="/path/to/model", prompt="a cat", image_path="input.png", cfg_scale=cfg_scale, seed=0) + + assert result == "result.mp4" + assert configs[0]["model_cls"] == "wan2.1" + assert configs[0]["model_path"] == "/path/to/model" + assert configs[0]["enable_cfg"] == (cfg_scale != 1) + assert requests[0].seed == 0 + assert requests[0].prompt == "a cat" + + +def test_gradio_releases_frozen_runner_before_loading_replacement(gradio_harness): + namespace, configs, _ = gradio_harness + namespace["cleanup_memory"] = gc.collect + infer = namespace["run_inference"] + kwargs = {"model_type_input": "Wan2.1", "model_path_input": "/models/wan", "prompt": "a cat", "image_path": "input.png"} + infer(**kwargs) + old_runner = weakref.ref(namespace["global_runner"]) + namespace["global_runner"].cycle = namespace["global_runner"] + namespace["global_runner"]._maybe_freeze_gc() + build_runner = namespace["build_runner"] + + def create_replacement(config): + assert namespace["global_runner"] is None + assert old_runner() is None + return build_runner(config) + + namespace["build_runner"] = create_replacement + infer(**kwargs, infer_steps=8) + + assert len(configs) == 2 + + +@pytest.mark.parametrize("model_type", ["Wan2.1", "Wan2.2"]) +def test_gradio_updates_strength_at_same_lora_path(gradio_harness, model_type): + namespace, configs, requests = gradio_harness + infer = namespace["run_inference"] + kwargs = {"model_type_input": model_type, "model_path_input": "/models/wan", "prompt": "a cat", "image_path": "input.png", "use_lora": True, "lora_path": "style.safetensors"} + infer(**kwargs, lora_strength=1.0) + runner = namespace["global_runner"] + + infer(**kwargs, lora_strength=0.5) + + path = "/models/wan/loras/style.safetensors" + if model_type == "Wan2.2": + runner.switch_lora.assert_called_once_with(high_lora_path=path, high_lora_strength=0.5, low_lora_path=path, low_lora_strength=0.5) + else: + runner.switch_lora.assert_called_once_with(path, 0.5) + infer(**{**kwargs, "prompt": "a dog", "image_path": "another.png", "seed": 7}, lora_strength=0.5, num_frames=49, aspect_ratio="16:9") + assert len(configs) == 1 + assert len(requests) == 3 + assert runner.switch_lora.call_count == 1 + + +@pytest.mark.parametrize( + ("model_type", "path_field"), + [("Wan2.1", "dit_path_input"), ("Wan2.2", "high_noise_path_input"), ("Wan2.2", "low_noise_path_input")], +) +def test_gradio_rebuilds_when_checkpoint_config_changes(gradio_harness, model_type, path_field): + namespace, configs, _ = gradio_harness + infer = namespace["run_inference"] + kwargs = {"model_type_input": model_type, "model_path_input": "/models/wan", "prompt": "a cat", "image_path": "input.png"} + infer(**kwargs, **{path_field: "original.safetensors"}) + original_runner = namespace["global_runner"] + + infer(**kwargs, **{path_field: "replacement.safetensors"}) + replacement_runner = namespace["global_runner"] + infer(**kwargs, **{path_field: "replacement.safetensors"}) + + assert len(configs) == 2 + assert replacement_runner is not original_runner + assert namespace["global_runner"] is replacement_runner + + +@pytest.mark.parametrize("model_type", ["Wan2.1", "Wan2.2"]) +def test_gradio_rebuilds_when_enabling_or_disabling_lora(gradio_harness, model_type): + namespace, configs, _ = gradio_harness + infer = namespace["run_inference"] + kwargs = {"model_type_input": model_type, "model_path_input": "/models/wan", "prompt": "a cat", "image_path": "input.png", "lora_path": "style.safetensors"} + infer(**kwargs, use_lora=False) + base_runner = namespace["global_runner"] + infer(**kwargs, use_lora=True) + lora_runner = namespace["global_runner"] + infer(**kwargs, use_lora=False) + + assert len(configs) == 3 + assert lora_runner is not base_runner + assert namespace["global_runner"] is not lora_runner + assert configs[0].get("lora_configs") is None + assert configs[1]["lora_configs"] + assert configs[2].get("lora_configs") is None + + +@pytest.mark.parametrize("branch", ["high", "low"]) +def test_gradio_updates_individual_moe_lora_strength(gradio_harness, branch): + namespace, configs, _ = gradio_harness + infer = namespace["run_inference"] + kwargs = { + "model_type_input": "Wan2.2", + "model_path_input": "/models/wan", + "prompt": "a cat", + "image_path": "input.png", + "use_lora": True, + "high_noise_lora_path": "high.safetensors", + "low_noise_lora_path": "low.safetensors", + } + infer(**kwargs) + runner = namespace["global_runner"] + + infer(**kwargs, **{f"{branch}_noise_lora_strength": 0.0}) + + runner.switch_lora.assert_called_once_with( + high_lora_path="/models/wan/loras/high.safetensors", + high_lora_strength=0.0 if branch == "high" else 1.0, + low_lora_path="/models/wan/loras/low.safetensors", + low_lora_strength=0.0 if branch == "low" else 1.0, + ) + assert len(configs) == 1 + + +@pytest.mark.parametrize("remaining_branch", ["high", "low"]) +@pytest.mark.parametrize("initial_both", [False, True]) +def test_gradio_loads_single_moe_lora_branch(gradio_harness, monkeypatch, remaining_branch, initial_both): + namespace, configs, _ = gradio_harness + namespace["cleanup_memory"] = gc.collect + monkeypatch.setattr(wan_runner, "WanModel", lambda **kwargs: kwargs) + build_runner = namespace["build_runner"] + + def load_models(config): + runner = build_runner(config) + runner.high_noise_model_path = "/models/wan/high_noise_model" + runner.low_noise_model_path = "/models/wan/low_noise_model" + runner.init_device = "cpu" + runner.distill_method = None + config["boundary"] = 0.9 + runner.model = runner.load_transformer() + return runner + + namespace["build_runner"] = load_models + infer = namespace["run_inference"] + kwargs = {"model_type_input": "Wan2.2", "model_path_input": "/models/wan", "prompt": "a cat", "image_path": "input.png", "use_lora": True} + if initial_both: + infer(**kwargs, high_noise_lora_path="high.safetensors", low_noise_lora_path="low.safetensors") + infer(**kwargs, **{f"{remaining_branch}_noise_lora_path": f"{remaining_branch}.safetensors"}) + + assert len(configs) == (2 if initial_both else 1) + models = namespace["global_runner"].model.model + assert len(models) == 2 + for model, branch in zip(models, ("high", "low")): + assert model["model_path"] == f"/models/wan/{branch}_noise_model" + if branch == remaining_branch: + assert model["lora_path"] == f"/models/wan/loras/{branch}.safetensors" + assert model["lora_strength"] == 1.0 + else: + assert "lora_path" not in model + assert "lora_strength" not in model + + +@pytest.mark.parametrize("model_type", ["wan2.1", "high_noise_model", "low_noise_model"]) +@pytest.mark.parametrize("dynamic", [False, True]) +def test_wan_lora_loading_preserves_branch_weights(monkeypatch, model_type, dynamic): + loras = [{"name": name, "path": f"{name}.safetensors", "strength": 0.0} for name in ("high_noise_model", "low_noise_model")] + expected = loras if model_type == "wan2.1" else [item for item in loras if item["name"] == model_type] + adapter = Mock() + adapter_factory = Mock(return_value=adapter) + monkeypatch.setattr(wan_runner, "LoraAdapter", adapter_factory) + + model = wan_runner.build_wan_model_with_lora(lambda **kwargs: kwargs, {"lora_dynamic_apply": dynamic}, {"model_path": "base"}, loras, model_type) + + if dynamic: + assert model == {"model_path": "base", "lora_path": expected[0]["path"], "lora_strength": 0.0} + adapter_factory.assert_not_called() + else: + assert model == {"model_path": "base"} + adapter.apply_lora.assert_called_once_with(expected, model_type=model_type) + + +@pytest.mark.parametrize("dynamic", [False, True]) +def test_wan_missing_branch_loads_base_model_without_lora(monkeypatch, dynamic): + adapter_factory = Mock() + monkeypatch.setattr(wan_runner, "LoraAdapter", adapter_factory) + loras = [{"name": "high_noise_model", "path": "high.safetensors", "strength": 1.0}] + + model = wan_runner.build_wan_model_with_lora(lambda **kwargs: kwargs, {"lora_dynamic_apply": dynamic}, {"model_path": "low"}, loras, "low_noise_model") + + assert model == {"model_path": "low"} + adapter_factory.assert_not_called() + + +@pytest.mark.parametrize("lora_path", [None, "adapter.safetensors"]) +@pytest.mark.parametrize("dynamic", [False, True]) +def test_weight_loading_registers_only_current_model_lora(lora_path, dynamic): + model = SimpleNamespace( + config={"lora_dynamic_apply": dynamic, "lora_configs": [{"name": "high_noise_model", "path": "adapter.safetensors", "strength": 0.0}]}, + lora_path=lora_path, + lora_strength=0.0, + original_weight_dict={"weight": "base"}, + pre_weight=Mock(), + transformer_weights=Mock(), + _register_lora=Mock(), + ) + + BaseTransformerModel._apply_weights(model) + + model.pre_weight.load.assert_called_once_with({"weight": "base"}) + model.transformer_weights.load.assert_called_once_with({"weight": "base"}) + if dynamic and lora_path is not None: + model._register_lora.assert_called_once_with(lora_path, 0.0) + else: + model._register_lora.assert_not_called() + + +def test_gradio_does_not_generate_or_cache_failed_lora_switch(gradio_harness): + namespace, configs, requests = gradio_harness + infer = namespace["run_inference"] + kwargs = {"model_type_input": "Wan2.1", "model_path_input": "/models/wan", "prompt": "a cat", "image_path": "input.png", "use_lora": True, "lora_path": "style.safetensors"} + infer(**kwargs) + runner = namespace["global_runner"] + runner.switch_lora.return_value = False + + with pytest.raises(RuntimeError, match="LoRA"): + infer(**kwargs, lora_strength=0.5) + + assert len(requests) == 1 + runner.switch_lora.return_value = True + infer(**kwargs, lora_strength=0.5) + assert runner.switch_lora.call_count == 2 + assert len(configs) == 1 + assert len(requests) == 2 diff --git a/tests/test_flux2_runner.py b/tests/test_flux2_runner.py new file mode 100644 index 000000000..f06968802 --- /dev/null +++ b/tests/test_flux2_runner.py @@ -0,0 +1,81 @@ +import pytest +import torch + +from lightx2v.models.networks.flux2.model import Flux2DevTransformerModel, Flux2KleinTransformerModel +from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.flux2.flux2_runner import Flux2Runner +from lightx2v.models.schedulers.flux2.feature_caching.scheduler import Flux2DevSchedulerCaching, Flux2SchedulerCaching +from lightx2v.models.schedulers.flux2.scheduler import Flux2DevScheduler, Flux2Scheduler +from lightx2v.utils.registry_factory import RUNNER_REGISTER + + +def test_flux2_runner_registration(): + assert RUNNER_REGISTER["flux2"] is Flux2Runner + assert "flux2_klein" not in RUNNER_REGISTER + assert "flux2_dev" not in RUNNER_REGISTER + + +@pytest.mark.parametrize( + ("model_variant", "transformer_class", "scheduler_class", "caching_scheduler_class"), + [ + ("klein", Flux2KleinTransformerModel, Flux2Scheduler, Flux2SchedulerCaching), + ("dev", Flux2DevTransformerModel, Flux2DevScheduler, Flux2DevSchedulerCaching), + ], +) +def test_flux2_variant_components(monkeypatch, model_variant, transformer_class, scheduler_class, caching_scheduler_class): + monkeypatch.setattr(DefaultRunner, "__init__", lambda self, config: setattr(self, "config", config)) + + runner = Flux2Runner({"model_variant": model_variant}) + + assert runner.transformer_class is transformer_class + assert runner.scheduler_class is scheduler_class + assert runner.caching_scheduler_class is caching_scheduler_class + + +@pytest.mark.parametrize("model_variant", [None, "unknown"]) +def test_flux2_model_variant_is_required(model_variant): + with pytest.raises(ValueError, match="Unsupported Flux2 model_variant"): + Flux2Runner({"model_variant": model_variant}) + + +class RecordingTextEncoder: + def __init__(self): + self.prompts = [] + + def infer(self, prompts): + self.prompts.append(prompts) + return [torch.zeros(2, 4)], None + + +@pytest.mark.parametrize( + ("model_variant", "enable_cfg", "sample_guide_scale", "uses_negative_prompt"), + [ + ("klein", True, 4.0, True), + ("klein", False, 4.0, False), + ("klein", False, 1.0, False), + ("klein", True, 0.5, True), + ("dev", True, 4.0, False), + ], +) +def test_flux2_text_encoder_matches_cfg_execution( + monkeypatch, + model_variant, + enable_cfg, + sample_guide_scale, + uses_negative_prompt, +): + monkeypatch.setattr("lightx2v.models.runners.flux2.flux2_runner.AI_DEVICE", "cpu") + text_encoder = RecordingTextEncoder() + runner = object.__new__(Flux2Runner) + runner.model_variant = model_variant + runner.config = { + "enable_cfg": enable_cfg, + "sample_guide_scale": sample_guide_scale, + } + runner.text_encoders = [text_encoder] + + output = runner.run_text_encoder("prompt") + + assert ("negative_prompt_embeds" in output) is uses_negative_prompt + expected_prompts = [["prompt"], [""]] if uses_negative_prompt else [["prompt"]] + assert text_encoder.prompts == expected_prompts diff --git a/tests/test_input_info_context.py b/tests/test_input_info_context.py new file mode 100644 index 000000000..3c1780597 --- /dev/null +++ b/tests/test_input_info_context.py @@ -0,0 +1,160 @@ +from dataclasses import fields +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from PIL import Image + +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.hidream_o1_image.hidream_o1_image_runner import HidreamO1ImageRunner +from lightx2v.models.runners.lingbot_video.lingbot_video_runner import LingBotVideoRunner +from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner +from lightx2v.models.runners.wan.wan_audio_runner import WanAudioRunner +from lightx2v.models.schedulers.flux2.scheduler import Flux2Scheduler +from lightx2v.shot_runner.rs2v_infer import ShotRS2VPipeline +from lightx2v.utils.input_info import INPUT_INFO_TYPES, UNSET, Cosmos3InputInfo, S2VInputInfo, SekoTalkInputs, T2IInputInfo, VaceInputInfo +from lightx2v.utils.va_controller import VAController + + +@pytest.mark.parametrize("media", ({}, {"video_path": None, "mask_path": None, "src_ref_images": None})) +def test_vace_missing_media_reaches_processor_as_none(media): + runner = object.__new__(DefaultRunner) + runner.input_info = VaceInputInfo(task="vace", target_shape=[480, 832], **media) + runner.get_target_size = Mock(return_value=(480, 832)) + runner.prepare_source = Mock(return_value=([None], [None], [None])) + runner.run_vae_encoder = Mock(return_value=("latents", [16, 21, 60, 104])) + runner.run_text_encoder = Mock(return_value={}) + runner.maybe_empty_cache = Mock() + runner.get_encoder_output_i2v = Mock() + runner._run_input_encoder_local_vace() + runner.prepare_source.assert_called_once_with([None], [None], [None], (832, 480)) + + +def test_vace_reference_paths_remain_a_sequence(): + runner = object.__new__(DefaultRunner) + runner.input_info = VaceInputInfo(task="vace", video_path="clip.mp4", mask_path="mask.mp4", src_ref_images="a.png,b.png") + runner.get_target_size = Mock(return_value=(480, 832)) + runner.prepare_source = Mock(return_value=([None], [None], [None])) + runner.run_vae_encoder = Mock(return_value=("latents", [16, 21, 60, 104])) + runner.run_text_encoder = Mock(return_value={}) + runner.maybe_empty_cache = Mock() + runner.get_encoder_output_i2v = Mock() + runner._run_input_encoder_local_vace() + runner.prepare_source.assert_called_once_with(["clip.mp4"], ["mask.mp4"], [["a.png", "b.png"]], (832, 480)) + + +@pytest.mark.parametrize("task", ("t2i", "t2v")) +@pytest.mark.parametrize("enable_cfg", (False, True)) +def test_lingbot_warmup_preserves_conditioning_and_latent_shape(task, enable_cfg): + runner = object.__new__(LingBotVideoRunner) + runner.config = {"task": task, "target_video_length": 9, "enable_cfg": enable_cfg} + runner.text_encoders = [SimpleNamespace(infer=Mock(side_effect=[{"prompt_embeds": torch.zeros(1, 3, 2), "prompt_mask": None}, {"prompt_embeds": torch.zeros(1, 2, 2), "prompt_mask": None}]))] + runner.maybe_empty_cache = Mock() + encoded = runner._prepare_warmup_inputs(256, 256) + assert encoded["prompt_embeds"].shape == (1, 3, 2) + assert ("negative_prompt_embeds" in encoded) == enable_cfg + if enable_cfg: + assert encoded["negative_prompt_embeds"].shape == (1, 2, 2) + assert runner.input_info.latent_shape == (1, 16, 1 if task == "t2i" else 3, 32, 32) + + +@pytest.mark.parametrize("task,needs_image,needs_last", (("t2av", False, False), ("i2av", True, False), ("l2av", False, True), ("fl2av", True, True), ("ref2av", True, False))) +def test_minimax_warmup_keeps_keyframe_placement(task, needs_image, needs_last): + runner = object.__new__(MiniMaxH3Runner) + runner.config = {"task": task} + runner._prepare_warmup_inputs(64, 64, 9) + assert type(runner.input_info) is INPUT_INFO_TYPES[task] + assert bool(getattr(runner.input_info, "image_path", "")) == needs_image + assert bool(getattr(runner.input_info, "last_frame_path", "")) == needs_last + assert runner.input_info.task == task + if task == "fl2av": + assert runner.input_info.image_path is not runner.input_info.last_frame_path + + +def test_mutable_context_defaults_are_not_shared(): + first, second = Cosmos3InputInfo(), Cosmos3InputInfo() + mutable_fields = [f.name for f in fields(first) if isinstance(getattr(first, f.name), (list, dict))] + assert mutable_fields + for name in mutable_fields: + assert getattr(first, name) is not getattr(second, name) + first.vision_condition_latents = torch.ones(1) + assert second.vision_condition_latents is None + assert second.action_latents is None + + +def test_shot_omitted_values_keep_unset_semantics(): + info = SekoTalkInputs() + info.update({"prompt": "shot", "seed": 0, "target_video_length": 17}) + assert info.prompt == "shot" + assert info.seed == 0 + assert info.target_video_length == 17 + for name in ("negative_prompt", "image_path", "audio_path", "stream_config", "overlap_latent", "audio_clip"): + assert getattr(info, name) is UNSET + + +@pytest.mark.parametrize("stream_config", (None, {}, {"enabled": True})) +def test_stream_controller_keeps_none_and_explicit_configs(stream_config): + controller = object.__new__(VAController) + controller.reader = None + controller.recorder = None + info = S2VInputInfo(target_shape=[480, 832], stream_config=stream_config) + controller.init_base({}, info, False, False) + assert controller.stream_config is stream_config + + +def test_flux2_clears_optional_text_positions_between_requests(monkeypatch): + monkeypatch.setattr("lightx2v.models.schedulers.flux2.scheduler.AI_DEVICE", "cpu") + monkeypatch.setattr("lightx2v.models.schedulers.flux2.scheduler.reset_scheduler_fls_state", Mock()) + scheduler = object.__new__(Flux2Scheduler) + scheduler.generator = None + scheduler.dtype = torch.float32 + scheduler.set_timesteps = Mock() + positions = torch.zeros(1, 3, 4) + scheduler.prepare(T2IInputInfo(seed=42, latent_shape=[1, 4, 2], txt_ids=positions)) + assert scheduler.txt_ids is positions + scheduler.prepare(T2IInputInfo(seed=42, latent_shape=[1, 4, 2])) + assert scheduler.txt_ids is None + + +@pytest.mark.parametrize("request_frames,expected_frames", ((None, 81), (0, 81), (33, 33))) +def test_wan_audio_image_encoding_uses_existing_frame_resolution(monkeypatch, request_frames, expected_frames): + monkeypatch.setattr("lightx2v.models.runners.wan.wan_audio_runner.AI_DEVICE", "cpu") + monkeypatch.setattr("lightx2v.models.runners.wan.wan_audio_runner.resize_image", lambda img, **kwargs: (img, 16, 16)) + runner = object.__new__(WanAudioRunner) + runner.config = {"target_video_length": 81, "vae_stride": [4, 8, 8], "patch_size": [1, 2, 2]} + runner.input_info = S2VInputInfo(target_video_length=request_frames) + runner._get_image_resize_kwargs = Mock(return_value={}) + runner._resolve_patched_spatial_size = Mock(return_value=([16, 16], 1, 1)) + _, latent_shape, _ = runner.read_image_input(Image.new("RGB", (16, 16))) + assert latent_shape == [16, (expected_frames - 1) // 4 + 1, 2, 2] + + +@pytest.mark.parametrize("shape,expected", ((None, None), ([], []), ([16, 4, 8, 8], [16, 9, 8, 8]))) +def test_shot_updates_only_a_prepared_latent_shape(shape, expected): + info = SekoTalkInputs(latent_shape=shape) + ShotRS2VPipeline._update_latent_shape(info, target_len=33, vae_stride=4) + assert info.latent_shape == expected + + +@pytest.mark.parametrize("layout", [UNSET, None, "", '[{"bbox": [0, 0, 32, 32]}]']) +def test_hidream_layout_request_overrides_config_before_encoding(monkeypatch, layout): + configured = '[{"bbox": [0, 0, 64, 64]}]' + runner = object.__new__(HidreamO1ImageRunner) + BaseRunner.__init__(runner, {"task": "i2i", "layout_bboxes": configured}) + runner.model = SimpleNamespace(model_config={}, device="cpu") + runner.tokenizer = runner.processor = None + runner.dtype = torch.float32 + runner._resolve_generation_config = Mock(return_value={"enable_cfg": False}) + build_samples = Mock(return_value={"samples": [], "tgt_image_len": 0}) + monkeypatch.setattr("lightx2v.models.networks.hidream_o1_image.i2i_utils.build_i2i_samples", build_samples) + request = {"image_path": "reference.png", "layout_bboxes": layout} + runner.input_info = runner.prepare_request(request) + + runner._run_input_encoder_local_i2i() + + expected = configured if layout is UNSET or layout is None else layout or None + assert build_samples.call_args.kwargs["layout_bboxes"] == expected + assert runner.config["layout_bboxes"] == configured + assert runner.prepare_request({}).layout_bboxes == configured diff --git a/tests/test_minimax_h3_scheduler.py b/tests/test_minimax_h3_scheduler.py new file mode 100644 index 000000000..d089afce5 --- /dev/null +++ b/tests/test_minimax_h3_scheduler.py @@ -0,0 +1,49 @@ +from unittest.mock import Mock + +import pytest +import torch + +from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner +from lightx2v.models.schedulers.minimax_h3.scheduler import MiniMaxH3Scheduler +from lightx2v.utils.profiler import no_sync_profiling + + +@pytest.fixture(autouse=True) +def cpu_scheduler(monkeypatch): + monkeypatch.setattr("lightx2v.models.schedulers.minimax_h3.scheduler.AI_DEVICE", "cpu") + + +@pytest.mark.parametrize("infer_steps", [1, 4, 5, 8, 29, 30]) +@pytest.mark.parametrize("step_update", ["reference_blend", "training_euler"]) +def test_runner_executes_configured_steps_and_reaches_terminal_zero(infer_steps, step_update): + scheduler = MiniMaxH3Scheduler({"infer_steps": infer_steps, "video_flow_shift": 6.0, "audio_flow_shift": 3.0, "h3_step_update": step_update}) + scheduler.prepare(seed=42, num_frames=124, height=32, width=32, text_token_tags=torch.tensor([1])) + initial_video = scheduler.video_latents.clone() + initial_audio = scheduler.audio_latents.clone() + + def predict_velocity(inputs): + scheduler.video_noise_pred = torch.ones_like(scheduler.video_latents) + scheduler.audio_noise_pred = torch.ones_like(scheduler.audio_latents) + + runner = MiniMaxH3Runner.__new__(MiniMaxH3Runner) + runner.scheduler = scheduler + runner.inputs = {} + runner.model = Mock(infer=Mock(side_effect=predict_velocity)) + runner.check_stop = Mock() + runner.progress_callback = Mock() + + with no_sync_profiling(): + video, audio = runner.run_segment() + + assert runner.model.infer.call_count == infer_steps + runner.progress_callback.assert_called_with(100, 100) + assert scheduler.video_sigmas[scheduler.step_index + 1] == 0 + assert scheduler.audio_sigmas[scheduler.step_index + 1] == 0 + torch.testing.assert_close(video, initial_video + 1, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(audio, initial_audio + 1, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("infer_steps", [0, -1]) +def test_scheduler_rejects_nonpositive_steps(infer_steps): + with pytest.raises(ValueError, match="infer_steps must be at least 1"): + MiniMaxH3Scheduler({"infer_steps": infer_steps}) diff --git a/tests/test_model_seed_contract.py b/tests/test_model_seed_contract.py new file mode 100644 index 000000000..35f985c3e --- /dev/null +++ b/tests/test_model_seed_contract.py @@ -0,0 +1,130 @@ +import random +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner +from lightx2v.models.schedulers.bagel.scheduler import BagelScheduler, get_flattened_position_ids_extrapolate +from lightx2v.models.schedulers.hunyuan3d.flow_match_euler import FlowMatchEulerDiscreteScheduler +from lightx2v.models.schedulers.hunyuan3d.scheduler import Hunyuan3DShapeScheduler +from lightx2v.models.schedulers.hunyuan_image3.scheduler import HunyuanImage3Scheduler +from lightx2v.models.schedulers.wan.infinitetalk.scheduler import InfiniteTalkScheduler + + +@pytest.fixture(params=[({}, 42), ({"seed": None}, 42), ({"seed": 0}, 0), ({"seed": 12}, 12)]) +def resolved_input(request): + request_data, expected_seed = request.param + runner = BaseRunner.__new__(BaseRunner) + runner.config = {"task": "t2i", "seed": 37} + input_info = runner.create_input_info({"task": "t2i", **request_data}) + assert input_info.seed == expected_seed + return input_info + + +def test_bagel_keeps_one_cpu_stream_for_multiple_images(resolved_input): + scheduler = BagelScheduler.__new__(BagelScheduler) + scheduler.config = {"seed": 999} + scheduler.latent_downsample = 2 + scheduler.max_latent_size = 8 + scheduler.latent_channel = 2 + scheduler.latent_patch_size = 1 + scheduler.get_flattened_position_ids = get_flattened_position_ids_extrapolate + args = { + "curr_kvlens": [3, 0], + "curr_rope": [4, 5], + "image_sizes": [(4, 4), (4, 8)], + "new_token_ids": {"start_of_image": 1, "end_of_image": 2}, + } + + result = scheduler.prepare_vae_latent(**args, seed=resolved_input.seed) + expected_generator = torch.Generator(device="cpu").manual_seed(resolved_input.seed) + expected = torch.cat([torch.randn(4, 2, generator=expected_generator), torch.randn(8, 2, generator=expected_generator)]) + assert scheduler.generator.device.type == "cpu" + assert torch.equal(result["packed_init_noises"], expected) + + scheduler.prepare_vae_latent(**args, seed=123) + repeated = scheduler.prepare_vae_latent(**args, seed=resolved_input.seed) + assert torch.equal(repeated["packed_init_noises"], expected) + + +def test_hunyuan_image3_uses_resolved_request_seed(resolved_input, monkeypatch): + monkeypatch.setattr("lightx2v.models.schedulers.hunyuan_image3.scheduler.AI_DEVICE", "cpu") + scheduler = HunyuanImage3Scheduler({"infer_steps": 1, "seed": 999}) + scheduler.prepare(resolved_input) + + expected_generator = torch.Generator(device="cpu").manual_seed(resolved_input.seed) + assert torch.equal(torch.randn(8, generator=scheduler.generator), torch.randn(8, generator=expected_generator)) + + +@pytest.mark.parametrize(("config_seed", "resolved_seed"), [(None, 0), (999, 27182)]) +@pytest.mark.parametrize(("text_config", "generation_options", "text_seed"), [({}, None, None), ({"text_seed": 7}, None, 7), ({"text_seed": 7}, {"text_seed": 13}, 13)]) +def test_hunyuan_image3_cot_uses_request_seed_with_text_overrides(config_seed, resolved_seed, text_config, generation_options, text_seed): + runner = HunyuanImage3Runner.__new__(HunyuanImage3Runner) + runner.config = {"task": "t2i", "seed": config_seed, "max_new_tokens": 3, "text_do_sample": True, **text_config} + runner._gc_frozen = True + runner._close_after_run = False + runner.hunyuan_generation_config = SimpleNamespace() + runner._hunyuan_text_kv_cache_enabled = lambda: False + runner._get_ar_cuda_graph_controller = lambda: SimpleNamespace(enabled=False) + runner._build_text_model_inputs = lambda *args, **kwargs: {} + runner._broadcast_parallel_tensor = lambda tensor: tensor + runner._is_output_rank = lambda: True + logits = torch.tensor([[[0.0, 0.5, 1.0, 1.5]]]) + runner.model = SimpleNamespace(infer=lambda inputs: {"logits": logits}) + plan = SimpleNamespace(stage_transitions=[], final_stop_tokens=set()) + observed_seeds = [] + sample_token = runner._sample_text_token + + def sample(logits, generator, generation_options=None): + observed_seeds.append(generator.initial_seed()) + return sample_token(logits, generator, generation_options) + + runner._sample_text_token = sample + runner.generate_t2i = lambda info: runner._generate_text_tokens(torch.tensor([[1]]), SimpleNamespace(), plan, generation_options=generation_options) + input_info = SimpleNamespace(seed=resolved_seed, return_result_tensor=True) + + result = runner.run_pipeline(input_info) + + expected_seed = resolved_seed if text_seed is None else text_seed + expected_generator = torch.Generator(device="cpu").manual_seed(expected_seed) + expected_tokens = [torch.multinomial(torch.softmax(logits[0, 0], dim=-1), 1, generator=expected_generator).item() for _ in range(3)] + assert runner.input_info is input_info + assert observed_seeds == [expected_seed] * 3 + assert result == {"image": expected_tokens} + + +def test_hunyuan3d_noise_uses_resolved_request_seed(resolved_input): + scheduler = Hunyuan3DShapeScheduler.__new__(Hunyuan3DShapeScheduler) + scheduler.config = {"infer_steps": 2, "seed": 999} + scheduler.flow_scheduler = FlowMatchEulerDiscreteScheduler() + scheduler.device = torch.device("cpu") + scheduler.dtype = torch.float32 + scheduler.prepare(resolved_input.seed, latent_shape=(1, 3, 4)) + + expected_generator = torch.Generator(device="cpu").manual_seed(resolved_input.seed) + expected = torch.randn((1, 3, 4), generator=expected_generator) + assert torch.equal(scheduler.latents, expected) + + +def test_infinitetalk_preserves_global_rng_and_determinism(resolved_input, monkeypatch): + cuda_seeds = [] + monkeypatch.setattr(torch.cuda, "manual_seed_all", cuda_seeds.append) + monkeypatch.setattr(torch.backends.cudnn, "deterministic", False) + scheduler = InfiniteTalkScheduler({"infer_steps": 1, "sample_shift": 5}) + python_state, numpy_state = random.getstate(), np.random.get_state() + + try: + with torch.random.fork_rng(devices=[]): + assert scheduler.seed_everything(resolved_input.seed) == resolved_input.seed + expected_generator = torch.Generator(device="cpu").manual_seed(resolved_input.seed) + assert torch.equal(torch.randn(8), torch.randn(8, generator=expected_generator)) + assert random.random() == random.Random(resolved_input.seed).random() + assert np.random.random() == np.random.RandomState(resolved_input.seed).random_sample() + assert cuda_seeds[-1] == resolved_input.seed + assert torch.backends.cudnn.deterministic + finally: + random.setstate(python_state) + np.random.set_state(numpy_state) diff --git a/tests/test_neopp_kv_lifecycle.py b/tests/test_neopp_kv_lifecycle.py new file mode 100644 index 000000000..c88b3778a --- /dev/null +++ b/tests/test_neopp_kv_lifecycle.py @@ -0,0 +1,46 @@ +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from lightx2v.models.runners.neopp.neopp_runner import NeoppRunner + + +@pytest.mark.parametrize("failure_stage", [None, "encoder", "main"]) +def test_pipeline_releases_injected_kv_on_success_and_failure(failure_stage): + runner = NeoppRunner.__new__(NeoppRunner) + runner._gc_frozen = True + runner.model = SimpleNamespace(transformer_infer=SimpleNamespace(kv_cache={})) + runner.past_key_values_cond = torch.ones(1) + runner.past_key_values_uncond = torch.zeros(1) + tensor_refs = [weakref.ref(runner.past_key_values_cond), weakref.ref(runner.past_key_values_uncond)] + failure = RuntimeError(f"{failure_stage} failed") + + def encode(): + if failure_stage == "encoder": + raise failure + return {"past_key_values_cond": runner.past_key_values_cond, "past_key_values_uncond": runner.past_key_values_uncond} + + def run_main(): + assert runner.inputs["past_key_values_cond"] is runner.past_key_values_cond + assert runner.inputs["past_key_values_uncond"] is runner.past_key_values_uncond + runner.model.transformer_infer.kv_cache.update(runner.inputs) + if failure_stage == "main": + raise failure + return b"encoded image" + + runner.run_input_encoder = encode + runner.run_main = run_main + if failure_stage is None: + assert runner.run_pipeline(SimpleNamespace()) == b"encoded image" + else: + with pytest.raises(RuntimeError) as raised: + runner.run_pipeline(SimpleNamespace()) + assert raised.value is failure + + assert runner.past_key_values_cond is None + assert runner.past_key_values_uncond is None + assert runner.inputs == {} + assert runner.model.transformer_infer.kv_cache == {} + assert all(tensor_ref() is None for tensor_ref in tensor_refs) diff --git a/tests/test_neopp_lightllm_contract.py b/tests/test_neopp_lightllm_contract.py new file mode 100644 index 000000000..262ee8bed --- /dev/null +++ b/tests/test_neopp_lightllm_contract.py @@ -0,0 +1,161 @@ +import base64 +import json +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +from lightx2v.models.runners.neopp.neopp_runner import NeoppRunner +from lightx2v.pipeline import LightX2VPipeline + + +@pytest.fixture +def create_pipeline(tmp_path, monkeypatch): + """Exercise the LightLLM adapter contract with real CPU request/scheduler methods.""" + monkeypatch.setattr("lightx2v.models.runners.neopp.neopp_runner.AI_DEVICE", "cpu") + monkeypatch.setattr("lightx2v.models.schedulers.neopp.scheduler.AI_DEVICE", "cpu") + monkeypatch.setattr("lightx2v.models.runners.base_runner.seed_all", torch.manual_seed) + + def load_model(runner): + infer = SimpleNamespace( + kv_cache={}, + fi_moe_autotune=SimpleNamespace(cache_rebuild_needed=lambda: False, session=lambda **kwargs: nullcontext()), + ) + model = SimpleNamespace(transformer_infer=infer, consumed_inputs=[]) + + def infer_step(inputs): + model.consumed_inputs.append(inputs) + infer.kv_cache["current"] = inputs["past_key_values_cond"] + + model.infer = infer_step + model.set_scheduler = lambda scheduler: setattr(model, "scheduler", scheduler) + runner.model = model + + monkeypatch.setattr(NeoppRunner, "load_model", load_model) + encode_base64 = NeoppRunner.process_images_after_vae_decoder + # LightLLM replaces this output hook to return encoded image bytes. + monkeypatch.setattr(NeoppRunner, "process_images_after_vae_decoder", lambda runner: base64.b64decode(encode_base64(runner))) + + config_path = tmp_path / "neopp.json" + config_path.write_text( + json.dumps( + { + "cpu_offload": True, + "infer_steps": 1, + "patch_size": 16, + "enable_cfg": True, + "seed": 37, + "save_result_for_debug": True, + "llm_config": {"head_dim": 32, "rope_theta": 10000, "rope_theta_hw": 10000}, + } + ), + encoding="utf-8", + ) + + def create(task=None): + task_args = {} if task is None else {"task": task} + pipeline = LightX2VPipeline(model_path=str(tmp_path), model_cls="neopp", support_tasks=["t2i", "i2i"], **task_args) + pipeline.create_generator(config_json=str(config_path)) + pipeline.runner.config.lock() + return pipeline + + with torch.random.fork_rng(devices=[]): + yield create + + +def generate_from_kv(pipeline, *, task, **request): + runner = pipeline.runner + runner.set_inference_params(index_offset_cond=9, index_offset_uncond=3, cfg_scale=3.5, timestep_shift=2.0, output_format="png") + runner.set_kvcache(torch.ones(1, 2, 3, 4), torch.zeros(1, 2, 3, 4)) + return pipeline.generate(task=task, save_result_path="", target_shape=[32, 32], **request) + + +@pytest.mark.parametrize(("task", "expected_task"), [(None, "t2i"), ("i2i", "i2i")]) +def test_support_tasks_selects_startup_task_without_setting_request_default(create_pipeline, task, expected_task): + pipeline = create_pipeline(task) + + assert pipeline.task == (task or "") + assert pipeline.runner.config["task"] == expected_task + assert "support_tasks" not in pipeline.runner.config + if task is None: + with pytest.raises(ValueError, match="task is required"): + pipeline.generate() + + +def test_modify_config_preserves_lightllm_bytes_output_and_kv_injection(create_pipeline): + pipeline = create_pipeline() + pipeline.modify_config({"load_kv_cache_in_pipeline_for_debug": False, "save_result_for_debug": False}) + + result = generate_from_kv(pipeline, task="t2i", seed=0) + runner = pipeline.runner + + assert isinstance(result, bytes) + assert result.startswith(b"\x89PNG\r\n\x1a\n") + assert runner.config.locked + assert runner.input_info.save_result_path == "" + inputs = runner.model.consumed_inputs[-1] + assert torch.equal(inputs["past_key_values_cond"], torch.ones(1, 2, 3, 4)) + assert torch.equal(inputs["past_key_values_uncond"], torch.zeros(1, 2, 3, 4)) + assert runner.model.cfg_scale == 3.5 + assert runner.scheduler.timestep_shift == 2.0 + assert runner.past_key_values_cond is None + assert runner.past_key_values_uncond is None + assert runner.model.transformer_infer.kv_cache == {} + + +def test_same_runner_accepts_t2i_and_i2i_encoded_kv_requests(create_pipeline): + pipeline = create_pipeline() + pipeline.modify_config({"save_result_for_debug": False}) + runner = pipeline.runner + + generate_from_kv(pipeline, task="t2i", seed=0) + first_input = runner.input_info + result = generate_from_kv(pipeline, task="i2i", seed=0) + + assert isinstance(result, bytes) + assert pipeline.runner is runner + assert runner.input_info is not first_input + assert (first_input.task, runner.input_info.task) == ("t2i", "i2i") + assert runner.config["task"] == "t2i" + + generate_from_kv(pipeline, task="t2i", seed=0) + assert runner.input_info.task == "t2i" + + +def test_omitted_seed_restores_code_default_after_explicit_zero(create_pipeline): + pipeline = create_pipeline() + pipeline.modify_config({"save_result_for_debug": False}) + + generate_from_kv(pipeline, task="t2i") + default_noise = pipeline.runner.scheduler.image_prediction.clone() + assert pipeline.runner.input_info.seed == 42 + + generate_from_kv(pipeline, task="t2i", seed=0) + assert pipeline.runner.input_info.seed == 0 + assert not torch.equal(pipeline.runner.scheduler.image_prediction, default_noise) + + generate_from_kv(pipeline, task="t2i") + assert pipeline.runner.input_info.seed == 42 + assert torch.equal(pipeline.runner.scheduler.image_prediction, default_noise) + assert pipeline.runner.config["seed"] == 37 + + +def test_explicit_none_continues_restored_session_rng(create_pipeline): + pipeline = create_pipeline() + pipeline.modify_config({"save_result_for_debug": False}) + + generate_from_kv(pipeline, task="t2i", seed=321) + session_state = torch.get_rng_state() + scheduler = pipeline.runner.scheduler + expected_next_noise = scheduler.noise_scale * torch.randn_like(scheduler.image_prediction) + + generate_from_kv(pipeline, task="t2i", seed=999) + torch.set_rng_state(session_state) + generate_from_kv(pipeline, task="t2i", seed=None) + + assert pipeline.runner.input_info.seed is None + assert torch.equal(scheduler.image_prediction, expected_next_noise) + + generate_from_kv(pipeline, task="t2i") + assert pipeline.runner.input_info.seed == 42 diff --git a/tests/test_optional_output_paths.py b/tests/test_optional_output_paths.py new file mode 100644 index 000000000..8c4a273e2 --- /dev/null +++ b/tests/test_optional_output_paths.py @@ -0,0 +1,376 @@ +import importlib +import importlib.util +import sys +import time +from contextlib import nullcontext +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch + +from lightx2v.models.runners.bagel import sensenova_vision_runner as sensenova +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.hunyuan3d.hunyuan3d_shape_runner import Hunyuan3DShapeRunner +from lightx2v.models.runners.motus import motus_runner as motus +from lightx2v.models.runners.swiftvr import swiftvr_runner as swiftvr +from lightx2v.models.runners.wan import wan_dreamzero_runner as dreamzero +from lightx2v.models.runners.wan.fastwam_runner import FastWAMRunner +from lightx2v.models.runners.wan.wan_lingbot_va_runner import LingbotVARunner +from lightx2v.models.runners.worldmirror import worldmirror_runner as worldmirror +from lightx2v.pipeline import LightX2VPipeline + + +def make_runner(cls, task, **config): + runner = object.__new__(cls) + BaseRunner.__init__(runner, {"task": task, **config}) + runner._gc_frozen = True + runner.progress_callback = None + return runner + + +@pytest.mark.parametrize("request_output", [{}, {"save_result_path": None}, {"save_result_path": ""}, {"save_result_path": "mesh.glb"}]) +def test_hunyuan3d_requires_mesh_output(request_output, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + runner = make_runner(Hunyuan3DShapeRunner, "i23d", enable_pbar=False) + runner.run_input_encoder = lambda: {"image_tensor": torch.zeros(1), "cond": {}, "guidance_scale": 1, "do_classifier_free_guidance": False} + mesh = Mock() + runner.vae_decoder = SimpleNamespace(vae=SimpleNamespace(latent_shape=(1,)), decode_mesh=Mock(return_value=mesh)) + runner.scheduler = SimpleNamespace(prepare=Mock(), infer_steps=1, latents=torch.zeros(1)) + runner.model = SimpleNamespace(transformer_infer=SimpleNamespace(fi_moe_autotune=SimpleNamespace(cache_rebuild_needed=lambda: False, session=lambda **kwargs: nullcontext()))) + runner._run_infer_step = Mock() + runner.end_run = Mock() + pipeline = object.__new__(LightX2VPipeline) + pipeline.runner = runner + pipeline.task = "i23d" + + if not request_output.get("save_result_path"): + with pytest.raises(ValueError, match="save_result_path must be set"): + pipeline.generate(image_path="input.png", **request_output) + mesh.export.assert_not_called() + return + + result = pipeline.generate(image_path="input.png", **request_output) + + runner._run_infer_step.assert_called_once() + runner.vae_decoder.decode_mesh.assert_called_once() + assert result == request_output.get("save_result_path") + assert mesh.export.call_count == (result is not None) + runner.end_run.assert_called_once() + + +@pytest.mark.parametrize("kind", ["image", "video"]) +def test_swiftvr_saves_restored_output(kind, tmp_path, monkeypatch): + runner = make_runner(swiftvr.SwiftVRRunner, "sr", clip_len=4) + runner.init_device = torch.device("cpu") + runner.copy_stream = None + runner.restorer = SimpleNamespace(reset=Mock()) + runner.check_stop = Mock() + runner.read_image_frame = lambda path: (torch.zeros(1, 3, 8, 8), 8, 8) + + class Reader: + def __init__(self, path): + pass + + def __len__(self): + return 9 + + def __getitem__(self, index): + return np.zeros((8, 8, 3)) + + def get_avg_fps(self): + return 24 + + def get_batch(self, indices): + return torch.zeros(len(indices), 8, 8, 3) + + monkeypatch.setattr(swiftvr, "VideoReader", Reader) + + def restore(frames, chunk, clip_latents, height, width, pad_height, pad_width, stage_marks=None): + if stage_marks is not None: + stage_marks.extend([time.perf_counter(), time.perf_counter()]) + return torch.zeros(1, chunk.frame_count, 3, height, width) + + runner.restore_frames = Mock(side_effect=restore) + image_save = Mock() + writer = Mock() + runner.open_video_writer = Mock(return_value=writer) + copy_frames = Mock(wraps=runner.copy_frames_to_cpu) + runner.copy_frames_to_cpu = copy_frames + mux = Mock() + monkeypatch.setattr(swiftvr, "save_to_image", image_save) + monkeypatch.setattr(swiftvr, "mux_audio_from_video", mux) + output_path = str(tmp_path / "output" / ("result.png" if kind == "image" else "result.mp4")) + + input_info = runner.prepare_request({"task": runner.config["task"], f"{kind}_path": "input", "save_result_path": output_path}) + result = runner.run_request(input_info) + + assert result["stats"]["output"] == output_path + assert result["stats"]["frames"] == (1 if kind == "image" else 9) + assert result["images" if kind == "image" else "video"] is None + assert runner.restore_frames.call_count == (1 if kind == "image" else 2) + assert runner.restorer.reset.call_count == 2 + assert image_save.call_count == (kind == "image") + assert runner.open_video_writer.call_count == (kind == "video") + assert mux.call_count == (kind == "video") + if kind == "video": + assert writer.append_data.call_count == 9 + writer.close.assert_called_once() + assert copy_frames.call_count == 2 + assert (tmp_path / "output").is_dir() + + +@pytest.mark.parametrize("kind", ["image", "video"]) +@pytest.mark.parametrize("request_output", [{}, {"save_result_path": None}, {"save_result_path": ""}]) +def test_swiftvr_requires_output_path_before_reading_input(kind, request_output, monkeypatch): + runner = make_runner(swiftvr.SwiftVRRunner, "sr") + runner.read_image_frame = Mock() + reader = Mock() + monkeypatch.setattr(swiftvr, "VideoReader", reader) + input_info = runner.prepare_request({f"{kind}_path": "input", **request_output}) + + with pytest.raises(ValueError, match="requires `save_result_path`"): + runner.run_request(input_info) + + runner.read_image_frame.assert_not_called() + reader.assert_not_called() + + +@pytest.mark.parametrize("source", ["request", "startup"]) +def test_swiftvr_video_rejects_tensor_output(source, monkeypatch): + config = {"return_result_tensor": True} if source == "startup" else {} + runner = make_runner(swiftvr.SwiftVRRunner, "sr", **config) + runner.read_image_frame = Mock() + reader = Mock() + monkeypatch.setattr(swiftvr, "VideoReader", reader) + request = {"video_path": "input", "save_result_path": "output"} + if source == "request": + request["return_result_tensor"] = True + + with pytest.raises(ValueError, match="return_result_tensor"): + input_info = runner.prepare_request(request) + runner.run_request(input_info) + + runner.read_image_frame.assert_not_called() + reader.assert_not_called() + + +@pytest.mark.parametrize("request_output", [{}, {"save_result_path": "unused.png"}]) +def test_swiftvr_image_returns_tensor_without_saving(request_output, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + runner = make_runner(swiftvr.SwiftVRRunner, "sr", clip_len=4) + runner.restorer = SimpleNamespace(reset=Mock()) + runner.read_image_frame = Mock(return_value=(torch.zeros(1, 3, 8, 8), 8, 8)) + restored = torch.linspace(0, 1, 3 * 16 * 16, dtype=torch.float16).reshape(1, 1, 3, 16, 16) + runner.restore_frames = Mock(return_value=restored) + image_save = Mock() + monkeypatch.setattr(swiftvr, "save_to_image", image_save) + input_info = runner.prepare_request({"image_path": "input", "return_result_tensor": True, **request_output}) + + result = runner.run_request(input_info) + + torch.testing.assert_close(result["images"], restored[0].permute(0, 2, 3, 1).float()) + assert result["images"].device.type == "cpu" + assert result["stats"]["output"] is None + runner.restore_frames.assert_called_once() + image_save.assert_not_called() + assert not list(tmp_path.iterdir()) + + +@pytest.mark.parametrize("paths", [{}, {"save_result_path": "output"}, {"strict_output_path": "exact"}]) +@pytest.mark.parametrize("rank", [0, 1]) +def test_worldmirror_saves_only_to_explicit_paths(paths, rank, monkeypatch): + runner = make_runner(worldmirror.WorldMirrorRunner, "recon", output_path="ignored-json-default", apply_sky_mask=False, apply_edge_mask=False, log_time=False) + runner.rank = rank + runner.is_distributed = True + runner._run_inference = Mock(return_value=({}, torch.zeros(1, 1, 3, 8, 8), 0.1)) + monkeypatch.setattr(worldmirror, "prepare_input", lambda *args, **kwargs: (["input.png"], "scene")) + monkeypatch.setattr(worldmirror, "compute_adaptive_target_size", lambda paths, size: size) + save = Mock() + monkeypatch.setattr(worldmirror, "save_results", save) + barrier = Mock() + monkeypatch.setattr(worldmirror.dist, "barrier", barrier) + + input_info = runner.prepare_request({"task": runner.config["task"], "input_path": "input.png", **paths}) + result = runner.run_request(input_info) + + runner._run_inference.assert_called_once() + assert save.call_count == (bool(paths) and rank == 0) + assert (result["output_dir"] is not None) == bool(paths) + if "strict_output_path" in paths: + assert result["output_dir"] == "exact" + barrier.assert_called_once() + + +@pytest.mark.parametrize("runner_cls", [dreamzero.WanDreamZeroRunner, LingbotVARunner]) +@pytest.mark.parametrize("output", [None, "video.mp4"]) +@pytest.mark.parametrize("return_tensor", [False, True]) +def test_action_video_output(runner_cls, output, return_tensor, tmp_path, monkeypatch): + runner = make_runner(runner_cls, "i2va") + runner.input_info = runner.prepare_request({"task": runner.config["task"], "save_result_path": str(tmp_path / output) if output else None, "return_result_tensor": return_tensor}) + runner.gen_video = torch.zeros(1) + runner.pred_action = torch.ones(1, 2, 3) + save = Mock() + monkeypatch.setattr(importlib.import_module(runner_cls.__module__), "save_to_video", save) + + if output is None: + with pytest.raises(ValueError, match="requires save_result_path"): + runner.process_images_after_vae_decoder() + save.assert_not_called() + assert not list(tmp_path.iterdir()) + return + + result = runner.process_images_after_vae_decoder() + + save.assert_called_once() + assert (tmp_path / "video.actions.npy").exists() + assert (result["video"] is runner.gen_video) == return_tensor + if return_tensor: + assert result["actions"] is runner.pred_action + + +@pytest.mark.parametrize("output", [None, "video.mp4"]) +@pytest.mark.parametrize("action_output", [None, "actions.json"]) +def test_motus_requires_video_output(output, action_output, tmp_path, monkeypatch): + runner = make_runner(motus.MotusRunner, "i2v") + request = {"image_path": "input.png", "state_path": "state.json", "prompt": "move", "save_result_path": str(tmp_path / output) if output else None} + if action_output: + request["save_action_path"] = str(tmp_path / action_output) + runner._load_state_value = lambda path: [0] + runner.run_input_encoder = lambda: {"motus_state": torch.zeros(1), "image_encoder_output": None} + runner.model = SimpleNamespace(prepare_runtime_inputs=lambda inputs, **kwargs: inputs, action_chunk_size=1, action_dim=1, postprocess_actions=lambda: torch.ones(1, 1)) + runner.scheduler = SimpleNamespace(prepare=Mock(), infer_steps=0, video_latents=torch.zeros(1, 1)) + runner.run_vae_decoder = Mock(return_value=torch.zeros(1)) + runner.end_run = Mock() + monkeypatch.setattr(motus, "wan_vae_to_comfy", lambda video: video) + save = Mock() + monkeypatch.setattr(motus, "save_to_video", save) + + input_info = runner.prepare_request({"task": runner.config["task"], **request}) + if output is None: + with pytest.raises(ValueError, match="requires `save_result_path`"): + runner.run_request(input_info) + save.assert_not_called() + runner.run_vae_decoder.assert_not_called() + assert not list(tmp_path.iterdir()) + return + + result = runner.run_request(input_info) + + save.assert_called_once() + expected_action = action_output or "video.actions.json" + assert [path.name for path in tmp_path.iterdir()] == [expected_action] + assert torch.equal(result["actions"], torch.ones(1, 1)) + runner.end_run.assert_called_once() + + +@pytest.mark.parametrize("paths", [{}, {"save_action_path": "actions.npy"}, {"save_result_path": "video.mp4"}]) +@pytest.mark.parametrize("return_tensor", [False, True]) +def test_fastwam_requires_action_output(paths, return_tensor, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + runner = make_runner(FastWAMRunner, "i2va") + runner._load_image_pair = lambda: (None, None) + runner._load_state = lambda: [0] * 8 + actions = np.ones((2, 8), dtype=np.float32) + runner.policy = SimpleNamespace(predict_action_chunk=Mock(return_value=actions)) + + input_info = runner.prepare_request({"task": runner.config["task"], "prompt": "move", "return_result_tensor": return_tensor, **paths}) + if not paths: + with pytest.raises(ValueError, match="requires `save_action_path` or `save_result_path`"): + runner.run_request(input_info) + assert not list(tmp_path.iterdir()) + return + + result = runner.run_request(input_info) + + assert result["actions"] is (actions if return_tensor else None) + expected_path = paths.get("save_action_path", "video.actions.npy") + np.testing.assert_array_equal(np.load(tmp_path / expected_path), actions) + + +@pytest.mark.parametrize("paths", [{}, {"save_result_path": "points.npy"}, {"raw_output_path": "points.npy", "glb_output_path": "scene.glb"}, {"glb_output_path": "scene.glb"}]) +def test_sensenova_reconstruction_optional_files(paths, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + runner = make_runner(sensenova.SenseNovaVisionRunner, "omni_vision_task", postprocess_predictions=True) + info = runner.prepare_request({"omni_vision_subtask": "recon3d", **paths}) + scene = Mock() + monkeypatch.setattr(sensenova, "load_official_postprocess", lambda source: lambda *args, **kwargs: scene) + pointmaps = np.zeros((1, 2, 2, 3)) + + points, result_scene, raw_path, glb_path = runner._postprocess_recon3d(pointmaps, SimpleNamespace(preprocessed_images=[None]), info, 1) + + assert result_scene is scene + assert np.array_equal(points, pointmaps) + has_raw = "save_result_path" in paths or "raw_output_path" in paths + assert (tmp_path / "points.npy").exists() == has_raw + assert raw_path == ("points.npy" if has_raw else None) + assert bool(glb_path) == bool(paths) + assert scene.export.call_count == bool(paths) + if not paths: + assert list(tmp_path.iterdir()) == [] + + +@pytest.mark.parametrize("request_data", [{}, {"negative_prompt": ""}, {"negative_prompt": "user negative"}]) +def test_dreamzero_uses_request_negative_prompt(request_data): + runner = make_runner(dreamzero.WanDreamZeroRunner, "i2va", enable_cfg=True) + runner.input_info = runner.prepare_request({"task": runner.config["task"], "prompt": "move", **request_data}) + consumed = [] + runner.run_text_encoder = lambda info: consumed.append(info.negative_prompt) + + runner._run_input_encoder_local_i2va() + + assert consumed == [request_data.get("negative_prompt", "")] + assert runner.input_info.prompt == "move" + + +def test_dreamzero_action_path_still_requires_video(tmp_path, monkeypatch): + runner = make_runner(dreamzero.WanDreamZeroRunner, "i2va") + runner.input_info = runner.prepare_request({"task": runner.config["task"], "save_action_path": str(tmp_path / "actions.npy")}) + runner.gen_video = torch.zeros(1) + runner.pred_action = torch.ones(1, 2) + save = Mock() + monkeypatch.setattr(dreamzero, "save_to_video", save) + + with pytest.raises(ValueError, match="requires save_result_path"): + runner.process_images_after_vae_decoder() + + save.assert_not_called() + assert not list(tmp_path.iterdir()) + + +@pytest.mark.parametrize("action_path", [None, "custom.actions.npy", "absolute"]) +def test_dreamzero_action_output_path(action_path, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + runner = make_runner(dreamzero.WanDreamZeroRunner, "i2va") + if action_path == "absolute": + action_path = str(tmp_path / "absolute.actions.npy") + runner.input_info = runner.prepare_request({"save_result_path": "output/video.mp4", "save_action_path": action_path}) + runner.gen_video = torch.zeros(1) + runner.pred_action = torch.ones(1, 2) + save = Mock() + monkeypatch.setattr(dreamzero, "save_to_video", save) + + runner.process_images_after_vae_decoder() + + save.assert_called_once() + expected_path = tmp_path / "output" / (action_path or "video.actions.npy") + np.testing.assert_array_equal(np.load(expected_path), runner.pred_action.numpy()) + + +@pytest.mark.parametrize("output_path", [None, "result.mp4"]) +def test_seko_client_downloads_only_saved_results(output_path, monkeypatch): + script = Path(__file__).resolve().parents[1] / "scripts/server/post_seko_talk_ar.py" + spec = importlib.util.spec_from_file_location("seko_client", script) + client = importlib.util.module_from_spec(spec) + spec.loader.exec_module(client) + monkeypatch.setattr(sys, "argv", [str(script)]) + monkeypatch.setattr(client, "submit_task", lambda args: "task1") + monkeypatch.setattr(client, "wait_task_done", lambda *args: {"status": "completed", "save_result_path": output_path}) + download = Mock() + monkeypatch.setattr(client, "download_result", download) + + client.main() + + assert download.call_count == (output_path is not None) diff --git a/tests/test_request_config.py b/tests/test_request_config.py new file mode 100644 index 000000000..5dd79c1ae --- /dev/null +++ b/tests/test_request_config.py @@ -0,0 +1,1219 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from loguru import logger + +from lightx2v import infer +from lightx2v.disagg.examples import infer as disagg_infer +from lightx2v.models.networks.worldplay.pose_utils import pose_to_input +from lightx2v.models.runners import runner_factory +from lightx2v.models.runners.bagel.sensenova_vision_runner import SenseNovaVisionRunner +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.ernie_image.ernie_image_runner import ErnieImageRunner +from lightx2v.models.runners.flux2.flux2_runner import Flux2Runner +from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner import HunyuanVideo15Runner +from lightx2v.models.runners.ltx2.ltx2_runner import LTX2Runner +from lightx2v.models.runners.ltx2.ltx25_runner import LTX25Runner +from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner +from lightx2v.models.runners.neopp.neopp_runner import NeoppRunner +from lightx2v.models.runners.qwen_image.qwen_image_runner import QwenImageRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, PROMPT_FIELDS, VIDEO_OUTPUT_FIELDS +from lightx2v.models.runners.seedvr.seedvr_runner import SeedVRRunner +from lightx2v.models.runners.wan.wan_animate_runner import WanAnimateRunner +from lightx2v.models.runners.wan.wan_audio_runner import WanAudioARRunner +from lightx2v.models.runners.wan.wan_dancer_runner import WanDancerRunner +from lightx2v.models.runners.wan.wan_runner import WanRunner +from lightx2v.models.schedulers.wan.scheduler import WanScheduler +from lightx2v.pipeline import LightX2VPipeline +from lightx2v.server.api import openai_images +from lightx2v.server.api.tasks import image as image_api +from lightx2v.server.schema import ImageTaskRequest, VideoTaskRequest +from lightx2v.server.services.generation.base import BaseGenerationService +from lightx2v.server.services.generation.image import ImageGenerationService +from lightx2v.server.services.inference.worker import TorchrunInferenceWorker +from lightx2v.utils import set_config as set_config_utils +from lightx2v.utils.input_info import ( + FL2AVInputInfo, + I2AVInputInfo, + I2IInputInfo, + I2VAInputInfo, + L2AVInputInfo, + NeoppInputInfo, + Ref2AVInputInfo, + SRInputInfo, + T2AVInputInfo, + T2IInputInfo, +) +from lightx2v.utils.lockable_dict import LockableDict +from lightx2v.utils.set_config import build_cli_inputs + + +class RecordingScheduler: + def __init__(self, config=None): + self.config = config + self.infer_steps = config["infer_steps"] if config else None + + +class RecordingRunner(BaseRunner): + supported_request_fields_by_task = { + "t2v": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | VIDEO_OUTPUT_FIELDS, + "t2av": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | VIDEO_OUTPUT_FIELDS, + "t2i": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | {"target_shape"}, + "sr": COMMON_REQUEST_FIELDS | {"image_path", "video_path", "sr_ratio", "target_shape"}, + "animate": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | {"src_face_path", "src_pose_path"}, + } + + def __init__(self, task="t2v"): + config = { + "task": task, + "target_video_length": 81, + "target_height": 720, + "target_width": 1280, + } + super().__init__(config) + self.requests = [] + self.pipeline_inputs = [] + self.lifecycle = [] + + def prepare_request(self, request_data): + input_info = super().prepare_request(request_data) + self.lifecycle.append("prepare") + return input_info + + def init_modules(self): + self.lifecycle.append("init") + + def create_input_info(self, request_data): + self.requests.append(request_data) + return super().create_input_info(request_data) + + def run_pipeline(self, input_info): + self.pipeline_inputs.append(input_info) + self.lifecycle.append("run") + return input_info + + +class MultitaskRecordingRunner(RecordingRunner): + supported_request_fields_by_task = { + "t2i": COMMON_REQUEST_FIELDS | {"aspect_ratio", "prompt", "target_shape"}, + "i2i": COMMON_REQUEST_FIELDS | {"image_path", "prompt", "target_shape"}, + } + + def get_supported_tasks(self): + return ("t2i", "i2i") + + +class FileOnlyRecordingRunner(RecordingRunner): + supported_request_fields_by_task = { + "i23d": COMMON_REQUEST_FIELDS - {"return_result_tensor"}, + } + + +def make_minimax_h3_runner(task): + runner = object.__new__(MiniMaxH3Runner) + runner.config = LockableDict( + { + "model_cls": "minimax_h3", + "task": task, + "target_video_length": 124, + "target_height": 544, + "target_width": 960, + } + ) + runner.loaded_transformer_partition = "transformer_ref" if task == "ref2av" else "transformer" + runner.supported_tasks = runner.get_supported_tasks() + return runner + + +class RequestConfigRunner(DefaultRunner): + supported_request_fields_by_task = { + "t2v": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | VIDEO_OUTPUT_FIELDS, + } + + def set_init_device(self): + pass + + def init_scheduler(self): + self.scheduler = RecordingScheduler(self.config) + + +class FixedShapeWanRunner(WanRunner): + supported_request_fields_by_task = { + "t2v": WanRunner.supported_request_fields_by_task["t2v"] - {"target_shape"}, + } + + def set_init_device(self): + pass + + def init_scheduler(self): + pass + + +class RecordingInferenceService: + def __init__(self): + self.requests = [] + self.worker = SimpleNamespace(runner=SimpleNamespace(config={})) + + async def submit_task_async(self, task_data): + self.requests.append(task_data) + return {"status": "success", "save_result_path": task_data["save_result_path"]} + + +class RecordingGenerationService(BaseGenerationService): + def get_output_extension(self): + return ".mp4" + + def get_task_type(self): + return "t2v" + + +def create_runner(): + config = LockableDict( + { + "infer_steps": 40, + "target_video_length": 81, + "target_height": 720, + "target_width": 1280, + "resize_mode": "adaptive", + "video_frame_interpolation": {"algo": "rife", "target_fps": 24}, + "task": "t2v", + } + ) + runner = RequestConfigRunner(config) + runner.config.lock() + return runner + + +def test_request_values_are_stored_in_input_info_only(): + runner = create_runner() + + input_info = runner.create_input_info( + { + "task": runner.config["task"], + "task_id": "task-1", + "prompt": "ignored by runner config", + "infer_steps": 20, + "target_video_length": 49, + "target_shape": [480, 832], + "resize_mode": "fixed_shape", + } + ) + + assert runner.config["infer_steps"] == 40 + assert runner.config["target_video_length"] == 81 + assert runner.config["target_height"] == 720 + assert runner.config["target_width"] == 1280 + assert runner.config["resize_mode"] == "adaptive" + assert "task_id" not in runner.config + assert "prompt" not in runner.config + assert input_info.target_video_length == 49 + assert input_info.target_shape == [480, 832] + assert input_info.prompt == "ignored by runner config" + assert input_info.negative_prompt == "" + assert not hasattr(input_info, "infer_steps") + assert not hasattr(input_info, "resize_mode") + + +def test_each_request_gets_startup_defaults_without_state_leakage(): + runner = create_runner() + first = runner.create_input_info({"task": runner.config["task"], "infer_steps": 20, "target_video_length": 49, "target_shape": [480, 832]}) + second = runner.create_input_info({"task": runner.config["task"], "prompt": "second request"}) + + assert runner.config["infer_steps"] == 40 + assert runner.config["target_video_length"] == 81 + assert runner.config["target_height"] == 720 + assert runner.config["target_width"] == 1280 + assert runner.config["resize_mode"] == "adaptive" + assert runner.scheduler.infer_steps == 40 + assert first.target_video_length == 49 + assert first.target_shape == [480, 832] + assert second.target_video_length == 81 + assert second.target_shape == [720, 1280] + assert first is not second + + +def test_explicit_aspect_ratio_takes_precedence_over_startup_dimensions(): + runner = MultitaskRecordingRunner(task="t2i") + + input_info = runner.prepare_request({"task": runner.config["task"], "aspect_ratio": "16:9"}) + + assert input_info.aspect_ratio == "16:9" + assert input_info.target_shape == [] + + +def test_ernie_uses_request_aspect_ratio_over_startup_dimensions(): + runner = object.__new__(ErnieImageRunner) + runner.config = { + "task": "t2i", + "target_height": 480, + "target_width": 832, + "vae_scale_factor": 16, + } + runner.supported_tasks = ("t2i",) + runner.resolution = 1024 + runner.input_info = runner.prepare_request({"task": runner.config["task"], "aspect_ratio": "16:9"}) + + runner.set_latent_shape() + + assert runner.input_info.target_shape == [768, 1360] + + +def test_qwen_uses_request_aspect_ratio_over_startup_dimensions(): + runner = object.__new__(QwenImageRunner) + runner.config = { + "task": "t2i", + "target_height": 480, + "target_width": 832, + } + runner.supported_tasks = ("t2i",) + runner.input_info = runner.prepare_request({"task": runner.config["task"], "aspect_ratio": "16:9"}) + + assert runner.get_custom_shape() == (1664, 928) + + +def test_video_frame_interpolation_is_startup_only(): + runner = create_runner() + + input_info = runner.create_input_info({"task": runner.config["task"], "video_frame_interpolation": {"algo": "rife", "target_fps": 30}}) + + assert runner.config["video_frame_interpolation"]["target_fps"] == 24 + assert not hasattr(input_info, "video_frame_interpolation") + with pytest.raises(TypeError, match="Dictionary is locked"): + runner.config["video_frame_interpolation"]["target_fps"] = 60 + + +def test_runner_rejects_startup_and_internal_fields(): + runner = create_runner() + + with pytest.raises(ValueError, match="infer_steps, latent_shape, resize_mode"): + runner.prepare_request( + { + "task": runner.config["task"], + "infer_steps": 20, + "latent_shape": [1, 2, 3], + "resize_mode": "fixed_shape", + } + ) + + +@pytest.mark.parametrize( + ("runner_cls", "model_cls", "task"), + [ + (WanRunner, "wan2.1", "omni_vision_task"), + (WanRunner, "wan2.1", "i2i"), + (SenseNovaVisionRunner, "sensenova_vision", "t2v"), + (NeoppRunner, "neopp", "t2v"), + (MiniMaxH3Runner, "minimax_h3", "t2i"), + ], +) +def test_runner_rejects_unsupported_startup_task_before_device_init(monkeypatch, runner_cls, model_cls, task): + device_init = Mock() + monkeypatch.setattr(runner_cls, "set_init_device", device_init) + + with pytest.raises(ValueError, match=f"{runner_cls.__name__} does not support task"): + runner_cls({"model_cls": model_cls, "task": task}) + + device_init.assert_not_called() + + +@pytest.mark.parametrize("task", [None, ""]) +@pytest.mark.parametrize("runner_cls", [RecordingRunner, NeoppRunner, MiniMaxH3Runner]) +def test_runner_requires_startup_task_even_with_task_selection_override(runner_cls, task): + with pytest.raises(ValueError, match="task must be set"): + if runner_cls is RecordingRunner: + runner_cls(task=task) + else: + runner_cls({"task": task}) + + +def test_task_selection_override_does_not_bypass_startup_validation(): + with pytest.raises(ValueError, match="MultitaskRecordingRunner does not support task 't2v'"): + MultitaskRecordingRunner(task="t2v") + + +@pytest.mark.parametrize("task", ["t2av", "i2av", "l2av", "fl2av", "ref2av"]) +def test_minimax_initializes_task_group_and_loaded_partition(monkeypatch, task): + monkeypatch.setattr(MiniMaxH3Runner, "set_init_device", lambda self: None) + monkeypatch.setattr(MiniMaxH3Runner, "init_scheduler", lambda self: None) + + runner = MiniMaxH3Runner({"task": task}) + + assert runner.supported_tasks == (("ref2av",) if task == "ref2av" else ("t2av", "i2av", "l2av", "fl2av")) + assert runner.loaded_transformer_partition == ("transformer_ref" if task == "ref2av" else "transformer") + + +def test_request_schema_has_no_generation_defaults(): + request = VideoTaskRequest() + + assert request.target_video_length is None + assert "infer_steps" not in VideoTaskRequest.model_fields + assert "target_fps" not in VideoTaskRequest.model_fields + assert "resize_mode" not in VideoTaskRequest.model_fields + assert "target_video_length" not in request.model_fields_set + + explicit_request = VideoTaskRequest(num_frames=49) + assert explicit_request.target_video_length == 49 + assert not hasattr(explicit_request, "infer_steps") + assert not hasattr(explicit_request, "target_fps") + assert not hasattr(explicit_request, "resize_mode") + with pytest.raises(ValueError, match="extra_forbidden"): + VideoTaskRequest(infer_steps=20, target_fps=30, resize_mode="fixed_shape") + + +def test_generation_service_forwards_only_explicit_generation_fields(tmp_path): + inference_service = RecordingInferenceService() + file_service = SimpleNamespace(get_output_path=lambda path: tmp_path / path) + service = RecordingGenerationService(file_service, inference_service) + message = VideoTaskRequest(prompt="request", num_frames=49, target_shape=[480, 832]) + + asyncio.run(service.generate_with_stop_event(message, asyncio.Event())) + + task_data = inference_service.requests[0] + assert task_data["target_video_length"] == 49 + assert task_data["target_shape"] == [480, 832] + assert "infer_steps" not in task_data + assert "resize_mode" not in task_data + assert "target_fps" not in task_data + + +def test_image_generation_does_not_expect_video_only_fields(tmp_path): + inference_service = RecordingInferenceService() + file_service = SimpleNamespace(get_output_path=lambda path: tmp_path / path) + service = RecordingGenerationService(file_service, inference_service) + + asyncio.run(service.generate_with_stop_event(ImageTaskRequest(prompt="image"), asyncio.Event())) + + assert inference_service.requests[0]["prompt"] == "image" + + +def test_image_generation_omits_disabled_memory_result(tmp_path): + inference_service = RecordingInferenceService() + file_service = SimpleNamespace(get_output_path=lambda path: tmp_path / path) + service = ImageGenerationService(file_service, inference_service) + + asyncio.run(service.generate_with_stop_event(ImageTaskRequest(prompt="image"), asyncio.Event())) + + assert "return_result_tensor" not in inference_service.requests[0] + + +def test_image_form_omits_unsubmitted_optional_fields(tmp_path, monkeypatch): + messages = [] + services = SimpleNamespace(file_service=SimpleNamespace(input_image_dir=tmp_path)) + monkeypatch.setattr(image_api, "get_services", lambda: services) + monkeypatch.setattr(image_api.task_manager, "create_task", lambda message: messages.append(message) or message.task_id) + + asyncio.run( + image_api.create_image_task_form( + request=SimpleNamespace(form=AsyncMock(return_value={})), + image_file=None, + prompt="", + save_result_path="", + negative_prompt="", + seed=42, + aspect_ratio=None, + ) + ) + + assert {"aspect_ratio", "image_path", "negative_prompt", "prompt"}.isdisjoint(messages[0].model_fields_set) + + +def test_openai_image_request_omits_empty_optional_fields(): + message = openai_images._build_image_task_request(prompt="image", task="t2i") + + assert {"image_mask_path", "image_path", "negative_prompt"}.isdisjoint(message.model_fields_set) + + +def test_server_requests_restore_startup_config_and_input_state(): + runner = create_runner() + requests = [] + + def run_pipeline(input_info): + requests.append( + { + "infer_steps": runner.config["infer_steps"], + "target_video_length": runner.config["target_video_length"], + "target_height": runner.config["target_height"], + "target_width": runner.config["target_width"], + "resize_mode": runner.config["resize_mode"], + "target_fps": runner.config["video_frame_interpolation"]["target_fps"], + "input_target_shape": input_info.target_shape, + "input_target_video_length": input_info.target_video_length, + "negative_prompt": input_info.negative_prompt, + "input_info": input_info, + } + ) + + runner.run_pipeline = run_pipeline + worker = TorchrunInferenceWorker() + worker.rank = 0 + worker.world_size = 1 + worker.runner = runner + + asyncio.run( + worker.process_request( + { + "task_id": "first", + "target_video_length": 49, + "target_shape": [480, 832], + } + ) + ) + asyncio.run(worker.process_request({"task_id": "second"})) + + assert requests[0]["infer_steps"] == 40 + assert requests[0]["target_video_length"] == 81 + assert requests[0]["target_height"] == 720 + assert requests[0]["target_width"] == 1280 + assert requests[0]["resize_mode"] == "adaptive" + assert requests[0]["target_fps"] == 24 + assert requests[0]["input_target_shape"] == [480, 832] + assert requests[0]["input_target_video_length"] == 49 + assert requests[1]["infer_steps"] == 40 + assert requests[1]["target_video_length"] == 81 + assert requests[1]["target_height"] == 720 + assert requests[1]["target_width"] == 1280 + assert requests[1]["resize_mode"] == "adaptive" + assert requests[1]["target_fps"] == 24 + assert requests[1]["input_target_shape"] == [720, 1280] + assert requests[1]["input_target_video_length"] == 81 + assert requests[0]["negative_prompt"] == "" + assert requests[1]["negative_prompt"] == "" + assert requests[0]["input_info"] is not requests[1]["input_info"] + + +def test_wan_i2v_reuse_key_includes_target_shape(): + runner = object.__new__(WanRunner) + runner.config = { + "task": "i2v", + "target_video_length": 81, + "target_height": 480, + "target_width": 832, + "resize_mode": None, + } + runner.input_info = SimpleNamespace(prompt="same", negative_prompt="", image_path="same.png", target_video_length=81, target_shape=[480, 832]) + first_key = runner.reuse_key() + + runner.input_info.target_shape = [720, 1280] + + assert runner.reuse_key() != first_key + + +def test_hunyuan_rejects_request_shape_and_uses_request_length_for_vsr(): + runner = object.__new__(HunyuanVideo15Runner) + runner.config = LockableDict( + { + "infer_steps": 50, + "target_video_length": 121, + "target_height": 480, + "target_width": 832, + "task": "i2v", + "vae_stride": (4, 16, 16), + } + ) + runner.scheduler = RecordingScheduler(runner.config) + runner.sr_version = "720p_sr_distilled" + runner.config_sr = { + "target_video_length": 121, + "vae_stride": (4, 16, 16), + "video_super_resolution": {"base_resolution": "480p"}, + } + runner.lq_latents_shape = (32, 31, 30, 40) + runner.build_bucket_map = lambda **kwargs: lambda size: (1280, 720) + runner.supported_tasks = ("i2v",) + + with pytest.raises(ValueError, match="target_shape"): + runner.prepare_request({"task": runner.config["task"], "target_shape": [720, 1280]}) + + runner.input_info = runner.create_input_info({"task": runner.config["task"], "target_video_length": 49}) + + assert runner.config["target_video_length"] == 121 + assert runner.get_sr_latent_shape_with_target_hw()[1] == 13 + + +def test_cli_keeps_json_defaults_separate_from_request_values(tmp_path): + model_path = tmp_path / "model" + model_path.mkdir() + (model_path / "config.json").write_text('{"hidden_size": 256}', encoding="utf-8") + config_path = tmp_path / "config.json" + config_path.write_text( + '{"aspect_ratio": "16:9", "infer_steps": 40, "target_video_length": 81, "target_height": 720, "target_width": 1280}', + encoding="utf-8", + ) + args = SimpleNamespace( + config_json=str(config_path), + model_cls="test_model", + model_path=str(model_path), + task="t2v", + target_video_length=49, + ) + + startup_config, request_data = build_cli_inputs(args) + + assert startup_config["hidden_size"] == 256 + assert startup_config["infer_steps"] == 40 + assert startup_config["target_video_length"] == 81 + assert startup_config["target_height"] == 720 + assert startup_config["target_width"] == 1280 + assert startup_config["aspect_ratio"] == "16:9" + assert request_data["target_video_length"] == 49 + + +def test_cli_uses_runner_as_request_field_authority(tmp_path): + model_path = tmp_path / "model" + model_path.mkdir() + config_path = tmp_path / "config.json" + config_path.write_text("{}", encoding="utf-8") + args = SimpleNamespace( + config_json=str(config_path), + model_cls="test_model", + model_path=str(model_path), + task="t2v", + model_specific_input="value", + ) + + startup_config, request_data = build_cli_inputs(args) + + assert "model_specific_input" not in startup_config + assert request_data["model_specific_input"] == "value" + + +@pytest.mark.parametrize( + ("cli", "frame_args", "expected_frames"), + [(infer, ["--num_frames", "49"], 49), (disagg_infer, [], 81)], + ids=["cli", "disagg_cli"], +) +def test_cli_passes_explicit_request_fields_to_runner(tmp_path, monkeypatch, cli, frame_args, expected_frames): + model_path = tmp_path / "model" + model_path.mkdir() + config_path = tmp_path / "config.json" + config_path.write_text( + '{"infer_steps": 40, "target_video_length": 81, "target_height": 720, "target_width": 1280}', + encoding="utf-8", + ) + runner = RecordingRunner() + monkeypatch.setattr( + "sys.argv", + [ + "lightx2v.infer", + "--model_cls", + "wan2.1", + "--task", + "t2v", + "--model_path", + str(model_path), + "--config_json", + str(config_path), + *frame_args, + "--target_shape", + "480", + "832", + ], + ) + monkeypatch.setattr(runner_factory, "RUNNER_REGISTER", {"wan2.1": lambda config: runner}) + monkeypatch.setattr(cli, "print_config", lambda config, **kwargs: None) + monkeypatch.setattr(cli, "print_request", lambda input_info, supported_request_fields: None) + monkeypatch.setattr(cli, "validate_config_paths", lambda config: None) + monkeypatch.setattr("lightx2v.models.runners.base_runner.seed_all", lambda seed: None) + + cli.main() + + assert runner.pipeline_inputs[0].target_video_length == expected_frames + assert runner.pipeline_inputs[0].target_shape == [480, 832] + assert runner.lifecycle == ["init", "prepare", "run"] + + +def test_print_request_logs_readable_effective_fields(monkeypatch): + input_info = I2VAInputInfo( + task="i2va", + prompt="move forward", + seed=7, + target_shape=[480, 832], + policy_image=object(), + policy_state=object(), + ) + logged = {} + + def record_config(config, title="config"): + logged["config"] = config + logged["title"] = title + + monkeypatch.setattr(set_config_utils, "print_config", record_config) + + set_config_utils.print_request(input_info, {"task", "policy_image", "policy_state", "prompt", "seed", "target_shape"}) + + assert logged == { + "config": { + "task": "i2va", + "seed": 7, + "prompt": "move forward", + "target_shape": [480, 832], + }, + "title": "Effective request", + } + + +@pytest.mark.parametrize(("requested_frames", "aligned_frames"), [(80, 81), (82, 81)]) +def test_wan_aligns_request_frames_with_warning(requested_frames, aligned_frames): + runner = object.__new__(WanRunner) + runner.config = {"task": "t2v", "vae_stride": (4, 8, 8)} + runner.supported_tasks = ("t2v",) + warnings = [] + sink_id = logger.add(lambda message: warnings.append(str(message)), level="WARNING") + + try: + input_info = runner.prepare_request({"task": runner.config["task"], "target_video_length": requested_frames}) + finally: + logger.remove(sink_id) + + assert input_info.target_video_length == aligned_frames + assert any(f"using {aligned_frames} instead of {requested_frames}" in message for message in warnings) + + +def test_wan_keeps_valid_request_frames_without_warning(): + runner = object.__new__(WanRunner) + runner.config = {"task": "t2v", "vae_stride": (4, 8, 8)} + runner.supported_tasks = ("t2v",) + warnings = [] + sink_id = logger.add(lambda message: warnings.append(str(message)), level="WARNING") + + try: + input_info = runner.prepare_request({"task": runner.config["task"], "target_video_length": 81}) + finally: + logger.remove(sink_id) + + assert input_info.target_video_length == 81 + assert not warnings + + +def test_wan_does_not_normalize_fixed_frame_count(): + runner = object.__new__(WanRunner) + runner.config = { + "self_attn_1_type": "radial_attn", + "task": "t2v", + "target_video_length": 80, + "vae_stride": (4, 8, 8), + } + runner.supported_tasks = ("t2v",) + + input_info = runner.prepare_request({"task": runner.config["task"]}) + + assert input_info.target_video_length == 80 + + +def test_fixed_frame_attention_preserves_runner_request_restrictions(): + runner = FixedShapeWanRunner( + { + "enable_cfg": False, + "model_cls": "test_wan", + "self_attn_1_type": "radial_attn", + "task": "t2v", + } + ) + + assert "target_shape" not in runner.get_supported_request_fields("t2v") + assert "target_video_length" not in runner.get_supported_request_fields("t2v") + + +@pytest.mark.parametrize( + ("runner_cls", "config", "task", "field", "supported"), + [ + (WanRunner, {"enable_cfg": True}, "t2v", "negative_prompt", True), + (WanRunner, {"enable_cfg": False}, "t2v", "negative_prompt", False), + (ErnieImageRunner, {"enable_cfg": True}, "t2i", "negative_prompt", True), + (ErnieImageRunner, {"enable_cfg": False}, "t2i", "negative_prompt", False), + (LTX2Runner, {"enable_cfg": True}, "t2av", "negative_prompt", True), + (LTX2Runner, {"enable_cfg": False}, "t2av", "negative_prompt", False), + (QwenImageRunner, {"enable_cfg": True}, "t2i", "negative_prompt", True), + (QwenImageRunner, {"enable_cfg": False}, "t2i", "negative_prompt", False), + (HunyuanVideo15Runner, {"enable_cfg": False}, "t2v", "negative_prompt", False), + (HunyuanVideo15Runner, {"enable_cfg": True}, "t2v", "negative_prompt", True), + (HunyuanVideo15Runner, {"enable_cfg": False, "video_super_resolution": {"enable_cfg": True}}, "t2v", "negative_prompt", True), + (Flux2Runner, {"enable_cfg": True}, "t2i", "negative_prompt", False), + (Flux2Runner, {"inpaint_mask_enabled": True}, "i2i", "inpaint_blur_sigma", True), + (Flux2Runner, {"inpaint_mask_enabled": False}, "i2i", "inpaint_blur_sigma", False), + (SeedVRRunner, {"seq_parallel": True}, "sr", "image_path", False), + (SeedVRRunner, {"seq_parallel": False}, "sr", "image_path", True), + (QwenImageRunner, {"layered": True}, "i2i", "i2i_denoise_strength", False), + (QwenImageRunner, {"layered": False}, "i2i", "i2i_denoise_strength", True), + (WanAnimateRunner, {"replace_flag": True}, "animate", "src_bg_path", True), + (WanAnimateRunner, {"replace_flag": False}, "animate", "src_bg_path", False), + (WanDancerRunner, {"dancer_stage": "local"}, "s2v", "video_path", True), + (WanDancerRunner, {"dancer_stage": "global"}, "s2v", "video_path", False), + (WanAudioARRunner, {}, "rs2v", "target_video_length", False), + (WanAudioARRunner, {}, "rs2v", "prompt", True), + (WanAudioARRunner, {"prompt_travel": {"prompt_travel_text": ["scene"]}}, "rs2v", "prompt", False), + ], +) +def test_request_fields_follow_startup_capabilities(runner_cls, config, task, field, supported): + runner = object.__new__(runner_cls) + BaseRunner.__init__(runner, {"task": task, **config}) + assert (field in runner.get_supported_request_fields(task)) is supported + + +def test_ltx2_upsampler_uses_final_request_shape_for_stage_one(): + runner = object.__new__(LTX2Runner) + runner.config = {"use_upsampler": True} + runner.input_info = SimpleNamespace(target_shape=[1024, 1536]) + + runner.prepare_stage1_target_shape() + + assert runner.input_info.target_shape == [512, 768] + + +def test_ltx25_upsampler_validates_and_converts_final_request_shape(): + runner = object.__new__(LTX25Runner) + runner.config = {"use_upsampler": True} + runner.input_info = SimpleNamespace(target_shape=[1024, 1536]) + + runner.prepare_stage1_target_shape() + + assert runner.input_info.target_shape == [512, 768] + + +def test_ltx2_v2av_keeps_source_shape_resolution_when_request_omits_shape(): + runner = object.__new__(LTX2Runner) + runner.config = { + "task": "v2av", + "target_height": 768, + "target_width": 1280, + } + runner.supported_tasks = ("v2av",) + + inferred_shape = runner.create_input_info({"task": runner.config["task"], "video_path": "control.mp4"}) + explicit_shape = runner.create_input_info({"task": runner.config["task"], "video_path": "control.mp4", "target_shape": [480, 832]}) + + assert inferred_shape.target_shape == [] + assert explicit_shape.target_shape == [480, 832] + + +def test_worldplay_rejects_pose_and_frame_count_mismatch(): + with pytest.raises(ValueError, match="pose corresponds to 1 frames, num_frames must be set to 1"): + pose_to_input({"0": {}}, latent_num=2) + + +def test_wan_scheduler_owns_guidance_scale_validation(): + config = { + "dim": 16, + "enable_cfg": False, + "infer_steps": 4, + "num_heads": 2, + "sample_shift": 5.0, + "seq_parallel": False, + } + + scheduler = WanScheduler(config) + assert scheduler.sample_guide_scale is None + + config["enable_cfg"] = True + with pytest.raises(ValueError, match="sample_guide_scale"): + WanScheduler(config) + + +def test_wan_validates_source_inputs_before_model_loading(): + runner = object.__new__(WanRunner) + runner.config = {"task": "flf2v"} + runner.supported_tasks = ("flf2v",) + + with pytest.raises(ValueError, match="last_frame_path"): + runner.prepare_request({"task": runner.config["task"], "image_path": "first.png", "last_frame_path": ""}) + + +@pytest.mark.parametrize("disagg_mode", ["transformer", "decode"]) +def test_wan_defers_source_validation_for_downstream_disagg_roles(disagg_mode): + runner = object.__new__(WanRunner) + runner.config = {"task": "i2v", "disagg_mode": disagg_mode} + runner.supported_tasks = ("i2v",) + + runner.prepare_request({"task": runner.config["task"], "image_path": ""}) + + +def test_pipeline_passes_each_request_directly_to_runner(): + pipeline = object.__new__(LightX2VPipeline) + pipeline.model_cls = "wan2.1" + pipeline.task = "t2v" + pipeline.runner = RecordingRunner(task="t2v") + + pipeline.generate(seed=None, num_frames=49, target_shape=[480, 832], return_result_tensor=True) + pipeline.generate(seed=None, return_result_tensor=True) + + first_request, second_request = pipeline.runner.requests + assert first_request["target_video_length"] == 49 + assert first_request["target_shape"] == [480, 832] + assert "target_video_length" not in second_request + assert "target_shape" not in second_request + + +def test_pipeline_omits_default_tensor_result_request(): + pipeline = object.__new__(LightX2VPipeline) + pipeline.model_cls = "hunyuan3d" + pipeline.task = "i23d" + pipeline.runner = FileOnlyRecordingRunner(task="i23d") + + pipeline.generate(seed=None) + + assert "return_result_tensor" not in pipeline.runner.requests[0] + + +def test_multitask_runner_uses_request_task_without_changing_startup_config(): + runner = MultitaskRecordingRunner(task="t2i") + + t2i_input = runner.prepare_request({"task": runner.config["task"], "prompt": "create an image"}) + runner.run_request(t2i_input) + i2i_input = runner.prepare_request({"task": "i2i", "image_path": "input.png", "prompt": "edit the image"}) + runner.run_request(i2i_input) + + assert isinstance(t2i_input, T2IInputInfo) + assert isinstance(i2i_input, I2IInputInfo) + assert t2i_input.task == "t2i" + assert i2i_input.task == "i2i" + assert runner.config["task"] == "t2i" + assert t2i_input is not i2i_input + + next_input = runner.prepare_request({"task": runner.config["task"], "prompt": "create another image"}) + runner.run_request(next_input) + assert isinstance(next_input, T2IInputInfo) + assert next_input.task == "t2i" + + +def test_multitask_runner_validates_fields_for_the_selected_task(): + runner = MultitaskRecordingRunner(task="t2i") + + with pytest.raises(ValueError, match="image_path"): + runner.prepare_request({"task": "t2i", "image_path": "input.png"}) + + with pytest.raises(ValueError, match="not supported by this runner"): + runner.prepare_request({"task": "vace"}) + + +@pytest.mark.parametrize( + ("startup_task", "request_task", "request_data", "input_info_cls"), + [ + ("t2av", "i2av", {"image_path": "first.png"}, I2AVInputInfo), + ("i2av", "l2av", {"last_frame_path": "last.png"}, L2AVInputInfo), + ("l2av", "fl2av", {"image_path": "first.png", "last_frame_path": "last.png"}, FL2AVInputInfo), + ("fl2av", "t2av", {}, T2AVInputInfo), + ], +) +def test_minimax_h3_base_tasks_share_one_runner(startup_task, request_task, request_data, input_info_cls): + runner = make_minimax_h3_runner(startup_task) + + input_info = runner.prepare_request({"task": request_task, **request_data}) + + assert runner.supported_tasks == ("t2av", "i2av", "l2av", "fl2av") + assert isinstance(input_info, input_info_cls) + assert input_info.task == request_task + assert runner.config["task"] == startup_task + + +def test_minimax_h3_rejects_unimplemented_image_controls(): + runner = make_minimax_h3_runner("t2av") + + with pytest.raises(ValueError, match="image_strength"): + runner.prepare_request({"task": "i2av", "image_path": "input.png", "image_strength": 0.5}) + + with pytest.raises(ValueError, match="image_frame_idx"): + runner.prepare_request({"task": "i2av", "image_path": "input.png", "image_frame_idx": [10]}) + + +def test_minimax_h3_reference_transformer_is_a_separate_task_group(): + base_runner = make_minimax_h3_runner("t2av") + reference_runner = make_minimax_h3_runner("ref2av") + + assert reference_runner.supported_tasks == ("ref2av",) + with pytest.raises(ValueError, match="Task 'ref2av' is not supported"): + base_runner.prepare_request({"task": "ref2av", "image_path": "reference.png"}) + with pytest.raises(ValueError, match="Task 't2av' is not supported"): + reference_runner.prepare_request({"task": "t2av"}) + + +def test_minimax_h3_runtime_uses_input_info_task(): + runner = make_minimax_h3_runner("t2av") + runner.input_info = I2AVInputInfo(task="i2av") + + with pytest.raises(ValueError, match="i2av requires exactly one"): + runner._prepare_keyframes() + + runner.input_info = Ref2AVInputInfo(task="ref2av") + with pytest.raises(ValueError, match="cannot switch between the base and reference"): + runner._run_input_encoder_local_h3() + + +def test_minimax_h3_warmup_input_keeps_startup_task(): + runner = make_minimax_h3_runner("i2av") + + runner._prepare_warmup_inputs(480, 480, 124) + + assert runner.input_info.task == "i2av" + + +def test_minimax_h3_releases_request_conditioning(): + runner = object.__new__(MiniMaxH3Runner) + runner.config = {} + runner.model = SimpleNamespace(scheduler=SimpleNamespace(clear=lambda: None)) + runner.input_info = T2AVInputInfo() + runner.inputs = {} + runner.condition_video_latents = [object()] + runner.condition_audio_latents = [object()] + runner.keyframe_anchors = ("first",) + runner.prepared_references = [object()] + runner.maybe_empty_cache = lambda **kwargs: False + + runner.end_run() + + assert runner.condition_video_latents == [] + assert runner.condition_audio_latents == [] + assert runner.keyframe_anchors == () + assert runner.prepared_references is None + assert runner.input_info is None + assert not hasattr(runner, "inputs") + + +def test_runner_rejects_task_switching_by_default(): + runner = RecordingRunner(task="t2v") + + with pytest.raises(ValueError, match="Task 'i2v' is not supported by this runner"): + runner.prepare_request({"task": "i2v", "image_path": "input.png"}) + + +@pytest.mark.parametrize("runner_cls", [RecordingRunner, MultitaskRecordingRunner]) +@pytest.mark.parametrize("request_data", [{}, {"task": None}]) +def test_runner_defaults_task_only_for_single_task_models(runner_cls, request_data): + runner = runner_cls(task="t2i") + original_request = request_data.copy() + + if runner_cls is MultitaskRecordingRunner: + with pytest.raises(ValueError, match="task is required"): + runner.prepare_request(request_data) + else: + assert runner.prepare_request(request_data).task == "t2i" + assert request_data == original_request + + +def test_pipeline_without_default_requires_task_on_every_call(): + pipeline = LightX2VPipeline(model_cls="test_model", support_tasks=["t2i", "i2i"]) + pipeline.runner = MultitaskRecordingRunner(task="t2i") + + with pytest.raises(ValueError, match="task is required"): + pipeline.generate() + assert pipeline.generate(task="i2i", image_path="input.png").task == "i2i" + with pytest.raises(ValueError, match="task is required"): + pipeline.generate() + assert pipeline.generate(task="t2i").task == "t2i" + assert pipeline.runner.config["task"] == "t2i" + + +def test_neopp_rejects_direct_inputs_encoded_outside_the_runner(): + runner = object.__new__(NeoppRunner) + runner.config = {"task": "i2i"} + runner.supported_tasks = ("i2i",) + + assert isinstance(runner.prepare_request({"task": "i2i"}), NeoppInputInfo) + assert runner.get_supported_request_fields("i2i") == {"task", "seed", "save_result_path", "target_shape"} + with pytest.raises(ValueError, match="image_path"): + runner.prepare_request({"task": runner.config["task"], "image_path": "input.png"}) + + +def test_neopp_rejects_non_image_tasks(): + with pytest.raises(ValueError, match="does not support task 't2v'"): + NeoppRunner({"task": "t2v"}) + + +def test_pipeline_leaves_minimax_h3_output_path_unset(): + pipeline = object.__new__(LightX2VPipeline) + pipeline.model_cls = "minimax_h3" + pipeline.task = "t2av" + pipeline.runner = RecordingRunner(task="t2av") + + result = pipeline.generate(seed=None) + + assert result.save_result_path is None + + +@pytest.mark.parametrize( + ("task", "video_path"), + [ + ("t2v", None), + ("t2i", None), + ("sr", "input.mp4"), + ("sr", None), + ], +) +def test_pipeline_leaves_output_path_unset_for_all_tasks(task, video_path): + pipeline = object.__new__(LightX2VPipeline) + pipeline.model_cls = "test_model" + pipeline.task = task + pipeline.runner = RecordingRunner(task=task) + + result = pipeline.generate(seed=None, video_path=video_path) + + assert result.save_result_path is None + + +@pytest.mark.parametrize("enable_cfg", [False, True]) +@pytest.mark.parametrize("request_data", [{}, {"negative_prompt": ""}, {"negative_prompt": "blur"}]) +def test_wan_negative_prompt_requires_cfg(enable_cfg, request_data): + runner = object.__new__(WanRunner) + BaseRunner.__init__(runner, {"task": "t2v", "enable_cfg": enable_cfg, "vae_stride": (4, 8, 8)}) + original_request = request_data.copy() + if not enable_cfg and "negative_prompt" in request_data: + with pytest.raises(ValueError, match="negative_prompt"): + runner.prepare_request({"task": runner.config["task"], **request_data}) + else: + assert runner.prepare_request({"task": runner.config["task"], **request_data}).negative_prompt == request_data.get("negative_prompt", "") + assert request_data == original_request + assert runner.config["enable_cfg"] is enable_cfg + + +@pytest.mark.parametrize("request_data", [{}, {"save_result_path": None}]) +def test_default_runner_skips_saving_without_output_path(request_data): + runner = object.__new__(DefaultRunner) + runner.config = {"task": "t2v"} + runner.input_info = runner.create_input_info({"task": "t2v", **request_data}) + runner.gen_video_final = object() + + assert runner.process_images_after_vae_decoder() == {"video": None} + assert runner.gen_video_final is None + + +def test_pipeline_forwards_request_task_without_changing_its_default(tmp_path, monkeypatch): + model_path = tmp_path / "model" + model_path.mkdir() + config_path = tmp_path / "config.json" + config_path.write_text("{}", encoding="utf-8") + pipeline = LightX2VPipeline( + model_path=str(model_path), + model_cls="test_model", + task="t2i", + ) + monkeypatch.setattr("lightx2v.pipeline.validate_config_paths", lambda config: None) + monkeypatch.setattr( + "lightx2v.pipeline.build_runner", + lambda config: MultitaskRecordingRunner(task=config["task"]), + ) + pipeline.create_generator(config_json=str(config_path)) + + default_result = pipeline.generate(seed=None, return_result_tensor=True) + result = pipeline.generate(task="i2i", image_path="input.png", seed=None, return_result_tensor=True) + next_result = pipeline.generate(seed=None, return_result_tensor=True) + + assert default_result.task == "t2i" + assert result.task == "i2i" + assert next_result.task == "t2i" + assert pipeline.task == "t2i" + assert pipeline.runner.config["task"] == "t2i" + assert "support_tasks" not in pipeline.startup_config + assert "support_tasks" not in pipeline.runner.config + + +def test_runner_applies_request_seed(monkeypatch): + seeds = [] + monkeypatch.setattr("lightx2v.models.runners.base_runner.seed_all", seeds.append) + + runner = RecordingRunner() + input_info = runner.prepare_request({"task": "t2v", "seed": 123}) + runner.run_request(input_info) + + assert seeds == [123] + + input_info = runner.prepare_request({"task": "t2v", "seed": 456}) + runner.run_request(input_info) + + assert seeds == [123, 456] + + +def test_seedvr_keeps_detected_fps_in_request_context(): + runner = object.__new__(SeedVRRunner) + runner.config = {"fps": 16} + first_request = SRInputInfo() + runner.input_info = first_request + + runner.set_output_fps(24) + + assert first_request.output_fps == 24 + assert runner.get_output_fps() == 24 + assert runner.config["fps"] == 16 + + runner.input_info = SRInputInfo() + + assert runner.get_output_fps() == 16 + + +def test_pipeline_forwards_task_specific_inputs(): + pipeline = object.__new__(LightX2VPipeline) + pipeline.model_cls = "wan2.2_animate" + pipeline.task = "animate" + pipeline.runner = RecordingRunner(task="animate") + + pipeline.generate(seed=None, src_pose_path="pose.mp4", src_face_path="face.mp4") + + request = pipeline.runner.requests[0] + assert request["src_pose_path"] == "pose.mp4" + assert request["src_face_path"] == "face.mp4" + + +def test_pipeline_copies_startup_defaults_to_input_info(): + pipeline = object.__new__(LightX2VPipeline) + pipeline.model_cls = "ltx2_5" + pipeline.task = "t2av" + pipeline.runner = RecordingRunner(task="t2av") + + input_info = pipeline.generate(seed=None, return_result_tensor=True) + + assert input_info.target_video_length == 81 + + +def test_pipeline_preserves_explicit_task_when_loading_model_config(tmp_path, monkeypatch): + model_path = tmp_path / "model" + model_path.mkdir() + json_model_path = tmp_path / "json-model" + json_model_path.mkdir() + config_path = tmp_path / "config.json" + config_path.write_text( + f'{{"model_cls": "json_model", "model_path": "{json_model_path}", "task": "i2v", "infer_steps": 20, "target_video_length": 49}}', + encoding="utf-8", + ) + pipeline = LightX2VPipeline( + model_path=str(model_path), + model_cls="constructor_model", + task="t2v", + ) + monkeypatch.setattr("lightx2v.pipeline.validate_config_paths", lambda config: None) + monkeypatch.setattr("lightx2v.pipeline.build_runner", lambda config: SimpleNamespace(config=config)) + + pipeline.create_generator(config_json=str(config_path)) + + assert pipeline.model_cls == "json_model" + assert pipeline.model_path == str(json_model_path) + assert pipeline.task == "t2v" + + with pytest.raises(RuntimeError, match="already been created"): + pipeline.create_generator() + + +def test_pipeline_manual_startup_fields(tmp_path, monkeypatch): + model_path = tmp_path / "model" + model_path.mkdir() + pipeline = LightX2VPipeline( + model_path=str(model_path), + model_cls="test_model", + task="t2v", + ) + monkeypatch.setattr("lightx2v.pipeline.validate_config_paths", lambda config: None) + monkeypatch.setattr("lightx2v.pipeline.build_runner", lambda config: SimpleNamespace(config=config)) + + pipeline.create_generator( + infer_steps=17, + num_frames=49, + height=480, + width=832, + resize_mode="adaptive", + ) + + assert pipeline.runner.config["infer_steps"] == 17 + assert pipeline.runner.config["target_video_length"] == 49 + assert pipeline.runner.config["target_height"] == 480 + assert pipeline.runner.config["target_width"] == 832 + assert pipeline.runner.config["resize_mode"] == "adaptive" diff --git a/tests/test_request_default_priority.py b/tests/test_request_default_priority.py new file mode 100644 index 000000000..f98d81dc3 --- /dev/null +++ b/tests/test_request_default_priority.py @@ -0,0 +1,254 @@ +import json +from unittest.mock import Mock + +import pytest + +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.cosmos3.cosmos3_runner import Cosmos3Runner +from lightx2v.models.runners.flux2.flux2_runner import Flux2Runner +from lightx2v.models.runners.longcat_image.longcat_image_runner import LongCatImageRunner +from lightx2v.models.runners.ltx2.ltx2_runner import LTX2Runner +from lightx2v.models.runners.qwen_image.qwen_image_runner import QwenImageRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS, PROMPT_FIELDS, VIDEO_OUTPUT_FIELDS +from lightx2v.models.runners.wan.wan_animate2_runner import WanAnimate2Runner +from lightx2v.models.runners.z_image.z_image_runner import ZImageRunner +from lightx2v.pipeline import LightX2VPipeline +from lightx2v.utils.input_info import UNSET + + +class PreparedInputRunner(BaseRunner): + supported_request_fields_by_task = { + "t2v": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | VIDEO_OUTPUT_FIELDS, + "t2i": COMMON_REQUEST_FIELDS | PROMPT_FIELDS | {"aspect_ratio", "target_shape"}, + "sr": COMMON_REQUEST_FIELDS | {"image_path", "video_path", "target_shape"}, + "animate": COMMON_REQUEST_FIELDS | {"image_path", "video_path"}, + } + + def __init__(self, config): + super().__init__(config) + self.requests = [] + + def create_input_info(self, request_data): + self.requests.append(request_data) + return super().create_input_info(request_data) + + def run_request(self, input_info): + return input_info + + +def make_pipeline(config, model_cls="test_model"): + pipeline = object.__new__(LightX2VPipeline) + pipeline.task = config["task"] + pipeline.model_cls = model_cls + pipeline.runner = PreparedInputRunner(config) + return pipeline + + +@pytest.mark.parametrize(("seed", "expected_seed"), [(UNSET, 42), (None, 42), (0, 0)]) +def test_prepare_request_handles_omitted_values_without_changing_explicit_values(seed, expected_seed): + runner = PreparedInputRunner({"task": "t2i", "target_shape": [512, 768], "enable_cfg": True}) + request_data = { + "task": "t2i", + "seed": seed, + "prompt": None, + "negative_prompt": "", + "target_shape": None, + "aspect_ratio": None, + "save_result_path": None, + "return_result_tensor": False, + } + + input_info = runner.prepare_request(request_data) + + assert input_info.target_shape == [512, 768] + assert input_info.seed == expected_seed + assert input_info.prompt == "" + assert input_info.negative_prompt == "" + assert runner.requests[0]["negative_prompt"] == "" + assert input_info.save_result_path is None + assert input_info.return_result_tensor is False + assert request_data["seed"] is seed + assert request_data["target_shape"] is None + + +def test_shape_request_overrides_json_then_next_request_restores_json(): + runner = PreparedInputRunner({"task": "t2i", "target_shape": [512, 768], "target_height": 720, "target_width": 1280}) + + first = runner.prepare_request({"task": runner.config["task"], "target_shape": [480, 832]}) + second = runner.prepare_request({"task": runner.config["task"]}) + + assert first.target_shape == [480, 832] + assert second.target_shape == [512, 768] + assert runner.config["target_shape"] == [512, 768] + + +@pytest.mark.parametrize("task", ("i2va", "v2av")) +@pytest.mark.parametrize( + "request_fields,expected_domain,expected_view", + [({}, "av", "ego_view"), ({"domain_name": "droid_lerobot", "view_point": "wrist_view"}, "droid_lerobot", "wrist_view")], +) +def test_cosmos_action_file_overrides_config_below_explicit_request(tmp_path, task, request_fields, expected_domain, expected_view): + config = {"task": task, "action_chunk_size": 16, "raw_action_dim": 29, "domain_name": "agibotworld", "view_point": "concat_view"} + action_path = tmp_path / "actions.json" + action_path.write_text(json.dumps({"action_chunk_size": 8, "raw_action_dim": 9, "domain_name": "av", "view_point": "ego_view"})) + runner = object.__new__(Cosmos3Runner) + BaseRunner.__init__(runner, config.copy()) + runner.input_info = runner.prepare_request({"action_path": str(action_path), **request_fields}) + + runner._prepare_action_context() + + assert runner.input_info.action_chunk_size == 8 + assert runner.input_info.target_video_length == 9 + assert runner._get_action_value("raw_action_dim") == 9 + assert runner._get_action_value("domain_name") == expected_domain + assert runner._get_action_value("view_point") == expected_view + assert runner.config == config + + +@pytest.mark.parametrize("task", ("i2va", "v2av")) +def test_cosmos_action_metadata_is_resolved_again_for_each_request(tmp_path, task): + config = {"task": task, "action_chunk_size": 16, "raw_action_dim": 29, "domain_name": "agibotworld", "view_point": "concat_view"} + action_path = tmp_path / "actions.json" + runner = object.__new__(Cosmos3Runner) + BaseRunner.__init__(runner, config.copy()) + + for spec in ( + {"action_chunk_size": 8, "raw_action_dim": 9, "domain_name": "av", "view_point": "ego_view"}, + {"action_chunk_size": 4, "raw_action_dim": 2, "domain_name": "pusht", "view_point": "third_person_view"}, + {}, + ): + action_path.write_text(json.dumps(spec)) + runner.input_info = runner.prepare_request({"action_path": str(action_path)}) + + runner._prepare_action_context() + + expected = spec or config + assert runner.input_info.target_video_length == expected["action_chunk_size"] + 1 + for name in ("action_chunk_size", "raw_action_dim", "domain_name", "view_point"): + assert runner._get_action_value(name) == expected[name] + assert runner.config == config + + +def test_explicit_aspect_ratio_can_replace_json_shape(): + runner = PreparedInputRunner({"task": "t2i", "target_shape": [512, 768], "target_height": 720, "target_width": 1280}) + + assert runner.prepare_request({"task": runner.config["task"], "aspect_ratio": "1:1"}).target_shape == [] + assert runner.prepare_request({"task": runner.config["task"]}).target_shape == [512, 768] + + +@pytest.mark.parametrize( + ("config_shape", "request_shape", "expected_shape", "should_probe"), + [ + (None, None, [512, 768], True), + ([640, 960], None, [640, 960], False), + ([640, 960], [480, 832], [480, 832], False), + (None, [480, 832], [480, 832], False), + ], +) +def test_ltx2_source_resolution_only_fills_missing_shape(config_shape, request_shape, expected_shape, should_probe): + runner = object.__new__(LTX2Runner) + config = {"task": "v2av", "target_height": 768, "target_width": 1280} + request = {"video_path": "control.mp4"} + if config_shape is not None: + config["target_shape"] = config_shape + if request_shape is not None: + request["target_shape"] = request_shape + BaseRunner.__init__(runner, config) + probes = [] + + def probe(path): + probes.append(path) + return 512, 768 + + runner._probe_video_hw = probe + runner._get_ref_downscale_factor = lambda: 1.0 + runner.input_info = runner.prepare_request({"task": runner.config["task"], **request}) + runner._override_target_hw_from_ref_video() + + assert runner.input_info.target_shape == expected_shape + assert probes == (["control.mp4"] if should_probe else []) + assert runner.config == config + + +def test_pipeline_uses_request_content_and_code_defaults(): + pipeline = make_pipeline({"task": "t2v"}) + + default = pipeline.generate() + explicit = pipeline.generate(seed=0, prompt="", save_result_path="explicit.mp4") + restored = pipeline.generate() + + assert (default.seed, default.prompt, default.save_result_path) == (42, "", None) + assert (explicit.seed, explicit.prompt, explicit.save_result_path) == (0, "", "explicit.mp4") + assert (restored.seed, restored.prompt, restored.save_result_path) == (42, "", None) + assert pipeline.runner.requests[0] == {"task": "t2v"} + assert pipeline.runner.requests[2] == {"task": "t2v"} + + +@pytest.mark.parametrize("task", ["t2i", "t2v"]) +def test_pipeline_preserves_explicit_none_output_path(task): + pipeline = make_pipeline({"task": task}) + + assert pipeline.generate(save_result_path=None).save_result_path is None + assert pipeline.generate(save_result_path="explicit.png").save_result_path == "explicit.png" + assert pipeline.generate().save_result_path is None + + +@pytest.mark.parametrize("runner_cls", [Flux2Runner, LongCatImageRunner, QwenImageRunner, ZImageRunner]) +@pytest.mark.parametrize("output_path", [None, "result.png"]) +def test_image_runners_save_only_with_an_output_path(runner_cls, output_path): + runner = object.__new__(runner_cls) + BaseRunner.__init__(runner, {"task": "t2i", "model_variant": "klein"}) + runner._gc_frozen = True + request_data = {} if output_path is None else {"save_result_path": output_path} + input_info = runner.prepare_request({"task": runner.config["task"], **request_data}) + image = Mock() + + if runner_cls is QwenImageRunner: + runner._save_images([image], input_info) + else: + runner.run_input_encoder = lambda: {} + runner.set_latent_shape = lambda: None + runner.run_dit = lambda: (None, None) + runner.run_vae_decoder = lambda latents: [image] + runner.end_run = lambda: None + runner.run_pipeline(input_info) + + if output_path is None: + image.save.assert_not_called() + else: + image.save.assert_called_once_with(output_path) + + +def test_pipeline_preserves_explicit_default_named_output_for_video(): + pipeline = make_pipeline({"task": "t2v"}) + + result = pipeline.generate(save_result_path="lightx2v_gen_result.png") + + assert result.save_result_path == "lightx2v_gen_result.png" + + +def test_pipeline_sr_does_not_infer_output_paths_from_media(): + pipeline = make_pipeline({"task": "sr"}) + + video = pipeline.generate(video_path="input.mp4") + image = pipeline.generate(video_path="", image_path="input.png") + restored = pipeline.generate(video_path="input.mp4") + + assert video.save_result_path is None + assert image.save_result_path is None + assert restored.save_result_path is None + + +def test_pipeline_animate2_uses_request_seed_and_code_default(): + pipeline = make_pipeline({"task": "animate", "seed": 71}, model_cls="wan2.2_animate2_distilled") + + assert pipeline.generate().seed == 42 + assert pipeline.generate(seed=None).seed == 42 + assert pipeline.generate(seed=0).seed == 0 + + +@pytest.mark.parametrize(("seed", "expected"), [(None, 42), (123, 123), (0, 0)]) +def test_animate2_receives_resolved_request_seed(seed, expected): + runner = object.__new__(WanAnimate2Runner) + BaseRunner.__init__(runner, {"task": "animate", "seed": 71}) + assert runner.prepare_request({"task": "animate", "seed": seed, "image_path": "reference.png", "video_path": "driver.mp4"}).seed == expected diff --git a/tests/test_request_seed.py b/tests/test_request_seed.py new file mode 100644 index 000000000..1d6af3712 --- /dev/null +++ b/tests/test_request_seed.py @@ -0,0 +1,57 @@ +import pytest + +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS +from lightx2v.pipeline import LightX2VPipeline + + +class SeedRunner(BaseRunner): + supported_request_fields_by_task = {"t2i": COMMON_REQUEST_FIELDS} + + def run_pipeline(self, input_info): + return input_info.seed + + +def test_request_preparation_does_not_apply_seed(monkeypatch): + seeded = [] + monkeypatch.setattr("lightx2v.models.runners.base_runner.seed_all", seeded.append) + runner = SeedRunner({"task": "t2i"}) + runner._gc_frozen = True + + input_info = runner.prepare_request({"task": "t2i", "seed": 7}) + assert seeded == [] + + assert runner.run_request(input_info) == 7 + assert seeded == [7] + + +@pytest.mark.parametrize("config", [{}, {"seed": None}, {"seed": 73}, {"seed": 0}]) +@pytest.mark.parametrize("request_data", [{}, {"seed": None}, {"seed": 0}, {"seed": 12}]) +def test_python_and_direct_requests_resolve_the_same_seed(config, request_data): + runner = SeedRunner({"task": "t2i", **config}) + pipeline = LightX2VPipeline(task="t2i", model_cls="test") + pipeline.runner = runner + expected = request_data.get("seed") + if expected is None: + expected = 42 + + assert runner.prepare_request({"task": runner.config["task"], **request_data}).seed == expected + runner.run_request = lambda input_info: input_info.seed + assert pipeline.generate(**request_data) == expected + assert runner.config == {"task": "t2i", **config} + + +def test_fixed_request_seeds_do_not_replace_json(monkeypatch): + seeded = [] + monkeypatch.setattr("lightx2v.models.runners.base_runner.seed_all", seeded.append) + runner = SeedRunner({"task": "t2i", "seed": 123}) + runner._gc_frozen = True + + first = runner.prepare_request({"task": runner.config["task"]}) + assert runner.run_request(first) == 42 + assert runner.run_request(first) == 42 + for request_data, expected in [({"seed": 456}, 456), ({"seed": 0}, 0), ({}, 42)]: + input_info = runner.prepare_request({"task": runner.config["task"], **request_data}) + assert runner.run_request(input_info) == expected + assert seeded == [42, 42, 456, 0, 42] + assert runner.config["seed"] == 123 diff --git a/tests/test_server_request_defaults.py b/tests/test_server_request_defaults.py new file mode 100644 index 000000000..b1d7ff1a9 --- /dev/null +++ b/tests/test_server_request_defaults.py @@ -0,0 +1,425 @@ +import asyncio +import base64 +import json +from io import BytesIO +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from PIL import Image +from fastapi import FastAPI, HTTPException, UploadFile +from fastapi.testclient import TestClient + +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS +from lightx2v.server.api import openai_images +from lightx2v.server.api.tasks import image as image_api +from lightx2v.server.api.tasks import video as video_api +from lightx2v.server.schema import ImageTaskRequest, SenseNovaVisionTaskRequest, TalkObject, VideoTaskRequest +from lightx2v.server.services.file_service import FileService +from lightx2v.server.services.generation.image import ImageGenerationService +from lightx2v.server.services.generation.sensenova_vision import SenseNovaVisionGenerationService +from lightx2v.server.services.generation.video import VideoGenerationService +from lightx2v.server.services.inference import worker as worker_module +from lightx2v.server.services.inference.worker import TorchrunInferenceWorker +from lightx2v.server.task_manager import TaskManager +from lightx2v.utils.input_info import SenseNovaVisionInputInfo +from lightx2v.utils.lockable_dict import LockableDict + + +class RecordingRunner(BaseRunner): + supported_request_fields_by_task = {task: COMMON_REQUEST_FIELDS | {"prompt", "negative_prompt", "target_shape"} for task in ("t2i", "t2v")} + + def __init__(self, config): + super().__init__(config) + self.requests = [] + + def run_pipeline(self, input_info): + self.requests.append(input_info) + + +class SenseNovaRecordingRunner(RecordingRunner): + input_info_cls_by_task = {"omni_vision_task": SenseNovaVisionInputInfo} + supported_request_fields_by_task = {"omni_vision_task": COMMON_REQUEST_FIELDS | {"prompt", "image_path", "omni_vision_subtask", "raw_output_path", "glb_output_path", "postprocess_predictions"}} + + def run_pipeline(self, input_info): + super().run_pipeline(input_info) + return {"text": "ready"} + + +def make_service(tmp_path, monkeypatch, task, config): + monkeypatch.setattr("lightx2v.models.runners.base_runner.seed_all", lambda seed: None) + runner_cls = SenseNovaRecordingRunner if task == "omni_vision_task" else RecordingRunner + runner = runner_cls(LockableDict({"task": task, **config})) + runner.config.lock() + worker = TorchrunInferenceWorker() + worker.world_size = 1 + worker.rank = 0 + worker.runner = runner + inference = SimpleNamespace(worker=worker, submit_task_async=worker.process_request) + cls = {"t2i": ImageGenerationService, "t2v": VideoGenerationService, "omni_vision_task": SenseNovaVisionGenerationService}[task] + if task == "omni_vision_task": + + async def resolve_images(self, sources): + return sources + + monkeypatch.setattr(cls, "_resolve_images", resolve_images) + return cls(FileService(tmp_path), inference), runner + + +@pytest.mark.parametrize(("task", "request_cls", "suffix"), [("t2i", ImageTaskRequest, ".png"), ("t2v", VideoTaskRequest, ".mp4")]) +def test_server_restores_output_specs_between_requests(tmp_path, monkeypatch, task, request_cls, suffix): + config = {"target_shape": [512, 768], "enable_cfg": True} + service, runner = make_service(tmp_path, monkeypatch, task, config) + first = request_cls(prompt="first", seed=0, negative_prompt="explicit", save_result_path="override" + suffix, target_shape=[480, 832]) + second = request_cls(prompt="second") + assert "seed" not in second.model_fields_set + assert "save_result_path" not in second.model_fields_set + + first_response = asyncio.run(service.generate_with_stop_event(first, asyncio.Event())) + second_response = asyncio.run(service.generate_with_stop_event(second, asyncio.Event())) + + assert [request.seed for request in runner.requests] == [0, 42] + assert [request.negative_prompt for request in runner.requests] == ["explicit", ""] + assert [request.target_shape for request in runner.requests] == [[480, 832], [512, 768]] + assert Path(first_response.save_result_path) == Path(runner.requests[0].save_result_path).absolute() + assert second_response.save_result_path is None + assert runner.requests[1].save_result_path is None + assert runner.config == {"task": task, **config} + + +@pytest.mark.parametrize(("task", "request_cls"), [("t2i", ImageTaskRequest), ("t2v", VideoTaskRequest)]) +@pytest.mark.parametrize("output_fields", [{}, {"save_result_path": None}]) +def test_server_does_not_invent_output_paths(tmp_path, monkeypatch, task, request_cls, output_fields): + service, runner = make_service(tmp_path, monkeypatch, task, {}) + message = request_cls(prompt="image", **output_fields) + + response = asyncio.run(service.generate_with_stop_event(message, asyncio.Event())) + + assert runner.requests[0].seed == 42 + assert runner.requests[0].save_result_path is None + assert response.save_result_path is None + assert response.task_status == "completed" + + +def test_sync_image_returns_memory_result_without_output_path(tmp_path, monkeypatch): + service, runner = make_service(tmp_path, monkeypatch, "t2i", {}) + original_run = runner.run_pipeline + + def run(input_info): + original_run(input_info) + return {"images": [Image.new("RGB", (4, 4), "red")]} + + monkeypatch.setattr(runner, "run_pipeline", run) + message = ImageTaskRequest(prompt="image", prefer_memory_result=True) + + response = asyncio.run(service.generate_with_stop_event(message, asyncio.Event())) + + assert runner.requests[0].save_result_path is None + assert runner.requests[0].return_result_tensor is True + assert response.save_result_path is None + assert response.result_png.startswith(b"\x89PNG") + assert not list(service.file_service.output_video_dir.iterdir()) + + +def test_single_task_post_uses_startup_task_and_rejects_conflicting_task(tmp_path, monkeypatch): + service, runner = make_service(tmp_path, monkeypatch, "t2i", {}) + + asyncio.run(service.generate_with_stop_event(ImageTaskRequest(prompt="image"), asyncio.Event())) + assert runner.requests[-1].task == "t2i" + with pytest.raises(RuntimeError, match="Task 't2v' is not supported"): + asyncio.run(service.generate_with_stop_event(ImageTaskRequest(task="t2v", prompt="image"), asyncio.Event())) + assert len(runner.requests) == 1 + + +def test_multitask_post_requires_and_preserves_explicit_task(tmp_path, monkeypatch): + service, runner = make_service(tmp_path, monkeypatch, "t2i", {}) + runner.supported_tasks = ("t2i", "t2v") + worker = service.inference_service.worker + worker.lora_dir = tmp_path + lora_changes = [] + monkeypatch.setattr(worker, "switch_lora", lambda name, strength: lora_changes.append((name, strength))) + + with pytest.raises(RuntimeError, match="task is required"): + asyncio.run(service.generate_with_stop_event(ImageTaskRequest(prompt="image"), asyncio.Event())) + assert lora_changes == [] + for task in ("t2v", "t2i"): + asyncio.run(service.generate_with_stop_event(ImageTaskRequest(task=task, prompt="image"), asyncio.Event())) + assert runner.requests[-1].task == task + assert len(lora_changes) == 2 + assert runner.config["task"] == "t2i" + + +@pytest.mark.parametrize("task", ["t2i", "t2v", "omni_vision_task"]) +@pytest.mark.parametrize( + ("config", "request_fields", "expected"), + [ + ({}, {}, 42), + ({}, {"seed": None}, 42), + ({"seed": None}, {}, 42), + ({"seed": None}, {"seed": None}, 42), + ({"seed": 73}, {}, 42), + ({"seed": 73}, {"seed": None}, 42), + ({"seed": 73}, {"seed": 0}, 0), + ], +) +def test_server_seed_defaults_are_resolved_by_runner(tmp_path, monkeypatch, task, config, request_fields, expected): + config = {"model_cls": "sensenova_vision", **config} if task == "omni_vision_task" else config + service, runner = make_service(tmp_path, monkeypatch, task, config) + if task == "omni_vision_task": + message = SenseNovaVisionTaskRequest(task="understanding", images=["reference.png"], visualize=False, **request_fields) + else: + request_cls = ImageTaskRequest if task == "t2i" else VideoTaskRequest + message = request_cls(prompt="image", **request_fields) + + asyncio.run(service.generate_with_stop_event(message, asyncio.Event())) + + assert runner.requests[0].seed == expected + assert runner.config == {"task": task, **config} + + +@pytest.mark.parametrize("task", ["t2i", "t2v", "omni_vision_task"]) +def test_server_forwards_explicit_seed_to_runner(tmp_path, monkeypatch, task): + config = {"seed": 73, "model_cls": "sensenova_vision"} if task == "omni_vision_task" else {"seed": 73} + service, runner = make_service(tmp_path, monkeypatch, task, config) + if task == "omni_vision_task": + message = SenseNovaVisionTaskRequest(task="understanding", images=["reference.png"], visualize=False, seed=12) + else: + request_cls = ImageTaskRequest if task == "t2i" else VideoTaskRequest + message = request_cls(prompt="image", seed=12) + seeded = [] + monkeypatch.setattr("lightx2v.models.runners.base_runner.seed_all", seeded.append) + submitted = [] + submit = service.inference_service.submit_task_async + + async def submit_task(task_data): + submitted.append(task_data["seed"]) + return await submit(task_data) + + service.inference_service.submit_task_async = submit_task + + asyncio.run(service.generate_with_stop_event(message, asyncio.Event())) + + assert submitted == [12] + assert runner.requests[0].seed == 12 + assert seeded == [12] + assert runner.config["seed"] == 73 + + +@pytest.mark.parametrize(("config", "expected"), [({}, False), ({"use_compile": True}, False), ({"use_compile": True, "warmup": True}, True)]) +def test_server_warmup_is_owned_by_startup_config(tmp_path, monkeypatch, config, expected): + config_path = tmp_path / "deployment.json" + config_path.write_text(json.dumps(config)) + captured = [] + monkeypatch.setattr(worker_module, "build_runner", lambda config: captured.append(config) or object()) + worker = TorchrunInferenceWorker() + worker.world_size = 1 + args = SimpleNamespace(model_path=str(tmp_path), config_json=str(config_path), model_cls="test_model", task="t2v") + + assert worker.init(args) + assert captured[0]["warmup"] is expected + assert "warmup" not in ImageTaskRequest.model_fields + assert "warmup" not in VideoTaskRequest.model_fields + + +@pytest.mark.parametrize("seed", [None, 0]) +@pytest.mark.parametrize("task", ["image", "video"]) +def test_form_preserves_seed_omission_and_zero(tmp_path, monkeypatch, task, seed): + api = image_api if task == "image" else video_api + messages = [] + monkeypatch.setattr(api, "get_services", lambda: SimpleNamespace(file_service=FileService(tmp_path))) + monkeypatch.setattr(api.task_manager, "create_task", lambda message: messages.append(message) or message.task_id) + fields = dict(request=SimpleNamespace(form=AsyncMock(return_value={"prompt": "prompt"})), image_file=None, prompt="prompt", save_result_path="", negative_prompt="", seed=seed) + if task == "image": + fields["aspect_ratio"] = None + asyncio.run(api.create_image_task_form(**fields)) + else: + fields.update(last_frame_file=None, target_video_length=None, audio_file=None, video_duration=None) + asyncio.run(api.create_video_task_form(**fields)) + + assert ("seed" in messages[0].model_fields_set) == (seed is not None) + assert messages[0].seed == seed + + +@pytest.mark.parametrize("task", ["image", "video"]) +def test_form_uploads_files_and_reports_full_queue(tmp_path, monkeypatch, task): + api = image_api if task == "image" else video_api + file_service = FileService(tmp_path) + manager = TaskManager(max_queue_size=1) + monkeypatch.setattr(api, "get_services", lambda: SimpleNamespace(file_service=file_service)) + monkeypatch.setattr(api, "task_manager", manager) + fields = dict( + request=SimpleNamespace(form=AsyncMock(return_value={"prompt": "prompt"})), + image_file=UploadFile(file=BytesIO(b"first"), filename="first.jpg"), + prompt="prompt", + save_result_path="", + negative_prompt="", + seed=None, + ) + expected = {"image_path": (file_service.input_image_dir, ".jpg", b"first")} + if task == "image": + fields["aspect_ratio"] = None + create_form = api.create_image_task_form + else: + fields.update( + last_frame_file=UploadFile(file=BytesIO(b"last"), filename="last.png"), + audio_file=UploadFile(file=BytesIO(b""), filename="audio.wav"), + target_video_length=None, + video_duration=None, + ) + expected.update(last_frame_path=(file_service.input_image_dir, ".png", b"last"), audio_path=(file_service.input_audio_dir, ".wav", b"")) + create_form = api.create_video_task_form + + response = asyncio.run(create_form(**fields)) + message = manager.get_task(response.task_id).message + assert response.task_status == "pending" + for field, (directory, suffix, content) in expected.items(): + path = Path(getattr(message, field)) + assert path.parent == directory + assert path.suffix == suffix + assert path.read_bytes() == content + + for field in ("image_file", "last_frame_file", "audio_file"): + if field in fields: + fields[field] = None + with pytest.raises(HTTPException) as exc: + asyncio.run(create_form(**fields)) + assert exc.value.status_code == 503 + assert "queue is full" in exc.value.detail + + +@pytest.mark.parametrize(("filename", "content", "suffix"), [("image.jpeg", b"image", ".jpeg"), ("image", b"image", ".png"), ("", b"image", None), ("image.png", b"", None)]) +def test_openai_edit_upload_contract(tmp_path, monkeypatch, filename, content, suffix): + file_service = FileService(tmp_path) + monkeypatch.setattr(openai_images, "get_services", lambda: SimpleNamespace(file_service=file_service)) + generate = AsyncMock(return_value=(b"result", None)) + monkeypatch.setattr(openai_images, "_run_sync_image_task", generate) + request = openai_images.create_openai_image_edit( + request=SimpleNamespace(form=AsyncMock(return_value={})), + image=[UploadFile(file=BytesIO(content), filename=filename)], + prompt="edit", + mask=None, + model=None, + n=1, + size=None, + response_format="b64_json", + user=None, + negative_prompt="", + seed=None, + i2i_denoise_strength=None, + ) + if suffix is None: + with pytest.raises(HTTPException) as exc: + asyncio.run(request) + assert exc.value.status_code == 400 + assert not list(file_service.input_image_dir.iterdir()) + generate.assert_not_awaited() + else: + response = asyncio.run(request) + path = Path(generate.await_args.args[1].image_path) + assert path.parent == file_service.input_image_dir + assert path.suffix == suffix + assert path.read_bytes() == content + assert response.data[0]["b64_json"] == base64.b64encode(b"result").decode("utf-8") + + +@pytest.mark.parametrize("task", ["image", "video"]) +@pytest.mark.parametrize( + "text_fields", [{}, {"prompt": "", "negative_prompt": "", "save_result_path": ""}, {"prompt": "cat", "negative_prompt": "blur", "save_result_path": "chosen.png"}, {"task": "i2v", "prompt": "cat"}] +) +def test_http_form_preserves_submitted_text(tmp_path, monkeypatch, task, text_fields): + api = image_api if task == "image" else video_api + messages = [] + monkeypatch.setattr(api, "get_services", lambda: SimpleNamespace(file_service=FileService(tmp_path))) + monkeypatch.setattr(api.task_manager, "create_task", lambda message: messages.append(message) or message.task_id) + app = FastAPI() + app.include_router(api.router) + + with TestClient(app) as client: + response = client.post("/form", data=text_fields, files={"image_file": ("input.png", b"image", "image/png")}) + + assert response.status_code == 200 + message = messages[0] + for key in ("task", "prompt", "negative_prompt", "save_result_path"): + assert (key in message.model_fields_set) == (key in text_fields) + assert getattr(message, key) == text_fields.get(key, None if key in {"task", "save_result_path"} else "") + + +@pytest.mark.parametrize("image_field", ["image", "image[]"]) +@pytest.mark.parametrize("negative_prompt", [None, "", "blur"]) +def test_http_openai_edit_preserves_negative_prompt(tmp_path, monkeypatch, image_field, negative_prompt): + monkeypatch.setattr(openai_images, "get_services", lambda: SimpleNamespace(file_service=FileService(tmp_path))) + generate = AsyncMock(return_value=(b"result", None)) + monkeypatch.setattr(openai_images, "_run_sync_image_task", generate) + app = FastAPI() + app.include_router(openai_images.router) + data = {"prompt": "edit", "response_format": "b64_json"} + if negative_prompt is not None: + data["negative_prompt"] = negative_prompt + + with TestClient(app) as client: + response = client.post("/edits", data=data, files={image_field: ("input.png", b"image", "image/png")}) + + assert response.status_code == 200 + message = generate.await_args.args[1] + assert message.task == "i2i" + assert ("negative_prompt" in message.model_fields_set) == (negative_prompt is not None) + assert message.negative_prompt == ("" if negative_prompt is None else negative_prompt) + + +def test_openai_generation_selects_task_without_a_public_task_parameter(monkeypatch): + generate = AsyncMock(return_value=(b"result", None)) + monkeypatch.setattr(openai_images, "_run_sync_image_task", generate) + app = FastAPI() + app.include_router(openai_images.router) + + with TestClient(app) as client: + response = client.post("/generations", json={"prompt": "cat", "response_format": "b64_json"}) + + assert response.status_code == 200 + assert generate.await_args.args[1].task == "t2i" + + +@pytest.mark.parametrize("source_type", ["local", "base64", "url"]) +def test_video_resolves_frames_and_talk_masks(tmp_path, source_type): + file_service = FileService(tmp_path) + image_path = tmp_path / "source.png" + Image.new("RGB", (4, 4)).save(image_path) + source = {"local": str(image_path), "base64": base64.b64encode(image_path.read_bytes()).decode("utf-8"), "url": "https://example.test/image.png"}[source_type] + file_service.download_image = AsyncMock(return_value=image_path) + submit = AsyncMock(return_value={"status": "success", "save_result_path": "output.mp4"}) + inference = SimpleNamespace(worker=SimpleNamespace(runner=SimpleNamespace(config={})), submit_task_async=submit) + service = VideoGenerationService(file_service, inference) + message = VideoTaskRequest(image_path=source, last_frame_path=source, talk_objects=[TalkObject(audio="audio.wav", mask=source)]) + + asyncio.run(service.generate_with_stop_event(message, asyncio.Event())) + + payload = submit.await_args.args[0] + talk_objects = json.loads((Path(payload["audio_path"]) / "config.json").read_text())["talk_objects"] + for path in (payload["image_path"], payload["last_frame_path"], talk_objects[0]["mask"]): + assert Path(path).read_bytes() == image_path.read_bytes() + assert talk_objects[0]["audio"] == "audio.wav" + assert "talk_objects" not in payload + + +def test_image_resolves_and_aligns_mask(tmp_path): + file_service = FileService(tmp_path) + image_path = tmp_path / "source.png" + mask_path = tmp_path / "mask.png" + Image.new("RGB", (4, 4)).save(image_path) + Image.new("RGB", (2, 2)).save(mask_path) + submit = AsyncMock(return_value={"status": "success", "save_result_path": "output.png"}) + inference = SimpleNamespace(worker=SimpleNamespace(runner=SimpleNamespace(config={})), submit_task_async=submit) + service = ImageGenerationService(file_service, inference) + message = ImageTaskRequest(image_path=str(image_path), image_mask_path=base64.b64encode(mask_path.read_bytes()).decode("utf-8")) + + asyncio.run(service.generate_with_stop_event(message, asyncio.Event())) + + payload = submit.await_args.args[0] + files = list(Path(payload["image_path"]).glob("*.png")) + assert len(files) == 2 + assert "image_mask_path" not in payload + for path in files: + with Image.open(path) as image: + assert image.size == (4, 4) diff --git a/tests/test_server_schema_fields.py b/tests/test_server_schema_fields.py new file mode 100644 index 000000000..8818c92c9 --- /dev/null +++ b/tests/test_server_schema_fields.py @@ -0,0 +1,253 @@ +import asyncio +import json +from dataclasses import fields +from importlib import import_module +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.cosmos3.cosmos3_runner import Cosmos3Runner +from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.flux2.flux2_runner import Flux2Runner +from lightx2v.models.runners.hidream_o1_image.hidream_o1_image_runner import HidreamO1ImageRunner +from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner +from lightx2v.models.runners.ltx2.ltx2_runner import LTX2Runner +from lightx2v.models.runners.motus.motus_runner import MotusRunner +from lightx2v.models.runners.runner_factory import RUNNER_MODULES +from lightx2v.models.runners.wan.wan_animate2_runner import WanAnimate2Runner +from lightx2v.models.runners.wan.wan_animate_runner import WanAnimateRunner +from lightx2v.models.runners.wan.wan_runner import LingbotRunner, WanRunner +from lightx2v.models.runners.wan.wan_s2v_runner import WanS2VRunner +from lightx2v.models.runners.wan.wan_vace_runner import WanVaceRunner +from lightx2v.server.api.tasks import image as image_api +from lightx2v.server.api.tasks import video as video_api +from lightx2v.server.schema import ImageTaskRequest, SenseNovaVisionTaskRequest, TalkObject, VideoTaskRequest +from lightx2v.server.services.file_service import FileService +from lightx2v.server.services.generation.image import ImageGenerationService +from lightx2v.server.services.generation.video import VideoGenerationService +from lightx2v.server.services.inference.worker import TorchrunInferenceWorker +from lightx2v.utils.input_info import INPUT_INFO_TYPES +from lightx2v.utils.registry_factory import RUNNER_REGISTER + + +@pytest.mark.parametrize("model_cls,module", RUNNER_MODULES.items()) +def test_native_schema_covers_runner_fields(monkeypatch, model_cls, module): + monkeypatch.setattr(DefaultRunner, "__del__", lambda self: None) + import_module(module) + runner_cls = RUNNER_REGISTER[model_cls] + image_tasks = {"t2i", "i2i", "ti2i"} + video_tasks = {"t2v", "i2v", "flf2v", "vace", "s2v", "rs2v", "animate", "sr", "t2av", "i2av", "l2av", "fl2av", "ref2av", "i2va", "v2av", "ltx2_s2v"} + internal_fields = {"return_result_tensor", "policy_image", "policy_state"} + variants = ( + {}, + {"inpaint_mask_enabled": True}, + {"image_size": [512, 512]}, + {"video_super_resolution": {"enable_cfg": True}}, + {"layered": True}, + {"seq_parallel": True}, + {"replace_flag": True}, + {"dancer_stage": "local"}, + {"prompt_travel": {"prompt_travel_text": ["scene"]}}, + {"self_attn_1_type": "radial_attn"}, + {"resize_mode": "fixed_shape"}, + ) + for task in runner_cls.supported_request_fields_by_task: + input_info_cls = runner_cls.input_info_cls_by_task.get(task) or INPUT_INFO_TYPES[task] + input_fields = {field.name for field in fields(input_info_cls)} + if task in image_tasks: + schema = ImageTaskRequest + elif task in video_tasks: + schema = VideoTaskRequest + else: + schema = None + for enable_cfg in (False, True): + for variant in variants: + runner = object.__new__(runner_cls) + BaseRunner.__init__(runner, {"task": task, "enable_cfg": enable_cfg, **variant}) + request_fields = runner.get_supported_request_fields(task) + assert not request_fields - input_fields, (model_cls, task, variant) + if schema is not None: + assert not request_fields - schema.model_fields.keys() - internal_fields, (model_cls, task, variant) + + +@pytest.mark.parametrize( + "runner_cls,task,payload", + [ + (WanS2VRunner, "s2v", {"src_pose_path": "pose.mp4"}), + (WanAnimateRunner, "animate", {"src_pose_path": "pose.mp4", "src_face_path": "face.mp4", "src_ref_images": ["reference.png"]}), + (WanAnimate2Runner, "animate", {"prompt_ref": "reference motion", "video_path": "driving.mp4", "src_pose_path": "pose.mp4"}), + (WanVaceRunner, "vace", {"video_path": "source.mp4", "mask_path": "mask.mp4", "src_ref_images": ["reference.png"]}), + ( + LTX2Runner, + "v2av", + { + "image_path": "first.png,last.png", + "image_strength": [0.0, 0.8], + "image_frame_idx": [0, 16], + "reference_video_strength": 0.5, + "reference_video_frame_cap": 17, + "mux_audio_video_path": "audio.mp4", + "video_path": "reference.mp4", + }, + ), + (LingbotRunner, "i2v", {"image_path": "input.png", "pose": "w-1", "action_path": "controls"}), + ( + Cosmos3Runner, + "i2va", + {"action_mode": "policy", "domain_name": "agibotworld", "view_point": "front", "action_path": "actions.npy", "state_path": "state.npy", "save_action_path": "actions.json"}, + ), + (MotusRunner, "i2v", {"image_path": "input.png", "state_path": "state.json", "save_action_path": "actions.json"}), + (HidreamO1ImageRunner, "i2i", {"keep_original_aspect": False, "layout_bboxes": '[{"bbox": [0, 0, 100, 100]}]'}), + (HunyuanImage3Runner, "i2i", {"infer_align_image_size": False}), + ], +) +def test_service_preserves_model_request_fields(tmp_path, monkeypatch, runner_cls, task, payload): + monkeypatch.setattr("lightx2v.models.runners.base_runner.seed_all", lambda seed: None) + runner = object.__new__(runner_cls) + BaseRunner.__init__(runner, {"task": task, "enable_cfg": True, "vae_stride": (4, 8, 8)}) + runner.run_pipeline = Mock(return_value=None) + worker = TorchrunInferenceWorker() + worker.rank = 0 + worker.world_size = 1 + worker.runner = runner + inference = SimpleNamespace(worker=worker, submit_task_async=worker.process_request) + is_image = task == "i2i" + schema = ImageTaskRequest if is_image else VideoTaskRequest + service_cls = ImageGenerationService if is_image else VideoGenerationService + service = service_cls(FileService(tmp_path), inference) + message = schema(task=task, **payload) + + response = asyncio.run(service.generate_with_stop_event(message, asyncio.Event())) + + assert response.task_status == "completed" + input_info = runner.run_pipeline.call_args.args[0] + for field, value in payload.items(): + assert getattr(input_info, field) == (",".join(value) if field == "src_ref_images" else value) + + +@pytest.mark.parametrize("field", ["src_pose_path", "image_strength"]) +def test_known_but_unsupported_fields_reach_runner_validation(tmp_path, monkeypatch, field): + runner = object.__new__(WanRunner) + BaseRunner.__init__(runner, {"task": "t2v", "vae_stride": (4, 8, 8)}) + message = VideoTaskRequest(**{field: "pose.mp4" if field == "src_pose_path" else 0.5}) + service = VideoGenerationService(FileService(tmp_path), None) + request_data = service.prepare_task_data(message) + request_data.pop("task_id") + + with pytest.raises(ValueError, match=field): + runner.prepare_request(request_data) + + +@pytest.fixture +def native_api(tmp_path, monkeypatch): + def build(api): + messages = [] + file_service = FileService(tmp_path) + monkeypatch.setattr(api, "get_services", lambda: SimpleNamespace(file_service=file_service)) + monkeypatch.setattr(api.task_manager, "create_task", lambda message: messages.append(message) or message.task_id) + app = FastAPI() + app.include_router(api.router) + return TestClient(app), messages + + return build + + +@pytest.mark.parametrize("encoding", ["json", "form"]) +@pytest.mark.parametrize( + "api,payload", + [ + (video_api, {"src_pose_path": "pose.mp4", "src_face_path": "face.mp4", "src_bg_path": "background.mp4", "mask_path": "mask.mp4", "prompt_ref": "motion"}), + (video_api, {"image_strength": [0.0, 0.8], "image_frame_idx": [0, 16], "reference_video_strength": 0.0, "reference_video_frame_cap": 17, "mux_audio_video_path": "audio.mp4"}), + (video_api, {"pose": {"0": {"camera": "forward"}}, "target_shape": [480, 832], "src_ref_images": ["reference.png"], "num_frames": 17, "video_duration": 2.5}), + (video_api, {"action_mode": "policy", "action_path": "actions.npy", "state_path": "state.json", "save_action_path": "output.npy", "domain_name": "robot", "view_point": "front"}), + (image_api, {"keep_original_aspect": False, "layout_bboxes": '[{"bbox": [0, 0, 1, 1]}]', "infer_align_image_size": False, "target_shape": [512, 768], "prompt": "[keep this text]"}), + ], +) +def test_http_preserves_schema_fields(native_api, encoding, api, payload): + client, messages = native_api(api) + with client: + if encoding == "json": + response = client.post("/", json=payload) + else: + form = {key: json.dumps(value) if isinstance(value, (list, dict)) else str(value) for key, value in payload.items()} + response = client.post("/form", data=form) + + assert response.status_code == 200, response.text + for field, value in payload.items(): + field = "target_video_length" if field == "num_frames" else field + assert getattr(messages[0], field) == value + assert field in messages[0].model_fields_set + assert messages[0].save_result_path is None + + +@pytest.mark.parametrize("encoding", ["json", "form"]) +@pytest.mark.parametrize("inpaint_enabled", [False, True]) +def test_http_flux2_inpaint_fields_follow_deployment(native_api, tmp_path, monkeypatch, encoding, inpaint_enabled): + client, messages = native_api(image_api) + payload = {"inpaint_blur_sigma": 0.5, "inpaint_blur_size": 3} + with client: + response = client.post("/", json=payload) if encoding == "json" else client.post("/form", data=payload) + assert response.status_code == 200, response.text + + monkeypatch.setattr("lightx2v.models.runners.base_runner.seed_all", lambda seed: None) + runner = object.__new__(Flux2Runner) + BaseRunner.__init__(runner, {"task": "i2i", "inpaint_mask_enabled": inpaint_enabled}) + runner.run_pipeline = Mock(return_value=None) + worker = TorchrunInferenceWorker() + worker.rank = 0 + worker.world_size = 1 + worker.runner = runner + inference = SimpleNamespace(worker=worker, submit_task_async=worker.process_request) + service = ImageGenerationService(FileService(tmp_path), inference) + + if inpaint_enabled: + response = asyncio.run(service.generate_with_stop_event(messages[0], asyncio.Event())) + assert response.task_status == "completed" + input_info = runner.run_pipeline.call_args.args[0] + assert input_info.inpaint_blur_sigma == 0.5 + assert input_info.inpaint_blur_size == 3 + else: + with pytest.raises(RuntimeError, match="inpaint_blur_sigma.*inpaint_blur_size") as error: + asyncio.run(service.generate_with_stop_event(messages[0], asyncio.Event())) + assert error.value.original_error_type == "ValueError" + runner.run_pipeline.assert_not_called() + + +@pytest.mark.parametrize("api", [image_api, video_api]) +@pytest.mark.parametrize("encoding", ["json", "form"]) +@pytest.mark.parametrize("field", ["unknown_option", "infer_steps", "warmup", "resize_mode", "return_result_tensor"]) +def test_http_rejects_unknown_fields(native_api, api, encoding, field): + client, messages = native_api(api) + with client: + response = client.post("/", json={field: True}) if encoding == "json" else client.post("/form", data={field: "true"}) + + assert response.status_code == 422, response.text + assert response.json()["detail"][0]["type"] == "extra_forbidden" + assert messages == [] + + +@pytest.mark.parametrize("api,payload", [(video_api, {"target_shape": "[480,"}), (image_api, {"keep_original_aspect": "maybe"})]) +def test_form_rejects_invalid_values(native_api, api, payload): + client, messages = native_api(api) + with client: + response = client.post("/form", data=payload) + assert response.status_code == 422, response.text + assert not messages + + +def test_form_rejects_unknown_upload(native_api): + client, messages = native_api(video_api) + with client: + response = client.post("/form", files={"unknown_file": ("file.mp4", b"video", "video/mp4")}) + assert response.status_code == 422, response.text + assert not messages + + +@pytest.mark.parametrize("schema,payload", [(TalkObject, {"audio": "audio.wav", "mask": "mask.png"}), (SenseNovaVisionTaskRequest, {"task": "depth"})]) +def test_nested_and_special_requests_reject_unknown_fields(schema, payload): + with pytest.raises(ValueError, match="extra_forbidden"): + schema(**payload, unknown_option=True) diff --git a/tests/test_shot_request_defaults.py b/tests/test_shot_request_defaults.py new file mode 100644 index 000000000..6a5d6905d --- /dev/null +++ b/tests/test_shot_request_defaults.py @@ -0,0 +1,115 @@ +import json +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.shot_runner import rs2v_infer, stream_infer +from lightx2v.shot_runner.shot_base import ShotPipeline, load_clip_configs + + +@pytest.mark.parametrize("task", ["s2v", "rs2v"]) +def test_shot_loads_startup_task_from_clip_json(tmp_path, task): + model_path = tmp_path / "model" + model_path.mkdir() + (model_path / "config.json").write_text('{"hidden_size": 256}') + (tmp_path / "clip.json").write_text(json.dumps({"model_cls": "seko_talk", "model_path": str(model_path), "task": task, "target_video_length": 81})) + main_path = tmp_path / "main.json" + main_path.write_text(json.dumps({"lightx2v_path": str(tmp_path), "clip_configs": [{"name": "clip", "path": "clip.json"}]})) + + clips = load_clip_configs(main_path) + + assert len(clips) == 1 + assert clips[0].name == "clip" + assert clips[0].config_json["task"] == task + assert clips[0].config_json["hidden_size"] == 256 + assert clips[0].config_json["target_video_length"] == 81 + assert clips[0].config_json["warmup"] is False + + +@pytest.mark.parametrize("request_fields", [{}, {"seed": 0, "target_shape": [480, 832], "save_result_path": "chosen.mp4"}]) +def test_shot_uses_flat_specs_and_caller_content(request_fields): + pipeline = object.__new__(ShotPipeline) + config = { + "task": "rs2v", + "target_video_length": 81, + "target_shape": [720, 1280], + } + args = SimpleNamespace(prompt="request", image_path="input.png", audio_path="input.wav", **request_fields) + + result = pipeline.prepare_input_info(args, config) + + assert result.task == "rs2v" + assert result.prompt == "request" + assert result.negative_prompt == "" + assert result.image_path == "input.png" + assert result.audio_path == "input.wav" + assert result.seed == request_fields.get("seed") + assert result.save_result_path == request_fields.get("save_result_path") + assert result.target_shape == request_fields.get("target_shape", [720, 1280]) + assert result.target_video_length == 81 + + +@pytest.mark.parametrize(("module", "pipeline_name"), [(rs2v_infer, "ShotRS2VPipeline"), (stream_infer, "ShotStreamPipeline")]) +@pytest.mark.parametrize("explicit_shape", [False, True]) +def test_shot_cli_keeps_json_shape_unless_explicitly_overridden(monkeypatch, module, pipeline_name, explicit_shape): + config = {"task": "rs2v" if module is rs2v_infer else "s2v", "target_shape": [720, 1280]} + pipeline = object.__new__(ShotPipeline) + inputs = [] + pipeline.generate = lambda args: inputs.append(pipeline.prepare_input_info(args, config)) + monkeypatch.setattr(module, "load_clip_configs", lambda path: []) + monkeypatch.setattr(module, pipeline_name, lambda configs: pipeline) + argv = ["shot", "--config_json", "deployment.json"] + if explicit_shape: + argv += ["--target_shape", "480", "832"] + monkeypatch.setattr("sys.argv", argv) + + module.main() + + assert inputs[0].target_shape == ([480, 832] if explicit_shape else [720, 1280]) + assert inputs[0].save_result_path is None + + +@pytest.mark.parametrize(("request_fields", "expected_frames"), [({}, 81), ({"video_duration": 1.0}, 17), ({"video_duration": 1.0, "target_video_length": 49}, 49)]) +def test_rs2v_keeps_duration_and_explicit_frame_priority(monkeypatch, request_fields, expected_frames): + pipeline = object.__new__(rs2v_infer.ShotRS2VPipeline) + runner = object.__new__(BaseRunner) + runner.config = {"task": "rs2v", "target_video_length": 81, "target_fps": 16, "audio_sr": 16000, "vae_stride": (4, 8, 8)} + observed_frames = [] + runner._run_input_encoder_local_rs2v_static = lambda: observed_frames.append(runner.input_info.target_video_length) + runner._run_input_encoder_local_rs2v_dynamic = lambda: None + runner.check_stop = lambda: None + runner.run_clip_main = lambda: (torch.zeros(1, 3, runner.input_info.target_video_length, 1, 1), runner.input_info.audio_clip, torch.zeros(1, 1, 1, 1)) + pipeline.clip_generators = {"rs2v_clip": runner} + pipeline.progress_callback = None + monkeypatch.setattr(rs2v_infer, "load_audio_file", lambda path: (torch.zeros(1, 16000), 16000)) + save = Mock() + monkeypatch.setattr(rs2v_infer, "save_to_video", save) + args = SimpleNamespace(prompt="request", image_path="input.png", audio_path="input.wav", **request_fields) + + pipeline.generate(args) + + assert observed_frames == [expected_frames] + assert runner.input_info.seed == 42 + assert runner.input_info.save_result_path is None + save.assert_not_called() + + +def test_shot_stream_omitted_output_does_not_write_temporary_files(monkeypatch): + pipeline = object.__new__(stream_infer.ShotStreamPipeline) + config = {"task": "s2v", "target_video_length": 33, "prev_frame_length": 1, "target_fps": 16, "audio_sr": 16000} + runner = object.__new__(BaseRunner) + runner.config = config + runner.prev_frame_length = 1 + runner.run_clip_pipeline = lambda inputs: (torch.zeros(1, 3, 33, 1, 1), inputs.audio_clip, None) + pipeline.clip_generators = {"s2v_clip": runner, "f2v_clip": runner} + monkeypatch.setattr(stream_infer, "load_audio_file", lambda path: (torch.zeros(1, 33000), 16000)) + save = Mock() + monkeypatch.setattr(stream_infer, "save_to_video", save) + monkeypatch.setattr(stream_infer, "save_audio", save) + + pipeline.generate(SimpleNamespace(prompt="request", image_path="input.png", audio_path="input.wav")) + + save.assert_not_called() diff --git a/tests/test_startup_config.py b/tests/test_startup_config.py new file mode 100644 index 000000000..3cbf7c678 --- /dev/null +++ b/tests/test_startup_config.py @@ -0,0 +1,254 @@ +import ast +import builtins +import json +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from lightx2v.models.runners.bagel.bagel_runner import BagelRunner +from lightx2v.models.runners.bagel.sensenova_vision_runner import SenseNovaVisionRunner +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.utils import set_config + + +@pytest.mark.parametrize("model_cls", ["cosmos3", "lingbot_va", "fastwam"]) +def test_ros_policy_uses_current_startup_api(tmp_path, model_cls): + root = Path(__file__).resolve().parents[1] + path = root / "lightx2v_ros/src/inference/inference" / f"{model_cls}_node/main.py" + tree = ast.parse(path.read_text()) + imports = [node for node in tree.body if isinstance(node, ast.ImportFrom) and node.module == "lightx2v.utils.set_config"] + cls = next(node for node in tree.body if isinstance(node, ast.ClassDef)) + method = next(node for node in cls.body if isinstance(node, ast.FunctionDef) and node.name in {"build_policy_config", "_build_policy_config"}) + namespace = {} + exec(compile(ast.Module(body=[*imports, method], type_ignores=[]), str(path), "exec"), namespace) + deployment = tmp_path / "deployment.json" + deployment.write_text('{"target_video_length": 81, "raw_action_dim": 8, "policy_prompt_format": "json"}') + values = {"config_json": str(deployment), "model_path": str(tmp_path), "prompt_format": "official_text", "seed": 7} + node = SimpleNamespace( + get_parameter=lambda key: SimpleNamespace(value=values[key]), + seed=7, + contract=SimpleNamespace(action_dim=8, state_dim=8, name="libero"), + get_logger=Mock(), + ) + + config = namespace[method.name](node) + + assert config["model_cls"] == model_cls + assert config["task"] == "i2va" + assert config["target_video_length"] == 81 + assert config["warmup"] is False + if model_cls == "cosmos3": + assert config["policy_prompt_format"] == "official_text" + + +@pytest.mark.parametrize( + ("preferred", "fallback"), + [("", "low_noise_model"), ("low_noise_model", "distill_models/low_noise_model"), ("distill_models/low_noise_model", "original"), ("original", "transformer")], +) +def test_checkpoint_directory_priority(tmp_path, preferred, fallback): + for subfolder, width in ((preferred, 128), (fallback, 256)): + folder = tmp_path / subfolder + folder.mkdir(parents=True, exist_ok=True) + (folder / "config.json").write_text(json.dumps({"hidden_size": width})) + deployment = tmp_path / "deployment.json" + deployment.write_text('{"hidden_size": 512, "target_video_length": 81, "infer_steps": 4}') + + config = set_config.build_startup_config({"model_cls": "wan2.1", "model_path": str(tmp_path), "task": "t2v", "config_json": str(deployment)}) + + assert config["hidden_size"] == 128 + assert config["target_video_length"] == 81 + assert config["infer_steps"] == 4 + + +@pytest.mark.parametrize("model_cls", ["hunyuan_video_1.5", "worldplay_ar", "worldplay_bi", "worldplay_distill"]) +def test_named_transformer_directory(tmp_path, model_cls): + transformer_path = tmp_path / "transformer" / "480p_i2v" + transformer_path.mkdir(parents=True) + (transformer_path / "config.json").write_text('{"hidden_size": 128}') + (tmp_path / "config.json").write_text('{"hidden_size": 256}') + + config = set_config.build_startup_config({"model_cls": model_cls, "model_path": str(tmp_path), "task": "i2v", "transformer_model_name": "480p_i2v"}) + + assert config["hidden_size"] == 128 + assert config["transformer_model_path"] == str(transformer_path) + + +def test_longcat_merges_both_model_configs(tmp_path): + (tmp_path / "config.json").write_text('{"root_setting": 1, "hidden_size": 128}') + (tmp_path / "transformer").mkdir() + (tmp_path / "transformer" / "config.json").write_text('{"transformer_setting": 2, "hidden_size": 256}') + + config = set_config.build_startup_config({"model_cls": "longcat_image", "model_path": str(tmp_path), "task": "t2i"}) + + assert config["root_setting"] == 1 + assert config["transformer_setting"] == 2 + assert config["hidden_size"] == 256 + + +@pytest.mark.parametrize("subfolder", ["", "transformer"]) +def test_ltx_preserves_runtime_rope_selector(tmp_path, subfolder): + folder = tmp_path / subfolder + folder.mkdir(exist_ok=True) + (folder / "config.json").write_text('{"rope_type": "split", "hidden_size": 128}') + + config = set_config.build_startup_config({"model_cls": "ltx2", "model_path": str(tmp_path), "task": "t2av", "rope_type": "torch"}) + + assert config["rope_type"] == "torch" + assert config["hidden_size"] == 128 + + +def test_cosmos_reads_vae_once_and_uses_its_scale_factors(tmp_path, monkeypatch): + vae_path = tmp_path / "vae" / "config.json" + vae_path.parent.mkdir() + vae_path.write_text('{"block_out_channels": [1, 2, 3], "scale_factor_spatial": 12, "scale_factor_temporal": 2}') + real_open = builtins.open + vae_reads = [] + + def record_open(path, *args, **kwargs): + if str(path) == str(vae_path): + vae_reads.append(path) + return real_open(path, *args, **kwargs) + + monkeypatch.setattr(builtins, "open", record_open) + + config = set_config.build_startup_config({"model_cls": "cosmos3", "model_path": str(tmp_path), "task": "i2va"}) + + assert len(vae_reads) == 1 + assert config["vae_scale_factor"] == 12 + assert config["vae_scale_factor_spatial"] == 12 + assert config["vae_scale_factor_temporal"] == 2 + + +@pytest.mark.parametrize(("task", "subfolder"), [("t2av", "transformer"), ("ref2av", "transformer_ref")]) +def test_minimax_selects_task_weights_and_keeps_spatial_scale(tmp_path, task, subfolder): + folder = tmp_path / subfolder + folder.mkdir() + (folder / "config.json").write_text("{}") + (tmp_path / "vae").mkdir() + (tmp_path / "vae" / "config.json").write_text('{"block_out_channels": [1, 2, 3, 4, 5, 6]}') + + config = set_config.build_startup_config({"model_cls": "minimax_h3", "model_path": str(tmp_path), "task": task}) + + assert config["dit_original_ckpt"] == str(folder) + assert config["vae_scale_factor"] == 16 + assert config["vae_spatial_scale_factor"] == 16 + assert config["enable_cfg"] is False + + +@pytest.mark.parametrize( + ("task", "subtask", "json_subtask", "expected", "error"), + [ + ("omni_vision_task", None, None, None, "Unsupported omni-vision subtask"), + ("omni_vision_task", "", "recon3d", None, "Unsupported omni-vision subtask"), + ("omni_vision_task", "invalid", None, None, "Unsupported omni-vision subtask"), + ("omni_vision_task", "recon3d", None, "recon3d", None), + ("omni_vision_task", None, "recon3d", "recon3d", None), + ("omni_vision_task", "depth", "recon3d", "depth", None), + ("t2i", "recon3d", None, None, "does not support request fields: omni_vision_subtask"), + ("t2i", "", None, None, "does not support request fields: omni_vision_subtask"), + ], +) +def test_cli_leaves_subtask_validation_to_runner(tmp_path, task, subtask, json_subtask, expected, error): + for filename in ("llm_config.json", "vit_config.json"): + (tmp_path / filename).write_text("{}") + deployment = tmp_path / "deployment.json" + deployment.write_text(json.dumps({"omni_vision_subtask": json_subtask} if json_subtask else {})) + args = SimpleNamespace( + model_cls="sensenova_vision" if task == "omni_vision_task" else "bagel", + model_path=str(tmp_path), + config_json=str(deployment), + task=task, + omni_vision_subtask=subtask, + ) + + config, request = set_config.build_cli_inputs(args) + runner_cls = SenseNovaVisionRunner if task == "omni_vision_task" else BagelRunner + runner = object.__new__(runner_cls) + BaseRunner.__init__(runner, config) + + if error: + with pytest.raises(ValueError, match=error): + runner.prepare_request(request) + else: + input_info = runner.prepare_request(request) + assert input_info.task == task + assert input_info.omni_vision_subtask == expected + + +def test_sensenova_normalizes_each_request_without_changing_defaults(): + config = {"task": "omni_vision_task", "omni_vision_subtask": "recon3d"} + runner = object.__new__(SenseNovaVisionRunner) + BaseRunner.__init__(runner, config) + request = {"omni_vision_subtask": " RAW-QUERY "} + + explicit = runner.prepare_request(request) + default = runner.prepare_request({}) + + assert explicit.omni_vision_subtask == "understanding" + assert default.omni_vision_subtask == "recon3d" + assert config["omni_vision_subtask"] == "recon3d" + assert request == {"omni_vision_subtask": " RAW-QUERY "} + assert runner.input_info is None + + +@pytest.mark.parametrize( + ("tp", "cfg", "sp", "enable_cfg", "shape", "names"), + [ + (1, 1, 1, False, (1, 1), ("cfg_p", "seq_p")), + (1, 2, 2, True, (2, 2), ("cfg_p", "seq_p")), + (2, 1, 1, False, (2,), ("tensor_p",)), + (2, 2, 4, True, (2, 4, 2), ("cfg_p", "seq_p", "tensor_p")), + (2, 2, 1, False, (2, 2), ("cfg_p", "tensor_p")), + ], +) +def test_parallel_topology_and_flags(monkeypatch, tp, cfg, sp, enable_cfg, shape, names): + mesh = Mock() + reduce = Mock() + monkeypatch.setattr(set_config, "AI_DEVICE", "cpu") + monkeypatch.setattr(set_config, "init_device_mesh", mesh) + monkeypatch.setattr(set_config.dist, "get_world_size", lambda: tp * cfg * sp) + monkeypatch.setattr(set_config.dist, "all_reduce", reduce) + config = {"parallel": {"tensor_p_size": tp, "cfg_p_size": cfg, "seq_p_size": sp}, "enable_cfg": enable_cfg} + + set_config.init_parallel(config) + + mesh.assert_called_once_with("cpu", shape, mesh_dim_names=names) + assert config["tensor_parallel"] == (tp > 1) + assert config["seq_parallel"] == (sp > 1) + assert config["cfg_parallel"] == (enable_cfg and cfg > 1) + reduce.assert_called_once() + + +def test_parallel_rejects_mismatched_world_size(monkeypatch): + mesh = Mock() + monkeypatch.setattr(set_config, "init_device_mesh", mesh) + monkeypatch.setattr(set_config.dist, "get_world_size", lambda: 2) + + with pytest.raises(ValueError, match="Parallel sizes must match"): + set_config.init_parallel({"parallel": {"tensor_p_size": 4}}) + + mesh.assert_not_called() + + +def test_phase_parallel_keeps_model_specific_flags(monkeypatch): + from lightx2v.models.networks.hunyuan_image3 import parallel + + def initialize(config): + config.update({"tensor_parallel": True, "seq_parallel": True, "cfg_parallel": False, "device_mesh": "phase_mesh"}) + + monkeypatch.setattr(parallel, "initialize_hunyuan_image3_parallel_runtime", initialize) + monkeypatch.setattr(set_config, "AI_DEVICE", "cpu") + monkeypatch.setattr(set_config.dist, "get_world_size", lambda: 4) + reduce = Mock() + monkeypatch.setattr(set_config.dist, "all_reduce", reduce) + config = {"model_cls": "hunyuan_image3", "parallel": {"phase_aware": True, "cfg_p_size": 4}, "enable_cfg": True} + + set_config.init_parallel(config) + + assert config["device_mesh"] == "phase_mesh" + assert config["tensor_parallel"] is True + assert config["seq_parallel"] is True + assert config["cfg_parallel"] is False + reduce.assert_called_once() diff --git a/tests/test_wan22_audio_restore.py b/tests/test_wan22_audio_restore.py new file mode 100644 index 000000000..49f9a88c3 --- /dev/null +++ b/tests/test_wan22_audio_restore.py @@ -0,0 +1,135 @@ +import inspect +import json +from pathlib import Path +from unittest.mock import Mock + +import pytest +import torch + +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.wan import wan_audio_runner +from lightx2v.models.runners.wan.wan_audio_runner import Wan22AudioRunner, WanAudioRunner +from lightx2v.models.schedulers.wan.audio import scheduler as audio_scheduler +from lightx2v.pipeline import LightX2VPipeline +from lightx2v.utils.lockable_dict import LockableDict +from lightx2v.utils.set_config import build_startup_config + + +def test_legacy_i2v_initializes_and_runs_audio_encoder(monkeypatch): + runner = object.__new__(Wan22AudioRunner) + BaseRunner.__init__( + runner, + LockableDict({"model_cls": "wan2.2_audio", "task": "i2v", "vae_stride": [4, 16, 16], "lazy_load": True, "cpu_offload": True, "use_image_encoder": False}), + ) + runner.init_modules() + runner.input_info = runner.prepare_request({"image_path": "portrait.png", "audio_path": "speech.wav"}) + runner.read_image_input = Mock(return_value=("image", [48, 5, 4, 4], [64, 64])) + runner.run_vae_encoder = Mock(return_value="image latents") + runner.read_audio_input = Mock(return_value=(["audio segment"], 17, None, 1)) + runner.run_text_encoder = Mock(return_value="text embeddings") + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + + result = runner.run_input_encoder() + + runner.read_image_input.assert_called_once_with("portrait.png") + runner.read_audio_input.assert_called_once_with("speech.wav") + assert result["audio_segments"] == ["audio segment"] + assert result["image_encoder_output"]["vae_encoder_out"] == "image latents" + assert runner.input_info.audio_num == 1 + assert runner.input_info.seed == 42 + + +@pytest.mark.parametrize("scheduler_cls", [audio_scheduler.EulerScheduler, audio_scheduler.ConsistencyModelScheduler]) +def test_audio_scheduler_preserves_previous_frames(monkeypatch, scheduler_cls): + monkeypatch.setattr(audio_scheduler, "AI_DEVICE", "cpu") + scheduler = scheduler_cls({"model_cls": "wan2.2_audio", "sample_shift": 1, "seq_parallel": False, "sample_guide_scale": 1, "dim": 8, "num_heads": 1, "parallel": None}) + shape = (2, 3, 4, 4) + scheduler.prepare(42, shape, infer_steps=2) + previous_latents = torch.full(shape, 7.0) + scheduler.reset(42, shape, {"prev_latents": previous_latents, "prev_len": 1}) + + torch.testing.assert_close(scheduler.latents[:, :1], previous_latents[:, :1]) + new_frames = scheduler.latents[:, 1:].clone() + scheduler.noise_pred = torch.ones(shape) + scheduler.step_post() + + torch.testing.assert_close(scheduler.latents[:, :1], previous_latents[:, :1]) + assert not torch.equal(scheduler.latents[:, 1:], new_frames) + + scheduler.clear() + assert scheduler.prev_latents is None + assert scheduler.prev_len == 0 + + scheduler.prepare(42, shape, infer_steps=2) + fresh_scheduler = scheduler_cls(scheduler.config) + fresh_scheduler.prepare(42, shape, infer_steps=2) + torch.testing.assert_close(scheduler.latents, fresh_scheduler.latents) + + scheduler.clear() + scheduler.prepare(42, (2, 3, 8, 8), infer_steps=2) + assert scheduler.latents.shape == (2, 3, 8, 8) + + +def test_rs2v_is_not_advertised_by_wan22_audio(): + runner = object.__new__(Wan22AudioRunner) + with pytest.raises(ValueError, match="Wan22AudioRunner does not support task 'rs2v'"): + BaseRunner.__init__(runner, {"model_cls": "wan2.2_audio", "task": "rs2v"}) + + seko_runner = object.__new__(WanAudioRunner) + BaseRunner.__init__(seko_runner, {"model_cls": "seko_talk", "task": "rs2v", "vae_stride": [4, 8, 8]}) + input_info = seko_runner.prepare_request({"image_path": "portrait.png", "audio_path": "speech.wav"}) + assert input_info.task == "rs2v" + + +@pytest.mark.parametrize("cpu_offload", [False, True]) +def test_vae_loading_preserves_independent_instances_and_options(tmp_path, monkeypatch, cpu_offload): + vae_path = tmp_path / "Wan2.2_VAE.pth" + vae_path.touch() + runner = object.__new__(Wan22AudioRunner) + BaseRunner.__init__(runner, {"model_cls": "wan2.2_audio", "task": "s2v", "model_path": str(tmp_path), "cpu_offload": cpu_offload, "vae_offload_cache": True, "dummy_model": True}) + signature = inspect.signature(wan_audio_runner.Wan2_2_VAE) + instances = [object(), object()] + constructor = Mock(side_effect=instances) + monkeypatch.setattr(wan_audio_runner, "Wan2_2_VAE", constructor) + + encoder, decoder = runner.load_vae() + + assert encoder is instances[0] + assert decoder is instances[1] + assert constructor.call_count == 2 + for call in constructor.call_args_list: + arguments = signature.bind(**call.kwargs) + arguments.apply_defaults() + assert arguments.arguments["vae_path"] == str(vae_path) + assert arguments.arguments["dtype"] == torch.float32 + assert arguments.arguments["device"] == torch.device("cpu" if cpu_offload else wan_audio_runner.AI_DEVICE) + assert arguments.arguments["cpu_offload"] is cpu_offload + assert arguments.arguments["offload_cache"] is True + assert arguments.arguments["dummy_model"] is True + + +@pytest.mark.parametrize("python_entry", [False, True]) +@pytest.mark.parametrize("overrides", [{}, {"vae_stride": [4, 32, 32], "num_channels_latents": 64, "use_image_encoder": True}]) +def test_wan22_audio_startup_uses_json_settings(tmp_path, python_entry, overrides): + startup = {"model_cls": "wan2.2_audio", "task": "s2v", "model_path": str(tmp_path)} + if python_entry: + startup = LightX2VPipeline(**startup).startup_config + preset = Path(__file__).resolve().parents[1] / "configs/seko_talk/seko_talk_08_5B_base.json" + settings = json.loads(preset.read_text()) + settings.update(overrides) + config_path = tmp_path / "deployment.json" + config_path.write_text(json.dumps(settings)) + config = build_startup_config({**startup, "config_json": str(config_path)}) + + assert config["vae_stride"] == settings["vae_stride"] + assert config["num_channels_latents"] == settings["num_channels_latents"] + assert config["use_image_encoder"] is settings["use_image_encoder"] + + +def test_wan22_audio_pipeline_defaults_without_json(tmp_path): + pipeline = LightX2VPipeline(model_cls="wan2.2_audio", task="s2v", model_path=str(tmp_path)) + config = build_startup_config(pipeline.startup_config) + + assert config["vae_stride"] == (4, 16, 16) + assert config["num_channels_latents"] == 48 + assert config["use_image_encoder"] is False diff --git a/tests/test_wan_distill_selection.py b/tests/test_wan_distill_selection.py new file mode 100644 index 000000000..61ccb91dc --- /dev/null +++ b/tests/test_wan_distill_selection.py @@ -0,0 +1,129 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightx2v.models.networks.wan.distill_model import WanDistillModel +from lightx2v.models.networks.wan.model import WanModel +from lightx2v.models.runners.wan.wan_runner import MultiModelStruct, WanRunner, get_wan_model_class +from lightx2v.models.schedulers.wan.scheduler import WanScheduler +from lightx2v.models.schedulers.wan.scheduler_factory import create_wan_scheduler, get_wan_distill_method +from lightx2v.models.schedulers.wan.step_distill.scheduler import Wan21MeanFlowStepDistillScheduler, WanStepDistillScheduler +from lightx2v.utils.input_info import T2VInputInfo +from lightx2v.utils.lockable_dict import LockableDict + + +def scheduler_config(**updates): + config = { + "model_cls": "wan2.1", + "feature_caching": "NoCaching", + "infer_steps": 4, + "target_video_length": 81, + "sample_shift": 5.0, + "sample_guide_scale": [3.5, 3.5], + "seq_parallel": False, + "dim": 5120, + "num_heads": 40, + "denoising_step_list": [1000, 750, 500, 250], + } + config.update(updates) + return config + + +def test_scheduler_selection(): + standard = create_wan_scheduler(scheduler_config()) + dmd2 = create_wan_scheduler(scheduler_config(model_cls="wan2.2_moe", distill_method="dmd2")) + mean_flow = create_wan_scheduler(scheduler_config(distill_method="mean_flow")) + + assert type(standard) is WanScheduler + assert isinstance(dmd2, WanStepDistillScheduler) + assert isinstance(mean_flow, Wan21MeanFlowStepDistillScheduler) + + +def test_distill_method_validation(): + assert get_wan_distill_method({"model_cls": "wan2.1"}) is None + assert get_wan_distill_method({"model_cls": "wan2.1", "distill_method": "mean_flow"}) == "mean_flow" + assert get_wan_distill_method({"model_cls": "wan2.2_moe", "distill_method": "dmd2"}) == "dmd2" + + with pytest.raises(NotImplementedError, match="wan2.2_moe does not support distill_method 'mean_flow'"): + get_wan_distill_method({"model_cls": "wan2.2_moe", "distill_method": "mean_flow"}) + with pytest.raises(NotImplementedError, match="wan2.1_vace does not support distill_method 'dmd2'"): + get_wan_distill_method({"model_cls": "wan2.1_vace", "distill_method": "dmd2"}) + with pytest.raises(NotImplementedError, match="wan2.1 does not support distill_method 'dmd-2'"): + create_wan_scheduler(scheduler_config(distill_method="dmd-2")) + + +def test_wan21_distill_scheduler_selection(): + runner = object.__new__(WanRunner) + runner.config = scheduler_config(model_cls="wan2.1", distill_method="mean_flow") + runner.init_scheduler() + assert isinstance(runner.scheduler, Wan21MeanFlowStepDistillScheduler) + + runner.config["distill_method"] = "unknown" + with pytest.raises(NotImplementedError, match="wan2.1 does not support distill_method 'unknown'"): + runner.init_scheduler() + + +def test_wan21_distill_model_selection(): + assert get_wan_model_class(None) is WanModel + assert get_wan_model_class("dmd2") is WanDistillModel + + +def test_multimodel_boundary_selection(): + normal = MultiModelStruct([None, None], {"boundary": 0.9}) + normal.scheduler = SimpleNamespace(step_index=0, timesteps=torch.tensor([1000, 750])) + assert normal.uses_high_noise_model() + assert normal.get_switch_step_index() == 1 + assert not hasattr(normal, "boundary_step_index") + + dmd2 = MultiModelStruct([None, None], {"model_cls": "wan2.2_moe", "distill_method": "dmd2", "boundary_step_index": 2}) + dmd2.scheduler = SimpleNamespace(step_index=1) + assert dmd2.uses_high_noise_model() + assert dmd2.get_switch_step_index() == 2 + assert not hasattr(dmd2, "boundary") + + +@pytest.mark.parametrize( + ("config", "field"), + [({}, "boundary"), ({"model_cls": "wan2.2_moe", "distill_method": "dmd2"}, "boundary_step_index")], +) +def test_multimodel_boundary_is_required(config, field): + with pytest.raises(KeyError, match=field): + MultiModelStruct([None, None], config) + + +def test_wan_input_info_excludes_infer_steps(): + input_info = T2VInputInfo() + input_info.update({"infer_steps": 8}) + + assert not hasattr(input_info, "infer_steps") + + +def test_wan_runner_keeps_scheduler_config_and_carries_request_frame_count(): + runner = object.__new__(WanRunner) + runner.config = LockableDict( + scheduler_config( + task="t2v", + target_height=720, + target_width=1280, + ) + ) + runner.supported_tasks = ("t2v",) + runner.scheduler = SimpleNamespace(infer_steps=4) + + first = runner.create_input_info({"task": runner.config["task"], "infer_steps": 8, "target_video_length": 49}) + second = runner.create_input_info({"task": runner.config["task"]}) + + assert runner.config["infer_steps"] == 4 + assert runner.scheduler.infer_steps == 4 + assert runner.config["target_video_length"] == 81 + assert first.target_video_length == 49 + assert second.target_video_length == 81 + + +def test_wan_runner_rejects_frame_length_cached_at_startup(): + runner = object.__new__(WanRunner) + runner.config = LockableDict(scheduler_config(task="t2v", target_height=720, target_width=1280, self_attn_1_type="svg_attn")) + runner.supported_tasks = ("t2v",) + with pytest.raises(ValueError, match="target_video_length"): + runner.prepare_request({"task": runner.config["task"], "target_video_length": 49}) diff --git a/tests/test_wan_request_reuse.py b/tests/test_wan_request_reuse.py new file mode 100644 index 000000000..3c3329770 --- /dev/null +++ b/tests/test_wan_request_reuse.py @@ -0,0 +1,129 @@ +import json +from functools import partial + +import pytest +import torch +from PIL import Image + +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.wan.wan_audio_runner import WanAudioARRunner +from lightx2v.models.runners.wan.wan_runner import Wan22DenseRunner, Wan22MoeRunner, WanRunner + + +@pytest.fixture +def cpu_runner(tmp_path, monkeypatch): + image_path = tmp_path / "square.png" + Image.new("RGB", (64, 64), "red").save(image_path) + config = { + "model_cls": "wan2.1", + "task": "i2v", + "target_height": 480, + "target_width": 832, + "target_video_length": 5, + "vae_stride": [4, 8, 8], + "patch_size": [1, 2, 2], + "enable_cfg": True, + "use_image_encoder": False, + "enable_reuse": True, + "reuse_cache_path": str(tmp_path / "cache"), + } + runner = object.__new__(WanRunner) + monkeypatch.setattr(runner, "set_init_device", lambda: setattr(runner, "init_device", torch.device("cpu"))) + monkeypatch.setattr(runner, "init_scheduler", lambda: None) + DefaultRunner.__init__(runner, config) + runner._gc_frozen = True + runner.run_input_encoder = runner._run_input_encoder_local_i2v + # Keep request preparation, image resizing, latent sizing and disk-cache IO real. + monkeypatch.setattr(runner, "get_vae_encoder_output", lambda *args: torch.ones(1)) + monkeypatch.setattr(runner, "run_text_encoder", lambda info: {"context": torch.ones(1)}) + monkeypatch.setattr(runner, "maybe_empty_cache", lambda: None) + monkeypatch.setattr(runner, "end_run", lambda: None) + monkeypatch.setattr(runner, "run_main", lambda: tuple(runner.input_info.latent_shape)) + monkeypatch.setattr(runner, "load_reuse_state", partial(runner.load_reuse_state, map_location="cpu")) + return runner + + +@pytest.mark.parametrize( + ("resize_config", "expected_key_shape"), + [({}, [480, 832]), ({"resize_mode": "adaptive"}, None), ({"resize_mode": "fixed_shape", "fixed_shape": [64, 96]}, None)], +) +def test_three_requests_reuse_after_real_image_preprocessing(cpu_runner, monkeypatch, resize_config, expected_key_shape, tmp_path): + runner = cpu_runner + runner.config.update(resize_config) + request = {"prompt": "same", "image_path": str(tmp_path / "square.png"), "save_result_path": str(tmp_path / "result.mp4")} + first = runner.prepare_request({"task": runner.config["task"], **request}) + assert first.target_shape == [480, 832] + first_result = runner.run_pipeline(first) + with open(runner.reuse_cache_dir + "/manifest.json") as file: + saved_key = json.load(file)["reuse_key"] + if expected_key_shape is None: + assert "target_shape" not in saved_key + else: + assert saved_key["target_shape"] == expected_key_shape + assert "resize_mode" not in saved_key + assert "resize_config" not in saved_key + if resize_config.get("resize_mode") == "adaptive": + assert first.target_shape == [480, 480] + elif resize_config.get("resize_mode") == "fixed_shape": + assert first.target_shape == [64, 96] + + def unexpected_encode(): + pytest.fail("Reusing a successful request must not rerun input encoders") + + monkeypatch.setattr(runner, "run_input_encoder", unexpected_encode) + runner.set_reuse(True) + for seed in (2, 3): + next_input = runner.prepare_request({"task": runner.config["task"], **request, "seed": seed}) + assert next_input.target_shape == [480, 832] + assert runner.run_pipeline(next_input) == first_result + assert next_input.target_shape == first.target_shape + assert next_input.latent_shape == first.latent_shape + assert next_input is not first + assert runner.config["target_height"] == 480 + assert runner.config["target_width"] == 832 + + +@pytest.mark.parametrize( + ("initial_config", "changed_request"), + [ + ({}, {"target_shape": [720, 1280]}), + ({"resize_mode": "adaptive"}, {"prompt": "changed"}), + ({"resize_mode": "adaptive"}, {"negative_prompt": "changed"}), + ({"resize_mode": "adaptive"}, {"target_video_length": 9}), + ({"resize_mode": "adaptive"}, {"image_path": "different.png"}), + ], +) +def test_reuse_rejects_changed_input_dependencies(cpu_runner, initial_config, changed_request, tmp_path): + runner = cpu_runner + runner.config.update(initial_config) + request = {"image_path": str(tmp_path / "square.png"), "save_result_path": str(tmp_path / "result.mp4")} + runner.run_pipeline(runner.prepare_request({"task": runner.config["task"], **request})) + runner.set_reuse(True) + with pytest.raises(ValueError, match="Reuse inputs must match"): + runner.run_pipeline(runner.prepare_request({"task": runner.config["task"], **request, **changed_request})) + + +@pytest.mark.parametrize("runner_cls", [WanRunner, Wan22DenseRunner, Wan22MoeRunner]) +@pytest.mark.parametrize("resize_mode", ["adaptive", "keep_ratio_fixed_area", "fixed_min_area", "fixed_max_area", "fixed_shape", "fixed_min_side"]) +def test_standard_wan_resize_modes_own_image_output_shape(runner_cls, resize_mode): + for task in runner_cls.supported_request_fields_by_task: + runner = object.__new__(runner_cls) + BaseRunner.__init__(runner, {"task": task, "resize_mode": resize_mode}) + fields = runner.get_supported_request_fields(task) + assert ("target_shape" in fields) == (task == "t2v") + del runner.config["resize_mode"] + assert "target_shape" in runner.get_supported_request_fields(task) + + +def test_custom_audio_image_preprocessing_keeps_its_shape_contract(): + runner = object.__new__(WanAudioARRunner) + BaseRunner.__init__(runner, {"task": "rs2v", "resize_mode": "adaptive"}) + fields = runner.get_supported_request_fields("rs2v") + assert "target_shape" in fields + + +def test_request_shape_is_rejected_when_startup_resize_owns_it(cpu_runner): + cpu_runner.config["resize_mode"] = "adaptive" + with pytest.raises(ValueError, match="target_shape"): + cpu_runner.prepare_request({"task": cpu_runner.config["task"], "target_shape": [720, 1280]}) diff --git a/tests/test_worldplay_request_pose.py b/tests/test_worldplay_request_pose.py new file mode 100644 index 000000000..2bb0ad4d8 --- /dev/null +++ b/tests/test_worldplay_request_pose.py @@ -0,0 +1,115 @@ +import json +from unittest.mock import Mock + +import pytest +import torch + +from lightx2v.models.networks.worldplay import pose_utils +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.worldplay.worldplay_ar_runner import WorldPlayARRunner +from lightx2v.models.runners.worldplay.worldplay_bi_runner import WorldPlayBIRunner +from lightx2v.models.runners.worldplay.worldplay_distill_runner import WorldPlayDistillRunner +from lightx2v.models.schedulers.hunyuan_video.scheduler import HunyuanVideo15Scheduler +from lightx2v.utils.input_info import UNSET +from lightx2v.utils.lockable_dict import LockableDict + + +@pytest.mark.parametrize("runner_cls", [WorldPlayARRunner, WorldPlayBIRunner, WorldPlayDistillRunner]) +def test_worldplay_t2v_reaches_existing_encoder_and_conditioning(monkeypatch, runner_cls): + monkeypatch.setattr(f"{runner_cls.__module__}.AI_DEVICE", "cpu") + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + instance = object.__new__(runner_cls) + config = LockableDict({"task": "t2v", "hidden_size": 4, "target_video_length": 17, "vae_stride": [4, 16, 16]}) + BaseRunner.__init__(instance, config) + instance.load_model = Mock() + instance.vision_num_semantic_tokens = 3 + instance.get_latent_shape_with_target_hw = Mock(return_value=[2, 5, 2, 2]) + instance.run_text_encoder = Mock(return_value={"prompt_embeds": torch.zeros(1, 2, 4)}) + DefaultRunner.init_modules(instance) + instance.input_info = instance.prepare_request({"prompt": "move forward", "pose": "w-4"}) + + encoded = instance.run_input_encoder() + + assert instance.input_info.task == "t2v" + assert instance.input_info.prompt == "move forward" + assert encoded["pose_output"]["action"].shape == (1, 5) + assert encoded["image_encoder_output"]["cond_latents"] is None + assert torch.count_nonzero(encoded["image_encoder_output"]["siglip_output"]) == 0 + scheduler = object.__new__(HunyuanVideo15Scheduler) + condition, mask = scheduler._prepare_cond_latents_and_mask("t2v", None, torch.ones(1, 2, 5, 2, 2), torch.zeros(5), False) + assert torch.count_nonzero(condition) == torch.count_nonzero(mask) == 0 + with pytest.raises(ValueError, match="image_path"): + instance.prepare_request({"image_path": "unused.png"}) + + +@pytest.fixture +def runner(): + instance = object.__new__(WorldPlayARRunner) + BaseRunner.__init__(instance, LockableDict({"task": "i2v", "model_cls": "worldplay_ar", "target_video_length": 125, "vae_stride": [4, 16, 16]})) + return instance + + +@pytest.mark.parametrize("frame_fields", [{}, {"target_video_length": None}, {"target_video_length": UNSET}, {"target_video_length": 61}]) +def test_pose_sets_request_length_without_changing_startup(runner, frame_fields): + request = {"pose": "w-15", **frame_fields} + original = request.copy() + + input_info = runner.prepare_request(request) + + assert input_info.target_video_length == 61 + assert len(input_info.pose) == 16 + assert request == original + assert runner.config["target_video_length"] == 125 + assert runner.prepare_request({}).target_video_length == 125 + assert runner.prepare_request({}).pose is None + + +@pytest.mark.parametrize("requested_frames", [62, 125]) +def test_explicit_frame_count_must_match_pose_exactly(runner, requested_frames): + with pytest.raises(ValueError, match=f"pose corresponds to 61 frames, but num_frames is {requested_frames}"): + runner.prepare_request({"pose": "w-15", "target_video_length": requested_frames}) + assert runner.config["target_video_length"] == 125 + + +def test_without_pose_preserves_request_frame_override(runner): + input_info = runner.prepare_request({"target_video_length": 121}) + assert input_info.target_video_length == 121 + assert input_info.pose is None + + +def test_prepared_pose_does_not_read_json_again(runner, tmp_path): + pose = pose_utils.pose_string_to_json("w-15") + path = tmp_path / "pose.json" + path.write_text(json.dumps(pose)) + + input_info = runner.prepare_request({"pose": str(path)}) + path.unlink() + viewmats, intrinsics, actions = pose_utils.pose_to_input(input_info.pose, latent_num=16) + + assert input_info.target_video_length == 61 + assert input_info.pose == pose + assert viewmats.shape == (16, 4, 4) + assert intrinsics.shape == (16, 3, 3) + assert actions.shape == (16,) + + +def test_command_pose_is_parsed_once(runner, monkeypatch): + parse = Mock(wraps=pose_utils.pose_string_to_json) + monkeypatch.setattr(pose_utils, "pose_string_to_json", parse) + + input_info = runner.prepare_request({"pose": "d-31"}) + pose_utils.pose_to_input(input_info.pose, latent_num=32) + + parse.assert_called_once_with("d-31") + assert input_info.target_video_length == 125 + + +def test_pose_dictionary_and_direct_shape_validation(runner): + pose = pose_utils.pose_string_to_json("w-15,s-15") + input_info = runner.prepare_request({"pose": pose, "target_video_length": 121}) + + assert input_info.pose is pose + assert input_info.target_video_length == 121 + with pytest.raises(ValueError, match="pose corresponds to 121 frames"): + pose_utils.pose_to_input(pose, latent_num=32)