





"""simple docstring"""




import os
import tempfile
import unittest

import numpy as np

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


if is_flax_available():
   import jax
   import jax.numpy as jnp
   from flax.jax_utils import replicate
   from flax.training.common_utils import shard

   from diffusers import FlaxDDIMScheduler, FlaxDiffusionPipeline, FlaxStableDiffusionPipeline

@require_flax
class  _UpperCAmelCase  (   unittest.TestCase		):



      '''simple docstring'''

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        with tempfile.TemporaryDirectory() as tmpdirname:
          # pipeline has Flax weights
          A_  :       List[Any]										=  FlaxDiffusionPipeline.from_pretrained(
              'hf-internal-testing/tiny-stable-diffusion-pipe'					,      safety_checker=snake_case_					,      cache_dir=snake_case_  )

          A_  :       Any										=  [t[-1] for t in os.walk(os.path.join(snake_case_					,      os.listdir(snake_case_  )[0]					,      'snapshots'  )  )]
          A_  :       Tuple										=  [item for sublist in all_root_files for item in sublist]

          # None of the downloaded files should be a PyTorch file even if we have some here:
          # https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_pytorch_model.bin
          assert not any(f.endswith('.bin'  ) for f in files  )


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



      '''simple docstring'''

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_						,       A_  :       str										=  FlaxStableDiffusionPipeline.from_pretrained(
            'hf-internal-testing/tiny-stable-diffusion-pipe'					,      safety_checker=snake_case_  )

        A_  :       str										=  (
            'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
            ' field, close up, split lighting, cinematic'
        )

        A_  :       Optional[int]										=  jax.random.PRNGKey(0  )
        A_  :       List[Any]										=  4

        A_  :       int										=  jax.device_count()
        A_  :       Any										=  num_samples * [prompt]
        A_  :       Union[str, Any]										=  pipeline.prepare_inputs(snake_case_  )

        # shard inputs and rng
        A_  :       str										=  replicate(snake_case_  )
        A_  :       List[Any]										=  jax.random.split(snake_case_					,      snake_case_  )
        A_  :       Any										=  shard(snake_case_  )

        A_  :       List[Any]										=  pipeline(snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      jit=snake_case_  ).images

        assert images.shape == (num_samples, 1, 6_4, 6_4, 3)
        if jax.device_count() == 8:
          assert np.abs(np.abs(images[0, 0, :2, :2, -2:]					,      dtype=np.floataa  ).sum() - 4.1_51_47_45  ) < 1E-3
          assert np.abs(np.abs(snake_case_					,      dtype=np.floataa  ).sum() - 4_99_47.8_75  ) < 5E-1

        A_  :       List[Any]										=  pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]  )  )  )
        assert len(snake_case_  ) == num_samples

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_						,       A_  :       Optional[int]										=  FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'					,      revision='flax'					,      safety_checker=snake_case_  )

        A_  :       Any										=  (
            'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
            ' field, close up, split lighting, cinematic'
        )

        A_  :       Optional[Any]										=  jax.random.PRNGKey(0  )
        A_  :       List[str]										=  5_0

        A_  :       Any										=  jax.device_count()
        A_  :       int										=  num_samples * [prompt]
        A_  :       str										=  pipeline.prepare_inputs(snake_case_  )

        # shard inputs and rng
        A_  :       str										=  replicate(snake_case_  )
        A_  :       Dict										=  jax.random.split(snake_case_					,      snake_case_  )
        A_  :       Optional[int]										=  shard(snake_case_  )

        A_  :       Any										=  pipeline(snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      jit=snake_case_  ).images

        assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
        if jax.device_count() == 8:
          assert np.abs((np.abs(images[0, 0, :2, :2, -2:]					,      dtype=np.floataa  ).sum() - 0.05_65_24_01)  ) < 1E-3
          assert np.abs((np.abs(snake_case_					,      dtype=np.floataa  ).sum() - 2_38_38_08.2)  ) < 5E-1

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_						,       A_  :       List[Any]										=  FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'					,      revision='bf16'					,      dtype=jnp.bfloataa					,      safety_checker=snake_case_  )

        A_  :       Optional[int]										=  (
            'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
            ' field, close up, split lighting, cinematic'
        )

        A_  :       Optional[Any]										=  jax.random.PRNGKey(0  )
        A_  :       int										=  5_0

        A_  :       List[Any]										=  jax.device_count()
        A_  :       Dict										=  num_samples * [prompt]
        A_  :       str										=  pipeline.prepare_inputs(snake_case_  )

        # shard inputs and rng
        A_  :       Dict										=  replicate(snake_case_  )
        A_  :       str										=  jax.random.split(snake_case_					,      snake_case_  )
        A_  :       Any										=  shard(snake_case_  )

        A_  :       List[str]										=  pipeline(snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      jit=snake_case_  ).images

        assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
        if jax.device_count() == 8:
          assert np.abs((np.abs(images[0, 0, :2, :2, -2:]					,      dtype=np.floataa  ).sum() - 0.04_00_39_06)  ) < 1E-3
          assert np.abs((np.abs(snake_case_					,      dtype=np.floataa  ).sum() - 2_37_35_16.75)  ) < 5E-1

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_						,       A_  :       Optional[int]										=  FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'					,      revision='bf16'					,      dtype=jnp.bfloataa  )

        A_  :       List[Any]										=  (
            'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
            ' field, close up, split lighting, cinematic'
        )

        A_  :       Dict										=  jax.random.PRNGKey(0  )
        A_  :       Any										=  5_0

        A_  :       List[Any]										=  jax.device_count()
        A_  :       Any										=  num_samples * [prompt]
        A_  :       int										=  pipeline.prepare_inputs(snake_case_  )

        # shard inputs and rng
        A_  :       str										=  replicate(snake_case_  )
        A_  :       str										=  jax.random.split(snake_case_					,      snake_case_  )
        A_  :       Optional[int]										=  shard(snake_case_  )

        A_  :       Tuple										=  pipeline(snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      jit=snake_case_  ).images

        assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
        if jax.device_count() == 8:
          assert np.abs((np.abs(images[0, 0, :2, :2, -2:]					,      dtype=np.floataa  ).sum() - 0.04_00_39_06)  ) < 1E-3
          assert np.abs((np.abs(snake_case_					,      dtype=np.floataa  ).sum() - 2_37_35_16.75)  ) < 5E-1

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Tuple										=  FlaxDDIMScheduler(
            beta_start=0.0_00_85					,      beta_end=0.0_12					,      beta_schedule='scaled_linear'					,      set_alpha_to_one=snake_case_					,      steps_offset=1					,      )

        A_						,       A_  :       Dict										=  FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'					,      revision='bf16'					,      dtype=jnp.bfloataa					,      scheduler=snake_case_					,      safety_checker=snake_case_					,      )
        A_  :       List[str]										=  scheduler.create_state()

        A_  :       Any										=  scheduler_state

        A_  :       Optional[int]										=  (
            'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
            ' field, close up, split lighting, cinematic'
        )

        A_  :       Union[str, Any]										=  jax.random.PRNGKey(0  )
        A_  :       Union[str, Any]										=  5_0

        A_  :       Union[str, Any]										=  jax.device_count()
        A_  :       List[Any]										=  num_samples * [prompt]
        A_  :       Union[str, Any]										=  pipeline.prepare_inputs(snake_case_  )

        # shard inputs and rng
        A_  :       Dict										=  replicate(snake_case_  )
        A_  :       List[str]										=  jax.random.split(snake_case_					,      snake_case_  )
        A_  :       Any										=  shard(snake_case_  )

        A_  :       Any										=  pipeline(snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      jit=snake_case_  ).images

        assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
        if jax.device_count() == 8:
          assert np.abs((np.abs(images[0, 0, :2, :2, -2:]					,      dtype=np.floataa  ).sum() - 0.0_45_04_39_45)  ) < 1E-3
          assert np.abs((np.abs(snake_case_					,      dtype=np.floataa  ).sum() - 2_34_76_93.5)  ) < 5E-1

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Optional[Any]										=  (
            'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
            ' field, close up, split lighting, cinematic'
        )

        A_  :       Dict										=  jax.device_count()
        A_  :       Any										=  num_samples * [prompt]
        A_  :       List[str]										=  jax.random.split(jax.random.PRNGKey(0  )					,      snake_case_  )

        A_						,       A_  :       str										=  FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'					,      revision='bf16'					,      dtype=jnp.bfloataa					,      safety_checker=snake_case_					,      )

        A_  :       Union[str, Any]										=  replicate(snake_case_  )
        A_  :       str										=  pipeline.prepare_inputs(snake_case_  )
        A_  :       str										=  shard(snake_case_  )
        A_  :       Any										=  pipeline(snake_case_					,      snake_case_					,      snake_case_					,      jit=snake_case_  ).images
        assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
        A_  :       Optional[int]										=  images[2, 0, 2_5_6, 1_0:1_7, 1]

        # With memory efficient attention
        A_						,       A_  :       Optional[int]										=  FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'					,      revision='bf16'					,      dtype=jnp.bfloataa					,      safety_checker=snake_case_					,      use_memory_efficient_attention=snake_case_					,      )

        A_  :       str										=  replicate(snake_case_  )
        A_  :       Any										=  pipeline.prepare_inputs(snake_case_  )
        A_  :       Optional[int]										=  shard(snake_case_  )
        A_  :       Optional[Any]										=  pipeline(snake_case_					,      snake_case_					,      snake_case_					,      jit=snake_case_  ).images
        assert images_eff.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
        A_  :       Optional[int]										=  images[2, 0, 2_5_6, 1_0:1_7, 1]

        # I checked the results visually and they are very similar. However, I saw that the max diff is `1` and the `sum`
        # over the 8 images is exactly `256`, which is very suspicious. Testing a random slice for now.
        assert abs(slice_eff - slice  ).max() < 1E-2






"""simple docstring"""




import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


lowerCamelCase_     :						Tuple						=					logging.get_logger(__name__)





def 						UpperCAmelCase__		(     _UpperCAmelCase					):






  """simple docstring"""





  A_  :       Optional[Any]										=  R'\w+[.]\d+'
  A_  :       int										=  re.findall(_UpperCAmelCase ,					_UpperCAmelCase					)
  for pat in pats:
    A_  :       Optional[int]										=  key.replace(_UpperCAmelCase ,					'_'.join(pat.split('.'					)					)					)
  return key





def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase					):






  """simple docstring"""





  A_  :       List[Any]										=  pt_tuple_key[:-1] + ('scale',)
  if (
      any('norm' in str_ for str_ in pt_tuple_key					)
      and (pt_tuple_key[-1] == "bias")
      and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
      and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
  ):
    A_  :       Union[str, Any]										=  pt_tuple_key[:-1] + ('scale',)
    return renamed_pt_tuple_key, pt_tensor
  elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
    A_  :       List[str]										=  pt_tuple_key[:-1] + ('scale',)
    return renamed_pt_tuple_key, pt_tensor

  # embedding
  if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
    A_  :       Optional[Any]										=  pt_tuple_key[:-1] + ('embedding',)
    return renamed_pt_tuple_key, pt_tensor

  # conv layer
  A_  :       int										=  pt_tuple_key[:-1] + ('kernel',)
  if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
    A_  :       str										=  pt_tensor.transpose(2 ,					3 ,					1 ,					0					)
    return renamed_pt_tuple_key, pt_tensor

  # linear layer
  A_  :       Optional[Any]										=  pt_tuple_key[:-1] + ('kernel',)
  if pt_tuple_key[-1] == "weight":
    A_  :       Optional[Any]										=  pt_tensor.T
    return renamed_pt_tuple_key, pt_tensor

  # old PyTorch layer norm weight
  A_  :       Tuple										=  pt_tuple_key[:-1] + ('weight',)
  if pt_tuple_key[-1] == "gamma":
    return renamed_pt_tuple_key, pt_tensor

  # old PyTorch layer norm bias
  A_  :       Optional[int]										=  pt_tuple_key[:-1] + ('bias',)
  if pt_tuple_key[-1] == "beta":
    return renamed_pt_tuple_key, pt_tensor

  return pt_tuple_key, pt_tensor





def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase=42					):






  """simple docstring"""





  A_  :       int										=  {k: v.numpy() for k, v in pt_state_dict.items()}

  # Step 2: Since the model is stateless, get random Flax params
  A_  :       Union[str, Any]										=  flax_model.init_weights(PRNGKey(_UpperCAmelCase					)					)

  A_  :       Optional[Any]										=  flatten_dict(_UpperCAmelCase					)
  A_  :       Tuple										=  {}

  # Need to change some parameters name to match Flax names
  for pt_key, pt_tensor in pt_state_dict.items():
    A_  :       Any										=  rename_key(_UpperCAmelCase					)
    A_  :       List[str]										=  tuple(renamed_pt_key.split('.'					)					)

    # Correctly rename weight parameters
    A_						,       A_  :       Union[str, Any]										=  rename_key_and_reshape_tensor(_UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase					)

    if flax_key in random_flax_state_dict:
      if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
        raise ValueError(
            f"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
            f"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""					)

        # also add unexpected weight so that warning is thrown
    A_  :       str										=  jnp.asarray(_UpperCAmelCase					)

  return unflatten_dict(_UpperCAmelCase					)