'''simple docstring'''





import gc
import unittest

import numpy as np
import torch
from torch.backends.cuda import sdp_kernel

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)
from diffusers.utils import randn_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_a, require_torch_gpu

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class        __A      ( UpperCamelCase__		,	unittest.TestCase	):
    a__       :						Any          =  ConsistencyModelPipeline
    a__       :						int          =  UNCONDITIONAL_IMAGE_GENERATION_PARAMS
    a__       :						str          =  UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

    # Override required_optional_params to remove num_images_per_prompt
    a__       :						List[Any]          =  frozenset(
        [
            """num_inference_steps""",
            """generator""",
            """latents""",
            """output_type""",
            """return_dict""",
            """callback""",
            """callback_steps""",
        ]	)



    @property
    def 					_lowercase						(self    :    int ):
         UpperCAmelCase_           =   UNetaDModel.from_pretrained(
             "diffusers/consistency-models-test"   ,			subfolder="test_unet"   ,			)
         return unet



    @property
    def 					_lowercase						(self    :    List[Any] ):
         UpperCAmelCase_           =   UNetaDModel.from_pretrained(
             "diffusers/consistency-models-test"   ,			subfolder="test_unet_class_cond"   ,			)
         return unet



    def 					_lowercase						(self    :    List[str]   ,			__a    :    Tuple=False ):
         if class_cond:
              UpperCAmelCase_           =   self.dummy_cond_unet
         else:
              UpperCAmelCase_           =   self.dummy_uncond_unet

         # Default to CM multistep sampler
         UpperCAmelCase_           =   CMStochasticIterativeScheduler(
             num_train_timesteps=40   ,			sigma_min=0.0_02   ,			sigma_max=80.0   ,			)

         UpperCAmelCase_           =   {
             "unet": unet,
             "scheduler": scheduler,
         }

         return components



    def 					_lowercase						(self    :    Any   ,			__a    :    List[str]   ,			__a    :    Tuple=0 ):
         if str(__a ).startswith("mps" ):
              UpperCAmelCase_           =   torch.manual_seed(__a )
         else:
              UpperCAmelCase_           =   torch.Generator(device=__a ).manual_seed(__a )

         UpperCAmelCase_           =   {
             "batch_size": 1,
             "num_inference_steps": None,
             "timesteps": [22, 0],
             "generator": generator,
             "output_type": "np",
         }

         return inputs



    def 					_lowercase						(self    :    Union[str, Any] ):
         UpperCAmelCase_           =   "cpu"  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase_           =   self.get_dummy_components()
         UpperCAmelCase_           =   ConsistencyModelPipeline(**__a )
         UpperCAmelCase_           =   pipe.to(__a )
         pipe.set_progress_bar_config(disable=__a )

         UpperCAmelCase_           =   self.get_dummy_inputs(__a )
         UpperCAmelCase_           =   pipe(**__a ).images
         assert image.shape == (1, 32, 32, 3)

         UpperCAmelCase_           =   image[0, -3:, -3:, -1]
         UpperCAmelCase_           =   np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-3



    def 					_lowercase						(self    :    Optional[int] ):
         UpperCAmelCase_           =   "cpu"  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase_           =   self.get_dummy_components(class_cond=__a )
         UpperCAmelCase_           =   ConsistencyModelPipeline(**__a )
         UpperCAmelCase_           =   pipe.to(__a )
         pipe.set_progress_bar_config(disable=__a )

         UpperCAmelCase_           =   self.get_dummy_inputs(__a )
         UpperCAmelCase_           =   0
         UpperCAmelCase_           =   pipe(**__a ).images
         assert image.shape == (1, 32, 32, 3)

         UpperCAmelCase_           =   image[0, -3:, -3:, -1]
         UpperCAmelCase_           =   np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-3



    def 					_lowercase						(self    :    str ):
         UpperCAmelCase_           =   "cpu"  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase_           =   self.get_dummy_components()
         UpperCAmelCase_           =   ConsistencyModelPipeline(**__a )
         UpperCAmelCase_           =   pipe.to(__a )
         pipe.set_progress_bar_config(disable=__a )

         UpperCAmelCase_           =   self.get_dummy_inputs(__a )
         UpperCAmelCase_           =   1
         UpperCAmelCase_           =   None
         UpperCAmelCase_           =   pipe(**__a ).images
         assert image.shape == (1, 32, 32, 3)

         UpperCAmelCase_           =   image[0, -3:, -3:, -1]
         UpperCAmelCase_           =   np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-3






    def 					_lowercase						(self    :    Optional[int] ):
         UpperCAmelCase_           =   "cpu"  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase_           =   self.get_dummy_components(class_cond=__a )
         UpperCAmelCase_           =   ConsistencyModelPipeline(**__a )
         UpperCAmelCase_           =   pipe.to(__a )
         pipe.set_progress_bar_config(disable=__a )

         UpperCAmelCase_           =   self.get_dummy_inputs(__a )
         UpperCAmelCase_           =   1
         UpperCAmelCase_           =   None
         UpperCAmelCase_           =   0
         UpperCAmelCase_           =   pipe(**__a ).images
         assert image.shape == (1, 32, 32, 3)

         UpperCAmelCase_           =   image[0, -3:, -3:, -1]
         UpperCAmelCase_           =   np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-3




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



    def 					_lowercase						(self    :    Any ):
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()



    def 					_lowercase						(self    :    Optional[Any]   ,			__a    :    str=0   ,			__a    :    str=False   ,			__a    :    List[str]="cpu"   ,			__a    :    str=torch.floataa   ,			__a    :    int=(1, 3, 64, 64) ):
         UpperCAmelCase_           =   torch.manual_seed(__a )

         UpperCAmelCase_           =   {
             "num_inference_steps": None,
             "timesteps": [22, 0],
             "class_labels": 0,
             "generator": generator,
             "output_type": "np",
         }

         if get_fixed_latents:
              UpperCAmelCase_           =   self.get_fixed_latents(seed=__a   ,			device=__a   ,			dtype=__a   ,			shape=__a )
              UpperCAmelCase_           =   latents

         return inputs



    def 					_lowercase						(self    :    Any   ,			__a    :    List[Any]=0   ,			__a    :    Tuple="cpu"   ,			__a    :    Optional[int]=torch.floataa   ,			__a    :    str=(1, 3, 64, 64) ):
         if type(__a ) == str:
              UpperCAmelCase_           =   torch.device(__a )
         UpperCAmelCase_           =   torch.Generator(device=__a ).manual_seed(__a )
         UpperCAmelCase_           =   randn_tensor(__a   ,			generator=__a   ,			device=__a   ,			dtype=__a )
         return latents



    def 					_lowercase						(self    :    str ):
         UpperCAmelCase_           =   UNetaDModel.from_pretrained("diffusers/consistency_models"   ,			subfolder="diffusers_cd_imagenet64_l2" )
         UpperCAmelCase_           =   CMStochasticIterativeScheduler(
             num_train_timesteps=40   ,			sigma_min=0.0_02   ,			sigma_max=80.0   ,			)
         UpperCAmelCase_           =   ConsistencyModelPipeline(unet=__a   ,			scheduler=__a )
         pipe.to(torch_device=__a )
         pipe.set_progress_bar_config(disable=__a )

         UpperCAmelCase_           =   self.get_inputs()
         UpperCAmelCase_           =   pipe(**__a ).images
         assert image.shape == (1, 64, 64, 3)

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

         UpperCAmelCase_           =   np.array([0.08_88, 0.08_81, 0.06_66, 0.04_79, 0.02_92, 0.01_95, 0.02_01, 0.01_63, 0.02_54] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() < 2E-2



    def 					_lowercase						(self    :    Union[str, Any] ):
         UpperCAmelCase_           =   UNetaDModel.from_pretrained("diffusers/consistency_models"   ,			subfolder="diffusers_cd_imagenet64_l2" )
         UpperCAmelCase_           =   CMStochasticIterativeScheduler(
             num_train_timesteps=40   ,			sigma_min=0.0_02   ,			sigma_max=80.0   ,			)
         UpperCAmelCase_           =   ConsistencyModelPipeline(unet=__a   ,			scheduler=__a )
         pipe.to(torch_device=__a )
         pipe.set_progress_bar_config(disable=__a )

         UpperCAmelCase_           =   self.get_inputs()
         UpperCAmelCase_           =   1
         UpperCAmelCase_           =   None
         UpperCAmelCase_           =   pipe(**__a ).images
         assert image.shape == (1, 64, 64, 3)

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

         UpperCAmelCase_           =   np.array([0.03_40, 0.01_52, 0.00_63, 0.02_67, 0.02_21, 0.01_07, 0.04_16, 0.01_86, 0.02_17] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() < 2E-2



    @require_torch_a
    def 					_lowercase						(self    :    int ):
         UpperCAmelCase_           =   UNetaDModel.from_pretrained("diffusers/consistency_models"   ,			subfolder="diffusers_cd_imagenet64_l2" )
         UpperCAmelCase_           =   CMStochasticIterativeScheduler(
             num_train_timesteps=40   ,			sigma_min=0.0_02   ,			sigma_max=80.0   ,			)
         UpperCAmelCase_           =   ConsistencyModelPipeline(unet=__a   ,			scheduler=__a )
         pipe.to(torch_device=__a   ,			torch_dtype=torch.floataa )
         pipe.set_progress_bar_config(disable=__a )

         UpperCAmelCase_           =   self.get_inputs(get_fixed_latents=__a   ,			device=__a )
         # Ensure usage of flash attention in torch 2.0
         with sdp_kernel(enable_flash=__a   ,			enable_math=__a   ,			enable_mem_efficient=__a ):
              UpperCAmelCase_           =   pipe(**__a ).images
         assert image.shape == (1, 64, 64, 3)

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

         UpperCAmelCase_           =   np.array([0.18_75, 0.14_28, 0.12_89, 0.21_51, 0.20_92, 0.14_77, 0.18_77, 0.16_41, 0.13_53] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-3






    @require_torch_a
    def 					_lowercase						(self    :    Any ):
         UpperCAmelCase_           =   UNetaDModel.from_pretrained("diffusers/consistency_models"   ,			subfolder="diffusers_cd_imagenet64_l2" )
         UpperCAmelCase_           =   CMStochasticIterativeScheduler(
             num_train_timesteps=40   ,			sigma_min=0.0_02   ,			sigma_max=80.0   ,			)
         UpperCAmelCase_           =   ConsistencyModelPipeline(unet=__a   ,			scheduler=__a )
         pipe.to(torch_device=__a   ,			torch_dtype=torch.floataa )
         pipe.set_progress_bar_config(disable=__a )

         UpperCAmelCase_           =   self.get_inputs(get_fixed_latents=__a   ,			device=__a )
         UpperCAmelCase_           =   1
         UpperCAmelCase_           =   None
         # Ensure usage of flash attention in torch 2.0
         with sdp_kernel(enable_flash=__a   ,			enable_math=__a   ,			enable_mem_efficient=__a ):
              UpperCAmelCase_           =   pipe(**__a ).images
         assert image.shape == (1, 64, 64, 3)

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

         UpperCAmelCase_           =   np.array([0.16_63, 0.19_48, 0.22_75, 0.16_80, 0.12_04, 0.12_45, 0.18_58, 0.13_38, 0.20_95] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-3





'''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 )




