accelerate
numpy==1.26.4 # Due to newer version of numpy changing dtype when multiplying
diffusers
torch
transformers
onnxruntime-directml

import onnxruntime as ort
from diffusers import OnnxStableDiffusionPipeline

model_dir = "D:\\Models\\stable-diffusion-v1-5_io32"

batch_size = 1
num_inference_steps = 30
image_size = 512
guidance_scale = 7.5
prompt = "a beautiful cabin in the mountains of Lake Tahoe"

ort.set_default_logger_severity(3)

sess_options = ort.SessionOptions()
sess_options.enable_mem_pattern = False

sess_options.add_free_dimension_override_by_name("unet_sample_batch", batch_size * 2)
sess_options.add_free_dimension_override_by_name("unet_sample_channels", 4)
sess_options.add_free_dimension_override_by_name("unet_sample_height", image_size // 8)
sess_options.add_free_dimension_override_by_name("unet_sample_width", image_size // 8)
sess_options.add_free_dimension_override_by_name("unet_time_batch", batch_size)
sess_options.add_free_dimension_override_by_name("unet_hidden_batch", batch_size * 2)
sess_options.add_free_dimension_override_by_name("unet_hidden_sequence", 77)

pipeline = OnnxStableDiffusionPipeline.from_pretrained(
    model_dir, provider="DmlExecutionProvider", sess_options=sess_options
)

result = pipeline(
        [prompt] * batch_size,
        num_inference_steps=num_inference_steps,
        callback=None,
        height=image_size,
        width=image_size,
        guidance_scale=guidance_scale,
        generator=None
    )

output_path = "output.png"
result.images[0].save(output_path)

print(f"Generated {output_path}")