"""Calibrated W8A8 E4M3 inference using native CUTLASS SM120 scaled GEMMs.

Built with Qwen. Quantization modifications by ProCreations, 2026.
The upstream model is subject to the included Qwen Research License.
"""
import json
from pathlib import Path
import torch
from torch import nn
import triton
import triton.language as tl
from safetensors.torch import load_file


@triton.jit
def _quantize_rows(X, SMOOTH, Q, SCALE, K: tl.constexpr, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    x = tl.load(X + row * K + cols, cols < K, 0).to(tl.float32)
    s = tl.load(SMOOTH + cols, cols < K, 1).to(tl.float32)
    z = x / s
    scale = tl.maximum(tl.max(tl.abs(z), 0) / 448.0, 1.e-12)
    tl.store(Q + row * K + cols, z / scale, cols < K)
    tl.store(SCALE + row, scale)


def quantize_activation(x, smooth):
    x = x.reshape(-1, x.shape[-1]).contiguous()
    q = torch.empty_like(x, dtype=torch.float8_e4m3fn)
    scale = torch.empty((x.shape[0], 1), device=x.device, dtype=torch.float32)
    _quantize_rows[(x.shape[0],)](x, smooth, q, scale, x.shape[1],
                                triton.next_power_of_2(x.shape[1]), num_warps=8)
    return q, scale


class CalibratedFP8Linear(nn.Module):
    def __init__(self, weight, scale, smooth):
        super().__init__()
        self.register_buffer("weight", weight)
        self.register_buffer("weight_scale", scale.float())
        self.register_buffer("smooth", smooth.float())
        self.in_features = weight.shape[1]
        self.out_features = weight.shape[0]
        self.bias = None

    def forward(self, x):
        shape = x.shape[:-1]
        q, scale = quantize_activation(x, self.smooth)
        y = torch._scaled_mm(q, self.weight.t(), scale_a=scale,
                             scale_b=self.weight_scale.t(),
                             out_dtype=torch.bfloat16, use_fast_accum=False)
        return y.reshape(*shape, self.out_features)


def eligible_modules(model):
    return {n:m for n,m in model.named_modules()
            if isinstance(m, nn.Linear) and n.startswith("transformer_blocks.")
            and m.bias is None}


def replace_module(root, name, replacement):
    parent_name, leaf = name.rsplit(".", 1)
    setattr(root.get_submodule(parent_name), leaf, replacement)


def load_fp8_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())
    # Construct without allocating an intermediate BF16 model.
    with torch.device("meta"):
        model = QwenImage21Transformer2DModel.from_config(config)
        for n, spec in meta["quantized_modules"].items():
            out_f, in_f = spec["shape"]
            replace_module(model, n, CalibratedFP8Linear(
                torch.empty((out_f,in_f),dtype=torch.float8_e4m3fn),
                torch.empty((out_f,1),dtype=torch.float32),
                torch.empty(in_f,dtype=torch.float32)))
    state = {}
    for path in sorted(directory.glob("model-*.safetensors")):
        state.update(load_file(str(path), device="cpu"))
    model.load_state_dict(state, strict=True, assign=True)
    # The original rotary module has nonpersistent buffers / Python tensor lists.
    # Reconstruct these by creating a lightweight ordinary position module.
    from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21Rope
    model.pos_embed = QwenImage21Rope(theta=10000, axes_dim=list(model.config.axes_dims_rope))
    from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21TemporalTimesteps
    model.time_text_embed.time_proj = QwenImage21TemporalTimesteps(timestep_dim=256)
    return model.eval().to(device=device)


def load_pipeline(base, quant=None):
    from diffusers import QwenImage21Pipeline
    args = {}
    if not Path(base).exists():
        args["revision"] = "b3179ad355be050328e483a9dfdd9e60cd62adfa"
    if quant:
        args["transformer"] = load_fp8_transformer(quant, device="cuda")
    pipe = QwenImage21Pipeline.from_pretrained(base, torch_dtype=torch.bfloat16, **args)
    # Do not call pipe.to(dtype=...) after FP8 loading; scales are FP32.
    pipe.to(device="cuda")
    pipe.set_progress_bar_config(disable=True)
    return pipe