yhzhai commited on
Commit
339b2b8
1 Parent(s): 0ebdef0

add zerogpu compatible

Browse files
Files changed (2) hide show
  1. app.py +3 -1
  2. requirements.txt +1 -0
app.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import os
2
  import random
3
  from datetime import datetime
@@ -15,7 +16,7 @@ from diffusers import (
15
  from diffusers.utils import export_to_video
16
  from peft import PeftModel
17
 
18
- device = "cuda" if torch.cuda.is_available() else "cpu"
19
  mcm_id = "yhzhai/mcm"
20
  basedir = os.getcwd()
21
  savedir = os.path.join(
@@ -152,6 +153,7 @@ cache_pipeline = {
152
  }
153
 
154
 
 
155
  def infer(
156
  base_model, variant, prompt, seed=0, randomize_seed=True, num_inference_steps=4
157
  ):
 
1
+ import spaces
2
  import os
3
  import random
4
  from datetime import datetime
 
16
  from diffusers.utils import export_to_video
17
  from peft import PeftModel
18
 
19
+ device = "cuda"
20
  mcm_id = "yhzhai/mcm"
21
  basedir = os.getcwd()
22
  savedir = os.path.join(
 
153
  }
154
 
155
 
156
+ @spaces.GPU
157
  def infer(
158
  base_model, variant, prompt, seed=0, randomize_seed=True, num_inference_steps=4
159
  ):
requirements.txt CHANGED
@@ -1,5 +1,6 @@
1
  # --extra-index-url https://download.pytorch.org/whl/cu118
2
  # torch==2.1.2
 
3
  torchvision==0.16.2
4
  git+https://github.com/yhZhai/diffusers.git
5
  transformers==4.36.2
 
1
  # --extra-index-url https://download.pytorch.org/whl/cu118
2
  # torch==2.1.2
3
+ spaces
4
  torchvision==0.16.2
5
  git+https://github.com/yhZhai/diffusers.git
6
  transformers==4.36.2