Update app.py

app.py

@@ -14,6 +14,7 @@ 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()
 
@@ -54,21 +55,20 @@ A quick overview of the currently supported merge methods:
 ## Citation
 This GUI is powered by [Arcee's MergeKit](https://arxiv.org/abs/2403.13257).
 If you use it in your research, please cite the following paper:
-```
 @article{goddard2024arcee,
   title={Arcee's MergeKit: A Toolkit for Merging Large Language Models},
   author={Goddard, Charles and Siriwardhana, Shamane and Ehghaghi, Malikeh and Meyers, Luke and Karpukhin, Vlad and Benedict, Brian and McQuade, Mark and Solawetz, Jacob},
   journal={arXiv preprint arXiv:2403.13257},
   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: str, hf_token: str, repo_name: str, profile_name: str) -> Iterable[List[Log]]:
+
+def merge(yaml_config, hf_token, repo_name, profile_name):
     runner = LogsViewRunner()
 
     if not yaml_config:
@@ -88,9 +88,7 @@ def merge(yaml_config: str, hf_token: str, repo_name: str, profile_name: str) ->
                 level="ERROR",
             )
             return
-        yield runner.log(
-            "No HF token provided. Your merged model will be uploaded to the https://huggingface.co/mergekit-community organization."
-        )
+        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
         if not COMMUNITY_HF_TOKEN:
             raise gr.Error("Cannot upload to community org: community token not set by Space owner.")
@@ -141,6 +139,7 @@ def merge(yaml_config: str, hf_token: str, repo_name: str, profile_name: str) ->
         )
         yield runner.log(f"Model successfully uploaded to HF: {repo_url.repo_id}")
 
+
 with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN_DESCRIPTION)
 
@@ -148,37 +147,17 @@ with gr.Blocks() as demo:
         filename = gr.Textbox(visible=False, label="filename")
         config = gr.Code(language="yaml", lines=10, label="config.yaml")
         with gr.Column():
-            token = gr.Textbox(
-                lines=1,
-                label="HF Write Token",
-                info="https://hf.co/settings/token",
-                type="password",
-                placeholder="Optional. Will upload merged model to MergeKit Community if empty.",
-            )
-            repo_name = gr.Textbox(
-                lines=1,
-                label="Repo name",
-                placeholder="Optional. Will create a random name if empty.",
-            )
-            profile_name = gr.Textbox(
-                lines=1,
-                label="Hugging Face Profile Name",
-                placeholder="Enter your Hugging Face profile name.",
-            )
+            token = gr.Textbox(lines=1, label="HF Write Token", info="https://hf.co/settings/token", type="password", placeholder="Optional. Will upload merged model to MergeKit Community if empty.")
+            repo_name = gr.Textbox(lines=1, label="Repo name", placeholder="Optional. Will create a random name if empty.")
+            profile_name = gr.Textbox(lines=1, label="Hugging Face Profile Name", placeholder="Enter your Hugging Face profile name.")
     button = gr.Button("Merge", variant="primary")
     logs = LogsView(label="Terminal output")
-    gr.Examples(
-        examples,
-        fn=lambda s: (s,),
-        run_on_click=True,
-        label="Examples",
-        inputs=[filename],
-        outputs=[config],
-    )
+    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])
 
+
 def _garbage_collect_every_hour():
     while True:
         try:
@@ -187,7 +166,15 @@ def _garbage_collect_every_hour():
             print("Error running garbage collection", e)
         time.sleep(3600)
 
+
 pool = ThreadPoolExecutor()
 pool.submit(_garbage_collect_every_hour)
 
-demo.queue(default_concurrency_limit=2).launch()
+
+@spaces.GPU
+def launch():
+    demo.queue(default_concurrency_limit=2).launch(share=True)
+
+
+if __name__ == "__main__":
+    launch()