import os, torch
os.environ.setdefault("DIFFSYNTH_ATTENTION_IMPLEMENTATION","torch")
import bitsandbytes as bnb
print("bnb", bnb.__version__)
dev="cuda:1"
# --- NF4 quantize + forward in FP16 on sm_75 ---
lin = torch.nn.Linear(5376, 14336, bias=False, dtype=torch.float16)
q = bnb.nn.Linear4bit(5376, 14336, bias=False, compute_dtype=torch.float16, quant_type="nf4", compress_statistics=True)
q.weight = bnb.nn.Params4bit(lin.weight.data, requires_grad=False, quant_type="nf4", compress_statistics=True)
q = q.to(dev)
x = torch.randn(1024, 5376, device=dev, dtype=torch.float16)
y = q(x)
ref = torch.nn.functional.linear(x, lin.weight.to(dev))
err = (y.float()-ref.float()).abs().mean().item()/ref.float().abs().mean().item()
print(f"[OK ] NF4 fp16 forward: out={tuple(y.shape)} dtype={y.dtype} finite={bool(torch.isfinite(y).all())} rel_err={err:.4f}")
# --- prequantized round-trip (what DiffSynth actually does when loading .safetensors) ---
qs = q.weight.quant_state.as_dict(packed=True)
w2 = bnb.nn.Params4bit.from_prequantized(data=q.weight.data, quantized_stats={k:(v.to(dev) if torch.is_tensor(v) else v) for k,v in qs.items()}, requires_grad=False, device=dev)
q2 = bnb.nn.Linear4bit(5376,14336,bias=False,compute_dtype=torch.float16,quant_type="nf4")
q2.weight = w2; q2 = q2.to(dev)
y2 = q2(x)
print(f"[OK ] NF4 prequantized reload: match={torch.allclose(y,y2)} dtype={y2.dtype}")
# --- DiffSynth attention backend selection ---
import sys
from diffsynth.core.attention import attention_forward
A = sys.modules[attention_forward.__module__]
print("ATTENTION_IMPLEMENTATION =", A.ATTENTION_IMPLEMENTATION,
      "| xformers_avail =", A.XFORMERS_AVAILABLE, "| flash2 =", A.FLASH_ATTN_2_AVAILABLE, "| sage =", A.SAGE_ATTN_AVAILABLE)
o = A.attention_forward(*[torch.randn(1,56,4096,128,device=dev,dtype=torch.float16)]*3)
print("[OK ] diffsynth attention_forward fp16:", tuple(o.shape), "finite=", bool(torch.isfinite(o).all()))