requirements

app.py

@@ -116,15 +116,77 @@ DEFAULT_AUDIO_PATH = os.path.join(
 )
 
 
-@spaces.GPU(duration=60)
-def load_all_models():
-    global \
-        keyframe_model, \
-        interpolation_model, \
-        vae_model, \
-        hubert_model, \
-        wavlm_model, \
-        landmarks_extractor
+# @spaces.GPU(duration=60)
+# def load_all_models():
+#     global \
+#         keyframe_model, \
+#         interpolation_model, \
+#         vae_model, \
+#         hubert_model, \
+#         wavlm_model, \
+#         landmarks_extractor
+#     vae_model = VaeWrapper("video")
+
+#     vae_model = vae_model.half()  # Convert to half precision
+#     try:
+#         vae_model = torch.compile(vae_model)
+#         print("Successfully compiled vae_model in FP16")
+#     except Exception as e:
+#         print(f"Warning: Failed to compile vae_model: {e}")
+
+#     hubert_model = HubertModel.from_pretrained("facebook/hubert-base-ls960").cuda()
+#     hubert_model = hubert_model.half()  # Convert to half precision
+#     try:
+#         hubert_model = torch.compile(hubert_model)
+#         print("Successfully compiled hubert_model in FP16")
+#     except Exception as e:
+#         print(f"Warning: Failed to compile hubert_model: {e}")
+
+#     wavlm_model = WavLM_wrapper(
+#         model_size="Base+",
+#         feed_as_frames=False,
+#         merge_type="None",
+#         model_path=os.path.join(repo_path, "checkpoints/WavLM-Base+.pt"),
+#     ).cuda()
+
+#     wavlm_model = wavlm_model.half()  # Convert to half precision
+#     try:
+#         wavlm_model = torch.compile(wavlm_model)
+#         print("Successfully compiled wavlm_model in FP16")
+#     except Exception as e:
+#         print(f"Warning: Failed to compile wavlm_model: {e}")
+
+#     landmarks_extractor = LandmarksExtractor()
+#     keyframe_model = load_model(
+#         config="keyframe.yaml",
+#         ckpt=os.path.join(repo_path, "checkpoints/keyframe_dub.pt"),
+#     )
+#     interpolation_model = load_model(
+#         config="interpolation.yaml",
+#         ckpt=os.path.join(repo_path, "checkpoints/interpolation_dub.pt"),
+#     )
+#     keyframe_model.en_and_decode_n_samples_a_time = 2
+#     interpolation_model.en_and_decode_n_samples_a_time = 2
+#     return (
+#         keyframe_model,
+#         interpolation_model,
+#         vae_model,
+#         hubert_model,
+#         wavlm_model,
+#         landmarks_extractor,
+#     )
+
+
+# (
+#     keyframe_model,
+#     interpolation_model,
+#     vae_model,
+#     hubert_model,
+#     wavlm_model,
+#     landmarks_extractor,
+# ) = load_all_models()
+
+with spaces.GPU(duration=60) as gpu:
     vae_model = VaeWrapper("video")
 
     vae_model = vae_model.half()  # Convert to half precision
@@ -167,24 +229,6 @@ def load_all_models():
     )
     keyframe_model.en_and_decode_n_samples_a_time = 2
     interpolation_model.en_and_decode_n_samples_a_time = 2
-    return (
-        keyframe_model,
-        interpolation_model,
-        vae_model,
-        hubert_model,
-        wavlm_model,
-        landmarks_extractor,
-    )
-
-
-(
-    keyframe_model,
-    interpolation_model,
-    vae_model,
-    hubert_model,
-    wavlm_model,
-    landmarks_extractor,
-) = load_all_models()
 
 
 @spaces.GPU(duration=60)