Create simple_alpha_model_merger.py

simple_alpha_model_merger.py

@@ -0,0 +1,68 @@
+# merge_two_fp8.py
+# Merges two FP8 .safetensors files with weighted average
+# Usage: python merge_two_fp8.py
+# Or with args: python merge_two_fp8.py --model_a path/to/existing_fp8.safetensors --model_b path/to/new_model_fp8.safetensors --alpha 0.7 --output merged_result.safetensors
+
+from safetensors.torch import load_file, save_file
+import torch
+import argparse
+import os
+
+def merge_weights(dict_a, dict_b, alpha):
+    """
+    alpha: weight for model_b (VBVR/new one). 0.0 = 100% model_a, 1.0 = 100% model_b
+    """
+    merged = {}
+    keys_a = set(dict_a.keys())
+    keys_b = set(dict_b.keys())
+    
+    all_keys = keys_a.union(keys_b)
+    
+    for key in all_keys:
+        if key in dict_a and key in dict_b:
+            # Weighted average for matching keys
+            merged[key] = (1 - alpha) * dict_a[key] + alpha * dict_b[key]
+        elif key in dict_a:
+            merged[key] = dict_a[key]
+        elif key in dict_b:
+            merged[key] = dict_b[key]
+        else:
+            continue  # shouldn't happen
+    
+    return merged
+
+def main():
+    parser = argparse.ArgumentParser(description="Weighted merge two FP8 safetensors files (Wan2.2 style)")
+    parser.add_argument("--model_a", required=True, help="Path to existing merged FP8 model (base/reference)")
+    parser.add_argument("--model_b", required=True, help="Path to new VBVR FP8 model (the one to blend in)")
+    parser.add_argument("--alpha", type=float, default=0.7, help="Blend strength for model_b (0.0–1.0, default 0.7 = 70%% VBVR + 30%% existing)")
+    parser.add_argument("--output", default="vbvr_merged_with_existing_fp8.safetensors", help="Output filename")
+    
+    args = parser.parse_args()
+    
+    if not (0 <= args.alpha <= 1):
+        raise ValueError("Alpha must be between 0.0 and 1.0")
+    
+    print(f"Merging:\n  A: {args.model_a}\n  B: {args.model_b}\n  Alpha (for B): {args.alpha}\n  Output: {args.output}")
+    
+    print("Loading model A...")
+    dict_a = load_file(args.model_a, device="cpu")
+    
+    print("Loading model B...")
+    dict_b = load_file(args.model_b, device="cpu")
+    
+    print("Merging weights...")
+    merged_dict = merge_weights(dict_a, dict_b, args.alpha)
+    
+    # Ensure output has .safetensors extension
+    if not args.output.endswith(".safetensors"):
+        args.output += ".safetensors"
+    
+    print(f"Saving merged FP8 model...")
+    save_file(merged_dict, args.output)
+    
+    print(f"Done! New file: {os.path.abspath(args.output)}")
+    print("Test it in ComfyUI or Diffusers — adjust alpha if motion/reasoning balance feels off.")
+
+if __name__ == "__main__":
+    main()