Add pipeline config
This commit is contained in:
JerryZhou54
@@ -1,3 +1 @@
|
||||
from fastvideo.v1.entrypoints.video_generator import VideoGenerator
|
||||
|
||||
__all__ = ["VideoGenerator"]
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
from fastvideo.v1.configs.hunyuan import HunyuanConfig, FastHunyuanConfig
|
||||
from fastvideo.v1.configs.wan import WanT2V480PConfig, WanI2V480PConfig
|
||||
from fastvideo.v1.configs.base import BaseConfig, SlidingTileAttnConfig
|
||||
from fastvideo.v1.configs.config import FastVideoConfig
|
||||
|
||||
__all__ = [
|
||||
"HunyuanConfig", "FastHunyuanConfig", "BaseConfig", "SlidingTileAttnConfig",
|
||||
"FastVideoConfig"
|
||||
"WanT2V480PConfig", "WanI2V480PConfig"
|
||||
]
|
||||
|
||||
@@ -12,19 +12,41 @@ class BaseConfig:
|
||||
num_frames: int = 125
|
||||
fps: int = 24
|
||||
|
||||
# Inference parameters
|
||||
num_inference_steps: int = 6
|
||||
# Video generation parameters
|
||||
num_inference_steps: int = 50
|
||||
guidance_scale: float = 1.0
|
||||
seed: int = 1024
|
||||
guidance_rescale: float = 0.0
|
||||
embedded_cfg_scale: float = 6.0
|
||||
flow_shift: Optional[float] = None
|
||||
use_cpu_offload: bool = False
|
||||
disable_autocast: bool = False
|
||||
|
||||
# Model configuration
|
||||
num_gpus: int = 1
|
||||
sp_size: Optional[int] = None # defaults to num_gpus
|
||||
tp_size: Optional[int] = None # defaults to num_gpus
|
||||
vae_sp: bool = True
|
||||
precision: str = "bf16"
|
||||
|
||||
# Output configuration
|
||||
output_path: str = "output_videos/"
|
||||
# VAE configuration
|
||||
vae_precision: str = "fp16"
|
||||
vae_tiling: bool = True
|
||||
vae_sp: bool = True
|
||||
vae_scale_factor: Optional[int] = None
|
||||
|
||||
# DiT configuration
|
||||
num_channels_latents: Optional[int] = None
|
||||
|
||||
# Image encoder configuration
|
||||
image_encoder_precision: str = "fp32"
|
||||
|
||||
# Text encoder configuration
|
||||
text_encoder_precision: str = "fp16"
|
||||
text_len: int = -1
|
||||
hidden_state_skip_layer: int = 0
|
||||
|
||||
# STA (Spatial-Temporal Attention) parameters
|
||||
mask_strategy_file_path: Optional[str] = None
|
||||
enable_torch_compile: bool = False
|
||||
|
||||
neg_prompt: Optional[str] = None
|
||||
|
||||
# Additional parameters can be added as a dict
|
||||
extra_params: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@@ -7,12 +7,25 @@ class HunyuanConfig(BaseConfig):
|
||||
"""Base configuration for HunYuan pipeline architecture."""
|
||||
|
||||
# HunyuanConfig-specific parameters with defaults
|
||||
# Denoising stage
|
||||
embedded_cfg_scale: int = 6
|
||||
flow_shift: int = 7
|
||||
|
||||
# Override some BaseConfig defaults
|
||||
num_inference_steps: int = 50
|
||||
|
||||
# Text encoding stage
|
||||
hidden_state_skip_layer: int = 2
|
||||
text_len: int = 256
|
||||
|
||||
# Precision for each component
|
||||
precision: str = "bf16"
|
||||
vae_precision: str = "fp16"
|
||||
text_encoder_precision: str = "fp16"
|
||||
|
||||
# HunyuanConfig-specific added parameters
|
||||
# Secondary text encoder
|
||||
text_encoder_precision_2: str = "fp16"
|
||||
text_len_2: int = 77
|
||||
|
||||
|
||||
@dataclass
|
||||
class FastHunyuanConfig(HunyuanConfig):
|
||||
|
||||
@@ -5,7 +5,7 @@ from typing import Dict, Type, Optional, Callable
|
||||
|
||||
from fastvideo.v1.configs.base import BaseConfig
|
||||
from fastvideo.v1.configs.hunyuan import HunyuanConfig, FastHunyuanConfig
|
||||
from fastvideo.v1.configs.wan import WanConfig
|
||||
from fastvideo.v1.configs.wan import WanT2V480PConfig, WanI2V480PConfig
|
||||
|
||||
from fastvideo.v1.utils import maybe_download_model_index, verify_model_config_and_directory
|
||||
from fastvideo.v1.logger import init_logger
|
||||
@@ -16,13 +16,16 @@ logger = init_logger(__name__)
|
||||
WEIGHT_CONFIG_REGISTRY: Dict[str, Type[BaseConfig]] = {
|
||||
"FastVideo/FastHunyuan-Diffusers": FastHunyuanConfig,
|
||||
"hunyuanvideo-community/HunyuanVideo": HunyuanConfig,
|
||||
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig,
|
||||
"Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig
|
||||
# Add other specific weight variants
|
||||
}
|
||||
|
||||
# For determining pipeline type from model ID
|
||||
PIPELINE_DETECTOR: Dict[str, Callable[[str], bool]] = {
|
||||
"hunyuan": lambda id: "hunyuan" in id.lower(),
|
||||
"wan": lambda id: "wan" in id.lower(),
|
||||
"wanpipeline": lambda id: "wanpipeline" in id.lower(),
|
||||
"wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
|
||||
# Add other pipeline architecture detectors
|
||||
}
|
||||
|
||||
@@ -30,7 +33,9 @@ PIPELINE_DETECTOR: Dict[str, Callable[[str], bool]] = {
|
||||
PIPELINE_FALLBACK_CONFIG: Dict[str, Type[BaseConfig]] = {
|
||||
"hunyuan":
|
||||
HunyuanConfig, # Base Hunyuan config as fallback for any Hunyuan variant
|
||||
"wan": WanConfig, # Base Wan config as fallback for any Wan variant
|
||||
"wanpipeline":
|
||||
WanT2V480PConfig, # Base Wan config as fallback for any Wan variant
|
||||
"wanimagetovideo": WanI2V480PConfig,
|
||||
# Other fallbacks by architecture
|
||||
}
|
||||
|
||||
@@ -60,13 +65,13 @@ def get_pipeline_config_for_name(
|
||||
|
||||
# If no match, try to use the fallback config
|
||||
fallback_config = None
|
||||
print(pipeline_name)
|
||||
# Try to determine pipeline architecture for fallback
|
||||
for pipeline_type, detector in PIPELINE_DETECTOR.items():
|
||||
if detector(pipeline_name.lower()):
|
||||
fallback_config = PIPELINE_FALLBACK_CONFIG.get(pipeline_type)
|
||||
break
|
||||
|
||||
logger.warning(
|
||||
f"No match found for pipeline {pipeline_name_or_path}, using fallback config {fallback_config}."
|
||||
)
|
||||
logger.warning("No match found for pipeline %s, using fallback config %s.",
|
||||
pipeline_name_or_path, fallback_config)
|
||||
return fallback_config
|
||||
|
||||
@@ -1,5 +1,44 @@
|
||||
from dataclasses import dataclass
|
||||
from fastvideo.v1.configs.base import BaseConfig
|
||||
|
||||
|
||||
class WanConfig(BaseConfig):
|
||||
pass
|
||||
@dataclass
|
||||
class WanT2V480PConfig(BaseConfig):
|
||||
"""Base configuration for Wan T2V 1.3B pipeline architecture."""
|
||||
|
||||
# WanConfig-specific parameters with defaults
|
||||
# Video parameters
|
||||
height: int = 480
|
||||
width: int = 832
|
||||
num_frames: int = 81
|
||||
fps: int = 16
|
||||
use_cpu_offload: bool = True
|
||||
|
||||
# Denoising stage
|
||||
guidance_scale: float = 3.0
|
||||
neg_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
|
||||
flow_shift: int = 3
|
||||
num_inference_steps: int = 50
|
||||
|
||||
# Text encoding stage
|
||||
text_len: int = 512
|
||||
|
||||
# Precision for each component
|
||||
precision: str = "bf16"
|
||||
vae_precision: str = "fp16"
|
||||
text_encoder_precision: str = "fp32"
|
||||
|
||||
# WanConfig-specific added parameters
|
||||
|
||||
|
||||
@dataclass
|
||||
class WanI2V480PConfig(WanT2V480PConfig):
|
||||
"""Base configuration for Wan I2V 14B 480P pipeline architecture."""
|
||||
|
||||
# WanConfig-specific parameters with defaults
|
||||
# Denoising stage
|
||||
guidance_scale: float = 5.0
|
||||
num_inference_steps: int = 40
|
||||
|
||||
# Precision for each component
|
||||
image_encoder_precision: str = "fp32"
|
||||
|
||||
@@ -1,405 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
VideoGenerator module for FastVideo.
|
||||
|
||||
This module provides a consolidated interface for generating videos using
|
||||
diffusion models.
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
import imageio
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision
|
||||
from einops import rearrange
|
||||
|
||||
from fastvideo.v1.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.v1.logger import init_logger
|
||||
from fastvideo.v1.pipelines import (ComposedPipelineBase, ForwardBatch)
|
||||
|
||||
from fastvideo.v1.utils import align_to
|
||||
from fastvideo.v1.executor.abstract import Executor
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# def initialize_distributed_and_parallelism(fastvideo_args: FastVideoArgs):
|
||||
# """Initialize distributed environment and model parallelism."""
|
||||
# local_rank = int(os.environ.get("LOCAL_RANK", 0))
|
||||
# rank = int(os.environ.get("RANK", 0))
|
||||
# world_size = int(os.environ.get("WORLD_SIZE", 1))
|
||||
# torch.cuda.set_device(local_rank)
|
||||
# init_distributed_environment(world_size=world_size,
|
||||
# rank=rank,
|
||||
# local_rank=local_rank)
|
||||
# device_str = f"cuda:{local_rank}"
|
||||
# fastvideo_args.device_str = device_str
|
||||
# fastvideo_args.device = torch.device(device_str)
|
||||
# initialize_model_parallel(
|
||||
# sequence_model_parallel_size=fastvideo_args.sp_size,
|
||||
# tensor_model_parallel_size=fastvideo_args.tp_size,
|
||||
# )
|
||||
|
||||
|
||||
class VideoGenerator:
|
||||
"""
|
||||
A unified class for generating videos using diffusion models.
|
||||
|
||||
This class provides a simple interface for video generation with rich
|
||||
customization options, similar to popular frameworks like HF Diffusers.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
executor_class: type[Executor],
|
||||
log_stats: bool):
|
||||
"""
|
||||
Initialize the video generator.
|
||||
|
||||
Args:
|
||||
pipeline: The pipeline to use for inference
|
||||
fastvideo_args: The inference arguments
|
||||
"""
|
||||
self.fastvideo_args = fastvideo_args
|
||||
self.executor = executor_class(fastvideo_args)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls,
|
||||
model_path: str,
|
||||
device: Optional[str] = None,
|
||||
torch_dtype: Optional[torch.dtype] = None,
|
||||
**kwargs) -> "VideoGenerator":
|
||||
"""
|
||||
Create a video generator from a pretrained model.
|
||||
|
||||
Args:
|
||||
model_path: Path or identifier for the pretrained model
|
||||
device: Device to load the model on (e.g., "cuda", "cuda:0", "cpu")
|
||||
torch_dtype: Data type for model weights (e.g., torch.float16)
|
||||
**kwargs: Additional arguments to customize model loading
|
||||
|
||||
Returns:
|
||||
The created video generator
|
||||
"""
|
||||
fastvideo_args = FastVideoArgs(
|
||||
model_path=model_path,
|
||||
device_str=device or "cuda" if torch.cuda.is_available() else "cpu",
|
||||
**kwargs)
|
||||
|
||||
if torch_dtype is not None:
|
||||
fastvideo_args.dtype = torch_dtype
|
||||
|
||||
return cls.from_fastvideo_args(fastvideo_args)
|
||||
|
||||
@classmethod
|
||||
def from_fastvideo_args(cls,
|
||||
fastvideo_args: FastVideoArgs) -> "VideoGenerator":
|
||||
"""
|
||||
Create a video generator with the specified arguments.
|
||||
|
||||
Args:
|
||||
fastvideo_args: The inference arguments
|
||||
|
||||
Returns:
|
||||
The created video generator
|
||||
"""
|
||||
# Initialize distributed environment if needed
|
||||
# initialize_distributed_and_parallelism(fastvideo_args)
|
||||
|
||||
executor_class = Executor.get_class(fastvideo_args)
|
||||
|
||||
return cls(
|
||||
fastvideo_args=fastvideo_args,
|
||||
executor_class=executor_class,
|
||||
log_stats=False, # TODO: implement
|
||||
)
|
||||
|
||||
def generate_video(
|
||||
self,
|
||||
prompt: str,
|
||||
negative_prompt: Optional[str] = None,
|
||||
output_path: Optional[str] = None,
|
||||
save_video: bool = True,
|
||||
return_frames: bool = False,
|
||||
num_inference_steps: Optional[int] = None,
|
||||
guidance_scale: Optional[float] = None,
|
||||
num_frames: Optional[int] = None,
|
||||
height: Optional[int] = None,
|
||||
width: Optional[int] = None,
|
||||
fps: Optional[int] = None,
|
||||
seed: Optional[int] = None,
|
||||
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
|
||||
callback_steps: int = 1,
|
||||
) -> Union[Dict[str, Any], List[np.ndarray]]:
|
||||
"""
|
||||
Generate a video based on the given prompt.
|
||||
|
||||
Args:
|
||||
prompt: The prompt to use for generation
|
||||
negative_prompt: The negative prompt to use (overrides the one in fastvideo_args)
|
||||
output_path: Path to save the video (overrides the one in fastvideo_args)
|
||||
save_video: Whether to save the video to disk
|
||||
return_frames: Whether to return the raw frames
|
||||
num_inference_steps: Number of denoising steps (overrides fastvideo_args)
|
||||
guidance_scale: Classifier-free guidance scale (overrides fastvideo_args)
|
||||
num_frames: Number of frames to generate (overrides fastvideo_args)
|
||||
height: Height of generated video (overrides fastvideo_args)
|
||||
width: Width of generated video (overrides fastvideo_args)
|
||||
fps: Frames per second for saved video (overrides fastvideo_args)
|
||||
seed: Random seed for generation (overrides fastvideo_args)
|
||||
callback: Callback function called after each step
|
||||
callback_steps: Number of steps between each callback
|
||||
|
||||
Returns:
|
||||
Either the output dictionary or the list of frames depending on return_frames
|
||||
"""
|
||||
# Create a copy of inference args to avoid modifying the original
|
||||
fastvideo_args = self.fastvideo_args.copy()
|
||||
|
||||
# Override parameters if provided
|
||||
if negative_prompt is not None:
|
||||
fastvideo_args.neg_prompt = negative_prompt
|
||||
if num_inference_steps is not None:
|
||||
fastvideo_args.num_inference_steps = num_inference_steps
|
||||
if guidance_scale is not None:
|
||||
fastvideo_args.guidance_scale = guidance_scale
|
||||
if num_frames is not None:
|
||||
fastvideo_args.num_frames = num_frames
|
||||
if height is not None:
|
||||
fastvideo_args.height = height
|
||||
if width is not None:
|
||||
fastvideo_args.width = width
|
||||
if fps is not None:
|
||||
fastvideo_args.fps = fps
|
||||
if seed is not None:
|
||||
fastvideo_args.seed = seed
|
||||
|
||||
# Store callback info
|
||||
fastvideo_args.callback = callback
|
||||
fastvideo_args.callback_steps = callback_steps
|
||||
|
||||
# Validate inputs
|
||||
if not isinstance(prompt, str):
|
||||
raise TypeError(
|
||||
f"`prompt` must be a string, but got {type(prompt)}")
|
||||
prompt = prompt.strip()
|
||||
|
||||
# Process negative prompt
|
||||
if fastvideo_args.neg_prompt is not None:
|
||||
fastvideo_args.neg_prompt = fastvideo_args.neg_prompt.strip()
|
||||
|
||||
# Validate dimensions
|
||||
if (fastvideo_args.height <= 0 or fastvideo_args.width <= 0
|
||||
or fastvideo_args.num_frames <= 0):
|
||||
raise ValueError(
|
||||
f"Height, width, and num_frames must be positive integers, got "
|
||||
f"height={fastvideo_args.height}, width={fastvideo_args.width}, "
|
||||
f"num_frames={fastvideo_args.num_frames}")
|
||||
|
||||
if (fastvideo_args.num_frames - 1) % 4 != 0:
|
||||
raise ValueError(
|
||||
f"num_frames-1 must be a multiple of 4, got {fastvideo_args.num_frames}"
|
||||
)
|
||||
|
||||
# Calculate sizes
|
||||
target_height = align_to(fastvideo_args.height, 16)
|
||||
target_width = align_to(fastvideo_args.width, 16)
|
||||
|
||||
# Calculate latent sizes
|
||||
latents_size = [(fastvideo_args.num_frames - 1) // 4 + 1,
|
||||
fastvideo_args.height // 8, fastvideo_args.width // 8]
|
||||
n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
|
||||
|
||||
# Log parameters
|
||||
debug_str = f"""
|
||||
height: {target_height}
|
||||
width: {target_width}
|
||||
video_length: {fastvideo_args.num_frames}
|
||||
prompt: {prompt}
|
||||
neg_prompt: {fastvideo_args.neg_prompt}
|
||||
seed: {fastvideo_args.seed}
|
||||
infer_steps: {fastvideo_args.num_inference_steps}
|
||||
num_videos_per_prompt: {fastvideo_args.num_videos}
|
||||
guidance_scale: {fastvideo_args.guidance_scale}
|
||||
n_tokens: {n_tokens}
|
||||
flow_shift: {fastvideo_args.flow_shift}
|
||||
embedded_guidance_scale: {fastvideo_args.embedded_cfg_scale}"""
|
||||
logger.info(debug_str)
|
||||
|
||||
# Prepare batch
|
||||
device = torch.device(fastvideo_args.device_str)
|
||||
batch = ForwardBatch(
|
||||
prompt=prompt,
|
||||
negative_prompt=fastvideo_args.neg_prompt,
|
||||
num_videos_per_prompt=fastvideo_args.num_videos,
|
||||
height=fastvideo_args.height,
|
||||
width=fastvideo_args.width,
|
||||
num_frames=fastvideo_args.num_frames,
|
||||
num_inference_steps=fastvideo_args.num_inference_steps,
|
||||
guidance_scale=fastvideo_args.guidance_scale,
|
||||
eta=0.0,
|
||||
n_tokens=n_tokens,
|
||||
data_type="video" if fastvideo_args.num_frames > 1 else "image",
|
||||
device=device,
|
||||
extra={},
|
||||
)
|
||||
|
||||
# Run inference
|
||||
start_time = time.time()
|
||||
samples = self.pipeline.forward(
|
||||
batch=batch,
|
||||
fastvideo_args=fastvideo_args,
|
||||
).output
|
||||
|
||||
gen_time = time.time() - start_time
|
||||
logger.info(f"Generated successfully in {gen_time:.2f} seconds")
|
||||
|
||||
# Process outputs
|
||||
videos = rearrange(samples, "b c t h w -> t b c h w")
|
||||
frames = []
|
||||
for x in videos:
|
||||
x = torchvision.utils.make_grid(x, nrow=6)
|
||||
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
|
||||
frames.append((x * 255).numpy().astype(np.uint8))
|
||||
|
||||
# Save video if requested
|
||||
if save_video:
|
||||
save_path = output_path or fastvideo_args.output_path
|
||||
if save_path:
|
||||
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
||||
video_path = os.path.join(save_path, f"{prompt[:100]}.mp4")
|
||||
imageio.mimsave(video_path, frames, fps=fastvideo_args.fps)
|
||||
logger.info(f"Saved video to {video_path}")
|
||||
else:
|
||||
logger.warning("No output path provided, video not saved")
|
||||
|
||||
if return_frames:
|
||||
return frames
|
||||
else:
|
||||
return {
|
||||
"samples": samples,
|
||||
"prompts": prompt,
|
||||
"size":
|
||||
(target_height, target_width, fastvideo_args.num_frames),
|
||||
"generation_time": gen_time
|
||||
}
|
||||
|
||||
def batch_generate(self,
|
||||
prompts: List[str],
|
||||
output_path: Optional[str] = None,
|
||||
**kwargs) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Generate videos for a batch of prompts.
|
||||
|
||||
Args:
|
||||
prompts: List of prompts to generate videos for
|
||||
output_path: Path to save the videos (overrides the one in fastvideo_args)
|
||||
**kwargs: Additional parameters to pass to generate_video
|
||||
|
||||
Returns:
|
||||
List of output dictionaries from each generation
|
||||
"""
|
||||
if output_path:
|
||||
self.fastvideo_args.output_path = output_path
|
||||
|
||||
results = []
|
||||
for prompt in prompts:
|
||||
result = self.generate_video(prompt=prompt, **kwargs)
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
def image_to_video(self,
|
||||
image: Union[str, torch.Tensor, np.ndarray],
|
||||
prompt: Optional[str] = None,
|
||||
strength: float = 0.8,
|
||||
**kwargs) -> Dict[str, Any]:
|
||||
"""
|
||||
Generate a video from an initial image.
|
||||
|
||||
Args:
|
||||
image: Input image (path, tensor, or numpy array)
|
||||
prompt: Text prompt to guide the generation
|
||||
strength: How much to transform the original image (0-1)
|
||||
**kwargs: Additional parameters to pass to generate_video
|
||||
|
||||
Returns:
|
||||
Output dictionary from the generation
|
||||
"""
|
||||
# Load image if path is provided
|
||||
if isinstance(image, str):
|
||||
# Implementation would depend on your image loading utilities
|
||||
# This is a placeholder for the concept
|
||||
image_tensor = self._load_image(image)
|
||||
elif isinstance(image, np.ndarray):
|
||||
# Convert numpy array to tensor
|
||||
image_tensor = torch.from_numpy(image).permute(2, 0, 1) / 255.0
|
||||
else:
|
||||
image_tensor = image
|
||||
|
||||
# Add image to inference args
|
||||
fastvideo_args = self.fastvideo_args.copy()
|
||||
fastvideo_args.init_image = image_tensor
|
||||
fastvideo_args.strength = strength
|
||||
|
||||
# Generate video
|
||||
return self.generate_video(prompt=prompt or "",
|
||||
fastvideo_args=fastvideo_args,
|
||||
**kwargs)
|
||||
|
||||
def to(self, device: Union[str, torch.device]) -> "VideoGenerator":
|
||||
"""
|
||||
Move the model to the specified device.
|
||||
|
||||
Args:
|
||||
device: The device to move the model to
|
||||
|
||||
Returns:
|
||||
Self for chaining
|
||||
"""
|
||||
device_str = str(device)
|
||||
self.fastvideo_args.device_str = device_str
|
||||
self.fastvideo_args.device = torch.device(device_str)
|
||||
|
||||
# Move pipeline components to device
|
||||
self.pipeline.to(device)
|
||||
|
||||
return self
|
||||
|
||||
def _load_image(self, image_path: str) -> torch.Tensor:
|
||||
"""
|
||||
Load an image from a path and convert to tensor.
|
||||
|
||||
Args:
|
||||
image_path: Path to the image
|
||||
|
||||
Returns:
|
||||
Tensor representation of the image
|
||||
"""
|
||||
# Placeholder implementation - would need to be implemented
|
||||
# based on your image loading utilities
|
||||
import PIL.Image
|
||||
from torchvision import transforms
|
||||
|
||||
image = PIL.Image.open(image_path).convert("RGB")
|
||||
transform = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
|
||||
])
|
||||
return transform(image).unsqueeze(0)
|
||||
|
||||
|
||||
def load_prompts_from_file(prompt_path: str) -> List[str]:
|
||||
"""
|
||||
Load prompts from a file.
|
||||
|
||||
Args:
|
||||
prompt_path: Path to the file containing prompts
|
||||
|
||||
Returns:
|
||||
List of prompts
|
||||
"""
|
||||
with open(prompt_path) as f:
|
||||
return [line.strip() for line in f.readlines()]
|
||||
@@ -1,15 +0,0 @@
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
# This will automatically handle distributed setup if num_gpus > 1
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/FastHunyuan-Diffusers",
|
||||
num_gpus=2,
|
||||
distributed_executor_backend="mp",
|
||||
)
|
||||
|
||||
# Generate videos with the same simple API, regardless of GPU count
|
||||
prompt = "A beautiful woman in a red dress walking down a street"
|
||||
video = generator.generate_video(prompt)
|
||||
|
||||
prompt2 = "A beautiful woman in a blue dress walking down a street"
|
||||
video2 = generator.generate_video(prompt2)
|
||||
@@ -1,251 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from typing import Union
|
||||
|
||||
from fastvideo.v1.fastvideo_args import FastVideoArgs
|
||||
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import (Any, Callable, Dict, List, Optional, Tuple)
|
||||
|
||||
import torch.nn as nn
|
||||
from typing_extensions import TypeVar
|
||||
|
||||
from fastvideo.v1.logger import init_logger
|
||||
from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
|
||||
# from vllm.lora.request import LoRARequest
|
||||
# from vllm.model_executor.layers.sampler import SamplerOutput
|
||||
from fastvideo.v1.worker.worker_base import WorkerBase
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_R = TypeVar("_R", default=Any)
|
||||
|
||||
|
||||
class Executor(ABC):
|
||||
"""Base class for all executors.
|
||||
|
||||
An executor is responsible for executing the model on one device,
|
||||
or it can be a distributed executor
|
||||
that can execute the model on multiple devices.
|
||||
"""
|
||||
|
||||
uses_ray: bool # whether the executor uses Ray for orchestration.
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
) -> None:
|
||||
self.fastvideo_args = fastvideo_args
|
||||
self._init_executor()
|
||||
self.is_sleeping = False
|
||||
|
||||
@abstractmethod
|
||||
def _init_executor(self) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def get_class(fastvideo_args: FastVideoArgs) -> type["Executor"]:
|
||||
executor_class: type[Executor]
|
||||
distributed_executor_backend = (
|
||||
fastvideo_args.distributed_executor_backend)
|
||||
# distributed_executor_backend must be set in VllmConfig.__post_init__
|
||||
if isinstance(distributed_executor_backend, type):
|
||||
if not issubclass(distributed_executor_backend, Executor):
|
||||
raise TypeError(
|
||||
"distributed_executor_backend must be a subclass of "
|
||||
f"Executor. Got {distributed_executor_backend}.")
|
||||
executor_class = distributed_executor_backend
|
||||
elif distributed_executor_backend == "ray":
|
||||
raise NotImplementedError
|
||||
from fastvideo.v1.executor.ray_distributed_executor import ( # noqa
|
||||
RayDistributedExecutor)
|
||||
executor_class = RayDistributedExecutor
|
||||
elif distributed_executor_backend == "mp":
|
||||
raise NotImplementedError
|
||||
from fastvideo.v1.executor.multiproc_executor import MultiprocExecutor
|
||||
executor_class = MultiprocExecutor
|
||||
elif distributed_executor_backend == "torch":
|
||||
from fastvideo.v1.executor.torchrun_executor import TorchRunExecutor
|
||||
executor_class = TorchRunExecutor
|
||||
else:
|
||||
raise ValueError("Unknown distributed executor backend: "
|
||||
f"{distributed_executor_backend}")
|
||||
return executor_class
|
||||
|
||||
@abstractmethod
|
||||
def collective_rpc(self,
|
||||
method: Union[str, Callable[..., _R]],
|
||||
timeout: Optional[float] = None,
|
||||
args: Tuple = (),
|
||||
kwargs: Optional[Dict[str, Any]] = None) -> List[_R]:
|
||||
"""
|
||||
Execute an RPC call on all workers.
|
||||
|
||||
Args:
|
||||
method: Name of the worker method to execute, or a callable that
|
||||
is serialized and sent to all workers to execute.
|
||||
|
||||
If the method is a callable, it should accept an additional
|
||||
`self` argument, in addition to the arguments passed in `args`
|
||||
and `kwargs`. The `self` argument will be the worker object.
|
||||
timeout: Maximum time in seconds to wait for execution. Raises a
|
||||
:exc:`TimeoutError` on timeout. `None` means wait indefinitely.
|
||||
args: Positional arguments to pass to the worker method.
|
||||
kwargs: Keyword arguments to pass to the worker method.
|
||||
|
||||
Returns:
|
||||
A list containing the results from each worker.
|
||||
|
||||
Note:
|
||||
It is recommended to use this API to only pass control messages,
|
||||
and set up data-plane communication to pass data.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
|
||||
"""
|
||||
Run a function directly on the model inside each worker,
|
||||
returning the result for each of them.
|
||||
"""
|
||||
|
||||
def rpc_func(worker: WorkerBase) -> _R:
|
||||
return func(worker.get_model())
|
||||
|
||||
return self.collective_rpc(rpc_func)
|
||||
|
||||
def execute_model(self,
|
||||
forward_batch: ForwardBatch) -> Optional[ForwardBatch]:
|
||||
output = self.collective_rpc("execute_model", args=(forward_batch, ))
|
||||
return output[0]
|
||||
|
||||
def stop_remote_worker_execution_loop(self) -> None:
|
||||
"""Releases parallel workers from model loop."""
|
||||
return
|
||||
|
||||
# TODO: add lora support
|
||||
# def add_lora(self, lora_request: LoRARequest) -> bool:
|
||||
# assert lora_request.lora_int_id > 0, "lora_id must be greater than 0."
|
||||
# return all(self.collective_rpc("add_lora", args=(lora_request, )))
|
||||
|
||||
# def remove_lora(self, lora_id: int) -> bool:
|
||||
# assert lora_id > 0, "lora_id must be greater than 0."
|
||||
# return all(self.collective_rpc("remove_lora", args=(lora_id, )))
|
||||
|
||||
# def pin_lora(self, lora_id: int) -> bool:
|
||||
# assert lora_id > 0, "lora_id must be greater than 0."
|
||||
# return all(self.collective_rpc("pin_lora", args=(lora_id, )))
|
||||
|
||||
# def list_loras(self) -> Set[int]:
|
||||
# sets = self.collective_rpc("list_loras")
|
||||
# for s in sets:
|
||||
# assert s == sets[0], "All workers should have the same LORAs."
|
||||
# return sets[0]
|
||||
|
||||
def profile(self, is_start: bool = True):
|
||||
self.collective_rpc("profile", args=(is_start, ))
|
||||
|
||||
# def start_profile(self) -> None:
|
||||
# self.collective_rpc("start_profile")
|
||||
|
||||
# def stop_profile(self) -> None:
|
||||
# self.collective_rpc("stop_profile")
|
||||
|
||||
def sleep(self, level: int = 1):
|
||||
if self.is_sleeping:
|
||||
logger.warning("Executor is already sleeping.")
|
||||
return
|
||||
time_before_sleep = time.perf_counter()
|
||||
self.collective_rpc("sleep", kwargs=dict(level=level))
|
||||
time_after_sleep = time.perf_counter()
|
||||
self.is_sleeping = True
|
||||
logger.info("It took %.6f seconds to fall asleep.",
|
||||
time_after_sleep - time_before_sleep)
|
||||
|
||||
def wake_up(self):
|
||||
if not self.is_sleeping:
|
||||
logger.warning("Executor is not sleeping.")
|
||||
return
|
||||
time_before_wakeup = time.perf_counter()
|
||||
self.collective_rpc("wake_up")
|
||||
time_after_wakeup = time.perf_counter()
|
||||
self.is_sleeping = False
|
||||
logger.info("It took %.6f seconds to wake up.",
|
||||
time_after_wakeup - time_before_wakeup)
|
||||
|
||||
def save_sharded_state(
|
||||
self,
|
||||
path: str,
|
||||
pattern: Optional[str] = None,
|
||||
max_size: Optional[int] = None,
|
||||
) -> None:
|
||||
self.collective_rpc("save_sharded_state",
|
||||
kwargs=dict(path=path,
|
||||
pattern=pattern,
|
||||
max_size=max_size))
|
||||
|
||||
@abstractmethod
|
||||
def check_health(self) -> None:
|
||||
"""Checks if the executor is healthy. If not, it should raise an
|
||||
exception."""
|
||||
raise NotImplementedError
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Shutdown the executor."""
|
||||
return
|
||||
|
||||
def __del__(self):
|
||||
self.shutdown()
|
||||
|
||||
# async def execute_model_async(
|
||||
# self,
|
||||
# execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]:
|
||||
# """Executes one model step on the given sequences."""
|
||||
# output = await make_async(self.execute_model)(execute_model_req)
|
||||
# return output
|
||||
|
||||
async def stop_remote_worker_execution_loop_async(self) -> None:
|
||||
"""Releases parallel workers from model loop."""
|
||||
return
|
||||
|
||||
async def check_health_async(self) -> None:
|
||||
"""Checks if the executor is healthy. If not, it should raise an
|
||||
exception."""
|
||||
self.check_health()
|
||||
|
||||
|
||||
# class Executor(ExecutorBase):
|
||||
# """
|
||||
# Abstract class for v1 executors, mainly define some methods for v1.
|
||||
# For methods shared by v0 and v1, define them in ExecutorBase"""
|
||||
|
||||
# def initialize_from_config(self,
|
||||
# kv_cache_configs: list[KVCacheConfig]) -> None:
|
||||
# """
|
||||
# Initialize the KV caches and begin the model execution loop of the
|
||||
# underlying workers.
|
||||
# """
|
||||
# self.collective_rpc("initialize_from_config", args=(kv_cache_configs, ))
|
||||
# self.collective_rpc("compile_or_warm_up_model")
|
||||
|
||||
# def determine_available_memory(self) -> list[int]: # in bytes
|
||||
# output = self.collective_rpc("determine_available_memory")
|
||||
# return output
|
||||
|
||||
# def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
|
||||
# output = self.collective_rpc("get_kv_cache_spec")
|
||||
# return output
|
||||
|
||||
# def execute_model(
|
||||
# self,
|
||||
# scheduler_output,
|
||||
# ) -> Union[ModelRunnerOutput, Future[ModelRunnerOutput]]:
|
||||
# output = self.collective_rpc("execute_model", args=(scheduler_output, ))
|
||||
# return output[0]
|
||||
|
||||
# @property
|
||||
# def max_concurrent_batches(self) -> int:
|
||||
# return 1
|
||||
|
||||
# def profile(self, is_start: bool = True):
|
||||
# self.collective_rpc("profile", args=(is_start, ))
|
||||
@@ -1,387 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import os
|
||||
import pickle
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
import weakref
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum, auto
|
||||
from functools import partial
|
||||
from multiprocessing.process import BaseProcess
|
||||
from typing import Any, Callable, Optional, Union
|
||||
|
||||
import cloudpickle
|
||||
import psutil
|
||||
import zmq
|
||||
|
||||
from fastvideo.v1.config import FastVideoArgs
|
||||
from fastvideo.v1.distributed import (destroy_distributed_environment,
|
||||
destroy_model_parallel)
|
||||
from vllm.distributed.device_communicators.shm_broadcast import (Handle,
|
||||
MessageQueue)
|
||||
from fastvideo.v1.executor.multiproc_worker_utils import (
|
||||
_add_prefix, set_multiprocessing_worker_envs)
|
||||
from fastvideo.v1.logger import init_logger
|
||||
from fastvideo.v1.utils import (get_distributed_init_method, get_mp_context,
|
||||
get_open_port, get_open_zmq_ipc_path, zmq_socket_ctx)
|
||||
from fastvideo.v1.executor.abstract import Executor
|
||||
from fastvideo.v1.worker.worker_base import WorkerWrapperBase
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
POLLING_TIMEOUT_MS = 5000
|
||||
POLLING_TIMEOUT_S = POLLING_TIMEOUT_MS // 1000
|
||||
|
||||
|
||||
class MultiprocExecutor(Executor):
|
||||
|
||||
def _init_executor(self) -> None:
|
||||
# Call self.shutdown at exit to clean up
|
||||
# and ensure workers will be terminated.
|
||||
self._finalizer = weakref.finalize(self, self.shutdown)
|
||||
|
||||
# The child processes will send SIGUSR1 when unrecoverable
|
||||
# errors happen.
|
||||
def sigusr1_handler(signum, frame):
|
||||
logger.fatal(
|
||||
"MulitprocExecutor got fatal signal from worker processes, "
|
||||
"shutting down. See stack trace above for root cause issue.")
|
||||
# Propagate error up to parent process.
|
||||
parent_process = psutil.Process().parent()
|
||||
parent_process.send_signal(signal.SIGUSR1)
|
||||
self.shutdown()
|
||||
|
||||
signal.signal(signal.SIGUSR1, sigusr1_handler)
|
||||
|
||||
self.world_size = self.parallel_config.world_size
|
||||
tensor_parallel_size = self.parallel_config.tensor_parallel_size
|
||||
assert self.world_size == tensor_parallel_size, (
|
||||
f"world_size ({self.world_size}) must be equal to the "
|
||||
f"tensor_parallel_size ({tensor_parallel_size}). "
|
||||
f"Pipeline parallelism is not yet implemented in v1")
|
||||
|
||||
# Set multiprocessing envs that are common to V0 and V1
|
||||
set_multiprocessing_worker_envs(self.parallel_config)
|
||||
|
||||
# Multiprocessing-based executor does not support multi-node setting.
|
||||
# Since it only works for single node, we can use the loopback address
|
||||
# 127.0.0.1 for communication.
|
||||
distributed_init_method = get_distributed_init_method(
|
||||
"127.0.0.1", get_open_port())
|
||||
|
||||
# Initialize worker and set up message queues for SchedulerOutputs
|
||||
# and ModelRunnerOutputs
|
||||
self.rpc_broadcast_mq = MessageQueue(self.world_size, self.world_size)
|
||||
scheduler_output_handle = self.rpc_broadcast_mq.export_handle()
|
||||
|
||||
# Create workers
|
||||
self.workers: list[WorkerProcHandle] = []
|
||||
for rank in range(self.world_size):
|
||||
worker = WorkerProc.make_worker_process(self.vllm_config, rank,
|
||||
rank,
|
||||
distributed_init_method,
|
||||
scheduler_output_handle)
|
||||
self.workers.append(worker)
|
||||
|
||||
# Ensure message queues are ready. Will deadlock if re-ordered
|
||||
# Must be kept consistent with the WorkerProc
|
||||
self.rpc_broadcast_mq.wait_until_ready()
|
||||
for w in self.workers:
|
||||
w.worker_response_mq.wait_until_ready()
|
||||
|
||||
def collective_rpc(self,
|
||||
method: Union[str, Callable],
|
||||
timeout: Optional[float] = None,
|
||||
args: tuple = (),
|
||||
kwargs: Optional[dict] = None) -> list[Any]:
|
||||
start_time = time.monotonic()
|
||||
kwargs = kwargs or {}
|
||||
|
||||
# NOTE: If the args are heterogeneous, then we pack them into a list,
|
||||
# and unpack them in the method of every worker, because every worker
|
||||
# knows their own rank.
|
||||
try:
|
||||
if isinstance(method, str):
|
||||
send_method = method
|
||||
else:
|
||||
send_method = cloudpickle.dumps(
|
||||
method, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
self.rpc_broadcast_mq.enqueue((send_method, args, kwargs))
|
||||
|
||||
responses = [None] * self.world_size
|
||||
for w in self.workers:
|
||||
dequeue_timeout = timeout - (time.monotonic() - start_time
|
||||
) if timeout is not None else None
|
||||
status, result = w.worker_response_mq.dequeue(
|
||||
timeout=dequeue_timeout)
|
||||
|
||||
if status != WorkerProc.ResponseStatus.SUCCESS:
|
||||
if isinstance(result, Exception):
|
||||
raise result
|
||||
else:
|
||||
raise RuntimeError("Worker failed")
|
||||
|
||||
responses[w.rank] = result
|
||||
|
||||
return responses
|
||||
except TimeoutError as e:
|
||||
raise TimeoutError(f"RPC call to {method} timed out.") from e
|
||||
except Exception as e:
|
||||
# Re-raise any other exceptions
|
||||
raise e
|
||||
|
||||
def _ensure_worker_termination(self):
|
||||
"""Ensure that all worker processes are terminated. Assumes workers have
|
||||
received termination requests. Waits for processing, then sends
|
||||
termination and kill signals if needed."""
|
||||
|
||||
def wait_for_termination(procs, timeout):
|
||||
if not time:
|
||||
# If we are in late stage shutdown, the interpreter may replace
|
||||
# `time` with `None`.
|
||||
return all(not proc.is_alive() for proc in procs)
|
||||
start_time = time.time()
|
||||
while time.time() - start_time < timeout:
|
||||
if all(not proc.is_alive() for proc in procs):
|
||||
return True
|
||||
time.sleep(0.1)
|
||||
return False
|
||||
|
||||
# Send SIGTERM if still running
|
||||
active_procs = [w.proc for w in self.workers if w.proc.is_alive()]
|
||||
for p in active_procs:
|
||||
p.terminate()
|
||||
if not wait_for_termination(active_procs, 4):
|
||||
# Send SIGKILL if still running
|
||||
active_procs = [p for p in active_procs if p.is_alive()]
|
||||
for p in active_procs:
|
||||
p.kill()
|
||||
|
||||
self._cleanup_sockets()
|
||||
|
||||
def _cleanup_sockets(self):
|
||||
for w in self.workers:
|
||||
# Remove the zmq ipc socket file
|
||||
socket_path = w.ready_path.replace("ipc://", "")
|
||||
if os and os.path.exists(socket_path):
|
||||
os.remove(socket_path)
|
||||
|
||||
def shutdown(self):
|
||||
"""Properly shut down the executor and its workers"""
|
||||
if not getattr(self, 'shutting_down', False):
|
||||
self.shutting_down = True
|
||||
for w in self.workers:
|
||||
w.worker_response_mq = None
|
||||
self._ensure_worker_termination()
|
||||
|
||||
self.rpc_broadcast_mq = None
|
||||
|
||||
def check_health(self) -> None:
|
||||
self.collective_rpc("check_health", timeout=10)
|
||||
return
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkerProcHandle:
|
||||
proc: BaseProcess
|
||||
rank: int
|
||||
ready_path: str
|
||||
worker_response_mq: MessageQueue # The worker process writes to this MQ
|
||||
|
||||
|
||||
class WorkerProc:
|
||||
"""Wrapper that runs one Worker in a separate process."""
|
||||
|
||||
READY_STR = "READY"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
local_rank: int,
|
||||
rank: int,
|
||||
distributed_init_method: str,
|
||||
input_shm_handle: Handle,
|
||||
ready_path: str,
|
||||
):
|
||||
self.rank = rank
|
||||
wrapper = WorkerWrapperBase(vllm_config=vllm_config, rpc_rank=rank)
|
||||
# TODO: move `init_worker` to executor level as a collective rpc call
|
||||
all_kwargs: list[dict] = [
|
||||
{} for _ in range(vllm_config.parallel_config.world_size)
|
||||
]
|
||||
all_kwargs[rank] = {
|
||||
"vllm_config": vllm_config,
|
||||
"local_rank": local_rank,
|
||||
"rank": rank,
|
||||
"distributed_init_method": distributed_init_method,
|
||||
"is_driver_worker": rank == 0,
|
||||
}
|
||||
wrapper.init_worker(all_kwargs)
|
||||
self.worker = wrapper
|
||||
|
||||
pid = os.getpid()
|
||||
_add_prefix(sys.stdout, f"FastVideoWorker rank={rank}", pid)
|
||||
_add_prefix(sys.stderr, f"FastVideoWorker rank={rank}", pid)
|
||||
|
||||
# Initialize MessageQueue for receiving SchedulerOutput
|
||||
self.rpc_broadcast_mq = MessageQueue.create_from_handle(
|
||||
input_shm_handle, self.worker.rank)
|
||||
|
||||
# Initializes a message queue for sending the model output
|
||||
self.worker_response_mq = MessageQueue(1, 1)
|
||||
worker_response_mq_handle = self.worker_response_mq.export_handle()
|
||||
|
||||
# Send Readiness signal to EngineCore process.
|
||||
# Set linger here because we want to ensure the message has
|
||||
# been sent before the context is closed.
|
||||
with zmq_socket_ctx(ready_path, zmq.constants.PUSH,
|
||||
linger=10000) as ready_socket:
|
||||
payload = pickle.dumps(worker_response_mq_handle,
|
||||
protocol=pickle.HIGHEST_PROTOCOL)
|
||||
ready_socket.send_string(WorkerProc.READY_STR)
|
||||
ready_socket.send(payload)
|
||||
|
||||
self.worker.init_device()
|
||||
self.worker.load_model()
|
||||
|
||||
@staticmethod
|
||||
def make_worker_process(
|
||||
vllm_config: VllmConfig,
|
||||
local_rank: int,
|
||||
rank: int,
|
||||
distributed_init_method: str,
|
||||
input_shm_handle, # Receive SchedulerOutput
|
||||
) -> WorkerProcHandle:
|
||||
context = get_mp_context()
|
||||
|
||||
# ZMQ path for worker to send ready message and shm_broadcast handle
|
||||
# back to core process.
|
||||
ready_path = get_open_zmq_ipc_path()
|
||||
|
||||
process_kwargs = {
|
||||
"vllm_config": vllm_config,
|
||||
"local_rank": local_rank,
|
||||
"rank": rank,
|
||||
"distributed_init_method": distributed_init_method,
|
||||
"input_shm_handle": input_shm_handle,
|
||||
"ready_path": ready_path,
|
||||
}
|
||||
# Run EngineCore busy loop in background process.
|
||||
proc = context.Process(target=WorkerProc.worker_main,
|
||||
kwargs=process_kwargs,
|
||||
daemon=True)
|
||||
|
||||
with zmq_socket_ctx(ready_path, zmq.constants.PULL) as ready_socket:
|
||||
proc.start()
|
||||
|
||||
# Wait for startup
|
||||
worker_response_mq_handle = WorkerProc.wait_for_startup(
|
||||
proc, ready_socket)
|
||||
|
||||
worker_response_mq = MessageQueue.create_from_handle(
|
||||
worker_response_mq_handle, 0)
|
||||
|
||||
return WorkerProcHandle(proc, rank, ready_path, worker_response_mq)
|
||||
|
||||
def shutdown(self):
|
||||
self.rpc_broadcast_mq = None
|
||||
self.worker_response_mq = None
|
||||
destroy_model_parallel()
|
||||
destroy_distributed_environment()
|
||||
|
||||
@staticmethod
|
||||
def worker_main(*args, **kwargs):
|
||||
""" Worker initialization and execution loops.
|
||||
This runs a background process """
|
||||
|
||||
# Signal handler used for graceful termination.
|
||||
# SystemExit exception is only raised once to allow this and worker
|
||||
# processes to terminate without error
|
||||
shutdown_requested = False
|
||||
|
||||
def signal_handler(signum, frame):
|
||||
nonlocal shutdown_requested
|
||||
if not shutdown_requested:
|
||||
shutdown_requested = True
|
||||
raise SystemExit()
|
||||
|
||||
# Either SIGTERM or SIGINT will terminate the worker
|
||||
signal.signal(signal.SIGTERM, signal_handler)
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
worker = None
|
||||
try:
|
||||
worker = WorkerProc(*args, **kwargs)
|
||||
|
||||
# Ensure message queues are ready. Will deadlock if re-ordered.
|
||||
# Must be kept consistent with the Executor
|
||||
worker.rpc_broadcast_mq.wait_until_ready()
|
||||
worker.worker_response_mq.wait_until_ready()
|
||||
|
||||
worker.worker_busy_loop()
|
||||
|
||||
except SystemExit:
|
||||
logger.debug("Worker interrupted.")
|
||||
|
||||
except Exception:
|
||||
# worker_busy_loop sends exceptions exceptons to Executor
|
||||
# for shutdown, but if there is an error in startup or an
|
||||
# error with IPC itself, we need to alert the parent.
|
||||
psutil.Process().parent().send_signal(signal.SIGUSR1)
|
||||
raise
|
||||
|
||||
finally:
|
||||
# Clean up once worker exits busy loop
|
||||
if worker is not None:
|
||||
worker.shutdown()
|
||||
worker = None
|
||||
|
||||
@staticmethod
|
||||
def wait_for_startup(
|
||||
proc: BaseProcess,
|
||||
ready_socket: zmq.Socket,
|
||||
) -> Optional[Handle]:
|
||||
"""Wait until the Worker is ready."""
|
||||
|
||||
# Wait for Worker to send READY.
|
||||
while ready_socket.poll(timeout=POLLING_TIMEOUT_MS) == 0:
|
||||
logger.debug("Waiting for WorkerProc to startup.")
|
||||
|
||||
if not proc.is_alive():
|
||||
raise RuntimeError("WorkerProc failed to start.")
|
||||
|
||||
message = ready_socket.recv_string()
|
||||
assert message == WorkerProc.READY_STR
|
||||
handle_frame = ready_socket.recv(copy=False)
|
||||
handle = pickle.loads(handle_frame.buffer)
|
||||
return handle
|
||||
|
||||
class ResponseStatus(Enum):
|
||||
SUCCESS = auto()
|
||||
FAILURE = auto()
|
||||
|
||||
def worker_busy_loop(self):
|
||||
"""Main busy loop for Multiprocessing Workers"""
|
||||
while True:
|
||||
method, args, kwargs = self.rpc_broadcast_mq.dequeue()
|
||||
|
||||
try:
|
||||
if isinstance(method, str):
|
||||
func = getattr(self.worker, method)
|
||||
elif isinstance(method, bytes):
|
||||
func = partial(cloudpickle.loads(method), self.worker)
|
||||
output = func(*args, **kwargs)
|
||||
except Exception as e:
|
||||
# Notes have been introduced in python 3.11
|
||||
if hasattr(e, "add_note"):
|
||||
e.add_note(traceback.format_exc())
|
||||
self.worker_response_mq.enqueue(
|
||||
(WorkerProc.ResponseStatus.FAILURE, e))
|
||||
logger.exception("WorkerProc hit an exception: %s", exc_info=e)
|
||||
continue
|
||||
|
||||
self.worker_response_mq.enqueue(
|
||||
(WorkerProc.ResponseStatus.SUCCESS, output))
|
||||
@@ -1,64 +0,0 @@
|
||||
import os
|
||||
from typing import TextIO
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.v1.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# ANSI color codes
|
||||
CYAN = '\033[1;36m'
|
||||
RESET = '\033[0;0m'
|
||||
|
||||
def _add_prefix(file: TextIO, worker_name: str, pid: int) -> None:
|
||||
"""Prepend each output line with process-specific prefix"""
|
||||
|
||||
prefix = f"{CYAN}({worker_name} pid={pid}){RESET} "
|
||||
file_write = file.write
|
||||
|
||||
def write_with_prefix(s: str):
|
||||
if not s:
|
||||
return
|
||||
if file.start_new_line: # type: ignore[attr-defined]
|
||||
file_write(prefix)
|
||||
idx = 0
|
||||
while (next_idx := s.find('\n', idx)) != -1:
|
||||
next_idx += 1
|
||||
file_write(s[idx:next_idx])
|
||||
if next_idx == len(s):
|
||||
file.start_new_line = True # type: ignore[attr-defined]
|
||||
return
|
||||
file_write(prefix)
|
||||
idx = next_idx
|
||||
file_write(s[idx:])
|
||||
file.start_new_line = False # type: ignore[attr-defined]
|
||||
|
||||
file.start_new_line = True # type: ignore[attr-defined]
|
||||
file.write = write_with_prefix # type: ignore[method-assign]
|
||||
|
||||
|
||||
def set_multiprocessing_worker_envs(parallel_config):
|
||||
""" Set up environment variables that should be used when there are workers
|
||||
in a multiprocessing environment. This should be called by the parent
|
||||
process before worker processes are created"""
|
||||
|
||||
_maybe_force_spawn()
|
||||
|
||||
# Configure thread parallelism if OMP_NUM_THREADS isn't set
|
||||
#
|
||||
# Helps to avoid CPU contention. The default of spawning a thread per
|
||||
# core combined with multiprocessing for each GPU can have a negative
|
||||
# impact on performance. The contention is amplified when running in a
|
||||
# container where CPU limits can cause throttling.
|
||||
default_omp_num_threads = 1
|
||||
if "OMP_NUM_THREADS" not in os.environ and (
|
||||
current_parallelism :=
|
||||
torch.get_num_threads()) > default_omp_num_threads:
|
||||
logger.warning(
|
||||
"Reducing Torch parallelism from %d threads to %d to avoid "
|
||||
"unnecessary CPU contention. Set OMP_NUM_THREADS in the "
|
||||
"external environment to tune this value as needed.",
|
||||
current_parallelism, default_omp_num_threads)
|
||||
os.environ["OMP_NUM_THREADS"] = str(default_omp_num_threads)
|
||||
torch.set_num_threads(default_omp_num_threads)
|
||||
@@ -1,61 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from concurrent.futures import Future
|
||||
from typing import Union
|
||||
|
||||
from vllm.executor.ray_distributed_executor import ( # noqa
|
||||
RayDistributedExecutor as RayDistributedExecutorV0)
|
||||
from vllm.v1.executor.abstract import Executor
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
|
||||
|
||||
class FutureWrapper(Future):
|
||||
"""A wrapper around a Ray output reference to meet the interface
|
||||
of .execute_model().
|
||||
"""
|
||||
|
||||
def __init__(self, ref):
|
||||
super().__init__()
|
||||
self.ref = ref
|
||||
|
||||
def result(self, timeout=None):
|
||||
if timeout is not None:
|
||||
raise NotImplementedError("timeout is not supported")
|
||||
return self.ref.get()
|
||||
|
||||
|
||||
class RayDistributedExecutor(RayDistributedExecutorV0, Executor):
|
||||
"""Ray distributed executor using Ray Compiled Graphs."""
|
||||
|
||||
@property
|
||||
def max_concurrent_batches(self) -> int:
|
||||
"""Ray distributed executor supports pipeline parallelism,
|
||||
meaning that it allows PP size batches to be executed concurrently.
|
||||
"""
|
||||
return self.parallel_config.pipeline_parallel_size
|
||||
|
||||
def execute_model(
|
||||
self,
|
||||
scheduler_output,
|
||||
) -> Union[ModelRunnerOutput, Future[ModelRunnerOutput]]:
|
||||
"""Execute the model on the Ray workers.
|
||||
|
||||
Args:
|
||||
scheduler_output: The scheduler output to execute.
|
||||
|
||||
Returns:
|
||||
The model runner output.
|
||||
"""
|
||||
# Build the compiled DAG for the first time.
|
||||
if self.forward_dag is None: # type: ignore
|
||||
self.forward_dag = self._compiled_ray_dag(enable_asyncio=False)
|
||||
|
||||
refs = self.forward_dag.execute(scheduler_output) # type: ignore
|
||||
|
||||
# When PP is not used, we block here until the result is available.
|
||||
if self.max_concurrent_batches == 1:
|
||||
return refs[0].get()
|
||||
|
||||
# When PP is used, we return a FutureWrapper immediately so that
|
||||
# the scheduler can yield to the next batch.
|
||||
return FutureWrapper(refs[0])
|
||||
@@ -1,10 +0,0 @@
|
||||
from fastvideo.v1.executor.abstract import Executor
|
||||
|
||||
|
||||
class TorchRunExecutor(Executor):
|
||||
|
||||
def _init_executor(self) -> None:
|
||||
pass
|
||||
|
||||
def collective_rpc(self, method: Union[str, Callable], timeout: Optional[float] = None, args: tuple = (), kwargs: Optional[dict] = None) -> list[Any]:
|
||||
pass
|
||||
@@ -11,7 +11,9 @@ from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
|
||||
from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
|
||||
from fastvideo.v1.pipelines.pipeline_registry import PipelineRegistry
|
||||
from fastvideo.v1.utils import (maybe_download_model,
|
||||
verify_model_config_and_directory)
|
||||
verify_model_config_and_directory, diff_keys,
|
||||
update_in_place)
|
||||
from fastvideo.v1.configs.registry import get_pipeline_config_for_name
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -40,6 +42,17 @@ def build_pipeline(fastvideo_args: FastVideoArgs) -> ComposedPipelineBase:
|
||||
pipeline_cls, pipeline_architecture = PipelineRegistry.resolve_pipeline_cls(
|
||||
pipeline_architecture)
|
||||
|
||||
# Get the user's input
|
||||
user_args = diff_keys(FastVideoArgs(model_path), fastvideo_args)
|
||||
# Get the optimal config
|
||||
pipeline_cfg_cls = get_pipeline_config_for_name(model_path)
|
||||
assert pipeline_cfg_cls is not None, f"Cannot get pipeline config for {model_path}"
|
||||
pipeline_cfg = pipeline_cfg_cls()
|
||||
# Override the fastvideo_args with pipeline config EXCEPT for the user-specified fields
|
||||
update_in_place(fastvideo_args,
|
||||
pipeline_cfg,
|
||||
ignore_fields=tuple(user_args))
|
||||
|
||||
# instantiate the pipeline
|
||||
pipeline = pipeline_cls(model_path, fastvideo_args, config)
|
||||
logger.info("Pipeline instantiated")
|
||||
|
||||
@@ -12,6 +12,7 @@ import sys
|
||||
import tempfile
|
||||
from functools import wraps, partial
|
||||
from typing import Any, Dict, List, Optional, Type, TypeVar, Union, cast, Callable
|
||||
from dataclasses import asdict, fields
|
||||
import cloudpickle
|
||||
|
||||
import filelock
|
||||
@@ -543,3 +544,13 @@ def run_method(obj: Any, method: Union[str, bytes, Callable], args: tuple[Any],
|
||||
else:
|
||||
func = partial(method, obj) # type: ignore
|
||||
return func(*args, **kwargs)
|
||||
|
||||
|
||||
def diff_keys(a, b):
|
||||
return [k for k in asdict(a) if asdict(a)[k] != asdict(b)[k]]
|
||||
|
||||
|
||||
def update_in_place(target, source, ignore_fields=()):
|
||||
for f in fields(target):
|
||||
if hasattr(source, f.name) and f.name not in list(ignore_fields):
|
||||
setattr(target, f.name, getattr(source, f.name))
|
||||
|
||||
@@ -1,155 +0,0 @@
|
||||
from typing import Optional
|
||||
from fastvideo.v1.worker.worker_base import WorkerBase
|
||||
from fastvideo.v1.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.v1.logger import init_logger
|
||||
import torch
|
||||
import os
|
||||
import gc
|
||||
from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
|
||||
from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
|
||||
from fastvideo.v1.pipelines import build_pipeline
|
||||
from fastvideo.v1.distributed import (init_distributed_environment,
|
||||
initialize_model_parallel)
|
||||
from fastvideo.v1.platforms import current_platform
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class Worker(WorkerBase):
|
||||
|
||||
def __init__(self,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
local_rank: int,
|
||||
rank: int,
|
||||
distributed_init_method: str,
|
||||
is_driver_worker: bool = False):
|
||||
super().__init__(fastvideo_args, local_rank, rank,
|
||||
distributed_init_method, is_driver_worker)
|
||||
|
||||
# Torch profiler. Enabled and configured through env vars:
|
||||
# FASTVIDEO_TORCH_PROFILER_DIR=/path/to/save/trace
|
||||
|
||||
# TODO: enable profiling
|
||||
# if envs.FASTVIDEO_TORCH_PROFILER_DIR:
|
||||
if False:
|
||||
# torch_profiler_trace_dir = envs.FASTVIDEO_TORCH_PROFILER_DIR
|
||||
logger.info("Profiling enabled. Traces will be saved to: %s",
|
||||
torch_profiler_trace_dir)
|
||||
self.profiler = torch.profiler.profile(
|
||||
activities=[
|
||||
torch.profiler.ProfilerActivity.CPU,
|
||||
torch.profiler.ProfilerActivity.CUDA,
|
||||
],
|
||||
with_stack=True,
|
||||
on_trace_ready=torch.profiler.tensorboard_trace_handler(
|
||||
torch_profiler_trace_dir, use_gzip=True))
|
||||
else:
|
||||
self.profiler = None
|
||||
|
||||
def init_device(self):
|
||||
if self.device_config.device.type == "cuda":
|
||||
# torch.distributed.all_reduce does not free the input tensor until
|
||||
# the synchronization point. This causes the memory usage to grow
|
||||
# as the number of all_reduce calls increases. This env var disables
|
||||
# this behavior.
|
||||
# Related issue:
|
||||
# https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
|
||||
os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
|
||||
|
||||
# This env var set by Ray causes exceptions with graph building.
|
||||
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
|
||||
self.device = torch.device(f"cuda:{self.local_rank}")
|
||||
torch.cuda.set_device(self.device)
|
||||
|
||||
_check_if_gpu_supports_dtype(self.model_config.dtype)
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
self.init_gpu_memory = torch.cuda.mem_get_info()[0]
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Not support device type: {self.device_config.device}")
|
||||
# Initialize the distributed environment.
|
||||
init_worker_distributed_environment(self.fastvideo_config, self.rank,
|
||||
self.distributed_init_method,
|
||||
self.local_rank)
|
||||
# Set random seed.
|
||||
# set_random_seed(self.model_config.seed)
|
||||
|
||||
# Construct the model runner
|
||||
self.pipeline: Optional[ComposedPipelineBase] = None
|
||||
|
||||
def load_pipeline(self) -> None:
|
||||
self.pipeline = build_pipeline(self.fastvideo_config)
|
||||
logger.info("Pipeline Ready")
|
||||
|
||||
def get_pipeline(self) -> ComposedPipelineBase:
|
||||
return self.pipeline
|
||||
|
||||
@torch.inference_mode()
|
||||
def execute_model(
|
||||
self,
|
||||
forward_batch: ForwardBatch,
|
||||
) -> Optional[ForwardBatch]:
|
||||
output = self.pipeline.forward(forward_batch)
|
||||
return output if self.is_driver_worker else None
|
||||
|
||||
def profile(self, is_start: bool = True):
|
||||
if self.profiler is None:
|
||||
raise RuntimeError("Profiler is not enabled.")
|
||||
if is_start:
|
||||
self.profiler.start()
|
||||
else:
|
||||
self.profiler.stop()
|
||||
|
||||
# def add_lora(self, lora_request: LoRARequest) -> bool:
|
||||
# return self.model_runner.add_lora(lora_request)
|
||||
|
||||
# def remove_lora(self, lora_id: int) -> bool:
|
||||
# return self.model_runner.remove_lora(lora_id)
|
||||
|
||||
# def list_loras(self) -> set[int]:
|
||||
# return self.model_runner.list_loras()
|
||||
|
||||
# def pin_lora(self, lora_id: int) -> bool:
|
||||
# return self.model_runner.pin_lora(lora_id)
|
||||
|
||||
def check_health(self) -> None:
|
||||
# worker will always be healthy as long as it's running.
|
||||
return
|
||||
|
||||
|
||||
def init_worker_distributed_environment(
|
||||
fastvideo_args: FastVideoArgs,
|
||||
rank: int,
|
||||
distributed_init_method: Optional[str] = None,
|
||||
local_rank: int = -1,
|
||||
) -> None:
|
||||
"""Initialize the distributed environment."""
|
||||
init_distributed_environment(fastvideo_args.world_size, rank,
|
||||
distributed_init_method, local_rank)
|
||||
|
||||
# ensure_model_parallel_initialized(fastvideo_args.tp_size,
|
||||
# fastvideo_args.sp_size)
|
||||
initialize_model_parallel(
|
||||
sequence_model_parallel_size=fastvideo_args.sp_size,
|
||||
tensor_model_parallel_size=fastvideo_args.tp_size)
|
||||
|
||||
|
||||
def _check_if_gpu_supports_dtype(torch_dtype: torch.dtype):
|
||||
# Check if the GPU supports the dtype.
|
||||
if torch_dtype == torch.bfloat16: # noqa: SIM102
|
||||
if not current_platform.has_device_capability(80):
|
||||
capability = current_platform.get_device_capability()
|
||||
gpu_name = current_platform.get_device_name()
|
||||
|
||||
if capability is None:
|
||||
compute_str = "does not have a compute capability"
|
||||
else:
|
||||
version_str = capability.as_version_str()
|
||||
compute_str = f"has compute capability {version_str}"
|
||||
|
||||
raise ValueError(
|
||||
"Bfloat16 is only supported on GPUs with compute capability "
|
||||
f"of at least 8.0. Your {gpu_name} GPU {compute_str}. "
|
||||
"You can use float16 instead by explicitly setting the "
|
||||
"`dtype` flag in CLI, for example: --dtype=half.")
|
||||
@@ -1,144 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
import os
|
||||
|
||||
from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
|
||||
from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
|
||||
from fastvideo.v1.fastvideo_args import FastVideoArgs, set_current_fastvideo_args
|
||||
from fastvideo.v1.utils import (update_environment_variables,
|
||||
resolve_obj_by_qualname, run_method)
|
||||
from fastvideo.v1.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class WorkerBase:
|
||||
|
||||
def __init__(self,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
local_rank: int,
|
||||
rank: int,
|
||||
distributed_init_method: str,
|
||||
is_driver_worker: bool = False):
|
||||
self.fastvideo_args = fastvideo_args
|
||||
self.local_rank = local_rank
|
||||
self.rank = rank
|
||||
self.distributed_init_method = distributed_init_method
|
||||
self.is_driver_worker = is_driver_worker
|
||||
|
||||
def init_device(self) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def get_pipeline(self) -> ComposedPipelineBase:
|
||||
raise NotImplementedError
|
||||
|
||||
def execute_pipeline(self, forward_batch: ForwardBatch) -> ForwardBatch:
|
||||
raise NotImplementedError
|
||||
|
||||
# TODO: add lora support
|
||||
# def add_lora(self, lora_request: LoRARequest) -> bool:
|
||||
# raise NotImplementedError
|
||||
|
||||
# def remove_lora(self, lora_id: int) -> bool:
|
||||
# raise NotImplementedError
|
||||
|
||||
# def pin_lora(self, lora_id: int) -> bool:
|
||||
# raise NotImplementedError
|
||||
|
||||
# def list_loras(self) -> Set[int]:
|
||||
# raise NotImplementedError
|
||||
|
||||
|
||||
class WorkerWrapperBase:
|
||||
"""
|
||||
This class represents one process in an executor/engine. It is responsible
|
||||
for lazily initializing the worker and handling the worker's lifecycle.
|
||||
We first instantiate the WorkerWrapper, which remembers the worker module
|
||||
and class name. Then, when we call `update_environment_variables`, and the
|
||||
real initialization happens in `init_worker`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
rpc_rank: int = 0,
|
||||
) -> None:
|
||||
"""
|
||||
Initialize the worker wrapper with the given fastvideo_args and
|
||||
rpc_rank. Note: rpc_rank is the rank of the worker in the executor. In
|
||||
most cases, it is also the rank of the worker in the distributed group.
|
||||
However, when multiple executors work together, they can be different.
|
||||
e.g. in the case of SPMD-style offline inference with TP=2, users can
|
||||
launch 2 engines/executors, each with only 1 worker. All workers have
|
||||
rpc_rank=0, but they have different ranks in the TP group.
|
||||
"""
|
||||
self.rpc_rank = rpc_rank
|
||||
self.worker: Optional[WorkerBase] = None
|
||||
|
||||
def adjust_rank(self, rank_mapping: Dict[int, int]) -> None:
|
||||
"""
|
||||
Adjust the rpc_rank based on the given mapping.
|
||||
It is only used during the initialization of the executor,
|
||||
to adjust the rpc_rank of workers after we create all workers.
|
||||
"""
|
||||
if self.rpc_rank in rank_mapping:
|
||||
self.rpc_rank = rank_mapping[self.rpc_rank]
|
||||
|
||||
def update_environment_variables(self, envs_list: List[Dict[str,
|
||||
str]]) -> None:
|
||||
envs = envs_list[self.rpc_rank]
|
||||
key = 'CUDA_VISIBLE_DEVICES'
|
||||
if key in envs and key in os.environ:
|
||||
# overwriting CUDA_VISIBLE_DEVICES is desired behavior
|
||||
# suppress the warning in `update_environment_variables`
|
||||
del os.environ[key]
|
||||
update_environment_variables(envs)
|
||||
|
||||
def init_worker(self, all_kwargs: List[Dict[str, Any]]) -> None:
|
||||
"""
|
||||
Here we inject some common logic before initializing the worker.
|
||||
Arguments are passed to the worker class constructor.
|
||||
"""
|
||||
kwargs = all_kwargs[self.rpc_rank]
|
||||
self.fastvideo_args = kwargs.get("fastvideo_args", None)
|
||||
assert self.fastvideo_args is not None, (
|
||||
"fastvideo_args is required to initialize the worker")
|
||||
# enable_trace_function_call_for_thread(self.fastvideo_args)
|
||||
|
||||
if isinstance(self.fastvideo_args.worker_cls, str):
|
||||
worker_class = resolve_obj_by_qualname(
|
||||
self.fastvideo_args.worker_cls)
|
||||
else:
|
||||
raise ValueError(
|
||||
"worker_cls must be a string of a qualified module name")
|
||||
|
||||
with set_current_fastvideo_args(self.fastvideo_args):
|
||||
# To make fastvideo config available during worker initialization
|
||||
self.worker = worker_class(**kwargs)
|
||||
assert self.worker is not None
|
||||
|
||||
def init_device(self):
|
||||
with set_current_fastvideo_args(self.fastvideo_args):
|
||||
# To make fastvideo config available during device initialization
|
||||
self.worker.init_device() # type: ignore
|
||||
|
||||
def execute_method(self, method: Union[str, bytes], *args, **kwargs):
|
||||
try:
|
||||
# method resolution order:
|
||||
# if a method is defined in this class, it will be called directly.
|
||||
# otherwise, since we define `__getattr__` and redirect attribute
|
||||
# query to `self.worker`, the method will be called on the worker.
|
||||
return run_method(self, method, args, kwargs)
|
||||
except Exception as e:
|
||||
# if the driver worker also execute methods,
|
||||
# exceptions in the rest worker may cause deadlock in rpc like ray
|
||||
# see https://github.com/vllm-project/vllm/issues/3455
|
||||
# print the error and inform the user to solve the error
|
||||
msg = (f"Error executing method {method!r}. "
|
||||
"This might cause deadlock in distributed execution.")
|
||||
logger.exception(msg)
|
||||
raise e
|
||||
|
||||
def __getattr__(self, attr):
|
||||
return getattr(self.worker, attr)
|
||||
@@ -11,17 +11,7 @@ torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
|
||||
fastvideo/v1/sample/v1_fastvideo_inference.py \
|
||||
--sp_size $num_gpus \
|
||||
--tp_size $num_gpus \
|
||||
--height 480 \
|
||||
--width 832 \
|
||||
--num_frames 77 \
|
||||
--num_inference_steps 50 \
|
||||
--fps 16 \
|
||||
--guidance_scale 3.0 \
|
||||
--prompt_path ./assets/prompt.txt \
|
||||
--neg_prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
|
||||
--seed 1024 \
|
||||
--output_path outputs_video/ \
|
||||
--model_path $MODEL_BASE \
|
||||
--vae-sp \
|
||||
--text-encoder-precision "fp32" \
|
||||
--use-cpu-offload
|
||||
--model_path $MODEL_BASE
|
||||
@@ -11,19 +11,8 @@ torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
|
||||
fastvideo/v1/sample/v1_fastvideo_inference.py \
|
||||
--sp_size $num_gpus \
|
||||
--tp_size $num_gpus \
|
||||
--height 480 \
|
||||
--width 832 \
|
||||
--num_frames 77 \
|
||||
--num_inference_steps 40 \
|
||||
--fps 16 \
|
||||
--flow_shift 3.0 \
|
||||
--guidance_scale 5.0 \
|
||||
--image_path "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" \
|
||||
--prompt "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot." \
|
||||
--neg_prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
|
||||
--seed 1024 \
|
||||
--output_path outputs_i2v/ \
|
||||
--model_path $MODEL_BASE \
|
||||
--vae-sp \
|
||||
--text-encoder-precision "fp32" \
|
||||
--use-cpu-offload
|
||||
--model_path $MODEL_BASE \
|
||||
Reference in New Issue
Block a user