model.py 21.5 KB
Newer Older
gushiqiao's avatar
gushiqiao committed
1
import gc
Yang Yong (雍洋)'s avatar
Yang Yong (雍洋) committed
2
import glob
3
4
import os

helloyongyang's avatar
helloyongyang committed
5
import torch
6
import torch.distributed as dist
helloyongyang's avatar
helloyongyang committed
7
import torch.nn.functional as F
PengGao's avatar
PengGao committed
8
9
10
from loguru import logger
from safetensors import safe_open

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
    WanTransformerInferFirstBlock,
Musisoul's avatar
Musisoul committed
17
    WanTransformerInferMagCaching,
PengGao's avatar
PengGao committed
18
19
20
    WanTransformerInferTaylorCaching,
    WanTransformerInferTeaCaching,
)
21
22
23
from lightx2v.models.networks.wan.infer.offload.transformer_infer import (
    WanOffloadTransformerInfer,
)
PengGao's avatar
PengGao committed
24
25
26
27
28
29
30
31
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.pre_weights import WanPreWeights
from lightx2v.models.networks.wan.weights.transformer_weights import (
    WanTransformerWeights,
32
)
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
33
from lightx2v.utils.custom_compiler import CompiledMethodsMixin, compiled_method
34
from lightx2v.utils.envs import *
yihuiwen's avatar
yihuiwen committed
35
from lightx2v.utils.ggml_tensor import load_gguf_sd_ckpt
36
from lightx2v.utils.utils import *
helloyongyang's avatar
helloyongyang committed
37
38


Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
39
class WanModel(CompiledMethodsMixin):
helloyongyang's avatar
helloyongyang committed
40
41
42
    pre_weight_class = WanPreWeights
    transformer_weight_class = WanTransformerWeights

43
    def __init__(self, model_path, config, device, model_type="wan2.1"):
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
44
        super().__init__()
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.model_type = model_type
50
51
52
53
        self.remove_keys = []
        self.lazy_load = self.config.get("lazy_load", False)
        if self.lazy_load:
            self.remove_keys.extend(["blocks."])
helloyongyang's avatar
helloyongyang committed
54
55
56
57
        if self.config["seq_parallel"]:
            self.seq_p_group = self.config.get("device_mesh").get_group(mesh_dim="seq_p")
        else:
            self.seq_p_group = None
58

gushiqiao's avatar
gushiqiao committed
59
        self.clean_cuda_cache = self.config.get("clean_cuda_cache", False)
60
        self.dit_quantized = self.config.get("dit_quantized", False)
61
        if self.dit_quantized:
62
63
64
65
66
67
68
69
70
71
72
73
74
75
            assert self.config.get("dit_quant_scheme", "Default") in [
                "Default-Force-FP32",
                "fp8-vllm",
                "int8-vllm",
                "fp8-q8f",
                "int8-q8f",
                "fp8-b128-deepgemm",
                "fp8-sgl",
                "int8-sgl",
                "int8-torchao",
                "nvfp4",
                "mxfp4",
                "mxfp6-mxfp8",
                "mxfp8",
Kane's avatar
Kane committed
76
                "int8-tmo",
yihuiwen's avatar
yihuiwen committed
77
78
79
80
81
82
83
84
85
86
87
88
                "gguf-Q8_0",
                "gguf-Q6_K",
                "gguf-Q5_K_S",
                "gguf-Q5_K_M",
                "gguf-Q5_0",
                "gguf-Q5_1",
                "gguf-Q4_K_S",
                "gguf-Q4_K_M",
                "gguf-Q4_0",
                "gguf-Q4_1",
                "gguf-Q3_K_S",
                "gguf-Q3_K_M",
89
            ]
gushiqiao's avatar
gushiqiao committed
90
        self.device = device
helloyongyang's avatar
helloyongyang committed
91
92
93
94
95
96
97
        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
98
99

        if self.config["feature_caching"] == "NoCaching":
100
            self.transformer_infer_class = WanTransformerInfer if not self.cpu_offload else WanOffloadTransformerInfer
helloyongyang's avatar
helloyongyang committed
101
102
103
104
105
106
107
108
109
110
111
112
113
114
        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
Musisoul's avatar
Musisoul committed
115
116
        elif self.config["feature_caching"] == "Mag":
            self.transformer_infer_class = WanTransformerInferMagCaching
helloyongyang's avatar
helloyongyang committed
117
        else:
helloyongyang's avatar
helloyongyang committed
118
            raise NotImplementedError(f"Unsupported feature_caching type: {self.config['feature_caching']}")
helloyongyang's avatar
helloyongyang committed
119

gushiqiao's avatar
gushiqiao committed
120
121
122
123
124
125
    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():
126
127
128
129
130
131
            if self.config.get("load_from_rank0", False):
                # Multi-GPU mode, only rank 0 loads
                if dist.get_rank() == 0:
                    logger.info(f"Loading weights from {self.model_path}")
                    return True
            else:
gushiqiao's avatar
gushiqiao committed
132
133
134
                return True
        return False

135
    def _should_init_empty_model(self):
136
        if self.config.get("lora_configs") and self.config["lora_configs"]:
137
138
139
140
141
142
143
144
145
146
147
148
            if self.model_type in ["wan2.1"]:
                return True
            if self.model_type in ["wan2.2_moe_high_noise"]:
                for lora_config in self.config["lora_configs"]:
                    if lora_config["name"] == "high_noise_model":
                        return True
            if self.model_type in ["wan2.2_moe_low_noise"]:
                for lora_config in self.config["lora_configs"]:
                    if lora_config["name"] == "low_noise_model":
                        return True
        return False

149
    def _load_safetensor_to_dict(self, file_path, unified_dtype, sensitive_layer):
150
151
        remove_keys = self.remove_keys if hasattr(self, "remove_keys") else []

152
        if self.device.type != "cpu" and dist.is_initialized():
153
            device = dist.get_rank()
154
        else:
155
            device = str(self.device)
156

157
        with safe_open(file_path, framework="pt", device=device) as f:
158
159
160
161
162
            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()))
                for key in f.keys()
                if not any(remove_key in key for remove_key in remove_keys)
            }
helloyongyang's avatar
helloyongyang committed
163

164
    def _load_ckpt(self, unified_dtype, sensitive_layer):
165
166
167
168
169
170
        if self.config.get("dit_original_ckpt", None):
            safetensors_path = self.config["dit_original_ckpt"]
        else:
            safetensors_path = self.model_path

        if os.path.isdir(safetensors_path):
Gu Shiqiao's avatar
Gu Shiqiao committed
171
172
173
174
175
176
177
178
179
            if self.lazy_load:
                self.lazy_load_path = safetensors_path
                non_block_file = os.path.join(safetensors_path, "non_block.safetensors")
                if os.path.exists(non_block_file):
                    safetensors_files = [non_block_file]
                else:
                    raise ValueError(f"Non-block file not found in {safetensors_path}. Please check the model path.")
            else:
                safetensors_files = glob.glob(os.path.join(safetensors_path, "*.safetensors"))
180
        else:
Gu Shiqiao's avatar
Gu Shiqiao committed
181
182
            if self.lazy_load:
                self.lazy_load_path = safetensors_path
183
            safetensors_files = [safetensors_path]
184

helloyongyang's avatar
helloyongyang committed
185
186
        weight_dict = {}
        for file_path in safetensors_files:
187
            if self.config.get("adapter_model_path", None) is not None:
188
                if self.config["adapter_model_path"] == file_path:
189
                    continue
190
            logger.info(f"Loading weights from {file_path}")
191
            file_weights = self._load_safetensor_to_dict(file_path, unified_dtype, sensitive_layer)
helloyongyang's avatar
helloyongyang committed
192
            weight_dict.update(file_weights)
193

helloyongyang's avatar
helloyongyang committed
194
195
        return weight_dict

196
    def _load_quant_ckpt(self, unified_dtype, sensitive_layer):
197
        remove_keys = self.remove_keys if hasattr(self, "remove_keys") else []
198
199
200
201
        if self.config.get("dit_quantized_ckpt", None):
            safetensors_path = self.config["dit_quantized_ckpt"]
        else:
            safetensors_path = self.model_path
gushiqiao's avatar
Fix  
gushiqiao committed
202

yihuiwen's avatar
yihuiwen committed
203
204
205
206
207
208
209
210
211
212
213
        if "gguf" in self.config.get("dit_quant_scheme", ""):
            gguf_path = ""
            if os.path.isdir(safetensors_path):
                gguf_type = self.config.get("dit_quant_scheme").replace("gguf-", "")
                gguf_files = list(filter(lambda x: gguf_type in x, glob.glob(os.path.join(safetensors_path, "*.gguf"))))
                gguf_path = gguf_files[0]
            else:
                gguf_path = safetensors_path
            weight_dict = self._load_gguf_ckpt(gguf_path)
            return weight_dict

214
        if os.path.isdir(safetensors_path):
Gu Shiqiao's avatar
Gu Shiqiao committed
215
216
217
218
219
220
221
222
223
            if self.lazy_load:
                self.lazy_load_path = safetensors_path
                non_block_file = os.path.join(safetensors_path, "non_block.safetensors")
                if os.path.exists(non_block_file):
                    safetensors_files = [non_block_file]
                else:
                    raise ValueError(f"Non-block file not found in {safetensors_path}. Please check the model path.")
            else:
                safetensors_files = glob.glob(os.path.join(safetensors_path, "*.safetensors"))
224
        else:
Gu Shiqiao's avatar
Gu Shiqiao committed
225
226
            if self.lazy_load:
                self.lazy_load_path = safetensors_path
227
            safetensors_files = [safetensors_path]
228
            safetensors_path = os.path.dirname(safetensors_path)
229

gushiqiao's avatar
Fix  
gushiqiao committed
230
        weight_dict = {}
231
232
233
234
        for safetensor_path in safetensors_files:
            if self.config.get("adapter_model_path", None) is not None:
                if self.config["adapter_model_path"] == safetensor_path:
                    continue
yihuiwen's avatar
yihuiwen committed
235

gushiqiao's avatar
Fix  
gushiqiao committed
236
237
238
            with safe_open(safetensor_path, framework="pt") as f:
                logger.info(f"Loading weights from {safetensor_path}")
                for k in f.keys():
239
240
                    if any(remove_key in k for remove_key in remove_keys):
                        continue
241
242
243
244
245
                    if f.get_tensor(k).dtype in [
                        torch.float16,
                        torch.bfloat16,
                        torch.float,
                    ]:
246
                        if unified_dtype or all(s not in k for s in sensitive_layer):
gushiqiao's avatar
gushiqiao committed
247
                            weight_dict[k] = f.get_tensor(k).to(GET_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
248
                        else:
gushiqiao's avatar
gushiqiao committed
249
                            weight_dict[k] = f.get_tensor(k).to(GET_SENSITIVE_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
250
                    else:
gushiqiao's avatar
gushiqiao committed
251
                        weight_dict[k] = f.get_tensor(k).to(self.device)
252

253
254
255
256
257
258
259
        if self.config.get("dit_quant_scheme", "Default") == "nvfp4":
            calib_path = os.path.join(safetensors_path, "calib.pt")
            logger.info(f"[CALIB] Loaded calibration data from: {calib_path}")
            calib_data = torch.load(calib_path, map_location="cpu")
            for k, v in calib_data["absmax"].items():
                weight_dict[k.replace(".weight", ".input_absmax")] = v.to(self.device)

260
261
        return weight_dict

yihuiwen's avatar
yihuiwen committed
262
263
264
    def _load_gguf_ckpt(self, gguf_path):
        state_dict = load_gguf_sd_ckpt(gguf_path, to_device=self.device)
        return state_dict
265

lijiaqi2's avatar
lijiaqi2 committed
266
    def _init_weights(self, weight_dict=None):
267
        unified_dtype = GET_DTYPE() == GET_SENSITIVE_DTYPE()
gushiqiao's avatar
Fix  
gushiqiao committed
268
        # Some layers run with float32 to achieve high accuracy
269
        sensitive_layer = {
gushiqiao's avatar
gushiqiao committed
270
271
272
273
274
275
            "norm",
            "embedding",
            "modulation",
            "time",
            "img_emb.proj.0",
            "img_emb.proj.4",
gushiqiao's avatar
gushiqiao committed
276
277
            "before_proj",  # vace
            "after_proj",  # vace
gushiqiao's avatar
gushiqiao committed
278
        }
279

lijiaqi2's avatar
lijiaqi2 committed
280
        if weight_dict is None:
gushiqiao's avatar
gushiqiao committed
281
            is_weight_loader = self._should_load_weights()
282
            if is_weight_loader:
283
                if not self.dit_quantized:
gushiqiao's avatar
gushiqiao committed
284
285
                    # Load original weights
                    weight_dict = self._load_ckpt(unified_dtype, sensitive_layer)
286
                else:
gushiqiao's avatar
gushiqiao committed
287
                    # Load quantized weights
288
                    weight_dict = self._load_quant_ckpt(unified_dtype, sensitive_layer)
289

290
291
            if self.config.get("device_mesh") is not None and self.config.get("load_from_rank0", False):
                weight_dict = self._load_weights_from_rank0(weight_dict, is_weight_loader)
292

293
294
295
            if hasattr(self, "adapter_weights_dict"):
                weight_dict.update(self.adapter_weights_dict)

gushiqiao's avatar
gushiqiao committed
296
            self.original_weight_dict = weight_dict
lijiaqi2's avatar
lijiaqi2 committed
297
298
        else:
            self.original_weight_dict = weight_dict
299

gushiqiao's avatar
gushiqiao committed
300
        # Initialize weight containers
helloyongyang's avatar
helloyongyang committed
301
        self.pre_weight = self.pre_weight_class(self.config)
302
303
304
305
        if self.lazy_load:
            self.transformer_weights = self.transformer_weight_class(self.config, self.lazy_load_path)
        else:
            self.transformer_weights = self.transformer_weight_class(self.config)
306
        if not self._should_init_empty_model():
307
            self._apply_weights()
gushiqiao's avatar
gushiqiao committed
308

309
310
311
312
313
    def _apply_weights(self, weight_dict=None):
        if weight_dict is not None:
            self.original_weight_dict = weight_dict
            del weight_dict
            gc.collect()
gushiqiao's avatar
gushiqiao committed
314
        # Load weights into containers
315
        self.pre_weight.load(self.original_weight_dict)
gushiqiao's avatar
gushiqiao committed
316
        self.transformer_weights.load(self.original_weight_dict)
helloyongyang's avatar
helloyongyang committed
317

gushiqiao's avatar
gushiqiao committed
318
319
320
321
        del self.original_weight_dict
        torch.cuda.empty_cache()
        gc.collect()

322
323
    def _load_weights_from_rank0(self, weight_dict, is_weight_loader):
        logger.info("Loading distributed weights")
gushiqiao's avatar
gushiqiao committed
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
        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()])

        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
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
            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
376
377

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

gushiqiao's avatar
gushiqiao committed
379
380
        return distributed_weight_dict

helloyongyang's avatar
helloyongyang committed
381
382
383
    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
384
        self.transformer_infer = self.transformer_infer_class(self.config)
385
        if hasattr(self.transformer_infer, "offload_manager"):
386
387
388
389
390
391
392
393
            self._init_offload_manager()

    def _init_offload_manager(self):
        self.transformer_infer.offload_manager.init_cuda_buffer(self.transformer_weights.offload_block_cuda_buffers, self.transformer_weights.offload_phase_cuda_buffers)
        if self.lazy_load:
            self.transformer_infer.offload_manager.init_cpu_buffer(self.transformer_weights.offload_block_cpu_buffers, self.transformer_weights.offload_phase_cpu_buffers)
            if self.config.get("warm_up_cpu_buffers", False):
                self.transformer_infer.offload_manager.warm_up_cpu_buffers(self.transformer_weights.blocks_num)
helloyongyang's avatar
helloyongyang committed
394
395
396

    def set_scheduler(self, scheduler):
        self.scheduler = scheduler
397
398
        self.pre_infer.set_scheduler(scheduler)
        self.post_infer.set_scheduler(scheduler)
helloyongyang's avatar
helloyongyang committed
399
400
        self.transformer_infer.set_scheduler(scheduler)

TorynCurtis's avatar
TorynCurtis committed
401
402
403
404
405
406
407
408
    def to_cpu(self):
        self.pre_weight.to_cpu()
        self.transformer_weights.to_cpu()

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

helloyongyang's avatar
helloyongyang committed
409
410
    @torch.no_grad()
    def infer(self, inputs):
411
        if self.cpu_offload:
412
            if self.offload_granularity == "model" and self.scheduler.step_index == 0 and "wan2.2_moe" not in self.config["model_cls"]:
413
414
415
                self.to_cuda()
            elif self.offload_granularity != "model":
                self.pre_weight.to_cuda()
gushiqiao's avatar
gushiqiao committed
416
                self.transformer_weights.non_block_weights_to_cuda()
417

418
        if self.config["enable_cfg"]:
helloyongyang's avatar
helloyongyang committed
419
420
421
422
423
424
425
            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:
helloyongyang's avatar
helloyongyang committed
426
                    noise_pred = self._infer_cond_uncond(inputs, infer_condition=True)
helloyongyang's avatar
helloyongyang committed
427
                else:
helloyongyang's avatar
helloyongyang committed
428
                    noise_pred = self._infer_cond_uncond(inputs, infer_condition=False)
helloyongyang's avatar
helloyongyang committed
429

helloyongyang's avatar
helloyongyang committed
430
431
432
433
434
435
                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 ====================
helloyongyang's avatar
helloyongyang committed
436
437
                noise_pred_cond = self._infer_cond_uncond(inputs, infer_condition=True)
                noise_pred_uncond = self._infer_cond_uncond(inputs, infer_condition=False)
gushiqiao's avatar
gushiqiao committed
438

helloyongyang's avatar
helloyongyang committed
439
440
441
            self.scheduler.noise_pred = noise_pred_uncond + self.scheduler.sample_guide_scale * (noise_pred_cond - noise_pred_uncond)
        else:
            # ==================== No CFG ====================
helloyongyang's avatar
helloyongyang committed
442
            self.scheduler.noise_pred = self._infer_cond_uncond(inputs, infer_condition=True)
443
444

        if self.cpu_offload:
445
            if self.offload_granularity == "model" and self.scheduler.step_index == self.scheduler.infer_steps - 1 and "wan2.2_moe" not in self.config["model_cls"]:
446
447
                self.to_cpu()
            elif self.offload_granularity != "model":
root's avatar
root committed
448
                self.pre_weight.to_cpu()
gushiqiao's avatar
gushiqiao committed
449
                self.transformer_weights.non_block_weights_to_cpu()
gushiqiao's avatar
gushiqiao committed
450

Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
451
    @compiled_method()
452
    @torch.no_grad()
helloyongyang's avatar
helloyongyang committed
453
454
455
456
    def _infer_cond_uncond(self, inputs, infer_condition=True):
        self.scheduler.infer_condition = infer_condition

        pre_infer_out = self.pre_infer.infer(self.pre_weight, inputs)
helloyongyang's avatar
helloyongyang committed
457
458
459
460
461
462
463
464
465

        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)

gushiqiao's avatar
gushiqiao committed
466
        noise_pred = self.post_infer.infer(x, pre_infer_out)[0]
helloyongyang's avatar
helloyongyang committed
467
468
469
470
471
472
473
474
475

        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):
helloyongyang's avatar
helloyongyang committed
476
        x = pre_infer_out.x
helloyongyang's avatar
helloyongyang committed
477
478
479
480
481
        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:
helloyongyang's avatar
helloyongyang committed
482
            x = F.pad(x, (0, 0, 0, padding_size))
helloyongyang's avatar
helloyongyang committed
483

helloyongyang's avatar
helloyongyang committed
484
        pre_infer_out.x = torch.chunk(x, world_size, dim=0)[cur_rank]
helloyongyang's avatar
helloyongyang committed
485

486
        if self.config["model_cls"] in ["wan2.2", "wan2.2_audio"] and self.config["task"] in ["i2v", "s2v"]:
helloyongyang's avatar
helloyongyang committed
487
488
489
490
491
492
493
            embed, embed0 = pre_infer_out.embed, pre_infer_out.embed0

            padding_size = (world_size - (embed.shape[0] % world_size)) % world_size
            if padding_size > 0:
                embed = F.pad(embed, (0, 0, 0, padding_size))
                embed0 = F.pad(embed0, (0, 0, 0, 0, 0, padding_size))

helloyongyang's avatar
helloyongyang committed
494
495
            pre_infer_out.embed = torch.chunk(embed, world_size, dim=0)[cur_rank]
            pre_infer_out.embed0 = torch.chunk(embed0, world_size, dim=0)[cur_rank]
helloyongyang's avatar
helloyongyang committed
496
497
498
499
500
501
502
503
504

        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)
        combined_output = torch.cat(gathered_x, dim=0)
helloyongyang's avatar
helloyongyang committed
505
        return combined_output