

import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







class 		__magic_name__     (				snake_case__						,			snake_case__						,			unittest.TestCase				):

   '''simple docstring'''




   lowerCamelCase__					:      Dict        =  StableDiffusionParadigmsPipeline
   lowerCamelCase__					:      List[Any]        =  TEXT_TO_IMAGE_PARAMS
   lowerCamelCase__					:      str        =  TEXT_TO_IMAGE_BATCH_PARAMS
   lowerCamelCase__					:      List[Any]        =  TEXT_TO_IMAGE_IMAGE_PARAMS
   lowerCamelCase__					:      Tuple        =  TEXT_TO_IMAGE_IMAGE_PARAMS



   def 					_UpperCAmelCase					(							self			)				->       Any:






     """simple docstring"""






     torch.manual_seed(0			)
     a__    =UNetaDConditionModel(
         block_out_channels=(32, 64),       layers_per_block=2,       sample_size=32,       in_channels=4,       out_channels=4,       down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D'''),       up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D'''),       cross_attention_dim=32,       attention_head_dim=(2, 4),       use_linear_projection=lowercase_,       )
     a__    =DDIMParallelScheduler(
         beta_start=0.00085,       beta_end=0.012,       beta_schedule='''scaled_linear''',       clip_sample=lowercase_,       set_alpha_to_one=lowercase_,       )
     torch.manual_seed(0			)
     a__    =AutoencoderKL(
         block_out_channels=[32, 64],       in_channels=3,       out_channels=3,       down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],       up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],       latent_channels=4,       sample_size=128,       )
     torch.manual_seed(0			)
     a__    =CLIPTextConfig(
         bos_token_id=0,       eos_token_id=2,       hidden_size=32,       intermediate_size=37,       layer_norm_eps=1E-05,       num_attention_heads=4,       num_hidden_layers=5,       pad_token_id=1,       vocab_size=1000,       hidden_act='''gelu''',       projection_dim=512,       )
     a__    =CLIPTextModel(lowercase_			)
     a__    =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''			)

     a__    ={
         'unet': unet,
         'scheduler': scheduler,
         'vae': vae,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'safety_checker': None,
         'feature_extractor': None,
     }
     return components



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_=0			)				->       Dict:






     """simple docstring"""






     if str(lowercase_			).startswith('''mps'''			):
       a__    =torch.manual_seed(lowercase_			)
     else:
       a__    =torch.Generator(device=lowercase_			).manual_seed(lowercase_			)
     a__    ={
         'prompt': 'a photograph of an astronaut riding a horse',
         'generator': generator,
         'num_inference_steps': 10,
         'guidance_scale': 6.0,
         'output_type': 'numpy',
         'parallel': 3,
         'debug': True,
     }
     return inputs



   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     a__    ='cpu'  # ensure determinism for the device-dependent torch.Generator
     a__    =self.get_dummy_components()
     a__    =StableDiffusionParadigmsPipeline(**lowercase_			)
     a__    =sd_pipe.to(lowercase_			)
     sd_pipe.set_progress_bar_config(disable=lowercase_			)

     a__    =self.get_dummy_inputs(lowercase_			)
     a__    =sd_pipe(**lowercase_			).images
     a__    =image[0, -3:, -3:, -1]
     assert image.shape == (1, 64, 64, 3)

     a__    =np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023]			)

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



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






     """simple docstring"""






     a__    ='cpu'  # ensure determinism for the device-dependent torch.Generator
     a__    =self.get_dummy_components()
     torch.manual_seed(0			)
     a__    =DDPMParallelScheduler()
     torch.manual_seed(0			)
     a__    =StableDiffusionParadigmsPipeline(**lowercase_			)
     a__    =sd_pipe.to(lowercase_			)
     sd_pipe.set_progress_bar_config(disable=lowercase_			)

     a__    =self.get_dummy_inputs(lowercase_			)
     a__    =sd_pipe(**lowercase_			).images
     a__    =image[0, -3:, -3:, -1]
     assert image.shape == (1, 64, 64, 3)

     a__    =np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468]			)

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



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






     """simple docstring"""






     super().test_inference_batch_consistent(batch_sizes=[1, 2]			)



   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     super().test_inference_batch_single_identical(batch_size=2,       expected_max_diff=3E-3			)



   def 					_UpperCAmelCase					(							self			)				->       Any:






     """simple docstring"""






     a__    ='cpu'  # ensure determinism for the device-dependent torch.Generator
     a__    =self.get_dummy_components()
     a__    =StableDiffusionParadigmsPipeline(**lowercase_			)
     a__    =sd_pipe.to(lowercase_			)
     sd_pipe.set_progress_bar_config(disable=lowercase_			)

     a__    =self.get_dummy_inputs(lowercase_			)
     a__    ='french fries'
     a__    =sd_pipe(**lowercase_,       negative_prompt=lowercase_			)
     a__    =output.images
     a__    =image[0, -3:, -3:, -1]

     assert image.shape == (1, 64, 64, 3)

     a__    =np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015]			)

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







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

   '''simple docstring'''




   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






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



   def 					_UpperCAmelCase					(							self,       lowercase_=0			)				->       int:






     """simple docstring"""






     a__    =torch.Generator(device=lowercase_			).manual_seed(lowercase_			)
     a__    ={
         'prompt': 'a photograph of an astronaut riding a horse',
         'generator': generator,
         'num_inference_steps': 10,
         'guidance_scale': 7.5,
         'output_type': 'numpy',
         'parallel': 3,
         'debug': True,
     }
     return inputs



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






     """simple docstring"""






     a__    ='stabilityai/stable-diffusion-2-base'
     a__    =DDIMParallelScheduler.from_pretrained(lowercase_,       subfolder='''scheduler'''			)
     a__    =StableDiffusionParadigmsPipeline.from_pretrained(lowercase_,       scheduler=lowercase_,       safety_checker=lowercase_			)
     pipe.to(lowercase_			)
     pipe.set_progress_bar_config(disable=lowercase_			)
     pipe.enable_attention_slicing()

     a__    =self.get_inputs()
     a__    =pipe(**lowercase_			).images
     a__    =image[0, -3:, -3:, -1].flatten()

     assert image.shape == (1, 512, 512, 3)

     a__    =np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741]			)

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



from typing import SupportsIndex

import numpy as np
from scipy.ndimage import convolve







def _a					(			SCREAMING_SNAKE_CASE						:     np.ndarray       ,   SCREAMING_SNAKE_CASE						:     np.ndarray       ,   SCREAMING_SNAKE_CASE						:     np.ndarray						)       ->       np.ndarray:

 """simple docstring"""






 __lowerCAmelCase:  Any        = np.stack((horizontal_flow, vertical_flow)       ,   2						)

 # Create a grid of all pixel coordinates and subtract the flow to get the
 # target pixels coordinates
 __lowerCAmelCase:  List[str]        = np.stack(
     np.meshgrid(np.arange(0       ,   image.shape[1]						)       ,   np.arange(0       ,   image.shape[0]						)						)       ,   2						)
 __lowerCAmelCase:  int        = np.round(grid - flow						).astype(np.intaa						)

 # Find the locations outside of the original image
 __lowerCAmelCase:  Dict        = (grid < 0) | (grid >= np.array([image.shape[1], image.shape[0]]						))
 __lowerCAmelCase:  str        = 0

 __lowerCAmelCase:  List[str]        = image[grid[:, :, 1], grid[:, :, 0]]

 # Set pixels at invalid locations to 0
 __lowerCAmelCase:  List[str]        = 0

 return warped







def _a					(			SCREAMING_SNAKE_CASE						:     np.ndarray       ,   SCREAMING_SNAKE_CASE						:     np.ndarray       ,   SCREAMING_SNAKE_CASE						:     SupportsIndex       ,   SCREAMING_SNAKE_CASE						:     float | None = None       ,   )       ->       tuple[np.ndarray, np.ndarray]:

 """simple docstring"""






 if alpha is None:
  __lowerCAmelCase:  Tuple        = 0.1

 # Initialize flow
 __lowerCAmelCase:  Optional[int]        = np.zeros_like(SCREAMING_SNAKE_CASE						)
 __lowerCAmelCase:  Dict        = np.zeros_like(SCREAMING_SNAKE_CASE						)

 # Prepare kernels for the calculation of the derivatives and the average velocity
 __lowerCAmelCase:  List[Any]        = np.array([[-1, 1], [-1, 1]]						) * 0.2_5
 __lowerCAmelCase:  int        = np.array([[-1, -1], [1, 1]]						) * 0.2_5
 __lowerCAmelCase:  Tuple        = np.array([[1, 1], [1, 1]]						) * 0.2_5
 __lowerCAmelCase:  Dict        = np.array(
     [[1 / 12, 1 / 6, 1 / 12], [1 / 6, 0, 1 / 6], [1 / 12, 1 / 6, 1 / 12]]						)

 # Iteratively refine the flow
 for _ in range(SCREAMING_SNAKE_CASE						):
  __lowerCAmelCase:  Optional[int]        = warp(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE						)
  __lowerCAmelCase:  List[str]        = convolve(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE						) + convolve(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE						)
  __lowerCAmelCase:  int        = convolve(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE						) + convolve(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE						)
  __lowerCAmelCase:  List[Any]        = convolve(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE						) + convolve(SCREAMING_SNAKE_CASE       ,   -kernel_t						)

  __lowerCAmelCase:  Any        = convolve(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE						)
  __lowerCAmelCase:  Dict        = convolve(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE						)

  # This updates the flow as proposed in the paper (Step 12)
  __lowerCAmelCase:  Tuple        = (
      derivative_x * avg_horizontal_velocity
      + derivative_y * avg_vertical_velocity
      + derivative_t
  )
  __lowerCAmelCase:  int        = update / (alpha**2 + derivative_x**2 + derivative_y**2)

  __lowerCAmelCase:  List[Any]        = avg_horizontal_velocity - derivative_x * update
  __lowerCAmelCase:  Union[str, Any]        = avg_vertical_velocity - derivative_y * update

 return horizontal_flow, vertical_flow


if __name__ == "__main__":
  import doctest

  doctest.testmod()





