
import math
import os

import cva
import numpy as np

lowerCamelCase_				:							Tuple       =    255.0
def        __lowercase(							__snake_case				:		float    ,__snake_case				:		float					)       ->   Optional[Any]:
       __snake_case						       =   np.mean((original - contrast) ** 2					)
       if mse == 0:
              return 1_00

       return 20 * math.logaa(PIXEL_MAX / math.sqrt(__A					)					)
def        __lowercase(							)       ->   Optional[int]:
       __snake_case						       =   os.path.dirname(os.path.realpath(__A					)					)
       # Loading images (original image and compressed image)
       __snake_case						       =   cva.imread(os.path.join(__A    ,'image_data/original_image.png'					)					)
       __snake_case						       =   cva.imread(os.path.join(__A    ,'image_data/compressed_image.png'					)    ,1					)

       __snake_case						       =   cva.imread(os.path.join(__A    ,'image_data/PSNR-example-base.png'					)					)
       __snake_case						       =   cva.imread(
           os.path.join(__A    ,'image_data/PSNR-example-comp-10.jpg'					)    ,1					)

       # Value expected: 29.73dB
       print('-- First Test --'					)
       print(f'''PSNR value is {peak_signal_to_noise_ratio(__A    ,__A					)} dB'''					)

       # # Value expected: 31.53dB (Wikipedia Example)
       print('\n-- Second Test --'					)
       print(f'''PSNR value is {peak_signal_to_noise_ratio(__A    ,__A					)} dB'''					)


if __name__ == "__main__":
 main()






import os
from pathlib import Path
from typing import Dict, List

import fire
import torch

from transformers import AutoModelForSeqaSeqLM, AutoTokenizer
from transformers.utils.logging import get_logger


__lowerCamelCase :    Any      						=      get_logger(__name__)







def 		lowercase__ (		__A:						str				,__A:						str					):





					'''simple docstring'''






					if text.startswith(__A					):
										return text[len(__A					) :]
					return text  # or whatever







def 		lowercase__ (		__A:						List[str]					):





					'''simple docstring'''






					return {remove_prefix(__A				,'''model.'''					): v for k, v in sd.items()}







def 		lowercase__ (		__A:						List[Dict[str, torch.Tensor]]					):





					'''simple docstring'''






					__magic_name__				:						List[Any]     				=				{}
					for k in state_dicts[0].keys():
										__magic_name__				:						str     				=				[sd[k] for sd in state_dicts]
										__magic_name__				:						List[Any]     				=				sum(__A					) / len(__A					)
										assert isinstance(__A				,torch.Tensor					)
										__magic_name__				:						Any     				=				new_t
					return new_sd







def 		lowercase__ (		__A:						str				,__A:						str				,__A:						str					):





					'''simple docstring'''






					__magic_name__				:						List[Any]     				=				AutoModelForSeqaSeqLM.from_pretrained(__A					)
					if os.path.isfile(__A					):
										__magic_name__				:						str     				=				[pl_ckpt_path]
					else:
										assert os.path.isdir(__A					)
										__magic_name__				:						List[str]     				=				list(Path(__A					).glob('''*.ckpt'''					)					)
										assert ckpt_files, F'''could not find any ckpt files inside the {pl_ckpt_path} directory'''

					if len(__A					) > 1:
										logger.info(F'''averaging the weights of {ckpt_files}'''					)

					__magic_name__				:						Union[str, Any]     				=				[sanitize(torch.load(__A				,map_location='''cpu'''					)['''state_dict''']					) for x in ckpt_files]
					__magic_name__				:						Dict     				=				average_state_dicts(__A					)

					__magic_name__				,     __magic_name__				:						Union[str, Any]     				=				hf_model.load_state_dict(__A				,strict=__A					)
					assert not missing, F'''missing keys: {missing}'''
					hf_model.save_pretrained(__A					)
					try:
										__magic_name__				:						Optional[Any]     				=				AutoTokenizer.from_pretrained(__A					)
										tok.save_pretrained(__A					)
					except Exception:
										pass
										# dont copy tokenizer if cant


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



