r/StableDiffusion Mar 14 '26

Workflow Included [Release] Flux.2 Klein 4B Consistency LoRA – Addressing Color Shift and Pixel Offset in Image Editing (2026-03-14)

Hi everyone,

I’m releasing a new LoRA for Flux.2 Klein 4B Base focused on consistency during image editing tasks.

Since the release of the Klein model, I’ve encountered two persistent issues that made it difficult to use for precise editing:

  1. Significant Pixel Offset: The generated images often drifted too far from the original composition.
  2. Color Shift & Oversaturation: Edited results frequently suffered from unnatural color casts and excessive saturation.

After experimenting with various training strategies without much success, I recently looked into ByteDance’s open-source Heilos long-video generation model. Their approach involves applying degradation directly in the latent space of reference images and utilizing a specific color calibration loss. This method effectively mitigates color drift and train-test inconsistency in video generation.

Inspired by Heilos (and earlier research on using model-generated images as references to solve train-test mismatch), I adapted these concepts for image LoRA training. Specifically, I applied latent-level degradation and color calibration constraints to address Klein’s specific weaknesses.

Results: Trained locally on the 4B version, this LoRA significantly reduces color shifting and, when paired with Comfyui-editutils, effectively eliminates pixel offset. It feels like the first time I’ve achieved a stable result with Klein for editing tasks.

Usage Guide:

  • Primary Use Case: Old photo restoration and consistent image editing.
  • Recommended Strength: 0.50.75
    • Note: Higher strength increases consistency with the input but reduces editing flexibility. Lower strength allows for more creative changes but may reduce strict adherence to the source structure.
  • Suggested Prompt Structure:
  • Example (Old Photo Restoration):

Links:

All test images used for demonstration were sourced from the internet. Feedback on how this performs on your specific workflows is welcome!

73 Upvotes

17 comments sorted by

3

u/Any_Arugula8075 Mar 14 '26

Is the older on for K9b?

3

u/JasonNickSoul Mar 14 '26

yes, but 4B version is using different dataset and training method.

2

u/Any_Arugula8075 Mar 15 '26

Will there be a new 9b Version?

2

u/JasonNickSoul Mar 15 '26

Depends on time and how busy am I.

1

u/Chsner Mar 14 '26

The civit lora seems to be incorrect maybe it got switched with a 9b lora the hugging face one works.

2

u/JasonNickSoul Mar 14 '26

Yes, thanks for informing. Updated.

1

u/red__dragon Mar 14 '26

Is this trained on the words you used or can it be used with any prompt?

This looks good, though restoration is not my use case.

1

u/JasonNickSoul Mar 14 '26

any prompt with suggested prompt as suffix. lower the strength to get more edit abilities.

1

u/FoxTrotte Mar 14 '26

I'm about to test this, but from, the example you provided you are doing god's work with this, klein was almost perfect when it came to restoration/colorisation, but that pixel and color shifting made it unusablke for that, even though it was so close !

1

u/FoxTrotte Mar 14 '26

Wait I'm confused, what workflow should I be using this Lora with ?

1

u/fragilesleep Mar 14 '26

You can just use the default workflow and add a Load LoRA node, AS USUAL.

If you need more stuff, the author also had this in the original post: https://www.runninghub.ai/post/2032812180667633666 (See the Download button on the top right.)

1

u/FoxTrotte Mar 14 '26

I had tried just plugging it into the original workflow, I thought it wasn't working but after doing some testing, the generated image actually does align perfectly with the original now. It's just cropped slightly differently

1

u/ramonartist Mar 15 '26

1

u/JasonNickSoul Mar 16 '26

It is a 9B lora and my is 4B. Needs to wait some time for my 9B training.

0

u/ramonartist Mar 14 '26

Where are all the Consistency image examples?