



"""simple docstring"""

import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()




class  _UpperCAmelCase (    lowerCAmelCase__					,  lowerCAmelCase__					,  lowerCAmelCase__					,  unittest.TestCase):
 _lowerCAmelCase     :    Dict	      =	StableUnCLIPPipeline
 _lowerCAmelCase     :    Optional[Any]	      =	TEXT_TO_IMAGE_PARAMS
 _lowerCAmelCase     :    int	      =	TEXT_TO_IMAGE_BATCH_PARAMS
 _lowerCAmelCase     :    Tuple	      =	TEXT_TO_IMAGE_IMAGE_PARAMS
 _lowerCAmelCase     :    List[Any]	      =	TEXT_TO_IMAGE_IMAGE_PARAMS

 # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
 _lowerCAmelCase     :    Union[str, Any]	      =	False

 def 		_snake_case      (  self	: int       ):
     snake_case_				:  Tuple       =			32
     snake_case_				:  Dict       =			embedder_hidden_size

     # prior components

     torch.manual_seed(0       )
     snake_case_				:  List[str]       =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       )

     torch.manual_seed(0       )
     snake_case_				:  List[Any]       =			CLIPTextModelWithProjection(
         CLIPTextConfig(
             bos_token_id=0						,  eos_token_id=2						,  hidden_size=lowercase_						,  projection_dim=lowercase_						,  intermediate_size=37						,  layer_norm_eps=1E-05						,  num_attention_heads=4						,  num_hidden_layers=5						,  pad_token_id=1						,  vocab_size=1000						,  )       )

     torch.manual_seed(0       )
     snake_case_				:  List[str]       =			PriorTransformer(
         num_attention_heads=2						,  attention_head_dim=12						,  embedding_dim=lowercase_						,  num_layers=1						,  )

     torch.manual_seed(0       )
     snake_case_				:  Tuple       =			DDPMScheduler(
         variance_type='''fixed_small_log'''						,  prediction_type='''sample'''						,  num_train_timesteps=1000						,  clip_sample=lowercase_						,  clip_sample_range=5.0						,  beta_schedule='''squaredcos_cap_v2'''						,  )

     # regular denoising components

     torch.manual_seed(0       )
     snake_case_				:  Optional[Any]       =			StableUnCLIPImageNormalizer(embedding_dim=lowercase_       )
     snake_case_				:  List[str]       =			DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''       )

     torch.manual_seed(0       )
     snake_case_				:  Optional[int]       =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       )

     torch.manual_seed(0       )
     snake_case_				:  Tuple       =			CLIPTextModel(
         CLIPTextConfig(
             bos_token_id=0						,  eos_token_id=2						,  hidden_size=lowercase_						,  projection_dim=32						,  intermediate_size=37						,  layer_norm_eps=1E-05						,  num_attention_heads=4						,  num_hidden_layers=5						,  pad_token_id=1						,  vocab_size=1000						,  )       )

     torch.manual_seed(0       )
     snake_case_				:  Optional[int]       =			UNetaDConditionModel(
         sample_size=32						,  in_channels=4						,  out_channels=4						,  down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''')						,  up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''')						,  block_out_channels=(32, 64)						,  attention_head_dim=(2, 4)						,  class_embed_type='''projection'''						,  projection_class_embeddings_input_dim=embedder_projection_dim * 2						,  cross_attention_dim=lowercase_						,  layers_per_block=1						,  upcast_attention=lowercase_						,  use_linear_projection=lowercase_						,  )

     torch.manual_seed(0       )
     snake_case_				:  Union[str, Any]       =			DDIMScheduler(
         beta_schedule='''scaled_linear'''						,  beta_start=0.0_00_85						,  beta_end=0.0_12						,  prediction_type='''v_prediction'''						,  set_alpha_to_one=lowercase_						,  steps_offset=1						,  )

     torch.manual_seed(0       )
     snake_case_				:  Optional[int]       =			AutoencoderKL()

     snake_case_				:  int       =			{
         # prior components
         '''prior_tokenizer''': prior_tokenizer,
         '''prior_text_encoder''': prior_text_encoder,
         '''prior''': prior,
         '''prior_scheduler''': prior_scheduler,
         # image noising components
         '''image_normalizer''': image_normalizer,
         '''image_noising_scheduler''': image_noising_scheduler,
         # regular denoising components
         '''tokenizer''': tokenizer,
         '''text_encoder''': text_encoder,
         '''unet''': unet,
         '''scheduler''': scheduler,
         '''vae''': vae,
     }

     return components

 def 		_snake_case      (  self	: Optional[Any]						,  lowercase_	: List[Any]						,  lowercase_	: Optional[int]=0       ):
     if str(lowercase_       ).startswith('''mps'''       ):
         snake_case_				:  Optional[Any]       =			torch.manual_seed(lowercase_       )
     else:
         snake_case_				:  List[str]       =			torch.Generator(device=lowercase_       ).manual_seed(lowercase_       )
     snake_case_				:  Tuple       =			{
         '''prompt''': '''A painting of a squirrel eating a burger''',
         '''generator''': generator,
         '''num_inference_steps''': 2,
         '''prior_num_inference_steps''': 2,
         '''output_type''': '''numpy''',
     }
     return inputs

 def 		_snake_case      (  self	: List[str]       ):
     snake_case_				:  Union[str, Any]       =			torch_device == '''cpu'''

     self._test_attention_slicing_forward_pass(test_max_difference=lowercase_       )






 def 		_snake_case      (  self	: Any       ):
     snake_case_				:  Any       =			torch_device in ['''cpu''', '''mps''']

     self._test_inference_batch_single_identical(test_max_difference=lowercase_       )





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

 def 		_snake_case      (  self	: List[str]       ):
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()

 def 		_snake_case      (  self	: List[str]       ):
     snake_case_				:  int       =			load_numpy(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'''       )

     snake_case_				:  List[Any]       =			StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''						,  torch_dtype=torch.floataa       )
     pipe.to(lowercase_       )
     pipe.set_progress_bar_config(disable=lowercase_       )
     # stable unclip will oom when integration tests are run on a V100,
     # so turn on memory savings
     pipe.enable_attention_slicing()
     pipe.enable_sequential_cpu_offload()

     snake_case_				:  Optional[int]       =			torch.Generator(device='''cpu'''       ).manual_seed(0       )
     snake_case_				:  Optional[Any]       =			pipe('''anime turle'''						,  generator=lowercase_						,  output_type='''np'''       )

     snake_case_				:  Dict       =			output.images[0]

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

     assert_mean_pixel_difference(lowercase_						,  lowercase_       )






 def 		_snake_case      (  self	: Tuple       ):
     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()

     snake_case_				:  List[str]       =			StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''						,  torch_dtype=torch.floataa       )
     snake_case_				:  Any       =			pipe.to(lowercase_       )
     pipe.set_progress_bar_config(disable=lowercase_       )
     pipe.enable_attention_slicing()
     pipe.enable_sequential_cpu_offload()

     snake_case_				:  Dict       =			pipe(
         '''anime turtle'''						,  prior_num_inference_steps=2						,  num_inference_steps=2						,  output_type='''np'''						,  )

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










"""simple docstring"""

import tensorflow as tf

from ...tf_utils import shape_list




class  _UpperCAmelCase (    tf.keras.layers.Layer):

 def __init__(  self	: int						,  lowercase_	: Tuple						,  lowercase_	: Union[str, Any]						,  lowercase_	: Tuple						,  lowercase_	: List[Any]						,  lowercase_	: Tuple=1						,  lowercase_	: List[str]=False						,  **lowercase_	: Optional[Any]       ):
     super().__init__(**lowercase_       )

     snake_case_				:  int       =			vocab_size
     snake_case_				:  Union[str, Any]       =			d_embed
     snake_case_				:  Optional[int]       =			d_proj

     snake_case_				:  int       =			cutoffs + [vocab_size]
     snake_case_				:  Optional[int]       =			[0] + self.cutoffs
     snake_case_				:  List[str]       =			div_val

     snake_case_				:  int       =			self.cutoffs[0]
     snake_case_				:  Optional[Any]       =			len(self.cutoffs       ) - 1
     snake_case_				:  Any       =			self.shortlist_size + self.n_clusters
     snake_case_				:  Dict       =			keep_order

     snake_case_				:  Tuple       =			[]
     snake_case_				:  Optional[Any]       =			[]

 def 		_snake_case      (  self	: Dict						,  lowercase_	: int       ):
     if self.n_clusters > 0:
         snake_case_				:  Tuple       =			self.add_weight(
             shape=(self.n_clusters, self.d_embed)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name='''cluster_weight'''       )
         snake_case_				:  List[str]       =			self.add_weight(
             shape=(self.n_clusters,)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name='''cluster_bias'''       )

     if self.div_val == 1:
         for i in range(len(self.cutoffs       )       ):
             if self.d_proj != self.d_embed:
                 snake_case_				:  List[Any]       =			self.add_weight(
                     shape=(self.d_embed, self.d_proj)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name=f"out_projs_._{i}"						,  )
                 self.out_projs.append(lowercase_       )
             else:
                 self.out_projs.append(lowercase_       )
             snake_case_				:  Any       =			self.add_weight(
                 shape=(self.vocab_size, self.d_embed)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name=f"out_layers_._{i}_._weight"						,  )
             snake_case_				:  Any       =			self.add_weight(
                 shape=(self.vocab_size,)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name=f"out_layers_._{i}_._bias"						,  )
             self.out_layers.append((weight, bias)       )
     else:
         for i in range(len(self.cutoffs       )       ):
             snake_case_,			snake_case_				:  str       =			self.cutoff_ends[i], self.cutoff_ends[i + 1]
             snake_case_				:  List[str]       =			self.d_embed // (self.div_val**i)

             snake_case_				:  Optional[Any]       =			self.add_weight(
                 shape=(d_emb_i, self.d_proj)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name=f"out_projs_._{i}"       )
             self.out_projs.append(lowercase_       )
             snake_case_				:  Any       =			self.add_weight(
                 shape=(r_idx - l_idx, d_emb_i)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name=f"out_layers_._{i}_._weight"						,  )
             snake_case_				:  List[Any]       =			self.add_weight(
                 shape=(r_idx - l_idx,)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name=f"out_layers_._{i}_._bias"						,  )
             self.out_layers.append((weight, bias)       )
     super().build(lowercase_       )

 @staticmethod
 def 		_snake_case      (  lowercase_	: Dict						,  lowercase_	: Union[str, Any]						,  lowercase_	: Optional[int]						,  lowercase_	: str=None       ):
     snake_case_				:  str       =			x
     if proj is not None:
         snake_case_				:  str       =			tf.einsum('''ibd,ed->ibe'''						,  lowercase_						,  lowercase_       )
     return tf.einsum('''ibd,nd->ibn'''						,  lowercase_						,  lowercase_       ) + b

 @staticmethod
 def 		_snake_case      (  lowercase_	: Tuple						,  lowercase_	: Union[str, Any]       ):
     snake_case_				:  Optional[int]       =			shape_list(lowercase_       )
     snake_case_				:  Any       =			tf.range(lp_size[0]						,  dtype=target.dtype       )
     snake_case_				:  Optional[int]       =			tf.stack([r, target]						,  1       )
     return tf.gather_nd(lowercase_						,  lowercase_       )






 def 		_snake_case      (  self	: Tuple						,  lowercase_	: Union[str, Any]						,  lowercase_	: int						,  lowercase_	: Tuple=True						,  lowercase_	: Tuple=False       ):
     snake_case_				:  Optional[int]       =			0
     if self.n_clusters == 0:
         snake_case_				:  List[str]       =			self._logit(lowercase_						,  self.out_layers[0][0]						,  self.out_layers[0][1]						,  self.out_projs[0]       )
         if target is not None:
             snake_case_				:  List[Any]       =			tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lowercase_						,  logits=lowercase_       )
         snake_case_				:  List[Any]       =			tf.nn.log_softmax(lowercase_						,  axis=-1       )
     else:
         snake_case_				:  Optional[int]       =			shape_list(lowercase_       )
         snake_case_				:  Any       =			[]
         snake_case_				:  List[str]       =			tf.zeros(hidden_sizes[:2]       )
         for i in range(len(self.cutoffs       )       ):
             snake_case_,			snake_case_				:  Union[str, Any]       =			self.cutoff_ends[i], self.cutoff_ends[i + 1]
             if target is not None:
                 snake_case_				:  Optional[Any]       =			(target >= l_idx) & (target < r_idx)
                 snake_case_				:  Any       =			tf.where(lowercase_       )
                 snake_case_				:  List[Any]       =			tf.boolean_mask(lowercase_						,  lowercase_       ) - l_idx

             if self.div_val == 1:
                 snake_case_				:  Optional[int]       =			self.out_layers[0][0][l_idx:r_idx]
                 snake_case_				:  List[Any]       =			self.out_layers[0][1][l_idx:r_idx]
             else:
                 snake_case_				:  Dict       =			self.out_layers[i][0]
                 snake_case_				:  str       =			self.out_layers[i][1]

             if i == 0:
                 snake_case_				:  Tuple       =			tf.concat([cur_W, self.cluster_weight]						,  0       )
                 snake_case_				:  List[str]       =			tf.concat([cur_b, self.cluster_bias]						,  0       )

                 snake_case_				:  List[Any]       =			self._logit(lowercase_						,  lowercase_						,  lowercase_						,  self.out_projs[0]       )
                 snake_case_				:  Union[str, Any]       =			tf.nn.log_softmax(lowercase_       )
                 out.append(head_logprob[..., : self.cutoffs[0]]       )
                 if target is not None:
                     snake_case_				:  int       =			tf.boolean_mask(lowercase_						,  lowercase_       )
                     snake_case_				:  str       =			self._gather_logprob(lowercase_						,  lowercase_       )
             else:
                 snake_case_				:  Union[str, Any]       =			self._logit(lowercase_						,  lowercase_						,  lowercase_						,  self.out_projs[i]       )
                 snake_case_				:  List[str]       =			tf.nn.log_softmax(lowercase_       )
                 snake_case_				:  Dict       =			self.cutoffs[0] + i - 1  # No probability for the head cluster
                 snake_case_				:  Union[str, Any]       =			head_logprob[..., cluster_prob_idx, None] + tail_logprob
                 out.append(lowercase_       )
                 if target is not None:
                     snake_case_				:  Optional[int]       =			tf.boolean_mask(lowercase_						,  lowercase_       )
                     snake_case_				:  Dict       =			tf.boolean_mask(lowercase_						,  lowercase_       )
                     snake_case_				:  Union[str, Any]       =			self._gather_logprob(lowercase_						,  lowercase_       )
                     cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
             if target is not None:
                 loss += tf.scatter_nd(lowercase_						,  -cur_logprob						,  shape_list(lowercase_       )       )
         snake_case_				:  Union[str, Any]       =			tf.concat(lowercase_						,  axis=-1       )

     if target is not None:
         if return_mean:
             snake_case_				:  Dict       =			tf.reduce_mean(lowercase_       )
         # Add the training-time loss value to the layer using `self.add_loss()`.
         self.add_loss(lowercase_       )

         # Log the loss as a metric (we could log arbitrary metrics,
         # including different metrics for training and inference.
         self.add_metric(lowercase_						,  name=self.name						,  aggregation='''mean''' if return_mean else ''''''       )

     return out





