owo

wfx.py

@@ -21,6 +21,7 @@ def parse_args():
     args.add_argument('--quantize-unet', action='store_true', default=False)
     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'])
     return args.parse_args()
 
 def quantize_unet(m):
@@ -91,8 +92,14 @@ class WFX():
         if args.quantize_unet:
             self.T2IPipeline.unet = quantize_unet(self.T2IPipeline.unet)
         
-        logger.info('compiling model...')
-        self.T2IPipeline = compile(self.T2IPipeline, self.compiler_config)
+        logger.info(f'compiling pipeline in {args.compile_mode} mode...')
+        if args.compile_mode == 'sfast':
+            self.T2IPipeline = compile(self.T2IPipeline, self.compiler_config)
+        elif args.compile_mode == 'torch':
+            logger.info('compiling unet...')
+            self.T2IPipeline.unet = torch.compile(self.T2IPipeline.unet, mode='max-autotune')
+            logger.info('compiling vae...')
+            self.T2IPipeline.vae = torch.compile(self.T2IPipeline.vae, mode='max-autotune')
         
         self.warmup()