

from typing import List, Optional, Tuple, Union

from . import logging
from .import_utils import is_torch_available, is_torch_version


if is_torch_available():
						import torch

_SCREAMING_SNAKE_CASE          = logging.get_logger(__name__)  # pylint: disable=invalid-name

try:
						from torch._dynamo import allow_in_graph as maybe_allow_in_graph
except (ImportError, ModuleNotFoundError):





						def    _snake_case		(cls				:    Optional[Any]     )       ->		Tuple:
								return cls





def    _snake_case		(snake_case_				:    Union[Tuple, List]						,					snake_case_				:    Optional[Union[List["torch.Generator"], "torch.Generator"]] = None						,					snake_case_				:    Optional["torch.device"] = None						,					snake_case_				:    Optional["torch.dtype"] = None						,					snake_case_				:    Optional["torch.layout"] = None						,					)       ->		Dict:
		__UpperCAmelCase  :     Union[str, Any]		=    device
		__UpperCAmelCase  :     str		=    shape[0]

		__UpperCAmelCase  :     Optional[Any]		=    layout or torch.strided
		__UpperCAmelCase  :     Union[str, Any]		=    device or torch.device('cpu'     )

		if generator is not None:
				__UpperCAmelCase  :     int		=    generator.device.type if not isinstance(_SCREAMING_SNAKE_CASE						,					_SCREAMING_SNAKE_CASE     ) else generator[0].device.type
				if gen_device_type != device.type and gen_device_type == "cpu":
						__UpperCAmelCase  :     List[str]		=    'cpu'
						if device != "mps":
								logger.info(
								    F"The passed generator was created on 'cpu' even though a tensor on {device} was expected."
								    F" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably"
								    F" slighly speed up this function by passing a generator that was created on the {device} device."     )
				elif gen_device_type != device.type and gen_device_type == "cuda":
						raise ValueError(F"Cannot generate a {device} tensor from a generator of type {gen_device_type}."     )

		if isinstance(_SCREAMING_SNAKE_CASE						,					_SCREAMING_SNAKE_CASE     ):
				__UpperCAmelCase  :     Optional[Any]		=    (1,) + shape[1:]
				__UpperCAmelCase  :     List[str]		=    [
				    torch.randn(_SCREAMING_SNAKE_CASE						,					generator=generator[i]						,					device=_SCREAMING_SNAKE_CASE						,					dtype=_SCREAMING_SNAKE_CASE						,					layout=_SCREAMING_SNAKE_CASE     )
				    for i in range(_SCREAMING_SNAKE_CASE     )
				]
				__UpperCAmelCase  :     Dict		=    torch.cat(_SCREAMING_SNAKE_CASE						,					dim=0     ).to(_SCREAMING_SNAKE_CASE     )
		else:
				__UpperCAmelCase  :     str		=    torch.randn(_SCREAMING_SNAKE_CASE						,					generator=_SCREAMING_SNAKE_CASE						,					device=_SCREAMING_SNAKE_CASE						,					dtype=_SCREAMING_SNAKE_CASE						,					layout=_SCREAMING_SNAKE_CASE     ).to(_SCREAMING_SNAKE_CASE     )

		return latents







def    _snake_case		(snake_case_				:    Optional[int]     )       ->		Any:
		if is_torch_version('<'						,					'2.0.0'     ) or not hasattr(_SCREAMING_SNAKE_CASE						,					'_dynamo'     ):
				return False
		return isinstance(_SCREAMING_SNAKE_CASE						,					torch._dynamo.eval_frame.OptimizedModule     )








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





