SD PNDMScheduler + Unet example through Turbine (#403)

TODO: Need to update the rest of the schedulers in diffusers upstream for e2e test to work. Xfailed for now.

core/shark_turbine/dynamo/passes.py

@@ -48,6 +48,7 @@
     torch.ops.aten._log_softmax_backward_data,
     torch.ops.aten.lift_fresh_copy.default,
     torch.ops.aten._unsafe_index.Tensor,
+    torch.ops.aten.unbind.int,
     # decompositions added manually in this file
     torch.ops.aten._scaled_dot_product_flash_attention.default,
 ]

models/turbine_models/custom_models/sd_inference/schedulers.py

@@ -0,0 +1,178 @@
+# Copyright 2024 Advanced Micro Devices, Inc
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+import os
+import sys
+
+import torch
+from torch.fx.experimental.proxy_tensor import make_fx
+from shark_turbine.aot import *
+from iree import runtime as ireert
+import iree.compiler as ireec
+from iree.compiler.ir import Context
+import numpy as np
+
+from turbine_models.custom_models.sd_inference import utils
+from diffusers import (
+    UNet2DConditionModel,
+)
+
+import safetensors
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--hf_auth_token", type=str, help="The Hugging Face auth token, required"
+)
+parser.add_argument(
+    "--hf_model_name",
+    type=str,
+    help="HF model name",
+    default="CompVis/stable-diffusion-v1-4",
+)
+parser.add_argument(
+    "--scheduler_id",
+    type=str,
+    help="Scheduler ID",
+    default="PNDM",
+)
+parser.add_argument(
+    "--num_inference_steps", type=int, default=50, help="Number of inference steps"
+)
+parser.add_argument(
+    "--batch_size", type=int, default=1, help="Batch size for inference"
+)
+parser.add_argument(
+    "--height", type=int, default=512, help="Height of Stable Diffusion"
+)
+parser.add_argument("--width", type=int, default=512, help="Width of Stable Diffusion")
+parser.add_argument("--compile_to", type=str, help="torch, linalg, vmfb")
+parser.add_argument("--external_weight_path", type=str, default="")
+parser.add_argument(
+    "--external_weights",
+    type=str,
+    default=None,
+    help="saves ir/vmfb without global weights for size and readability, options [safetensors]",
+)
+parser.add_argument("--device", type=str, default="cpu", help="cpu, cuda, vulkan, rocm")
+# TODO: Bring in detection for target triple
+parser.add_argument(
+    "--iree_target_triple",
+    type=str,
+    default="",
+    help="Specify vulkan target triple or rocm/cuda target device.",
+)
+parser.add_argument("--vulkan_max_allocation", type=str, default="4294967296")
+
+
+class Scheduler(torch.nn.Module):
+    def __init__(self, hf_model_name, num_inference_steps, scheduler):
+        super().__init__()
+        self.scheduler = scheduler
+        self.scheduler.set_timesteps(num_inference_steps)
+        self.unet = UNet2DConditionModel.from_pretrained(
+            hf_model_name,
+            subfolder="unet",
+        )
+        self.guidance_scale = 7.5
+
+    def forward(self, latents, encoder_hidden_states) -> torch.FloatTensor:
+        latents = latents * self.scheduler.init_noise_sigma
+        for t in self.scheduler.timesteps:
+            latent_model_input = torch.cat([latents] * 2)
+            t = t.unsqueeze(0)
+            latent_model_input = self.scheduler.scale_model_input(
+                latent_model_input, timestep=t
+            )
+            unet_out = self.unet.forward(
+                latent_model_input, t, encoder_hidden_states, return_dict=False
+            )[0]
+            noise_pred_uncond, noise_pred_text = unet_out.chunk(2)
+            noise_pred = noise_pred_uncond + self.guidance_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
+            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+        return latents
+
+
+def export_scheduler(
+    scheduler,
+    hf_model_name,
+    batch_size,
+    height,
+    width,
+    hf_auth_token=None,
+    compile_to="torch",
+    external_weights=None,
+    external_weight_path=None,
+    device=None,
+    target_triple=None,
+    max_alloc=None,
+):
+    mapper = {}
+    utils.save_external_weights(
+        mapper, scheduler, external_weights, external_weight_path
+    )
+
+    encoder_hidden_states_sizes = (2, 77, 768)
+    if hf_model_name == "stabilityai/stable-diffusion-2-1-base":
+        encoder_hidden_states_sizes = (2, 77, 1024)
+
+    sample = (batch_size, 4, height // 8, width // 8)
+
+    class CompiledScheduler(CompiledModule):
+        if external_weights:
+            params = export_parameters(
+                scheduler, external=True, external_scope="", name_mapper=mapper.get
+            )
+        else:
+            params = export_parameters(scheduler)
+
+        def main(
+            self,
+            sample=AbstractTensor(*sample, dtype=torch.float32),
+            encoder_hidden_states=AbstractTensor(
+                *encoder_hidden_states_sizes, dtype=torch.float32
+            ),
+        ):
+            return jittable(scheduler.forward)(sample, encoder_hidden_states)
+
+    import_to = "INPUT" if compile_to == "linalg" else "IMPORT"
+    inst = CompiledScheduler(context=Context(), import_to=import_to)
+
+    module_str = str(CompiledModule.get_mlir_module(inst))
+    safe_name = utils.create_safe_name(hf_model_name, "-scheduler")
+    if compile_to != "vmfb":
+        return module_str
+    else:
+        utils.compile_to_vmfb(module_str, device, target_triple, max_alloc, safe_name)
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    schedulers = utils.get_schedulers(args.hf_model_name)
+    scheduler = schedulers[args.scheduler_id]
+    scheduler_module = Scheduler(
+        args.hf_model_name, args.num_inference_steps, scheduler
+    )
+    mod_str = export_scheduler(
+        scheduler_module,
+        args.hf_model_name,
+        args.batch_size,
+        args.height,
+        args.width,
+        args.hf_auth_token,
+        args.compile_to,
+        args.external_weights,
+        args.external_weight_path,
+        args.device,
+        args.iree_target_triple,
+        args.vulkan_max_allocation,
+    )
+    safe_name = utils.create_safe_name(args.hf_model_name, "-scheduler")
+    with open(f"{safe_name}.mlir", "w+") as f:
+        f.write(mod_str)
+    print("Saved to", safe_name + ".mlir")

models/turbine_models/custom_models/sd_inference/schedulers_runner.py

@@ -0,0 +1,172 @@
+# Copyright 2024 Advanced Micro Devices, Inc
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+import argparse
+from turbine_models.model_runner import vmfbRunner
+from iree import runtime as ireert
+import torch
+from diffusers import (
+    PNDMScheduler,
+    UNet2DConditionModel,
+)
+
+parser = argparse.ArgumentParser()
+
+# TODO move common runner flags to generic flag file
+parser.add_argument(
+    "--scheduler_id",
+    type=str,
+    help="Scheduler ID",
+    default="PNDM",
+)
+parser.add_argument(
+    "--num_inference_steps", type=int, default=50, help="Number of inference steps"
+)
+parser.add_argument(
+    "--vmfb_path", type=str, default="", help="path to vmfb containing compiled module"
+)
+parser.add_argument(
+    "--external_weight_path",
+    type=str,
+    default="",
+    help="path to external weight parameters if model compiled without them",
+)
+parser.add_argument(
+    "--compare_vs_torch",
+    action="store_true",
+    help="Runs both turbine vmfb and a torch model to compare results",
+)
+parser.add_argument(
+    "--hf_model_name",
+    type=str,
+    help="HF model name",
+    default="CompVis/stable-diffusion-v1-4",
+)
+parser.add_argument(
+    "--hf_auth_token",
+    type=str,
+    help="The Hugging face auth token, required for some models",
+)
+parser.add_argument(
+    "--device",
+    type=str,
+    default="local-task",
+    help="local-sync, local-task, cuda, vulkan, rocm",
+)
+parser.add_argument(
+    "--batch_size", type=int, default=1, help="Batch size for inference"
+)
+parser.add_argument(
+    "--height", type=int, default=512, help="Height of Stable Diffusion"
+)
+parser.add_argument("--width", type=int, default=512, help="Width of Stable Diffusion")
+
+
+def run_scheduler(
+    device,
+    sample,
+    encoder_hidden_states,
+    vmfb_path,
+    hf_model_name,
+    hf_auth_token,
+    external_weight_path,
+):
+    runner = vmfbRunner(device, vmfb_path, external_weight_path)
+
+    inputs = [
+        ireert.asdevicearray(runner.config.device, sample),
+        ireert.asdevicearray(runner.config.device, encoder_hidden_states),
+    ]
+    results = runner.ctx.modules.compiled_scheduler["main"](*inputs)
+    return results
+
+
+def run_torch_scheduler(
+    hf_model_name, scheduler, num_inference_steps, sample, encoder_hidden_states
+):
+    class Scheduler(torch.nn.Module):
+        def __init__(self, hf_model_name, num_inference_steps, scheduler):
+            super().__init__()
+            self.scheduler = scheduler
+            self.scheduler.set_timesteps(num_inference_steps)
+            self.unet = UNet2DConditionModel.from_pretrained(
+                hf_model_name,
+                subfolder="unet",
+            )
+            self.guidance_scale = 7.5
+
+        def forward(self, latents, encoder_hidden_states) -> torch.FloatTensor:
+            latents = latents * self.scheduler.init_noise_sigma
+            for t in self.scheduler.timesteps:
+                latent_model_input = torch.cat([latents] * 2)
+                t = t.unsqueeze(0)
+                latent_model_input = self.scheduler.scale_model_input(
+                    latent_model_input, timestep=t
+                )
+                unet_out = self.unet.forward(
+                    latent_model_input, t, encoder_hidden_states, return_dict=False
+                )[0]
+                noise_pred_uncond, noise_pred_text = unet_out.chunk(2)
+                noise_pred = noise_pred_uncond + self.guidance_scale * (
+                    noise_pred_text - noise_pred_uncond
+                )
+                latents = self.scheduler.step(
+                    noise_pred, t, latents, return_dict=False
+                )[0]
+            return latents
+
+    scheduler_module = Scheduler(hf_model_name, num_inference_steps, scheduler)
+    results = scheduler_module.forward(sample, encoder_hidden_states)
+    np_torch_output = results.detach().cpu().numpy()
+    return np_torch_output
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    sample = torch.rand(
+        args.batch_size, 4, args.height // 8, args.width // 8, dtype=torch.float32
+    )
+    if args.hf_model_name == "CompVis/stable-diffusion-v1-4":
+        encoder_hidden_states = torch.rand(2, 77, 768, dtype=torch.float32)
+    elif args.hf_model_name == "stabilityai/stable-diffusion-2-1-base":
+        encoder_hidden_states = torch.rand(2, 77, 1024, dtype=torch.float32)
+
+    turbine_output = run_scheduler(
+        args.device,
+        sample,
+        encoder_hidden_states,
+        args.vmfb_path,
+        args.hf_model_name,
+        args.hf_auth_token,
+        args.external_weight_path,
+    )
+    print(
+        "TURBINE OUTPUT:",
+        turbine_output.to_host(),
+        turbine_output.to_host().shape,
+        turbine_output.to_host().dtype,
+    )
+
+    if args.compare_vs_torch:
+        print("generating torch output: ")
+        from turbine_models.custom_models.sd_inference import utils
+
+        schedulers = utils.get_schedulers(args.hf_model_name)
+        scheduler = schedulers[args.scheduler_id]
+        torch_output = run_torch_scheduler(
+            args.hf_model_name,
+            scheduler,
+            args.num_inference_steps,
+            sample,
+            encoder_hidden_states,
+        )
+        print("TORCH OUTPUT:", torch_output, torch_output.shape, torch_output.dtype)
+        err = utils.largest_error(torch_output, turbine_output)
+        print("Largest Error: ", err)
+        assert err < 9e-3
+
+    # TODO: Figure out why we occasionally segfault without unlinking output variables
+    turbine_output = None

models/turbine_models/custom_models/sd_inference/utils.py

@@ -2,6 +2,9 @@
 import numpy as np
 import safetensors
 import re
+from diffusers import (
+    PNDMScheduler,
+)
 
 
 def save_external_weights(
@@ -35,6 +38,7 @@ def compile_to_vmfb(module_str, device, target_triple, max_alloc, safe_name):
         "--iree-llvmcpu-target-triple=x86_64-linux-gnu",
         "--iree-stream-resource-index-bits=64",
         "--iree-vm-target-index-bits=64",
+        "--iree-flow-inline-constants-max-byte-length=1",
     ]
     if device == "cpu":
         flags.append("--iree-llvmcpu-enable-ukernels=all")
@@ -86,3 +90,21 @@ def create_safe_name(hf_model_name, model_name_str):
     safe_name = hf_model_name.split("/")[-1].strip() + model_name_str
     safe_name = re.sub("-", "_", safe_name)
     return safe_name
+
+
+def get_schedulers(model_id):
+    # TODO: Robust scheduler setup on pipeline creation -- if we don't
+    # set batch_size here, the SHARK schedulers will
+    # compile with batch size = 1 regardless of whether the model
+    # outputs latents of a larger batch size, e.g. SDXL.
+    # However, obviously, searching for whether the base model ID
+    # contains "xl" is not very robust.
+
+    batch_size = 2 if "xl" in model_id.lower() else 1
+
+    schedulers = dict()
+    schedulers["PNDM"] = PNDMScheduler.from_pretrained(
+        model_id,
+        subfolder="scheduler",
+    )
+    return schedulers