DCAE state key mapping temporary fix

[zergzzlun]

Hi @city96 , thanks for all the awesome work you've done here.

I’ve encountered the same grey/black output issue reported in #93 , which appears to be related to the VAE weight loading process. I tried downloading multiple versions of the weight files, including the one mentioned in the current README (as well as from #87 ), but unfortunately, none of them worked.

I’ve written a temporary key mapping solution for the weight model available here. This successfully resolved the issue, and you can check out the mapping script here along with the modified loader.

I'd be happy to submit a pull request to offer this as a temporary fix before the coming up major rewrite you're working on. Please let me know if you see that helpful.

[city96]

Hey. Thank you, this is super useful. If you send a PR aimed at the main branch I can most likely just include/cherry pick it on the rewrite branch.

Also kinda planning to just making the rewrite a PR to upstream ComfyUI so I can turn this repo back to what I originally used it for lol (completely random, small research models.) - probably better for PixArt / Sana to have first class support than through a node pack.

[zergzzlun]

Hi @city96 Thanks for the positive feedback.

I’ve just submitted a PR aimed at the main branch as you suggested: #103 . Hope it helps and I'm happy to adjust anything if needed.

Pushing the rewrite upstream to ComfyUI sounds like a really good idea, I can imagine its gonna make life much easier for users of PixArt / Sana & ComfyUI community. Please let me know if there's anything I can do to help.

[city96]

I guess one part that may need adjustment is whether we should run the convert logic or not (as to not break weights that are already in the right format).

i.e. just something simple like this (key for detection may be wrong)

if "decoder.project_out.op_list.0.bias" in sd:
    from .models import dcae_key_mapping
    sd = dcae_key_mapping.convert_sd(sd)

That + I think you could simplify the loop if you did sth like this, but I can sort that out, the explicit list is fine for now.

for k in range(6):
    mapping[f"encoder.stages.{k}.op_list.0.main.conv1.conv.bias"] = f"encoder.stages.{k}.0.conv1.conv.bias",
    ...

[zergzzlun]

Thanks for the thoughtful feedback.

That check does make a lot of sense. I completely agree it’s a good idea to check if the weights are already in the correct format before running the conversion logic. I've already added the key check in the PR. I've double-checked the key is correct, so no worries there. I'm happy to make any further adjustment if needed.

Regarding the simplification, I definitely see where you’re going, it is really good idea with clean and elegant code style. I made some early attempts on it and find out the block structures differs enough that a fully generalized approach might not be as straightforward as expected. Given that I figure we might use explicit list as a temporary solution for now. It makes sense to revisit this later, and I’d be happy to give it another look. Please let me know if I can be of help.