import gradio as gr
import subprocess
import os
from huggingface_hub import hf_hub_download
import glob

# 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
):
    # 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()