Spaces:
Sleeping
Sleeping
File size: 2,805 Bytes
488eece 7d436c5 195b61e a40e0c1 195b61e b506075 195b61e b506075 195b61e b506075 195b61e b506075 195b61e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import os
import sys
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
# 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)
# 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')
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'
)
# 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."
)
if __name__ == "__main__":
iface.launch()
|