



"""simple docstring"""
import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()


class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase      ,  __UpperCAmelCase      ,  __UpperCAmelCase      ,  unittest.TestCase   ):
    lowercase__      =   StableUnCLIPImgaImgPipeline
    lowercase__      =   TEXT_GUIDED_IMAGE_VARIATION_PARAMS
    lowercase__      =   TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    lowercase__      =   frozenset(
        []   )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
    lowercase__      =   frozenset([]   )

    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      lowercase_								= 3_2
      lowercase_								= embedder_hidden_size

      # image encoding components

      lowercase_								= CLIPImageProcessor(crop_size=3_2 ,  size=3_2)

      torch.manual_seed(0)
      lowercase_								= CLIPVisionModelWithProjection(
          CLIPVisionConfig(
              hidden_size=lowerCAmelCase_ ,  projection_dim=lowerCAmelCase_ ,  num_hidden_layers=5 ,  num_attention_heads=4 ,  image_size=3_2 ,  intermediate_size=3_7 ,  patch_size=1 ,  ))

      # regular denoising components

      torch.manual_seed(0)
      lowercase_								= StableUnCLIPImageNormalizer(embedding_dim=lowerCAmelCase_)
      lowercase_								= DDPMScheduler(beta_schedule="""squaredcos_cap_v2""")

      torch.manual_seed(0)
      lowercase_								= CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""")

      torch.manual_seed(0)
      lowercase_								= CLIPTextModel(
          CLIPTextConfig(
              bos_token_id=0 ,  eos_token_id=2 ,  hidden_size=lowerCAmelCase_ ,  projection_dim=3_2 ,  intermediate_size=3_7 ,  layer_norm_eps=1E-05 ,  num_attention_heads=4 ,  num_hidden_layers=5 ,  pad_token_id=1 ,  vocab_size=1_0_0_0 ,  ))

      torch.manual_seed(0)
      lowercase_								= UNetaDConditionModel(
          sample_size=3_2 ,  in_channels=4 ,  out_channels=4 ,  down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""") ,  up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""") ,  block_out_channels=(3_2, 6_4) ,  attention_head_dim=(2, 4) ,  class_embed_type="""projection""" ,  projection_class_embeddings_input_dim=embedder_projection_dim * 2 ,  cross_attention_dim=lowerCAmelCase_ ,  layers_per_block=1 ,  upcast_attention=lowerCAmelCase_ ,  use_linear_projection=lowerCAmelCase_ ,  )

      torch.manual_seed(0)
      lowercase_								= DDIMScheduler(
          beta_schedule="""scaled_linear""" ,  beta_start=0.00_085 ,  beta_end=0.012 ,  prediction_type="""v_prediction""" ,  set_alpha_to_one=lowerCAmelCase_ ,  steps_offset=1 ,  )

      torch.manual_seed(0)
      lowercase_								= AutoencoderKL()

      lowercase_								= {
          # image encoding components
          """feature_extractor""": feature_extractor,
          """image_encoder""": image_encoder.eval(),
          # image noising components
          """image_normalizer""": image_normalizer.eval(),
          """image_noising_scheduler""": image_noising_scheduler,
          # regular denoising components
          """tokenizer""": tokenizer,
          """text_encoder""": text_encoder.eval(),
          """unet""": unet.eval(),
          """scheduler""": scheduler,
          """vae""": vae.eval(),
      }

      return components

    def 					_UpperCAmelCase	(				self			:				str ,  lowerCAmelCase_			:				Optional[int] ,  lowerCAmelCase_			:				List[str]=0 ,  lowerCAmelCase_			:				str=True):







      """simple docstring"""







      if str(lowerCAmelCase_).startswith("""mps"""):
        lowercase_								= torch.manual_seed(lowerCAmelCase_)
      else:
        lowercase_								= torch.Generator(device=lowerCAmelCase_).manual_seed(lowerCAmelCase_)

      lowercase_								= floats_tensor((1, 3, 3_2, 3_2) ,  rng=random.Random(lowerCAmelCase_)).to(lowerCAmelCase_)

      if pil_image:
        lowercase_								= input_image * 0.5 + 0.5
        lowercase_								= input_image.clamp(0 ,  1)
        lowercase_								= input_image.cpu().permute(0 ,  2 ,  3 ,  1).float().numpy()
        lowercase_								= DiffusionPipeline.numpy_to_pil(lowerCAmelCase_)[0]

      return {
          "prompt": "An anime racoon running a marathon",
          "image": input_image,
          "generator": generator,
          "num_inference_steps": 2,
          "output_type": "np",
      }

    @skip_mps
    def 					_UpperCAmelCase	(				self			:				str):







      """simple docstring"""







      lowercase_								= """cpu"""  # ensure determinism for the device-dependent torch.Generator
      lowercase_								= self.get_dummy_components()
      lowercase_								= StableUnCLIPImgaImgPipeline(**lowerCAmelCase_)
      lowercase_								= sd_pipe.to(lowerCAmelCase_)
      sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_)

      lowercase_								= self.get_dummy_inputs(lowerCAmelCase_)
      inputs.update({"""image_embeds""": None})
      lowercase_								= sd_pipe(**lowerCAmelCase_).images
      lowercase_								= image[0, -3:, -3:, -1]

      assert image.shape == (1, 3_2, 3_2, 3)
      lowercase_								= np.array([0.3_872, 0.7_224, 0.5_601, 0.4_741, 0.6_872, 0.5_814, 0.4_636, 0.3_867, 0.5_078])

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

    def 					_UpperCAmelCase	(				self			:				Union[str, Any]):







      """simple docstring"""







      lowercase_								= torch_device in ["""cpu""", """mps"""]

      self._test_attention_slicing_forward_pass(test_max_difference=lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      lowercase_								= torch_device in ["""cpu""", """mps"""]

      self._test_inference_batch_single_identical(test_max_difference=lowerCAmelCase_)




    @unittest.skipIf(
        torch_device != """cuda""" or not is_xformers_available() ,  reason="""XFormers attention is only available with CUDA and `xformers` installed""" ,  )
    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      self._test_xformers_attention_forwardGenerator_pass(test_max_difference=lowerCAmelCase_)







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

    def 					_UpperCAmelCase	(				self			:				Tuple):







      """simple docstring"""







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

    def 					_UpperCAmelCase	(				self			:				Tuple):







      """simple docstring"""







      lowercase_								= load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png""")

      lowercase_								= load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy""")

      lowercase_								= StableUnCLIPImgaImgPipeline.from_pretrained(
          """fusing/stable-unclip-2-1-l-img2img""" ,  torch_dtype=torch.floataa)
      pipe.to(lowerCAmelCase_)
      pipe.set_progress_bar_config(disable=lowerCAmelCase_)
      # stable unclip will oom when integration tests are run on a V100,
      # so turn on memory savings
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      lowercase_								= torch.Generator(device="""cpu""").manual_seed(0)
      lowercase_								= pipe(lowerCAmelCase_ ,  """anime turle""" ,  generator=lowerCAmelCase_ ,  output_type="""np""")

      lowercase_								= output.images[0]

      assert image.shape == (7_6_8, 7_6_8, 3)

      assert_mean_pixel_difference(lowerCAmelCase_ ,  lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				List[str]):







      """simple docstring"""







      lowercase_								= load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png""")

      lowercase_								= load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy""")

      lowercase_								= StableUnCLIPImgaImgPipeline.from_pretrained(
          """fusing/stable-unclip-2-1-h-img2img""" ,  torch_dtype=torch.floataa)
      pipe.to(lowerCAmelCase_)
      pipe.set_progress_bar_config(disable=lowerCAmelCase_)
      # stable unclip will oom when integration tests are run on a V100,
      # so turn on memory savings
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      lowercase_								= torch.Generator(device="""cpu""").manual_seed(0)
      lowercase_								= pipe(lowerCAmelCase_ ,  """anime turle""" ,  generator=lowerCAmelCase_ ,  output_type="""np""")

      lowercase_								= output.images[0]

      assert image.shape == (7_6_8, 7_6_8, 3)

      assert_mean_pixel_difference(lowerCAmelCase_ ,  lowerCAmelCase_)




    def 					_UpperCAmelCase	(				self			:				Dict):







      """simple docstring"""







      lowercase_								= load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png""")

      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

      lowercase_								= StableUnCLIPImgaImgPipeline.from_pretrained(
          """fusing/stable-unclip-2-1-h-img2img""" ,  torch_dtype=torch.floataa)
      lowercase_								= pipe.to(lowerCAmelCase_)
      pipe.set_progress_bar_config(disable=lowerCAmelCase_)
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      lowercase_								= pipe(
          lowerCAmelCase_ ,  """anime turtle""" ,  num_inference_steps=2 ,  output_type="""np""" ,  )

      lowercase_								= torch.cuda.max_memory_allocated()
      # make sure that less than 7 GB is allocated
      assert mem_bytes < 7 * 1_0**9






"""simple docstring"""
from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase						:		List[Any]			= logging.get_logger(__name__)

UpperCAmelCase						:		Optional[Any]			= {
    "alibaba-damo/mgp-str-base": "https://huggingface.co/alibaba-damo/mgp-str-base/resolve/main/config.json",
}


class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase   ):
    lowercase__      =   "mgp-str"

    def __init__(				self			:				Optional[Any] ,  lowerCAmelCase_			:				Optional[Any]=[3_2, 1_2_8] ,  lowerCAmelCase_			:				Union[str, Any]=4 ,  lowerCAmelCase_			:				Dict=3 ,  lowerCAmelCase_			:				List[str]=2_7 ,  lowerCAmelCase_			:				str=3_8 ,  lowerCAmelCase_			:				Optional[int]=5_0_2_5_7 ,  lowerCAmelCase_			:				Any=3_0_5_2_2 ,  lowerCAmelCase_			:				str=7_6_8 ,  lowerCAmelCase_			:				Any=1_2 ,  lowerCAmelCase_			:				Optional[Any]=1_2 ,  lowerCAmelCase_			:				Tuple=4.0 ,  lowerCAmelCase_			:				str=True ,  lowerCAmelCase_			:				List[str]=False ,  lowerCAmelCase_			:				List[Any]=1E-5 ,  lowerCAmelCase_			:				Any=0.0 ,  lowerCAmelCase_			:				int=0.0 ,  lowerCAmelCase_			:				str=0.0 ,  lowerCAmelCase_			:				Dict=False ,  lowerCAmelCase_			:				List[str]=0.02 ,  **lowerCAmelCase_			:				str ,  ):







      """simple docstring"""







      super().__init__(**lowerCAmelCase_)

      lowercase_								= image_size
      lowercase_								= patch_size
      lowercase_								= num_channels
      lowercase_								= max_token_length
      lowercase_								= num_character_labels
      lowercase_								= num_bpe_labels
      lowercase_								= num_wordpiece_labels
      lowercase_								= hidden_size
      lowercase_								= num_hidden_layers
      lowercase_								= num_attention_heads
      lowercase_								= mlp_ratio
      lowercase_								= distilled
      lowercase_								= layer_norm_eps
      lowercase_								= drop_rate
      lowercase_								= qkv_bias
      lowercase_								= attn_drop_rate
      lowercase_								= drop_path_rate
      lowercase_								= output_aa_attentions
      lowercase_								= initializer_range

