Update decorators.py

decorators.py

@@ -1,48 +1,12 @@
 # decorators.py
 
-import functools
-import torch
-import threading
-import time
-
-class spaces:
-    @staticmethod
-    def GPU(duration=0):
-        def decorator(func):
-            @functools.wraps(func)
-            def wrapper(*args, **kwargs):
-                try:
-                    # Verifica si la GPU está disponible, si no lanza un error controlado
-                    if torch.cuda.is_available():
-                        device = torch.device('cuda')
-                        print("GPU disponible. Ejecutando en GPU.")
-                    else:
-                        raise RuntimeError("No se detecta la GPU. Asegúrate de que los drivers CUDA están instalados y configurados correctamente.")
-                
-                    # Pasar el dispositivo a la función como argumento
-                    kwargs['device'] = device
-
-                    result = [None]
-                    exception = [None]
-
-                    def target():
-                        try:
-                            result[0] = func(*args, **kwargs)
-                        except Exception as e:
-                            exception[0] = e
-
-                    # Ejecutar la función en un hilo separado
-                    thread = threading.Thread(target=target)
-                    thread.start()
-                    thread.join(duration)
-
-                    if thread.is_alive():
-                        raise TimeoutError(f"La ejecución de la función excedió {duration} segundos.")
-                    if exception[0]:
-                        raise exception[0]
-                    return result[0]
-                except RuntimeError as e:
-                    print(f"Error: {str(e)}")
-                    raise e  # Lanzar el error de GPU para que se maneje adecuadamente
-            return wrapper
-        return decorator
+import spaces
+
+# Decorador personalizado que usa GPU si está disponible
+def gpu_decorator(duration=100):
+    def decorator(func):
+        @spaces.GPU(duration=duration)
+        def wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+        return wrapper
+    return decorator