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

@torch.library.custom_op('image21_nvfp4::linear',mutates_args=())
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')

@_linear.register_fake
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