
'''simple docstring'''
import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPanoramaPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






@skip_mps
class    UpperCAmelCase_				(	snake_case__    ,						snake_case__    ,						unittest.TestCase					):
 UpperCAmelCase_		    =						StableDiffusionPanoramaPipeline
 UpperCAmelCase_		    =						TEXT_TO_IMAGE_PARAMS
 UpperCAmelCase_		    =						TEXT_TO_IMAGE_BATCH_PARAMS
 UpperCAmelCase_		    =						TEXT_TO_IMAGE_IMAGE_PARAMS
 UpperCAmelCase_		    =						TEXT_TO_IMAGE_IMAGE_PARAMS





 def        snake_case__		(    self):
     torch.manual_seed(0)
     snake_case_ :    int     			=						UNetaDConditionModel(
         block_out_channels=(32, 64) ,       layers_per_block=1 ,       sample_size=32 ,       in_channels=4 ,       out_channels=4 ,       down_block_types=("DownBlock2D", "CrossAttnDownBlock2D") ,       up_block_types=("CrossAttnUpBlock2D", "UpBlock2D") ,       cross_attention_dim=32 ,       )
     snake_case_ :    Any     			=						DDIMScheduler()
     torch.manual_seed(0)
     snake_case_ :    Optional[Any]     			=						AutoencoderKL(
         block_out_channels=[32, 64] ,       in_channels=3 ,       out_channels=3 ,       down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"] ,       up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"] ,       latent_channels=4 ,       )
     torch.manual_seed(0)
     snake_case_ :    Dict     			=						CLIPTextConfig(
         bos_token_id=0 ,       eos_token_id=2 ,       hidden_size=32 ,       intermediate_size=37 ,       layer_norm_eps=1E-05 ,       num_attention_heads=4 ,       num_hidden_layers=5 ,       pad_token_id=1 ,       vocab_size=10_00 ,       )
     snake_case_ :    List[Any]     			=						CLIPTextModel(lowercase_)
     snake_case_ :    List[str]     			=						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

     snake_case_ :    Optional[Any]     			=						{
         "unet": unet,
         "scheduler": scheduler,
         "vae": vae,
         "text_encoder": text_encoder,
         "tokenizer": tokenizer,
         "safety_checker": None,
         "feature_extractor": None,
     }
     return components





 def        snake_case__		(    self ,       lowercase_ ,       lowercase_=0):
     snake_case_ :    str     			=						torch.manual_seed(lowercase_)
     snake_case_ :    Optional[int]     			=						{
         "prompt": "a photo of the dolomites",
         "generator": generator,
         # Setting height and width to None to prevent OOMs on CPU.
         "height": None,
         "width": None,
         "num_inference_steps": 1,
         "guidance_scale": 6.0,
         "output_type": "numpy",
     }
     return inputs





 def        snake_case__		(    self):
     snake_case_ :    Optional[Any]     			=						"cpu"  # ensure determinism for the device-dependent torch.Generator
     snake_case_ :    Dict     			=						self.get_dummy_components()
     snake_case_ :    List[Any]     			=						StableDiffusionPanoramaPipeline(**lowercase_)
     snake_case_ :    Optional[int]     			=						sd_pipe.to(lowercase_)
     sd_pipe.set_progress_bar_config(disable=lowercase_)

     snake_case_ :    Optional[int]     			=						self.get_dummy_inputs(lowercase_)
     snake_case_ :    Optional[Any]     			=						sd_pipe(**lowercase_).images
     snake_case_ :    Union[str, Any]     			=						image[0, -3:, -3:, -1]
     assert image.shape == (1, 64, 64, 3)

     snake_case_ :    Tuple     			=						np.array([0.6_186, 0.5_374, 0.4_915, 0.4_135, 0.4_114, 0.4_563, 0.5_128, 0.4_977, 0.4_757])

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





 def        snake_case__		(    self):
     super().test_inference_batch_consistent(batch_sizes=[1, 2])





 def        snake_case__		(    self):
     super().test_inference_batch_single_identical(batch_size=2 ,       expected_max_diff=3.25E-3)





 def        snake_case__		(    self):
     snake_case_ :    Optional[Any]     			=						"cpu"  # ensure determinism for the device-dependent torch.Generator
     snake_case_ :    Optional[Any]     			=						self.get_dummy_components()
     snake_case_ :    Any     			=						StableDiffusionPanoramaPipeline(**lowercase_)
     snake_case_ :    List[str]     			=						sd_pipe.to(lowercase_)
     sd_pipe.set_progress_bar_config(disable=lowercase_)

     snake_case_ :    Union[str, Any]     			=						self.get_dummy_inputs(lowercase_)
     snake_case_ :    Any     			=						"french fries"
     snake_case_ :    List[Any]     			=						sd_pipe(**lowercase_ ,       negative_prompt=lowercase_)
     snake_case_ :    List[str]     			=						output.images
     snake_case_ :    str     			=						image[0, -3:, -3:, -1]

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

     snake_case_ :    Optional[int]     			=						np.array([0.6_187, 0.5_375, 0.4_915, 0.4_136, 0.4_114, 0.4_563, 0.5_128, 0.4_976, 0.4_757])

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





 def        snake_case__		(    self):
     snake_case_ :    List[str]     			=						"cpu"  # ensure determinism for the device-dependent torch.Generator
     snake_case_ :    int     			=						self.get_dummy_components()
     snake_case_ :    int     			=						StableDiffusionPanoramaPipeline(**lowercase_)
     snake_case_ :    List[Any]     			=						sd_pipe.to(lowercase_)
     sd_pipe.set_progress_bar_config(disable=lowercase_)

     snake_case_ :    Union[str, Any]     			=						self.get_dummy_inputs(lowercase_)
     snake_case_ :    Any     			=						sd_pipe(**lowercase_ ,       view_batch_size=2)
     snake_case_ :    Any     			=						output.images
     snake_case_ :    Tuple     			=						image[0, -3:, -3:, -1]

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

     snake_case_ :    str     			=						np.array([0.6_187, 0.5_375, 0.4_915, 0.4_136, 0.4_114, 0.4_563, 0.5_128, 0.4_976, 0.4_757])

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





 def        snake_case__		(    self):
     snake_case_ :    List[Any]     			=						"cpu"  # ensure determinism for the device-dependent torch.Generator
     snake_case_ :    Any     			=						self.get_dummy_components()
     snake_case_ :    List[Any]     			=						EulerAncestralDiscreteScheduler(
         beta_start=0.00_085 ,       beta_end=0.012 ,       beta_schedule="scaled_linear")
     snake_case_ :    str     			=						StableDiffusionPanoramaPipeline(**lowercase_)
     snake_case_ :    Any     			=						sd_pipe.to(lowercase_)
     sd_pipe.set_progress_bar_config(disable=lowercase_)

     snake_case_ :    Any     			=						self.get_dummy_inputs(lowercase_)
     snake_case_ :    List[str]     			=						sd_pipe(**lowercase_).images
     snake_case_ :    List[Any]     			=						image[0, -3:, -3:, -1]

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

     snake_case_ :    int     			=						np.array([0.4_024, 0.6_510, 0.4_901, 0.5_378, 0.5_813, 0.5_622, 0.4_795, 0.4_467, 0.4_952])

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





 def        snake_case__		(    self):
     snake_case_ :    Union[str, Any]     			=						"cpu"  # ensure determinism for the device-dependent torch.Generator
     snake_case_ :    List[str]     			=						self.get_dummy_components()
     snake_case_ :    int     			=						PNDMScheduler(
         beta_start=0.00_085 ,       beta_end=0.012 ,       beta_schedule="scaled_linear" ,       skip_prk_steps=lowercase_)
     snake_case_ :    Dict     			=						StableDiffusionPanoramaPipeline(**lowercase_)
     snake_case_ :    List[str]     			=						sd_pipe.to(lowercase_)
     sd_pipe.set_progress_bar_config(disable=lowercase_)

     snake_case_ :    List[str]     			=						self.get_dummy_inputs(lowercase_)
     snake_case_ :    str     			=						sd_pipe(**lowercase_).images
     snake_case_ :    str     			=						image[0, -3:, -3:, -1]

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

     snake_case_ :    int     			=						np.array([0.6_391, 0.6_291, 0.4_861, 0.5_134, 0.5_552, 0.4_578, 0.5_032, 0.5_023, 0.4_539])

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






@slow
@require_torch_gpu
class    UpperCAmelCase_				(	unittest.TestCase					):





 def        snake_case__		(    self):
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()





 def        snake_case__		(    self ,       lowercase_=0):
     snake_case_ :    Tuple     			=						torch.manual_seed(lowercase_)
     snake_case_ :    str     			=						{
         "prompt": "a photo of the dolomites",
         "generator": generator,
         "num_inference_steps": 3,
         "guidance_scale": 7.5,
         "output_type": "numpy",
     }
     return inputs





 def        snake_case__		(    self):
     snake_case_ :    Optional[Any]     			=						"stabilityai/stable-diffusion-2-base"
     snake_case_ :    Dict     			=						DDIMScheduler.from_pretrained(lowercase_ ,       subfolder="scheduler")
     snake_case_ :    Any     			=						StableDiffusionPanoramaPipeline.from_pretrained(lowercase_ ,       scheduler=lowercase_ ,       safety_checker=lowercase_)
     pipe.to(lowercase_)
     pipe.set_progress_bar_config(disable=lowercase_)
     pipe.enable_attention_slicing()

     snake_case_ :    List[Any]     			=						self.get_inputs()
     snake_case_ :    Optional[int]     			=						pipe(**lowercase_).images
     snake_case_ :    Any     			=						image[0, -3:, -3:, -1].flatten()

     assert image.shape == (1, 5_12, 20_48, 3)

     snake_case_ :    int     			=						np.array(
         [
             0.36_968_392,
             0.27_025_372,
             0.32_446_766,
             0.28_379_387,
             0.36_363_274,
             0.30_733_347,
             0.27_100_027,
             0.27_054_125,
             0.25_536_096,
         ])

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





 def        snake_case__		(    self):
     snake_case_ :    int     			=						StableDiffusionPanoramaPipeline.from_pretrained(
         "stabilityai/stable-diffusion-2-base" ,       safety_checker=lowercase_)
     snake_case_ :    Dict     			=						LMSDiscreteScheduler.from_config(pipe.scheduler.config)
     pipe.to(lowercase_)
     pipe.set_progress_bar_config(disable=lowercase_)
     pipe.enable_attention_slicing()

     snake_case_ :    int     			=						self.get_inputs()
     snake_case_ :    str     			=						pipe(**lowercase_).images
     snake_case_ :    Tuple     			=						image[0, -3:, -3:, -1].flatten()

     assert image.shape == (1, 5_12, 20_48, 3)

     snake_case_ :    List[Any]     			=						np.array(
         [
             [
                 0.0,
                 0.0,
                 0.0,
                 0.0,
                 0.0,
                 0.0,
                 0.0,
                 0.0,
                 0.0,
             ]
         ])

     assert np.abs(expected_slice - image_slice).max() < 1E-3





 def        snake_case__		(    self):
     snake_case_ :    int     			=						0

     def callback_fn(lowercase_ ,       lowercase_ ,       lowercase_) -> None:
         snake_case_ :    Dict     			=						True
         nonlocal number_of_steps
         number_of_steps += 1
         if step == 1:
             snake_case_ :    List[str]     			=						latents.detach().cpu().numpy()
             assert latents.shape == (1, 4, 64, 2_56)
             snake_case_ :    Tuple     			=						latents[0, -3:, -3:, -1]

             snake_case_ :    int     			=						np.array(
                 [
                     0.18_681_869,
                     0.33_907_816,
                     0.5_361_276,
                     0.14_432_865,
                     -0.02_856_611,
                     -0.73_941_123,
                     0.23_397_987,
                     0.47_322_682,
                     -0.37_823_164,
                 ])
             assert np.abs(latents_slice.flatten() - expected_slice).max() < 5E-2
         elif step == 2:
             snake_case_ :    Any     			=						latents.detach().cpu().numpy()
             assert latents.shape == (1, 4, 64, 2_56)
             snake_case_ :    Union[str, Any]     			=						latents[0, -3:, -3:, -1]

             snake_case_ :    Optional[Any]     			=						np.array(
                 [
                     0.18_539_645,
                     0.33_987_248,
                     0.5_378_559,
                     0.14_437_142,
                     -0.02_455_261,
                     -0.7_338_317,
                     0.23_990_755,
                     0.47_356_272,
                     -0.3_786_505,
                 ])

             assert np.abs(latents_slice.flatten() - expected_slice).max() < 5E-2

     snake_case_ :    Optional[int]     			=						False

     snake_case_ :    str     			=						"stabilityai/stable-diffusion-2-base"
     snake_case_ :    Optional[Any]     			=						DDIMScheduler.from_pretrained(lowercase_ ,       subfolder="scheduler")
     snake_case_ :    Union[str, Any]     			=						StableDiffusionPanoramaPipeline.from_pretrained(lowercase_ ,       scheduler=lowercase_ ,       safety_checker=lowercase_)
     snake_case_ :    Tuple     			=						pipe.to(lowercase_)
     pipe.set_progress_bar_config(disable=lowercase_)
     pipe.enable_attention_slicing()

     snake_case_ :    Union[str, Any]     			=						self.get_inputs()
     pipe(**lowercase_ ,       callback=lowercase_ ,       callback_steps=1)
     assert callback_fn.has_been_called
     assert number_of_steps == 3





 def        snake_case__		(    self):
     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()

     snake_case_ :    Union[str, Any]     			=						"stabilityai/stable-diffusion-2-base"
     snake_case_ :    Union[str, Any]     			=						DDIMScheduler.from_pretrained(lowercase_ ,       subfolder="scheduler")
     snake_case_ :    Optional[Any]     			=						StableDiffusionPanoramaPipeline.from_pretrained(lowercase_ ,       scheduler=lowercase_ ,       safety_checker=lowercase_)
     snake_case_ :    Tuple     			=						pipe.to(lowercase_)
     pipe.set_progress_bar_config(disable=lowercase_)
     pipe.enable_attention_slicing(1)
     pipe.enable_sequential_cpu_offload()

     snake_case_ :    List[str]     			=						self.get_inputs()
     snake_case_ :    str     			=						pipe(**lowercase_)

     snake_case_ :    List[Any]     			=						torch.cuda.max_memory_allocated()
     # make sure that less than 5.2 GB is allocated
     assert mem_bytes < 5.5 * 10**9






'''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


a_											=						logging.get_logger(__name__)

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    snake_case_ :    List[Any]     			=						r"\w+[.]\d+"
    snake_case_ :    Tuple     			=						re.findall(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       )
    for pat in pats:
        snake_case_ :    Union[str, Any]     			=						key.replace(__SCREAMING_SNAKE_CASE,		"_".join(pat.split("."       )       )       )
    return key

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    snake_case_ :    Tuple     			=						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)
    ):
        snake_case_ :    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:
        snake_case_ :    Dict     			=						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:
        snake_case_ :    List[Any]     			=						pt_tuple_key[:-1] + ("embedding",)
        return renamed_pt_tuple_key, pt_tensor

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

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

    # old PyTorch layer norm weight
    snake_case_ :    Union[str, Any]     			=						pt_tuple_key[:-1] + ("weight",)
    if pt_tuple_key[-1] == "gamma":
        return renamed_pt_tuple_key, pt_tensor

    # old PyTorch layer norm bias
    snake_case_ :    Union[str, Any]     			=						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_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE=4_2       ):
    """simple docstring"""






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

    # Step 2: Since the model is stateless, get random Flax params
    snake_case_ :    Optional[int]     			=						flax_model.init_weights(PRNGKey(__SCREAMING_SNAKE_CASE       )       )

    snake_case_ :    Union[str, Any]     			=						flatten_dict(__SCREAMING_SNAKE_CASE       )
    snake_case_ :    str     			=						{}

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

        # Correctly rename weight parameters
        snake_case_			, snake_case_ :    Dict     			=						rename_key_and_reshape_tensor(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       )

        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
        snake_case_ :    Any     			=						jnp.asarray(__SCREAMING_SNAKE_CASE       )

    return unflatten_dict(__SCREAMING_SNAKE_CASE       )




