import os
from rtmlib import RTMO

class RTMO_GPU(RTMO):

    def __init__(self,
                 onnx_model: str = None,
                 model_input_size: tuple = (640, 640),
                 mean: tuple = None,
                 std: tuple = None,
                 to_openpose: bool = False,
                 backend: str = 'onnxruntime',
                 device: str = 'cuda'):
        
        if backend == 'onnxruntime':

            if not os.path.exists(onnx_model):
                from rtmlib.tools.file import download_checkpoint
                onnx_model = download_checkpoint(onnx_model)

            import onnxruntime as ort

            providers = {'cpu': 'CPUExecutionProvider',
                         'cuda': [
                                      ('CUDAExecutionProvider', {
                                        'device_id': 0,
                                        'arena_extend_strategy': 'kNextPowerOfTwo',
                                        'gpu_mem_limit': 2 * 1024 * 1024 * 1024,
                                        'cudnn_conv_algo_search': 'DEFAULT',
                                        'do_copy_in_default_stream': True,
                                        'enable_cuda_graph': False
                                    }),
                                  'CPUExecutionProvider']}
            
            self.session = ort.InferenceSession(path_or_bytes=onnx_model,
                                                providers=providers[device])

            print(f'load {onnx_model} with {backend} backend')

            self.onnx_model = onnx_model
            self.model_input_size = model_input_size
            self.mean = mean
            self.std = std
            self.backend = backend
            self.device = device
            self.to_openpose = to_openpose

        else:
            super().__init__(onnx_model,
                             model_input_size, 
                             mean,
                             std, 
                             to_openpose,
                             backend, 
                             device)