Some other enhancement for stable diffusion

optimum/exporters/neuron/__main__.py

@@ -108,13 +108,17 @@ def normalize_input_shapes(task: str, args: argparse.Namespace) -> Dict[str, int
 
 
 def normalize_stable_diffusion_input_shapes(
-    task: str,
     args: argparse.Namespace,
 ) -> Dict[str, Dict[str, int]]:
     args = vars(args) if isinstance(args, argparse.Namespace) else args
     mandatory_axes = set(getattr(inspect.getfullargspec(build_stable_diffusion_components_mandatory_shapes), "args"))
     # Remove `sequence_length` as diffusers will pad it to the max and remove number of channels .
-    mandatory_axes = mandatory_axes - {"sequence_length", "unet_num_channels", "vae_num_channels"}
+    mandatory_axes = mandatory_axes - {
+        "sequence_length",
+        "unet_num_channels",
+        "vae_encoder_num_channels",
+        "vae_decoder_num_channels",
+    }
     if not mandatory_axes.issubset(set(args.keys())):
         raise AttributeError(
             f"Shape of {mandatory_axes} are mandatory for neuron compilation, while {mandatory_axes.difference(args.keys())} are not given."
@@ -274,7 +278,7 @@ def main():
     compiler_kwargs = infer_compiler_kwargs(args)
 
     if task == "stable-diffusion":
-        input_shapes = normalize_stable_diffusion_input_shapes(task, args)
+        input_shapes = normalize_stable_diffusion_input_shapes(args)
     else:
         input_shapes = normalize_input_shapes(task, args)
 

optimum/exporters/neuron/convert.py

@@ -21,6 +21,7 @@
 
 import numpy as np
 import torch
+from packaging import version
 from transformers import PretrainedConfig
 
 from ...exporters.error_utils import OutputMatchError, ShapeError
@@ -434,7 +435,7 @@ def export_neuronx(
         neuron_model = neuronx.dynamic_batch(neuron_model)
 
     # diffusers specific
-    async_stable_diffusion_loading(config, neuron_model)
+    improve_stable_diffusion_loading(config, neuron_model)
 
     torch.jit.save(neuron_model, output)
     del neuron_model
@@ -453,14 +454,15 @@ def add_stable_diffusion_compiler_args(config, compiler_args):
     return compiler_args
 
 
-def async_stable_diffusion_loading(config, neuron_model):
-    if hasattr(config._config, "_name_or_path"):
-        sd_components = ["text_encoder", "unet", "vae", "vae_encoder", "vae_decoder"]
-        if any(component in config._config._name_or_path.lower() for component in sd_components):
-            neuronx.async_load(neuron_model)
-        # unet
-        if "unet" in config._config._name_or_path.lower():
-            neuronx.lazy_load(neuron_model)
+def improve_stable_diffusion_loading(config, neuron_model):
+    if version.parse(neuronx.__version__) >= version.parse("1.13.1.1.9.0"):
+        if hasattr(config._config, "_name_or_path"):
+            sd_components = ["text_encoder", "unet", "vae", "vae_encoder", "vae_decoder"]
+            if any(component in config._config._name_or_path.lower() for component in sd_components):
+                neuronx.async_load(neuron_model)
+            # unet
+            if "unet" in config._config._name_or_path.lower():
+                neuronx.lazy_load(neuron_model)
 
 
 def export_neuron(

optimum/exporters/neuron/utils.py

@@ -85,14 +85,21 @@ def build_stable_diffusion_components_mandatory_shapes(
     batch_size: Optional[int] = None,
     sequence_length: Optional[int] = None,
     unet_num_channels: Optional[int] = None,
-    vae_num_channels: Optional[int] = None,
+    vae_encoder_num_channels: Optional[int] = None,
+    vae_decoder_num_channels: Optional[int] = None,
     height: Optional[int] = None,
     width: Optional[int] = None,
 ):
     text_encoder_input_shapes = {"batch_size": batch_size, "sequence_length": sequence_length}
-    vae_encoder_input_shapes = vae_decoder_input_shapes = {
+    vae_encoder_input_shapes = {
         "batch_size": batch_size,
-        "num_channels": vae_num_channels,
+        "num_channels": vae_encoder_num_channels,
+        "height": height,
+        "width": width,
+    }
+    vae_decoder_input_shapes = {
+        "batch_size": batch_size,
+        "num_channels": vae_decoder_num_channels,
         "height": height,
         "width": width,
     }

optimum/neuron/modeling_diffusion.py

@@ -78,6 +78,7 @@ def __init__(
         tokenizer: CLIPTokenizer,
         scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
         feature_extractor: Optional[CLIPFeatureExtractor] = None,
+        device_ids: Optional[List[int]] = [],
         configs: Optional[Dict[str, "PretrainedConfig"]] = None,
         neuron_configs: Optional[Dict[str, "NeuronConfig"]] = None,
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
@@ -103,6 +104,8 @@ def __init__(
                 A scheduler to be used in combination with the U-NET component to denoise the encoded image latents.
             feature_extractor (`Optional[CLIPFeatureExtractor]`, defaults to `None`):
                 A model extracting features from generated images to be used as inputs for the `safety_checker`
+            device_ids (Optional[List[int]], defaults to `[]`):
+                A list of integers that specify the NeuronCores to use for parallelization
             configs (Optional[Dict[str, "PretrainedConfig"]], defaults to `None`):
                 A dictionary configurations for components of the pipeline.
             neuron_configs (Optional["NeuronConfig"], defaults to `None`):
@@ -114,6 +117,7 @@ def __init__(
         """
 
         self._internal_dict = config
+        self.device_ids = device_ids
         self.configs = configs
         self.neuron_configs = neuron_configs
         self.dynamic_batch_size = all(
@@ -332,6 +336,7 @@ def _from_pretrained(
             tokenizer=sub_models["tokenizer"],
             scheduler=sub_models["scheduler"],
             feature_extractor=sub_models.pop("feature_extractor", None),
+            device_ids=device_ids,
             configs=configs,
             neuron_configs=neuron_configs,
             model_save_dir=model_save_dir,

optimum/neuron/pipelines/diffusers/pipeline_stable_diffusion.py

@@ -196,7 +196,7 @@ def __call__(
         # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
         # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
         # corresponds to doing no classifier free guidance.
-        do_classifier_free_guidance = guidance_scale > 1.0 and self.dynamic_batch_size
+        do_classifier_free_guidance = guidance_scale > 1.0 and (self.dynamic_batch_size or len(self.device_ids) == 2)
 
         # 3. Encode input prompt
         text_encoder_lora_scale = (