Update

1 files changed, 7 insertions, 5 deletions

diff --git a/infer.py b/infer.py
index 6197aa3..a542534 100644
--- a/infer.py
+++ b/infer.py
@@ -5,12 +5,11 @@ import sys
 import shlex
 import cmd
 from pathlib import Path
-from torch import autocast
 import torch
 import json
 from PIL import Image
-from diffusers import StableDiffusionPipeline, AutoencoderKL, UNet2DConditionModel, PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler
-from transformers import CLIPTextModel, CLIPTokenizer, CLIPFeatureExtractor
+from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler
+from transformers import CLIPTextModel, CLIPTokenizer
 from slugify import slugify
 from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 from schedulers.scheduling_euler_a import EulerAScheduler
@@ -22,7 +21,7 @@ torch.backends.cuda.matmul.allow_tf32 = True
 default_args = {
     "model": None,
     "scheduler": "euler_a",
-    "precision": "bf16",
+    "precision": "fp16",
     "embeddings_dir": "embeddings",
     "output_dir": "output/inference",
     "config": None,
@@ -260,7 +259,7 @@ def generate(output_dir, pipeline, args):
     else:
         init_image = None
 
-    with autocast("cuda"):
+    with torch.autocast("cuda"), torch.inference_mode():
         for i in range(args.batch_num):
             pipeline.set_progress_bar_config(
                 desc=f"Batch {i + 1} of {args.batch_num}",
@@ -313,6 +312,9 @@ class CmdParse(cmd.Cmd):
             args = run_parser(self.parser, default_cmds, elements)
         except SystemExit:
             self.parser.print_help()
+        except Exception as e:
+            print(e)
+            return
 
         if len(args.prompt) == 0:
             print('Try again with a prompt!')