fix

app.py

@@ -246,30 +246,32 @@ def extract_video_frames(video_path: str):
 
 
 # ==================== Feature Extraction ====================
-
-def extract_features(clip_chunk: torch.Tensor, sync_chunk: torch.Tensor, caption: str) -> dict:
-    """Reuses globally loaded FeaturesUtils — no reload per call."""
+@spaces.GPU
+def extract_features_gpu(clip_chunk, sync_chunk, caption):
     model = _MODELS["feature_extractor"]
-    assert model is not None, "FeaturesUtils not initialized."
-
     info = {}
     with torch.no_grad():
+
+        model.t5_model.to('cuda')
         text_features = model.encode_t5_text([caption])
         info['text_features'] = text_features[0].cpu()
+        model.t5_model.to('cpu')
 
-        clip_input = torch.from_numpy(clip_chunk).unsqueeze(0)
-        video_feat, frame_embed, _, text_feat = \
-            model.encode_video_and_text_with_videoprism(clip_input, [caption])
 
-        info['global_video_features'] = torch.tensor(np.array(video_feat)).squeeze(0).cpu()
-        info['video_features']        = torch.tensor(np.array(frame_embed)).squeeze(0).cpu()
-        info['global_text_features']  = torch.tensor(np.array(text_feat)).squeeze(0).cpu()
-
-        sync_input = sync_chunk.unsqueeze(0).to(DEVICE)
+        model.synchformer.to('cuda')
+        sync_input = sync_chunk.unsqueeze(0).to('cuda')
         info['sync_features'] = model.encode_video_with_sync(sync_input)[0].cpu()
+        model.synchformer.to('cpu') 
 
     return info
 
+def extract_features(clip_chunk, sync_chunk, caption):
+
+    info = extract_features_cpu(clip_chunk, sync_chunk, caption)
+
+    info.update(extract_features_gpu(clip_chunk, sync_chunk, caption))
+    return info
+
 
 # ==================== Build Meta ====================
 
@@ -288,7 +290,7 @@ def build_meta(info: dict, duration: float, caption: str):
 
 
 # ==================== Diffusion Sampling ====================
-
+@spaces.GPU
 def run_diffusion(audio_latent: torch.Tensor, meta: dict, duration: float) -> torch.Tensor:
     """Reuses globally loaded diffusion model — no reload per call."""
     from PrismAudio.inference.sampling import sample, sample_discrete_euler
@@ -296,20 +298,22 @@ def run_diffusion(audio_latent: torch.Tensor, meta: dict, duration: float) -> to
 
     diffusion    = _MODELS["diffusion"]
     model_config = _MODELS["model_config"]
+    device = 'cuda'
+    diffusion.to("cuda")
     assert diffusion is not None, "Diffusion model not initialized."
 
     diffusion_objective = model_config["model"]["diffusion"]["diffusion_objective"]
     latent_length       = round(SAMPLE_RATE * duration / 2048)
 
     meta_on_device = {
-        k: v.to(DEVICE) if isinstance(v, torch.Tensor) else v
+        k: v.to(device) if isinstance(v, torch.Tensor) else v
         for k, v in meta.items()
     }
     metadata = (meta_on_device,)
 
     with torch.no_grad():
         with torch.amp.autocast('cuda'):
-            conditioning = diffusion.conditioner(metadata, DEVICE)
+            conditioning = diffusion.conditioner(metadata, device)
 
         video_exist = torch.stack([item['video_exist'] for item in metadata], dim=0)
         if 'metaclip_features' in conditioning:
@@ -320,7 +324,7 @@ def run_diffusion(audio_latent: torch.Tensor, meta: dict, duration: float) -> to
                 diffusion.model.model.empty_sync_feat
 
         cond_inputs = diffusion.get_conditioning_inputs(conditioning)
-        noise       = torch.randn([1, diffusion.io_channels, latent_length]).to(DEVICE)
+        noise       = torch.randn([1, diffusion.io_channels, latent_length]).to(device)
 
         with torch.amp.autocast('cuda'):
             if diffusion_objective == "v":
@@ -339,6 +343,7 @@ def run_diffusion(audio_latent: torch.Tensor, meta: dict, duration: float) -> to
             if diffusion.pretransform is not None:
                 fakes = diffusion.pretransform.decode(fakes)
 
+    diffusion.to('cpu')
     return (
         fakes.to(torch.float32)
              .div(torch.max(torch.abs(fakes)))
@@ -351,14 +356,9 @@ def run_diffusion(audio_latent: torch.Tensor, meta: dict, duration: float) -> to
 
 # ==================== Full Inference Pipeline ====================
 
-@spaces.GPU(duration=120)
+
 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)
-    
     total_start_time = time.time()
 
     if video_file is None: