Upload 2 files

README.md

@@ -1,3 +1,4 @@
+---
 title: mergekit-gui
 emoji: 🔀
 colorFrom: yellow
@@ -7,5 +8,6 @@ sdk_version: 4.44.1
 app_file: app.py
 pinned: false
 license: apache-2.0
------
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -17,34 +17,12 @@ from clean_community_org import garbage_collect_empty_models
 from apscheduler.schedulers.background import BackgroundScheduler
 from datetime import timezone
 
-# HF Spaces: needed for ZeroGPU. Safe on CPU Spaces too.
-try:
-    import spaces  # type: ignore
-except Exception:
-    spaces = None
-import yaml
-from gradio_logsview.logsview import Log, LogsView, LogsViewRunner
-from mergekit.config import MergeConfiguration
-
-from clean_community_org import garbage_collect_empty_models
-
-def _build_cli(use_cuda: bool) -> str:
-    """Build mergekit CLI string.
-
-    NOTE: On HF ZeroGPU, torch.cuda.is_available() is only True inside a @spaces.GPU call.
-    """
-    base = "mergekit-yaml config.yaml merge --copy-tokenizer --allow-crimes"
-    if use_cuda:
-        return base + " --cuda --low-cpu-memory --write-model-card"
-    return base + " --out-shard-size 5B --lazy-unpickle --trust-remote-code"
-
-
 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 = """
@@ -104,86 +82,6 @@ examples = [[str(f)] for f in pathlib.Path("examples").glob("*.yaml")]
 COMMUNITY_HF_TOKEN = os.getenv("COMMUNITY_HF_TOKEN")
 
 
-def _merge_impl(yaml_config: str, hf_token: str, repo_name: str, *, force_cuda: bool = False) -> Iterable[List[Log]]:
-    runner = LogsViewRunner()
-
-    # Decide CUDA at call-time. On ZeroGPU, CUDA is only available inside a @spaces.GPU function.
-    use_cuda = bool(force_cuda) or torch.cuda.is_available()
-    cli = _build_cli(use_cuda)
-
-    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:
-        yield 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",
-            )
-            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
-        if not COMMUNITY_HF_TOKEN:
-            raise gr.Error("Cannot upload to community org: community token not set by Space owner.")
-        hf_token = COMMUNITY_HF_TOKEN
-
-    api = huggingface_hub.HfApi(token=hf_token)
-
-    with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname:
-        tmpdir = pathlib.Path(tmpdirname)
-        merged_path = tmpdir / "merged"
-        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}")
-
-        if not repo_name:
-            yield runner.log("No repo name provided. Generating a random one.")
-            repo_name = f"mergekit-{merge_config.merge_method}"
-            # Make repo_name "unique" (no need to be extra careful on uniqueness)
-            repo_name += "-" + "".join(random.choices(string.ascii_lowercase, k=7))
-            repo_name = repo_name.replace("/", "-").strip("-")
-
-        if is_community_model and not repo_name.startswith("mergekit-community/"):
-            repo_name = f"mergekit-community/{repo_name}"
-
-        try:
-            yield 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}")
-        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)
-
-        if runner.exit_code != 0:
-            yield 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}")
-
-
 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()
@@ -193,7 +91,7 @@ def run_merge_cpu(runner: LogsViewRunner, cli: str, merged_path: str, tmpdirname
     yield ("done", runner.exit_code)
 
 
-@spaces.GPU(duration=60 * 2)
+@spaces.GPU(duration=60 * 5)
 def run_merge_gpu(runner: LogsViewRunner, cli: str, merged_path: str, tmpdirname: str):
     yield from run_merge_cpu(
         runner,
@@ -341,6 +239,8 @@ def merge(yaml_config: str, hf_token: str, repo_name: str, private: bool) -> Ite
 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
@@ -404,4 +304,4 @@ with gr.Blocks() as demo:
     button.click(fn=merge, inputs=[config, token, repo_name, private], outputs=[logs])
 
 
-demo.queue(default_concurrency_limit=1).launch()
+demo.queue(default_concurrency_limit=1).launch()