# Architecture Documentation
## System Overview
The MCP ComfyUI Flux system is a fully containerized, optimized architecture consisting of services that communicate via WebSocket and HTTP protocols. The system has been optimized for performance, reduced image size, and faster build times while maintaining high availability and ease of deployment.
## Component Architecture
### 1. MCP Server Container
**Purpose**: Implements the Model Context Protocol to bridge Claude/MCP clients with ComfyUI.
**Technology Stack**:
- Node.js 20 (Alpine Linux - 581MB image)
- MCP SDK (@modelcontextprotocol/sdk)
- WebSocket client (ws)
- Claude Code CLI (pre-installed)
- Development tools (ripgrep, git, bash)
**Key Responsibilities**:
- MCP protocol implementation with auto-connect
- Tool registration and handling
- WebSocket connection management
- Image data encoding/decoding
- Error handling and automatic recovery
- File path-based responses for large images
**Design Patterns**:
- Singleton pattern for ComfyUI client
- Command pattern for tool handlers
- Observer pattern for WebSocket events
- Auto-reconnect with exponential backoff
### 2. ComfyUI Container (Optimized)
**Purpose**: Provides the ML inference engine for Flux model execution with optimized performance.
**Technology Stack**:
- Ubuntu 22.04 with CUDA 12.1 (base layer cached)
- Python 3.11 (no venv - Docker provides isolation)
- PyTorch 2.5.1 with CUDA support (latest stable)
- ComfyUI framework with custom nodes
- Hugging Face libraries with transformers cache
- Claude Code CLI (for in-container development)
**Optimization Details**:
- **Multi-stage build**: Separates build and runtime dependencies
- **BuildKit cache mounts**: Reduces I/O operations in WSL2
- **No virtual environment**: Docker IS the isolation (saves ~2GB)
- **PyTorch 2.5.1**: Native RMSNorm support, better performance
- **Image size**: 10.9GB (down from 14.6GB - 25% reduction)
- **Build time**: 40% faster with cache mounts
**Custom Nodes Included**:
- ComfyUI-Manager: Node management and updates
- ComfyUI-KJNodes: Advanced processing nodes
- ComfyUI-RMBG: Background removal (RMBG-2.0 model)
**Key Responsibilities**:
- Model loading and management (FP8 quantized models)
- Workflow execution with 4-step generation
- GPU memory management with CUDA 12.1
- Image generation, upscaling, and processing
- Custom node support with dependency management
**Design Decisions**:
- Direct system Python (no venv overhead)
- CUDA 12.1 for RTX 40-series optimization
- FP8 models for 50% memory reduction
- Sequential git clones for reliability
- Non-root user (comfyuser) for security
### 3. Docker Orchestration
**Purpose**: Container lifecycle management with BuildKit optimizations.
**Components**:
- `docker-compose.yml`: Service definitions with health checks
- `Dockerfile.comfyui`: Optimized multi-stage build
- `Dockerfile.mcp`: Lightweight MCP server build
- `docker-entrypoint.sh`: Smart initialization with model detection
- `build.sh`: BuildKit-optimized build script
- `install.sh`: Automated setup with GPU detection
**Network Architecture**:
```
┌─────────────────────────┐
│ Host Network │
│ Port 8188 (ComfyUI) │
└───────────┬─────────────┘
│
┌───────────▼─────────────┐
│ Bridge Network │
│ (mcp-network) │
│ MTU: 1450 (WSL2 opt) │
├─────────────────────────┤
│ • comfyui:8188 │
│ • mcp-server:internal │
└─────────────────────────┘
```
**Volume Management**:
- `models/`: Persistent model storage with symlinks
- `output/`: Generated images
- `input/`: Input images for processing
- `pycache/`: Python bytecode cache (faster startup)
- No custom_nodes volume (built into image)
## Data Flow Architecture
### Image Generation Flow (Optimized)
```mermaid
sequenceDiagram
participant Client as MCP Client
participant MCP as MCP Server
participant WS as WebSocket
participant ComfyUI as ComfyUI Server
participant GPU as GPU/CUDA
Client->>MCP: generate_image(prompt)
MCP->>WS: Auto-connect to ComfyUI
WS->>ComfyUI: Queue workflow (FP8 models)
ComfyUI->>GPU: Load FLUX schnell fp8
GPU->>GPU: 4-step inference (2-4s)
GPU->>ComfyUI: Generated latents
ComfyUI->>ComfyUI: VAE decode (fp8)
ComfyUI->>WS: Execution complete
WS->>MCP: File path + metadata
MCP->>Client: Return file path (not base64)
```
### Error Recovery Flow (Enhanced)
```mermaid
flowchart TD
A[Request] --> B{Connection OK?}
B -->|No| C[Auto-Reconnect]
C --> D{Max Retries?}
D -->|No| C
D -->|Yes| E[Return Error]
B -->|Yes| F[Execute Workflow]
F --> G{Execution OK?}
G -->|No| H[Log Error]
H --> I[Clean Resources]
I --> J{Retry?}
J -->|Yes| C
J -->|No| E
G -->|Yes| K[Save to output/]
K --> L[Return File Path]
```
## Workflow System (FLUX Optimized)
### FLUX Schnell FP8 Workflow
The system uses optimized FP8 quantized models for 50% memory reduction:
**Model Configuration**:
```javascript
{
"unet_loader": {
"class_type": "UNETLoader",
"inputs": {
"unet_name": "flux1-schnell-fp8-e4m3fn.safetensors",
"weight_dtype": "fp8_e4m3fn"
}
},
"dual_clip_loader": {
"class_type": "DualCLIPLoader",
"inputs": {
"clip_name1": "t5xxl_fp8_e4m3fn_scaled.safetensors",
"clip_name2": "clip_l.safetensors",
"type": "flux"
}
},
"vae_loader": {
"class_type": "VAELoader",
"inputs": {
"vae_name": "ae.safetensors"
}
},
"sampler": {
"class_type": "KSampler",
"inputs": {
"steps": 4, // Schnell optimized for 4 steps
"cfg": 1.0, // Low guidance for schnell
"sampler_name": "euler",
"scheduler": "simple",
"denoise": 1.0
}
}
}
```
### Performance Metrics
- **Generation Time**: 2-4 seconds per image
- **Batch Generation**: ~1.5s per additional image
- **VRAM Usage**: ~10GB base + 1GB per batch
- **Model Loading**: One-time 5-10s initialization
- **Quality**: 95% of FP16 at 50% memory
## Security Architecture
### Container Security
- **User Isolation**: Non-root user (comfyuser:1000)
- **Network Isolation**: Internal bridge network
- **Build-time Security**: No secrets in layers
- **Runtime Security**: Read-only where possible
- **Resource Limits**: Memory and CPU constraints
### Secret Management
```yaml
# Environment variables for sensitive data
environment:
- HF_TOKEN=${HF_TOKEN:-} # Optional for gated models
- COMFYUI_API_KEY=${COMFYUI_API_KEY:-} # Optional API key
# Docker secrets (production)
secrets:
hf_token:
external: true
```
## Performance Optimization
### GPU Memory Management (Enhanced)
```python
# Optimized memory allocation
PYTORCH_CUDA_ALLOC_CONF = {
"max_split_size_mb": 512,
"garbage_collection_threshold": 0.7
}
# FP8 quantization settings
MODEL_CONFIGS = {
"flux-schnell-fp8": {
"precision": "fp8_e4m3fn",
"memory": "10GB",
"speed": "2-4s",
"quality": "95%"
},
"flux-dev-fp16": {
"precision": "fp16",
"memory": "24GB",
"speed": "10-20s",
"quality": "100%"
}
}
```
### Build Optimization (BuildKit)
```dockerfile
# Cache mount for pip packages
RUN --mount=type=cache,target=/root/.cache/pip,sharing=locked \
python3.11 -m pip install --user -r requirements.txt
# Cache mount for apt packages
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
--mount=type=cache,target=/var/lib/apt,sharing=locked \
apt-get update && apt-get install -y packages
# Pre-compile Python bytecode
RUN python3.11 -m compileall -q /root/.local/lib/python3.11/site-packages/
```
### Caching Strategy
1. **Docker Layer Cache**: Optimized instruction ordering
2. **BuildKit Cache Mounts**: Persistent package caches
3. **Python Bytecode Cache**: Pre-compiled .pyc files
4. **Model Cache**: Persistent volume with symlinks
5. **Transformers Cache**: Cached model configs
### WSL2-Specific Optimizations
```yaml
# docker-compose.yml
services:
comfyui:
# WSL2 optimizations
shm_size: "16g" # Shared memory for PyTorch
environment:
- LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4 # Better memory allocation
networks:
mcp-network:
driver_opts:
com.docker.network.driver.mtu: 1450 # Optimize for WSL2
```
## Health Monitoring
### Container Health Checks
```yaml
healthcheck:
test: ["CMD", "curl", "-f", "-H", "Host: localhost", "http://localhost:8188/system_stats"]
interval: 30s
timeout: 10s
retries: 3
start_period: 60s
```
### Service Monitoring
```javascript
// MCP Server health endpoint
app.get('/health', (req, res) => {
res.json({
status: 'healthy',
uptime: process.uptime(),
memory: process.memoryUsage(),
comfyui: comfyuiClient.isConnected(),
version: process.env.npm_package_version
});
});
// ComfyUI connection monitoring
class ComfyUIClient {
constructor() {
this.reconnectAttempts = 0;
this.maxReconnectAttempts = 5;
this.reconnectDelay = 1000;
}
async autoConnect() {
while (this.reconnectAttempts < this.maxReconnectAttempts) {
try {
await this.connect();
console.log('Auto-connected to ComfyUI');
this.reconnectAttempts = 0;
return;
} catch (error) {
this.reconnectAttempts++;
const delay = this.reconnectDelay * Math.pow(2, this.reconnectAttempts);
await new Promise(resolve => setTimeout(resolve, delay));
}
}
throw new Error('Failed to connect after maximum attempts');
}
}
```
## Development Workflow
### Local Development
```bash
# Quick rebuild with cache
./build.sh --start
# Development with hot reload
docker exec -it mcp-comfyui-flux-comfyui-1 bash
cd /app/ComfyUI
python main.py --listen --preview-method auto
# Access Claude Code in container
docker exec -it mcp-comfyui-flux-mcp-server-1 claude
```
### Testing Workflow
```javascript
// Test MCP integration
const { generateImage } = require('./src/tools');
describe('Image Generation', () => {
it('should generate with FLUX schnell', async () => {
const result = await generateImage({
prompt: 'test image',
steps: 4,
cfg_scale: 1.0
});
expect(result.filename).toMatch(/flux_output_\d+\.png/);
});
});
```
### Performance Testing
```bash
# Monitor GPU usage during generation
docker exec mcp-comfyui-flux-comfyui-1 nvidia-smi -l 1
# Check memory usage
docker stats mcp-comfyui-flux-comfyui-1
# Benchmark generation time
time docker exec mcp-comfyui-flux-mcp-server-1 node -e "
const client = require('/app/src/comfyui-client.js');
client.generateImage({prompt: 'benchmark test'}).then(console.log);
"
```
## Deployment Patterns
### Production Deployment
```yaml
# docker-compose.prod.yml
services:
comfyui:
image: mcp-comfyui-comfyui:latest
deploy:
resources:
reservations:
devices:
- driver: nvidia
count: 1
capabilities: [gpu]
limits:
memory: 20G
restart: unless-stopped
logging:
driver: "json-file"
options:
max-size: "100m"
max-file: "3"
```
### Multi-GPU Setup
```yaml
# Scale across multiple GPUs
services:
comfyui:
deploy:
replicas: 2
environment:
- CUDA_VISIBLE_DEVICES=${GPU_ID} # 0 or 1
```
## Troubleshooting Guide
### Common Issues and Solutions
1. **Out of Memory (OOM)**
- Use FP8 models instead of FP16
- Reduce batch size
- Lower resolution (768x768 instead of 1024x1024)
2. **Slow Generation**
- Ensure using schnell model (4 steps)
- Check cfg_scale is 1.0 (not 7.0)
- Verify PyTorch 2.5.1 is installed
3. **Container Won't Start**
- Check port 8188 availability
- Verify models are in correct directories
- Check Docker resources (WSL2 memory)
4. **yarl/aiohttp Errors**
- Fixed in optimized build with pinned versions
- aiohttp==3.10.10, yarl==1.17.2
## Migration Path
### From Original to Optimized
```bash
# 1. Backup existing setup
cp docker-compose.yml .backup/
cp Dockerfile.comfyui .backup/
# 2. Stop old containers
docker-compose -p mcp-comfyui-flux down
# 3. Build optimized version
./build.sh --start
# 4. Verify functionality
curl http://localhost:8188/system_stats
```
## Future Enhancements
### Planned Optimizations
1. **Further Size Reduction**
- Alpine-based Python image (experimental)
- Distroless runtime containers
- Multi-arch builds (ARM64 support)
2. **Performance Improvements**
- TensorRT optimization
- ONNX model conversion
- Triton Inference Server integration
3. **Scalability Features**
- Kubernetes operators
- Horizontal pod autoscaling
- Distributed model serving
4. **Developer Experience**
- VS Code Dev Containers
- Jupyter notebook integration
- Real-time collaboration features
## Architecture Decisions Log
### Decision: Remove Python venv (2024)
- **Context**: venv added 2GB and caused WSL2 chown crashes
- **Decision**: Use system Python in Docker
- **Outcome**: 25% size reduction, faster builds, no crashes
### Decision: Upgrade to PyTorch 2.5.1
- **Context**: Warnings about old PyTorch, missing native RMSNorm
- **Decision**: Update from 2.2.0 to 2.5.1
- **Outcome**: Better performance, no warnings, native optimizations
### Decision: Use FP8 Quantized Models
- **Context**: 24GB VRAM requirement limiting accessibility
- **Decision**: Default to FP8 models (schnell)
- **Outcome**: 50% memory reduction, 2-4s generation, 95% quality
### Decision: BuildKit Cache Mounts
- **Context**: WSL2 I/O performance issues
- **Decision**: Implement cache mounts for packages
- **Outcome**: 40% faster rebuilds, reduced I/O operations
## Conclusion
The optimized MCP ComfyUI Flux architecture delivers:
- **Performance**: 2-4s generation with FP8 models
- **Efficiency**: 25% smaller images, 40% faster builds
- **Reliability**: Auto-reconnect, health checks, error recovery
- **Scalability**: BuildKit caching, multi-GPU ready
- **Security**: Non-root execution, isolated networks
- **Developer Experience**: Claude Code integration, fast rebuilds
The system maintains production reliability while significantly improving resource utilization and build times.
