zhiweili
commited on
Commit
Β·
dfb0452
1
Parent(s):
90660b6
import torch_tensorrt
Browse files- app_tensorrt.py +1 -5
app_tensorrt.py
CHANGED
@@ -9,9 +9,6 @@ from diffusers import (
|
|
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 |
|
@@ -35,8 +32,6 @@ base_pipe.scheduler = DDPMScheduler.from_pretrained(
|
|
35 |
|
36 |
backend = "torch_tensorrt"
|
37 |
|
38 |
-
|
39 |
-
|
40 |
# print('Loading compiled model...')
|
41 |
# loadedModel = torch_tensorrt.load("compiled_pipe.ep").module()
|
42 |
# print('Compiled model loaded!')
|
@@ -50,6 +45,7 @@ def create_demo() -> gr.Blocks:
|
|
50 |
prompt:str,
|
51 |
steps:int,
|
52 |
):
|
|
|
53 |
print('Compiling model...')
|
54 |
compiledModel = torch.compile(
|
55 |
base_pipe.unet,
|
|
|
9 |
|
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 |
|
|
|
32 |
|
33 |
backend = "torch_tensorrt"
|
34 |
|
|
|
|
|
35 |
# print('Loading compiled model...')
|
36 |
# loadedModel = torch_tensorrt.load("compiled_pipe.ep").module()
|
37 |
# print('Compiled model loaded!')
|
|
|
45 |
prompt:str,
|
46 |
steps:int,
|
47 |
):
|
48 |
+
import torch_tensorrt
|
49 |
print('Compiling model...')
|
50 |
compiledModel = torch.compile(
|
51 |
base_pipe.unet,
|