
"""simple docstring"""




import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()







class 						lowercase(_lowercase						,	unittest.TestCase					):
     __snake_case:  Union[str, Any]  				=    KandinskyPipeline
     __snake_case:  Optional[Any]  				=    [
         'prompt',
         'image_embeds',
         'negative_image_embeds',
     ]
     __snake_case:  str  				=    ['prompt', 'negative_prompt', 'image_embeds', 'negative_image_embeds']
     __snake_case:  Union[str, Any]  				=    [
         'generator',
         'height',
         'width',
         'latents',
         'guidance_scale',
         'negative_prompt',
         'num_inference_steps',
         'return_dict',
         'guidance_scale',
         'num_images_per_prompt',
         'output_type',
         'return_dict',
     ]
     __snake_case:  Any  				=    False



     @property
     def  lowercase__     (       self      )		->    Tuple:







         """simple docstring"""



         return 3_2



     @property
     def  lowercase__     (       self      )		->    str:







         """simple docstring"""



         return 3_2



     @property
     def  lowercase__     (       self      )		->    List[str]:







         """simple docstring"""



         return self.time_input_dim



     @property
     def  lowercase__     (       self      )		->    Union[str, Any]:







         """simple docstring"""



         return self.time_input_dim * 4



     @property
     def  lowercase__     (       self      )		->    Any:







         """simple docstring"""



         return 1_0_0



     @property
     def  lowercase__     (       self      )		->    Optional[Any]:







         """simple docstring"""



         a__    = XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base'      )
         return tokenizer



     @property
     def  lowercase__     (       self      )		->    Dict:







         """simple docstring"""



         torch.manual_seed(0      )
         a__    = MCLIPConfig(
             numDims=self.cross_attention_dim    ,							transformerDimensions=self.text_embedder_hidden_size    ,							hidden_size=self.text_embedder_hidden_size    ,							intermediate_size=3_7    ,							num_attention_heads=4    ,							num_hidden_layers=5    ,							vocab_size=1_0_0_5    ,							)

         a__    = MultilingualCLIP(__SCREAMING_SNAKE_CASE      )
         a__    = text_encoder.eval()

         return text_encoder



     @property
     def  lowercase__     (       self      )		->    Dict:







         """simple docstring"""



         torch.manual_seed(0      )

         a__    = {
             'in_channels': 4,
             # Out channels is double in channels because predicts mean and variance
             'out_channels': 8,
             'addition_embed_type': 'text_image',
             'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
             'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
             'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
             'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
             'layers_per_block': 1,
             'encoder_hid_dim': self.text_embedder_hidden_size,
             'encoder_hid_dim_type': 'text_image_proj',
             'cross_attention_dim': self.cross_attention_dim,
             'attention_head_dim': 4,
             'resnet_time_scale_shift': 'scale_shift',
             'class_embed_type': None,
         }

         a__    = UNetaDConditionModel(**__SCREAMING_SNAKE_CASE      )
         return model



     @property
     def  lowercase__     (       self      )		->    Tuple:







         """simple docstring"""



         return {
             "block_out_channels": [3_2, 6_4],
             "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
             "in_channels": 3,
             "latent_channels": 4,
             "layers_per_block": 1,
             "norm_num_groups": 8,
             "norm_type": "spatial",
             "num_vq_embeddings": 1_2,
             "out_channels": 3,
             "up_block_types": [
                 "AttnUpDecoderBlock2D",
                 "UpDecoderBlock2D",
             ],
             "vq_embed_dim": 4,
         }



     @property
     def  lowercase__     (       self      )		->    Any:







         """simple docstring"""



         torch.manual_seed(0      )
         a__    = VQModel(**self.dummy_movq_kwargs      )
         return model



     def  lowercase__     (       self      )		->    Union[str, Any]:







         """simple docstring"""



         a__    = self.dummy_text_encoder
         a__    = self.dummy_tokenizer
         a__    = self.dummy_unet
         a__    = self.dummy_movq

         a__    = DDIMScheduler(
             num_train_timesteps=1_0_0_0    ,							beta_schedule='linear'    ,							beta_start=0.0_00_85    ,							beta_end=0.0_12    ,							clip_sample=__SCREAMING_SNAKE_CASE    ,							set_alpha_to_one=__SCREAMING_SNAKE_CASE    ,							steps_offset=1    ,							prediction_type='epsilon'    ,							thresholding=__SCREAMING_SNAKE_CASE    ,							)

         a__    = {
             'text_encoder': text_encoder,
             'tokenizer': tokenizer,
             'unet': unet,
             'scheduler': scheduler,
             'movq': movq,
         }
         return components



     def  lowercase__     (       self    ,							__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE=0      )		->    str:







         """simple docstring"""



         a__    = floats_tensor((1, self.cross_attention_dim)    ,							rng=random.Random(__SCREAMING_SNAKE_CASE      )      ).to(__SCREAMING_SNAKE_CASE      )
         a__    = floats_tensor((1, self.cross_attention_dim)    ,							rng=random.Random(seed + 1      )      ).to(__SCREAMING_SNAKE_CASE      )
         if str(__SCREAMING_SNAKE_CASE      ).startswith('mps'      ):
             a__    = torch.manual_seed(__SCREAMING_SNAKE_CASE      )
         else:
             a__    = torch.Generator(device=__SCREAMING_SNAKE_CASE      ).manual_seed(__SCREAMING_SNAKE_CASE      )
         a__    = {
             'prompt': 'horse',
             'image_embeds': image_embeds,
             'negative_image_embeds': negative_image_embeds,
             'generator': generator,
             'height': 6_4,
             'width': 6_4,
             'guidance_scale': 4.0,
             'num_inference_steps': 2,
             'output_type': 'np',
         }
         return inputs




     def  lowercase__     (       self      )		->    List[Any]:







         """simple docstring"""



         a__    = 'cpu'

         a__    = self.get_dummy_components()

         a__    = self.pipeline_class(**__SCREAMING_SNAKE_CASE      )
         a__    = pipe.to(__SCREAMING_SNAKE_CASE      )

         pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE      )

         a__    = pipe(**self.get_dummy_inputs(__SCREAMING_SNAKE_CASE      )      )
         a__    = output.images

         a__    = pipe(
             **self.get_dummy_inputs(__SCREAMING_SNAKE_CASE      )    ,							return_dict=__SCREAMING_SNAKE_CASE    ,							)[0]

         a__    = image[0, -3:, -3:, -1]
         a__    = image_from_tuple[0, -3:, -3:, -1]

         assert image.shape == (1, 6_4, 6_4, 3)

         a__    = np.array(
             [0.32_86_63, 1.0, 0.23_21_68_73, 1.0, 0.92_71_75_64, 0.4_63_90_46, 0.96_89_47_77, 0.31_71_33_78, 0.6_29_39_53]      )

         assert (
             np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2
         ), f' expected_slice {expected_slice}, but got {image_slice.flatten()}'
         assert (
             np.abs(image_from_tuple_slice.flatten() - expected_slice      ).max() < 1e-2
         ), f' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'







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



     def  lowercase__     (       self      )		->    str:







         """simple docstring"""



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




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







         """simple docstring"""



         a__    = load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
             '/kandinsky/kandinsky_text2img_cat_fp16.npy'      )

         a__    = KandinskyPriorPipeline.from_pretrained(
             'kandinsky-community/kandinsky-2-1-prior'    ,							torch_dtype=torch.floataa      )
         pipe_prior.to(__SCREAMING_SNAKE_CASE      )

         a__    = KandinskyPipeline.from_pretrained('kandinsky-community/kandinsky-2-1'    ,							torch_dtype=torch.floataa      )
         a__    = pipeline.to(__SCREAMING_SNAKE_CASE      )
         pipeline.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE      )

         a__    = 'red cat, 4k photo'

         a__    = torch.Generator(device='cuda'      ).manual_seed(0      )
         a__  ,    a__    = pipe_prior(
             __SCREAMING_SNAKE_CASE    ,							generator=__SCREAMING_SNAKE_CASE    ,							num_inference_steps=5    ,							negative_prompt=''    ,							).to_tuple()

         a__    = torch.Generator(device='cuda'      ).manual_seed(0      )
         a__    = pipeline(
             __SCREAMING_SNAKE_CASE    ,							image_embeds=__SCREAMING_SNAKE_CASE    ,							negative_image_embeds=__SCREAMING_SNAKE_CASE    ,							generator=__SCREAMING_SNAKE_CASE    ,							num_inference_steps=1_0_0    ,							output_type='np'    ,							)

         a__    = output.images[0]

         assert image.shape == (5_1_2, 5_1_2, 3)

         assert_mean_pixel_difference(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE      )




"""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


