



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  (						lowerCAmelCase__	,						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					):



    '''simple docstring'''
    return 32




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    return 32




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    return self.time_input_dim * 4




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    return 8




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    torch.manual_seed(0					)
    __lowerCamelCase    				=     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    				=     CLIPVisionModel(__UpperCAmelCase					)
    return model




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     CLIPImageProcessor(
        crop_size=224						,					do_center_crop=__UpperCAmelCase						,					do_normalize=__UpperCAmelCase						,					do_resize=__UpperCAmelCase						,					image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073]						,					image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711]						,					resample=3						,					size=224						,					)

    return image_processor




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    torch.manual_seed(0					)

    __lowerCamelCase    				=     {
        '''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    				=     PriorTransformer(**__UpperCAmelCase					)
    return model




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    torch.manual_seed(0					)

    __lowerCamelCase    				=     {
        '''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    				=     ShapERenderer(**__UpperCAmelCase					)
    return model




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.dummy_prior
    __lowerCamelCase    				=     self.dummy_image_encoder
    __lowerCamelCase    				=     self.dummy_image_processor
    __lowerCamelCase    				=     self.dummy_renderer

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

    return components




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase=0					):



    '''simple docstring'''
    __lowerCamelCase    				=     floats_tensor((1, 3, 64, 64)						,					rng=random.Random(__UpperCAmelCase					)					).to(__UpperCAmelCase					)

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




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     '''cpu'''

    __lowerCamelCase    				=     self.get_dummy_components()

    __lowerCamelCase    				=     self.pipeline_class(**__UpperCAmelCase					)
    __lowerCamelCase    				=     pipe.to(__UpperCAmelCase					)

    pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

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

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

    __lowerCamelCase    				=     np.array(
        [
            0.00_039_216,
            0.00_039_216,
            0.00_039_216,
            0.00_039_216,
            0.00_039_216,
            0.00_039_216,
            0.00_039_216,
            0.00_039_216,
            0.00_039_216,
        ]					)

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




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    # 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					):



    '''simple docstring'''
    __lowerCamelCase    				=     torch_device == '''cpu'''
    __lowerCamelCase    				=     True
    self._test_inference_batch_single_identical(
        batch_size=2						,					test_max_difference=__UpperCAmelCase						,					relax_max_difference=__UpperCAmelCase						,					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.get_dummy_components()
    __lowerCamelCase    				=     self.pipeline_class(**__UpperCAmelCase					)
    __lowerCamelCase    				=     pipe.to(__UpperCAmelCase					)
    pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

    __lowerCamelCase    				=     1
    __lowerCamelCase    				=     2

    __lowerCamelCase    				=     self.get_dummy_inputs(__UpperCAmelCase					)

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

    __lowerCamelCase    				=     pipe(**__UpperCAmelCase						,					num_images_per_prompt=__UpperCAmelCase					)[0]

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



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




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()




 def 							lowerCamelCase   (    self					):



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

    __lowerCamelCase    				=     torch.Generator(device=__UpperCAmelCase					).manual_seed(0					)

    __lowerCamelCase    				=     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					)











from typing import Callable, Optional

from .. import Features
from ..packaged_modules.generator.generator import Generator
from .abc import AbstractDatasetInputStream



class 			__lowerCAmelCase  (						lowerCAmelCase__ ):




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase = None						,					__UpperCAmelCase = None						,					__UpperCAmelCase = False						,					__UpperCAmelCase = False						,					__UpperCAmelCase = None						,					__UpperCAmelCase = None						,					**__UpperCAmelCase						,					):



    '''simple docstring'''
    super().__init__(
        features=__UpperCAmelCase						,					cache_dir=__UpperCAmelCase						,					keep_in_memory=__UpperCAmelCase						,					streaming=__UpperCAmelCase						,					num_proc=__UpperCAmelCase						,					**__UpperCAmelCase						,					)
    __lowerCamelCase    				=     Generator(
        cache_dir=__UpperCAmelCase						,					features=__UpperCAmelCase						,					generator=__UpperCAmelCase						,					gen_kwargs=__UpperCAmelCase						,					**__UpperCAmelCase						,					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    # Build iterable dataset
    if self.streaming:
       __lowerCamelCase    				=     self.builder.as_streaming_dataset(split='''train'''					)
    # Build regular (map-style) dataset
    else:
       __lowerCamelCase    				=     None
       __lowerCamelCase    				=     None
       __lowerCamelCase    				=     None
       __lowerCamelCase    				=     None

       self.builder.download_and_prepare(
           download_config=__UpperCAmelCase						,					download_mode=__UpperCAmelCase						,					verification_mode=__UpperCAmelCase						,					base_path=__UpperCAmelCase						,					num_proc=self.num_proc						,					)
       __lowerCamelCase    				=     self.builder.as_dataset(
           split='''train'''						,					verification_mode=__UpperCAmelCase						,					in_memory=self.keep_in_memory					)
    return dataset






