args

wfx.py

@@ -22,6 +22,8 @@ def parse_args():
     args.add_argument('--model', type=str, required=True)
     args.add_argument('--custom-pipeline', type=str, default=None)
     args.add_argument('--compile-mode', default='sfast', type=str, choices=['sfast', 'torch', 'no-compile'])
+    args.add_argument('--enable-cuda-graph', action='store_true', default=False)
+    args.add_argument('--disable-prefer-lowp-gemm', action='store_true', default=False)
     return args.parse_args()
 
 def quantize_unet(m):
@@ -68,7 +70,10 @@ class WFX():
             except ImportError:
                 logger.warning('triton not found, disabling triton')
         
-        self.compiler_config.enable_cuda_graph = True
+        self.compiler_config.enable_cuda_graph = args.enable_cuda_graph
+        
+        if args.disable_prefer_lowp_gemm:
+            self.compiler_config.prefer_lowp_gemm = False
         
         for key in self.compiler_config.__dict__:
             logger.info(f'cc - {key}: {self.compiler_config.__dict__[key]}')