From 61b08d4ba65fec37070376bf50da3ec3c534e859 Mon Sep 17 00:00:00 2001
From: chaObserv <154517000+chaObserv@users.noreply.github.com>
Date: Thu, 31 Jul 2025 07:25:56 +0800
Subject: [PATCH] Replace manual x * sigmoid(x) with torch silu in VAE
 nonlinearity (#9057)

---
 comfy/ldm/cosmos/cosmos_tokenizer/utils.py  | 3 ++-
 comfy/ldm/modules/diffusionmodules/model.py | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/comfy/ldm/cosmos/cosmos_tokenizer/utils.py b/comfy/ldm/cosmos/cosmos_tokenizer/utils.py
index 3af8d0d05..ca993006f 100644
--- a/comfy/ldm/cosmos/cosmos_tokenizer/utils.py
+++ b/comfy/ldm/cosmos/cosmos_tokenizer/utils.py
@@ -58,7 +58,8 @@ def is_odd(n: int) -> bool:
 
 
 def nonlinearity(x):
-    return x * torch.sigmoid(x)
+    # x * sigmoid(x)
+    return torch.nn.functional.silu(x)
 
 
 def Normalize(in_channels, num_groups=32):
diff --git a/comfy/ldm/modules/diffusionmodules/model.py b/comfy/ldm/modules/diffusionmodules/model.py
index 8162742cf..5c0373b74 100644
--- a/comfy/ldm/modules/diffusionmodules/model.py
+++ b/comfy/ldm/modules/diffusionmodules/model.py
@@ -36,7 +36,7 @@ def get_timestep_embedding(timesteps, embedding_dim):
 
 def nonlinearity(x):
     # swish
-    return x*torch.sigmoid(x)
+    return torch.nn.functional.silu(x)
 
 
 def Normalize(in_channels, num_groups=32):