Switch to Zero GPU

README.md

@@ -4,7 +4,7 @@ emoji: 🔀
 colorFrom: yellow
 colorTo: yellow
 sdk: gradio
-sdk_version: 4.24.0
+sdk_version: 4.44.1
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -3,10 +3,9 @@ import pathlib
 import random
 import string
 import tempfile
-import time
-from concurrent.futures import ThreadPoolExecutor
 from typing import Iterable, List
 
+import spaces
 import gradio as gr
 import huggingface_hub
 import torch
@@ -16,44 +15,14 @@ from mergekit.config import MergeConfiguration
 
 from clean_community_org import garbage_collect_empty_models
 from apscheduler.schedulers.background import BackgroundScheduler
-from datetime import datetime, timezone
-
-has_gpu = torch.cuda.is_available()
-
-# Running directly from Python doesn't work well with Gradio+run_process because of:
-# Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method
-# Let's use the CLI instead.
-#
-# import mergekit.merge
-# from mergekit.common import parse_kmb
-# from mergekit.options import MergeOptions
-#
-# merge_options = (
-#     MergeOptions(
-#         copy_tokenizer=True,
-#         cuda=True,
-#         low_cpu_memory=True,
-#         write_model_card=True,
-#     )
-#     if has_gpu
-#     else MergeOptions(
-#         allow_crimes=True,
-#         out_shard_size=parse_kmb("1B"),
-#         lazy_unpickle=True,
-#         write_model_card=True,
-#     )
-# )
-
-cli = "mergekit-yaml config.yaml merge --copy-tokenizer" + (
-    " --cuda --low-cpu-memory --allow-crimes" if has_gpu else " --allow-crimes --out-shard-size 1B --lazy-unpickle"
-)
+from datetime import timezone
 
 MARKDOWN_DESCRIPTION = """
 # mergekit-gui
 
 The fastest way to perform a model merge 🔥
 
-Specify a YAML configuration file (see examples below) and a HF token and this app will perform the merge and upload the merged model to your user profile.
+Specify a YAML configuration file (see examples below) and a HF token and this app will perform the merge and upload the merged model to your user profile. Uses Zero GPU quota to perform the merge.
 """
 
 MARKDOWN_ARTICLE = """
@@ -113,11 +82,56 @@ 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) -> Iterable[List[Log]]:
+def run_merge_cpu(runner: LogsViewRunner, cli: str, merged_path: str, tmpdirname: str):
+    # Set tmp HF_HOME to avoid filling up disk Space
+    tmp_env = os.environ.copy()
+    tmp_env["HF_HOME"] = f"{tmpdirname}/.cache"
+    full_cli = cli + f" --lora-merge-cache {tmpdirname}/.lora_cache --transformers-cache {tmpdirname}/.cache"
+    yield from runner.run_command(full_cli.split(), cwd=merged_path, env=tmp_env)
+    yield ("done", runner.exit_code)
+
+
+@spaces.GPU(duration=60 * 5)
+def run_merge_gpu(runner: LogsViewRunner, cli: str, merged_path: str, tmpdirname: str):
+    yield from run_merge_cpu(
+        runner,
+        cli,
+        merged_path,
+        tmpdirname,
+    )
+
+
+def run_merge(
+    runner: LogsViewRunner,
+    cli: str,
+    merged_path: str,
+    tmpdirname: str,
+    use_gpu: bool,
+):
+    if use_gpu:
+        yield from run_merge_gpu(runner, cli, merged_path, tmpdirname)
+    else:
+        yield from run_merge_cpu(runner, cli, merged_path, tmpdirname)
+
+
+def prefetch_models(
+    runner: LogsViewRunner,
+    merge_config: MergeConfiguration,
+    hf_home: str,
+    lora_merge_cache: str,
+):
+    for model in merge_config.referenced_models():
+        yield runner.log(f"Downloading {model}...")
+        model = model.merged(cache_dir=lora_merge_cache, trust_remote_code=False)
+        local_path = model.local_path(cache_dir=hf_home)
+        yield runner.log(f"\tDownloaded {model} to {local_path}")
+
+
+def merge(yaml_config: str, hf_token: str, repo_name: str, private: bool) -> Iterable[List[Log]]:
     runner = LogsViewRunner()
 
     if not yaml_config:
-        yield runner.log("Empty yaml, pick an example below", level="ERROR")
+        yield runner.log("Empty yaml, enter your config or pick an example below", level="ERROR")
         return
     try:
         merge_config = MergeConfiguration.model_validate(yaml.safe_load(yaml_config))
@@ -127,6 +141,12 @@ def merge(yaml_config: str, hf_token: str, repo_name: str) -> Iterable[List[Log]
 
     is_community_model = False
     if not hf_token:
+        if private:
+            yield runner.log(
+                "Cannot upload model as private without a token. Please provide a HF token.",
+                level="ERROR",
+            )
+            return
         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.",
@@ -142,6 +162,10 @@ def merge(yaml_config: str, hf_token: str, repo_name: str) -> Iterable[List[Log]
         hf_token = COMMUNITY_HF_TOKEN
 
     api = huggingface_hub.HfApi(token=hf_token)
+    has_gpu = torch.cuda.is_available()
+    cli = "mergekit-yaml config.yaml merge --copy-tokenizer --allow-crimes -v" + (
+        " --cuda --low-cpu-memory --read-to-gpu" if (has_gpu) else " --out-shard-size 1B --lazy-unpickle"
+    )
 
     with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname:
         tmpdir = pathlib.Path(tmpdirname)
@@ -163,19 +187,42 @@ def merge(yaml_config: str, hf_token: str, repo_name: str) -> Iterable[List[Log]
 
         try:
             yield runner.log(f"Creating repo {repo_name}")
-            repo_url = api.create_repo(repo_name, exist_ok=True)
+            repo_url = api.create_repo(repo_name, exist_ok=True, private=private, repo_type="model")
             yield runner.log(f"Repo created: {repo_url}")
         except Exception as e:
             yield runner.log(f"Error creating repo {e}", level="ERROR")
             return
 
-        # Set tmp HF_HOME to avoid filling up disk Space
-        tmp_env = os.environ.copy()  # taken from https://stackoverflow.com/a/4453495
-        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)
+        # Prefetch models to avoid downloading them with scarce GPU time
+        yield runner.log("Prefetching models...")
+        yield from prefetch_models(
+            runner,
+            merge_config,
+            hf_home=f"{tmpdirname}/.cache",
+            lora_merge_cache=f"{tmpdirname}/.lora_cache",
+        )
+        yield runner.log("Models prefetched. Starting merge.")
+        exit_code = None
+        try:
+            for ev in run_merge(
+                runner,
+                cli,
+                merged_path,
+                tmpdirname,
+                use_gpu=has_gpu,
+            ):
+                if isinstance(ev, tuple) and ev[0] == "done":
+                    exit_code = ev[1]
+                    continue
+                yield ev
+        except Exception as e:
+            yield runner.log(f"Error running merge {e}", level="ERROR")
+            yield runner.log("Merge failed. Deleting repo as no model is uploaded.", level="ERROR")
+            api.delete_repo(repo_url.repo_id)
+            return
 
-        if runner.exit_code != 0:
+        if exit_code != 0:
+            yield runner.log(f"Exit code: {exit_code}")
             yield runner.log("Merge failed. Deleting repo as no model is uploaded.", level="ERROR")
             api.delete_repo(repo_url.repo_id)
             return
@@ -188,9 +235,18 @@ def merge(yaml_config: str, hf_token: str, repo_name: str) -> Iterable[List[Log]
         )
         yield runner.log(f"Model successfully uploaded to HF: {repo_url.repo_id}")
 
+
+merge.zerogpu = True
+run_merge.zerogpu = True
+
+
 # This is workaround. As the space always getting stuck.
 def _restart_space():
-    huggingface_hub.HfApi().restart_space(repo_id="arcee-ai/mergekit-gui", token=COMMUNITY_HF_TOKEN, factory_reboot=False)
+    huggingface_hub.HfApi().restart_space(
+        repo_id="arcee-ai/mergekit-gui", token=COMMUNITY_HF_TOKEN, factory_reboot=False
+    )
+
+
 # Run garbage collection every hour to keep the community org clean.
 # Empty models might exists if the merge fails abruptly (e.g. if user leaves the Space).
 def _garbage_remover():
@@ -199,6 +255,7 @@ def _garbage_remover():
     except Exception as e:
         print("Error running garbage collection", e)
 
+
 scheduler = BackgroundScheduler()
 restart_space_job = scheduler.add_job(_restart_space, "interval", seconds=21600)
 garbage_remover_job = scheduler.add_job(_garbage_remover, "interval", seconds=3600)
@@ -210,7 +267,7 @@ NEXT_RESTART = f"Next Restart: {next_run_time_utc.strftime('%Y-%m-%d %H:%M:%S')}
 with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN_DESCRIPTION)
     gr.Markdown(NEXT_RESTART)
-    
+
     with gr.Row():
         filename = gr.Textbox(visible=False, label="filename")
         config = gr.Code(language="yaml", lines=10, label="config.yaml")
@@ -227,6 +284,11 @@ with gr.Blocks() as demo:
                 label="Repo name",
                 placeholder="Optional. Will create a random name if empty.",
             )
+            private = gr.Checkbox(
+                label="Private",
+                value=False,
+                info="Upload the model as private. If not checked, will be public. Must provide a token.",
+            )
     button = gr.Button("Merge", variant="primary")
     logs = LogsView(label="Terminal output")
     gr.Examples(
@@ -239,8 +301,7 @@ with gr.Blocks() as demo:
     )
     gr.Markdown(MARKDOWN_ARTICLE)
 
-    button.click(fn=merge, inputs=[config, token, repo_name], outputs=[logs])
-
+    button.click(fn=merge, inputs=[config, token, repo_name, private], outputs=[logs])
 
 
 demo.queue(default_concurrency_limit=1).launch()

requirements.txt

@@ -1,5 +1,5 @@
 apscheduler
-torch
+torch==2.5.1
 git+https://github.com/arcee-ai/mergekit.git
 # see https://huggingface.co/spaces/Wauplin/gradio_logsview
 gradio_logsview@https://huggingface.co/spaces/Wauplin/gradio_logsview/resolve/main/gradio_logsview-0.0.5-py3-none-any.whl