app.py

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.primitives import padding as sym_padding # Use symmetric padding
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. Includes encryption for the `mergekit_config.yml` and a tool to decrypt it.
"""

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 _prepare_key(key: str) -> bytes:
    """Pads or truncates the key to 32 bytes (256 bits) for AES."""
    key_bytes = key.encode('utf-8')
    if len(key_bytes) < 32:
        return key_bytes + b'\0' * (32 - len(key_bytes))
    else:
        return key_bytes[:32]

def encrypt_file(file_path, key: str) -> bool:
    """
    Encrypt the contents of a file using AES-256-CBC encryption with the provided key.
    The output is Base64 encoded.

    Args:
        file_path: Path to the file to encrypt (pathlib.Path or string)
        key: Encryption key string.

    Returns:
        bool: True if encryption was successful, False otherwise
    """
    try:
        file_path = pathlib.Path(file_path)
        if not file_path.exists():
            print(f"Encryption error: File not found at {file_path}")
            return False

        key_bytes = _prepare_key(key)

        # Generate a random IV (Initialization Vector) - 16 bytes for AES
        iv = os.urandom(16)

        # Create an AES cipher instance with CBC mode
        cipher = Cipher(algorithms.AES(key_bytes), modes.CBC(iv), backend=default_backend())
        encryptor = cipher.encryptor()

        # Use PKCS7 padding
        padder = sym_padding.PKCS7(algorithms.AES.block_size).padder()

        with open(file_path, 'rb') as f:
            plaintext = f.read()

        # Pad the data
        padded_data = padder.update(plaintext) + padder.finalize()

        # Encrypt the padded data
        ciphertext = encryptor.update(padded_data) + encryptor.finalize()

        # Prepend the IV to the ciphertext and base64 encode the result
        encrypted_data_with_iv = base64.b64encode(iv + ciphertext)

        # Write the base64 encoded encrypted data back to the file
        with open(file_path, 'wb') as f:
            f.write(encrypted_data_with_iv)

        return True
    except Exception as e:
        print(f"Encryption error: {e}")
        return False

def decrypt_file_content(file_input, key: str) -> str:
    """
    Decrypts the content of an uploaded file using AES-256-CBC and returns the result.
    Assumes the file content is Base64 encoded IV + ciphertext.

    Args:
        file_input: Gradio File component output (temporary file object).
        key: Decryption key string.

    Returns:
        str: Decrypted content as a UTF-8 string, or an error message.
    """
    if file_input is None:
        return "Error: No file provided for decryption."
    if not key:
        return "Error: Decryption key cannot be empty."

    try:
        file_path = file_input.name # Get the temporary file path from Gradio
        key_bytes = _prepare_key(key)

        with open(file_path, 'rb') as f:
            base64_encoded_data = f.read()

        # Decode from Base64
        encrypted_data_with_iv = base64.b64decode(base64_encoded_data)

        # Extract the IV (first 16 bytes)
        iv = encrypted_data_with_iv[:16]
        # Extract the ciphertext (the rest)
        ciphertext = encrypted_data_with_iv[16:]

        # Create an AES cipher instance with CBC mode for decryption
        cipher = Cipher(algorithms.AES(key_bytes), modes.CBC(iv), backend=default_backend())
        decryptor = cipher.decryptor()

        # Decrypt the data
        padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize()

        # Unpad the data using PKCS7
        unpadder = sym_padding.PKCS7(algorithms.AES.block_size).unpadder()
        plaintext = unpadder.update(padded_plaintext) + unpadder.finalize()

        # Decode the plaintext from bytes to string (assuming UTF-8)
        return plaintext.decode('utf-8')

    except (ValueError, TypeError) as e:
         # Catches Base64 decoding errors, incorrect key type errors
        return f"Decryption Error: Invalid input data or key format. ({e})"
    except Exception as e:
        # Catches padding errors (often due to wrong key), or other crypto issues
        print(f"Decryption error details: {e}")
        return f"Decryption Failed: Likely incorrect key or corrupted file. Error: {type(e).__name__}"


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

    # Use default key if none provided, but log a warning
    if not cipher_key:
        cipher_key = "default_insecure_key" # Make default explicitely insecure sounding
        yield runner.log("No cipher key provided. Using a default, insecure key. Please provide your own key for security.", level="WARNING")
    elif cipher_key == "Default": # Check against the placeholder value
         cipher_key = "default_insecure_key" # Treat placeholder as no key provided
         yield runner.log("Default placeholder key detected. Using an insecure key. Please provide your own key.", 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")
            try:
                api.delete_repo(repo_url.repo_id)
                yield runner.log(f"Repo {repo_url.repo_id} deleted.")
            except Exception as delete_e:
                 yield runner.log(f"Failed to delete repo {repo_url.repo_id}: {delete_e}", level="WARNING")
            return

        yield runner.log("Model merged successfully. Preparing for upload.")

        # ---- Encryption Step ----
        merge_dir = merged_path / "merge"
        config_yml_path = merge_dir / "mergekit_config.yml"

        if config_yml_path.exists():
            yield runner.log(f"Found {config_yml_path.name}. Encrypting...")
            if encrypt_file(config_yml_path, cipher_key):
                yield runner.log(f"Successfully encrypted {config_yml_path.name} with provided key.")
            else:
                yield runner.log(f"Failed to encrypt {config_yml_path.name}. Uploading unencrypted.", level="ERROR")
        else:
            yield runner.log(f"{config_yml_path.name} not found in merge output, nothing to encrypt.", level="INFO")
        # ---- End Encryption Step ----

        # Delete Readme.md if it exists (case-insensitive check) before upload
        readme_deleted = False
        try:
            for file in merge_dir.glob("*"):
                if file.name.lower() == "readme.md":
                    file.unlink()
                    readme_deleted = True
                    yield runner.log(f"Deleted {file.name} file before upload")
                    break # Assume only one readme
        except Exception as e:
            yield runner.log(f"Error deleting Readme.md: {e}", level="WARNING")

        if not readme_deleted:
            yield runner.log("No Readme.md file found to delete.", level="INFO")

        yield runner.log("Uploading merged model files to HF.")
        yield from runner.run_python(
            api.upload_folder,
            repo_id=repo_url.repo_id,
            folder_path=merge_dir, # Upload from the 'merge' subdirectory
        )
        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.Tabs():
        with gr.TabItem("Merge Model"):
            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",
                        info="Key used to encrypt the generated mergekit_config.yml file before upload. Leave blank or 'Default' for no encryption (or insecure default).",
                        placeholder="Enter your secret key here",
                        value="Default" # Set a default placeholder
                    )
            button = gr.Button("Merge and Upload", variant="primary")
            logs = LogsView(label="Merge Progress / Terminal output")
            gr.Examples(
                examples,
                fn=lambda s: (s,),
                run_on_click=True,
                label="Merge Examples",
                inputs=[filename],
                outputs=[config],
            )
            gr.Markdown(MARKDOWN_ARTICLE)

            button.click(fn=merge, inputs=[config, token, repo_name, cipher_key], outputs=[logs])

        with gr.TabItem("Decrypt Configuration"):
            gr.Markdown("Upload an encrypted `mergekit_config.yml` file and provide the key to decrypt it.")
            with gr.Row():
                decrypt_file_input = gr.File(label="Upload Encrypted mergekit_config.yml")
                decrypt_key_input = gr.Textbox(
                    lines=1,
                    label="Decryption Key",
                    type="password",
                    placeholder="Enter the key used for encryption",
                )
            decrypt_button = gr.Button("Decrypt File", variant="secondary")
            decrypted_output = gr.Code(language="yaml", label="Decrypted Configuration", lines=15, interactive=False)

            decrypt_button.click(
                fn=decrypt_file_content,
                inputs=[decrypt_file_input, decrypt_key_input],
                outputs=[decrypted_output]
            )


demo.queue(default_concurrency_limit=1).launch()