

import gc
import unittest

from transformers import MODEL_FOR_MASKED_LM_MAPPING, TF_MODEL_FOR_MASKED_LM_MAPPING, FillMaskPipeline, pipeline
from transformers.pipelines import PipelineException
from transformers.testing_utils import (
    is_pipeline_test,
    is_torch_available,
    nested_simplify,
    require_tf,
    require_torch,
    require_torch_gpu,
    slow,
)

from .test_pipelines_common import ANY





@is_pipeline_test
class        lowerCAmelCase__  (					unittest.TestCase			):



       '''simple docstring'''







       lowercase_						=    MODEL_FOR_MASKED_LM_MAPPING
       lowercase_						=    TF_MODEL_FOR_MASKED_LM_MAPPING

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             super().tearDown()
             # clean-up as much as possible GPU memory occupied by PyTorch
             gc.collect()
             if is_torch_available():
                   import torch

                   torch.cuda.empty_cache()

       @require_tf
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =pipeline(task='''fill-mask''' ,   model='''sshleifer/tiny-distilroberta-base''' ,   top_k=2 ,   framework='''tf'''     )
             __A    =unmasker('''My name is <mask>'''     )
             self.assertEqual(
                 nested_simplify(lowercase__ ,   decimals=6     ) ,   [
                     {'''sequence''': '''My name is grouped''', '''score''': 2.1E-05, '''token''': 3_8_0_1_5, '''token_str''': ''' grouped'''},
                     {'''sequence''': '''My name is accuser''', '''score''': 2.1E-05, '''token''': 2_5_5_0_6, '''token_str''': ''' accuser'''},
                 ] ,   )

             __A    =unmasker('''The largest city in France is <mask>'''     )
             self.assertEqual(
                 nested_simplify(lowercase__ ,   decimals=6     ) ,   [
                     {
                         '''sequence''': '''The largest city in France is grouped''',
                         '''score''': 2.1E-05,
                         '''token''': 3_8_0_1_5,
                         '''token_str''': ''' grouped''',
                     },
                     {
                         '''sequence''': '''The largest city in France is accuser''',
                         '''score''': 2.1E-05,
                         '''token''': 2_5_5_0_6,
                         '''token_str''': ''' accuser''',
                     },
                 ] ,   )

             __A    =unmasker('''My name is <mask>''' ,   targets=[''' Patrick''', ''' Clara''', ''' Teven'''] ,   top_k=3     )
             self.assertEqual(
                 nested_simplify(lowercase__ ,   decimals=6     ) ,   [
                     {'''sequence''': '''My name is Clara''', '''score''': 2E-05, '''token''': 1_3_6_0_6, '''token_str''': ''' Clara'''},
                     {'''sequence''': '''My name is Patrick''', '''score''': 2E-05, '''token''': 3_4_9_9, '''token_str''': ''' Patrick'''},
                     {'''sequence''': '''My name is Te''', '''score''': 1.9E-05, '''token''': 2_9_4_1, '''token_str''': ''' Te'''},
                 ] ,   )

       @require_torch
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =pipeline(task='''fill-mask''' ,   model='''sshleifer/tiny-distilroberta-base''' ,   top_k=2 ,   framework='''pt'''     )

             __A    =unmasker('''My name is <mask>'''     )
             self.assertEqual(
                 nested_simplify(lowercase__ ,   decimals=6     ) ,   [
                     {'''sequence''': '''My name is Maul''', '''score''': 2.2E-05, '''token''': 3_5_6_7_6, '''token_str''': ''' Maul'''},
                     {'''sequence''': '''My name isELS''', '''score''': 2.2E-05, '''token''': 1_6_4_1_6, '''token_str''': '''ELS'''},
                 ] ,   )

             __A    =unmasker('''The largest city in France is <mask>'''     )
             self.assertEqual(
                 nested_simplify(lowercase__ ,   decimals=6     ) ,   [
                     {
                         '''sequence''': '''The largest city in France is Maul''',
                         '''score''': 2.2E-05,
                         '''token''': 3_5_6_7_6,
                         '''token_str''': ''' Maul''',
                     },
                     {'''sequence''': '''The largest city in France isELS''', '''score''': 2.2E-05, '''token''': 1_6_4_1_6, '''token_str''': '''ELS'''},
                 ] ,   )

             __A    =unmasker('''My name is <mask>''' ,   targets=[''' Patrick''', ''' Clara''', ''' Teven'''] ,   top_k=3     )
             self.assertEqual(
                 nested_simplify(lowercase__ ,   decimals=6     ) ,   [
                     {'''sequence''': '''My name is Patrick''', '''score''': 2.1E-05, '''token''': 3_4_9_9, '''token_str''': ''' Patrick'''},
                     {'''sequence''': '''My name is Te''', '''score''': 2E-05, '''token''': 2_9_4_1, '''token_str''': ''' Te'''},
                     {'''sequence''': '''My name is Clara''', '''score''': 2E-05, '''token''': 1_3_6_0_6, '''token_str''': ''' Clara'''},
                 ] ,   )

             __A    =unmasker('''My name is <mask> <mask>''' ,   top_k=2     )

             self.assertEqual(
                 nested_simplify(lowercase__ ,   decimals=6     ) ,   [
                     [
                         {
                             '''score''': 2.2E-05,
                             '''token''': 3_5_6_7_6,
                             '''token_str''': ''' Maul''',
                             '''sequence''': '''<s>My name is Maul<mask></s>''',
                         },
                         {'''score''': 2.2E-05, '''token''': 1_6_4_1_6, '''token_str''': '''ELS''', '''sequence''': '''<s>My name isELS<mask></s>'''},
                     ],
                     [
                         {
                             '''score''': 2.2E-05,
                             '''token''': 3_5_6_7_6,
                             '''token_str''': ''' Maul''',
                             '''sequence''': '''<s>My name is<mask> Maul</s>''',
                         },
                         {'''score''': 2.2E-05, '''token''': 1_6_4_1_6, '''token_str''': '''ELS''', '''sequence''': '''<s>My name is<mask>ELS</s>'''},
                     ],
                 ] ,   )

       @require_torch_gpu
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =pipeline('''fill-mask''' ,   model='''hf-internal-testing/tiny-random-distilbert''' ,   device=0 ,   framework='''pt'''     )

             # convert model to fp16
             pipe.model.half()

             __A    =pipe('''Paris is the [MASK] of France.'''     )
             # We actually don't care about the result, we just want to make sure
             # it works, meaning the float16 tensor got casted back to float32
             # for postprocessing.
             self.assertIsInstance(lowercase__ ,   lowercase__     )

       @slow
       @require_torch
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =pipeline(task='''fill-mask''' ,   model='''distilroberta-base''' ,   top_k=2 ,   framework='''pt'''     )
             self.run_large_test(lowercase__     )

       @slow
       @require_tf
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =pipeline(task='''fill-mask''' ,   model='''distilroberta-base''' ,   top_k=2 ,   framework='''tf'''     )
             self.run_large_test(lowercase__     )

       def     __UpperCamelCase     (							self ,   lowercase__     ):







             '''simple docstring'''



             __A    =unmasker('''My name is <mask>'''     )
             self.assertEqual(
                 nested_simplify(lowercase__     ) ,   [
                     {'''sequence''': '''My name is John''', '''score''': 0.008, '''token''': 6_1_0, '''token_str''': ''' John'''},
                     {'''sequence''': '''My name is Chris''', '''score''': 0.007, '''token''': 1_5_7_3, '''token_str''': ''' Chris'''},
                 ] ,   )
             __A    =unmasker('''The largest city in France is <mask>'''     )
             self.assertEqual(
                 nested_simplify(lowercase__     ) ,   [
                     {
                         '''sequence''': '''The largest city in France is Paris''',
                         '''score''': 0.251,
                         '''token''': 2_2_0_1,
                         '''token_str''': ''' Paris''',
                     },
                     {
                         '''sequence''': '''The largest city in France is Lyon''',
                         '''score''': 0.214,
                         '''token''': 1_2_7_9_0,
                         '''token_str''': ''' Lyon''',
                     },
                 ] ,   )

             __A    =unmasker('''My name is <mask>''' ,   targets=[''' Patrick''', ''' Clara''', ''' Teven'''] ,   top_k=3     )
             self.assertEqual(
                 nested_simplify(lowercase__     ) ,   [
                     {'''sequence''': '''My name is Patrick''', '''score''': 0.005, '''token''': 3_4_9_9, '''token_str''': ''' Patrick'''},
                     {'''sequence''': '''My name is Clara''', '''score''': 0.000, '''token''': 1_3_6_0_6, '''token_str''': ''' Clara'''},
                     {'''sequence''': '''My name is Te''', '''score''': 0.000, '''token''': 2_9_4_1, '''token_str''': ''' Te'''},
                 ] ,   )

       @require_torch
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =pipeline(task='''fill-mask''' ,   model='''sshleifer/tiny-distilroberta-base''' ,   framework='''pt'''     )
             __A    =None
             __A    =None
             self.run_pipeline_test(lowercase__ ,   []     )

       @require_tf
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =pipeline(task='''fill-mask''' ,   model='''sshleifer/tiny-distilroberta-base''' ,   framework='''tf'''     )
             __A    =None
             __A    =None
             self.run_pipeline_test(lowercase__ ,   []     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             if tokenizer is None or tokenizer.mask_token_id is None:
                   self.skipTest('''The provided tokenizer has no mask token, (probably reformer or wav2vec2)'''     )

             __A    =FillMaskPipeline(model=lowercase__ ,   tokenizer=lowercase__     )
             __A    =[
                 f'''This is another {tokenizer.mask_token} test''',
             ]
             return fill_masker, examples

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =fill_masker.tokenizer
             __A    =fill_masker.model

             __A    =fill_masker(
                 f'''This is a {tokenizer.mask_token}''' ,   )
             self.assertEqual(
                 lowercase__ ,   [
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                 ] ,   )

             __A    =fill_masker([f'''This is a {tokenizer.mask_token}''']     )
             self.assertEqual(
                 lowercase__ ,   [
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                 ] ,   )

             __A    =fill_masker([f'''This is a {tokenizer.mask_token}''', f'''Another {tokenizer.mask_token} great test.''']     )
             self.assertEqual(
                 lowercase__ ,   [
                     [
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     ],
                     [
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     ],
                 ] ,   )

             with self.assertRaises(lowercase__     ):
                   fill_masker([None]     )
             # No mask_token is not supported
             with self.assertRaises(lowercase__     ):
                   fill_masker('''This is'''     )

             self.run_test_top_k(lowercase__ ,   lowercase__     )
             self.run_test_targets(lowercase__ ,   lowercase__     )
             self.run_test_top_k_targets(lowercase__ ,   lowercase__     )
             self.fill_mask_with_duplicate_targets_and_top_k(lowercase__ ,   lowercase__     )
             self.fill_mask_with_multiple_masks(lowercase__ ,   lowercase__     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =tokenizer.get_vocab()
             __A    =sorted(vocab.keys()     )[:2]
             # Pipeline argument
             __A    =FillMaskPipeline(model=lowercase__ ,   tokenizer=lowercase__ ,   targets=lowercase__     )
             __A    =fill_masker(f'''This is a {tokenizer.mask_token}'''     )
             self.assertEqual(
                 lowercase__ ,   [
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                 ] ,   )
             __A    ={vocab[el] for el in targets}
             self.assertEqual({el['''token'''] for el in outputs} ,   lowercase__     )
             __A    =[tokenizer.decode([x]     ) for x in target_ids]
             self.assertEqual({el['''token_str'''] for el in outputs} ,   set(lowercase__     )     )

             # Call argument
             __A    =FillMaskPipeline(model=lowercase__ ,   tokenizer=lowercase__     )
             __A    =fill_masker(f'''This is a {tokenizer.mask_token}''' ,   targets=lowercase__     )
             self.assertEqual(
                 lowercase__ ,   [
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                 ] ,   )
             __A    ={vocab[el] for el in targets}
             self.assertEqual({el['''token'''] for el in outputs} ,   lowercase__     )
             __A    =[tokenizer.decode([x]     ) for x in target_ids]
             self.assertEqual({el['''token_str'''] for el in outputs} ,   set(lowercase__     )     )

             # Score equivalence
             __A    =fill_masker(f'''This is a {tokenizer.mask_token}''' ,   targets=lowercase__     )
             __A    =[top_mask['''token_str'''] for top_mask in outputs]
             __A    =[top_mask['''score'''] for top_mask in outputs]

             # For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
             if set(lowercase__     ) == set(lowercase__     ):
                   __A    =fill_masker(f'''This is a {tokenizer.mask_token}''' ,   targets=lowercase__     )
                   __A    =[top_mask['''score'''] for top_mask in unmasked_targets]
                   self.assertEqual(nested_simplify(lowercase__     ) ,   nested_simplify(lowercase__     )     )

             # Raises with invalid
             with self.assertRaises(lowercase__     ):
                   __A    =fill_masker(f'''This is a {tokenizer.mask_token}''' ,   targets=[]     )
             # For some tokenizers, `""` is actually in the vocabulary and the expected error won't raised
             if "" not in tokenizer.get_vocab():
                   with self.assertRaises(lowercase__     ):
                         __A    =fill_masker(f'''This is a {tokenizer.mask_token}''' ,   targets=['''''']     )
                   with self.assertRaises(lowercase__     ):
                         __A    =fill_masker(f'''This is a {tokenizer.mask_token}''' ,   targets=''''''     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =FillMaskPipeline(model=lowercase__ ,   tokenizer=lowercase__ ,   top_k=2     )
             __A    =fill_masker(f'''This is a {tokenizer.mask_token}'''     )
             self.assertEqual(
                 lowercase__ ,   [
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                 ] ,   )

             __A    =FillMaskPipeline(model=lowercase__ ,   tokenizer=lowercase__     )
             __A    =fill_masker(f'''This is a {tokenizer.mask_token}''' ,   top_k=2     )
             self.assertEqual(
                 lowercase__ ,   [
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                 ] ,   )
             self.assertEqual(nested_simplify(lowercase__     ) ,   nested_simplify(lowercase__     )     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =tokenizer.get_vocab()
             __A    =FillMaskPipeline(model=lowercase__ ,   tokenizer=lowercase__     )

             # top_k=2, ntargets=3
             __A    =sorted(vocab.keys()     )[:3]
             __A    =fill_masker(f'''This is a {tokenizer.mask_token}''' ,   top_k=2 ,   targets=lowercase__     )

             # If we use the most probably targets, and filter differently, we should still
             # have the same results
             __A    =[el['''token_str'''] for el in sorted(lowercase__ ,   key=lambda lowercase__     : x["score"] ,   reverse=lowercase__     )]
             # For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
             if set(lowercase__     ).issubset(lowercase__     ):
                   __A    =fill_masker(f'''This is a {tokenizer.mask_token}''' ,   top_k=3 ,   targets=lowercase__     )
                   # They should yield exactly the same result
                   self.assertEqual(nested_simplify(lowercase__     ) ,   nested_simplify(lowercase__     )     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =FillMaskPipeline(model=lowercase__ ,   tokenizer=lowercase__     )
             __A    =tokenizer.get_vocab()
             # String duplicates + id duplicates
             __A    =sorted(vocab.keys()     )[:3]
             __A    =[targets[0], targets[1], targets[0], targets[2], targets[1]]
             __A    =fill_masker(f'''My name is {tokenizer.mask_token}''' ,   targets=lowercase__ ,   top_k=1_0     )

             # The target list contains duplicates, so we can't output more
             # than them
             self.assertEqual(len(lowercase__     ) ,   3     )


       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =FillMaskPipeline(model=lowercase__ ,   tokenizer=lowercase__     )

             __A    =fill_masker(
                 f'''This is a {tokenizer.mask_token} {tokenizer.mask_token} {tokenizer.mask_token}''' ,   top_k=2     )
             self.assertEqual(
                 lowercase__ ,   [
                     [
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     ],
                     [
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     ],
                     [
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                         {'''sequence''': ANY(lowercase__     ), '''score''': ANY(lowercase__     ), '''token''': ANY(lowercase__     ), '''token_str''': ANY(lowercase__     )},
                     ],
                 ] ,   )





import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()





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



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



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

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =StableDiffusionKDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''     )
             __A    =sd_pipe.to(lowercase__     )
             sd_pipe.set_progress_bar_config(disable=lowercase__     )

             sd_pipe.set_scheduler('''sample_euler'''     )

             __A    ='''A painting of a squirrel eating a burger'''
             __A    =torch.manual_seed(0     )
             __A    =sd_pipe([prompt] ,   generator=lowercase__ ,   guidance_scale=9.0 ,   num_inference_steps=2_0 ,   output_type='''np'''     )

             __A    =output.images

             __A    =image[0, -3:, -3:, -1]

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             __A    =np.array([0.0447, 0.0492, 0.0468, 0.0408, 0.0383, 0.0408, 0.0354, 0.0380, 0.0339]     )

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

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''     )
             __A    =sd_pipe.to(lowercase__     )
             sd_pipe.set_progress_bar_config(disable=lowercase__     )

             sd_pipe.set_scheduler('''sample_euler'''     )

             __A    ='''A painting of a squirrel eating a burger'''
             __A    =torch.manual_seed(0     )
             __A    =sd_pipe([prompt] ,   generator=lowercase__ ,   guidance_scale=9.0 ,   num_inference_steps=2_0 ,   output_type='''np'''     )

             __A    =output.images

             __A    =image[0, -3:, -3:, -1]

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             __A    =np.array([0.1237, 0.1320, 0.1438, 0.1359, 0.1390, 0.1132, 0.1277, 0.1175, 0.1112]     )

             assert np.abs(image_slice.flatten() - expected_slice     ).max() < 5E-1


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''     )
             __A    =sd_pipe.to(lowercase__     )
             sd_pipe.set_progress_bar_config(disable=lowercase__     )

             sd_pipe.set_scheduler('''sample_dpmpp_2m'''     )

             __A    ='''A painting of a squirrel eating a burger'''
             __A    =torch.manual_seed(0     )
             __A    =sd_pipe(
                 [prompt] ,   generator=lowercase__ ,   guidance_scale=7.5 ,   num_inference_steps=1_5 ,   output_type='''np''' ,   use_karras_sigmas=lowercase__ ,   )

             __A    =output.images

             __A    =image[0, -3:, -3:, -1]

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             __A    =np.array(
                 [0.1138_1689, 0.1211_2921, 0.138_9457, 0.1254_9606, 0.124_4964, 0.1083_1517, 0.1156_2866, 0.1086_7816, 0.1049_9048]     )

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


