Update app.py

app.py

@@ -1,57 +1,83 @@
-import gradio as gr
-import subprocess
 import os
 import sys
-os.environ["PYOPENGL_PLATFORM"] = "egl"
+import gradio as gr
+import torch
+import numpy as np
+from transformers import Wav2Vec2Processor
+from visualise.rendering import RenderTool
+from data_utils import torch_data
+from trainer.options import parse_args
+from trainer.config import load_JsonConfig
+from nets import init_model, infer  # Ensure these functions are properly defined
 
-def run_talkshow_model(audio_file):
-    # Path to the TalkSHOW demo script
-    demo_script = 'scripts/demo.py'
-    
-    # Configuration and model parameters
-    config_file = './config/LS3DCG.json'
-    body_model_name = 's2g_LS3DCG'
-    body_model_path = 'experiments/2022-10-19-smplx_S2G-LS3DCG/ckpt-99.pth'
-    
-    # Path of the uploaded audio file
-    audio_file_path = audio_file
-    
-    # Path where the output .mp4 video will be saved
-    output_video_path = './visualise/result.mp4'
+# Set environment variables
+os.environ['PYOPENGL_PLATFORM'] = 'egl'
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+sys.path.append(os.getcwd())
+
+# Load the model and configuration
+def load_models(config_file, face_model_name, face_model_path, body_model_name, body_model_path):
+    args = parse_args()
+    config = load_JsonConfig(config_file)
     
-    # Run the demo.py script with the necessary arguments
-    command = [
-        'python', demo_script,
-        '--config_file', config_file,
-        '--infer',
-        '--audio_file', audio_file_path,
-        '--body_model_name', body_model_name,
-        '--body_model_path', body_model_path,
-        '--id', '0',
-        '--output', output_video_path 
-    ]
+    # Initialize models
+    generator_face = init_model(face_model_name, face_model_path, args, config)
+    generator_body = init_model(body_model_name, body_model_path, args, config)
+
+    # Initialize SMPL-X model
+    smplx_model_params = {
+        'model_path': './visualise/',
+        'model_type': 'smplx',
+        'create_global_orient': True,
+        'create_body_pose': True,
+        'create_betas': True,
+        'num_betas': 300,
+        'create_left_hand_pose': True,
+        'create_right_hand_pose': True,
+        'use_pca': False,
+        'flat_hand_mean': False,
+        'create_expression': True,
+        'num_expression_coeffs': 100,
+        'num_pca_comps': 12,
+        'create_jaw_pose': True,
+        'create_leye_pose': True,
+        'create_reye_pose': True,
+        'create_transl': False,
+        'dtype': torch.float64,
+    }
+    smplx_model = smpl.create(**smplx_model_params).to('cuda')
     
-    try:
-        # Run the subprocess and capture output
-        subprocess.run(command, check=True, capture_output=True, text=True)
-        
-        # Check if the .mp4 file is generated
-        if os.path.exists(output_video_path):
-            return output_video_path  # Return the path of the generated video
-        else:
-            return "Error: Output video not generated."
+    return generator_face, generator_body, smplx_model, config
+
+# Inference function
+def run_inference(audio_file):
+    # Load models
+    generator_face, generator_body, smplx_model, config = load_models(
+        './config/LS3DCG.json',
+        's2g_LS3DCG',
+        'experiments/2022-10-19-smplx_S2G-LS3DCG/ckpt-99.pth',
+        's2g_LS3DCG',
+        'experiments/2022-10-19-smplx_S2G-LS3DCG/ckpt-99.pth'
+    )
     
-    except subprocess.CalledProcessError as e:
-        return f"Error running the model: {e.stderr}"  # Return the error message
-
-# Set up the Gradio interface
-interface = gr.Interface(
-    fn=run_talkshow_model, 
-    inputs=gr.Audio(source="upload", type="filepath"), 
-    outputs=gr.Video(),  # Use gr.Video to output the generated .mp4 video
-    title="TalkSHOW: Audio to Mesh"
+    # Initialize rendering tool
+    rendertool = RenderTool('visualise/video/' + config.Log.name)
+
+    # Inference
+    infer(generator_body, generator_face, smplx_model, rendertool, config, audio_file)
+
+    # Provide output (e.g., path to the rendered video)
+    output_video_path = f'visualise/video/{config.Log.name}/{audio_file.split("/")[-1].split(".")[0]}.npy'
+    return output_video_path
+
+# Gradio interface
+iface = gr.Interface(
+    fn=run_inference,
+    inputs=gr.inputs.Audio(source="upload", type="filepath", label="Upload Audio File"),
+    outputs=gr.outputs.Textbox(label="Output Video Path"),
+    title="Audio to 3D Model Renderer",
+    description="Upload an audio file to generate a 3D model rendering."
 )
 
-# Launch the interface
 if __name__ == "__main__":
-    interface.launch()
+    iface.launch()