



"""simple docstring"""





import gc
import random
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()
class  lowercase__  (       unittest.TestCase ):







       '''simple docstring'''


       def 	lowercase__							(				self     :    Optional[Any]       )					->		int:


             '''simple docstring'''




             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()



       @property
       def 	lowercase__							(				self     :    List[str]       )					->		Optional[int]:


             '''simple docstring'''




             UpperCAmelCase_     					=				1
             UpperCAmelCase_     					=				3
             UpperCAmelCase_     					=				(32, 32)

             UpperCAmelCase_     					=				floats_tensor((batch_size, num_channels) + sizes						,							rng=random.Random(0       )       ).to(_UpperCAmelCase       )
             return image



       @property
       def 	lowercase__							(				self     :    Optional[int]       )					->		int:


             '''simple docstring'''




             torch.manual_seed(0       )
             UpperCAmelCase_     					=				UNetaDConditionModel(
                 block_out_channels=(32, 32, 64)						,							layers_per_block=2						,							sample_size=32						,							in_channels=7						,							out_channels=4						,							down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D")						,							up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D")						,							cross_attention_dim=32						,							attention_head_dim=8						,							use_linear_projection=_UpperCAmelCase						,							only_cross_attention=(True, True, False)						,							num_class_embeds=100						,							)
             return model



       @property
       def 	lowercase__							(				self     :    List[Any]       )					->		List[str]:


             '''simple docstring'''




             torch.manual_seed(0       )
             UpperCAmelCase_     					=				AutoencoderKL(
                 block_out_channels=[32, 32, 64]						,							in_channels=3						,							out_channels=3						,							down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"]						,							up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"]						,							latent_channels=4						,							)
             return model



       @property
       def 	lowercase__							(				self     :    int       )					->		Any:


             '''simple docstring'''




             torch.manual_seed(0       )
             UpperCAmelCase_     					=				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=1000						,							hidden_act="gelu"						,							projection_dim=512						,							)
             return CLIPTextModel(_UpperCAmelCase       )



       def 	lowercase__							(				self     :    str       )					->		Any:


             '''simple docstring'''




             UpperCAmelCase_     					=				"cpu"  # ensure determinism for the device-dependent torch.Generator
             UpperCAmelCase_     					=				self.dummy_cond_unet_upscale
             UpperCAmelCase_     					=				DDPMScheduler()
             UpperCAmelCase_     					=				DDIMScheduler(prediction_type="v_prediction"       )
             UpperCAmelCase_     					=				self.dummy_vae
             UpperCAmelCase_     					=				self.dummy_text_encoder
             UpperCAmelCase_     					=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

             UpperCAmelCase_     					=				self.dummy_image.cpu().permute(0						,							2						,							3						,							1       )[0]
             UpperCAmelCase_     					=				Image.fromarray(np.uinta(_UpperCAmelCase       )       ).convert("RGB"       ).resize((64, 64)       )

             # make sure here that pndm scheduler skips prk
             UpperCAmelCase_     					=				StableDiffusionUpscalePipeline(
                 unet=_UpperCAmelCase						,							low_res_scheduler=_UpperCAmelCase						,							scheduler=_UpperCAmelCase						,							vae=_UpperCAmelCase						,							text_encoder=_UpperCAmelCase						,							tokenizer=_UpperCAmelCase						,							max_noise_level=350						,							)
             UpperCAmelCase_     					=				sd_pipe.to(_UpperCAmelCase       )
             sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase       )

             UpperCAmelCase_     					=				"A painting of a squirrel eating a burger"
             UpperCAmelCase_     					=				torch.Generator(device=_UpperCAmelCase       ).manual_seed(0       )
             UpperCAmelCase_     					=				sd_pipe(
                 [prompt]						,							image=_UpperCAmelCase						,							generator=_UpperCAmelCase						,							guidance_scale=6.0						,							noise_level=20						,							num_inference_steps=2						,							output_type="np"						,							)

             UpperCAmelCase_     					=				output.images

             UpperCAmelCase_     					=				torch.Generator(device=_UpperCAmelCase       ).manual_seed(0       )
             UpperCAmelCase_     					=				sd_pipe(
                 [prompt]						,							image=_UpperCAmelCase						,							generator=_UpperCAmelCase						,							guidance_scale=6.0						,							noise_level=20						,							num_inference_steps=2						,							output_type="np"						,							return_dict=_UpperCAmelCase						,							)[0]

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

             UpperCAmelCase_     					=				low_res_image.size[0] * 4
             assert image.shape == (1, expected_height_width, expected_height_width, 3)
             UpperCAmelCase_     					=				np.array([0.3113, 0.3910, 0.4272, 0.4859, 0.5061, 0.4652, 0.5362, 0.5715, 0.5661]       )

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



       def 	lowercase__							(				self     :    Dict       )					->		List[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				"cpu"  # ensure determinism for the device-dependent torch.Generator
             UpperCAmelCase_     					=				self.dummy_cond_unet_upscale
             UpperCAmelCase_     					=				DDPMScheduler()
             UpperCAmelCase_     					=				DDIMScheduler(prediction_type="v_prediction"       )
             UpperCAmelCase_     					=				self.dummy_vae
             UpperCAmelCase_     					=				self.dummy_text_encoder
             UpperCAmelCase_     					=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

             UpperCAmelCase_     					=				self.dummy_image.cpu().permute(0						,							2						,							3						,							1       )[0]
             UpperCAmelCase_     					=				Image.fromarray(np.uinta(_UpperCAmelCase       )       ).convert("RGB"       ).resize((64, 64)       )

             # make sure here that pndm scheduler skips prk
             UpperCAmelCase_     					=				StableDiffusionUpscalePipeline(
                 unet=_UpperCAmelCase						,							low_res_scheduler=_UpperCAmelCase						,							scheduler=_UpperCAmelCase						,							vae=_UpperCAmelCase						,							text_encoder=_UpperCAmelCase						,							tokenizer=_UpperCAmelCase						,							max_noise_level=350						,							)
             UpperCAmelCase_     					=				sd_pipe.to(_UpperCAmelCase       )
             sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase       )

             UpperCAmelCase_     					=				"A painting of a squirrel eating a burger"
             UpperCAmelCase_     					=				sd_pipe(
                 2 * [prompt]						,							image=2 * [low_res_image]						,							guidance_scale=6.0						,							noise_level=20						,							num_inference_steps=2						,							output_type="np"						,							)
             UpperCAmelCase_     					=				output.images
             assert image.shape[0] == 2

             UpperCAmelCase_     					=				torch.Generator(device=_UpperCAmelCase       ).manual_seed(0       )
             UpperCAmelCase_     					=				sd_pipe(
                 [prompt]						,							image=_UpperCAmelCase						,							generator=_UpperCAmelCase						,							num_images_per_prompt=2						,							guidance_scale=6.0						,							noise_level=20						,							num_inference_steps=2						,							output_type="np"						,							)
             UpperCAmelCase_     					=				output.images
             assert image.shape[0] == 2







       @unittest.skipIf(torch_device != "cuda"						,							"This test requires a GPU"       )
       def 	lowercase__							(				self     :    List[Any]       )					->		List[str]:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.dummy_cond_unet_upscale
             UpperCAmelCase_     					=				DDPMScheduler()
             UpperCAmelCase_     					=				DDIMScheduler(prediction_type="v_prediction"       )
             UpperCAmelCase_     					=				self.dummy_vae
             UpperCAmelCase_     					=				self.dummy_text_encoder
             UpperCAmelCase_     					=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

             UpperCAmelCase_     					=				self.dummy_image.cpu().permute(0						,							2						,							3						,							1       )[0]
             UpperCAmelCase_     					=				Image.fromarray(np.uinta(_UpperCAmelCase       )       ).convert("RGB"       ).resize((64, 64)       )

             # put models in fp16, except vae as it overflows in fp16
             UpperCAmelCase_     					=				unet.half()
             UpperCAmelCase_     					=				text_encoder.half()

             # make sure here that pndm scheduler skips prk
             UpperCAmelCase_     					=				StableDiffusionUpscalePipeline(
                 unet=_UpperCAmelCase						,							low_res_scheduler=_UpperCAmelCase						,							scheduler=_UpperCAmelCase						,							vae=_UpperCAmelCase						,							text_encoder=_UpperCAmelCase						,							tokenizer=_UpperCAmelCase						,							max_noise_level=350						,							)
             UpperCAmelCase_     					=				sd_pipe.to(_UpperCAmelCase       )
             sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase       )

             UpperCAmelCase_     					=				"A painting of a squirrel eating a burger"
             UpperCAmelCase_     					=				torch.manual_seed(0       )
             UpperCAmelCase_     					=				sd_pipe(
                 [prompt]						,							image=_UpperCAmelCase						,							generator=_UpperCAmelCase						,							num_inference_steps=2						,							output_type="np"						,							).images

             UpperCAmelCase_     					=				low_res_image.size[0] * 4
             assert image.shape == (1, expected_height_width, expected_height_width, 3)



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







       '''simple docstring'''


       def 	lowercase__							(				self     :    Optional[int]       )					->		int:


             '''simple docstring'''




             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()



       def 	lowercase__							(				self     :    List[Any]       )					->		Optional[int]:


             '''simple docstring'''




             UpperCAmelCase_     					=				load_image(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
                 "/sd2-upscale/low_res_cat.png"       )
             UpperCAmelCase_     					=				load_numpy(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
                 "/upsampled_cat.npy"       )

             UpperCAmelCase_     					=				"stabilityai/stable-diffusion-x4-upscaler"
             UpperCAmelCase_     					=				StableDiffusionUpscalePipeline.from_pretrained(_UpperCAmelCase       )
             pipe.to(_UpperCAmelCase       )
             pipe.set_progress_bar_config(disable=_UpperCAmelCase       )
             pipe.enable_attention_slicing()

             UpperCAmelCase_     					=				"a cat sitting on a park bench"

             UpperCAmelCase_     					=				torch.manual_seed(0       )
             UpperCAmelCase_     					=				pipe(
                 prompt=_UpperCAmelCase						,							image=_UpperCAmelCase						,							generator=_UpperCAmelCase						,							output_type="np"						,							)
             UpperCAmelCase_     					=				output.images[0]

             assert image.shape == (512, 512, 3)
             assert np.abs(expected_image - image       ).max() < 1e-3



       def 	lowercase__							(				self     :    str       )					->		List[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				load_image(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
                 "/sd2-upscale/low_res_cat.png"       )
             UpperCAmelCase_     					=				load_numpy(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
                 "/upsampled_cat_fp16.npy"       )

             UpperCAmelCase_     					=				"stabilityai/stable-diffusion-x4-upscaler"
             UpperCAmelCase_     					=				StableDiffusionUpscalePipeline.from_pretrained(
                 _UpperCAmelCase						,							torch_dtype=torch.floataa						,							)
             pipe.to(_UpperCAmelCase       )
             pipe.set_progress_bar_config(disable=_UpperCAmelCase       )
             pipe.enable_attention_slicing()

             UpperCAmelCase_     					=				"a cat sitting on a park bench"

             UpperCAmelCase_     					=				torch.manual_seed(0       )
             UpperCAmelCase_     					=				pipe(
                 prompt=_UpperCAmelCase						,							image=_UpperCAmelCase						,							generator=_UpperCAmelCase						,							output_type="np"						,							)
             UpperCAmelCase_     					=				output.images[0]

             assert image.shape == (512, 512, 3)
             assert np.abs(expected_image - image       ).max() < 5e-1







       def 	lowercase__							(				self     :    int       )					->		Optional[Any]:


             '''simple docstring'''




             torch.cuda.empty_cache()
             torch.cuda.reset_max_memory_allocated()
             torch.cuda.reset_peak_memory_stats()

             UpperCAmelCase_     					=				load_image(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
                 "/sd2-upscale/low_res_cat.png"       )

             UpperCAmelCase_     					=				"stabilityai/stable-diffusion-x4-upscaler"
             UpperCAmelCase_     					=				StableDiffusionUpscalePipeline.from_pretrained(
                 _UpperCAmelCase						,							torch_dtype=torch.floataa						,							)
             pipe.to(_UpperCAmelCase       )
             pipe.set_progress_bar_config(disable=_UpperCAmelCase       )
             pipe.enable_attention_slicing(1       )
             pipe.enable_sequential_cpu_offload()

             UpperCAmelCase_     					=				"a cat sitting on a park bench"

             UpperCAmelCase_     					=				torch.manual_seed(0       )
             UpperCAmelCase_     					=				pipe(
                 prompt=_UpperCAmelCase						,							image=_UpperCAmelCase						,							generator=_UpperCAmelCase						,							num_inference_steps=5						,							output_type="np"						,							)

             UpperCAmelCase_     					=				torch.cuda.max_memory_allocated()
             # make sure that less than 2.9 GB is allocated
             assert mem_bytes < 2.9 * 10**9











"""simple docstring"""


import numpy as np




def        lowercase__     (       lowercase_		,lowercase_		,lowercase_ = 1e-12		,lowercase_ = 100		,)       ->		tuple[float, np.ndarray]:




      """simple docstring"""

      assert np.shape(lowercase_      )[0] == np.shape(lowercase_      )[1]
      # Ensure proper dimensionality.
      assert np.shape(lowercase_      )[0] == np.shape(lowercase_      )[0]
      # Ensure inputs are either both complex or both real
      assert np.iscomplexobj(lowercase_      ) == np.iscomplexobj(lowercase_      )
      _UpperCamelCase   :						Dict    =							np.iscomplexobj(lowercase_      )
      if is_complex:
            # Ensure complex input_matrix is Hermitian
            assert np.array_equal(lowercase_		,input_matrix.conj().T      )

      # Set convergence to False. Will define convergence when we exceed max_iterations
      # or when we have small changes from one iteration to next.

      _UpperCamelCase   :						List[Any]    =							False
      _UpperCamelCase   :						Optional[int]    =							0
      _UpperCamelCase   :						int    =							0
      _UpperCamelCase   :						Optional[int]    =							1e12

      while not convergence:
            # Multiple matrix by the vector.
            _UpperCamelCase   :						Optional[int]    =							np.dot(lowercase_		,lowercase_      )
            # Normalize the resulting output vector.
            _UpperCamelCase   :						List[Any]    =							w / np.linalg.norm(lowercase_      )
            # Find rayleigh quotient
            # (faster than usual b/c we know vector is normalized already)
            _UpperCamelCase   :						Optional[Any]    =							vector.conj().T if is_complex else vector.T
            _UpperCamelCase   :						Dict    =							np.dot(lowercase_		,np.dot(lowercase_		,lowercase_      )      )

            # Check convergence.
            _UpperCamelCase   :						Optional[int]    =							np.abs(lambda_ - lambda_previous      ) / lambda_
            iterations += 1

            if error <= error_tol or iterations >= max_iterations:
                  _UpperCamelCase   :						List[Any]    =							True

            _UpperCamelCase   :						Tuple    =							lambda_

      if is_complex:
            _UpperCamelCase   :						Tuple    =							np.real(lambda_      )

      return lambda_, vector




def        lowercase__     (       )       ->		None:




      """simple docstring"""

      _UpperCamelCase   :						Dict    =							np.array([[41, 4, 20], [4, 26, 30], [20, 30, 50]]      )
      _UpperCamelCase   :						Optional[int]    =							np.array([41, 4, 20]      )
      _UpperCamelCase   :						Union[str, Any]    =							real_input_matrix.astype(np.complexaaa      )
      _UpperCamelCase   :						int    =							np.triu(1J * complex_input_matrix		,1      )
      complex_input_matrix += imag_matrix
      complex_input_matrix += -1 * imag_matrix.T
      _UpperCamelCase   :						Union[str, Any]    =							np.array([41, 4, 20]      ).astype(np.complexaaa      )

      for problem_type in ["real", "complex"]:
            if problem_type == "real":
                  _UpperCamelCase   :						Dict    =							real_input_matrix
                  _UpperCamelCase   :						Any    =							real_vector
            elif problem_type == "complex":
                  _UpperCamelCase   :						int    =							complex_input_matrix
                  _UpperCamelCase   :						Union[str, Any]    =							complex_vector

            # Our implementation.
            _UpperCamelCase,	_UpperCamelCase   :						int    =							power_iteration(lowercase_		,lowercase_      )

            # Numpy implementation.

            # Get eigenvalues and eigenvectors using built-in numpy
            # eigh (eigh used for symmetric or hermetian matrices).
            _UpperCamelCase,	_UpperCamelCase   :						str    =							np.linalg.eigh(lowercase_      )
            # Last eigenvalue is the maximum one.
            _UpperCamelCase   :						Any    =							eigen_values[-1]
            # Last column in this matrix is eigenvector corresponding to largest eigenvalue.
            _UpperCamelCase   :						Tuple    =							eigen_vectors[:, -1]

            # Check our implementation and numpy gives close answers.
            assert np.abs(eigen_value - eigen_value_max      ) <= 1e-6
            # Take absolute values element wise of each eigenvector.
            # as they are only unique to a minus sign.
            assert np.linalg.norm(np.abs(lowercase_      ) - np.abs(lowercase_      )      ) <= 1e-6


if __name__ == "__main__":
  import doctest

  doctest.testmod()
  test_power_iteration()
