[Tests] fix some fast gpu tests. (huggingface#9379)

fix some fast gpu tests.

Author: sayakpaul

examples/dreambooth/train_dreambooth_lora_flux.py

@@ -1597,6 +1597,7 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
                             tokenizers=[None, None],
                             text_input_ids_list=[tokens_one, tokens_two],
                             max_sequence_length=args.max_sequence_length,
+                            device=accelerator.device,
                             prompt=prompts,
                         )
                 else:
@@ -1606,6 +1607,7 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
                             tokenizers=[None, None],
                             text_input_ids_list=[tokens_one, tokens_two],
                             max_sequence_length=args.max_sequence_length,
+                            device=accelerator.device,
                             prompt=args.instance_prompt,
                         )
 

src/diffusers/models/transformers/transformer_flux.py

@@ -465,6 +465,7 @@ def forward(
                 "Please remove the batch dimension and pass it as a 2d torch Tensor"
             )
             img_ids = img_ids[0]
+
         ids = torch.cat((txt_ids, img_ids), dim=0)
         image_rotary_emb = self.pos_embed(ids)
 

tests/pipelines/flux/test_pipeline_flux_img2img.py

@@ -18,11 +18,11 @@
 enable_full_determinism()
 
 
-@unittest.skipIf(torch_device == "mps", "Flux has a float64 operation which is not supported in MPS.")
 class FluxImg2ImgPipelineFastTests(unittest.TestCase, PipelineTesterMixin):
     pipeline_class = FluxImg2ImgPipeline
     params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"])
     batch_params = frozenset(["prompt"])
+    test_xformers_attention = False
 
     def get_dummy_components(self):
         torch.manual_seed(0)

tests/pipelines/flux/test_pipeline_flux_inpaint.py

@@ -18,11 +18,11 @@
 enable_full_determinism()
 
 
-@unittest.skipIf(torch_device == "mps", "Flux has a float64 operation which is not supported in MPS.")
 class FluxInpaintPipelineFastTests(unittest.TestCase, PipelineTesterMixin):
     pipeline_class = FluxInpaintPipeline
     params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"])
     batch_params = frozenset(["prompt"])
+    test_xformers_attention = False
 
     def get_dummy_components(self):
         torch.manual_seed(0)