"src/lib/components/icons/QuestionMarkCircle.svelte" did not exist on "d28925de0589b5caf79a26e0014ffe3e94f03833"
gradio_demo.py 47.9 KB
Newer Older
gushiqiao's avatar
gushiqiao committed
1
2
3
4
5
6
7
8
9
10
import os
import gradio as gr
import argparse
import json
import torch
import gc
from easydict import EasyDict
from datetime import datetime
from loguru import logger

gushiqiao's avatar
gushiqiao committed
11
12
import importlib.util
import psutil
gushiqiao's avatar
gushiqiao committed
13
import random
gushiqiao's avatar
gushiqiao committed
14
15
16
17
18
19
20
21
22
23

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

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


def generate_random_seed():
    return random.randint(0, MAX_NUMPY_SEED)

gushiqiao's avatar
gushiqiao committed
30

gushiqiao's avatar
gushiqiao committed
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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
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
84
85
86
87
88
89
    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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
    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()
            total_memory = memory_info[1] / (1024**3)  # Convert bytes to GB
            return total_memory
    except Exception as e:
        logger.warning(f"Failed to get GPU memory: {e}")
        return 0


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


gushiqiao's avatar
gushiqiao committed
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
def cleanup_memory():
    gc.collect()

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

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


gushiqiao's avatar
gushiqiao committed
129
130
def generate_unique_filename(output_dir):
    os.makedirs(output_dir, exist_ok=True)
gushiqiao's avatar
gushiqiao committed
131
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
gushiqiao's avatar
gushiqiao committed
132
    return os.path.join(output_dir, f"{model_cls}_{timestamp}.mp4")
gushiqiao's avatar
gushiqiao committed
133
134


gushiqiao's avatar
gushiqiao committed
135
136
137
138
139
140
141
142
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
143
144
145
146
147
148
149
150
151
152
153
154
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
155
156
global_runner = None
current_config = None
gushiqiao's avatar
gushiqiao committed
157
158
159
160
161
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
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177

available_quant_ops = get_available_quant_ops()
quant_op_choices = []
for op_name, is_installed in available_quant_ops:
    status_text = "✅ Installed" if is_installed else "❌ Not Installed"
    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 = "✅ Installed" if is_installed else "❌ Not Installed"
    display_text = f"{op_name} ({status_text})"
    attn_op_choices.append((op_name, display_text))


gushiqiao's avatar
gushiqiao committed
178
179
180
181
182
183
184
185
186
187
188
189
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
190
    use_ret_steps,
gushiqiao's avatar
gushiqiao committed
191
192
193
194
195
196
197
198
199
200
201
202
    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
203
    offload_ratio,
gushiqiao's avatar
gushiqiao committed
204
205
    t5_cpu_offload,
    unload_modules,
gushiqiao's avatar
gushiqiao committed
206
207
208
209
    t5_offload_granularity,
    attention_type,
    quant_op,
    rotary_chunk,
gushiqiao's avatar
gushiqiao committed
210
    rotary_chunk_size,
gushiqiao's avatar
gushiqiao committed
211
    clean_cuda_cache,
gushiqiao's avatar
gushiqiao committed
212
    image_path=None,
gushiqiao's avatar
gushiqiao committed
213
):
gushiqiao's avatar
gushiqiao committed
214
215
    cleanup_memory()

gushiqiao's avatar
gushiqiao committed
216
217
218
    quant_op = quant_op.split("(")[0].strip()
    attention_type = attention_type.split("(")[0].strip()

gushiqiao's avatar
gushiqiao committed
219
    global global_runner, current_config, model_path, task
gushiqiao's avatar
gushiqiao committed
220
    global cur_dit_quant_scheme, cur_clip_quant_scheme, cur_t5_quant_scheme, cur_precision_mode, cur_enable_teacache
gushiqiao's avatar
gushiqiao committed
221
222
223
224
225
226

    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)

    if task == "t2v":
gushiqiao's avatar
gushiqiao committed
227
        if model_size == "1.3b":
gushiqiao's avatar
gushiqiao committed
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
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
            # 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
303
    save_video_path = generate_unique_filename(output_dir)
gushiqiao's avatar
gushiqiao committed
304
305
306
307

    is_dit_quant = dit_quant_scheme != "bf16"
    is_t5_quant = t5_quant_scheme != "bf16"
    if is_t5_quant:
gushiqiao's avatar
gushiqiao committed
308
309
        t5_path = os.path.join(model_path, t5_quant_scheme)
        t5_quant_ckpt = os.path.join(t5_path, f"models_t5_umt5-xxl-enc-{t5_quant_scheme}.pth")
gushiqiao's avatar
gushiqiao committed
310
311
312
    else:
        t5_quant_ckpt = None

gushiqiao's avatar
gushiqiao committed
313
    is_clip_quant = clip_quant_scheme != "fp16"
gushiqiao's avatar
gushiqiao committed
314
    if is_clip_quant:
gushiqiao's avatar
gushiqiao committed
315
316
        clip_path = os.path.join(model_path, clip_quant_scheme)
        clip_quant_ckpt = os.path.join(clip_path, f"clip-{clip_quant_scheme}.pth")
gushiqiao's avatar
gushiqiao committed
317
318
319
    else:
        clip_quant_ckpt = None

gushiqiao's avatar
gushiqiao committed
320
321
    needs_reinit = (
        lazy_load
gushiqiao's avatar
gushiqiao committed
322
        or unload_modules
gushiqiao's avatar
gushiqiao committed
323
324
325
326
327
328
329
330
331
332
333
334
335
        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
336
337
338
339
340
341
342
343
344
345
346
347
348
349

    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
350
351
352
353
            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
354
355
        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
356
357
            t5_quant_scheme = f"{t5_quant_scheme}-q8f"
            clip_quant_scheme = f"{clip_quant_scheme}-q8f"
gushiqiao's avatar
gushiqiao committed
358
359

        dit_quantized_ckpt = os.path.join(model_path, dit_quant_scheme)
gushiqiao's avatar
gushiqiao committed
360
361
362
        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
363
364
        else:
            quant_model_config = {}
gushiqiao's avatar
gushiqiao committed
365
366
    else:
        mm_type = "Default"
gushiqiao's avatar
gushiqiao committed
367
        dit_quantized_ckpt = None
gushiqiao's avatar
gushiqiao committed
368
        quant_model_config = {}
gushiqiao's avatar
gushiqiao committed
369
370
371
372
373
374

    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
375
376
377
        "self_attn_1_type": attention_type,
        "cross_attn_1_type": attention_type,
        "cross_attn_2_type": attention_type,
gushiqiao's avatar
gushiqiao committed
378
379
380
381
382
383
        "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
384
        "offload_ratio": offload_ratio,
gushiqiao's avatar
gushiqiao committed
385
        "t5_offload_granularity": t5_offload_granularity,
gushiqiao's avatar
gushiqiao committed
386
        "dit_quantized_ckpt": dit_quantized_ckpt,
gushiqiao's avatar
gushiqiao committed
387
388
389
390
391
        "mm_config": {
            "mm_type": mm_type,
        },
        "fps": fps,
        "feature_caching": "Tea" if enable_teacache else "NoCaching",
gushiqiao's avatar
gushiqiao committed
392
393
        "coefficients": coefficient[0] if use_ret_steps else coefficient[1],
        "use_ret_steps": use_ret_steps,
gushiqiao's avatar
gushiqiao committed
394
        "teacache_thresh": teacache_thresh,
gushiqiao's avatar
gushiqiao committed
395
396
        "t5_cpu_offload": t5_cpu_offload,
        "unload_modules": unload_modules,
gushiqiao's avatar
gushiqiao committed
397
398
399
400
401
402
403
        "t5_quantized": is_t5_quant,
        "t5_quantized_ckpt": t5_quant_ckpt,
        "t5_quant_scheme": t5_quant_scheme,
        "clip_quantized": is_clip_quant,
        "clip_quantized_ckpt": clip_quant_ckpt,
        "clip_quant_scheme": clip_quant_scheme,
        "use_tiling_vae": use_tiling_vae,
helloyongyang's avatar
helloyongyang committed
404
        "use_tiny_vae": use_tiny_vae,
gushiqiao's avatar
gushiqiao committed
405
406
407
408
409
410
411
412
413
414
415
416
417
        "tiny_vae_path": (os.path.join(model_path, "taew2_1.pth") if use_tiny_vae else None),
        "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
418
        "rotary_chunk_size": rotary_chunk_size,
gushiqiao's avatar
gushiqiao committed
419
        "clean_cuda_cache": clean_cuda_cache,
gushiqiao's avatar
gushiqiao committed
420
        "denoising_step_list": [1000, 750, 500, 250],
gushiqiao's avatar
gushiqiao committed
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
    }

    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
437
    config.update(quant_model_config)
gushiqiao's avatar
gushiqiao committed
438
439
440
441

    logger.info(f"Using model: {model_path}")
    logger.info(f"Inference configuration:\n{json.dumps(config, indent=4, ensure_ascii=False)}")

gushiqiao's avatar
gushiqiao committed
442
    # Initialize or reuse the runner
gushiqiao's avatar
gushiqiao committed
443
444
445
446
447
448
449
    runner = global_runner
    if needs_reinit:
        if runner is not None:
            del runner
            torch.cuda.empty_cache()
            gc.collect()

gushiqiao's avatar
gushiqiao committed
450
451
        from lightx2v.infer import init_runner  # noqa

gushiqiao's avatar
gushiqiao committed
452
453
        runner = init_runner(config)
        current_config = config
gushiqiao's avatar
gushiqiao committed
454
455
456
457
458
        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
459
460
461

        if not lazy_load:
            global_runner = runner
gushiqiao's avatar
gushiqiao committed
462
463
    else:
        runner.config = config
gushiqiao's avatar
gushiqiao committed
464

465
    runner.run_pipeline()
gushiqiao's avatar
gushiqiao committed
466

gushiqiao's avatar
gushiqiao committed
467
468
469
470
471
472
473
474
475
    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
476
477
478
479

    return save_video_path


gushiqiao's avatar
gushiqiao committed
480
481
482
483
484
485
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
486
487
488
489
490
491
492
493
494
    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
495
496
        "t5_cpu_offload_val": False,
        "unload_modules_val": False,
gushiqiao's avatar
gushiqiao committed
497
498
499
500
501
502
503
504
505
506
507
508
509
        "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
510

gushiqiao's avatar
gushiqiao committed
511
512
513
514
515
516
517
518
519
520
521
    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
522
    attn_priority = ["sage_attn2", "flash_attn3", "flash_attn2", "torch_sdpa"]
gushiqiao's avatar
gushiqiao committed
523
524
525
526
527

    if is_ada_architecture_gpu():
        quant_op_priority = ["q8f", "vllm", "sgl"]
    else:
        quant_op_priority = ["sgl", "vllm", "q8f"]
gushiqiao's avatar
gushiqiao committed
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556

    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
557
    if model_size == "14b":
gushiqiao's avatar
gushiqiao committed
558
559
560
561
562
563
564
        is_14b = True
    else:
        is_14b = False

    if res == "720p" and is_14b:
        gpu_rules = [
            (80, {}),
gushiqiao's avatar
gushiqiao committed
565
566
567
            (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
568
569
570
571
            (
                24,
                {
                    "cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
572
                    "t5_cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
573
574
575
576
577
578
579
580
581
582
                    "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
583
                    "t5_cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
584
585
586
587
588
589
590
591
592
593
594
595
596
                    "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
597
                    "t5_cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
598
599
600
601
602
603
604
605
                    "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
606
                    "use_tiny_vae_val": True,
gushiqiao's avatar
gushiqiao committed
607
608
609
610
611
612
                },
            ),
            (
                8,
                {
                    "cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
613
                    "t5_cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
614
615
616
617
618
619
620
621
622
623
624
625
                    "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
626
                    "unload_modules_val": True,
gushiqiao's avatar
gushiqiao committed
627
                    "use_tiny_vae_val": True,
gushiqiao's avatar
gushiqiao committed
628
629
630
631
632
633
634
                },
            ),
        ]

    elif is_14b:
        gpu_rules = [
            (80, {}),
gushiqiao's avatar
gushiqiao committed
635
636
637
            (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
638
639
640
641
            (
                16,
                {
                    "cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
642
                    "t5_cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
643
644
645
646
647
648
649
650
651
652
653
654
                    "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
655
                        "t5_cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
656
657
658
659
660
661
662
663
664
                        "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
665
                        "unload_modules_val": True,
gushiqiao's avatar
gushiqiao committed
666
667
                        "rotary_chunk_val": True,
                        "rotary_chunk_size_val": 10000,
gushiqiao's avatar
gushiqiao committed
668
                        "use_tiny_vae_val": True,
gushiqiao's avatar
gushiqiao committed
669
670
671
672
                    }
                    if res == "540p"
                    else {
                        "cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
673
                        "t5_cpu_offload_val": True,
gushiqiao's avatar
gushiqiao committed
674
675
676
677
678
679
680
681
682
                        "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
683
                        "unload_modules_val": True,
gushiqiao's avatar
gushiqiao committed
684
                        "use_tiny_vae_val": True,
gushiqiao's avatar
gushiqiao committed
685
686
687
688
                    }
                ),
            ),
        ]
gushiqiao's avatar
gushiqiao committed
689

gushiqiao's avatar
gushiqiao committed
690
    else:
gushiqiao's avatar
gushiqiao committed
691
692
693
694
695
696
697
698
699
700
701
        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
702

gushiqiao's avatar
gushiqiao committed
703
704
705
706
707
708
709
    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
710
711
712
713
714
                {
                    "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
715
                    "unload_modules_val": True,
gushiqiao's avatar
gushiqiao committed
716
                },
gushiqiao's avatar
gushiqiao committed
717
718
            ),
        ]
gushiqiao's avatar
gushiqiao committed
719
    else:
gushiqiao's avatar
gushiqiao committed
720
721
722
723
724
725
726
727
728
729
730
        cpu_rules = [
            (64, {}),
            (
                16,
                {
                    "t5_quant_scheme_val": quant_type,
                    "unload_modules_val": True,
                    "use_tiny_vae_val": True,
                },
            ),
        ]
gushiqiao's avatar
gushiqiao committed
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745

    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
746
    def toggle_image_input(task):
gushiqiao's avatar
gushiqiao committed
747
        return gr.update(visible=(task == "i2v"))
gushiqiao's avatar
gushiqiao committed
748
749

    with gr.Blocks(
gushiqiao's avatar
gushiqiao committed
750
        title="Lightx2v (Lightweight Video Inference and Generation Engine)",
gushiqiao's avatar
gushiqiao committed
751
752
753
754
755
756
        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
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
        .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
781
782
783
    """,
    ) as demo:
        gr.Markdown(f"# 🎬 {model_cls} Video Generator")
gushiqiao's avatar
gushiqiao committed
784
        gr.Markdown(f"### Using Model: {model_path}")
gushiqiao's avatar
gushiqiao committed
785
786
787
788
789
790
791
792

        with gr.Tabs() as tabs:
            with gr.Tab("Basic Settings", id=1):
                with gr.Row():
                    with gr.Column(scale=4):
                        with gr.Group():
                            gr.Markdown("## 📥 Input Parameters")

gushiqiao's avatar
gushiqiao committed
793
794
795
796
797
798
799
800
801
                            if task == "i2v":
                                with gr.Row():
                                    image_path = gr.Image(
                                        label="Input Image",
                                        type="filepath",
                                        height=300,
                                        interactive=True,
                                        visible=True,
                                    )
gushiqiao's avatar
gushiqiao committed
802
803
804
805
806
807
808
809
810
811
812
813
814

                            with gr.Row():
                                with gr.Column():
                                    prompt = gr.Textbox(
                                        label="Prompt",
                                        lines=3,
                                        placeholder="Describe the video content...",
                                        max_lines=5,
                                    )
                                with gr.Column():
                                    negative_prompt = gr.Textbox(
                                        label="Negative Prompt",
                                        lines=3,
gushiqiao's avatar
gushiqiao committed
815
                                        placeholder="What you don't want to appear in the video...",
gushiqiao's avatar
gushiqiao committed
816
                                        max_lines=5,
gushiqiao's avatar
gushiqiao committed
817
                                        value="镜头晃动,色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走",
gushiqiao's avatar
gushiqiao committed
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
                                    )
                                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"),
                                        ],
gushiqiao's avatar
gushiqiao committed
840
841
                                        value="832x480",
                                        label="Maximum Resolution",
gushiqiao's avatar
gushiqiao committed
842
                                    )
gushiqiao's avatar
gushiqiao committed
843
844

                                with gr.Column():
gushiqiao's avatar
gushiqiao committed
845
846
847
848
849
850
851
852
                                    with gr.Group():
                                        gr.Markdown("### 🚀 **Smart Configuration Recommendation**", elem_classes=["auto-config-title"])
                                        enable_auto_config = gr.Checkbox(
                                            label="🎯 **Auto-configure Inference Options**",
                                            value=False,
                                            info="💡 **Automatically optimize GPU settings to match the current resolution. After changing the resolution, please re-check this option to prevent potential performance degradation or runtime errors.**",
                                            elem_classes=["auto-config-checkbox"],
                                        )
gushiqiao's avatar
gushiqiao committed
853
                                with gr.Column(scale=9):
gushiqiao's avatar
gushiqiao committed
854
855
                                    seed = gr.Slider(
                                        label="Random Seed",
gushiqiao's avatar
gushiqiao committed
856
857
                                        minimum=0,
                                        maximum=MAX_NUMPY_SEED,
gushiqiao's avatar
gushiqiao committed
858
                                        step=1,
gushiqiao's avatar
gushiqiao committed
859
                                        value=generate_random_seed(),
gushiqiao's avatar
gushiqiao committed
860
                                    )
gushiqiao's avatar
gushiqiao committed
861
862
863
864
865
866
                                with gr.Column(scale=1):
                                    randomize_btn = gr.Button("🎲 Randomize", variant="secondary")

                                randomize_btn.click(fn=generate_random_seed, inputs=None, outputs=seed)

                                with gr.Column():
gushiqiao's avatar
gushiqiao committed
867
                                    # Set default inference steps based on model class
gushiqiao's avatar
gushiqiao committed
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
                                    if model_cls == "wan2.1_distill":
                                        infer_steps = gr.Slider(
                                            label="Inference Steps",
                                            minimum=4,
                                            maximum=4,
                                            step=1,
                                            value=4,
                                            interactive=False,
                                            info="Inference steps fixed at 4 for optimal performance for distill model.",
                                        )
                                    elif model_cls == "wan2.1":
                                        if task == "i2v":
                                            infer_steps = gr.Slider(
                                                label="Inference Steps",
                                                minimum=1,
                                                maximum=100,
                                                step=1,
                                                value=40,
                                                info="Number of inference steps for video generation. Increasing steps may improve quality but reduce speed.",
                                            )
                                        elif task == "t2v":
                                            infer_steps = gr.Slider(
                                                label="Inference Steps",
                                                minimum=1,
                                                maximum=100,
                                                step=1,
                                                value=50,
                                                info="Number of inference steps for video generation. Increasing steps may improve quality but reduce speed.",
                                            )
gushiqiao's avatar
gushiqiao committed
897

gushiqiao's avatar
gushiqiao committed
898
899
                            # Set default CFG based on model class
                            default_enable_cfg = False if model_cls == "wan2.1_distill" else True
gushiqiao's avatar
gushiqiao committed
900
901
                            enable_cfg = gr.Checkbox(
                                label="Enable Classifier-Free Guidance",
gushiqiao's avatar
gushiqiao committed
902
                                value=default_enable_cfg,
gushiqiao's avatar
gushiqiao committed
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
                                info="Enable classifier-free guidance to control prompt strength",
                            )
                            cfg_scale = gr.Slider(
                                label="CFG Scale Factor",
                                minimum=1,
                                maximum=10,
                                step=1,
                                value=5,
                                info="Controls the influence strength of the prompt. Higher values give more influence to the prompt.",
                            )
                            sample_shift = gr.Slider(
                                label="Distribution Shift",
                                value=5,
                                minimum=0,
                                maximum=10,
                                step=1,
                                info="Controls the degree of distribution shift for samples. Larger values indicate more significant shifts.",
gushiqiao's avatar
gushiqiao committed
920
921
                            )

gushiqiao's avatar
gushiqiao committed
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
                            fps = gr.Slider(
                                label="Frames Per Second (FPS)",
                                minimum=8,
                                maximum=30,
                                step=1,
                                value=16,
                                info="Frames per second of the video. Higher FPS results in smoother videos.",
                            )
                            num_frames = gr.Slider(
                                label="Total Frames",
                                minimum=16,
                                maximum=120,
                                step=1,
                                value=81,
                                info="Total number of frames in the video. More frames result in longer videos.",
                            )
gushiqiao's avatar
gushiqiao committed
938

gushiqiao's avatar
gushiqiao committed
939
940
                        save_video_path = gr.Textbox(
                            label="Output Video Path",
gushiqiao's avatar
gushiqiao committed
941
                            value=generate_unique_filename(output_dir),
gushiqiao's avatar
gushiqiao committed
942
943
                            info="Must include .mp4 extension. If left blank or using the default value, a unique filename will be automatically generated.",
                        )
gushiqiao's avatar
gushiqiao committed
944
945
946
947
948
949
950
951
952
953
                    with gr.Column(scale=6):
                        gr.Markdown("## 📤 Generated Video")
                        output_video = gr.Video(
                            label="Result",
                            height=624,
                            width=360,
                            autoplay=True,
                            elem_classes=["output-video"],
                        )

gushiqiao's avatar
gushiqiao committed
954
                        infer_btn = gr.Button("Generate Video", variant="primary", size="lg")
gushiqiao's avatar
gushiqiao committed
955

gushiqiao's avatar
gushiqiao committed
956
957
            with gr.Tab("⚙️ Advanced Options", id=2):
                with gr.Group(elem_classes="advanced-options"):
gushiqiao's avatar
gushiqiao committed
958
                    gr.Markdown("### GPU Memory Optimization")
gushiqiao's avatar
gushiqiao committed
959
                    with gr.Row():
gushiqiao's avatar
gushiqiao committed
960
961
                        rotary_chunk = gr.Checkbox(
                            label="Chunked Rotary Position Embedding",
gushiqiao's avatar
gushiqiao committed
962
                            value=False,
gushiqiao's avatar
gushiqiao committed
963
                            info="When enabled, processes rotary position embeddings in chunks to save GPU memory.",
gushiqiao's avatar
gushiqiao committed
964
965
                        )

gushiqiao's avatar
gushiqiao committed
966
967
968
969
970
971
972
                        rotary_chunk_size = gr.Slider(
                            label="Rotary Embedding Chunk Size",
                            value=100,
                            minimum=100,
                            maximum=10000,
                            step=100,
                            info="Controls the chunk size for applying rotary embeddings. Larger values may improve performance but increase memory usage. Only effective if 'rotary_chunk' is checked.",
gushiqiao's avatar
gushiqiao committed
973
974
                        )

gushiqiao's avatar
gushiqiao committed
975
976
977
978
979
                        unload_modules = gr.Checkbox(
                            label="Unload Modules",
                            value=False,
                            info="Unload modules (T5, CLIP, DIT, etc.) after inference to reduce GPU/CPU memory usage",
                        )
gushiqiao's avatar
gushiqiao committed
980
981
982
                        clean_cuda_cache = gr.Checkbox(
                            label="Clean CUDA Memory Cache",
                            value=False,
gushiqiao's avatar
gushiqiao committed
983
                            info="When enabled, frees up GPU memory promptly but slows down inference.",
gushiqiao's avatar
gushiqiao committed
984
985
                        )

gushiqiao's avatar
gushiqiao committed
986
                    gr.Markdown("### Asynchronous Offloading")
gushiqiao's avatar
gushiqiao committed
987
988
                    with gr.Row():
                        cpu_offload = gr.Checkbox(
gushiqiao's avatar
gushiqiao committed
989
990
991
992
993
994
995
                            label="CPU Offloading",
                            value=False,
                            info="Offload parts of the model computation from GPU to CPU to reduce GPU memory usage",
                        )

                        lazy_load = gr.Checkbox(
                            label="Enable Lazy Loading",
gushiqiao's avatar
gushiqiao committed
996
                            value=False,
gushiqiao's avatar
gushiqiao committed
997
                            info="Lazy load model components during inference. Requires CPU loading and DIT quantization.",
gushiqiao's avatar
gushiqiao committed
998
                        )
gushiqiao's avatar
gushiqiao committed
999

gushiqiao's avatar
gushiqiao committed
1000
1001
1002
                        offload_granularity = gr.Dropdown(
                            label="Dit Offload Granularity",
                            choices=["block", "phase"],
gushiqiao's avatar
gushiqiao committed
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
                            value="phase",
                            info="Sets Dit model offloading granularity: blocks or computational phases",
                        )
                        offload_ratio = gr.Slider(
                            label="Offload ratio for Dit model",
                            minimum=0.0,
                            maximum=1.0,
                            step=0.1,
                            value=1.0,
                            info="Controls how much of the Dit model is offloaded to the CPU",
gushiqiao's avatar
gushiqiao committed
1013
                        )
gushiqiao's avatar
gushiqiao committed
1014
1015
1016
1017
1018
1019
                        t5_cpu_offload = gr.Checkbox(
                            label="T5 CPU Offloading",
                            value=False,
                            info="Offload the T5 Encoder model to CPU to reduce GPU memory usage",
                        )

gushiqiao's avatar
gushiqiao committed
1020
1021
1022
                        t5_offload_granularity = gr.Dropdown(
                            label="T5 Encoder Offload Granularity",
                            choices=["model", "block"],
gushiqiao's avatar
gushiqiao committed
1023
1024
                            value="model",
                            info="Controls the granularity when offloading the T5 Encoder model to CPU",
gushiqiao's avatar
gushiqiao committed
1025
1026
1027
1028
                        )

                    gr.Markdown("### Low-Precision Quantization")
                    with gr.Row():
gushiqiao's avatar
gushiqiao committed
1029
1030
1031
1032
1033
1034
                        torch_compile = gr.Checkbox(
                            label="Torch Compile",
                            value=False,
                            info="Use torch.compile to accelerate the inference process",
                        )

gushiqiao's avatar
gushiqiao committed
1035
1036
                        attention_type = gr.Dropdown(
                            label="Attention Operator",
gushiqiao's avatar
gushiqiao committed
1037
1038
1039
                            choices=[op[1] for op in attn_op_choices],
                            value=attn_op_choices[0][1],
                            info="Use appropriate attention operators to accelerate inference",
gushiqiao's avatar
gushiqiao committed
1040
1041
                        )
                        quant_op = gr.Dropdown(
gushiqiao's avatar
gushiqiao committed
1042
1043
1044
1045
1046
                            label="Quantization Matmul Operator",
                            choices=[op[1] for op in quant_op_choices],
                            value=quant_op_choices[0][1],
                            info="Select the quantization matrix multiplication operator to accelerate inference",
                            interactive=True,
gushiqiao's avatar
gushiqiao committed
1047
1048
1049
1050
1051
                        )
                        dit_quant_scheme = gr.Dropdown(
                            label="Dit",
                            choices=["fp8", "int8", "bf16"],
                            value="bf16",
gushiqiao's avatar
gushiqiao committed
1052
                            info="Quantization precision for the Dit model",
gushiqiao's avatar
gushiqiao committed
1053
1054
1055
1056
1057
                        )
                        t5_quant_scheme = gr.Dropdown(
                            label="T5 Encoder",
                            choices=["fp8", "int8", "bf16"],
                            value="bf16",
gushiqiao's avatar
gushiqiao committed
1058
                            info="Quantization precision for the T5 Encoder model",
gushiqiao's avatar
gushiqiao committed
1059
1060
1061
1062
1063
                        )
                        clip_quant_scheme = gr.Dropdown(
                            label="Clip Encoder",
                            choices=["fp8", "int8", "fp16"],
                            value="fp16",
gushiqiao's avatar
gushiqiao committed
1064
                            info="Quantization precision for the Clip Encoder",
gushiqiao's avatar
gushiqiao committed
1065
1066
                        )
                        precision_mode = gr.Dropdown(
gushiqiao's avatar
gushiqiao committed
1067
                            label="Precision Mode for Sensitive Layers",
gushiqiao's avatar
gushiqiao committed
1068
                            choices=["fp32", "bf16"],
gushiqiao's avatar
gushiqiao committed
1069
                            value="fp32",
gushiqiao's avatar
gushiqiao committed
1070
                            info="Select the numerical precision for critical model components like normalization and embedding layers. FP32 offers higher accuracy, while BF16 improves performance on compatible hardware.",
gushiqiao's avatar
gushiqiao committed
1071
1072
1073
1074
1075
                        )

                    gr.Markdown("### Variational Autoencoder (VAE)")
                    with gr.Row():
                        use_tiny_vae = gr.Checkbox(
gushiqiao's avatar
gushiqiao committed
1076
                            label="Use Tiny VAE",
gushiqiao's avatar
gushiqiao committed
1077
1078
1079
1080
                            value=False,
                            info="Use a lightweight VAE model to accelerate the decoding process",
                        )
                        use_tiling_vae = gr.Checkbox(
gushiqiao's avatar
gushiqiao committed
1081
                            label="VAE Tiling Inference",
gushiqiao's avatar
gushiqiao committed
1082
                            value=False,
gushiqiao's avatar
gushiqiao committed
1083
                            info="Use VAE tiling inference to reduce GPU memory usage",
gushiqiao's avatar
gushiqiao committed
1084
1085
1086
1087
1088
                        )

                    gr.Markdown("### Feature Caching")
                    with gr.Row():
                        enable_teacache = gr.Checkbox(
gushiqiao's avatar
gushiqiao committed
1089
                            label="Tea Cache",
gushiqiao's avatar
gushiqiao committed
1090
1091
1092
1093
1094
1095
1096
1097
                            value=False,
                            info="Cache features during inference to reduce the number of inference steps",
                        )
                        teacache_thresh = gr.Slider(
                            label="Tea Cache Threshold",
                            value=0.26,
                            minimum=0,
                            maximum=1,
gushiqiao's avatar
gushiqiao committed
1098
1099
1100
1101
1102
1103
                            info="Higher acceleration may result in lower quality —— Setting to 0.1 provides ~2.0x acceleration, setting to 0.2 provides ~3.0x acceleration",
                        )
                        use_ret_steps = gr.Checkbox(
                            label="Cache Only Key Steps",
                            value=False,
                            info="When checked, cache is written only at key steps where the scheduler returns results; when unchecked, cache is written at all steps to ensure the highest quality",
gushiqiao's avatar
gushiqiao committed
1104
1105
                        )

gushiqiao's avatar
gushiqiao committed
1106
1107
                enable_auto_config.change(
                    fn=auto_configure,
gushiqiao's avatar
gushiqiao committed
1108
                    inputs=[enable_auto_config, resolution],
gushiqiao's avatar
gushiqiao committed
1109
1110
1111
1112
1113
1114
1115
1116
1117
                    outputs=[
                        torch_compile,
                        lazy_load,
                        rotary_chunk,
                        rotary_chunk_size,
                        clean_cuda_cache,
                        cpu_offload,
                        offload_granularity,
                        offload_ratio,
gushiqiao's avatar
gushiqiao committed
1118
1119
                        t5_cpu_offload,
                        unload_modules,
gushiqiao's avatar
gushiqiao committed
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
                        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
1134
1135
1136
1137
1138
1139

                lazy_load.change(
                    fn=handle_lazy_load_change,
                    inputs=[lazy_load],
                    outputs=[unload_modules],
                )
gushiqiao's avatar
gushiqiao committed
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
        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
1169
1170
                    t5_cpu_offload,
                    unload_modules,
gushiqiao's avatar
gushiqiao committed
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
                    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
1210
1211
                    t5_cpu_offload,
                    unload_modules,
gushiqiao's avatar
gushiqiao committed
1212
1213
1214
1215
1216
1217
1218
1219
1220
                    t5_offload_granularity,
                    attention_type,
                    quant_op,
                    rotary_chunk,
                    rotary_chunk_size,
                    clean_cuda_cache,
                ],
                outputs=output_video,
            )
gushiqiao's avatar
gushiqiao committed
1221

gushiqiao's avatar
gushiqiao committed
1222
    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
1223
1224
1225


if __name__ == "__main__":
gushiqiao's avatar
gushiqiao committed
1226
1227
1228
1229
1230
    parser = argparse.ArgumentParser(description="Light Video Generation")
    parser.add_argument("--model_path", type=str, required=True, help="Model folder path")
    parser.add_argument(
        "--model_cls",
        type=str,
gushiqiao's avatar
gushiqiao committed
1231
        choices=["wan2.1", "wan2.1_distill"],
gushiqiao's avatar
gushiqiao committed
1232
        default="wan2.1",
gushiqiao's avatar
gushiqiao committed
1233
        help="Model class to use (wan2.1: standard model, wan2.1_distill: distilled model for faster inference)",
gushiqiao's avatar
gushiqiao committed
1234
    )
gushiqiao's avatar
gushiqiao committed
1235
    parser.add_argument("--model_size", type=str, required=True, choices=["14b", "1.3b"], help="Model type to use")
gushiqiao's avatar
gushiqiao committed
1236
    parser.add_argument("--task", type=str, required=True, choices=["i2v", "t2v"], help="Specify the task type. 'i2v' for image-to-video translation, 't2v' for text-to-video generation.")
gushiqiao's avatar
gushiqiao committed
1237
1238
    parser.add_argument("--server_port", type=int, default=7862, help="Server port")
    parser.add_argument("--server_name", type=str, default="0.0.0.0", help="Server ip")
gushiqiao's avatar
gushiqiao committed
1239
    parser.add_argument("--output_dir", type=str, default="./outputs", help="Output video save directory")
gushiqiao's avatar
gushiqiao committed
1240
1241
    args = parser.parse_args()

gushiqiao's avatar
gushiqiao committed
1242
    global model_path, model_cls, model_size, output_dir
gushiqiao's avatar
gushiqiao committed
1243
1244
    model_path = args.model_path
    model_cls = args.model_cls
gushiqiao's avatar
gushiqiao committed
1245
    model_size = args.model_size
gushiqiao's avatar
gushiqiao committed
1246
    task = args.task
gushiqiao's avatar
gushiqiao committed
1247
    output_dir = args.output_dir
gushiqiao's avatar
gushiqiao committed
1248

gushiqiao's avatar
gushiqiao committed
1249
    main()