Update app.py

app.py

@@ -4,27 +4,28 @@ from huggingface_hub import HfApi
 import spaces
 
 @spaces.GPU
-
-def write_repo(base_model, model_to_merge):
-    with open("repo.txt", "w") as repo:
-        repo.writelines(base_model, model_to_merge)
-
-def merge_and_upload(weight_drop_prob, scaling_factor, repo_name, token):
+def merge_and_upload(weight_drop_prob, scaling_factor, repo_name, token, base_model, model_to_merge):
     # Define a fixed output path
     output_path = "/home/user/project/output"
+    
+    # Write the base model and model to merge to the repo.txt file
+    with open("repo.txt", "w") as repo:
+        repo.write(base_model + "\n" + 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
+        "-repo_file", "repo.txt",
+        "-output_path", 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:
+    if result.returncode!= 0:
         return f"Error in merging models: {result.stderr}"
     
     # Upload the result to Hugging Face Hub
@@ -35,28 +36,25 @@ def merge_and_upload(weight_drop_prob, scaling_factor, repo_name, token):
         
         # Upload the file
         api.upload_file(
-            path_or_fileobj=output_path,
-            path_in_repo="merged_model.safetensors",
+            path_or_fileobj=output_path + "/merged_model.safetensor",
+            path_in_repo="merged_model.safetensor",
             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)}" 
+        return f"Error uploading to Hugging Face Hub: {str(e)}"
 
 # Define the Gradio interface
 iface = gr.Interface(
-    fn=write_repo,
-    inputs=[
-        gr.Textbox(label="Models"),
-        gr.Textbox(label="Model to Merge")
-    ],
     fn=merge_and_upload,
     inputs=[
         gr.Slider(minimum=0, maximum=1, value=0.3, label="Weight Drop Probability"),
         gr.Number(value=3.0, label="Scaling Factor"),
         gr.Textbox(label="Hugging Face Token", type="password"),
-        gr.Textbox(label="Repo Name")
+        gr.Textbox(label="Repo Name"),
+        gr.Textbox(label="Base Model"),
+        gr.Textbox(label="Model to Merge")
     ],
     outputs=gr.Textbox(label="Output"),
     title="Model Merger and Uploader",