Spaces:

Reubencf
/

Kiku-Engine

Running on Zero

Kiku-Engine / native_gpu.py

Deploy Kiku-Engine

052cbee verified 29 days ago

505 Bytes

"""Serialize the shared CUDA handlers on Modal without ZeroGPU leases."""
from functools import wraps
from threading import RLock

lock = RLock()
after_call = None


def GPU(*, duration=None):
    def decorate(function):
        @wraps(function)
        def invoke(*args, **kwargs):
            with lock:
                result = function(*args, **kwargs)
                if after_call is not None:
                    after_call()
                return result
        return invoke
    return decorate

	"""Serialize the shared CUDA handlers on Modal without ZeroGPU leases."""
	from functools import wraps
	from threading import RLock

	lock = RLock()
	after_call = None


	def GPU(*, duration=None):
	def decorate(function):
	@wraps(function)
	def invoke(args, *kwargs):
	with lock:
	result = function(args, *kwargs)
	if after_call is not None:
	after_call()
	return result
	return invoke
	return decorate