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

@spaces.GPU
def merge_and_upload(weight_drop_prob, scaling_factor, base_model, model_to_merge, output_path, repo_name, token):
    # Construct the command to run hf_merge.py
    command = [
        "python3", "hf_merge.py",
        "-p", str(weight_drop_prob),
        "-lambda", str(scaling_factor),
        base_model, model_to_merge, output_path
    ]
    
    # Run the command and capture the output
    result = subprocess.run(command, capture_output=True, text=True)
    
    # 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=output_path.split('/')[-1],
            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)}"

# Define the Gradio interface
iface = gr.Interface(
    fn=merge_and_upload,
    inputs=[
        gr.Slider(minimum=0, maximum=1, value=0.13, label="Weight Drop Probability"),
        gr.Number(value=3.0, label="Scaling Factor"),
        gr.Textbox(label="Base Model File/Folder"),
        gr.Textbox(label="Model to Merge"),
        gr.Textbox(label="Output Path"),
        gr.Textbox(label="Hugging Face Repo Name"),
        gr.Textbox(label="Hugging Face Token", type="password")
    ],
    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."
)

# Launch the interface
iface.launch()