diff --git a/src/diffusers/models/transformers/transformer_flux2.py b/src/diffusers/models/transformers/transformer_flux2.py index c3fa6ac141f3..17c8bd0ffd52 100644 --- a/src/diffusers/models/transformers/transformer_flux2.py +++ b/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 diff --git a/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py b/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py index 9a3468525c0c..d768e6127f26 100644 --- a/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py +++ b/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: diff --git a/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py b/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py index 11eaeaca7fc0..b98248f3ca7b 100644 --- a/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py +++ b/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 diff --git a/tests/pipelines/flux2/test_pipeline_flux2_klein.py b/tests/pipelines/flux2/test_pipeline_flux2_klein.py index 377f02dc9aa1..76528ef466df 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2_klein.py +++ b/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]", + ) diff --git a/tests/pipelines/pixart_alpha/test_pixart.py b/tests/pipelines/pixart_alpha/test_pixart.py index 86fe673a8c7d..399bcdf75609 100644 --- a/tests/pipelines/pixart_alpha/test_pixart.py +++ b/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]", + ) diff --git a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py index c9afdc3209cd..7cf1f7dec3bc 100644 --- a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py +++ b/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]", + )