Update app.py

app.py

@@ -6,6 +6,7 @@ import torch
 import torch.nn.functional as F
 from matplotlib import pyplot as plt
 from PIL import Image
+import spaces
 
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -23,7 +24,7 @@ def image_gen(click,rows = 4,cols = 4):
     return grid
     #return 'button clicked'
     
-
+@spaces.GPU
 def image_gen_modified(rows=4,cols=4):
   pic_hold = []
   model_output = pipeline(batch_size=16).images