update

src/create_modules.py

@@ -35,6 +35,7 @@ class Processer():
     def __init__(self):
         self.create_models()
 
+    @spaces.GPU
     def create_models(self):
 
         self.lmk_extractor = LMKExtractor()
@@ -50,8 +51,8 @@ class Processer():
         audio_infer_config = OmegaConf.load(config.audio_inference_config)
         # prepare model
         self.a2m_model = Audio2MeshModel(audio_infer_config['a2m_model'])
-        self.a2m_model.load_state_dict(torch.load(audio_infer_config['pretrained_model']['a2m_ckpt']), strict=False)
-        self.a2m_model.cuda().eval()
+        self.a2m_model.load_state_dict(torch.load(audio_infer_config['pretrained_model']['a2m_ckpt'], map_location="cpu"), strict=False)
+        self.a2m_model.to("cuda").eval()
 
         self.vae = AutoencoderKL.from_pretrained(
             config.pretrained_vae_path,