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..8d3c59720 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,72 @@ std::vector BPETokenizer::apply_bpe(const std::vector& return current_tokens; } +std::vector BPETokenizer::clip_word_split(const std::string& text) const { + 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 +535,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 +589,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..795e1466e 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; @@ -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); @@ -485,6 +501,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 +704,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 +1059,9 @@ 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; + 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..246dda746 --- /dev/null +++ b/cactus-engine/src/generate_image.cpp @@ -0,0 +1,252 @@ +#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; + +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..1cb364a13 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -173,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; @@ -766,6 +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 (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"; @@ -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(); } @@ -2440,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(); @@ -3281,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) { @@ -3312,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; @@ -5550,6 +5565,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_t2i") 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..228fe0efd --- /dev/null +++ b/cactus-engine/tests/test_diffusion.cpp @@ -0,0 +1,136 @@ +#include "test_utils.h" +#include "../src/engine.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_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) { + 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() { + TestUtils::TestRunner runner("Diffusion Tests"); + 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; +} diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 9475f4293..f7c401bd1 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, @@ -1237,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/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/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/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/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/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-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_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; } 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..f00396d91 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)]], @@ -2058,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/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..ed1bce976 100644 --- a/cactus-kernels/src/metal_backend.mm +++ b/cactus-kernels/src/metal_backend.mm @@ -59,8 +59,8 @@ 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 psoBilinear=nil, psoConv1dGen=nil, psoConv1dNlcDw=nil, psoConv2d=nil; + id psoConcat2=nil, psoGatherF32Idx=nil, psoRopeFull=nil, psoMaxpool1d=nil, psoUpsampleNearest2d=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; @@ -122,8 +122,10 @@ 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"); + 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"); @@ -2067,6 +2069,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; @@ -2121,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]; 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/cactus-kernels/tests/test_conv.cpp b/cactus-kernels/tests/test_conv.cpp index c6ae32682..4435ea4cc 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; @@ -169,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; } @@ -179,6 +236,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(); 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 diff --git a/python/cactus/bindings/cactus.py b/python/cactus/bindings/cactus.py index fc9115713..7c1c65779 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", [ @@ -839,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) @@ -1304,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)() @@ -2135,6 +2174,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/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"") + ) diff --git a/python/cactus/cli/model.py b/python/cactus/cli/model.py index fd930632e..750d5da9b 100644 --- a/python/cactus/cli/model.py +++ b/python/cactus/cli/model.py @@ -47,9 +47,22 @@ def _convert_from_source(model_id, *, bits, token, weights_dir, skip_model_load= err = convert_toolchain_error() if err: raise RuntimeError(err) - print_color(YELLOW, f"Converting {model_id} from HuggingFace source...") + if token: + os.environ["HF_TOKEN"] = token + os.environ["HUGGING_FACE_HUB_TOKEN"] = token + from .transpiler import profile_for_model_id + + 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), @@ -57,15 +70,79 @@ 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, profile, *, bits, token, weights_dir): + """Convert a pipeline model one component at a time into a single weights dir.""" + import json + import tempfile + + 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 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)]) + 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 {source.mode!r} both produce tensor file {tensor_file.name!r}" + ) + shutil.move(str(tensor_file), destination) + converted.append(source.mode) + (weights_dir / "weights_manifest.json").write_text( + json.dumps({"weights": merged_records}, indent=2), 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 _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 a pipeline bundle needs at runtime.""" + 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)) + 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): from .download import get_bundle_dir diff --git a/python/cactus/cli/transpiler.py b/python/cactus/cli/transpiler.py index d332c24f4..1f0cee6f6 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,12 +168,34 @@ def build_transpiled_bundle( return bundle_path +def diffusion_scheduler_metadata(weights_dir: Path) -> dict[str, str]: + 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 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, *, @@ -200,6 +225,7 @@ def resolve_transpile_config( profile=profile, modalities=default_modalities(profile), profile_source="registered", + 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..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: @@ -599,3 +600,58 @@ 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) + 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/cli.py b/python/cactus/convert/cli.py index 60a85bce8..151522f00 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,6 +138,9 @@ def _bits_for_component(component: str, args: argparse.Namespace) -> int: return int(args.bits) +CHECKPOINT_STEMS = ("model", "diffusion_pytorch_model") + + 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: @@ -145,11 +150,16 @@ def _load_checkpoint_state_dict(model_id_or_path: str) -> dict[str, Any] | None: try: from huggingface_hub import hf_hub_download from safetensors.torch import load_file - try: - model_file = hf_hub_download(model_id_or_path, "model.safetensors", cache_dir=_hf_cache_dir()) - return load_file(model_file) - except Exception: - index_file = hf_hub_download(model_id_or_path, "model.safetensors.index.json", cache_dir=_hf_cache_dir()) + 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: @@ -162,6 +172,7 @@ def _load_checkpoint_state_dict(model_id_or_path: str) -> dict[str, Any] | None: raise RuntimeError(f"duplicate tensor key {key!r} across checkpoint shards") state[key] = tensor return state + return None except RuntimeError: raise except Exception: @@ -457,6 +468,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), @@ -523,6 +535,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) diff --git a/python/cactus/convert/model_adapters/adapters.py b/python/cactus/convert/model_adapters/adapters.py index 517e4fa5c..ee4df1961 100644 --- a/python/cactus/convert/model_adapters/adapters.py +++ b/python/cactus/convert/model_adapters/adapters.py @@ -502,6 +502,93 @@ def build_calibration_inputs(self, row: dict[str, Any], processor: Any, modality return None +def _verbatim_name_tensor(source_name: str, component: str) -> NameMatch: + 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" + + 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: + 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", "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" @@ -884,6 +971,9 @@ def scale_factor(self, output_name: str) -> float: "whisper": WhisperAdapter(), "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 69cb2aa09..a5e78928a 100644 --- a/python/cactus/convert/model_adapters/detection.py +++ b/python/cactus/convert/model_adapters/detection.py @@ -3,12 +3,18 @@ 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", "sd_unet", "clip_text", "generic"} def detect_family(config, requested: str = "auto") -> str: if requested != "auto": 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" diff --git a/python/cactus/py.typed b/python/cactus/py.typed new file mode 100644 index 000000000..e69de29bb 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..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", @@ -1783,6 +1784,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 +1980,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..909f43da1 100644 --- a/python/cactus/transpile/lower.py +++ b/python/cactus/transpile/lower.py @@ -2153,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]) 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/cactus/transpiler/Converter/diffusion.py b/python/cactus/transpiler/Converter/diffusion.py new file mode 100644 index 000000000..28f779d5e --- /dev/null +++ b/python/cactus/transpiler/Converter/diffusion.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from typing import Any + +import torch + + +class TextEncoderLastHiddenState(torch.nn.Module): + #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 + + 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, autoencoder: torch.nn.Module): + super().__init__() + self.decoder = autoencoder.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 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 + + +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..bdb274187 100644 --- a/python/cactus/transpiler/Converter/input_processor.py +++ b/python/cactus/transpiler/Converter/input_processor.py @@ -91,6 +91,31 @@ 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", {}) + 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 +124,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 64ee5b6ce..1002347aa 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -13,10 +13,12 @@ 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 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, @@ -185,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, @@ -343,33 +354,40 @@ 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, - ) - for model_class in candidate_classes: - if model_class in seen_classes: + 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 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] @@ -380,7 +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.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) + +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}") @@ -424,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[str, ...] = () + 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) @@ -442,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/Converter/overrides.py b/python/cactus/transpiler/Converter/overrides.py index 40deb13be..01c2d5999 100644 --- a/python/cactus/transpiler/Converter/overrides.py +++ b/python/cactus/transpiler/Converter/overrides.py @@ -20,6 +20,18 @@ 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 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: + 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/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/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py index b71be4a27..2a2a8b140 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_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_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, 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)), + "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/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_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/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/Generator/models.py b/python/cactus/transpiler/Generator/models.py index 274f164d0..4947851ea 100644 --- a/python/cactus/transpiler/Generator/models.py +++ b/python/cactus/transpiler/Generator/models.py @@ -434,12 +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 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 diff --git a/python/cactus/transpiler/IR/constants.py b/python/cactus/transpiler/IR/constants.py index 11b59ec5e..f1658bee2 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"}, @@ -108,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/match_utils.py b/python/cactus/transpiler/IR/match_utils.py index 6eb9068be..64589cb2d 100644 --- a/python/cactus/transpiler/IR/match_utils.py +++ b/python/cactus/transpiler/IR/match_utils.py @@ -368,6 +368,7 @@ 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 + 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) 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/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..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, @@ -42,3 +45,9 @@ 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/models.py b/python/cactus/transpiler/ModelProfiles/models.py index 54586e306..92e57a736 100644 --- a/python/cactus/transpiler/ModelProfiles/models.py +++ b/python/cactus/transpiler/ModelProfiles/models.py @@ -28,6 +28,13 @@ 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 + preserved_ops:tuple[str, ...] = () #Aten ops this component keeps whole through run_decompositions + @dataclass(slots=True, frozen=True) class PromptContract: style: str = "auto" @@ -127,6 +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 + 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 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..08723db30 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -4,12 +4,14 @@ MULTIMODAL_COMPONENTS, SPEECH_SEQ2SEQ_COMPONENTS, SPEECH_TRANSCRIBER_COMPONENTS, + T2I_COMPONENTS, TEXT_COMPONENTS, VISION_LANGUAGE_COMPONENTS, ) from .models import ( AliasContract, CacheContract, + ComponentSource, MediaContract, ModelProfile, PromptContract, @@ -27,6 +29,7 @@ LFM_VLM_INFERENCE_PATTERNS, PARAKEET_INFERENCE_PATTERNS, QWEN2_5_INFERENCE_PATTERNS, + SD15_T2I_INFERENCE_PATTERNS, WHISPER_INFERENCE_PATTERNS, ) @@ -204,6 +207,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, @@ -566,6 +589,37 @@ def decoder_kv_state(producer: str = "decoder_prefill_chunk") -> StateContract: runtime_contract=WHISPER_RUNTIME_CONTRACT, ) +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="synthetic", + export_patches=("clip_position_ids",), + load_strategy="auto", + component_sources=( + 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"), + ), + 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,8 +630,25 @@ 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 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/cactus/transpiler/ModelProfiles/routes.py b/python/cactus/transpiler/ModelProfiles/routes.py index f91f6cc10..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, @@ -45,3 +47,7 @@ TEXT_PREFILL.name: TEXT_PREFILL, TEXT_DECODE.name: TEXT_DECODE, } + +SD15_T2I_INFERENCE_PATTERNS = { + T2I_GENERATE.name: T2I_GENERATE, +} diff --git a/python/tests/test_transpile_clip.py b/python/tests/test_transpile_clip.py new file mode 100644 index 000000000..dd50c6631 --- /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 + + +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}%" diff --git a/python/tests/test_transpile_sd_unet.py b/python/tests/test_transpile_sd_unet.py new file mode 100644 index 000000000..c12ef2bd2 --- /dev/null +++ b/python/tests/test_transpile_sd_unet.py @@ -0,0 +1,97 @@ +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) + + 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 + 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}%" diff --git a/python/tests/test_transpile_t2i.py b/python/tests/test_transpile_t2i.py new file mode 100644 index 000000000..5511168b9 --- /dev/null +++ b/python/tests/test_transpile_t2i.py @@ -0,0 +1,179 @@ +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) + + 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 + + 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" diff --git a/python/tests/test_transpile_taesd.py b/python/tests/test_transpile_taesd.py new file mode 100644 index 000000000..0b1cf3cff --- /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 + +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}" 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 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() diff --git a/python/tests/test_transpiler_sd_profile.py b/python/tests/test_transpiler_sd_profile.py new file mode 100644 index 000000000..e79e184c7 --- /dev/null +++ b/python/tests/test_transpiler_sd_profile.py @@ -0,0 +1,121 @@ +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.Converter import input_processor +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_resolve_to_a_repository_and_subfolder(self): + profile = profiles.profile_for_model_id(MODEL_ID) + 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( + profiles.component_repo_and_subfolder(vae, MODEL_ID), ("madebyollin/taesd", "") + ) + with self.assertRaises(ValueError): + 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_synthetic_inputs_follow_the_configs(self): + configs = { + "unet/config.json": { + "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 = input_processor.clip_text_inputs(configs) + self.assertEqual(ids["input_ids"].shape, (1, 77)) + self.assertEqual(ids["input_ids"].dtype, torch.int64) + + 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 = input_processor.taesd_decoder_inputs(configs) + self.assertEqual(latent["x"].shape, (1, 4, 64, 64)) + + 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() 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()