test: pixart compile mode on neuron

Author: JingyaHuang

tests/pipelines/pixart_alpha/test_pixart.py

@@ -381,3 +381,47 @@ def test_pixart_512_without_resolution_binning(self):
         no_res_bin_image_slice = no_res_bin_image[0, -3:, -3:, -1]
 
         assert not np.allclose(image_slice, no_res_bin_image_slice, atol=1e-4, rtol=1e-4)
+
+    @unittest.skipUnless(is_torch_neuronx_available(), "torch_neuronx not available")
+    def test_pixart_512_neuron_compile(self):
+        """
+        Smoke-test PixArtAlphaPipeline under torch.compile(backend="neuron") at 512×512.
+        """
+        import torch_neuronx  # noqa: F401 — registers torch.neuron
+        from torch_neuronx.neuron_dynamo_backend import set_model_name
+
+        device = torch.neuron.current_device()
+        generator = torch.Generator("cpu").manual_seed(0)
+
+        pipe = PixArtAlphaPipeline.from_pretrained(self.ckpt_id_512, torch_dtype=torch.bfloat16)
+        pipe = pipe.to(device)
+        # Flush pending lazy-XLA parameter-copy ops before compiling.
+        torch.neuron.synchronize()
+
+        pipe.transformer.eval()
+        pipe.vae.eval()
+        pipe.text_encoder.eval()
+
+        set_model_name("pixart_text_encoder")
+        pipe.text_encoder = torch.compile(pipe.text_encoder, backend="neuron", fullgraph=True)
+        set_model_name("pixart_transformer")
+        pipe.transformer = torch.compile(pipe.transformer, backend="neuron", fullgraph=True)
+        # VAE must be compiled after pipeline __init__ (which reads vae.config.block_out_channels).
+        set_model_name("pixart_vae")
+        pipe.vae = torch.compile(pipe.vae, backend="neuron", fullgraph=True)
+
+        image = pipe(
+            self.prompt,
+            generator=generator,
+            height=512,
+            width=512,
+            num_inference_steps=2,
+            output_type="np",
+        ).images
+
+        self.assertEqual(image.shape, (1, 512, 512, 3))
+        self.assertFalse(np.isnan(image).any(), "Output contains NaN values")
+        self.assertTrue(
+            (image >= 0.0).all() and (image <= 1.0).all(),
+            "Output pixel values outside [0, 1]",
+        )