






"""simple docstring"""



from typing import Callable, List, Optional, Union

import torch

from ...models import UNetaDModel
from ...schedulers import CMStochasticIterativeScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


_lowerCAmelCase  :  str   =			logging.get_logger(__name__)  # pylint: disable=invalid-name


_lowerCAmelCase  :  Any   =			"""
    Examples:
        ```py
        >>> import torch

        >>> from diffusers import ConsistencyModelPipeline

        >>> device = \"cuda\"
        >>> # Load the cd_imagenet64_l2 checkpoint.
        >>> model_id_or_path = \"openai/diffusers-cd_imagenet64_l2\"
        >>> pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
        >>> pipe.to(device)

        >>> # Onestep Sampling
        >>> image = pipe(num_inference_steps=1).images[0]
        >>> image.save(\"cd_imagenet64_l2_onestep_sample.png\")

        >>> # Onestep sampling, class-conditional image generation
        >>> # ImageNet-64 class label 145 corresponds to king penguins
        >>> image = pipe(num_inference_steps=1, class_labels=145).images[0]
        >>> image.save(\"cd_imagenet64_l2_onestep_sample_penguin.png\")

        >>> # Multistep sampling, class-conditional image generation
        >>> # Timesteps can be explicitly specified; the particular timesteps below are from the original Github repo:
        >>> # https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L77
        >>> image = pipe(num_inference_steps=None, timesteps=[22, 0], class_labels=145).images[0]
        >>> image.save(\"cd_imagenet64_l2_multistep_sample_penguin.png\")
        ```
"""




class  lowerCAmelCase__ (						__magic_name__      ):
   def __init__(		self			:  Any       , snake_case__			:  UNetaDModel       , snake_case__			:  CMStochasticIterativeScheduler       ):


        '''simple docstring'''
        super().__init__()

        self.register_modules(
            unet=snake_case__       , scheduler=snake_case__       , )

        UpperCAmelCase__	:						Any       =       None
   def 			__a       (		self			:  str       , snake_case__			:  Union[str, Any]=0       ):


        '''simple docstring'''
        if is_accelerate_available() and is_accelerate_version(">="       , "0.14.0"       ):
             from accelerate import cpu_offload
        else:
             raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher"       )

        UpperCAmelCase__	:						str       =       torch.device(f'cuda:{gpu_id}'       )

        if self.device.type != "cpu":
             self.to("cpu"       , silence_dtype_warnings=snake_case__       )
             torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

        for cpu_offloaded_model in [self.unet]:
             cpu_offload(snake_case__       , snake_case__       )

        if self.safety_checker is not None:
             cpu_offload(self.safety_checker       , execution_device=snake_case__       , offload_buffers=snake_case__       )
   def 			__a       (		self			:  Optional[Any]       , snake_case__			:  Any=0       ):


        '''simple docstring'''
        if is_accelerate_available() and is_accelerate_version(">="       , "0.17.0.dev0"       ):
             from accelerate import cpu_offload_with_hook
        else:
             raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher."       )

        UpperCAmelCase__	:						str       =       torch.device(f'cuda:{gpu_id}'       )

        if self.device.type != "cpu":
             self.to("cpu"       , silence_dtype_warnings=snake_case__       )
             torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

        UpperCAmelCase__	:						List[Any]       =       None
        for cpu_offloaded_model in [self.unet]:
             UpperCAmelCase__       ,							UpperCAmelCase__	:						int       =       cpu_offload_with_hook(snake_case__       , snake_case__       , prev_module_hook=snake_case__       )

        if self.safety_checker is not None:
             UpperCAmelCase__       ,							UpperCAmelCase__	:						str       =       cpu_offload_with_hook(self.safety_checker       , snake_case__       , prev_module_hook=snake_case__       )

        # We'll offload the last model manually.
        UpperCAmelCase__	:						Dict       =       hook
   @property
   # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
   def 			__a       (		self			:  Union[str, Any]       ):


        '''simple docstring'''
        if not hasattr(self.unet       , "_hf_hook"       ):
             return self.device
        for module in self.unet.modules():
             if (
                 hasattr(snake_case__       , "_hf_hook"       )
                 and hasattr(module._hf_hook       , "execution_device"       )
                 and module._hf_hook.execution_device is not None
             ):
                  return torch.device(module._hf_hook.execution_device       )
        return self.device
   def 			__a       (		self			:  Union[str, Any]       , snake_case__			:  str       , snake_case__			:  Tuple       , snake_case__			:  Any       , snake_case__			:  Optional[int]       , snake_case__			:  Dict       , snake_case__			:  str       , snake_case__			:  List[Any]       , snake_case__			:  List[str]=None       ):


        '''simple docstring'''
        UpperCAmelCase__	:						Dict       =       (batch_size, num_channels, height, width)
        if isinstance(snake_case__       , snake_case__       ) and len(snake_case__       ) != batch_size:
             raise ValueError(
                 f'You have passed a list of generators of length {len(snake_case__       )}, but requested an effective batch'
                 f' size of {batch_size}. Make sure the batch size matches the length of the generators.'       )

        if latents is None:
             UpperCAmelCase__	:						Any       =       randn_tensor(snake_case__       , generator=snake_case__       , device=snake_case__       , dtype=snake_case__       )
        else:
             UpperCAmelCase__	:						Dict       =       latents.to(device=snake_case__       , dtype=snake_case__       )

        # scale the initial noise by the standard deviation required by the scheduler
        UpperCAmelCase__	:						List[Any]       =       latents * self.scheduler.init_noise_sigma
        return latents
   def 			__a       (		self			:  Optional[Any]       , snake_case__			:  torch.FloatTensor       , snake_case__			:  str = "pil"       ):


        '''simple docstring'''
        if output_type not in ["pt", "np", "pil"]:
             raise ValueError(
                 f'output_type={output_type} is not supported. Make sure to choose one of [\'pt\', \'np\', or \'pil\']'       )

        # Equivalent to diffusers.VaeImageProcessor.denormalize
        UpperCAmelCase__	:						List[str]       =       (sample / 2 + 0.5).clamp(0       , 1       )
        if output_type == "pt":
             return sample

        # Equivalent to diffusers.VaeImageProcessor.pt_to_numpy
        UpperCAmelCase__	:						Optional[int]       =       sample.cpu().permute(0       , 2       , 3       , 1       ).numpy()
        if output_type == "np":
             return sample

        # Output_type must be 'pil'
        UpperCAmelCase__	:						Tuple       =       self.numpy_to_pil(snake_case__       )
        return sample
   def 			__a       (		self			:  Dict       , snake_case__			:  List[str]       , snake_case__			:  Optional[Any]       , snake_case__			:  List[Any]=None       ):


        '''simple docstring'''
        if self.unet.config.num_class_embeds is not None:
             if isinstance(snake_case__       , snake_case__       ):
                  UpperCAmelCase__	:						List[Any]       =       torch.tensor(snake_case__       , dtype=torch.int       )
             elif isinstance(snake_case__       , snake_case__       ):
                  assert batch_size == 1, "Batch size must be 1 if classes is an int"
                  UpperCAmelCase__	:						int       =       torch.tensor([class_labels]       , dtype=torch.int       )
             elif class_labels is None:
                  # Randomly generate batch_size class labels
                  # TODO: should use generator here? int analogue of randn_tensor is not exposed in ...utils
                  UpperCAmelCase__	:						str       =       torch.randint(0       , self.unet.config.num_class_embeds       , size=(batch_size,)       )
             UpperCAmelCase__	:						Optional[int]       =       class_labels.to(snake_case__       )
        else:
             UpperCAmelCase__	:						Optional[int]       =       None
        return class_labels
   def 			__a       (		self			:  Optional[Any]       , snake_case__			:  int       , snake_case__			:  int       , snake_case__			:  List[str]       , snake_case__			:  Optional[Any]       , snake_case__			:  Optional[int]       , snake_case__			:  Optional[Any]       ):


        '''simple docstring'''
        if num_inference_steps is None and timesteps is None:
             raise ValueError("Exactly one of `num_inference_steps` or `timesteps` must be supplied."       )

        if num_inference_steps is not None and timesteps is not None:
             logger.warning(
                 f'Both `num_inference_steps`: {num_inference_steps} and `timesteps`: {timesteps} are supplied;'
                 " `timesteps` will be used over `num_inference_steps`."       )

        if latents is not None:
             UpperCAmelCase__	:						List[str]       =       (batch_size, 3, img_size, img_size)
             if latents.shape != expected_shape:
                  raise ValueError(f'The shape of latents is {latents.shape} but is expected to be {expected_shape}.'       )

        if (callback_steps is None) or (
            callback_steps is not None and (not isinstance(snake_case__       , snake_case__       ) or callback_steps <= 0)
        ):
             raise ValueError(
                 f'`callback_steps` has to be a positive integer but is {callback_steps} of type'
                 f' {type(snake_case__       )}.'       )




   @torch.no_grad()
   @replace_example_docstring(snake_case__       )
   def __call__(		self			:  Union[str, Any]       , snake_case__			:  int = 1       , snake_case__			:  Optional[Union[torch.Tensor, List[int], int]] = None       , snake_case__			:  int = 1       , snake_case__			:  List[int] = None       , snake_case__			:  Optional[Union[torch.Generator, List[torch.Generator]]] = None       , snake_case__			:  Optional[torch.FloatTensor] = None       , snake_case__			:  Optional[str] = "pil"       , snake_case__			:  bool = True       , snake_case__			:  Optional[Callable[[int, int, torch.FloatTensor], None]] = None       , snake_case__			:  int = 1       , ):


        '''simple docstring'''
        UpperCAmelCase__	:						Optional[int]       =       self.unet.config.sample_size
        UpperCAmelCase__	:						Tuple       =       self._execution_device

        # 1. Check inputs
        self.check_inputs(snake_case__       , snake_case__       , snake_case__       , snake_case__       , snake_case__       , snake_case__       )

        # 2. Prepare image latents
        # Sample image latents x_0 ~ N(0, sigma_0^2 * I)
        UpperCAmelCase__	:						Dict       =       self.prepare_latents(
            batch_size=snake_case__       , num_channels=self.unet.config.in_channels       , height=snake_case__       , width=snake_case__       , dtype=self.unet.dtype       , device=snake_case__       , generator=snake_case__       , latents=snake_case__       , )

        # 3. Handle class_labels for class-conditional models
        UpperCAmelCase__	:						Union[str, Any]       =       self.prepare_class_labels(snake_case__       , snake_case__       , class_labels=snake_case__       )

        # 4. Prepare timesteps
        if timesteps is not None:
             self.scheduler.set_timesteps(timesteps=snake_case__       , device=snake_case__       )
             UpperCAmelCase__	:						Union[str, Any]       =       self.scheduler.timesteps
             UpperCAmelCase__	:						Any       =       len(snake_case__       )
        else:
             self.scheduler.set_timesteps(snake_case__       )
             UpperCAmelCase__	:						Optional[int]       =       self.scheduler.timesteps

        # 5. Denoising loop
        # Multistep sampling: implements Algorithm 1 in the paper
        with self.progress_bar(total=snake_case__       ) as progress_bar:
             for i, t in enumerate(snake_case__       ):
                  UpperCAmelCase__	:						Union[str, Any]       =       self.scheduler.scale_model_input(snake_case__       , snake_case__       )
                  UpperCAmelCase__	:						Tuple       =       self.unet(snake_case__       , snake_case__       , class_labels=snake_case__       , return_dict=snake_case__       )[0]

                  UpperCAmelCase__	:						List[str]       =       self.scheduler.step(snake_case__       , snake_case__       , snake_case__       , generator=snake_case__       )[0]

                  # call the callback, if provided
                  progress_bar.update()
                  if callback is not None and i % callback_steps == 0:
                       callback(snake_case__       , snake_case__       , snake_case__       )

        # 6. Post-process image sample
        UpperCAmelCase__	:						Tuple       =       self.postprocess_image(snake_case__       , output_type=snake_case__       )

        # Offload last model to CPU
        if hasattr(self       , "final_offload_hook"       ) and self.final_offload_hook is not None:
             self.final_offload_hook.offload()

        if not return_dict:
             return (image,)

        return ImagePipelineOutput(images=snake_case__       )










"""simple docstring"""



#
# Sort large text files in a minimum amount of memory
#
import argparse
import os




class  lowerCAmelCase__ :
   SCREAMING_SNAKE_CASE_         ='''block_{0}.dat'''
   def __init__(		self			:  Tuple       , snake_case__			:  List[str]       ):


        '''simple docstring'''
        UpperCAmelCase__	:						List[Any]       =       filename
        UpperCAmelCase__	:						List[Any]       =       []
   def 			__a       (		self			:  Optional[Any]       , snake_case__			:  Union[str, Any]       , snake_case__			:  Union[str, Any]       ):


        '''simple docstring'''
        UpperCAmelCase__	:						Dict       =       self.BLOCK_FILENAME_FORMAT.format(snake_case__       )
        with open(snake_case__       , "w"       ) as file:
             file.write(snake_case__       )
        self.block_filenames.append(snake_case__       )
   def 			__a       (		self			:  List[str]       ):


        '''simple docstring'''
        return self.block_filenames
   def 			__a       (		self			:  Optional[Any]       , snake_case__			:  List[str]       , snake_case__			:  str=None       ):


        '''simple docstring'''
        UpperCAmelCase__	:						int       =       0
        with open(self.filename       ) as file:
             while True:
                  UpperCAmelCase__	:						Union[str, Any]       =       file.readlines(snake_case__       )

                  if lines == []:
                       break

                  if sort_key is None:
                       lines.sort()
                  else:
                       lines.sort(key=snake_case__       )

                  self.write_block("".join(snake_case__       )       , snake_case__       )
                  i += 1




   def 			__a       (		self			:  Optional[Any]       ):


        '''simple docstring'''
        map(os.remove       , self.block_filenames       )




class  lowerCAmelCase__ :
   def 			__a       (		self			:  Union[str, Any]       , snake_case__			:  Optional[Any]       ):


        '''simple docstring'''
        UpperCAmelCase__	:						Any       =       -1
        UpperCAmelCase__	:						Tuple       =       None

        for i in range(len(snake_case__       )       ):
             if min_str is None or choices[i] < min_str:
                  UpperCAmelCase__	:						List[str]       =       i

        return min_index




class  lowerCAmelCase__ :
   def __init__(		self			:  Dict       , snake_case__			:  Optional[Any]       ):


        '''simple docstring'''
        UpperCAmelCase__	:						Optional[Any]       =       files
        UpperCAmelCase__	:						str       =       set()
        UpperCAmelCase__	:						Any       =       len(snake_case__       )
        UpperCAmelCase__	:						Tuple       =       {i: None for i in range(self.num_buffers       )}
   def 			__a       (		self			:  List[str]       ):


        '''simple docstring'''
        return {
            i: self.buffers[i] for i in range(self.num_buffers       ) if i not in self.empty
        }
   def 			__a       (		self			:  str       ):


        '''simple docstring'''
        for i in range(self.num_buffers       ):
             if self.buffers[i] is None and i not in self.empty:
                  UpperCAmelCase__	:						Tuple       =       self.files[i].readline()

                  if self.buffers[i] == "":
                       self.empty.add(snake_case__       )
                       self.files[i].close()

        if len(self.empty       ) == self.num_buffers:
             return False

        return True




   def 			__a       (		self			:  List[str]       , snake_case__			:  Tuple       ):


        '''simple docstring'''
        UpperCAmelCase__	:						Any       =       self.buffers[index]
        UpperCAmelCase__	:						Union[str, Any]       =       None

        return value




class  lowerCAmelCase__ :
   def __init__(		self			:  Optional[Any]       , snake_case__			:  int       ):


        '''simple docstring'''
        UpperCAmelCase__	:						str       =       merge_strategy
   def 			__a       (		self			:  Dict       , snake_case__			:  Dict       , snake_case__			:  Optional[int]       , snake_case__			:  Optional[int]       ):


        '''simple docstring'''
        UpperCAmelCase__	:						Dict       =       FilesArray(self.get_file_handles(snake_case__       , snake_case__       )       )
        with open(snake_case__       , "w"       , snake_case__       ) as outfile:
             while buffers.refresh():
                  UpperCAmelCase__	:						Any       =       self.merge_strategy.select(buffers.get_dict()       )
                  outfile.write(buffers.unshift(snake_case__       )       )




   def 			__a       (		self			:  Tuple       , snake_case__			:  Dict       , snake_case__			:  List[str]       ):


        '''simple docstring'''
        UpperCAmelCase__	:						str       =       {}

        for i in range(len(snake_case__       )       ):
             UpperCAmelCase__	:						List[Any]       =       open(filenames[i]       , "r"       , snake_case__       )  # noqa: UP015

        return files




class  lowerCAmelCase__ :
   def __init__(		self			:  str       , snake_case__			:  Optional[Any]       ):


        '''simple docstring'''
        UpperCAmelCase__	:						List[str]       =       block_size
   def 			__a       (		self			:  Optional[Any]       , snake_case__			:  Optional[int]       , snake_case__			:  Any=None       ):


        '''simple docstring'''
        UpperCAmelCase__	:						Union[str, Any]       =       self.get_number_blocks(snake_case__       , self.block_size       )
        UpperCAmelCase__	:						Union[str, Any]       =       FileSplitter(snake_case__       )
        splitter.split(self.block_size       , snake_case__       )

        UpperCAmelCase__	:						int       =       FileMerger(NWayMerge()       )
        UpperCAmelCase__	:						int       =       self.block_size / (num_blocks + 1)
        merger.merge(splitter.get_block_filenames()       , filename + ".out"       , snake_case__       )

        splitter.cleanup()




   def 			__a       (		self			:  Union[str, Any]       , snake_case__			:  Optional[int]       , snake_case__			:  Optional[Any]       ):


        '''simple docstring'''
        return (os.stat(snake_case__       ).st_size / block_size) + 1




def   SCREAMING_SNAKE_CASE__					(     snake_case					:   Tuple							)->		Union[str, Any]:






     '''simple docstring'''



     if string[-1].lower() == "k":
          return int(string[:-1]							) * 1024
     elif string[-1].lower() == "m":
          return int(string[:-1]							) * 1024 * 1024
     elif string[-1].lower() == "g":
          return int(string[:-1]							) * 1024 * 1024 * 1024
     else:
          return int(snake_case							)


def   SCREAMING_SNAKE_CASE__					(     )->		Any:






     '''simple docstring'''



     UpperCAmelCase__	:						Any       =       argparse.ArgumentParser()
     parser.add_argument(
         "-m" ,				"--mem" ,				help="amount of memory to use for sorting" ,				default="100M"							)
     parser.add_argument(
         "filename" ,				metavar="<filename>" ,				nargs=1 ,				help="name of file to sort"							)
     UpperCAmelCase__	:						int       =       parser.parse_args()

     UpperCAmelCase__	:						Optional[int]       =       ExternalSort(parse_memory(args.mem							)							)
     sorter.sort(args.filename[0]							)


if __name__ == "__main__":
 main()


