






import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()




class       snake_case__ (				lowercase_						,     lowercase_						,     lowercase_						,     unittest.TestCase):

       '''simple docstring'''
       lowerCamelCase				:					Optional[int]								=			StableUnCLIPImgaImgPipeline
       lowerCamelCase				:					Optional[int]								=			TEXT_GUIDED_IMAGE_VARIATION_PARAMS
       lowerCamelCase				:					Optional[int]								=			TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
       lowerCamelCase				:					Optional[int]								=			frozenset(
           [])  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
       lowerCamelCase				:					Dict								=			frozenset([])



       def        __lowercase    (				self     )			->   List[Any]:






            '''simple docstring'''





            __snake_case      :Any 							=    32
            __snake_case      :List[str] 							=    embedder_hidden_size

            # image encoding components

            __snake_case      :Optional[Any] 							=    CLIPImageProcessor(crop_size=32      ,    size=32     )

            torch.manual_seed(0     )
            __snake_case      :List[str] 							=    CLIPVisionModelWithProjection(
                CLIPVisionConfig(
                    hidden_size=a__      ,    projection_dim=a__      ,    num_hidden_layers=5      ,    num_attention_heads=4      ,    image_size=32      ,    intermediate_size=37      ,    patch_size=1      ,    )     )

            # regular denoising components

            torch.manual_seed(0     )
            __snake_case      :int 							=    StableUnCLIPImageNormalizer(embedding_dim=a__     )
            __snake_case      :Dict 							=    DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""     )

            torch.manual_seed(0     )
            __snake_case      :str 							=    CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

            torch.manual_seed(0     )
            __snake_case      :Tuple 							=    CLIPTextModel(
                CLIPTextConfig(
                    bos_token_id=0      ,    eos_token_id=2      ,    hidden_size=a__      ,    projection_dim=32      ,    intermediate_size=37      ,    layer_norm_eps=1e-05      ,    num_attention_heads=4      ,    num_hidden_layers=5      ,    pad_token_id=1      ,    vocab_size=10_00      ,    )     )

            torch.manual_seed(0     )
            __snake_case      :Optional[Any] 							=    UNetaDConditionModel(
                sample_size=32      ,    in_channels=4      ,    out_channels=4      ,    down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")      ,    up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")      ,    block_out_channels=(32, 64)      ,    attention_head_dim=(2, 4)      ,    class_embed_type="""projection"""      ,    projection_class_embeddings_input_dim=embedder_projection_dim * 2      ,    cross_attention_dim=a__      ,    layers_per_block=1      ,    upcast_attention=a__      ,    use_linear_projection=a__      ,    )

            torch.manual_seed(0     )
            __snake_case      :str 							=    DDIMScheduler(
                beta_schedule="""scaled_linear"""      ,    beta_start=0.0_00_85      ,    beta_end=0.0_12      ,    prediction_type="""v_prediction"""      ,    set_alpha_to_one=a__      ,    steps_offset=1      ,    )

            torch.manual_seed(0     )
            __snake_case      :List[str] 							=    AutoencoderKL()

            __snake_case      :Optional[Any] 							=    {
                # image encoding components
                """feature_extractor""": feature_extractor,
                """image_encoder""": image_encoder.eval(),
                # image noising components
                """image_normalizer""": image_normalizer.eval(),
                """image_noising_scheduler""": image_noising_scheduler,
                # regular denoising components
                """tokenizer""": tokenizer,
                """text_encoder""": text_encoder.eval(),
                """unet""": unet.eval(),
                """scheduler""": scheduler,
                """vae""": vae.eval(),
            }

            return components



       def        __lowercase    (				self      ,    a__      ,    a__=0      ,    a__=True     )			->   Any:






            '''simple docstring'''





            if str(a__     ).startswith("""mps"""     ):
                 __snake_case      :Dict 							=    torch.manual_seed(a__     )
            else:
                 __snake_case      :str 							=    torch.Generator(device=a__     ).manual_seed(a__     )

            __snake_case      :List[str] 							=    floats_tensor((1, 3, 32, 32)      ,    rng=random.Random(a__     )     ).to(a__     )

            if pil_image:
                 __snake_case      :List[Any] 							=    input_image * 0.5 + 0.5
                 __snake_case      :str 							=    input_image.clamp(0      ,    1     )
                 __snake_case      :Tuple 							=    input_image.cpu().permute(0      ,    2      ,    3      ,    1     ).float().numpy()
                 __snake_case      :Union[str, Any] 							=    DiffusionPipeline.numpy_to_pil(a__     )[0]

            return {
                "prompt": "An anime racoon running a marathon",
                "image": input_image,
                "generator": generator,
                "num_inference_steps": 2,
                "output_type": "np",
            }



       @skip_mps
       def        __lowercase    (				self     )			->   List[Any]:






            '''simple docstring'''





            __snake_case      :List[str] 							=    """cpu"""  # ensure determinism for the device-dependent torch.Generator
            __snake_case      :List[str] 							=    self.get_dummy_components()
            __snake_case      :Dict 							=    StableUnCLIPImgaImgPipeline(**a__     )
            __snake_case      :Union[str, Any] 							=    sd_pipe.to(a__     )
            sd_pipe.set_progress_bar_config(disable=a__     )

            __snake_case      :Tuple 							=    self.get_dummy_inputs(a__     )
            inputs.update({"""image_embeds""": None}     )
            __snake_case      :Optional[int] 							=    sd_pipe(**a__     ).images
            __snake_case      :Optional[Any] 							=    image[0, -3:, -3:, -1]

            assert image.shape == (1, 32, 32, 3)
            __snake_case      :Optional[int] 							=    np.array([0.38_72, 0.72_24, 0.56_01, 0.47_41, 0.68_72, 0.58_14, 0.46_36, 0.38_67, 0.50_78]     )

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



       def        __lowercase    (				self     )			->   List[Any]:






            '''simple docstring'''





            __snake_case      :int 							=    torch_device in ["""cpu""", """mps"""]

            self._test_attention_slicing_forward_pass(test_max_difference=a__     )



       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            __snake_case      :List[str] 							=    torch_device in ["""cpu""", """mps"""]

            self._test_inference_batch_single_identical(test_max_difference=a__     )



       @unittest.skipIf(
           torch_device != """cuda""" or not is_xformers_available()      ,    reason="""XFormers attention is only available with CUDA and `xformers` installed"""      ,    )
       def        __lowercase    (				self     )			->   Union[str, Any]:






            '''simple docstring'''





            self._test_xformers_attention_forwardGenerator_pass(test_max_difference=a__     )




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

       '''simple docstring'''



       def        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





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



       def        __lowercase    (				self     )			->   str:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    load_image(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png"""     )

            __snake_case      :Union[str, Any] 							=    load_numpy(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy"""     )

            __snake_case      :Union[str, Any] 							=    StableUnCLIPImgaImgPipeline.from_pretrained(
                """fusing/stable-unclip-2-1-l-img2img"""      ,    torch_dtype=torch.floataa     )
            pipe.to(a__     )
            pipe.set_progress_bar_config(disable=a__     )
            # stable unclip will oom when integration tests are run on a V100,
            # so turn on memory savings
            pipe.enable_attention_slicing()
            pipe.enable_sequential_cpu_offload()

            __snake_case      :Dict 							=    torch.Generator(device="""cpu"""     ).manual_seed(0     )
            __snake_case      :Union[str, Any] 							=    pipe(a__      ,    """anime turle"""      ,    generator=a__      ,    output_type="""np"""     )

            __snake_case      :Any 							=    output.images[0]

            assert image.shape == (7_68, 7_68, 3)

            assert_mean_pixel_difference(a__      ,    a__     )



       def        __lowercase    (				self     )			->   Union[str, Any]:






            '''simple docstring'''





            __snake_case      :Any 							=    load_image(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png"""     )

            __snake_case      :str 							=    load_numpy(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy"""     )

            __snake_case      :str 							=    StableUnCLIPImgaImgPipeline.from_pretrained(
                """fusing/stable-unclip-2-1-h-img2img"""      ,    torch_dtype=torch.floataa     )
            pipe.to(a__     )
            pipe.set_progress_bar_config(disable=a__     )
            # stable unclip will oom when integration tests are run on a V100,
            # so turn on memory savings
            pipe.enable_attention_slicing()
            pipe.enable_sequential_cpu_offload()

            __snake_case      :Optional[Any] 							=    torch.Generator(device="""cpu"""     ).manual_seed(0     )
            __snake_case      :Any 							=    pipe(a__      ,    """anime turle"""      ,    generator=a__      ,    output_type="""np"""     )

            __snake_case      :List[str] 							=    output.images[0]

            assert image.shape == (7_68, 7_68, 3)

            assert_mean_pixel_difference(a__      ,    a__     )



       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            __snake_case      :str 							=    load_image(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png"""     )

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

            __snake_case      :List[Any] 							=    StableUnCLIPImgaImgPipeline.from_pretrained(
                """fusing/stable-unclip-2-1-h-img2img"""      ,    torch_dtype=torch.floataa     )
            __snake_case      :Any 							=    pipe.to(a__     )
            pipe.set_progress_bar_config(disable=a__     )
            pipe.enable_attention_slicing()
            pipe.enable_sequential_cpu_offload()

            __snake_case      :List[Any] 							=    pipe(
                a__      ,    """anime turtle"""      ,    num_inference_steps=2      ,    output_type="""np"""      ,    )

            __snake_case      :Optional[int] 							=    torch.cuda.max_memory_allocated()
            # make sure that less than 7 GB is allocated
            assert mem_bytes < 7 * 10**9










from __future__ import annotations




class       snake_case__ :

       '''simple docstring'''



       def __init__(				self      ,    a__     )			->   None:






            '''simple docstring'''





            __snake_case      :str 							=    order

            # a_{0} ... a_{k}
            __snake_case      :int 							=    [1.0] + [0.0] * order
            # b_{0} ... b_{k}
            __snake_case      :Any 							=    [1.0] + [0.0] * order

            # x[n-1] ... x[n-k]
            __snake_case      :List[Any] 							=    [0.0] * self.order
            # y[n-1] ... y[n-k]
            __snake_case      :int 							=    [0.0] * self.order



       def        __lowercase    (				self      ,    a__      ,    a__     )			->   None:






            '''simple docstring'''





            if len(a__     ) < self.order:
                 __snake_case      :Optional[int] 							=    [1.0, *a_coeffs]

            if len(a__     ) != self.order + 1:
                 __snake_case      :Dict 							=    (
                     F'''Expected a_coeffs to have {self.order + 1} elements '''
                     F'''for {self.order}-order filter, got {len(a__     )}'''
                 )
                 raise ValueError(a__     )

            if len(a__     ) != self.order + 1:
                 __snake_case      :Optional[Any] 							=    (
                     F'''Expected b_coeffs to have {self.order + 1} elements '''
                     F'''for {self.order}-order filter, got {len(a__     )}'''
                 )
                 raise ValueError(a__     )

            __snake_case      :str 							=    a_coeffs
            __snake_case      :str 							=    b_coeffs



       def        __lowercase    (				self      ,    a__     )			->   float:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    0.0

            # Start at index 1 and do index 0 at the end.
            for i in range(1      ,    self.order + 1     ):
                 result += (
                     self.b_coeffs[i] * self.input_history[i - 1]
                     - self.a_coeffs[i] * self.output_history[i - 1]
                 )

            __snake_case      :Dict 							=    (result + self.b_coeffs[0] * sample) / self.a_coeffs[0]

            __snake_case      :Optional[int] 							=    self.input_history[:-1]
            __snake_case      :Optional[Any] 							=    self.output_history[:-1]

            __snake_case      :Union[str, Any] 							=    sample
            __snake_case      :Optional[Any] 							=    result

            return result


