





import shutil
import tempfile
import unittest

from transformers import SPIECE_UNDERLINE, BatchEncoding, MBartTokenizer, MBartTokenizerFast, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
)

from ...test_tokenization_common import TokenizerTesterMixin


_a     :    Tuple         =						get_tests_dir('fixtures/test_sentencepiece.model')


if is_torch_available():
			from transformers.models.mbart.modeling_mbart import shift_tokens_right

_a     :    Dict         =						250_004
_a     :    List[Any]         =						250_020
@require_sentencepiece
@require_tokenizers
class   a_ (			a		,							unittest.TestCase      ):
			A__      :			int       	=  MBartTokenizer
			A__      :			int       	=  MBartTokenizerFast
			A__      :			List[str]       	=  True
			A__      :			List[str]       	=  True

			def      lowerCAmelCase(   self   :  str						):






							"""simple docstring"""
							super().setUp()

							# We have a SentencePiece fixture for testing
							snake_case   :					Union[str, Any]    						=				MBartTokenizer(__UpperCAmelCase     , keep_accents=__UpperCAmelCase						)
							tokenizer.save_pretrained(self.tmpdirname						)

			def      lowerCAmelCase(   self   :  List[str]						):






							"""simple docstring"""
							snake_case   :					Dict    						=				MBartTokenizer(__UpperCAmelCase     , keep_accents=__UpperCAmelCase						)

							snake_case   :					Optional[int]    						=				tokenizer.tokenize('''This is a test'''						)
							self.assertListEqual(__UpperCAmelCase     , ['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est''']						)

							self.assertListEqual(
							    tokenizer.convert_tokens_to_ids(__UpperCAmelCase						)     , [value + tokenizer.fairseq_offset for value in [285, 46, 10, 170, 382]]     , )

							snake_case   :					Optional[int]    						=				tokenizer.tokenize('''I was born in 92000, and this is falsé.'''						)
							self.assertListEqual(
							    __UpperCAmelCase     , [
							        SPIECE_UNDERLINE + '''I''',
							        SPIECE_UNDERLINE + '''was''',
							        SPIECE_UNDERLINE + '''b''',
							        '''or''',
							        '''n''',
							        SPIECE_UNDERLINE + '''in''',
							        SPIECE_UNDERLINE + '''''',
							        '''9''',
							        '''2''',
							        '''0''',
							        '''0''',
							        '''0''',
							        ''',''',
							        SPIECE_UNDERLINE + '''and''',
							        SPIECE_UNDERLINE + '''this''',
							        SPIECE_UNDERLINE + '''is''',
							        SPIECE_UNDERLINE + '''f''',
							        '''al''',
							        '''s''',
							        '''é''',
							        '''.''',
							    ]     , )
							snake_case   :					Any    						=				tokenizer.convert_tokens_to_ids(__UpperCAmelCase						)
							self.assertListEqual(
							    __UpperCAmelCase     , [
							        value + tokenizer.fairseq_offset
							        for value in [8, 21, 84, 55, 24, 19, 7, 2, 602, 347, 347, 347, 3, 12, 66, 46, 72, 80, 6, 2, 4]
							        #                                       ^ unk: 2 + 1 = 3                  unk: 2 + 1 = 3 ^
							    ]     , )

							snake_case   :					List[Any]    						=				tokenizer.convert_ids_to_tokens(__UpperCAmelCase						)
							self.assertListEqual(
							    __UpperCAmelCase     , [
							        SPIECE_UNDERLINE + '''I''',
							        SPIECE_UNDERLINE + '''was''',
							        SPIECE_UNDERLINE + '''b''',
							        '''or''',
							        '''n''',
							        SPIECE_UNDERLINE + '''in''',
							        SPIECE_UNDERLINE + '''''',
							        '''<unk>''',
							        '''2''',
							        '''0''',
							        '''0''',
							        '''0''',
							        ''',''',
							        SPIECE_UNDERLINE + '''and''',
							        SPIECE_UNDERLINE + '''this''',
							        SPIECE_UNDERLINE + '''is''',
							        SPIECE_UNDERLINE + '''f''',
							        '''al''',
							        '''s''',
							        '''<unk>''',
							        '''.''',
							    ]     , )



			def      lowerCAmelCase(   self   :  Tuple						):






							"""simple docstring"""


							if not self.test_slow_tokenizer:
											# as we don't have a slow version, we can't compare the outputs between slow and fast versions
											return

							snake_case   :					List[str]    						=				(self.rust_tokenizer_class, '''hf-internal-testing/tiny-random-mbart''', {})
							for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
											with self.subTest(F"{tokenizer.__class__.__name__} ({pretrained_name})"						):
															snake_case   :					Optional[Any]    						=				self.rust_tokenizer_class.from_pretrained(__UpperCAmelCase     , **__UpperCAmelCase						)
															snake_case   :					Any    						=				self.tokenizer_class.from_pretrained(__UpperCAmelCase     , **__UpperCAmelCase						)

															snake_case   :					Union[str, Any]    						=				tempfile.mkdtemp()

															snake_case   :					Optional[Any]    						=				tokenizer_r.save_pretrained(__UpperCAmelCase						)
															snake_case   :					str    						=				tokenizer_p.save_pretrained(__UpperCAmelCase						)

															# Checks it save with the same files + the tokenizer.json file for the fast one
															self.assertTrue(any('''tokenizer.json''' in f for f in tokenizer_r_files						)						)
															snake_case   :					Dict    						=				tuple(f for f in tokenizer_r_files if '''tokenizer.json''' not in f						)
															self.assertSequenceEqual(__UpperCAmelCase     , __UpperCAmelCase						)

															# Checks everything loads correctly in the same way
															snake_case   :					Union[str, Any]    						=				tokenizer_r.from_pretrained(__UpperCAmelCase						)
															snake_case   :					Any    						=				tokenizer_p.from_pretrained(__UpperCAmelCase						)

															# Check special tokens are set accordingly on Rust and Python
															for key in tokenizer_pp.special_tokens_map:
																			self.assertTrue(hasattr(__UpperCAmelCase     , __UpperCAmelCase						)						)
																			# self.assertEqual(getattr(tokenizer_rp, key), getattr(tokenizer_pp, key))
																			# self.assertEqual(getattr(tokenizer_rp, key + "_id"), getattr(tokenizer_pp, key + "_id"))

															shutil.rmtree(__UpperCAmelCase						)

															# Save tokenizer rust, legacy_format=True
															snake_case   :					Union[str, Any]    						=				tempfile.mkdtemp()

															snake_case   :					List[Any]    						=				tokenizer_r.save_pretrained(__UpperCAmelCase     , legacy_format=__UpperCAmelCase						)
															snake_case   :					str    						=				tokenizer_p.save_pretrained(__UpperCAmelCase						)

															# Checks it save with the same files
															self.assertSequenceEqual(__UpperCAmelCase     , __UpperCAmelCase						)

															# Checks everything loads correctly in the same way
															snake_case   :					Any    						=				tokenizer_r.from_pretrained(__UpperCAmelCase						)
															snake_case   :					int    						=				tokenizer_p.from_pretrained(__UpperCAmelCase						)

															# Check special tokens are set accordingly on Rust and Python
															for key in tokenizer_pp.special_tokens_map:
																			self.assertTrue(hasattr(__UpperCAmelCase     , __UpperCAmelCase						)						)

															shutil.rmtree(__UpperCAmelCase						)

															# Save tokenizer rust, legacy_format=False
															snake_case   :					Any    						=				tempfile.mkdtemp()

															snake_case   :					Dict    						=				tokenizer_r.save_pretrained(__UpperCAmelCase     , legacy_format=__UpperCAmelCase						)
															snake_case   :					Optional[Any]    						=				tokenizer_p.save_pretrained(__UpperCAmelCase						)

															# Checks it saved the tokenizer.json file
															self.assertTrue(any('''tokenizer.json''' in f for f in tokenizer_r_files						)						)

															# Checks everything loads correctly in the same way
															snake_case   :					Tuple    						=				tokenizer_r.from_pretrained(__UpperCAmelCase						)
															snake_case   :					Optional[Any]    						=				tokenizer_p.from_pretrained(__UpperCAmelCase						)

															# Check special tokens are set accordingly on Rust and Python
															for key in tokenizer_pp.special_tokens_map:
																			self.assertTrue(hasattr(__UpperCAmelCase     , __UpperCAmelCase						)						)

															shutil.rmtree(__UpperCAmelCase						)
@require_torch
@require_sentencepiece
@require_tokenizers
class   a_ (			unittest.TestCase      ):
			A__      :			int       	=  'facebook/mbart-large-en-ro'
			A__      :			Dict       	=  [
			    ' UN Chief Says There Is No Military Solution in Syria',
			    ' Secretary-General Ban Ki-moon says his response to Russia\'s stepped up military support for Syria is that "there is no military solution" to the nearly five-year conflict and more weapons will only worsen the violence and misery for millions of people.',
			]
			A__      :			Optional[int]       	=  [
			    'Şeful ONU declară că nu există o soluţie militară în Siria',
			    'Secretarul General Ban Ki-moon declară că răspunsul său la intensificarea sprijinului militar al Rusiei'
			    ' pentru Siria este că "nu există o soluţie militară" la conflictul de aproape cinci ani şi că noi arme nu vor'
			    ' face decât să înrăutăţească violenţele şi mizeria pentru milioane de oameni.',
			]
			A__      :			Any       	=  [8274, 12_7873, 2_5916, 7, 8622, 2071, 438, 6_7485, 53, 18_7895, 23, 5_1712, 2, EN_CODE]

			@classmethod
			def      lowerCAmelCase(   cls   :  int						):






							"""simple docstring"""
							snake_case   :					Tuple    						=				MBartTokenizer.from_pretrained(
							    cls.checkpoint_name     , src_lang='''en_XX'''     , tgt_lang='''ro_RO'''						)
							snake_case   :					Dict    						=				1
							return cls

			def      lowerCAmelCase(   self   :  Dict						):






							"""simple docstring"""
							self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''ar_AR''']     , 250_001						)
							self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''en_EN''']     , 250_004						)
							self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''ro_RO''']     , 250_020						)

			def      lowerCAmelCase(   self   :  List[Any]						):






							"""simple docstring"""
							snake_case   :					Optional[Any]    						=				self.tokenizer.batch_encode_plus(self.src_text						).input_ids[0]
							self.assertListEqual(self.expected_src_tokens     , __UpperCAmelCase						)

			def      lowerCAmelCase(   self   :  Optional[Any]						):






							"""simple docstring"""
							self.assertIn(__UpperCAmelCase     , self.tokenizer.all_special_ids						)
							snake_case   :					Optional[int]    						=				[RO_CODE, 884, 9_019, 96, 9, 916, 86_792, 36, 18_743, 15_596, 5, 2]
							snake_case   :					List[Any]    						=				self.tokenizer.decode(__UpperCAmelCase     , skip_special_tokens=__UpperCAmelCase						)
							snake_case   :					Optional[int]    						=				self.tokenizer.decode(generated_ids[1:]     , skip_special_tokens=__UpperCAmelCase						)
							self.assertEqual(__UpperCAmelCase     , __UpperCAmelCase						)
							self.assertNotIn(self.tokenizer.eos_token     , __UpperCAmelCase						)

			def      lowerCAmelCase(   self   :  Any						):






							"""simple docstring"""
							snake_case   :					Union[str, Any]    						=				['''this is gunna be a long sentence ''' * 20]
							assert isinstance(src_text[0]     , __UpperCAmelCase						)
							snake_case   :					int    						=				10
							snake_case   :					Optional[Any]    						=				self.tokenizer(__UpperCAmelCase     , max_length=__UpperCAmelCase     , truncation=__UpperCAmelCase						).input_ids[0]
							self.assertEqual(ids[-2]     , 2						)
							self.assertEqual(ids[-1]     , __UpperCAmelCase						)
							self.assertEqual(len(__UpperCAmelCase						)     , __UpperCAmelCase						)

			def      lowerCAmelCase(   self   :  Any						):






							"""simple docstring"""
							self.assertListEqual(self.tokenizer.convert_tokens_to_ids(['''<mask>''', '''ar_AR''']						)     , [250_026, 250_001]						)

			def      lowerCAmelCase(   self   :  List[Any]						):






							"""simple docstring"""
							snake_case   :					List[Any]    						=				tempfile.mkdtemp()
							snake_case   :					str    						=				self.tokenizer.fairseq_tokens_to_ids
							self.tokenizer.save_pretrained(__UpperCAmelCase						)
							snake_case   :					int    						=				MBartTokenizer.from_pretrained(__UpperCAmelCase						)
							self.assertDictEqual(new_tok.fairseq_tokens_to_ids     , __UpperCAmelCase						)

			@require_torch
			def      lowerCAmelCase(   self   :  Dict						):






							"""simple docstring"""
							snake_case   :					Optional[Any]    						=				self.tokenizer(self.src_text     , text_target=self.tgt_text     , padding=__UpperCAmelCase     , return_tensors='''pt'''						)
							snake_case   :					Tuple    						=				shift_tokens_right(batch['''labels''']     , self.tokenizer.pad_token_id						)

							# fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
							assert batch.input_ids[1][-2:].tolist() == [2, EN_CODE]
							assert batch.decoder_input_ids[1][0].tolist() == RO_CODE
							assert batch.decoder_input_ids[1][-1] == 2
							assert batch.labels[1][-2:].tolist() == [2, RO_CODE]

			@require_torch
			def      lowerCAmelCase(   self   :  Optional[Any]						):






							"""simple docstring"""
							snake_case   :					str    						=				self.tokenizer(
							    self.src_text     , text_target=self.tgt_text     , padding=__UpperCAmelCase     , truncation=__UpperCAmelCase     , max_length=len(self.expected_src_tokens						)     , return_tensors='''pt'''     , )

							snake_case   :					Union[str, Any]    						=				shift_tokens_right(batch['''labels''']     , self.tokenizer.pad_token_id						)

							self.assertIsInstance(__UpperCAmelCase     , __UpperCAmelCase						)

							self.assertEqual((2, 14)     , batch.input_ids.shape						)
							self.assertEqual((2, 14)     , batch.attention_mask.shape						)
							snake_case   :					Tuple    						=				batch.input_ids.tolist()[0]
							self.assertListEqual(self.expected_src_tokens     , __UpperCAmelCase						)
							self.assertEqual(2     , batch.decoder_input_ids[0, -1]						)  # EOS
							# Test that special tokens are reset
							self.assertEqual(self.tokenizer.prefix_tokens     , []						)
							self.assertEqual(self.tokenizer.suffix_tokens     , [self.tokenizer.eos_token_id, EN_CODE]						)

			def      lowerCAmelCase(   self   :  Optional[int]						):






							"""simple docstring"""
							snake_case   :					str    						=				self.tokenizer(self.src_text     , padding=__UpperCAmelCase     , truncation=__UpperCAmelCase     , max_length=3     , return_tensors='''pt'''						)
							snake_case   :					List[str]    						=				self.tokenizer(
							    text_target=self.tgt_text     , padding=__UpperCAmelCase     , truncation=__UpperCAmelCase     , max_length=10     , return_tensors='''pt'''						)
							snake_case   :					Optional[int]    						=				targets['''input_ids''']
							snake_case   :					Tuple    						=				shift_tokens_right(__UpperCAmelCase     , self.tokenizer.pad_token_id						)

							self.assertEqual(batch.input_ids.shape[1]     , 3						)
							self.assertEqual(batch.decoder_input_ids.shape[1]     , 10						)



			@require_torch
			def      lowerCAmelCase(   self   :  Dict						):






							"""simple docstring"""
							snake_case   :					List[str]    						=				self.tokenizer._build_translation_inputs(
							    '''A test'''     , return_tensors='''pt'''     , src_lang='''en_XX'''     , tgt_lang='''ar_AR'''						)

							self.assertEqual(
							    nested_simplify(__UpperCAmelCase						)     , {
							        # A, test, EOS, en_XX
							        '''input_ids''': [[62, 3_034, 2, 250_004]],
							        '''attention_mask''': [[1, 1, 1, 1]],
							        # ar_AR
							        '''forced_bos_token_id''': 250_001,
							    }     , )








import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


UpperCAmelCase__							  =						False






class 						lowercase_   (      unittest.TestCase  ):






      '''simple docstring'''




      def  __lowerCAmelCase   ( self :		Dict    )       ->str:




       """simple docstring"""


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()







      @property
      def  __lowerCAmelCase   ( self :		List[str]    )       ->Any:




       """simple docstring"""


       return 12







      @property
      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->Dict:




       """simple docstring"""


       return 12







      @property
      def  __lowerCAmelCase   ( self :		Any    )       ->Any:




       """simple docstring"""


       return 32







      @property
      def  __lowerCAmelCase   ( self :		int    )       ->List[str]:




       """simple docstring"""


       torch.manual_seed(0    )
       a						    = VQModel(
           block_out_channels=[32, 64]					,   in_channels=3					,   out_channels=3					,   down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']					,   up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']					,   latent_channels=3					,   num_vq_embeddings=self.num_embed					,   vq_embed_dim=3					,   )
       return model







      @property
      def  __lowerCAmelCase   ( self :		str    )       ->str:




       """simple docstring"""


       a						    = CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    )
       return tokenizer







      @property
      def  __lowerCAmelCase   ( self :		str    )       ->Union[str, Any]:




       """simple docstring"""


       torch.manual_seed(0    )
       a						    = CLIPTextConfig(
           bos_token_id=0					,   eos_token_id=2					,   hidden_size=self.text_embedder_hidden_size					,   intermediate_size=37					,   layer_norm_eps=1e-0_5					,   num_attention_heads=4					,   num_hidden_layers=5					,   pad_token_id=1					,   vocab_size=1_000					,   )
       return CLIPTextModel(__UpperCAmelCase    )







      @property
      def  __lowerCAmelCase   ( self :		Any    )       ->Tuple:




       """simple docstring"""


       torch.manual_seed(0    )

       a						    = 12
       a						    = 12

       a						    = {
           '''attention_bias''': True,
           '''cross_attention_dim''': 32,
           '''attention_head_dim''': height * width,
           '''num_attention_heads''': 1,
           '''num_vector_embeds''': self.num_embed,
           '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
           '''norm_num_groups''': 32,
           '''sample_size''': width,
           '''activation_fn''': '''geglu-approximate''',
       }

       a						    = TransformeraDModel(**__UpperCAmelCase    )
       return model







      def  __lowerCAmelCase   ( self :		Dict    )       ->Any:




       """simple docstring"""


       a						    = '''cpu'''

       a						    = self.dummy_vqvae
       a						    = self.dummy_text_encoder
       a						    = self.dummy_tokenizer
       a						    = self.dummy_transformer
       a						    = VQDiffusionScheduler(self.num_embed    )
       a						    = LearnedClassifierFreeSamplingEmbeddings(learnable=__UpperCAmelCase    )

       a						    = VQDiffusionPipeline(
           vqvae=__UpperCAmelCase					,   text_encoder=__UpperCAmelCase					,   tokenizer=__UpperCAmelCase					,   transformer=__UpperCAmelCase					,   scheduler=__UpperCAmelCase					,   learned_classifier_free_sampling_embeddings=__UpperCAmelCase					,   )
       a						    = pipe.to(__UpperCAmelCase    )
       pipe.set_progress_bar_config(disable=__UpperCAmelCase    )

       a						    = '''teddy bear playing in the pool'''

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe([prompt]					,   generator=__UpperCAmelCase					,   num_inference_steps=2					,   output_type='''np'''    )
       a						    = output.images

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe(
           [prompt]					,   generator=__UpperCAmelCase					,   output_type='''np'''					,   return_dict=__UpperCAmelCase					,   num_inference_steps=2    )[0]

       a						    = image[0, -3:, -3:, -1]
       a						    = image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 24, 24, 3)

       a						    = np.array([0.6551, 0.6168, 0.5008, 0.5676, 0.5659, 0.4295, 0.6073, 0.5599, 0.4992]    )

       assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-2
       assert np.abs(image_from_tuple_slice.flatten() - expected_slice    ).max() < 1e-2







      def  __lowerCAmelCase   ( self :		Tuple    )       ->str:




       """simple docstring"""


       a						    = '''cpu'''

       a						    = self.dummy_vqvae
       a						    = self.dummy_text_encoder
       a						    = self.dummy_tokenizer
       a						    = self.dummy_transformer
       a						    = VQDiffusionScheduler(self.num_embed    )
       a						    = LearnedClassifierFreeSamplingEmbeddings(
           learnable=__UpperCAmelCase					,   hidden_size=self.text_embedder_hidden_size					,   length=tokenizer.model_max_length    )

       a						    = VQDiffusionPipeline(
           vqvae=__UpperCAmelCase					,   text_encoder=__UpperCAmelCase					,   tokenizer=__UpperCAmelCase					,   transformer=__UpperCAmelCase					,   scheduler=__UpperCAmelCase					,   learned_classifier_free_sampling_embeddings=__UpperCAmelCase					,   )
       a						    = pipe.to(__UpperCAmelCase    )
       pipe.set_progress_bar_config(disable=__UpperCAmelCase    )

       a						    = '''teddy bear playing in the pool'''

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe([prompt]					,   generator=__UpperCAmelCase					,   num_inference_steps=2					,   output_type='''np'''    )
       a						    = output.images

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe(
           [prompt]					,   generator=__UpperCAmelCase					,   output_type='''np'''					,   return_dict=__UpperCAmelCase					,   num_inference_steps=2    )[0]

       a						    = image[0, -3:, -3:, -1]
       a						    = image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 24, 24, 3)

       a						    = np.array([0.6693, 0.6075, 0.4959, 0.5701, 0.5583, 0.4333, 0.6171, 0.5684, 0.4988]    )

       assert np.abs(image_slice.flatten() - expected_slice    ).max() < 2.0
       assert np.abs(image_from_tuple_slice.flatten() - expected_slice    ).max() < 1e-2






@slow
@require_torch_gpu
class 						lowercase_   (      unittest.TestCase  ):






      '''simple docstring'''




      def  __lowerCAmelCase   ( self :		Union[str, Any]    )       ->Union[str, Any]:




       """simple docstring"""


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()







      def  __lowerCAmelCase   ( self :		Dict    )       ->int:




       """simple docstring"""


       a						    = load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''    )

       a						    = VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''    )
       a						    = pipeline.to(__UpperCAmelCase    )
       pipeline.set_progress_bar_config(disable=__UpperCAmelCase    )

       # requires GPU generator for gumbel softmax
       # don't use GPU generator in tests though
       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipeline(
           '''teddy bear playing in the pool'''					,   num_images_per_prompt=1					,   generator=__UpperCAmelCase					,   output_type='''np'''					,   )

       a						    = output.images[0]

       assert image.shape == (256, 256, 3)
       assert np.abs(expected_image - image    ).max() < 2.0


