From 6cec284a6ff8d88cd53c74d7eadc8167481fa262 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Tue, 5 Dec 2023 11:49:59 -0500
Subject: [PATCH] cleanup compile

---
 modules/sd_models_compile.py | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py
index 67fc32740..ece694085 100644
--- a/modules/sd_models_compile.py
+++ b/modules/sd_models_compile.py
@@ -77,6 +77,7 @@ def compile_stablefast(sd_model):
     config.enable_cuda_graph = shared.opts.cuda_compile_fullgraph
     config.enable_jit_freeze = shared.opts.diffusers_eval
     config.memory_format = torch.channels_last if shared.opts.opt_channelslast else torch.contiguous_format
+    # config.trace_scheduler = False
     # config.enable_cnn_optimization
     # config.prefer_lowp_gemm
     try:
@@ -108,12 +109,17 @@ def compile_torch(sd_model):
         torch._dynamo.config.suppress_errors = shared.opts.cuda_compile_errors # pylint: disable=protected-access
         t0 = time.time()
         if shared.opts.cuda_compile:
-            sd_model.unet = torch.compile(sd_model.unet, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
+            if shared.opts.cuda_compile and (not hasattr(sd_model, 'unet') or not hasattr(sd_model.unet, 'config')):
+                shared.log.warning('Model compile enabled but model has no Unet')
+            else:
+                sd_model.unet = torch.compile(sd_model.unet, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
         if shared.opts.cuda_compile_vae:
-            if hasattr(sd_model, 'vae'):
+            if hasattr(sd_model, 'vae') and hasattr(sd_model.vae, 'decode'):
                 sd_model.vae.decode = torch.compile(sd_model.vae.decode, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
-            if hasattr(sd_model, 'movq'):
+            elif hasattr(sd_model, 'movq') and hasattr(sd_model.movq, 'decode'):
                 sd_model.movq.decode = torch.compile(sd_model.movq.decode, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
+            else:
+                shared.log.warning('Model compile enabled but model has no VAE')
         setup_logging() # compile messes with logging so reset is needed
         if shared.opts.cuda_compile_precompile:
             sd_model("dummy prompt")
@@ -127,14 +133,10 @@ def compile_torch(sd_model):
 def compile_diffusers(sd_model):
     if not (shared.opts.cuda_compile or shared.opts.cuda_compile_vae or shared.opts.cuda_compile_upscaler):
         return sd_model
-    if not hasattr(sd_model, 'unet') or not hasattr(sd_model.unet, 'config'):
-        shared.log.warning('Model compile enabled but model has no Unet')
-        return sd_model
     if shared.opts.cuda_compile_backend == 'none':
         shared.log.warning('Model compile enabled but no backend specified')
         return sd_model
-    size = 8*getattr(sd_model.unet.config, 'sample_size', 0)
-    shared.log.info(f"Model compile: pipeline={sd_model.__class__.__name__} shape={size} mode={shared.opts.cuda_compile_mode} backend={shared.opts.cuda_compile_backend} fullgraph={shared.opts.cuda_compile_fullgraph} unet={shared.opts.cuda_compile} vae={shared.opts.cuda_compile_vae} upscaler={shared.opts.cuda_compile_upscaler}")
+    shared.log.info(f"Model compile: pipeline={sd_model.__class__.__name__} mode={shared.opts.cuda_compile_mode} backend={shared.opts.cuda_compile_backend} fullgraph={shared.opts.cuda_compile_fullgraph} unet={shared.opts.cuda_compile} vae={shared.opts.cuda_compile_vae} upscaler={shared.opts.cuda_compile_upscaler}")
     if shared.opts.cuda_compile_backend == 'stable-fast':
         sd_model = compile_stablefast(sd_model)
     else: