

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






     '''simple docstring'''
     _SCREAMING_SNAKE_CASE						:Union[str, Any]       = ShapEImgaImgPipeline
     _SCREAMING_SNAKE_CASE						:Tuple       = ["""image"""]
     _SCREAMING_SNAKE_CASE						:Any       = ["""image"""]
     _SCREAMING_SNAKE_CASE						:int       = [
         """num_images_per_prompt""",
         """num_inference_steps""",
         """generator""",
         """latents""",
         """guidance_scale""",
         """frame_size""",
         """output_type""",
         """return_dict""",
     ]
     _SCREAMING_SNAKE_CASE						:Tuple       = False

     @property
     def 					_a			(						self  )  ->      Dict:


           """simple docstring"""


           return 32

     @property
     def 					_a			(						self  )  ->      Tuple:


           """simple docstring"""


           return 32

     @property
     def 					_a			(						self  )  ->      Any:


           """simple docstring"""


           return self.time_input_dim * 4

     @property
     def 					_a			(						self  )  ->      Any:


           """simple docstring"""


           return 8

     @property
     def 					_a			(						self  )  ->      int:


           """simple docstring"""


           torch.manual_seed(0  )
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	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	,    )

           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	CLIPVisionModel(_a  )
           return model

     @property
     def 					_a			(						self  )  ->      Dict:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Any      		=	CLIPImageProcessor(
               crop_size=224	,    do_center_crop=_a	,    do_normalize=_a	,    do_resize=_a	,    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 					_a			(						self  )  ->      List[str]:


           """simple docstring"""


           torch.manual_seed(0  )

           SCREAMING_SNAKE_CASE__   :      List[str]      		=	{
               """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,
           }

           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	PriorTransformer(**_a  )
           return model

     @property
     def 					_a			(						self  )  ->      Dict:


           """simple docstring"""


           torch.manual_seed(0  )

           SCREAMING_SNAKE_CASE__   :      Tuple      		=	{
               """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,
               ),
           }
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	ShapERenderer(**_a  )
           return model

     def 					_a			(						self  )  ->      List[str]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Any      		=	self.dummy_prior
           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	self.dummy_image_encoder
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	self.dummy_image_processor
           SCREAMING_SNAKE_CASE__   :      Dict      		=	self.dummy_renderer

           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	HeunDiscreteScheduler(
               beta_schedule="""exp"""	,    num_train_timesteps=1_024	,    prediction_type="""sample"""	,    use_karras_sigmas=_a	,    clip_sample=_a	,    clip_sample_range=1.0	,    )
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	{
               """prior""": prior,
               """image_encoder""": image_encoder,
               """image_processor""": image_processor,
               """renderer""": renderer,
               """scheduler""": scheduler,
           }

           return components

     def 					_a			(						self	,    _a	,    _a=0  )  ->      Union[str, Any]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	floats_tensor((1, 3, 64, 64)	,    rng=random.Random(_a  )  ).to(_a  )

           if str(_a  ).startswith("""mps"""  ):
                 SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	torch.manual_seed(_a  )
           else:
                 SCREAMING_SNAKE_CASE__   :      Tuple      		=	torch.Generator(device=_a  ).manual_seed(_a  )
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	{
               """image""": input_image,
               """generator""": generator,
               """num_inference_steps""": 1,
               """frame_size""": 32,
               """output_type""": """np""",
           }
           return inputs

     def 					_a			(						self  )  ->      int:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	"""cpu"""

           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	self.get_dummy_components()

           SCREAMING_SNAKE_CASE__   :      int      		=	self.pipeline_class(**_a  )
           SCREAMING_SNAKE_CASE__   :      Dict      		=	pipe.to(_a  )

           pipe.set_progress_bar_config(disable=_a  )

           SCREAMING_SNAKE_CASE__   :      List[str]      		=	pipe(**self.get_dummy_inputs(_a  )  )
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	output.images[0]
           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	image[0, -3:, -3:, -1]

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

           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	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 					_a			(						self  )  ->      Optional[int]:


           """simple docstring"""


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

     def 					_a			(						self  )  ->      Union[str, Any]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      str      		=	torch_device == """cpu"""
           SCREAMING_SNAKE_CASE__   :      Dict      		=	True
           self._test_inference_batch_single_identical(
               batch_size=2	,    test_max_difference=_a	,    relax_max_difference=_a	,    )





     def 					_a			(						self  )  ->      str:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Dict      		=	self.get_dummy_components()
           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	self.pipeline_class(**_a  )
           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	pipe.to(_a  )
           pipe.set_progress_bar_config(disable=_a  )

           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	1
           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	2

           SCREAMING_SNAKE_CASE__   :      Any      		=	self.get_dummy_inputs(_a  )

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

           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	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):






     '''simple docstring'''

     def 					_a			(						self  )  ->      Optional[int]:


           """simple docstring"""


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





     def 					_a			(						self  )  ->      Optional[int]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	load_image(
               """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/shap_e/corgi.png"""  )
           SCREAMING_SNAKE_CASE__   :      Dict      		=	load_numpy(
               """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
               """/shap_e/test_shap_e_img2img_out.npy"""  )
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	ShapEImgaImgPipeline.from_pretrained("""openai/shap-e-img2img"""  )
           SCREAMING_SNAKE_CASE__   :      Any      		=	pipe.to(_a  )
           pipe.set_progress_bar_config(disable=_a  )

           SCREAMING_SNAKE_CASE__   :      Any      		=	torch.Generator(device=_a  ).manual_seed(0  )

           SCREAMING_SNAKE_CASE__   :      Dict      		=	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 multiprocessing
from typing import TYPE_CHECKING, Optional, Union

from .. import Dataset, Features, config
from ..formatting import query_table
from ..packaged_modules.sql.sql import Sql
from ..utils import logging
from .abc import AbstractDatasetInputStream


if TYPE_CHECKING:
  import sqlitea

  import sqlalchemy

class      __a							(UpperCamelCase_):






     '''simple docstring'''

     def __init__(						self	,    _a	,    _a	,    _a = None	,    _a = None	,    _a = False	,    **_a	,    )  ->      Union[str, Any]:


           """simple docstring"""


           super().__init__(features=_a	,    cache_dir=_a	,    keep_in_memory=_a	,    **_a  )
           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	Sql(
               cache_dir=_a	,    features=_a	,    sql=_a	,    con=_a	,    **_a	,    )





     def 					_a			(						self  )  ->      int:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	None
           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	None
           SCREAMING_SNAKE_CASE__   :      Dict      		=	None
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	None

           self.builder.download_and_prepare(
               download_config=_a	,    download_mode=_a	,    verification_mode=_a	,    base_path=_a	,    )

           # Build dataset for splits
           SCREAMING_SNAKE_CASE__   :      str      		=	self.builder.as_dataset(
               split="""train"""	,    verification_mode=_a	,    in_memory=self.keep_in_memory  )
           return dataset





class      __a							:






     '''simple docstring'''

     def __init__(						self	,    _a	,    _a	,    _a	,    _a = None	,    _a = None	,    **_a	,    )  ->      Any:


           """simple docstring"""


           if num_proc is not None and num_proc <= 0:
                 raise ValueError(f'''num_proc {num_proc} must be an integer > 0.'''  )

           SCREAMING_SNAKE_CASE__   :      int      		=	dataset
           SCREAMING_SNAKE_CASE__   :      Any      		=	name
           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	con
           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
           SCREAMING_SNAKE_CASE__   :      int      		=	num_proc
           SCREAMING_SNAKE_CASE__   :      int      		=	to_sql_kwargs

     def 					_a			(						self  )  ->      int:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      str      		=	self.to_sql_kwargs.pop("""sql"""	,    _a  )
           SCREAMING_SNAKE_CASE__   :      Tuple      		=	self.to_sql_kwargs.pop("""con"""	,    _a  )
           SCREAMING_SNAKE_CASE__   :      Tuple      		=	self.to_sql_kwargs.pop("""index"""	,    _a  )

           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	self._write(index=_a	,    **self.to_sql_kwargs  )
           return written

     def 					_a			(						self	,    _a  )  ->      Union[str, Any]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__		,		SCREAMING_SNAKE_CASE__		,		SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	args
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	{**to_sql_kwargs, """if_exists""": """append"""} if offset > 0 else to_sql_kwargs
           SCREAMING_SNAKE_CASE__   :      Any      		=	query_table(
               table=self.dataset.data	,    key=slice(_a	,    offset + self.batch_size  )	,    indices=self.dataset._indices	,    )
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	batch.to_pandas()
           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	df.to_sql(self.name	,    self.con	,    index=_a	,    **_a  )
           return num_rows or len(_a  )





     def 					_a			(						self	,    _a	,    **_a  )  ->      int:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[str]      		=	0

           if self.num_proc is None or self.num_proc == 1:
                 for offset in logging.tqdm(
                     range(0	,    len(self.dataset  )	,    self.batch_size  )	,    unit="""ba"""	,    disable=not logging.is_progress_bar_enabled()	,    desc="""Creating SQL from Arrow format"""	,    ):
                       written += self._batch_sql((offset, index, to_sql_kwargs)  )
           else:
                 SCREAMING_SNAKE_CASE__		,		SCREAMING_SNAKE_CASE__   :      str      		=	len(self.dataset  ), self.batch_size
                 with multiprocessing.Pool(self.num_proc  ) as pool:
                       for num_rows in logging.tqdm(
                           pool.imap(
                               self._batch_sql	,    [(offset, index, to_sql_kwargs) for offset in range(0	,    _a	,    _a  )]	,    )	,    total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size	,    unit="""ba"""	,    disable=not logging.is_progress_bar_enabled()	,    desc="""Creating SQL from Arrow format"""	,    ):
                             written += num_rows

           return written




