




import unittest

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

from diffusers import DDIMScheduler, LDMPipeline, UNetaDModel, VQModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()




class    __lowerCAmelCase       (							unittest.TestCase):






       @property
       def 					SCREAMING_SNAKE_CASE    (   self:					str				):
         torch.manual_seed(0				)
         lowercase    :str      	=						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






       @property
       def 					SCREAMING_SNAKE_CASE    (   self:					Union[str, Any]				):
         torch.manual_seed(0				)
         lowercase    :Optional[int]      	=						VQModel(
             block_out_channels=[32, 64]       ,					in_channels=3       ,					out_channels=3       ,					down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]       ,					up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]       ,					latent_channels=3       ,					)
         return model






       @property
       def 					SCREAMING_SNAKE_CASE    (   self:					Tuple				):
         torch.manual_seed(0				)
         lowercase    :Tuple      	=						CLIPTextConfig(
             bos_token_id=0       ,					eos_token_id=2       ,					hidden_size=32       ,					intermediate_size=37       ,					layer_norm_eps=1e-0_5       ,					num_attention_heads=4       ,					num_hidden_layers=5       ,					pad_token_id=1       ,					vocab_size=10_00       ,					)
         return CLIPTextModel(_lowerCAmelCase				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Tuple				):
         lowercase    :int      	=						self.dummy_uncond_unet
         lowercase    :List[Any]      	=						DDIMScheduler()
         lowercase    :List[Any]      	=						self.dummy_vq_model

         lowercase    :List[str]      	=						LDMPipeline(unet=_lowerCAmelCase       ,					vqvae=_lowerCAmelCase       ,					scheduler=_lowerCAmelCase				)
         ldm.to(_lowerCAmelCase				)
         ldm.set_progress_bar_config(disable=_lowerCAmelCase				)

         lowercase    :Optional[Any]      	=						torch.manual_seed(0				)
         lowercase    :List[str]      	=						ldm(generator=_lowerCAmelCase       ,					num_inference_steps=2       ,					output_type="numpy"				).images

         lowercase    :List[Any]      	=						torch.manual_seed(0				)
         lowercase    :int      	=						ldm(generator=_lowerCAmelCase       ,					num_inference_steps=2       ,					output_type="numpy"       ,					return_dict=_lowerCAmelCase				)[0]

         lowercase    :Any      	=						image[0, -3:, -3:, -1]
         lowercase    :Optional[Any]      	=						image_from_tuple[0, -3:, -3:, -1]

         assert image.shape == (1, 64, 64, 3)
         lowercase    :List[str]      	=						np.array([0.85_12, 0.8_18, 0.64_11, 0.68_08, 0.44_65, 0.56_18, 0.46, 0.62_31, 0.51_72]				)
         lowercase    :str      	=						1e-2 if torch_device != "mps" else 3e-2

         assert np.abs(image_slice.flatten() - expected_slice				).max() < tolerance
         assert np.abs(image_from_tuple_slice.flatten() - expected_slice				).max() < tolerance







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






       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]				):
         lowercase    :Dict      	=						LDMPipeline.from_pretrained("CompVis/ldm-celebahq-256"				)
         ldm.to(_lowerCAmelCase				)
         ldm.set_progress_bar_config(disable=_lowerCAmelCase				)

         lowercase    :Any      	=						torch.manual_seed(0				)
         lowercase    :Optional[int]      	=						ldm(generator=_lowerCAmelCase       ,					num_inference_steps=5       ,					output_type="numpy"				).images

         lowercase    :Tuple      	=						image[0, -3:, -3:, -1]

         assert image.shape == (1, 2_56, 2_56, 3)
         lowercase    :Tuple      	=						np.array([0.43_99, 0.4_49_75, 0.4_68_25, 0.4_74, 0.43_59, 0.45_81, 0.4_50_95, 0.43_41, 0.44_47]				)
         lowercase    :Optional[Any]      	=						1e-2 if torch_device != "mps" else 3e-2

         assert np.abs(image_slice.flatten() - expected_slice				).max() < tolerance









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


_UpperCAmelCase	:						List[Any]			=     logging.get_logger(__name__)



def       UpperCAmelCase__  (     lowerCamelCase   ):
  lowercase    :int      	=						r"\w+[.]\d+"
  lowercase    :Tuple      	=						re.findall(lowerCamelCase,			lowerCamelCase   )
  for pat in pats:
    lowercase    :List[str]      	=						key.replace(lowerCamelCase,			"_".join(pat.split("."   )   )   )
  return key



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase,			lowerCamelCase   ):
  lowercase    :List[str]      	=						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)
  ):
    lowercase    :Optional[int]      	=						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:
    lowercase    :int      	=						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:
    lowercase    :List[str]      	=						pt_tuple_key[:-1] + ("embedding",)
    return renamed_pt_tuple_key, pt_tensor

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

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

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

  # old PyTorch layer norm bias
  lowercase    :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__  (     lowerCamelCase,			lowerCamelCase,			lowerCamelCase=42   ):
  # Step 1: Convert pytorch tensor to numpy
  lowercase    :Optional[int]      	=						{k: v.numpy() for k, v in pt_state_dict.items()}

  # Step 2: Since the model is stateless, get random Flax params
  lowercase    :str      	=						flax_model.init_weights(PRNGKey(lowerCamelCase   )   )

  lowercase    :Tuple      	=						flatten_dict(lowerCamelCase   )
  lowercase    :Optional[Any]      	=						{}

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

    # Correctly rename weight parameters
    lowercase	,    lowercase    :List[str]      	=						rename_key_and_reshape_tensor(lowerCamelCase,			lowerCamelCase,			lowerCamelCase   )

    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
    lowercase    :List[str]      	=						jnp.asarray(lowerCamelCase   )

  return unflatten_dict(lowerCamelCase   )



