#!/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 , , ... python example.py --image scene.png ref.png \ --prompt "Replace the woman in with the woman from ; keep 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 = , , ... (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 "1 the prompt must reference , , ...", 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()