File size: 4,692 Bytes
aad6de7 | 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 | """Reference BF16 reconstruction of source-bound AXQ NVFP4 diffusion storage."""
import hashlib
import importlib.util
import sys
import json
from pathlib import Path, PurePosixPath
def digest(path):
value=hashlib.sha256()
with path.open('rb') as stream:
while data:=stream.read(8*1024*1024): value.update(data)
return value.hexdigest()
def verify(directory):
root=Path(directory).resolve()
manifest=json.loads((root/'axq-diffusion-manifest.json').read_text())
schema_path=Path(__file__).with_name('axq_diffusion_schema.py')
module_spec=importlib.util.spec_from_file_location('axq_pack_diffusion_schema',schema_path)
schema=importlib.util.module_from_spec(module_spec);sys.modules[module_spec.name]=schema;module_spec.loader.exec_module(schema)
schema.DiffusionManifest.model_validate(manifest)
if manifest['schema_version']!='axq-diffusion.reference/1' or manifest['format']!='nvfp4':
raise ValueError('Unsupported AXQ diffusion storage')
for name,entry in manifest['files'].items():
relative=PurePosixPath(name)
if relative.is_absolute() or '..' in relative.parts: raise ValueError('Unsafe manifest member')
path=root/name
if not path.is_file() or path.is_symlink() or not path.resolve().is_relative_to(root):
raise ValueError('Use a regular-file snapshot_download local_dir; symlinks are not accepted')
if path.stat().st_size!=entry['size_bytes'] or digest(path)!=entry['sha256']:
raise ValueError('Pack digest mismatch: '+name)
return root,manifest
def load_transformer(directory, component='transformer', *, verify_files=True):
import torch
from safetensors.torch import load_file
from axquant.nvfp4 import Nvfp4Tensor,dequantize_nvfp4
import diffusers
from accelerate import init_empty_weights
root=Path(directory).resolve()
root,manifest=verify(root) if verify_files else (root,json.loads((root/'axq-diffusion-manifest.json').read_text()))
if component not in manifest['components']: raise ValueError('Component is not packaged')
spec=manifest['components'][component]
allowed={'QwenImage21Transformer2DModel','MiniMaxH3Transformer3DModel'}
if spec['class_name'] not in allowed: raise ValueError('Unsupported transformer class')
config=json.loads((root/component/'config.json').read_text())
cls=getattr(diffusers,spec['class_name'])
with init_empty_weights(include_buffers=False): model=cls.from_config(config)
expected=set(model.state_dict())
if set(spec['tensors'])!=expected: raise ValueError('Model state does not match pinned component inventory')
state={}
for filename in sorted(set(x['file'] for x in spec['tensors'].values())):
stored=load_file(str(root/filename),device='cpu')
consumed=set()
for name,entry in spec['tensors'].items():
if entry['file']!=filename: continue
if entry['method']=='preserve':
value=stored[name];consumed.add(name)
elif entry['method']=='nvfp4':
keys=[name+'.axq_packed',name+'.axq_scales']
consumed.update(keys)
shape=tuple(entry['shape'])
value=torch.empty(shape,dtype=torch.bfloat16)
for start in range(0,shape[0],128):
stop=min(start+128,shape[0])
decoded=dequantize_nvfp4(Nvfp4Tensor(stored[keys[0]][start:stop].numpy(),stored[keys[1]][start:stop].numpy(),entry['global_scale'],(stop-start,shape[1]),'numpy-reference'))
value[start:stop]=torch.from_numpy(decoded).to(torch.bfloat16)
else: raise ValueError('Unknown tensor method')
if list(value.shape)!=entry['shape']: raise ValueError('Shape mismatch')
if str(value.dtype).removeprefix('torch.')!=entry['dtype']: raise ValueError('Unexpected tensor dtype')
state[name]=value
if consumed!=set(stored): raise ValueError('Unconsumed packed tensor keys')
model.load_state_dict(state,strict=True,assign=True)
del state
model.eval()
if any(x.is_meta for x in model.parameters()) or any(x.is_meta for x in model.buffers()): raise ValueError('Model has unloaded tensors')
return model
def load_qwen_image_pipeline(directory):
import torch
from diffusers import QwenImage21Pipeline
root,manifest=verify(directory)
if manifest['source_model']!='Qwen/Qwen-Image-2.1-Turbo': raise ValueError('Wrong pipeline family')
model=load_transformer(root,verify_files=False)
return QwenImage21Pipeline.from_pretrained(str(root),transformer=model,torch_dtype=torch.bfloat16,local_files_only=True)
|