[Neuron] Enable torch.compile compatibility with Neuron device

src/diffusers/models/transformers/transformer_flux2.py

@@ -292,8 +292,8 @@ def __init__(self):
  self.gate_fn = nn.SiLU()
 
  def forward(self, x: torch.Tensor) -> torch.Tensor:
-x1, x2 = x.chunk(2, dim=-1)
- x = self.gate_fn(x1) * x2
+half = x.shape[-1] // 2
+ x = self.gate_fn(x[..., :half]) * x[..., half:]
  return x
 
 

src/diffusers/pipelines/flux2/pipeline_flux2_klein.py

@@ -24,7 +24,7 @@
 from ...models import AutoencoderKLFlux2, Flux2Transformer2DModel
 from ...schedulers import FlowMatchEulerDiscreteScheduler
 from ...utils import is_torch_xla_available, logging, replace_example_docstring
-from ...utils.torch_utils import randn_tensor
+from ...utils.torch_utils import maybe_adjust_dtype_for_device, randn_tensor
 from ..pipeline_utils import DiffusionPipeline
 from .image_processor import Flux2ImageProcessor
 from .pipeline_output import Flux2PipelineOutput
@@ -405,8 +405,9 @@ def _unpack_latents_with_ids(
  x_list = []
  for data, pos in zip(x, x_ids):
  _, ch = data.shape # noqa: F841
- h_ids = pos[:, 1].to(torch.int64)
- w_ids = pos[:, 2].to(torch.int64)
+ idx_dtype = maybe_adjust_dtype_for_device(torch.int64, data.device)
+ h_ids = pos[:, 1].to(idx_dtype)
+ w_ids = pos[:, 2].to(idx_dtype)
 
  # Use provided height/width to avoid DtoH sync from torch.max().item()
  h = height if height is not None else torch.max(h_ids) + 1
@@ -826,7 +827,8 @@ def __call__(
  # 7. Denoising loop
  # We set the index here to remove DtoH sync, helpful especially during compilation.
  # Check out more details here: https://github.com/huggingface/diffusers/pull/11696
- self.scheduler.set_begin_index(0)
+ if hasattr(self.scheduler, "set_begin_index"):
+ self.scheduler.set_begin_index(0)
  with self.progress_bar(total=num_inference_steps) as progress_bar:
  for i, t in enumerate(timesteps):
  if self.interrupt:

src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py

@@ -861,6 +861,10 @@ def __call__(
  prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
  prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
 
+ prompt_attention_mask = prompt_attention_mask.to(
+ maybe_adjust_dtype_for_device(prompt_attention_mask.dtype, prompt_attention_mask.device)
+ )
+
  # 4. Prepare timesteps
  is_neuron_device = device.type == "neuron"
  if XLA_AVAILABLE or is_neuron_device:
@@ -903,7 +907,8 @@ def __call__(
 
  # 7. Denoising loop
  num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
-
+ if hasattr(self.scheduler, "set_begin_index"):
+ self.scheduler.set_begin_index(0)
  with self.progress_bar(total=num_inference_steps) as progress_bar:
  for i, t in enumerate(timesteps):
  latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents

tests/pipelines/flux2/test_pipeline_flux2_klein.py

@@ -13,10 +13,10 @@
  Flux2KleinPipeline,
  Flux2Transformer2DModel,
 )
-from diffusers.utils.import_utils import is_torch_neuronx_available
 
 from ...testing_utils import (
  backend_empty_cache,
+ backend_synchronize,
  require_torch_neuron,
  torch_device,
 )
@@ -198,13 +198,12 @@ class Flux2KleinPipelineIntegrationTests(unittest.TestCase):
  def setUp(self):
  super().setUp()
  self._saved_env = {}
- if is_torch_neuronx_available():
- neff_cache_dir = "/tmp/neff_cache"
- os.makedirs(neff_cache_dir, exist_ok=True)
- for key in ("TORCH_NEURONX_NEFF_CACHE_DIR", "TORCH_NEURONX_ENABLE_NKI_SDPA"):
- self._saved_env[key] = os.environ.get(key)
- os.environ["TORCH_NEURONX_NEFF_CACHE_DIR"] = neff_cache_dir
- os.environ.setdefault("TORCH_NEURONX_ENABLE_NKI_SDPA", "0")
+ neff_cache_dir = "/tmp/neff_cache"
+ os.makedirs(neff_cache_dir, exist_ok=True)
+ for key in ("TORCH_NEURONX_NEFF_CACHE_DIR", "TORCH_NEURONX_ENABLE_NKI_SDPA"):
+ self._saved_env[key] = os.environ.get(key)
+ os.environ["TORCH_NEURONX_NEFF_CACHE_DIR"] = neff_cache_dir
+ os.environ.setdefault("TORCH_NEURONX_ENABLE_NKI_SDPA", "0")
  gc.collect()
  backend_empty_cache(torch_device)
 
@@ -223,8 +222,7 @@ def test_flux2_klein_inference_512(self):
 
  pipe = Flux2KleinPipeline.from_pretrained(self.ckpt_id, torch_dtype=torch.bfloat16)
  pipe.to(torch_device)
- if is_torch_neuronx_available():
- torch.neuron.synchronize()
+ backend_synchronize(torch_device)
  pipe.set_progress_bar_config(disable=None)
 
  image = pipe(
@@ -242,3 +240,44 @@ def test_flux2_klein_inference_512(self):
  self.assertTrue(np.all((image >= 0.0) & (image <= 1.0)), "Pixel values must be in [0, 1]")
  expected_slice = np.array([0.3652, 0.3574, 0.3633, 0.4102, 0.4062, 0.4043, 0.4453, 0.4355, 0.4570])
  self.assertLess(np.abs(image_slice.flatten() - expected_slice).max(), 5e-2)
+
+ @require_torch_neuron
+ def test_flux2_klein_neuron_compile_128(self):
+ from torch_neuronx.neuron_dynamo_backend import set_model_name
+
+ device = torch.neuron.current_device()
+ generator = torch.Generator("cpu").manual_seed(0)
+
+ pipe = Flux2KleinPipeline.from_pretrained(self.ckpt_id, torch_dtype=torch.bfloat16)
+ pipe = pipe.to(device)
+ backend_synchronize(torch_device)
+
+ pipe.transformer.eval()
+ pipe.vae.eval()
+ pipe.text_encoder.eval()
+
+ # Keep the text encoder eager: it reads intermediate hidden_states, which
+ # transformers only materializes outside of torch.compile(fullgraph=True).
+ # It runs once per generation, so leaving it uncompiled is negligible.
+ set_model_name("flux2_klein_transformer")
+ pipe.transformer = torch.compile(pipe.transformer, backend="neuron", fullgraph=True)
+
+ set_model_name("flux2_klein_vae")
+ pipe.vae = torch.compile(pipe.vae, backend="neuron", fullgraph=True)
+
+ image = pipe(
+ prompt=self.prompt,
+ height=128,
+ width=128,
+ num_inference_steps=4,
+ guidance_scale=1.0,
+ generator=generator,
+ output_type="np",
+ ).images
+
+ self.assertEqual(image.shape, (1, 128, 128, 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]",
+ )

tests/pipelines/pixart_alpha/test_pixart.py

@@ -31,9 +31,11 @@
 
 from ...testing_utils import (
  backend_empty_cache,
+ backend_synchronize,
  enable_full_determinism,
  numpy_cosine_similarity_distance,
  require_torch_accelerator,
+ require_torch_neuron,
  slow,
  torch_device,
 )
@@ -381,3 +383,45 @@ 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)
+
+ @require_torch_neuron
+ def test_pixart_512_neuron_compile(self):
+ """
+ Smoke-test PixArtAlphaPipeline under torch.compile(backend="neuron") at 512×512.
+ """
+ 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)
+ backend_synchronize(torch_device)
+
+ 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]",
+ )

tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py

@@ -36,10 +36,10 @@
  UNet2DConditionModel,
  UniPCMultistepScheduler,
 )
-from diffusers.utils.import_utils import is_torch_neuronx_available
 
 from ...testing_utils import (
  backend_empty_cache,
+ backend_synchronize,
  enable_full_determinism,
  load_image,
  numpy_cosine_similarity_distance,
@@ -987,10 +987,8 @@ class StableDiffusionXLTurboPipelineIntegrationTests(unittest.TestCase):
 
  def setUp(self):
  super().setUp()
- self._saved_env = {}
- if is_torch_neuronx_available():
- self._saved_env["TORCH_NEURONX_ENABLE_NKI_SDPA"] = os.environ.get("TORCH_NEURONX_ENABLE_NKI_SDPA")
- os.environ.setdefault("TORCH_NEURONX_ENABLE_NKI_SDPA", "0")
+ self._saved_env = {"TORCH_NEURONX_ENABLE_NKI_SDPA": os.environ.get("TORCH_NEURONX_ENABLE_NKI_SDPA")}
+ os.environ.setdefault("TORCH_NEURONX_ENABLE_NKI_SDPA", "0")
  gc.collect()
  backend_empty_cache(torch_device)
 
@@ -1009,8 +1007,7 @@ def test_sdxl_turbo_512(self):
 
  pipe = AutoPipelineForText2Image.from_pretrained(self.ckpt_id, torch_dtype=torch.float16, variant="fp16")
  pipe.to(torch_device)
- if is_torch_neuronx_available():
- torch.neuron.synchronize()
+ backend_synchronize(torch_device)
  pipe.set_progress_bar_config(disable=None)
 
  image = pipe(
@@ -1026,3 +1023,59 @@ def test_sdxl_turbo_512(self):
  self.assertTrue(np.all((image >= 0.0) & (image <= 1.0)), "Pixel values must be in [0, 1]")
  expected_slice = np.array([0.3524, 0.3160, 0.3652, 0.3316, 0.3376, 0.3315, 0.3042, 0.3102, 0.3449])
  self.assertLess(np.abs(image_slice.flatten() - expected_slice).max(), 5e-2)
+
+ @require_torch_neuron
+ def test_sdxl_turbo_neuron_compile_256(self):
+ from torch_neuronx.neuron_dynamo_backend import set_model_name
+ from transformers.utils.output_capturing import install_all_output_capturing_hooks
+
+ device = torch.neuron.current_device()
+ generator = torch.Generator("cpu").manual_seed(0)
+
+ pipe = AutoPipelineForText2Image.from_pretrained(self.ckpt_id, torch_dtype=torch.bfloat16, variant="fp16")
+ pipe = pipe.to(device)
+ backend_synchronize(torch_device)
+
+ pipe.unet.eval()
+ pipe.vae.eval()
+ pipe.text_encoder.eval()
+ pipe.text_encoder_2.eval()
+
+ install_all_output_capturing_hooks(pipe.text_encoder)
+ set_model_name("sdxl_turbo_text_encoder")
+ pipe.text_encoder = torch.compile(pipe.text_encoder, backend="neuron", fullgraph=True)
+
+ install_all_output_capturing_hooks(pipe.text_encoder_2)
+ set_model_name("sdxl_turbo_text_encoder_2")
+ pipe.text_encoder_2 = torch.compile(pipe.text_encoder_2, backend="neuron", fullgraph=True)
+
+ set_model_name("sdxl_turbo_unet")
+ pipe.unet = torch.compile(pipe.unet, backend="neuron", fullgraph=True)
+
+ # Pre-warm text encoders and copy ops for 256×256 (latent: 32×32).
+ tok_kwargs = {"padding": "max_length", "max_length": 77, "truncation": True, "return_tensors": "pt"}
+ with torch.no_grad():
+ _ids = pipe.tokenizer("warmup", **tok_kwargs).input_ids.to(device)
+ _ = pipe.text_encoder(_ids, output_hidden_states=True)
+ _ids2 = pipe.tokenizer_2("warmup", **tok_kwargs).input_ids.to(device)
+ _ = pipe.text_encoder_2(_ids2, output_hidden_states=True)
+ for _shape, _dtype in [((1, 4, 32, 32), torch.bfloat16), ((1, 6), torch.bfloat16)]:
+ _ = torch.zeros(_shape, dtype=_dtype).to(device)
+ backend_synchronize(torch_device)
+
+ image = pipe(
+ self.prompt,
+ height=256,
+ width=256,
+ num_inference_steps=1,
+ guidance_scale=0.0,
+ generator=generator,
+ output_type="np",
+ ).images
+
+ self.assertEqual(image.shape, (1, 256, 256, 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]",
+ )