Nó customizado do ComfyUI que acelera a inferência do Z-Image em 20–30%

Olá. Impulsionado pelo interesse que vocês demonstraram no meu projeto anterior de otimização de kernel Triton para o Qwen3-TTS, compartilho agora meu segundo projeto open source.

É um nó customizado do ComfyUI que acelera em 20–30% a inferência do popular modelo de geração de imagens Z-Image S3-DiT (6.15B), que registra mais de 3,6 milhões de downloads mensais no HuggingFace.

1. Por que eu criei isso? (contexto e características)
Atualmente existe o Nunchaku(SVDQuant) como ferramenta para acelerar o Z-Image, mas ele só oferece suporte ao modelo Z-Image 'Turbo'. Achei que era necessária uma solução de aceleração em nível de kernel para o modelo Base.

Além disso, do ponto de vista do usuário, baixar novamente modelos quantizados com dezenas de GB (GGUF etc.) é um grande peso. Por isso, fiz com que fosse possível quantizar imediatamente em tempo de execução (On-the-fly Quantization) o modelo BF16 safetensors que você já usa e utilizá-lo diretamente.

• Pode ser instalado com um clique pelo ComfyUI Manager ou facilmente com pip install. (Sem o estresse de builds customizados de CUDA complicados ou de compatibilizar versões.)
• Basta adicionar apenas 1 nó ao workflow existente, com compatibilidade perfeita com LoRA e ControlNet.

2. Benchmark de desempenho (RTX 5090, com base em 30 steps)

• T2I Baseline: 18.9s → Triton + INT8: 15.3s (ganho de velocidade de 1.24x)
• LoRA Baseline: 19.0s → LoRA + Triton + INT8: 14.6s (ganho de velocidade de 1.30x)
• Economia de VRAM: total de 23GB → 19.5GB (redução de cerca de 3.5GB)

3. Preservação de qualidade verificada visualmente (sem cherry-picking)
No projeto anterior de TTS, havia o incômodo de precisar baixar e ouvir diretamente os áudios para verificar os resultados, mas desta vez é possível comparar a qualidade imediatamente na web.

Não houve nenhum cherry-picking para a comparação de desempenho. Devido às características da fusão de kernels e da quantização, existem pequenas mudanças em nível de pixel, mas a composição geral e a qualidade dos detalhes foram preservadas de forma excelente. Confira você mesmo, no link abaixo, as imagens originais de comparação de todos os cenários.

• Ver resultados do benchmark: https://github.com/newgrit1004/ComfyUI-ZImage-Triton/…

4. Pontos de engenharia
Neste código de kernel também contei ativamente com a ajuda do Claude Code, e em troca concentrei toda a minha energia em benchmarking rigoroso e validação de qualidade.

• Aplicação de 6 kernels de fusão Triton (RMSNorm, SwiGLU, QK-Norm+RoPE, Norm+Gate+Residual, AdaLN, RoPE 3D)
• Manutenção máxima da qualidade de quantização ao dispersar outliers por meio de W8A8 + Hadamard Rotation (baseado em QuaRot / ConvRot, NeurIPS 2024)

5. Aviso de atualização do projeto anterior
Aproveitando para comentar, também pretendo em breve levar o projeto qwen3-tts-triton, que publiquei anteriormente, para um nó customizado do ComfyUI. (atualização v0.2.0: mitigação de pronúncia embolada por meio de híbrido Triton+PyTorch, aplicação do TurboQuant, troca da ferramenta de avaliação da Cohere etc.)

No momento, concluí os testes no meu ambiente pessoal (RTX 5090). Se vocês puderem testar em GPUs da série 30/40 ou em outros ambientes e me enviar feedback por GitHub Issues ou comentários, isso ajudaria muito. Obrigado!

• GitHub: https://github.com/newgrit1004/ComfyUI-ZImage-Triton