Fix formatting

Author: monorimet

models/turbine_models/custom_models/sd3_inference/sd3_mmdit.py

@@ -52,10 +52,11 @@ def forward(
             return_dict=False,
         )[0]
         return noise_pred
-    
+
+
 class MMDiTAttention(torch.nn.Module):
     def __init__(
-            self,
+        self,
     ):
         super().__init__()
 
@@ -84,7 +85,7 @@ def export_attn(
 
     if dtype == torch.float16:
         attn_module = attn_module.half()
-    
+
     example_qkv = [
         torch.empty(qkv_shape, dtype=dtype),
         torch.empty(qkv_shape, dtype=dtype),
@@ -134,6 +135,7 @@ class CompiledAttn(CompiledModule):
         )
     return vmfb_path
 
+
 @torch.no_grad()
 def export_mmdit_model(
     mmdit_model,

models/turbine_models/custom_models/sd3_inference/sd3_mmdit_runner.py

@@ -57,6 +57,7 @@ def run_diffusers_mmdit(
 
     return noise_pred.numpy()
 
+
 def run_attn_turbine(q, k, v, args):
     attn_runner = vmfbRunner(
         args.device,
@@ -73,6 +74,7 @@ def run_attn_turbine(q, k, v, args):
     ).to_host()
     return attn_output
 
+
 @torch.no_grad()
 def run_attn_torch(q, k, v, args):
     from turbine_models.custom_models.sd3_inference.sd3_mmdit import MMDiTAttention
@@ -86,21 +88,27 @@ def run_attn_torch(q, k, v, args):
 
     return attn_output.numpy()
 
+
 def find_errs(turbine_output, torch_output, dim=[], failed_dims=[], errs=[]):
     if not np.allclose(turbine_output, torch_output, rtol=4e-2, atol=4e-2):
         if turbine_output.ndim > 0:
             orig_dim = dim
             for idx, i in enumerate(torch_output):
                 dim = [*orig_dim, idx]
                 try:
-                    np.testing.assert_allclose(turbine_output[idx], torch_output[idx], rtol=4e-2, atol=4e-2)
+                    np.testing.assert_allclose(
+                        turbine_output[idx], torch_output[idx], rtol=4e-2, atol=4e-2
+                    )
                 except Exception as e:
                     err = np.abs(turbine_output[idx] - torch_output[idx])
                     failed_dims.append(dim)
                     errs.append([err, turbine_output[idx], torch_output[idx]])
-                    failed_dims, errs = find_errs(turbine_output[idx], torch_output[idx], dim, failed_dims, errs)
+                    failed_dims, errs = find_errs(
+                        turbine_output[idx], torch_output[idx], dim, failed_dims, errs
+                    )
     return (failed_dims, errs)
 
+
 if __name__ == "__main__":
     from turbine_models.custom_models.sd3_inference.sd3_cmd_opts import args
     import numpy as np
@@ -137,8 +145,8 @@ def find_errs(turbine_output, torch_output, dim=[], failed_dims=[], errs=[]):
                 print("Torch output: ", errs[idx][2])
         print(torch_output.shape)
         exit()
-    
-    batch_size = args.batch_size * 2 #do classifier free guidance
+
+    batch_size = args.batch_size * 2  # do classifier free guidance
     hidden_states = torch.randn(
         (batch_size, 16, args.height // 8, args.width // 8), dtype=dtype
     )

models/turbine_models/custom_models/sd3_inference/sd3_pipeline.py

@@ -75,38 +75,42 @@ def __init__(
         self.num_inference_steps = num_inference_steps
         self.devices = {}
         if isinstance(device, dict):
-            assert isinstance(iree_target_triple, dict), "Device and target triple must be both dicts or both strings."
+            assert isinstance(
+                iree_target_triple, dict
+            ), "Device and target triple must be both dicts or both strings."
             self.devices["clip"] = {
                 "device": device["clip"],
                 "driver": utils.iree_device_map(device["clip"]),
-                "target": iree_target_triple["clip"]
+                "target": iree_target_triple["clip"],
             }
             self.devices["mmdit"] = {
                 "device": device["mmdit"],
                 "driver": utils.iree_device_map(device["mmdit"]),
-                "target": iree_target_triple["mmdit"]
+                "target": iree_target_triple["mmdit"],
             }
             self.devices["vae"] = {
                 "device": device["vae"],
                 "driver": utils.iree_device_map(device["vae"]),
-                "target": iree_target_triple["vae"]
+                "target": iree_target_triple["vae"],
             }
         else:
-            assert isinstance(iree_target_triple, str), "Device and target triple must be both dicts or both strings."
+            assert isinstance(
+                iree_target_triple, str
+            ), "Device and target triple must be both dicts or both strings."
             self.devices["clip"] = {
                 "device": device,
                 "driver": utils.iree_device_map(device),
-                "target": iree_target_triple
+                "target": iree_target_triple,
             }
             self.devices["mmdit"] = {
                 "device": device,
                 "driver": utils.iree_device_map(device),
-                "target": iree_target_triple
+                "target": iree_target_triple,
             }
             self.devices["vae"] = {
                 "device": device,
                 "driver": utils.iree_device_map(device),
-                "target": iree_target_triple
+                "target": iree_target_triple,
             }
         self.iree_target_triple = iree_target_triple
         self.ireec_flags = ireec_flags if ireec_flags else EMPTY_FLAGS
@@ -645,7 +649,8 @@ def generate_images(
                 image.save(img_path)
                 print(img_path, "saved")
         return
-    
+
+
 def run_diffusers_cpu(
     hf_model_name,
     prompt,
@@ -658,7 +663,9 @@ def run_diffusers_cpu(
 ):
     from diffusers import StableDiffusion3Pipeline
 
-    pipe = StableDiffusion3Pipeline.from_pretrained(hf_model_name, torch_dtype=torch.float32)
+    pipe = StableDiffusion3Pipeline.from_pretrained(
+        hf_model_name, torch_dtype=torch.float32
+    )
     pipe = pipe.to("cpu")
     generator = torch.Generator().manual_seed(int(seed))
 
@@ -703,7 +710,9 @@ def run_diffusers_cpu(
             x for x in [args.clip_target, args.mmdit_target, args.vae_target]
         ), "Please specify target triple for all submodels or pass --iree_target_triple for all submodels."
         args.device = "hybrid"
-        args.iree_target_triple = "_".join([args.clip_target, args.mmdit_target, args.vae_target])
+        args.iree_target_triple = "_".join(
+            [args.clip_target, args.mmdit_target, args.vae_target]
+        )
     else:
         args.clip_device = args.device
         args.mmdit_device = args.device
@@ -785,7 +794,11 @@ def run_diffusers_cpu(
     else:
         extra_device_args = {}
     sd3_pipe.load_pipeline(
-        vmfbs, weights, args.compiled_pipeline, args.split_scheduler, extra_device_args=extra_device_args
+        vmfbs,
+        weights,
+        args.compiled_pipeline,
+        args.split_scheduler,
+        extra_device_args=extra_device_args,
     )
     sd3_pipe.generate_images(
         args.prompt,

models/turbine_models/custom_models/sd3_inference/sd3_schedulers.py

@@ -66,7 +66,7 @@ def __init__(
     def initialize(self, sample):
         step_count = torch.tensor(len(self.timesteps))
         timesteps = self.model.timesteps
-        #ops.trace_tensor("sample", sample[:,:,0,0])
+        # ops.trace_tensor("sample", sample[:,:,0,0])
         return (
             sample,
             step_count,
@@ -93,6 +93,7 @@ def step(self, noise_pred, t, sample, guidance_scale, i):
         sample = self.model.step(noise_pred, t, sample, return_dict=False)[0]
         return sample.type(self.dtype)
 
+
 # Wraps a diffusers scheduler running on native pytorch+cpu.
 # This allows us to use it interchangeably with compiled schedulers in our pipeline(s).
 class TorchCPUFlowSchedulerCompat:

models/turbine_models/custom_models/sd_inference/utils.py

@@ -92,6 +92,7 @@
     ],
 }
 
+
 def iree_device_map(device):
     uri_parts = device.split("://", 2)
     iree_driver = (
@@ -106,6 +107,7 @@ def iree_device_map(device):
     else:
         return f"{iree_driver}://{uri_parts[1]}"
 
+
 def compile_to_vmfb(
     module_str,
     device,

models/turbine_models/custom_models/sdxl_inference/sdxl_compiled_pipeline.py

@@ -78,38 +78,42 @@ def __init__(
         self.num_inference_steps = num_inference_steps
         self.devices = {}
         if isinstance(device, dict):
-            assert isinstance(iree_target_triple, dict), "Device and target triple must be both dicts or both strings."
+            assert isinstance(
+                iree_target_triple, dict
+            ), "Device and target triple must be both dicts or both strings."
             self.devices["clip"] = {
                 "device": device["clip"],
                 "driver": utils.iree_device_map(device["clip"]),
-                "target": iree_target_triple["clip"]
+                "target": iree_target_triple["clip"],
             }
             self.devices["unet"] = {
                 "device": device["unet"],
                 "driver": utils.iree_device_map(device["unet"]),
-                "target": iree_target_triple["unet"]
+                "target": iree_target_triple["unet"],
             }
             self.devices["vae"] = {
                 "device": device["vae"],
                 "driver": utils.iree_device_map(device["vae"]),
-                "target": iree_target_triple["vae"]
+                "target": iree_target_triple["vae"],
             }
         else:
-            assert isinstance(iree_target_triple, str), "Device and target triple must be both dicts or both strings."
+            assert isinstance(
+                iree_target_triple, str
+            ), "Device and target triple must be both dicts or both strings."
             self.devices["clip"] = {
                 "device": device,
                 "driver": utils.iree_device_map(device),
-                "target": iree_target_triple
+                "target": iree_target_triple,
             }
             self.devices["unet"] = {
                 "device": device,
                 "driver": utils.iree_device_map(device),
-                "target": iree_target_triple
+                "target": iree_target_triple,
             }
             self.devices["vae"] = {
                 "device": device,
                 "driver": utils.iree_device_map(device),
-                "target": iree_target_triple
+                "target": iree_target_triple,
             }
         self.ireec_flags = ireec_flags if ireec_flags else EMPTY_FLAGS
         self.attn_spec = attn_spec
@@ -548,8 +552,14 @@ def load_pipeline(
             clip_loaded = time.time()
             print("\n[LOG] CLIP loaded in ", clip_loaded - vae_loaded, "sec")
         elif compiled_pipeline:
-            assert self.devices["unet"]["device"] == self.devices["clip"]["device"] == self.devices["vae"]["device"], "Compiled pipeline requires all submodels to be on the same device."
-            assert self.precision == self.vae_precision, "Compiled pipeline requires all submodels to have the same precision for now."
+            assert (
+                self.devices["unet"]["device"]
+                == self.devices["clip"]["device"]
+                == self.devices["vae"]["device"]
+            ), "Compiled pipeline requires all submodels to be on the same device."
+            assert (
+                self.precision == self.vae_precision
+            ), "Compiled pipeline requires all submodels to have the same precision for now."
             runners["pipe"] = vmfbRunner(
                 self.devices["unet"]["driver"],
                 [
@@ -796,9 +806,14 @@ def generate_images(
                     latents = self.runners["pipe"].ctx.modules.sdxl_compiled_pipeline[
                         "produce_image_latents"
                     ](samples[i], prompt_embeds, add_text_embeds, guidance_scale)
-                if self.devices["unet"]["driver"] != self.devices["vae"]["driver"] or self.precision != self.vae_precision:
+                if (
+                    self.devices["unet"]["driver"] != self.devices["vae"]["driver"]
+                    or self.precision != self.vae_precision
+                ):
                     latents = ireert.asdevicearray(
-                        self.runners["vae_decode"].config.device, latents.to_host(), dtype=self.vae_dtype
+                        self.runners["vae_decode"].config.device,
+                        latents.to_host(),
+                        dtype=self.vae_dtype,
                     )
                 vae_start = time.time()
                 vae_out = self.runners["vae_decode"].ctx.modules.compiled_vae["main"](
@@ -906,7 +921,9 @@ def numpy_to_pil_image(images):
             x for x in [args.clip_target, args.unet_target, args.vae_target]
         ), "Please specify target triple for all submodels or pass --iree_target_triple for all submodels."
         args.device = "hybrid"
-        args.iree_target_triple = "_".join([args.clip_target, args.unet_target, args.vae_target])
+        args.iree_target_triple = "_".join(
+            [args.clip_target, args.unet_target, args.vae_target]
+        )
     else:
         args.clip_device = args.device
         args.unet_device = args.device
@@ -987,7 +1004,11 @@ def numpy_to_pil_image(images):
     else:
         extra_device_args = {}
     sdxl_pipe.load_pipeline(
-        vmfbs, weights, args.compiled_pipeline, args.split_scheduler, extra_device_args,
+        vmfbs,
+        weights,
+        args.compiled_pipeline,
+        args.split_scheduler,
+        extra_device_args,
     )
     sdxl_pipe.generate_images(
         args.prompt,

models/turbine_models/custom_models/sdxl_inference/vae.py

@@ -107,7 +107,7 @@ def export_vae_model(
 
     if device == "cpu":
         decomp_attn = True
-    
+
     dtype = torch.float16 if precision == "fp16" else torch.float32
     if precision == "fp16":
         vae_model = vae_model.half()
@@ -119,7 +119,7 @@ def export_vae_model(
     )
     if weights_only:
         return external_weight_path
-    
+
     input_image_shape = (height, width, 3)
     input_latents_shape = (batch_size, 4, height // 8, width // 8)
     encode_args = [