

'''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 							lowercase_						:

  '''simple docstring'''





  def     snake_case_      (  self  )				->					str:


    """simple docstring"""
    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=3_2 ,			layers_per_block=1 ,			block_out_channels=[3_2, 6_4] ,			down_block_types=[
            'ResnetDownsampleBlock2D',
            'SimpleCrossAttnDownBlock2D',
        ] ,			mid_block_type='UNetMidBlock2DSimpleCrossAttn' ,			up_block_types=['SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'] ,			in_channels=3 ,			out_channels=6 ,			cross_attention_dim=3_2 ,			encoder_hid_dim=3_2 ,			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=1_0_0_0 ,			beta_schedule='squaredcos_cap_v2' ,			beta_start=0.0001 ,			beta_end=0.02 ,			thresholding=a_ ,			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     snake_case_      (  self  )				->					Tuple:


    """simple docstring"""
    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=3_2 ,			layers_per_block=[1, 2] ,			block_out_channels=[3_2, 6_4] ,			down_block_types=[
            'ResnetDownsampleBlock2D',
            'SimpleCrossAttnDownBlock2D',
        ] ,			mid_block_type='UNetMidBlock2DSimpleCrossAttn' ,			up_block_types=['SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'] ,			in_channels=6 ,			out_channels=6 ,			cross_attention_dim=3_2 ,			encoder_hid_dim=3_2 ,			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=3_2 ,			)
    unet.set_attn_processor(AttnAddedKVProcessor()  )  # For reproducibility tests

    torch.manual_seed(0  )
    UpperCAmelCase    				=						DDPMScheduler(
        num_train_timesteps=1_0_0_0 ,			beta_schedule='squaredcos_cap_v2' ,			beta_start=0.0001 ,			beta_end=0.02 ,			thresholding=a_ ,			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=1_0_0_0 ,			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     snake_case_      (  self  )				->					List[str]:


    """simple docstring"""
    UpperCAmelCase    				=						self.get_dummy_components()
    UpperCAmelCase    				=						self.pipeline_class(**a_  )
    pipe.to(a_  )
    pipe.set_progress_bar_config(disable=a_  )

    UpperCAmelCase    				=						self.get_dummy_inputs(a_  )

    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(a_  )

    # 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(a_ ,			a_ ,			a_  )

    UpperCAmelCase    				=						pipe(**a_  )[0]

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

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

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

    UpperCAmelCase    				=						self.get_dummy_inputs(a_  )

    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(**a_  )[0]

    UpperCAmelCase    				=						np.abs(to_np(a_  ) - to_np(a_  )  ).max()
    self.assertLess(a_ ,			1E-4  )





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


    """simple docstring"""
    UpperCAmelCase    				=						self.get_dummy_components()
    UpperCAmelCase    				=						self.pipeline_class(**a_  )
    pipe.to(a_  )
    pipe.set_progress_bar_config(disable=a_  )

    UpperCAmelCase    				=						self.get_dummy_inputs(a_  )
    UpperCAmelCase    				=						pipe(**a_  )[0]

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

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

    UpperCAmelCase    				=						self.get_dummy_inputs(a_  )
    UpperCAmelCase    				=						pipe_loaded(**a_  )[0]

    UpperCAmelCase    				=						np.abs(to_np(a_  ) - to_np(a_  )  ).max()
    self.assertLess(a_ ,			1E-4  )





'''simple docstring'''




import subprocess
import sys

from transformers import BertConfig, BertModel, BertTokenizer, pipeline
from transformers.testing_utils import TestCasePlus, require_torch



class 							lowercase_						(		a			):

  '''simple docstring'''





  @require_torch
  def     snake_case_      (  self  )				->					Optional[Any]:


    """simple docstring"""
    # this test is a bit tricky since TRANSFORMERS_OFFLINE can only be changed before
    # `transformers` is loaded, and it's too late for inside pytest - so we are changing it
    # while running an external program

    # python one-liner segments

    # this must be loaded before socket.socket is monkey-patched
    UpperCAmelCase    				=						'\nfrom transformers import BertConfig, BertModel, BertTokenizer, pipeline\n        '

    UpperCAmelCase    				=						'\nmname = "hf-internal-testing/tiny-random-bert"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nBertTokenizer.from_pretrained(mname)\npipe = pipeline(task="fill-mask", model=mname)\nprint("success")\n        '

    UpperCAmelCase    				=						'\nimport socket\ndef offline_socket(*args, **kwargs): raise RuntimeError("Offline mode is enabled, we shouldn\'t access internet")\nsocket.socket = offline_socket\n        '

    # Force fetching the files so that we can use the cache
    UpperCAmelCase    				=						'hf-internal-testing/tiny-random-bert'
    BertConfig.from_pretrained(a_  )
    BertModel.from_pretrained(a_  )
    BertTokenizer.from_pretrained(a_  )
    pipeline(task='fill-mask' ,			model=a_  )

    # baseline - just load from_pretrained with normal network
    UpperCAmelCase    				=						[sys.executable, '-c', '\n'.join([load, run, mock]  )]

    # should succeed
    UpperCAmelCase    				=						self.get_env()
    # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
    UpperCAmelCase    				=						'1'
    UpperCAmelCase    				=						subprocess.run(a_ ,			env=a_ ,			check=a_ ,			capture_output=a_  )
    self.assertEqual(result.returncode ,			0 ,			result.stderr  )
    self.assertIn('success' ,			result.stdout.decode()  )
  @require_torch
  def     snake_case_      (  self  )				->					Union[str, Any]:


    """simple docstring"""
    # python one-liner segments
    # this must be loaded before socket.socket is monkey-patched
    UpperCAmelCase    				=						'\nfrom transformers import BertConfig, BertModel, BertTokenizer, pipeline\n        '

    UpperCAmelCase    				=						'\nmname = "hf-internal-testing/tiny-random-bert"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nBertTokenizer.from_pretrained(mname)\npipe = pipeline(task="fill-mask", model=mname)\nprint("success")\n        '

    UpperCAmelCase    				=						'\nimport socket\ndef offline_socket(*args, **kwargs): raise socket.error("Faking flaky internet")\nsocket.socket = offline_socket\n        '

    # Force fetching the files so that we can use the cache
    UpperCAmelCase    				=						'hf-internal-testing/tiny-random-bert'
    BertConfig.from_pretrained(a_  )
    BertModel.from_pretrained(a_  )
    BertTokenizer.from_pretrained(a_  )
    pipeline(task='fill-mask' ,			model=a_  )

    # baseline - just load from_pretrained with normal network
    UpperCAmelCase    				=						[sys.executable, '-c', '\n'.join([load, run, mock]  )]

    # should succeed
    UpperCAmelCase    				=						self.get_env()
    UpperCAmelCase    				=						subprocess.run(a_ ,			env=a_ ,			check=a_ ,			capture_output=a_  )
    self.assertEqual(result.returncode ,			0 ,			result.stderr  )
    self.assertIn('success' ,			result.stdout.decode()  )
  @require_torch
  def     snake_case_      (  self  )				->					List[str]:


    """simple docstring"""
    # this test is a bit tricky since TRANSFORMERS_OFFLINE can only be changed before
    # `transformers` is loaded, and it's too late for inside pytest - so we are changing it
    # while running an external program

    # python one-liner segments

    # this must be loaded before socket.socket is monkey-patched
    UpperCAmelCase    				=						'\nfrom transformers import BertConfig, BertModel, BertTokenizer\n        '

    UpperCAmelCase    				=						'\nmname = "hf-internal-testing/tiny-random-bert-sharded"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nprint("success")\n        '

    UpperCAmelCase    				=						'\nimport socket\ndef offline_socket(*args, **kwargs): raise ValueError("Offline mode is enabled")\nsocket.socket = offline_socket\n        '

    # baseline - just load from_pretrained with normal network
    UpperCAmelCase    				=						[sys.executable, '-c', '\n'.join([load, run]  )]

    # should succeed
    UpperCAmelCase    				=						self.get_env()
    UpperCAmelCase    				=						subprocess.run(a_ ,			env=a_ ,			check=a_ ,			capture_output=a_  )
    self.assertEqual(result.returncode ,			0 ,			result.stderr  )
    self.assertIn('success' ,			result.stdout.decode()  )

    # next emulate no network
    UpperCAmelCase    				=						[sys.executable, '-c', '\n'.join([load, mock, run]  )]

    # Doesn't fail anymore since the model is in the cache due to other tests, so commenting this.
    # env["TRANSFORMERS_OFFLINE"] = "0"
    # result = subprocess.run(cmd, env=env, check=False, capture_output=True)
    # self.assertEqual(result.returncode, 1, result.stderr)

    # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
    UpperCAmelCase    				=						'1'
    UpperCAmelCase    				=						subprocess.run(a_ ,			env=a_ ,			check=a_ ,			capture_output=a_  )
    self.assertEqual(result.returncode ,			0 ,			result.stderr  )
    self.assertIn('success' ,			result.stdout.decode()  )
  @require_torch
  def     snake_case_      (  self  )				->					Optional[Any]:


    """simple docstring"""
    UpperCAmelCase    				=						'\nfrom transformers import pipeline\n        '
    UpperCAmelCase    				=						'\nmname = "hf-internal-testing/tiny-random-bert"\npipe = pipeline(model=mname)\n        '

    UpperCAmelCase    				=						'\nimport socket\ndef offline_socket(*args, **kwargs): raise socket.error("Offline mode is enabled")\nsocket.socket = offline_socket\n        '
    UpperCAmelCase    				=						self.get_env()
    UpperCAmelCase    				=						'1'
    UpperCAmelCase    				=						[sys.executable, '-c', '\n'.join([load, mock, run]  )]
    UpperCAmelCase    				=						subprocess.run(a_ ,			env=a_ ,			check=a_ ,			capture_output=a_  )
    self.assertEqual(result.returncode ,			1 ,			result.stderr  )
    self.assertIn(
        'You cannot infer task automatically within `pipeline` when using offline mode' ,			result.stderr.decode().replace('\n' ,			''  ) ,			)





  @require_torch
  def     snake_case_      (  self  )				->					str:


    """simple docstring"""
    UpperCAmelCase    				=						'\nfrom transformers import AutoModel\n        '
    UpperCAmelCase    				=						'\nmname = "hf-internal-testing/test_dynamic_model"\nAutoModel.from_pretrained(mname, trust_remote_code=True)\nprint("success")\n        '

    # baseline - just load from_pretrained with normal network
    UpperCAmelCase    				=						[sys.executable, '-c', '\n'.join([load, run]  )]

    # should succeed
    UpperCAmelCase    				=						self.get_env()
    UpperCAmelCase    				=						subprocess.run(a_ ,			env=a_ ,			check=a_ ,			capture_output=a_  )
    self.assertEqual(result.returncode ,			0 ,			result.stderr  )
    self.assertIn('success' ,			result.stdout.decode()  )

    # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
    UpperCAmelCase    				=						'1'
    UpperCAmelCase    				=						subprocess.run(a_ ,			env=a_ ,			check=a_ ,			capture_output=a_  )
    self.assertEqual(result.returncode ,			0 ,			result.stderr  )
    self.assertIn('success' ,			result.stdout.decode()  )


