app.py

import gc
import gradio as gr
import torch
from huggingface_hub import hf_hub_download, HfApi, login, list_repo_files
from safetensors import safe_open
from safetensors.torch import save_file, load_file
import os
import shutil
import json

api = HfApi()

def info_fn(text):
    gr.Info(text)

def warning_fn(text):
    gr.Warning(text)

def load_lora_state(lora_model_name):
    """Download and load LoRA adapter weights"""
    temp_lora_dir = "/tmp/lora_adapter"
    os.makedirs(temp_lora_dir, exist_ok=True)

    # Download adapter config
    config_path = hf_hub_download(
        repo_id=lora_model_name,
        filename="adapter_config.json",
        local_dir=temp_lora_dir,
        local_dir_use_symlinks=False
    )

    with open(config_path, 'r') as f:
        lora_config = json.load(f)

    scale = lora_config['lora_alpha'] / lora_config['r']

    # Download adapter weights
    try:
        adapter_path = hf_hub_download(
            repo_id=lora_model_name,
            filename="adapter_model.safetensors",
            local_dir=temp_lora_dir,
            local_dir_use_symlinks=False
        )
        lora_state = load_file(adapter_path, device='cpu')
    except:
        adapter_path = hf_hub_download(
            repo_id=lora_model_name,
            filename="adapter_model.bin",
            local_dir=temp_lora_dir,
            local_dir_use_symlinks=False
        )
        lora_state = torch.load(adapter_path, map_location='cpu')

    return lora_state, scale, temp_lora_dir

def find_lora_weights(lora_state, key):
    """Find corresponding LoRA A and B weights for a given key"""
    lora_A = None
    lora_B = None

    # Remove .weight suffix and handle potential prefixes
    clean_key = key.replace('.weight', '')

    for lora_key, lora_weight in lora_state.items():
        if clean_key in lora_key or clean_key.replace('language_model.', '') in lora_key:
            if 'lora_A' in lora_key:
                lora_A = lora_weight
            elif 'lora_B' in lora_key:
                lora_B = lora_weight

    # Both should be None or both should have values
    if (lora_A is None) != (lora_B is None):
        return None, None

    return lora_A, lora_B

def download_and_upload_non_model_files(base_model_name, output_repo_name):
    """Download and upload non-model files (config, tokenizer, etc.)"""
    temp_config_dir = "/tmp/config_files"
    os.makedirs(temp_config_dir, exist_ok=True)

    try:
        # List all files in the repository
        files = list_repo_files(repo_id=base_model_name)

        # Filter non-model files
        non_model_files = [
            f for f in files
            if not (f.startswith('model') and f.endswith('.safetensors'))
        ]

        # Download and upload each non-model file
        for filename in non_model_files:
            if filename.endswith(('.gguf', '.bin')) and 'model' in filename:
                continue  # Skip other model formats

            try:
                file_path = hf_hub_download(
                    repo_id=base_model_name,
                    filename=filename,
                    local_dir=temp_config_dir,
                    local_dir_use_symlinks=False
                )

                # Upload to output repo
                api.upload_file(
                    path_or_fileobj=file_path,
                    path_in_repo=filename,
                    repo_id=output_repo_name,
                    repo_type="model"
                )

            except Exception as e:
                info_fn(f"Skipping {filename}: {e}")

    finally:
        shutil.rmtree(temp_config_dir, ignore_errors=True)

def merge_lora_efficient(hf_token, base_model_name, lora_model_name, output_repo_name,
                         multiplicative_lora, progress=gr.Progress()):
    temp_lora_dir = None
    try:
        login(hf_token)

        progress(0.1, desc="Loading LoRA adapter...")
        info_fn("Loading LoRA adapter...")

        # Load LoRA state (this downloads the adapter)
        lora_state, scale, temp_lora_dir = load_lora_state(lora_model_name)
        info_fn(f"Using LoRA scale: {scale}")

        progress(0.2, desc="Creating output repository...")

        # Create repository
        try:
            repo_url = api.create_repo(repo_id=output_repo_name, exist_ok=True)
            info_fn(f"Repository created/updated: {repo_url}")
        except Exception as e:
            warning_fn(f"Repository might already exist: {e}")

        progress(0.3, desc="Uploading configuration files...")
        info_fn("Uploading configuration files...")

        # Download and upload non-model files
        download_and_upload_non_model_files(base_model_name, output_repo_name)

        progress(0.4, desc="Finding model shards...")
        info_fn("Finding model shards...")

        # Get list of all safetensors files
        all_files = list_repo_files(repo_id=base_model_name)
        shard_files = [f for f in all_files if f.startswith('model') and f.endswith('.safetensors')]

        if not shard_files:
            raise FileNotFoundError("No model safetensors files found in the repository")

        info_fn(f"Found {len(shard_files)} model shards to process")

        merged_tensors = 0
        total_shards = len(shard_files)

        # Process each shard individually
        for i, shard_filename in enumerate(shard_files):
            progress(0.4 + (i / total_shards) * 0.5,
                    desc=f"Processing {shard_filename} ({i+1}/{total_shards})")
            info_fn(f"Processing shard {i+1}/{total_shards}: {shard_filename}")

            # Create temporary directory for this shard only
            temp_shard_dir = f"/tmp/shard_{i}"
            os.makedirs(temp_shard_dir, exist_ok=True)

            try:
                # Download the current shard
                shard_path = hf_hub_download(
                    repo_id=base_model_name,
                    filename=shard_filename,
                    local_dir=temp_shard_dir,
                    local_dir_use_symlinks=False
                )

                # Process the shard
                tensors = {}
                shard_merged_count = 0

                with safe_open(shard_path, framework='pt', device='cpu') as f:
                    # Get metadata if available
                    metadata = f.metadata() if hasattr(f, 'metadata') else {}

                    for key in f.keys():
                        tensor = f.get_tensor(key)

                        # Try to find corresponding LoRA weights
                        lora_A, lora_B = find_lora_weights(lora_state, key)

                        if lora_A is not None and lora_B is not None:
                            lora_type = "Multiplicative" if multiplicative_lora else "Additive"
                            info_fn(f"Merging {lora_type} LoRA weights for {key}")
                            shard_merged_count += 1
                            merged_tensors += 1

                            # Convert to float32 for computation
                            original_dtype = tensor.dtype
                            tensor_f32 = tensor.to(torch.float32)
                            lora_A_f32 = lora_A.to(torch.float32)
                            lora_B_f32 = lora_B.to(torch.float32)

                            if multiplicative_lora:
                                # Apply Multiplicative-LoRA: W = W + scale * B @ A @ W
                                tensor_f32 += scale * lora_B_f32 @ lora_A_f32 @ tensor_f32
                            else:
                                # Apply standard LoRA: W = W + scale * B @ A
                                tensor_f32 += scale * lora_B_f32 @ lora_A_f32

                            # Convert back to original dtype
                            tensor = tensor_f32.to(original_dtype)

                            # Clean up intermediate tensors
                            del tensor_f32, lora_A_f32, lora_B_f32
                            if torch.cuda.is_available():
                                torch.cuda.empty_cache()

                        tensors[key] = tensor

                # Save processed shard to temporary file
                output_shard_path = os.path.join(temp_shard_dir, f"processed_{shard_filename}")
                save_file(tensors, output_shard_path, metadata=metadata)

                info_fn(f"Shard {shard_filename}: Merged {shard_merged_count} tensors")

                # Upload the processed shard
                api.upload_file(
                    path_or_fileobj=output_shard_path,
                    path_in_repo=shard_filename,
                    repo_id=output_repo_name,
                    repo_type="model"
                )

                # Clean up this shard's data
                del tensors
                gc.collect()

            finally:
                # Always clean up the temporary shard directory
                shutil.rmtree(temp_shard_dir, ignore_errors=True)

        progress(1.0, desc="Upload completed!")

        success_msg = f"✓ Successfully merged and uploaded model!\nModel URL: https://huggingface.co/{output_repo_name}\nProcessed {total_shards} shards\nMerged {merged_tensors} layers with LoRA weights"
        info_fn("Merge completed successfully!")

        return success_msg

    except Exception as e:
        error_msg = f"✗ Error during merge: {str(e)}"
        warning_fn(error_msg)
        return error_msg

    finally:
        # Cleanup LoRA directory
        if temp_lora_dir and os.path.exists(temp_lora_dir):
            shutil.rmtree(temp_lora_dir, ignore_errors=True)
        gc.collect()

INTRODUCTION_TEXT = """
## Memory-Efficient LoRA Merge

This tool merges LoRA (Low-Rank Adaptation) adapters with base models using a memory-efficient approach that processes model files individually, significantly reducing memory requirements compared to traditional methods.

### Key Features
- **Minimal Memory Usage**: Processes one model shard at a time instead of loading the entire model
- **Streaming Processing**: Downloads → Processes → Uploads → Deletes each shard sequentially
- **Automatic Cleanup**: Temporary files are automatically removed after processing
- **Progress Tracking**: Real-time status updates throughout the merge process
- **Advanced Options**: Multiplicative LoRA support
"""

DETAILS_TEXT = """
### How It Works
LoRA enables efficient fine-tuning by adding small adapter weights rather than modifying the entire model. This tool applies the LoRA transformation:

- **Standard Additive-LoRA**: `W_new = W + scale × B^T @ A`
- **Multiplicative LoRA**: `W_new = W + scale × B^T @ A @ W`

### Memory Efficiency
- **Traditional approach**: Loads entire model (~15GB+ for 7B parameter models)
- **This approach**: Peak usage determined by largest shard size, not total model size
- **Result**: Enables merging of much larger models on limited hardware

### Example Usage
- **Base Model:** `microsoft/DialoGPT-medium`
- **LoRA Adapter:** `username/my-trained-lora`
- **Output Name:** `username/dialogpt-merged`

### Attribution
This tool builds upon excellent work from the community:

- **Base implementation:** [Weyaxi/merge-lora](https://huggingface.co/spaces/Weyaxi/merge-lora)
- **Memory-efficient method:** [qlora-pipe](https://github.com/tdrussell/qlora-pipe/blob/main/tools/merge_lora.py) by tdrussell
"""

with gr.Blocks(title="Memory-Efficient LoRA Merge", theme=gr.themes.Soft()) as demo:
    gr.Markdown(INTRODUCTION_TEXT)

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Configuration")
            hf_token = gr.Textbox(
                label="Hugging Face Token",
                placeholder="hf_...",
                type="password",
                info="Token with write access to create repositories"
            )
            base_model_name = gr.Textbox(
                label="Base Model Repository",
                placeholder="microsoft/DialoGPT-medium",
                info="The original model to merge LoRA into"
            )
            lora_model_name = gr.Textbox(
                label="LoRA Adapter Repository",
                placeholder="username/my-lora-adapter",
                info="Repository containing adapter_model.safetensors"
            )
            output_repo_name = gr.Textbox(
                label="Output Repository Name",
                placeholder="username/my-merged-model",
                info="Name for the new merged model repository"
            )

            gr.Markdown("### Advanced Options")
            multiplicative_lora = gr.Checkbox(
                label="Multiplicative LoRA",
                value=False,
                info="Apply a \"multiplicative-LoRA\" instead of a standard \"additive-LoRA\""
            )

        with gr.Column(scale=1):
            gr.Markdown("### Status")
            output_text = gr.Textbox(
                label="Merge Progress & Results",
                lines=20,
                interactive=False,
                show_copy_button=True
            )

    with gr.Row():
        submit_btn = gr.Button("Start LoRA Merge", variant="primary", size="lg")

    submit_btn.click(
        fn=merge_lora_efficient,
        inputs=[hf_token, base_model_name, lora_model_name, output_repo_name, multiplicative_lora],
        outputs=output_text
    )

    gr.Markdown(DETAILS_TEXT)

demo.queue()
demo.launch(show_error=True)