renyuxi commited on
Commit
c491c89
1 Parent(s): a57b3cf

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import argparse
2
  import os
3
  import time
@@ -54,6 +55,7 @@ with gr.Blocks() as demo:
54
  with gr.Column():
55
  output = gr.Gallery(height=1024)
56
 
 
57
  def process_image(num_images, height, width, prompt, seed):
58
  global pipe
59
  with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
 
1
+ import spaces
2
  import argparse
3
  import os
4
  import time
 
55
  with gr.Column():
56
  output = gr.Gallery(height=1024)
57
 
58
+ @spaces.GPU
59
  def process_image(num_images, height, width, prompt, seed):
60
  global pipe
61
  with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):