zhiweili
commited on
Commit
·
9966900
1
Parent(s):
dfb0452
fix device
Browse files- app_tensorrt.py +2 -1
app_tensorrt.py
CHANGED
|
@@ -10,7 +10,8 @@ from diffusers import (
|
|
| 10 |
os.system("pip install torch_tensorrt==2.4.0")
|
| 11 |
|
| 12 |
BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
|
| 13 |
-
device = "cuda"
|
|
|
|
| 14 |
|
| 15 |
vae = AutoencoderKL.from_pretrained(
|
| 16 |
"madebyollin/sdxl-vae-fp16-fix",
|
|
|
|
| 10 |
os.system("pip install torch_tensorrt==2.4.0")
|
| 11 |
|
| 12 |
BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
|
| 13 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
| 14 |
+
print(f"--------->Device: {device}")
|
| 15 |
|
| 16 |
vae = AutoencoderKL.from_pretrained(
|
| 17 |
"madebyollin/sdxl-vae-fp16-fix",
|