visualtorch-mcp

<div align="center"> <h1>🔥 VisualTorch MCP 🔥</h1> </div>

MCP server for rendering PyTorch architecture diagrams with VisualTorch.

Documentation surfaced by this MCP points to the upstream VisualTorch docs in willyfh/visualtorch/docs. The MCP does not maintain a separate API guide.

The rendering styles are the upstream VisualTorch styles documented here:

• Render API
• Flow View
• Graph View
• LeNet Style View

The MCP accepts layered as an alias for flow to tolerate older VisualTorch releases and docs that used the layered naming.

Install

python -m venv .venv
.\.venv\Scripts\python -m pip install -U pip
.\.venv\Scripts\python -m pip install -e ".[render]"

The render extra installs VisualTorch from willyfh/visualtorch so the newer flow and render(...) APIs are available. If you already manage VisualTorch separately, install this MCP with pip install -e . instead.

MCP config

Use the console script after installing in your environment:

{
  "mcpServers": {
    "visualtorch": {
      "command": "visualtorch-mcp"
    }
  }
}

Or call the module through a specific Python interpreter:

{
  "mcpServers": {
    "visualtorch": {
      "command": "D:\\Florent\\Desktop\\visualtorch-mcp\\.venv\\Scripts\\python.exe",
      "args": ["-m", "visualtorch_mcp.server"]
    }
  }
}

Tools

visualize_model

Renders a model and returns the output file path plus image metadata.

Arguments:

• source: Python source code that defines the model and imports its dependencies.
• input_shape: shape including batch dimension, for example [1, 3, 224, 224]. For multi-input models, pass one shape per positional input, for example [[1, 3, 224, 224], [1, 10]].
• style: graph, flow, layered, or lenet.
• model_expression: expression evaluated after source runs. Defaults to model; use Net() or build_model() when the source defines a class/factory.
• output_path: optional absolute or relative image path. Defaults to a generated PNG in visualtorch_outputs.
• options: VisualTorch render options such as {"palette": "dracula", "show_dimension": true}.
• workdir: optional working directory used while executing the model source.
• timeout_seconds: render subprocess timeout.

Example tool payload:

{
  "source": "import torch\nfrom torch import nn\nmodel = nn.Sequential(nn.Conv2d(3, 8, 3, padding=1), nn.ReLU(), nn.Flatten(), nn.Linear(8 * 16 * 16, 10))",
  "input_shape": [1, 3, 16, 16],
  "style": "graph",
  "options": {
    "show_dimension": true,
    "palette": "okabe_ito"
  }
}

visualtorch_reference

Returns upstream VisualTorch documentation links for the requested style, including the GitHub docs source and Read the Docs page. Style aliases listed by this tool are MCP aliases, not separate VisualTorch documentation.

Development

python -m pip install -e ".[test]"
pytest

The render worker executes supplied Python source in a subprocess with a timeout. Treat it as a local developer tool: it intentionally runs model code so it can instantiate arbitrary PyTorch models, but the subprocess boundary prevents a render exception from crashing the MCP server.

Citation

If this MCP helps you use VisualTorch in research or published work, cite VisualTorch:

Hendria, W. F., & Gavrikov, P. (2024). VisualTorch: Streamlining visualization for PyTorch neural network architectures. Journal of Open Source Software, 9(102), 6678. https://doi.org/10.21105/joss.06678

@article{hendria2024visualtorch,
  author = {Hendria, Willy Fitra and Gavrikov, Paul},
  title = {VisualTorch: Streamlining Visualization for PyTorch Neural Network Architectures},
  journal = {Journal of Open Source Software},
  year = {2024},
  volume = {9},
  number = {102},
  pages = {6678},
  doi = {10.21105/joss.06678},
  url = {https://doi.org/10.21105/joss.06678}
}

License

This MCP server is released under the MIT License. VisualTorch is also released under the MIT License by Paul Gavrikov and Willy Fitra Hendria; see the upstream VisualTorch license.