




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  _lowerCamelCase						(			_SCREAMING_SNAKE_CASE							,			_SCREAMING_SNAKE_CASE							,			_SCREAMING_SNAKE_CASE							,			unittest.TestCase   ):

      """simple docstring"""

      lowerCAmelCase__			: int    =					StableUnCLIPImgaImgPipeline
      lowerCAmelCase__			: Tuple    =					TEXT_GUIDED_IMAGE_VARIATION_PARAMS
      lowerCAmelCase__			: List[Any]    =					TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
      lowerCAmelCase__			: Optional[Any]    =					frozenset(
          []   )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
      lowerCAmelCase__			: List[Any]    =					frozenset([]   )







      def       snake_case	(     self					: Dict				):
          __UpperCamelCase  		=						32
          __UpperCamelCase  		=						embedder_hidden_size

          # image encoding components

          __UpperCamelCase  		=						CLIPImageProcessor(crop_size=32 ,       size=32				)

          torch.manual_seed(0				)
          __UpperCamelCase  		=						CLIPVisionModelWithProjection(
              CLIPVisionConfig(
                  hidden_size=snake_case ,       projection_dim=snake_case ,       num_hidden_layers=5 ,       num_attention_heads=4 ,       image_size=32 ,       intermediate_size=37 ,       patch_size=1 ,       )				)

          # regular denoising components

          torch.manual_seed(0				)
          __UpperCamelCase  		=						StableUnCLIPImageNormalizer(embedding_dim=snake_case				)
          __UpperCamelCase  		=						DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''				)

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

          torch.manual_seed(0				)
          __UpperCamelCase  		=						CLIPTextModel(
              CLIPTextConfig(
                  bos_token_id=0 ,       eos_token_id=2 ,       hidden_size=snake_case ,       projection_dim=32 ,       intermediate_size=37 ,       layer_norm_eps=1E-05 ,       num_attention_heads=4 ,       num_hidden_layers=5 ,       pad_token_id=1 ,       vocab_size=1000 ,       )				)

          torch.manual_seed(0				)
          __UpperCamelCase  		=						UNetaDConditionModel(
              sample_size=32 ,       in_channels=4 ,       out_channels=4 ,       down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''') ,       up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''') ,       block_out_channels=(32, 64) ,       attention_head_dim=(2, 4) ,       class_embed_type='''projection''' ,       projection_class_embeddings_input_dim=embedder_projection_dim * 2 ,       cross_attention_dim=snake_case ,       layers_per_block=1 ,       upcast_attention=snake_case ,       use_linear_projection=snake_case ,       )

          torch.manual_seed(0				)
          __UpperCamelCase  		=						DDIMScheduler(
              beta_schedule='''scaled_linear''' ,       beta_start=0.0_00_85 ,       beta_end=0.0_12 ,       prediction_type='''v_prediction''' ,       set_alpha_to_one=snake_case ,       steps_offset=1 ,       )

          torch.manual_seed(0				)
          __UpperCamelCase  		=						AutoencoderKL()

          __UpperCamelCase  		=						{
              # 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       snake_case	(     self					: Dict ,       snake_case					: Optional[int] ,       snake_case					: List[str]=0 ,       snake_case					: List[str]=True				):
          if str(snake_case				).startswith('''mps'''				):
              __UpperCamelCase  		=						torch.manual_seed(snake_case				)
          else:
              __UpperCamelCase  		=						torch.Generator(device=snake_case				).manual_seed(snake_case				)

          __UpperCamelCase  		=						floats_tensor((1, 3, 32, 32) ,       rng=random.Random(snake_case				)				).to(snake_case				)

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

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







      @skip_mps
      def       snake_case	(     self					: Optional[Any]				):
          __UpperCamelCase  		=						'''cpu'''  # ensure determinism for the device-dependent torch.Generator
          __UpperCamelCase  		=						self.get_dummy_components()
          __UpperCamelCase  		=						StableUnCLIPImgaImgPipeline(**snake_case				)
          __UpperCamelCase  		=						sd_pipe.to(snake_case				)
          sd_pipe.set_progress_bar_config(disable=snake_case				)

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

          assert image.shape == (1, 32, 32, 3)
          __UpperCamelCase  		=						np.array([0.38_72, 0.72_24, 0.56_01, 0.47_41, 0.68_72, 0.58_14, 0.46_36, 0.38_67, 0.50_78]				)

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







      def       snake_case	(     self					: List[Any]				):
          __UpperCamelCase  		=						torch_device in ['''cpu''', '''mps''']

          self._test_attention_slicing_forward_pass(test_max_difference=snake_case				)







      def       snake_case	(     self					: List[str]				):
          __UpperCamelCase  		=						torch_device in ['''cpu''', '''mps''']

          self._test_inference_batch_single_identical(test_max_difference=snake_case				)







      @unittest.skipIf(
          torch_device != '''cuda''' or not is_xformers_available() ,       reason='''XFormers attention is only available with CUDA and `xformers` installed''' ,       )
      def       snake_case	(     self					: Dict				):
          self._test_xformers_attention_forwardGenerator_pass(test_max_difference=snake_case				)







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

      """simple docstring"""

      def       snake_case	(     self					: Union[str, Any]				):
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()







      def       snake_case	(     self					: Tuple				):
          __UpperCamelCase  		=						load_image(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''				)

          __UpperCamelCase  		=						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'''				)

          __UpperCamelCase  		=						StableUnCLIPImgaImgPipeline.from_pretrained(
              '''fusing/stable-unclip-2-1-l-img2img''' ,       torch_dtype=torch.floataa				)
          pipe.to(snake_case				)
          pipe.set_progress_bar_config(disable=snake_case				)
          # 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()

          __UpperCamelCase  		=						torch.Generator(device='''cpu'''				).manual_seed(0				)
          __UpperCamelCase  		=						pipe(snake_case ,       '''anime turle''' ,       generator=snake_case ,       output_type='''np'''				)

          __UpperCamelCase  		=						output.images[0]

          assert image.shape == (768, 768, 3)

          assert_mean_pixel_difference(snake_case ,       snake_case				)







      def       snake_case	(     self					: Any				):
          __UpperCamelCase  		=						load_image(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''				)

          __UpperCamelCase  		=						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'''				)

          __UpperCamelCase  		=						StableUnCLIPImgaImgPipeline.from_pretrained(
              '''fusing/stable-unclip-2-1-h-img2img''' ,       torch_dtype=torch.floataa				)
          pipe.to(snake_case				)
          pipe.set_progress_bar_config(disable=snake_case				)
          # 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()

          __UpperCamelCase  		=						torch.Generator(device='''cpu'''				).manual_seed(0				)
          __UpperCamelCase  		=						pipe(snake_case ,       '''anime turle''' ,       generator=snake_case ,       output_type='''np'''				)

          __UpperCamelCase  		=						output.images[0]

          assert image.shape == (768, 768, 3)

          assert_mean_pixel_difference(snake_case ,       snake_case				)







      def       snake_case	(     self					: List[str]				):
          __UpperCamelCase  		=						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()

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

          __UpperCamelCase  		=						pipe(
              snake_case ,       '''anime turtle''' ,       num_inference_steps=2 ,       output_type='''np''' ,       )

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












from pathlib import PurePosixPath
from typing import Optional

import fsspec
from fsspec import AbstractFileSystem
from huggingface_hub.hf_api import DatasetInfo

from ..utils.file_utils import get_authentication_headers_for_url
from ..utils.hub import hf_hub_url

class  _lowerCamelCase						(			_SCREAMING_SNAKE_CASE   ):

      """simple docstring"""

      lowerCAmelCase__			: Any    =					""
      lowerCAmelCase__			: Any    =					"hf-legacy"  # "hf://"" is reserved for hffs







      def __init__(     self					: str ,       snake_case					: Optional[DatasetInfo] = None ,       snake_case					: Optional[str] = None ,       **snake_case					: List[Any] ,       ):
          super().__init__(self ,       **snake_case				)
          __UpperCamelCase  		=						repo_info
          __UpperCamelCase  		=						token
          __UpperCamelCase  		=						None







      def       snake_case	(     self					: List[Any]				):
          if self.dir_cache is None:
              __UpperCamelCase  		=						{}
              for hf_file in self.repo_info.siblings:
                  # TODO(QL): add sizes
                  __UpperCamelCase  		=						{
                      '''name''': hf_file.rfilename,
                      '''size''': None,
                      '''type''': '''file''',
                  }
                  self.dir_cache.update(
                      {
                          str(snake_case				): {'''name''': str(snake_case				), '''size''': None, '''type''': '''directory'''}
                          for d in list(PurePosixPath(hf_file.rfilename				).parents				)[:-1]
                      }				)







      def       snake_case	(     self					: Dict ,       snake_case					: str ,       snake_case					: str = "rb" ,       **snake_case					: Union[str, Any] ,       ):
          if not isinstance(self.repo_info ,       snake_case				):
              raise NotImplementedError(F"Open is only implemented for dataset repositories, but got {self.repo_info}"				)
          __UpperCamelCase  		=						hf_hub_url(self.repo_info.id ,       snake_case ,       revision=self.repo_info.sha				)
          return fsspec.open(
              snake_case ,       mode=snake_case ,       headers=get_authentication_headers_for_url(snake_case ,       use_auth_token=self.token				) ,       client_kwargs={'''trust_env''': True} ,       ).open()







      def       snake_case	(     self					: Optional[Any] ,       snake_case					: Tuple ,       **snake_case					: List[Any]				):
          self._get_dirs()
          __UpperCamelCase  		=						self._strip_protocol(snake_case				)
          if path in self.dir_cache:
              return self.dir_cache[path]
          else:
              raise FileNotFoundError(snake_case				)







      def       snake_case	(     self					: List[str] ,       snake_case					: int ,       snake_case					: Tuple=False ,       **snake_case					: Dict				):
          self._get_dirs()
          __UpperCamelCase  		=						PurePosixPath(path.strip('''/'''				)				)
          __UpperCamelCase  		=						{}
          for p, f in self.dir_cache.items():
              __UpperCamelCase  		=						PurePosixPath(p.strip('''/'''				)				)
              __UpperCamelCase  		=						p.parent
              if root == path:
                  __UpperCamelCase  		=						f
          __UpperCamelCase  		=						list(paths.values()				)
          if detail:
              return out
          else:
              return sorted(f['''name'''] for f in out				)






