import sys
import warnings
import os

warnings.filterwarnings("ignore", category=UserWarning, module="pydantic")

def _monkey_patch_for_mergekit():
    
    import pydantic
    import pydantic_core
    import torch
    
    original_validate = pydantic_core.core_schema.is_instance_schema
    
    def patched_is_instance_schema(cls, *args, **kwargs):
        if cls is torch.Tensor:
            # Return a simple any_schema for torch.Tensor
            return pydantic_core.core_schema.any_schema()
        return original_validate(cls, *args, **kwargs)
    
    pydantic_core.core_schema.is_instance_schema = patched_is_instance_schema
    
    original_create_model = pydantic.create_model
    
    def patched_create_model(__model_name, __config__=None, **kwargs):
        if __config__ is None:
            __config__ = pydantic.ConfigDict(arbitrary_types_allowed=True)
        elif isinstance(__config__, dict):
            __config__["arbitrary_types_allowed"] = True
        elif hasattr(__config__, "arbitrary_types_allowed"):
            __config__.arbitrary_types_allowed = True
        
        return original_create_model(__model_name, __config__=__config__, **kwargs)
    
    pydantic.create_model = patched_create_model
    
    return True

if _monkey_patch_for_mergekit():
    print("Applied monkey patches for MergeKit compatibility")

import gradio as gr
import torch
import gc
import shutil
import requests
import json
import struct
import numpy as np
import yaml
import subprocess
from pathlib import Path
from typing import Dict, Any, Optional, List
from huggingface_hub import HfApi, hf_hub_download, list_repo_files, login
from safetensors.torch import load_file, save_file
from tqdm import tqdm


# --- Memory Efficient Safetensors ---
class MemoryEfficientSafeOpen:
    def __init__(self, filename):
        self.filename = filename
        self.file = open(filename, "rb")
        self.header, self.header_size = self._read_header()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.file.close()

    def keys(self) -> list[str]:
        return [k for k in self.header.keys() if k != "__metadata__"]

    def metadata(self) -> Dict[str, str]:
        return self.header.get("__metadata__", {})

    def get_tensor(self, key):
        if key not in self.header:
            raise KeyError(f"Tensor '{key}' not found")
        metadata = self.header[key]
        offset_start, offset_end = metadata["data_offsets"]
        self.file.seek(self.header_size + 8 + offset_start)
        tensor_bytes = self.file.read(offset_end - offset_start)
        return self._deserialize_tensor(tensor_bytes, metadata)

    def _read_header(self):
        header_size = struct.unpack("<Q", self.file.read(8))[0]
        header_json = self.file.read(header_size).decode("utf-8")
        return json.loads(header_json), header_size

    def _deserialize_tensor(self, tensor_bytes, metadata):
        dtype_map = {
            "F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16,
            "I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
            "U8": torch.uint8, "BOOL": torch.bool
        }
        dtype = dtype_map[metadata["dtype"]]
        shape = metadata["shape"]
        return torch.frombuffer(tensor_bytes, dtype=torch.uint8).view(dtype).reshape(shape)

# --- Constants & Setup ---
try:
    TempDir = Path("/tmp/temp_tool")
    os.makedirs(TempDir, exist_ok=True)
except:
    TempDir = Path("./temp_tool")
    os.makedirs(TempDir, exist_ok=True)

api = HfApi()

def cleanup_temp():
    if TempDir.exists():
        shutil.rmtree(TempDir)
    os.makedirs(TempDir, exist_ok=True)
    gc.collect()

def get_key_stem(key):
    key = key.replace(".weight", "").replace(".bias", "")
    key = key.replace(".lora_down", "").replace(".lora_up", "")
    key = key.replace(".lora_A", "").replace(".lora_B", "")
    key = key.replace(".alpha", "")
    prefixes = [
        "model.diffusion_model.", "diffusion_model.", "model.", 
        "transformer.", "text_encoder.", "lora_unet_", "lora_te_", "base_model.model."
    ]
    changed = True
    while changed:
        changed = False
        for p in prefixes:
            if key.startswith(p):
                key = key[len(p):]
                changed = True
    return key

# =================================================================================
# TAB 1-4: ORIGINAL UNCHANGED LOGIC 
# =================================================================================

def parse_hf_url(url):
    if "huggingface.co" in url and "resolve" in url:
        try:
            parts = url.split("huggingface.co/")[-1].split("/")
            repo_id = f"{parts[0]}/{parts[1]}"
            filename = "/".join(parts[4:]).split("?")[0]
            return repo_id, filename
        except:
            return None, None
    return None, None

def download_lora_smart(input_str, token):
    local_path = TempDir / "adapter.safetensors"
    if local_path.exists(): os.remove(local_path)
    
    repo_id, filename = parse_hf_url(input_str)
    if repo_id and filename:
        try:
            hf_hub_download(repo_id=repo_id, filename=filename, token=token, local_dir=TempDir)
            found = list(TempDir.rglob(filename.split("/")[-1]))[0]
            if found != local_path: shutil.move(found, local_path)
            return local_path
        except: pass

    try:
        if ".safetensors" in input_str and input_str.count("/") >= 2:
            parts = input_str.split("/")
            repo_id = f"{parts[0]}/{parts[1]}"
            filename = "/".join(parts[2:])
            hf_hub_download(repo_id=repo_id, filename=filename, token=token, local_dir=TempDir)
            found = list(TempDir.rglob(filename.split("/")[-1]))[0]
            if found != local_path: shutil.move(found, local_path)
            return local_path
            
        candidates = ["adapter_model.safetensors", "model.safetensors"]
        files = list_repo_files(repo_id=input_str, token=token)
        target = next((f for f in files if f in candidates), None)
        if not target:
            safes = [f for f in files if f.endswith(".safetensors")]
            if safes: target = safes[0]
        
        if not target: raise ValueError("No safetensors found")
        
        hf_hub_download(repo_id=input_str, filename=target, token=token, local_dir=TempDir)
        found = list(TempDir.rglob(target.split("/")[-1]))[0]
        if found != local_path: shutil.move(found, local_path)
        return local_path
        
    except Exception as e:
        if input_str.startswith("http"):
            try:
                headers = {"Authorization": f"Bearer {token}"} if token else {}
                r = requests.get(input_str, stream=True, headers=headers, timeout=60)
                r.raise_for_status()
                with open(local_path, 'wb') as f:
                    for chunk in r.iter_content(chunk_size=8192): f.write(chunk)
                return local_path
            except: pass
        raise e

def load_lora_to_memory(lora_path, precision_dtype=torch.bfloat16):
    state_dict = load_file(lora_path, device="cpu")
    pairs = {} 
    alphas = {}
    for k, v in state_dict.items():
        stem = get_key_stem(k)
        if "alpha" in k:
            alphas[stem] = v.item() if isinstance(v, torch.Tensor) else v
        else:
            if stem not in pairs: pairs[stem] = {}
            if "lora_down" in k or "lora_A" in k:
                pairs[stem]["down"] = v.to(dtype=precision_dtype)
                pairs[stem]["rank"] = v.shape[0]
            elif "lora_up" in k or "lora_B" in k:
                pairs[stem]["up"] = v.to(dtype=precision_dtype)
    for stem in pairs:
        pairs[stem]["alpha"] = alphas.get(stem, float(pairs[stem].get("rank", 1.0)))
    return pairs

class ShardBuffer:
    def __init__(self, max_size_gb, output_dir, output_repo, subfolder, hf_token, filename_prefix="model"):
        self.max_bytes = int(max_size_gb * 1024**3)
        self.output_dir = output_dir
        self.output_repo = output_repo
        self.subfolder = subfolder
        self.hf_token = hf_token
        self.filename_prefix = filename_prefix
        self.buffer = []
        self.current_bytes = 0
        self.shard_count = 0
        self.index_map = {}
        self.total_size = 0

    def add_tensor(self, key, tensor):
        if tensor.dtype == torch.bfloat16:
            raw_bytes = tensor.view(torch.int16).numpy().tobytes()
            dtype_str = "BF16"
        elif tensor.dtype == torch.float16:
            raw_bytes = tensor.numpy().tobytes()
            dtype_str = "F16"
        else:
            raw_bytes = tensor.numpy().tobytes()
            dtype_str = "F32"
        size = len(raw_bytes)
        self.buffer.append({"key": key, "data": raw_bytes, "dtype": dtype_str, "shape": tensor.shape})
        self.current_bytes += size
        self.total_size += size
        if self.current_bytes >= self.max_bytes: self.flush()
            
    def flush(self):
        if not self.buffer: return
        self.shard_count += 1
        filename = f"{self.filename_prefix}-{self.shard_count:05d}.safetensors"
        path_in_repo = f"{self.subfolder}/{filename}" if self.subfolder else filename
        header = {"__metadata__": {"format": "pt"}}
        current_offset = 0
        for item in self.buffer:
            header[item["key"]] = {"dtype": item["dtype"], "shape": item["shape"], "data_offsets": [current_offset, current_offset + len(item["data"])]}
            current_offset += len(item["data"])
            self.index_map[item["key"]] = filename
        header_json = json.dumps(header).encode('utf-8')
        out_path = self.output_dir / filename
        with open(out_path, 'wb') as f:
            f.write(struct.pack('<Q', len(header_json)))
            f.write(header_json)
            for item in self.buffer: f.write(item["data"])
        api.upload_file(path_or_fileobj=out_path, path_in_repo=path_in_repo, repo_id=self.output_repo, token=self.hf_token)
        os.remove(out_path)
        self.buffer = []
        self.current_bytes = 0
        gc.collect()

def streaming_copy_structure(token, src_repo, dst_repo, ignore_prefix=None, is_root_merge=False):
    try:
        files = api.list_repo_files(repo_id=src_repo, token=token)
        for f in tqdm(files, desc="Copying Structure"):
            if ignore_prefix and f.startswith(ignore_prefix): continue
            if is_root_merge:
                if any(f.endswith(ext) for ext in ['.safetensors', '.bin', '.pt', '.pth']): continue
            try:
                local = hf_hub_download(repo_id=src_repo, filename=f, token=token, local_dir=TempDir)
                api.upload_file(path_or_fileobj=local, path_in_repo=f, repo_id=dst_repo, token=token)
                if os.path.exists(local): os.remove(local)
            except: pass
    except: pass

def task_merge(hf_token, base_repo, base_subfolder, lora_input, scale, precision, shard_size, output_repo, structure_repo, private, progress=gr.Progress()):
    cleanup_temp()
    if not hf_token: return "Error: HF Token required."
    login(hf_token.strip())
    
    try: api.create_repo(repo_id=output_repo, private=private, exist_ok=True, token=hf_token)
    except Exception as e: return f"Error creating repo: {e}"
    
    output_subfolder = base_subfolder if base_subfolder else "" 
    if structure_repo:
        ignore = output_subfolder if output_subfolder else None
        streaming_copy_structure(hf_token, structure_repo, output_repo, ignore_prefix=ignore, is_root_merge=not bool(output_subfolder))

    progress(0.1, desc="Downloading Input Model...")
    files = list_repo_files(repo_id=base_repo, token=hf_token)
    input_shards = []
    
    for f in files:
        if f.endswith(".safetensors"):
            if output_subfolder and not f.startswith(output_subfolder): continue
            local = TempDir / "inputs" / os.path.basename(f)
            os.makedirs(local.parent, exist_ok=True)
            hf_hub_download(repo_id=base_repo, filename=f, token=hf_token, local_dir=local.parent, local_dir_use_symlinks=False)
            found = list(local.parent.rglob(os.path.basename(f)))
            if found: input_shards.append(found[0])

    if not input_shards: return "No safetensors found."
    input_shards.sort()

    filename_prefix = "diffusion_pytorch_model" if (output_subfolder in ["transformer", "unet"] or "diffusion_pytorch_model" in os.path.basename(input_shards[0])) else "model"
    index_filename = f"{filename_prefix}.safetensors.index.json"

    dtype = torch.bfloat16 if precision == "bf16" else torch.float16 if precision == "fp16" else torch.float32
    try:
        progress(0.15, desc="Downloading LoRA...")
        lora_path = download_lora_smart(lora_input, hf_token)
        lora_pairs = load_lora_to_memory(lora_path, precision_dtype=dtype)
    except Exception as e: return f"Error loading LoRA: {e}"

    buffer = ShardBuffer(shard_size, TempDir, output_repo, output_subfolder, hf_token, filename_prefix=filename_prefix)
    
    for i, shard_file in enumerate(input_shards):
        progress(0.2 + (0.7 * i / len(input_shards)), desc=f"Processing {os.path.basename(shard_file)}")
        with MemoryEfficientSafeOpen(shard_file) as f:
            for k in f.keys():
                v = f.get_tensor(k)
                base_stem = get_key_stem(k)
                match = lora_pairs.get(base_stem)
                if not match:
                     if "to_q" in base_stem: match = lora_pairs.get(base_stem.replace("to_q", "qkv"))
                     elif "to_k" in base_stem: match = lora_pairs.get(base_stem.replace("to_k", "qkv"))
                     elif "to_v" in base_stem: match = lora_pairs.get(base_stem.replace("to_v", "qkv"))

                if match:
                    down, up = match["down"], match["up"]
                    scaling = scale * (match["alpha"] / match["rank"])
                    if len(v.shape) == 4 and len(down.shape) == 2:
                        down, up = down.unsqueeze(-1).unsqueeze(-1), up.unsqueeze(-1).unsqueeze(-1)
                    try:
                        delta = (up.squeeze() @ down.squeeze()).reshape(up.shape[0], down.shape[1], 1, 1) if len(up.shape) == 4 else up @ down
                    except: delta = up.T @ down 
                    delta = delta * scaling
                    if delta.shape == v.shape: v = v.to(dtype).add_(delta.to(dtype))
                    del delta

                buffer.add_tensor(k, v.to(dtype))
                del v
        os.remove(shard_file)
        gc.collect()

    buffer.flush()
    index_data = {"metadata": {"total_size": buffer.total_size}, "weight_map": buffer.index_map}
    path_in_repo = f"{output_subfolder}/{index_filename}" if output_subfolder else index_filename
    with open(TempDir / index_filename, "w") as f: json.dump(index_data, f, indent=4)
    api.upload_file(path_or_fileobj=TempDir / index_filename, path_in_repo=path_in_repo, repo_id=output_repo, token=hf_token)
    cleanup_temp()
    return f"Done! Merged {buffer.shard_count} shards to {output_repo}"

def identify_and_download_model(input_str, token):
    repo_id, filename = parse_hf_url(input_str)
    if repo_id and filename:
        local_path = TempDir / os.path.basename(filename)
        if local_path.exists(): os.remove(local_path)
        hf_hub_download(repo_id=repo_id, filename=filename, token=token, local_dir=TempDir)
        return list(TempDir.rglob(os.path.basename(filename)))[0]

    files = list_repo_files(repo_id=input_str, token=token)
    priorities = ["transformer/diffusion_pytorch_model.safetensors", "unet/diffusion_pytorch_model.safetensors", "model.safetensors"]
    target_file = next((f for f in priorities if f in files), next((f for f in files if f.endswith(".safetensors") and "lora" not in f), None))
    if not target_file: raise ValueError("No model file found")
    
    hf_hub_download(repo_id=input_str, filename=target_file, token=token, local_dir=TempDir)
    return list(TempDir.rglob(os.path.basename(target_file)))[0]

def extract_lora_layer_by_layer(model_org, model_tuned, rank, clamp):
    org = MemoryEfficientSafeOpen(model_org)
    tuned = MemoryEfficientSafeOpen(model_tuned)
    lora_sd = {}
    keys = set(org.keys()).intersection(set(tuned.keys()))
    
    for key in tqdm(keys, desc="Extracting"):
        if "num_batches_tracked" in key or "running_mean" in key or "running_var" in key: continue
        mat_org = org.get_tensor(key).float()
        mat_tuned = tuned.get_tensor(key).float()
        if mat_org.shape != mat_tuned.shape: continue
        diff = mat_tuned - mat_org
        if torch.max(torch.abs(diff)) < 1e-4: continue
        
        out_dim, in_dim = diff.shape[0], diff.shape[1] if len(diff.shape) > 1 else 1
        r = min(rank, in_dim, out_dim)
        is_conv = len(diff.shape) == 4
        if is_conv: diff = diff.flatten(start_dim=1)
        elif len(diff.shape) == 1: diff = diff.unsqueeze(1)
            
        U, S, V = torch.svd_lowrank(diff, q=r+4, niter=4)
        Vh = V.t()
        U, S, Vh = U[:, :r], S[:r], Vh[:r, :]
        U = U @ torch.diag(S)
        
        dist = torch.cat([U.flatten(), Vh.flatten()])
        hi_val = torch.quantile(torch.abs(dist), clamp)
        if hi_val > 0:
            U, Vh = U.clamp(-hi_val, hi_val), Vh.clamp(-hi_val, hi_val)
            
        if is_conv:
            U = U.reshape(out_dim, r, 1, 1)
            Vh = Vh.reshape(r, in_dim, mat_org.shape[2], mat_org.shape[3])
        else:
            U = U.reshape(out_dim, r)
            Vh = Vh.reshape(r, in_dim)
            
        stem = key.replace(".weight", "")
        lora_sd[f"{stem}.lora_up.weight"] = U.contiguous()
        lora_sd[f"{stem}.lora_down.weight"] = Vh.contiguous()
        lora_sd[f"{stem}.alpha"] = torch.tensor(r).float()
            
    out = TempDir / "extracted.safetensors"
    save_file(lora_sd, out)
    return str(out)

def task_extract(hf_token, org, tun, rank, out):
    cleanup_temp()
    if hf_token: login(hf_token.strip())
    try:
        p1 = identify_and_download_model(org, hf_token)
        p2 = identify_and_download_model(tun, hf_token)
        f = extract_lora_layer_by_layer(p1, p2, int(rank), 0.99)
        api.create_repo(repo_id=out, exist_ok=True, token=hf_token)
        api.upload_file(path_or_fileobj=f, path_in_repo="extracted_lora.safetensors", repo_id=out, token=hf_token)
        return "Done! Extracted to " + out
    except Exception as e: return f"Error: {e}"

def load_full_state_dict(path):
    raw = load_file(path, device="cpu")
    cleaned = {}
    for k, v in raw.items():
        if "lora_A" in k: new_k = k.replace("lora_A", "lora_down")
        elif "lora_B" in k: new_k = k.replace("lora_B", "lora_up")
        else: new_k = k
        cleaned[new_k] = v.float()
    return cleaned

def task_merge_adapters_advanced(hf_token, inputs_text, method, weight_str, beta, sigma_rel, target_rank, out_repo, private):
    cleanup_temp()
    if hf_token: login(hf_token.strip())
    urls = [line.strip() for line in inputs_text.replace(" ", "\n").split('\n') if line.strip()]
    if len(urls) < 2: return "Error: Please provide at least 2 adapters."
    
    try:
        weights = [float(w.strip()) for w in weight_str.split(',')] if weight_str.strip() else [1.0] * len(urls)
        if len(weights) < len(urls): weights += [1.0] * (len(urls) - len(weights))
    except: return "Error parsing weights."

    paths = []
    try:
        for url in tqdm(urls, desc="Downloading Adapters"): paths.append(download_lora_smart(url, hf_token))
    except Exception as e: return f"Download Error: {e}"

    merged = None
    if "Iterative EMA" in method:
        base_sd = load_file(paths[0], device="cpu")
        for k in base_sd:
             if base_sd[k].dtype.is_floating_point: base_sd[k] = base_sd[k].float()
        
        gamma = None
        if sigma_rel > 0:
            t_val = sigma_rel**-2
            roots = np.roots([1, 7, 16 - t_val, 12 - t_val])
            gamma = roots[np.isreal(roots) & (roots.real >= 0)].real.max()
            
        for i, path in enumerate(paths[1:]):
            current_beta = (1 - 1 / (i + 1)) ** (gamma + 1) if gamma is not None else beta
            curr = load_file(path, device="cpu")
            for k in base_sd:
                if k in curr and "alpha" not in k:
                    base_sd[k] = base_sd[k] * current_beta + curr[k].float() * (1 - current_beta)
        merged = base_sd
    else:
        states = [load_full_state_dict(p) for p in paths]
        merged = {}
        all_stems = set()
        for s in states:
            for k in s: 
                if "lora_" in k: all_stems.add(k.split(".lora_")[0])

        for stem in tqdm(all_stems):
            down_list, up_list = [], []
            alpha_sum = 0.0
            total_delta = None
            
            for i, state in enumerate(states):
                w = weights[i]
                dk, uk, ak = f"{stem}.lora_down.weight", f"{stem}.lora_up.weight", f"{stem}.alpha"
                if dk in state and uk in state:
                    d, u = state[dk], state[uk]
                    alpha_sum += state[ak].item() if ak in state else d.shape[0]
                    
                    if "Concatenation" in method:
                        down_list.append(d)
                        up_list.append(u * w)
                    elif "SVD" in method:
                        rank, alpha = d.shape[0], state[ak].item() if ak in state else d.shape[0]
                        scale = (alpha / rank) * w
                        delta = ((u.flatten(1) @ d.flatten(1)).reshape(u.shape[0], d.shape[1], d.shape[2], d.shape[3]) if len(d.shape)==4 else u @ d) * scale
                        total_delta = delta if total_delta is None else total_delta + delta

            if "Concatenation" in method and down_list:
                merged[f"{stem}.lora_down.weight"] = torch.cat(down_list, dim=0).contiguous()
                merged[f"{stem}.lora_up.weight"] = torch.cat(up_list, dim=1).contiguous()
                merged[f"{stem}.alpha"] = torch.tensor(alpha_sum)
            elif "SVD" in method and total_delta is not None:
                tr = int(target_rank)
                flat = total_delta.flatten(1) if len(total_delta.shape)==4 else total_delta
                try:
                    U, S, V = torch.svd_lowrank(flat, q=tr + 4, niter=4)
                    Vh = V.t()
                    U, S, Vh = U[:, :tr], S[:tr], Vh[:tr, :]
                    U = U @ torch.diag(S)
                    
                    if len(total_delta.shape) == 4:
                        U = U.reshape(total_delta.shape[0], tr, 1, 1)
                        Vh = Vh.reshape(tr, total_delta.shape[1], total_delta.shape[2], total_delta.shape[3])
                    else:
                        U, Vh = U.reshape(total_delta.shape[0], tr), Vh.reshape(tr, total_delta.shape[1])
                        
                    merged[f"{stem}.lora_down.weight"] = Vh.contiguous()
                    merged[f"{stem}.lora_up.weight"] = U.contiguous()
                    merged[f"{stem}.alpha"] = torch.tensor(tr).float()
                except: pass

    out = TempDir / "merged_adapters.safetensors"
    save_file(merged, out)
    api.create_repo(repo_id=out_repo, private=private, exist_ok=True, token=hf_token)
    api.upload_file(path_or_fileobj=out, path_in_repo="merged_adapters.safetensors", repo_id=out_repo, token=hf_token)
    return f"Success! Merged to {out_repo}"

def task_resize(hf_token, lora_input, new_rank, dynamic_method, dynamic_param, out_repo):
    cleanup_temp()
    if hf_token: login(hf_token.strip())
    path = download_lora_smart(lora_input, hf_token)
    state = load_file(path, device="cpu")
    new_state = {}
    groups = {}
    
    for k in state:
        simple = k.split(".lora_")[0]
        if simple not in groups: groups[simple] = {}
        if "lora_down" in k or "lora_A" in k: groups[simple]["down"] = state[k]
        if "lora_up" in k or "lora_B" in k: groups[simple]["up"] = state[k]
        if "alpha" in k: groups[simple]["alpha"] = state[k]

    target_rank_limit = int(new_rank)
    for stem, g in tqdm(groups.items()):
        if "down" in g and "up" in g:
            down, up = g["down"].float(), g["up"].float()
            merged = (up.squeeze() @ down.squeeze()).reshape(up.shape[0], down.shape[1], down.shape[2], down.shape[3]) if len(down.shape)==4 else up @ down
            flat = merged.flatten(1)
            
            U, S, V = torch.svd_lowrank(flat, q=target_rank_limit + 32)
            Vh = V.t() 
            
            calc_rank = target_rank_limit
            if dynamic_method == "sv_ratio":
                calc_rank = int(torch.sum(S > (S[0] / dynamic_param)).item())
            elif dynamic_method == "sv_cumulative":
                calc_rank = int(torch.searchsorted(torch.cumsum(S, 0) / torch.sum(S), dynamic_param)) + 1
            elif dynamic_method == "sv_fro":
                calc_rank = int(torch.searchsorted(torch.cumsum(S.pow(2), 0) / torch.sum(S.pow(2)), dynamic_param**2)) + 1
                
            final_rank = max(1, min(calc_rank, target_rank_limit, S.shape[0]))
            
            U = U[:, :final_rank] @ torch.diag(S[:final_rank])
            Vh = Vh[:final_rank, :]
            
            if len(down.shape) == 4:
                U = U.reshape(up.shape[0], final_rank, 1, 1)
                Vh = Vh.reshape(final_rank, down.shape[1], down.shape[2], down.shape[3])
                
            new_state[f"{stem}.lora_down.weight"] = Vh.contiguous()
            new_state[f"{stem}.lora_up.weight"] = U.contiguous()
            new_state[f"{stem}.alpha"] = torch.tensor(final_rank).float()

    out = TempDir / "shrunken.safetensors"
    save_file(new_state, out)
    api.create_repo(repo_id=out_repo, exist_ok=True, token=hf_token)
    api.upload_file(path_or_fileobj=out, path_in_repo="shrunken.safetensors", repo_id=out_repo, token=hf_token)
    return "Done"

# =================================================================================
# MERGEKIT CLI HELPERS
# =================================================================================

def run_mergekit_yaml(config_str, output_path, hf_token):
    """Run mergekit-yaml CLI subprocess"""
    config_file = TempDir / "config.yaml"
    with open(config_file, "w") as f:
        f.write(config_str)
    
    env = os.environ.copy()
    if hf_token:
        env["HF_TOKEN"] = hf_token.strip()

    cmd = [
        "mergekit-yaml",
        str(config_file),
        str(output_path),
        "--allow-crimes",
        "--lazy-unpickle",
        "--copy-tokenizer"
    ]
    
    print(f"Running: {' '.join(cmd)}")
    result = subprocess.run(cmd, env=env, capture_output=True, text=True)
    
    if result.returncode != 0:
        raise RuntimeError(f"MergeKit failed:\n{result.stderr}")
    
    return str(output_path)

def run_mergekit_moe(config_str, output_path, hf_token):
    """Run mergekit-moe CLI subprocess"""
    config_file = TempDir / "moe_config.yaml"
    with open(config_file, "w") as f:
        f.write(config_str)
    
    env = os.environ.copy()
    if hf_token:
        env["HF_TOKEN"] = hf_token.strip()

    cmd = [
        "mergekit-moe",
        str(config_file),
        str(output_path),
        "--copy-tokenizer"
    ]
    
    print(f"Running: {' '.join(cmd)}")
    result = subprocess.run(cmd, env=env, capture_output=True, text=True)
    
    if result.returncode != 0:
        raise RuntimeError(f"MergeKit MoE failed:\n{result.stderr}")
    
    return str(output_path)

def upload_folder_to_hf(folder, repo_id, token, private=True):
    """Upload entire folder to HuggingFace"""
    api.create_repo(repo_id=repo_id, private=private, exist_ok=True, token=token)
    api.upload_folder(folder_path=folder, repo_id=repo_id, token=token)
    return f"Success! Uploaded to {repo_id}"

def parse_gradient_value(value_str):
    """Parse gradient values like '[0, 0.3, 0.7, 1]' or '0.5'"""
    value_str = value_str.strip()
    if not value_str:
        return 1.0
    if value_str.startswith('[') and value_str.endswith(']'):
        try:
            return json.loads(value_str)
        except:
            pass
    try:
        return float(value_str)
    except:
        return 1.0

# =================================================================================
# TAB 5: AMPHINTERPOLATIVE (SLERP, NUSLERP, MULTISLERP, KARCHER)
# =================================================================================

def task_amphinterpolative(
    hf_token, method, dtype, base_model,
    model1, weight1, model2, weight2, model3, weight3, model4, weight4, model5, weight5,
    t_val, normalize, int8_mask, tokenizer_source,
    nuslerp_flatten, nuslerp_row_wise, eps, max_iter, tol,
    out_repo, private
):
    cleanup_temp()
    if not hf_token:
        return "Error: Token required"
    login(hf_token.strip())
    
    # Build model list
    models = []
    if model1.strip():
        models.append({
            "model": model1.strip(),
            "parameters": {"weight": parse_gradient_value(weight1)}
        })
    if model2.strip():
        models.append({
            "model": model2.strip(),
            "parameters": {"weight": parse_gradient_value(weight2)}
        })
    if model3.strip():
        models.append({
            "model": model3.strip(),
            "parameters": {"weight": parse_gradient_value(weight3)}
        })
    if model4.strip():
        models.append({
            "model": model4.strip(),
            "parameters": {"weight": parse_gradient_value(weight4)}
        })
    if model5.strip():
        models.append({
            "model": model5.strip(),
            "parameters": {"weight": parse_gradient_value(weight5)}
        })
    
    if len(models) < 2:
        return "Error: At least 2 models required"
    
    # Validate method requirements
    if method == "slerp" and not base_model.strip():
        return "Error: slerp requires base_model"
    if method == "slerp" and len(models) != 2:
        return "Error: slerp requires exactly 2 models"
    if method == "nuslerp" and len(models) != 2:
        return "Error: nuslerp requires exactly 2 models"
    
    # Build config
    config = {
        "merge_method": method,
        "dtype": dtype
    }
    
    if base_model.strip():
        config["base_model"] = base_model.strip()
    
    # Use slices format for slerp/nuslerp (per MergeKit docs)
    if method in ["slerp", "nuslerp"]:
        config["slices"] = [{
            "sources": [
                {"model": models[0]["model"]},
                {"model": models[1]["model"]}
            ],
            "parameters": {"t": parse_gradient_value(t_val)}
        }]
    else:
        # Use models format for multislerp/karcher
        config["models"] = models
    
    # Add global parameters
    params = {}
    if method in ["slerp", "nuslerp", "multislerp"]:
        params["t"] = parse_gradient_value(t_val)
    if normalize:
        params["normalize"] = True
    if int8_mask:
        params["int8_mask"] = True
    
    # Method-specific parameters
    if method == "nuslerp":
        if nuslerp_flatten:
            params["nuslerp_flatten"] = True
        if nuslerp_row_wise:
            params["nuslerp_row_wise"] = True
    elif method == "multislerp":
        try:
            params["eps"] = float(eps)
        except:
            params["eps"] = 1e-8
    elif method == "karcher":
        try:
            params["max_iter"] = int(max_iter)
            params["tol"] = float(tol)
        except:
            params["max_iter"] = 10
            params["tol"] = 1e-5
    
    if params:
        config["parameters"] = params
    
    if tokenizer_source != "base":
        config["tokenizer_source"] = tokenizer_source
    
    yaml_str = yaml.dump(config, sort_keys=False)
    out_path = TempDir / "out_interp"
    
    try:
        run_mergekit_yaml(yaml_str, out_path, hf_token)
        return upload_folder_to_hf(str(out_path), out_repo, hf_token, private)
    except Exception as e:
        return f"Error: {e}"

# =================================================================================
# TAB 6: STIR/TIE BASES (TASK VECTOR FAMILY)
# =================================================================================

def task_stir_tie(
    hf_token, method, dtype, base_model,
    model1, weight1, density1, gamma1, epsilon1,
    model2, weight2, density2, gamma2, epsilon2,
    model3, weight3, density3, gamma3, epsilon3,
    model4, weight4, density4, gamma4, epsilon4,
    normalize, int8_mask, lambda_val, rescale, select_topk,
    tokenizer_source, out_repo, private
):
    cleanup_temp()
    if not hf_token:
        return "Error: Token required"
    login(hf_token.strip())
    
    if not base_model.strip():
        return "Error: base_model required for Task Vector methods"
    
    # Build models list with per-model parameters
    models = []
    for i, (m, w, d, g, e) in enumerate([
        (model1, weight1, density1, gamma1, epsilon1),
        (model2, weight2, density2, gamma2, epsilon2),
        (model3, weight3, density3, gamma3, epsilon3),
        (model4, weight4, density4, gamma4, epsilon4)
    ]):
        if not m.strip():
            continue
        
        params = {"weight": parse_gradient_value(w)}
        
        # Add density for DARE/TIES methods
        if method in ["ties", "dare_ties", "dare_linear"] and d.strip():
            params["density"] = parse_gradient_value(d)
        
        # Add gamma for breadcrumbs
        if method in ["breadcrumbs", "breadcrumbs_ties"] and g.strip():
            try:
                params["gamma"] = float(g)
            except:
                params["gamma"] = 0.01
        
        # Add epsilon for DELLA
        if method in ["della", "della_linear"] and e.strip():
            try:
                params["epsilon"] = float(e)
            except:
                params["epsilon"] = 0.15
        
        models.append({"model": m.strip(), "parameters": params})
    
    if len(models) < 1:
        return "Error: At least 1 model required (in addition to base_model)"
    
    # Build global parameters
    global_params = {"normalize": normalize}
    if int8_mask:
        global_params["int8_mask"] = True
    if lambda_val.strip():
        try:
            global_params["lambda"] = float(lambda_val)
        except:
            pass
    if method in ["dare_linear"] and rescale:
        global_params["rescale"] = True
    if method == "sce" and select_topk.strip():
        try:
            global_params["select_topk"] = float(select_topk)
        except:
            global_params["select_topk"] = 1.0
    
    config = {
        "models": models,
        "merge_method": method,
        "base_model": base_model.strip(),
        "parameters": global_params,
        "dtype": dtype
    }
    
    if tokenizer_source != "base":
        config["tokenizer_source"] = tokenizer_source
    
    yaml_str = yaml.dump(config, sort_keys=False)
    out_path = TempDir / "out_task_vec"
    
    try:
        run_mergekit_yaml(yaml_str, out_path, hf_token)
        return upload_folder_to_hf(str(out_path), out_repo, hf_token, private)
    except Exception as e:
        return f"Error: {e}"

# =================================================================================
# TAB 7: SPECIOUS (SPECIALIZED METHODS)
# =================================================================================

def task_specious(
    hf_token, method, dtype, base_model,
    model1, weight1, filter1,
    model2, weight2, filter2,
    model3, weight3,
    model4, weight4,
    model5, weight5,
    t_val, normalize, int8_mask, filter_wise, tokenizer_source,
    out_repo, private
):
    cleanup_temp()
    if not hf_token:
        return "Error: Token required"
    login(hf_token.strip())
    
    # Build models list
    models = []
    
    if model1.strip():
        params = {"weight": parse_gradient_value(weight1)}
        # For passthrough, add filter
        if method == "passthrough" and filter1.strip():
            params["filter"] = filter1.strip()
        models.append({"model": model1.strip(), "parameters": params})
    
    if model2.strip():
        params = {"weight": parse_gradient_value(weight2)}
        if method == "passthrough" and filter2.strip():
            params["filter"] = filter2.strip()
        models.append({"model": model2.strip(), "parameters": params})
    
    if model3.strip():
        models.append({
            "model": model3.strip(),
            "parameters": {"weight": parse_gradient_value(weight3)}
        })
    if model4.strip():
        models.append({
            "model": model4.strip(),
            "parameters": {"weight": parse_gradient_value(weight4)}
        })
    if model5.strip():
        models.append({
            "model": model5.strip(),
            "parameters": {"weight": parse_gradient_value(weight5)}
        })
    
    # Validate method requirements
    if method == "passthrough" and len(models) != 1:
        return "Error: passthrough requires exactly 1 model"
    if method in ["nearswap", "arcee_fusion"] and len(models) != 2:
        return f"Error: {method} requires exactly 2 models"
    if method == "model_stock" and len(models) < 3:
        return "Error: model_stock requires at least 3 models"
    
    if not models:
        return "Error: At least 1 model required"
    
    # Build config
    config = {
        "models": models,
        "merge_method": method,
        "dtype": dtype
    }
    
    if base_model.strip():
        config["base_model"] = base_model.strip()
    
    # Add parameters
    params = {}
    if normalize:
        params["normalize"] = True
    if int8_mask:
        params["int8_mask"] = True
    if method == "nearswap" and t_val.strip():
        try:
            params["t"] = parse_gradient_value(t_val)
        except:
            params["t"] = 0.5
    if method == "model_stock" and filter_wise:
        params["filter_wise"] = True
    
    if params:
        config["parameters"] = params
    
    if tokenizer_source != "base":
        config["tokenizer_source"] = tokenizer_source
    
    yaml_str = yaml.dump(config, sort_keys=False)
    out_path = TempDir / "out_specious"
    
    try:
        run_mergekit_yaml(yaml_str, out_path, hf_token)
        return upload_folder_to_hf(str(out_path), out_repo, hf_token, private)
    except Exception as e:
        return f"Error: {e}"

# =================================================================================
# TAB 8: MOER (MIXTURE OF EXPERTS)
# =================================================================================

def task_moer(
    hf_token, base_model, dtype,
    expert1, prompt1,
    expert2, prompt2,
    expert3, prompt3,
    expert4, prompt4,
    expert5, prompt5,
    gate_mode, tokenizer_source,
    out_repo, private
):
    cleanup_temp()
    if not hf_token:
        return "Error: Token required"
    login(hf_token.strip())
    
    if not base_model.strip():
        return "Error: base_model required"
    
    # Build experts list
    experts = []
    for exp, pmt in [
        (expert1, prompt1), (expert2, prompt2), (expert3, prompt3),
        (expert4, prompt4), (expert5, prompt5)
    ]:
        if exp.strip():
            expert_entry = {"source_model": exp.strip()}
            # Parse prompts (comma-separated)
            if pmt.strip():
                prompts = [p.strip() for p in pmt.split(',') if p.strip()]
                expert_entry["positive_prompts"] = prompts
            else:
                expert_entry["positive_prompts"] = [""]
            experts.append(expert_entry)
    
    if len(experts) < 2:
        return "Error: At least 2 experts required"
    
    # Build config for MoE
    config = {
        "base_model": base_model.strip(),
        "gate_mode": gate_mode,
        "dtype": dtype,
        "experts": experts
    }
    
    if tokenizer_source != "base":
        config["tokenizer_source"] = tokenizer_source
    
    yaml_str = yaml.dump(config, sort_keys=False)
    out_path = TempDir / "out_moe"
    
    try:
        run_mergekit_moe(yaml_str, out_path, hf_token)
        return upload_folder_to_hf(str(out_path), out_repo, hf_token, private)
    except Exception as e:
        return f"Error: {e}"

# =================================================================================
# TAB 9: RAWER (RAW PYTORCH / NON-TRANSFORMER)
# =================================================================================

def task_rawer(
    hf_token, models_text, method, dtype, 
    tokenizer_source, out_repo, private
):
    cleanup_temp()
    if not hf_token:
        return "Error: Token required"
    login(hf_token.strip())
    
    models = [m.strip() for m in models_text.split('\n') if m.strip()]
    if not models:
        return "Error: No models listed"
    
    # For raw/passthrough, simple linear structure
    config = {
        "models": [{"model": m, "parameters": {"weight": 1.0}} for m in models],
        "merge_method": method,
        "dtype": dtype
    }
    
    if tokenizer_source != "base":
        config["tokenizer_source"] = tokenizer_source
    
    yaml_str = yaml.dump(config, sort_keys=False)
    out_path = TempDir / "out_raw"
    
    try:
        run_mergekit_yaml(yaml_str, out_path, hf_token)
        return upload_folder_to_hf(str(out_path), out_repo, hf_token, private)
    except Exception as e:
        return f"Error: {e}"

# =================================================================================
# TAB 10: MARIO,DARE! (CUSTOM DARE IMPLEMENTATION)
# =================================================================================

def task_dare_soonr(hf_token, base_model, ft_model, ratio, mask_rate, out_repo, private):
    cleanup_temp()
    if not hf_token:
        return "Error: Token required"
    login(hf_token.strip())

    try:
        print("Downloading Base Model...")
        base_path = identify_and_download_model(base_model, hf_token)
        print("Downloading Fine-Tuned Model...")
        ft_path = identify_and_download_model(ft_model, hf_token)
        
        print("Loading Tensors...")
        base_sd = load_file(base_path, device="cpu")
        ft_sd = load_file(ft_path, device="cpu")
        
        merged_sd = {}
        common_keys = set(base_sd.keys()).intersection(set(ft_sd.keys()))
        
        print("DARE Merging...")
        for key in tqdm(common_keys, desc="DARE Merge"):
            base_t = base_sd[key]
            ft_t = ft_sd[key]
            
            if base_t.dtype != ft_t.dtype or base_t.shape != ft_t.shape:
                merged_sd[key] = ft_t
                continue
            
            # DARE Algorithm:
            # 1. Compute delta = fine_tuned - base
            delta = ft_t.float() - base_t.float()
            
            # 2. Apply mask (drop) via Bernoulli sampling
            if mask_rate > 0.0:
                # Create Bernoulli mask (1 = keep, 0 = drop)
                mask = torch.bernoulli(torch.full_like(delta, 1.0 - mask_rate))
                # Rescale to compensate for dropped elements
                rescale_factor = 1.0 / (1.0 - mask_rate) if mask_rate < 1.0 else 1.0
                delta = delta * mask * rescale_factor
            
            # 3. Apply ratio and add to base
            merged_t = base_t.float() + (delta * float(ratio))
            
            # 4. Cast back to original dtype
            if base_t.dtype == torch.bfloat16:
                merged_sd[key] = merged_t.bfloat16()
            elif base_t.dtype == torch.float16:
                merged_sd[key] = merged_t.half()
            else:
                merged_sd[key] = merged_t
        
        # Save and upload
        out_path = TempDir / "dare_merged.safetensors"
        save_file(merged_sd, out_path)
        
        api.create_repo(repo_id=out_repo, private=private, exist_ok=True, token=hf_token)
        api.upload_file(
            path_or_fileobj=out_path,
            path_in_repo="model.safetensors",
            repo_id=out_repo,
            token=hf_token
        )
        
        cleanup_temp()
        return f"Success! DARE-merged model uploaded to {out_repo}"
        
    except Exception as e:
        return f"DARE Error: {e}"

# =================================================================================
# UI DEFINITION
# =================================================================================

css = ".container { max-width: 1100px; margin: auto; }"

with gr.Blocks() as demo:
    title = gr.HTML(
        """<h1><img src="https://huggingface.co/spaces/AlekseyCalvin/Soon_Merger/resolve/main/SMerger3.png" alt="SOONmerge®"> Transform Transformers for FREE!</h1>""",
        elem_id="title",
    )
    gr.Markdown("# 🧰 Training-Free CPU-run Model Creation Toolkit")
    
    with gr.Tabs():
        # TAB 1: Merge into Base Model (UNCHANGED)
        with gr.Tab("Merge into Base Model"):
            with gr.Row():
                t1_token = gr.Textbox(label="Token", type="password")
            with gr.Row():
                t1_base = gr.Textbox(label="Base Repo", value="name/repo")
            t1_sub = gr.Textbox(label="Subfolder (Optional)", value="")
            t1_lora = gr.Textbox(label="LoRA Direct Link or Repo", value="https://huggingface.co/GuangyuanSD/Z-Image-Re-Turbo-LoRA/resolve/main/Z-image_re_turbo_lora_8steps_rank_32_v1_fp16.safetensors")
            with gr.Row():
                t1_scale = gr.Slider(label="Scale", value=1.0, minimum=0, maximum=3.0, step=0.1)
                t1_prec = gr.Radio(["bf16", "fp16", "float32"], value="bf16", label="Precision")
                t1_shard = gr.Slider(label="Max Shard Size (GB)", value=2.0, minimum=0.1, maximum=10.0, step=0.1)
            t1_out = gr.Textbox(label="Output Repo")
            t1_struct = gr.Textbox(label="Extras Source (copies configs/components/etc)", value="name/repo")
            t1_priv = gr.Checkbox(label="Private", value=True)
            t1_btn = gr.Button("Merge")
            t1_res = gr.Textbox(label="Result")
            t1_btn.click(task_merge, [t1_token, t1_base, t1_sub, t1_lora, t1_scale, t1_prec, t1_shard, t1_out, t1_struct, t1_priv], t1_res)

        # TAB 2: Extract Adapter (UNCHANGED)
        with gr.Tab("Extract Adapter"):
            t2_token = gr.Textbox(label="Token", type="password")
            t2_org = gr.Textbox(label="Original Model")
            t2_tun = gr.Textbox(label="Tuned or Homologous Model")
            t2_rank = gr.Number(label="Extract At Rank", value=32, minimum=1, maximum=1024, step=1)
            t2_out = gr.Textbox(label="Output Repo")
            t2_btn = gr.Button("Extract")
            t2_res = gr.Textbox(label="Result")
            t2_btn.click(task_extract, [t2_token, t2_org, t2_tun, t2_rank, t2_out], t2_res)
            
        # TAB 3: Merge Adapters (UNCHANGED)
        with gr.Tab("Merge Adapters"):
            gr.Markdown("### Batch Adapter Merging")
            t3_token = gr.Textbox(label="Token", type="password")
            t3_urls = gr.TextArea(label="Adapter URLs/Repos (one per line, or space-separated)", placeholder="user/lora1\nhttps://hf.co/user/lora2.safetensors\n...")
            with gr.Row():
                t3_method = gr.Dropdown(
                    ["Iterative EMA (Linear w/ Beta/Sigma coefficient)", "Concatenation (MOE-like weights-stack)", "SVD Fusion (Task Arithmetic/Compressed)"], 
                    value="Iterative EMA (Linear w/ Beta/Sigma coefficient)", 
                    label="Merge Method"
                )
            with gr.Row():
                t3_weights = gr.Textbox(label="Weights (comma-separated) – for Concat/SVD", placeholder="1.0, 0.5, 0.8...")
                t3_rank = gr.Number(label="Target Rank – For SVD only", value=128, minimum=1, maximum=1024)
            with gr.Row():
                t3_beta = gr.Slider(label="Beta – for linear/post-hoc EMA", value=0.95, minimum=0.01, maximum=1.00, step=0.01)
                t3_sigma = gr.Slider(label="Sigma Rel – for linear/post-hoc EMA", value=0.21, minimum=0.01, maximum=1.00, step=0.01)
            t3_out = gr.Textbox(label="Output Repo")
            t3_priv = gr.Checkbox(label="Private Output", value=True)
            t3_btn = gr.Button("Merge")
            t3_res = gr.Textbox(label="Result")
            t3_btn.click(task_merge_adapters_advanced, [t3_token, t3_urls, t3_method, t3_weights, t3_beta, t3_sigma, t3_rank, t3_out, t3_priv], t3_res)
        # TAB 4: Resize Adapter (UNCHANGED)
        with gr.Tab("Resize Adapter"):
            t4_token = gr.Textbox(label="Token", type="password")
            t4_in = gr.Textbox(label="LoRA")
            with gr.Row():
                t4_rank = gr.Number(label="To Rank (Safety Ceiling)", value=8, minimum=1, maximum=512, step=1)
                t4_method = gr.Dropdown(["None", "sv_ratio", "sv_fro", "sv_cumulative"], value="None", label="Dynamic Method")
                t4_param = gr.Number(label="Dynamic Param", value=0.9)
                
            gr.Markdown(
                """
                ### 📉 Dynamic Resizing Guide
                These methods intelligently determine the best rank per layer. 
                * **sv_ratio (Relative Strength):** Keeps features that are at least `1/Param` as strong as the main feature. **Param must be >= 2**. (e.g. 2 = keep features half as strong as top).
                * **sv_fro (Visual Information Density):** Preserves `Param%` of the total information content (Frobenius Norm) of the layer. **Param between 0.0 and 1.0** (e.g. 0.9 = 90% info retention).
                * **sv_cumulative (Cumulative Sum):** Preserves weights that sum up to `Param%` of the total strength. **Param between 0.0 and 1.0**.
                * **⚠️ Safety Ceiling:** The **"To Rank"** slider acts as a hard limit. Even if a dynamic method wants a higher rank, it will be cut down to this number to keep file sizes small.
                """
            )
            t4_out = gr.Textbox(label="Output")
            t4_btn = gr.Button("Resize")
            t4_res = gr.Textbox(label="Result")
            t4_btn.click(task_resize, [t4_token, t4_in, t4_rank, t4_method, t4_param, t4_out], t4_res)
        
        # TAB 5: AMPHINTERPOLATIVE
        with gr.Tab("Amphinterpolative"):
            gr.Markdown("### Spherical Interpolation Methods: slerp, nuslerp, multislerp, karcher")
            
            t5_token = gr.Textbox(label="HF Token", type="password")
            
            with gr.Row():
                t5_method = gr.Dropdown(
                    ["slerp", "nuslerp", "multislerp", "karcher"],
                    value="slerp",
                    label="Merge Method"
                )
                t5_dtype = gr.Dropdown(
                    ["float16", "bfloat16", "float32"],
                    value="bfloat16",
                    label="dtype"
                )
            
            t5_base = gr.Textbox(label="Base Model (required for slerp)", placeholder="org/model")
            
            gr.Markdown("#### Models (at least 2 required)")
            with gr.Row():
                t5_model1 = gr.Textbox(label="Model 1", placeholder="org/model1")
                t5_weight1 = gr.Textbox(label="Weight 1", value="1.0", placeholder="1.0 or [0, 0.5, 1]")
            
            with gr.Row():
                t5_model2 = gr.Textbox(label="Model 2", placeholder="org/model2")
                t5_weight2 = gr.Textbox(label="Weight 2", value="1.0", placeholder="1.0 or [0, 0.5, 1]")
            
            with gr.Row():
                t5_model3 = gr.Textbox(label="Model 3 (optional, for multislerp/karcher)", placeholder="org/model3")
                t5_weight3 = gr.Textbox(label="Weight 3", value="1.0")
            
            with gr.Row():
                t5_model4 = gr.Textbox(label="Model 4 (optional)", placeholder="org/model4")
                t5_weight4 = gr.Textbox(label="Weight 4", value="1.0")
            
            with gr.Row():
                t5_model5 = gr.Textbox(label="Model 5 (optional)", placeholder="org/model5")
                t5_weight5 = gr.Textbox(label="Weight 5", value="1.0")
            
            with gr.Accordion("Advanced Parameters", open=False):
                t5_t = gr.Textbox(label="t (interpolation factor)", value="0.5", placeholder="0.5 or [0, 0.3, 0.7, 1]")
                t5_normalize = gr.Checkbox(label="Normalize Weights", value=True)
                t5_int8_mask = gr.Checkbox(label="int8_mask", value=False)
                t5_tokenizer_source = gr.Dropdown(
                    ["base", "union"],
                    value="base",
                    label="Tokenizer Source"
                )
                
                gr.Markdown("**NuSlerp Specific:**")
                t5_nuslerp_flatten = gr.Checkbox(label="nuslerp_flatten (row/column-wise interpolation)", value=False)
                t5_nuslerp_row_wise = gr.Checkbox(label="nuslerp_row_wise", value=False)
                
                gr.Markdown("**MultiSlerp Specific:**")
                t5_eps = gr.Textbox(label="eps (numerical constant)", value="1e-8")
                
                gr.Markdown("**Karcher Specific:**")
                t5_max_iter = gr.Number(label="max_iter", value=10)
                t5_tol = gr.Textbox(label="tol (convergence tolerance)", value="1e-5")
            
            t5_out = gr.Textbox(label="Output Repo")
            t5_priv = gr.Checkbox(label="Private", value=True)
            t5_btn = gr.Button("Run MergeKit Interpolation")
            t5_res = gr.Textbox(label="Result")
            
            t5_btn.click(
                task_amphinterpolative,
                [t5_token, t5_method, t5_dtype, t5_base,
                 t5_model1, t5_weight1, t5_model2, t5_weight2, t5_model3, t5_weight3, t5_model4, t5_weight4, t5_model5, t5_weight5,
                 t5_t, t5_normalize, t5_int8_mask, t5_tokenizer_source,
                 t5_nuslerp_flatten, t5_nuslerp_row_wise, t5_eps, t5_max_iter, t5_tol,
                 t5_out, t5_priv],
                t5_res
            )
            
        # TAB 6: STIR/TIE BASES
        with gr.Tab("Stir/Tie Bases"):
            gr.Markdown("### Task Vector Methods: task_arithmetic, ties, dare_ties, dare_linear, della, della_linear, breadcrumbs, breadcrumbs_ties, sce")
            
            t6_token = gr.Textbox(label="HF Token", type="password")
            
            with gr.Row():
                t6_method = gr.Dropdown(
                    ["task_arithmetic", "ties", "dare_ties", "dare_linear", "della", "della_linear", "breadcrumbs", "breadcrumbs_ties", "sce"],
                    value="ties",
                    label="Merge Method"
                )
                t6_dtype = gr.Dropdown(
                    ["float16", "bfloat16", "float32"],
                    value="bfloat16",
                    label="dtype"
                )
            
            t6_base = gr.Textbox(label="Base Model (required)", placeholder="org/base-model")
            
            gr.Markdown("#### Models (at least 1 required)")
            with gr.Row():
                with gr.Column():
                    t6_model1 = gr.Textbox(label="Model 1", placeholder="org/model1")
                    t6_weight1 = gr.Textbox(label="Weight", value="1.0", placeholder="1.0 or [0, 0.5, 1]")
                with gr.Column():
                    t6_density1 = gr.Textbox(label="Density (DARE/TIES)", value="0.5", placeholder="0.5 or [0.3, 0.7]")
                    t6_gamma1 = gr.Textbox(label="Gamma (breadcrumbs)", value="0.01")
                    t6_epsilon1 = gr.Textbox(label="Epsilon (DELLA)", value="0.15")
            
            with gr.Row():
                with gr.Column():
                    t6_model2 = gr.Textbox(label="Model 2 (optional)", placeholder="org/model2")
                    t6_weight2 = gr.Textbox(label="Weight", value="1.0")
                with gr.Column():
                    t6_density2 = gr.Textbox(label="Density", value="0.5")
                    t6_gamma2 = gr.Textbox(label="Gamma", value="0.01")
                    t6_epsilon2 = gr.Textbox(label="Epsilon", value="0.15")
            
            with gr.Row():
                with gr.Column():
                    t6_model3 = gr.Textbox(label="Model 3 (optional)", placeholder="org/model3")
                    t6_weight3 = gr.Textbox(label="Weight", value="1.0")
                with gr.Column():
                    t6_density3 = gr.Textbox(label="Density", value="0.5")
                    t6_gamma3 = gr.Textbox(label="Gamma", value="0.01")
                    t6_epsilon3 = gr.Textbox(label="Epsilon", value="0.15")
            
            with gr.Row():
                with gr.Column():
                    t6_model4 = gr.Textbox(label="Model 4 (optional)", placeholder="org/model4")
                    t6_weight4 = gr.Textbox(label="Weight", value="1.0")
                with gr.Column():
                    t6_density4 = gr.Textbox(label="Density", value="0.5")
                    t6_gamma4 = gr.Textbox(label="Gamma", value="0.01")
                    t6_epsilon4 = gr.Textbox(label="Epsilon", value="0.15")
            
            with gr.Accordion("Global Parameters", open=False):
                t6_normalize = gr.Checkbox(label="Normalize", value=True)
                t6_int8_mask = gr.Checkbox(label="int8_mask", value=False)
                t6_lambda = gr.Textbox(label="lambda", value="1.0")
                t6_rescale = gr.Checkbox(label="rescale (for dare_linear)", value=True)
                t6_select_topk = gr.Textbox(label="select_topk (for sce)", value="1.0", placeholder="0-1.0")
                t6_tokenizer_source = gr.Dropdown(["base", "union"], value="base", label="Tokenizer Source")
            
            t6_out = gr.Textbox(label="Output Repo")
            t6_priv = gr.Checkbox(label="Private", value=True)
            t6_btn = gr.Button("Run MergeKit Task Vector")
            t6_res = gr.Textbox(label="Result")
            
            t6_btn.click(
                task_stir_tie,
                [t6_token, t6_method, t6_dtype, t6_base,
                 t6_model1, t6_weight1, t6_density1, t6_gamma1, t6_epsilon1,
                 t6_model2, t6_weight2, t6_density2, t6_gamma2, t6_epsilon2,
                 t6_model3, t6_weight3, t6_density3, t6_gamma3, t6_epsilon3,
                 t6_model4, t6_weight4, t6_density4, t6_gamma4, t6_epsilon4,
                 t6_normalize, t6_int8_mask, t6_lambda, t6_rescale, t6_select_topk,
                 t6_tokenizer_source, t6_out, t6_priv],
                t6_res
            )
            
        # TAB 7: SPECIOUS
        with gr.Tab("Specious"):
            gr.Markdown("### Specialized Methods: model_stock, nearswap, arcee_fusion, passthrough")
            
            t7_token = gr.Textbox(label="HF Token", type="password")
            
            with gr.Row():
                t7_method = gr.Dropdown(
                    ["model_stock", "nearswap", "arcee_fusion", "passthrough"],
                    value="model_stock",
                    label="Merge Method"
                )
                t7_dtype = gr.Dropdown(
                    ["float16", "bfloat16", "float32"],
                    value="bfloat16",
                    label="dtype"
                )
            
            t7_base = gr.Textbox(label="Base Model (required for nearswap/arcee_fusion/model_stock)", placeholder="org/base-model")
            
            gr.Markdown("#### Models")
            gr.Markdown("**passthrough:** 1 model | **nearswap/arcee_fusion:** 2 models | **model_stock:** 3+ models")
            
            with gr.Row():
                with gr.Column():
                    t7_model1 = gr.Textbox(label="Model 1", placeholder="org/model1")
                    t7_weight1 = gr.Textbox(label="Weight", value="1.0")
                with gr.Column():
                    t7_filter1 = gr.Textbox(label="Filter Model Component")
            
            with gr.Row():
                with gr.Column():
                    t7_model2 = gr.Textbox(label="Model 2 (optional)", placeholder="org/model2")
                    t7_weight2 = gr.Textbox(label="Weight", value="1.0")
                with gr.Column():
                    t7_filter2 = gr.Textbox(label="Filter Model Component")
            
            with gr.Row():
                t7_model3 = gr.Textbox(label="Model 3 (optional for model_stock)", placeholder="org/model3")
                t7_weight3 = gr.Textbox(label="Weight", value="1.0")
            
            with gr.Row():
                t7_model4 = gr.Textbox(label="Model 4 (optional)", placeholder="org/model4")
                t7_weight4 = gr.Textbox(label="Weight", value="1.0")
            
            with gr.Row():
                t7_model5 = gr.Textbox(label="Model 5 (optional)", placeholder="org/model5")
                t7_weight5 = gr.Textbox(label="Weight", value="1.0")
            
            with gr.Accordion("Parameters", open=False):
                t7_t = gr.Textbox(label="t (for nearswap)", value="0.5")
                t7_normalize = gr.Checkbox(label="Normalize", value=True)
                t7_int8_mask = gr.Checkbox(label="int8_mask", value=False)
                t7_filter_wise = gr.Checkbox(label="filter_wise (for model_stock)", value=False)
                t7_tokenizer_source = gr.Dropdown(["base", "union"], value="base", label="Tokenizer Source")
            
            t7_out = gr.Textbox(label="Output Repo")
            t7_priv = gr.Checkbox(label="Private", value=True)
            t7_btn = gr.Button("Run MergeKit Specialized")
            t7_res = gr.Textbox(label="Result")
            
            t7_btn.click(
                task_specious,
                [t7_token, t7_method, t7_dtype, t7_base,
                 t7_model1, t7_weight1, t7_filter1,
                 t7_model2, t7_weight2, t7_filter2,
                 t7_model3, t7_weight3,
                 t7_model4, t7_weight4,
                 t7_model5, t7_weight5,
                 t7_t, t7_normalize, t7_int8_mask, t7_filter_wise, t7_tokenizer_source,
                 t7_out, t7_priv],
                t7_res
            )
            
        # TAB 8: MOER
        with gr.Tab("MoEr"):
            gr.Markdown("### Mixture of Experts Construction")
            
            t8_token = gr.Textbox(label="HF Token", type="password")
            
            with gr.Row():
                t8_dtype = gr.Dropdown(
                    ["float16", "bfloat16", "float32"],
                    value="bfloat16",
                    label="dtype"
                )
                t8_gate = gr.Dropdown(
                    ["cheap_embed", "random", "hidden"],
                    value="cheap_embed",
                    label="Gate Mode"
                )
            
            t8_base = gr.Textbox(label="Base Model (required)", placeholder="org/base-model")
            
            gr.Markdown("#### Experts (at least 2 required)")
            gr.Markdown("Prompts are comma-separated descriptors for each expert")
            
            with gr.Row():
                t8_expert1 = gr.Textbox(label="Expert 1", placeholder="org/expert1")
                t8_prompt1 = gr.Textbox(label="Positive Prompts", placeholder="math, reasoning, logic")
            
            with gr.Row():
                t8_expert2 = gr.Textbox(label="Expert 2", placeholder="org/expert2")
                t8_prompt2 = gr.Textbox(label="Positive Prompts", placeholder="creative, writing, storytelling")
            
            with gr.Row():
                t8_expert3 = gr.Textbox(label="Expert 3 (optional)", placeholder="org/expert3")
                t8_prompt3 = gr.Textbox(label="Positive Prompts", placeholder="code, programming")
            
            with gr.Row():
                t8_expert4 = gr.Textbox(label="Expert 4 (optional)", placeholder="org/expert4")
                t8_prompt4 = gr.Textbox(label="Positive Prompts", placeholder="")
            
            with gr.Row():
                t8_expert5 = gr.Textbox(label="Expert 5 (optional)", placeholder="org/expert5")
                t8_prompt5 = gr.Textbox(label="Positive Prompts", placeholder="")
            
            with gr.Accordion("Parameters", open=False):
                t8_tokenizer_source = gr.Dropdown(["base", "union"], value="base", label="Tokenizer Source")
            
            t8_out = gr.Textbox(label="Output Repo")
            t8_priv = gr.Checkbox(label="Private", value=True)
            t8_btn = gr.Button("Build MoE")
            t8_res = gr.Textbox(label="Result")
            
            t8_btn.click(
                task_moer,
                [t8_token, t8_base, t8_dtype,
                 t8_expert1, t8_prompt1,
                 t8_expert2, t8_prompt2,
                 t8_expert3, t8_prompt3,
                 t8_expert4, t8_prompt4,
                 t8_expert5, t8_prompt5,
                 t8_gate, t8_tokenizer_source,
                 t8_out, t8_priv],
                t8_res
            )
            
        # TAB 9: RAWER
        with gr.Tab("Rawer"):
            gr.Markdown("### Raw PyTorch MergeKit / Non-pipeline-classed Models")
            
            t9_token = gr.Textbox(label="HF Token", type="password")
            
            t9_models = gr.TextArea(
                label="Models (one per line)",
                placeholder="org/model1\norg/model2\norg/model3",
                lines=5
            )
            
            with gr.Row():
                t9_method = gr.Dropdown(
                    ["linear", "passthrough"],
                    value="linear",
                    label="Method"
                )
                t9_dtype = gr.Dropdown(
                    ["float32", "float16", "bfloat16"],
                    value="float32",
                    label="dtype"
                )
            
            with gr.Accordion("Parameters", open=False):
                t9_tokenizer_source = gr.Dropdown(["base", "union"], value="base", label="Tokenizer Source")
            
            t9_out = gr.Textbox(label="Output Repo")
            t9_priv = gr.Checkbox(label="Private", value=True)
            t9_btn = gr.Button("Merge Raw")
            t9_res = gr.Textbox(label="Result")
            
            t9_btn.click(
                task_rawer,
                [t9_token, t9_models, t9_method, t9_dtype, t9_tokenizer_source, t9_out, t9_priv],
                t9_res
            )
            
        # TAB 10: MARIO,DARE!
        with gr.Tab("Mario,DARE!"):
            gr.Markdown("### Custom DARE Implementation (Drop And REscale)")
            gr.Markdown("From [sft-merger by Martyn Garcia](https://github.com/martyn)")
            
            t10_token = gr.Textbox(label="HF Token", type="password")
            
            with gr.Row():
                t10_base = gr.Textbox(label="Base Model", placeholder="org/base-model")
                t10_ft = gr.Textbox(label="Fine-Tuned Model", placeholder="org/fine-tuned-model")
            
            with gr.Row():
                t10_ratio = gr.Slider(
                    label="Merge Ratio (delta weight)",
                    value=1.0,
                    minimum=0.0,
                    maximum=2.0,
                    step=0.1
                )
                t10_mask = gr.Slider(
                    label="Mask Rate (drop probability)",
                    value=0.5,
                    minimum=0.0,
                    maximum=0.99,
                    step=0.01
                )
            
            gr.Markdown(
                """
                ### How DARE Works:
                1. **Compute Delta**: Difference between fine-tuned and base weights
                2. **Drop Elements**: Randomly mask out delta values based on mask rate
                3. **Rescale**: Compensate for dropped elements by rescaling remaining values
                4. **Apply**: Add scaled delta back to base model
                
                **Mask Rate**: 0.5 = drop 50% of delta values, 0.9 = drop 90% (more aggressive sparsification)
                """
            )
            
            t10_out = gr.Textbox(label="Output Repo")
            t10_priv = gr.Checkbox(label="Private", value=True)
            t10_btn = gr.Button("Run DARE Merge")
            t10_res = gr.Textbox(label="Result")
            
            t10_btn.click(
                task_dare_soonr,
                [t10_token, t10_base, t10_ft, t10_ratio, t10_mask, t10_out, t10_priv],
                t10_res
            )

if __name__ == "__main__":
    demo.queue().launch(css=css, ssr_mode=False, mcp_server=True)