Royir commited on
Commit
6bbdd23
1 Parent(s): db3cb87

Update run.py

Browse files
Files changed (1) hide show
  1. run.py +1 -0
run.py CHANGED
@@ -14,6 +14,7 @@ def main(prompt, seed, output_directory, model_path):
14
  def load_model(model_path):
15
  device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
16
  pipe = SynGenDiffusionPipeline.from_pretrained(model_path).to(device)
 
17
 
18
  return pipe
19
 
 
14
  def load_model(model_path):
15
  device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
16
  pipe = SynGenDiffusionPipeline.from_pretrained(model_path).to(device)
17
+ pipe.disable_xformers_memory_efficient_attention()
18
 
19
  return pipe
20