zhiweili commited on
Commit
9b4449a
Β·
1 Parent(s): 2e18836

fix torch_tensorrt

Browse files
Files changed (2) hide show
  1. app_tensorrt.py +6 -1
  2. requirements.txt +1 -2
app_tensorrt.py CHANGED
@@ -1,5 +1,5 @@
1
  import torch
2
- import torch_tensorrt
3
 
4
  from diffusers import (
5
  DDPMScheduler,
@@ -7,6 +7,11 @@ from diffusers import (
7
  AutoencoderKL,
8
  )
9
 
 
 
 
 
 
10
  BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
11
  device = "cuda"
12
 
 
1
  import torch
2
+ import os
3
 
4
  from diffusers import (
5
  DDPMScheduler,
 
7
  AutoencoderKL,
8
  )
9
 
10
+ os.system("pip install torch_tensorrt==2.4.0")
11
+
12
+
13
+ import torch_tensorrt
14
+
15
  BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
16
  device = "cuda"
17
 
requirements.txt CHANGED
@@ -1,6 +1,5 @@
1
  gradio
2
- torch==2.4.0
3
- torch_tensorrt==2.4.0
4
  torchvision
5
  diffusers
6
  transformers
 
1
  gradio
2
+ torch
 
3
  torchvision
4
  diffusers
5
  transformers