


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


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


__snake_case     =		'''
    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	(_a		):







       def __init__(	self:				Optional[Any],A_:				UNetaDModel,A_:				CMStochasticIterativeScheduler			):






            '''simple docstring'''




            super().__init__()

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

            __UpperCamelCase				=	None







       def 						snake_case_					(	self:				Optional[int],A_:				Dict=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				=	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 						snake_case_					(	self:				List[str],A_:				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				=	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)

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

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

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







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






            '''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 						snake_case_					(	self:				Any,A_:				Dict,A_:				List[str],A_:				Any,A_:				Any,A_:				Optional[Any],A_:				List[str],A_:				List[Any],A_:				int=None			):






            '''simple docstring'''




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

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







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

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

            # Output_type must be 'pil'
            __UpperCamelCase				=	self.numpy_to_pil(A_			)
            return sample







       def 						snake_case_					(	self:				Optional[Any],A_:				Optional[Any],A_:				Optional[int],A_:				str=None			):






            '''simple docstring'''




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







       def 						snake_case_					(	self:				Tuple,A_:				Dict,A_:				Union[str, Any],A_:				str,A_:				Any,A_:				Optional[int],A_:				Union[str, 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				=	(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:				Dict,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'''




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

                      __UpperCamelCase				=	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
            __UpperCamelCase				=	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 inspect
import unittest

from transformers import ResNetConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch
    from torch import nn

    from transformers import ResNetBackbone, ResNetForImageClassification, ResNetModel
    from transformers.models.resnet.modeling_resnet import RESNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import AutoImageProcessor



class 					__lowerCamelCase	:







       def __init__(	self:				Dict,A_:				Optional[Any],A_:				str=3,A_:				Any=32,A_:				Optional[Any]=3,A_:				Optional[Any]=10,A_:				str=[10, 20, 30, 40],A_:				Optional[int]=[1, 1, 2, 1],A_:				Dict=True,A_:				Any=True,A_:				Dict="relu",A_:				Union[str, Any]=3,A_:				Optional[int]=None,A_:				Optional[int]=["stage2", "stage3", "stage4"],A_:				Tuple=[2, 3, 4],):






            '''simple docstring'''




            __UpperCamelCase				=	parent
            __UpperCamelCase				=	batch_size
            __UpperCamelCase				=	image_size
            __UpperCamelCase				=	num_channels
            __UpperCamelCase				=	embeddings_size
            __UpperCamelCase				=	hidden_sizes
            __UpperCamelCase				=	depths
            __UpperCamelCase				=	is_training
            __UpperCamelCase				=	use_labels
            __UpperCamelCase				=	hidden_act
            __UpperCamelCase				=	num_labels
            __UpperCamelCase				=	scope
            __UpperCamelCase				=	len(A_			)
            __UpperCamelCase				=	out_features
            __UpperCamelCase				=	out_indices







       def 						snake_case_					(	self:				str			):






            '''simple docstring'''




            __UpperCamelCase				=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

            __UpperCamelCase				=	None
            if self.use_labels:
                 __UpperCamelCase				=	ids_tensor([self.batch_size],self.num_labels			)

            __UpperCamelCase				=	self.get_config()

            return config, pixel_values, labels







       def 						snake_case_					(	self:				Optional[Any]			):






            '''simple docstring'''




            return ResNetConfig(
                num_channels=self.num_channels,embeddings_size=self.embeddings_size,hidden_sizes=self.hidden_sizes,depths=self.depths,hidden_act=self.hidden_act,num_labels=self.num_labels,out_features=self.out_features,out_indices=self.out_indices,)







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






            '''simple docstring'''




            __UpperCamelCase				=	ResNetModel(config=A_			)
            model.to(A_			)
            model.eval()
            __UpperCamelCase				=	model(A_			)
            # expected last hidden states: B, C, H // 32, W // 32
            self.parent.assertEqual(
                result.last_hidden_state.shape,(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32),)







       def 						snake_case_					(	self:				int,A_:				int,A_:				List[str],A_:				Optional[Any]			):






            '''simple docstring'''




            __UpperCamelCase				=	self.num_labels
            __UpperCamelCase				=	ResNetForImageClassification(A_			)
            model.to(A_			)
            model.eval()
            __UpperCamelCase				=	model(A_,labels=A_			)
            self.parent.assertEqual(result.logits.shape,(self.batch_size, self.num_labels)			)







       def 						snake_case_					(	self:				Optional[Any],A_:				Any,A_:				str,A_:				Dict			):






            '''simple docstring'''




            __UpperCamelCase				=	ResNetBackbone(config=A_			)
            model.to(A_			)
            model.eval()
            __UpperCamelCase				=	model(A_			)

            # verify feature maps
            self.parent.assertEqual(len(result.feature_maps			),len(config.out_features			)			)
            self.parent.assertListEqual(list(result.feature_maps[0].shape			),[self.batch_size, self.hidden_sizes[1], 4, 4]			)

            # verify channels
            self.parent.assertEqual(len(model.channels			),len(config.out_features			)			)
            self.parent.assertListEqual(model.channels,config.hidden_sizes[1:]			)

            # verify backbone works with out_features=None
            __UpperCamelCase				=	None
            __UpperCamelCase				=	ResNetBackbone(config=A_			)
            model.to(A_			)
            model.eval()
            __UpperCamelCase				=	model(A_			)

            # verify feature maps
            self.parent.assertEqual(len(result.feature_maps			),1			)
            self.parent.assertListEqual(list(result.feature_maps[0].shape			),[self.batch_size, self.hidden_sizes[-1], 1, 1]			)

            # verify channels
            self.parent.assertEqual(len(model.channels			),1			)
            self.parent.assertListEqual(model.channels,[config.hidden_sizes[-1]]			)







       def 						snake_case_					(	self:				int			):






            '''simple docstring'''




            __UpperCamelCase				=	self.prepare_config_and_inputs()
            __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase				=	config_and_inputs
            __UpperCamelCase				=	{'pixel_values': pixel_values}
            return config, inputs_dict



@require_torch
class 					__lowerCamelCase	(_a							,     _a							,     unittest.TestCase		):
       _lowercase    							=      (
           (
               ResNetModel,
               ResNetForImageClassification,
               ResNetBackbone,
           )
           if is_torch_available()
           else ()
       )
       _lowercase    							=      (
           {"""feature-extraction""": ResNetModel, """image-classification""": ResNetForImageClassification}
           if is_torch_available()
           else {}
       )

       _lowercase    							=      True
       _lowercase    							=      False
       _lowercase    							=      False
       _lowercase    							=      False
       _lowercase    							=      False







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






            '''simple docstring'''




            __UpperCamelCase				=	ResNetModelTester(self			)
            __UpperCamelCase				=	ConfigTester(self,config_class=A_,has_text_modality=A_			)







       def 						snake_case_					(	self:				Optional[int]			):






            '''simple docstring'''




            self.create_and_test_config_common_properties()
            self.config_tester.create_and_test_config_to_json_string()
            self.config_tester.create_and_test_config_to_json_file()
            self.config_tester.create_and_test_config_from_and_save_pretrained()
            self.config_tester.create_and_test_config_with_num_labels()
            self.config_tester.check_config_can_be_init_without_params()
            self.config_tester.check_config_arguments_init()







       def 						snake_case_					(	self:				Optional[int]			):






            '''simple docstring'''




            return







       @unittest.skip(reason='ResNet does not use inputs_embeds'			)
       def 						snake_case_					(	self:				List[str]			):






            '''simple docstring'''




            pass







       @unittest.skip(reason='ResNet does not support input and output embeddings'			)
       def 						snake_case_					(	self:				Union[str, Any]			):






            '''simple docstring'''




            pass







       def 						snake_case_					(	self:				int			):






            '''simple docstring'''




            __UpperCamelCase,  __UpperCamelCase				=	self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                 __UpperCamelCase				=	model_class(A_			)
                 __UpperCamelCase				=	inspect.signature(model.forward			)
                 # signature.parameters is an OrderedDict => so arg_names order is deterministic
                 __UpperCamelCase				=	[*signature.parameters.keys()]

                 __UpperCamelCase				=	['pixel_values']
                 self.assertListEqual(arg_names[:1],A_			)







       def 						snake_case_					(	self:				List[Any]			):






            '''simple docstring'''




            __UpperCamelCase				=	self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*A_			)







       def 						snake_case_					(	self:				Dict			):






            '''simple docstring'''




            __UpperCamelCase				=	self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_backbone(*A_			)







       def 						snake_case_					(	self:				List[Any]			):






            '''simple docstring'''




            __UpperCamelCase,  __UpperCamelCase				=	self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                 __UpperCamelCase				=	model_class(config=A_			)
                 for name, module in model.named_modules():
                      if isinstance(A_,(nn.BatchNormad, nn.GroupNorm)			):
                           self.assertTrue(
                               torch.all(module.weight == 1			),msg=F'''Parameter {name} of model {model_class} seems not properly initialized''',)
                           self.assertTrue(
                               torch.all(module.bias == 0			),msg=F'''Parameter {name} of model {model_class} seems not properly initialized''',)







       def 						snake_case_					(	self:				Dict			):






            '''simple docstring'''




            def check_hidden_states_output(A_:				Tuple,A_:				Tuple,A_:				Optional[int]			):
                 __UpperCamelCase				=	model_class(A_			)
                 model.to(A_			)
                 model.eval()

                 with torch.no_grad():
                      __UpperCamelCase				=	model(**self._prepare_for_class(A_,A_			)			)

                 __UpperCamelCase				=	outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                 __UpperCamelCase				=	self.model_tester.num_stages
                 self.assertEqual(len(A_			),expected_num_stages + 1			)

                 # ResNet's feature maps are of shape (batch_size, num_channels, height, width)
                 self.assertListEqual(
                     list(hidden_states[0].shape[-2:]			),[self.model_tester.image_size // 4, self.model_tester.image_size // 4],)

            __UpperCamelCase,  __UpperCamelCase				=	self.model_tester.prepare_config_and_inputs_for_common()
            __UpperCamelCase				=	['basic', 'bottleneck']
            for model_class in self.all_model_classes:
                 for layer_type in layers_type:
                      __UpperCamelCase				=	layer_type
                      __UpperCamelCase				=	True
                      check_hidden_states_output(A_,A_,A_			)

                      # check that output_hidden_states also work using config
                      del inputs_dict["output_hidden_states"]
                      __UpperCamelCase				=	True

                      check_hidden_states_output(A_,A_,A_			)







       @unittest.skip(reason='ResNet does not use feedforward chunking'			)
       def 						snake_case_					(	self:				List[str]			):






            '''simple docstring'''




            pass







       def 						snake_case_					(	self:				Any			):






            '''simple docstring'''




            __UpperCamelCase				=	self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_image_classification(*A_			)







       @slow
       def 						snake_case_					(	self:				Union[str, Any]			):






            '''simple docstring'''




            for model_name in RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 __UpperCamelCase				=	ResNetModel.from_pretrained(A_			)
                 self.assertIsNotNone(A_			)






def 						_A						(						)			->				int:


     """simple docstring"""


     __UpperCamelCase				=	Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'				)
     return image



@require_torch
@require_vision
class 					__lowerCamelCase	(unittest.TestCase		):







       @cached_property
       def 						snake_case_					(	self:				Optional[Any]			):






            '''simple docstring'''




            return (
                AutoImageProcessor.from_pretrained(RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]			)
                if is_vision_available()
                else None
            )







       @slow
       def 						snake_case_					(	self:				Optional[int]			):






            '''simple docstring'''




            __UpperCamelCase				=	ResNetForImageClassification.from_pretrained(RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]			).to(A_			)

            __UpperCamelCase				=	self.default_image_processor
            __UpperCamelCase				=	prepare_img()
            __UpperCamelCase				=	image_processor(images=A_,return_tensors='pt'			).to(A_			)

            # forward pass
            with torch.no_grad():
                 __UpperCamelCase				=	model(**A_			)

            # verify the logits
            __UpperCamelCase				=	torch.Size((1, 1000)			)
            self.assertEqual(outputs.logits.shape,A_			)

            __UpperCamelCase				=	torch.tensor([-1_1.1_0_6_9, -9.7_8_7_7, -8.3_7_7_7]			).to(A_			)

            self.assertTrue(torch.allclose(outputs.logits[0, :3],A_,atol=1E-4			)			)



@require_torch
class 					__lowerCamelCase	(_a							,     unittest.TestCase		):
       _lowercase    							=      (ResNetBackbone,) if is_torch_available() else ()
       _lowercase    							=      False
       _lowercase    							=      ResNetConfig







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






            '''simple docstring'''




            __UpperCamelCase				=	ResNetModelTester(self			)


