'''simple docstring'''




import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class       __UpperCAmelCase							(							_UpperCamelCase		, _UpperCamelCase		, unittest.TestCase	):





       '''simple docstring'''



       __lowercase   :   Any  					=							StableDiffusionXLPipeline
       __lowercase   :   Any  					=							TEXT_TO_IMAGE_PARAMS
       __lowercase   :   Union[str, Any]  					=							TEXT_TO_IMAGE_BATCH_PARAMS
       __lowercase   :   List[Any]  					=							TEXT_TO_IMAGE_IMAGE_PARAMS
       __lowercase   :   Union[str, Any]  					=							TEXT_TO_IMAGE_IMAGE_PARAMS

       def    __A			( self					)						->	Tuple:
           torch.manual_seed(0					)
           A_      =  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''')       ,	attention_head_dim=(2, 4)       ,	use_linear_projection=_SCREAMING_SNAKE_CASE       ,	addition_embed_type='''text_time'''       ,	addition_time_embed_dim=8       ,	transformer_layers_per_block=(1, 2)       ,	projection_class_embeddings_input_dim=80       ,	cross_attention_dim=64       ,	)
           A_      =  EulerDiscreteScheduler(
               beta_start=0.00_085       ,	beta_end=0.012       ,	steps_offset=1       ,	beta_schedule='''scaled_linear'''       ,	timestep_spacing='''leading'''       ,	)
           torch.manual_seed(0					)
           A_      =  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       ,	sample_size=128       ,	)
           torch.manual_seed(0					)
           A_      =  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       ,	hidden_act='''gelu'''       ,	projection_dim=32       ,	)
           A_      =  CLIPTextModel(_SCREAMING_SNAKE_CASE					)
           A_      =  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       ,	local_files_only=_SCREAMING_SNAKE_CASE					)

           A_      =  CLIPTextModelWithProjection(_SCREAMING_SNAKE_CASE					)
           A_      =  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       ,	local_files_only=_SCREAMING_SNAKE_CASE					)

           A_      =  {
               '''unet''': unet,
               '''scheduler''': scheduler,
               '''vae''': vae,
               '''text_encoder''': text_encoder,
               '''tokenizer''': tokenizer,
               '''text_encoder_2''': text_encoder_a,
               '''tokenizer_2''': tokenizer_a,
               # "safety_checker": None,
               # "feature_extractor": None,
           }
           return components

       def    __A			( self       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE=0					)						->	Optional[int]:
           if str(_SCREAMING_SNAKE_CASE					).startswith('''mps'''					):
               A_      =  torch.manual_seed(_SCREAMING_SNAKE_CASE					)
           else:
               A_      =  torch.Generator(device=_SCREAMING_SNAKE_CASE					).manual_seed(_SCREAMING_SNAKE_CASE					)
           A_      =  {
               '''prompt''': '''A painting of a squirrel eating a burger''',
               '''generator''': generator,
               '''num_inference_steps''': 2,
               '''guidance_scale''': 5.0,
               '''output_type''': '''numpy''',
           }
           return inputs

       def    __A			( self					)						->	Any:
           A_      =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           A_      =  self.get_dummy_components()
           A_      =  StableDiffusionXLPipeline(**_SCREAMING_SNAKE_CASE					)
           A_      =  sd_pipe.to(_SCREAMING_SNAKE_CASE					)
           sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

           A_      =  self.get_dummy_inputs(_SCREAMING_SNAKE_CASE					)
           A_      =  sd_pipe(**_SCREAMING_SNAKE_CASE					).images
           A_      =  image[0, -3:, -3:, -1]

           assert image.shape == (1, 64, 64, 3)
           A_      =  np.array([0.5_873, 0.6_128, 0.4_797, 0.5_122, 0.5_674, 0.4_639, 0.5_227, 0.5_149, 0.4_747]					)

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

       def    __A			( self					)						->	Any:
           A_      =  self.get_dummy_components()
           A_      =  StableDiffusionXLPipeline(**_SCREAMING_SNAKE_CASE					)
           A_      =  sd_pipe.to(_SCREAMING_SNAKE_CASE					)
           A_      =  sd_pipe.to(_SCREAMING_SNAKE_CASE					)
           sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

           # forward without prompt embeds
           A_      =  self.get_dummy_inputs(_SCREAMING_SNAKE_CASE					)
           A_      =  3 * ['''this is a negative prompt''']
           A_      =  negative_prompt
           A_      =  3 * [inputs['''prompt''']]

           A_      =  sd_pipe(**_SCREAMING_SNAKE_CASE					)
           A_      =  output.images[0, -3:, -3:, -1]

           # forward with prompt embeds
           A_      =  self.get_dummy_inputs(_SCREAMING_SNAKE_CASE					)
           A_      =  3 * ['''this is a negative prompt''']
           A_      =  3 * [inputs.pop('''prompt'''					)]

           (
               (
               A_
           )  ,(
               A_
           )  ,(
               A_
           )  ,(
               A_
           )  ,
           )      =  sd_pipe.encode_prompt(_SCREAMING_SNAKE_CASE       ,	negative_prompt=_SCREAMING_SNAKE_CASE					)

           A_      =  sd_pipe(
               **_SCREAMING_SNAKE_CASE       ,	prompt_embeds=_SCREAMING_SNAKE_CASE       ,	negative_prompt_embeds=_SCREAMING_SNAKE_CASE       ,	pooled_prompt_embeds=_SCREAMING_SNAKE_CASE       ,	negative_pooled_prompt_embeds=_SCREAMING_SNAKE_CASE       ,	)
           A_      =  output.images[0, -3:, -3:, -1]

           # make sure that it's equal
           assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()					).max() < 1E-4

       def    __A			( self					)						->	Optional[int]:
           super().test_attention_slicing_forward_pass(expected_max_diff=3E-3					)

       def    __A			( self					)						->	str:
           super().test_inference_batch_single_identical(expected_max_diff=3E-3					)




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





       '''simple docstring'''



       def    __A			( self					)						->	List[Any]:
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()

       def    __A			( self       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE="cpu"       ,	_SCREAMING_SNAKE_CASE=torch.floataa       ,	_SCREAMING_SNAKE_CASE=0					)						->	List[Any]:
           A_      =  torch.Generator(device=_SCREAMING_SNAKE_CASE					).manual_seed(_SCREAMING_SNAKE_CASE					)
           A_      =  np.random.RandomState(_SCREAMING_SNAKE_CASE					).standard_normal((1, 4, 64, 64)					)
           A_      =  torch.from_numpy(_SCREAMING_SNAKE_CASE					).to(device=_SCREAMING_SNAKE_CASE       ,	dtype=_SCREAMING_SNAKE_CASE					)
           A_      =  {
               '''prompt''': '''a photograph of an astronaut riding a horse''',
               '''latents''': latents,
               '''generator''': generator,
               '''num_inference_steps''': 3,
               '''guidance_scale''': 7.5,
               '''output_type''': '''numpy''',
           }
           return inputs

       def    __A			( self					)						->	Tuple:
           A_      =  DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''					)
           pipe.to(_SCREAMING_SNAKE_CASE					)
           pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

           A_      =  self.get_inputs(_SCREAMING_SNAKE_CASE					)
           A_      =  pipe(**_SCREAMING_SNAKE_CASE					).images
           A_      =  image[0, -3:, -3:, -1].flatten()

           assert image.shape == (1, 512, 512, 3)
           A_      =  np.array([0.49_493, 0.47_896, 0.40_798, 0.54_214, 0.53_212, 0.48_202, 0.47_656, 0.46_329, 0.48_506]					)
           assert np.abs(image_slice - expected_slice					).max() < 7E-3







'''simple docstring'''




import argparse
import logging
import sys
import time

import tensorflow as tf
from datasets import load_dataset

from transformers import AutoTokenizer, TFAutoModelForSequenceClassification


if __name__ == "__main__":
  __snake_case  : Any		       =    argparse.ArgumentParser()

  # Hyperparameters sent by the client are passed as command-line arguments to the script.
  parser.add_argument('--epochs', type=int, default=1)
  parser.add_argument('--per_device_train_batch_size', type=int, default=16)
  parser.add_argument('--per_device_eval_batch_size', type=int, default=8)
  parser.add_argument('--model_name_or_path', type=str)
  parser.add_argument('--learning_rate', type=str, default=5E-5)
  parser.add_argument('--do_train', type=bool, default=True)
  parser.add_argument('--do_eval', type=bool, default=True)
  parser.add_argument('--output_dir', type=str)

  __snake_case       ,       __snake_case  : Dict		       =    parser.parse_known_args()

  # overwrite batch size until we have tf_glue.py
  __snake_case  : Any		       =    16
  __snake_case  : Optional[Any]		       =    16

  # Set up logging
  __snake_case  : Union[str, Any]		       =    logging.getLogger(__name__)

  logging.basicConfig(
      level=logging.getLevelName('INFO'),
      handlers=[logging.StreamHandler(sys.stdout)],
      format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
  )

  # Load model and tokenizer
  __snake_case  : Optional[Any]		       =    TFAutoModelForSequenceClassification.from_pretrained(args.model_name_or_path)
  __snake_case  : Optional[Any]		       =    AutoTokenizer.from_pretrained(args.model_name_or_path)

  # Load dataset
  __snake_case       ,       __snake_case  : Optional[int]		       =    load_dataset('imdb', split=['train', 'test'])
  __snake_case  : List[str]		       =    train_dataset.shuffle().select(range(5_000))  # smaller the size for train dataset to 5k
  __snake_case  : List[str]		       =    test_dataset.shuffle().select(range(500))  # smaller the size for test dataset to 500

  # Preprocess train dataset
  __snake_case  : int		       =    train_dataset.map(
      lambda e: tokenizer(e['text'], truncation=True, padding='max_length'), batched=True
  )
  train_dataset.set_format(type='tensorflow', columns=['input_ids', 'attention_mask', 'label'])

  __snake_case  : Dict		       =    {
      x: train_dataset[x].to_tensor(default_value=0, shape=[None, tokenizer.model_max_length])
      for x in ['input_ids', 'attention_mask']
  }
  __snake_case  : str		       =    tf.data.Dataset.from_tensor_slices((train_features, train_dataset['label'])).batch(
      args.per_device_train_batch_size
  )

  # Preprocess test dataset
  __snake_case  : Tuple		       =    test_dataset.map(
      lambda e: tokenizer(e['text'], truncation=True, padding='max_length'), batched=True
  )
  test_dataset.set_format(type='tensorflow', columns=['input_ids', 'attention_mask', 'label'])

  __snake_case  : Optional[int]		       =    {
      x: test_dataset[x].to_tensor(default_value=0, shape=[None, tokenizer.model_max_length])
      for x in ['input_ids', 'attention_mask']
  }
  __snake_case  : List[str]		       =    tf.data.Dataset.from_tensor_slices((test_features, test_dataset['label'])).batch(
      args.per_device_eval_batch_size
  )

  # fine optimizer and loss
  __snake_case  : Optional[int]		       =    tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
  __snake_case  : Union[str, Any]		       =    tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
  __snake_case  : int		       =    [tf.keras.metrics.SparseCategoricalAccuracy()]
  model.compile(optimizer=optimizer, loss=loss, metrics=metrics)

  __snake_case  : Any		       =    time.time()
  __snake_case  : int		       =    model.fit(tf_train_dataset, epochs=args.epochs, batch_size=args.per_device_train_batch_size)
  __snake_case  : Union[str, Any]		       =    time.time() - start_train_time

  logger.info('*** Train ***')
  logger.info(F"""train_runtime = {end_train_time}""")
  for key, value in train_results.history.items():
    logger.info(F"""  {key} = {value}""")






