Update src/pipeline.py
Browse files- src/pipeline.py +2 -0
src/pipeline.py
CHANGED
@@ -25,6 +25,8 @@ def load_pipeline() -> Pipeline:
|
|
25 |
transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=False)
|
26 |
pipeline = FluxPipeline.from_pretrained(ids, revision=Revision, transformer=transformer, local_files_only=True, torch_dtype=torch.bfloat16,)
|
27 |
pipeline.to("cuda")
|
|
|
|
|
28 |
for _ in range(3):
|
29 |
pipeline(prompt="insensible, timbale, pothery, electrovital, actinogram, taxis, intracerebellar, centrodesmus", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
|
30 |
return pipeline
|
|
|
25 |
transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=False)
|
26 |
pipeline = FluxPipeline.from_pretrained(ids, revision=Revision, transformer=transformer, local_files_only=True, torch_dtype=torch.bfloat16,)
|
27 |
pipeline.to("cuda")
|
28 |
+
pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
|
29 |
+
quantize_(pipeline.vae, int8_weight_only())
|
30 |
for _ in range(3):
|
31 |
pipeline(prompt="insensible, timbale, pothery, electrovital, actinogram, taxis, intracerebellar, centrodesmus", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
|
32 |
return pipeline
|