Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/dog.jpg filter=lfs diff=lfs merge=lfs -text
+examples/man_beach.jpg filter=lfs diff=lfs merge=lfs -text
+examples/tent.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,27 @@
 ---
-title: Krea 2 Depth Controlnet
-emoji: 🚀
-colorFrom: green
-colorTo: purple
+title: Krea-2 Depth ControlNet
+emoji: 🏔️
+colorFrom: gray
+colorTo: yellow
 sdk: gradio
 sdk_version: 6.19.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Depth-controlled image generation with Krea-2 Turbo
+python_version: "3.12"
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Krea-2 Depth ControlNet-LoRA
+
+Depth-conditioned image generation for [Krea-2](https://huggingface.co/krea/Krea-2-Turbo).
+Give it any image and a prompt — it extracts the depth map with
+**Depth-Anything-V2** and generates a new image with the **same 3D structure and
+composition**, but whatever content and style you ask for.
+
+- Base: [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo) (8-step)
+- Control: [Patil/Krea-2-depth-controlnet](https://huggingface.co/Patil/Krea-2-depth-controlnet)
+  (rank-64 LoRA + expanded input projection)
+- Depth: `depth-anything/Depth-Anything-V2-Large-hf`
+
+Model weights are subject to the
+[Krea 2 community license](https://www.krea.ai/krea-2-licensing).

app.py

@@ -0,0 +1,173 @@
+import os
+
+# Neutralize torch.compile decorators inside mmdit.py (not supported on ZeroGPU
+# forked workers) and reduce allocator fragmentation for the 13B DiT.
+os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import random
+
+import spaces  # noqa: E402  MUST come before torch / CUDA-touching imports
+import torch  # noqa: E402
+import gradio as gr  # noqa: E402
+from huggingface_hub import hf_hub_download  # noqa: E402
+
+from pipeline import DepthLoRAPipeline  # noqa: E402
+
+MAX_SEED = 2**31 - 1
+
+# --------------------------------------------------------------------- loading
+# Turbo base (8-step, no CFG) — the author's recommended fast configuration.
+# Krea-2 base checkpoint (~26GB) + depth-control LoRA + Qwen3-VL-4B text encoder
+# + Qwen-Image VAE + Depth-Anything-V2-Large. Loaded once at module scope so
+# ZeroGPU packs the weights and streams them to VRAM on the first GPU call.
+print("Resolving Krea-2-Turbo base checkpoint...")
+BASE_CKPT = os.path.realpath(hf_hub_download("krea/Krea-2-Turbo", "turbo.safetensors"))
+LORA_CKPT = os.path.realpath(
+    hf_hub_download("Patil/Krea-2-depth-controlnet", "depth-control-lora.safetensors")
+)
+
+print("Building DepthLoRAPipeline (13B DiT + Qwen3-VL-4B + VAE + DepthAnything)...")
+pipe = DepthLoRAPipeline(BASE_CKPT, LORA_CKPT, device="cuda")
+print("Pipeline ready.")
+
+
+# --------------------------------------------------------------------- inference
+@spaces.GPU(duration=120)
+def generate(
+    image,
+    prompt: str = "",
+    steps: int = 8,
+    lora_scale: float = 1.0,
+    seed: int = 0,
+    randomize_seed: bool = True,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Generate a new image that keeps the 3D structure of an input image.
+
+    Extracts a depth map from the input image with Depth-Anything-V2 and
+    generates a new image following the same depth/composition but with the
+    content and style described by the prompt (Krea-2-Turbo, 8-step).
+
+    Args:
+        image: The input image whose depth/structure is preserved.
+        prompt: What to generate. Leave empty for depth-only generation.
+        steps: Number of sampling steps (8 recommended for Turbo).
+        lora_scale: Control strength. <1.0 relaxes structure adherence.
+        seed: RNG seed for reproducibility.
+        randomize_seed: If True, pick a random seed each run.
+
+    Returns:
+        A tuple of (generated image, extracted depth map, used seed).
+    """
+    if image is None:
+        raise gr.Error("Please provide an input image.")
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed)
+
+    # Turbo config: cfg=0.0, mu=1.15. lora_scale is applied to the loaded LoRA
+    # layers in place (they were built with scale=1.0), so scale the effective
+    # weight by mutating each LoRALinear's scale before the run.
+    from pipeline import LoRALinear
+
+    for module in pipe.model.modules():
+        if isinstance(module, LoRALinear):
+            module.scale = (64 / 64) * float(lora_scale)
+
+    out, depth = pipe(
+        image,
+        prompt=prompt or "",
+        steps=int(steps),
+        cfg=0.0,
+        mu=1.15,
+        seed=seed,
+    )
+    return out, depth, seed
+
+
+# --------------------------------------------------------------------- UI
+CSS = """
+#col-container { max-width: 1200px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            """
+            # Krea-2 Depth ControlNet-LoRA
+
+            Give it any image and a prompt — it extracts the depth map with
+            **Depth-Anything-V2** and generates a new image with the **same 3D
+            structure and composition**, but whatever content and style you ask
+            for. Powered by [Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo)
+            (8-step) + the
+            [depth-control LoRA](https://huggingface.co/Patil/Krea-2-depth-controlnet).
+
+            *Best with photos / renders that have real perspective. Flat 2D
+            illustrations give weak control. Empty prompt = depth-only generation.*
+            """
+        )
+        with gr.Row():
+            with gr.Column():
+                image = gr.Image(label="Input image (depth source)", type="pil")
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    placeholder="a futuristic spaceship interior, cinematic lighting",
+                    lines=2,
+                )
+                run = gr.Button("Generate", variant="primary")
+                with gr.Accordion("Advanced settings", open=False):
+                    steps = gr.Slider(
+                        4, 16, value=8, step=1, label="Sampling steps"
+                    )
+                    lora_scale = gr.Slider(
+                        0.3,
+                        1.4,
+                        value=1.0,
+                        step=0.05,
+                        label="Control strength (LoRA scale)",
+                    )
+                    randomize_seed = gr.Checkbox(
+                        label="Randomize seed", value=True
+                    )
+                    seed = gr.Slider(
+                        0, MAX_SEED, value=0, step=1, label="Seed"
+                    )
+            with gr.Column():
+                output = gr.Image(label="Generated image")
+                depth_out = gr.Image(label="Extracted depth map")
+
+        gr.Examples(
+            examples=[
+                ["examples/dog.jpg", "a majestic lion, golden hour, photorealistic"],
+                [
+                    "examples/landscape.jpg",
+                    "an alien planet landscape, purple sky, sci-fi",
+                ],
+                [
+                    "examples/man_beach.jpg",
+                    "an astronaut on the moon, cinematic lighting",
+                ],
+                [
+                    "examples/tent.jpg",
+                    "a cozy cabin in a snowy forest at dusk",
+                ],
+            ],
+            inputs=[image, prompt],
+            outputs=[output, depth_out, seed],
+            fn=generate,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+    run.click(
+        fn=generate,
+        inputs=[image, prompt, steps, lora_scale, seed, randomize_seed],
+        outputs=[output, depth_out, seed],
+        api_name="generate",
+    )
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

mmdit.py

@@ -0,0 +1,417 @@
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Tensor
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+
+def rope(pos: Tensor, dim: int, theta: float = 1e4, ntk: float = 1.0) -> 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 ropeapply(xq: Tensor, xk: Tensor, freqs: Tensor) -> tuple[Tensor, 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 attention(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    mask: Tensor | None = None,
+    scale: float | None = None,
+    gqa: bool = False,
+) -> Tensor:
+    with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
+        x = F.scaled_dot_product_attention(
+            q, k, v, attn_mask=mask, scale=scale, enable_gqa=gqa
+        )
+    return rearrange(x, "B H L D -> B L (H D)")
+
+
+def _mask(mask: Tensor) -> 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)
+
+
+def temb(
+    t: Tensor,
+    dim: int,
+    period: float = 1e4,
+    tfactor: float = 1e3,
+    device: torch.device = None,
+    dtype: torch.dtype = None,
+) -> Tensor:
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(period)
+        * torch.arange(half, dtype=torch.float32, device=device)
+        / half
+    )
+    # t: (B,) -> args: (B, 1, half), so the embedding broadcasts as a per-sample vec.
+    args = (t.float() * tfactor)[:, None, None] * freqs
+    sin, cos = torch.sin(args), torch.cos(args)
+    return torch.cat((cos, sin), dim=-1).to(dtype=dtype)
+
+
+@dataclass
+class SingleMMDiTConfig:
+    features: int
+    tdim: int
+    txtdim: int
+    heads: int
+    multiplier: int
+    layers: int
+    patch: int
+    channels: int
+    bias: bool = False
+    theta: float = 1e3
+    kvheads: int | None = None
+    txtlayers: int = 1
+    txtheads: int = 20
+    txtkvheads: int = 20
+
+
+class SimpleModulation(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.lin = torch.nn.Parameter(torch.zeros(2, dim))
+        self.multiplier = 2
+
+    # vec (b d)
+    def forward(self, vec: Tensor):
+        out = vec + rearrange(self.lin, "two d -> 1 two d")
+        scale, shift = out.chunk(self.multiplier, dim=1)
+        return scale, shift
+
+
+class DoubleSharedModulation(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.lin = torch.nn.Parameter(torch.zeros(6 * dim))
+
+    # vec (b (6 d))
+    def forward(self, vec: Tensor):
+        out = vec + self.lin
+        prescale, preshift, pregate, postscale, postshift, postgate = out.chunk(
+            6, dim=-1
+        )
+        return prescale, preshift, pregate, postscale, postshift, postgate
+
+
+class PositionalEncoding(torch.nn.Module):
+    def __init__(self, dim, axdims: list[int], theta: float = 1e2, ntk: float = 1.0):
+        super().__init__()
+        self.axdims = axdims  # how to split the head dimension across the position axes
+        self.theta = theta
+        self.ntk = ntk
+
+    @torch.compile(fullgraph=True)
+    def forward(self, pos: Tensor) -> Tensor:
+        return torch.cat(
+            [
+                rope(pos[..., i], d, self.theta, self.ntk)
+                for i, d in enumerate(self.axdims)
+            ],
+            dim=-3,
+        )
+
+
+class QKNorm(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.qnorm = RMSNorm(dim)
+        self.knorm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        return self.qnorm(q), self.knorm(k), v
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, features: int, eps: float = 1e-05, device: torch.device = None):
+        super().__init__()
+        self.features = features
+        self.eps = eps
+        self.scale = torch.nn.Parameter(
+            torch.zeros(features, device=device, dtype=torch.float32)
+        )
+
+    @torch.compile(fullgraph=True)
+    def forward(self, x: Tensor) -> Tensor:
+        t, dtype = x.float(), x.dtype
+        t = F.rms_norm(
+            t, (self.features,), eps=self.eps, weight=(self.scale.float() + 1.0)
+        )
+        return t.to(dtype)
+
+
+class SwiGLU(torch.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 = torch.nn.Linear(features, mlpdim, bias=bias)
+        self.up = torch.nn.Linear(features, mlpdim, bias=bias)
+        self.down = torch.nn.Linear(mlpdim, features, bias=bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class Attention(torch.nn.Module):
+    def __init__(self, dim: int, heads: int, kvheads: int = 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.wq = torch.nn.Linear(dim, self.headdim * self.heads, bias=bias)
+        self.wk = torch.nn.Linear(dim, self.headdim * self.kvheads, bias=bias)
+        self.wv = torch.nn.Linear(dim, self.headdim * self.kvheads, bias=bias)
+        self.gate = torch.nn.Linear(dim, dim, bias=bias)
+        self.qknorm = QKNorm(self.headdim)
+        self.gqa = self.heads != self.kvheads
+        self.wo = torch.nn.Linear(dim, dim, bias=bias)
+
+    def forward(
+        self, qkv: Tensor, freqs: Tensor | None = None, mask: Tensor | None = None
+    ) -> Tensor:
+        q, k, v, gate = self.wq(qkv), self.wk(qkv), self.wv(qkv), self.gate(qkv)
+
+        q, k, v = (
+            rearrange(q, "B L (H D) -> B H L D", H=self.heads),
+            rearrange(k, "B L (H D) -> B H L D", H=self.kvheads),
+            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 = ropeapply(q, k, freqs)
+        out = self.wo(attention(q, k, v, mask=mask, gqa=self.gqa) * F.sigmoid(gate))
+
+        return out
+
+
+class LastLayer(torch.nn.Module):
+    def __init__(self, features: int, patch: int, channels: int):
+        super().__init__()
+        self.norm = RMSNorm(features)
+        self.linear = torch.nn.Linear(features, patch * patch * channels, bias=True)
+        self.modulation = SimpleModulation(features)
+
+    @torch.compile(fullgraph=True)
+    def forward(self, x: Tensor, tvec: Tensor) -> Tensor:
+        scale, shift = self.modulation(tvec)
+        x = (1 + scale) * self.norm(x) + shift
+        x = self.linear(x)
+        return x
+
+
+class TextFusionBlock(torch.nn.Module):
+    def __init__(
+        self,
+        features: int,
+        heads: int,
+        multiplier: int,
+        bias: bool = False,
+        kvheads: int = 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: Tensor, mask: Tensor | None = None) -> Tensor:
+        x = x + self.attn(self.prenorm(x), mask=mask)
+        x = x + self.mlp(self.postnorm(x))
+
+        return x
+
+
+class TextFusionTransformer(torch.nn.Module):
+    # num_txt_layers is the number of selected encoder hidden-state layers fed in
+    # (projected down to 1), NOT the transformer depth — that's fixed at 2 + 2 blocks.
+    def __init__(
+        self,
+        num_txt_layers: int,
+        txt_dim: int,
+        heads: int,
+        multiplier: int,
+        bias: bool = False,
+        kvheads: int = None,
+    ):
+        super().__init__()
+        self.layerwise_blocks = torch.nn.ModuleList(
+            [
+                TextFusionBlock(txt_dim, heads, multiplier, bias, kvheads)
+                for _ in range(2)
+            ]
+        )
+        self.projector = torch.nn.Linear(num_txt_layers, 1, bias=False)
+        self.refiner_blocks = torch.nn.ModuleList(
+            [
+                TextFusionBlock(txt_dim, heads, multiplier, bias, kvheads)
+                for _ in range(2)
+            ]
+        )
+
+    def forward(self, x: Tensor, mask: Tensor | None = None) -> 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)
+        x = 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,
+    ):
+        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: Tensor, vec: Tensor, freqs: Tensor, mask: Tensor | None = None
+    ) -> 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 SingleStreamDiT(nn.Module):
+    def __init__(self, config: SingleMMDiTConfig):
+        super().__init__()
+        self.config = config
+
+        headdim = config.features // config.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(
+            config.features, axes, theta=config.theta, ntk=1.0
+        )
+        self.first = nn.Linear(
+            config.channels * config.patch**2, config.features, bias=True
+        )
+
+        self.blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(
+                    config.features,
+                    config.heads,
+                    config.multiplier,
+                    config.bias,
+                    config.kvheads,
+                )
+                for _ in range(config.layers)
+            ]
+        )
+        self.tmlp = nn.Sequential(
+            nn.Linear(config.tdim, config.features),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(config.features, config.features),
+        )
+        self.txtfusion = TextFusionTransformer(
+            config.txtlayers,
+            config.txtdim,
+            config.txtheads,
+            config.multiplier,
+            config.bias,
+            config.txtkvheads,
+        )
+        self.txtmlp = nn.Sequential(
+            RMSNorm(config.txtdim),
+            nn.Linear(config.txtdim, config.features),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(config.features, config.features),
+        )
+        self.last = LastLayer(config.features, config.patch, config.channels)
+
+        self.tproj = nn.Sequential(
+            nn.GELU(approximate="tanh"), nn.Linear(config.features, config.features * 6)
+        )
+
+    def forward(
+        self,
+        img: Tensor,
+        context: Tensor,
+        t: Tensor,
+        pos: Tensor,
+        mask: Tensor | None = None,
+    ) -> Tensor:
+        img = self.first(img)
+        t = self.tmlp(temb(t, self.config.tdim, device=img.device, dtype=img.dtype))
+        tvec = self.tproj(t)
+
+        txtmask = _mask(mask[:, : context.shape[1]])
+
+        context = self.txtfusion(context, mask=txtmask)
+        context = self.txtmlp(context)
+
+        txtlen, imglen = context.shape[1], img.shape[1]
+        combined = torch.cat((context, img), dim=1)
+
+        # Pad combined sequence to a multiple of 256 to stabilize compiled kernel shapes.
+        fulllen = combined.shape[1]
+        _padlen = (-fulllen) % 256
+        if _padlen > 0:
+            combined = F.pad(combined, (0, 0, 0, _padlen))
+            mask = F.pad(mask, (0, _padlen), value=False)
+            pos = F.pad(pos, (0, 0, 0, _padlen))
+
+        mask = _mask(mask)
+
+        freqs = self.posemb(pos)
+
+        for block in self.blocks:
+            combined = block(combined, tvec, freqs, mask)
+
+        final = self.last(combined, t)
+        output = final[:, txtlen : txtlen + imglen, :]
+
+        return output

pipeline.py

@@ -0,0 +1,313 @@
+"""Standalone depth-ControlNet-LoRA pipeline for Krea-2.
+
+Everything needed for inference: LoRA surgery, text conditioning, VAE,
+depth estimation, and the flow-matching sampler with control injection.
+`mmdit.py` is the unmodified DiT definition from the official krea-2 repo.
+"""
+
+import math
+import os
+
+os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from PIL import Image
+from safetensors.torch import load_file
+
+from mmdit import SingleMMDiTConfig, SingleStreamDiT, _mask, temb
+
+K2_CONFIG = SingleMMDiTConfig(
+    features=6144, tdim=256, txtdim=2560, heads=48, kvheads=12, multiplier=4,
+    layers=28, patch=2, channels=16, txtheads=20, txtkvheads=20, txtlayers=12,
+)
+
+BASE_CKPTS = {"raw": ("krea/Krea-2-Raw", "raw.safetensors"),
+              "turbo": ("krea/Krea-2-Turbo", "turbo.safetensors")}
+
+LORA_TARGETS = ("attn.wq", "attn.wk", "attn.wv", "attn.wo", "attn.gate",
+                "mlp.gate", "mlp.up", "mlp.down")
+
+# timestep-shift interpolation endpoints (image seq len -> mu), from the
+# scheduler config Krea-2 was trained with
+MU_X1, MU_Y1, MU_X2, MU_Y2 = 256, 0.5, 6400, 1.15
+
+BUCKETS = [(1024, 1024), (896, 1152), (1152, 896), (832, 1216), (1216, 832),
+           (768, 1344), (1344, 768), (704, 1472), (1472, 704)]
+
+
+# ---------------------------------------------------------------- model surgery
+
+class LoRALinear(nn.Module):
+    """y = Wx + scale * B(Ax). A: rank x in, B: out x rank."""
+
+    def __init__(self, base: nn.Linear, rank: int, alpha: float, scale: float = 1.0):
+        super().__init__()
+        self.base = base
+        self.scale = (alpha / rank) * scale
+        self.A = nn.Parameter(torch.zeros(rank, base.in_features, dtype=torch.float32))
+        self.B = nn.Parameter(torch.zeros(base.out_features, rank, dtype=torch.float32))
+
+    def forward(self, x):
+        lora = (x @ self.A.T.to(x.dtype)) @ self.B.T.to(x.dtype)
+        return self.base(x) + lora * self.scale
+
+
+class ControlInputLayer(nn.Module):
+    """Replaces the DiT input projection: in_features doubled (64 -> 128) to
+    accept [noisy latent patches ; depth latent patches] concatenated on the
+    channel dim. Trained weights are loaded from the LoRA checkpoint."""
+
+    def __init__(self, pretrained: nn.Linear):
+        super().__init__()
+        in_f, out_f = pretrained.in_features, pretrained.out_features
+        self.weight = nn.Parameter(torch.zeros(out_f, in_f * 2, dtype=torch.float32))
+        self.bias = nn.Parameter(pretrained.bias.detach().float().clone())
+        with torch.no_grad():
+            self.weight[:, :in_f] = pretrained.weight.detach().float()
+
+    def forward(self, x):
+        return F.linear(x, self.weight.to(x.dtype), self.bias.to(x.dtype))
+
+
+def _get(root, path):
+    for p in path.split("."):
+        root = getattr(root, p)
+    return root
+
+
+def _set(root, path, new):
+    parts = path.split(".")
+    setattr(_get(root, ".".join(parts[:-1])) if len(parts) > 1 else root,
+            parts[-1], new)
+
+
+def build_model(base_ckpt: str, lora_ckpt: str, rank: int = 64,
+                lora_scale: float = 1.0, device: str = "cuda",
+                dtype: torch.dtype = torch.bfloat16) -> SingleStreamDiT:
+    with torch.device("meta"):
+        model = SingleStreamDiT(K2_CONFIG)
+    model.load_state_dict(load_file(base_ckpt), strict=True, assign=True)
+    model = model.to(device=device, dtype=dtype).requires_grad_(False)
+
+    model.first = ControlInputLayer(model.first).to(device)
+    for i in range(K2_CONFIG.layers):
+        for t in LORA_TARGETS:
+            path = f"blocks.{i}.{t}"
+            _set(model, path, LoRALinear(_get(model, path), rank, rank,
+                                         lora_scale).to(device))
+
+    sd = load_file(lora_ckpt)
+    missing, unexpected = model.load_state_dict(sd, strict=False)
+    assert not unexpected, f"unexpected keys: {unexpected[:5]}"
+    return model.eval()
+
+
+# ---------------------------------------------------------------- conditioning
+
+class TextConditioner(nn.Module):
+    """Qwen3-VL-4B encoder exactly as Krea-2 uses it: hidden states from 12
+    selected layers, stacked, with the chat-template prefix sliced off."""
+
+    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"
+    )
+    SUFFIX = "<|im_end|>\n<|im_start|>assistant\n"
+    PREFIX_IDX = 34
+    SELECT_LAYERS = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35)
+
+    def __init__(self, model_id="Qwen/Qwen3-VL-4B-Instruct", max_length=512,
+                 device="cuda", dtype=torch.bfloat16):
+        super().__init__()
+        from transformers import AutoTokenizer, Qwen3VLForConditionalGeneration
+
+        self.qwen = (Qwen3VLForConditionalGeneration
+                     .from_pretrained(model_id, torch_dtype=dtype)
+                     .to(device).eval().requires_grad_(False))
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        self.max_length = max_length
+        self.device = device
+
+    @torch.no_grad()
+    def forward(self, prompts):
+        text = [self.PREFIX + p for p in prompts]
+        inputs = self.tokenizer(
+            text, truncation=True, padding="longest",
+            max_length=self.max_length + self.PREFIX_IDX,
+            return_tensors="pt", padding_side="right").to(self.device)
+        suffix = self.tokenizer([self.SUFFIX] * len(prompts),
+                                return_tensors="pt").to(self.device)
+        ids = torch.cat([inputs["input_ids"], suffix["input_ids"]], dim=1)
+        mask = torch.cat([inputs["attention_mask"].bool(),
+                          suffix["attention_mask"].bool()], dim=1)
+        states = self.qwen(input_ids=ids, attention_mask=mask,
+                           output_hidden_states=True)
+        hiddens = torch.stack([states.hidden_states[i]
+                               for i in self.SELECT_LAYERS], dim=2)
+        return hiddens[:, self.PREFIX_IDX:], mask[:, self.PREFIX_IDX:]
+
+
+class VAE(nn.Module):
+    """Qwen-Image VAE (f8, 16ch) with Krea-2's latent normalization."""
+
+    def __init__(self, device="cuda", dtype=torch.bfloat16):
+        super().__init__()
+        from diffusers import AutoencoderKLQwenImage
+
+        self.ae = (AutoencoderKLQwenImage
+                   .from_pretrained("Qwen/Qwen-Image", subfolder="vae",
+                                    torch_dtype=dtype)
+                   .to(device).eval().requires_grad_(False))
+        self.mean = torch.tensor(self.ae.config.latents_mean,
+                                 device=device).view(1, -1, 1, 1, 1)
+        self.std = torch.tensor(self.ae.config.latents_std,
+                                device=device).view(1, -1, 1, 1, 1)
+
+    @torch.no_grad()
+    def encode(self, x):  # (b,3,h,w) in [-1,1] -> (b,16,h/8,w/8) normalized
+        z = self.ae.encode(x.unsqueeze(2)).latent_dist.sample()
+        return ((z - self.mean) / self.std).squeeze(2)
+
+    @torch.no_grad()
+    def decode(self, z):  # normalized latent -> (b,3,h,w) in [-1,1]
+        z = (z.unsqueeze(2) * self.std + self.mean).to(next(self.ae.parameters()).dtype)
+        return rearrange(self.ae.decode(z).sample, "b c 1 h w -> b c h w")
+
+
+class DepthEstimator:
+    """Depth-Anything-V2-Large. Returns inverse depth in [0,1], near = 1."""
+
+    def __init__(self, device="cuda"):
+        from transformers import AutoImageProcessor, AutoModelForDepthEstimation
+
+        mid = "depth-anything/Depth-Anything-V2-Large-hf"
+        self.processor = AutoImageProcessor.from_pretrained(mid)
+        self.model = (AutoModelForDepthEstimation
+                      .from_pretrained(mid, torch_dtype=torch.float16)
+                      .to(device).eval().requires_grad_(False))
+        self.device = device
+
+    @torch.no_grad()
+    def __call__(self, image: Image.Image) -> torch.Tensor:
+        inputs = self.processor(images=[image], return_tensors="pt").to(self.device)
+        d = self.model(**inputs).predicted_depth[None].float()
+        d = F.interpolate(d, size=(image.height, image.width),
+                          mode="bilinear", align_corners=False)[0, 0]
+        return (d - d.min()) / (d.max() - d.min() + 1e-6)
+
+
+# ---------------------------------------------------------------- sampling
+
+def pick_bucket(w, h):
+    ar = math.log(w / h)
+    return min(BUCKETS, key=lambda b: abs(math.log(b[0] / b[1]) - ar))
+
+
+def resize_center_crop(img, tw, th):
+    w, h = img.size
+    s = max(tw / w, th / h)
+    img = img.resize((round(w * s), round(h * s)), Image.LANCZOS)
+    w, h = img.size
+    l, t = (w - tw) // 2, (h - th) // 2
+    return img.crop((l, t, l + tw, t + th))
+
+
+def prepare(img, txtlen, patch, txtmask):
+    """Patchify a latent and build combined text+image position/mask tensors."""
+    b, _, h, w = img.shape
+    h_, w_ = h // patch, w // patch
+    ids = torch.zeros((h_, w_, 3), device=img.device)
+    ids[..., 1] = torch.arange(h_, device=img.device)[:, None]
+    ids[..., 2] = torch.arange(w_, device=img.device)[None, :]
+    pos = repeat(ids, "h w c -> b (h w) c", b=b)
+    imgmask = torch.ones(b, h_ * w_, device=img.device, dtype=torch.bool)
+    img = rearrange(img, "b c (h p) (w q) -> b (h w) (c p q)", p=patch, q=patch)
+    txtpos = torch.zeros(b, txtlen, 3, device=img.device)
+    return img, torch.cat((txtpos, pos), 1), torch.cat((txtmask, imgmask), 1)
+
+
+def timesteps(seq_len, steps, mu=None):
+    """Resolution-shifted flow schedule, t: 1 -> 0."""
+    ts = torch.linspace(1, 0, steps + 1)
+    if mu is None:
+        slope = (MU_Y2 - MU_Y1) / (MU_X2 - MU_X1)
+        mu = slope * seq_len + (MU_Y1 - slope * MU_X1)
+    return (math.exp(mu) / (math.exp(mu) + (1.0 / ts - 1.0))).tolist()
+
+
+def forward_control(model, img, ctrl, context, t, pos, mask):
+    """DiT forward with the depth latent concatenated on the channel dim."""
+    x = model.first(torch.cat([img, ctrl], dim=-1))
+    tv = model.tmlp(temb(t, model.config.tdim, device=x.device, dtype=x.dtype))
+    tvec = model.tproj(tv)
+
+    txtmask = _mask(mask[:, : context.shape[1]])
+    context = model.txtmlp(model.txtfusion(context, mask=txtmask))
+
+    txtlen, imglen = context.shape[1], x.shape[1]
+    combined = torch.cat((context, x), dim=1)
+    pad = (-combined.shape[1]) % 256
+    if pad:
+        combined = F.pad(combined, (0, 0, 0, pad))
+        mask = F.pad(mask, (0, pad), value=False)
+        pos = F.pad(pos, (0, 0, 0, pad))
+
+    mask, freqs = _mask(mask), model.posemb(pos)
+    for block in model.blocks:
+        combined = block(combined, tvec, freqs, mask)
+    return model.last(combined, tv)[:, txtlen: txtlen + imglen]
+
+
+class DepthLoRAPipeline:
+    def __init__(self, base_ckpt, lora_ckpt, rank=64, lora_scale=1.0, device="cuda"):
+        self.device = device
+        self.model = build_model(base_ckpt, lora_ckpt, rank, lora_scale, device)
+        self.text = TextConditioner(device=device)
+        self.vae = VAE(device=device)
+        self.depth = DepthEstimator(device=device)
+
+    @torch.no_grad()
+    def __call__(self, image: Image.Image, prompt: str = "", steps: int = 8,
+                 cfg: float = 0.0, mu: float | None = None, seed: int = 0):
+        """Returns (output PIL, depth PIL). Turbo: steps=8 cfg=0 mu=1.15;
+        Raw: steps=28-52 cfg=3.5 mu=None."""
+        bw, bh = pick_bucket(*image.size)
+        image = resize_center_crop(image.convert("RGB"), bw, bh)
+
+        d = self.depth(image)
+        depth_img = Image.fromarray((d.cpu().numpy() * 255).astype(np.uint8))
+        depth_rgb = (d[None, None].repeat(1, 3, 1, 1).to(self.device) * 2 - 1)
+        ctrl_lat = self.vae.encode(depth_rgb.to(torch.bfloat16))
+
+        patch = self.model.config.patch
+        noise = torch.randn(ctrl_lat.shape, device=self.device, dtype=torch.bfloat16,
+                            generator=torch.Generator(self.device).manual_seed(seed))
+
+        txt, tmask = self.text([prompt])
+        x, pos, mask = prepare(noise, txt.shape[1], patch, tmask)
+        ctrl, _, _ = prepare(ctrl_lat.to(torch.bfloat16), txt.shape[1], patch, tmask)
+        if cfg > 0:
+            untxt, unmask_t = self.text([""])
+            _, unpos, unmask = prepare(noise, untxt.shape[1], patch, unmask_t)
+
+        ts = timesteps(x.shape[1], steps, mu)
+        img = x
+        for tc, tp in zip(ts[:-1], ts[1:]):
+            t = torch.full((1,), tc, dtype=img.dtype, device=self.device)
+            v = forward_control(self.model, img, ctrl, txt, t, pos, mask)
+            if cfg > 0:
+                un = forward_control(self.model, img, ctrl, untxt, t, unpos, unmask)
+                v = v + cfg * (v - un)
+            img = img + (tp - tc) * v
+
+        h, w = ctrl_lat.shape[-2:]
+        img = rearrange(img, "b (h w) (c p q) -> b c (h p) (w q)",
+                        p=patch, q=patch, h=h // patch, w=w // patch)
+        px = (self.vae.decode(img).clamp(-1, 1) * 0.5 + 0.5) * 255
+        out = Image.fromarray(px[0].permute(1, 2, 0).float().cpu().byte().numpy())
+        return out, depth_img

requirements.txt

@@ -0,0 +1,9 @@
+torchvision
+transformers>=5.0
+diffusers>=0.36
+accelerate
+safetensors
+einops
+sentencepiece
+pillow
+numpy