Update app.py

app.py

@@ -230,8 +230,7 @@ def flex_seq(input_seq, input_file):
         print("Falling back to CPU execution. This may take a while...")
         return core_flex_seq(input_seq, input_file, force_cpu=True)
 
-@spaces.GPU
-def flex_3d(input_file):        
+def core_flex_3d(input_file):        
     if not input_file:
         return None, "Provide a file or a input sequence"
     
@@ -293,10 +292,16 @@ def flex_3d(input_file):
     config = yaml.load(open('configs/train_config.yaml', 'r'), Loader=yaml.FullLoader)
     class_config=ClassConfig(config)
     class_config.adaptor_architecture = 'conv'
-    config['inference_args']['device'] = config['inference_args']['device'] if torch.cuda.is_available() else 'cpu'
-    model, tokenizer = PT5_classification_model(half_precision=config['mixed_precision'], class_config=class_config)
 
+    if force_cpu:
+        target_device = 'cpu'
+    else:
+        target_device = config['inference_args']['device'] if torch.cuda.is_available() else 'cpu'
+    config['inference_args']['device'] = target_device
+
+    model, tokenizer = PT5_classification_model(half_precision=config['mixed_precision'], class_config=class_config)
     model.to(config['inference_args']['device'])
+    
     repo_id = "Honzus24/Flexpert_weights"
     print("Loading 3D model from {}".format(config['inference_args']['3d_model_path']))
     file_weights = config['inference_args']['3d_model_path']
@@ -402,6 +407,26 @@ def flex_3d(input_file):
 
     return output_files, output_message, output_files_enm
 
+@spaces.GPU
+def flex_3d_gpu(input_seq, input_file):
+    """Wrapper that requests ZeroGPU."""
+    return core_flex_3d(input_seq, input_file, force_cpu=False)
+
+def flex_3d(input_seq, input_file):
+    """
+    Main entry point for Gradio. 
+    Tries GPU first, falls back to CPU if quota is exceeded or time limit runs out.
+    """
+    try:
+        print("Attempting to run on ZeroGPU...")
+        return flex_3d_gpu(input_seq, input_file)
+    except Exception as e:
+        # ZeroGPU exceptions (like SpaceTaskError or timeouts) are caught here
+        print(f"ZeroGPU failed or timed out. Reason: {e}")
+        print("Falling back to CPU execution. This may take longer...")
+        return core_flex_3d(input_seq, input_file, force_cpu=True)
+
+    
 def rescale_bfactors(pdb_file):
     base, ext = os.path.splitext(pdb_file)
     # Create the new filename