





"""simple docstring"""





import gc
import unittest

from parameterized import parameterized

from diffusers import FlaxUNetaDConditionModel
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import load_hf_numpy, require_flax, slow


if is_flax_available():
       import jax
       import jax.numpy as jnp

@slow
@require_flax
class 						lowercase   (   unittest.TestCase   ):





      def        _snake_case     (			self ,		lowercase ,		lowercase					)						->	List[str]:
            return f'gaussian_noise_s={seed}_shape={"_".join([str(lowercase					) for s in shape]					)}.npy'





      def        _snake_case     (			self					)						->	List[Any]:
            # clean up the VRAM after each test
            super().tearDown()
            gc.collect()





      def        _snake_case     (			self ,		lowercase=0 ,		lowercase=(4, 4, 64, 64) ,		lowercase=False					)						->	Union[str, Any]:
            lowerCAmelCase    		=			jnp.bfloataa if fpaa else jnp.floataa
            lowerCAmelCase    		=			jnp.array(load_hf_numpy(self.get_file_format(lowercase ,		lowercase					)					) ,		dtype=lowercase					)
            return image





      def        _snake_case     (			self ,		lowercase=False ,		lowercase="CompVis/stable-diffusion-v1-4"					)						->	Optional[Any]:
            lowerCAmelCase    		=			jnp.bfloataa if fpaa else jnp.floataa
            lowerCAmelCase    		=			"""bf16""" if fpaa else None

            lowerCAmelCase		,						lowerCAmelCase    		=			FlaxUNetaDConditionModel.from_pretrained(
                lowercase ,		subfolder="""unet""" ,		dtype=lowercase ,		revision=lowercase					)
            return model, params





      def        _snake_case     (			self ,		lowercase=0 ,		lowercase=(4, 77, 768) ,		lowercase=False					)						->	str:
            lowerCAmelCase    		=			jnp.bfloataa if fpaa else jnp.floataa
            lowerCAmelCase    		=			jnp.array(load_hf_numpy(self.get_file_format(lowercase ,		lowercase					)					) ,		dtype=lowercase					)
            return hidden_states





      @parameterized.expand(
          [
              # fmt: off
              [83, 4, [-0.2_323, -0.1_304, 0.0_813, -0.3_093, -0.0_919, -0.1_571, -0.1_125, -0.5_806]],
              [17, 0.55, [-0.0_831, -0.2_443, 0.0_901, -0.0_919, 0.3_396, 0.0_103, -0.3_743, 0.0_701]],
              [8, 0.89, [-0.4_863, 0.0_859, 0.0_875, -0.1_658, 0.9_199, -0.0_114, 0.4_839, 0.4_639]],
              [3, 1_000, [-0.5_649, 0.2_402, -0.5_518, 0.1_248, 1.1_328, -0.2_443, -0.0_325, -1.0_078]],
              # fmt: on
          ]					)
      def        _snake_case     (			self ,		lowercase ,		lowercase ,		lowercase					)						->	int:
            lowerCAmelCase		,						lowerCAmelCase    		=			self.get_unet_model(model_id="""CompVis/stable-diffusion-v1-4""" ,		fpaa=lowercase					)
            lowerCAmelCase    		=			self.get_latents(lowercase ,		fpaa=lowercase					)
            lowerCAmelCase    		=			self.get_encoder_hidden_states(lowercase ,		fpaa=lowercase					)

            lowerCAmelCase    		=			model.apply(
                {"""params""": params} ,		lowercase ,		jnp.array(lowercase ,		dtype=jnp.intaa					) ,		encoder_hidden_states=lowercase ,		).sample

            assert sample.shape == latents.shape

            lowerCAmelCase    		=			jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())					) ,		dtype=jnp.floataa					)
            lowerCAmelCase    		=			jnp.array(lowercase ,		dtype=jnp.floataa					)

            # Found torch (float16) and flax (bfloat16) outputs to be within this tolerance, in the same hardware
            assert jnp.allclose(lowercase ,		lowercase ,		atol=1e-2					)





      @parameterized.expand(
          [
              # fmt: off
              [83, 4, [0.1_514, 0.0_807, 0.1_624, 0.1_016, -0.1_896, 0.0_263, 0.0_677, 0.2_310]],
              [17, 0.55, [0.1_164, -0.0_216, 0.0_170, 0.1_589, -0.3_120, 0.1_005, -0.0_581, -0.1_458]],
              [8, 0.89, [-0.1_758, -0.0_169, 0.1_004, -0.1_411, 0.1_312, 0.1_103, -0.1_996, 0.2_139]],
              [3, 1_000, [0.1_214, 0.0_352, -0.0_731, -0.1_562, -0.0_994, -0.0_906, -0.2_340, -0.0_539]],
              # fmt: on
          ]					)
      def        _snake_case     (			self ,		lowercase ,		lowercase ,		lowercase					)						->	Dict:
            lowerCAmelCase		,						lowerCAmelCase    		=			self.get_unet_model(model_id="""stabilityai/stable-diffusion-2""" ,		fpaa=lowercase					)
            lowerCAmelCase    		=			self.get_latents(lowercase ,		shape=(4, 4, 96, 96) ,		fpaa=lowercase					)
            lowerCAmelCase    		=			self.get_encoder_hidden_states(lowercase ,		shape=(4, 77, 1_024) ,		fpaa=lowercase					)

            lowerCAmelCase    		=			model.apply(
                {"""params""": params} ,		lowercase ,		jnp.array(lowercase ,		dtype=jnp.intaa					) ,		encoder_hidden_states=lowercase ,		).sample

            assert sample.shape == latents.shape

            lowerCAmelCase    		=			jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())					) ,		dtype=jnp.floataa					)
            lowerCAmelCase    		=			jnp.array(lowercase ,		dtype=jnp.floataa					)

            # Found torch (float16) and flax (bfloat16) outputs to be within this tolerance, on the same hardware
            assert jnp.allclose(lowercase ,		lowercase ,		atol=1e-2					)













"""simple docstring"""





import unittest

import numpy as np
import torch

from diffusers import ScoreSdeVePipeline, ScoreSdeVeScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()

class 						lowercase   (   unittest.TestCase   ):





      @property
      def        _snake_case     (			self					)						->	List[str]:
            torch.manual_seed(0					)
            lowerCAmelCase    		=			UNetaDModel(
                block_out_channels=(32, 64) ,		layers_per_block=2 ,		sample_size=32 ,		in_channels=3 ,		out_channels=3 ,		down_block_types=("""DownBlock2D""", """AttnDownBlock2D""") ,		up_block_types=("""AttnUpBlock2D""", """UpBlock2D""") ,		)
            return model





      def        _snake_case     (			self					)						->	List[Any]:
            lowerCAmelCase    		=			self.dummy_uncond_unet
            lowerCAmelCase    		=			ScoreSdeVeScheduler()

            lowerCAmelCase    		=			ScoreSdeVePipeline(unet=lowercase ,		scheduler=lowercase					)
            sde_ve.to(lowercase					)
            sde_ve.set_progress_bar_config(disable=lowercase					)

            lowerCAmelCase    		=			torch.manual_seed(0					)
            lowerCAmelCase    		=			sde_ve(num_inference_steps=2 ,		output_type="""numpy""" ,		generator=lowercase					).images

            lowerCAmelCase    		=			torch.manual_seed(0					)
            lowerCAmelCase    		=			sde_ve(num_inference_steps=2 ,		output_type="""numpy""" ,		generator=lowercase ,		return_dict=lowercase					)[
                0
            ]

            lowerCAmelCase    		=			image[0, -3:, -3:, -1]
            lowerCAmelCase    		=			image_from_tuple[0, -3:, -3:, -1]

            assert image.shape == (1, 32, 32, 3)
            lowerCAmelCase    		=			np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]					)

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







@slow
@require_torch
class 						lowercase   (   unittest.TestCase   ):





      def        _snake_case     (			self					)						->	Optional[int]:
            lowerCAmelCase    		=			"""google/ncsnpp-church-256"""
            lowerCAmelCase    		=			UNetaDModel.from_pretrained(lowercase					)

            lowerCAmelCase    		=			ScoreSdeVeScheduler.from_pretrained(lowercase					)

            lowerCAmelCase    		=			ScoreSdeVePipeline(unet=lowercase ,		scheduler=lowercase					)
            sde_ve.to(lowercase					)
            sde_ve.set_progress_bar_config(disable=lowercase					)

            lowerCAmelCase    		=			torch.manual_seed(0					)
            lowerCAmelCase    		=			sde_ve(num_inference_steps=10 ,		output_type="""numpy""" ,		generator=lowercase					).images

            lowerCAmelCase    		=			image[0, -3:, -3:, -1]

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

            lowerCAmelCase    		=			np.array([0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0]					)

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






