
"""simple docstring"""




from dataclasses import dataclass
from typing import Optional, Tuple, Union

import flax
import jax.numpy as jnp
from jax import random

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils_flax import FlaxSchedulerMixin, FlaxSchedulerOutput, broadcast_to_shape_from_left




@flax.struct.dataclass
class 	UpperCamelCase      :

				"""simple docstring"""



				UpperCAmelCase_							: Optional[jnp.ndarray]		    =							None
				UpperCAmelCase_							: Optional[jnp.ndarray]		    =							None
				UpperCAmelCase_							: Optional[jnp.ndarray]		    =							None


				@classmethod
				def 				A   (							cls     )       ->       int:
									"""simple docstring"""







									return cls()






@dataclass
class 	UpperCamelCase      (			lowerCAmelCase_ ):

				"""simple docstring"""



				UpperCAmelCase_							: ScoreSdeVeSchedulerState
				UpperCAmelCase_							: jnp.ndarray
				UpperCAmelCase_							: Optional[jnp.ndarray]		    =							None






class 	UpperCamelCase      (			lowerCAmelCase_    , lowerCAmelCase_ ):

				"""simple docstring"""



				@property
				def 				A   (							self     )       ->       Optional[Any]:
									"""simple docstring"""







									return True


				@register_to_config
				def __init__(							self						,      lowercase__ = 2000						,      lowercase__ = 0.15						,      lowercase__ = 0.01						,      lowercase__ = 1348.0						,      lowercase__ = 1E-5						,      lowercase__ = 1						,      )       ->       int:
									"""simple docstring"""







									pass


				def 				A   (							self     )       ->       str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =      ScoreSdeVeSchedulerState.create()
									return self.set_sigmas(
									    lowerCamelCase_						,      self.config.num_train_timesteps						,      self.config.sigma_min						,      self.config.sigma_max						,      self.config.sampling_eps						,      )


				def 				A   (							self						,      lowercase__						,      lowercase__						,      lowercase__ = ()						,      lowercase__ = None     )       ->       Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =      sampling_eps if sampling_eps is not None else self.config.sampling_eps

									SCREAMING_SNAKE_CASE           =      jnp.linspace(1						,      lowerCamelCase_						,      lowerCamelCase_     )
									return state.replace(timesteps=lowerCamelCase_     )


				def 				A   (							self						,      lowercase__						,      lowercase__						,      lowercase__ = None						,      lowercase__ = None						,      lowercase__ = None						,      )       ->       str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =      sigma_min if sigma_min is not None else self.config.sigma_min
									SCREAMING_SNAKE_CASE           =      sigma_max if sigma_max is not None else self.config.sigma_max
									SCREAMING_SNAKE_CASE           =      sampling_eps if sampling_eps is not None else self.config.sampling_eps
									if state.timesteps is None:
														SCREAMING_SNAKE_CASE           =      self.set_timesteps(lowerCamelCase_						,      lowerCamelCase_						,      lowerCamelCase_     )

									SCREAMING_SNAKE_CASE           =      jnp.exp(jnp.linspace(jnp.log(lowerCamelCase_     )						,      jnp.log(lowerCamelCase_     )						,      lowerCamelCase_     )     )
									SCREAMING_SNAKE_CASE           =      jnp.array([sigma_min * (sigma_max / sigma_min) ** t for t in state.timesteps]     )

									return state.replace(discrete_sigmas=lowerCamelCase_						,      sigmas=lowerCamelCase_     )


				def 				A   (							self						,      lowercase__						,      lowercase__						,      lowercase__     )       ->       Tuple:
									"""simple docstring"""







									return jnp.where(timesteps == 0						,      jnp.zeros_like(lowerCamelCase_     )						,      state.discrete_sigmas[timesteps - 1]     )


				def 				A   (							self						,      lowercase__						,      lowercase__						,      lowercase__						,      lowercase__						,      lowercase__						,      lowercase__ = True						,      )       ->       List[Any]:
									"""simple docstring"""







									if state.timesteps is None:
														raise ValueError(
														    '`state.timesteps` is not set, you need to run \'set_timesteps\' after creating the scheduler'     )

									SCREAMING_SNAKE_CASE           =      timestep * jnp.ones(
									    sample.shape[0]						,      )
									SCREAMING_SNAKE_CASE           =      (timestep * (len(state.timesteps     ) - 1)).long()

									SCREAMING_SNAKE_CASE           =      state.discrete_sigmas[timesteps]
									SCREAMING_SNAKE_CASE           =      self.get_adjacent_sigma(lowerCamelCase_						,      lowerCamelCase_						,      lowerCamelCase_     )
									SCREAMING_SNAKE_CASE           =      jnp.zeros_like(lowerCamelCase_     )
									SCREAMING_SNAKE_CASE           =      (sigma**2 - adjacent_sigma**2) ** 0.5

									# equation 6 in the paper: the model_output modeled by the network is grad_x log pt(x)
									# also equation 47 shows the analog from SDE models to ancestral sampling methods
									SCREAMING_SNAKE_CASE           =      diffusion.flatten()
									SCREAMING_SNAKE_CASE           =      broadcast_to_shape_from_left(lowerCamelCase_						,      sample.shape     )
									SCREAMING_SNAKE_CASE           =      drift - diffusion**2 * model_output

									#  equation 6: sample noise for the diffusion term of
									SCREAMING_SNAKE_CASE           =      random.split(lowerCamelCase_						,      num=1     )
									SCREAMING_SNAKE_CASE           =      random.normal(key=lowerCamelCase_						,      shape=sample.shape     )
									SCREAMING_SNAKE_CASE           =      sample - drift  # subtract because `dt` is a small negative timestep
									# TODO is the variable diffusion the correct scaling term for the noise?
									SCREAMING_SNAKE_CASE           =      prev_sample_mean + diffusion * noise  # add impact of diffusion field g

									if not return_dict:
														return (prev_sample, prev_sample_mean, state)

									return FlaxSdeVeOutput(prev_sample=lowerCamelCase_						,      prev_sample_mean=lowerCamelCase_						,      state=lowerCamelCase_     )


				def 				A   (							self						,      lowercase__						,      lowercase__						,      lowercase__						,      lowercase__						,      lowercase__ = True						,      )       ->       Optional[Any]:
									"""simple docstring"""







									if state.timesteps is None:
														raise ValueError(
														    '`state.timesteps` is not set, you need to run \'set_timesteps\' after creating the scheduler'     )

									# For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z"
									# sample noise for correction
									SCREAMING_SNAKE_CASE           =      random.split(lowerCamelCase_						,      num=1     )
									SCREAMING_SNAKE_CASE           =      random.normal(key=lowerCamelCase_						,      shape=sample.shape     )

									# compute step size from the model_output, the noise, and the snr
									SCREAMING_SNAKE_CASE           =      jnp.linalg.norm(lowerCamelCase_     )
									SCREAMING_SNAKE_CASE           =      jnp.linalg.norm(lowerCamelCase_     )
									SCREAMING_SNAKE_CASE           =      (self.config.snr * noise_norm / grad_norm) ** 2 * 2
									SCREAMING_SNAKE_CASE           =      step_size * jnp.ones(sample.shape[0]     )

									# compute corrected sample: model_output term and noise term
									SCREAMING_SNAKE_CASE           =      step_size.flatten()
									SCREAMING_SNAKE_CASE           =      broadcast_to_shape_from_left(lowerCamelCase_						,      sample.shape     )
									SCREAMING_SNAKE_CASE           =      sample + step_size * model_output
									SCREAMING_SNAKE_CASE           =      prev_sample_mean + ((step_size * 2) ** 0.5) * noise

									if not return_dict:
														return (prev_sample, state)

									return FlaxSdeVeOutput(prev_sample=lowerCamelCase_						,      state=lowerCamelCase_     )




				def __len__(							self     )       ->       Optional[Any]:
									"""simple docstring"""







									return self.config.num_train_timesteps


"""simple docstring"""







import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()






class      __a      (       lowerCAmelCase_					,						unittest.TestCase      ):
       _lowerCamelCase				:				int   			=      UnCLIPPipeline
       _lowerCamelCase				:				Any   			=      TEXT_TO_IMAGE_PARAMS - {
           'negative_prompt',
           'height',
           'width',
           'negative_prompt_embeds',
           'guidance_scale',
           'prompt_embeds',
           'cross_attention_kwargs',
       }
       _lowerCamelCase				:				Optional[int]   			=      TEXT_TO_IMAGE_BATCH_PARAMS
       _lowerCamelCase				:				Dict   			=      [
           'generator',
           'return_dict',
           'prior_num_inference_steps',
           'decoder_num_inference_steps',
           'super_res_num_inference_steps',
       ]
       _lowerCamelCase				:				Any   			=      False







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          return 32







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          return 32







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          return self.time_input_dim







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          return self.time_input_dim * 4







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          return 100







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          _UpperCAmelCase 						=       CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)
          return tokenizer







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          torch.manual_seed(0		)
          _UpperCAmelCase 						=       CLIPTextConfig(
              bos_token_id=0				,eos_token_id=2				,hidden_size=self.text_embedder_hidden_size				,projection_dim=self.text_embedder_hidden_size				,intermediate_size=37				,layer_norm_eps=1e-05				,num_attention_heads=4				,num_hidden_layers=5				,pad_token_id=1				,vocab_size=1000				,)
          return CLIPTextModelWithProjection(lowerCamelCase_		)







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          torch.manual_seed(0		)

          _UpperCAmelCase 						=       {
              "num_attention_heads": 2,
              "attention_head_dim": 12,
              "embedding_dim": self.text_embedder_hidden_size,
              "num_layers": 1,
          }

          _UpperCAmelCase 						=       PriorTransformer(**lowerCamelCase_		)
          return model







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          torch.manual_seed(0		)

          _UpperCAmelCase 						=       {
              "clip_embeddings_dim": self.text_embedder_hidden_size,
              "time_embed_dim": self.time_embed_dim,
              "cross_attention_dim": self.cross_attention_dim,
          }

          _UpperCAmelCase 						=       UnCLIPTextProjModel(**lowerCamelCase_		)
          return model







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          torch.manual_seed(0		)

          _UpperCAmelCase 						=       {
              "sample_size": 32,
              # RGB in channels
              "in_channels": 3,
              # Out channels is double in channels because predicts mean and variance
              "out_channels": 6,
              "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
              "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
              "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
              "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
              "layers_per_block": 1,
              "cross_attention_dim": self.cross_attention_dim,
              "attention_head_dim": 4,
              "resnet_time_scale_shift": "scale_shift",
              "class_embed_type": "identity",
          }

          _UpperCAmelCase 						=       UNetaDConditionModel(**lowerCamelCase_		)
          return model







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          return {
              "sample_size": 64,
              "layers_per_block": 1,
              "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
              "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
              "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
              "in_channels": 6,
              "out_channels": 3,
          }







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          torch.manual_seed(0		)

          _UpperCAmelCase 						=       UNetaDModel(**self.dummy_super_res_kwargs		)
          return model







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          torch.manual_seed(1		)

          _UpperCAmelCase 						=       UNetaDModel(**self.dummy_super_res_kwargs		)
          return model







       def  _snake_case      (       self		):



          '''simple docstring'''

          _UpperCAmelCase 						=       self.dummy_prior
          _UpperCAmelCase 						=       self.dummy_decoder
          _UpperCAmelCase 						=       self.dummy_text_proj
          _UpperCAmelCase 						=       self.dummy_text_encoder
          _UpperCAmelCase 						=       self.dummy_tokenizer
          _UpperCAmelCase 						=       self.dummy_super_res_first
          _UpperCAmelCase 						=       self.dummy_super_res_last

          _UpperCAmelCase 						=       UnCLIPScheduler(
              variance_type="fixed_small_log"				,prediction_type="sample"				,num_train_timesteps=1000				,clip_sample_range=5.0				,)

          _UpperCAmelCase 						=       UnCLIPScheduler(
              variance_type="learned_range"				,prediction_type="epsilon"				,num_train_timesteps=1000				,)

          _UpperCAmelCase 						=       UnCLIPScheduler(
              variance_type="fixed_small_log"				,prediction_type="epsilon"				,num_train_timesteps=1000				,)

          _UpperCAmelCase 						=       {
              "prior": prior,
              "decoder": decoder,
              "text_proj": text_proj,
              "text_encoder": text_encoder,
              "tokenizer": tokenizer,
              "super_res_first": super_res_first,
              "super_res_last": super_res_last,
              "prior_scheduler": prior_scheduler,
              "decoder_scheduler": decoder_scheduler,
              "super_res_scheduler": super_res_scheduler,
          }

          return components







       def  _snake_case      (       self				,lowerCamelCase_				,lowerCamelCase_=0		):



          '''simple docstring'''

          if str(lowerCamelCase_		).startswith("mps"		):
             _UpperCAmelCase 						=       torch.manual_seed(lowerCamelCase_		)
          else:
             _UpperCAmelCase 						=       torch.Generator(device=lowerCamelCase_		).manual_seed(lowerCamelCase_		)
          _UpperCAmelCase 						=       {
              "prompt": "horse",
              "generator": generator,
              "prior_num_inference_steps": 2,
              "decoder_num_inference_steps": 2,
              "super_res_num_inference_steps": 2,
              "output_type": "numpy",
          }
          return inputs







       def  _snake_case      (       self		):



          '''simple docstring'''

          _UpperCAmelCase 						=       "cpu"

          _UpperCAmelCase 						=       self.get_dummy_components()

          _UpperCAmelCase 						=       self.pipeline_class(**lowerCamelCase_		)
          _UpperCAmelCase 						=       pipe.to(lowerCamelCase_		)

          pipe.set_progress_bar_config(disable=lowerCamelCase_		)

          _UpperCAmelCase 						=       pipe(**self.get_dummy_inputs(lowerCamelCase_		)		)
          _UpperCAmelCase 						=       output.images

          _UpperCAmelCase 						=       pipe(
              **self.get_dummy_inputs(lowerCamelCase_		)				,return_dict=lowerCamelCase_				,)[0]

          _UpperCAmelCase 						=       image[0, -3:, -3:, -1]
          _UpperCAmelCase 						=       image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 64, 64, 3)

          _UpperCAmelCase 						=       np.array(
              [
                  0.9_9_9_7,
                  0.9_9_8_8,
                  0.0_0_2_8,
                  0.9_9_9_7,
                  0.9_9_8_4,
                  0.9_9_6_5,
                  0.0_0_2_9,
                  0.9_9_8_6,
                  0.0_0_2_5,
              ]		)

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







       def  _snake_case      (       self		):



          '''simple docstring'''

          _UpperCAmelCase 						=       torch.device("cpu"		)






          class      __a      :
                 _lowerCamelCase				:				Optional[Any]   			=      1



          _UpperCAmelCase 						=       self.get_dummy_components()

          _UpperCAmelCase 						=       self.pipeline_class(**lowerCamelCase_		)
          _UpperCAmelCase 						=       pipe.to(lowerCamelCase_		)

          _UpperCAmelCase 						=       components["prior"]
          _UpperCAmelCase 						=       components["decoder"]
          _UpperCAmelCase 						=       components["super_res_first"]
          _UpperCAmelCase 						=       components["tokenizer"]
          _UpperCAmelCase 						=       components["text_encoder"]

          _UpperCAmelCase 						=       torch.Generator(device=lowerCamelCase_		).manual_seed(0		)
          _UpperCAmelCase 						=       prior.dtype
          _UpperCAmelCase 						=       1

          _UpperCAmelCase 						=       (batch_size, prior.config.embedding_dim)
          _UpperCAmelCase 						=       pipe.prepare_latents(
              lowerCamelCase_				,dtype=lowerCamelCase_				,device=lowerCamelCase_				,generator=lowerCamelCase_				,latents=lowerCamelCase_				,scheduler=DummyScheduler()		)
          _UpperCAmelCase 						=       (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
          _UpperCAmelCase 						=       pipe.prepare_latents(
              lowerCamelCase_				,dtype=lowerCamelCase_				,device=lowerCamelCase_				,generator=lowerCamelCase_				,latents=lowerCamelCase_				,scheduler=DummyScheduler()		)

          _UpperCAmelCase 						=       (
              batch_size,
              super_res_first.config.in_channels // 2,
              super_res_first.config.sample_size,
              super_res_first.config.sample_size,
          )
          _UpperCAmelCase 						=       pipe.prepare_latents(
              lowerCamelCase_				,dtype=lowerCamelCase_				,device=lowerCamelCase_				,generator=lowerCamelCase_				,latents=lowerCamelCase_				,scheduler=DummyScheduler()		)

          pipe.set_progress_bar_config(disable=lowerCamelCase_		)

          _UpperCAmelCase 						=       "this is a prompt example"

          _UpperCAmelCase 						=       torch.Generator(device=lowerCamelCase_		).manual_seed(0		)
          _UpperCAmelCase 						=       pipe(
              [prompt]				,generator=lowerCamelCase_				,prior_num_inference_steps=2				,decoder_num_inference_steps=2				,super_res_num_inference_steps=2				,prior_latents=lowerCamelCase_				,decoder_latents=lowerCamelCase_				,super_res_latents=lowerCamelCase_				,output_type="np"				,)
          _UpperCAmelCase 						=       output.images

          _UpperCAmelCase 						=       tokenizer(
              lowerCamelCase_				,padding="max_length"				,max_length=tokenizer.model_max_length				,return_tensors="pt"				,)
          _UpperCAmelCase 						=       text_encoder(text_inputs.input_ids		)
          _UpperCAmelCase 						=       text_inputs.attention_mask

          _UpperCAmelCase 						=       torch.Generator(device=lowerCamelCase_		).manual_seed(0		)
          _UpperCAmelCase 						=       pipe(
              generator=lowerCamelCase_				,prior_num_inference_steps=2				,decoder_num_inference_steps=2				,super_res_num_inference_steps=2				,prior_latents=lowerCamelCase_				,decoder_latents=lowerCamelCase_				,super_res_latents=lowerCamelCase_				,text_model_output=lowerCamelCase_				,text_attention_mask=lowerCamelCase_				,output_type="np"				,)[0]

          # make sure passing text embeddings manually is identical
          assert np.abs(image - image_from_text		).max() < 1e-4







       @skip_mps
       def  _snake_case      (       self		):



          '''simple docstring'''

          _UpperCAmelCase 						=       torch_device == "cpu"

          self._test_attention_slicing_forward_pass(test_max_difference=lowerCamelCase_				,expected_max_diff=0.0_1		)







       @skip_mps
       def  _snake_case      (       self		):



          '''simple docstring'''

          _UpperCAmelCase 						=       torch_device == "cpu"
          _UpperCAmelCase 						=       True
          _UpperCAmelCase 						=       [
              "prior_num_inference_steps",
              "decoder_num_inference_steps",
              "super_res_num_inference_steps",
          ]

          self._test_inference_batch_single_identical(
              test_max_difference=lowerCamelCase_				,relax_max_difference=lowerCamelCase_				,additional_params_copy_to_batched_inputs=lowerCamelCase_				,)







       def  _snake_case      (       self		):



          '''simple docstring'''

          _UpperCAmelCase 						=       [
              "prior_num_inference_steps",
              "decoder_num_inference_steps",
              "super_res_num_inference_steps",
          ]

          if torch_device == "mps":
             # TODO: MPS errors with larger batch sizes
             _UpperCAmelCase 						=       [2, 3]
             self._test_inference_batch_consistent(
                 batch_sizes=lowerCamelCase_				,additional_params_copy_to_batched_inputs=lowerCamelCase_				,)
          else:
             self._test_inference_batch_consistent(
                 additional_params_copy_to_batched_inputs=lowerCamelCase_		)







       @skip_mps
       def  _snake_case      (       self		):



          '''simple docstring'''

          return super().test_dict_tuple_outputs_equivalent()







       @skip_mps
       def  _snake_case      (       self		):



          '''simple docstring'''

          return super().test_save_load_local()







       @skip_mps
       def  _snake_case      (       self		):



          '''simple docstring'''

          return super().test_save_load_optional_components()






@nightly
class      __a      (       unittest.TestCase      ):







       def  _snake_case      (       self		):



          '''simple docstring'''

          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()







       def  _snake_case      (       self		):



          '''simple docstring'''

          _UpperCAmelCase 						=       load_numpy(
              "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
              "/unclip/karlo_v1_alpha_horse_cpu.npy"		)

          _UpperCAmelCase 						=       UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"		)
          pipeline.set_progress_bar_config(disable=lowerCamelCase_		)

          _UpperCAmelCase 						=       torch.manual_seed(0		)
          _UpperCAmelCase 						=       pipeline(
              "horse"				,num_images_per_prompt=1				,generator=lowerCamelCase_				,output_type="np"				,)

          _UpperCAmelCase 						=       output.images[0]

          assert image.shape == (256, 256, 3)
          assert np.abs(expected_image - image		).max() < 1e-1






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







       def  _snake_case      (       self		):



          '''simple docstring'''

          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()







       def  _snake_case      (       self		):



          '''simple docstring'''

          _UpperCAmelCase 						=       load_numpy(
              "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
              "/unclip/karlo_v1_alpha_horse_fp16.npy"		)

          _UpperCAmelCase 						=       UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"				,torch_dtype=torch.floataa		)
          _UpperCAmelCase 						=       pipeline.to(lowerCamelCase_		)
          pipeline.set_progress_bar_config(disable=lowerCamelCase_		)

          _UpperCAmelCase 						=       torch.Generator(device="cpu"		).manual_seed(0		)
          _UpperCAmelCase 						=       pipeline(
              "horse"				,generator=lowerCamelCase_				,output_type="np"				,)

          _UpperCAmelCase 						=       output.images[0]

          assert image.shape == (256, 256, 3)

          assert_mean_pixel_difference(lowerCamelCase_				,lowerCamelCase_		)







       def  _snake_case      (       self		):



          '''simple docstring'''

          torch.cuda.empty_cache()
          torch.cuda.reset_max_memory_allocated()
          torch.cuda.reset_peak_memory_stats()

          _UpperCAmelCase 						=       UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"				,torch_dtype=torch.floataa		)
          _UpperCAmelCase 						=       pipe.to(lowerCamelCase_		)
          pipe.set_progress_bar_config(disable=lowerCamelCase_		)
          pipe.enable_attention_slicing()
          pipe.enable_sequential_cpu_offload()

          _UpperCAmelCase 						=       pipe(
              "horse"				,num_images_per_prompt=1				,prior_num_inference_steps=2				,decoder_num_inference_steps=2				,super_res_num_inference_steps=2				,output_type="np"				,)

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

