Update sh.py

sh.py

@@ -1,22 +1,26 @@
- 
 import torch
 from safetensors.torch import load_file
 
-
 # Load both models
 model1 = load_file('merged_model_016.safetensors')
 model2 = load_file('diffusion_pytorch_model-00001-of-00003.safetensors')
 
-# Iterate through the tensor names and shapes
-for name in model1.keys():
-    if name in model2:
-        shape1 = model1[name].shape
-        shape2 = model2[name].shape
-        if shape1 != shape2:
-            print(f"Tensor '{name}' has different shapes: Model 1: {shape1}, Model 2: {shape2}")
-    else:
-        print(f"Tensor '{name}' is not present in model 2.")
+# Open log.txt for writing the output
+with open('log.txt', 'w') as log_file:
+    
+    # Iterate through the tensor names and shapes of model1
+    for name in model1.keys():
+        if name in model2:
+            shape1 = model1[name].shape
+            shape2 = model2[name].shape
+            if shape1 != shape2:
+                log_file.write(f"Tensor '{name}' has different shapes: Model 1: {shape1}, Model 2: {shape2}\n")
+        else:
+            log_file.write(f"Tensor '{name}' is not present in model 2.\n")
+
+    # Iterate through the tensor names and shapes of model2 to find ones not in model1
+    for name in model2.keys():
+        if name not in model1:
+            log_file.write(f"Tensor '{name}' is not present in model 1.\n")
 
-for name in model2.keys():
-    if name not in model1:
-        print(f"Tensor '{name}' is not present in model 1.")
+print("Comparison complete. Check log.txt for details.")