Update gradio-app.py

gradio-app.py

@@ -33,12 +33,42 @@ pipe.set_progress_bar_config(disable=True)
 
 if TORCH_COMPILE:
     #optmode = 'max-autotune'
-    optmode = 'reduce-overhead'
-    pipe.text_encoder = torch.compile(pipe.text_encoder, mode=optmode)
-    pipe.tokenizer = torch.compile(pipe.tokenizer, mode=optmode)
-    pipe.unet = torch.compile(pipe.unet, mode=optmode)
-    pipe.vae = torch.compile(pipe.vae, mode=optmode)
+    #optmode = 'reduce-overhead'
+    #pipe.text_encoder = torch.compile(pipe.text_encoder, mode=optmode)
+    #pipe.tokenizer = torch.compile(pipe.tokenizer, mode=optmode)
+    #pipe.unet = torch.compile(pipe.unet, mode=optmode)
+    #pipe.vae = torch.compile(pipe.vae, mode=optmode)
+    doCompile = True
+    if doCompile:
+        config = CompilationConfig.Default()
 
+        try:
+            import xformers
+            config.enable_xformers = True
+        except ImportError:
+            print('xformers not installed, skipping')
+        try:
+            import triton
+            config.enable_triton = True
+        except ImportError:
+            print('Triton not installed, skipping')
+
+        config.enable_cuda_graph = True
+        config.enable_jit = True
+        config.enable_jit_freeze = True
+        config.trace_scheduler = False#True   CHECK THIS AGAIN
+        config.enable_cnn_optimization = True
+        config.preserve_parameters = False
+        config.prefer_lowp_gemm = True
+        config.enable_fused_linear_geglu = True
+
+        torch.jit.optimize_for_inference = True
+        torch.jit.enable_onednn_fusion = True
+        torch.jit.set_fusion_strategy([('STATIC', 1), ('DYNAMIC', 1)])
+
+        for p in pipe.text_encoder.parameters(): p.requires_grad=False
+        for p in pipe.vae.decoder.parameters(): p.requires_grad=False
+        for p in pipe.unet.parameters(): p.requires_grad=False
 
 def predict(prompt1, prompt2, merge_ratio, guidance, steps, sharpness, seed=1231231):
     torch.manual_seed(seed)