
'''simple docstring'''







from __future__ import annotations

import unittest

from transformers import AutoTokenizer, GPTJConfig, is_tf_available
from transformers.testing_utils import require_tf, slow, tooslow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin
from ...utils.test_modeling_tf_core import TFCoreModelTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers.models.gptj.modeling_tf_gptj import (
        TFGPTJForCausalLM,
        TFGPTJForQuestionAnswering,
        TFGPTJForSequenceClassification,
        TFGPTJModel,
        shape_list,
    )







class 		snake_case_       :



   """simple docstring"""





   def __init__( self			,   UpperCamelCase):
          lowerCamelCase__				   =      parent
          lowerCamelCase__				   =      13
          lowerCamelCase__				   =      7
          lowerCamelCase__				   =      True
          lowerCamelCase__				   =      True
          lowerCamelCase__				   =      True
          lowerCamelCase__				   =      True
          lowerCamelCase__				   =      True
          lowerCamelCase__				   =      99
          lowerCamelCase__				   =      32
          lowerCamelCase__				   =      4
          lowerCamelCase__				   =      2
          lowerCamelCase__				   =      4
          lowerCamelCase__				   =      37
          lowerCamelCase__				   =      "gelu"
          lowerCamelCase__				   =      0.1
          lowerCamelCase__				   =      0.1
          lowerCamelCase__				   =      5_12
          lowerCamelCase__				   =      16
          lowerCamelCase__				   =      2
          lowerCamelCase__				   =      0.0_2
          lowerCamelCase__				   =      3
          lowerCamelCase__				   =      4
          lowerCamelCase__				   =      None
          lowerCamelCase__				   =      self.vocab_size - 1
          lowerCamelCase__				   =      self.vocab_size - 1
          lowerCamelCase__				   =      self.vocab_size - 1







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      ids_tensor([self.batch_size, self.seq_length]			,   self.vocab_size)

          lowerCamelCase__				   =      None
          if self.use_input_mask:
                 lowerCamelCase__				   =      random_attention_mask([self.batch_size, self.seq_length])

          lowerCamelCase__				   =      None
          if self.use_token_type_ids:
                 lowerCamelCase__				   =      ids_tensor([self.batch_size, self.seq_length]			,   self.type_vocab_size)

          lowerCamelCase__				   =      None
          if self.use_mc_token_ids:
                 lowerCamelCase__				   =      ids_tensor([self.batch_size, self.num_choices]			,   self.seq_length)

          lowerCamelCase__				   =      None
          lowerCamelCase__				   =      None
          lowerCamelCase__				   =      None
          if self.use_labels:
                 lowerCamelCase__				   =      ids_tensor([self.batch_size]			,   self.type_sequence_label_size)
                 lowerCamelCase__				   =      ids_tensor([self.batch_size, self.seq_length]			,   self.num_labels)
                 lowerCamelCase__				   =      ids_tensor([self.batch_size]			,   self.num_choices)

          lowerCamelCase__				   =      GPTJConfig(
              vocab_size=self.vocab_size			,   n_embd=self.hidden_size			,   n_layer=self.num_hidden_layers			,   n_head=self.num_attention_heads			,   intermediate_size=self.intermediate_size			,   hidden_act=self.hidden_act			,   hidden_dropout_prob=self.hidden_dropout_prob			,   attention_probs_dropout_prob=self.attention_probs_dropout_prob			,   n_positions=self.max_position_embeddings			,   type_vocab_size=self.type_vocab_size			,   initializer_range=self.initializer_range			,   bos_token_id=self.bos_token_id			,   eos_token_id=self.eos_token_id			,   pad_token_id=self.pad_token_id			,   rotary_dim=self.rotary_dim			,   return_dict=UpperCamelCase			,   )

          lowerCamelCase__				   =      ids_tensor([self.num_hidden_layers, self.num_attention_heads]			,   2)

          return (
              config,
              input_ids,
              input_mask,
              head_mask,
              token_type_ids,
              mc_token_ids,
              sequence_labels,
              token_labels,
              choice_labels,
          )







   def     __UpperCAmelCase      ( self			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   *UpperCamelCase):
          lowerCamelCase__				   =      TFGPTJModel(config=UpperCamelCase)
          lowerCamelCase__				   =      {
              "input_ids": input_ids,
              "attention_mask": input_mask,
              "token_type_ids": token_type_ids,
          }
          lowerCamelCase__				   =      model(UpperCamelCase)

          lowerCamelCase__				   =      [input_ids, None, input_mask]  # None is the input for 'past'
          lowerCamelCase__				   =      model(UpperCamelCase)

          lowerCamelCase__				   =      model(UpperCamelCase)

          self.parent.assertEqual(result.last_hidden_state.shape			,   (self.batch_size, self.seq_length, self.hidden_size))







   def     __UpperCAmelCase      ( self			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   *UpperCamelCase):
          lowerCamelCase__				   =      TFGPTJModel(config=UpperCamelCase)

          # first forward pass
          lowerCamelCase__				   =      model(UpperCamelCase			,   token_type_ids=UpperCamelCase			,   use_cache=UpperCamelCase)
          lowerCamelCase__				   =      model(UpperCamelCase			,   token_type_ids=UpperCamelCase)
          lowerCamelCase__				   =      model(UpperCamelCase			,   token_type_ids=UpperCamelCase			,   use_cache=UpperCamelCase)

          self.parent.assertTrue(len(UpperCamelCase) == len(UpperCamelCase))
          self.parent.assertTrue(len(UpperCamelCase) == len(UpperCamelCase) + 1)

          lowerCamelCase__   ,		lowerCamelCase__				   =      outputs.to_tuple()

          # create hypothetical next token and extent to next_input_ids
          lowerCamelCase__				   =      ids_tensor((self.batch_size, 1)			,   config.vocab_size)
          lowerCamelCase__				   =      ids_tensor([self.batch_size, 1]			,   self.type_vocab_size)

          # append to next input_ids and token_type_ids
          lowerCamelCase__				   =      tf.concat([input_ids, next_tokens]			,   axis=-1)
          lowerCamelCase__				   =      tf.concat([token_type_ids, next_token_types]			,   axis=-1)

          lowerCamelCase__				   =      model(UpperCamelCase			,   token_type_ids=UpperCamelCase)["last_hidden_state"]
          lowerCamelCase__				   =      model(UpperCamelCase			,   token_type_ids=UpperCamelCase			,   past_key_values=UpperCamelCase)[
              "last_hidden_state"
          ]

          # select random slice
          lowerCamelCase__				   =      int(ids_tensor((1,)			,   shape_list(UpperCamelCase)[-1]))
          lowerCamelCase__				   =      output_from_no_past[:, -1, random_slice_idx]
          lowerCamelCase__				   =      output_from_past[:, 0, random_slice_idx]

          # test that outputs are equal for slice
          tf.debugging.assert_near(UpperCamelCase			,   UpperCamelCase			,   rtol=1E-6)







   def     __UpperCAmelCase      ( self			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   *UpperCamelCase):
          lowerCamelCase__				   =      TFGPTJModel(config=UpperCamelCase)

          # create attention mask
          lowerCamelCase__				   =      self.seq_length // 2
          lowerCamelCase__				   =      tf.ones((self.batch_size, half_seq_length)			,   dtype=tf.intaa)
          lowerCamelCase__				   =      tf.zeros((self.batch_size, self.seq_length - half_seq_length)			,   dtype=tf.intaa)
          lowerCamelCase__				   =      tf.concat([attn_mask_begin, attn_mask_end]			,   axis=1)

          # first forward pass
          lowerCamelCase__   ,		lowerCamelCase__				   =      model(UpperCamelCase			,   attention_mask=UpperCamelCase).to_tuple()

          # create hypothetical next token and extent to next_input_ids
          lowerCamelCase__				   =      ids_tensor((self.batch_size, 1)			,   config.vocab_size)

          # change a random masked slice from input_ids
          lowerCamelCase__				   =      ids_tensor((1,)			,   UpperCamelCase).numpy() + 1
          lowerCamelCase__				   =      ids_tensor((self.batch_size, self.seq_length)			,   config.vocab_size)
          lowerCamelCase__				   =      tf.range(self.seq_length) == (self.seq_length - random_seq_idx_to_change)
          lowerCamelCase__				   =      tf.transpose(
              tf.broadcast_to(tf.expand_dims(UpperCamelCase			,   -1)			,   (self.seq_length, self.batch_size)))
          lowerCamelCase__				   =      tf.where(UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase)

          # append to next input_ids and attn_mask
          lowerCamelCase__				   =      tf.concat([input_ids, next_tokens]			,   axis=-1)
          lowerCamelCase__				   =      tf.concat([attn_mask, tf.ones((shape_list(UpperCamelCase)[0], 1)			,   dtype=tf.intaa)]			,   axis=1)

          # get two different outputs
          lowerCamelCase__				   =      model(UpperCamelCase			,   attention_mask=UpperCamelCase)["last_hidden_state"]
          lowerCamelCase__				   =      model(UpperCamelCase			,   past_key_values=UpperCamelCase			,   attention_mask=UpperCamelCase)[
              "last_hidden_state"
          ]

          # select random slice
          lowerCamelCase__				   =      int(ids_tensor((1,)			,   shape_list(UpperCamelCase)[-1]))
          lowerCamelCase__				   =      output_from_no_past[:, -1, random_slice_idx]
          lowerCamelCase__				   =      output_from_past[:, 0, random_slice_idx]

          # test that outputs are equal for slice
          tf.debugging.assert_near(UpperCamelCase			,   UpperCamelCase			,   rtol=1E-1_2)







   def     __UpperCAmelCase      ( self			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   *UpperCamelCase):
          lowerCamelCase__				   =      TFGPTJModel(config=UpperCamelCase)

          lowerCamelCase__				   =      input_ids[:1, :]
          lowerCamelCase__				   =      input_mask[:1, :]
          lowerCamelCase__				   =      token_type_ids[:1, :]
          lowerCamelCase__				   =      1

          # first forward pass
          lowerCamelCase__				   =      model(UpperCamelCase			,   attention_mask=UpperCamelCase			,   token_type_ids=UpperCamelCase			,   use_cache=UpperCamelCase)

          lowerCamelCase__   ,		lowerCamelCase__				   =      outputs.to_tuple()

          # create hypothetical next token and extent to next_input_ids
          lowerCamelCase__				   =      ids_tensor((self.batch_size, 3)			,   config.vocab_size)
          lowerCamelCase__				   =      ids_tensor((self.batch_size, 3)			,   2)
          lowerCamelCase__				   =      ids_tensor((self.batch_size, 3)			,   self.type_vocab_size)

          # append to next input_ids and token_type_ids
          lowerCamelCase__				   =      tf.concat([input_ids, next_tokens]			,   axis=-1)
          lowerCamelCase__				   =      tf.concat([input_mask, next_attn_mask]			,   axis=-1)
          lowerCamelCase__				   =      tf.concat([token_type_ids, next_token_types]			,   axis=-1)

          lowerCamelCase__				   =      model(
              UpperCamelCase			,   token_type_ids=UpperCamelCase			,   attention_mask=UpperCamelCase)["last_hidden_state"]
          lowerCamelCase__				   =      model(
              UpperCamelCase			,   token_type_ids=UpperCamelCase			,   attention_mask=UpperCamelCase			,   past_key_values=UpperCamelCase			,   )["last_hidden_state"]
          self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1])

          # select random slice
          lowerCamelCase__				   =      int(ids_tensor((1,)			,   shape_list(UpperCamelCase)[-1]))
          lowerCamelCase__				   =      output_from_no_past[:, -3:, random_slice_idx]
          lowerCamelCase__				   =      output_from_past[:, :, random_slice_idx]

          # test that outputs are equal for slice
          tf.debugging.assert_near(UpperCamelCase			,   UpperCamelCase			,   rtol=1E-3)







   def     __UpperCAmelCase      ( self			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   *UpperCamelCase):
          lowerCamelCase__				   =      TFGPTJForCausalLM(config=UpperCamelCase)
          lowerCamelCase__				   =      {
              "input_ids": input_ids,
              "attention_mask": input_mask,
              "token_type_ids": token_type_ids,
          }
          lowerCamelCase__				   =      model(UpperCamelCase)
          self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.seq_length, self.vocab_size))







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.prepare_config_and_inputs()

          (
              (
              lowerCamelCase__
          )   ,		(
              lowerCamelCase__
          )   ,		(
              lowerCamelCase__
          )   ,		(
              lowerCamelCase__
          )   ,		(
              lowerCamelCase__
          )   ,		(
              lowerCamelCase__
          )   ,		(
              lowerCamelCase__
          )   ,		(
              lowerCamelCase__
          )   ,		(
              lowerCamelCase__
          )   ,		
          )				   =      config_and_inputs

          lowerCamelCase__				   =      {
              "input_ids": input_ids,
              "token_type_ids": token_type_ids,
              "attention_mask": input_mask,
          }
          return config, inputs_dict







@require_tf
class 		snake_case_       (						A__		,	A__		,	A__		,	unittest.TestCase						):



   """simple docstring"""





   __lowerCAmelCase	:							Union[str, Any]								=(
       (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
       if is_tf_available()
       else ()
   )

   __lowerCAmelCase	:							Dict								=(TFGPTJForCausalLM,) if is_tf_available() else ()
   __lowerCAmelCase	:							Optional[Any]								=(
       {
           '''feature-extraction''': TFGPTJModel,
           '''question-answering''': TFGPTJForQuestionAnswering,
           '''text-classification''': TFGPTJForSequenceClassification,
           '''text-generation''': TFGPTJForCausalLM,
           '''zero-shot''': TFGPTJForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   __lowerCAmelCase	:							Dict								=False
   __lowerCAmelCase	:							List[Any]								=False
   __lowerCAmelCase	:							int								=False
   __lowerCAmelCase	:							Tuple								=False







   def     __UpperCAmelCase      ( self			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase):
          if (
              pipeline_test_casse_name == "QAPipelineTests"
              and tokenizer_name is not None
              and not tokenizer_name.endswith("Fast")
          ):
                 # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
                 # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
                 # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
                 return True

          return False







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      TFGPTJModelTester(self)
          lowerCamelCase__				   =      ConfigTester(self			,   config_class=UpperCamelCase			,   n_embd=37)







   def     __UpperCAmelCase      ( self):
          self.config_tester.run_common_tests()







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_gptj_model(*UpperCamelCase)







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_gptj_model_past(*UpperCamelCase)







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_gptj_model_attention_mask_past(*UpperCamelCase)







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_gptj_model_past_large_inputs(*UpperCamelCase)







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_gptj_lm_head_model(*UpperCamelCase)







   @slow
   @unittest.skipIf(
       not is_tf_available() or len(tf.config.list_physical_devices("GPU")) > 0			,   "skip testing on GPU for now to avoid GPU OOM."			,   )
   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      TFGPTJModel.from_pretrained("EleutherAI/gpt-j-6B"			,   from_pt=UpperCamelCase)
          self.assertIsNotNone(UpperCamelCase)







   @unittest.skip(reason="Currently, model embeddings are going to undergo a major refactor.")
   def     __UpperCAmelCase      ( self):
          super().test_resize_token_embeddings()







@require_tf
@tooslow
# Marked as @tooslow due to GPU OOM -- but still useful to run locally. Requires ~39GB of RAM.
class 		snake_case_       (						unittest.TestCase						):



   """simple docstring"""





   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"			,   from_pt=UpperCamelCase)
          lowerCamelCase__				   =      tf.convert_to_tensor([[4_64, 32_90]]			,   dtype=tf.intaa)  # The dog
          # fmt: off
          # The dog is a man's best friend. It is a loyal companion, and it is a friend
          lowerCamelCase__				   =      [4_64, 32_90, 3_18, 2_57, 5_82, 3_38, 12_66, 15_45, 13, 6_32, 3_18, 2_57, 91_12, 1_51_85, 11, 2_90, 3_40, 3_18, 2_57, 15_45]
          # fmt: on
          lowerCamelCase__				   =      model.generate(UpperCamelCase			,   do_sample=UpperCamelCase)
          self.assertListEqual(output_ids[0].numpy().tolist()			,   UpperCamelCase)







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"			,   revision="float16")
          lowerCamelCase__				   =      TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"			,   revision="float16"			,   from_pt=UpperCamelCase)

          lowerCamelCase__				   =      tokenizer("Today is a nice day and"			,   return_tensors="tf")
          # forces the generation to happen on CPU, to avoid GPU-related quirks
          with tf.device(":/CPU:0"):
                 lowerCamelCase__				   =      model.generate(**UpperCamelCase			,   do_sample=UpperCamelCase			,   seed=[42, 0])
          lowerCamelCase__				   =      tokenizer.decode(output_ids[0]			,   skip_special_tokens=UpperCamelCase)

          lowerCamelCase__				   =      "Today is a nice day and I’m going to go for a walk. I’"
          self.assertEqual(UpperCamelCase			,   UpperCamelCase)







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"			,   revision="float16"			,   from_pt=UpperCamelCase)
          lowerCamelCase__				   =      AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"			,   revision="float16")

          lowerCamelCase__				   =      "left"

          # Define PAD Token = EOS Token = 50256
          lowerCamelCase__				   =      tokenizer.eos_token
          lowerCamelCase__				   =      model.config.eos_token_id

          # use different length sentences to test batching
          lowerCamelCase__				   =      [
              "Hello, my dog is a little",
              "Today, I",
          ]
          lowerCamelCase__				   =      [
              "Hello, my dog is a little over a year old and has been diagnosed with hip dysplasia",
              "Today, I’m going to be talking about a topic that’",
          ]
          return model, tokenizer, sentences, expected_output_sentences







   def     __UpperCAmelCase      ( self):
          # Confirms that we get the expected results with left-padded beam search
          lowerCamelCase__   ,		lowerCamelCase__   ,		lowerCamelCase__   ,		lowerCamelCase__				   =      self._get_beam_search_test_objects()

          lowerCamelCase__				   =      tokenizer(UpperCamelCase			,   return_tensors="tf"			,   padding=UpperCamelCase)
          lowerCamelCase__				   =      model.generate(**UpperCamelCase			,   do_sample=UpperCamelCase			,   num_beams=2)
          lowerCamelCase__				   =      tokenizer.batch_decode(UpperCamelCase			,   skip_special_tokens=UpperCamelCase)
          self.assertListEqual(UpperCamelCase			,   UpperCamelCase)







   def     __UpperCAmelCase      ( self):
          # Confirms that left-padding is working properly
          lowerCamelCase__   ,		lowerCamelCase__   ,		lowerCamelCase__   ,		lowerCamelCase__				   =      self._get_beam_search_test_objects()

          lowerCamelCase__				   =      tokenizer(UpperCamelCase			,   return_tensors="tf"			,   padding=UpperCamelCase)
          lowerCamelCase__				   =      tokenizer(sentences[0]			,   return_tensors="tf")
          lowerCamelCase__				   =      model.generate(**UpperCamelCase			,   do_sample=UpperCamelCase			,   num_beams=2)
          lowerCamelCase__				   =      (
              shape_list(inputs_non_padded["input_ids"])[-1]
              - tf.reduce_sum(tf.cast(inputs["attention_mask"][-1]			,   tf.intaa)).numpy()
          )
          lowerCamelCase__				   =      tokenizer(sentences[1]			,   return_tensors="tf")
          lowerCamelCase__				   =      model.generate(
              **UpperCamelCase			,   do_sample=UpperCamelCase			,   num_beams=2			,   max_length=model.config.max_length - num_paddings)
          lowerCamelCase__				   =      tokenizer.decode(output_non_padded[0]			,   skip_special_tokens=UpperCamelCase)
          lowerCamelCase__				   =      tokenizer.decode(output_padded[0]			,   skip_special_tokens=UpperCamelCase)
          self.assertListEqual(UpperCamelCase			,   [non_padded_sentence, padded_sentence])







   def     __UpperCAmelCase      ( self):
          # Confirms that XLA is working properly
          lowerCamelCase__   ,		lowerCamelCase__   ,		lowerCamelCase__   ,		lowerCamelCase__				   =      self._get_beam_search_test_objects()

          lowerCamelCase__				   =      tokenizer(UpperCamelCase			,   return_tensors="tf"			,   padding=UpperCamelCase)
          lowerCamelCase__				   =      tf.function(model.generate			,   jit_compile=UpperCamelCase)
          lowerCamelCase__				   =      xla_generate(**UpperCamelCase			,   do_sample=UpperCamelCase			,   num_beams=2)
          lowerCamelCase__				   =      tokenizer.batch_decode(UpperCamelCase			,   skip_special_tokens=UpperCamelCase)
          self.assertListEqual(UpperCamelCase			,   UpperCamelCase)



'''simple docstring'''







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 		snake_case_       (						unittest.TestCase						):



   """simple docstring"""





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







   @property
   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      1
          lowerCamelCase__				   =      3
          lowerCamelCase__				   =      (32, 32)

          lowerCamelCase__				   =      floats_tensor((batch_size, num_channels) + sizes			,   rng=random.Random(0)).to(UpperCamelCase)
          return image







   @property
   def     __UpperCAmelCase      ( self):
          torch.manual_seed(0)
          lowerCamelCase__				   =      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     __UpperCAmelCase      ( self):
          torch.manual_seed(0)
          lowerCamelCase__				   =      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     __UpperCAmelCase      ( self):
          torch.manual_seed(0)
          lowerCamelCase__				   =      CLIPTextConfig(
              bos_token_id=0			,   eos_token_id=2			,   hidden_size=32			,   intermediate_size=37			,   layer_norm_eps=1E-0_5			,   num_attention_heads=4			,   num_hidden_layers=5			,   pad_token_id=1			,   vocab_size=10_00			,   )
          return CLIPTextModel(UpperCamelCase)







   @property
   def     __UpperCAmelCase      ( self):
          def extract(*UpperCamelCase			,   **UpperCamelCase):







                 class 		snake_case_       :



                    """simple docstring"""





                    def __init__( self):
                           lowerCamelCase__				   =      torch.ones([0])







                    def     __UpperCAmelCase      ( self			,   UpperCamelCase):
                           self.pixel_values.to(UpperCamelCase)
                           return self






                 return Out()

          return extract







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cpu"  # ensure determinism for the device-dependent torch.Generator
          lowerCamelCase__				   =      self.dummy_cond_unet
          lowerCamelCase__				   =      DDIMScheduler(
              beta_start=0.0_0_0_8_5			,   beta_end=0.0_1_2			,   beta_schedule="scaled_linear"			,   clip_sample=UpperCamelCase			,   set_alpha_to_one=UpperCamelCase			,   )

          lowerCamelCase__				   =      self.dummy_vae
          lowerCamelCase__				   =      self.dummy_text_encoder
          lowerCamelCase__				   =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

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

          lowerCamelCase__				   =      "A painting of a squirrel eating a burger"

          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      sd_pipe([prompt]			,   generator=UpperCamelCase			,   guidance_scale=6.0			,   num_inference_steps=2			,   output_type="np")
          lowerCamelCase__				   =      output.images

          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      sd_pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=6.0			,   num_inference_steps=2			,   output_type="np"			,   return_dict=UpperCamelCase			,   )[0]

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

          assert image.shape == (1, 64, 64, 3)
          lowerCamelCase__				   =      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     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cpu"  # ensure determinism for the device-dependent torch.Generator
          lowerCamelCase__				   =      self.dummy_cond_unet
          lowerCamelCase__				   =      PNDMScheduler(skip_prk_steps=UpperCamelCase)
          lowerCamelCase__				   =      self.dummy_vae
          lowerCamelCase__				   =      self.dummy_text_encoder
          lowerCamelCase__				   =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

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

          lowerCamelCase__				   =      "A painting of a squirrel eating a burger"
          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      sd_pipe([prompt]			,   generator=UpperCamelCase			,   guidance_scale=6.0			,   num_inference_steps=2			,   output_type="np")

          lowerCamelCase__				   =      output.images

          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      sd_pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=6.0			,   num_inference_steps=2			,   output_type="np"			,   return_dict=UpperCamelCase			,   )[0]

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

          assert image.shape == (1, 64, 64, 3)
          lowerCamelCase__				   =      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     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained(
              "hf-internal-testing/tiny-stable-diffusion-lms-pipe"			,   safety_checker=UpperCamelCase)
          assert isinstance(UpperCamelCase			,   UpperCamelCase)
          assert isinstance(pipe.scheduler			,   UpperCamelCase)
          assert pipe.safety_checker is None

          lowerCamelCase__				   =      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(UpperCamelCase)
                 lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained(UpperCamelCase)

          # sanity check that the pipeline still works
          assert pipe.safety_checker is None
          lowerCamelCase__				   =      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     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.dummy_cond_unet
          lowerCamelCase__				   =      PNDMScheduler(skip_prk_steps=UpperCamelCase)
          lowerCamelCase__				   =      self.dummy_vae
          lowerCamelCase__				   =      self.dummy_text_encoder
          lowerCamelCase__				   =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

          # put models in fp16
          lowerCamelCase__				   =      unet.half()
          lowerCamelCase__				   =      vae.half()
          lowerCamelCase__				   =      bert.half()

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

          lowerCamelCase__				   =      "A painting of a squirrel eating a burger"
          lowerCamelCase__				   =      sd_pipe([prompt]			,   num_inference_steps=2			,   output_type="np").images

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







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



   """simple docstring"""





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







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cuda"
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)
          pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "a photo of a cat"
          lowerCamelCase__				   =      {
              "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,
          }

          lowerCamelCase__				   =      3
          lowerCamelCase__				   =      7

          # no sega enabled
          lowerCamelCase__				   =      torch.Generator(UpperCamelCase)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              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, 5_12, 5_12, 3)

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   **UpperCamelCase			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              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, 5_12, 5_12, 3)

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







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cuda"
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)
          pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "an image of a crowded boulevard, realistic, 4k"
          lowerCamelCase__				   =      {
              "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,
          }

          lowerCamelCase__				   =      9
          lowerCamelCase__				   =      7

          # no sega enabled
          lowerCamelCase__				   =      torch.Generator(UpperCamelCase)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              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, 5_12, 5_12, 3)

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   **UpperCamelCase			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              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, 5_12, 5_12, 3)

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







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cuda"
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)
          pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "a castle next to a river"
          lowerCamelCase__				   =      {
              "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,
          }

          lowerCamelCase__				   =      48
          lowerCamelCase__				   =      7

          # no sega enabled
          lowerCamelCase__				   =      torch.Generator(UpperCamelCase)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              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, 5_12, 5_12, 3)

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   **UpperCamelCase			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              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, 5_12, 5_12, 3)

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







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cuda"
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"			,   torch_dtype=torch.floataa)
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)
          pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "a photo of a cat"
          lowerCamelCase__				   =      {
              "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,
          }

          lowerCamelCase__				   =      3
          lowerCamelCase__				   =      7

          # no sega enabled
          lowerCamelCase__				   =      torch.Generator(UpperCamelCase)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              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, 5_12, 5_12, 3)

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   **UpperCamelCase			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              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, 5_12, 5_12, 3)

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

