AiArtLab/zen-image-edit

1#!/usr/bin/env python32"""zen-image-edit inference: no native text encoder (Qwen3-VL-8B, 17.5 GB) anywhere.3 4On the GPU: Qwen3.5-0.8B (~1.7 GB), the DiT with the adapter inside (~14.5 GB) and the VAE5(~1.4 GB, fp32).6 7  # text-to-image8  python example.py --prompt "a red fox in a snowy forest at dusk, cinematic, 85mm" --out fox.png9 10  # editing: 1..N condition images. The FIRST one is the edit target, the rest are references;11  # the prompt refers to them as <image1>, <image2>, ...12  python example.py --image scene.png ref.png \13      --prompt "Replace the woman in <image1> with the woman from <image2>; keep <image1> pose, \\14clothing and background unchanged." --out swap.png15 16  # a batch from a text file: one prompt per line, '#' starts a comment, blank lines are skipped17  python example.py --prompts-file prompts.txt --out gens --size 1024 --steps 3018 19  # non-square, and classifier-free guidance with a negative prompt20  python example.py --prompt "..." --width 1280 --height 768 --out wide.png21  python example.py --prompt "..." --negative "low quality, blurry, watermark" --cfg 3 --out cfg.png22 23  # scheduler A/B: the same seed and prompt rendered twice — the shipped static shift versus24  # Qwen-Image-2.1's original dynamic-shift schedule — glued side by side with labels25  python example.py --prompt "..." --scheduler-test --shift 5 --out ab.png26 27The pipeline is loaded once, so a batch pays the ~17 GB load a single time; every prompt uses the28same `--seed`, so a rerun reproduces the same set.29"""30import argparse31import os32import sys33 34import torch35from PIL import Image as PILImage36from PIL import ImageDraw, ImageFont37 38sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))39from pipeline import ZenImageEditPipeline                        # noqa: E40240 41HERE = os.path.dirname(os.path.abspath(__file__))42 43 44def read_prompts(path):45    """One prompt per line; '#' comments and blank lines are skipped."""46    with open(path, encoding="utf-8") as handle:47        lines = [line.strip() for line in handle]48    return [line for line in lines if line and not line.startswith("#")]49 50 51def _scheduler_config(pipe):52    """Scheduler config as plain values, with the service key dropped.53 54    A loaded config carries `_use_default_values`, and `ConfigMixin.extract_init_dict` *removes* those55    keys from a dict passed to `from_config`. Left in, every field we set afterwards (base_shift,56    max_shift, shift_terminal, ...) would be silently dropped and replaced by library defaults.57    """58    return {k: v for k, v in dict(pipe.scheduler.config).items() if k != "_use_default_values"}59 60 61def static_scheduler(pipe, shift):62    """Pipeline scheduler with a plain static shift — the shipped default (sdxs-micro uses 5.0).63 64    sdxs-micro's config is exactly `{shift: 5.0, use_dynamic_shifting: false}`, so `shift_terminal`65    (which stretches the schedule to end at a fixed sigma) is switched off as well.66    """67    from diffusers import FlowMatchEulerDiscreteScheduler68 69    config = _scheduler_config(pipe)70    config.update(use_dynamic_shifting=False, shift=shift, shift_terminal=None)71    return FlowMatchEulerDiscreteScheduler.from_config(config)72 73 74def dynamic_scheduler(pipe):75    """Qwen-Image-2.1's original schedule (dynamic shifting), kept for the `--scheduler-test` A/B."""76    from diffusers import FlowMatchEulerDiscreteScheduler77 78    config = _scheduler_config(pipe)79    config.update(use_dynamic_shifting=True, shift=1.0, shift_terminal=0.02, base_shift=0.5,80                  max_shift=0.9, base_image_seq_len=256, max_image_seq_len=8192,81                  time_shift_type="exponential")82    return FlowMatchEulerDiscreteScheduler.from_config(config)83 84 85def run(pipe, scheduler, args, prompt, call):86    """One generation on a fresh generator with the same seed; the scheduler is swapped for the call."""87    previous = pipe.scheduler88    pipe.scheduler = scheduler89    try:90        generator = torch.Generator(args.device).manual_seed(args.seed)91        return pipe(prompt=prompt, generator=generator, **call).images[0]92    finally:93        pipe.scheduler = previous94 95 96def label_font(size):97    for path in ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",98                 "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"):99        if os.path.exists(path):100            return ImageFont.truetype(path, size)101    return ImageFont.load_default()102 103 104def side_by_side(left, right, left_label, right_label):105    """Glue two frames horizontally with a label above each."""106    gap, bar = 8, 28107    canvas = PILImage.new("RGB", (left.width + right.width + gap, left.height + bar), (16, 16, 16))108    canvas.paste(left.convert("RGB"), (0, bar))109    canvas.paste(right.convert("RGB"), (left.width + gap, bar))110    draw = ImageDraw.Draw(canvas)111    font = label_font(20)112    draw.text((8, 4), left_label, font=font, fill=(240, 240, 240))113    draw.text((left.width + gap + 8, 4), right_label, font=font, fill=(240, 240, 240))114    return canvas115 116 117def main():118    ap = argparse.ArgumentParser(description="Qwen-Image-2.1 with Qwen3.5-0.8B and the adapter inside the DiT")119    ap.add_argument("--prompt", help="a single prompt")120    ap.add_argument("--prompts-file", help="text file with one prompt per line ('#' = comment)")121    ap.add_argument("--image", nargs="*", default=[],122                    help="condition images, order = <image1>, <image2>, ... (apply to every prompt)")123    ap.add_argument("--out", help="output image, or output folder together with --prompts-file")124    ap.add_argument("--model", default=HERE, help="model folder (the layout shipped in this repo)")125    ap.add_argument("--size", type=int, default=1024, help="output_resolution (frame side, square)")126    ap.add_argument("--width", type=int, help="output width in px; overrides --size, must be a multiple of 32")127    ap.add_argument("--height", type=int, help="output height in px; overrides --size, must be a multiple of 32")128    ap.add_argument("--negative", default=None,129                    help="negative prompt; only used when --cfg > 1")130    ap.add_argument("--cfg", type=float, default=1.0,131                    help="true_cfg_scale: 1.0 = no guidance, which is how this model is meant to run")132    ap.add_argument("--scheduler-test", action="store_true",133                    help="also render Qwen-Image-2.1's original dynamic-shift schedule and glue the pair")134    ap.add_argument("--shift", type=float, default=5.0,135                    help="static shift of the shipped scheduler; sdxs-micro uses 5.0")136    ap.add_argument("--steps", type=int, default=30)137    ap.add_argument("--seed", type=int, default=1234)138    ap.add_argument("--device", default="cuda")139    ap.add_argument("--no-offload", action="store_true",140                    help="keep every component on the device (needs a large GPU)")141    args = ap.parse_args()142 143    if bool(args.prompt) == bool(args.prompts_file):144        ap.error("pass exactly one of --prompt or --prompts-file")145    if args.prompts_file:146        prompts = read_prompts(args.prompts_file)147        if not prompts:148            ap.error(f"no prompts in {args.prompts_file}")149    else:150        prompts = [args.prompt]151    batch = args.prompts_file is not None152    out = args.out or ("gens" if batch else "out.png")153    if batch:154        os.makedirs(out, exist_ok=True)155 156    condition = [PILImage.open(path) for path in args.image] or None157    if condition and len(condition) > 1 and "<image" not in prompts[0]:158        print("WARNING: with N>1 the prompt must reference <image1>, <image2>, ...", flush=True)159 160    pipe = ZenImageEditPipeline.from_pretrained(args.model, dtype=torch.float16)161    pipe.set_progress_bar_config(disable=True)162    # Phase-by-phase offload by default: the 14.5 GB fp16 DiT and the fp32 VAE decoder do not fit163    # an 32 GB card at the same time. Keeping everything resident needs roughly 40 GB.164    if args.device.startswith("cuda") and not args.no_offload:165        pipe.enable_model_cpu_offload(device=args.device)166    else:167        pipe.to(args.device)168 169    static = static_scheduler(pipe, args.shift)170    dynamic = dynamic_scheduler(pipe) if args.scheduler_test else None171    call = dict(image=condition, negative_prompt=args.negative, output_resolution=args.size,172                height=args.height, width=args.width, num_inference_steps=args.steps,173                true_cfg_scale=args.cfg, output_type="pil")174 175    for index, prompt in enumerate(prompts, start=1):176        image = run(pipe, static, args, prompt, call)177        if dynamic is not None:178            image = side_by_side(image, run(pipe, dynamic, args, prompt, call),179                                 f"static shift {args.shift:g} (default)", "dynamic shift (Qwen 2.1)")180        path = os.path.join(out, f"{index:04d}.png") if batch else out181        image.save(path)182        print(f"[{index}/{len(prompts)}] {path} -> {image.size}  {prompt[:70]}", flush=True)183 184 185if __name__ == "__main__":186    main()187