Create app.py

app.py

@@ -0,0 +1,656 @@
+import gradio as gr
+import torch
+import os
+import gc
+import shutil
+import requests
+import json
+import struct
+import numpy as np
+import re
+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:
+    """
+    Reads safetensors metadata and tensors without mmap, keeping RAM usage low.
+    """
+    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 in the file")
+        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 download_file(input_path, token, filename=None):
+    local_path = TempDir / (filename if filename else "model.safetensors")
+    if input_path.startswith("http"):
+        print(f"Downloading {filename} from URL...")
+        try:
+            response = requests.get(input_path, stream=True, timeout=30)
+            response.raise_for_status()
+            with open(local_path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+        except Exception as e: raise ValueError(f"Download failed: {e}")
+    else:
+        print(f"Downloading {filename} from Hub...")
+        if not filename:
+            try:
+                files = list_repo_files(repo_id=input_path, token=token)
+                safetensors = [f for f in files if f.endswith(".safetensors")]
+                filename = safetensors[0] if safetensors else "adapter_model.safetensors"
+            except: filename = "adapter_model.safetensors"
+        
+        try:
+            hf_hub_download(repo_id=input_path, filename=filename, token=token, local_dir=TempDir, local_dir_use_symlinks=False)
+            if not (TempDir / filename).exists():
+                found = list(TempDir.rglob(filename))
+                if found: shutil.move(found[0], local_path)
+        except Exception as e: raise ValueError(f"Hub download failed: {e}")
+    
+    return local_path
+
+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: MERGE & RESHARD (Fixes Folder Structure & Aux Files)
+# =================================================================================
+
+def load_lora_to_memory(lora_path, precision_dtype=torch.bfloat16):
+    print(f"Loading LoRA from {lora_path}...")
+    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  # Dynamic prefix (e.g. 'diffusion_pytorch_model' or 'model')
+        self.buffer = []
+        self.current_bytes = 0
+        self.shard_count = 0
+        self.index_map = {}
+        self.total_model_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_model_size += size
+        
+        if self.current_bytes >= self.max_bytes:
+            self.flush()
+            
+    def flush(self):
+        if not self.buffer: return
+        self.shard_count += 1
+        
+        # ADAPTIVE NAMING: Uses the prefix detected from the base model
+        filename = f"{self.filename_prefix}-{self.shard_count:05d}.safetensors"
+        
+        # Proper Subfolder Handling
+        path_in_repo = f"{self.subfolder}/{filename}" if self.subfolder else filename
+        
+        print(f"Flushing {path_in_repo} ({self.current_bytes / 1024**3:.2f} GB)...")
+        
+        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"])
+                
+        print(f"Uploading {path_in_repo}...")
+        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="transformer"):
+    """
+    Copies files one-by-one from source to dest, skipping 'ignore_prefix'.
+    Does NOT skip .safetensors/.bin if they are outside the ignore folder.
+    """
+    print(f"Scanning {src_repo} for auxiliary files...")
+    try:
+        files = api.list_repo_files(repo_id=src_repo, token=token)
+        
+        for f in tqdm(files, desc="Copying Structure"):
+            # 1. Skip the folder we are replacing (e.g., transformer/)
+            if ignore_prefix and f.startswith(ignore_prefix):
+                continue
+            
+            # 2. Skip hidden/system files
+            if f.startswith("."):
+                continue
+                
+            # 3. Download -> Upload -> Delete loop
+            # This ensures we get VAE/TextEnc weights without disk overflow
+            try:
+                print(f"Copying {f}...")
+                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 Exception as e:
+                print(f"Failed to copy {f}: {e}")
+                
+    except Exception as e:
+        print(f"Structure cloning error: {e}")
+
+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()
+    login(hf_token)
+    
+    # 1. Output Setup
+    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}"
+    
+    # 2. Server-Side Structure Clone
+    if structure_repo:
+        ignore = base_subfolder if base_subfolder else None
+        streaming_copy_structure(hf_token, structure_repo, output_repo, ignore)
+
+    # 3. Load LoRA
+    dtype = torch.bfloat16 if precision == "bf16" else torch.float16 if precision == "fp16" else torch.float32
+    try:
+        progress(0.1, desc="Downloading LoRA...")
+        lora_path = download_file(lora_input, hf_token, filename="adapter.safetensors")
+        lora_pairs = load_lora_to_memory(lora_path, precision_dtype=dtype)
+    except Exception as e: return f"Error loading LoRA: {e}"
+
+    # 4. Stream Process
+    progress(0.2, desc="Fetching File List...")
+    files = list_repo_files(repo_id=base_repo, token=hf_token)
+    
+    # Identify valid shards in the target folder
+    input_shards = []
+    for f in files:
+        if not f.endswith(".safetensors"): continue
+        if base_subfolder and not f.startswith(base_subfolder): continue
+        input_shards.append(f)
+        
+    if not input_shards: return "No base safetensors found in specified location."
+    
+    input_shards.sort()
+    
+    # --- AUTO-DETECT NAMING CONVENTION ---
+    # We look at the first file to decide the naming scheme.
+    # Common schemes: 
+    #   "diffusion_pytorch_model-00001..." -> prefix: "diffusion_pytorch_model"
+    #   "model-00001..." -> prefix: "model"
+    #   "model.safetensors" -> prefix: "model"
+    
+    first_file = os.path.basename(input_shards[0])
+    
+    if first_file.startswith("diffusion_pytorch_model"):
+        filename_prefix = "diffusion_pytorch_model"
+        index_filename = "diffusion_pytorch_model.safetensors.index.json"
+    else:
+        # Default for LLMs, Text Encoders, etc.
+        filename_prefix = "model"
+        index_filename = "model.safetensors.index.json"
+        
+    print(f"Detected naming convention: {filename_prefix} (Index: {index_filename})")
+
+    # Initialize Buffer with detected prefix
+    buffer = ShardBuffer(shard_size, TempDir, output_repo, base_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 {shard_file}")
+        
+        local_shard = hf_hub_download(repo_id=base_repo, filename=shard_file, token=hf_token, local_dir=TempDir)
+        
+        with MemoryEfficientSafeOpen(local_shard) as f:
+            keys = f.keys()
+            for k in keys:
+                v = f.get_tensor(k)
+                base_stem = get_key_stem(k)
+                lora_keys = set(lora_pairs.keys())
+                match = None
+                
+                # Matching Logic (Exact + Heuristic for QKV)
+                if base_stem in lora_keys:
+                    match = lora_pairs[base_stem]
+                else:
+                    if "to_q" in base_stem:
+                        qkv_stem = base_stem.replace("to_q", "qkv")
+                        if qkv_stem in lora_keys: match = lora_pairs[qkv_stem]
+                    elif "to_k" in base_stem:
+                        qkv_stem = base_stem.replace("to_k", "qkv")
+                        if qkv_stem in lora_keys: match = lora_pairs[qkv_stem]
+                    elif "to_v" in base_stem:
+                        qkv_stem = base_stem.replace("to_v", "qkv")
+                        if qkv_stem in lora_keys: match = lora_pairs[qkv_stem]
+
+                if match and "down" in match and "up" in match:
+                    down = match["down"]
+                    up = match["up"]
+                    alpha = match["alpha"]
+                    rank = match["rank"]
+                    scaling = scale * (alpha / rank)
+                    
+                    # Handle Conv 1x1 squeeze
+                    if len(v.shape) == 4 and len(down.shape) == 2:
+                        down = down.unsqueeze(-1).unsqueeze(-1)
+                        up = up.unsqueeze(-1).unsqueeze(-1)
+                    
+                    try:
+                        if len(up.shape) == 4:
+                            delta = (up.squeeze() @ down.squeeze()).reshape(up.shape[0], down.shape[1], 1, 1)
+                        else:
+                            delta = up @ down
+                    except:
+                        delta = up.T @ down 
+                    
+                    delta = delta * scaling
+                    valid_delta = True
+                    
+                    # Shape Slicing Logic
+                    if delta.shape == v.shape:
+                        pass
+                    elif delta.shape[0] == v.shape[0] * 3:
+                        chunk = v.shape[0]
+                        if "to_q" in k: delta = delta[0:chunk, ...]
+                        elif "to_k" in k: delta = delta[chunk:2*chunk, ...]
+                        elif "to_v" in k: delta = delta[2*chunk:, ...]
+                        else: valid_delta = False
+                    elif delta.numel() == v.numel():
+                        delta = delta.reshape(v.shape)
+                    else:
+                        valid_delta = False
+                        
+                    if valid_delta:
+                        v = v.to(dtype)
+                        delta = delta.to(dtype)
+                        v.add_(delta)
+                        del delta
+
+                if v.dtype != dtype: v = v.to(dtype)
+                buffer.add_tensor(k, v)
+                del v
+        
+        os.remove(local_shard)
+        gc.collect()
+
+    buffer.flush()
+    
+    # Upload Index (Using the dynamically determined index filename)
+    print(f"Uploading Index: {index_filename}")
+    
+    index_data = {
+        "metadata": {"total_size": buffer.total_model_size}, 
+        "weight_map": buffer.index_map
+    }
+    
+    with open(TempDir / index_filename, "w") as f:
+        json.dump(index_data, f, indent=4)
+        
+    path_in_repo = f"{base_subfolder}/{index_filename}" if base_subfolder else index_filename
+    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 into {buffer.shard_count} shards at {output_repo}"
+
+# =================================================================================
+# TAB 2: EXTRACT LORA
+# =================================================================================
+
+def extract_lora_layer_by_layer(model_org, model_tuned, rank, clamp):
+    org = MemoryEfficientSafeOpen(model_org)
+    tuned = MemoryEfficientSafeOpen(model_tuned)
+    lora_sd = {}
+    print("Calculating diffs...")
+    for key in tqdm(org.keys()):
+        if key not in tuned.keys(): continue
+        mat_org = org.get_tensor(key).float()
+        mat_tuned = tuned.get_tensor(key).float()
+        diff = mat_tuned - mat_org
+        if torch.max(torch.abs(diff)) < 1e-4: continue
+        
+        out_dim, in_dim = diff.shape[:2]
+        r = min(rank, in_dim, out_dim)
+        is_conv = len(diff.shape) == 4
+        if is_conv: diff = diff.flatten(start_dim=1)
+            
+        try:
+            U, S, Vh = torch.linalg.svd(diff, full_matrices=False)
+            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(dist, clamp)
+            U = U.clamp(-hi_val, hi_val)
+            Vh = 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
+            lora_sd[f"{stem}.lora_down.weight"] = Vh
+            lora_sd[f"{stem}.alpha"] = torch.tensor(r).float()
+        except: pass
+    out = TempDir / "extracted.safetensors"
+    save_file(lora_sd, out)
+    return str(out)
+
+def task_extract(hf_token, org, tun, rank, out):
+    cleanup_temp()
+    login(hf_token)
+    try:
+        p1 = download_file(org, hf_token, filename="org.safetensors")
+        p2 = download_file(tun, hf_token, filename="tun.safetensors")
+        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.safetensors", repo_id=out, token=hf_token)
+        return "Done"
+    except Exception as e: return f"Error: {e}"
+
+# =================================================================================
+# TAB 3: MERGE ADAPTERS (EMA) with Sigma Rel
+# =================================================================================
+
+def sigma_rel_to_gamma(sigma_rel):
+    t = sigma_rel**-2
+    coeffs = [1, 7, 16 - t, 12 - t]
+    roots = np.roots(coeffs)
+    gamma = roots[np.isreal(roots) & (roots.real >= 0)].real.max()
+    return gamma
+
+def task_merge_adapters(hf_token, lora_urls, beta, sigma_rel, out_repo):
+    cleanup_temp()
+    login(hf_token)
+    
+    urls = [u.strip() for u in lora_urls.split(",") if u.strip()]
+    paths = []
+    try:
+        for i, url in enumerate(urls):
+            paths.append(download_file(url, hf_token, filename=f"a_{i}.safetensors"))
+    except Exception as e: return f"Download Error: {e}"
+        
+    if not paths: return "No models found"
+    
+    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:
+        gamma = sigma_rel_to_gamma(sigma_rel)
+        
+    for i, path in enumerate(paths[1:]):
+        print(f"Merging {path}")
+        if gamma is not None:
+            t = i + 1
+            current_beta = (1 - 1 / t) ** (gamma + 1)
+        else:
+            current_beta = 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)
+                
+    out = TempDir / "merged_adapters.safetensors"
+    save_file(base_sd, out)
+    api.create_repo(repo_id=out_repo, 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 "Done"
+
+# =================================================================================
+# TAB 4: RESIZE
+# =================================================================================
+
+def index_sv_ratio(S, target):
+    max_sv = S[0]
+    min_sv = max_sv / target
+    index = int(torch.sum(S > min_sv).item())
+    return max(1, min(index, len(S) - 1))
+
+def task_resize(hf_token, lora_input, new_rank, dynamic_method, dynamic_param, out_repo):
+    cleanup_temp()
+    login(hf_token)
+    try:
+        path = download_file(lora_input, hf_token)
+    except Exception as e: return f"Error: {e}"
+
+    state = load_file(path, device="cpu")
+    new_state = {}
+    
+    groups = {}
+    for k in state:
+        stem = get_key_stem(k)
+        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]
+
+    for stem, g in tqdm(groups.items()):
+        if "down" in g and "up" in g:
+            down, up = g["down"].float(), g["up"].float()
+            
+            if len(down.shape) == 4:
+                merged = (up.squeeze() @ down.squeeze()).reshape(up.shape[0], down.shape[1], down.shape[2], down.shape[3])
+                flat = merged.flatten(1)
+            else:
+                merged = up @ down
+                flat = merged
+            
+            U, S, Vh = torch.linalg.svd(flat, full_matrices=False)
+            
+            target_rank = int(new_rank)
+            if dynamic_method == "sv_ratio":
+                target_rank = index_sv_ratio(S, dynamic_param)
+            
+            target_rank = min(target_rank, S.shape[0])
+            
+            U = U[:, :target_rank]
+            S = S[:target_rank]
+            U = U @ torch.diag(S)
+            Vh = Vh[:target_rank, :]
+            
+            if len(down.shape) == 4:
+                U = U.reshape(up.shape[0], target_rank, 1, 1)
+                Vh = Vh.reshape(target_rank, down.shape[1], down.shape[2], down.shape[3])
+                
+            new_state[f"{stem}.lora_down.weight"] = Vh
+            new_state[f"{stem}.lora_up.weight"] = U
+            new_state[f"{stem}.alpha"] = torch.tensor(target_rank).float()
+
+    out = TempDir / "resized.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="resized.safetensors", repo_id=out_repo, token=hf_token)
+    return "Done"
+
+# =================================================================================
+# UI
+# =================================================================================
+
+css = ".container { max-width: 900px; margin: auto; }"
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 🧰SOONmerge® LoRA Toolkit")
+    
+    with gr.Tabs():
+        with gr.Tab("Merge to Base + Reshard Output"):
+            t1_token = gr.Textbox(label="Token", type="password")
+            t1_base = gr.Textbox(label="Base Repo (Diffusers)", value="ostris/Z-Image-De-Turbo")
+            t1_sub = gr.Textbox(label="Subfolder", value="transformer")
+            t1_lora = gr.Textbox(label="LoRA Direct Link", 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="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="Diffusers Extras (Copies VAE/TextEnc/etc)", value="Tongyi-MAI/Z-Image-Turbo")
+            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)
+
+        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 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)
+
+        with gr.Tab("Merge Multiple Adapters"):
+            t3_token = gr.Textbox(label="Token", type="password")
+            t3_urls = gr.Textbox(label="URLs")
+            with gr.Row():
+                t3_beta = gr.Slider(label="Beta", value=0.95, minimum=0.01, maximum=1.00, step=0.01)
+                t3_sigma = gr.Slider(label="Sigma Rel (Overrides Beta)", value=0.21, minimum=0.01, maximum=1.00, step=0.01)
+            t3_out = gr.Textbox(label="Output Repo")
+            t3_btn = gr.Button("Merge")
+            t3_res = gr.Textbox(label="Result")
+            t3_btn.click(task_merge_adapters, [t3_token, t3_urls, t3_beta, t3_sigma, t3_out], t3_res)
+
+        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 (Lower Only!)", value=8, minimum=1, maximum=256, step=1)
+                t4_method = gr.Dropdown(["None", "sv_ratio"], value="None", label="Dynamic Method")
+                t4_param = gr.Number(label="Dynamic Param", value=4.0)
+            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)
+
+if __name__ == "__main__":
+    demo.queue().launch(css=css, ssr_mode=False)