app.py

import gradio as gr
import subprocess
from huggingface_hub import HfApi
import spaces
import shutil
import logging

@spaces.GPU

def write_repo(base_model, model_to_merge):
    with open("repo.txt", "w") as repo:
        repo.write(base_model + "\n" + model_to_merge)

def merge_and_upload(base_model, model_to_merge, scaling_factor, weight_drop_prob, repo_name, token):
    # Define a fixed output path
    output_path = "/tmp/output"
    
    # Write the base model and model to merge to the repo.txt file
    write_repo(base_model, model_to_merge)
    
    # Construct the command to run hf_merge.py
    command = [
        "python3", "hf_merge.py",
        "-p", str(weight_drop_prob),
        "-lambda", str(scaling_factor),
        "repo.txt", output_path
    ]
    
    # Set up logging
    logging.basicConfig(level=logging.INFO)
    
    # Run the command and capture the output
    result = subprocess.run(command, capture_output=True, text=True)
    
    # Log the output
    logging.info(result.stdout)
    logging.error(result.stderr)
    
    # Check if the merge was successful
    if result.returncode!= 0:
        return f"Error in merging models: {result.stderr}"
    
    # Upload the result to Hugging Face Hub
    api = HfApi()
    try:
        # Create a new repo or update an existing one
        api.create_repo(repo_id=repo_name, token=token, exist_ok=True)
        
        # Upload the file
        api.upload_file(
            path_or_fileobj=output_path,
            path_in_repo="merged_model.safetensors",
            repo_id=repo_name,
            token=token
        )
        return f"Model merged and uploaded successfully to {repo_name}!"
    except Exception as e:
        return f"Error uploading to Hugging Face Hub: {str(e)}"

    # Wipe /tmp repo after each use
    shutil.rmtree("/tmp")

# Define the Gradio interface
iface = gr.Interface(
    fn=merge_and_upload,
    inputs=[
        gr.Textbox(label="Base Model"),
        gr.Textbox(label="Model to Merge"),
        gr.Dropdown(
            label="Settings",
            choices=["Scaling Factor", "Weight Drop Probability"],
            value="Scaling Factor",
            info="Select a setting to adjust"
        ),
        gr.Slider(minimum=0, maximum=1, value=0.3, label="Value", visible=False),
        gr.Textbox(label="Hugging Face Token", type="password"),
        gr.Textbox(label="Repo Name")
    ],
    outputs=gr.Textbox(label="Output"),
    title="Model Merger and Uploader",
    description="Merge two models using the Super Mario merge method and upload to Hugging Face Hub."
)

# Define a callback to toggle the visibility of the slider
def toggle_slider(setting):
    if setting == "Scaling Factor":
        return gr.Slider(minimum=0, maximum=10, value=3.0, label="Scaling Factor")
    elif setting == "Weight Drop Probability":
        return gr.Slider(minimum=0, maximum=1, value=0.3, label="Weight Drop Probability")
)

# Launch the interface
iface.launch()