jokerbit commited on
Commit
34d338c
·
verified ·
1 Parent(s): c31a6b2

Upload src/pipeline.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. src/pipeline.py +1 -2
src/pipeline.py CHANGED
@@ -10,7 +10,6 @@ from pipelines.models import TextToImageRequest
10
  from torch import Generator
11
  from torchao.quantization import quantize_, int8_weight_only
12
  from transformers import T5EncoderModel, CLIPTextModel
13
- import torch_tensorrt
14
 
15
 
16
  Pipeline: TypeAlias = FluxPipeline
@@ -52,7 +51,7 @@ def load_pipeline() -> Pipeline:
52
 
53
  pipeline.transformer.to(memory_format=torch.channels_last)
54
  # quantize_(pipeline.vae, int8_weight_only())
55
- pipeline.vae = torch.compile(pipeline.vae, backend="tensorrt")
56
  pipeline.to("cuda")
57
 
58
  for _ in range(2):
 
10
  from torch import Generator
11
  from torchao.quantization import quantize_, int8_weight_only
12
  from transformers import T5EncoderModel, CLIPTextModel
 
13
 
14
 
15
  Pipeline: TypeAlias = FluxPipeline
 
51
 
52
  pipeline.transformer.to(memory_format=torch.channels_last)
53
  # quantize_(pipeline.vae, int8_weight_only())
54
+ pipeline.vae = torch.compile(pipeline.vae)
55
  pipeline.to("cuda")
56
 
57
  for _ in range(2):