

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


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


__UpperCAmelCase     					=  '\n    Examples:\n        ```py\n        >>> import torch\n\n        >>> from diffusers import ConsistencyModelPipeline\n\n        >>> device = "cuda"\n        >>> # Load the cd_imagenet64_l2 checkpoint.\n        >>> model_id_or_path = "openai/diffusers-cd_imagenet64_l2"\n        >>> pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)\n        >>> pipe.to(device)\n\n        >>> # Onestep Sampling\n        >>> image = pipe(num_inference_steps=1).images[0]\n        >>> image.save("cd_imagenet64_l2_onestep_sample.png")\n\n        >>> # Onestep sampling, class-conditional image generation\n        >>> # ImageNet-64 class label 145 corresponds to king penguins\n        >>> image = pipe(num_inference_steps=1, class_labels=145).images[0]\n        >>> image.save("cd_imagenet64_l2_onestep_sample_penguin.png")\n\n        >>> # Multistep sampling, class-conditional image generation\n        >>> # Timesteps can be explicitly specified; the particular timesteps below are from the original Github repo:\n        >>> # https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L77\n        >>> image = pipe(num_inference_steps=None, timesteps=[22, 0], class_labels=145).images[0]\n        >>> image.save("cd_imagenet64_l2_multistep_sample_penguin.png")\n        ```\n'





class A__     (			A			):



 """simple docstring"""







 def __init__(		self			:						Dict  ,    A_			:						UNetaDModel  ,    A_			:						CMStochasticIterativeScheduler							):






      '''simple docstring'''




      super().__init__()

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

      _lowerCAmelCase :       List[Any]	   =     None

 def  __magic_name__							(		self			:						int  ,    A_			:						Tuple=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"							)

      _lowerCAmelCase :       int	   =     torch.device(F'''cuda:{gpu_id}'''							)

      if self.device.type != "cpu":
           self.to("cpu"  ,    silence_dtype_warnings=A_							)
           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(A_  ,    A_							)

      if self.safety_checker is not None:
           cpu_offload(self.safety_checker  ,    execution_device=A_  ,    offload_buffers=A_							)

 def  __magic_name__							(		self			:						int  ,    A_			:						Optional[int]=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."							)

      _lowerCAmelCase :       List[Any]	   =     torch.device(F'''cuda:{gpu_id}'''							)

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

      _lowerCAmelCase :       Optional[Any]	   =     None
      for cpu_offloaded_model in [self.unet]:
           _lowerCAmelCase   ,      _lowerCAmelCase :       int	   =     cpu_offload_with_hook(A_  ,    A_  ,    prev_module_hook=A_							)

      if self.safety_checker is not None:
           _lowerCAmelCase   ,      _lowerCAmelCase :       Any	   =     cpu_offload_with_hook(self.safety_checker  ,    A_  ,    prev_module_hook=A_							)

      # We'll offload the last model manually.
      _lowerCAmelCase :       Tuple	   =     hook

 @property
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
 def  __magic_name__							(		self			:						List[str]							):






      '''simple docstring'''




      if not hasattr(self.unet  ,    "_hf_hook"							):
           return self.device
      for module in self.unet.modules():
           if (
               hasattr(A_  ,    "_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  __magic_name__							(		self			:						Union[str, Any]  ,    A_			:						Tuple  ,    A_			:						int  ,    A_			:						str  ,    A_			:						Union[str, Any]  ,    A_			:						Any  ,    A_			:						str  ,    A_			:						Union[str, Any]  ,    A_			:						List[str]=None							):






      '''simple docstring'''




      _lowerCAmelCase :       Optional[Any]	   =     (batch_size, num_channels, height, width)
      if isinstance(A_  ,    A_							) and len(A_							) != batch_size:
           raise ValueError(
               F'''You have passed a list of generators of length {len(A_							)}, 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:
           _lowerCAmelCase :       List[str]	   =     randn_tensor(A_  ,    generator=A_  ,    device=A_  ,    dtype=A_							)
      else:
           _lowerCAmelCase :       Union[str, Any]	   =     latents.to(device=A_  ,    dtype=A_							)

      # scale the initial noise by the standard deviation required by the scheduler
      _lowerCAmelCase :       str	   =     latents * self.scheduler.init_noise_sigma
      return latents

 def  __magic_name__							(		self			:						Any  ,    A_			:						torch.FloatTensor  ,    A_			:						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
      _lowerCAmelCase :       List[str]	   =     (sample / 2 + 0.5).clamp(0  ,    1							)
      if output_type == "pt":
           return sample

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

      # Output_type must be 'pil'
      _lowerCAmelCase :       Optional[int]	   =     self.numpy_to_pil(A_							)
      return sample

 def  __magic_name__							(		self			:						Dict  ,    A_			:						Optional[int]  ,    A_			:						str  ,    A_			:						List[Any]=None							):






      '''simple docstring'''




      if self.unet.config.num_class_embeds is not None:
           if isinstance(A_  ,    A_							):
                _lowerCAmelCase :       Tuple	   =     torch.tensor(A_  ,    dtype=torch.int							)
           elif isinstance(A_  ,    A_							):
                assert batch_size == 1, "Batch size must be 1 if classes is an int"
                _lowerCAmelCase :       List[Any]	   =     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
                _lowerCAmelCase :       Optional[Any]	   =     torch.randint(0  ,    self.unet.config.num_class_embeds  ,    size=(batch_size,)							)
           _lowerCAmelCase :       Union[str, Any]	   =     class_labels.to(A_							)
      else:
           _lowerCAmelCase :       Optional[Any]	   =     None
      return class_labels

 def  __magic_name__							(		self			:						List[str]  ,    A_			:						Any  ,    A_			:						Optional[Any]  ,    A_			:						Union[str, Any]  ,    A_			:						Optional[int]  ,    A_			:						Optional[int]  ,    A_			:						Tuple							):






      '''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:
           _lowerCAmelCase :       List[Any]	   =     (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(A_  ,    A_							) or callback_steps <= 0)
      ):
           raise ValueError(
               F'''`callback_steps` has to be a positive integer but is {callback_steps} of type'''
               F''' {type(A_							)}.'''							)


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






      '''simple docstring'''




      _lowerCAmelCase :       Tuple	   =     self.unet.config.sample_size
      _lowerCAmelCase :       str	   =     self._execution_device

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

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

      # 3. Handle class_labels for class-conditional models
      _lowerCAmelCase :       List[Any]	   =     self.prepare_class_labels(A_  ,    A_  ,    class_labels=A_							)

      # 4. Prepare timesteps
      if timesteps is not None:
           self.scheduler.set_timesteps(timesteps=A_  ,    device=A_							)
           _lowerCAmelCase :       Tuple	   =     self.scheduler.timesteps
           _lowerCAmelCase :       List[Any]	   =     len(A_							)
      else:
           self.scheduler.set_timesteps(A_							)
           _lowerCAmelCase :       Tuple	   =     self.scheduler.timesteps

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

                _lowerCAmelCase :       List[str]	   =     self.scheduler.step(A_  ,    A_  ,    A_  ,    generator=A_							)[0]

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

        # 6. Post-process image sample
      _lowerCAmelCase :       List[str]	   =     self.postprocess_image(A_  ,    output_type=A_							)

      # 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=A_							)









import unittest
from timeit import timeit







def 						_snake_case		(     SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE					)							->   int:

     """simple docstring"""







     _lowerCAmelCase :       List[str]	   =     first_num if first_num >= second_num else second_num
     _lowerCAmelCase :       List[str]	   =     max_num
     while (common_mult % first_num > 0) or (common_mult % second_num > 0):
          common_mult += max_num
     return common_mult







def 						_snake_case		(     SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE					)							->   int:

     """simple docstring"""







     return b if a == 0 else greatest_common_divisor(b % a		,				SCREAMING_SNAKE_CASE					)







def 						_snake_case		(     SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE					)							->   int:

     """simple docstring"""







     return first_num // greatest_common_divisor(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE					) * second_num







def 						_snake_case		(     )							->   str:

     """simple docstring"""







     _lowerCAmelCase :       List[Any]	   =     (
         "from __main__ import least_common_multiple_slow, least_common_multiple_fast"
     )
     print(
         "least_common_multiple_slow():"		,				timeit("least_common_multiple_slow(1000, 999)"		,				setup=SCREAMING_SNAKE_CASE					)		,				)
     print(
         "least_common_multiple_fast():"		,				timeit("least_common_multiple_fast(1000, 999)"		,				setup=SCREAMING_SNAKE_CASE					)		,				)





class A__     (			unittest.TestCase			):



 """simple docstring"""







 _lowercase  :       Union[str, Any]  			=    (
     (10, 20),
     (13, 15),
     (4, 31),
     (10, 42),
     (43, 34),
     (5, 12),
     (12, 25),
     (10, 25),
     (6, 9),
 )
 _lowercase  :       Dict  			=    (20, 195, 124, 210, 1_462, 60, 300, 50, 18)

 def  __magic_name__							(		self			:						int							):






      '''simple docstring'''




      for i, (first_num, second_num) in enumerate(self.test_inputs							):
           _lowerCAmelCase :       Optional[int]	   =     least_common_multiple_slow(A_  ,    A_							)
           _lowerCAmelCase :       Optional[int]	   =     least_common_multiple_fast(A_  ,    A_							)
           with self.subTest(i=A_							):
                self.assertEqual(A_  ,    self.expected_results[i]							)
                self.assertEqual(A_  ,    self.expected_results[i]							)


if __name__ == "__main__":
  benchmark()
  unittest.main()






