Update app.py

app.py

@@ -16,7 +16,13 @@ from safetensors.torch import load_file, save_file
 from tqdm import tqdm
 
 # --- Import Helpers ---
-from merge_utils import execute_mergekit_config, build_full_merge_config, build_moe_config
+from merge_utils import (
+    execute_mergekit_config, 
+    execute_raw_pytorch,
+    build_full_merge_config, 
+    build_moe_config,
+    build_raw_config
+)
 from dare_utils import task_dare_custom
 
 # --- Memory Efficient Safetensors ---
@@ -601,16 +607,17 @@ def task_full_mergekit_merge(hf_token, models_text, method, dtype, base_model, w
 # TAB 6: MOE CREATION
 # =================================================================================
 
-def task_moe_create(hf_token, base_model, experts_text, gate_mode, dtype, tok_source, shard_size, out_repo, private):
+def task_moe_create(hf_token, base_model, experts_text, prompts_text, gate_mode, dtype, tok_source, shard_size, out_repo, private):
     cleanup_temp()
     if not hf_token or not out_repo: return "Error: Token and Output Repo required."
     login(hf_token.strip())
     
     experts = [e.strip() for e in experts_text.split('\n') if e.strip()]
+    prompts = [p.strip() for p in prompts_text.split('\n') if p.strip()]
     
     # 1. Build Config
     config = build_moe_config(
-        base_model=base_model, experts=experts, gate_mode=gate_mode,
+        base_model=base_model, experts=experts, prompts=prompts, gate_mode=gate_mode,
         dtype=dtype, tokenizer_source=tok_source
     )
     
@@ -624,6 +631,23 @@ def task_moe_create(hf_token, base_model, experts_text, gate_mode, dtype, tok_so
     except Exception as e:
         return f"MoE Error: {e}"
 
+# --- TAB 8: Raw PyTorch (New) ---
+def task_raw_pytorch(hf_token, models_text, method, dtype, base_model, weights, shard_size, out_repo, private):
+    cleanup_temp()
+    if not hf_token or not out_repo: return "Error: Token and Output Repo required."
+    login(hf_token.strip())
+
+    models = [m.strip() for m in models_text.split('\n') if m.strip()]
+    config = build_raw_config(method, models, base_model, dtype, weights)
+
+    out_path = TempDir / "raw_merged"
+    try:
+        execute_raw_pytorch(config, str(out_path), shard_size)
+        api.create_repo(repo_id=out_repo, private=private, exist_ok=True, token=hf_token)
+        api.upload_folder(folder_path=str(out_path), repo_id=out_repo, token=hf_token)
+        return f"Success! Raw merge uploaded to {out_repo}"
+    except Exception as e: return f"Raw Merge Error: {e}"
+
 # =================================================================================
 # UI
 # =================================================================================
@@ -760,6 +784,22 @@ with gr.Blocks() as demo:
             t7_res = gr.Textbox(label="Result")
             
             t7_btn.click(task_dare_custom, [t7_token, t7_base, t7_ft, t7_ratio, t7_mask, t7_out, t7_priv], t7_res)
+            
+        with gr.Tab("Raw PyTorch Merge"):
+            gr.Markdown("### 🧠 Raw Weight Merging (Non-Transformers)")
+            t8_token = gr.Textbox(label="HF Token", type="password")
+            t8_method = gr.Dropdown(["Linear", "TIES", "Task_Arithmetic"], value="Linear", label="Method")
+            t8_models = gr.TextArea(label="Models (Path/Repo)")
+            with gr.Row():
+                t8_base = gr.Textbox(label="Base Model (Optional)")
+                t8_dtype = gr.Dropdown(["float32", "float16", "bfloat16"], value="float32", label="Dtype")
+            t8_weights = gr.Textbox(label="Weights")
+            t8_shard = gr.Slider(0.5, 10, 2.0, label="Shard Size (GB)")
+            t8_out = gr.Textbox(label="Output Repo")
+            t8_priv = gr.Checkbox(label="Private", value=True)
+            t8_btn = gr.Button("Merge Raw Weights")
+            t8_res = gr.Textbox(label="Result")
+            t8_btn.click(task_raw_pytorch, [t8_token, t8_models, t8_method, t8_dtype, t8_base, t8_weights, t8_shard, t8_out, t8_priv], t8_res)
 
 if __name__ == "__main__":
     demo.queue().launch(css=css, ssr_mode=False)