fix:apr_02 beta

Author: JingyaHuang

src/diffusers/models/transformers/transformer_flux2.py

@@ -961,7 +961,8 @@ def forward(self, ids: torch.Tensor) -> torch.Tensor:
         pos = ids.float()
         is_mps = ids.device.type == "mps"
         is_npu = ids.device.type == "npu"
-        freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
+        is_neuron = ids.device.type == "neuron"
+        freqs_dtype = torch.float32 if (is_mps or is_npu or is_neuron) else torch.float64
         # Unlike Flux 1, loop over len(self.axes_dim) rather than ids.shape[-1]
         for i in range(len(self.axes_dim)):
             cos, sin = get_1d_rotary_pos_embed(

src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py

@@ -29,6 +29,7 @@
     deprecate,
     is_bs4_available,
     is_ftfy_available,
+    is_torch_neuronx_available,
     is_torch_xla_available,
     logging,
     replace_example_docstring,
@@ -862,7 +863,7 @@ def __call__(
             prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
 
         # 4. Prepare timesteps
-        if XLA_AVAILABLE:
+        if XLA_AVAILABLE or is_torch_neuronx_available():
             timestep_device = "cpu"
         else:
             timestep_device = device
@@ -914,10 +915,11 @@ def __call__(
                     # This would be a good case for the `match` statement (Python 3.10+)
                     is_mps = latent_model_input.device.type == "mps"
                     is_npu = latent_model_input.device.type == "npu"
+                    is_neuron = latent_model_input.device.type == "neuron"
                     if isinstance(current_timestep, float):
-                        dtype = torch.float32 if (is_mps or is_npu) else torch.float64
+                        dtype = torch.float32 if (is_mps or is_npu or is_neuron) else torch.float64
                     else:
-                        dtype = torch.int32 if (is_mps or is_npu) else torch.int64
+                        dtype = torch.int32 if (is_mps or is_npu or is_neuron) else torch.int64
                     current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device)
                 elif len(current_timestep.shape) == 0:
                     current_timestep = current_timestep[None].to(latent_model_input.device)