feat(tools): add export_scroll_decoder for ONNX export

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