
"""simple docstring"""




from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types



def 						__magic_name__		(				UpperCamelCase    :		torch.Tensor       ,		UpperCamelCase    :		torch.Tensor       ,		UpperCamelCase    :		bool    )    ->	torch.Tensor:
    # The first operation in a checkpoint can't be in-place, but it's
    # nice to have in-place addition during inference. Thus...
    if not inplace:
        a__    = ma + ma
    else:
        ma += ma

    return ma



def 						__magic_name__		(				UpperCamelCase    :		torch.Tensor       ,		UpperCamelCase    :		List[int]    )    ->	torch.Tensor:
    a__    = -1 * len(UpperCamelCase    )
    a__    = list(range(len(tensor.shape[:zero_index]    )    )    )
    return tensor.permute(first_inds + [zero_index + i for i in inds]    )



def 						__magic_name__		(				UpperCamelCase    :		torch.Tensor       ,		UpperCamelCase    :		int    )    ->	torch.Tensor:
    return t.reshape(t.shape[:-no_dims] + (-1,)    )



def 						__magic_name__		(				UpperCamelCase    :		torch.Tensor       ,		UpperCamelCase    :		torch.Tensor       ,		UpperCamelCase    :		int       ,		UpperCamelCase    :		float = 1E-4    )    ->	torch.Tensor:
    a__    = mask.expand(*value.shape    )
    return torch.sum(mask * value       ,		dim=UpperCamelCase    ) / (eps + torch.sum(UpperCamelCase       ,		dim=UpperCamelCase    ))



def 						__magic_name__		(				UpperCamelCase    :		torch.Tensor       ,		UpperCamelCase    :		torch.types.Number = 2.3125       ,		UpperCamelCase    :		torch.types.Number = 21.6875       ,		UpperCamelCase    :		int = 64    )    ->	torch.Tensor:
    a__    = torch.linspace(UpperCamelCase       ,		UpperCamelCase       ,		no_bins - 1       ,		device=pts.device    )
    a__    = torch.sqrt(torch.sum((pts.unsqueeze(-2    ) - pts.unsqueeze(-3    )) ** 2       ,		dim=-1    )    )
    return torch.bucketize(UpperCamelCase       ,		UpperCamelCase    )



def 						__magic_name__		(				UpperCamelCase    :		Callable[[list], Any]       ,		UpperCamelCase    :		List[dict]    )    ->	dict:
    a__    = dicts[0]
    a__    = {}
    for k, v in first.items():
        a__    = [d[k] for d in dicts]
        if isinstance(UpperCamelCase       ,		UpperCamelCase    ):
            a__    = dict_multimap(UpperCamelCase       ,		UpperCamelCase    )
        else:
            a__    = fn(UpperCamelCase    )

    return new_dict



def 						__magic_name__		(				UpperCamelCase    :		torch.Tensor       ,		UpperCamelCase    :		torch.Tensor    )    ->	torch.Tensor:
    a__    = v_bins.view(((1,) * len(x.shape    )) + (len(UpperCamelCase    ),)    )
    a__    = x[..., None] - reshaped_bins
    a__    = torch.argmin(torch.abs(UpperCamelCase    )       ,		dim=-1    )
    return nn.functional.one_hot(UpperCamelCase       ,		num_classes=len(UpperCamelCase    )    ).float()



def 						__magic_name__		(				UpperCamelCase    :		torch.Tensor       ,		UpperCamelCase    :		torch.Tensor       ,		UpperCamelCase    :		int = 0       ,		UpperCamelCase    :		int = 0    )    ->	torch.Tensor:
    a__    = []
    for i, s in enumerate(data.shape[:no_batch_dims]    ):
        a__    = torch.arange(UpperCamelCase    )
        a__    = r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape    ) - i - 1)))    )
        ranges.append(UpperCamelCase    )

    a__    = [slice(UpperCamelCase    ) for _ in range(len(data.shape    ) - no_batch_dims    )]
    a__    = inds
    ranges.extend(UpperCamelCase    )
    # Matt note: Editing this to get around the behaviour of using a list as an array index changing
    # in recent Numpy versions
    return data[tuple(UpperCamelCase    )]


a		:					str    	=					TypeVar('T')



def 						__magic_name__		(				UpperCamelCase    :		Callable[[T], Any]       ,		UpperCamelCase    :		Dict[Any, Union[dict, list, tuple, T]]       ,		UpperCamelCase    :		Type[T]    )    ->	Dict[Any, Union[dict, list, tuple, Any]]:
    a__    = {}
    for k, v in dic.items():
        if isinstance(UpperCamelCase       ,		UpperCamelCase    ):
            a__    = dict_map(UpperCamelCase       ,		UpperCamelCase       ,		UpperCamelCase    )
        else:
            a__    = tree_map(UpperCamelCase       ,		UpperCamelCase       ,		UpperCamelCase    )

    return new_dict



@overload
def 						__magic_name__		(				UpperCamelCase    :		Callable[[T], Any]       ,		UpperCamelCase    :		T       ,		UpperCamelCase    :		Type[T]    )    ->	Any:
    ...



@overload
def 						__magic_name__		(				UpperCamelCase    :		Callable[[T], Any]       ,		UpperCamelCase    :		dict       ,		UpperCamelCase    :		Type[T]    )    ->	dict:
    ...



@overload
def 						__magic_name__		(				UpperCamelCase    :		Callable[[T], Any]       ,		UpperCamelCase    :		list       ,		UpperCamelCase    :		Type[T]    )    ->	list:
    ...



@overload
def 						__magic_name__		(				UpperCamelCase    :		Callable[[T], Any]       ,		UpperCamelCase    :		tuple       ,		UpperCamelCase    :		Type[T]    )    ->	tuple:
    ...



def 						__magic_name__		(				UpperCamelCase    :		Optional[Any]       ,		UpperCamelCase    :		Optional[Any]       ,		UpperCamelCase    :		Optional[Any]    )    ->	Optional[Any]:
    if isinstance(UpperCamelCase       ,		UpperCamelCase    ):
        return dict_map(UpperCamelCase       ,		UpperCamelCase       ,		UpperCamelCase    )
    elif isinstance(UpperCamelCase       ,		UpperCamelCase    ):
        return [tree_map(UpperCamelCase       ,		UpperCamelCase       ,		UpperCamelCase    ) for x in tree]
    elif isinstance(UpperCamelCase       ,		UpperCamelCase    ):
        return tuple(tree_map(UpperCamelCase       ,		UpperCamelCase       ,		UpperCamelCase    ) for x in tree    )
    elif isinstance(UpperCamelCase       ,		UpperCamelCase    ):
        return fn(UpperCamelCase    )
    else:
        print(type(UpperCamelCase    )    )
        raise ValueError('Not supported'    )



a		:					List[Any]    	=					partial(tree_map, leaf_type=torch.Tensor)




"""simple docstring"""




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


a		:					str    	=					False







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



     def  lowercase__     (       self      )		->    Any:







         """simple docstring"""



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



     def  lowercase__     (       self      )		->    Optional[int]:







         """simple docstring"""



         a__    = VersatileDiffusionDualGuidedPipeline.from_pretrained('shi-labs/versatile-diffusion'      )
         # remove text_unet
         pipe.remove_unused_weights()
         pipe.to(__SCREAMING_SNAKE_CASE      )
         pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE      )

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

         a__    = torch.manual_seed(0      )
         a__    = pipe(
             prompt='first prompt'    ,							image=__SCREAMING_SNAKE_CASE    ,							text_to_image_strength=0.75    ,							generator=__SCREAMING_SNAKE_CASE    ,							guidance_scale=7.5    ,							num_inference_steps=2    ,							output_type='numpy'    ,							).images

         with tempfile.TemporaryDirectory() as tmpdirname:
             pipe.save_pretrained(__SCREAMING_SNAKE_CASE      )
             a__    = VersatileDiffusionDualGuidedPipeline.from_pretrained(__SCREAMING_SNAKE_CASE      )

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

         a__    = generator.manual_seed(0      )
         a__    = pipe(
             prompt='first prompt'    ,							image=__SCREAMING_SNAKE_CASE    ,							text_to_image_strength=0.75    ,							generator=__SCREAMING_SNAKE_CASE    ,							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  lowercase__     (       self      )		->    Union[str, Any]:







         """simple docstring"""



         a__    = VersatileDiffusionDualGuidedPipeline.from_pretrained('shi-labs/versatile-diffusion'      )
         pipe.remove_unused_weights()
         pipe.to(__SCREAMING_SNAKE_CASE      )
         pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE      )

         a__    = 'cyberpunk 2077'
         a__    = load_image(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'      )
         a__    = torch.manual_seed(0      )
         a__    = pipe(
             prompt=__SCREAMING_SNAKE_CASE    ,							image=__SCREAMING_SNAKE_CASE    ,							text_to_image_strength=0.75    ,							generator=__SCREAMING_SNAKE_CASE    ,							guidance_scale=7.5    ,							num_inference_steps=5_0    ,							output_type='numpy'    ,							).images

         a__    = image[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         a__    = np.array([0.07_87, 0.08_49, 0.08_26, 0.08_12, 0.08_07, 0.07_95, 0.08_18, 0.07_98, 0.07_79]      )

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


