

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





       """simple docstring"""


       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        torch.manual_seed(0   )
        lowerCAmelCase__   :				str 			=				TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"   )

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

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

        torch.manual_seed(0   )
        lowerCAmelCase__   :				Optional[Any] 			=				IFWatermarker()

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






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

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

        torch.manual_seed(0   )
        lowerCAmelCase__   :				int 			=				AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"   )

        torch.manual_seed(0   )
        lowerCAmelCase__   :				Dict 			=				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   )
        lowerCAmelCase__   :				str 			=				DDPMScheduler(
            num_train_timesteps=1_000			,					beta_schedule="squaredcos_cap_v2"			,					beta_start=0.0_001			,					beta_end=0.02			,					thresholding=snake_case			,					dynamic_thresholding_ratio=0.95			,					sample_max_value=1.0			,					prediction_type="epsilon"			,					variance_type="learned_range"			,					)

        torch.manual_seed(0   )
        lowerCAmelCase__   :				Tuple 			=				DDPMScheduler(
            num_train_timesteps=1_000			,					beta_schedule="squaredcos_cap_v2"			,					beta_start=0.0_001			,					beta_end=0.02			,					)

        torch.manual_seed(0   )
        lowerCAmelCase__   :				List[str] 			=				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 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				List[str] 			=				self.get_dummy_components()
        lowerCAmelCase__   :				List[str] 			=				self.pipeline_class(**snake_case   )
        pipe.to(snake_case   )
        pipe.set_progress_bar_config(disable=snake_case   )

        lowerCAmelCase__   :				int 			=				self.get_dummy_inputs(snake_case   )

        lowerCAmelCase__   :				Any 			=				inputs["prompt"]
        lowerCAmelCase__   :				List[str] 			=				inputs["generator"]
        lowerCAmelCase__   :				Optional[int] 			=				inputs["num_inference_steps"]
        lowerCAmelCase__   :				int 			=				inputs["output_type"]

        if "image" in inputs:
         lowerCAmelCase__   :				Optional[int] 			=				inputs["image"]
        else:
         lowerCAmelCase__   :				Dict 			=				None

        if "mask_image" in inputs:
         lowerCAmelCase__   :				List[str] 			=				inputs["mask_image"]
        else:
         lowerCAmelCase__   :				List[str] 			=				None

        if "original_image" in inputs:
         lowerCAmelCase__   :				int 			=				inputs["original_image"]
        else:
         lowerCAmelCase__   :				List[Any] 			=				None

        lowerCAmelCase__			,						lowerCAmelCase__   :				Union[str, Any] 			=				pipe.encode_prompt(snake_case   )

        # inputs with prompt converted to embeddings
        lowerCAmelCase__   :				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:
         lowerCAmelCase__   :				List[Any] 			=				image

        if mask_image is not None:
         lowerCAmelCase__   :				str 			=				mask_image

        if original_image is not None:
         lowerCAmelCase__   :				Optional[Any] 			=				original_image

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

        lowerCAmelCase__   :				Dict 			=				pipe(**snake_case   )[0]

        with tempfile.TemporaryDirectory() as tmpdir:
         pipe.save_pretrained(snake_case   )
         lowerCAmelCase__   :				Optional[int] 			=				self.pipeline_class.from_pretrained(snake_case   )
         pipe_loaded.to(snake_case   )
         pipe_loaded.set_progress_bar_config(disable=snake_case   )

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

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

        lowerCAmelCase__   :				Optional[Any] 			=				self.get_dummy_inputs(snake_case   )

        lowerCAmelCase__   :				List[str] 			=				inputs["generator"]
        lowerCAmelCase__   :				List[Any] 			=				inputs["num_inference_steps"]
        lowerCAmelCase__   :				Any 			=				inputs["output_type"]

        # inputs with prompt converted to embeddings
        lowerCAmelCase__   :				Union[str, Any] 			=				{
            "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:
         lowerCAmelCase__   :				Any 			=				image

        if mask_image is not None:
         lowerCAmelCase__   :				Optional[int] 			=				mask_image

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

        lowerCAmelCase__   :				Dict 			=				pipe_loaded(**snake_case   )[0]

        lowerCAmelCase__   :				List[str] 			=				np.abs(to_np(snake_case   ) - to_np(snake_case   )   ).max()
        self.assertLess(snake_case			,					1e-4   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Union[str, Any] 			=				self.get_dummy_components()
        lowerCAmelCase__   :				Optional[int] 			=				self.pipeline_class(**snake_case   )
        pipe.to(snake_case   )
        pipe.set_progress_bar_config(disable=snake_case   )

        lowerCAmelCase__   :				Optional[int] 			=				self.get_dummy_inputs(snake_case   )
        lowerCAmelCase__   :				List[Any] 			=				pipe(**snake_case   )[0]

        with tempfile.TemporaryDirectory() as tmpdir:
         pipe.save_pretrained(snake_case   )
         lowerCAmelCase__   :				Any 			=				self.pipeline_class.from_pretrained(snake_case   )
         pipe_loaded.to(snake_case   )
         pipe_loaded.set_progress_bar_config(disable=snake_case   )

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

        lowerCAmelCase__   :				Any 			=				self.get_dummy_inputs(snake_case   )
        lowerCAmelCase__   :				Optional[Any] 			=				pipe_loaded(**snake_case   )[0]

        lowerCAmelCase__   :				str 			=				np.abs(to_np(snake_case   ) - to_np(snake_case   )   ).max()
        self.assertLess(snake_case			,					1e-4   )









"""simple docstring"""

from typing import Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
    from PIL import Image

    from ..image_utils import load_image

if is_torch_available():
    from ..models.auto.modeling_auto import MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING

_UpperCamelCase		     =    logging.get_logger(__name__)






@add_end_docstrings(__magic_name__       )
class 			__a					(						__magic_name__       ):





       """simple docstring"""


       def __init__(   self			,					*snake_case			,					**snake_case   ):





        """simple docstring"""

        super().__init__(*snake_case			,					**snake_case   )
        self.check_model_type(snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case=None			,					snake_case=None			,					snake_case=None			,					**snake_case   ):





        """simple docstring"""

        lowerCAmelCase__			,						lowerCAmelCase__   :				List[Any] 			=				{}, {}
        if padding is not None:
         lowerCAmelCase__   :				Dict 			=				padding
        if truncation is not None:
         lowerCAmelCase__   :				Tuple 			=				truncation
        if top_k is not None:
         lowerCAmelCase__   :				str 			=				top_k
        return preprocess_params, {}, postprocess_params






       def __call__(   self			,					snake_case			,					snake_case = None			,					**snake_case   ):





        """simple docstring"""

        if isinstance(snake_case			,					(Image.Image, str)   ) and isinstance(snake_case			,					snake_case   ):
         lowerCAmelCase__   :				int 			=				{"image": image, "question": question}
        else:
         lowerCAmelCase__   :				Optional[int] 			=				image
        lowerCAmelCase__   :				Optional[Any] 			=				super().__call__(snake_case			,					**snake_case   )
        return results






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case=False			,					snake_case=False   ):





        """simple docstring"""

        lowerCAmelCase__   :				Union[str, Any] 			=				load_image(inputs["image"]   )
        lowerCAmelCase__   :				str 			=				self.tokenizer(
            inputs["question"]			,					return_tensors=self.framework			,					padding=snake_case			,					truncation=snake_case   )
        lowerCAmelCase__   :				int 			=				self.image_processor(images=snake_case			,					return_tensors=self.framework   )
        model_inputs.update(snake_case   )
        return model_inputs






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[int] 			=				self.model(**snake_case   )
        return model_outputs






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case=5   ):





        """simple docstring"""

        if top_k > self.model.config.num_labels:
         lowerCAmelCase__   :				List[Any] 			=				self.model.config.num_labels

        if self.framework == "pt":
         lowerCAmelCase__   :				Optional[int] 			=				model_outputs.logits.sigmoid()[0]
         lowerCAmelCase__			,						lowerCAmelCase__   :				Optional[int] 			=				probs.topk(snake_case   )
        else:
         raise ValueError(F"""Unsupported framework: {self.framework}"""   )

        lowerCAmelCase__   :				int 			=				scores.tolist()
        lowerCAmelCase__   :				str 			=				ids.tolist()
        return [{"score": score, "answer": self.model.config.idalabel[_id]} for score, _id in zip(snake_case			,					snake_case   )]






