Fix compiled pipeline

Author: gpetters-amd

models/turbine_models/custom_models/pipeline_base.py

@@ -437,6 +437,7 @@ def prepare_all(
         vmfbs: dict = {},
         weights: dict = {},
         interactive: bool = False,
+        num_steps: int = 20,
     ):
         ready = self.is_prepared(vmfbs, weights)
         match ready:
@@ -463,7 +464,9 @@ def prepare_all(
                     if not self.map[submodel].get("weights") and self.map[submodel][
                         "export_args"
                     ].get("external_weights"):
-                        self.export_submodel(submodel, weights_only=True)
+                        self.export_submodel(
+                            submodel, weights_only=True, num_steps=num_steps
+                        )
                 return self.prepare_all(mlirs, vmfbs, weights, interactive)
 
     def is_prepared(self, vmfbs, weights):
@@ -581,6 +584,7 @@ def export_submodel(
         submodel: str,
         input_mlir: str = None,
         weights_only: bool = False,
+        num_steps: int = 20,
     ):
         if not os.path.exists(self.pipeline_dir):
             os.makedirs(self.pipeline_dir)
@@ -670,7 +674,9 @@ def export_submodel(
                     self.map[submodel]["export_args"]["precision"],
                     self.map[submodel]["export_args"]["batch_size"],
                     self.map[submodel]["export_args"]["max_length"],
-                    "tokens_to_image",
+                    "produce_img_split",
+                    unet_module_name=self.map["unet"]["module_name"],
+                    num_steps=num_steps,
                 )
                 dims = [
                     self.map[submodel]["export_args"]["width"],
@@ -699,8 +705,8 @@ def export_submodel(
                     return_path=True,
                     mlir_source="str",
                 )
-                self.map[submodel]["vmfb"] = vmfb_path
-                self.map[submodel]["weights"] = None
+                self.map[submodel]["vmfb"] = [vmfb_path]
+                self.map[submodel]["weights"] = []
             case _:
                 export_args = self.map[submodel].get("export_args", {})
                 if weights_only:
@@ -721,10 +727,24 @@ def export_submodel(
 
     # LOAD
     def load_map(self):
-        for submodel in self.map.keys():
+        # Make sure fullpipeline is imported last
+        submodels = list(self.map.keys() - {"fullpipeline"})
+        submodels += ["fullpipeline"] if "fullpipeline" in self.map.keys() else []
+        for submodel in submodels:
             if not self.map[submodel]["load"]:
-                self.printer.print("Skipping load for ", submodel)
+                self.printer.print(f"Skipping load for {submodel}")
                 continue
+            elif self.map[submodel].get("wraps"):
+                vmfbs = []
+                weights = []
+                for wrapped in self.map[submodel]["wraps"]:
+                    vmfbs.append(self.map[wrapped]["vmfb"])
+                    if "weights" in self.map[wrapped]:
+                        weights.append(self.map[wrapped]["weights"])
+                self.map[submodel]["vmfb"] = vmfbs + self.map[submodel]["vmfb"]
+                self.map[submodel]["weights"] = weights + self.map[submodel]["weights"]
+
+            print(f"Loading {submodel}")
             self.load_submodel(submodel)
 
     def load_submodel(self, submodel):
@@ -751,6 +771,10 @@ def load_submodel(self, submodel):
 
     def unload_submodel(self, submodel):
         self.map[submodel]["runner"].unload()
+        self.map[submodel]["vmfb"] = None
+        self.map[submodel]["mlir"] = None
+        self.map[submodel]["weights"] = None
+        self.map[submodel]["export_args"]["input_mlir"] = None
         setattr(self, submodel, None)
 
 

models/turbine_models/custom_models/sd_inference/sd_pipeline.py

@@ -118,23 +118,11 @@
             "decomp_attn": None,
         },
     },
-    "unetloop": {
-        "module_name": "sdxl_compiled_pipeline",
-        "load": False,
-        "keywords": ["unetloop"],
-        "wraps": ["unet", "scheduler"],
-        "export_args": {
-            "batch_size": 1,
-            "height": 1024,
-            "width": 1024,
-            "max_length": 64,
-        },
-    },
     "fullpipeline": {
         "module_name": "sdxl_compiled_pipeline",
-        "load": False,
+        "load": True,
         "keywords": ["fullpipeline"],
-        "wraps": ["text_encoder", "unet", "scheduler", "vae"],
+        "wraps": ["unet", "scheduler", "vae"],
         "export_args": {
             "batch_size": 1,
             "height": 1024,
@@ -190,6 +178,7 @@ def get_sd_model_map(hf_model_name):
         "stabilityai/sdxl-turbo",
         "stabilityai/stable-diffusion-xl-base-1.0",
         "/models/SDXL/official_pytorch/fp16/stable_diffusion_fp16/checkpoint_pipe",
+        "/models/SDXL/official_pytorch/fp16/stable_diffusion_fp16//checkpoint_pipe",
     ]:
         return sdxl_model_map
     elif "stabilityai/stable-diffusion-3" in name:
@@ -233,6 +222,7 @@ def __init__(
         benchmark: bool | dict[bool] = False,
         verbose: bool = False,
         batch_prompts: bool = False,
+        compiled_pipeline: bool = False,
     ):
         common_export_args = {
             "hf_model_name": None,
@@ -243,11 +233,11 @@ def __init__(
             "exit_on_vmfb": False,
             "pipeline_dir": pipeline_dir,
             "input_mlir": None,
-            "attn_spec": None,
+            "attn_spec": attn_spec,
             "external_weights": None,
             "external_weight_path": None,
         }
-        sd_model_map = get_sd_model_map(hf_model_name)
+        sd_model_map = copy.deepcopy(get_sd_model_map(hf_model_name))
         for submodel in sd_model_map:
             if "load" not in sd_model_map[submodel]:
                 sd_model_map[submodel]["load"] = True
@@ -311,6 +301,7 @@ def __init__(
         self.scheduler = None
 
         self.split_scheduler = True
+        self.compiled_pipeline = compiled_pipeline
 
         self.base_model_name = (
             hf_model_name
@@ -321,11 +312,6 @@ def __init__(
         self.is_sdxl = "xl" in self.base_model_name.lower()
         self.is_sd3 = "stable-diffusion-3" in self.base_model_name
         if self.is_sdxl:
-            if self.split_scheduler:
-                if self.map.get("unetloop"):
-                    self.map.pop("unetloop")
-                if self.map.get("fullpipeline"):
-                    self.map.pop("fullpipeline")
             self.tokenizers = [
                 CLIPTokenizer.from_pretrained(
                     self.base_model_name, subfolder="tokenizer"
@@ -339,6 +325,20 @@ def __init__(
             self.scheduler_device = self.map["unet"]["device"]
             self.scheduler_driver = self.map["unet"]["driver"]
             self.scheduler_target = self.map["unet"]["target"]
+            if not self.compiled_pipeline:
+                if self.map.get("unetloop"):
+                    self.map.pop("unetloop")
+                if self.map.get("fullpipeline"):
+                    self.map.pop("fullpipeline")
+            elif self.compiled_pipeline:
+                self.map["unet"]["load"] = False
+                self.map["vae"]["load"] = False
+                self.load_scheduler(
+                    scheduler_id,
+                    num_inference_steps,
+                )
+                self.map["scheduler"]["runner"].unload()
+                self.map["scheduler"]["load"] = False
         elif not self.is_sd3:
             self.tokenizer = CLIPTokenizer.from_pretrained(
                 self.base_model_name, subfolder="tokenizer"
@@ -351,23 +351,27 @@ def __init__(
 
         self.latents_dtype = torch_dtypes[self.latents_precision]
         self.use_i8_punet = self.use_punet = use_i8_punet
+        if self.use_punet:
+            self.setup_punet()
+        else:
+            self.map["unet"]["keywords"].append("!punet")
+            self.map["unet"]["function_name"] = "run_forward"
+
+    def setup_punet(self):
         if self.use_i8_punet:
             self.map["unet"]["export_args"]["precision"] = "i8"
-            self.map["unet"]["export_args"]["use_punet"] = True
-            self.map["unet"]["use_weights_for_export"] = True
-            self.map["unet"]["keywords"].append("punet")
-            self.map["unet"]["module_name"] = "compiled_punet"
-            self.map["unet"]["function_name"] = "main"
             self.map["unet"]["export_args"]["external_weight_path"] = (
                 utils.create_safe_name(self.base_model_name) + "_punet_dataset_i8.irpa"
             )
             for idx, word in enumerate(self.map["unet"]["keywords"]):
                 if word in ["fp32", "fp16"]:
                     self.map["unet"]["keywords"][idx] = "i8"
                     break
-        else:
-            self.map["unet"]["keywords"].append("!punet")
-            self.map["unet"]["function_name"] = "run_forward"
+        self.map["unet"]["export_args"]["use_punet"] = True
+        self.map["unet"]["use_weights_for_export"] = True
+        self.map["unet"]["keywords"].append("punet")
+        self.map["unet"]["module_name"] = "compiled_punet"
+        self.map["unet"]["function_name"] = "main"
 
     # LOAD
 
@@ -376,10 +380,6 @@ def load_scheduler(
         scheduler_id: str,
         steps: int = 30,
     ):
-        if self.is_sd3:
-            scheduler_device = self.mmdit.device
-        else:
-            scheduler_device = self.unet.device
         if not self.cpu_scheduling:
             self.map["scheduler"] = {
                 "module_name": "compiled_scheduler",
@@ -425,7 +425,11 @@ def load_scheduler(
             except:
                 print("JIT export of scheduler failed. Loading CPU scheduler.")
                 self.cpu_scheduling = True
-        if self.cpu_scheduling:
+        elif self.cpu_scheduling:
+            if self.is_sd3:
+                scheduler_device = self.mmdit.device
+            else:
+                scheduler_device = self.unet.device
             scheduler = schedulers.get_scheduler(self.base_model_name, scheduler_id)
             self.scheduler = schedulers.SharkSchedulerCPUWrapper(
                 scheduler,
@@ -461,13 +465,10 @@ def encode_prompts_sdxl(self, prompt, negative_prompt):
             text_input_ids_list += text_inputs.input_ids.unsqueeze(0)
             uncond_input_ids_list += uncond_input.input_ids.unsqueeze(0)
 
-        if self.compiled_pipeline:
-            return text_input_ids_list, uncond_input_ids_list
-        else:
-            prompt_embeds, add_text_embeds = self.text_encoder(
-                "encode_prompts", [*text_input_ids_list, *uncond_input_ids_list]
-            )
-            return prompt_embeds, add_text_embeds
+        prompt_embeds, add_text_embeds = self.text_encoder(
+            "encode_prompts", [*text_input_ids_list, *uncond_input_ids_list]
+        )
+        return prompt_embeds, add_text_embeds
 
     def prepare_latents(
         self,
@@ -565,9 +566,11 @@ def _produce_latents_sdxl(
             [guidance_scale],
             dtype=self.map["unet"]["np_dtype"],
         )
+        # Disable progress bar if we aren't in verbose mode or if we're printing
+        # benchmark latencies for unet.
         for i, t in tqdm(
             enumerate(timesteps),
-            disable=(self.map["unet"].get("benchmark") and self.verbose),
+            disable=(self.map["unet"].get("benchmark") or not self.verbose),
         ):
             if self.cpu_scheduling:
                 latent_model_input, t = self.scheduler.scale_model_input(
@@ -608,6 +611,75 @@ def _produce_latents_sdxl(
                 latents = self.scheduler("run_step", [noise_pred, t, latents])
         return latents
 
+    def produce_images_compiled(
+        self,
+        sample,
+        prompt_embeds,
+        text_embeds,
+        guidance_scale,
+    ):
+        pipe_inputs = [
+            sample,
+            prompt_embeds,
+            text_embeds,
+            torch.as_tensor([guidance_scale], dtype=sample.dtype),
+        ]
+        # image = self.compiled_pipeline("produce_img_latents", pipe_inputs)
+        image = self.map["fullpipeline"]["runner"]("produce_image_latents", pipe_inputs)
+        return image
+
+    def prepare_sampling_inputs(
+        self,
+        prompt: str,
+        negative_prompt: str = "",
+        steps: int = 30,
+        batch_count: int = 1,
+        guidance_scale: float = 7.5,
+        seed: float = -1,
+        cpu_scheduling: bool = True,
+        scheduler_id: str = "EulerDiscrete",
+        return_imgs: bool = False,
+    ):
+        needs_new_scheduler = (
+            (steps and steps != self.num_inference_steps)
+            or (cpu_scheduling != self.cpu_scheduling)
+            and self.split_scheduler
+        )
+        if not self.scheduler and not self.compiled_pipeline:
+            needs_new_scheduler = True
+
+        if guidance_scale == 0:
+            negative_prompt = prompt
+            prompt = ""
+
+        self.cpu_scheduling = cpu_scheduling
+        if steps and needs_new_scheduler:
+            self.num_inference_steps = steps
+            self.load_scheduler(scheduler_id, steps)
+
+        pipe_start = time.time()
+        numpy_images = []
+
+        samples = self.get_rand_latents(seed, batch_count)
+
+        # Tokenize prompt and negative prompt.
+        if self.is_sdxl:
+            prompt_embeds, negative_embeds = self.encode_prompts_sdxl(
+                prompt, negative_prompt
+            )
+        else:
+            prompt_embeds, negative_embeds = encode_prompt(
+                self, prompt, negative_prompt
+            )
+        produce_latents_input = [
+            samples[0],
+            prompt_embeds,
+            negative_embeds,
+            steps,
+            guidance_scale,
+        ]
+        return produce_latents_input
+
     def generate_images(
         self,
         prompt: str,
@@ -653,18 +725,23 @@ def generate_images(
             )
 
         for i in range(batch_count):
-            produce_latents_input = [
-                samples[i],
-                prompt_embeds,
-                negative_embeds,
-                steps,
-                guidance_scale,
-            ]
-            if self.is_sdxl:
-                latents = self._produce_latents_sdxl(*produce_latents_input)
+            if self.compiled_pipeline:
+                image = self.produce_images_compiled(
+                    samples[i], prompt_embeds, negative_embeds, guidance_scale
+                ).to_host()
             else:
-                latents = self._produce_latents_sd(*produce_latents_input)
-            image = self.vae("decode", [latents])
+                produce_latents_input = [
+                    samples[i],
+                    prompt_embeds,
+                    negative_embeds,
+                    steps,
+                    guidance_scale,
+                ]
+                if self.is_sdxl:
+                    latents = self._produce_latents_sdxl(*produce_latents_input)
+                else:
+                    latents = self._produce_latents_sd(*produce_latents_input)
+                image = self.vae("decode", [latents])
             numpy_images.append(image)
             pipe_end = time.time()
 
@@ -750,8 +827,10 @@ def numpy_to_pil_image(images):
         args.use_i8_punet,
         benchmark,
         args.verbose,
+        False,
+        args.compiled_pipeline,
     )
-    sd_pipe.prepare_all()
+    sd_pipe.prepare_all(num_steps=args.num_inference_steps)
     sd_pipe.load_map()
     sd_pipe.generate_images(
         args.prompt,