





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  __UpperCamelCase						(       A__  ,   unittest.TestCase						):
     __A  :     List[Any]            =      UnCLIPPipeline
     __A  :     Union[str, Any]            =      TEXT_TO_IMAGE_PARAMS - {
         """negative_prompt""",
         """height""",
         """width""",
         """negative_prompt_embeds""",
         """guidance_scale""",
         """prompt_embeds""",
         """cross_attention_kwargs""",
     }
     __A  :     List[str]            =      TEXT_TO_IMAGE_BATCH_PARAMS
     __A  :     Any            =      [
         """generator""",
         """return_dict""",
         """prior_num_inference_steps""",
         """decoder_num_inference_steps""",
         """super_res_num_inference_steps""",
     ]
     __A  :     List[str]            =      False





     @property
     def 		UpperCamelCase(    self				):
        return 32





     @property
     def 		UpperCamelCase(    self				):
        return 32





     @property
     def 		UpperCamelCase(    self				):
        return self.time_input_dim





     @property
     def 		UpperCamelCase(    self				):
        return self.time_input_dim * 4





     @property
     def 		UpperCamelCase(    self				):
        return 100





     @property
     def 		UpperCamelCase(    self				):
        _UpperCAmelCase											=      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)
        return tokenizer





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





     @property
     def 		UpperCamelCase(    self				):
        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(**_UpperCamelCase				)
        return model





     @property
     def 		UpperCamelCase(    self				):
        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(**_UpperCamelCase				)
        return model





     @property
     def 		UpperCamelCase(    self				):
        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(**_UpperCamelCase				)
        return model





     @property
     def 		UpperCamelCase(    self				):
        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				):
        torch.manual_seed(0				)

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





     @property
     def 		UpperCamelCase(    self				):
        # seeded differently to get different unet than `self.dummy_super_res_first`
        torch.manual_seed(1				)

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





     def 		UpperCamelCase(    self				):
        _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 		UpperCamelCase(    self							,       _UpperCamelCase							,       _UpperCamelCase=0				):
        if str(_UpperCamelCase				).startswith('''mps'''				):
           _UpperCAmelCase											=      torch.manual_seed(_UpperCamelCase				)
        else:
           _UpperCAmelCase											=      torch.Generator(device=_UpperCamelCase				).manual_seed(_UpperCamelCase				)
        _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 		UpperCamelCase(    self				):
        _UpperCAmelCase											=      '''cpu'''

        _UpperCAmelCase											=      self.get_dummy_components()

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

        pipe.set_progress_bar_config(disable=_UpperCamelCase				)

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

        _UpperCAmelCase											=      pipe(
            **self.get_dummy_inputs(_UpperCamelCase				)							,       return_dict=_UpperCamelCase							,       )[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.9997,
                0.9988,
                0.0028,
                0.9997,
                0.9984,
                0.9965,
                0.0029,
                0.9986,
                0.0025,
            ]				)

        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				):
        _UpperCAmelCase											=      torch.device('''cpu'''				)




        class  __UpperCamelCase						:
             __A  :     int            =      1




        _UpperCAmelCase											=      self.get_dummy_components()

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

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

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

        _UpperCAmelCase											=      (batch_size, prior.config.embedding_dim)
        _UpperCAmelCase											=      pipe.prepare_latents(
            _UpperCamelCase							,       dtype=_UpperCamelCase							,       device=_UpperCamelCase							,       generator=_UpperCamelCase							,       latents=_UpperCamelCase							,       scheduler=DummyScheduler()				)
        _UpperCAmelCase											=      (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
        _UpperCAmelCase											=      pipe.prepare_latents(
            _UpperCamelCase							,       dtype=_UpperCamelCase							,       device=_UpperCamelCase							,       generator=_UpperCamelCase							,       latents=_UpperCamelCase							,       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(
            _UpperCamelCase							,       dtype=_UpperCamelCase							,       device=_UpperCamelCase							,       generator=_UpperCamelCase							,       latents=_UpperCamelCase							,       scheduler=DummyScheduler()				)

        pipe.set_progress_bar_config(disable=_UpperCamelCase				)

        _UpperCAmelCase											=      '''this is a prompt example'''

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

        _UpperCAmelCase											=      tokenizer(
            _UpperCamelCase							,       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=_UpperCamelCase				).manual_seed(0				)
        _UpperCAmelCase											=      pipe(
            generator=_UpperCamelCase							,       prior_num_inference_steps=2							,       decoder_num_inference_steps=2							,       super_res_num_inference_steps=2							,       prior_latents=_UpperCamelCase							,       decoder_latents=_UpperCamelCase							,       super_res_latents=_UpperCamelCase							,       text_model_output=_UpperCamelCase							,       text_attention_mask=_UpperCamelCase							,       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				):
        _UpperCAmelCase											=      torch_device == '''cpu'''

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





     @skip_mps
     def 		UpperCamelCase(    self				):
        _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=_UpperCamelCase							,       relax_max_difference=_UpperCamelCase							,       additional_params_copy_to_batched_inputs=_UpperCamelCase							,       )





     def 		UpperCamelCase(    self				):
        _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=_UpperCamelCase							,       additional_params_copy_to_batched_inputs=_UpperCamelCase							,       )
        else:
           self._test_inference_batch_consistent(
               additional_params_copy_to_batched_inputs=_UpperCamelCase				)





     @skip_mps
     def 		UpperCamelCase(    self				):
        return super().test_dict_tuple_outputs_equivalent()





     @skip_mps
     def 		UpperCamelCase(    self				):
        return super().test_save_load_local()





     @skip_mps
     def 		UpperCamelCase(    self				):
        return super().test_save_load_optional_components()




@nightly
class  __UpperCamelCase						(       unittest.TestCase						):





     def 		UpperCamelCase(    self				):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()





     def 		UpperCamelCase(    self				):
        _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=_UpperCamelCase				)

        _UpperCAmelCase											=      torch.manual_seed(0				)
        _UpperCAmelCase											=      pipeline(
            '''horse'''							,       num_images_per_prompt=1							,       generator=_UpperCamelCase							,       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  __UpperCamelCase						(       unittest.TestCase						):





     def 		UpperCamelCase(    self				):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()





     def 		UpperCamelCase(    self				):
        _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(_UpperCamelCase				)
        pipeline.set_progress_bar_config(disable=_UpperCamelCase				)

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

        _UpperCAmelCase											=      output.images[0]

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

        assert_mean_pixel_difference(_UpperCamelCase							,       _UpperCamelCase				)





     def 		UpperCamelCase(    self				):
        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(_UpperCamelCase				)
        pipe.set_progress_bar_config(disable=_UpperCamelCase				)
        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






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


UpperCAmelCase_											=  {
    "configuration_oneformer": ["ONEFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP", "OneFormerConfig"],
    "processing_oneformer": ["OneFormerProcessor"],
}

try:
   if not is_vision_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase_											=  ["OneFormerImageProcessor"]

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase_											=  [
       "ONEFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
       "OneFormerForUniversalSegmentation",
       "OneFormerModel",
       "OneFormerPreTrainedModel",
   ]

if TYPE_CHECKING:
   from .configuration_oneformer import ONEFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, OneFormerConfig
   from .processing_oneformer import OneFormerProcessor

   try:
      if not is_vision_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .image_processing_oneformer import OneFormerImageProcessor
   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_oneformer import (
          ONEFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
          OneFormerForUniversalSegmentation,
          OneFormerModel,
          OneFormerPreTrainedModel,
      )


else:
   import sys

   UpperCAmelCase_											=  _LazyModule(__name__, globals()["__file__"], _import_structure)