add tests

optimum/exporters/neuron/convert.py

@@ -273,7 +273,7 @@ def validate_model_outputs(
                 logger.info(f"\t\t-[✓] {output.shape} matches {ref_output.shape}")
 
             # Values
-            if not torch.allclose(ref_output, output, atol=atol):
+            if not torch.allclose(ref_output, output.to(ref_output.dtype), atol=atol):
                 max_diff = torch.max(torch.abs(ref_output - output))
                 logger.error(f"\t\t-[x] values not close enough, max diff: {max_diff} (atol: {atol})")
                 value_failures.append((name, max_diff))

optimum/neuron/modeling_diffusion.py

@@ -1120,6 +1120,9 @@ def __call__(self, *args, **kwargs):
         kwargs.pop("width", None)
         if kwargs.get("image", None):
             kwargs["image"] = self.image_processor.preprocess(kwargs["image"], height=height, width=width)
+        # Override default `max_sequence_length`, eg. pixart
+        if "max_sequence_length" in inspect.signature(self.auto_model_class.__call__).parameters:
+            kwargs["max_sequence_length"] = self.text_encoder.config.neuron.get("static_sequence_length", None)
         return self.auto_model_class.__call__(self, height=height, width=width, *args, **kwargs)
 
 

tests/inference/test_stable_diffusion_pipeline.py

@@ -19,13 +19,15 @@
 import cv2
 import numpy as np
 import PIL
+import torch
 from compel import Compel, ReturnedEmbeddingsType
 from diffusers import UniPCMultistepScheduler
 from diffusers.utils import load_image
 from parameterized import parameterized
 
 from optimum.neuron import (
     NeuronLatentConsistencyModelPipeline,
+    NeuronPixArtAlphaPipeline,
     NeuronStableDiffusionControlNetPipeline,
     NeuronStableDiffusionImg2ImgPipeline,
     NeuronStableDiffusionInpaintPipeline,
@@ -38,6 +40,7 @@
 from optimum.neuron.modeling_diffusion import (
     NeuronControlNetModel,
     NeuronModelTextEncoder,
+    NeuronModelTransformer,
     NeuronModelUnet,
     NeuronModelVaeDecoder,
     NeuronModelVaeEncoder,
@@ -435,3 +438,40 @@ def test_from_pipe(self, model_arch):
         prompt = "a dog running, lake, moat"
         image = img2img_pipeline(prompt=prompt, image=init_image).images[0]
         self.assertIsInstance(image, PIL.Image.Image)
+
+
+is_inferentia_test
+
+
+@requires_neuronx
+@require_diffusers
+class NeuronPixArtAlphaPipelineIntegrationTest(unittest.TestCase):
+    ATOL_FOR_VALIDATION = 1e-3
+
+    def test_export_and_inference_non_dyn(self):
+        model_id = "hf-internal-testing/tiny-pixart-alpha-pipe"
+        compiler_args = {"auto_cast": "none"}
+        input_shapes = {"batch_size": 1, "height": 64, "width": 64, "sequence_length": 32}
+        neuron_pipeline = NeuronPixArtAlphaPipeline.from_pretrained(
+            model_id,
+            export=True,
+            torch_dtype=torch.bfloat16,
+            dynamic_batch_size=False,
+            disable_neuron_cache=True,
+            **input_shapes,
+            **compiler_args,
+        )
+        self.assertIsInstance(neuron_pipeline.text_encoder, NeuronModelTextEncoder)
+        self.assertIsInstance(neuron_pipeline.transformer, NeuronModelTransformer)
+        self.assertIsInstance(neuron_pipeline.vae_encoder, NeuronModelVaeEncoder)
+        self.assertIsInstance(neuron_pipeline.vae_decoder, NeuronModelVaeDecoder)
+
+        prompt = "Mario eating hamburgers."
+
+        neuron_pipeline.transformer.config.sample_size = (
+            32  # Skip the sample size check because the dummy model uses a smaller sample size (8).
+        )
+        image = neuron_pipeline(prompt=prompt, use_resolution_binning=False).images[
+            0
+        ]  # Set `use_resolution_binning=False` to prevent resizing.
+        self.assertIsInstance(image, PIL.Image.Image)