






import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, XLMRobertaTokenizer

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




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

       '''simple docstring'''
       lowerCamelCase				:					Optional[Any]								=			AltDiffusionPipeline
       lowerCamelCase				:					str								=			TEXT_TO_IMAGE_PARAMS
       lowerCamelCase				:					str								=			TEXT_TO_IMAGE_BATCH_PARAMS
       lowerCamelCase				:					Optional[int]								=			TEXT_TO_IMAGE_IMAGE_PARAMS
       lowerCamelCase				:					List[str]								=			TEXT_TO_IMAGE_IMAGE_PARAMS



       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            torch.manual_seed(0     )
            __snake_case      :str 							=    UNetaDConditionModel(
                block_out_channels=(32, 64)      ,    layers_per_block=2      ,    sample_size=32      ,    in_channels=4      ,    out_channels=4      ,    down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")      ,    up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")      ,    cross_attention_dim=32      ,    )
            __snake_case      :int 							=    DDIMScheduler(
                beta_start=0.0_00_85      ,    beta_end=0.0_12      ,    beta_schedule="""scaled_linear"""      ,    clip_sample=a__      ,    set_alpha_to_one=a__      ,    )
            torch.manual_seed(0     )
            __snake_case      :Any 							=    AutoencoderKL(
                block_out_channels=[32, 64]      ,    in_channels=3      ,    out_channels=3      ,    down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]      ,    up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]      ,    latent_channels=4      ,    )

            # TODO: address the non-deterministic text encoder (fails for save-load tests)
            # torch.manual_seed(0)
            # text_encoder_config = RobertaSeriesConfig(
            #     hidden_size=32,
            #     project_dim=32,
            #     intermediate_size=37,
            #     layer_norm_eps=1e-05,
            #     num_attention_heads=4,
            #     num_hidden_layers=5,
            #     vocab_size=5002,
            # )
            # text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

            torch.manual_seed(0     )
            __snake_case      :Optional[Any] 							=    CLIPTextConfig(
                bos_token_id=0      ,    eos_token_id=2      ,    hidden_size=32      ,    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=50_02      ,    )
            __snake_case      :Any 							=    CLIPTextModel(a__     )

            __snake_case      :Optional[Any] 							=    XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""     )
            __snake_case      :List[Any] 							=    77

            __snake_case      :Optional[int] 							=    {
                """unet""": unet,
                """scheduler""": scheduler,
                """vae""": vae,
                """text_encoder""": text_encoder,
                """tokenizer""": tokenizer,
                """safety_checker""": None,
                """feature_extractor""": None,
            }
            return components



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






            '''simple docstring'''





            if str(a__     ).startswith("""mps"""     ):
                 __snake_case      :int 							=    torch.manual_seed(a__     )
            else:
                 __snake_case      :Tuple 							=    torch.Generator(device=a__     ).manual_seed(a__     )
            __snake_case      :List[Any] 							=    {
                """prompt""": """A painting of a squirrel eating a burger""",
                """generator""": generator,
                """num_inference_steps""": 2,
                """guidance_scale""": 6.0,
                """output_type""": """numpy""",
            }
            return inputs



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






            '''simple docstring'''





            super().test_attention_slicing_forward_pass(expected_max_diff=3e-3     )



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






            '''simple docstring'''





            super().test_inference_batch_single_identical(expected_max_diff=3e-3     )



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






            '''simple docstring'''





            __snake_case      :List[Any] 							=    """cpu"""  # ensure determinism for the device-dependent torch.Generator

            __snake_case      :List[str] 							=    self.get_dummy_components()
            torch.manual_seed(0     )
            __snake_case      :Union[str, Any] 							=    RobertaSeriesConfig(
                hidden_size=32      ,    project_dim=32      ,    intermediate_size=37      ,    layer_norm_eps=1e-05      ,    num_attention_heads=4      ,    num_hidden_layers=5      ,    vocab_size=50_02      ,    )
            # TODO: remove after fixing the non-deterministic text encoder
            __snake_case      :Tuple 							=    RobertaSeriesModelWithTransformation(a__     )
            __snake_case      :str 							=    text_encoder

            __snake_case      :Any 							=    AltDiffusionPipeline(**a__     )
            __snake_case      :Tuple 							=    alt_pipe.to(a__     )
            alt_pipe.set_progress_bar_config(disable=a__     )

            __snake_case      :Any 							=    self.get_dummy_inputs(a__     )
            __snake_case      :str 							=    """A photo of an astronaut"""
            __snake_case      :Optional[Any] 							=    alt_pipe(**a__     )
            __snake_case      :int 							=    output.images
            __snake_case      :Any 							=    image[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            __snake_case      :Optional[Any] 							=    np.array(
                [0.5_74_81_62, 0.60_44_71_45, 0.48_82_12_17, 0.50_10_06_36, 0.5_43_11_85, 0.45_76_36_83, 0.49_65_76_96, 0.48_13_27_33, 0.47_57_30_93]     )

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



       def        __lowercase    (				self     )			->   int:






            '''simple docstring'''





            __snake_case      :Tuple 							=    """cpu"""  # ensure determinism for the device-dependent torch.Generator

            __snake_case      :Tuple 							=    self.get_dummy_components()
            __snake_case      :Optional[Any] 							=    PNDMScheduler(skip_prk_steps=a__     )
            torch.manual_seed(0     )
            __snake_case      :List[Any] 							=    RobertaSeriesConfig(
                hidden_size=32      ,    project_dim=32      ,    intermediate_size=37      ,    layer_norm_eps=1e-05      ,    num_attention_heads=4      ,    num_hidden_layers=5      ,    vocab_size=50_02      ,    )
            # TODO: remove after fixing the non-deterministic text encoder
            __snake_case      :List[str] 							=    RobertaSeriesModelWithTransformation(a__     )
            __snake_case      :Any 							=    text_encoder
            __snake_case      :Optional[Any] 							=    AltDiffusionPipeline(**a__     )
            __snake_case      :Dict 							=    alt_pipe.to(a__     )
            alt_pipe.set_progress_bar_config(disable=a__     )

            __snake_case      :str 							=    self.get_dummy_inputs(a__     )
            __snake_case      :int 							=    alt_pipe(**a__     )
            __snake_case      :int 							=    output.images
            __snake_case      :int 							=    image[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            __snake_case      :Optional[int] 							=    np.array(
                [0.51_60_50_93, 0.5_70_72_41, 0.47_36_55_07, 0.50_57_88_86, 0.5_63_38_77, 0.4_64_25_03, 0.5_18_20_81, 0.48_76_34_84, 0.49_08_42_37]     )

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




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

       '''simple docstring'''



       def        __lowercase    (				self     )			->   Tuple:






            '''simple docstring'''





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



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






            '''simple docstring'''





            __snake_case      :List[str] 							=    AltDiffusionPipeline.from_pretrained("""BAAI/AltDiffusion"""      ,    safety_checker=a__     )
            __snake_case      :Optional[int] 							=    alt_pipe.to(a__     )
            alt_pipe.set_progress_bar_config(disable=a__     )

            __snake_case      :List[Any] 							=    """A painting of a squirrel eating a burger"""
            __snake_case      :List[str] 							=    torch.manual_seed(0     )
            __snake_case      :Optional[Any] 							=    alt_pipe([prompt]      ,    generator=a__      ,    guidance_scale=6.0      ,    num_inference_steps=20      ,    output_type="""np"""     )

            __snake_case      :List[Any] 							=    output.images

            __snake_case      :Tuple 							=    image[0, -3:, -3:, -1]

            assert image.shape == (1, 5_12, 5_12, 3)
            __snake_case      :int 							=    np.array([0.10_10, 0.08_00, 0.07_94, 0.08_85, 0.08_43, 0.07_62, 0.07_69, 0.07_29, 0.05_86]     )

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



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






            '''simple docstring'''





            __snake_case      :List[str] 							=    DDIMScheduler.from_pretrained("""BAAI/AltDiffusion"""      ,    subfolder="""scheduler"""     )

            __snake_case      :Dict 							=    AltDiffusionPipeline.from_pretrained("""BAAI/AltDiffusion"""      ,    scheduler=a__      ,    safety_checker=a__     )
            __snake_case      :Optional[int] 							=    alt_pipe.to(a__     )
            alt_pipe.set_progress_bar_config(disable=a__     )

            __snake_case      :List[Any] 							=    """A painting of a squirrel eating a burger"""
            __snake_case      :Any 							=    torch.manual_seed(0     )

            __snake_case      :str 							=    alt_pipe([prompt]      ,    generator=a__      ,    num_inference_steps=2      ,    output_type="""numpy"""     )
            __snake_case      :Optional[Any] 							=    output.images

            __snake_case      :Tuple 							=    image[0, -3:, -3:, -1]

            assert image.shape == (1, 5_12, 5_12, 3)
            __snake_case      :Optional[int] 							=    np.array([0.40_19, 0.40_52, 0.38_10, 0.41_19, 0.39_16, 0.39_82, 0.46_51, 0.41_95, 0.53_23]     )

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










# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch

import math
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, randn_tensor
from .scheduling_utils import SchedulerMixin, SchedulerOutput




@dataclass
class       snake_case__ (				lowercase_):

       '''simple docstring'''
       lowerCamelCase				:					torch.FloatTensor
       lowerCamelCase				:					torch.FloatTensor




class       snake_case__ (				lowercase_						,     lowercase_):

       '''simple docstring'''
       lowerCamelCase				:					int								=			1



       @register_to_config
       def __init__(				self      ,    a__ = 20_00      ,    a__ = 0.15      ,    a__ = 0.01      ,    a__ = 13_48.0      ,    a__ = 1e-5      ,    a__ = 1      ,    )			->   Any:






            '''simple docstring'''





            __snake_case      :int 							=    sigma_max

            # setable values
            __snake_case      :Any 							=    None

            self.set_sigmas(a__      ,    a__      ,    a__      ,    a__     )



       def        __lowercase    (				self      ,    a__      ,    a__ = None     )			->   torch.FloatTensor:






            '''simple docstring'''





            return sample



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






            '''simple docstring'''





            __snake_case      :List[Any] 							=    sampling_eps if sampling_eps is not None else self.config.sampling_eps

            __snake_case      :int 							=    torch.linspace(1      ,    a__      ,    a__      ,    device=a__     )



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






            '''simple docstring'''





            __snake_case      :Union[str, Any] 							=    sigma_min if sigma_min is not None else self.config.sigma_min
            __snake_case      :Any 							=    sigma_max if sigma_max is not None else self.config.sigma_max
            __snake_case      :Any 							=    sampling_eps if sampling_eps is not None else self.config.sampling_eps
            if self.timesteps is None:
                 self.set_timesteps(a__      ,    a__     )

            __snake_case      :int 							=    sigma_min * (sigma_max / sigma_min) ** (self.timesteps / sampling_eps)
            __snake_case      :List[str] 							=    torch.exp(torch.linspace(math.log(a__     )      ,    math.log(a__     )      ,    a__     )     )
            __snake_case      :Any 							=    torch.tensor([sigma_min * (sigma_max / sigma_min) ** t for t in self.timesteps]     )



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






            '''simple docstring'''





            return torch.where(
                timesteps == 0      ,    torch.zeros_like(t.to(timesteps.device     )     )      ,    self.discrete_sigmas[timesteps - 1].to(timesteps.device     )      ,    )



       def        __lowercase    (				self      ,    a__      ,    a__      ,    a__      ,    a__ = None      ,    a__ = True      ,    )			->   Union[SdeVeOutput, Tuple]:






            '''simple docstring'''





            if self.timesteps is None:
                 raise ValueError(
                     """`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"""     )

            __snake_case      :Optional[int] 							=    timestep * torch.ones(
                sample.shape[0]      ,    device=sample.device     )  # torch.repeat_interleave(timestep, sample.shape[0])
            __snake_case      :Tuple 							=    (timestep * (len(self.timesteps     ) - 1)).long()

            # mps requires indices to be in the same device, so we use cpu as is the default with cuda
            __snake_case      :str 							=    timesteps.to(self.discrete_sigmas.device     )

            __snake_case      :str 							=    self.discrete_sigmas[timesteps].to(sample.device     )
            __snake_case      :Union[str, Any] 							=    self.get_adjacent_sigma(a__      ,    a__     ).to(sample.device     )
            __snake_case      :Dict 							=    torch.zeros_like(a__     )
            __snake_case      :str 							=    (sigma**2 - adjacent_sigma**2) ** 0.5

            # equation 6 in the paper: the model_output modeled by the network is grad_x log pt(x)
            # also equation 47 shows the analog from SDE models to ancestral sampling methods
            __snake_case      :Optional[Any] 							=    diffusion.flatten()
            while len(diffusion.shape     ) < len(sample.shape     ):
                 __snake_case      :int 							=    diffusion.unsqueeze(-1     )
            __snake_case      :Optional[Any] 							=    drift - diffusion**2 * model_output

            #  equation 6: sample noise for the diffusion term of
            __snake_case      :str 							=    randn_tensor(
                sample.shape      ,    layout=sample.layout      ,    generator=a__      ,    device=sample.device      ,    dtype=sample.dtype     )
            __snake_case      :str 							=    sample - drift  # subtract because `dt` is a small negative timestep
            # TODO is the variable diffusion the correct scaling term for the noise?
            __snake_case      :str 							=    prev_sample_mean + diffusion * noise  # add impact of diffusion field g

            if not return_dict:
                 return (prev_sample, prev_sample_mean)

            return SdeVeOutput(prev_sample=a__      ,    prev_sample_mean=a__     )



       def        __lowercase    (				self      ,    a__      ,    a__      ,    a__ = None      ,    a__ = True      ,    )			->   Union[SchedulerOutput, Tuple]:






            '''simple docstring'''





            if self.timesteps is None:
                 raise ValueError(
                     """`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"""     )

            # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z"
            # sample noise for correction
            __snake_case      :Any 							=    randn_tensor(sample.shape      ,    layout=sample.layout      ,    generator=a__     ).to(sample.device     )

            # compute step size from the model_output, the noise, and the snr
            __snake_case      :List[Any] 							=    torch.norm(model_output.reshape(model_output.shape[0]      ,    -1     )      ,    dim=-1     ).mean()
            __snake_case      :Tuple 							=    torch.norm(noise.reshape(noise.shape[0]      ,    -1     )      ,    dim=-1     ).mean()
            __snake_case      :Dict 							=    (self.config.snr * noise_norm / grad_norm) ** 2 * 2
            __snake_case      :int 							=    step_size * torch.ones(sample.shape[0]     ).to(sample.device     )
            # self.repeat_scalar(step_size, sample.shape[0])

            # compute corrected sample: model_output term and noise term
            __snake_case      :Tuple 							=    step_size.flatten()
            while len(step_size.shape     ) < len(sample.shape     ):
                 __snake_case      :Optional[int] 							=    step_size.unsqueeze(-1     )
            __snake_case      :List[str] 							=    sample + step_size * model_output
            __snake_case      :Tuple 							=    prev_sample_mean + ((step_size * 2) ** 0.5) * noise

            if not return_dict:
                 return (prev_sample,)

            return SchedulerOutput(prev_sample=a__     )



       def        __lowercase    (				self      ,    a__      ,    a__      ,    a__      ,    )			->   torch.FloatTensor:






            '''simple docstring'''





            __snake_case      :Tuple 							=    timesteps.to(original_samples.device     )
            __snake_case      :Optional[Any] 							=    self.discrete_sigmas.to(original_samples.device     )[timesteps]
            __snake_case      :int 							=    (
                noise * sigmas[:, None, None, None]
                if noise is not None
                else torch.randn_like(a__     ) * sigmas[:, None, None, None]
            )
            __snake_case      :str 							=    noise + original_samples
            return noisy_samples



       def __len__(				self     )			->   Optional[int]:






            '''simple docstring'''





            return self.config.num_train_timesteps


