This adaptation of RealRebelAI/LLaDa-Image_ComfyUI adds independent diffusion, GGUF text-encoder and VAE loaders, with a staged memory runtime for the native LLaDA architecture.
Status: experimental. CPU numerical and reduced-model integration tests are provided. Full-checkpoint CUDA generation, image-editing quality, peak VRAM and speed have not been verified.
See performance comparison instructions. The assembler now has an optional linear_backend setting. dequantize preserves the existing weight-only computation; comfy_kitchen opts into Kitchen INT8 linear computation with dynamically quantized activations. Small block parameters can remain resident during denoising, and logs provide stage timings and cache counters. GPU speedup and full-model image quality need local comparison.
Follow the CachyOS setup guide for installation with the existing runtime protected. Run tools/preflight.py check using ComfyUI's Python to inspect dependencies and optional backend imports.
INT8 linear layers use floating-point dequantization and PyTorch linear operations. Selecting comfy_kitchen switches those layers to Comfy Kitchen int8_linear; the default remains dequantize. Diffusion attention follows the selected ComfyUI backend, including Sage, Flash or Comfy Kitchen INT8 attention when supported by the host. Masked Flash calls can fall back to SDPA. Text-encoder attention retains its own implementation.
Full BF16 transformer checkpoints are accepted by the same diffusion loader. workflows/modular-bf16.json uses the full BF16 diffusion file with a GGUF text encoder and the same VAE. The INT8 workflow remains the starting point for lower memory use.
Load workflows/modular-int8.json in ComfyUI. It connects:
- LLaDA Load Diffusion Model →
LLADA_MODEL - LLaDA Load GGUF Text Encoder (CLIP) →
LLADA_CLIP - LLaDA Load VAE →
LLADA_VAE
These feed LLaDA Assemble Pipeline, then LLaDA Image Text to Image, then native Save Image. workflows/modular-int8-edit.json adds a native Load Image node for editing. Select an actual source image before queuing it. An API-format text workflow is also included.
The sockets deliberately identify LLaDA-specific objects. The text encoder is LLaDA2-MoE, and requires QueryFormer and text projection. These outputs cannot be wired to native KSampler, CLIP Text Encode or VAE Decode. Separate loader nodes provide independent Comfy graph caching; they do not establish compatibility with native MODEL/CLIP/VAE wrappers.
Use the actual Hugging Face file list. The model card contains older filenames.
| Location inside ComfyUI | Current filename | Approximate file size |
|---|---|---|
models/diffusion_models/ |
LLaDA-Image-Turbo-INT8.safetensors |
6.59 GB |
models/text_encoders/ |
LLaDA-Image-Turbo-text_encoder-Q4_K_M.gguf |
9.18 GB |
models/vae/ |
LLaDa_VAE.safetensors |
168 MB |
Registered extra model directories are supported. The INT8 transformer contains int8_tensorwise markers, packed weights and scales. Computation expands individual weights to the selected floating dtype; this is weight-only INT8, with no claim of native INT8 matrix-multiplication acceleration.
The first assembly downloads tokenizer, scheduler, QueryFormer and text-projection assets from inclusionAI/LLaDA-Image-Turbo. The staged path downloads and instantiates SigVQ only when editing requests it. Hugging Face's normal cache and offline settings apply. Model-file sizes do not equal total working RAM: auxiliary models, dequantization and host paging also matter.
Place this repository in ComfyUI/custom_nodes/ and install City96 ComfyUI-GGUF. Avoid keeping two copies of this LLaDA integration enabled, since they register the same legacy node identifiers.
From this repository, using ComfyUI's Python environment:
python -m pip install -r requirements.txtRestart ComfyUI. The requirements intentionally omit PyTorch, torchvision, CUDA and attention-kernel packages to preserve the host installation. Dependency versions used for testing are recorded in docs/VALIDATION.md.
Use matching bfloat16 on all loaders, VAE tiling On, 512 × 512, four steps, guidance 1.0. The assembler selects the staged runtime. Its reserve is at least ComfyUI's inference-memory allowance and includes ComfyUI's additional reserved memory. Increase reserve_mb when activation pressure causes allocation failures. Higher resolutions and guidance above 1 increase memory and execution cost.
Measure a first run separately from subsequent runs: matching prompts can reuse cached conditioning. Compare the same seed, prompt and settings when measuring performance. A change to negative conditioning, sequence length or other encoder inputs invalidates that cache.
- Uses ComfyUI's selected execution device and asks its memory manager to release managed models before allocating an active stage.
- Streams ModuleList blocks, retaining the remaining active-stage tensors on the execution device. Quantized storage remains CPU-owned.
- Retains a bounded subset of packed transformer weights during denoising. Expansion uses an in-place scale operation to avoid a second full floating-weight allocation.
- Dequantizes only unique requested embedding rows. Expert banks retain upstream selected-expert dequantization.
- Routes diffusion attention through ComfyUI's attention selector. Supported attention backends follow the host configuration. Text-encoder attention retains its architecture-specific implementation.
- Keeps one detached CPU conditioning-cache entry per pipeline. The legacy pipeline registry uses weak references.
- Reports step progress and checks interruption at stage/block boundaries. Generation cleanup runs on normal completion and exceptions. GPU kernels already running finish before cancellation can take effect.
- Enables the Diffusers VAE's tiled decode when selected. SigVQ stays unloaded during text generation.
This is cooperative memory management around a custom Diffusers pipeline. It is not a complete ModelPatcher port. Native LoRA patching, every ComfyUI model-management option, asynchronous weight streaming and native sampling nodes are not implemented. The runtime does not enable torch.compile or change global TF32 settings. Adding every optimization without verifying architectural compatibility would be unsafe.
The legacy combined loader remains available. Its default for new nodes is comfy_staged; existing workflows retain their saved offload selection. The old CUDA and Accelerate paths remain legacy paths and do not receive the staged runtime's guarantees. Prefer the modular examples.
python -m unittest discover -s tests -vRead docs/VALIDATION.md for coverage and outstanding GPU checks. Memory savings and throughput must be measured with real checkpoints. If a run fails, capture the complete traceback and selected settings. Preserve the original environment until the adapted workflow passes locally.
The Unload node clears device allocations and the conditioning cache. Comfy's graph cache may still hold the separate loaded CPU components. Clearing that graph cache or restarting ComfyUI releases those references.
- inclusionAI/LLaDA-Image: original model and architecture.
- RealRebelAI: original integration and quantized runtime.
- City96 ComfyUI-GGUF: GGUF loading/dequantization dependency.
Existing source copyright and license notices are preserved. The upstream model and weight terms continue to apply. This adaptation does not grant additional rights to the original code or weights.
Auxiliary loading never downloads or checks the Hub during assembly. See offline setup for explicit first-install downloads.