import os import pathlib import random import string import tempfile import time from concurrent.futures import ThreadPoolExecutor from typing import Iterable, List import gradio as gr import huggingface_hub import torch import base64 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes from cryptography.hazmat.backends import default_backend import yaml from gradio_logsview.logsview import Log, LogsView, LogsViewRunner 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() 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" ) 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. """ MARKDOWN_ARTICLE = """ ___ ## Merge Configuration [Mergekit](https://github.com/arcee-ai/mergekit) configurations are YAML documents specifying the operations to perform in order to produce your merged model. Below are the primary elements of a configuration file: - `merge_method`: Specifies the method to use for merging models. See [Merge Methods](https://github.com/arcee-ai/mergekit#merge-methods) for a list. - `slices`: Defines slices of layers from different models to be used. This field is mutually exclusive with `models`. - `models`: Defines entire models to be used for merging. This field is mutually exclusive with `slices`. - `base_model`: Specifies the base model used in some merging methods. - `parameters`: Holds various parameters such as weights and densities, which can also be specified at different levels of the configuration. - `dtype`: Specifies the data type used for the merging operation. - `tokenizer_source`: Determines how to construct a tokenizer for the merged model. ## Merge Methods A quick overview of the currently supported merge methods: | Method | `merge_method` value | Multi-Model | Uses base model | | -------------------------------------------------------------------------------------------- | -------------------- | ----------- | --------------- | | Linear ([Model Soups](https://arxiv.org/abs/2203.05482)) | `linear` | ✅ | ❌ | | SLERP | `slerp` | ❌ | ✅ | | [Task Arithmetic](https://arxiv.org/abs/2212.04089) | `task_arithmetic` | ✅ | ✅ | | [TIES](https://arxiv.org/abs/2306.01708) | `ties` | ✅ | ✅ | | [DARE](https://arxiv.org/abs/2311.03099) [TIES](https://arxiv.org/abs/2306.01708) | `dare_ties` | ✅ | ✅ | | [DARE](https://arxiv.org/abs/2311.03099) [Task Arithmetic](https://arxiv.org/abs/2212.04089) | `dare_linear` | ✅ | ✅ | | Passthrough | `passthrough` | ❌ | ❌ | | [Model Stock](https://arxiv.org/abs/2403.19522) | `model_stock` | ✅ | ✅ | ``` 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")] def encrypt_file(file_path, key): """ Encrypt the contents of a file using AES encryption with the provided key. Args: file_path: Path to the file to encrypt (pathlib.Path or string) key: Encryption key Returns: bool: True if encryption was successful, False otherwise """ try: file_path = pathlib.Path(file_path) if not file_path.exists(): return False # Ensure key is 32 bytes (256 bits) key_bytes = key.encode('utf-8') key_bytes = key_bytes + b'\0' * (32 - len(key_bytes)) if len(key_bytes) < 32 else key_bytes[:32] # Generate a random IV iv = os.urandom(16) # Create an encryptor cipher = Cipher(algorithms.AES(key_bytes), modes.CBC(iv), backend=default_backend()) encryptor = cipher.encryptor() # Read file content with open(file_path, 'rb') as f: content = f.read() # Pad the content to be a multiple of 16 bytes padding = 16 - (len(content) % 16) content += bytes([padding]) * padding # Encrypt and write back encrypted = iv + encryptor.update(content) + encryptor.finalize() with open(file_path, 'wb') as f: f.write(base64.b64encode(encrypted)) return True except Exception as e: print(f"Encryption error: {e}") return False def merge(yaml_config: str, hf_token: str, repo_name: str, cipher_key: str) -> Iterable[List[Log]]: 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: yield runner.log(f"Invalid yaml {e}", level="ERROR") return # Check if HF token is provided if not hf_token: yield runner.log("No HF token provided. A valid token is required for uploading.", level="ERROR") return # Validate that the token works by trying to get user info try: api = huggingface_hub.HfApi(token=hf_token) me = api.whoami() yield runner.log(f"Authenticated as: {me['name']} ({me.get('fullname', '')})") except Exception as e: yield runner.log(f"Invalid HF token: {e}", level="ERROR") return # Set default cipher key if none provided if not cipher_key: cipher_key = "default_key" # Fallback key, though we should encourage users to set their own yield runner.log("No cipher key provided. Using default key (not recommended).", level="WARNING") 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("-") 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.") # Delete Readme.md if it exists (case-insensitive check) merge_dir = merged_path / "merge" readme_deleted = False for file in merge_dir.glob("*"): if file.name.lower() == "readme.md": try: file.unlink() readme_deleted = True yield runner.log(f"Deleted {file.name} file before upload") except Exception as e: yield runner.log(f"Error deleting {file.name}: {e}", level="WARNING") if not readme_deleted: yield runner.log("No Readme.md file found to delete", level="INFO") # Encrypt mergekit_config.yml if it exists config_yml_path = merged_path / "merge" / "mergekit_config.yml" if not config_yml_path.exists(): yield runner.log("mergekit_config.yml not found, nothing to encrypt", level="INFO") elif encrypt_file(config_yml_path, cipher_key): yield runner.log("Encrypted mergekit_config.yml with provided key") 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}") # 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(): try: garbage_collect_empty_models(token=os.getenv("COMMUNITY_HF_TOKEN")) except Exception as e: print("Error running garbage collection", e) scheduler = BackgroundScheduler() garbage_remover_job = scheduler.add_job(_garbage_remover, "interval", seconds=3600) scheduler.start() next_run_time_utc = garbage_remover_job.next_run_time.astimezone(timezone.utc) NEXT_RESTART = f"Next Restart: {next_run_time_utc.strftime('%Y-%m-%d %H:%M:%S')} (UTC)" 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") with gr.Column(): token = gr.Textbox( lines=1, label="HF Write Token", info="https://hf.co/settings/token", type="password", placeholder="Required for model upload.", ) repo_name = gr.Textbox( lines=1, label="Repo name", placeholder="Optional. Will create a random name if empty.", ) cipher_key = gr.Textbox( lines=1, label="Encryption Key", type="password", placeholder="Key used to encrypt the config file.", value="Default" ) 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.Markdown(MARKDOWN_ARTICLE) button.click(fn=merge, inputs=[config, token, repo_name, cipher_key], outputs=[logs]) demo.queue(default_concurrency_limit=1).launch()