Co-authored-by: Aryan Kumar <156166681+aryan5v@users.noreply.github.com>
74 lines
3.0 KiB
Python
74 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Qwen-Image-2.1 generation, editing, references and transparent PNG output."""
|
|
|
|
import argparse
|
|
|
|
from fastvideo import VideoGenerator
|
|
from fastvideo.api import (
|
|
EngineConfig,
|
|
GenerationRequest,
|
|
GeneratorConfig,
|
|
InputConfig,
|
|
OffloadConfig,
|
|
OutputConfig,
|
|
ParallelismConfig,
|
|
PipelineSelection,
|
|
SamplingConfig,
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--model-path", default="Qwen/Qwen-Image-2.1")
|
|
parser.add_argument("--revision")
|
|
parser.add_argument("--prompt", default="A red ceramic teapot, isolated on a transparent background.")
|
|
parser.add_argument("--reference", action="append", default=[], help="Ordered reference image; repeat up to ten times")
|
|
parser.add_argument("--negative-prompt", default=None)
|
|
parser.add_argument("--true-cfg-scale", type=float, default=1.0)
|
|
parser.add_argument("--height", type=int, default=1024)
|
|
parser.add_argument("--width", type=int, default=1024)
|
|
parser.add_argument("--reference-resolution", type=int, default=1024)
|
|
parser.add_argument("--steps", type=int, default=40)
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument("--use-kv-cache", action=argparse.BooleanOptionalAction, default=True)
|
|
parser.add_argument("--output", default="outputs/qwen_image21/image.png")
|
|
args = parser.parse_args()
|
|
if len(args.reference) > 10:
|
|
parser.error("At most ten ordered reference images are supported")
|
|
workload = "i2i" if args.reference else "t2i"
|
|
generator = VideoGenerator.from_config(
|
|
GeneratorConfig(
|
|
model_path=args.model_path,
|
|
revision=args.revision,
|
|
engine=EngineConfig(
|
|
num_gpus=1,
|
|
parallelism=ParallelismConfig(tp_size=1, sp_size=1),
|
|
offload=OffloadConfig(dit=True, dit_layerwise=True, text_encoder=True, vae=True,
|
|
lazy_module_load=True),
|
|
),
|
|
pipeline=PipelineSelection(
|
|
workload_type=workload,
|
|
experimental={
|
|
"reference_resolution": args.reference_resolution,
|
|
"use_kv_cache": args.use_kv_cache,
|
|
},
|
|
),
|
|
))
|
|
try:
|
|
generator.generate(
|
|
GenerationRequest(
|
|
prompt=args.prompt,
|
|
negative_prompt=args.negative_prompt,
|
|
inputs=InputConfig(references=args.reference or None),
|
|
sampling=SamplingConfig(height=args.height, width=args.width, num_frames=1, fps=1,
|
|
num_inference_steps=args.steps, guidance_scale=1.0,
|
|
true_cfg_scale=args.true_cfg_scale, seed=args.seed),
|
|
output=OutputConfig(output_path=args.output, save_video=True, return_frames=False),
|
|
))
|
|
finally:
|
|
generator.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|