





'''simple docstring'''

import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

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


@skip_mps
class       a_  (       _SCREAMING_SNAKE_CASE   ,					_SCREAMING_SNAKE_CASE   ,					_SCREAMING_SNAKE_CASE   ,					unittest.TestCase ):

    '''simple docstring'''
    UpperCamelCase									=				StableDiffusionModelEditingPipeline
    UpperCamelCase									=				TEXT_TO_IMAGE_PARAMS
    UpperCamelCase									=				TEXT_TO_IMAGE_BATCH_PARAMS
    UpperCamelCase									=				TEXT_TO_IMAGE_IMAGE_PARAMS
    UpperCamelCase									=				TEXT_TO_IMAGE_IMAGE_PARAMS


    def     snake_case_(	self )	->     int:
        torch.manual_seed(0 )
        _SCREAMING_SNAKE_CASE             =    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				,	)
        _SCREAMING_SNAKE_CASE             =    DDIMScheduler()
        torch.manual_seed(0 )
        _SCREAMING_SNAKE_CASE             =    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				,	)
        torch.manual_seed(0 )
        _SCREAMING_SNAKE_CASE             =    CLIPTextConfig(
            bos_token_id=0				,	eos_token_id=2				,	hidden_size=32				,	intermediate_size=37				,	layer_norm_eps=1e-05				,	num_attention_heads=4				,	num_hidden_layers=5				,	pad_token_id=1				,	vocab_size=1000				,	)
        _SCREAMING_SNAKE_CASE             =    CLIPTextModel(UpperCamelCase__ )
        _SCREAMING_SNAKE_CASE             =    CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""" )

        _SCREAMING_SNAKE_CASE             =    {
            """unet""": unet,
            """scheduler""": scheduler,
            """vae""": vae,
            """text_encoder""": text_encoder,
            """tokenizer""": tokenizer,
            """safety_checker""": None,
            """feature_extractor""": None,
        }
        return components


    def     snake_case_(	self				,	A				,	A=0 )	->     Any:
        _SCREAMING_SNAKE_CASE             =    torch.manual_seed(UpperCamelCase__ )
        _SCREAMING_SNAKE_CASE             =    {
            """prompt""": """A field of roses""",
            """generator""": generator,
            # Setting height and width to None to prevent OOMs on CPU.
            """height""": None,
            """width""": None,
            """num_inference_steps""": 2,
            """guidance_scale""": 6.0,
            """output_type""": """numpy""",
        }
        return inputs


    def     snake_case_(	self )	->     Union[str, Any]:
        _SCREAMING_SNAKE_CASE             =    """cpu"""  # ensure determinism for the device-dependent torch.Generator
        _SCREAMING_SNAKE_CASE             =    self.get_dummy_components()
        _SCREAMING_SNAKE_CASE             =    StableDiffusionModelEditingPipeline(**UpperCamelCase__ )
        _SCREAMING_SNAKE_CASE             =    sd_pipe.to(UpperCamelCase__ )
        sd_pipe.set_progress_bar_config(disable=UpperCamelCase__ )

        _SCREAMING_SNAKE_CASE             =    self.get_dummy_inputs(UpperCamelCase__ )
        _SCREAMING_SNAKE_CASE             =    sd_pipe(**UpperCamelCase__ ).images
        _SCREAMING_SNAKE_CASE             =    image[0, -3:, -3:, -1]
        assert image.shape == (1, 64, 64, 3)

        _SCREAMING_SNAKE_CASE             =    np.array([0.4755, 0.5132, 0.4976, 0.3904, 0.3554, 0.4765, 0.5139, 0.5158, 0.4889] )

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


    def     snake_case_(	self )	->     Any:
        _SCREAMING_SNAKE_CASE             =    """cpu"""  # ensure determinism for the device-dependent torch.Generator
        _SCREAMING_SNAKE_CASE             =    self.get_dummy_components()
        _SCREAMING_SNAKE_CASE             =    StableDiffusionModelEditingPipeline(**UpperCamelCase__ )
        _SCREAMING_SNAKE_CASE             =    sd_pipe.to(UpperCamelCase__ )
        sd_pipe.set_progress_bar_config(disable=UpperCamelCase__ )

        _SCREAMING_SNAKE_CASE             =    self.get_dummy_inputs(UpperCamelCase__ )
        _SCREAMING_SNAKE_CASE             =    """french fries"""
        _SCREAMING_SNAKE_CASE             =    sd_pipe(**UpperCamelCase__				,	negative_prompt=UpperCamelCase__ )
        _SCREAMING_SNAKE_CASE             =    output.images
        _SCREAMING_SNAKE_CASE             =    image[0, -3:, -3:, -1]

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

        _SCREAMING_SNAKE_CASE             =    np.array([0.4992, 0.5101, 0.5004, 0.3949, 0.3604, 0.4735, 0.5216, 0.5204, 0.4913] )

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


    def     snake_case_(	self )	->     List[Any]:
        _SCREAMING_SNAKE_CASE             =    """cpu"""  # ensure determinism for the device-dependent torch.Generator
        _SCREAMING_SNAKE_CASE             =    self.get_dummy_components()
        _SCREAMING_SNAKE_CASE             =    EulerAncestralDiscreteScheduler(
            beta_start=0.0_0085				,	beta_end=0.012				,	beta_schedule="""scaled_linear""" )
        _SCREAMING_SNAKE_CASE             =    StableDiffusionModelEditingPipeline(**UpperCamelCase__ )
        _SCREAMING_SNAKE_CASE             =    sd_pipe.to(UpperCamelCase__ )
        sd_pipe.set_progress_bar_config(disable=UpperCamelCase__ )

        _SCREAMING_SNAKE_CASE             =    self.get_dummy_inputs(UpperCamelCase__ )
        _SCREAMING_SNAKE_CASE             =    sd_pipe(**UpperCamelCase__ ).images
        _SCREAMING_SNAKE_CASE             =    image[0, -3:, -3:, -1]

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

        _SCREAMING_SNAKE_CASE             =    np.array([0.4747, 0.5372, 0.4779, 0.4982, 0.5543, 0.4816, 0.5238, 0.4904, 0.5027] )

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


    def     snake_case_(	self )	->     Tuple:
        _SCREAMING_SNAKE_CASE             =    """cpu"""  # ensure determinism for the device-dependent torch.Generator
        _SCREAMING_SNAKE_CASE             =    self.get_dummy_components()
        _SCREAMING_SNAKE_CASE             =    PNDMScheduler()
        _SCREAMING_SNAKE_CASE             =    StableDiffusionModelEditingPipeline(**UpperCamelCase__ )
        _SCREAMING_SNAKE_CASE             =    sd_pipe.to(UpperCamelCase__ )
        sd_pipe.set_progress_bar_config(disable=UpperCamelCase__ )

        _SCREAMING_SNAKE_CASE             =    self.get_dummy_inputs(UpperCamelCase__ )
        # the pipeline does not expect pndm so test if it raises error.
        with self.assertRaises(UpperCamelCase__ ):
            _SCREAMING_SNAKE_CASE             =    sd_pipe(**UpperCamelCase__ ).images


    def     snake_case_(	self )	->     str:
        super().test_inference_batch_single_identical(expected_max_diff=5e-3 )





    def     snake_case_(	self )	->     List[Any]:
        super().test_attention_slicing_forward_pass(expected_max_diff=5e-3 )




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

    '''simple docstring'''


    def     snake_case_(	self )	->     int:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()


    def     snake_case_(	self				,	A=0 )	->     Optional[int]:
        _SCREAMING_SNAKE_CASE             =    torch.manual_seed(UpperCamelCase__ )
        _SCREAMING_SNAKE_CASE             =    {
            """prompt""": """A field of roses""",
            """generator""": generator,
            """num_inference_steps""": 3,
            """guidance_scale""": 7.5,
            """output_type""": """numpy""",
        }
        return inputs


    def     snake_case_(	self )	->     Tuple:
        _SCREAMING_SNAKE_CASE             =    """CompVis/stable-diffusion-v1-4"""
        _SCREAMING_SNAKE_CASE             =    StableDiffusionModelEditingPipeline.from_pretrained(UpperCamelCase__				,	safety_checker=UpperCamelCase__ )
        pipe.to(UpperCamelCase__ )
        pipe.set_progress_bar_config(disable=UpperCamelCase__ )
        pipe.enable_attention_slicing()

        _SCREAMING_SNAKE_CASE             =    self.get_inputs()
        _SCREAMING_SNAKE_CASE             =    pipe(**UpperCamelCase__ ).images
        _SCREAMING_SNAKE_CASE             =    image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 512, 512, 3)

        _SCREAMING_SNAKE_CASE             =    np.array(
            [0.674_9496, 0.638_6453, 0.5144_3267, 0.6609_4905, 0.6192_1215, 0.549_1332, 0.574_4417, 0.5807_5106, 0.517_4658] )

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

        # make sure image changes after editing
        pipe.edit_model("""A pack of roses"""				,	"""A pack of blue roses""" )

        _SCREAMING_SNAKE_CASE             =    pipe(**UpperCamelCase__ ).images
        _SCREAMING_SNAKE_CASE             =    image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 512, 512, 3)

        assert np.abs(expected_slice - image_slice ).max() > 1e-1





    def     snake_case_(	self )	->     List[Any]:
        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        _SCREAMING_SNAKE_CASE             =    """CompVis/stable-diffusion-v1-4"""
        _SCREAMING_SNAKE_CASE             =    DDIMScheduler.from_pretrained(UpperCamelCase__				,	subfolder="""scheduler""" )
        _SCREAMING_SNAKE_CASE             =    StableDiffusionModelEditingPipeline.from_pretrained(
            UpperCamelCase__				,	scheduler=UpperCamelCase__				,	safety_checker=UpperCamelCase__ )
        _SCREAMING_SNAKE_CASE             =    pipe.to(UpperCamelCase__ )
        pipe.set_progress_bar_config(disable=UpperCamelCase__ )
        pipe.enable_attention_slicing(1 )
        pipe.enable_sequential_cpu_offload()

        _SCREAMING_SNAKE_CASE             =    self.get_inputs()
        _SCREAMING_SNAKE_CASE             =    pipe(**UpperCamelCase__ )

        _SCREAMING_SNAKE_CASE             =    torch.cuda.max_memory_allocated()
        # make sure that less than 4.4 GB is allocated
        assert mem_bytes < 4.4 * 10**9










'''simple docstring'''

# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
# and https://github.com/hojonathanho/diffusion

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

import flax
import jax.numpy as jnp

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils_flax import (
    CommonSchedulerState,
    FlaxKarrasDiffusionSchedulers,
    FlaxSchedulerMixin,
    FlaxSchedulerOutput,
    add_noise_common,
    get_velocity_common,
)


@flax.struct.dataclass
class       a_  :

    '''simple docstring'''
    UpperCamelCase									=				42
    UpperCamelCase									=				42

    # setable values
    UpperCamelCase									=				42
    UpperCamelCase									=				42
    UpperCamelCase									=				None


    @classmethod
    def     snake_case_(	cls				,	A				,	A				,	A				,	A				,	)	->     List[Any]:
        return cls(
            common=A				,	final_alpha_cumprod=A				,	init_noise_sigma=A				,	timesteps=A				,	)




@dataclass
class       a_  (       snake_case_ ):

    '''simple docstring'''
    UpperCamelCase									=				42




class       a_  (       snake_case_   ,					snake_case_ ):

    '''simple docstring'''
    UpperCamelCase									=				[e.name for e in FlaxKarrasDiffusionSchedulers]

    UpperCamelCase									=				42


    @property
    def     snake_case_(	self )	->     List[Any]:
        return True


    @register_to_config
    def __init__(	self				,	A = 1000				,	A = 0.0001				,	A = 0.02				,	A = "linear"				,	A = None				,	A = True				,	A = 0				,	A = "epsilon"				,	A = jnp.floataa				,	)	->     int:
        _SCREAMING_SNAKE_CASE             =    dtype


    def     snake_case_(	self				,	A = None )	->     DDIMSchedulerState:
        if common is None:
            _SCREAMING_SNAKE_CASE             =    CommonSchedulerState.create(self )

        # At every step in ddim, we are looking into the previous alphas_cumprod
        # For the final step, there is no previous alphas_cumprod because we are already at 0
        # `set_alpha_to_one` decides whether we set this parameter simply to one or
        # whether we use the final alpha of the "non-previous" one.
        _SCREAMING_SNAKE_CASE             =    (
            jnp.array(1.0				,	dtype=self.dtype ) if self.config.set_alpha_to_one else common.alphas_cumprod[0]
        )

        # standard deviation of the initial noise distribution
        _SCREAMING_SNAKE_CASE             =    jnp.array(1.0				,	dtype=self.dtype )

        _SCREAMING_SNAKE_CASE             =    jnp.arange(0				,	self.config.num_train_timesteps ).round()[::-1]

        return DDIMSchedulerState.create(
            common=A				,	final_alpha_cumprod=A				,	init_noise_sigma=A				,	timesteps=A				,	)


    def     snake_case_(	self				,	A				,	A				,	A = None )	->     jnp.ndarray:
        return sample


    def     snake_case_(	self				,	A				,	A				,	A = () )	->     DDIMSchedulerState:
        _SCREAMING_SNAKE_CASE             =    self.config.num_train_timesteps // num_inference_steps
        # creates integer timesteps by multiplying by ratio
        # rounding to avoid issues when num_inference_step is power of 3
        _SCREAMING_SNAKE_CASE             =    (jnp.arange(0				,	A ) * step_ratio).round()[::-1] + self.config.steps_offset

        return state.replace(
            num_inference_steps=A				,	timesteps=A				,	)


    def     snake_case_(	self				,	A				,	A				,	A )	->     int:
        _SCREAMING_SNAKE_CASE             =    state.common.alphas_cumprod[timestep]
        _SCREAMING_SNAKE_CASE             =    jnp.where(
            prev_timestep >= 0				,	state.common.alphas_cumprod[prev_timestep]				,	state.final_alpha_cumprod )
        _SCREAMING_SNAKE_CASE             =    1 - alpha_prod_t
        _SCREAMING_SNAKE_CASE             =    1 - alpha_prod_t_prev

        _SCREAMING_SNAKE_CASE             =    (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)

        return variance


    def     snake_case_(	self				,	A				,	A				,	A				,	A				,	A = 0.0				,	A = True				,	)	->     Union[FlaxDDIMSchedulerOutput, Tuple]:
        if state.num_inference_steps is None:
            raise ValueError(
                """Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler""" )

        # See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf
        # Ideally, read DDIM paper in-detail understanding

        # Notation (<variable name> -> <name in paper>
        # - pred_noise_t -> e_theta(x_t, t)
        # - pred_original_sample -> f_theta(x_t, t) or x_0
        # - std_dev_t -> sigma_t
        # - eta -> η
        # - pred_sample_direction -> "direction pointing to x_t"
        # - pred_prev_sample -> "x_t-1"

        # 1. get previous step value (=t-1)
        _SCREAMING_SNAKE_CASE             =    timestep - self.config.num_train_timesteps // state.num_inference_steps

        _SCREAMING_SNAKE_CASE             =    state.common.alphas_cumprod
        _SCREAMING_SNAKE_CASE             =    state.final_alpha_cumprod

        # 2. compute alphas, betas
        _SCREAMING_SNAKE_CASE             =    alphas_cumprod[timestep]
        _SCREAMING_SNAKE_CASE             =    jnp.where(prev_timestep >= 0				,	alphas_cumprod[prev_timestep]				,	A )

        _SCREAMING_SNAKE_CASE             =    1 - alpha_prod_t

        # 3. compute predicted original sample from predicted noise also called
        # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
        if self.config.prediction_type == "epsilon":
            _SCREAMING_SNAKE_CASE             =    (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
            _SCREAMING_SNAKE_CASE             =    model_output
        elif self.config.prediction_type == "sample":
            _SCREAMING_SNAKE_CASE             =    model_output
            _SCREAMING_SNAKE_CASE             =    (sample - alpha_prod_t ** 0.5 * pred_original_sample) / beta_prod_t ** 0.5
        elif self.config.prediction_type == "v_prediction":
            _SCREAMING_SNAKE_CASE             =    (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
            _SCREAMING_SNAKE_CASE             =    (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
        else:
            raise ValueError(
                f'prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or'
                """ `v_prediction`""" )

        # 4. compute variance: "sigma_t(η)" -> see formula (16)
        # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
        _SCREAMING_SNAKE_CASE             =    self._get_variance(A				,	A				,	A )
        _SCREAMING_SNAKE_CASE             =    eta * variance ** 0.5

        # 5. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
        _SCREAMING_SNAKE_CASE             =    (1 - alpha_prod_t_prev - std_dev_t**2) ** 0.5 * pred_epsilon

        # 6. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
        _SCREAMING_SNAKE_CASE             =    alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction

        if not return_dict:
            return (prev_sample, state)

        return FlaxDDIMSchedulerOutput(prev_sample=A				,	state=A )


    def     snake_case_(	self				,	A				,	A				,	A				,	A				,	)	->     jnp.ndarray:
        return add_noise_common(state.common				,	A				,	A				,	A )


    def     snake_case_(	self				,	A				,	A				,	A				,	A				,	)	->     jnp.ndarray:
        return get_velocity_common(state.common				,	A				,	A				,	A )





    def __len__(	self )	->     Any:
        return self.config.num_train_timesteps



