#!/usr/bin/env python3 """Export the Qwen material adapter for inference-only distribution. The source is preserved. Optimizer/RNG/data paths are omitted, and tensor values are compared directly after saving; no content digests are computed. """ from __future__ import annotations import argparse import json import math from pathlib import Path import numpy as np import torch from safetensors import safe_open from safetensors.torch import save_file BASE_MODEL = "Qwen/Qwen-Image-2.1" ADAPTER_FORMAT = "supratex_qwen_material_adapter" SCHEMA_VERSION = 1 SAFETENSORS_DTYPES = { torch.float64: "F64", torch.float32: "F32", torch.float16: "F16", torch.bfloat16: "BF16", } INFERENCE_DEFAULTS = { "lora_rank": 32, "lora_alpha": 32, "resolution": 512, "noise_seed": 42, "lora_output_projection": False, "omit_all_valid_mask": False, "bf16_autocast": False, "attention_backend": "torch", "deterministic_algorithms": False, } def load_source(path: Path) -> dict: # Legacy trainer checkpoints include NumPy RNG arrays. Allow only their # concrete reconstruction types, without unrestricted pickle execution. from numpy.core.multiarray import _reconstruct with torch.serialization.safe_globals([_reconstruct, (_reconstruct, "numpy.core.multiarray._reconstruct"), (_reconstruct, "numpy._core.multiarray._reconstruct"), np.ndarray, np.dtype, type(np.dtype("uint32"))]): return torch.load(path, map_location="cpu", weights_only=True, mmap=True) def export(source: Path, output_dir: Path, *, format: str = "safetensors") -> dict: source, output_dir = Path(source).resolve(), Path(output_dir).resolve() if format not in {"safetensors", "pt"}: raise ValueError("format must be safetensors or pt") if output_dir.exists(): raise FileExistsError(f"select a new export directory: {output_dir}") marker = json.loads(source.with_suffix(".complete.json").read_text()) if (marker.get("status") != "complete" or marker.get("step") != 100000 or marker.get("bytes") != source.stat().st_size): raise ValueError("source requires a matching 100k completion record") state = load_source(source) if state.get("format") != "qwen_pbr_endpoint_v1" or state.get("step") != 100000: raise ValueError("export requires the Qwen 100k endpoint checkpoint") config, model = state.get("config"), state.get("model") if not isinstance(config, dict) or not isinstance(model, dict) or not model: raise ValueError("checkpoint requires config and nonempty trainable tensors") if not {"model_root", "resolution", "lora_rank", "lora_alpha", "noise_seed"} <= config.keys(): raise ValueError("checkpoint lacks required inference configuration") public_config = {key: config.get(key, default) for key, default in INFERENCE_DEFAULTS.items()} for key, value in public_config.items(): if type(value) not in {int, float, bool, str}: raise ValueError(f"invalid inference config value: {key}") if (type(public_config["resolution"]) is not int or public_config["resolution"] != 1024 or type(public_config["noise_seed"]) is not int or public_config["noise_seed"] != 42 or public_config["attention_backend"] != "torch"): raise ValueError("the released adapter requires 1024 resolution, seed 42 and torch attention") for key in ("lora_rank", "lora_alpha"): if (type(public_config[key]) not in {int, float} or not math.isfinite(public_config[key]) or public_config[key] <= 0): raise ValueError(f"invalid {key}") if type(public_config["lora_rank"]) is not int: raise ValueError("lora_rank must be a positive integer") for key in ("lora_output_projection", "omit_all_valid_mask", "bf16_autocast", "deterministic_algorithms"): if type(public_config[key]) is not bool: raise ValueError(f"{key} must be boolean") public_config["model_root"] = BASE_MODEL public_model = {} for name, tensor in model.items(): if not isinstance(name, str) or not name or "/" in name or "\\" in name: raise ValueError("trainable tensor names must be parameter names") if (not isinstance(tensor, torch.Tensor) or tensor.dtype not in SAFETENSORS_DTYPES or not torch.isfinite(tensor).all()): raise ValueError(f"invalid/nonfinite trainable tensor: {name}") # Avoid retaining storage beyond this tensor's logical contents. public_model[name] = tensor.detach().cpu().contiguous().clone() payload = {"format": state["format"], "step": state["step"], "config": public_config, "model": public_model} output_dir.mkdir(parents=True) if format == "pt": output = output_dir / "checkpoint-0100000.pt" torch.save(payload, output) restored = torch.load(output, map_location="cpu", weights_only=True, mmap=True) if restored["config"] != public_config or restored["model"].keys() != model.keys(): raise RuntimeError("export config or tensor key mismatch") for name, expected in model.items(): actual = restored["model"][name] if actual.shape != expected.shape or actual.dtype != expected.dtype or not torch.equal(actual, expected): raise RuntimeError(f"export tensor mismatch: {name}") else: output = output_dir / "adapter_model.safetensors" # This is the native PBR state (LoRA plus learned modality embeddings), # not a standard PEFT adapter. Preserve every native parameter name. metadata = {"format": "pt", "adapter_format": ADAPTER_FORMAT, "schema_version": str(SCHEMA_VERSION)} save_file(public_model, str(output), metadata=metadata) with safe_open(str(output), framework="pt", device="cpu") as restored: if restored.metadata() != metadata or set(restored.keys()) != model.keys(): raise RuntimeError("export metadata or tensor key mismatch") for name, expected in model.items(): actual = restored.get_tensor(name) if actual.shape != expected.shape or actual.dtype != expected.dtype or not torch.equal(actual, expected): raise RuntimeError(f"export tensor mismatch: {name}") adapter_config = { "format": ADAPTER_FORMAT, "schema_version": SCHEMA_VERSION, "base_model_name_or_path": BASE_MODEL, "adapter_type": "qwen_image_21_pbr_endpoint", "weight_file": output.name, "source_step": state["step"], "inference_config": {key: value for key, value in public_config.items() if key != "model_root"}, "tensors": {name: {"shape": list(tensor.shape), "dtype": SAFETENSORS_DTYPES[tensor.dtype]} for name, tensor in public_model.items()}, } report = {"status": "verified", "source_step": 100000, "base_model": BASE_MODEL, "source_bytes": source.stat().st_size, "export_bytes": output.stat().st_size, "weight_format": format, "weight_file": output.name, "tensor_count": len(model), "tensor_values_equal": True, "removed_training_state": sorted(set(state) - set(payload)), "retained_config_keys": sorted(public_config), "verification": "direct tensor key, shape, dtype and value comparison; no content digests"} (output_dir / "EXPORT_REPORT.json").write_text(json.dumps(report, indent=2) + "\n") # Publish the config/legacy marker last, after a reload and exact comparison. if format == "pt": output.with_suffix(".complete.json").write_text(json.dumps( {"status": "complete", "step": 100000, "bytes": output.stat().st_size}, indent=2) + "\n") else: config_path = output_dir / "adapter_config.json" config_path.write_text(json.dumps(adapter_config, indent=2) + "\n") if json.loads(config_path.read_text()) != adapter_config: raise RuntimeError("export adapter config mismatch") return report def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--source", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--format", choices=("safetensors", "pt"), default="safetensors", help="Native safetensors adapter (default) or legacy inference PT") args = parser.parse_args() print(json.dumps(export(args.source, args.output_dir, format=args.format), indent=2)) if __name__ == "__main__": main()