Adding "Image Align with Rife" and "Wavelet Color Fix" Nodes

[pifroggi]

Hey! I was asked if I could contribute my "Image Align with Rife" node to the official chaiNNer repository. For this purpose I've created this pull request. In addition it also adds a "Wavelet Color Fix" node. The code was tested with Alpha v0.22.2 on Windows.
I've originally not contributed as I'm unfamiliar with the process (this is my first time using github) and I have only very limited coding abilities, for which I would like to apologize already in case I've done something wrong.

Image Align with Rife

Aligns an Image with a Reference Image using a modified Rife. Images should have vague alignment before using this Node. Output Image will have the same dimensions as Reference Image. Resize Reference Image to get desired output scale.

Examples (1.Input 2.Reference 3.Output): https://slow.pics/c/rqeq3D97

Wavelet Color Fix

Correct for upscaling model color shift by first separating the image into wavelets of different frequencies, then matching the average color of the Input Image to that of a Reference Image. In general produces better results than the Average Color Fix at the cost of more computation. The Wavelet Color Fix functions are from sd-webui-stablesr.

Example transferring DVD colors to upscaled image: https://imgsli.com/MjM5NzM5/0/2

Example Chain

An example chain that produces good alignment in general can be found here:
https://github.com/pifroggi/custom_chainner_nodes/tree/main/Examples

I hope everything makes sense!

[joeyballentine]

The binary file you included for the rife model would increase the overall size of chainner's installer by 20%. Would it be possible for you to have the code download the model to a folder in chainner's python folder if it doesn't exist, rather than bundling it?

[pifroggi]

Yes, I can try that! The official download would be a zip file hosted on google drive that contains additional files as well that won't be needed.
https://drive.google.com/file/d/1BjuEY7CHZv1wzmwXSQP9ZTj0mLWu_4xy/view
Should we use that, or host it here on github maybe?

[joeyballentine]

The Google drive link should work, you'll just have to convert it to a direct download link first

[pifroggi]

The model now gets automatically downloaded into chainner's "python/models/rife_v4.14" if it's not already there. Let me know if this path works.

[joeyballentine]

Sorry for the delay, will review this when i have some time

[pifroggi]

No worries, thanks! :)

[joeyballentine]

The code for this looks fine to me. I'll test it out today and just verify everything works. thanks again for this

[pifroggi]

Great, thank you as well!

[joeyballentine]

Please run the ruff formatter on the code as well. Thanks

[pifroggi]

I'll look into the ruff formatter you mentioned and do the other changes likely later today.

[Splendide-Imaginarius]

The Wavelet Color Fix functions are from sd-webui-stablesr.

Does this PR also include the adain_color_fix function? (I don't see it here but I wasn't sure if I missed it.) If not, is there any reason it's excluded?

[Splendide-Imaginarius]

How does the Wavelet Color Fix handle fully transparent pixels? E.g. if you take an image that has some fully transparent regions, and run magick in.png -background lime -alpha background out.png on it, and upscale the result, what does Wavelet Color Fix look like? (See #1907 for more discussion.)

[pifroggi]

Hey, that sounds great about your homography alignment! :) The wavelet color fix node currently only takes 3 channel inputs, so if your images have transparency, you would have to use for example a "Split Transparency" node and then recombine it with the output again.

I've not included the Adain Color Fix as it didn't produce good results on my use case, but I didn't do a lot of testing on it outside of that, sorry about that. It works quite differently, more similarly to the "Color Transfer" node in chainner. As in the content of the two images can be completely different and it tries to match the overall color palette.
While the Wavelet Color Fix is very similar to the Average Color Fix, as in the content of the two images need to be in the same place spatially.

@joeyballentine Since the two alignment nodes will stay separate for now, I'm guessing the Average Color Fix and Wavelet Color Fix (which requires pytorch) will as well. Regarding the "Image Align with Rife" name, let me know if you would still prefer to use a different naming scheme.

We could maybe either describe it a bit more broadly with something like:
Fast Image Alignment
Refine Image Alignment

Or keep it more specific to which methods are used with:
Homography Image Alignment
Rife Image Alignment

[joeyballentine]

I would prefer to keep the implementation detail out of the node name. We don't say "resize with chainner-rs" or "inpaint with opencv". That's useless information to someone just trying to do the thing

[pifroggi]

I agree, sometimes it's just hard to find good words to differentiate them when there is more than one and I was unsure if "fast" and "refine" are good fits here. For now I've left it with "Align Image to Reference" and when @Splendide-Imaginarius's node comes it, we can change that if needed.

Other changes:
removed all the commented out unneeded code
added minimum of 1 where needed
download now goes into appdata/roaming/chaiNNer/python/rife_v4.14/weights
ran ruff

[Splendide-Imaginarius]

Is this using a standard RIFE model, or has the model been modified in some way? Is there a reason why you're not just letting the user provide the model file, so that the user can use different RIFE models per their preference? I'd be a lot more comfortable with having a node that can load a RIFE model from an arbitrary file (just like loading upscale models), so that the user can pick their favorite model (or train their own); the behavior here feels way too magic and inflexible for my tastes.

[Splendide-Imaginarius]

The wavelet color fix node currently only takes 3 channel inputs, so if your images have transparency, you would have to use for example a "Split Transparency" node and then recombine it with the output again.

Yeah I know, I was mainly asking whether the end result looks OK when the alpha channel has been split out and recombined. Average Color Fix needed a specific workaround (masking out the transparent pixels when calculating averages) to avoid contaminating the image with the RGB data from the transparent pixels. It sounds like Wavelet Color Fix may be less susceptible to that issue, but I was just curious if you had tried it. It's not a big deal if you haven't tried it, I can test it myself later, and it's not a reason to block the PR.

[pifroggi]

Yeah I know, I was mainly asking whether the end result looks OK when the alpha channel has been split out and recombined. Average Color Fix needed a specific workaround (masking out the transparent pixels when calculating averages) to avoid contaminating the image with the RGB data from the transparent pixels.

Ah I think I understand the discussion from the other issue now, sorry for the misinterpretation. I didn't think of that, but now that you say that, that is likely the case if the two images have different RGB values on the transparent pixels.

Is this using a standard RIFE model, or has the model been modified in some way? Is there a reason why you're not just letting the user provide the model file, so that the user can use different RIFE models per their preference?

The RIFE model itself has not been modified, but the model architecture in this file, which is needed to run the model, has been modified:
https://github.com/chaiNNer-org/chaiNNer/blob/fa8b9b5040ad477b149ed4aee0b9c1c6369fc303/backend/src/packages/chaiNNer_pytorch/pytorch/processing/rife/IFNet_HDv3_v4_14_align.py
Unfortunately the RIFE model architecture changes quite frequently, so similar changes would have to be made to the older or potentially newer versions.

[pifroggi]

Hey, I just wanted to check if something would still be missing from my side. I saw there is still some issue with ruff (or is this from the older commits?), but when I do "ruff check --fix filename" I'm getting "All checks passed!".

Sorry if I'm misinterpreting the Github UI.

[joeyballentine]

I'll just fix that later when I'm and to and then this should be good I think

[joeyballentine]

Thanks again for this and sorry it took me so long to do the remaining fixes and review.

[pifroggi]

Awesome and no worries, I'm glad it worked out! Looking forward to the next release! :)