Interactive, zero-config PyTorch architecture visualization with local-first MCP tracing.
ModelViz natively brings high-definition computational graphs straight to your VS Code editor. Utilizing rigorous topological analysis via the Model Context Protocol (MCP), it perfectly bridges your raw PyTorch modeling logic into an interactive React-Flow diagram.
🚀 Key Features
Zero-Config ML Environment: ModelViz automatically downloads and manages a minimal, standalone Python environment natively on the first run. No need to mess with venv, conda, or global dependency configurations manually.
Local-First Privacy via MCP: The engine operates via a local Model Context Protocol (MCP) server. Your proprietary architectures and code layers never leave your secure host machine.
Deep Symbolic Tracing: Uses torch.fx.symbolic_trace by default for true execution-path rendering, capturing explicit multidimensional tensors, layer parameter accumulations, and branched operational linkages.
Dynamic UI-Driven Retracing: A beautiful Glassmorphism webview where you can directly modify your target input tensor dimensions (e.g., [1, 1, 28, 28]) and dynamically trigger immediate graph reconstructions.
📋 Prerequisites
Visual Studio Code version 1.90.0 or higher.
Executable Rule: An internet connection is required exclusively for the first launch. ModelViz automatically downloads a minimal Python runtime and necessary dependencies (torch, fastmcp, etc.) into its local bin/ folder for executing local AST and symbolic analysis safely and securely without messing up your global configurations.
🛠️ Usage Workflows
ModelViz natively connects to your files through two very distinct entry points.
1. The VS Code Chat Participant (@modelviz)
You can ask @modelviz to review and visualize an architecture directly alongside your code:
Open the GitHub Copilot Chat panel.
Tag @modelviz in your prompt.
Paste or reference your complete torch.nn.Module class and submit!
2. The Command Palette (Direct Execution)
For an immediate trace on whatever is currently loaded in your window:
Open up your model's source file (.py).
Press Ctrl+Shift+P (or Cmd+Shift+P on macOS) to open the Command Palette.
Use the command: ModelViz: Visualize Model Architecture.
The Extension will analyze the payload and instantly project your graph!
PyTorch, the PyTorch logo and any related marks are trademarks of The Linux Foundation. This extension is an independent project and is not affiliated with, sponsored by, or endorsed by the PyTorch Foundation.
