model.py 17.9 KB
Newer Older
1
import json
2
3
import os

helloyongyang's avatar
helloyongyang committed
4
import torch
5
import torch.distributed as dist
PengGao's avatar
PengGao committed
6
7
8
from loguru import logger
from safetensors import safe_open

helloyongyang's avatar
helloyongyang committed
9
from lightx2v.common.ops.attn import MaskMap
10
from lightx2v.models.networks.wan.infer.feature_caching.transformer_infer import (
11
12
    WanTransformerInferAdaCaching,
    WanTransformerInferCustomCaching,
Rongjin Yang's avatar
Rongjin Yang committed
13
14
    WanTransformerInferDualBlock,
    WanTransformerInferDynamicBlock,
PengGao's avatar
PengGao committed
15
16
17
18
19
20
21
22
23
24
25
26
27
    WanTransformerInferFirstBlock,
    WanTransformerInferTaylorCaching,
    WanTransformerInferTeaCaching,
)
from lightx2v.models.networks.wan.infer.post_infer import WanPostInfer
from lightx2v.models.networks.wan.infer.pre_infer import WanPreInfer
from lightx2v.models.networks.wan.infer.transformer_infer import (
    WanTransformerInfer,
)
from lightx2v.models.networks.wan.weights.post_weights import WanPostWeights
from lightx2v.models.networks.wan.weights.pre_weights import WanPreWeights
from lightx2v.models.networks.wan.weights.transformer_weights import (
    WanTransformerWeights,
28
)
29
from lightx2v.utils.envs import *
30
from lightx2v.utils.utils import *
helloyongyang's avatar
helloyongyang committed
31

32
33
34
35
36
try:
    import gguf
except ImportError:
    gguf = None

helloyongyang's avatar
helloyongyang committed
37
38
39
40
41
42

class WanModel:
    pre_weight_class = WanPreWeights
    post_weight_class = WanPostWeights
    transformer_weight_class = WanTransformerWeights

43
    def __init__(self, model_path, config, device, seq_p_group=None):
helloyongyang's avatar
helloyongyang committed
44
45
        self.model_path = model_path
        self.config = config
46
47
        self.cpu_offload = self.config.get("cpu_offload", False)
        self.offload_granularity = self.config.get("offload_granularity", "block")
48
        self.seq_p_group = seq_p_group
49

gushiqiao's avatar
gushiqiao committed
50
        self.clean_cuda_cache = self.config.get("clean_cuda_cache", False)
51
        self.dit_quantized = self.config.mm_config.get("mm_type", "Default") != "Default"
52

gushiqiao's avatar
gushiqiao committed
53
54
        if self.dit_quantized:
            dit_quant_scheme = self.config.mm_config.get("mm_type").split("-")[1]
gushiqiao's avatar
gushiqiao committed
55
56
            if self.config.model_cls == "wan2.1_distill":
                dit_quant_scheme = "distill_" + dit_quant_scheme
57
58
59
60
            if dit_quant_scheme == "gguf":
                self.dit_quantized_ckpt = find_gguf_model_path(config, "dit_quantized_ckpt", subdir=dit_quant_scheme)
                self.config.use_gguf = True
            else:
helloyongyang's avatar
helloyongyang committed
61
                self.dit_quantized_ckpt = find_hf_model_path(config, self.model_path, "dit_quantized_ckpt", subdir=dit_quant_scheme)
gushiqiao's avatar
Fix bug  
gushiqiao committed
62
63
64
65
66
            quant_config_path = os.path.join(self.dit_quantized_ckpt, "config.json")
            if os.path.exists(quant_config_path):
                with open(quant_config_path, "r") as f:
                    quant_model_config = json.load(f)
                self.config.update(quant_model_config)
gushiqiao's avatar
gushiqiao committed
67
68
        else:
            self.dit_quantized_ckpt = None
69
70
            assert not self.config.get("lazy_load", False)

gushiqiao's avatar
gushiqiao committed
71
        self.config.dit_quantized_ckpt = self.dit_quantized_ckpt
gushiqiao's avatar
gushiqiao committed
72

73
74
75
76
        self.weight_auto_quant = self.config.mm_config.get("weight_auto_quant", False)
        if self.dit_quantized:
            assert self.weight_auto_quant or self.dit_quantized_ckpt is not None

gushiqiao's avatar
gushiqiao committed
77
        self.device = device
helloyongyang's avatar
helloyongyang committed
78
79
80
81
82
83
84
        self._init_infer_class()
        self._init_weights()
        self._init_infer()

    def _init_infer_class(self):
        self.pre_infer_class = WanPreInfer
        self.post_infer_class = WanPostInfer
helloyongyang's avatar
helloyongyang committed
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

        if self.config["feature_caching"] == "NoCaching":
            self.transformer_infer_class = WanTransformerInfer
        elif self.config["feature_caching"] == "Tea":
            self.transformer_infer_class = WanTransformerInferTeaCaching
        elif self.config["feature_caching"] == "TaylorSeer":
            self.transformer_infer_class = WanTransformerInferTaylorCaching
        elif self.config["feature_caching"] == "Ada":
            self.transformer_infer_class = WanTransformerInferAdaCaching
        elif self.config["feature_caching"] == "Custom":
            self.transformer_infer_class = WanTransformerInferCustomCaching
        elif self.config["feature_caching"] == "FirstBlock":
            self.transformer_infer_class = WanTransformerInferFirstBlock
        elif self.config["feature_caching"] == "DualBlock":
            self.transformer_infer_class = WanTransformerInferDualBlock
        elif self.config["feature_caching"] == "DynamicBlock":
            self.transformer_infer_class = WanTransformerInferDynamicBlock
helloyongyang's avatar
helloyongyang committed
102
        else:
helloyongyang's avatar
helloyongyang committed
103
            raise NotImplementedError(f"Unsupported feature_caching type: {self.config['feature_caching']}")
helloyongyang's avatar
helloyongyang committed
104

gushiqiao's avatar
gushiqiao committed
105
106
107
108
109
110
111
112
113
114
115
116
    def _should_load_weights(self):
        """Determine if current rank should load weights from disk."""
        if self.config.get("device_mesh") is None:
            # Single GPU mode
            return True
        elif dist.is_initialized():
            # Multi-GPU mode, only rank 0 loads
            if dist.get_rank() == 0:
                logger.info(f"Loading weights from {self.model_path}")
                return True
        return False

117
    def _load_safetensor_to_dict(self, file_path, unified_dtype, sensitive_layer):
helloyongyang's avatar
helloyongyang committed
118
        with safe_open(file_path, framework="pt") as f:
119
120
121
122
            return {
                key: (f.get_tensor(key).to(GET_DTYPE()) if unified_dtype or all(s not in key for s in sensitive_layer) else f.get_tensor(key).to(GET_SENSITIVE_DTYPE())).pin_memory().to(self.device)
                for key in f.keys()
            }
helloyongyang's avatar
helloyongyang committed
123

124
    def _load_ckpt(self, unified_dtype, sensitive_layer):
helloyongyang's avatar
helloyongyang committed
125
        safetensors_path = find_hf_model_path(self.config, self.model_path, "dit_original_ckpt", subdir="original")
126
        safetensors_files = glob.glob(os.path.join(safetensors_path, "*.safetensors"))
helloyongyang's avatar
helloyongyang committed
127
128
        weight_dict = {}
        for file_path in safetensors_files:
129
            file_weights = self._load_safetensor_to_dict(file_path, unified_dtype, sensitive_layer)
helloyongyang's avatar
helloyongyang committed
130
131
132
            weight_dict.update(file_weights)
        return weight_dict

133
    def _load_quant_ckpt(self, unified_dtype, sensitive_layer):
gushiqiao's avatar
gushiqiao committed
134
        ckpt_path = self.dit_quantized_ckpt
135
        logger.info(f"Loading quant dit model from {ckpt_path}")
136

gushiqiao's avatar
Fix  
gushiqiao committed
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
        index_files = [f for f in os.listdir(ckpt_path) if f.endswith(".index.json")]
        if not index_files:
            raise FileNotFoundError(f"No *.index.json found in {ckpt_path}")

        index_path = os.path.join(ckpt_path, index_files[0])
        logger.info(f" Using safetensors index: {index_path}")

        with open(index_path, "r") as f:
            index_data = json.load(f)

        weight_dict = {}
        for filename in set(index_data["weight_map"].values()):
            safetensor_path = os.path.join(ckpt_path, filename)
            with safe_open(safetensor_path, framework="pt") as f:
                logger.info(f"Loading weights from {safetensor_path}")
                for k in f.keys():
153
                    if f.get_tensor(k).dtype in [torch.float16, torch.bfloat16, torch.float]:
154
155
                        if unified_dtype or all(s not in k for s in sensitive_layer):
                            weight_dict[k] = f.get_tensor(k).pin_memory().to(GET_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
156
                        else:
157
                            weight_dict[k] = f.get_tensor(k).pin_memory().to(GET_SENSITIVE_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
158
159
                    else:
                        weight_dict[k] = f.get_tensor(k).pin_memory().to(self.device)
160

161
162
        return weight_dict

163
    def _load_quant_split_ckpt(self, unified_dtype, sensitive_layer):
gushiqiao's avatar
gushiqiao committed
164
        lazy_load_model_path = self.dit_quantized_ckpt
165
        logger.info(f"Loading splited quant model from {lazy_load_model_path}")
gushiqiao's avatar
gushiqiao committed
166
        pre_post_weight_dict = {}
167
168

        safetensor_path = os.path.join(lazy_load_model_path, "non_block.safetensors")
gushiqiao's avatar
gushiqiao committed
169
        with safe_open(safetensor_path, framework="pt", device="cpu") as f:
170
            for k in f.keys():
171
                if f.get_tensor(k).dtype in [torch.float16, torch.bfloat16, torch.float]:
172
173
                    if unified_dtype or all(s not in k for s in sensitive_layer):
                        pre_post_weight_dict[k] = f.get_tensor(k).pin_memory().to(GET_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
174
                    else:
175
                        pre_post_weight_dict[k] = f.get_tensor(k).pin_memory().to(GET_SENSITIVE_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
176
177
                else:
                    pre_post_weight_dict[k] = f.get_tensor(k).pin_memory().to(self.device)
178

gushiqiao's avatar
gushiqiao committed
179
        return pre_post_weight_dict
180

181
182
183
184
185
186
187
188
    def _load_gguf_ckpt(self):
        gguf_path = self.dit_quantized_ckpt
        logger.info(f"Loading gguf-quant dit model from {gguf_path}")
        reader = gguf.GGUFReader(gguf_path)
        for tensor in reader.tensors:
            # TODO: implement _load_gguf_ckpt
            pass

lijiaqi2's avatar
lijiaqi2 committed
189
    def _init_weights(self, weight_dict=None):
190
        unified_dtype = GET_DTYPE() == GET_SENSITIVE_DTYPE()
gushiqiao's avatar
Fix  
gushiqiao committed
191
        # Some layers run with float32 to achieve high accuracy
192
        sensitive_layer = {
gushiqiao's avatar
gushiqiao committed
193
194
195
196
197
198
199
            "norm",
            "embedding",
            "modulation",
            "time",
            "img_emb.proj.0",
            "img_emb.proj.4",
        }
200

lijiaqi2's avatar
lijiaqi2 committed
201
        if weight_dict is None:
gushiqiao's avatar
gushiqiao committed
202
            is_weight_loader = self._should_load_weights()
203
204
            if is_weight_loader:
                if not self.dit_quantized or self.weight_auto_quant:
gushiqiao's avatar
gushiqiao committed
205
206
                    # Load original weights
                    weight_dict = self._load_ckpt(unified_dtype, sensitive_layer)
207
                else:
gushiqiao's avatar
gushiqiao committed
208
                    # Load quantized weights
209
                    if not self.config.get("lazy_load", False):
gushiqiao's avatar
gushiqiao committed
210
                        weight_dict = self._load_quant_ckpt(unified_dtype, sensitive_layer)
211
                    else:
gushiqiao's avatar
gushiqiao committed
212
                        weight_dict = self._load_quant_split_ckpt(unified_dtype, sensitive_layer)
213

gushiqiao's avatar
gushiqiao committed
214
            if self.config.get("device_mesh") is not None:
gushiqiao's avatar
gushiqiao committed
215
                weight_dict = self._load_weights_distribute(weight_dict, is_weight_loader)
216

gushiqiao's avatar
gushiqiao committed
217
            self.original_weight_dict = weight_dict
lijiaqi2's avatar
lijiaqi2 committed
218
219
        else:
            self.original_weight_dict = weight_dict
220

gushiqiao's avatar
gushiqiao committed
221
        # Initialize weight containers
helloyongyang's avatar
helloyongyang committed
222
        self.pre_weight = self.pre_weight_class(self.config)
TorynCurtis's avatar
TorynCurtis committed
223
        self.post_weight = self.post_weight_class(self.config)
helloyongyang's avatar
helloyongyang committed
224
        self.transformer_weights = self.transformer_weight_class(self.config)
gushiqiao's avatar
gushiqiao committed
225
226

        # Load weights into containers
227
228
        self.pre_weight.load(self.original_weight_dict)
        self.post_weight.load(self.original_weight_dict)
gushiqiao's avatar
gushiqiao committed
229
        self.transformer_weights.load(self.original_weight_dict)
helloyongyang's avatar
helloyongyang committed
230

231
232
233
        del self.original_weight_dict
        torch.cuda.empty_cache()

gushiqiao's avatar
gushiqiao committed
234
    def _load_weights_distribute(self, weight_dict, is_weight_loader):
gushiqiao's avatar
gushiqiao committed
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
        global_src_rank = 0
        target_device = "cpu" if self.cpu_offload else "cuda"

        if is_weight_loader:
            meta_dict = {}
            for key, tensor in weight_dict.items():
                meta_dict[key] = {"shape": tensor.shape, "dtype": tensor.dtype}

            obj_list = [meta_dict]
            dist.broadcast_object_list(obj_list, src=global_src_rank)
            synced_meta_dict = obj_list[0]
        else:
            obj_list = [None]
            dist.broadcast_object_list(obj_list, src=global_src_rank)
            synced_meta_dict = obj_list[0]

        distributed_weight_dict = {}
        for key, meta in synced_meta_dict.items():
            distributed_weight_dict[key] = torch.empty(meta["shape"], dtype=meta["dtype"], device=target_device)

        if target_device == "cuda":
            dist.barrier(device_ids=[torch.cuda.current_device()])
        else:
            dist.barrier()

        for key in sorted(synced_meta_dict.keys()):
            if is_weight_loader:
                distributed_weight_dict[key].copy_(weight_dict[key], non_blocking=True)

gushiqiao's avatar
gushiqiao committed
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
            if target_device == "cpu":
                if is_weight_loader:
                    gpu_tensor = distributed_weight_dict[key].cuda()
                    dist.broadcast(gpu_tensor, src=global_src_rank)
                    distributed_weight_dict[key].copy_(gpu_tensor.cpu(), non_blocking=True)
                    del gpu_tensor
                    torch.cuda.empty_cache()
                else:
                    gpu_tensor = torch.empty_like(distributed_weight_dict[key], device="cuda")
                    dist.broadcast(gpu_tensor, src=global_src_rank)
                    distributed_weight_dict[key].copy_(gpu_tensor.cpu(), non_blocking=True)
                    del gpu_tensor
                    torch.cuda.empty_cache()

                if distributed_weight_dict[key].is_pinned():
                    distributed_weight_dict[key].copy_(distributed_weight_dict[key], non_blocking=True)
            else:
                dist.broadcast(distributed_weight_dict[key], src=global_src_rank)

        if target_device == "cuda":
            torch.cuda.synchronize()
        else:
            for tensor in distributed_weight_dict.values():
                if tensor.is_pinned():
                    tensor.copy_(tensor, non_blocking=False)
gushiqiao's avatar
gushiqiao committed
289
290
291
292

        logger.info(f"Weights distributed across {dist.get_world_size()} devices on {target_device}")
        return distributed_weight_dict

helloyongyang's avatar
helloyongyang committed
293
294
295
    def _init_infer(self):
        self.pre_infer = self.pre_infer_class(self.config)
        self.post_infer = self.post_infer_class(self.config)
helloyongyang's avatar
helloyongyang committed
296
        self.transformer_infer = self.transformer_infer_class(self.config)
helloyongyang's avatar
helloyongyang committed
297
298
299

    def set_scheduler(self, scheduler):
        self.scheduler = scheduler
300
301
        self.pre_infer.set_scheduler(scheduler)
        self.post_infer.set_scheduler(scheduler)
helloyongyang's avatar
helloyongyang committed
302
303
        self.transformer_infer.set_scheduler(scheduler)

TorynCurtis's avatar
TorynCurtis committed
304
305
306
307
308
309
310
311
312
313
    def to_cpu(self):
        self.pre_weight.to_cpu()
        self.post_weight.to_cpu()
        self.transformer_weights.to_cpu()

    def to_cuda(self):
        self.pre_weight.to_cuda()
        self.post_weight.to_cuda()
        self.transformer_weights.to_cuda()

helloyongyang's avatar
helloyongyang committed
314
315
    @torch.no_grad()
    def infer(self, inputs):
316
317
318
319
320
321
322
        if self.cpu_offload:
            if self.offload_granularity == "model" and self.scheduler.step_index == 0:
                self.to_cuda()
            elif self.offload_granularity != "model":
                self.pre_weight.to_cuda()
                self.post_weight.to_cuda()

323
324
325
326
327
        if self.transformer_infer.mask_map is None:
            _, c, h, w = self.scheduler.latents.shape
            video_token_num = c * (h // 2) * (w // 2)
            self.transformer_infer.mask_map = MaskMap(video_token_num, c)

328
        if self.config["enable_cfg"]:
helloyongyang's avatar
helloyongyang committed
329
330
331
332
333
334
335
336
337
338
            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:
                    noise_pred = self._infer_cond_uncond(inputs, positive=True)
                else:
                    noise_pred = self._infer_cond_uncond(inputs, positive=False)
helloyongyang's avatar
helloyongyang committed
339

helloyongyang's avatar
helloyongyang committed
340
341
342
343
344
345
346
347
                noise_pred_list = [torch.zeros_like(noise_pred) for _ in range(2)]
                dist.all_gather(noise_pred_list, noise_pred, group=cfg_p_group)
                noise_pred_cond = noise_pred_list[0]  # cfg_p_rank == 0
                noise_pred_uncond = noise_pred_list[1]  # cfg_p_rank == 1
            else:
                # ==================== CFG Processing ====================
                noise_pred_cond = self._infer_cond_uncond(inputs, positive=True)
                noise_pred_uncond = self._infer_cond_uncond(inputs, positive=False)
gushiqiao's avatar
gushiqiao committed
348

helloyongyang's avatar
helloyongyang committed
349
350
351
352
            self.scheduler.noise_pred = noise_pred_uncond + self.scheduler.sample_guide_scale * (noise_pred_cond - noise_pred_uncond)
        else:
            # ==================== No CFG ====================
            self.scheduler.noise_pred = self._infer_cond_uncond(inputs, positive=True)
353
354
355
356
357

        if self.cpu_offload:
            if self.offload_granularity == "model" and self.scheduler.step_index == self.scheduler.infer_steps - 1:
                self.to_cpu()
            elif self.offload_granularity != "model":
root's avatar
root committed
358
359
                self.pre_weight.to_cpu()
                self.post_weight.to_cpu()
gushiqiao's avatar
gushiqiao committed
360

361
    @torch.no_grad()
helloyongyang's avatar
helloyongyang committed
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
    def _infer_cond_uncond(self, inputs, positive=True):
        pre_infer_out = self.pre_infer.infer(self.pre_weight, inputs, positive=positive)

        if self.config["seq_parallel"]:
            pre_infer_out = self._seq_parallel_pre_process(pre_infer_out)

        x = self.transformer_infer.infer(self.transformer_weights, pre_infer_out)

        if self.config["seq_parallel"]:
            x = self._seq_parallel_post_process(x)

        noise_pred = self.post_infer.infer(self.post_weight, x, pre_infer_out)[0]

        if self.clean_cuda_cache:
            del x, pre_infer_out
            torch.cuda.empty_cache()

        return noise_pred

    @torch.no_grad()
    def _seq_parallel_pre_process(self, pre_infer_out):
        embed, x, embed0 = pre_infer_out.embed, pre_infer_out.x, pre_infer_out.embed0

        world_size = dist.get_world_size(self.seq_p_group)
        cur_rank = dist.get_rank(self.seq_p_group)

        padding_size = (world_size - (x.shape[0] % world_size)) % world_size

        if padding_size > 0:
            # 使用 F.pad 填充第一维
            x = F.pad(x, (0, 0, 0, padding_size))  # (后维度填充, 前维度填充)

        x = torch.chunk(x, world_size, dim=0)[cur_rank]
        if self.config["model_cls"].startswith("wan2.2"):
            padding_size = (world_size - (embed0.shape[0] % world_size)) % world_size
            if padding_size > 0:
                embed0 = F.pad(embed0, (0, 0, 0, 0, 0, padding_size))  # (后维度填充, 前维度填充)
                embed = F.pad(embed, (0, 0, 0, padding_size))

        pre_infer_out.x = x
        pre_infer_out.embed = embed
        pre_infer_out.embed0 = embed0

        return pre_infer_out

    @torch.no_grad()
    def _seq_parallel_post_process(self, x):
        world_size = dist.get_world_size(self.seq_p_group)

        # 创建一个列表,用于存储所有进程的输出
        gathered_x = [torch.empty_like(x) for _ in range(world_size)]

        # 收集所有进程的输出
        dist.all_gather(gathered_x, x, group=self.seq_p_group)
416

helloyongyang's avatar
helloyongyang committed
417
418
        # 在指定的维度上合并所有进程的输出
        combined_output = torch.cat(gathered_x, dim=0)
419

helloyongyang's avatar
helloyongyang committed
420
        return combined_output  # 返回合并后的输出