From 9e006c3fa38ae538aab5489d224dec77b8e0d5e4 Mon Sep 17 00:00:00 2001
From: lodestone-rock <lodestone.rock@gmail.com>
Date: Tue, 9 Dec 2025 09:02:02 +0700
Subject: [PATCH] add chroma-radiance-x0 mode

---
 comfy/ldm/chroma_radiance/model.py | 21 +++++++++++++++++++--
 comfy/model_detection.py           |  3 +++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/comfy/ldm/chroma_radiance/model.py b/comfy/ldm/chroma_radiance/model.py
index e643b4414..e90b3abd2 100644
--- a/comfy/ldm/chroma_radiance/model.py
+++ b/comfy/ldm/chroma_radiance/model.py
@@ -37,7 +37,7 @@ class ChromaRadianceParams(ChromaParams):
     nerf_final_head_type: str
     # None means use the same dtype as the model.
     nerf_embedder_dtype: Optional[torch.dtype]
-
+    use_x0: bool
 
 class ChromaRadiance(Chroma):
     """
@@ -159,6 +159,10 @@ class ChromaRadiance(Chroma):
         self.skip_dit = []
         self.lite = False
 
+        if params.use_x0:
+            print("the model is using x0 prediction")
+            self.register_buffer("__x0__", torch.tensor([]))
+
     @property
     def _nerf_final_layer(self) -> nn.Module:
         if self.params.nerf_final_head_type == "linear":
@@ -276,6 +280,12 @@ class ChromaRadiance(Chroma):
         params_dict |= overrides
         return params.__class__(**params_dict)
 
+    def _apply_x0_residual(self, predicted, noisy, timesteps):
+
+        # non zero during training to prevent 0 div
+        eps = 0.0
+        return (noisy - predicted) / (timesteps.view(-1,1,1,1) + eps)
+    
     def _forward(
         self,
         x: Tensor,
@@ -316,4 +326,11 @@ class ChromaRadiance(Chroma):
             transformer_options,
             attn_mask=kwargs.get("attention_mask", None),
         )
-        return self.forward_nerf(img, img_out, params)[:, :, :h, :w]
+
+        out = self.forward_nerf(img, img_out, params)[:, :, :h, :w]
+
+        # If x0 variant → v-pred, just return this instead
+        if hasattr(self, "__x0__"):
+            out = self._apply_x0_residual(out, img, timestep)
+        return out
+
diff --git a/comfy/model_detection.py b/comfy/model_detection.py
index 74c547427..475b43d90 100644
--- a/comfy/model_detection.py
+++ b/comfy/model_detection.py
@@ -257,6 +257,9 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
                 dit_config["nerf_tile_size"] = 512
                 dit_config["nerf_final_head_type"] = "conv" if f"{key_prefix}nerf_final_layer_conv.norm.scale" in state_dict_keys else "linear"
                 dit_config["nerf_embedder_dtype"] = torch.float32
+            if f"__x0__" in state_dict_keys: # x0 pred
+                print("radiance in x0 mode")
+                dit_config["use_x0"] = True
         else:
             dit_config["guidance_embed"] = "{}guidance_in.in_layer.weight".format(key_prefix) in state_dict_keys
             dit_config["yak_mlp"] = '{}double_blocks.0.img_mlp.gate_proj.weight'.format(key_prefix) in state_dict_keys