From 8cae5e470818c486da06862185e05ab60edc9a01 Mon Sep 17 00:00:00 2001 From: William Lin Date: Sun, 10 Aug 2025 15:34:03 -0700 Subject: [PATCH] [feature] add Gradio live serving demo code (#727) Co-authored-by: RandNMR73 --- examples/inference/gradio/README.md | 59 -- examples/inference/gradio/gradio_demo.py | 169 ---- examples/inference/gradio/gradio_frontend.py | 730 ++++++++++++++++++ .../inference/gradio/ray_serve_backend.py | 379 +++++++++ examples/inference/gradio/start.sh | 3 + .../inference/gradio/start_ray_serve_app.py | 257 ++++++ fastvideo/entrypoints/video_generator.py | 11 +- fastvideo/pipelines/pipeline_batch_info.py | 41 + fastvideo/pipelines/stages/base.py | 2 + fastvideo/worker/gpu_worker.py | 9 +- fastvideo/worker/multiproc_executor.py | 20 +- 11 files changed, 1446 insertions(+), 234 deletions(-) delete mode 100644 examples/inference/gradio/README.md delete mode 100644 examples/inference/gradio/gradio_demo.py create mode 100644 examples/inference/gradio/gradio_frontend.py create mode 100644 examples/inference/gradio/ray_serve_backend.py create mode 100644 examples/inference/gradio/start.sh create mode 100644 examples/inference/gradio/start_ray_serve_app.py diff --git a/examples/inference/gradio/README.md b/examples/inference/gradio/README.md deleted file mode 100644 index c7ff9f5fd..000000000 --- a/examples/inference/gradio/README.md +++ /dev/null @@ -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 diff --git a/examples/inference/gradio/gradio_demo.py b/examples/inference/gradio/gradio_demo.py deleted file mode 100644 index 06340bae6..000000000 --- a/examples/inference/gradio/gradio_demo.py +++ /dev/null @@ -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) diff --git a/examples/inference/gradio/gradio_frontend.py b/examples/inference/gradio/gradio_frontend.py new file mode 100644 index 000000000..fab308800 --- /dev/null +++ b/examples/inference/gradio/gradio_frontend.py @@ -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""" +
+

⏱️ Timing Breakdown

+
+
+
πŸš€
+
DiT Denoising
+
{dit_denoising_time}
+
+
+
🧠
+
E2E (w. vae/text encoder)
+
{inference_time:.2f}s
+
+
+
🎬
+
Video Encoding
+
{encoding_time:.2f}s
+
+
+
🌐
+
Network Transfer
+
{network_time:.2f}s
+
+
+
πŸ“Š
+
Total Processing
+
{total_time:.2f}s
+
+
""" + + if inference_time > 0: + fps = num_frames / inference_time + timing_html += f""" +
+ Generation Speed: + {fps:.1f} frames/second +
""" + + return timing_html + "
" + + +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(""" +
+

Make Video Generation Go Blurrrrrrr

+

Code | Blog | Docs

+
+ """) + + with gr.Accordion("πŸŽ₯ What Is FastVideo?", open=False): + gr.HTML(""" +
+

+ 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. +

+
+ """) + + 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("
Advanced Options
") + 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(""" + + """) + + 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(""" +
+

The compute for this demo is generously provided by GMI Cloud. 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.

+
+ """) + + 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""" + + + + + + + FastWan + + + + + + + + + + + + + + + + + + + + + + + + + + + + """ + + 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() \ No newline at end of file diff --git a/examples/inference/gradio/ray_serve_backend.py b/examples/inference/gradio/ray_serve_backend.py new file mode 100644 index 000000000..67192fd80 --- /dev/null +++ b/examples/inference/gradio/ray_serve_backend.py @@ -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) \ No newline at end of file diff --git a/examples/inference/gradio/start.sh b/examples/inference/gradio/start.sh new file mode 100644 index 000000000..6eabdd9ce --- /dev/null +++ b/examples/inference/gradio/start.sh @@ -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" \ No newline at end of file diff --git a/examples/inference/gradio/start_ray_serve_app.py b/examples/inference/gradio/start_ray_serve_app.py new file mode 100644 index 000000000..437121132 --- /dev/null +++ b/examples/inference/gradio/start_ray_serve_app.py @@ -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() \ No newline at end of file diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 2e6f63696..a24ea4554 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -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, diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index ba79b95d0..eb89f62f5 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -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.""" diff --git a/fastvideo/pipelines/stages/base.py b/fastvideo/pipelines/stages/base.py index 10b4f690d..97d818044 100644 --- a/fastvideo/pipelines/stages/base.py +++ b/fastvideo/pipelines/stages/base.py @@ -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", diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index f3eebcbb7..6f89d1772 100644 --- a/fastvideo/worker/gpu_worker.py +++ b/fastvideo/worker/gpu_worker.py @@ -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'] diff --git a/fastvideo/worker/multiproc_executor.py b/fastvideo/worker/multiproc_executor.py index 1d770d322..b3b1398df 100644 --- a/fastvideo/worker/multiproc_executor.py +++ b/fastvideo/worker/multiproc_executor.py @@ -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,