Update app.py

app.py

@@ -5,39 +5,22 @@ import time
 from functools import wraps
 import sys
 import multimolecule  # Importar para registrar los modelos de multimolecule
+import spaces  # Importar el módulo spaces
 
-# Decorador mejorado para medir el tiempo de ejecución
+# Decorador para medir el tiempo de ejecución
 def medir_tiempo(func):
     @wraps(func)
     def wrapper(*args, **kwargs):
-        try:
-            inicio = time.time()
-            resultado = func(*args, **kwargs)
-            fin = time.time()
-            tiempo_transcurrido = fin - inicio
-            print(f"Tiempo de ejecución de '{func.__name__}': {tiempo_transcurrido:.2f} segundos")
-            return resultado
-        except RecursionError as re:
-            print(f"RecursionError en '{func.__name__}': {re}")
-            raise re
-        except Exception as e:
-            print(f"Error en '{func.__name__}': {e}")
-            raise e
+        inicio = time.time()
+        resultado = func(*args, **kwargs)
+        fin = time.time()
+        tiempo_transcurrido = fin - inicio
+        print(f"Tiempo de ejecución de '{func.__name__}': {tiempo_transcurrido:.2f} segundos")
+        return resultado
     return wrapper
 
-# Configurar el dispositivo
-device = 0 if torch.cuda.is_available() else -1
-if device == -1:
-    print("Advertencia: CUDA no está disponible. Se usará la CPU, lo que puede ser lento.")
-
-# Cargar el pipeline de Fill-Mask
-try:
-    print("Cargando el pipeline de Fill-Mask...")
-    fill_mask = pipeline('fill-mask', model='multimolecule/mrnafm', device=device)
-except Exception as e:
-    print(f"Error al cargar el pipeline de Fill-Mask: {e}")
-    sys.exit(1)
-
+# Decorador para asignar GPU (según la funcionalidad proporcionada por spaces)
+@spaces.GPU
 @medir_tiempo
 def predecir_fill_mask(secuencias):
     """
@@ -77,6 +60,19 @@ def predecir_fill_mask(secuencias):
         print(f"Error durante la predicción: {e}")
         return f"Error al realizar la predicción: {e}"
 
+# Configurar el dispositivo (aunque el decorador @spaces.GPU debería manejar esto)
+device = 0 if torch.cuda.is_available() else -1
+if device == -1:
+    print("Advertencia: CUDA no está disponible. Se usará la CPU, lo que puede ser lento.")
+
+# Cargar el pipeline de Fill-Mask
+try:
+    print("Cargando el pipeline de Fill-Mask...")
+    fill_mask = pipeline('fill-mask', model='multimolecule/mrnafm', device=device)
+except Exception as e:
+    print(f"Error al cargar el pipeline de Fill-Mask: {e}")
+    sys.exit(1)
+
 # Definir la interfaz de Gradio
 titulo = "OmniGenome: Predicción de Fill-Mask para Secuencias de ARN"
 descripcion = (