code: 
  
  
  """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 __snake_case		(					_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       unittest.TestCase      ):
  
  
  
  
  
  
  
     """simple docstring"""
  
  
  
  
     lowerCAmelCase_				:   int     	=			StableDiffusionXLPipeline
     lowerCAmelCase_				:   List[Any]     	=			TEXT_TO_IMAGE_PARAMS
     lowerCAmelCase_				:   int     	=			TEXT_TO_IMAGE_BATCH_PARAMS
     lowerCAmelCase_				:   Any     	=			TEXT_TO_IMAGE_IMAGE_PARAMS
     lowerCAmelCase_				:   str     	=			TEXT_TO_IMAGE_IMAGE_PARAMS
  
  
  
  
     def SCREAMING_SNAKE_CASE_      (		self							:List[Any]       ):
            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=UpperCamelCase__				,  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.00085				,  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=1_000				,  hidden_act="gelu"				,  projection_dim=32				,  )
            _a          =				CLIPTextModel(UpperCamelCase__       )
            _a          =				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				,  local_files_only=UpperCamelCase__       )
  
            _a          =				CLIPTextModelWithProjection(UpperCamelCase__       )
            _a          =				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				,  local_files_only=UpperCamelCase__       )
  
            _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 SCREAMING_SNAKE_CASE_      (		self							:Union[str, Any]				,  UpperCamelCase__							:int				,  UpperCamelCase__							:Optional[int]=0       ):
            if str(UpperCamelCase__       ).startswith("mps"       ):
                   _a          =				torch.manual_seed(UpperCamelCase__       )
            else:
                   _a          =				torch.Generator(device=UpperCamelCase__       ).manual_seed(UpperCamelCase__       )
            _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 SCREAMING_SNAKE_CASE_      (		self							:Optional[Any]       ):
            _a          =				"cpu"  # ensure determinism for the device-dependent torch.Generator
            _a          =				self.get_dummy_components()
            _a          =				StableDiffusionXLPipeline(**UpperCamelCase__       )
            _a          =				sd_pipe.to(UpperCamelCase__       )
            sd_pipe.set_progress_bar_config(disable=UpperCamelCase__       )
  
            _a          =				self.get_dummy_inputs(UpperCamelCase__       )
            _a          =				sd_pipe(**UpperCamelCase__       ).images
            _a          =				image[0, -3:, -3:, -1]
  
            assert image.shape == (1, 64, 64, 3)
            _a          =				np.array([0.5873, 0.6128, 0.4797, 0.5122, 0.5674, 0.4639, 0.5227, 0.5149, 0.4747]       )
  
            assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
  
  
  
  
     def SCREAMING_SNAKE_CASE_      (		self							:Union[str, Any]       ):
            _a          =				self.get_dummy_components()
            _a          =				StableDiffusionXLPipeline(**UpperCamelCase__       )
            _a          =				sd_pipe.to(UpperCamelCase__       )
            _a          =				sd_pipe.to(UpperCamelCase__       )
            sd_pipe.set_progress_bar_config(disable=UpperCamelCase__       )
  
            # forward without prompt embeds
            _a          =				self.get_dummy_inputs(UpperCamelCase__       )
            _a          =				3 * ["this is a negative prompt"]
            _a          =				negative_prompt
            _a          =				3 * [inputs["prompt"]]
  
            _a          =				sd_pipe(**UpperCamelCase__       )
            _a          =				output.images[0, -3:, -3:, -1]
  
            # forward with prompt embeds
            _a          =				self.get_dummy_inputs(UpperCamelCase__       )
            _a          =				3 * ["this is a negative prompt"]
            _a          =				3 * [inputs.pop("prompt"       )]
  
            (
                (
                _a 
            )   , (
                _a 
            )   , (
                _a 
            )   , (
                _a 
            )   ,  
            )          =				sd_pipe.encode_prompt(UpperCamelCase__				,  negative_prompt=UpperCamelCase__       )
  
            _a          =				sd_pipe(
                **UpperCamelCase__				,  prompt_embeds=UpperCamelCase__				,  negative_prompt_embeds=UpperCamelCase__				,  pooled_prompt_embeds=UpperCamelCase__				,  negative_pooled_prompt_embeds=UpperCamelCase__				,  )
            _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 SCREAMING_SNAKE_CASE_      (		self							:Any       ):
            super().test_attention_slicing_forward_pass(expected_max_diff=3E-3       )
  
  
  
  
     def SCREAMING_SNAKE_CASE_      (		self							:List[Any]       ):
            super().test_inference_batch_single_identical(expected_max_diff=3E-3       )
  
  
  
  
  
  
  @slow
  @require_torch_gpu
  class __snake_case		(					unittest.TestCase      ):
  
  
  
  
  
  
  
     """simple docstring"""
  
  
  
  
     def SCREAMING_SNAKE_CASE_      (		self							:List[str]       ):
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()
  
  
  
  
     def SCREAMING_SNAKE_CASE_      (		self							:Optional[int]				,  UpperCamelCase__							:Tuple				,  UpperCamelCase__							:List[str]="cpu"				,  UpperCamelCase__							:Dict=torch.floataa				,  UpperCamelCase__							:Dict=0       ):
            _a          =				torch.Generator(device=UpperCamelCase__       ).manual_seed(UpperCamelCase__       )
            _a          =				np.random.RandomState(UpperCamelCase__       ).standard_normal((1, 4, 64, 64)       )
            _a          =				torch.from_numpy(UpperCamelCase__       ).to(device=UpperCamelCase__				,  dtype=UpperCamelCase__       )
            _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 SCREAMING_SNAKE_CASE_      (		self							:Optional[Any]       ):
            _a          =				DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-base"       )
            pipe.to(UpperCamelCase__       )
            pipe.set_progress_bar_config(disable=UpperCamelCase__       )
  
            _a          =				self.get_inputs(UpperCamelCase__       )
            _a          =				pipe(**UpperCamelCase__       ).images
            _a          =				image[0, -3:, -3:, -1].flatten()
  
            assert image.shape == (1, 512, 512, 3)
            _a          =				np.array([0.49493, 0.47896, 0.40798, 0.54214, 0.53212, 0.48202, 0.47656, 0.46329, 0.48506]       )
            assert np.abs(image_slice - expected_slice       ).max() < 7E-3
  
style_context: 
  
  
  """simple docstring"""
  
  
  
  
  import argparse
  import time
  
  import torch
  import transformers
  from measures_util import end_measure, log_measures, start_measure
  from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForSeqaSeqLM, AutoTokenizer
  
  from accelerate.utils import compute_module_sizes
  
  
  __SCREAMING_SNAKE_CASE					=						{
      """gpt-j-6b""": {"""is_causal""": True, """model""": """sgugger/sharded-gpt-j-6B""", """tokenizer""": """EleutherAI/gpt-j-6B"""},
      """gpt-neox""": {"""is_causal""": True, """model""": """EleutherAI/gpt-neox-20b"""},
      """opt""": {"""is_causal""": True, """model""": """facebook/opt-30b"""},
      """T0pp""": {"""is_causal""": False, """model""": """bigscience/T0pp""", """model_revision""": """sharded"""},
  }
  
  __SCREAMING_SNAKE_CASE					=						[
      """Hello, my name is""",
      """Are unicorns real? Unicorns are""",
      """For the first time in several years,""",
      """My name is Julien and I am""",
      """The goal of life is""",
      """Whenever I'm sad, I like to""",
  ]
  
  
  
  def 						__a      (  ):
  
  
  
  
  
         """simple docstring"""
  
  
  
  
  
  
  
         _a          =				argparse.ArgumentParser(description="Run and time generations on a big model using Accelerate."					)
         parser.add_argument("model_name",		type=a,		default=a,		help="The name of the model to try."					)
         parser.add_argument(
             "--tokenizer_name",		type=a,		default=a,		help="The name of the tokenizer (if different from the model."					)
         parser.add_argument("--is_causal",		type=a,		default=a,		help="Whether or not the model is causal."					)
         parser.add_argument(
             "--model_revision",		type=a,		default=a,		help="The revision to use for the model checkpoint."					)
         parser.add_argument("--torch_dtype",		type=a,		default=a,		help="The dtype for the model."					)
         parser.add_argument("--disk_offload",		action="store_true"					)
  
         _a          =				parser.parse_args()
  
         # Sanitize args
         if args.model_name in DEFAULT_MODELS:
                _a          =				DEFAULT_MODELS[args.model_name]
                _a          =				defaults["model"]
                if args.tokenizer_name is None:
                       _a          =				defaults.get("tokenizer",		args.model_name					)
                if args.is_causal is None:
                       _a          =				defaults["is_causal"]
                if args.model_revision is None:
                       _a          =				defaults.get("model_revision",		"main"					)
  
         if args.is_causal is None:
                raise ValueError("Could not infer the default for `--is_causal`, pass either True or False for it."					)
         if args.tokenizer_name is None:
                _a          =				args.model_name
         if args.model_revision is None:
                _a          =				"main"
  
         return args
  
  
  
  def 						__a      (  ):
  
  
  
  
  
         """simple docstring"""
  
  
  
  
  
  
  
         transformers.utils.logging.set_verbosity_error()
         _a          =				parse_args()
  
         if args.torch_dtype is None:
                _a          =				AutoConfig.from_pretrained(args.model_name					)
                _a          =				getattr(a,		"torch_dtype",		torch.floataa					)
         else:
                _a          =				getattr(a,		args.torch_dtype					)
         _a          =				AutoModelForCausalLM if args.is_causal else AutoModelForSeqaSeqLM
         _a          =				{
             "torch_dtype": torch_dtype,
             "revision": args.model_revision,
         }
         if args.disk_offload:
                _a          =				"tmp_offload"
                _a          =				True
  
         _a          =				start_measure()
         _a          =				model_cls.from_pretrained(args.model_name,		device_map="auto",		**a					)
         _a          =				end_measure(a					)
         log_measures(a,		"Model loading"					)
  
         _a          =				compute_module_sizes(a					)
         _a          =				{v: 0 for v in model.hf_device_map.values()}
         for module, device in model.hf_device_map.items():
                device_size[device] += module_sizes[module]
         _a          =				"\n".join([F'- {device}: {size // 2**2_0}MiB' for device, size in device_size.items()]					)
         print(F'\nTheoretical use:\n{message}'					)
  
         _a          =				AutoTokenizer.from_pretrained(args.tokenizer_name					)
  
         _a          =				start_measure()
         _a          =				[]
         _a          =				[]
         _a          =				[]
         for prompt in PROMPTS:
                _a          =				tokenizer(a,		return_tensors="pt"					).to(0					)
                _a          =				inputs["input_ids"][0].tolist()
                _a          =				time.time()
                _a          =				model.generate(inputs["input_ids"]					)
                _a          =				time.time()
                _a          =				outputs[0].tolist()
                _a          =				len(a					) if outputs[: len(a					)] != tokens else len(a					) - len(a					)
                _a          =				after_generate - before_generate
  
                _a          =				tokenizer.decode(a,		skip_special_tokens=a					)
                texts_outs.append(a					)
                generation_times.append(a					)
                gen_tokens.append(a					)
                print(F'Prompt: {prompt}\nGeneration {text_out}\nIn {generation_time:.2f}s for {num_gen_tokens} tokens\n'					)
  
         _a          =				end_measure(a					)
         log_measures(a,		"Model generation"					)
  
         _a          =				[gen / tok for gen, tok in zip(a,		a					)]
         _a          =				sum(a					) / len(a					)
         print(F'Average time of generation per token: {avg_gen:.2f}s'					)
         print(F'First generation (avg time per token): {generation_times_per_token[0]:.2f}s'					)
         _a          =				sum(generation_times_per_token[1:]					) / (len(a					) - 1)
         print(F'Average time of generation per token (excluding the first): {avg_gen:.2f}s'					)
  
  
  if __name__ == "__main__":
     main()
  