Skip to content

About

C++ / CUDA implementation of StreamDiffusion

Topics

Resources

Stars

39 stars

Watchers

0 watching

Forks

Repository files navigation

librediffusion

A C++ / CUDA / TensorRT implementation of StreamDiffusion

Implemented in ossia score

Benchmarks

On a RTX 5090 at 1 step:

SDXL Turbo 1024x1024: stable 26 fps

sdxl

SD Turbo 512x512: stable 96 fps

sdturbo

SDXS: above 600 fps

sdxs

Models need to be converted to TensorRT through the Python script [train-lora.py] beforehand:

$ uv run train-lora.py --model stabilityai/sd-turbo --min-batch 1 --max-batch 1 --opt-batch 1 --min-resolution 512 --max-resolution 1024 --output ./engines-sd-turbo

Engine build recipes (validated)

All commands run as uv run python train-lora.py <args>. Resolution flags are --min-resolution/--max-resolution/--opt-width/--opt-height (512 for SD1.5, 1024 for SDXL); add -l REPO (or -l "REPO|file.safetensors") to fuse a LoRA, repeatable to stack LoRAs.

Batch flags decide how many denoising steps a bundle can run. With denoise-batch on, the UNet runs one row per step, so --max-batch must be >= the step count you intend to use; a static bundle (--min-batch 1 --max-batch 1) is 1-step only, and is ~10-13% faster at batch 1 than a dynamic one. Non-square geometry uses --min-width/--max-width/--min-height/--max-height instead of --min-resolution/--max-resolution.

On Windows: pass --cache-dir on the build volume

uv's cache defaults to %LOCALAPPDATA%\uv\cache on C:. The engine-build venv is ~10 GB of torch/TensorRT wheels, and if you build anywhere other than C: — which you will, because the engines and the HF snapshot need tens of GB — uv cannot hardlink out of its cache across the volume boundary and copies every file instead:

warning: Failed to hardlink files; falling back to full copy. This may lead to degraded performance.
         If the cache and target directories are on different filesystems, hardlinking may not be supported.

So you pay for the wheels twice on disk and wait through the copy. Point the cache at the volume you are building on, and keep passing it on every later uv call in the same tree:

uv sync    --cache-dir D:\tmp\uvcache
uv run     --cache-dir D:\tmp\uvcache python train-lora.py --type sd15 --model stabilityai/sd-turbo ...

set UV_CACHE_DIR=D:\tmp\uvcache does the same thing for a whole shell session. The HF snapshot has its own cache: set HF_HOME=D:\tmp\hf (or hf download --cache-dir D:\tmp\hf\hub).

Koaia's packaged builder passes --cache-dir c:\uv on Windows for both the sync and the build (qml/koaia/Views/ModelView.qml) — a short root-level path, which also stays clear of MAX_PATH on machines without long paths enabled. That cache is on C:, so it only avoids the copy for a venv that is also on C:; if you build on another volume, use a cache on that volume rather than sharing koaia's.

Base models

# SD1.5 family (--type sd15)
train-lora.py --type sd15 --model stabilityai/sd-turbo                 --min-resolution 512 --max-resolution 512  --output engines/sd-turbo
train-lora.py --type sd15 --model SimianLuo/LCM_Dreamshaper_v7         --min-resolution 512 --max-resolution 512  --output engines/lcm-dreamshaper
train-lora.py --type sd15 --model Lykon/dreamshaper-8-lcm              --min-resolution 512 --max-resolution 512  --output engines/dreamshaper-8-lcm
train-lora.py --type sd15 --model IDKiro/sdxs-512-dreamshaper          --min-resolution 512 --max-resolution 512  --output engines/sdxs
train-lora.py --type sd15 --model runwayml/stable-diffusion-v1-5 -l latent-consistency/lcm-lora-sdv1-5            --min-resolution 512 --max-resolution 512 --output engines/sd15-lcm-lora
train-lora.py --type sd15 --model runwayml/stable-diffusion-v1-5 -l "ByteDance/Hyper-SD|Hyper-SD15-4steps-lora.safetensors" --min-resolution 512 --max-resolution 512 --max-batch 4 --output engines/hyper-sd15
train-lora.py --type sd15 --model runwayml/stable-diffusion-v1-5 -l "ByteDance/Hyper-SD|Hyper-SD15-1step-lora.safetensors"  --min-resolution 512 --max-resolution 512 --min-batch 1 --max-batch 1 --opt-batch 1 --output engines/hyper-sd15-1step
train-lora.py --type sd15 --model runwayml/stable-diffusion-v1-5 -l "ByteDance/Hyper-SD|Hyper-SD15-2steps-lora.safetensors" --min-resolution 512 --max-resolution 512 --max-batch 2 --output engines/hyper-sd15-2step
# SD2.1 shares the sd15 path but encodes to 1024 (the bundle manifest carries embedding_dim; the node reads it)
train-lora.py --type sd15 --model sd-research/stable-diffusion-2-1-base --min-resolution 512 --max-resolution 512 --min-batch 1 --max-batch 1 --opt-batch 1 --output engines/sd21-base-1step
train-lora.py --type sd15 --model sd-research/stable-diffusion-2-1-base --min-resolution 512 --max-resolution 512 --max-batch 4 --output engines/sd21-base-4step
# non-square / larger SD-Turbo (one bundle per geometry: the profile is baked in)
train-lora.py --type sd15 --model stabilityai/sd-turbo --min-width 512 --max-width 512 --min-height 768 --max-height 768 --min-batch 1 --max-batch 1 --opt-batch 1 --output engines/sd-turbo-512x768
train-lora.py --type sd15 --model stabilityai/sd-turbo --min-resolution 768 --max-resolution 768 --min-batch 1 --max-batch 1 --opt-batch 1 --output engines/sd-turbo-768

# SDXL family (--type sdxl)
train-lora.py --type sdxl --model stabilityai/sdxl-turbo                            --min-resolution 1024 --max-resolution 1024 --output engines/sdxl-turbo
train-lora.py --type sdxl --model stabilityai/sdxl-turbo                            --min-resolution 512  --max-resolution 512  --min-batch 1 --max-batch 1 --opt-batch 1 --output engines/sdxl-turbo-512
train-lora.py --type sdxl --model stabilityai/stable-diffusion-xl-base-1.0 -l latent-consistency/lcm-lora-sdxl   --min-resolution 1024 --max-resolution 1024 --output engines/sdxl-lcm-lora
train-lora.py --type sdxl --model stabilityai/stable-diffusion-xl-base-1.0 -l "ByteDance/Hyper-SD|Hyper-SDXL-4steps-lora.safetensors" --min-resolution 1024 --max-resolution 1024 --output engines/hyper-sdxl
train-lora.py --type sdxl --model stabilityai/stable-diffusion-xl-base-1.0 -l "ByteDance/SDXL-Lightning|sdxl_lightning_4step_lora.safetensors" --min-resolution 1024 --max-resolution 1024 --output engines/sdxl-lightning
train-lora.py --type sdxl --model segmind/Segmind-Vega -l segmind/Segmind-VegaRT    --min-resolution 1024 --max-resolution 1024 --output engines/vega-rt
# SDXL with a distilled replacement UNet (--unet "REPO|file", a local file or a diffusers dir; loaded strictly)
# DMD2 4-step UNet: the DMD2 for video. img2img at 1-2 steps (Timesteps "32, 45"); its 4 steps are for txt2img
# ("0, 12, 25, 37", its trained 999/749/499/249), and need --max-batch 4 with denoise-batch. CC-BY-NC-4.0.
train-lora.py --type sdxl --model stabilityai/stable-diffusion-xl-base-1.0 --unet "tianweiy/DMD2|dmd2_sdxl_4step_unet_fp16.bin" --min-resolution 1024 --max-resolution 1024 --max-batch 4 --opt-batch 4 --output engines/dmd2-sdxl-4step
# DMD2 1-step UNet: txt2img only (trained from pure noise at t=399; img2img input comes out blurred)
train-lora.py --type sdxl --model stabilityai/stable-diffusion-xl-base-1.0 --unet "tianweiy/DMD2|dmd2_sdxl_1step_unet_fp16.bin" --min-resolution 1024 --max-resolution 1024 --min-batch 1 --max-batch 1 --opt-batch 1 --output engines/dmd2-sdxl-1step
# NitroFusion: the UNet was distilled on a shifted timestep; --unet-timestep-scale records it in bundle.json for the host
train-lora.py --type sdxl --model stabilityai/stable-diffusion-xl-base-1.0 --unet "ChenDY/NitroFusion|nitrosd-realism_unet.safetensors" --unet-timestep-scale 0.25 --min-resolution 1024 --max-resolution 1024 --max-batch 2 --output engines/nitrosd-realism
train-lora.py --type sdxl --model stabilityai/stable-diffusion-xl-base-1.0 --unet "ChenDY/NitroFusion|nitrosd-vibrant_unet.safetensors" --unet-timestep-scale 0.5  --min-resolution 1024 --max-resolution 1024 --max-batch 2 --output engines/nitrosd-vibrant

# img2img-turbo (GaParmar pix2pix-turbo skip-VAE)
train-lora.py --type img2img-turbo --model edge_to_image --min-resolution 512 --max-resolution 512 --output engines/img2img-turbo
train-lora.py --type img2img-turbo --model sketch_to_image_stochastic --min-resolution 512 --max-resolution 512 --output engines/img2img-turbo-sketch

FLUX.2-klein-4B (--type klein)

Download the weights first; the single-file .safetensors is a second copy of the same transformer that this export does not read, so excluding it saves 7.3 GB of the 23 GB:

hf download black-forest-labs/FLUX.2-klein-4B --exclude flux-2-klein-4b.safetensors
train-lora.py --type klein --model black-forest-labs/FLUX.2-klein-4B --output engines/klein
# klein at a non-square geometry needs its own VAE pair; transformer/qwen/rife are geometry-agnostic and can be reused
train-lora.py --type klein --model black-forest-labs/FLUX.2-klein-4B --opt-width 576 --opt-height 320 --output engines/klein-landscape
# klein + a LoRA (fused into a transformer-only copy of the model next to the bundle, then exported
# from it, FP8 calibration included). SANA-Sprint makes 1 step usable (Timesteps "1.0"):
train-lora.py --type klein --model black-forest-labs/FLUX.2-klein-4B -l radames/FLUX.2-klein-Sana-Sprint --output engines/klein-sprint

One command, same as the models above: it runs the reference pipeline once to get the FP8 calibration activations and the VAE normalisation constants, exports the three ONNX graphs, patches them (fix_qwen_complex for the Qwen COMPLEX128 casts, fix_klein_dynamic_seq for the baked sequence length), builds the bf16 and FP8 transformers, the Qwen encoder and both VAE halves, builds RIFE from the vendored ONNX, and stages a bundle the node loads directly. It refuses to finish if any piece is missing. Nothing outside this repo and the HF snapshot is needed; --model also takes a local snapshot directory.

Expect ~18 GB of engines (the bf16 transformer alone is 7.8 GB) plus the ONNX intermediates, and a long build — the transformer and Qwen engines dominate. Steps are idempotent enough to resume: a rerun reuses an existing calibration directory.

  • --klein-quality picks which transformer the bundle gets: speed the FP8 one, quality the bf16 one, both (default) both. Only staging is selective — the bf16 transformer is built either way, so speed costs the same as both minus the FP8 calibration. FP8 needs an Ada (SM 8.9) or newer GPU — TensorRT refuses to build FP8 Q/DQ engines on Ampere and earlier (a 3090 is SM 8.6, a 4090 is 8.9). On such a card the default quietly drops to the bf16 transformer; pass --klein-quality quality to say so explicitly, or speed/both to make the mismatch an error instead.
  • Resolution defaults to klein's native 320x576. --opt-width/--opt-height (multiples of 16) build another geometry; cost scales with (W/16)*(H/16) tokens, so 1024x1024 is ~5.7x the 320x576 work.
  • --hw-compat ampere_plus makes the engines portable across SM 8.0+ instead of locked to the build GPU.

Frame interpolation (RIFE)

The node's Interpolation exp control runs RIFE between consecutive diffused frames (2^exp displayed frames per diffused one). klein bundles always ship rife_ifnet_fp16.plan. RIFE works on decoded RGB frames, so one engine serves every model: add --rife to an SD / SDXL / img2img-turbo export to put it in that bundle, or build it once into the folder that holds the bundles (the node looks in the bundle folder, then in its parent):

python src/streamdiffusion/klein/build_rife.py src/streamdiffusion/klein/assets/rife_ifnet_fp16.onnx engines/rife_ifnet_fp16.plan

The engine covers 64..1024 px per side (KLEIN_HW_COMPAT=ampere_plus for a portable one). RIFE runs after each diffusion on the same producer thread, so it lowers the diffused rate; at 1024 on a 4090 it costs ~7-13 ms per frame at exp 1 and ~21-26 ms at exp 2 (validation_cpp/sd_video_harness --rife PLAN --rife-exp N measures both):

SDXL 1024, img2img diffused fps displayed fps
sdxl-turbo 1 step, exp 1 12.6 25
sdxl-turbo 2 steps, exp 2 7.8 31
DMD2 2 steps, exp 2 8.0 32
DMD2 4 steps, exp 2 4.8 19

ControlNet (--controlnet REPO)

The ControlNet must match the base UNet's block/residual count: standard SD1.5/SDXL ControlNets for standard models, and the model-native ControlNet for reduced architectures (e.g. sdxs).

# SD1.5 (12+mid residuals) — works on any standard SD1.5 base (swap --model/-l)
train-lora.py --type sd15 --model SimianLuo/LCM_Dreamshaper_v7 --controlnet lllyasviel/control_v11p_sd15_canny --min-resolution 512 --max-resolution 512 --output engines/lcm-dreamshaper-canny
# other SD1.5 ControlNets: control_v11f1p_sd15_depth, control_v11p_sd15_openpose, control_v11p_sd15_softedge

# SDXL (works on any SDXL base, incl. LoRA-fused)
train-lora.py --type sdxl --model stabilityai/stable-diffusion-xl-base-1.0 -l latent-consistency/lcm-lora-sdxl --controlnet diffusers/controlnet-canny-sdxl-1.0 --min-resolution 1024 --max-resolution 1024 --output engines/sdxl-lcm-lora-canny

# sdxs needs its NATIVE sketch ControlNet (7 residuals) — generic SD1.5 canny (13) is incompatible
train-lora.py --type sd15 --model IDKiro/sdxs-512-dreamshaper --controlnet IDKiro/sdxs-512-dreamshaper-sketch --min-resolution 512 --max-resolution 512 --output engines/sdxs-sketch

IP-Adapter (--ipadapter CKPT)

--ipadapter produces a COMPLETE bundle: the IP cross-attention baked into unet.engine and the on-device image side (clip_image_encoder.engine + ip_image_proj.engine) that turns the node's Control/Style texture into IP tokens. Without the image side the node silently falls back to externally-fed tokens and the style image appears to do nothing.

# SD1.5 (h94 ip-adapter_sd15.bin) — works on standard SD1.5 bases (NOT sdxs, whose attention differs)
train-lora.py --type sd15 --model SimianLuo/LCM_Dreamshaper_v7 --ipadapter /path/to/h94/IP-Adapter/models/ip-adapter_sd15.bin --min-resolution 512 --max-resolution 512 --output engines/lcm-dreamshaper-ip
# SDXL (h94 sdxl_models/ip-adapter_sdxl_vit-h.bin)
train-lora.py --type sdxl --model stabilityai/sdxl-turbo --ipadapter /path/to/h94/IP-Adapter/sdxl_models/ip-adapter_sdxl_vit-h.bin --min-resolution 1024 --max-resolution 1024 --output engines/sdxl-turbo-ip

Merged / stacked LoRAs (-l repeated)

Any base + accel LoRA + arbitrary style LoRA(s); LoRAs are fused at export time.

# accel (LCM) + a style LoRA, stacked
train-lora.py --type sdxl --model stabilityai/stable-diffusion-xl-base-1.0 \
  -l latent-consistency/lcm-lora-sdxl \
  -l "ostris/crayon_style_lora_sdxl|crayons_v1_sdxl.safetensors" \
  --min-resolution 1024 --max-resolution 1024 --output engines/sdxl-lcm-crayon

# single style LoRA on a turbo base
train-lora.py --type sdxl --model stabilityai/sdxl-turbo \
  -l "ostris/ikea-instructions-lora-sdxl|ikea_instructions_xl_v1_5.safetensors" \
  --min-resolution 1024 --max-resolution 1024 --output engines/sdxlturbo-ikea

Runtime LoRA (-l PATH:runtime)

Fused LoRAs are baked at export time and cannot be changed later. Suffix a LoRA with :runtime instead to expose it as a lora_scale engine input, which the node's "LoRA scale" knob drives live (1 = fully on, matching a fused build). Costs ~2%.

train-lora.py --type sdxl --model stabilityai/stable-diffusion-xl-base-1.0 \
  -l "/path/to/crayons_v1_sdxl.safetensors:runtime" \
  --min-resolution 1024 --max-resolution 1024 --output engines/crayon-rt

Optional flags

  • --fp8 — FP8 (e4m3) UNet quantization (SDXL/SD; diffusion recipe).
  • --v2v / --v2v-inject — StreamV2V kvo extended-attention UNet (SD1.5 only).
  • --hw-compat ampere_plus — portable engine across SM 8.0+ GPUs (~5–15% slower, larger).
  • -l "REPO:WEIGHT" — set LoRA fusion scale; --lora-scale for the global scale.

About

C++ / CUDA implementation of StreamDiffusion

Topics

Resources

Stars

39 stars

Watchers

0 watching

Forks

Releases

Sponsor this project

Contributors

Languages