Update app.py

app.py

@@ -4,7 +4,7 @@ import random
 import string
 import tempfile
 import time
-from concurrent.futures import ThreadPoolExecutor
+import threading
 from typing import Iterable, List
 
 import gradio as gr
@@ -13,7 +13,7 @@ import torch
 import yaml
 from gradio_logsview.logsview import Log, LogsView, LogsViewRunner
 from mergekit.config import MergeConfiguration
-import spaces
+from huggingface_hub import spaces
 
 has_gpu = torch.cuda.is_available()
 
@@ -66,25 +66,25 @@ This Space is heavily inspired by LazyMergeKit by Maxime Labonne (see [Colab](ht
 examples = [[str(f)] for f in pathlib.Path("examples").glob("*.yaml")]
 COMMUNITY_HF_TOKEN = os.getenv("COMMUNITY_HF_TOKEN")
 
-def merge(yaml_config, hf_token, repo_name, profile_name):
-    runner = LogsViewRunner()
+def merge_process(yaml_config, hf_token, repo_name, profile_name, logs_queue):
+    runner = LogsViewRunner(logs_queue)
 
     if not yaml_config:
-        yield runner.log("Empty yaml, pick an example below", level="ERROR")
+        runner.log("Empty yaml, pick an example below", level="ERROR")
         return
 
     try:
         merge_config = MergeConfiguration.model_validate(yaml.safe_load(yaml_config))
     except Exception as e:
-        yield runner.log(f"Invalid yaml {e}", level="ERROR")
+        runner.log(f"Invalid yaml {e}", level="ERROR")
         return
 
     is_community_model = False
     if not hf_token:
         if "/" in repo_name and not repo_name.startswith("mergekit-community/"):
-            yield runner.log(f"Cannot upload merge model to namespace {repo_name.split('/')[0]}: you must provide a valid token.", level="ERROR")
+            runner.log(f"Cannot upload merge model to namespace {repo_name.split('/')[0]}: you must provide a valid token.", level="ERROR")
             return
-        yield runner.log("No HF token provided. Your merged model will be uploaded to the https://huggingface.co/mergekit-community organization.")
+        runner.log("No HF token provided. Your merged model will be uploaded to the https://huggingface.co/mergekit-community organization.")
         is_community_model = True
         if not COMMUNITY_HF_TOKEN:
             raise gr.Error("Cannot upload to community org: community token not set by Space owner.")
@@ -98,10 +98,10 @@ def merge(yaml_config, hf_token, repo_name, profile_name):
         merged_path.mkdir(parents=True, exist_ok=True)
         config_path = merged_path / "config.yaml"
         config_path.write_text(yaml_config)
-        yield runner.log(f"Merge configuration saved in {config_path}")
+        runner.log(f"Merge configuration saved in {config_path}")
 
         if not repo_name:
-            yield runner.log("No repo name provided. Generating a random one.")
+            runner.log("No repo name provided. Generating a random one.")
             repo_name = f"{profile_name}/mergekit-{merge_config.merge_method}" if profile_name else f"mergekit-{merge_config.merge_method}"
             repo_name += "-" + "".join(random.choices(string.ascii_lowercase, k=7))
             repo_name = repo_name.replace("/", "-").strip("-")
@@ -110,26 +110,42 @@ def merge(yaml_config, hf_token, repo_name, profile_name):
             repo_name = f"mergekit-community/{repo_name}"
 
         try:
-            yield runner.log(f"Creating repo {repo_name}")
+            runner.log(f"Creating repo {repo_name}")
             repo_url = api.create_repo(repo_name, exist_ok=True)
-            yield runner.log(f"Repo created: {repo_url}")
+            runner.log(f"Repo created: {repo_url}")
         except Exception as e:
-            yield runner.log(f"Error creating repo {e}", level="ERROR")
+            runner.log(f"Error creating repo {e}", level="ERROR")
             return
 
         tmp_env = os.environ.copy()
         tmp_env["HF_HOME"] = f"{tmpdirname}/.cache"
         full_cli = cli + f" --lora-merge-cache {tmpdirname}/.lora_cache"
-        yield from runner.run_command(full_cli.split(), cwd=merged_path, env=tmp_env)
+        runner.run_command(full_cli.split(), cwd=merged_path, env=tmp_env) # Removed yield from
 
         if runner.exit_code != 0:
-            yield runner.log("Merge failed. Deleting repo as no model is uploaded.", level="ERROR")
+            runner.log("Merge failed. Deleting repo as no model is uploaded.", level="ERROR")
             api.delete_repo(repo_url.repo_id)
             return
 
-        yield runner.log("Model merged successfully. Uploading to HF.")
-        yield from runner.run_python(api.upload_folder, repo_id=repo_url.repo_id, folder_path=merged_path / "merge")
-        yield runner.log(f"Model successfully uploaded to HF: {repo_url.repo_id}")
+        runner.log("Model merged successfully. Uploading to HF.")
+        runner.run_python(api.upload_folder, repo_id=repo_url.repo_id, folder_path=merged_path / "merge") # Removed yield from
+        runner.log(f"Model successfully uploaded to HF: {repo_url.repo_id}")
+
+
+def merge(yaml_config, hf_token, repo_name, profile_name):
+    logs_queue = []
+    thread = threading.Thread(target=merge_process, args=(yaml_config, hf_token, repo_name, profile_name, logs_queue))
+    thread.start()
+
+    while thread.is_alive():
+        if logs_queue:
+            yield logs_queue[:]
+            logs_queue.clear()
+        time.sleep(0.1)
+
+    if logs_queue:
+        yield logs_queue
+
 
 with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN_DESCRIPTION)
@@ -146,7 +162,7 @@ with gr.Blocks() as demo:
     gr.Examples(examples, fn=lambda s: (s,), run_on_click=True, label="Examples", inputs=[filename], outputs=[config])
     gr.Markdown(MARKDOWN_ARTICLE)
 
-    button.click(fn=merge, inputs=[config, token, repo_name, profile_name], outputs=[logs], queue=False)
+    button.click(fn=merge, inputs=[config, token, repo_name, profile_name], outputs=[logs])
 
 
 @spaces.GPU