






"""simple docstring"""




import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()




class  lowercase_     (_UpperCAmelCase,     unittest.TestCase  ):
      A__						:		Tuple							      =					KandinskyPipeline
      A__						:		Optional[int]							      =					[
          '''prompt''',
          '''image_embeds''',
          '''negative_image_embeds''',
      ]
      A__						:		Dict							      =					['''prompt''', '''negative_prompt''', '''image_embeds''', '''negative_image_embeds''']
      A__						:		Any							      =					[
          '''generator''',
          '''height''',
          '''width''',
          '''latents''',
          '''guidance_scale''',
          '''negative_prompt''',
          '''num_inference_steps''',
          '''return_dict''',
          '''guidance_scale''',
          '''num_images_per_prompt''',
          '''output_type''',
          '''return_dict''',
      ]
      A__						:		str							      =					False






      @property
      def       lowerCamelCase__     (     self						)						->Dict:







        '''simple docstring'''


        return 3_2






      @property
      def       lowerCamelCase__     (     self						)						->Union[str, Any]:







        '''simple docstring'''


        return 3_2






      @property
      def       lowerCamelCase__     (     self						)						->Any:







        '''simple docstring'''


        return self.time_input_dim






      @property
      def       lowerCamelCase__     (     self						)						->Dict:







        '''simple docstring'''


        return self.time_input_dim * 4






      @property
      def       lowerCamelCase__     (     self						)						->str:







        '''simple docstring'''


        return 1_0_0






      @property
      def       lowerCamelCase__     (     self						)						->Optional[int]:







        '''simple docstring'''


        _a					      =						XLMRobertaTokenizerFast.from_pretrained("YiYiXu/tiny-random-mclip-base"						)
        return tokenizer






      @property
      def       lowerCamelCase__     (     self						)						->Dict:







        '''simple docstring'''


        torch.manual_seed(0						)
        _a					      =						MCLIPConfig(
            numDims=self.cross_attention_dim  ,  transformerDimensions=self.text_embedder_hidden_size  ,  hidden_size=self.text_embedder_hidden_size  ,  intermediate_size=3_7  ,  num_attention_heads=4  ,  num_hidden_layers=5  ,  vocab_size=1_0_0_5  ,  )

        _a					      =						MultilingualCLIP(a_						)
        _a					      =						text_encoder.eval()

        return text_encoder






      @property
      def       lowerCamelCase__     (     self						)						->List[str]:







        '''simple docstring'''


        torch.manual_seed(0						)

        _a					      =						{
            "in_channels": 4,
            # Out channels is double in channels because predicts mean and variance
            "out_channels": 8,
            "addition_embed_type": "text_image",
            "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
            "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
            "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
            "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
            "layers_per_block": 1,
            "encoder_hid_dim": self.text_embedder_hidden_size,
            "encoder_hid_dim_type": "text_image_proj",
            "cross_attention_dim": self.cross_attention_dim,
            "attention_head_dim": 4,
            "resnet_time_scale_shift": "scale_shift",
            "class_embed_type": None,
        }

        _a					      =						UNetaDConditionModel(**a_						)
        return model






      @property
      def       lowerCamelCase__     (     self						)						->int:







        '''simple docstring'''


        return {
            "block_out_channels": [3_2, 6_4],
            "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
            "in_channels": 3,
            "latent_channels": 4,
            "layers_per_block": 1,
            "norm_num_groups": 8,
            "norm_type": "spatial",
            "num_vq_embeddings": 1_2,
            "out_channels": 3,
            "up_block_types": [
                "AttnUpDecoderBlock2D",
                "UpDecoderBlock2D",
            ],
            "vq_embed_dim": 4,
        }






      @property
      def       lowerCamelCase__     (     self						)						->Dict:







        '''simple docstring'''


        torch.manual_seed(0						)
        _a					      =						VQModel(**self.dummy_movq_kwargs						)
        return model






      def       lowerCamelCase__     (     self						)						->Any:







        '''simple docstring'''


        _a					      =						self.dummy_text_encoder
        _a					      =						self.dummy_tokenizer
        _a					      =						self.dummy_unet
        _a					      =						self.dummy_movq

        _a					      =						DDIMScheduler(
            num_train_timesteps=1_0_0_0  ,  beta_schedule="linear"  ,  beta_start=0.00_085  ,  beta_end=0.012  ,  clip_sample=a_  ,  set_alpha_to_one=a_  ,  steps_offset=1  ,  prediction_type="epsilon"  ,  thresholding=a_  ,  )

        _a					      =						{
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "unet": unet,
            "scheduler": scheduler,
            "movq": movq,
        }
        return components






      def       lowerCamelCase__     (     self  ,  a_  ,  a_=0						)						->Union[str, Any]:







        '''simple docstring'''


        _a					      =						floats_tensor((1, self.cross_attention_dim)  ,  rng=random.Random(a_						)						).to(a_						)
        _a					      =						floats_tensor((1, self.cross_attention_dim)  ,  rng=random.Random(seed + 1						)						).to(a_						)
        if str(a_						).startswith("mps"						):
          _a					      =						torch.manual_seed(a_						)
        else:
          _a					      =						torch.Generator(device=a_						).manual_seed(a_						)
        _a					      =						{
            "prompt": "horse",
            "image_embeds": image_embeds,
            "negative_image_embeds": negative_image_embeds,
            "generator": generator,
            "height": 6_4,
            "width": 6_4,
            "guidance_scale": 4.0,
            "num_inference_steps": 2,
            "output_type": "np",
        }
        return inputs






      def       lowerCamelCase__     (     self						)						->Dict:







        '''simple docstring'''


        _a					      =						"cpu"

        _a					      =						self.get_dummy_components()

        _a					      =						self.pipeline_class(**a_						)
        _a					      =						pipe.to(a_						)

        pipe.set_progress_bar_config(disable=a_						)

        _a					      =						pipe(**self.get_dummy_inputs(a_						)						)
        _a					      =						output.images

        _a					      =						pipe(
            **self.get_dummy_inputs(a_						)  ,  return_dict=a_  ,  )[0]

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

        assert image.shape == (1, 6_4, 6_4, 3)

        _a					      =						np.array(
            [0.328_663, 1.0, 0.23_216_873, 1.0, 0.92_717_564, 0.4_639_046, 0.96_894_777, 0.31_713_378, 0.6_293_953]						)

        assert (
            np.abs(image_slice.flatten() - expected_slice						).max() < 1E-2
        ), f''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''
        assert (
            np.abs(image_from_tuple_slice.flatten() - expected_slice						).max() < 1E-2
        ), f''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''







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






      def       lowerCamelCase__     (     self						)						->Optional[int]:







        '''simple docstring'''


        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()






      def       lowerCamelCase__     (     self						)						->int:







        '''simple docstring'''


        _a					      =						load_numpy(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
            "/kandinsky/kandinsky_text2img_cat_fp16.npy"						)

        _a					      =						KandinskyPriorPipeline.from_pretrained(
            "kandinsky-community/kandinsky-2-1-prior"  ,  torch_dtype=torch.floataa						)
        pipe_prior.to(a_						)

        _a					      =						KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1"  ,  torch_dtype=torch.floataa						)
        _a					      =						pipeline.to(a_						)
        pipeline.set_progress_bar_config(disable=a_						)

        _a					      =						"red cat, 4k photo"

        _a					      =						torch.Generator(device="cuda"						).manual_seed(0						)
        _a      ,						_a					      =						pipe_prior(
            a_  ,  generator=a_  ,  num_inference_steps=5  ,  negative_prompt=""  ,  ).to_tuple()

        _a					      =						torch.Generator(device="cuda"						).manual_seed(0						)
        _a					      =						pipeline(
            a_  ,  image_embeds=a_  ,  negative_image_embeds=a_  ,  generator=a_  ,  num_inference_steps=1_0_0  ,  output_type="np"  ,  )

        _a					      =						output.images[0]

        assert image.shape == (5_1_2, 5_1_2, 3)

        assert_mean_pixel_difference(a_  ,  a_						)










"""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


UpperCamelCase											=     {
    """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"""},
}

UpperCamelCase											=     [
    """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 							lowerCAmelCase			(   )							->				Optional[int]:







  '''simple docstring'''







  _a					      =						argparse.ArgumentParser(description="Run and time generations on a big model using Accelerate."       )
  parser.add_argument("model_name"  ,			type=UpperCamelCase_  ,			default=UpperCamelCase_  ,			help="The name of the model to try."       )
  parser.add_argument(
      "--tokenizer_name"  ,			type=UpperCamelCase_  ,			default=UpperCamelCase_  ,			help="The name of the tokenizer (if different from the model."       )
  parser.add_argument("--is_causal"  ,			type=UpperCamelCase_  ,			default=UpperCamelCase_  ,			help="Whether or not the model is causal."       )
  parser.add_argument(
      "--model_revision"  ,			type=UpperCamelCase_  ,			default=UpperCamelCase_  ,			help="The revision to use for the model checkpoint."       )
  parser.add_argument("--torch_dtype"  ,			type=UpperCamelCase_  ,			default=UpperCamelCase_  ,			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 							lowerCAmelCase			(   )							->				str:







  '''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(UpperCamelCase_  ,			"torch_dtype"  ,			torch.floataa       )
  else:
    _a					      =						getattr(UpperCamelCase_  ,			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"  ,			**UpperCamelCase_       )
  _a					      =						end_measure(UpperCamelCase_       )
  log_measures(UpperCamelCase_  ,			"Model loading"       )

  _a					      =						compute_module_sizes(UpperCamelCase_       )
  _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**20}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(UpperCamelCase_  ,			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(UpperCamelCase_       ) if outputs[: len(UpperCamelCase_       )] != tokens else len(UpperCamelCase_       ) - len(UpperCamelCase_       )
    _a					      =						after_generate - before_generate

    _a					      =						tokenizer.decode(UpperCamelCase_  ,			skip_special_tokens=UpperCamelCase_       )
    texts_outs.append(UpperCamelCase_       )
    generation_times.append(UpperCamelCase_       )
    gen_tokens.append(UpperCamelCase_       )
    print(f'''Prompt: {prompt}\nGeneration {text_out}\nIn {generation_time:.2f}s for {num_gen_tokens} tokens\n'''       )

  _a					      =						end_measure(UpperCamelCase_       )
  log_measures(UpperCamelCase_  ,			"Model generation"       )

  _a					      =						[gen / tok for gen, tok in zip(UpperCamelCase_  ,			UpperCamelCase_       )]
  _a					      =						sum(UpperCamelCase_       ) / len(UpperCamelCase_       )
  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(UpperCamelCase_       ) - 1)
  print(f'''Average time of generation per token (excluding the first): {avg_gen:.2f}s'''       )


if __name__ == "__main__":
     main()


