()
| 55 | from lightx2v.utils.utils import seed_all, validate_config_paths |
| 56 | from lightx2v_platform.registry_factory import PLATFORM_DEVICE_REGISTER |
| 57 | |
| 58 | |
| 59 | def init_runner(config): |
| 60 | torch.set_grad_enabled(False) |
| 61 | runner = RUNNER_REGISTER[config["model_cls"]](config) |
| 62 | runner.init_modules() |
| 63 | return runner |
| 64 | |
| 65 | |
| 66 | def distributed_barrier(): |
| 67 | import torch.distributed as dist |
| 68 | |
| 69 | if not dist.is_available() or not dist.is_initialized() or dist.get_world_size() <= 1: |
| 70 | return False |
| 71 | |
| 72 | from lightx2v_platform.base.global_var import AI_DEVICE |
| 73 | |
| 74 | if AI_DEVICE == "cuda" and torch.cuda.is_available(): |
| 75 | torch.cuda.synchronize() |
| 76 | dist.barrier(device_ids=[torch.cuda.current_device()]) |
| 77 | else: |
| 78 | dist.barrier() |
| 79 | |
| 80 | from loguru import logger |
| 81 | |
| 82 | logger.info(f"[Barrier] synchronized all ranks") |
| 83 | return True |
| 84 | |
| 85 | |
| 86 | def main(): |
| 87 | parser = argparse.ArgumentParser() |
| 88 | parser.add_argument("--seed", type=int, default=42, help="The seed for random generator") |
| 89 | parser.add_argument( |
| 90 | "--model_cls", |
| 91 | type=str, |
| 92 | required=True, |
| 93 | choices=[ |
| 94 | "wan2.1", |
| 95 | "wan2.1_distill", |
| 96 | "wan2.1_mean_flow_distill", |
| 97 | "wan_dancer", |
| 98 | "wan2.1_vace", |
| 99 | "wan2.1_sf", |
| 100 | "wan2.1_sf_mtxg2", |
| 101 | "seko_talk", |
| 102 | "seko_talk_ar", |
| 103 | "wan2.2_moe", |
| 104 | "lingbot_world", |
| 105 | "wan2.2", |
| 106 | "wan2.2_matrix_game3", |
| 107 | "wan2.2_moe_audio", |
| 108 | "wan2.2_audio", |
| 109 | "wan2.2_moe_distill", |
| 110 | "wan2.2_moe_vace", |
| 111 | "qwen_image", |
| 112 | "ernie_image", |
| 113 | "ernie_image_turbo", |
| 114 | "hidream_o1_image", |
no test coverage detected