Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import argparse
 import os
 import tempfile
@@ -96,6 +97,7 @@ def adjust_location(x0, y0, x1, y1, input_image):
     draw.rectangle([(x0,y0),(x1,y1)], outline="red", width=5)
     return x_0, y_0, x_1, y_1, concat_img
 
+@spaces.GPU
 def prepare_data(device, input_image, x0, y0, x1, y1, polar, azimuth, text):
     if input_image.size[0] != 256 or input_image.size[1] != 256:
         input_image = input_image.resize((256, 256))
@@ -127,6 +129,7 @@ def prepare_data(device, input_image, x0, y0, x1, y1, polar, azimuth, text):
     return batch
 
 
+@spaces.GPU
 @torch.no_grad()
 def run_generation(sampler, model, device, input_image, x0, y0, x1, y1, polar, azimuth, text, seed):
     seed_everything(seed)
@@ -182,6 +185,7 @@ def load_example(input_image, x0, y0, x1, y1, polar, azimuth, prompt):
     # print(type(polar))
     return input_image, x0, y0, x1, y1, polar, azimuth, prompt
 
+@spaces.GPU
 @torch.no_grad()
 def main(args):
     # load model