OzzyGT/krea2_preview_blocks

Krea 2 prediction-preview modular blocks

Custom Modular Diffusers blocks that show what the model is aiming at while it denoises. A preview built from the loop's latents is noise that slowly resolves and will stay full of noise for the first half of a run. A preview built from the x0 prediction, what the model thinks the image will be judged from the current step, is a whole image from step one, which then sharpens on each step.

what the loop holds	what the model is aiming at

Same run, same seed, 8 steps of Krea 2 Turbo.

preview_callback is called as preview_callback(step, total, image), once per previewed step, and preview_mode chooses how much work goes into each one:

preview_mode="rgb"        16x3 linear projection of the latent channels. No weights, under 1 ms at 1024x1024
preview_mode="taew"       TAEW 2.1 tiny autoencoder. Real detail, ~21 ms at 1024x1024
preview_mode="taew_fast"  the same decoder at quarter size, ~6 ms
preview_mode="latents"    no decode; the raw prediction, for callers that decode themselves

The taew modes download the decoder weights, 22 MB, the first time they run.

The blocksets are the stock Krea 2 ones with a preview step added to the denoise loop, so with no preview_callback they generate exactly what the stock pipeline does.

Loading & running

import sdnq  # needed to load the quantized text encoder
import torch

from diffusers import ModularPipeline


pipe = ModularPipeline.from_pretrained(
    "OzzyGT/krea2_preview_blocks", trust_remote_code=True
)
pipe.load_components(dtype=torch.bfloat16)
pipe.to("cuda")


def on_preview(step, total, image):
    image.save(f"preview_{step:02d}.png")


image = pipe(
    prompt="a photograph of a cat wearing a wool hat",
    height=1024,
    width=1024,
    num_inference_steps=8,
    preview_callback=on_preview,
    preview_mode="rgb",
    output="images",
)[0]
image.save("result.png")

These blocks only work with Krea 2, Turbo or raw.

For loading a different checkpoint or swapping components, see Modular pipeline in the diffusers docs.

Previews

preview_every=N previews every Nth step; the final step is always previewed. The callback runs inside the denoising loop, so whatever it does is time the loop is not spending on the model. rgb is effectively free and the one I'd recommend; the tiny decoders cost real time on every previewed step.

Raising from the callback aborts the run: the exception propagates out of pipe(...) and the pipeline is reusable afterwards. That is the cancel button.

class Cancelled(Exception):
    pass


def on_preview(step, total, image):
    if stop_requested:
        raise Cancelled

The loop logs a traceback before re-raising, so a clean cancel still prints one.

References

• madebyollin/taehv -- TAEW 2.1, vendored here as taehv.py (MIT, (c) 2025 Ollin Boer Bohan)
• ComfyUI latent_formats.py -- source of the Wan 2.1 latent-to-RGB factors