example.py

#!/usr/bin/env python3
"""zen-image-edit inference: no native text encoder (Qwen3-VL-8B, 17.5 GB) anywhere.

On the GPU: Qwen3.5-0.8B (~1.7 GB), the DiT with the adapter inside (~14.5 GB) and the VAE
(~1.4 GB, fp32).

  # text-to-image
  python example.py --prompt "a red fox in a snowy forest at dusk, cinematic, 85mm" --out fox.png

  # editing: 1..N condition images. The FIRST one is the edit target, the rest are references;
  # the prompt refers to them as <image1>, <image2>, ...
  python example.py --image scene.png ref.png \
      --prompt "Replace the woman in <image1> with the woman from <image2>; keep <image1> pose, \\
clothing and background unchanged." --out swap.png

  # a batch from a text file: one prompt per line, '#' starts a comment, blank lines are skipped
  python example.py --prompts-file prompts.txt --out gens --size 1024 --steps 30

  # non-square, and classifier-free guidance with a negative prompt
  python example.py --prompt "..." --width 1280 --height 768 --out wide.png
  python example.py --prompt "..." --negative "low quality, blurry, watermark" --cfg 3 --out cfg.png

  # scheduler A/B: the same seed and prompt rendered twice — the shipped static shift versus
  # Qwen-Image-2.1's original dynamic-shift schedule — glued side by side with labels
  python example.py --prompt "..." --scheduler-test --shift 5 --out ab.png

The pipeline is loaded once, so a batch pays the ~17 GB load a single time; every prompt uses the
same `--seed`, so a rerun reproduces the same set.
"""
import argparse
import os
import sys

import torch
from PIL import Image as PILImage
from PIL import ImageDraw, ImageFont

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from pipeline import ZenImageEditPipeline                        # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))


def read_prompts(path):
    """One prompt per line; '#' comments and blank lines are skipped."""
    with open(path, encoding="utf-8") as handle:
        lines = [line.strip() for line in handle]
    return [line for line in lines if line and not line.startswith("#")]


def _scheduler_config(pipe):
    """Scheduler config as plain values, with the service key dropped.

    A loaded config carries `_use_default_values`, and `ConfigMixin.extract_init_dict` *removes* those
    keys from a dict passed to `from_config`. Left in, every field we set afterwards (base_shift,
    max_shift, shift_terminal, ...) would be silently dropped and replaced by library defaults.
    """
    return {k: v for k, v in dict(pipe.scheduler.config).items() if k != "_use_default_values"}


def static_scheduler(pipe, shift):
    """Pipeline scheduler with a plain static shift — the shipped default (sdxs-micro uses 5.0).

    sdxs-micro's config is exactly `{shift: 5.0, use_dynamic_shifting: false}`, so `shift_terminal`
    (which stretches the schedule to end at a fixed sigma) is switched off as well.
    """
    from diffusers import FlowMatchEulerDiscreteScheduler

    config = _scheduler_config(pipe)
    config.update(use_dynamic_shifting=False, shift=shift, shift_terminal=None)
    return FlowMatchEulerDiscreteScheduler.from_config(config)


def dynamic_scheduler(pipe):
    """Qwen-Image-2.1's original schedule (dynamic shifting), kept for the `--scheduler-test` A/B."""
    from diffusers import FlowMatchEulerDiscreteScheduler

    config = _scheduler_config(pipe)
    config.update(use_dynamic_shifting=True, shift=1.0, shift_terminal=0.02, base_shift=0.5,
                  max_shift=0.9, base_image_seq_len=256, max_image_seq_len=8192,
                  time_shift_type="exponential")
    return FlowMatchEulerDiscreteScheduler.from_config(config)


def run(pipe, scheduler, args, prompt, call):
    """One generation on a fresh generator with the same seed; the scheduler is swapped for the call."""
    previous = pipe.scheduler
    pipe.scheduler = scheduler
    try:
        generator = torch.Generator(args.device).manual_seed(args.seed)
        return pipe(prompt=prompt, generator=generator, **call).images[0]
    finally:
        pipe.scheduler = previous


def label_font(size):
    for path in ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
                 "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"):
        if os.path.exists(path):
            return ImageFont.truetype(path, size)
    return ImageFont.load_default()


def side_by_side(left, right, left_label, right_label):
    """Glue two frames horizontally with a label above each."""
    gap, bar = 8, 28
    canvas = PILImage.new("RGB", (left.width + right.width + gap, left.height + bar), (16, 16, 16))
    canvas.paste(left.convert("RGB"), (0, bar))
    canvas.paste(right.convert("RGB"), (left.width + gap, bar))
    draw = ImageDraw.Draw(canvas)
    font = label_font(20)
    draw.text((8, 4), left_label, font=font, fill=(240, 240, 240))
    draw.text((left.width + gap + 8, 4), right_label, font=font, fill=(240, 240, 240))
    return canvas


def main():
    ap = argparse.ArgumentParser(description="Qwen-Image-2.1 with Qwen3.5-0.8B and the adapter inside the DiT")
    ap.add_argument("--prompt", help="a single prompt")
    ap.add_argument("--prompts-file", help="text file with one prompt per line ('#' = comment)")
    ap.add_argument("--image", nargs="*", default=[],
                    help="condition images, order = <image1>, <image2>, ... (apply to every prompt)")
    ap.add_argument("--out", help="output image, or output folder together with --prompts-file")
    ap.add_argument("--model", default=HERE, help="model folder (the layout shipped in this repo)")
    ap.add_argument("--size", type=int, default=1024, help="output_resolution (frame side, square)")
    ap.add_argument("--width", type=int, help="output width in px; overrides --size, must be a multiple of 32")
    ap.add_argument("--height", type=int, help="output height in px; overrides --size, must be a multiple of 32")
    ap.add_argument("--negative", default=None,
                    help="negative prompt; only used when --cfg > 1")
    ap.add_argument("--cfg", type=float, default=1.0,
                    help="true_cfg_scale: 1.0 = no guidance, which is how this model is meant to run")
    ap.add_argument("--scheduler-test", action="store_true",
                    help="also render Qwen-Image-2.1's original dynamic-shift schedule and glue the pair")
    ap.add_argument("--shift", type=float, default=5.0,
                    help="static shift of the shipped scheduler; sdxs-micro uses 5.0")
    ap.add_argument("--steps", type=int, default=30)
    ap.add_argument("--seed", type=int, default=1234)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--no-offload", action="store_true",
                    help="keep every component on the device (needs a large GPU)")
    args = ap.parse_args()

    if bool(args.prompt) == bool(args.prompts_file):
        ap.error("pass exactly one of --prompt or --prompts-file")
    if args.prompts_file:
        prompts = read_prompts(args.prompts_file)
        if not prompts:
            ap.error(f"no prompts in {args.prompts_file}")
    else:
        prompts = [args.prompt]
    batch = args.prompts_file is not None
    out = args.out or ("gens" if batch else "out.png")
    if batch:
        os.makedirs(out, exist_ok=True)

    condition = [PILImage.open(path) for path in args.image] or None
    if condition and len(condition) > 1 and "<image" not in prompts[0]:
        print("WARNING: with N>1 the prompt must reference <image1>, <image2>, ...", flush=True)

    pipe = ZenImageEditPipeline.from_pretrained(args.model, dtype=torch.float16)
    pipe.set_progress_bar_config(disable=True)
    # Phase-by-phase offload by default: the 14.5 GB fp16 DiT and the fp32 VAE decoder do not fit
    # an 32 GB card at the same time. Keeping everything resident needs roughly 40 GB.
    if args.device.startswith("cuda") and not args.no_offload:
        pipe.enable_model_cpu_offload(device=args.device)
    else:
        pipe.to(args.device)

    static = static_scheduler(pipe, args.shift)
    dynamic = dynamic_scheduler(pipe) if args.scheduler_test else None
    call = dict(image=condition, negative_prompt=args.negative, output_resolution=args.size,
                height=args.height, width=args.width, num_inference_steps=args.steps,
                true_cfg_scale=args.cfg, output_type="pil")

    for index, prompt in enumerate(prompts, start=1):
        image = run(pipe, static, args, prompt, call)
        if dynamic is not None:
            image = side_by_side(image, run(pipe, dynamic, args, prompt, call),
                                 f"static shift {args.shift:g} (default)", "dynamic shift (Qwen 2.1)")
        path = os.path.join(out, f"{index:04d}.png") if batch else out
        image.save(path)
        print(f"[{index}/{len(prompts)}] {path} -> {image.size}  {prompt[:70]}", flush=True)


if __name__ == "__main__":
    main()