X-Portrait / app.py
fffiloni's picture
Update app.py
0027dc5 verified
raw
history blame
3.02 kB
import gradio as gr
import subprocess
import os
from huggingface_hub import hf_hub_download
import glob
from datetime import datetime
# Ensure 'checkpoint' directory exists
os.makedirs("checkpoint", exist_ok=True)
hf_hub_download(
repo_id="fffiloni/X-Portrait",
filename="model_state-415001.th",
local_dir="checkpoint"
)
# Define a function to run your script with selected inputs
def run_xportrait(
model_config,
output_dir,
resume_dir,
seed,
uc_scale,
source_image,
driving_video,
best_frame,
out_frames,
num_mix,
ddim_steps
):
# Check if the model weights are in place
if not os.path.exists(resume_dir):
return "Model weights not found in checkpoint directory. Please download them first."
# Create a unique output directory name based on current date and time
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
output_dir = os.path.join(output_dir_base, f"output_{timestamp}")
os.makedirs(output_dir, exist_ok=True)
# Construct the command
command = [
"python3", "core/test_xportrait.py",
"--model_config", model_config,
"--output_dir", output_dir,
"--resume_dir", resume_dir,
"--seed", str(seed),
"--uc_scale", str(uc_scale),
"--source_image", source_image,
"--driving_video", driving_video,
"--best_frame", str(best_frame),
"--out_frames", str(out_frames),
"--num_mix", str(num_mix),
"--ddim_steps", str(ddim_steps)
]
# Run the command
try:
subprocess.run(command, check=True)
# Find the generated video file in the output directory
video_files = glob.glob(os.path.join(output_dir, "*.mp4")) + glob.glob(os.path.join(output_dir, "*.avi"))
print(video_files)
if video_files:
return f"Output video saved at: {video_files[0]}", video_files[0]
else:
return "No video file was found in the output directory.", None
except subprocess.CalledProcessError as e:
return f"An error occurred: {e}", None
# Set up Gradio interface
app = gr.Interface(
fn=run_xportrait,
inputs=[
gr.Textbox(value="config/cldm_v15_appearance_pose_local_mm.yaml", label="Model Config Path"),
gr.Textbox(value="outputs", label="Output Directory"),
gr.Textbox(value="checkpoint/model_state-415001.th", label="Resume Directory"),
gr.Number(value=999, label="Seed"),
gr.Number(value=5, label="UC Scale"),
gr.Image(label="Source Image", type="filepath"),
gr.Video(label="Driving Video"),
gr.Number(value=36, label="Best Frame"),
gr.Number(value=-1, label="Out Frames"),
gr.Number(value=4, label="Number of Mix"),
gr.Number(value=30, label="DDIM Steps")
],
outputs=["text", "video"],
title="XPortrait Model Runner",
description="Run XPortrait with customizable parameters."
)
# Launch the Gradio app
app.launch()