

import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin


if is_torch_available():
      import torch

if is_vision_available():
      from PIL import Image

      from transformers import SwinaSRImageProcessor
      from transformers.image_transforms import get_image_size


class     lowerCamelCase_   ( unittest.TestCase):




 def __init__( self     ,     _UpperCAmelCase     ,     _UpperCAmelCase=7     ,     _UpperCAmelCase=3     ,     _UpperCAmelCase=18     ,     _UpperCAmelCase=30     ,     _UpperCAmelCase=400     ,     _UpperCAmelCase=True     ,     _UpperCAmelCase=1 / 255     ,     _UpperCAmelCase=True     ,     _UpperCAmelCase=8     ,     ):
   __UpperCAmelCase  :     Optional[int]		=    parent
   __UpperCAmelCase  :     Dict		=    batch_size
   __UpperCAmelCase  :     Dict		=    num_channels
   __UpperCAmelCase  :     List[str]		=    image_size
   __UpperCAmelCase  :     Dict		=    min_resolution
   __UpperCAmelCase  :     Any		=    max_resolution
   __UpperCAmelCase  :     List[str]		=    do_rescale
   __UpperCAmelCase  :     str		=    rescale_factor
   __UpperCAmelCase  :     Any		=    do_pad
   __UpperCAmelCase  :     Dict		=    pad_size




 def 						a					( self       ):
   return {
       "do_rescale": self.do_rescale,
       "rescale_factor": self.rescale_factor,
       "do_pad": self.do_pad,
       "pad_size": self.pad_size,
   }




 def 						a					( self     ,     _UpperCAmelCase=False     ,     _UpperCAmelCase=False     ,     _UpperCAmelCase=False       ):
   assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

   if equal_resolution:
     __UpperCAmelCase  :     Optional[int]		=    []
     for i in range(self.batch_size       ):
       image_inputs.append(
           np.random.randint(
               255     ,     size=(self.num_channels, self.max_resolution, self.max_resolution)     ,     dtype=np.uinta       )       )
   else:
     __UpperCAmelCase  :     Optional[int]		=    []
     for i in range(self.batch_size       ):
       __UpperCAmelCase, __UpperCAmelCase  :     List[str]		=    np.random.choice(np.arange(self.min_resolution     ,     self.max_resolution       )     ,     2       )
       image_inputs.append(np.random.randint(255     ,     size=(self.num_channels, width, height)     ,     dtype=np.uinta       )       )

   if not numpify and not torchify:
     # PIL expects the channel dimension as last dimension
     __UpperCAmelCase  :     Optional[Any]		=    [Image.fromarray(np.moveaxis(_UpperCAmelCase     ,     0     ,     -1       )       ) for x in image_inputs]

   if torchify:
     __UpperCAmelCase  :     List[Any]		=    [torch.from_numpy(_UpperCAmelCase       ) for x in image_inputs]

   return image_inputs



@require_torch
@require_vision
class     lowerCamelCase_   ( _lowerCAmelCase  ,							unittest.TestCase):
 _lowerCAmelCase       		=  SwinaSRImageProcessor if is_vision_available() else None




 def 						a					( self       ):
   __UpperCAmelCase  :     Optional[int]		=    SwinaSRImageProcessingTester(self       )




 @property
 def 						a					( self       ):
   return self.image_processor_tester.prepare_image_processor_dict()




 def 						a					( self       ):
   __UpperCAmelCase  :     Union[str, Any]		=    self.image_processing_class(**self.image_processor_dict       )
   self.assertTrue(hasattr(_UpperCAmelCase     ,     'do_rescale'       )       )
   self.assertTrue(hasattr(_UpperCAmelCase     ,     'rescale_factor'       )       )
   self.assertTrue(hasattr(_UpperCAmelCase     ,     'do_pad'       )       )
   self.assertTrue(hasattr(_UpperCAmelCase     ,     'pad_size'       )       )




 def 						a					( self       ):
   pass




 def 						a					( self     ,     _UpperCAmelCase       ):
   __UpperCAmelCase, __UpperCAmelCase  :     str		=    get_image_size(_UpperCAmelCase       )
   __UpperCAmelCase  :     Union[str, Any]		=    self.image_processor_tester.pad_size

   __UpperCAmelCase  :     List[str]		=    (old_height // size + 1) * size - old_height
   __UpperCAmelCase  :     List[Any]		=    (old_width // size + 1) * size - old_width
   return old_height + pad_height, old_width + pad_width




 def 						a					( self       ):
   # Initialize image_processor
   __UpperCAmelCase  :     str		=    self.image_processing_class(**self.image_processor_dict       )
   # create random PIL images
   __UpperCAmelCase  :     int		=    self.image_processor_tester.prepare_inputs(equal_resolution=_UpperCAmelCase       )
   for image in image_inputs:
     self.assertIsInstance(_UpperCAmelCase     ,     Image.Image       )

   # Test not batched input
   __UpperCAmelCase  :     Optional[int]		=    image_processor(image_inputs[0]     ,     return_tensors='pt'       ).pixel_values
   __UpperCAmelCase, __UpperCAmelCase  :     Any		=    self.calculate_expected_size(np.array(image_inputs[0]       )       )
   self.assertEqual(
       encoded_images.shape     ,     (
           1,
           self.image_processor_tester.num_channels,
           expected_height,
           expected_width,
       )     ,     )




 def 						a					( self       ):
   # Initialize image_processor
   __UpperCAmelCase  :     Any		=    self.image_processing_class(**self.image_processor_dict       )
   # create random numpy tensors
   __UpperCAmelCase  :     Optional[Any]		=    self.image_processor_tester.prepare_inputs(equal_resolution=_UpperCAmelCase     ,     numpify=_UpperCAmelCase       )
   for image in image_inputs:
     self.assertIsInstance(_UpperCAmelCase     ,     np.ndarray       )

   # Test not batched input
   __UpperCAmelCase  :     str		=    image_processor(image_inputs[0]     ,     return_tensors='pt'       ).pixel_values
   __UpperCAmelCase, __UpperCAmelCase  :     List[str]		=    self.calculate_expected_size(image_inputs[0]       )
   self.assertEqual(
       encoded_images.shape     ,     (
           1,
           self.image_processor_tester.num_channels,
           expected_height,
           expected_width,
       )     ,     )




 def 						a					( self       ):
   # Initialize image_processor
   __UpperCAmelCase  :     Union[str, Any]		=    self.image_processing_class(**self.image_processor_dict       )
   # create random PyTorch tensors
   __UpperCAmelCase  :     str		=    self.image_processor_tester.prepare_inputs(equal_resolution=_UpperCAmelCase     ,     torchify=_UpperCAmelCase       )
   for image in image_inputs:
     self.assertIsInstance(_UpperCAmelCase     ,     torch.Tensor       )

   # Test not batched input
   __UpperCAmelCase  :     Dict		=    image_processor(image_inputs[0]     ,     return_tensors='pt'       ).pixel_values
   __UpperCAmelCase, __UpperCAmelCase  :     Optional[int]		=    self.calculate_expected_size(image_inputs[0]       )
   self.assertEqual(
       encoded_images.shape     ,     (
           1,
           self.image_processor_tester.num_channels,
           expected_height,
           expected_width,
       )     ,     )








import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionDualGuidedPipeline
from diffusers.utils.testing_utils import load_image, nightly, require_torch_gpu, torch_device


_SCREAMING_SNAKE_CASE          = False


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




 def 						a					( self       ):
   # clean up the VRAM after each test
   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()




 def 						a					( self       ):
   __UpperCAmelCase  :     Dict		=    VersatileDiffusionDualGuidedPipeline.from_pretrained('shi-labs/versatile-diffusion'       )
   # remove text_unet
   pipe.remove_unused_weights()
   pipe.to(_UpperCAmelCase       )
   pipe.set_progress_bar_config(disable=_UpperCAmelCase       )

   __UpperCAmelCase  :     Any		=    load_image(
       'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'       )

   __UpperCAmelCase  :     Optional[Any]		=    torch.manual_seed(0       )
   __UpperCAmelCase  :     Union[str, Any]		=    pipe(
       prompt='first prompt'     ,     image=_UpperCAmelCase     ,     text_to_image_strength=0.7_5     ,     generator=_UpperCAmelCase     ,     guidance_scale=7.5     ,     num_inference_steps=2     ,     output_type='numpy'     ,     ).images

   with tempfile.TemporaryDirectory() as tmpdirname:
     pipe.save_pretrained(_UpperCAmelCase       )
     __UpperCAmelCase  :     List[str]		=    VersatileDiffusionDualGuidedPipeline.from_pretrained(_UpperCAmelCase       )

   pipe.to(_UpperCAmelCase       )
   pipe.set_progress_bar_config(disable=_UpperCAmelCase       )

   __UpperCAmelCase  :     List[str]		=    generator.manual_seed(0       )
   __UpperCAmelCase  :     Optional[int]		=    pipe(
       prompt='first prompt'     ,     image=_UpperCAmelCase     ,     text_to_image_strength=0.7_5     ,     generator=_UpperCAmelCase     ,     guidance_scale=7.5     ,     num_inference_steps=2     ,     output_type='numpy'     ,     ).images

   assert np.abs(image - new_image       ).sum() < 1E-5, "Models don't have the same forward pass"




 def 						a					( self       ):
   __UpperCAmelCase  :     Optional[Any]		=    VersatileDiffusionDualGuidedPipeline.from_pretrained('shi-labs/versatile-diffusion'       )
   pipe.remove_unused_weights()
   pipe.to(_UpperCAmelCase       )
   pipe.set_progress_bar_config(disable=_UpperCAmelCase       )

   __UpperCAmelCase  :     Tuple		=    'cyberpunk 2077'
   __UpperCAmelCase  :     Dict		=    load_image(
       'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'       )
   __UpperCAmelCase  :     Optional[Any]		=    torch.manual_seed(0       )
   __UpperCAmelCase  :     List[Any]		=    pipe(
       prompt=_UpperCAmelCase     ,     image=_UpperCAmelCase     ,     text_to_image_strength=0.7_5     ,     generator=_UpperCAmelCase     ,     guidance_scale=7.5     ,     num_inference_steps=50     ,     output_type='numpy'     ,     ).images

   __UpperCAmelCase  :     Tuple		=    image[0, 253:256, 253:256, -1]

   assert image.shape == (1, 512, 512, 3)
   __UpperCAmelCase  :     Tuple		=    np.array([0.0_7_8_7, 0.0_8_4_9, 0.0_8_2_6, 0.0_8_1_2, 0.0_8_0_7, 0.0_7_9_5, 0.0_8_1_8, 0.0_7_9_8, 0.0_7_7_9]       )

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





