Small perf improvements

1 files changed, 3 insertions, 2 deletions

diff --git a/infer.py b/infer.py
index b15b17f..3dc0f32 100644
--- a/infer.py
+++ b/infer.py
@@ -16,6 +16,9 @@ from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 from schedulers.scheduling_euler_a import EulerAScheduler
 
 
+torch.backends.cuda.matmul.allow_tf32 = True
+
+
 default_args = {
     "model": None,
     "scheduler": "euler_a",
@@ -166,7 +169,6 @@ def create_pipeline(model, scheduler, dtype):
     text_encoder = CLIPTextModel.from_pretrained(model + '/text_encoder', torch_dtype=dtype)
     vae = AutoencoderKL.from_pretrained(model + '/vae', torch_dtype=dtype)
     unet = UNet2DConditionModel.from_pretrained(model + '/unet', torch_dtype=dtype)
-    feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32", torch_dtype=dtype)
 
     if scheduler == "plms":
         scheduler = PNDMScheduler(
@@ -191,7 +193,6 @@ def create_pipeline(model, scheduler, dtype):
         unet=unet,
         tokenizer=tokenizer,
         scheduler=scheduler,
-        feature_extractor=feature_extractor
     )
     # pipeline.enable_attention_slicing()
     pipeline.to("cuda")