| """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) |
|
|
