gradio_demo_zh.py

import argparse
import gc
import glob
import importlib.util
import json
import os
import random
from datetime import datetime

import gradio as gr
import psutil
import torch
from loguru import logger

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

MAX_NUMPY_SEED = 2**32 - 1


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.")


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


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))

    torch_installed = is_module_installed("torch")
    if torch_installed:
        available_ops.append(("torch_sdpa", True))
    else:
        available_ops.append(("torch_sdpa", False))

    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"获取GPU内存失败: {e}")
        return 0


def get_cpu_memory():
    available_bytes = psutil.virtual_memory().available
    return available_bytes / 1024**3


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


def generate_unique_filename(output_dir):
    os.makedirs(output_dir, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    return os.path.join(output_dir, f"{model_cls}_{timestamp}.mp4")


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)


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


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)

        # 如果没有子目录但有原始文件，添加原始类型
        if has_original_file:
            if not choices or "original" not in choices:
                choices.append(original_type)

        # 如果没有任何选项，使用默认值
        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}


global_runner = None
current_config = None
cur_dit_quant_scheme = None
cur_clip_quant_scheme = None
cur_t5_quant_scheme = None
cur_precision_mode = None
cur_enable_teacache = None

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))


def run_inference(
    prompt,
    negative_prompt,
    save_result_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_tae,
    use_tiling_vae,
    lazy_load,
    precision_mode,
    cpu_offload,
    offload_granularity,
    offload_ratio,
    t5_cpu_offload,
    unload_modules,
    t5_offload_granularity,
    attention_type,
    quant_op,
    rotary_chunk,
    rotary_chunk_size,
    clean_cuda_cache,
    image_path=None,
):
    cleanup_memory()

    quant_op = quant_op.split("(")[0].strip()
    attention_type = attention_type.split("(")[0].strip()

    global global_runner, current_config, model_path, task
    global cur_dit_quant_scheme, cur_clip_quant_scheme, cur_t5_quant_scheme, cur_precision_mode, cur_enable_teacache

    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)
    else:
        model_config = {}

    if task == "t2v":
        if model_size == "1.3b":
            # 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,
                ],
            ]

    save_result_path = generate_unique_filename(output_dir)

    is_dit_quant = dit_quant_scheme != "bf16"
    is_t5_quant = t5_quant_scheme != "bf16"

    if is_t5_quant:
        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
    else:
        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")

    is_clip_quant = clip_quant_scheme != "fp16"

    if is_clip_quant:
        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
    else:
        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")

    needs_reinit = (
        lazy_load
        or unload_modules
        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
    )

    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":
            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"
        elif quant_op == "q8f":
            mm_type = f"W-{dit_quant_scheme}-channel-sym-A-{dit_quant_scheme}-channel-sym-dynamic-Q8F"
            t5_quant_scheme = f"{t5_quant_scheme}-q8f"
            clip_quant_scheme = f"{clip_quant_scheme}-q8f"

        dit_quantized_ckpt = find_hf_model_path(model_path, dit_quant_scheme)
        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)
        else:
            quant_model_config = {}
    else:
        mm_type = "Default"
        dit_quantized_ckpt = None
        quant_model_config = {}

    config = {
        "infer_steps": infer_steps,
        "target_video_length": num_frames,
        "target_width": int(resolution.split("x")[0]),
        "target_height": int(resolution.split("x")[1]),
        "self_attn_1_type": attention_type,
        "cross_attn_1_type": attention_type,
        "cross_attn_2_type": attention_type,
        "seed": seed,
        "enable_cfg": enable_cfg,
        "sample_guide_scale": cfg_scale,
        "sample_shift": sample_shift,
        "cpu_offload": cpu_offload,
        "offload_granularity": offload_granularity,
        "offload_ratio": offload_ratio,
        "t5_offload_granularity": t5_offload_granularity,
        "dit_quantized_ckpt": dit_quantized_ckpt,
        "mm_config": {
            "mm_type": mm_type,
        },
        "fps": fps,
        "feature_caching": "Tea" if enable_teacache else "NoCaching",
        "coefficients": coefficient[0] if use_ret_steps else coefficient[1],
        "use_ret_steps": use_ret_steps,
        "teacache_thresh": teacache_thresh,
        "t5_original_ckpt": t5_original_ckpt,
        "t5_cpu_offload": t5_cpu_offload,
        "unload_modules": unload_modules,
        "t5_quantized": is_t5_quant,
        "t5_quantized_ckpt": t5_quantized_ckpt,
        "t5_quant_scheme": t5_quant_scheme,
        "clip_original_ckpt": clip_original_ckpt,
        "clip_quantized": is_clip_quant,
        "clip_quantized_ckpt": clip_quantized_ckpt,
        "clip_quant_scheme": clip_quant_scheme,
        "vae_path": find_torch_model_path(model_path, "Wan2.1_VAE.pth"),
        "use_tiling_vae": use_tiling_vae,
        "use_tae": use_tae,
        "tae_pth": (find_torch_model_path(model_path, "taew2_1.pth") if use_tae 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,
        "rotary_chunk_size": rotary_chunk_size,
        "clean_cuda_cache": clean_cuda_cache,
        "denoising_step_list": [1000, 750, 500, 250],
    }

    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_result_path=save_result_path,
    )

    config.update({k: v for k, v in vars(args).items()})
    config.update(model_config)
    config.update(quant_model_config)

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

    # Initialize or reuse the runner
    runner = global_runner
    if needs_reinit:
        if runner is not None:
            del runner
            torch.cuda.empty_cache()
            gc.collect()

        from lightx2v.infer import init_runner  # noqa

        runner = init_runner(config)
        current_config = config
        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

        if not lazy_load:
            global_runner = runner
    else:
        runner.config = config

    runner.run_pipeline()

    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()

    return save_result_path


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):
    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,
        "t5_cpu_offload_val": False,
        "unload_modules_val": False,
        "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_tae_val": False,
        "use_tiling_vae_val": False,
        "enable_teacache_val": False,
        "teacache_thresh_val": 0.26,
        "use_ret_steps_val": False,
    }

    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"

    attn_priority = ["sage_attn2", "flash_attn3", "flash_attn2", "torch_sdpa"]

    if is_ada_architecture_gpu():
        quant_op_priority = ["q8f", "vllm", "sgl"]
    else:
        quant_op_priority = ["sgl", "vllm", "q8f"]

    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"

    if model_size == "14b":
        is_14b = True
    else:
        is_14b = False

    if res == "720p" and is_14b:
        gpu_rules = [
            (80, {}),
            (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}),
            (
                24,
                {
                    "cpu_offload_val": True,
                    "t5_cpu_offload_val": True,
                    "offload_ratio_val": 1,
                    "t5_offload_granularity_val": "block",
                    "precision_mode_val": "bf16",
                    "use_tiling_vae_val": True,
                },
            ),
            (
                16,
                {
                    "cpu_offload_val": True,
                    "t5_cpu_offload_val": True,
                    "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,
                    "t5_cpu_offload_val": True,
                    "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,
                    "use_tae_val": True,
                },
            ),
            (
                8,
                {
                    "cpu_offload_val": True,
                    "t5_cpu_offload_val": True,
                    "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,
                    "unload_modules_val": True,
                    "use_tae_val": True,
                },
            ),
        ]

    elif is_14b:
        gpu_rules = [
            (80, {}),
            (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}),
            (
                16,
                {
                    "cpu_offload_val": True,
                    "t5_cpu_offload_val": True,
                    "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,
                        "t5_cpu_offload_val": True,
                        "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,
                        "unload_modules_val": True,
                        "rotary_chunk_val": True,
                        "rotary_chunk_size_val": 10000,
                        "use_tae_val": True,
                    }
                    if res == "540p"
                    else {
                        "cpu_offload_val": True,
                        "t5_cpu_offload_val": True,
                        "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,
                        "unload_modules_val": True,
                        "use_tae_val": True,
                    }
                ),
            ),
        ]

    else:
        gpu_rules = [
            (24, {}),
            (
                8,
                {
                    "t5_cpu_offload_val": True,
                    "t5_offload_granularity_val": "block",
                    "t5_quant_scheme_val": quant_type,
                },
            ),
        ]

    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,
                {
                    "dit_quant_scheme_val": quant_type,
                    "t5_quant_scheme_val": quant_type,
                    "clip_quant_scheme_val": quant_type,
                    "lazy_load_val": True,
                    "unload_modules_val": True,
                },
            ),
        ]
    else:
        cpu_rules = [
            (64, {}),
            (
                16,
                {
                    "t5_quant_scheme_val": quant_type,
                    "unload_modules_val": True,
                    "use_tae_val": True,
                },
            ),
        ]

    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():
    with gr.Blocks(
        title="Lightx2v (轻量级视频推理和生成引擎)",
        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; }
        .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;
        }
    """,
    ) as demo:
        gr.Markdown(f"# 🎬 {model_cls} 视频生成器")
        gr.Markdown(f"### 使用模型: {model_path}")

        with gr.Tabs() as tabs:
            with gr.Tab("基本设置", id=1):
                with gr.Row():
                    with gr.Column(scale=4):
                        with gr.Group():
                            gr.Markdown("## 📥 输入参数")

                            if task == "i2v":
                                with gr.Row():
                                    image_path = gr.Image(
                                        label="输入图像",
                                        type="filepath",
                                        height=300,
                                        interactive=True,
                                        visible=True,
                                    )

                            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,
                                        placeholder="不希望出现在视频中的内容...",
                                        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="最大分辨率",
                                    )

                                with gr.Column():
                                    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"],
                                        )
                                with gr.Column(scale=9):
                                    seed = gr.Slider(
                                        label="随机种子",
                                        minimum=0,
                                        maximum=MAX_NUMPY_SEED,
                                        step=1,
                                        value=generate_random_seed(),
                                    )
                                with gr.Column(scale=1):
                                    randomize_btn = gr.Button("🎲 随机化", variant="secondary")

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

                                with gr.Column():
                                    # 根据模型类别设置默认推理步数
                                    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="视频生成的推理步数。增加步数可能提高质量但降低速度。",
                                            )

                            # 根据模型类别设置默认CFG
                            default_enable_cfg = False if model_cls == "wan2.1_distill" else True
                            enable_cfg = gr.Checkbox(
                                label="启用无分类器引导",
                                value=default_enable_cfg,
                                info="启用无分类器引导以控制提示词强度",
                            )
                            cfg_scale = gr.Slider(
                                label="CFG缩放因子",
                                minimum=1,
                                maximum=10,
                                step=1,
                                value=5,
                                info="控制提示词的影响强度。值越高，提示词的影响越大。",
                            )
                            sample_shift = gr.Slider(
                                label="分布偏移",
                                value=5,
                                minimum=0,
                                maximum=10,
                                step=1,
                                info="控制样本分布偏移的程度。值越大表示偏移越明显。",
                            )

                            fps = gr.Slider(
                                label="每秒帧数(FPS)",
                                minimum=8,
                                maximum=30,
                                step=1,
                                value=16,
                                info="视频的每秒帧数。较高的FPS会产生更流畅的视频。",
                            )
                            num_frames = gr.Slider(
                                label="总帧数",
                                minimum=16,
                                maximum=120,
                                step=1,
                                value=81,
                                info="视频中的总帧数。更多帧数会产生更长的视频。",
                            )

                        save_result_path = gr.Textbox(
                            label="输出视频路径",
                            value=generate_unique_filename(output_dir),
                            info="必须包含.mp4扩展名。如果留空或使用默认值，将自动生成唯一文件名。",
                        )
                    with gr.Column(scale=6):
                        gr.Markdown("## 📤 生成的视频")
                        output_video = gr.Video(
                            label="结果",
                            height=624,
                            width=360,
                            autoplay=True,
                            elem_classes=["output-video"],
                        )

                        infer_btn = gr.Button("生成视频", variant="primary", size="lg")

            with gr.Tab("⚙️ 高级选项", id=2):
                with gr.Group(elem_classes="advanced-options"):
                    gr.Markdown("### GPU内存优化")
                    with gr.Row():
                        rotary_chunk = gr.Checkbox(
                            label="分块旋转位置编码",
                            value=False,
                            info="启用时，将旋转位置编码分块处理以节省GPU内存。",
                        )

                        rotary_chunk_size = gr.Slider(
                            label="旋转编码块大小",
                            value=100,
                            minimum=100,
                            maximum=10000,
                            step=100,
                            info="控制应用旋转编码的块大小。较大的值可能提高性能但增加内存使用。仅在'rotary_chunk'勾选时有效。",
                        )
                        unload_modules = gr.Checkbox(
                            label="卸载模块",
                            value=False,
                            info="推理后卸载模块（T5、CLIP、DIT等）以减少GPU/CPU内存使用",
                        )
                        clean_cuda_cache = gr.Checkbox(
                            label="清理CUDA内存缓存",
                            value=False,
                            info="启用时，及时释放GPU内存但会减慢推理速度。",
                        )

                    gr.Markdown("### 异步卸载")
                    with gr.Row():
                        cpu_offload = gr.Checkbox(
                            label="CPU卸载",
                            value=False,
                            info="将模型计算的一部分从GPU卸载到CPU以减少GPU内存使用",
                        )

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

                        offload_granularity = gr.Dropdown(
                            label="Dit卸载粒度",
                            choices=["block", "phase"],
                            value="phase",
                            info="设置Dit模型卸载粒度：块或计算阶段",
                        )
                        offload_ratio = gr.Slider(
                            label="Dit模型卸载比例",
                            minimum=0.0,
                            maximum=1.0,
                            step=0.1,
                            value=1.0,
                            info="控制将多少Dit模型卸载到CPU",
                        )
                        t5_cpu_offload = gr.Checkbox(
                            label="T5 CPU卸载",
                            value=False,
                            info="将T5编码器模型卸载到CPU以减少GPU内存使用",
                        )
                        t5_offload_granularity = gr.Dropdown(
                            label="T5编码器卸载粒度",
                            choices=["model", "block"],
                            value="model",
                            info="控制将T5编码器模型卸载到CPU时的粒度",
                        )

                    gr.Markdown("### 低精度量化")
                    with gr.Row():
                        torch_compile = gr.Checkbox(
                            label="Torch编译",
                            value=False,
                            info="使用torch.compile加速推理过程",
                        )

                        attention_type = gr.Dropdown(
                            label="注意力算子",
                            choices=[op[1] for op in attn_op_choices],
                            value=attn_op_choices[0][1],
                            info="使用适当的注意力算子加速推理",
                        )
                        quant_op = gr.Dropdown(
                            label="量化矩阵乘法算子",
                            choices=[op[1] for op in quant_op_choices],
                            value=quant_op_choices[0][1],
                            info="选择量化矩阵乘法算子以加速推理",
                            interactive=True,
                        )
                        # 获取动态量化选项
                        quant_options = get_quantization_options(model_path)

                        dit_quant_scheme = gr.Dropdown(
                            label="Dit",
                            choices=quant_options["dit_choices"],
                            value=quant_options["dit_default"],
                            info="Dit模型的量化精度",
                        )
                        t5_quant_scheme = gr.Dropdown(
                            label="T5编码器",
                            choices=quant_options["t5_choices"],
                            value=quant_options["t5_default"],
                            info="T5编码器模型的量化精度",
                        )
                        clip_quant_scheme = gr.Dropdown(
                            label="Clip编码器",
                            choices=quant_options["clip_choices"],
                            value=quant_options["clip_default"],
                            info="Clip编码器的量化精度",
                        )
                        precision_mode = gr.Dropdown(
                            label="敏感层精度模式",
                            choices=["fp32", "bf16"],
                            value="fp32",
                            info="选择用于关键模型组件（如归一化和嵌入层）的数值精度。FP32提供更高精度，而BF16在兼容硬件上提高性能。",
                        )

                    gr.Markdown("### 变分自编码器(VAE)")
                    with gr.Row():
                        use_tae = gr.Checkbox(
                            label="使用轻量级VAE",
                            value=False,
                            info="使用轻量级VAE模型加速解码过程",
                        )
                        use_tiling_vae = gr.Checkbox(
                            label="VAE分块推理",
                            value=False,
                            info="使用VAE分块推理以减少GPU内存使用",
                        )

                    gr.Markdown("### 特征缓存")
                    with gr.Row():
                        enable_teacache = gr.Checkbox(
                            label="Tea Cache",
                            value=False,
                            info="在推理过程中缓存特征以减少推理步数",
                        )
                        teacache_thresh = gr.Slider(
                            label="Tea Cache阈值",
                            value=0.26,
                            minimum=0,
                            maximum=1,
                            info="较高的加速可能导致质量下降 —— 设置为0.1提供约2.0倍加速，设置为0.2提供约3.0倍加速",
                        )
                        use_ret_steps = gr.Checkbox(
                            label="仅缓存关键步骤",
                            value=False,
                            info="勾选时，仅在调度器返回结果的关键步骤写入缓存；未勾选时，在所有步骤写入缓存以确保最高质量",
                        )

                enable_auto_config.change(
                    fn=auto_configure,
                    inputs=[enable_auto_config, resolution],
                    outputs=[
                        torch_compile,
                        lazy_load,
                        rotary_chunk,
                        rotary_chunk_size,
                        clean_cuda_cache,
                        cpu_offload,
                        offload_granularity,
                        offload_ratio,
                        t5_cpu_offload,
                        unload_modules,
                        t5_offload_granularity,
                        attention_type,
                        quant_op,
                        dit_quant_scheme,
                        t5_quant_scheme,
                        clip_quant_scheme,
                        precision_mode,
                        use_tae,
                        use_tiling_vae,
                        enable_teacache,
                        teacache_thresh,
                        use_ret_steps,
                    ],
                )

                lazy_load.change(
                    fn=handle_lazy_load_change,
                    inputs=[lazy_load],
                    outputs=[unload_modules],
                )
        if task == "i2v":
            infer_btn.click(
                fn=run_inference,
                inputs=[
                    prompt,
                    negative_prompt,
                    save_result_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_tae,
                    use_tiling_vae,
                    lazy_load,
                    precision_mode,
                    cpu_offload,
                    offload_granularity,
                    offload_ratio,
                    t5_cpu_offload,
                    unload_modules,
                    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_result_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_tae,
                    use_tiling_vae,
                    lazy_load,
                    precision_mode,
                    cpu_offload,
                    offload_granularity,
                    offload_ratio,
                    t5_cpu_offload,
                    unload_modules,
                    t5_offload_granularity,
                    attention_type,
                    quant_op,
                    rotary_chunk,
                    rotary_chunk_size,
                    clean_cuda_cache,
                ],
                outputs=output_video,
            )

    demo.launch(share=True, server_port=args.server_port, server_name=args.server_name, inbrowser=True, allowed_paths=[output_dir])


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

    global model_path, model_cls, model_size, output_dir
    model_path = args.model_path
    model_cls = args.model_cls
    model_size = args.model_size
    task = args.task
    output_dir = args.output_dir

    main()