[feature] add Gradio live serving demo code (#727)

Co-authored-by: RandNMR73 <notomatthew31@gmail.com>
This commit is contained in:
William Lin
2025-08-10 15:34:03 -07:00
committed by GitHub
co-authored by RandNMR73
parent 2a773fa34e
commit 8cae5e4708
11 changed files with 1446 additions and 234 deletions
-59
View File
@@ -1,59 +0,0 @@
# FastVideo Gradio Demo
This is a Gradio-based web interface for generating videos using the FastVideo framework. The demo allows users to create videos from text prompts with various customization options.
## Overview
The demo uses the FastVideo framework to generate videos based on text prompts. It provides a simple web interface built with Gradio that allows users to:
- Enter text prompts to generate videos
- Customize video parameters (dimensions, number of frames, etc.)
- Use negative prompts to guide the generation process
- Set or randomize seeds for reproducibility
---
## Usage
Run the demo with:
```bash
python examples/inference/gradio/gradio_demo.py
```
This will start a web server at `http://0.0.0.0:7860` where you can access the interface.
---
## Model Initialization
This demo initializes a `VideoGenerator` with the minimum required arguments for inference. Users can seamlessly adjust inference options between generations, including prompts, resolution, video length, or even the number of inference steps, *without ever needing to reload the model*.
## Video Generation
The core functionality is in the `generate_video` function, which:
1. Processes user inputs
2. Uses the FastVideo VideoGenerator from earlier to run inference (`generator.generate_video()`)
3. Returns an output path that Gradio uses to display the generated video
## Gradio Interface
The interface is built with several components:
- A text input for the prompt
- A video display for the result
- Inference options in a collapsible accordion:
- Height and width sliders
- Number of frames slider
- Guidance scale slider
- Inference steps slider
- Negative prompt options
- Seed controls
### Inference Options
- **Height/Width**: Control the resolution of the generated video
- **Number of Frames**: Set how many frames to generate
- **Guidance Scale**: Control how closely the generation follows the prompt
- **Inference Steps**: More steps can improve quality but take longer
- **Negative Prompt**: Specify what you don't want to see in the video
- **Seed**: Control randomness for reproducible results
-169
View File
@@ -1,169 +0,0 @@
import argparse
import os
from copy import deepcopy
import gradio as gr
import torch
from fastvideo import VideoGenerator
from fastvideo.configs.sample.base import SamplingParam
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="FastVideo Gradio Demo")
parser.add_argument("--model_path",
type=str,
default="FastVideo/FastHunyuan-diffusers",
help="Path to the model")
parser.add_argument("--num_gpus",
type=int,
default=1,
help="Number of GPUs to use")
parser.add_argument("--output_path",
type=str,
default="outputs",
help="Path to save generated videos")
parsed_args = parser.parse_args()
# args = FastVideoArgs(model_path="FastVideo/FastHunyuan-Diffusers", num_gpus=2)
generator = VideoGenerator.from_pretrained(
model_path=parsed_args.model_path, num_gpus=parsed_args.num_gpus)
default_params = SamplingParam.from_pretrained(parsed_args.model_path)
def generate_video(
prompt,
negative_prompt,
use_negative_prompt,
seed,
guidance_scale,
num_frames,
height,
width,
num_inference_steps,
randomize_seed=False,
):
params = deepcopy(default_params)
params.prompt = prompt
params.negative_prompt = negative_prompt
params.seed = seed
params.guidance_scale = guidance_scale
params.num_frames = num_frames
params.height = height
params.width = width
params.num_inference_steps = num_inference_steps
if randomize_seed:
params.seed = torch.randint(0, 1000000, (1, )).item()
if not use_negative_prompt:
params.negative_prompt = None
generator.generate_video(prompt=prompt, sampling_param=params)
output_path = os.path.join(parsed_args.output_path,
f"{params.prompt[:100]}.mp4")
return output_path, params.seed
examples = [
"A hand enters the frame, pulling a sheet of plastic wrap over three balls of dough placed on a wooden surface. The plastic wrap is stretched to cover the dough more securely. The hand adjusts the wrap, ensuring that it is tight and smooth over the dough. The scene focuses on the hand’s movements as it secures the edges of the plastic wrap. No new objects appear, and the camera remains stationary, focusing on the action of covering the dough.",
"A vintage train snakes through the mountains, its plume of white steam rising dramatically against the jagged peaks. The cars glint in the late afternoon sun, their deep crimson and gold accents lending a touch of elegance. The tracks carve a precarious path along the cliffside, revealing glimpses of a roaring river far below. Inside, passengers peer out the large windows, their faces lit with awe as the landscape unfolds.",
"A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background. Strings of fairy lights hang above, casting a warm, golden glow over the scene. Groups of people gather around high tables, their laughter blending with the soft rhythm of live jazz. The aroma of freshly mixed cocktails and charred appetizers wafts through the air, mingling with the cool night breeze.",
]
with gr.Blocks() as demo:
gr.Markdown("# FastVideo Inference Demo")
with gr.Group():
with gr.Row():
prompt = gr.Text(
label="Prompt",
show_label=False,
max_lines=1,
placeholder="Enter your prompt",
container=False,
)
run_button = gr.Button("Run", scale=0)
result = gr.Video(label="Result", show_label=False)
with gr.Accordion("Advanced options", open=False):
with gr.Group():
with gr.Row():
height = gr.Slider(
label="Height",
minimum=256,
maximum=1024,
step=32,
value=default_params.height,
)
width = gr.Slider(label="Width",
minimum=256,
maximum=1024,
step=32,
value=default_params.width)
with gr.Row():
num_frames = gr.Slider(
label="Number of Frames",
minimum=21,
maximum=163,
value=default_params.num_frames,
)
guidance_scale = gr.Slider(
label="Guidance Scale",
minimum=1,
maximum=12,
value=default_params.guidance_scale,
)
num_inference_steps = gr.Slider(
label="Inference Steps",
minimum=4,
maximum=100,
value=default_params.num_inference_steps,
)
with gr.Row():
use_negative_prompt = gr.Checkbox(
label="Use negative prompt", value=False)
negative_prompt = gr.Text(
label="Negative prompt",
max_lines=1,
placeholder="Enter a negative prompt",
visible=False,
)
seed = gr.Slider(label="Seed",
minimum=0,
maximum=1000000,
step=1,
value=default_params.seed)
randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
seed_output = gr.Number(label="Used Seed")
gr.Examples(examples=examples, inputs=prompt)
use_negative_prompt.change(
fn=lambda x: gr.update(visible=x),
inputs=use_negative_prompt,
outputs=default_params.negative_prompt,
)
run_button.click(
fn=generate_video,
inputs=[
prompt,
negative_prompt,
use_negative_prompt,
seed,
guidance_scale,
num_frames,
height,
width,
num_inference_steps,
randomize_seed,
],
outputs=[result, seed_output],
)
demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)
@@ -0,0 +1,730 @@
import argparse
import os
import requests
import base64
import time
import gradio as gr
from fastvideo.configs.sample.base import SamplingParam
MODEL_PATH_MAPPING = {
"FastWan2.1-T2V-1.3B": "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
"FastWan2.2-TI2V-5B-FullAttn": "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
}
class RayServeClient:
def __init__(self, backend_url: str):
self.backend_url = backend_url
self.session = requests.Session()
def check_health(self) -> bool:
try:
response = self.session.get(f"{self.backend_url}/health", timeout=5)
return response.status_code == 200
except requests.exceptions.RequestException:
return False
def generate_video(self, request_data: dict) -> dict:
start_time = time.time()
try:
headers = {"Content-Type": "application/json"}
response = self.session.post(
f"{self.backend_url}/generate_video",
json=request_data,
headers=headers,
timeout=300
)
round_trip_time = time.time() - start_time
if response.status_code == 200:
result = response.json()
backend_total = result.get("total_time", 0)
network_time = round_trip_time - backend_total
result["network_time"] = network_time
return result
else:
return {"success": False, "error_message": f"HTTP {response.status_code}: {response.text}"}
except requests.exceptions.RequestException as e:
return {"success": False, "error_message": f"Request failed: {str(e)}"}
def save_video_from_base64(video_data: str, output_dir: str, prompt: str) -> str:
if not video_data:
return None
try:
if video_data.startswith('data:video/'):
video_data = video_data.split(',')[1]
video_bytes = base64.b64decode(video_data)
safe_prompt = prompt[:50].replace(' ', '_').replace('/', '_').replace('\\', '_')
video_filename = f"{safe_prompt}.mp4"
video_path = os.path.join(output_dir, video_filename)
os.makedirs(output_dir, exist_ok=True)
with open(video_path, 'wb') as f:
f.write(video_bytes)
return video_path
except Exception as e:
print(f"Failed to save video: {e}")
return None
def create_timing_display(inference_time, encoding_time, network_time, total_time, stage_execution_times, num_frames):
dit_denoising_time = f"{stage_execution_times[5]:.2f}s" if len(stage_execution_times) > 5 else "N/A"
timing_html = f"""
<div style="margin: 10px 0;">
<h3 style="text-align: center; margin-bottom: 10px;">⏱️ Timing Breakdown</h3>
<div style="display: grid; grid-template-columns: repeat(5, 1fr); gap: 10px; margin-bottom: 10px;">
<div class="timing-card timing-card-highlight">
<div style="font-size: 20px;">🚀</div>
<div style="font-weight: bold; margin: 3px 0; font-size: 14px;">DiT Denoising</div>
<div style="font-size: 16px; color: #ffa200; font-weight: bold;">{dit_denoising_time}</div>
</div>
<div class="timing-card">
<div style="font-size: 20px;">🧠</div>
<div style="font-weight: bold; margin: 3px 0; font-size: 14px;">E2E (w. vae/text encoder)</div>
<div style="font-size: 16px; color: #2563eb;">{inference_time:.2f}s</div>
</div>
<div class="timing-card">
<div style="font-size: 20px;">🎬</div>
<div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Video Encoding</div>
<div style="font-size: 16px; color: #dc2626;">{encoding_time:.2f}s</div>
</div>
<div class="timing-card">
<div style="font-size: 20px;">🌐</div>
<div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Network Transfer</div>
<div style="font-size: 16px; color: #059669;">{network_time:.2f}s</div>
</div>
<div class="timing-card">
<div style="font-size: 20px;">📊</div>
<div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Total Processing</div>
<div style="font-size: 18px; color: #0277bd;">{total_time:.2f}s</div>
</div>
</div>"""
if inference_time > 0:
fps = num_frames / inference_time
timing_html += f"""
<div class="performance-card" style="margin-top: 15px;">
<span style="font-weight: bold;">Generation Speed: </span>
<span style="font-size: 18px; color: #6366f1; font-weight: bold;">{fps:.1f} frames/second</span>
</div>"""
return timing_html + "</div>"
def load_example_prompts():
def contains_chinese(text):
return any('\u4e00' <= char <= '\u9fff' for char in text)
def load_from_file(filepath):
prompts, labels = [], []
try:
with open(filepath, "r", encoding='utf-8') as f:
for line in f:
line = line.strip()
if line and not contains_chinese(line):
label = line[:100] + "..." if len(line) > 100 else line
labels.append(label)
prompts.append(line)
except Exception as e:
print(f"Warning: Could not read {filepath}: {e}")
return prompts, labels
examples, example_labels = load_from_file("prompts/prompts_final.txt")
if not examples:
examples = ["A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background."]
example_labels = ["Crowded rooftop bar at night"]
return examples, example_labels
def create_gradio_interface(backend_url: str, default_params: dict[str, SamplingParam]):
client = RayServeClient(backend_url)
def generate_video(
prompt, negative_prompt, use_negative_prompt, seed, guidance_scale,
num_frames, height, width, randomize_seed, model_selection, progress
):
if not client.check_health():
return None, f"Backend is not available. Please check if Ray Serve is running at {backend_url}", ""
# Validate dimensions
max_pixels = 720 * 1280
if height * width > max_pixels:
return None, f"Video dimensions too large. Maximum: 720x1280 pixels", ""
if progress:
progress(0.1, desc="Checking backend health...")
request_data = {
"prompt": prompt,
"negative_prompt": negative_prompt,
"use_negative_prompt": use_negative_prompt,
"seed": seed,
"guidance_scale": guidance_scale,
"num_frames": num_frames,
"height": height,
"width": width,
"randomize_seed": randomize_seed,
"return_frames": False,
"image_path": None,
"model_path": MODEL_PATH_MAPPING.get(model_selection, "FastVideo/FastWan2.1-T2V-1.3B-Diffusers")
}
if progress:
progress(0.4, desc="Generating video...")
response = client.generate_video(request_data)
if progress:
progress(0.8, desc="Processing response...")
if response.get("success", False):
video_data = response.get("video_data", "")
used_seed = response.get("seed", seed)
inference_time = response.get("inference_time", 0.0)
encoding_time = response.get("encoding_time", 0.0)
total_time = response.get("total_time", 0.0)
network_time = response.get("network_time", 0.0)
stage_execution_times = response.get("stage_execution_times", [])
timing_details = create_timing_display(
inference_time, encoding_time, network_time, total_time,
stage_execution_times, num_frames
)
if video_data:
if progress:
progress(0.9, desc="Saving video...")
video_path = save_video_from_base64(video_data, "outputs", prompt)
if progress:
progress(1.0, desc="Generation complete!")
if video_path and os.path.exists(video_path):
return video_path, used_seed, timing_details
else:
return None, "Failed to save video", ""
else:
return None, "No video data received from backend", ""
else:
error_msg = response.get("error_message", "Unknown error occurred")
return None, f"Generation failed: {error_msg}", ""
examples, example_labels = load_example_prompts()
theme = gr.themes.Base().set(
button_primary_background_fill="#2563eb",
button_primary_background_fill_hover="#1d4ed8",
button_primary_text_color="white",
slider_color="#2563eb",
checkbox_background_color_selected="#2563eb",
)
def get_default_values(model_name):
model_path = MODEL_PATH_MAPPING.get(model_name)
if model_path and model_path in default_params:
params = default_params[model_path]
return {
'height': params.height,
'width': params.width,
'num_frames': params.num_frames,
'guidance_scale': params.guidance_scale,
'seed': params.seed,
}
return {
'height': 448,
'width': 832,
'num_frames': 61,
'guidance_scale': 3.0,
'seed': 1024,
}
initial_values = get_default_values("FastWan2.1-T2V-1.3B")
with gr.Blocks(title="FastWan", theme=theme) as demo:
gr.Image("fastvideo-logos/main/svg/full.svg", show_label=False, container=False, height=80)
gr.HTML("""
<div style="text-align: center; margin-bottom: 10px;">
<p style="font-size: 18px;"> Make Video Generation Go Blurrrrrrr </p>
<p style="font-size: 18px;"> <a href="https://github.com/hao-ai-lab/FastVideo/tree/main" target="_blank">Code</a> | <a href="https://hao-ai-lab.github.io/blogs/fastvideo_post_training/" target="_blank">Blog</a> | <a href="https://hao-ai-lab.github.io/FastVideo/" target="_blank">Docs</a> </p>
</div>
""")
with gr.Accordion("🎥 What Is FastVideo?", open=False):
gr.HTML("""
<div style="padding: 20px; line-height: 1.6;">
<p style="font-size: 16px; margin-bottom: 15px;">
FastVideo is an inference and post-training framework for diffusion models. It features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.
</p>
</div>
""")
with gr.Row():
model_selection = gr.Dropdown(
choices=list(MODEL_PATH_MAPPING.keys()),
value="FastWan2.1-T2V-1.3B",
label="Select Model",
interactive=True
)
with gr.Row():
example_dropdown = gr.Dropdown(
choices=example_labels,
label="Example Prompts",
value=None,
interactive=True,
allow_custom_value=False
)
with gr.Row():
with gr.Column(scale=6):
prompt = gr.Text(
label="Prompt",
show_label=False,
max_lines=3,
placeholder="Describe your scene...",
container=False,
lines=3,
autofocus=True,
)
with gr.Column(scale=1, min_width=120, elem_classes="center-button"):
run_button = gr.Button("Run", variant="primary", size="lg")
with gr.Row():
with gr.Column():
error_output = gr.Text(label="Error", visible=False)
timing_display = gr.Markdown(label="Timing Breakdown", visible=False)
with gr.Row(equal_height=True, elem_classes="main-content-row"):
with gr.Column(scale=1, elem_classes="advanced-options-column"):
with gr.Group():
gr.HTML("<div style='margin: 0 0 15px 0; text-align: center; font-size: 16px;'>Advanced Options</div>")
with gr.Row():
height = gr.Number(
label="Height",
value=initial_values['height'],
interactive=False,
container=True
)
width = gr.Number(
label="Width",
value=initial_values['width'],
interactive=False,
container=True
)
with gr.Row():
num_frames = gr.Number(
label="Number of Frames",
value=initial_values['num_frames'],
interactive=False,
container=True
)
guidance_scale = gr.Slider(
label="Guidance Scale",
minimum=1,
maximum=12,
value=initial_values['guidance_scale'],
)
with gr.Row():
use_negative_prompt = gr.Checkbox(
label="Use negative prompt", value=False)
negative_prompt = gr.Text(
label="Negative prompt",
max_lines=3,
lines=3,
placeholder="Enter a negative prompt",
visible=False,
)
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=1000000,
step=1,
value=initial_values['seed'],
)
randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
seed_output = gr.Number(label="Used Seed")
with gr.Column(scale=1, elem_classes="video-column"):
result = gr.Video(
label="Generated Video",
show_label=True,
height=466,
width=600,
container=True,
elem_classes="video-component"
)
gr.HTML("""
<style>
.center-button {
display: flex !important;
justify-content: center !important;
height: 100% !important;
padding-top: 1.4em !important;
}
.gradio-container {
max-width: 1200px !important;
margin: 0 auto !important;
}
.main {
max-width: 1200px !important;
margin: 0 auto !important;
}
.gr-form, .gr-box, .gr-group {
max-width: 1200px !important;
}
.gr-video {
max-width: 500px !important;
margin: 0 auto !important;
}
.main-content-row {
display: flex !important;
align-items: flex-start !important;
min-height: 500px !important;
gap: 20px !important;
}
.advanced-options-column,
.video-column {
display: flex !important;
flex-direction: column !important;
flex: 1 !important;
min-height: 400px !important;
align-items: stretch !important;
}
.video-column > * {
margin-top: 0 !important;
}
.video-column .gr-video,
.video-component {
margin-top: 0 !important;
padding-top: 0 !important;
}
.video-column .gr-video .gr-form {
margin-top: 0 !important;
}
.advanced-options-column .gr-group,
.video-column .gr-video {
margin-top: 0 !important;
vertical-align: top !important;
}
.advanced-options-column > *:last-child,
.video-column > *:last-child {
flex-grow: 0 !important;
}
@media (max-width: 1400px) {
.main-content-row {
min-height: 600px !important;
}
.advanced-options-column,
.video-column {
min-height: 600px !important;
}
}
@media (max-width: 1200px) {
.main-content-row {
flex-direction: column !important;
align-items: stretch !important;
}
.advanced-options-column,
.video-column {
min-height: auto !important;
width: 100% !important;
}
}
.timing-card {
background: var(--background-fill-secondary) !important;
border: 1px solid var(--border-color-primary) !important;
color: var(--body-text-color) !important;
padding: 10px;
border-radius: 8px;
text-align: center;
min-height: 80px;
display: flex;
flex-direction: column;
justify-content: center;
}
.timing-card-highlight {
background: var(--background-fill-primary) !important;
border: 2px solid var(--color-accent) !important;
}
.performance-card {
background: var(--background-fill-secondary) !important;
border: 1px solid var(--border-color-primary) !important;
color: var(--body-text-color) !important;
padding: 10px;
border-radius: 6px;
text-align: center;
}
.gr-number input[readonly] {
background-color: var(--background-fill-secondary) !important;
border: 1px solid var(--border-color-primary) !important;
color: var(--body-text-color-subdued) !important;
cursor: default !important;
text-align: center !important;
font-weight: 500 !important;
}
</style>
""")
def on_example_select(example_label):
if example_label and example_label in example_labels:
index = example_labels.index(example_label)
return examples[index]
return ""
example_dropdown.change(
fn=on_example_select,
inputs=example_dropdown,
outputs=prompt,
)
gr.HTML("""
<div style="text-align: center; margin-top: 10px; margin-bottom: 15px;">
<p style="font-size: 16px; margin: 0;">The compute for this demo is generously provided by <a href="https://www.gmicloud.ai/" target="_blank">GMI Cloud</a>. Note that this demo is meant to showcase FastWan's quality and that under a large number of requests, generation speed may be affected. We are also rate-limiting users to 3 requests per minute.</p>
</div>
""")
use_negative_prompt.change(
fn=lambda x: gr.update(visible=x),
inputs=use_negative_prompt,
outputs=negative_prompt,
)
def on_model_selection_change(selected_model):
if not selected_model:
selected_model = "FastWan2.1-T2V-1.3B"
model_path = MODEL_PATH_MAPPING.get(selected_model)
if model_path and model_path in default_params:
params = default_params[model_path]
return (
gr.update(value=params.height),
gr.update(value=params.width),
gr.update(value=params.num_frames),
gr.update(value=params.guidance_scale),
gr.update(value=params.seed),
)
return (
gr.update(value=448),
gr.update(value=832),
gr.update(value=61),
gr.update(value=3.0),
gr.update(value=1024),
)
model_selection.change(
fn=on_model_selection_change,
inputs=model_selection,
outputs=[height, width, num_frames, guidance_scale, seed],
)
def handle_generation(*args, progress=None, request: gr.Request = None):
model_selection, prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, num_frames, height, width, randomize_seed = args
result_path, seed_or_error, timing_details = generate_video(
prompt, negative_prompt, use_negative_prompt, seed, guidance_scale,
num_frames, height, width, randomize_seed, model_selection, progress
)
if result_path and os.path.exists(result_path):
return (
result_path,
seed_or_error,
gr.update(visible=False),
gr.update(visible=True, value=timing_details),
)
else:
return (
None,
seed_or_error,
gr.update(visible=True, value=seed_or_error),
gr.update(visible=False),
)
run_button.click(
fn=handle_generation,
inputs=[
model_selection,
prompt,
negative_prompt,
use_negative_prompt,
seed,
guidance_scale,
num_frames,
height,
width,
randomize_seed,
],
outputs=[result, seed_output, error_output, timing_display],
concurrency_limit=20,
)
return demo
def main():
parser = argparse.ArgumentParser(description="FastVideo Gradio Frontend")
parser.add_argument("--backend_url", type=str, default="http://localhost:8000",
help="URL of the Ray Serve backend")
parser.add_argument("--t2v_model_paths", type=str,
default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.1-T2V-14B-Diffusers",
help="Comma separated list of paths to the T2V model(s)")
parser.add_argument("--host", type=str, default="0.0.0.0",
help="Host to bind to")
parser.add_argument("--port", type=int, default=7860,
help="Port to bind to")
args = parser.parse_args()
default_params = {}
model_paths = args.t2v_model_paths.split(",")
for model_path in model_paths:
default_params[model_path] = SamplingParam.from_pretrained(model_path)
demo = create_gradio_interface(args.backend_url, default_params)
print(f"Starting Gradio frontend at http://{args.host}:{args.port}")
print(f"Backend URL: {args.backend_url}")
print(f"T2V Models: {args.t2v_model_paths}")
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import HTMLResponse, FileResponse
import uvicorn
app = FastAPI()
@app.get("/logo.png")
def get_logo():
return FileResponse(
"fastvideo-logos/main/svg/full.svg",
media_type="image/svg+xml",
headers={
"Cache-Control": "public, max-age=3600",
"Access-Control-Allow-Origin": "*"
}
)
@app.get("/favicon.ico")
def get_favicon():
favicon_path = "fastvideo-logos/main/svg/icon-simple.svg"
if os.path.exists(favicon_path):
return FileResponse(
favicon_path,
media_type="image/svg+xml",
headers={
"Cache-Control": "public, max-age=3600",
"Access-Control-Allow-Origin": "*"
}
)
else:
raise HTTPException(status_code=404, detail="Favicon not found")
@app.get("/", response_class=HTMLResponse)
def index(request: Request):
base_url = str(request.base_url).rstrip('/')
return f"""
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
<title>FastWan</title>
<meta name="title" content="FastWan">
<meta name="description" content="Make video generation go blurrrrrrr">
<meta name="keywords" content="FastVideo, video generation, AI, machine learning, FastWan">
<meta property="og:type" content="website">
<meta property="og:url" content="{base_url}/">
<meta property="og:title" content="FastWan">
<meta property="og:description" content="Make video generation go blurrrrrrr">
<meta property="og:image" content="{base_url}/logo.png">
<meta property="og:image:width" content="1200">
<meta property="og:image:height" content="630">
<meta property="og:site_name" content="FastWan">
<meta property="twitter:card" content="summary_large_image">
<meta property="twitter:url" content="{base_url}/">
<meta property="twitter:title" content="FastWan">
<meta property="twitter:description" content="Make video generation go blurrrrrrr">
<meta property="twitter:image" content="{base_url}/logo.png">
<link rel="icon" type="image/png" sizes="32x32" href="/favicon.ico">
<link rel="icon" type="image/png" sizes="16x16" href="/favicon.ico">
<link rel="apple-touch-icon" href="/favicon.ico">
<style>
body, html {{
margin: 0;
padding: 0;
height: 100%;
overflow: hidden;
}}
iframe {{
width: 100%;
height: 100vh;
border: none;
}}
</style>
</head>
<body>
<iframe src="/gradio" width="100%" height="100%" style="border: none;"></iframe>
</body>
</html>
"""
app = gr.mount_gradio_app(
app,
demo,
path="/gradio",
allowed_paths=[os.path.abspath("outputs"), os.path.abspath("fastvideo-logos")]
)
uvicorn.run(app, host=args.host, port=args.port)
if __name__ == "__main__":
main()
@@ -0,0 +1,379 @@
import time
import os
import torch
import base64
import io
from copy import deepcopy
from typing import Dict, Any, Optional, List
import signal
import sys
import ray
from ray import serve
from fastapi import FastAPI, Request, Response
from pydantic import BaseModel
import numpy as np
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
import imageio
from ray.serve.handle import DeploymentHandle
from prometheus_client import Counter, Histogram, generate_latest
NUM_GPUS = 16
DEFAULT_FPS = 16
SEED_RANGE_MAX = 1_000_000
SUPPORTED_MODELS = [
"FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
"FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
]
MODEL_CONFIGS = {
"1.3B": {
"num_cpus": 2,
"text_encoder_cpu_offload": False,
"dit_cpu_offload": False,
"vae_cpu_offload": False,
"VSA_sparsity": 0.8,
},
"14B": {
"num_cpus": 16,
"text_encoder_cpu_offload": True,
"dit_cpu_offload": True,
"vae_cpu_offload": False,
"VSA_sparsity": 0.9,
}
}
class VideoGenerationRequest(BaseModel):
prompt: str
negative_prompt: Optional[str] = None
use_negative_prompt: bool = False
seed: int = 42
guidance_scale: float = 7.5
num_frames: int = 21
height: int = 448
width: int = 832
randomize_seed: bool = False
return_frames: bool = False
model_path: Optional[str] = None
class VideoGenerationResponse(BaseModel):
video_data: Optional[str] = None
seed: int
success: bool
error_message: Optional[str] = None
generation_time: Optional[float] = None
model_load_time: Optional[float] = None
inference_time: Optional[float] = None
encoding_time: Optional[float] = None
total_time: Optional[float] = None
stage_names: Optional[List[str]] = None
stage_execution_times: Optional[List[float]] = None
def encode_video_to_base64(frames: List[np.ndarray], fps: int = DEFAULT_FPS) -> str:
if not frames:
return ""
try:
buffer = io.BytesIO()
imageio.mimsave(buffer, frames, fps=fps, format="mp4")
buffer.seek(0)
video_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
return f"data:video/mp4;base64,{video_base64}"
except Exception as e:
print(f"Warning: Failed to encode video: {e}")
return ""
def setup_model_environment(model_path: str) -> None:
if "fullattn" in model_path.lower():
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
else:
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
os.environ["FASTVIDEO_STAGE_LOGGING"] = "1"
def process_generation_result(result: Any) -> tuple[List[np.ndarray], float, List[str], List[float]]:
frames = result if isinstance(result, list) else result.get("frames", [])
generation_time = result.get("generation_time", 0.0) if isinstance(result, dict) else 0.0
logging_info = result.get("logging_info", None)
if logging_info:
stage_names = logging_info.get_execution_order()
stage_execution_times = [
logging_info.get_stage_info(stage_name).get("execution_time", 0.0)
for stage_name in stage_names
]
else:
stage_names = []
stage_execution_times = []
return frames, generation_time, stage_names, stage_execution_times
def prepare_sampling_params(video_request: VideoGenerationRequest, default_params: Any) -> Any:
params = deepcopy(default_params)
params.prompt = video_request.prompt
if video_request.use_negative_prompt:
params.negative_prompt = video_request.negative_prompt
params.seed = (video_request.seed if not video_request.randomize_seed
else torch.randint(0, SEED_RANGE_MAX, (1,)).item())
params.randomize_seed = video_request.randomize_seed
params.guidance_scale = video_request.guidance_scale
params.num_frames = video_request.num_frames
params.height = video_request.height
params.width = video_request.width
params.save_video = False
params.return_frames = False
return params
class BaseModelDeployment:
def __init__(self, model_path: str, output_path: str = "outputs"):
self.model_path = model_path
self.output_path = output_path
self.generator = None
self.default_params = None
os.makedirs(self.output_path, exist_ok=True)
setup_model_environment(self.model_path)
def _initialize_generator(self, config: Dict[str, Any]) -> None:
from fastvideo.entrypoints.video_generator import VideoGenerator
from fastvideo.configs.sample.base import SamplingParam
print(f"Initializing model: {self.model_path}")
self.generator = VideoGenerator.from_pretrained(
model_path=self.model_path,
num_gpus=1,
use_fsdp_inference=True,
text_encoder_cpu_offload=config["text_encoder_cpu_offload"],
dit_cpu_offload=config["dit_cpu_offload"],
vae_cpu_offload=config["vae_cpu_offload"],
VSA_sparsity=config["VSA_sparsity"],
enable_stage_verification=False,
)
self.default_params = SamplingParam.from_pretrained(self.model_path)
def generate_video(self, video_request: VideoGenerationRequest) -> VideoGenerationResponse:
total_start_time = time.time()
params = prepare_sampling_params(video_request, self.default_params)
inference_start_time = time.time()
result = self.generator.generate_video(
prompt=video_request.prompt,
sampling_param=params,
save_video=False,
return_frames=False,
)
inference_time = time.time() - inference_start_time
frames, generation_time, stage_names, stage_execution_times = process_generation_result(result)
encoding_start_time = time.time()
video_data = encode_video_to_base64(frames, fps=DEFAULT_FPS)
encoding_time = time.time() - encoding_start_time
total_time = time.time() - total_start_time
return VideoGenerationResponse(
video_data=video_data,
seed=params.seed,
success=True,
generation_time=generation_time,
inference_time=inference_time,
encoding_time=encoding_time,
total_time=total_time,
stage_names=stage_names,
stage_execution_times=stage_execution_times,
)
@serve.deployment(
ray_actor_options={"num_cpus": 2, "num_gpus": 1, "runtime_env": {"conda": "fv"}},
)
class T2VModelDeployment(BaseModelDeployment):
def __init__(self, t2v_model_path: str, output_path: str = "outputs"):
super().__init__(t2v_model_path, output_path)
self._initialize_generator(MODEL_CONFIGS["1.3B"])
print("✅ T2V model initialized successfully")
@serve.deployment(
ray_actor_options={"num_cpus": 16, "num_gpus": 1, "runtime_env": {"conda": "fv"}},
)
class T2V14BModelDeployment(BaseModelDeployment):
def __init__(self, t2v_14b_model_path: str, output_path: str = "outputs"):
super().__init__(t2v_14b_model_path, output_path)
# Override environment for 14B model
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
self._initialize_generator(MODEL_CONFIGS["14B"])
print("✅ T2V 14B model initialized successfully")
app = FastAPI()
limiter = Limiter(key_func=get_remote_address)
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
@serve.deployment(num_replicas=50, ray_actor_options={"num_cpus": 2})
@serve.ingress(app)
class FastVideoAPI:
def __init__(self, t2v_deployments: Dict[str, DeploymentHandle]):
self.t2v_deployments = t2v_deployments
# Initialize Prometheus metrics
self.request_count = Counter('fastvideo_requests_total', 'Total FastVideo requests', ['model_type', 'status'])
self.request_duration = Histogram('fastvideo_request_duration_seconds', 'FastVideo request duration', ['model_type'])
self.video_generation_time = Histogram('fastvideo_video_generation_seconds', 'Video generation time', ['model_type'])
def _get_model_name(self, model_path: Optional[str]) -> str:
return model_path.split('/')[-1] if model_path else "unknown"
def _record_metrics(self, model_name: str, status: str, duration: float, response: Optional[VideoGenerationResponse] = None) -> None:
self.request_count.labels(model_type=model_name, status=status).inc()
self.request_duration.labels(model_type=model_name).observe(duration)
if response and hasattr(response, 'generation_time') and response.generation_time:
self.video_generation_time.labels(model_type=model_name).observe(response.generation_time)
@app.post("/generate_video", response_model=VideoGenerationResponse)
@limiter.limit("10/minute")
async def generate_video(self, request: Request, video_request: VideoGenerationRequest) -> VideoGenerationResponse:
"""Route the request to the appropriate model deployment based on model_path."""
start_time = time.time()
model_name = self._get_model_name(video_request.model_path)
try:
if video_request.model_path not in self.t2v_deployments:
raise ValueError(f"Model {video_request.model_path} not found")
response_ref = self.t2v_deployments[video_request.model_path].generate_video.remote(video_request)
response = await response_ref
self._record_metrics(model_name, "success", time.time() - start_time, response)
return response
except Exception as e:
self._record_metrics(model_name, "error", time.time() - start_time)
return VideoGenerationResponse(
video_data=None,
seed=video_request.seed,
success=False,
error_message=str(e),
generation_time=0,
inference_time=0,
encoding_time=0,
total_time=0,
)
@app.get("/health")
@limiter.limit("10/minute")
async def health_check(self, request: Request) -> Dict[str, str]:
return {"status": "healthy"}
@app.get("/metrics")
async def metrics(self) -> Response:
return Response(generate_latest(), media_type="text/plain")
def validate_configuration(model_paths: List[str], replicas: List[int]) -> None:
assert len(model_paths) == len(replicas), "Number of models and replicas must match"
assert sum(replicas) <= NUM_GPUS, f"Total replicas ({sum(replicas)}) must be <= {NUM_GPUS}"
for model, replica_count in zip(model_paths, replicas):
assert model in SUPPORTED_MODELS, f"Model {model} not supported"
assert replica_count > 0, f"Replicas must be greater than 0"
def start_ray_serve(
*,
t2v_model_paths: str,
t2v_model_replicas: str,
output_path: str = "outputs",
host: str = "0.0.0.0",
port: int = 8000,
) -> None:
if not ray.is_initialized():
ray.init()
model_paths = t2v_model_paths.split(",")
replicas = [int(r) for r in t2v_model_replicas.split(",")]
validate_configuration(model_paths, replicas)
t2v_deps = {}
for model_path, replica_count in zip(model_paths, replicas):
t2v_dep = T2VModelDeployment.options(num_replicas=replica_count).bind(model_path, output_path)
t2v_deps[model_path] = t2v_dep
api = FastVideoAPI.bind(t2v_deps)
serve.run(api, route_prefix="/", name="fast_video")
print(f"Ray Serve backend started at http://{host}:{port}")
for model_path, replica_count in zip(model_paths, replicas):
print(f"T2V Model: {model_path} | Replicas: {replica_count}")
print(f"Health check: http://{host}:{port}/health")
print(f"Video generation endpoint: http://{host}:{port}/generate_video")
def setup_signal_handlers() -> None:
signal.signal(signal.SIGINT, lambda *_: sys.exit(0))
signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="FastVideo Ray Serve Backend")
parser.add_argument("--t2v_model_paths",
type=str,
default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
help="Comma separated list of paths to the T2V model(s)")
parser.add_argument("--t2v_model_replicas",
type=str,
default="4,4",
help="Comma separated list of number of replicas for the T2V model(s)")
parser.add_argument("--output_path",
type=str,
default="outputs",
help="Path to save generated videos")
parser.add_argument("--host",
type=str,
default="0.0.0.0",
help="Host to bind to")
parser.add_argument("--port",
type=int,
default=8000,
help="Port to bind to")
args = parser.parse_args()
model_paths = args.t2v_model_paths.split(",")
replicas = [int(r) for r in args.t2v_model_replicas.split(",")]
validate_configuration(model_paths, replicas)
start_ray_serve(
t2v_model_paths=args.t2v_model_paths,
t2v_model_replicas=args.t2v_model_replicas,
output_path=args.output_path,
host=args.host,
port=args.port,
)
setup_signal_handlers()
print("✅ FastVideo backend is running. Press Ctrl-C to stop.")
while True:
time.sleep(3600)
+3
View File
@@ -0,0 +1,3 @@
python examples/inference/gradio/start_ray_serve_app.py \
--t2v_model_paths "FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers" \
--t2v_model_replicas "4,4"
@@ -0,0 +1,257 @@
"""
Startup script for FastVideo with Ray Serve backend and Gradio frontend.
This script starts both the backend and frontend services.
"""
import argparse
import os
import subprocess
import sys
import time
import threading
import signal
import requests
from pathlib import Path
from typing import Dict, Any, Optional
DEFAULT_BACKEND_HOST = "0.0.0.0"
DEFAULT_BACKEND_PORT = 8000
DEFAULT_FRONTEND_HOST = "0.0.0.0"
DEFAULT_FRONTEND_PORT = 7860
DEFAULT_OUTPUT_PATH = "outputs"
DEFAULT_T2V_MODELS = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
DEFAULT_T2V_REPLICAS = "4,4"
HEALTH_CHECK_TIMEOUT = 5
HEALTH_CHECK_MAX_RETRIES = 100
HEALTH_CHECK_INTERVAL = 2
PROCESS_SHUTDOWN_TIMEOUT = 5
PROCESS_MONITOR_INTERVAL = 1
PROJECT_ROOT = Path(__file__).parent.parent.parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
class ServiceManager:
def __init__(self, args: argparse.Namespace):
self.args = args
self.backend_process: Optional[subprocess.Popen] = None
self.frontend_process: Optional[subprocess.Popen] = None
self.backend_url = f"http://{args.backend_host}:{args.backend_port}"
def check_backend_health(self, max_retries: int = HEALTH_CHECK_MAX_RETRIES) -> bool:
health_url = f"{self.backend_url}/health"
for attempt in range(max_retries):
try:
response = requests.get(health_url, timeout=HEALTH_CHECK_TIMEOUT)
if response.status_code == 200:
print(f"✅ Backend is healthy at {self.backend_url}")
return True
except requests.exceptions.RequestException:
pass
if attempt < max_retries - 1:
print(f"⏳ Waiting for backend to start... ({attempt + 1}/{max_retries})")
time.sleep(HEALTH_CHECK_INTERVAL)
print(f"❌ Backend failed to start within {max_retries * HEALTH_CHECK_INTERVAL} seconds")
return False
def _create_monitor_thread(self, process: subprocess.Popen, service_name: str) -> threading.Thread:
def monitor():
if process.stdout:
for line in process.stdout:
print(f"[{service_name}] {line.rstrip()}")
thread = threading.Thread(target=monitor, daemon=True)
thread.start()
return thread
def _start_service(self, script_name: str, args_dict: Dict[str, Any], service_name: str) -> subprocess.Popen:
script_path = Path(__file__).parent / script_name
cmd = [sys.executable, str(script_path)]
for key, value in args_dict.items():
cmd.extend([f"--{key}", str(value)])
print(f"🚀 Starting {service_name}...")
print(f"Command: {' '.join(cmd)}")
process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
universal_newlines=True,
bufsize=1
)
self._create_monitor_thread(process, service_name.upper())
return process
def start_backend(self) -> subprocess.Popen:
backend_args = {
"t2v_model_paths": self.args.t2v_model_paths,
"t2v_model_replicas": self.args.t2v_model_replicas,
"output_path": self.args.output_path,
"host": self.args.backend_host,
"port": self.args.backend_port
}
self.backend_process = self._start_service("ray_serve_backend.py", backend_args, "backend")
return self.backend_process
def start_frontend(self) -> subprocess.Popen:
frontend_args = {
"backend_url": self.backend_url,
"t2v_model_paths": self.args.t2v_model_paths,
"host": self.args.frontend_host,
"port": self.args.frontend_port
}
self.frontend_process = self._start_service("gradio_frontend.py", frontend_args, "frontend")
return self.frontend_process
def shutdown_services(self) -> None:
print("\n🛑 Shutting down services...")
processes = []
if self.frontend_process:
self.frontend_process.terminate()
processes.append(("frontend", self.frontend_process))
if self.backend_process:
self.backend_process.terminate()
processes.append(("backend", self.backend_process))
for name, process in processes:
try:
process.wait(timeout=PROCESS_SHUTDOWN_TIMEOUT)
print(f"✅ {name.capitalize()} stopped gracefully")
except subprocess.TimeoutExpired:
print(f"⚠️ Force killing {name} process...")
process.kill()
print("✅ All services stopped")
def monitor_processes(self) -> None:
if not self.backend_process or not self.frontend_process:
print("❌ Processes not properly initialized")
return
try:
while True:
if self.frontend_process.poll() is not None:
print("❌ Frontend process died unexpectedly")
break
if self.backend_process.poll() is not None:
print("❌ Backend process died unexpectedly")
break
time.sleep(PROCESS_MONITOR_INTERVAL)
except KeyboardInterrupt:
pass
self.shutdown_services()
def setup_signal_handlers(service_manager: ServiceManager) -> None:
def signal_handler(signum: int, frame: Any) -> None:
service_manager.shutdown_services()
sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
def print_startup_info(args: argparse.Namespace) -> None:
print("🎬 FastVideo Ray Serve App")
print("=" * 50)
print(f"T2V Models: {args.t2v_model_paths}")
print(f"T2V Model Replicas: {args.t2v_model_replicas}")
print(f"Output: {args.output_path}")
print(f"Backend: http://{args.backend_host}:{args.backend_port}")
print(f"Frontend: http://{args.frontend_host}:{args.frontend_port}")
print("=" * 50)
def parse_arguments() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="FastVideo Ray Serve App")
parser.add_argument("--t2v_model_paths",
type=str,
default=DEFAULT_T2V_MODELS,
help="Comma separated list of paths to the T2V model(s)")
parser.add_argument("--t2v_model_replicas",
type=str,
default=DEFAULT_T2V_REPLICAS,
help="Comma separated list of number of replicas for the T2V model(s)")
parser.add_argument("--output_path",
type=str,
default=DEFAULT_OUTPUT_PATH,
help="Path to save generated videos")
parser.add_argument("--backend_host",
type=str,
default=DEFAULT_BACKEND_HOST,
help="Backend host to bind to")
parser.add_argument("--backend_port",
type=int,
default=DEFAULT_BACKEND_PORT,
help="Backend port to bind to")
parser.add_argument("--frontend_host",
type=str,
default=DEFAULT_FRONTEND_HOST,
help="Frontend host to bind to")
parser.add_argument("--frontend_port",
type=int,
default=DEFAULT_FRONTEND_PORT,
help="Frontend port to bind to")
parser.add_argument("--skip_backend_check",
action="store_true",
help="Skip backend health check")
return parser.parse_args()
def main() -> None:
args = parse_arguments()
os.makedirs(args.output_path, exist_ok=True)
print_startup_info(args)
service_manager = ServiceManager(args)
setup_signal_handlers(service_manager)
try:
service_manager.start_backend()
if not args.skip_backend_check:
if not service_manager.check_backend_health():
print("❌ Backend failed to start. Terminating...")
service_manager.shutdown_services()
sys.exit(1)
service_manager.start_frontend()
print("\n🎉 Both services are starting up!")
print(f"📺 Frontend will be available at: http://{args.frontend_host}:{args.frontend_port}")
print(f"🔧 Backend API will be available at: http://{args.backend_host}:{args.backend_port}")
print("\nPress Ctrl+C to stop both services...")
service_manager.monitor_processes()
except Exception as e:
print(f"❌ Unexpected error: {e}")
service_manager.shutdown_services()
sys.exit(1)
if __name__ == "__main__":
main()
+9 -2
View File
@@ -9,6 +9,7 @@ diffusion models.
import math
import os
import time
from copy import deepcopy
from typing import Any
import imageio
@@ -202,6 +203,8 @@ class VideoGenerator:
if sampling_param is None:
sampling_param = SamplingParam.from_pretrained(
fastvideo_args.model_path)
else:
sampling_param = deepcopy(sampling_param)
kwargs["prompt"] = prompt
sampling_param.update(kwargs)
@@ -275,6 +278,7 @@ class VideoGenerator:
width: {target_width}
video_length: {sampling_param.num_frames}
prompt: {prompt}
image_path: {sampling_param.image_path}
neg_prompt: {sampling_param.negative_prompt}
seed: {sampling_param.seed}
infer_steps: {sampling_param.num_inference_steps}
@@ -304,7 +308,8 @@ class VideoGenerator:
# Run inference
start_time = time.perf_counter()
output_batch = self.executor.execute_forward(batch, fastvideo_args)
samples = output_batch
samples = output_batch.output
logging_info = output_batch.logging_info
gen_time = time.perf_counter() - start_time
logger.info("Generated successfully in %.2f seconds", gen_time)
@@ -334,9 +339,11 @@ class VideoGenerator:
else:
return {
"samples": samples,
"frames": frames,
"prompts": prompt,
"size": (target_height, target_width, batch.num_frames),
"generation_time": gen_time
"generation_time": gen_time,
"logging_info": logging_info,
}
def set_lora_adapter(self,
@@ -17,10 +17,47 @@ import torch
if TYPE_CHECKING:
from torchcodec.decoders import VideoDecoder
import time
from collections import OrderedDict
from fastvideo.attention import AttentionMetadata
from fastvideo.configs.sample.teacache import TeaCacheParams, WanTeaCacheParams
class PipelineLoggingInfo:
"""Simple approach using OrderedDict to track stage metrics."""
def __init__(self):
# OrderedDict preserves insertion order and allows easy access
self.stages: OrderedDict[str, dict[str, Any]] = OrderedDict()
def add_stage_execution_time(self, stage_name: str, execution_time: float):
"""Add execution time for a stage."""
if stage_name not in self.stages:
self.stages[stage_name] = {}
self.stages[stage_name]['execution_time'] = execution_time
self.stages[stage_name]['timestamp'] = time.time()
def add_stage_metric(self, stage_name: str, metric_name: str, value: Any):
"""Add any metric for a stage."""
if stage_name not in self.stages:
self.stages[stage_name] = {}
self.stages[stage_name][metric_name] = value
def get_stage_info(self, stage_name: str) -> dict[str, Any]:
"""Get all info for a specific stage."""
return self.stages.get(stage_name, {})
def get_execution_order(self) -> list[str]:
"""Get stages in execution order."""
return list(self.stages.keys())
def get_total_execution_time(self) -> float:
"""Get total pipeline execution time."""
return sum(
stage.get('execution_time', 0) for stage in self.stages.values())
@dataclass
class ForwardBatch:
"""
@@ -132,6 +169,10 @@ class ForwardBatch:
# VSA parameters
VSA_sparsity: float = 0.0
# Logging info
logging_info: PipelineLoggingInfo = field(
default_factory=PipelineLoggingInfo)
def __post_init__(self):
"""Initialize dependent fields after dataclass initialization."""
+2
View File
@@ -155,6 +155,8 @@ class PipelineStage(ABC):
execution_time = time.perf_counter() - start_time
logger.info("[%s] Execution completed in %s ms", stage_name,
execution_time * 1000)
batch.logging_info.add_stage_execution_time(
stage_name, execution_time)
except Exception as e:
execution_time = time.perf_counter() - start_time
logger.error("[%s] Error during execution after %s ms: %s",
+8 -1
View File
@@ -11,6 +11,7 @@ from typing import Any, TextIO, cast
import psutil
import torch
import fastvideo.envs as envs
from fastvideo.distributed import (
cleanup_dist_env_and_memory,
maybe_init_distributed_environment_and_model_parallel)
@@ -139,7 +140,13 @@ class Worker:
fastvideo_args = recv_rpc['kwargs']['fastvideo_args']
output_batch = self.execute_forward(forward_batch,
fastvideo_args)
self.pipe.send({"output_batch": output_batch.output.cpu()})
logging_info = None
if envs.FASTVIDEO_STAGE_LOGGING:
logging_info = output_batch.logging_info
self.pipe.send({
"output_batch": output_batch.output.cpu(),
"logging_info": logging_info
})
elif method_name == 'set_lora_adapter':
lora_nickname = recv_rpc['kwargs']['lora_nickname']
lora_path = recv_rpc['kwargs']['lora_path']
+17 -3
View File
@@ -8,8 +8,9 @@ import socket
import time
from collections.abc import Callable
from multiprocessing.process import BaseProcess
from typing import Any, cast
from typing import Any
import fastvideo.envs as envs
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.logger import init_logger
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
@@ -40,7 +41,8 @@ class MultiprocExecutor(Executor):
logger.info("Using provided master port: %s", self.master_port)
else:
# Auto-find available port
for port in range(29503, 65535):
import random
for port in range(29503 + random.randint(0, 10000), 65535):
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
if s.connect_ex(('localhost', port)) != 0:
self.master_port = port
@@ -80,7 +82,19 @@ class MultiprocExecutor(Executor):
"forward_batch": forward_batch,
"fastvideo_args": fastvideo_args
})
return cast(ForwardBatch, responses[0]["output_batch"])
output = responses[0]["output_batch"]
logging_info = None
if envs.FASTVIDEO_STAGE_LOGGING:
logging_info = responses[0]["logging_info"]
else:
logging_info = None
result_batch = ForwardBatch(data_type=forward_batch.data_type,
output=output,
logging_info=logging_info)
return result_batch
def set_lora_adapter(self,
lora_nickname: str,