Update app.py

app.py

@@ -454,42 +454,111 @@ def task_merge(hf_token, base_repo, base_subfolder, lora_input, scale, precision
 # TAB 2: EXTRACT LORA
 # =================================================================================
 
+def identify_and_download_model(repo_id, token):
+    """
+    Smart download: checks for diffusers format (unet/transformer) vs standard safetensors.
+    """
+    print(f"Scanning {repo_id} for model weights...")
+    files = list_repo_files(repo_id=repo_id, token=token)
+    
+    # Priority list for diffusers vs single file
+    priorities = [
+        "transformer/diffusion_pytorch_model.safetensors",
+        "unet/diffusion_pytorch_model.safetensors",
+        "model.safetensors",
+        # Fallback to any safetensors that isn't an adapter or lora
+        lambda f: f.endswith(".safetensors") and "lora" not in f and "adapter" not in f and "extracted" not in f
+    ]
+    
+    target_file = None
+    for p in priorities:
+        if callable(p):
+            candidates = [f for f in files if p(f)]
+            if candidates:
+                target_file = candidates[0]
+                break
+        elif p in files:
+            target_file = p
+            break
+            
+    if not target_file:
+        raise ValueError(f"Could not find a valid model weight file in {repo_id}. Ensure it contains .safetensors weights.")
+        
+    print(f"Downloading main weight file: {target_file}")
+    hf_hub_download(repo_id=repo_id, filename=target_file, token=token, local_dir=TempDir)
+    
+    # Locate actual path
+    found = list(TempDir.rglob(os.path.basename(target_file)))[0]
+    return found
+
 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
+    print("Calculating diffs & extracting LoRA...")
+    
+    # Get intersection of keys
+    keys = set(org.keys()).intersection(set(tuned.keys()))
+    
+    for key in tqdm(keys, desc="Extracting"):
+        # Skip integer buffers/metadata
+        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()
+        
+        # Skip if shapes mismatch (shouldn't happen if models match)
+        if mat_org.shape != mat_tuned.shape: continue
+        
         diff = mat_tuned - mat_org
+        
+        # Skip if no difference
         if torch.max(torch.abs(diff)) < 1e-4: continue
         
-        out_dim, in_dim = diff.shape[:2]
+        out_dim = diff.shape[0]
+        in_dim = 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) # Handle biases if needed, though rarely lora'd
             
         try:
-            U, S, Vh = torch.linalg.svd(diff, full_matrices=False)
-            U, S, Vh = U[:, :r], S[:r], Vh[:r, :]
+            # Use svd_lowrank for massive speedup on CPU vs linalg.svd
+            U, S, V = torch.svd_lowrank(diff, q=r+4, niter=4)
+            Vh = V.t()
+            
+            U = U[:, :r]
+            S = S[:r]
+            Vh = Vh[:r, :]
+            
+            # Merge S into U for standard LoRA format
             U = U @ torch.diag(S)
+            
+            # Clamp outliers
             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)
+            hi_val = torch.quantile(torch.abs(dist), clamp)
+            if hi_val > 0:
+                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}.lora_up.weight"] = U.contiguous()
+            lora_sd[f"{stem}.lora_down.weight"] = Vh.contiguous()
             lora_sd[f"{stem}.alpha"] = torch.tensor(r).float()
-        except: pass
+        except Exception as e:
+            print(f"Skipping {key} due to error: {e}")
+            pass
+            
     out = TempDir / "extracted.safetensors"
     save_file(lora_sd, out)
     return str(out)
@@ -498,12 +567,16 @@ 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")
+        print("Downloading Original Model...")
+        p1 = identify_and_download_model(org, hf_token)
+        print("Downloading Tuned Model...")
+        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.safetensors", repo_id=out, token=hf_token)
-        return "Done"
+        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}"
 
 # =================================================================================