Implement mock spaces module for Zero GPU fallback and improve error handling for GPU decorator application

app.py

@@ -21,25 +21,36 @@ import gc
 # Check if Zero GPU should be disabled via environment variable
 DISABLE_ZERO_GPU = os.environ.get('DISABLE_ZERO_GPU', 'false').lower() == 'true'
 
+# Create a mock spaces module for fallback
+class MockSpaces:
+    """Mock spaces module for when Zero GPU is not available"""
+    def GPU(self, *args, **kwargs):
+        """Mock GPU decorator that just returns the function as-is"""
+        def decorator(func):
+            return func
+        return decorator
+
 # Try to import spaces for Zero GPU support
 if DISABLE_ZERO_GPU:
     SPACES_AVAILABLE = False
+    spaces = MockSpaces()
     print("Zero GPU disabled via environment variable")
 else:
     try:
-        import spaces
+        import spaces as real_spaces
+        # Test if spaces is working properly
+        test_decorator = real_spaces.GPU()
+        spaces = real_spaces
         SPACES_AVAILABLE = True
-        # Check if we're actually running on Hugging Face Spaces with Zero GPU
-        if hasattr(spaces, 'config') and hasattr(spaces.config, 'Config'):
-            print("Zero GPU support detected and available")
-        else:
-            print("Spaces module imported but Zero GPU may not be fully configured")
+        print("Zero GPU support detected and available")
     except ImportError:
         SPACES_AVAILABLE = False
+        spaces = MockSpaces()
         print("Warning: spaces module not available. Running without Zero GPU support.")
     except Exception as e:
         SPACES_AVAILABLE = False
-        print(f"Warning: Error importing spaces module: {e}. Running without Zero GPU support.")
+        spaces = MockSpaces()
+        print(f"Warning: Error with spaces module: {e}. Running without Zero GPU support.")
 
 # Add current directory to path
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
@@ -354,22 +365,16 @@ def score_and_create_matrix_all_singles_impl(sequence,mutation_range_start=None,
       if torch.cuda.is_available():
         torch.cuda.empty_cache()
 
-# Apply Zero GPU decorator if available
-if SPACES_AVAILABLE:
+# Apply Zero GPU decorator - will use real decorator if available, mock otherwise
+try:
+    score_and_create_matrix_all_singles = spaces.GPU(duration=300)(score_and_create_matrix_all_singles_impl)
+except Exception as e:
+    print(f"Warning: Failed to apply GPU decorator with duration: {e}")
     try:
-        # Try with explicit duration
-        score_and_create_matrix_all_singles = spaces.GPU(duration=300)(score_and_create_matrix_all_singles_impl)
-    except Exception as e:
-        print(f"Warning: Failed to apply Zero GPU decorator with duration: {e}")
-        try:
-            # Try without duration parameter
-            score_and_create_matrix_all_singles = spaces.GPU()(score_and_create_matrix_all_singles_impl)
-        except Exception as e2:
-            print(f"Warning: Failed to apply Zero GPU decorator: {e2}")
-            print("Falling back to CPU mode")
-            score_and_create_matrix_all_singles = score_and_create_matrix_all_singles_impl
-else:
-    score_and_create_matrix_all_singles = score_and_create_matrix_all_singles_impl
+        score_and_create_matrix_all_singles = spaces.GPU()(score_and_create_matrix_all_singles_impl)
+    except Exception as e2:
+        print(f"Warning: Failed to apply GPU decorator: {e2}")
+        score_and_create_matrix_all_singles = score_and_create_matrix_all_singles_impl
 
 def extract_sequence(protein_id, taxon, sequence):
   return sequence
@@ -484,6 +489,7 @@ if __name__ == "__main__":
     )
     
     # Launch with appropriate settings for HF Spaces
+    # The spaces module is already handled above (either real or mock)
     tranception_design.launch(
         max_threads=2,  # Limit concurrent threads
         show_error=True,