zhiweili commited on
Commit
dfb0452
Β·
1 Parent(s): 90660b6

import torch_tensorrt

Browse files
Files changed (1) hide show
  1. 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,