diff --git a/backend/args.py b/backend/args.py index f86af015..917f90e3 100644 --- a/backend/args.py +++ b/backend/args.py @@ -103,6 +103,7 @@ parser.add_argument("--mmap-torch-files", action="store_true", help="Use mmap wh parser.add_argument("--disable-mmap", action="store_true", help="Don't use mmap when loading safetensors") parser.add_argument("--tiled-conv2d", type=int, default=0, metavar="TILE_SIZE", choices=[0, 64, 128, 256, 512], help="reduce VAE memory usage ; increase processing time") +parser.add_argument("--enable-triton-backend", action="store_true", help="Enable the use of Triton backend in comfy-kitchen") class SageAttentionFuncs(enum.Enum): diff --git a/backend/quant_ops.py b/backend/quant_ops.py index d29f25c4..c3f8236e 100644 --- a/backend/quant_ops.py +++ b/backend/quant_ops.py @@ -22,9 +22,10 @@ else: if cuda_version < (13,): ck.registry.disable("cuda") +from backend.args import args -# https://github.com/Comfy-Org/ComfyUI/blob/v0.16.4/comfy/quant_ops.py#L24 -ck.registry.disable("triton") +if not args.enable_triton_backend: + ck.registry.disable("triton") # region FP8 Layouts