fix

app.py

@@ -59,7 +59,8 @@ CKPT_PATH             = "ckpts/prismaudio.ckpt"
 VAE_CKPT_PATH         = "ckpts/vae.ckpt"
 VAE_CONFIG_PATH       = "PrismAudio/configs/model_configs/stable_audio_2_0_vae.json"
 SYNCHFORMER_CKPT_PATH = "ckpts/synchformer_state_dict.pth"
-DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+DEVICE = 'cpu'  # 启动时用CPU
+
 
 # ==================== Global Model Registry ====================
 _MODELS = {
@@ -96,7 +97,7 @@ def load_all_models():
         enable_conditions=True,
         synchformer_ckpt=SYNCHFORMER_CKPT_PATH,
     )
-    feature_extractor = feature_extractor.eval().to(DEVICE)
+    feature_extractor = feature_extractor.eval()
     _MODELS["feature_extractor"] = feature_extractor
     log.info("✅ FeaturesUtils loaded")
 
@@ -114,7 +115,7 @@ def load_all_models():
     vae_state = load_ckpt_state_dict(VAE_CKPT_PATH, prefix='autoencoder.')
     diffusion.pretransform.load_state_dict(vae_state)
 
-    diffusion = diffusion.eval().to(DEVICE)
+    diffusion = diffusion.eval()
     _MODELS["diffusion"] = diffusion
     log.info("✅ Diffusion model loaded")
 
@@ -353,7 +354,10 @@ def run_diffusion(audio_latent: torch.Tensor, meta: dict, duration: float) -> to
 @spaces.GPU
 def generate_audio_core(video_file, caption):
     
-
+    global DEVICE
+    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+    _MODELS["feature_extractor"].to(DEVICE)
+    _MODELS["diffusion"].to(DEVICE)
     start_time =time.time()
 
     """
@@ -379,11 +383,9 @@ def generate_audio_core(video_file, caption):
         return "\n".join(logs)
 
     # ---- Working directory (auto-cleaned on exit) ----
-    work_dir = tempfile.mkdtemp(dir=os.environ["GRADIO_TEMP_DIR"], prefix="PrismAudio_")
+    work_dir = tempfile.mkdtemp(prefix="PrismAudio_")
 
     try:
-        if _MODELS["diffusion"] is None:
-            load_all_models()
         # ---- Step 1: Convert / copy to mp4 ----
         status = log_step("📹 Step 1: Preparing video...")
 
@@ -619,7 +621,7 @@ if __name__ == "__main__":
         log.info("✅ All model files found.")
 
     # ⭐ Load all models once at startup
-    #load_all_models()
+    load_all_models()
 
     demo = build_ui()
     demo.queue(max_size=3)