teganmosi commited on
Commit
900bbc0
·
verified ·
1 Parent(s): 749af2a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -3
app.py CHANGED
@@ -1,6 +1,6 @@
1
  import os
2
  import pathlib
3
- import spaces # Added for ZeroGPU
4
  import gradio as gr
5
  import torch
6
  from PIL import Image
@@ -36,12 +36,13 @@ def update_style_image(style_name: str) -> dict:
36
  text = get_style_image_markdown_text(style_name)
37
  return gr.Markdown.update(value=text)
38
 
39
- @spaces.GPU(duration=120) # Added for ZeroGPU, set duration for long inference
40
  def inference(img, vid):
41
  if not os.path.exists('temp'):
42
  os.system('mkdir temp')
43
 
44
  img.save("temp/image.jpg", "JPEG")
 
45
  if torch.cuda.is_available():
46
  os.system(f"python demo.py --config config/vox-256.yaml --checkpoint ./checkpoints/vox.pth.tar --source_image 'temp/image.jpg' --driving_video {vid} --result_video './temp/result.mp4'")
47
  else:
@@ -86,7 +87,7 @@ def main():
86
 
87
  with gr.Box():
88
  gr.Markdown('''## Step 3 (Generate Animated Image based on the Video)
89
- - Hit the **Generate** button. (Note: On cpu-basic, it takes ~ 10 minutes to generate final results.)
90
  ''')
91
  with gr.Row():
92
  with gr.Column():
 
1
  import os
2
  import pathlib
3
+ import spaces # Required for ZeroGPU
4
  import gradio as gr
5
  import torch
6
  from PIL import Image
 
36
  text = get_style_image_markdown_text(style_name)
37
  return gr.Markdown.update(value=text)
38
 
39
+ @spaces.GPU(duration=300) # Increased duration for animation (5 minutes)
40
  def inference(img, vid):
41
  if not os.path.exists('temp'):
42
  os.system('mkdir temp')
43
 
44
  img.save("temp/image.jpg", "JPEG")
45
+ # Check CUDA availability before running
46
  if torch.cuda.is_available():
47
  os.system(f"python demo.py --config config/vox-256.yaml --checkpoint ./checkpoints/vox.pth.tar --source_image 'temp/image.jpg' --driving_video {vid} --result_video './temp/result.mp4'")
48
  else:
 
87
 
88
  with gr.Box():
89
  gr.Markdown('''## Step 3 (Generate Animated Image based on the Video)
90
+ - Hit the **Generate** button. (Note: It may take a few minutes to generate results.)
91
  ''')
92
  with gr.Row():
93
  with gr.Column():