Enhanced Model Detection, UNet Workflow Support (z-image workflows) and VSCode mcp config

[eduardojaime]

Summary

This PR adds support for modern UNet-based and GGUF-format diffusion models (Flux, SD3, z_image, etc.) by:

1. Categorizing models by loader type - Detects checkpoints, UNet models, and diffusion models separately
2. New GGUF workflow - Adds generate_image_unet.json for GGUF models requiring specialized loaders
3. Enhanced list_models tool - Returns categorized model list with counts and clear organization
4. Comprehensive tests - Adds unit tests for categorized model detection

Motivation

ComfyUI supports multiple model architectures that use different loader nodes:

• CheckpointLoaderSimple: Traditional checkpoints (SD 1.5, SDXL)
• UNETLoader: Modern transformer models in safetensors format (Flux, SD3, Wanxiang)
• LoaderGGUF: GGUF-quantized models (z_image, specialized models)
• DiffusionModelLoader: Specialized diffusion transformers

Previously, the server only detected checkpoint models, making it impossible to use newer architectures. This PR enables full support for all model types including GGUF-format models.

Changes

1. Enhanced Model Detection (comfyui_client.py)

Added _get_available_models_categorized() method that:

• Queries three ComfyUI API endpoints (/object_info/CheckpointLoaderSimple, /object_info/UNETLoader, /object_info/DiffusionModelLoader)
• Returns a dictionary with three categories: checkpoints, unet, diffusion_models
• Maintains backward compatibility: available_models still defaults to checkpoint list

Key Code:

def _get_available_models_categorized(self):
    """Fetch categorized list of all model types from ComfyUI"""
    all_models = {
        "checkpoints": [],
        "unet": [],
        "diffusion_models": []
    }
    # Query each endpoint and populate categories
    # ...
    return all_models

2. Updated list_models Tool (tools/configuration.py)

Modified to return categorized structure:

Before:

{
  "models": ["model1.ckpt", "model2.safetensors"],
  "count": 2,
  "default": "model1.ckpt"
}

After:

{
  "models": {
    "checkpoints": ["model1.ckpt"],
    "unet": ["flux_dev.safetensors", "z_image_turbo_bf16.safetensors"],
    "diffusion_models": []
  },
  "counts": {
    "checkpoints": 1,
    "unet": 2,
    "diffusion_models": 0,
    "total": 3
  },
  "default_checkpoint": "model1.ckpt"
}

3. New Workflow: generate_image_unet.json

Adds support for GGUF-format models with specialized component loading:

Key Nodes:

• LoaderGGUF (node 4): Loads GGUF-format transformer models
• ClipLoaderGGUF (node 10): Loads CLIP encoder for GGUF models (hardcoded: qwen_3_4b.safetensors)
• VAELoader (node 3): Loads VAE separately
• EmptySD3LatentImage (node 7): Creates SD3-compatible latent space
• KSampler, CLIPTextEncode, VAEDecode, SaveImage: Standard generation pipeline

Parameters:

• PARAM_MODEL: GGUF model name (e.g., z_image_turbo-Q8_0.gguf)
• PARAM_VAE: VAE model (e.g., ae_zimage.safetensors)
• Standard parameters: PARAM_PROMPT, PARAM_NEGATIVE_PROMPT, PARAM_INT_WIDTH, etc.
• Fixed parameters: cfg=1.0, steps=8 (optimized for z_image turbo models)

Usage Example:

generate_image_unet(
    prompt="a mighty mexican jaguar, powerful and majestic",
    model="z_image_turbo-Q8_0.gguf",
    vae="ae_zimage.safetensors",
    width=1024,
    height=1024,
    steps=8,
    cfg=1.0
)

4. Comprehensive Tests (tests/test_model_categorization.py)

Adds test coverage for:

• Categorized model detection across all three loader types
• Graceful degradation when some endpoints are unavailable
• Model refresh updates categorized data correctly
• list_models tool returns correct structure with counts

Test Classes:

• TestModelCategorization: Tests ComfyUI client behavior
• TestListModelsTool: Tests MCP tool response format

5. Documentation Updates

docs/REFERENCE.md:

• Updated list_models documentation with new return structure
• Added model category explanations
• Included usage examples for discovering modern models

README.md:

• Updated Configuration Tools section to mention enhanced model support
• Added "Included Workflows" subsection listing all three default workflows
• Clarified which workflow to use for different model types

Backward Compatibility

✅ Fully backward compatible

• comfyui_client.available_models still contains checkpoint models (maintains existing behavior)
• Existing code using list_models will receive categorized structure with additional information
• No breaking changes to existing workflows or tool signatures

Testing

Manual Testing

# Run all tests
pytest tests/ -v

# Run model categorization tests specifically
pytest tests/test_model_categorization.py -v

Integration Testing

1. Start ComfyUI with multiple model types installed
2. Run MCP server: python server.py
3. Call list_models tool - should show all three categories
4. Use generate_image_unet with a UNet model

Model Compatibility

Supported in generate_image.json (CheckpointLoaderSimple)

• Stable Diffusion 1.5
• Stable Diffusion XL
• Any traditional checkpoint model

Supported in generate_image_unet.json (LoaderGGUF)

• GGUF Models:
  • z_image (z_image_turbo-Q8_0.gguf, z_image_bf16.gguf)
  • Other GGUF-quantized models requiring LoaderGGUF

Note: Models in safetensors format (Flux, SD3, Wanxiang) using UNETLoader require a different workflow structure not included in this PR.

Files Changed

modified:   comfyui_client.py
modified:   tools/configuration.py
modified:   docs/REFERENCE.md
modified:   README.md
new file:   workflows/generate_image_unet.json
new file:   tests/test_model_categorization.py

Future Enhancements

Potential follow-ups (not in this PR):

• Auto-detect appropriate workflow based on model type and format
• Additional workflow for safetensors UNet models (Flux, SD3) using UNETLoader + DualCLIPLoader
• Default VAE model suggestions based on model architecture
• LoRA and ControlNet categorization

Checklist

• Code changes implemented
• Tests added and passing
• Documentation updated
• Backward compatibility maintained
• Example workflow included
• No breaking changes

Questions for Reviewers

1. Should we add automatic workflow selection based on model category?
2. Should CLIP/VAE parameters have sensible defaults, or require explicit specification?
3. Is the categorization structure clear enough for AI agents to understand?