"examples/trials/mnist-tfv1/config.yml" did not exist on "a587648999af697474fc48aff00b870736bfd618"
infer.py 3.45 KB
Newer Older
helloyongyang's avatar
helloyongyang committed
1
import argparse
helloyongyang's avatar
update  
helloyongyang committed
2
import json
3

PengGao's avatar
PengGao committed
4
5
import torch.distributed as dist
from loguru import logger
6

PengGao's avatar
PengGao committed
7
from lightx2v.common.ops import *
PengGao's avatar
PengGao committed
8
from lightx2v.models.runners.cogvideox.cogvidex_runner import CogvideoxRunner  # noqa: F401
PengGao's avatar
PengGao committed
9
from lightx2v.models.runners.graph_runner import GraphRunner
PengGao's avatar
PengGao committed
10
from lightx2v.models.runners.hunyuan.hunyuan_runner import HunyuanRunner  # noqa: F401
11
from lightx2v.models.runners.wan.wan_audio_runner import Wan22AudioRunner, Wan22MoeAudioRunner, WanAudioRunner  # noqa: F401
PengGao's avatar
PengGao committed
12
13
14
15
from lightx2v.models.runners.wan.wan_causvid_runner import WanCausVidRunner  # noqa: F401
from lightx2v.models.runners.wan.wan_distill_runner import WanDistillRunner  # noqa: F401
from lightx2v.models.runners.wan.wan_runner import Wan22MoeRunner, WanRunner  # noqa: F401
from lightx2v.models.runners.wan.wan_skyreels_v2_df_runner import WanSkyreelsV2DFRunner  # noqa: F401
PengGao's avatar
PengGao committed
16
17
18
from lightx2v.utils.envs import *
from lightx2v.utils.profiler import ProfilingContext
from lightx2v.utils.registry_factory import RUNNER_REGISTER
helloyongyang's avatar
helloyongyang committed
19
from lightx2v.utils.set_config import set_config, set_parallel_config
PengGao's avatar
PengGao committed
20
from lightx2v.utils.utils import seed_all
lijiaqi2's avatar
lijiaqi2 committed
21
22


helloyongyang's avatar
helloyongyang committed
23
24
25
26
27
def init_runner(config):
    seed_all(config.seed)

    if CHECK_ENABLE_GRAPH_MODE():
        default_runner = RUNNER_REGISTER[config.model_cls](config)
helloyongyang's avatar
helloyongyang committed
28
        default_runner.init_modules()
helloyongyang's avatar
helloyongyang committed
29
30
31
        runner = GraphRunner(default_runner)
    else:
        runner = RUNNER_REGISTER[config.model_cls](config)
gushiqiao's avatar
gushiqiao committed
32
        runner.init_modules()
helloyongyang's avatar
helloyongyang committed
33
34
35
    return runner


36
def main():
helloyongyang's avatar
helloyongyang committed
37
    parser = argparse.ArgumentParser()
wangshankun's avatar
wangshankun committed
38
    parser.add_argument(
39
40
41
        "--model_cls",
        type=str,
        required=True,
42
43
44
45
46
47
48
49
50
51
52
53
54
55
        choices=[
            "wan2.1",
            "hunyuan",
            "wan2.1_distill",
            "wan2.1_causvid",
            "wan2.1_skyreels_v2_df",
            "cogvideox",
            "wan2.1_audio",
            "wan2.2_moe",
            "wan2.2_moe_audio",
            "wan2.2_audio",
            "wan2.2",
            "wan2.2_moe_distill",
        ],
56
        default="wan2.1",
wangshankun's avatar
wangshankun committed
57
    )
helloyongyang's avatar
helloyongyang committed
58

helloyongyang's avatar
helloyongyang committed
59
60
    parser.add_argument("--task", type=str, choices=["t2v", "i2v"], default="t2v")
    parser.add_argument("--model_path", type=str, required=True)
helloyongyang's avatar
helloyongyang committed
61
    parser.add_argument("--config_json", type=str, required=True)
62
    parser.add_argument("--use_prompt_enhancer", action="store_true")
helloyongyang's avatar
helloyongyang committed
63

wangshankun's avatar
wangshankun committed
64
    parser.add_argument("--prompt", type=str, default="", help="The input prompt for text-to-video generation")
helloyongyang's avatar
helloyongyang committed
65
    parser.add_argument("--negative_prompt", type=str, default="")
helloyongyang's avatar
helloyongyang committed
66
67
68
69

    parser.add_argument("--image_path", type=str, default="", help="The path to input image file for image-to-video (i2v) task")
    parser.add_argument("--audio_path", type=str, default="", help="The path to input audio file for audio-to-video (a2v) task")

70
    parser.add_argument("--save_video_path", type=str, default="./output_lightx2v.mp4", help="The path to save video path/file")
helloyongyang's avatar
helloyongyang committed
71
    args = parser.parse_args()
PengGao's avatar
PengGao committed
72

73
    logger.info(f"args: {args}")
Dongz's avatar
Dongz committed
74

helloyongyang's avatar
helloyongyang committed
75
76
77
78
    # set config
    config = set_config(args)
    logger.info(f"config:\n{json.dumps(config, ensure_ascii=False, indent=4)}")

helloyongyang's avatar
fix bug  
helloyongyang committed
79
    if config.parallel:
helloyongyang's avatar
helloyongyang committed
80
81
82
83
        dist.init_process_group(backend="nccl")
        torch.cuda.set_device(dist.get_rank())
        set_parallel_config(config)

helloyongyang's avatar
helloyongyang committed
84
    with ProfilingContext("Total Cost"):
helloyongyang's avatar
helloyongyang committed
85
        runner = init_runner(config)
86
        runner.run_pipeline()
Xinchi Huang's avatar
Xinchi Huang committed
87

helloyongyang's avatar
helloyongyang committed
88
89
90
91
92
    # Clean up distributed process group
    if dist.is_initialized():
        dist.destroy_process_group()
        logger.info("Distributed process group cleaned up")

Xinchi Huang's avatar
Xinchi Huang committed
93
94

if __name__ == "__main__":
95
    main()