| |
| """Generate one image with the INT4 OpenVINO build of stable-diffusion-3.5-medium. |
| |
| Inference parameters follow the ORIGINAL model card, not the FLUX.2 reference: |
| |
| stabilityai/stable-diffusion-3.5-medium |
| 1024x1024, num_inference_steps=40, guidance_scale=4.5 |
| |
| (FLUX.2-klein-4B-OpenVINO-INT4 uses 4 steps / guidance 1.0 because FLUX.2 is a |
| 4-step distilled model. SD3.5-medium is a full flow-matching model, so the step |
| count is 10x higher. The scheduler shift of 3.0 comes from the repo's |
| scheduler/scheduler_config.json and is picked up automatically.) |
| |
| Usage |
| ----- |
| python inference_int4.py --prompt "Astronaut in a jungle" --seed 43 --output out.png |
| python inference_int4.py --prompt "..." --steps 28 # faster, lower quality |
| python inference_int4.py --prompt "..." --resize 512 # cheap preview |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import os |
| import time |
| from pathlib import Path |
|
|
| REPO_ID = "HelloSun/SD3.5-medium-OpenVINO-INT4" |
|
|
| |
| DEFAULT_STEPS = 40 |
| DEFAULT_GUIDANCE = 4.5 |
| DEFAULT_SIZE = 1024 |
|
|
|
|
| def load_pipeline(repo_id: str, device: str, compile_model: bool): |
| import torch |
| from optimum.intel import OVDiffusionPipeline |
|
|
| print(f"loading {repo_id} ...", flush=True) |
| t0 = time.perf_counter() |
| pipe = OVDiffusionPipeline.from_pretrained(repo_id, device=device, compile=compile_model) |
| print(f"load+compile {time.perf_counter() - t0:.2f}s", flush=True) |
| return pipe, torch |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("--prompt", required=True) |
| parser.add_argument("--negative-prompt", default="") |
| parser.add_argument("--model", default=REPO_ID, help="local path or Hub id") |
| parser.add_argument("--output", type=Path, default=Path("out.png")) |
| parser.add_argument("--steps", type=int, default=DEFAULT_STEPS) |
| parser.add_argument("--guidance", type=float, default=DEFAULT_GUIDANCE) |
| parser.add_argument("--width", type=int, default=DEFAULT_SIZE) |
| parser.add_argument("--height", type=int, default=DEFAULT_SIZE) |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument("--stepsize", type=int, default=8, help="scheduler shift; 8 is common for SD3 speed/quality") |
| parser.add_argument("--resize", type=int, default=None, help="also save a downscaled copy at this size") |
| parser.add_argument("--device", default="CPU") |
| parser.add_argument("--no-compile", action="store_true", help="skip OV graph compilation") |
| parser.add_argument("--num-images", type=int, default=1) |
| args = parser.parse_args() |
|
|
| pipe, torch = load_pipeline(args.model, args.device, not args.no_compile) |
|
|
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| started = time.perf_counter() |
| image = pipe( |
| prompt=args.prompt, |
| negative_prompt=args.negative_prompt, |
| num_inference_steps=args.steps, |
| guidance_scale=args.guidance, |
| width=args.width, |
| height=args.height, |
| num_images_per_prompt=args.num_images, |
| generator=torch.Generator().manual_seed(args.seed), |
| ).images[0] |
| elapsed = time.perf_counter() - started |
|
|
| image.save(args.output) |
| print(f"saved {args.output} {image.size[0]}x{image.size[1]} {elapsed:.2f}s " |
| f"({elapsed / max(args.steps, 1):.2f}s/step)", flush=True) |
|
|
| if args.resize: |
| preview = args.output.with_name(f"{args.output.stem}_{args.resize}.png") |
| try: |
| from PIL import Image as _Image |
|
|
| _resample = _Image.Resampling.LANCZOS |
| except (ImportError, AttributeError): |
| _resample = image.LANCZOS |
| image.resize((args.resize, args.resize), _resample).save(preview) |
| print(f"saved {preview}", flush=True) |
|
|
| peak_mb = 0.0 |
| try: |
| |
| for line in Path("/proc/self/status").read_text().splitlines(): |
| if line.startswith("VmHWM:"): |
| peak_mb = int(line.split()[1]) / 1024 |
| break |
| else: |
| import psutil |
|
|
| peak_mb = psutil.Process(os.getpid()).memory_info().rss / 1024**2 |
| except (OSError, ValueError, IndexError, ImportError): |
| pass |
| if peak_mb: |
| print(f"peak RSS: {peak_mb:.0f} MB", flush=True) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
