
"""simple docstring"""







from __future__ import annotations

import sys




def 			A_			(					lowercase						, lowercase      )     ->							tuple[int, int]:


     """simple docstring"""







     if abs(lowercase      ) == 1:
          return a, 0
     elif abs(lowercase      ) == 1:
          return 0, b

     UpperCAmelCase_					,UpperCAmelCase_						:    Union[str, Any]    					=  a, b
     UpperCAmelCase_					,UpperCAmelCase_						:    List[Any]    					=  1, 0
     UpperCAmelCase_					,UpperCAmelCase_						:    int    					=  0, 1

     while remainder != 0:
          UpperCAmelCase_						:    Tuple    					=  old_remainder // remainder
          UpperCAmelCase_					,UpperCAmelCase_						:    int    					=  remainder, old_remainder - quotient * remainder
          UpperCAmelCase_					,UpperCAmelCase_						:    List[str]    					=  coeff_a, old_coeff_a - quotient * coeff_a
          UpperCAmelCase_					,UpperCAmelCase_						:    int    					=  coeff_b, old_coeff_b - quotient * coeff_b

     # sign correction for negative numbers
     if a < 0:
          UpperCAmelCase_						:    Any    					=  -old_coeff_a
     if b < 0:
          UpperCAmelCase_						:    Union[str, Any]    					=  -old_coeff_b

     return old_coeff_a, old_coeff_b




def 			A_			(					)     ->							Dict:


     """simple docstring"""







     if len(sys.argv      ) < 3:
          print("""2 integer arguments required"""      )
          return 1
     UpperCAmelCase_						:    Tuple    					=  int(sys.argv[1]      )
     UpperCAmelCase_						:    Union[str, Any]    					=  int(sys.argv[2]      )
     print(extended_euclidean_algorithm(lowercase						, lowercase      )      )
     return 0


if __name__ == "__main__":
       raise SystemExit(main())



"""simple docstring"""







import gc
import unittest

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

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, 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_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()



class 		UpperCAmelCase_	(lowerCamelCase_  , unittest.TestCase						):
 """simple docstring"""




 UpperCamelCase_						:     List[str]				   =    UnCLIPPipeline
 UpperCamelCase_						:     Any				   =    TEXT_TO_IMAGE_PARAMS - {
     """negative_prompt""",
     """height""",
     """width""",
     """negative_prompt_embeds""",
     """guidance_scale""",
     """prompt_embeds""",
     """cross_attention_kwargs""",
 }
 UpperCamelCase_						:     int				   =    TEXT_TO_IMAGE_BATCH_PARAMS
 UpperCamelCase_						:     int				   =    [
     """generator""",
     """return_dict""",
     """prior_num_inference_steps""",
     """decoder_num_inference_steps""",
     """super_res_num_inference_steps""",
 ]
 UpperCamelCase_						:     Union[str, Any]				   =    False




 @property
 def   a      ( self       : Tuple )->					str:



      """simple docstring"""

      return 32




 @property
 def   a      ( self       : Tuple )->					Any:



      """simple docstring"""

      return 32




 @property
 def   a      ( self       : Union[str, Any] )->					str:



      """simple docstring"""

      return self.time_input_dim




 @property
 def   a      ( self       : int )->					Dict:



      """simple docstring"""

      return self.time_input_dim * 4




 @property
 def   a      ( self       : List[str] )->					Tuple:



      """simple docstring"""

      return 1_00




 @property
 def   a      ( self       : Tuple )->					Optional[Any]:



      """simple docstring"""

      UpperCAmelCase_						:    List[Any]    					=  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""" )
      return tokenizer




 @property
 def   a      ( self       : Optional[Any] )->					Dict:



      """simple docstring"""

      torch.manual_seed(0 )
      UpperCAmelCase_						:    Any    					=  CLIPTextConfig(
          bos_token_id=0     ,   eos_token_id=2     ,   hidden_size=self.text_embedder_hidden_size     ,   projection_dim=self.text_embedder_hidden_size     ,   intermediate_size=37     ,   layer_norm_eps=1E-05     ,   num_attention_heads=4     ,   num_hidden_layers=5     ,   pad_token_id=1     ,   vocab_size=10_00     ,   )
      return CLIPTextModelWithProjection(a_ )




 @property
 def   a      ( self       : str )->					Optional[Any]:



      """simple docstring"""

      torch.manual_seed(0 )

      UpperCAmelCase_						:    Optional[Any]    					=  {
          """num_attention_heads""": 2,
          """attention_head_dim""": 12,
          """embedding_dim""": self.text_embedder_hidden_size,
          """num_layers""": 1,
      }

      UpperCAmelCase_						:    Tuple    					=  PriorTransformer(**a_ )
      return model




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



      """simple docstring"""

      torch.manual_seed(0 )

      UpperCAmelCase_						:    Optional[int]    					=  {
          """clip_embeddings_dim""": self.text_embedder_hidden_size,
          """time_embed_dim""": self.time_embed_dim,
          """cross_attention_dim""": self.cross_attention_dim,
      }

      UpperCAmelCase_						:    Tuple    					=  UnCLIPTextProjModel(**a_ )
      return model




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



      """simple docstring"""

      torch.manual_seed(0 )

      UpperCAmelCase_						:    Tuple    					=  {
          """sample_size""": 32,
          # RGB in channels
          """in_channels""": 3,
          # Out channels is double in channels because predicts mean and variance
          """out_channels""": 6,
          """down_block_types""": ("""ResnetDownsampleBlock2D""", """SimpleCrossAttnDownBlock2D"""),
          """up_block_types""": ("""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""),
          """mid_block_type""": """UNetMidBlock2DSimpleCrossAttn""",
          """block_out_channels""": (self.block_out_channels_a, self.block_out_channels_a * 2),
          """layers_per_block""": 1,
          """cross_attention_dim""": self.cross_attention_dim,
          """attention_head_dim""": 4,
          """resnet_time_scale_shift""": """scale_shift""",
          """class_embed_type""": """identity""",
      }

      UpperCAmelCase_						:    List[str]    					=  UNetaDConditionModel(**a_ )
      return model




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



      """simple docstring"""

      return {
          "sample_size": 64,
          "layers_per_block": 1,
          "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
          "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
          "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
          "in_channels": 6,
          "out_channels": 3,
      }




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



      """simple docstring"""

      torch.manual_seed(0 )

      UpperCAmelCase_						:    Tuple    					=  UNetaDModel(**self.dummy_super_res_kwargs )
      return model




 @property
 def   a      ( self       : Union[str, Any] )->					Union[str, Any]:



      """simple docstring"""

      torch.manual_seed(1 )

      UpperCAmelCase_						:    Dict    					=  UNetaDModel(**self.dummy_super_res_kwargs )
      return model




 def   a      ( self       : Tuple )->					str:



      """simple docstring"""

      UpperCAmelCase_						:    List[str]    					=  self.dummy_prior
      UpperCAmelCase_						:    Any    					=  self.dummy_decoder
      UpperCAmelCase_						:    Union[str, Any]    					=  self.dummy_text_proj
      UpperCAmelCase_						:    int    					=  self.dummy_text_encoder
      UpperCAmelCase_						:    Union[str, Any]    					=  self.dummy_tokenizer
      UpperCAmelCase_						:    Tuple    					=  self.dummy_super_res_first
      UpperCAmelCase_						:    List[str]    					=  self.dummy_super_res_last

      UpperCAmelCase_						:    Dict    					=  UnCLIPScheduler(
          variance_type="""fixed_small_log"""     ,   prediction_type="""sample"""     ,   num_train_timesteps=10_00     ,   clip_sample_range=5.0     ,   )

      UpperCAmelCase_						:    List[str]    					=  UnCLIPScheduler(
          variance_type="""learned_range"""     ,   prediction_type="""epsilon"""     ,   num_train_timesteps=10_00     ,   )

      UpperCAmelCase_						:    List[str]    					=  UnCLIPScheduler(
          variance_type="""fixed_small_log"""     ,   prediction_type="""epsilon"""     ,   num_train_timesteps=10_00     ,   )

      UpperCAmelCase_						:    List[str]    					=  {
          """prior""": prior,
          """decoder""": decoder,
          """text_proj""": text_proj,
          """text_encoder""": text_encoder,
          """tokenizer""": tokenizer,
          """super_res_first""": super_res_first,
          """super_res_last""": super_res_last,
          """prior_scheduler""": prior_scheduler,
          """decoder_scheduler""": decoder_scheduler,
          """super_res_scheduler""": super_res_scheduler,
      }

      return components




 def   a      ( self       : Union[str, Any]     ,   a_       : Optional[Any]     ,   a_       : Dict=0 )->					Union[str, Any]:



      """simple docstring"""

      if str(a_ ).startswith("""mps""" ):
           UpperCAmelCase_						:    Union[str, Any]    					=  torch.manual_seed(a_ )
      else:
           UpperCAmelCase_						:    List[str]    					=  torch.Generator(device=a_ ).manual_seed(a_ )
      UpperCAmelCase_						:    Dict    					=  {
          """prompt""": """horse""",
          """generator""": generator,
          """prior_num_inference_steps""": 2,
          """decoder_num_inference_steps""": 2,
          """super_res_num_inference_steps""": 2,
          """output_type""": """numpy""",
      }
      return inputs




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



      """simple docstring"""

      UpperCAmelCase_						:    List[Any]    					=  """cpu"""

      UpperCAmelCase_						:    int    					=  self.get_dummy_components()

      UpperCAmelCase_						:    str    					=  self.pipeline_class(**a_ )
      UpperCAmelCase_						:    Any    					=  pipe.to(a_ )

      pipe.set_progress_bar_config(disable=a_ )

      UpperCAmelCase_						:    int    					=  pipe(**self.get_dummy_inputs(a_ ) )
      UpperCAmelCase_						:    List[str]    					=  output.images

      UpperCAmelCase_						:    Optional[int]    					=  pipe(
          **self.get_dummy_inputs(a_ )     ,   return_dict=a_     ,   )[0]

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

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

      UpperCAmelCase_						:    Any    					=  np.array(
          [
              0.9997,
              0.9988,
              0.0028,
              0.9997,
              0.9984,
              0.9965,
              0.0029,
              0.9986,
              0.0025,
          ] )

      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   a      ( self       : int )->					str:



      """simple docstring"""

      UpperCAmelCase_						:    str    					=  torch.device("""cpu""" )



      class 		UpperCAmelCase_	:
       """simple docstring"""




       UpperCamelCase_						:     Optional[Any]				   =    1

      UpperCAmelCase_						:    List[str]    					=  self.get_dummy_components()

      UpperCAmelCase_						:    str    					=  self.pipeline_class(**a_ )
      UpperCAmelCase_						:    Any    					=  pipe.to(a_ )

      UpperCAmelCase_						:    List[Any]    					=  components["""prior"""]
      UpperCAmelCase_						:    str    					=  components["""decoder"""]
      UpperCAmelCase_						:    str    					=  components["""super_res_first"""]
      UpperCAmelCase_						:    Optional[int]    					=  components["""tokenizer"""]
      UpperCAmelCase_						:    List[str]    					=  components["""text_encoder"""]

      UpperCAmelCase_						:    Any    					=  torch.Generator(device=a_ ).manual_seed(0 )
      UpperCAmelCase_						:    Optional[Any]    					=  prior.dtype
      UpperCAmelCase_						:    List[Any]    					=  1

      UpperCAmelCase_						:    int    					=  (batch_size, prior.config.embedding_dim)
      UpperCAmelCase_						:    Optional[Any]    					=  pipe.prepare_latents(
          a_     ,   dtype=a_     ,   device=a_     ,   generator=a_     ,   latents=a_     ,   scheduler=DummyScheduler() )
      UpperCAmelCase_						:    List[str]    					=  (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
      UpperCAmelCase_						:    int    					=  pipe.prepare_latents(
          a_     ,   dtype=a_     ,   device=a_     ,   generator=a_     ,   latents=a_     ,   scheduler=DummyScheduler() )

      UpperCAmelCase_						:    List[str]    					=  (
          batch_size,
          super_res_first.config.in_channels // 2,
          super_res_first.config.sample_size,
          super_res_first.config.sample_size,
      )
      UpperCAmelCase_						:    int    					=  pipe.prepare_latents(
          a_     ,   dtype=a_     ,   device=a_     ,   generator=a_     ,   latents=a_     ,   scheduler=DummyScheduler() )

      pipe.set_progress_bar_config(disable=a_ )

      UpperCAmelCase_						:    str    					=  """this is a prompt example"""

      UpperCAmelCase_						:    Any    					=  torch.Generator(device=a_ ).manual_seed(0 )
      UpperCAmelCase_						:    int    					=  pipe(
          [prompt]     ,   generator=a_     ,   prior_num_inference_steps=2     ,   decoder_num_inference_steps=2     ,   super_res_num_inference_steps=2     ,   prior_latents=a_     ,   decoder_latents=a_     ,   super_res_latents=a_     ,   output_type="""np"""     ,   )
      UpperCAmelCase_						:    Tuple    					=  output.images

      UpperCAmelCase_						:    int    					=  tokenizer(
          a_     ,   padding="""max_length"""     ,   max_length=tokenizer.model_max_length     ,   return_tensors="""pt"""     ,   )
      UpperCAmelCase_						:    Union[str, Any]    					=  text_encoder(text_inputs.input_ids )
      UpperCAmelCase_						:    Optional[Any]    					=  text_inputs.attention_mask

      UpperCAmelCase_						:    List[str]    					=  torch.Generator(device=a_ ).manual_seed(0 )
      UpperCAmelCase_						:    List[str]    					=  pipe(
          generator=a_     ,   prior_num_inference_steps=2     ,   decoder_num_inference_steps=2     ,   super_res_num_inference_steps=2     ,   prior_latents=a_     ,   decoder_latents=a_     ,   super_res_latents=a_     ,   text_model_output=a_     ,   text_attention_mask=a_     ,   output_type="""np"""     ,   )[0]

      # make sure passing text embeddings manually is identical
      assert np.abs(image - image_from_text ).max() < 1E-4




 @skip_mps
 def   a      ( self       : Tuple )->					Tuple:



      """simple docstring"""

      UpperCAmelCase_						:    int    					=  torch_device == """cpu"""

      self._test_attention_slicing_forward_pass(test_max_difference=a_     ,   expected_max_diff=0.01 )




 @skip_mps
 def   a      ( self       : Dict )->					Union[str, Any]:



      """simple docstring"""

      UpperCAmelCase_						:    Optional[Any]    					=  torch_device == """cpu"""
      UpperCAmelCase_						:    List[str]    					=  True
      UpperCAmelCase_						:    Any    					=  [
          """prior_num_inference_steps""",
          """decoder_num_inference_steps""",
          """super_res_num_inference_steps""",
      ]

      self._test_inference_batch_single_identical(
          test_max_difference=a_     ,   relax_max_difference=a_     ,   additional_params_copy_to_batched_inputs=a_     ,   )




 def   a      ( self       : Any )->					Optional[Any]:



      """simple docstring"""

      UpperCAmelCase_						:    Union[str, Any]    					=  [
          """prior_num_inference_steps""",
          """decoder_num_inference_steps""",
          """super_res_num_inference_steps""",
      ]

      if torch_device == "mps":
           # TODO: MPS errors with larger batch sizes
           UpperCAmelCase_						:    List[Any]    					=  [2, 3]
           self._test_inference_batch_consistent(
               batch_sizes=a_     ,   additional_params_copy_to_batched_inputs=a_     ,   )
      else:
           self._test_inference_batch_consistent(
               additional_params_copy_to_batched_inputs=a_ )




 @skip_mps
 def   a      ( self       : Optional[int] )->					Any:



      """simple docstring"""

      return super().test_dict_tuple_outputs_equivalent()




 @skip_mps
 def   a      ( self       : int )->					Any:



      """simple docstring"""

      return super().test_save_load_local()






 @skip_mps
 def   a      ( self       : List[Any] )->					List[Any]:



      """simple docstring"""

      return super().test_save_load_optional_components()



@nightly
class 		UpperCAmelCase_	(unittest.TestCase						):
 """simple docstring"""




 def   a      ( self       : Optional[int] )->					Union[str, Any]:



      """simple docstring"""

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






 def   a      ( self       : int )->					Tuple:



      """simple docstring"""

      UpperCAmelCase_						:    List[Any]    					=  load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/unclip/karlo_v1_alpha_horse_cpu.npy""" )

      UpperCAmelCase_						:    List[Any]    					=  UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha""" )
      pipeline.set_progress_bar_config(disable=a_ )

      UpperCAmelCase_						:    List[str]    					=  torch.manual_seed(0 )
      UpperCAmelCase_						:    Any    					=  pipeline(
          """horse"""     ,   num_images_per_prompt=1     ,   generator=a_     ,   output_type="""np"""     ,   )

      UpperCAmelCase_						:    Optional[int]    					=  output.images[0]

      assert image.shape == (2_56, 2_56, 3)
      assert np.abs(expected_image - image ).max() < 1E-1



@slow
@require_torch_gpu
class 		UpperCAmelCase_	(unittest.TestCase						):
 """simple docstring"""




 def   a      ( self       : str )->					Any:



      """simple docstring"""

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




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



      """simple docstring"""

      UpperCAmelCase_						:    Tuple    					=  load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/unclip/karlo_v1_alpha_horse_fp16.npy""" )

      UpperCAmelCase_						:    Optional[int]    					=  UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""     ,   torch_dtype=torch.floataa )
      UpperCAmelCase_						:    Optional[int]    					=  pipeline.to(a_ )
      pipeline.set_progress_bar_config(disable=a_ )

      UpperCAmelCase_						:    Any    					=  torch.Generator(device="""cpu""" ).manual_seed(0 )
      UpperCAmelCase_						:    Union[str, Any]    					=  pipeline(
          """horse"""     ,   generator=a_     ,   output_type="""np"""     ,   )

      UpperCAmelCase_						:    Optional[Any]    					=  output.images[0]

      assert image.shape == (2_56, 2_56, 3)

      assert_mean_pixel_difference(a_     ,   a_ )






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



      """simple docstring"""

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

      UpperCAmelCase_						:    Optional[Any]    					=  UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""     ,   torch_dtype=torch.floataa )
      UpperCAmelCase_						:    List[str]    					=  pipe.to(a_ )
      pipe.set_progress_bar_config(disable=a_ )
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      UpperCAmelCase_						:    int    					=  pipe(
          """horse"""     ,   num_images_per_prompt=1     ,   prior_num_inference_steps=2     ,   decoder_num_inference_steps=2     ,   super_res_num_inference_steps=2     ,   output_type="""np"""     ,   )

      UpperCAmelCase_						:    Optional[Any]    					=  torch.cuda.max_memory_allocated()
      # make sure that less than 7 GB is allocated
      assert mem_bytes < 7 * 10**9

