ZeroGPU

app.py

@@ -1,3 +1,4 @@
+import spaces
 import contextlib
 import gc
 import json
@@ -147,6 +148,7 @@ examples = [
 global pipeline
 global MultiResNetModel
 
+@spaces.GPU
 def load_ckpt(input_style):
     global pipeline
     global MultiResNetModel
@@ -236,7 +238,7 @@ load_ckpt(cur_input_style)
 cur_input_style = "GrayImage(ScreenStyle)"
 load_ckpt(cur_input_style)
 
-
+@spaces.GPU
 def fix_random_seeds(seed):
     random.seed(seed)
     np.random.seed(seed)
@@ -252,6 +254,7 @@ def process_multi_images(files):
         imgs.append(img)
     return imgs 
 
+@spaces.GPU
 def extract_lines(image):
     src = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
 
@@ -276,6 +279,7 @@ def extract_lines(image):
     torch.cuda.empty_cache()
     return outimg
 
+@spaces.GPU
 def to_screen_image(input_image):
     global opt
     global ScreenModel
@@ -291,6 +295,7 @@ def to_screen_image(input_image):
     torch.cuda.empty_cache()
     return image_pil
 
+@spaces.GPU
 def extract_line_image(query_image_, input_style, resolution):
     if resolution == "640x640":
         tar_width = 640
@@ -317,6 +322,7 @@ def extract_line_image(query_image_, input_style, resolution):
     torch.cuda.empty_cache()
     return input_context, extracted_line, input_context  
 
+@spaces.GPU(duration=180)
 def colorize_image(VAE_input, input_context, reference_images, resolution, seed, input_style, num_inference_steps):
     if VAE_input is None or input_context is None:
         gr.Info("Please preprocess the image first")