visualtorch-mcp

<div align="center"> <h1>🔥 VisualTorch MCP 🔥</h1> </div> MCP server for rendering PyTorch architecture diagrams with [VisualTorch](https://visualtorch.readthedocs.io/en/latest/). Documentation surfaced by this MCP points to the upstream VisualTorch docs in [`willyfh/visualtorch/docs`](https://github.com/willyfh/visualtorch/tree/main/docs). The MCP does not maintain a separate API guide. The rendering styles are the upstream VisualTorch styles documented here: - [Render API](https://github.com/willyfh/visualtorch/blob/main/docs/source/markdown/api_references/render.md) - [Flow View](https://github.com/willyfh/visualtorch/blob/main/docs/source/markdown/api_references/flow.md) - [Graph View](https://github.com/willyfh/visualtorch/blob/main/docs/source/markdown/api_references/graph.md) - [LeNet Style View](https://github.com/willyfh/visualtorch/blob/main/docs/source/markdown/api_references/lenet_style.md) The MCP accepts `layered` as an alias for `flow` to tolerate older VisualTorch releases and docs that used the layered naming. ## Install ```powershell 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: ```json { "mcpServers": { "visualtorch": { "command": "visualtorch-mcp" } } } ``` Or call the module through a specific Python interpreter: ```json { "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: ```json { "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 ```powershell 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 ```bibtex @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](https://github.com/willyfh/visualtorch/blob/main/LICENSE).