




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 				lowerCAmelCase_       (					lowerCamelCase_						,     lowerCamelCase_						,     unittest.TestCase							):
				__a     :       Optional[int]						=     IFPipeline
				__a     :       str						=     TEXT_TO_IMAGE_PARAMS - {"width", "height", "latents"}
				__a     :       Dict						=     TEXT_TO_IMAGE_BATCH_PARAMS
				__a     :       Optional[Any]						=     PipelineTesterMixin.required_optional_params - {"latents"}






				def 		snake_case	(			self ):
							return self._get_dummy_components()






				def 		snake_case	(			self      ,snake_case__      ,snake_case__=0 ):
							if str(snake_case__ ).startswith('mps' ):
										SCREAMING_SNAKE_CASE_   : int					  =      torch.manual_seed(snake_case__ )
							else:
										SCREAMING_SNAKE_CASE_   : List[Any]					  =      torch.Generator(device=snake_case__ ).manual_seed(snake_case__ )

							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      {
							    'prompt': 'A painting of a squirrel eating a burger',
							    'generator': generator,
							    'num_inference_steps': 2,
							    'output_type': 'numpy',
							}

							return inputs






				def 		snake_case	(			self ):
							self._test_save_load_optional_components()






				@unittest.skipIf(torch_device != 'cuda'      ,reason='float16 requires CUDA' )
				def 		snake_case	(			self ):
							# 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 ):
							self._test_attention_slicing_forward_pass(expected_max_diff=1E-2 )






				def 		snake_case	(			self ):
							self._test_save_load_local()






				def 		snake_case	(			self ):
							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 ):
							self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3 )







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






				def 		snake_case	(			self ):
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()






				def 		snake_case	(			self ):
							# if

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

							SCREAMING_SNAKE_CASE_   : Any					  =      IFSuperResolutionPipeline.from_pretrained(
							    'DeepFloyd/IF-II-L-v1.0'      ,variant='fp16'      ,torch_dtype=torch.floataa      ,text_encoder=snake_case__      ,tokenizer=snake_case__ )

							# pre compute text embeddings and remove T5 to save memory

							pipe_a.text_encoder.to('cuda' )

							SCREAMING_SNAKE_CASE_  ,    SCREAMING_SNAKE_CASE_   : Any					  =      pipe_a.encode_prompt('anime turtle'      ,device='cuda' )

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

							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      None
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      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(snake_case__      ,snake_case__      ,snake_case__      ,snake_case__ )

							pipe_a.remove_all_hooks()
							pipe_a.remove_all_hooks()

							# img2img

							SCREAMING_SNAKE_CASE_   : Tuple					  =      IFImgaImgPipeline(**pipe_a.components )
							SCREAMING_SNAKE_CASE_   : Any					  =      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(snake_case__      ,snake_case__      ,snake_case__      ,snake_case__ )

							pipe_a.remove_all_hooks()
							pipe_a.remove_all_hooks()

							# inpainting

							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      IFInpaintingPipeline(**pipe_a.components )
							SCREAMING_SNAKE_CASE_   : Tuple					  =      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(snake_case__      ,snake_case__      ,snake_case__      ,snake_case__ )






				def 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__ ):
							# pipeline 1

							_start_torch_memory_measurement()

							SCREAMING_SNAKE_CASE_   : int					  =      torch.Generator(device='cpu' ).manual_seed(0 )
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      pipe_a(
							    prompt_embeds=snake_case__      ,negative_prompt_embeds=snake_case__      ,num_inference_steps=2      ,generator=snake_case__      ,output_type='np'      ,)

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      output.images[0]

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

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      torch.cuda.max_memory_allocated()
							assert mem_bytes < 13 * 10**9

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

							# pipeline 2

							_start_torch_memory_measurement()

							SCREAMING_SNAKE_CASE_   : Tuple					  =      torch.Generator(device='cpu' ).manual_seed(0 )

							SCREAMING_SNAKE_CASE_   : List[str]					  =      floats_tensor((1, 3, 64, 64)      ,rng=random.Random(0 ) ).to(snake_case__ )

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      pipe_a(
							    prompt_embeds=snake_case__      ,negative_prompt_embeds=snake_case__      ,image=snake_case__      ,generator=snake_case__      ,num_inference_steps=2      ,output_type='np'      ,)

							SCREAMING_SNAKE_CASE_   : int					  =      output.images[0]

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

							SCREAMING_SNAKE_CASE_   : List[str]					  =      torch.cuda.max_memory_allocated()
							assert mem_bytes < 4 * 10**9

							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      load_numpy(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy' )
							assert_mean_pixel_difference(snake_case__      ,snake_case__ )






				def 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__ ):
							# pipeline 1

							_start_torch_memory_measurement()

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      floats_tensor((1, 3, 64, 64)      ,rng=random.Random(0 ) ).to(snake_case__ )

							SCREAMING_SNAKE_CASE_   : str					  =      torch.Generator(device='cpu' ).manual_seed(0 )

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      pipe_a(
							    prompt_embeds=snake_case__      ,negative_prompt_embeds=snake_case__      ,image=snake_case__      ,num_inference_steps=2      ,generator=snake_case__      ,output_type='np'      ,)

							SCREAMING_SNAKE_CASE_   : str					  =      output.images[0]

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

							SCREAMING_SNAKE_CASE_   : Dict					  =      torch.cuda.max_memory_allocated()
							assert mem_bytes < 10 * 10**9

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

							# pipeline 2

							_start_torch_memory_measurement()

							SCREAMING_SNAKE_CASE_   : str					  =      torch.Generator(device='cpu' ).manual_seed(0 )

							SCREAMING_SNAKE_CASE_   : int					  =      floats_tensor((1, 3, 256, 256)      ,rng=random.Random(0 ) ).to(snake_case__ )
							SCREAMING_SNAKE_CASE_   : int					  =      floats_tensor((1, 3, 64, 64)      ,rng=random.Random(0 ) ).to(snake_case__ )

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      pipe_a(
							    prompt_embeds=snake_case__      ,negative_prompt_embeds=snake_case__      ,image=snake_case__      ,original_image=snake_case__      ,generator=snake_case__      ,num_inference_steps=2      ,output_type='np'      ,)

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      output.images[0]

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

							SCREAMING_SNAKE_CASE_   : List[str]					  =      torch.cuda.max_memory_allocated()
							assert mem_bytes < 4 * 10**9

							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      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(snake_case__      ,snake_case__ )






				def 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__ ):
							# pipeline 1

							_start_torch_memory_measurement()

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      floats_tensor((1, 3, 64, 64)      ,rng=random.Random(0 ) ).to(snake_case__ )
							SCREAMING_SNAKE_CASE_   : int					  =      floats_tensor((1, 3, 64, 64)      ,rng=random.Random(1 ) ).to(snake_case__ )

							SCREAMING_SNAKE_CASE_   : Tuple					  =      torch.Generator(device='cpu' ).manual_seed(0 )
							SCREAMING_SNAKE_CASE_   : Dict					  =      pipe_a(
							    prompt_embeds=snake_case__      ,negative_prompt_embeds=snake_case__      ,image=snake_case__      ,mask_image=snake_case__      ,num_inference_steps=2      ,generator=snake_case__      ,output_type='np'      ,)

							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      output.images[0]

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

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      torch.cuda.max_memory_allocated()
							assert mem_bytes < 10 * 10**9

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

							# pipeline 2

							_start_torch_memory_measurement()

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      torch.Generator(device='cpu' ).manual_seed(0 )

							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      floats_tensor((1, 3, 64, 64)      ,rng=random.Random(0 ) ).to(snake_case__ )
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      floats_tensor((1, 3, 256, 256)      ,rng=random.Random(0 ) ).to(snake_case__ )
							SCREAMING_SNAKE_CASE_   : str					  =      floats_tensor((1, 3, 256, 256)      ,rng=random.Random(1 ) ).to(snake_case__ )

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      pipe_a(
							    prompt_embeds=snake_case__      ,negative_prompt_embeds=snake_case__      ,image=snake_case__      ,mask_image=snake_case__      ,original_image=snake_case__      ,generator=snake_case__      ,num_inference_steps=2      ,output_type='np'      ,)

							SCREAMING_SNAKE_CASE_   : Tuple					  =      output.images[0]

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

							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      torch.cuda.max_memory_allocated()
							assert mem_bytes < 4 * 10**9

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      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(snake_case__      ,snake_case__ )





def        __UpperCAmelCase						(							)       ->  List[str]:







			"""simple docstring"""





			torch.cuda.empty_cache()
			torch.cuda.reset_max_memory_allocated()
			torch.cuda.reset_peak_memory_stats()






"""simple docstring"""
# Copyright 2022 The HuggingFace Team and The OpenBMB Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


_A							=  {
    """configuration_cpmant""": ["""CPMANT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """CpmAntConfig"""],
    """tokenization_cpmant""": ["""CpmAntTokenizer"""],
}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _A							=  [
        """CPMANT_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """CpmAntForCausalLM""",
        """CpmAntModel""",
        """CpmAntPreTrainedModel""",
    ]


if TYPE_CHECKING:
    from .configuration_cpmant import CPMANT_PRETRAINED_CONFIG_ARCHIVE_MAP, CpmAntConfig
    from .tokenization_cpmant import CpmAntTokenizer

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_cpmant import (
            CPMANT_PRETRAINED_MODEL_ARCHIVE_LIST,
            CpmAntForCausalLM,
            CpmAntModel,
            CpmAntPreTrainedModel,
        )


else:
    import sys

    _A							=  _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

