

import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModelWithProjection

from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    PNDMScheduler,
    StableDiffusionImageVariationPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class 						lowercase_					(						__lowercase     ,							__lowercase     ,							__lowercase     ,							unittest.TestCase  ):
						UpperCamelCase_    :	Union[str, Any]		   =       StableDiffusionImageVariationPipeline
						UpperCamelCase_    :	Any		   =       IMAGE_VARIATION_PARAMS
						UpperCamelCase_    :	Optional[int]		   =       IMAGE_VARIATION_BATCH_PARAMS
						UpperCamelCase_    :	Dict		   =       frozenset([]  )
						# TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
						UpperCamelCase_    :	List[Any]		   =       frozenset([]  )



						def 		UpperCamelCase_	(    self   :       Any	)					->  List[str]:
								torch.manual_seed(0	)
								_snake_case		=							UNetaDConditionModel(
								    block_out_channels=(32, 64)    ,      layers_per_block=2    ,      sample_size=32    ,      in_channels=4    ,      out_channels=4    ,      down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')    ,      up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')    ,      cross_attention_dim=32    ,      )
								_snake_case		=							PNDMScheduler(skip_prk_steps=A__	)
								torch.manual_seed(0	)
								_snake_case		=							AutoencoderKL(
								    block_out_channels=[32, 64]    ,      in_channels=3    ,      out_channels=3    ,      down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']    ,      up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']    ,      latent_channels=4    ,      )
								torch.manual_seed(0	)
								_snake_case		=							CLIPVisionConfig(
								    hidden_size=32    ,      projection_dim=32    ,      intermediate_size=37    ,      layer_norm_eps=1e-05    ,      num_attention_heads=4    ,      num_hidden_layers=5    ,      image_size=32    ,      patch_size=4    ,      )
								_snake_case		=							CLIPVisionModelWithProjection(A__	)
								_snake_case		=							CLIPImageProcessor(crop_size=32    ,      size=32	)

								_snake_case		=							{
								    '''unet''': unet,
								    '''scheduler''': scheduler,
								    '''vae''': vae,
								    '''image_encoder''': image_encoder,
								    '''feature_extractor''': feature_extractor,
								    '''safety_checker''': None,
								}
								return components



						def 		UpperCamelCase_	(    self   :       Optional[Any]    ,      A__   :       Union[str, Any]    ,      A__   :       Any=0	)					->  Optional[int]:
								_snake_case		=							floats_tensor((1, 3, 32, 32)    ,      rng=random.Random(A__	)	)
								_snake_case		=							image.cpu().permute(0    ,      2    ,      3    ,      1	)[0]
								_snake_case		=							Image.fromarray(np.uinta(A__	)	).convert('''RGB'''	).resize((32, 32)	)
								if str(A__	).startswith('''mps'''	):
										_snake_case		=							torch.manual_seed(A__	)
								else:
										_snake_case		=							torch.Generator(device=A__	).manual_seed(A__	)
								_snake_case		=							{
								    '''image''': image,
								    '''generator''': generator,
								    '''num_inference_steps''': 2,
								    '''guidance_scale''': 6.0,
								    '''output_type''': '''numpy''',
								}
								return inputs



						def 		UpperCamelCase_	(    self   :       Optional[Any]	)					->  Optional[Any]:
								_snake_case		=							'''cpu'''  # ensure determinism for the device-dependent torch.Generator
								_snake_case		=							self.get_dummy_components()
								_snake_case		=							StableDiffusionImageVariationPipeline(**A__	)
								_snake_case		=							sd_pipe.to(A__	)
								sd_pipe.set_progress_bar_config(disable=A__	)

								_snake_case		=							self.get_dummy_inputs(A__	)
								_snake_case		=							sd_pipe(**A__	).images
								_snake_case		=							image[0, -3:, -3:, -1]

								assert image.shape == (1, 64, 64, 3)
								_snake_case		=							np.array([0.5239, 0.5723, 0.4796, 0.5049, 0.5550, 0.4685, 0.5329, 0.4891, 0.4921]	)

								assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-3



						def 		UpperCamelCase_	(    self   :       Optional[Any]	)					->  str:
								_snake_case		=							'''cpu'''  # ensure determinism for the device-dependent torch.Generator
								_snake_case		=							self.get_dummy_components()
								_snake_case		=							StableDiffusionImageVariationPipeline(**A__	)
								_snake_case		=							sd_pipe.to(A__	)
								sd_pipe.set_progress_bar_config(disable=A__	)

								_snake_case		=							self.get_dummy_inputs(A__	)
								_snake_case		=							2 * [inputs['''image''']]
								_snake_case		=							sd_pipe(**A__	)

								_snake_case		=							output.images

								_snake_case		=							image[-1, -3:, -3:, -1]

								assert image.shape == (2, 64, 64, 3)
								_snake_case		=							np.array([0.6892, 0.5637, 0.5836, 0.5771, 0.6254, 0.6409, 0.5580, 0.5569, 0.5289]	)

								assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-3







						def 		UpperCamelCase_	(    self   :       str	)					->  Tuple:
								super().test_inference_batch_single_identical(expected_max_diff=3e-3	)



@slow
@require_torch_gpu
class 						lowercase_					(						unittest.TestCase  ):



						def 		UpperCamelCase_	(    self   :       Optional[int]	)					->  List[str]:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()



						def 		UpperCamelCase_	(    self   :       Union[str, Any]    ,      A__   :       Dict    ,      A__   :       Any="cpu"    ,      A__   :       Optional[Any]=torch.floataa    ,      A__   :       Tuple=0	)					->  List[Any]:
								_snake_case		=							torch.Generator(device=A__	).manual_seed(A__	)
								_snake_case		=							load_image(
								    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'''
								    '''/stable_diffusion_imgvar/input_image_vermeer.png'''	)
								_snake_case		=							np.random.RandomState(A__	).standard_normal((1, 4, 64, 64)	)
								_snake_case		=							torch.from_numpy(A__	).to(device=A__    ,      dtype=A__	)
								_snake_case		=							{
								    '''image''': init_image,
								    '''latents''': latents,
								    '''generator''': generator,
								    '''num_inference_steps''': 3,
								    '''guidance_scale''': 7.5,
								    '''output_type''': '''numpy''',
								}
								return inputs



						def 		UpperCamelCase_	(    self   :       Optional[Any]	)					->  Optional[Any]:
								_snake_case		=							StableDiffusionImageVariationPipeline.from_pretrained(
								    '''lambdalabs/sd-image-variations-diffusers'''    ,      safety_checker=A__	)
								_snake_case		=							sd_pipe.to(A__	)
								sd_pipe.set_progress_bar_config(disable=A__	)

								_snake_case		=							self.get_inputs(A__	)
								_snake_case		=							sd_pipe(**A__	).images
								_snake_case		=							image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 512, 512, 3)
								_snake_case		=							np.array([0.8_4491, 0.9_0789, 0.7_5708, 0.7_8734, 0.8_3485, 0.7_0099, 0.6_6938, 0.6_8727, 0.6_1379]	)
								assert np.abs(image_slice - expected_slice	).max() < 6e-3



						def 		UpperCamelCase_	(    self   :       int	)					->  List[str]:
								_snake_case		=							0

								def callback_fn(A__   :       int    ,      A__   :       int    ,      A__   :       torch.FloatTensor	) -> None:
										_snake_case		=							True
										nonlocal number_of_steps
										number_of_steps += 1
										if step == 1:
												_snake_case		=							latents.detach().cpu().numpy()
												assert latents.shape == (1, 4, 64, 64)
												_snake_case		=							latents[0, -3:, -3:, -1]
												_snake_case		=							np.array(
												    [-0.1621, 0.2837, -0.7979, -0.1221, -1.3057, 0.7681, -2.1191, 0.0464, 1.6309]	)

												assert np.abs(latents_slice.flatten() - expected_slice	).max() < 5e-2
										elif step == 2:
												_snake_case		=							latents.detach().cpu().numpy()
												assert latents.shape == (1, 4, 64, 64)
												_snake_case		=							latents[0, -3:, -3:, -1]
												_snake_case		=							np.array([0.6299, 1.7500, 1.1992, -2.1582, -1.8994, 0.7334, -0.7090, 1.0137, 1.5273]	)

												assert np.abs(latents_slice.flatten() - expected_slice	).max() < 5e-2

								_snake_case		=							False

								_snake_case		=							StableDiffusionImageVariationPipeline.from_pretrained(
								    '''fusing/sd-image-variations-diffusers'''    ,      safety_checker=A__    ,      torch_dtype=torch.floataa    ,      )
								pipe.to(A__	)
								pipe.set_progress_bar_config(disable=A__	)
								pipe.enable_attention_slicing()

								_snake_case		=							self.get_inputs(A__    ,      dtype=torch.floataa	)
								pipe(**A__    ,      callback=A__    ,      callback_steps=1	)
								assert callback_fn.has_been_called
								assert number_of_steps == inputs["num_inference_steps"]







						def 		UpperCamelCase_	(    self   :       str	)					->  Tuple:
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

								_snake_case		=							'''fusing/sd-image-variations-diffusers'''
								_snake_case		=							StableDiffusionImageVariationPipeline.from_pretrained(
								    A__    ,      safety_checker=A__    ,      torch_dtype=torch.floataa	)
								_snake_case		=							pipe.to(A__	)
								pipe.set_progress_bar_config(disable=A__	)
								pipe.enable_attention_slicing(1	)
								pipe.enable_sequential_cpu_offload()

								_snake_case		=							self.get_inputs(A__    ,      dtype=torch.floataa	)
								_snake_case		=							pipe(**A__	)

								_snake_case		=							torch.cuda.max_memory_allocated()
								# make sure that less than 2.6 GB is allocated
								assert mem_bytes < 2.6 * 10**9



@nightly
@require_torch_gpu
class 						lowercase_					(						unittest.TestCase  ):



						def 		UpperCamelCase_	(    self   :       Dict	)					->  List[str]:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()



						def 		UpperCamelCase_	(    self   :       Optional[int]    ,      A__   :       Optional[Any]    ,      A__   :       Optional[int]="cpu"    ,      A__   :       List[Any]=torch.floataa    ,      A__   :       Optional[Any]=0	)					->  int:
								_snake_case		=							torch.Generator(device=A__	).manual_seed(A__	)
								_snake_case		=							load_image(
								    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'''
								    '''/stable_diffusion_imgvar/input_image_vermeer.png'''	)
								_snake_case		=							np.random.RandomState(A__	).standard_normal((1, 4, 64, 64)	)
								_snake_case		=							torch.from_numpy(A__	).to(device=A__    ,      dtype=A__	)
								_snake_case		=							{
								    '''image''': init_image,
								    '''latents''': latents,
								    '''generator''': generator,
								    '''num_inference_steps''': 50,
								    '''guidance_scale''': 7.5,
								    '''output_type''': '''numpy''',
								}
								return inputs



						def 		UpperCamelCase_	(    self   :       List[Any]	)					->  int:
								_snake_case		=							StableDiffusionImageVariationPipeline.from_pretrained('''fusing/sd-image-variations-diffusers'''	)
								sd_pipe.to(A__	)
								sd_pipe.set_progress_bar_config(disable=A__	)

								_snake_case		=							self.get_inputs(A__	)
								_snake_case		=							sd_pipe(**A__	).images[0]

								_snake_case		=							load_numpy(
								    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'''
								    '''/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy'''	)
								_snake_case		=							np.abs(expected_image - image	).max()
								assert max_diff < 1e-3







						def 		UpperCamelCase_	(    self   :       Union[str, Any]	)					->  Dict:
								_snake_case		=							StableDiffusionImageVariationPipeline.from_pretrained('''fusing/sd-image-variations-diffusers'''	)
								_snake_case		=							DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config	)
								sd_pipe.to(A__	)
								sd_pipe.set_progress_bar_config(disable=A__	)

								_snake_case		=							self.get_inputs(A__	)
								_snake_case		=							25
								_snake_case		=							sd_pipe(**A__	).images[0]

								_snake_case		=							load_numpy(
								    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'''
								    '''/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy'''	)
								_snake_case		=							np.abs(expected_image - image	).max()
								assert max_diff < 1e-3




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 						lowercase_					(						__lowercase  ):



      def __init__(    self   :       str    ,      A__   :       UNetaDModel    ,      A__   :       CMStochasticIterativeScheduler	)					->  None:
        super().__init__()

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

        _snake_case		=							None



      def 		UpperCamelCase_	(    self   :       Union[str, Any]    ,      A__   :       Optional[int]=0	)					->  List[str]:
        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'''	)

        _snake_case		=							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 		UpperCamelCase_	(    self   :       Optional[int]    ,      A__   :       Optional[Any]=0	)					->  Union[str, Any]:
        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.'''	)

        _snake_case		=							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)

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

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

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



      @property
      # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
      def 		UpperCamelCase_	(    self   :       int	)					->  List[str]:
        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 		UpperCamelCase_	(    self   :       List[str]    ,      A__   :       Tuple    ,      A__   :       str    ,      A__   :       Tuple    ,      A__   :       Union[str, Any]    ,      A__   :       Union[str, Any]    ,      A__   :       Union[str, Any]    ,      A__   :       Tuple    ,      A__   :       Union[str, Any]=None	)					->  Optional[int]:
        _snake_case		=							(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:
          _snake_case		=							randn_tensor(A__    ,      generator=A__    ,      device=A__    ,      dtype=A__	)
        else:
          _snake_case		=							latents.to(device=A__    ,      dtype=A__	)

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



      def 		UpperCamelCase_	(    self   :       Optional[Any]    ,      A__   :       torch.FloatTensor    ,      A__   :       str = "pil"	)					->  Optional[int]:
        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
        _snake_case		=							(sample / 2 + 0.5).clamp(0    ,      1	)
        if output_type == "pt":
          return sample

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

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



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



      def 		UpperCamelCase_	(    self   :       List[str]    ,      A__   :       List[str]    ,      A__   :       List[str]    ,      A__   :       str    ,      A__   :       List[Any]    ,      A__   :       Optional[Any]    ,      A__   :       Dict	)					->  Optional[Any]:
        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:
          _snake_case		=							(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   :       Any    ,      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    ,      )					->  Dict:
        _snake_case		=							self.unet.config.sample_size
        _snake_case		=							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)
        _snake_case		=							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
        _snake_case		=							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__	)
          _snake_case		=							self.scheduler.timesteps
          _snake_case		=							len(A__	)
        else:
          self.scheduler.set_timesteps(A__	)
          _snake_case		=							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__	):
            _snake_case		=							self.scheduler.scale_model_input(A__    ,      A__	)
            _snake_case		=							self.unet(A__    ,      A__    ,      class_labels=A__    ,      return_dict=A__	)[0]

            _snake_case		=							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
        _snake_case		=							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__	)

