From 5fc9b2c038f7ee03c48fdd806f6b63c94f0887b5 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sun, 2 Aug 2026 13:33:40 -0700 Subject: [PATCH 01/34] add upsample_nearest2d op and TAESD VAE decoder test 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 --- cactus-graph/cactus_graph.h | 4 +- cactus-graph/src/builder.cpp | 20 +++ cactus-graph/src/execute.cpp | 14 +- cactus-graph/src/ops_conv.cpp | 19 +++ cactus-graph/tests/export_taesd.py | 103 ++++++++++++++ cactus-graph/tests/test_metal_parity.cpp | 8 ++ cactus-graph/tests/test_vae.cpp | 161 ++++++++++++++++++++++ cactus-kernels/cactus_kernels.h | 8 ++ cactus-kernels/metal_backend.h | 2 + cactus-kernels/src/cactus_kernels.metal | 11 ++ cactus-kernels/src/conv2d.cpp | 40 ++++++ cactus-kernels/src/metal_backend.mm | 15 +- cactus-kernels/src/metal_backend_stub.cpp | 1 + docs/cactus_graph.md | 6 + 14 files changed, 408 insertions(+), 4 deletions(-) create mode 100644 cactus-graph/tests/export_taesd.py create mode 100644 cactus-graph/tests/test_vae.cpp diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 9475f4293..655121e77 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -147,7 +147,8 @@ enum class OpType { QKV_TQ_FUSED, PROJECTION_PAIR_TQ_FUSED, CONV1D_CAUSAL_CHANNEL_FIRST, - LOGITS_TQ_SOFTCAP + LOGITS_TQ_SOFTCAP, + UPSAMPLE_NEAREST2D }; struct PrecisionTraits { @@ -746,6 +747,7 @@ class CactusGraph { size_t bilinear_interpolation(size_t pos_embeds, size_t dst_height, size_t dst_width, bool align_corners = true, ComputeBackend backend = cactus_default_backend()); size_t maxpool1d(size_t input, size_t kernel_size, size_t stride, ComputeBackend backend = cactus_default_backend()); + size_t upsample_nearest2d(size_t input, size_t scale_factor, ComputeBackend backend = cactus_default_backend()); size_t lstm_cell( size_t input, size_t h_prev, size_t c_prev, diff --git a/cactus-graph/src/builder.cpp b/cactus-graph/src/builder.cpp index 0dd3796e4..21624bea5 100644 --- a/cactus-graph/src/builder.cpp +++ b/cactus-graph/src/builder.cpp @@ -1481,6 +1481,26 @@ size_t CactusGraph::maxpool1d(size_t input, size_t kernel_size, size_t stride, C return tag_backend(add_node(OpType::MAXPOOL1D, {input}, {batch, channels, output_length}, params), backend); } +size_t CactusGraph::upsample_nearest2d(size_t input, size_t scale_factor, ComputeBackend backend) { + const auto& in_buf = get_output_buffer(input); + if (in_buf.shape.size() != 4) { + throw std::runtime_error("upsample_nearest2d expects input [N, C, H, W]"); + } + if (scale_factor == 0) { + throw std::runtime_error("upsample_nearest2d scale_factor must be > 0"); + } + + size_t batch = in_buf.shape[0]; + size_t channels = in_buf.shape[1]; + size_t height = in_buf.shape[2]; + size_t width = in_buf.shape[3]; + + OpParams params; + params.stride = scale_factor; + return tag_backend(add_node(OpType::UPSAMPLE_NEAREST2D, {input}, + {batch, channels, height * scale_factor, width * scale_factor}, params), backend); +} + size_t CactusGraph::conv2d_k3s1p1(size_t input, size_t weight, ComputeBackend backend) { const auto& xin = get_output_buffer(input); const auto& w = get_output_buffer(weight); diff --git a/cactus-graph/src/execute.cpp b/cactus-graph/src/execute.cpp index 56830324e..827701511 100644 --- a/cactus-graph/src/execute.cpp +++ b/cactus-graph/src/execute.cpp @@ -76,6 +76,7 @@ DECLARE_COMPUTE(compute_gaussian_topk_node); DECLARE_COMPUTE(compute_maxpool1d_node); DECLARE_COMPUTE(compute_bilstm_sequence_node); DECLARE_COMPUTE(compute_conv2d_k3s1p1_node); +DECLARE_COMPUTE(compute_upsample_nearest2d_node); DECLARE_COMPUTE(compute_stats_pool_node); DECLARE_COMPUTE(compute_weighted_stats_pool_node); DECLARE_COMPUTE(compute_transpose_node); @@ -117,7 +118,7 @@ DECLARE_COMPUTE(compute_conv1d_causal_channel_first_node); extern void shrink_thread_local_buffers(); #undef DECLARE_COMPUTE -static constexpr int OP_TYPE_COUNT = static_cast(OpType::LOGITS_TQ_SOFTCAP) + 1; +static constexpr int OP_TYPE_COUNT = static_cast(OpType::UPSAMPLE_NEAREST2D) + 1; static_assert(OP_TYPE_COUNT <= 256, "OpType dispatch table overflow"); static ComputeFn dispatch_flat[OP_TYPE_COUNT] = {}; @@ -227,6 +228,7 @@ static bool init_dispatch() { dispatch_flat[static_cast(OpType::ALTUP_CORRECT)] = compute_altup_correct_node; dispatch_flat[static_cast(OpType::GAUSSIAN_TOPK)] = compute_gaussian_topk_node; dispatch_flat[static_cast(OpType::MAXPOOL1D)] = compute_maxpool1d_node; + dispatch_flat[static_cast(OpType::UPSAMPLE_NEAREST2D)] = compute_upsample_nearest2d_node; dispatch_flat[static_cast(OpType::BILSTM_SEQUENCE)] = compute_bilstm_sequence_node; dispatch_flat[static_cast(OpType::STATS_POOL)] = compute_stats_pool_node; dispatch_flat[static_cast(OpType::WEIGHTED_STATS_POOL)] = compute_weighted_stats_pool_node; @@ -311,7 +313,8 @@ static const char* op_type_names[] = { "EXPAND", "MASKED_SELECT_PREFIX", "QKV_TQ_FUSED", "PROJECTION_PAIR_TQ_FUSED", "CONV1D_CAUSAL_CHANNEL_FIRST", - "LOGITS_TQ_SOFTCAP" + "LOGITS_TQ_SOFTCAP", + "UPSAMPLE_NEAREST2D" }; static const char* get_op_name(OpType op) { @@ -864,6 +867,13 @@ static bool try_encode_metal(GraphNode& node, const nodes_vector& nodes, const n (uint32_t)(in.shape[0] * in.shape[1]), (uint32_t)in.shape[2], (uint32_t)out.shape[2], (uint32_t)node.params.kernel_size, (uint32_t)node.params.stride); } + case OpType::UPSAMPLE_NEAREST2D: { + const auto& in = get_input(node, 0, nodes, map); + if (!fp16(in) || !fp16(out) || in.shape.size() != 4 || out.shape.size() != 4) return false; + return cactus_metal_encode_upsample_nearest2d(out.get_data(), in.get_data(), + (uint32_t)(in.shape[0] * in.shape[1]), (uint32_t)in.shape[2], + (uint32_t)in.shape[3], (uint32_t)node.params.stride); + } case OpType::BILINEAR_INTERPOLATION: { const auto& in = get_input(node, 0, nodes, map); if (!fp16(in) || !fp16(out) || in.shape.size() != 2) return false; diff --git a/cactus-graph/src/ops_conv.cpp b/cactus-graph/src/ops_conv.cpp index 8850123e7..3b477382f 100644 --- a/cactus-graph/src/ops_conv.cpp +++ b/cactus-graph/src/ops_conv.cpp @@ -910,6 +910,25 @@ void compute_maxpool1d_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& input = get_input(node, 0, nodes, node_index_map); + + if (input.shape.size() != 4) { + throw std::runtime_error("upsample_nearest2d expects input [N, C, H, W]"); + } + if (input.precision != Precision::FP16) { + throw std::runtime_error("upsample_nearest2d only supports FP16"); + } + + cactus_upsample_nearest2d_f16( + input.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + input.shape[0] * input.shape[1], + input.shape[2], input.shape[3], + node.params.stride); +} + void compute_conv2d_k3s1p1_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { const auto& X = get_input(node, 0, nodes, node_index_map); diff --git a/cactus-graph/tests/export_taesd.py b/cactus-graph/tests/export_taesd.py new file mode 100644 index 000000000..eac656b15 --- /dev/null +++ b/cactus-graph/tests/export_taesd.py @@ -0,0 +1,103 @@ +"""Export TAESD decoder weights and reference artifacts for test_vae.cpp. + +Usage: python export_taesd.py [output_dir] [test_image] +Then: CACTUS_TAESD_DIR= ./test_vae + +Requires torch, safetensors, and Pillow. Downloads the MIT-licensed TAESD +weights (https://huggingface.co/madebyollin/taesd) on first run. +""" + +import json +import sys +import urllib.request +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +from PIL import Image +from safetensors.torch import load_file + +OUT = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(__file__).parent / "taesd" +ASSET = Path(sys.argv[2]) if len(sys.argv) > 2 else \ + Path(__file__).resolve().parents[2] / "cactus-engine" / "tests" / "assets" / "test_monkey.png" +OUT.mkdir(parents=True, exist_ok=True) + +HF = "https://huggingface.co/madebyollin/taesd/resolve/main" + + +def fetch(name): + dst = OUT / name + if not dst.exists(): + print(f"downloading {name}...") + urllib.request.urlretrieve(f"{HF}/{name}", dst) + return dst + + +def conv(n_in, n_out, **kwargs): + return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs) + + +class Clamp(nn.Module): + def forward(self, x): + return torch.tanh(x / 3) * 3 + + +class Block(nn.Module): + def __init__(self, n_in, n_out): + super().__init__() + self.conv = nn.Sequential(conv(n_in, n_out), nn.ReLU(), conv(n_out, n_out), nn.ReLU(), conv(n_out, n_out)) + self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() + self.fuse = nn.ReLU() + + def forward(self, x): + return self.fuse(self.conv(x) + self.skip(x)) + + +def encoder(): + return nn.Sequential( + conv(3, 64), Block(64, 64), + conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64), + conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64), + conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64), + conv(64, 4), + ) + + +def decoder(): + return nn.Sequential( + Clamp(), conv(4, 64), nn.ReLU(), + Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False), + Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False), + Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False), + Block(64, 64), conv(64, 3), + ) + + +enc = encoder() +enc.load_state_dict(load_file(fetch("taesd_encoder.safetensors"))) +dec = decoder() +dec.load_state_dict(load_file(fetch("taesd_decoder.safetensors"))) +enc.eval() +dec.eval() + +img = Image.open(ASSET).convert("RGB").resize((512, 512), Image.LANCZOS) +x = torch.from_numpy(np.array(img)).float().permute(2, 0, 1).unsqueeze(0) / 255.0 + +with torch.no_grad(): + latent = enc(x) + ref = dec(latent).clamp(0, 1) + +latent.to(torch.float16).numpy().tofile(OUT / "latent.f16.bin") +ref.to(torch.float32).numpy().tofile(OUT / "ref_image.f32.bin") +Image.fromarray((ref[0].permute(1, 2, 0).numpy() * 255).round().astype(np.uint8)).save(OUT / "ref_image.png") + +manifest = [] +with open(OUT / "decoder.f16.bin", "wb") as f: + for name, p in dec.named_parameters(): + f.write(p.detach().to(torch.float16).numpy().tobytes()) + manifest.append({"name": name, "shape": list(p.shape)}) +(OUT / "manifest.json").write_text(json.dumps(manifest, indent=1)) + +print(f"latent shape {list(latent.shape)}, range [{latent.min():.3f}, {latent.max():.3f}]") +print(f"wrote {len(manifest)} tensors to {OUT}") diff --git a/cactus-graph/tests/test_metal_parity.cpp b/cactus-graph/tests/test_metal_parity.cpp index ae8f2d548..3d173a5ad 100644 --- a/cactus-graph/tests/test_metal_parity.cpp +++ b/cactus-graph/tests/test_metal_parity.cpp @@ -210,6 +210,13 @@ bool parity_maxpool1d() { return c.check(); } +bool parity_upsample_nearest2d() { + ParityCase c; + c.add_input({2, 3, 7, 5}); + c.build = [](CactusGraph& g, const std::vector& in) { return g.upsample_nearest2d(in[0], 2); }; + return c.check(); +} + bool parity_bilinear() { ParityCase c; c.add_input({64, 12}); @@ -544,6 +551,7 @@ int main() { runner.run_test("rope", parity_rope(false)); runner.run_test("rope_gptj", parity_rope(true)); runner.run_test("maxpool1d", parity_maxpool1d()); + runner.run_test("upsample_nearest2d", parity_upsample_nearest2d()); runner.run_test("bilinear_interpolation", parity_bilinear()); runner.run_test("conv1d", parity_conv1d()); runner.run_test("conv1d_k7s3", parity_conv1d_k7s3()); diff --git a/cactus-graph/tests/test_vae.cpp b/cactus-graph/tests/test_vae.cpp new file mode 100644 index 000000000..65383e625 --- /dev/null +++ b/cactus-graph/tests/test_vae.cpp @@ -0,0 +1,161 @@ +#include "test_utils.h" +#include +#include +#include +#include +#include + +using namespace TestUtils; + +namespace { + +constexpr size_t LATENT_C = 4, LATENT_HW = 64, IMAGE_C = 3, IMAGE_HW = 512, CH = 64; + +std::vector<__fp16> read_fp16_file(const std::string& path) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) return {}; + size_t bytes = static_cast(f.tellg()); + std::vector<__fp16> data(bytes / sizeof(__fp16)); + f.seekg(0); + f.read(reinterpret_cast(data.data()), bytes); + return data; +} + +std::vector read_fp32_file(const std::string& path) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) return {}; + size_t bytes = static_cast(f.tellg()); + std::vector data(bytes / sizeof(float)); + f.seekg(0); + f.read(reinterpret_cast(data.data()), bytes); + return data; +} + +struct WeightStream { + const std::vector<__fp16>& blob; + size_t offset = 0; + + size_t next(CactusGraph& g, const std::vector& shape) { + size_t n = 1; + for (size_t d : shape) n *= d; + size_t id = g.input(shape, Precision::FP16); + g.set_input(id, blob.data() + offset, Precision::FP16); + offset += n; + return id; + } +}; + +size_t taesd_conv(CactusGraph& g, WeightStream& w, size_t x, size_t c_in, size_t c_out, bool bias = true) { + size_t weight = w.next(g, {c_out, c_in, 3, 3}); + if (!bias) return g.conv2d_k3s1p1(x, weight); + size_t b = w.next(g, {c_out}); + return g.conv2d_k3s1p1(x, weight, b); +} + +size_t taesd_block(CactusGraph& g, WeightStream& w, size_t x) { + size_t h = g.relu(taesd_conv(g, w, x, CH, CH)); + h = g.relu(taesd_conv(g, w, h, CH, CH)); + h = taesd_conv(g, w, h, CH, CH); + return g.relu(g.add(h, x)); +} + +size_t build_taesd_decoder(CactusGraph& g, WeightStream& w, size_t latent) { + size_t x = g.scalar_multiply(g.tanh(g.scalar_divide(latent, 3.0f)), 3.0f); + x = g.relu(taesd_conv(g, w, x, LATENT_C, CH)); + for (int stage = 0; stage < 3; ++stage) { + x = taesd_block(g, w, x); + x = taesd_block(g, w, x); + x = taesd_block(g, w, x); + x = g.upsample_nearest2d(x, 2); + x = taesd_conv(g, w, x, CH, CH, false); + } + x = taesd_block(g, w, x); + x = taesd_conv(g, w, x, CH, IMAGE_C); + return g.clamp(x, 0.0f, 1.0f); +} + +void write_ppm(const std::string& path, const __fp16* chw) { + std::ofstream f(path, std::ios::binary); + f << "P6\n" << IMAGE_HW << " " << IMAGE_HW << "\n255\n"; + const size_t plane = IMAGE_HW * IMAGE_HW; + for (size_t i = 0; i < plane; ++i) { + for (size_t c = 0; c < IMAGE_C; ++c) { + float v = static_cast(chw[c * plane + i]) * 255.0f; + f.put(static_cast(v < 0.0f ? 0 : (v > 255.0f ? 255 : static_cast(v + 0.5f)))); + } + } +} + +float psnr_vs_reference(const __fp16* out, const std::vector& ref) { + double se = 0.0; + for (size_t i = 0; i < ref.size(); ++i) { + double d = static_cast(out[i]) - static_cast(ref[i]); + se += d * d; + } + double mse = se / ref.size(); + return mse > 0.0 ? static_cast(10.0 * std::log10(1.0 / mse)) : 99.0f; +} + +bool decode_and_check(const std::string& dir, const char* backend, const std::vector<__fp16>& weights, + const std::vector<__fp16>& latent, const std::vector& ref) { + if (cactus_backend_select(backend) != 0) { + std::cout << " [" << backend << "] unavailable, skipping" << std::endl; + return true; + } + + CactusGraph g; + size_t latent_id = g.input({1, LATENT_C, LATENT_HW, LATENT_HW}, Precision::FP16); + WeightStream w{weights}; + size_t out_id = build_taesd_decoder(g, w, latent_id); + if (w.offset != weights.size()) { + std::cout << " weight blob mismatch: consumed " << w.offset << " of " << weights.size() << std::endl; + return false; + } + g.set_input(latent_id, latent.data(), Precision::FP16); + + g.execute(); + constexpr int runs = 5; + Timer timer; + for (int i = 0; i < runs; ++i) g.execute(); + double ms = timer.elapsed_ms() / runs; + + const __fp16* out = static_cast(g.get_output(out_id)); + float psnr = psnr_vs_reference(out, ref); + write_ppm(dir + "/cactus_decode_" + backend + ".ppm", out); + std::cout << " [" << backend << "] " << ms << " ms/decode, PSNR vs torch fp32: " << psnr << " dB" << std::endl; + + g.hard_reset(); + return psnr > 30.0f; +} + +bool test_taesd_decoder() { + const char* env = std::getenv("CACTUS_TAESD_DIR"); + if (!env) { + std::cout << " skipped: set CACTUS_TAESD_DIR to a dir with decoder.f16.bin/latent.f16.bin/ref_image.f32.bin" + << std::endl; + return true; + } + std::string dir(env); + + auto weights = read_fp16_file(dir + "/decoder.f16.bin"); + auto latent = read_fp16_file(dir + "/latent.f16.bin"); + auto ref = read_fp32_file(dir + "/ref_image.f32.bin"); + if (weights.empty() || latent.size() != LATENT_C * LATENT_HW * LATENT_HW || + ref.size() != IMAGE_C * IMAGE_HW * IMAGE_HW) { + std::cout << " missing or malformed TAESD artifacts in " << dir << std::endl; + return false; + } + + bool ok = decode_and_check(dir, "cpu", weights, latent, ref); + ok = decode_and_check(dir, "metal", weights, latent, ref) && ok; + return ok; +} + +} + +int main() { + TestRunner runner("VAE Decoder Tests"); + runner.run_test("taesd_decoder", test_taesd_decoder()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-kernels/cactus_kernels.h b/cactus-kernels/cactus_kernels.h index eddcc5bd9..96ccd288b 100644 --- a/cactus-kernels/cactus_kernels.h +++ b/cactus-kernels/cactus_kernels.h @@ -659,6 +659,14 @@ void cactus_maxpool1d_f16( size_t kernel_size, size_t stride); +void cactus_upsample_nearest2d_f16( + const __fp16* input, + __fp16* output, + size_t num_planes, + size_t height, + size_t width, + size_t scale); + void cactus_lstm_cell_f16( const __fp16* x_input, const __fp16* h_prev, diff --git a/cactus-kernels/metal_backend.h b/cactus-kernels/metal_backend.h index 9505aad48..e7982495a 100644 --- a/cactus-kernels/metal_backend.h +++ b/cactus-kernels/metal_backend.h @@ -156,6 +156,8 @@ bool cactus_metal_encode_rope_full(void* out, const void* in, uint32_t tokens, u float theta, int gptj); bool cactus_metal_encode_maxpool1d(void* out, const void* in, uint32_t NC, uint32_t L, uint32_t Lout, uint32_t K, uint32_t stride); +bool cactus_metal_encode_upsample_nearest2d(void* out, const void* in, uint32_t NC, + uint32_t H, uint32_t W, uint32_t scale); bool cactus_metal_encode_bilinear(void* out, const void* in, uint32_t sh, uint32_t sw, uint32_t dh, uint32_t dw, uint32_t E, int align); bool cactus_metal_encode_conv1d_gen(void* out, const void* x, const void* w, const void* bias, diff --git a/cactus-kernels/src/cactus_kernels.metal b/cactus-kernels/src/cactus_kernels.metal index 1c4a724c5..fedb02efe 100644 --- a/cactus-kernels/src/cactus_kernels.metal +++ b/cactus-kernels/src/cactus_kernels.metal @@ -1993,6 +1993,17 @@ kernel void maxpool1d_f16(device const half* x [[buffer(0)]], device half* y [[b y[(size_t)nc * Lout + lo] = (half)r; } +kernel void upsample_nearest2d_f16(device const half* x [[buffer(0)]], device half* y [[buffer(1)]], + constant uint& H [[buffer(2)]], constant uint& W [[buffer(3)]], + constant uint& scale [[buffer(4)]], + uint2 gid [[thread_position_in_grid]]) { + uint pix = gid.x, nc = gid.y; + uint Wo = W * scale, Ho = H * scale; + if (pix >= Ho * Wo) return; + uint oy = pix / Wo, ox = pix % Wo; + y[(size_t)nc * Ho * Wo + pix] = x[(size_t)nc * H * W + (size_t)(oy / scale) * W + ox / scale]; +} + kernel void bilinear_f16(device const half* in [[buffer(0)]], device half* y [[buffer(1)]], constant uint& sh [[buffer(2)]], constant uint& sw [[buffer(3)]], constant uint& dh [[buffer(4)]], constant uint& dw [[buffer(5)]], diff --git a/cactus-kernels/src/conv2d.cpp b/cactus-kernels/src/conv2d.cpp index 768d85b1c..0f2c8918a 100644 --- a/cactus-kernels/src/conv2d.cpp +++ b/cactus-kernels/src/conv2d.cpp @@ -617,3 +617,43 @@ void cactus_maxpool1d_f16( } }); } + +void cactus_upsample_nearest2d_f16( + const __fp16* input, + __fp16* output, + size_t num_planes, + size_t height, + size_t width, + size_t scale +) { + if (scale == 1) { + std::memcpy(output, input, num_planes * height * width * sizeof(__fp16)); + return; + } + + const size_t out_width = width * scale; + const size_t out_plane = height * scale * out_width; + + CactusThreading::parallel_for(num_planes, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start, size_t end) { + for (size_t p = start; p < end; ++p) { + const __fp16* src_plane = input + p * height * width; + __fp16* dst_plane = output + p * out_plane; + + for (size_t y = 0; y < height; ++y) { + const __fp16* src_row = src_plane + y * width; + __fp16* dst_row = dst_plane + y * scale * out_width; + + for (size_t x = 0; x < width; ++x) { + const __fp16 val = src_row[x]; + for (size_t k = 0; k < scale; ++k) { + dst_row[x * scale + k] = val; + } + } + for (size_t k = 1; k < scale; ++k) { + std::memcpy(dst_row + k * out_width, dst_row, out_width * sizeof(__fp16)); + } + } + } + }); +} diff --git a/cactus-kernels/src/metal_backend.mm b/cactus-kernels/src/metal_backend.mm index 829a11fb6..e77533cf6 100644 --- a/cactus-kernels/src/metal_backend.mm +++ b/cactus-kernels/src/metal_backend.mm @@ -59,7 +59,7 @@ id psoLayerNorm=nil, psoSoftmaxR=nil, psoConv1dK3=nil, psoAttnDense=nil, psoAttnD64=nil; id psoGemmBatch=nil, psoGemmBatchF32A=nil, psoConvDw=nil, psoStridedRows=nil, psoMaskFill=nil; id psoReduceF16=nil, psoReduceF32=nil, psoCumsumF16=nil, psoCumsumF32=nil; - id psoConcat2=nil, psoGatherF32Idx=nil, psoRopeFull=nil, psoMaxpool1d=nil; + id psoConcat2=nil, psoGatherF32Idx=nil, psoRopeFull=nil, psoMaxpool1d=nil, psoUpsampleNearest2d=nil; id psoBilinear=nil, psoConv1dGen=nil, psoConv1dNlcDw=nil, psoConv2d=nil; id psoBatchnorm=nil, psoGroupnorm=nil, psoBiasAddRows=nil, psoEwChain=nil; id psoAttnFlash=nil, psoRmsSimd=nil, psoScatterRows=nil; @@ -122,6 +122,7 @@ psoCumsumF16=pso("cumsum_f16"); psoCumsumF32=pso("cumsum_f32"); psoConcat2=pso("concat2_f16"); psoGatherF32Idx=pso("gather_f32idx_f16"); psoRopeFull=pso("rope_full_f16"); psoMaxpool1d=pso("maxpool1d_f16"); + psoUpsampleNearest2d=pso("upsample_nearest2d_f16"); psoBilinear=pso("bilinear_f16"); psoConv1dGen=pso("conv1d_gen_f16"); psoConv1dNlcDw=pso("conv1d_nlc_dw_f16"); psoConv2d=pso("conv2d_f16"); psoBatchnorm=pso("batchnorm_f16"); psoGroupnorm=pso("groupnorm_f16"); @@ -2067,6 +2068,18 @@ bool cactus_metal_encode_maxpool1d(void* out, const void* in, uint32_t NC, uint3 [g_enc dispatchThreads:MTLSizeMake(Lout, NC, 1) threadsPerThreadgroup:MTLSizeMake(std::min(Lout, 64u), 4, 1)]; return true; } +bool cactus_metal_encode_upsample_nearest2d(void* out, const void* in, uint32_t NC, + uint32_t H, uint32_t W, uint32_t scale) { + if (!ctx().ok || !ctx().psoUpsampleNearest2d || scale == 0) return false; + ensureEncoder(); + uint32_t pixels = H * scale * W * scale; + [g_enc setComputePipelineState:ctx().psoUpsampleNearest2d]; + setBufAt(in, (size_t)NC*H*W*2, 0); setBufAt(out, (size_t)NC*pixels*2, 1); + [g_enc setBytes:&H length:4 atIndex:2]; [g_enc setBytes:&W length:4 atIndex:3]; + [g_enc setBytes:&scale length:4 atIndex:4]; + [g_enc dispatchThreads:MTLSizeMake(pixels, NC, 1) threadsPerThreadgroup:MTLSizeMake(std::min(pixels, 64u), 4, 1)]; + return true; +} bool cactus_metal_encode_bilinear(void* out, const void* in, uint32_t sh, uint32_t sw, uint32_t dh, uint32_t dw, uint32_t E, int align) { if (!ctx().ok || !ctx().psoBilinear) return false; diff --git a/cactus-kernels/src/metal_backend_stub.cpp b/cactus-kernels/src/metal_backend_stub.cpp index 9ecafc05b..b7eb7214d 100644 --- a/cactus-kernels/src/metal_backend_stub.cpp +++ b/cactus-kernels/src/metal_backend_stub.cpp @@ -79,6 +79,7 @@ bool cactus_metal_encode_concat2(void*, const void*, const void*, uint32_t, uint bool cactus_metal_encode_gather_f32idx(void*, const void*, const void*, uint32_t, uint32_t, size_t) { return false; } bool cactus_metal_encode_rope_full(void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, float, int) { return false; } bool cactus_metal_encode_maxpool1d(void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t) { return false; } +bool cactus_metal_encode_upsample_nearest2d(void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t) { return false; } bool cactus_metal_encode_bilinear(void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, int) { return false; } bool cactus_metal_encode_conv1d_gen(void*, const void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, int) { return false; } bool cactus_metal_encode_conv1d_nlc_dw(void*, const void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t) { return false; } diff --git a/docs/cactus_graph.md b/docs/cactus_graph.md index a07a48d51..83aadd91a 100644 --- a/docs/cactus_graph.md +++ b/docs/cactus_graph.md @@ -306,6 +306,12 @@ size_t correction = graph.altup_correct(coefs, innovation, predictions, num_pred size_t interpolated = graph.bilinear_interpolation(pos_embeds, dst_height, dst_width); ``` +#### Nearest-Neighbor Upsampling +```cpp +// Input [N, C, H, W] -> output [N, C, H*scale, W*scale] +size_t upsampled = graph.upsample_nearest2d(input, /*scale_factor=*/2); +``` + #### Sampling ```cpp // Defaults: temperature=0.6, top_p=0.95, top_k=20 From 8fa51be8e24b30f4231b82ca7fb42aead0ade9bf Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sun, 2 Aug 2026 23:29:14 -0700 Subject: [PATCH 02/34] add simdgroup-MMA Metal fast path for conv2d_k3s1p1 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 --- cactus-kernels/src/cactus_kernels.metal | 61 +++++++++++++++++++++++++ cactus-kernels/src/metal_backend.mm | 15 +++++- 2 files changed, 75 insertions(+), 1 deletion(-) diff --git a/cactus-kernels/src/cactus_kernels.metal b/cactus-kernels/src/cactus_kernels.metal index fedb02efe..f00396d91 100644 --- a/cactus-kernels/src/cactus_kernels.metal +++ b/cactus-kernels/src/cactus_kernels.metal @@ -2069,6 +2069,67 @@ kernel void conv1d_nlc_dw_f16(device const half* x [[buffer(0)]], device const h y[((size_t)n*L + l)*C + c] = (half)acc; } +kernel void conv2d_k3s1p1_mma_f16(device const half* x [[buffer(0)]], device const half* w [[buffer(1)]], + device const half* bias [[buffer(2)]], device half* y [[buffer(3)]], + constant uint& Cin [[buffer(4)]], constant uint& H [[buffer(5)]], + constant uint& W [[buffer(6)]], constant uint& Cout [[buffer(7)]], + constant uint& has_bias [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], + uint tl [[thread_index_in_threadgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup half sa[32*8]; + threadgroup half sb[8*64]; + threadgroup float Cs[32*64]; + const uint P = H * W, Kt = Cin * 9u; + const uint p0 = tg.x * 64u, c0 = tg.y * 32u, n = tg.z; + device const half* xin = x + (size_t)n * Cin * P; + + simdgroup_matrix ma; + simdgroup_matrix mb; + simdgroup_matrix mc[8]; + for (uint i = 0; i < 8u; ++i) mc[i] = make_filled_simdgroup_matrix(0.f); + + for (uint k0 = 0; k0 < Kt; k0 += 8u) { + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint idx = tl; idx < 32u*8u; idx += 128u) { + uint m = idx / 8u, kk = idx % 8u; + sa[idx] = (c0 + m < Cout && k0 + kk < Kt) ? w[(size_t)(c0 + m) * Kt + k0 + kk] : (half)0; + } + for (uint idx = tl; idx < 8u*64u; idx += 128u) { + uint kk = idx / 64u, nn = idx % 64u; + uint k = k0 + kk, p = p0 + nn; + half v = (half)0; + if (k < Kt && p < P) { + uint ci = k / 9u, r = k % 9u; + int iy = (int)(p / W) + (int)(r / 3u) - 1; + int ix = (int)(p % W) + (int)(r % 3u) - 1; + if (iy >= 0 && iy < (int)H && ix >= 0 && ix < (int)W) + v = xin[(size_t)ci * P + (size_t)iy * W + ix]; + } + sb[idx] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + simdgroup_load(ma, sa + sg * 8u * 8u, 8u); + for (uint c = 0; c < 8u; ++c) { + simdgroup_load(mb, sb + c * 8u, 64u); + simdgroup_multiply_accumulate(mc[c], ma, mb, mc[c]); + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint c = 0; c < 8u; ++c) + simdgroup_store(mc[c], &Cs[(sg * 8u) * 64u + c * 8u], 64u); + threadgroup_barrier(mem_flags::mem_threadgroup); + device half* yout = y + (size_t)n * Cout * P; + for (uint idx = tl; idx < 32u*64u; idx += 128u) { + uint m = idx / 64u, nn = idx % 64u; + if (c0 + m < Cout && p0 + nn < P) { + float v = Cs[idx] + (has_bias ? (float)bias[c0 + m] : 0.f); + yout[(size_t)(c0 + m) * P + p0 + nn] = (half)v; + } + } +} + kernel void conv2d_f16(device const half* x [[buffer(0)]], device const half* w [[buffer(1)]], device const half* bias [[buffer(2)]], device half* y [[buffer(3)]], constant uint& Cin [[buffer(4)]], constant uint& H [[buffer(5)]], diff --git a/cactus-kernels/src/metal_backend.mm b/cactus-kernels/src/metal_backend.mm index e77533cf6..ed1bce976 100644 --- a/cactus-kernels/src/metal_backend.mm +++ b/cactus-kernels/src/metal_backend.mm @@ -60,7 +60,7 @@ id psoGemmBatch=nil, psoGemmBatchF32A=nil, psoConvDw=nil, psoStridedRows=nil, psoMaskFill=nil; id psoReduceF16=nil, psoReduceF32=nil, psoCumsumF16=nil, psoCumsumF32=nil; id psoConcat2=nil, psoGatherF32Idx=nil, psoRopeFull=nil, psoMaxpool1d=nil, psoUpsampleNearest2d=nil; - id psoBilinear=nil, psoConv1dGen=nil, psoConv1dNlcDw=nil, psoConv2d=nil; + id psoBilinear=nil, psoConv1dGen=nil, psoConv1dNlcDw=nil, psoConv2d=nil, psoConv2dK3S1P1Mma=nil; id psoBatchnorm=nil, psoGroupnorm=nil, psoBiasAddRows=nil, psoEwChain=nil; id psoAttnFlash=nil, psoRmsSimd=nil, psoScatterRows=nil; id psoTranspose2d=nil, psoBcastRows=nil, psoRmsAddSimd=nil; @@ -125,6 +125,7 @@ psoUpsampleNearest2d=pso("upsample_nearest2d_f16"); psoBilinear=pso("bilinear_f16"); psoConv1dGen=pso("conv1d_gen_f16"); psoConv1dNlcDw=pso("conv1d_nlc_dw_f16"); psoConv2d=pso("conv2d_f16"); + psoConv2dK3S1P1Mma=pso("conv2d_k3s1p1_mma_f16"); psoBatchnorm=pso("batchnorm_f16"); psoGroupnorm=pso("groupnorm_f16"); psoBiasAddRows=pso("bias_add_rows_f16"); psoEwChain=pso("elemwise_chain_f16"); psoAttnFlash=pso("attn_flash_f16"); psoRmsSimd=pso("rms_norm_simd_f16"); @@ -2134,6 +2135,18 @@ bool cactus_metal_encode_conv2d(void* out, const void* x, const void* w, const v ensureEncoder(); uint32_t hb = bias ? 1u : 0u, dwf = dw ? 1u : 0u; size_t wbytes = dw ? (size_t)Cout*K*K*2 : (size_t)Cout*Cin*K*K*2; + if (K == 3 && stride == 1 && pad == 1 && !dw && ctx().psoConv2dK3S1P1Mma) { + [g_enc setComputePipelineState:ctx().psoConv2dK3S1P1Mma]; + setBufAt(x, (size_t)N*Cin*H*W*2, 0); setBufAt(w, wbytes, 1); + if (bias) setBufAt(bias, (size_t)Cout*2, 2); else [g_enc setBuffer:ctx().dummy offset:0 atIndex:2]; + setBufAt(out, (size_t)N*Cout*Ho*Wo*2, 3); + [g_enc setBytes:&Cin length:4 atIndex:4]; [g_enc setBytes:&H length:4 atIndex:5]; + [g_enc setBytes:&W length:4 atIndex:6]; [g_enc setBytes:&Cout length:4 atIndex:7]; + [g_enc setBytes:&hb length:4 atIndex:8]; + [g_enc dispatchThreadgroups:MTLSizeMake((H*W + 63)/64, (Cout + 31)/32, N) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + return true; + } [g_enc setComputePipelineState:ctx().psoConv2d]; setBufAt(x, (size_t)N*Cin*H*W*2, 0); setBufAt(w, wbytes, 1); if (bias) setBufAt(bias, (size_t)Cout*2, 2); else [g_enc setBuffer:ctx().dummy offset:0 atIndex:2]; From 780027c82b39cfce81205b2c0a4a63330cf8517f Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Fri, 7 Aug 2026 02:30:16 -0700 Subject: [PATCH 03/34] convert TAESD through the standard weights pipeline 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 --weights-only Signed-off-by: Neel Shejwalkar --- cactus-graph/tests/export_taesd.py | 103 --------- cactus-graph/tests/test_vae.cpp | 196 +++++++++--------- python/cactus/convert/cli.py | 74 +++++-- .../cactus/convert/model_adapters/adapters.py | 32 +++ .../convert/model_adapters/detection.py | 7 +- .../convert/tests/test_taesd_adapter.py | 32 +++ 6 files changed, 226 insertions(+), 218 deletions(-) delete mode 100644 cactus-graph/tests/export_taesd.py create mode 100644 python/cactus/convert/tests/test_taesd_adapter.py diff --git a/cactus-graph/tests/export_taesd.py b/cactus-graph/tests/export_taesd.py deleted file mode 100644 index eac656b15..000000000 --- a/cactus-graph/tests/export_taesd.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Export TAESD decoder weights and reference artifacts for test_vae.cpp. - -Usage: python export_taesd.py [output_dir] [test_image] -Then: CACTUS_TAESD_DIR= ./test_vae - -Requires torch, safetensors, and Pillow. Downloads the MIT-licensed TAESD -weights (https://huggingface.co/madebyollin/taesd) on first run. -""" - -import json -import sys -import urllib.request -from pathlib import Path - -import numpy as np -import torch -import torch.nn as nn -from PIL import Image -from safetensors.torch import load_file - -OUT = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(__file__).parent / "taesd" -ASSET = Path(sys.argv[2]) if len(sys.argv) > 2 else \ - Path(__file__).resolve().parents[2] / "cactus-engine" / "tests" / "assets" / "test_monkey.png" -OUT.mkdir(parents=True, exist_ok=True) - -HF = "https://huggingface.co/madebyollin/taesd/resolve/main" - - -def fetch(name): - dst = OUT / name - if not dst.exists(): - print(f"downloading {name}...") - urllib.request.urlretrieve(f"{HF}/{name}", dst) - return dst - - -def conv(n_in, n_out, **kwargs): - return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs) - - -class Clamp(nn.Module): - def forward(self, x): - return torch.tanh(x / 3) * 3 - - -class Block(nn.Module): - def __init__(self, n_in, n_out): - super().__init__() - self.conv = nn.Sequential(conv(n_in, n_out), nn.ReLU(), conv(n_out, n_out), nn.ReLU(), conv(n_out, n_out)) - self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() - self.fuse = nn.ReLU() - - def forward(self, x): - return self.fuse(self.conv(x) + self.skip(x)) - - -def encoder(): - return nn.Sequential( - conv(3, 64), Block(64, 64), - conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64), - conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64), - conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64), - conv(64, 4), - ) - - -def decoder(): - return nn.Sequential( - Clamp(), conv(4, 64), nn.ReLU(), - Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False), - Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False), - Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False), - Block(64, 64), conv(64, 3), - ) - - -enc = encoder() -enc.load_state_dict(load_file(fetch("taesd_encoder.safetensors"))) -dec = decoder() -dec.load_state_dict(load_file(fetch("taesd_decoder.safetensors"))) -enc.eval() -dec.eval() - -img = Image.open(ASSET).convert("RGB").resize((512, 512), Image.LANCZOS) -x = torch.from_numpy(np.array(img)).float().permute(2, 0, 1).unsqueeze(0) / 255.0 - -with torch.no_grad(): - latent = enc(x) - ref = dec(latent).clamp(0, 1) - -latent.to(torch.float16).numpy().tofile(OUT / "latent.f16.bin") -ref.to(torch.float32).numpy().tofile(OUT / "ref_image.f32.bin") -Image.fromarray((ref[0].permute(1, 2, 0).numpy() * 255).round().astype(np.uint8)).save(OUT / "ref_image.png") - -manifest = [] -with open(OUT / "decoder.f16.bin", "wb") as f: - for name, p in dec.named_parameters(): - f.write(p.detach().to(torch.float16).numpy().tobytes()) - manifest.append({"name": name, "shape": list(p.shape)}) -(OUT / "manifest.json").write_text(json.dumps(manifest, indent=1)) - -print(f"latent shape {list(latent.shape)}, range [{latent.min():.3f}, {latent.max():.3f}]") -print(f"wrote {len(manifest)} tensors to {OUT}") diff --git a/cactus-graph/tests/test_vae.cpp b/cactus-graph/tests/test_vae.cpp index 65383e625..121b3f09d 100644 --- a/cactus-graph/tests/test_vae.cpp +++ b/cactus-graph/tests/test_vae.cpp @@ -1,4 +1,5 @@ #include "test_utils.h" +#include "../../cactus-kernels/cactus_kernels.h" #include #include #include @@ -9,109 +10,125 @@ using namespace TestUtils; namespace { -constexpr size_t LATENT_C = 4, LATENT_HW = 64, IMAGE_C = 3, IMAGE_HW = 512, CH = 64; +constexpr size_t IMAGE_HW = 512, LATENT_HW = 64, LATENT_C = 4, IMAGE_C = 3, CH = 64; -std::vector<__fp16> read_fp16_file(const std::string& path) { - std::ifstream f(path, std::ios::binary | std::ios::ate); - if (!f) return {}; - size_t bytes = static_cast(f.tellg()); - std::vector<__fp16> data(bytes / sizeof(__fp16)); - f.seekg(0); - f.read(reinterpret_cast(data.data()), bytes); - return data; +struct TaesdWeights { + CactusGraph& graph; + std::string dir; + + size_t weight(const std::string& name) { return graph.mmap_weights(dir + "/" + name + ".weights"); } + size_t bias(const std::string& name) { return graph.mmap_weights(dir + "/" + name + ".bias"); } +}; + +size_t conv_biased(TaesdWeights& w, size_t x, const std::string& name) { + return w.graph.conv2d_k3s1p1(x, w.weight(name), w.bias(name)); } -std::vector read_fp32_file(const std::string& path) { - std::ifstream f(path, std::ios::binary | std::ios::ate); - if (!f) return {}; - size_t bytes = static_cast(f.tellg()); - std::vector data(bytes / sizeof(float)); - f.seekg(0); - f.read(reinterpret_cast(data.data()), bytes); - return data; +size_t block(TaesdWeights& w, size_t x, const std::string& prefix) { + size_t h = w.graph.relu(conv_biased(w, x, prefix + "_conv_0")); + h = w.graph.relu(conv_biased(w, h, prefix + "_conv_2")); + h = conv_biased(w, h, prefix + "_conv_4"); + return w.graph.relu(w.graph.add(h, x)); } -struct WeightStream { - const std::vector<__fp16>& blob; - size_t offset = 0; - - size_t next(CactusGraph& g, const std::vector& shape) { - size_t n = 1; - for (size_t d : shape) n *= d; - size_t id = g.input(shape, Precision::FP16); - g.set_input(id, blob.data() + offset, Precision::FP16); - offset += n; - return id; +size_t build_encoder(TaesdWeights& w, size_t image) { + size_t x = conv_biased(w, image, "encoder_layers_0"); + x = block(w, x, "encoder_layers_1"); + for (int stage = 0; stage < 3; ++stage) { + const int down = 2 + stage * 4; + x = w.graph.conv2d_k3s2p1(x, w.weight("encoder_layers_" + std::to_string(down))); + for (int j = 1; j <= 3; ++j) x = block(w, x, "encoder_layers_" + std::to_string(down + j)); } -}; - -size_t taesd_conv(CactusGraph& g, WeightStream& w, size_t x, size_t c_in, size_t c_out, bool bias = true) { - size_t weight = w.next(g, {c_out, c_in, 3, 3}); - if (!bias) return g.conv2d_k3s1p1(x, weight); - size_t b = w.next(g, {c_out}); - return g.conv2d_k3s1p1(x, weight, b); + return conv_biased(w, x, "encoder_layers_14"); } -size_t taesd_block(CactusGraph& g, WeightStream& w, size_t x) { - size_t h = g.relu(taesd_conv(g, w, x, CH, CH)); - h = g.relu(taesd_conv(g, w, h, CH, CH)); - h = taesd_conv(g, w, h, CH, CH); - return g.relu(g.add(h, x)); +size_t build_decoder(TaesdWeights& w, size_t latent) { + size_t x = w.graph.scalar_multiply(w.graph.tanh(w.graph.scalar_divide(latent, 3.0f)), 3.0f); + x = w.graph.relu(conv_biased(w, x, "decoder_layers_0")); + for (int stage = 0; stage < 3; ++stage) { + const int first = 2 + stage * 5; + for (int j = 0; j < 3; ++j) x = block(w, x, "decoder_layers_" + std::to_string(first + j)); + x = w.graph.upsample_nearest2d(x, 2); + x = w.graph.conv2d_k3s1p1(x, w.weight("decoder_layers_" + std::to_string(first + 4))); + } + x = block(w, x, "decoder_layers_17"); + return w.graph.clamp(conv_biased(w, x, "decoder_layers_18"), 0.0f, 1.0f); } -size_t build_taesd_decoder(CactusGraph& g, WeightStream& w, size_t latent) { - size_t x = g.scalar_multiply(g.tanh(g.scalar_divide(latent, 3.0f)), 3.0f); - x = g.relu(taesd_conv(g, w, x, LATENT_C, CH)); - for (int stage = 0; stage < 3; ++stage) { - x = taesd_block(g, w, x); - x = taesd_block(g, w, x); - x = taesd_block(g, w, x); - x = g.upsample_nearest2d(x, 2); - x = taesd_conv(g, w, x, CH, CH, false); +std::vector<__fp16> load_image_chw(const std::string& path, std::vector& rgb) { + int w = 0, h = 0, c = 0; + unsigned char* pixels = cactus_image_load(path.c_str(), &w, &h, &c, 3); + if (!pixels) return {}; + rgb.assign(IMAGE_HW * IMAGE_HW * IMAGE_C, 0); + if (w == (int)IMAGE_HW && h == (int)IMAGE_HW) { + std::copy(pixels, pixels + rgb.size(), rgb.begin()); + } else { + cactus_image_resize_uint8(pixels, w, h, rgb.data(), IMAGE_HW, IMAGE_HW, IMAGE_C); } - x = taesd_block(g, w, x); - x = taesd_conv(g, w, x, CH, IMAGE_C); - return g.clamp(x, 0.0f, 1.0f); + cactus_image_free(pixels); + + std::vector<__fp16> chw(rgb.size()); + const size_t plane = IMAGE_HW * IMAGE_HW; + for (size_t i = 0; i < plane; ++i) + for (size_t ch = 0; ch < IMAGE_C; ++ch) + chw[ch * plane + i] = static_cast<__fp16>(rgb[i * IMAGE_C + ch] / 255.0f); + return chw; } void write_ppm(const std::string& path, const __fp16* chw) { std::ofstream f(path, std::ios::binary); f << "P6\n" << IMAGE_HW << " " << IMAGE_HW << "\n255\n"; const size_t plane = IMAGE_HW * IMAGE_HW; - for (size_t i = 0; i < plane; ++i) { + for (size_t i = 0; i < plane; ++i) for (size_t c = 0; c < IMAGE_C; ++c) { float v = static_cast(chw[c * plane + i]) * 255.0f; f.put(static_cast(v < 0.0f ? 0 : (v > 255.0f ? 255 : static_cast(v + 0.5f)))); } - } } -float psnr_vs_reference(const __fp16* out, const std::vector& ref) { +float psnr_vs_source(const __fp16* out, const std::vector& rgb) { + const size_t plane = IMAGE_HW * IMAGE_HW; double se = 0.0; - for (size_t i = 0; i < ref.size(); ++i) { - double d = static_cast(out[i]) - static_cast(ref[i]); - se += d * d; - } - double mse = se / ref.size(); + for (size_t i = 0; i < plane; ++i) + for (size_t c = 0; c < IMAGE_C; ++c) { + double d = static_cast(out[c * plane + i]) - rgb[i * IMAGE_C + c] / 255.0; + se += d * d; + } + double mse = se / (plane * IMAGE_C); return mse > 0.0 ? static_cast(10.0 * std::log10(1.0 / mse)) : 99.0f; } -bool decode_and_check(const std::string& dir, const char* backend, const std::vector<__fp16>& weights, - const std::vector<__fp16>& latent, const std::vector& ref) { +bool roundtrip(const std::string& dir, const std::string& image_path, const char* backend) { if (cactus_backend_select(backend) != 0) { std::cout << " [" << backend << "] unavailable, skipping" << std::endl; return true; } - CactusGraph g; - size_t latent_id = g.input({1, LATENT_C, LATENT_HW, LATENT_HW}, Precision::FP16); - WeightStream w{weights}; - size_t out_id = build_taesd_decoder(g, w, latent_id); - if (w.offset != weights.size()) { - std::cout << " weight blob mismatch: consumed " << w.offset << " of " << weights.size() << std::endl; + std::vector rgb; + auto image = load_image_chw(image_path, rgb); + if (image.empty()) { + std::cout << " could not read " << image_path << ": " << cactus_image_failure_reason() << std::endl; return false; } - g.set_input(latent_id, latent.data(), Precision::FP16); + + std::vector<__fp16> latent(LATENT_C * LATENT_HW * LATENT_HW); + { + CactusGraph g; + TaesdWeights w{g, dir}; + size_t input = g.input({1, IMAGE_C, IMAGE_HW, IMAGE_HW}, Precision::FP16); + size_t out = build_encoder(w, input); + g.set_input(input, image.data(), Precision::FP16); + g.execute(); + const __fp16* data = static_cast(g.get_output(out)); + std::copy(data, data + latent.size(), latent.begin()); + g.hard_reset(); + } + + CactusGraph g; + TaesdWeights w{g, dir}; + size_t input = g.input({1, LATENT_C, LATENT_HW, LATENT_HW}, Precision::FP16); + size_t out = build_decoder(w, input); + g.set_input(input, latent.data(), Precision::FP16); g.execute(); constexpr int runs = 5; @@ -119,43 +136,36 @@ bool decode_and_check(const std::string& dir, const char* backend, const std::ve for (int i = 0; i < runs; ++i) g.execute(); double ms = timer.elapsed_ms() / runs; - const __fp16* out = static_cast(g.get_output(out_id)); - float psnr = psnr_vs_reference(out, ref); - write_ppm(dir + "/cactus_decode_" + backend + ".ppm", out); - std::cout << " [" << backend << "] " << ms << " ms/decode, PSNR vs torch fp32: " << psnr << " dB" << std::endl; + const __fp16* decoded = static_cast(g.get_output(out)); + float quality = psnr_vs_source(decoded, rgb); + write_ppm(dir + "/taesd_roundtrip_" + backend + ".ppm", decoded); + std::cout << " [" << backend << "] decode " << ms << " ms, roundtrip PSNR " << quality << " dB" << std::endl; g.hard_reset(); - return psnr > 30.0f; + return quality > 20.0f; } -bool test_taesd_decoder() { - const char* env = std::getenv("CACTUS_TAESD_DIR"); - if (!env) { - std::cout << " skipped: set CACTUS_TAESD_DIR to a dir with decoder.f16.bin/latent.f16.bin/ref_image.f32.bin" - << std::endl; +bool test_taesd_roundtrip() { + const char* bundle = std::getenv("CACTUS_TAESD_DIR"); + if (!bundle) { + std::cout << " skipped: convert a bundle first, then set CACTUS_TAESD_DIR\n" + << " cactus convert madebyollin/taesd --weights-only" << std::endl; return true; } - std::string dir(env); - - auto weights = read_fp16_file(dir + "/decoder.f16.bin"); - auto latent = read_fp16_file(dir + "/latent.f16.bin"); - auto ref = read_fp32_file(dir + "/ref_image.f32.bin"); - if (weights.empty() || latent.size() != LATENT_C * LATENT_HW * LATENT_HW || - ref.size() != IMAGE_C * IMAGE_HW * IMAGE_HW) { - std::cout << " missing or malformed TAESD artifacts in " << dir << std::endl; - return false; - } + const char* assets = std::getenv("CACTUS_TEST_ASSETS"); + std::string image_path = assets ? std::string(assets) + "/test_monkey.png" + : "../../cactus-engine/tests/assets/test_monkey.png"; - bool ok = decode_and_check(dir, "cpu", weights, latent, ref); - ok = decode_and_check(dir, "metal", weights, latent, ref) && ok; + bool ok = roundtrip(bundle, image_path, "cpu"); + ok = roundtrip(bundle, image_path, "metal") && ok; return ok; } } int main() { - TestRunner runner("VAE Decoder Tests"); - runner.run_test("taesd_decoder", test_taesd_decoder()); + TestRunner runner("VAE Tests"); + runner.run_test("taesd_encode_decode_roundtrip", test_taesd_roundtrip()); runner.print_summary(); return runner.all_passed() ? 0 : 1; } diff --git a/python/cactus/convert/cli.py b/python/cactus/convert/cli.py index 60a85bce8..b5f008236 100644 --- a/python/cactus/convert/cli.py +++ b/python/cactus/convert/cli.py @@ -78,6 +78,8 @@ def _load_hf(model_id_or_path: str, device: str, *, skip_model_load: bool = Fals if skip_model_load: return cfg, processor, None model_cls = adapter.model_class(cfg) + if model_cls is None: + return cfg, processor, None model_type = str(cfg_get(cfg, "model_type", "") or "").lower() if isinstance(cfg, dict) and model_type == "parakeet_tdt": return cfg, processor, None @@ -136,36 +138,53 @@ def _bits_for_component(component: str, args: argparse.Namespace) -> int: return int(args.bits) +CHECKPOINT_STEMS = ("model", "diffusion_pytorch_model") + + +def _download_checkpoint_state_dict(model_id_or_path: str, stem: str) -> dict[str, Any] | None: + from huggingface_hub import hf_hub_download + from safetensors.torch import load_file + + try: + return load_file(hf_hub_download(model_id_or_path, f"{stem}.safetensors", cache_dir=_hf_cache_dir())) + except RuntimeError: + raise + except Exception: + pass + try: + index_file = hf_hub_download(model_id_or_path, f"{stem}.safetensors.index.json", cache_dir=_hf_cache_dir()) + except Exception: + return None + index = json.loads(Path(index_file).read_text(encoding="utf-8")) + shard_names = sorted(set(index.get("weight_map", {}).values())) + if not shard_names: + return None + state: dict[str, Any] = {} + for shard in shard_names: + shard_path = hf_hub_download(model_id_or_path, shard, cache_dir=_hf_cache_dir()) + for key, tensor in load_file(shard_path).items(): + if key in state: + raise RuntimeError(f"duplicate tensor key {key!r} across checkpoint shards") + state[key] = tensor + return state + + def _load_checkpoint_state_dict(model_id_or_path: str) -> dict[str, Any] | None: nemo_export = ensure_parakeet_tdt_nemo_source(model_id_or_path, cache_dir=_hf_cache_dir()) if nemo_export is not None: model_id_or_path = nemo_export root = Path(model_id_or_path) if not root.exists() or not root.is_dir(): - try: - from huggingface_hub import hf_hub_download - from safetensors.torch import load_file + for stem in CHECKPOINT_STEMS: try: - model_file = hf_hub_download(model_id_or_path, "model.safetensors", cache_dir=_hf_cache_dir()) - return load_file(model_file) + state = _download_checkpoint_state_dict(model_id_or_path, stem) + except RuntimeError: + raise except Exception: - index_file = hf_hub_download(model_id_or_path, "model.safetensors.index.json", cache_dir=_hf_cache_dir()) - index = json.loads(Path(index_file).read_text(encoding="utf-8")) - shard_names = sorted(set(index.get("weight_map", {}).values())) - if not shard_names: - return None - state: dict[str, Any] = {} - for shard in shard_names: - shard_path = hf_hub_download(model_id_or_path, shard, cache_dir=_hf_cache_dir()) - for key, tensor in load_file(shard_path).items(): - if key in state: - raise RuntimeError(f"duplicate tensor key {key!r} across checkpoint shards") - state[key] = tensor + return None + if state is not None: return state - except RuntimeError: - raise - except Exception: - return None + return None safetensor_files = sorted(root.glob("*.safetensors")) if safetensor_files: try: @@ -604,6 +623,19 @@ def convert(args: argparse.Namespace) -> None: "hessian_missing_reason": hessian_missing_reason, }) + written: dict[str, str] = {} + for row in rows: + output_file = str(row.get("output_file") or "") + if not output_file: + continue + previous = written.get(output_file) + if previous is not None: + raise RuntimeError( + f"tensors {previous!r} and {row['source_name']!r} both map to {output_file!r}; " + "the adapter for this family must give them distinct names" + ) + written[output_file] = str(row["source_name"]) + summary = write_reports(out_dir, rows) print_summary(summary) unrecognized = [r["source_name"] for r in rows if r["status"] == "unrecognized"] diff --git a/python/cactus/convert/model_adapters/adapters.py b/python/cactus/convert/model_adapters/adapters.py index 517e4fa5c..047bb240a 100644 --- a/python/cactus/convert/model_adapters/adapters.py +++ b/python/cactus/convert/model_adapters/adapters.py @@ -502,6 +502,37 @@ def build_calibration_inputs(self, row: dict[str, Any], processor: Any, modality return None +class TaesdAdapter(FamilyAdapter): + family = "taesd" + + def runtime_config(self, cfg: Any) -> dict[str, Any]: + return { + "latent_channels": int(_cfg_get(cfg, "latent_channels", 4) or 4), + "latent_magnitude": float(_cfg_get(cfg, "latent_magnitude", 3) or 3), + "latent_shift": float(_cfg_get(cfg, "latent_shift", 0.5) or 0.5), + "upsampling_scaling_factor": int(_cfg_get(cfg, "upsampling_scaling_factor", 2) or 2), + } + + def model_class(self, cfg: Any): + return None + + def load_processor(self, model_id_or_path: str): + return None + + def name_tensor(self, source_name: str, _tensor: Any, _num_layers: int | None) -> NameMatch: + component = "vision" + if not (source_name.endswith(".weight") or source_name.endswith(".bias")): + return NameMatch(source_name, None, component, False, hf_name=source_name, adapter_name=source_name) + generic = source_name.replace(".", "_") + output = generic[:-7] + ".weights" if generic.endswith("_weight") else generic[:-5] + ".bias" + return NameMatch(source_name, output, component, True, hf_name=source_name, adapter_name=source_name) + + def policy(self, match: NameMatch, _shape: tuple[int, ...], _requested_bits: int) -> TensorPolicy: + if match.output_name is None: + return TensorPolicy("ignored", "none", None, match.component, False, "none", "no output filename") + return TensorPolicy("fallback", "FP16", None, match.component, False, "none", "tiny autoencoder conv tensor") + + class ParakeetAdapter(FamilyAdapter): family = "parakeet" @@ -884,6 +915,7 @@ def scale_factor(self, output_name: str) -> float: "whisper": WhisperAdapter(), "parakeet": ParakeetAdapter(), "parakeet_tdt": ParakeetTDTAdapter(), + "taesd": TaesdAdapter(), } diff --git a/python/cactus/convert/model_adapters/detection.py b/python/cactus/convert/model_adapters/detection.py index 69cb2aa09..050e1f5a3 100644 --- a/python/cactus/convert/model_adapters/detection.py +++ b/python/cactus/convert/model_adapters/detection.py @@ -3,12 +3,17 @@ from ..cactus_adapters.config_utils import cfg_get, detect_model_type -SUPPORTED_FAMILIES = {"auto", "gemma3", "gemma4", "qwen", "lfm2", "whisper", "parakeet", "parakeet_tdt", "moonshine", "nomic", "needle", "generic"} +SUPPORTED_FAMILIES = {"auto", "gemma3", "gemma4", "qwen", "lfm2", "whisper", "parakeet", "parakeet_tdt", "moonshine", "nomic", "needle", "taesd", "generic"} + +DIFFUSERS_FAMILIES = {"AutoencoderTiny": "taesd"} def detect_family(config, requested: str = "auto") -> str: if requested != "auto": return requested + diffusers_class = str(cfg_get(config, "_class_name", "") or "") + if diffusers_class in DIFFUSERS_FAMILIES: + return DIFFUSERS_FAMILIES[diffusers_class] text_config = cfg_get(config, "text_config", None) base = text_config if text_config is not None else config detected = detect_model_type(base, config) diff --git a/python/cactus/convert/tests/test_taesd_adapter.py b/python/cactus/convert/tests/test_taesd_adapter.py new file mode 100644 index 000000000..4a46bb56f --- /dev/null +++ b/python/cactus/convert/tests/test_taesd_adapter.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +import torch + +from cactus.convert.model_adapters.adapters import adapter_for_family +from cactus.convert.model_adapters.detection import detect_family + + +def test_diffusers_autoencoder_tiny_config_selects_taesd(): + assert detect_family({"_class_name": "AutoencoderTiny", "latent_channels": 4}) == "taesd" + + +def test_taesd_converts_from_checkpoint_without_a_transformers_model_class(): + adapter = adapter_for_family("taesd") + assert adapter.model_class({"_class_name": "AutoencoderTiny"}) is None + assert adapter.load_processor("madebyollin/taesd") is None + + +def test_taesd_keeps_encoder_and_decoder_names_distinct(): + adapter = adapter_for_family("taesd") + tensor = torch.zeros(64, 64, 3, 3) + encoder = adapter.name_tensor("encoder.layers.12.conv.0.weight", tensor, None) + decoder = adapter.name_tensor("decoder.layers.12.conv.0.weight", tensor, None) + assert encoder.output_name == "encoder_layers_12_conv_0.weights" + assert decoder.output_name == "decoder_layers_12_conv_0.weights" + + +def test_taesd_conv_tensors_stay_fp16(): + adapter = adapter_for_family("taesd") + match = adapter.name_tensor("decoder.layers.0.weight", torch.zeros(64, 4, 3, 3), None) + policy = adapter.policy(match, (64, 4, 3, 3), 4) + assert policy.precision == "FP16" From 4bb0426b2b1b2d5dcd2dd38b25ca2ba67e4a7388 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Fri, 7 Aug 2026 18:34:59 -0700 Subject: [PATCH 04/34] prepare VAE test weights through cactus test 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 --- cactus-graph/test.sh | 4 ++++ cactus-graph/tests/test_vae.cpp | 8 ++++---- python/cactus/cli/common.py | 1 + python/cactus/cli/test.py | 20 +++++++++++++++++++- 4 files changed, 28 insertions(+), 5 deletions(-) diff --git a/cactus-graph/test.sh b/cactus-graph/test.sh index 4bf78b6d8..abec2ef32 100755 --- a/cactus-graph/test.sh +++ b/cactus-graph/test.sh @@ -7,10 +7,14 @@ SUITE="" while [[ $# -gt 0 ]]; do case $1 in --suite) SUITE="${2:?--suite needs an argument}"; shift 2 ;; + --vae-model) CACTUS_TEST_VAE_MODEL="${2:?--vae-model needs an argument}"; shift 2 ;; *) echo "Unknown arg: $1" >&2; exit 2 ;; esac done +export CACTUS_TEST_VAE_MODEL="${CACTUS_TEST_VAE_MODEL:-}" +export CACTUS_TEST_ASSETS="$(pwd)/../cactus-engine/tests/assets" + echo "Building and testing cactus-graph..." rm -rf build diff --git a/cactus-graph/tests/test_vae.cpp b/cactus-graph/tests/test_vae.cpp index 121b3f09d..2ad8031c8 100644 --- a/cactus-graph/tests/test_vae.cpp +++ b/cactus-graph/tests/test_vae.cpp @@ -146,10 +146,10 @@ bool roundtrip(const std::string& dir, const std::string& image_path, const char } bool test_taesd_roundtrip() { - const char* bundle = std::getenv("CACTUS_TAESD_DIR"); - if (!bundle) { - std::cout << " skipped: convert a bundle first, then set CACTUS_TAESD_DIR\n" - << " cactus convert madebyollin/taesd --weights-only" << std::endl; + const char* bundle = std::getenv("CACTUS_TEST_VAE_MODEL"); + if (!bundle || !*bundle) { + std::cout << " skipped: no VAE weights. Run graph tests via 'cactus test', which prepares them." + << std::endl; return true; } const char* assets = std::getenv("CACTUS_TEST_ASSETS"); diff --git a/python/cactus/cli/common.py b/python/cactus/cli/common.py index f39ef435e..214c79f3d 100644 --- a/python/cactus/cli/common.py +++ b/python/cactus/cli/common.py @@ -31,6 +31,7 @@ def weights_root() -> Path: DEFAULT_TRANSCRIPTION_MODEL_ID = "nvidia/parakeet-tdt-0.6b-v3" DEFAULT_TEST_MODEL_ID = DEFAULT_MODEL_ID DEFAULT_TEST_TRANSCRIPTION_MODEL_ID = DEFAULT_TRANSCRIPTION_MODEL_ID +DEFAULT_TEST_VAE_MODEL_ID = "madebyollin/taesd" RED = '\033[0;31m' diff --git a/python/cactus/cli/test.py b/python/cactus/cli/test.py index 720f9268b..f97150638 100644 --- a/python/cactus/cli/test.py +++ b/python/cactus/cli/test.py @@ -3,7 +3,8 @@ from .common import ( BLUE, DEFAULT_TEST_MODEL_ID, DEFAULT_TEST_TRANSCRIPTION_MODEL_ID, - PROJECT_ROOT, RED, YELLOW, apply_cloud_api_key_env, print_color, + DEFAULT_TEST_VAE_MODEL_ID, PROJECT_ROOT, RED, YELLOW, + apply_cloud_api_key_env, print_color, ) COMPONENTS = ("kernels", "graph", "engine", "all") @@ -28,6 +29,8 @@ def _component_args(component, args): cmd = [str(PROJECT_ROOT / f"cactus-{component}" / "test.sh")] if args.suite: cmd.extend(["--suite", args.suite]) + if component == "graph" and getattr(args, "vae_model_id", None): + cmd.extend(["--vae-model", args.vae_model_id]) if component == "engine": cmd.extend(["--model", args.model_id]) cmd.extend(["--transcription-model", args.transcription_model_id]) @@ -74,6 +77,21 @@ def cmd_test(args): else: targets = ("kernels", "graph", "engine") if args.component == "all" else (args.component,) + if "graph" in targets: + from .model import ensure_weights + + try: + args.vae_model_id = str(ensure_weights( + getattr(args, "vae_model_id", None) or DEFAULT_TEST_VAE_MODEL_ID, + bits=getattr(args, "bits", 4), + token=getattr(args, "token", None), + reconvert=getattr(args, "reconvert", False), + skip_model_load=True, + )) + except (RuntimeError, OSError, ValueError) as exc: + print_color(YELLOW, f"VAE weights unavailable ({exc}); the vae suite will be skipped") + args.vae_model_id = None + if "engine" in targets: from .model import prepare_bundle model_dir = prepare_bundle(args, model_id=args.model_id or DEFAULT_TEST_MODEL_ID, From d4b854380a73415e27aae8198ab5db6a89111f23 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Fri, 7 Aug 2026 18:49:52 -0700 Subject: [PATCH 05/34] test upsample_nearest2d against a hand-computed reference 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 --- cactus-kernels/tests/test_conv.cpp | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/cactus-kernels/tests/test_conv.cpp b/cactus-kernels/tests/test_conv.cpp index c6ae32682..2086b095f 100644 --- a/cactus-kernels/tests/test_conv.cpp +++ b/cactus-kernels/tests/test_conv.cpp @@ -119,6 +119,34 @@ bool test_maxpool1d() { return true; } +bool test_upsample_nearest2d() { + const size_t planes = 6, height = 7, width = 5; + std::vector<__fp16> input(planes * height * width); + fill_random_fp16(input, -5.0f, 5.0f); + + for (size_t scale = 1; scale <= 3; scale++) { + const size_t out_h = height * scale, out_w = width * scale; + std::vector<__fp16> output(planes * out_h * out_w, static_cast<__fp16>(0)); + + cactus_upsample_nearest2d_f16(input.data(), output.data(), planes, height, width, scale); + + for (size_t p = 0; p < planes; p++) { + for (size_t y = 0; y < out_h; y++) { + for (size_t x = 0; x < out_w; x++) { + float expected = static_cast(input[p * height * width + (y / scale) * width + x / scale]); + float actual = static_cast(output[p * out_h * out_w + y * out_w + x]); + if (actual != expected) { + std::cerr << " upsample_nearest2d mismatch [scale=" << scale << ",p=" << p + << ",y=" << y << ",x=" << x << "]: " << actual << " vs " << expected << "\n"; + return false; + } + } + } + } + } + return true; +} + bool run_benchmarks() { { const size_t N = 1, L = 3000, C_in = 80, C_out = 512, stride = 1; @@ -179,6 +207,7 @@ int main() { runner.run_test("conv1d_causal_depthwise_channel_first", test_conv1d_causal_depthwise_channel_first()); runner.run_test("stft_complex", test_stft_complex()); runner.run_test("maxpool1d", test_maxpool1d()); + runner.run_test("upsample_nearest2d", test_upsample_nearest2d()); runner.print_benchmarks_header(); runner.run_bench("benchmarks", run_benchmarks()); runner.print_summary(); From bf4e6d293dcd9af4e5282f8ebacb2e7a22db506e Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Fri, 7 Aug 2026 20:22:06 -0700 Subject: [PATCH 06/34] inline the diffusers family check 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 --- python/cactus/convert/model_adapters/detection.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/python/cactus/convert/model_adapters/detection.py b/python/cactus/convert/model_adapters/detection.py index 050e1f5a3..1b8720a72 100644 --- a/python/cactus/convert/model_adapters/detection.py +++ b/python/cactus/convert/model_adapters/detection.py @@ -5,15 +5,12 @@ SUPPORTED_FAMILIES = {"auto", "gemma3", "gemma4", "qwen", "lfm2", "whisper", "parakeet", "parakeet_tdt", "moonshine", "nomic", "needle", "taesd", "generic"} -DIFFUSERS_FAMILIES = {"AutoencoderTiny": "taesd"} - def detect_family(config, requested: str = "auto") -> str: if requested != "auto": return requested - diffusers_class = str(cfg_get(config, "_class_name", "") or "") - if diffusers_class in DIFFUSERS_FAMILIES: - return DIFFUSERS_FAMILIES[diffusers_class] + if cfg_get(config, "_class_name", "") == "AutoencoderTiny": + return "taesd" text_config = cfg_get(config, "text_config", None) base = text_config if text_config is not None else config detected = detect_model_type(base, config) From 4e483fae0c62105cd3b7e2f3ceaa9bd216c6f2a1 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Fri, 7 Aug 2026 22:02:16 -0700 Subject: [PATCH 07/34] try every checkpoint stem when one fails 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 --- python/cactus/convert/cli.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/cactus/convert/cli.py b/python/cactus/convert/cli.py index b5f008236..01642942c 100644 --- a/python/cactus/convert/cli.py +++ b/python/cactus/convert/cli.py @@ -147,8 +147,6 @@ def _download_checkpoint_state_dict(model_id_or_path: str, stem: str) -> dict[st try: return load_file(hf_hub_download(model_id_or_path, f"{stem}.safetensors", cache_dir=_hf_cache_dir())) - except RuntimeError: - raise except Exception: pass try: @@ -181,7 +179,7 @@ def _load_checkpoint_state_dict(model_id_or_path: str) -> dict[str, Any] | None: except RuntimeError: raise except Exception: - return None + continue if state is not None: return state return None From c17def519aed3f028fe858fd270bec2cf7f83c21 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Fri, 7 Aug 2026 22:11:51 -0700 Subject: [PATCH 08/34] reject colliding tensor names before writing, not after 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 --- python/cactus/convert/cli.py | 22 +++++++++------------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/python/cactus/convert/cli.py b/python/cactus/convert/cli.py index 01642942c..c52b7f29b 100644 --- a/python/cactus/convert/cli.py +++ b/python/cactus/convert/cli.py @@ -474,6 +474,7 @@ def convert(args: argparse.Namespace) -> None: "embedding": args.max_embedding_examples, } rows: list[dict[str, Any]] = [] + written: dict[str, str] = {} pending: list[tuple[str, Any, Any, Any]] = [] num_layers = max( int(model_config.get("num_layers", 0) or 0), @@ -540,6 +541,14 @@ def convert(args: argparse.Namespace) -> None: for emission in emissions: qdq_restore = getattr(emission, "qdq_restore", None) or (provenance.qdq_restore if provenance else "hf_key") out_path = out_dir / emission.output_name if emission.output_name else None + if out_path is not None: + claimed_by = written.get(out_path.name) + if claimed_by is not None: + raise RuntimeError( + f"tensors {claimed_by!r} and {name!r} both map to {out_path.name!r}; " + "the adapter for this family must give them distinct names" + ) + written[out_path.name] = name emit_tensor = emission.tensor emit_match = replace(match, output_name=emission.output_name) requested_bits = _bits_for_component(emit_match.component, args) @@ -621,19 +630,6 @@ def convert(args: argparse.Namespace) -> None: "hessian_missing_reason": hessian_missing_reason, }) - written: dict[str, str] = {} - for row in rows: - output_file = str(row.get("output_file") or "") - if not output_file: - continue - previous = written.get(output_file) - if previous is not None: - raise RuntimeError( - f"tensors {previous!r} and {row['source_name']!r} both map to {output_file!r}; " - "the adapter for this family must give them distinct names" - ) - written[output_file] = str(row["source_name"]) - summary = write_reports(out_dir, rows) print_summary(summary) unrecognized = [r["source_name"] for r in rows if r["status"] == "unrecognized"] From 8e69c9813a0077ac8a8d5492ee97a66633ea93de Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Fri, 7 Aug 2026 22:43:06 -0700 Subject: [PATCH 09/34] pass the VAE weights through the test environment 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 --- cactus-graph/test.sh | 4 ---- python/cactus/cli/test.py | 4 ++-- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/cactus-graph/test.sh b/cactus-graph/test.sh index abec2ef32..4bf78b6d8 100755 --- a/cactus-graph/test.sh +++ b/cactus-graph/test.sh @@ -7,14 +7,10 @@ SUITE="" while [[ $# -gt 0 ]]; do case $1 in --suite) SUITE="${2:?--suite needs an argument}"; shift 2 ;; - --vae-model) CACTUS_TEST_VAE_MODEL="${2:?--vae-model needs an argument}"; shift 2 ;; *) echo "Unknown arg: $1" >&2; exit 2 ;; esac done -export CACTUS_TEST_VAE_MODEL="${CACTUS_TEST_VAE_MODEL:-}" -export CACTUS_TEST_ASSETS="$(pwd)/../cactus-engine/tests/assets" - echo "Building and testing cactus-graph..." rm -rf build diff --git a/python/cactus/cli/test.py b/python/cactus/cli/test.py index f97150638..11603e64d 100644 --- a/python/cactus/cli/test.py +++ b/python/cactus/cli/test.py @@ -29,8 +29,6 @@ def _component_args(component, args): cmd = [str(PROJECT_ROOT / f"cactus-{component}" / "test.sh")] if args.suite: cmd.extend(["--suite", args.suite]) - if component == "graph" and getattr(args, "vae_model_id", None): - cmd.extend(["--vae-model", args.vae_model_id]) if component == "engine": cmd.extend(["--model", args.model_id]) cmd.extend(["--transcription-model", args.transcription_model_id]) @@ -111,6 +109,8 @@ def cmd_test(args): env.pop("CACTUS_NO_CLOUD_TELE", None) else: env["CACTUS_NO_CLOUD_TELE"] = "1" + if getattr(args, "vae_model_id", None): + env["CACTUS_TEST_VAE_MODEL"] = args.vae_model_id for c in targets: cmd = _component_args(c, args) From ba6db4c580173ba763b010afbeb3ea82e0c48567 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sat, 8 Aug 2026 04:06:06 -0700 Subject: [PATCH 10/34] transpile nearest-neighbour upsampling 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 --- cactus-graph/cactus_graph.h | 4 +++ cactus-graph/src/graph_ffi.cpp | 26 +++++++++++++++ python/cactus/bindings/cactus.py | 15 +++++++++ python/cactus/transpile/aten_ops.py | 1 + python/cactus/transpile/importers.py | 32 +++++++++++++++++++ python/cactus/transpile/lower.py | 4 +++ python/cactus/transpile/ops.py | 1 + python/tests/test_transpile_upsample.py | 42 +++++++++++++++++++++++++ 8 files changed, 125 insertions(+) create mode 100644 python/tests/test_transpile_upsample.py diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 655121e77..f7c401bd1 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -1239,6 +1239,10 @@ CACTUS_FFI_EXPORT int cactus_graph_conv2d_depthwise_k3s2p1( cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_conv2d_pointwise_1x1( cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_conv2d_k3s1p1( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_upsample_nearest2d( + cactus_graph_t graph, cactus_node_t input, size_t scale_factor, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_lstm_cell( cactus_graph_t graph, cactus_node_t input, cactus_node_t h_prev, cactus_node_t c_prev, cactus_node_t weight_ih, cactus_node_t weight_hh, cactus_node_t bias_ih, cactus_node_t bias_hh, cactus_node_t* out); diff --git a/cactus-graph/src/graph_ffi.cpp b/cactus-graph/src/graph_ffi.cpp index 223ce8c84..bb87e062a 100644 --- a/cactus-graph/src/graph_ffi.cpp +++ b/cactus-graph/src/graph_ffi.cpp @@ -1574,6 +1574,32 @@ int cactus_graph_conv2d_pointwise_1x1(cactus_graph_t graph, cactus_node_t input, } } +int cactus_graph_conv2d_k3s1p1(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv2d_k3s1p1"); + try { + if (has_bias) { + *out = static_cast(as_graph(graph)->graph.conv2d_k3s1p1(static_cast(input), static_cast(weight), static_cast(bias))); + } else { + *out = static_cast(as_graph(graph)->graph.conv2d_k3s1p1(static_cast(input), static_cast(weight))); + } + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_upsample_nearest2d(cactus_graph_t graph, cactus_node_t input, size_t scale_factor, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_upsample_nearest2d"); + try { + *out = static_cast(as_graph(graph)->graph.upsample_nearest2d(static_cast(input), scale_factor)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_lstm_cell(cactus_graph_t graph, cactus_node_t input, cactus_node_t h_prev, cactus_node_t c_prev, cactus_node_t weight_ih, cactus_node_t weight_hh, cactus_node_t bias_ih, cactus_node_t bias_hh, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_lstm_cell"); try { diff --git a/python/cactus/bindings/cactus.py b/python/cactus/bindings/cactus.py index fc9115713..60805c0a0 100644 --- a/python/cactus/bindings/cactus.py +++ b/python/cactus/bindings/cactus.py @@ -660,6 +660,11 @@ class cactus_tensor_info_t(ctypes.Structure): [cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_bool, cactus_node_t, ctypes.POINTER(cactus_node_t)], ctypes.c_int, ) +_bind_optional( + "cactus_graph_upsample_nearest2d", + [cactus_graph_t, cactus_node_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) _bind_optional( "cactus_graph_conv2d", [ @@ -2135,6 +2140,16 @@ def bind_mmap_weights(self, tensor, filename): if rc != 0: raise RuntimeError(_err("graph_bind_mmap_weights failed")) + def upsample_nearest2d(self, x, scale_factor, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_upsample_nearest2d( + self.h, cactus_node_t(x.id), ctypes.c_size_t(int(scale_factor)), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(_err("graph_upsample_nearest2d failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + def bilinear_interpolation(self, pos_embeds, dst_height, dst_width, backend=None): pos_embeds = self._ensure_tensor(pos_embeds) out = cactus_node_t() diff --git a/python/cactus/transpile/aten_ops.py b/python/cactus/transpile/aten_ops.py index 9096b8515..025f58ec4 100644 --- a/python/cactus/transpile/aten_ops.py +++ b/python/cactus/transpile/aten_ops.py @@ -152,6 +152,7 @@ def canonical_torch_op(target: Any) -> str: AtenOpRule("aten.embedding", "embedding"), AtenOpRule("aten.conv1d", "conv1d"), AtenOpRule("aten.conv2d", "conv2d"), + AtenOpRule("aten.upsample_nearest2d", "upsample_nearest2d"), AtenOpRule("aten.pow", "pow"), AtenOpRule("aten.layer_norm", "layer_norm"), AtenOpRule("aten.native_layer_norm", "layer_norm"), diff --git a/python/cactus/transpile/importers.py b/python/cactus/transpile/importers.py index 9666f7c03..f08ddd70d 100644 --- a/python/cactus/transpile/importers.py +++ b/python/cactus/transpile/importers.py @@ -1783,6 +1783,37 @@ def import_batch_norm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tupl register_node(ir, ir_node, shape=shape, dtype=dtype) +def import_upsample_nearest2d(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + scale_factors = extract_literals(node.args[2]) if len(node.args) > 2 else None + output_size = extract_literals(node.args[1]) if len(node.args) > 1 else None + + scale = None + if isinstance(scale_factors, (list, tuple)) and len(scale_factors) == 2: + if float(scale_factors[0]) == float(scale_factors[1]) and float(scale_factors[0]).is_integer(): + scale = int(scale_factors[0]) + elif isinstance(output_size, (list, tuple)) and len(output_size) == 2: + source = _extract_fx_tensor_shape(node.args[0]) + if source is not None and len(source) == 4: + height, width = int(source[2]), int(source[3]) + height_scale = int(output_size[0]) // height if height else 0 + width_scale = int(output_size[1]) // width if width else 0 + if height_scale == width_scale and height_scale * height == int(output_size[0]) and width_scale * width == int(output_size[1]): + scale = height_scale + + if scale is None or scale < 1: + ctx.fail(f"unsupported upsample_nearest2d for {torch_op}: output_size={output_size!r} scale_factors={scale_factors!r}") + + ir_node = IRNode( + id=node_id(node), + op="upsample_nearest2d", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"scale_factor": scale}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + def import_contiguous(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: ir_node = IRNode( id=node_id(node), @@ -1948,6 +1979,7 @@ def import_gemma4_moe_layer_gated(ir, node, ctx, *, shape, dtype, torch_op): "embedding": import_embedding, "conv1d": import_conv1d, "conv2d": import_conv2d, + "upsample_nearest2d": import_upsample_nearest2d, "layer_norm": import_layer_norm, "rms_norm": import_rms_norm, "group_norm": import_group_norm, diff --git a/python/cactus/transpile/lower.py b/python/cactus/transpile/lower.py index fab23d6ad..d09df9e6e 100644 --- a/python/cactus/transpile/lower.py +++ b/python/cactus/transpile/lower.py @@ -2086,6 +2086,10 @@ def _lower_ir_node(g: Graph, node: IRNode, env: dict[str, Any], ir: IRGraph) -> return [g.cat(outputs, axis=1)] return [g.conv1d(x, weight, bias=bias, stride=stride)] + if op == "upsample_nearest2d": + x = _tensor(env, node.inputs[0]) + return [g.upsample_nearest2d(x, int(node.attrs["scale_factor"]))] + if op == "conv2d": x = _tensor(env, node.inputs[0]) weight = _tensor(env, node.inputs[1]) diff --git a/python/cactus/transpile/ops.py b/python/cactus/transpile/ops.py index 5e3ee5436..dc1535438 100644 --- a/python/cactus/transpile/ops.py +++ b/python/cactus/transpile/ops.py @@ -102,6 +102,7 @@ class OpSchema: OpSchema("batch_norm", num_inputs=5, attrs=("axis", "eps"), backend_op="batch_norm"), OpSchema("conv1d", num_inputs=2, attrs=("stride", "padding", "dilation", "groups"), backend_op="conv1d"), OpSchema("conv2d", num_inputs=2, attrs=("stride", "padding", "dilation", "groups"), backend_op="conv2d"), + OpSchema("upsample_nearest2d", num_inputs=1, attrs=("scale_factor",), backend_op="upsample_nearest2d"), OpSchema("pad", num_inputs=1, attrs=("pads", "value", "mode")), OpSchema("lstm_cell", num_inputs=7, backend_op="lstm_cell"), diff --git a/python/tests/test_transpile_upsample.py b/python/tests/test_transpile_upsample.py new file mode 100644 index 000000000..5f724c569 --- /dev/null +++ b/python/tests/test_transpile_upsample.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import torch + +from cactus.transpile.capture_pytorch import capture_model + + +class NearestUpsample(torch.nn.Module): + def __init__(self, scale_factor: int) -> None: + super().__init__() + self.up = torch.nn.Upsample(scale_factor=scale_factor, mode="nearest") + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.up(x) + + +def _nodes(captured): + nodes = captured.ir_graph.nodes + return list(nodes.values()) if isinstance(nodes, dict) else list(nodes) + + +def _upsample_node(scale_factor: int): + captured = capture_model(NearestUpsample(scale_factor).eval(), (torch.randn(1, 3, 5, 7),)) + matches = [n for n in _nodes(captured) if n.op == "upsample_nearest2d"] + assert len(matches) == 1, f"expected one upsample node, got {[n.op for n in _nodes(captured)]}" + return matches[0] + + +def test_scale_factor_is_recovered_from_the_exported_graph(): + assert _upsample_node(2).attrs["scale_factor"] == 2 + assert _upsample_node(3).attrs["scale_factor"] == 3 + + +def test_output_size_form_is_recovered_from_the_input_shape(): + class ExplicitSize(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.interpolate(x, size=(10, 14), mode="nearest") + + captured = capture_model(ExplicitSize().eval(), (torch.randn(1, 3, 5, 7),)) + matches = [n for n in _nodes(captured) if n.op == "upsample_nearest2d"] + assert len(matches) == 1 + assert matches[0].attrs["scale_factor"] == 2 From 216f911462ee82014061eaeeec899f46d2af1bbb Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sat, 8 Aug 2026 04:24:37 -0700 Subject: [PATCH 11/34] drop the hand-built VAE test for kernel benchmarks 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 --- cactus-graph/tests/test_vae.cpp | 171 ----------------------------- cactus-kernels/tests/test_conv.cpp | 29 +++++ python/cactus/cli/common.py | 1 - python/cactus/cli/test.py | 20 +--- 4 files changed, 30 insertions(+), 191 deletions(-) delete mode 100644 cactus-graph/tests/test_vae.cpp diff --git a/cactus-graph/tests/test_vae.cpp b/cactus-graph/tests/test_vae.cpp deleted file mode 100644 index 2ad8031c8..000000000 --- a/cactus-graph/tests/test_vae.cpp +++ /dev/null @@ -1,171 +0,0 @@ -#include "test_utils.h" -#include "../../cactus-kernels/cactus_kernels.h" -#include -#include -#include -#include -#include - -using namespace TestUtils; - -namespace { - -constexpr size_t IMAGE_HW = 512, LATENT_HW = 64, LATENT_C = 4, IMAGE_C = 3, CH = 64; - -struct TaesdWeights { - CactusGraph& graph; - std::string dir; - - size_t weight(const std::string& name) { return graph.mmap_weights(dir + "/" + name + ".weights"); } - size_t bias(const std::string& name) { return graph.mmap_weights(dir + "/" + name + ".bias"); } -}; - -size_t conv_biased(TaesdWeights& w, size_t x, const std::string& name) { - return w.graph.conv2d_k3s1p1(x, w.weight(name), w.bias(name)); -} - -size_t block(TaesdWeights& w, size_t x, const std::string& prefix) { - size_t h = w.graph.relu(conv_biased(w, x, prefix + "_conv_0")); - h = w.graph.relu(conv_biased(w, h, prefix + "_conv_2")); - h = conv_biased(w, h, prefix + "_conv_4"); - return w.graph.relu(w.graph.add(h, x)); -} - -size_t build_encoder(TaesdWeights& w, size_t image) { - size_t x = conv_biased(w, image, "encoder_layers_0"); - x = block(w, x, "encoder_layers_1"); - for (int stage = 0; stage < 3; ++stage) { - const int down = 2 + stage * 4; - x = w.graph.conv2d_k3s2p1(x, w.weight("encoder_layers_" + std::to_string(down))); - for (int j = 1; j <= 3; ++j) x = block(w, x, "encoder_layers_" + std::to_string(down + j)); - } - return conv_biased(w, x, "encoder_layers_14"); -} - -size_t build_decoder(TaesdWeights& w, size_t latent) { - size_t x = w.graph.scalar_multiply(w.graph.tanh(w.graph.scalar_divide(latent, 3.0f)), 3.0f); - x = w.graph.relu(conv_biased(w, x, "decoder_layers_0")); - for (int stage = 0; stage < 3; ++stage) { - const int first = 2 + stage * 5; - for (int j = 0; j < 3; ++j) x = block(w, x, "decoder_layers_" + std::to_string(first + j)); - x = w.graph.upsample_nearest2d(x, 2); - x = w.graph.conv2d_k3s1p1(x, w.weight("decoder_layers_" + std::to_string(first + 4))); - } - x = block(w, x, "decoder_layers_17"); - return w.graph.clamp(conv_biased(w, x, "decoder_layers_18"), 0.0f, 1.0f); -} - -std::vector<__fp16> load_image_chw(const std::string& path, std::vector& rgb) { - int w = 0, h = 0, c = 0; - unsigned char* pixels = cactus_image_load(path.c_str(), &w, &h, &c, 3); - if (!pixels) return {}; - rgb.assign(IMAGE_HW * IMAGE_HW * IMAGE_C, 0); - if (w == (int)IMAGE_HW && h == (int)IMAGE_HW) { - std::copy(pixels, pixels + rgb.size(), rgb.begin()); - } else { - cactus_image_resize_uint8(pixels, w, h, rgb.data(), IMAGE_HW, IMAGE_HW, IMAGE_C); - } - cactus_image_free(pixels); - - std::vector<__fp16> chw(rgb.size()); - const size_t plane = IMAGE_HW * IMAGE_HW; - for (size_t i = 0; i < plane; ++i) - for (size_t ch = 0; ch < IMAGE_C; ++ch) - chw[ch * plane + i] = static_cast<__fp16>(rgb[i * IMAGE_C + ch] / 255.0f); - return chw; -} - -void write_ppm(const std::string& path, const __fp16* chw) { - std::ofstream f(path, std::ios::binary); - f << "P6\n" << IMAGE_HW << " " << IMAGE_HW << "\n255\n"; - const size_t plane = IMAGE_HW * IMAGE_HW; - for (size_t i = 0; i < plane; ++i) - for (size_t c = 0; c < IMAGE_C; ++c) { - float v = static_cast(chw[c * plane + i]) * 255.0f; - f.put(static_cast(v < 0.0f ? 0 : (v > 255.0f ? 255 : static_cast(v + 0.5f)))); - } -} - -float psnr_vs_source(const __fp16* out, const std::vector& rgb) { - const size_t plane = IMAGE_HW * IMAGE_HW; - double se = 0.0; - for (size_t i = 0; i < plane; ++i) - for (size_t c = 0; c < IMAGE_C; ++c) { - double d = static_cast(out[c * plane + i]) - rgb[i * IMAGE_C + c] / 255.0; - se += d * d; - } - double mse = se / (plane * IMAGE_C); - return mse > 0.0 ? static_cast(10.0 * std::log10(1.0 / mse)) : 99.0f; -} - -bool roundtrip(const std::string& dir, const std::string& image_path, const char* backend) { - if (cactus_backend_select(backend) != 0) { - std::cout << " [" << backend << "] unavailable, skipping" << std::endl; - return true; - } - - std::vector rgb; - auto image = load_image_chw(image_path, rgb); - if (image.empty()) { - std::cout << " could not read " << image_path << ": " << cactus_image_failure_reason() << std::endl; - return false; - } - - std::vector<__fp16> latent(LATENT_C * LATENT_HW * LATENT_HW); - { - CactusGraph g; - TaesdWeights w{g, dir}; - size_t input = g.input({1, IMAGE_C, IMAGE_HW, IMAGE_HW}, Precision::FP16); - size_t out = build_encoder(w, input); - g.set_input(input, image.data(), Precision::FP16); - g.execute(); - const __fp16* data = static_cast(g.get_output(out)); - std::copy(data, data + latent.size(), latent.begin()); - g.hard_reset(); - } - - CactusGraph g; - TaesdWeights w{g, dir}; - size_t input = g.input({1, LATENT_C, LATENT_HW, LATENT_HW}, Precision::FP16); - size_t out = build_decoder(w, input); - g.set_input(input, latent.data(), Precision::FP16); - - g.execute(); - constexpr int runs = 5; - Timer timer; - for (int i = 0; i < runs; ++i) g.execute(); - double ms = timer.elapsed_ms() / runs; - - const __fp16* decoded = static_cast(g.get_output(out)); - float quality = psnr_vs_source(decoded, rgb); - write_ppm(dir + "/taesd_roundtrip_" + backend + ".ppm", decoded); - std::cout << " [" << backend << "] decode " << ms << " ms, roundtrip PSNR " << quality << " dB" << std::endl; - - g.hard_reset(); - return quality > 20.0f; -} - -bool test_taesd_roundtrip() { - const char* bundle = std::getenv("CACTUS_TEST_VAE_MODEL"); - if (!bundle || !*bundle) { - std::cout << " skipped: no VAE weights. Run graph tests via 'cactus test', which prepares them." - << std::endl; - return true; - } - const char* assets = std::getenv("CACTUS_TEST_ASSETS"); - std::string image_path = assets ? std::string(assets) + "/test_monkey.png" - : "../../cactus-engine/tests/assets/test_monkey.png"; - - bool ok = roundtrip(bundle, image_path, "cpu"); - ok = roundtrip(bundle, image_path, "metal") && ok; - return ok; -} - -} - -int main() { - TestRunner runner("VAE Tests"); - runner.run_test("taesd_encode_decode_roundtrip", test_taesd_roundtrip()); - runner.print_summary(); - return runner.all_passed() ? 0 : 1; -} diff --git a/cactus-kernels/tests/test_conv.cpp b/cactus-kernels/tests/test_conv.cpp index 2086b095f..571c5e6b5 100644 --- a/cactus-kernels/tests/test_conv.cpp +++ b/cactus-kernels/tests/test_conv.cpp @@ -148,6 +148,35 @@ bool test_upsample_nearest2d() { } bool run_benchmarks() { + { + const size_t N = 1, C_in = 64, C_out = 64, H = 512, W = 512; + std::vector<__fp16> input(N * C_in * H * W), weight(C_out * C_in * 9), output(N * C_out * H * W); + fill_random_fp16(input, -0.5f, 0.5f); + fill_random_fp16(weight, -0.5f, 0.5f); + cactus_conv2d_f16_k3s1p1_nchw(input.data(), weight.data(), nullptr, output.data(), N, C_in, H, W, C_out); + Timer t; + for (int i = 0; i < 10; i++) + cactus_conv2d_f16_k3s1p1_nchw(input.data(), weight.data(), nullptr, output.data(), N, C_in, H, W, C_out); + double ms = t.elapsed_ms() / 10.0; + double gflops = (2.0 * N * C_out * H * W * 9 * C_in) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << "conv2d_k3s1p1 64x512x512" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gflops << " GFLOPS\n"; + } + { + const size_t planes = 64, H = 256, W = 256, scale = 2; + std::vector<__fp16> input(planes * H * W), output(planes * H * scale * W * scale); + fill_random_fp16(input, -0.5f, 0.5f); + cactus_upsample_nearest2d_f16(input.data(), output.data(), planes, H, W, scale); + Timer t; + for (int i = 0; i < 50; i++) + cactus_upsample_nearest2d_f16(input.data(), output.data(), planes, H, W, scale); + double ms = t.elapsed_ms() / 50.0; + double gbs = (output.size() * sizeof(__fp16)) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << "upsample_nearest2d 64x256^2" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gbs << " GB/s\n"; + } { const size_t N = 1, L = 3000, C_in = 80, C_out = 512, stride = 1; const size_t L_out = ((L - 1) / stride) + 1; diff --git a/python/cactus/cli/common.py b/python/cactus/cli/common.py index 214c79f3d..f39ef435e 100644 --- a/python/cactus/cli/common.py +++ b/python/cactus/cli/common.py @@ -31,7 +31,6 @@ def weights_root() -> Path: DEFAULT_TRANSCRIPTION_MODEL_ID = "nvidia/parakeet-tdt-0.6b-v3" DEFAULT_TEST_MODEL_ID = DEFAULT_MODEL_ID DEFAULT_TEST_TRANSCRIPTION_MODEL_ID = DEFAULT_TRANSCRIPTION_MODEL_ID -DEFAULT_TEST_VAE_MODEL_ID = "madebyollin/taesd" RED = '\033[0;31m' diff --git a/python/cactus/cli/test.py b/python/cactus/cli/test.py index 11603e64d..720f9268b 100644 --- a/python/cactus/cli/test.py +++ b/python/cactus/cli/test.py @@ -3,8 +3,7 @@ from .common import ( BLUE, DEFAULT_TEST_MODEL_ID, DEFAULT_TEST_TRANSCRIPTION_MODEL_ID, - DEFAULT_TEST_VAE_MODEL_ID, PROJECT_ROOT, RED, YELLOW, - apply_cloud_api_key_env, print_color, + PROJECT_ROOT, RED, YELLOW, apply_cloud_api_key_env, print_color, ) COMPONENTS = ("kernels", "graph", "engine", "all") @@ -75,21 +74,6 @@ def cmd_test(args): else: targets = ("kernels", "graph", "engine") if args.component == "all" else (args.component,) - if "graph" in targets: - from .model import ensure_weights - - try: - args.vae_model_id = str(ensure_weights( - getattr(args, "vae_model_id", None) or DEFAULT_TEST_VAE_MODEL_ID, - bits=getattr(args, "bits", 4), - token=getattr(args, "token", None), - reconvert=getattr(args, "reconvert", False), - skip_model_load=True, - )) - except (RuntimeError, OSError, ValueError) as exc: - print_color(YELLOW, f"VAE weights unavailable ({exc}); the vae suite will be skipped") - args.vae_model_id = None - if "engine" in targets: from .model import prepare_bundle model_dir = prepare_bundle(args, model_id=args.model_id or DEFAULT_TEST_MODEL_ID, @@ -109,8 +93,6 @@ def cmd_test(args): env.pop("CACTUS_NO_CLOUD_TELE", None) else: env["CACTUS_NO_CLOUD_TELE"] = "1" - if getattr(args, "vae_model_id", None): - env["CACTUS_TEST_VAE_MODEL"] = args.vae_model_id for c in targets: cmd = _component_args(c, args) From e1c830328ef17968966095b8a399f772804075e3 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sat, 8 Aug 2026 04:47:40 -0700 Subject: [PATCH 12/34] persist the upsample scale factor across save/load 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 --- cactus-graph/src/io.cpp | 4 +-- cactus-graph/src/param_io.cpp | 1 + cactus-graph/tests/test_io.cpp | 53 +++++++++++++++++++++++++++ cactus-kernels/tests/test_conv.cpp | 58 +++++++++++++++--------------- python/cactus/transpile/lower.py | 8 ++--- 5 files changed, 89 insertions(+), 35 deletions(-) diff --git a/cactus-graph/src/io.cpp b/cactus-graph/src/io.cpp index 2d90133a4..94f4605e0 100644 --- a/cactus-graph/src/io.cpp +++ b/cactus-graph/src/io.cpp @@ -278,10 +278,10 @@ namespace { GraphFile::NodeEntry node; node.index = read_u32(in); uint32_t op_type_val = read_u32(in); - if (op_type_val > static_cast(OpType::LOGITS_TQ_SOFTCAP)) { + if (op_type_val > static_cast(OpType::UPSAMPLE_NEAREST2D)) { throw std::runtime_error( "Graph file corrupted: invalid op type " + std::to_string(op_type_val) + - " > " + std::to_string(static_cast(OpType::LOGITS_TQ_SOFTCAP))); + " > " + std::to_string(static_cast(OpType::UPSAMPLE_NEAREST2D))); } node.op_type = static_cast(op_type_val); node.inputs = read_u32_vector(in); diff --git a/cactus-graph/src/param_io.cpp b/cactus-graph/src/param_io.cpp index 2c58b6020..d72a2444a 100644 --- a/cactus-graph/src/param_io.cpp +++ b/cactus-graph/src/param_io.cpp @@ -251,6 +251,7 @@ const Schema& op_schema(OpType op_type) { {OpType::ALTUP_PREDICT, {{ParamField::NumAltupInputs, FieldPersistence::Persistent}}}, {OpType::ALTUP_CORRECT, {{ParamField::NumAltupInputs, FieldPersistence::Persistent}}}, {OpType::MAXPOOL1D, {{ParamField::KernelSize, FieldPersistence::Persistent}, {ParamField::Stride, FieldPersistence::Persistent}}}, + {OpType::UPSAMPLE_NEAREST2D, {{ParamField::Stride, FieldPersistence::Persistent}}}, {OpType::CONV1D_CAUSAL, {{ParamField::Dilation, FieldPersistence::Persistent}}}, {OpType::CONV1D_CAUSAL_CHANNEL_FIRST, {{ParamField::Dilation, FieldPersistence::Persistent}}}, {OpType::LOGITS_TQ_SOFTCAP, {{ParamField::Scalar, FieldPersistence::Persistent}, {ParamField::Scale, FieldPersistence::Persistent}}}, diff --git a/cactus-graph/tests/test_io.cpp b/cactus-graph/tests/test_io.cpp index d792b6102..9cacea1c4 100644 --- a/cactus-graph/tests/test_io.cpp +++ b/cactus-graph/tests/test_io.cpp @@ -202,6 +202,58 @@ bool test_graph_save_load_roundtrip_execution() { } } +bool test_save_load_preserves_upsample_scale_factor() { + try { + const std::string filename = "test_upsample_scale_roundtrip.cg"; + constexpr size_t scale = 3; + + CactusGraph original; + size_t input_id = original.input({1, 1, 2, 2}, Precision::FP16); + size_t upsampled = original.upsample_nearest2d(input_id, scale); + std::vector<__fp16> data = {1, 2, 3, 4}; + original.set_input(input_id, data.data(), Precision::FP16); + original.execute(); + original.save(filename); + + CactusGraph loaded = CactusGraph::load(filename); + GraphFile::SerializedGraph loaded_graph = GraphFile::load_graph(filename); + const size_t loaded_input = runtime_id_from_serialized_index(loaded_graph.graph_inputs[0]); + const size_t loaded_output = runtime_id_from_serialized_index(loaded_graph.graph_outputs[0]); + + loaded.set_input(loaded_input, data.data(), Precision::FP16); + loaded.execute(); + + const auto& shape = loaded.get_output_buffer(loaded_output).shape; + if (shape.size() != 4 || shape[2] != 2 * scale || shape[3] != 2 * scale) { + std::cout << "[upsample_scale_roundtrip] scale factor lost: output is ["; + for (size_t dim : shape) std::cout << dim << ","; + std::cout << "] not [1,1,6,6]" << std::endl; + std::remove(filename.c_str()); + return false; + } + + const __fp16* out = static_cast(loaded.get_output(loaded_output)); + for (size_t y = 0; y < 2 * scale; ++y) { + for (size_t x = 0; x < 2 * scale; ++x) { + float expected = static_cast(data[(y / scale) * 2 + x / scale]); + float got = static_cast(out[y * 2 * scale + x]); + if (got != expected) { + std::cout << "[upsample_scale_roundtrip] mismatch at (" << y << "," << x + << "): got=" << got << " expected=" << expected << std::endl; + std::remove(filename.c_str()); + return false; + } + } + } + + std::remove(filename.c_str()); + return true; + } catch (const std::exception& e) { + std::cout << "[upsample_scale_roundtrip] exception: " << e.what() << std::endl; + return false; + } +} + bool test_graph_save_load_supported_ops_roundtrip() { try { const std::string filename = "test_graph_supported_roundtrip.cg"; @@ -512,6 +564,7 @@ int main() { runner.run_test("Graph Save/Load", test_graph_save_load()); runner.run_test("Graph Save/Load Roundtrip Execution", test_graph_save_load_roundtrip_execution()); runner.run_test("Graph Save/Load Supported Ops Roundtrip", test_graph_save_load_supported_ops_roundtrip()); + runner.run_test("Save/Load Preserves Upsample Scale", test_save_load_preserves_upsample_scale_factor()); runner.run_test("Graph Save For Inspection", test_graph_save_for_inspection()); runner.run_test("Save/Load Preserves Recurrent Cache Persistence", test_save_load_preserves_recurrent_cache_persistence()); diff --git a/cactus-kernels/tests/test_conv.cpp b/cactus-kernels/tests/test_conv.cpp index 571c5e6b5..4435ea4cc 100644 --- a/cactus-kernels/tests/test_conv.cpp +++ b/cactus-kernels/tests/test_conv.cpp @@ -148,35 +148,6 @@ bool test_upsample_nearest2d() { } bool run_benchmarks() { - { - const size_t N = 1, C_in = 64, C_out = 64, H = 512, W = 512; - std::vector<__fp16> input(N * C_in * H * W), weight(C_out * C_in * 9), output(N * C_out * H * W); - fill_random_fp16(input, -0.5f, 0.5f); - fill_random_fp16(weight, -0.5f, 0.5f); - cactus_conv2d_f16_k3s1p1_nchw(input.data(), weight.data(), nullptr, output.data(), N, C_in, H, W, C_out); - Timer t; - for (int i = 0; i < 10; i++) - cactus_conv2d_f16_k3s1p1_nchw(input.data(), weight.data(), nullptr, output.data(), N, C_in, H, W, C_out); - double ms = t.elapsed_ms() / 10.0; - double gflops = (2.0 * N * C_out * H * W * 9 * C_in) / (ms * 1e6); - std::cout << " ⚡ " << std::left << std::setw(28) << "conv2d_k3s1p1 64x512x512" - << std::fixed << std::setprecision(3) << ms << "ms " - << std::setprecision(1) << gflops << " GFLOPS\n"; - } - { - const size_t planes = 64, H = 256, W = 256, scale = 2; - std::vector<__fp16> input(planes * H * W), output(planes * H * scale * W * scale); - fill_random_fp16(input, -0.5f, 0.5f); - cactus_upsample_nearest2d_f16(input.data(), output.data(), planes, H, W, scale); - Timer t; - for (int i = 0; i < 50; i++) - cactus_upsample_nearest2d_f16(input.data(), output.data(), planes, H, W, scale); - double ms = t.elapsed_ms() / 50.0; - double gbs = (output.size() * sizeof(__fp16)) / (ms * 1e6); - std::cout << " ⚡ " << std::left << std::setw(28) << "upsample_nearest2d 64x256^2" - << std::fixed << std::setprecision(3) << ms << "ms " - << std::setprecision(1) << gbs << " GB/s\n"; - } { const size_t N = 1, L = 3000, C_in = 80, C_out = 512, stride = 1; const size_t L_out = ((L - 1) / stride) + 1; @@ -226,6 +197,35 @@ bool run_benchmarks() { << std::fixed << std::setprecision(3) << ms << "ms " << std::setprecision(1) << ms << " ms\n"; } + { + const size_t N = 1, C_in = 64, C_out = 64, H = 512, W = 512; + std::vector<__fp16> input(N * C_in * H * W), weight(C_out * C_in * 9), output(N * C_out * H * W); + fill_random_fp16(input, -0.5f, 0.5f); + fill_random_fp16(weight, -0.5f, 0.5f); + cactus_conv2d_f16_k3s1p1_nchw(input.data(), weight.data(), nullptr, output.data(), N, C_in, H, W, C_out); + Timer t; + for (int i = 0; i < 10; i++) + cactus_conv2d_f16_k3s1p1_nchw(input.data(), weight.data(), nullptr, output.data(), N, C_in, H, W, C_out); + double ms = t.elapsed_ms() / 10.0; + double gflops = (2.0 * N * C_out * H * W * 9 * C_in) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << "conv2d_k3s1p1 64x512x512" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gflops << " GFLOPS\n"; + } + { + const size_t planes = 64, H = 256, W = 256, scale = 2; + std::vector<__fp16> input(planes * H * W), output(planes * H * scale * W * scale); + fill_random_fp16(input, -0.5f, 0.5f); + cactus_upsample_nearest2d_f16(input.data(), output.data(), planes, H, W, scale); + Timer t; + for (int i = 0; i < 50; i++) + cactus_upsample_nearest2d_f16(input.data(), output.data(), planes, H, W, scale); + double ms = t.elapsed_ms() / 50.0; + double gbs = (output.size() * sizeof(__fp16)) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << "upsample_nearest2d 64x256^2" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gbs << " GB/s\n"; + } return true; } diff --git a/python/cactus/transpile/lower.py b/python/cactus/transpile/lower.py index d09df9e6e..909f43da1 100644 --- a/python/cactus/transpile/lower.py +++ b/python/cactus/transpile/lower.py @@ -2086,10 +2086,6 @@ def _lower_ir_node(g: Graph, node: IRNode, env: dict[str, Any], ir: IRGraph) -> return [g.cat(outputs, axis=1)] return [g.conv1d(x, weight, bias=bias, stride=stride)] - if op == "upsample_nearest2d": - x = _tensor(env, node.inputs[0]) - return [g.upsample_nearest2d(x, int(node.attrs["scale_factor"]))] - if op == "conv2d": x = _tensor(env, node.inputs[0]) weight = _tensor(env, node.inputs[1]) @@ -2157,6 +2153,10 @@ def _lower_ir_node(g: Graph, node: IRNode, env: dict[str, Any], ir: IRGraph) -> return [g.conv2d(x, weight, bias=bias, stride=stride, padding=padding, dilation=dilation, groups=groups)] + if op == "upsample_nearest2d": + x = _tensor(env, node.inputs[0]) + return [g.upsample_nearest2d(x, int(node.attrs["scale_factor"]))] + if op == "layer_norm": x = _tensor(env, node.inputs[0]) weight = _tensor(env, node.inputs[1]) From a94815b6adcf4d80b5d2634f0b95b367291d2a61 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sat, 8 Aug 2026 05:47:25 -0700 Subject: [PATCH 13/34] benchmark conv2d and upsample at graph level 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 --- cactus-graph/tests/test_nn.cpp | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/cactus-graph/tests/test_nn.cpp b/cactus-graph/tests/test_nn.cpp index 67519ca28..0a37180e6 100644 --- a/cactus-graph/tests/test_nn.cpp +++ b/cactus-graph/tests/test_nn.cpp @@ -705,6 +705,29 @@ bool run_benchmarks() { g.set_input(ii, in.data(), Precision::FP16); bench("softmax 1024x1024", []{}, [&]{ g.execute(); }); } + { + const size_t channels = 64, height = 512, width = 512; + std::vector<__fp16> in(channels * height * width), w(channels * channels * 9); + TestUtils::fill_random_fp16(in); + TestUtils::fill_random_fp16(w); + CactusGraph g; + size_t ii = g.input({1, channels, height, width}, Precision::FP16); + size_t iw = g.input({channels, channels, 3, 3}, Precision::FP16); + g.conv2d_k3s1p1(ii, iw); + g.set_input(ii, in.data(), Precision::FP16); + g.set_input(iw, w.data(), Precision::FP16); + bench("conv2d_k3s1p1 64x512x512", []{}, [&]{ g.execute(); }); + } + { + const size_t channels = 64, height = 256, width = 256; + std::vector<__fp16> in(channels * height * width); + TestUtils::fill_random_fp16(in); + CactusGraph g; + size_t ii = g.input({1, channels, height, width}, Precision::FP16); + g.upsample_nearest2d(ii, 2); + g.set_input(ii, in.data(), Precision::FP16); + bench("upsample_nearest2d 64x256^2", []{}, [&]{ g.execute(); }); + } return true; } From 44f210f30f060236dd0bdb9d169bf81b3b0c4d1a Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sat, 8 Aug 2026 05:50:56 -0700 Subject: [PATCH 14/34] test the transpiled decoder end to end MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- python/tests/test_transpile_taesd.py | 79 ++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) create mode 100644 python/tests/test_transpile_taesd.py diff --git a/python/tests/test_transpile_taesd.py b/python/tests/test_transpile_taesd.py new file mode 100644 index 000000000..d09e9736a --- /dev/null +++ b/python/tests/test_transpile_taesd.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import time + +import numpy as np +import pytest +import torch + +diffusers = pytest.importorskip("diffusers") + +try: + from cactus.transpile.capture_pytorch import capture_model + from cactus.transpile.lower import transpile_ir +except Exception as exc: # pragma: no cover + pytest.skip(f"cactus runtime library unavailable: {exc}", allow_module_level=True) + +MODEL_ID = "madebyollin/taesd" +LATENT_SHAPE = (1, 4, 64, 64) + + +def _decoder(): + from diffusers import AutoencoderTiny + + try: + vae = AutoencoderTiny.from_pretrained(MODEL_ID, torch_dtype=torch.float32, local_files_only=True) + except Exception as exc: + pytest.skip(f"{MODEL_ID} not cached: {exc}") + return vae.eval().decoder + + +def _ir_nodes(captured): + nodes = captured.ir_graph.nodes + return list(nodes.values()) if isinstance(nodes, dict) else list(nodes) + + +def test_taesd_decoder_transpiles_to_the_expected_ops(): + torch.manual_seed(0) + captured = capture_model(_decoder(), (torch.randn(*LATENT_SHAPE),)) + counts: dict[str, int] = {} + for node in _ir_nodes(captured): + counts[node.op] = counts.get(node.op, 0) + 1 + + assert counts["conv2d"] == 35 + assert counts["upsample_nearest2d"] == 3 + assert counts["relu"] == 31 + assert set(counts) == { + "conv2d", "relu", "add", "upsample_nearest2d", + "scalar_multiply", "scalar_divide", "tanh", "scalar_subtract", + } + + +def test_taesd_decode_matches_torch_and_reports_latency(): + torch.manual_seed(0) + decoder = _decoder() + latent = torch.randn(*LATENT_SHAPE) + + transpiled = transpile_ir(capture_model(decoder, (latent,)).ir_graph) + transpiled.set_input(0, latent.numpy().astype(np.float16)) + decoded = np.asarray(transpiled.execute()[0].numpy(), dtype=np.float32) + + runs = 5 + start = time.perf_counter() + for _ in range(runs): + transpiled.execute() + elapsed_ms = (time.perf_counter() - start) * 1000.0 / runs + + with torch.no_grad(): + reference = decoder(latent).float().numpy() + + assert decoded.shape == reference.shape == (1, 3, 512, 512) + diff = np.abs(reference - decoded) + relative = diff.mean() / np.abs(reference).mean() + print( + f"\n taesd decode 512x512: {elapsed_ms:.1f} ms" + f" max_abs_diff={diff.max():.4f} relative_mean_err={relative * 100:.3f}%" + ) + + assert relative < 0.01, f"transpiled decode drifted from torch: {relative * 100:.3f}%" + assert diff.max() < 0.2, f"transpiled decode has an outlier: {diff.max():.4f}" From a44b4cbefb346dcae1f44c6663d04fef7e4ca7e4 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sat, 8 Aug 2026 14:34:56 -0700 Subject: [PATCH 15/34] inline the checkpoint stem loop 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 --- python/cactus/convert/cli.py | 64 ++++++++++++++++-------------------- 1 file changed, 29 insertions(+), 35 deletions(-) diff --git a/python/cactus/convert/cli.py b/python/cactus/convert/cli.py index c52b7f29b..151522f00 100644 --- a/python/cactus/convert/cli.py +++ b/python/cactus/convert/cli.py @@ -141,48 +141,42 @@ def _bits_for_component(component: str, args: argparse.Namespace) -> int: CHECKPOINT_STEMS = ("model", "diffusion_pytorch_model") -def _download_checkpoint_state_dict(model_id_or_path: str, stem: str) -> dict[str, Any] | None: - from huggingface_hub import hf_hub_download - from safetensors.torch import load_file - - try: - return load_file(hf_hub_download(model_id_or_path, f"{stem}.safetensors", cache_dir=_hf_cache_dir())) - except Exception: - pass - try: - index_file = hf_hub_download(model_id_or_path, f"{stem}.safetensors.index.json", cache_dir=_hf_cache_dir()) - except Exception: - return None - index = json.loads(Path(index_file).read_text(encoding="utf-8")) - shard_names = sorted(set(index.get("weight_map", {}).values())) - if not shard_names: - return None - state: dict[str, Any] = {} - for shard in shard_names: - shard_path = hf_hub_download(model_id_or_path, shard, cache_dir=_hf_cache_dir()) - for key, tensor in load_file(shard_path).items(): - if key in state: - raise RuntimeError(f"duplicate tensor key {key!r} across checkpoint shards") - state[key] = tensor - return state - - def _load_checkpoint_state_dict(model_id_or_path: str) -> dict[str, Any] | None: nemo_export = ensure_parakeet_tdt_nemo_source(model_id_or_path, cache_dir=_hf_cache_dir()) if nemo_export is not None: model_id_or_path = nemo_export root = Path(model_id_or_path) if not root.exists() or not root.is_dir(): - for stem in CHECKPOINT_STEMS: - try: - state = _download_checkpoint_state_dict(model_id_or_path, stem) - except RuntimeError: - raise - except Exception: - continue - if state is not None: + try: + from huggingface_hub import hf_hub_download + from safetensors.torch import load_file + for stem in CHECKPOINT_STEMS: + try: + model_file = hf_hub_download(model_id_or_path, f"{stem}.safetensors", cache_dir=_hf_cache_dir()) + return load_file(model_file) + except Exception: + pass + try: + index_file = hf_hub_download(model_id_or_path, f"{stem}.safetensors.index.json", cache_dir=_hf_cache_dir()) + except Exception: + continue + index = json.loads(Path(index_file).read_text(encoding="utf-8")) + shard_names = sorted(set(index.get("weight_map", {}).values())) + if not shard_names: + return None + state: dict[str, Any] = {} + for shard in shard_names: + shard_path = hf_hub_download(model_id_or_path, shard, cache_dir=_hf_cache_dir()) + for key, tensor in load_file(shard_path).items(): + if key in state: + raise RuntimeError(f"duplicate tensor key {key!r} across checkpoint shards") + state[key] = tensor return state - return None + return None + except RuntimeError: + raise + except Exception: + return None safetensor_files = sorted(root.glob("*.safetensors")) if safetensor_files: try: From 4dd57e0eee8a12e1a70c5322947a5d301feb3adb Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sat, 8 Aug 2026 15:01:20 -0700 Subject: [PATCH 16/34] drop the unused importorskip binding The module gets imported where it is used, so the guard does not need to bind it. Signed-off-by: Neel Shejwalkar --- python/tests/test_transpile_taesd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tests/test_transpile_taesd.py b/python/tests/test_transpile_taesd.py index d09e9736a..0b1cf3cff 100644 --- a/python/tests/test_transpile_taesd.py +++ b/python/tests/test_transpile_taesd.py @@ -6,7 +6,7 @@ import pytest import torch -diffusers = pytest.importorskip("diffusers") +pytest.importorskip("diffusers") try: from cactus.transpile.capture_pytorch import capture_model From 417ecb4f4292c237259a48c9adc4ce62ed61ed62 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sun, 9 Aug 2026 14:38:50 -0700 Subject: [PATCH 17/34] restore the py.typed marker the packaging test expects Signed-off-by: Neel Shejwalkar --- .../convert/tests/test_taesd_adapter.py | 32 ------------------- python/cactus/py.typed | 0 2 files changed, 32 deletions(-) delete mode 100644 python/cactus/convert/tests/test_taesd_adapter.py create mode 100644 python/cactus/py.typed diff --git a/python/cactus/convert/tests/test_taesd_adapter.py b/python/cactus/convert/tests/test_taesd_adapter.py deleted file mode 100644 index 4a46bb56f..000000000 --- a/python/cactus/convert/tests/test_taesd_adapter.py +++ /dev/null @@ -1,32 +0,0 @@ -from __future__ import annotations - -import torch - -from cactus.convert.model_adapters.adapters import adapter_for_family -from cactus.convert.model_adapters.detection import detect_family - - -def test_diffusers_autoencoder_tiny_config_selects_taesd(): - assert detect_family({"_class_name": "AutoencoderTiny", "latent_channels": 4}) == "taesd" - - -def test_taesd_converts_from_checkpoint_without_a_transformers_model_class(): - adapter = adapter_for_family("taesd") - assert adapter.model_class({"_class_name": "AutoencoderTiny"}) is None - assert adapter.load_processor("madebyollin/taesd") is None - - -def test_taesd_keeps_encoder_and_decoder_names_distinct(): - adapter = adapter_for_family("taesd") - tensor = torch.zeros(64, 64, 3, 3) - encoder = adapter.name_tensor("encoder.layers.12.conv.0.weight", tensor, None) - decoder = adapter.name_tensor("decoder.layers.12.conv.0.weight", tensor, None) - assert encoder.output_name == "encoder_layers_12_conv_0.weights" - assert decoder.output_name == "decoder_layers_12_conv_0.weights" - - -def test_taesd_conv_tensors_stay_fp16(): - adapter = adapter_for_family("taesd") - match = adapter.name_tensor("decoder.layers.0.weight", torch.zeros(64, 4, 3, 3), None) - policy = adapter.policy(match, (64, 4, 3, 3), 4) - assert policy.precision == "FP16" diff --git a/python/cactus/py.typed b/python/cactus/py.typed new file mode 100644 index 000000000..e69de29bb From 69518454a48bde9525feb92337e73066e5c0f42b Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sun, 9 Aug 2026 14:38:50 -0700 Subject: [PATCH 18/34] tolerate group_norm exports that omit eps 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 --- python/cactus/transpile/importers.py | 3 +- python/tests/test_transpile_sd_unet.py | 99 ++++++++++++++++++++++++++ 2 files changed, 101 insertions(+), 1 deletion(-) create mode 100644 python/tests/test_transpile_sd_unet.py diff --git a/python/cactus/transpile/importers.py b/python/cactus/transpile/importers.py index f08ddd70d..120b9e822 100644 --- a/python/cactus/transpile/importers.py +++ b/python/cactus/transpile/importers.py @@ -1751,7 +1751,8 @@ def import_rms_norm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[ def import_group_norm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: inputs = [value_id(node.args[0], ctx), value_id(node.args[2], ctx), value_id(node.args[3], ctx)] - attrs = {"num_groups": int(node.args[1]), "eps": float(node.args[4])} + eps = node.args[4] if len(node.args) > 4 else node.kwargs.get("eps", 1e-5) + attrs = {"num_groups": int(node.args[1]), "eps": float(eps)} ir_node = IRNode( id=node_id(node), op="group_norm", diff --git a/python/tests/test_transpile_sd_unet.py b/python/tests/test_transpile_sd_unet.py new file mode 100644 index 000000000..211665775 --- /dev/null +++ b/python/tests/test_transpile_sd_unet.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import time + +import numpy as np +import pytest +import torch + +pytest.importorskip("diffusers") + +try: + from cactus.transpile.capture_pytorch import capture_model + from cactus.transpile.lower import transpile_ir +except Exception as exc: # pragma: no cover + pytest.skip(f"cactus runtime library unavailable: {exc}", allow_module_level=True) + +MODEL_ID = "SimianLuo/LCM_Dreamshaper_v7" +LATENT_SHAPE = (1, 4, 64, 64) +TEXT_EMB_SHAPE = (1, 77, 768) +GUIDANCE_EMB_SHAPE = (1, 256) + + +class _Denoiser(torch.nn.Module): + def __init__(self, unet): + super().__init__() + self.unet = unet + + def forward(self, sample, timestep, encoder_hidden_states, timestep_cond): + return self.unet( + sample, + timestep, + encoder_hidden_states=encoder_hidden_states, + timestep_cond=timestep_cond, + ).sample + + +def _denoiser(dtype): + from diffusers import UNet2DConditionModel + + try: + unet = UNet2DConditionModel.from_pretrained( + MODEL_ID, subfolder="unet", torch_dtype=dtype, local_files_only=True + ) + except Exception as exc: + pytest.skip(f"{MODEL_ID} not cached: {exc}") + return _Denoiser(unet.eval()) + + +def _example_inputs(): + torch.manual_seed(0) + return ( + torch.randn(*LATENT_SHAPE), + torch.tensor([999.0]), + torch.randn(*TEXT_EMB_SHAPE), + torch.randn(*GUIDANCE_EMB_SHAPE), + ) + + +def test_sd_unet_denoise_matches_torch_and_reports_latency(): + inputs = _example_inputs() + half = tuple(a.to(torch.float16) for a in inputs) + + # One capture only -- exporting the 860M UNet takes minutes, so the op + # histogram is asserted on the same capture the numerics run on. + captured = capture_model(_denoiser(torch.float16), half) + nodes = captured.ir_graph.nodes + nodes = list(nodes.values()) if isinstance(nodes, dict) else list(nodes) + counts: dict[str, int] = {} + for node in nodes: + counts[node.op] = counts.get(node.op, 0) + 1 + assert counts["conv2d"] == 98 + assert counts["attention"] == 32 # 16 self + 16 cross + assert counts["group_norm"] == 61 + assert counts["upsample_nearest2d"] == 3 + assert counts["silu"] == 68 + + transpiled = transpile_ir(captured.ir_graph) + for i, tensor in enumerate(half): + transpiled.set_input(i, tensor.numpy()) + eps = np.asarray(transpiled.execute()[0].numpy(), dtype=np.float32) + + runs = 3 + start = time.perf_counter() + for _ in range(runs): + transpiled.execute() + elapsed_ms = (time.perf_counter() - start) * 1000.0 / runs + + with torch.no_grad(): + reference = _denoiser(torch.float32)(*[a.float() for a in half]).float().numpy() + + assert eps.shape == reference.shape == LATENT_SHAPE + diff = np.abs(reference - eps) + relative = diff.mean() / np.abs(reference).mean() + print( + f"\n sd unet denoise 64x64: {elapsed_ms:.0f} ms" + f" max_abs_diff={diff.max():.4f} relative_mean_err={relative * 100:.3f}%" + ) + + assert relative < 0.01, f"transpiled denoise drifted from torch: {relative * 100:.3f}%" From f78c6d66e654cc91b3155e95b0d8fc4c365f68ec Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sun, 9 Aug 2026 14:38:50 -0700 Subject: [PATCH 19/34] test the transpiled CLIP text encoder end to end Signed-off-by: Neel Shejwalkar --- python/tests/test_transpile_clip.py | 88 +++++++++++++++++++++++++++++ 1 file changed, 88 insertions(+) create mode 100644 python/tests/test_transpile_clip.py diff --git a/python/tests/test_transpile_clip.py b/python/tests/test_transpile_clip.py new file mode 100644 index 000000000..261ebdf37 --- /dev/null +++ b/python/tests/test_transpile_clip.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +import time + +import numpy as np +import pytest +import torch + +pytest.importorskip("transformers") + +try: + from cactus.transpile.capture_pytorch import capture_model + from cactus.transpile.lower import transpile_ir +except Exception as exc: # pragma: no cover + pytest.skip(f"cactus runtime library unavailable: {exc}", allow_module_level=True) + +MODEL_ID = "SimianLuo/LCM_Dreamshaper_v7" +BOS, EOS, MAX_LEN = 49406, 49407, 77 + + +class _LastHiddenState(torch.nn.Module): + def __init__(self, text_model): + super().__init__() + self.text_model = text_model + + def forward(self, input_ids): + return self.text_model(input_ids).last_hidden_state + + +def _text_encoder(dtype): + from transformers import CLIPTextModel + + try: + encoder = CLIPTextModel.from_pretrained( + MODEL_ID, subfolder="text_encoder", torch_dtype=dtype, local_files_only=True + ) + except Exception as exc: + pytest.skip(f"{MODEL_ID} not cached: {exc}") + return _LastHiddenState(encoder.eval()) + + +def _token_ids(): + torch.manual_seed(0) + ids = torch.randint(0, BOS, (1, MAX_LEN)) + ids[0, 0] = BOS + ids[0, -1] = EOS + return ids + + +def test_clip_text_encoder_transpiles_to_the_expected_ops(): + captured = capture_model(_text_encoder(torch.float32), (_token_ids(),)) + nodes = captured.ir_graph.nodes + nodes = list(nodes.values()) if isinstance(nodes, dict) else list(nodes) + counts: dict[str, int] = {} + for node in nodes: + counts[node.op] = counts.get(node.op, 0) + 1 + + assert counts["linear"] == 72 + assert counts["layer_norm"] == 25 + assert counts["attention"] == 12 + assert counts["embedding"] == 2 + assert counts["sigmoid"] == 12 # quick_gelu lowers as x * sigmoid(1.702 x) + + +def test_clip_text_encode_matches_torch_and_reports_latency(): + ids = _token_ids() + transpiled = transpile_ir(capture_model(_text_encoder(torch.float16), (ids,)).ir_graph) + transpiled.set_input(0, ids.numpy()) + encoded = np.asarray(transpiled.execute()[0].numpy(), dtype=np.float32) + + runs = 5 + start = time.perf_counter() + for _ in range(runs): + transpiled.execute() + elapsed_ms = (time.perf_counter() - start) * 1000.0 / runs + + with torch.no_grad(): + reference = _text_encoder(torch.float32)(ids).float().numpy() + + assert encoded.shape == reference.shape == (1, MAX_LEN, 768) + diff = np.abs(reference - encoded) + relative = diff.mean() / np.abs(reference).mean() + print( + f"\n clip text encode 77 tokens: {elapsed_ms:.1f} ms" + f" max_abs_diff={diff.max():.4f} relative_mean_err={relative * 100:.3f}%" + ) + + assert relative < 0.01, f"transpiled encode drifted from torch: {relative * 100:.3f}%" From dde57efaeb099e2dd0355481b18e9a4c11e51a27 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sun, 9 Aug 2026 14:38:50 -0700 Subject: [PATCH 20/34] test the transpiled text-to-image pipeline end to end Signed-off-by: Neel Shejwalkar --- python/tests/test_transpile_t2i.py | 187 +++++++++++++++++++++++++++++ 1 file changed, 187 insertions(+) create mode 100644 python/tests/test_transpile_t2i.py diff --git a/python/tests/test_transpile_t2i.py b/python/tests/test_transpile_t2i.py new file mode 100644 index 000000000..e457d15b7 --- /dev/null +++ b/python/tests/test_transpile_t2i.py @@ -0,0 +1,187 @@ +"""End-to-end text-to-image through transpiled graphs only. + +CLIP text encoder, LCM UNet, and TAESD decoder each go through +capture_model -> transpile_ir; the LCM scheduler step stays host-side numpy. +The reference is the identical loop run in torch fp32. +""" +from __future__ import annotations + +import time + +import numpy as np +import pytest +import torch + +pytest.importorskip("diffusers") +pytest.importorskip("transformers") + +try: + from cactus.transpile.capture_pytorch import capture_model + from cactus.transpile.lower import transpile_ir +except Exception as exc: # pragma: no cover + pytest.skip(f"cactus runtime library unavailable: {exc}", allow_module_level=True) + +SD_MODEL_ID = "SimianLuo/LCM_Dreamshaper_v7" +TAESD_MODEL_ID = "madebyollin/taesd" +PROMPT = "a photograph of an astronaut riding a horse" +STEPS = 4 +GUIDANCE = 8.5 +SEED = 1234 +TIMESTEP_SCALING = 10.0 +SIGMA_DATA = 0.5 + + +def _load(cls_name, model_id, subfolder, dtype): + if cls_name == "CLIPTextModel": + from transformers import CLIPTextModel as cls + elif cls_name == "UNet2DConditionModel": + from diffusers import UNet2DConditionModel as cls + else: + from diffusers import AutoencoderTiny as cls + try: + kwargs = {"subfolder": subfolder} if subfolder else {} + return cls.from_pretrained(model_id, torch_dtype=dtype, local_files_only=True, **kwargs).eval() + except Exception as exc: + pytest.skip(f"{model_id} not cached: {exc}") + + +class _LastHiddenState(torch.nn.Module): + def __init__(self, text_model): + super().__init__() + self.text_model = text_model + + def forward(self, input_ids): + return self.text_model(input_ids).last_hidden_state + + +class _Denoiser(torch.nn.Module): + def __init__(self, unet): + super().__init__() + self.unet = unet + + def forward(self, sample, timestep, encoder_hidden_states, timestep_cond): + return self.unet( + sample, + timestep, + encoder_hidden_states=encoder_hidden_states, + timestep_cond=timestep_cond, + ).sample + + +def _guidance_embedding(w: float, embedding_dim: int = 256) -> torch.Tensor: + scaled = torch.tensor([w * 1000.0]) + half_dim = embedding_dim // 2 + freq = torch.exp(torch.arange(half_dim, dtype=torch.float32) * (-np.log(10000.0) / (half_dim - 1))) + angles = scaled[:, None] * freq[None, :] + return torch.cat([torch.sin(angles), torch.cos(angles)], dim=1) + + +def _scheduler_scalars(): + from diffusers import LCMScheduler + + sched = LCMScheduler.from_pretrained(SD_MODEL_ID, subfolder="scheduler", local_files_only=True) + assert sched.config.prediction_type == "epsilon" + assert float(sched.config.timestep_scaling) == TIMESTEP_SCALING + sched.set_timesteps(STEPS) + timesteps = sched.timesteps + rows = [] + for i, t in enumerate(timesteps): + prev_t = timesteps[i + 1] if i + 1 < len(timesteps) else t + a_t, a_prev = sched.alphas_cumprod[t], sched.alphas_cumprod[prev_t] + st = float(t) * TIMESTEP_SCALING + c_skip = SIGMA_DATA**2 / (st * st + SIGMA_DATA**2) + c_out = st / (st * st + SIGMA_DATA**2) ** 0.5 + rows.append((float(t), a_t.sqrt().item(), (1 - a_t).sqrt().item(), c_skip, c_out, + a_prev.sqrt().item(), (1 - a_prev).sqrt().item())) + return rows + + +def _lcm_loop(scalars, latent, noises, denoise): + x = latent + for i, (t, sa, sb, c_skip, c_out, sap, sbp) in enumerate(scalars): + eps = denoise(x, t) + x0 = (x - sb * eps) / sa + denoised = c_out * x0 + c_skip * x + x = sap * denoised + sbp * noises[i] if i < len(scalars) - 1 else denoised + return x + + +def test_t2i_pipeline_matches_torch(): + from transformers import CLIPTokenizer + + try: + tokenizer = CLIPTokenizer.from_pretrained(SD_MODEL_ID, subfolder="tokenizer", local_files_only=True) + except Exception as exc: + pytest.skip(f"{SD_MODEL_ID} not cached: {exc}") + + ids = tokenizer( + PROMPT, padding="max_length", max_length=77, truncation=True, return_tensors="pt" + ).input_ids + + scalars = _scheduler_scalars() + generator = torch.Generator("cpu").manual_seed(SEED) + latent = torch.randn(1, 4, 64, 64, generator=generator) + noises = [torch.randn(1, 4, 64, 64, generator=generator) for _ in range(STEPS - 1)] + w_emb = _guidance_embedding(GUIDANCE - 1.0) + + # --- transpiled pipeline --- + text_graph = transpile_ir( + capture_model(_LastHiddenState(_load("CLIPTextModel", SD_MODEL_ID, "text_encoder", torch.float16)), (ids,)).ir_graph + ) + start = time.perf_counter() + text_graph.set_input(0, ids.numpy()) + text_emb = np.asarray(text_graph.execute()[0].numpy(), dtype=np.float16) + clip_ms = (time.perf_counter() - start) * 1000.0 + + unet_example = ( + latent.to(torch.float16), + torch.tensor([scalars[0][0]], dtype=torch.float16), + torch.from_numpy(text_emb), + w_emb.to(torch.float16), + ) + unet_graph = transpile_ir( + capture_model(_Denoiser(_load("UNet2DConditionModel", SD_MODEL_ID, "unet", torch.float16)), unet_example).ir_graph + ) + + def denoise_cactus(x, t): + unet_graph.set_input(0, x.astype(np.float16)) + unet_graph.set_input(1, np.array([t], dtype=np.float16)) + unet_graph.set_input(2, text_emb) + unet_graph.set_input(3, w_emb.numpy().astype(np.float16)) + return np.asarray(unet_graph.execute()[0].numpy(), dtype=np.float32) + + start = time.perf_counter() + final_latent = _lcm_loop(scalars, latent.numpy().astype(np.float32), [n.numpy() for n in noises], denoise_cactus) + unet_ms = (time.perf_counter() - start) * 1000.0 / STEPS + + decoder = _load("AutoencoderTiny", TAESD_MODEL_ID, None, torch.float32).decoder + decoder_graph = transpile_ir(capture_model(decoder, (torch.randn(1, 4, 64, 64),)).ir_graph) + start = time.perf_counter() + decoder_graph.set_input(0, final_latent.astype(np.float16)) + image = np.asarray(decoder_graph.execute()[0].numpy(), dtype=np.float32).clip(0.0, 1.0) + decode_ms = (time.perf_counter() - start) * 1000.0 + + # --- torch fp32 reference, same math --- + ref_text = _LastHiddenState(_load("CLIPTextModel", SD_MODEL_ID, "text_encoder", torch.float32)) + ref_unet = _Denoiser(_load("UNet2DConditionModel", SD_MODEL_ID, "unet", torch.float32)) + with torch.no_grad(): + ref_emb = ref_text(ids) + + def denoise_torch(x, t): + with torch.no_grad(): + return ref_unet(torch.from_numpy(x), torch.tensor([t]), ref_emb, w_emb).numpy() + + ref_latent = _lcm_loop(scalars, latent.numpy().astype(np.float32), [n.numpy() for n in noises], denoise_torch) + ref_image = decoder(torch.from_numpy(ref_latent)).clamp(0.0, 1.0).numpy() + + assert image.shape == ref_image.shape == (1, 3, 512, 512) + latent_rel = np.abs(final_latent - ref_latent).mean() / np.abs(ref_latent).mean() + mse = np.mean((image - ref_image) ** 2) + psnr = 10.0 * np.log10(1.0 / mse) if mse > 0 else np.inf + print( + f"\n t2i 512x512 {STEPS}-step: clip {clip_ms:.0f} ms unet {unet_ms:.0f} ms/step decode {decode_ms:.0f} ms" + f" latent_rel_err={latent_rel * 100:.2f}% psnr={psnr:.1f} dB" + ) + + assert latent_rel < 0.05, f"final latent drifted from torch: {latent_rel * 100:.2f}%" + assert psnr > 35.0, f"decoded image diverged from torch: {psnr:.1f} dB" From 1bb943fec08704a8821bd5e7214ec81f524aeb8f Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sun, 9 Aug 2026 14:38:50 -0700 Subject: [PATCH 21/34] convert the SD text encoder and UNet through family adapters 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 --- .../cactus/convert/model_adapters/adapters.py | 71 ++++++++++++++++-- .../convert/model_adapters/detection.py | 6 +- .../cactus/convert/tests/test_sd_adapters.py | 74 +++++++++++++++++++ 3 files changed, 144 insertions(+), 7 deletions(-) create mode 100644 python/cactus/convert/tests/test_sd_adapters.py diff --git a/python/cactus/convert/model_adapters/adapters.py b/python/cactus/convert/model_adapters/adapters.py index 047bb240a..0e17bbb83 100644 --- a/python/cactus/convert/model_adapters/adapters.py +++ b/python/cactus/convert/model_adapters/adapters.py @@ -502,6 +502,15 @@ def build_calibration_inputs(self, row: dict[str, Any], processor: Any, modality return None +def _verbatim_name_tensor(source_name: str, component: str) -> NameMatch: + """Keep the full dotted checkpoint name, so parallel stacks never collide.""" + if not (source_name.endswith(".weight") or source_name.endswith(".bias")): + return NameMatch(source_name, None, component, False, hf_name=source_name, adapter_name=source_name) + generic = source_name.replace(".", "_") + output = generic[:-7] + ".weights" if generic.endswith("_weight") else generic[:-5] + ".bias" + return NameMatch(source_name, output, component, True, hf_name=source_name, adapter_name=source_name) + + class TaesdAdapter(FamilyAdapter): family = "taesd" @@ -520,12 +529,7 @@ def load_processor(self, model_id_or_path: str): return None def name_tensor(self, source_name: str, _tensor: Any, _num_layers: int | None) -> NameMatch: - component = "vision" - if not (source_name.endswith(".weight") or source_name.endswith(".bias")): - return NameMatch(source_name, None, component, False, hf_name=source_name, adapter_name=source_name) - generic = source_name.replace(".", "_") - output = generic[:-7] + ".weights" if generic.endswith("_weight") else generic[:-5] + ".bias" - return NameMatch(source_name, output, component, True, hf_name=source_name, adapter_name=source_name) + return _verbatim_name_tensor(source_name, "vision") def policy(self, match: NameMatch, _shape: tuple[int, ...], _requested_bits: int) -> TensorPolicy: if match.output_name is None: @@ -533,6 +537,59 @@ def policy(self, match: NameMatch, _shape: tuple[int, ...], _requested_bits: int return TensorPolicy("fallback", "FP16", None, match.component, False, "none", "tiny autoencoder conv tensor") +class SdUnetAdapter(FamilyAdapter): + family = "sd_unet" + + def runtime_config(self, cfg: Any) -> dict[str, Any]: + return { + "in_channels": int(_cfg_get(cfg, "in_channels", 4) or 4), + "sample_size": int(_cfg_get(cfg, "sample_size", 64) or 64), + "cross_attention_dim": int(_cfg_get(cfg, "cross_attention_dim", 768) or 768), + "block_out_channels": [int(c) for c in _cfg_get(cfg, "block_out_channels", [320, 640, 1280, 1280])], + "time_cond_proj_dim": int(_cfg_get(cfg, "time_cond_proj_dim", 0) or 0), + } + + def model_class(self, cfg: Any): + return None + + def load_processor(self, model_id_or_path: str): + return None + + def name_tensor(self, source_name: str, _tensor: Any, _num_layers: int | None) -> NameMatch: + return _verbatim_name_tensor(source_name, "vision") + + def policy(self, match: NameMatch, _shape: tuple[int, ...], _requested_bits: int) -> TensorPolicy: + if match.output_name is None: + return TensorPolicy("ignored", "none", None, match.component, False, "none", "no output filename") + return TensorPolicy("fallback", "FP16", None, match.component, False, "none", "diffusion unet tensor") + + +class ClipTextAdapter(FamilyAdapter): + family = "clip_text" + + def runtime_config(self, cfg: Any) -> dict[str, Any]: + return { + "hidden_size": int(_cfg_get(cfg, "hidden_size", 768) or 768), + "num_hidden_layers": int(_cfg_get(cfg, "num_hidden_layers", 12) or 12), + "max_position_embeddings": int(_cfg_get(cfg, "max_position_embeddings", 77) or 77), + "vocab_size": int(_cfg_get(cfg, "vocab_size", 49408) or 49408), + } + + def model_class(self, cfg: Any): + return None + + def load_processor(self, model_id_or_path: str): + return None + + def name_tensor(self, source_name: str, _tensor: Any, _num_layers: int | None) -> NameMatch: + return _verbatim_name_tensor(source_name, "language") + + def policy(self, match: NameMatch, _shape: tuple[int, ...], _requested_bits: int) -> TensorPolicy: + if match.output_name is None: + return TensorPolicy("ignored", "none", None, match.component, False, "none", "no output filename") + return TensorPolicy("fallback", "FP16", None, match.component, False, "none", "clip text encoder tensor") + + class ParakeetAdapter(FamilyAdapter): family = "parakeet" @@ -916,6 +973,8 @@ def scale_factor(self, output_name: str) -> float: "parakeet": ParakeetAdapter(), "parakeet_tdt": ParakeetTDTAdapter(), "taesd": TaesdAdapter(), + "sd_unet": SdUnetAdapter(), + "clip_text": ClipTextAdapter(), } diff --git a/python/cactus/convert/model_adapters/detection.py b/python/cactus/convert/model_adapters/detection.py index 1b8720a72..a5e78928a 100644 --- a/python/cactus/convert/model_adapters/detection.py +++ b/python/cactus/convert/model_adapters/detection.py @@ -3,7 +3,7 @@ from ..cactus_adapters.config_utils import cfg_get, detect_model_type -SUPPORTED_FAMILIES = {"auto", "gemma3", "gemma4", "qwen", "lfm2", "whisper", "parakeet", "parakeet_tdt", "moonshine", "nomic", "needle", "taesd", "generic"} +SUPPORTED_FAMILIES = {"auto", "gemma3", "gemma4", "qwen", "lfm2", "whisper", "parakeet", "parakeet_tdt", "moonshine", "nomic", "needle", "taesd", "sd_unet", "clip_text", "generic"} def detect_family(config, requested: str = "auto") -> str: @@ -11,6 +11,10 @@ def detect_family(config, requested: str = "auto") -> str: return requested if cfg_get(config, "_class_name", "") == "AutoencoderTiny": return "taesd" + if cfg_get(config, "_class_name", "") == "UNet2DConditionModel": + return "sd_unet" + if str(cfg_get(config, "model_type", "") or "").lower() == "clip_text_model": + return "clip_text" text_config = cfg_get(config, "text_config", None) base = text_config if text_config is not None else config detected = detect_model_type(base, config) diff --git a/python/cactus/convert/tests/test_sd_adapters.py b/python/cactus/convert/tests/test_sd_adapters.py new file mode 100644 index 000000000..4777d5827 --- /dev/null +++ b/python/cactus/convert/tests/test_sd_adapters.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +import torch + +from cactus.convert.model_adapters.adapters import adapter_for_family +from cactus.convert.model_adapters.detection import detect_family + + +def test_diffusers_autoencoder_tiny_config_selects_taesd(): + assert detect_family({"_class_name": "AutoencoderTiny", "latent_channels": 4}) == "taesd" + + +def test_taesd_converts_from_checkpoint_without_a_transformers_model_class(): + adapter = adapter_for_family("taesd") + assert adapter.model_class({"_class_name": "AutoencoderTiny"}) is None + assert adapter.load_processor("madebyollin/taesd") is None + + +def test_taesd_keeps_encoder_and_decoder_names_distinct(): + adapter = adapter_for_family("taesd") + tensor = torch.zeros(64, 64, 3, 3) + encoder = adapter.name_tensor("encoder.layers.12.conv.0.weight", tensor, None) + decoder = adapter.name_tensor("decoder.layers.12.conv.0.weight", tensor, None) + assert encoder.output_name == "encoder_layers_12_conv_0.weights" + assert decoder.output_name == "decoder_layers_12_conv_0.weights" + + +def test_taesd_conv_tensors_stay_fp16(): + adapter = adapter_for_family("taesd") + match = adapter.name_tensor("decoder.layers.0.weight", torch.zeros(64, 4, 3, 3), None) + policy = adapter.policy(match, (64, 4, 3, 3), 4) + assert policy.precision == "FP16" + + +def test_diffusers_unet_config_selects_sd_unet(): + assert detect_family({"_class_name": "UNet2DConditionModel", "cross_attention_dim": 768}) == "sd_unet" + + +def test_clip_text_config_selects_clip_text(): + assert detect_family({"model_type": "clip_text_model", "hidden_size": 768}) == "clip_text" + + +def test_sd_unet_converts_from_checkpoint_without_a_transformers_model_class(): + adapter = adapter_for_family("sd_unet") + assert adapter.model_class({"_class_name": "UNet2DConditionModel"}) is None + assert adapter.load_processor("SimianLuo/LCM_Dreamshaper_v7") is None + + +def test_sd_unet_keeps_parallel_block_names_distinct(): + adapter = adapter_for_family("sd_unet") + tensor = torch.zeros(320, 320, 3, 3) + down = adapter.name_tensor("down_blocks.0.resnets.0.conv1.weight", tensor, None) + up = adapter.name_tensor("up_blocks.0.resnets.0.conv1.weight", tensor, None) + assert down.output_name == "down_blocks_0_resnets_0_conv1.weights" + assert up.output_name == "up_blocks_0_resnets_0_conv1.weights" + + +def test_sd_unet_tensors_stay_fp16(): + adapter = adapter_for_family("sd_unet") + match = adapter.name_tensor("mid_block.attentions.0.proj_out.weight", torch.zeros(1280, 1280), None) + policy = adapter.policy(match, (1280, 1280), 4) + assert policy.precision == "FP16" + + +def test_clip_text_keeps_layer_names_distinct_and_fp16(): + adapter = adapter_for_family("clip_text") + tensor = torch.zeros(768, 768) + q0 = adapter.name_tensor("text_model.encoder.layers.0.self_attn.q_proj.weight", tensor, None) + q1 = adapter.name_tensor("text_model.encoder.layers.1.self_attn.q_proj.weight", tensor, None) + assert q0.output_name == "text_model_encoder_layers_0_self_attn_q_proj.weights" + assert q1.output_name == "text_model_encoder_layers_1_self_attn_q_proj.weights" + policy = adapter.policy(q0, (768, 768), 4) + assert policy.precision == "FP16" + assert policy.component == "language" From 82e01c06d1be8a29436a9eda6012ab526880ab73 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sun, 9 Aug 2026 14:58:37 -0700 Subject: [PATCH 22/34] lower upsample_nearest2d in the replacement transpiler 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 --- .../cactus/transpiler/Generator/constants.py | 6 + .../Generator/lowerings/lowering_core.py | 3 +- .../Generator/lowerings/lowering_nn_ops.py | 26 ++++ python/cactus/transpiler/IR/constants.py | 2 + python/tests/test_transpiler_upsample.py | 143 ++++++++++++++++++ 5 files changed, 179 insertions(+), 1 deletion(-) create mode 100644 python/tests/test_transpiler_upsample.py diff --git a/python/cactus/transpiler/Generator/constants.py b/python/cactus/transpiler/Generator/constants.py index 1bdff2dc9..620145d47 100644 --- a/python/cactus/transpiler/Generator/constants.py +++ b/python/cactus/transpiler/Generator/constants.py @@ -359,6 +359,12 @@ "aten.conv2d.default", } +UPSAMPLE_TARGETS = { + "cactus.upsample_nearest2d", + "aten.upsample_nearest2d.default", + "aten.upsample_nearest2d.vec", +} + ATTENTION_TARGETS = { "cactus.attention", "cactus.attention_cached", diff --git a/python/cactus/transpiler/Generator/lowerings/lowering_core.py b/python/cactus/transpiler/Generator/lowerings/lowering_core.py index 024d3e6c8..8c817f9b3 100644 --- a/python/cactus/transpiler/Generator/lowerings/lowering_core.py +++ b/python/cactus/transpiler/Generator/lowerings/lowering_core.py @@ -9,7 +9,7 @@ from .lowering_basic_ops import * from .lowering_cache import cache_attention_generation_plan, lower_cache, lower_cache_placeholder, should_lower_cache_placeholder_as_state, lower_attention from .lowering_constant_ops import * -from .lowering_nn_ops import lower_conv, lower_norm +from .lowering_nn_ops import lower_conv, lower_norm, lower_upsample from .lowering_special_ops import lower_moe, lower_special_cactus, lower_unsupported_semantic from .lowering_utils import * from .lowering_weights import bind_weight_placeholder, lower_lfm_grouped_moe_placeholder, should_dequantize_int8_weight_placeholder @@ -79,6 +79,7 @@ def build_lowering_rules( add_rules(rules, constants.CLAMP_TARGETS, lower_clamp) add_rules(rules, constants.NORM_TARGETS, lower_norm) add_rules(rules, constants.CONV_TARGETS, lower_conv) + add_rules(rules, constants.UPSAMPLE_TARGETS, lower_upsample) add_rules(rules, constants.ATTENTION_TARGETS, lower_attention) add_rules(rules, constants.CACHE_TARGETS, lower_cache) add_rules(rules, constants.MOE_TARGETS, lower_moe) diff --git a/python/cactus/transpiler/Generator/lowerings/lowering_nn_ops.py b/python/cactus/transpiler/Generator/lowerings/lowering_nn_ops.py index 3a3ca4325..44dd6eb5d 100644 --- a/python/cactus/transpiler/Generator/lowerings/lowering_nn_ops.py +++ b/python/cactus/transpiler/Generator/lowerings/lowering_nn_ops.py @@ -179,6 +179,32 @@ def lower_specialized_conv2d(context: models.GenerationContext, node: IRModels.N return context.graph.conv2d_pointwise_1x1(inputs[0], inputs[1], bias=bias) return None +def lower_upsample(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + require_len(node, inputs, 1) + return context.graph.upsample_nearest2d(inputs[0], upsample_scale(node)) + +def upsample_scale(node: IRModels.Node) -> int: + scale_factors = node.attrs.get("scale_factors") + if scale_factors is None and node.attrs.get("scale_factor") is not None: + scale_factors = (node.attrs["scale_factor"], node.attrs["scale_factor"]) + if scale_factors is None and node.attrs.get("scales_h") is not None and node.attrs.get("scales_w") is not None: + scale_factors = (node.attrs["scales_h"], node.attrs["scales_w"]) + if scale_factors is None: + output_size = node.attrs.get("output_size") + input_shape = meta_shape(node.parents[0]) if node.parents else () + if isinstance(output_size, (list, tuple)) and len(output_size) == 2 and len(input_shape) == 4: + height, width = int(input_shape[2]), int(input_shape[3]) + if height > 0 and width > 0 and int(output_size[0]) % height == 0 and int(output_size[1]) % width == 0: + scale_factors = (int(output_size[0]) // height, int(output_size[1]) // width) + if isinstance(scale_factors, (list, tuple)) and len(scale_factors) == 2: + first, second = float(scale_factors[0]), float(scale_factors[1]) + if first == second and first.is_integer() and first >= 1: + return int(first) + raise UnsupportedLoweringError( + f"{node.name}: upsample_nearest2d requires one integer scale for both axes; attrs={node.attrs}" + ) + def require_plain_conv1d_attrs(node: IRModels.Node, method: str) -> None: padding = first_int(node.attrs.get("padding"), 0) dilation = first_int(node.attrs.get("dilation"), 1) diff --git a/python/cactus/transpiler/IR/constants.py b/python/cactus/transpiler/IR/constants.py index 11b59ec5e..cb6fb44f5 100644 --- a/python/cactus/transpiler/IR/constants.py +++ b/python/cactus/transpiler/IR/constants.py @@ -79,6 +79,8 @@ }, "aten.conv1d.default": {3: "stride", 4: "padding", 5: "dilation", 6: "groups"}, "aten.conv2d.default": {3: "stride", 4: "padding", 5: "dilation", 6: "groups"}, + "aten.upsample_nearest2d.default": {1: "output_size", 2: "scales_h", 3: "scales_w"}, + "aten.upsample_nearest2d.vec": {1: "output_size", 2: "scale_factors"}, "aten.native_group_norm.default": {3: "batch_size", 4: "channels", 5: "height_width", 6: "num_groups", 7: "eps"}, "aten.group_norm.default": {1: "num_groups", 4: "eps"}, "aten.native_batch_norm.default": {5: "training", 6: "momentum", 7: "eps"}, diff --git a/python/tests/test_transpiler_upsample.py b/python/tests/test_transpiler_upsample.py new file mode 100644 index 000000000..533d0c76d --- /dev/null +++ b/python/tests/test_transpiler_upsample.py @@ -0,0 +1,143 @@ +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +from cactus.transpiler.Converter import models as CModels +from cactus.transpiler.Fusions import models as FModels +from cactus.transpiler.Generator import models as GModels +from cactus.transpiler.Generator.lowerings import lowering_core +from cactus.transpiler.Generator.lowerings.lowering_nn_ops import upsample_scale +from cactus.transpiler.IR import models as IRModels + + +def upsample_graph(target: str, args: list, attrs: dict | None = None) -> IRModels.Graph: + x = IRModels.Node( + index=0, + name="x", + node_type="placeholder", + target="x", + args=[], + kwargs={}, + users=(), + tensor_output_meta={"shape": [1, 2, 2, 3], "dtype": "torch.float16"}, + module_stack=None, + value_kind=FModels.ValueKind.USER_INPUT, + ir_metadata={"logical_input": "x"}, + ) + up = IRModels.Node( + index=1, + name="up", + node_type="call_function", + target=target, + args=args, + kwargs={}, + users=(), + tensor_output_meta={"shape": [1, 2, 4, 6], "dtype": "torch.float16"}, + module_stack=None, + value_kind=FModels.ValueKind.ACTIVATION, + attrs=attrs or {}, + ) + output = IRModels.Node( + index=2, + name="output", + node_type="output", + target="output", + args=[[{"node": "up"}]], + kwargs={}, + users=(), + tensor_output_meta=None, + module_stack=None, + value_kind=FModels.ValueKind.OUTPUT, + ) + return IRModels.rebuild_graph( + (x, up, output), + IRModels.Graph(source=None, sources=(), outputs=(), nodes=(), nodes_map={}), + ) + + +class TestTranspilerUpsample(unittest.TestCase): + def test_ir_extracts_attrs_for_both_aten_overloads(self): + records = [ + CModels.LayerRecord( + index=0, name="x", node_type="placeholder", target="x", args=[], kwargs={}, + users=[], tensor_output_meta={"shape": [1, 2, 2, 3], "dtype": "torch.float16"}, + module_stack=None, + ), + CModels.LayerRecord( + index=1, name="up_vec", node_type="call_function", + target="aten.upsample_nearest2d.vec", + args=[{"node": "x"}, None, [2.0, 2.0]], kwargs={}, users=[], + tensor_output_meta={"shape": [1, 2, 4, 6], "dtype": "torch.float16"}, + module_stack=None, + ), + CModels.LayerRecord( + index=2, name="up_size", node_type="call_function", + target="aten.upsample_nearest2d.default", + args=[{"node": "up_vec"}, [8, 12]], kwargs={}, users=[], + tensor_output_meta={"shape": [1, 2, 8, 12], "dtype": "torch.float16"}, + module_stack=None, + ), + ] + layer_map = CModels.LayerMap( + model_name="upsample-attrs-test", task="prefill_no_cache", + graph_signature="", range_constants="", nodes=records, + ) + graph = IRModels.Graph.from_map(layer_map) + + self.assertEqual(graph.nodes_map["up_vec"].attrs.get("scale_factors"), [2.0, 2.0]) + self.assertEqual(graph.nodes_map["up_size"].attrs.get("output_size"), [8, 12]) + + def test_lowering_rules_cover_every_upsample_target(self): + rules = lowering_core.build_lowering_rules() + for target in ( + "cactus.upsample_nearest2d", + "aten.upsample_nearest2d.default", + "aten.upsample_nearest2d.vec", + ): + self.assertIn(target, rules) + + def test_scale_derives_from_output_size_and_rejects_anisotropic(self): + graph = upsample_graph("aten.upsample_nearest2d.vec", [{"node": "x"}], {"output_size": [4, 6]}) + self.assertEqual(upsample_scale(graph.nodes_map["up"]), 2) + + from cactus.transpiler.Generator.errors import UnsupportedLoweringError + + bad = upsample_graph("aten.upsample_nearest2d.vec", [{"node": "x"}], {"scale_factors": [2.0, 3.0]}) + with self.assertRaises(UnsupportedLoweringError): + upsample_scale(bad.nodes_map["up"]) + + def test_lowered_graph_executes_exact_nearest_neighbour(self): + try: + from cactus.bindings.cactus import Tensor + except Exception as exc: # pragma: no cover + self.skipTest(f"cactus runtime library unavailable: {exc}") + + graph = upsample_graph( + "aten.upsample_nearest2d.vec", + [{"node": "x"}, None, [2.0, 2.0]], + {"output_size": None, "scale_factors": [2.0, 2.0]}, + ) + with tempfile.TemporaryDirectory() as tmp: + config = GModels.GeneratorConfig(output_dir=Path(tmp)) + component = GModels.ComponentGraph.from_ir("upsample", graph, config) + lowering_core.lower_component(component, config, lowering_core.build_lowering_rules()) + + self.assertEqual(component.unsupported_nodes, []) + self.assertEqual(len(component.runtime_input_ids), 1) + self.assertEqual(len(component.output_node_ids), 1) + + g = component.graph + g.retain_outputs(component.output_node_ids) + data = np.arange(12, dtype=np.float16).reshape(1, 2, 2, 3) + g.set_input(Tensor(g, component.runtime_input_ids[0], (1, 2, 2, 3), g.FP16), data) + g.execute() + result = Tensor(g, component.output_node_ids[0], (1, 2, 4, 6), g.FP16).numpy() + + expected = data.repeat(2, axis=2).repeat(2, axis=3) + np.testing.assert_array_equal(np.asarray(result, dtype=np.float16), expected) + + +if __name__ == "__main__": + unittest.main() From 92aa1fcae463982332359dde98f60f3a219bf795 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sun, 9 Aug 2026 15:28:59 -0700 Subject: [PATCH 23/34] distinguish conv1d and conv2d fusions by kernel rank 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 --- python/cactus/transpiler/Fusions/fusions.py | 18 +++++++++--------- python/cactus/transpiler/IR/match_utils.py | 3 +++ 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/python/cactus/transpiler/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py index b71be4a27..35f0f05fb 100644 --- a/python/cactus/transpiler/Fusions/fusions.py +++ b/python/cactus/transpiler/Fusions/fusions.py @@ -94,18 +94,18 @@ metadata={"special_matcher": "lfm_bmm_masked_attention"}, ), "rope": _definition("rope", "rope", ROPE_GRAPH, fusion_fields=("generic", "rope", "gemma4_rope", "qwen2_5_rope")), - "conv": _definition("conv", "conv1d", CONV_GRAPH, fusion_fields=("generic", "conv", "audio", "vision"), metadata=_required_attrs(padding=0, dilation=1, groups=1)), - "conv_bias": _definition("conv_bias", "conv1d", CONV_BIAS_GRAPH, fusion_fields=("generic", "conv", "audio", "vision")), + "conv": _definition("conv", "conv1d", CONV_GRAPH, fusion_fields=("generic", "conv", "audio", "vision"), metadata=_required_attrs(padding=0, dilation=1, groups=1, spatial_dims=1)), + "conv_bias": _definition("conv_bias", "conv1d", CONV_BIAS_GRAPH, fusion_fields=("generic", "conv", "audio", "vision"), metadata=_required_attrs(spatial_dims=1)), "lfm_short_conv_prefill": _definition("lfm_short_conv_prefill", "conv1d_causal", LFM_SHORT_CONV_PREFILL_GRAPH, fusion_fields=("generic", "conv", "cache", "lfm_moe"), supported_inference_modes=("prefill_with_cache",)), "lfm_short_conv_decode": _definition("lfm_short_conv_decode", "lfm_short_conv_decode", LFM_SHORT_CONV_DECODE_GRAPH, fusion_fields=("generic", "conv", "cache", "lfm_moe"), supported_inference_modes=("decode_with_cache",)), - "conv1d_k3": _definition("conv1d_k3", "conv1d_k3", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=3, stride=1, padding=0, dilation=1, groups=1)), - "conv1d_k7s3": _definition("conv1d_k7s3", "conv1d_k7s3", CONV_BIAS_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=7, stride=3)), + "conv1d_k3": _definition("conv1d_k3", "conv1d_k3", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=3, stride=1, padding=0, dilation=1, groups=1, spatial_dims=1)), + "conv1d_k7s3": _definition("conv1d_k7s3", "conv1d_k7s3", CONV_BIAS_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=7, stride=3, spatial_dims=1)), "conv1d_causal": _definition("conv1d_causal", "conv1d_causal", CONV_GRAPH, fusion_fields=("generic", "conv", "cache"), metadata=_required_attrs(causal=True)), - "conv1d_pointwise": _definition("conv1d_pointwise", "conv1d_pointwise", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=1)), - "conv1d_same_depthwise_k9": _definition("conv1d_same_depthwise_k9", "conv1d_same_depthwise_k9", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=9, depthwise=True)), - "conv2d_k3s2p1": _definition("conv2d_k3s2p1", "conv2d_k3s2p1", CONV_GRAPH, fusion_fields=("generic", "conv", "vision"), metadata=_required_attrs(kernel_size=3, stride=2, padding=1)), - "conv2d_depthwise_k3s2p1": _definition("conv2d_depthwise_k3s2p1", "conv2d_depthwise_k3s2p1", CONV_GRAPH, fusion_fields=("generic", "conv", "vision"), metadata=_required_attrs(kernel_size=3, stride=2, padding=1, depthwise=True)), - "conv2d_pointwise_1x1": _definition("conv2d_pointwise_1x1", "conv2d_pointwise_1x1", CONV_GRAPH, fusion_fields=("generic", "conv", "vision"), metadata=_required_attrs(kernel_size=1)), + "conv1d_pointwise": _definition("conv1d_pointwise", "conv1d_pointwise", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=1, spatial_dims=1)), + "conv1d_same_depthwise_k9": _definition("conv1d_same_depthwise_k9", "conv1d_same_depthwise_k9", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=9, depthwise=True, spatial_dims=1)), + "conv2d_k3s2p1": _definition("conv2d_k3s2p1", "conv2d_k3s2p1", CONV_GRAPH, fusion_fields=("generic", "conv", "vision"), metadata=_required_attrs(kernel_size=3, stride=2, padding=1, spatial_dims=2)), + "conv2d_depthwise_k3s2p1": _definition("conv2d_depthwise_k3s2p1", "conv2d_depthwise_k3s2p1", CONV_GRAPH, fusion_fields=("generic", "conv", "vision"), metadata=_required_attrs(kernel_size=3, stride=2, padding=1, depthwise=True, spatial_dims=2)), + "conv2d_pointwise_1x1": _definition("conv2d_pointwise_1x1", "conv2d_pointwise_1x1", CONV_GRAPH, fusion_fields=("generic", "conv", "vision"), metadata=_required_attrs(kernel_size=1, spatial_dims=2)), "kv_cache_append": _definition("kv_cache_append", "kv_cache_append", KV_CACHE_APPEND_GRAPH, fusion_fields=("generic", "cache"), supported_inference_modes=("prefill_with_cache", "decode_with_cache")), "kv_cache_initial_append": _definition("kv_cache_initial_append", "kv_cache_append", KV_CACHE_INITIAL_APPEND_GRAPH, fusion_fields=("generic", "cache"), supported_inference_modes=("prefill_with_cache",)), "attention_cached": _definition("attention_cached", "attention_cached", ATTENTION_CACHED_GRAPH, fusion_fields=("generic", "attention", "cache"), supported_inference_modes=("decode_with_cache",)), diff --git a/python/cactus/transpiler/IR/match_utils.py b/python/cactus/transpiler/IR/match_utils.py index 6eb9068be..05c319ec7 100644 --- a/python/cactus/transpiler/IR/match_utils.py +++ b/python/cactus/transpiler/IR/match_utils.py @@ -368,6 +368,9 @@ def infer_conv_attrs(fusion: FModels.FusionGraph, bindings: dict[str, models.Nod if len(weight_shape) >= 3: kernel_shape = weight_shape[2:] attrs["kernel_size"] = kernel_shape[0] if all(values_equal(dim, kernel_shape[0]) for dim in kernel_shape) else kernel_shape + # A square kernel collapses to one scalar, so conv1d and conv2d + # definitions need the rank to stay distinguishable. + attrs["spatial_dims"] = len(kernel_shape) if len(weight_shape) >= 2: attrs["depthwise"] = not values_equal(groups, 1) and values_equal(weight_shape[1], 1) From aa931c9fc2b5bc963ccf79ede7a7b4ae8448355b Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sun, 9 Aug 2026 15:28:59 -0700 Subject: [PATCH 24/34] wire SD text-to-image through the replacement transpiler 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 --- python/cactus/cli/model.py | 55 +++++++ python/cactus/cli/transpiler.py | 1 + .../cactus/transpiler/Converter/constants.py | 1 + .../cactus/transpiler/Converter/diffusion.py | 144 ++++++++++++++++++ python/cactus/transpiler/Converter/models.py | 6 +- .../cactus/transpiler/Converter/overrides.py | 13 ++ python/cactus/transpiler/Generator/models.py | 8 + .../transpiler/ModelProfiles/combinations.py | 2 + .../transpiler/ModelProfiles/components.py | 10 ++ .../cactus/transpiler/ModelProfiles/models.py | 2 + .../transpiler/ModelProfiles/profiles.py | 53 +++++++ .../cactus/transpiler/ModelProfiles/routes.py | 9 ++ python/tests/test_transpiler_sd_profile.py | 112 ++++++++++++++ 13 files changed, 415 insertions(+), 1 deletion(-) create mode 100644 python/cactus/transpiler/Converter/diffusion.py create mode 100644 python/tests/test_transpiler_sd_profile.py diff --git a/python/cactus/cli/model.py b/python/cactus/cli/model.py index fd930632e..dd1af4c5f 100644 --- a/python/cactus/cli/model.py +++ b/python/cactus/cli/model.py @@ -47,6 +47,13 @@ def _convert_from_source(model_id, *, bits, token, weights_dir, skip_model_load= err = convert_toolchain_error() if err: raise RuntimeError(err) + from .transpiler import profile_for_model_id + + component_sources = tuple(getattr(profile_for_model_id(model_id), "component_sources", ()) or ()) + if component_sources: + return _convert_component_sources( + model_id, component_sources, bits=bits, token=token, weights_dir=weights_dir + ) print_color(YELLOW, f"Converting {model_id} from HuggingFace source...") from ..convert.cli import main as cq_main @@ -66,6 +73,54 @@ def _convert_from_source(model_id, *, bits, token, weights_dir, skip_model_load= return weights_dir +def _convert_component_sources(model_id, component_sources, *, bits, token, weights_dir): + """Convert a multi-repo pipeline (diffusion) component by component into one + weights dir, merging each run's weights manifest.""" + import json + import tempfile + + from huggingface_hub import snapshot_download + + from ..convert.cli import main as cq_main + + weights_dir = Path(weights_dir) + weights_dir.mkdir(parents=True, exist_ok=True) + merged_records = [] + converted = [] + for mode, spec in component_sources: + _kind, _, source = spec.partition(":") + if "/" in source: + component_src = source + else: + snapshot = snapshot_download(model_id, allow_patterns=[f"{source}/*"], token=token) + component_src = str(Path(snapshot) / source) + print_color(YELLOW, f"Converting {mode} component from {component_src}...") + with tempfile.TemporaryDirectory() as tmp: + component_out = Path(tmp) / "out" + cq_main(["convert", "--model", component_src, "--out", str(component_out), "--bits", str(bits)]) + manifest_path = component_out / "weights_manifest.json" + if manifest_path.exists(): + merged_records.extend(json.loads(manifest_path.read_text(encoding="utf-8")).get("weights", [])) + for tensor_file in sorted(component_out.iterdir()): + if tensor_file.suffix not in {".weights", ".bias"}: + continue + destination = weights_dir / tensor_file.name + if destination.exists(): + raise RuntimeError( + f"components {converted} and {mode!r} both produce tensor file {tensor_file.name!r}" + ) + shutil.move(str(tensor_file), destination) + converted.append(mode) + (weights_dir / "weights_manifest.json").write_text( + json.dumps({"weights": merged_records}, indent=2), encoding="utf-8" + ) + (weights_dir / "config.txt").write_text( + "\n".join(f"component={mode}" for mode in converted) + "\n", encoding="utf-8" + ) + print_color(GREEN, f"Model converted and ready at {weights_dir}") + return weights_dir + + def ensure_weights(model_id, *, bits=4, token=None, reconvert=False, output_dir=None, skip_model_load=False): from .download import get_bundle_dir diff --git a/python/cactus/cli/transpiler.py b/python/cactus/cli/transpiler.py index d332c24f4..9410e2ac2 100644 --- a/python/cactus/cli/transpiler.py +++ b/python/cactus/cli/transpiler.py @@ -200,6 +200,7 @@ def resolve_transpile_config( profile=profile, modalities=default_modalities(profile), profile_source="registered", + inference_modes=tuple(getattr(profile, "export_modes", ()) or DEFAULT_TRANSPILER_MODES), ) resolved_task = generic_task or "causal-lm" diff --git a/python/cactus/transpiler/Converter/constants.py b/python/cactus/transpiler/Converter/constants.py index a330101cd..ca67ab695 100644 --- a/python/cactus/transpiler/Converter/constants.py +++ b/python/cactus/transpiler/Converter/constants.py @@ -18,6 +18,7 @@ DROP_MULTIMODAL_ON_DECODE_POLICY = "drop_multimodal_on_decode" SYNTHETIC_INPUT_STRATEGY = "synthetic" +DIFFUSION_INPUT_STRATEGY = "diffusion_components" DECODE_WITH_CACHE_MODE = "decode_with_cache" CONVERTER_JSON_DIR = Path(__file__).resolve().parent / "jsons" diff --git a/python/cactus/transpiler/Converter/diffusion.py b/python/cactus/transpiler/Converter/diffusion.py new file mode 100644 index 000000000..92aaf2e84 --- /dev/null +++ b/python/cactus/transpiler/Converter/diffusion.py @@ -0,0 +1,144 @@ +"""Diffusion pipeline components: one torch.export per component. + +A diffusion family declares `export_modes` and `component_sources` on its +profile; every mode loads its own module and synthesizes tensor inputs, so +nothing here goes through an HF processor. Model-specific facts (repo ids, +subfolders) stay in the profile -- this module only knows component kinds. +""" +from __future__ import annotations + +from typing import Any + +import torch + +from . import constants +from ..ModelProfiles import models as MP_Models + + +class TextEncoderLastHiddenState(torch.nn.Module): + """Wraps the inner CLIPTextTransformer so exported parameter targets match + the checkpoint keys (text_model.*) instead of doubling the prefix.""" + + def __init__(self, clip_text_model: torch.nn.Module): + super().__init__() + self.text_model = clip_text_model.text_model + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.text_model(input_ids).last_hidden_state + + +class TaesdDecode(torch.nn.Module): + """The attribute name keeps exported parameter targets on the checkpoint's + decoder.* keys.""" + + def __init__(self, decoder: torch.nn.Module): + super().__init__() + self.decoder = decoder + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.decoder(x) + + +class UnetEpsilon(torch.nn.Module): + def __init__(self, unet: torch.nn.Module): + super().__init__() + self.unet = unet + + def forward( + self, + sample: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep_cond: torch.Tensor, + ) -> torch.Tensor: + return self.unet( + sample, + timestep, + encoder_hidden_states=encoder_hidden_states, + timestep_cond=timestep_cond, + ).sample + + +def create_diffusion_model( + mp: MP_Models.ModelProfile, + model_id: str, + inference_mode: str, + input_cls: Any, + model_cls: Any, +) -> Any: + from .models import EXPORT_PATCHES, load_configs + + kind, source = component_source(mp, inference_mode) + configs = load_configs(mp, model_id) + for patch in mp.export_patches: + EXPORT_PATCHES[patch]() + module = load_component(kind, source, model_id) + input_ = input_cls( + args=(), + kwargs=component_export_inputs(kind, configs), + modalities=("text",), + inference_mode=inference_mode, + ) + return model_cls(name=model_id, model_profile=mp, input=input_, model=module) + + +def component_source(mp: MP_Models.ModelProfile, inference_mode: str) -> tuple[str, str]: + for mode, spec in getattr(mp, "component_sources", ()): + if mode == inference_mode: + kind, _, source = spec.partition(":") + return kind, source + raise ValueError( + f"profile {mp.model_profiles!r} declares no component source for mode {inference_mode!r}" + ) + + +def load_component(kind: str, source: str, model_id: str) -> torch.nn.Module: + if kind == "clip_text": + from transformers import CLIPTextModel + + return TextEncoderLastHiddenState(from_pretrained_local_first(CLIPTextModel, model_id, subfolder=source)) + if kind == "sd_unet": + from diffusers import UNet2DConditionModel + + return UnetEpsilon(from_pretrained_local_first(UNet2DConditionModel, model_id, subfolder=source)) + if kind == "taesd_decoder": + from diffusers import AutoencoderTiny + + return TaesdDecode(from_pretrained_local_first(AutoencoderTiny, source).decoder) + raise ValueError(f"unknown diffusion component kind {kind!r}") + + +def from_pretrained_local_first(model_class: Any, repo_id: str, **kwargs: Any) -> torch.nn.Module: + if constants.token is not None: + kwargs.setdefault("token", constants.token) + try: + model = model_class.from_pretrained(repo_id, torch_dtype=torch.float16, local_files_only=True, **kwargs) + except Exception: + model = model_class.from_pretrained(repo_id, torch_dtype=torch.float16, **kwargs) + return model.eval() + + +def component_export_inputs(kind: str, configs: dict[str, dict[str, Any]]) -> dict[str, torch.Tensor]: + unet_config = configs.get("unet/config.json", {}) + text_config = configs.get("text_encoder/config.json", {}) + tokens = int(text_config.get("max_position_embeddings", 77)) + latent_channels = int(unet_config.get("in_channels", 4)) + # Graphs are shape-specialized; export at 64x64 latents (512x512 images) + # rather than the config's training sample_size. + latent_size = 64 + if kind == "clip_text": + return {"input_ids": torch.zeros((1, tokens), dtype=torch.int64)} + if kind == "sd_unet": + return { + "sample": torch.zeros((1, latent_channels, latent_size, latent_size), dtype=torch.float16), + "timestep": torch.full((1,), 999.0, dtype=torch.float16), + "encoder_hidden_states": torch.zeros( + (1, tokens, int(unet_config.get("cross_attention_dim", 768))), dtype=torch.float16 + ), + "timestep_cond": torch.zeros( + (1, int(unet_config.get("time_cond_proj_dim", 256))), dtype=torch.float16 + ), + } + if kind == "taesd_decoder": + return {"x": torch.zeros((1, latent_channels, latent_size, latent_size), dtype=torch.float16)} + raise ValueError(f"unknown diffusion component kind {kind!r}") diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index 64ee5b6ce..fccde391d 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -17,6 +17,7 @@ token = constants.token EXPORT_PATCHES = { + "clip_position_ids": OV.patch_clip_position_ids_for_export, "gemma4_audio_mask": OV.patch_gemma4_audio_mask_for_export, "transformers_moe_grouped_mm_fallback": OV.patch_transformers_moe_grouped_mm_for_export, "lfm2_vl_image_features": OV.patch_lfm2_vl_image_features_for_export, @@ -380,7 +381,10 @@ def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch raise RuntimeError(f"Unable to load model {model_id}") from last_error def create_model(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], model_id: str, inference_mode: str = "prefill_no_cache") -> Model: - + if mp.input_strategy == constants.DIFFUSION_INPUT_STRATEGY: + from .diffusion import create_diffusion_model + + return create_diffusion_model(mp=mp, model_id=model_id, inference_mode=inference_mode, input_cls=Input, model_cls=Model) input_ = build_input(mp, input_modalities, Input, model_id, inference_mode) if input_ is None: raise ValueError(f"Could not build input for modalities {input_modalities}") diff --git a/python/cactus/transpiler/Converter/overrides.py b/python/cactus/transpiler/Converter/overrides.py index 40deb13be..04b3851f8 100644 --- a/python/cactus/transpiler/Converter/overrides.py +++ b/python/cactus/transpiler/Converter/overrides.py @@ -20,6 +20,19 @@ def exportable_bidirectional_mask(config, inputs_embeds, attention_mask=None, an return mask.expand(batch_size, 1, seq_len, seq_len) gemma4_modeling.create_bidirectional_mask = exportable_bidirectional_mask +def patch_clip_position_ids_for_export() -> None: + import transformers.models.clip.modeling_clip as clip_modeling + #CLIP registers position_ids as a non-persistent buffer, which torch.export lifts + #into a weight placeholder no checkpoint can bind; recompute it in the graph instead. + original_forward = clip_modeling.CLIPTextEmbeddings.forward + def forward_with_computed_position_ids(self, input_ids=None, position_ids=None, inputs_embeds=None): + if position_ids is None: + length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange(length, device=device).unsqueeze(0) + return original_forward(self, input_ids=input_ids, position_ids=position_ids, inputs_embeds=inputs_embeds) + clip_modeling.CLIPTextEmbeddings.forward = forward_with_computed_position_ids + def patch_transformers_moe_grouped_mm_for_export() -> None: import transformers.integrations.moe as moe moe._can_use_grouped_mm = lambda input, weight, offs: False diff --git a/python/cactus/transpiler/Generator/models.py b/python/cactus/transpiler/Generator/models.py index 274f164d0..d0e1741f1 100644 --- a/python/cactus/transpiler/Generator/models.py +++ b/python/cactus/transpiler/Generator/models.py @@ -438,6 +438,14 @@ def resolve_source_weight_record(resolver: WeightResolver, source_target: str) - record = resolver.records_by_name.get(variant) if record is not None: return record + # Export wrappers add one module attribute in front of the checkpoint key + # (e.g. unet.down_blocks...). Only consulted when every exact variant missed. + head, sep, rest = source_target.partition(".") + if sep and rest: + for variant in weight_name_variants(rest): + record = resolver.records_by_name.get(variant) + if record is not None: + return record return None def weight_name_variants(name: str) -> tuple[str, ...]: diff --git a/python/cactus/transpiler/ModelProfiles/combinations.py b/python/cactus/transpiler/ModelProfiles/combinations.py index ac09b5b05..cad4e73f5 100644 --- a/python/cactus/transpiler/ModelProfiles/combinations.py +++ b/python/cactus/transpiler/ModelProfiles/combinations.py @@ -10,3 +10,5 @@ AUDIO_TEXT_TO_DECODER = Combinations(input=("audio_encoder", "text_embed"), output="text_decoder") AUDIO_ENCODER_TO_ASR_DECODER = Combinations(input=("audio_encoder",), output="asr_decoder") AUDIO_ASR_DECODER_TO_ASR_HEAD = Combinations(input=("audio_encoder", "asr_decoder"), output="asr_head") +TEXT_ENCODER_TO_UNET = Combinations(input=("text_encoder",), output="unet") +UNET_TO_VAE_DECODER = Combinations(input=("unet",), output="vae_decoder") diff --git a/python/cactus/transpiler/ModelProfiles/components.py b/python/cactus/transpiler/ModelProfiles/components.py index 789bf5875..b077de88c 100644 --- a/python/cactus/transpiler/ModelProfiles/components.py +++ b/python/cactus/transpiler/ModelProfiles/components.py @@ -36,6 +36,16 @@ TOKEN_MERGE.name: TOKEN_MERGE, } +TEXT_ENCODER = Components(name="text_encoder", patterns=("text_encoder", "text_model")) +UNET = Components(name="unet", patterns=("unet",)) +VAE_DECODER = Components(name="vae_decoder", patterns=("vae", "decoder")) + +T2I_COMPONENTS = { + TEXT_ENCODER.name: TEXT_ENCODER, + UNET.name: UNET, + VAE_DECODER.name: VAE_DECODER, +} + MULTIMODAL_COMPONENTS = { **TEXT_COMPONENTS, VISION_ENCODER.name: VISION_ENCODER, diff --git a/python/cactus/transpiler/ModelProfiles/models.py b/python/cactus/transpiler/ModelProfiles/models.py index 54586e306..198983e88 100644 --- a/python/cactus/transpiler/ModelProfiles/models.py +++ b/python/cactus/transpiler/ModelProfiles/models.py @@ -127,6 +127,8 @@ class ModelProfile: input_strategy: str #Specifies what functions/procedures to use when generating sample input for torch.export export_patches: tuple[str, ...] #Specifies what patches/masks to apply to sample inputs load_strategy: str #Specifies which Hugging Face AutoModel class/loading path to prefer + export_modes: tuple[str, ...] = () #Per-component export modes when the family is not prefill/decode shaped; empty means the default modes + component_sources: tuple[tuple[str, str], ...] = () #Mode -> "kind:source" specs for families whose components load separately (diffusion pipelines) disabled_fusion_fields: tuple[str, ...] = () #Fusion groups to skip for this profile when a family is known unsafe disabled_fusions: tuple[str, ...] = () #Individual fusion names or cactus ops to skip for this profile prompt_contract: PromptContract = field(default_factory=PromptContract) diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index 3e047f4f3..b770f3848 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -4,6 +4,7 @@ MULTIMODAL_COMPONENTS, SPEECH_SEQ2SEQ_COMPONENTS, SPEECH_TRANSCRIBER_COMPONENTS, + T2I_COMPONENTS, TEXT_COMPONENTS, VISION_LANGUAGE_COMPONENTS, ) @@ -27,6 +28,7 @@ LFM_VLM_INFERENCE_PATTERNS, PARAKEET_INFERENCE_PATTERNS, QWEN2_5_INFERENCE_PATTERNS, + SD15_T2I_INFERENCE_PATTERNS, WHISPER_INFERENCE_PATTERNS, ) @@ -566,6 +568,56 @@ def decoder_kv_state(producer: str = "decoder_prefill_chunk") -> StateContract: runtime_contract=WHISPER_RUNTIME_CONTRACT, ) +SD15_T2I_RUNTIME_CONTRACT = RuntimeContract( + plan_name="sd15_text_to_image", + execution_strategy="iterative_denoise", + state_owner="request", + cache_persistence="none", + cache_transfer_policy="none", + states=( + StateContract( + name="text_embeddings", kind="activation", producer="text_encoder", + consumers=("unet",), lifetime="request", transfer="copy_or_alias", + metadata=(("outputs", "text_embeddings"),), + ), + StateContract( + name="latents", kind="activation", producer="unet", + consumers=("unet", "vae_decoder"), lifetime="request", transfer="move", + metadata=(("outputs", "latents"),), + ), + ), +) + +LCM_DREAMSHAPER_V7_PROFILE = ModelProfile( + model_profiles="sd15_t2i", + components=T2I_COMPONENTS, + inference_type=SD15_T2I_INFERENCE_PATTERNS, + cache_type=(), + cache_policy=(), + files=( + "model_index.json", + "text_encoder/config.json", + "unet/config.json", + "scheduler/scheduler_config.json", + "tokenizer/tokenizer_config.json", + "tokenizer/vocab.json", + "tokenizer/merges.txt", + ), + fusion_fields=("generic", "conv", "vision", "normalization", "attention", "linear", "embedding"), + supported_modalties=("text",), + input_strategy="diffusion_components", + export_patches=("clip_position_ids",), + load_strategy="diffusion_components", + export_modes=("text_encoder", "unet", "vae_decoder"), + component_sources=( + ("text_encoder", "clip_text:text_encoder"), + ("unet", "sd_unet:unet"), + ("vae_decoder", "taesd_decoder:madebyollin/taesd"), + ), + prompt_contract=PromptContract(style="raw", template_source="none"), + runtime_contract=SD15_T2I_RUNTIME_CONTRACT, +) + MODEL_ID_MAP = { "google/gemma-4-E2B": GEMMA4_E2B_PROFILE, "google/gemma-4-E2B-it": GEMMA4_E2B_IT_PROFILE, @@ -576,6 +628,7 @@ def decoder_kv_state(producer: str = "decoder_prefill_chunk") -> StateContract: "LiquidAI/LFM2-VL-3B": LFM_VLM_PROFILE, "Qwen/Qwen2.5-0.5B": QWEN2_5_0_5B_PROFILE, "LiquidAI/LFM2.5-8B-A1B": LFM_MOE_PROFILE, + "SimianLuo/LCM_Dreamshaper_v7": LCM_DREAMSHAPER_V7_PROFILE, } def profile_for_model_id(model_id: str) -> ModelProfile | None: diff --git a/python/cactus/transpiler/ModelProfiles/routes.py b/python/cactus/transpiler/ModelProfiles/routes.py index f91f6cc10..21a9e4ab1 100644 --- a/python/cactus/transpiler/ModelProfiles/routes.py +++ b/python/cactus/transpiler/ModelProfiles/routes.py @@ -45,3 +45,12 @@ TEXT_PREFILL.name: TEXT_PREFILL, TEXT_DECODE.name: TEXT_DECODE, } + +# The denoise loop is decode-shaped: the unet edge runs once per scheduler +# step with the latent as persistent state; encode and decode run once. +T2I_GENERATE_ROUTE = (combinations.TEXT_ENCODER_TO_UNET, combinations.UNET_TO_VAE_DECODER) +T2I_GENERATE = InferencePattern(name="t2i_generate", route=T2I_GENERATE_ROUTE) + +SD15_T2I_INFERENCE_PATTERNS = { + T2I_GENERATE.name: T2I_GENERATE, +} diff --git a/python/tests/test_transpiler_sd_profile.py b/python/tests/test_transpiler_sd_profile.py new file mode 100644 index 000000000..713ceeefa --- /dev/null +++ b/python/tests/test_transpiler_sd_profile.py @@ -0,0 +1,112 @@ +import unittest +from pathlib import Path + +import torch + +from cactus.cli import transpiler as cli_transpiler +from cactus.transpiler.Converter import diffusion +from cactus.transpiler.ModelProfiles import profiles +from cactus.transpiler.RuntimePlan import models as RPModels + +MODEL_ID = "SimianLuo/LCM_Dreamshaper_v7" + + +class TestSdProfile(unittest.TestCase): + def test_model_id_resolves_to_registered_profile_case_insensitively(self): + profile = profiles.profile_for_model_id(MODEL_ID) + self.assertIsNotNone(profile) + self.assertEqual(profile.model_profiles, "sd15_t2i") + self.assertIs(profiles.profile_for_model_id(MODEL_ID.lower()), profile) + + def test_resolved_config_exports_one_mode_per_component(self): + resolved = cli_transpiler.resolve_transpile_config(MODEL_ID) + self.assertEqual(resolved.profile_source, "registered") + self.assertEqual(resolved.inference_modes, ("text_encoder", "unet", "vae_decoder")) + + def test_registered_profile_rejects_generic_flags(self): + with self.assertRaises(RuntimeError): + cli_transpiler.resolve_transpile_config(MODEL_ID, generic_task="causal-lm") + + def test_component_sources_parse_kind_and_source(self): + profile = profiles.profile_for_model_id(MODEL_ID) + self.assertEqual(diffusion.component_source(profile, "unet"), ("sd_unet", "unet")) + self.assertEqual( + diffusion.component_source(profile, "vae_decoder"), + ("taesd_decoder", "madebyollin/taesd"), + ) + with self.assertRaises(ValueError): + diffusion.component_source(profile, "decode_with_cache") + + def test_component_export_inputs_follow_the_configs(self): + configs = { + "unet/config.json": { + "in_channels": 4, "sample_size": 64, + "cross_attention_dim": 768, "time_cond_proj_dim": 256, + }, + "text_encoder/config.json": {"max_position_embeddings": 77}, + } + ids = diffusion.component_export_inputs("clip_text", configs) + self.assertEqual(ids["input_ids"].shape, (1, 77)) + self.assertEqual(ids["input_ids"].dtype, torch.int64) + + unet = diffusion.component_export_inputs("sd_unet", configs) + self.assertEqual(unet["sample"].shape, (1, 4, 64, 64)) + self.assertEqual(unet["timestep"].shape, (1,)) + self.assertEqual(unet["encoder_hidden_states"].shape, (1, 77, 768)) + self.assertEqual(unet["timestep_cond"].shape, (1, 256)) + self.assertTrue(all(t.dtype == torch.float16 for t in unet.values())) + + latent = diffusion.component_export_inputs("taesd_decoder", configs) + self.assertEqual(latent["x"].shape, (1, 4, 64, 64)) + + with self.assertRaises(ValueError): + diffusion.component_export_inputs("vqgan", configs) + + def test_runtime_plan_declares_the_t2i_route_and_denoise_strategy(self): + profile = profiles.profile_for_model_id(MODEL_ID) + routes = RPModels.runtime_routes_from_model_profile(profile) + by_name = {route.name: route for route in routes} + self.assertIn("t2i_generate", by_name) + edges = by_name["t2i_generate"].edges + self.assertEqual([edge.output for edge in edges], ["unet", "vae_decoder"]) + self.assertEqual(edges[0].inputs, ("text_encoder",)) + + metadata = RPModels.runtime_plan_metadata_from_model_profile(profile) + self.assertEqual(metadata.get("runtime_execution_strategy"), "iterative_denoise") + self.assertEqual(metadata.get("runtime_plan_name"), "sd15_text_to_image") + + def test_full_bundle_builds_from_the_cached_model(self): + import tempfile + + try: + import diffusers # noqa: F401 + except Exception: + self.skipTest("diffusers unavailable") + try: + from transformers import CLIPTextModel + + CLIPTextModel.from_pretrained(MODEL_ID, subfolder="text_encoder", local_files_only=True) + except Exception as exc: + self.skipTest(f"{MODEL_ID} not cached: {exc}") + + from cactus.cli.model import ensure_weights + + with tempfile.TemporaryDirectory() as tmp: + weights_dir = ensure_weights(MODEL_ID, output_dir=Path(tmp) / "weights") + self.assertTrue((weights_dir / "weights_manifest.json").exists()) + + bundle = cli_transpiler.build_transpiled_bundle( + MODEL_ID, weights_dir=weights_dir, output_dir=Path(tmp) / "bundle" + ) + components = {path.stem for path in (bundle / "components").glob("*.cactus")} + self.assertEqual(components, {"text_encoder", "unet", "vae_decoder"}) + self.assertTrue((bundle / "runtime_plan.json").exists()) + + import json + + plan = json.loads((bundle / "runtime_plan.json").read_text()) + self.assertEqual(plan["metadata"].get("runtime_execution_strategy"), "iterative_denoise") + + +if __name__ == "__main__": + unittest.main() From 6ed7fcb3c0702133a5d219f8e26d25502a08236b Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sun, 9 Aug 2026 16:27:27 -0700 Subject: [PATCH 25/34] materialize diffusion runtime assets and fix weight-name aliasing 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 --- python/cactus/cli/model.py | 47 +++++++++++++++++++- python/cactus/cli/transpiler.py | 38 ++++++++++++---- python/cactus/transpiler/Generator/models.py | 27 +++++++---- 3 files changed, 94 insertions(+), 18 deletions(-) diff --git a/python/cactus/cli/model.py b/python/cactus/cli/model.py index dd1af4c5f..b9fcece52 100644 --- a/python/cactus/cli/model.py +++ b/python/cactus/cli/model.py @@ -114,13 +114,58 @@ def _convert_component_sources(model_id, component_sources, *, bits, token, weig (weights_dir / "weights_manifest.json").write_text( json.dumps({"weights": merged_records}, indent=2), encoding="utf-8" ) + _materialize_diffusion_assets(model_id, weights_dir, token=token) (weights_dir / "config.txt").write_text( - "\n".join(f"component={mode}" for mode in converted) + "\n", encoding="utf-8" + "model_type=sd15\n" + "\n".join(f"component={mode}" for mode in converted) + "\n", encoding="utf-8" ) print_color(GREEN, f"Model converted and ready at {weights_dir}") return weights_dir +def _materialize_diffusion_assets(model_id, weights_dir, *, token=None): + """Emit the engine-format tokenizer sidecars and the scheduler config the + diffusion runtime needs alongside the converted weights.""" + import json + + from huggingface_hub import snapshot_download + + snapshot = Path(snapshot_download(model_id, allow_patterns=["tokenizer/*", "scheduler/*"], token=token)) + vocab = json.loads((snapshot / "tokenizer" / "vocab.json").read_text(encoding="utf-8")) + with open(weights_dir / "vocab.txt", "w", encoding="utf-8") as f: + for token_str, token_id in sorted(vocab.items(), key=lambda kv: kv[1]): + f.write(f"{token_id}\t{token_str}\n") + shutil.copy2(snapshot / "tokenizer" / "merges.txt", weights_dir / "merges.txt") + (weights_dir / "tokenizer_config.txt").write_text( + "vocab_size=49408\n" + "bos_token_id=49406\n" + "eos_token_id=49407\n" + "pad_token_id=49407\n" + "model_max_length=77\n" + "tokenizer_type=bpe\n" + "vocab_format=id_tab_token\n" + "normalizer=clip\n" + "decoder=clip\n" + "byte_fallback=false\n" + "has_chat_template=false\n", + encoding="utf-8", + ) + (weights_dir / "special_tokens.json").write_text( + json.dumps( + { + "bos_token_id": 49406, + "eos_token_id": 49407, + "pad_token_id": 49407, + "vocab_size": 49408, + "model_max_length": 77, + "special_tokens": {"49406": "<|startoftext|>", "49407": "<|endoftext|>"}, + }, + indent=2, + ), + encoding="utf-8", + ) + shutil.copy2(snapshot / "scheduler" / "scheduler_config.json", weights_dir / "scheduler_config.json") + + def ensure_weights(model_id, *, bits=4, token=None, reconvert=False, output_dir=None, skip_model_load=False): from .download import get_bundle_dir diff --git a/python/cactus/cli/transpiler.py b/python/cactus/cli/transpiler.py index 9410e2ac2..b581b442e 100644 --- a/python/cactus/cli/transpiler.py +++ b/python/cactus/cli/transpiler.py @@ -139,20 +139,23 @@ def build_transpiled_bundle( ) print_color(CYAN, f"\nRunning Generator (Caesar v{generator_version})") + bundle_metadata = { + "model_id": str(profile_model_id or model_id), + "source_model_id": str(model_id), + "input_modalities": ",".join(modalities), + "profile_source": profile_source, + "generic_task": resolved.generic_task, + "cache_style": resolved.cache_style, + "fusion_groups": ",".join(profile.fusion_fields), + } + bundle_metadata.update(diffusion_scheduler_metadata(weights_path)) + result = generate_bundle( simplified_maps, bundle_path, model_profile=profile, weights_dir=weights_path, - metadata={ - "model_id": str(profile_model_id or model_id), - "source_model_id": str(model_id), - "input_modalities": ",".join(modalities), - "profile_source": profile_source, - "generic_task": resolved.generic_task, - "cache_style": resolved.cache_style, - "fusion_groups": ",".join(profile.fusion_fields), - }, + metadata=bundle_metadata, strict=strict, allow_unsupported_ops=allow_unsupported_ops, ) @@ -165,6 +168,23 @@ def build_transpiled_bundle( return bundle_path +def diffusion_scheduler_metadata(weights_dir: Path) -> dict[str, str]: + """Flatten the diffusers scheduler config into engine-manifest metadata keys.""" + scheduler_path = weights_dir / "scheduler_config.json" + if not scheduler_path.exists(): + return {} + scheduler = json.loads(scheduler_path.read_text(encoding="utf-8")) + return { + "diffusion_beta_start": str(scheduler.get("beta_start", 0.00085)), + "diffusion_beta_end": str(scheduler.get("beta_end", 0.012)), + "diffusion_beta_schedule": str(scheduler.get("beta_schedule", "scaled_linear")), + "diffusion_num_train_timesteps": str(scheduler.get("num_train_timesteps", 1000)), + "diffusion_original_inference_steps": str(scheduler.get("original_inference_steps", 50)), + "diffusion_timestep_scaling": str(scheduler.get("timestep_scaling", 10.0)), + "diffusion_prediction_type": str(scheduler.get("prediction_type", "epsilon")), + } + + def profile_for_model_id(model_id: str): from cactus.transpiler.ModelProfiles import profiles diff --git a/python/cactus/transpiler/Generator/models.py b/python/cactus/transpiler/Generator/models.py index d0e1741f1..4947851ea 100644 --- a/python/cactus/transpiler/Generator/models.py +++ b/python/cactus/transpiler/Generator/models.py @@ -434,20 +434,31 @@ def resolve_weight_record(resolver: WeightResolver, placeholder_name: str) -> We return resolve_source_weight_record(resolver, source_target) def resolve_source_weight_record(resolver: WeightResolver, source_target: str) -> WeightRecord | None: + candidates = list(exact_weight_name_variants(source_target)) + head, sep, rest = source_target.partition(".") + if sep and rest: + candidates.extend(exact_weight_name_variants(rest)) + for variant in candidates: + record = resolver.records_by_name.get(variant) + if record is not None: + return record for variant in weight_name_variants(source_target): record = resolver.records_by_name.get(variant) if record is not None: return record - # Export wrappers add one module attribute in front of the checkpoint key - # (e.g. unet.down_blocks...). Only consulted when every exact variant missed. - head, sep, rest = source_target.partition(".") - if sep and rest: - for variant in weight_name_variants(rest): - record = resolver.records_by_name.get(variant) - if record is not None: - return record return None + +def exact_weight_name_variants(name: str) -> tuple[str, ...]: + variants = [name] + current = name + while current.startswith("model."): + current = current[len("model."):] + variants.append(current) + if name.startswith("_orig_mod."): + variants.append(name[len("_orig_mod."):]) + return tuple(unique_strings(variants)) + def weight_name_variants(name: str) -> tuple[str, ...]: variants = [name] current = name From ca410ab99b660dbc54f5fc0db1256185ff10b513 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sun, 9 Aug 2026 16:27:27 -0700 Subject: [PATCH 26/34] run text-to-image bundles in the engine New ITERATIVE_DENOISE route for text_encoder+unet+vae_decoder bundles: Model::generate_image tokenizes with a new CLIP BPE mode (lowercase, word split, 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 --- cactus-engine/CMakeLists.txt | 1 + cactus-engine/cactus_engine.h | 12 ++ cactus-engine/src/bpe.cpp | 84 ++++++++++ cactus-engine/src/engine.h | 24 ++- cactus-engine/src/generate_image.cpp | 51 ++++++ cactus-engine/src/model.cpp | 216 +++++++++++++++++++++++++ cactus-engine/src/tokenizer.cpp | 2 + cactus-engine/tests/test_diffusion.cpp | 86 ++++++++++ 8 files changed, 472 insertions(+), 4 deletions(-) create mode 100644 cactus-engine/src/generate_image.cpp create mode 100644 cactus-engine/tests/test_diffusion.cpp diff --git a/cactus-engine/CMakeLists.txt b/cactus-engine/CMakeLists.txt index bba63c225..5ad67a2ab 100644 --- a/cactus-engine/CMakeLists.txt +++ b/cactus-engine/CMakeLists.txt @@ -27,6 +27,7 @@ set(ENGINE_SOURCES src/cloud.cpp src/init.cpp src/embed.cpp + src/generate_image.cpp src/complete.cpp src/transcribe.cpp src/stream.cpp diff --git a/cactus-engine/cactus_engine.h b/cactus-engine/cactus_engine.h index a60e4aec2..29f951919 100644 --- a/cactus-engine/cactus_engine.h +++ b/cactus-engine/cactus_engine.h @@ -140,6 +140,18 @@ CACTUS_FFI_EXPORT int cactus_stream_transcribe_stop( size_t buffer_size ); +CACTUS_FFI_EXPORT int cactus_generate_image( + cactus_model_t model, + const char* prompt, + uint8_t* rgb_buffer, + size_t buffer_size, + unsigned int* image_width, + unsigned int* image_height, + int steps, + float guidance_scale, + unsigned long long seed +); + CACTUS_FFI_EXPORT int cactus_embed( cactus_model_t model, const char* text, diff --git a/cactus-engine/src/bpe.cpp b/cactus-engine/src/bpe.cpp index 41d49c795..54e5ed4c3 100644 --- a/cactus-engine/src/bpe.cpp +++ b/cactus-engine/src/bpe.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -455,6 +456,73 @@ std::vector BPETokenizer::apply_bpe(const std::vector& return current_tokens; } +std::vector BPETokenizer::clip_word_split(const std::string& text) const { + // ASCII matches CLIP's regex exactly; non-ASCII UTF-8 is treated as letters. + static const char* kContractions[] = {"'s", "'t", "'re", "'ve", "'m", "'ll", "'d"}; + auto is_letter = [](unsigned char c) { return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || c >= 0x80; }; + auto is_digit = [](unsigned char c) { return c >= '0' && c <= '9'; }; + auto is_space = [](unsigned char c) { return c == ' ' || c == '\t' || c == '\n' || c == '\r' || c == '\f' || c == '\v'; }; + + std::vector words; + size_t i = 0; + while (i < text.size()) { + const unsigned char c = static_cast(text[i]); + if (is_space(c)) { + ++i; + continue; + } + if (c == '\'') { + bool matched = false; + for (const char* contraction : kContractions) { + const size_t len = std::strlen(contraction); + if (text.compare(i, len, contraction) == 0) { + words.emplace_back(text.substr(i, len)); + i += len; + matched = true; + break; + } + } + if (matched) continue; + } + if (is_letter(c)) { + size_t start = i; + while (i < text.size() && is_letter(static_cast(text[i]))) ++i; + words.emplace_back(text.substr(start, i - start)); + continue; + } + if (is_digit(c)) { + words.emplace_back(text.substr(i, 1)); + ++i; + continue; + } + size_t start = i; + while (i < text.size()) { + const unsigned char p = static_cast(text[i]); + if (is_space(p) || is_letter(p) || is_digit(p) || p == '\'') break; + ++i; + } + if (i == start) ++i; + else words.emplace_back(text.substr(start, i - start)); + } + return words; +} + +void BPETokenizer::encode_clip_segment(const std::string& segment, std::vector& token_ids) const { + std::string lowered = segment; + for (char& c : lowered) { + if (c >= 'A' && c <= 'Z') c += 'a' - 'A'; + } + for (const auto& word : clip_word_split(lowered)) { + std::vector chars = utf8_split(bytes_to_unicode(word)); + if (chars.empty()) continue; + chars.back() += ""; + for (const auto& token : apply_bpe(chars)) { + auto it = token_to_id_.find(token); + token_ids.push_back(it != token_to_id_.end() ? it->second : unk_token_id_); + } + } +} + std::vector BPETokenizer::encode(const std::string& text) const { if (text.empty()) return {}; @@ -468,6 +536,8 @@ std::vector BPETokenizer::encode(const std::string& text) const { auto special_it = special_tokens_.find(segment); if (special_it != special_tokens_.end()) { token_ids.push_back(special_it->second); + } else if (runtime_config_.normalizer == TokenizerRuntimeConfig::Normalizer::CLIP) { + encode_clip_segment(segment, token_ids); } else { std::string normalized_segment = segment; std::vector chars; @@ -520,6 +590,20 @@ std::vector BPETokenizer::encode(const std::string& text) const { } std::string BPETokenizer::decode(const std::vector& tokens) const { + if (runtime_config_.decoder == TokenizerRuntimeConfig::Decoder::CLIP) { + std::string unicode_result; + unicode_result.reserve(tokens.size() * 6); + for (uint32_t token_id : tokens) { + if (token_id >= id_to_token_.size()) continue; + std::string tok = id_to_token_[token_id]; + size_t pos = tok.find(""); + if (pos != std::string::npos) tok = tok.substr(0, pos) + " "; + unicode_result += tok; + } + std::string result = unicode_to_bytes(unicode_result); + while (!result.empty() && result.back() == ' ') result.pop_back(); + return result; + } if (runtime_config_.decoder == TokenizerRuntimeConfig::Decoder::REPLACE_METASPACE) { std::string result; result.reserve(tokens.size() * 4); diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 0ee887c3d..3be1cb78b 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -149,7 +149,7 @@ struct Config { float rope_scaling_factor = 1.0f; float rope_mscale_all_dim = 0.0f; - enum class ModelType {QWEN = 0, GEMMA = 1, NOMIC = 3, LFM2 = 5, SIGLIP2 = 6, WHISPER = 7, MOONSHINE = 8, PARAKEET = 10, QWEN3P5 = 11, PARAKEET_TDT = 12, GEMMA3N = 13, YOUTU = 14, GEMMA4 = 15, NEEDLE = 18, GENERIC = 19}; + enum class ModelType {QWEN = 0, GEMMA = 1, NOMIC = 3, LFM2 = 5, SIGLIP2 = 6, WHISPER = 7, MOONSHINE = 8, PARAKEET = 10, QWEN3P5 = 11, PARAKEET_TDT = 12, GEMMA3N = 13, YOUTU = 14, GEMMA4 = 15, NEEDLE = 18, GENERIC = 19, SD15 = 20}; uint32_t predictor_hidden_dim = 0; uint32_t predictor_num_layers = 0; uint32_t tdt_joint_dim = 0; @@ -287,8 +287,8 @@ struct ToolConstraintSpec { struct TokenizerRuntimeConfig { enum class TokenizerType { UNKNOWN, BPE, SENTENCEPIECE }; enum class VocabFormat { UNKNOWN, ID_TAB_TOKEN, LINE_TOKEN }; - enum class Normalizer { NONE, METASPACE, BYTE_LEVEL }; - enum class Decoder { NONE, REPLACE_METASPACE, BYTE_LEVEL }; + enum class Normalizer { NONE, METASPACE, BYTE_LEVEL, CLIP }; + enum class Decoder { NONE, REPLACE_METASPACE, BYTE_LEVEL, CLIP }; TokenizerType tokenizer_type = TokenizerType::UNKNOWN; VocabFormat vocab_format = VocabFormat::UNKNOWN; @@ -485,6 +485,8 @@ class BPETokenizer : public Tokenizer { std::string unicode_to_bytes(const std::string& text) const; std::vector byte_level_split(const std::string& text) const; std::vector utf8_split(const std::string& text) const; + std::vector clip_word_split(const std::string& text) const; + void encode_clip_segment(const std::string& segment, std::vector& token_ids) const; void cleanup_mmap(); @@ -686,6 +688,10 @@ class Model { uint32_t decode(const std::vector& tokens, float temperature = -1.0f, float top_p = -1.0f, size_t top_k = 0, const std::string& profile_file = "", float* out_entropy = nullptr, float min_p = 0.15f, float repetition_penalty = 1.1f); + + int generate_image(const std::string& prompt, uint8_t* out_rgb, size_t out_capacity, + uint32_t* out_width, uint32_t* out_height, + int steps, float guidance_scale, uint64_t seed); void set_sample_seed(uint64_t seed) { sample_rng_.seed(static_cast(seed)); } bool prefill_and_sample_first_token(const std::vector& tokens, uint32_t& out_token, float* out_uncertainty = nullptr); @@ -1037,8 +1043,18 @@ class Model { Component* decoder_prefill_ = nullptr; Component* decoder_embed_ = nullptr; Component* prefill_encoder_ = nullptr; - enum class DecodeRoute { CACHED_STEP, DIRECT_DECODER_STEP, FULL_CONTEXT_TEXT, ENCODER_CROSS_KV_STEP }; + enum class DecodeRoute { CACHED_STEP, DIRECT_DECODER_STEP, FULL_CONTEXT_TEXT, ENCODER_CROSS_KV_STEP, ITERATIVE_DENOISE }; DecodeRoute decode_route_ = DecodeRoute::CACHED_STEP; + struct DiffusionParams { + float beta_start = 0.00085f; + float beta_end = 0.012f; + std::string beta_schedule = "scaled_linear"; + uint32_t num_train_timesteps = 1000; + uint32_t original_inference_steps = 50; + float timestep_scaling = 10.0f; + std::string prediction_type = "epsilon"; + }; + DiffusionParams diffusion_; Component* source_encoder_ = nullptr; Component* decoder_cross_kv_ = nullptr; Component* vision_encoder_ = nullptr; diff --git a/cactus-engine/src/generate_image.cpp b/cactus-engine/src/generate_image.cpp new file mode 100644 index 000000000..c64b7ef40 --- /dev/null +++ b/cactus-engine/src/generate_image.cpp @@ -0,0 +1,51 @@ +#include "../cactus_engine.h" +#include "utils.h" +#include +#include + +using namespace cactus::engine; +using namespace cactus::ffi; + +extern "C" { + +int cactus_generate_image( + cactus_model_t model, + const char* prompt, + uint8_t* rgb_buffer, + size_t buffer_size, + unsigned int* image_width, + unsigned int* image_height, + int steps, + float guidance_scale, + unsigned long long seed +) { + if (!model || !prompt || !rgb_buffer || buffer_size == 0) { + CACTUS_LOG_ERROR("generate_image", "Invalid parameters for image generation"); + return -1; + } + + try { + auto* handle = static_cast(model); + std::lock_guard lock(handle->model_mutex); + uint32_t width = 0; + uint32_t height = 0; + const int result = handle->model->generate_image( + prompt, rgb_buffer, buffer_size, &width, &height, + steps, guidance_scale, static_cast(seed)); + if (result > 0) { + if (image_width) *image_width = width; + if (image_height) *image_height = height; + } + return result; + } catch (const std::exception& e) { + last_error_message = e.what(); + CACTUS_LOG_ERROR("generate_image", "Exception: " << e.what()); + return -1; + } catch (...) { + last_error_message = "Unknown error during image generation"; + CACTUS_LOG_ERROR("generate_image", last_error_message); + return -1; + } +} + +} diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index d7f1213c5..19f10a86a 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -766,6 +767,9 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, if (has_metadata_encoder_cross_kv_route) { decode_route_ = DecodeRoute::ENCODER_CROSS_KV_STEP; required_components = {decoder_name}; + } else if (components_.count("text_encoder") && components_.count("unet") && components_.count("vae_decoder")) { + decode_route_ = DecodeRoute::ITERATIVE_DENOISE; + required_components = {"text_encoder", "unet", "vae_decoder"}; } else if (has_chunked_prefill) { encoder_name = "lm_encoder_step"; decoder_name = "decoder_media_step"; @@ -815,6 +819,7 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, : decoder_prefill_chunk_; } else if (decode_route_ != DecodeRoute::ENCODER_CROSS_KV_STEP && decode_route_ != DecodeRoute::FULL_CONTEXT_TEXT + && decode_route_ != DecodeRoute::ITERATIVE_DENOISE && !components_.count("audio_encoder")) { CACTUS_LOG_WARN("model", "Bundle has no decoder_prefill_chunk component; prompts will prefill token-by-token (prefill speed ~= decode speed)."); } @@ -1020,6 +1025,21 @@ bool Model::load_manifest() { image_feature_candidates_ = parse_csv_strings(read_string_key("media_image_feature_names")); audio_feature_candidates_ = parse_csv_strings(read_string_key("media_audio_feature_names")); + const std::string beta_start = read_string_key("diffusion_beta_start"); + if (!beta_start.empty()) diffusion_.beta_start = std::stof(beta_start); + const std::string beta_end = read_string_key("diffusion_beta_end"); + if (!beta_end.empty()) diffusion_.beta_end = std::stof(beta_end); + const std::string beta_schedule = read_string_key("diffusion_beta_schedule"); + if (!beta_schedule.empty()) diffusion_.beta_schedule = beta_schedule; + diffusion_.num_train_timesteps = parse_uint32_or_zero(read_string_key("diffusion_num_train_timesteps")); + if (diffusion_.num_train_timesteps == 0) diffusion_.num_train_timesteps = 1000; + diffusion_.original_inference_steps = parse_uint32_or_zero(read_string_key("diffusion_original_inference_steps")); + if (diffusion_.original_inference_steps == 0) diffusion_.original_inference_steps = 50; + const std::string timestep_scaling = read_string_key("diffusion_timestep_scaling"); + if (!timestep_scaling.empty()) diffusion_.timestep_scaling = std::stof(timestep_scaling); + const std::string prediction_type = read_string_key("diffusion_prediction_type"); + if (!prediction_type.empty()) diffusion_.prediction_type = prediction_type; + if (obj.count("npu_audio_encoder") && obj.at("npu_audio_encoder").is()) { npu_audio_encoder_mlpackage_ = obj.at("npu_audio_encoder").get(); } @@ -5449,6 +5469,201 @@ std::vector Model::get_embeddings(const std::vector& tokens, bo return result; } +namespace { + +std::vector diffusion_alphas_cumprod(float beta_start, float beta_end, uint32_t num_train_timesteps) { + std::vector alphas_cumprod(num_train_timesteps); + const double sqrt_start = std::sqrt(static_cast(beta_start)); + const double sqrt_end = std::sqrt(static_cast(beta_end)); + double cumprod = 1.0; + for (uint32_t i = 0; i < num_train_timesteps; ++i) { + const double frac = num_train_timesteps > 1 ? static_cast(i) / (num_train_timesteps - 1) : 0.0; + const double sqrt_beta = sqrt_start + (sqrt_end - sqrt_start) * frac; + cumprod *= 1.0 - sqrt_beta * sqrt_beta; + alphas_cumprod[i] = static_cast(cumprod); + } + return alphas_cumprod; +} + +std::vector diffusion_lcm_timesteps(uint32_t num_train_timesteps, uint32_t original_steps, uint32_t steps) { + const uint32_t k = num_train_timesteps / original_steps; + std::vector timesteps(steps); + for (uint32_t s = 0; s < steps; ++s) { + const uint32_t index = static_cast( + std::floor(static_cast(s) * original_steps / steps)); + timesteps[s] = (original_steps - index) * k - 1; + } + return timesteps; +} + +std::vector diffusion_guidance_embedding(float w, size_t dim) { + const size_t half = dim / 2; + std::vector embedding(dim, 0.0f); + const double scaled = static_cast(w) * 1000.0; + for (size_t i = 0; i < half; ++i) { + const double freq = std::exp(-std::log(10000.0) * static_cast(i) / (half - 1)); + embedding[i] = static_cast(std::sin(scaled * freq)); + embedding[half + i] = static_cast(std::cos(scaled * freq)); + } + return embedding; +} + +} // namespace + +int Model::generate_image(const std::string& prompt, uint8_t* out_rgb, size_t out_capacity, + uint32_t* out_width, uint32_t* out_height, + int steps, float guidance_scale, uint64_t seed) { + if (decode_route_ != DecodeRoute::ITERATIVE_DENOISE) { + CACTUS_LOG_ERROR("model", "generate_image requires a text-to-image bundle"); + return -1; + } + if (steps < 1 || steps > static_cast(diffusion_.original_inference_steps)) { + CACTUS_LOG_ERROR("model", "generate_image steps must be in [1, " << diffusion_.original_inference_steps << "]"); + return -1; + } + Component& text_encoder = components_.at("text_encoder"); + Component& unet = components_.at("unet"); + Component& vae_decoder = components_.at("vae_decoder"); + if (!text_encoder.graph || !unet.graph || !vae_decoder.graph || !tokenizer_) return -1; + + const int ids_idx = input_index(text_encoder, "input_ids"); + const int sample_idx = input_index(unet, "sample"); + const int timestep_idx = input_index(unet, "timestep"); + const int hidden_idx = input_index(unet, "encoder_hidden_states"); + const int cond_idx = input_index(unet, "timestep_cond"); + const int latent_idx = input_index(vae_decoder, "x"); + if (ids_idx < 0 || sample_idx < 0 || timestep_idx < 0 || hidden_idx < 0 || cond_idx < 0 || latent_idx < 0) { + CACTUS_LOG_ERROR("model", "text-to-image bundle is missing expected component inputs"); + return -1; + } + + const auto& ids_desc = text_encoder.graph->get_output_buffer( + static_cast(text_encoder.runtime_input_node_ids[ids_idx])); + const size_t max_tokens = ids_desc.total_size; + const uint32_t bos = tokenizer_->get_bos_token(); + const uint32_t eos = tokenizer_->get_eos_token(); + auto tokens = tokenizer_->encode(prompt); + if (tokens.size() > max_tokens - 2) tokens.resize(max_tokens - 2); + write_int_input_at(text_encoder, "input_ids", 0, static_cast(bos)); + for (size_t i = 0; i < tokens.size(); ++i) { + write_int_input_at(text_encoder, "input_ids", i + 1, static_cast(tokens[i])); + } + for (size_t i = tokens.size() + 1; i < max_tokens; ++i) { + write_int_input_at(text_encoder, "input_ids", i, static_cast(eos)); + } + text_encoder.graph->execute(); + + const size_t hidden_node = static_cast(text_encoder.output_node_ids[0]); + const auto& hidden_desc = text_encoder.graph->get_output_buffer(hidden_node); + write_typed_buffer( + unet.input_buffers[hidden_idx], + unet.graph->get_output_buffer(static_cast(unet.runtime_input_node_ids[hidden_idx])).precision, + text_encoder.graph->get_output(hidden_node), + hidden_desc.byte_size, + hidden_desc.precision); + + const auto& cond_desc = unet.graph->get_output_buffer( + static_cast(unet.runtime_input_node_ids[cond_idx])); + const std::vector guidance = diffusion_guidance_embedding(guidance_scale - 1.0f, cond_desc.total_size); + write_typed_buffer(unet.input_buffers[cond_idx], cond_desc.precision, + guidance.data(), guidance.size() * sizeof(float), Precision::FP32); + + const auto& sample_desc = unet.graph->get_output_buffer( + static_cast(unet.runtime_input_node_ids[sample_idx])); + const size_t latent_count = sample_desc.total_size; + std::mt19937_64 rng(seed); + std::normal_distribution normal(0.0f, 1.0f); + std::vector latents(latent_count); + for (float& value : latents) value = normal(rng); + + const std::vector alphas_cumprod = + diffusion_alphas_cumprod(diffusion_.beta_start, diffusion_.beta_end, diffusion_.num_train_timesteps); + const std::vector timesteps = diffusion_lcm_timesteps( + diffusion_.num_train_timesteps, diffusion_.original_inference_steps, static_cast(steps)); + + const auto& timestep_desc = unet.graph->get_output_buffer( + static_cast(unet.runtime_input_node_ids[timestep_idx])); + std::vector eps(latent_count); + for (int s = 0; s < steps; ++s) { + const uint32_t t = timesteps[s]; + write_typed_buffer(unet.input_buffers[sample_idx], sample_desc.precision, + latents.data(), latents.size() * sizeof(float), Precision::FP32); + const float timestep_value = static_cast(t); + write_typed_buffer(unet.input_buffers[timestep_idx], timestep_desc.precision, + ×tep_value, sizeof(float), Precision::FP32); + unet.graph->execute(); + + const size_t eps_node = static_cast(unet.output_node_ids[0]); + const auto& eps_desc = unet.graph->get_output_buffer(eps_node); + if (eps_desc.precision == Precision::FP16) { + const __fp16* data = static_cast(unet.graph->get_output(eps_node)); + for (size_t i = 0; i < latent_count; ++i) eps[i] = static_cast(data[i]); + } else { + std::memcpy(eps.data(), unet.graph->get_output(eps_node), latent_count * sizeof(float)); + } + + const float acp_t = alphas_cumprod[t]; + const float sqrt_acp = std::sqrt(acp_t); + const float sqrt_one_minus = std::sqrt(1.0f - acp_t); + const float scaled_t = static_cast(t) * diffusion_.timestep_scaling; + const float c_skip = 0.25f / (scaled_t * scaled_t + 0.25f); + const float c_out = scaled_t / std::sqrt(scaled_t * scaled_t + 0.25f); + const bool final_step = s == steps - 1; + const float acp_prev = final_step ? 1.0f : alphas_cumprod[timesteps[s + 1]]; + const float sqrt_acp_prev = std::sqrt(acp_prev); + const float sqrt_one_minus_prev = std::sqrt(1.0f - acp_prev); + for (size_t i = 0; i < latent_count; ++i) { + const float x0 = (latents[i] - sqrt_one_minus * eps[i]) / sqrt_acp; + const float denoised = c_out * x0 + c_skip * latents[i]; + latents[i] = final_step ? denoised + : sqrt_acp_prev * denoised + sqrt_one_minus_prev * normal(rng); + } + } + + const auto& latent_desc = vae_decoder.graph->get_output_buffer( + static_cast(vae_decoder.runtime_input_node_ids[latent_idx])); + write_typed_buffer(vae_decoder.input_buffers[latent_idx], latent_desc.precision, + latents.data(), latents.size() * sizeof(float), Precision::FP32); + vae_decoder.graph->execute(); + + const size_t image_node = static_cast(vae_decoder.output_node_ids[0]); + const auto& image_desc = vae_decoder.graph->get_output_buffer(image_node); + if (image_desc.shape.size() != 4 || image_desc.shape[1] != 3) { + CACTUS_LOG_ERROR("model", "vae_decoder output is not an NCHW RGB image"); + return -1; + } + const size_t height = image_desc.shape[2]; + const size_t width = image_desc.shape[3]; + const size_t out_bytes = height * width * 3; + if (out_capacity < out_bytes) { + CACTUS_LOG_ERROR("model", "generate_image buffer too small: need " << out_bytes << " bytes"); + return -2; + } + + const size_t plane = height * width; + auto write_pixels = [&](auto value_at) { + for (size_t y = 0; y < height; ++y) { + for (size_t x = 0; x < width; ++x) { + const size_t pixel = y * width + x; + for (size_t c = 0; c < 3; ++c) { + const float value = std::min(1.0f, std::max(0.0f, value_at(c * plane + pixel))); + out_rgb[pixel * 3 + c] = static_cast(std::lround(value * 255.0f)); + } + } + } + }; + if (image_desc.precision == Precision::FP16) { + const __fp16* data = static_cast(vae_decoder.graph->get_output(image_node)); + write_pixels([data](size_t i) { return static_cast(data[i]); }); + } else { + const float* data = static_cast(vae_decoder.graph->get_output(image_node)); + write_pixels([data](size_t i) { return data[i]; }); + } + if (out_width) *out_width = static_cast(width); + if (out_height) *out_height = static_cast(height); + return static_cast(out_bytes); +} + bool Config::from_json(const std::string& config_path) { std::ifstream file(config_path); if (!file) { @@ -5550,6 +5765,7 @@ bool Config::from_json(const std::string& config_path) { else if (mt == "needle") model_type = ModelType::NEEDLE; else if (mt == "bert" || mt == "nomic") model_type = ModelType::NOMIC; else if (mt == "generic") model_type = ModelType::GENERIC; + else if (mt == "sd15") model_type = ModelType::SD15; else model_type = ModelType::GEMMA4; } else if (key == "model_variant") { diff --git a/cactus-engine/src/tokenizer.cpp b/cactus-engine/src/tokenizer.cpp index 370c61f3e..f084e769e 100644 --- a/cactus-engine/src/tokenizer.cpp +++ b/cactus-engine/src/tokenizer.cpp @@ -92,12 +92,14 @@ TokenizerRuntimeConfig::VocabFormat parse_vocab_format(const std::string& value) TokenizerRuntimeConfig::Normalizer parse_normalizer(const std::string& value) { if (value == "metaspace") return TokenizerRuntimeConfig::Normalizer::METASPACE; if (value == "byte_level") return TokenizerRuntimeConfig::Normalizer::BYTE_LEVEL; + if (value == "clip") return TokenizerRuntimeConfig::Normalizer::CLIP; return TokenizerRuntimeConfig::Normalizer::NONE; } TokenizerRuntimeConfig::Decoder parse_decoder(const std::string& value) { if (value == "replace_metaspace") return TokenizerRuntimeConfig::Decoder::REPLACE_METASPACE; if (value == "byte_level") return TokenizerRuntimeConfig::Decoder::BYTE_LEVEL; + if (value == "clip") return TokenizerRuntimeConfig::Decoder::CLIP; return TokenizerRuntimeConfig::Decoder::NONE; } diff --git a/cactus-engine/tests/test_diffusion.cpp b/cactus-engine/tests/test_diffusion.cpp new file mode 100644 index 000000000..40c47bb85 --- /dev/null +++ b/cactus-engine/tests/test_diffusion.cpp @@ -0,0 +1,86 @@ +#include "test_utils.h" +#include +#include +#include +#include +#include + +static const char* g_bundle_path = std::getenv("CACTUS_SD_BUNDLE"); + +static bool check_tokens(cactus_model_t model, const char* text, const std::vector& expected) { + uint32_t buffer[128] = {0}; + size_t count = 0; + int rc = cactus_tokenize(model, text, buffer, 128, &count); + if (rc < 0 || count != expected.size()) { + std::cerr << "[✗] tokenize \"" << text << "\": rc=" << rc << " count=" << count + << " expected " << expected.size() << "\n"; + return false; + } + for (size_t i = 0; i < expected.size(); ++i) { + if (buffer[i] != expected[i]) { + std::cerr << "[✗] tokenize \"" << text << "\": token[" << i << "]=" << buffer[i] + << " expected " << expected[i] << "\n"; + return false; + } + } + return true; +} + +static bool test_clip_tokenizer(cactus_model_t model) { + bool ok = check_tokens(model, "a photograph of an astronaut riding a horse", + {320, 8853, 539, 550, 18376, 6765, 320, 4558}); + ok = check_tokens(model, "Hello, WORLD!! (test-123)", + {3306, 267, 1002, 748, 263, 1628, 268, 272, 273, 274, 264}) && ok; + std::cout << "├─ CLIP tokenizer parity: " << (ok ? "ok" : "FAILED") << "\n"; + return ok; +} + +static bool test_generate(cactus_model_t model) { + std::vector rgb(512 * 512 * 3, 0); + unsigned int width = 0; + unsigned int height = 0; + int rc = cactus_generate_image(model, "a photograph of an astronaut riding a horse", + rgb.data(), rgb.size(), &width, &height, 2, 8.5f, 1234ULL); + if (rc != static_cast(rgb.size()) || width != 512 || height != 512) { + std::cerr << "[✗] generate_image rc=" << rc << " " << width << "x" << height << "\n"; + return false; + } + double sum = 0.0; + double sum_sq = 0.0; + for (uint8_t v : rgb) { + sum += v; + sum_sq += static_cast(v) * v; + } + const double mean = sum / rgb.size(); + const double stddev = std::sqrt(sum_sq / rgb.size() - mean * mean); + std::cout << "├─ image 512x512: mean=" << mean << " stddev=" << stddev << "\n"; + if (mean < 5.0 || mean > 250.0 || stddev < 10.0) { + std::cerr << "[✗] generated image looks degenerate\n"; + return false; + } + return true; +} + +static bool test_diffusion() { + cactus_model_t model = cactus_init(g_bundle_path, nullptr, false); + if (!model) { + std::cerr << "[✗] Failed to initialize bundle: " << g_bundle_path << "\n"; + return false; + } + bool passed = test_clip_tokenizer(model); + passed = test_generate(model) && passed; + cactus_destroy(model); + std::cout << "└─ Status: " << (passed ? "PASSED ✓" : "FAILED ✗") << "\n"; + return passed; +} + +int main() { + if (!g_bundle_path || !*g_bundle_path) { + std::cout << "CACTUS_SD_BUNDLE not set; skipping diffusion tests\n"; + return 0; + } + TestUtils::TestRunner runner("Diffusion Tests"); + runner.run_test("text_to_image", test_diffusion()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} From 08c68935c359f8ddadf1969bbcc41f6721f2a4b3 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Sun, 9 Aug 2026 16:44:29 -0700 Subject: [PATCH 27/34] keep declarations grouped by kind in the profile modules Signed-off-by: Neel Shejwalkar --- python/cactus/cli/transpiler.py | 1 - .../cactus/transpiler/Converter/diffusion.py | 18 ++------- python/cactus/transpiler/IR/match_utils.py | 2 - .../transpiler/ModelProfiles/components.py | 19 +++++---- .../transpiler/ModelProfiles/profiles.py | 40 +++++++++---------- .../cactus/transpiler/ModelProfiles/routes.py | 7 +--- 6 files changed, 34 insertions(+), 53 deletions(-) diff --git a/python/cactus/cli/transpiler.py b/python/cactus/cli/transpiler.py index b581b442e..b750e5a9f 100644 --- a/python/cactus/cli/transpiler.py +++ b/python/cactus/cli/transpiler.py @@ -169,7 +169,6 @@ def build_transpiled_bundle( def diffusion_scheduler_metadata(weights_dir: Path) -> dict[str, str]: - """Flatten the diffusers scheduler config into engine-manifest metadata keys.""" scheduler_path = weights_dir / "scheduler_config.json" if not scheduler_path.exists(): return {} diff --git a/python/cactus/transpiler/Converter/diffusion.py b/python/cactus/transpiler/Converter/diffusion.py index 92aaf2e84..b50fbf980 100644 --- a/python/cactus/transpiler/Converter/diffusion.py +++ b/python/cactus/transpiler/Converter/diffusion.py @@ -1,10 +1,3 @@ -"""Diffusion pipeline components: one torch.export per component. - -A diffusion family declares `export_modes` and `component_sources` on its -profile; every mode loads its own module and synthesizes tensor inputs, so -nothing here goes through an HF processor. Model-specific facts (repo ids, -subfolders) stay in the profile -- this module only knows component kinds. -""" from __future__ import annotations from typing import Any @@ -16,9 +9,7 @@ class TextEncoderLastHiddenState(torch.nn.Module): - """Wraps the inner CLIPTextTransformer so exported parameter targets match - the checkpoint keys (text_model.*) instead of doubling the prefix.""" - + #Wraps the inner transformer so exported parameter targets match the checkpoint's text_model.* keys def __init__(self, clip_text_model: torch.nn.Module): super().__init__() self.text_model = clip_text_model.text_model @@ -28,9 +19,7 @@ def forward(self, input_ids: torch.Tensor) -> torch.Tensor: class TaesdDecode(torch.nn.Module): - """The attribute name keeps exported parameter targets on the checkpoint's - decoder.* keys.""" - + #The attribute name keeps exported parameter targets on the checkpoint's decoder.* keys def __init__(self, decoder: torch.nn.Module): super().__init__() self.decoder = decoder @@ -123,8 +112,7 @@ def component_export_inputs(kind: str, configs: dict[str, dict[str, Any]]) -> di text_config = configs.get("text_encoder/config.json", {}) tokens = int(text_config.get("max_position_embeddings", 77)) latent_channels = int(unet_config.get("in_channels", 4)) - # Graphs are shape-specialized; export at 64x64 latents (512x512 images) - # rather than the config's training sample_size. + #Graphs are shape-specialized; export 64x64 latents (512x512 images) rather than the config's training sample_size latent_size = 64 if kind == "clip_text": return {"input_ids": torch.zeros((1, tokens), dtype=torch.int64)} diff --git a/python/cactus/transpiler/IR/match_utils.py b/python/cactus/transpiler/IR/match_utils.py index 05c319ec7..64589cb2d 100644 --- a/python/cactus/transpiler/IR/match_utils.py +++ b/python/cactus/transpiler/IR/match_utils.py @@ -368,8 +368,6 @@ def infer_conv_attrs(fusion: FModels.FusionGraph, bindings: dict[str, models.Nod if len(weight_shape) >= 3: kernel_shape = weight_shape[2:] attrs["kernel_size"] = kernel_shape[0] if all(values_equal(dim, kernel_shape[0]) for dim in kernel_shape) else kernel_shape - # A square kernel collapses to one scalar, so conv1d and conv2d - # definitions need the rank to stay distinguishable. attrs["spatial_dims"] = len(kernel_shape) if len(weight_shape) >= 2: diff --git a/python/cactus/transpiler/ModelProfiles/components.py b/python/cactus/transpiler/ModelProfiles/components.py index b077de88c..eda39fa6c 100644 --- a/python/cactus/transpiler/ModelProfiles/components.py +++ b/python/cactus/transpiler/ModelProfiles/components.py @@ -10,6 +10,9 @@ TOKEN_MERGE = Components(name="token_merge", patterns=("multi_modal", "multimodal", "projector", "merge")) TEXT_DECODER = Components(name="text_decoder", patterns=("language_model", "text", "decoder", "model.layers")) LM_HEAD = Components(name="lm_head", patterns=("lm_head", "logits"),) +TEXT_ENCODER = Components(name="text_encoder", patterns=("text_encoder", "text_model")) +UNET = Components(name="unet", patterns=("unet",)) +VAE_DECODER = Components(name="vae_decoder", patterns=("vae", "decoder")) TEXT_COMPONENTS = { TEXT_EMBED.name: TEXT_EMBED, @@ -36,19 +39,15 @@ TOKEN_MERGE.name: TOKEN_MERGE, } -TEXT_ENCODER = Components(name="text_encoder", patterns=("text_encoder", "text_model")) -UNET = Components(name="unet", patterns=("unet",)) -VAE_DECODER = Components(name="vae_decoder", patterns=("vae", "decoder")) - -T2I_COMPONENTS = { - TEXT_ENCODER.name: TEXT_ENCODER, - UNET.name: UNET, - VAE_DECODER.name: VAE_DECODER, -} - MULTIMODAL_COMPONENTS = { **TEXT_COMPONENTS, VISION_ENCODER.name: VISION_ENCODER, AUDIO_ENCODER.name: AUDIO_ENCODER, TOKEN_MERGE.name: TOKEN_MERGE, } + +T2I_COMPONENTS = { + TEXT_ENCODER.name: TEXT_ENCODER, + UNET.name: UNET, + VAE_DECODER.name: VAE_DECODER, +} diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index b770f3848..acb672fa5 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -206,6 +206,26 @@ def decoder_kv_state(producer: str = "decoder_prefill_chunk") -> StateContract: ), ) +SD15_T2I_RUNTIME_CONTRACT = RuntimeContract( + plan_name="sd15_text_to_image", + execution_strategy="iterative_denoise", + state_owner="request", + cache_persistence="none", + cache_transfer_policy="none", + states=( + StateContract( + name="text_embeddings", kind="activation", producer="text_encoder", + consumers=("unet",), lifetime="request", transfer="copy_or_alias", + metadata=(("outputs", "text_embeddings"),), + ), + StateContract( + name="latents", kind="activation", producer="unet", + consumers=("unet", "vae_decoder"), lifetime="request", transfer="move", + metadata=(("outputs", "latents"),), + ), + ), +) + WHISPER_SUPPRESS_TOKEN_IDS = ( 1, 2, 7, 8, 9, 10, 14, 25, 26, 27, 28, 29, 31, 58, 59, 60, 61, 62, 63, 90, 91, 92, 93, 359, 503, 522, 542, 873, 893, 902, 918, 922, 931, 1350, @@ -568,26 +588,6 @@ def decoder_kv_state(producer: str = "decoder_prefill_chunk") -> StateContract: runtime_contract=WHISPER_RUNTIME_CONTRACT, ) -SD15_T2I_RUNTIME_CONTRACT = RuntimeContract( - plan_name="sd15_text_to_image", - execution_strategy="iterative_denoise", - state_owner="request", - cache_persistence="none", - cache_transfer_policy="none", - states=( - StateContract( - name="text_embeddings", kind="activation", producer="text_encoder", - consumers=("unet",), lifetime="request", transfer="copy_or_alias", - metadata=(("outputs", "text_embeddings"),), - ), - StateContract( - name="latents", kind="activation", producer="unet", - consumers=("unet", "vae_decoder"), lifetime="request", transfer="move", - metadata=(("outputs", "latents"),), - ), - ), -) - LCM_DREAMSHAPER_V7_PROFILE = ModelProfile( model_profiles="sd15_t2i", components=T2I_COMPONENTS, diff --git a/python/cactus/transpiler/ModelProfiles/routes.py b/python/cactus/transpiler/ModelProfiles/routes.py index 21a9e4ab1..59fe3fe06 100644 --- a/python/cactus/transpiler/ModelProfiles/routes.py +++ b/python/cactus/transpiler/ModelProfiles/routes.py @@ -9,6 +9,7 @@ SPEECH_SEQ2SEQ_PREFILL_ROUTE = (combinations.AUDIO_TEXT_TO_DECODER, combinations.DECODER_TO_LM_HEAD) SPEECH_SEQ2SEQ_DECODE_ROUTE = (combinations.AUDIO_TEXT_TO_DECODER, combinations.DECODER_TO_LM_HEAD) SPEECH_TRANSCRIBE_ROUTE = (combinations.AUDIO_ENCODER_TO_ASR_DECODER, combinations.AUDIO_ASR_DECODER_TO_ASR_HEAD) +T2I_GENERATE_ROUTE = (combinations.TEXT_ENCODER_TO_UNET, combinations.UNET_TO_VAE_DECODER) TEXT_PREFILL = InferencePattern(name="text_prefill", route=TEXT_PREFILL_ROUTE,) TEXT_VISION_PREFILL = InferencePattern(name="text_vision_prefill", route=TEXT_VISION_PREFILL_ROUTE) TEXT_AUDIO_PREFILL = InferencePattern(name="text_audio_prefill", route=TEXT_AUDIO_PREFILL_ROUTE) @@ -17,6 +18,7 @@ SPEECH_SEQ2SEQ_PREFILL = InferencePattern(name="speech_seq2seq_prefill", route=SPEECH_SEQ2SEQ_PREFILL_ROUTE,) SPEECH_SEQ2SEQ_DECODE = InferencePattern(name="speech_seq2seq_decode", route=SPEECH_SEQ2SEQ_DECODE_ROUTE) SPEECH_TRANSCRIBE = InferencePattern(name="speech_transcribe", route=SPEECH_TRANSCRIBE_ROUTE) +T2I_GENERATE = InferencePattern(name="t2i_generate", route=T2I_GENERATE_ROUTE) GEMMA4_INFERENCE_PATTERNS = { TEXT_PREFILL.name: TEXT_PREFILL, @@ -46,11 +48,6 @@ TEXT_DECODE.name: TEXT_DECODE, } -# The denoise loop is decode-shaped: the unet edge runs once per scheduler -# step with the latent as persistent state; encode and decode run once. -T2I_GENERATE_ROUTE = (combinations.TEXT_ENCODER_TO_UNET, combinations.UNET_TO_VAE_DECODER) -T2I_GENERATE = InferencePattern(name="t2i_generate", route=T2I_GENERATE_ROUTE) - SD15_T2I_INFERENCE_PATTERNS = { T2I_GENERATE.name: T2I_GENERATE, } From 4c0a701dde9520cd64892215d1f1b6faad848bb0 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Tue, 11 Aug 2026 20:26:14 -0700 Subject: [PATCH 28/34] reuse the Converter registries for pipeline components 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 --- cactus-engine/src/model.cpp | 2 +- python/cactus/cli/model.py | 111 +++++++---------- python/cactus/cli/transpiler.py | 8 +- .../convert/cactus_adapters/tokenizer.py | 57 +++++++++ .../cactus/convert/model_adapters/adapters.py | 1 - .../cactus/transpiler/Converter/constants.py | 1 - .../cactus/transpiler/Converter/diffusion.py | 113 ++++-------------- .../transpiler/Converter/input_processor.py | 32 +++++ python/cactus/transpiler/Converter/models.py | 71 ++++++++--- .../cactus/transpiler/ModelProfiles/models.py | 8 +- .../transpiler/ModelProfiles/profiles.py | 28 ++++- python/tests/test_transpiler_sd_profile.py | 35 ++++-- 12 files changed, 272 insertions(+), 195 deletions(-) diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 19f10a86a..2f8144cbd 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -5765,7 +5765,7 @@ bool Config::from_json(const std::string& config_path) { else if (mt == "needle") model_type = ModelType::NEEDLE; else if (mt == "bert" || mt == "nomic") model_type = ModelType::NOMIC; else if (mt == "generic") model_type = ModelType::GENERIC; - else if (mt == "sd15") model_type = ModelType::SD15; + else if (mt == "sd15_t2i") model_type = ModelType::SD15; else model_type = ModelType::GEMMA4; } else if (key == "model_variant") { diff --git a/python/cactus/cli/model.py b/python/cactus/cli/model.py index b9fcece52..ffd21980a 100644 --- a/python/cactus/cli/model.py +++ b/python/cactus/cli/model.py @@ -47,16 +47,22 @@ def _convert_from_source(model_id, *, bits, token, weights_dir, skip_model_load= err = convert_toolchain_error() if err: raise RuntimeError(err) + if token: + os.environ["HF_TOKEN"] = token + os.environ["HUGGING_FACE_HUB_TOKEN"] = token from .transpiler import profile_for_model_id - component_sources = tuple(getattr(profile_for_model_id(model_id), "component_sources", ()) or ()) - if component_sources: - return _convert_component_sources( - model_id, component_sources, bits=bits, token=token, weights_dir=weights_dir - ) - print_color(YELLOW, f"Converting {model_id} from HuggingFace source...") + profile = profile_for_model_id(model_id) + if profile is not None and profile.component_sources: + return _convert_component_sources(model_id, profile, bits=bits, token=token, weights_dir=weights_dir) + return _convert_whole_model(model_id, bits=bits, weights_dir=weights_dir, skip_model_load=skip_model_load) + + +def _convert_whole_model(model_id, *, bits, weights_dir, skip_model_load): + """Convert a model that ships as a single checkpoint.""" from ..convert.cli import main as cq_main + print_color(YELLOW, f"Converting {model_id} from HuggingFace source...") cq_args = [ "convert", "--model", model_id, "--out", str(weights_dir), @@ -64,37 +70,28 @@ def _convert_from_source(model_id, *, bits, token, weights_dir, skip_model_load= ] if skip_model_load: cq_args.append("--skip-model-load") - if token: - os.environ["HF_TOKEN"] = token - os.environ["HUGGING_FACE_HUB_TOKEN"] = token cq_main(cq_args) print_color(GREEN, f"Model converted and ready at {weights_dir}") return weights_dir -def _convert_component_sources(model_id, component_sources, *, bits, token, weights_dir): - """Convert a multi-repo pipeline (diffusion) component by component into one - weights dir, merging each run's weights manifest.""" +def _convert_component_sources(model_id, profile, *, bits, token, weights_dir): + """Convert a pipeline model component by component into one weights dir, + merging each run's weights manifest.""" import json import tempfile - from huggingface_hub import snapshot_download - from ..convert.cli import main as cq_main + from ..convert.export.files import write_config_txt weights_dir = Path(weights_dir) weights_dir.mkdir(parents=True, exist_ok=True) merged_records = [] converted = [] - for mode, spec in component_sources: - _kind, _, source = spec.partition(":") - if "/" in source: - component_src = source - else: - snapshot = snapshot_download(model_id, allow_patterns=[f"{source}/*"], token=token) - component_src = str(Path(snapshot) / source) - print_color(YELLOW, f"Converting {mode} component from {component_src}...") + for source in profile.component_sources: + component_src = _component_source_path(source, model_id, token=token) + print_color(YELLOW, f"Converting {source.mode} component from {component_src}...") with tempfile.TemporaryDirectory() as tmp: component_out = Path(tmp) / "out" cq_main(["convert", "--model", component_src, "--out", str(component_out), "--bits", str(bits)]) @@ -107,63 +104,45 @@ def _convert_component_sources(model_id, component_sources, *, bits, token, weig destination = weights_dir / tensor_file.name if destination.exists(): raise RuntimeError( - f"components {converted} and {mode!r} both produce tensor file {tensor_file.name!r}" + f"components {converted} and {source.mode!r} both produce tensor file {tensor_file.name!r}" ) shutil.move(str(tensor_file), destination) - converted.append(mode) + converted.append(source.mode) (weights_dir / "weights_manifest.json").write_text( json.dumps({"weights": merged_records}, indent=2), encoding="utf-8" ) - _materialize_diffusion_assets(model_id, weights_dir, token=token) - (weights_dir / "config.txt").write_text( - "model_type=sd15\n" + "\n".join(f"component={mode}" for mode in converted) + "\n", encoding="utf-8" - ) + _materialize_pipeline_assets(model_id, profile, weights_dir, token=token) + write_config_txt({"model_type": profile.model_profiles, "components": converted}, weights_dir) print_color(GREEN, f"Model converted and ready at {weights_dir}") return weights_dir -def _materialize_diffusion_assets(model_id, weights_dir, *, token=None): - """Emit the engine-format tokenizer sidecars and the scheduler config the - diffusion runtime needs alongside the converted weights.""" - import json +def _component_source_path(source, model_id, *, token=None): + from huggingface_hub import snapshot_download + + from cactus.transpiler.ModelProfiles.profiles import component_repo_and_subfolder + repo_id, subfolder = component_repo_and_subfolder(source, model_id) + if not subfolder: + return repo_id + snapshot = snapshot_download(repo_id, allow_patterns=[f"{subfolder}/*"], token=token) + return str(Path(snapshot) / subfolder) + + +def _materialize_pipeline_assets(model_id, profile, weights_dir, *, token=None): + """Emit the tokenizer sidecars and scheduler config the runtime reads from a + pipeline bundle, alongside the converted weights.""" from huggingface_hub import snapshot_download + from ..convert.cactus_adapters.tokenizer import convert_clip_tokenizer + + strategies = {source.load_strategy for source in profile.component_sources} snapshot = Path(snapshot_download(model_id, allow_patterns=["tokenizer/*", "scheduler/*"], token=token)) - vocab = json.loads((snapshot / "tokenizer" / "vocab.json").read_text(encoding="utf-8")) - with open(weights_dir / "vocab.txt", "w", encoding="utf-8") as f: - for token_str, token_id in sorted(vocab.items(), key=lambda kv: kv[1]): - f.write(f"{token_id}\t{token_str}\n") - shutil.copy2(snapshot / "tokenizer" / "merges.txt", weights_dir / "merges.txt") - (weights_dir / "tokenizer_config.txt").write_text( - "vocab_size=49408\n" - "bos_token_id=49406\n" - "eos_token_id=49407\n" - "pad_token_id=49407\n" - "model_max_length=77\n" - "tokenizer_type=bpe\n" - "vocab_format=id_tab_token\n" - "normalizer=clip\n" - "decoder=clip\n" - "byte_fallback=false\n" - "has_chat_template=false\n", - encoding="utf-8", - ) - (weights_dir / "special_tokens.json").write_text( - json.dumps( - { - "bos_token_id": 49406, - "eos_token_id": 49407, - "pad_token_id": 49407, - "vocab_size": 49408, - "model_max_length": 77, - "special_tokens": {"49406": "<|startoftext|>", "49407": "<|endoftext|>"}, - }, - indent=2, - ), - encoding="utf-8", - ) - shutil.copy2(snapshot / "scheduler" / "scheduler_config.json", weights_dir / "scheduler_config.json") + if "clip_text" in strategies: + convert_clip_tokenizer(snapshot / "tokenizer", weights_dir) + scheduler_config = snapshot / "scheduler" / "scheduler_config.json" + if scheduler_config.exists(): + shutil.copy2(scheduler_config, weights_dir / "scheduler_config.json") def ensure_weights(model_id, *, bits=4, token=None, reconvert=False, output_dir=None, skip_model_load=False): diff --git a/python/cactus/cli/transpiler.py b/python/cactus/cli/transpiler.py index b750e5a9f..1f0cee6f6 100644 --- a/python/cactus/cli/transpiler.py +++ b/python/cactus/cli/transpiler.py @@ -190,6 +190,12 @@ def profile_for_model_id(model_id: str): return profiles.profile_for_model_id(model_id) +def export_modes(profile) -> tuple[str, ...]: + from cactus.transpiler.ModelProfiles import profiles + + return profiles.export_modes_for_profile(profile) or DEFAULT_TRANSPILER_MODES + + def resolve_transpile_config( model_id: str, *, @@ -219,7 +225,7 @@ def resolve_transpile_config( profile=profile, modalities=default_modalities(profile), profile_source="registered", - inference_modes=tuple(getattr(profile, "export_modes", ()) or DEFAULT_TRANSPILER_MODES), + inference_modes=export_modes(profile), ) resolved_task = generic_task or "causal-lm" diff --git a/python/cactus/convert/cactus_adapters/tokenizer.py b/python/cactus/convert/cactus_adapters/tokenizer.py index 940a1e7b2..96bbfbdbf 100644 --- a/python/cactus/convert/cactus_adapters/tokenizer.py +++ b/python/cactus/convert/cactus_adapters/tokenizer.py @@ -599,3 +599,60 @@ def convert_sentencepiece_tokenizer(tokenizer_path, output_dir, model_max_length pieces = parse_sentencepiece_pieces(tokenizer_path) meta = _build_sentencepiece_metadata(pieces, model_max_length) _write_sentencepiece_files(output_dir, pieces, meta) + + +def _clip_special_token(tokenizer_config, special_tokens_map, name, default): + value = special_tokens_map.get(name, tokenizer_config.get(name, default)) + if isinstance(value, dict): + return value.get("content", default) + return value + + +def convert_clip_tokenizer(tokenizer_dir, output_dir): + tokenizer_dir, output_dir = Path(tokenizer_dir), Path(output_dir) + import shutil + + vocab = json.loads((tokenizer_dir / "vocab.json").read_text(encoding="utf-8")) + tokenizer_config = json.loads((tokenizer_dir / "tokenizer_config.json").read_text(encoding="utf-8")) + special_tokens_map = json.loads((tokenizer_dir / "special_tokens_map.json").read_text(encoding="utf-8")) + + id_to_token = [""] * (max(vocab.values()) + 1) + for token_str, token_id in vocab.items(): + id_to_token[token_id] = token_str + with open(output_dir / "vocab.txt", "w", encoding="utf-8") as f: + for token_id, token_str in enumerate(id_to_token): + f.write(f"{token_id}\t{token_str}\n") + shutil.copy2(tokenizer_dir / "merges.txt", output_dir / "merges.txt") + + bos_token = _clip_special_token(tokenizer_config, special_tokens_map, "bos_token", "<|startoftext|>") + eos_token = _clip_special_token(tokenizer_config, special_tokens_map, "eos_token", "<|endoftext|>") + pad_token = _clip_special_token(tokenizer_config, special_tokens_map, "pad_token", eos_token) + unk_token = _clip_special_token(tokenizer_config, special_tokens_map, "unk_token", eos_token) + special_token_ids = { + 'bos_token_id': int(vocab[bos_token]), + 'eos_token_id': int(vocab[eos_token]), + 'pad_token_id': int(vocab[pad_token]), + 'unk_token_id': int(vocab[unk_token]), + } + model_max_length = int(tokenizer_config.get("model_max_length", 77)) + + with open(output_dir / "special_tokens.json", 'w', encoding='utf-8') as f: + json.dump({ + **special_token_ids, + "vocab_size": len(id_to_token), + "model_max_length": model_max_length, + "special_tokens": {str(vocab[bos_token]): bos_token, str(vocab[eos_token]): eos_token}, + "additional_special_tokens": [], + }, f, indent=2, ensure_ascii=False) + + with open(output_dir / "tokenizer_config.txt", 'w') as f: + f.write(f"vocab_size={len(id_to_token)}\n") + for key, value in special_token_ids.items(): + f.write(f"{key}={value}\n") + f.write(f"model_max_length={model_max_length}\n") + f.write("tokenizer_type=bpe\n") + f.write("vocab_format=id_tab_token\n") + f.write("normalizer=clip\n") + f.write("decoder=clip\n") + f.write("byte_fallback=false\n") + f.write("has_chat_template=false\n") diff --git a/python/cactus/convert/model_adapters/adapters.py b/python/cactus/convert/model_adapters/adapters.py index 0e17bbb83..ee4df1961 100644 --- a/python/cactus/convert/model_adapters/adapters.py +++ b/python/cactus/convert/model_adapters/adapters.py @@ -503,7 +503,6 @@ def build_calibration_inputs(self, row: dict[str, Any], processor: Any, modality def _verbatim_name_tensor(source_name: str, component: str) -> NameMatch: - """Keep the full dotted checkpoint name, so parallel stacks never collide.""" if not (source_name.endswith(".weight") or source_name.endswith(".bias")): return NameMatch(source_name, None, component, False, hf_name=source_name, adapter_name=source_name) generic = source_name.replace(".", "_") diff --git a/python/cactus/transpiler/Converter/constants.py b/python/cactus/transpiler/Converter/constants.py index ca67ab695..a330101cd 100644 --- a/python/cactus/transpiler/Converter/constants.py +++ b/python/cactus/transpiler/Converter/constants.py @@ -18,7 +18,6 @@ DROP_MULTIMODAL_ON_DECODE_POLICY = "drop_multimodal_on_decode" SYNTHETIC_INPUT_STRATEGY = "synthetic" -DIFFUSION_INPUT_STRATEGY = "diffusion_components" DECODE_WITH_CACHE_MODE = "decode_with_cache" CONVERTER_JSON_DIR = Path(__file__).resolve().parent / "jsons" diff --git a/python/cactus/transpiler/Converter/diffusion.py b/python/cactus/transpiler/Converter/diffusion.py index b50fbf980..031d2b7d8 100644 --- a/python/cactus/transpiler/Converter/diffusion.py +++ b/python/cactus/transpiler/Converter/diffusion.py @@ -4,9 +4,6 @@ import torch -from . import constants -from ..ModelProfiles import models as MP_Models - class TextEncoderLastHiddenState(torch.nn.Module): #Wraps the inner transformer so exported parameter targets match the checkpoint's text_model.* keys @@ -20,9 +17,9 @@ def forward(self, input_ids: torch.Tensor) -> torch.Tensor: class TaesdDecode(torch.nn.Module): #The attribute name keeps exported parameter targets on the checkpoint's decoder.* keys - def __init__(self, decoder: torch.nn.Module): + def __init__(self, autoencoder: torch.nn.Module): super().__init__() - self.decoder = decoder + self.decoder = autoencoder.decoder def forward(self, x: torch.Tensor) -> torch.Tensor: return self.decoder(x) @@ -48,85 +45,27 @@ def forward( ).sample -def create_diffusion_model( - mp: MP_Models.ModelProfile, - model_id: str, - inference_mode: str, - input_cls: Any, - model_cls: Any, -) -> Any: - from .models import EXPORT_PATCHES, load_configs - - kind, source = component_source(mp, inference_mode) - configs = load_configs(mp, model_id) - for patch in mp.export_patches: - EXPORT_PATCHES[patch]() - module = load_component(kind, source, model_id) - input_ = input_cls( - args=(), - kwargs=component_export_inputs(kind, configs), - modalities=("text",), - inference_mode=inference_mode, - ) - return model_cls(name=model_id, model_profile=mp, input=input_, model=module) - - -def component_source(mp: MP_Models.ModelProfile, inference_mode: str) -> tuple[str, str]: - for mode, spec in getattr(mp, "component_sources", ()): - if mode == inference_mode: - kind, _, source = spec.partition(":") - return kind, source - raise ValueError( - f"profile {mp.model_profiles!r} declares no component source for mode {inference_mode!r}" - ) - - -def load_component(kind: str, source: str, model_id: str) -> torch.nn.Module: - if kind == "clip_text": - from transformers import CLIPTextModel - - return TextEncoderLastHiddenState(from_pretrained_local_first(CLIPTextModel, model_id, subfolder=source)) - if kind == "sd_unet": - from diffusers import UNet2DConditionModel - - return UnetEpsilon(from_pretrained_local_first(UNet2DConditionModel, model_id, subfolder=source)) - if kind == "taesd_decoder": - from diffusers import AutoencoderTiny - - return TaesdDecode(from_pretrained_local_first(AutoencoderTiny, source).decoder) - raise ValueError(f"unknown diffusion component kind {kind!r}") - - -def from_pretrained_local_first(model_class: Any, repo_id: str, **kwargs: Any) -> torch.nn.Module: - if constants.token is not None: - kwargs.setdefault("token", constants.token) - try: - model = model_class.from_pretrained(repo_id, torch_dtype=torch.float16, local_files_only=True, **kwargs) - except Exception: - model = model_class.from_pretrained(repo_id, torch_dtype=torch.float16, **kwargs) - return model.eval() - - -def component_export_inputs(kind: str, configs: dict[str, dict[str, Any]]) -> dict[str, torch.Tensor]: - unet_config = configs.get("unet/config.json", {}) - text_config = configs.get("text_encoder/config.json", {}) - tokens = int(text_config.get("max_position_embeddings", 77)) - latent_channels = int(unet_config.get("in_channels", 4)) - #Graphs are shape-specialized; export 64x64 latents (512x512 images) rather than the config's training sample_size - latent_size = 64 - if kind == "clip_text": - return {"input_ids": torch.zeros((1, tokens), dtype=torch.int64)} - if kind == "sd_unet": - return { - "sample": torch.zeros((1, latent_channels, latent_size, latent_size), dtype=torch.float16), - "timestep": torch.full((1,), 999.0, dtype=torch.float16), - "encoder_hidden_states": torch.zeros( - (1, tokens, int(unet_config.get("cross_attention_dim", 768))), dtype=torch.float16 - ), - "timestep_cond": torch.zeros( - (1, int(unet_config.get("time_cond_proj_dim", 256))), dtype=torch.float16 - ), - } - if kind == "taesd_decoder": - return {"x": torch.zeros((1, latent_channels, latent_size, latent_size), dtype=torch.float16)} - raise ValueError(f"unknown diffusion component kind {kind!r}") +def clip_text_class() -> Any: + from transformers import CLIPTextModel + + return CLIPTextModel + + +def sd_unet_class() -> Any: + from diffusers import UNet2DConditionModel + + return UNet2DConditionModel + + +def taesd_class() -> Any: + from diffusers import AutoencoderTiny + + return AutoencoderTiny + + +#diffusers is an optional dependency, so the classes resolve on use rather than joining LOAD_STRATEGIES +COMPONENT_LOADERS = { + "clip_text": (clip_text_class, TextEncoderLastHiddenState), + "sd_unet": (sd_unet_class, UnetEpsilon), + "taesd_decoder": (taesd_class, TaesdDecode), +} diff --git a/python/cactus/transpiler/Converter/input_processor.py b/python/cactus/transpiler/Converter/input_processor.py index 617b5d42f..52149d764 100644 --- a/python/cactus/transpiler/Converter/input_processor.py +++ b/python/cactus/transpiler/Converter/input_processor.py @@ -91,6 +91,32 @@ def __call__(self, audio=None, return_tensors: str = "pt", sampling_rate: int = def parakeet_processor(model_id: str, configs: dict[str, dict[str, Any]], model_profile: str): return ParakeetAudioProcessor(configs) +def _latent_shape(configs: dict[str, dict[str, Any]]) -> tuple[int, int, int, int]: + unet_config = configs.get("unet/config.json", {}) + #Graphs are shape-specialized; export 64x64 latents (512x512 images) rather than the config's training sample_size + return (1, int(unet_config.get("in_channels", 4)), 64, 64) + +def _prompt_token_count(configs: dict[str, dict[str, Any]]) -> int: + return int(configs.get("text_encoder/config.json", {}).get("max_position_embeddings", 77)) + +def clip_text_inputs(configs: dict[str, dict[str, Any]]) -> dict[str, Any]: + import torch + return {"input_ids": torch.zeros((1, _prompt_token_count(configs)), dtype=torch.int64)} + +def sd_unet_inputs(configs: dict[str, dict[str, Any]]) -> dict[str, Any]: + import torch + unet_config = configs.get("unet/config.json", {}) + return { + "sample": torch.zeros(_latent_shape(configs), dtype=torch.float16), + "timestep": torch.full((1,), float(int(configs.get("scheduler/scheduler_config.json", {}).get("num_train_timesteps", 1000)) - 1), dtype=torch.float16), + "encoder_hidden_states": torch.zeros((1, _prompt_token_count(configs), int(unet_config.get("cross_attention_dim", 768))), dtype=torch.float16), + "timestep_cond": torch.zeros((1, int(unet_config.get("time_cond_proj_dim", 256))), dtype=torch.float16), + } + +def taesd_decoder_inputs(configs: dict[str, dict[str, Any]]) -> dict[str, Any]: + import torch + return {"x": torch.zeros(_latent_shape(configs), dtype=torch.float16)} + PROCESSOR_MAP = { "google/gemma-4-E2B": gemma4_processor, "openai/whisper-tiny": default_processor, @@ -99,3 +125,9 @@ def parakeet_processor(model_id: str, configs: dict[str, dict[str, Any]], model_ "Qwen/Qwen2.5-0.5B": default_processor, "LiquidAI/LFM2.5-8B-A1B": text_tokenizer_processor, } + +SYNTHETIC_INPUT_BUILDERS = { + "clip_text": clip_text_inputs, + "sd_unet": sd_unet_inputs, + "taesd_decoder": taesd_decoder_inputs, +} diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index fccde391d..a985e12be 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -13,6 +13,7 @@ from . import input_processor as IP from . import overrides as OV from ..ModelProfiles import models as MP_Models +from ..ModelProfiles import profiles as MP_Profiles from . import constants token = constants.token @@ -186,19 +187,28 @@ def load_configs(mp: MP_Models.ModelProfile, model_id: str | None) -> dict[str, configs[filename] = json.load(f) return configs -def build_input(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], input_cls: Any, model_id: str | None = None, inference_mode: str = "prefill_no_cache") -> Any | None: +def build_input(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], input_cls: Any, model_id: str | None = None, inference_mode: str = "prefill_no_cache", component_source: MP_Models.ComponentSource | None = None) -> Any | None: if not all(modality in mp.supported_modalties for modality in input_modalities): return None configs = load_configs(mp, model_id) if mp.input_strategy == constants.SYNTHETIC_INPUT_STRATEGY: - raise ValueError("Synthetic input creation is currently disabled") + kwargs = build_synthetic_kwargs(mp, component_source, configs) + else: + kwargs = build_processor_kwargs(model_id, input_modalities, configs, mp.model_profiles) return input_cls( args=(), - kwargs=build_processor_kwargs(model_id, input_modalities, configs, mp.model_profiles), + kwargs=kwargs, modalities=input_modalities, inference_mode=inference_mode, ) +def build_synthetic_kwargs(mp: MP_Models.ModelProfile, component_source: MP_Models.ComponentSource | None, configs: dict[str, dict[str, Any]]) -> dict[str, Any]: + strategy = component_source.load_strategy if component_source is not None else mp.load_strategy + builder = IP.SYNTHETIC_INPUT_BUILDERS.get(strategy) + if builder is None: + raise ValueError(f"No synthetic input builder registered for load strategy {strategy!r}") + return builder(configs) + def build_decode_with_cache_input( model: torch.nn.Module, input_: Any, @@ -344,26 +354,33 @@ def pad_tensor_dim(tensor: torch.Tensor, dim: int, target: int, value: int | flo padding = torch.full(shape, value, dtype=tensor.dtype, device=tensor.device) return torch.cat((tensor, padding), dim=dim) -def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch.nn.Module: - if mp is None: +def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None, model_class: Any | None = None, load_kwargs: dict[str, Any] | None = None) -> torch.nn.Module: + export_patches = mp.export_patches if mp is not None else () + if model_class is not None: + candidate_classes = (model_class,) + elif mp is None: candidate_classes = (AutoModel,) - export_patches = () else: candidate_classes = constants.LOAD_STRATEGIES.get(mp.load_strategy, (AutoModel,)) - export_patches = mp.export_patches seen_classes: set[Any] = set() base_kwargs: dict[str, Any] = {"trust_remote_code": True} last_error: Exception | None = None if constants.token is not None: base_kwargs["token"] = constants.token - load_attempts = ( - {**base_kwargs, "dtype": "auto", "low_cpu_mem_usage": True, "local_files_only": True}, - {**base_kwargs, "torch_dtype": "auto", "low_cpu_mem_usage": True, "local_files_only": True}, - {**base_kwargs, "local_files_only": True}, - {**base_kwargs, "dtype": "auto", "low_cpu_mem_usage": True}, - {**base_kwargs, "torch_dtype": "auto", "low_cpu_mem_usage": True}, - base_kwargs, - ) + if load_kwargs: + load_attempts = ( + {**base_kwargs, **load_kwargs, "local_files_only": True}, + {**base_kwargs, **load_kwargs}, + ) + else: + load_attempts = ( + {**base_kwargs, "dtype": "auto", "low_cpu_mem_usage": True, "local_files_only": True}, + {**base_kwargs, "torch_dtype": "auto", "low_cpu_mem_usage": True, "local_files_only": True}, + {**base_kwargs, "local_files_only": True}, + {**base_kwargs, "dtype": "auto", "low_cpu_mem_usage": True}, + {**base_kwargs, "torch_dtype": "auto", "low_cpu_mem_usage": True}, + base_kwargs, + ) for model_class in candidate_classes: if model_class in seen_classes: continue @@ -381,10 +398,28 @@ def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch raise RuntimeError(f"Unable to load model {model_id}") from last_error def create_model(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], model_id: str, inference_mode: str = "prefill_no_cache") -> Model: - if mp.input_strategy == constants.DIFFUSION_INPUT_STRATEGY: - from .diffusion import create_diffusion_model + if mp.component_sources: + return create_component_model(mp, input_modalities, model_id, inference_mode) + return create_whole_model(mp, input_modalities, model_id, inference_mode) + +def create_component_model(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], model_id: str, inference_mode: str) -> Model: + from .diffusion import COMPONENT_LOADERS + + source = MP_Profiles.component_source_for_mode(mp, inference_mode) + if source.load_strategy not in COMPONENT_LOADERS: + raise ValueError(f"No component loader registered for load strategy {source.load_strategy!r}") + resolve_class, wrap_for_export = COMPONENT_LOADERS[source.load_strategy] + repo_id, subfolder = MP_Profiles.component_repo_and_subfolder(source, model_id) + load_kwargs: dict[str, Any] = {"torch_dtype": torch.float16} + if subfolder: + load_kwargs["subfolder"] = subfolder + loaded_model = wrap_for_export(load_model(repo_id, mp, model_class=resolve_class(), load_kwargs=load_kwargs)) + input_ = build_input(mp, input_modalities, Input, model_id, inference_mode, component_source=source) + if input_ is None: + raise ValueError(f"Could not build input for modalities {input_modalities}") + return Model(name=model_id, model_profile=mp, input=input_, model=loaded_model) - return create_diffusion_model(mp=mp, model_id=model_id, inference_mode=inference_mode, input_cls=Input, model_cls=Model) +def create_whole_model(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], model_id: str, inference_mode: str) -> Model: input_ = build_input(mp, input_modalities, Input, model_id, inference_mode) if input_ is None: raise ValueError(f"Could not build input for modalities {input_modalities}") diff --git a/python/cactus/transpiler/ModelProfiles/models.py b/python/cactus/transpiler/ModelProfiles/models.py index 198983e88..120d56055 100644 --- a/python/cactus/transpiler/ModelProfiles/models.py +++ b/python/cactus/transpiler/ModelProfiles/models.py @@ -28,6 +28,12 @@ class InferencePattern: name:str route:tuple[Combinations, ...] +@dataclass(slots=True, frozen=True) +class ComponentSource: + mode:str #Export mode this component is captured under + load_strategy:str #Key into Converter LOAD_STRATEGIES/SYNTHETIC_INPUT_BUILDERS for this component + source:str = "" #Repository subfolder, or a full repository id when the component ships separately + @dataclass(slots=True, frozen=True) class PromptContract: style: str = "auto" @@ -128,7 +134,7 @@ class ModelProfile: export_patches: tuple[str, ...] #Specifies what patches/masks to apply to sample inputs load_strategy: str #Specifies which Hugging Face AutoModel class/loading path to prefer export_modes: tuple[str, ...] = () #Per-component export modes when the family is not prefill/decode shaped; empty means the default modes - component_sources: tuple[tuple[str, str], ...] = () #Mode -> "kind:source" specs for families whose components load separately (diffusion pipelines) + component_sources: tuple[ComponentSource, ...] = () #Components of a pipeline model that load and export separately disabled_fusion_fields: tuple[str, ...] = () #Fusion groups to skip for this profile when a family is known unsafe disabled_fusions: tuple[str, ...] = () #Individual fusion names or cactus ops to skip for this profile prompt_contract: PromptContract = field(default_factory=PromptContract) diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index acb672fa5..2a29ffa78 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -11,6 +11,7 @@ from .models import ( AliasContract, CacheContract, + ComponentSource, MediaContract, ModelProfile, PromptContract, @@ -605,14 +606,13 @@ def decoder_kv_state(producer: str = "decoder_prefill_chunk") -> StateContract: ), fusion_fields=("generic", "conv", "vision", "normalization", "attention", "linear", "embedding"), supported_modalties=("text",), - input_strategy="diffusion_components", + input_strategy="synthetic", export_patches=("clip_position_ids",), - load_strategy="diffusion_components", - export_modes=("text_encoder", "unet", "vae_decoder"), + load_strategy="auto", component_sources=( - ("text_encoder", "clip_text:text_encoder"), - ("unet", "sd_unet:unet"), - ("vae_decoder", "taesd_decoder:madebyollin/taesd"), + ComponentSource(mode="text_encoder", load_strategy="clip_text", source="text_encoder"), + ComponentSource(mode="unet", load_strategy="sd_unet", source="unet"), + ComponentSource(mode="vae_decoder", load_strategy="taesd_decoder", source="madebyollin/taesd"), ), prompt_contract=PromptContract(style="raw", template_source="none"), runtime_contract=SD15_T2I_RUNTIME_CONTRACT, @@ -631,6 +631,22 @@ def decoder_kv_state(producer: str = "decoder_prefill_chunk") -> StateContract: "SimianLuo/LCM_Dreamshaper_v7": LCM_DREAMSHAPER_V7_PROFILE, } +def export_modes_for_profile(profile: ModelProfile) -> tuple[str, ...]: + if profile.export_modes: + return profile.export_modes + return tuple(source.mode for source in profile.component_sources) + +def component_source_for_mode(profile: ModelProfile, mode: str) -> ComponentSource: + for source in profile.component_sources: + if source.mode == mode: + return source + raise ValueError(f"profile {profile.model_profiles!r} declares no component source for mode {mode!r}") + +def component_repo_and_subfolder(source: ComponentSource, model_id: str) -> tuple[str, str]: + if "/" in source.source: + return source.source, "" + return model_id, source.source + def profile_for_model_id(model_id: str) -> ModelProfile | None: """Return only explicitly registered optimized profiles. Unknown models intentionally return None. Their contract must be built by diff --git a/python/tests/test_transpiler_sd_profile.py b/python/tests/test_transpiler_sd_profile.py index 713ceeefa..e79e184c7 100644 --- a/python/tests/test_transpiler_sd_profile.py +++ b/python/tests/test_transpiler_sd_profile.py @@ -5,6 +5,7 @@ from cactus.cli import transpiler as cli_transpiler from cactus.transpiler.Converter import diffusion +from cactus.transpiler.Converter import input_processor from cactus.transpiler.ModelProfiles import profiles from cactus.transpiler.RuntimePlan import models as RPModels @@ -27,41 +28,49 @@ def test_registered_profile_rejects_generic_flags(self): with self.assertRaises(RuntimeError): cli_transpiler.resolve_transpile_config(MODEL_ID, generic_task="causal-lm") - def test_component_sources_parse_kind_and_source(self): + def test_component_sources_resolve_to_a_repository_and_subfolder(self): profile = profiles.profile_for_model_id(MODEL_ID) - self.assertEqual(diffusion.component_source(profile, "unet"), ("sd_unet", "unet")) + unet = profiles.component_source_for_mode(profile, "unet") + self.assertEqual(unet.load_strategy, "sd_unet") + self.assertEqual(profiles.component_repo_and_subfolder(unet, MODEL_ID), (MODEL_ID, "unet")) + + vae = profiles.component_source_for_mode(profile, "vae_decoder") self.assertEqual( - diffusion.component_source(profile, "vae_decoder"), - ("taesd_decoder", "madebyollin/taesd"), + profiles.component_repo_and_subfolder(vae, MODEL_ID), ("madebyollin/taesd", "") ) with self.assertRaises(ValueError): - diffusion.component_source(profile, "decode_with_cache") + profiles.component_source_for_mode(profile, "decode_with_cache") + + def test_every_component_source_has_a_loader_and_an_input_builder(self): + profile = profiles.profile_for_model_id(MODEL_ID) + for source in profile.component_sources: + self.assertIn(source.load_strategy, diffusion.COMPONENT_LOADERS) + self.assertIn(source.load_strategy, input_processor.SYNTHETIC_INPUT_BUILDERS) - def test_component_export_inputs_follow_the_configs(self): + def test_synthetic_inputs_follow_the_configs(self): configs = { "unet/config.json": { - "in_channels": 4, "sample_size": 64, + "in_channels": 4, "sample_size": 96, "cross_attention_dim": 768, "time_cond_proj_dim": 256, }, "text_encoder/config.json": {"max_position_embeddings": 77}, + "scheduler/scheduler_config.json": {"num_train_timesteps": 1000}, } - ids = diffusion.component_export_inputs("clip_text", configs) + ids = input_processor.clip_text_inputs(configs) self.assertEqual(ids["input_ids"].shape, (1, 77)) self.assertEqual(ids["input_ids"].dtype, torch.int64) - unet = diffusion.component_export_inputs("sd_unet", configs) + unet = input_processor.sd_unet_inputs(configs) self.assertEqual(unet["sample"].shape, (1, 4, 64, 64)) self.assertEqual(unet["timestep"].shape, (1,)) + self.assertEqual(float(unet["timestep"][0]), 999.0) self.assertEqual(unet["encoder_hidden_states"].shape, (1, 77, 768)) self.assertEqual(unet["timestep_cond"].shape, (1, 256)) self.assertTrue(all(t.dtype == torch.float16 for t in unet.values())) - latent = diffusion.component_export_inputs("taesd_decoder", configs) + latent = input_processor.taesd_decoder_inputs(configs) self.assertEqual(latent["x"].shape, (1, 4, 64, 64)) - with self.assertRaises(ValueError): - diffusion.component_export_inputs("vqgan", configs) - def test_runtime_plan_declares_the_t2i_route_and_denoise_strategy(self): profile = profiles.profile_for_model_id(MODEL_ID) routes = RPModels.runtime_routes_from_model_profile(profile) From e6a59b904ac063e07eba331b60c49614c702f6a9 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Wed, 12 Aug 2026 13:17:47 -0700 Subject: [PATCH 29/34] give the diffusion scheduler its own translation unit and test 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 --- cactus-engine/src/engine.h | 25 ++- cactus-engine/src/generate_image.cpp | 201 +++++++++++++++++++++++ cactus-engine/src/model.cpp | 218 +------------------------ cactus-engine/tests/test_diffusion.cpp | 60 ++++++- 4 files changed, 281 insertions(+), 223 deletions(-) diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 3be1cb78b..795e1466e 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -298,6 +298,22 @@ struct TokenizerRuntimeConfig { bool has_chat_template = false; }; +struct DiffusionParams { + float beta_start = 0.00085f; + float beta_end = 0.012f; + std::string beta_schedule = "scaled_linear"; + uint32_t num_train_timesteps = 1000; + uint32_t original_inference_steps = 50; + float timestep_scaling = 10.0f; + std::string prediction_type = "epsilon"; +}; + +std::vector diffusion_alphas_cumprod(const DiffusionParams& params); +std::vector diffusion_lcm_timesteps(const DiffusionParams& params, uint32_t steps); +std::vector diffusion_guidance_embedding(float guidance_embedding_scale, size_t dim); +void write_typed_buffer(std::vector& buf, Precision dst_prec, + const void* src_data, size_t src_bytes, Precision src_prec); + TokenizerRuntimeConfig load_tokenizer_runtime_config(const std::string& config_file); void load_special_tokens_map(const std::string& config_file, std::unordered_map& special_tokens); std::vector split_with_special_tokens(const std::string& text, const std::unordered_map& special_tokens); @@ -1045,15 +1061,6 @@ class Model { Component* prefill_encoder_ = nullptr; enum class DecodeRoute { CACHED_STEP, DIRECT_DECODER_STEP, FULL_CONTEXT_TEXT, ENCODER_CROSS_KV_STEP, ITERATIVE_DENOISE }; DecodeRoute decode_route_ = DecodeRoute::CACHED_STEP; - struct DiffusionParams { - float beta_start = 0.00085f; - float beta_end = 0.012f; - std::string beta_schedule = "scaled_linear"; - uint32_t num_train_timesteps = 1000; - uint32_t original_inference_steps = 50; - float timestep_scaling = 10.0f; - std::string prediction_type = "epsilon"; - }; DiffusionParams diffusion_; Component* source_encoder_ = nullptr; Component* decoder_cross_kv_ = nullptr; diff --git a/cactus-engine/src/generate_image.cpp b/cactus-engine/src/generate_image.cpp index c64b7ef40..246dda746 100644 --- a/cactus-engine/src/generate_image.cpp +++ b/cactus-engine/src/generate_image.cpp @@ -1,7 +1,208 @@ #include "../cactus_engine.h" +#include "engine.h" #include "utils.h" +#include #include #include +#include + +namespace cactus { +namespace engine { + +std::vector diffusion_alphas_cumprod(const DiffusionParams& params) { + std::vector alphas_cumprod(params.num_train_timesteps); + const double sqrt_start = std::sqrt(static_cast(params.beta_start)); + const double sqrt_end = std::sqrt(static_cast(params.beta_end)); + double cumprod = 1.0; + for (uint32_t i = 0; i < params.num_train_timesteps; ++i) { + const double frac = params.num_train_timesteps > 1 + ? static_cast(i) / (params.num_train_timesteps - 1) : 0.0; + const double sqrt_beta = sqrt_start + (sqrt_end - sqrt_start) * frac; + cumprod *= 1.0 - sqrt_beta * sqrt_beta; + alphas_cumprod[i] = static_cast(cumprod); + } + return alphas_cumprod; +} + +std::vector diffusion_lcm_timesteps(const DiffusionParams& params, uint32_t steps) { + if (steps == 0 || params.original_inference_steps == 0) return {}; + const uint32_t k = params.num_train_timesteps / params.original_inference_steps; + std::vector timesteps(steps); + for (uint32_t s = 0; s < steps; ++s) { + const uint32_t index = static_cast( + std::floor(static_cast(s) * params.original_inference_steps / steps)); + timesteps[s] = (params.original_inference_steps - index) * k - 1; + } + return timesteps; +} + +std::vector diffusion_guidance_embedding(float guidance_embedding_scale, size_t dim) { + const size_t half = dim / 2; + std::vector embedding(dim, 0.0f); + const double scaled = static_cast(guidance_embedding_scale) * 1000.0; + const double span = half > 1 ? static_cast(half - 1) : 1.0; + for (size_t i = 0; i < half; ++i) { + const double freq = std::exp(-std::log(10000.0) * static_cast(i) / span); + embedding[i] = static_cast(std::sin(scaled * freq)); + embedding[half + i] = static_cast(std::cos(scaled * freq)); + } + return embedding; +} + +int Model::generate_image(const std::string& prompt, uint8_t* out_rgb, size_t out_capacity, + uint32_t* out_width, uint32_t* out_height, + int steps, float guidance_scale, uint64_t seed) { + if (decode_route_ != DecodeRoute::ITERATIVE_DENOISE) { + CACTUS_LOG_ERROR("model", "generate_image requires a text-to-image bundle"); + return -1; + } + if (steps < 1 || steps > static_cast(diffusion_.original_inference_steps)) { + CACTUS_LOG_ERROR("model", "generate_image steps must be in [1, " << diffusion_.original_inference_steps << "]"); + return -1; + } + Component& text_encoder = components_.at("text_encoder"); + Component& unet = components_.at("unet"); + Component& vae_decoder = components_.at("vae_decoder"); + if (!text_encoder.graph || !unet.graph || !vae_decoder.graph || !tokenizer_) return -1; + + const int ids_idx = input_index(text_encoder, "input_ids"); + const int sample_idx = input_index(unet, "sample"); + const int timestep_idx = input_index(unet, "timestep"); + const int hidden_idx = input_index(unet, "encoder_hidden_states"); + const int cond_idx = input_index(unet, "timestep_cond"); + const int latent_idx = input_index(vae_decoder, "x"); + if (ids_idx < 0 || sample_idx < 0 || timestep_idx < 0 || hidden_idx < 0 || cond_idx < 0 || latent_idx < 0) { + CACTUS_LOG_ERROR("model", "text-to-image bundle is missing expected component inputs"); + return -1; + } + + const auto& ids_desc = text_encoder.graph->get_output_buffer( + static_cast(text_encoder.runtime_input_node_ids[ids_idx])); + const size_t max_tokens = ids_desc.total_size; + const uint32_t bos = tokenizer_->get_bos_token(); + const uint32_t eos = tokenizer_->get_eos_token(); + auto tokens = tokenizer_->encode(prompt); + if (tokens.size() > max_tokens - 2) tokens.resize(max_tokens - 2); + write_int_input_at(text_encoder, "input_ids", 0, static_cast(bos)); + for (size_t i = 0; i < tokens.size(); ++i) { + write_int_input_at(text_encoder, "input_ids", i + 1, static_cast(tokens[i])); + } + for (size_t i = tokens.size() + 1; i < max_tokens; ++i) { + write_int_input_at(text_encoder, "input_ids", i, static_cast(eos)); + } + text_encoder.graph->execute(); + + const size_t hidden_node = static_cast(text_encoder.output_node_ids[0]); + const auto& hidden_desc = text_encoder.graph->get_output_buffer(hidden_node); + write_typed_buffer( + unet.input_buffers[hidden_idx], + unet.graph->get_output_buffer(static_cast(unet.runtime_input_node_ids[hidden_idx])).precision, + text_encoder.graph->get_output(hidden_node), + hidden_desc.byte_size, + hidden_desc.precision); + + const auto& cond_desc = unet.graph->get_output_buffer( + static_cast(unet.runtime_input_node_ids[cond_idx])); + const std::vector guidance = diffusion_guidance_embedding(guidance_scale - 1.0f, cond_desc.total_size); + write_typed_buffer(unet.input_buffers[cond_idx], cond_desc.precision, + guidance.data(), guidance.size() * sizeof(float), Precision::FP32); + + const auto& sample_desc = unet.graph->get_output_buffer( + static_cast(unet.runtime_input_node_ids[sample_idx])); + const size_t latent_count = sample_desc.total_size; + std::mt19937_64 rng(seed); + std::normal_distribution normal(0.0f, 1.0f); + std::vector latents(latent_count); + for (float& value : latents) value = normal(rng); + + const std::vector alphas_cumprod = diffusion_alphas_cumprod(diffusion_); + const std::vector timesteps = diffusion_lcm_timesteps(diffusion_, static_cast(steps)); + + const auto& timestep_desc = unet.graph->get_output_buffer( + static_cast(unet.runtime_input_node_ids[timestep_idx])); + std::vector eps(latent_count); + for (int s = 0; s < steps; ++s) { + const uint32_t t = timesteps[s]; + write_typed_buffer(unet.input_buffers[sample_idx], sample_desc.precision, + latents.data(), latents.size() * sizeof(float), Precision::FP32); + const float timestep_value = static_cast(t); + write_typed_buffer(unet.input_buffers[timestep_idx], timestep_desc.precision, + ×tep_value, sizeof(float), Precision::FP32); + unet.graph->execute(); + + const size_t eps_node = static_cast(unet.output_node_ids[0]); + const auto& eps_desc = unet.graph->get_output_buffer(eps_node); + if (eps_desc.precision == Precision::FP16) { + const __fp16* data = static_cast(unet.graph->get_output(eps_node)); + for (size_t i = 0; i < latent_count; ++i) eps[i] = static_cast(data[i]); + } else { + std::memcpy(eps.data(), unet.graph->get_output(eps_node), latent_count * sizeof(float)); + } + + const float acp_t = alphas_cumprod[t]; + const float sqrt_acp = std::sqrt(acp_t); + const float sqrt_one_minus = std::sqrt(1.0f - acp_t); + const float scaled_t = static_cast(t) * diffusion_.timestep_scaling; + const float c_skip = 0.25f / (scaled_t * scaled_t + 0.25f); + const float c_out = scaled_t / std::sqrt(scaled_t * scaled_t + 0.25f); + const bool final_step = s == steps - 1; + const float acp_prev = final_step ? 1.0f : alphas_cumprod[timesteps[s + 1]]; + const float sqrt_acp_prev = std::sqrt(acp_prev); + const float sqrt_one_minus_prev = std::sqrt(1.0f - acp_prev); + for (size_t i = 0; i < latent_count; ++i) { + const float x0 = (latents[i] - sqrt_one_minus * eps[i]) / sqrt_acp; + const float denoised = c_out * x0 + c_skip * latents[i]; + latents[i] = final_step ? denoised + : sqrt_acp_prev * denoised + sqrt_one_minus_prev * normal(rng); + } + } + + const auto& latent_desc = vae_decoder.graph->get_output_buffer( + static_cast(vae_decoder.runtime_input_node_ids[latent_idx])); + write_typed_buffer(vae_decoder.input_buffers[latent_idx], latent_desc.precision, + latents.data(), latents.size() * sizeof(float), Precision::FP32); + vae_decoder.graph->execute(); + + const size_t image_node = static_cast(vae_decoder.output_node_ids[0]); + const auto& image_desc = vae_decoder.graph->get_output_buffer(image_node); + if (image_desc.shape.size() != 4 || image_desc.shape[1] != 3) { + CACTUS_LOG_ERROR("model", "vae_decoder output is not an NCHW RGB image"); + return -1; + } + const size_t height = image_desc.shape[2]; + const size_t width = image_desc.shape[3]; + const size_t out_bytes = height * width * 3; + if (out_capacity < out_bytes) { + CACTUS_LOG_ERROR("model", "generate_image buffer too small: need " << out_bytes << " bytes"); + return -2; + } + + const size_t plane = height * width; + auto write_pixels = [&](auto value_at) { + for (size_t y = 0; y < height; ++y) { + for (size_t x = 0; x < width; ++x) { + const size_t pixel = y * width + x; + for (size_t c = 0; c < 3; ++c) { + const float value = std::min(1.0f, std::max(0.0f, value_at(c * plane + pixel))); + out_rgb[pixel * 3 + c] = static_cast(std::lround(value * 255.0f)); + } + } + } + }; + if (image_desc.precision == Precision::FP16) { + const __fp16* data = static_cast(vae_decoder.graph->get_output(image_node)); + write_pixels([data](size_t i) { return static_cast(data[i]); }); + } else { + const float* data = static_cast(vae_decoder.graph->get_output(image_node)); + write_pixels([data](size_t i) { return data[i]; }); + } + if (out_width) *out_width = static_cast(width); + if (out_height) *out_height = static_cast(height); + return static_cast(out_bytes); +} + +} // namespace engine +} // namespace cactus using namespace cactus::engine; using namespace cactus::ffi; diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 2f8144cbd..1cb364a13 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -10,7 +10,6 @@ #include #include #include -#include #include #include #include @@ -174,11 +173,6 @@ std::map json_string_map(const picojson::object& obj, return out; } -namespace { -void write_typed_buffer(std::vector& buf, Precision dst_prec, - const void* src_data, size_t src_bytes, Precision src_prec); -} // namespace - size_t parse_ceil_divisor(const std::string& value) { const std::string prefix = "ceil_div:"; if (value.rfind(prefix, 0) != 0) return 0; @@ -767,9 +761,15 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, if (has_metadata_encoder_cross_kv_route) { decode_route_ = DecodeRoute::ENCODER_CROSS_KV_STEP; required_components = {decoder_name}; - } else if (components_.count("text_encoder") && components_.count("unet") && components_.count("vae_decoder")) { + } else if (runtime_execution_strategy_ == "iterative_denoise") { decode_route_ = DecodeRoute::ITERATIVE_DENOISE; required_components = {"text_encoder", "unet", "vae_decoder"}; + for (const std::string& name : required_components) { + if (!components_.count(name)) { + CACTUS_LOG_ERROR("model", "iterative_denoise bundle is missing the " << name << " component"); + return false; + } + } } else if (has_chunked_prefill) { encoder_name = "lm_encoder_step"; decoder_name = "decoder_media_step"; @@ -2460,11 +2460,6 @@ void Model::run_media_step(size_t position, const uint8_t* feature_row, size_t f run_step(static_cast(config_.pad_token_id), position, false); } -namespace { -void write_typed_buffer(std::vector& buf, Precision dst_prec, - const void* src_data, size_t src_bytes, Precision src_prec); -} // namespace - void Model::run_vision_encoder(const std::string& image_path) { if (!vision_encoder_) return; const std::string strategy = image_preprocess_strategy(); @@ -3301,8 +3296,6 @@ void Model::prefill_with_audio(const std::vector& tokens, prefill_with_media(tokens, {}, audio_features_per_message, profile_file); } -namespace { - void write_typed_buffer(std::vector& buf, Precision dst_prec, const void* src_data, size_t src_bytes, Precision src_prec) { if (dst_prec == src_prec) { @@ -3332,6 +3325,8 @@ void write_typed_buffer(std::vector& buf, Precision dst_prec, } } +namespace { + static inline void write_int_element(uint8_t* buf, Precision prec, size_t index, int64_t v) { switch (prec) { case Precision::FP32: reinterpret_cast(buf)[index] = static_cast(v); break; @@ -5469,201 +5464,6 @@ std::vector Model::get_embeddings(const std::vector& tokens, bo return result; } -namespace { - -std::vector diffusion_alphas_cumprod(float beta_start, float beta_end, uint32_t num_train_timesteps) { - std::vector alphas_cumprod(num_train_timesteps); - const double sqrt_start = std::sqrt(static_cast(beta_start)); - const double sqrt_end = std::sqrt(static_cast(beta_end)); - double cumprod = 1.0; - for (uint32_t i = 0; i < num_train_timesteps; ++i) { - const double frac = num_train_timesteps > 1 ? static_cast(i) / (num_train_timesteps - 1) : 0.0; - const double sqrt_beta = sqrt_start + (sqrt_end - sqrt_start) * frac; - cumprod *= 1.0 - sqrt_beta * sqrt_beta; - alphas_cumprod[i] = static_cast(cumprod); - } - return alphas_cumprod; -} - -std::vector diffusion_lcm_timesteps(uint32_t num_train_timesteps, uint32_t original_steps, uint32_t steps) { - const uint32_t k = num_train_timesteps / original_steps; - std::vector timesteps(steps); - for (uint32_t s = 0; s < steps; ++s) { - const uint32_t index = static_cast( - std::floor(static_cast(s) * original_steps / steps)); - timesteps[s] = (original_steps - index) * k - 1; - } - return timesteps; -} - -std::vector diffusion_guidance_embedding(float w, size_t dim) { - const size_t half = dim / 2; - std::vector embedding(dim, 0.0f); - const double scaled = static_cast(w) * 1000.0; - for (size_t i = 0; i < half; ++i) { - const double freq = std::exp(-std::log(10000.0) * static_cast(i) / (half - 1)); - embedding[i] = static_cast(std::sin(scaled * freq)); - embedding[half + i] = static_cast(std::cos(scaled * freq)); - } - return embedding; -} - -} // namespace - -int Model::generate_image(const std::string& prompt, uint8_t* out_rgb, size_t out_capacity, - uint32_t* out_width, uint32_t* out_height, - int steps, float guidance_scale, uint64_t seed) { - if (decode_route_ != DecodeRoute::ITERATIVE_DENOISE) { - CACTUS_LOG_ERROR("model", "generate_image requires a text-to-image bundle"); - return -1; - } - if (steps < 1 || steps > static_cast(diffusion_.original_inference_steps)) { - CACTUS_LOG_ERROR("model", "generate_image steps must be in [1, " << diffusion_.original_inference_steps << "]"); - return -1; - } - Component& text_encoder = components_.at("text_encoder"); - Component& unet = components_.at("unet"); - Component& vae_decoder = components_.at("vae_decoder"); - if (!text_encoder.graph || !unet.graph || !vae_decoder.graph || !tokenizer_) return -1; - - const int ids_idx = input_index(text_encoder, "input_ids"); - const int sample_idx = input_index(unet, "sample"); - const int timestep_idx = input_index(unet, "timestep"); - const int hidden_idx = input_index(unet, "encoder_hidden_states"); - const int cond_idx = input_index(unet, "timestep_cond"); - const int latent_idx = input_index(vae_decoder, "x"); - if (ids_idx < 0 || sample_idx < 0 || timestep_idx < 0 || hidden_idx < 0 || cond_idx < 0 || latent_idx < 0) { - CACTUS_LOG_ERROR("model", "text-to-image bundle is missing expected component inputs"); - return -1; - } - - const auto& ids_desc = text_encoder.graph->get_output_buffer( - static_cast(text_encoder.runtime_input_node_ids[ids_idx])); - const size_t max_tokens = ids_desc.total_size; - const uint32_t bos = tokenizer_->get_bos_token(); - const uint32_t eos = tokenizer_->get_eos_token(); - auto tokens = tokenizer_->encode(prompt); - if (tokens.size() > max_tokens - 2) tokens.resize(max_tokens - 2); - write_int_input_at(text_encoder, "input_ids", 0, static_cast(bos)); - for (size_t i = 0; i < tokens.size(); ++i) { - write_int_input_at(text_encoder, "input_ids", i + 1, static_cast(tokens[i])); - } - for (size_t i = tokens.size() + 1; i < max_tokens; ++i) { - write_int_input_at(text_encoder, "input_ids", i, static_cast(eos)); - } - text_encoder.graph->execute(); - - const size_t hidden_node = static_cast(text_encoder.output_node_ids[0]); - const auto& hidden_desc = text_encoder.graph->get_output_buffer(hidden_node); - write_typed_buffer( - unet.input_buffers[hidden_idx], - unet.graph->get_output_buffer(static_cast(unet.runtime_input_node_ids[hidden_idx])).precision, - text_encoder.graph->get_output(hidden_node), - hidden_desc.byte_size, - hidden_desc.precision); - - const auto& cond_desc = unet.graph->get_output_buffer( - static_cast(unet.runtime_input_node_ids[cond_idx])); - const std::vector guidance = diffusion_guidance_embedding(guidance_scale - 1.0f, cond_desc.total_size); - write_typed_buffer(unet.input_buffers[cond_idx], cond_desc.precision, - guidance.data(), guidance.size() * sizeof(float), Precision::FP32); - - const auto& sample_desc = unet.graph->get_output_buffer( - static_cast(unet.runtime_input_node_ids[sample_idx])); - const size_t latent_count = sample_desc.total_size; - std::mt19937_64 rng(seed); - std::normal_distribution normal(0.0f, 1.0f); - std::vector latents(latent_count); - for (float& value : latents) value = normal(rng); - - const std::vector alphas_cumprod = - diffusion_alphas_cumprod(diffusion_.beta_start, diffusion_.beta_end, diffusion_.num_train_timesteps); - const std::vector timesteps = diffusion_lcm_timesteps( - diffusion_.num_train_timesteps, diffusion_.original_inference_steps, static_cast(steps)); - - const auto& timestep_desc = unet.graph->get_output_buffer( - static_cast(unet.runtime_input_node_ids[timestep_idx])); - std::vector eps(latent_count); - for (int s = 0; s < steps; ++s) { - const uint32_t t = timesteps[s]; - write_typed_buffer(unet.input_buffers[sample_idx], sample_desc.precision, - latents.data(), latents.size() * sizeof(float), Precision::FP32); - const float timestep_value = static_cast(t); - write_typed_buffer(unet.input_buffers[timestep_idx], timestep_desc.precision, - ×tep_value, sizeof(float), Precision::FP32); - unet.graph->execute(); - - const size_t eps_node = static_cast(unet.output_node_ids[0]); - const auto& eps_desc = unet.graph->get_output_buffer(eps_node); - if (eps_desc.precision == Precision::FP16) { - const __fp16* data = static_cast(unet.graph->get_output(eps_node)); - for (size_t i = 0; i < latent_count; ++i) eps[i] = static_cast(data[i]); - } else { - std::memcpy(eps.data(), unet.graph->get_output(eps_node), latent_count * sizeof(float)); - } - - const float acp_t = alphas_cumprod[t]; - const float sqrt_acp = std::sqrt(acp_t); - const float sqrt_one_minus = std::sqrt(1.0f - acp_t); - const float scaled_t = static_cast(t) * diffusion_.timestep_scaling; - const float c_skip = 0.25f / (scaled_t * scaled_t + 0.25f); - const float c_out = scaled_t / std::sqrt(scaled_t * scaled_t + 0.25f); - const bool final_step = s == steps - 1; - const float acp_prev = final_step ? 1.0f : alphas_cumprod[timesteps[s + 1]]; - const float sqrt_acp_prev = std::sqrt(acp_prev); - const float sqrt_one_minus_prev = std::sqrt(1.0f - acp_prev); - for (size_t i = 0; i < latent_count; ++i) { - const float x0 = (latents[i] - sqrt_one_minus * eps[i]) / sqrt_acp; - const float denoised = c_out * x0 + c_skip * latents[i]; - latents[i] = final_step ? denoised - : sqrt_acp_prev * denoised + sqrt_one_minus_prev * normal(rng); - } - } - - const auto& latent_desc = vae_decoder.graph->get_output_buffer( - static_cast(vae_decoder.runtime_input_node_ids[latent_idx])); - write_typed_buffer(vae_decoder.input_buffers[latent_idx], latent_desc.precision, - latents.data(), latents.size() * sizeof(float), Precision::FP32); - vae_decoder.graph->execute(); - - const size_t image_node = static_cast(vae_decoder.output_node_ids[0]); - const auto& image_desc = vae_decoder.graph->get_output_buffer(image_node); - if (image_desc.shape.size() != 4 || image_desc.shape[1] != 3) { - CACTUS_LOG_ERROR("model", "vae_decoder output is not an NCHW RGB image"); - return -1; - } - const size_t height = image_desc.shape[2]; - const size_t width = image_desc.shape[3]; - const size_t out_bytes = height * width * 3; - if (out_capacity < out_bytes) { - CACTUS_LOG_ERROR("model", "generate_image buffer too small: need " << out_bytes << " bytes"); - return -2; - } - - const size_t plane = height * width; - auto write_pixels = [&](auto value_at) { - for (size_t y = 0; y < height; ++y) { - for (size_t x = 0; x < width; ++x) { - const size_t pixel = y * width + x; - for (size_t c = 0; c < 3; ++c) { - const float value = std::min(1.0f, std::max(0.0f, value_at(c * plane + pixel))); - out_rgb[pixel * 3 + c] = static_cast(std::lround(value * 255.0f)); - } - } - } - }; - if (image_desc.precision == Precision::FP16) { - const __fp16* data = static_cast(vae_decoder.graph->get_output(image_node)); - write_pixels([data](size_t i) { return static_cast(data[i]); }); - } else { - const float* data = static_cast(vae_decoder.graph->get_output(image_node)); - write_pixels([data](size_t i) { return data[i]; }); - } - if (out_width) *out_width = static_cast(width); - if (out_height) *out_height = static_cast(height); - return static_cast(out_bytes); -} - bool Config::from_json(const std::string& config_path) { std::ifstream file(config_path); if (!file) { diff --git a/cactus-engine/tests/test_diffusion.cpp b/cactus-engine/tests/test_diffusion.cpp index 40c47bb85..228fe0efd 100644 --- a/cactus-engine/tests/test_diffusion.cpp +++ b/cactus-engine/tests/test_diffusion.cpp @@ -1,4 +1,5 @@ #include "test_utils.h" +#include "../src/engine.h" #include #include #include @@ -61,6 +62,54 @@ static bool test_generate(cactus_model_t model) { return true; } +static bool test_lcm_scheduler() { + cactus::engine::DiffusionParams params; + const std::vector expected_four = {999, 759, 499, 259}; + if (cactus::engine::diffusion_lcm_timesteps(params, 4) != expected_four) { + std::cerr << "[✗] 4-step LCM schedule does not match diffusers\n"; + return false; + } + const std::vector expected_two = {999, 499}; + if (cactus::engine::diffusion_lcm_timesteps(params, 2) != expected_two) { + std::cerr << "[✗] 2-step LCM schedule does not match diffusers\n"; + return false; + } + if (!cactus::engine::diffusion_lcm_timesteps(params, 0).empty()) { + std::cerr << "[✗] zero-step schedule should be empty\n"; + return false; + } + + const std::vector alphas = cactus::engine::diffusion_alphas_cumprod(params); + if (alphas.size() != params.num_train_timesteps) { + std::cerr << "[✗] alphas_cumprod has the wrong length\n"; + return false; + } + if (!(alphas.front() < 1.0f && alphas.back() > 0.0f && alphas.back() < alphas.front())) { + std::cerr << "[✗] alphas_cumprod is not a decreasing schedule inside (0, 1)\n"; + return false; + } + if (std::abs(alphas[999] - 0.00466f) > 1e-4f) { + std::cerr << "[✗] alphas_cumprod[999]=" << alphas[999] << " expected ~0.00466\n"; + return false; + } + + for (size_t dim : {2u, 256u}) { + const std::vector embedding = cactus::engine::diffusion_guidance_embedding(7.5f, dim); + if (embedding.size() != dim) { + std::cerr << "[✗] guidance embedding has the wrong length for dim=" << dim << "\n"; + return false; + } + for (float value : embedding) { + if (!std::isfinite(value)) { + std::cerr << "[✗] guidance embedding is not finite for dim=" << dim << "\n"; + return false; + } + } + } + std::cout << "├─ LCM scheduler: ok\n"; + return true; +} + static bool test_diffusion() { cactus_model_t model = cactus_init(g_bundle_path, nullptr, false); if (!model) { @@ -75,12 +124,13 @@ static bool test_diffusion() { } int main() { - if (!g_bundle_path || !*g_bundle_path) { - std::cout << "CACTUS_SD_BUNDLE not set; skipping diffusion tests\n"; - return 0; - } TestUtils::TestRunner runner("Diffusion Tests"); - runner.run_test("text_to_image", test_diffusion()); + runner.run_test("lcm_scheduler", test_lcm_scheduler()); + if (g_bundle_path && *g_bundle_path) { + runner.run_test("text_to_image", test_diffusion()); + } else { + std::cout << "CACTUS_SD_BUNDLE not set; skipping text-to-image generation\n"; + } runner.print_summary(); return runner.all_passed() ? 0 : 1; } From 5b18b7207f8093ccf6e1666506f6739969fc16c7 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Wed, 12 Aug 2026 13:38:26 -0700 Subject: [PATCH 30/34] drop comments the surrounding files do not carry Signed-off-by: Neel Shejwalkar --- cactus-engine/src/bpe.cpp | 1 - python/cactus/cli/model.py | 6 ++---- python/cactus/transpiler/Converter/diffusion.py | 1 - python/cactus/transpiler/Converter/input_processor.py | 1 - python/cactus/transpiler/Converter/overrides.py | 3 +-- python/tests/test_transpile_clip.py | 2 +- python/tests/test_transpile_sd_unet.py | 4 +--- python/tests/test_transpile_t2i.py | 8 -------- 8 files changed, 5 insertions(+), 21 deletions(-) diff --git a/cactus-engine/src/bpe.cpp b/cactus-engine/src/bpe.cpp index 54e5ed4c3..8d3c59720 100644 --- a/cactus-engine/src/bpe.cpp +++ b/cactus-engine/src/bpe.cpp @@ -457,7 +457,6 @@ std::vector BPETokenizer::apply_bpe(const std::vector& } std::vector BPETokenizer::clip_word_split(const std::string& text) const { - // ASCII matches CLIP's regex exactly; non-ASCII UTF-8 is treated as letters. static const char* kContractions[] = {"'s", "'t", "'re", "'ve", "'m", "'ll", "'d"}; auto is_letter = [](unsigned char c) { return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || c >= 0x80; }; auto is_digit = [](unsigned char c) { return c >= '0' && c <= '9'; }; diff --git a/python/cactus/cli/model.py b/python/cactus/cli/model.py index ffd21980a..750d5da9b 100644 --- a/python/cactus/cli/model.py +++ b/python/cactus/cli/model.py @@ -77,8 +77,7 @@ def _convert_whole_model(model_id, *, bits, weights_dir, skip_model_load): def _convert_component_sources(model_id, profile, *, bits, token, weights_dir): - """Convert a pipeline model component by component into one weights dir, - merging each run's weights manifest.""" + """Convert a pipeline model one component at a time into a single weights dir.""" import json import tempfile @@ -130,8 +129,7 @@ def _component_source_path(source, model_id, *, token=None): def _materialize_pipeline_assets(model_id, profile, weights_dir, *, token=None): - """Emit the tokenizer sidecars and scheduler config the runtime reads from a - pipeline bundle, alongside the converted weights.""" + """Emit the tokenizer sidecars and scheduler config a pipeline bundle needs at runtime.""" from huggingface_hub import snapshot_download from ..convert.cactus_adapters.tokenizer import convert_clip_tokenizer diff --git a/python/cactus/transpiler/Converter/diffusion.py b/python/cactus/transpiler/Converter/diffusion.py index 031d2b7d8..28f779d5e 100644 --- a/python/cactus/transpiler/Converter/diffusion.py +++ b/python/cactus/transpiler/Converter/diffusion.py @@ -63,7 +63,6 @@ def taesd_class() -> Any: return AutoencoderTiny -#diffusers is an optional dependency, so the classes resolve on use rather than joining LOAD_STRATEGIES COMPONENT_LOADERS = { "clip_text": (clip_text_class, TextEncoderLastHiddenState), "sd_unet": (sd_unet_class, UnetEpsilon), diff --git a/python/cactus/transpiler/Converter/input_processor.py b/python/cactus/transpiler/Converter/input_processor.py index 52149d764..bdb274187 100644 --- a/python/cactus/transpiler/Converter/input_processor.py +++ b/python/cactus/transpiler/Converter/input_processor.py @@ -93,7 +93,6 @@ def parakeet_processor(model_id: str, configs: dict[str, dict[str, Any]], model_ def _latent_shape(configs: dict[str, dict[str, Any]]) -> tuple[int, int, int, int]: unet_config = configs.get("unet/config.json", {}) - #Graphs are shape-specialized; export 64x64 latents (512x512 images) rather than the config's training sample_size return (1, int(unet_config.get("in_channels", 4)), 64, 64) def _prompt_token_count(configs: dict[str, dict[str, Any]]) -> int: diff --git a/python/cactus/transpiler/Converter/overrides.py b/python/cactus/transpiler/Converter/overrides.py index 04b3851f8..01c2d5999 100644 --- a/python/cactus/transpiler/Converter/overrides.py +++ b/python/cactus/transpiler/Converter/overrides.py @@ -22,8 +22,7 @@ def exportable_bidirectional_mask(config, inputs_embeds, attention_mask=None, an def patch_clip_position_ids_for_export() -> None: import transformers.models.clip.modeling_clip as clip_modeling - #CLIP registers position_ids as a non-persistent buffer, which torch.export lifts - #into a weight placeholder no checkpoint can bind; recompute it in the graph instead. + #CLIP registers position_ids as a non-persistent buffer that no checkpoint can bind; recompute it in the graph original_forward = clip_modeling.CLIPTextEmbeddings.forward def forward_with_computed_position_ids(self, input_ids=None, position_ids=None, inputs_embeds=None): if position_ids is None: diff --git a/python/tests/test_transpile_clip.py b/python/tests/test_transpile_clip.py index 261ebdf37..dd50c6631 100644 --- a/python/tests/test_transpile_clip.py +++ b/python/tests/test_transpile_clip.py @@ -59,7 +59,7 @@ def test_clip_text_encoder_transpiles_to_the_expected_ops(): assert counts["layer_norm"] == 25 assert counts["attention"] == 12 assert counts["embedding"] == 2 - assert counts["sigmoid"] == 12 # quick_gelu lowers as x * sigmoid(1.702 x) + assert counts["sigmoid"] == 12 def test_clip_text_encode_matches_torch_and_reports_latency(): diff --git a/python/tests/test_transpile_sd_unet.py b/python/tests/test_transpile_sd_unet.py index 211665775..c12ef2bd2 100644 --- a/python/tests/test_transpile_sd_unet.py +++ b/python/tests/test_transpile_sd_unet.py @@ -60,8 +60,6 @@ def test_sd_unet_denoise_matches_torch_and_reports_latency(): inputs = _example_inputs() half = tuple(a.to(torch.float16) for a in inputs) - # One capture only -- exporting the 860M UNet takes minutes, so the op - # histogram is asserted on the same capture the numerics run on. captured = capture_model(_denoiser(torch.float16), half) nodes = captured.ir_graph.nodes nodes = list(nodes.values()) if isinstance(nodes, dict) else list(nodes) @@ -69,7 +67,7 @@ def test_sd_unet_denoise_matches_torch_and_reports_latency(): for node in nodes: counts[node.op] = counts.get(node.op, 0) + 1 assert counts["conv2d"] == 98 - assert counts["attention"] == 32 # 16 self + 16 cross + assert counts["attention"] == 32 assert counts["group_norm"] == 61 assert counts["upsample_nearest2d"] == 3 assert counts["silu"] == 68 diff --git a/python/tests/test_transpile_t2i.py b/python/tests/test_transpile_t2i.py index e457d15b7..5511168b9 100644 --- a/python/tests/test_transpile_t2i.py +++ b/python/tests/test_transpile_t2i.py @@ -1,9 +1,3 @@ -"""End-to-end text-to-image through transpiled graphs only. - -CLIP text encoder, LCM UNet, and TAESD decoder each go through -capture_model -> transpile_ir; the LCM scheduler step stays host-side numpy. -The reference is the identical loop run in torch fp32. -""" from __future__ import annotations import time @@ -124,7 +118,6 @@ def test_t2i_pipeline_matches_torch(): noises = [torch.randn(1, 4, 64, 64, generator=generator) for _ in range(STEPS - 1)] w_emb = _guidance_embedding(GUIDANCE - 1.0) - # --- transpiled pipeline --- text_graph = transpile_ir( capture_model(_LastHiddenState(_load("CLIPTextModel", SD_MODEL_ID, "text_encoder", torch.float16)), (ids,)).ir_graph ) @@ -161,7 +154,6 @@ def denoise_cactus(x, t): image = np.asarray(decoder_graph.execute()[0].numpy(), dtype=np.float32).clip(0.0, 1.0) decode_ms = (time.perf_counter() - start) * 1000.0 - # --- torch fp32 reference, same math --- ref_text = _LastHiddenState(_load("CLIPTextModel", SD_MODEL_ID, "text_encoder", torch.float32)) ref_unet = _Denoiser(_load("UNet2DConditionModel", SD_MODEL_ID, "unet", torch.float32)) with torch.no_grad(): From cc4cb469a579d5278eb119b2787ad079ea509f33 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Wed, 12 Aug 2026 14:45:16 -0700 Subject: [PATCH 31/34] add a generate-image command 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 --- python/cactus/bindings/cactus.py | 34 +++++++++++++ python/cactus/cli/__init__.py | 29 +++++++++++ python/cactus/cli/common.py | 1 + python/cactus/cli/generate_image.py | 75 +++++++++++++++++++++++++++++ 4 files changed, 139 insertions(+) create mode 100644 python/cactus/cli/generate_image.py diff --git a/python/cactus/bindings/cactus.py b/python/cactus/bindings/cactus.py index 60805c0a0..7c1c65779 100644 --- a/python/cactus/bindings/cactus.py +++ b/python/cactus/bindings/cactus.py @@ -844,6 +844,13 @@ class cactus_tensor_info_t(ctypes.Structure): ] _lib.cactus_image_embed.restype = ctypes.c_int +_lib.cactus_generate_image.argtypes = [ + ctypes.c_void_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_uint8), + ctypes.c_size_t, ctypes.POINTER(ctypes.c_uint), ctypes.POINTER(ctypes.c_uint), + ctypes.c_int, ctypes.c_float, ctypes.c_ulonglong +] +_lib.cactus_generate_image.restype = ctypes.c_int + _lib.cactus_audio_embed.argtypes = [ ctypes.c_void_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_float), ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t) @@ -1309,6 +1316,33 @@ def cactus_image_embed(model, image_path): return list(buf[:dim.value]) +def cactus_generate_image(model, prompt, steps=4, guidance_scale=8.5, seed=0, max_side=1024): + """Generate an image from a text prompt with a text-to-image model. + + Args: + model: Model handle from cactus_init(). + prompt: The text prompt. + steps: Denoising steps (default 4, an LCM model's sweet spot). + guidance_scale: Classifier-free guidance embedding scale (default 8.5). + seed: Seed for the initial latents. + max_side: Largest image side the result buffer must hold. + + Returns: + (pixels, width, height) where pixels is packed 8-bit RGB, row-major. + """ + buf = (ctypes.c_uint8 * (max_side * max_side * 3))() + width = ctypes.c_uint() + height = ctypes.c_uint() + rc = _lib.cactus_generate_image( + model, _enc(prompt), buf, ctypes.sizeof(buf), ctypes.byref(width), ctypes.byref(height), + int(steps), float(guidance_scale), int(seed)) + if rc == -2: + raise RuntimeError(f"Image is larger than max_side={max_side}; raise it and retry") + if rc < 0: + raise RuntimeError(_err("Image generation failed")) + return bytes(buf[:rc]), width.value, height.value + + def cactus_audio_embed(model, audio_path): """Compute an audio embedding. Returns a list of floats.""" buf = (ctypes.c_float * 4096)() diff --git a/python/cactus/cli/__init__.py b/python/cactus/cli/__init__.py index e841b7d36..21e36e235 100644 --- a/python/cactus/cli/__init__.py +++ b/python/cactus/cli/__init__.py @@ -5,12 +5,14 @@ from .. import __version__ from .common import ( DEFAULT_MODEL_ID, + DEFAULT_IMAGE_MODEL_ID, DEFAULT_TRANSCRIPTION_MODEL_ID, DEFAULT_TEST_MODEL_ID, DEFAULT_TEST_TRANSCRIPTION_MODEL_ID, ) from .download import cmd_download from .compile import cmd_build +from .generate_image import cmd_generate_image from .serve import cmd_serve from .transcribe import cmd_transcribe from .test import cmd_test, cmd_benchmark, COMPONENTS @@ -150,6 +152,16 @@ def create_parser(): --token HuggingFace token (gated models) --reconvert refresh cached bundle + cactus generate-image [model] generate an image (default: {DEFAULT_IMAGE_MODEL_ID}) + --prompt text prompt to render + --out output PNG path (default: image.png) + --steps denoising steps (default: 4) + --guidance-scale guidance scale (default: 8.5) + --seed seed for the initial latents (default: 0) + --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) + --token HuggingFace token (gated models) + --reconvert refresh cached bundle + cactus download [model] fetch a prebuilt bundle (default: {DEFAULT_MODEL_ID}) --bits 1|2|3|4 CQ quantization (default: 4) --token HuggingFace token (gated models) @@ -294,6 +306,22 @@ def create_parser(): transcribe_parser.add_argument("--language", default="en", help="Language code (default: en)") + generate_image_parser = subparsers.add_parser("generate-image", help="Generate an image from a text prompt", + parents=[_telemetry_parent(), _build_parent()]) + generate_image_parser.add_argument("model_id", nargs="?", default=DEFAULT_IMAGE_MODEL_ID, + type=_hf_id_or_path, + help=f"HuggingFace model id (default: {DEFAULT_IMAGE_MODEL_ID})") + generate_image_parser.add_argument("--prompt", required=True, + help="Text prompt to render") + generate_image_parser.add_argument("--out", dest="output", default="image.png", + help="Output PNG path (default: image.png)") + generate_image_parser.add_argument("--steps", type=int, default=4, + help="Denoising steps (default: 4)") + generate_image_parser.add_argument("--guidance-scale", type=float, default=8.5, + help="Guidance scale (default: 8.5)") + generate_image_parser.add_argument("--seed", type=int, default=0, + help="Seed for the initial latents (default: 0)") + serve_parser = subparsers.add_parser("serve", help="OpenAI-compatible local HTTP server", parents=[_telemetry_parent(), _build_parent()]) serve_parser.add_argument("model_id", nargs="?", default=None, @@ -434,6 +462,7 @@ def create_parser(): "run": cmd_run, "serve": cmd_serve, "transcribe": cmd_transcribe, + "generate-image": cmd_generate_image, "test": cmd_test, "benchmark": cmd_benchmark, "list": cmd_list, diff --git a/python/cactus/cli/common.py b/python/cactus/cli/common.py index f39ef435e..0dc606a72 100644 --- a/python/cactus/cli/common.py +++ b/python/cactus/cli/common.py @@ -29,6 +29,7 @@ def weights_root() -> Path: DEFAULT_MODEL_ID = "google/gemma-4-E2B-it" DEFAULT_TRANSCRIPTION_MODEL_ID = "nvidia/parakeet-tdt-0.6b-v3" +DEFAULT_IMAGE_MODEL_ID = "SimianLuo/LCM_Dreamshaper_v7" DEFAULT_TEST_MODEL_ID = DEFAULT_MODEL_ID DEFAULT_TEST_TRANSCRIPTION_MODEL_ID = DEFAULT_TRANSCRIPTION_MODEL_ID diff --git a/python/cactus/cli/generate_image.py b/python/cactus/cli/generate_image.py new file mode 100644 index 000000000..e705b067a --- /dev/null +++ b/python/cactus/cli/generate_image.py @@ -0,0 +1,75 @@ +import json +import struct +import time +import zlib +from pathlib import Path + +from .common import apply_runtime_env, print_color, GREEN, RED + + +def cmd_generate_image(args): + from .model import prepare_bundle + + apply_runtime_env(args) + + bundle_dir = prepare_bundle(args) + if bundle_dir is None: + return 1 + + if bundle_execution_strategy(bundle_dir) != "iterative_denoise": + print_color(RED, f"{args.model_id} is not a text-to-image model") + return 1 + + from ..bindings.cactus import cactus_destroy, cactus_generate_image, cactus_init + + output_path = Path(args.output).expanduser() + print_color(GREEN, f"Generating {args.steps}-step image with model: {args.model_id}") + print() + + model = cactus_init(str(bundle_dir)) + try: + started = time.perf_counter() + pixels, width, height = cactus_generate_image( + model, + args.prompt, + steps=args.steps, + guidance_scale=args.guidance_scale, + seed=args.seed, + ) + elapsed = time.perf_counter() - started + except RuntimeError as exc: + print_color(RED, str(exc)) + return 1 + finally: + cactus_destroy(model) + + write_png(output_path, pixels, width, height) + print_color(GREEN, f"{width}x{height} image written to {output_path} in {elapsed:.1f}s") + return 0 + + +def bundle_execution_strategy(bundle_dir: Path) -> str: + manifest = Path(bundle_dir) / "components" / "manifest.json" + if not manifest.is_file(): + return "" + try: + return json.loads(manifest.read_text(encoding="utf-8")).get("runtime_execution_strategy", "") + except (OSError, ValueError): + return "" + + +def write_png(path: Path, pixels: bytes, width: int, height: int) -> None: + stride = width * 3 + raw = b"".join(b"\x00" + pixels[row * stride:(row + 1) * stride] for row in range(height)) + + def chunk(tag: bytes, data: bytes) -> bytes: + return (struct.pack(">I", len(data)) + tag + data + + struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF)) + + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes( + b"\x89PNG\r\n\x1a\n" + + chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0)) + + chunk(b"IDAT", zlib.compress(raw, 6)) + + chunk(b"IEND", b"") + ) From 2ea22e8a809eb901511d217453acd9500a676b20 Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Wed, 12 Aug 2026 15:26:10 -0700 Subject: [PATCH 32/34] keep attention whole through export so the fusion can match it 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 --- python/cactus/transpiler/Converter/models.py | 18 +++++++++++++++++- .../Fusions/fusion_defs/fusion_attention.py | 7 ++++++- python/cactus/transpiler/IR/constants.py | 4 ++++ python/cactus/transpiler/IR/models.py | 2 ++ .../cactus/transpiler/ModelProfiles/models.py | 3 +++ .../transpiler/ModelProfiles/profiles.py | 3 ++- 6 files changed, 34 insertions(+), 3 deletions(-) diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index a985e12be..afd76245e 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -463,6 +463,20 @@ def pad_no_cache_full_context_input(input_: Input, input_cls: Any, capacity: int inference_mode=input_.inference_mode, ) +def preserved_decomposition_table(mp: MP_Models.ModelProfile, inference_mode: str) -> dict[Any, Any] | None: + preserved = tuple(mp.preserved_ops) + for source in mp.component_sources: + if source.mode == inference_mode: + preserved += tuple(source.preserved_ops) + if not preserved: + return None + table = torch.export.default_decompositions() + for name in preserved: + overload_packet = getattr(torch.ops.aten, name, None) + if overload_packet is not None: + table.pop(overload_packet.default, None) + return table + def export_(model: Model, input: Input) -> LayerMap: should_use_cache = input.inference_mode in constants.CACHE_INFERENCE_MODES CU.configure_model_for_export(model.model, should_use_cache=should_use_cache) @@ -481,8 +495,10 @@ def export_(model: Model, input: Input) -> LayerMap: OV.prepare_model_input_hints_for_export(model.model, export_kwargs) export_kwargs["use_cache"] = should_use_cache export_kwargs = CU.filter_forward_kwargs(export_model, export_kwargs) + decomposition_table = preserved_decomposition_table(model.model_profile, input.inference_mode) with torch.no_grad(): - exported = torch.export.export(export_model, args=input.args, kwargs=export_kwargs, strict=False).run_decompositions() + exported = torch.export.export(export_model, args=input.args, kwargs=export_kwargs, strict=False) + exported = exported.run_decompositions() if decomposition_table is None else exported.run_decompositions(decomposition_table) graph = exported.graph_module.graph if hasattr(exported, "graph_module") else exported.graph records = [LayerRecord.from_node(i, node) for i, node in enumerate(graph.nodes)] return LayerMap.from_data(x=exported, name=model.name, model_task=input.inference_mode, nodes_list=records) diff --git a/python/cactus/transpiler/Fusions/fusion_defs/fusion_attention.py b/python/cactus/transpiler/Fusions/fusion_defs/fusion_attention.py index dd78635f4..ed288dd50 100644 --- a/python/cactus/transpiler/Fusions/fusion_defs/fusion_attention.py +++ b/python/cactus/transpiler/Fusions/fusion_defs/fusion_attention.py @@ -9,7 +9,12 @@ "attn_sdp", ("attn_sdp",), inputs=(_input("query", "attn_sdp", 0), _input("key", "attn_sdp", 1), _input("value", "attn_sdp", 2), _input("mask", "attn_sdp", 3, optional=True)), - attr_captures=(M.AttrCapture("scale", "attn_sdp", "scale", required=False), M.AttrCapture("is_causal", "attn_sdp", "is_causal", default=True, required=False)), + attr_captures=( + M.AttrCapture("scale", "attn_sdp", "scale", required=False), + M.AttrCapture("is_causal", "attn_sdp", "is_causal", default=False, required=False), + M.AttrCapture("input_layout", default="bhqd_bhsd_bhsd", required=False), + M.AttrCapture("output_layout", default="bhqd", required=False), + ), ) ATTENTION_CORE_GRAPH = _graph( diff --git a/python/cactus/transpiler/IR/constants.py b/python/cactus/transpiler/IR/constants.py index cb6fb44f5..f1658bee2 100644 --- a/python/cactus/transpiler/IR/constants.py +++ b/python/cactus/transpiler/IR/constants.py @@ -110,3 +110,7 @@ "aten.fft_rfft.default": {1: "n", 2: "dim", 3: "norm"}, "aten.fft_irfft.default": {1: "n", 2: "dim", 3: "norm"}, } + +OMITTED_ATTR_DEFAULTS: dict[str, dict[str, object]] = { + "aten.scaled_dot_product_attention.default": {"is_causal": False}, +} diff --git a/python/cactus/transpiler/IR/models.py b/python/cactus/transpiler/IR/models.py index 2f2480c2e..8d67dc5c2 100644 --- a/python/cactus/transpiler/IR/models.py +++ b/python/cactus/transpiler/IR/models.py @@ -229,6 +229,8 @@ def extract_attrs(record: CModels.LayerRecord) -> dict[str, Any]: if extract_node_refs(value): continue attrs.setdefault(f"arg_{index}", value) + for attr_name, value in constants.OMITTED_ATTR_DEFAULTS.get(record.target, {}).items(): + attrs.setdefault(attr_name, value) return attrs def generate_node(record: CModels.LayerRecord) -> Node: diff --git a/python/cactus/transpiler/ModelProfiles/models.py b/python/cactus/transpiler/ModelProfiles/models.py index 120d56055..451a56ca2 100644 --- a/python/cactus/transpiler/ModelProfiles/models.py +++ b/python/cactus/transpiler/ModelProfiles/models.py @@ -33,6 +33,7 @@ class ComponentSource: mode:str #Export mode this component is captured under load_strategy:str #Key into Converter LOAD_STRATEGIES/SYNTHETIC_INPUT_BUILDERS for this component source:str = "" #Repository subfolder, or a full repository id when the component ships separately + preserved_ops:tuple[str, ...] = () #Aten ops this component keeps whole through run_decompositions @dataclass(slots=True, frozen=True) class PromptContract: @@ -133,7 +134,9 @@ class ModelProfile: input_strategy: str #Specifies what functions/procedures to use when generating sample input for torch.export export_patches: tuple[str, ...] #Specifies what patches/masks to apply to sample inputs load_strategy: str #Specifies which Hugging Face AutoModel class/loading path to prefer + preserved_ops: tuple[str, ...] = () #Aten ops kept whole through run_decompositions so a fusion can match them directly export_modes: tuple[str, ...] = () #Per-component export modes when the family is not prefill/decode shaped; empty means the default modes + component_sources: tuple[ComponentSource, ...] = () #Components of a pipeline model that load and export separately disabled_fusion_fields: tuple[str, ...] = () #Fusion groups to skip for this profile when a family is known unsafe disabled_fusions: tuple[str, ...] = () #Individual fusion names or cactus ops to skip for this profile diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index 2a29ffa78..3256200bb 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -611,7 +611,8 @@ def decoder_kv_state(producer: str = "decoder_prefill_chunk") -> StateContract: load_strategy="auto", component_sources=( ComponentSource(mode="text_encoder", load_strategy="clip_text", source="text_encoder"), - ComponentSource(mode="unet", load_strategy="sd_unet", source="unet"), + ComponentSource(mode="unet", load_strategy="sd_unet", source="unet", + preserved_ops=("scaled_dot_product_attention",)), ComponentSource(mode="vae_decoder", load_strategy="taesd_decoder", source="madebyollin/taesd"), ), prompt_contract=PromptContract(style="raw", template_source="none"), From c6d1e5b59801e9aa3b58b2559a9b1819d54145bf Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Wed, 12 Aug 2026 16:32:00 -0700 Subject: [PATCH 33/34] normalize the attention mask rank so masked attention can fuse 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 --- .../Generator/lowerings/lowering_cache.py | 15 ++++++++++++++- .../cactus/transpiler/ModelProfiles/profiles.py | 3 ++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/python/cactus/transpiler/Generator/lowerings/lowering_cache.py b/python/cactus/transpiler/Generator/lowerings/lowering_cache.py index b97dc5bab..4eedc56d0 100644 --- a/python/cactus/transpiler/Generator/lowerings/lowering_cache.py +++ b/python/cactus/transpiler/Generator/lowerings/lowering_cache.py @@ -180,7 +180,7 @@ def lower_attention(context: models.GenerationContext, node: IRModels.Node) -> A query = cast_to_precision(context, query, context.graph.FP16) key = cast_to_precision(context, key, context.graph.FP16) value = cast_to_precision(context, value, context.graph.FP16) - mask = cast_to_precision(context, mask, context.graph.FP16) if mask is not None else None + mask = attention_mask_for_kernel(context, mask, query) if mask is not None else None output = context.graph.attention( query, key, @@ -380,6 +380,19 @@ def attention_inputs_for_layout(context: models.GenerationContext, node: IRModel ) return inputs[0], inputs[1], inputs[2], mask +def attention_mask_for_kernel(context: models.GenerationContext, mask: Any, query: Any) -> Any: + mask = cast_to_precision(context, mask, context.graph.FP16) + mask_shape = tuple(int(dim) for dim in getattr(mask, "shape", ())) + query_shape = tuple(int(dim) for dim in getattr(query, "shape", ())) + if len(query_shape) != 4: + return mask + batch_size, seq_len = query_shape[0], query_shape[1] + if len(mask_shape) == 2 and batch_size == 1 and mask_shape[0] == seq_len: + return context.graph.reshape(mask, (1, mask_shape[0], mask_shape[1])) + if len(mask_shape) == 4 and mask_shape[0] == batch_size and mask_shape[1] == 1 and mask_shape[2] == seq_len: + return context.graph.reshape(mask, (mask_shape[0], mask_shape[2], mask_shape[3])) + return mask + def attention_output_for_layout(context: models.GenerationContext, node: IRModels.Node, output: Any) -> Any: output_layout = node.attrs.get("output_layout") if output_layout == "bthd_flat": diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index 3256200bb..08723db30 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -610,7 +610,8 @@ def decoder_kv_state(producer: str = "decoder_prefill_chunk") -> StateContract: export_patches=("clip_position_ids",), load_strategy="auto", component_sources=( - ComponentSource(mode="text_encoder", load_strategy="clip_text", source="text_encoder"), + ComponentSource(mode="text_encoder", load_strategy="clip_text", source="text_encoder", + preserved_ops=("scaled_dot_product_attention",)), ComponentSource(mode="unet", load_strategy="sd_unet", source="unet", preserved_ops=("scaled_dot_product_attention",)), ComponentSource(mode="vae_decoder", load_strategy="taesd_decoder", source="madebyollin/taesd"), From 2ba9f46c367e62e31157cf362a6dceda291e114f Mon Sep 17 00:00:00 2001 From: Neel Shejwalkar Date: Wed, 12 Aug 2026 17:25:35 -0700 Subject: [PATCH 34/34] pin the preserved-attention contract and tighten the conv discriminator 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 --- .../convert/cactus_adapters/tokenizer.py | 3 +- python/cactus/transpiler/Converter/models.py | 12 ++-- python/cactus/transpiler/Fusions/fusions.py | 2 +- .../cactus/transpiler/ModelProfiles/models.py | 1 - .../tests/test_transpiler_attention_fusion.py | 60 +++++++++++++++++++ 5 files changed, 68 insertions(+), 10 deletions(-) create mode 100644 python/tests/test_transpiler_attention_fusion.py diff --git a/python/cactus/convert/cactus_adapters/tokenizer.py b/python/cactus/convert/cactus_adapters/tokenizer.py index 96bbfbdbf..f18f977da 100644 --- a/python/cactus/convert/cactus_adapters/tokenizer.py +++ b/python/cactus/convert/cactus_adapters/tokenizer.py @@ -1,4 +1,5 @@ import json +import shutil from pathlib import Path try: @@ -610,8 +611,6 @@ def _clip_special_token(tokenizer_config, special_tokens_map, name, default): def convert_clip_tokenizer(tokenizer_dir, output_dir): tokenizer_dir, output_dir = Path(tokenizer_dir), Path(output_dir) - import shutil - vocab = json.loads((tokenizer_dir / "vocab.json").read_text(encoding="utf-8")) tokenizer_config = json.loads((tokenizer_dir / "tokenizer_config.json").read_text(encoding="utf-8")) special_tokens_map = json.loads((tokenizer_dir / "special_tokens_map.json").read_text(encoding="utf-8")) diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index afd76245e..1002347aa 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -381,13 +381,13 @@ def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None, model_cl {**base_kwargs, "torch_dtype": "auto", "low_cpu_mem_usage": True}, base_kwargs, ) - for model_class in candidate_classes: - if model_class in seen_classes: + for candidate_class in candidate_classes: + if candidate_class in seen_classes: continue - seen_classes.add(model_class) - for load_kwargs in load_attempts: + seen_classes.add(candidate_class) + for attempt_kwargs in load_attempts: try: - model = model_class.from_pretrained(model_id, **load_kwargs) + model = candidate_class.from_pretrained(model_id, **attempt_kwargs) model.eval() for patch in export_patches: patch_fn = EXPORT_PATCHES[patch] @@ -464,7 +464,7 @@ def pad_no_cache_full_context_input(input_: Input, input_cls: Any, capacity: int ) def preserved_decomposition_table(mp: MP_Models.ModelProfile, inference_mode: str) -> dict[Any, Any] | None: - preserved = tuple(mp.preserved_ops) + preserved: tuple[str, ...] = () for source in mp.component_sources: if source.mode == inference_mode: preserved += tuple(source.preserved_ops) diff --git a/python/cactus/transpiler/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py index 35f0f05fb..2a2a8b140 100644 --- a/python/cactus/transpiler/Fusions/fusions.py +++ b/python/cactus/transpiler/Fusions/fusions.py @@ -100,7 +100,7 @@ "lfm_short_conv_decode": _definition("lfm_short_conv_decode", "lfm_short_conv_decode", LFM_SHORT_CONV_DECODE_GRAPH, fusion_fields=("generic", "conv", "cache", "lfm_moe"), supported_inference_modes=("decode_with_cache",)), "conv1d_k3": _definition("conv1d_k3", "conv1d_k3", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=3, stride=1, padding=0, dilation=1, groups=1, spatial_dims=1)), "conv1d_k7s3": _definition("conv1d_k7s3", "conv1d_k7s3", CONV_BIAS_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=7, stride=3, spatial_dims=1)), - "conv1d_causal": _definition("conv1d_causal", "conv1d_causal", CONV_GRAPH, fusion_fields=("generic", "conv", "cache"), metadata=_required_attrs(causal=True)), + "conv1d_causal": _definition("conv1d_causal", "conv1d_causal", CONV_GRAPH, fusion_fields=("generic", "conv", "cache"), metadata=_required_attrs(causal=True, spatial_dims=1)), "conv1d_pointwise": _definition("conv1d_pointwise", "conv1d_pointwise", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=1, spatial_dims=1)), "conv1d_same_depthwise_k9": _definition("conv1d_same_depthwise_k9", "conv1d_same_depthwise_k9", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=9, depthwise=True, spatial_dims=1)), "conv2d_k3s2p1": _definition("conv2d_k3s2p1", "conv2d_k3s2p1", CONV_GRAPH, fusion_fields=("generic", "conv", "vision"), metadata=_required_attrs(kernel_size=3, stride=2, padding=1, spatial_dims=2)), diff --git a/python/cactus/transpiler/ModelProfiles/models.py b/python/cactus/transpiler/ModelProfiles/models.py index 451a56ca2..92e57a736 100644 --- a/python/cactus/transpiler/ModelProfiles/models.py +++ b/python/cactus/transpiler/ModelProfiles/models.py @@ -134,7 +134,6 @@ class ModelProfile: input_strategy: str #Specifies what functions/procedures to use when generating sample input for torch.export export_patches: tuple[str, ...] #Specifies what patches/masks to apply to sample inputs load_strategy: str #Specifies which Hugging Face AutoModel class/loading path to prefer - preserved_ops: tuple[str, ...] = () #Aten ops kept whole through run_decompositions so a fusion can match them directly export_modes: tuple[str, ...] = () #Per-component export modes when the family is not prefill/decode shaped; empty means the default modes component_sources: tuple[ComponentSource, ...] = () #Components of a pipeline model that load and export separately diff --git a/python/tests/test_transpiler_attention_fusion.py b/python/tests/test_transpiler_attention_fusion.py new file mode 100644 index 000000000..d0fffcf3b --- /dev/null +++ b/python/tests/test_transpiler_attention_fusion.py @@ -0,0 +1,60 @@ +import json +import tempfile +import unittest +from pathlib import Path + +import torch + +from cactus.transpiler.Converter import models as CModels +from cactus.transpiler.IR import simplify_ir + +FUSION_FIELDS = ("generic", "attention") + + +class _Sdpa(torch.nn.Module): + def __init__(self, **kwargs): + super().__init__() + self.kwargs = kwargs + + def forward(self, query, key, value): + return torch.nn.functional.scaled_dot_product_attention(query, key, value, **self.kwargs) + + +def _fused_attention_nodes(module): + args = (torch.randn(1, 8, 32, 64), torch.randn(1, 8, 32, 64), torch.randn(1, 8, 32, 64)) + table = torch.export.default_decompositions() + table.pop(torch.ops.aten.scaled_dot_product_attention.default, None) + exported = torch.export.export(module.eval(), args, strict=False).run_decompositions(table) + records = [CModels.LayerRecord.from_node(i, node) for i, node in enumerate(exported.graph.nodes)] + layer_map = CModels.LayerMap.from_data( + x=exported, name="sdpa-fusion-test", model_task="prefill_no_cache", nodes_list=records + ) + with tempfile.TemporaryDirectory() as tmp: + simplified = Path(tmp) / "simplified.json" + simplify_ir.write_simplified_json( + layer_map, simplified, input_modalities=("text",), fusion_fields=FUSION_FIELDS, + disabled_fusion_fields=(), disabled_fusions=(), + ) + nodes = json.loads(simplified.read_text(encoding="utf-8"))["nodes"] + return [node for node in nodes if node["target"] == "cactus.attention"] + + +class TestAttentionDirectFusion(unittest.TestCase): + def test_preserved_attention_fuses_with_the_aten_layout(self): + fused = _fused_attention_nodes(_Sdpa()) + self.assertEqual(len(fused), 1) + kwargs = fused[0]["kwargs"] + self.assertEqual(kwargs["input_layout"], "bhqd_bhsd_bhsd") + self.assertEqual(kwargs["output_layout"], "bhqd") + + def test_omitted_is_causal_follows_the_aten_default(self): + fused = _fused_attention_nodes(_Sdpa()) + self.assertIs(fused[0]["kwargs"]["is_causal"], False) + + def test_explicit_is_causal_is_carried_through(self): + fused = _fused_attention_nodes(_Sdpa(is_causal=True)) + self.assertIs(fused[0]["kwargs"]["is_causal"], True) + + +if __name__ == "__main__": + unittest.main()