Update app.py

app.py

@@ -608,9 +608,22 @@ def task_extract(hf_token, org, tun, rank, out):
     except Exception as e: return f"Error: {e}"
 
 # =================================================================================
-# TAB 3: MERGE ADAPTERS (EMA)
+# TAB 3: MERGE ADAPTERS (Multi-Method)
 # =================================================================================
 
+def load_full_state_dict(path):
+    """Loads a safetensor file and cleans keys for easier processing."""
+    raw = load_file(path, device="cpu")
+    cleaned = {}
+    for k, v in raw.items():
+        # Map common keys to standard "lora_up/lora_down"
+        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
+
+# --- Original EMA Method ---
 def sigma_rel_to_gamma(sigma_rel):
     t = sigma_rel**-2
     coeffs = [1, 7, 16 - t, 12 - t]
@@ -618,19 +631,8 @@ def sigma_rel_to_gamma(sigma_rel):
     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"
-    
+def merge_lora_iterative_ema(paths, beta, sigma_rel):
+    print("Executing Iterative EMA Merge (Original 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()
@@ -651,12 +653,210 @@ def task_merge_adapters(hf_token, lora_urls, beta, sigma_rel, out_repo):
         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)
+    return base_sd
+
+# --- New Concatenation Method (DiffSynth) ---
+def merge_lora_concatenation(adapter_states, weights):
+    """
+    DiffSynth Method: Concatenates ranks.
+    New Rank = sum(ranks). Lossless merging.
+    """
+    print("Executing Concatenation Merge (Rank Summation)...")
+    merged_state = {}
+    
+    # Identify all stems (layers) present across all adapters
+    all_stems = set()
+    for state in adapter_states:
+        for k in state.keys():
+            stem = k.split(".lora_")[0]
+            if "lora_" in k: all_stems.add(stem)
+    
+    for stem in tqdm(all_stems, desc="Concatenating Layers"):
+        down_list = []
+        up_list = []
+        alpha_sum = 0.0
+        
+        for i, state in enumerate(adapter_states):
+            w = weights[i]
+            down_key = f"{stem}.lora_down.weight"
+            up_key = f"{stem}.lora_up.weight"
+            alpha_key = f"{stem}.alpha"
+            
+            if down_key in state and up_key in state:
+                d = state[down_key]
+                u = state[up_key] * w # weighted contribution applied to UP
+                
+                down_list.append(d)
+                up_list.append(u)
+                
+                if alpha_key in state:
+                    alpha_sum += state[alpha_key].item()
+                else:
+                    alpha_sum += d.shape[0] 
+                    
+        if down_list and up_list:
+            # Concat Down (A) along dim 0 (output of A, input to B) - Wait, lora_A is (rank, in)
+            # Concat Up (B) along dim 1 (input of B) - lora_B is (out, rank)
+            # Reference: DiffSynth code: lora_A = concat(tensors_A, dim=0), lora_B = concat(tensors_B, dim=1)
+            
+            new_down = torch.cat(down_list, dim=0) # (sum_rank, in)
+            new_up = torch.cat(up_list, dim=1)     # (out, sum_rank)
+            
+            merged_state[f"{stem}.lora_down.weight"] = new_down.contiguous()
+            merged_state[f"{stem}.lora_up.weight"] = new_up.contiguous()
+            merged_state[f"{stem}.alpha"] = torch.tensor(alpha_sum)
+            
+    return merged_state
+
+# --- New SVD/Task Arithmetic Method ---
+def merge_lora_svd(adapter_states, weights, target_rank):
+    """
+    SVD / Task Arithmetic Method:
+    1. Calculate Delta W for each adapter: dW = B @ A
+    2. Sum Delta Ws: Total dW = sum(weight_i * dW_i)
+    3. SVD(Total dW) -> New B, New A at target_rank
+    """
+    print(f"Executing SVD Merge (Target Rank: {target_rank})...")
+    merged_state = {}
+    
+    all_stems = set()
+    for state in adapter_states:
+        for k in state.keys():
+            stem = k.split(".lora_")[0]
+            if "lora_" in k: all_stems.add(stem)
+            
+    for stem in tqdm(all_stems, desc="SVD Merging Layers"):
+        total_delta = None
+        valid_layer = False
+        
+        for i, state in enumerate(adapter_states):
+            w = weights[i]
+            down_key = f"{stem}.lora_down.weight"
+            up_key = f"{stem}.lora_up.weight"
+            alpha_key = f"{stem}.alpha"
+            
+            if down_key in state and up_key in state:
+                down = state[down_key]
+                up = state[up_key]
+                alpha = state[alpha_key].item() if alpha_key in state else down.shape[0]
+                rank = down.shape[0]
+                
+                scale = (alpha / rank) * w
+                
+                # Reconstruct Delta
+                if len(down.shape) == 4: # Conv2d
+                    d_flat = down.flatten(start_dim=1)
+                    u_flat = up.flatten(start_dim=1)
+                    delta = (u_flat @ d_flat).reshape(up.shape[0], down.shape[1], down.shape[2], down.shape[3])
+                else:
+                    delta = up @ down
+                    
+                delta = delta * scale
+                
+                if total_delta is None:
+                    total_delta = delta
+                    valid_layer = True
+                else:
+                    if total_delta.shape == delta.shape:
+                        total_delta += delta
+                    else:
+                        print(f"Shape mismatch in {stem}, skipping.")
+        
+        if valid_layer and total_delta is not None:
+            out_dim = total_delta.shape[0]
+            in_dim = total_delta.shape[1]
+            is_conv = len(total_delta.shape) == 4
+            
+            if is_conv:
+                flat_delta = total_delta.flatten(start_dim=1)
+            else:
+                flat_delta = total_delta
+                
+            try:
+                U, S, V = torch.svd_lowrank(flat_delta, q=target_rank + 4, niter=4)
+                Vh = V.t()
+                
+                U = U[:, :target_rank]
+                S = S[:target_rank]
+                Vh = Vh[:target_rank, :]
                 
+                U = U @ torch.diag(S)
+                
+                if is_conv:
+                    U = U.reshape(out_dim, target_rank, 1, 1)
+                    Vh = Vh.reshape(target_rank, in_dim, total_delta.shape[2], total_delta.shape[3])
+                else:
+                    U = U.reshape(out_dim, target_rank)
+                    Vh = Vh.reshape(target_rank, in_dim)
+                    
+                merged_state[f"{stem}.lora_down.weight"] = Vh.contiguous()
+                merged_state[f"{stem}.lora_up.weight"] = U.contiguous()
+                merged_state[f"{stem}.alpha"] = torch.tensor(target_rank).float()
+            except Exception as e:
+                print(f"SVD Failed for {stem}: {e}")
+
+    return merged_state
+
+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())
+    
+    if not out_repo or not out_repo.strip():
+        return "Error: Output Repo cannot be empty."
+
+    # 1. Parse Inputs (Multi-line support)
+    raw_lines = inputs_text.replace(" ", "\n").split('\n')
+    urls = [line.strip() for line in raw_lines if line.strip()]
+    if len(urls) < 2: return "Error: Please provide at least 2 adapters."
+    
+    # 2. Parse Weights (for SVD/Concatenation)
+    try:
+        if not weight_str.strip():
+            weights = [1.0] * len(urls)
+        else:
+            weights = [float(w.strip()) for w in weight_str.split(',')]
+            # Broadcast or Truncate
+            if len(weights) < len(urls):
+                weights += [1.0] * (len(urls) - len(weights))
+            else:
+                weights = weights[:len(urls)]
+    except:
+        return "Error parsing weights. Use format: 1.0, 0.5, 0.8"
+
+    # 3. Download All
+    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
+
+    # 4. Execute Selected Method
+    if "Iterative EMA" in method:
+        # Calls the original method logic exactly
+        merged = merge_lora_iterative_ema(paths, beta, sigma_rel)
+    
+    else:
+        # For new methods, we load everything upfront
+        states = [load_full_state_dict(p) for p in paths]
+        
+        if "Concatenation" in method:
+            merged = merge_lora_concatenation(states, weights)
+        elif "SVD" in method:
+            merged = merge_lora_svd(states, weights, int(target_rank))
+
+    if not merged: return "Merge failed (Result empty)."
+
+    # 5. Save & Upload
     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"
+    save_file(merged, out)
+    
+    try:
+        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}"
+    except Exception as e: return f"Upload Error: {e}"
 
 # =================================================================================
 # TAB 4: RESIZE (CPU Optimized)
@@ -770,18 +970,34 @@ with gr.Blocks() as demo:
             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"):
+            gr.Markdown("### Batch Adapter Merging")
             t3_token = gr.Textbox(label="Token", type="password")
-            t3_urls = gr.Textbox(label="URLs")
+            t3_urls = gr.TextArea(label="Adapter URLs/Repos (One per line, or space separated)", placeholder="ostris/lora1\nhttps://hf.co/user/lora2.safetensors\n...")
+            
+            with gr.Row():
+                t3_method = gr.Dropdown(
+                    ["Iterative EMA (Original Beta/Sigma)", "Concatenation (DiffSynth - Lossless)", "SVD Merge (Task Arithmetic/Compressed)"], 
+                    value="Iterative EMA (Original Beta/Sigma)", 
+                    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=4, maximum=1024)
+            
             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_beta = gr.Slider(label="Beta - For EMA only", value=0.95, minimum=0.01, maximum=1.00, step=0.01)
+                t3_sigma = gr.Slider(label="Sigma Rel - For EMA only", 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_priv = gr.Checkbox(label="Private Output", value=True)
+            t3_btn = gr.Button("Merge Adapters")
             t3_res = gr.Textbox(label="Result")
-            t3_btn.click(task_merge_adapters, [t3_token, t3_urls, t3_beta, t3_sigma, t3_out], t3_res)
-
+            
+            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)
+            
         with gr.Tab("Resize Adapter"):
             t4_token = gr.Textbox(label="Token", type="password")
             t4_in = gr.Textbox(label="LoRA")