| |
| """Standalone runner shipped with this Hugging Face model repository. |
| |
| Loads real TorchAO quantized DiT shards, pins the exact original SFT encoder, |
| tokenizer, scheduler, model code and VAE by immutable source commit. |
| Never silently falls back to BF16, a 4-step model or on-the-fly quantization. |
| """ |
| import argparse |
| import json |
| import os |
| from pathlib import Path |
| import sys |
|
|
| SOURCE_REPO = "OpenMOSS-Team/MOSS-Image-SFT" |
| SOURCE_SHA = "fc5de0d56fc910cbca6f6625a70f5f582b26deb1" |
| SRC_ALLOW_PATTERNS = [ |
| "LICENSE", "model_index.json", "pipeline_moss_image.py", |
| "text_encoder/*", "tokenizer/*", "vae/*", "scheduler/*", |
| ] |
|
|
| def main(): |
| ap=argparse.ArgumentParser(description=__doc__) |
| ap.add_argument("--prompt",type=str,help="Plain caption or MOSS structured layout as a JSON string") |
| ap.add_argument("--prompt-file",type=Path,help="File containing a UTF-8 literal prompt, optionally JSON") |
| ap.add_argument("--output",type=Path,default=Path("image.png")) |
| ap.add_argument("--source",type=Path,default=None, |
| help="Original MOSS-SFT components: immutable HF cache snapshot or complete local receipt") |
| ap.add_argument("--encoder-root",type=Path,default=None, |
| help="Optional local complete, SHA-verified original Qwen3.5 TorchAO FP8WO encoder") |
| ap.add_argument("--encoder-repo",type=str,default=None, |
| help="Optional Hugging Face model ID with verified FP8WO Qwen") |
| ap.add_argument("--vae-root",type=Path,default=None, |
| help="Optional original MOSS-SFT FP16/BF16 Flux2 VAE export directory") |
| ap.add_argument("--vae-repo",type=str,default=None, |
| help="Optional Hugging Face VAE precision variant repo ID") |
| ap.add_argument("--component-revision",type=str,default=None, |
| help="Optional immutable Git SHA for auxiliary HF component repos") |
| ap.add_argument("--gpu-index",type=str,default=None, |
| help="Physical nvidia-smi index, sets CUDA_VISIBLE_DEVICES before importing torch") |
| ap.add_argument("--seed",type=int,default=42) |
| ap.add_argument("--steps",type=int,default=50) |
| ap.add_argument("--cfg",type=float,default=3.0) |
| ap.add_argument("--width",type=int,default=1024) |
| ap.add_argument("--height",type=int,default=1024) |
| ap.add_argument("--experimental-high-res",action="store_true") |
| ap.add_argument("--allow-unsupported-device",action="store_true", |
| help="Try loading on unsupported hardware experimentally; kernels may still fail") |
| ap.add_argument("--offline",action="store_true",help="Do not access HF; use local source snapshot cache") |
| ap.add_argument("--verify-only",action="store_true",help="SHA-256-verify all local transformer shards and exit") |
| args=ap.parse_args() |
| if args.encoder_root is not None and args.encoder_repo is not None: |
| ap.error("Specify only one Qwen encoder source: --encoder-root or --encoder-repo") |
| if args.vae_root is not None and args.vae_repo is not None: |
| ap.error("Specify only one VAE source: --vae-root or --vae-repo") |
| if args.offline: |
| os.environ["HF_HUB_OFFLINE"]="1" |
| if args.encoder_repo: |
| from huggingface_hub import snapshot_download |
| args.encoder_root=Path(snapshot_download( |
| repo_id=args.encoder_repo,revision=args.component_revision, |
| local_files_only=args.offline)) |
| if args.vae_repo: |
| from huggingface_hub import snapshot_download |
| args.vae_root=Path(snapshot_download( |
| repo_id=args.vae_repo,revision=args.component_revision, |
| local_files_only=args.offline)) |
| if args.gpu_index is not None: |
| if not args.gpu_index.isdigit(): |
| raise SystemExit("--gpu-index must be an integer. Set before CUDA imports.") |
| os.environ["CUDA_VISIBLE_DEVICES"]=args.gpu_index |
| from run_quantized_sft import check_artifact,main as run_original |
| repo=Path(__file__).resolve().parent |
| receipt=check_artifact(repo) |
| if receipt["source_revision"] != SOURCE_SHA or not receipt["pure_sft"]: |
| raise RuntimeError("Checkpoint is not the pinned original full MOSS SFT.") |
| if args.verify_only: |
| from compact_components import inspect_qwen, inspect_vae |
| verified={"status":"SHA256_VERIFIED","variant":receipt["variant"], |
| "layers":receipt["selected_quantized_linears"], |
| "encoder":"original_BF16" if args.encoder_root is None |
| else inspect_qwen(args.encoder_root)["quantization"], |
| "vae":"original_FP32" if args.vae_root is None |
| else inspect_vae(args.vae_root)[0]["dtype"]} |
| print(json.dumps(verified,ensure_ascii=False)) |
| return |
| if bool(args.prompt) == bool(args.prompt_file): |
| ap.error("Provide exactly one of --prompt or --prompt-file") |
| if args.prompt_file is not None: |
| raw=args.prompt_file.read_text(encoding="utf-8").strip() |
| try: |
| payload=json.loads(raw) |
| prompt=json.dumps(payload,ensure_ascii=False,separators=(",",":")) if isinstance(payload,(dict,list)) else str(payload) |
| except json.JSONDecodeError: |
| prompt=raw |
| else: |
| prompt=args.prompt |
| if args.source is not None: |
| source=args.source.resolve() |
| else: |
| from huggingface_hub import snapshot_download |
| if args.offline:os.environ["HF_HUB_OFFLINE"]="1" |
| patterns=list(SRC_ALLOW_PATTERNS) |
| if args.encoder_root is not None: |
| patterns.remove("text_encoder/*") |
| patterns.append("text_encoder/config.json") |
| if args.vae_root is not None: |
| patterns.remove("vae/*") |
| patterns.append("vae/config.json") |
| print("Downloading pinned original MOSS-SFT code/tokenizer/scheduler and missing original components only; never downloading unneeded original BF16 DiT weights.",flush=True) |
| source=Path(snapshot_download( |
| repo_id=SOURCE_REPO,revision=SOURCE_SHA, |
| allow_patterns=patterns, |
| local_files_only=args.offline)) |
| argv=["run_quantized_sft.py","--source",str(source),"--variant",str(repo), |
| "--prompt",prompt,"--out",str(args.output),"--seed",str(args.seed), |
| "--steps",str(args.steps),"--cfg",str(args.cfg), |
| "--width",str(args.width),"--height",str(args.height)] |
| if args.encoder_root is not None: |
| argv.extend(["--encoder-root",str(args.encoder_root.resolve())]) |
| if args.vae_root is not None: |
| argv.extend(["--vae-root",str(args.vae_root.resolve())]) |
| if args.experimental_high_res:argv.append("--experimental-high-res") |
| if args.allow_unsupported_device:argv.append("--allow-unsupported-device") |
| old=sys.argv |
| try: |
| sys.argv=argv |
| run_original() |
| finally: |
| sys.argv=old |
|
|
| if __name__=="__main__":main() |
|
|
