radames commited on
Commit
72bd0bb
1 Parent(s): 15889a4

try torch compile

Browse files
Files changed (1) hide show
  1. app.py +7 -0
app.py CHANGED
@@ -47,6 +47,13 @@ if torch.cuda.is_available():
47
  pipe.vae = AutoencoderTiny.from_pretrained(
48
  "madebyollin/taesd3", torch_dtype=torch.float16
49
  )
 
 
 
 
 
 
 
50
  pipe.vae.config.shift_factor = 0.0
51
 
52
  pipe = pipe.to(device)
 
47
  pipe.vae = AutoencoderTiny.from_pretrained(
48
  "madebyollin/taesd3", torch_dtype=torch.float16
49
  )
50
+
51
+ pipe.vae.decoder.layers = torch.compile(
52
+ pipe.vae.decoder.layers,
53
+ fullgraph=True,
54
+ dynamic=False,
55
+ mode="max-autotune-no-cudagraphs",
56
+ )
57
  pipe.vae.config.shift_factor = 0.0
58
 
59
  pipe = pipe.to(device)