






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






    def 							_UpperCamelCase				(					self   ):
          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=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   )
          UpperCAmelCase							     =			IFWatermarker()

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






    def 							_UpperCamelCase				(					self   ):
          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=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   )
          UpperCAmelCase							     =			DDPMScheduler(
              num_train_timesteps=1_000      ,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   ):
          UpperCAmelCase							     =			self.get_dummy_components()
          UpperCAmelCase							     =			self.pipeline_class(**_lowerCamelCase   )
          pipe.to(_lowerCamelCase   )
          pipe.set_progress_bar_config(disable=_lowerCamelCase   )

          UpperCAmelCase							     =			self.get_dummy_inputs(_lowerCamelCase   )

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

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

          UpperCAmelCase							     =			pipe(**_lowerCamelCase   )[0]

          with tempfile.TemporaryDirectory() as tmpdir:
                pipe.save_pretrained(_lowerCamelCase   )
                UpperCAmelCase							     =			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.'''      ,)

          UpperCAmelCase							     =			self.get_dummy_inputs(_lowerCamelCase   )

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

          UpperCAmelCase							     =			np.abs(to_np(_lowerCamelCase   ) - to_np(_lowerCamelCase   )   ).max()
          self.assertLess(_lowerCamelCase      ,1e-4   )






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			self.get_dummy_components()
          UpperCAmelCase							     =			self.pipeline_class(**_lowerCamelCase   )
          pipe.to(_lowerCamelCase   )
          pipe.set_progress_bar_config(disable=_lowerCamelCase   )

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

          with tempfile.TemporaryDirectory() as tmpdir:
                pipe.save_pretrained(_lowerCamelCase   )
                UpperCAmelCase							     =			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

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

          UpperCAmelCase							     =			np.abs(to_np(_lowerCamelCase   ) - to_np(_lowerCamelCase   )   ).max()
          self.assertLess(_lowerCamelCase      ,1e-4   )




"""simple docstring"""


from .imports import is_rich_available


if is_rich_available():
	from rich.traceback import install

	install(show_locals=False)

else:
	raise ModuleNotFoundError('To use the rich extension, install rich with `pip install rich`')


