# CLAUDE.md
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
## Project Overview
This is a fully containerized MCP (Model Context Protocol) server that integrates with ComfyUI to generate images using FLUX schnell fp8 models. The entire system runs in optimized Docker containers with BuildKit cache mounts, PyTorch 2.5.1, automatic setup, GPU acceleration, and Claude Desktop integration.
## Quick Start for Development
```bash
# Build containers (with optional flags)
./build.sh # Build only
./build.sh --start # Build and start
./build.sh --start --cleanup # Build, start, and cleanup old images
# Or use the one-command setup
./install.sh
# Access the development environment
docker exec -it mcp-comfyui-flux-mcp-server-1 bash
cd /app
claude # Claude Code is pre-installed
# Or work directly with the code
docker exec -it mcp-comfyui-flux-mcp-server-1 claude --dir /app
```
## MCP Server Usage
### Available MCP Tools
#### 1. generate_image
Generate high-quality images using FLUX schnell fp8 model.
**Function Signature:**
```javascript
generate_image({
prompt: string, // Required: Description of image to generate
negative_prompt?: string, // Optional: What to avoid (default: '')
width?: number, // Optional: Image width (default: 1024)
height?: number, // Optional: Image height (default: 1024)
steps?: number, // Optional: Generation steps (default: 4)
cfg_scale?: number, // Optional: Guidance scale (default: 1.0)
seed?: number, // Optional: Random seed (default: -1 for random)
sampler_name?: string, // Optional: Sampler method (default: 'euler')
scheduler?: string, // Optional: Scheduler type (default: 'simple')
batch_size?: number // Optional: Number of images to generate (default: 1, max: 8)
})
```
**Optimized Defaults for FLUX schnell:**
- `steps: 4` - FLUX schnell is specifically designed for 4-step generation
- `cfg_scale: 1.0` - Low guidance works best with schnell's distilled model
- `scheduler: 'simple'` - Optimal for schnell's architecture
- `sampler_name: 'euler'` - Most stable and predictable results
- `batch_size: 1` - Can be increased up to 8 for multiple variations
**Batch Generation:**
- Native ComfyUI batch support for efficient parallel generation
- Generates multiple images in a single GPU pass (much faster than sequential)
- Each image in batch gets a unique seed (unless seed is specified)
- With specified seed, each image gets seed+index (e.g., seed:100 → 100, 101, 102...)
**Usage Examples:**
```javascript
// Simple generation with defaults (recommended)
generate_image({
prompt: "a serene japanese garden with cherry blossoms"
})
// Detailed prompt with style guidance
generate_image({
prompt: "cyberpunk city at night, neon lights, rain, reflections, detailed architecture, flying cars",
negative_prompt: "blurry, low quality, distorted"
})
// Portrait with specific dimensions
generate_image({
prompt: "portrait of a wizard, detailed face, magical aura, fantasy art style",
width: 768,
height: 1024
})
// Reproducible generation with seed
generate_image({
prompt: "majestic mountain landscape",
seed: 42 // Same seed = same image
})
// Batch generation - create multiple variations
generate_image({
prompt: "fantasy dragon in flight, epic lighting",
batch_size: 4 // Generate 4 images at once
})
// Batch with fixed seed (incremental seeds for variations)
generate_image({
prompt: "steampunk airship",
seed: 100, // Images will use seeds: 100, 101, 102
batch_size: 3
})
```
#### 2. upscale_image
Upscale images using AI models to 4x resolution (e.g., 1024x1024 → 4096x4096).
**Function Signature:**
```javascript
upscale_image({
image_path: string, // Required: Path to image file (e.g., "flux_output_00001_.png")
model?: string, // Optional: Upscaling model to use (default: 'ultrasharp')
// Options: 'ultrasharp', 'animesharp'
scale_factor?: number, // Optional: Additional scaling factor (default: 1.0)
// Range: 0.5 to 2.0 (1.0 = 4x native upscaling)
content_type?: string // Optional: Auto-select model based on content
// Options: 'general', 'anime', 'artwork', 'illustration'
})
```
**Usage Examples:**
```javascript
// General purpose upscaling with UltraSharp
upscale_image({
image_path: "flux_output_00001_.png",
model: "ultrasharp"
})
// Anime/illustration upscaling with AnimeSharp
upscale_image({
image_path: "bg_removed_00001_.png",
model: "animesharp"
})
// Auto-select model based on content type
upscale_image({
image_path: "output/my_image.png",
content_type: "anime" // Will use animesharp model
})
```
**Available Models:**
- **4x-UltraSharp**: General purpose upscaler, excellent for photos and realistic images
- **4x-AnimeSharp**: Optimized for anime, illustrations, and artwork
**Features:**
- 4x native upscaling (1024x1024 → 4096x4096)
- Preserves image quality and enhances details
- Smart model selection based on content type
- Returns file path (no base64 to avoid memory issues with large images)
- Outputs to `output/upscaled_*.png`
#### 3. remove_background
Remove background from images using RMBG-2.0 AI model.
**Function Signature:**
```javascript
remove_background({
image_path: string, // Required: Path to image file (e.g., "output/image.png")
alpha_matting?: boolean, // Optional: Use alpha matting for edges (default: true)
output_format?: string // Optional: Output format 'png' or 'webp' (default: 'png')
})
```
**Usage Examples:**
```javascript
// Simple background removal
remove_background({
image_path: "flux_output_00001_.png"
})
// Without alpha matting for sharp edges
remove_background({
image_path: "output/my_image.png",
alpha_matting: false
})
```
**Features:**
- Uses state-of-the-art RMBG-2.0 model for accurate segmentation
- Supports images from output directory (generated images)
- Automatic alpha channel for transparency
- Optional alpha matting for better edge quality
- Outputs to `output/bg_removed_*.png`
#### 4. check_models
Verify available models in ComfyUI.
**Function Signature:**
```javascript
check_models()
```
Returns list of available FLUX models and their status.
#### 4. connect_comfyui
Manually connect to ComfyUI server (usually auto-connects).
**Function Signature:**
```javascript
connect_comfyui({
server_address?: string // Optional: ComfyUI address (default: uses Docker network)
})
```
#### 5. disconnect_comfyui
Disconnect from ComfyUI server.
**Function Signature:**
```javascript
disconnect_comfyui()
```
## Docker Environment
### Container Architecture
1. **ComfyUI Container** (`mcp-comfyui-comfyui-1`)
- Ubuntu 22.04 with CUDA 12.1 (optimized base image)
- Python 3.11, PyTorch 2.5.1 (latest stable)
- FLUX schnell fp8 model (11GB)
- T5-XXL fp8 encoder (4.9GB)
- CLIP-L encoder (235MB)
- VAE decoder (320MB)
- Custom nodes:
- ComfyUI-Manager - Node management UI
- ComfyUI-KJNodes - Advanced image processing nodes
- ComfyUI-RMBG - Background removal with RMBG-2.0 model
- Upscaling models:
- 4x-UltraSharp (64MB) - General purpose
- 4x-AnimeSharp (64MB) - Anime/illustrations
- Runs with `--highvram` flag for optimal GPU usage
- Multi-stage build with BuildKit cache mounts for faster rebuilds
- Image size: ~10.9GB (optimized from 14.6GB)
2. **MCP Server Container** (`mcp-comfyui-flux-mcp-server-1`)
- Node.js 20 Alpine
- MCP protocol implementation
- Auto-connects to ComfyUI on startup
- WebSocket client for ComfyUI communication
### Development Commands
```bash
# Build and manage containers
./build.sh --help # Show all build options
./build.sh # Build containers
./build.sh --start # Build and start
./build.sh --no-cache # Rebuild from scratch
# Docker commands for development
docker-compose -p mcp-comfyui up -d # Start services
docker-compose -p mcp-comfyui logs -f # View logs
docker-compose -p mcp-comfyui restart # Restart after changes
docker-compose -p mcp-comfyui down # Stop services
# Access containers
docker exec -it mcp-comfyui-comfyui-1 bash # ComfyUI container
docker exec -it mcp-comfyui-flux-mcp-server-1 sh # MCP server container
# Run tests
docker exec mcp-comfyui-flux-mcp-server-1 npm test
docker exec mcp-comfyui-flux-mcp-server-1 node examples/example.js
# Health check
./scripts/health-check.sh
```
## Code Architecture
### Core Components
1. **MCP Server (`src/index.js`)**: Main entry point
- Implements MCP protocol with auto-connect
- Optimized defaults for FLUX schnell
- Returns base64-encoded images
- Uses `getFluxDiffusersWorkflow` for fp8 models
2. **ComfyUI Client (`src/comfyui-client.js`)**: WebSocket client
- Persistent WebSocket connection
- Queue management with prompt tracking
- Image retrieval and base64 encoding
- Error recovery with retries
3. **Workflow Templates**: ComfyUI workflows
- **FLUX Generation** (`src/flux-workflow.js`):
- `getFluxDiffusersWorkflow`: Used for fp8 models with separate loaders
- Configured for schnell: `flux1-schnell-fp8-e4m3fn.safetensors`
- T5-XXL: `t5xxl_fp8_e4m3fn_scaled.safetensors`
- CLIP-L: `clip_l.safetensors`
- **Background Removal** (`src/workflows/background-removal.js`):
- Uses RMBG-2.0 model for segmentation
- Optional alpha matting for edge refinement
- **Image Upscaling** (`src/workflows/upscaling.js`):
- Native ComfyUI upscaling nodes
- 4x resolution enhancement
- Model-specific optimization
### Communication Flow
```
Claude Desktop → MCP Server (via WSL2)
↓
WebSocket (Docker network)
↓
ComfyUI Container → GPU/CUDA → Generated/Processed Image
↓
File saved to output/ → MCP Response → Claude Desktop
↓
Returns: File path + metadata (avoids base64 for large images)
```
## Key Implementation Details
### Docker Networking
- Services communicate via `mcp-network` bridge
- ComfyUI accessible as `comfyui:8188` from MCP server
- Host access via `localhost:8188`
### Memory Configuration
- WSL2: 20GB RAM allocated (via .wslconfig)
- Docker: 16GB shared memory (`shm_size: "16g"`)
- GPU: 24GB VRAM (RTX 4090)
- Models use ~10GB VRAM with fp8 quantization
### Environment Variables
```bash
COMFYUI_HOST=comfyui # Docker service name
COMFYUI_PORT=8188 # ComfyUI port
PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512 # Memory optimization
MODEL_PRECISION=fp16 # Precision for operations
```
### FLUX Schnell Optimization
- **4-step generation**: Schnell is distilled for ultra-fast 4-step generation
- **Low CFG (1.0)**: Works best without strong guidance
- **Simple scheduler**: Matches the distilled model's training
- **FP8 quantization**: Reduces VRAM from 24GB to ~10GB with minimal quality loss
## Prompt Engineering Tips
### For Best Results with FLUX Schnell
1. **Be Descriptive**: Include details about style, lighting, mood
```
Good: "serene japanese garden, cherry blossoms, koi pond, morning mist, soft sunlight"
Poor: "garden"
```
2. **Specify Art Style**: FLUX responds well to style cues
```
"oil painting style", "anime art style", "photorealistic", "watercolor illustration"
```
3. **Use Quality Modifiers**: Add these to improve output
```
"high quality", "detailed", "sharp focus", "professional photography"
```
4. **Negative Prompts**: Tell it what to avoid
```
"blurry, low quality, distorted, ugly, deformed, bad anatomy"
```
5. **Composition Terms**: Guide the framing
```
"portrait", "full body", "close-up", "wide angle", "aerial view"
```
## Troubleshooting
### Common Issues
1. **Out of Memory (OOM)**:
- Reduce image dimensions (try 768x768)
- Ensure using fp8 models not fp16
- Check WSL2 memory allocation
2. **Blurry Images**:
- Ensure steps=4 (not 20)
- Use cfg_scale=1.0 (not 7.0)
- Add "sharp focus, detailed" to prompt
3. **Container Connection Issues**:
```bash
docker-compose -p mcp-comfyui restart
docker logs mcp-comfyui-flux-mcp-server-1 --tail 50
```
4. **MCP Not Detected in Claude**:
- Ensure config is in Windows path: `%APPDATA%\Claude\claude_desktop_config.json`
- Restart Claude Desktop completely
- Check with `/mcp` command
### Debug Commands
```bash
# Check GPU usage
docker exec mcp-comfyui-comfyui-1 nvidia-smi
# View MCP server logs
docker logs mcp-comfyui-flux-mcp-server-1 -f
# Test MCP directly
docker exec mcp-comfyui-flux-mcp-server-1 node examples/example.js
# Check model files
docker exec mcp-comfyui-comfyui-1 ls -la /app/ComfyUI/models/unet/
```
## Performance Metrics
### FLUX Schnell FP8 Performance
- **Generation Time**: ~2-4 seconds per image
- **Batch Generation**: ~1.5 seconds per additional image in batch
- **VRAM Usage**: ~10GB base + ~1GB per additional batch image
- **Quality**: 95% of fp16 quality at 50% memory usage
- **Optimal Resolution**: 1024x1024 (can do 768-1536px)
- **Max Batch Size**: 8 images (limited by VRAM)
### System Requirements
- **Minimum**: 16GB RAM, 12GB VRAM
- **Recommended**: 20GB+ RAM, 16GB+ VRAM
- **Current Setup**: 20GB RAM (WSL2), 24GB VRAM (RTX 4090)
### Docker Optimizations (BuildKit)
- **Multi-stage builds**: Separate build and runtime layers
- **Cache mounts**: Persistent pip and apt caches across builds
- **Layer optimization**: Minimized layer count and size
- **PyTorch 2.5.1**: Latest stable version with native RMSNorm support
- **Image size**: ~10.9GB (reduced from 14.6GB)
## Best Practices
1. **Always use relative symlinks** for models in Docker volumes
2. **Keep containers running** for instant generation
3. **Use seeds** for reproducible results during testing
4. **Monitor VRAM** with `nvidia-smi` during generation
5. **Restart containers** after code changes to `src/`
6. **Use `./build.sh`** for consistent builds with optimizations
7. **Avoid recursive chown** on large directories in WSL2
8. **Use BuildKit cache mounts** for faster Python package installs
## Files to Focus On
When working on this project, these are the key files:
1. **Core Logic**:
- `src/index.js` - MCP server with auto-connect and optimized defaults
- `src/comfyui-client.js` - ComfyUI WebSocket client
- `src/flux-workflow.js` - Workflow templates for fp8 models
2. **Configuration**:
- `docker-compose.yml` - Service orchestration
- `.env` - Environment variables
- `Dockerfile.comfyui` - Optimized ComfyUI container with BuildKit
- `Dockerfile.mcp` - MCP server container
- `requirements.txt` - Python dependencies (pinned versions)
3. **Scripts**:
- `build.sh` - Build script with options and optimizations
- `install.sh` - Automated setup with model detection
- `scripts/download-models.sh` - Model management
- `scripts/health-check.sh` - System validation
## Example Claude Desktop Usage
```
User: "Generate an image of a cyberpunk cat"
Claude: I'll generate a cyberpunk cat image for you.
[Calls generate_image with optimized defaults]
User: "Make it more neon with rain"
Claude: I'll regenerate with more neon and rain effects.
[Calls generate_image with enhanced prompt]
User: "I want the same image but higher quality"
Claude: I'll regenerate with the same seed for consistency.
[Uses same seed with quality modifiers in prompt]
User: "Give me 4 different versions of a space station"
Claude: I'll generate 4 variations of a space station for you.
[Uses batch_size: 4 for efficient parallel generation]
User: "Upscale that image to 4K resolution"
Claude: I'll upscale your image to 4096x4096 resolution.
[Calls upscale_image with ultrasharp model]
User: "Can you upscale this anime artwork?"
Claude: I'll upscale your anime artwork using the specialized anime model.
[Calls upscale_image with animesharp model for optimal results]
User: "Remove the background and then upscale it"
Claude: I'll first remove the background and then upscale the result.
[Calls remove_background, then upscale_image on the output]
```
## Recent Improvements
- **PyTorch 2.5.1**: Upgraded from 2.2.0, native RMSNorm support
- **BuildKit optimizations**: Multi-stage builds with cache mounts
- **Image size reduction**: 10.9GB (down from 14.6GB, 25% reduction)
- **Fixed dependencies**: All Python packages pinned for reproducibility
- **Custom nodes integrated**: KJNodes, RMBG, Manager built into image
- **WSL2 optimizations**: Avoided recursive chown issues
- **aiohttp/yarl compatibility**: Fixed version incompatibilities
- **Build script**: New `build.sh` with options and monitoring
Remember: Everything runs in optimized containers with fp8 models for fast, efficient generation!
