





'''simple docstring'''






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


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

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

  _UpperCamelCase       			=multiprocessing.Process(target=__SCREAMING_SNAKE_CASE  ,      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 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

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

    _UpperCamelCase       			=shutil.rmtree
    _UpperCamelCase       			=os.rmdir
    _UpperCamelCase       			=os.chdir

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

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

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


@contextlib.contextmanager
def 					_a   (__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

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

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


@contextlib.contextmanager
def 					_a   ():


  """simple docstring"""

  _UpperCamelCase       			=WriteOnlyStringIO()
  with contextlib.redirect_stdout(__SCREAMING_SNAKE_CASE   ):
    with contextlib.redirect_stderr(__SCREAMING_SNAKE_CASE   ):
      with redirect_stdin(__SCREAMING_SNAKE_CASE   ):
        yield


@contextlib.contextmanager
def 					_a   ():


  """simple docstring"""

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







class 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







  pass







class 		UpperCAmelCase   (				io.StringIO):


  """simple docstring"""







  def   UpperCamelCase__      (				self    :					Tuple							,	*UpperCamelCase__    :					Optional[Any]							,	**UpperCamelCase__    :					Tuple  )     ->						Tuple:
    raise OSError


  def   UpperCamelCase__      (				self    :					Optional[Any]							,	*UpperCamelCase__    :					Optional[int]							,	**UpperCamelCase__    :					Union[str, Any]  )     ->						int:
    raise OSError


  def   UpperCamelCase__      (				self    :					Union[str, Any]							,	*UpperCamelCase__    :					Dict							,	**UpperCamelCase__    :					List[str]  )     ->						List[Any]:
    raise OSError




  def   UpperCamelCase__      (				self    :					Any							,	*UpperCamelCase__    :					Optional[Any]							,	**UpperCamelCase__    :					str  )     ->						Tuple:
    return False







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


  """simple docstring"""







  lowerCAmelCase_					=				"""stdin"""




@contextlib.contextmanager
def 					_a   (__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

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


def 					_a   (__SCREAMING_SNAKE_CASE=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

  _UpperCamelCase       			=None
  _UpperCamelCase       			=None

  import os

  _UpperCamelCase       			='''1'''

  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None

  import shutil

  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None

  import subprocess

  _UpperCamelCase       			=None  # type: ignore

  _UpperCamelCase       			=None

  import sys

  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None
  _UpperCamelCase       			=None













'''simple docstring'''






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   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Union[str, Any]:
    torch.manual_seed(0  )
    _UpperCamelCase       			=TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )

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

    torch.manual_seed(0  )
    _UpperCamelCase       			=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  )
    _UpperCamelCase       			=DDPMScheduler(
        num_train_timesteps=1000							,	beta_schedule='''squaredcos_cap_v2'''							,	beta_start=0.0001							,	beta_end=0.02							,	thresholding=UpperCamelCase__							,	dynamic_thresholding_ratio=0.95							,	sample_max_value=1.0							,	prediction_type='''epsilon'''							,	variance_type='''learned_range'''							,	)

    torch.manual_seed(0  )
    _UpperCamelCase       			=IFWatermarker()

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


  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						Optional[int]:
    torch.manual_seed(0  )
    _UpperCamelCase       			=TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )

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

    torch.manual_seed(0  )
    _UpperCamelCase       			=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  )
    _UpperCamelCase       			=DDPMScheduler(
        num_train_timesteps=1000							,	beta_schedule='''squaredcos_cap_v2'''							,	beta_start=0.0001							,	beta_end=0.02							,	thresholding=UpperCamelCase__							,	dynamic_thresholding_ratio=0.95							,	sample_max_value=1.0							,	prediction_type='''epsilon'''							,	variance_type='''learned_range'''							,	)

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

    torch.manual_seed(0  )
    _UpperCamelCase       			=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   UpperCamelCase__      (				self    :					Any  )     ->						Any:
    _UpperCamelCase       			=self.get_dummy_components()
    _UpperCamelCase       			=self.pipeline_class(**UpperCamelCase__  )
    pipe.to(UpperCamelCase__  )
    pipe.set_progress_bar_config(disable=UpperCamelCase__  )

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )

    _UpperCamelCase       			=inputs['''prompt''']
    _UpperCamelCase       			=inputs['''generator''']
    _UpperCamelCase       			=inputs['''num_inference_steps''']
    _UpperCamelCase       			=inputs['''output_type''']

    if "image" in inputs:
      _UpperCamelCase       			=inputs['''image''']
    else:
      _UpperCamelCase       			=None

    if "mask_image" in inputs:
      _UpperCamelCase       			=inputs['''mask_image''']
    else:
      _UpperCamelCase       			=None

    if "original_image" in inputs:
      _UpperCamelCase       			=inputs['''original_image''']
    else:
      _UpperCamelCase       			=None

    _UpperCamelCase			,     _UpperCamelCase       			=pipe.encode_prompt(UpperCamelCase__  )

    # inputs with prompt converted to embeddings
    _UpperCamelCase       			={
        '''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:
      _UpperCamelCase       			=image

    if mask_image is not None:
      _UpperCamelCase       			=mask_image

    if original_image is not None:
      _UpperCamelCase       			=original_image

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

    _UpperCamelCase       			=pipe(**UpperCamelCase__  )[0]

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

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

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

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )

    _UpperCamelCase       			=inputs['''generator''']
    _UpperCamelCase       			=inputs['''num_inference_steps''']
    _UpperCamelCase       			=inputs['''output_type''']

    # inputs with prompt converted to embeddings
    _UpperCamelCase       			={
        '''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:
      _UpperCamelCase       			=image

    if mask_image is not None:
      _UpperCamelCase       			=mask_image

    if original_image is not None:
      _UpperCamelCase       			=original_image

    _UpperCamelCase       			=pipe_loaded(**UpperCamelCase__  )[0]

    _UpperCamelCase       			=np.abs(to_np(UpperCamelCase__  ) - to_np(UpperCamelCase__  )  ).max()
    self.assertLess(UpperCamelCase__							,	1E-4  )




  def   UpperCamelCase__      (				self    :					int  )     ->						str:
    _UpperCamelCase       			=self.get_dummy_components()
    _UpperCamelCase       			=self.pipeline_class(**UpperCamelCase__  )
    pipe.to(UpperCamelCase__  )
    pipe.set_progress_bar_config(disable=UpperCamelCase__  )

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )
    _UpperCamelCase       			=pipe(**UpperCamelCase__  )[0]

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

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

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )
    _UpperCamelCase       			=pipe_loaded(**UpperCamelCase__  )[0]

    _UpperCamelCase       			=np.abs(to_np(UpperCamelCase__  ) - to_np(UpperCamelCase__  )  ).max()
    self.assertLess(UpperCamelCase__							,	1E-4  )






