import torch
from sdxl_unet2_trt import MyStableDiffusionXLPipeline
from PIL import Image

"""
With the current pipeline it just use to compare the Unet Diffusers vs Unet Onnx (calibrate)
Cannot get the result image  
"""

pipe = MyStableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float32
)
# pipe = pipe.to("cuda")

"""
For compare the Unet Diffusers vs Unet Onnx (calibrate)
"""
prompt = 'Portrait shot of a woman, yellow shirt, photograph'
image = pipe(prompt)

"""
Uncommand for get the result image
"""
# prompt = 'Portrait shot of a woman, yellow shirt, photograph'
# image = pipe(prompt).images[0]

# Save the image
# save_path = "/home/tiennv/trang/model_unet_quantize_full_onnx_2/image_sdxl_trt_b4.png"
# image.save(save_path)