From a5afb5950c68ebf572b57ca6efce388d2f69b824 Mon Sep 17 00:00:00 2001
From: thot-experiment <thot@thiic.cc>
Date: Mon, 10 Mar 2025 17:41:05 -0700
Subject: [PATCH] fix model_sampling kludges

---
 comfy/model_base.py                  | 14 +-------------
 comfy/supported_models.py            |  4 ++--
 comfy_extras/nodes_model_advanced.py |  8 +++++++-
 3 files changed, 10 insertions(+), 16 deletions(-)

diff --git a/comfy/model_base.py b/comfy/model_base.py
index a15a9f6e5..2e3a157b0 100644
--- a/comfy/model_base.py
+++ b/comfy/model_base.py
@@ -140,10 +140,6 @@ class BaseModel(torch.nn.Module):
         sigma = t
         xc = self.model_sampling.calculate_input(sigma, x)
         
-        #this is probably not the right way to do this
-        if isinstance(self,LotusModel): 
-            xc = x#torch.cat([x], dim=1)
-
         if c_concat is not None:
             xc = torch.cat([xc] + [c_concat], dim=1)
 
@@ -605,7 +601,7 @@ class SDXL_instructpix2pix(IP2P, SDXL):
         else:
             self.process_ip2p_image_in = lambda image: image #diffusers ip2p
 
-class LotusModel(BaseModel):
+class Lotus(BaseModel):
     def extra_conds(self, **kwargs):
         out = {}
         cross_attn = kwargs.get("cross_attn", None)
@@ -616,14 +612,6 @@ class LotusModel(BaseModel):
         out['y'] = comfy.conds.CONDRegular(task_emb)
         return out
 
-    def process_latent_out(self, latent):
-        #TODO FIX
-        #there is some scaling issue that happens during diffusion and this is a kludge to fix it
-        #my best guess is that the scaling is somehow fucked up by the 999 sigma which is why
-        #i'm using it here to scale things back, however this is definitely not correct and is just
-        #patching an error that comes from the actual model sampling being done incorrectly
-        return latent/-(999*0.18215)
-
     def __init__(self, model_config, model_type=ModelType.EPS, device=None):
         super().__init__(model_config, model_type, device=device)
 
diff --git a/comfy/supported_models.py b/comfy/supported_models.py
index 850da68d6..5642310c4 100644
--- a/comfy/supported_models.py
+++ b/comfy/supported_models.py
@@ -511,6 +511,7 @@ class LotusD(SD20):
         "model_channels": 320,
         "use_linear_in_transformer": True,
         "use_temporal_attention": False,
+        "adm_in_channels": 4,
         "in_channels": 4,
     }
 
@@ -519,8 +520,7 @@ class LotusD(SD20):
     }
 
     def get_model(self, state_dict, prefix="", device=None):
-        print('identified lotus-d model')
-        return model_base.LotusModel(self, device=device)
+        return model_base.Lotus(self, device=device)
 
 class SD3(supported_models_base.BASE):
     unet_config = {
diff --git a/comfy_extras/nodes_model_advanced.py b/comfy_extras/nodes_model_advanced.py
index ceac5654b..2b805c1ee 100644
--- a/comfy_extras/nodes_model_advanced.py
+++ b/comfy_extras/nodes_model_advanced.py
@@ -24,6 +24,10 @@ class X0(comfy.model_sampling.EPS):
     def calculate_denoised(self, sigma, model_output, model_input):
         return model_output
 
+class Lotus(X0):
+    def calculate_input(self, sigma, noise):
+        return noise
+
 class ModelSamplingDiscreteDistilled(comfy.model_sampling.ModelSamplingDiscrete):
     original_timesteps = 50
 
@@ -56,7 +60,7 @@ class ModelSamplingDiscrete:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": { "model": ("MODEL",),
-                              "sampling": (["eps", "v_prediction", "lcm", "x0"],),
+                              "sampling": (["eps", "v_prediction", "lcm", "x0", "lotus"],),
                               "zsnr": ("BOOLEAN", {"default": False}),
                               }}
 
@@ -78,6 +82,8 @@ class ModelSamplingDiscrete:
             sampling_base = ModelSamplingDiscreteDistilled
         elif sampling == "x0":
             sampling_type = X0
+        elif sampling == "lotus":
+            sampling_type = Lotus
 
         class ModelSamplingAdvanced(sampling_base, sampling_type):
             pass