From ce1c679f2222657096d8161e04e1bffa83119958 Mon Sep 17 00:00:00 2001 From: blepping Date: Sat, 23 Aug 2025 23:34:06 -0600 Subject: [PATCH] Allow setting NeRF embedder dtype for Radiance Bump Radiance nerf tile size to 32 Support EasyCache/LazyCache on Radiance (maybe) --- comfy/ldm/chroma/layers_dct.py | 30 +++++++++++++++++++----------- comfy/ldm/chroma/model_dct.py | 24 ++++++++++++++++++++---- comfy/model_detection.py | 3 ++- 3 files changed, 41 insertions(+), 16 deletions(-) diff --git a/comfy/ldm/chroma/layers_dct.py b/comfy/ldm/chroma/layers_dct.py index 8a2824e02..f3130a78f 100644 --- a/comfy/ldm/chroma/layers_dct.py +++ b/comfy/ldm/chroma/layers_dct.py @@ -15,7 +15,7 @@ class NerfEmbedder(nn.Module): patch size, and enriches it with positional information before projecting it to a new hidden size. """ - def __init__(self, in_channels, hidden_size_input, max_freqs, dtype=None, device=None, operations=None): + def __init__(self, in_channels, hidden_size_input, max_freqs, dtype=None, device=None, operations=None, *, embedder_dtype=None): """ Initializes the NerfEmbedder. @@ -29,6 +29,7 @@ class NerfEmbedder(nn.Module): super().__init__() self.max_freqs = max_freqs self.hidden_size_input = hidden_size_input + self.embedder_dtype = embedder_dtype # A linear layer to project the concatenated input features and # positional encodings to the final output dimension. @@ -37,7 +38,7 @@ class NerfEmbedder(nn.Module): ) @lru_cache(maxsize=4) - def fetch_pos(self, patch_size, device, dtype): + def fetch_pos(self, patch_size: int, device, dtype) -> torch.Tensor: """ Generates and caches 2D DCT-like positional embeddings for a given patch size. @@ -91,7 +92,7 @@ class NerfEmbedder(nn.Module): return dct - def forward(self, inputs): + def forward(self, inputs: torch.Tensor) -> torch.Tensor: """ Forward pass for the embedder. @@ -107,26 +108,34 @@ class NerfEmbedder(nn.Module): # Infer the patch side length from the number of pixels (P^2). patch_size = int(P2 ** 0.5) + # Possibly run the operation with a different dtype. + orig_dtype = inputs.dtype + if self.embedder_dtype is not None and self.embedder_dtype != orig_dtype: + embedder = self.embedder.to(dtype=self.embedder_dtype) + else: + embedder = self.embedder + # Fetch the pre-computed or cached positional embeddings. - dct = self.fetch_pos(patch_size, inputs.device, inputs.dtype) + dct = self.fetch_pos(patch_size, inputs.device, self.embedder_dtype or inputs.dtype) # Repeat the positional embeddings for each item in the batch. dct = dct.repeat(B, 1, 1) # Concatenate the original input features with the positional embeddings # along the feature dimension. - inputs = torch.cat([inputs, dct], dim=-1) + inputs = torch.cat((inputs, dct), dim=-1) # Project the combined tensor to the target hidden size. - inputs = self.embedder(inputs) + inputs = embedder(inputs) - return inputs + # No-op if already the same dtype. + return inputs.to(dtype=orig_dtype) class NerfGLUBlock(nn.Module): """ A NerfBlock using a Gated Linear Unit (GLU) like MLP. """ - def __init__(self, hidden_size_s, hidden_size_x, mlp_ratio, dtype=None, device=None, operations=None): + def __init__(self, hidden_size_s: int, hidden_size_x: int, mlp_ratio, dtype=None, device=None, operations=None): super().__init__() # The total number of parameters for the MLP is increased to accommodate # the gate, value, and output projection matrices. @@ -137,7 +146,7 @@ class NerfGLUBlock(nn.Module): self.mlp_ratio = mlp_ratio - def forward(self, x, s): + def forward(self, x: torch.Tensor, s: torch.Tensor) -> torch.Tensor: batch_size, num_x, hidden_size_x = x.shape mlp_params = self.param_generator(s) @@ -160,8 +169,7 @@ class NerfGLUBlock(nn.Module): # Apply the final output projection. x = torch.bmm(torch.nn.functional.silu(torch.bmm(x, fc1_gate)) * torch.bmm(x, fc1_value), fc2) - x = x + res_x - return x + return x + res_x class NerfFinalLayer(nn.Module): diff --git a/comfy/ldm/chroma/model_dct.py b/comfy/ldm/chroma/model_dct.py index 57a1939f7..dfcd6d1de 100644 --- a/comfy/ldm/chroma/model_dct.py +++ b/comfy/ldm/chroma/model_dct.py @@ -3,6 +3,7 @@ # Chroma Radiance adaption referenced from https://github.com/lodestone-rock/flow from dataclasses import dataclass +from typing import Optional import torch from torch import Tensor, nn @@ -37,6 +38,7 @@ class ChromaRadianceParams(chroma_model.ChromaParams): nerf_max_freqs: int nerf_tile_size: int nerf_final_head_type: str + nerf_embedder_dtype: Optional[torch.dtype] class ChromaRadiance(chroma_model.Chroma): @@ -101,7 +103,12 @@ class ChromaRadiance(chroma_model.Chroma): self.single_blocks = nn.ModuleList( [ - SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio, dtype=dtype, device=device, operations=operations) + SingleStreamBlock( + self.hidden_size, + self.num_heads, + mlp_ratio=params.mlp_ratio, + dtype=dtype, device=device, operations=operations, + ) for _ in range(params.depth_single_blocks) ] ) @@ -114,6 +121,7 @@ class ChromaRadiance(chroma_model.Chroma): dtype=dtype, device=device, operations=operations, + embedder_dtype=params.nerf_embedder_dtype, ) self.nerf_blocks = nn.ModuleList([ @@ -135,7 +143,6 @@ class ChromaRadiance(chroma_model.Chroma): device=device, operations=operations, ) - self._nerf_final_layer = self.nerf_final_layer elif params.nerf_final_head_type == "conv": self.nerf_final_layer_conv = NerfFinalLayerConv( params.nerf_hidden_size, @@ -144,14 +151,23 @@ class ChromaRadiance(chroma_model.Chroma): device=device, operations=operations, ) - self._nerf_final_layer = self.nerf_final_layer_conv else: errstr = f"Unsupported nerf_final_head_type {params.nerf_final_head_type}" + raise ValueError(errstr) self.skip_mmdit = [] self.skip_dit = [] self.lite = False + @property + def _nerf_final_layer(self) -> nn.Module: + if self.params.nerf_final_head_type == "linear": + return self.nerf_final_layer + if self.params.nerf_final_head_type == "conv": + return self.nerf_final_layer_conv + # Impossible to get here as we raise an error on unexpected types on initialization. + raise NotImplementedError + def forward_tiled_nerf( self, nerf_hidden: Tensor, @@ -337,7 +353,7 @@ class ChromaRadiance(chroma_model.Chroma): ) return self._nerf_final_layer(img_dct) - def forward(self, x, timestep, context, guidance, control=None, transformer_options={}, **kwargs): + def _forward(self, x, timestep, context, guidance, control=None, transformer_options={}, **kwargs): bs, c, h, w = x.shape img = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size)) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index 14b2033f5..bfc6a188a 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -213,8 +213,9 @@ def detect_unet_config(state_dict, key_prefix, metadata=None): dit_config["nerf_mlp_ratio"] = 4 dit_config["nerf_depth"] = 4 dit_config["nerf_max_freqs"] = 8 - dit_config["nerf_tile_size"] = 16 + dit_config["nerf_tile_size"] = 32 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 else: dit_config["guidance_embed"] = "{}guidance_in.in_layer.weight".format(key_prefix) in state_dict_keys return dit_config