





"""simple docstring"""






from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from transformers import CLIPImageProcessor, CLIPVisionModel

from ...models import PriorTransformer
from ...pipelines import DiffusionPipeline
from ...schedulers import HeunDiscreteScheduler
from ...utils import (
    BaseOutput,
    is_accelerate_available,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from .renderer import ShapERenderer


lowerCAmelCase_			=							logging.get_logger(__name__)  # pylint: disable=invalid-name

lowerCAmelCase_			=							"\n    Examples:\n        ```py\n        >>> from PIL import Image\n        >>> import torch\n        >>> from diffusers import DiffusionPipeline\n        >>> from diffusers.utils import export_to_gif, load_image\n\n        >>> device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n        >>> repo = \"openai/shap-e-img2img\"\n        >>> pipe = DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.float16)\n        >>> pipe = pipe.to(device)\n\n        >>> guidance_scale = 3.0\n        >>> image_url = \"https://hf.co/datasets/diffusers/docs-images/resolve/main/shap-e/corgi.png\"\n        >>> image = load_image(image_url).convert(\"RGB\")\n\n        >>> images = pipe(\n        ...     image,\n        ...     guidance_scale=guidance_scale,\n        ...     num_inference_steps=64,\n        ...     frame_size=256,\n        ... ).images\n\n        >>> gif_path = export_to_gif(images[0], \"corgi_3d.gif\")\n        ```\n"



@dataclass
class 	lowerCAmelCase					(      snake_case							):
       lowerCAmelCase__						 =  42



class 	lowerCAmelCase					(      snake_case							):






       def __init__(   self    ,  a__    ,  a__    ,  a__    ,  a__    ,  a__    ,  ):
              super().__init__()

              self.register_modules(
                  prior=UpperCAmelCase__    ,  image_encoder=UpperCAmelCase__    ,  image_processor=UpperCAmelCase__    ,  scheduler=UpperCAmelCase__    ,  renderer=UpperCAmelCase__    ,  )






       def 							__A (   self    ,  a__    ,  a__    ,  a__    ,  a__    ,  a__    ,  a__     ):
              if latents is None:
                     _UpperCAmelCase     =     randn_tensor(UpperCAmelCase__    ,  generator=UpperCAmelCase__    ,  device=UpperCAmelCase__    ,  dtype=UpperCAmelCase__     )
              else:
                     if latents.shape != shape:
                            raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {shape}"""     )
                     _UpperCAmelCase     =     latents.to(UpperCAmelCase__     )

              _UpperCAmelCase     =     latents * scheduler.init_noise_sigma
              return latents






       def 							__A (   self    ,  a__=0     ):
              if is_accelerate_available():
                     from accelerate import cpu_offload
              else:
                     raise ImportError('Please install accelerate via `pip install accelerate`'     )

              _UpperCAmelCase     =     torch.device(f"""cuda:{gpu_id}"""     )

              _UpperCAmelCase     =     [self.image_encoder, self.prior]
              for cpu_offloaded_model in models:
                     if cpu_offloaded_model is not None:
                            cpu_offload(UpperCAmelCase__    ,  UpperCAmelCase__     )






       @property
       def 							__A (   self     ):
              if self.device != torch.device('meta'     ) or not hasattr(self.image_encoder    ,  '_hf_hook'     ):
                     return self.device
              for module in self.image_encoder.modules():
                     if (
                         hasattr(UpperCAmelCase__    ,  '_hf_hook'     )
                         and hasattr(module._hf_hook    ,  'execution_device'     )
                         and module._hf_hook.execution_device is not None
                     ):
                            return torch.device(module._hf_hook.execution_device     )
              return self.device






       def 							__A (   self    ,  a__    ,  a__    ,  a__    ,  a__    ,  ):
              if isinstance(UpperCAmelCase__    ,  UpperCAmelCase__     ) and isinstance(image[0]    ,  torch.Tensor     ):
                     _UpperCAmelCase     =     torch.cat(UpperCAmelCase__    ,  axis=0     ) if image[0].ndim == 4 else torch.stack(UpperCAmelCase__    ,  axis=0     )

              if not isinstance(UpperCAmelCase__    ,  torch.Tensor     ):
                     _UpperCAmelCase     =     self.image_processor(UpperCAmelCase__    ,  return_tensors='pt'     ).pixel_values[0].unsqueeze(0     )

              _UpperCAmelCase     =     image.to(dtype=self.image_encoder.dtype    ,  device=UpperCAmelCase__     )

              _UpperCAmelCase     =     self.image_encoder(UpperCAmelCase__     )['last_hidden_state']
              _UpperCAmelCase     =     image_embeds[:, 1:, :].contiguous()  # batch_size, dim, 256

              _UpperCAmelCase     =     image_embeds.repeat_interleave(UpperCAmelCase__    ,  dim=0     )

              if do_classifier_free_guidance:
                     _UpperCAmelCase     =     torch.zeros_like(UpperCAmelCase__     )

                     # For classifier free guidance, we need to do two forward passes.
                     # Here we concatenate the unconditional and text embeddings into a single batch
                     # to avoid doing two forward passes
                     _UpperCAmelCase     =     torch.cat([negative_image_embeds, image_embeds]     )

              return image_embeds






       @torch.no_grad()
       @replace_example_docstring(UpperCAmelCase__     )
       def __call__(   self    ,  a__    ,  a__ = 1    ,  a__ = 25    ,  a__ = None    ,  a__ = None    ,  a__ = 4.0    ,  a__ = 64    ,  a__ = "pil"    ,  a__ = True    ,  ):

              if isinstance(UpperCAmelCase__    ,  PIL.Image.Image     ):
                     _UpperCAmelCase     =     1
              elif isinstance(UpperCAmelCase__    ,  torch.Tensor     ):
                     _UpperCAmelCase     =     image.shape[0]
              elif isinstance(UpperCAmelCase__    ,  UpperCAmelCase__     ) and isinstance(image[0]    ,  (torch.Tensor, PIL.Image.Image)     ):
                     _UpperCAmelCase     =     len(UpperCAmelCase__     )
              else:
                     raise ValueError(
                         f"""`image` has to be of type `PIL.Image.Image`, `torch.Tensor`, `List[PIL.Image.Image]` or `List[torch.Tensor]` but is {type(UpperCAmelCase__     )}"""     )

              _UpperCAmelCase     =     self._execution_device

              _UpperCAmelCase     =     batch_size * num_images_per_prompt

              _UpperCAmelCase     =     guidance_scale > 1.0
              _UpperCAmelCase     =     self._encode_image(UpperCAmelCase__    ,  UpperCAmelCase__    ,  UpperCAmelCase__    ,  UpperCAmelCase__     )

              # prior

              self.scheduler.set_timesteps(UpperCAmelCase__    ,  device=UpperCAmelCase__     )
              _UpperCAmelCase     =     self.scheduler.timesteps

              _UpperCAmelCase     =     self.prior.config.num_embeddings
              _UpperCAmelCase     =     self.prior.config.embedding_dim

              _UpperCAmelCase     =     self.prepare_latents(
                  (batch_size, num_embeddings * embedding_dim)    ,  image_embeds.dtype    ,  UpperCAmelCase__    ,  UpperCAmelCase__    ,  UpperCAmelCase__    ,  self.scheduler    ,  )

              # YiYi notes: for testing only to match ldm, we can directly create a latents with desired shape: batch_size, num_embeddings, embedding_dim
              _UpperCAmelCase     =     latents.reshape(latents.shape[0]    ,  UpperCAmelCase__    ,  UpperCAmelCase__     )

              for i, t in enumerate(self.progress_bar(UpperCAmelCase__     )     ):
                     # expand the latents if we are doing classifier free guidance
                     _UpperCAmelCase     =     torch.cat([latents] * 2     ) if do_classifier_free_guidance else latents
                     _UpperCAmelCase     =     self.scheduler.scale_model_input(UpperCAmelCase__    ,  UpperCAmelCase__     )

                     _UpperCAmelCase     =     self.prior(
                         UpperCAmelCase__    ,  timestep=UpperCAmelCase__    ,  proj_embedding=UpperCAmelCase__    ,  ).predicted_image_embedding

                     # remove the variance
                     _UpperCAmelCase    ,							_UpperCAmelCase     =     noise_pred.split(
                         scaled_model_input.shape[2]    ,  dim=2     )  # batch_size, num_embeddings, embedding_dim

                     if do_classifier_free_guidance is not None:
                            _UpperCAmelCase    ,							_UpperCAmelCase     =     noise_pred.chunk(2     )
                            _UpperCAmelCase     =     noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)

                     _UpperCAmelCase     =     self.scheduler.step(
                         UpperCAmelCase__    ,  timestep=UpperCAmelCase__    ,  sample=UpperCAmelCase__    ,  ).prev_sample

              if output_type == "latent":
                     return ShapEPipelineOutput(images=UpperCAmelCase__     )

              _UpperCAmelCase     =     []
              for i, latent in enumerate(UpperCAmelCase__     ):
                     print()
                     _UpperCAmelCase     =     self.renderer.decode(
                         latent[None, :]    ,  UpperCAmelCase__    ,  size=UpperCAmelCase__    ,  ray_batch_size=40_96    ,  n_coarse_samples=64    ,  n_fine_samples=1_28    ,  )

                     images.append(UpperCAmelCase__     )

              _UpperCAmelCase     =     torch.stack(UpperCAmelCase__     )

              if output_type not in ["np", "pil"]:
                     raise ValueError(f"""Only the output types `pil` and `np` are supported not output_type={output_type}"""     )

              _UpperCAmelCase     =     images.cpu().numpy()

              if output_type == "pil":
                     _UpperCAmelCase     =     [self.numpy_to_pil(UpperCAmelCase__     ) for image in images]

              # Offload last model to CPU
              if hasattr(self    ,  'final_offload_hook'     ) and self.final_offload_hook is not None:
                     self.final_offload_hook.offload()

              if not return_dict:
                     return (images,)

              return ShapEPipelineOutput(images=UpperCAmelCase__     )













"""simple docstring"""






import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin


if is_torch_available():
    import torch

if is_vision_available():
    from PIL import Image

    from transformers import ChineseCLIPImageProcessor



class 	lowerCAmelCase					(      unittest.TestCase							):






       def __init__(   self    ,  a__    ,  a__=7    ,  a__=3    ,  a__=18    ,  a__=30    ,  a__=4_00    ,  a__=True    ,  a__=None    ,  a__=True    ,  a__=None    ,  a__=True    ,  a__=[0.48_145_466, 0.4_578_275, 0.40_821_073]    ,  a__=[0.26_862_954, 0.26_130_258, 0.27_577_711]    ,  a__=True    ,  ):
              _UpperCAmelCase     =     size if size is not None else {'height': 2_24, 'width': 2_24}
              _UpperCAmelCase     =     crop_size if crop_size is not None else {'height': 18, 'width': 18}
              _UpperCAmelCase     =     parent
              _UpperCAmelCase     =     batch_size
              _UpperCAmelCase     =     num_channels
              _UpperCAmelCase     =     image_size
              _UpperCAmelCase     =     min_resolution
              _UpperCAmelCase     =     max_resolution
              _UpperCAmelCase     =     do_resize
              _UpperCAmelCase     =     size
              _UpperCAmelCase     =     do_center_crop
              _UpperCAmelCase     =     crop_size
              _UpperCAmelCase     =     do_normalize
              _UpperCAmelCase     =     image_mean
              _UpperCAmelCase     =     image_std
              _UpperCAmelCase     =     do_convert_rgb






       def 							__A (   self     ):
              return {
                  "do_resize": self.do_resize,
                  "size": self.size,
                  "do_center_crop": self.do_center_crop,
                  "crop_size": self.crop_size,
                  "do_normalize": self.do_normalize,
                  "image_mean": self.image_mean,
                  "image_std": self.image_std,
                  "do_convert_rgb": self.do_convert_rgb,
              }






       def 							__A (   self    ,  a__=False    ,  a__=False    ,  a__=False     ):
              assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

              if equal_resolution:
                     _UpperCAmelCase     =     []
                     for i in range(self.batch_size     ):
                            image_inputs.append(
                                np.random.randint(
                                    2_55    ,  size=(self.num_channels, self.max_resolution, self.max_resolution)    ,  dtype=np.uinta     )     )
              else:
                     _UpperCAmelCase     =     []
                     for i in range(self.batch_size     ):
                            _UpperCAmelCase    ,							_UpperCAmelCase     =     np.random.choice(np.arange(self.min_resolution    ,  self.max_resolution     )    ,  2     )
                            image_inputs.append(np.random.randint(2_55    ,  size=(self.num_channels, width, height)    ,  dtype=np.uinta     )     )

              if not numpify and not torchify:
                     # PIL expects the channel dimension as last dimension
                     _UpperCAmelCase     =     [Image.fromarray(np.moveaxis(a__    ,  0    ,  -1     )     ) for x in image_inputs]

              if torchify:
                     _UpperCAmelCase     =     [torch.from_numpy(a__     ) for x in image_inputs]

              return image_inputs



@require_torch
@require_vision
class 	lowerCAmelCase					(      snake_case			,	unittest.TestCase							):
       lowerCAmelCase__						 =  ChineseCLIPImageProcessor if is_vision_available() else None






       def 							__A (   self     ):
              _UpperCAmelCase     =     ChineseCLIPImageProcessingTester(self    ,  do_center_crop=a__     )






       @property
       def 							__A (   self     ):
              return self.image_processor_tester.prepare_image_processor_dict()






       def 							__A (   self     ):
              _UpperCAmelCase     =     self.image_processing_class(**self.image_processor_dict     )
              self.assertTrue(hasattr(a__    ,  'do_resize'     )     )
              self.assertTrue(hasattr(a__    ,  'size'     )     )
              self.assertTrue(hasattr(a__    ,  'do_center_crop'     )     )
              self.assertTrue(hasattr(a__    ,  'center_crop'     )     )
              self.assertTrue(hasattr(a__    ,  'do_normalize'     )     )
              self.assertTrue(hasattr(a__    ,  'image_mean'     )     )
              self.assertTrue(hasattr(a__    ,  'image_std'     )     )
              self.assertTrue(hasattr(a__    ,  'do_convert_rgb'     )     )






       def 							__A (   self     ):
              _UpperCAmelCase     =     self.image_processing_class.from_dict(self.image_processor_dict     )
              self.assertEqual(image_processor.size    ,  {'height': 2_24, 'width': 2_24}     )
              self.assertEqual(image_processor.crop_size    ,  {'height': 18, 'width': 18}     )

              _UpperCAmelCase     =     self.image_processing_class.from_dict(self.image_processor_dict    ,  size=42    ,  crop_size=84     )
              self.assertEqual(image_processor.size    ,  {'shortest_edge': 42}     )
              self.assertEqual(image_processor.crop_size    ,  {'height': 84, 'width': 84}     )






       def 							__A (   self     ):
              pass






       def 							__A (   self     ):
              # Initialize image_processing
              _UpperCAmelCase     =     self.image_processing_class(**self.image_processor_dict     )
              # create random PIL images
              _UpperCAmelCase     =     self.image_processor_tester.prepare_inputs(equal_resolution=a__     )
              for image in image_inputs:
                     self.assertIsInstance(a__    ,  Image.Image     )

              # Test not batched input
              _UpperCAmelCase     =     image_processing(image_inputs[0]    ,  return_tensors='pt'     ).pixel_values
              self.assertEqual(
                  encoded_images.shape    ,  (
                      1,
                      self.image_processor_tester.num_channels,
                      self.image_processor_tester.crop_size['height'],
                      self.image_processor_tester.crop_size['width'],
                  )    ,  )

              # Test batched
              _UpperCAmelCase     =     image_processing(a__    ,  return_tensors='pt'     ).pixel_values
              self.assertEqual(
                  encoded_images.shape    ,  (
                      self.image_processor_tester.batch_size,
                      self.image_processor_tester.num_channels,
                      self.image_processor_tester.crop_size['height'],
                      self.image_processor_tester.crop_size['width'],
                  )    ,  )






       def 							__A (   self     ):
              # Initialize image_processing
              _UpperCAmelCase     =     self.image_processing_class(**self.image_processor_dict     )
              # create random numpy tensors
              _UpperCAmelCase     =     self.image_processor_tester.prepare_inputs(equal_resolution=a__    ,  numpify=a__     )
              for image in image_inputs:
                     self.assertIsInstance(a__    ,  np.ndarray     )

              # Test not batched input
              _UpperCAmelCase     =     image_processing(image_inputs[0]    ,  return_tensors='pt'     ).pixel_values
              self.assertEqual(
                  encoded_images.shape    ,  (
                      1,
                      self.image_processor_tester.num_channels,
                      self.image_processor_tester.crop_size['height'],
                      self.image_processor_tester.crop_size['width'],
                  )    ,  )

              # Test batched
              _UpperCAmelCase     =     image_processing(a__    ,  return_tensors='pt'     ).pixel_values
              self.assertEqual(
                  encoded_images.shape    ,  (
                      self.image_processor_tester.batch_size,
                      self.image_processor_tester.num_channels,
                      self.image_processor_tester.crop_size['height'],
                      self.image_processor_tester.crop_size['width'],
                  )    ,  )






       def 							__A (   self     ):
              # Initialize image_processing
              _UpperCAmelCase     =     self.image_processing_class(**self.image_processor_dict     )
              # create random PyTorch tensors
              _UpperCAmelCase     =     self.image_processor_tester.prepare_inputs(equal_resolution=a__    ,  torchify=a__     )
              for image in image_inputs:
                     self.assertIsInstance(a__    ,  torch.Tensor     )

              # Test not batched input
              _UpperCAmelCase     =     image_processing(image_inputs[0]    ,  return_tensors='pt'     ).pixel_values
              self.assertEqual(
                  encoded_images.shape    ,  (
                      1,
                      self.image_processor_tester.num_channels,
                      self.image_processor_tester.crop_size['height'],
                      self.image_processor_tester.crop_size['width'],
                  )    ,  )

              # Test batched
              _UpperCAmelCase     =     image_processing(a__    ,  return_tensors='pt'     ).pixel_values
              self.assertEqual(
                  encoded_images.shape    ,  (
                      self.image_processor_tester.batch_size,
                      self.image_processor_tester.num_channels,
                      self.image_processor_tester.crop_size['height'],
                      self.image_processor_tester.crop_size['width'],
                  )    ,  )



@require_torch
@require_vision
class 	lowerCAmelCase					(      snake_case			,	unittest.TestCase							):
       lowerCAmelCase__						 =  ChineseCLIPImageProcessor if is_vision_available() else None






       def 							__A (   self     ):
              _UpperCAmelCase     =     ChineseCLIPImageProcessingTester(self    ,  num_channels=4    ,  do_center_crop=a__     )
              _UpperCAmelCase     =     3






       @property
       def 							__A (   self     ):
              return self.image_processor_tester.prepare_image_processor_dict()






       def 							__A (   self     ):
              _UpperCAmelCase     =     self.image_processing_class(**self.image_processor_dict     )
              self.assertTrue(hasattr(a__    ,  'do_resize'     )     )
              self.assertTrue(hasattr(a__    ,  'size'     )     )
              self.assertTrue(hasattr(a__    ,  'do_center_crop'     )     )
              self.assertTrue(hasattr(a__    ,  'center_crop'     )     )
              self.assertTrue(hasattr(a__    ,  'do_normalize'     )     )
              self.assertTrue(hasattr(a__    ,  'image_mean'     )     )
              self.assertTrue(hasattr(a__    ,  'image_std'     )     )
              self.assertTrue(hasattr(a__    ,  'do_convert_rgb'     )     )






       def 							__A (   self     ):
              pass






       def 							__A (   self     ):
              # Initialize image_processing
              _UpperCAmelCase     =     self.image_processing_class(**self.image_processor_dict     )
              # create random PIL images
              _UpperCAmelCase     =     self.image_processor_tester.prepare_inputs(equal_resolution=a__     )
              for image in image_inputs:
                     self.assertIsInstance(a__    ,  Image.Image     )

              # Test not batched input
              _UpperCAmelCase     =     image_processing(image_inputs[0]    ,  return_tensors='pt'     ).pixel_values
              self.assertEqual(
                  encoded_images.shape    ,  (
                      1,
                      self.expected_encoded_image_num_channels,
                      self.image_processor_tester.crop_size['height'],
                      self.image_processor_tester.crop_size['width'],
                  )    ,  )

              # Test batched
              _UpperCAmelCase     =     image_processing(a__    ,  return_tensors='pt'     ).pixel_values
              self.assertEqual(
                  encoded_images.shape    ,  (
                      self.image_processor_tester.batch_size,
                      self.expected_encoded_image_num_channels,
                      self.image_processor_tester.crop_size['height'],
                      self.image_processor_tester.crop_size['width'],
                  )    ,  )






