manbeast3b commited on
Commit
6a0ebda
·
verified ·
1 Parent(s): e18b8f9

Update src/pipeline.py

Browse files
Files changed (1) hide show
  1. src/pipeline.py +1 -1
src/pipeline.py CHANGED
@@ -86,7 +86,7 @@ def load_pipeline() -> Pipeline:
86
  pipeline.vae.to(memory_format=torch.channels_last)
87
  pipeline.vae = torch.compile(pipeline.vae)
88
 
89
- pipeline._exclude_from_cpu_offload = ["vae"]
90
  pipeline.enable_sequential_cpu_offload()
91
  for _ in range(2):
92
  pipeline(prompt="onomancy, aftergo, spirantic, Platyhelmia, modificator, drupaceous, jobbernowl, hereness", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
 
86
  pipeline.vae.to(memory_format=torch.channels_last)
87
  pipeline.vae = torch.compile(pipeline.vae)
88
 
89
+ pipeline._exclude_from_cpu_offload = ["vae", "transformer"]
90
  pipeline.enable_sequential_cpu_offload()
91
  for _ in range(2):
92
  pipeline(prompt="onomancy, aftergo, spirantic, Platyhelmia, modificator, drupaceous, jobbernowl, hereness", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)