Nodo personalizado de ComfyUI que acelera la inferencia de Z-Image entre 20% y 30%

Hola. Gracias al interés que mostraron en el proyecto anterior de optimización de kernels Triton para Qwen3-TTS, comparto mi segundo proyecto open source.

Se trata de un nodo personalizado de ComfyUI que acelera entre 20% y 30% la inferencia de Z-Image S3-DiT (6.15B), un popular modelo de generación de imágenes que registra más de 3.6 millones de descargas mensuales en HuggingFace.

1. ¿Por qué lo hice? (contexto y características)
Actualmente existe Nunchaku(SVDQuant) como herramienta para acelerar Z-Image, pero solo es compatible con el modelo Z-Image "Turbo". Me pareció que hacía falta una solución de aceleración a nivel kernel para el modelo Base.

Además, desde la perspectiva del usuario, volver a descargar modelos cuantizados de decenas de GB (como GGUF) supone una carga importante. Por eso lo hice de forma que el modelo BF16 en safetensors que ya usabas pueda cuantizarse al instante en tiempo de ejecución (On-the-fly Quantization) y seguir utilizándose tal cual.

• Se puede instalar con un solo clic desde ComfyUI Manager o fácilmente con pip install. (Sin el estrés de builds personalizados de CUDA ni de hacer coincidir versiones.)
• Solo hace falta añadir 1 nodo al workflow existente, y es totalmente compatible con LoRA y ControlNet.

2. Benchmark de rendimiento (RTX 5090, basado en 30 steps)

• T2I Baseline: 18.9s → Triton + INT8: 15.3s (mejora de velocidad de 1.24x)
• LoRA Baseline: 19.0s → LoRA + Triton + INT8: 14.6s (mejora de velocidad de 1.30x)
• Ahorro de VRAM: total de 23GB → 19.5GB (reducción de aprox. 3.5GB)

3. Conservación de calidad verificable a simple vista (sin cherry-picking)
En el proyecto anterior de TTS, para comprobar los resultados había que descargar y escuchar el audio manualmente, lo cual era incómodo; esta vez, la comparación de calidad puede verse al instante en la web.

No hubo ningún tipo de cherry-picking para comparar el rendimiento. Por la naturaleza de la fusión de kernels y la cuantización, existen cambios mínimos a nivel de píxel, pero la composición general y la calidad del detalle se conservan de forma excelente. Revisa en el siguiente enlace las imágenes originales comparativas de todos los escenarios.

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

4. Puntos de ingeniería
Este código de kernels también recibió una fuerte ayuda de Claude Code; por mi parte, enfoqué toda mi energía en el benchmarking riguroso y la validación de calidad.

• Aplicación de 6 kernels fusionados de Triton (RMSNorm, SwiGLU, QK-Norm+RoPE, Norm+Gate+Residual, AdaLN, RoPE 3D)
• Con W8A8 + Hadamard Rotation (basado en QuaRot / ConvRot de NeurIPS 2024), se dispersan los outliers para preservar al máximo la calidad de la cuantización

5. Próxima actualización del proyecto anterior
Además, planeo llevar pronto el proyecto qwen3-tts-triton, que publiqué anteriormente, a un nodo personalizado de ComfyUI. (Actualización v0.2.0: mitigación del habla pastosa mediante híbrido Triton+PyTorch, aplicación de TurboQuant, reemplazo de la herramienta de evaluación de Cohere, etc.)

Ya terminé las pruebas en mi entorno personal (RTX 5090). Si lo prueban en GPUs serie 30/40 o en otros entornos y me dejan feedback por GitHub Issues o en los comentarios, sería de muchísima ayuda. ¡Gracias!

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