






"""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		:






						def      __snake_case      (		self  :int				)      ->Optional[Any]:
										torch.manual_seed(0				)
										lowercase    :      str       = TaEncoderModel.from_pretrained("""hf-internal-testing/tiny-random-t5"""				)

										torch.manual_seed(0				)
										lowercase    :      Optional[Any]       = AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-t5"""				)

										torch.manual_seed(0				)
										lowercase    :      Tuple       = 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				)
										lowercase    :      int       = DDPMScheduler(
										    num_train_timesteps=1_000			, beta_schedule="""squaredcos_cap_v2"""			, beta_start=0.0001			, beta_end=0.02			, thresholding=__lowerCAmelCase			, dynamic_thresholding_ratio=0.95			, sample_max_value=1.0			, prediction_type="""epsilon"""			, variance_type="""learned_range"""			, )

										torch.manual_seed(0				)
										lowercase    :      str       = IFWatermarker()

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






						def      __snake_case      (		self  :Optional[int]				)      ->Optional[Any]:
										torch.manual_seed(0				)
										lowercase    :      Optional[Any]       = TaEncoderModel.from_pretrained("""hf-internal-testing/tiny-random-t5"""				)

										torch.manual_seed(0				)
										lowercase    :      List[Any]       = AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-t5"""				)

										torch.manual_seed(0				)
										lowercase    :      Optional[Any]       = 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				)
										lowercase    :      Tuple       = DDPMScheduler(
										    num_train_timesteps=1_000			, beta_schedule="""squaredcos_cap_v2"""			, beta_start=0.0001			, beta_end=0.02			, thresholding=__lowerCAmelCase			, dynamic_thresholding_ratio=0.95			, sample_max_value=1.0			, prediction_type="""epsilon"""			, variance_type="""learned_range"""			, )

										torch.manual_seed(0				)
										lowercase    :      List[str]       = DDPMScheduler(
										    num_train_timesteps=1_000			, beta_schedule="""squaredcos_cap_v2"""			, beta_start=0.0001			, beta_end=0.02			, )

										torch.manual_seed(0				)
										lowercase    :      Tuple       = 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  :Dict				)      ->List[Any]:
										lowercase    :      List[Any]       = self.get_dummy_components()
										lowercase    :      Optional[Any]       = self.pipeline_class(**__lowerCAmelCase				)
										pipe.to(__lowerCAmelCase				)
										pipe.set_progress_bar_config(disable=__lowerCAmelCase				)

										lowercase    :      Optional[int]       = self.get_dummy_inputs(__lowerCAmelCase				)

										lowercase    :      Optional[Any]       = inputs["""prompt"""]
										lowercase    :      Tuple       = inputs["""generator"""]
										lowercase    :      Union[str, Any]       = inputs["""num_inference_steps"""]
										lowercase    :      List[Any]       = inputs["""output_type"""]

										if "image" in inputs:
														lowercase    :      Optional[Any]       = inputs["""image"""]
										else:
														lowercase    :      List[Any]       = None

										if "mask_image" in inputs:
														lowercase    :      Union[str, Any]       = inputs["""mask_image"""]
										else:
														lowercase    :      Tuple       = None

										if "original_image" in inputs:
														lowercase    :      int       = inputs["""original_image"""]
										else:
														lowercase    :      Dict       = None

										lowercase      ,      lowercase    :      Optional[int]       = pipe.encode_prompt(__lowerCAmelCase				)

										# inputs with prompt converted to embeddings
										lowercase    :      Tuple       = {
										    """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:
														lowercase    :      Optional[Any]       = image

										if mask_image is not None:
														lowercase    :      List[str]       = mask_image

										if original_image is not None:
														lowercase    :      List[Any]       = original_image

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

										lowercase    :      Optional[int]       = pipe(**__lowerCAmelCase				)[0]

										with tempfile.TemporaryDirectory() as tmpdir:
														pipe.save_pretrained(__lowerCAmelCase				)
														lowercase    :      Dict       = self.pipeline_class.from_pretrained(__lowerCAmelCase				)
														pipe_loaded.to(__lowerCAmelCase				)
														pipe_loaded.set_progress_bar_config(disable=__lowerCAmelCase				)

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

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

										lowercase    :      Union[str, Any]       = self.get_dummy_inputs(__lowerCAmelCase				)

										lowercase    :      Optional[int]       = inputs["""generator"""]
										lowercase    :      Optional[Any]       = inputs["""num_inference_steps"""]
										lowercase    :      List[Any]       = inputs["""output_type"""]

										# inputs with prompt converted to embeddings
										lowercase    :      Optional[int]       = {
										    """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:
														lowercase    :      int       = image

										if mask_image is not None:
														lowercase    :      str       = mask_image

										if original_image is not None:
														lowercase    :      Optional[int]       = original_image

										lowercase    :      List[Any]       = pipe_loaded(**__lowerCAmelCase				)[0]

										lowercase    :      Optional[Any]       = np.abs(to_np(__lowerCAmelCase				) - to_np(__lowerCAmelCase				)				).max()
										self.assertLess(__lowerCAmelCase			, 1E-4				)






						def      __snake_case      (		self  :Optional[int]				)      ->Optional[Any]:
										lowercase    :      Tuple       = self.get_dummy_components()
										lowercase    :      Tuple       = self.pipeline_class(**__lowerCAmelCase				)
										pipe.to(__lowerCAmelCase				)
										pipe.set_progress_bar_config(disable=__lowerCAmelCase				)

										lowercase    :      Dict       = self.get_dummy_inputs(__lowerCAmelCase				)
										lowercase    :      str       = pipe(**__lowerCAmelCase				)[0]

										with tempfile.TemporaryDirectory() as tmpdir:
														pipe.save_pretrained(__lowerCAmelCase				)
														lowercase    :      Dict       = self.pipeline_class.from_pretrained(__lowerCAmelCase				)
														pipe_loaded.to(__lowerCAmelCase				)
														pipe_loaded.set_progress_bar_config(disable=__lowerCAmelCase				)

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

										lowercase    :      Any       = self.get_dummy_inputs(__lowerCAmelCase				)
										lowercase    :      Optional[int]       = pipe_loaded(**__lowerCAmelCase				)[0]

										lowercase    :      List[Any]       = np.abs(to_np(__lowerCAmelCase				) - to_np(__lowerCAmelCase				)				).max()
										self.assertLess(__lowerCAmelCase			, 1E-4				)

"""simple docstring"""






from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Features, Sequence, Value
from .base import TaskTemplate





@dataclass(frozen=lowerCAmelCase      )
class 					_A		(			lowerCAmelCase      ):
      # `task` is not a ClassVar since we want it to be part of the `asdict` output for JSON serialization
      snake_case__     :  str				    =      field(default='question-answering-extractive'    ,       metadata={'include_in_asdict_even_if_is_default': True}      )
      snake_case__     :  ClassVar[Features]				    =      Features({'question': Value('string'      ), 'context': Value('string'      )}      )
      snake_case__     :  ClassVar[Features]				    =      Features(
          {
              'answers': Sequence(
                  {
                      'text': Value('string'      ),
                      'answer_start': Value('int32'      ),
                  }      )
          }      )
      snake_case__     :  str				    =      "question"
      snake_case__     :  str				    =      "context"
      snake_case__     :  str				    =      "answers"

      @property
      def 							A__			(		self		):







            """simple docstring"""



            return {self.question_column: "question", self.context_column: "context", self.answers_column: "answers"}



