Update gradio_demo.py

gradio_demo.py

@@ -1,3 +1,4 @@
+import spaces
 import os
 import math
 import gradio as gr
@@ -255,7 +256,7 @@ def encode_prompt_inner(txt: str):
 
     return conds
 
-
+@spaces.GPU(duration=60)
 @torch.inference_mode()
 def encode_prompt_pair(positive_prompt, negative_prompt):
     c = encode_prompt_inner(positive_prompt)
@@ -276,7 +277,7 @@ def encode_prompt_pair(positive_prompt, negative_prompt):
 
     return c, uc
 
-
+@spaces.GPU(duration=60)
 @torch.inference_mode()
 def pytorch2numpy(imgs, quant=True):
     results = []
@@ -293,7 +294,7 @@ def pytorch2numpy(imgs, quant=True):
         results.append(y)
     return results
 
-
+@spaces.GPU(duration=60)
 @torch.inference_mode()
 def numpy2pytorch(imgs):
     h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.0 - 1.0  # so that 127 must be strictly 0.0
@@ -321,7 +322,7 @@ def resize_without_crop(image, target_width, target_height):
     resized_image = pil_image.resize((target_width, target_height), Image.LANCZOS)
     return np.array(resized_image)
 
-
+@spaces.GPU(duration=60)
 @torch.inference_mode()
 def run_rmbg(img, sigma=0.0):
     # Convert RGBA to RGB if needed
@@ -346,6 +347,8 @@ def run_rmbg(img, sigma=0.0):
     rgba = np.dstack((img, alpha * 255)).astype(np.uint8)
     result = 127 + (img.astype(np.float32) - 127 + sigma) * alpha
     return result.clip(0, 255).astype(np.uint8), rgba
+
+@spaces.GPU(duration=60)
 @torch.inference_mode()
 def process(input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source):
     clear_memory()
@@ -465,6 +468,8 @@ def process(input_fg, prompt, image_width, image_height, num_samples, seed, step
 
     return pixels
 
+
+@spaces.GPU(duration=60)
 @torch.inference_mode()
 def process_bg(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source):
     clear_memory()
@@ -558,7 +563,7 @@ def process_bg(input_fg, input_bg, prompt, image_width, image_height, num_sample
     clear_memory()
     return pixels, [fg, bg]
 
-
+@spaces.GPU(duration=60)
 @torch.inference_mode()
 def process_relight(input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source):
     input_fg, matting = run_rmbg(input_fg)
@@ -566,7 +571,7 @@ def process_relight(input_fg, prompt, image_width, image_height, num_samples, se
     return input_fg, results
 
 
-
+@spaces.GPU(duration=60)
 @torch.inference_mode()
 def process_relight_bg(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source):
     bg_source = BGSource(bg_source)