GPU

app.py

@@ -15,6 +15,7 @@ import shutil
 import time
 import torch.nn as nn
 import torch.nn.functional as F
+import spaces
 from collections import OrderedDict
 
 
@@ -263,6 +264,7 @@ def swap_t2i_unet(pipe, target_unet_path):
 #         if not os.path.exists(example_filepath):
 #             generated_image.save(example_filepath); print(f"✓ New T2I example saved: {example_filepath}")
 #     return generated_image
+@spaces.GPU(duration=120)
 def generate_t2i(prompt, num_inference_steps):
     global t2i_pipe
     if t2i_pipe is None: raise gr.Error("Text-to-Image model is not loaded.")
@@ -282,6 +284,7 @@ def generate_t2i(prompt, num_inference_steps):
             print(f"✓ New T2I example saved: {example_filepath}")
     return generated_image
 
+@spaces.GPU(duration=120)
 def run_mask_to_image_generation(mask_file_obj, cell_type, num_images, steps, seed):
     if controlnet_pipe is None: raise gr.Error("ControlNet pipeline is not loaded.")
     if mask_file_obj is None: raise gr.Error("Please upload a segmentation mask TIF file.")
@@ -322,6 +325,7 @@ def run_mask_to_image_generation(mask_file_obj, cell_type, num_images, steps, se
         
     return input_display_image, generated_images_pil
 
+@spaces.GPU(duration=120)
 def run_super_resolution(low_res_file_obj, controlnet_model_name, prompt, steps, seed):
     if controlnet_pipe is None: raise gr.Error("ControlNet pipeline is not loaded.")
     if low_res_file_obj is None: raise gr.Error("Please upload a low-resolution TIF file.")
@@ -358,6 +362,7 @@ def run_super_resolution(low_res_file_obj, controlnet_model_name, prompt, steps,
         
     return input_display_image, numpy_to_pil(output_np)
 
+@spaces.GPU(duration=120)
 def run_denoising(noisy_image_np, image_type, steps, seed):
     if controlnet_pipe is None: raise gr.Error("ControlNet pipeline is not loaded.")
     if noisy_image_np is None: raise gr.Error("Please upload a noisy image.")
@@ -387,6 +392,7 @@ def run_denoising(noisy_image_np, image_type, steps, seed):
         
     return numpy_to_pil(noisy_image_np, "L"), numpy_to_pil(output_np)
 
+@spaces.GPU(duration=120)
 def run_segmentation(input_image_np, model_name, diameter, flow_threshold, cellprob_threshold):
     """
     Runs cell segmentation and creates a dark red overlay.
@@ -459,6 +465,7 @@ def run_segmentation(input_image_np, model_name, diameter, flow_threshold, cellp
     
     return numpy_to_pil(input_image_np, "L"), numpy_to_pil(blended_image_np, "RGB")
 
+@spaces.GPU(duration=120)
 def run_classification(input_image_np, model_name):
     """
     Runs classification on a single image using a pre-trained ResNet50 model.