# How to Implement VAE Encoding and Decoding Nodes in ComfyUI

> Learn to implement VAE encoding and decoding nodes in ComfyUI using its declarative node architecture. Optimize for high-resolution images and video with standard and tiled operations.

- Repository: [Comfy Org/ComfyUI](https://github.com/Comfy-Org/ComfyUI)
- Tags: how-to-guide
- Published: 2026-02-26

---

**ComfyUI implements VAE encoding and decoding through a declarative node architecture in [`nodes.py`](https://github.com/Comfy-Org/ComfyUI/blob/main/nodes.py) that wraps the core `VAE` class in [`comfy/sd.py`](https://github.com/Comfy-Org/ComfyUI/blob/main/comfy/sd.py), supporting both standard and tiled operations for high-resolution images and video.**

Implementing **VAE encoding and decoding nodes** in ComfyUI requires understanding how the framework treats Variational Auto-Encoders as first-class data types that flow between nodes. The implementation spans declarative node definitions in [`nodes.py`](https://github.com/Comfy-Org/ComfyUI/blob/main/nodes.py) and the heavy-duty tensor processing logic in [`comfy/sd.py`](https://github.com/Comfy-Org/ComfyUI/blob/main/comfy/sd.py). This guide breaks down the exact implementation patterns used in the Comfy-Org/ComfyUI repository.

## Understanding the VAE Node Architecture in nodes.py

All VAE nodes in ComfyUI follow a strict declarative protocol that separates UI metadata from execution logic.

### The Declarative Node Protocol

Every VAE node class defines three critical class attributes that the ComfyUI engine uses to render sockets and route execution:

```python
class VAEDecode:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "samples": ("LATENT", {"tooltip": "The latent to be decoded."}),
                "vae": ("VAE", {"tooltip": "The VAE model used for decoding the latent."})
            }
        }
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "decode"
    CATEGORY = "latent"

```

**`INPUT_TYPES`** declares the input sockets, **`RETURN_TYPES`** specifies outputs, and **`FUNCTION`** maps the node to a Python method. The actual `decode` method remains a thin wrapper around the `VAE` object:

```python
def decode(self, vae, samples):
    latent = samples["samples"]
    if latent.is_nested:  # handle batched latent tensors

        latent = latent.unbind()[0]
    images = vae.decode(latent)  # delegate to VAE class

    if len(images.shape) == 5:   # flatten video batches

        images = images.reshape(-1, images.shape[-3],
                               images.shape[-2], images.shape[-1])
    return (images,)

```

### Core VAE Nodes and Their Implementation

The [`nodes.py`](https://github.com/Comfy-Org/ComfyUI/blob/main/nodes.py) file (lines 293-384) defines five primary VAE nodes that map directly to methods in the `VAE` class:

- **VAEDecode**: Calls `vae.decode(latent)` to convert latent tensors to RGB images
- **VAEEncode**: Calls `vae.encode(pixels)` to compress images into latent space  
- **VAEDecodeTiled**: Invokes `vae.decode_tiled(...)` for memory-efficient decoding of large images
- **VAEEncodeTiled**: Invokes `vae.encode_tiled(...)` for tiled encoding operations
- **VAEEncodeForInpaint**: Wraps `vae.encode()` with mask preprocessing for inpainting workflows (lines 94-124)

These nodes are registered in the global `NODE_CLASS_MAPPINGS` dictionary near line 2040, making them discoverable by the UI and API.

## The VAE Class Implementation in comfy/sd.py

The `VAE` class in [`comfy/sd.py`](https://github.com/Comfy-Org/ComfyUI/blob/main/comfy/sd.py) (lines 370-440) encapsulates all model-specific logic, checkpoint handling, and tensor transformations.

### Checkpoint Handling and Model Detection

The constructor automatically detects VAE formats by probing for signature keys:

```python
class VAE:
    def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None):
        # Auto-convert Diffusers format checkpoints

        if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys():
            sd = diffusers_convert.convert_vae_state_dict(sd)
        
        self.downscale_ratio = 8
        self.upscale_ratio = 8
        self.latent_channels = 4
        self.output_channels = 3
        self.process_input = lambda img: img * 2.0 - 1.0  # Scale to [-1, 1]

        self.process_output = lambda img: torch.clamp((img + 1.0) / 2.0, 0.0, 1.0)

```

The class recognizes multiple VAE families including standard Stable Diffusion, TAESD (Tiny AutoEncoder), and Stable Cascade by checking for keys like `"taesd_decoder.1.weight"` or `"decoder.conv_in.weight"`.

### Encode and Decode Methods

The core methods handle the forward and inverse transformations between pixel and latent space:

```python
def encode(self, pixels):
    latent = self.first_stage_model.encode(self.process_input(pixels))
    return latent

def decode(self, latent):
    images = self.first_stage_model.decode(latent)
    return self.process_output(images)

```

Where `self.first_stage_model` is an instance of `AutoencodingEngine` or a TAESD wrapper. The `process_input` lambda scales images from [0, 1] to [-1, 1] range expected by the encoder, while `process_output` clamps decoder outputs back to valid RGB ranges.

### Tiled Operations for High-Resolution Processing

For memory-constrained environments, the VAE class provides tiled variants that process images in overlapping patches:

```python
def encode_tiled(self, pixel_samples, tile_x=None, tile_y=None, overlap=None,
                 tile_t=None, overlap_t=None):
    # Compute optimal tile size based on compression ratio

    return self.first_stage_model.encode_tiled(...)

def decode_tiled(self, latent, tile_x, tile_y, overlap, tile_t=None, overlap_t=None):
    # Handle spatial and temporal tiling for video batches

    return self.first_stage_model.decode_tiled(...)

```

These methods enable processing of 4K images or long video sequences without out-of-memory errors by splitting tensors along spatial dimensions (and temporal dimension `tile_t` for video).

## Loading VAE Models with VAELoader

The `VAELoader` node (lines 728-776 in [`nodes.py`](https://github.com/Comfy-Org/ComfyUI/blob/main/nodes.py)) handles checkpoint discovery and initialization:

```python
class VAELoader:
    @staticmethod
    def vae_list(s):
        # Merge exact checkpoint files with approximate TAESD families

        return files + taesd_files
    
    @staticmethod
    def load_taesd(name):
        # Assemble encoder/decoder from separate files (taesd_encoder.pt, taesd_decoder.pt)

        # Inject scale/shift constants required by the engine

```

When executed, the loader calls `comfy.sd.VAE(sd=sd, metadata=metadata)`, passing the merged state dictionary. The node outputs a `VAE` type that subsequent encoding/decoding nodes consume.

## Practical Implementation Examples

### API Workflow JSON

Below is a complete round-trip workflow demonstrating VAE encoding and decoding via the ComfyUI API:

```json
{
  "1": {
    "class_type": "VAELoader",
    "inputs": { "vae_name": "taesd" }
  },
  "2": {
    "class_type": "LoadImage",
    "inputs": { "image_path": "inputs/example.png" }
  },
  "3": {
    "class_type": "VAEEncode",
    "inputs": {
      "pixels": ["2", 0],
      "vae": ["1", 0]
    }
  },
  "4": {
    "class_type": "VAEDecode",
    "inputs": {
      "samples": ["3", 0],
      "vae": ["1", 0]
    }
  },
  "5": {
    "class_type": "SaveImage",
    "inputs": {
      "images": ["4", 0],
      "filename_prefix": "vae_roundtrip"
    }
  }
}

```

This JSON structure corresponds to the example in [`script_examples/basic_api_example.py`](https://github.com/Comfy-Org/ComfyUI/blob/main/script_examples/basic_api_example.py) (lines 73-84).

### Python Script Implementation

For programmatic use outside the node system:

```python
import torch
from comfy.sd import VAE
from comfy.utils import load_image, save_image, load_torch_file

# 1. Load VAE checkpoint

sd = load_torch_file("models/vae/taesdxl.pt", safe_load=True)
vae = VAE(sd=sd)  # Initialize the engine

# 2. Load source image [1, H, W, 3] in [0,1] range

image = load_image("inputs/portrait.png")

# 3. Encode to latent space

latent = vae.encode(image)

# 4. Decode back to RGB

reconstructed = vae.decode(latent)

# 5. Save output

save_image(reconstructed, "outputs/reconstructed.png")

```

### Tiled Video Processing

For video workflows with 5D tensors `[B, T, C, H, W]`:

```python

# Decode with both spatial and temporal tiling

decoded = vae.decode_tiled(
    latent_video,
    tile_x=512, tile_y=512, overlap=64,
    tile_t=64, overlap_t=8
)

```

## Summary

- **VAE nodes** in ComfyUI follow a declarative pattern using `INPUT_TYPES`, `RETURN_TYPES`, and `FUNCTION` class attributes defined in [`nodes.py`](https://github.com/Comfy-Org/ComfyUI/blob/main/nodes.py).
- The **VAE class** in [`comfy/sd.py`](https://github.com/Comfy-Org/ComfyUI/blob/main/comfy/sd.py) handles format detection, tensor scaling, and provides both standard and tiled encode/decode methods.
- **Tiled operations** (`encode_tiled`, `decode_tiled`) prevent out-of-memory errors when processing high-resolution images or video sequences.
- **VAELoader** manages checkpoint discovery including special handling for TAESD models that split encoder and decoder into separate files.
- Nodes pass VAE objects as typed references, allowing different VAE models to be chained within the same workflow graph.

## Frequently Asked Questions

### What is the difference between VAEEncode and VAEEncodeTiled?

**VAEEncode** processes entire images in a single forward pass through `vae.encode()`, requiring sufficient GPU memory for the full resolution. **VAEEncodeTiled** splits large images into overlapping tiles (configurable via `tile_x`, `tile_y`, and `overlap` parameters) and processes them sequentially through `vae.encode_tiled()`, significantly reducing memory usage at the cost of slightly longer processing time.

### How does ComfyUI handle different VAE formats like TAESD?

The `VAE` class constructor detects TAESD checkpoints by probing for keys like `"taesd_decoder.1.weight"`. The `VAELoader` node assembles TAESD models from separate encoder and decoder files (e.g., `taesd_encoder.pt` and `taesd_decoder.pt`) and injects required scale/shift constants. The framework also auto-converts Diffusers-format checkpoints by detecting keys like `"decoder.up_blocks.0.resnets.0.norm1.weight"` and applying `diffusers_convert.convert_vae_state_dict()`.

### Can I use different VAEs for encoding and decoding in the same workflow?

Yes. Because ComfyUI treats VAE as a distinct data type that flows through graph edges, you can route the output of one `VAELoader` into a `VAEEncode` node and a different `VAELoader` into a subsequent `VAEDecode` node. This enables workflows such as encoding with a lightweight TAESD for speed, then decoding with a full-quality SD-VAE for final output quality.

### What are the input and output tensor formats for VAE nodes?

**Inputs**: `VAEEncode` expects `IMAGE` tensors with shape `[batch, height, width, 3]` and values in `[0, 1]` range. `VAEDecode` expects `LATENT` dictionaries containing a `"samples"` key with tensors of shape `[batch, channels, height//8, width//8]`. **Outputs**: `VAEEncode` returns `LATENT` dictionaries, while `VAEDecode` returns `IMAGE` tensors of shape `[batch, height, width, 3]`. Video inputs use 5D tensors `[batch, frames, channels, height, width]` which the nodes flatten appropriately.