import numpy as np
import onnx
from onnxruntime import InferenceSession
from optimum.onnx.graph_transformations import embedd_loop_in_unet

path_unet = "stable_diffusion/unet/model.onnx"
path_unet_with_loop = "unet_with_loop.onnx"

unet_model = onnx.load(path_unet)
unet_with_loop = embedd_loop_in_unet(unet_model)
onnx.save(unet_with_loop, path_unet_with_loop)

unet_sess = InferenceSession(path_unet, providers=["CPUExecutionProvider"])
unet_with_loop_sess = InferenceSession("unet_with_loop.onnx", providers=["CPUExecutionProvider"])


batch_size = 16
initial_sample = np.random.uniform(size=(batch_size, 4, 32, 32)).astype(np.float32)
hidden_states = np.random.uniform(size=(batch_size, 25, 32)).astype(np.float32)

def outside_loop(iterations):
    sample = initial_sample
    for i in range(iterations):
        timestep = np.array([i] * batch_size)
        out  = unet_sess.run(
            ["out_sample"],
            {
                "sample": sample, 
                "timestep": timestep,
                "encoder_hidden_states": hidden_states
            }
        )
        sample = out[0]
    return sample

def embedded_loop(iterations):
    num_iterations = np.array(iterations)
    outputs = unet_with_loop_sess.run(
        ["sample"], 
        {
            "num_iterations": num_iterations,
            "initial_sample": initial_sample, 
            "encoder_hidden_states": hidden_states
        }
    )
    outputs = outputs[0]
    return outputs

outputs = outside_loop(10)
outputs_with_loop = embedded_loop(10)

print(np.max(np.abs(outputs - outputs_with_loop)))