import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoImageProcessor, ViTImageProcessor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / 'utils'))

from test_module.custom_image_processing import CustomImageProcessor  # noqa E402


UpperCamelCase__    							= get_tests_dir('fixtures')





class 					A							(    unittest.TestCase					):


   def       lowercase_				(self					:						Dict		)						->    Optional[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  mock.Mock()
         UpperCAmelCase__								=  5_0_0
         UpperCAmelCase__								=  {}
         UpperCAmelCase__								=  HTTPError
         UpperCAmelCase__								=  {}

         # Download this model to make sure it's in the cache.
         UpperCAmelCase__								=  ViTImageProcessor.from_pretrained("hf-internal-testing/tiny-random-vit"		)
         # Under the mock environment we get a 500 error when trying to reach the model.
         with mock.patch("requests.Session.request"	,  return_value=__UpperCAmelCase		) as mock_head:
               UpperCAmelCase__								=  ViTImageProcessor.from_pretrained("hf-internal-testing/tiny-random-vit"		)
               # This check we did call the fake head request
               mock_head.assert_called()


   def       lowercase_				(self					:						int		)						->    int:
         """simple docstring"""



         UpperCAmelCase__								=  ViTImageProcessor.from_pretrained(
             "https://huggingface.co/hf-internal-testing/tiny-random-vit/resolve/main/preprocessor_config.json"		)





   def       lowercase_				(self					:						int		)						->    Optional[Any]:
         """simple docstring"""



         with self.assertRaises(__UpperCAmelCase		):
               # config is in subfolder, the following should not work without specifying the subfolder
               UpperCAmelCase__								=  AutoImageProcessor.from_pretrained("hf-internal-testing/stable-diffusion-all-variants"		)

         UpperCAmelCase__								=  AutoImageProcessor.from_pretrained(
             "hf-internal-testing/stable-diffusion-all-variants"	,  subfolder="feature_extractor"		)

         self.assertIsNotNone(__UpperCAmelCase		)





@is_staging_test
class 					A							(    unittest.TestCase					):


   @classmethod
   def       lowercase_				(cls					:						Tuple		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  TOKEN
         HfFolder.save_token(__UpperCAmelCase		)


   @classmethod
   def       lowercase_				(cls					:						int		)						->    int:
         """simple docstring"""



         try:
               delete_repo(token=cls._token	,  repo_id="test-image-processor"		)
         except HTTPError:
               pass

         try:
               delete_repo(token=cls._token	,  repo_id="valid_org/test-image-processor-org"		)
         except HTTPError:
               pass

         try:
               delete_repo(token=cls._token	,  repo_id="test-dynamic-image-processor"		)
         except HTTPError:
               pass


   def       lowercase_				(self					:						Any		)						->    str:
         """simple docstring"""



         UpperCAmelCase__								=  ViTImageProcessor.from_pretrained(__UpperCAmelCase		)
         image_processor.push_to_hub("test-image-processor"	,  use_auth_token=self._token		)

         UpperCAmelCase__								=  ViTImageProcessor.from_pretrained(f"""{USER}/test-image-processor"""		)
         for k, v in image_processor.__dict__.items():
               self.assertEqual(__UpperCAmelCase	,  getattr(__UpperCAmelCase	,  __UpperCAmelCase		)		)

         # Reset repo
         delete_repo(token=self._token	,  repo_id="test-image-processor"		)

         # Push to hub via save_pretrained
         with tempfile.TemporaryDirectory() as tmp_dir:
               image_processor.save_pretrained(
                   __UpperCAmelCase	,  repo_id="test-image-processor"	,  push_to_hub=__UpperCAmelCase	,  use_auth_token=self._token		)

         UpperCAmelCase__								=  ViTImageProcessor.from_pretrained(f"""{USER}/test-image-processor"""		)
         for k, v in image_processor.__dict__.items():
               self.assertEqual(__UpperCAmelCase	,  getattr(__UpperCAmelCase	,  __UpperCAmelCase		)		)


   def       lowercase_				(self					:						Dict		)						->    Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__								=  ViTImageProcessor.from_pretrained(__UpperCAmelCase		)
         image_processor.push_to_hub("valid_org/test-image-processor"	,  use_auth_token=self._token		)

         UpperCAmelCase__								=  ViTImageProcessor.from_pretrained("valid_org/test-image-processor"		)
         for k, v in image_processor.__dict__.items():
               self.assertEqual(__UpperCAmelCase	,  getattr(__UpperCAmelCase	,  __UpperCAmelCase		)		)

         # Reset repo
         delete_repo(token=self._token	,  repo_id="valid_org/test-image-processor"		)

         # Push to hub via save_pretrained
         with tempfile.TemporaryDirectory() as tmp_dir:
               image_processor.save_pretrained(
                   __UpperCAmelCase	,  repo_id="valid_org/test-image-processor-org"	,  push_to_hub=__UpperCAmelCase	,  use_auth_token=self._token		)

         UpperCAmelCase__								=  ViTImageProcessor.from_pretrained("valid_org/test-image-processor-org"		)
         for k, v in image_processor.__dict__.items():
               self.assertEqual(__UpperCAmelCase	,  getattr(__UpperCAmelCase	,  __UpperCAmelCase		)		)





   def       lowercase_				(self					:						Optional[Any]		)						->    str:
         """simple docstring"""



         CustomImageProcessor.register_for_auto_class()
         UpperCAmelCase__								=  CustomImageProcessor.from_pretrained(__UpperCAmelCase		)

         image_processor.push_to_hub("test-dynamic-image-processor"	,  use_auth_token=self._token		)

         # This has added the proper auto_map field to the config
         self.assertDictEqual(
             image_processor.auto_map	,  {"AutoImageProcessor": "custom_image_processing.CustomImageProcessor"}	,  )

         UpperCAmelCase__								=  AutoImageProcessor.from_pretrained(
             f"""{USER}/test-dynamic-image-processor"""	,  trust_remote_code=__UpperCAmelCase		)
         # Can't make an isinstance check because the new_image_processor is from the CustomImageProcessor class of a dynamic module
         self.assertEqual(new_image_processor.__class__.__name__	,  "CustomImageProcessor"		)




from typing import Any, Callable, Dict, List, Optional, Union

import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker


UpperCamelCase__    							= 'CompVis/stable-diffusion-v1-1'
UpperCamelCase__    							= 'CompVis/stable-diffusion-v1-2'
UpperCamelCase__    							= 'CompVis/stable-diffusion-v1-3'
UpperCamelCase__    							= 'CompVis/stable-diffusion-v1-4'





class 					A							(    UpperCAmelCase_					):


   def __init__(self					:						Union[str, Any]	,  __UpperCAmelCase					:						AutoencoderKL	,  __UpperCAmelCase					:						CLIPTextModel	,  __UpperCAmelCase					:						CLIPTokenizer	,  __UpperCAmelCase					:						UNetaDConditionModel	,  __UpperCAmelCase					:						Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]	,  __UpperCAmelCase					:						StableDiffusionSafetyChecker	,  __UpperCAmelCase					:						CLIPImageProcessor	,  __UpperCAmelCase					:						bool = True	,  )						->    Tuple:
         """simple docstring"""



         super()._init_()

         UpperCAmelCase__								=  StableDiffusionPipeline.from_pretrained(__UpperCAmelCase		)
         UpperCAmelCase__								=  StableDiffusionPipeline.from_pretrained(__UpperCAmelCase		)
         UpperCAmelCase__								=  StableDiffusionPipeline.from_pretrained(__UpperCAmelCase		)
         UpperCAmelCase__								=  StableDiffusionPipeline(
             vae=__UpperCAmelCase	,  text_encoder=__UpperCAmelCase	,  tokenizer=__UpperCAmelCase	,  unet=__UpperCAmelCase	,  scheduler=__UpperCAmelCase	,  safety_checker=__UpperCAmelCase	,  feature_extractor=__UpperCAmelCase	,  requires_safety_checker=__UpperCAmelCase	,  )

         self.register_modules(pipelinea=self.pipea	,  pipelinea=self.pipea	,  pipelinea=self.pipea	,  pipelinea=self.pipea		)


   @property
   def       lowercase_				(self					:						int		)						->    Dict[str, Any]:
         """simple docstring"""



         return {k: getattr(self	,  __UpperCAmelCase		) for k in self.config.keys() if not k.startswith("_"		)}


   def       lowercase_				(self					:						Union[str, Any]	,  __UpperCAmelCase					:						Optional[Union[str, int]] = "auto"		)						->    Optional[int]:
         """simple docstring"""



         if slice_size == "auto":
               # half the attention head size is usually a good trade-off between
               # speed and memory
               UpperCAmelCase__								=  self.unet.config.attention_head_dim // 2
         self.unet.set_attention_slice(__UpperCAmelCase		)


   def       lowercase_				(self					:						int		)						->    Optional[int]:
         """simple docstring"""



         self.enable_attention_slicing(__UpperCAmelCase		)


   @torch.no_grad()
   def       lowercase_				(self					:						Optional[int]	,  __UpperCAmelCase					:						Union[str, List[str]]	,  __UpperCAmelCase					:						int = 5_1_2	,  __UpperCAmelCase					:						int = 5_1_2	,  __UpperCAmelCase					:						int = 5_0	,  __UpperCAmelCase					:						float = 7.5	,  __UpperCAmelCase					:						Optional[Union[str, List[str]]] = None	,  __UpperCAmelCase					:						Optional[int] = 1	,  __UpperCAmelCase					:						float = 0.0	,  __UpperCAmelCase					:						Optional[torch.Generator] = None	,  __UpperCAmelCase					:						Optional[torch.FloatTensor] = None	,  __UpperCAmelCase					:						Optional[str] = "pil"	,  __UpperCAmelCase					:						bool = True	,  __UpperCAmelCase					:						Optional[Callable[[int, int, torch.FloatTensor], None]] = None	,  __UpperCAmelCase					:						int = 1	,  **__UpperCAmelCase					:						Any	,  )						->    Dict:
         """simple docstring"""



         return self.pipea(
             prompt=__UpperCAmelCase	,  height=__UpperCAmelCase	,  width=__UpperCAmelCase	,  num_inference_steps=__UpperCAmelCase	,  guidance_scale=__UpperCAmelCase	,  negative_prompt=__UpperCAmelCase	,  num_images_per_prompt=__UpperCAmelCase	,  eta=__UpperCAmelCase	,  generator=__UpperCAmelCase	,  latents=__UpperCAmelCase	,  output_type=__UpperCAmelCase	,  return_dict=__UpperCAmelCase	,  callback=__UpperCAmelCase	,  callback_steps=__UpperCAmelCase	,  **__UpperCAmelCase	,  )


   @torch.no_grad()
   def       lowercase_				(self					:						List[str]	,  __UpperCAmelCase					:						Union[str, List[str]]	,  __UpperCAmelCase					:						int = 5_1_2	,  __UpperCAmelCase					:						int = 5_1_2	,  __UpperCAmelCase					:						int = 5_0	,  __UpperCAmelCase					:						float = 7.5	,  __UpperCAmelCase					:						Optional[Union[str, List[str]]] = None	,  __UpperCAmelCase					:						Optional[int] = 1	,  __UpperCAmelCase					:						float = 0.0	,  __UpperCAmelCase					:						Optional[torch.Generator] = None	,  __UpperCAmelCase					:						Optional[torch.FloatTensor] = None	,  __UpperCAmelCase					:						Optional[str] = "pil"	,  __UpperCAmelCase					:						bool = True	,  __UpperCAmelCase					:						Optional[Callable[[int, int, torch.FloatTensor], None]] = None	,  __UpperCAmelCase					:						int = 1	,  **__UpperCAmelCase					:						Optional[Any]	,  )						->    Any:
         """simple docstring"""



         return self.pipea(
             prompt=__UpperCAmelCase	,  height=__UpperCAmelCase	,  width=__UpperCAmelCase	,  num_inference_steps=__UpperCAmelCase	,  guidance_scale=__UpperCAmelCase	,  negative_prompt=__UpperCAmelCase	,  num_images_per_prompt=__UpperCAmelCase	,  eta=__UpperCAmelCase	,  generator=__UpperCAmelCase	,  latents=__UpperCAmelCase	,  output_type=__UpperCAmelCase	,  return_dict=__UpperCAmelCase	,  callback=__UpperCAmelCase	,  callback_steps=__UpperCAmelCase	,  **__UpperCAmelCase	,  )


   @torch.no_grad()
   def       lowercase_				(self					:						List[str]	,  __UpperCAmelCase					:						Union[str, List[str]]	,  __UpperCAmelCase					:						int = 5_1_2	,  __UpperCAmelCase					:						int = 5_1_2	,  __UpperCAmelCase					:						int = 5_0	,  __UpperCAmelCase					:						float = 7.5	,  __UpperCAmelCase					:						Optional[Union[str, List[str]]] = None	,  __UpperCAmelCase					:						Optional[int] = 1	,  __UpperCAmelCase					:						float = 0.0	,  __UpperCAmelCase					:						Optional[torch.Generator] = None	,  __UpperCAmelCase					:						Optional[torch.FloatTensor] = None	,  __UpperCAmelCase					:						Optional[str] = "pil"	,  __UpperCAmelCase					:						bool = True	,  __UpperCAmelCase					:						Optional[Callable[[int, int, torch.FloatTensor], None]] = None	,  __UpperCAmelCase					:						int = 1	,  **__UpperCAmelCase					:						Any	,  )						->    Dict:
         """simple docstring"""



         return self.pipea(
             prompt=__UpperCAmelCase	,  height=__UpperCAmelCase	,  width=__UpperCAmelCase	,  num_inference_steps=__UpperCAmelCase	,  guidance_scale=__UpperCAmelCase	,  negative_prompt=__UpperCAmelCase	,  num_images_per_prompt=__UpperCAmelCase	,  eta=__UpperCAmelCase	,  generator=__UpperCAmelCase	,  latents=__UpperCAmelCase	,  output_type=__UpperCAmelCase	,  return_dict=__UpperCAmelCase	,  callback=__UpperCAmelCase	,  callback_steps=__UpperCAmelCase	,  **__UpperCAmelCase	,  )


   @torch.no_grad()
   def       lowercase_				(self					:						Union[str, Any]	,  __UpperCAmelCase					:						Union[str, List[str]]	,  __UpperCAmelCase					:						int = 5_1_2	,  __UpperCAmelCase					:						int = 5_1_2	,  __UpperCAmelCase					:						int = 5_0	,  __UpperCAmelCase					:						float = 7.5	,  __UpperCAmelCase					:						Optional[Union[str, List[str]]] = None	,  __UpperCAmelCase					:						Optional[int] = 1	,  __UpperCAmelCase					:						float = 0.0	,  __UpperCAmelCase					:						Optional[torch.Generator] = None	,  __UpperCAmelCase					:						Optional[torch.FloatTensor] = None	,  __UpperCAmelCase					:						Optional[str] = "pil"	,  __UpperCAmelCase					:						bool = True	,  __UpperCAmelCase					:						Optional[Callable[[int, int, torch.FloatTensor], None]] = None	,  __UpperCAmelCase					:						int = 1	,  **__UpperCAmelCase					:						Union[str, Any]	,  )						->    List[str]:
         """simple docstring"""



         return self.pipea(
             prompt=__UpperCAmelCase	,  height=__UpperCAmelCase	,  width=__UpperCAmelCase	,  num_inference_steps=__UpperCAmelCase	,  guidance_scale=__UpperCAmelCase	,  negative_prompt=__UpperCAmelCase	,  num_images_per_prompt=__UpperCAmelCase	,  eta=__UpperCAmelCase	,  generator=__UpperCAmelCase	,  latents=__UpperCAmelCase	,  output_type=__UpperCAmelCase	,  return_dict=__UpperCAmelCase	,  callback=__UpperCAmelCase	,  callback_steps=__UpperCAmelCase	,  **__UpperCAmelCase	,  )





   @torch.no_grad()
   def       lowercase_				(self					:						int	,  __UpperCAmelCase					:						Union[str, List[str]]	,  __UpperCAmelCase					:						int = 5_1_2	,  __UpperCAmelCase					:						int = 5_1_2	,  __UpperCAmelCase					:						int = 5_0	,  __UpperCAmelCase					:						float = 7.5	,  __UpperCAmelCase					:						Optional[Union[str, List[str]]] = None	,  __UpperCAmelCase					:						Optional[int] = 1	,  __UpperCAmelCase					:						float = 0.0	,  __UpperCAmelCase					:						Optional[torch.Generator] = None	,  __UpperCAmelCase					:						Optional[torch.FloatTensor] = None	,  __UpperCAmelCase					:						Optional[str] = "pil"	,  __UpperCAmelCase					:						bool = True	,  __UpperCAmelCase					:						Optional[Callable[[int, int, torch.FloatTensor], None]] = None	,  __UpperCAmelCase					:						int = 1	,  **__UpperCAmelCase					:						Any	,  )						->    Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__								=  "cuda" if torch.cuda.is_available() else "cpu"
         self.to(__UpperCAmelCase		)

         # Checks if the height and width are divisible by 8 or not
         if height % 8 != 0 or width % 8 != 0:
               raise ValueError(f"""`height` and `width` must be divisible by 8 but are {height} and {width}."""		)

         # Get first result from Stable Diffusion Checkpoint v1.1
         UpperCAmelCase__								=  self.textaimg_sda_a(
             prompt=__UpperCAmelCase	,  height=__UpperCAmelCase	,  width=__UpperCAmelCase	,  num_inference_steps=__UpperCAmelCase	,  guidance_scale=__UpperCAmelCase	,  negative_prompt=__UpperCAmelCase	,  num_images_per_prompt=__UpperCAmelCase	,  eta=__UpperCAmelCase	,  generator=__UpperCAmelCase	,  latents=__UpperCAmelCase	,  output_type=__UpperCAmelCase	,  return_dict=__UpperCAmelCase	,  callback=__UpperCAmelCase	,  callback_steps=__UpperCAmelCase	,  **__UpperCAmelCase	,  )

         # Get first result from Stable Diffusion Checkpoint v1.2
         UpperCAmelCase__								=  self.textaimg_sda_a(
             prompt=__UpperCAmelCase	,  height=__UpperCAmelCase	,  width=__UpperCAmelCase	,  num_inference_steps=__UpperCAmelCase	,  guidance_scale=__UpperCAmelCase	,  negative_prompt=__UpperCAmelCase	,  num_images_per_prompt=__UpperCAmelCase	,  eta=__UpperCAmelCase	,  generator=__UpperCAmelCase	,  latents=__UpperCAmelCase	,  output_type=__UpperCAmelCase	,  return_dict=__UpperCAmelCase	,  callback=__UpperCAmelCase	,  callback_steps=__UpperCAmelCase	,  **__UpperCAmelCase	,  )

         # Get first result from Stable Diffusion Checkpoint v1.3
         UpperCAmelCase__								=  self.textaimg_sda_a(
             prompt=__UpperCAmelCase	,  height=__UpperCAmelCase	,  width=__UpperCAmelCase	,  num_inference_steps=__UpperCAmelCase	,  guidance_scale=__UpperCAmelCase	,  negative_prompt=__UpperCAmelCase	,  num_images_per_prompt=__UpperCAmelCase	,  eta=__UpperCAmelCase	,  generator=__UpperCAmelCase	,  latents=__UpperCAmelCase	,  output_type=__UpperCAmelCase	,  return_dict=__UpperCAmelCase	,  callback=__UpperCAmelCase	,  callback_steps=__UpperCAmelCase	,  **__UpperCAmelCase	,  )

         # Get first result from Stable Diffusion Checkpoint v1.4
         UpperCAmelCase__								=  self.textaimg_sda_a(
             prompt=__UpperCAmelCase	,  height=__UpperCAmelCase	,  width=__UpperCAmelCase	,  num_inference_steps=__UpperCAmelCase	,  guidance_scale=__UpperCAmelCase	,  negative_prompt=__UpperCAmelCase	,  num_images_per_prompt=__UpperCAmelCase	,  eta=__UpperCAmelCase	,  generator=__UpperCAmelCase	,  latents=__UpperCAmelCase	,  output_type=__UpperCAmelCase	,  return_dict=__UpperCAmelCase	,  callback=__UpperCAmelCase	,  callback_steps=__UpperCAmelCase	,  **__UpperCAmelCase	,  )

         # Get all result images into a single list and pass it via StableDiffusionPipelineOutput for final result
         return StableDiffusionPipelineOutput([resa[0], resa[0], resa[0], resa[0]]		)



