





'''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 			_lowerCAmelCase		(			A__							,     unittest.TestCase	):






     """simple docstring"""






     snake_case_						      =   ShapEImgaImgPipeline
     snake_case_						      =   ["image"]
     snake_case_						      =   ["image"]
     snake_case_						      =   [
         "num_images_per_prompt",
         "num_inference_steps",
         "generator",
         "latents",
         "guidance_scale",
         "frame_size",
         "output_type",
         "return_dict",
     ]
     snake_case_						      =   False



     @property
     def      lowerCAmelCase			(   self							:			Optional[int]						)->	Dict:
        return 32



     @property
     def      lowerCAmelCase			(   self							:			Tuple						)->	Any:
        return 32



     @property
     def      lowerCAmelCase			(   self							:			Optional[int]						)->	Any:
        return self.time_input_dim * 4



     @property
     def      lowerCAmelCase			(   self							:			Union[str, Any]						)->	Optional[int]:
        return 8



     @property
     def      lowerCAmelCase			(   self							:			List[Any]						)->	Optional[int]:
        torch.manual_seed(0						)
        snake_case 						=   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   ,       )

        snake_case 						=   CLIPVisionModel(__snake_case						)
        return model



     @property
     def      lowerCAmelCase			(   self							:			Any						)->	Dict:
        snake_case 						=   CLIPImageProcessor(
            crop_size=2_24   ,       do_center_crop=__snake_case   ,       do_normalize=__snake_case   ,       do_resize=__snake_case   ,       image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]   ,       image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]   ,       resample=3   ,       size=2_24   ,       )

        return image_processor



     @property
     def      lowerCAmelCase			(   self							:			Tuple						)->	List[Any]:
        torch.manual_seed(0						)

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

        snake_case 						=   PriorTransformer(**__snake_case						)
        return model



     @property
     def      lowerCAmelCase			(   self							:			Tuple						)->	Dict:
        torch.manual_seed(0						)

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



     def      lowerCAmelCase			(   self							:			Optional[Any]						)->	Dict:
        snake_case 						=   self.dummy_prior
        snake_case 						=   self.dummy_image_encoder
        snake_case 						=   self.dummy_image_processor
        snake_case 						=   self.dummy_renderer

        snake_case 						=   HeunDiscreteScheduler(
            beta_schedule="""exp"""   ,       num_train_timesteps=10_24   ,       prediction_type="""sample"""   ,       use_karras_sigmas=__snake_case   ,       clip_sample=__snake_case   ,       clip_sample_range=1.0   ,       )
        snake_case 						=   {
            """prior""": prior,
            """image_encoder""": image_encoder,
            """image_processor""": image_processor,
            """renderer""": renderer,
            """scheduler""": scheduler,
        }

        return components



     def      lowerCAmelCase			(   self							:			Optional[Any]   ,       __snake_case							:			Optional[Any]   ,       __snake_case							:			List[str]=0						)->	int:
        snake_case 						=   floats_tensor((1, 3, 64, 64)   ,       rng=random.Random(__snake_case						)						).to(__snake_case						)

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



     def      lowerCAmelCase			(   self							:			Optional[Any]						)->	Any:
        snake_case 						=   """cpu"""

        snake_case 						=   self.get_dummy_components()

        snake_case 						=   self.pipeline_class(**__snake_case						)
        snake_case 						=   pipe.to(__snake_case						)

        pipe.set_progress_bar_config(disable=__snake_case						)

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

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

        snake_case 						=   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      lowerCAmelCase			(   self							:			List[str]						)->	List[Any]:
        # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
        self._test_inference_batch_consistent(batch_sizes=[1, 2]						)



     def      lowerCAmelCase			(   self							:			Optional[Any]						)->	Dict:
        snake_case 						=   torch_device == """cpu"""
        snake_case 						=   True
        self._test_inference_batch_single_identical(
            batch_size=2   ,       test_max_difference=__snake_case   ,       relax_max_difference=__snake_case   ,       )



     def      lowerCAmelCase			(   self							:			Tuple						)->	Any:
        snake_case 						=   self.get_dummy_components()
        snake_case 						=   self.pipeline_class(**__snake_case						)
        snake_case 						=   pipe.to(__snake_case						)
        pipe.set_progress_bar_config(disable=__snake_case						)

        snake_case 						=   1
        snake_case 						=   2

        snake_case 						=   self.get_dummy_inputs(__snake_case						)

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

        snake_case 						=   pipe(**__snake_case   ,       num_images_per_prompt=__snake_case						)[0]

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




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






     """simple docstring"""






     def      lowerCAmelCase			(   self							:			Union[str, Any]						)->	List[Any]:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()



     def      lowerCAmelCase			(   self							:			Any						)->	int:
        snake_case 						=   load_image(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/shap_e/corgi.png"""						)
        snake_case 						=   load_numpy(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/shap_e/test_shap_e_img2img_out.npy"""						)
        snake_case 						=   ShapEImgaImgPipeline.from_pretrained("""openai/shap-e-img2img"""						)
        snake_case 						=   pipe.to(__snake_case						)
        pipe.set_progress_bar_config(disable=__snake_case						)

        snake_case 						=   torch.Generator(device=__snake_case						).manual_seed(0						)

        snake_case 						=   pipe(
            __snake_case   ,       generator=__snake_case   ,       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(__snake_case   ,       __snake_case						)








'''simple docstring'''




from . import (
    albert,
    align,
    altclip,
    audio_spectrogram_transformer,
    auto,
    autoformer,
    bark,
    bart,
    barthez,
    bartpho,
    beit,
    bert,
    bert_generation,
    bert_japanese,
    bertweet,
    big_bird,
    bigbird_pegasus,
    biogpt,
    bit,
    blenderbot,
    blenderbot_small,
    blip,
    blip_a,
    bloom,
    bridgetower,
    byta,
    camembert,
    canine,
    chinese_clip,
    clap,
    clip,
    clipseg,
    codegen,
    conditional_detr,
    convbert,
    convnext,
    convnextva,
    cpm,
    cpmant,
    ctrl,
    cvt,
    dataavec,
    deberta,
    deberta_va,
    decision_transformer,
    deformable_detr,
    deit,
    deprecated,
    deta,
    detr,
    dialogpt,
    dinat,
    distilbert,
    dit,
    donut,
    dpr,
    dpt,
    efficientformer,
    efficientnet,
    electra,
    encodec,
    encoder_decoder,
    ernie,
    ernie_m,
    esm,
    falcon,
    flaubert,
    flava,
    fnet,
    focalnet,
    fsmt,
    funnel,
    git,
    glpn,
    gpta,
    gpt_bigcode,
    gpt_neo,
    gpt_neox,
    gpt_neox_japanese,
    gpt_swa,
    gptj,
    gptsan_japanese,
    graphormer,
    groupvit,
    herbert,
    hubert,
    ibert,
    imagegpt,
    informer,
    instructblip,
    jukebox,
    layoutlm,
    layoutlmva,
    layoutlmva,
    layoutxlm,
    led,
    levit,
    lilt,
    llama,
    longformer,
    longta,
    luke,
    lxmert,
    mam_aaa,
    marian,
    markuplm,
    maskaformer,
    maskformer,
    mbart,
    mbartaa,
    mega,
    megatron_bert,
    megatron_gpta,
    mgp_str,
    mluke,
    mobilebert,
    mobilenet_va,
    mobilenet_va,
    mobilevit,
    mobilevitva,
    mpnet,
    mra,
    mta,
    musicgen,
    mvp,
    nat,
    nezha,
    nllb,
    nllb_moe,
    nystromformer,
    oneformer,
    open_llama,
    openai,
    opt,
    owlvit,
    pegasus,
    pegasus_x,
    perceiver,
    phobert,
    pixastruct,
    plbart,
    poolformer,
    prophetnet,
    qdqbert,
    rag,
    realm,
    reformer,
    regnet,
    rembert,
    resnet,
    roberta,
    roberta_prelayernorm,
    roc_bert,
    roformer,
    rwkv,
    sam,
    segformer,
    sew,
    sew_d,
    speech_encoder_decoder,
    speech_to_text,
    speech_to_text_a,
    speechta,
    splinter,
    squeezebert,
    swiftformer,
    swin,
    swinasr,
    swinva,
    switch_transformers,
    ta,
    table_transformer,
    tapas,
    time_series_transformer,
    timesformer,
    timm_backbone,
    transfo_xl,
    trocr,
    tvlt,
    umta,
    unispeech,
    unispeech_sat,
    upernet,
    videomae,
    vilt,
    vision_encoder_decoder,
    vision_text_dual_encoder,
    visual_bert,
    vit,
    vit_hybrid,
    vit_mae,
    vit_msn,
    vivit,
    wavaveca,
    wavaveca_conformer,
    wavaveca_phoneme,
    wavaveca_with_lm,
    wavlm,
    whisper,
    x_clip,
    xglm,
    xlm,
    xlm_prophetnet,
    xlm_roberta,
    xlm_roberta_xl,
    xlnet,
    xmod,
    yolos,
    yoso,
)

