File size: 6,921 Bytes
efa7c77 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | #!/usr/bin/env python3
"""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()
|