





import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    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 PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




@skip_mps
class  __UpperCamelCase						(       __SCREAMING_SNAKE_CASE  ,   __SCREAMING_SNAKE_CASE  ,   __SCREAMING_SNAKE_CASE  ,   unittest.TestCase						):
     __A  :     Any            =      StableDiffusionModelEditingPipeline
     __A  :     Dict            =      TEXT_TO_IMAGE_PARAMS
     __A  :     Dict            =      TEXT_TO_IMAGE_BATCH_PARAMS
     __A  :     Optional[int]            =      TEXT_TO_IMAGE_IMAGE_PARAMS
     __A  :     Optional[int]            =      TEXT_TO_IMAGE_IMAGE_PARAMS





     def 		UpperCamelCase(    self				):
        torch.manual_seed(0				)
        _UpperCAmelCase											=      UNetaDConditionModel(
            block_out_channels=(32, 64)							,       layers_per_block=2							,       sample_size=32							,       in_channels=4							,       out_channels=4							,       down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')							,       up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')							,       cross_attention_dim=32							,       )
        _UpperCAmelCase											=      DDIMScheduler()
        torch.manual_seed(0				)
        _UpperCAmelCase											=      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				)
        _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							,       )
        _UpperCAmelCase											=      CLIPTextModel(UpperCamelCase__				)
        _UpperCAmelCase											=      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)

        _UpperCAmelCase											=      {
            '''unet''': unet,
            '''scheduler''': scheduler,
            '''vae''': vae,
            '''text_encoder''': text_encoder,
            '''tokenizer''': tokenizer,
            '''safety_checker''': None,
            '''feature_extractor''': None,
        }
        return components





     def 		UpperCamelCase(    self							,       _UpperCamelCase							,       _UpperCamelCase=0				):
        _UpperCAmelCase											=      torch.manual_seed(UpperCamelCase__				)
        _UpperCAmelCase											=      {
            '''prompt''': '''A field of roses''',
            '''generator''': generator,
            # Setting height and width to None to prevent OOMs on CPU.
            '''height''': None,
            '''width''': None,
            '''num_inference_steps''': 2,
            '''guidance_scale''': 6.0,
            '''output_type''': '''numpy''',
        }
        return inputs





     def 		UpperCamelCase(    self				):
        _UpperCAmelCase											=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        _UpperCAmelCase											=      self.get_dummy_components()
        _UpperCAmelCase											=      StableDiffusionModelEditingPipeline(**UpperCamelCase__				)
        _UpperCAmelCase											=      sd_pipe.to(UpperCamelCase__				)
        sd_pipe.set_progress_bar_config(disable=UpperCamelCase__				)

        _UpperCAmelCase											=      self.get_dummy_inputs(UpperCamelCase__				)
        _UpperCAmelCase											=      sd_pipe(**UpperCamelCase__				).images
        _UpperCAmelCase											=      image[0, -3:, -3:, -1]
        assert image.shape == (1, 64, 64, 3)

        _UpperCAmelCase											=      np.array([0.4755, 0.5132, 0.4976, 0.3904, 0.3554, 0.4765, 0.5139, 0.5158, 0.4889]				)

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





     def 		UpperCamelCase(    self				):
        _UpperCAmelCase											=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        _UpperCAmelCase											=      self.get_dummy_components()
        _UpperCAmelCase											=      StableDiffusionModelEditingPipeline(**UpperCamelCase__				)
        _UpperCAmelCase											=      sd_pipe.to(UpperCamelCase__				)
        sd_pipe.set_progress_bar_config(disable=UpperCamelCase__				)

        _UpperCAmelCase											=      self.get_dummy_inputs(UpperCamelCase__				)
        _UpperCAmelCase											=      '''french fries'''
        _UpperCAmelCase											=      sd_pipe(**UpperCamelCase__							,       negative_prompt=UpperCamelCase__				)
        _UpperCAmelCase											=      output.images
        _UpperCAmelCase											=      image[0, -3:, -3:, -1]

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

        _UpperCAmelCase											=      np.array([0.4992, 0.5101, 0.5004, 0.3949, 0.3604, 0.4735, 0.5216, 0.5204, 0.4913]				)

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





     def 		UpperCamelCase(    self				):
        _UpperCAmelCase											=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        _UpperCAmelCase											=      self.get_dummy_components()
        _UpperCAmelCase											=      EulerAncestralDiscreteScheduler(
            beta_start=0.00085							,       beta_end=0.012							,       beta_schedule='''scaled_linear'''				)
        _UpperCAmelCase											=      StableDiffusionModelEditingPipeline(**UpperCamelCase__				)
        _UpperCAmelCase											=      sd_pipe.to(UpperCamelCase__				)
        sd_pipe.set_progress_bar_config(disable=UpperCamelCase__				)

        _UpperCAmelCase											=      self.get_dummy_inputs(UpperCamelCase__				)
        _UpperCAmelCase											=      sd_pipe(**UpperCamelCase__				).images
        _UpperCAmelCase											=      image[0, -3:, -3:, -1]

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

        _UpperCAmelCase											=      np.array([0.4747, 0.5372, 0.4779, 0.4982, 0.5543, 0.4816, 0.5238, 0.4904, 0.5027]				)

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





     def 		UpperCamelCase(    self				):
        _UpperCAmelCase											=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        _UpperCAmelCase											=      self.get_dummy_components()
        _UpperCAmelCase											=      PNDMScheduler()
        _UpperCAmelCase											=      StableDiffusionModelEditingPipeline(**UpperCamelCase__				)
        _UpperCAmelCase											=      sd_pipe.to(UpperCamelCase__				)
        sd_pipe.set_progress_bar_config(disable=UpperCamelCase__				)

        _UpperCAmelCase											=      self.get_dummy_inputs(UpperCamelCase__				)
        # the pipeline does not expect pndm so test if it raises error.
        with self.assertRaises(UpperCamelCase__				):
           _UpperCAmelCase											=      sd_pipe(**UpperCamelCase__				).images





     def 		UpperCamelCase(    self				):
        super().test_inference_batch_single_identical(expected_max_diff=5e-3				)





     def 		UpperCamelCase(    self				):
        super().test_attention_slicing_forward_pass(expected_max_diff=5e-3				)




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





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





     def 		UpperCamelCase(    self							,       _UpperCamelCase=0				):
        _UpperCAmelCase											=      torch.manual_seed(UpperCamelCase__				)
        _UpperCAmelCase											=      {
            '''prompt''': '''A field of roses''',
            '''generator''': generator,
            '''num_inference_steps''': 3,
            '''guidance_scale''': 7.5,
            '''output_type''': '''numpy''',
        }
        return inputs





     def 		UpperCamelCase(    self				):
        _UpperCAmelCase											=      '''CompVis/stable-diffusion-v1-4'''
        _UpperCAmelCase											=      StableDiffusionModelEditingPipeline.from_pretrained(UpperCamelCase__							,       safety_checker=UpperCamelCase__				)
        pipe.to(UpperCamelCase__				)
        pipe.set_progress_bar_config(disable=UpperCamelCase__				)
        pipe.enable_attention_slicing()

        _UpperCAmelCase											=      self.get_inputs()
        _UpperCAmelCase											=      pipe(**UpperCamelCase__				).images
        _UpperCAmelCase											=      image[0, -3:, -3:, -1].flatten()

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

        _UpperCAmelCase											=      np.array(
            [0.6749496, 0.6386453, 0.51443267, 0.66094905, 0.61921215, 0.5491332, 0.5744417, 0.58075106, 0.5174658]				)

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

        # make sure image changes after editing
        pipe.edit_model('''A pack of roses'''							,       '''A pack of blue roses'''				)

        _UpperCAmelCase											=      pipe(**UpperCamelCase__				).images
        _UpperCAmelCase											=      image[0, -3:, -3:, -1].flatten()

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

        assert np.abs(expected_slice - image_slice				).max() > 1e-1





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

        _UpperCAmelCase											=      '''CompVis/stable-diffusion-v1-4'''
        _UpperCAmelCase											=      DDIMScheduler.from_pretrained(UpperCamelCase__							,       subfolder='''scheduler'''				)
        _UpperCAmelCase											=      StableDiffusionModelEditingPipeline.from_pretrained(
            UpperCamelCase__							,       scheduler=UpperCamelCase__							,       safety_checker=UpperCamelCase__				)
        _UpperCAmelCase											=      pipe.to(UpperCamelCase__				)
        pipe.set_progress_bar_config(disable=UpperCamelCase__				)
        pipe.enable_attention_slicing(1				)
        pipe.enable_sequential_cpu_offload()

        _UpperCAmelCase											=      self.get_inputs()
        _UpperCAmelCase											=      pipe(**UpperCamelCase__				)

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






def     A__					(						SCREAMING_SNAKE_CASE_      :			str							,  SCREAMING_SNAKE_CASE_      :			str	)     ->     str:






   """simple docstring"""







   if not isinstance(SCREAMING_SNAKE_CASE_							,  SCREAMING_SNAKE_CASE_	):
      raise TypeError('''plaintext must be a string'''	)
   if not isinstance(SCREAMING_SNAKE_CASE_							,  SCREAMING_SNAKE_CASE_	):
      raise TypeError('''key must be a string'''	)

   if not plaintext:
      raise ValueError('''plaintext is empty'''	)
   if not key:
      raise ValueError('''key is empty'''	)

   key += plaintext
   _UpperCAmelCase											=      plaintext.lower()
   _UpperCAmelCase											=      key.lower()
   _UpperCAmelCase											=      0
   _UpperCAmelCase											=      0
   _UpperCAmelCase											=      ''''''
   while plaintext_iterator < len(SCREAMING_SNAKE_CASE_	):
      if (
          ord(plaintext[plaintext_iterator]	) < 97
          or ord(plaintext[plaintext_iterator]	) > 1_22
      ):
         ciphertext += plaintext[plaintext_iterator]
         plaintext_iterator += 1
      elif ord(key[key_iterator]	) < 97 or ord(key[key_iterator]	) > 1_22:
         key_iterator += 1
      else:
         ciphertext += chr(
             (
                 (ord(plaintext[plaintext_iterator]	) - 97 + ord(key[key_iterator]	))
                 - 97
             )
             % 26
             + 97	)
         key_iterator += 1
         plaintext_iterator += 1
   return ciphertext
def     A__					(						SCREAMING_SNAKE_CASE_      :			str							,  SCREAMING_SNAKE_CASE_      :			str	)     ->     str:






   """simple docstring"""







   if not isinstance(SCREAMING_SNAKE_CASE_							,  SCREAMING_SNAKE_CASE_	):
      raise TypeError('''ciphertext must be a string'''	)
   if not isinstance(SCREAMING_SNAKE_CASE_							,  SCREAMING_SNAKE_CASE_	):
      raise TypeError('''key must be a string'''	)

   if not ciphertext:
      raise ValueError('''ciphertext is empty'''	)
   if not key:
      raise ValueError('''key is empty'''	)

   _UpperCAmelCase											=      key.lower()
   _UpperCAmelCase											=      0
   _UpperCAmelCase											=      0
   _UpperCAmelCase											=      ''''''
   while ciphertext_iterator < len(SCREAMING_SNAKE_CASE_	):
      if (
          ord(ciphertext[ciphertext_iterator]	) < 97
          or ord(ciphertext[ciphertext_iterator]	) > 1_22
      ):
         plaintext += ciphertext[ciphertext_iterator]
      else:
         plaintext += chr(
             (ord(ciphertext[ciphertext_iterator]	) - ord(key[key_iterator]	)) % 26
             + 97	)
         key += chr(
             (ord(ciphertext[ciphertext_iterator]	) - ord(key[key_iterator]	)) % 26
             + 97	)
         key_iterator += 1
      ciphertext_iterator += 1
   return plaintext


if __name__ == "__main__":
   import doctest

   doctest.testmod()
   UpperCAmelCase_											=  int(input("Type 1 to encrypt or 2 to decrypt:"))
   if operation == 1:
      UpperCAmelCase_											=  input("Typeplaintext to be encrypted:\n")
      UpperCAmelCase_											=  input("Type the key:\n")
      print(encrypt(plaintext, key))
   elif operation == 2:
      UpperCAmelCase_											=  input("Type the ciphertext to be decrypted:\n")
      UpperCAmelCase_											=  input("Type the key:\n")
      print(decrypt(ciphertext, key))
   decrypt("jsqqs avvwo", "coffee")