applogsview2.py

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

# --- Essential Imports ---
from gradio_logsview import LogsView, LogsViewRunner
from mergekit.config import MergeConfiguration
try:
    from mergekit.config import MergeConfiguration
except ImportError:
    # Fallback if installation fails temporarily
    class MergeConfiguration:
        @staticmethod
        def model_validate(config): pass

# --- 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()

# =================================================================================
# SHARED HELPERS
# =================================================================================

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

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

# =================================================================================
# TABS 1-4 LOGIC (RESTORED)
# =================================================================================

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]
        start, end = metadata["data_offsets"]
        self.file.seek(self.header_size + 8 + start)
        return self._deserialize_tensor(self.file.read(end - start), metadata)
    def _read_header(self):
        header_size = struct.unpack("<Q", self.file.read(8))[0]
        return json.loads(self.file.read(header_size).decode("utf-8")), 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)

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

def task_merge_legacy(hf_token, base, sub, lora, scale, prec, shard, out, struct_s, priv, progress=gr.Progress()):
    cleanup_temp()
    if hf_token: login(hf_token.strip())
    try: api.create_repo(repo_id=out, private=priv, exist_ok=True, token=hf_token)
    except Exception as e: return f"Error: {e}"
    if struct_s:
        try:
            files = api.list_repo_files(repo_id=struct_s, token=hf_token)
            for f in tqdm(files, desc="Copying Structure"):
                if sub and f.startswith(sub): continue 
                if not sub and any(f.endswith(x) for x in ['.safetensors', '.bin', '.pt', '.pth']): continue
                l = hf_hub_download(repo_id=struct_s, filename=f, token=hf_token, local_dir=TempDir)
                api.upload_file(path_or_fileobj=l, path_in_repo=f, repo_id=out, token=hf_token)
        except: pass
    
    files = [f for f in list_repo_files(repo_id=base, token=hf_token) if f.endswith(".safetensors")]
    if sub: files = [f for f in files if f.startswith(sub)]
    if not files: return "No safetensors found"
    
    prefix = "diffusion_pytorch_model" if (sub in ["transformer", "unet"] or "diffusion_pytorch_model" in os.path.basename(files[0])) else "model"
    dtype = torch.bfloat16 if prec == "bf16" else torch.float16 if prec == "fp16" else torch.float32
    try: lora_pairs = load_lora_to_memory(download_lora_smart(lora, hf_token), dtype)
    except Exception as e: return f"LoRA Error: {e}"

    buf = ShardBuffer(shard, TempDir, out, sub, hf_token, prefix)
    for i, fpath in enumerate(files):
        local = hf_hub_download(repo_id=base, filename=fpath, token=hf_token, local_dir=TempDir)
        with MemoryEfficientSafeOpen(local) as f:
            for k in f.keys():
                v = f.get_tensor(k)
                stem = get_key_stem(k)
                match = lora_pairs.get(stem) or lora_pairs.get(stem.replace("to_q", "qkv")) or lora_pairs.get(stem.replace("to_k", "qkv")) or lora_pairs.get(stem.replace("to_v", "qkv"))
                if match:
                    d, u = match["down"], match["up"]
                    s = scale * (match["alpha"] / match["rank"])
                    if len(v.shape)==4 and len(d.shape)==2: d, u = d.unsqueeze(-1).unsqueeze(-1), u.unsqueeze(-1).unsqueeze(-1)
                    delta = (up.squeeze() @ down.squeeze()).reshape(up.shape[0], down.shape[1], 1, 1) if len(up.shape)==4 else u @ d
                    v = v.to(dtype).add_((delta * s).to(dtype))
                buf.add_tensor(k, v.to(dtype))
        os.remove(local)
    buf.flush()
    idx = {"metadata": {"total_size": buf.total_size}, "weight_map": buf.index_map}
    idx_n = f"{prefix}.safetensors.index.json"
    with open(TempDir/idx_n, "w") as f: json.dump(idx, f, indent=4)
    api.upload_file(path_or_fileobj=TempDir/idx_n, path_in_repo=f"{sub}/{idx_n}" if sub else idx_n, repo_id=out, token=hf_token)
    return "Done"

def task_extract(hf_token, org, tun, rank, out):
    cleanup_temp()
    if hf_token: login(hf_token.strip())
    try:
        p1 = download_lora_smart(org, hf_token)
        p2 = download_lora_smart(tun, hf_token)
        org_f, tun_f = MemoryEfficientSafeOpen(p1), MemoryEfficientSafeOpen(p2)
        lora_sd = {}
        common = set(org_f.keys()) & set(tun_f.keys())
        for k in tqdm(common, desc="Extracting"):
            if "num_batches_tracked" in k or "running_mean" in k or "running_var" in k: continue
            m1, m2 = org_f.get_tensor(k).float(), tun_f.get_tensor(k).float()
            if m1.shape != m2.shape: continue
            diff = m2 - m1
            if torch.max(torch.abs(diff)) < 1e-4: continue
            out_d, in_d = diff.shape[0], diff.shape[1] if len(diff.shape) > 1 else 1
            r = min(int(rank), in_d, out_d)
            if len(diff.shape)==4: diff = diff.flatten(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), 0.99)
            if hi_val > 0: U, Vh = U.clamp(-hi_val, hi_val), Vh.clamp(-hi_val, hi_val)
            if len(m1.shape)==4:
                U = U.reshape(out_d, r, 1, 1)
                Vh = Vh.reshape(r, in_d, m1.shape[2], m1.shape[3])
            else:
                U, Vh = U.reshape(out_d, r), Vh.reshape(r, in_d)
            stem = k.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_f = TempDir/"extracted.safetensors"
        save_file(lora_sd, out_f)
        api.create_repo(repo_id=out, exist_ok=True, token=hf_token)
        api.upload_file(path_or_fileobj=out_f, path_in_repo="extracted_lora.safetensors", repo_id=out, token=hf_token)
        return "Done"
    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: 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)
    except: return "Error parsing weights."
    if len(weights) < len(urls): weights += [1.0] * (len(urls) - len(weights))
    
    paths = []
    for url in tqdm(urls, desc="Downloading"): paths.append(download_lora_smart(url, hf_token))
    
    merged = {}
    if "Iterative EMA" in method:
        base_sd = load_file(paths[0], device="cpu")
        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].float() * current_beta + curr[k].float() * (1 - current_beta)
        merged = base_sd
    else:
        states = [load_full_state_dict(p) for p in paths]
        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, total_delta = 0.0, 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 = d.shape[0]
                        alpha = state[ak].item() if ak in state else rank
                        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"
    if merged: 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 & LOGSVIEW (TABS 5-9) - CLI LOGIC
# =================================================================================

def parse_weight(w_str):
    if not w_str.strip(): return 1.0
    try:
        if "[" in w_str: return yaml.safe_load(w_str)
        return float(w_str)
    except: return 1.0

def run_mergekit_logic(config_dict, token, out_repo, private, shard_size, output_precision, tokenizer_source, chat_template, program="mergekit-yaml"):
    runner = LogsViewRunner()
    cleanup_temp()
    
    if chat_template and chat_template.strip():
        config_dict["chat_template"] = chat_template.strip()
    
    try:
        if program != "mergekit-moe":
            MergeConfiguration.model_validate(config_dict)
    except Exception as e:
        yield runner.log(f"Invalid Config: {e}", level="ERROR")
        return

    if token: 
        login(token.strip())
        os.environ["HF_TOKEN"] = token.strip()
    
    if "dtype" not in config_dict: config_dict["dtype"] = output_precision
    if "tokenizer_source" not in config_dict and tokenizer_source != "base":
        config_dict["tokenizer_source"] = tokenizer_source

    config_path = TempDir / "config.yaml"
    with open(config_path, "w") as f: yaml.dump(config_dict, f, sort_keys=False)
    
    yield runner.log(f"Config saved to {config_path}")
    yield runner.log(f"YAML:\n{yaml.dump(config_dict, sort_keys=False)}")

    try:
        api.create_repo(repo_id=out_repo, private=private, exist_ok=True, token=token)
        yield runner.log(f"Repo {out_repo} ready.")
    except Exception as e:
        yield runner.log(f"Repo Error: {e}", level="ERROR")
        return

    out_path = TempDir / "merge_output"
    shard_arg = f"{int(float(shard_size) * 1024)}M"
    
    cmd = [
        program, 
        str(config_path), 
        str(out_path), 
        "--allow-crimes", 
        "--copy-tokenizer", 
        "--out-shard-size", shard_arg,
        "--lazy-unpickle"
    ]
    
    if torch.cuda.is_available():
        cmd.extend(["--cuda", "--low-cpu-memory"])
    
    yield runner.log(f"Executing: {' '.join(cmd)}")
    env = os.environ.copy()
    env["HF_HOME"] = str(TempDir / ".cache")
    
    yield from runner.run_command(cmd, env=env)
    
    if runner.exit_code != 0:
        yield runner.log("Merge failed.", level="ERROR")
        return

    yield runner.log(f"Uploading to {out_repo}...")
    yield from runner.run_python(api.upload_folder, repo_id=out_repo, folder_path=out_path)
    yield runner.log("Upload Complete!")

# --- UI Wrappers for Tabs 5-9 ---

def wrapper_amphinterpolative(token, method, base, t, norm, i8, flat, row, eps, m_iter, tol, m1, w1, m2, w2, m3, w3, m4, w4, m5, w5, out, priv, shard, prec, tok_src, chat_t):
    params = {"normalize": norm, "int8_mask": i8}
    if method in ["slerp", "nuslerp"]: params["t"] = float(t)
    if method == "nuslerp": params.update({"flatten": flat, "row_wise": row})
    if method == "multislerp": params["eps"] = float(eps)
    if method == "karcher": params.update({"max_iter": int(m_iter), "tol": float(tol)})
    
    config = {"merge_method": method}
    
    if method in ["slerp", "nuslerp"]:
        if not base.strip(): yield runner.log("Error: Base model required", level="ERROR"); return
        config["base_model"] = base.strip()
        sources = [{"model": m, "parameters": {"weight": parse_weight(w)}} for m, w in [(m1,w1), (m2,w2)] if m.strip()]
        config["slices"] = [{"sources": sources, "parameters": params}]
    else:
        if base.strip() and method == "multislerp": config["base_model"] = base.strip()
        models = [{"model": m, "parameters": {"weight": parse_weight(w)}} for m, w in [(m1,w1), (m2,w2), (m3,w3), (m4,w4), (m5,w5)] if m.strip()]
        config["models"] = models
        config["parameters"] = params

    yield from run_mergekit_logic(config, token, out, priv, shard, prec, tok_src, chat_t, program="mergekit-yaml")

def wrapper_stirtie(token, method, base, norm, i8, lamb, resc, topk, m1, w1, d1, g1, e1, m2, w2, d2, g2, e2, m3, w3, d3, g3, e3, m4, w4, d4, g4, e4, out, priv, shard, prec, tok_src, chat_t):
    models = []
    # Explicit loop over the 4 sets of model inputs
    for m, w, d, g, e in [
        (m1, w1, d1, g1, e1),
        (m2, w2, d2, g2, e2),
        (m3, w3, d3, g3, e3),
        (m4, w4, d4, g4, e4)
    ]:
        if not m.strip(): continue
        p = {"weight": parse_weight(w)}
        if method in ["ties", "dare_ties", "dare_linear", "breadcrumbs_ties"]: p["density"] = parse_weight(d)
        if "breadcrumbs" in method: p["gamma"] = float(g)
        if "della" in method: p["epsilon"] = float(e)
        models.append({"model": m, "parameters": p})
    
    g_params = {"normalize": norm, "int8_mask": i8}
    if method != "sce": g_params["lambda"] = float(lamb)
    if method == "dare_linear": g_params["rescale"] = resc
    if method == "sce": g_params["select_topk"] = float(topk)

    config = {
        "merge_method": method,
        "base_model": base.strip() if base.strip() else models[0]["model"],
        "parameters": g_params,
        "models": models
    }
    yield from run_mergekit_logic(config, token, out, priv, shard, prec, tok_src, chat_t, program="mergekit-yaml")

def wrapper_specious(token, method, base, norm, i8, t, filt_w, m1, w1, f1, m2, w2, m3, w3, m4, w4, m5, w5, out, priv, shard, prec, tok_src, chat_t):
    models = []
    if method == "passthrough":
        p = {"weight": parse_weight(w1)}
        if f1.strip(): p["filter"] = f1.strip()
        models.append({"model": m1, "parameters": p})
    else:
        models = [{"model": m, "parameters": {"weight": parse_weight(w)}} for m, w in [(m1,w1), (m2,w2), (m3,w3), (m4,w4), (m5,w5)] if m.strip()]

    config = {"merge_method": method, "parameters": {"normalize": norm, "int8_mask": i8}}
    if base.strip(): config["base_model"] = base.strip()
    if method == "nearswap": config["parameters"]["t"] = float(t)
    if method == "model_stock": config["parameters"]["filter_wise"] = filt_w
    config["models"] = models

    yield from run_mergekit_logic(config, token, out, priv, shard, prec, tok_src, chat_t, program="mergekit-yaml")

def wrapper_moer(token, base, experts, gate, dtype, out, priv, shard, prec, tok_src, chat_t):
    formatted = [{"source_model": e.strip(), "positive_prompts": ["chat", "assist"]} for e in experts.split('\n') if e.strip()]
    config = {
        "base_model": base.strip() if base.strip() else formatted[0]["source_model"],
        "gate_mode": gate,
        "dtype": dtype, 
        "experts": formatted
    }
    # Uses mergekit-moe CLI
    yield from run_mergekit_logic(config, token, out, priv, shard, prec, tok_src, chat_t, program="mergekit-moe")

def wrapper_rawer(token, models, method, dtype, out, priv, shard, prec, tok_src, chat_t):
    models_list = [{"model": m.strip(), "parameters": {"weight": 1.0}} for m in models.split('\n') if m.strip()]
    config = {
        "models": models_list,
        "merge_method": method,
        "dtype": dtype
    }
    yield from run_mergekit_logic(config, token, out, priv, shard, prec, tok_src, chat_t, program="mergekit-yaml")

# --- TAB 10 (Custom DARE) Logic ---
def task_dare_custom(token, base, ft, ratio, mask, out, priv):
    cleanup_temp()
    if token: login(token.strip())
    try:
        b_path = download_lora_smart(base, token) 
        f_path = download_lora_smart(ft, token)
        b_sd = load_file(b_path, device="cpu")
        f_sd = load_file(f_path, device="cpu")
        merged = {}
        common = set(b_sd.keys()) & set(f_sd.keys())
        for k in tqdm(common, desc="Merging"):
            tb, tf = b_sd[k], f_sd[k]
            if tb.shape != tf.shape: 
                merged[k] = tf
                continue
            delta = tf.float() - tb.float()
            if mask > 0:
                m = torch.bernoulli(torch.full_like(delta, 1.0 - mask))
                delta = (delta * m) / (1.0 - mask)
            merged[k] = (tb.float() + ratio * delta).to(tb.dtype)
        
        out_f = TempDir / "model.safetensors"
        save_file(merged, out_f)
        api.create_repo(repo_id=out, private=priv, exist_ok=True, token=token)
        api.upload_file(path_or_fileobj=out_f, path_in_repo="model.safetensors", repo_id=out, token=token)
        return f"Done! {out}"
    except Exception as e: return str(e)

# =================================================================================
# UI GENERATION
# =================================================================================

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

with gr.Blocks() as demo:
    gr.HTML("""<h1><img src="https://huggingface.co/spaces/AlekseyCalvin/Soon_Merger/resolve/main/SMerger3.png" alt="SOONmerge®"> Transform Transformers for FREE!</h1>""")
    gr.Markdown("# 🧰Training-Free CPU-run Model Creation Toolkit")
    
    with gr.Tabs():
        # --- TAB 1: RESTORED ---
        with gr.Tab("Merge to Base Model + Reshard Output"):
            t1_token = gr.Textbox(label="Token", type="password")
            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)
            gr.Button("Merge").click(task_merge_legacy, [t1_token, t1_base, t1_sub, t1_lora, t1_scale, t1_prec, t1_shard, t1_out, t1_struct, t1_priv], gr.Textbox(label="Result"))

        # --- TAB 2: RESTORED ---
        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")
            gr.Button("Extract").click(task_extract, [t2_token, t2_org, t2_tun, t2_rank, t2_out], gr.Textbox(label="Result"))

        # --- TAB 3: RESTORED ---
        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)")
            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")
                t3_rank = gr.Number(label="Target Rank – For SVD only", value=128)
            with gr.Row():
                t3_beta = gr.Slider(label="Beta – for linear/post-hoc EMA", value=0.95, minimum=0.01, maximum=1.00)
                t3_sigma = gr.Slider(label="Sigma Rel – for linear/post-hoc EMA", value=0.21, minimum=0.01, maximum=1.00)
            t3_out = gr.Textbox(label="Output Repo")
            t3_priv = gr.Checkbox(label="Private Output", value=True)
            gr.Button("Merge").click(task_merge_adapters_advanced, [t3_token, t3_urls, t3_method, t3_weights, t3_beta, t3_sigma, t3_rank, t3_out, t3_priv], gr.Textbox(label="Result"))

        # --- TAB 4: RESTORED ---
        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)
                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\nThese methods intelligently determine the best rank per layer.\n- **sv_ratio (Relative Strength):** Keeps features that are at least `1/Param` as strong as the main feature. **Param must be >= 2**.\n- **sv_fro (Visual Information Density):** Preserves `Param%` of total information content. **Param between 0.0 and 1.0**.\n- **sv_cumulative (Cumulative Sum):** Preserves weights that sum up to `Param%` of total strength. **Param between 0.0 and 1.0**.\n- **⚠️ Safety Ceiling:** The **'To Rank'** slider acts as a hard limit.")
            t4_out = gr.Textbox(label="Output")
            gr.Button("Resize").click(task_resize, [t4_token, t4_in, t4_rank, t4_method, t4_param, t4_out], gr.Textbox(label="Result"))

        # --- TAB 5 ---
        with gr.Tab("Amphinterpolative"):
            gr.Markdown("### Spherical Interpolation Family")
            t5_token = gr.Textbox(label="HF Token", type="password")
            t5_method = gr.Dropdown(["slerp", "nuslerp", "multislerp", "karcher"], value="slerp", label="Method")
            with gr.Row():
                t5_shard = gr.Slider(label="Max Shard Size (GB)", value=5.0, minimum=1.0, maximum=20.0)
                t5_prec = gr.Dropdown(["float16", "bfloat16", "float32"], value="bfloat16", label="Output Precision")
                t5_tok = gr.Dropdown(["base", "union", "model:path"], value="base", label="Tokenizer Source")
                t5_chat = gr.Textbox(label="Chat Template (write-in, default: auto)", placeholder="auto")
            with gr.Row():
                t5_base = gr.Textbox(label="Base Model")
                t5_t = gr.Slider(0, 1, 0.5, label="t")
            with gr.Row():
                t5_norm = gr.Checkbox(label="Normalize", value=True); t5_i8 = gr.Checkbox(label="Int8 Mask", value=False); t5_flat = gr.Checkbox(label="NuSlerp Flatten", value=False); t5_row = gr.Checkbox(label="NuSlerp Row Wise", value=False)
            with gr.Row():
                t5_eps = gr.Textbox(label="eps", value="1e-8"); t5_iter = gr.Number(label="max_iter", value=10); t5_tol = gr.Textbox(label="tol", value="1e-5")
            m1, w1 = gr.Textbox(label="Model 1"), gr.Textbox(label="Weight 1", value="1.0"); m2, w2 = gr.Textbox(label="Model 2"), gr.Textbox(label="Weight 2", value="1.0")
            with gr.Accordion("More", open=False):
                m3, w3 = gr.Textbox(label="Model 3"), gr.Textbox(label="Weight 3", value="1.0"); m4, w4 = gr.Textbox(label="Model 4"), gr.Textbox(label="Weight 4", value="1.0"); m5, w5 = gr.Textbox(label="Model 5"), gr.Textbox(label="Weight 5", value="1.0")
            t5_out = gr.Textbox(label="Output Repo"); t5_priv = gr.Checkbox(label="Private", value=True)
            t5_btn = gr.Button("Execute")
            t5_logs = LogsView()
            t5_btn.click(wrapper_amphinterpolative, [t5_token, t5_method, t5_base, t5_t, t5_norm, t5_i8, t5_flat, t5_row, t5_eps, t5_iter, t5_tol, m1, w1, m2, w2, m3, w3, m4, w4, m5, w5, t5_out, t5_priv, t5_shard, t5_prec, t5_tok, t5_chat], t5_logs)

        # --- TAB 6 ---
        with gr.Tab("Stir/Tie Bases"):
            gr.Markdown("### Task Vector Family")
            t6_token = gr.Textbox(label="Token", type="password")
            t6_method = gr.Dropdown(["task_arithmetic", "ties", "dare_ties", "dare_linear", "della", "della_linear", "breadcrumbs", "breadcrumbs_ties", "sce"], value="ties", label="Method")
            with gr.Row():
                t6_shard = gr.Slider(label="Max Shard Size (GB)", value=5.0, minimum=1.0, maximum=20.0); t6_prec = gr.Dropdown(["float16", "bfloat16", "float32"], value="bfloat16", label="Output Precision"); t6_tok = gr.Dropdown(["base", "union", "model:path"], value="base", label="Tokenizer Source"); t6_chat = gr.Textbox(label="Chat Template", placeholder="auto")
            t6_base = gr.Textbox(label="Base Model")
            with gr.Row():
                t6_norm = gr.Checkbox(label="Normalize", value=True); t6_i8 = gr.Checkbox(label="Int8 Mask", value=False); t6_resc = gr.Checkbox(label="Rescale", value=True); t6_lamb = gr.Number(label="Lambda", value=1.0); t6_topk = gr.Slider(0, 1, 1.0, label="Select TopK")
            m1_6, w1_6 = gr.Textbox(label="Model 1"), gr.Textbox(label="Weight 1", value="1.0"); d1_6, g1_6, e1_6 = gr.Textbox(label="Density", value="1.0"), gr.Number(label="Gamma", value=0.01), gr.Number(label="Epsilon", value=0.15)
            with gr.Accordion("More", open=False):
                m2_6, w2_6 = gr.Textbox(label="Model 2"), gr.Textbox(label="Weight 2", value="1.0"); d2_6, g2_6, e2_6 = gr.Textbox(label="Density", value="1.0"), gr.Number(label="Gamma", value=0.01), gr.Number(label="Epsilon", value=0.15)
                # Corrected UI Layout: Added missing textboxes for Models 3 & 4
                m3_6, w3_6 = gr.Textbox(label="Model 3"), gr.Textbox(label="Weight 3", value="1.0"); d3_6, g3_6, e3_6 = gr.Textbox(label="Density", value="1.0"), gr.Number(label="Gamma", value=0.01), gr.Number(label="Epsilon", value=0.15)
                m4_6, w4_6 = gr.Textbox(label="Model 4"), gr.Textbox(label="Weight 4", value="1.0"); d4_6, g4_6, e4_6 = gr.Textbox(label="Density", value="1.0"), gr.Number(label="Gamma", value=0.01), gr.Number(label="Epsilon", value=0.15)
            t6_out = gr.Textbox(label="Output Repo"); t6_priv = gr.Checkbox(label="Private", value=True)
            t6_btn = gr.Button("Execute")
            t6_logs = LogsView()
            # The .click list now matches the wrapper function signature perfectly (34 arguments)
            t6_btn.click(wrapper_stirtie, [t6_token, t6_method, t6_base, t6_norm, t6_i8, t6_lamb, t6_resc, t6_topk, m1_6, w1_6, d1_6, g1_6, e1_6, m2_6, w2_6, d2_6, g2_6, e2_6, m3_6, w3_6, d3_6, g3_6, e3_6, m4_6, w4_6, d4_6, g4_6, e4_6, t6_out, t6_priv, t6_shard, t6_prec, t6_tok, t6_chat], t6_logs)

        # --- TAB 7 ---
        with gr.Tab("Specious"):
            gr.Markdown("### Specialized Methods")
            t7_token = gr.Textbox(label="Token", type="password")
            t7_method = gr.Dropdown(["model_stock", "nearswap", "arcee_fusion", "passthrough", "linear"], value="model_stock", label="Method")
            with gr.Row():
                t7_shard = gr.Slider(label="Max Shard Size (GB)", value=5.0, minimum=1.0, maximum=20.0); t7_prec = gr.Dropdown(["float16", "bfloat16", "float32"], value="bfloat16", label="Output Precision"); t7_tok = gr.Dropdown(["base", "union", "model:path"], value="base", label="Tokenizer Source"); t7_chat = gr.Textbox(label="Chat Template", placeholder="auto")
            t7_base = gr.Textbox(label="Base Model")
            with gr.Row():
                t7_norm = gr.Checkbox(label="Normalize", value=True); t7_i8 = gr.Checkbox(label="Int8 Mask", value=False); t7_t = gr.Slider(0, 1, 0.5, label="t"); t7_filt_w = gr.Checkbox(label="Filter Wise", value=False)
            m1_7, w1_7, f1_7 = gr.Textbox(label="Model 1"), gr.Textbox(label="Weight 1", value="1.0"), gr.Textbox(label="Filter (Passthrough)")
            m2_7, w2_7 = gr.Textbox(label="Model 2"), gr.Textbox(label="Weight 2", value="1.0")
            with gr.Accordion("More", open=False):
                m3_7, w3_7 = gr.Textbox(label="Model 3"), gr.Textbox(label="Weight 3", value="1.0"); m4_7, w4_7 = gr.Textbox(label="Model 4"), gr.Textbox(label="Weight 4", value="1.0"); m5_7, w5_7 = gr.Textbox(label="Model 5"), gr.Textbox(label="Weight 5", value="1.0")
            t7_out = gr.Textbox(label="Output Repo"); t7_priv = gr.Checkbox(label="Private", value=True)
            t7_btn = gr.Button("Execute")
            t7_logs = LogsView()
            t7_btn.click(wrapper_specious, [t7_token, t7_method, t7_base, t7_norm, t7_i8, t7_t, t7_filt_w, m1_7, w1_7, f1_7, m2_7, w2_7, m3_7, w3_7, m4_7, w4_7, m5_7, w5_7, t7_out, t7_priv, t7_shard, t7_prec, t7_tok, t7_chat], t7_logs)

        # --- TAB 8 (MoEr) ---
        with gr.Tab("MoEr"):
            gr.Markdown("### Mixture of Experts")
            t8_token = gr.Textbox(label="Token", type="password")
            with gr.Row():
                t8_shard = gr.Slider(label="Max Shard Size (GB)", value=5.0, minimum=1.0, maximum=20.0); t8_prec = gr.Dropdown(["float16", "bfloat16", "float32"], value="bfloat16", label="Output Precision"); t8_tok = gr.Dropdown(["base", "union", "model:path"], value="base", label="Tokenizer Source"); t8_chat = gr.Textbox(label="Chat Template", placeholder="auto")
            t8_base = gr.Textbox(label="Base Model"); t8_experts = gr.TextArea(label="Experts List"); t8_gate = gr.Dropdown(["cheap_embed", "random", "hidden"], value="cheap_embed", label="Gate Mode"); t8_dtype = gr.Dropdown(["float16", "bfloat16"], value="bfloat16", label="Internal Dtype")
            t8_out = gr.Textbox(label="Output Repo"); t8_priv = gr.Checkbox(label="Private", value=True)
            t8_btn = gr.Button("Build MoE")
            t8_logs = LogsView()
            t8_btn.click(wrapper_moer, [t8_token, t8_base, t8_experts, t8_gate, t8_dtype, t8_out, t8_priv, t8_shard, t8_prec, t8_tok, t8_chat], t8_logs)

        # --- TAB 9 (Rawer) ---
        with gr.Tab("Rawer"):
            gr.Markdown("### Raw PyTorch / Non-Transformer")
            t9_token = gr.Textbox(label="Token", type="password"); t9_models = gr.TextArea(label="Models (one per line)")
            with gr.Row():
                t9_shard = gr.Slider(label="Max Shard Size (GB)", value=5.0, minimum=1.0, maximum=20.0); t9_prec = gr.Dropdown(["float16", "bfloat16", "float32"], value="bfloat16", label="Output Precision"); t9_tok = gr.Dropdown(["base", "union", "model:path"], value="base", label="Tokenizer Source"); t9_chat = gr.Textbox(label="Chat Template", placeholder="auto")
            t9_method = gr.Dropdown(["linear", "passthrough"], value="linear", label="Method"); t9_dtype = gr.Dropdown(["float32", "float16", "bfloat16"], value="float32", label="Config Dtype")
            t9_out = gr.Textbox(label="Output Repo"); t9_priv = gr.Checkbox(label="Private", value=True)
            t9_btn = gr.Button("Merge Raw")
            t9_logs = LogsView()
            t9_btn.click(wrapper_rawer, [t9_token, t9_models, t9_method, t9_dtype, t9_out, t9_priv, t9_shard, t9_prec, t9_tok, t9_chat], t9_logs)

        # --- TAB 10 ---
        with gr.Tab("Mario,DARE!"):
            t10_token = gr.Textbox(label="Token", type="password")
            with gr.Row():
                t10_base = gr.Textbox(label="Base Model"); t10_ft = gr.Textbox(label="Fine-Tuned Model")
            with gr.Row():
                t10_ratio = gr.Slider(0, 5, 1.0, label="Ratio"); t10_mask = gr.Slider(0, 0.99, 0.5, label="Mask Rate")
            t10_out = gr.Textbox(label="Output Repo"); t10_priv = gr.Checkbox(label="Private", value=True)
            gr.Button("Run").click(task_dare_custom, [t10_token, t10_base, t10_ft, t10_ratio, t10_mask, t10_out, t10_priv], gr.Textbox(label="Result"))

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