




"""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_	(					__lowerCAmelCase      ,     unittest.TestCase							):



 '''simple docstring'''
 UpperCAmelCase   :					int				  =     ShapEImgaImgPipeline
 UpperCAmelCase   :					Dict				  =     ['''image''']
 UpperCAmelCase   :					str				  =     ['''image''']
 UpperCAmelCase   :					List[Any]				  =     [
     '''num_images_per_prompt''',
     '''num_inference_steps''',
     '''generator''',
     '''latents''',
     '''guidance_scale''',
     '''frame_size''',
     '''output_type''',
     '''return_dict''',
 ]
 UpperCAmelCase   :					int				  =     False

 @property
 def       lowerCAmelCase_							(			self	:	Union[str, Any]		):
       return 32

 @property
 def       lowerCAmelCase_							(			self	:	Optional[Any]		):
       return 32

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

 @property
 def       lowerCAmelCase_							(			self	:	List[Any]		):
       return 8

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

       _A        =       CLIPVisionModel(_UpperCAmelCase		)
       return model

 @property
 def       lowerCAmelCase_							(			self	:	List[Any]		):
       _A        =       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       lowerCAmelCase_							(			self	:	str		):
       torch.manual_seed(0		)

       _A        =       {
           '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,
       }

       _A        =       PriorTransformer(**_UpperCAmelCase		)
       return model

 @property
 def       lowerCAmelCase_							(			self	:	Any		):
       torch.manual_seed(0		)

       _A        =       {
           '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,
           ),
       }
       _A        =       ShapERenderer(**_UpperCAmelCase		)
       return model

 def       lowerCAmelCase_							(			self	:	List[str]		):
       _A        =       self.dummy_prior
       _A        =       self.dummy_image_encoder
       _A        =       self.dummy_image_processor
       _A        =       self.dummy_renderer

       _A        =       HeunDiscreteScheduler(
           beta_schedule='exp'   ,							num_train_timesteps=1_024   ,							prediction_type='sample'   ,							use_karras_sigmas=_UpperCAmelCase   ,							clip_sample=_UpperCAmelCase   ,							clip_sample_range=1.0   ,							)
       _A        =       {
           'prior': prior,
           'image_encoder': image_encoder,
           'image_processor': image_processor,
           'renderer': renderer,
           'scheduler': scheduler,
       }

       return components

 def       lowerCAmelCase_							(			self	:	Optional[int]   ,							_UpperCAmelCase	:	Union[str, Any]   ,							_UpperCAmelCase	:	Optional[int]=0		):
       _A        =       floats_tensor((1, 3, 64, 64)   ,							rng=random.Random(_UpperCAmelCase		)		).to(_UpperCAmelCase		)

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

 def       lowerCAmelCase_							(			self	:	Dict		):
       _A        =       'cpu'

       _A        =       self.get_dummy_components()

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

       pipe.set_progress_bar_config(disable=_UpperCAmelCase		)

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

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

       _A        =       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       lowerCAmelCase_							(			self	:	Dict		):
       # 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	:	List[Any]		):
       _A        =       torch_device == 'cpu'
       _A        =       True
       self._test_inference_batch_single_identical(
           batch_size=2   ,							test_max_difference=_UpperCAmelCase   ,							relax_max_difference=_UpperCAmelCase   ,							)



 def       lowerCAmelCase_							(			self	:	int		):
       _A        =       self.get_dummy_components()
       _A        =       self.pipeline_class(**_UpperCAmelCase		)
       _A        =       pipe.to(_UpperCAmelCase		)
       pipe.set_progress_bar_config(disable=_UpperCAmelCase		)

       _A        =       1
       _A        =       2

       _A        =       self.get_dummy_inputs(_UpperCAmelCase		)

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

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

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





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



 '''simple docstring'''

 def       lowerCAmelCase_							(			self	:	str		):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



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

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

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







"""simple docstring"""

import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Union

from .generation.configuration_utils import GenerationConfig
from .training_args import TrainingArguments
from .utils import add_start_docstrings


a         =      logging.getLogger(__name__)





@dataclass
@add_start_docstrings(TrainingArguments.__doc__							)
class     lowercase_	(					__lowerCAmelCase							):



 '''simple docstring'''
 UpperCAmelCase   :					bool				  =     field(default=__lowerCAmelCase      ,     metadata={'''help''': '''Whether to use SortishSampler or not.'''}							)
 UpperCAmelCase   :					bool				  =     field(
     default=__lowerCAmelCase      ,     metadata={'''help''': '''Whether to use generate to calculate generative metrics (ROUGE, BLEU).'''}							)
 UpperCAmelCase   :					Optional[int]				  =     field(
     default=__lowerCAmelCase      ,     metadata={
         '''help''': (
             '''The `max_length` to use on each evaluation loop when `predict_with_generate=True`. Will default '''
             '''to the `max_length` value of the model configuration.'''
         )
     }      ,     )
 UpperCAmelCase   :					Optional[int]				  =     field(
     default=__lowerCAmelCase      ,     metadata={
         '''help''': (
             '''The `num_beams` to use on each evaluation loop when `predict_with_generate=True`. Will default '''
             '''to the `num_beams` value of the model configuration.'''
         )
     }      ,     )
 UpperCAmelCase   :					Optional[Union[str, Path, GenerationConfig]]				  =     field(
     default=__lowerCAmelCase      ,     metadata={
         '''help''': '''Model id, file path or url pointing to a GenerationConfig json file, to use during prediction.'''
     }      ,     )

 def       lowerCAmelCase_							(			self	:	int		):
       _A        =       super().to_dict()
       for k, v in d.items():
             if isinstance(_UpperCAmelCase   ,							_UpperCAmelCase		):
                   _A        =       v.to_dict()
       return d

