1 files changed, 13 insertions, 8 deletions

diff --git a/infer.py b/infer.py
index 9bc9efe..9b0ec1f 100644
--- a/infer.py
+++ b/infer.py
@@ -8,11 +8,10 @@ from pathlib import Path
 import torch
 import json
 from PIL import Image
-from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler
+from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler, DPMSolverMultistepScheduler, DDIMScheduler, LMSDiscreteScheduler, EulerAncestralDiscreteScheduler
 from transformers import CLIPTextModel, CLIPTokenizer
 from slugify import slugify
 
-from schedulers.scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler
 from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 
 
@@ -21,7 +20,7 @@ torch.backends.cuda.matmul.allow_tf32 = True
 
 default_args = {
     "model": None,
-    "scheduler": "euler_a",
+    "scheduler": "dpmpp",
     "precision": "fp32",
     "ti_embeddings_dir": "embeddings_ti",
     "output_dir": "output/inference",
@@ -65,7 +64,7 @@ def create_args_parser():
     parser.add_argument(
         "--scheduler",
         type=str,
-        choices=["plms", "ddim", "klms", "euler_a"],
+        choices=["plms", "ddim", "klms", "dpmpp", "euler_a"],
     )
     parser.add_argument(
         "--precision",
@@ -222,6 +221,10 @@ def create_pipeline(model, scheduler, ti_embeddings_dir, dtype):
         scheduler = DDIMScheduler(
             beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False
         )
+    elif scheduler == "dpmpp":
+        scheduler = DPMSolverMultistepScheduler(
+            beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
+        )
     else:
         scheduler = EulerAncestralDiscreteScheduler(
             beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
@@ -282,7 +285,8 @@ def generate(output_dir, pipeline, args):
             ).images
 
             for j, image in enumerate(images):
-                image.save(output_dir.joinpath(f"{args.seed + i}_{j}.jpg"))
+                image.save(output_dir.joinpath(f"{args.seed + i}_{j}.png"))
+                image.save(output_dir.joinpath(f"{args.seed + i}_{j}.jpg"), quality=85)
 
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
@@ -312,15 +316,16 @@ class CmdParse(cmd.Cmd):
 
         try:
             args = run_parser(self.parser, default_cmds, elements)
+
+            if len(args.prompt) == 0:
+                print('Try again with a prompt!')
+                return
         except SystemExit:
             self.parser.print_help()
         except Exception as e:
             print(e)
             return
 
-        if len(args.prompt) == 0:
-            print('Try again with a prompt!')
-
         try:
             generate(self.output_dir, self.pipeline, args)
         except KeyboardInterrupt: