| """Minimal inference with the OpenVINO INT4 FLUX.1-schnell pipeline. |
| |
| Usage: |
| python inference_int4_flux.py --prompt "a cat" --output out.png |
| """ |
|
|
| import argparse |
| import time |
|
|
| import torch |
| from optimum.intel import OVFluxPipeline |
|
|
| MODEL_PATH = "/home/user/app/flux-schnell-ov-int4" |
|
|
|
|
| def load_pipeline(model_path: str = MODEL_PATH, device: str = "CPU"): |
| t0 = time.perf_counter() |
| pipe = OVFluxPipeline.from_pretrained(model_path, compile=True, device=device) |
| load_s = time.perf_counter() - t0 |
| return pipe, load_s |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model_path", type=str, default=MODEL_PATH) |
| parser.add_argument("--prompt", type=str, default="A cinematic photo of a mountain lake at sunrise") |
| parser.add_argument("--negative_prompt", type=str, default="") |
| parser.add_argument("--output", type=str, default="output.png") |
| parser.add_argument("--width", type=int, default=1024) |
| parser.add_argument("--height", type=int, default=1024) |
| parser.add_argument("--steps", type=int, default=4) |
| parser.add_argument("--guidance_scale", type=float, default=0.0) |
| parser.add_argument("--max_sequence_length", type=int, default=256) |
| parser.add_argument("--seed", type=int, default=42) |
| args = parser.parse_args() |
|
|
| pipe, load_s = load_pipeline(args.model_path) |
| print(f"load+compile: {load_s:.2f}s") |
|
|
| generator = torch.Generator(device="cpu").manual_seed(args.seed) |
|
|
| t0 = time.perf_counter() |
| result = pipe( |
| prompt=args.prompt, |
| negative_prompt=args.negative_prompt, |
| width=args.width, |
| height=args.height, |
| num_inference_steps=args.steps, |
| guidance_scale=args.guidance_scale, |
| max_sequence_length=args.max_sequence_length, |
| generator=generator, |
| ) |
| elapsed = time.perf_counter() - t0 |
|
|
| result.images[0].save(args.output) |
| print(f"generated in {elapsed:.2f}s ({elapsed / args.steps:.2f}s/step) -> {args.output}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
