ShoufaChen commited on
Commit
925c081
1 Parent(s): d9a6181

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -1,3 +1,4 @@
 
1
  from PIL import Image
2
  import gradio as gr
3
  from huggingface_hub import hf_hub_download, snapshot_download
@@ -99,6 +100,7 @@ def load_model(args):
99
  return t5_model, vq_model, gpt_model, image_size
100
 
101
 
 
102
  def infer(cfg_scale, top_k, top_p, temperature, prompt, seed):
103
  prompts = [prompt for _ in range(4)]
104
  caption_embs, emb_masks = t5_model.get_text_embeddings(prompts)
 
1
+ import spaces
2
  from PIL import Image
3
  import gradio as gr
4
  from huggingface_hub import hf_hub_download, snapshot_download
 
100
  return t5_model, vq_model, gpt_model, image_size
101
 
102
 
103
+ @spaces.GPU
104
  def infer(cfg_scale, top_k, top_p, temperature, prompt, seed):
105
  prompts = [prompt for _ in range(4)]
106
  caption_embs, emb_masks = t5_model.get_text_embeddings(prompts)