nvidia--tensorrt
wehub-resource-sync
129 行
4.9 KiB
Python
129 行
4.9 KiB
Python
#
|
||
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||
# SPDX-License-Identifier: Apache-2.0
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
#
|
||
|
||
# Configure dependencies before any external imports
|
||
from demo_diffusion import deps
|
||
deps.configure("cosmos")
|
||
|
||
import argparse
|
||
|
||
from cuda.bindings import runtime as cudart
|
||
|
||
from demo_diffusion import dd_argparse
|
||
from demo_diffusion import pipeline as pipeline_module
|
||
|
||
|
||
def parseArgs():
|
||
parser = argparse.ArgumentParser(description="Options for Wan 2.2 Txt2Vid Demo", conflict_handler='resolve')
|
||
parser = dd_argparse.add_arguments(parser)
|
||
parser.add_argument('--version', type=str, default="wan2.2-t2v-a14b", help="Version of Wan")
|
||
parser.add_argument('--guidance-scale', type=float, default=4.0, help="Guidance scale for high-noise stage (Wan default: 4.0)")
|
||
parser.add_argument('--guidance-scale-2', type=float, default=3.0, help="Guidance scale for low-noise stage (Wan default: 3.0)")
|
||
parser.add_argument('--boundary-ratio', type=float, default=0.875, help="Boundary ratio for two-stage denoising (default: 0.875)")
|
||
parser.add_argument('--denoising-steps', type=int, default=40, help="Number of denoising steps (Wan default: 40)")
|
||
parser.add_argument('--num-warmup-runs', type=int, default=1, help="Number of warmup runs before benchmarking")
|
||
parser.add_argument(
|
||
'--negative-prompt',
|
||
nargs='*',
|
||
default= (
|
||
"vivid colors, overexposed, static, blurry details, subtitles, style, "
|
||
"work of art, painting, picture, still, overall grayish, worst quality, "
|
||
"low quality, JPEG artifacts, ugly, deformed, extra fingers, poorly drawn hands, "
|
||
"poorly drawn face, deformed, disfigured, deformed limbs, fused fingers, "
|
||
"static image, cluttered background, three legs, many people in the background, "
|
||
"walking backwards"
|
||
),
|
||
help="Negative prompt (Wan team default, English translation)"
|
||
)
|
||
parser.add_argument(
|
||
"--onnx-opset",
|
||
type=int,
|
||
default=23,
|
||
choices=range(7, 24),
|
||
help="Select ONNX opset version to target for exported models",
|
||
)
|
||
return parser.parse_args()
|
||
|
||
|
||
def process_demo_args(args):
|
||
args.height = 720
|
||
args.width = 1280
|
||
args.num_frames = 81
|
||
args.max_sequence_length = 512
|
||
|
||
# require static batch = 1
|
||
if args.batch_size > 1:
|
||
raise ValueError(f"Batch size {args.batch_size} is larger than allowed (max=1 for Wan).")
|
||
args.batch_size = 1
|
||
args.build_static_batch = True
|
||
args.build_dynamic_shape = False
|
||
|
||
print(f"[I] Building Wan 2.2 T2V with fixed resolution: {args.height}×{args.width}, {args.num_frames} frames")
|
||
|
||
negative_prompt = args.negative_prompt
|
||
if isinstance(negative_prompt, list):
|
||
negative_prompt = ' '.join(negative_prompt) if negative_prompt else ""
|
||
|
||
kwargs_run_demo = {
|
||
'prompt': args.prompt,
|
||
'height': args.height,
|
||
'width': args.width,
|
||
'num_frames': args.num_frames,
|
||
'batch_size': args.batch_size,
|
||
'batch_count': args.batch_count,
|
||
'num_warmup_runs': args.num_warmup_runs,
|
||
'use_cuda_graph': args.use_cuda_graph,
|
||
'negative_prompt': negative_prompt,
|
||
'num_inference_steps': args.denoising_steps,
|
||
}
|
||
|
||
return kwargs_run_demo
|
||
|
||
|
||
if __name__ == "__main__":
|
||
print("[I] Initializing Wan 2.2 txt2vid demo using TensorRT")
|
||
args = parseArgs()
|
||
|
||
kwargs_run_demo = process_demo_args(args)
|
||
|
||
_, kwargs_load_engine, _ = dd_argparse.process_pipeline_args(args)
|
||
|
||
# Initialize demo
|
||
demo = pipeline_module.WanPipeline.FromArgs(args, pipeline_type=pipeline_module.PIPELINE_TYPE.TXT2VID)
|
||
|
||
# Load TensorRT engines and pytorch modules
|
||
demo.load_engines(
|
||
framework_model_dir=args.framework_model_dir,
|
||
**kwargs_load_engine,
|
||
)
|
||
|
||
# In low-vram mode we allocate the required device memory individually before each model is run
|
||
if args.low_vram:
|
||
demo.device_memory_sizes = demo.get_device_memory_sizes()
|
||
else:
|
||
_, shared_device_memory = cudart.cudaMalloc(demo.calculate_max_device_memory())
|
||
demo.activate_engines(shared_device_memory)
|
||
|
||
# Load resources
|
||
demo.load_resources(args.height, args.width, args.batch_size, args.seed)
|
||
|
||
# Run inference
|
||
demo.run(**kwargs_run_demo)
|
||
|
||
demo.teardown()
|
||
|