diffusers-modular/krea2-edit

Krea 2 reference-image edit — Modular Diffusers blocks

Custom Modular Diffusers blocks that reproduce the ostris/Krea2OstrisEdit reference-image ("edit") workflow for Krea 2, loadable as remote code on top of stock diffusers.

import torch
from transformers import Qwen3VLProcessor
from diffusers import ClassifierFreeGuidance
from diffusers.modular_pipelines import ModularPipelineBlocks

blocks = ModularPipelineBlocks.from_pretrained("diffusers-modular/krea2-edit", trust_remote_code=True)
pipe = blocks.init_pipeline("krea/Krea-2-Turbo")        # weights from the base repo
pipe.load_components(torch_dtype=torch.bfloat16)
pipe.update_components(processor=Qwen3VLProcessor.from_pretrained("Qwen/Qwen3-VL-4B-Instruct"))
pipe.update_components(guider=ClassifierFreeGuidance(guidance_scale=0.0, use_original_formulation=True))
pipe.to("cuda")

from PIL import Image
image = pipe(
    prompt="a white yeti with horns reading a book",
    image=Image.open("reference.png"),   # one or more reference images
    num_inference_steps=8, mu=1.15, output="images",
)[0]

What it does

Reference images condition generation two ways (matching how the Ostris AI-Toolkit edit LoRAs train):

1. Qwen3-VL prompt embedding — a coarse view of each reference is embedded into the text conditioning through the vision tower.
2. Clean VAE latents at flow time t=0 — each reference is VAE-encoded and appended to the transformer sequence as clean tokens on its own rotary frame axis, so the noisy image tokens attend to it at every block.

The bundled Krea2Transformer2DModel adds a small, backward-compatible ref_seq_len argument to the Krea 2 transformer forward (t=0 modulation of the reference span; those tokens are excluded from the predicted velocity). With ref_seq_len=0 it is numerically identical to plain Krea 2 text-to-image.

Files

• block.py — entry point (Krea2EditBlocks), referenced by config.json's auto_map.
• transformer_krea2.py — the Krea 2 transformer (with the ref_seq_len edit path).
• modular_blocks_krea2*.py, encoders.py, before_denoise.py, denoise.py, decoders.py, inputs.py, modular_pipeline.py — the modular blocks.