jiaweir commited on
Commit
b601d28
β€’
1 Parent(s): 6f64936
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -38,7 +38,8 @@ def check_video_input(image_block: Image.Image):
38
  if not os.path.exists(os.path.join('tmp_data', f'{img_hash}_rgba_generated.mp4')):
39
  raise gr.Error("Please generate a video first")
40
 
41
- @spaces.GPU
 
42
  def optimize_stage_1(image_block: Image.Image, preprocess_chk: bool, seed_slider: int):
43
  if not os.path.exists('tmp_data'):
44
  os.makedirs('tmp_data')
@@ -59,7 +60,7 @@ def optimize_stage_1(image_block: Image.Image, preprocess_chk: bool, seed_slider
59
  # return [os.path.join('logs', 'tmp_rgba_model.ply')]
60
  return os.path.join('tmp_data', f'{img_hash}_rgba_generated.mp4')
61
 
62
- @spaces.GPU
63
  def optimize_stage_2(image_block: Image.Image, seed_slider: int):
64
  img_hash = hashlib.sha256(image_block.tobytes()).hexdigest()
65
  subprocess.run(f'python lgm/infer.py big --resume {ckpt_path} --test_path tmp_data/{img_hash}_rgba.png', shell=True)
 
38
  if not os.path.exists(os.path.join('tmp_data', f'{img_hash}_rgba_generated.mp4')):
39
  raise gr.Error("Please generate a video first")
40
 
41
+
42
+ @spaces.GPU(duration=120)
43
  def optimize_stage_1(image_block: Image.Image, preprocess_chk: bool, seed_slider: int):
44
  if not os.path.exists('tmp_data'):
45
  os.makedirs('tmp_data')
 
60
  # return [os.path.join('logs', 'tmp_rgba_model.ply')]
61
  return os.path.join('tmp_data', f'{img_hash}_rgba_generated.mp4')
62
 
63
+ @spaces.GPU(duration=180)
64
  def optimize_stage_2(image_block: Image.Image, seed_slider: int):
65
  img_hash = hashlib.sha256(image_block.tobytes()).hexdigest()
66
  subprocess.run(f'python lgm/infer.py big --resume {ckpt_path} --test_path tmp_data/{img_hash}_rgba.png', shell=True)