Update app.py

app.py

@@ -1,102 +1,93 @@
 import gradio as gr
-from huggingface_hub import HfApi
-import spaces
+import os
 import shutil
 import logging
 import subprocess
 from pathlib import Path
+from merge_script import ModelMerger, get_max_vocab_size, download_json_files
+import spaces
 
 @spaces.GPU
-def write_repo(base_model, model_to_merge):
-    with open("repo.txt", "w") as repo:       
-        repo.write(base_model + "\n" + model_to_merge)
-
-def merge_and_upload(base_model, model_to_merge, scaling_factor, weight_drop_prob, repo_name, token):
-    # Define a fixed output path
-    outpath = Path('/tmp/output')
-    write_repo(base_model, model_to_merge)
-
-    # Construct the command to run hf_merge.py
-    command = [
-        "python3", "hf_merge.py",
-        "-p", str(weight_drop_prob),
-        "-lambda", str(scaling_factor),
-        "repo.txt", str(outpath)
-    ]
+def merge_models(base_model, model_to_merge, scaling_factor, weight_drop_prob, repo_name, token, commit_message):
+    # Define staging and output paths
+    staging_path = "/tmp/staging"
+    output_path = "/tmp/output"
+    os.makedirs(staging_path, exist_ok=True)
+    os.makedirs(output_path, exist_ok=True)
 
-    # Set up logging
-    logging.basicConfig(level=logging.INFO)
-    log_output = ""
+    # Initialize ModelMerger and prepare base model
+    model_merger = ModelMerger(staging_path, repo_name, token)
+    model_merger.prepare_base_model(base_model, os.path.join(staging_path, "base_model"))
 
-    # Run the command and capture the output
-    result = subprocess.run(command, capture_output=True, text=True)
-    
+    # Merge models and handle progress updates
+    for repo_name in [base_model, model_to_merge]:
+        model_merger.merge_repo(repo_name, os.path.join(staging_path, "staging_model"), weight_drop_prob, scaling_factor)
+        yield 0.25, f"Merged {repo_name}"
 
-    # Log the output
-    log_output += result.stdout + "\n"
-    log_output += result.stderr + "\n"
-    logging.info(result.stdout)
-    logging.error(result.stderr)
+    # Finalize merge and handle vocab size
+    model_merger.finalize_merge(output_path)
+    yield 0.5, "Finalizing merge and handling vocab size..."
+    max_vocab_size, repo_with_max_vocab = get_max_vocab_size([base_model, model_to_merge])
+    if max_vocab_size > 0:
+        download_json_files(repo_with_max_vocab, ['config.json', 'special_tokens_map.json', 'tokenizer.json', 'tokenizer_config.json'], output_path)
 
-    # Check if the merge was successful
-    if result.returncode != 0:
-        return None, f"Error in merging models: {result.stderr}", log_output
-
-    # Update progress bar
-    yield 0.5, "Merging completed. Uploading to Hugging Face Hub..."
-
-    # Upload the result to Hugging Face Hub
-    api = HfApi(token=token)
-    try:
-        # Get the username of the user who is logged in
-        user = api.whoami(token=token)["name"]
+    # Upload merged model to Hugging Face Hub
+    if repo_name:
+        model_merger.upload_model(output_path, repo_name, commit_message)
+        yield 0.75, "Uploading merged model to Hugging Face Hub..."
+        repo_url = f"https://huggingface.co/{repo_name}"
+        yield 1.0, f"Model merged and uploaded successfully! {repo_url}"
+    else:
+        yield 1.0, "Model merged successfully! No upload performed."
 
-        # Autofill the repo name if none is provided
-        if not repo_name:
-            repo_name = f"{user}/default-repo"
+def get_model_type_info(model_name):
+    model_types = {
+        "base_model": "Base model should be in .safetensors format.",
+        "model_to_merge": "Model to merge can be in .safetensors or .bin format."
+    }
+    return model_types.get(model_name, "No specific info available.")
 
-        # Create a new repo or update an existing one
-        api.create_repo(repo_id=repo_name, token=token, exist_ok=True)
+def validate_model_format(model_name, model_path):
+    if model_name == "base_model":
+        if not model_path.endswith(".safetensors"):
+            return False, "Base model must be in .safetensors format."
+    elif model_name == "model_to_merge":
+        if not model_path.endswith((".safetensors", ".bin")):
+            return False, "Model to merge must be in .safetensors or .bin format."
+    return True, None
 
-        # Upload the file
-        api.upload_folder(
-            folder_path=str(outpath),
-            repo_id=repo_name,
-            repo_type="model",
-            token=token
-        )
-        repo_url = f"https://huggingface.co/{repo_name}"
-        yield 1.0, "Upload completed."
-        return repo_url, "Model merged and uploaded successfully!", log_output
-    except Exception as e:
-        return None, f"Error uploading to Hugging Face Hub: {str(e)}", log_output
+def merge_and_upload_interface():
+    with gr.Blocks(theme="Ytheme/XRainbow", fill_width=True) as demo:
+        gr.Markdown("# Model Merger and Uploader")
+        gr.Markdown("Combine and upload models with real-time progress updates.")
+        gr.Markdown("**Model Compatibility:**")
+        gr.Markdown("Combine any two models using a Super Mario merge.")
+        gr.Markdown("Works with:")
+        gr.Markdown("* Stable Diffusion (1.5, XL/XL Turbo)")
+        gr.Markdown("* LLMs (Mistral, Llama, etc)")
+        gr.Markdown("* LoRas (must be same size)")
+        gr.Markdown("* Any two homologous models")
 
-# Define the Gradio interface
-with gr.Blocks(theme="Ytheme/XRainbow", fill_width=True, delete_cache=(60, 3600)) as demo:
-    gr.Markdown("# Model Merger and Uploader")
-    gr.Markdown("Combine any two models using a Super Mario merge(DARE) as described in the linked whitepaper.")
-    gr.Markdown("Works with:")
-    gr.Markdown("* Stable Diffusion (1.5, XL/XL Turbo)")
-    gr.Markdown("* LLMs (Mistral, Llama, etc)")
-    gr.Markdown("* LoRas (must be same size)")
-    gr.Markdown("* Any two homologous models")
+        with gr.Column():
+            token = gr.Textbox(label="HuggingFace Token")
+            base_model = gr.Textbox(label="Base Model")
+            base_model_info = gr.HTML(get_model_type_info("base_model"))
+            model_to_merge = gr.Textbox(label="Model to Merge")
+            model_to_merge_info = gr.HTML(get_model_type_info("model_to_merge"))
+            repo_name = gr.Textbox(label="New Model Name")
+            scaling_factor = gr.Slider(minimum=0, maximum=10, label="Scaling Factor")
+            weight_drop_prob = gr.Slider(minimum=0, maximum=1, label="Weight Drop Probability")
+            gr.Button("Merge and Upload").click(
+                merge_models,
+                inputs=[base_model, model_to_merge, scaling_factor, weight_drop_prob, repo_name, token, commit_message],
+                outputs=[
+                    gr.Textbox(label="Progress/Output"),
+                    gr.Textbox(label="Log")
+                ],
+                pre_process=validate_model_format
+            )
 
-    with gr.Column():
-        with gr.Row():
-            token = gr.Textbox(label="Your HF write token", placeholder="hf_...", value="", max_lines=1)
-        with gr.Row():
-            base_model = gr.Textbox(label="Base Model", placeholder=".safetensors")
-        with gr.Row():
-            model_to_merge = gr.Textbox(label="Merge Model", placeholder=".bin/.safetensors")
-        with gr.Row():
-            repo_name = gr.Textbox(label="New Model", placeholder="SDXL-", info="If empty, auto-complete", value="", max_lines=1)
-        with gr.Row():
-            scaling_factor = gr.Slider(minimum=0, maximum=10, value=3.0, label="Scaling Factor")
-        with gr.Row():
-            weight_drop_prob = gr.Slider(minimum=0, maximum=1, value=0.3, label="Weight Drop Probability")
-        gr.Button("Merge").click(
-            merge_and_upload,
-            inputs=[base_model, model_to_merge, scaling_factor, weight_drop_prob, repo_name, token]
-        )
+    demo.launch()
 
-demo.launch()
+if __name__ == "__main__":
+    merge_and_upload_interface()