




import gc
import unittest

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

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import 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       a__   ( snake_case__			,    unittest.TestCase				):
						_a							:   Optional[Any] =   ShapEPipeline
						_a							:   Union[str, Any] =   ["""prompt"""]
						_a							:   str =   ["""prompt"""]
						_a							:   Optional[int] =   [
						    """num_images_per_prompt""",
						    """num_inference_steps""",
						    """generator""",
						    """latents""",
						    """guidance_scale""",
						    """frame_size""",
						    """output_type""",
						    """return_dict""",
						]
						_a							:   List[Any] =   False






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return 3_2






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return 3_2






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return self.time_input_dim * 4






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return 8






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"      )
											return tokenizer






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											torch.manual_seed(0      )
											__lowerCAmelCase           =	CLIPTextConfig(
											    bos_token_id=0		,      eos_token_id=2		,      hidden_size=self.text_embedder_hidden_size		,      projection_dim=self.text_embedder_hidden_size		,      intermediate_size=3_7		,      layer_norm_eps=1E-0_5		,      num_attention_heads=4		,      num_hidden_layers=5		,      pad_token_id=1		,      vocab_size=1_0_0_0		,      )
											return CLIPTextModelWithProjection(_A      )






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											torch.manual_seed(0      )

											__lowerCAmelCase           =	{
											    "num_attention_heads": 2,
											    "attention_head_dim": 1_6,
											    "embedding_dim": self.time_input_dim,
											    "num_embeddings": 3_2,
											    "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",
											    "encoder_hid_proj_type": None,
											    "added_emb_type": None,
											}

											__lowerCAmelCase           =	PriorTransformer(**_A      )
											return model






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											torch.manual_seed(0      )

											__lowerCAmelCase           =	{
											    "param_shapes": (
											        (self.renderer_dim, 9_3),
											        (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": 1_2,
											    "background": (
											        0.1,
											        0.1,
											        0.1,
											    ),
											}
											__lowerCAmelCase           =	ShapERenderer(**_A      )
											return model






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.dummy_prior
											__lowerCAmelCase           =	self.dummy_text_encoder
											__lowerCAmelCase           =	self.dummy_tokenizer
											__lowerCAmelCase           =	self.dummy_renderer

											__lowerCAmelCase           =	HeunDiscreteScheduler(
											    beta_schedule="exp"		,      num_train_timesteps=1_0_2_4		,      prediction_type="sample"		,      use_karras_sigmas=_A		,      clip_sample=_A		,      clip_sample_range=1.0		,      )
											__lowerCAmelCase           =	{
											    "prior": prior,
											    "text_encoder": text_encoder,
											    "tokenizer": tokenizer,
											    "renderer": renderer,
											    "scheduler": scheduler,
											}

											return components






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A=0      ):




											"""simple docstring"""



											if str(_A      ).startswith("mps"      ):
																__lowerCAmelCase           =	torch.manual_seed(_A      )
											else:
																__lowerCAmelCase           =	torch.Generator(device=_A      ).manual_seed(_A      )
											__lowerCAmelCase           =	{
											    "prompt": "horse",
											    "generator": generator,
											    "num_inference_steps": 1,
											    "frame_size": 3_2,
											    "output_type": "np",
											}
											return inputs






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	"cpu"

											__lowerCAmelCase           =	self.get_dummy_components()

											__lowerCAmelCase           =	self.pipeline_class(**_A      )
											__lowerCAmelCase           =	pipe.to(_A      )

											pipe.set_progress_bar_config(disable=_A      )

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

											assert image.shape == (2_0, 3_2, 3_2, 3)

											__lowerCAmelCase           =	np.array(
											    [
											        0.00_03_92_16,
											        0.00_03_92_16,
											        0.00_03_92_16,
											        0.00_03_92_16,
											        0.00_03_92_16,
											        0.00_03_92_16,
											        0.00_03_92_16,
											        0.00_03_92_16,
											        0.00_03_92_16,
											    ]      )

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






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



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






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	torch_device == "cpu"
											__lowerCAmelCase           =	True

											self._test_inference_batch_single_identical(
											    batch_size=2		,      test_max_difference=_A		,      relax_max_difference=_A		,      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.get_dummy_components()
											__lowerCAmelCase           =	self.pipeline_class(**_A      )
											__lowerCAmelCase           =	pipe.to(_A      )
											pipe.set_progress_bar_config(disable=_A      )

											__lowerCAmelCase           =	1
											__lowerCAmelCase           =	2

											__lowerCAmelCase           =	self.get_dummy_inputs(_A      )

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

											__lowerCAmelCase           =	pipe(**_A		,      num_images_per_prompt=_A      )[0]

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






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






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



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






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	load_numpy(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    "/shap_e/test_shap_e_np_out.npy"      )
											__lowerCAmelCase           =	ShapEPipeline.from_pretrained("openai/shap-e"      )
											__lowerCAmelCase           =	pipe.to(_A      )
											pipe.set_progress_bar_config(disable=_A      )

											__lowerCAmelCase           =	torch.Generator(device=_A      ).manual_seed(0      )

											__lowerCAmelCase           =	pipe(
											    "a shark"		,      generator=_A		,      guidance_scale=15.0		,      num_inference_steps=6_4		,      frame_size=6_4		,      output_type="np"		,      ).images[0]

											assert images.shape == (2_0, 6_4, 6_4, 3)

											assert_mean_pixel_difference(_A		,      _A      )










import unittest

import numpy as np
import torch

from diffusers import ScoreSdeVePipeline, ScoreSdeVeScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()
class       a__   ( unittest.TestCase				):






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											torch.manual_seed(0      )
											__lowerCAmelCase           =	UNetaDModel(
											    block_out_channels=(3_2, 6_4)		,      layers_per_block=2		,      sample_size=3_2		,      in_channels=3		,      out_channels=3		,      down_block_types=("DownBlock2D", "AttnDownBlock2D")		,      up_block_types=("AttnUpBlock2D", "UpBlock2D")		,      )
											return model






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.dummy_uncond_unet
											__lowerCAmelCase           =	ScoreSdeVeScheduler()

											__lowerCAmelCase           =	ScoreSdeVePipeline(unet=_A		,      scheduler=_A      )
											sde_ve.to(_A      )
											sde_ve.set_progress_bar_config(disable=_A      )

											__lowerCAmelCase           =	torch.manual_seed(0      )
											__lowerCAmelCase           =	sde_ve(num_inference_steps=2		,      output_type="numpy"		,      generator=_A      ).images

											__lowerCAmelCase           =	torch.manual_seed(0      )
											__lowerCAmelCase           =	sde_ve(num_inference_steps=2		,      output_type="numpy"		,      generator=_A		,      return_dict=_A      )[
											    0
											]

											__lowerCAmelCase           =	image[0, -3:, -3:, -1]
											__lowerCAmelCase           =	image_from_tuple[0, -3:, -3:, -1]

											assert image.shape == (1, 3_2, 3_2, 3)
											__lowerCAmelCase           =	np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]      )

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






@slow
@require_torch
class       a__   ( unittest.TestCase				):






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	"google/ncsnpp-church-256"
											__lowerCAmelCase           =	UNetaDModel.from_pretrained(_A      )

											__lowerCAmelCase           =	ScoreSdeVeScheduler.from_pretrained(_A      )

											__lowerCAmelCase           =	ScoreSdeVePipeline(unet=_A		,      scheduler=_A      )
											sde_ve.to(_A      )
											sde_ve.set_progress_bar_config(disable=_A      )

											__lowerCAmelCase           =	torch.manual_seed(0      )
											__lowerCAmelCase           =	sde_ve(num_inference_steps=1_0		,      output_type="numpy"		,      generator=_A      ).images

											__lowerCAmelCase           =	image[0, -3:, -3:, -1]

											assert image.shape == (1, 2_5_6, 2_5_6, 3)

											__lowerCAmelCase           =	np.array([0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0]      )

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




