add decorators to visualizers instead of main function

app.py

@@ -225,7 +225,8 @@ def run_structure_prediction(protein: ESMProtein, temperature: float, num_steps:
     except Exception as e:
         print(f"Error during structure prediction: {str(e)}")
         return None
-
+    
+@spaces.GPU
 def align_after_prediction(protein: ESMProtein, structure_prediction: ESMProtein) -> tuple[ESMProtein, float]:
     if structure_prediction is None:
         return None, float('inf')
@@ -261,7 +262,8 @@ def align_after_prediction(protein: ESMProtein, structure_prediction: ESMProtein
     except Exception as e:
         print(f"Unexpected error during alignment: {str(e)}")
         return None, float('inf')
-    
+
+@spaces.GPU
 def visualize_after_pred(protein: ESMProtein, aligned: ESMProtein):
     if aligned is None:
         return None
@@ -284,6 +286,7 @@ def protein_to_pdb(protein: ESMProtein):
                 pdb_str += f"ATOM  {i*37+j+1:5d}  {atom:3s} {aa:3s} A{i+1:4d}    {x:8.3f}{y:8.3f}{z:8.3f}\n"
     return pdb_str
 
+@spaces.GPU
 def prediction_visualization(pdb_file, num_runs: int, noise_level: float, num_frames: int, temperature: float, num_steps: int, progress=gr.Progress()):
     protein = get_protein(pdb_file)
     runs = []
@@ -310,7 +313,6 @@ def prediction_visualization(pdb_file, num_runs: int, noise_level: float, num_fr
     view_data = visualize_after_pred(protein, best_aligned[1])
     return view_data, f"Best cRMSD: {best_aligned[0]:.4f}"
 
-@spaces.GPU
 def run_prediction(pdb_file, num_runs, noise_level, num_frames, temperature, num_steps, progress=gr.Progress()):
     try:
         if pdb_file is None: