Update app.py

app.py

@@ -1,41 +1,73 @@
 import gradio as gr
 import subprocess
 import os
+from pydub import AudioSegment
+import tempfile
 
-# Define the command template
-COMMAND_TEMPLATE = "python scripts/demo.py --config_file ./config/LS3DCG.json --infer --audio_file {audio_path} --body_model_name s2g_LS3DCG --body_model_path experiments/2022-10-19-smplx_S2G-LS3DCG/ckpt-99.pth --id 0"
+# Set OpenGL platform
+os.environ["PYOPENGL_PLATFORM"] = "egl"
+os.environ["MESA_GL_VERSION_OVERRIDE"] = "4.1"
+os.system('pip install /home/user/app/pyrender')
 
-# Define the function to process the audio and generate the video
-def process_audio(audio):
-    # Save uploaded audio to a specific file path
-    audio_path = "demo_audio/style.wav"
-    with open(audio_path, "wb") as f:
-        f.write(audio.read())
+# Paths to TalkShow demo script and configuration
+DEMO_SCRIPT_PATH = "scripts/demo.py"
+CONFIG_FILE = "./config/LS3DCG.json"
+BODY_MODEL_PATH = "experiments/2022-10-19-smplx_S2G-LS3DCG/ckpt-99.pth"
+OUTPUT_MP4_PATH = "./output/demo_output.mp4"
 
-    # Run the script command with the specified audio file
-    command = COMMAND_TEMPLATE.format(audio_path=audio_path)
-    subprocess.run(command, shell=True)
 
-    # Specify the path to the output video
-    output_video_path = "./output_video.mp4"  # Update if output location is different
+# Ensure the output directory exists
+os.makedirs(os.path.dirname(OUTPUT_MP4_PATH), exist_ok=True)
 
-    # Check if the video file exists and return it for display
-    if os.path.exists(output_video_path):
-        return output_video_path
-    else:
-        return "Error: Video file not generated."
 
-# Set up Gradio interface
-audio_input = gr.Audio(label="Upload a .wav audio file", type="file")
-video_output = gr.Video(label="Generated Video Output")
+def convert_audio_to_wav(input_audio):
+    # Use a temporary file to save the .wav output
+    temp_wav = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
+    temp_wav_path = temp_wav.name
 
+    # Convert the input audio to .wav format
+    audio = AudioSegment.from_file(input_audio)
+    audio.export(temp_wav_path, format="wav")
+    return temp_wav_path
+
+
+def run_demo_and_get_video(input_audio):
+    # Convert the input audio to a temporary .wav file
+    wav_path = convert_audio_to_wav(input_audio)
+
+    try:
+        # Run the demo script with the specified arguments
+        command = [
+            "python", DEMO_SCRIPT_PATH,
+            "--config_file", CONFIG_FILE,
+            "--infer",
+            "--audio_file", wav_path,
+            "--body_model_name", "s2g_LS3DCG",
+            "--body_model_path", BODY_MODEL_PATH,
+            "--id", "0"
+        ]
+        subprocess.run(command, check=True)
+
+        # Check if the output file exists
+        if os.path.exists(OUTPUT_MP4_PATH):
+            return OUTPUT_MP4_PATH
+        else:
+            return "Error: Output video not generated."
+    except subprocess.CalledProcessError as e:
+        return f"Error: {str(e)}"
+    finally:
+        # Clean up the temporary .wav file after processing
+        os.remove(wav_path)
+
+# Define the Gradio interface
 interface = gr.Interface(
-    fn=process_audio,
-    inputs=audio_input,
-    outputs=video_output,
-    title="Audio-to-Video Generator",
-    description="Upload a .wav file to generate an animated video output."
+    fn=run_demo_and_get_video,
+    inputs=gr.Audio(source="upload", type="filepath"),
+    outputs=gr.Video(type="file"),
+    title="TalkSHOW Demo",
+    description="Generate a video from audio input using TalkSHOW model."
 )
 
-# Launch the Gradio app
-interface.launch()
+# Launch the app
+if __name__ == "__main__":
+    interface.launch()