





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


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


a_				=    """
    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 	__snake_case      (	SCREAMING_SNAKE_CASE__  ):





 """simple docstring"""


 def __init__(					self      ,		__lowerCamelCase      ,		__lowerCamelCase						):



     '''simple docstring'''




     super().__init__()

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

     __A					:		Union[str, Any]						=  None

 def       UpperCamelCase__(					self      ,		__lowerCamelCase=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					:		Optional[int]						=  torch.device(F"""cuda:{gpu_id}"""						)

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

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

 def       UpperCamelCase__(					self      ,		__lowerCamelCase=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					:		Optional[int]						=  torch.device(F"""cuda:{gpu_id}"""						)

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

     __A					:		Optional[Any]						=  None
     for cpu_offloaded_model in [self.unet]:
         __A				, __A					:		Dict						=  cpu_offload_with_hook(__lowerCamelCase      ,		__lowerCamelCase      ,		prev_module_hook=__lowerCamelCase						)

     if self.safety_checker is not None:
         __A				, __A					:		Optional[int]						=  cpu_offload_with_hook(self.safety_checker      ,		__lowerCamelCase      ,		prev_module_hook=__lowerCamelCase						)

     # We'll offload the last model manually.
     __A					:		int						=  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(__lowerCamelCase      ,		'''_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      ,		__lowerCamelCase      ,		__lowerCamelCase      ,		__lowerCamelCase      ,		__lowerCamelCase      ,		__lowerCamelCase      ,		__lowerCamelCase      ,		__lowerCamelCase      ,		__lowerCamelCase=None						):



     '''simple docstring'''




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

     # scale the initial noise by the standard deviation required by the scheduler
     __A					:		List[str]						=  latents * self.scheduler.init_noise_sigma
     return latents

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

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

     # Output_type must be 'pil'
     __A					:		Optional[Any]						=  self.numpy_to_pil(__lowerCamelCase						)
     return sample

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



     '''simple docstring'''




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

 def       UpperCamelCase__(					self      ,		__lowerCamelCase      ,		__lowerCamelCase      ,		__lowerCamelCase      ,		__lowerCamelCase      ,		__lowerCamelCase      ,		__lowerCamelCase						):



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







 @torch.no_grad()
 @replace_example_docstring(__lowerCamelCase						)
 def __call__(					self      ,		__lowerCamelCase = 1      ,		__lowerCamelCase = None      ,		__lowerCamelCase = 1      ,		__lowerCamelCase = None      ,		__lowerCamelCase = None      ,		__lowerCamelCase = None      ,		__lowerCamelCase = "pil"      ,		__lowerCamelCase = True      ,		__lowerCamelCase = None      ,		__lowerCamelCase = 1      ,		):



     '''simple docstring'''




     __A					:		Optional[Any]						=  self.unet.config.sample_size
     __A					:		str						=  self._execution_device

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

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

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

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

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

             __A					:		Union[str, Any]						=  self.scheduler.step(__lowerCamelCase      ,		__lowerCamelCase      ,		__lowerCamelCase      ,		generator=__lowerCamelCase						)[0]

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

        # 6. Post-process image sample
     __A					:		Optional[Any]						=  self.postprocess_image(__lowerCamelCase      ,		output_type=__lowerCamelCase						)

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











"""simple docstring"""



import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_video_inputs


if is_torch_available():
       import torch

if is_vision_available():
       from PIL import Image

       from transformers import VideoMAEImageProcessor






class 	__snake_case      (	unittest.TestCase  ):





 """simple docstring"""


 def __init__(					self      ,		__lowerCamelCase      ,		__lowerCamelCase=7      ,		__lowerCamelCase=3      ,		__lowerCamelCase=10      ,		__lowerCamelCase=18      ,		__lowerCamelCase=30      ,		__lowerCamelCase=400      ,		__lowerCamelCase=True      ,		__lowerCamelCase=None      ,		__lowerCamelCase=True      ,		__lowerCamelCase=[0.5, 0.5, 0.5]      ,		__lowerCamelCase=[0.5, 0.5, 0.5]      ,		__lowerCamelCase=None      ,		):



     '''simple docstring'''




     __A					:		Optional[Any]						=  size if size is not None else {'''shortest_edge''': 18}
     __A					:		Optional[int]						=  crop_size if crop_size is not None else {'''height''': 18, '''width''': 18}

     __A					:		Any						=  parent
     __A					:		Union[str, Any]						=  batch_size
     __A					:		Optional[int]						=  num_channels
     __A					:		int						=  num_frames
     __A					:		Optional[int]						=  image_size
     __A					:		Union[str, Any]						=  min_resolution
     __A					:		int						=  max_resolution
     __A					:		Tuple						=  do_resize
     __A					:		str						=  size
     __A					:		Any						=  do_normalize
     __A					:		Union[str, Any]						=  image_mean
     __A					:		List[str]						=  image_std
     __A					:		Union[str, Any]						=  crop_size







 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return {
         "image_mean": self.image_mean,
         "image_std": self.image_std,
         "do_normalize": self.do_normalize,
         "do_resize": self.do_resize,
         "size": self.size,
         "crop_size": self.crop_size,
     }






@require_torch
@require_vision
class 	__snake_case      (	SCREAMING_SNAKE_CASE__							,							unittest.TestCase  ):





 """simple docstring"""


 _lowerCamelCase				 =  VideoMAEImageProcessor if is_vision_available() else None

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		List[str]						=  VideoMAEImageProcessingTester(self						)

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return self.image_processor_tester.prepare_image_processor_dict()

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		Dict						=  self.image_processing_class(**self.image_processor_dict						)
     self.assertTrue(hasattr(__lowerCamelCase      ,		'''image_mean'''						)						)
     self.assertTrue(hasattr(__lowerCamelCase      ,		'''image_std'''						)						)
     self.assertTrue(hasattr(__lowerCamelCase      ,		'''do_normalize'''						)						)
     self.assertTrue(hasattr(__lowerCamelCase      ,		'''do_resize'''						)						)
     self.assertTrue(hasattr(__lowerCamelCase      ,		'''do_center_crop'''						)						)
     self.assertTrue(hasattr(__lowerCamelCase      ,		'''size'''						)						)

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		Optional[int]						=  self.image_processing_class.from_dict(self.image_processor_dict						)
     self.assertEqual(image_processor.size      ,		{'''shortest_edge''': 18}						)
     self.assertEqual(image_processor.crop_size      ,		{'''height''': 18, '''width''': 18}						)

     __A					:		Dict						=  self.image_processing_class.from_dict(self.image_processor_dict      ,		size=42      ,		crop_size=84						)
     self.assertEqual(image_processor.size      ,		{'''shortest_edge''': 42}						)
     self.assertEqual(image_processor.crop_size      ,		{'''height''': 84, '''width''': 84}						)

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     pass

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		List[Any]						=  self.image_processing_class(**self.image_processor_dict						)
     # create random PIL videos
     __A					:		Optional[int]						=  prepare_video_inputs(self.image_processor_tester      ,		equal_resolution=__lowerCamelCase						)
     for video in video_inputs:
         self.assertIsInstance(__lowerCamelCase      ,		__lowerCamelCase						)
         self.assertIsInstance(video[0]      ,		Image.Image						)

     # Test not batched input
     __A					:		Any						=  image_processing(video_inputs[0]      ,		return_tensors='''pt'''						).pixel_values
     self.assertEqual(
         encoded_videos.shape      ,		(
             1,
             self.image_processor_tester.num_frames,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.crop_size['''height'''],
             self.image_processor_tester.crop_size['''width'''],
         )      ,		)

     # Test batched
     __A					:		int						=  image_processing(__lowerCamelCase      ,		return_tensors='''pt'''						).pixel_values
     self.assertEqual(
         encoded_videos.shape      ,		(
             self.image_processor_tester.batch_size,
             self.image_processor_tester.num_frames,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.crop_size['''height'''],
             self.image_processor_tester.crop_size['''width'''],
         )      ,		)

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		List[Any]						=  self.image_processing_class(**self.image_processor_dict						)
     # create random numpy tensors
     __A					:		List[str]						=  prepare_video_inputs(self.image_processor_tester      ,		equal_resolution=__lowerCamelCase      ,		numpify=__lowerCamelCase						)
     for video in video_inputs:
         self.assertIsInstance(__lowerCamelCase      ,		__lowerCamelCase						)
         self.assertIsInstance(video[0]      ,		np.ndarray						)

     # Test not batched input
     __A					:		List[str]						=  image_processing(video_inputs[0]      ,		return_tensors='''pt'''						).pixel_values
     self.assertEqual(
         encoded_videos.shape      ,		(
             1,
             self.image_processor_tester.num_frames,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.crop_size['''height'''],
             self.image_processor_tester.crop_size['''width'''],
         )      ,		)

     # Test batched
     __A					:		int						=  image_processing(__lowerCamelCase      ,		return_tensors='''pt'''						).pixel_values
     self.assertEqual(
         encoded_videos.shape      ,		(
             self.image_processor_tester.batch_size,
             self.image_processor_tester.num_frames,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.crop_size['''height'''],
             self.image_processor_tester.crop_size['''width'''],
         )      ,		)







 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		Optional[int]						=  self.image_processing_class(**self.image_processor_dict						)
     # create random PyTorch tensors
     __A					:		int						=  prepare_video_inputs(self.image_processor_tester      ,		equal_resolution=__lowerCamelCase      ,		torchify=__lowerCamelCase						)
     for video in video_inputs:
         self.assertIsInstance(__lowerCamelCase      ,		__lowerCamelCase						)
         self.assertIsInstance(video[0]      ,		torch.Tensor						)

     # Test not batched input
     __A					:		Optional[int]						=  image_processing(video_inputs[0]      ,		return_tensors='''pt'''						).pixel_values
     self.assertEqual(
         encoded_videos.shape      ,		(
             1,
             self.image_processor_tester.num_frames,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.crop_size['''height'''],
             self.image_processor_tester.crop_size['''width'''],
         )      ,		)

     # Test batched
     __A					:		str						=  image_processing(__lowerCamelCase      ,		return_tensors='''pt'''						).pixel_values
     self.assertEqual(
         encoded_videos.shape      ,		(
             self.image_processor_tester.batch_size,
             self.image_processor_tester.num_frames,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.crop_size['''height'''],
             self.image_processor_tester.crop_size['''width'''],
         )      ,		)




