Merge pull request #4962 from vladmandic/feat/krea2

feat(krea2): add Krea 2 (K2) image model support

data/reference-distilled.json

@@ -9,6 +9,16 @@
     "extras": "steps: 4, cfg_scale: 0.0",
     "size": 20.81
   },
+  "Krea 2 Turbo": {
+    "path": "CalamitousFelicitousness/Krea-2-Turbo-Diffusers",
+    "preview": "CalamitousFelicitousness--Krea-2-Turbo-Diffusers.jpg",
+    "desc": "Krea 2 (K2) Turbo is the 8-step distilled inference model of the Krea 2 family, trained from scratch by Krea. A 12.9B-parameter single-stream flow-matching DiT that uses a Qwen3-VL-4B vision-language model as its text encoder and the Qwen-Image VAE. Runs without classifier-free guidance; LoRAs trained on Krea 2 Base apply directly.",
+    "skip": true,
+    "tags": "distilled",
+    "extras": "sampler: Default, cfg_scale: 0.0, steps: 8, width: 1024, height: 1024",
+    "size": 34.0,
+    "date": "2026 June"
+  },
   "StabilityAI Stable Cascade Lite": {
     "path": "huggingface/stabilityai/stable-cascade-lite",
     "skip": true,

data/reference.json

@@ -193,6 +193,15 @@
     "size": 53.58,
     "date": "2026 June"
   },
+  "Krea 2 Base": {
+    "path": "CalamitousFelicitousness/Krea-2-Base-Diffusers",
+    "preview": "CalamitousFelicitousness--Krea-2-Base-Diffusers.jpg",
+    "desc": "Krea 2 (K2) Base is the undistilled foundation model of the Krea 2 family, trained from scratch by Krea. A 12.9B-parameter single-stream flow-matching DiT that uses a Qwen3-VL-4B vision-language model as its text encoder and the Qwen-Image VAE. The base checkpoint is intended for fine-tuning and LoRA training; LoRAs trained on it apply to Krea 2 Turbo.",
+    "skip": true,
+    "extras": "sampler: Default, cfg_scale: 3.5, steps: 52, width: 1024, height: 1024",
+    "size": 34.0,
+    "date": "2026 June"
+  },
   "Baidu ERNIE-Image": {
     "path": "baidu/ERNIE-Image",
     "preview": "baidu--ERNIE-Image.jpg",

modules/lora/lora_load.py

@@ -27,6 +27,7 @@ _NATIVE_DISPATCH = {
     'ernieimage': 'pipelines.ernie.ernie_lora',
     'f2':         'pipelines.flux.flux2_lora',
     'anima':      'pipelines.anima.anima_lora',
+    'krea2':      'pipelines.krea2.krea2_lora',
 }
 
 

modules/lora/lora_overrides.py

@@ -31,6 +31,7 @@ allow_native = [
     'zimage',
     'anima',
     'ernieimage',
+    'krea2',
 ]
 
 

modules/modeldata.py

@@ -57,6 +57,8 @@ def get_model_type(pipe):
         model_type = 'zimage'
     elif "Ideogram4" in name:
         model_type = 'ideogram4'
+    elif "Krea2" in name:
+        model_type = 'krea2'
     elif "LuminaDiMOO" in name:
         model_type = 'luminadimoo'
     elif "Lumina2" in name:

modules/processing_vae.py

@@ -130,8 +130,8 @@ def full_vae_decode(latents, model):
         latents = latents + shift_factor
 
     # check dims
-    if model.vae.__class__.__name__ in ['AutoencoderKLWan'] and latents.ndim == 4:
-        latents = latents.unsqueeze(2) # wan is __nhw
+    if model.vae.__class__.__name__ in ['AutoencoderKLWan', 'AutoencoderKLQwenImage'] and latents.ndim == 4:
+        latents = latents.unsqueeze(2) # video VAEs (wan, qwen-image) expect a frame axis
 
     # handle quants
     if getattr(model.vae, "post_quant_conv", None) is not None:

modules/sd_detect.py

@@ -151,6 +151,8 @@ def guess_by_name(fn, current_guess):
         new_guess = 'NucleusImage'
     elif 'z-image' in fn.lower() or 'z_image' in fn.lower():
         new_guess = 'ZImage'
+    elif 'krea-2' in fn.lower() or 'krea2' in fn.lower():
+        new_guess = 'Krea2'
     elif 'ideogram' in fn.lower():
         new_guess = 'Ideogram4'
     elif 'longcat-image' in fn.lower():

modules/sd_models.py

@@ -540,6 +540,10 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf
             from pipelines.model_z_image import load_z_image
             sd_model = load_z_image(checkpoint_info, diffusers_load_config)
             allow_post_quant = False
+        elif model_type in ['Krea2']:
+            from pipelines.model_krea2 import load_krea2
+            sd_model = load_krea2(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
         elif model_type in ['Ideogram4']:
             from pipelines.model_ideogram4 import load_ideogram4
             sd_model = load_ideogram4(checkpoint_info, diffusers_load_config)

modules/sd_samplers_common.py

@@ -11,7 +11,7 @@ from modules.image import convert
 
 SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options'])
 approximation_indexes = { "Simple": 0, "Approximate": 1, "TAESD": 2, "Full VAE": 3 }
-flow_models = ['f1', 'f2', 'sd3', 'lumina', 'auraflow', 'sana', 'zimage', 'lumina2', 'cogview4', 'h1', 'cosmos', 'anima', 'chroma', 'omnigen', 'omnigen2', 'longcat', 'ideogram4']
+flow_models = ['f1', 'f2', 'sd3', 'lumina', 'auraflow', 'sana', 'zimage', 'lumina2', 'cogview4', 'h1', 'cosmos', 'anima', 'chroma', 'omnigen', 'omnigen2', 'longcat', 'ideogram4', 'krea2']
 warned = False
 queue_lock = threading.Lock()
 

modules/shared_items.py

@@ -70,6 +70,7 @@ pipelines = {
     'HiDreamO1': None,
     'HunyuanImage3': None,
     'Ideogram4': None,
+    'Krea2': None,
     'Lens': None,
     'LuminaDiMOO': None,
     'Meissonic': None,

modules/ui_extra_networks_lora.py

@@ -69,6 +69,7 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage):
         'f1': 'Flux', 'sd1': 'SD 1.5', 'sd2': 'SD 2', 'xl': 'SDXL',
         'sd3': 'SD3', 'sc': 'Cascade', 'hv': 'HunyuanVideo',
         'chroma': 'Chroma', 'zimage': 'zImage', 'qwen': 'Qwen',
+        'krea2': 'Krea 2',
     }
 
     def cleanup_version(self, dct, lora):

modules/vae/sd_vae_taesd.py

@@ -76,7 +76,7 @@ def get_model(model_cls, variant=None):
         variant = 'TAE FLUX.2'
     elif model_cls in {'sd3'}:
         variant = 'TAE SD3'
-    elif model_cls in {'wanai', 'qwen', 'chrono', 'cosmos', 'anima', 'fibo', 'joy'}:
+    elif model_cls in {'wanai', 'qwen', 'chrono', 'cosmos', 'anima', 'fibo', 'joy', 'krea2'}:
         variant = 'TAE WanVideo'
     else:
         warn_once(f'cls={shared.sd_model.__class__.__name__} type={shared.sd_model_type} unsuppported', variant=variant)

pipelines/generic_shared.py

@@ -98,6 +98,11 @@ shared_te_map = {
         'target_subfolder': 'text_encoder',
     },
 
+    'Qwen3-VL 4B Base': {  # Krea 2 base+turbo share one canonical 4B copy
+        'cls': transformers.Qwen3VLModel,
+        'identifier': 'krea',
+        'target_repo': 'Qwen/Qwen3-VL-4B-Instruct',
+    },
     'Qwen3-VL 8B SDNQ-UInt4': {
         'cls': transformers.Qwen3VLModel,
         'identifier': 'uint4',

pipelines/krea2/__init__.py

@@ -0,0 +1,8 @@
+from pipelines.krea2.transformer_krea2 import Krea2Transformer2DModel
+from pipelines.native_transformer import TransformerSpec
+
+
+# Checkpoint keys are bare (`first.`, `blocks.N.`, `txtfusion.`, ...) and the transformer's
+# module tree mirrors them exactly, so no state-dict conversion is needed. The model has no
+# rope/pos buffers, so the default acceptable-missing set is sufficient.
+KREA2_SPEC = TransformerSpec(cls=Krea2Transformer2DModel, converter=None)

pipelines/krea2/krea2_lora.py

@@ -0,0 +1,71 @@
+"""Krea 2 native adapter loader.
+
+Runs when :func:`modules.lora.lora_overrides.get_method` returns ``'native'``
+(``lora_force_diffusers`` off and ``krea2`` in ``allow_native``).
+
+The transformer module tree mirrors the checkpoint (``blocks.N.attn.{wq,wk,wv,wo,gate}``,
+``blocks.N.mlp.{gate,up,down}``, ``txtfusion.*``, ``first``, ``last`` ...), so dotted keys
+bind verbatim with no name rewrite and no fused-QKV split. Kohya flat-underscore keys are
+reconstructed back to dotted paths, protecting the two compound module names
+(``layerwise_blocks``, ``refiner_blocks``).
+"""
+
+from modules.lora import native_adapter
+
+
+KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT
+
+# Top-level module names that a bare diffusers-format LoRA key can start with.
+BARE_DIFFUSERS_PREFIXES = ("blocks.", "txtfusion.", "first.", "last.", "tmlp.", "tproj.", "txtmlp.")
+
+
+def resolve_targets(prefix_used, base):
+    """Return ``[(diffusers_path, None), ...]`` for a parsed group key.
+
+    K2's diffusers module names equal the checkpoint names, so dotted keys map verbatim.
+    Universal passthrough prefixes are handled upstream by
+    :func:`native_adapter.resolve_group_targets`.
+    """
+    if prefix_used in (None, "diffusion_model.", "transformer."):
+        return [(base, None)]
+    if prefix_used in ("lora_unet_", "lora_transformer_"):
+        return _underscore_to_dotted(base)
+    return []
+
+
+def _underscore_to_dotted(base):
+    """Rebuild a dotted path from a kohya flat-underscore base, keeping compound names intact."""
+    protected = base.replace("layerwise_blocks", "layerwise@blocks").replace("refiner_blocks", "refiner@blocks")
+    return [(protected.replace("_", ".").replace("@", "_"), None)]
+
+
+_BIND_KWARGS = dict(
+    resolve_targets=resolve_targets,
+    prefixes=KNOWN_PREFIXES,
+    bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES,
+    arch_name="krea2",
+)
+
+
+def try_load_lora(name, network_on_disk, lora_scale):
+    return native_adapter.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS)
+
+
+def try_load_lokr(name, network_on_disk, lora_scale):
+    return native_adapter.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS)
+
+
+def try_load_loha(name, network_on_disk, lora_scale):
+    return native_adapter.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS)
+
+
+def try_load_oft(name, network_on_disk, lora_scale):
+    return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS)
+
+
+def try_load(name, network_on_disk, lora_scale):
+    """Run every Krea 2 family loader, merge any that match."""
+    return native_adapter.try_load_chain(
+        name, network_on_disk, lora_scale,
+        family_loaders=(try_load_lora, try_load_lokr, try_load_loha, try_load_oft),
+    )

pipelines/krea2/pipeline_krea2.py

@@ -0,0 +1,273 @@
+"""Krea 2 (K2) text-to-image pipeline.
+
+A single-stream flow-matching pipeline that conditions a custom DiT on stacked Qwen3-VL
+hidden states and decodes with the Qwen-Image VAE. The encode, packing, denoise and decode
+steps mirror the reference K2 inference code. This module imports only diffusers/transformers
+so the repos can ship it for standalone use; SD.Next-specific wiring lives in the loader.
+"""
+
+import torch
+from einops import rearrange, repeat
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.loaders import FromSingleFileMixin
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+
+class Krea2Pipeline(DiffusionPipeline, FromSingleFileMixin):
+    r"""Text-to-image generation with Krea 2.
+
+    Args:
+        transformer (`Krea2Transformer2DModel`): single-stream flow-matching DiT.
+        text_encoder (`Qwen3VLModel`): multimodal backbone tapped for stacked hidden states.
+        tokenizer (`Qwen2Tokenizer`): tokenizer paired with `text_encoder`.
+        vae (`AutoencoderKLQwenImage`): f8/16-channel latent autoencoder.
+        scheduler (`FlowMatchEulerDiscreteScheduler`): exponential-shift flow-matching scheduler.
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+    _callback_tensor_inputs = ["latents"]
+
+    # Conditioning template and layer taps (reference encoder.py / inference.py).
+    PROMPT_TEMPLATE_PREFIX = (
+        "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, "
+        "quantity, text, spatial relationships of the objects and background:<|im_end|>\n"
+        "<|im_start|>user\n"
+    )
+    PROMPT_TEMPLATE_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n"
+    PREFIX_TOKENS = 34
+    SUFFIX_START = 5
+    MAX_LENGTH = 512
+    SELECT_LAYERS = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35)
+    MIN_RES = 256
+    MAX_RES = 1280
+
+    def __init__(self, transformer, text_encoder, tokenizer, vae, scheduler):
+        super().__init__()
+        self.register_modules(
+            transformer=transformer,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            vae=vae,
+            scheduler=scheduler,
+        )
+        self.patch = int(getattr(transformer.config, "patch", 2))
+        self.latent_channels = int(getattr(transformer.config, "channels", 16))
+        self.vae_compression = 8  # AutoencoderKLQwenImage is f8
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_compression)
+        self._interrupt = False
+        self._guidance_scale = None
+        self._num_timesteps = 0
+
+    # --- text conditioning: port of encoder.py Qwen3VLConditioner.forward ---
+    def encode_prompt(self, prompts: list[str], device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return `(hidden, mask)` where hidden is `(B, L, len(SELECT_LAYERS), txtdim)`.
+
+        The user prompt is wrapped in the image-description chat template, padded to a fixed
+        length, and the assistant suffix appended; the system prefix tokens are then dropped.
+        """
+        texts = [self.PROMPT_TEMPLATE_PREFIX + p for p in prompts]
+        suffix = self.tokenizer([self.PROMPT_TEMPLATE_SUFFIX] * len(texts), return_tensors="pt").to(device)
+        main = self.tokenizer(
+            texts,
+            truncation=True,
+            padding="max_length",
+            max_length=self.MAX_LENGTH + self.PREFIX_TOKENS - self.SUFFIX_START,
+            return_tensors="pt",
+        ).to(device)
+
+        input_ids = torch.cat([main.input_ids, suffix.input_ids], dim=1)
+        mask = torch.cat([main.attention_mask.bool(), suffix.attention_mask.bool()], dim=1)
+
+        out = self.text_encoder(input_ids=input_ids, attention_mask=mask, output_hidden_states=True)
+        hidden = torch.stack([out.hidden_states[i] for i in self.SELECT_LAYERS], dim=2)
+        return hidden[:, self.PREFIX_TOKENS:], mask[:, self.PREFIX_TOKENS:]
+
+    # --- packing: port of sampling.prepare ---
+    def pack_sequence(self, latent: torch.Tensor, text_mask: torch.Tensor):
+        """Patchify the latent and build the joint text+image `(tokens, position_ids, mask)`."""
+        batch, _, height, width = latent.shape
+        patch = self.patch
+        grid_h, grid_w = height // patch, width // patch
+        device = latent.device
+
+        img_ids = torch.zeros(grid_h, grid_w, 3, device=device)
+        img_ids[..., 1] = torch.arange(grid_h, device=device)[:, None]
+        img_ids[..., 2] = torch.arange(grid_w, device=device)[None, :]
+        img_pos = repeat(img_ids, "h w c -> b (h w) c", b=batch)
+        img_mask = torch.ones(batch, grid_h * grid_w, dtype=torch.bool, device=device)
+        img_tokens = rearrange(latent, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch)
+
+        txt_pos = torch.zeros(batch, text_mask.shape[1], 3, device=device)
+        pos = torch.cat([txt_pos, img_pos], dim=1)
+        mask = torch.cat([text_mask, img_mask], dim=1)
+        return img_tokens, pos, mask
+
+    def prepare_latents(self, batch_size, height, width, dtype, device, generator, latents=None):
+        shape = (batch_size, self.latent_channels, height // self.vae_compression, width // self.vae_compression)
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+        return randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+
+    @staticmethod
+    def calculate_shift(image_seq_len, base_seq_len, max_seq_len, base_shift, max_shift):
+        slope = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+        return image_seq_len * slope + (base_shift - slope * base_seq_len)
+
+    def decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        """Denormalize and decode latents with the Qwen-Image VAE (port of autoencoder.decode)."""
+        cfg = self.vae.config
+        mean = torch.tensor(cfg.latents_mean, device=latents.device, dtype=latents.dtype).view(1, -1, 1, 1, 1)
+        std = torch.tensor(cfg.latents_std, device=latents.device, dtype=latents.dtype).view(1, -1, 1, 1, 1)
+        latents = latents.unsqueeze(2) * std + mean  # (B, C, 1, H, W); the VAE treats images as 1-frame video
+        image = self.vae.decode(latents).sample
+        return image.squeeze(2)
+
+    def encode_image(self, image, height, width, dtype, device):
+        """Encode a pixel image to a normalized latent (inverse of decode_latents)."""
+        pixels = self.image_processor.preprocess(image, height=height, width=width)
+        pixels = pixels.to(device=device, dtype=self.vae.dtype).unsqueeze(2)  # (B, C, 1, H, W)
+        cfg = self.vae.config
+        mean = torch.tensor(cfg.latents_mean, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
+        std = torch.tensor(cfg.latents_std, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
+        raw = self.vae.encode(pixels).latent_dist.mode()
+        return ((raw.to(dtype) - mean) / std).squeeze(2)
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = None,
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: int = 28,
+        guidance_scale: float = 4.5,
+        num_images_per_prompt: int = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        image=None,
+        strength: float = 0.6,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        attention_kwargs: dict | None = None,
+        callback_on_step_end=None,
+        callback_on_step_end_tensor_inputs: list[str] | None = None,
+    ):
+        align = self.vae_compression * self.patch
+        height = (height // align) * align
+        width = (width // align) * align
+
+        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
+        device = self._execution_device
+        dtype = self.transformer.dtype
+        self._guidance_scale = guidance_scale
+        self._interrupt = False
+
+        is_distilled = bool(getattr(self.transformer.config, "is_distilled", False))
+        do_cfg = guidance_scale is not None and guidance_scale > 0 and not is_distilled
+
+        text, text_mask = self.encode_prompt(prompts, device)
+        text = text.to(dtype)
+        if do_cfg:
+            negatives = [negative_prompt or ""] * len(prompts) if not isinstance(negative_prompt, list) else negative_prompt
+            uncond, uncond_mask = self.encode_prompt(negatives, device)
+            uncond = uncond.to(dtype)
+
+        batch = len(prompts) * num_images_per_prompt
+        text = text.repeat_interleave(num_images_per_prompt, dim=0)
+        text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0)
+        if do_cfg:
+            uncond = uncond.repeat_interleave(num_images_per_prompt, dim=0)
+            uncond_mask = uncond_mask.repeat_interleave(num_images_per_prompt, dim=0)
+
+        cfg = self.scheduler.config
+        grid_h = height // (self.vae_compression * self.patch)
+        grid_w = width // (self.vae_compression * self.patch)
+        mu = self.calculate_shift(
+            grid_h * grid_w,
+            cfg.get("base_image_seq_len", 256),
+            cfg.get("max_image_seq_len", 6400),
+            cfg.get("base_shift", 0.5),
+            cfg.get("max_shift", 1.15),
+        )
+        self.scheduler.set_timesteps(num_inference_steps, device=device, mu=mu)
+        timesteps = self.scheduler.timesteps
+
+        if image is not None:
+            clean = self.encode_image(image, height, width, dtype, device)
+            clean = clean.repeat_interleave(num_images_per_prompt, dim=0)
+            init_steps = min(int(num_inference_steps * strength), num_inference_steps)
+            t_start = max(num_inference_steps - init_steps, 0)
+            timesteps = timesteps[t_start:]
+            noise = self.prepare_latents(batch, height, width, dtype, device, generator)
+            latents = self.scheduler.scale_noise(clean, timesteps[:1], noise)
+        else:
+            latents = self.prepare_latents(batch, height, width, dtype, device, generator, latents)
+
+        img, pos, mask = self.pack_sequence(latents, text_mask)
+        if do_cfg:
+            _, uncond_pos, uncond_full_mask = self.pack_sequence(latents, uncond_mask)
+        self._num_timesteps = len(timesteps)
+        num_train = cfg.get("num_train_timesteps", 1000)
+
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+                model_t = (t.float() / num_train).reshape(1).expand(batch).to(device=device, dtype=img.dtype)
+                cond = self.transformer(
+                    hidden_states=img, encoder_hidden_states=text, timestep=model_t,
+                    position_ids=pos, attention_mask=mask, return_dict=False,
+                )[0]
+                if do_cfg:
+                    neg = self.transformer(
+                        hidden_states=img, encoder_hidden_states=uncond, timestep=model_t,
+                        position_ids=uncond_pos, attention_mask=uncond_full_mask, return_dict=False,
+                    )[0]
+                    velocity = cond + guidance_scale * (cond - neg)
+                else:
+                    velocity = cond
+                img = self.scheduler.step(velocity, t, img, return_dict=False)[0]
+
+                if callback_on_step_end is not None:
+                    # Unpack the packed tokens to a standard [B, C, h, w] latent for the callback's preview.
+                    cb_kwargs = {}
+                    if "latents" in (callback_on_step_end_tensor_inputs or ["latents"]):
+                        cb_kwargs["latents"] = rearrange(img, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=grid_h, w=grid_w, ph=self.patch, pw=self.patch)
+                    callback_on_step_end(self, i, t, cb_kwargs)
+                progress_bar.update()
+
+        latent = rearrange(img, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=grid_h, w=grid_w, ph=self.patch, pw=self.patch)
+
+        if output_type == "latent":
+            image = latent
+        else:
+            image = self.decode_latents(latent.to(dtype))
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        self.maybe_free_model_hooks()
+        if not return_dict:
+            return (image,)
+        return ImagePipelineOutput(images=image)
+
+
+class Krea2Img2ImgPipeline(Krea2Pipeline):
+    """Image-to-image task variant.
+
+    The denoise path is identical to the base pipeline (which already accepts `image` + `strength`);
+    this distinct class exists so the diffusers AUTO maps and `get_diffusers_task` can tell the two
+    tasks apart, following the per-task-class convention (e.g. ChromaPipeline / ChromaImg2ImgPipeline).
+    """

pipelines/krea2/transformer_krea2.py

@@ -0,0 +1,349 @@
+"""Krea 2 (K2) single-stream DiT, ported to the diffusers ModelMixin contract.
+
+The module tree mirrors the original `SingleStreamDiT` checkpoint exactly (``first``,
+``blocks.N.attn.{wq,wk,wv,gate,wo}``, ``txtfusion.*``, ``last`` ...), so the safetensors
+load is an identity map (no key conversion). Architecture specifics: gated attention,
+grouped-query attention, per-head QK RMSNorm, 3-axis rotary embedding, a shared+per-block
+modulation, and a text-fusion stage that collapses several text-encoder hidden-state layers
+into one conditioning stream.
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+
+
+def rope(pos: torch.Tensor, dim: int, theta: float = 1e4, ntk: float = 1.0) -> torch.Tensor:
+    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+    omega = 1.0 / ((theta * ntk) ** scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
+    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+    return out.float()
+
+
+def rope_apply(xq: torch.Tensor, xk: torch.Tensor, freqs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    freqs = freqs[:, None, :, :, :]
+    xq_ = freqs[..., 0] * xq_[..., 0] + freqs[..., 1] * xq_[..., 1]
+    xk_ = freqs[..., 0] * xk_[..., 0] + freqs[..., 1] * xk_[..., 1]
+    return xq_.reshape(*xq.shape).to(xq.dtype), xk_.reshape(*xk.shape).to(xk.dtype)
+
+
+def time_embed(t: torch.Tensor, dim: int, period: float = 1e4, tfactor: float = 1e3, device=None, dtype=None) -> torch.Tensor:
+    half = dim // 2
+    freqs = torch.exp(-math.log(period) * torch.arange(half, dtype=torch.float32, device=device) / half)
+    # t: (B,) -> (B, 1, half) so the embedding broadcasts as a per-sample vector.
+    args = (t.float() * tfactor)[:, None, None] * freqs
+    return torch.cat((torch.cos(args), torch.sin(args)), dim=-1).to(dtype=dtype)
+
+
+def expand_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """Repeat each KV head ``n_rep`` times so grouped-query attention runs on any SDPA backend."""
+    if n_rep == 1:
+        return x
+    return x.repeat_interleave(n_rep, dim=1)  # x: (B, kvheads, L, D) -> (B, heads, L, D)
+
+
+def segment_mask(mask: torch.Tensor) -> torch.Tensor:
+    """Expand a (B, L) key-padding mask into a (B, 1, L, L) attention mask."""
+    return mask.unsqueeze(1).unsqueeze(2) * mask.unsqueeze(1).unsqueeze(3)
+
+
+class RMSNorm(nn.Module):
+    """RMSNorm whose stored weight is a zero-init delta applied as ``scale + 1`` and computed in fp32."""
+
+    def __init__(self, features: int, eps: float = 1e-05):
+        super().__init__()
+        self.features = features
+        self.eps = eps
+        self.scale = nn.Parameter(torch.zeros(features, dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        t = F.rms_norm(x.float(), (self.features,), eps=self.eps, weight=self.scale.float() + 1.0)
+        return t.to(dtype)
+
+
+class QKNorm(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.qnorm = RMSNorm(dim)
+        self.knorm = RMSNorm(dim)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return self.qnorm(q), self.knorm(k), v
+
+
+class SimpleModulation(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.lin = nn.Parameter(torch.zeros(2, dim))
+        self.multiplier = 2
+
+    def forward(self, vec: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        out = vec + rearrange(self.lin, "two d -> 1 two d")
+        scale, shift = out.chunk(self.multiplier, dim=1)
+        return scale, shift
+
+
+class DoubleSharedModulation(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.lin = nn.Parameter(torch.zeros(6 * dim))
+
+    def forward(self, vec: torch.Tensor):
+        out = vec + self.lin
+        return out.chunk(6, dim=-1)
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, axdims: list[int], theta: float = 1e2, ntk: float = 1.0):
+        super().__init__()
+        self.axdims = axdims  # split of the head dimension across the position axes
+        self.theta = theta
+        self.ntk = ntk
+
+    def forward(self, pos: torch.Tensor) -> torch.Tensor:
+        return torch.cat([rope(pos[..., i], d, self.theta, self.ntk) for i, d in enumerate(self.axdims)], dim=-3)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, features: int, multiplier: int, bias: bool = False, multiple: int = 128):
+        super().__init__()
+        mlpdim = int(2 * features / 3) * multiplier
+        mlpdim = multiple * ((mlpdim + multiple - 1) // multiple)
+        self.gate = nn.Linear(features, mlpdim, bias=bias)
+        self.up = nn.Linear(features, mlpdim, bias=bias)
+        self.down = nn.Linear(mlpdim, features, bias=bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class Attention(nn.Module):
+    """Gated grouped-query attention with per-head QK RMSNorm and optional 3-axis RoPE."""
+
+    def __init__(self, dim: int, heads: int, kvheads: int | None = None, bias: bool = False):
+        super().__init__()
+        self.heads = heads
+        self.kvheads = kvheads if kvheads is not None else heads
+        self.headdim = dim // self.heads
+        self.n_rep = self.heads // self.kvheads
+
+        self.wq = nn.Linear(dim, self.headdim * self.heads, bias=bias)
+        self.wk = nn.Linear(dim, self.headdim * self.kvheads, bias=bias)
+        self.wv = nn.Linear(dim, self.headdim * self.kvheads, bias=bias)
+        self.gate = nn.Linear(dim, dim, bias=bias)
+        self.qknorm = QKNorm(self.headdim)
+        self.wo = nn.Linear(dim, dim, bias=bias)
+
+    def forward(self, x: torch.Tensor, freqs: torch.Tensor | None = None, mask: torch.Tensor | None = None) -> torch.Tensor:
+        q, k, v, gate = self.wq(x), self.wk(x), self.wv(x), self.gate(x)
+        q = rearrange(q, "B L (H D) -> B H L D", H=self.heads)
+        k = rearrange(k, "B L (H D) -> B H L D", H=self.kvheads)
+        v = rearrange(v, "B L (H D) -> B H L D", H=self.kvheads)
+
+        q, k, v = self.qknorm(q, k, v)
+        if freqs is not None:
+            q, k = rope_apply(q, k, freqs)
+        k, v = expand_kv(k, self.n_rep), expand_kv(v, self.n_rep)
+
+        # dispatch_attention_fn expects (B, L, H, D); the q/k/v math above stays in (B, H, L, D).
+        out = dispatch_attention_fn(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_mask=mask)
+        # A fully-masked query row (padding token) yields NaN on the CUDA SDPA backends; zero it so
+        # it cannot propagate through the next layer's `0 * NaN`. cuDNN returns 0 here already.
+        out = torch.nan_to_num(out)
+        out = rearrange(out, "B L H D -> B L (H D)")
+        return self.wo(out * F.sigmoid(gate))
+
+
+class LastLayer(nn.Module):
+    def __init__(self, features: int, patch: int, channels: int):
+        super().__init__()
+        self.norm = RMSNorm(features)
+        self.linear = nn.Linear(features, patch * patch * channels, bias=True)
+        self.modulation = SimpleModulation(features)
+        self.down = nn.Linear(features, features, bias=False)
+        self.up = nn.Linear(features, features, bias=False)
+
+    def forward(self, x: torch.Tensor, tvec: torch.Tensor) -> torch.Tensor:
+        scale, shift = self.modulation(tvec)
+        x = (1 + scale) * self.norm(x) + shift + self.up(self.down(x))
+        return self.linear(x)
+
+
+class TextFusionBlock(nn.Module):
+    def __init__(self, features: int, heads: int, multiplier: int, bias: bool = False, kvheads: int | None = None):
+        super().__init__()
+        self.prenorm = RMSNorm(features)
+        self.postnorm = RMSNorm(features)
+        self.attn = Attention(dim=features, heads=heads, bias=bias, kvheads=kvheads)
+        self.mlp = SwiGLU(features, multiplier, bias)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
+        x = x + self.attn(self.prenorm(x), mask=mask)
+        x = x + self.mlp(self.postnorm(x))
+        return x
+
+
+class TextFusionTransformer(nn.Module):
+    """Fuse ``num_txt_layers`` stacked text-encoder hidden states into a single conditioning stream.
+
+    ``num_txt_layers`` is the count of selected encoder layers fed in (projected down to 1),
+    not the transformer depth, which is fixed at 2 layerwise + 2 refiner blocks.
+    """
+
+    def __init__(self, num_txt_layers: int, txt_dim: int, heads: int, multiplier: int, bias: bool = False, kvheads: int | None = None):
+        super().__init__()
+        self.layerwise_blocks = nn.ModuleList([TextFusionBlock(txt_dim, heads, multiplier, bias, kvheads) for _ in range(2)])
+        self.projector = nn.Linear(num_txt_layers, 1, bias=False)
+        self.refiner_blocks = nn.ModuleList([TextFusionBlock(txt_dim, heads, multiplier, bias, kvheads) for _ in range(2)])
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
+        b, l, n, d = x.shape
+        x = x.reshape(b * l, n, d)
+        for block in self.layerwise_blocks:
+            x = block(x.contiguous(), mask=None)
+        x = rearrange(x, "(b l) n d -> b l d n", b=b, l=l)
+        x = self.projector(x).squeeze(-1)
+        for block in self.refiner_blocks:
+            x = block(x, mask=mask)
+        return x
+
+
+class SingleStreamBlock(nn.Module):
+    def __init__(self, features: int, heads: int, multiplier: int, bias: bool = False, kvheads: int | None = None):
+        super().__init__()
+        self.mod = DoubleSharedModulation(features)
+        self.prenorm = RMSNorm(features)
+        self.postnorm = RMSNorm(features)
+        self.attn = Attention(dim=features, heads=heads, bias=bias, kvheads=kvheads)
+        self.mlp = SwiGLU(features, multiplier, bias)
+
+    def forward(self, x: torch.Tensor, vec: torch.Tensor, freqs: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
+        prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec)
+        x = x + pregate * self.attn((1 + prescale) * self.prenorm(x) + preshift, freqs, mask)
+        x = x + postgate * self.mlp((1 + postscale) * self.postnorm(x) + postshift)
+        return x
+
+
+class Krea2Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+    r"""Single-stream flow-matching DiT backbone for Krea 2.
+
+    The transformer consumes patchified noisy image tokens plus stacked text-encoder hidden
+    states, fuses the text layers internally, concatenates text and image tokens into one
+    stream, and predicts the flow-matching velocity for the image-token positions.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["SingleStreamBlock", "TextFusionBlock"]
+    _repeated_blocks = ["SingleStreamBlock"]
+
+    @register_to_config
+    def __init__(
+        self,
+        features: int = 6144,
+        tdim: int = 256,
+        txtdim: int = 2560,
+        heads: int = 48,
+        kvheads: int = 12,
+        multiplier: int = 4,
+        layers: int = 28,
+        patch: int = 2,
+        channels: int = 16,
+        bias: bool = False,
+        theta: float = 1e3,
+        txtlayers: int = 12,
+        txtheads: int = 20,
+        txtkvheads: int = 20,
+        is_distilled: bool = False,
+    ):
+        super().__init__()
+        self.gradient_checkpointing = False
+        self.tdim = tdim
+        self.is_distilled = is_distilled
+
+        headdim = features // heads
+        axes = [headdim - 12 * (headdim // 16), 6 * (headdim // 16), 6 * (headdim // 16)]
+        assert sum(axes) == headdim, f"sum(axes)={sum(axes)} != headdim={headdim}"
+        assert all(a % 2 == 0 for a in axes), f"axes={axes}"
+
+        self.posemb = PositionalEncoding(axes, theta=theta, ntk=1.0)
+        self.first = nn.Linear(channels * patch**2, features, bias=True)
+        self.blocks = nn.ModuleList(
+            [SingleStreamBlock(features, heads, multiplier, bias, kvheads) for _ in range(layers)]
+        )
+        self.tmlp = nn.Sequential(
+            nn.Linear(tdim, features),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(features, features),
+        )
+        self.txtfusion = TextFusionTransformer(txtlayers, txtdim, txtheads, multiplier, bias, txtkvheads)
+        self.txtmlp = nn.Sequential(
+            RMSNorm(txtdim),
+            nn.Linear(txtdim, features),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(features, features),
+        )
+        self.last = LastLayer(features, patch, channels)
+        self.tproj = nn.Sequential(nn.GELU(approximate="tanh"), nn.Linear(features, features * 6))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        position_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        return_dict: bool = True,
+    ):
+        r"""
+        Args:
+            hidden_states: `(B, L_img, channels * patch ** 2)` patchified noisy image tokens.
+            encoder_hidden_states: `(B, L_txt, txtlayers, txtdim)` stacked text-encoder hidden states.
+            timestep: `(B,)` flow-matching time in `[0, 1]`.
+            position_ids: `(B, L_txt + L_img, 3)` `(t, h, w)` coordinates for the 3-axis RoPE.
+            attention_mask: `(B, L_txt + L_img)` boolean key-padding mask (True = valid token).
+        """
+        img = self.first(hidden_states)
+        t = self.tmlp(time_embed(timestep, self.tdim, device=img.device, dtype=img.dtype))
+        tvec = self.tproj(t)
+
+        txtmask = segment_mask(attention_mask[:, : encoder_hidden_states.shape[1]])
+        context = self.txtfusion(encoder_hidden_states, mask=txtmask)
+        context = self.txtmlp(context)
+
+        txtlen, imglen = context.shape[1], img.shape[1]
+        combined = torch.cat((context, img), dim=1)
+
+        # Pad the joint sequence to a multiple of 256 to keep compiled attention kernel shapes stable.
+        padlen = (-combined.shape[1]) % 256
+        if padlen > 0:
+            combined = F.pad(combined, (0, 0, 0, padlen))
+            attention_mask = F.pad(attention_mask, (0, padlen), value=False)
+            position_ids = F.pad(position_ids, (0, 0, 0, padlen))
+
+        mask = segment_mask(attention_mask)
+        freqs = self.posemb(position_ids)
+
+        for block in self.blocks:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                combined = self._gradient_checkpointing_func(block, combined, tvec, freqs, mask)
+            else:
+                combined = block(combined, tvec, freqs, mask)
+
+        output = self.last(combined, t)[:, txtlen : txtlen + imglen, :]
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)

pipelines/model_krea2.py

@@ -0,0 +1,53 @@
+import diffusers
+import transformers
+from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
+from modules.logger import log
+from pipelines import generic
+
+
+def load_krea2(checkpoint_info, diffusers_load_config=None):
+    if diffusers_load_config is None:
+        diffusers_load_config = {}
+    repo_id = sd_models.path_to_repo(checkpoint_info)
+    sd_models.hf_auth_check(checkpoint_info)
+    load_args, _ = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
+    log.debug(f'Load model: type=Krea2 repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
+
+    from pipelines.krea2.transformer_krea2 import Krea2Transformer2DModel
+    from pipelines.krea2.pipeline_krea2 import Krea2Pipeline, Krea2Img2ImgPipeline
+    from pipelines.krea2 import KREA2_SPEC
+    diffusers.Krea2Transformer2DModel = Krea2Transformer2DModel
+    diffusers.Krea2Pipeline = Krea2Pipeline
+    diffusers.Krea2Img2ImgPipeline = Krea2Img2ImgPipeline
+    generic.set_pipeline('Krea2', Krea2Pipeline)
+    # One class per task so get_diffusers_task defaults to text2image and set_diffuser_pipe switches
+    # to the img2img variant cleanly (matches the Chroma/Qwen per-task-class pattern).
+    from diffusers.pipelines import auto_pipeline
+    auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['krea2'] = Krea2Pipeline
+    auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING['krea2'] = Krea2Img2ImgPipeline
+    if repo_id is None or repo_id.lower() == 'none':
+        return None
+
+    # Keep small/sensitive layers in compute dtype. `first` (in=64) and `txtfusion.projector`
+    # (in=12) are below the int8 GEMM's minimum K; `last` is the output projection; `tmlp`/`tproj`
+    # produce the global per-block modulation, too int8-sensitive to quantize (its error compounds
+    # across blocks and steps). All are tiny next to the 28 blocks, so the memory cost is small.
+    transformer = generic.load_transformer(repo_id, cls_name=Krea2Transformer2DModel, load_config=diffusers_load_config, native_spec=KREA2_SPEC, modules_to_not_convert=['first', 'last', 'projector', 'tmlp', 'tproj'])
+    text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3VLModel, load_config=diffusers_load_config)
+
+    pipe = Krea2Pipeline.from_pretrained(
+        repo_id,
+        cache_dir=shared.opts.diffusers_dir,
+        transformer=transformer,
+        text_encoder=text_encoder,
+        **load_args,
+    )
+
+    generic.load_vae_override(pipe, diffusers_load_config)
+
+    del transformer
+    del text_encoder
+    sd_hijack_te.init_hijack(pipe)
+    sd_hijack_vae.init_hijack(pipe)
+    devices.torch_gc(force=True, reason='load')
+    return pipe

test/test-krea2-transformer.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python
+"""Offline parity test: Krea2Transformer2DModel vs the reference SingleStreamDiT.
+
+Builds both models from one tiny config, copies the reference state dict into the diffusers
+port, runs identical inputs, and asserts the forward outputs match. No server, no checkpoint.
+
+The reference checkpoint repo (mmdit.py) is expected at $KREA2_REF_DIR
+(default /home/ohiom/database/watering-hole).
+"""
+
+import importlib.util
+import os
+import sys
+from contextlib import nullcontext
+
+import torch
+
+REF_DIR = os.environ.get("KREA2_REF_DIR", "/home/ohiom/database/watering-hole")
+
+
+def load_reference():
+    sys.path.insert(0, REF_DIR)
+    import mmdit
+    # The reference pins the cuDNN SDPA backend; neutralize it so both models use the same
+    # default kernel and the test can run on CPU.
+    mmdit.sdpa_kernel = lambda *a, **k: nullcontext()
+    return mmdit
+
+
+def load_port():
+    path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "pipelines", "krea2", "transformer_krea2.py"))
+    spec = importlib.util.spec_from_file_location("transformer_krea2", path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def main():
+    mmdit = load_reference()
+    port = load_port()
+
+    cfg = dict(
+        features=128, tdim=32, txtdim=64, heads=4, kvheads=2, multiplier=4,
+        layers=2, patch=2, channels=4, bias=False, theta=1e3,
+        txtlayers=3, txtheads=2, txtkvheads=2,
+    )
+
+    torch.manual_seed(0)
+    ref = mmdit.SingleStreamDiT(mmdit.SingleMMDiTConfig(**cfg)).float().eval()
+    mine = port.Krea2Transformer2DModel(**cfg).float().eval()
+    missing, unexpected = mine.load_state_dict(ref.state_dict(), strict=False)
+    assert not missing, f"missing keys when loading reference weights: {missing}"
+    assert not unexpected, f"unexpected keys when loading reference weights: {unexpected}"
+
+    batch, txtlen, imglen = 2, 5, 9
+    cdim = cfg["channels"] * cfg["patch"] ** 2
+    seq = txtlen + imglen
+    gen = torch.Generator().manual_seed(1)
+    img = torch.randn(batch, imglen, cdim, generator=gen)
+    context = torch.randn(batch, txtlen, cfg["txtlayers"], cfg["txtdim"], generator=gen)
+    timestep = torch.rand(batch, generator=gen)
+    pos = torch.randint(0, 16, (batch, seq, 3), generator=gen).float()
+    mask = torch.ones(batch, seq, dtype=torch.bool)
+    mask[0, -2:] = False  # exercise the key-padding path
+
+    with torch.no_grad():
+        out_ref = ref(img, context, timestep, pos, mask)
+        out_mine = mine(
+            hidden_states=img,
+            encoder_hidden_states=context,
+            timestep=timestep,
+            position_ids=pos,
+            attention_mask=mask,
+            return_dict=False,
+        )[0]
+
+    assert out_ref.shape == out_mine.shape, f"shape mismatch: {out_ref.shape} vs {out_mine.shape}"
+    diff = (out_ref - out_mine).abs().max().item()
+    rel = diff / (out_ref.abs().max().item() + 1e-8)
+    print(f"output shape: {tuple(out_mine.shape)}")
+    print(f"max abs diff: {diff:.3e}   max rel diff: {rel:.3e}")
+    tol = 1e-4
+    assert diff < tol, f"PARITY FAILED: max abs diff {diff:.3e} >= {tol}"
+    print("PARITY OK")
+
+
+if __name__ == "__main__":
+    main()