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
@@ -10,9 +11,6 @@ from skimage.transform import resize
 from PIL import Image, ImageEnhance, ImageFilter
 from tqdm import tqdm
 
-# Ensure TensorFlow runs on CPU
-os.environ['CUDA_VISIBLE_DEVICES'] = '0'
-
 class SwarmAgent:
     def __init__(self, position, velocity):
         self.position = position
@@ -85,6 +83,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)]
         
@@ -101,6 +100,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
@@ -156,6 +156,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=1000):
         for agent in self.agents:
             agent.position = np.random.randn(*self.image_shape)
@@ -166,6 +167,7 @@ class SwarmNeuralNetwork:
         self.generate_image()
         return self.generated_image
 
+    @spaces.GPU(duration=60)
     def apply_super_resolution(self, image):
         import cv2
         sr = cv2.dnn_superres.DnnSuperResImpl_create()
@@ -177,6 +179,7 @@ class SwarmNeuralNetwork:
         return upscaled / 255.0
 
 # Gradio Interface
+@spaces.GPU(duration=120)
 def train_snn(image, num_agents, epochs, brightness, contrast, color):
     snn = SwarmNeuralNetwork(num_agents=num_agents, image_shape=(256, 256, 3), target_image=image)
     
@@ -192,6 +195,7 @@ def train_snn(image, num_agents, epochs, brightness, contrast, color):
     upscaled_image = snn.apply_super_resolution(generated_image)
     return upscaled_image
 
+@spaces.GPU(duration=120)
 def generate_new_image():
     snn = SwarmNeuralNetwork(num_agents=2000, image_shape=(256, 256, 3), target_image=None)
     snn.load_model('snn_model.npy')