diff --git a/tools/export_onnx.py b/tools/export_onnx.py index 0cd4ac1..054979d 100644 --- a/tools/export_onnx.py +++ b/tools/export_onnx.py @@ -93,3 +93,79 @@ def export_flow_model(ckpt_dir: Path, out_dir: Path) -> Path: ) logger.info("Wrote %s (%.1f MB)", out_path, out_path.stat().st_size / 1e6) return out_path + + +class _ScrollDecoder(torch.nn.Module): + """Wraps ScrollCVAE.decode for ONNX export. + + The full ScrollCVAE is encoder+decoder; inference only needs decoder. + """ + + def __init__(self, dec_h0, dec_gru, dec_out, seq_len: int, hidden: int): + super().__init__() + self.dec_h0 = dec_h0 + self.dec_gru = dec_gru + self.dec_out = dec_out + self.seq_len = seq_len + self.hidden = hidden + + def forward(self, z: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: + b = z.shape[0] + zc = torch.cat([z, cond], dim=-1) + h0_flat = self.dec_h0(zc) + h0 = h0_flat.view(b, 2, self.hidden).permute(1, 0, 2).contiguous() + inp = zc.unsqueeze(1).expand(b, self.seq_len, -1) + out, _ = self.dec_gru(inp, h0) + return self.dec_out(out) + + +def export_scroll_decoder(ckpt_dir: Path, out_dir: Path) -> Path: + """Export ScrollCVAE decoder to ONNX.""" + from tools.scroll.models import ScrollCVAE + + config_path = ckpt_dir / "scroll_config.json" + cfg = json.loads(config_path.read_text()) + seq_len = int(cfg["seq_len"]) + latent_dim = int(cfg["latent_dim"]) + hidden = int(cfg["hidden"]) + cond_dim = int(cfg["cond_dim"]) + + full = ScrollCVAE( + seq_len=seq_len, latent_dim=latent_dim, hidden=hidden, cond_dim=cond_dim + ) + state = torch.load(ckpt_dir / "scroll_model.pt", map_location="cpu", weights_only=True) + full.load_state_dict(state) + full.eval() + + decoder = _ScrollDecoder( + dec_h0=full.dec_h0, + dec_gru=full.dec_gru, + dec_out=full.dec_out, + seq_len=seq_len, + hidden=hidden, + ) + decoder.eval() + + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / "scroll_decoder.onnx" + + dummy_z = torch.zeros(1, latent_dim, dtype=torch.float32) + dummy_cond = torch.zeros(1, cond_dim, dtype=torch.float32) + + torch.onnx.export( + decoder, + (dummy_z, dummy_cond), + str(out_path), + input_names=["z", "cond"], + output_names=["seq"], + dynamic_axes={ + "z": {0: "batch"}, + "cond": {0: "batch"}, + "seq": {0: "batch"}, + }, + opset_version=17, + do_constant_folding=True, + dynamo=False, + ) + logger.info("Wrote %s (%.1f KB)", out_path, out_path.stat().st_size / 1e3) + return out_path