You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
125 lines
4.0 KiB
Python
125 lines
4.0 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
import os
|
|
import subprocess
|
|
from fastapi import FastAPI, File, UploadFile
|
|
from fastapi.responses import FileResponse
|
|
from datetime import datetime
|
|
|
|
app = FastAPI()
|
|
|
|
def save_upload_file(upload_file: UploadFile, filename: str):
|
|
with open(filename, "wb") as buffer:
|
|
buffer.write(upload_file.file.read())
|
|
|
|
def generate_video_command(result_dir: str, img_path: str, audio_path: str, video_path: str):
|
|
return [
|
|
"python", "script.py",
|
|
"--source_image", img_path,
|
|
"--result_dir", result_dir,
|
|
"--driven_audio", audio_path,
|
|
"--ref_eyeblink", video_path,
|
|
]
|
|
|
|
def get_latest_sub_dir(result_dir: str):
|
|
sub_dirs = [os.path.join(result_dir, d) for d in os.listdir(result_dir) if os.path.isdir(os.path.join(result_dir, d))]
|
|
if not sub_dirs:
|
|
return None
|
|
return max(sub_dirs, key=os.path.getmtime)
|
|
|
|
def get_video_duration(video_path: str):
|
|
video_duration_command = [
|
|
"ffprobe",
|
|
"-v", "error",
|
|
"-show_entries", "format=duration",
|
|
"-of", "default=noprint_wrappers=1:nokey=1",
|
|
video_path
|
|
]
|
|
result = subprocess.run(video_duration_command, capture_output=True, text=True)
|
|
return float(result.stdout.strip())
|
|
|
|
def trim_video(input_video_path: str, output_video_path: str, duration: float):
|
|
trim_command = [
|
|
"ffmpeg",
|
|
"-i", input_video_path,
|
|
"-t", str(duration - 2),
|
|
"-c", "copy",
|
|
output_video_path
|
|
]
|
|
subprocess.run(trim_command, check=True)
|
|
|
|
def remove_audio(input_video_path: str, output_video_path: str):
|
|
remove_audio_command = [
|
|
"ffmpeg",
|
|
"-i", input_video_path,
|
|
"-an",
|
|
"-vcodec", "copy",
|
|
output_video_path
|
|
]
|
|
subprocess.run(remove_audio_command, check=True)
|
|
|
|
@app.post("/dynamic-video")
|
|
async def generate_video(
|
|
image: UploadFile = File(...),
|
|
):
|
|
img_path = "dynamic/dynamic_image.png"
|
|
save_upload_file(image, img_path)
|
|
|
|
audio_path = "./examples/driven_audio/dynamic_audio.wav"
|
|
|
|
video_path = "./examples/ref_video/dynamic.mp4"
|
|
|
|
result_dir = os.path.join("results")
|
|
os.makedirs(result_dir, exist_ok=True)
|
|
|
|
command = generate_video_command(result_dir, img_path, audio_path, video_path)
|
|
subprocess.run(command, check=True)
|
|
|
|
latest_sub_dir = get_latest_sub_dir(result_dir)
|
|
if not latest_sub_dir:
|
|
return {"error": "No subdirectory found in result directory"}
|
|
|
|
result_video_path = os.path.join(latest_sub_dir, "dynamic_image##dynamic_audio_enhanced.mp4")
|
|
silent_video_path = os.path.join(latest_sub_dir, "dynamic_image##dynamic_audio_enhanced_dynamic.mp4")
|
|
|
|
if os.path.exists(result_video_path):
|
|
remove_audio(result_video_path, silent_video_path)
|
|
return FileResponse(silent_video_path, media_type='video/mp4')
|
|
else:
|
|
return {"error": "Video file not found"}
|
|
|
|
@app.post("/silent-video")
|
|
async def generate_and_trim_video(
|
|
image: UploadFile = File(...),
|
|
):
|
|
img_path = "silent/silent_image.png"
|
|
save_upload_file(image, img_path)
|
|
|
|
audio_path = "./examples/driven_audio/silent_audio.wav"
|
|
|
|
video_path = "./examples/ref_video/silent.mp4"
|
|
|
|
result_dir = os.path.join("results")
|
|
os.makedirs(result_dir, exist_ok=True)
|
|
|
|
command = generate_video_command(result_dir, img_path, audio_path, video_path)
|
|
subprocess.run(command, check=True)
|
|
|
|
latest_sub_dir = get_latest_sub_dir(result_dir)
|
|
if not latest_sub_dir:
|
|
return {"error": "No subdirectory found in result directory"}
|
|
|
|
result_video_path = os.path.join(latest_sub_dir, "silent_image##silent_audio_enhanced.mp4")
|
|
trimmed_video_path = os.path.join(latest_sub_dir, "silent_image##silent_audio_enhanced_trimmed.mp4")
|
|
|
|
if os.path.exists(result_video_path):
|
|
video_duration = get_video_duration(result_video_path)
|
|
trim_video(result_video_path, trimmed_video_path, video_duration)
|
|
return FileResponse(trimmed_video_path, media_type='video/mp4')
|
|
else:
|
|
return {"error": "Video file not found"}
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|