Spaces:
Sleeping
Sleeping
Update run.py
Browse files
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 |
|