From 499b6e04b486b490236471b9491fe3223045f5fc Mon Sep 17 00:00:00 2001 From: chflame163 Date: Wed, 6 Nov 2024 14:41:40 +0800 Subject: [PATCH] fix bug of separate_cfg_infer set to False --- py/OmniGen/pipeline.py | 1 + pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/py/OmniGen/pipeline.py b/py/OmniGen/pipeline.py index d3b0c9a..b2a4957 100644 --- a/py/OmniGen/pipeline.py +++ b/py/OmniGen/pipeline.py @@ -110,6 +110,7 @@ class OmniGenPipeline: self.device = device def vae_encode(self, x, dtype): + x = x.to(self.vae.config.get("dtype", torch.bfloat16)) if self.vae.config.shift_factor is not None: x = self.vae.encode(x).latent_dist.sample() x = (x - self.vae.config.shift_factor) * self.vae.config.scaling_factor diff --git a/pyproject.toml b/pyproject.toml index 6c2ff3e..09cab58 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "ComfyUI_OmniGen_Wrapper" description = "ComfyUI custom node of OmniGen project." -version = "1.0.1" +version = "1.0.2" license = "MIT" dependencies = ["transformers", "diffusers", "numpy", "datasets", accelerate", "timm", "peft", "safetensors", "huggingface_hub"]