121 lines
3.8 KiB
Python
121 lines
3.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Run full Flux2 text-to-image generation through FastVideo.
|
|
|
|
User story:
|
|
"I have a local or HF Diffusers-format full Flux2 checkpoint and want a
|
|
minimal text-to-image generation command that uses embedded guidance."
|
|
"""
|
|
import argparse
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from fastvideo import VideoGenerator
|
|
from fastvideo.api import (
|
|
ComponentConfig,
|
|
EngineConfig,
|
|
GenerationRequest,
|
|
GeneratorConfig,
|
|
OffloadConfig,
|
|
OutputConfig,
|
|
ParallelismConfig,
|
|
PipelineSelection,
|
|
SamplingConfig,
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Run full Flux2 text-to-image generation.")
|
|
parser.add_argument(
|
|
"--model-path",
|
|
default="black-forest-labs/FLUX.2-dev",
|
|
help="HF id or local diffusers-format full Flux2 weights directory.",
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
default="outputs/flux2/flux2.png",
|
|
help="Output PNG path.",
|
|
)
|
|
parser.add_argument(
|
|
"--prompt",
|
|
default="a photo of a banana on a wooden table, studio lighting",
|
|
help="Text prompt.",
|
|
)
|
|
parser.add_argument("--height", type=int, default=1024)
|
|
parser.add_argument("--width", type=int, default=1024)
|
|
parser.add_argument("--steps", type=int, default=50)
|
|
parser.add_argument("--guidance-scale", type=float, default=4.0)
|
|
parser.add_argument("--max-sequence-length", type=int, default=None)
|
|
parser.add_argument("--seed", type=int, default=0)
|
|
parser.add_argument("--num-gpus", type=int, default=1)
|
|
parser.add_argument("--tp-size", type=int, default=None)
|
|
parser.add_argument("--sp-size", type=int, default=None)
|
|
parser.add_argument(
|
|
"--backend",
|
|
default=None,
|
|
help="Set FASTVIDEO_ATTENTION_BACKEND, for example TORCH_SDPA.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if args.backend:
|
|
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = args.backend
|
|
|
|
output = Path(args.output)
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
|
|
sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
|
|
|
|
generator_config = GeneratorConfig(
|
|
model_path=args.model_path,
|
|
engine=EngineConfig(
|
|
num_gpus=args.num_gpus,
|
|
parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
|
|
use_fsdp_inference=False,
|
|
offload=OffloadConfig(
|
|
dit=False,
|
|
vae=True,
|
|
text_encoder=True,
|
|
pin_cpu_memory=False,
|
|
),
|
|
),
|
|
pipeline=PipelineSelection(
|
|
workload_type="t2i",
|
|
components=ComponentConfig(override_pipeline_cls_name="Flux2Pipeline"),
|
|
),
|
|
)
|
|
|
|
generator = VideoGenerator.from_config(generator_config)
|
|
try:
|
|
sampling = SamplingConfig(
|
|
height=args.height,
|
|
width=args.width,
|
|
num_frames=1,
|
|
fps=1,
|
|
num_inference_steps=args.steps,
|
|
guidance_scale=args.guidance_scale,
|
|
seed=args.seed,
|
|
)
|
|
extensions = {}
|
|
if args.max_sequence_length is not None:
|
|
extensions["max_sequence_length"] = args.max_sequence_length
|
|
|
|
request = GenerationRequest(
|
|
prompt=args.prompt,
|
|
sampling=sampling,
|
|
output=OutputConfig(
|
|
output_path=str(output),
|
|
save_video=True,
|
|
return_frames=False,
|
|
),
|
|
extensions=extensions,
|
|
)
|
|
generator.generate(request)
|
|
finally:
|
|
generator.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|