memory_decorator.py

import inspect
import os
from functools import wraps
from typing import Any
from typing import Callable
from typing import cast
from typing import TypeVar


def get_function_name(func: Callable) -> str:
    """Get file and function name."""
    module_name = func.__module__.split('.')[-1]
    if module_name == "__main__":
        module_name = os.path.splitext(os.path.basename(inspect.getfile(func)))[0]
    return f"{module_name}::{func.__name__}"


FuncT = TypeVar('FuncT', bound=Callable[..., Any])


def memory_peak_profile(func: FuncT) -> FuncT:
    """Memory peak Profiling decorator (Both RAM and VRAM)."""
    @wraps(func)
    def wrapper(*args, **kwargs):
        with MemoryScope(get_function_name(func)):
            retval = func(*args, **kwargs)

        return retval

    return cast(FuncT, wrapper)