add warning if using non jit unet

This commit is contained in:
terrance committed 2023-02-03 22:36:04 +00:00
1 parent fd5b8be160
commit 3dffd956fe
1 file changed
+3
@@ -109,6 +109,9 @@ class StableDiffusion(nn.Module):
self.unet = TracedUNet()
del pipe.unet
else:
CONSOLE.print(
"Warning: Loading UNet without JIT acceleration. Run trace_stablediff.py in the scripts folder for a speedup!"
)
self.unet = pipe.unet
self.unet.to(memory_format=torch.channels_last)