File size: 4,143 Bytes
2945355
 
 
30b0683
0aa3e03
0027dc5
2945355
 
 
 
30b0683
 
 
 
 
2945355
4264aae
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2945355
 
 
8ab8669
2945355
 
 
 
 
 
 
 
 
 
0027dc5
 
 
 
 
 
 
 
 
2945355
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0aa3e03
 
 
 
 
 
 
 
2945355
0aa3e03
2945355
 
 
 
 
 
 
 
 
 
3cc9c1d
47baa10
2945355
 
 
 
 
0aa3e03
2945355
 
 
 
 
 
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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
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"
)

def extract_frames_with_labels(video_path, output_dir="frames"):
    # Ensure output directory exists
    os.makedirs(output_dir, exist_ok=True)
    
    # Open the video file
    video_capture = cv2.VideoCapture(video_path)
    if not video_capture.isOpened():
        raise ValueError(f"Cannot open video file: {video_path}")
    
    frame_data = []
    frame_index = 0
    
    # Loop through the video frames
    while True:
        ret, frame = video_capture.read()
        if not ret:
            break  # Exit the loop if there are no frames left to read

        # Zero-padded frame index for filename and label
        frame_label = f"{frame_index:04}"
        frame_filename = os.path.join(output_dir, f"frame_{frame_label}.jpg")
        
        # Save the frame as a .jpg file
        cv2.imwrite(frame_filename, frame)
        
        # Append the tuple (filename, label) to the list
        frame_data.append((frame_filename, frame_label))
        
        # Increment frame index
        frame_index += 1
    
    # Release the video capture object
    video_capture.release()
    
    return frame_data

# Define a function to run your script with selected inputs
def run_xportrait(
    model_config, 
    output_dir_base, 
    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()