Update hf_gradio_app.py

hf_gradio_app.py

@@ -73,34 +73,31 @@ with torch.inference_mode():
     pipeline = VideoPipeline(vae=vae, reference_net=reference_net, diffusion_net=diffusion_net, scheduler=noise_scheduler, image_proj=image_proj)
     pipeline.to(device=device, dtype=weight_dtype)
 
-def process_audio(file_path):
-    # Create a temporary directory
-    with tempfile.TemporaryDirectory() as temp_dir:
-        # Load the audio file
-        audio = AudioSegment.from_file(file_path)
-        
-        # Check and cut the audio if longer than 4 seconds
-        max_duration = 4 * 1000  # 4 seconds in milliseconds
-        if len(audio) > max_duration:
-            audio = audio[:max_duration]
-        
-        # Save the processed audio in the temporary directory
-        output_path = os.path.join(temp_dir, "trimmed_audio.wav")
-        audio.export(output_path, format="wav")
-        
-        # Temporary file is available here for use
-        print(f"Processed audio saved at: {output_path}")
-        
-        # Return the path for reference (optional)
-        return output_path
+def process_audio(file_path, temp_dir):
+    # Load the audio file
+    audio = AudioSegment.from_file(file_path)
+    
+    # Check and cut the audio if longer than 4 seconds
+    max_duration = 4 * 1000  # 4 seconds in milliseconds
+    if len(audio) > max_duration:
+        audio = audio[:max_duration]
+    
+    # Save the processed audio in the temporary directory
+    output_path = os.path.join(temp_dir, "trimmed_audio.wav")
+    audio.export(output_path, format="wav")
+    
+    # Return the path to the trimmed file
+    print(f"Processed audio saved at: {output_path}")
+    return output_path
 
 @torch.inference_mode()
 def generate(input_video, input_audio, seed, progress=gr.Progress(track_tqdm=True)):
 
     is_shared_ui = True if "fffiloni/MEMO" in os.environ['SPACE_ID'] else False
-
+    temp_dir = None
     if is_shared_ui:
-        input_audio = process_audio(input_audio)
+        temp_dir = tempfile.mkdtemp()
+        input_audio = process_audio(input_audio, temp_dir)
         print(f"Processed file was stored temporarily at: {input_audio}")
     
     resolution = 512
@@ -125,6 +122,11 @@ def generate(input_video, input_audio, seed, progress=gr.Progress(track_tqdm=Tru
     os.makedirs(cache_dir, exist_ok=True)
     input_audio = resample_audio(input_audio, os.path.join(cache_dir, f"{os.path.basename(input_audio).split('.')[0]}-16k.wav"))
 
+    # Clean up the temporary directory
+    if os.path.exists(temp_dir):
+        shutil.rmtree(temp_dir)
+        print(f"Temporary directory {temp_dir} deleted.")
+
     audio_emb, audio_length = preprocess_audio(
         wav_path=input_audio,
         num_generated_frames_per_clip=num_generated_frames_per_clip,