Update mvsepless/gradio_helper.py

mvsepless/gradio_helper.py

@@ -3,6 +3,7 @@ import gradio as gr
 import zipfile
 from datetime import timezone, timedelta
 import platform
+from functools import wraps
 try:
     import spaces
 except:
@@ -250,14 +251,18 @@ else:
     if zerogpu_is:
         zerogpu_available = True
         def hf_spaces_gpu(*args, device="cpu", **kwargs):
-            # Поддержка разных вариантов вызова
             device = "cuda:0"
-            if len(args) == 1 and callable(args[0]):
-                # @hf_spaces_gpu
-                return spaces.GPU(args[0], **kwargs)
-            else:
-                # @hf_spaces_gpu(duration=30)
-                return spaces.GPU(*args, **kwargs)
+            def decorator(func):
+                # Для методов класса используем обертку
+                @wraps(func)
+                def wrapper(*args, **kwargs):
+                    # Создаем временную функцию для spaces.GPU
+                    @spaces.GPU(duration=kwargs.get('duration', 180))
+                    def gpu_wrapper():
+                        return func(*args, **kwargs)
+                    return gpu_wrapper()
+                return wrapper
+            return decorator(*args, **kwargs) if args and callable(args[0]) else decorator
     else:
         def hf_spaces_gpu(*args, **kwargs):
             if len(args) == 1 and callable(args[0]):