jokerbit commited on
Commit
18611d2
·
verified ·
1 Parent(s): d6fa8be
Files changed (1) hide show
  1. src/pipeline.py +1 -3
src/pipeline.py CHANGED
@@ -10,8 +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, logging
13
- import torch._dynamo
14
- torch._dynamo.config.suppress_errors = True
15
 
16
  Pipeline: TypeAlias = FluxPipeline
17
 
@@ -48,7 +46,7 @@ def load_pipeline() -> Pipeline:
48
  ).to("cuda")
49
 
50
  pipeline.to(memory_format=torch.channels_last)
51
- pipeline.transformer = torch.compile(pipeline.transformer)
52
  quantize_(pipeline.vae, int8_weight_only())
53
 
54
  PROMPT = 'semiconformity, peregrination, quip, twineless, emotionless, tawa, depickle'
 
10
  from torch import Generator
11
  from torchao.quantization import quantize_, int8_weight_only
12
  from transformers import T5EncoderModel, CLIPTextModel, logging
 
 
13
 
14
  Pipeline: TypeAlias = FluxPipeline
15
 
 
46
  ).to("cuda")
47
 
48
  pipeline.to(memory_format=torch.channels_last)
49
+ pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
50
  quantize_(pipeline.vae, int8_weight_only())
51
 
52
  PROMPT = 'semiconformity, peregrination, quip, twineless, emotionless, tawa, depickle'