Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import spaces
 import gradio as gr
 import numpy as np
 import tensorflow as tf
@@ -86,6 +87,7 @@ class SwarmNeuralNetwork:
 
         return np.array(losses)
 
+    @spaces.GPU(duration=120)
     def update_agents(self, timestep):
         noise_level = self.noise_schedule[min(timestep, len(self.noise_schedule) - 1)]
         
@@ -102,6 +104,7 @@ class SwarmNeuralNetwork:
             # Clip values
             agent.position = np.clip(agent.position, -1, 1)
 
+    @spaces.GPU(duration=120)
     def generate_image(self):
         self.generated_image = np.mean([agent.position for agent in self.agents], axis=0)
         # Normalize to [0, 1] range for display
@@ -113,6 +116,7 @@ class SwarmNeuralNetwork:
         image_pil = image_pil.filter(ImageFilter.SHARPEN)
         self.generated_image = np.array(image_pil) / 255.0
 
+    @spaces.GPU(duration=120)
     def train(self, epochs):
         logging.basicConfig(filename='training.log', level=logging.INFO)
 
@@ -156,6 +160,7 @@ class SwarmNeuralNetwork:
         self.generated_image = model_state['generated_image']
         self.current_epoch = model_state['current_epoch']
 
+    @spaces.GPU(duration=120)
     def generate_new_image(self, num_steps=500):  # Optimized number of steps
         for agent in self.agents:
             agent.position = np.random.randn(*self.image_shape)