





"""simple docstring"""
def 		_SCREAMING_SNAKE_CASE (_UpperCAmelCase				:						str ):
 lowerCAmelCase													=   str(_UpperCAmelCase ).strip()
 if not bin_string:
  raise ValueError('Empty string was passed to the function' )
 lowerCAmelCase													=   bin_string[0] == '-'
 if is_negative:
  lowerCAmelCase													=   bin_string[1:]
 if not all(char in '01' for char in bin_string ):
  raise ValueError('Non-binary value was passed to the function' )
 lowerCAmelCase													=   0
 for char in bin_string:
  lowerCAmelCase													=   2 * decimal_number + int(_UpperCAmelCase )
 return -decimal_number if is_negative else decimal_number


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

       testmod()












"""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       (      a__      ,   unittest.TestCase     ):
 snake_case__								=					UnCLIPPipeline
 snake_case__								=					TEXT_TO_IMAGE_PARAMS - {
     '''negative_prompt''',
     '''height''',
     '''width''',
     '''negative_prompt_embeds''',
     '''guidance_scale''',
     '''prompt_embeds''',
     '''cross_attention_kwargs''',
 }
 snake_case__								=					TEXT_TO_IMAGE_BATCH_PARAMS
 snake_case__								=					[
     '''generator''',
     '''return_dict''',
     '''prior_num_inference_steps''',
     '''decoder_num_inference_steps''',
     '''super_res_num_inference_steps''',
 ]
 snake_case__								=					False



 @property
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  return 32



 @property
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  return 32



 @property
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  return self.time_input_dim



 @property
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  return self.time_input_dim * 4



 @property
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  return 1_00



 @property
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






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



 @property
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  torch.manual_seed(0    )
  lowerCAmelCase													=   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=10_00							,					)
  return CLIPTextModelWithProjection(_snake_case    )



 @property
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  torch.manual_seed(0    )

  lowerCAmelCase													=   {
      'num_attention_heads': 2,
      'attention_head_dim': 12,
      'embedding_dim': self.text_embedder_hidden_size,
      'num_layers': 1,
  }

  lowerCAmelCase													=   PriorTransformer(**_snake_case    )
  return model



 @property
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  torch.manual_seed(0    )

  lowerCAmelCase													=   {
      'clip_embeddings_dim': self.text_embedder_hidden_size,
      'time_embed_dim': self.time_embed_dim,
      'cross_attention_dim': self.cross_attention_dim,
  }

  lowerCAmelCase													=   UnCLIPTextProjModel(**_snake_case    )
  return model



 @property
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  torch.manual_seed(0    )

  lowerCAmelCase													=   {
      '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',
  }

  lowerCAmelCase													=   UNetaDConditionModel(**_snake_case    )
  return model



 @property
 def       UpperCamelCase__     (    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       UpperCamelCase__     (    self    ):







  """simple docstring"""






  torch.manual_seed(0    )

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



 @property
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  torch.manual_seed(1    )

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



 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






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

  lowerCAmelCase													=   UnCLIPScheduler(
      variance_type='fixed_small_log'							,					prediction_type='sample'							,					num_train_timesteps=10_00							,					clip_sample_range=5.0							,					)

  lowerCAmelCase													=   UnCLIPScheduler(
      variance_type='learned_range'							,					prediction_type='epsilon'							,					num_train_timesteps=10_00							,					)

  lowerCAmelCase													=   UnCLIPScheduler(
      variance_type='fixed_small_log'							,					prediction_type='epsilon'							,					num_train_timesteps=10_00							,					)

  lowerCAmelCase													=   {
      '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       UpperCamelCase__     (    self							,					_snake_case							,					_snake_case=0    ):







  """simple docstring"""






  if str(_snake_case    ).startswith('mps'    ):
   lowerCAmelCase													=   torch.manual_seed(_snake_case    )
  else:
   lowerCAmelCase													=   torch.Generator(device=_snake_case    ).manual_seed(_snake_case    )
  lowerCAmelCase													=   {
      '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       UpperCamelCase__     (    self    ):







  """simple docstring"""






  lowerCAmelCase													=   'cpu'

  lowerCAmelCase													=   self.get_dummy_components()

  lowerCAmelCase													=   self.pipeline_class(**_snake_case    )
  lowerCAmelCase													=   pipe.to(_snake_case    )

  pipe.set_progress_bar_config(disable=_snake_case    )

  lowerCAmelCase													=   pipe(**self.get_dummy_inputs(_snake_case    )    )
  lowerCAmelCase													=   output.images

  lowerCAmelCase													=   pipe(
      **self.get_dummy_inputs(_snake_case    )							,					return_dict=_snake_case							,					)[0]

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

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

  lowerCAmelCase													=   np.array(
      [
          0.9_997,
          0.9_988,
          0.0_028,
          0.9_997,
          0.9_984,
          0.9_965,
          0.0_029,
          0.9_986,
          0.0_025,
      ]    )

  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       UpperCamelCase__     (    self    ):







  """simple docstring"""






  lowerCAmelCase													=   torch.device('cpu'    )







  class 			a       :
   snake_case__								=					1


  lowerCAmelCase													=   self.get_dummy_components()

  lowerCAmelCase													=   self.pipeline_class(**_snake_case    )
  lowerCAmelCase													=   pipe.to(_snake_case    )

  lowerCAmelCase													=   components['prior']
  lowerCAmelCase													=   components['decoder']
  lowerCAmelCase													=   components['super_res_first']
  lowerCAmelCase													=   components['tokenizer']
  lowerCAmelCase													=   components['text_encoder']

  lowerCAmelCase													=   torch.Generator(device=_snake_case    ).manual_seed(0    )
  lowerCAmelCase													=   prior.dtype
  lowerCAmelCase													=   1

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

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

  pipe.set_progress_bar_config(disable=_snake_case    )

  lowerCAmelCase													=   'this is a prompt example'

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

  lowerCAmelCase													=   tokenizer(
      _snake_case							,					padding='max_length'							,					max_length=tokenizer.model_max_length							,					return_tensors='pt'							,					)
  lowerCAmelCase													=   text_encoder(text_inputs.input_ids    )
  lowerCAmelCase													=   text_inputs.attention_mask

  lowerCAmelCase													=   torch.Generator(device=_snake_case    ).manual_seed(0    )
  lowerCAmelCase													=   pipe(
      generator=_snake_case							,					prior_num_inference_steps=2							,					decoder_num_inference_steps=2							,					super_res_num_inference_steps=2							,					prior_latents=_snake_case							,					decoder_latents=_snake_case							,					super_res_latents=_snake_case							,					text_model_output=_snake_case							,					text_attention_mask=_snake_case							,					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       UpperCamelCase__     (    self    ):







  """simple docstring"""






  lowerCAmelCase													=   torch_device == 'cpu'

  self._test_attention_slicing_forward_pass(test_max_difference=_snake_case							,					expected_max_diff=0.01    )



 @skip_mps
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  lowerCAmelCase													=   torch_device == 'cpu'
  lowerCAmelCase													=   True
  lowerCAmelCase													=   [
      'prior_num_inference_steps',
      'decoder_num_inference_steps',
      'super_res_num_inference_steps',
  ]

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



 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  lowerCAmelCase													=   [
      'prior_num_inference_steps',
      'decoder_num_inference_steps',
      'super_res_num_inference_steps',
  ]

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



 @skip_mps
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  return super().test_dict_tuple_outputs_equivalent()



 @skip_mps
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  return super().test_save_load_local()



 @skip_mps
 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






  return super().test_save_load_optional_components()







@nightly
class 			a       (      unittest.TestCase     ):



 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






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



 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






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

  lowerCAmelCase													=   UnCLIPPipeline.from_pretrained('kakaobrain/karlo-v1-alpha'    )
  pipeline.set_progress_bar_config(disable=_snake_case    )

  lowerCAmelCase													=   torch.manual_seed(0    )
  lowerCAmelCase													=   pipeline(
      'horse'							,					num_images_per_prompt=1							,					generator=_snake_case							,					output_type='np'							,					)

  lowerCAmelCase													=   output.images[0]

  assert image.shape == (2_56, 2_56, 3)
  assert np.abs(expected_image - image    ).max() < 1E-1







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



 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






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



 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






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

  lowerCAmelCase													=   UnCLIPPipeline.from_pretrained('kakaobrain/karlo-v1-alpha'							,					torch_dtype=torch.floataa    )
  lowerCAmelCase													=   pipeline.to(_snake_case    )
  pipeline.set_progress_bar_config(disable=_snake_case    )

  lowerCAmelCase													=   torch.Generator(device='cpu'    ).manual_seed(0    )
  lowerCAmelCase													=   pipeline(
      'horse'							,					generator=_snake_case							,					output_type='np'							,					)

  lowerCAmelCase													=   output.images[0]

  assert image.shape == (2_56, 2_56, 3)

  assert_mean_pixel_difference(_snake_case							,					_snake_case    )



 def       UpperCamelCase__     (    self    ):







  """simple docstring"""






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

  lowerCAmelCase													=   UnCLIPPipeline.from_pretrained('kakaobrain/karlo-v1-alpha'							,					torch_dtype=torch.floataa    )
  lowerCAmelCase													=   pipe.to(_snake_case    )
  pipe.set_progress_bar_config(disable=_snake_case    )
  pipe.enable_attention_slicing()
  pipe.enable_sequential_cpu_offload()

  lowerCAmelCase													=   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'							,					)

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





