1 files changed, 24 insertions, 0 deletions

diff --git a/train_ti.py b/train_ti.py
index 26f7941..6fd974e 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -5,13 +5,16 @@ from functools import partial
 from pathlib import Path
 from typing import Union
 import math
+import warnings
 
 import torch
 import torch.utils.checkpoint
+import hidet
 
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import LoggerType, set_seed
+from diffusers.models.attention_processor import AttnProcessor
 from timm.models import create_model
 import transformers
 
@@ -28,10 +31,18 @@ from training.util import AverageMeter, save_args
 
 logger = get_logger(__name__)
 
+warnings.filterwarnings('ignore')
+
 
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.benchmark = True
 
+# torch._dynamo.config.log_level = logging.WARNING
+
+hidet.torch.dynamo_config.use_tensor_core(True)
+hidet.torch.dynamo_config.use_attention(True)
+hidet.torch.dynamo_config.search_space(0)
+
 
 def parse_args():
     parser = argparse.ArgumentParser(
@@ -706,6 +717,19 @@ def main():
     if args.use_xformers:
         vae.set_use_memory_efficient_attention_xformers(True)
         unet.enable_xformers_memory_efficient_attention()
+    elif args.compile_unet:
+        unet.mid_block.attentions[0].transformer_blocks[0].attn1._use_2_0_attn = False
+    
+        proc = AttnProcessor()
+    
+        def fn_recursive_set_proc(module: torch.nn.Module):
+            if hasattr(module, "processor"):
+                module.processor = proc
+    
+            for child in module.children():
+                fn_recursive_set_proc(child)
+    
+        fn_recursive_set_proc(unet)
 
     if args.gradient_checkpointing:
         unet.enable_gradient_checkpointing()