Update app.py

app.py

@@ -11,7 +11,7 @@ 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):
+def merge_and_upload(base_model, model_to_merge, scaling_factor, weight_drop_prob, repo_name, token, progress=gr.Progress()):
     # Define a fixed output path
     outpath = Path('/tmp/output')
     if outpath.exists() and outpath.is_dir():
@@ -29,17 +29,23 @@ def merge_and_upload(base_model, model_to_merge, scaling_factor, weight_drop_pro
 
     # Set up logging
     logging.basicConfig(level=logging.INFO)
+    log_output = ""
 
     # Run the command and capture the output
     result = subprocess.run(command, capture_output=True, text=True)
 
     # Log the output
+    log_output += result.stdout + "\n"
+    log_output += result.stderr + "\n"
     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}"
+        return log_output, None, f"Error in merging models: {result.stderr}"
+
+    # Update progress bar
+    progress(0.5, desc="Merging completed. Uploading to Hugging Face Hub...")
 
     # Upload the result to Hugging Face Hub
     api = HfApi(token=token)
@@ -61,9 +67,10 @@ def merge_and_upload(base_model, model_to_merge, scaling_factor, weight_drop_pro
             repo_type="model",
             token=token
         )
-        return f"Model merged and uploaded successfully to {repo_name}!"
+        repo_url = f"https://huggingface.co/{repo_name}"
+        return log_output, repo_url, "Model merged and uploaded successfully!"
     except Exception as e:
-        return f"Error uploading to Hugging Face Hub: {str(e)}"
+        return log_output, None, f"Error uploading to Hugging Face Hub: {str(e)}"
 
 # Define the Gradio interface
 with gr.Blocks() as demo:
@@ -88,13 +95,16 @@ with gr.Blocks() as demo:
         with gr.Row():
             weight_drop_prob = gr.Slider(minimum=0, maximum=1, value=0.3, label="Weight Drop Probability")
 
+        progress = gr.Progress()
         gr.Button("Merge and Upload").click(
             merge_and_upload,
             inputs=[base_model, model_to_merge, scaling_factor, weight_drop_prob, repo_name, token],
-            result=output
+            outputs=[log_output, repo_url, output]
         )
 
     with gr.Column():
+        log_output = gr.Textbox(label="Log Output")
+        repo_url = gr.Markdown(label="Repository URL")
         output = gr.Textbox(label="Output")
 
 demo.launch()