

'''simple docstring'''




import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin






@skip_mps
class     _SCREAMING_SNAKE_CASE       (       lowerCAmelCase_						,lowerCAmelCase_						,unittest.TestCase      ):
 __SCREAMING_SNAKE_CASE	:int     			=	IFPipeline
 __SCREAMING_SNAKE_CASE	:Dict     			=	TEXT_TO_IMAGE_PARAMS - {"""width""", """height""", """latents"""}
 __SCREAMING_SNAKE_CASE	:Union[str, Any]     			=	TEXT_TO_IMAGE_BATCH_PARAMS
 __SCREAMING_SNAKE_CASE	:int     			=	PipelineTesterMixin.required_optional_params - {"""latents"""}



 def 					snake_case__    (       self						:       str     ):
      return self._get_dummy_components()



 def 					snake_case__    (       self						:       Union[str, Any]					, a__						:       int					, a__						:       int=0     ):
      if str(__lowerCamelCase     ).startswith('''mps'''     ):
           __magic_name__		  =							torch.manual_seed(__lowerCamelCase     )
      else:
           __magic_name__		  =							torch.Generator(device=__lowerCamelCase     ).manual_seed(__lowerCamelCase     )

      __magic_name__		  =							{
          'prompt': 'A painting of a squirrel eating a burger',
          'generator': generator,
          'num_inference_steps': 2,
          'output_type': 'numpy',
      }

      return inputs



 def 					snake_case__    (       self						:       Dict     ):
      self._test_save_load_optional_components()



 @unittest.skipIf(torch_device != '''cuda'''					, reason='''float16 requires CUDA'''     )
 def 					snake_case__    (       self						:       Union[str, Any]     ):
      # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
      super().test_save_load_floataa(expected_max_diff=1E-1     )



 def 					snake_case__    (       self						:       Optional[Any]     ):
      self._test_attention_slicing_forward_pass(expected_max_diff=1E-2     )



 def 					snake_case__    (       self						:       Any     ):
      self._test_save_load_local()



 def 					snake_case__    (       self						:       Optional[Any]     ):
      self._test_inference_batch_single_identical(
          expected_max_diff=1E-2					, )



 @unittest.skipIf(
     torch_device != '''cuda''' or not is_xformers_available()					, reason='''XFormers attention is only available with CUDA and `xformers` installed'''					, )
 def 					snake_case__    (       self						:       Any     ):
      self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3     )






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



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



 def 					snake_case__    (       self						:       Tuple     ):
      # if

      __magic_name__		  =							IFPipeline.from_pretrained('''DeepFloyd/IF-I-XL-v1.0'''					, variant='''fp16'''					, torch_dtype=torch.floataa     )

      __magic_name__		  =							IFSuperResolutionPipeline.from_pretrained(
          '''DeepFloyd/IF-II-L-v1.0'''					, variant='''fp16'''					, torch_dtype=torch.floataa					, text_encoder=__lowerCamelCase					, tokenizer=__lowerCamelCase     )

      # pre compute text embeddings and remove T5 to save memory

      pipe_a.text_encoder.to('''cuda'''     )

      __magic_name__		  =							pipe_a.encode_prompt('''anime turtle'''					, device='''cuda'''     )

      del pipe_a.tokenizer
      del pipe_a.text_encoder
      gc.collect()

      __magic_name__		  =							None
      __magic_name__		  =							None

      pipe_a.enable_model_cpu_offload()
      pipe_a.enable_model_cpu_offload()

      pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()     )
      pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()     )

      self._test_if(__lowerCamelCase					, __lowerCamelCase					, __lowerCamelCase					, __lowerCamelCase     )

      pipe_a.remove_all_hooks()
      pipe_a.remove_all_hooks()

      # img2img

      __magic_name__		  =							IFImgaImgPipeline(**pipe_a.components     )
      __magic_name__		  =							IFImgaImgSuperResolutionPipeline(**pipe_a.components     )

      pipe_a.enable_model_cpu_offload()
      pipe_a.enable_model_cpu_offload()

      pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()     )
      pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()     )

      self._test_if_imgaimg(__lowerCamelCase					, __lowerCamelCase					, __lowerCamelCase					, __lowerCamelCase     )

      pipe_a.remove_all_hooks()
      pipe_a.remove_all_hooks()

      # inpainting

      __magic_name__		  =							IFInpaintingPipeline(**pipe_a.components     )
      __magic_name__		  =							IFInpaintingSuperResolutionPipeline(**pipe_a.components     )

      pipe_a.enable_model_cpu_offload()
      pipe_a.enable_model_cpu_offload()

      pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()     )
      pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()     )

      self._test_if_inpainting(__lowerCamelCase					, __lowerCamelCase					, __lowerCamelCase					, __lowerCamelCase     )



 def 					snake_case__    (       self						:       Union[str, Any]					, a__						:       Optional[int]					, a__						:       List[Any]					, a__						:       Union[str, Any]					, a__						:       Dict     ):
      # pipeline 1

      _start_torch_memory_measurement()

      __magic_name__		  =							torch.Generator(device='''cpu'''     ).manual_seed(0     )
      __magic_name__		  =							pipe_a(
          prompt_embeds=__lowerCamelCase					, negative_prompt_embeds=__lowerCamelCase					, num_inference_steps=2					, generator=__lowerCamelCase					, output_type='''np'''					, )

      __magic_name__		  =							output.images[0]

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

      __magic_name__		  =							torch.cuda.max_memory_allocated()
      assert mem_bytes < 13 * 10**9

      __magic_name__		  =							load_numpy(
          '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy'''     )
      assert_mean_pixel_difference(__lowerCamelCase					, __lowerCamelCase     )

      # pipeline 2

      _start_torch_memory_measurement()

      __magic_name__		  =							torch.Generator(device='''cpu'''     ).manual_seed(0     )

      __magic_name__		  =							floats_tensor((1, 3, 64, 64)					, rng=random.Random(0     )     ).to(__lowerCamelCase     )

      __magic_name__		  =							pipe_a(
          prompt_embeds=__lowerCamelCase					, negative_prompt_embeds=__lowerCamelCase					, image=__lowerCamelCase					, generator=__lowerCamelCase					, num_inference_steps=2					, output_type='''np'''					, )

      __magic_name__		  =							output.images[0]

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

      __magic_name__		  =							torch.cuda.max_memory_allocated()
      assert mem_bytes < 4 * 10**9

      __magic_name__		  =							load_numpy(
          '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy'''     )
      assert_mean_pixel_difference(__lowerCamelCase					, __lowerCamelCase     )



 def 					snake_case__    (       self						:       List[str]					, a__						:       Dict					, a__						:       List[str]					, a__						:       Tuple					, a__						:       Any     ):
      # pipeline 1

      _start_torch_memory_measurement()

      __magic_name__		  =							floats_tensor((1, 3, 64, 64)					, rng=random.Random(0     )     ).to(__lowerCamelCase     )

      __magic_name__		  =							torch.Generator(device='''cpu'''     ).manual_seed(0     )

      __magic_name__		  =							pipe_a(
          prompt_embeds=__lowerCamelCase					, negative_prompt_embeds=__lowerCamelCase					, image=__lowerCamelCase					, num_inference_steps=2					, generator=__lowerCamelCase					, output_type='''np'''					, )

      __magic_name__		  =							output.images[0]

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

      __magic_name__		  =							torch.cuda.max_memory_allocated()
      assert mem_bytes < 10 * 10**9

      __magic_name__		  =							load_numpy(
          '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy'''     )
      assert_mean_pixel_difference(__lowerCamelCase					, __lowerCamelCase     )

      # pipeline 2

      _start_torch_memory_measurement()

      __magic_name__		  =							torch.Generator(device='''cpu'''     ).manual_seed(0     )

      __magic_name__		  =							floats_tensor((1, 3, 256, 256)					, rng=random.Random(0     )     ).to(__lowerCamelCase     )
      __magic_name__		  =							floats_tensor((1, 3, 64, 64)					, rng=random.Random(0     )     ).to(__lowerCamelCase     )

      __magic_name__		  =							pipe_a(
          prompt_embeds=__lowerCamelCase					, negative_prompt_embeds=__lowerCamelCase					, image=__lowerCamelCase					, original_image=__lowerCamelCase					, generator=__lowerCamelCase					, num_inference_steps=2					, output_type='''np'''					, )

      __magic_name__		  =							output.images[0]

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

      __magic_name__		  =							torch.cuda.max_memory_allocated()
      assert mem_bytes < 4 * 10**9

      __magic_name__		  =							load_numpy(
          '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy'''     )
      assert_mean_pixel_difference(__lowerCamelCase					, __lowerCamelCase     )



 def 					snake_case__    (       self						:       List[str]					, a__						:       Optional[Any]					, a__						:       Optional[Any]					, a__						:       Dict					, a__						:       List[Any]     ):
      # pipeline 1

      _start_torch_memory_measurement()

      __magic_name__		  =							floats_tensor((1, 3, 64, 64)					, rng=random.Random(0     )     ).to(__lowerCamelCase     )
      __magic_name__		  =							floats_tensor((1, 3, 64, 64)					, rng=random.Random(1     )     ).to(__lowerCamelCase     )

      __magic_name__		  =							torch.Generator(device='''cpu'''     ).manual_seed(0     )
      __magic_name__		  =							pipe_a(
          prompt_embeds=__lowerCamelCase					, negative_prompt_embeds=__lowerCamelCase					, image=__lowerCamelCase					, mask_image=__lowerCamelCase					, num_inference_steps=2					, generator=__lowerCamelCase					, output_type='''np'''					, )

      __magic_name__		  =							output.images[0]

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

      __magic_name__		  =							torch.cuda.max_memory_allocated()
      assert mem_bytes < 10 * 10**9

      __magic_name__		  =							load_numpy(
          '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy'''     )
      assert_mean_pixel_difference(__lowerCamelCase					, __lowerCamelCase     )

      # pipeline 2

      _start_torch_memory_measurement()

      __magic_name__		  =							torch.Generator(device='''cpu'''     ).manual_seed(0     )

      __magic_name__		  =							floats_tensor((1, 3, 64, 64)					, rng=random.Random(0     )     ).to(__lowerCamelCase     )
      __magic_name__		  =							floats_tensor((1, 3, 256, 256)					, rng=random.Random(0     )     ).to(__lowerCamelCase     )
      __magic_name__		  =							floats_tensor((1, 3, 256, 256)					, rng=random.Random(1     )     ).to(__lowerCamelCase     )

      __magic_name__		  =							pipe_a(
          prompt_embeds=__lowerCamelCase					, negative_prompt_embeds=__lowerCamelCase					, image=__lowerCamelCase					, mask_image=__lowerCamelCase					, original_image=__lowerCamelCase					, generator=__lowerCamelCase					, num_inference_steps=2					, output_type='''np'''					, )

      __magic_name__		  =							output.images[0]

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

      __magic_name__		  =							torch.cuda.max_memory_allocated()
      assert mem_bytes < 4 * 10**9

      __magic_name__		  =							load_numpy(
          '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy'''     )
      assert_mean_pixel_difference(__lowerCamelCase					, __lowerCamelCase     )




def 		UpperCamelCase			(					) ->      Tuple:






     '''simple docstring'''
     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()












'''simple docstring'''






import logging
import os
from typing import List, Tuple

import numpy as np
import psutil
import torch
import torch.distributed as dist

from transformers import RagRetriever


A_			      =		logging.getLogger(__name__)






class        lowercase_      (						lowerCAmelCase_							):





     def __init__(     self     :					List[Any]     ,							__lowerCamelCase     :					Tuple     ,							__lowerCamelCase     :					str     ,							__lowerCamelCase     :					Tuple     ,							__lowerCamelCase     :					int=None			):
           super().__init__(
               __lowerCamelCase     ,							question_encoder_tokenizer=__lowerCamelCase     ,							generator_tokenizer=__lowerCamelCase     ,							index=__lowerCamelCase     ,							init_retrieval=__lowerCamelCase     ,							)
           snake_case__       :   Optional[int]	      =    None





     def 			_lowerCAmelCase      (     self     :					List[str]     ,							__lowerCamelCase     :					int			):
           logger.info('initializing retrieval'			)

           # initializing a separate process group for retrieval as the default
           # nccl backend doesn't support gather/scatter operations while gloo
           # is too slow to replace nccl for the core gpu communication
           if dist.is_initialized():
                 logger.info('dist initialized'			)
                 # needs to be set manually
                 snake_case__       :   Any	      =    self._infer_socket_ifname()
                 # avoid clash with the NCCL port
                 snake_case__       :   int	      =    str(distributed_port + 1			)
                 snake_case__       :   Any	      =    dist.new_group(ranks=__lowerCamelCase     ,							backend='gloo'			)

           # initialize retriever only on the main worker
           if not dist.is_initialized() or self._is_main():
                 logger.info('dist not initialized / main'			)
                 self.index.init_index()

           # all processes wait untill the retriever is initialized by the main process
           if dist.is_initialized():
                 torch.distributed.barrier(group=self.process_group			)





     def 			_lowerCAmelCase      (     self     :					List[Any]			):
           return dist.get_rank(group=self.process_group			) == 0





     def 			_lowerCAmelCase      (     self     :					Any     ,							__lowerCamelCase     :					Dict     ,							__lowerCamelCase     :					str     ,							__lowerCamelCase     :					Union[str, Any]=torch.floataa			):
           snake_case__       :   Optional[int]	      =    torch.empty(__lowerCamelCase     ,							dtype=__lowerCamelCase			)
           dist.scatter(__lowerCamelCase     ,							src=0     ,							scatter_list=__lowerCamelCase     ,							group=self.process_group			)
           return target_tensor





     def 			_lowerCAmelCase      (     self     :					Dict			):
           snake_case__       :   Union[str, Any]	      =    psutil.net_if_addrs()
           # a hacky way to deal with varying network interface names
           snake_case__       :   Any	      =    next((addr for addr in addrs if addr.startswith('e'			))     ,							__lowerCamelCase			)
           return ifname





     def 			_lowerCAmelCase      (     self     :					Tuple     ,							__lowerCamelCase     :					np.ndarray     ,							__lowerCamelCase     :					int			):

           # single GPU training
           if not dist.is_initialized():
                 snake_case__      ,							snake_case__       :   Optional[int]	      =    self._main_retrieve(__lowerCamelCase     ,							__lowerCamelCase			)
                 return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(__lowerCamelCase			)

           # distributed training
           snake_case__       :   Optional[Any]	      =    dist.get_world_size(group=self.process_group			)

           # gather logic
           snake_case__       :   Optional[int]	      =    None
           if self._is_main():
                 snake_case__       :   Optional[int]	      =    [torch.empty(question_hidden_states.shape     ,							dtype=torch.floataa			) for _ in range(__lowerCamelCase			)]
           dist.gather(torch.tensor(__lowerCamelCase			)     ,							dst=0     ,							gather_list=__lowerCamelCase     ,							group=self.process_group			)

           # scatter logic
           snake_case__       :   Union[str, Any]	      =    question_hidden_states.shape[0]
           snake_case__       :   Tuple	      =    []
           snake_case__       :   List[Any]	      =    []
           if self._is_main():
                 assert len(__lowerCamelCase			) == world_size
                 snake_case__      ,							snake_case__       :   Optional[int]	      =    self._main_retrieve(torch.cat(__lowerCamelCase			).numpy()     ,							__lowerCamelCase			)
                 snake_case__      ,							snake_case__       :   Optional[Any]	      =    torch.tensor(__lowerCamelCase			), torch.tensor(__lowerCamelCase			)
                 snake_case__       :   str	      =    self._chunk_tensor(__lowerCamelCase     ,							__lowerCamelCase			)
                 snake_case__       :   Dict	      =    self._chunk_tensor(__lowerCamelCase     ,							__lowerCamelCase			)
           snake_case__       :   Tuple	      =    self._scattered(__lowerCamelCase     ,							[n_queries, n_docs]     ,							target_type=torch.intaa			)
           snake_case__       :   Dict	      =    self._scattered(__lowerCamelCase     ,							[n_queries, n_docs, question_hidden_states.shape[1]]			)

           return retrieved_doc_embeds.numpy(), doc_ids.numpy(), self.index.get_doc_dicts(__lowerCamelCase			)
