converter.py

import argparse
import gc
import glob
import json
import math
import multiprocessing
import os
import re
import shutil
import sys
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from tempfile import TemporaryDirectory

import torch
from loguru import logger
from safetensors import safe_open
from safetensors import torch as st
from tqdm import tqdm

sys.path.append(str(Path(__file__).parent.parent.parent))
quant_path = str(Path(__file__).parent / "quant")
if quant_path not in sys.path:
    sys.path.insert(0, quant_path)

from h3_fp8_f16_accum import (  # noqa: E402
    FP8_F16_ACCUM_QUANTIZATION_PROFILE,
    create_h3_fp8_f16_accum_quantization,
)
from quant import *  # noqa: E402
from qwen_image_21_fp8_f16_accum import (  # noqa: E402
    QUANTIZATION_PROFILE as QWEN_IMAGE_21_FP8_F16_ACCUM_PROFILE,
)
from qwen_image_21_fp8_f16_accum import (  # noqa: E402
    QwenImage21FP8F16AccumQuantization,
)

from lightx2v.utils.lora_loader import LoRALoader  # noqa: E402
from lightx2v.utils.registry_factory import CONVERT_WEIGHT_REGISTER  # noqa: E402
from tools.convert.h3_video_vae_encoder import (  # noqa: E402
    FP8_ENCODER_CONV_MODES,
    convert_h3_video_vae_encoder_fp8,
)
from tools.convert.key_mapping import get_key_mapping_rules  # noqa: E402
from tools.convert.qwen_image_21_vae_decoder import convert_qwen_image_21_vae_decoder_fp8  # noqa: E402

dtype_mapping = {
    "int8": torch.int8,
    "int8-convrot": torch.int8,
    "fp8": torch.float8_e4m3fn,
}


def quantize_model(
    weights,
    w_bit=8,
    target_keys=["attn", "ffn"],
    adapter_keys=None,
    key_idx=2,
    ignore_key=None,
    ignore_quant_keys=None,
    linear_type="int8",
    non_linear_dtype=torch.float,
    preserve_non_quant_dtype=False,
    comfyui_mode=False,
    comfyui_keys=[],
    quantization_policy=None,
):
    """
    Quantize model weights in-place

    Args:
        weights: Model state dictionary
        w_bit: Quantization bit width
        target_keys: List of module names to quantize

    Returns:
        Modified state dictionary with quantized weights and scales
    """
    total_quantized = 0
    original_size = 0
    quantized_size = 0
    non_quantized_size = 0
    keys = list(weights.keys())

    with tqdm(keys, desc="Quantizing weights") as pbar:
        for key in pbar:
            pbar.set_postfix(current_key=key, refresh=False)

            if ignore_key is not None and any(ig_key in key for ig_key in ignore_key):
                del weights[key]
                continue

            tensor = weights[key]

            # Skip non-tensors and non-2D tensors
            if not isinstance(tensor, torch.Tensor) or tensor.dim() != 2:
                if not preserve_non_quant_dtype and tensor.dtype != non_linear_dtype:
                    weights[key] = tensor.to(non_linear_dtype)
                    non_quantized_size += weights[key].numel() * weights[key].element_size()
                else:
                    non_quantized_size += tensor.numel() * tensor.element_size()
                continue

            # Check if key matches target modules
            parts = key.split(".")

            if comfyui_mode and (comfyui_keys is not None and key in comfyui_keys):
                pass
            elif len(parts) < key_idx + 1 or parts[key_idx] not in target_keys:
                if adapter_keys is None:
                    if not preserve_non_quant_dtype and tensor.dtype != non_linear_dtype:
                        weights[key] = tensor.to(non_linear_dtype)
                        non_quantized_size += weights[key].numel() * weights[key].element_size()
                    else:
                        non_quantized_size += tensor.numel() * tensor.element_size()
                elif not any(adapter_key in parts for adapter_key in adapter_keys):
                    if not preserve_non_quant_dtype and tensor.dtype != non_linear_dtype:
                        weights[key] = tensor.to(non_linear_dtype)
                        non_quantized_size += weights[key].numel() * weights[key].element_size()
                    else:
                        non_quantized_size += tensor.numel() * tensor.element_size()
                else:
                    non_quantized_size += tensor.numel() * tensor.element_size()
                continue

            # ignore_quant_keys: keep tensor but skip quantization when key matches.
            if ignore_quant_keys is not None and any(ig_q in key for ig_q in ignore_quant_keys):
                original_tensor_size = tensor.numel() * tensor.element_size()
                original_size += original_tensor_size
                if not preserve_non_quant_dtype and tensor.dtype != non_linear_dtype:
                    weights[key] = tensor.to(non_linear_dtype)
                    non_quantized_size += weights[key].numel() * weights[key].element_size()
                else:
                    non_quantized_size += tensor.numel() * tensor.element_size()
                continue

            # try:
            original_tensor_size = tensor.numel() * tensor.element_size()
            original_size += original_tensor_size

            # Quantize tensor and store results
            quantizer = CONVERT_WEIGHT_REGISTER[linear_type](tensor)
            if quantization_policy is None:
                w_q, scales, extra = quantizer.weight_quant_func(tensor, comfyui_mode)
            else:
                w_q, scales, extra = quantization_policy.quantize_weight(key, tensor, quantizer.weight_quant_func)
            weight_global_scale = extra.get("weight_global_scale", None)  # For nvfp4
            convrot_groupsize = extra.get("convrot_groupsize", None)

            # Replace original tensor and store scales
            weights[key] = w_q
            if comfyui_mode:
                weights[key.replace(".weight", ".scale_weight")] = scales
            else:
                weights[key + "_scale"] = scales
            if weight_global_scale:
                weights[key + "_global_scale"] = weight_global_scale
            if convrot_groupsize is not None:
                weights[key.removesuffix(".weight") + ".convrot_groupsize"] = convrot_groupsize

            quantized_tensor_size = w_q.numel() * w_q.element_size()
            scale_size = scales.numel() * scales.element_size()
            extra_size = 0 if convrot_groupsize is None else convrot_groupsize.numel() * convrot_groupsize.element_size()
            quantized_size += quantized_tensor_size + scale_size + extra_size

            total_quantized += 1
            del w_q, scales

            # except Exception as e:
            #     logger.error(f"Error quantizing {key}: {str(e)}")

            gc.collect()

    original_size_mb = original_size / (1024**2)
    quantized_size_mb = quantized_size / (1024**2)
    non_quantized_size_mb = non_quantized_size / (1024**2)
    total_final_size_mb = (quantized_size + non_quantized_size) / (1024**2)
    size_reduction_mb = original_size_mb - quantized_size_mb

    logger.info(f"Quantized {total_quantized} tensors")
    logger.info(f"Original quantized tensors size: {original_size_mb:.2f} MB")
    logger.info(f"After quantization size: {quantized_size_mb:.2f} MB (includes scales)")
    logger.info(f"Non-quantized tensors size: {non_quantized_size_mb:.2f} MB")
    logger.info(f"Total final model size: {total_final_size_mb:.2f} MB")
    logger.info(f"Size reduction in quantized tensors: {size_reduction_mb:.2f} MB ({size_reduction_mb / original_size_mb * 100:.1f}%)")

    if quantization_policy is not None:
        quantization_policy.validate()
    if comfyui_mode:
        weights["scaled_fp8"] = torch.zeros(2, dtype=torch.float8_e4m3fn)

    return weights


def _validate_lora_merge(lora_path, weight_dict, lora_weights, lora_pairs, lora_diffs, alpha, require_alpha):
    """Validate a LoRA completely before modifying any base-model tensors."""
    if not lora_pairs and not lora_diffs:
        raise ValueError(f"No supported LoRA weights found in: {lora_path}")

    consumed_keys = set()
    for pair_info in lora_pairs.values():
        consumed_keys.update((pair_info["up_key"], pair_info["down_key"]))
        if pair_info["mid_key"] is not None:
            consumed_keys.add(pair_info["mid_key"])
    consumed_keys.update(diff_info["diff_key"] for diff_info in lora_diffs.values())

    tensor_keys = {key for key in lora_weights if not key.endswith(".alpha")}
    unsupported_keys = sorted(tensor_keys - consumed_keys)
    missing_model_keys = sorted((set(lora_pairs) | set(lora_diffs)) - set(weight_dict))
    missing_alpha_keys = sorted(model_key for model_key, pair_info in lora_pairs.items() if pair_info["alpha"] is None)

    shape_mismatches = []
    for model_key, pair_info in lora_pairs.items():
        if model_key not in weight_dict:
            continue
        param = weight_dict[model_key]
        lora_up = lora_weights[pair_info["up_key"]]
        lora_down = lora_weights[pair_info["down_key"]]
        if lora_up.dim() != 2 or lora_down.dim() != 2:
            shape_mismatches.append(f"{model_key}: up={tuple(lora_up.shape)}, down={tuple(lora_down.shape)}")
            continue
        expected_shape = (lora_up.shape[0], lora_down.shape[1])
        if lora_up.shape[1] != lora_down.shape[0] or tuple(param.shape) != expected_shape:
            shape_mismatches.append(f"{model_key}: base={tuple(param.shape)}, up={tuple(lora_up.shape)}, down={tuple(lora_down.shape)}")

    for model_key, diff_info in lora_diffs.items():
        if model_key in weight_dict and tuple(weight_dict[model_key].shape) != tuple(lora_weights[diff_info["diff_key"]].shape):
            shape_mismatches.append(f"{model_key}: base={tuple(weight_dict[model_key].shape)}, diff={tuple(lora_weights[diff_info['diff_key']].shape)}")

    problems = []
    if unsupported_keys:
        problems.append(f"unsupported tensors ({len(unsupported_keys)}): {unsupported_keys[:3]}")
    if missing_model_keys:
        problems.append(f"missing base-model keys ({len(missing_model_keys)}): {missing_model_keys[:3]}")
    if shape_mismatches:
        problems.append(f"shape mismatches ({len(shape_mismatches)}): {shape_mismatches[:3]}")
    if require_alpha and alpha is None and missing_alpha_keys:
        problems.append(f"alpha is missing for {len(missing_alpha_keys)} LoRA pairs; pass --lora_alpha with the training alpha (8 for the MiniMax-H3 Turbo LoRA)")
    if problems:
        raise ValueError(f"LoRA validation failed for {lora_path}: " + "; ".join(problems))


def load_loras(lora_path, weight_dict, alpha, key_mapping_rules=None, strength=1.0, strict=False, require_alpha=False):
    """
    Load and apply LoRA weights to model weights using the LoRALoader class.

    Args:
        lora_path: Path to LoRA safetensors file
        weight_dict: Model weights dictionary (will be modified in place)
        alpha: Global alpha scaling factor
        key_mapping_rules: Optional list of (pattern, replacement) regex rules for key mapping
        strength: Additional strength factor for LoRA deltas
        strict: Fail instead of producing a partially merged checkpoint
        require_alpha: Require a global or per-layer alpha for every LoRA pair

    Returns:
        Number of LoRA weight adjustments applied
    """
    if alpha is not None and not math.isfinite(alpha):
        raise ValueError(f"LoRA alpha must be finite, got {alpha}")
    if strength is not None and not math.isfinite(strength):
        raise ValueError(f"LoRA strength must be finite, got {strength}")
    logger.info(f"Loading LoRA from: {lora_path} with alpha={alpha}, strength={strength}")

    # Load LoRA weights from safetensors file
    with safe_open(lora_path, framework="pt") as f:
        lora_weights = {k: f.get_tensor(k) for k in f.keys()}

    # Create LoRA loader with key mapping rules
    lora_loader = LoRALoader(key_mapping_rules=key_mapping_rules)

    lora_pairs = lora_loader.extract_lora_pairs(lora_weights)
    lora_diffs = lora_loader.extract_lora_diffs(lora_weights)
    if strict:
        _validate_lora_merge(
            lora_path=lora_path,
            weight_dict=weight_dict,
            lora_weights=lora_weights,
            lora_pairs=lora_pairs,
            lora_diffs=lora_diffs,
            alpha=alpha,
            require_alpha=require_alpha,
        )

    # Apply LoRA weights to model
    applied_count = lora_loader.apply_lora(
        weight_dict=weight_dict,
        lora_weights=lora_weights,
        alpha=alpha,
        strength=strength,
    )
    expected_count = len(lora_pairs) + len(lora_diffs)
    if strict and applied_count != expected_count:
        raise RuntimeError(f"LoRA merge was incomplete for {lora_path}: applied {applied_count} of {expected_count} adjustments")
    return applied_count


H3_TEXT_ENCODER_LAYER_COUNT = 50
QWEN3VL_TEXT_ENCODER_PREFIX = "model.language_model"
QWEN3VL_TEXT_ENCODER_LINEAR_SUFFIXES = (
    "self_attn.q_proj.weight",
    "self_attn.k_proj.weight",
    "self_attn.v_proj.weight",
    "self_attn.o_proj.weight",
    "mlp.gate_proj.weight",
    "mlp.up_proj.weight",
    "mlp.down_proj.weight",
)
QWEN3VL_TEXT_ENCODER_NORM_SUFFIXES = (
    "input_layernorm.weight",
    "post_attention_layernorm.weight",
    "self_attn.q_norm.weight",
    "self_attn.k_norm.weight",
)


def qwen3vl_text_encoder_weight_names(layer_count):
    layer_prefixes = [f"{QWEN3VL_TEXT_ENCODER_PREFIX}.layers.{index}" for index in range(layer_count)]
    quantized_names = {f"{prefix}.{suffix}" for prefix in layer_prefixes for suffix in QWEN3VL_TEXT_ENCODER_LINEAR_SUFFIXES}
    required_names = quantized_names | {f"{prefix}.{suffix}" for prefix in layer_prefixes for suffix in QWEN3VL_TEXT_ENCODER_NORM_SUFFIXES}
    required_names.add(f"{QWEN3VL_TEXT_ENCODER_PREFIX}.embed_tokens.weight")
    return required_names, quantized_names


def h3_text_encoder_weight_names(layer_count=H3_TEXT_ENCODER_LAYER_COUNT):
    return qwen3vl_text_encoder_weight_names(layer_count)


def qwen_image_21_text_encoder_weight_names(source_root):
    with (source_root / "config.json").open("r", encoding="utf-8") as handle:
        model_config = json.load(handle)
    return qwen3vl_text_encoder_weight_names(model_config["text_config"]["num_hidden_layers"])


def _convert_qwen3vl_text_encoder_fp8(args, required_names, quantized_names, model_label, metadata_model):
    output_root = Path(args.output)
    output_path = output_root / f"{args.output_name}{args.output_ext}"
    if output_path.exists():
        raise FileExistsError(f"Refusing to overwrite an existing output: {output_path}")

    source_root = Path(args.source)
    index_path = source_root / "model.safetensors.index.json"
    if not index_path.is_file():
        raise FileNotFoundError(f"{model_label} index was not found: {index_path}")
    with index_path.open("r", encoding="utf-8") as handle:
        weight_map = json.load(handle).get("weight_map")
    if not isinstance(weight_map, dict):
        raise ValueError(f"Invalid safetensors index without a weight_map: {index_path}")

    missing = sorted(required_names - weight_map.keys())
    if missing:
        raise KeyError(f"{model_label} is missing {len(missing)} required tensors: {missing[:8]}")

    names_by_shard = defaultdict(list)
    for name in sorted(required_names):
        names_by_shard[weight_map[name]].append(name)
    missing_shards = [name for name in names_by_shard if not (source_root / name).is_file()]
    if missing_shards:
        raise FileNotFoundError(f"{model_label} is missing checkpoint shards: {missing_shards}")
    output_root.mkdir(parents=True, exist_ok=True)

    converted_weights = {}
    device = torch.device(args.device)
    quantizer_type = CONVERT_WEIGHT_REGISTER["fp8"]
    for shard_index, shard_name in enumerate(sorted(names_by_shard), start=1):
        shard_path = source_root / shard_name
        logger.info("Processing {} shard {}/{}: {}", model_label, shard_index, len(names_by_shard), shard_path.name)
        with safe_open(shard_path, framework="pt", device="cpu") as checkpoint:
            for name in tqdm(names_by_shard[shard_name], desc=shard_path.name, leave=False):
                tensor = checkpoint.get_tensor(name)
                if tensor.dtype != torch.bfloat16:
                    raise ValueError(f"Expected BF16 source tensor for {name}, got {tensor.dtype}")
                if name not in quantized_names:
                    converted_weights[name] = tensor.clone()
                else:
                    weight = tensor.to(device=device, dtype=torch.float32)
                    quantizer = quantizer_type(weight)
                    weight_fp8, weight_scale, _ = quantizer.weight_quant_func(weight)
                    converted_weights[name] = weight_fp8.cpu().contiguous()
                    converted_weights[f"{name}_scale"] = weight_scale.float().cpu().contiguous()
                    del weight, quantizer, weight_fp8, weight_scale
                del tensor
        gc.collect()
        if device.type == "cuda":
            torch.cuda.empty_cache()

    total_size = sum(tensor.numel() * tensor.element_size() for tensor in converted_weights.values())
    logger.info(
        "Saving {} FP8 checkpoint: {} tensors, {:.2f} GiB -> {}",
        model_label,
        len(converted_weights),
        total_size / (1024**3),
        output_path,
    )
    with TemporaryDirectory(prefix=f".{args.output_name}.", dir=output_root) as temporary_dir:
        temporary_path = Path(temporary_dir) / output_path.name
        st.save_file(
            converted_weights,
            temporary_path,
            metadata={
                "format": "pt",
                "model": metadata_model,
                "quantization": "fp8-sgl",
            },
        )
        os.replace(temporary_path, output_path)
    logger.info("{} FP8 checkpoint saved successfully: {}", model_label, output_path)


def convert_minimax_h3_text_encoder_fp8(args):
    """Build the LightX2V MiniMax-H3 50-layer text-prefix FP8 checkpoint."""
    if not args.quantized or args.linear_type != "fp8":
        raise ValueError("h3_text_encoder conversion requires --quantized --linear_type fp8")
    if args.output_ext != ".safetensors" or not args.single_file:
        raise ValueError("h3_text_encoder conversion requires --output_ext .safetensors --single_file")
    if args.direction is not None or args.lora_path is not None:
        raise ValueError("h3_text_encoder FP8 conversion does not support key conversion or LoRA merging")
    required_names, quantized_names = h3_text_encoder_weight_names()
    _convert_qwen3vl_text_encoder_fp8(
        args,
        required_names,
        quantized_names,
        "MiniMax-H3 text encoder",
        "MiniMax-H3 text encoder prefix (layers 0-49)",
    )


def convert_qwen_image_21_text_encoder_fp8(args):
    """Build the Qwen-Image-2.1 QwenVL language-stack FP8 checkpoint.

    Example:
        python tools/convert/converter.py \
            --source /path/to/Qwen-Image-2.1/text_encoder \
            --output /path/to/Qwen-Image-2.1-qwenvl-language-fp8-sgl \
            --output_name qwen_image_21_qwenvl_language_fp8_sgl \
            --model_type qwen_image_21_text_encoder \
            --quantized --linear_type fp8 --device cuda:0 --single_file
    """
    if not args.quantized or args.linear_type != "fp8":
        raise ValueError("qwen_image_21_text_encoder conversion requires --quantized --linear_type fp8")
    if args.output_ext != ".safetensors" or not args.single_file:
        raise ValueError("qwen_image_21_text_encoder conversion requires --output_ext .safetensors --single_file")
    if args.direction is not None or args.lora_path is not None:
        raise ValueError("Qwen-Image-2.1 QwenVL FP8 conversion does not support key conversion or LoRA merging")
    required_names, quantized_names = qwen_image_21_text_encoder_weight_names(Path(args.source))
    _convert_qwen3vl_text_encoder_fp8(
        args,
        required_names,
        quantized_names,
        "Qwen-Image-2.1 QwenVL language stack",
        "Qwen-Image-2.1 QwenVL language stack",
    )


def convert_weights(args):
    if args.model_type == "h3_text_encoder":
        convert_minimax_h3_text_encoder_fp8(args)
        return
    if args.model_type == "qwen_image_21_text_encoder":
        convert_qwen_image_21_text_encoder_fp8(args)
        return
    if args.model_type == "qwen_image_21_vae_decoder":
        convert_qwen_image_21_vae_decoder_fp8(args)
        return
    if args.model_type == "h3_video_vae_encoder":
        convert_h3_video_vae_encoder_fp8(args)
        return

    if os.path.isdir(args.source):
        src_files = sorted(glob.glob(os.path.join(args.source, "*.safetensors"), recursive=True))
    elif args.source.endswith((".pth", ".safetensors", "pt")):
        src_files = [args.source]
    else:
        raise ValueError("Invalid input path")
    if not src_files:
        raise FileNotFoundError(f"No model weight files found under: {args.source}")

    merged_weights = {}
    logger.info(f"Processing source files: {src_files}")

    # Optimize loading for better memory usage
    for file_path in tqdm(src_files, desc="Loading weights"):
        logger.info(f"Loading weights from: {file_path}")
        if file_path.endswith(".pt") or file_path.endswith(".pth"):
            weights = torch.load(file_path, map_location=args.device, weights_only=True)
            if args.model_type == "hunyuan_dit":
                weights = weights["module"]
            elif args.model_type == "self_forcing":
                weights = weights["generator_ema"]
        elif file_path.endswith(".safetensors"):
            # Use lazy loading for safetensors to reduce memory usage
            with safe_open(file_path, framework="pt", device=args.device) as f:
                # Only load tensors when needed (lazy loading)
                weights = {}
                keys = f.keys()

                # For large files, show progress
                if len(keys) > 100:
                    for k in tqdm(keys, desc=f"Loading {os.path.basename(file_path)}", leave=False):
                        weights[k] = f.get_tensor(k)
                else:
                    weights = {k: f.get_tensor(k) for k in keys}

        duplicate_keys = set(weights.keys()) & set(merged_weights.keys())
        if duplicate_keys:
            raise ValueError(f"Duplicate keys found: {duplicate_keys} in file {file_path}")

        # Update weights more efficiently
        merged_weights.update(weights)

        # Clear weights dict to free memory
        del weights
        if len(src_files) > 1:
            gc.collect()  # Force garbage collection between files

    if args.direction is not None:
        rules = get_key_mapping_rules(args.direction, args.model_type)
        converted_weights = {}
        logger.info("Converting keys...")

        # Pre-compile regex patterns for better performance
        compiled_rules = [(re.compile(pattern), replacement) for pattern, replacement in rules]

        def convert_key(key):
            """Convert a single key using compiled rules"""
            new_key = key
            for pattern, replacement in compiled_rules:
                new_key = pattern.sub(replacement, new_key)
            return new_key

        # Batch convert keys using list comprehension (faster than loop)
        keys_list = list(merged_weights.keys())

        # Use parallel processing for large models
        if len(keys_list) > 1000 and args.parallel:
            logger.info(f"Using parallel processing for {len(keys_list)} keys")
            # Use ThreadPoolExecutor for I/O bound regex operations
            num_workers = min(8, multiprocessing.cpu_count())

            with ThreadPoolExecutor(max_workers=num_workers) as executor:
                # Submit all conversion tasks
                future_to_key = {executor.submit(convert_key, key): key for key in keys_list}

                # Process results as they complete with progress bar
                for future in tqdm(as_completed(future_to_key), total=len(keys_list), desc="Converting keys (parallel)"):
                    original_key = future_to_key[future]
                    new_key = future.result()
                    converted_weights[new_key] = merged_weights[original_key]
        else:
            # For smaller models, use simple loop with less overhead
            for key in tqdm(keys_list, desc="Converting keys"):
                new_key = convert_key(key)
                converted_weights[new_key] = merged_weights[key]
    else:
        converted_weights = merged_weights

    # Apply LoRA AFTER key conversion to ensure proper key matching
    if args.lora_path is not None:
        # Handle alpha list - if single alpha, replicate for all LoRAs
        if args.lora_alpha is not None:
            if len(args.lora_alpha) == 1 and len(args.lora_path) > 1:
                args.lora_alpha = args.lora_alpha * len(args.lora_path)
            elif len(args.lora_alpha) != len(args.lora_path):
                raise ValueError(f"Number of lora_alpha ({len(args.lora_alpha)}) must match number of lora_path ({len(args.lora_path)}) or be 1")

        # Normalize strength list
        if args.lora_strength is not None:
            if len(args.lora_strength) == 1 and len(args.lora_path) > 1:
                args.lora_strength = args.lora_strength * len(args.lora_path)
            elif len(args.lora_strength) != len(args.lora_path):
                raise ValueError(f"Number of strength ({len(args.lora_strength)}) must match number of lora_path ({len(args.lora_path)}) or be 1")

        # Determine if we should apply key mapping rules to LoRA keys
        key_mapping_rules = None
        if args.lora_key_convert == "convert" and args.direction is not None:
            # Apply same conversion as model
            key_mapping_rules = get_key_mapping_rules(args.direction, args.model_type)
            logger.info("Applying key conversion to LoRA weights")
        elif args.lora_key_convert == "same":
            # Don't convert LoRA keys
            logger.info("Using original LoRA keys without conversion")
        else:  # auto
            # Auto-detect: if model was converted, try with conversion first
            if args.direction is not None:
                key_mapping_rules = get_key_mapping_rules(args.direction, args.model_type)
                logger.info("Auto mode: will try with key conversion first")

        for idx, path in enumerate(args.lora_path):
            # Pass key mapping rules to handle converted keys properly
            strength = args.lora_strength[idx] if args.lora_strength is not None else 1.0
            alpha = args.lora_alpha[idx] if args.lora_alpha is not None else None
            strict_lora = args.model_type == "h3"
            load_loras(
                path,
                converted_weights,
                alpha,
                key_mapping_rules,
                strength=strength,
                strict=strict_lora,
                require_alpha=strict_lora,
            )

    if args.quantized:
        if args.full_quantized and args.comfyui_mode:
            logger.info("Quant all tensors...")
            assert args.linear_dtype, f"Error: only support 'torch.int8' and 'torch.float8_e4m3fn'."
            for k in converted_weights.keys():
                converted_weights[k] = converted_weights[k].float().to(args.linear_dtype)
        else:
            converted_weights = quantize_model(
                converted_weights,
                w_bit=args.bits,
                target_keys=args.target_keys,
                adapter_keys=args.adapter_keys,
                key_idx=args.key_idx,
                ignore_key=args.ignore_key,
                ignore_quant_keys=args.ignore_quant_keys,
                linear_type=args.linear_type,
                non_linear_dtype=args.non_linear_dtype,
                preserve_non_quant_dtype=getattr(args, "preserve_non_quant_dtype", False),
                comfyui_mode=args.comfyui_mode,
                comfyui_keys=args.comfyui_keys,
                quantization_policy=args.quantization_policy,
            )
            if args.model_type == "h3_video_vae_decoder" and args.linear_type == "fp8":
                # FP8 decoder linears require FP16 bias; preserve other VAE tensors.
                for key in converted_weights:
                    if key.startswith("decoder.") and key.endswith(".weight_scale"):
                        bias_key = key.removesuffix(".weight_scale") + ".bias"
                        if bias_key in converted_weights:
                            converted_weights[bias_key] = converted_weights[bias_key].to(torch.float16)

    os.makedirs(args.output, exist_ok=True)

    if args.output_ext == ".pth":
        torch.save(converted_weights, os.path.join(args.output, args.output_name + ".pth"))

    else:
        index = {"metadata": {"total_size": 0}, "weight_map": {}}
        if args.single_file:
            output_filename = f"{args.output_name}.safetensors"
            output_path = os.path.join(args.output, output_filename)
            logger.info(f"Saving model to single file: {output_path}")

            # For memory efficiency with large models
            try:
                # If model is very large (over threshold), consider warning
                total_size = sum(tensor.numel() * tensor.element_size() for tensor in converted_weights.values())
                total_size_gb = total_size / (1024**3)

                if total_size_gb > 10:  # Warn if model is larger than 10GB
                    logger.warning(f"Model size is {total_size_gb:.2f}GB. This will require significant memory to save as a single file.")
                    logger.warning("Consider using --save_by_block or default chunked saving for better memory efficiency.")

                # Save the entire model as a single file
                metadata = args.quantization_policy.metadata if args.quantization_policy is not None else None
                st.save_file(converted_weights, output_path, metadata=metadata)
                logger.info(f"Model saved successfully to: {output_path} ({total_size_gb:.2f}GB)")

            except MemoryError:
                logger.error("Memory error while saving. The model is too large to save as a single file.")
                logger.error("Please use --save_by_block or remove --single_file to use chunked saving.")
                raise
            except Exception as e:
                logger.error(f"Error saving model: {e}")
                raise
        elif args.save_by_block:
            logger.info("Backward conversion: grouping weights by block")
            block_groups = defaultdict(dict)
            non_block_weights = {}
            block_pattern = re.compile(r"blocks\.(\d+)\.")

            for key, tensor in converted_weights.items():
                match = block_pattern.search(key)
                if match:
                    block_idx = match.group(1)
                    if args.model_type == "wan_animate_dit" and "face_adapter" in key:
                        block_idx = str(int(block_idx) * 5)
                    block_groups[block_idx][key] = tensor
                else:
                    non_block_weights[key] = tensor

            for block_idx, weights_dict in tqdm(block_groups.items(), desc="Saving block chunks"):
                output_filename = f"block_{block_idx}.safetensors"
                output_path = os.path.join(args.output, output_filename)
                st.save_file(weights_dict, output_path)
                for key in weights_dict:
                    index["weight_map"][key] = output_filename
                index["metadata"]["total_size"] += os.path.getsize(output_path)

            if non_block_weights:
                output_filename = f"non_block.safetensors"
                output_path = os.path.join(args.output, output_filename)
                st.save_file(non_block_weights, output_path)
                for key in non_block_weights:
                    index["weight_map"][key] = output_filename
                index["metadata"]["total_size"] += os.path.getsize(output_path)

        else:
            chunk_idx = 0
            current_chunk = {}
            for idx, (k, v) in tqdm(enumerate(converted_weights.items()), desc="Saving chunks"):
                current_chunk[k] = v
                if args.chunk_size > 0 and (idx + 1) % args.chunk_size == 0:
                    output_filename = f"{args.output_name}_part{chunk_idx}.safetensors"
                    output_path = os.path.join(args.output, output_filename)
                    logger.info(f"Saving chunk to: {output_path}")
                    st.save_file(current_chunk, output_path)
                    for key in current_chunk:
                        index["weight_map"][key] = output_filename
                    index["metadata"]["total_size"] += os.path.getsize(output_path)
                    current_chunk = {}
                    chunk_idx += 1

            if current_chunk:
                output_filename = f"{args.output_name}_part{chunk_idx}.safetensors"
                output_path = os.path.join(args.output, output_filename)
                logger.info(f"Saving final chunk to: {output_path}")
                st.save_file(current_chunk, output_path)
                for key in current_chunk:
                    index["weight_map"][key] = output_filename
                index["metadata"]["total_size"] += os.path.getsize(output_path)

        # Save index file
        if not args.single_file:
            index_path = os.path.join(args.output, "diffusion_pytorch_model.safetensors.index.json")
            with open(index_path, "w", encoding="utf-8") as f:
                json.dump(index, f, indent=2)
            logger.info(f"Index file written to: {index_path}")

    if os.path.isdir(args.source) and args.copy_no_weight_files:
        copy_non_weight_files(args.source, args.output)


def copy_non_weight_files(source_dir, target_dir):
    ignore_extensions = [".pth", ".pt", ".safetensors", ".index.json"]

    logger.info(f"Start copying non-weighted files and subdirectories...")

    for item in tqdm(os.listdir(source_dir), desc="copy non-weighted file"):
        source_item = os.path.join(source_dir, item)
        target_item = os.path.join(target_dir, item)

        try:
            if os.path.isdir(source_item):
                os.makedirs(target_item, exist_ok=True)
                copy_non_weight_files(source_item, target_item)
            elif os.path.isfile(source_item) and not any(source_item.endswith(ext) for ext in ignore_extensions):
                shutil.copy2(source_item, target_item)
                logger.debug(f"copy file: {source_item} -> {target_item}")
        except Exception as e:
            logger.error(f"copy {source_item} : {str(e)}")

    logger.info(f"Non-weight files and subdirectories copied")


def main():
    parser = argparse.ArgumentParser(description="Model weight format converter")
    parser.add_argument("-s", "--source", required=True, help="Input path (file or directory)")
    parser.add_argument("-o_e", "--output_ext", default=".safetensors", choices=[".pth", ".safetensors"])
    parser.add_argument("-o_n", "--output_name", type=str, default="converted", help="Output file name")
    parser.add_argument("-o", "--output", required=True, help="Output directory path")
    parser.add_argument(
        "-d",
        "--direction",
        choices=[None, "forward", "backward"],
        default=None,
        help="Conversion direction: forward = 'lightx2v' -> 'Diffusers', backward = reverse",
    )
    parser.add_argument(
        "-c",
        "--chunk-size",
        type=int,
        default=100,
        help="Chunk size for saving (only applies to forward), 0 = no chunking",
    )
    parser.add_argument(
        "-t",
        "--model_type",
        choices=[
            "wan_dit",
            "h3",
            "h3_video_vae_decoder",
            "h3_video_vae_encoder",
            "qwen_image_21_vae_decoder",
            "h3_text_encoder",
            "qwen_image_21_text_encoder",
            "hunyuan_dit",
            "wan_t5",
            "wan_clip",
            "wan_animate_dit",
            "qwen_image_dit",
            "qwen_image_21_dit",
            "qwen25vl_llm",
            "z_image_dit",
            "self_forcing",
        ],
        default="wan_dit",
        help="Model type",
    )
    parser.add_argument("-b", "--save_by_block", action="store_true")

    # Quantization module selection overrides.
    # If provided, these override the per-model_type defaults inside the quantized branch.
    parser.add_argument("--key-idx", dest="key_idx_override", type=int, default=None, help="Override quantize key_idx selection")
    parser.add_argument(
        "--target-keys",
        dest="target_keys_override",
        type=str,
        default=None,
        help="Override quantize target_keys (comma-separated, e.g. 'self_attn,cross_attn,ffn')",
    )
    parser.add_argument(
        "--ignore-keys",
        dest="ignore_keys_override",
        type=str,
        default=None,
        help="Override quantize ignore_key (comma-separated). Use 'none' to disable ignoring.",
    )

    parser.add_argument(
        "--ignore-quant-keys",
        dest="ignore_quant_keys_override",
        type=str,
        default=None,
        help=("Comma-separated substrings: if a tensor key contains any of these substrings, it will NOT be quantized but will still be kept and saved."),
    )

    # Quantization
    parser.add_argument("--comfyui_mode", action="store_true")
    parser.add_argument("--full_quantized", action="store_true")
    parser.add_argument("--quantized", action="store_true")
    parser.add_argument("--quantization_profile", choices=[FP8_F16_ACCUM_QUANTIZATION_PROFILE, QWEN_IMAGE_21_FP8_F16_ACCUM_PROFILE])
    parser.add_argument("--bits", type=int, default=8, choices=[8], help="Quantization bit width")
    parser.add_argument("--vae_encoder_conv_mode", choices=FP8_ENCODER_CONV_MODES)
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        help="Device to use for quantization (cpu/cuda)",
    )
    parser.add_argument(
        "--linear_type",
        type=str,
        choices=["int8", "int8-convrot", "fp8", "nvfp4", "mxfp4", "mxfp6", "mxfp8"],
        help="Quant type for linear",
    )
    parser.add_argument(
        "--non_linear_dtype",
        type=str,
        default="torch.float32",
        choices=["torch.bfloat16", "torch.float16"],
        help="Data type for non-linear",
    )
    parser.add_argument("--lora_path", type=str, nargs="+", help="Path(s) to LoRA file(s). Can specify multiple paths separated by spaces.")
    parser.add_argument(
        "--lora_alpha",
        type=float,
        nargs="+",
        default=None,
        help="Alpha used for LoRA scaling (alpha/rank). Required for MiniMax-H3.",
    )
    parser.add_argument(
        "--lora_strength",
        type=float,
        nargs="+",
        help="Additional strength factor(s) for LoRA deltas; default 1.0",
    )
    parser.add_argument("--copy_no_weight_files", action="store_true")
    parser.add_argument("--single_file", action="store_true", help="Save as a single safetensors file instead of chunking (warning: requires loading entire model in memory)")
    parser.add_argument(
        "--lora_key_convert",
        choices=["auto", "same", "convert"],
        default="auto",
        help="How to handle LoRA key conversion: 'auto' (detect from LoRA), 'same' (use original keys), 'convert' (apply same conversion as model)",
    )
    parser.add_argument("--parallel", action="store_true", default=True, help="Use parallel processing for faster conversion (default: True)")
    parser.add_argument("--no-parallel", dest="parallel", action="store_false", help="Disable parallel processing")
    args = parser.parse_args()

    # Validate conflicting arguments
    if args.single_file and args.save_by_block:
        parser.error("--single_file and --save_by_block cannot be used together. Choose one saving strategy.")

    if args.single_file and args.chunk_size > 0 and args.chunk_size != 100:
        logger.warning("--chunk_size is ignored when using --single_file option.")

    if args.quantized and args.linear_type is None:
        parser.error("--linear_type is required when --quantized is enabled (use --linear_type fp8 for MiniMax-H3 FP8)")
    if args.linear_type == "int8-convrot" and args.comfyui_mode:
        parser.error("--linear_type int8-convrot produces a LightX2V checkpoint and cannot be combined with --comfyui_mode")

    if args.model_type == "h3" and args.lora_path and args.lora_alpha is None:
        parser.error("MiniMax-H3 LoRA merging requires --lora_alpha; use --lora_alpha 8 for the MiniMax-H3 Turbo LoRA")

    def _parse_csv_override(v: str | None) -> list[str] | None:
        if v is None:
            return None
        v = v.strip()
        if v == "" or v.lower() in {"none", "null"}:
            return None
        return [x.strip() for x in v.split(",") if x.strip()]

    if args.quantized and args.model_type not in {"h3_text_encoder", "qwen_image_21_text_encoder", "h3_video_vae_encoder", "qwen_image_21_vae_decoder"}:
        args.linear_dtype = dtype_mapping.get(args.linear_type, None)
        args.non_linear_dtype = eval(args.non_linear_dtype)

        model_type_keys_map = {
            "z_image_dit": {"key_idx": 2, "target_keys": ["attention", "feed_forward", "adaLN_modulation", "linear"], "ignore_key": None},
            "qwen_image_21_dit": {
                "key_idx": 2,
                "target_keys": ["attn", "img_mlp"],
                "ignore_key": None,
                "preserve_non_quant_dtype": True,
            },
            "qwen_image_dit": {
                "key_idx": 2,
                "target_keys": ["attn", "img_mlp", "txt_mlp", "txt_mod", "img_mod"],
                "ignore_key": None,
                "comfyui_keys": [
                    "time_text_embed.timestep_embedder.linear_1.weight",
                    "time_text_embed.timestep_embedder.linear_2.weight",
                    "img_in.weight",
                    "txt_in.weight",
                    "norm_out.linear.weight",
                    "proj_out.weight",
                ],
            },
            "wan_dit": {
                "key_idx": 2,
                "target_keys": ["self_attn", "cross_attn", "ffn"],
                "ignore_key": ["ca", "audio"],
            },
            "h3": {
                "key_idx": 2,
                "target_keys": ["attn", "ff", "adaln_proj"],
                "ignore_key": None,
                # H3 deliberately mixes BF16 transformer tensors with FP32
                # input/time/output projections. Preserve that precision for
                # every tensor outside the quantized block linears.
                "preserve_non_quant_dtype": True,
            },
            "h3_video_vae_decoder": {
                "key_idx": 1,
                "target_keys": ["transformer_blocks", "proj_out"],
                "ignore_key": None,
                "preserve_non_quant_dtype": True,
            },
            "self_forcing": {
                "key_idx": 3,
                "target_keys": ["self_attn", "cross_attn", "ffn"],
                "ignore_key": None,
            },
            "wan_animate_dit": {"key_idx": 2, "target_keys": ["self_attn", "cross_attn", "ffn"], "adapter_keys": ["linear1_kv", "linear1_q", "linear2"], "ignore_key": None},
            "hunyuan_dit": {
                "key_idx": 2,
                "target_keys": [
                    "img_mod",
                    "img_attn_q",
                    "img_attn_k",
                    "img_attn_v",
                    "img_attn_proj",
                    "img_mlp",
                    "txt_mod",
                    "txt_attn_q",
                    "txt_attn_k",
                    "txt_attn_v",
                    "txt_attn_proj",
                    "txt_mlp",
                ],
                "ignore_key": None,
            },
            "wan_t5": {"key_idx": 2, "target_keys": ["attn", "ffn"], "ignore_key": None},
            "wan_clip": {
                "key_idx": 3,
                "target_keys": ["attn", "mlp"],
                "ignore_key": ["textual"],
            },
            "qwen25vl_llm": {
                "key_idx": 3,
                "target_keys": ["self_attn", "mlp"],
                "ignore_key": ["visual"],
            },
        }

        # If model_type wasn't provided by caller, keep user-provided overrides.
        # (This matters for programmatic usage / custom wrappers where args.model_type can be None.)
        if args.model_type is not None:
            if args.model_type not in model_type_keys_map:
                raise ValueError(f"Unsupported model_type for quantization overrides: {args.model_type}")

            args.target_keys = model_type_keys_map[args.model_type]["target_keys"]
            args.adapter_keys = model_type_keys_map[args.model_type]["adapter_keys"] if "adapter_keys" in model_type_keys_map[args.model_type] else None
            args.key_idx = model_type_keys_map[args.model_type]["key_idx"]
            args.ignore_key = model_type_keys_map[args.model_type]["ignore_key"]
            args.comfyui_keys = model_type_keys_map[args.model_type]["comfyui_keys"] if "comfyui_keys" in model_type_keys_map[args.model_type] else None
            args.preserve_non_quant_dtype = model_type_keys_map[args.model_type].get("preserve_non_quant_dtype", False)
        else:
            args.target_keys = None
            args.adapter_keys = None
            args.key_idx = None
            args.ignore_key = None
            args.comfyui_keys = None
            args.preserve_non_quant_dtype = False

        # Apply CLI overrides (if provided)
        if args.key_idx_override is not None:
            args.key_idx = args.key_idx_override

        target_keys_ov = _parse_csv_override(args.target_keys_override)
        if target_keys_ov is not None:
            args.target_keys = target_keys_ov

        ignore_keys_ov = _parse_csv_override(args.ignore_keys_override)
        if args.ignore_keys_override is not None:
            # Explicit override even if user passes 'none'
            args.ignore_key = ignore_keys_ov

        ignore_quant_keys_ov = _parse_csv_override(args.ignore_quant_keys_override)
        if args.ignore_quant_keys_override is not None:
            args.ignore_quant_keys = ignore_quant_keys_ov
        else:
            args.ignore_quant_keys = None

    args.quantization_policy = None
    if args.quantization_profile is not None:
        if not args.quantized or args.linear_type != "fp8" or not args.single_file or args.output_ext != ".safetensors" or args.comfyui_mode:
            parser.error("FP8-F16 accumulation conversion requires --quantized --linear_type fp8 --output_ext .safetensors --single_file without --comfyui_mode")
        try:
            if args.quantization_profile == QWEN_IMAGE_21_FP8_F16_ACCUM_PROFILE:
                args.quantization_policy = QwenImage21FP8F16AccumQuantization(args.model_type)
            else:
                args.quantization_policy = create_h3_fp8_f16_accum_quantization(args.quantization_profile, args.model_type)
        except ValueError as profile_error:
            parser.error(str(profile_error))

    if os.path.isfile(args.output):
        raise ValueError("Output path must be a directory, not a file")

    logger.info("Starting model weight conversion...")
    convert_weights(args)
    logger.info(f"Conversion completed! Files saved to: {args.output}")


if __name__ == "__main__":
    main()