set_config.py 4.31 KB
Newer Older
1
import json
2
import os
PengGao's avatar
PengGao committed
3
4

import torch.distributed as dist
5
from loguru import logger
6
from torch.distributed.tensor.device_mesh import init_device_mesh
7

8
9
10
from lightx2v.utils.input_info import ALL_INPUT_INFO_KEYS
from lightx2v.utils.lockable_dict import LockableDict

11

helloyongyang's avatar
helloyongyang committed
12
13
14
15
16
17
18
19
20
21
22
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,
23
        "lora_configs": None,  # List of dicts with 'path' and 'strength' keys
helloyongyang's avatar
helloyongyang committed
24
        "use_prompt_enhancer": False,
25
        "parallel": False,
helloyongyang's avatar
fix bug  
helloyongyang committed
26
27
        "seq_parallel": False,
        "cfg_parallel": False,
28
        "enable_cfg": False,
gushiqiao's avatar
gushiqiao committed
29
        "use_image_encoder": True,
helloyongyang's avatar
helloyongyang committed
30
    }
31
    default_config = LockableDict(default_config)
helloyongyang's avatar
helloyongyang committed
32
33
34
    return default_config


35
def set_config(args):
helloyongyang's avatar
helloyongyang committed
36
    config = get_default_config()
37
    config.update({k: v for k, v in vars(args).items() if k not in ALL_INPUT_INFO_KEYS})
38

39
40
41
42
43
    if config.get("config_json", None) 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)
44

45
46
    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:
helloyongyang's avatar
helloyongyang committed
47
48
            model_config = json.load(f)
        config.update(model_config)
49
50
    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:
51
52
            model_config = json.load(f)
        config.update(model_config)
53
54
    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:
gushiqiao's avatar
gushiqiao committed
55
56
            model_config = json.load(f)
        config.update(model_config)
57
58
    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:
gushiqiao's avatar
gushiqiao committed
59
60
61
            model_config = json.load(f)
        config.update(model_config)
    # load quantized config
gushiqiao's avatar
gushiqiao committed
62
    if config.get("dit_quantized_ckpt", None) is not None:
63
        config_path = os.path.join(config["dit_quantized_ckpt"], "config.json")
gushiqiao's avatar
gushiqiao committed
64
65
66
67
68
        if os.path.exists(config_path):
            with open(config_path, "r") as f:
                model_config = json.load(f)
            config.update(model_config)

69
70
71
72
    if config["task"] in ["i2v", "s2v"]:
        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
LiangLiu's avatar
LiangLiu committed
73

74
    return config
75
76
77


def set_parallel_config(config):
78
79
80
    if config["parallel"]:
        cfg_p_size = config["parallel"].get("cfg_p_size", 1)
        seq_p_size = config["parallel"].get("seq_p_size", 1)
81
82
83
        assert cfg_p_size * seq_p_size == dist.get_world_size(), f"cfg_p_size * seq_p_size must be equal to world_size"
        config["device_mesh"] = init_device_mesh("cuda", (cfg_p_size, seq_p_size), mesh_dim_names=("cfg_p", "seq_p"))

84
        if config["parallel"] and config["parallel"].get("seq_p_size", False) and config["parallel"]["seq_p_size"] > 1:
85
86
            config["seq_parallel"] = True

87
        if config.get("enable_cfg", False) and config["parallel"] and config["parallel"].get("cfg_p_size", False) and config["parallel"]["cfg_p_size"] > 1:
88
            config["cfg_parallel"] = True
helloyongyang's avatar
helloyongyang committed
89
90
91
92
93


def print_config(config):
    config_to_print = config.copy()
    config_to_print.pop("device_mesh", None)
94
    if config["parallel"]:
helloyongyang's avatar
helloyongyang committed
95
96
97
98
        if dist.get_rank() == 0:
            logger.info(f"config:\n{json.dumps(config_to_print, ensure_ascii=False, indent=4)}")
    else:
        logger.info(f"config:\n{json.dumps(config_to_print, ensure_ascii=False, indent=4)}")