File size: 2,155 Bytes
926d100 ae3512a 926d100 ae3512a 926d100 b1b85b3 926d100 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | #!/usr/bin/env python3
"""Turbo-Image-2.1 OrbitQuant W4A4: text-to-image and editing in 6 steps.
python run_inference.py --prompt "A red fox in a snowy forest at dusk" --output fox.png
python run_inference.py --image photo.png --prompt "Replace the background of <image1> with a beach" --output edit.png
python run_inference.py --prompt "..." --steps 8 --output text.png
"""
import argparse
import os
os.environ.setdefault("ORBITQUANT_STRICT_PACKED", "1")
import orbitquant
import torch
from diffusers import QwenImage21Pipeline
from PIL import Image
SIGMAS = {
4: [1.0, 0.75, 0.5, 0.25],
5: [1.0, 0.875, 0.75, 0.5, 0.25],
6: [1.0, 0.9375, 0.875, 0.75, 0.5, 0.25],
7: [1.0, 0.9583, 0.9167, 0.875, 0.75, 0.5, 0.25],
8: [1.0, 0.9375, 0.875, 0.75, 0.625, 0.5, 0.25, 0.125],
}
def main():
ap = argparse.ArgumentParser(description="Turbo-Image-2.1 OrbitQuant W4A4")
ap.add_argument("--model", default="WaveCut/Turbo-Image-2.1-OrbitQuant-W4A4")
ap.add_argument("--prompt", required=True)
ap.add_argument("--image", nargs="*", default=[])
ap.add_argument("--output", default="out.png")
ap.add_argument("--width", type=int)
ap.add_argument("--height", type=int)
ap.add_argument("--size", type=int, default=1024)
ap.add_argument("--steps", type=int, default=6, choices=sorted(SIGMAS))
ap.add_argument("--seed", type=int, default=0)
args = ap.parse_args()
pipe = QwenImage21Pipeline.from_pretrained(args.model, dtype=torch.float16).to("cuda")
pipe.set_progress_bar_config(disable=True)
call = dict(prompt=args.prompt, num_inference_steps=args.steps, sigmas=SIGMAS[args.steps],
output_resolution=args.size, generator=torch.Generator("cuda").manual_seed(args.seed))
if args.image:
call["image"] = [Image.open(p) for p in args.image]
if args.width and args.height:
call.update(width=args.width, height=args.height)
elif not args.image:
call.update(width=args.size, height=args.size)
image = pipe(**call).images[0]
image.save(args.output)
print(f"{args.output} {image.size}")
if __name__ == "__main__":
main()
|