pms-nvidia-processor 2.0.0__tar.gz → 2.1.0__tar.gz

{pms_nvidia_processor-2.0.0 → pms_nvidia_processor-2.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pms-nvidia-processor
-Version: 2.0.0
+Version: 2.1.0
 Summary: lib for pms nvidia processors
 Author: HyeongSeok Kim
 Author-email: tiryul@gmail.com

pms_nvidia_processor-2.1.0/pms_nvidia_processor/__init__.py

@@ -0,0 +1,13 @@
+from .processor.dpir.processor_v2 import DPIRProcessor
+from .processor.dru_rbpn_sr.processor_f3_v2 import DRURBPNSRF3Processor
+from .processor.dru_rbpn_sr.processor_f5_v2 import DRURBPNSRF5Processor
+from .processor.dru_asm_sr.processor_f3_v2 import DRUASMSRF3Processor
+from .processor.dru_rbpn_deinter.processor_f3_gloss_v2 import (
+    DRURBPNDEINTERF3GLOSSProcessor,
+)
+from .processor.fisf.processor_v2 import FISFProcessor
+from .processor.color_resnet.processor_pre_v2 import COLORRESNETPREProcessor
+from .processor.color_resnet.processor_post_v2 import COLORRESNETPOSTProcessor
+from .processor.gg.processor_v2 import GGProcessor
+
+__version__ = "2.1.0"

pms_nvidia_processor-2.1.0/pms_nvidia_processor/processor/base/config.py

@@ -0,0 +1,69 @@
+from .dependency import *
+
+
+@dataclass
+class PatcherIOConfig:
+    patch_size: int
+    upscale_ratio: int
+    number_of_input_channels: int
+    number_of_output_channels: int
+    input_overlab_length: int
+
+    @property
+    def output_overlab_length(self) -> int:
+        return self.input_overlab_length * self.upscale_ratio
+
+    @property
+    def input_shape(self) -> tuple[int, int, int]:
+        return (
+            self.patch_size,
+            self.patch_size,
+            self.number_of_input_channels,
+        )
+
+    @property
+    def output_shape(self) -> tuple[int, int, int]:
+        output_patch_size = (
+            self.patch_size - self.input_overlab_length * 2
+        ) * self.upscale_ratio
+        return (
+            output_patch_size,
+            output_patch_size,
+            self.number_of_output_channels,
+        )
+
+    def build_patcher_params(
+        self,
+        input_vector: np.ndarray,
+        output_vector: np.ndarray,
+    ) -> dict[str, object]:
+        input_vector_shape = input_vector.shape
+        output_vector_shape = output_vector.shape
+        return {
+            "input_vector_shape": input_vector_shape,  # type: ignore
+            "input_patch_shape": self.input_shape,
+            "input_overlap_length": self.input_overlab_length,
+            "output_vector_shape": output_vector_shape,  # type: ignore
+            "output_patch_shape": self.output_shape,
+            "output_overlap_length": self.output_overlab_length,
+        }
+
+
+@dataclass
+class TRTIOConfig:
+    patch_size: int
+    upscale_ratio: int
+    number_of_input_channels: int
+    number_of_output_channels: int
+
+    @property
+    def input_shape(self) -> tuple[int, int, int]:
+        return (self.number_of_input_channels, self.patch_size, self.patch_size)
+
+    @property
+    def output_shape(self) -> tuple[int, int, int]:
+        return (
+            self.number_of_output_channels,
+            self.patch_size * self.upscale_ratio,
+            self.patch_size * self.upscale_ratio,
+        )

pms_nvidia_processor-2.0.0/pms_nvidia_processor/_dependency.py → pms_nvidia_processor-2.1.0/pms_nvidia_processor/processor/base/dependency.py

@@ -1,23 +1,18 @@
 from typing import (
-    List,
-    Optional,
-    Tuple,
-    Iterable,
     Literal,
-    Dict,
-    Any,
     TypeVar,
-    Union,
     Generic,
 )
-from abc import ABCMeta, abstractmethod
+from abc import ABCMeta, abstractmethod, ABC
 from dataclasses import dataclass
 import os
 import time
 import uvloop
 import asyncio
 import numpy as np
-from loguru import logger
+import loguru
+import enum
+from itertools import chain
 import cv2
 import pms_tensorrt as TRT
 from pms_inference_engine import IEngineProcessor, EngineIOData, register

pms_nvidia_processor-2.1.0/pms_nvidia_processor/processor/base/logger.py

@@ -0,0 +1,42 @@
+from .dependency import *
+
+
+class LogLevel(enum.Enum):
+    TRACE = enum.auto()
+    DEBUG = enum.auto()
+    INFO = enum.auto()
+    WARNING = enum.auto()
+    ERROR = enum.auto()
+    CRITICAL = enum.auto()
+
+
+class ILogger(ABC):
+
+    @abstractmethod
+    def log(self, level: LogLevel, message: str, **kwargs): ...
+
+    def trace(self, message: str, **kwargs):
+        self.log(LogLevel.TRACE, message, **kwargs)
+
+    def debug(self, message: str, **kwargs):
+        self.log(LogLevel.DEBUG, message, **kwargs)
+
+    def info(self, message: str, **kwargs):
+        self.log(LogLevel.INFO, message, **kwargs)
+
+    def warning(self, message: str, **kwargs):
+        self.log(LogLevel.WARNING, message, **kwargs)
+
+    def error(self, message: str, **kwargs):
+        self.log(LogLevel.ERROR, message, **kwargs)
+
+    def critical(self, message: str, **kwargs):
+        self.log(LogLevel.CRITICAL, message, **kwargs)
+
+
+class LoguruLogger(ILogger):
+    def __init__(self):
+        self._logger = loguru.logger
+
+    def log(self, level: LogLevel, message: str):
+        return self._logger.log(level.name, message)

pms_nvidia_processor-2.0.0/pms_nvidia_processor/_nvidia_processor_base.py → pms_nvidia_processor-2.1.0/pms_nvidia_processor/processor/base/processor.py

@@ -1,6 +1,5 @@
-from ._dependency import *
-from ._const import *
-from . import patcher
+from .dependency import *
+from .logger import LoguruLogger, LogLevel, ILogger
 
 
 class NVIDIAProcessorBase(IEngineProcessor[InputTypeT, OutputTypeT], metaclass=ABCMeta):
@@ -10,20 +9,22 @@ class NVIDIAProcessorBase(IEngineProcessor[InputTypeT, OutputTypeT], metaclass=A
         index: int,
         concurrency: int,
         model_path: str,
-        device: Union[Literal["auto"], int] = "auto",
+        device: Literal["auto"] | int = "auto",
+        logger: ILogger | None = None,
     ):
         # set member var
         self.index = index
         self.model_path = model_path
         self._concurrency = concurrency
         self._session = None
-        self.logger = logger
+
+        self.logger = LoguruLogger() if logger == None else logger
         # set loop policy
         asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
         # set device_id
         device_count = TRT.get_device_count()
-        if device is "auto":
+        if device == "auto":
             if "CUDA_VISIBLE_DEVICES" in os.environ:
                 devices = [
                     int(s) for s in os.environ["CUDA_VISIBLE_DEVICES"].split(",")
@@ -33,7 +34,7 @@ class NVIDIAProcessorBase(IEngineProcessor[InputTypeT, OutputTypeT], metaclass=A
             else:
                 device_id = index % device_count
         else:
-            assert type(device) is int, f"ERROR, The device's type is NOT int."
+            assert type(device) == int, f"ERROR, The device's type is NOT int."
             assert device > -1, f"ERROR, The device must be 'x > -1' but {device}."
             device_id = device
         self.device_id = device_id
@@ -47,7 +48,7 @@ class NVIDIAProcessorBase(IEngineProcessor[InputTypeT, OutputTypeT], metaclass=A
     def initialize_trt_session(
         self,
         required_batch_size: int,
-        io_shape: Dict[str, Tuple[List[int], np.dtype]],
+        io_shape: dict[str, tuple[list[int], np.dtype]],
     ):
         model_path = self.model_path
         device_id = self.device_id

pms_nvidia_processor-2.1.0/pms_nvidia_processor/processor/color_resnet/config.py

@@ -0,0 +1,42 @@
+from ..base.config import *
+
+
+class ColorResnetPreConfig:
+    NUMBER_OF_INPUT_CHANNELS: int = 3
+    NUMBER_OF_OUTPUT_CHANNELS: int = 30
+    UPSCALE_RATIO: int = 1
+    PATCH_SIZE = 256
+    MAX_BATCH_SIZE = 1
+    MIN_BATCH_SIZE = 1
+    OPT_BATCH_SIZE = 1
+
+    TRT_SHAPE_INPUT = (NUMBER_OF_INPUT_CHANNELS, PATCH_SIZE, PATCH_SIZE)
+    TRT_SHAPE_OUTPUT = (NUMBER_OF_OUTPUT_CHANNELS,)
+
+
+class ColorResnetPostConfig:
+    NUMBER_OF_FRAMES = 1
+    NUMBER_OF_INPUT_CHANNELS: int = 3 * NUMBER_OF_FRAMES
+    NUMBER_OF_MODEL_OUTPUT_CHANNELS: int = 30
+    NUMBER_OF_OUTPUT_CHANNELS: int = 3
+    UPSCALE_RATIO: int = 1
+    PATCH_SIZE = 512
+    MAX_BATCH_SIZE = 1
+    MIN_BATCH_SIZE = 1
+    OPT_BATCH_SIZE = 1
+    INPUT_OVERLAB_LENGTH = 16 * 2
+
+    PATCHER_CONFIG = PatcherIOConfig(
+        patch_size=PATCH_SIZE,
+        upscale_ratio=UPSCALE_RATIO,
+        number_of_input_channels=NUMBER_OF_INPUT_CHANNELS,
+        number_of_output_channels=NUMBER_OF_OUTPUT_CHANNELS,
+        input_overlab_length=INPUT_OVERLAB_LENGTH,
+    )
+    TRT_CONFIG = TRTIOConfig(
+        patch_size=PATCH_SIZE,
+        upscale_ratio=UPSCALE_RATIO,
+        number_of_input_channels=NUMBER_OF_INPUT_CHANNELS,
+        number_of_output_channels=NUMBER_OF_OUTPUT_CHANNELS,
+    )
+    TRT_SHAPE_MODEL_OUTPUT = (NUMBER_OF_MODEL_OUTPUT_CHANNELS,)

pms_nvidia_processor-2.1.0/pms_nvidia_processor/processor/color_resnet/context.py

@@ -0,0 +1,41 @@
+from .config import *
+from ..base.dependency import *
+from ..base.processor import *
+from ...utility import patcher, normalizer
+
+
+class ColorResnetPreContext:
+
+    def __init__(
+        self,
+        input_vector: np.ndarray,
+        output_buffer: np.ndarray,
+    ) -> None:
+        config = ColorResnetPreConfig
+        self.input_vector = input_vector
+        self.output_vector: np.ndarray = np.zeros_like(output_buffer, np.float32)
+
+
+class ColorResnetPostContext:
+    def __init__(
+        self,
+        input_vector: np.ndarray,
+        input_patcher: patcher.Patcher,
+    ) -> None:
+        config = ColorResnetPostConfig
+
+        self.input_vector = input_vector
+        self.input_pathcer = input_patcher
+        self.padded_input_image = patcher.pad_vector(
+            vector=self.input_vector,
+            overlap_length=config.PATCHER_CONFIG.input_overlab_length,
+            mode="edge",
+        )
+        self.input_patches: list[np.ndarray] = self.input_pathcer.slice(
+            self.padded_input_image
+        )
+        self.output_patches: np.ndarray = np.zeros(
+            (len(self.input_patches), config.PATCH_SIZE, config.PATCH_SIZE, 3),
+            np.float32,
+        )
+        self.output_vector: np.ndarray = np.zeros_like(self.input_vector, np.float32)

pms_nvidia_processor-2.0.0/pms_nvidia_processor/_color_resnet_post_processor.py → pms_nvidia_processor-2.1.0/pms_nvidia_processor/processor/color_resnet/processor_post.py

@@ -1,12 +1,11 @@
-from datetime import datetime
-from ._dependency import *
-from ._const import *
-from . import patcher
-from ._nvidia_processor_base import NVIDIAProcessorBase
+from .config import *
+from ..base.dependency import *
+from ..base.processor import *
+from ...utility import patcher
 
 
 def _pre_processing(
-    batch_input_images: List[np.ndarray],
+    batch_input_images: list[np.ndarray],
     input_buffer: np.ndarray,
 ) -> None:
     b = len(batch_input_images)
@@ -40,7 +39,7 @@ def _post_processing_merged(
     return np.clip(np.multiply(output_image, 255), 0, 255).astype(np.uint8)
 
 
-def _get_scale_factor(output_image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+def _get_scale_factor(output_image: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     def _rgb2grayscale(rgb: np.ndarray):
         return (
             0.299 * rgb[..., :, :, 0:1]
@@ -78,7 +77,7 @@ class COLORRESNETPOSTProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
         index: int,
         model_path: str,
         filter_coefficient: np.ndarray,
-        device: Union[Literal["auto"], int] = "auto",
+        device: Literal["auto"] | int = "auto",
     ):
         # super
         super().__init__(
@@ -92,7 +91,7 @@ class COLORRESNETPOSTProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
         self.is_scale_factor_calculated = False
         assert True
 
-    async def inference(self, batch_input_data: List[np.ndarray]) -> List[np.ndarray]:
+    async def inference(self, batch_input_data: list[np.ndarray]) -> list[np.ndarray]:
         session = self.session
         patch_size = self.config.PATCH_SIZE
         # batch_input_data = [np.zeros_like(batch_input_data[0]), *batch_input_data]

pms_nvidia_processor-2.1.0/pms_nvidia_processor/processor/color_resnet/processor_post_v2.py

@@ -0,0 +1,209 @@
+from pms_nvidia_processor.processor.color_resnet.context import ColorResnetPostContext
+from .config import *
+from ..base.dependency import *
+from ..base.processor import *
+from ...utility import patcher, normalizer, batch, caster
+
+
+def _pre_processing(
+    batch_input_images: list[np.ndarray],
+    input_buffer: np.ndarray,
+) -> None:
+    b = len(batch_input_images)
+    for batch_idx in range(b):
+        image = batch_input_images[batch_idx]
+        h, w, c = image.shape
+        for channel_idx in range(c):
+            input_buffer[batch_idx, channel_idx, :h, :w] = image[:, :, channel_idx]
+
+
+def _post_processing_patch(
+    output_buffer: np.ndarray,  # BxCxHxW
+    output_image: np.ndarray,  # BxHxWxC
+) -> None:
+    b, h, w, c = output_image.shape
+    pred = output_buffer[:b, :, :h, :w]
+    for i in range(3):
+        np.copyto(src=pred[:, i, :, :], dst=output_image[:, :, :, i])
+
+
+def _post_processing_merged(
+    output_image: np.ndarray,
+    scale_min: np.ndarray,
+    scale_max: np.ndarray,
+):
+    np.subtract(output_image, scale_min, out=output_image)
+    np.divide(output_image, scale_max - scale_min, out=output_image)
+
+
+def _get_scale_factor(output_image: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    def _rgb2grayscale(rgb: np.ndarray):
+        return (
+            0.299 * rgb[..., :, :, 0:1]
+            + 0.587 * rgb[..., :, :, 1:2]
+            + 0.114 * rgb[..., :, :, 2:3]
+        )
+
+    scale_max = (
+        np.amax(
+            _rgb2grayscale(output_image),
+            # axis=(-3, -2, -1),
+            axis=(-2, -1, -3),
+            keepdims=True,
+        )
+        + 0.0015
+    )
+    scale_min = (
+        np.amin(
+            _rgb2grayscale(output_image),
+            # axis=(-3, -2, -1),
+            axis=(-2, -1, -3),
+            keepdims=True,
+        )
+        - 0.0015
+    )
+    return scale_min, scale_max
+
+
+@register
+class COLORRESNETPOSTProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
+
+    def __init__(
+        self,
+        concurrency: int,
+        index: int,
+        model_path: str,
+        filter_coefficient: np.ndarray,
+        device: Literal["auto"] | int = "auto",
+        normalization: str = "color-basic",
+    ):
+        # super
+        super().__init__(
+            concurrency=concurrency,
+            index=index,
+            model_path=model_path,
+            device=device,
+        )
+        self.filter_coefficient = filter_coefficient
+        self.config = ColorResnetPostConfig
+        self.is_scale_factor_calculated = False
+        self.normalizer = normalizer.NormalizerFactory.create_multi(normalization)
+
+    async def inference(
+        self,
+        batch_input_data: list[np.ndarray],
+        batch_output_data: np.ndarray,
+    ):
+        session = self.session
+
+        _pre_processing(
+            batch_input_images=batch_input_data,
+            input_buffer=self.input_buffer,
+        )
+        session.run()
+        _post_processing_patch(
+            output_buffer=self.output_buffer,
+            output_image=batch_output_data,
+        )
+
+    async def _run(self, input_data: EngineIOData) -> EngineIOData:
+        max_batch_size = 1
+        # 여기서 patching
+        input_vector = self.caster.cast(input_data.frame)
+        self.normalizer.normalize(input_vector, 2)
+        context = ColorResnetPostContext(input_vector, self.patcher)
+
+        for batch_input_patches, batch_output_patches in zip(
+            batch.BatchTool.batch_list(context.input_patches, max_batch_size),
+            batch.BatchTool.batch_vector(context.output_patches, max_batch_size),
+        ):
+            await self.inference(
+                batch_input_data=batch_input_patches,
+                batch_output_data=batch_output_patches,
+            )
+
+        self.patcher.merge(
+            output_vector=context.output_vector,
+            patches=context.output_patches,
+        )
+
+        if not self.is_scale_factor_calculated:
+            self.scale_min, self.scale_max = _get_scale_factor(context.output_vector)
+            self.is_scale_factor_calculated = True
+        _post_processing_merged(
+            output_image=context.output_vector,
+            scale_min=self.scale_min,
+            scale_max=self.scale_max,
+        )
+        self.normalizer.denormalize(context.output_vector, 2)
+        return EngineIOData(
+            frame_id=input_data.frame_id,
+            frame=self.caster.uncast(context.output_vector),
+        )
+
+    def _ready_processor(self) -> bool:
+        return True
+
+    def _bind_io(self, input_data: EngineIOData):
+        self.caster = caster.Caster(input_data.frame)
+
+        patcher_config = self.config.PATCHER_CONFIG
+        trt_config = self.config.TRT_CONFIG
+
+        input_image: np.ndarray = input_data.frame  # type: ignore
+        padded_input_image = patcher.pad_vector(
+            input_image,
+            overlap_length=patcher_config.input_overlab_length,
+        )
+        output_image: np.ndarray = np.zeros(input_image.shape, np.float32)
+        self.patcher = patcher.Patcher(
+            **patcher_config.build_patcher_params(
+                input_vector=padded_input_image,
+                output_vector=output_image,
+            )
+        )
+        n_patches = len(self.patcher.slice(input_vector=padded_input_image))
+
+        # set io shape
+        # self.batch_size = min(n_patches + 1, self.config.MAX_BATCH_SIZE - 1)
+        self.batch_size = self.config.MAX_BATCH_SIZE
+        self.io_shapes = {
+            "input": (
+                [self.batch_size, *trt_config.input_shape],
+                np.float32,
+            ),
+            "model_output": (
+                [1, *self.config.TRT_SHAPE_MODEL_OUTPUT],
+                np.float32,
+            ),
+            "output": (
+                [self.batch_size, *trt_config.output_shape],
+                np.float32,
+            ),
+        }
+
+        # init trt engine
+        self.initialize_trt_session(
+            required_batch_size=self.batch_size,
+            io_shape=self.io_shapes,
+        )
+
+        # set io buffer
+        self.input_buffer = self.session._input_bindings[0].host_buffer.reshape(
+            self.io_shapes["input"][0]
+        )
+        self.filter_coefficient_buffer = self.session._input_bindings[
+            1
+        ].host_buffer.reshape(self.io_shapes["model_output"][0])
+        self.output_buffer = self.session._output_bindings[0].host_buffer.reshape(
+            self.io_shapes["output"][0]
+        )
+        self.filter_coefficient_buffer[:] = self.filter_coefficient[:]
+        self.session.run()
+        return True
+
+    def _get_live(self) -> bool:
+        return True
+
+    def _get_concurrency(self) -> int:
+        return self._concurrency

pms_nvidia_processor-2.0.0/pms_nvidia_processor/_color_resnet_pre_processor.py → pms_nvidia_processor-2.1.0/pms_nvidia_processor/processor/color_resnet/processor_pre.py

@@ -1,8 +1,7 @@
-from typing import Any
-from ._dependency import *
-from ._const import *
-from . import patcher
-from ._nvidia_processor_base import NVIDIAProcessorBase
+from .config import *
+from ..base.dependency import *
+from ..base.processor import *
+from ...utility import patcher
 
 
 def _pre_processing(
@@ -35,7 +34,7 @@ class COLORRESNETPREProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
         concurrency: int,
         index: int,
         model_path: str,
-        device: Union[Literal["auto"], int] = "auto",
+        device: Literal["auto"] | int = "auto",
     ):
         # super
         super().__init__(

pms_nvidia_processor-2.1.0/pms_nvidia_processor/processor/color_resnet/processor_pre_v2.py

@@ -0,0 +1,114 @@
+from .config import *
+from ..base.dependency import *
+from ..base.processor import *
+from .context import ColorResnetPreContext
+from ...utility import normalizer, caster
+
+
+def _pre_processing(
+    input_image: np.ndarray, input_buffer: np.ndarray, input_size: int
+) -> None:
+    resize = cv2.resize(
+        input_image, (input_size, input_size), interpolation=cv2.INTER_CUBIC
+    )
+    h, w, c = resize.shape
+    for channel_idx in range(c):
+        input_buffer[0, channel_idx, :h, :w] = resize[:, :, channel_idx]
+
+
+def _post_processing(
+    output_buffer: np.ndarray,  # BxCxHxW
+    output_data: np.ndarray,  # BxHxWxC
+) -> None:
+    output_data[:] = output_buffer[:]
+
+
+@register
+class COLORRESNETPREProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
+
+    def __init__(
+        self,
+        concurrency: int,
+        index: int,
+        model_path: str,
+        device: Literal["auto"] | int = "auto",
+        normalization: str = "color-basic",
+    ):
+        # super
+        super().__init__(
+            concurrency=concurrency,
+            index=index,
+            model_path=model_path,
+            device=device,
+        )
+
+        self.config = ColorResnetPreConfig
+        self.normalizer = normalizer.NormalizerFactory.create_multi(normalization)
+
+    async def inference(
+        self, input_data: np.ndarray, output_data: np.ndarray
+    ) -> np.ndarray:
+        session = self.session
+        _pre_processing(
+            input_image=input_data,
+            input_buffer=self.input_buffer,
+            input_size=self.config.PATCH_SIZE,
+        )
+        session.run()
+        _post_processing(
+            output_buffer=self.output_buffer,
+            output_data=output_data,
+        )
+
+    async def _run(self, input_data: EngineIOData) -> EngineIOData:
+        input_vector = self.caster.cast(input_data.frame)
+        self.normalizer.normalize(input_vector, 2)
+        context = ColorResnetPreContext(input_vector, self.output_buffer)
+
+        # batch inference
+        await self.inference(
+            input_data=context.input_vector, output_data=context.output_vector
+        )
+
+        return EngineIOData(frame_id=input_data.frame_id, frame=context.output_vector)
+
+    def _ready_processor(self) -> bool:
+        return True
+
+    def _bind_io(self, input_data: EngineIOData):
+        self.caster = caster.Caster(input_data.frame)
+
+        # set io shape
+        self.batch_size = self.config.MAX_BATCH_SIZE
+        self.io_shapes = {
+            "input": (
+                [self.batch_size, *self.config.TRT_SHAPE_INPUT],
+                np.float32,
+            ),
+            "output": (
+                [self.batch_size, *self.config.TRT_SHAPE_OUTPUT],
+                np.float32,
+            ),
+        }
+
+        # init trt engine
+        self.initialize_trt_session(
+            required_batch_size=self.batch_size,
+            io_shape=self.io_shapes,
+        )
+
+        # set io buffer
+        self.input_buffer = self.session._input_bindings[0].host_buffer.reshape(
+            *self.io_shapes["input"][0]
+        )
+        self.output_buffer = self.session._output_bindings[0].host_buffer.reshape(
+            *self.io_shapes["output"][0]
+        )
+
+        return True
+
+    def _get_live(self) -> bool:
+        return True
+
+    def _get_concurrency(self) -> int:
+        return self._concurrency