

"""simple docstring"""






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        _lowercase					(     UpperCamelCase_   , unittest.TestCase     ):
     _lowerCamelCase     =       ShapEImgaImgPipeline
     _lowerCamelCase     =       ["""image"""]
     _lowerCamelCase     =       ["""image"""]
     _lowerCamelCase     =       [
         """num_images_per_prompt""",
         """num_inference_steps""",
         """generator""",
         """latents""",
         """guidance_scale""",
         """frame_size""",
         """output_type""",
         """return_dict""",
     ]
     _lowerCamelCase     =       False







     @property
     def 	lowerCAmelCase__    (      self       ):
      return 32







     @property
     def 	lowerCAmelCase__    (      self       ):
      return 32







     @property
     def 	lowerCAmelCase__    (      self       ):
      return self.time_input_dim * 4







     @property
     def 	lowerCAmelCase__    (      self       ):
      return 8







     @property
     def 	lowerCAmelCase__    (      self       ):
      torch.manual_seed(0       )
      __magic_name__     =      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 ,						)

      __magic_name__     =      CLIPVisionModel(_a       )
      return model







     @property
     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      CLIPImageProcessor(
          crop_size=224 ,						do_center_crop=_a ,						do_normalize=_a ,						do_resize=_a ,						image_mean=[0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3] ,						image_std=[0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1] ,						resample=3 ,						size=224 ,						)

      return image_processor







     @property
     def 	lowerCAmelCase__    (      self       ):
      torch.manual_seed(0       )

      __magic_name__     =      {
          """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,
      }

      __magic_name__     =      PriorTransformer(**_a       )
      return model







     @property
     def 	lowerCAmelCase__    (      self       ):
      torch.manual_seed(0       )

      __magic_name__     =      {
          """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,
          ),
      }
      __magic_name__     =      ShapERenderer(**_a       )
      return model







     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      self.dummy_prior
      __magic_name__     =      self.dummy_image_encoder
      __magic_name__     =      self.dummy_image_processor
      __magic_name__     =      self.dummy_renderer

      __magic_name__     =      HeunDiscreteScheduler(
          beta_schedule='''exp''' ,						num_train_timesteps=1024 ,						prediction_type='''sample''' ,						use_karras_sigmas=_a ,						clip_sample=_a ,						clip_sample_range=1.0 ,						)
      __magic_name__     =      {
          """prior""": prior,
          """image_encoder""": image_encoder,
          """image_processor""": image_processor,
          """renderer""": renderer,
          """scheduler""": scheduler,
      }

      return components







     def 	lowerCAmelCase__    (      self ,						UpperCamelCase_ ,						UpperCamelCase_=0       ):
      __magic_name__     =      floats_tensor((1, 3, 64, 64) ,						rng=random.Random(_a       )       ).to(_a       )

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







     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      """cpu"""

      __magic_name__     =      self.get_dummy_components()

      __magic_name__     =      self.pipeline_class(**_a       )
      __magic_name__     =      pipe.to(_a       )

      pipe.set_progress_bar_config(disable=_a       )

      __magic_name__     =      pipe(**self.get_dummy_inputs(_a       )       )
      __magic_name__     =      output.images[0]
      __magic_name__     =      image[0, -3:, -3:, -1]

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

      __magic_name__     =      np.array(
          [
              0.0_0_0_3_9_2_1_6,
              0.0_0_0_3_9_2_1_6,
              0.0_0_0_3_9_2_1_6,
              0.0_0_0_3_9_2_1_6,
              0.0_0_0_3_9_2_1_6,
              0.0_0_0_3_9_2_1_6,
              0.0_0_0_3_9_2_1_6,
              0.0_0_0_3_9_2_1_6,
              0.0_0_0_3_9_2_1_6,
          ]       )

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







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







     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      torch_device == """cpu"""
      __magic_name__     =      True
      self._test_inference_batch_single_identical(
          batch_size=2 ,						test_max_difference=_a ,						relax_max_difference=_a ,						)







     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      self.get_dummy_components()
      __magic_name__     =      self.pipeline_class(**_a       )
      __magic_name__     =      pipe.to(_a       )
      pipe.set_progress_bar_config(disable=_a       )

      __magic_name__     =      1
      __magic_name__     =      2

      __magic_name__     =      self.get_dummy_inputs(_a       )

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

      __magic_name__     =      pipe(**_a ,						num_images_per_prompt=_a       )[0]

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


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







     def 	lowerCAmelCase__    (      self       ):
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()







     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      load_image(
          '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/shap_e/corgi.png'''       )
      __magic_name__     =      load_numpy(
          '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
          '''/shap_e/test_shap_e_img2img_out.npy'''       )
      __magic_name__     =      ShapEImgaImgPipeline.from_pretrained('''openai/shap-e-img2img'''       )
      __magic_name__     =      pipe.to(_a       )
      pipe.set_progress_bar_config(disable=_a       )

      __magic_name__     =      torch.Generator(device=_a       ).manual_seed(0       )

      __magic_name__     =      pipe(
          _a ,						generator=_a ,						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(_a ,						_a       )




"""simple docstring"""






import torch

from diffusers import EulerDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest
class        _lowercase					(     __UpperCAmelCase     ):
     _lowerCamelCase     =       (EulerDiscreteScheduler,)
     _lowerCamelCase     =       10







     def 	lowerCAmelCase__    (      self ,						**UpperCamelCase_       ):
      __magic_name__     =      {
          '''num_train_timesteps''': 1100,
          '''beta_start''': 0.0_0_0_1,
          '''beta_end''': 0.0_2,
          '''beta_schedule''': '''linear''',
      }

      config.update(**UpperCamelCase_       )
      return config







     def 	lowerCAmelCase__    (      self       ):
      for timesteps in [10, 50, 100, 1000]:
       self.check_over_configs(num_train_timesteps=UpperCamelCase_       )







     def 	lowerCAmelCase__    (      self       ):
      for beta_start, beta_end in zip([0.0_0_0_0_1, 0.0_0_0_1, 0.0_0_1] ,						[0.0_0_0_2, 0.0_0_2, 0.0_2]       ):
       self.check_over_configs(beta_start=UpperCamelCase_ ,						beta_end=UpperCamelCase_       )







     def 	lowerCAmelCase__    (      self       ):
      for schedule in ["linear", "scaled_linear"]:
       self.check_over_configs(beta_schedule=UpperCamelCase_       )







     def 	lowerCAmelCase__    (      self       ):
      for prediction_type in ["epsilon", "v_prediction"]:
       self.check_over_configs(prediction_type=UpperCamelCase_       )







     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      self.scheduler_classes[0]
      __magic_name__     =      self.get_scheduler_config()
      __magic_name__     =      scheduler_class(**UpperCamelCase_       )

      scheduler.set_timesteps(self.num_inference_steps       )

      __magic_name__     =      torch.manual_seed(0       )

      __magic_name__     =      self.dummy_model()
      __magic_name__     =      self.dummy_sample_deter * scheduler.init_noise_sigma
      __magic_name__     =      sample.to(UpperCamelCase_       )

      for i, t in enumerate(scheduler.timesteps       ):
       __magic_name__     =      scheduler.scale_model_input(UpperCamelCase_ ,						UpperCamelCase_       )

       __magic_name__     =      model(UpperCamelCase_ ,						UpperCamelCase_       )

       __magic_name__     =      scheduler.step(UpperCamelCase_ ,						UpperCamelCase_ ,						UpperCamelCase_ ,						generator=UpperCamelCase_       )
       __magic_name__     =      output.prev_sample

      __magic_name__     =      torch.sum(torch.abs(UpperCamelCase_       )       )
      __magic_name__     =      torch.mean(torch.abs(UpperCamelCase_       )       )

      assert abs(result_sum.item() - 1_0.0_8_0_7       ) < 1E-2
      assert abs(result_mean.item() - 0.0_1_3_1       ) < 1E-3







     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      self.scheduler_classes[0]
      __magic_name__     =      self.get_scheduler_config(prediction_type='''v_prediction'''       )
      __magic_name__     =      scheduler_class(**UpperCamelCase_       )

      scheduler.set_timesteps(self.num_inference_steps       )

      __magic_name__     =      torch.manual_seed(0       )

      __magic_name__     =      self.dummy_model()
      __magic_name__     =      self.dummy_sample_deter * scheduler.init_noise_sigma
      __magic_name__     =      sample.to(UpperCamelCase_       )

      for i, t in enumerate(scheduler.timesteps       ):
       __magic_name__     =      scheduler.scale_model_input(UpperCamelCase_ ,						UpperCamelCase_       )

       __magic_name__     =      model(UpperCamelCase_ ,						UpperCamelCase_       )

       __magic_name__     =      scheduler.step(UpperCamelCase_ ,						UpperCamelCase_ ,						UpperCamelCase_ ,						generator=UpperCamelCase_       )
       __magic_name__     =      output.prev_sample

      __magic_name__     =      torch.sum(torch.abs(UpperCamelCase_       )       )
      __magic_name__     =      torch.mean(torch.abs(UpperCamelCase_       )       )

      assert abs(result_sum.item() - 0.0_0_0_2       ) < 1E-2
      assert abs(result_mean.item() - 2.2_6_7_6E-0_6       ) < 1E-3







     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      self.scheduler_classes[0]
      __magic_name__     =      self.get_scheduler_config()
      __magic_name__     =      scheduler_class(**UpperCamelCase_       )

      scheduler.set_timesteps(self.num_inference_steps ,						device=UpperCamelCase_       )

      __magic_name__     =      torch.manual_seed(0       )

      __magic_name__     =      self.dummy_model()
      __magic_name__     =      self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
      __magic_name__     =      sample.to(UpperCamelCase_       )

      for t in scheduler.timesteps:
       __magic_name__     =      scheduler.scale_model_input(UpperCamelCase_ ,						UpperCamelCase_       )

       __magic_name__     =      model(UpperCamelCase_ ,						UpperCamelCase_       )

       __magic_name__     =      scheduler.step(UpperCamelCase_ ,						UpperCamelCase_ ,						UpperCamelCase_ ,						generator=UpperCamelCase_       )
       __magic_name__     =      output.prev_sample

      __magic_name__     =      torch.sum(torch.abs(UpperCamelCase_       )       )
      __magic_name__     =      torch.mean(torch.abs(UpperCamelCase_       )       )

      assert abs(result_sum.item() - 1_0.0_8_0_7       ) < 1E-2
      assert abs(result_mean.item() - 0.0_1_3_1       ) < 1E-3







     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      self.scheduler_classes[0]
      __magic_name__     =      self.get_scheduler_config()
      __magic_name__     =      scheduler_class(**UpperCamelCase_ ,						use_karras_sigmas=UpperCamelCase_       )

      scheduler.set_timesteps(self.num_inference_steps ,						device=UpperCamelCase_       )

      __magic_name__     =      torch.manual_seed(0       )

      __magic_name__     =      self.dummy_model()
      __magic_name__     =      self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
      __magic_name__     =      sample.to(UpperCamelCase_       )

      for t in scheduler.timesteps:
       __magic_name__     =      scheduler.scale_model_input(UpperCamelCase_ ,						UpperCamelCase_       )

       __magic_name__     =      model(UpperCamelCase_ ,						UpperCamelCase_       )

       __magic_name__     =      scheduler.step(UpperCamelCase_ ,						UpperCamelCase_ ,						UpperCamelCase_ ,						generator=UpperCamelCase_       )
       __magic_name__     =      output.prev_sample

      __magic_name__     =      torch.sum(torch.abs(UpperCamelCase_       )       )
      __magic_name__     =      torch.mean(torch.abs(UpperCamelCase_       )       )

      assert abs(result_sum.item() - 1_2_4.5_2_2_9_9_4_9_9_5_1_1_7_1_9       ) < 1E-2
      assert abs(result_mean.item() - 0.1_6_2_1_3_9_3_2_6_3_3_3_9_9_9_6_3       ) < 1E-3

