import torch
from safetensors.torch import load_file

# Load both models
model1 = load_file('merged_model.safetensors')
model2 = load_file('diffusion_pytorch_model-00002-of-00003.safetensors')

# 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")

print("Comparison complete. Check log.txt for details.")