





import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np

class 		UpperCamelCase__ :



     '''simple docstring'''

     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[Any]       )     ->		Tuple:




           '''simple docstring'''


           torch.manual_seed(0       )
           SCREAMING_SNAKE_CASE		=						TaEncoderModel.from_pretrained("""hf-internal-testing/tiny-random-t5"""       )

           torch.manual_seed(0       )
           SCREAMING_SNAKE_CASE		=						AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-t5"""       )

           torch.manual_seed(0       )
           SCREAMING_SNAKE_CASE		=						UNetaDConditionModel(
               sample_size=32						,layers_per_block=1						,block_out_channels=[32, 64]						,down_block_types=[
                   """ResnetDownsampleBlock2D""",
                   """SimpleCrossAttnDownBlock2D""",
               ]						,mid_block_type="""UNetMidBlock2DSimpleCrossAttn"""						,up_block_types=["""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""]						,in_channels=3						,out_channels=6						,cross_attention_dim=32						,encoder_hid_dim=32						,attention_head_dim=8						,addition_embed_type="""text"""						,addition_embed_type_num_heads=2						,cross_attention_norm="""group_norm"""						,resnet_time_scale_shift="""scale_shift"""						,act_fn="""gelu"""						,)
           unet.set_attn_processor(AttnAddedKVProcessor()       )  # For reproducibility tests

           torch.manual_seed(0       )
           SCREAMING_SNAKE_CASE		=						DDPMScheduler(
               num_train_timesteps=1000						,beta_schedule="""squaredcos_cap_v2"""						,beta_start=0.0001						,beta_end=0.02						,thresholding=lowerCamelCase__						,dynamic_thresholding_ratio=0.95						,sample_max_value=1.0						,prediction_type="""epsilon"""						,variance_type="""learned_range"""						,)

           torch.manual_seed(0       )
           SCREAMING_SNAKE_CASE		=						IFWatermarker()

           return {
               "text_encoder": text_encoder,
               "tokenizer": tokenizer,
               "unet": unet,
               "scheduler": scheduler,
               "watermarker": watermarker,
               "safety_checker": None,
               "feature_extractor": None,
           }






     def 				SCREAMING_SNAKE_CASE__  (							self						:    List[Any]       )     ->		Any:




           '''simple docstring'''


           torch.manual_seed(0       )
           SCREAMING_SNAKE_CASE		=						TaEncoderModel.from_pretrained("""hf-internal-testing/tiny-random-t5"""       )

           torch.manual_seed(0       )
           SCREAMING_SNAKE_CASE		=						AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-t5"""       )

           torch.manual_seed(0       )
           SCREAMING_SNAKE_CASE		=						UNetaDConditionModel(
               sample_size=32						,layers_per_block=[1, 2]						,block_out_channels=[32, 64]						,down_block_types=[
                   """ResnetDownsampleBlock2D""",
                   """SimpleCrossAttnDownBlock2D""",
               ]						,mid_block_type="""UNetMidBlock2DSimpleCrossAttn"""						,up_block_types=["""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""]						,in_channels=6						,out_channels=6						,cross_attention_dim=32						,encoder_hid_dim=32						,attention_head_dim=8						,addition_embed_type="""text"""						,addition_embed_type_num_heads=2						,cross_attention_norm="""group_norm"""						,resnet_time_scale_shift="""scale_shift"""						,act_fn="""gelu"""						,class_embed_type="""timestep"""						,mid_block_scale_factor=1.414						,time_embedding_act_fn="""gelu"""						,time_embedding_dim=32						,)
           unet.set_attn_processor(AttnAddedKVProcessor()       )  # For reproducibility tests

           torch.manual_seed(0       )
           SCREAMING_SNAKE_CASE		=						DDPMScheduler(
               num_train_timesteps=1000						,beta_schedule="""squaredcos_cap_v2"""						,beta_start=0.0001						,beta_end=0.02						,thresholding=lowerCamelCase__						,dynamic_thresholding_ratio=0.95						,sample_max_value=1.0						,prediction_type="""epsilon"""						,variance_type="""learned_range"""						,)

           torch.manual_seed(0       )
           SCREAMING_SNAKE_CASE		=						DDPMScheduler(
               num_train_timesteps=1000						,beta_schedule="""squaredcos_cap_v2"""						,beta_start=0.0001						,beta_end=0.02						,)

           torch.manual_seed(0       )
           SCREAMING_SNAKE_CASE		=						IFWatermarker()

           return {
               "text_encoder": text_encoder,
               "tokenizer": tokenizer,
               "unet": unet,
               "scheduler": scheduler,
               "image_noising_scheduler": image_noising_scheduler,
               "watermarker": watermarker,
               "safety_checker": None,
               "feature_extractor": None,
           }






     def 				SCREAMING_SNAKE_CASE__  (							self						:    List[str]       )     ->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.get_dummy_components()
           SCREAMING_SNAKE_CASE		=						self.pipeline_class(**lowerCamelCase__       )
           pipe.to(lowerCamelCase__       )
           pipe.set_progress_bar_config(disable=lowerCamelCase__       )

           SCREAMING_SNAKE_CASE		=						self.get_dummy_inputs(lowerCamelCase__       )

           SCREAMING_SNAKE_CASE		=						inputs["""prompt"""]
           SCREAMING_SNAKE_CASE		=						inputs["""generator"""]
           SCREAMING_SNAKE_CASE		=						inputs["""num_inference_steps"""]
           SCREAMING_SNAKE_CASE		=						inputs["""output_type"""]

           if "image" in inputs:
                 SCREAMING_SNAKE_CASE		=						inputs["""image"""]
           else:
                 SCREAMING_SNAKE_CASE		=						None

           if "mask_image" in inputs:
                 SCREAMING_SNAKE_CASE		=						inputs["""mask_image"""]
           else:
                 SCREAMING_SNAKE_CASE		=						None

           if "original_image" in inputs:
                 SCREAMING_SNAKE_CASE		=						inputs["""original_image"""]
           else:
                 SCREAMING_SNAKE_CASE		=						None

           SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE		=						pipe.encode_prompt(lowerCamelCase__       )

           # inputs with prompt converted to embeddings
           SCREAMING_SNAKE_CASE		=						{
               """prompt_embeds""": prompt_embeds,
               """negative_prompt_embeds""": negative_prompt_embeds,
               """generator""": generator,
               """num_inference_steps""": num_inference_steps,
               """output_type""": output_type,
           }

           if image is not None:
                 SCREAMING_SNAKE_CASE		=						image

           if mask_image is not None:
                 SCREAMING_SNAKE_CASE		=						mask_image

           if original_image is not None:
                 SCREAMING_SNAKE_CASE		=						original_image

           # set all optional components to None
           for optional_component in pipe._optional_components:
                 setattr(lowerCamelCase__						,lowerCamelCase__						,lowerCamelCase__       )

           SCREAMING_SNAKE_CASE		=						pipe(**lowerCamelCase__       )[0]

           with tempfile.TemporaryDirectory() as tmpdir:
                 pipe.save_pretrained(lowerCamelCase__       )
                 SCREAMING_SNAKE_CASE		=						self.pipeline_class.from_pretrained(lowerCamelCase__       )
                 pipe_loaded.to(lowerCamelCase__       )
                 pipe_loaded.set_progress_bar_config(disable=lowerCamelCase__       )

           pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()       )  # For reproducibility tests

           for optional_component in pipe._optional_components:
                 self.assertTrue(
                     getattr(lowerCamelCase__						,lowerCamelCase__       ) is None						,F"""`{optional_component}` did not stay set to None after loading."""						,)

           SCREAMING_SNAKE_CASE		=						self.get_dummy_inputs(lowerCamelCase__       )

           SCREAMING_SNAKE_CASE		=						inputs["""generator"""]
           SCREAMING_SNAKE_CASE		=						inputs["""num_inference_steps"""]
           SCREAMING_SNAKE_CASE		=						inputs["""output_type"""]

           # inputs with prompt converted to embeddings
           SCREAMING_SNAKE_CASE		=						{
               """prompt_embeds""": prompt_embeds,
               """negative_prompt_embeds""": negative_prompt_embeds,
               """generator""": generator,
               """num_inference_steps""": num_inference_steps,
               """output_type""": output_type,
           }

           if image is not None:
                 SCREAMING_SNAKE_CASE		=						image

           if mask_image is not None:
                 SCREAMING_SNAKE_CASE		=						mask_image

           if original_image is not None:
                 SCREAMING_SNAKE_CASE		=						original_image

           SCREAMING_SNAKE_CASE		=						pipe_loaded(**lowerCamelCase__       )[0]

           SCREAMING_SNAKE_CASE		=						np.abs(to_np(lowerCamelCase__       ) - to_np(lowerCamelCase__       )       ).max()
           self.assertLess(lowerCamelCase__						,1e-4       )






     def 				SCREAMING_SNAKE_CASE__  (							self						:    Dict       )     ->		Optional[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.get_dummy_components()
           SCREAMING_SNAKE_CASE		=						self.pipeline_class(**lowerCamelCase__       )
           pipe.to(lowerCamelCase__       )
           pipe.set_progress_bar_config(disable=lowerCamelCase__       )

           SCREAMING_SNAKE_CASE		=						self.get_dummy_inputs(lowerCamelCase__       )
           SCREAMING_SNAKE_CASE		=						pipe(**lowerCamelCase__       )[0]

           with tempfile.TemporaryDirectory() as tmpdir:
                 pipe.save_pretrained(lowerCamelCase__       )
                 SCREAMING_SNAKE_CASE		=						self.pipeline_class.from_pretrained(lowerCamelCase__       )
                 pipe_loaded.to(lowerCamelCase__       )
                 pipe_loaded.set_progress_bar_config(disable=lowerCamelCase__       )

           pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()       )  # For reproducibility tests

           SCREAMING_SNAKE_CASE		=						self.get_dummy_inputs(lowerCamelCase__       )
           SCREAMING_SNAKE_CASE		=						pipe_loaded(**lowerCamelCase__       )[0]

           SCREAMING_SNAKE_CASE		=						np.abs(to_np(lowerCamelCase__       ) - to_np(lowerCamelCase__       )       ).max()
           self.assertLess(lowerCamelCase__						,1e-4       )









import math
import unittest



def    __lowercase   (      _SCREAMING_SNAKE_CASE	) ->   bool:


      '''simple docstring'''


      assert isinstance(_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE	) and (
          number >= 0
      ), "'number' must been an int and positive"

      if 1 < number < 4:
            # 2 and 3 are primes
            return True
      elif number < 2 or number % 2 == 0 or number % 3 == 0:
            # Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
            return False

      # All primes number are in format of 6k +/- 1
      for i in range(5 ,						int(math.sqrt(_SCREAMING_SNAKE_CASE	) + 1	) ,						6	):
            if number % i == 0 or number % (i + 2) == 0:
                  return False
      return True

class 		UpperCamelCase__ (      unittest.TestCase   ):



     '''simple docstring'''

     def 				SCREAMING_SNAKE_CASE__  (							self						:    Dict       )     ->		Tuple:




           '''simple docstring'''


           self.assertTrue(is_prime(2       )       )
           self.assertTrue(is_prime(3       )       )
           self.assertTrue(is_prime(5       )       )
           self.assertTrue(is_prime(7       )       )
           self.assertTrue(is_prime(11       )       )
           self.assertTrue(is_prime(13       )       )
           self.assertTrue(is_prime(17       )       )
           self.assertTrue(is_prime(19       )       )
           self.assertTrue(is_prime(23       )       )
           self.assertTrue(is_prime(29       )       )






     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[int]       )     ->		Tuple:




           '''simple docstring'''


           with self.assertRaises(lowerCamelCase__       ):
                 is_prime(-19       )
           self.assertFalse(
               is_prime(0       )						,"""Zero doesn't have any positive factors, primes must have exactly two."""						,)
           self.assertFalse(
               is_prime(1       )						,"""One only has 1 positive factor, primes must have exactly two."""						,)
           self.assertFalse(is_prime(2 * 2       )       )
           self.assertFalse(is_prime(2 * 3       )       )
           self.assertFalse(is_prime(3 * 3       )       )
           self.assertFalse(is_prime(3 * 5       )       )
           self.assertFalse(is_prime(3 * 5 * 7       )       )


if __name__ == "__main__":
      unittest.main()


