


import gc
import unittest

from transformers import CTRLConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
      CTRLForSequenceClassification,
      CTRLLMHeadModel,
      CTRLModel,
  )







class   __lowerCAmelCase			:





   def __init__( self					,			snake_case					,			snake_case=14					,			snake_case=7					,			snake_case=True					,			snake_case=True					,			snake_case=True					,			snake_case=True					,			snake_case=True					,			snake_case=99					,			snake_case=32					,			snake_case=5					,			snake_case=4					,			snake_case=37					,			snake_case="gelu"					,			snake_case=0.1					,			snake_case=0.1					,			snake_case=512					,			snake_case=16					,			snake_case=2					,			snake_case=0.02					,			snake_case=3					,			snake_case=4					,			snake_case=None					,			)				->  Dict:
      """simple docstring"""





      a__	:						Any      		=  parent
      a__	:						Union[str, Any]      		=  batch_size
      a__	:						Union[str, Any]      		=  seq_length
      a__	:						List[Any]      		=  is_training
      a__	:						Tuple      		=  use_token_type_ids
      a__	:						Optional[Any]      		=  use_input_mask
      a__	:						Optional[Any]      		=  use_labels
      a__	:						List[str]      		=  use_mc_token_ids
      a__	:						List[Any]      		=  vocab_size
      a__	:						Tuple      		=  hidden_size
      a__	:						str      		=  num_hidden_layers
      a__	:						Optional[int]      		=  num_attention_heads
      a__	:						Any      		=  intermediate_size
      a__	:						List[Any]      		=  hidden_act
      a__	:						List[Any]      		=  hidden_dropout_prob
      a__	:						int      		=  attention_probs_dropout_prob
      a__	:						Union[str, Any]      		=  max_position_embeddings
      a__	:						Union[str, Any]      		=  type_vocab_size
      a__	:						List[str]      		=  type_sequence_label_size
      a__	:						Any      		=  initializer_range
      a__	:						List[Any]      		=  num_labels
      a__	:						Optional[Any]      		=  num_choices
      a__	:						List[Any]      		=  scope
      a__	:						List[str]      		=  self.vocab_size - 1





   def 	_snake_case ( self      )				->  int:
      """simple docstring"""





      a__	:						Tuple      		=  ids_tensor([self.batch_size, self.seq_length]					,			self.vocab_size      )

      a__	:						int      		=  None
      if self.use_input_mask:
         a__	:						Optional[Any]      		=  random_attention_mask([self.batch_size, self.seq_length]      )

      a__	:						int      		=  None
      if self.use_token_type_ids:
         a__	:						List[str]      		=  ids_tensor([self.batch_size, self.seq_length]					,			self.type_vocab_size      )

      a__	:						List[str]      		=  None
      if self.use_mc_token_ids:
         a__	:						List[Any]      		=  ids_tensor([self.batch_size, self.num_choices]					,			self.seq_length      )

      a__	:						int      		=  None
      a__	:						str      		=  None
      a__	:						str      		=  None
      if self.use_labels:
         a__	:						Dict      		=  ids_tensor([self.batch_size]					,			self.type_sequence_label_size      )
         a__	:						Optional[Any]      		=  ids_tensor([self.batch_size, self.seq_length]					,			self.num_labels      )
         a__	:						Optional[int]      		=  ids_tensor([self.batch_size]					,			self.num_choices      )

      a__	:						Optional[Any]      		=  self.get_config()

      a__	:						List[str]      		=  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 	_snake_case ( self      )				->  Union[str, Any]:
      """simple docstring"""





      return CTRLConfig(
          vocab_size=self.vocab_size					,			n_embd=self.hidden_size					,			n_layer=self.num_hidden_layers					,			n_head=self.num_attention_heads					,			n_positions=self.max_position_embeddings					,			pad_token_id=self.pad_token_id					,			)





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			*snake_case      )				->  Dict:
      """simple docstring"""





      a__	:						int      		=  CTRLModel(config=snake_case      )
      model.to(snake_case      )
      model.eval()

      model(snake_case					,			token_type_ids=snake_case					,			head_mask=snake_case      )
      model(snake_case					,			token_type_ids=snake_case      )
      a__	:						List[str]      		=  model(snake_case      )
      self.parent.assertEqual(result.last_hidden_state.shape					,			(self.batch_size, self.seq_length, self.hidden_size)      )
      self.parent.assertEqual(len(result.past_key_values      )					,			config.n_layer      )





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			*snake_case      )				->  Dict:
      """simple docstring"""





      a__	:						Tuple      		=  CTRLLMHeadModel(snake_case      )
      model.to(snake_case      )
      model.eval()

      a__	:						Union[str, Any]      		=  model(snake_case					,			token_type_ids=snake_case					,			labels=snake_case      )
      self.parent.assertEqual(result.loss.shape					,			()      )
      self.parent.assertEqual(result.logits.shape					,			(self.batch_size, self.seq_length, self.vocab_size)      )





   def 	_snake_case ( self      )				->  List[str]:
      """simple docstring"""





      a__	:						Any      		=  self.prepare_config_and_inputs()

      (
          a__     
      )	:						str      		=  config_and_inputs

      a__	:						Tuple      		=  {"input_ids": input_ids, "token_type_ids": token_type_ids, "head_mask": head_mask}

      return config, inputs_dict





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			*snake_case      )				->  int:
      """simple docstring"""





      a__	:						int      		=  self.num_labels
      a__	:						Dict      		=  CTRLForSequenceClassification(snake_case      )
      model.to(snake_case      )
      model.eval()
      a__	:						str      		=  ids_tensor([self.batch_size]					,			self.type_sequence_label_size      )
      a__	:						str      		=  model(snake_case					,			token_type_ids=snake_case					,			labels=snake_case      )
      self.parent.assertEqual(result.logits.shape					,			(self.batch_size, self.num_labels)      )







@require_torch
class   __lowerCAmelCase			(					_UpperCamelCase		,_UpperCamelCase		,_UpperCamelCase		,unittest.TestCase  ):
   _UpperCamelCase :						int     		=				(CTRLModel, CTRLLMHeadModel, CTRLForSequenceClassification) if is_torch_available() else ()
   _UpperCamelCase :						int     		=				(CTRLLMHeadModel,) if is_torch_available() else ()
   _UpperCamelCase :						List[str]     		=				(
       {
           """feature-extraction""": CTRLModel,
           """text-classification""": CTRLForSequenceClassification,
           """text-generation""": CTRLLMHeadModel,
           """zero-shot""": CTRLForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   _UpperCamelCase :						Dict     		=				True
   _UpperCamelCase :						str     		=				False
   _UpperCamelCase :						int     		=				False





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case      )				->  Tuple:
      """simple docstring"""





      if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
         # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
         # `CTRLConfig` was never used in pipeline tests, either because of a missing checkpoint or because a tiny
         # config could not be created.
         return True

      return False





   def 	_snake_case ( self      )				->  int:
      """simple docstring"""





      a__	:						Dict      		=  CTRLModelTester(self      )
      a__	:						Union[str, Any]      		=  ConfigTester(self					,			config_class=snake_case					,			n_embd=37      )





   def 	_snake_case ( self      )				->  Union[str, Any]:
      """simple docstring"""





      super().tearDown()
      # clean-up as much as possible GPU memory occupied by PyTorch
      gc.collect()
      torch.cuda.empty_cache()





   def 	_snake_case ( self      )				->  int:
      """simple docstring"""





      self.config_tester.run_common_tests()





   def 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





      a__	:						Optional[Any]      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_ctrl_model(*snake_case      )





   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      a__	:						List[str]      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_lm_head_model(*snake_case      )





   @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."      )
   def 	_snake_case ( self      )				->  Union[str, Any]:
      """simple docstring"""





      pass





   @slow
   def 	_snake_case ( self      )				->  Union[str, Any]:
      """simple docstring"""





      for model_name in CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         a__	:						List[Any]      		=  CTRLModel.from_pretrained(snake_case      )
         self.assertIsNotNone(snake_case      )





   @unittest.skip("The model doesn't support left padding"      )  # and it's not used enough to be worth fixing :)
   def 	_snake_case ( self      )				->  List[Any]:
      """simple docstring"""





      pass







@require_torch
class   __lowerCAmelCase			(					unittest.TestCase  ):





   def 	_snake_case ( self      )				->  List[str]:
      """simple docstring"""





      super().tearDown()
      # clean-up as much as possible GPU memory occupied by PyTorch
      gc.collect()
      torch.cuda.empty_cache()





   @slow
   def 	_snake_case ( self      )				->  List[Any]:
      """simple docstring"""





      a__	:						Optional[int]      		=  CTRLLMHeadModel.from_pretrained("ctrl"      )
      model.to(snake_case      )
      a__	:						int      		=  torch.tensor(
          [[11_859, 0, 1_611, 8]]					,			dtype=torch.long					,			device=snake_case      )  # Legal the president is
      a__	:						List[Any]      		=  [
          11_859,
          0,
          1_611,
          8,
          5,
          150,
          26_449,
          2,
          19,
          348,
          469,
          3,
          2_595,
          48,
          20_740,
          246_533,
          246_533,
          19,
          30,
          5,
      ]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

      a__	:						Any      		=  model.generate(snake_case					,			do_sample=snake_case      )
      self.assertListEqual(output_ids[0].tolist()					,			snake_case      )








from typing import Any, Callable, Dict, List, Optional, Union

import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker


SCREAMING_SNAKE_CASE__       :			Dict			      =						"""CompVis/stable-diffusion-v1-1"""
SCREAMING_SNAKE_CASE__       :			Dict			      =						"""CompVis/stable-diffusion-v1-2"""
SCREAMING_SNAKE_CASE__       :			Tuple			      =						"""CompVis/stable-diffusion-v1-3"""
SCREAMING_SNAKE_CASE__       :			str			      =						"""CompVis/stable-diffusion-v1-4"""







class   __lowerCAmelCase			(					_UpperCamelCase  ):





   def __init__( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case = True					,			)				->  Any:
      """simple docstring"""





      super()._init_()

      a__	:						Union[str, Any]      		=  StableDiffusionPipeline.from_pretrained(snake_case      )
      a__	:						Optional[Any]      		=  StableDiffusionPipeline.from_pretrained(snake_case      )
      a__	:						Optional[Any]      		=  StableDiffusionPipeline.from_pretrained(snake_case      )
      a__	:						int      		=  StableDiffusionPipeline(
          vae=snake_case					,			text_encoder=snake_case					,			tokenizer=snake_case					,			unet=snake_case					,			scheduler=snake_case					,			safety_checker=snake_case					,			feature_extractor=snake_case					,			requires_safety_checker=snake_case					,			)

      self.register_modules(pipelinea=self.pipea					,			pipelinea=self.pipea					,			pipelinea=self.pipea					,			pipelinea=self.pipea      )





   @property
   def 	_snake_case ( self      )				->  Dict[str, Any]:
      """simple docstring"""





      return {k: getattr(self					,			snake_case      ) for k in self.config.keys() if not k.startswith("_"      )}





   def 	_snake_case ( self					,			snake_case = "auto"      )				->  Optional[Any]:
      """simple docstring"""





      if slice_size == "auto":
         # half the attention head size is usually a good trade-off between
         # speed and memory
         a__	:						List[Any]      		=  self.unet.config.attention_head_dim // 2
      self.unet.set_attention_slice(snake_case      )





   def 	_snake_case ( self      )				->  Tuple:
      """simple docstring"""





      self.enable_attention_slicing(snake_case      )





   @torch.no_grad()
   def 	_snake_case ( self					,			snake_case					,			snake_case = 512					,			snake_case = 512					,			snake_case = 50					,			snake_case = 7.5					,			snake_case = None					,			snake_case = 1					,			snake_case = 0.0					,			snake_case = None					,			snake_case = None					,			snake_case = "pil"					,			snake_case = True					,			snake_case = None					,			snake_case = 1					,			**snake_case					,			)				->  Optional[int]:
      """simple docstring"""





      return self.pipea(
          prompt=snake_case					,			height=snake_case					,			width=snake_case					,			num_inference_steps=snake_case					,			guidance_scale=snake_case					,			negative_prompt=snake_case					,			num_images_per_prompt=snake_case					,			eta=snake_case					,			generator=snake_case					,			latents=snake_case					,			output_type=snake_case					,			return_dict=snake_case					,			callback=snake_case					,			callback_steps=snake_case					,			**snake_case					,			)





   @torch.no_grad()
   def 	_snake_case ( self					,			snake_case					,			snake_case = 512					,			snake_case = 512					,			snake_case = 50					,			snake_case = 7.5					,			snake_case = None					,			snake_case = 1					,			snake_case = 0.0					,			snake_case = None					,			snake_case = None					,			snake_case = "pil"					,			snake_case = True					,			snake_case = None					,			snake_case = 1					,			**snake_case					,			)				->  Dict:
      """simple docstring"""





      return self.pipea(
          prompt=snake_case					,			height=snake_case					,			width=snake_case					,			num_inference_steps=snake_case					,			guidance_scale=snake_case					,			negative_prompt=snake_case					,			num_images_per_prompt=snake_case					,			eta=snake_case					,			generator=snake_case					,			latents=snake_case					,			output_type=snake_case					,			return_dict=snake_case					,			callback=snake_case					,			callback_steps=snake_case					,			**snake_case					,			)





   @torch.no_grad()
   def 	_snake_case ( self					,			snake_case					,			snake_case = 512					,			snake_case = 512					,			snake_case = 50					,			snake_case = 7.5					,			snake_case = None					,			snake_case = 1					,			snake_case = 0.0					,			snake_case = None					,			snake_case = None					,			snake_case = "pil"					,			snake_case = True					,			snake_case = None					,			snake_case = 1					,			**snake_case					,			)				->  Union[str, Any]:
      """simple docstring"""





      return self.pipea(
          prompt=snake_case					,			height=snake_case					,			width=snake_case					,			num_inference_steps=snake_case					,			guidance_scale=snake_case					,			negative_prompt=snake_case					,			num_images_per_prompt=snake_case					,			eta=snake_case					,			generator=snake_case					,			latents=snake_case					,			output_type=snake_case					,			return_dict=snake_case					,			callback=snake_case					,			callback_steps=snake_case					,			**snake_case					,			)





   @torch.no_grad()
   def 	_snake_case ( self					,			snake_case					,			snake_case = 512					,			snake_case = 512					,			snake_case = 50					,			snake_case = 7.5					,			snake_case = None					,			snake_case = 1					,			snake_case = 0.0					,			snake_case = None					,			snake_case = None					,			snake_case = "pil"					,			snake_case = True					,			snake_case = None					,			snake_case = 1					,			**snake_case					,			)				->  Dict:
      """simple docstring"""





      return self.pipea(
          prompt=snake_case					,			height=snake_case					,			width=snake_case					,			num_inference_steps=snake_case					,			guidance_scale=snake_case					,			negative_prompt=snake_case					,			num_images_per_prompt=snake_case					,			eta=snake_case					,			generator=snake_case					,			latents=snake_case					,			output_type=snake_case					,			return_dict=snake_case					,			callback=snake_case					,			callback_steps=snake_case					,			**snake_case					,			)





   @torch.no_grad()
   def 	_snake_case ( self					,			snake_case					,			snake_case = 512					,			snake_case = 512					,			snake_case = 50					,			snake_case = 7.5					,			snake_case = None					,			snake_case = 1					,			snake_case = 0.0					,			snake_case = None					,			snake_case = None					,			snake_case = "pil"					,			snake_case = True					,			snake_case = None					,			snake_case = 1					,			**snake_case					,			)				->  Dict:
      """simple docstring"""





      a__	:						Any      		=  "cuda" if torch.cuda.is_available() else "cpu"
      self.to(snake_case      )

      # Checks if the height and width are divisible by 8 or not
      if height % 8 != 0 or width % 8 != 0:
         raise ValueError(F"""`height` and `width` must be divisible by 8 but are {height} and {width}."""      )

      # Get first result from Stable Diffusion Checkpoint v1.1
      a__	:						Any      		=  self.textaimg_sda_a(
          prompt=snake_case					,			height=snake_case					,			width=snake_case					,			num_inference_steps=snake_case					,			guidance_scale=snake_case					,			negative_prompt=snake_case					,			num_images_per_prompt=snake_case					,			eta=snake_case					,			generator=snake_case					,			latents=snake_case					,			output_type=snake_case					,			return_dict=snake_case					,			callback=snake_case					,			callback_steps=snake_case					,			**snake_case					,			)

      # Get first result from Stable Diffusion Checkpoint v1.2
      a__	:						List[Any]      		=  self.textaimg_sda_a(
          prompt=snake_case					,			height=snake_case					,			width=snake_case					,			num_inference_steps=snake_case					,			guidance_scale=snake_case					,			negative_prompt=snake_case					,			num_images_per_prompt=snake_case					,			eta=snake_case					,			generator=snake_case					,			latents=snake_case					,			output_type=snake_case					,			return_dict=snake_case					,			callback=snake_case					,			callback_steps=snake_case					,			**snake_case					,			)

      # Get first result from Stable Diffusion Checkpoint v1.3
      a__	:						Optional[Any]      		=  self.textaimg_sda_a(
          prompt=snake_case					,			height=snake_case					,			width=snake_case					,			num_inference_steps=snake_case					,			guidance_scale=snake_case					,			negative_prompt=snake_case					,			num_images_per_prompt=snake_case					,			eta=snake_case					,			generator=snake_case					,			latents=snake_case					,			output_type=snake_case					,			return_dict=snake_case					,			callback=snake_case					,			callback_steps=snake_case					,			**snake_case					,			)

      # Get first result from Stable Diffusion Checkpoint v1.4
      a__	:						Dict      		=  self.textaimg_sda_a(
          prompt=snake_case					,			height=snake_case					,			width=snake_case					,			num_inference_steps=snake_case					,			guidance_scale=snake_case					,			negative_prompt=snake_case					,			num_images_per_prompt=snake_case					,			eta=snake_case					,			generator=snake_case					,			latents=snake_case					,			output_type=snake_case					,			return_dict=snake_case					,			callback=snake_case					,			callback_steps=snake_case					,			**snake_case					,			)

      # Get all result images into a single list and pass it via StableDiffusionPipelineOutput for final result
      return StableDiffusionPipelineOutput([resa[0], resa[0], resa[0], resa[0]]      )

