__main__.py 17.2 KB
Newer Older
helloyongyang's avatar
helloyongyang committed
1
import argparse
lijiaqi2's avatar
lijiaqi2 committed
2
from contextlib import contextmanager
helloyongyang's avatar
helloyongyang committed
3
4
5
6
7
8
import torch
import torch.distributed as dist
import os
import time
import gc
import json
helloyongyang's avatar
helloyongyang committed
9
import torchvision
helloyongyang's avatar
helloyongyang committed
10
11
12
13
14
15
import torchvision.transforms.functional as TF
import numpy as np
from PIL import Image
from lightx2v.text2v.models.text_encoders.hf.llama.model import TextEncoderHFLlamaModel
from lightx2v.text2v.models.text_encoders.hf.clip.model import TextEncoderHFClipModel
from lightx2v.text2v.models.text_encoders.hf.t5.model import T5EncoderModel
helloyongyang's avatar
helloyongyang committed
16
from lightx2v.text2v.models.text_encoders.hf.llava.model import TextEncoderHFLlavaModel
helloyongyang's avatar
helloyongyang committed
17
18
19
20
21
22
23
24

from lightx2v.text2v.models.schedulers.hunyuan.scheduler import HunyuanScheduler
from lightx2v.text2v.models.schedulers.hunyuan.feature_caching.scheduler import HunyuanSchedulerFeatureCaching
from lightx2v.text2v.models.schedulers.wan.scheduler import WanScheduler
from lightx2v.text2v.models.schedulers.wan.feature_caching.scheduler import WanSchedulerFeatureCaching

from lightx2v.text2v.models.networks.hunyuan.model import HunyuanModel
from lightx2v.text2v.models.networks.wan.model import WanModel
lijiaqi2's avatar
lijiaqi2 committed
25
26
from lightx2v.text2v.models.networks.wan.lora_adapter import WanLoraWrapper

helloyongyang's avatar
helloyongyang committed
27
28
29
30
31
32
33
34

from lightx2v.text2v.models.video_encoders.hf.autoencoder_kl_causal_3d.model import VideoEncoderKLCausal3DModel
from lightx2v.text2v.models.video_encoders.hf.wan.vae import WanVAE
from lightx2v.utils.utils import save_videos_grid, seed_all, cache_video
from lightx2v.common.ops import *
from lightx2v.image2v.models.wan.model import CLIPModel


lijiaqi2's avatar
lijiaqi2 committed
35
36
37
38
39
40
41
42
@contextmanager
def time_duration(label: str = ""):
    start_time = time.time()
    yield
    end_time = time.time()
    print(f"==> {label} start:{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))} cost {end_time - start_time:.2f} seconds")


helloyongyang's avatar
helloyongyang committed
43
def load_models(args, model_config):
Xinchi Huang's avatar
Xinchi Huang committed
44
    if model_config["parallel_attn_type"]:
helloyongyang's avatar
helloyongyang committed
45
46
47
48
49
50
51
52
53
        cur_rank = dist.get_rank()  # 获取当前进程的 rank
        torch.cuda.set_device(cur_rank)  # 设置当前进程的 CUDA 设备
    image_encoder = None
    if args.cpu_offload:
        init_device = torch.device("cpu")
    else:
        init_device = torch.device("cuda")

    if args.model_cls == "hunyuan":
helloyongyang's avatar
helloyongyang committed
54
55
56
57
        if args.task == "t2v":
            text_encoder_1 = TextEncoderHFLlamaModel(os.path.join(args.model_path, "text_encoder"), init_device)
        else:
            text_encoder_1 = TextEncoderHFLlavaModel(os.path.join(args.model_path, "text_encoder_i2v"), init_device)
helloyongyang's avatar
helloyongyang committed
58
59
        text_encoder_2 = TextEncoderHFClipModel(os.path.join(args.model_path, "text_encoder_2"), init_device)
        text_encoders = [text_encoder_1, text_encoder_2]
helloyongyang's avatar
helloyongyang committed
60
61
        model = HunyuanModel(args.model_path, model_config, init_device, args)
        vae_model = VideoEncoderKLCausal3DModel(args.model_path, dtype=torch.float16, device=init_device, args=args)
helloyongyang's avatar
helloyongyang committed
62
63
64
65
66

    elif args.model_cls == "wan2.1":
        text_encoder = T5EncoderModel(
            text_len=model_config["text_len"],
            dtype=torch.bfloat16,
TorynCurtis's avatar
TorynCurtis committed
67
            device=init_device,
helloyongyang's avatar
helloyongyang committed
68
69
70
71
72
            checkpoint_path=os.path.join(args.model_path, "models_t5_umt5-xxl-enc-bf16.pth"),
            tokenizer_path=os.path.join(args.model_path, "google/umt5-xxl"),
            shard_fn=None,
        )
        text_encoders = [text_encoder]
lijiaqi2's avatar
lijiaqi2 committed
73
74
75
76
77
78
79
80
81
82
83
84
85

        with time_duration("Load Wan Model"):
            model = WanModel(args.model_path, model_config, init_device)

        if args.lora_path:
            lora_wrapper = WanLoraWrapper(model)
            with time_duration("Load LoRA Model"):
                lora_name = lora_wrapper.load_lora(args.lora_path)
                lora_wrapper.apply_lora(lora_name, args.strength_model)
                print(f"Loaded LoRA: {lora_name}")

        with time_duration("Load WAN VAE Model"):
            vae_model = WanVAE(vae_pth=os.path.join(args.model_path, "Wan2.1_VAE.pth"), device=init_device, parallel=args.parallel_vae)
Dongz's avatar
Dongz committed
86
        if args.task == "i2v":
lijiaqi2's avatar
lijiaqi2 committed
87
88
89
90
91
92
93
            with time_duration("Load Image Encoder"):
                image_encoder = CLIPModel(
                    dtype=torch.float16,
                    device=init_device,
                    checkpoint_path=os.path.join(args.model_path, "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"),
                    tokenizer_path=os.path.join(args.model_path, "xlm-roberta-large"),
                )
helloyongyang's avatar
helloyongyang committed
94
95
96
97
98
99
    else:
        raise NotImplementedError(f"Unsupported model class: {args.model_cls}")

    return model, text_encoders, vae_model, image_encoder


helloyongyang's avatar
helloyongyang committed
100
def set_target_shape(args, image_encoder_output):
Dongz's avatar
Dongz committed
101
    if args.model_cls == "hunyuan":
helloyongyang's avatar
helloyongyang committed
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
        if args.task == "t2v":
            vae_scale_factor = 2 ** (4 - 1)
            args.target_shape = (
                1,
                16,
                (args.target_video_length - 1) // 4 + 1,
                int(args.target_height) // vae_scale_factor,
                int(args.target_width) // vae_scale_factor,
            )
        elif args.task == "i2v":
            vae_scale_factor = 2 ** (4 - 1)
            args.target_shape = (
                1,
                16,
                (args.target_video_length - 1) // 4 + 1,
                int(image_encoder_output["target_height"]) // vae_scale_factor,
                int(image_encoder_output["target_width"]) // vae_scale_factor,
            )
Dongz's avatar
Dongz committed
120
121
122
123
    elif args.model_cls == "wan2.1":
        if args.task == "i2v":
            args.target_shape = (16, 21, args.lat_h, args.lat_w)
        elif args.task == "t2v":
helloyongyang's avatar
helloyongyang committed
124
125
126
127
128
129
130
131
            args.target_shape = (
                16,
                (args.target_video_length - 1) // 4 + 1,
                int(args.target_height) // args.vae_stride[1],
                int(args.target_width) // args.vae_stride[2],
            )


helloyongyang's avatar
helloyongyang committed
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def generate_crop_size_list(base_size=256, patch_size=32, max_ratio=4.0):
    num_patches = round((base_size / patch_size) ** 2)
    assert max_ratio >= 1.0
    crop_size_list = []
    wp, hp = num_patches, 1
    while wp > 0:
        if max(wp, hp) / min(wp, hp) <= max_ratio:
            crop_size_list.append((wp * patch_size, hp * patch_size))
        if (hp + 1) * wp <= num_patches:
            hp += 1
        else:
            wp -= 1
    return crop_size_list


def get_closest_ratio(height: float, width: float, ratios: list, buckets: list):
    aspect_ratio = float(height) / float(width)
    diff_ratios = ratios - aspect_ratio

    if aspect_ratio >= 1:
        indices = [(index, x) for index, x in enumerate(diff_ratios) if x <= 0]
    else:
        indices = [(index, x) for index, x in enumerate(diff_ratios) if x > 0]

    closest_ratio_id = min(indices, key=lambda pair: abs(pair[1]))[0]
    closest_size = buckets[closest_ratio_id]
    closest_ratio = ratios[closest_ratio_id]

    return closest_size, closest_ratio


helloyongyang's avatar
helloyongyang committed
163
164
def run_image_encoder(args, image_encoder, vae_model):
    if args.model_cls == "hunyuan":
helloyongyang's avatar
helloyongyang committed
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
        img = Image.open(args.image_path).convert("RGB")
        origin_size = img.size

        i2v_resolution = "720p"
        if i2v_resolution == "720p":
            bucket_hw_base_size = 960
        elif i2v_resolution == "540p":
            bucket_hw_base_size = 720
        elif i2v_resolution == "360p":
            bucket_hw_base_size = 480
        else:
            raise ValueError(f"i2v_resolution: {i2v_resolution} must be in [360p, 540p, 720p]")

        crop_size_list = generate_crop_size_list(bucket_hw_base_size, 32)
        aspect_ratios = np.array([round(float(h) / float(w), 5) for h, w in crop_size_list])
        closest_size, closest_ratio = get_closest_ratio(origin_size[1], origin_size[0], aspect_ratios, crop_size_list)

        resize_param = min(closest_size)
        center_crop_param = closest_size

        ref_image_transform = torchvision.transforms.Compose(
            [torchvision.transforms.Resize(resize_param), torchvision.transforms.CenterCrop(center_crop_param), torchvision.transforms.ToTensor(), torchvision.transforms.Normalize([0.5], [0.5])]
        )

        semantic_image_pixel_values = [ref_image_transform(img)]
        semantic_image_pixel_values = torch.cat(semantic_image_pixel_values).unsqueeze(0).unsqueeze(2).to(torch.float16).to(torch.device("cuda"))

        img_latents = vae_model.encode(semantic_image_pixel_values, args).mode()

        scaling_factor = 0.476986
        img_latents.mul_(scaling_factor)

        target_height, target_width = closest_size

        return {"img": img, "img_latents": img_latents, "target_height": target_height, "target_width": target_width}

Dongz's avatar
Dongz committed
201
    elif args.model_cls == "wan2.1":
helloyongyang's avatar
helloyongyang committed
202
203
        img = Image.open(args.image_path).convert("RGB")
        img = TF.to_tensor(img).sub_(0.5).div_(0.5).cuda()
gushiqiao's avatar
gushiqiao committed
204
        clip_encoder_out = image_encoder.visual([img[:, None, :, :]], args).squeeze(0).to(torch.bfloat16)
helloyongyang's avatar
helloyongyang committed
205
206
207
        h, w = img.shape[1:]
        aspect_ratio = h / w
        max_area = args.target_height * args.target_width
Dongz's avatar
Dongz committed
208
209
        lat_h = round(np.sqrt(max_area * aspect_ratio) // args.vae_stride[1] // args.patch_size[1] * args.patch_size[1])
        lat_w = round(np.sqrt(max_area / aspect_ratio) // args.vae_stride[2] // args.patch_size[2] * args.patch_size[2])
helloyongyang's avatar
helloyongyang committed
210
211
212
213
214
        h = lat_h * args.vae_stride[1]
        w = lat_w * args.vae_stride[2]

        args.lat_h = lat_h
        args.lat_w = lat_w
Dongz's avatar
Dongz committed
215
216

        msk = torch.ones(1, 81, lat_h, lat_w, device=torch.device("cuda"))
helloyongyang's avatar
helloyongyang committed
217
        msk[:, 1:] = 0
Dongz's avatar
Dongz committed
218
        msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1)
helloyongyang's avatar
helloyongyang committed
219
220
        msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)
        msk = msk.transpose(1, 2)[0]
gushiqiao's avatar
gushiqiao committed
221
222
223
        vae_encode_out = vae_model.encode(
            [torch.concat([torch.nn.functional.interpolate(img[None].cpu(), size=(h, w), mode="bicubic").transpose(0, 1), torch.zeros(3, 80, h, w)], dim=1).cuda()], args
        )[0]
helloyongyang's avatar
helloyongyang committed
224
225
226
227
        vae_encode_out = torch.concat([msk, vae_encode_out]).to(torch.bfloat16)
        return {"clip_encoder_out": clip_encoder_out, "vae_encode_out": vae_encode_out}

    else:
gushiqiao's avatar
gushiqiao committed
228
        raise NotImplementedError(f"Unsupported model class: {args.model_cls}")
helloyongyang's avatar
helloyongyang committed
229
230


helloyongyang's avatar
helloyongyang committed
231
def run_text_encoder(args, text, text_encoders, model_config, image_encoder_output):
helloyongyang's avatar
helloyongyang committed
232
233
234
    text_encoder_output = {}
    if args.model_cls == "hunyuan":
        for i, encoder in enumerate(text_encoders):
helloyongyang's avatar
helloyongyang committed
235
236
237
238
            if args.task == "i2v" and i == 0:
                text_state, attention_mask = encoder.infer(text, image_encoder_output["img"], args)
            else:
                text_state, attention_mask = encoder.infer(text, args)
Dongz's avatar
Dongz committed
239
240
            text_encoder_output[f"text_encoder_{i + 1}_text_states"] = text_state.to(dtype=torch.bfloat16)
            text_encoder_output[f"text_encoder_{i + 1}_attention_mask"] = attention_mask
helloyongyang's avatar
helloyongyang committed
241
242
243
244
245
246
247
248
249
250
251
252
253
254

    elif args.model_cls == "wan2.1":
        n_prompt = model_config.get("sample_neg_prompt", "")
        context = text_encoders[0].infer([text], args)
        context_null = text_encoders[0].infer([n_prompt if n_prompt else ""], args)
        text_encoder_output["context"] = context
        text_encoder_output["context_null"] = context_null

    else:
        raise NotImplementedError(f"Unsupported model type: {args.model_cls}")

    return text_encoder_output


helloyongyang's avatar
helloyongyang committed
255
def init_scheduler(args, image_encoder_output):
helloyongyang's avatar
helloyongyang committed
256
257
    if args.model_cls == "hunyuan":
        if args.feature_caching == "NoCaching":
helloyongyang's avatar
helloyongyang committed
258
            scheduler = HunyuanScheduler(args, image_encoder_output)
helloyongyang's avatar
helloyongyang committed
259
        elif args.feature_caching == "TaylorSeer":
helloyongyang's avatar
helloyongyang committed
260
            scheduler = HunyuanSchedulerFeatureCaching(args, image_encoder_output)
helloyongyang's avatar
helloyongyang committed
261
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
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
        else:
            raise NotImplementedError(f"Unsupported feature_caching type: {args.feature_caching}")

    elif args.model_cls == "wan2.1":
        if args.feature_caching == "NoCaching":
            scheduler = WanScheduler(args)
        elif args.feature_caching == "Tea":
            scheduler = WanSchedulerFeatureCaching(args)
        else:
            raise NotImplementedError(f"Unsupported feature_caching type: {args.feature_caching}")

    else:
        raise NotImplementedError(f"Unsupported model class: {args.model_cls}")
    return scheduler


def run_main_inference(args, model, text_encoder_output, image_encoder_output):
    for step_index in range(model.scheduler.infer_steps):
        torch.cuda.synchronize()
        time1 = time.time()

        model.scheduler.step_pre(step_index=step_index)

        torch.cuda.synchronize()
        time2 = time.time()

        model.infer(text_encoder_output, image_encoder_output, args)

        torch.cuda.synchronize()
        time3 = time.time()

        model.scheduler.step_post()

        torch.cuda.synchronize()
        time4 = time.time()

        print(f"step {step_index} infer time: {time3 - time2}")
        print(f"step {step_index} all time: {time4 - time1}")
        print("*" * 10)

    return model.scheduler.latents, model.scheduler.generator


def run_vae(latents, generator, args):
    images = vae_model.decode(latents, generator=generator, args=args)
    return images


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_cls", type=str, required=True, choices=["wan2.1", "hunyuan"], default="hunyuan")
    parser.add_argument("--task", type=str, choices=["t2v", "i2v"], default="t2v")
    parser.add_argument("--model_path", type=str, required=True)
    parser.add_argument("--config_path", type=str, default=None)
    parser.add_argument("--image_path", type=str, default=None)
Dongz's avatar
Dongz committed
316
    parser.add_argument("--save_video_path", type=str, default="./output_ligthx2v.mp4")
helloyongyang's avatar
helloyongyang committed
317
318
319
320
321
322
323
324
325
    parser.add_argument("--prompt", type=str, required=True)
    parser.add_argument("--infer_steps", type=int, required=True)
    parser.add_argument("--target_video_length", type=int, required=True)
    parser.add_argument("--target_width", type=int, required=True)
    parser.add_argument("--target_height", type=int, required=True)
    parser.add_argument("--attention_type", type=str, required=True)
    parser.add_argument("--sample_neg_prompt", type=str, default="")
    parser.add_argument("--sample_guide_scale", type=float, default=5.0)
    parser.add_argument("--sample_shift", type=float, default=5.0)
Dongz's avatar
Dongz committed
326
327
328
329
330
    parser.add_argument("--do_mm_calib", action="store_true")
    parser.add_argument("--cpu_offload", action="store_true")
    parser.add_argument("--feature_caching", choices=["NoCaching", "TaylorSeer", "Tea"], default="NoCaching")
    parser.add_argument("--mm_config", default=None)
    parser.add_argument("--seed", type=int, default=42)
Xinchi Huang's avatar
Xinchi Huang committed
331
    parser.add_argument("--parallel_attn_type", default=None, choices=["ulysses", "ring"])
Dongz's avatar
Dongz committed
332
333
334
335
    parser.add_argument("--parallel_vae", action="store_true")
    parser.add_argument("--max_area", action="store_true")
    parser.add_argument("--vae_stride", default=(4, 8, 8))
    parser.add_argument("--patch_size", default=(1, 2, 2))
helloyongyang's avatar
helloyongyang committed
336
337
    parser.add_argument("--teacache_thresh", type=float, default=0.26)
    parser.add_argument("--use_ret_steps", action="store_true", default=False)
lijiaqi2's avatar
lijiaqi2 committed
338
339
340
341
    parser.add_argument("--use_bfloat16", action="store_true", default=True)
    parser.add_argument("--lora_path", type=str, default=None)
    parser.add_argument("--strength_model", type=float, default=1.0)

helloyongyang's avatar
helloyongyang committed
342
343
344
345
    args = parser.parse_args()

    start_time = time.time()
    print(f"args: {args}")
Dongz's avatar
Dongz committed
346

helloyongyang's avatar
helloyongyang committed
347
348
    seed_all(args.seed)

Xinchi Huang's avatar
Xinchi Huang committed
349
    if args.parallel_attn_type:
Dongz's avatar
Dongz committed
350
        dist.init_process_group(backend="nccl")
helloyongyang's avatar
helloyongyang committed
351
352
353
354
355
356
357

    if args.mm_config:
        mm_config = json.loads(args.mm_config)
    else:
        mm_config = None

    model_config = {
358
        "model_cls": args.model_cls,
helloyongyang's avatar
helloyongyang committed
359
360
361
362
363
364
365
        "task": args.task,
        "attention_type": args.attention_type,
        "sample_neg_prompt": args.sample_neg_prompt,
        "mm_config": mm_config,
        "do_mm_calib": args.do_mm_calib,
        "cpu_offload": args.cpu_offload,
        "feature_caching": args.feature_caching,
Xinchi Huang's avatar
Xinchi Huang committed
366
        "parallel_attn_type": args.parallel_attn_type,
Dongz's avatar
Dongz committed
367
        "parallel_vae": args.parallel_vae,
lijiaqi2's avatar
lijiaqi2 committed
368
        "use_bfloat16": args.use_bfloat16,
helloyongyang's avatar
helloyongyang committed
369
370
371
372
373
374
375
376
377
    }

    if args.config_path is not None:
        with open(args.config_path, "r") as f:
            config = json.load(f)
        model_config.update(config)

    print(f"model_config: {model_config}")

lijiaqi2's avatar
lijiaqi2 committed
378
379
    with time_duration("Load models"):
        model, text_encoders, vae_model, image_encoder = load_models(args, model_config)
380

Dongz's avatar
Dongz committed
381
    if args.task in ["i2v"]:
helloyongyang's avatar
helloyongyang committed
382
383
384
385
        image_encoder_output = run_image_encoder(args, image_encoder, vae_model)
    else:
        image_encoder_output = {"clip_encoder_out": None, "vae_encode_out": None}

helloyongyang's avatar
helloyongyang committed
386
    text_encoder_output = run_text_encoder(args, args.prompt, text_encoders, model_config, image_encoder_output)
helloyongyang's avatar
helloyongyang committed
387

helloyongyang's avatar
helloyongyang committed
388
389
    set_target_shape(args, image_encoder_output)
    scheduler = init_scheduler(args, image_encoder_output)
helloyongyang's avatar
helloyongyang committed
390
391
392
393
394
395
396

    model.set_scheduler(scheduler)

    gc.collect()
    torch.cuda.empty_cache()
    latents, generator = run_main_inference(args, model, text_encoder_output, image_encoder_output)

397
398
399
400
    if args.cpu_offload:
        scheduler.clear()
        del text_encoder_output, image_encoder_output, model, text_encoders, scheduler
        torch.cuda.empty_cache()
helloyongyang's avatar
helloyongyang committed
401
402
403

    images = run_vae(latents, generator, args)

Xinchi Huang's avatar
Xinchi Huang committed
404
    if not args.parallel_attn_type or (args.parallel_attn_type and dist.get_rank() == 0):
405
        save_video_st = time.time()
Xinchi Huang's avatar
Xinchi Huang committed
406
407
408
409
        if args.model_cls == "wan2.1":
            cache_video(tensor=images, save_file=args.save_video_path, fps=16, nrow=1, normalize=True, value_range=(-1, 1))
        else:
            save_videos_grid(images, args.save_video_path, fps=24)
410
411
        save_video_et = time.time()
        print(f"Save video cost: {save_video_et - save_video_st}")
helloyongyang's avatar
helloyongyang committed
412
413

    end_time = time.time()
414
    print(f"Total cost: {end_time - start_time}")