
import torch

from diffusers import EulerDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest



class 		UpperCamelCase__  (       __lowercase    ):
   _SCREAMING_SNAKE_CASE      :							int	       =      (EulerDiscreteScheduler,)
   _SCREAMING_SNAKE_CASE      :							Any	       =      10
   def       lowerCAmelCase			(self :					Any					,				**snake_case_ :					Any   ):
    __a :  Dict						=  {
        '''num_train_timesteps''': 1_1_0_0,
        '''beta_start''': 0.0001,
        '''beta_end''': 0.02,
        '''beta_schedule''': '''linear''',
    }

    config.update(**snake_case_   )
    return config
   def       lowerCAmelCase			(self :					Optional[Any]   ):
    for timesteps in [1_0, 5_0, 1_0_0, 1_0_0_0]:
     self.check_over_configs(num_train_timesteps=snake_case_   )
   def       lowerCAmelCase			(self :					Any   ):
    for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001]					,				[0.0002, 0.002, 0.02]   ):
     self.check_over_configs(beta_start=snake_case_					,				beta_end=snake_case_   )
   def       lowerCAmelCase			(self :					List[str]   ):
    for schedule in ["linear", "scaled_linear"]:
     self.check_over_configs(beta_schedule=snake_case_   )
   def       lowerCAmelCase			(self :					List[Any]   ):
    for prediction_type in ["epsilon", "v_prediction"]:
     self.check_over_configs(prediction_type=snake_case_   )
   def       lowerCAmelCase			(self :					Tuple   ):
    __a :  List[str]						=  self.scheduler_classes[0]
    __a :  List[Any]						=  self.get_scheduler_config()
    __a :  Optional[Any]						=  scheduler_class(**snake_case_   )

    scheduler.set_timesteps(self.num_inference_steps   )

    __a :  Dict						=  torch.manual_seed(0   )

    __a :  Tuple						=  self.dummy_model()
    __a :  str						=  self.dummy_sample_deter * scheduler.init_noise_sigma
    __a :  Tuple						=  sample.to(snake_case_   )

    for i, t in enumerate(scheduler.timesteps   ):
     __a :  Any						=  scheduler.scale_model_input(snake_case_					,				snake_case_   )

     __a :  Any						=  model(snake_case_					,				snake_case_   )

     __a :  str						=  scheduler.step(snake_case_					,				snake_case_					,				snake_case_					,				generator=snake_case_   )
     __a :  Union[str, Any]						=  output.prev_sample

    __a :  Optional[int]						=  torch.sum(torch.abs(snake_case_   )   )
    __a :  Optional[int]						=  torch.mean(torch.abs(snake_case_   )   )

    assert abs(result_sum.item() - 10.0807   ) < 1E-2
    assert abs(result_mean.item() - 0.0131   ) < 1E-3
   def       lowerCAmelCase			(self :					Tuple   ):
    __a :  List[str]						=  self.scheduler_classes[0]
    __a :  List[Any]						=  self.get_scheduler_config(prediction_type='''v_prediction'''   )
    __a :  List[str]						=  scheduler_class(**snake_case_   )

    scheduler.set_timesteps(self.num_inference_steps   )

    __a :  Dict						=  torch.manual_seed(0   )

    __a :  List[Any]						=  self.dummy_model()
    __a :  str						=  self.dummy_sample_deter * scheduler.init_noise_sigma
    __a :  List[Any]						=  sample.to(snake_case_   )

    for i, t in enumerate(scheduler.timesteps   ):
     __a :  Tuple						=  scheduler.scale_model_input(snake_case_					,				snake_case_   )

     __a :  str						=  model(snake_case_					,				snake_case_   )

     __a :  Optional[int]						=  scheduler.step(snake_case_					,				snake_case_					,				snake_case_					,				generator=snake_case_   )
     __a :  List[str]						=  output.prev_sample

    __a :  int						=  torch.sum(torch.abs(snake_case_   )   )
    __a :  int						=  torch.mean(torch.abs(snake_case_   )   )

    assert abs(result_sum.item() - 0.0002   ) < 1E-2
    assert abs(result_mean.item() - 2.2676E-06   ) < 1E-3
   def       lowerCAmelCase			(self :					List[str]   ):
    __a :  List[Any]						=  self.scheduler_classes[0]
    __a :  Optional[int]						=  self.get_scheduler_config()
    __a :  str						=  scheduler_class(**snake_case_   )

    scheduler.set_timesteps(self.num_inference_steps					,				device=snake_case_   )

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

    __a :  List[Any]						=  self.dummy_model()
    __a :  int						=  self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
    __a :  str						=  sample.to(snake_case_   )

    for t in scheduler.timesteps:
     __a :  List[str]						=  scheduler.scale_model_input(snake_case_					,				snake_case_   )

     __a :  Optional[Any]						=  model(snake_case_					,				snake_case_   )

     __a :  List[str]						=  scheduler.step(snake_case_					,				snake_case_					,				snake_case_					,				generator=snake_case_   )
     __a :  List[str]						=  output.prev_sample

    __a :  List[str]						=  torch.sum(torch.abs(snake_case_   )   )
    __a :  Tuple						=  torch.mean(torch.abs(snake_case_   )   )

    assert abs(result_sum.item() - 10.0807   ) < 1E-2
    assert abs(result_mean.item() - 0.0131   ) < 1E-3







   def       lowerCAmelCase			(self :					str   ):
    __a :  Union[str, Any]						=  self.scheduler_classes[0]
    __a :  Tuple						=  self.get_scheduler_config()
    __a :  Any						=  scheduler_class(**snake_case_					,				use_karras_sigmas=snake_case_   )

    scheduler.set_timesteps(self.num_inference_steps					,				device=snake_case_   )

    __a :  str						=  torch.manual_seed(0   )

    __a :  Any						=  self.dummy_model()
    __a :  str						=  self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
    __a :  List[Any]						=  sample.to(snake_case_   )

    for t in scheduler.timesteps:
     __a :  Any						=  scheduler.scale_model_input(snake_case_					,				snake_case_   )

     __a :  str						=  model(snake_case_					,				snake_case_   )

     __a :  List[Any]						=  scheduler.step(snake_case_					,				snake_case_					,				snake_case_					,				generator=snake_case_   )
     __a :  Dict						=  output.prev_sample

    __a :  List[str]						=  torch.sum(torch.abs(snake_case_   )   )
    __a :  List[str]						=  torch.mean(torch.abs(snake_case_   )   )

    assert abs(result_sum.item() - 124.52_2994_9951_1719   ) < 1E-2
    assert abs(result_mean.item() - 0.1_6213_9326_3339_9963   ) < 1E-3


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 		UpperCamelCase__  (       unittest.TestCase    ):
   def       lowerCAmelCase			(self :					Optional[Any]   ):
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
   def       lowerCAmelCase			(self :					int   ):
    __a				,       __a :  Optional[Any]						=  FlaxStableDiffusionPipeline.from_pretrained(
        '''stabilityai/stable-diffusion-2'''					,				revision='''bf16'''					,				dtype=jnp.bfloataa					,				)

    __a :  Any						=  '''A painting of a squirrel eating a burger'''
    __a :  Dict						=  jax.device_count()
    __a :  Optional[int]						=  num_samples * [prompt]
    __a :  Optional[Any]						=  sd_pipe.prepare_inputs(snake_case_   )

    __a :  Optional[Any]						=  replicate(snake_case_   )
    __a :  Optional[int]						=  shard(snake_case_   )

    __a :  int						=  jax.random.PRNGKey(0   )
    __a :  str						=  jax.random.split(snake_case_					,				jax.device_count()   )

    __a :  int						=  sd_pipe(snake_case_					,				snake_case_					,				snake_case_					,				num_inference_steps=2_5					,				jit=snake_case_   )[0]
    assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

    __a :  Optional[int]						=  images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]   )
    __a :  List[Any]						=  images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

    __a :  Dict						=  jnp.asarray(jax.device_get(image_slice.flatten()   )   )
    __a :  str						=  jnp.array([0.4238, 0.4414, 0.4395, 0.4453, 0.4629, 0.4590, 0.4531, 0.4_5508, 0.4512]   )
    print(f"output_slice: {output_slice}"   )
    assert jnp.abs(output_slice - expected_slice   ).max() < 1E-2







   def       lowerCAmelCase			(self :					Tuple   ):
    __a :  Optional[Any]						=  '''stabilityai/stable-diffusion-2'''
    __a				,       __a :  Dict						=  FlaxDPMSolverMultistepScheduler.from_pretrained(snake_case_					,				subfolder='''scheduler'''   )
    __a				,       __a :  List[Any]						=  FlaxStableDiffusionPipeline.from_pretrained(
        snake_case_					,				scheduler=snake_case_					,				revision='''bf16'''					,				dtype=jnp.bfloataa					,				)
    __a :  Union[str, Any]						=  scheduler_params

    __a :  List[Any]						=  '''A painting of a squirrel eating a burger'''
    __a :  Any						=  jax.device_count()
    __a :  Any						=  num_samples * [prompt]
    __a :  List[Any]						=  sd_pipe.prepare_inputs(snake_case_   )

    __a :  Tuple						=  replicate(snake_case_   )
    __a :  Dict						=  shard(snake_case_   )

    __a :  Dict						=  jax.random.PRNGKey(0   )
    __a :  Dict						=  jax.random.split(snake_case_					,				jax.device_count()   )

    __a :  str						=  sd_pipe(snake_case_					,				snake_case_					,				snake_case_					,				num_inference_steps=2_5					,				jit=snake_case_   )[0]
    assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

    __a :  List[Any]						=  images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]   )
    __a :  Any						=  images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

    __a :  List[Any]						=  jnp.asarray(jax.device_get(image_slice.flatten()   )   )
    __a :  List[Any]						=  jnp.array([0.4336, 0.4_2969, 0.4453, 0.4199, 0.4297, 0.4531, 0.4434, 0.4434, 0.4297]   )
    print(f"output_slice: {output_slice}"   )
    assert jnp.abs(output_slice - expected_slice   ).max() < 1E-2
