[Misc] Remove InferenceEngine (#455)
This commit is contained in:
@@ -79,7 +79,6 @@ def main():
|
||||
sampling_param = SamplingParam.from_pretrained(model_name)
|
||||
sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
|
||||
sampling_param.num_frames = 107
|
||||
sampling_param.image_strength = 0.8 # How much to preserve the original image (0-1)
|
||||
|
||||
# Generate video based on the image
|
||||
prompt = "A photograph coming to life with gentle movement"
|
||||
|
||||
@@ -1,207 +0,0 @@
|
||||
# type: ignore
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Inference module for diffusion models.
|
||||
|
||||
This module provides classes and functions for running inference with diffusion models.
|
||||
"""
|
||||
|
||||
import time
|
||||
from typing import Any, Dict
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.v1.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.v1.logger import init_logger
|
||||
from fastvideo.v1.pipelines import (ComposedPipelineBase, ForwardBatch,
|
||||
build_pipeline)
|
||||
# TODO(will): remove, check if this is hunyuan specific
|
||||
from fastvideo.v1.utils import align_to
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class InferenceEngine:
|
||||
"""
|
||||
Engine for running inference with diffusion models.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pipeline: ComposedPipelineBase,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
):
|
||||
"""
|
||||
Initialize the inference engine.
|
||||
|
||||
Args:
|
||||
pipeline: The pipeline to use for inference.
|
||||
fastvideo_args: The inference arguments.
|
||||
default_negative_prompt: The default negative prompt to use.
|
||||
"""
|
||||
self.pipeline = pipeline
|
||||
self.fastvideo_args = fastvideo_args
|
||||
|
||||
@classmethod
|
||||
def create_engine(
|
||||
cls,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
) -> "InferenceEngine":
|
||||
"""
|
||||
Create an inference engine with the specified arguments.
|
||||
|
||||
Args:
|
||||
fastvideo_args: The inference arguments.
|
||||
model_loader_cls: The model loader class to use. If None, it will be
|
||||
determined from the model type.
|
||||
pipeline_type: The type of pipeline to create. If None, it will be
|
||||
determined from the model type.
|
||||
|
||||
Returns:
|
||||
The created inference engine.
|
||||
|
||||
Raises:
|
||||
ValueError: If the model type is not recognized or if the pipeline type
|
||||
is not recognized.
|
||||
"""
|
||||
|
||||
logger.info("Building pipeline...")
|
||||
|
||||
# TODO(will): I don't really like this api.
|
||||
# it should be something closer to pipeline_cls.from_pretrained(...)
|
||||
# this way for training we can just do pipeline_cls.from_pretrained(
|
||||
# checkpoint_path) and have it handle everything.
|
||||
# TODO(Peiyuan): Then maybe we should only pass in model path and device, not the entire inference args?
|
||||
pipeline = build_pipeline(fastvideo_args)
|
||||
logger.info("Pipeline Ready")
|
||||
|
||||
# Create the inference engine
|
||||
return cls(pipeline, fastvideo_args)
|
||||
|
||||
def run(
|
||||
self,
|
||||
prompt: str,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Run inference with the pipeline.
|
||||
|
||||
Args:
|
||||
prompt: The prompt to use for generation.
|
||||
negative_prompt: The negative prompt to use. If None, the default will be used.
|
||||
seed: The random seed to use. If None, a random seed will be used.
|
||||
**kwargs: Additional arguments to pass to the pipeline.
|
||||
|
||||
Returns:
|
||||
A dictionary containing the generated videos and metadata.
|
||||
"""
|
||||
out_dict: Dict[str, Any] = dict()
|
||||
|
||||
num_videos_per_prompt = fastvideo_args.num_videos
|
||||
seed = fastvideo_args.seed
|
||||
height = fastvideo_args.height
|
||||
width = fastvideo_args.width
|
||||
video_length = fastvideo_args.num_frames
|
||||
negative_prompt = fastvideo_args.neg_prompt
|
||||
infer_steps = fastvideo_args.num_inference_steps
|
||||
guidance_scale = fastvideo_args.guidance_scale
|
||||
flow_shift = fastvideo_args.flow_shift
|
||||
embedded_guidance_scale = fastvideo_args.embedded_cfg_scale
|
||||
image_path = fastvideo_args.image_path
|
||||
|
||||
# ========================================================================
|
||||
# Arguments: target_width, target_height, target_video_length
|
||||
# ========================================================================
|
||||
if width <= 0 or height <= 0 or video_length <= 0:
|
||||
raise ValueError(
|
||||
f"`height` and `width` and `video_length` must be positive integers, got height={height}, width={width}, video_length={video_length}"
|
||||
)
|
||||
if (video_length - 1) % 4 != 0:
|
||||
raise ValueError(
|
||||
f"`video_length-1` must be a multiple of 4, got {video_length}")
|
||||
|
||||
target_height = align_to(height, 16)
|
||||
target_width = align_to(width, 16)
|
||||
target_video_length = video_length
|
||||
|
||||
out_dict["size"] = (target_height, target_width, target_video_length)
|
||||
|
||||
# ========================================================================
|
||||
# Arguments: prompt, new_prompt, negative_prompt
|
||||
# ========================================================================
|
||||
if not isinstance(prompt, str):
|
||||
raise TypeError(
|
||||
f"`prompt` must be a string, but got {type(prompt)}")
|
||||
prompt = prompt.strip()
|
||||
|
||||
# negative prompt
|
||||
if negative_prompt is not None:
|
||||
negative_prompt = negative_prompt.strip()
|
||||
|
||||
# TODO(PY): move to hunyuan stage
|
||||
latents_size = [(video_length - 1) // 4 + 1, height // 8, width // 8]
|
||||
n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
|
||||
|
||||
# ========================================================================
|
||||
# Print infer args
|
||||
# ========================================================================
|
||||
debug_str = f"""
|
||||
height: {target_height}
|
||||
width: {target_width}
|
||||
video_length: {target_video_length}
|
||||
prompt: {prompt}
|
||||
neg_prompt: {negative_prompt}
|
||||
seed: {seed}
|
||||
infer_steps: {infer_steps}
|
||||
num_videos_per_prompt: {num_videos_per_prompt}
|
||||
guidance_scale: {guidance_scale}
|
||||
n_tokens: {n_tokens}
|
||||
flow_shift: {flow_shift}
|
||||
embedded_guidance_scale: {embedded_guidance_scale}"""
|
||||
logger.info(debug_str)
|
||||
# return
|
||||
# sp_group = get_sp_group()
|
||||
# local_rank = sp_group.rank
|
||||
device = torch.device(fastvideo_args.device_str)
|
||||
batch = ForwardBatch(
|
||||
image_path=image_path,
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
num_videos_per_prompt=num_videos_per_prompt,
|
||||
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,
|
||||
# generator=generator,
|
||||
eta=0.0,
|
||||
n_tokens=n_tokens,
|
||||
data_type="video" if fastvideo_args.num_frames > 1 else "image",
|
||||
device=device,
|
||||
extra={}, # Any additional parameters
|
||||
)
|
||||
|
||||
print('===============================================')
|
||||
print(batch)
|
||||
print('===============================================')
|
||||
print('===============================================')
|
||||
print(fastvideo_args)
|
||||
|
||||
# ========================================================================
|
||||
# Pipeline inference
|
||||
# ========================================================================
|
||||
start_time = time.perf_counter()
|
||||
samples = self.pipeline.forward(
|
||||
batch=batch,
|
||||
fastvideo_args=fastvideo_args,
|
||||
).output
|
||||
# TODO(will): fix and move to hunyuan stage
|
||||
# out_dict["seeds"] = batch.seeds
|
||||
out_dict["samples"] = samples
|
||||
out_dict["prompts"] = prompt
|
||||
|
||||
gen_time = time.perf_counter() - start_time
|
||||
logger.info("Success, time: %s", gen_time)
|
||||
|
||||
return out_dict
|
||||
@@ -1,77 +0,0 @@
|
||||
# type: ignore
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
import imageio
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision
|
||||
from einops import rearrange
|
||||
|
||||
from fastvideo.v1.distributed import (init_distributed_environment,
|
||||
initialize_model_parallel)
|
||||
from fastvideo.v1.fastvideo_args import FastVideoArgs, prepare_fastvideo_args
|
||||
# Fix the import path
|
||||
from fastvideo.v1.inference_engine import InferenceEngine
|
||||
|
||||
|
||||
def initialize_distributed_and_parallelism(fastvideo_args: FastVideoArgs):
|
||||
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)
|
||||
assert fastvideo_args.sp_size is not None
|
||||
assert fastvideo_args.tp_size is not None
|
||||
initialize_model_parallel(
|
||||
sequence_model_parallel_size=fastvideo_args.sp_size,
|
||||
tensor_model_parallel_size=fastvideo_args.tp_size,
|
||||
)
|
||||
|
||||
|
||||
def main(fastvideo_args: FastVideoArgs):
|
||||
initialize_distributed_and_parallelism(fastvideo_args)
|
||||
engine = InferenceEngine.create_engine(fastvideo_args, )
|
||||
|
||||
if fastvideo_args.prompt_path is not None:
|
||||
with open(fastvideo_args.prompt_path) as f:
|
||||
prompts = [line.strip() for line in f.readlines()]
|
||||
else:
|
||||
if fastvideo_args.prompt is None:
|
||||
raise ValueError("prompt or prompt_path is required")
|
||||
prompts = [fastvideo_args.prompt]
|
||||
|
||||
# Process each prompt
|
||||
for prompt in prompts:
|
||||
outputs = engine.run(
|
||||
prompt=prompt,
|
||||
fastvideo_args=fastvideo_args,
|
||||
)
|
||||
|
||||
# Process outputs
|
||||
videos = rearrange(outputs["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
|
||||
os.makedirs(os.path.dirname(fastvideo_args.output_path), exist_ok=True)
|
||||
imageio.mimsave(os.path.join(fastvideo_args.output_path,
|
||||
f"{prompt[:100]}.mp4"),
|
||||
frames,
|
||||
fps=fastvideo_args.fps)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
fastvideo_args = prepare_fastvideo_args(sys.argv[1:])
|
||||
main(fastvideo_args)
|
||||
Reference in New Issue
Block a user