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
helloyongyang's avatar
helloyongyang committed
6
import torch.nn.functional as F
PengGao's avatar
PengGao committed
7
8
9
from loguru import logger
from safetensors import safe_open

helloyongyang's avatar
helloyongyang committed
10
from lightx2v.common.ops.attn import MaskMap
11
from lightx2v.models.networks.wan.infer.feature_caching.transformer_infer import (
12
13
    WanTransformerInferAdaCaching,
    WanTransformerInferCustomCaching,
Rongjin Yang's avatar
Rongjin Yang committed
14
15
    WanTransformerInferDualBlock,
    WanTransformerInferDynamicBlock,
PengGao's avatar
PengGao committed
16
17
18
19
20
21
22
23
24
25
26
27
28
    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,
29
)
30
from lightx2v.utils.envs import *
31
from lightx2v.utils.utils import *
helloyongyang's avatar
helloyongyang committed
32

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

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

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

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

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

gushiqiao's avatar
gushiqiao committed
54
55
        if self.dit_quantized:
            dit_quant_scheme = self.config.mm_config.get("mm_type").split("-")[1]
gushiqiao's avatar
gushiqiao committed
56
57
            if self.config.model_cls == "wan2.1_distill":
                dit_quant_scheme = "distill_" + dit_quant_scheme
58
59
60
61
            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
62
                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
63
64
65
66
67
            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
68
69
        else:
            self.dit_quantized_ckpt = None
70
71
            assert not self.config.get("lazy_load", False)

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

74
75
76
77
        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
78
        self.device = device
helloyongyang's avatar
helloyongyang committed
79
80
81
82
83
84
85
        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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

        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
103
        else:
helloyongyang's avatar
helloyongyang committed
104
            raise NotImplementedError(f"Unsupported feature_caching type: {self.config['feature_caching']}")
helloyongyang's avatar
helloyongyang committed
105

gushiqiao's avatar
gushiqiao committed
106
107
108
109
110
111
112
113
114
115
116
117
    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

118
    def _load_safetensor_to_dict(self, file_path, unified_dtype, sensitive_layer):
helloyongyang's avatar
helloyongyang committed
119
        with safe_open(file_path, framework="pt") as f:
120
121
122
123
            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
124

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

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

gushiqiao's avatar
Fix  
gushiqiao committed
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
        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():
154
                    if f.get_tensor(k).dtype in [torch.float16, torch.bfloat16, torch.float]:
155
156
                        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
157
                        else:
158
                            weight_dict[k] = f.get_tensor(k).pin_memory().to(GET_SENSITIVE_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
159
160
                    else:
                        weight_dict[k] = f.get_tensor(k).pin_memory().to(self.device)
161

162
163
        return weight_dict

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

        safetensor_path = os.path.join(lazy_load_model_path, "non_block.safetensors")
gushiqiao's avatar
gushiqiao committed
170
        with safe_open(safetensor_path, framework="pt", device="cpu") as f:
171
            for k in f.keys():
172
                if f.get_tensor(k).dtype in [torch.float16, torch.bfloat16, torch.float]:
173
174
                    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
175
                    else:
176
                        pre_post_weight_dict[k] = f.get_tensor(k).pin_memory().to(GET_SENSITIVE_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
177
178
                else:
                    pre_post_weight_dict[k] = f.get_tensor(k).pin_memory().to(self.device)
179

gushiqiao's avatar
gushiqiao committed
180
        return pre_post_weight_dict
181

182
183
184
185
186
187
188
189
    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
190
    def _init_weights(self, weight_dict=None):
191
        unified_dtype = GET_DTYPE() == GET_SENSITIVE_DTYPE()
gushiqiao's avatar
Fix  
gushiqiao committed
192
        # Some layers run with float32 to achieve high accuracy
193
        sensitive_layer = {
gushiqiao's avatar
gushiqiao committed
194
195
196
197
198
199
200
            "norm",
            "embedding",
            "modulation",
            "time",
            "img_emb.proj.0",
            "img_emb.proj.4",
        }
201

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

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

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

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

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

gushiqiao's avatar
gushiqiao committed
232
    def _load_weights_distribute(self, weight_dict, is_weight_loader):
gushiqiao's avatar
gushiqiao committed
233
234
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
        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
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
            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
287
288
289
290

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

helloyongyang's avatar
helloyongyang committed
291
292
293
    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
294
        self.transformer_infer = self.transformer_infer_class(self.config)
helloyongyang's avatar
helloyongyang committed
295
296
297

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

TorynCurtis's avatar
TorynCurtis committed
302
303
304
305
306
307
308
309
310
311
    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
312
313
    @torch.no_grad()
    def infer(self, inputs):
314
315
316
317
318
319
320
        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()

321
322
323
324
325
        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)

326
        if self.config["enable_cfg"]:
helloyongyang's avatar
helloyongyang committed
327
328
329
330
331
332
333
334
335
336
            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
337

helloyongyang's avatar
helloyongyang committed
338
339
340
341
342
343
344
345
                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
346

helloyongyang's avatar
helloyongyang committed
347
348
349
350
            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)
351
352
353
354
355

        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
356
357
                self.pre_weight.to_cpu()
                self.post_weight.to_cpu()
gushiqiao's avatar
gushiqiao committed
358

359
    @torch.no_grad()
helloyongyang's avatar
helloyongyang committed
360
361
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
    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)
414

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

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