Update app.py

app.py

@@ -450,6 +450,113 @@ def task_merge(hf_token, base_repo, base_subfolder, lora_input, scale, precision
     cleanup_temp()
     return f"Done! Merged {buffer.shard_count} shards to {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()
+    if hf_token: login(hf_token.strip())
+    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)
+# =================================================================================
+
+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()
+    if hf_token: login(hf_token.strip())
+    
+    urls = [u.strip() for u in lora_urls.split(",") if u.strip()]
+    paths = []
+    try:
+        for i, url in enumerate(urls):
+            paths.append(download_lora_smart(url, hf_token))
+    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 (CPU Optimized)
 # =================================================================================