





import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference



class    _SCREAMING_SNAKE_CASE						(	__SCREAMING_SNAKE_CASE    ,							unittest.TestCase    ):







 '''simple docstring'''
 lowercase_		=			ShapEImgaImgPipeline
 lowercase_		=			["image"]
 lowercase_		=			["image"]
 lowercase_		=			[
     "num_images_per_prompt",
     "num_inference_steps",
     "generator",
     "latents",
     "guidance_scale",
     "frame_size",
     "output_type",
     "return_dict",
 ]
 lowercase_		=			False

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



   '''simple docstring'''



   return 32

 @property
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->int:



   '''simple docstring'''



   return 32

 @property
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->Dict:



   '''simple docstring'''



   return self.time_input_dim * 4

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



   '''simple docstring'''



   return 8

 @property
 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->List[Any]:



   '''simple docstring'''



   torch.manual_seed(0)
   lowerCamelCase__:      List[str]			  =CLIPVisionConfig(
       hidden_size=self.text_embedder_hidden_size       , image_size=64       , projection_dim=self.text_embedder_hidden_size       , intermediate_size=37       , num_attention_heads=4       , num_channels=3       , num_hidden_layers=5       , patch_size=1       , )

   lowerCamelCase__:      int			  =CLIPVisionModel(UpperCAmelCase_)
   return model

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



   '''simple docstring'''



   lowerCamelCase__:      Union[str, Any]			  =CLIPImageProcessor(
       crop_size=224       , do_center_crop=UpperCAmelCase_       , do_normalize=UpperCAmelCase_       , do_resize=UpperCAmelCase_       , image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]       , image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]       , resample=3       , size=224       , )

   return image_processor

 @property
 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int])       ->int:



   '''simple docstring'''



   torch.manual_seed(0)

   lowerCamelCase__:      Union[str, Any]			  ={
       "num_attention_heads": 2,
       "attention_head_dim": 16,
       "embedding_dim": self.time_input_dim,
       "num_embeddings": 32,
       "embedding_proj_dim": self.text_embedder_hidden_size,
       "time_embed_dim": self.time_embed_dim,
       "num_layers": 1,
       "clip_embed_dim": self.time_input_dim * 2,
       "additional_embeddings": 0,
       "time_embed_act_fn": "gelu",
       "norm_in_type": "layer",
       "embedding_proj_norm_type": "layer",
       "encoder_hid_proj_type": None,
       "added_emb_type": None,
   }

   lowerCamelCase__:      Tuple			  =PriorTransformer(**UpperCAmelCase_)
   return model

 @property
 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->List[Any]:



   '''simple docstring'''



   torch.manual_seed(0)

   lowerCamelCase__:      Any			  ={
       "param_shapes": (
           (self.renderer_dim, 93),
           (self.renderer_dim, 8),
           (self.renderer_dim, 8),
           (self.renderer_dim, 8),
       ),
       "d_latent": self.time_input_dim,
       "d_hidden": self.renderer_dim,
       "n_output": 12,
       "background": (
           0.1,
           0.1,
           0.1,
       ),
   }
   lowerCamelCase__:      Optional[Any]			  =ShapERenderer(**UpperCAmelCase_)
   return model

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->List[str]:



   '''simple docstring'''



   lowerCamelCase__:      Dict			  =self.dummy_prior
   lowerCamelCase__:      Optional[Any]			  =self.dummy_image_encoder
   lowerCamelCase__:      Optional[int]			  =self.dummy_image_processor
   lowerCamelCase__:      Tuple			  =self.dummy_renderer

   lowerCamelCase__:      Any			  =HeunDiscreteScheduler(
       beta_schedule="exp"       , num_train_timesteps=1_024       , prediction_type="sample"       , use_karras_sigmas=UpperCAmelCase_       , clip_sample=UpperCAmelCase_       , clip_sample_range=1.0       , )
   lowerCamelCase__:      str			  ={
       "prior": prior,
       "image_encoder": image_encoder,
       "image_processor": image_processor,
       "renderer": renderer,
       "scheduler": scheduler,
   }

   return components

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str]       , UpperCAmelCase_				:   Tuple       , UpperCAmelCase_				:   List[Any]=0)       ->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:      Optional[Any]			  =floats_tensor((1, 3, 64, 64)       , rng=random.Random(UpperCAmelCase_)).to(UpperCAmelCase_)

   if str(UpperCAmelCase_).startswith("mps"):
     lowerCamelCase__:      int			  =torch.manual_seed(UpperCAmelCase_)
   else:
     lowerCamelCase__:      int			  =torch.Generator(device=UpperCAmelCase_).manual_seed(UpperCAmelCase_)
   lowerCamelCase__:      str			  ={
       "image": input_image,
       "generator": generator,
       "num_inference_steps": 1,
       "frame_size": 32,
       "output_type": "np",
   }
   return inputs

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:      List[str]			  ="cpu"

   lowerCamelCase__:      int			  =self.get_dummy_components()

   lowerCamelCase__:      int			  =self.pipeline_class(**UpperCAmelCase_)
   lowerCamelCase__:      List[str]			  =pipe.to(UpperCAmelCase_)

   pipe.set_progress_bar_config(disable=UpperCAmelCase_)

   lowerCamelCase__:      int			  =pipe(**self.get_dummy_inputs(UpperCAmelCase_))
   lowerCamelCase__:      int			  =output.images[0]
   lowerCamelCase__:      Tuple			  =image[0, -3:, -3:, -1]

   assert image.shape == (20, 32, 32, 3)

   lowerCamelCase__:      Dict			  =np.array(
       [
           0.0003_9216,
           0.0003_9216,
           0.0003_9216,
           0.0003_9216,
           0.0003_9216,
           0.0003_9216,
           0.0003_9216,
           0.0003_9216,
           0.0003_9216,
       ])

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

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



   '''simple docstring'''



   self._test_inference_batch_consistent(batch_sizes=[1, 2])

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



   '''simple docstring'''



   lowerCamelCase__:      List[Any]			  =torch_device == "cpu"
   lowerCamelCase__:      List[str]			  =True
   self._test_inference_batch_single_identical(
       batch_size=2       , test_max_difference=UpperCAmelCase_       , relax_max_difference=UpperCAmelCase_       , )







 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:      List[Any]			  =self.get_dummy_components()
   lowerCamelCase__:      List[Any]			  =self.pipeline_class(**UpperCAmelCase_)
   lowerCamelCase__:      List[str]			  =pipe.to(UpperCAmelCase_)
   pipe.set_progress_bar_config(disable=UpperCAmelCase_)

   lowerCamelCase__:      List[Any]			  =1
   lowerCamelCase__:      int			  =2

   lowerCamelCase__:      Tuple			  =self.get_dummy_inputs(UpperCAmelCase_)

   for key in inputs.keys():
     if key in self.batch_params:
       lowerCamelCase__:      Dict			  =batch_size * [inputs[key]]

   lowerCamelCase__:      Dict			  =pipe(**UpperCAmelCase_       , num_images_per_prompt=UpperCAmelCase_)[0]

   assert images.shape[0] == batch_size * num_images_per_prompt



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







 '''simple docstring'''

 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->Dict:



   '''simple docstring'''



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







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



   '''simple docstring'''



   lowerCamelCase__:      Any			  =load_image(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/shap_e/corgi.png")
   lowerCamelCase__:      Tuple			  =load_numpy(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
       "/shap_e/test_shap_e_img2img_out.npy")
   lowerCamelCase__:      str			  =ShapEImgaImgPipeline.from_pretrained("openai/shap-e-img2img")
   lowerCamelCase__:      int			  =pipe.to(UpperCAmelCase_)
   pipe.set_progress_bar_config(disable=UpperCAmelCase_)

   lowerCamelCase__:      int			  =torch.Generator(device=UpperCAmelCase_).manual_seed(0)

   lowerCamelCase__:      Any			  =pipe(
       UpperCAmelCase_       , generator=UpperCAmelCase_       , guidance_scale=3.0       , num_inference_steps=64       , frame_size=64       , output_type="np"       , ).images[0]

   assert images.shape == (20, 64, 64, 3)

   assert_mean_pixel_difference(UpperCAmelCase_       , UpperCAmelCase_)








import torch

from diffusers import DDPMParallelScheduler

from .test_schedulers import SchedulerCommonTest



class    _SCREAMING_SNAKE_CASE						(	__SCREAMING_SNAKE_CASE    ):







 '''simple docstring'''
 lowercase_		=			(DDPMParallelScheduler,)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Any       , **UpperCAmelCase_				:   Any)       ->Any:



   '''simple docstring'''



   lowerCamelCase__:      Any			  ={
       "num_train_timesteps": 1_000,
       "beta_start": 0.0001,
       "beta_end": 0.02,
       "beta_schedule": "linear",
       "variance_type": "fixed_small",
       "clip_sample": True,
   }

   config.update(**UpperCAmelCase_)
   return config

 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->Dict:



   '''simple docstring'''



   for timesteps in [1, 5, 100, 1_000]:
     self.check_over_configs(num_train_timesteps=UpperCAmelCase_)

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



   '''simple docstring'''



   for beta_start, beta_end in zip([0.0001, 0.001, 0.01, 0.1]       , [0.002, 0.02, 0.2, 2]):
     self.check_over_configs(beta_start=UpperCAmelCase_       , beta_end=UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->Any:



   '''simple docstring'''



   for schedule in ["linear", "squaredcos_cap_v2"]:
     self.check_over_configs(beta_schedule=UpperCAmelCase_)

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



   '''simple docstring'''



   for variance in ["fixed_small", "fixed_large", "other"]:
     self.check_over_configs(variance_type=UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int])       ->Optional[Any]:



   '''simple docstring'''



   for clip_sample in [True, False]:
     self.check_over_configs(clip_sample=UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->Tuple:



   '''simple docstring'''



   self.check_over_configs(thresholding=UpperCAmelCase_)
   for threshold in [0.5, 1.0, 2.0]:
     for prediction_type in ["epsilon", "sample", "v_prediction"]:
       self.check_over_configs(
           thresholding=UpperCAmelCase_       , prediction_type=UpperCAmelCase_       , sample_max_value=UpperCAmelCase_       , )

 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->Optional[int]:



   '''simple docstring'''



   for prediction_type in ["epsilon", "sample", "v_prediction"]:
     self.check_over_configs(prediction_type=UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->int:



   '''simple docstring'''



   for t in [0, 500, 999]:
     self.check_over_forward(time_step=UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->str:



   '''simple docstring'''



   lowerCamelCase__:      Dict			  =self.scheduler_classes[0]
   lowerCamelCase__:      Tuple			  =self.get_scheduler_config()
   lowerCamelCase__:      Any			  =scheduler_class(**UpperCAmelCase_)

   assert torch.sum(torch.abs(scheduler._get_variance(0) - 0.0)) < 1E-5
   assert torch.sum(torch.abs(scheduler._get_variance(487) - 0.0_0979)) < 1E-5
   assert torch.sum(torch.abs(scheduler._get_variance(999) - 0.02)) < 1E-5

 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->str:



   '''simple docstring'''



   lowerCamelCase__:      int			  =self.scheduler_classes[0]
   lowerCamelCase__:      Tuple			  =self.get_scheduler_config()
   lowerCamelCase__:      Tuple			  =scheduler_class(**UpperCAmelCase_)

   lowerCamelCase__:      str			  =len(UpperCAmelCase_)

   lowerCamelCase__:      Optional[int]			  =self.dummy_model()
   lowerCamelCase__:      int			  =self.dummy_sample_deter
   lowerCamelCase__:      Union[str, Any]			  =self.dummy_sample_deter + 0.1
   lowerCamelCase__:      Optional[Any]			  =self.dummy_sample_deter - 0.1

   lowerCamelCase__:      Optional[Any]			  =samplea.shape[0]
   lowerCamelCase__:      List[Any]			  =torch.stack([samplea, samplea, samplea]       , dim=0)
   lowerCamelCase__:      Union[str, Any]			  =torch.arange(UpperCAmelCase_)[0:3, None].repeat(1       , UpperCAmelCase_)

   lowerCamelCase__:      Optional[int]			  =model(samples.flatten(0       , 1)       , timesteps.flatten(0       , 1))
   lowerCamelCase__:      Tuple			  =scheduler.batch_step_no_noise(UpperCAmelCase_       , timesteps.flatten(0       , 1)       , samples.flatten(0       , 1))

   lowerCamelCase__:      List[str]			  =torch.sum(torch.abs(UpperCAmelCase_))
   lowerCamelCase__:      Any			  =torch.mean(torch.abs(UpperCAmelCase_))

   assert abs(result_sum.item() - 1153.1833) < 1E-2
   assert abs(result_mean.item() - 0.5005) < 1E-3

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



   '''simple docstring'''



   lowerCamelCase__:      Any			  =self.scheduler_classes[0]
   lowerCamelCase__:      Optional[Any]			  =self.get_scheduler_config()
   lowerCamelCase__:      Optional[int]			  =scheduler_class(**UpperCAmelCase_)

   lowerCamelCase__:      Union[str, Any]			  =len(UpperCAmelCase_)

   lowerCamelCase__:      Union[str, Any]			  =self.dummy_model()
   lowerCamelCase__:      List[Any]			  =self.dummy_sample_deter
   lowerCamelCase__:      int			  =torch.manual_seed(0)

   for t in reversed(range(UpperCAmelCase_)):
     # 1. predict noise residual
     lowerCamelCase__:      Tuple			  =model(UpperCAmelCase_       , UpperCAmelCase_)

     # 2. predict previous mean of sample x_t-1
     lowerCamelCase__:      Optional[Any]			  =scheduler.step(UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , generator=UpperCAmelCase_).prev_sample

     lowerCamelCase__:      Any			  =pred_prev_sample

   lowerCamelCase__:      Any			  =torch.sum(torch.abs(UpperCAmelCase_))
   lowerCamelCase__:      List[str]			  =torch.mean(torch.abs(UpperCAmelCase_))

   assert abs(result_sum.item() - 258.9606) < 1E-2
   assert abs(result_mean.item() - 0.3372) < 1E-3

 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->Any:



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =self.scheduler_classes[0]
   lowerCamelCase__:      Any			  =self.get_scheduler_config(prediction_type="v_prediction")
   lowerCamelCase__:      Any			  =scheduler_class(**UpperCAmelCase_)

   lowerCamelCase__:      str			  =len(UpperCAmelCase_)

   lowerCamelCase__:      str			  =self.dummy_model()
   lowerCamelCase__:      str			  =self.dummy_sample_deter
   lowerCamelCase__:      Dict			  =torch.manual_seed(0)

   for t in reversed(range(UpperCAmelCase_)):
     # 1. predict noise residual
     lowerCamelCase__:      Union[str, Any]			  =model(UpperCAmelCase_       , UpperCAmelCase_)

     # 2. predict previous mean of sample x_t-1
     lowerCamelCase__:      Dict			  =scheduler.step(UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , generator=UpperCAmelCase_).prev_sample

     lowerCamelCase__:      List[str]			  =pred_prev_sample

   lowerCamelCase__:      List[Any]			  =torch.sum(torch.abs(UpperCAmelCase_))
   lowerCamelCase__:      Tuple			  =torch.mean(torch.abs(UpperCAmelCase_))

   assert abs(result_sum.item() - 202.0296) < 1E-2
   assert abs(result_mean.item() - 0.2631) < 1E-3

 def 							SCREAMING_SNAKE_CASE_					(self				:   Tuple)       ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:      str			  =self.scheduler_classes[0]
   lowerCamelCase__:      Union[str, Any]			  =self.get_scheduler_config()
   lowerCamelCase__:      Any			  =scheduler_class(**UpperCAmelCase_)

   lowerCamelCase__:      List[Any]			  =[100, 87, 50, 1, 0]

   scheduler.set_timesteps(timesteps=UpperCAmelCase_)

   lowerCamelCase__:      Union[str, Any]			  =scheduler.timesteps

   for i, timestep in enumerate(UpperCAmelCase_):
     if i == len(UpperCAmelCase_) - 1:
       lowerCamelCase__:      Dict			  =-1
     else:
       lowerCamelCase__:      Union[str, Any]			  =timesteps[i + 1]

     lowerCamelCase__:      Tuple			  =scheduler.previous_timestep(UpperCAmelCase_)
     lowerCamelCase__:      str			  =prev_t.item()

     self.assertEqual(UpperCAmelCase_       , UpperCAmelCase_)

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



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =self.scheduler_classes[0]
   lowerCamelCase__:      List[Any]			  =self.get_scheduler_config()
   lowerCamelCase__:      Dict			  =scheduler_class(**UpperCAmelCase_)

   lowerCamelCase__:      Optional[Any]			  =[100, 87, 50, 51, 0]

   with self.assertRaises(UpperCAmelCase_       , msg="`custom_timesteps` must be in descending order."):
     scheduler.set_timesteps(timesteps=UpperCAmelCase_)

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



   '''simple docstring'''



   lowerCamelCase__:      Dict			  =self.scheduler_classes[0]
   lowerCamelCase__:      Any			  =self.get_scheduler_config()
   lowerCamelCase__:      int			  =scheduler_class(**UpperCAmelCase_)

   lowerCamelCase__:      Optional[int]			  =[100, 87, 50, 1, 0]
   lowerCamelCase__:      int			  =len(UpperCAmelCase_)

   with self.assertRaises(UpperCAmelCase_       , msg="Can only pass one of `num_inference_steps` or `custom_timesteps`."):
     scheduler.set_timesteps(num_inference_steps=UpperCAmelCase_       , timesteps=UpperCAmelCase_)







 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->Any:



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =self.scheduler_classes[0]
   lowerCamelCase__:      Optional[Any]			  =self.get_scheduler_config()
   lowerCamelCase__:      Optional[Any]			  =scheduler_class(**UpperCAmelCase_)

   lowerCamelCase__:      Dict			  =[scheduler.config.num_train_timesteps]

   with self.assertRaises(
       UpperCAmelCase_       , msg="`timesteps` must start before `self.config.train_timesteps`: {scheduler.config.num_train_timesteps}}"       , ):
     scheduler.set_timesteps(timesteps=UpperCAmelCase_)

