File size: 3,291 Bytes
4684d79 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | # Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from diffusers.utils import logging
from diffusers.modular_pipelines.modular_pipeline import SequentialPipelineBlocks
from diffusers.modular_pipelines.modular_pipeline_utils import InsertableDict, OutputParam
from .before_denoise import (
Krea2EditRoPEInputsStep,
Krea2PrepareLatentsStep,
Krea2SetTimestepsStep,
)
from .decoders import Krea2AfterDenoiseStep
from .denoise import Krea2EditDenoiseStep
from .encoders import (
Krea2EditReferenceLatentsStep,
Krea2EditTextEncoderStep,
)
from .inputs import Krea2TextInputsStep
from .modular_blocks_krea2 import Krea2DecodeStep
logger = logging.get_logger(__name__)
# Krea 2 (reference-image edit)
class Krea2EditCoreDenoiseStep(SequentialPipelineBlocks):
model_name = "krea2"
block_classes = [
Krea2TextInputsStep(),
Krea2PrepareLatentsStep(),
Krea2SetTimestepsStep(),
Krea2EditRoPEInputsStep(),
Krea2EditDenoiseStep(),
Krea2AfterDenoiseStep(),
]
block_names = [
"input",
"prepare_latents",
"set_timesteps",
"prepare_rope_inputs",
"denoise",
"after_denoise",
]
@property
def description(self):
return (
"step that denoises noise into an image for the reference-image edit task. It includes the denoise loop "
"(with the clean reference tokens appended each step), as well as preparing the inputs (timesteps, "
"latents, rope inputs etc.)."
)
@property
def outputs(self):
return [
OutputParam.template("latents"),
]
# Preset: reference-image edit. `image` carries the reference image(s); without it, use the text2image preset.
EDIT_BLOCKS = InsertableDict(
[
("text_encoder", Krea2EditTextEncoderStep()),
("reference_latents", Krea2EditReferenceLatentsStep()),
("denoise", Krea2EditCoreDenoiseStep()),
("decode", Krea2DecodeStep()),
]
)
class Krea2EditBlocks(SequentialPipelineBlocks):
"""
Modular pipeline for the Krea 2 reference-image edit task.
Pass one or more reference images as `image`; they condition generation both through the Qwen3-VL text encoder
and as clean VAE latents appended to the transformer sequence at flow time t=0 (matching the Ostris AI-Toolkit
edit LoRAs). Load an edit/style LoRA into the transformer to control the effect.
"""
model_name = "krea2"
block_classes = EDIT_BLOCKS.values()
block_names = EDIT_BLOCKS.keys()
@property
def description(self):
return "Auto Modular pipeline for the reference-image edit task using Krea 2. Requires `prompt` and `image`."
|