Skip to content

Stable Diffusion 1.5 - #779

Open
nshejwalkar wants to merge 34 commits into
cactus-compute:mainfrom
nshejwalkar:neel-diffusion
Open

Stable Diffusion 1.5#779
nshejwalkar wants to merge 34 commits into
cactus-compute:mainfrom
nshejwalkar:neel-diffusion

Conversation

@nshejwalkar

@nshejwalkar nshejwalkar commented Aug 8, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

This PR adds Stable Diffusion 1.5 (LCM) end to end. This is the first feature with text->image functionality in Cactus, as requested in issue 216. It implements a pipeline of three models: a CLIP encoder, U-net, and Tiny AutoEncoder for Stable Diffusion as the decoder.

cactus-kernels + cactus-graph

  • Added upsample_nearest2d: a new graph op with NEON and Metal kernels.
  • Added conv2d_k3s1p1: Metal kernel, 6.3× over the generic conv on TAESD decode.
  • Exported both through the graph FFI so the transpiler can emit them.

cactus-engine

  • Added cactus_generate_image to the public C API, and a new translation unit that drives the pipeline: encode the prompt, run the denoising loop, decode the final latents, and pack the result as an RGB image.
  • Added an iterative_denoise execution route.
  • Implemented the LCM sampler on the host — timestep schedule, alpha table, and guidance embedding. The three graphs are pure tensor functions; all of the loop control lives here, because the original pipeline has Python between each stage.
  • Scheduler settings (beta schedule, training timesteps, timestep scaling) are read from the bundle manifest, carried through from the model's own scheduler_config.json rather than hardcoded.
  • Added CLIP's tokenization to the BPE tokenizer, which needs its own word splitting, normalizer, and decoder.

python/convert

  • Added support for diffusers checkpoints. These differ from transformers repos in three ways that all had to be handled: components live in separate subfolders, the weights file has a different name, and the model classes can't be constructed by AutoModel at all.
  • Added family adapters for TAESD, the SD U-Net, and the CLIP text encoder. They keep the full checkpoint tensor name, so parallel stacks like encoder.* and decoder.* can't collapse onto the same output file.
  • Conversion now fails loudly when two tensors would be written to the same filename. Previously the second one silently overwrote the first for TAESD.
  • Added a CLIP tokenizer converter that emits the engine's sidecar files from the model's own tokenizer files.

python/transpiler

  • Added an SD1.5 profile declaring the three components, how they're wired, and that the runtime should drive them as a denoising loop.
  • This is the first Cactus model that is a pipeline of independent models rather than a single checkpoint. Existing multi-component models like LFM2-VL are one checkpoint exported once and then sliced into components afterwards; here, each component is its own checkpoint with its own inputs, and the decoder comes from a different repository altogether. Added ComponentSource so a profile can say that: each part is loaded and exported on its own, reusing the existing loading and sample-input machinery rather than introducing a parallel path for diffusion.
  • Mapped upsample_nearest2d through the IR and lowering, and taught the convolution fusions to tell 1D from 2D by kernel rank.
  • Two export fixes: CLIP registers position ids as a buffer no checkpoint can supply, so it's recomputed in the graph; and attention is kept whole through decomposition so the existing attention fusion still matches (improved decode from 33s -> 8.2s).

python/cli

  • Added cactus generate-image, which prepares the bundle, generates, and writes a PNG.
  • cactus convert now handles pipeline models, converting each component into a single shared weights directory. One component (the TAESD decoder) comes from a different repository entirely.
  • Added the Python binding for the new image generation API.

Tests (correctness and performance)

  • kernels — upsample against a hand-computed reference, covering all three code paths in the kernel, plus benchmarks for the new convolution and upsample.
  • graph — save/load round trip proving the scale factor survives serialization, and Metal/CPU parity.
  • engine — unit tests for the sampler math (timestep schedule, alpha table, guidance embedding), including empty and edge-case inputs.
  • convert — adapter selection and tensor naming, including that encoder and decoder names stay distinct.
  • transpiler — each component transpiled and compared against PyTorch (CLIP 0.30%, U-Net 0.02%, decoder 0.17% relative error), plus the full text→image pipeline end to end.

Example

cactus convert SimianLuo/LCM_Dreamshaper_v7
cactus generate-image --prompt "a red fox in snow" --out fox.png
fox

512×512 on M1 Pro w/ Metal: 8.2 s at 4 steps, 4.0 s at 2 steps.

Passes all correctness tests.

@nshejwalkar
add upsample_nearest2d op and TAESD VAE decoder test
5fc9b2c 
Adds a nearest-neighbor 2D upsample graph op (CPU NEON path + Metal
kernel, with CPU/Metal parity test), the one op missing for VAE
decoders. Includes test_vae.cpp which builds the TAESD decoder as a
CactusGraph and decodes a 512x512 image from a 4x64x64 latent
(74dB PSNR vs torch fp32 reference; weights exported at test time by
tests/export_taesd.py, skipped when CACTUS_TAESD_DIR is unset).

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
add simdgroup-MMA Metal fast path for conv2d_k3s1p1
8fa51be 
Implicit-GEMM tiled kernel (32 output channels x 64 pixels per
threadgroup, fp16 MMA with fp32 accumulate, in-kernel im2col staging).
TAESD 512x512 VAE decode on M1 Pro drops from 1321ms to 230ms on Metal
(CPU path unchanged at ~760ms); CPU/Metal parity suite passes.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
convert TAESD through the standard weights pipeline
780027c 
The VAE test shipped its own downloader and raw fp16 blob format, which
sidestepped the conversion pipeline every other model uses. Instead:

- load diffusers checkpoints (diffusion_pytorch_model.safetensors) in
  the remote path, mirroring the local directory glob
- let an adapter report no transformers model class, so families that
  convert from checkpoint tensors alone skip the model load
- add a taesd family adapter (AutoencoderTiny) that keeps encoder and
  decoder tensor names distinct; the generic layer-suffix namer strips
  the prefix and silently collapsed 42 of 134 tensors onto each other
- fail conversion when two tensors map to the same output file

test_vae now runs both the encoder and decoder from a converted bundle
via mmap_weights and round-trips a repo asset, so it needs no torch and
no reference dumps: 31.1dB round-trip on M1 Pro (torch scores 30.7dB on
the same image), decode 238ms Metal / 846ms CPU.

Prepare with: cactus convert madebyollin/taesd <dir> --weights-only

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
prepare VAE test weights through cactus test
4bb0426 
The vae suite required a manual convert plus a hand-set env var, and
silently skipped otherwise, unlike every other model-backed test. Follow
the engine pattern instead: cactus test prepares the weights itself via
ensure_weights and passes --vae-model to the graph test.sh, which
exports CACTUS_TEST_VAE_MODEL along with CACTUS_TEST_ASSETS.

cactus test --component graph now converts madebyollin/taesd on first
run and executes the round-trip with no setup. Running test.sh directly
still skips, since it has no way to fetch a model.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
test upsample_nearest2d against a hand-computed reference
d4b8543 
The only coverage was a CPU/Metal parity case, which asserts the two
backends agree rather than that either is right, and which the Linux CI
leg skips entirely for want of Metal. It also only exercised scale=2,
leaving the scale==1 memcpy short-circuit and the scale>=3 row
replication loop unrun.

Assert bit-equality with input[(y/s)*W + x/s] over scales 1..3, matching
the maxpool1d reference test; nearest-neighbour is an exact copy, so no
tolerance is needed. Verified by injecting a scale-3-only fault, which
this catches and the parity case does not.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
inline the diffusers family check
bf4e6d2 
A one-entry lookup table read as a registry other code consults, but
detect_family was its only caller and the file already dispatches with
plain conditionals.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
try every checkpoint stem when one fails
4e483fa 
Two flaws in the stem loop:

The single-file attempt re-raised RuntimeError instead of falling
through to the index path, which the nested handlers it replaced did
not do. It guarded nothing either way: the duplicate-key raise is
outside any try and propagates on its own.

The loop returned None as soon as a stem raised, so a repo whose
model.safetensors.index.json is corrupt never reached
diffusion_pytorch_model.safetensors. TAESD only worked because a
missing index returns None rather than raising.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
reject colliding tensor names before writing, not after
c17def5 
The guard ran after every tensor had been saved, so a colliding name had
already truncated the earlier tensor via the 'wb' open in
save_tensor_with_header, and raising before write_reports left a weights
directory with no manifest.

Move it to where out_path is resolved, which every write path funnels
through and which sits outside the per-tensor try, so the collision
aborts before the second tensor overwrites the first. Verified by
converting TAESD under --model-family generic, whose layer-suffix namer
collides 42 ways: it now stops at the first collision with the earlier
file intact.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
pass the VAE weights through the test environment
8e69c98 
The test runner already copies an environment into each component
subprocess, so routing the path through a test.sh flag only to re-export
it added a hop. Set it beside CACTUS_NO_CLOUD_TELE instead and leave
test.sh untouched. CACTUS_TEST_ASSETS was redundant too: the test falls
back to the repo-relative asset path, which resolves from build/.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
transpile nearest-neighbour upsampling
ba6db4c 
Export conv2d_k3s1p1 and upsample_nearest2d through the C FFI: the
Python bindings already declared conv2d_k3s1p1 optionally and lower.py
already called it, so the graph builder was unreachable from the
transpiler and any model needing it fell over.

Map aten.upsample_nearest2d through the op tables. The .vec overload
carries either scale factors or an output size, so the importer takes
the factors when present and otherwise divides the output size by the
input shape, failing loudly on anything non-integral or anisotropic.

The TAESD decoder now transpiles end to end from diffusers with no
hand-written graph: 35 conv2d, 31 relu, 10 add, 3 upsample_nearest2d,
matching torch to 0.166% relative mean error (max abs 0.028) at fp16.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
drop the hand-built VAE test for kernel benchmarks
216f911 
test_vae was the only graph test that loaded a trained model and
restated an architecture in C++; every other test there builds what it
needs. Transpiling AutoencoderTiny now validates the same numerics
against the real diffusers module, which is stronger evidence than a
hand-transcribed graph, so the reimplementation has no remaining claim.

Its perf signal moves to the kernel benchmarks, where conv2d_k3s1p1 and
upsample_nearest2d are timed at the sizes a 512px decode hits. The test
CLI plumbing that only existed to feed it goes back to upstream.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
persist the upsample scale factor across save/load
e1c8303 
A new OpType has to be registered in two places I missed. Without a
param_io schema entry no fields are written at all, so a graph saved
with scale 3 reloaded as 1, the OpParams default. And read_node_entry
bounds-checks against the last enumerator by name, so any saved graph
containing the op was rejected as corrupt.

Both only bite through graph.save/load, which is the path a transpiled
bundle takes, so nothing in the hand-built tests could have caught it.
Covered by a round-trip test in the serialization suite, verified to
fail without the schema entry.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
benchmark conv2d and upsample at graph level
a94815b 
The graph benchmarks inherit cactus_default_backend(), so they time
Metal on Apple and CPU on the Linux leg. That is where the MMA conv
fast path shows up, and nothing was measuring it: no conv of any kind
had a graph-level benchmark, so a change that silently dropped the fast
path would still pass every test.

conv2d_k3s1p1 64x512x512 on M1 Pro: 176.40ms through the generic
conv2d_f16, 27.97ms through conv2d_k3s1p1_mma_f16.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
test the transpiled decoder end to end
44f210f 
The op counts and the torch comparison for TAESD lived in a scratch
script, so the numbers behind the port were not reproducible from the
tree. This runs the same path as a test: capture AutoencoderTiny's
decoder, lower it, execute, and compare against the torch module.

Guarded the way the repo already guards optional work — importorskip on
diffusers, local_files_only so it uses a cached checkout rather than
downloading, and a skip when the runtime library is absent. It asserts
on the relative error and reports latency rather than asserting it,
matching how the benchmarks report.

512x512 decode on M1 Pro: 234ms, 0.166% relative mean error.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
inline the checkpoint stem loop
a44b4cb 
Keeps the loop in _load_checkpoint_state_dict rather than extracting a
per-stem helper. Same behaviour either way; this trades a smaller
conceptual diff for one more level of indentation on the shard block.

Revert this commit to get the extracted version back.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
drop the unused importorskip binding
4dd57e0 
The module gets imported where it is used, so the guard does not need to
bind it.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
nshejwalkar force-pushed the neel-diffusion branch 2 times, most recently from cc585ac to 4dd57e0 Compare August 8, 2026 22:17
@nshejwalkar

nshejwalkar commented Aug 8, 2026

Copy link
Copy Markdown
Contributor Author

The failed ci python test is because of e9773b1, unrelated to this PR

@nshejwalkar
restore the py.typed marker the packaging test expects
417ecb4 
Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
tolerate group_norm exports that omit eps
6951845 
torch.export drops trailing defaulted args, so the SD1.5 UNet's
group_norm nodes carry only (input, num_groups, weight, bias). The
transpiled UNet test fails to capture without this.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
test the transpiled CLIP text encoder end to end
f78c6d6 
Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
test the transpiled text-to-image pipeline end to end
dde57ef 
Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
convert the SD text encoder and UNet through family adapters
1bb943f 
Verbatim full-name tensor naming shared with taesd, so parallel
blocks never collide. Folds the taesd adapter tests into the one
SD adapter suite.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
lower upsample_nearest2d in the replacement transpiler
82e01c0 
Both aten overloads plus the cactus.* target: attrs from positional
args, one integer scale derived from scale_factors, scales_h/w, or
output_size against the input shape. The specialized conv2d lowerings
(k3s1p1 MMA, k3s2p1, pointwise, depthwise) were already wired upstream.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
distinguish conv1d and conv2d fusions by kernel rank
92aa1fc 
A square kernel collapses kernel_size to one scalar, so a 1x1 conv2d
satisfied conv1d_pointwise's requirements and lowered as rank-3; the SD
UNet's 46 pointwise convs hit the builder's shape check. Infer
spatial_dims from the weight shape and require it on every conv
definition.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
wire SD text-to-image through the replacement transpiler
aa931c9 
LCM_Dreamshaper_v7 gets a registered sd15_t2i profile: per-component
export modes (text_encoder, unet, vae_decoder), a diffusion Converter
path that loads each component from its subfolder or external repo and
synthesizes latent/timestep inputs, a t2i_generate route with an
iterative_denoise runtime contract, and multi-component weight
conversion merging one weights_manifest per bundle. The CLIP wrapper
exports the inner transformer and an export patch recomputes its
non-persistent position_ids buffer in-graph; a resolver fallback drops
one leading wrapper attribute when exact weight names miss.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
materialize diffusion runtime assets and fix weight-name aliasing
6ed7fcb 
SD conversion now emits the engine tokenizer sidecars (vocab.txt,
merges.txt, tokenizer_config.txt with the clip normalizer, special
tokens) and scheduler_config.json; the bundle build flattens the
scheduler constants into manifest metadata. Weight resolution tries
exact name variants, including the dropped wrapper attribute, before
the tied-lm-head aliases -- every UNet attention proj_out previously
resolved to the first block's tensor through the embed_tokens alias.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
run text-to-image bundles in the engine
ca410ab 
New ITERATIVE_DENOISE route for text_encoder+unet+vae_decoder bundles:
Model::generate_image tokenizes with a new CLIP BPE mode (lowercase,
word split, </w> suffix), runs the text encoder, drives the LCM denoise
loop host-side from the manifest's scheduler metadata, decodes latents,
and returns u8 RGB through the new cactus_generate_image C API.
test_diffusion covers tokenizer parity against reference ids and a
seeded 2-step generation, skipping unless CACTUS_SD_BUNDLE is set.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
keep declarations grouped by kind in the profile modules
08c6893 
Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
nshejwalkar marked this pull request as draft August 9, 2026 23:48
@nshejwalkar
reuse the Converter registries for pipeline components
4c0a701 
Component sources become typed ComponentSource records instead of
colon-packed strings parsed in two modules, and the export modes derive
from them rather than being written twice. Loading now goes through
load_model's retry ladder with an explicit class, input synthesis joins
SYNTHETIC_INPUT_BUILDERS beside the processor map (un-disabling the
synthetic strategy the profiles already named), and the CLIP tokenizer
sidecars are written by the convert layer from the downloaded files
rather than from constants inlined in the CLI. create_model and
_convert_from_source are dispatchers over a named branch per shape, so
diffusion.py keeps only the export wrappers.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
give the diffusion scheduler its own translation unit and test
e6a59b9 
Model::generate_image and the LCM helpers move to generate_image.cpp
beside the FFI shim, so model.cpp keeps only manifest parsing and the
route. The helpers take DiffusionParams and are declared in engine.h,
which makes the numerics unit-testable: test_diffusion now checks the
4- and 2-step schedules against diffusers and the alphas_cumprod tail
without needing a bundle, so they run when CACTUS_SD_BUNDLE is unset.
The route follows the plan's declared iterative_denoise strategy and
reports which component is missing rather than sniffing component
names, and the guidance embedding no longer divides by zero when
time_cond_proj_dim is 2.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
drop comments the surrounding files do not carry
5b18b72 
Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
add a generate-image command
cc4cb46 
Wraps cactus_generate_image in the bindings and drives it from the CLI:
resolve or convert the bundle, reject a model whose plan is not
iterative_denoise, then write the returned RGB as a PNG with no new
dependency.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
keep attention whole through export so the fusion can match it
2ea22e8 
run_decompositions shatters scaled_dot_product_attention into bmm,
softmax and a NaN guard, and the re-fusion patterns only cover the LLM
shape, so the SD UNet's 32 attention blocks stayed decomposed: 4835
executed nodes against the 1759 the old pipeline produces. A component
can now name aten ops to keep whole, and the direct SDPA fusion carries
the layout and causality the aten op actually implies -- without them
the fused graph is fast and wrong, since q/k/v reach the kernel as
BHSD and is_causal defaults to true.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
normalize the attention mask rank so masked attention can fuse
c6d1e5b 
The kernel takes a rank-3 mask, but HF hands scaled_dot_product_attention
a (batch, 1, query, key) one, and the replacement transpiler passed it
through untouched -- so the text encoder produced a coherent image for
the wrong prompt. Reshape it the way the graph builder path already
does, which lets the text encoder keep its attention whole too.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
pin the preserved-attention contract and tighten the conv discriminator
2ba9f46 
The direct attention fusion now has a test covering the layout it
assumes and both causality paths, since an omitted is_causal has to
follow aten rather than the decoder-shaped guess it used to make.
conv1d_causal shares CONV_GRAPH with the conv2d fusions, so it needs
the same rank discriminator; the short-conv fusions match distinct
patterns and one carries no conv node to infer it from. Drops the
profile-level preserved_ops nothing sets and the shadowed loop
variable in load_model.

Signed-off-by: Neel Shejwalkar <nshejwalkar7@gmail.com>
@nshejwalkar
nshejwalkar marked this pull request as ready for review August 13, 2026 03:37
@nshejwalkar nshejwalkar changed the title TAESD Stable Diffusion 1.5 Aug 13, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@nshejwalkar