






"""simple docstring"""







# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
import difflib
from collections.abc import Mapping
from typing import Any, Callable, Dict





class      UpperCamelCase_       (				a_     ):
  pass






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase=None		,					lowerCAmelCase=False						):







       '''simple docstring'''
       try:
              UpperCAmelCase					=	ast.parse(lowerCAmelCase						)
       except SyntaxError as e:
              print("""The code generated by the agent is not valid.\n"""		,					lowerCAmelCase						)
              return
       if state is None:
              UpperCAmelCase					=	{}
       UpperCAmelCase					=	None
       for idx, node in enumerate(expression.body						):
              try:
                     UpperCAmelCase					=	evaluate_ast(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						)
              except InterpretorError as e:
                     UpperCAmelCase					=	F'''Evaluation of the code stopped at line {idx} before the end because of the following error'''
                     if chat_mode:
                            msg += (
                                F'''. Copy paste the following error message and send it back to the agent:\nI get an error: \'{e}\''''
                            )
                     else:
                            msg += F''':\n{e}'''
                     print(lowerCAmelCase						)
                     break
              if line_result is not None:
                     UpperCAmelCase					=	line_result

       return result






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       if isinstance(lowerCAmelCase		,					ast.Assign						):
              # Assignement -> we evaluate the assignement which should update the state
              # We return the variable assigned as it may be used to determine the final result.
              return evaluate_assign(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						)
       elif isinstance(lowerCAmelCase		,					ast.Call						):
              # Function call -> we return the value of the function call
              return evaluate_call(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						)
       elif isinstance(lowerCAmelCase		,					ast.Constant						):
              # Constant -> just return the value
              return expression.value
       elif isinstance(lowerCAmelCase		,					ast.Dict						):
              # Dict -> evaluate all keys and values
              UpperCAmelCase					=	[evaluate_ast(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						) for k in expression.keys]
              UpperCAmelCase					=	[evaluate_ast(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						) for v in expression.values]
              return dict(zip(lowerCAmelCase		,					lowerCAmelCase						)						)
       elif isinstance(lowerCAmelCase		,					ast.Expr						):
              # Expression -> evaluate the content
              return evaluate_ast(expression.value		,					lowerCAmelCase		,					lowerCAmelCase						)
       elif isinstance(lowerCAmelCase		,					ast.For						):
              # For loop -> execute the loop
              return evaluate_for(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						)
       elif isinstance(lowerCAmelCase		,					ast.FormattedValue						):
              # Formatted value (part of f-string) -> evaluate the content and return
              return evaluate_ast(expression.value		,					lowerCAmelCase		,					lowerCAmelCase						)
       elif isinstance(lowerCAmelCase		,					ast.If						):
              # If -> execute the right branch
              return evaluate_if(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						)
       elif hasattr(lowerCAmelCase		,					"""Index"""						) and isinstance(lowerCAmelCase		,					ast.Index						):
              return evaluate_ast(expression.value		,					lowerCAmelCase		,					lowerCAmelCase						)
       elif isinstance(lowerCAmelCase		,					ast.JoinedStr						):
              return "".join([str(evaluate_ast(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						)						) for v in expression.values]						)
       elif isinstance(lowerCAmelCase		,					ast.List						):
              # List -> evaluate all elements
              return [evaluate_ast(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						) for elt in expression.elts]
       elif isinstance(lowerCAmelCase		,					ast.Name						):
              # Name -> pick up the value in the state
              return evaluate_name(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						)
       elif isinstance(lowerCAmelCase		,					ast.Subscript						):
              # Subscript -> return the value of the indexing
              return evaluate_subscript(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						)
       else:
              # For now we refuse anything else. Let's add things as we need them.
              raise InterpretorError(F'''{expression.__class__.__name__} is not supported.'''						)






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       UpperCAmelCase					=	assign.targets
       UpperCAmelCase					=	evaluate_ast(assign.value		,					lowerCAmelCase		,					lowerCAmelCase						)

       if len(lowerCAmelCase						) == 1:
              UpperCAmelCase					=	result
       else:
              if len(lowerCAmelCase						) != len(lowerCAmelCase						):
                     raise InterpretorError(F'''Expected {len(lowerCAmelCase						)} values but got {len(lowerCAmelCase						)}.'''						)
              for var_name, r in zip(lowerCAmelCase		,					lowerCAmelCase						):
                     UpperCAmelCase					=	r
       return result






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       if not isinstance(call.func		,					ast.Name						):
              raise InterpretorError(
                  F'''It is not permitted to evaluate other functions than the provided tools (tried to execute {call.func} of '''
                  F'''type {type(call.func						)}.'''						)
       UpperCAmelCase					=	call.func.id
       if func_name not in tools:
              raise InterpretorError(
                  F'''It is not permitted to evaluate other functions than the provided tools (tried to execute {call.func.id}).'''						)

       UpperCAmelCase					=	tools[func_name]
       # Todo deal with args
       UpperCAmelCase					=	[evaluate_ast(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						) for arg in call.args]
       UpperCAmelCase					=	{keyword.arg: evaluate_ast(keyword.value		,					lowerCAmelCase		,					lowerCAmelCase						) for keyword in call.keywords}
       return func(*lowerCAmelCase		,					**lowerCAmelCase						)






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       UpperCAmelCase					=	evaluate_ast(subscript.slice		,					lowerCAmelCase		,					lowerCAmelCase						)
       UpperCAmelCase					=	evaluate_ast(subscript.value		,					lowerCAmelCase		,					lowerCAmelCase						)
       if isinstance(lowerCAmelCase		,					(list, tuple)						):
              return value[int(lowerCAmelCase						)]
       if index in value:
              return value[index]
       if isinstance(lowerCAmelCase		,					lowerCAmelCase						) and isinstance(lowerCAmelCase		,					lowerCAmelCase						):
              UpperCAmelCase					=	difflib.get_close_matches(lowerCAmelCase		,					list(value.keys()						)						)
              if len(lowerCAmelCase						) > 0:
                     return value[close_matches[0]]

       raise InterpretorError(F'''Could not index {value} with \'{index}\'.'''						)






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       if name.id in state:
              return state[name.id]
       UpperCAmelCase					=	difflib.get_close_matches(name.id		,					list(state.keys()						)						)
       if len(lowerCAmelCase						) > 0:
              return state[close_matches[0]]
       raise InterpretorError(F'''The variable `{name.id}` is not defined.'''						)






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       if len(condition.ops						) > 1:
              raise InterpretorError("""Cannot evaluate conditions with multiple operators"""						)

       UpperCAmelCase					=	evaluate_ast(condition.left		,					lowerCAmelCase		,					lowerCAmelCase						)
       UpperCAmelCase					=	condition.ops[0]
       UpperCAmelCase					=	evaluate_ast(condition.comparators[0]		,					lowerCAmelCase		,					lowerCAmelCase						)

       if isinstance(lowerCAmelCase		,					ast.Eq						):
              return left == right
       elif isinstance(lowerCAmelCase		,					ast.NotEq						):
              return left != right
       elif isinstance(lowerCAmelCase		,					ast.Lt						):
              return left < right
       elif isinstance(lowerCAmelCase		,					ast.LtE						):
              return left <= right
       elif isinstance(lowerCAmelCase		,					ast.Gt						):
              return left > right
       elif isinstance(lowerCAmelCase		,					ast.GtE						):
              return left >= right
       elif isinstance(lowerCAmelCase		,					ast.Is						):
              return left is right
       elif isinstance(lowerCAmelCase		,					ast.IsNot						):
              return left is not right
       elif isinstance(lowerCAmelCase		,					ast.In						):
              return left in right
       elif isinstance(lowerCAmelCase		,					ast.NotIn						):
              return left not in right
       else:
              raise InterpretorError(F'''Operator not supported: {comparator}'''						)






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       UpperCAmelCase					=	None
       if evaluate_condition(if_statement.test		,					lowerCAmelCase		,					lowerCAmelCase						):
              for line in if_statement.body:
                     UpperCAmelCase					=	evaluate_ast(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						)
                     if line_result is not None:
                            UpperCAmelCase					=	line_result
       else:
              for line in if_statement.orelse:
                     UpperCAmelCase					=	evaluate_ast(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						)
                     if line_result is not None:
                            UpperCAmelCase					=	line_result
       return result






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       UpperCAmelCase					=	None
       UpperCAmelCase					=	evaluate_ast(for_loop.iter		,					lowerCAmelCase		,					lowerCAmelCase						)
       for counter in iterator:
              UpperCAmelCase					=	counter
              for expression in for_loop.body:
                     UpperCAmelCase					=	evaluate_ast(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						)
                     if line_result is not None:
                            UpperCAmelCase					=	line_result
       return result


import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class      lowerCamelCase		(  A_							, unittest.TestCase  ):
     UpperCAmelCase__				:							Tuple           =		KandinskyPriorPipeline
     UpperCAmelCase__				:							List[str]           =		["prompt"]
     UpperCAmelCase__				:							List[str]           =		["prompt", "negative_prompt"]
     UpperCAmelCase__				:							List[str]           =		[
         "num_images_per_prompt",
         "generator",
         "num_inference_steps",
         "latents",
         "negative_prompt",
         "guidance_scale",
         "output_type",
         "return_dict",
     ]
     UpperCAmelCase__				:							str           =		False





     @property
     def 		UpperCAmelCase(self      :    List[str]   )						->						Union[str, Any]:
         return 3_2





     @property
     def 		UpperCAmelCase(self      :    Any   )						->						Any:
         return 3_2





     @property
     def 		UpperCAmelCase(self      :    int   )						->						Dict:
         return self.time_input_dim





     @property
     def 		UpperCAmelCase(self      :    str   )						->						Tuple:
         return self.time_input_dim * 4





     @property
     def 		UpperCAmelCase(self      :    Dict   )						->						Dict:
         return 1_0_0





     @property
     def 		UpperCAmelCase(self      :    Any   )						->						int:
         snake_case      =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )
         return tokenizer





     @property
     def 		UpperCAmelCase(self      :    Optional[Any]   )						->						Optional[Any]:
         torch.manual_seed(0   )
         snake_case      =					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=3_7				,    layer_norm_eps=1E-05				,    num_attention_heads=4				,    num_hidden_layers=5				,    pad_token_id=1				,    vocab_size=1_0_0_0				,    )
         return CLIPTextModelWithProjection(_A   )





     @property
     def 		UpperCAmelCase(self      :    List[Any]   )						->						Tuple:
         torch.manual_seed(0   )

         snake_case      =					{
             "num_attention_heads": 2,
             "attention_head_dim": 1_2,
             "embedding_dim": self.text_embedder_hidden_size,
             "num_layers": 1,
         }

         snake_case      =					PriorTransformer(**_A   )
         # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
         snake_case      =					nn.Parameter(torch.ones(model.clip_std.shape   )   )
         return model





     @property
     def 		UpperCAmelCase(self      :    str   )						->						Union[str, Any]:
         torch.manual_seed(0   )
         snake_case      =					CLIPVisionConfig(
             hidden_size=self.text_embedder_hidden_size				,    image_size=2_2_4				,    projection_dim=self.text_embedder_hidden_size				,    intermediate_size=3_7				,    num_attention_heads=4				,    num_channels=3				,    num_hidden_layers=5				,    patch_size=1_4				,    )

         snake_case      =					CLIPVisionModelWithProjection(_A   )
         return model





     @property
     def 		UpperCAmelCase(self      :    List[Any]   )						->						int:
         snake_case      =					CLIPImageProcessor(
             crop_size=2_2_4				,    do_center_crop=_A				,    do_normalize=_A				,    do_resize=_A				,    image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]				,    image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]				,    resample=3				,    size=2_2_4				,    )

         return image_processor





     def 		UpperCAmelCase(self      :    int   )						->						Union[str, Any]:
         snake_case      =					self.dummy_prior
         snake_case      =					self.dummy_image_encoder
         snake_case      =					self.dummy_text_encoder
         snake_case      =					self.dummy_tokenizer
         snake_case      =					self.dummy_image_processor

         snake_case      =					UnCLIPScheduler(
             variance_type="fixed_small_log"				,    prediction_type="sample"				,    num_train_timesteps=1_0_0_0				,    clip_sample=_A				,    clip_sample_range=10.0				,    )

         snake_case      =					{
             "prior": prior,
             "image_encoder": image_encoder,
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "scheduler": scheduler,
             "image_processor": image_processor,
         }

         return components





     def 		UpperCAmelCase(self      :    Optional[int]				,    _A      :    Tuple				,    _A      :    str=0   )						->						Union[str, Any]:
         if str(_A   ).startswith("mps"   ):
             snake_case      =					torch.manual_seed(_A   )
         else:
             snake_case      =					torch.Generator(device=_A   ).manual_seed(_A   )
         snake_case      =					{
             "prompt": "horse",
             "generator": generator,
             "guidance_scale": 4.0,
             "num_inference_steps": 2,
             "output_type": "np",
         }
         return inputs





     def 		UpperCAmelCase(self      :    List[Any]   )						->						int:
         snake_case      =					"cpu"

         snake_case      =					self.get_dummy_components()

         snake_case      =					self.pipeline_class(**_A   )
         snake_case      =					pipe.to(_A   )

         pipe.set_progress_bar_config(disable=_A   )

         snake_case      =					pipe(**self.get_dummy_inputs(_A   )   )
         snake_case      =					output.image_embeds

         snake_case      =					pipe(
             **self.get_dummy_inputs(_A   )				,    return_dict=_A				,    )[0]

         snake_case      =					image[0, -1_0:]
         snake_case      =					image_from_tuple[0, -1_0:]

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

         snake_case      =					np.array(
             [-0.05_32, 1.71_20, 0.36_56, -1.08_52, -0.89_46, -1.17_56, 0.43_48, 0.24_82, 0.51_46, -0.11_56]   )

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





     @skip_mps
     def 		UpperCAmelCase(self      :    List[Any]   )						->						str:
         snake_case      =					torch_device == "cpu"
         snake_case      =					True
         snake_case      =					False

         self._test_inference_batch_single_identical(
             test_max_difference=_A				,    relax_max_difference=_A				,    test_mean_pixel_difference=_A				,    )





     @skip_mps
     def 		UpperCAmelCase(self      :    Optional[int]   )						->						List[str]:
         snake_case      =					torch_device == "cpu"
         snake_case      =					False

         self._test_attention_slicing_forward_pass(
             test_max_difference=_A				,    test_mean_pixel_difference=_A				,    )

