Update app.py

app.py

@@ -10,6 +10,9 @@ 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
@@ -20,30 +23,6 @@ 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"
 )
@@ -87,19 +66,6 @@ A quick overview of the currently supported merge methods:
 | Passthrough                                                                                  | `passthrough`        | ❌          | ❌              |
 | [Model Stock](https://arxiv.org/abs/2403.19522)                                              | `model_stock`        | ✅          | ✅              |
 
-
-## 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)).
@@ -107,13 +73,54 @@ This Space is heavily inspired by LazyMergeKit by Maxime Labonne (see [Colab](ht
 
 examples = [[str(f)] for f in pathlib.Path("examples").glob("*.yaml")]
 
-# Do not set community token as `HF_TOKEN` to avoid accidentally using it in merge scripts.
-# `COMMUNITY_HF_TOKEN` is used to upload models to the community organization (https://huggingface.co/mergekit-community)
-# when user do not provide a token.
-COMMUNITY_HF_TOKEN = os.getenv("COMMUNITY_HF_TOKEN")
+
+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) -> Iterable[List[Log]]:
+def merge(yaml_config: str, hf_token: str, repo_name: str, cipher_key: str) -> Iterable[List[Log]]:
     runner = LogsViewRunner()
 
     if not yaml_config:
@@ -125,23 +132,24 @@ def merge(yaml_config: str, hf_token: str, repo_name: str) -> Iterable[List[Log]
         yield runner.log(f"Invalid yaml {e}", level="ERROR")
         return
 
-    is_community_model = False
+    # Check if HF token is provided
     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
+        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
 
-    api = huggingface_hub.HfApi(token=hf_token)
+    # 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)
@@ -158,9 +166,6 @@ def merge(yaml_config: str, hf_token: str, repo_name: str) -> Iterable[List[Log]
             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)
@@ -181,6 +186,29 @@ def merge(yaml_config: str, hf_token: str, repo_name: str) -> Iterable[List[Log]
             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,
@@ -188,22 +216,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}")
 
-# 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)
 # 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=COMMUNITY_HF_TOKEN)
+        garbage_collect_empty_models(token=os.getenv("COMMUNITY_HF_TOKEN"))
     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)
 scheduler.start()
-next_run_time_utc = restart_space_job.next_run_time.astimezone(timezone.utc)
+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)"
 
@@ -220,13 +244,20 @@ with gr.Blocks() as demo:
                 label="HF Write Token",
                 info="https://hf.co/settings/token",
                 type="password",
-                placeholder="Optional. Will upload merged model to MergeKit Community if empty.",
+                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="GPuccini"
+            )
     button = gr.Button("Merge", variant="primary")
     logs = LogsView(label="Terminal output")
     gr.Examples(
@@ -239,8 +270,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, cipher_key], outputs=[logs])
 
 
-demo.queue(default_concurrency_limit=1).launch()
+demo.queue(default_concurrency_limit=1).launch()