Release calibrated Image2.1 NVFP4 transformer with dynamic scaling and BF16 rank correction, native SM120 runtime, quality evidence and real-time demo
1961af5 verified | """Calibrated NVFP4 W4A4 with BF16 rank-128 correction on SM120. | |
| Built with Qwen. Quantization modifications by ProCreations, 2026. | |
| Requires the pinned FlashInfer source revision shipped in requirements. | |
| """ | |
| import json | |
| from pathlib import Path | |
| import torch | |
| from torch import nn | |
| from safetensors.torch import load_file | |
| def _linear(x:torch.Tensor,weight:torch.Tensor,sf:torch.Tensor,pre:torch.Tensor, | |
| gx:torch.Tensor,alpha:torch.Tensor,down:torch.Tensor,up:torch.Tensor)->torch.Tensor: | |
| from flashinfer.gemm import nvfp4_quantize_smooth,mm_nvfp4_svdquant | |
| from dynamic_scale import scale | |
| x=x.contiguous() | |
| gx,alpha,up=scale(x,pre,gx,alpha,up) | |
| q,sc=nvfp4_quantize_smooth(x,pre,gx,backend='cute-dsl') | |
| d=torch.mm(x,down) | |
| return mm_nvfp4_svdquant(q,weight,sc,sf,alpha,d,up,backend='cute-dsl') | |
| def _linear_fake(x,weight,sf,pre,gx,alpha,down,up): | |
| return x.new_empty((x.shape[0],weight.shape[0]),dtype=torch.bfloat16) | |
| class CalibratedNVFP4Linear(nn.Module): | |
| def __init__(self,state): | |
| super().__init__() | |
| for key in ['weight','sf','pre','gx','alpha','down','up']: | |
| self.register_buffer(key,state[key]) | |
| self.out_features,self.in_features=self.weight.shape[0],self.weight.shape[1]*2 | |
| self.bias=None | |
| def forward(self,x): | |
| shape=x.shape[:-1] | |
| return _linear(x.reshape(-1,self.in_features),self.weight,self.sf,self.pre, | |
| self.gx,self.alpha,self.down,self.up).reshape(*shape,self.out_features) | |
| def replace_module(model,name,replacement): | |
| parent,leaf=name.rsplit('.',1) | |
| setattr(model.get_submodule(parent),leaf,replacement) | |
| def load_nvfp4_transformer(directory,device='cuda'): | |
| from diffusers import QwenImage21Transformer2DModel | |
| directory=Path(directory) | |
| config=json.loads((directory/'config.json').read_text()) | |
| meta=json.loads((directory/'quantization_config.json').read_text()) | |
| with torch.device('meta'): | |
| model=QwenImage21Transformer2DModel.from_config(config) | |
| for name,spec in meta['quantized_modules'].items(): | |
| n,k=spec['shape'];r=spec['rank'] | |
| state={ | |
| 'weight':torch.empty((n,k//2),dtype=torch.uint8), | |
| 'sf':torch.empty(n*k//16,dtype=torch.uint8), | |
| 'pre':torch.empty(k,dtype=torch.bfloat16), | |
| 'gx':torch.empty(1,dtype=torch.float32), | |
| 'alpha':torch.empty(1,dtype=torch.float32), | |
| 'down':torch.empty((k,r),dtype=torch.bfloat16), | |
| 'up':torch.empty((n,r),dtype=torch.bfloat16), | |
| } | |
| replace_module(model,name,CalibratedNVFP4Linear(state)) | |
| for name,spec in meta.get('fp8_modules',{}).items(): | |
| from fp8_runtime import CalibratedFP8Linear | |
| n,k=spec['shape'] | |
| replace_module(model,name,CalibratedFP8Linear( | |
| torch.empty((n,k),dtype=torch.float8_e4m3fn), | |
| torch.empty((n,1),dtype=torch.float32),torch.empty(k,dtype=torch.float32))) | |
| state={} | |
| for p in sorted(directory.glob('model-*.safetensors')):state.update(load_file(str(p))) | |
| model.load_state_dict(state,strict=True,assign=True) | |
| from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21Rope,QwenImage21TemporalTimesteps | |
| model.pos_embed=QwenImage21Rope(theta=10000,axes_dim=list(model.config.axes_dims_rope)) | |
| model.time_text_embed.time_proj=QwenImage21TemporalTimesteps(timestep_dim=256) | |
| return model.eval().to(device=device) | |
| def load_pipeline(base,quant): | |
| from diffusers import QwenImage21Pipeline | |
| kwargs={} if Path(base).exists() else {'revision':'b3179ad355be050328e483a9dfdd9e60cd62adfa'} | |
| p=QwenImage21Pipeline.from_pretrained(base,transformer=load_nvfp4_transformer(quant),torch_dtype=torch.bfloat16,**kwargs) | |
| p.to(device='cuda');p.set_progress_bar_config(disable=True) | |
| return p | |