





import unittest

from datasets import load_dataset

from transformers import load_tool

from .test_tools_common import ToolTesterMixin





class 	lowerCAmelCase__(       unittest.TestCase     ,   __lowercase			):






       '''simple docstring'''


       def 			UpperCamelCase_	(					self							)	->   Tuple:
            _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							load_tool("document-question-answering"							)
            self.tool.setup()
            _SCREAMING_SNAKE_CASE				:		Dict  			=							load_tool("document-question-answering"			,  remote=__lowerCamelCase							)




       def 			UpperCamelCase_	(					self							)	->   Any:
            _SCREAMING_SNAKE_CASE				:		List[Any]  			=							load_dataset("hf-internal-testing/example-documents"			,  split="test"							)
            _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							dataset[0]["image"]

            _SCREAMING_SNAKE_CASE				:		Tuple  			=							self.tool(__lowerCamelCase			,  "When is the coffee break?"							)
            self.assertEqual(__lowerCamelCase			,  "11-14 to 11:39 a.m."							)




       def 			UpperCamelCase_	(					self							)	->   Any:
            _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							load_dataset("hf-internal-testing/example-documents"			,  split="test"							)
            _SCREAMING_SNAKE_CASE				:		List[Any]  			=							dataset[0]["image"]

            _SCREAMING_SNAKE_CASE				:		Tuple  			=							self.remote_tool(__lowerCamelCase			,  "When is the coffee break?"							)
            self.assertEqual(__lowerCamelCase			,  "11-14 to 11:39 a.m."							)




       def 			UpperCamelCase_	(					self							)	->   Dict:
            _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							load_dataset("hf-internal-testing/example-documents"			,  split="test"							)
            _SCREAMING_SNAKE_CASE				:		List[Any]  			=							dataset[0]["image"]

            self.tool(document=__lowerCamelCase			,  question="When is the coffee break?"							)







       def 			UpperCamelCase_	(					self							)	->   Optional[Any]:
            _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							load_dataset("hf-internal-testing/example-documents"			,  split="test"							)
            _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							dataset[0]["image"]

            _SCREAMING_SNAKE_CASE				:		Any  			=							self.remote_tool(document=__lowerCamelCase			,  question="When is the coffee break?"							)
            self.assertEqual(__lowerCamelCase			,  "11-14 to 11:39 a.m."							)





import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModelWithProjection

from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    PNDMScheduler,
    StableDiffusionImageVariationPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
class       __a					( __lowerCamelCase							,     __lowerCamelCase							,     __lowerCamelCase							,     unittest.TestCase     ):


     """simple docstring"""


     _A				:    int									=  StableDiffusionImageVariationPipeline
     _A				:    List[str]									=  IMAGE_VARIATION_PARAMS
     _A				:    List[str]									=  IMAGE_VARIATION_BATCH_PARAMS
     _A				:    Any									=  frozenset([]     )
     # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
     _A				:    Tuple									=  frozenset([]     )




     def   __A				(	self			: List[str]						)		->   Optional[int]:


         '''simple docstring'''





         torch.manual_seed(0						)
         SCREAMING_SNAKE_CASE__								=UNetaDConditionModel(
             block_out_channels=(3_2, 6_4)				,layers_per_block=2				,sample_size=3_2				,in_channels=4				,out_channels=4				,down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")				,up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")				,cross_attention_dim=3_2				,)
         SCREAMING_SNAKE_CASE__								=PNDMScheduler(skip_prk_steps=_UpperCamelCase						)
         torch.manual_seed(0						)
         SCREAMING_SNAKE_CASE__								=AutoencoderKL(
             block_out_channels=[3_2, 6_4]				,in_channels=3				,out_channels=3				,down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]				,up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]				,latent_channels=4				,)
         torch.manual_seed(0						)
         SCREAMING_SNAKE_CASE__								=CLIPVisionConfig(
             hidden_size=3_2				,projection_dim=3_2				,intermediate_size=3_7				,layer_norm_eps=1e-05				,num_attention_heads=4				,num_hidden_layers=5				,image_size=3_2				,patch_size=4				,)
         SCREAMING_SNAKE_CASE__								=CLIPVisionModelWithProjection(_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__								=CLIPImageProcessor(crop_size=3_2				,size=3_2						)

         SCREAMING_SNAKE_CASE__								={
             """unet""": unet,
             """scheduler""": scheduler,
             """vae""": vae,
             """image_encoder""": image_encoder,
             """feature_extractor""": feature_extractor,
             """safety_checker""": None,
         }
         return components




     def   __A				(	self			: str				,_UpperCamelCase			: Optional[int]				,_UpperCamelCase			: Union[str, Any]=0						)		->   Optional[Any]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__								=floats_tensor((1, 3, 3_2, 3_2)				,rng=random.Random(_UpperCamelCase						)						)
         SCREAMING_SNAKE_CASE__								=image.cpu().permute(0				,2				,3				,1						)[0]
         SCREAMING_SNAKE_CASE__								=Image.fromarray(np.uinta(_UpperCamelCase						)						).convert("""RGB"""						).resize((3_2, 3_2)						)
         if str(_UpperCamelCase						).startswith("""mps"""						):
             SCREAMING_SNAKE_CASE__								=torch.manual_seed(_UpperCamelCase						)
         else:
             SCREAMING_SNAKE_CASE__								=torch.Generator(device=_UpperCamelCase						).manual_seed(_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__								={
             """image""": image,
             """generator""": generator,
             """num_inference_steps""": 2,
             """guidance_scale""": 6.0,
             """output_type""": """numpy""",
         }
         return inputs




     def   __A				(	self			: int						)		->   List[Any]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__								="""cpu"""  # ensure determinism for the device-dependent torch.Generator
         SCREAMING_SNAKE_CASE__								=self.get_dummy_components()
         SCREAMING_SNAKE_CASE__								=StableDiffusionImageVariationPipeline(**_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__								=sd_pipe.to(_UpperCamelCase						)
         sd_pipe.set_progress_bar_config(disable=_UpperCamelCase						)

         SCREAMING_SNAKE_CASE__								=self.get_dummy_inputs(_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__								=sd_pipe(**_UpperCamelCase						).images
         SCREAMING_SNAKE_CASE__								=image[0, -3:, -3:, -1]

         assert image.shape == (1, 6_4, 6_4, 3)
         SCREAMING_SNAKE_CASE__								=np.array([0.5239, 0.5723, 0.4796, 0.5049, 0.5550, 0.4685, 0.5329, 0.4891, 0.4921]						)

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




     def   __A				(	self			: Tuple						)		->   Dict:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__								="""cpu"""  # ensure determinism for the device-dependent torch.Generator
         SCREAMING_SNAKE_CASE__								=self.get_dummy_components()
         SCREAMING_SNAKE_CASE__								=StableDiffusionImageVariationPipeline(**_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__								=sd_pipe.to(_UpperCamelCase						)
         sd_pipe.set_progress_bar_config(disable=_UpperCamelCase						)

         SCREAMING_SNAKE_CASE__								=self.get_dummy_inputs(_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__								=2 * [inputs["""image"""]]
         SCREAMING_SNAKE_CASE__								=sd_pipe(**_UpperCamelCase						)

         SCREAMING_SNAKE_CASE__								=output.images

         SCREAMING_SNAKE_CASE__								=image[-1, -3:, -3:, -1]

         assert image.shape == (2, 6_4, 6_4, 3)
         SCREAMING_SNAKE_CASE__								=np.array([0.6892, 0.5637, 0.5836, 0.5771, 0.6254, 0.6409, 0.5580, 0.5569, 0.5289]						)

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




     def   __A				(	self			: Dict						)		->   Dict:


         '''simple docstring'''





         super().test_inference_batch_single_identical(expected_max_diff=3e-3						)






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


     """simple docstring"""


     def   __A				(	self			: str						)		->   List[str]:


         '''simple docstring'''





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




     def   __A				(	self			: Optional[int]				,_UpperCamelCase			: Dict				,_UpperCamelCase			: Union[str, Any]="cpu"				,_UpperCamelCase			: Tuple=torch.floataa				,_UpperCamelCase			: List[Any]=0						)		->   List[Any]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__								=torch.Generator(device=_UpperCamelCase						).manual_seed(_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__								=load_image(
             """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
             """/stable_diffusion_imgvar/input_image_vermeer.png"""						)
         SCREAMING_SNAKE_CASE__								=np.random.RandomState(_UpperCamelCase						).standard_normal((1, 4, 6_4, 6_4)						)
         SCREAMING_SNAKE_CASE__								=torch.from_numpy(_UpperCamelCase						).to(device=_UpperCamelCase				,dtype=_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__								={
             """image""": init_image,
             """latents""": latents,
             """generator""": generator,
             """num_inference_steps""": 3,
             """guidance_scale""": 7.5,
             """output_type""": """numpy""",
         }
         return inputs




     def   __A				(	self			: List[Any]						)		->   int:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__								=StableDiffusionImageVariationPipeline.from_pretrained(
             """lambdalabs/sd-image-variations-diffusers"""				,safety_checker=_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__								=sd_pipe.to(_UpperCamelCase						)
         sd_pipe.set_progress_bar_config(disable=_UpperCamelCase						)

         SCREAMING_SNAKE_CASE__								=self.get_inputs(_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__								=sd_pipe(**_UpperCamelCase						).images
         SCREAMING_SNAKE_CASE__								=image[0, -3:, -3:, -1].flatten()

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         SCREAMING_SNAKE_CASE__								=np.array([0.8_4491, 0.9_0789, 0.7_5708, 0.7_8734, 0.8_3485, 0.7_0099, 0.6_6938, 0.6_8727, 0.6_1379]						)
         assert np.abs(image_slice - expected_slice						).max() < 6e-3




     def   __A				(	self			: Optional[int]						)		->   Dict:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__								=0

         def callback_fn(_UpperCamelCase			: int				,_UpperCamelCase			: int				,_UpperCamelCase			: torch.FloatTensor						) -> None:
             SCREAMING_SNAKE_CASE__								=True
             nonlocal number_of_steps
             number_of_steps += 1
             if step == 1:
                 SCREAMING_SNAKE_CASE__								=latents.detach().cpu().numpy()
                 assert latents.shape == (1, 4, 6_4, 6_4)
                 SCREAMING_SNAKE_CASE__								=latents[0, -3:, -3:, -1]
                 SCREAMING_SNAKE_CASE__								=np.array(
                     [-0.1621, 0.2837, -0.7979, -0.1221, -1.3057, 0.7681, -2.1191, 0.0464, 1.6309]						)

                 assert np.abs(latents_slice.flatten() - expected_slice						).max() < 5e-2
             elif step == 2:
                 SCREAMING_SNAKE_CASE__								=latents.detach().cpu().numpy()
                 assert latents.shape == (1, 4, 6_4, 6_4)
                 SCREAMING_SNAKE_CASE__								=latents[0, -3:, -3:, -1]
                 SCREAMING_SNAKE_CASE__								=np.array([0.6299, 1.7500, 1.1992, -2.1582, -1.8994, 0.7334, -0.7090, 1.0137, 1.5273]						)

                 assert np.abs(latents_slice.flatten() - expected_slice						).max() < 5e-2

         SCREAMING_SNAKE_CASE__								=False

         SCREAMING_SNAKE_CASE__								=StableDiffusionImageVariationPipeline.from_pretrained(
             """fusing/sd-image-variations-diffusers"""				,safety_checker=_UpperCamelCase				,torch_dtype=torch.floataa				,)
         pipe.to(_UpperCamelCase						)
         pipe.set_progress_bar_config(disable=_UpperCamelCase						)
         pipe.enable_attention_slicing()

         SCREAMING_SNAKE_CASE__								=self.get_inputs(_UpperCamelCase				,dtype=torch.floataa						)
         pipe(**_UpperCamelCase				,callback=_UpperCamelCase				,callback_steps=1						)
         assert callback_fn.has_been_called
         assert number_of_steps == inputs["num_inference_steps"]




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


         '''simple docstring'''





         torch.cuda.empty_cache()
         torch.cuda.reset_max_memory_allocated()
         torch.cuda.reset_peak_memory_stats()

         SCREAMING_SNAKE_CASE__								="""fusing/sd-image-variations-diffusers"""
         SCREAMING_SNAKE_CASE__								=StableDiffusionImageVariationPipeline.from_pretrained(
             _UpperCamelCase				,safety_checker=_UpperCamelCase				,torch_dtype=torch.floataa						)
         SCREAMING_SNAKE_CASE__								=pipe.to(_UpperCamelCase						)
         pipe.set_progress_bar_config(disable=_UpperCamelCase						)
         pipe.enable_attention_slicing(1						)
         pipe.enable_sequential_cpu_offload()

         SCREAMING_SNAKE_CASE__								=self.get_inputs(_UpperCamelCase				,dtype=torch.floataa						)
         SCREAMING_SNAKE_CASE__								=pipe(**_UpperCamelCase						)

         SCREAMING_SNAKE_CASE__								=torch.cuda.max_memory_allocated()
         # make sure that less than 2.6 GB is allocated
         assert mem_bytes < 2.6 * 1_0**9






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


     """simple docstring"""


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


         '''simple docstring'''





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




     def   __A				(	self			: Dict				,_UpperCamelCase			: List[str]				,_UpperCamelCase			: Union[str, Any]="cpu"				,_UpperCamelCase			: Any=torch.floataa				,_UpperCamelCase			: Tuple=0						)		->   Any:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__								=torch.Generator(device=_UpperCamelCase						).manual_seed(_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__								=load_image(
             """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
             """/stable_diffusion_imgvar/input_image_vermeer.png"""						)
         SCREAMING_SNAKE_CASE__								=np.random.RandomState(_UpperCamelCase						).standard_normal((1, 4, 6_4, 6_4)						)
         SCREAMING_SNAKE_CASE__								=torch.from_numpy(_UpperCamelCase						).to(device=_UpperCamelCase				,dtype=_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__								={
             """image""": init_image,
             """latents""": latents,
             """generator""": generator,
             """num_inference_steps""": 5_0,
             """guidance_scale""": 7.5,
             """output_type""": """numpy""",
         }
         return inputs




     def   __A				(	self			: Dict						)		->   str:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__								=StableDiffusionImageVariationPipeline.from_pretrained("""fusing/sd-image-variations-diffusers"""						)
         sd_pipe.to(_UpperCamelCase						)
         sd_pipe.set_progress_bar_config(disable=_UpperCamelCase						)

         SCREAMING_SNAKE_CASE__								=self.get_inputs(_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__								=sd_pipe(**_UpperCamelCase						).images[0]

         SCREAMING_SNAKE_CASE__								=load_numpy(
             """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
             """/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy"""						)
         SCREAMING_SNAKE_CASE__								=np.abs(expected_image - image						).max()
         assert max_diff < 1e-3




     def   __A				(	self			: List[str]						)		->   Tuple:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__								=StableDiffusionImageVariationPipeline.from_pretrained("""fusing/sd-image-variations-diffusers"""						)
         SCREAMING_SNAKE_CASE__								=DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config						)
         sd_pipe.to(_UpperCamelCase						)
         sd_pipe.set_progress_bar_config(disable=_UpperCamelCase						)

         SCREAMING_SNAKE_CASE__								=self.get_inputs(_UpperCamelCase						)
         SCREAMING_SNAKE_CASE__								=2_5
         SCREAMING_SNAKE_CASE__								=sd_pipe(**_UpperCamelCase						).images[0]

         SCREAMING_SNAKE_CASE__								=load_numpy(
             """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
             """/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy"""						)
         SCREAMING_SNAKE_CASE__								=np.abs(expected_image - image						).max()
         assert max_diff < 1e-3






