from safetensors.torch import load_file

# Step 1: Load the safetensors file
checkpoint_path = 'flowgram.safetensors'  # Replace with your actual file path
checkpoint = load_file(checkpoint_path)

# Step 2: Open a log file to save the output
with open("log.txt", "w") as log_file:
    # Step 3: Write the size (shape) of each tensor to the file
    for tensor_name, tensor in checkpoint.items():
        log_file.write(f"Tensor Name: {tensor_name}, Size: {tensor.shape}\n")

print("Tensor sizes saved to log.txt")