Start fixing tests

Author: monorimet

models/turbine_models/custom_models/sd_inference/schedulers.py

@@ -41,7 +41,9 @@ def __init__(self, rt_device, vmfb):
         self.runner = vmfbRunner(rt_device, vmfb, None)
 
     def initialize(self, sample):
-        sample, time_ids, steps, timesteps = self.runner.ctx.modules.compiled_scheduler["run_initialize"](sample)
+        sample, time_ids, steps, timesteps = self.runner.ctx.modules.compiled_scheduler[
+            "run_initialize"
+        ](sample)
         return sample, time_ids, steps.to_host(), timesteps
 
     def scale_model_input(self, sample, t, timesteps):
@@ -50,11 +52,6 @@ def scale_model_input(self, sample, t, timesteps):
         )
 
     def step(self, noise_pred, t, sample, guidance_scale, step_index):
-        print(
-            noise_pred.to_host()[:,:,0,2],
-            t,
-            sample.to_host()[:,:,0,2],
-        )
         return self.runner.ctx.modules.compiled_scheduler["run_step"](
             noise_pred, t, sample, guidance_scale, step_index
         )
@@ -128,6 +125,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)
 
+
 class SharkSchedulerCPUWrapper:
     @torch.no_grad()
     def __init__(
@@ -183,11 +181,13 @@ def step(self, noise_pred, t, latents, guidance_scale, i):
         noise_pred = torch.tensor(noise_pred.to_host())
         if self.do_classifier_free_guidance:
             noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+            noise_pred = noise_pred_uncond + guidance_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
         print(
-            noise_pred[:,:,0,2],
+            noise_pred[:, :, 0, 2],
             t,
-            latents[:,:,0,2],
+            latents[:, :, 0, 2],
         )
         return self.module.step(
             noise_pred,

models/turbine_models/custom_models/sdxl_inference/sdxl_compiled_pipeline.py

@@ -132,7 +132,9 @@ def check_prepared(
                         vmfbs[submodel] = vmfb
                         if weights[submodel] is None:
                             weights[submodel] = weight
-                    elif weights[submodel] is None and not any(x in submodel for x in ["pipeline", "scheduler"]):
+                    elif weights[submodel] is None and not any(
+                        x in submodel for x in ["pipeline", "scheduler"]
+                    ):
                         _, weight = self.export_submodel(submodel, weights_only=True)
                         weights[submodel] = weight
                 ready, vmfbs, weights = self.is_prepared(vmfbs, weights)
@@ -157,7 +159,7 @@ def is_prepared(self, vmfbs, weights):
                 default_filepath = os.path.join(self.pipeline_dir, val + ".vmfb")
             elif key == "scheduler":
                 val = None
-                default_filepath=None
+                default_filepath = None
                 continue
             else:
                 val = vmfbs[key]
@@ -494,7 +496,9 @@ def load_pipeline(
                 )
             else:
                 print("\n[LOG] Running scheduler on CPU. This will affect performance.")
-                scheduler = schedulers.get_scheduler(args.hf_model_name, args.scheduler_id)
+                scheduler = schedulers.get_scheduler(
+                    args.hf_model_name, args.scheduler_id
+                )
                 runners["scheduler"] = schedulers.SharkSchedulerCPUWrapper(
                     scheduler,
                     args.batch_size,
@@ -535,7 +539,9 @@ def load_pipeline(
                 ],
             )
             pipe_loaded = time.time()
-            print("\n[LOG] Compiled Pipeline loaded in ", pipe_loaded - load_start, "sec")
+            print(
+                "\n[LOG] Compiled Pipeline loaded in ", pipe_loaded - load_start, "sec"
+            )
 
         else:
             runners["pipe"] = vmfbRunner(
@@ -556,7 +562,9 @@ def load_pipeline(
             runners["vae_decode"] = runners["pipe"]
             runners["prompt_encoder"] = runners["pipe"]
             pipe_loaded = time.time()
-            print("\n[LOG] Compiled Pipeline loaded in ", pipe_loaded - load_start, "sec")
+            print(
+                "\n[LOG] Compiled Pipeline loaded in ", pipe_loaded - load_start, "sec"
+            )
         tok_start = time.time()
         runners["tokenizer_1"] = CLIPTokenizer.from_pretrained(
             self.hf_model_name,
@@ -704,11 +712,17 @@ def generate_images(
             for i in range(batch_count):
                 unet_start = time.time()
                 if self.runners["scheduler"]:
-                    sample, time_ids, steps, timesteps = self.runners["scheduler"].initialize(samples[i])
+                    sample, time_ids, steps, timesteps = self.runners[
+                        "scheduler"
+                    ].initialize(samples[i])
                     iree_inputs = [
                         sample,
-                        ireert.asdevicearray(self.runners["pipe"].config.device, prompt_embeds),
-                        ireert.asdevicearray(self.runners["pipe"].config.device, add_text_embeds),
+                        ireert.asdevicearray(
+                            self.runners["pipe"].config.device, prompt_embeds
+                        ),
+                        ireert.asdevicearray(
+                            self.runners["pipe"].config.device, add_text_embeds
+                        ),
                         time_ids,
                         None,
                     ]
@@ -717,13 +731,19 @@ def generate_images(
                         if self.cpu_scheduling:
                             step_index = s
                         else:
-                            step_index = ireert.asdevicearray(self.runners["scheduler"].runner.config.device, torch.tensor([s]), "int64")
+                            step_index = ireert.asdevicearray(
+                                self.runners["scheduler"].runner.config.device,
+                                torch.tensor([s]),
+                                "int64",
+                            )
                         latents, t = self.runners["scheduler"].scale_model_input(
                             sample,
                             step_index,
                             timesteps,
                         )
-                        noise_pred = self.runners["pipe"].ctx.modules.compiled_unet["run_forward"](
+                        noise_pred = self.runners["pipe"].ctx.modules.compiled_unet[
+                            "run_forward"
+                        ](
                             latents,
                             t,
                             iree_inputs[1],
@@ -738,9 +758,13 @@ def generate_images(
                             step_index,
                         )
                     if isinstance(sample, torch.Tensor):
-                        #TODO: pipe an option for vae_dtype
+                        # TODO: pipe an option for vae_dtype
                         vae_dtype = "float32" if self.precision == "fp32" else "float16"
-                        latents = ireert.asdevicearray(self.runners["vae_decode"].config.device, sample, dtype=vae_dtype)
+                        latents = ireert.asdevicearray(
+                            self.runners["vae_decode"].config.device,
+                            sample,
+                            dtype=vae_dtype,
+                        )
                     else:
                         latents = sample
                 else:
@@ -833,6 +857,7 @@ def numpy_to_pil_image(images):
 
 if __name__ == "__main__":
     from turbine_models.custom_models.sdxl_inference.sdxl_cmd_opts import args
+
     map = empty_pipe_dict
     if args.split_scheduler:
         map["scheduler"] = None
@@ -894,13 +919,15 @@ def numpy_to_pil_image(images):
         args.external_weights_dir,
         args.external_weights,
         args.vae_decomp_attn,
-        custom_vae = None,
-        cpu_scheduling = args.cpu_scheduling,
+        custom_vae=None,
+        cpu_scheduling=args.cpu_scheduling,
     )
     vmfbs, weights = sdxl_pipe.check_prepared(mlirs, vmfbs, weights)
     if args.cpu_scheduling:
         vmfbs["scheduler"] = None
-    sdxl_pipe.load_pipeline(vmfbs, weights, args.rt_device, args.compiled_pipeline, args.split_scheduler)
+    sdxl_pipe.load_pipeline(
+        vmfbs, weights, args.rt_device, args.compiled_pipeline, args.split_scheduler
+    )
     sdxl_pipe.generate_images(
         args.prompt,
         args.negative_prompt,

models/turbine_models/custom_models/sdxl_inference/sdxl_prompt_encoder_runner.py

@@ -5,58 +5,18 @@
 import numpy as np
 
 
-def run_torch_clip(hf_model_name, hf_auth_token, prompt, max_length=64):
-    # TODO: Integrate with HFTransformerBuilder
-    from turbine_models.custom_models.sdxl_inference.clip import ClipModel
-
-    model_1 = ClipModel(hf_model_name, hf_auth_token, index=1)
-    model_2 = ClipModel(hf_model_name, hf_auth_token, index=2)
-    tokenizer_1 = CLIPTokenizer.from_pretrained(
-        hf_model_name,
-        subfolder="tokenizer",
-        token=hf_auth_token,
-    )
-    tokenizer_2 = CLIPTokenizer.from_pretrained(
-        hf_model_name,
-        subfolder="tokenizer_2",
-        token=hf_auth_token,
-    )
-    text_input_1 = tokenizer_1(
-        prompt,
-        padding="max_length",
-        max_length=max_length,
-        truncation=True,
-        return_tensors="pt",
-    )
-    text_input_2 = tokenizer_2(
-        prompt,
-        padding="max_length",
-        max_length=max_length,
-        truncation=True,
-        return_tensors="pt",
-    )
-    example_input_1 = text_input_1.input_ids
-    example_input_2 = text_input_2.input_ids
-
-    results_1 = model_1.forward(example_input_1)
-    results_2 = model_2.forward(example_input_2)
-    np_torch_output_1 = results_1[0].detach().cpu().numpy().astype(np.float16)
-    np_torch_output_2 = results_2[0].detach().cpu().numpy().astype(np.float16)
-    return np_torch_output_1, np_torch_output_2
-
-
 def run_prompt_encoder(
-    args,
+    vmfb_path,
+    device,
+    external_weight_path,
     input_ids,
     uncond_input_ids,
 ):
-    prompt_encoder_runner = vmfbRunner(
-        args.device, args.vmfb_path, args.external_weight_path
-    )
-    np.save("input0.npy", input_ids[0].numpy())
-    np.save("input1.npy", input_ids[1].numpy())
-    np.save("input2.npy", uncond_input_ids[0].numpy())
-    np.save("input3.npy", uncond_input_ids[1].numpy())
+    prompt_encoder_runner = vmfbRunner(device, vmfb_path, external_weight_path)
+    # np.save("input0.npy", input_ids[0].numpy())
+    # np.save("input1.npy", input_ids[1].numpy())
+    # np.save("input2.npy", uncond_input_ids[0].numpy())
+    # np.save("input3.npy", uncond_input_ids[1].numpy())
     prompt_encoder_inputs = [
         ireert.asdevicearray(prompt_encoder_runner.config.device, input_ids[0]),
         ireert.asdevicearray(prompt_encoder_runner.config.device, input_ids[1]),
@@ -66,40 +26,36 @@ def run_prompt_encoder(
     encoded_outputs = prompt_encoder_runner.ctx.modules.compiled_clip["encode_prompts"](
         *prompt_encoder_inputs
     )
+    for i in encoded_outputs:
+        i = i.to_host()
     del prompt_encoder_inputs
     return encoded_outputs
 
 
-if __name__ == "__main__":
-    from turbine_models.custom_models.sdxl_inference.sdxl_cmd_opts import args
-
-    tokenizer_1 = CLIPTokenizer.from_pretrained(
-        args.hf_model_name,
-        subfolder="tokenizer",
-        token=args.hf_auth_token,
-    )
-    tokenizer_2 = CLIPTokenizer.from_pretrained(
-        args.hf_model_name,
-        subfolder="tokenizer_2",
-        token=args.hf_auth_token,
-    )
+def run_tokenize(
+    tokenizer_1,
+    tokenizer_2,
+    prompt,
+    negative_prompt,
+    max_length=64,
+):
     text_input_ids_list = []
     uncond_input_ids_list = []
 
     # Tokenize prompt and negative prompt.
     tokenizers = [tokenizer_1, tokenizer_2]
     for tokenizer in tokenizers:
         text_inputs = tokenizer(
-            args.prompt,
+            prompt,
             padding="max_length",
-            max_length=args.max_length,
+            max_length=max_length,
             truncation=True,
             return_tensors="pt",
         )
         uncond_input = tokenizer(
-            args.negative_prompt,
+            negative_prompt,
             padding="max_length",
-            max_length=args.max_length,
+            max_length=max_length,
             truncation=True,
             return_tensors="pt",
         )
@@ -108,9 +64,34 @@ def run_prompt_encoder(
 
         text_input_ids_list.extend([text_input_ids])
         uncond_input_ids_list.extend([uncond_input_ids])
+    return text_input_ids_list, uncond_input_ids_list
+
 
+if __name__ == "__main__":
+    from turbine_models.custom_models.sdxl_inference.sdxl_cmd_opts import args
+
+    tokenizer_1 = CLIPTokenizer.from_pretrained(
+        args.hf_model_name,
+        subfolder="tokenizer",
+        token=args.hf_auth_token,
+    )
+    tokenizer_2 = CLIPTokenizer.from_pretrained(
+        args.hf_model_name,
+        subfolder="tokenizer_2",
+        token=args.hf_auth_token,
+    )
+
+    text_input_ids_list, uncond_input_ids_list = run_tokenize(
+        tokenizer_1,
+        tokenizer_2,
+        args.prompt,
+        args.negative_prompt,
+        args.max_length,
+    )
     turbine_output1, turbine_output2 = run_prompt_encoder(
-        args,
+        args.vmfb_path,
+        args.rt_device,
+        args.external_weight_path,
         text_input_ids_list,
         uncond_input_ids_list,
     )

models/turbine_models/custom_models/sdxl_inference/unet.py

@@ -94,9 +94,7 @@ def export_unet_model(
     weights_only=False,
 ):
     if pipeline_dir:
-        safe_name = os.path.join(
-            pipeline_dir, f"unet"
-        )
+        safe_name = os.path.join(pipeline_dir, f"unet")
     else:
         safe_name = utils.create_safe_name(
             hf_model_name,