Update app.py

app.py

@@ -1,37 +1,48 @@
 import gradio as gr
 import spaces
-import torch
 import subprocess
 import os
+from PIL import Image
 import ffmpeg
+from pydub import AudioSegment
 
-zero = torch.Tensor([0]).cuda()
-print(zero.device) # <-- 'cpu' 🤔
+import numpy as np
+import soundfile as sf
+
+def save_audio_mp3(audio_tuple, filename):
+    sampling_rate, audio_data = audio_tuple
+    audio_bytes = np.array(audio_data, dtype=np.int16).tobytes()
+    audio_segment = AudioSegment(audio_bytes, sample_width=2, frame_rate=sampling_rate, channels=1)
+    audio_segment.export(filename, format="mp3")
+    
+    return f"Audio saved successfully as {filename}"
 
-@spaces.GPU
-def greet(n):
-    print(zero.device) # <-- 'cuda:0' 🤗
-    return f"Hello {zero + n} Tensor"
 
 def audio_video():
-    print("started =========================")
     input_video = ffmpeg.input('results/result_voice.mp4')
-
-    input_audio = ffmpeg.input('sample_data/sir.mp3')
+    input_audio = ffmpeg.input('sample_data/uploaded_audio.mp3')
     os.system(f"rm -rf results/final_output.mp4")
     ffmpeg.concat(input_video, input_audio, v=1, a=1).output('results/final_output.mp4').run()
     
     return "results/final_output.mp4"
 
-def run_infrence(input_video,input_audio):
-    audio = "sample_data/sir.mp3"
-    video = "sample_data/spark_input.mp4"
-    command = f'python3 inference.py --checkpoint_path checkpoints/wav2lip_gan.pth --face sample_data/spark.png --audio sample_data/sir.mp3'
-    print("running ")
-    # Execute the command
-    process = subprocess.Popen(command, stdout=subprocess.PIPE, shell=True)
+@spaces.GPU
+def run_infrence(input_image,input_audio):
+    pil_image = Image.fromarray(input_image.astype(np.uint8))
 
-    # Get the output
+    save_dir = "sample_data"
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+    
+    # Save input image 
+    filename = os.path.join(save_dir, "uploaded_image.png")
+    pil_image.save(filename)
+
+    #Save input audio
+    save_audio_mp3(input_audio, "sample_data/uploaded_audio.mp3")
+
+    command = f'python3 inference.py --checkpoint_path checkpoints/wav2lip_gan.pth --face sample_data/uploaded_image.png --audio sample_data/uploaded_audio.mp3'
+    process = subprocess.Popen(command, stdout=subprocess.PIPE, shell=True)
     output, error = process.communicate()
 
     return audio_video()
@@ -40,17 +51,14 @@ def run():
   with gr.Blocks(css=".gradio-container {background-color: lightgray} #radio_div {background-color: #FFD8B4; font-size: 40px;}") as demo:
     gr.Markdown("<h1 style='text-align: center;'>"+ "One Shot Talking Face from Text" + "</h1><br/><br/>")
     with gr.Group():
-      # with gr.Box():
         with gr.Row():
-        # with gr.Row().style(equal_height=True):
-            input_video = gr.Video(label="Input Video")
+            input_image = gr.Image(label="Input Image")
             input_audio = gr.Audio(label="Input Audio")
             video_out = gr.Video(show_label=True,label="Output")
         with gr.Row():
             btn = gr.Button("Generate")   
-
-    btn.click(run_infrence,inputs=[input_video,input_audio], outputs=[video_out])
-    # btn.click(run_infrence,inputs=[input_video,input_audio])
+    
+    btn.click(run_infrence,inputs=[input_image,input_audio], outputs=[video_out])
     demo.queue()
     demo.launch(server_name="0.0.0.0", server_port=7860)