| |
| |
| |
| |
| |
| |
| |
| """Main entry point for exporting the Z-Image-Turbo pipeline to standalone ONNX models.""" |
|
|
| import argparse |
| import os |
| import shutil |
| import subprocess |
| import sys |
|
|
| import build_helper_models |
| import build_safety_checker |
| import build_text_encoder |
| import build_transformer |
| import build_vae_decoder |
|
|
| ALL_COMPONENTS = ("transformer", "text_encoder", "vae_decoder", "helper_models", "safety_checker") |
|
|
| |
| |
| TOKENIZER_FILES = ("merges.txt", "tokenizer.json", "tokenizer_config.json", "vocab.json") |
|
|
|
|
| def get_args(): |
| parser = argparse.ArgumentParser(description="Export the Z-Image-Turbo pipeline to ONNX.") |
| parser.add_argument("input_path", help="Path to the Z-Image-Turbo checkpoint (repo root or component subfolder).") |
| parser.add_argument( |
| "-o", "--output_dir", default=build_transformer.DEFAULT_OUTPUT_DIR, |
| help=f"Directory to write the ONNX model(s) to (under an `onnx/` subdir). Default: {build_transformer.DEFAULT_OUTPUT_DIR}", |
| ) |
| parser.add_argument( |
| "-m", "--model", default="all", |
| choices=["transformer", "vae_decoder", "helper_models", "text_encoder", "safety_checker", "all"], |
| help="Which component to build. Default: all.", |
| ) |
| parser.add_argument( |
| "--extra_options", nargs="*", default=[], |
| help="Extra key=value options passed through to the component's builder.", |
| ) |
| parser.add_argument( |
| "--safety_checker_checkpoint", default="", |
| help="Path to a local CompVis/stable-diffusion-safety-checker-compatible checkpoint " |
| "(e.g. via huggingface_hub.snapshot_download('CompVis/stable-diffusion-safety-checker', " |
| "local_dir=...)). This is a separate download from input_path. Required for " |
| "-m safety_checker; for -m all, safety_checker is skipped if this isn't given.", |
| ) |
| return parser.parse_args() |
|
|
|
|
| def build_one(model, input_path, output_dir, extra_options, safety_checker_checkpoint=""): |
| if model == "transformer": |
| transformer_input = input_path |
| if os.path.isdir(os.path.join(input_path, "transformer")): |
| transformer_input = os.path.join(input_path, "transformer") |
| return build_transformer.build( |
| transformer_input, output_dir, extra_options=build_transformer.parse_extra_options(extra_options) |
| ) |
| if model == "text_encoder": |
| text_encoder_input = input_path |
| if os.path.isdir(os.path.join(input_path, "text_encoder")): |
| text_encoder_input = os.path.join(input_path, "text_encoder") |
| return build_text_encoder.build( |
| text_encoder_input, output_dir, extra_options=build_text_encoder.parse_extra_options(extra_options) |
| ) |
| if model == "vae_decoder": |
| vae_input = input_path |
| if os.path.isdir(os.path.join(input_path, "vae")): |
| vae_input = os.path.join(input_path, "vae") |
| return build_vae_decoder.build( |
| vae_input, output_dir, extra_options=build_vae_decoder.parse_extra_options(extra_options) |
| ) |
| if model == "helper_models": |
| return build_helper_models.build( |
| input_path, output_dir, extra_options=build_helper_models.parse_extra_options(extra_options) |
| ) |
| if model == "safety_checker": |
| if not safety_checker_checkpoint: |
| raise NotImplementedError( |
| "-m safety_checker requires --safety_checker_checkpoint (a separate local " |
| "CompVis/stable-diffusion-safety-checker-compatible checkpoint, not input_path)." |
| ) |
| return build_safety_checker.build( |
| safety_checker_checkpoint, output_dir, extra_options=build_safety_checker.parse_extra_options(extra_options) |
| ) |
| raise ValueError(f"Unknown -m/--model value: {model}") |
|
|
|
|
| def resolve_tokenizer_dir(input_path): |
| text_encoder_dir = input_path |
| if os.path.isdir(os.path.join(input_path, "text_encoder")): |
| text_encoder_dir = os.path.join(input_path, "text_encoder") |
| return os.path.join(os.path.dirname(os.path.normpath(text_encoder_dir)), "tokenizer") |
|
|
|
|
| def _link_or_copy(src, dst): |
| if os.path.exists(dst): |
| os.remove(dst) |
| try: |
| os.link(src, dst) |
| except OSError: |
| shutil.copy2(src, dst) |
|
|
|
|
| def copy_tokenizer(input_path, output_dir): |
| tokenizer_dir = resolve_tokenizer_dir(input_path) |
| if not os.path.isdir(tokenizer_dir): |
| print(f"Skipping tokenizer copy: {tokenizer_dir} not found", file=sys.stderr) |
| return |
|
|
| dest_dir = os.path.join(output_dir, "tokenizer") |
| os.makedirs(dest_dir, exist_ok=True) |
| missing = [] |
| for fname in TOKENIZER_FILES: |
| src = os.path.join(tokenizer_dir, fname) |
| if os.path.isfile(src): |
| _link_or_copy(src, os.path.join(dest_dir, fname)) |
| else: |
| missing.append(fname) |
|
|
| if missing: |
| print(f"Warning: tokenizer files not found in {tokenizer_dir}: {missing}", file=sys.stderr) |
| else: |
| print(f"Copied tokenizer to {dest_dir}") |
|
|
|
|
| def build_component_subprocess(model, args): |
| cmd = [sys.executable, os.path.abspath(__file__), args.input_path, "-m", model, "-o", args.output_dir] |
| if args.extra_options: |
| cmd += ["--extra_options", *args.extra_options] |
| if args.safety_checker_checkpoint: |
| cmd += ["--safety_checker_checkpoint", args.safety_checker_checkpoint] |
| return subprocess.run(cmd).returncode |
|
|
|
|
| def build_all(args): |
| results = {} |
| for model in ALL_COMPONENTS: |
| if model == "safety_checker" and not args.safety_checker_checkpoint: |
| print(f"\nSkipping {model}: no --safety_checker_checkpoint given") |
| results[model] = "skipped" |
| continue |
| print(f"\n=== Building {model} (isolated subprocess) ===") |
| returncode = build_component_subprocess(model, args) |
| results[model] = "ok" if returncode == 0 else f"FAILED (exit {returncode})" |
|
|
| print("\n=== Summary ===") |
| for model in ALL_COMPONENTS: |
| print(f" {model:<14} {results.get(model, 'skipped')}") |
|
|
| failed = [m for m, status in results.items() if status.startswith("FAILED")] |
| return results, failed |
|
|
|
|
| def main(): |
| args = get_args() |
| if args.model == "all": |
| _results, failed = build_all(args) |
| copy_tokenizer(args.input_path, args.output_dir) |
| onnx_dir = os.path.join(args.output_dir, "onnx") |
| if failed: |
| print(f"\nFailure: {len(failed)} component(s) failed: {', '.join(failed)}", file=sys.stderr) |
| sys.exit(1) |
| print(f"\nSuccess: all exported to {onnx_dir}") |
| else: |
| onnx_dir = build_one( |
| args.model, args.input_path, args.output_dir, args.extra_options, args.safety_checker_checkpoint |
| ) |
| copy_tokenizer(args.input_path, args.output_dir) |
| print(f"\nSuccess: {args.model} exported to {onnx_dir}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
