Update app.py

app.py

@@ -4,7 +4,7 @@ import random
 import string
 import tempfile
 import time
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Iterable, List
 
 import gradio as gr
@@ -14,7 +14,6 @@ import yaml
 from gradio_logsview.logsview import Log, LogsView, LogsViewRunner
 from mergekit.config import MergeConfiguration
 from clean_community_org import garbage_collect_empty_models
-import spaces
 
 has_gpu = torch.cuda.is_available()
 
@@ -62,18 +61,19 @@ If you use it in your research, please cite the following paper:
   year={2024}
 }
 This Space is heavily inspired by LazyMergeKit by Maxime Labonne (see [Colab](https://colab.research.google.com/drive/1obulZ1ROXHjYLn6PPZJwRR6GzgQogxxb)).
+
 """
 
 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()
 
     if not yaml_config:
         yield 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:
@@ -83,10 +83,7 @@ def merge(yaml_config, hf_token, repo_name, profile_name):
     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",
-            )
+            yield 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.")
         is_community_model = True
@@ -132,11 +129,7 @@ def merge(yaml_config, hf_token, repo_name, profile_name):
             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 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}")
 
 
@@ -158,22 +151,25 @@ with gr.Blocks() as demo:
     button.click(fn=merge, inputs=[config, token, repo_name, profile_name], outputs=[logs])
 
 
-def _garbage_collect_every_hour():
-    while True:
-        try:
-            garbage_collect_empty_models(token=COMMUNITY_HF_TOKEN)
-        except Exception as e:
-            print("Error running garbage collection", e)
-        time.sleep(3600)
+def garbage_collect():
+    try:
+        garbage_collect_empty_models(token=COMMUNITY_HF_TOKEN)
+    except Exception as e:
+        print(f"Error running garbage collection: {e}")
 
 
-pool = ThreadPoolExecutor()
-pool.submit(_garbage_collect_every_hour)
+def schedule_garbage_collection(interval_seconds=3600):
+    while True:
+        time.sleep(interval_seconds)
+        garbage_collect()
 
 
 @spaces.GPU
 def launch():
-    demo.queue(default_concurrency_limit=2).launch(share=True)
+    pool = ThreadPoolExecutor(max_workers=1)
+    pool.submit(schedule_garbage_collection)
+
+    demo.launch(share=True)
 
 
 if __name__ == "__main__":