Update app.py

app.py

@@ -24,7 +24,7 @@ class SwarmAgent:
 class SwarmNeuralNetwork:
     def __init__(self, num_agents, image_shape, target_image_path):
         self.image_shape = image_shape
-        self.resized_shape = (256, 256, 3)  # High resolution
+        self.resized_shape = (128, 128, 3)  # Reduced resolution
         self.agents = [SwarmAgent(self.random_position(), self.random_velocity()) for _ in range(num_agents)]
         self.target_image = self.load_target_image(target_image_path)
         self.generated_image = np.random.randn(*image_shape)  # Start with noise
@@ -127,7 +127,7 @@ class SwarmNeuralNetwork:
             mse = np.mean(((self.generated_image * 2 - 1) - self.target_image)**2)
             logging.info(f"Epoch {epoch}, MSE: {mse}")
 
-            if epoch % 5 == 0:
+            if epoch % 2 == 0:  # Display more frequently for faster feedback
                 print(f"Epoch {epoch}, MSE: {mse}")
                 self.display_image(self.generated_image, title=f'Epoch {epoch}')
             self.current_epoch += 1
@@ -161,7 +161,7 @@ class SwarmNeuralNetwork:
         self.current_epoch = model_state['current_epoch']
 
     @spaces.GPU(duration=120)
-    def generate_new_image(self, num_steps=500):  # Optimized number of steps
+    def generate_new_image(self, num_steps=200):  # Reduced number of steps
         for agent in self.agents:
             agent.position = np.random.randn(*self.image_shape)
         
@@ -184,7 +184,7 @@ class SwarmNeuralNetwork:
 
 # Gradio Interface
 def train_snn(image_path, num_agents, epochs, arm_position, leg_position, brightness, contrast, color):
-    snn = SwarmNeuralNetwork(num_agents=num_agents, image_shape=(256, 256, 3), target_image_path=image_path)  # High resolution
+    snn = SwarmNeuralNetwork(num_agents=num_agents, image_shape=(128, 128, 3), target_image_path=image_path)  # Reduced resolution
     
     # Apply user-specified adjustments to the target image
     image = Image.open(image_path)
@@ -203,7 +203,7 @@ def train_snn(image_path, num_agents, epochs, arm_position, leg_position, bright
     return generated_image
 
 def generate_new_image():
-    snn = SwarmNeuralNetwork(num_agents=2000, image_shape=(256, 256, 3), target_image_path=None)  # High resolution and optimized number of agents
+    snn = SwarmNeuralNetwork(num_agents=1000, image_shape=(128, 128, 3), target_image_path=None)  # Reduced number of agents
     snn.load_model('snn_model.npy')
     new_image = snn.generate_new_image()
     return new_image
@@ -212,8 +212,8 @@ interface = gr.Interface(
     fn=train_snn,
     inputs=[
         gr.Image(type="filepath", label="Upload Target Image"),
-        gr.Slider(minimum=500, maximum=2000, value=1000, label="Number of Agents"),  # Adjusted range for number of agents
-        gr.Slider(minimum=10, maximum=100, value=50, label="Number of Epochs"),  # Adjusted range for number of epochs
+        gr.Slider(minimum=100, maximum=1000, value=500, label="Number of Agents"),  # Further reduced range for number of agents
+        gr.Slider(minimum=5, maximum=20, value=10, label="Number of Epochs"),  # Further reduced range for number of epochs
         gr.Slider(minimum=-100, maximum=100, value=0, label="Arm Position"),
         gr.Slider(minimum=-100, maximum=100, value=0, label="Leg Position"),
         gr.Slider(minimum=0.5, maximum=2.0, value=1.0, label="Brightness"),