Update app.py

app.py

@@ -12,6 +12,9 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 import torch
+import hidet
+import diffusers
+
 from diffusers import AutoencoderKL, StableDiffusionXLPipeline
 #from diffusers import EulerAncestralDiscreteScheduler
 from typing import Tuple
@@ -100,6 +103,10 @@ def load_and_prepare_model():
     #pipe.vae.vae_scale_factor=8
     #pipe.to(device=device, dtype=torch.bfloat16)
     pipe.to(device)
+    
+    #Make sure the unet is contiguous
+    pipe.unet = pipe.unet.to(memory_format=torch.contiguous_format)
+
     pipe.to(torch.bfloat16)
     pipe.vae.set_default_attn_processor()
     print(f'init noise scale: {pipe.scheduler.init_noise_sigma}')
@@ -110,6 +117,25 @@ def load_and_prepare_model():
 # Preload and compile both models
 pipe = load_and_prepare_model()
 
+
+          # for compile
+hidet.option.parallel_build(True)
+torch._dynamo.config.suppress_errors = True
+torch._dynamo.disallow_in_graph(diffusers.models.attention.BasicTransformerBlock)
+# 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.int64).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' "