fix bug of separate_cfg_infer set to False

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

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"]