def merge_lora_weights(cfg, base_model, lora_weights):
    """Merge LoRA weights into the base model.

    Args:
        base_model: MetisStage1 model instance
        lora_weights: dict of LoRA weights (from safetensors or torch.load)

    Returns:
        MetisStage1: Model with merged weights
    """
    if isinstance(base_model, LoraModel):
        base_model = base_model.model  # Get the underlying model if it's wrapped

    # Create temporary LoRA model
    lora_config = LoraConfig(
        task_type="SEQ_2_SEQ_LM",
        r=cfg.lora_r,
        lora_alpha=cfg.lora_alpha,
        target_modules=[
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
        ],
        lora_dropout=0.01,
    )

    temp_model = LoraModel(base_model, lora_config, adapter_name="default")

    # Load LoRA weights
    temp_model.load_state_dict(lora_weights, strict=False)

    # Merge weights
    merged_model = temp_model.merge_and_unload()

    return merged_model
Merge LoRA weights into the base model.

Args:
base_model: MetisStage1 model instance
lora_weights: dict of LoRA weights (from safetensors or torch.load)

Returns:
MetisStage1: Model with merged weights
models/tts/metis/metis.py