jokerbit commited on
Commit
0a1e10b
·
verified ·
1 Parent(s): 34f1ca3

Compile text_encoder

Browse files
Files changed (1) hide show
  1. src/pipeline.py +2 -1
src/pipeline.py CHANGED
@@ -49,7 +49,8 @@ def load_pipeline() -> Pipeline:
49
  torch_dtype=torch.bfloat16,
50
  )
51
 
52
- pipeline.transformer.to(memory_format=torch.channels_last)
 
53
  # quantize_(pipeline.vae, int8_weight_only())
54
  # pipeline.vae = torch.compile(pipeline.vae, mode="reduce-overhead")
55
  pipeline.to("cuda")
 
49
  torch_dtype=torch.bfloat16,
50
  )
51
 
52
+ pipeline.transformer.to(memory_format=torch.channels_last)
53
+ pipeline.text_encoder = torch.compile(pipeline.text_encoder, mode="reduce-overhead")
54
  # quantize_(pipeline.vae, int8_weight_only())
55
  # pipeline.vae = torch.compile(pipeline.vae, mode="reduce-overhead")
56
  pipeline.to("cuda")