# This code is adapted from OpenAI's release
# https://github.com/openai/human-eval/blob/master/human_eval/execution.py

import contextlib
import faulthandler
import io
import multiprocessing
import os
import platform
import signal
import tempfile





def 							UpperCAmelCase		(       lowercase				,     lowercase				,     lowercase				,     lowercase				):







				"""simple docstring"""



				__lowercase					     =				multiprocessing.Manager()
				__lowercase					     =				manager.list()

				__lowercase					     =				multiprocessing.Process(target=lowercase				,     args=(check_program, result, timeout)				)
				p.start()
				p.join(timeout=timeout + 1				)
				if p.is_alive():
								p.kill()

				if not result:
								result.append('''timed out'''				)

				return {
				    "task_id": task_id,
				    "passed": result[0] == "passed",
				    "result": result[0],
				    "completion_id": completion_id,
				}





def 							UpperCAmelCase		(       lowercase				,     lowercase				,     lowercase				):







				"""simple docstring"""



				with create_tempdir():
								# These system calls are needed when cleaning up tempdir.
								import os
								import shutil

								__lowercase					     =				shutil.rmtree
								__lowercase					     =				os.rmdir
								__lowercase					     =				os.chdir

								# Disable functionalities that can make destructive changes to the test.
								reliability_guard()

								# Run program.
								try:
												__lowercase					     =				{}
												with swallow_io():
																with time_limit(lowercase				):
																				exec(lowercase				,     lowercase				)
												result.append('''passed'''				)
								except TimeoutException:
												result.append('''timed out'''				)
								except BaseException as e:
												result.append(F"failed: {e}"				)

								# Needed for cleaning up.
								__lowercase					     =				rmtree
								__lowercase					     =				rmdir
								__lowercase					     =				chdir





@contextlib.contextmanager
def 							UpperCAmelCase		(       lowercase				):







				"""simple docstring"""



				def signal_handler(lowercase				,     lowercase				):
								raise TimeoutException('''Timed out!'''				)

				signal.setitimer(signal.ITIMER_REAL				,     lowercase				)
				signal.signal(signal.SIGALRM				,     lowercase				)
				try:
								yield
				finally:
								signal.setitimer(signal.ITIMER_REAL				,     0				)





@contextlib.contextmanager
def 							UpperCAmelCase		(       ):







				"""simple docstring"""



				__lowercase					     =				WriteOnlyStringIO()
				with contextlib.redirect_stdout(lowercase				):
								with contextlib.redirect_stderr(lowercase				):
												with redirect_stdin(lowercase				):
																yield





@contextlib.contextmanager
def 							UpperCAmelCase		(       ):







				"""simple docstring"""



				with tempfile.TemporaryDirectory() as dirname:
								with chdir(lowercase				):
												yield dirname




class   _UpperCamelCase (		_UpperCAmelCase					):


			"""simple docstring"""






			pass




class   _UpperCamelCase (		io.StringIO					):


			"""simple docstring"""






			def 			_SCREAMING_SNAKE_CASE     (  self ,  *lowerCAmelCase__ ,  **lowerCAmelCase__		)	->   str:

							'''simple docstring'''


							raise OSError





			def 			_SCREAMING_SNAKE_CASE     (  self ,  *lowerCAmelCase__ ,  **lowerCAmelCase__		)	->   Optional[Any]:

							'''simple docstring'''


							raise OSError





			def 			_SCREAMING_SNAKE_CASE     (  self ,  *lowerCAmelCase__ ,  **lowerCAmelCase__		)	->   Optional[Any]:

							'''simple docstring'''


							raise OSError





			def 			_SCREAMING_SNAKE_CASE     (  self ,  *lowerCAmelCase__ ,  **lowerCAmelCase__		)	->   Optional[Any]:

							'''simple docstring'''


							return False




class   _UpperCamelCase (		contextlib._RedirectStream					):  # type: ignore


			"""simple docstring"""






			__a    :					List[Any]     					=		'''stdin'''





@contextlib.contextmanager
def 							UpperCAmelCase		(       lowercase				):







				"""simple docstring"""



				if root == ".":
								yield
								return
				__lowercase					     =				os.getcwd()
				os.chdir(lowercase				)
				try:
								yield
				except BaseException as exc:
								raise exc
				finally:
								os.chdir(lowercase				)





def 							UpperCAmelCase		(       lowercase=None				):







				"""simple docstring"""



				if maximum_memory_bytes is not None:
								import resource

								resource.setrlimit(resource.RLIMIT_AS				,     (maximum_memory_bytes, maximum_memory_bytes)				)
								resource.setrlimit(resource.RLIMIT_DATA				,     (maximum_memory_bytes, maximum_memory_bytes)				)
								if not platform.uname().system == "Darwin":
												resource.setrlimit(resource.RLIMIT_STACK				,     (maximum_memory_bytes, maximum_memory_bytes)				)

				faulthandler.disable()

				import builtins

				__lowercase					     =				None
				__lowercase					     =				None

				import os

				__lowercase					     =				'''1'''

				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None

				import shutil

				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None

				import subprocess

				__lowercase					     =				None  # type: ignore

				__lowercase					     =				None

				import sys

				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
				__lowercase					     =				None
import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin




class   _UpperCamelCase (		_UpperCAmelCase		,_UpperCAmelCase					):


			"""simple docstring"""






			@register_to_config
			def __init__(  self ,  *,
			    lowerCAmelCase__ = 4 ,  lowerCAmelCase__ = 7_68 ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  )	->   Dict:

							'''simple docstring'''


							super().__init__()

							__lowercase					     =				nn.Parameter(torch.zeros(lowerCAmelCase__		)		)

							# parameters for additional clip time embeddings
							__lowercase					     =				nn.Linear(lowerCAmelCase__ ,  lowerCAmelCase__		)
							__lowercase					     =				nn.Linear(lowerCAmelCase__ ,  lowerCAmelCase__		)

							# parameters for encoder hidden states
							__lowercase					     =				clip_extra_context_tokens
							__lowercase					     =				nn.Linear(
							    lowerCAmelCase__ ,  self.clip_extra_context_tokens * cross_attention_dim		)
							__lowercase					     =				nn.Linear(lowerCAmelCase__ ,  lowerCAmelCase__		)
							__lowercase					     =				nn.LayerNorm(lowerCAmelCase__		)





			def 			_SCREAMING_SNAKE_CASE     (  self ,  *, lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   List[str]:

							'''simple docstring'''


							if do_classifier_free_guidance:
											# Add the classifier free guidance embeddings to the image embeddings
											__lowercase					     =				image_embeddings.shape[0]
											__lowercase					     =				self.learned_classifier_free_guidance_embeddings.unsqueeze(0		)
											__lowercase					     =				classifier_free_guidance_embeddings.expand(
											    lowerCAmelCase__ ,  -1		)
											__lowercase					     =				torch.cat([classifier_free_guidance_embeddings, image_embeddings] ,  dim=0		)

							# The image embeddings batch size and the text embeddings batch size are equal
							assert image_embeddings.shape[0] == prompt_embeds.shape[0]

							__lowercase					     =				prompt_embeds.shape[0]

							# "Specifically, we modify the architecture described in Nichol et al. (2021) by projecting and
							# adding CLIP embeddings to the existing timestep embedding, ...
							__lowercase					     =				self.embedding_proj(lowerCAmelCase__		)
							__lowercase					     =				self.clip_image_embeddings_project_to_time_embeddings(lowerCAmelCase__		)
							__lowercase					     =				time_projected_image_embeddings + time_projected_prompt_embeds

							# ... and by projecting CLIP embeddings into four
							# extra tokens of context that are concatenated to the sequence of outputs from the GLIDE text encoder"
							__lowercase					     =				self.clip_extra_context_tokens_proj(lowerCAmelCase__		)
							__lowercase					     =				clip_extra_context_tokens.reshape(lowerCAmelCase__ ,  -1 ,  self.clip_extra_context_tokens		)
							__lowercase					     =				clip_extra_context_tokens.permute(0 ,  2 ,  1		)

							__lowercase					     =				self.encoder_hidden_states_proj(lowerCAmelCase__		)
							__lowercase					     =				self.text_encoder_hidden_states_norm(lowerCAmelCase__		)
							__lowercase					     =				torch.cat([clip_extra_context_tokens, text_encoder_hidden_states] ,  dim=1		)

							return text_encoder_hidden_states, additive_clip_time_embeddings