| |
| """Run full MOSS-Image-SFT quantized WEIGHTS (zero LoRA). |
| |
| Usage: use downloaded exact original SFT components plus a quantized |
| transformer directory produced by create_full_sft_checkpoints.py. |
| MOSS teacher/default is 50 denoising steps and CFG=3; no 4-step student. |
| |
| Hardware requirements depend on the quantization format, library and device kernels. |
| """ |
| import argparse |
| import gc |
| import hashlib |
| import importlib.util |
| import json |
| import os |
| from pathlib import Path |
| import sys |
| import time |
| import torch |
|
|
| def sha256(path): |
| h=hashlib.sha256() |
| with open(path,"rb") as f: |
| for chunk in iter(lambda:f.read(4*1024*1024),b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
| def pinned_loader(root): |
| spec=importlib.util.spec_from_file_location("moss_sft_native_model",root/"pipeline_moss_image.py") |
| module=importlib.util.module_from_spec(spec) |
| sys.modules[spec.name]=module |
| spec.loader.exec_module(module) |
| return module |
|
|
| def check_artifact(variant): |
| manifest=json.loads((variant/"QUANTIZATION_RECEIPT.json").read_text()) |
| done=json.loads((variant/"COMPLETE.json").read_text()) |
| verified = done.get("verified_all_disk_weights", done.get("reload_verified", False)) |
| if not verified or done["variant"]!=manifest["variant"]: |
| raise RuntimeError("Incomplete or unverified quantized weight export") |
| if manifest["adapter"] is not None or manifest["lora"] is not None or manifest["student"] is not None: |
| raise RuntimeError("An adapter is present, expected full SFT only") |
| for fn, digest in done["checkpoint_sha256"].items(): |
| f=variant/"transformer"/fn |
| if not f.is_file() or sha256(f)!=digest: |
| raise RuntimeError(f"Missing/modified full SFT checkpoint shard: {fn}") |
| return manifest |
|
|
|
|
| def verify_pinned_source(source, expected_revision): |
| """Verify original model identity for a verified local snapshot or an immutable HF cache snapshot. |
| |
| A bare directory with unverified model files must not be accepted. The |
| official HF snapshot is identified by the exact commit in its cache path, |
| not a mutable branch/ref name. A local model_receipt.json is also valid. |
| """ |
| source=Path(source).resolve() |
| receipt=source/"model_receipt.json" |
| if receipt.is_file(): |
| data=json.loads(receipt.read_text()) |
| if data.get("revision")!=expected_revision or not data.get("complete"): |
| raise RuntimeError("Original SFT local receipt revision/completeness mismatch") |
| elif not (source.name==expected_revision and source.parent.name=="snapshots"): |
| raise RuntimeError("Unverified original components. Use HF snapshot_download with exact revision " |
| +expected_revision+" or original complete local model_receipt.json") |
| required=("pipeline_moss_image.py","text_encoder/config.json", |
| "tokenizer/tokenizer_config.json","vae/config.json","scheduler/scheduler_config.json") |
| missing=[part for part in required if not (source/part).is_file()] |
| if missing: |
| raise RuntimeError("Original SFT source components missing: "+", ".join(missing)) |
| return source |
|
|
|
|
| class StreamedQwenText: |
| """BF16 Qwen text encoder streamed via Accelerate; no LoRA.""" |
| def __init__(self,root,device="cuda",cache_size=16,encoder_root=None): |
| from transformers import AutoModel,AutoTokenizer |
| from accelerate import cpu_offload |
| from compact_components import inspect_qwen |
| self.device=torch.device(device) |
| self.tokenizer=AutoTokenizer.from_pretrained(root/"tokenizer",local_files_only=True) |
| if encoder_root is None: |
| |
| enc=AutoModel.from_pretrained(root/"text_encoder",dtype=torch.bfloat16, |
| local_files_only=True).eval() |
| self.quantization="original_bf16" |
| else: |
| |
| |
| encoder_root=Path(encoder_root).resolve() |
| info=inspect_qwen(encoder_root,verify_sha=True) |
| import torchao |
| import torchao.quantization |
| enc=AutoModel.from_pretrained(encoder_root/"text_encoder", |
| dtype=torch.bfloat16,use_safetensors=False,local_files_only=True).eval() |
| real=sum(type(m.weight).__name__=="AffineQuantizedTensor" |
| for n,m in enc.named_modules() |
| if n.startswith("language_model.") and isinstance(m,torch.nn.Linear)) |
| if real!=info["converted_linear_weights"]: |
| raise RuntimeError(f"Packed Qwen was not actually rehydrated: {real}") |
| self.quantization="original_qwen3.5_fp8_weight_only_248_linear" |
| if hasattr(enc,"visual"): |
| |
| del enc.visual |
| gc.collect() |
| |
| |
| |
| weights={k:v.detach() for k,v in enc.state_dict().items()} |
| cpu_offload(enc,execution_device=self.device,offload_buffers=True, |
| state_dict=weights) |
| del weights |
| self.enc=enc |
| self.cache={} |
| self.cache_size=cache_size |
|
|
| @torch.no_grad() |
| def encode(self,prompt): |
| if prompt in self.cache: |
| return self.cache[prompt] |
| formatted=self.tokenizer.apply_chat_template( |
| [{"role":"user","content":prompt}],tokenize=False, |
| add_generation_prompt=True,enable_thinking=True) |
| inp=self.tokenizer(formatted,return_tensors="pt",truncation=True, |
| max_length=4096).to(self.device) |
| result=self.enc(**inp,use_cache=False,return_dict=True) |
| embedding=result.last_hidden_state[:1] |
| data=(embedding,torch.ones(embedding.shape[:2], |
| device=embedding.device,dtype=torch.bool)) |
| if len(self.cache)>=self.cache_size: |
| self.cache.clear() |
| self.cache[prompt]=data |
| return data |
|
|
|
|
| def main(): |
| ap=argparse.ArgumentParser() |
| ap.add_argument("--variant",type=Path,required=True,help="folder with transformer/*.bin and COMPLETE.json") |
| ap.add_argument("--source",type=Path,required=True, |
| help="original pinned MOSS-Image-SFT, for tokenizer/text encoder/VAE/scheduler") |
| ap.add_argument("--prompt",required=True) |
| ap.add_argument("--out",type=Path,required=True) |
| ap.add_argument("--seed",type=int,default=123) |
| 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("--encoder-root",type=Path,default=None, |
| help="Optional separately packaged original Qwen3.5 FP8WO encoder; validates all four shards") |
| ap.add_argument("--vae-root",type=Path,default=None, |
| help="Optional locally verified original Flux2 VAE FP16 or BF16 component") |
| ap.add_argument("--experimental-high-res",action="store_true") |
| ap.add_argument("--allow-unsupported-device",action="store_true", |
| help="Experimental override for unsupported hardware; does NOT guarantee execution") |
| args=ap.parse_args() |
| if args.steps<1 or not 0<=args.cfg<=10: |
| raise ValueError("invalid inference steps/guidance") |
| if args.width%16 or args.height%16 or min(args.width,args.height)<512: |
| raise ValueError("width and height must be multiples of 16, at least 512") |
| if (args.height,args.width)!=(1024,1024) and not args.experimental_high_res: |
| raise ValueError("non-1024 sizes require --experimental-high-res (not quality-certified)") |
| if not torch.cuda.is_available(): |
| raise RuntimeError("This inference CLI needs a CUDA GPU") |
| gpu=torch.cuda.get_device_name(0) |
| cc=torch.cuda.get_device_capability(0) |
|
|
| args.source=args.source.resolve() |
| args.variant=args.variant.resolve() |
| receipt=check_artifact(args.variant) |
| method=receipt.get("quantization") |
| required_cc=((8,9) if method in ("fp8dynamic-row","fp8dynamic-tensor") |
| else (10,0) if method=="nvfp4" else (0,0)) |
| if cc < required_cc and not args.allow_unsupported_device: |
| raise RuntimeError(f"GPU {gpu} compute capability {cc} below minimum for {method} ({required_cc}); " |
| "use a compatible accelerator/runtime or explicit experimental override.") |
| verify_pinned_source(args.source, receipt["source_revision"]) |
| import torchao |
| import torchao.quantization |
| import torchao.prototype.mx_formats |
| from diffusers import AutoencoderKLFlux2 |
| module=pinned_loader(args.source) |
| start=time.perf_counter() |
| transformer=module.MOSSImageTransformer.from_pretrained( |
| args.variant/"transformer",torch_dtype=torch.bfloat16, |
| use_safetensors=False,local_files_only=True).eval() |
| selected=receipt.get("selected_quantized_linears", |
| receipt.get("conversion",{}).get("selected_linear_count")) |
| if selected is None or selected < 200: |
| raise RuntimeError("Quantized weight count missing or implausibly small") |
| converted=sum(type(m.weight).__name__!="Parameter" for m in transformer.modules() |
| if isinstance(m,torch.nn.Linear)) |
| if converted!=selected: |
| raise RuntimeError(f"Rehydration mismatch: {converted} != {selected}; will not fallback") |
| transformer=transformer.to("cuda") |
| print(json.dumps({"loaded_gpu":gpu, "sm":cc, "variant":receipt["variant"], |
| "selected":converted, "load_seconds":round(time.perf_counter()-start,2)}),flush=True) |
| |
| |
| text=StreamedQwenText(args.source,encoder_root=args.encoder_root) |
| positive=text.encode(args.prompt) |
| negative=text.encode("") if args.cfg!=1.0 else None |
| if args.vae_root is None: |
| vae=AutoencoderKLFlux2.from_pretrained(args.source/"vae", |
| torch_dtype=torch.float32,local_files_only=True).eval().to("cuda") |
| vae_kind="original_fp32" |
| else: |
| from compact_components import inspect_vae |
| vae_info,dtype=inspect_vae(args.vae_root,verify_sha=True) |
| vae=AutoencoderKLFlux2.from_pretrained(Path(args.vae_root)/"vae", |
| torch_dtype=dtype,local_files_only=True).eval().to("cuda") |
| vae_kind=vae_info["dtype"] |
| scheduler=module.MOSSImageScheduler.from_pretrained(args.source/"scheduler", |
| local_files_only=True) |
| pipe=module.MOSSImagePipeline(transformer,scheduler,None,None,vae) |
| generator=torch.Generator("cpu").manual_seed(args.seed) |
| kw={} |
| if negative is not None: |
| kw["negative_prompt_embeds"]=negative[0] |
| kw["negative_prompt_attention_mask"]=negative[1] |
| torch.cuda.reset_peak_memory_stats() |
| torch.cuda.synchronize() |
| t0=time.perf_counter() |
| with torch.inference_mode(): |
| im=pipe(prompt_embeds=positive[0],prompt_attention_mask=positive[1], |
| width=args.width,height=args.height, |
| num_inference_steps=args.steps,guidance_scale=args.cfg, |
| uncond_skip_layers=[22],generator=generator,**kw).images[0] |
| torch.cuda.synchronize() |
| args.out.parent.mkdir(parents=True,exist_ok=True) |
| im.save(args.out) |
| print(json.dumps({"image":str(args.out),"variant":receipt["variant"],"lora":False, |
| "encoder_precision":text.quantization,"vae_precision":vae_kind, |
| "seconds_image":round(time.perf_counter()-t0,2), |
| "peak_cuda_alloc_gib":round(torch.cuda.max_memory_allocated()/2**30,3), |
| "peak_cuda_reserved_gib":round(torch.cuda.max_memory_reserved()/2**30,3)}),flush=True) |
|
|
| if __name__=="__main__": |
| main() |
|
|
