Update app.py

app.py

@@ -11,6 +11,8 @@ import uuid
 import gradio as gr
 import numpy as np
 from PIL import Image
+import hidet
+
 import diffusers
 from diffusers import AutoencoderKL, StableDiffusionXLPipeline
 from diffusers import EulerAncestralDiscreteScheduler
@@ -20,6 +22,7 @@ import datetime
 import cyper
 from image_gen_aux import UpscaleWithModel
 import torch
+import torch._dynamo
 
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
@@ -148,19 +151,39 @@ def load_and_prepare_model():
     pipe.watermark=None
     pipe.safety_checker=None  
     
-            # Freeze vae and unet
+    '''       # Freeze vae and unet
     pipe.vae.requires_grad_(False)
     pipe.unet.requires_grad_(False)
     pipe.text_encoder.requires_grad_(False)
     pipe.unet.eval()
     pipe.vae.eval()
     pipe.text_encoder.eval()
+    '''
+
+    pipe.unet = pipe.unet.to(memory_format=torch.contiguous_format)
 
     return pipe
     
+hidet.option.parallel_build(True)
+torch._dynamo.config.suppress_errors = True
+torch._dynamo.disallow_in_graph(diffusers.models.attention.BasicTransformerBlock)
 # Preload and compile both models
 pipe = load_and_prepare_model()
 
+
+# more search
+hidet.torch.dynamo_config.search_space(0)
+#hidet.torch.dynamo_config.dump_graph_ir("./local_graph")
+hidet.option.cache_dir("local_cache")
+# automatically transform the model to use float16 data type
+#hidet.torch.dynamo_config.use_fp16(True)
+# use float16 data type as the accumulate data type in operators with reduction
+#hidet.torch.dynamo_config.use_fp16_reduction(True)
+# use tensorcore
+hidet.torch.dynamo_config.use_tensor_core()
+pipe.unet = torch.compile(pipe.unet, backend="hidet")
+
+
 MAX_SEED = np.iinfo(np.int32).max
 
 neg_prompt_2 = " 'non-photorealistic':1.5, 'unrealistic skin','unattractive face':1.3, 'low quality':1.1, ('dull color scheme', 'dull colors', 'digital noise':1.2),'amateurish', 'poorly drawn face':1.3, 'poorly drawn', 'distorted face', 'low resolution', 'simplistic' "