Update app.py

app.py

@@ -239,27 +239,21 @@ def merge_shard_logic(base_path, lora_pairs, scale, output_path):
                 valid_delta = False
                 
             if valid_delta:
-                # IN-PLACE MERGE to save memory
-                # 1. Promote to float32
-                # 2. Add delta
-                # 3. Cast back to original dtype
-                # 4. Replace in dict
-                orig_dtype = v.dtype
+                # Optimized In-Place Addition
+                # We do NOT cast base to float32. We trust bf16/fp16 is sufficient for merging.
                 
-                # Perform add in float32 to avoid overflow/precision issues
-                # Create temp float tensor
-                v_float = v.to(torch.float32)
-                v_float.add_(delta) # In-place add
+                # If base is float32 (rare for new models), we respect it. 
+                # If base is bf16, we add bf16 delta.
+                if v.dtype != delta.dtype:
+                    delta = delta.to(v.dtype)
                 
-                # Cast back and replace in dict
-                base_state[k] = v_float.to(orig_dtype)
+                # In-place add
+                v.add_(delta)
                 
                 # Explicit cleanup
-                del v_float
                 del delta
-                # del v # v is a reference to base_state[k], which we just overwrote
         
-        # Periodic GC to prevent fragmentation OOM
+        # Periodic GC
         if len(keys_to_process) > 100 and keys_to_process.index(k) % 50 == 0:
             gc.collect()
 
@@ -269,6 +263,16 @@ def merge_shard_logic(base_path, lora_pairs, scale, output_path):
 def task_merge(hf_token, base_repo, base_subfolder, lora_input, scale, output_repo, structure_repo, private, progress=gr.Progress()):
     cleanup_temp()
     login(hf_token)
+
+        # Determine Dtype
+    if precision == "bf16":
+        dtype = torch.bfloat16
+    elif precision == "fp16":
+        dtype = torch.float16
+    else:
+        dtype = torch.float32
+    
+    print(f"Selected Precision: {dtype}")
     
     try:
         api.create_repo(repo_id=output_repo, private=private, exist_ok=True, token=hf_token)
@@ -288,9 +292,13 @@ def task_merge(hf_token, base_repo, base_subfolder, lora_input, scale, output_re
         except Exception as e:
             print(f"Structure clone warning: {e}")
 
-    progress(0.1, desc="Loading LoRA...")
-    lora_path = download_file(lora_input, hf_token)
-    lora_pairs = load_lora_to_memory(lora_path)
+    try:
+        progress(0.1, desc="Downloading LoRA...")
+        lora_path = download_file(lora_input, hf_token)
+        # Load LoRA in target precision to save RAM immediately
+        lora_pairs = load_lora_to_memory(lora_path, precision_dtype=dtype)
+    except Exception as e:
+        return f"CRITICAL ERROR: {str(e)}"
     
     files = list_repo_files(repo_id=base_repo, token=hf_token)
     shards = [f for f in files if f.endswith(".safetensors")]
@@ -303,9 +311,9 @@ def task_merge(hf_token, base_repo, base_subfolder, lora_input, scale, output_re
         progress(0.2 + (0.8 * i/len(shards)), desc=f"Merging {shard}")
         local_shard = hf_hub_download(repo_id=base_repo, filename=shard, token=hf_token, local_dir=TempDir)
         merged_path = TempDir / "merged.safetensors"
-        
-        # Merge Logic
-        merge_shard_logic(local_shard, lora_pairs, scale, merged_path)
+
+                # Pass precision preference
+        merge_shard_logic(local_shard, lora_pairs, scale, merged_path, precision_dtype=dtype)
         
         # Upload
         api.upload_file(path_or_fileobj=merged_path, path_in_repo=shard, repo_id=output_repo, token=hf_token)