


'''simple docstring'''


import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()





class     lowercase__  (						_snake_case     ,     _snake_case     ,     _snake_case     ,     unittest.TestCase					):





							'''simple docstring'''


							A_							:      Tuple      	=  StableUnCLIPImgaImgPipeline
							A_							:      Union[str, Any]      	=  TEXT_GUIDED_IMAGE_VARIATION_PARAMS
							A_							:      Union[str, Any]      	=  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
							A_							:      Tuple      	=  frozenset(
							    []					)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
							A_							:      List[Any]      	=  frozenset([]					)






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       List[str]     		=      32
											_SCREAMING_SNAKE_CASE :       Tuple     		=      embedder_hidden_size

											# image encoding components

											_SCREAMING_SNAKE_CASE :       int     		=      CLIPImageProcessor(crop_size=32    ,					size=32  )

											torch.manual_seed(0  )
											_SCREAMING_SNAKE_CASE :       Tuple     		=      CLIPVisionModelWithProjection(
											    CLIPVisionConfig(
											        hidden_size=__snake_case    ,					projection_dim=__snake_case    ,					num_hidden_layers=5    ,					num_attention_heads=4    ,					image_size=32    ,					intermediate_size=37    ,					patch_size=1    ,					)  )

											# regular denoising components

											torch.manual_seed(0  )
											_SCREAMING_SNAKE_CASE :       str     		=      StableUnCLIPImageNormalizer(embedding_dim=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""  )

											torch.manual_seed(0  )
											_SCREAMING_SNAKE_CASE :       str     		=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""  )

											torch.manual_seed(0  )
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      CLIPTextModel(
											    CLIPTextConfig(
											        bos_token_id=0    ,					eos_token_id=2    ,					hidden_size=__snake_case    ,					projection_dim=32    ,					intermediate_size=37    ,					layer_norm_eps=1e-05    ,					num_attention_heads=4    ,					num_hidden_layers=5    ,					pad_token_id=1    ,					vocab_size=1000    ,					)  )

											torch.manual_seed(0  )
											_SCREAMING_SNAKE_CASE :       int     		=      UNetaDConditionModel(
											    sample_size=32    ,					in_channels=4    ,					out_channels=4    ,					down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")    ,					up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")    ,					block_out_channels=(32, 64)    ,					attention_head_dim=(2, 4)    ,					class_embed_type="""projection"""    ,					projection_class_embeddings_input_dim=embedder_projection_dim * 2    ,					cross_attention_dim=__snake_case    ,					layers_per_block=1    ,					upcast_attention=__snake_case    ,					use_linear_projection=__snake_case    ,					)

											torch.manual_seed(0  )
											_SCREAMING_SNAKE_CASE :       Tuple     		=      DDIMScheduler(
											    beta_schedule="""scaled_linear"""    ,					beta_start=0.00085    ,					beta_end=0.012    ,					prediction_type="""v_prediction"""    ,					set_alpha_to_one=__snake_case    ,					steps_offset=1    ,					)

											torch.manual_seed(0  )
											_SCREAMING_SNAKE_CASE :       int     		=      AutoencoderKL()

											_SCREAMING_SNAKE_CASE :       Dict     		=      {
											    # image encoding components
											    """feature_extractor""": feature_extractor,
											    """image_encoder""": image_encoder.eval(),
											    # image noising components
											    """image_normalizer""": image_normalizer.eval(),
											    """image_noising_scheduler""": image_noising_scheduler,
											    # regular denoising components
											    """tokenizer""": tokenizer,
											    """text_encoder""": text_encoder.eval(),
											    """unet""": unet.eval(),
											    """scheduler""": scheduler,
											    """vae""": vae.eval(),
											}

											return components






							def UpperCAmelCase_	(    self    ,					__snake_case    ,					__snake_case=0    ,					__snake_case=True  ):
											if str(__snake_case  ).startswith("""mps"""  ):
															_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      torch.manual_seed(__snake_case  )
											else:
															_SCREAMING_SNAKE_CASE :       int     		=      torch.Generator(device=__snake_case  ).manual_seed(__snake_case  )

											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      floats_tensor((1, 3, 32, 32)    ,					rng=random.Random(__snake_case  )  ).to(__snake_case  )

											if pil_image:
															_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      input_image * 0.5 + 0.5
															_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      input_image.clamp(0    ,					1  )
															_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      input_image.cpu().permute(0    ,					2    ,					3    ,					1  ).float().numpy()
															_SCREAMING_SNAKE_CASE :       List[Any]     		=      DiffusionPipeline.numpy_to_pil(__snake_case  )[0]

											return {
											    "prompt": "An anime racoon running a marathon",
											    "image": input_image,
											    "generator": generator,
											    "num_inference_steps": 2,
											    "output_type": "np",
											}






							@skip_mps
							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       List[str]     		=      """cpu"""  # ensure determinism for the device-dependent torch.Generator
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      self.get_dummy_components()
											_SCREAMING_SNAKE_CASE :       Dict     		=      StableUnCLIPImgaImgPipeline(**__snake_case  )
											_SCREAMING_SNAKE_CASE :       Dict     		=      sd_pipe.to(__snake_case  )
											sd_pipe.set_progress_bar_config(disable=__snake_case  )

											_SCREAMING_SNAKE_CASE :       str     		=      self.get_dummy_inputs(__snake_case  )
											inputs.update({"""image_embeds""": None}  )
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      sd_pipe(**__snake_case  ).images
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      image[0, -3:, -3:, -1]

											assert image.shape == (1, 32, 32, 3)
											_SCREAMING_SNAKE_CASE :       Dict     		=      np.array([0.3872, 0.7224, 0.5601, 0.4741, 0.6872, 0.5814, 0.4636, 0.3867, 0.5078]  )

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






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Tuple     		=      torch_device in ["""cpu""", """mps"""]

											self._test_attention_slicing_forward_pass(test_max_difference=__snake_case  )






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      torch_device in ["""cpu""", """mps"""]

											self._test_inference_batch_single_identical(test_max_difference=__snake_case  )






							@unittest.skipIf(
							    torch_device != """cuda""" or not is_xformers_available()    ,					reason="""XFormers attention is only available with CUDA and `xformers` installed"""    ,					)
							def UpperCAmelCase_	(    self  ):
											self._test_xformers_attention_forwardGenerator_pass(test_max_difference=__snake_case  )





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





							'''simple docstring'''


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






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       str     		=      load_image(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png"""  )

											_SCREAMING_SNAKE_CASE :       List[str]     		=      load_numpy(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy"""  )

											_SCREAMING_SNAKE_CASE :       Tuple     		=      StableUnCLIPImgaImgPipeline.from_pretrained(
											    """fusing/stable-unclip-2-1-l-img2img"""    ,					torch_dtype=torch.floataa  )
											pipe.to(__snake_case  )
											pipe.set_progress_bar_config(disable=__snake_case  )
											# stable unclip will oom when integration tests are run on a V100,
											# so turn on memory savings
											pipe.enable_attention_slicing()
											pipe.enable_sequential_cpu_offload()

											_SCREAMING_SNAKE_CASE :       str     		=      torch.Generator(device="""cpu"""  ).manual_seed(0  )
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      pipe(__snake_case    ,					"""anime turle"""    ,					generator=__snake_case    ,					output_type="""np"""  )

											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      output.images[0]

											assert image.shape == (768, 768, 3)

											assert_mean_pixel_difference(__snake_case    ,					__snake_case  )






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      load_image(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png"""  )

											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      load_numpy(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy"""  )

											_SCREAMING_SNAKE_CASE :       int     		=      StableUnCLIPImgaImgPipeline.from_pretrained(
											    """fusing/stable-unclip-2-1-h-img2img"""    ,					torch_dtype=torch.floataa  )
											pipe.to(__snake_case  )
											pipe.set_progress_bar_config(disable=__snake_case  )
											# stable unclip will oom when integration tests are run on a V100,
											# so turn on memory savings
											pipe.enable_attention_slicing()
											pipe.enable_sequential_cpu_offload()

											_SCREAMING_SNAKE_CASE :       Tuple     		=      torch.Generator(device="""cpu"""  ).manual_seed(0  )
											_SCREAMING_SNAKE_CASE :       Any     		=      pipe(__snake_case    ,					"""anime turle"""    ,					generator=__snake_case    ,					output_type="""np"""  )

											_SCREAMING_SNAKE_CASE :       List[str]     		=      output.images[0]

											assert image.shape == (768, 768, 3)

											assert_mean_pixel_difference(__snake_case    ,					__snake_case  )






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      load_image(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png"""  )

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

											_SCREAMING_SNAKE_CASE :       str     		=      StableUnCLIPImgaImgPipeline.from_pretrained(
											    """fusing/stable-unclip-2-1-h-img2img"""    ,					torch_dtype=torch.floataa  )
											_SCREAMING_SNAKE_CASE :       int     		=      pipe.to(__snake_case  )
											pipe.set_progress_bar_config(disable=__snake_case  )
											pipe.enable_attention_slicing()
											pipe.enable_sequential_cpu_offload()

											_SCREAMING_SNAKE_CASE :       List[Any]     		=      pipe(
											    __snake_case    ,					"""anime turtle"""    ,					num_inference_steps=2    ,					output_type="""np"""    ,					)

											_SCREAMING_SNAKE_CASE :       Dict     		=      torch.cuda.max_memory_allocated()
											# make sure that less than 7 GB is allocated
											assert mem_bytes < 7 * 10**9








'''simple docstring'''


from typing import List, Optional, Tuple, Union

import torch

from ...models import UNetaDModel
from ...schedulers import ScoreSdeVeScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput





class     lowercase__  (						_snake_case					):





							'''simple docstring'''


							A_							:      UNetaDModel
							A_							:      ScoreSdeVeScheduler






							def __init__(    self    ,					__snake_case    ,					__snake_case  ):
											super().__init__()
											self.register_modules(unet=__snake_case    ,					scheduler=__snake_case  )






							@torch.no_grad()
							def __call__(    self    ,					__snake_case = 1    ,					__snake_case = 2000    ,					__snake_case = None    ,					__snake_case = "pil"    ,					__snake_case = True    ,					**__snake_case    ,					):
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      self.unet.config.sample_size
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      (batch_size, 3, img_size, img_size)

											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      self.unet

											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      randn_tensor(__snake_case    ,					generator=__snake_case  ) * self.scheduler.init_noise_sigma
											_SCREAMING_SNAKE_CASE :       Any     		=      sample.to(self.device  )

											self.scheduler.set_timesteps(__snake_case  )
											self.scheduler.set_sigmas(__snake_case  )

											for i, t in enumerate(self.progress_bar(self.scheduler.timesteps  )  ):
															_SCREAMING_SNAKE_CASE :       List[Any]     		=      self.scheduler.sigmas[i] * torch.ones(shape[0]    ,					device=self.device  )

															# correction step
															for _ in range(self.scheduler.config.correct_steps  ):
																			_SCREAMING_SNAKE_CASE :       Any     		=      self.unet(__snake_case    ,					__snake_case  ).sample
																			_SCREAMING_SNAKE_CASE :       Any     		=      self.scheduler.step_correct(__snake_case    ,					__snake_case    ,					generator=__snake_case  ).prev_sample

															# prediction step
															_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      model(__snake_case    ,					__snake_case  ).sample
															_SCREAMING_SNAKE_CASE :       Tuple     		=      self.scheduler.step_pred(__snake_case    ,					__snake_case    ,					__snake_case    ,					generator=__snake_case  )

															_SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE :       int     		=      output.prev_sample, output.prev_sample_mean

											_SCREAMING_SNAKE_CASE :       Tuple     		=      sample_mean.clamp(0    ,					1  )
											_SCREAMING_SNAKE_CASE :       str     		=      sample.cpu().permute(0    ,					2    ,					3    ,					1  ).numpy()
											if output_type == "pil":
															_SCREAMING_SNAKE_CASE :       Optional[int]     		=      self.numpy_to_pil(__snake_case  )

											if not return_dict:
															return (sample,)

											return ImagePipelineOutput(images=__snake_case  )




