init

tools/onnx_optimize.sh

+if [[ $1 = "--help" ]] || [[ $1 = "-h" ]]
+then
+    echo "Optimize stable diffusion ONNX models by using ONNXRuntime."
+    echo "Usage: ./onnx_optimize.sh <original_onnx_path> <converted_onnx_path>"
+    exit 0
+fi
+
+if [ $# -ne 2 ]
+then
+    echo "Error: Incorrect number of arguments"
+	echo "Usage: ./onnx_optimize.sh <original_onnx_path> <converted_onnx_path>"
+    exit 1
+fi
+
+set -e
+
+original_onnx_path=$1
+converted_onnx_path=$2
+tmp_save_path="./temp"
+
+# for unet
+python -m onnxruntime.transformers.models.stable_diffusion.optimize_pipeline \
+-i ${original_onnx_path} \
+-o ${tmp_save_path} \
+--disable_bias_gelu \
+--disable_bias_add \
+--disable_bias_splitgelu \
+--disable_nhwc_conv \
+--use_group_norm_channels_first \
+--use_multi_head_attention \
+--float16
+
+# for text_encoders and vae_decoder
+python -m onnxruntime.transformers.models.stable_diffusion.optimize_pipeline \
+-i ${original_onnx_path} \
+-o ${converted_onnx_path} \
+--disable_bias_gelu \
+--disable_bias_add \
+--disable_bias_splitgelu \
+--disable_nhwc_conv \
+--disable_attention \
+--use_group_norm_channels_first \
+--use_multi_head_attention \
+--float16
+
+rm -r ${converted_onnx_path}/unet
+mv ${tmp_save_path}/unet ${converted_onnx_path}/
+rm -r ${tmp_save_path}
+
+echo "Finish to optimize ONNX models!"
+echo "optimized ONNX models are saved in ${converted_onnx_path}"

tools/run_examples.py

+import json
+import os
+import os.path as osp
+
+from diffusers import DiffusionPipeline
+import migraphx_diffusers
+import torch
+
+
+def parse_args():
+    from argparse import ArgumentParser
+    parser = ArgumentParser(description="SDXL inference with migraphx backend")
+
+    #=========================== mdoel load and compile ========================
+    parser.add_argument(
+        "-m", 
+        "--model-dir",
+        type=str,
+        required=True,
+        help="Path to local model directory.",
+    )
+    parser.add_argument(
+        "--force-compile",
+        action="store_true",
+        default=False,
+        help="Ignore existing .mxr files and override them",
+    )
+    parser.add_argument(
+        "--num-images-per-prompt",
+        type=int,
+        default=1,
+        help="The number of images to generate per prompt."
+    )
+    parser.add_argument(
+        "--img-size",
+        type=int,
+        default=None,
+        help="output image size",
+    )
+    # --------------------------------------------------------------------------
+
+    # =============================== generation ===============================
+    parser.add_argument(
+        "-t",
+        "--num-inference-steps",
+        type=int,
+        default=50,
+        help="Number of iteration steps",
+    )
+    parser.add_argument(
+        "-s",
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed",
+    )
+    # --------------------------------------------------------------------------
+
+    parser.add_argument(
+        "--examples-json",
+        type=str,
+        default="./examples/prompts_and_negative_prompts.json",
+        help="Prompts and negative prompts data path",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default=None,
+        help="Path to save images",
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+def get_name_and_migraphx_config(model_dir):
+    model_index_json = osp.join(model_dir, "model_index.json")
+    with open(model_index_json, "r") as f:
+        pipe_cfg = json.load(f)
+    
+    if pipe_cfg["_class_name"] == "StableDiffusionXLPipeline":
+        return 'sdxl', migraphx_diffusers.DEFAULT_ARGS['sdxl']
+    elif pipe_cfg["_class_name"] == "StableDiffusionPipeline":
+        return 'sd2.1', migraphx_diffusers.DEFAULT_ARGS['sd2.1']
+    else:
+        raise NotImplementedError(
+            f"{pipe_cfg['_class_name']} has not been adapted yet")
+
+def parse_prompts(examples_json):
+    with open(examples_json, 'r') as f:
+        prompt_data = json.load(f)
+    return prompt_data
+
+def main():
+    args = parse_args()
+    name, migraphx_config = get_name_and_migraphx_config(args.model_dir)
+    if args.output_dir is None:
+        args.output_dir = f"./examples/{name}-images-{args.img_size}"
+
+    if args.img_size is not None:
+        migraphx_config['common_args']['img_size'] = args.img_size
+    migraphx_config['common_args'].update(dict(
+        batch=args.num_images_per_prompt,
+        force_compile=args.force_compile,
+    ))
+
+    pipe = DiffusionPipeline.from_pretrained(
+        args.model_dir,
+        torch_dtype=torch.float16, 
+        migraphx_config=migraphx_config
+    )
+    pipe.to("cuda")
+
+    prompt_data = parse_prompts(args.examples_json)
+
+    cnt = 0
+    for i, d in enumerate(prompt_data):
+        theme = d["theme"]
+        pairs = d["examples"]
+        sub_dir = osp.join(args.output_dir, 
+                           f"{i}-{theme.title().replace(' ', '')}")
+        os.makedirs(sub_dir, exist_ok=True)
+
+        for j, pair in enumerate(pairs):
+            print(f"Generating image {cnt}...")
+            prompt = pair["prompt"]
+            negative_prompt = pair["negative_prompt"]
+            print(f"Prompt: {prompt}")
+            print(f"negative Prompt: {negative_prompt}")
+
+            images = pipe(
+                prompt=prompt, 
+                negative_prompt=negative_prompt, 
+                num_inference_steps=args.num_inference_steps,
+                generator=torch.Generator("cuda").manual_seed(args.seed)
+            ).images
+
+            for k, image in enumerate(images):
+                save_path = osp.join(
+                    sub_dir, f"theme_{i}_example_{j}_image_{k}.png")
+                image.save(save_path)
+                print(f"Image saved: {save_path}")
+                cnt += 1
+    
+    print(f"Total {cnt} images Generated!")
+
+
+if __name__ == "__main__":
+    main()

tools/run_pipe.py

+import json
+import os.path as osp
+
+from diffusers import DiffusionPipeline
+import migraphx_diffusers
+import torch
+
+
+def parse_args():
+    from argparse import ArgumentParser
+    parser = ArgumentParser(description="SDXL inference with migraphx backend")
+
+    #=========================== mdoel load and compile ========================
+    parser.add_argument(
+        "-m", 
+        "--model-dir",
+        type=str,
+        required=True,
+        help="Path to local model directory.",
+    )
+    parser.add_argument(
+        "--force-compile",
+        action="store_true",
+        default=False,
+        help="Ignore existing .mxr files and override them",
+    )
+    parser.add_argument(
+        "--img-size",
+        type=int,
+        default=None,
+        help="output image size",
+    )
+    parser.add_argument(
+        "--num-images-per-prompt",
+        type=int,
+        default=1,
+        help="The number of images to generate per prompt."
+    )
+    # --------------------------------------------------------------------------
+
+    # =============================== generation ===============================
+    parser.add_argument(
+        "-p",
+        "--prompt",
+        type=str,
+        required=True,
+        help="Prompt for describe image content, style and so on."
+    )
+    parser.add_argument(
+        "-n",
+        "--negative-prompt",
+        type=str,
+        default=None,
+        help="Negative prompt",
+    )
+    parser.add_argument(
+        "-t",
+        "--num-inference-steps",
+        type=int,
+        default=50,
+        help="Number of iteration steps",
+    )
+    parser.add_argument(
+        "--save-prefix",
+        type=str,
+        default=None,
+        help="Prefix of path for saving results",
+    )
+    parser.add_argument(
+        "-s",
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed",
+    )
+    # --------------------------------------------------------------------------
+
+    args = parser.parse_args()
+    return args
+
+
+def get_name_and_migraphx_config(model_dir):
+    model_index_json = osp.join(model_dir, "model_index.json")
+    with open(model_index_json, "r") as f:
+        pipe_cfg = json.load(f)
+    
+    if pipe_cfg["_class_name"] == "StableDiffusionXLPipeline":
+        return 'sdxl', migraphx_diffusers.DEFAULT_ARGS['sdxl']
+    elif pipe_cfg["_class_name"] == "StableDiffusionPipeline":
+        return 'sd2.1', migraphx_diffusers.DEFAULT_ARGS['sd2.1']
+    else:
+        raise NotImplementedError(
+            f"{pipe_cfg['_class_name']} has not been adapted yet")
+
+
+def main():
+    args = parse_args()
+    name, migraphx_config = get_name_and_migraphx_config(args.model_dir)
+    if args.save_prefix is None:
+        args.save_prefix = f"./{name}_output"
+
+    if args.img_size is not None:
+        migraphx_config['common_args']['img_size'] = args.img_size
+    migraphx_config['common_args'].update(dict(
+        batch=args.num_images_per_prompt,
+        force_compile=args.force_compile,
+    ))
+
+    pipe = DiffusionPipeline.from_pretrained(
+        args.model_dir,
+        torch_dtype=torch.float16,
+        migraphx_config=migraphx_config
+    )
+    pipe.to("cuda")
+
+    print("Generating image...")
+    images = pipe(
+        prompt=args.prompt, 
+        negative_prompt=args.negative_prompt, 
+        num_inference_steps=args.num_inference_steps,
+        generator=torch.Generator("cuda").manual_seed(args.seed)
+    ).images
+
+    for i, image in enumerate(images):
+        save_path = f"{args.save_prefix}_{i}.png"
+        image.save(save_path)
+        print(f"Generated image: {save_path}")
+
+
+if __name__ == "__main__":
+    main()

tools/run_sdxl_with_custom_components.py

+import copy
+from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
+from migraphx_diffusers import (MIGraphXAutoencoderKL, MIGraphXCLIPTextModel,
+                                MIGraphXCLIPTextModelWithProjection, 
+                                MIGraphXUNet2DConditionModel)
+from transformers import AutoTokenizer
+import numpy as np
+import torch
+
+
+def parse_args():
+    from argparse import ArgumentParser
+    parser = ArgumentParser(description="SDXL inference with migraphx backend")
+
+    #=========================== mdoel load and compile ========================
+    parser.add_argument(
+        "-m", 
+        "--model-dir",
+        type=str,
+        required=True,
+        help="Path to local model directory.",
+    )
+    parser.add_argument(
+        "--force-compile",
+        action="store_true",
+        default=False,
+        help="Ignore existing .mxr files and override them",
+    )
+    parser.add_argument(
+        "--img-size",
+        type=int,
+        default=1024,
+        help="output image size",
+    )
+    parser.add_argument(
+        "--num-images-per-prompt",
+        type=int,
+        default=1,
+        help="The number of images to generate per prompt."
+    )
+    # --------------------------------------------------------------------------
+
+    # =============================== generation ===============================
+    parser.add_argument(
+        "-p",
+        "--prompt",
+        type=str,
+        required=True,
+        help="Prompt for describe image content, style and so on."
+    )
+    parser.add_argument(
+        "-n",
+        "--negative-prompt",
+        type=str,
+        default=None,
+        help="Negative prompt",
+    )
+    parser.add_argument(
+        "-t",
+        "--num-inference-steps",
+        type=int,
+        default=50,
+        help="Number of iteration steps",
+    )
+    parser.add_argument(
+        "--save-prefix",
+        type=str,
+        default="sdxl_output",
+        help="Prefix of path for saving results",
+    )
+    parser.add_argument(
+        "-s",
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed",
+    )
+    # --------------------------------------------------------------------------
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+
+    pipeline_dir = args.model_dir
+    common_args = dict(
+        batch=args.num_images_per_prompt,
+        img_size=args.img_size,
+        model_dtype='fp16',
+        force_compile=args.force_compile,
+    )
+    text_encoder_args = copy.deepcopy(common_args)
+    text_encoder_args['batch'] = 1
+
+    # ========================== load migraphx mdoels ==========================
+    text_encoder = MIGraphXCLIPTextModel.from_pretrained(
+        pipeline_dir, subfolder="text_encoder", **text_encoder_args)
+    text_encoder_2 = MIGraphXCLIPTextModelWithProjection.from_pretrained(
+        pipeline_dir, subfolder="text_encoder_2", **text_encoder_args)
+    unet = MIGraphXUNet2DConditionModel.from_pretrained(
+        pipeline_dir, subfolder="unet", **common_args, 
+        pipeline_class=StableDiffusionXLPipeline)
+    vae = MIGraphXAutoencoderKL.from_pretrained(
+        pipeline_dir, subfolder="vae_decoder", **common_args)
+    # --------------------------------------------------------------------------
+
+    # ============================ load torch models ===========================
+    scheduler = EulerDiscreteScheduler.from_pretrained(
+        pipeline_dir, subfolder="scheduler")
+    tokenizer = AutoTokenizer.from_pretrained(
+        pipeline_dir, subfolder="tokenizer")
+    tokenizer_2 = AutoTokenizer.from_pretrained(
+        pipeline_dir, subfolder="tokenizer_2")
+    # --------------------------------------------------------------------------
+
+    # create pipeline
+    pipe = StableDiffusionXLPipeline(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            unet=unet,
+            scheduler=scheduler,
+            force_zeros_for_empty_prompt=True,
+            add_watermarker=None,
+    )
+    pipe.to("cuda")
+    pipe.to(torch.float16)
+
+    # register configuration
+    pipe.register_to_config(
+        _mgx_models=["text_encoder", "text_encoder_2", "unet", "vae"])
+    pipe.register_to_config(_batch=args.num_images_per_prompt)
+    pipe.register_to_config(_img_height=args.img_size)
+    pipe.register_to_config(_img_width=args.img_size)
+
+    # generate images
+    print("Generating image...")
+    images = pipe(
+        prompt=args.prompt, 
+        negative_prompt=args.negative_prompt, 
+        num_inference_steps=args.num_inference_steps,
+        generator=torch.Generator("cuda").manual_seed(args.seed)
+    ).images
+
+    for i, image in enumerate(images):
+        save_path = f"{args.save_prefix}_{i}.png"
+        image.save(save_path)
+        print(f"Generated image: {save_path}")
+
+
+if __name__ == "__main__":
+    main()

tools/time_count.py

+
+import json
+import os.path as osp
+import time
+
+from diffusers import DiffusionPipeline
+import migraphx_diffusers
+from migraphx_diffusers import AutoTimer
+import torch
+
+
+def parse_args():
+    date_str = time.strftime("%Y%m%d-%H%M%S", time.localtime())
+
+    from argparse import ArgumentParser
+    parser = ArgumentParser(description="SDXL inference with migraphx backend")
+
+    #=========================== mdoel load and compile ========================
+    parser.add_argument(
+        "-m", 
+        "--model-dir",
+        type=str,
+        required=True,
+        help="Path to local model directory.",
+    )
+    parser.add_argument(
+        "--force-compile",
+        action="store_true",
+        default=False,
+        help="Ignore existing .mxr files and override them",
+    )
+    parser.add_argument(
+        "--img-size",
+        type=int,
+        default=None,
+        help="output image size",
+    )
+    parser.add_argument(
+        "--num-images-per-prompt",
+        type=int,
+        default=1,
+        help="The number of images to generate per prompt."
+    )
+    # --------------------------------------------------------------------------
+
+    # =============================== generation ===============================
+    parser.add_argument(
+        "-t",
+        "--num-inference-steps",
+        type=int,
+        default=50,
+        help="Number of iteration steps",
+    )
+    parser.add_argument(
+        "--out-csv-file",
+        type=str,
+        default=f"./perf-{date_str}.csv",
+        help="Prefix of path for saving results",
+    )
+    # --------------------------------------------------------------------------
+
+    # =============================== time count ===============================
+    parser.add_argument(
+        "--num-warmup-loops",
+        type=int,
+        default=1,
+        help="warmup loops",
+    )
+    parser.add_argument(
+        "--num-count-loops",
+        type=int,
+        default=100,
+        help="time count loops",
+    )
+    # --------------------------------------------------------------------------
+
+    args = parser.parse_args()
+    return args
+
+
+def get_name_and_migraphx_config(model_dir):
+    model_index_json = osp.join(model_dir, "model_index.json")
+    with open(model_index_json, "r") as f:
+        pipe_cfg = json.load(f)
+    
+    if pipe_cfg["_class_name"] == "StableDiffusionXLPipeline":
+        return 'sdxl', migraphx_diffusers.DEFAULT_ARGS['sdxl']
+    elif pipe_cfg["_class_name"] == "StableDiffusionPipeline":
+        return 'sd2.1', migraphx_diffusers.DEFAULT_ARGS['sd2.1']
+    else:
+        raise NotImplementedError(
+            f"{pipe_cfg['_class_name']} has not been adapted yet")
+
+
+def main():
+    args = parse_args()
+    pipe_name, migraphx_config = get_name_and_migraphx_config(args.model_dir)
+    assert pipe_name in ['sdxl', 'sd2.1'], "Only support SDXL or SD2.1!"
+
+    if args.img_size is not None:
+        migraphx_config['common_args']['img_size'] = args.img_size
+    migraphx_config['common_args'].update(dict(
+        batch=args.num_images_per_prompt,
+        force_compile=args.force_compile,
+    ))
+
+    pipe = DiffusionPipeline.from_pretrained(
+        args.model_dir,
+        torch_dtype=torch.float16, 
+        migraphx_config=migraphx_config
+    )
+    pipe.to("cuda")
+
+    t = AutoTimer()
+    t.add_targets([
+        (pipe, "end2end"), 
+        (pipe.text_encoder, "text_encoder"), 
+        (pipe.unet, "unet"),
+        (pipe.vae.decode, "vae_decoder")
+    ])
+    if hasattr(pipe, "text_encoder_2"):
+        t.add_target(pipe.text_encoder_2, key="text_encoder_2")
+    
+    for i in range(args.num_warmup_loops + args.num_count_loops):
+        if i == args.num_warmup_loops:
+            t.start_work()
+        pipe(prompt="the ocean in dream", 
+             negative_prompt=None, 
+             num_inference_steps=args.num_inference_steps)
+
+    table = t.summary(batchsize=migraphx_config['common_args']['batch'])
+    t.clear()
+
+    with open(args.out_csv_file, 'w') as f:
+        f.write(table.get_csv_string())
+
+
+if __name__ == "__main__":
+    main()

tools/vae_decoder_export.py

+import argparse
+import os.path as osp
+from diffusers import AutoencoderKL, DiffusionPipeline
+from migraphx_diffusers import ONNXModifier
+import numpy as np
+import onnx
+import torch
+
+
+def export_vae_decoder(pipeline_dir):
+    save_path = osp.join(pipeline_dir, "vae_decoder", "model.onnx")
+    vae = AutoencoderKL.from_pretrained(pipeline_dir, subfolder="vae")
+    # pipe = DiffusionPipeline.from_pretrained(pipeline_dir, use_safetensors=True)
+
+    input_names = ["latent_sample"]
+    output_names = ["sample"]
+    dynamic_axes = {
+        'latent_sample': {
+            0: 'batch_size',
+            2: 'latent_height',
+            3: 'latent_width'
+        },
+        'latent': {
+            0: 'batch_size',
+            2: 'image_height',
+            3: 'image_width'
+        }
+    }
+
+    vae.forward = vae.decode
+    # pipe.vae.forward = pipe.vae.decode
+
+    torch.onnx.export(
+        vae,
+        # pipe.vae, 
+        torch.randn(1, vae.config["latent_channels"], 128, 128), 
+        # torch.randn(1, pipe.vae.config["latent_channels"], 128, 128), 
+        save_path ,
+        export_params=True,
+        input_names=input_names,
+        output_names=output_names,
+        dynamic_axes=dynamic_axes
+    )
+
+    if osp.isfile(save_path):
+        print(f"Successfully exported vae decoder to ONNX: {save_path}")
+    else:
+        raise RuntimeError(f"Failed to export vae decoder to ONNX.")
+
+
+    return save_path
+
+
+def modify_vae_decoder(onnx_path, save_path=None):
+    om = ONNXModifier(onnx_path)
+
+    shape_node = om.get_node("/decoder/mid_block/attentions.0/Shape")
+    C_gather = om.get_node("/decoder/mid_block/attentions.0/Gather_1")
+    H_gather = om.get_node("/decoder/mid_block/attentions.0/Gather_2")
+    W_gather = om.get_node("/decoder/mid_block/attentions.0/Gather_3")
+    C_gather.set_input(0, shape_node.outputs[0])
+    H_gather.set_input(0, shape_node.outputs[0])
+    W_gather.set_input(0, shape_node.outputs[0])
+
+    matmul1 = om.get_node("/decoder/mid_block/attentions.0/to_k/MatMul")
+    matmul1.set_input(0, matmul1.inputs[1])
+    matmul1.set_input(1, "/decoder/mid_block/attentions.0/group_norm/Add_output_0")
+    matmul1_ini = om.get_initializer("onnx::MatMul_918")
+    orig_dtype = onnx.helper.tensor_dtype_to_np_dtype(matmul1_ini.data_type)
+    orig_shape = matmul1_ini.dims[:]
+    new_value = np.frombuffer(matmul1_ini.raw_data, 
+                              dtype=orig_dtype).reshape(orig_shape).T
+    matmul1_ini.raw_data = new_value.tobytes()
+    matmul1_ini.dims[:] = orig_shape[::-1]
+    add_ini = om.get_initializer("decoder.mid_block.attentions.0.to_k.bias")
+    add_ini.dims[:] = [512, 1]
+
+    mul1 = om.get_node("/decoder/mid_block/attentions.0/Mul_1")
+    mul2 = om.get_node("/decoder/mid_block/attentions.0/Mul_2")
+    mul1.set_input(0, "/decoder/mid_block/attentions.0/to_q/Add_output_0")
+    mul2.set_input(0, "/decoder/mid_block/attentions.0/to_k/Add_output_0")
+
+    value_B = om.create_initializer("/decoder/mid_block/attentions.0/mul_B", 
+                                    np.sqrt(1 / np.sqrt(512)).astype(np.float32)[None])
+    mul1.set_input(1, value_B.name)
+    mul2.set_input(1, value_B.name)
+
+    matmul2 = om.get_node("/decoder/mid_block/attentions.0/MatMul_1")
+    matmul2.set_input(1, "/decoder/mid_block/attentions.0/to_v/Add_output_0")
+
+    matmul3 = om.get_node("/decoder/mid_block/attentions.0/to_out.0/MatMul")
+    matmul3.set_input(0, matmul2.outputs[0])
+
+    reshape_node = om.get_node("/decoder/mid_block/attentions.0/Reshape_5")
+    reshape_node.set_input(1, shape_node.outputs[0])
+
+    om.update_map()
+    om.remove_trash()
+    save_path = onnx_path if save_path is None else save_path
+    om.save(save_path)
+
+
+def main():
+    parser = argparse.ArgumentParser("Export vae decoder to ONNX")
+    parser.add_argument("--pipeline-dir", 
+                        type=str, 
+                        required=True, 
+                        help="The path to the sdxl pipeline directory.")
+    args = parser.parse_args()
+
+    onnx_path = export_vae_decoder(args.pipeline_dir)
+    modify_vae_decoder(onnx_path)
+
+
+if __name__ == "__main__":
+    main()