
import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LDMTextToImagePipeline, UNetaDConditionModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_numpy,
    nightly,
    require_torch_gpu,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class       _snake_case   (	UpperCAmelCase_  ,						unittest.TestCase							):
 __lowerCAmelCase  :     Any        =			LDMTextToImagePipeline
 __lowerCAmelCase  :     Any        =			TEXT_TO_IMAGE_PARAMS - {
     'negative_prompt',
     'negative_prompt_embeds',
     'cross_attention_kwargs',
     'prompt_embeds',
 }
 __lowerCAmelCase  :     Optional[Any]        =			PipelineTesterMixin.required_optional_params - {
     'num_images_per_prompt',
     'callback',
     'callback_steps',
 }
 __lowerCAmelCase  :     int        =			TEXT_TO_IMAGE_BATCH_PARAMS
 __lowerCAmelCase  :     Any        =			False


 def      lowercase__     (					self):


   '''simple docstring'''




   torch.manual_seed(0)
   lowercase__  :   int    		=      UNetaDConditionModel(
       block_out_channels=(32, 64)					,  layers_per_block=2					,  sample_size=32					,  in_channels=4					,  out_channels=4					,  down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")					,  up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")					,  cross_attention_dim=32					,  )
   lowercase__  :   Optional[Any]    		=      DDIMScheduler(
       beta_start=0.0_0_0_8_5					,  beta_end=0.0_1_2					,  beta_schedule="""scaled_linear"""					,  clip_sample=SCREAMING_SNAKE_CASE_					,  set_alpha_to_one=SCREAMING_SNAKE_CASE_					,  )
   torch.manual_seed(0)
   lowercase__  :   Optional[Any]    		=      AutoencoderKL(
       block_out_channels=(32, 64)					,  in_channels=3					,  out_channels=3					,  down_block_types=("""DownEncoderBlock2D""", """DownEncoderBlock2D""")					,  up_block_types=("""UpDecoderBlock2D""", """UpDecoderBlock2D""")					,  latent_channels=4					,  )
   torch.manual_seed(0)
   lowercase__  :   int    		=      CLIPTextConfig(
       bos_token_id=0					,  eos_token_id=2					,  hidden_size=32					,  intermediate_size=37					,  layer_norm_eps=1E-05					,  num_attention_heads=4					,  num_hidden_layers=5					,  pad_token_id=1					,  vocab_size=10_00					,  )
   lowercase__  :   Dict    		=      CLIPTextModel(SCREAMING_SNAKE_CASE_)
   lowercase__  :   Optional[Any]    		=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""")

   lowercase__  :   List[Any]    		=      {
       """unet""": unet,
       """scheduler""": scheduler,
       """vqvae""": vae,
       """bert""": text_encoder,
       """tokenizer""": tokenizer,
   }
   return components


 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_=0):


   '''simple docstring'''




   if str(SCREAMING_SNAKE_CASE_).startswith("""mps"""):
     lowercase__  :   Optional[Any]    		=      torch.manual_seed(SCREAMING_SNAKE_CASE_)
   else:
     lowercase__  :   str    		=      torch.Generator(device=SCREAMING_SNAKE_CASE_).manual_seed(SCREAMING_SNAKE_CASE_)
   lowercase__  :   Optional[Any]    		=      {
       """prompt""": """A painting of a squirrel eating a burger""",
       """generator""": generator,
       """num_inference_steps""": 2,
       """guidance_scale""": 6.0,
       """output_type""": """numpy""",
   }
   return inputs



 def      lowercase__     (					self):


   '''simple docstring'''




   lowercase__  :   Tuple    		=      """cpu"""  # ensure determinism for the device-dependent torch.Generator

   lowercase__  :   Dict    		=      self.get_dummy_components()
   lowercase__  :   List[str]    		=      LDMTextToImagePipeline(**SCREAMING_SNAKE_CASE_)
   pipe.to(SCREAMING_SNAKE_CASE_)
   pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_)

   lowercase__  :   Union[str, Any]    		=      self.get_dummy_inputs(SCREAMING_SNAKE_CASE_)
   lowercase__  :   Optional[int]    		=      pipe(**SCREAMING_SNAKE_CASE_).images
   lowercase__  :   List[str]    		=      image[0, -3:, -3:, -1]

   assert image.shape == (1, 16, 16, 3)
   lowercase__  :   Optional[Any]    		=      np.array([0.6_1_0_1, 0.6_1_5_6, 0.5_6_2_2, 0.4_8_9_5, 0.6_6_6_1, 0.3_8_0_4, 0.5_7_4_8, 0.6_1_3_6, 0.5_0_1_4])

   assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-3



@slow
@require_torch_gpu
class       _snake_case   (	unittest.TestCase							):


 def      lowercase__     (					self):


   '''simple docstring'''




   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()


 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_=torch.floataa					,  SCREAMING_SNAKE_CASE_=0):


   '''simple docstring'''




   lowercase__  :   Any    		=      torch.manual_seed(SCREAMING_SNAKE_CASE_)
   lowercase__  :   Any    		=      np.random.RandomState(SCREAMING_SNAKE_CASE_).standard_normal((1, 4, 32, 32))
   lowercase__  :   Optional[Any]    		=      torch.from_numpy(SCREAMING_SNAKE_CASE_).to(device=SCREAMING_SNAKE_CASE_					,  dtype=SCREAMING_SNAKE_CASE_)
   lowercase__  :   Union[str, Any]    		=      {
       """prompt""": """A painting of a squirrel eating a burger""",
       """latents""": latents,
       """generator""": generator,
       """num_inference_steps""": 3,
       """guidance_scale""": 6.0,
       """output_type""": """numpy""",
   }
   return inputs



 def      lowercase__     (					self):


   '''simple docstring'''




   lowercase__  :   int    		=      LDMTextToImagePipeline.from_pretrained("""CompVis/ldm-text2im-large-256""").to(SCREAMING_SNAKE_CASE_)
   pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_)

   lowercase__  :   Optional[Any]    		=      self.get_inputs(SCREAMING_SNAKE_CASE_)
   lowercase__  :   Union[str, Any]    		=      pipe(**SCREAMING_SNAKE_CASE_).images
   lowercase__  :   List[Any]    		=      image[0, -3:, -3:, -1].flatten()

   assert image.shape == (1, 2_56, 2_56, 3)
   lowercase__  :   Any    		=      np.array([0.5_1_8_2_5, 0.5_2_8_5_0, 0.5_2_5_4_3, 0.5_4_2_5_8, 0.5_2_3_0_4, 0.5_2_5_6_9, 0.5_4_3_6_3, 0.5_5_2_7_6, 0.5_6_8_7_8])
   lowercase__  :   Optional[Any]    		=      np.abs(expected_slice - image_slice).max()
   assert max_diff < 1E-3



@nightly
@require_torch_gpu
class       _snake_case   (	unittest.TestCase							):


 def      lowercase__     (					self):


   '''simple docstring'''




   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()


 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_=torch.floataa					,  SCREAMING_SNAKE_CASE_=0):


   '''simple docstring'''




   lowercase__  :   Any    		=      torch.manual_seed(SCREAMING_SNAKE_CASE_)
   lowercase__  :   Tuple    		=      np.random.RandomState(SCREAMING_SNAKE_CASE_).standard_normal((1, 4, 32, 32))
   lowercase__  :   Union[str, Any]    		=      torch.from_numpy(SCREAMING_SNAKE_CASE_).to(device=SCREAMING_SNAKE_CASE_					,  dtype=SCREAMING_SNAKE_CASE_)
   lowercase__  :   Optional[Any]    		=      {
       """prompt""": """A painting of a squirrel eating a burger""",
       """latents""": latents,
       """generator""": generator,
       """num_inference_steps""": 50,
       """guidance_scale""": 6.0,
       """output_type""": """numpy""",
   }
   return inputs



 def      lowercase__     (					self):


   '''simple docstring'''




   lowercase__  :   Optional[int]    		=      LDMTextToImagePipeline.from_pretrained("""CompVis/ldm-text2im-large-256""").to(SCREAMING_SNAKE_CASE_)
   pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_)

   lowercase__  :   Dict    		=      self.get_inputs(SCREAMING_SNAKE_CASE_)
   lowercase__  :   Optional[Any]    		=      pipe(**SCREAMING_SNAKE_CASE_).images[0]

   lowercase__  :   str    		=      load_numpy(
       """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/ldm_text2img/ldm_large_256_ddim.npy""")
   lowercase__  :   List[Any]    		=      np.abs(expected_image - image).max()
   assert max_diff < 1E-3








import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..bit import BitConfig


lowerCamelCase__    :		List[str]          =   logging.get_logger(__name__)

lowerCamelCase__    :		str          =   {
    """Intel/dpt-large""": """https://huggingface.co/Intel/dpt-large/resolve/main/config.json""",
    # See all DPT models at https://huggingface.co/models?filter=dpt
}



class       _snake_case   (	UpperCAmelCase_							):
 __lowerCAmelCase  :     Optional[int]        =			'dpt'


 def __init__(					self					,  SCREAMING_SNAKE_CASE_=7_68					,  SCREAMING_SNAKE_CASE_=12					,  SCREAMING_SNAKE_CASE_=12					,  SCREAMING_SNAKE_CASE_=30_72					,  SCREAMING_SNAKE_CASE_="gelu"					,  SCREAMING_SNAKE_CASE_=0.0					,  SCREAMING_SNAKE_CASE_=0.0					,  SCREAMING_SNAKE_CASE_=0.0_2					,  SCREAMING_SNAKE_CASE_=1E-12					,  SCREAMING_SNAKE_CASE_=3_84					,  SCREAMING_SNAKE_CASE_=16					,  SCREAMING_SNAKE_CASE_=3					,  SCREAMING_SNAKE_CASE_=False					,  SCREAMING_SNAKE_CASE_=True					,  SCREAMING_SNAKE_CASE_=[2, 5, 8, 11]					,  SCREAMING_SNAKE_CASE_="project"					,  SCREAMING_SNAKE_CASE_=[4, 2, 1, 0.5]					,  SCREAMING_SNAKE_CASE_=[96, 1_92, 3_84, 7_68]					,  SCREAMING_SNAKE_CASE_=2_56					,  SCREAMING_SNAKE_CASE_=-1					,  SCREAMING_SNAKE_CASE_=False					,  SCREAMING_SNAKE_CASE_=True					,  SCREAMING_SNAKE_CASE_=0.4					,  SCREAMING_SNAKE_CASE_=2_55					,  SCREAMING_SNAKE_CASE_=0.1					,  SCREAMING_SNAKE_CASE_=[1, 10_24, 24, 24]					,  SCREAMING_SNAKE_CASE_=[0, 1]					,  SCREAMING_SNAKE_CASE_=None					,  **SCREAMING_SNAKE_CASE_					,  ):


   '''simple docstring'''




   super().__init__(**SCREAMING_SNAKE_CASE_)

   lowercase__  :   List[str]    		=      hidden_size
   lowercase__  :   str    		=      is_hybrid

   if self.is_hybrid:
     if backbone_config is None:
       logger.info("""Initializing the config with a `BiT` backbone.""")
       lowercase__  :   Any    		=      {
           """global_padding""": """same""",
           """layer_type""": """bottleneck""",
           """depths""": [3, 4, 9],
           """out_features""": ["""stage1""", """stage2""", """stage3"""],
           """embedding_dynamic_padding""": True,
       }
       lowercase__  :   Optional[Any]    		=      BitConfig(**SCREAMING_SNAKE_CASE_)
     elif isinstance(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_):
       logger.info("""Initializing the config with a `BiT` backbone.""")
       lowercase__  :   List[Any]    		=      BitConfig(**SCREAMING_SNAKE_CASE_)
     elif isinstance(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_):
       lowercase__  :   int    		=      backbone_config
     else:
       raise ValueError(
           f'backbone_config must be a dictionary or a `PretrainedConfig`, got {backbone_config.__class__}.')

     lowercase__  :   Optional[Any]    		=      backbone_featmap_shape
     lowercase__  :   List[str]    		=      neck_ignore_stages

     if readout_type != "project":
       raise ValueError("""Readout type must be 'project' when using `DPT-hybrid` mode.""")
   else:
     lowercase__  :   Optional[int]    		=      None
     lowercase__  :   Union[str, Any]    		=      None
     lowercase__  :   Tuple    		=      []

   lowercase__  :   List[str]    		=      num_hidden_layers
   lowercase__  :   Union[str, Any]    		=      num_attention_heads
   lowercase__  :   List[Any]    		=      intermediate_size
   lowercase__  :   Optional[Any]    		=      hidden_act
   lowercase__  :   int    		=      hidden_dropout_prob
   lowercase__  :   str    		=      attention_probs_dropout_prob
   lowercase__  :   Dict    		=      initializer_range
   lowercase__  :   Tuple    		=      layer_norm_eps
   lowercase__  :   List[str]    		=      image_size
   lowercase__  :   Tuple    		=      patch_size
   lowercase__  :   Tuple    		=      num_channels
   lowercase__  :   Dict    		=      qkv_bias
   lowercase__  :   List[Any]    		=      backbone_out_indices
   if readout_type not in ["ignore", "add", "project"]:
     raise ValueError("""Readout_type must be one of ['ignore', 'add', 'project']""")
   lowercase__  :   Dict    		=      readout_type
   lowercase__  :   Tuple    		=      reassemble_factors
   lowercase__  :   Any    		=      neck_hidden_sizes
   lowercase__  :   Dict    		=      fusion_hidden_size
   lowercase__  :   Tuple    		=      head_in_index
   lowercase__  :   Tuple    		=      use_batch_norm_in_fusion_residual
   # auxiliary head attributes (semantic segmentation)
   lowercase__  :   Tuple    		=      use_auxiliary_head
   lowercase__  :   int    		=      auxiliary_loss_weight
   lowercase__  :   Union[str, Any]    		=      semantic_loss_ignore_index
   lowercase__  :   str    		=      semantic_classifier_dropout



 def      lowercase__     (					self):


   '''simple docstring'''




   lowercase__  :   Optional[int]    		=      copy.deepcopy(self.__dict__)

   if output["backbone_config"] is not None:
     lowercase__  :   Optional[Any]    		=      self.backbone_config.to_dict()

   lowercase__  :   Dict    		=      self.__class__.model_type
   return output






