

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


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

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 					__lowerCamelCase			(      )							->   Any:
      """simple docstring"""

      SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       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"    )

      SCREAMING_SNAKE_CASE_      :			str     =       parser.parse_args()

      # Sanitize args
      if args.model_name in DEFAULT_MODELS:
            SCREAMING_SNAKE_CASE_      :			Tuple     =       DEFAULT_MODELS[args.model_name]
            SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       defaults["model"]
            if args.tokenizer_name is None:
                  SCREAMING_SNAKE_CASE_      :			str     =       defaults.get("tokenizer"  ,  args.model_name    )
            if args.is_causal is None:
                  SCREAMING_SNAKE_CASE_      :			List[Any]     =       defaults["is_causal"]
            if args.model_revision is None:
                  SCREAMING_SNAKE_CASE_      :			Dict     =       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:
            SCREAMING_SNAKE_CASE_      :			List[Any]     =       args.model_name
      if args.model_revision is None:
            SCREAMING_SNAKE_CASE_      :			int     =       "main"

      return args


def 					__lowerCamelCase			(      )							->   Tuple:
      """simple docstring"""

      transformers.utils.logging.set_verbosity_error()
      SCREAMING_SNAKE_CASE_      :			str     =       parse_args()

      if args.torch_dtype is None:
            SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       AutoConfig.from_pretrained(args.model_name    )
            SCREAMING_SNAKE_CASE_      :			Optional[int]     =       getattr(a__  ,  "torch_dtype"  ,  torch.floataa    )
      else:
            SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       getattr(a__  ,  args.torch_dtype    )
      SCREAMING_SNAKE_CASE_      :			Dict     =       AutoModelForCausalLM if args.is_causal else AutoModelForSeqaSeqLM
      SCREAMING_SNAKE_CASE_      :			Dict     =       {
          "torch_dtype": torch_dtype,
          "revision": args.model_revision,
      }
      if args.disk_offload:
            SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       "tmp_offload"
            SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       True

      SCREAMING_SNAKE_CASE_      :			Optional[int]     =       start_measure()
      SCREAMING_SNAKE_CASE_      :			List[Any]     =       model_cls.from_pretrained(args.model_name  ,  device_map="auto"  ,  **a__    )
      SCREAMING_SNAKE_CASE_      :			int     =       end_measure(a__    )
      log_measures(a__  ,  "Model loading"    )

      SCREAMING_SNAKE_CASE_      :			Dict     =       compute_module_sizes(a__    )
      SCREAMING_SNAKE_CASE_      :			Tuple     =       {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]
      SCREAMING_SNAKE_CASE_      :			Dict     =       "\n".join([F"- {device}: {size // 2**2_0}MiB" for device, size in device_size.items()]    )
      print(F"\nTheoretical use:\n{message}"    )

      SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       AutoTokenizer.from_pretrained(args.tokenizer_name    )

      SCREAMING_SNAKE_CASE_      :			str     =       start_measure()
      SCREAMING_SNAKE_CASE_      :			Tuple     =       []
      SCREAMING_SNAKE_CASE_      :			Any     =       []
      SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       []
      for prompt in PROMPTS:
            SCREAMING_SNAKE_CASE_      :			Optional[int]     =       tokenizer(a__  ,  return_tensors="pt"    ).to(0    )
            SCREAMING_SNAKE_CASE_      :			Optional[int]     =       inputs["input_ids"][0].tolist()
            SCREAMING_SNAKE_CASE_      :			List[str]     =       time.time()
            SCREAMING_SNAKE_CASE_      :			List[Any]     =       model.generate(inputs["input_ids"]    )
            SCREAMING_SNAKE_CASE_      :			List[str]     =       time.time()
            SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       outputs[0].tolist()
            SCREAMING_SNAKE_CASE_      :			Dict     =       len(a__    ) if outputs[: len(a__    )] != tokens else len(a__    ) - len(a__    )
            SCREAMING_SNAKE_CASE_      :			List[Any]     =       after_generate - before_generate

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

      SCREAMING_SNAKE_CASE_      :			List[str]     =       end_measure(a__    )
      log_measures(a__  ,  "Model generation"    )

      SCREAMING_SNAKE_CASE_      :			Dict     =       [gen / tok for gen, tok in zip(a__  ,  a__    )]
      SCREAMING_SNAKE_CASE_      :			Tuple     =       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"    )
      SCREAMING_SNAKE_CASE_      :			List[Any]     =       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()







'''simple docstring'''
import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    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        SCREAMING_SNAKE_CASE(				__A   ,      __A   ,      unittest.TestCase     ):
       snake_case_ :     str    =					StableDiffusionParadigmsPipeline
       snake_case_ :     Optional[Any]    =					TEXT_TO_IMAGE_PARAMS
       snake_case_ :     Tuple    =					TEXT_TO_IMAGE_BATCH_PARAMS
       snake_case_ :     Optional[Any]    =					TEXT_TO_IMAGE_IMAGE_PARAMS
       snake_case_ :     str    =					TEXT_TO_IMAGE_IMAGE_PARAMS

       def     snake_case__			(    self     )    ->					int:

           """simple docstring"""



           torch.manual_seed(0     )
           __lowercase							=						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					,		attention_head_dim=(2, 4)					,		use_linear_projection=lowerCamelCase__					,		)
           __lowercase							=						DDIMParallelScheduler(
               beta_start=0.0_00_85					,		beta_end=0.0_12					,		beta_schedule="""scaled_linear"""					,		clip_sample=lowerCamelCase__					,		set_alpha_to_one=lowerCamelCase__					,		)
           torch.manual_seed(0     )
           __lowercase							=						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     )
           __lowercase							=						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=1000					,		hidden_act="""gelu"""					,		projection_dim=512					,		)
           __lowercase							=						CLIPTextModel(lowerCamelCase__     )
           __lowercase							=						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

           __lowercase							=						{
               """unet""": unet,
               """scheduler""": scheduler,
               """vae""": vae,
               """text_encoder""": text_encoder,
               """tokenizer""": tokenizer,
               """safety_checker""": None,
               """feature_extractor""": None,
           }
           return components

       def     snake_case__			(    self					,		lowerCamelCase__					,		lowerCamelCase__=0     )    ->					Any:

           """simple docstring"""



           if str(lowerCamelCase__     ).startswith("""mps"""     ):
               __lowercase							=						torch.manual_seed(lowerCamelCase__     )
           else:
               __lowercase							=						torch.Generator(device=lowerCamelCase__     ).manual_seed(lowerCamelCase__     )
           __lowercase							=						{
               """prompt""": """a photograph of an astronaut riding a horse""",
               """generator""": generator,
               """num_inference_steps""": 10,
               """guidance_scale""": 6.0,
               """output_type""": """numpy""",
               """parallel""": 3,
               """debug""": True,
           }
           return inputs

       def     snake_case__			(    self     )    ->					List[str]:

           """simple docstring"""



           __lowercase							=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
           __lowercase							=						self.get_dummy_components()
           __lowercase							=						StableDiffusionParadigmsPipeline(**lowerCamelCase__     )
           __lowercase							=						sd_pipe.to(lowerCamelCase__     )
           sd_pipe.set_progress_bar_config(disable=lowerCamelCase__     )

           __lowercase							=						self.get_dummy_inputs(lowerCamelCase__     )
           __lowercase							=						sd_pipe(**lowerCamelCase__     ).images
           __lowercase							=						image[0, -3:, -3:, -1]
           assert image.shape == (1, 64, 64, 3)

           __lowercase							=						np.array([0.47_73, 0.54_17, 0.47_23, 0.49_25, 0.56_31, 0.47_52, 0.52_40, 0.49_35, 0.50_23]     )

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

       def     snake_case__			(    self     )    ->					str:

           """simple docstring"""



           __lowercase							=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
           __lowercase							=						self.get_dummy_components()
           torch.manual_seed(0     )
           __lowercase							=						DDPMParallelScheduler()
           torch.manual_seed(0     )
           __lowercase							=						StableDiffusionParadigmsPipeline(**lowerCamelCase__     )
           __lowercase							=						sd_pipe.to(lowerCamelCase__     )
           sd_pipe.set_progress_bar_config(disable=lowerCamelCase__     )

           __lowercase							=						self.get_dummy_inputs(lowerCamelCase__     )
           __lowercase							=						sd_pipe(**lowerCamelCase__     ).images
           __lowercase							=						image[0, -3:, -3:, -1]
           assert image.shape == (1, 64, 64, 3)

           __lowercase							=						np.array([0.35_73, 0.44_20, 0.49_60, 0.47_99, 0.37_96, 0.38_79, 0.48_19, 0.43_65, 0.44_68]     )

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

       def     snake_case__			(    self     )    ->					Any:

           """simple docstring"""



           super().test_inference_batch_consistent(batch_sizes=[1, 2]     )

       def     snake_case__			(    self     )    ->					Optional[Any]:

           """simple docstring"""



           super().test_inference_batch_single_identical(batch_size=2					,		expected_max_diff=3E-3     )




       def     snake_case__			(    self     )    ->					str:

           """simple docstring"""



           __lowercase							=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
           __lowercase							=						self.get_dummy_components()
           __lowercase							=						StableDiffusionParadigmsPipeline(**lowerCamelCase__     )
           __lowercase							=						sd_pipe.to(lowerCamelCase__     )
           sd_pipe.set_progress_bar_config(disable=lowerCamelCase__     )

           __lowercase							=						self.get_dummy_inputs(lowerCamelCase__     )
           __lowercase							=						"""french fries"""
           __lowercase							=						sd_pipe(**lowerCamelCase__					,		negative_prompt=lowerCamelCase__     )
           __lowercase							=						output.images
           __lowercase							=						image[0, -3:, -3:, -1]

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

           __lowercase							=						np.array([0.47_71, 0.54_20, 0.46_83, 0.49_18, 0.56_36, 0.47_25, 0.52_30, 0.49_23, 0.50_15]     )

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






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

       def     snake_case__			(    self     )    ->					str:

           """simple docstring"""



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

       def     snake_case__			(    self					,		lowerCamelCase__=0     )    ->					int:

           """simple docstring"""



           __lowercase							=						torch.Generator(device=lowerCamelCase__     ).manual_seed(lowerCamelCase__     )
           __lowercase							=						{
               """prompt""": """a photograph of an astronaut riding a horse""",
               """generator""": generator,
               """num_inference_steps""": 10,
               """guidance_scale""": 7.5,
               """output_type""": """numpy""",
               """parallel""": 3,
               """debug""": True,
           }
           return inputs




       def     snake_case__			(    self     )    ->					Optional[Any]:

           """simple docstring"""



           __lowercase							=						"""stabilityai/stable-diffusion-2-base"""
           __lowercase							=						DDIMParallelScheduler.from_pretrained(lowerCamelCase__					,		subfolder="""scheduler"""     )
           __lowercase							=						StableDiffusionParadigmsPipeline.from_pretrained(lowerCamelCase__					,		scheduler=lowerCamelCase__					,		safety_checker=lowerCamelCase__     )
           pipe.to(lowerCamelCase__     )
           pipe.set_progress_bar_config(disable=lowerCamelCase__     )
           pipe.enable_attention_slicing()

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

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

           __lowercase							=						np.array([0.96_22, 0.96_02, 0.97_48, 0.95_91, 0.96_30, 0.96_91, 0.96_61, 0.96_31, 0.97_41]     )

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




