Update app.py

app.py

@@ -1,9 +1,8 @@
 import gradio as gr
 import subprocess
-import spaces
+from huggingface_hub import HfApi, HfFolder
 
-@spaces.GPU
-def merge_models(weight_drop_prob, scaling_factor, base_model, model_to_merge, output_path):
+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",
@@ -15,22 +14,42 @@ def merge_models(weight_drop_prob, scaling_factor, base_model, model_to_merge, o
     # Run the command and capture the output
     result = subprocess.run(command, capture_output=True, text=True)
     
-    # Return the output of the command
-    return result.stdout
+    # 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_models,
+    fn=merge_and_upload,
     inputs=[
         gr.inputs.Slider(minimum=0, maximum=1, default=0.13, label="Weight Drop Probability"),
         gr.inputs.Number(default=3.0, label="Scaling Factor"),
         gr.inputs.Textbox(label="Base Model File/Folder"),
         gr.inputs.Textbox(label="Model to Merge"),
-        gr.inputs.Textbox(label="Output Path")
+        gr.inputs.Textbox(label="Output Path"),
+        gr.inputs.Textbox(label="Hugging Face Repo Name"),
+        gr.inputs.Textbox(label="Hugging Face Token", type="password")
     ],
     outputs="text",
-    title="Model Merger",
-    description="Merge two models using the Super Mario merge method."
+    title="Model Merger and Uploader",
+    description="Merge two models using the Super Mario merge method and upload to Hugging Face Hub."
 )
 
 # Launch the interface