





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  __A(						__UpperCAmelCase						,			__UpperCAmelCase						,			unittest.TestCase ):
     __A	  =       IFPipeline
     __A	  =       TEXT_TO_IMAGE_PARAMS - {"width", "height", "latents"}
     __A	  =       TEXT_TO_IMAGE_BATCH_PARAMS
     __A	  =       PipelineTesterMixin.required_optional_params - {"latents"}

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            return self._get_dummy_components()

     def 	_UpperCamelCase     (      self,       A,       A=0      ):





            """simple docstring"""

            if str(A      ).startswith('''mps'''      ):
                   _UpperCamelCase              =						torch.manual_seed(A      )
            else:
                   _UpperCamelCase              =						torch.Generator(device=A      ).manual_seed(A      )

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

            return inputs

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            self._test_save_load_optional_components()

     @unittest.skipIf(torch_device != '''cuda''',       reason='''float16 requires CUDA'''      )
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            super().test_save_load_floataa(expected_max_diff=1E-1      )

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            self._test_attention_slicing_forward_pass(expected_max_diff=1E-2      )

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            self._test_save_load_local()

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

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





            """simple docstring"""

            self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3      )






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

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

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

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

            # pre compute text embeddings and remove T5 to save memory

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

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

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

            _UpperCamelCase              =						None
            _UpperCamelCase              =						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(A,       A,       A,       A      )

            pipe_a.remove_all_hooks()
            pipe_a.remove_all_hooks()

            # img2img

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

            pipe_a.remove_all_hooks()
            pipe_a.remove_all_hooks()

            # inpainting

            _UpperCamelCase              =						IFInpaintingPipeline(**pipe_a.components      )
            _UpperCamelCase              =						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(A,       A,       A,       A      )

     def 	_UpperCamelCase     (      self,       A,       A,       A,       A      ):





            """simple docstring"""

            _start_torch_memory_measurement()

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

            _UpperCamelCase              =						output.images[0]

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

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

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

            # pipeline 2

            _start_torch_memory_measurement()

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

            _UpperCamelCase              =						floats_tensor((1, 3, 64, 64),       rng=random.Random(0      )      ).to(A      )

            _UpperCamelCase              =						pipe_a(
                prompt_embeds=A,       negative_prompt_embeds=A,       image=A,       generator=A,       num_inference_steps=2,       output_type='''np''',       )

            _UpperCamelCase              =						output.images[0]

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

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

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

     def 	_UpperCamelCase     (      self,       A,       A,       A,       A      ):





            """simple docstring"""

            _start_torch_memory_measurement()

            _UpperCamelCase              =						floats_tensor((1, 3, 64, 64),       rng=random.Random(0      )      ).to(A      )

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

            _UpperCamelCase              =						pipe_a(
                prompt_embeds=A,       negative_prompt_embeds=A,       image=A,       num_inference_steps=2,       generator=A,       output_type='''np''',       )

            _UpperCamelCase              =						output.images[0]

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

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

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

            # pipeline 2

            _start_torch_memory_measurement()

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

            _UpperCamelCase              =						floats_tensor((1, 3, 256, 256),       rng=random.Random(0      )      ).to(A      )
            _UpperCamelCase              =						floats_tensor((1, 3, 64, 64),       rng=random.Random(0      )      ).to(A      )

            _UpperCamelCase              =						pipe_a(
                prompt_embeds=A,       negative_prompt_embeds=A,       image=A,       original_image=A,       generator=A,       num_inference_steps=2,       output_type='''np''',       )

            _UpperCamelCase              =						output.images[0]

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

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

            _UpperCamelCase              =						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(A,       A      )



     def 	_UpperCamelCase     (      self,       A,       A,       A,       A      ):





            """simple docstring"""

            _start_torch_memory_measurement()

            _UpperCamelCase              =						floats_tensor((1, 3, 64, 64),       rng=random.Random(0      )      ).to(A      )
            _UpperCamelCase              =						floats_tensor((1, 3, 64, 64),       rng=random.Random(1      )      ).to(A      )

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

            _UpperCamelCase              =						output.images[0]

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

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

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

            # pipeline 2

            _start_torch_memory_measurement()

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

            _UpperCamelCase              =						floats_tensor((1, 3, 64, 64),       rng=random.Random(0      )      ).to(A      )
            _UpperCamelCase              =						floats_tensor((1, 3, 256, 256),       rng=random.Random(0      )      ).to(A      )
            _UpperCamelCase              =						floats_tensor((1, 3, 256, 256),       rng=random.Random(1      )      ).to(A      )

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

            _UpperCamelCase              =						output.images[0]

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

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

            _UpperCamelCase              =						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(A,       A      )






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













import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
       import jax
       import jax.numpy as jnp
       from flax.jax_utils import replicate
       from flax.training.common_utils import shard

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

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            super().tearDown()
            gc.collect()

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase       ,		_UpperCamelCase              =						FlaxStableDiffusionPipeline.from_pretrained(
                '''stabilityai/stable-diffusion-2''',       revision='''bf16''',       dtype=jnp.bfloataa,       )

            _UpperCamelCase              =						'''A painting of a squirrel eating a burger'''
            _UpperCamelCase              =						jax.device_count()
            _UpperCamelCase              =						num_samples * [prompt]
            _UpperCamelCase              =						sd_pipe.prepare_inputs(A      )

            _UpperCamelCase              =						replicate(A      )
            _UpperCamelCase              =						shard(A      )

            _UpperCamelCase              =						jax.random.PRNGKey(0      )
            _UpperCamelCase              =						jax.random.split(A,       jax.device_count()      )

            _UpperCamelCase              =						sd_pipe(A,       A,       A,       num_inference_steps=25,       jit=A      )[0]
            assert images.shape == (jax.device_count(), 1, 768, 768, 3)

            _UpperCamelCase              =						images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]      )
            _UpperCamelCase              =						images[0, 253:256, 253:256, -1]

            _UpperCamelCase              =						jnp.asarray(jax.device_get(image_slice.flatten()      )      )
            _UpperCamelCase              =						jnp.array([0.4_238, 0.4_414, 0.4_395, 0.4_453, 0.4_629, 0.4_590, 0.4_531, 0.45_508, 0.4_512]      )
            print(F'''output_slice: {output_slice}'''      )
            assert jnp.abs(output_slice - expected_slice      ).max() < 1E-2



     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						'''stabilityai/stable-diffusion-2'''
            _UpperCamelCase       ,		_UpperCamelCase              =						FlaxDPMSolverMultistepScheduler.from_pretrained(A,       subfolder='''scheduler'''      )
            _UpperCamelCase       ,		_UpperCamelCase              =						FlaxStableDiffusionPipeline.from_pretrained(
                A,       scheduler=A,       revision='''bf16''',       dtype=jnp.bfloataa,       )
            _UpperCamelCase              =						scheduler_params

            _UpperCamelCase              =						'''A painting of a squirrel eating a burger'''
            _UpperCamelCase              =						jax.device_count()
            _UpperCamelCase              =						num_samples * [prompt]
            _UpperCamelCase              =						sd_pipe.prepare_inputs(A      )

            _UpperCamelCase              =						replicate(A      )
            _UpperCamelCase              =						shard(A      )

            _UpperCamelCase              =						jax.random.PRNGKey(0      )
            _UpperCamelCase              =						jax.random.split(A,       jax.device_count()      )

            _UpperCamelCase              =						sd_pipe(A,       A,       A,       num_inference_steps=25,       jit=A      )[0]
            assert images.shape == (jax.device_count(), 1, 768, 768, 3)

            _UpperCamelCase              =						images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]      )
            _UpperCamelCase              =						images[0, 253:256, 253:256, -1]

            _UpperCamelCase              =						jnp.asarray(jax.device_get(image_slice.flatten()      )      )
            _UpperCamelCase              =						jnp.array([0.4_336, 0.42_969, 0.4_453, 0.4_199, 0.4_297, 0.4_531, 0.4_434, 0.4_434, 0.4_297]      )
            print(F'''output_slice: {output_slice}'''      )
            assert jnp.abs(output_slice - expected_slice      ).max() < 1E-2






