wan_runner.py 22.2 KB
Newer Older
1
import gc
PengGao's avatar
PengGao committed
2
3
import os

helloyongyang's avatar
helloyongyang committed
4
5
import numpy as np
import torch
6
import torch.distributed as dist
PengGao's avatar
PengGao committed
7
import torchvision.transforms.functional as TF
helloyongyang's avatar
helloyongyang committed
8
from PIL import Image
PengGao's avatar
PengGao committed
9
10
11
12
13
from loguru import logger

from lightx2v.models.input_encoders.hf.t5.model import T5EncoderModel
from lightx2v.models.input_encoders.hf.xlm_roberta.model import CLIPModel
from lightx2v.models.networks.wan.lora_adapter import WanLoraWrapper
helloyongyang's avatar
helloyongyang committed
14
from lightx2v.models.networks.wan.model import WanModel
helloyongyang's avatar
helloyongyang committed
15
from lightx2v.models.runners.default_runner import DefaultRunner
gushiqiao's avatar
gushiqiao committed
16
from lightx2v.models.schedulers.wan.changing_resolution.scheduler import (
17
    WanScheduler4ChangingResolutionInterface,
gushiqiao's avatar
gushiqiao committed
18
)
19
from lightx2v.models.schedulers.wan.feature_caching.scheduler import (
20
    WanSchedulerCaching,
21
    WanSchedulerTaylorCaching,
22
)
PengGao's avatar
PengGao committed
23
from lightx2v.models.schedulers.wan.scheduler import WanScheduler
helloyongyang's avatar
helloyongyang committed
24
from lightx2v.models.video_encoders.hf.wan.vae import WanVAE
25
from lightx2v.models.video_encoders.hf.wan.vae_2_2 import Wan2_2_VAE
26
from lightx2v.models.video_encoders.hf.wan.vae_tiny import Wan2_2_VAE_tiny, WanVAE_tiny
27
from lightx2v.utils.envs import *
PengGao's avatar
PengGao committed
28
29
30
from lightx2v.utils.registry_factory import RUNNER_REGISTER
from lightx2v.utils.utils import *
from lightx2v.utils.utils import best_output_size, cache_video
helloyongyang's avatar
helloyongyang committed
31
32
33
34
35
36


@RUNNER_REGISTER("wan2.1")
class WanRunner(DefaultRunner):
    def __init__(self, config):
        super().__init__(config)
37
38
39
40
        self.vae_cls = WanVAE
        self.tiny_vae_cls = WanVAE_tiny
        self.vae_name = "Wan2.1_VAE.pth"
        self.tiny_vae_name = "taew2_1.pth"
helloyongyang's avatar
helloyongyang committed
41

42
43
44
45
46
47
    def load_transformer(self):
        model = WanModel(
            self.config.model_path,
            self.config,
            self.init_device,
        )
48
        if self.config.get("lora_configs") and self.config.lora_configs:
49
            assert not self.config.get("dit_quantized", False) or self.config.mm_config.get("weight_auto_quant", False)
50
            lora_wrapper = WanLoraWrapper(model)
51
52
53
            for lora_config in self.config.lora_configs:
                lora_path = lora_config["path"]
                strength = lora_config.get("strength", 1.0)
GoatWu's avatar
GoatWu committed
54
                lora_name = lora_wrapper.load_lora(lora_path)
55
56
                lora_wrapper.apply_lora(lora_name, strength)
                logger.info(f"Loaded LoRA: {lora_name} with strength: {strength}")
57
58
        return model

59
    def load_image_encoder(self):
helloyongyang's avatar
helloyongyang committed
60
        image_encoder = None
gushiqiao's avatar
gushiqiao committed
61
        if self.config.task in ["i2v", "flf2v"] and self.config.get("use_image_encoder", True):
gushiqiao's avatar
gushiqiao committed
62
63
64
65
66
67
            # offload config
            clip_offload = self.config.get("clip_cpu_offload", self.config.get("cpu_offload", False))
            if clip_offload:
                clip_device = torch.device("cpu")
            else:
                clip_device = torch.device("cuda")
gushiqiao's avatar
gushiqiao committed
68
69
70
71
72
            # quant_config
            clip_quantized = self.config.get("clip_quantized", False)
            if clip_quantized:
                clip_quant_scheme = self.config.get("clip_quant_scheme", None)
                assert clip_quant_scheme is not None
gushiqiao's avatar
gushiqiao committed
73
                tmp_clip_quant_scheme = clip_quant_scheme.split("-")[0]
74
                clip_model_name = f"clip-{tmp_clip_quant_scheme}.pth"
75
                clip_quantized_ckpt = find_torch_model_path(self.config, "clip_quantized_ckpt", clip_model_name)
76
                clip_original_ckpt = None
gushiqiao's avatar
gushiqiao committed
77
78
79
            else:
                clip_quantized_ckpt = None
                clip_quant_scheme = None
80
                clip_model_name = "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"
81
                clip_original_ckpt = find_torch_model_path(self.config, "clip_original_ckpt", clip_model_name)
gushiqiao's avatar
gushiqiao committed
82

83
84
            image_encoder = CLIPModel(
                dtype=torch.float16,
gushiqiao's avatar
gushiqiao committed
85
                device=clip_device,
86
                checkpoint_path=clip_original_ckpt,
gushiqiao's avatar
gushiqiao committed
87
88
89
                clip_quantized=clip_quantized,
                clip_quantized_ckpt=clip_quantized_ckpt,
                quant_scheme=clip_quant_scheme,
gushiqiao's avatar
gushiqiao committed
90
                cpu_offload=clip_offload,
gushiqiao's avatar
gushiqiao committed
91
                use_31_block=self.config.get("use_31_block", True),
92
                load_from_rank0=self.config.get("load_from_rank0", False),
93
            )
94

95
        return image_encoder
helloyongyang's avatar
helloyongyang committed
96

97
    def load_text_encoder(self):
gushiqiao's avatar
gushiqiao committed
98
        # offload config
99
        t5_offload = self.config.get("t5_cpu_offload", self.config.get("cpu_offload"))
gushiqiao's avatar
gushiqiao committed
100
101
102
103
        if t5_offload:
            t5_device = torch.device("cpu")
        else:
            t5_device = torch.device("cuda")
gushiqiao's avatar
gushiqiao committed
104
105
106
107
108
109

        # quant_config
        t5_quantized = self.config.get("t5_quantized", False)
        if t5_quantized:
            t5_quant_scheme = self.config.get("t5_quant_scheme", None)
            assert t5_quant_scheme is not None
110
111
            tmp_t5_quant_scheme = t5_quant_scheme.split("-")[0]
            t5_model_name = f"models_t5_umt5-xxl-enc-{tmp_t5_quant_scheme}.pth"
112
            t5_quantized_ckpt = find_torch_model_path(self.config, "t5_quantized_ckpt", t5_model_name)
113
            t5_original_ckpt = None
gushiqiao's avatar
gushiqiao committed
114
            tokenizer_path = os.path.join(os.path.dirname(t5_quantized_ckpt), "google/umt5-xxl")
gushiqiao's avatar
gushiqiao committed
115
116
117
        else:
            t5_quant_scheme = None
            t5_quantized_ckpt = None
118
            t5_model_name = "models_t5_umt5-xxl-enc-bf16.pth"
119
            t5_original_ckpt = find_torch_model_path(self.config, "t5_original_ckpt", t5_model_name)
gushiqiao's avatar
gushiqiao committed
120
            tokenizer_path = os.path.join(os.path.dirname(t5_original_ckpt), "google/umt5-xxl")
gushiqiao's avatar
Fix  
gushiqiao committed
121

helloyongyang's avatar
helloyongyang committed
122
123
124
        text_encoder = T5EncoderModel(
            text_len=self.config["text_len"],
            dtype=torch.bfloat16,
gushiqiao's avatar
gushiqiao committed
125
            device=t5_device,
126
            checkpoint_path=t5_original_ckpt,
gushiqiao's avatar
gushiqiao committed
127
            tokenizer_path=tokenizer_path,
helloyongyang's avatar
helloyongyang committed
128
            shard_fn=None,
gushiqiao's avatar
gushiqiao committed
129
            cpu_offload=t5_offload,
130
            offload_granularity=self.config.get("t5_offload_granularity", "model"),  # support ['model', 'block']
gushiqiao's avatar
gushiqiao committed
131
132
133
            t5_quantized=t5_quantized,
            t5_quantized_ckpt=t5_quantized_ckpt,
            quant_scheme=t5_quant_scheme,
134
            load_from_rank0=self.config.get("load_from_rank0", False),
helloyongyang's avatar
helloyongyang committed
135
136
        )
        text_encoders = [text_encoder]
137
        return text_encoders
helloyongyang's avatar
helloyongyang committed
138

139
    def load_vae_encoder(self):
140
141
142
143
144
145
146
        # offload config
        vae_offload = self.config.get("vae_cpu_offload", self.config.get("cpu_offload"))
        if vae_offload:
            vae_device = torch.device("cpu")
        else:
            vae_device = torch.device("cuda")

147
        vae_config = {
148
            "vae_pth": find_torch_model_path(self.config, "vae_pth", self.vae_name),
149
            "device": vae_device,
150
            "parallel": self.config.parallel,
151
            "use_tiling": self.config.get("use_tiling_vae", False),
152
            "cpu_offload": vae_offload,
153
            "dtype": GET_DTYPE(),
154
            "load_from_rank0": self.config.get("load_from_rank0", False),
155
        }
gushiqiao's avatar
gushiqiao committed
156
        if self.config.task not in ["i2v", "flf2v", "vace"]:
157
158
            return None
        else:
159
            return self.vae_cls(**vae_config)
160
161

    def load_vae_decoder(self):
162
163
164
165
166
167
168
        # offload config
        vae_offload = self.config.get("vae_cpu_offload", self.config.get("cpu_offload"))
        if vae_offload:
            vae_device = torch.device("cpu")
        else:
            vae_device = torch.device("cuda")

169
        vae_config = {
170
            "vae_pth": find_torch_model_path(self.config, "vae_pth", self.vae_name),
171
            "device": vae_device,
172
            "parallel": self.config.parallel,
173
            "use_tiling": self.config.get("use_tiling_vae", False),
174
            "cpu_offload": vae_offload,
175
            "dtype": GET_DTYPE(),
176
            "load_from_rank0": self.config.get("load_from_rank0", False),
177
        }
helloyongyang's avatar
helloyongyang committed
178
        if self.config.get("use_tiny_vae", False):
179
180
            tiny_vae_path = find_torch_model_path(self.config, "tiny_vae_path", self.tiny_vae_name)
            vae_decoder = self.tiny_vae_cls(vae_pth=tiny_vae_path, device=self.init_device, need_scaled=self.config.get("need_scaled", False)).to("cuda")
181
        else:
182
            vae_decoder = self.vae_cls(**vae_config)
183
        return vae_decoder
helloyongyang's avatar
helloyongyang committed
184

185
    def load_vae(self):
gushiqiao's avatar
gushiqiao committed
186
        vae_encoder = self.load_vae_encoder()
helloyongyang's avatar
helloyongyang committed
187
        if vae_encoder is None or self.config.get("use_tiny_vae", False):
gushiqiao's avatar
gushiqiao committed
188
189
190
191
            vae_decoder = self.load_vae_decoder()
        else:
            vae_decoder = vae_encoder
        return vae_encoder, vae_decoder
helloyongyang's avatar
helloyongyang committed
192
193

    def init_scheduler(self):
194
195
196
197
        if self.config.feature_caching == "NoCaching":
            scheduler_class = WanScheduler
        elif self.config.feature_caching == "TaylorSeer":
            scheduler_class = WanSchedulerTaylorCaching
Musisoul's avatar
Musisoul committed
198
        elif self.config.feature_caching in ["Tea", "Ada", "Custom", "FirstBlock", "DualBlock", "DynamicBlock", "Mag"]:
199
200
201
202
            scheduler_class = WanSchedulerCaching
        else:
            raise NotImplementedError(f"Unsupported feature_caching type: {self.config.feature_caching}")

203
        if self.config.get("changing_resolution", False):
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
204
            self.scheduler = WanScheduler4ChangingResolutionInterface(scheduler_class, self.config)
helloyongyang's avatar
helloyongyang committed
205
        else:
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
206
            self.scheduler = scheduler_class(self.config)
helloyongyang's avatar
helloyongyang committed
207

gushiqiao's avatar
gushiqiao committed
208
    def run_text_encoder(self, text, img=None):
gushiqiao's avatar
gushiqiao committed
209
        if self.config.get("lazy_load", False) or self.config.get("unload_modules", False):
210
            self.text_encoders = self.load_text_encoder()
211
        n_prompt = self.config.get("negative_prompt", "")
212
213
214
215
216
217

        if self.config["cfg_parallel"]:
            cfg_p_group = self.config["device_mesh"].get_group(mesh_dim="cfg_p")
            cfg_p_rank = dist.get_rank(cfg_p_group)
            if cfg_p_rank == 0:
                context = self.text_encoders[0].infer([text])
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
218
                context = torch.stack([torch.cat([u, u.new_zeros(self.config["text_len"] - u.size(0), u.size(1))]) for u in context])
219
220
221
                text_encoder_output = {"context": context}
            else:
                context_null = self.text_encoders[0].infer([n_prompt])
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
222
                context_null = torch.stack([torch.cat([u, u.new_zeros(self.config["text_len"] - u.size(0), u.size(1))]) for u in context_null])
223
224
225
                text_encoder_output = {"context_null": context_null}
        else:
            context = self.text_encoders[0].infer([text])
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
226
            context = torch.stack([torch.cat([u, u.new_zeros(self.config["text_len"] - u.size(0), u.size(1))]) for u in context])
227
            context_null = self.text_encoders[0].infer([n_prompt])
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
228
            context_null = torch.stack([torch.cat([u, u.new_zeros(self.config["text_len"] - u.size(0), u.size(1))]) for u in context_null])
229
230
231
232
233
            text_encoder_output = {
                "context": context,
                "context_null": context_null,
            }

gushiqiao's avatar
gushiqiao committed
234
        if self.config.get("lazy_load", False) or self.config.get("unload_modules", False):
235
236
237
            del self.text_encoders[0]
            torch.cuda.empty_cache()
            gc.collect()
238

helloyongyang's avatar
helloyongyang committed
239
240
        return text_encoder_output

gushiqiao's avatar
gushiqiao committed
241
    def run_image_encoder(self, first_frame, last_frame=None):
gushiqiao's avatar
gushiqiao committed
242
        if self.config.get("lazy_load", False) or self.config.get("unload_modules", False):
243
            self.image_encoder = self.load_image_encoder()
gushiqiao's avatar
gushiqiao committed
244
        if last_frame is None:
helloyongyang's avatar
helloyongyang committed
245
            clip_encoder_out = self.image_encoder.visual([first_frame]).squeeze(0).to(GET_DTYPE())
gushiqiao's avatar
gushiqiao committed
246
        else:
helloyongyang's avatar
helloyongyang committed
247
            clip_encoder_out = self.image_encoder.visual([first_frame, last_frame]).squeeze(0).to(GET_DTYPE())
gushiqiao's avatar
gushiqiao committed
248
        if self.config.get("lazy_load", False) or self.config.get("unload_modules", False):
249
250
251
            del self.image_encoder
            torch.cuda.empty_cache()
            gc.collect()
252
253
        return clip_encoder_out

gushiqiao's avatar
gushiqiao committed
254
    def run_vae_encoder(self, first_frame, last_frame=None):
helloyongyang's avatar
helloyongyang committed
255
        h, w = first_frame.shape[2:]
helloyongyang's avatar
helloyongyang committed
256
        aspect_ratio = h / w
257
258
259
        max_area = self.config.target_height * self.config.target_width
        lat_h = round(np.sqrt(max_area * aspect_ratio) // self.config.vae_stride[1] // self.config.patch_size[1] * self.config.patch_size[1])
        lat_w = round(np.sqrt(max_area / aspect_ratio) // self.config.vae_stride[2] // self.config.patch_size[2] * self.config.patch_size[2])
260
261

        if self.config.get("changing_resolution", False):
gushiqiao's avatar
gushiqiao committed
262
            assert last_frame is None
263
            self.config.lat_h, self.config.lat_w = lat_h, lat_w
264
265
            vae_encode_out_list = []
            for i in range(len(self.config["resolution_rate"])):
266
267
268
269
                lat_h, lat_w = (
                    int(self.config.lat_h * self.config.resolution_rate[i]) // 2 * 2,
                    int(self.config.lat_w * self.config.resolution_rate[i]) // 2 * 2,
                )
gushiqiao's avatar
gushiqiao committed
270
271
                vae_encode_out_list.append(self.get_vae_encoder_output(first_frame, lat_h, lat_w))
            vae_encode_out_list.append(self.get_vae_encoder_output(first_frame, self.config.lat_h, self.config.lat_w))
272
            return vae_encode_out_list
273
        else:
gushiqiao's avatar
gushiqiao committed
274
            if last_frame is not None:
helloyongyang's avatar
helloyongyang committed
275
276
                first_frame_size = first_frame.shape[2:]
                last_frame_size = last_frame.shape[2:]
gushiqiao's avatar
gushiqiao committed
277
278
279
280
281
282
283
                if first_frame_size != last_frame_size:
                    last_frame_resize_ratio = max(first_frame_size[0] / last_frame_size[0], first_frame_size[1] / last_frame_size[1])
                    last_frame_size = [
                        round(last_frame_size[0] * last_frame_resize_ratio),
                        round(last_frame_size[1] * last_frame_resize_ratio),
                    ]
                    last_frame = TF.center_crop(last_frame, last_frame_size)
284
            self.config.lat_h, self.config.lat_w = lat_h, lat_w
gushiqiao's avatar
gushiqiao committed
285
            vae_encoder_out = self.get_vae_encoder_output(first_frame, lat_h, lat_w, last_frame)
286
            return vae_encoder_out
287

gushiqiao's avatar
gushiqiao committed
288
    def get_vae_encoder_output(self, first_frame, lat_h, lat_w, last_frame=None):
289
290
        h = lat_h * self.config.vae_stride[1]
        w = lat_w * self.config.vae_stride[2]
291
292
293
294
295
296
297
        msk = torch.ones(
            1,
            self.config.target_video_length,
            lat_h,
            lat_w,
            device=torch.device("cuda"),
        )
gushiqiao's avatar
gushiqiao committed
298
299
300
301
302
        if last_frame is not None:
            msk[:, 1:-1] = 0
        else:
            msk[:, 1:] = 0

helloyongyang's avatar
helloyongyang committed
303
304
305
        msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1)
        msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)
        msk = msk.transpose(1, 2)[0]
gushiqiao's avatar
gushiqiao committed
306

gushiqiao's avatar
gushiqiao committed
307
        if self.config.get("lazy_load", False) or self.config.get("unload_modules", False):
308
            self.vae_encoder = self.load_vae_encoder()
gushiqiao's avatar
gushiqiao committed
309
310
311
312

        if last_frame is not None:
            vae_input = torch.concat(
                [
helloyongyang's avatar
helloyongyang committed
313
                    torch.nn.functional.interpolate(first_frame.cpu(), size=(h, w), mode="bicubic").transpose(0, 1),
gushiqiao's avatar
gushiqiao committed
314
                    torch.zeros(3, self.config.target_video_length - 2, h, w),
helloyongyang's avatar
helloyongyang committed
315
                    torch.nn.functional.interpolate(last_frame.cpu(), size=(h, w), mode="bicubic").transpose(0, 1),
gushiqiao's avatar
gushiqiao committed
316
317
318
319
320
321
                ],
                dim=1,
            ).cuda()
        else:
            vae_input = torch.concat(
                [
helloyongyang's avatar
helloyongyang committed
322
                    torch.nn.functional.interpolate(first_frame.cpu(), size=(h, w), mode="bicubic").transpose(0, 1),
gushiqiao's avatar
gushiqiao committed
323
324
325
326
327
                    torch.zeros(3, self.config.target_video_length - 1, h, w),
                ],
                dim=1,
            ).cuda()

328
        vae_encoder_out = self.vae_encoder.encode(vae_input.unsqueeze(0).to(GET_DTYPE()))
gushiqiao's avatar
gushiqiao committed
329

gushiqiao's avatar
gushiqiao committed
330
        if self.config.get("lazy_load", False) or self.config.get("unload_modules", False):
331
332
333
            del self.vae_encoder
            torch.cuda.empty_cache()
            gc.collect()
334
        vae_encoder_out = torch.concat([msk, vae_encoder_out]).to(GET_DTYPE())
335
        return vae_encoder_out
336

gushiqiao's avatar
gushiqiao committed
337
    def get_encoder_output_i2v(self, clip_encoder_out, vae_encoder_out, text_encoder_output, img=None):
338
339
        image_encoder_output = {
            "clip_encoder_out": clip_encoder_out,
340
            "vae_encoder_out": vae_encoder_out,
341
        }
342
343
344
345
        return {
            "text_encoder_output": text_encoder_output,
            "image_encoder_output": image_encoder_output,
        }
helloyongyang's avatar
helloyongyang committed
346
347

    def set_target_shape(self):
348
        num_channels_latents = self.config.get("num_channels_latents", 16)
gushiqiao's avatar
gushiqiao committed
349
        if self.config.task in ["i2v", "flf2v"]:
350
351
            self.config.target_shape = (
                num_channels_latents,
352
                (self.config.target_video_length - 1) // self.config.vae_stride[0] + 1,
353
354
355
                self.config.lat_h,
                self.config.lat_w,
            )
helloyongyang's avatar
helloyongyang committed
356
357
        elif self.config.task == "t2v":
            self.config.target_shape = (
358
                num_channels_latents,
359
                (self.config.target_video_length - 1) // self.config.vae_stride[0] + 1,
helloyongyang's avatar
helloyongyang committed
360
361
362
                int(self.config.target_height) // self.config.vae_stride[1],
                int(self.config.target_width) // self.config.vae_stride[2],
            )
363
364

    def save_video_func(self, images):
365
366
367
368
369
370
371
372
        cache_video(
            tensor=images,
            save_file=self.config.save_video_path,
            fps=self.config.get("fps", 16),
            nrow=1,
            normalize=True,
            value_range=(-1, 1),
        )
helloyongyang's avatar
helloyongyang committed
373
374
375
376
377
378
379
380
381
382
383
384


class MultiModelStruct:
    def __init__(self, model_list, config, boundary=0.875, num_train_timesteps=1000):
        self.model = model_list  # [high_noise_model, low_noise_model]
        assert len(self.model) == 2, "MultiModelStruct only supports 2 models now."
        self.config = config
        self.boundary = boundary
        self.boundary_timestep = self.boundary * num_train_timesteps
        self.cur_model_index = -1
        logger.info(f"boundary: {self.boundary}, boundary_timestep: {self.boundary_timestep}")

wangshankun's avatar
wangshankun committed
385
386
387
388
    @property
    def device(self):
        return self.model[self.cur_model_index].device

helloyongyang's avatar
helloyongyang committed
389
390
391
392
393
394
395
396
397
398
399
400
401
    def set_scheduler(self, shared_scheduler):
        self.scheduler = shared_scheduler
        for model in self.model:
            model.set_scheduler(shared_scheduler)

    def infer(self, inputs):
        self.get_current_model_index()
        self.model[self.cur_model_index].infer(inputs)

    def get_current_model_index(self):
        if self.scheduler.timesteps[self.scheduler.step_index] >= self.boundary_timestep:
            logger.info(f"using - HIGH - noise model at step_index {self.scheduler.step_index + 1}")
            self.scheduler.sample_guide_scale = self.config.sample_guide_scale[0]
402
403
404
405
406
407
            if self.config.get("cpu_offload", False) and self.config.get("offload_granularity", "block") == "model":
                if self.cur_model_index == -1:
                    self.to_cuda(model_index=0)
                elif self.cur_model_index == 1:  # 1 -> 0
                    self.offload_cpu(model_index=1)
                    self.to_cuda(model_index=0)
helloyongyang's avatar
helloyongyang committed
408
409
410
411
            self.cur_model_index = 0
        else:
            logger.info(f"using - LOW - noise model at step_index {self.scheduler.step_index + 1}")
            self.scheduler.sample_guide_scale = self.config.sample_guide_scale[1]
412
413
414
415
416
417
            if self.config.get("cpu_offload", False) and self.config.get("offload_granularity", "block") == "model":
                if self.cur_model_index == -1:
                    self.to_cuda(model_index=1)
                elif self.cur_model_index == 0:  # 0 -> 1
                    self.offload_cpu(model_index=0)
                    self.to_cuda(model_index=1)
helloyongyang's avatar
helloyongyang committed
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
            self.cur_model_index = 1

    def offload_cpu(self, model_index):
        self.model[model_index].to_cpu()

    def to_cuda(self, model_index):
        self.model[model_index].to_cuda()


@RUNNER_REGISTER("wan2.2_moe")
class Wan22MoeRunner(WanRunner):
    def __init__(self, config):
        super().__init__(config)

    def load_transformer(self):
        # encoder -> high_noise_model -> low_noise_model -> vae -> video_output
helloyongyang's avatar
helloyongyang committed
434
        high_noise_model = WanModel(
helloyongyang's avatar
helloyongyang committed
435
436
437
438
            os.path.join(self.config.model_path, "high_noise_model"),
            self.config,
            self.init_device,
        )
helloyongyang's avatar
helloyongyang committed
439
        low_noise_model = WanModel(
helloyongyang's avatar
helloyongyang committed
440
441
442
443
            os.path.join(self.config.model_path, "low_noise_model"),
            self.config,
            self.init_device,
        )
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464

        if self.config.get("lora_configs") and self.config.lora_configs:
            assert not self.config.get("dit_quantized", False) or self.config.mm_config.get("weight_auto_quant", False)

            for lora_config in self.config.lora_configs:
                lora_path = lora_config["path"]
                strength = lora_config.get("strength", 1.0)
                base_name = os.path.basename(lora_path)
                if base_name.startswith("high"):
                    lora_wrapper = WanLoraWrapper(high_noise_model)
                    lora_name = lora_wrapper.load_lora(lora_path)
                    lora_wrapper.apply_lora(lora_name, strength)
                    logger.info(f"Loaded LoRA: {lora_name} with strength: {strength}")
                elif base_name.startswith("low"):
                    lora_wrapper = WanLoraWrapper(low_noise_model)
                    lora_name = lora_wrapper.load_lora(lora_path)
                    lora_wrapper.apply_lora(lora_name, strength)
                    logger.info(f"Loaded LoRA: {lora_name} with strength: {strength}")
                else:
                    raise ValueError(f"Unsupported LoRA path: {lora_path}")

helloyongyang's avatar
helloyongyang committed
465
        return MultiModelStruct([high_noise_model, low_noise_model], self.config, self.config.boundary)
466
467
468
469
470
471


@RUNNER_REGISTER("wan2.2")
class Wan22DenseRunner(WanRunner):
    def __init__(self, config):
        super().__init__(config)
472
        self.vae_encoder_need_img_original = True
473
474
475
476
        self.vae_cls = Wan2_2_VAE
        self.tiny_vae_cls = Wan2_2_VAE_tiny
        self.vae_name = "Wan2.2_VAE.pth"
        self.tiny_vae_name = "taew2_2.pth"
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500

    def run_vae_encoder(self, img):
        max_area = self.config.target_height * self.config.target_width
        ih, iw = img.height, img.width
        dh, dw = self.config.patch_size[1] * self.config.vae_stride[1], self.config.patch_size[2] * self.config.vae_stride[2]
        ow, oh = best_output_size(iw, ih, dw, dh, max_area)

        scale = max(ow / iw, oh / ih)
        img = img.resize((round(iw * scale), round(ih * scale)), Image.LANCZOS)

        # center-crop
        x1 = (img.width - ow) // 2
        y1 = (img.height - oh) // 2
        img = img.crop((x1, y1, x1 + ow, y1 + oh))
        assert img.width == ow and img.height == oh

        # to tensor
        img = TF.to_tensor(img).sub_(0.5).div_(0.5).cuda().unsqueeze(1)
        vae_encoder_out = self.get_vae_encoder_output(img)
        self.config.lat_w, self.config.lat_h = ow // self.config.vae_stride[2], oh // self.config.vae_stride[1]

        return vae_encoder_out

    def get_vae_encoder_output(self, img):
501
        z = self.vae_encoder.encode(img.to(GET_DTYPE()))
502
        return z