






'''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_							      =      logging.get_logger(__name__)  # pylint: disable=invalid-name


lowerCamelCase_							      =      '''
    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      _UpperCAmelCase    (  snake_case_		):




   """simple docstring"""

   def __init__(	self				:			List[Any]   ,   __UpperCAmelCase				:			UNetaDModel   ,   __UpperCAmelCase				:			CMStochasticIterativeScheduler	):





         '''simple docstring'''





         super().__init__()

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

         _A      =			None

   def 		lowerCAmelCase				(	self				:			List[str]   ,   __UpperCAmelCase				:			str=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"	)

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

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

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

   def 		lowerCAmelCase				(	self				:			int   ,   __UpperCAmelCase				:			Union[str, 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."	)

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

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

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

         if self.safety_checker is not None:
               _A  ,			_A      =			cpu_offload_with_hook(self.safety_checker   ,   __UpperCAmelCase   ,   prev_module_hook=__UpperCAmelCase	)

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

   @property
   # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
   def 		lowerCAmelCase				(	self				:			Union[str, Any]	):





         '''simple docstring'''





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





         '''simple docstring'''





         _A      =			(batch_size, num_channels, height, width)
         if isinstance(__UpperCAmelCase   ,   __UpperCAmelCase	) and len(__UpperCAmelCase	) != batch_size:
               raise ValueError(
                   f'''You have passed a list of generators of length {len(__UpperCAmelCase	)}, 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:
               _A      =			randn_tensor(__UpperCAmelCase   ,   generator=__UpperCAmelCase   ,   device=__UpperCAmelCase   ,   dtype=__UpperCAmelCase	)
         else:
               _A      =			latents.to(device=__UpperCAmelCase   ,   dtype=__UpperCAmelCase	)

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

   def 		lowerCAmelCase				(	self				:			Tuple   ,   __UpperCAmelCase				:			torch.FloatTensor   ,   __UpperCAmelCase				:			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
         _A      =			(sample / 2 + 0.5).clamp(0   ,   1	)
         if output_type == "pt":
               return sample

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

         # Output_type must be 'pil'
         _A      =			self.numpy_to_pil(__UpperCAmelCase	)
         return sample

   def 		lowerCAmelCase				(	self				:			str   ,   __UpperCAmelCase				:			Dict   ,   __UpperCAmelCase				:			Optional[Any]   ,   __UpperCAmelCase				:			Dict=None	):





         '''simple docstring'''





         if self.unet.config.num_class_embeds is not None:
               if isinstance(__UpperCAmelCase   ,   __UpperCAmelCase	):
                     _A      =			torch.tensor(__UpperCAmelCase   ,   dtype=torch.int	)
               elif isinstance(__UpperCAmelCase   ,   __UpperCAmelCase	):
                     assert batch_size == 1, "Batch size must be 1 if classes is an int"
                     _A      =			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
                     _A      =			torch.randint(0   ,   self.unet.config.num_class_embeds   ,   size=(batch_size,)	)
               _A      =			class_labels.to(__UpperCAmelCase	)
         else:
               _A      =			None
         return class_labels

   def 		lowerCAmelCase				(	self				:			Optional[Any]   ,   __UpperCAmelCase				:			int   ,   __UpperCAmelCase				:			str   ,   __UpperCAmelCase				:			Optional[int]   ,   __UpperCAmelCase				:			Dict   ,   __UpperCAmelCase				:			Dict   ,   __UpperCAmelCase				:			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:
               _A      =			(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(__UpperCAmelCase   ,   __UpperCAmelCase	) or callback_steps <= 0)
         ):
               raise ValueError(
                   f'''`callback_steps` has to be a positive integer but is {callback_steps} of type'''
                   f''' {type(__UpperCAmelCase	)}.'''	)

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





         '''simple docstring'''





         _A      =			self.unet.config.sample_size
         _A      =			self._execution_device

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

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

         # 3. Handle class_labels for class-conditional models
         _A      =			self.prepare_class_labels(__UpperCAmelCase   ,   __UpperCAmelCase   ,   class_labels=__UpperCAmelCase	)

         # 4. Prepare timesteps
         if timesteps is not None:
               self.scheduler.set_timesteps(timesteps=__UpperCAmelCase   ,   device=__UpperCAmelCase	)
               _A      =			self.scheduler.timesteps
               _A      =			len(__UpperCAmelCase	)
         else:
               self.scheduler.set_timesteps(__UpperCAmelCase	)
               _A      =			self.scheduler.timesteps

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

                     _A      =			self.scheduler.step(__UpperCAmelCase   ,   __UpperCAmelCase   ,   __UpperCAmelCase   ,   generator=__UpperCAmelCase	)[0]

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

        # 6. Post-process image sample
         _A      =			self.postprocess_image(__UpperCAmelCase   ,   output_type=__UpperCAmelCase	)

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













'''simple docstring'''




from __future__ import annotations

from collections.abc import Iterator
from typing import Any

class      _UpperCAmelCase    :




   """simple docstring"""

   def __init__(	self				:			Any   ,   __UpperCAmelCase				:			Any	):





         '''simple docstring'''





         _A      =			data
         _A      =			None

   def __str__(	self				:			str	):





         '''simple docstring'''





         return f'''{self.data}'''





class      _UpperCAmelCase    :




   """simple docstring"""

   def __init__(	self				:			Optional[int]	):





         '''simple docstring'''





         _A      =			None
         _A      =			None

   def __iter__(	self				:			List[Any]	):





         '''simple docstring'''





         _A      =			self.front
         while node:
               yield node.data
               _A      =			node.next

   def __len__(	self				:			Union[str, Any]	):





         '''simple docstring'''





         return len(tuple(iter(self	)	)	)

   def __str__(	self				:			int	):





         '''simple docstring'''





         return " <- ".join(str(__UpperCAmelCase	) for item in self	)

   def 		lowerCAmelCase				(	self				:			List[Any]	):





         '''simple docstring'''





         return len(self	) == 0

   def 		lowerCAmelCase				(	self				:			int   ,   __UpperCAmelCase				:			Any	):





         '''simple docstring'''





         _A      =			Node(__UpperCAmelCase	)
         if self.is_empty():
               _A      =			_A      =			node
         else:
               assert isinstance(self.rear   ,   __UpperCAmelCase	)
               _A      =			node
               _A      =			node

   def 		lowerCAmelCase				(	self				:			str	):





         '''simple docstring'''





         if self.is_empty():
               raise IndexError("dequeue from empty queue"	)
         assert isinstance(self.front   ,   __UpperCAmelCase	)
         _A      =			self.front
         _A      =			self.front.next
         if self.front is None:
               _A      =			None
         return node.data

   def 		lowerCAmelCase				(	self				:			int	):





         '''simple docstring'''





         _A      =			_A      =			None


if __name__ == "__main__":
    from doctest import testmod

    testmod()





