


'''simple docstring'''






import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionTextToImagePipeline
from diffusers.utils.testing_utils import nightly, require_torch_gpu, torch_device


__A   :    Dict											= False






class   __snake_case			(  unittest.TestCase):







    """simple docstring"""



    pass







@nightly
@require_torch_gpu
class   __snake_case			(  unittest.TestCase):







    """simple docstring"""



    def __lowercase  (  self       :    List[str]						)  ->    Any:
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()

    def __lowercase  (  self       :    List[str]						)  ->    Any:
      lowerCAmelCase_     : Dict     =  VersatileDiffusionTextToImagePipeline.from_pretrained("""shi-labs/versatile-diffusion"""						)
      # remove text_unet
      pipe.remove_unused_weights()
      pipe.to(lowerCamelCase						)
      pipe.set_progress_bar_config(disable=lowerCamelCase						)

      lowerCAmelCase_     : Any     =  """A painting of a squirrel eating a burger """
      lowerCAmelCase_     : Dict     =  torch.manual_seed(0						)
      lowerCAmelCase_     : Union[str, Any]     =  pipe(
          prompt=lowerCamelCase       ,		generator=lowerCamelCase       ,		guidance_scale=7.5       ,		num_inference_steps=2       ,		output_type="""numpy"""						).images

      with tempfile.TemporaryDirectory() as tmpdirname:
        pipe.save_pretrained(lowerCamelCase						)
        lowerCAmelCase_     : str     =  VersatileDiffusionTextToImagePipeline.from_pretrained(lowerCamelCase						)
      pipe.to(lowerCamelCase						)
      pipe.set_progress_bar_config(disable=lowerCamelCase						)

      lowerCAmelCase_     : Union[str, Any]     =  generator.manual_seed(0						)
      lowerCAmelCase_     : Any     =  pipe(
          prompt=lowerCamelCase       ,		generator=lowerCamelCase       ,		guidance_scale=7.5       ,		num_inference_steps=2       ,		output_type="""numpy"""						).images

      assert np.abs(image - new_image						).sum() < 1E-5, "Models don't have the same forward pass"





    def __lowercase  (  self       :    Union[str, Any]						)  ->    int:
      lowerCAmelCase_     : Union[str, Any]     =  VersatileDiffusionTextToImagePipeline.from_pretrained(
          """shi-labs/versatile-diffusion"""       ,		torch_dtype=torch.floataa						)
      pipe.to(lowerCamelCase						)
      pipe.set_progress_bar_config(disable=lowerCamelCase						)

      lowerCAmelCase_     : Optional[Any]     =  """A painting of a squirrel eating a burger """
      lowerCAmelCase_     : Union[str, Any]     =  torch.manual_seed(0						)
      lowerCAmelCase_     : str     =  pipe(
          prompt=lowerCamelCase       ,		generator=lowerCamelCase       ,		guidance_scale=7.5       ,		num_inference_steps=50       ,		output_type="""numpy"""						).images

      lowerCAmelCase_     : Tuple     =  image[0, 2_53:2_56, 2_53:2_56, -1]

      assert image.shape == (1, 5_12, 5_12, 3)
      lowerCAmelCase_     : int     =  np.array([0.3_367, 0.3_169, 0.2_656, 0.3_870, 0.4_790, 0.3_796, 0.4_009, 0.4_878, 0.4_778]						)

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







'''simple docstring'''






from typing import Union

import fire
import torch
from tqdm import tqdm







def 							UpperCamelCase_     (       A__   :  str				,       A__   :  str = "cpu"				,       A__   :  Union[str, None] = None   ):
  '''simple docstring'''





  lowerCAmelCase_     : Dict     =  torch.load(A__				,       map_location=A__   )
  for k, v in tqdm(state_dict.items()   ):
    if not isinstance(A__				,       torch.Tensor   ):
      raise TypeError("""FP16 conversion only works on paths that are saved state dicts, like pytorch_model.bin"""   )
    lowerCAmelCase_     : List[str]     =  v.half()
  if save_path is None:  # overwrite src_path
    lowerCAmelCase_     : List[str]     =  src_path
  torch.save(A__				,       A__   )


if __name__ == "__main__":
  fire.Fire(convert)



