dont pass null vae to from_pretrained

painebenjamin · Jan 19, 2024 · 47380df · 47380df
1 parent 7d82f32
commit 47380df
Showing 1 changed file with 26 additions and 0 deletions.
diff --git a/src/python/enfugue/diffusion/manager.py b/src/python/enfugue/diffusion/manager.py
@@ -3357,6 +3357,7 @@ def pipeline(self) -> EnfugueStableDiffusionPipeline:
 
                 if self.model_diffusers_cache_dir is None:
                     raise IOError("Couldn't create engine cache, check logs.")
+
                 if not self.is_sdxl:
                     kwargs["tokenizer_2"] = None
                     kwargs["text_encoder_2"] = None
@@ -3365,9 +3366,13 @@ def pipeline(self) -> EnfugueStableDiffusionPipeline:
                     kwargs["build_half"] = True
                     kwargs["variant"] = "fp16"
 
+                if kwargs["vae"] is None:
+                    kwargs.pop("vae")
+
                 logger.debug(
                     f"Initializing TensorRT pipeline from diffusers cache directory at {self.model_diffusers_cache_dir}. Arguments are {redact_for_log(kwargs)}"
                 )
+
                 pipeline = self.pipeline_class.from_pretrained(
                     self.model_diffusers_cache_dir,
                     local_files_only=self.offline,
@@ -3383,6 +3388,9 @@ def pipeline(self) -> EnfugueStableDiffusionPipeline:
                 if "16" in str(self.dtype):
                     kwargs["variant"] = "fp16"
 
+                if kwargs["vae"] is None:
+                    kwargs.pop("vae")
+
                 logger.debug(
                     f"Initializing pipeline from diffusers cache directory at {self.model_diffusers_cache_dir}. Arguments are {redact_for_log(kwargs)}"
                 )
@@ -3521,6 +3529,9 @@ def refiner_pipeline(self) -> EnfugueStableDiffusionPipeline:
                     kwargs["build_half"] = True
                     kwargs["variant"] = "fp16"
 
+                if kwargs["vae"] is None:
+                    kwargs.pop("vae")
+
                 logger.debug(
                     f"Initializing refiner TensorRT pipeline from diffusers cache directory at {self.refiner_diffusers_cache_dir}. Arguments are {redact_for_log(kwargs)}"
                 )
@@ -3544,6 +3555,9 @@ def refiner_pipeline(self) -> EnfugueStableDiffusionPipeline:
                 if "16" in str(self.dtype):
                     kwargs["variant"] = "fp16"
 
+                if kwargs["vae"] is None:
+                    kwargs.pop("vae")
+
                 logger.debug(
                     f"Initializing refiner pipeline from diffusers cache directory at {self.refiner_diffusers_cache_dir}. Arguments are {redact_for_log(kwargs)}"
                 )
@@ -3700,6 +3714,9 @@ def inpainter_pipeline(self) -> EnfugueStableDiffusionPipeline:
                     kwargs["variant"] = "fp16"
                     kwargs["build_half"] = True
 
+                if kwargs["vae"] is None:
+                    kwargs.pop("vae")
+
                 logger.debug(
                     f"Initializing inpainter TensorRT pipeline from diffusers cache directory at {self.inpainter_diffusers_cache_dir}. Arguments are {redact_for_log(kwargs)}"
                 )
@@ -3719,6 +3736,9 @@ def inpainter_pipeline(self) -> EnfugueStableDiffusionPipeline:
                     kwargs["tokenizer_2"] = None
                 if "16" in str(self.dtype):
                     kwargs["variant"] = "fp16"
+
+                if kwargs["vae"] is None:
+                    kwargs.pop("vae")
 
                 logger.debug(
                     f"Initializing inpainter pipeline from diffusers cache directory at {self.inpainter_diffusers_cache_dir}. Arguments are {redact_for_log(kwargs)}"
@@ -3875,6 +3895,9 @@ def animator_pipeline(self) -> EnfugueAnimateStableDiffusionPipeline:
                     kwargs["variant"] = "fp16"
                     kwargs["build_half"] = True
 
+                if kwargs["vae"] is None:
+                    kwargs.pop("vae")
+
                 logger.debug(
                     f"Initializing animator TensorRT pipeline from diffusers cache directory at {self.animator_diffusers_cache_dir}. Arguments are {redact_for_log(kwargs)}"
                 )
@@ -3893,6 +3916,9 @@ def animator_pipeline(self) -> EnfugueAnimateStableDiffusionPipeline:
                     kwargs["tokenizer_2"] = None
                 if "16" in str(self.dtype):
                     kwargs["variant"] = "fp16"
+
+                if kwargs["vae"] is None:
+                    kwargs.pop("vae")
 
                 logger.debug(
                     f"Initializing animator pipeline from diffusers cache directory at {self.animator_diffusers_cache_dir}. Arguments are {redact_for_log(kwargs)}"