Skip to content
 
 

Repository files navigation

Krea-2 Depth ControlNet-LoRA — Forge Classic 2.28.1 fork

Forge Classic 2.28.1 compatibility

Why this fork?

The original Krea 2 depth-control project is a standalone inference and training toolkit, not a Stable Diffusion WebUI extension. Cloning it into extensions therefore did not add a Forge panel or connect it to Forge's model lifecycle, low-VRAM loading, Hires fix, Krea Edit references, or normal generation flow.

I found the depth-control tool too useful to leave outside WebUI, so I turned it into a native integration for sd-webui-forge-classic 2.28.1 and tested it through Forge's real generation lifecycle, with a little help from AI. I hope this fork saves other Forge users some time and proves useful to the community.

This fork adds a native, always-visible Forge panel while keeping the original standalone scripts and training code available. It has been tested on Windows, Python 3.13, Gradio 4.40 and a Krea 2 checkpoint using the Qwen/Krea image VAE.

Install in Forge Classic 2.28.1

Clone this fork into Forge's extensions directory, then restart Forge:

git clone -b forge-classic-2.28.1 \
  https://github.com/fabiencomte/Krea-2-controlnet.git

Open Krea 2 Depth ControlNet-LoRA below txt2img or img2img. Click Download / verify model (862 MB) once, add an image, keep depth_anything_v2 selected, and enable the accordion before Generate. The official checkpoint is stored outside Git at:

models/ControlNet/Krea2/depth-control-lora.safetensors

The Depth Anything V2 preprocessor downloads its own model through Forge on first use. You can also choose None (already a depth map) to supply a map directly. White means near by default; use Invert depth map when the input uses the opposite convention.

Why this Forge integration is safe

  • It clones Forge's current UnetPatcher for each sampling pass. The base model and Forge core files are not changed.
  • It keeps Forge's real Krea input layer registered, so cold-load and low-VRAM moves remain under Forge's control. The depth checkpoint intentionally replaces that layer's output for the noisy image tokens, matching its official inference code; a separate LoRA aimed specifically at first.weight therefore does not combine with depth control on that one call.
  • During each denoising call, a temporary one-shot hook applies the checkpoint's complete trained input projection (first.weight and first.bias) to the noisy image tokens only. Krea Edit reference tokens keep their normal path.
  • All 224 expected LoRA pairs are shape-checked before sampling. A partial or incompatible checkpoint is rejected instead of being applied silently.
  • A control error installs a fail-closed sampling guard. Forge cannot quietly continue with an uncontrolled image.
  • Generate, Skip and Interrupt remain entirely owned by Forge. The temporary hook is always removed in finally, including after an interruption.
  • The WebUI API accepts the same base64 image format as /sdapi/v1/txt2img. Its booleans, resolution and strength are validated strictly, so an ambiguous or out-of-range request fails closed instead of silently changing meaning.
  • The model download is pinned to the official file revision and checked against SHA-256 fb80547ed79b47c1e3fea7bb9d36297e3917b2115fab6700ca1501350f9f483c. Normal generation verifies the same hash before loading, while the download button reuses a valid local file and replaces a corrupt one.
  • Hires fix rebuilds the control map at the second pass's real dimensions, so a changed aspect ratio is not first cropped square and then stretched.

Verified cases

  • 41 automated tests: official checkpoint layout, completeness and tensor shapes, full projection weight/bias, CFG batch repetition, non-square and changed-ratio Hires resize, image/API contracts, previous-wrapper multi-call composition, one-shot Krea Edit behavior, low-VRAM dtype casts, strict API controls, pinned local/download integrity, and cleanup after an interrupt-like exception;
  • real Krea 2 generation at 128×128 and an 8-step 512×512 generation using Depth Anything V2;
  • deterministic enabled/disabled comparison proving that control changes the denoising result;
  • Skip during batch 1/2, continuation and successful batch 2/2;
  • Interrupt during sampling followed by a clean controlled generation;
  • Ruff, compileall, git diff --check, UI load and API infotext metadata.

Run the tests from this repository:

python -m pytest -q tests
python -m ruff check forge_krea2_depth scripts tests
python -m compileall -q forge_krea2_depth scripts tests

API example

Use the always-on script name Krea 2 Depth ControlNet-LoRA and these six arguments:

{
  "alwayson_scripts": {
    "Krea 2 Depth ControlNet-LoRA": {
      "args": [true, "data:image/png;base64,...", "depth_anything_v2", 768, false, 1.0]
    }
  }
}

For API calls, keep the two checkbox values as JSON booleans, the resolution as an integer from 256 to 2048, and the strength as a finite number from 0 to 2.

The upstream repository currently contains no code license file. This fork does not invent or change one. Model weights remain subject to the Krea 2 community license.


Original standalone project

Depth-conditioned generation for Krea-2. 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.

  • Trained on Krea-2-Raw, works on both Raw and Krea-2-Turbo (8-step)
  • Single 862MB LoRA file (rank 64 + expanded input projection), base stays frozen
  • Depth consistency (Pearson corr. of input depth vs. depth of generated image): 0.98 with no prompt, 0.99 with prompts

Each strip: init image → extracted depth → generated output.

Examples

robot example

cat to tiger example

Checkpoint

file base trained on size
depth-control-lora.safetensors krea/Krea-2-Raw 862MB

Setup

git clone https://github.com/Tanmaypatil123/Krea-2-controlnet.git
cd Krea-2-controlnet
pip install -r requirements.txt

hf download Patil/Krea-2-depth-controlnet depth-control-lora.safetensors --local-dir .

Inference

# Turbo base — fast, recommended (8 steps, no CFG)
python inference.py photo.jpg -p "a futuristic spaceship interior, cinematic lighting" \
    --lora depth-control-lora.safetensors

# Raw base — undistilled (28-52 steps, CFG 3.5)
python inference.py photo.jpg -p "..." --lora depth-control-lora.safetensors \
    --base raw

# No prompt: the depth map is the only signal
python inference.py photo.jpg --lora depth-control-lora.safetensors --save-strip

# Weaker structure adherence (more creative freedom)
python inference.py photo.jpg -p "..." --lora depth-control-lora.safetensors --lora-scale 0.6
flag default notes
-p / --prompt "" empty = depth-only generation
--base turbo turbo or raw
--steps 8 turbo / 28 raw
--cfg 0 turbo / 3.5 raw classifier-free guidance
--mu 1.15 turbo / auto raw timestep shift
--lora-scale 1.0 control-strength dial
--seed 0
--save-strip off also saves input|depth|output comparison

Python API

from PIL import Image
from huggingface_hub import hf_hub_download
from pipeline import DepthLoRAPipeline

base = hf_hub_download("krea/Krea-2-Turbo", "turbo.safetensors")
pipe = DepthLoRAPipeline(base, "depth-control-lora.safetensors")

out, depth = pipe(Image.open("photo.jpg"),
                  prompt="a cozy cabin interior at dusk",
                  steps=8, cfg=0.0, mu=1.15, seed=0)
out.save("output.png")

How it works (inference path)

  1. The init image is resized to the nearest ~1MP aspect bucket and run through Depth-Anything-V2-Large → inverse depth map (near = white).
  2. The depth map is encoded with the same Qwen-Image VAE the model uses for images, so control lives in latent space.
  3. At every denoising step, the depth latent is concatenated channel-wise to the noisy latent (each DiT token: 64 → 128 dims). The expanded input projection + rank-64 LoRA on all 28 blocks (both included in the checkpoint) steer generation to follow the depth structure.
  4. Standard Krea-2 flow-matching Euler sampling otherwise — same recipe as BFL's Flux.1-Depth-dev-lora.

Tips & limitations

  • Best inputs: photos / renders with real perspective. Flat 2D illustrations produce nearly-uniform depth maps, so control will be weak (garbage in, garbage out).
  • Empty-prompt generation works (0.98 depth consistency) — useful for testing how much structure the control alone carries.
  • --lora-scale below 1.0 relaxes structure adherence; above 1.0 tightens it at some quality cost.
  • Krea-2-Raw generates up to ~1K resolution; outputs are capped at the ~1MP buckets.

Training your own ControlNet-LoRA

The training code is in trainer/ and is control-type agnostic: the same recipe trains depth, canny, tile, gray, or any custom pixel-aligned control signal. See trainer/README.md for data preparation (local folder or HF dataset), the shard format, and training instructions.

# example: canny ControlNet-LoRA from a folder of captioned images
python trainer/prepare_data.py --source folder --input-dir ./my_images \
    --out-dir ./data --control-type canny
python trainer/train_control_lora.py --data-dir ./data --ckpt-dir ./ckpts \
    --raw-ckpt raw.safetensors --control-type canny

Files

  • inference.py — CLI
  • pipeline.py — full pipeline: LoRA surgery, Qwen3-VL conditioner, VAE, depth estimator, flow sampler with control injection
  • mmdit.py — unmodified DiT definition from the krea-2 repo
  • k2_lora.py — model surgery: expanded input projection + LoRA injection (used by the trainer)
  • trainer/ — data prep + training for any control type (docs)

Model weights are subject to the Krea 2 community license.

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages