ALSv/self-forcing

1# ---- INT8 (optional) ----2from demo_utils.vae import (3    VAEDecoderWrapperSingle,                         # main nn.Module4    ZERO_VAE_CACHE           # helper constants shipped with your code base5)6import pycuda.driver as cuda          # ← add7import pycuda.autoinit  # noqa8 9import sys10from pathlib import Path11 12import torch13import tensorrt as trt14 15from utils.dataset import ShardingLMDBDataset16 17data_path = "/mnt/localssd/wanx_14B_shift-3.0_cfg-5.0_lmdb_oneshard"18dataset = ShardingLMDBDataset(data_path, max_pair=int(1e8))19dataloader = torch.utils.data.DataLoader(20    dataset,21    batch_size=1,22    num_workers=023)24 25# ─────────────────────────────────────────────────────────26# 1️⃣  Bring the PyTorch model into scope27#     (all code you pasted lives in `vae_decoder.py`)28# ─────────────────────────────────────────────────────────29 30# --- dummy tensors (exact shapes you posted) ---31dummy_input = torch.randn(1, 1, 16, 60, 104).half().cuda()32is_first_frame = torch.tensor([1.0], device="cuda", dtype=torch.float16)33dummy_cache_input = [34    torch.randn(*s.shape).half().cuda() if isinstance(s, torch.Tensor) else s35    for s in ZERO_VAE_CACHE               # keep exactly the same ordering36]37inputs = [dummy_input, is_first_frame, *dummy_cache_input]38 39# ─────────────────────────────────────────────────────────40# 2️⃣  Export → ONNX41# ─────────────────────────────────────────────────────────42model = VAEDecoderWrapperSingle().half().cuda().eval()43 44vae_state_dict = torch.load('wan_models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth', map_location="cpu")45decoder_state_dict = {}46for key, value in vae_state_dict.items():47    if 'decoder.' in key or 'conv2' in key:48        decoder_state_dict[key] = value49model.load_state_dict(decoder_state_dict)50model = model.half().cuda().eval()                          # only batch dim dynamic51 52onnx_path = Path("vae_decoder.onnx")53feat_names = [f"vae_cache_{i}" for i in range(len(dummy_cache_input))]54all_inputs_names = ["z", "use_cache"] + feat_names55 56with torch.inference_mode():57    torch.onnx.export(58        model,59        tuple(inputs),                                        # must be a tuple60        onnx_path.as_posix(),61        input_names=all_inputs_names,62        output_names=["rgb_out", "cache_out"],63        opset_version=17,64        do_constant_folding=True,65        dynamo=True66    )67print(f"✅  ONNX graph saved to {onnx_path.resolve()}")68 69# (Optional) quick sanity-check with ONNX-Runtime70try:71    import onnxruntime as ort72    sess = ort.InferenceSession(onnx_path.as_posix(),73                                providers=["CUDAExecutionProvider"])74    ort_inputs = {n: t.cpu().numpy() for n, t in zip(all_inputs_names, inputs)}75    _ = sess.run(None, ort_inputs)76    print("✅  ONNX graph is executable")77except Exception as e:78    print("⚠️  ONNX check failed:", e)79 80# ─────────────────────────────────────────────────────────81# 3️⃣  Build the TensorRT engine82# ─────────────────────────────────────────────────────────83TRT_LOGGER = trt.Logger(trt.Logger.WARNING)84builder = trt.Builder(TRT_LOGGER)85network = builder.create_network(86    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))87parser = trt.OnnxParser(network, TRT_LOGGER)88 89with open(onnx_path, "rb") as f:90    if not parser.parse(f.read()):91        for i in range(parser.num_errors):92            print(parser.get_error(i))93        sys.exit("❌  ONNX → TRT parsing failed")94 95config = builder.create_builder_config()96 97 98def set_workspace(config, bytes_):99    """Version-agnostic workspace limit."""100    if hasattr(config, "max_workspace_size"):                # TRT 8 / 9101        config.max_workspace_size = bytes_102    else:                                                    # TRT 10+103        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, bytes_)104 105 106# …107config = builder.create_builder_config()108set_workspace(config, 4 << 30)          # 4 GB109# 4 GB110 111if builder.platform_has_fast_fp16:112    config.set_flag(trt.BuilderFlag.FP16)113 114# ---- INT8 (optional) ----115# provide a calibrator if you need an INT8 engine; comment this116# block if you only care about FP16.117# ─────────────────────────────────────────────────────────118# helper: version-agnostic workspace limit119# ─────────────────────────────────────────────────────────120 121 122def set_workspace(config: trt.IBuilderConfig, bytes_: int = 4 << 30):123    """124    TRT < 10.x  →  config.max_workspace_size125    TRT ≥ 10.x  →  config.set_memory_pool_limit(...)126    """127    if hasattr(config, "max_workspace_size"):                     # TRT 8 / 9128        config.max_workspace_size = bytes_129    else:                                                         # TRT 10+130        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE,131                                     bytes_)132 133# ─────────────────────────────────────────────────────────134# (optional) INT-8 calibrator135# ─────────────────────────────────────────────────────────136# ‼ Only keep this block if you really need INT-8 ‼                      # gracefully skip if PyCUDA not present137 138 139class VAECalibrator(trt.IInt8EntropyCalibrator2):140    def __init__(self, loader, cache="calibration.cache", max_batches=10):141        super().__init__()142        self.loader = iter(loader)143        self.batch_size = loader.batch_size or 1144        self.max_batches = max_batches145        self.count = 0146        self.cache_file = cache147        self.stream = cuda.Stream()148        self.dev_ptrs = {}149 150    # --- TRT 10 needs BOTH spellings ---151    def get_batch_size(self):152        return self.batch_size153 154    def getBatchSize(self):155        return self.batch_size156 157    def get_batch(self, names):158        if self.count >= self.max_batches:159            return None160 161        # Randomly sample a number from 1 to 10162        import random163        vae_idx = random.randint(0, 10)164        data = next(self.loader)165 166        latent = data['ode_latent'][0][:, :1]167        is_first_frame = torch.tensor([1.0], device="cuda", dtype=torch.float16)168        feat_cache = ZERO_VAE_CACHE169        for i in range(vae_idx):170            inputs = [latent, is_first_frame, *feat_cache]171            with torch.inference_mode():172                outputs = model(*inputs)173            latent = data['ode_latent'][0][:, i + 1:i + 2]174            is_first_frame = torch.tensor([0.0], device="cuda", dtype=torch.float16)175            feat_cache = outputs[1:]176 177        # -------- ensure context is current --------178        z_np = latent.cpu().numpy().astype('float32')179 180        ptrs = []                # list[int] – one entry per name181        for name in names:         # <-- match TRT's binding order182            if name == "z":183                arr = z_np184            elif name == "use_cache":185                arr = is_first_frame.cpu().numpy().astype('float32')186            else:187                idx = int(name.split('_')[-1])   # "vae_cache_17" -> 17188                arr = feat_cache[idx].cpu().numpy().astype('float32')189 190            if name not in self.dev_ptrs:191                self.dev_ptrs[name] = cuda.mem_alloc(arr.nbytes)192 193            cuda.memcpy_htod_async(self.dev_ptrs[name], arr, self.stream)194            ptrs.append(int(self.dev_ptrs[name]))   # ***int() is required***195 196        self.stream.synchronize()197        self.count += 1198        print(f"Calibration batch {self.count}/{self.max_batches}")199        return ptrs200 201    # --- calibration-cache helpers (both spellings) ---202    def read_calibration_cache(self):203        try:204            with open(self.cache_file, "rb") as f:205                return f.read()206        except FileNotFoundError:207            return None208 209    def readCalibrationCache(self):210        return self.read_calibration_cache()211 212    def write_calibration_cache(self, cache):213        with open(self.cache_file, "wb") as f:214            f.write(cache)215 216    def writeCalibrationCache(self, cache):217        self.write_calibration_cache(cache)218 219 220# ─────────────────────────────────────────────────────────221# Builder-config + optimisation profile222# ─────────────────────────────────────────────────────────223config = builder.create_builder_config()224set_workspace(config, 4 << 30)                    # 4 GB225 226# ► enable FP16 if possible227if builder.platform_has_fast_fp16:228    config.set_flag(trt.BuilderFlag.FP16)229 230# ► enable INT-8  (delete this block if you don’t need it)231if cuda is not None:232    config.set_flag(trt.BuilderFlag.INT8)233    # supply any representative batch you like – here we reuse the latent z234    calib = VAECalibrator(dataloader)235    # TRT-10 renamed the setter:236    if hasattr(config, "set_int8_calibrator"):    # TRT 10+237        config.set_int8_calibrator(calib)238    else:                                         # TRT ≤ 9239        config.int8_calibrator = calib240 241# ---- optimisation profile ----242profile = builder.create_optimization_profile()243profile.set_shape(all_inputs_names[0],            # latent z244                  min=(1, 1, 16, 60, 104),245                  opt=(1, 1, 16, 60, 104),246                  max=(1, 1, 16, 60, 104))247profile.set_shape("use_cache",               # scalar flag248                  min=(1,), opt=(1,), max=(1,))249for name, tensor in zip(all_inputs_names[2:], dummy_cache_input):250    profile.set_shape(name, tensor.shape, tensor.shape, tensor.shape)251 252config.add_optimization_profile(profile)253 254# ─────────────────────────────────────────────────────────255# Build the engine  (API changed in TRT-10)256# ─────────────────────────────────────────────────────────257print("⚙️  Building engine … (can take a minute)")258 259if hasattr(builder, "build_serialized_network"):          # TRT 10+260    serialized_engine = builder.build_serialized_network(network, config)261    assert serialized_engine is not None, "build_serialized_network() failed"262    plan_path = Path("checkpoints/vae_decoder_int8.trt")263    plan_path.write_bytes(serialized_engine)264    engine_bytes = serialized_engine                      # keep for smoke-test265else:                                                     # TRT ≤ 9266    engine = builder.build_engine(network, config)267    assert engine is not None, "build_engine() returned None"268    plan_path = Path("checkpoints/vae_decoder_int8.trt")269    plan_path.write_bytes(engine.serialize())270    engine_bytes = engine.serialize()271 272print(f"✅  TensorRT engine written to {plan_path.resolve()}")273 274# ─────────────────────────────────────────────────────────275# 4️⃣  Quick smoke-test with the brand-new engine276# ─────────────────────────────────────────────────────────277with trt.Runtime(TRT_LOGGER) as rt:278    engine = rt.deserialize_cuda_engine(engine_bytes)279    context = engine.create_execution_context()280    stream = torch.cuda.current_stream().cuda_stream281 282    # pre-allocate device buffers once283    device_buffers, outputs = {}, []284    dtype_map = {trt.float32: torch.float32,285                 trt.float16: torch.float16,286                 trt.int8:    torch.int8,287                 trt.int32:   torch.int32}288 289    for name, tensor in zip(all_inputs_names, inputs):290        if -1 in engine.get_tensor_shape(name):            # dynamic input291            context.set_input_shape(name, tensor.shape)292        context.set_tensor_address(name, int(tensor.data_ptr()))293        device_buffers[name] = tensor294 295    context.infer_shapes()                                 # propagate ⇢ outputs296    for i in range(engine.num_io_tensors):297        name = engine.get_tensor_name(i)298        if engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT:299            shape = tuple(context.get_tensor_shape(name))300            dtype = dtype_map[engine.get_tensor_dtype(name)]301            out = torch.empty(shape, dtype=dtype, device="cuda")302            context.set_tensor_address(name, int(out.data_ptr()))303            outputs.append(out)304            print(f"output {name} shape: {shape}")305 306    context.execute_async_v3(stream_handle=stream)307    torch.cuda.current_stream().synchronize()308    print("✅  TRT execution OK – first output shape:", outputs[0].shape)309