update

app.py

@@ -11,7 +11,7 @@
 # // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # // See the License for the specific language governing permissions and
 # // limitations under the License.
-from spaces import GPU
+import spaces
 
 import os
 import torch
@@ -128,6 +128,7 @@ def configure_sequence_parallel(sp_size):
     if sp_size > 1:
         init_sequence_parallel(sp_size)
 
+@spaces.GPU(duration=120)
 def configure_runner(sp_size):
     config_path = os.path.join('./configs_3b', 'main.yaml')
     config = load_config(config_path)
@@ -143,6 +144,7 @@ def configure_runner(sp_size):
         runner.vae.set_memory_limit(**runner.config.vae.memory_limit)
     return runner
 
+@spaces.GPU(duration=120)
 def generation_step(runner, text_embeds_dict, cond_latents):
     def _move_to_cuda(x):
         return [i.to(get_device()) for i in x]
@@ -199,7 +201,7 @@ def generation_step(runner, text_embeds_dict, cond_latents):
 
     return samples
 
-@GPU
+@spaces.GPU(duration=120)
 def generation_loop(video_path='./test_videos', output_dir='./results', seed=666, batch_size=1, cfg_scale=1.0, cfg_rescale=0.0, sample_steps=1, res_h=1280, res_w=720, sp_size=1):
     runner = configure_runner(1)
     output_dir = 'output/out.mp4'