


'''simple docstring'''




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import 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 (	__lowercase						,					unittest.TestCase							):
   _SCREAMING_SNAKE_CASE  :  Union[str, Any]          =	ShapEPipeline
   _SCREAMING_SNAKE_CASE  :  Union[str, Any]          =	["prompt"]
   _SCREAMING_SNAKE_CASE  :  Any          =	["prompt"]
   _SCREAMING_SNAKE_CASE  :  str          =	[
       "num_images_per_prompt",
       "num_inference_steps",
       "generator",
       "latents",
       "guidance_scale",
       "frame_size",
       "output_type",
       "return_dict",
   ]
   _SCREAMING_SNAKE_CASE  :  Optional[int]          =	False
   @property
   def a  (						self       :      Any      ) ->							Optional[int]:
        return 32
   @property
   def a  (						self       :      List[Any]      ) ->							List[Any]:
        return 32
   @property
   def a  (						self       :      Tuple      ) ->							List[str]:
        return self.time_input_dim * 4
   @property
   def a  (						self       :      Dict      ) ->							Union[str, Any]:
        return 8
   @property
   def a  (						self       :      List[Any]      ) ->							Optional[Any]:
        __snake_case           =						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      )
        return tokenizer
   @property
   def a  (						self       :      Dict      ) ->							Any:
        torch.manual_seed(0      )
        __snake_case           =						CLIPTextConfig(
            bos_token_id=0       ,   eos_token_id=2       ,   hidden_size=self.text_embedder_hidden_size       ,   projection_dim=self.text_embedder_hidden_size       ,   intermediate_size=37       ,   layer_norm_eps=1e-05       ,   num_attention_heads=4       ,   num_hidden_layers=5       ,   pad_token_id=1       ,   vocab_size=1000       ,   )
        return CLIPTextModelWithProjection(SCREAMING_SNAKE_CASE_      )
   @property
   def a  (						self       :      str      ) ->							Dict:
        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',
            'encoder_hid_proj_type': None,
            'added_emb_type': None,
        }

        __snake_case           =						PriorTransformer(**SCREAMING_SNAKE_CASE_      )
        return model
   @property
   def a  (						self       :      Optional[Any]      ) ->							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(**SCREAMING_SNAKE_CASE_      )
        return model
   def a  (						self       :      Tuple      ) ->							Dict:
        __snake_case           =						self.dummy_prior
        __snake_case           =						self.dummy_text_encoder
        __snake_case           =						self.dummy_tokenizer
        __snake_case           =						self.dummy_renderer

        __snake_case           =						HeunDiscreteScheduler(
            beta_schedule='exp'       ,   num_train_timesteps=1024       ,   prediction_type='sample'       ,   use_karras_sigmas=SCREAMING_SNAKE_CASE_       ,   clip_sample=SCREAMING_SNAKE_CASE_       ,   clip_sample_range=1.0       ,   )
        __snake_case           =						{
            'prior': prior,
            'text_encoder': text_encoder,
            'tokenizer': tokenizer,
            'renderer': renderer,
            'scheduler': scheduler,
        }

        return components
   def a  (						self       :      Any       ,   SCREAMING_SNAKE_CASE_       :      str       ,   SCREAMING_SNAKE_CASE_       :      Optional[int]=0      ) ->							Union[str, Any]:
        if str(SCREAMING_SNAKE_CASE_      ).startswith('mps'      ):
             __snake_case           =						torch.manual_seed(SCREAMING_SNAKE_CASE_      )
        else:
             __snake_case           =						torch.Generator(device=SCREAMING_SNAKE_CASE_      ).manual_seed(SCREAMING_SNAKE_CASE_      )
        __snake_case           =						{
            'prompt': 'horse',
            'generator': generator,
            'num_inference_steps': 1,
            'frame_size': 32,
            'output_type': 'np',
        }
        return inputs
   def a  (						self       :      Optional[Any]      ) ->							str:
        __snake_case           =						'cpu'

        __snake_case           =						self.get_dummy_components()

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

        pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_      )

        __snake_case           =						pipe(**self.get_dummy_inputs(SCREAMING_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.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
            ]      )

        assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2
   def a  (						self       :      int      ) ->							List[str]:
        # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
        self._test_inference_batch_consistent(batch_sizes=[1, 2]      )
   def a  (						self       :      Dict      ) ->							Any:
        __snake_case           =						torch_device == 'cpu'
        __snake_case           =						True

        self._test_inference_batch_single_identical(
            batch_size=2       ,   test_max_difference=SCREAMING_SNAKE_CASE_       ,   relax_max_difference=SCREAMING_SNAKE_CASE_       ,   )





   def a  (						self       :      Union[str, Any]      ) ->							str:
        __snake_case           =						self.get_dummy_components()
        __snake_case           =						self.pipeline_class(**SCREAMING_SNAKE_CASE_      )
        __snake_case           =						pipe.to(SCREAMING_SNAKE_CASE_      )
        pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_      )

        __snake_case           =						1
        __snake_case           =						2

        __snake_case           =						self.get_dummy_inputs(SCREAMING_SNAKE_CASE_      )

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

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

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




@slow
@require_torch_gpu
class 			_lowercase (	unittest.TestCase							):
   def a  (						self       :      Optional[int]      ) ->							Optional[Any]:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()





   def a  (						self       :      Union[str, Any]      ) ->							Optional[Any]:
        __snake_case           =						load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/shap_e/test_shap_e_np_out.npy'      )
        __snake_case           =						ShapEPipeline.from_pretrained('openai/shap-e'      )
        __snake_case           =						pipe.to(SCREAMING_SNAKE_CASE_      )
        pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_      )

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

        __snake_case           =						pipe(
            'a shark'       ,   generator=SCREAMING_SNAKE_CASE_       ,   guidance_scale=1_5.0       ,   num_inference_steps=64       ,   frame_size=64       ,   output_type='np'       ,   ).images[0]

        assert images.shape == (20, 64, 64, 3)

        assert_mean_pixel_difference(SCREAMING_SNAKE_CASE_       ,   SCREAMING_SNAKE_CASE_      )







import argparse
import json
import os

import torch
from transformers.file_utils import has_file

from diffusers import UNetaDConditionModel, UNetaDModel


lowercase_   							=  False
lowercase_   							=  True
lowercase_   							=  False


if __name__ == "__main__":
  lowercase_   							=  argparse.ArgumentParser()

  parser.add_argument(
      '--repo_path',
      default=None,
      type=str,
      required=True,
      help='The config json file corresponding to the architecture.',
  )

  parser.add_argument('--dump_path', default=None, type=str, required=True, help='Path to the output model.')

  lowercase_   							=  parser.parse_args()

  lowercase_   							=  {
      'image_size': 'sample_size',
      'num_res_blocks': 'layers_per_block',
      'block_channels': 'block_out_channels',
      'down_blocks': 'down_block_types',
      'up_blocks': 'up_block_types',
      'downscale_freq_shift': 'freq_shift',
      'resnet_num_groups': 'norm_num_groups',
      'resnet_act_fn': 'act_fn',
      'resnet_eps': 'norm_eps',
      'num_head_channels': 'attention_head_dim',
  }

  lowercase_   							=  {
      'time_steps': 'time_proj',
      'mid': 'mid_block',
      'downsample_blocks': 'down_blocks',
      'upsample_blocks': 'up_blocks',
  }

  lowercase_   							=  '' if has_file(args.repo_path, 'config.json') else 'unet'

  with open(os.path.join(args.repo_path, subfolder, 'config.json'), 'r', encoding='utf-8') as reader:
    lowercase_   							=  reader.read()
    lowercase_   							=  json.loads(text)

  if do_only_config:
    for key in config_parameters_to_change.keys():
      config.pop(key, None)

  if has_file(args.repo_path, 'config.json'):
    lowercase_   							=  UNetaDModel(**config)
  else:
    lowercase_   							=  UNetaDConditionModel if 'ldm-text2im-large-256' in args.repo_path else UNetaDModel
    lowercase_   							=  class_name(**config)

  if do_only_config:
    model.save_config(os.path.join(args.repo_path, subfolder))

  lowercase_   							=  dict(model.config)

  if do_only_renaming:
    for key, value in config_parameters_to_change.items():
      if key in config:
        lowercase_   							=  config[key]
        del config[key]

    lowercase_   							=  [k.replace('UNetRes', '') for k in config['down_block_types']]
    lowercase_   							=  [k.replace('UNetRes', '') for k in config['up_block_types']]

  if do_only_weights:
    lowercase_   							=  torch.load(os.path.join(args.repo_path, subfolder, 'diffusion_pytorch_model.bin'))

    lowercase_   							=  {}
    for param_key, param_value in state_dict.items():
      if param_key.endswith('.op.bias') or param_key.endswith('.op.weight'):
        continue
      lowercase_   							=  False
      for key, new_key in key_parameters_to_change.items():
        if not has_changed and param_key.split('.')[0] == key:
          lowercase_   							=  param_value
          lowercase_   							=  True
      if not has_changed:
        lowercase_   							=  param_value

    model.load_state_dict(new_state_dict)
    model.save_pretrained(os.path.join(args.repo_path, subfolder))


