540 lines
22 KiB
Python
540 lines
22 KiB
Python
# 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 re
|
|
import threading
|
|
import time
|
|
from copy import deepcopy
|
|
from typing import Any
|
|
|
|
import imageio
|
|
import numpy as np
|
|
import torch
|
|
import torchvision
|
|
from einops import rearrange
|
|
import shutil
|
|
import tempfile
|
|
|
|
from fastvideo.configs.sample import SamplingParam
|
|
from fastvideo.fastvideo_args import FastVideoArgs
|
|
from fastvideo.logger import init_logger
|
|
from fastvideo.pipelines import ForwardBatch
|
|
from fastvideo.utils import align_to, shallow_asdict
|
|
from fastvideo.worker.executor import Executor
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _infer_latent_batch_size(batch: ForwardBatch) -> int:
|
|
if isinstance(batch.prompt, list):
|
|
latent_batch_size = len(batch.prompt)
|
|
elif batch.prompt is not None:
|
|
latent_batch_size = 1
|
|
elif batch.prompt_embeds is not None and len(batch.prompt_embeds) > 0:
|
|
latent_batch_size = batch.prompt_embeds[0].shape[0]
|
|
else:
|
|
raise ValueError("Cannot infer batch size from batch; no prompt or prompt_embeds found")
|
|
latent_batch_size *= batch.num_videos_per_prompt
|
|
return latent_batch_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:
|
|
fastvideo_args: The inference arguments
|
|
executor_class: The executor class to use for inference
|
|
"""
|
|
self.fastvideo_args = fastvideo_args
|
|
self.executor = executor_class(fastvideo_args)
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, model_path: str, **kwargs) -> "VideoGenerator":
|
|
"""
|
|
Create a video generator from a pretrained model.
|
|
|
|
Args:
|
|
model_path: Path or identifier for the pretrained model
|
|
pipeline_config: Pipeline config to use for inference
|
|
**kwargs: Additional arguments to customize model loading, set any FastVideoArgs or PipelineConfig attributes here.
|
|
|
|
Returns:
|
|
The created video generator
|
|
|
|
Priority level: Default pipeline config < User's pipeline config < User's kwargs
|
|
"""
|
|
# If users also provide some kwargs, it will override the FastVideoArgs and PipelineConfig.
|
|
kwargs['model_path'] = model_path
|
|
fastvideo_args = FastVideoArgs.from_kwargs(**kwargs)
|
|
|
|
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 | None = None,
|
|
sampling_param: SamplingParam | None = None,
|
|
# Action control inputs (Matrix-Game)
|
|
mouse_cond: torch.Tensor | None = None,
|
|
keyboard_cond: torch.Tensor | None = None,
|
|
grid_sizes: tuple[int, int, int] | list[int] | torch.Tensor
|
|
| None = None,
|
|
**kwargs,
|
|
) -> dict[str, Any] | list[dict[str, Any]]:
|
|
"""
|
|
Generate a video based on the given prompt.
|
|
|
|
Args:
|
|
prompt: The prompt to use for generation (optional if prompt_txt is provided)
|
|
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)
|
|
prompt_path: Path to prompt file
|
|
save_video: Whether to save the video to disk
|
|
return_frames: Whether to include raw frames in the result dict
|
|
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:
|
|
A metadata dictionary for single-prompt generation, or a list of
|
|
metadata dictionaries for prompt-file batch generation.
|
|
"""
|
|
# Handle batch processing from text file
|
|
if sampling_param is None:
|
|
sampling_param = SamplingParam.from_pretrained(self.fastvideo_args.model_path)
|
|
|
|
# Add action control inputs to kwargs if provided
|
|
if mouse_cond is not None:
|
|
kwargs['mouse_cond'] = mouse_cond
|
|
if keyboard_cond is not None:
|
|
kwargs['keyboard_cond'] = keyboard_cond
|
|
if grid_sizes is not None:
|
|
kwargs['grid_sizes'] = grid_sizes
|
|
|
|
sampling_param.update(kwargs)
|
|
|
|
if self.fastvideo_args.prompt_txt is not None or sampling_param.prompt_path is not None:
|
|
prompt_txt_path = sampling_param.prompt_path or self.fastvideo_args.prompt_txt
|
|
if not os.path.exists(prompt_txt_path):
|
|
raise FileNotFoundError(f"Prompt text file not found: {prompt_txt_path}")
|
|
|
|
# Read prompts from file
|
|
with open(prompt_txt_path, encoding='utf-8') as f:
|
|
prompts = [line.strip() for line in f if line.strip()]
|
|
|
|
if not prompts:
|
|
raise ValueError(f"No prompts found in file: {prompt_txt_path}")
|
|
|
|
logger.info("Found %d prompts in %s", len(prompts), prompt_txt_path)
|
|
|
|
results = []
|
|
for i, batch_prompt in enumerate(prompts):
|
|
logger.info("Processing prompt %d/%d: %s...", i + 1, len(prompts), batch_prompt[:100])
|
|
try:
|
|
# Generate video for this prompt using the same logic below
|
|
output_path = self._prepare_output_path(sampling_param.output_path, batch_prompt)
|
|
kwargs["output_path"] = output_path
|
|
result = self._generate_single_video(prompt=batch_prompt, sampling_param=sampling_param, **kwargs)
|
|
|
|
# Add prompt info to result
|
|
result["prompt_index"] = i
|
|
result["prompt"] = batch_prompt
|
|
|
|
results.append(result)
|
|
logger.info("Successfully generated video for prompt %d", i + 1)
|
|
|
|
except Exception as e:
|
|
logger.error("Failed to generate video for prompt %d: %s", i + 1, e)
|
|
continue
|
|
|
|
logger.info("Completed batch processing. Generated %d videos successfully.", len(results))
|
|
return results
|
|
|
|
# Single prompt generation (original behavior)
|
|
if prompt is None:
|
|
raise ValueError("Either prompt or prompt_txt must be provided")
|
|
output_path = self._prepare_output_path(sampling_param.output_path, prompt)
|
|
kwargs["output_path"] = output_path
|
|
return self._generate_single_video(prompt=prompt, sampling_param=sampling_param, **kwargs)
|
|
|
|
def _is_image_workload(self) -> bool:
|
|
"""Return True when the workload produces a single image (t2i, i2i …)."""
|
|
args = getattr(self, "fastvideo_args", None)
|
|
if args is None:
|
|
return False
|
|
return args.workload_type.value.endswith("2i")
|
|
|
|
def _prepare_output_path(
|
|
self,
|
|
output_path: str,
|
|
prompt: str,
|
|
) -> str:
|
|
"""Build a unique, sanitized output file path.
|
|
|
|
The file extension is chosen automatically based on the workload type:
|
|
``.png`` for image workloads (``t2i``, ``i2i``, …) and ``.mp4`` for
|
|
video workloads.
|
|
|
|
- If ``output_path`` already carries the correct extension, treat it
|
|
as a file path.
|
|
- Otherwise, treat ``output_path`` as a directory and derive the
|
|
filename from the prompt.
|
|
- Invalid filename characters are removed; if the name changes, a
|
|
warning is logged.
|
|
- If the target path already exists, a numeric suffix is appended.
|
|
"""
|
|
target_ext = ".png" if self._is_image_workload() else ".mp4"
|
|
|
|
def _sanitize_filename_component(name: str) -> str:
|
|
# Remove characters invalid on common filesystems, strip spaces/dots
|
|
sanitized = re.sub(r'[\\/:*?"<>|]', '', name)
|
|
sanitized = sanitized.strip().strip('.')
|
|
sanitized = re.sub(r'\s+', ' ', sanitized)
|
|
return sanitized or "output"
|
|
|
|
base_path, extension = os.path.splitext(output_path)
|
|
extension_lower = extension.lower()
|
|
|
|
if extension_lower == target_ext:
|
|
output_dir = os.path.dirname(output_path)
|
|
base_name = os.path.basename(base_path) # filename without extension
|
|
sanitized_base = _sanitize_filename_component(base_name)
|
|
if sanitized_base != base_name:
|
|
logger.warning(
|
|
"The output name '%s' contained invalid characters. "
|
|
"It has been renamed to '%s%s'",
|
|
os.path.basename(output_path),
|
|
sanitized_base,
|
|
target_ext,
|
|
)
|
|
out_name = f"{sanitized_base}{target_ext}"
|
|
else:
|
|
# Treat as directory; inform if an unexpected extension was
|
|
# provided.
|
|
if extension:
|
|
logger.info(
|
|
"Output path '%s' has extension '%s' which does not "
|
|
"match the target '%s'; treating it as a directory",
|
|
output_path,
|
|
extension,
|
|
target_ext,
|
|
)
|
|
output_dir = output_path
|
|
prompt_component = _sanitize_filename_component(prompt[:100])
|
|
out_name = f"{prompt_component}{target_ext}"
|
|
|
|
if output_dir:
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
new_output_path = os.path.join(output_dir, out_name)
|
|
counter = 1
|
|
while os.path.exists(new_output_path):
|
|
name_part, ext_part = os.path.splitext(out_name)
|
|
new_name = f"{name_part}_{counter}{ext_part}"
|
|
new_output_path = os.path.join(output_dir, new_name)
|
|
counter += 1
|
|
return new_output_path
|
|
|
|
def _generate_single_video(
|
|
self,
|
|
prompt: str,
|
|
sampling_param: SamplingParam | None = None,
|
|
**kwargs,
|
|
) -> dict[str, Any]:
|
|
"""Internal method for single video generation"""
|
|
# Create a copy of inference args to avoid modifying the original
|
|
fastvideo_args = self.fastvideo_args
|
|
|
|
# Validate inputs
|
|
if not isinstance(prompt, str):
|
|
raise TypeError(f"`prompt` must be a string, but got {type(prompt)}")
|
|
prompt = prompt.strip()
|
|
sampling_param = deepcopy(sampling_param)
|
|
output_path = kwargs["output_path"]
|
|
sampling_param.prompt = prompt
|
|
# Process negative prompt
|
|
if sampling_param.negative_prompt is not None:
|
|
sampling_param.negative_prompt = sampling_param.negative_prompt.strip()
|
|
|
|
# Validate dimensions
|
|
if (sampling_param.height <= 0 or sampling_param.width <= 0 or sampling_param.num_frames <= 0):
|
|
raise ValueError(f"Height, width, and num_frames must be positive integers, got "
|
|
f"height={sampling_param.height}, width={sampling_param.width}, "
|
|
f"num_frames={sampling_param.num_frames}")
|
|
|
|
# Calculate sizes
|
|
target_height = align_to(sampling_param.height, 16)
|
|
target_width = align_to(sampling_param.width, 16)
|
|
|
|
# Calculate latent sizes
|
|
latents_size = [(sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.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: {sampling_param.num_frames}
|
|
prompt: {sampling_param.prompt}
|
|
image_path: {sampling_param.image_path}
|
|
neg_prompt: {sampling_param.negative_prompt}
|
|
seed: {sampling_param.seed}
|
|
infer_steps: {sampling_param.num_inference_steps}
|
|
num_videos_per_prompt: {sampling_param.num_videos_per_prompt}
|
|
guidance_scale: {sampling_param.guidance_scale}
|
|
n_tokens: {n_tokens}
|
|
flow_shift: {fastvideo_args.pipeline_config.flow_shift}
|
|
embedded_guidance_scale: {fastvideo_args.pipeline_config.embedded_cfg_scale}
|
|
save_video: {sampling_param.save_video}
|
|
output_path: {output_path}
|
|
""" # type: ignore[attr-defined]
|
|
logger.info(debug_str)
|
|
|
|
# Prepare batch
|
|
batch = ForwardBatch(
|
|
**shallow_asdict(sampling_param),
|
|
eta=0.0,
|
|
n_tokens=n_tokens,
|
|
VSA_sparsity=fastvideo_args.VSA_sparsity,
|
|
)
|
|
|
|
# Run inference
|
|
start_time = time.perf_counter()
|
|
|
|
# Execute forward pass in a new thread for non-blocking tensor
|
|
# allocation. Capture thread exceptions so we can surface the true
|
|
# failure in the main thread instead of later hitting None outputs.
|
|
result_container = {"output_batch": ForwardBatch(data_type=batch.data_type)}
|
|
thread_error: dict[str, BaseException | None] = {"error": None}
|
|
thread_error_traceback: dict[str, str] = {"traceback": ""}
|
|
|
|
def execute_forward_thread():
|
|
import traceback
|
|
try:
|
|
result_container["output_batch"] = self.executor.execute_forward(batch, fastvideo_args)
|
|
except BaseException as error: # noqa: BLE001
|
|
thread_error["error"] = error
|
|
thread_error_traceback["traceback"] = traceback.format_exc()
|
|
|
|
thread = threading.Thread(target=execute_forward_thread)
|
|
thread.start()
|
|
latent_batch_size = _infer_latent_batch_size(batch)
|
|
samples = torch.empty(
|
|
(latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width),
|
|
device='cpu',
|
|
pin_memory=fastvideo_args.pin_cpu_memory)
|
|
thread.join()
|
|
|
|
if thread_error["error"] is not None:
|
|
raise RuntimeError("Forward execution thread failed.\n"
|
|
f"{thread_error_traceback['traceback']}") from thread_error["error"]
|
|
|
|
output_batch = result_container["output_batch"]
|
|
if output_batch.output is None:
|
|
raise RuntimeError("Forward execution returned no output tensor. "
|
|
"This usually means the executor/pipeline failed earlier.")
|
|
|
|
if output_batch.output.shape == samples.shape:
|
|
samples.copy_(output_batch.output)
|
|
else:
|
|
logger.warning("Output shape %s does not match expected shape %s; use slow path", output_batch.output.shape,
|
|
samples.shape)
|
|
samples = output_batch.output.cpu()
|
|
logging_info = output_batch.logging_info
|
|
|
|
gen_time = time.perf_counter() - start_time
|
|
logger.info("Generated successfully in %.2f seconds", gen_time)
|
|
|
|
# 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.permute(1, 2, 0).squeeze(-1)
|
|
x = (x * 255).to(torch.uint8)
|
|
frames.append(x.cpu().numpy())
|
|
|
|
# Save output if requested
|
|
if batch.save_video:
|
|
if self._is_image_workload():
|
|
# Image workloads (t2i, i2i, …): save the first frame as PNG.
|
|
imageio.imwrite(output_path, frames[0])
|
|
logger.info("Saved image to %s", output_path)
|
|
else:
|
|
imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4")
|
|
logger.info("Saved video to %s", output_path)
|
|
audio = output_batch.extra.get("audio")
|
|
audio_sample_rate = output_batch.extra.get("audio_sample_rate")
|
|
if (audio is not None and audio_sample_rate is not None
|
|
and not self._mux_audio(output_path, audio, audio_sample_rate)):
|
|
logger.warning("Audio mux failed; saved video without audio.")
|
|
|
|
result: dict[str, Any] = {
|
|
"prompts": prompt,
|
|
"samples": samples if batch.return_frames else None,
|
|
"frames": frames if batch.return_frames else None,
|
|
"audio": output_batch.extra.get("audio") if batch.return_frames else None,
|
|
"size": (target_height, target_width, batch.num_frames),
|
|
"generation_time": gen_time,
|
|
"logging_info": logging_info,
|
|
"trajectory": output_batch.trajectory_latents,
|
|
"trajectory_timesteps": output_batch.trajectory_timesteps,
|
|
"trajectory_decoded": output_batch.trajectory_decoded,
|
|
"video_path": output_path if batch.save_video else None,
|
|
"peak_memory_mb": output_batch.extra.get("peak_memory_mb"),
|
|
}
|
|
|
|
return result
|
|
|
|
@staticmethod
|
|
def _mux_audio(
|
|
video_path: str,
|
|
audio: torch.Tensor | np.ndarray,
|
|
sample_rate: int,
|
|
) -> bool:
|
|
"""Mux audio into video using PyAV."""
|
|
try:
|
|
import av
|
|
except ImportError:
|
|
logger.warning("PyAV not installed; cannot mux audio. "
|
|
"Install with: pip install av")
|
|
return False
|
|
|
|
if torch.is_tensor(audio):
|
|
audio_np = audio.detach().cpu().float().numpy()
|
|
else:
|
|
audio_np = np.asarray(audio, dtype=np.float32)
|
|
|
|
if audio_np.ndim == 1:
|
|
audio_np = audio_np[:, None]
|
|
elif audio_np.ndim == 2:
|
|
if audio_np.shape[0] <= 8 and audio_np.shape[1] > audio_np.shape[0]:
|
|
audio_np = audio_np.T
|
|
else:
|
|
logger.warning("Unexpected audio shape %s; skipping mux.", audio_np.shape)
|
|
return False
|
|
|
|
audio_np = np.clip(audio_np, -1.0, 1.0)
|
|
audio_int16 = (audio_np * 32767.0).astype(np.int16)
|
|
num_channels = audio_int16.shape[1]
|
|
layout = "stereo" if num_channels == 2 else "mono"
|
|
|
|
try:
|
|
import wave
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
out_path = os.path.join(tmpdir, "muxed.mp4")
|
|
wav_path = os.path.join(tmpdir, "audio.wav")
|
|
|
|
# Write audio to WAV file
|
|
with wave.open(wav_path, "wb") as wav_file:
|
|
wav_file.setnchannels(num_channels)
|
|
wav_file.setsampwidth(2)
|
|
wav_file.setframerate(sample_rate)
|
|
wav_file.writeframes(audio_int16.tobytes())
|
|
|
|
# Open input video and audio
|
|
input_video = av.open(video_path)
|
|
input_audio = av.open(wav_path)
|
|
|
|
# Create output with both streams
|
|
output = av.open(out_path, mode="w")
|
|
|
|
# Add video stream (copy codec from input)
|
|
in_video_stream = input_video.streams.video[0]
|
|
out_video_stream = output.add_stream(
|
|
codec_name=in_video_stream.codec_context.name,
|
|
rate=in_video_stream.average_rate,
|
|
)
|
|
out_video_stream.width = in_video_stream.width
|
|
out_video_stream.height = in_video_stream.height
|
|
out_video_stream.pix_fmt = in_video_stream.pix_fmt
|
|
|
|
# Add audio stream (AAC)
|
|
out_audio_stream = output.add_stream("aac", rate=sample_rate)
|
|
out_audio_stream.layout = layout
|
|
|
|
# Remux video (decode and re-encode to be safe)
|
|
for frame in input_video.decode(video=0):
|
|
for packet in out_video_stream.encode(frame):
|
|
output.mux(packet)
|
|
for packet in out_video_stream.encode():
|
|
output.mux(packet)
|
|
|
|
# Encode audio
|
|
for frame in input_audio.decode(audio=0):
|
|
frame.pts = None # Let encoder assign PTS
|
|
for packet in out_audio_stream.encode(frame):
|
|
output.mux(packet)
|
|
for packet in out_audio_stream.encode():
|
|
output.mux(packet)
|
|
|
|
input_video.close()
|
|
input_audio.close()
|
|
output.close()
|
|
shutil.move(out_path, video_path)
|
|
return True
|
|
except Exception as e:
|
|
logger.warning("Audio mux failed: %s", e)
|
|
return False
|
|
|
|
def set_lora_adapter(self, lora_nickname: str, lora_path: str | None = None) -> None:
|
|
self.executor.set_lora_adapter(lora_nickname, lora_path)
|
|
|
|
def unmerge_lora_weights(self) -> None:
|
|
"""
|
|
Use unmerged weights for inference to produce videos that align with
|
|
validation videos generated during training.
|
|
"""
|
|
self.executor.unmerge_lora_weights()
|
|
|
|
def merge_lora_weights(self) -> None:
|
|
self.executor.merge_lora_weights()
|
|
|
|
def shutdown(self):
|
|
"""
|
|
Shutdown the video generator.
|
|
"""
|
|
self.executor.shutdown()
|
|
del self.executor
|