



import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import OPTConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
 import os

 # The slow tests are often failing with OOM error on GPU
 # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
 # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
 a			:     Any     = 'platform'

 import jax
 import jax.numpy as jnp

 from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer







def    lowerCAmelCase_  (lowerCAmelCase__:		Tuple		,	lowerCAmelCase__:		str		,	lowerCAmelCase__:		Optional[int]=None		,	lowerCAmelCase__:		List[str]=None     ):





   """simple docstring"""


   if attention_mask is None:
      UpperCAmelCase_:		Optional[int]       			=    np.where(input_ids != config.pad_token_id		,	1		,	0     )
   return {
       "input_ids": input_ids,
       "attention_mask": attention_mask,
   }







@require_flax
class 							_a							:

   def __init__(self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_=13,      SCREAMING_SNAKE_CASE_=7,      SCREAMING_SNAKE_CASE_=True,      SCREAMING_SNAKE_CASE_=False,      SCREAMING_SNAKE_CASE_=99,      SCREAMING_SNAKE_CASE_=16,      SCREAMING_SNAKE_CASE_=2,      SCREAMING_SNAKE_CASE_=4,      SCREAMING_SNAKE_CASE_=4,      SCREAMING_SNAKE_CASE_="gelu",      SCREAMING_SNAKE_CASE_=0.1,      SCREAMING_SNAKE_CASE_=0.1,      SCREAMING_SNAKE_CASE_=20,      SCREAMING_SNAKE_CASE_=2,      SCREAMING_SNAKE_CASE_=1,      SCREAMING_SNAKE_CASE_=0,      SCREAMING_SNAKE_CASE_=16,      SCREAMING_SNAKE_CASE_=16,      SCREAMING_SNAKE_CASE_=0.0_2,      )					->      str:
      UpperCAmelCase_:		Optional[Any]       			=    parent
      UpperCAmelCase_:		Optional[int]       			=    batch_size
      UpperCAmelCase_:		Any       			=    seq_length
      UpperCAmelCase_:		Any       			=    is_training
      UpperCAmelCase_:		Tuple       			=    use_labels
      UpperCAmelCase_:		Dict       			=    vocab_size
      UpperCAmelCase_:		str       			=    hidden_size
      UpperCAmelCase_:		Tuple       			=    num_hidden_layers
      UpperCAmelCase_:		Dict       			=    num_attention_heads
      UpperCAmelCase_:		str       			=    intermediate_size
      UpperCAmelCase_:		Any       			=    hidden_act
      UpperCAmelCase_:		Tuple       			=    hidden_dropout_prob
      UpperCAmelCase_:		List[str]       			=    attention_probs_dropout_prob
      UpperCAmelCase_:		Optional[Any]       			=    max_position_embeddings
      UpperCAmelCase_:		int       			=    eos_token_id
      UpperCAmelCase_:		Any       			=    pad_token_id
      UpperCAmelCase_:		List[Any]       			=    bos_token_id
      UpperCAmelCase_:		Dict       			=    embed_dim
      UpperCAmelCase_:		List[Any]       			=    word_embed_proj_dim
      UpperCAmelCase_:		str       			=    initializer_range
      UpperCAmelCase_:		Tuple       			=    False

   def     __snake_case					(self				)					->      int:
      UpperCAmelCase_:		Tuple       			=    np.clip(ids_tensor([self.batch_size, self.seq_length - 1],      self.vocab_size				),      3,      self.vocab_size				)
      UpperCAmelCase_:		Union[str, Any]       			=    np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1),      dtype=np.intaa				)),      -1				)

      UpperCAmelCase_:		List[str]       			=    OPTConfig(
          vocab_size=self.vocab_size,      hidden_size=self.hidden_size,      num_hidden_layers=self.num_hidden_layers,      num_attention_heads=self.num_attention_heads,      ffn_dim=self.intermediate_size,      dropout=self.hidden_dropout_prob,      attention_dropout=self.attention_probs_dropout_prob,      max_position_embeddings=self.max_position_embeddings,      eos_token_id=self.eos_token_id,      bos_token_id=self.bos_token_id,      pad_token_id=self.pad_token_id,      embed_dim=self.embed_dim,      is_encoder_decoder=lowercase_,      word_embed_proj_dim=self.word_embed_proj_dim,      initializer_range=self.initializer_range,      use_cache=lowercase_,      )
      UpperCAmelCase_:		Optional[int]       			=    prepare_opt_inputs_dict(lowercase_,      lowercase_				)
      return config, inputs_dict

   def     __snake_case					(self				)					->      Tuple:
      UpperCAmelCase_:		str       			=    self.prepare_config_and_inputs()
      return config, inputs_dict

   def     __snake_case					(self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_				)					->      Optional[Any]:
      UpperCAmelCase_:		str       			=    20
      UpperCAmelCase_:		Any       			=    model_class_name(lowercase_				)

      UpperCAmelCase_:		Tuple       			=    inputs_dict['''input_ids''']
      UpperCAmelCase_:		List[Any]       			=    inputs_dict['''attention_mask''']

      UpperCAmelCase_:		str       			=    model.init_cache(input_ids.shape[0],      lowercase_				)
      UpperCAmelCase_:		str       			=    jnp.ones((input_ids.shape[0], max_length),      dtype="""i4"""				)

      UpperCAmelCase_:		int       			=    jnp.broadcast_to(
          jnp.arange(input_ids.shape[-1] - 1				)[None, :],      (input_ids.shape[0], input_ids.shape[-1] - 1),      )
      UpperCAmelCase_:		List[str]       			=    model(
          input_ids[:, :-1],      attention_mask=lowercase_,      past_key_values=lowercase_,      position_ids=lowercase_,      )

      UpperCAmelCase_:		List[str]       			=    jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]],      dtype="""i4"""				)
      UpperCAmelCase_:		Any       			=    model(
          input_ids[:, -1:],      attention_mask=lowercase_,      past_key_values=outputs_cache.past_key_values,      position_ids=lowercase_,      )

      UpperCAmelCase_:		List[Any]       			=    model(lowercase_				)

      UpperCAmelCase_:		List[Any]       			=    np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])				)				)
      self.parent.assertTrue(diff < 1E-3,      msg=f'Max diff is {diff}'				)

   def     __snake_case					(self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_				)					->      Tuple:
      UpperCAmelCase_:		Optional[int]       			=    20
      UpperCAmelCase_:		Optional[int]       			=    model_class_name(lowercase_				)

      UpperCAmelCase_:		Optional[Any]       			=    (
          inputs_dict['''input_ids'''],
          inputs_dict['''attention_mask'''],
      )

      UpperCAmelCase_:		List[str]       			=    jnp.concatenate(
          [
              attention_mask,
              jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])				),
          ],      axis=-1,      )

      UpperCAmelCase_:		int       			=    model.init_cache(input_ids.shape[0],      lowercase_				)
      UpperCAmelCase_:		str       			=    jnp.broadcast_to(
          jnp.arange(input_ids.shape[-1] - 1				)[None, :],      (input_ids.shape[0], input_ids.shape[-1] - 1),      )

      UpperCAmelCase_:		Optional[int]       			=    model(
          input_ids[:, :-1],      attention_mask=lowercase_,      past_key_values=lowercase_,      position_ids=lowercase_,      )
      UpperCAmelCase_:		str       			=    jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]],      dtype="""i4"""				)
      UpperCAmelCase_:		List[str]       			=    model(
          input_ids[:, -1:],      past_key_values=outputs_cache.past_key_values,      attention_mask=lowercase_,      position_ids=lowercase_,      )

      UpperCAmelCase_:		Optional[Any]       			=    model(lowercase_,      attention_mask=lowercase_				)

      UpperCAmelCase_:		Optional[int]       			=    np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])				)				)
      self.parent.assertTrue(diff < 1E-3,      msg=f'Max diff is {diff}'				)







@require_flax
class 							_a							(      __UpperCAmelCase   ,		unittest.TestCase   ,		__UpperCAmelCase				):
   A					=  (FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
   A					=  () if is_flax_available() else ()

   def     __snake_case					(self				)					->      Optional[Any]:
      UpperCAmelCase_:		Tuple       			=    FlaxOPTModelTester(self				)

   def     __snake_case					(self				)					->      Dict:
      UpperCAmelCase_:		Optional[Any]       			=    self.model_tester.prepare_config_and_inputs()
      for model_class in self.all_model_classes:
         self.model_tester.check_use_cache_forward(lowercase_,      lowercase_,      lowercase_				)

   def     __snake_case					(self				)					->      Union[str, Any]:
      UpperCAmelCase_:		int       			=    self.model_tester.prepare_config_and_inputs()
      for model_class in self.all_model_classes:
         self.model_tester.check_use_cache_forward_with_attn_mask(lowercase_,      lowercase_,      lowercase_				)

   @slow
   def     __snake_case					(self				)					->      Optional[Any]:

      for model_class_name in self.all_model_classes:
         UpperCAmelCase_:		int       			=    model_class_name.from_pretrained("""facebook/opt-125m"""				)
         UpperCAmelCase_:		List[str]       			=    np.ones((1, 1)				) * model.config.eos_token_id
         UpperCAmelCase_:		Tuple       			=    model(lowercase_				)
         self.assertIsNotNone(lowercase_				)







@require_sentencepiece
@require_flax
class 							_a							(      unittest.TestCase				):

   @slow
   def     __snake_case					(self				)					->      List[str]:
      UpperCAmelCase_:		List[Any]       			=    FlaxOPTModel.from_pretrained("""facebook/opt-350m"""				)
      UpperCAmelCase_:		List[Any]       			=    jnp.array([[0, 31414, 232, 328, 740, 1140, 12695, 69, 46078, 1588, 2]]				)
      UpperCAmelCase_:		Tuple       			=    model(input_ids=lowercase_				).last_hidden_state
      UpperCAmelCase_:		int       			=    (1, 11, 512)
      self.assertEqual(output.shape,      lowercase_				)
      UpperCAmelCase_:		List[Any]       			=    jnp.array(
          [[-0.2_8_6_7, -1.9_2_5_6, -0.3_0_6_2], [-1.2_7_1_1, -0.1_3_3_7, -0.1_8_9_7], [0.4_1_0_9, 0.1_1_8_7, -1.3_1_4_2]]				)
      self.assertTrue(jnp.allclose(output[:, :3, :3],      lowercase_,      atol=4E-2				)				)







@require_flax
@slow
class 							_a							(      unittest.TestCase				):

   def     __snake_case					(self				)					->      Union[str, Any]:
      super().setUp()
      UpperCAmelCase_:		Optional[int]       			=    '''facebook/opt-350m'''

   def     __snake_case					(self				)					->      Union[str, Any]:
      UpperCAmelCase_:		Any       			=    FlaxOPTForCausalLM.from_pretrained(self.path_model				)
      UpperCAmelCase_:		List[Any]       			=    GPTaTokenizer.from_pretrained(self.path_model				)

      UpperCAmelCase_:		Dict       			=    [
          '''Today is a beautiful day and I want to''',
          '''In the city of''',
          '''Paris is the capital of France and''',
          '''Computers and mobile phones have taken''',
      ]
      # verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
      UpperCAmelCase_:		Optional[Any]       			=    tokenizer(lowercase_,      return_tensors="""jax""",      padding=lowercase_,      add_special_tokens=lowercase_				)
      UpperCAmelCase_:		Optional[int]       			=    model(inputs.input_ids,      attention_mask=inputs.attention_mask				)[0].mean(axis=-1				)
      UpperCAmelCase_:		Optional[int]       			=    jnp.array(
          [
              [1.3_8_5_1, -1_3.8_9_2_3, -1_0.5_2_2_9, -1_0.7_5_3_3, -0.2_3_0_9, -1_0.2_3_8_4, -0.5_3_6_5, -9.0_9_4_7, -5.1_6_7_0],
              [-4.7_0_7_3, -1_0.6_2_7_6, -3.9_4_1_5, -2_1.5_2_4_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2],
              [0.6_2_4_7, -3.4_2_2_9, -8.9_1_7_9, -1.4_2_9_7, -1_4.1_6_5_0, 1.4_1_4_6, -9.0_2_1_8, -0.2_7_0_3, -0.2_7_0_3],
              [6.4_7_8_3, -1.9_9_1_3, -1_0.7_9_2_6, -2.3_3_3_6, 1.5_0_9_2, -0.9_9_7_4, -6.8_2_1_3, 1.3_4_7_7, 1.3_4_7_7],
          ]				)
      self.assertTrue(jnp.allclose(lowercase_,      lowercase_,      atol=4E-2				)				)

      UpperCAmelCase_:		Optional[int]       			=    jax.jit(lowercase_				)
      UpperCAmelCase_:		Tuple       			=    model(inputs.input_ids,      attention_mask=inputs.attention_mask				)[0].mean(axis=-1				)
      self.assertTrue(jnp.allclose(lowercase_,      lowercase_,      atol=4E-2				)				)







@require_flax
@slow
class 							_a							(      unittest.TestCase				):

   @property
   def     __snake_case					(self				)					->      Optional[Any]:
      return [
          "Today is a beautiful day and I want",
          "In the city of",
          "Paris is the capital of France and",
          "Computers and mobile phones have taken",
      ]

   def     __snake_case					(self				)					->      Any:
      UpperCAmelCase_:		Dict       			=    '''facebook/opt-125m'''

      UpperCAmelCase_:		List[Any]       			=    [
          '''Today is a beautiful day and I want to''',
          '''In the city of New York, the city''',
          '''Paris is the capital of France and the capital''',
          '''Computers and mobile phones have taken over the''',
      ]

      UpperCAmelCase_:		Union[str, Any]       			=    []

      UpperCAmelCase_:		Dict       			=    FlaxOPTForCausalLM.from_pretrained(lowercase_				)
      UpperCAmelCase_:		int       			=    GPTaTokenizer.from_pretrained(lowercase_				)

      for prompt in self.prompts:
         UpperCAmelCase_:		int       			=    tokenizer(lowercase_,      return_tensors="""jax"""				).input_ids

         UpperCAmelCase_:		Optional[int]       			=    model.generate(lowercase_,      max_length=10				)
         UpperCAmelCase_:		Dict       			=    generated_ids[0]

         UpperCAmelCase_:		str       			=    tokenizer.batch_decode(lowercase_,      skip_special_tokens=lowercase_				)
         predicted_outputs += generated_string

      self.assertListEqual(lowercase_,      lowercase_				)

   def     __snake_case					(self				)					->      int:
      UpperCAmelCase_:		int       			=    '''facebook/opt-350m'''

      UpperCAmelCase_:		List[Any]       			=    [
          '''Today is a beautiful day and I want to''',
          '''In the city of San Francisco, the city''',
          '''Paris is the capital of France and the capital''',
          '''Computers and mobile phones have taken over the''',
      ]

      UpperCAmelCase_:		Union[str, Any]       			=    []
      UpperCAmelCase_:		Any       			=    FlaxOPTForCausalLM.from_pretrained(lowercase_				)
      UpperCAmelCase_:		Optional[Any]       			=    GPTaTokenizer.from_pretrained(lowercase_				)

      for prompt in self.prompts:
         UpperCAmelCase_:		str       			=    tokenizer(lowercase_,      return_tensors="""jax"""				).input_ids

         UpperCAmelCase_:		int       			=    model.generate(lowercase_,      max_length=10				)
         UpperCAmelCase_:		Any       			=    generated_ids[0]

         UpperCAmelCase_:		Tuple       			=    tokenizer.batch_decode(lowercase_,      skip_special_tokens=lowercase_				)
         predicted_outputs += generated_string

      self.assertListEqual(lowercase_,      lowercase_				)

   def     __snake_case					(self				)					->      int:
      UpperCAmelCase_:		List[Any]       			=    '''facebook/opt-125m'''
      UpperCAmelCase_:		Any       			=    [
          '''Today is a beautiful day and I want to thank''',
          '''In the city of Rome Canaver Canaver Canaver Canaver''',
      ]
      UpperCAmelCase_:		Tuple       			=    FlaxOPTForCausalLM.from_pretrained(lowercase_				)
      UpperCAmelCase_:		Any       			=    GPTaTokenizer.from_pretrained(lowercase_				)
      UpperCAmelCase_:		Any       			=    tokenizer(
          [
              """Today is a beautiful day and I want to""",
              """In the city of""",
          ],      return_tensors="""jax""",      padding=lowercase_,      )

      UpperCAmelCase_:		Any       			=    jax.jit(model.generate				)

      UpperCAmelCase_:		List[str]       			=    jit_generate(inputs["""input_ids"""],      attention_mask=inputs["""attention_mask"""]				).sequences

      UpperCAmelCase_:		List[Any]       			=    tokenizer.batch_decode(lowercase_,      skip_special_tokens=lowercase_				)

      self.assertIsNotNone(lowercase_,      lowercase_				)

   def     __snake_case					(self				)					->      Tuple:
      UpperCAmelCase_:		str       			=    '''facebook/opt-350m'''

      UpperCAmelCase_:		Optional[Any]       			=    GPTaTokenizer.from_pretrained(lowercase_				)
      UpperCAmelCase_:		Union[str, Any]       			=    FlaxOPTForCausalLM.from_pretrained(lowercase_				)

      UpperCAmelCase_:		int       			=    '''left'''

      # use different length sentences to test batching
      UpperCAmelCase_:		Tuple       			=    [
          '''Hello, my dog is a little''',
          '''Today, I''',
      ]

      UpperCAmelCase_:		str       			=    tokenizer(lowercase_,      return_tensors="""jax""",      padding=lowercase_				)
      UpperCAmelCase_:		List[str]       			=    inputs['''input_ids''']

      UpperCAmelCase_:		Union[str, Any]       			=    model.generate(input_ids=lowercase_,      attention_mask=inputs["""attention_mask"""],      trace=lowercase_				)

      UpperCAmelCase_:		List[Any]       			=    tokenizer(sentences[0],      return_tensors="""jax"""				).input_ids
      UpperCAmelCase_:		int       			=    model.generate(input_ids=lowercase_				)

      UpperCAmelCase_:		List[str]       			=    inputs_non_padded.shape[-1] - inputs['''attention_mask'''][-1].sum()
      UpperCAmelCase_:		Any       			=    tokenizer(sentences[1],      return_tensors="""jax"""				).input_ids
      UpperCAmelCase_:		Optional[Any]       			=    model.generate(input_ids=lowercase_,      max_length=model.config.max_length - num_paddings				)

      UpperCAmelCase_:		int       			=    tokenizer.batch_decode(outputs[0],      skip_special_tokens=lowercase_				)
      UpperCAmelCase_:		List[str]       			=    tokenizer.decode(output_non_padded[0][0],      skip_special_tokens=lowercase_				)
      UpperCAmelCase_:		List[Any]       			=    tokenizer.decode(output_padded[0][0],      skip_special_tokens=lowercase_				)

      UpperCAmelCase_:		List[str]       			=    [
          '''Hello, my dog is a little bit of a dork.\nI\'m a little bit''',
          '''Today, I was in the middle of a conversation with a friend about the''',
      ]
      self.assertListEqual(lowercase_,      lowercase_				)
      self.assertListEqual(lowercase_,      [non_padded_sentence, padded_sentence]				)











import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







class 							_a							(      unittest.TestCase				):

			def     __snake_case					(self				)					->      Any:
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()

			@property
			def     __snake_case					(self				)					->      Tuple:
						UpperCAmelCase_:		Tuple       			=    1
						UpperCAmelCase_:		Dict       			=    3
						UpperCAmelCase_:		Optional[int]       			=    (32, 32)

						UpperCAmelCase_:		int       			=    floats_tensor((batch_size, num_channels) + sizes,      rng=random.Random(0				)				).to(SCREAMING_SNAKE_CASE_				)
						return image

			@property
			def     __snake_case					(self				)					->      Optional[Any]:
						torch.manual_seed(0				)
						UpperCAmelCase_:		Tuple       			=    UNetaDConditionModel(
						    block_out_channels=(32, 64),      layers_per_block=2,      sample_size=32,      in_channels=4,      out_channels=4,      down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D"""),      up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D"""),      cross_attention_dim=32,      )
						return model

			@property
			def     __snake_case					(self				)					->      str:
						torch.manual_seed(0				)
						UpperCAmelCase_:		int       			=    AutoencoderKL(
						    block_out_channels=[32, 64],      in_channels=3,      out_channels=3,      down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],      up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],      latent_channels=4,      )
						return model

			@property
			def     __snake_case					(self				)					->      List[str]:
						torch.manual_seed(0				)
						UpperCAmelCase_:		Tuple       			=    CLIPTextConfig(
						    bos_token_id=0,      eos_token_id=2,      hidden_size=32,      intermediate_size=37,      layer_norm_eps=1E-05,      num_attention_heads=4,      num_hidden_layers=5,      pad_token_id=1,      vocab_size=1000,      )
						return CLIPTextModel(SCREAMING_SNAKE_CASE_				)

			@property
			def     __snake_case					(self				)					->      str:
						def extract(*SCREAMING_SNAKE_CASE_,      **SCREAMING_SNAKE_CASE_				):







									class 							_a							:

												def __init__(self				)					->      Any:
															UpperCAmelCase_:		Optional[Any]       			=    torch.ones([0]				)

												def     __snake_case					(self,      SCREAMING_SNAKE_CASE_				)					->      int:
															self.pixel_values.to(SCREAMING_SNAKE_CASE_				)
															return self






									return Out()

						return extract

			def     __snake_case					(self				)					->      List[Any]:
						UpperCAmelCase_:		Any       			=    """cpu"""  # ensure determinism for the device-dependent torch.Generator
						UpperCAmelCase_:		Dict       			=    self.dummy_cond_unet
						UpperCAmelCase_:		Union[str, Any]       			=    DDIMScheduler(
						    beta_start=0.0_0_0_8_5,      beta_end=0.0_1_2,      beta_schedule="""scaled_linear""",      clip_sample=SCREAMING_SNAKE_CASE_,      set_alpha_to_one=SCREAMING_SNAKE_CASE_,      )

						UpperCAmelCase_:		List[str]       			=    self.dummy_vae
						UpperCAmelCase_:		Union[str, Any]       			=    self.dummy_text_encoder
						UpperCAmelCase_:		Dict       			=    CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)

						# make sure here that pndm scheduler skips prk
						UpperCAmelCase_:		str       			=    StableDiffusionPipeline(
						    unet=SCREAMING_SNAKE_CASE_,      scheduler=SCREAMING_SNAKE_CASE_,      vae=SCREAMING_SNAKE_CASE_,      text_encoder=SCREAMING_SNAKE_CASE_,      tokenizer=SCREAMING_SNAKE_CASE_,      safety_checker=SCREAMING_SNAKE_CASE_,      feature_extractor=self.dummy_extractor,      )
						UpperCAmelCase_:		str       			=    sd_pipe.to(SCREAMING_SNAKE_CASE_				)
						sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		int       			=    """A painting of a squirrel eating a burger"""

						UpperCAmelCase_:		Tuple       			=    torch.Generator(device=SCREAMING_SNAKE_CASE_				).manual_seed(0				)
						UpperCAmelCase_:		Dict       			=    sd_pipe([prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=6.0,      num_inference_steps=2,      output_type="""np"""				)
						UpperCAmelCase_:		Tuple       			=    output.images

						UpperCAmelCase_:		List[Any]       			=    torch.Generator(device=SCREAMING_SNAKE_CASE_				).manual_seed(0				)
						UpperCAmelCase_:		List[str]       			=    sd_pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=6.0,      num_inference_steps=2,      output_type="""np""",      return_dict=SCREAMING_SNAKE_CASE_,      )[0]

						UpperCAmelCase_:		int       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		List[str]       			=    image_from_tuple[0, -3:, -3:, -1]

						assert image.shape == (1, 64, 64, 3)
						UpperCAmelCase_:		Dict       			=    np.array([0.5_7_5_3, 0.6_1_1_4, 0.5_0_0_1, 0.5_0_3_4, 0.5_4_7_0, 0.4_7_2_9, 0.4_9_7_1, 0.4_8_6_7, 0.4_8_6_7]				)

						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     __snake_case					(self				)					->      Any:
						UpperCAmelCase_:		List[Any]       			=    """cpu"""  # ensure determinism for the device-dependent torch.Generator
						UpperCAmelCase_:		Tuple       			=    self.dummy_cond_unet
						UpperCAmelCase_:		Optional[int]       			=    PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Union[str, Any]       			=    self.dummy_vae
						UpperCAmelCase_:		Any       			=    self.dummy_text_encoder
						UpperCAmelCase_:		List[Any]       			=    CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)

						# make sure here that pndm scheduler skips prk
						UpperCAmelCase_:		List[str]       			=    StableDiffusionPipeline(
						    unet=SCREAMING_SNAKE_CASE_,      scheduler=SCREAMING_SNAKE_CASE_,      vae=SCREAMING_SNAKE_CASE_,      text_encoder=SCREAMING_SNAKE_CASE_,      tokenizer=SCREAMING_SNAKE_CASE_,      safety_checker=SCREAMING_SNAKE_CASE_,      feature_extractor=self.dummy_extractor,      )
						UpperCAmelCase_:		List[str]       			=    sd_pipe.to(SCREAMING_SNAKE_CASE_				)
						sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		int       			=    """A painting of a squirrel eating a burger"""
						UpperCAmelCase_:		Union[str, Any]       			=    torch.Generator(device=SCREAMING_SNAKE_CASE_				).manual_seed(0				)
						UpperCAmelCase_:		List[Any]       			=    sd_pipe([prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=6.0,      num_inference_steps=2,      output_type="""np"""				)

						UpperCAmelCase_:		Optional[int]       			=    output.images

						UpperCAmelCase_:		Any       			=    torch.Generator(device=SCREAMING_SNAKE_CASE_				).manual_seed(0				)
						UpperCAmelCase_:		Tuple       			=    sd_pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=6.0,      num_inference_steps=2,      output_type="""np""",      return_dict=SCREAMING_SNAKE_CASE_,      )[0]

						UpperCAmelCase_:		Optional[Any]       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		Tuple       			=    image_from_tuple[0, -3:, -3:, -1]

						assert image.shape == (1, 64, 64, 3)
						UpperCAmelCase_:		Tuple       			=    np.array([0.5_1_2_2, 0.5_7_1_2, 0.4_8_2_5, 0.5_0_5_3, 0.5_6_4_6, 0.4_7_6_9, 0.5_1_7_9, 0.4_8_9_4, 0.4_9_9_4]				)

						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     __snake_case					(self				)					->      List[str]:
						UpperCAmelCase_:		Dict       			=    StableDiffusionPipeline.from_pretrained(
						    """hf-internal-testing/tiny-stable-diffusion-lms-pipe""",      safety_checker=SCREAMING_SNAKE_CASE_				)
						assert isinstance(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_				)
						assert isinstance(pipe.scheduler,      SCREAMING_SNAKE_CASE_				)
						assert pipe.safety_checker is None

						UpperCAmelCase_:		Optional[Any]       			=    pipe("""example prompt""",      num_inference_steps=2				).images[0]
						assert image is not None

						# check that there's no error when saving a pipeline with one of the models being None
						with tempfile.TemporaryDirectory() as tmpdirname:
									pipe.save_pretrained(SCREAMING_SNAKE_CASE_				)
									UpperCAmelCase_:		Optional[int]       			=    StableDiffusionPipeline.from_pretrained(SCREAMING_SNAKE_CASE_				)

						# sanity check that the pipeline still works
						assert pipe.safety_checker is None
						UpperCAmelCase_:		Union[str, Any]       			=    pipe("""example prompt""",      num_inference_steps=2				).images[0]
						assert image is not None

			@unittest.skipIf(torch_device != """cuda""",      """This test requires a GPU"""				)
			def     __snake_case					(self				)					->      Optional[Any]:
						UpperCAmelCase_:		Dict       			=    self.dummy_cond_unet
						UpperCAmelCase_:		Any       			=    PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		List[str]       			=    self.dummy_vae
						UpperCAmelCase_:		Tuple       			=    self.dummy_text_encoder
						UpperCAmelCase_:		Dict       			=    CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)

						# put models in fp16
						UpperCAmelCase_:		Dict       			=    unet.half()
						UpperCAmelCase_:		Optional[Any]       			=    vae.half()
						UpperCAmelCase_:		Union[str, Any]       			=    bert.half()

						# make sure here that pndm scheduler skips prk
						UpperCAmelCase_:		str       			=    StableDiffusionPipeline(
						    unet=SCREAMING_SNAKE_CASE_,      scheduler=SCREAMING_SNAKE_CASE_,      vae=SCREAMING_SNAKE_CASE_,      text_encoder=SCREAMING_SNAKE_CASE_,      tokenizer=SCREAMING_SNAKE_CASE_,      safety_checker=SCREAMING_SNAKE_CASE_,      feature_extractor=self.dummy_extractor,      )
						UpperCAmelCase_:		Optional[Any]       			=    sd_pipe.to(SCREAMING_SNAKE_CASE_				)
						sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		str       			=    """A painting of a squirrel eating a burger"""
						UpperCAmelCase_:		Optional[Any]       			=    sd_pipe([prompt],      num_inference_steps=2,      output_type="""np"""				).images

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







@nightly
@require_torch_gpu
class 							_a							(      unittest.TestCase				):

			def     __snake_case					(self				)					->      Dict:
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()

			def     __snake_case					(self				)					->      str:
						UpperCAmelCase_:		Optional[int]       			=    """cuda"""
						UpperCAmelCase_:		List[str]       			=    StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""				)
						UpperCAmelCase_:		Union[str, Any]       			=    pipe.to(SCREAMING_SNAKE_CASE_				)
						pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		int       			=    """a photo of a cat"""
						UpperCAmelCase_:		Tuple       			=    {
						    """editing_prompt""": ["""sunglasses"""],
						    """reverse_editing_direction""": [False],
						    """edit_warmup_steps""": 10,
						    """edit_guidance_scale""": 6,
						    """edit_threshold""": 0.9_5,
						    """edit_momentum_scale""": 0.5,
						    """edit_mom_beta""": 0.6,
						}

						UpperCAmelCase_:		List[str]       			=    3
						UpperCAmelCase_:		str       			=    7

						# no sega enabled
						UpperCAmelCase_:		Tuple       			=    torch.Generator(SCREAMING_SNAKE_CASE_				)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Any       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      )

						UpperCAmelCase_:		Optional[Any]       			=    output.images
						UpperCAmelCase_:		List[Any]       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		Optional[int]       			=    [
						    0.3_4_6_7_3_1_1_3,
						    0.3_8_4_9_2_7_3_3,
						    0.3_7_5_9_7_3_5_2,
						    0.3_4_0_8_6_3_3_5,
						    0.3_5_6_5_0_7_4_8,
						    0.3_5_5_7_9_2_0_5,
						    0.3_3_8_4_7_6_3,
						    0.3_4_3_4_0_2_3_6,
						    0.3_5_7_3_2_7_1,
						]

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

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

						# with sega enabled
						# generator = torch.manual_seed(seed)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Optional[Any]       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      **SCREAMING_SNAKE_CASE_,      )

						UpperCAmelCase_:		Any       			=    output.images
						UpperCAmelCase_:		Dict       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		Optional[int]       			=    [
						    0.4_1_8_8_7_8_2_6,
						    0.3_7_7_2_8_7_6_6,
						    0.3_0_1_3_8_2_7_2,
						    0.4_1_4_1_6_3_3_5,
						    0.4_1_6_6_4_9_8_5,
						    0.3_6_2_8_3_3_9_2,
						    0.3_6_1_9_1_2_4_6,
						    0.4_3_3_6_4_4_6_5,
						    0.4_3_0_0_1_7_3_2,
						]

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

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

			def     __snake_case					(self				)					->      Optional[Any]:
						UpperCAmelCase_:		List[Any]       			=    """cuda"""
						UpperCAmelCase_:		Tuple       			=    StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""				)
						UpperCAmelCase_:		List[Any]       			=    pipe.to(SCREAMING_SNAKE_CASE_				)
						pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		Dict       			=    """an image of a crowded boulevard, realistic, 4k"""
						UpperCAmelCase_:		Optional[Any]       			=    {
						    """editing_prompt""": """crowd, crowded, people""",
						    """reverse_editing_direction""": True,
						    """edit_warmup_steps""": 10,
						    """edit_guidance_scale""": 8.3,
						    """edit_threshold""": 0.9,
						    """edit_momentum_scale""": 0.5,
						    """edit_mom_beta""": 0.6,
						}

						UpperCAmelCase_:		Dict       			=    9
						UpperCAmelCase_:		Optional[int]       			=    7

						# no sega enabled
						UpperCAmelCase_:		List[Any]       			=    torch.Generator(SCREAMING_SNAKE_CASE_				)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		str       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      )

						UpperCAmelCase_:		Optional[Any]       			=    output.images
						UpperCAmelCase_:		List[Any]       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		int       			=    [
						    0.4_3_4_9_7_9_9_8,
						    0.9_1_8_1_4_0_6_5,
						    0.7_5_4_0_7_3_9,
						    0.5_5_5_8_0_2_0_5,
						    0.8_4_6_7_2_6_5,
						    0.5_3_8_9_6_9_1,
						    0.6_2_5_7_4_5_0_6,
						    0.5_8_8_9_7_7_6_3,
						    0.5_0_9_2_6_7_5_7,
						]

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

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

						# with sega enabled
						# generator = torch.manual_seed(seed)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Optional[int]       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      **SCREAMING_SNAKE_CASE_,      )

						UpperCAmelCase_:		int       			=    output.images
						UpperCAmelCase_:		Any       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		str       			=    [
						    0.3_0_8_9_7_1_9,
						    0.3_0_5_0_0_1_4_4,
						    0.2_9_0_1_6_0_4_2,
						    0.3_0_6_3_0_9_6_4,
						    0.3_2_5_6_8_7,
						    0.2_9_4_1_9_2_2_5,
						    0.2_9_0_8_0_9_1,
						    0.2_8_7_2_3_5_9_8,
						    0.2_7_6_9_6_2_9_4,
						]

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

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

			def     __snake_case					(self				)					->      Tuple:
						UpperCAmelCase_:		Union[str, Any]       			=    """cuda"""
						UpperCAmelCase_:		Optional[Any]       			=    StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""				)
						UpperCAmelCase_:		Optional[Any]       			=    pipe.to(SCREAMING_SNAKE_CASE_				)
						pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		Tuple       			=    """a castle next to a river"""
						UpperCAmelCase_:		Optional[Any]       			=    {
						    """editing_prompt""": ["""boat on a river, boat""", """monet, impression, sunrise"""],
						    """reverse_editing_direction""": False,
						    """edit_warmup_steps""": [15, 18],
						    """edit_guidance_scale""": 6,
						    """edit_threshold""": [0.9, 0.8],
						    """edit_momentum_scale""": 0.5,
						    """edit_mom_beta""": 0.6,
						}

						UpperCAmelCase_:		str       			=    48
						UpperCAmelCase_:		Optional[int]       			=    7

						# no sega enabled
						UpperCAmelCase_:		int       			=    torch.Generator(SCREAMING_SNAKE_CASE_				)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Union[str, Any]       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      )

						UpperCAmelCase_:		Tuple       			=    output.images
						UpperCAmelCase_:		Tuple       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		Tuple       			=    [
						    0.7_5_1_6_3_5_5_5,
						    0.7_6_0_3_7_1_4_5,
						    0.6_1_7_8_5,
						    0.9_1_8_9_6_7_3,
						    0.8_6_2_7_7_0_1,
						    0.8_5_1_8_9_6_9_4,
						    0.8_5_1_2_8_1_3,
						    0.8_7_0_1_2_0_7_6,
						    0.8_3_1_2_8_5_7,
						]

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

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

						# with sega enabled
						# generator = torch.manual_seed(seed)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Dict       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      **SCREAMING_SNAKE_CASE_,      )

						UpperCAmelCase_:		List[str]       			=    output.images
						UpperCAmelCase_:		Optional[int]       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		Optional[Any]       			=    [
						    0.7_3_5_5_3_3_6_5,
						    0.7_5_3_7_2_7_1,
						    0.7_4_3_4_1_9_0_5,
						    0.6_6_4_8_0_3_5_6,
						    0.6_4_7_2_9_2_5,
						    0.6_3_0_3_9_4_1_6,
						    0.6_4_8_1_2_9_0_5,
						    0.6_7_4_9_7_1_7,
						    0.6_5_1_7_1_0_2,
						]

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

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

			def     __snake_case					(self				)					->      Optional[int]:
						UpperCAmelCase_:		List[Any]       			=    """cuda"""
						UpperCAmelCase_:		Union[str, Any]       			=    StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5""",      torch_dtype=torch.floataa				)
						UpperCAmelCase_:		Optional[int]       			=    pipe.to(SCREAMING_SNAKE_CASE_				)
						pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		Dict       			=    """a photo of a cat"""
						UpperCAmelCase_:		Dict       			=    {
						    """editing_prompt""": ["""sunglasses"""],
						    """reverse_editing_direction""": [False],
						    """edit_warmup_steps""": 10,
						    """edit_guidance_scale""": 6,
						    """edit_threshold""": 0.9_5,
						    """edit_momentum_scale""": 0.5,
						    """edit_mom_beta""": 0.6,
						}

						UpperCAmelCase_:		List[Any]       			=    3
						UpperCAmelCase_:		Dict       			=    7

						# no sega enabled
						UpperCAmelCase_:		Tuple       			=    torch.Generator(SCREAMING_SNAKE_CASE_				)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		List[Any]       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      )

						UpperCAmelCase_:		str       			=    output.images
						UpperCAmelCase_:		Optional[Any]       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		List[Any]       			=    [
						    0.3_4_8_8_7_6_9_5,
						    0.3_8_7_6_9_5_3,
						    0.3_7_5,
						    0.3_4_4_2_3_8_2_8,
						    0.3_5_8_1_5_4_3,
						    0.3_5_7_1_7_7_7_3,
						    0.3_3_8_3_7_8_9,
						    0.3_4_5_7_0_3_1_2,
						    0.3_5_9_3_7_5,
						]

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

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

						# with sega enabled
						# generator = torch.manual_seed(seed)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Any       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      **SCREAMING_SNAKE_CASE_,      )

						UpperCAmelCase_:		List[Any]       			=    output.images
						UpperCAmelCase_:		List[str]       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		List[Any]       			=    [
						    0.4_2_2_8_5_1_5_6,
						    0.3_6_9_1_4_0_6_2,
						    0.2_9_0_7_7_1_4_8,
						    0.4_2_0_4_1_0_1_6,
						    0.4_1_9_1_8_9_4_5,
						    0.3_5_4_9_8_0_4_7,
						    0.3_6_1_8_1_6_4,
						    0.4_4_2_3_8_2_8,
						    0.4_3_1_1_5_2_3_4,
						]

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

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






