'''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 lowercase_       (		_UpperCamelCase    ,    _UpperCamelCase    ,    unittest.TestCase			):







   """simple docstring"""

   __lowerCAmelCase         =   IFPipeline
   __lowerCAmelCase         =   TEXT_TO_IMAGE_PARAMS - {"width", "height", "latents"}
   __lowerCAmelCase         =   TEXT_TO_IMAGE_BATCH_PARAMS
   __lowerCAmelCase         =   PipelineTesterMixin.required_optional_params - {"latents"}
   def 	__UpperCAmelCase   (      self      :			int     )  ->			Tuple:
        return self._get_dummy_components()
   def 	__UpperCAmelCase   (      self      :			Dict, UpperCamelCase__      :			str, UpperCamelCase__      :			Any=0     )  ->			List[str]:
        if str(UpperCamelCase__     ).startswith('mps'     ):
             _A    		=  torch.manual_seed(UpperCamelCase__     )
        else:
             _A    		=  torch.Generator(device=UpperCamelCase__     ).manual_seed(UpperCamelCase__     )

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

        return inputs
   def 	__UpperCAmelCase   (      self      :			List[Any]     )  ->			int:
        self._test_save_load_optional_components()
   @unittest.skipIf(torch_device != 'cuda', reason='float16 requires CUDA'     )
   def 	__UpperCAmelCase   (      self      :			List[str]     )  ->			Optional[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 	__UpperCAmelCase   (      self      :			Optional[int]     )  ->			List[str]:
        self._test_attention_slicing_forward_pass(expected_max_diff=1e-2     )
   def 	__UpperCAmelCase   (      self      :			str     )  ->			str:
        self._test_save_load_local()
   def 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			Optional[int]:
        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 	__UpperCAmelCase   (      self      :			Optional[int]     )  ->			List[str]:
        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3     )







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







   """simple docstring"""

   def 	__UpperCAmelCase   (      self      :			List[str]     )  ->			Tuple:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()
   def 	__UpperCAmelCase   (      self      :			Dict     )  ->			Dict:
        # if

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

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

        # pre compute text embeddings and remove T5 to save memory

        pipe_a.text_encoder.to('cuda'     )

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

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

        _A    		=  None
        _A    		=  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(UpperCamelCase__, UpperCamelCase__, UpperCamelCase__, UpperCamelCase__     )

        pipe_a.remove_all_hooks()
        pipe_a.remove_all_hooks()

        # img2img

        _A    		=  IFImgaImgPipeline(**pipe_a.components     )
        _A    		=  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(UpperCamelCase__, UpperCamelCase__, UpperCamelCase__, UpperCamelCase__     )

        pipe_a.remove_all_hooks()
        pipe_a.remove_all_hooks()

        # inpainting

        _A    		=  IFInpaintingPipeline(**pipe_a.components     )
        _A    		=  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(UpperCamelCase__, UpperCamelCase__, UpperCamelCase__, UpperCamelCase__     )
   def 	__UpperCAmelCase   (      self      :			Tuple, UpperCamelCase__      :			Dict, UpperCamelCase__      :			int, UpperCamelCase__      :			Dict, UpperCamelCase__      :			Dict     )  ->			List[Any]:
        # pipeline 1

        _start_torch_memory_measurement()

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

        _A    		=  output.images[0]

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

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

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

        # pipeline 2

        _start_torch_memory_measurement()

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

        _A    		=  floats_tensor((1, 3, 64, 64), rng=random.Random(0     )     ).to(UpperCamelCase__     )

        _A    		=  pipe_a(
            prompt_embeds=UpperCamelCase__, negative_prompt_embeds=UpperCamelCase__, image=UpperCamelCase__, generator=UpperCamelCase__, num_inference_steps=2, output_type='np', )

        _A    		=  output.images[0]

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

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

        _A    		=  load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy'     )
        assert_mean_pixel_difference(UpperCamelCase__, UpperCamelCase__     )
   def 	__UpperCAmelCase   (      self      :			List[Any], UpperCamelCase__      :			List[Any], UpperCamelCase__      :			List[str], UpperCamelCase__      :			Optional[Any], UpperCamelCase__      :			int     )  ->			str:
        # pipeline 1

        _start_torch_memory_measurement()

        _A    		=  floats_tensor((1, 3, 64, 64), rng=random.Random(0     )     ).to(UpperCamelCase__     )

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

        _A    		=  pipe_a(
            prompt_embeds=UpperCamelCase__, negative_prompt_embeds=UpperCamelCase__, image=UpperCamelCase__, num_inference_steps=2, generator=UpperCamelCase__, output_type='np', )

        _A    		=  output.images[0]

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

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

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

        # pipeline 2

        _start_torch_memory_measurement()

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

        _A    		=  floats_tensor((1, 3, 2_56, 2_56), rng=random.Random(0     )     ).to(UpperCamelCase__     )
        _A    		=  floats_tensor((1, 3, 64, 64), rng=random.Random(0     )     ).to(UpperCamelCase__     )

        _A    		=  pipe_a(
            prompt_embeds=UpperCamelCase__, negative_prompt_embeds=UpperCamelCase__, image=UpperCamelCase__, original_image=UpperCamelCase__, generator=UpperCamelCase__, num_inference_steps=2, output_type='np', )

        _A    		=  output.images[0]

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

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

        _A    		=  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(UpperCamelCase__, UpperCamelCase__     )

   def 	__UpperCAmelCase   (      self      :			Optional[Any], UpperCamelCase__      :			Union[str, Any], UpperCamelCase__      :			Union[str, Any], UpperCamelCase__      :			Union[str, Any], UpperCamelCase__      :			List[str]     )  ->			Tuple:
        # pipeline 1

        _start_torch_memory_measurement()

        _A    		=  floats_tensor((1, 3, 64, 64), rng=random.Random(0     )     ).to(UpperCamelCase__     )
        _A    		=  floats_tensor((1, 3, 64, 64), rng=random.Random(1     )     ).to(UpperCamelCase__     )

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

        _A    		=  output.images[0]

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

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

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

        # pipeline 2

        _start_torch_memory_measurement()

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

        _A    		=  floats_tensor((1, 3, 64, 64), rng=random.Random(0     )     ).to(UpperCamelCase__     )
        _A    		=  floats_tensor((1, 3, 2_56, 2_56), rng=random.Random(0     )     ).to(UpperCamelCase__     )
        _A    		=  floats_tensor((1, 3, 2_56, 2_56), rng=random.Random(1     )     ).to(UpperCamelCase__     )

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

        _A    		=  output.images[0]

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

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

        _A    		=  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(UpperCamelCase__, UpperCamelCase__     )






def  _SCREAMING_SNAKE_CASE		(      ):
     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()




'''simple docstring'''






import platform
from argparse import ArgumentParser

import huggingface_hub

from .. import __version__ as version
from ..utils import is_accelerate_available, is_torch_available, is_transformers_available, is_xformers_available
from . import BaseDiffusersCLICommand





def  _SCREAMING_SNAKE_CASE		(      __snake_case							:			Union[str, Any]				):
     return EnvironmentCommand()







class lowercase_       (		_UpperCamelCase			):







   """simple docstring"""

   @staticmethod
   def 	__UpperCAmelCase   (      UpperCamelCase__      :			ArgumentParser     )  ->			List[str]:
        _A    		=  parser.add_parser('env'     )
        download_parser.set_defaults(func=UpperCamelCase__     )
   def 	__UpperCAmelCase   (      self      :			Optional[int]     )  ->			Optional[Any]:
        _A    		=  huggingface_hub.__version__

        _A    		=  'not installed'
        _A    		=  'NA'
        if is_torch_available():
             import torch

             _A    		=  torch.__version__
             _A    		=  torch.cuda.is_available()

        _A    		=  'not installed'
        if is_transformers_available():
             import transformers

             _A    		=  transformers.__version__

        _A    		=  'not installed'
        if is_accelerate_available():
             import accelerate

             _A    		=  accelerate.__version__

        _A    		=  'not installed'
        if is_xformers_available():
             import xformers

             _A    		=  xformers.__version__

        _A    		=  {
            '`diffusers` version': version,
            'Platform': platform.platform(),
            'Python version': platform.python_version(),
            'PyTorch version (GPU?)': f'{pt_version} ({pt_cuda_available})',
            'Huggingface_hub version': hub_version,
            'Transformers version': transformers_version,
            'Accelerate version': accelerate_version,
            'xFormers version': xformers_version,
            'Using GPU in script?': '<fill in>',
            'Using distributed or parallel set-up in script?': '<fill in>',
        }

        print('\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n'     )
        print(self.format_dict(UpperCamelCase__     )     )

        return info

   @staticmethod
   def 	__UpperCAmelCase   (      UpperCamelCase__      :			List[Any]     )  ->			Dict:
        return "\n".join([f'- {prop}: {val}' for prop, val in d.items()]     ) + "\n"



