cleanup compile

This commit is contained in:
Vladimir Mandic
2023-12-05 11:49:59 -05:00
parent c8970ab361
commit 6cec284a6f
+10 -8
View File
@@ -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: