From 6aadb34af4fe5ca2dfc92fae8eee87610a5848ad Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Sat, 8 Oct 2022 21:56:54 +0200
Subject: Update

---
 infer.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

(limited to 'infer.py')

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!')
-- 
cgit v1.2.3-54-g00ecf