'''simple docstring'''





import gc
import random
import unittest

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

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()




class        __A      ( UpperCamelCase__		,	unittest.TestCase	):
    a__       :						List[str]          =  UnCLIPImageVariationPipeline
    a__       :						str          =  IMAGE_VARIATION_PARAMS - {"""height""", """width""", """guidance_scale"""}
    a__       :						Union[str, Any]          =  IMAGE_VARIATION_BATCH_PARAMS

    a__       :						List[Any]          =  [
        """generator""",
        """return_dict""",
        """decoder_num_inference_steps""",
        """super_res_num_inference_steps""",
    ]
    a__       :						Optional[int]          =  False



    @property
    def 					_lowercase						(self    :    List[str] ):
         return 32



    @property
    def 					_lowercase						(self    :    Union[str, Any] ):
         return 32



    @property
    def 					_lowercase						(self    :    Optional[Any] ):
         return self.time_input_dim



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



    @property
    def 					_lowercase						(self    :    List[str] ):
         return 100



    @property
    def 					_lowercase						(self    :    Union[str, Any] ):
         UpperCAmelCase_           =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )
         return tokenizer



    @property
    def 					_lowercase						(self    :    Dict ):
         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(__a )



    @property
    def 					_lowercase						(self    :    List[str] ):
         torch.manual_seed(0 )
         UpperCAmelCase_           =   CLIPVisionConfig(
             hidden_size=self.text_embedder_hidden_size   ,			projection_dim=self.text_embedder_hidden_size   ,			num_hidden_layers=5   ,			num_attention_heads=4   ,			image_size=32   ,			intermediate_size=37   ,			patch_size=1   ,			)
         return CLIPVisionModelWithProjection(__a )



    @property
    def 					_lowercase						(self    :    List[Any] ):
         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(**__a )
         return model



    @property
    def 					_lowercase						(self    :    str ):
         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(**__a )
         return model



    @property
    def 					_lowercase						(self    :    Tuple ):
         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 					_lowercase						(self    :    List[str] ):
         torch.manual_seed(0 )

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



    @property
    def 					_lowercase						(self    :    List[Any] ):
         # 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 					_lowercase						(self    :    int ):
         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="learned_range"   ,			prediction_type="epsilon"   ,			num_train_timesteps=1000   ,			)

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

         UpperCAmelCase_           =   CLIPImageProcessor(crop_size=32   ,			size=32 )

         UpperCAmelCase_           =   self.dummy_image_encoder

         return {
             "decoder": decoder,
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "text_proj": text_proj,
             "feature_extractor": feature_extractor,
             "image_encoder": image_encoder,
             "super_res_first": super_res_first,
             "super_res_last": super_res_last,
             "decoder_scheduler": decoder_scheduler,
             "super_res_scheduler": super_res_scheduler,
         }



    def 					_lowercase						(self    :    Union[str, Any]   ,			__a    :    Tuple   ,			__a    :    Optional[int]=0   ,			__a    :    int=True ):
         UpperCAmelCase_           =   floats_tensor((1, 3, 32, 32)   ,			rng=random.Random(__a ) ).to(__a )
         if str(__a ).startswith("mps" ):
              UpperCAmelCase_           =   torch.manual_seed(__a )
         else:
              UpperCAmelCase_           =   torch.Generator(device=__a ).manual_seed(__a )

         if pil_image:
              UpperCAmelCase_           =   input_image * 0.5 + 0.5
              UpperCAmelCase_           =   input_image.clamp(0   ,			1 )
              UpperCAmelCase_           =   input_image.cpu().permute(0   ,			2   ,			3   ,			1 ).float().numpy()
              UpperCAmelCase_           =   DiffusionPipeline.numpy_to_pil(__a )[0]

         return {
             "image": input_image,
             "generator": generator,
             "decoder_num_inference_steps": 2,
             "super_res_num_inference_steps": 2,
             "output_type": "np",
         }



    def 					_lowercase						(self    :    Tuple ):
         UpperCAmelCase_           =   "cpu"

         UpperCAmelCase_           =   self.get_dummy_components()

         UpperCAmelCase_           =   self.pipeline_class(**__a )
         UpperCAmelCase_           =   pipe.to(__a )

         pipe.set_progress_bar_config(disable=__a )

         UpperCAmelCase_           =   self.get_dummy_inputs(__a   ,			pil_image=__a )

         UpperCAmelCase_           =   pipe(**__a )
         UpperCAmelCase_           =   output.images

         UpperCAmelCase_           =   self.get_dummy_inputs(__a   ,			pil_image=__a )

         UpperCAmelCase_           =   pipe(
             **__a   ,			return_dict=__a   ,			)[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.99_97,
                 0.00_02,
                 0.99_97,
                 0.99_97,
                 0.99_69,
                 0.00_23,
                 0.99_97,
                 0.99_69,
                 0.99_70,
             ] )

         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 					_lowercase						(self    :    Any ):
         UpperCAmelCase_           =   "cpu"

         UpperCAmelCase_           =   self.get_dummy_components()

         UpperCAmelCase_           =   self.pipeline_class(**__a )
         UpperCAmelCase_           =   pipe.to(__a )

         pipe.set_progress_bar_config(disable=__a )

         UpperCAmelCase_           =   self.get_dummy_inputs(__a   ,			pil_image=__a )

         UpperCAmelCase_           =   pipe(**__a )
         UpperCAmelCase_           =   output.images

         UpperCAmelCase_           =   self.get_dummy_inputs(__a   ,			pil_image=__a )

         UpperCAmelCase_           =   pipe(
             **__a   ,			return_dict=__a   ,			)[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.99_97, 0.00_03, 0.99_97, 0.99_97, 0.99_70, 0.00_24, 0.99_97, 0.99_71, 0.99_71] )

         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 					_lowercase						(self    :    List[str] ):
         UpperCAmelCase_           =   "cpu"

         UpperCAmelCase_           =   self.get_dummy_components()

         UpperCAmelCase_           =   self.pipeline_class(**__a )
         UpperCAmelCase_           =   pipe.to(__a )

         pipe.set_progress_bar_config(disable=__a )

         UpperCAmelCase_           =   self.get_dummy_inputs(__a   ,			pil_image=__a )
         UpperCAmelCase_           =   [
             pipeline_inputs["image"],
             pipeline_inputs["image"],
         ]

         UpperCAmelCase_           =   pipe(**__a )
         UpperCAmelCase_           =   output.images

         UpperCAmelCase_           =   self.get_dummy_inputs(__a   ,			pil_image=__a )
         UpperCAmelCase_           =   [
             tuple_pipeline_inputs["image"],
             tuple_pipeline_inputs["image"],
         ]

         UpperCAmelCase_           =   pipe(
             **__a   ,			return_dict=__a   ,			)[0]

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

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

         UpperCAmelCase_           =   np.array(
             [
                 0.99_97,
                 0.99_89,
                 0.00_08,
                 0.00_21,
                 0.99_60,
                 0.00_18,
                 0.00_14,
                 0.00_02,
                 0.99_33,
             ] )

         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 					_lowercase						(self    :    Dict ):
         UpperCAmelCase_           =   torch.device("cpu" )




         class        __A      :
             a__       :						Tuple          =  1




         UpperCAmelCase_           =   self.get_dummy_components()

         UpperCAmelCase_           =   self.pipeline_class(**__a )
         UpperCAmelCase_           =   pipe.to(__a )

         pipe.set_progress_bar_config(disable=__a )

         UpperCAmelCase_           =   torch.Generator(device=__a ).manual_seed(0 )
         UpperCAmelCase_           =   pipe.decoder.dtype
         UpperCAmelCase_           =   1

         UpperCAmelCase_           =   (
             batch_size,
             pipe.decoder.config.in_channels,
             pipe.decoder.config.sample_size,
             pipe.decoder.config.sample_size,
         )
         UpperCAmelCase_           =   pipe.prepare_latents(
             __a   ,			dtype=__a   ,			device=__a   ,			generator=__a   ,			latents=__a   ,			scheduler=DummyScheduler() )

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

         UpperCAmelCase_           =   self.get_dummy_inputs(__a   ,			pil_image=__a )

         UpperCAmelCase_           =   pipe(
             **__a   ,			decoder_latents=__a   ,			super_res_latents=__a ).images

         UpperCAmelCase_           =   self.get_dummy_inputs(__a   ,			pil_image=__a )
         # Don't pass image, instead pass embedding
         UpperCAmelCase_           =   pipeline_inputs.pop("image" )
         UpperCAmelCase_           =   pipe.image_encoder(__a ).image_embeds

         UpperCAmelCase_           =   pipe(
             **__a   ,			decoder_latents=__a   ,			super_res_latents=__a   ,			image_embeddings=__a   ,			).images

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



    @skip_mps
    def 					_lowercase						(self    :    Optional[int] ):
         UpperCAmelCase_           =   torch_device == "cpu"

         # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
         UpperCAmelCase_           =   1E-2

         self._test_attention_slicing_forward_pass(
             test_max_difference=__a   ,			expected_max_diff=__a )



    @skip_mps
    def 					_lowercase						(self    :    Tuple ):
         UpperCAmelCase_           =   torch_device == "cpu"
         UpperCAmelCase_           =   True
         UpperCAmelCase_           =   [
             "decoder_num_inference_steps",
             "super_res_num_inference_steps",
         ]

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



    def 					_lowercase						(self    :    Dict ):
         UpperCAmelCase_           =   [
             "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=__a   ,			additional_params_copy_to_batched_inputs=__a   ,			)
         else:
              self._test_inference_batch_consistent(
                  additional_params_copy_to_batched_inputs=__a )



    @skip_mps
    def 					_lowercase						(self    :    int ):
         return super().test_dict_tuple_outputs_equivalent()



    @skip_mps
    def 					_lowercase						(self    :    str ):
         return super().test_save_load_local()






    @skip_mps
    def 					_lowercase						(self    :    List[str] ):
         return super().test_save_load_optional_components()




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



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






    def 					_lowercase						(self    :    Optional[Any] ):
         UpperCAmelCase_           =   load_image(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png" )
         UpperCAmelCase_           =   load_numpy(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
             "/unclip/karlo_v1_alpha_cat_variation_fp16.npy" )

         UpperCAmelCase_           =   UnCLIPImageVariationPipeline.from_pretrained(
             "kakaobrain/karlo-v1-alpha-image-variations"   ,			torch_dtype=torch.floataa )
         UpperCAmelCase_           =   pipeline.to(__a )
         pipeline.set_progress_bar_config(disable=__a )

         UpperCAmelCase_           =   torch.Generator(device="cpu" ).manual_seed(0 )
         UpperCAmelCase_           =   pipeline(
             __a   ,			generator=__a   ,			output_type="np"   ,			)

         UpperCAmelCase_           =   output.images[0]

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

         assert_mean_pixel_difference(__a   ,			__a   ,			15 )





'''simple docstring'''





import json
import logging
import os
import sys
from pathlib import Path

import finetune_rag

from transformers.file_utils import is_apex_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    require_ray,
    require_torch_gpu,
    require_torch_multi_gpu,
)


logging.basicConfig(level=logging.DEBUG)
SCREAMING_SNAKE_CASE_:     Optional[int]  =logging.getLogger()

SCREAMING_SNAKE_CASE_:     Union[str, Any]  =logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)




class        __A      ( UpperCamelCase__	):



    def 					_lowercase						(self    :    Optional[Any]   ,			__a    :    str ):
         os.makedirs(__a   ,			exist_ok=__a )
         UpperCAmelCase_           =   {"source": "What is love ?", "target": "life"}
         UpperCAmelCase_           =   {"train": 12, "val": 2, "test": 2}
         for split in ["train", "test", "val"]:
              for field in ["source", "target"]:
                   UpperCAmelCase_           =   "\n".join([contents[field]] * n_lines[split] )
                   with open(os.path.join(__a   ,			f"""{split}.{field}""" )   ,			"w" ) as f:
                        f.write(__a )



    def 					_lowercase						(self    :    Optional[int]   ,			__a    :    int   ,			__a    :    str = "pytorch" ):
         UpperCAmelCase_           =   self.get_auto_remove_tmp_dir()
         UpperCAmelCase_           =   os.path.join(__a   ,			"output" )
         UpperCAmelCase_           =   os.path.join(__a   ,			"data" )
         self._create_dummy_data(data_dir=__a )

         UpperCAmelCase_           =   f"""
                --data_dir {data_dir} \
                --output_dir {output_dir} \
                --model_name_or_path facebook/rag-sequence-base \
                --model_type rag_sequence \
                --do_train \
                --do_predict \
                --n_val -1 \
                --val_check_interval 1.0 \
                --train_batch_size 2 \
                --eval_batch_size 1 \
                --max_source_length 25 \
                --max_target_length 25 \
                --val_max_target_length 25 \
                --test_max_target_length 25 \
                --label_smoothing 0.1 \
                --dropout 0.1 \
                --attention_dropout 0.1 \
                --weight_decay 0.001 \
                --adam_epsilon 1e-08 \
                --max_grad_norm 0.1 \
                --lr_scheduler polynomial \
                --learning_rate 3e-04 \
                --num_train_epochs 1 \
                --warmup_steps 4 \
                --gradient_accumulation_steps 1 \
                --distributed-port 8787 \
                --use_dummy_dataset 1 \
                --distributed_retriever {distributed_retriever} \
            """.split()

         if gpus > 0:
              testargs.append(f"""--gpus={gpus}""" )
              if is_apex_available():
                   testargs.append("--fp16" )
         else:
              testargs.append("--gpus=0" )
              testargs.append("--distributed_backend=ddp_cpu" )
              testargs.append("--num_processes=2" )

         UpperCAmelCase_           =   [sys.executable, str(Path(finetune_rag.__file__ ).resolve() )] + testargs
         execute_subprocess_async(__a   ,			env=self.get_env() )

         UpperCAmelCase_           =   os.path.join(__a   ,			"metrics.json" )
         with open(__a ) as f:
              UpperCAmelCase_           =   json.load(__a )
         return result



    @require_torch_gpu
    def 					_lowercase						(self    :    Optional[int] ):
         UpperCAmelCase_           =   self._run_finetune(gpus=1 )
         self.assertGreaterEqual(result["test"][0]["test_avg_em"]   ,			0.2 )



    @require_torch_multi_gpu
    def 					_lowercase						(self    :    Dict ):
         UpperCAmelCase_           =   self._run_finetune(gpus=2 )
         self.assertGreaterEqual(result["test"][0]["test_avg_em"]   ,			0.2 )



    @require_torch_gpu
    @require_ray
    def 					_lowercase						(self    :    Optional[int] ):
         UpperCAmelCase_           =   self._run_finetune(gpus=1   ,			distributed_retriever="ray" )
         self.assertGreaterEqual(result["test"][0]["test_avg_em"]   ,			0.2 )






    @require_torch_multi_gpu
    @require_ray
    def 					_lowercase						(self    :    Any ):
         UpperCAmelCase_           =   self._run_finetune(gpus=1   ,			distributed_retriever="ray" )
         self.assertGreaterEqual(result["test"][0]["test_avg_em"]   ,			0.2 )




