Update app.py

app.py

@@ -46,15 +46,12 @@ def download_lora(lora_input, hf_token):
     else:
         # Repo ID download
         print(f"Downloading LoRA from Repo: {lora_input}")
-        # Try finding the safetensors file
         try:
             return hf_hub_download(repo_id=lora_input, filename="adapter_model.safetensors", token=hf_token, local_dir=TempDir)
         except:
-            # Fallback for diffusion models which might use different names
             files = list_repo_files(repo_id=lora_input, token=hf_token)
             safe_files = [f for f in files if f.endswith(".safetensors") and "adapter" in f]
             if not safe_files:
-                # Last ditch: grab the first safetensors
                 safe_files = [f for f in files if f.endswith(".safetensors")]
             
             if not safe_files:
@@ -63,28 +60,11 @@ def download_lora(lora_input, hf_token):
             return hf_hub_download(repo_id=lora_input, filename=safe_files[0], token=hf_token, local_dir=TempDir)
 
 def load_lora_weights(path):
-    """Loads LoRA weights and attempts to determine rank/alpha."""
     tensors = load_file(path, device="cpu")
-    # Basic metadata extraction could happen here if needed, 
-    # but for raw merging we mainly need the state dict.
     return tensors
 
 def match_keys(base_key, lora_keys):
-    """
-    Heuristic matching. 
-    1. Exact match (rare for LoRA).
-    2. LoRA naming conventions (lora_A, lora_B, lora_down, etc).
-    """
-    # Common LoRA naming patterns
-    # pattern: base_key.lora_A.weight
-    # pattern: base_key + ".0.lora_B.weight" (sometimes happens)
-    
     matches = {}
-    
-    # Cleaning the keys for comparison
-    # If base is "transformer.blocks.0.weight"
-    # LoRA might be "transformer.blocks.0.lora_A.weight"
-    
     candidates = [k for k in lora_keys if base_key in k]
     
     pair_A = None
@@ -99,11 +79,9 @@ def match_keys(base_key, lora_keys):
     return pair_A, pair_B
 
 def copy_auxiliary_files(src_repo, tgt_repo, token, subfolder=""):
-    """Copies config/tokenizer/scheduler files from source to target."""
     print(f"Copying infrastructure from {src_repo} to {tgt_repo}...")
     files = list_repo_files(repo_id=src_repo, token=token)
     
-    # Filter out heavy weights
     files_to_copy = [
         f for f in files 
         if not f.endswith(".safetensors") 
@@ -116,7 +94,6 @@ def copy_auxiliary_files(src_repo, tgt_repo, token, subfolder=""):
 
     for f in tqdm(files_to_copy, desc="Copying configs"):
         try:
-            # We download to memory/temp and upload immediately
             local = hf_hub_download(repo_id=src_repo, filename=f, token=token)
             api.upload_file(
                 path_or_fileobj=local,
@@ -154,9 +131,9 @@ def run_merge(
         except Exception as e:
             return "\n".join(logs) + f"\nError creating repo: {e}"
 
-        # 2. Replicate Structure (If requested)
+        # 2. Replicate Structure
         if structure_repo.strip():
-            progress(0.1, desc="Cloning Model Structure (Configs)...")
+            progress(0.1, desc="Cloning Model Structure...")
             logs.append(f"Cloning configuration from {structure_repo}...")
             copy_auxiliary_files(structure_repo, output_repo, hf_token)
             logs.append("Configuration files copied.")
@@ -173,25 +150,19 @@ def run_merge(
         progress(0.3, desc="Analyzing Base Model...")
         all_files = list_repo_files(repo_id=base_repo, token=hf_token)
         
-        # Filter for safetensors in the specific subfolder (if provided)
         target_shards = []
         for f in all_files:
             if not f.endswith(".safetensors"):
                 continue
-            
-            # Check subfolder constraint
-            if base_subfolder.strip():
-                # Normalize paths
-                if not f.startswith(base_subfolder.strip("/")):
-                    continue
-            
+            if base_subfolder.strip() and not f.startswith(base_subfolder.strip("/")):
+                continue
             target_shards.append(f)
             
         logs.append(f"Found {len(target_shards)} matching safetensors shards in base.")
         if not target_shards:
             raise ValueError("No safetensors found in the specified base repo/subfolder.")
 
-        # 5. Process Shards (Streamed)
+        # 5. Process Shards
         total_shards = len(target_shards)
         merged_count = 0
         
@@ -199,28 +170,16 @@ def run_merge(
             progress(0.3 + (0.6 * (idx / total_shards)), desc=f"Processing Shard {idx+1}/{total_shards}")
             logs.append(f"--- Processing {shard_file} ---")
             
-            # Download Shard
             local_shard = hf_hub_download(repo_id=base_repo, filename=shard_file, token=hf_token, local_dir=TempDir)
-            
-            # Load and Merge
-            # We use safe_open to read metadata, but load_file for the dict to modify
-            # load_file loads to CPU RAM.
             base_tensors = load_file(local_shard, device="cpu")
             modified_tensors = {}
             has_changes = False
             
             for key, tensor in base_tensors.items():
-                # Match LoRA
-                # Handle architectural prefix mismatches (e.g. Ostris repo might rely on folder structure, 
-                # while LoRA expects "transformer." prefix)
-                
-                # Try exact match first (unlikely for LoRA)
                 pair_A, pair_B = match_keys(key, lora_keys)
                 
-                # If not found, try adding/removing common prefixes
                 if not pair_A:
-                     # Attempt to match "blocks.1..." to "transformer.blocks.1..."
-                     matches = [k for k in lora_keys if key in k] # Simple substring check
+                     matches = [k for k in lora_keys if key in k]
                      for k in matches:
                         if "lora_A" in k or "lora_down" in k:
                             pair_A = k
@@ -228,24 +187,16 @@ def run_merge(
                             pair_B = k
 
                 if pair_A and pair_B:
-                    # Apply Merge
                     w_a = lora_state[pair_A].float()
                     w_b = lora_state[pair_B].float()
-                    
-                    # Target tensor
                     current_tensor = tensor.float()
-                    
-                    # Dimension Check
-                    # LoRA = B @ A. Shape should match current_tensor.
-                    # Sometimes LoRA weights are transposed relative to base depending on training lib.
                     delta = (w_b @ w_a) * scale
                     
                     if delta.shape != current_tensor.shape:
-                        # Try transposing matches
                         if delta.T.shape == current_tensor.shape:
                             delta = delta.T
                         else:
-                            logs.append(f"Warning: Shape mismatch for {key}. Base: {current_tensor.shape}, LoRA Delta: {delta.shape}. Skipping.")
+                            logs.append(f"Warning: Shape mismatch for {key}. Skipping.")
                             modified_tensors[key] = tensor
                             continue
                             
@@ -255,33 +206,16 @@ def run_merge(
                 else:
                     modified_tensors[key] = tensor
 
-            # Save and Upload
             if has_changes:
                 logs.append(f"Merging complete for shard. Saving...")
                 output_path = TempDir / "processed.safetensors"
                 save_file(modified_tensors, output_path)
-                
-                api.upload_file(
-                    path_or_fileobj=output_path,
-                    path_in_repo=shard_file, # Keep original structure
-                    repo_id=output_repo,
-                    repo_type="model",
-                    token=hf_token
-                )
+                api.upload_file(path_or_fileobj=output_path, path_in_repo=shard_file, repo_id=output_repo, repo_type="model", token=hf_token)
                 logs.append(f"Uploaded {shard_file}")
             else:
-                # If no changes, just copy the original file to the new repo
-                # This saves re-saving the tensor dict
                 logs.append(f"No LoRA matches in this shard. Copying original...")
-                api.upload_file(
-                    path_or_fileobj=local_shard,
-                    path_in_repo=shard_file,
-                    repo_id=output_repo,
-                    repo_type="model",
-                    token=hf_token
-                )
+                api.upload_file(path_or_fileobj=local_shard, path_in_repo=shard_file, repo_id=output_repo, repo_type="model", token=hf_token)
             
-            # Cleanup Memory immediately
             del base_tensors
             del modified_tensors
             if 'delta' in locals(): del delta
@@ -301,10 +235,9 @@ def run_merge(
     
     finally:
         cleanup_temp()
-        
+    
     return "\n".join(logs)
 
-
 # --- UI ---
 
 css = """
@@ -312,7 +245,8 @@ css = """
 .header { text-align: center; margin-bottom: 20px; }
 """
 
-with gr.Blocks(css=css) as demo:
+# NOTE: Removed 'css' and 'theme' from gr.Blocks() to be compatible with latest Gradio versions.
+with gr.Blocks() as demo:
     gr.Markdown(
         """
         # ⚡ Universal LoRA Merger & Reconstructor
@@ -357,4 +291,5 @@ with gr.Blocks(css=css) as demo:
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=1).launch()
+    # CSS is now passed here in the launch method
+    demo.queue(max_size=1).launch(css=css)