Update app.py

app.py

@@ -11,6 +11,9 @@ from skimage.transform import resize
 from PIL import Image, ImageEnhance, ImageFilter
 from tqdm import tqdm
 
+# Disable GPU usage by default
+os.environ['CUDA_VISIBLE_DEVICES'] = ''
+
 class SwarmAgent:
     def __init__(self, position, velocity):
         self.position = position
@@ -21,7 +24,7 @@ class SwarmAgent:
 class SwarmNeuralNetwork:
     def __init__(self, num_agents, image_shape, target_image):
         self.image_shape = image_shape
-        self.resized_shape = (256, 256, 3)  # Increased resolution
+        self.resized_shape = (256, 256, 3)  # High resolution
         self.agents = [SwarmAgent(self.random_position(), self.random_velocity()) for _ in range(num_agents)]
         self.target_image = self.load_target_image(target_image)
         self.generated_image = np.random.randn(*image_shape)  # Start with noise
@@ -83,7 +86,7 @@ class SwarmNeuralNetwork:
 
         return np.array(losses)
 
-    @spaces.GPU(duration=120)
+    @spaces.GPU(duration=90)
     def update_agents(self, timestep):
         noise_level = self.noise_schedule[min(timestep, len(self.noise_schedule) - 1)]
         
@@ -100,7 +103,7 @@ class SwarmNeuralNetwork:
             # Clip values
             agent.position = np.clip(agent.position, -1, 1)
 
-    @spaces.GPU(duration=120)
+    @spaces.GPU(duration=90)
     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
@@ -110,7 +113,6 @@ class SwarmNeuralNetwork:
         # Apply sharpening filter
         image_pil = Image.fromarray((self.generated_image * 255).astype(np.uint8))
         image_pil = image_pil.filter(ImageFilter.SHARPEN)
-        image_pil = image_pil.filter(ImageFilter.DETAIL)
         self.generated_image = np.array(image_pil) / 255.0
 
     def train(self, epochs):
@@ -123,7 +125,7 @@ class SwarmNeuralNetwork:
             mse = np.mean(((self.generated_image * 2 - 1) - self.target_image)**2)
             logging.info(f"Epoch {epoch}, MSE: {mse}")
 
-            if epoch % 10 == 0:
+            if epoch % 5 == 0:
                 print(f"Epoch {epoch}, MSE: {mse}")
                 self.display_image(self.generated_image, title=f'Epoch {epoch}')
             self.current_epoch += 1
@@ -156,8 +158,8 @@ 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):
+    @spaces.GPU(duration=90)
+    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)
         
@@ -167,21 +169,10 @@ 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()
-        path = "EDSR_x3.pb"  # Path to a pre-trained super-resolution model file (download it from OpenCV's repository)
-        sr.readModel(path)
-        sr.setModel("edsr", 3)  # Use EDSR model with a scale factor of 3
-        image = (image * 255).astype(np.uint8)
-        upscaled = sr.upsample(image)
-        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)
+    snn = SwarmNeuralNetwork(num_agents=num_agents, image_shape=(256, 256, 3), target_image=image)  # High resolution
     
     # Apply user-specified adjustments to the target image
     image = ImageEnhance.Brightness(image).enhance(brightness)
@@ -192,23 +183,21 @@ def train_snn(image, num_agents, epochs, brightness, contrast, color):
     snn.train(epochs=epochs)
     snn.save_model('snn_model.npy')
     generated_image = snn.generated_image
-    upscaled_image = snn.apply_super_resolution(generated_image)
-    return upscaled_image
+    return generated_image
 
 @spaces.GPU(duration=120)
 def generate_new_image():
-    snn = SwarmNeuralNetwork(num_agents=2000, image_shape=(256, 256, 3), target_image=None)
+    snn = SwarmNeuralNetwork(num_agents=2000, image_shape=(256, 256, 3), target_image=None)  # High resolution and optimized number of agents
     snn.load_model('snn_model.npy')
     new_image = snn.generate_new_image()
-    upscaled_image = snn.apply_super_resolution(new_image)
-    return upscaled_image
+    return new_image
 
 interface = gr.Interface(
     fn=train_snn,
     inputs=[
         gr.Image(type="pil", label="Upload Target Image"),
-        gr.Slider(minimum=500, maximum=3000, value=2000, label="Number of Agents"),
-        gr.Slider(minimum=10, maximum=200, value=100, label="Number of Epochs"),
+        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=0.5, maximum=2.0, value=1.0, label="Brightness"),
         gr.Slider(minimum=0.5, maximum=2.0, value=1.0, label="Contrast"),
         gr.Slider(minimum=0.5, maximum=2.0, value=1.0, label="Color Balance")
@@ -218,4 +207,4 @@ interface = gr.Interface(
     description="Upload an image and set the number of agents and epochs to train the Swarm Neural Network to generate a new image. Adjust brightness, contrast, and color balance for personalization."
 )
 
-interface.launch()
+interface.launch()