# Copyright 2026 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. """Custom modular-diffusers blocks adding reference-image conditioning to Krea 2. Krea 2's text encoder is the full Qwen3-VL VLM, vision tower included, so reference images can condition generation through the encoder's vision path. Two community edit LoRAs also feed the references as *clean VAE latents* inside the transformer sequence. All three paths live here, selected by `reference_mode`: * `"off"` (default) -- vision path only. Works on the stock base checkpoint, no LoRA. * `"append"` -- vision (coarse view) + clean VAE reference tokens **after** the noisy target, modulated at flow time `t=0`. The Ostris AI-Toolkit reference/edit LoRA recipe. * `"prepend"` -- vision (grounded view) + the clean VAE source **before** the noisy target under the plain uniform timestep modulation, distinguished only by its RoPE frame axis. The Identity-Edit recipe (ai-toolkit `predict_velocity_edit`); its unconditional branch is grounded on the same image. With `reference_mode="off"` and no images the blocks reduce to stock text-to-image. No mode patches the `Krea2Transformer2DModel` source. """ import math from typing import Any import numpy as np import PIL.Image import torch import torch.nn.functional as F from transformers import Qwen3VLProcessor from diffusers.configuration_utils import FrozenDict from diffusers.guiders import ClassifierFreeGuidance from diffusers.models import AutoencoderKLQwenImage from diffusers.models.transformers.transformer_krea2 import Krea2Transformer2DModel from diffusers.modular_pipelines.krea2.before_denoise import ( Krea2PrepareLatentsStep, Krea2PreparePositionIdsStep, Krea2SetTimestepsStep, Krea2TurboSetTimestepsStep, ) from diffusers.modular_pipelines.krea2.decoders import Krea2DecodeStep from diffusers.modular_pipelines.krea2.denoise import ( Krea2DenoiseLoopWrapper, Krea2LoopAfterDenoiser, Krea2LoopBeforeDenoiser, ) from diffusers.modular_pipelines.krea2.encoders import ( _PROMPT_TEMPLATE_ENCODE_PREFIX, _PROMPT_TEMPLATE_ENCODE_START_IDX, _PROMPT_TEMPLATE_ENCODE_SUFFIX, KREA2_TEXT_ENCODER_SELECT_LAYERS, Krea2TextEncoderStep, ) from diffusers.modular_pipelines.krea2.modular_pipeline import Krea2ModularPipeline from diffusers.modular_pipelines.modular_pipeline import ( BlockState, ConditionalPipelineBlocks, ModularPipelineBlocks, PipelineState, SequentialPipelineBlocks, ) from diffusers.modular_pipelines.modular_pipeline_utils import ( ComponentSpec, InputParam, InsertableDict, OutputParam, ) from diffusers.utils import logging from .krea2_vision_attention import build_patch_keeps, vision_key_masks logger = logging.get_logger(__name__) # pylint: disable=invalid-name REFERENCE_MODE_OFF = "off" REFERENCE_MODE_APPEND = "append" REFERENCE_MODE_PREPEND = "prepend" REFERENCE_MODES = (REFERENCE_MODE_OFF, REFERENCE_MODE_APPEND, REFERENCE_MODE_PREPEND) # What a painted mask does to the Qwen3-VL vision path. Zeroing output tokens alone cannot remove content # from a reference (see `krea2_vision_attention`), so `deemphasize` attenuates rather than excludes. REFERENCE_MASK_MODE_EXCLUDE_BLANK = "exclude_blank" REFERENCE_MASK_MODE_EXCLUDE = "exclude" REFERENCE_MASK_MODE_DEEMPHASIZE = "deemphasize" REFERENCE_MASK_MODES = ( REFERENCE_MASK_MODE_EXCLUDE_BLANK, REFERENCE_MASK_MODE_EXCLUDE, REFERENCE_MASK_MODE_DEEMPHASIZE, ) # Qwen3-VL vision-path budgets. `off` matches the reference ComfyUI encoder's default; `append` matches the # AI-Toolkit recipe. `prepend` sizes its view from `grounding_px` instead. VISION_MAX_PIXELS = 1024 * 1024 VISION_EDIT_MAX_PIXELS = 384 * 384 REFERENCE_LATENTS_MAX_PIXELS = 1024 * 1024 # Longest-side cap (px) for the `prepend` grounding view; 0 means native. The Identity-Edit LoRA trained on # 384-768px jitter. DEFAULT_GROUNDING_PX = 768 def _coerce_reference_mode(value: Any) -> str: """Normalize a reference mode, falling back to `"off"` for anything unrecognized.""" if isinstance(value, str) and value in REFERENCE_MODES: return value if value is not None and value != REFERENCE_MODE_OFF: logger.warning(f"Unknown `reference_mode` {value!r}; falling back to {REFERENCE_MODE_OFF!r}.") return REFERENCE_MODE_OFF def _coerce_reference_mask_mode(value: Any) -> str: """Normalize a reference-mask mode, falling back to `"exclude_blank"` for anything unrecognized.""" if isinstance(value, str) and value in REFERENCE_MASK_MODES: return value if value is not None: logger.warning( f"Unknown `reference_mask_mode` {value!r}; falling back to {REFERENCE_MASK_MODE_EXCLUDE_BLANK!r}." ) return REFERENCE_MASK_MODE_EXCLUDE_BLANK def _as_list(value: Any) -> list: """Wrap a bare value into a single-element list; `None` becomes an empty list.""" if value is None: return [] if isinstance(value, (list, tuple)): return list(value) return [value] def _broadcast_per_image(value: Any, num_images: int, default: float) -> list[float]: """Expand a scalar (or short list) of per-image strengths to one float per reference image.""" values = _as_list(value) if not values: return [default] * num_images if len(values) == 1: return [float(values[0])] * num_images values = [float(v) for v in values[:num_images]] return values + [default] * (num_images - len(values)) def _to_rgb(image) -> PIL.Image.Image: if not isinstance(image, PIL.Image.Image): image = PIL.Image.fromarray(np.asarray(image)) return image.convert("RGB") def _cap_area(image, max_pixels: int) -> PIL.Image.Image: """Aspect-preserving downscale so the image covers at most `max_pixels`. Never upscales.""" image = _to_rgb(image) pixels = image.width * image.height if pixels > max_pixels: scale = math.sqrt(max_pixels / pixels) image = image.resize((max(1, round(image.width * scale)), max(1, round(image.height * scale)))) return image def _cap_longest_side(image, longest_side: int) -> PIL.Image.Image: """Downscale so the longest side is at most `longest_side` px (0 = native). Never upscales.""" image = _to_rgb(image) longest = max(image.width, image.height) if longest_side and longest > longest_side: scale = longest_side / longest image = image.resize((max(1, round(image.width * scale)), max(1, round(image.height * scale)))) return image def _mask_to_keep_grid(mask, grid_height: int, grid_width: int): """Rasterize a painted mask onto a `(grid_height, grid_width)` bool grid, `True` = keep. Painted is the part to *use*. Shared by every consumer so the polarity means the same thing on the pre-merge patch grid, the merged vision-token grid and the VAE latent patch grid. `RGBA` masks are read from their alpha channel, anything else from luminance. Returns `None` when the mask is absent or paints nothing -- an empty mask read literally would drop every token and silently disable the reference. """ if mask is None or grid_height <= 0 or grid_width <= 0: return None if not isinstance(mask, PIL.Image.Image): mask = PIL.Image.fromarray(np.asarray(mask)) painted = np.array(mask)[:, :, 3] if mask.mode == "RGBA" else np.array(mask.convert("L")) painted = PIL.Image.fromarray(painted.astype(np.uint8)).resize((grid_width, grid_height), PIL.Image.BILINEAR) keep = np.asarray(painted, dtype=np.float32) / 255.0 >= 0.5 if not keep.any(): logger.warning( "A reference mask paints nothing and is being ignored. `RGBA` masks are read from their alpha " "channel, so loading one through a helper that flattens to RGB (such as `load_image`) discards it." ) return None return keep def _mask_to_keep_vector(mask, grid_height: int, grid_width: int, temporal: int = 1): """Row-major flattened `_mask_to_keep_grid`, tiled over `temporal` frames. Row-major is correct at merged-token resolution; the pre-merge patch sequence is not row-major and uses `krea2_vision_attention.patch_keep_vector` instead. """ keep = _mask_to_keep_grid(mask, grid_height, grid_width) if keep is None: return None flat = keep.astype(np.float32).reshape(-1) if temporal > 1: flat = np.tile(flat, temporal) return torch.from_numpy(flat) def _build_keep_vectors(masks: list, image_grid_thw, merge_size: int) -> list: """Build one keep vector per row of `image_grid_thw` (or `None` where the slot has no effective mask).""" if image_grid_thw is None: return [] keep_vectors = [] for i in range(int(image_grid_thw.shape[0])): mask = masks[i] if i < len(masks) else None t, h, w = (int(x) for x in image_grid_thw[i].tolist()) keep_vectors.append(_mask_to_keep_vector(mask, h // merge_size, w // merge_size, t)) return keep_vectors def _scaled_get_image_features(original_fn, style_strengths, subject_strengths, merge_size, keep_vectors=None): """Wrap `Qwen3VLModel.get_image_features` to scale each reference's vision features per channel. Qwen3-VL injects each image twice: as scattered `<|image_pad|>` embeddings (`pooler_output`, a per-image list) and as deepstack features (per-layer tensors re-injected at deeper LLM layers). `subject` scales the former, `style` the latter; 0 removes that channel, 1.0 leaves it at full effect. `keep_vectors` (one entry per image, or `None`) additionally zeroes the vision tokens outside the painted region in both channels. """ def strength_for(values, index): return float(values[index]) if index < len(values) else 1.0 def wrapper(pixel_values, image_grid_thw, **kwargs): outputs = original_fn(pixel_values, image_grid_thw, **kwargs) def keep_for(index, like): if not keep_vectors or index >= len(keep_vectors) or keep_vectors[index] is None: return None return keep_vectors[index].to(device=like.device, dtype=like.dtype) pooled = getattr(outputs, "pooler_output", None) if isinstance(pooled, (list, tuple)): scaled = [] for i, features in enumerate(pooled): features = features * strength_for(subject_strengths, i) keep = keep_for(i, features) if keep is not None: features = features * keep.unsqueeze(-1) scaled.append(features) outputs.pooler_output = type(pooled)(scaled) deepstack = getattr(outputs, "deepstack_features", None) if deepstack: reference = deepstack[0] counts = [ int(image_grid_thw[i].prod().item()) // (merge_size * merge_size) for i in range(image_grid_thw.shape[0]) ] blocks = [] for i, count in enumerate(counts): block = torch.full( (count,), strength_for(style_strengths, i), device=reference.device, dtype=reference.dtype ) keep = keep_for(i, reference) if keep is not None: block = block * keep blocks.append(block) per_token = torch.cat(blocks).unsqueeze(-1) outputs.deepstack_features = type(deepstack)(layer * per_token for layer in deepstack) return outputs return wrapper def _build_reference_processor(text_encoder, tokenizer): """Build a `Qwen3VLProcessor` for reference-image conditioning. Krea 2 ships no image-processor config, so one is derived here. Qwen3-VL reuses the Qwen2-VL image processor; the video processor is unused but the processor class requires a non-`None` instance. The patch/merge/temporal sizes MUST come from the model's own vision config (Qwen3-VL uses `patch_size=16`, not the Qwen2-VL default of 14) or the patch-embed reshape fails. """ from transformers import Qwen3VLProcessor, Qwen3VLVideoProcessor from transformers.models.qwen2_vl.image_processing_qwen2_vl_fast import Qwen2VLImageProcessorFast vision_config = text_encoder.config.vision_config image_processor = Qwen2VLImageProcessorFast( patch_size=vision_config.patch_size, merge_size=vision_config.spatial_merge_size, temporal_patch_size=vision_config.temporal_patch_size, ) return Qwen3VLProcessor( image_processor=image_processor, tokenizer=tokenizer, video_processor=Qwen3VLVideoProcessor() ) def _pack_latents(latents: torch.Tensor, patch_size: int) -> torch.Tensor: """Pack spatial latents `(B, C, H, W)` into Krea 2's sequence `(B, (H/p) * (W/p), C * p * p)`.""" batch_size, channels, height, width = latents.shape latents = latents.view(batch_size, channels, height // patch_size, patch_size, width // patch_size, patch_size) latents = latents.permute(0, 2, 4, 1, 3, 5) return latents.reshape( batch_size, (height // patch_size) * (width // patch_size), channels * patch_size * patch_size ) def _reference_position_ids(index: int, grid_height: int, grid_width: int) -> torch.Tensor: """`(grid_h * grid_w, 3)` rotary coordinates for reference `index`: frame axis `index + 1`, own y/x grid.""" ids = torch.zeros(grid_height, grid_width, 3) ids[..., 0] = index + 1 ids[..., 1] = torch.arange(grid_height)[:, None] ids[..., 2] = torch.arange(grid_width)[None, :] return ids.reshape(-1, 3) def _register_zero_time_reference_hooks(transformer, timestep, batch_size, split, total): """Modulate the trailing `total - split` reference tokens at flow time `t=0` for the `append` mode. A `forward_pre_hook` swaps each block's broadcast `(B, 1, 6 * hidden)` timestep embedding for a per-token `(B, total, 6 * hidden)` one carrying the real timestep on the text and target rows and the `t=0` embedding on the reference rows. The block's `temb.unflatten(-1, (6, -1)) + scale_shift_table` already broadcasts over the sequence axis, so no source patching is needed. Returns the hook handles; the caller must remove them once the forward pass is done. """ dtype = transformer.dtype def time_mod(t): return transformer.time_mod_proj(F.gelu(transformer.time_embed(t, dtype=dtype), approximate="tanh")) temb = time_mod(timestep) temb_zero = time_mod(torch.zeros_like(timestep)) per_token = torch.cat([temb.expand(batch_size, split, -1), temb_zero.expand(batch_size, total - split, -1)], dim=1) def pre_hook(module, args, kwargs): if len(args) >= 2: return (args[0], per_token, *args[2:]), kwargs return args, {**kwargs, "temb": per_token} return [block.register_forward_pre_hook(pre_hook, with_kwargs=True) for block in transformer.transformer_blocks] def _concat_reference_latents(latents, reference_latents, mode, dtype): """Splice the clean reference tokens into the noisy latents in the order the mode dictates. Returns `(latents, reference_seq_len)`; `reference_seq_len` is 0 when the mode carries no latent channel. """ if mode == REFERENCE_MODE_OFF or reference_latents is None: return latents, 0 reference_latents = reference_latents.to(device=latents.device, dtype=dtype) if reference_latents.shape[0] != latents.shape[0]: reference_latents = reference_latents.expand(latents.shape[0], -1, -1) if mode == REFERENCE_MODE_PREPEND: return torch.cat([reference_latents, latents], dim=1), reference_latents.shape[1] return torch.cat([latents, reference_latents], dim=1), reference_latents.shape[1] def _slice_reference_rows(noise_pred, mode, reference_seq_len, target_seq_len): """Drop the reference rows from a prediction, wherever the mode put them.""" if not reference_seq_len: return noise_pred if mode == REFERENCE_MODE_PREPEND: return noise_pred[:, reference_seq_len:] return noise_pred[:, :target_seq_len] # auto_docstring class Krea2ReferenceImagesStep(ModularPipelineBlocks): """Base reference collector. Not used directly -- pick the subclass for the mode.""" model_name = "krea2" uses_reference_latents = False def _vision_view(self, image, block_state): raise NotImplementedError @property def description(self) -> str: return ( "Collect the reference images and size them for this mode's Qwen3-VL vision view, normalizing the " "per-image style/subject strengths and masks." ) @property def inputs(self) -> list[InputParam]: return [ InputParam( name="reference_images", type_hint=PIL.Image.Image | list[PIL.Image.Image], description="Reference image(s) to condition on. A single image or a list.", ), InputParam( name="reference_mode", type_hint=str, default=REFERENCE_MODE_OFF, description=( "How references condition the model: 'off' (Qwen3-VL vision path only, no LoRA), 'append' " "(vision + clean VAE tokens after the target, Ostris edit LoRA) or 'prepend' (vision + " "clean VAE source before the target, Identity-Edit LoRA)." ), ), InputParam( name="reference_style_strength", type_hint=float | list[float], default=1.0, description=( "Per-image scale on the deepstack vision features (low-level texture/style). A scalar " "applies to every reference; 0 removes the channel." ), ), InputParam( name="reference_subject_strength", type_hint=float | list[float], default=1.0, description=( "Per-image scale on the in-sequence image tokens (content/subject). A scalar applies to " "every reference; 0 removes the channel." ), ), InputParam( name="reference_masks", type_hint=PIL.Image.Image | list[PIL.Image.Image], description=( "Optional per-image masks restricting each reference to part of the picture: painted is " "the part to use. `reference_mask_mode` decides what happens to the rest. Index-aligned " "to `reference_images`; use `None` for slots without a mask." ), ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam(name="reference_mode", type_hint=str, description="The normalized reference mode."), OutputParam( name="vision_reference_images", type_hint=list, description="References sized for the Qwen3-VL vision path (empty when there are none).", ), OutputParam( name="vae_reference_images", type_hint=list, description="References for the VAE reference-latent channel (empty outside the latent modes).", ), OutputParam( name="reference_style_strengths", type_hint=list, description="One style scale per reference." ), OutputParam( name="reference_subject_strengths", type_hint=list, description="One subject scale per reference." ), OutputParam(name="reference_masks", type_hint=list, description="One mask (or None) per reference."), ] @torch.no_grad() def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) images = [image for image in _as_list(block_state.reference_images) if image is not None] block_state.reference_mode = self.reference_mode block_state.reference_style_strengths = _broadcast_per_image( block_state.reference_style_strength, len(images), 1.0 ) block_state.reference_subject_strengths = _broadcast_per_image( block_state.reference_subject_strength, len(images), 1.0 ) masks = _as_list(block_state.reference_masks) block_state.reference_masks = [masks[i] if i < len(masks) else None for i in range(len(images))] block_state.vision_reference_images = [self._vision_view(image, block_state) for image in images] block_state.vae_reference_images = [_to_rgb(image) for image in images] if self.uses_reference_latents else [] self.set_block_state(state, block_state) return components, state # auto_docstring class Krea2VisionReferenceImagesStep(Krea2ReferenceImagesStep): """Reference collector for the vision-only mode: a ~1 MP Qwen3-VL view, no VAE latent channel.""" reference_mode = REFERENCE_MODE_OFF uses_reference_latents = False @property def description(self) -> str: return "Collect the references and give each a ~1 MP Qwen3-VL view. Vision-only: no VAE latent channel." def _vision_view(self, image, block_state): return _cap_area(image, VISION_MAX_PIXELS) # auto_docstring class Krea2AppendReferenceImagesStep(Krea2ReferenceImagesStep): """Reference collector for the Ostris `append` mode: a coarse ~384px Qwen3-VL view plus the VAE channel.""" reference_mode = REFERENCE_MODE_APPEND uses_reference_latents = True @property def description(self) -> str: return ( "Collect the references with a coarse ~384px Qwen3-VL view; the detail rides the VAE " "reference-latent channel instead." ) def _vision_view(self, image, block_state): return _cap_area(image, VISION_EDIT_MAX_PIXELS) # auto_docstring class Krea2PrependReferenceImagesStep(Krea2ReferenceImagesStep): """ Reference collector for the Identity-Edit `prepend` mode: a `grounding_px` longest-side Qwen3-VL view plus the VAE channel. """ reference_mode = REFERENCE_MODE_PREPEND uses_reference_latents = True @property def description(self) -> str: return ( "Collect the references with a `grounding_px` longest-side Qwen3-VL view, plus the images for the " "VAE reference-latent channel." ) @property def inputs(self) -> list[InputParam]: return super().inputs + [ InputParam( name="grounding_px", type_hint=int, default=DEFAULT_GROUNDING_PX, description=( "Longest-side cap (px) on the Qwen3-VL view; the identity-vs-adherence dial. Lower is a " "stronger edit, higher is stronger identity, 0 means native resolution." ), ) ] def _vision_view(self, image, block_state): return _cap_longest_side(image, int(block_state.grounding_px or 0)) # auto_docstring class Krea2ReferenceTextEncoderStep(Krea2TextEncoderStep): """ Text encoder step that feeds the reference images through Qwen3-VL's vision tower alongside the prompt, keeping Krea 2's descriptor system template and 12-layer tap but inserting vision placeholders in the user turn. Falls back to the stock text-only encode when there are no references. """ model_name = "krea2" # Identity-Edit trains its unconditional as the same source image with an empty instruction, so only that # mode grounds the negative branch. grounds_negative = False def __init__(self): super().__init__() self._processor = None self._processor_tokenizer = None @property def description(self) -> str: return ( "Text encoder step that feeds the reference images through Qwen3-VL's vision tower alongside the " "prompt, scaling each reference's vision features by its per-image style/subject strengths and " "keeping only what a painted mask selects. Falls back to the stock text-only encode with no " "references." ) @property def expected_components(self) -> list[ComponentSpec]: return super().expected_components + [ ComponentSpec("processor", Qwen3VLProcessor, description="The Qwen3-VL vision processor.") ] @property def inputs(self) -> list[InputParam]: return super().inputs + [ InputParam( name="reference_mode", type_hint=str, default=REFERENCE_MODE_OFF, description="The normalized reference mode from the reference-images step.", ), InputParam( name="vision_reference_images", type_hint=list, description="References sized for the Qwen3-VL vision path.", ), InputParam(name="reference_style_strengths", type_hint=list, description="One style scale per reference."), InputParam( name="reference_subject_strengths", type_hint=list, description="One subject scale per reference." ), InputParam(name="reference_masks", type_hint=list, description="One mask (or None) per reference."), InputParam( name="reference_mask_mode", type_hint=str, default=REFERENCE_MASK_MODE_EXCLUDE_BLANK, description=( "What a painted mask does to the vision path: 'exclude_blank' (default, mask the vision " "attention and blank the masked-out tokens), 'exclude' (mask the attention only, leaving " "the masked-out tokens in the sequence) or 'deemphasize' (blank the tokens only, which " "attenuates the region rather than removing it)." ), ), ] def _get_processor(self, components): """The declared `processor` component, or one derived from the text encoder when the repo has none.""" if components.processor is not None: return components.processor tokenizer = components.tokenizer if self._processor is None or self._processor_tokenizer is not tokenizer: self._processor = _build_reference_processor(components.text_encoder, tokenizer) self._processor_tokenizer = tokenizer return self._processor def _image_prompt(self, images) -> str: """The vision-token layout for this mode's training template.""" raise NotImplementedError def _encode_prompt_with_vision( self, components, prompts, images, style_strengths, subject_strengths, masks, device, mask_mode ): """Encode `prompts` with `images` in the user turn, returning `(hidden_states, attention_mask)`. Unlike the text-only path this is variable length: the processor expands each `<|image_pad|>` to the image's token count and right-pads the batch. The vision-token layout comes from `_image_prompt`. """ processor = self._get_processor(components) text_encoder = components.text_encoder prefix_idx = _PROMPT_TEMPLATE_ENCODE_START_IDX image_prompt = self._image_prompt(images) # The stock prefix/suffix keep the leading block byte-identical to the text-only path, so the same # `prefix_idx` drop still holds; the vision tokens sit after it and survive. text = [ _PROMPT_TEMPLATE_ENCODE_PREFIX + image_prompt + (p or "") + _PROMPT_TEMPLATE_ENCODE_SUFFIX for p in prompts ] # Every prompt in the batch gets its own copy of the reference images, in placeholder order. model_inputs = processor(text=text, images=list(images) * len(prompts), padding=True, return_tensors="pt") model_inputs = model_inputs.to(device) # Skip the `get_image_features` wrap entirely when every strength is 1.0 and nothing is masked, so the # plain full-strength path stays byte-identical to an unwrapped encode. repeats = len(prompts) style = list(style_strengths) * repeats subject = list(subject_strengths) * repeats masks = list(masks) * repeats has_pixels = model_inputs.get("pixel_values") is not None needs_scaling = any(float(s) != 1.0 for s in style) or any(float(s) != 1.0 for s in subject) has_masks = any(m is not None for m in masks) # The mode selects which of the two independent layers a mask acts on. mask_attention = has_masks and mask_mode != REFERENCE_MASK_MODE_DEEMPHASIZE zero_masked_tokens = mask_mode != REFERENCE_MASK_MODE_EXCLUDE keep_vectors = [] if has_pixels and zero_masked_tokens and (needs_scaling or has_masks): keep_vectors = _build_keep_vectors( masks, model_inputs.get("image_grid_thw"), processor.image_processor.merge_size ) apply_scaling = has_pixels and (needs_scaling or any(k is not None for k in keep_vectors)) # Only the attention layer can keep unpainted content out of the reference; the output tokens the # wrapper zeroes are whole-image summaries, so zeroing them attenuates but never excludes. patch_keeps = [] if has_pixels and mask_attention: patch_keeps = build_patch_keeps( masks, model_inputs.get("image_grid_thw"), processor.image_processor.merge_size, _mask_to_keep_grid, ) original_get_image_features = None if apply_scaling: original_get_image_features = text_encoder.get_image_features text_encoder.get_image_features = _scaled_get_image_features( original_get_image_features, style, subject, processor.image_processor.merge_size, keep_vectors ) try: # No `position_ids`: with `pixel_values` present Qwen3-VL derives its mRoPE positions internally. with vision_key_masks(patch_keeps): outputs = text_encoder( input_ids=model_inputs["input_ids"], attention_mask=model_inputs["attention_mask"], pixel_values=model_inputs.get("pixel_values"), image_grid_thw=model_inputs.get("image_grid_thw"), mm_token_type_ids=model_inputs.get("mm_token_type_ids"), output_hidden_states=True, ) finally: if original_get_image_features is not None: text_encoder.get_image_features = original_get_image_features hidden_states = torch.stack([outputs.hidden_states[i] for i in KREA2_TEXT_ENCODER_SELECT_LAYERS], dim=2) attention_mask = model_inputs["attention_mask"].bool() return hidden_states[:, prefix_idx:], attention_mask[:, prefix_idx:] @torch.no_grad() def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device prompts = [block_state.prompt] if isinstance(block_state.prompt, str) else list(block_state.prompt) images = list(block_state.vision_reference_images or []) mask_mode = _coerce_reference_mask_mode(block_state.reference_mask_mode) if images: block_state.prompt_embeds, block_state.prompt_embeds_mask = self._encode_prompt_with_vision( components, prompts, images, block_state.reference_style_strengths or [], block_state.reference_subject_strengths or [], block_state.reference_masks or [], device, mask_mode, ) else: block_state.prompt_embeds, block_state.prompt_embeds_mask = self._encode_prompt( components, prompts, block_state.max_sequence_length, device ) block_state.negative_prompt_embeds = None block_state.negative_prompt_embeds_mask = None if components.requires_unconditional_embeds: negative_prompt = block_state.negative_prompt if negative_prompt is None: negative_prompt = "" if isinstance(negative_prompt, str): negative_prompt = [negative_prompt] * len(prompts) if self.grounds_negative and images: # The negative MUST get the same masks as the positive: an unmasked negative sees the whole # reference, so (cond - uncond) pushes the unpainted region away instead of ignoring it. # Strengths stay at full -- that is the training condition. block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = ( self._encode_prompt_with_vision( components, negative_prompt, images, [], [], block_state.reference_masks or [], device, mask_mode, ) ) else: block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = self._encode_prompt( components, negative_prompt, block_state.max_sequence_length, device ) self.set_block_state(state, block_state) return components, state # auto_docstring class Krea2VisionReferenceTextEncoderStep(Krea2ReferenceTextEncoderStep): """ Text encoder for the vision-only mode. A single reference goes in as a bare vision block; several are labelled "Picture N:". The negative branch stays text-only. """ grounds_negative = False @property def description(self) -> str: return "Text encoder for the vision-only mode: references enter the prompt through Qwen3-VL's vision tower." def _image_prompt(self, images) -> str: if len(images) > 1: return "".join(f"Picture {i + 1}: <|vision_start|><|image_pad|><|vision_end|>" for i in range(len(images))) return "<|vision_start|><|image_pad|><|vision_end|>" # auto_docstring class Krea2AppendReferenceTextEncoderStep(Krea2ReferenceTextEncoderStep): """ Text encoder for the Ostris `append` mode, which labels every reference "Picture N:" even when there is only one. The negative branch stays text-only. """ grounds_negative = False @property def description(self) -> str: return 'Text encoder for the Ostris append mode: every reference is labelled "Picture N:".' def _image_prompt(self, images) -> str: return "".join(f"Picture {i + 1}: <|vision_start|><|image_pad|><|vision_end|>" for i in range(len(images))) # auto_docstring class Krea2PrependReferenceTextEncoderStep(Krea2ReferenceTextEncoderStep): """Text encoder for the Identity-Edit `prepend` mode: unlabelled vision blocks, grounded negative branch.""" grounds_negative = True @property def description(self) -> str: return ( "Text encoder for the Identity-Edit prepend mode: bare vision blocks and a negative branch grounded " "on the same source, as the LoRA was trained." ) def _image_prompt(self, images) -> str: return "".join("<|vision_start|><|image_pad|><|vision_end|>" for _ in images) # auto_docstring class Krea2ReferenceTextInputsStep(ModularPipelineBlocks): """ Input step that determines `batch_size`/`dtype` from the per-prompt `prompt_embeds` and replicates the text conditioning (and the negative branch) to `batch_size * num_images_per_prompt`. Unlike the stock step it lets the two branches carry different sequence lengths. """ model_name = "krea2" @property def description(self) -> str: return ( "Input step that determines `batch_size`/`dtype` and batch-expands the text conditioning, allowing " "the positive and negative branches to carry different sequence lengths (the positive gains vision " "tokens the text-only negative does not have)." ) @property def inputs(self) -> list[InputParam]: return [ InputParam.template("num_images_per_prompt", default=1), InputParam.template("prompt_embeds"), InputParam.template("prompt_embeds_mask"), InputParam.template("negative_prompt_embeds"), InputParam.template("negative_prompt_embeds_mask"), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( name="batch_size", type_hint=int, description="Effective batch size (num prompts * num_images_per_prompt).", ), OutputParam(name="dtype", type_hint=torch.dtype, description="The dtype of the text features."), OutputParam(name="prompt_embeds", type_hint=torch.Tensor, description="Text features, batch-expanded."), OutputParam(name="prompt_embeds_mask", type_hint=torch.Tensor, description="Text mask, batch-expanded."), OutputParam( name="negative_prompt_embeds", type_hint=torch.Tensor, description="Negative text features, batch-expanded.", ), OutputParam( name="negative_prompt_embeds_mask", type_hint=torch.Tensor, description="Negative text mask, batch-expanded.", ), ] @staticmethod def _expand(embeds, mask, num_images_per_prompt): prompt_batch, seq_len, num_layers, dim = embeds.shape n = num_images_per_prompt embeds = embeds.repeat(1, n, 1, 1).view(prompt_batch * n, seq_len, num_layers, dim) mask = mask.repeat(1, n).view(prompt_batch * n, seq_len) return embeds, mask @torch.no_grad() def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) n = block_state.num_images_per_prompt block_state.dtype = block_state.prompt_embeds.dtype block_state.batch_size = block_state.prompt_embeds.shape[0] * n block_state.prompt_embeds, block_state.prompt_embeds_mask = self._expand( block_state.prompt_embeds, block_state.prompt_embeds_mask, n ) if block_state.negative_prompt_embeds is not None: block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = self._expand( block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask, n ) self.set_block_state(state, block_state) return components, state # auto_docstring class Krea2ReferenceLatentsStep(ModularPipelineBlocks): """ VAE-encode the reference images into clean latent tokens for the `append`/`prepend` conditioning paths. Each reference is encoded at up to ~1 MP with its own aspect ratio, normalized with the VAE's per-channel statistics, patch-packed exactly like the noise latents and given Kontext-style rotary coordinates: the i-th reference sits on rotary frame axis `i + 1` with its own y/x grid. Reference masks reach this channel only when `mask_reference_latents` is set. A no-op in `off` mode or with no references. """ model_name = "krea2" @property def description(self) -> str: return ( "VAE-encode the reference images into clean, patch-packed latent tokens with frame-axis rotary " "coordinates for the `append`/`prepend` conditioning paths. A no-op in `off` mode." ) @property def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("vae", AutoencoderKLQwenImage)] @property def inputs(self) -> list[InputParam]: return [ InputParam( name="reference_mode", type_hint=str, default=REFERENCE_MODE_OFF, description="The normalized reference mode from the reference-images step.", ), InputParam( name="vae_reference_images", type_hint=list, description="References for the VAE reference-latent channel.", ), InputParam(name="reference_masks", type_hint=list, description="One mask (or None) per reference."), InputParam( name="mask_reference_latents", type_hint=bool, default=False, description=( "Apply the reference masks to this channel too, on top of the vision tokens they always " "mask. Off (default) is vision-only masking, so the edit LoRA still sees the whole " "reference at full detail -- which is why a masked edit can show hints of the unpainted " "parts. On drops the unpainted latent patch tokens from the sequence entirely." ), ), InputParam.template("dtype"), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( name="reference_latents", type_hint=torch.Tensor, description="Clean packed reference tokens (1, reference_seq_len, in_channels), or None.", ), OutputParam( name="reference_position_ids", type_hint=torch.Tensor, description="Rotary coordinates for the reference tokens (reference_seq_len, 3), or None.", ), OutputParam( name="reference_seq_len", type_hint=int, description="Number of reference tokens (0 when unused)." ), ] @staticmethod def _to_snapped_tensor(image, snap: int) -> torch.Tensor: """Reference image -> normalized `(1, 3, 1, H, W)` tensor for the temporal VAE. Aspect-preserving downscale to fit `REFERENCE_LATENTS_MAX_PIXELS` (never upscaled), then snap each side to a multiple of `snap` so the latent grid is patchifiable. """ image = _to_rgb(image) width, height = image.width, image.height if height * width > REFERENCE_LATENTS_MAX_PIXELS: ratio = height / width new_height = math.sqrt(REFERENCE_LATENTS_MAX_PIXELS * ratio) new_width = math.sqrt(REFERENCE_LATENTS_MAX_PIXELS / ratio) else: new_height, new_width = float(height), float(width) new_height = max(snap, int(round(new_height / snap)) * snap) new_width = max(snap, int(round(new_width / snap)) * snap) array = np.asarray(image, dtype=np.float32) / 255.0 tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0) * 2.0 - 1.0 if (new_height, new_width) != (height, width): tensor = F.interpolate(tensor, size=(new_height, new_width), mode="bilinear", align_corners=False) return tensor.unsqueeze(2) @torch.no_grad() def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) images = list(block_state.vae_reference_images or []) if not images: raise ValueError( "The append and prepend modes need at least one reference image; pass `reference_images` or " "switch to the vision-only mode." ) block_state.reference_latents = None block_state.reference_position_ids = None block_state.reference_seq_len = 0 vae = components.vae patch_size = components.patch_size snap = components.vae_scale_factor * patch_size z_dim = vae.config.z_dim latents_mean = torch.tensor(vae.config.latents_mean).view(1, z_dim, 1, 1, 1) latents_std = torch.tensor(vae.config.latents_std).view(1, z_dim, 1, 1, 1) masks = list(block_state.reference_masks or []) mask_latents = bool(block_state.mask_reference_latents) tokens = [] position_ids = [] for i, image in enumerate(images): pixels = self._to_snapped_tensor(image, snap).to(vae.device, vae.dtype) raw = vae.encode(pixels).latent_dist.mode() mean = latents_mean.to(raw.device, raw.dtype) std = latents_std.to(raw.device, raw.dtype) latents = ((raw - mean) / std)[:, :, 0] # drop the temporal axis -> (1, C, lat_h, lat_w) image_tokens = _pack_latents(latents, patch_size).to(block_state.dtype) _, _, latent_height, latent_width = latents.shape grid_height, grid_width = latent_height // patch_size, latent_width // patch_size image_position_ids = _reference_position_ids(i, grid_height, grid_width) if mask_latents: keep = _mask_to_keep_vector(masks[i] if i < len(masks) else None, grid_height, grid_width) if keep is not None: # Dropped rather than zeroed: a zero token is the mean latent (flat mid-grey), which the # model can still attend to and reproduce. Position ids go with them, so the kept tokens # keep their true coordinates. index = torch.nonzero(keep > 0.5, as_tuple=False).squeeze(-1) if index.numel() == 0: continue image_tokens = image_tokens[:, index] image_position_ids = image_position_ids[index] tokens.append(image_tokens) position_ids.append(image_position_ids) block_state.reference_latents = torch.cat(tokens, dim=1) block_state.reference_position_ids = torch.cat(position_ids, dim=0) block_state.reference_seq_len = int(block_state.reference_latents.shape[1]) self.set_block_state(state, block_state) return components, state # auto_docstring class Krea2ReferencePreparePositionIdsStep(ModularPipelineBlocks): """ Build the rotary position ids for both guidance branches, which carry different text lengths once the positive gains vision tokens. Reference-latent coordinates are spliced in before the target for `prepend` and after it for `append`. """ model_name = "krea2" @property def description(self) -> str: return ( "Build the rotary position ids for both guidance branches (the branches carry different text " "lengths once the positive gains vision tokens), splicing in the reference-latent coordinates " "before the target for `prepend` and after it for `append`." ) @property def inputs(self) -> list[InputParam]: return [ InputParam.template("height", default=1024), InputParam.template("width", default=1024), InputParam.template("prompt_embeds"), InputParam.template("negative_prompt_embeds"), InputParam( name="reference_mode", type_hint=str, default=REFERENCE_MODE_OFF, description="The normalized reference mode from the reference-images step.", ), InputParam( name="reference_position_ids", type_hint=torch.Tensor, description="Rotary coordinates for the reference tokens, when the latent modes are active.", ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( name="position_ids", type_hint=torch.Tensor, description="Rotary coordinates for the conditional branch.", ), OutputParam( name="negative_position_ids", type_hint=torch.Tensor, description="Rotary coordinates for the unconditional branch.", ), ] @torch.no_grad() def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device patch_size = components.patch_size grid_h = block_state.height // (components.vae_scale_factor * patch_size) grid_w = block_state.width // (components.vae_scale_factor * patch_size) reference_ids = block_state.reference_position_ids if reference_ids is not None: reference_ids = reference_ids.to(device) def build(text_seq_len: int) -> torch.Tensor: ids = Krea2PreparePositionIdsStep.prepare_position_ids(text_seq_len, grid_h, grid_w, device) if reference_ids is None or self.reference_mode == REFERENCE_MODE_OFF: return ids if self.reference_mode == REFERENCE_MODE_PREPEND: # [text | reference | target] return torch.cat([ids[:text_seq_len], reference_ids, ids[text_seq_len:]], dim=0) # [text | target | reference] return torch.cat([ids, reference_ids], dim=0) text_seq_len = block_state.prompt_embeds.shape[1] block_state.position_ids = build(text_seq_len) negative_prompt_embeds = block_state.negative_prompt_embeds if negative_prompt_embeds is None or negative_prompt_embeds.shape[1] == text_seq_len: block_state.negative_position_ids = block_state.position_ids else: block_state.negative_position_ids = build(negative_prompt_embeds.shape[1]) self.set_block_state(state, block_state) return components, state class Krea2ReferenceLoopDenoiser(ModularPipelineBlocks): model_name = "krea2" @property def description(self) -> str: return ( "Within the denoising loop: concatenate the clean reference tokens onto the noisy latents in the " "order the mode dictates, run the `transformer` per guidance branch with that branch's own " "position ids, and slice the reference rows back off the prediction. Compose into " "`Krea2ReferenceDenoiseStep`." ) @property def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "guider", ClassifierFreeGuidance, # Krea 2 uses cond-anchored CFG (`cond + scale * (cond - uncond)`), which is the # `use_original_formulation` branch. config=FrozenDict({"guidance_scale": 4.5, "use_original_formulation": True}), default_creation_method="from_config", ), ComponentSpec("transformer", Krea2Transformer2DModel), ] @property def inputs(self) -> list[InputParam]: return [ InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Packed image latents."), InputParam.template("num_inference_steps", required=True), InputParam.template("prompt_embeds"), InputParam.template("prompt_embeds_mask"), InputParam.template("negative_prompt_embeds"), InputParam.template("negative_prompt_embeds_mask"), InputParam( name="position_ids", required=True, type_hint=torch.Tensor, description="Rotary coordinates for the conditional branch.", ), InputParam( name="negative_position_ids", type_hint=torch.Tensor, description="Rotary coordinates for the unconditional branch.", ), InputParam( name="reference_mode", type_hint=str, default=REFERENCE_MODE_OFF, description="The normalized reference mode from the reference-images step.", ), InputParam( name="reference_latents", type_hint=torch.Tensor, description="Clean packed reference tokens, when the latent modes are active.", ), InputParam.template("attention_kwargs"), ] @torch.no_grad() def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): transformer = components.transformer latents = block_state.latents.to(transformer.dtype) timestep = block_state.timestep.to(transformer.dtype) target_seq_len = latents.shape[1] latents, reference_seq_len = _concat_reference_latents( latents, block_state.reference_latents, self.reference_mode, transformer.dtype ) negative_position_ids = block_state.negative_position_ids if negative_position_ids is None: negative_position_ids = block_state.position_ids guider_inputs = { "encoder_hidden_states": ( block_state.prompt_embeds.to(transformer.dtype), block_state.negative_prompt_embeds.to(transformer.dtype) if block_state.negative_prompt_embeds is not None else None, ), "encoder_attention_mask": ( block_state.prompt_embeds_mask, block_state.negative_prompt_embeds_mask, ), # A tuple is indexed per guidance branch, so each pass gets coordinates for its own text length. "position_ids": (block_state.position_ids, negative_position_ids), } components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) guider_state = components.guider.prepare_inputs(guider_inputs) for guider_state_batch in guider_state: components.guider.prepare_models(transformer) cond_kwargs = {name: getattr(guider_state_batch, name) for name in guider_inputs} handles = [] if reference_seq_len and self.reference_mode == REFERENCE_MODE_APPEND: # The split index is branch-dependent: the conditional text block carries the vision tokens. text_seq_len = cond_kwargs["encoder_hidden_states"].shape[1] handles = _register_zero_time_reference_hooks( transformer, timestep, batch_size=latents.shape[0], split=text_seq_len + target_seq_len, total=text_seq_len + target_seq_len + reference_seq_len, ) try: noise_pred = transformer( hidden_states=latents, timestep=timestep, attention_kwargs=block_state.attention_kwargs, return_dict=False, **cond_kwargs, )[0] finally: for handle in handles: handle.remove() guider_state_batch.noise_pred = _slice_reference_rows( noise_pred, self.reference_mode, reference_seq_len, target_seq_len ) components.guider.cleanup_models(transformer) block_state.noise_pred = components.guider(guider_state).pred return components, block_state # auto_docstring class Krea2ReferenceDenoiseStep(Krea2DenoiseLoopWrapper): """ Denoising loop that iteratively denoises the packed image latents over `timesteps` with reference conditioning: the clean reference tokens are concatenated in the order the mode dictates, each guidance branch runs with its own position ids, and the reference rows are sliced back off the prediction before the scheduler step. """ model_name = "krea2" block_classes = [Krea2LoopBeforeDenoiser, Krea2ReferenceLoopDenoiser, Krea2LoopAfterDenoiser] block_names = ["before_denoiser", "denoiser", "after_denoiser"] @property def description(self) -> str: return ( "Denoising loop with reference conditioning: concatenates the clean reference tokens, runs each " "guidance branch with its own position ids, and slices the reference rows back off the prediction. " "Identical to the stock loop when no reference latents are active." ) # ===================================================================================================== # Per-mode leaf blocks # ===================================================================================================== # auto_docstring class Krea2VisionReferencePositionIdsStep(Krea2ReferencePreparePositionIdsStep): """Rotary coordinates for the vision-only mode: per-branch, with no reference rows to splice in.""" reference_mode = REFERENCE_MODE_OFF @property def description(self) -> str: return "Per-branch rotary coordinates for the [text | image] sequence (vision-only: no reference rows)." # auto_docstring class Krea2AppendReferencePositionIdsStep(Krea2ReferencePreparePositionIdsStep): """Rotary coordinates for the Ostris `append` mode: the reference rows go after the target.""" reference_mode = REFERENCE_MODE_APPEND @property def description(self) -> str: return "Per-branch rotary coordinates with the reference rows appended after the target." # auto_docstring class Krea2PrependReferencePositionIdsStep(Krea2ReferencePreparePositionIdsStep): """Rotary coordinates for the Identity-Edit `prepend` mode: the reference rows go before the target.""" reference_mode = REFERENCE_MODE_PREPEND @property def description(self) -> str: return "Per-branch rotary coordinates with the reference rows spliced in before the target." class Krea2VisionReferenceLoopDenoiser(Krea2ReferenceLoopDenoiser): reference_mode = REFERENCE_MODE_OFF @property def description(self) -> str: return ( "Within the denoising loop: run the `transformer` per guidance branch with that branch's own " "position ids. Vision-only, so there are no reference tokens in the sequence." ) class Krea2AppendReferenceLoopDenoiser(Krea2ReferenceLoopDenoiser): reference_mode = REFERENCE_MODE_APPEND @property def description(self) -> str: return ( "Within the denoising loop: append the clean reference tokens after the noisy target, modulate them " "at flow time t=0, and slice them back off the prediction." ) class Krea2PrependReferenceLoopDenoiser(Krea2ReferenceLoopDenoiser): reference_mode = REFERENCE_MODE_PREPEND @property def description(self) -> str: return ( "Within the denoising loop: prepend the clean reference tokens before the noisy target under the " "plain uniform modulation, and slice them back off the prediction." ) # auto_docstring class Krea2VisionReferenceDenoiseStep(Krea2DenoiseLoopWrapper): """Denoising loop for the vision-only mode, with symmetric CFG.""" model_name = "krea2" block_classes = [Krea2LoopBeforeDenoiser, Krea2VisionReferenceLoopDenoiser, Krea2LoopAfterDenoiser] block_names = ["before_denoiser", "denoiser", "after_denoiser"] @property def description(self) -> str: return "Denoising loop for the vision-only mode, with per-branch position ids." # auto_docstring class Krea2AppendReferenceDenoiseStep(Krea2DenoiseLoopWrapper): """Denoising loop for the Ostris `append` mode: reference tokens at the tail, modulated at t=0.""" model_name = "krea2" block_classes = [Krea2LoopBeforeDenoiser, Krea2AppendReferenceLoopDenoiser, Krea2LoopAfterDenoiser] block_names = ["before_denoiser", "denoiser", "after_denoiser"] @property def description(self) -> str: return "Denoising loop for the Ostris append mode: clean reference tokens at the tail, modulated at t=0." # auto_docstring class Krea2PrependReferenceDenoiseStep(Krea2DenoiseLoopWrapper): """Denoising loop for the Identity-Edit `prepend` mode: the clean source sits before the target.""" model_name = "krea2" block_classes = [Krea2LoopBeforeDenoiser, Krea2PrependReferenceLoopDenoiser, Krea2LoopAfterDenoiser] block_names = ["before_denoiser", "denoiser", "after_denoiser"] @property def description(self) -> str: return "Denoising loop for the Identity-Edit prepend mode: the clean source before the noisy target." class _Krea2TurboTextEncoderMixin: """Drops the negative branch from a reference text encoder for the distilled checkpoint.""" @property def expected_components(self) -> list[ComponentSpec]: return [spec for spec in super().expected_components if spec.name != "guider"] @property def inputs(self) -> list[InputParam]: return [param for param in super().inputs if param.name != "negative_prompt"] @property def intermediate_outputs(self) -> list[OutputParam]: return [ param for param in super().intermediate_outputs if param.name not in ("negative_prompt_embeds", "negative_prompt_embeds_mask") ] @torch.no_grad() def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device prompts = [block_state.prompt] if isinstance(block_state.prompt, str) else list(block_state.prompt) images = list(block_state.vision_reference_images or []) mask_mode = _coerce_reference_mask_mode(block_state.reference_mask_mode) if images: block_state.prompt_embeds, block_state.prompt_embeds_mask = self._encode_prompt_with_vision( components, prompts, images, block_state.reference_style_strengths or [], block_state.reference_subject_strengths or [], block_state.reference_masks or [], device, mask_mode, ) else: block_state.prompt_embeds, block_state.prompt_embeds_mask = self._encode_prompt( components, prompts, block_state.max_sequence_length, device ) self.set_block_state(state, block_state) return components, state # auto_docstring class Krea2TurboVisionReferenceTextEncoderStep(_Krea2TurboTextEncoderMixin, Krea2VisionReferenceTextEncoderStep): """Vision-only text encoder for the distilled checkpoint: no negative branch, no guider.""" # auto_docstring class Krea2TurboAppendReferenceTextEncoderStep(_Krea2TurboTextEncoderMixin, Krea2AppendReferenceTextEncoderStep): """Ostris append text encoder for the distilled checkpoint: no negative branch, no guider.""" # auto_docstring class Krea2TurboPrependReferenceTextEncoderStep(_Krea2TurboTextEncoderMixin, Krea2PrependReferenceTextEncoderStep): """Identity-Edit prepend text encoder for the distilled checkpoint: no negative branch, no guider.""" class Krea2TurboReferenceLoopDenoiser(ModularPipelineBlocks): model_name = "krea2" @property def description(self) -> str: return ( "Within the denoising loop: run the `transformer` on the conditional text features with the clean " "reference tokens spliced in, then slice them back off. The distilled checkpoint runs without " "classifier-free guidance, so there is no negative branch or guider." ) @property def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("transformer", Krea2Transformer2DModel)] @property def inputs(self) -> list[InputParam]: return [ InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Packed image latents."), InputParam.template("prompt_embeds"), InputParam.template("prompt_embeds_mask"), InputParam( name="position_ids", required=True, type_hint=torch.Tensor, description="Rotary coordinates for the [text | image] sequence.", ), InputParam( name="reference_mode", type_hint=str, default=REFERENCE_MODE_OFF, description="The normalized reference mode from the reference-images step.", ), InputParam( name="reference_latents", type_hint=torch.Tensor, description="Clean packed reference tokens, when the latent modes are active.", ), InputParam.template("attention_kwargs"), ] @torch.no_grad() def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): transformer = components.transformer latents = block_state.latents.to(transformer.dtype) timestep = block_state.timestep.to(transformer.dtype) target_seq_len = latents.shape[1] latents, reference_seq_len = _concat_reference_latents( latents, block_state.reference_latents, self.reference_mode, transformer.dtype ) handles = [] if reference_seq_len and self.reference_mode == REFERENCE_MODE_APPEND: text_seq_len = block_state.prompt_embeds.shape[1] handles = _register_zero_time_reference_hooks( transformer, timestep, batch_size=latents.shape[0], split=text_seq_len + target_seq_len, total=text_seq_len + target_seq_len + reference_seq_len, ) try: noise_pred = transformer( hidden_states=latents, timestep=timestep, position_ids=block_state.position_ids, attention_kwargs=block_state.attention_kwargs, encoder_hidden_states=block_state.prompt_embeds.to(transformer.dtype), encoder_attention_mask=block_state.prompt_embeds_mask, return_dict=False, )[0] finally: for handle in handles: handle.remove() block_state.noise_pred = _slice_reference_rows( noise_pred, self.reference_mode, reference_seq_len, target_seq_len ) return components, block_state # auto_docstring class Krea2TurboReferenceDenoiseStep(Krea2DenoiseLoopWrapper): """ Denoising loop for the distilled Krea 2 Turbo checkpoint with reference conditioning: the clean reference tokens are spliced into the sequence and sliced back off the prediction. No classifier-free guidance. """ model_name = "krea2" block_classes = [Krea2LoopBeforeDenoiser, Krea2TurboReferenceLoopDenoiser, Krea2LoopAfterDenoiser] block_names = ["before_denoiser", "denoiser", "after_denoiser"] @property def description(self) -> str: return ( "Denoising loop for the distilled Krea 2 Turbo checkpoint with reference conditioning, without " "classifier-free guidance." ) class Krea2TurboVisionReferenceLoopDenoiser(Krea2TurboReferenceLoopDenoiser): reference_mode = REFERENCE_MODE_OFF class Krea2TurboAppendReferenceLoopDenoiser(Krea2TurboReferenceLoopDenoiser): reference_mode = REFERENCE_MODE_APPEND class Krea2TurboPrependReferenceLoopDenoiser(Krea2TurboReferenceLoopDenoiser): reference_mode = REFERENCE_MODE_PREPEND def _turbo_denoise_step(name, denoiser, summary): """Build a guidance-free denoise loop for one mode.""" return type( name, (Krea2DenoiseLoopWrapper,), { "__doc__": summary, "model_name": "krea2", "block_classes": [Krea2LoopBeforeDenoiser, denoiser, Krea2LoopAfterDenoiser], "block_names": ["before_denoiser", "denoiser", "after_denoiser"], "description": property(lambda self, _s=summary: _s), }, ) Krea2TurboVisionReferenceDenoiseStep = _turbo_denoise_step( "Krea2TurboVisionReferenceDenoiseStep", Krea2TurboVisionReferenceLoopDenoiser, "Guidance-free denoising loop for the vision-only mode on the distilled checkpoint.", ) Krea2TurboAppendReferenceDenoiseStep = _turbo_denoise_step( "Krea2TurboAppendReferenceDenoiseStep", Krea2TurboAppendReferenceLoopDenoiser, "Guidance-free denoising loop for the Ostris append mode on the distilled checkpoint.", ) Krea2TurboPrependReferenceDenoiseStep = _turbo_denoise_step( "Krea2TurboPrependReferenceDenoiseStep", Krea2TurboPrependReferenceLoopDenoiser, "Guidance-free denoising loop for the Identity-Edit prepend mode on the distilled checkpoint.", ) # ===================================================================================================== # Workflows: one flat blockset per mode, per checkpoint. # ===================================================================================================== def _workflow(name, images, encoder, timesteps, position_ids, denoise, latents, summary): core = InsertableDict( [("input", Krea2ReferenceTextInputsStep())] + ([("reference_latents", Krea2ReferenceLatentsStep())] if latents else []) + [ ("prepare_latents", Krea2PrepareLatentsStep()), ("set_timesteps", timesteps()), ("prepare_position_ids", position_ids()), ("denoise", denoise()), ] ) core_cls = type( f"{name}CoreDenoiseStep", (SequentialPipelineBlocks,), { "__doc__": f"Core denoising workflow: {summary}", "model_name": "krea2", "block_classes": list(core.values()), "block_names": list(core.keys()), "description": property(lambda self, _s=summary: f"Core denoising workflow: {_s}"), "outputs": property( lambda self: [ OutputParam.template( "latents", description="The denoised packed latents (B, image_seq_len, in_channels)." ) ] ), }, ) return type( name, (SequentialPipelineBlocks,), { "__doc__": summary, "model_name": "krea2", "block_classes": [images, encoder, core_cls, Krea2DecodeStep], "block_names": ["reference_images", "text_encoder", "denoise", "decode"], "description": property(lambda self, _s=summary: _s), "outputs": property(lambda self: [OutputParam.template("images")]), }, ) Krea2VisionReferenceWorkflow = _workflow( "Krea2VisionReferenceWorkflow", Krea2VisionReferenceImagesStep, Krea2VisionReferenceTextEncoderStep, Krea2SetTimestepsStep, Krea2VisionReferencePositionIdsStep, Krea2VisionReferenceDenoiseStep, latents=False, summary="vision-only references on the stock checkpoint, no LoRA needed.", ) Krea2AppendReferenceWorkflow = _workflow( "Krea2AppendReferenceWorkflow", Krea2AppendReferenceImagesStep, Krea2AppendReferenceTextEncoderStep, Krea2SetTimestepsStep, Krea2AppendReferencePositionIdsStep, Krea2AppendReferenceDenoiseStep, latents=True, summary="vision plus clean VAE reference tokens after the target, for the Ostris edit LoRA.", ) Krea2PrependReferenceWorkflow = _workflow( "Krea2PrependReferenceWorkflow", Krea2PrependReferenceImagesStep, Krea2PrependReferenceTextEncoderStep, Krea2SetTimestepsStep, Krea2PrependReferencePositionIdsStep, Krea2PrependReferenceDenoiseStep, latents=True, summary="vision plus the clean VAE source before the target, for the Identity-Edit LoRA.", ) Krea2TurboVisionReferenceWorkflow = _workflow( "Krea2TurboVisionReferenceWorkflow", Krea2VisionReferenceImagesStep, Krea2TurboVisionReferenceTextEncoderStep, Krea2TurboSetTimestepsStep, Krea2VisionReferencePositionIdsStep, Krea2TurboVisionReferenceDenoiseStep, latents=False, summary="vision-only references on the distilled checkpoint.", ) Krea2TurboAppendReferenceWorkflow = _workflow( "Krea2TurboAppendReferenceWorkflow", Krea2AppendReferenceImagesStep, Krea2TurboAppendReferenceTextEncoderStep, Krea2TurboSetTimestepsStep, Krea2AppendReferencePositionIdsStep, Krea2TurboAppendReferenceDenoiseStep, latents=True, summary="vision plus clean VAE reference tokens after the target, for the Ostris Turbo edit LoRA.", ) Krea2TurboPrependReferenceWorkflow = _workflow( "Krea2TurboPrependReferenceWorkflow", Krea2PrependReferenceImagesStep, Krea2TurboPrependReferenceTextEncoderStep, Krea2TurboSetTimestepsStep, Krea2PrependReferencePositionIdsStep, Krea2TurboPrependReferenceDenoiseStep, latents=True, summary="vision plus the clean VAE source before the target, on the distilled checkpoint.", ) # auto_docstring class Krea2ReferenceAutoBlocks(ConditionalPipelineBlocks): """ Modular pipeline for Krea 2 with reference-image conditioning. `reference_mode` picks the workflow: `"off"` (vision path only, stock checkpoint), `"append"` (Ostris edit LoRA) or `"prepend"` (Identity-Edit LoRA). With no `reference_images` this is plain text-to-image. """ model_name = "krea2" block_classes = [ Krea2VisionReferenceWorkflow, Krea2AppendReferenceWorkflow, Krea2PrependReferenceWorkflow, ] block_names = [REFERENCE_MODE_OFF, REFERENCE_MODE_APPEND, REFERENCE_MODE_PREPEND] block_trigger_inputs = ["reference_mode"] default_block_name = REFERENCE_MODE_OFF @property def description(self) -> str: return ( "Krea 2 with reference-image conditioning: `reference_mode` selects the vision-only path, the " "Ostris append path or the Identity-Edit prepend path." ) def select_block(self, **kwargs) -> str: return _coerce_reference_mode(kwargs.get("reference_mode")) # auto_docstring class Krea2TurboReferenceAutoBlocks(Krea2ReferenceAutoBlocks): """ Modular pipeline for the distilled Krea 2 Turbo checkpoint with reference-image conditioning. The same three modes on the distilled schedule and without classifier-free guidance, so it takes no negative prompt and carries no guider. """ model_name = "krea2" block_classes = [ Krea2TurboVisionReferenceWorkflow, Krea2TurboAppendReferenceWorkflow, Krea2TurboPrependReferenceWorkflow, ] block_names = [REFERENCE_MODE_OFF, REFERENCE_MODE_APPEND, REFERENCE_MODE_PREPEND] block_trigger_inputs = ["reference_mode"] default_block_name = REFERENCE_MODE_OFF @property def description(self) -> str: return ( "Krea 2 Turbo with reference-image conditioning: the distilled schedule, no CFG, and the same " "three reference modes." )