




import random
import unittest

import numpy as np
import torch
from PIL import Image
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorEmbaEmbPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()
class 				__a (			__lowerCamelCase							,					unittest.TestCase					):


     """simple docstring"""


     _A    :   str     =  KandinskyVaaPriorEmbaEmbPipeline
     _A    :   Dict     =  ["prompt", "image"]
     _A    :   List[str]     =  ["prompt", "image"]
     _A    :   Union[str, Any]     =  [
         "num_images_per_prompt",
         "strength",
         "generator",
         "num_inference_steps",
         "latents",
         "negative_prompt",
         "guidance_scale",
         "output_type",
         "return_dict",
     ]
     _A    :   int     =  False




     @property
     def   __A    (     self      :			Dict						)		->       Dict:


         '''simple docstring'''





         return 3_2




     @property
     def   __A    (     self      :			Any						)		->       List[str]:


         '''simple docstring'''





         return 3_2




     @property
     def   __A    (     self      :			Optional[Any]						)		->       Optional[int]:


         '''simple docstring'''





         return self.time_input_dim




     @property
     def   __A    (     self      :			List[str]						)		->       Dict:


         '''simple docstring'''





         return self.time_input_dim * 4




     @property
     def   __A    (     self      :			Optional[int]						)		->       Any:


         '''simple docstring'''





         return 1_0_0




     @property
     def   __A    (     self      :			Any						)		->       str:


         '''simple docstring'''





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




     @property
     def   __A    (     self      :			str						)		->       Optional[Any]:


         '''simple docstring'''





         torch.manual_seed(0						)
         SCREAMING_SNAKE_CASE__       =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=3_7    ,layer_norm_eps=1e-05    ,num_attention_heads=4    ,num_hidden_layers=5    ,pad_token_id=1    ,vocab_size=1_0_0_0    ,)
         return CLIPTextModelWithProjection(_UpperCamelCase						)




     @property
     def   __A    (     self      :			List[Any]						)		->       str:


         '''simple docstring'''





         torch.manual_seed(0						)

         SCREAMING_SNAKE_CASE__       ={
             """num_attention_heads""": 2,
             """attention_head_dim""": 1_2,
             """embedding_dim""": self.text_embedder_hidden_size,
             """num_layers""": 1,
         }

         SCREAMING_SNAKE_CASE__       =PriorTransformer(**_UpperCamelCase						)
         # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
         SCREAMING_SNAKE_CASE__       =nn.Parameter(torch.ones(model.clip_std.shape						)						)
         return model




     @property
     def   __A    (     self      :			Optional[int]						)		->       Optional[int]:


         '''simple docstring'''





         torch.manual_seed(0						)
         SCREAMING_SNAKE_CASE__       =CLIPVisionConfig(
             hidden_size=self.text_embedder_hidden_size    ,image_size=2_2_4    ,projection_dim=self.text_embedder_hidden_size    ,intermediate_size=3_7    ,num_attention_heads=4    ,num_channels=3    ,num_hidden_layers=5    ,patch_size=1_4    ,)

         SCREAMING_SNAKE_CASE__       =CLIPVisionModelWithProjection(_UpperCamelCase						)
         return model




     @property
     def   __A    (     self      :			Union[str, Any]						)		->       Any:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       =CLIPImageProcessor(
             crop_size=2_2_4    ,do_center_crop=_UpperCamelCase    ,do_normalize=_UpperCamelCase    ,do_resize=_UpperCamelCase    ,image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]    ,image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]    ,resample=3    ,size=2_2_4    ,)

         return image_processor




     def   __A    (     self      :			str						)		->       Any:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       =self.dummy_prior
         SCREAMING_SNAKE_CASE__       =self.dummy_image_encoder
         SCREAMING_SNAKE_CASE__       =self.dummy_text_encoder
         SCREAMING_SNAKE_CASE__       =self.dummy_tokenizer
         SCREAMING_SNAKE_CASE__       =self.dummy_image_processor

         SCREAMING_SNAKE_CASE__       =UnCLIPScheduler(
             variance_type="""fixed_small_log"""    ,prediction_type="""sample"""    ,num_train_timesteps=1_0_0_0    ,clip_sample=_UpperCamelCase    ,clip_sample_range=10.0    ,)

         SCREAMING_SNAKE_CASE__       ={
             """prior""": prior,
             """image_encoder""": image_encoder,
             """text_encoder""": text_encoder,
             """tokenizer""": tokenizer,
             """scheduler""": scheduler,
             """image_processor""": image_processor,
         }

         return components




     def   __A    (     self      :			Dict    ,_UpperCamelCase      :			Dict    ,_UpperCamelCase      :			Union[str, Any]=0						)		->       List[Any]:


         '''simple docstring'''





         if str(_UpperCamelCase						).startswith("""mps"""						):
             SCREAMING_SNAKE_CASE__       =torch.manual_seed(_UpperCamelCase						)
         else:
             SCREAMING_SNAKE_CASE__       =torch.Generator(device=_UpperCamelCase						).manual_seed(_UpperCamelCase						)

         SCREAMING_SNAKE_CASE__       =floats_tensor((1, 3, 6_4, 6_4)    ,rng=random.Random(_UpperCamelCase						)						).to(_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__       =image.cpu().permute(0    ,2    ,3    ,1						)[0]
         SCREAMING_SNAKE_CASE__       =Image.fromarray(np.uinta(_UpperCamelCase						)						).convert("""RGB"""						).resize((2_5_6, 2_5_6)						)

         SCREAMING_SNAKE_CASE__       ={
             """prompt""": """horse""",
             """image""": init_image,
             """strength""": 0.5,
             """generator""": generator,
             """guidance_scale""": 4.0,
             """num_inference_steps""": 2,
             """output_type""": """np""",
         }
         return inputs




     def   __A    (     self      :			int						)		->       Dict:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       ="""cpu"""

         SCREAMING_SNAKE_CASE__       =self.get_dummy_components()

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

         pipe.set_progress_bar_config(disable=_UpperCamelCase						)

         SCREAMING_SNAKE_CASE__       =pipe(**self.get_dummy_inputs(_UpperCamelCase						)						)
         SCREAMING_SNAKE_CASE__       =output.image_embeds

         SCREAMING_SNAKE_CASE__       =pipe(
             **self.get_dummy_inputs(_UpperCamelCase						)    ,return_dict=_UpperCamelCase    ,)[0]

         SCREAMING_SNAKE_CASE__       =image[0, -1_0:]
         SCREAMING_SNAKE_CASE__       =image_from_tuple[0, -1_0:]

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

         SCREAMING_SNAKE_CASE__       =np.array(
             [
                 0.107_1284,
                 1.333_0271,
                 0.6126_0223,
                 -0.669_1065,
                 -0.384_6852,
                 -1.030_3661,
                 0.2271_6111,
                 0.0334_8901,
                 0.3004_0675,
                 -0.2480_5029,
             ]						)

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




     @skip_mps
     def   __A    (     self      :			Optional[int]						)		->       Optional[int]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       =torch_device == """cpu"""
         SCREAMING_SNAKE_CASE__       =True
         SCREAMING_SNAKE_CASE__       =False

         self._test_inference_batch_single_identical(
             test_max_difference=_UpperCamelCase    ,relax_max_difference=_UpperCamelCase    ,test_mean_pixel_difference=_UpperCamelCase    ,)




     @skip_mps
     def   __A    (     self      :			Optional[int]						)		->       Optional[int]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       =torch_device == """cpu"""
         SCREAMING_SNAKE_CASE__       =False

         self._test_attention_slicing_forward_pass(
             test_max_difference=_UpperCamelCase    ,test_mean_pixel_difference=_UpperCamelCase    ,)












import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only

from utils import save_json



def      UpperCAmelCase_   (     __UpperCamelCase     ):
    SCREAMING_SNAKE_CASE__       =filter(lambda __UpperCamelCase     : p.requires_grad,    model.parameters()     )
    SCREAMING_SNAKE_CASE__       =sum([np.prod(p.size()     ) for p in model_parameters]     )
    return params


lowerCamelCase_  						=       logging.getLogger(__name__)
class 				__a (			pl.Callback					):


     """simple docstring"""


     def   __A    (     self      :			Optional[int]    ,_UpperCamelCase      :			Optional[Any]    ,_UpperCamelCase      :			Any						)		->       int:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       ={f"""lr_group_{i}""": param["""lr"""] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups						)}
         pl_module.logger.log_metrics(_UpperCamelCase						)




     @rank_zero_only
     def   __A    (     self      :			Optional[int]    ,_UpperCamelCase      :			pl.Trainer    ,_UpperCamelCase      :			pl.LightningModule    ,_UpperCamelCase      :			str    ,_UpperCamelCase      :			Tuple=True						)		->       None:


         '''simple docstring'''





         logger.info(f"""***** {type_path} results at step {trainer.global_step:05d} *****"""						)
         SCREAMING_SNAKE_CASE__       =trainer.callback_metrics
         trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ["""log""", """progress_bar""", """preds"""]}						)
         # Log results
         SCREAMING_SNAKE_CASE__       =Path(pl_module.hparams.output_dir						)
         if type_path == "test":
             SCREAMING_SNAKE_CASE__       =od / """test_results.txt"""
             SCREAMING_SNAKE_CASE__       =od / """test_generations.txt"""
         else:
             # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
             # If people want this it will be easy enough to add back.
             SCREAMING_SNAKE_CASE__       =od / f"""{type_path}_results/{trainer.global_step:05d}.txt"""
             SCREAMING_SNAKE_CASE__       =od / f"""{type_path}_generations/{trainer.global_step:05d}.txt"""
             results_file.parent.mkdir(exist_ok=_UpperCamelCase						)
             generations_file.parent.mkdir(exist_ok=_UpperCamelCase						)
         with open(_UpperCamelCase    ,"""a+"""						) as writer:
             for key in sorted(_UpperCamelCase						):
                 if key in ["log", "progress_bar", "preds"]:
                     continue
                 SCREAMING_SNAKE_CASE__       =metrics[key]
                 if isinstance(_UpperCamelCase    ,torch.Tensor						):
                     SCREAMING_SNAKE_CASE__       =val.item()
                 SCREAMING_SNAKE_CASE__       =f"""{key}: {val:.6f}\n"""
                 writer.write(_UpperCamelCase						)

         if not save_generations:
             return

         if "preds" in metrics:
             SCREAMING_SNAKE_CASE__       ="""\n""".join(metrics["""preds"""]						)
             generations_file.open("""w+"""						).write(_UpperCamelCase						)




     @rank_zero_only
     def   __A    (     self      :			Tuple    ,_UpperCamelCase      :			int    ,_UpperCamelCase      :			Dict						)		->       Dict:


         '''simple docstring'''





         try:
             SCREAMING_SNAKE_CASE__       =pl_module.model.model.num_parameters()
         except AttributeError:
             SCREAMING_SNAKE_CASE__       =pl_module.model.num_parameters()

         SCREAMING_SNAKE_CASE__       =count_trainable_parameters(_UpperCamelCase						)
         # mp stands for million parameters
         trainer.logger.log_metrics({"""n_params""": npars, """mp""": npars / 1e6, """grad_mp""": n_trainable_pars / 1e6}						)




     @rank_zero_only
     def   __A    (     self      :			Optional[Any]    ,_UpperCamelCase      :			pl.Trainer    ,_UpperCamelCase      :			pl.LightningModule						)		->       Any:


         '''simple docstring'''





         save_json(pl_module.metrics    ,pl_module.metrics_save_path						)
         return self._write_logs(_UpperCamelCase    ,_UpperCamelCase    ,"""test"""						)




     @rank_zero_only
     def   __A    (     self      :			Optional[Any]    ,_UpperCamelCase      :			pl.Trainer    ,_UpperCamelCase      :			List[Any]						)		->       Dict:


         '''simple docstring'''





         save_json(pl_module.metrics    ,pl_module.metrics_save_path						)
         # Uncommenting this will save val generations
         # return self._write_logs(trainer, pl_module, "valid")






def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase=1,    __UpperCamelCase=False     ):
    if metric == "rouge2":
        SCREAMING_SNAKE_CASE__       ="""{val_avg_rouge2:.4f}-{step_count}"""
    elif metric == "bleu":
        SCREAMING_SNAKE_CASE__       ="""{val_avg_bleu:.4f}-{step_count}"""
    elif metric == "loss":
        SCREAMING_SNAKE_CASE__       ="""{val_avg_loss:.4f}-{step_count}"""
    else:
        raise NotImplementedError(
            f"""seq2seq callbacks only support rouge2, bleu and loss, got {metric}, You can make your own by adding to"""
            """ this function."""     )

    SCREAMING_SNAKE_CASE__       =ModelCheckpoint(
        dirpath=__UpperCamelCase,    filename=__UpperCamelCase,    monitor=f"""val_{metric}""",    mode="""min""" if """loss""" in metric else """max""",    save_top_k=__UpperCamelCase,    )
    return checkpoint_callback




def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase     ):
    return EarlyStopping(
        monitor=f"""val_{metric}""",    mode="""min""" if """loss""" in metric else """max""",    patience=__UpperCamelCase,    verbose=__UpperCamelCase,    )






