File size: 3,701 Bytes
4e424ea ca753f0 4e424ea 7bedcdd 4e424ea ca753f0 73566e5 2562fab 73566e5 935512c 417bb8d 935512c 417bb8d 935512c 4e424ea f0f4c78 4e424ea f0f4c78 4e424ea f0f4c78 4e424ea ca753f0 6c641ac 935512c 417bb8d 935512c f20624c 935512c f20624c 1cfe5df 935512c 417bb8d 0561c55 417bb8d a0044b5 8612a65 417bb8d 935512c 3217fc0 417bb8d 73566e5 ca753f0 417bb8d 935512c f0f4c78 2562fab f20624c 935512c 2562fab f0f4c78 4e424ea d701afa 4e424ea f0f4c78 4e424ea c81f025 4e424ea |
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 |
import gradio as gr
import re
import subprocess
from tqdm import tqdm
from huggingface_hub import snapshot_download
#Download model
snapshot_download(
repo_id = "Wan-AI/Wan2.1-T2V-1.3B",
local_dir = "./Wan2.1-T2V-1.3B"
)
def infer(prompt, progress=gr.Progress(track_tqdm=True)):
total_process_steps = 11
irrelevant_steps = 4
relevant_steps = total_process_steps - irrelevant_steps # 7 steps
# Create the overall progress bar for the process steps.
overall_bar = tqdm(total=relevant_steps, desc="Overall Process", position=1, dynamic_ncols=True, leave=True)
processed_steps = 0
# Regex for detecting video generation progress lines (e.g., "10%|...| 5/50")
progress_pattern = re.compile(r"(\d+)%\|.*\| (\d+)/(\d+)")
gen_progress_bar = None
command = [
"python", "-u", "-m", "generate", # using -u for unbuffered output and omitting .py extension
"--task", "t2v-1.3B",
"--size", "832*480",
"--ckpt_dir", "./Wan2.1-T2V-1.3B",
"--sample_shift", "8",
"--sample_guide_scale", "6",
"--prompt", prompt,
"--save_file", "generated_video.mp4"
]
# Start the process with unbuffered output and combine stdout and stderr.
process = subprocess.Popen(
command,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1 # line-buffered
)
for line in iter(process.stdout.readline, ''):
stripped_line = line.strip()
if not stripped_line:
continue
# Check if this is a video generation progress line.
progress_match = progress_pattern.search(stripped_line)
if progress_match:
current = int(progress_match.group(2))
total = int(progress_match.group(3))
if gen_progress_bar is None:
gen_progress_bar = tqdm(total=total, desc="Video Generation", position=0, dynamic_ncols=True, leave=True)
gen_progress_bar.update(current - gen_progress_bar.n)
gen_progress_bar.refresh()
continue
# Check for INFO lines.
if "INFO:" in stripped_line:
# Split the line at "INFO:" to extract the message.
parts = stripped_line.split("INFO:", 1)
msg = parts[1].strip() if len(parts) > 1 else ""
# Print the log line.
print(stripped_line)
# Skip updating the overall bar for the first three (irrelevant) steps.
if processed_steps < irrelevant_steps:
processed_steps += 1
else:
overall_bar.update(1)
percentage = (overall_bar.n / overall_bar.total) * 100
# Directly assign to the description property.
overall_bar.desc = f"Overall Process - {percentage:.1f}% | {msg}"
overall_bar.refresh()
else:
# Print any other lines.
print(stripped_line)
process.wait()
if gen_progress_bar:
gen_progress_bar.close()
overall_bar.close()
if process.returncode == 0:
print("Command executed successfully.")
return "generated_video.mp4"
else:
print("Error executing command.")
raise Exception("Error executing command")
with gr.Blocks() as demo:
with gr.Column():
gr.Markdown("# Wan 2.1")
prompt = gr.Textbox(label="Prompt")
submit_btn = gr.Button("Submit")
video_res = gr.Video(label="Generated Video")
submit_btn.click(
fn = infer,
inputs = [prompt],
outputs = [video_res]
)
demo.queue().launch(show_error=True, show_api=False, ssr_mode=False) |