gradio_demo_zh.py 51.1 KB
Newer Older
gushiqiao's avatar
gushiqiao committed
1
2
import argparse
import gc
PengGao's avatar
PengGao committed
3
4
5
6
7
import glob
import importlib.util
import json
import os
import random
gushiqiao's avatar
gushiqiao committed
8
9
from datetime import datetime

PengGao's avatar
PengGao committed
10
import gradio as gr
gushiqiao's avatar
gushiqiao committed
11
import psutil
PengGao's avatar
PengGao committed
12
13
14
import torch
from easydict import EasyDict
from loguru import logger
gushiqiao's avatar
gushiqiao committed
15
16
17
18
19
20
21
22
23
24

logger.add(
    "inference_logs.log",
    rotation="100 MB",
    encoding="utf-8",
    enqueue=True,
    backtrace=True,
    diagnose=True,
)

gushiqiao's avatar
gushiqiao committed
25
26
27
MAX_NUMPY_SEED = 2**32 - 1


gushiqiao's avatar
gushiqiao committed
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def find_hf_model_path(model_path, subdir=["original", "fp8", "int8"]):
    paths_to_check = [model_path]
    if isinstance(subdir, list):
        for sub in subdir:
            paths_to_check.append(os.path.join(model_path, sub))
    else:
        paths_to_check.append(os.path.join(model_path, subdir))

    for path in paths_to_check:
        safetensors_pattern = os.path.join(path, "*.safetensors")
        safetensors_files = glob.glob(safetensors_pattern)
        if safetensors_files:
            logger.info(f"Found Hugging Face model files in: {path}")
            return path
    raise FileNotFoundError(f"No Hugging Face model files (.safetensors) found.\nPlease download the model from: https://huggingface.co/lightx2v/ or specify the model path in the configuration file.")


def find_torch_model_path(model_path, filename=None, subdir=["original", "fp8", "int8"]):
    paths_to_check = [
        os.path.join(model_path, filename),
    ]
    if isinstance(subdir, list):
        for sub in subdir:
            paths_to_check.append(os.path.join(model_path, sub, filename))
    else:
        paths_to_check.append(os.path.join(model_path, subdir, filename))
    print(paths_to_check)
    for path in paths_to_check:
        if os.path.exists(path):
            logger.info(f"Found PyTorch model checkpoint: {path}")
            return path
    raise FileNotFoundError(f"PyTorch model file '{filename}' not found.\nPlease download the model from https://huggingface.co/lightx2v/ or specify the model path in the configuration file.")


gushiqiao's avatar
gushiqiao committed
62
63
64
def generate_random_seed():
    return random.randint(0, MAX_NUMPY_SEED)

gushiqiao's avatar
gushiqiao committed
65

gushiqiao's avatar
gushiqiao committed
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def is_module_installed(module_name):
    try:
        spec = importlib.util.find_spec(module_name)
        return spec is not None
    except ModuleNotFoundError:
        return False


def get_available_quant_ops():
    available_ops = []

    vllm_installed = is_module_installed("vllm")
    if vllm_installed:
        available_ops.append(("vllm", True))
    else:
        available_ops.append(("vllm", False))

    sgl_installed = is_module_installed("sgl_kernel")
    if sgl_installed:
        available_ops.append(("sgl", True))
    else:
        available_ops.append(("sgl", False))

    q8f_installed = is_module_installed("q8_kernels")
    if q8f_installed:
        available_ops.append(("q8f", True))
    else:
        available_ops.append(("q8f", False))

    return available_ops


def get_available_attn_ops():
    available_ops = []

    vllm_installed = is_module_installed("flash_attn")
    if vllm_installed:
        available_ops.append(("flash_attn2", True))
    else:
        available_ops.append(("flash_attn2", False))

    sgl_installed = is_module_installed("flash_attn_interface")
    if sgl_installed:
        available_ops.append(("flash_attn3", True))
    else:
        available_ops.append(("flash_attn3", False))

    q8f_installed = is_module_installed("sageattention")
    if q8f_installed:
        available_ops.append(("sage_attn2", True))
    else:
        available_ops.append(("sage_attn2", False))

gushiqiao's avatar
gushiqiao committed
119
120
121
122
123
124
    torch_installed = is_module_installed("torch")
    if torch_installed:
        available_ops.append(("torch_sdpa", True))
    else:
        available_ops.append(("torch_sdpa", False))

gushiqiao's avatar
gushiqiao committed
125
126
127
128
129
130
131
132
133
    return available_ops


def get_gpu_memory(gpu_idx=0):
    if not torch.cuda.is_available():
        return 0
    try:
        with torch.cuda.device(gpu_idx):
            memory_info = torch.cuda.mem_get_info()
gushiqiao's avatar
gushiqiao committed
134
            total_memory = memory_info[1] / (1024**3)  # Convert bytes to GB
gushiqiao's avatar
gushiqiao committed
135
136
137
138
139
140
141
142
143
            return total_memory
    except Exception as e:
        logger.warning(f"获取GPU内存失败: {e}")
        return 0


def get_cpu_memory():
    available_bytes = psutil.virtual_memory().available
    return available_bytes / 1024**3
gushiqiao's avatar
gushiqiao committed
144
145


gushiqiao's avatar
gushiqiao committed
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
def cleanup_memory():
    gc.collect()

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

    try:
        import psutil

        if hasattr(psutil, "virtual_memory"):
            if os.name == "posix":
                try:
                    os.system("sync")
                except:  # noqa
                    pass
    except:  # noqa
        pass


gushiqiao's avatar
gushiqiao committed
166
167
def generate_unique_filename(output_dir):
    os.makedirs(output_dir, exist_ok=True)
gushiqiao's avatar
gushiqiao committed
168
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
gushiqiao's avatar
gushiqiao committed
169
    return os.path.join(output_dir, f"{model_cls}_{timestamp}.mp4")
gushiqiao's avatar
gushiqiao committed
170
171


gushiqiao's avatar
gushiqiao committed
172
173
174
175
176
177
178
179
def is_fp8_supported_gpu():
    if not torch.cuda.is_available():
        return False
    compute_capability = torch.cuda.get_device_capability(0)
    major, minor = compute_capability
    return (major == 8 and minor == 9) or (major >= 9)


gushiqiao's avatar
gushiqiao committed
180
181
182
183
184
185
186
187
188
189
190
191
def is_ada_architecture_gpu():
    if not torch.cuda.is_available():
        return False
    try:
        gpu_name = torch.cuda.get_device_name(0).upper()
        ada_keywords = ["RTX 40", "RTX40", "4090", "4080", "4070", "4060"]
        return any(keyword in gpu_name for keyword in ada_keywords)
    except Exception as e:
        logger.warning(f"Failed to get GPU name: {e}")
        return False


gushiqiao's avatar
gushiqiao committed
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
def get_quantization_options(model_path):
    """根据model_path动态获取量化选项"""
    import os

    # 检查子目录
    subdirs = ["original", "fp8", "int8"]
    has_subdirs = {subdir: os.path.exists(os.path.join(model_path, subdir)) for subdir in subdirs}

    # 检查根目录下的原始文件
    t5_bf16_exists = os.path.exists(os.path.join(model_path, "models_t5_umt5-xxl-enc-bf16.pth"))
    clip_fp16_exists = os.path.exists(os.path.join(model_path, "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"))

    # 生成选项
    def get_choices(has_subdirs, original_type, fp8_type, int8_type, fallback_type, has_original_file=False):
        choices = []
        if has_subdirs["original"]:
            choices.append(original_type)
        if has_subdirs["fp8"]:
            choices.append(fp8_type)
        if has_subdirs["int8"]:
            choices.append(int8_type)

        # 如果没有子目录但有原始文件,添加原始类型
gushiqiao's avatar
gushiqiao committed
215
216
217
        if has_original_file:
            if not choices or "original" not in choices:
                choices.append(original_type)
gushiqiao's avatar
gushiqiao committed
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236

        # 如果没有任何选项,使用默认值
        if not choices:
            choices = [fallback_type]

        return choices, choices[0]

    # DIT选项
    dit_choices, dit_default = get_choices(has_subdirs, "bf16", "fp8", "int8", "bf16")

    # T5选项 - 检查是否有原始文件
    t5_choices, t5_default = get_choices(has_subdirs, "bf16", "fp8", "int8", "bf16", t5_bf16_exists)

    # CLIP选项 - 检查是否有原始文件
    clip_choices, clip_default = get_choices(has_subdirs, "fp16", "fp8", "int8", "fp16", clip_fp16_exists)

    return {"dit_choices": dit_choices, "dit_default": dit_default, "t5_choices": t5_choices, "t5_default": t5_default, "clip_choices": clip_choices, "clip_default": clip_default}


gushiqiao's avatar
gushiqiao committed
237
238
global_runner = None
current_config = None
gushiqiao's avatar
gushiqiao committed
239
240
241
242
243
cur_dit_quant_scheme = None
cur_clip_quant_scheme = None
cur_t5_quant_scheme = None
cur_precision_mode = None
cur_enable_teacache = None
gushiqiao's avatar
gushiqiao committed
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259

available_quant_ops = get_available_quant_ops()
quant_op_choices = []
for op_name, is_installed in available_quant_ops:
    status_text = "✅ 已安装" if is_installed else "❌ 未安装"
    display_text = f"{op_name} ({status_text})"
    quant_op_choices.append((op_name, display_text))

available_attn_ops = get_available_attn_ops()
attn_op_choices = []
for op_name, is_installed in available_attn_ops:
    status_text = "✅ 已安装" if is_installed else "❌ 未安装"
    display_text = f"{op_name} ({status_text})"
    attn_op_choices.append((op_name, display_text))


gushiqiao's avatar
gushiqiao committed
260
261
262
263
264
265
266
267
268
269
270
271
def run_inference(
    prompt,
    negative_prompt,
    save_video_path,
    torch_compile,
    infer_steps,
    num_frames,
    resolution,
    seed,
    sample_shift,
    enable_teacache,
    teacache_thresh,
gushiqiao's avatar
gushiqiao committed
272
    use_ret_steps,
gushiqiao's avatar
gushiqiao committed
273
274
275
276
277
278
279
280
281
282
283
284
    enable_cfg,
    cfg_scale,
    dit_quant_scheme,
    t5_quant_scheme,
    clip_quant_scheme,
    fps,
    use_tiny_vae,
    use_tiling_vae,
    lazy_load,
    precision_mode,
    cpu_offload,
    offload_granularity,
gushiqiao's avatar
gushiqiao committed
285
    offload_ratio,
gushiqiao's avatar
gushiqiao committed
286
287
    t5_cpu_offload,
    unload_modules,
gushiqiao's avatar
gushiqiao committed
288
289
290
291
    t5_offload_granularity,
    attention_type,
    quant_op,
    rotary_chunk,
gushiqiao's avatar
gushiqiao committed
292
    rotary_chunk_size,
gushiqiao's avatar
gushiqiao committed
293
    clean_cuda_cache,
gushiqiao's avatar
gushiqiao committed
294
    image_path=None,
gushiqiao's avatar
gushiqiao committed
295
):
gushiqiao's avatar
gushiqiao committed
296
297
    cleanup_memory()

gushiqiao's avatar
gushiqiao committed
298
299
300
    quant_op = quant_op.split("(")[0].strip()
    attention_type = attention_type.split("(")[0].strip()

gushiqiao's avatar
gushiqiao committed
301
    global global_runner, current_config, model_path, task
gushiqiao's avatar
gushiqiao committed
302
    global cur_dit_quant_scheme, cur_clip_quant_scheme, cur_t5_quant_scheme, cur_precision_mode, cur_enable_teacache
gushiqiao's avatar
gushiqiao committed
303
304
305
306

    if os.path.exists(os.path.join(model_path, "config.json")):
        with open(os.path.join(model_path, "config.json"), "r") as f:
            model_config = json.load(f)
gushiqiao's avatar
gushiqiao committed
307
308
    else:
        model_config = {}
gushiqiao's avatar
gushiqiao committed
309
310

    if task == "t2v":
gushiqiao's avatar
gushiqiao committed
311
        if model_size == "1.3b":
gushiqiao's avatar
gushiqiao committed
312
313
314
315
316
317
318
319
320
321
322
323
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
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
376
377
378
379
380
381
382
383
384
385
386
            # 1.3B
            coefficient = [
                [
                    -5.21862437e04,
                    9.23041404e03,
                    -5.28275948e02,
                    1.36987616e01,
                    -4.99875664e-02,
                ],
                [
                    2.39676752e03,
                    -1.31110545e03,
                    2.01331979e02,
                    -8.29855975e00,
                    1.37887774e-01,
                ],
            ]
        else:
            # 14B
            coefficient = [
                [
                    -3.03318725e05,
                    4.90537029e04,
                    -2.65530556e03,
                    5.87365115e01,
                    -3.15583525e-01,
                ],
                [
                    -5784.54975374,
                    5449.50911966,
                    -1811.16591783,
                    256.27178429,
                    -13.02252404,
                ],
            ]
    elif task == "i2v":
        if resolution in [
            "1280x720",
            "720x1280",
            "1280x544",
            "544x1280",
            "1104x832",
            "832x1104",
            "960x960",
        ]:
            # 720p
            coefficient = [
                [
                    8.10705460e03,
                    2.13393892e03,
                    -3.72934672e02,
                    1.66203073e01,
                    -4.17769401e-02,
                ],
                [-114.36346466, 65.26524496, -18.82220707, 4.91518089, -0.23412683],
            ]
        else:
            # 480p
            coefficient = [
                [
                    2.57151496e05,
                    -3.54229917e04,
                    1.40286849e03,
                    -1.35890334e01,
                    1.32517977e-01,
                ],
                [
                    -3.02331670e02,
                    2.23948934e02,
                    -5.25463970e01,
                    5.87348440e00,
                    -2.01973289e-01,
                ],
            ]

gushiqiao's avatar
gushiqiao committed
387
    save_video_path = generate_unique_filename(output_dir)
gushiqiao's avatar
gushiqiao committed
388
389
390

    is_dit_quant = dit_quant_scheme != "bf16"
    is_t5_quant = t5_quant_scheme != "bf16"
gushiqiao's avatar
gushiqiao committed
391

gushiqiao's avatar
gushiqiao committed
392
    if is_t5_quant:
gushiqiao's avatar
gushiqiao committed
393
394
395
        t5_model_name = f"models_t5_umt5-xxl-enc-{t5_quant_scheme}.pth"
        t5_quantized_ckpt = find_torch_model_path(model_path, t5_model_name, t5_quant_scheme)
        t5_original_ckpt = None
gushiqiao's avatar
gushiqiao committed
396
    else:
gushiqiao's avatar
gushiqiao committed
397
398
399
        t5_quantized_ckpt = None
        t5_model_name = "models_t5_umt5-xxl-enc-bf16.pth"
        t5_original_ckpt = find_torch_model_path(model_path, t5_model_name, "original")
gushiqiao's avatar
gushiqiao committed
400

gushiqiao's avatar
gushiqiao committed
401
    is_clip_quant = clip_quant_scheme != "fp16"
gushiqiao's avatar
gushiqiao committed
402

gushiqiao's avatar
gushiqiao committed
403
    if is_clip_quant:
gushiqiao's avatar
gushiqiao committed
404
405
406
        clip_model_name = f"clip-{t5_quant_scheme}.pth"
        clip_quantized_ckpt = find_torch_model_path(model_path, clip_model_name, clip_quant_scheme)
        clip_original_ckpt = None
gushiqiao's avatar
gushiqiao committed
407
    else:
gushiqiao's avatar
gushiqiao committed
408
409
410
        clip_quantized_ckpt = None
        clip_model_name = "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"
        clip_original_ckpt = find_torch_model_path(model_path, clip_model_name, "original")
gushiqiao's avatar
gushiqiao committed
411

gushiqiao's avatar
gushiqiao committed
412
413
    needs_reinit = (
        lazy_load
gushiqiao's avatar
gushiqiao committed
414
        or unload_modules
gushiqiao's avatar
gushiqiao committed
415
416
417
418
419
420
421
422
423
424
425
426
427
        or global_runner is None
        or current_config is None
        or cur_dit_quant_scheme is None
        or cur_dit_quant_scheme != dit_quant_scheme
        or cur_clip_quant_scheme is None
        or cur_clip_quant_scheme != clip_quant_scheme
        or cur_t5_quant_scheme is None
        or cur_t5_quant_scheme != t5_quant_scheme
        or cur_precision_mode is None
        or cur_precision_mode != precision_mode
        or cur_enable_teacache is None
        or cur_enable_teacache != enable_teacache
    )
gushiqiao's avatar
gushiqiao committed
428
429
430
431
432
433
434
435
436
437
438
439
440
441

    if torch_compile:
        os.environ["ENABLE_GRAPH_MODE"] = "true"
    else:
        os.environ["ENABLE_GRAPH_MODE"] = "false"
    if precision_mode == "bf16":
        os.environ["DTYPE"] = "BF16"
    else:
        os.environ.pop("DTYPE", None)

    if is_dit_quant:
        if quant_op == "vllm":
            mm_type = f"W-{dit_quant_scheme}-channel-sym-A-{dit_quant_scheme}-channel-sym-dynamic-Vllm"
        elif quant_op == "sgl":
gushiqiao's avatar
gushiqiao committed
442
443
444
445
            if dit_quant_scheme == "int8":
                mm_type = f"W-{dit_quant_scheme}-channel-sym-A-{dit_quant_scheme}-channel-sym-dynamic-Sgl-ActVllm"
            else:
                mm_type = f"W-{dit_quant_scheme}-channel-sym-A-{dit_quant_scheme}-channel-sym-dynamic-Sgl"
gushiqiao's avatar
gushiqiao committed
446
447
        elif quant_op == "q8f":
            mm_type = f"W-{dit_quant_scheme}-channel-sym-A-{dit_quant_scheme}-channel-sym-dynamic-Q8F"
gushiqiao's avatar
Fix  
gushiqiao committed
448
449
            t5_quant_scheme = f"{t5_quant_scheme}-q8f"
            clip_quant_scheme = f"{clip_quant_scheme}-q8f"
gushiqiao's avatar
gushiqiao committed
450

gushiqiao's avatar
gushiqiao committed
451
        dit_quantized_ckpt = find_hf_model_path(model_path, dit_quant_scheme)
gushiqiao's avatar
gushiqiao committed
452
453
454
        if os.path.exists(os.path.join(dit_quantized_ckpt, "config.json")):
            with open(os.path.join(dit_quantized_ckpt, "config.json"), "r") as f:
                quant_model_config = json.load(f)
gushiqiao's avatar
gushiqiao committed
455
456
        else:
            quant_model_config = {}
gushiqiao's avatar
gushiqiao committed
457
458
    else:
        mm_type = "Default"
gushiqiao's avatar
gushiqiao committed
459
        dit_quantized_ckpt = None
gushiqiao's avatar
gushiqiao committed
460
        quant_model_config = {}
gushiqiao's avatar
gushiqiao committed
461
462
463
464
465
466

    config = {
        "infer_steps": infer_steps,
        "target_video_length": num_frames,
        "target_width": int(resolution.split("x")[0]),
        "target_height": int(resolution.split("x")[1]),
gushiqiao's avatar
gushiqiao committed
467
468
469
        "self_attn_1_type": attention_type,
        "cross_attn_1_type": attention_type,
        "cross_attn_2_type": attention_type,
gushiqiao's avatar
gushiqiao committed
470
471
472
473
474
475
        "seed": seed,
        "enable_cfg": enable_cfg,
        "sample_guide_scale": cfg_scale,
        "sample_shift": sample_shift,
        "cpu_offload": cpu_offload,
        "offload_granularity": offload_granularity,
gushiqiao's avatar
gushiqiao committed
476
        "offload_ratio": offload_ratio,
gushiqiao's avatar
gushiqiao committed
477
        "t5_offload_granularity": t5_offload_granularity,
gushiqiao's avatar
gushiqiao committed
478
        "dit_quantized_ckpt": dit_quantized_ckpt,
gushiqiao's avatar
gushiqiao committed
479
480
481
482
483
        "mm_config": {
            "mm_type": mm_type,
        },
        "fps": fps,
        "feature_caching": "Tea" if enable_teacache else "NoCaching",
gushiqiao's avatar
gushiqiao committed
484
485
        "coefficients": coefficient[0] if use_ret_steps else coefficient[1],
        "use_ret_steps": use_ret_steps,
gushiqiao's avatar
gushiqiao committed
486
        "teacache_thresh": teacache_thresh,
gushiqiao's avatar
gushiqiao committed
487
        "t5_original_ckpt": t5_original_ckpt,
gushiqiao's avatar
gushiqiao committed
488
489
        "t5_cpu_offload": t5_cpu_offload,
        "unload_modules": unload_modules,
gushiqiao's avatar
gushiqiao committed
490
        "t5_quantized": is_t5_quant,
gushiqiao's avatar
gushiqiao committed
491
        "t5_quantized_ckpt": t5_quantized_ckpt,
gushiqiao's avatar
gushiqiao committed
492
        "t5_quant_scheme": t5_quant_scheme,
gushiqiao's avatar
gushiqiao committed
493
        "clip_original_ckpt": clip_original_ckpt,
gushiqiao's avatar
gushiqiao committed
494
        "clip_quantized": is_clip_quant,
gushiqiao's avatar
gushiqiao committed
495
        "clip_quantized_ckpt": clip_quantized_ckpt,
gushiqiao's avatar
gushiqiao committed
496
        "clip_quant_scheme": clip_quant_scheme,
gushiqiao's avatar
gushiqiao committed
497
        "vae_path": find_torch_model_path(model_path, "Wan2.1_VAE.pth"),
gushiqiao's avatar
gushiqiao committed
498
        "use_tiling_vae": use_tiling_vae,
helloyongyang's avatar
helloyongyang committed
499
        "use_tiny_vae": use_tiny_vae,
gushiqiao's avatar
gushiqiao committed
500
        "tiny_vae_path": (find_torch_model_path(model_path, "taew2_1.pth") if use_tiny_vae else None),
gushiqiao's avatar
gushiqiao committed
501
502
503
504
505
506
507
508
509
510
511
512
        "lazy_load": lazy_load,
        "do_mm_calib": False,
        "parallel_attn_type": None,
        "parallel_vae": False,
        "max_area": False,
        "vae_stride": (4, 8, 8),
        "patch_size": (1, 2, 2),
        "lora_path": None,
        "strength_model": 1.0,
        "use_prompt_enhancer": False,
        "text_len": 512,
        "rotary_chunk": rotary_chunk,
gushiqiao's avatar
gushiqiao committed
513
        "rotary_chunk_size": rotary_chunk_size,
gushiqiao's avatar
gushiqiao committed
514
        "clean_cuda_cache": clean_cuda_cache,
gushiqiao's avatar
gushiqiao committed
515
        "denoising_step_list": [1000, 750, 500, 250],
gushiqiao's avatar
gushiqiao committed
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
    }

    args = argparse.Namespace(
        model_cls=model_cls,
        task=task,
        model_path=model_path,
        prompt_enhancer=None,
        prompt=prompt,
        negative_prompt=negative_prompt,
        image_path=image_path,
        save_video_path=save_video_path,
    )

    config.update({k: v for k, v in vars(args).items()})
    config = EasyDict(config)
    config.update(model_config)
gushiqiao's avatar
gushiqiao committed
532
    config.update(quant_model_config)
gushiqiao's avatar
gushiqiao committed
533
534
535
536

    logger.info(f"使用模型: {model_path}")
    logger.info(f"推理配置:\n{json.dumps(config, indent=4, ensure_ascii=False)}")

gushiqiao's avatar
gushiqiao committed
537
    # Initialize or reuse the runner
gushiqiao's avatar
gushiqiao committed
538
539
540
541
542
543
544
    runner = global_runner
    if needs_reinit:
        if runner is not None:
            del runner
            torch.cuda.empty_cache()
            gc.collect()

gushiqiao's avatar
gushiqiao committed
545
546
        from lightx2v.infer import init_runner  # noqa

gushiqiao's avatar
gushiqiao committed
547
548
        runner = init_runner(config)
        current_config = config
gushiqiao's avatar
gushiqiao committed
549
550
551
552
553
        cur_dit_quant_scheme = dit_quant_scheme
        cur_clip_quant_scheme = clip_quant_scheme
        cur_t5_quant_scheme = t5_quant_scheme
        cur_precision_mode = precision_mode
        cur_enable_teacache = enable_teacache
gushiqiao's avatar
gushiqiao committed
554
555
556

        if not lazy_load:
            global_runner = runner
gushiqiao's avatar
gushiqiao committed
557
558
    else:
        runner.config = config
gushiqiao's avatar
gushiqiao committed
559

560
    runner.run_pipeline()
gushiqiao's avatar
gushiqiao committed
561

gushiqiao's avatar
gushiqiao committed
562
563
564
565
566
567
568
569
570
    del config, args, model_config, quant_model_config
    if "dit_quantized_ckpt" in locals():
        del dit_quantized_ckpt
    if "t5_quant_ckpt" in locals():
        del t5_quant_ckpt
    if "clip_quant_ckpt" in locals():
        del clip_quant_ckpt

    cleanup_memory()
gushiqiao's avatar
gushiqiao committed
571
572
573
574

    return save_video_path


gushiqiao's avatar
gushiqiao committed
575
576
577
578
579
580
def handle_lazy_load_change(lazy_load_enabled):
    """Handle lazy_load checkbox change to automatically enable unload_modules"""
    return gr.update(value=lazy_load_enabled)


def auto_configure(enable_auto_config, resolution):
gushiqiao's avatar
gushiqiao committed
581
582
583
584
585
586
587
588
589
    default_config = {
        "torch_compile_val": False,
        "lazy_load_val": False,
        "rotary_chunk_val": False,
        "rotary_chunk_size_val": 100,
        "clean_cuda_cache_val": False,
        "cpu_offload_val": False,
        "offload_granularity_val": "block",
        "offload_ratio_val": 1,
gushiqiao's avatar
gushiqiao committed
590
591
        "t5_cpu_offload_val": False,
        "unload_modules_val": False,
gushiqiao's avatar
gushiqiao committed
592
593
594
595
596
597
598
599
600
601
602
603
604
        "t5_offload_granularity_val": "model",
        "attention_type_val": attn_op_choices[0][1],
        "quant_op_val": quant_op_choices[0][1],
        "dit_quant_scheme_val": "bf16",
        "t5_quant_scheme_val": "bf16",
        "clip_quant_scheme_val": "fp16",
        "precision_mode_val": "fp32",
        "use_tiny_vae_val": False,
        "use_tiling_vae_val": False,
        "enable_teacache_val": False,
        "teacache_thresh_val": 0.26,
        "use_ret_steps_val": False,
    }
gushiqiao's avatar
gushiqiao committed
605

gushiqiao's avatar
gushiqiao committed
606
607
608
609
610
611
612
613
614
615
616
    if not enable_auto_config:
        return tuple(gr.update(value=default_config[key]) for key in default_config)

    gpu_memory = round(get_gpu_memory())
    cpu_memory = round(get_cpu_memory())

    if is_fp8_supported_gpu():
        quant_type = "fp8"
    else:
        quant_type = "int8"

gushiqiao's avatar
gushiqiao committed
617
    attn_priority = ["sage_attn2", "flash_attn3", "flash_attn2", "torch_sdpa"]
gushiqiao's avatar
gushiqiao committed
618
619
620
621
622

    if is_ada_architecture_gpu():
        quant_op_priority = ["q8f", "vllm", "sgl"]
    else:
        quant_op_priority = ["sgl", "vllm", "q8f"]
gushiqiao's avatar
gushiqiao committed
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651

    for op in attn_priority:
        if dict(available_attn_ops).get(op):
            default_config["attention_type_val"] = dict(attn_op_choices)[op]
            break

    for op in quant_op_priority:
        if dict(available_quant_ops).get(op):
            default_config["quant_op_val"] = dict(quant_op_choices)[op]
            break

    if resolution in [
        "1280x720",
        "720x1280",
        "1280x544",
        "544x1280",
        "1104x832",
        "832x1104",
        "960x960",
    ]:
        res = "720p"
    elif resolution in [
        "960x544",
        "544x960",
    ]:
        res = "540p"
    else:
        res = "480p"

gushiqiao's avatar
gushiqiao committed
652
    if model_size == "14b":
gushiqiao's avatar
gushiqiao committed
653
654
655
656
657
658
659
        is_14b = True
    else:
        is_14b = False

    if res == "720p" and is_14b:
        gpu_rules = [
            (80, {}),
gushiqiao's avatar
gushiqiao committed
660
661
662
            (48, {"cpu_offload_val": True, "offload_ratio_val": 0.5, "t5_cpu_offload_val": True}),
            (40, {"cpu_offload_val": True, "offload_ratio_val": 0.8, "t5_cpu_offload_val": True}),
            (32, {"cpu_offload_val": True, "offload_ratio_val": 1, "t5_cpu_offload_val": True}),
gushiqiao's avatar
gushiqiao committed
663
664
665
666
            (
                24,
                {
                    "cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
667
                    "t5_cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
668
669
670
671
672
673
674
675
676
677
                    "offload_ratio_val": 1,
                    "t5_offload_granularity_val": "block",
                    "precision_mode_val": "bf16",
                    "use_tiling_vae_val": True,
                },
            ),
            (
                16,
                {
                    "cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
678
                    "t5_cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
679
680
681
682
683
684
685
686
687
688
689
690
691
                    "offload_ratio_val": 1,
                    "t5_offload_granularity_val": "block",
                    "precision_mode_val": "bf16",
                    "use_tiling_vae_val": True,
                    "offload_granularity_val": "phase",
                    "rotary_chunk_val": True,
                    "rotary_chunk_size_val": 100,
                },
            ),
            (
                12,
                {
                    "cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
692
                    "t5_cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
693
694
695
696
697
698
699
700
                    "offload_ratio_val": 1,
                    "t5_offload_granularity_val": "block",
                    "precision_mode_val": "bf16",
                    "use_tiling_vae_val": True,
                    "offload_granularity_val": "phase",
                    "rotary_chunk_val": True,
                    "rotary_chunk_size_val": 100,
                    "clean_cuda_cache_val": True,
gushiqiao's avatar
gushiqiao committed
701
                    "use_tiny_vae_val": True,
gushiqiao's avatar
gushiqiao committed
702
703
704
705
706
707
                },
            ),
            (
                8,
                {
                    "cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
708
                    "t5_cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
709
710
711
712
713
714
715
716
717
718
719
720
                    "offload_ratio_val": 1,
                    "t5_offload_granularity_val": "block",
                    "precision_mode_val": "bf16",
                    "use_tiling_vae_val": True,
                    "offload_granularity_val": "phase",
                    "rotary_chunk_val": True,
                    "rotary_chunk_size_val": 100,
                    "clean_cuda_cache_val": True,
                    "t5_quant_scheme_val": quant_type,
                    "clip_quant_scheme_val": quant_type,
                    "dit_quant_scheme_val": quant_type,
                    "lazy_load_val": True,
gushiqiao's avatar
gushiqiao committed
721
                    "unload_modules_val": True,
gushiqiao's avatar
gushiqiao committed
722
                    "use_tiny_vae_val": True,
gushiqiao's avatar
gushiqiao committed
723
724
725
                },
            ),
        ]
gushiqiao's avatar
gushiqiao committed
726

gushiqiao's avatar
gushiqiao committed
727
728
729
    elif is_14b:
        gpu_rules = [
            (80, {}),
gushiqiao's avatar
gushiqiao committed
730
731
732
            (48, {"cpu_offload_val": True, "offload_ratio_val": 0.2, "t5_cpu_offload_val": True}),
            (40, {"cpu_offload_val": True, "offload_ratio_val": 0.5, "t5_cpu_offload_val": True}),
            (24, {"cpu_offload_val": True, "offload_ratio_val": 0.8, "t5_cpu_offload_val": True}),
gushiqiao's avatar
gushiqiao committed
733
734
735
736
            (
                16,
                {
                    "cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
737
                    "t5_cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
738
739
740
741
742
743
744
745
746
747
748
749
                    "offload_ratio_val": 1,
                    "t5_offload_granularity_val": "block",
                    "precision_mode_val": "bf16",
                    "use_tiling_vae_val": True,
                    "offload_granularity_val": "block",
                },
            ),
            (
                8,
                (
                    {
                        "cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
750
                        "t5_cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
751
752
753
754
755
756
757
758
759
                        "offload_ratio_val": 1,
                        "t5_offload_granularity_val": "block",
                        "precision_mode_val": "bf16",
                        "use_tiling_vae_val": True,
                        "offload_granularity_val": "phase",
                        "t5_quant_scheme_val": quant_type,
                        "clip_quant_scheme_val": quant_type,
                        "dit_quant_scheme_val": quant_type,
                        "lazy_load_val": True,
gushiqiao's avatar
gushiqiao committed
760
                        "unload_modules_val": True,
gushiqiao's avatar
gushiqiao committed
761
762
                        "rotary_chunk_val": True,
                        "rotary_chunk_size_val": 10000,
gushiqiao's avatar
gushiqiao committed
763
                        "use_tiny_vae_val": True,
gushiqiao's avatar
gushiqiao committed
764
765
766
767
                    }
                    if res == "540p"
                    else {
                        "cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
768
                        "t5_cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
769
770
771
772
773
774
775
776
777
                        "offload_ratio_val": 1,
                        "t5_offload_granularity_val": "block",
                        "precision_mode_val": "bf16",
                        "use_tiling_vae_val": True,
                        "offload_granularity_val": "phase",
                        "t5_quant_scheme_val": quant_type,
                        "clip_quant_scheme_val": quant_type,
                        "dit_quant_scheme_val": quant_type,
                        "lazy_load_val": True,
gushiqiao's avatar
gushiqiao committed
778
                        "unload_modules_val": True,
gushiqiao's avatar
gushiqiao committed
779
                        "use_tiny_vae_val": True,
gushiqiao's avatar
gushiqiao committed
780
781
782
783
784
                    }
                ),
            ),
        ]

gushiqiao's avatar
gushiqiao committed
785
    else:
gushiqiao's avatar
gushiqiao committed
786
787
788
789
790
791
792
793
794
795
796
        gpu_rules = [
            (24, {}),
            (
                8,
                {
                    "t5_cpu_offload_val": True,
                    "t5_offload_granularity_val": "block",
                    "t5_quant_scheme_val": quant_type,
                },
            ),
        ]
gushiqiao's avatar
gushiqiao committed
797

gushiqiao's avatar
gushiqiao committed
798
799
800
801
802
803
804
    if is_14b:
        cpu_rules = [
            (128, {}),
            (64, {"dit_quant_scheme_val": quant_type}),
            (32, {"dit_quant_scheme_val": quant_type, "lazy_load_val": True}),
            (
                16,
gushiqiao's avatar
gushiqiao committed
805
806
807
808
809
                {
                    "dit_quant_scheme_val": quant_type,
                    "t5_quant_scheme_val": quant_type,
                    "clip_quant_scheme_val": quant_type,
                    "lazy_load_val": True,
gushiqiao's avatar
gushiqiao committed
810
                    "unload_modules_val": True,
gushiqiao's avatar
gushiqiao committed
811
                },
gushiqiao's avatar
gushiqiao committed
812
813
            ),
        ]
gushiqiao's avatar
gushiqiao committed
814
    else:
gushiqiao's avatar
gushiqiao committed
815
816
817
818
819
820
821
822
823
824
825
        cpu_rules = [
            (64, {}),
            (
                16,
                {
                    "t5_quant_scheme_val": quant_type,
                    "unload_modules_val": True,
                    "use_tiny_vae_val": True,
                },
            ),
        ]
gushiqiao's avatar
gushiqiao committed
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840

    for threshold, updates in gpu_rules:
        if gpu_memory >= threshold:
            default_config.update(updates)
            break

    for threshold, updates in cpu_rules:
        if cpu_memory >= threshold:
            default_config.update(updates)
            break

    return tuple(gr.update(value=default_config[key]) for key in default_config)


def main():
gushiqiao's avatar
gushiqiao committed
841
    with gr.Blocks(
gushiqiao's avatar
gushiqiao committed
842
        title="Lightx2v (轻量级视频推理和生成引擎)",
gushiqiao's avatar
gushiqiao committed
843
844
845
846
847
848
        css="""
        .main-content { max-width: 1400px; margin: auto; }
        .output-video { max-height: 650px; }
        .warning { color: #ff6b6b; font-weight: bold; }
        .advanced-options { background: #f9f9ff; border-radius: 10px; padding: 15px; }
        .tab-button { font-size: 16px; padding: 10px 20px; }
gushiqiao's avatar
gushiqiao committed
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
        .auto-config-title {
            background: linear-gradient(45deg, #ff6b6b, #4ecdc4);
            background-clip: text;
            -webkit-background-clip: text;
            color: transparent;
            text-align: center;
            margin: 0 !important;
            padding: 8px;
            border: 2px solid #4ecdc4;
            border-radius: 8px;
            background-color: #f0f8ff;
        }
        .auto-config-checkbox {
            border: 2px solid #ff6b6b !important;
            border-radius: 8px !important;
            padding: 10px !important;
            background: linear-gradient(135deg, #fff5f5, #f0fff0) !important;
            box-shadow: 0 2px 8px rgba(255, 107, 107, 0.2) !important;
        }
        .auto-config-checkbox label {
            font-size: 16px !important;
            font-weight: bold !important;
            color: #2c3e50 !important;
        }
gushiqiao's avatar
gushiqiao committed
873
874
875
876
877
878
    """,
    ) as demo:
        gr.Markdown(f"# 🎬 {model_cls} 视频生成器")
        gr.Markdown(f"### 使用模型: {model_path}")

        with gr.Tabs() as tabs:
gushiqiao's avatar
gushiqiao committed
879
            with gr.Tab("基本设置", id=1):
gushiqiao's avatar
gushiqiao committed
880
881
882
883
884
                with gr.Row():
                    with gr.Column(scale=4):
                        with gr.Group():
                            gr.Markdown("## 📥 输入参数")

gushiqiao's avatar
gushiqiao committed
885
886
887
888
889
890
891
892
893
                            if task == "i2v":
                                with gr.Row():
                                    image_path = gr.Image(
                                        label="输入图像",
                                        type="filepath",
                                        height=300,
                                        interactive=True,
                                        visible=True,
                                    )
gushiqiao's avatar
gushiqiao committed
894
895
896
897
898
899
900
901
902
903
904
905
906

                            with gr.Row():
                                with gr.Column():
                                    prompt = gr.Textbox(
                                        label="提示词",
                                        lines=3,
                                        placeholder="描述视频内容...",
                                        max_lines=5,
                                    )
                                with gr.Column():
                                    negative_prompt = gr.Textbox(
                                        label="负向提示词",
                                        lines=3,
gushiqiao's avatar
gushiqiao committed
907
                                        placeholder="不希望出现在视频中的内容...",
gushiqiao's avatar
gushiqiao committed
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
                                        max_lines=5,
                                        value="镜头晃动,色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走",
                                    )
                                with gr.Column():
                                    resolution = gr.Dropdown(
                                        choices=[
                                            # 720p
                                            ("1280x720 (16:9, 720p)", "1280x720"),
                                            ("720x1280 (9:16, 720p)", "720x1280"),
                                            ("1280x544 (21:9, 720p)", "1280x544"),
                                            ("544x1280 (9:21, 720p)", "544x1280"),
                                            ("1104x832 (4:3, 720p)", "1104x832"),
                                            ("832x1104 (3:4, 720p)", "832x1104"),
                                            ("960x960 (1:1, 720p)", "960x960"),
                                            # 480p
                                            ("960x544 (16:9, 540p)", "960x544"),
                                            ("544x960 (9:16, 540p)", "544x960"),
                                            ("832x480 (16:9, 480p)", "832x480"),
                                            ("480x832 (9:16, 480p)", "480x832"),
                                            ("832x624 (4:3, 480p)", "832x624"),
                                            ("624x832 (3:4, 480p)", "624x832"),
                                            ("720x720 (1:1, 480p)", "720x720"),
                                            ("512x512 (1:1, 480p)", "512x512"),
                                        ],
                                        value="832x480",
                                        label="最大分辨率",
                                    )
gushiqiao's avatar
gushiqiao committed
935
936

                                with gr.Column():
gushiqiao's avatar
gushiqiao committed
937
938
939
940
941
942
943
944
                                    with gr.Group():
                                        gr.Markdown("### 🚀 **智能配置推荐**", elem_classes=["auto-config-title"])
                                        enable_auto_config = gr.Checkbox(
                                            label="🎯 **自动配置推理选项**",
                                            value=False,
                                            info="💡 **智能优化GPU设置以匹配当前分辨率。修改分辨率后,请重新勾选此选项,否则可能导致性能下降或运行失败。**",
                                            elem_classes=["auto-config-checkbox"],
                                        )
gushiqiao's avatar
gushiqiao committed
945
                                with gr.Column(scale=9):
gushiqiao's avatar
gushiqiao committed
946
947
                                    seed = gr.Slider(
                                        label="随机种子",
gushiqiao's avatar
gushiqiao committed
948
949
                                        minimum=0,
                                        maximum=MAX_NUMPY_SEED,
gushiqiao's avatar
gushiqiao committed
950
                                        step=1,
gushiqiao's avatar
gushiqiao committed
951
                                        value=generate_random_seed(),
gushiqiao's avatar
gushiqiao committed
952
                                    )
gushiqiao's avatar
gushiqiao committed
953
                                with gr.Column(scale=1):
gushiqiao's avatar
gushiqiao committed
954
                                    randomize_btn = gr.Button("🎲 随机化", variant="secondary")
gushiqiao's avatar
gushiqiao committed
955
956

                                randomize_btn.click(fn=generate_random_seed, inputs=None, outputs=seed)
gushiqiao's avatar
gushiqiao committed
957

gushiqiao's avatar
gushiqiao committed
958
                                with gr.Column():
gushiqiao's avatar
gushiqiao committed
959
                                    # 根据模型类别设置默认推理步数
gushiqiao's avatar
gushiqiao committed
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
                                    if model_cls == "wan2.1_distill":
                                        infer_steps = gr.Slider(
                                            label="推理步数",
                                            minimum=4,
                                            maximum=4,
                                            step=1,
                                            value=4,
                                            interactive=False,
                                            info="推理步数固定为4,以获得最佳性能(对于蒸馏模型)。",
                                        )
                                    elif model_cls == "wan2.1":
                                        if task == "i2v":
                                            infer_steps = gr.Slider(
                                                label="推理步数",
                                                minimum=1,
                                                maximum=100,
                                                step=1,
                                                value=40,
                                                info="视频生成的推理步数。增加步数可能提高质量但降低速度。",
                                            )
                                        elif task == "t2v":
                                            infer_steps = gr.Slider(
                                                label="推理步数",
                                                minimum=1,
                                                maximum=100,
                                                step=1,
                                                value=50,
                                                info="视频生成的推理步数。增加步数可能提高质量但降低速度。",
                                            )
gushiqiao's avatar
gushiqiao committed
989

gushiqiao's avatar
gushiqiao committed
990
991
                            # 根据模型类别设置默认CFG
                            default_enable_cfg = False if model_cls == "wan2.1_distill" else True
gushiqiao's avatar
gushiqiao committed
992
993
                            enable_cfg = gr.Checkbox(
                                label="启用无分类器引导",
gushiqiao's avatar
gushiqiao committed
994
                                value=default_enable_cfg,
gushiqiao's avatar
gushiqiao committed
995
996
997
998
999
1000
1001
1002
                                info="启用无分类器引导以控制提示词强度",
                            )
                            cfg_scale = gr.Slider(
                                label="CFG缩放因子",
                                minimum=1,
                                maximum=10,
                                step=1,
                                value=5,
gushiqiao's avatar
gushiqiao committed
1003
                                info="控制提示词的影响强度。值越高,提示词的影响越大。",
gushiqiao's avatar
gushiqiao committed
1004
1005
1006
1007
1008
1009
1010
                            )
                            sample_shift = gr.Slider(
                                label="分布偏移",
                                value=5,
                                minimum=0,
                                maximum=10,
                                step=1,
gushiqiao's avatar
gushiqiao committed
1011
                                info="控制样本分布偏移的程度。值越大表示偏移越明显。",
gushiqiao's avatar
gushiqiao committed
1012
1013
                            )

gushiqiao's avatar
gushiqiao committed
1014
1015
1016
1017
1018
1019
                            fps = gr.Slider(
                                label="每秒帧数(FPS)",
                                minimum=8,
                                maximum=30,
                                step=1,
                                value=16,
gushiqiao's avatar
gushiqiao committed
1020
                                info="视频的每秒帧数。较高的FPS会产生更流畅的视频。",
gushiqiao's avatar
gushiqiao committed
1021
1022
1023
1024
1025
1026
1027
                            )
                            num_frames = gr.Slider(
                                label="总帧数",
                                minimum=16,
                                maximum=120,
                                step=1,
                                value=81,
gushiqiao's avatar
gushiqiao committed
1028
                                info="视频中的总帧数。更多帧数会产生更长的视频。",
gushiqiao's avatar
gushiqiao committed
1029
                            )
gushiqiao's avatar
gushiqiao committed
1030

gushiqiao's avatar
gushiqiao committed
1031
1032
                        save_video_path = gr.Textbox(
                            label="输出视频路径",
gushiqiao's avatar
gushiqiao committed
1033
                            value=generate_unique_filename(output_dir),
gushiqiao's avatar
gushiqiao committed
1034
1035
                            info="必须包含.mp4扩展名。如果留空或使用默认值,将自动生成唯一文件名。",
                        )
gushiqiao's avatar
gushiqiao committed
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
                    with gr.Column(scale=6):
                        gr.Markdown("## 📤 生成的视频")
                        output_video = gr.Video(
                            label="结果",
                            height=624,
                            width=360,
                            autoplay=True,
                            elem_classes=["output-video"],
                        )

gushiqiao's avatar
gushiqiao committed
1046
                        infer_btn = gr.Button("生成视频", variant="primary", size="lg")
gushiqiao's avatar
gushiqiao committed
1047

gushiqiao's avatar
gushiqiao committed
1048
1049
            with gr.Tab("⚙️ 高级选项", id=2):
                with gr.Group(elem_classes="advanced-options"):
gushiqiao's avatar
gushiqiao committed
1050
                    gr.Markdown("### GPU内存优化")
gushiqiao's avatar
gushiqiao committed
1051
                    with gr.Row():
gushiqiao's avatar
gushiqiao committed
1052
1053
                        rotary_chunk = gr.Checkbox(
                            label="分块旋转位置编码",
gushiqiao's avatar
gushiqiao committed
1054
                            value=False,
gushiqiao's avatar
gushiqiao committed
1055
                            info="启用时,将旋转位置编码分块处理以节省GPU内存。",
gushiqiao's avatar
gushiqiao committed
1056
1057
                        )

gushiqiao's avatar
gushiqiao committed
1058
1059
1060
1061
1062
1063
                        rotary_chunk_size = gr.Slider(
                            label="旋转编码块大小",
                            value=100,
                            minimum=100,
                            maximum=10000,
                            step=100,
gushiqiao's avatar
gushiqiao committed
1064
                            info="控制应用旋转编码的块大小。较大的值可能提高性能但增加内存使用。仅在'rotary_chunk'勾选时有效。",
gushiqiao's avatar
gushiqiao committed
1065
                        )
gushiqiao's avatar
gushiqiao committed
1066
1067
1068
1069
1070
                        unload_modules = gr.Checkbox(
                            label="卸载模块",
                            value=False,
                            info="推理后卸载模块(T5、CLIP、DIT等)以减少GPU/CPU内存使用",
                        )
gushiqiao's avatar
gushiqiao committed
1071
                        clean_cuda_cache = gr.Checkbox(
gushiqiao's avatar
gushiqiao committed
1072
                            label="清理CUDA内存缓存",
gushiqiao's avatar
gushiqiao committed
1073
                            value=False,
gushiqiao's avatar
gushiqiao committed
1074
                            info="启用时,及时释放GPU内存但会减慢推理速度。",
gushiqiao's avatar
gushiqiao committed
1075
1076
                        )

gushiqiao's avatar
gushiqiao committed
1077
                    gr.Markdown("### 异步卸载")
gushiqiao's avatar
gushiqiao committed
1078
1079
1080
1081
                    with gr.Row():
                        cpu_offload = gr.Checkbox(
                            label="CPU卸载",
                            value=False,
gushiqiao's avatar
gushiqiao committed
1082
                            info="将模型计算的一部分从GPU卸载到CPU以减少GPU内存使用",
gushiqiao's avatar
gushiqiao committed
1083
                        )
gushiqiao's avatar
gushiqiao committed
1084
1085
1086
1087

                        lazy_load = gr.Checkbox(
                            label="启用延迟加载",
                            value=False,
gushiqiao's avatar
gushiqiao committed
1088
                            info="在推理过程中延迟加载模型组件。需要CPU加载和DIT量化。",
gushiqiao's avatar
gushiqiao committed
1089
1090
                        )

gushiqiao's avatar
gushiqiao committed
1091
                        offload_granularity = gr.Dropdown(
gushiqiao's avatar
gushiqiao committed
1092
                            label="Dit卸载粒度",
gushiqiao's avatar
gushiqiao committed
1093
                            choices=["block", "phase"],
gushiqiao's avatar
gushiqiao committed
1094
                            value="phase",
gushiqiao's avatar
gushiqiao committed
1095
                            info="设置Dit模型卸载粒度:块或计算阶段",
gushiqiao's avatar
gushiqiao committed
1096
1097
1098
1099
1100
1101
1102
1103
                        )
                        offload_ratio = gr.Slider(
                            label="Dit模型卸载比例",
                            minimum=0.0,
                            maximum=1.0,
                            step=0.1,
                            value=1.0,
                            info="控制将多少Dit模型卸载到CPU",
gushiqiao's avatar
gushiqiao committed
1104
                        )
gushiqiao's avatar
gushiqiao committed
1105
1106
1107
1108
1109
                        t5_cpu_offload = gr.Checkbox(
                            label="T5 CPU卸载",
                            value=False,
                            info="将T5编码器模型卸载到CPU以减少GPU内存使用",
                        )
gushiqiao's avatar
gushiqiao committed
1110
                        t5_offload_granularity = gr.Dropdown(
gushiqiao's avatar
gushiqiao committed
1111
                            label="T5编码器卸载粒度",
gushiqiao's avatar
gushiqiao committed
1112
                            choices=["model", "block"],
gushiqiao's avatar
gushiqiao committed
1113
1114
                            value="model",
                            info="控制将T5编码器模型卸载到CPU时的粒度",
gushiqiao's avatar
gushiqiao committed
1115
1116
1117
1118
                        )

                    gr.Markdown("### 低精度量化")
                    with gr.Row():
gushiqiao's avatar
gushiqiao committed
1119
1120
1121
1122
                        torch_compile = gr.Checkbox(
                            label="Torch编译",
                            value=False,
                            info="使用torch.compile加速推理过程",
gushiqiao's avatar
gushiqiao committed
1123
1124
                        )

gushiqiao's avatar
gushiqiao committed
1125
1126
1127
1128
1129
1130
                        attention_type = gr.Dropdown(
                            label="注意力算子",
                            choices=[op[1] for op in attn_op_choices],
                            value=attn_op_choices[0][1],
                            info="使用适当的注意力算子加速推理",
                        )
gushiqiao's avatar
gushiqiao committed
1131
                        quant_op = gr.Dropdown(
gushiqiao's avatar
gushiqiao committed
1132
1133
1134
1135
1136
                            label="量化矩阵乘法算子",
                            choices=[op[1] for op in quant_op_choices],
                            value=quant_op_choices[0][1],
                            info="选择量化矩阵乘法算子以加速推理",
                            interactive=True,
gushiqiao's avatar
gushiqiao committed
1137
                        )
gushiqiao's avatar
gushiqiao committed
1138
1139
1140
                        # 获取动态量化选项
                        quant_options = get_quantization_options(model_path)

gushiqiao's avatar
gushiqiao committed
1141
1142
                        dit_quant_scheme = gr.Dropdown(
                            label="Dit",
gushiqiao's avatar
gushiqiao committed
1143
1144
                            choices=quant_options["dit_choices"],
                            value=quant_options["dit_default"],
gushiqiao's avatar
gushiqiao committed
1145
                            info="Dit模型的量化精度",
gushiqiao's avatar
gushiqiao committed
1146
1147
                        )
                        t5_quant_scheme = gr.Dropdown(
gushiqiao's avatar
gushiqiao committed
1148
                            label="T5编码器",
gushiqiao's avatar
gushiqiao committed
1149
1150
                            choices=quant_options["t5_choices"],
                            value=quant_options["t5_default"],
gushiqiao's avatar
gushiqiao committed
1151
                            info="T5编码器模型的量化精度",
gushiqiao's avatar
gushiqiao committed
1152
1153
                        )
                        clip_quant_scheme = gr.Dropdown(
gushiqiao's avatar
gushiqiao committed
1154
                            label="Clip编码器",
gushiqiao's avatar
gushiqiao committed
1155
1156
                            choices=quant_options["clip_choices"],
                            value=quant_options["clip_default"],
gushiqiao's avatar
gushiqiao committed
1157
                            info="Clip编码器的量化精度",
gushiqiao's avatar
gushiqiao committed
1158
1159
                        )
                        precision_mode = gr.Dropdown(
gushiqiao's avatar
gushiqiao committed
1160
                            label="敏感层精度模式",
gushiqiao's avatar
gushiqiao committed
1161
                            choices=["fp32", "bf16"],
gushiqiao's avatar
gushiqiao committed
1162
                            value="fp32",
gushiqiao's avatar
gushiqiao committed
1163
                            info="选择用于关键模型组件(如归一化和嵌入层)的数值精度。FP32提供更高精度,而BF16在兼容硬件上提高性能。",
gushiqiao's avatar
gushiqiao committed
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
                        )

                    gr.Markdown("### 变分自编码器(VAE)")
                    with gr.Row():
                        use_tiny_vae = gr.Checkbox(
                            label="使用轻量级VAE",
                            value=False,
                            info="使用轻量级VAE模型加速解码过程",
                        )
                        use_tiling_vae = gr.Checkbox(
gushiqiao's avatar
gushiqiao committed
1174
                            label="VAE分块推理",
gushiqiao's avatar
gushiqiao committed
1175
                            value=False,
gushiqiao's avatar
gushiqiao committed
1176
                            info="使用VAE分块推理以减少GPU内存使用",
gushiqiao's avatar
gushiqiao committed
1177
1178
1179
1180
1181
                        )

                    gr.Markdown("### 特征缓存")
                    with gr.Row():
                        enable_teacache = gr.Checkbox(
gushiqiao's avatar
gushiqiao committed
1182
                            label="Tea Cache",
gushiqiao's avatar
gushiqiao committed
1183
1184
1185
1186
1187
1188
1189
1190
                            value=False,
                            info="在推理过程中缓存特征以减少推理步数",
                        )
                        teacache_thresh = gr.Slider(
                            label="Tea Cache阈值",
                            value=0.26,
                            minimum=0,
                            maximum=1,
gushiqiao's avatar
gushiqiao committed
1191
1192
1193
1194
1195
1196
                            info="较高的加速可能导致质量下降 —— 设置为0.1提供约2.0倍加速,设置为0.2提供约3.0倍加速",
                        )
                        use_ret_steps = gr.Checkbox(
                            label="仅缓存关键步骤",
                            value=False,
                            info="勾选时,仅在调度器返回结果的关键步骤写入缓存;未勾选时,在所有步骤写入缓存以确保最高质量",
gushiqiao's avatar
gushiqiao committed
1197
1198
                        )

gushiqiao's avatar
gushiqiao committed
1199
1200
                enable_auto_config.change(
                    fn=auto_configure,
gushiqiao's avatar
gushiqiao committed
1201
                    inputs=[enable_auto_config, resolution],
gushiqiao's avatar
gushiqiao committed
1202
1203
1204
1205
1206
1207
1208
1209
1210
                    outputs=[
                        torch_compile,
                        lazy_load,
                        rotary_chunk,
                        rotary_chunk_size,
                        clean_cuda_cache,
                        cpu_offload,
                        offload_granularity,
                        offload_ratio,
gushiqiao's avatar
gushiqiao committed
1211
1212
                        t5_cpu_offload,
                        unload_modules,
gushiqiao's avatar
gushiqiao committed
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
                        t5_offload_granularity,
                        attention_type,
                        quant_op,
                        dit_quant_scheme,
                        t5_quant_scheme,
                        clip_quant_scheme,
                        precision_mode,
                        use_tiny_vae,
                        use_tiling_vae,
                        enable_teacache,
                        teacache_thresh,
                        use_ret_steps,
                    ],
                )
gushiqiao's avatar
gushiqiao committed
1227
1228
1229
1230
1231
1232

                lazy_load.change(
                    fn=handle_lazy_load_change,
                    inputs=[lazy_load],
                    outputs=[unload_modules],
                )
gushiqiao's avatar
gushiqiao committed
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
        if task == "i2v":
            infer_btn.click(
                fn=run_inference,
                inputs=[
                    prompt,
                    negative_prompt,
                    save_video_path,
                    torch_compile,
                    infer_steps,
                    num_frames,
                    resolution,
                    seed,
                    sample_shift,
                    enable_teacache,
                    teacache_thresh,
                    use_ret_steps,
                    enable_cfg,
                    cfg_scale,
                    dit_quant_scheme,
                    t5_quant_scheme,
                    clip_quant_scheme,
                    fps,
                    use_tiny_vae,
                    use_tiling_vae,
                    lazy_load,
                    precision_mode,
                    cpu_offload,
                    offload_granularity,
                    offload_ratio,
gushiqiao's avatar
gushiqiao committed
1262
1263
                    t5_cpu_offload,
                    unload_modules,
gushiqiao's avatar
gushiqiao committed
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
                    t5_offload_granularity,
                    attention_type,
                    quant_op,
                    rotary_chunk,
                    rotary_chunk_size,
                    clean_cuda_cache,
                    image_path,
                ],
                outputs=output_video,
            )
        else:
            infer_btn.click(
                fn=run_inference,
                inputs=[
                    prompt,
                    negative_prompt,
                    save_video_path,
                    torch_compile,
                    infer_steps,
                    num_frames,
                    resolution,
                    seed,
                    sample_shift,
                    enable_teacache,
                    teacache_thresh,
                    use_ret_steps,
                    enable_cfg,
                    cfg_scale,
                    dit_quant_scheme,
                    t5_quant_scheme,
                    clip_quant_scheme,
                    fps,
                    use_tiny_vae,
                    use_tiling_vae,
                    lazy_load,
                    precision_mode,
                    cpu_offload,
                    offload_granularity,
                    offload_ratio,
gushiqiao's avatar
gushiqiao committed
1303
1304
                    t5_cpu_offload,
                    unload_modules,
gushiqiao's avatar
gushiqiao committed
1305
1306
1307
1308
1309
1310
1311
1312
1313
                    t5_offload_granularity,
                    attention_type,
                    quant_op,
                    rotary_chunk,
                    rotary_chunk_size,
                    clean_cuda_cache,
                ],
                outputs=output_video,
            )
gushiqiao's avatar
gushiqiao committed
1314

gushiqiao's avatar
gushiqiao committed
1315
    demo.launch(share=True, server_port=args.server_port, server_name=args.server_name, inbrowser=True, allowed_paths=[output_dir])
gushiqiao's avatar
gushiqiao committed
1316
1317
1318


if __name__ == "__main__":
gushiqiao's avatar
gushiqiao committed
1319
1320
1321
1322
1323
    parser = argparse.ArgumentParser(description="轻量级视频生成")
    parser.add_argument("--model_path", type=str, required=True, help="模型文件夹路径")
    parser.add_argument(
        "--model_cls",
        type=str,
gushiqiao's avatar
gushiqiao committed
1324
        choices=["wan2.1", "wan2.1_distill"],
gushiqiao's avatar
gushiqiao committed
1325
        default="wan2.1",
gushiqiao's avatar
gushiqiao committed
1326
        help="要使用的模型类别 (wan2.1: 标准模型, wan2.1_distill: 蒸馏模型,推理更快)",
gushiqiao's avatar
gushiqiao committed
1327
    )
gushiqiao's avatar
gushiqiao committed
1328
    parser.add_argument("--model_size", type=str, required=True, choices=["14b", "1.3b"], help="模型大小:14b 或 1.3b")
gushiqiao's avatar
gushiqiao committed
1329
    parser.add_argument("--task", type=str, required=True, choices=["i2v", "t2v"], help="指定任务类型。'i2v'用于图像到视频转换,'t2v'用于文本到视频生成。")
gushiqiao's avatar
gushiqiao committed
1330
1331
    parser.add_argument("--server_port", type=int, default=7862, help="服务器端口")
    parser.add_argument("--server_name", type=str, default="0.0.0.0", help="服务器IP")
gushiqiao's avatar
gushiqiao committed
1332
    parser.add_argument("--output_dir", type=str, default="./outputs", help="输出视频保存目录")
gushiqiao's avatar
gushiqiao committed
1333
1334
    args = parser.parse_args()

gushiqiao's avatar
gushiqiao committed
1335
    global model_path, model_cls, model_size, output_dir
gushiqiao's avatar
gushiqiao committed
1336
1337
    model_path = args.model_path
    model_cls = args.model_cls
gushiqiao's avatar
gushiqiao committed
1338
    model_size = args.model_size
gushiqiao's avatar
gushiqiao committed
1339
    task = args.task
gushiqiao's avatar
gushiqiao committed
1340
    output_dir = args.output_dir
gushiqiao's avatar
gushiqiao committed
1341

gushiqiao's avatar
gushiqiao committed
1342
    main()