"""Krea 2 single-stream diffusion model.
The architecture is adapted from Krea's Apache-2.0 reference implementation
in ``~/programs/krea-2/mmdit.py``. It is modified to use local checkpoint
loading, strict state-key accounting, and the local scaled-FP8 layer.
"""
from __future__ import annotations
from dataclasses import dataclass
import math
import re
from typing import Iterable
import torch
from torch import nn
from torch.nn import functional as F
from diffusion_cli.checkpoint import loadStateDict
from diffusion_cli.config import ModelSource
from diffusion_cli.errors import DiffusionCliError
from diffusion_cli.quantization import (
QuantizationManifest,
ScaledFp8Linear,
manifestFromStateDict,
loadQuantizationManifest,
moveModulePreservingQuantization,
validateQuantizedStateDict,
)
KREA2_FEATURES = 6144
KREA2_TIMESTEP_WIDTH = 256
KREA2_TEXT_WIDTH = 2560
KREA2_HEADS = 48
KREA2_KV_HEADS = 12
KREA2_BLOCKS = 28
KREA2_MLP_MULTIPLIER = 4
KREA2_PATCH_SIZE = 2
KREA2_LATENT_CHANNELS = 16
KREA2_TEXT_LAYERS = 12
KREA2_TEXT_HEADS = 20
KREA2_TEXT_KV_HEADS = 20
KREA2_ROPE_THETA = 1000.0
@dataclass(frozen=True)
class Krea2Conditioning:
"""Twelve Qwen hidden-state taps and their token mask."""
hidden_states: torch.Tensor
attention_mask: torch.Tensor
@dataclass(frozen=True)
class Krea2Config:
"""Architecture settings for one Krea 2 diffusion model."""
features: int = KREA2_FEATURES
timestep_width: int = KREA2_TIMESTEP_WIDTH
text_width: int = KREA2_TEXT_WIDTH
heads: int = KREA2_HEADS
kv_heads: int = KREA2_KV_HEADS
blocks: int = KREA2_BLOCKS
mlp_multiplier: int = KREA2_MLP_MULTIPLIER
patch_size: int = KREA2_PATCH_SIZE
latent_channels: int = KREA2_LATENT_CHANNELS
text_layers: int = KREA2_TEXT_LAYERS
text_heads: int = KREA2_TEXT_HEADS
text_kv_heads: int = KREA2_TEXT_KV_HEADS
rope_theta: float = KREA2_ROPE_THETA
bias: bool = False
def _finite(value: torch.Tensor, label: str) -> None:
"""Reject non-finite tensors at a model boundary."""
if not torch.isfinite(value).all():
raise DiffusionCliError(f"Krea diffusion {label} contains NaN or Inf")
class KreaRmsNorm(nn.Module):
"""RMS normalization with the reference implementation's additive scale."""
def __init__(self, features: int, *, device=None, dtype=torch.float32):
super().__init__()
self.features = features
self.eps = 1e-5
self.scale = nn.Parameter(
torch.zeros(features, device=device, dtype=dtype),
requires_grad=False,
)
def forward(self, value: torch.Tensor) -> torch.Tensor:
"""Normalize the last feature dimension."""
original_dtype = value.dtype
normalized = F.rms_norm(
value.float(),
(self.features,),
eps=self.eps,
weight=self.scale.float() + 1.0,
)
return normalized.to(original_dtype)
class KreaAttention(nn.Module):
"""Grouped-query attention used by text fusion and image blocks."""
def __init__(
self,
features: int,
heads: int,
kv_heads: int,
*,
bias: bool,
device=None,
dtype=None,
) -> None:
super().__init__()
if features % heads != 0 or heads % kv_heads != 0:
raise ValueError("Krea attention heads must divide evenly")
self.heads = heads
self.kv_heads = kv_heads
self.head_dim = features // heads
self.wq = nn.Linear(
features,
features,
bias=bias,
device=device,
dtype=dtype,
)
self.wk = nn.Linear(
features,
self.head_dim * kv_heads,
bias=bias,
device=device,
dtype=dtype,
)
self.wv = nn.Linear(
features,
self.head_dim * kv_heads,
bias=bias,
device=device,
dtype=dtype,
)
self.gate = nn.Linear(
features,
features,
bias=bias,
device=device,
dtype=dtype,
)
self.qknorm = KreaQkNorm(
self.head_dim,
device=device,
)
self.wo = nn.Linear(
features,
features,
bias=bias,
device=device,
dtype=dtype,
)
def forward(
self,
value: torch.Tensor,
frequencies: tuple[torch.Tensor, torch.Tensor] | None = None,
mask: torch.Tensor | None = None,
) -> torch.Tensor:
"""Apply normalized, optionally rotary, grouped-query attention."""
batch, sequence, _ = value.shape
query = self.wq(value).reshape(
batch, sequence, self.heads, self.head_dim
).transpose(1, 2)
key = self.wk(value).reshape(
batch, sequence, self.kv_heads, self.head_dim
).transpose(1, 2)
val = self.wv(value).reshape(
batch, sequence, self.kv_heads, self.head_dim
).transpose(1, 2)
query, key = self.qknorm(query, key)
if frequencies is not None:
query = _applyRotary(query, frequencies)
key = _applyRotary(key, frequencies)
if self.heads != self.kv_heads:
repeat = self.heads // self.kv_heads
key = key.repeat_interleave(repeat, dim=1)
val = val.repeat_interleave(repeat, dim=1)
attention_mask = _attentionMask(mask)
attended = F.scaled_dot_product_attention(
query,
key,
val,
attn_mask=attention_mask,
)
attended = attended.transpose(1, 2).reshape(batch, sequence, -1)
return self.wo(attended * torch.sigmoid(self.gate(value)))
class KreaQkNorm(nn.Module):
"""Apply independent RMS normalization to query and key heads."""
def __init__(self, head_dim: int, *, device=None):
super().__init__()
self.qnorm = KreaRmsNorm(head_dim, device=device)
self.knorm = KreaRmsNorm(head_dim, device=device)
def forward(
self,
query: torch.Tensor,
key: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Normalize query and key while preserving their shapes."""
return self.qnorm(query), self.knorm(key)
class KreaSwiGlu(nn.Module):
"""SwiGLU MLP with the Krea checkpoint's dimensions."""
def __init__(
self,
features: int,
multiplier: int,
*,
bias: bool,
device=None,
dtype=None,
) -> None:
super().__init__()
hidden = int(2 * features / 3) * multiplier
hidden = 128 * ((hidden + 127) // 128)
self.gate = nn.Linear(features, hidden, bias=bias, device=device, dtype=dtype)
self.up = nn.Linear(features, hidden, bias=bias, device=device, dtype=dtype)
self.down = nn.Linear(hidden, features, bias=bias, device=device, dtype=dtype)
def forward(self, value: torch.Tensor) -> torch.Tensor:
"""Apply gated SiLU projection."""
return self.down(F.silu(self.gate(value)) * self.up(value))
class KreaSimpleModulation(nn.Module):
"""Final two-way AdaLN modulation."""
def __init__(self, features: int, *, device=None):
super().__init__()
self.lin = nn.Parameter(
torch.zeros(2, features, device=device, dtype=torch.float32),
requires_grad=False,
)
def forward(self, value: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Return final scale and shift."""
value = value[:, None, :] + self.lin.to(value.dtype)[None, :, :]
scale, shift = value.unbind(dim=1)
return scale, shift
class KreaDoubleModulation(nn.Module):
"""Six-way modulation shared by one single-stream block."""
def __init__(self, features: int, *, device=None):
super().__init__()
self.lin = nn.Parameter(
torch.zeros(6 * features, device=device, dtype=torch.float32),
requires_grad=False,
)
def forward(self, value: torch.Tensor) -> tuple[torch.Tensor, ...]:
"""Return pre/post scale, shift, and residual gates."""
return tuple((value + self.lin.to(value.dtype)).chunk(6, dim=-1))
class KreaTextFusionBlock(nn.Module):
"""One attention and MLP block in the text fusion adapter."""
def __init__(self, config: Krea2Config, *, device=None, dtype=None):
super().__init__()
self.prenorm = KreaRmsNorm(config.text_width, device=device)
self.postnorm = KreaRmsNorm(config.text_width, device=device)
self.attn = KreaAttention(
config.text_width,
config.text_heads,
config.text_kv_heads,
bias=config.bias,
device=device,
dtype=dtype,
)
self.mlp = KreaSwiGlu(
config.text_width,
config.mlp_multiplier,
bias=config.bias,
device=device,
dtype=dtype,
)
def forward(self, value: torch.Tensor, mask=None) -> torch.Tensor:
"""Apply a residual text-fusion block."""
value = value + self.attn(self.prenorm(value), mask=mask)
return value + self.mlp(self.postnorm(value))
class KreaTextFusion(nn.Module):
"""Fuse twelve Qwen layer taps into one text sequence."""
def __init__(self, config: Krea2Config, *, device=None, dtype=None):
super().__init__()
self.layerwise_blocks = nn.ModuleList(
[
KreaTextFusionBlock(config, device=device, dtype=dtype)
for _ in range(2)
]
)
self.projector = nn.Linear(
config.text_layers,
1,
bias=False,
device=device,
dtype=dtype,
)
self.refiner_blocks = nn.ModuleList(
[
KreaTextFusionBlock(config, device=device, dtype=dtype)
for _ in range(2)
]
)
def forward(self, value: torch.Tensor, mask=None) -> torch.Tensor:
"""Project the layer axis and refine the resulting text sequence."""
batch, sequence, layers, width = value.shape
value = value.reshape(batch * sequence, layers, width)
for block in self.layerwise_blocks:
value = block(value)
value = value.reshape(batch, sequence, layers, width).permute(0, 1, 3, 2)
value = self.projector(value).squeeze(-1)
for block in self.refiner_blocks:
value = block(value, mask=mask)
return value
class KreaSingleBlock(nn.Module):
"""One single-stream MMDiT block."""
def __init__(self, config: Krea2Config, *, device=None, dtype=None):
super().__init__()
self.mod = KreaDoubleModulation(config.features, device=device)
self.prenorm = KreaRmsNorm(config.features, device=device)
self.postnorm = KreaRmsNorm(config.features, device=device)
self.attn = KreaAttention(
config.features,
config.heads,
config.kv_heads,
bias=config.bias,
device=device,
dtype=dtype,
)
self.mlp = KreaSwiGlu(
config.features,
config.mlp_multiplier,
bias=config.bias,
device=device,
dtype=dtype,
)
def forward(self, value, modulation, frequencies, mask=None):
"""Apply modulated attention and MLP residuals."""
pre_scale, pre_shift, pre_gate, post_scale, post_shift, post_gate = (
self.mod(modulation)
)
attention_input = (1 + pre_scale) * self.prenorm(value) + pre_shift
value = value + pre_gate * self.attn(
attention_input,
frequencies,
mask,
)
mlp_input = (1 + post_scale) * self.postnorm(value) + post_shift
return value + post_gate * self.mlp(mlp_input)
class KreaLastLayer(nn.Module):
"""Final normalized patch projection with checkpoint-compatible names."""
def __init__(self, config: Krea2Config, *, device=None, dtype=None):
super().__init__()
self.norm = KreaRmsNorm(config.features, device=device)
self.linear = nn.Linear(
config.features,
config.patch_size ** 2 * config.latent_channels,
bias=True,
device=device,
dtype=dtype,
)
self.modulation = KreaSimpleModulation(
config.features,
device=device,
)
def forward(self, value: torch.Tensor, timestep_value: torch.Tensor):
"""Apply final AdaLN modulation and patch projection."""
scale, shift = self.modulation(timestep_value)
return self.linear((1 + scale[:, None, :]) * self.norm(value) + shift[:, None, :])
def _attentionMask(mask: torch.Tensor | None) -> torch.Tensor | None:
"""Expand a token mask into the boolean mask expected by SDPA."""
if mask is None:
return None
return mask[:, None, :, None] & mask[:, None, None, :]
def _rotaryFrequencies(
positions: torch.Tensor,
axis_dimensions: Iterable[int],
theta: float,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Build three-axis rotary frequencies for image and text positions."""
cos_parts = []
sin_parts = []
for axis, dimension in enumerate(axis_dimensions):
frequencies = torch.arange(
0,
dimension,
2,
device=positions.device,
dtype=torch.float32,
)
frequencies = 1.0 / theta ** (frequencies / dimension)
angles = positions[..., axis, None].float() * frequencies
cos_parts.append(torch.cos(angles))
sin_parts.append(torch.sin(angles))
return torch.cat(cos_parts, dim=-1), torch.cat(sin_parts, dim=-1)
def _applyRotary(
value: torch.Tensor,
frequencies: tuple[torch.Tensor, torch.Tensor],
) -> torch.Tensor:
"""Apply pairwise rotary rotation to a head-major tensor."""
cos, sin = frequencies
original_dtype = value.dtype
value = value.float().reshape(*value.shape[:-1], -1, 2)
cos = cos[:, None, :, :].unsqueeze(-1)
sin = sin[:, None, :, :].unsqueeze(-1)
first, second = value.unbind(dim=-1)
rotated = torch.stack((first * cos[..., 0] - second * sin[..., 0],
first * sin[..., 0] + second * cos[..., 0]), dim=-1)
return rotated.reshape(*rotated.shape[:-2], -1).to(original_dtype)
def _patchify(value: torch.Tensor, patch: int) -> torch.Tensor:
"""Convert an NCHW latent into row-major patch tokens."""
batch, channels, height, width = value.shape
return value.reshape(
batch,
channels,
height // patch,
patch,
width // patch,
patch,
).permute(0, 2, 4, 1, 3, 5).reshape(
batch,
(height // patch) * (width // patch),
channels * patch * patch,
)
def _unpatchify(value: torch.Tensor, height: int, width: int, patch: int, channels: int):
"""Convert row-major patch tokens back to an NCHW latent."""
batch = value.shape[0]
return value.reshape(
batch,
height // patch,
width // patch,
channels,
patch,
patch,
).permute(0, 3, 1, 4, 2, 5).reshape(batch, channels, height, width)
def patchify(value: torch.Tensor, patch: int = KREA2_PATCH_SIZE) -> torch.Tensor:
"""Convert an NCHW latent into row-major Krea patch tokens."""
return _patchify(value, patch)
def unpatchify(
value: torch.Tensor,
height: int,
width: int,
channels: int = KREA2_LATENT_CHANNELS,
patch: int = KREA2_PATCH_SIZE,
) -> torch.Tensor:
"""Convert row-major Krea patch tokens back to an NCHW latent."""
return _unpatchify(value, height, width, patch, channels)
class Krea2Model(nn.Module):
"""Krea 2 MMDiT wrapper accepting NCHW latents and conditioning."""
def __init__(
self,
config: Krea2Config = Krea2Config(),
*,
device=None,
dtype=torch.bfloat16,
quantization_manifest: QuantizationManifest | None = None,
) -> None:
super().__init__()
self.config = config
self._manifest = quantization_manifest or QuantizationManifest({})
self._device = device
self.first = nn.Linear(
config.latent_channels * config.patch_size ** 2,
config.features,
bias=True,
device=device,
dtype=dtype,
)
self.blocks = nn.ModuleList(
[
KreaSingleBlock(config, device=device, dtype=dtype)
for _ in range(config.blocks)
]
)
self.tmlp = nn.Sequential(
nn.Linear(config.timestep_width, config.features, device=device, dtype=dtype),
nn.GELU(approximate="tanh"),
nn.Linear(config.features, config.features, device=device, dtype=dtype),
)
self.txtfusion = KreaTextFusion(config, device=device, dtype=dtype)
self.txtmlp = nn.Sequential(
KreaRmsNorm(config.text_width, device=device),
nn.Linear(config.text_width, config.features, device=device, dtype=dtype),
nn.GELU(approximate="tanh"),
nn.Linear(config.features, config.features, device=device, dtype=dtype),
)
self.tproj = nn.Sequential(
nn.GELU(approximate="tanh"),
nn.Linear(config.features, config.features * 6, device=device, dtype=dtype),
)
self.last = KreaLastLayer(config, device=device, dtype=dtype)
def _replaceQuantizedLinear(self) -> None:
"""Replace manifest-selected ordinary linears with scaled FP8 linears."""
for name in self._manifest.layers:
module = self.get_submodule(name)
if not isinstance(module, nn.Linear):
raise DiffusionCliError(
f"Quantized Krea layer is not a linear module: {name}"
)
replacement = ScaledFp8Linear(
module.in_features,
module.out_features,
bias=module.bias is not None,
device=module.weight.device,
dtype=module.bias.dtype if module.bias is not None else torch.bfloat16,
)
parent_name, _, child_name = name.rpartition(".")
parent = self.get_submodule(parent_name) if parent_name else self
setattr(parent, child_name, replacement)
def forward(
self,
latent: torch.Tensor,
timestep: torch.Tensor,
conditioning: Krea2Conditioning,
) -> torch.Tensor:
"""Return flow velocity with the same NCHW shape as ``latent``."""
config = self.config
if latent.ndim != 4 or latent.shape[1] != config.latent_channels:
raise DiffusionCliError(
f"Krea latent must be NCHW with {config.latent_channels} channels"
)
if latent.shape[2] % config.patch_size or latent.shape[3] % config.patch_size:
raise DiffusionCliError("Krea latent dimensions must be patch aligned")
hidden = conditioning.hidden_states
mask = conditioning.attention_mask
if hidden.ndim != 4:
raise DiffusionCliError("Krea conditioning must have four dimensions")
if hidden.shape[2] != config.text_layers or hidden.shape[3] != config.text_width:
raise DiffusionCliError(
"Krea conditioning must contain 12 layers of 2560 features"
)
if mask.ndim != 2 or mask.shape[:2] != hidden.shape[:2]:
raise DiffusionCliError("Krea conditioning mask shape does not match")
if latent.shape[0] != hidden.shape[0]:
raise DiffusionCliError("Krea latent and conditioning batch mismatch")
timestep = timestep.reshape(-1)
if timestep.numel() == 1:
timestep = timestep.expand(latent.shape[0])
if timestep.numel() != latent.shape[0]:
raise DiffusionCliError("Krea timestep and latent batch mismatch")
hidden = hidden.to(device=latent.device)
mask = mask.to(device=latent.device, dtype=torch.bool)
_finite(latent, "input")
_finite(hidden, "conditioning")
image = self.first(_patchify(latent, config.patch_size))
time_embedding = _timestepEmbedding(
timestep,
config.timestep_width,
image.dtype,
)
modulation = self.tproj(self.tmlp(time_embedding))
text = self.txtfusion(hidden.to(image.dtype), mask=mask)
text = self.txtmlp(text)
batch, _, height, width = latent.shape
image_height = height // config.patch_size
image_width = width // config.patch_size
text_length = text.shape[1]
text_positions = torch.zeros(
batch,
text_length,
3,
device=latent.device,
dtype=torch.float32,
)
image_positions = torch.zeros(
batch,
image_height,
image_width,
3,
device=latent.device,
dtype=torch.float32,
)
image_positions[..., 1] = torch.arange(
image_height,
device=latent.device,
)[None, :, None]
image_positions[..., 2] = torch.arange(
image_width,
device=latent.device,
)[None, None, :]
image_positions = image_positions.reshape(batch, -1, 3)
value = torch.cat((text, image), dim=1)
positions = torch.cat((text_positions, image_positions), dim=1)
mask = torch.cat(
(
mask,
torch.ones(batch, image.shape[1], device=latent.device, dtype=torch.bool),
),
dim=1,
)
pad = (-value.shape[1]) % 256
if pad:
value = F.pad(value, (0, 0, 0, pad))
positions = F.pad(positions, (0, 0, 0, pad))
mask = F.pad(mask, (0, pad), value=False)
axes = _axisDimensions(config)
frequencies = _rotaryFrequencies(positions, axes, config.rope_theta)
for block in self.blocks:
value = block(value, modulation, frequencies, mask)
value = value[:, text_length : text_length + image.shape[1]]
output = self.last(value, self.tmlp(time_embedding))
output = _unpatchify(
output,
height,
width,
config.patch_size,
config.latent_channels,
)
_finite(output, "output")
return output
def toDevice(self, device, dtype=None) -> None:
"""Move model components while preserving quantized storage dtypes."""
moveModulePreservingQuantization(
self,
device,
dtype or torch.bfloat16,
)
self._device = device
def toCpu(self) -> None:
"""Move the model back to CPU without expanding FP8 weights."""
self.toDevice(torch.device("cpu"))
def _axisDimensions(config: Krea2Config) -> tuple[int, int, int]:
"""Return the three RoPE axis widths from the official layout."""
head_dim = config.features // config.heads
unit = head_dim // 16
return head_dim - 12 * unit, 6 * unit, 6 * unit
def _timestepEmbedding(
timestep: torch.Tensor,
width: int,
dtype: torch.dtype,
) -> torch.Tensor:
"""Build the official sinusoidal timestep embedding."""
timestep = timestep.reshape(-1)
half = width // 2
frequencies = torch.exp(
-math.log(10000.0)
* torch.arange(half, device=timestep.device, dtype=torch.float32)
/ half
)
angles = timestep.float()[:, None] * 1000.0 * frequencies[None]
return torch.cat((torch.cos(angles), torch.sin(angles)), dim=-1).to(dtype)
def normalizeKreaStateDict(
state_dict: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
"""Strip known ComfyUI diffusion prefixes and metadata tensors."""
normalized = {}
for key, value in state_dict.items():
if key.endswith(".comfy_quant"):
continue
for prefix in ("model.diffusion_model.", "diffusion_model.", "model."):
if key.startswith(prefix):
key = key.removeprefix(prefix)
break
normalized[key] = value
return normalized
def detectKrea2Config(state_dict: dict[str, torch.Tensor]) -> Krea2Config:
"""Validate supported Krea checkpoint shapes and infer block count."""
state_dict = normalizeKreaStateDict(state_dict)
try:
first = state_dict["first.weight"]
text_projector = state_dict["txtfusion.projector.weight"]
except KeyError as exc:
raise DiffusionCliError(
"Checkpoint is not a Krea 2 diffusion model"
) from exc
blocks = sorted(
{
int(match.group(1))
for key in state_dict
if (match := re.match(r"blocks\.(\d+)\.", key))
}
)
if (
text_projector.ndim != 2
or tuple(text_projector.shape) != (1, KREA2_TEXT_LAYERS)
):
raise DiffusionCliError("Unsupported Krea text-fusion projector shape")
config = Krea2Config(
features=first.shape[0],
latent_channels=KREA2_LATENT_CHANNELS,
patch_size=KREA2_PATCH_SIZE,
text_layers=text_projector.shape[1],
blocks=len(blocks),
)
if first.shape[1] != config.latent_channels * config.patch_size ** 2:
raise DiffusionCliError("Unsupported Krea latent patch shape")
if blocks != list(range(KREA2_BLOCKS)):
raise DiffusionCliError(
"Unsupported Krea block numbering: expected blocks 0 through 27"
)
if config.features != KREA2_FEATURES or config.blocks != KREA2_BLOCKS:
raise DiffusionCliError(
"Unsupported Krea 2 architecture constants: expected 6144 features "
"and 28 blocks"
)
return config
def loadKrea2StateDict(
model: Krea2Model,
state_dict: dict[str, torch.Tensor],
manifest: QuantizationManifest | None = None,
) -> Krea2Model:
"""Strictly load a Krea state dict into an existing, including reduced, model."""
normalized = normalizeKreaStateDict(state_dict)
if manifest is None:
manifest = manifestFromStateDict(normalized)
manifest = validateQuantizedStateDict(normalized, manifest)
model._manifest = manifest
model._replaceQuantizedLinear()
missing, unexpected = model.load_state_dict(
normalized,
strict=False,
assign=True,
)
if missing or unexpected:
raise DiffusionCliError(
"Krea diffusion checkpoint key mismatch: "
f"missing={list(missing)[:3]}, unexpected={list(unexpected)[:3]}"
)
return model
def buildKrea2Model(
source: ModelSource,
*,
device=None,
dtype=torch.bfloat16,
) -> Krea2Model:
"""Load a local Krea diffusion checkpoint with strict key accounting."""
state_dict = loadStateDict(source)
manifest = validateQuantizedStateDict(
normalizeKreaStateDict(state_dict),
loadQuantizationManifest(source),
)
config = detectKrea2Config(state_dict)
with torch.device("meta"):
model = Krea2Model(
config,
device="meta",
dtype=dtype,
quantization_manifest=manifest,
)
loadKrea2StateDict(model, state_dict, manifest)
model.toDevice(device or torch.device("cpu"), dtype)
model.eval()
return model
# Descriptive aliases used by integration callers.
loadKrea2Model = buildKrea2Model