Merge branch 'main' into sdxl-inpaint-multicontrolnet-batching

Author: abhay23-AI

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/modular_pipelines/ideogram4/modular_pipeline.py

@@ -12,11 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
+from ...loaders import Ideogram4LoraLoaderMixin
 from ..modular_pipeline import ModularPipeline
 
 
-class Ideogram4ModularPipeline(ModularPipeline):
+class Ideogram4ModularPipeline(ModularPipeline, Ideogram4LoraLoaderMixin):
     """
     A ModularPipeline for Ideogram4.
 

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]",
+        )