Add support for Chroma Radiance #9682

blepping · 2025-09-02T14:13:24Z

This patch adds support for Chroma Radiance, a pixel-space model adapted from Chroma, which is showing very promising results (although it is relatively early in training). The performance is quite close to Chroma.

I've tried to reuse existing code as much as possible so this is fairly self-contained. A pretty good chunk of it is copied from the implementation at https://github.com/lodestone-rock/flow - I can't take any credit for those parts.

The model can be downloaded at: https://huggingface.co/lodestones/chroma-debug-development-only/tree/main/radiance

No difference from a normal Chroma workflow except you need to use the ChromaRadianceStubVAE node instead of loading a VAE. It's not a real VAE, it just moves the channel dimension and scales the input image values to -1..1 for encode and the reverse for decode. There are also ChromaRadianceLatentToImage and ChromaRadianceImageToLatent nodes that do the same thing which could be removed. Lodestone Rock (Chroma and Radiance creator) said he preferred leaving them, I don't have a preference.

There is also a ChromaRadianceOptions node that allows overriding some parameters. The main one of interest is likely nerf_tile_size: setting it as high as possible without OOMing is beneficial for performance. I've been told that Blackwell benefits from having it set to 64, the default is currently 32.

I have a few more things I want to do before I'd consider it ready to merge (mostly minor stuff like better tooltips for the nodes) so this is starting out as a draft pull. Any comments/feedback are very welcome!

edit: Going to move this out of draft status now. Possible things to resolve before merging:

This subclasses the existing Chroma module instead of just copying the code. I try to use code when possible, but this does mean future modifications to Chroma need to take Radiance using it into account. Also, I don't think there is an existing precedent for doing it this way in the existing code. I can change this if necessary.
Possibly remove the ChromaRadianceLatentToImage and ChromaRadianceImageToLatent nodes to clean up the node list. They aren't doing anything unique, on the other hand they are just wrappers around the existing conversion code so maintaining them doesn't add much of a burden and removing them would possibly break existing workflows.
ChromaRadianceStubVAE - I'm not good at naming stuff so this may not be the best name. It would be kind of a pain to change now since it would break the existing workflows (though changing the display name would be free). If necessary this can be changed.
The memory usage factor for the model is a guesstimate that seems to work fairly well.
Default values for NeRF tile size (32) and NeRF head dtype (overriding it to float32 was recommended by Lodestone, which is what does currently).
Code style/implementation details? I mostly tried to stick with what the existing code does.

Let me know if anything needs to be changed to make this mergeable. General comments/feedback are also welcome, of course.

yamatazen · 2025-09-09T04:36:58Z

Can you show some generation examples?

without-ordinary · 2025-09-09T08:11:59Z

Can you show some generation examples?

yamatazen · 2025-09-09T08:16:40Z

Can you show some generation examples?

The image is not loading

yamatazen · 2025-09-09T13:42:21Z

Can you show some generation examples?

The image is not loading

Oh it loaded now

lodestone-rock · 2025-09-09T23:04:57Z

blepping · 2025-09-12T22:52:40Z

~~This merges but probably needs to be updated for #9639 to work properly. I will try to do that today. Not very likely it would have been merged in the next few hours, but mentioning this just in case.~~

edit: No changes were needed since Radiance delegates to the Chroma forward_orig.

…e device

Add batch dimension and fix channels when necessary in ChromaRadianceImageToLatent node

… usage factor

Bump Radiance nerf tile size to 32 Support EasyCache/LazyCache on Radiance (maybe)

…be in line with existing VAE behavior

Cleanups/refactoring to reduce code duplication with Chroma.

Minor code cleanups and tooltip improvements

comfyanonymous · 2025-09-13T04:26:31Z

comfy/model_detection.py

+                dit_config["nerf_max_freqs"] = 8
+                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


does this need to be in fp32?

Thanks for the review! Lodestone though it would be beneficial to run the embedder in 32-bit precision so that's why I did it that way. I let him know about your question, hopefully he will show up here to give you a better answer.

yep, that's just positional embedding part, the rest of the nerf head is in lower precision

comfyanonymous · 2025-09-13T04:33:07Z

Instead of adding a new node it might be best to just add the chroma radiance "vae" in the list of the "Load VAE" node

Then you can remove ChromaRadianceLatentToImage, ChromaRadianceImageToLatent and ChromaRadianceStubVAE.

blepping · 2025-09-13T04:53:50Z

Instead of adding a new node it might be best to just add the chroma radiance "vae" in the list of the "Load VAE" node

I can definitely do it that way. It might make sense to name it something non-Chroma Radiance specific since presumably any future pixel-space models will be B, 3, H, W and values between -1 and 1. Either way, what would you like me to call it in the list? Should the actual implementation be in nodes.py or maybe comfy/sd.py? comfy_extras/nodes_chroma_radiance.py imports the base nodes.py so I think I have to move it to avoid a circular dependency.

We actually could remove the EmptyChromaRadianceLatentImage node as well and just VAE "encode" the output from the builtin EmptyImage node:

Might be unintuitive for users though.

Add a chroma_radiance option to the VAELoader builtin node which uses comfy.sd.PixelspaceConversionVAE Add a PixelspaceConversionVAE to comfy.sd for converting BHWC 0..1 <-> BCHW -1..1

blepping · 2025-09-13T05:43:59Z

@comfyanonymous Hopefully the changes in the last commit are in right direction, if not please let me know!

comfyanonymous · 2025-09-13T22:01:33Z

Did some changes: #9851

The VAE should behave like an actual VAE meaning you can save it in the checkpoint and it should auto load.

I also made any resolution work with the model.

I also changed the name of the "vae" to "pixel_space"

blepping · 2025-09-14T02:41:17Z

@comfyanonymous

The VAE should behave like an actual VAE meaning you can save it in the checkpoint and it should auto load.

Nice, I didn't understand the VAE code there enough to mess with it and I wasn't sure if affecting existing code by trying to make things generic would get accepted. That's better than the approach I used.

I also made any resolution work with the model.

Your change made it so it doesn't crash when the resolution isn't a multiple of 16 but... It definitely does not seem to like resolutions that aren't divisible by 16.

I don't know if @lodestone-rock has any insights on how to actually make it work properly with resolutions that aren't multiples of 16, but based on what he told me in the past that is a requirement.

lodestone-rock · 2025-09-14T05:17:37Z

hmm yeah the model works on 16x16 patch so rounding to the nearest multiple of 16 is the intended behaviour.
if the resolution is arbitrary, zero padding to the nearest multiple of 16 should be sufficient i think

the reason it sorta works is because im using "conv" to do ViT like patch computation instead of flattening each patch and do matmul on it.

* Initial Chroma Radiance support * Minor Chroma Radiance cleanups * Update Radiance nodes to ensure latents/images are on the intermediate device * Fix Chroma Radiance memory estimation. * Increase Chroma Radiance memory usage factor * Increase Chroma Radiance memory usage factor once again * Ensure images are multiples of 16 for Chroma Radiance Add batch dimension and fix channels when necessary in ChromaRadianceImageToLatent node * Tile Chroma Radiance NeRF to reduce memory consumption, update memory usage factor * Update Radiance to support conv nerf final head type. * Allow setting NeRF embedder dtype for Radiance Bump Radiance nerf tile size to 32 Support EasyCache/LazyCache on Radiance (maybe) * Add ChromaRadianceStubVAE node * Crop Radiance image inputs to multiples of 16 instead of erroring to be in line with existing VAE behavior * Convert Chroma Radiance nodes to V3 schema. * Add ChromaRadianceOptions node and backend support. Cleanups/refactoring to reduce code duplication with Chroma. * Fix overriding the NeRF embedder dtype for Chroma Radiance * Minor Chroma Radiance cleanups * Move Chroma Radiance to its own directory in ldm Minor code cleanups and tooltip improvements * Fix Chroma Radiance embedder dtype overriding * Remove Radiance dynamic nerf_embedder dtype override feature * Unbork Radiance NeRF embedder init * Remove Chroma Radiance image conversion and stub VAE nodes Add a chroma_radiance option to the VAELoader builtin node which uses comfy.sd.PixelspaceConversionVAE Add a PixelspaceConversionVAE to comfy.sd for converting BHWC 0..1 <-> BCHW -1..1

blepping force-pushed the feat_support_chroma_radiance branch 2 times, most recently from 70a17b1 to 85bfc39 Compare September 5, 2025 02:38

Green-Sky mentioned this pull request Sep 9, 2025

Chroma Radiance (v0.4) leejet/stable-diffusion.cpp#807

Open

blepping force-pushed the feat_support_chroma_radiance branch from 7e4045e to 747abae Compare September 9, 2025 20:12

blepping force-pushed the feat_support_chroma_radiance branch from b64cc0b to 19bd9b8 Compare September 11, 2025 01:43

blepping marked this pull request as ready for review September 11, 2025 02:26

blepping requested review from Kosinkadink, christian-byrne, comfyanonymous, guill, ltdrdata, pythongosssss, robinjhuang, webfiltered and yoland68 as code owners September 11, 2025 02:26

blepping force-pushed the feat_support_chroma_radiance branch 2 times, most recently from 6e5f987 to ca5d3ff Compare September 12, 2025 21:32

blepping added 7 commits September 12, 2025 20:48

Initial Chroma Radiance support

52acaa6

Minor Chroma Radiance cleanups

d15a96e

Update Radiance nodes to ensure latents/images are on the intermediat…

1b0b2b7

…e device

Fix Chroma Radiance memory estimation.

913fe08

Increase Chroma Radiance memory usage factor

6e44dc8

Increase Chroma Radiance memory usage factor once again

0f584ea

Ensure images are multiples of 16 for Chroma Radiance

d1a3a5d

Add batch dimension and fix channels when necessary in ChromaRadianceImageToLatent node

blepping added 13 commits September 12, 2025 20:48

Tile Chroma Radiance NeRF to reduce memory consumption, update memory…

53fc2f0

… usage factor

Update Radiance to support conv nerf final head type.

93d9933

Allow setting NeRF embedder dtype for Radiance

ce1c679

Bump Radiance nerf tile size to 32 Support EasyCache/LazyCache on Radiance (maybe)

Add ChromaRadianceStubVAE node

e5aee26

Crop Radiance image inputs to multiples of 16 instead of erroring to …

731a892

…be in line with existing VAE behavior

Convert Chroma Radiance nodes to V3 schema.

a3e5850

Add ChromaRadianceOptions node and backend support.

e7073b5

Cleanups/refactoring to reduce code duplication with Chroma.

Fix overriding the NeRF embedder dtype for Chroma Radiance

50f3b65

Minor Chroma Radiance cleanups

a46078a

Move Chroma Radiance to its own directory in ldm

f1f5b7d

Minor code cleanups and tooltip improvements

Fix Chroma Radiance embedder dtype overriding

42349be

Remove Radiance dynamic nerf_embedder dtype override feature

cc6e7d6

Unbork Radiance NeRF embedder init

0828916

blepping force-pushed the feat_support_chroma_radiance branch from ca5d3ff to 0828916 Compare September 13, 2025 02:58

comfyanonymous reviewed Sep 13, 2025

View reviewed changes

Remove Chroma Radiance image conversion and stub VAE nodes

1bc45d3

Add a chroma_radiance option to the VAELoader builtin node which uses comfy.sd.PixelspaceConversionVAE Add a PixelspaceConversionVAE to comfy.sd for converting BHWC 0..1 <-> BCHW -1..1

comfyanonymous merged commit c1297f4 into comfyanonymous:master Sep 13, 2025
9 checks passed

Add support for Chroma Radiance #9682

Add support for Chroma Radiance #9682

Conversation

blepping commented Sep 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yamatazen commented Sep 9, 2025

Uh oh!

without-ordinary commented Sep 9, 2025

Uh oh!

yamatazen commented Sep 9, 2025

Uh oh!

yamatazen commented Sep 9, 2025

Uh oh!

lodestone-rock commented Sep 9, 2025

Uh oh!

blepping commented Sep 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

comfyanonymous Sep 13, 2025

Choose a reason for hiding this comment

Uh oh!

blepping Sep 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lodestone-rock Sep 13, 2025

Choose a reason for hiding this comment

Uh oh!

comfyanonymous commented Sep 13, 2025

Uh oh!

blepping commented Sep 13, 2025

Uh oh!

blepping commented Sep 13, 2025

Uh oh!

Uh oh!

comfyanonymous commented Sep 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

blepping commented Sep 14, 2025

Uh oh!

lodestone-rock commented Sep 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

blepping commented Sep 2, 2025 •

edited

Loading

blepping commented Sep 12, 2025 •

edited

Loading

blepping Sep 13, 2025 •

edited

Loading

comfyanonymous commented Sep 13, 2025 •

edited

Loading

lodestone-rock commented Sep 14, 2025 •

edited

Loading