





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


__UpperCamelCase			:		Optional[Any]		    =   logging.get_logger(__name__)  # pylint: disable=invalid-name


__UpperCamelCase			:		Optional[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 	a				(					a__   ):



 def __init__(			self							,    _snake_case							,    _snake_case			):







  """simple docstring"""






  super().__init__()

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

  lowerCAmelCase  							=      None



 def 					UpperCamelCase__   (			self							,    _snake_case=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  							=      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 					UpperCamelCase__   (			self							,    _snake_case=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  							=      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)

  lowerCAmelCase  							=      None
  for cpu_offloaded_model in [self.unet]:
   lowerCAmelCase							,lowerCAmelCase  							=      cpu_offload_with_hook(_snake_case							,    _snake_case							,    prev_module_hook=_snake_case			)

  if self.safety_checker is not None:
   lowerCAmelCase							,lowerCAmelCase  							=      cpu_offload_with_hook(self.safety_checker							,    _snake_case							,    prev_module_hook=_snake_case			)

  # We'll offload the last model manually.
  lowerCAmelCase  							=      hook



 @property
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
 def 					UpperCamelCase__   (			self			):







  """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 					UpperCamelCase__   (			self							,    _snake_case							,    _snake_case							,    _snake_case							,    _snake_case							,    _snake_case							,    _snake_case							,    _snake_case							,    _snake_case=None			):







  """simple docstring"""






  lowerCAmelCase  							=      (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:
   lowerCAmelCase  							=      randn_tensor(_snake_case							,    generator=_snake_case							,    device=_snake_case							,    dtype=_snake_case			)
  else:
   lowerCAmelCase  							=      latents.to(device=_snake_case							,    dtype=_snake_case			)

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



 def 					UpperCamelCase__   (			self							,    _snake_case							,    _snake_case = "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  							=      (sample / 2 + 0.5).clamp(0							,    1			)
  if output_type == "pt":
   return sample

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

  # Output_type must be 'pil'
  lowerCAmelCase  							=      self.numpy_to_pil(_snake_case			)
  return sample



 def 					UpperCamelCase__   (			self							,    _snake_case							,    _snake_case							,    _snake_case=None			):







  """simple docstring"""






  if self.unet.config.num_class_embeds is not None:
   if isinstance(_snake_case							,    _snake_case			):
    lowerCAmelCase  							=      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"
    lowerCAmelCase  							=      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  							=      torch.randint(0							,    self.unet.config.num_class_embeds							,    size=(batch_size,)			)
   lowerCAmelCase  							=      class_labels.to(_snake_case			)
  else:
   lowerCAmelCase  							=      None
  return class_labels



 def 					UpperCamelCase__   (			self							,    _snake_case							,    _snake_case							,    _snake_case							,    _snake_case							,    _snake_case							,    _snake_case			):







  """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  							=      (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							,    _snake_case = 1							,    _snake_case = None							,    _snake_case = 1							,    _snake_case = None							,    _snake_case = None							,    _snake_case = None							,    _snake_case = "pil"							,    _snake_case = True							,    _snake_case = None							,    _snake_case = 1							,    ):







  """simple docstring"""






  lowerCAmelCase  							=      self.unet.config.sample_size
  lowerCAmelCase  							=      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)
  lowerCAmelCase  							=      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
  lowerCAmelCase  							=      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			)
   lowerCAmelCase  							=      self.scheduler.timesteps
   lowerCAmelCase  							=      len(_snake_case			)
  else:
   self.scheduler.set_timesteps(_snake_case			)
   lowerCAmelCase  							=      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			):
    lowerCAmelCase  							=      self.scheduler.scale_model_input(_snake_case							,    _snake_case			)
    lowerCAmelCase  							=      self.unet(_snake_case							,    _snake_case							,    class_labels=_snake_case							,    return_dict=_snake_case			)[0]

    lowerCAmelCase  							=      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
  lowerCAmelCase  							=      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"""
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


__UpperCamelCase			:		List[Any]		    =   {'''configuration_yoso''': ['''YOSO_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''YosoConfig''']}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __UpperCamelCase			:		Dict		    =   [
           '''YOSO_PRETRAINED_MODEL_ARCHIVE_LIST''',
           '''YosoForMaskedLM''',
           '''YosoForMultipleChoice''',
           '''YosoForQuestionAnswering''',
           '''YosoForSequenceClassification''',
           '''YosoForTokenClassification''',
           '''YosoLayer''',
           '''YosoModel''',
           '''YosoPreTrainedModel''',
       ]


if TYPE_CHECKING:
       from .configuration_yoso import YOSO_PRETRAINED_CONFIG_ARCHIVE_MAP, YosoConfig

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_yoso import (
                  YOSO_PRETRAINED_MODEL_ARCHIVE_LIST,
                  YosoForMaskedLM,
                  YosoForMultipleChoice,
                  YosoForQuestionAnswering,
                  YosoForSequenceClassification,
                  YosoForTokenClassification,
                  YosoLayer,
                  YosoModel,
                  YosoPreTrainedModel,
              )


else:
       import sys

       __UpperCamelCase			:		Any		    =   _LazyModule(__name__, globals()['''__file__'''], _import_structure)





