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)

    # 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, lora_config, 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 for matching
    clean_key = key.strip('.weight')

    for lora_key, lora_weight in lora_state.items():
        if clean_key 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,
                         scale_factor, progress=gr.Progress()):
    temp_lora_dir = None
    try:
        # Validate scale factor
        if not (-2 <= scale_factor <= 2):
            error_msg = "Scale factor must be in the range [-2, 2]"
            warning_fn(error_msg)
            return f"✗ Error: {error_msg}"

        login(hf_token)

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

        # Load LoRA state (this downloads the adapter)
        lora_state, lora_config, temp_lora_dir = load_lora_state(lora_model_name)
        
        # Calculate scale with user factor
        base_scale = lora_config['lora_alpha'] / lora_config['r']
        scale = base_scale * scale_factor
        info_fn(f"Using LoRA scale: {scale} (base: {base_scale:.3f} × factor: {scale_factor})")

        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:
                            info_fn(f"Merging LoRA weights for {key}")
                            shard_merged_count += 1
                            merged_tensors += 1

                            # Convert to float32 for computation
                            original_dtype = tensor.dtype
                            tensor = tensor.to(torch.float32)
                            lora_delta = scale * lora_B.to(torch.float32) @ lora_A.to(torch.float32)

                            # Validate dimensions for additive LoRA
                            if lora_delta.shape != tensor.shape:
                                raise ValueError(f"Additive LoRA dimension mismatch for {key}: {lora_delta.shape} vs {tensor.shape}")
                            
                            tensor += lora_delta

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

                            # Clean up intermediate tensors
                            del lora_delta
                            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}\nScale factor: {scale_factor}\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**: Custom scale factors (including negative values)
"""

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

`W_new = W + scale × (α/r) × B @ A`

### Scale Factor
The scale factor (-2 ≤ scale ≤ 2) controls the strength of the LoRA merge:
- **1.0**: Full strength (default)
- **0.5**: Half strength
- **-1.0**: Reverse effect (removes LoRA impact)

### 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")
            scale_factor = gr.Slider(
                minimum=-2.0,
                maximum=2.0,
                value=1.0,
                step=0.01,
                label="Scale Factor",
                info="Strength of LoRA merge (-2 ≤ scale ≤ 2)"
            )

        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, 
                scale_factor],
        outputs=output_text
    )

    gr.Markdown(DETAILS_TEXT)

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