





"""simple docstring"""


import unittest

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

from diffusers import KandinskyPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class      lowercase_				(     __lowerCAmelCase				,unittest.TestCase							):
       snake_case__				:							List[Any]											=     KandinskyPriorPipeline
       snake_case__				:							Any											=     ["""prompt"""]
       snake_case__				:							int											=     ["""prompt""", """negative_prompt"""]
       snake_case__				:							str											=     [
           """num_images_per_prompt""",
           """generator""",
           """num_inference_steps""",
           """latents""",
           """negative_prompt""",
           """guidance_scale""",
           """output_type""",
           """return_dict""",
       ]
       snake_case__				:							Union[str, Any]											=     False

       @property
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        return 32

       @property
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        return 32

       @property
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        return self.time_input_dim

       @property
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        return self.time_input_dim * 4

       @property
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        return 100

       @property
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )
        return tokenizer

       @property
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        torch.manual_seed(0     )
        snake_case__    							=				CLIPTextConfig(
            bos_token_id=0   ,					eos_token_id=2   ,					hidden_size=self.text_embedder_hidden_size   ,					projection_dim=self.text_embedder_hidden_size   ,					intermediate_size=37   ,					layer_norm_eps=1e-05   ,					num_attention_heads=4   ,					num_hidden_layers=5   ,					pad_token_id=1   ,					vocab_size=1000   ,					)
        return CLIPTextModelWithProjection(__SCREAMING_SNAKE_CASE     )

       @property
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        torch.manual_seed(0     )

        snake_case__    							=				{
            "num_attention_heads": 2,
            "attention_head_dim": 12,
            "embedding_dim": self.text_embedder_hidden_size,
            "num_layers": 1,
        }

        snake_case__    							=				PriorTransformer(**__SCREAMING_SNAKE_CASE     )
        # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
        snake_case__    							=				nn.Parameter(torch.ones(model.clip_std.shape     )     )
        return model

       @property
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        torch.manual_seed(0     )
        snake_case__    							=				CLIPVisionConfig(
            hidden_size=self.text_embedder_hidden_size   ,					image_size=224   ,					projection_dim=self.text_embedder_hidden_size   ,					intermediate_size=37   ,					num_attention_heads=4   ,					num_channels=3   ,					num_hidden_layers=5   ,					patch_size=14   ,					)

        snake_case__    							=				CLIPVisionModelWithProjection(__SCREAMING_SNAKE_CASE     )
        return model

       @property
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				CLIPImageProcessor(
            crop_size=224   ,					do_center_crop=__SCREAMING_SNAKE_CASE   ,					do_normalize=__SCREAMING_SNAKE_CASE   ,					do_resize=__SCREAMING_SNAKE_CASE   ,					image_mean=[0.48145466, 0.4578275, 0.40821073]   ,					image_std=[0.26862954, 0.26130258, 0.27577711]   ,					resample=3   ,					size=224   ,					)

        return image_processor

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				self.dummy_prior
        snake_case__    							=				self.dummy_image_encoder
        snake_case__    							=				self.dummy_text_encoder
        snake_case__    							=				self.dummy_tokenizer
        snake_case__    							=				self.dummy_image_processor

        snake_case__    							=				UnCLIPScheduler(
            variance_type="fixed_small_log"   ,					prediction_type="sample"   ,					num_train_timesteps=1000   ,					clip_sample=__SCREAMING_SNAKE_CASE   ,					clip_sample_range=10.0   ,					)

        snake_case__    							=				{
            "prior": prior,
            "image_encoder": image_encoder,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "scheduler": scheduler,
            "image_processor": image_processor,
        }

        return components

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE=0     ):






        '''simple docstring'''

        if str(__SCREAMING_SNAKE_CASE     ).startswith("mps"     ):
         snake_case__    							=				torch.manual_seed(__SCREAMING_SNAKE_CASE     )
        else:
         snake_case__    							=				torch.Generator(device=__SCREAMING_SNAKE_CASE     ).manual_seed(__SCREAMING_SNAKE_CASE     )
        snake_case__    							=				{
            "prompt": "horse",
            "generator": generator,
            "guidance_scale": 4.0,
            "num_inference_steps": 2,
            "output_type": "np",
        }
        return inputs

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				"cpu"

        snake_case__    							=				self.get_dummy_components()

        snake_case__    							=				self.pipeline_class(**__SCREAMING_SNAKE_CASE     )
        snake_case__    							=				pipe.to(__SCREAMING_SNAKE_CASE     )

        pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

        snake_case__    							=				pipe(**self.get_dummy_inputs(__SCREAMING_SNAKE_CASE     )     )
        snake_case__    							=				output.image_embeds

        snake_case__    							=				pipe(
            **self.get_dummy_inputs(__SCREAMING_SNAKE_CASE     )   ,					return_dict=__SCREAMING_SNAKE_CASE   ,					)[0]

        snake_case__    							=				image[0, -10:]
        snake_case__    							=				image_from_tuple[0, -10:]

        assert image.shape == (1, 32)

        snake_case__    							=				np.array(
            [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156]     )

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

       @skip_mps
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				torch_device == "cpu"
        snake_case__    							=				True
        snake_case__    							=				False

        self._test_inference_batch_single_identical(
            test_max_difference=__SCREAMING_SNAKE_CASE   ,					relax_max_difference=__SCREAMING_SNAKE_CASE   ,					test_mean_pixel_difference=__SCREAMING_SNAKE_CASE   ,					)

       @skip_mps
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				torch_device == "cpu"
        snake_case__    							=				False

        self._test_attention_slicing_forward_pass(
            test_max_difference=__SCREAMING_SNAKE_CASE   ,					test_mean_pixel_difference=__SCREAMING_SNAKE_CASE   ,					)








"""simple docstring"""


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

import flax
import jax.numpy as jnp
from scipy import integrate

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



@flax.struct.dataclass
class      lowercase_				:
       snake_case__				:							CommonSchedulerState

       # setable values
       snake_case__				:							jnp.ndarray
       snake_case__				:							jnp.ndarray
       snake_case__				:							jnp.ndarray
       snake_case__				:							Optional[int]											=     None

       # running values
       snake_case__				:							Optional[jnp.ndarray]											=     None

       @classmethod
       def 		_SCREAMING_SNAKE_CASE(			cls   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     ):






        '''simple docstring'''

        return cls(common=__SCREAMING_SNAKE_CASE   ,					init_noise_sigma=__SCREAMING_SNAKE_CASE   ,					timesteps=__SCREAMING_SNAKE_CASE   ,					sigmas=__SCREAMING_SNAKE_CASE     )



@dataclass
class      lowercase_				(     __lowerCAmelCase							):
       snake_case__				:							LMSDiscreteSchedulerState



class      lowercase_				(     __lowerCAmelCase				,__lowerCAmelCase							):
       snake_case__				:							Any											=     [e.name for e in FlaxKarrasDiffusionSchedulers]

       snake_case__				:							jnp.dtype

       @property
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        return True

       @register_to_config
       def __init__(			self   ,					__SCREAMING_SNAKE_CASE = 1000   ,					__SCREAMING_SNAKE_CASE = 0.0001   ,					__SCREAMING_SNAKE_CASE = 0.02   ,					__SCREAMING_SNAKE_CASE = "linear"   ,					__SCREAMING_SNAKE_CASE = None   ,					__SCREAMING_SNAKE_CASE = "epsilon"   ,					__SCREAMING_SNAKE_CASE = jnp.floataa   ,					):






        '''simple docstring'''

        snake_case__    							=				dtype

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE = None     ):






        '''simple docstring'''

        if common is None:
         snake_case__    							=				CommonSchedulerState.create(self     )

        snake_case__    							=				jnp.arange(0   ,					self.config.num_train_timesteps     ).round()[::-1]
        snake_case__    							=				((1 - common.alphas_cumprod) / common.alphas_cumprod) ** 0.5

        # standard deviation of the initial noise distribution
        snake_case__    							=				sigmas.max()

        return LMSDiscreteSchedulerState.create(
            common=__SCREAMING_SNAKE_CASE   ,					init_noise_sigma=__SCREAMING_SNAKE_CASE   ,					timesteps=__SCREAMING_SNAKE_CASE   ,					sigmas=__SCREAMING_SNAKE_CASE   ,					)

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     ):






        '''simple docstring'''

        ((snake_case__)							,   )    							=				jnp.where(state.timesteps == timestep   ,					size=1     )
        snake_case__    							=				step_index[0]

        snake_case__    							=				state.sigmas[step_index]
        snake_case__    							=				sample / ((sigma**2 + 1) ** 0.5)
        return sample

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     ):






        '''simple docstring'''

        def lms_derivative(__SCREAMING_SNAKE_CASE     ):
         snake_case__    							=				1.0
         for k in range(__SCREAMING_SNAKE_CASE     ):
          if current_order == k:
           continue
          prod *= (tau - state.sigmas[t - k]) / (state.sigmas[t - current_order] - state.sigmas[t - k])
         return prod

        snake_case__    							=				integrate.quad(__SCREAMING_SNAKE_CASE   ,					state.sigmas[t]   ,					state.sigmas[t + 1]   ,					epsrel=1e-4     )[0]

        return integrated_coeff

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE = ()     ):






        '''simple docstring'''

        snake_case__    							=				jnp.linspace(self.config.num_train_timesteps - 1   ,					0   ,					__SCREAMING_SNAKE_CASE   ,					dtype=self.dtype     )

        snake_case__    							=				jnp.floor(__SCREAMING_SNAKE_CASE     ).astype(jnp.intaa     )
        snake_case__    							=				jnp.ceil(__SCREAMING_SNAKE_CASE     ).astype(jnp.intaa     )

        snake_case__    							=				jnp.mod(__SCREAMING_SNAKE_CASE   ,					1.0     )

        snake_case__    							=				((1 - state.common.alphas_cumprod) / state.common.alphas_cumprod) ** 0.5
        snake_case__    							=				(1 - frac) * sigmas[low_idx] + frac * sigmas[high_idx]
        snake_case__    							=				jnp.concatenate([sigmas, jnp.array([0.0]   ,					dtype=self.dtype     )]     )

        snake_case__    							=				timesteps.astype(jnp.intaa     )

        # initial running values
        snake_case__    							=				jnp.zeros((0,) + shape   ,					dtype=self.dtype     )

        return state.replace(
            timesteps=__SCREAMING_SNAKE_CASE   ,					sigmas=__SCREAMING_SNAKE_CASE   ,					num_inference_steps=__SCREAMING_SNAKE_CASE   ,					derivatives=__SCREAMING_SNAKE_CASE   ,					)

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE = 4   ,					__SCREAMING_SNAKE_CASE = True   ,					):






        '''simple docstring'''

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

        snake_case__    							=				state.sigmas[timestep]

        # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
        if self.config.prediction_type == "epsilon":
         snake_case__    							=				sample - sigma * model_output
        elif self.config.prediction_type == "v_prediction":
         # * c_out + input * c_skip
         snake_case__    							=				model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))
        else:
         raise ValueError(
             F"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`"     )

        # 2. Convert to an ODE derivative
        snake_case__    							=				(sample - pred_original_sample) / sigma
        snake_case__    							=				state.replace(derivatives=jnp.append(state.derivatives   ,					__SCREAMING_SNAKE_CASE     )     )
        if len(state.derivatives     ) > order:
         snake_case__    							=				state.replace(derivatives=jnp.delete(state.derivatives   ,					0     )     )

        # 3. Compute linear multistep coefficients
        snake_case__    							=				min(timestep + 1   ,					__SCREAMING_SNAKE_CASE     )
        snake_case__    							=				[self.get_lms_coefficient(__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     ) for curr_order in range(__SCREAMING_SNAKE_CASE     )]

        # 4. Compute previous sample based on the derivatives path
        snake_case__    							=				sample + sum(
            coeff * derivative for coeff, derivative in zip(__SCREAMING_SNAKE_CASE   ,					reversed(state.derivatives     )     )     )

        if not return_dict:
         return (prev_sample, state)

        return FlaxLMSSchedulerOutput(prev_sample=__SCREAMING_SNAKE_CASE   ,					state=__SCREAMING_SNAKE_CASE     )

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					):






        '''simple docstring'''

        snake_case__    							=				state.sigmas[timesteps].flatten()
        snake_case__    							=				broadcast_to_shape_from_left(__SCREAMING_SNAKE_CASE   ,					noise.shape     )

        snake_case__    							=				original_samples + noise * sigma

        return noisy_samples

       def __len__(			self     ):






        '''simple docstring'''

        return self.config.num_train_timesteps

