#include "diffusion_engine.h"

#include <algorithm>
#include <cinttypes>
#include <cmath>
#include <cstdlib>
#include <list>
#include <mutex>
#include <set>
#include <tuple>
#include <type_traits>
#include <unordered_set>
#include <utility>
#include <vector>

#include "core/ggml_extend_backend.h"
#include "core/ggml_graph_cut.h"
#include "core/ggml_runner.h"
#include "core/ggml_tensor_utils.h"
#include "core/layer_split_partition.h"
#include "core/probe/conv_direct.h"
#include "core/probe/flash_attn.h"
#include "model.h"
#include "model_definition.h"

#include "core/rng.hpp"
#include "core/rng_mt19937.hpp"
#include "core/rng_philox.hpp"
#include "core/util.h"
#include "model_builders.h"
#include "model_loader.h"
#include "model_manager.h"
#include "stable-diffusion.h"

#include "conditioning/conditioner.hpp"
#include "conditioning/conditioning_cache.h"
#include "core/backend_fit.h"
#include "extensions/generation_extension.h"
#include "model/adapter/ip_adapter.hpp"
#include "model/adapter/lora.hpp"
#include "model/audio/wav2vec2.hpp"
#include "model/diffusion/animatediff.hpp"
#include "model/diffusion/control.hpp"
#include "model/diffusion/model.hpp"
#include "model/diffusion/z_image.hpp"
#include "model/vae/audio_vae.hpp"
#include "model/vae/ltx_vae.hpp"
#include "model/vae/vae.hpp"
#include "runtime/audio_processing.h"
#include "runtime/denoiser.hpp"
#include "runtime/guidance.h"
#include "runtime/preview_interval.h"
#include "runtime/sample-cache.h"

#include "name_conversion.h"
#include "runtime/latent-preview.h"

#include <atomic>

void calculate_alphas_cumprod(float* alphas_cumprod,
                              float linear_start = 0.00085f,
                              float linear_end   = 0.0120f,
                              int timesteps      = TIMESTEPS) {
    float ls_sqrt = sqrtf(linear_start);
    float le_sqrt = sqrtf(linear_end);
    float amount  = le_sqrt - ls_sqrt;
    float product = 1.0f;
    for (int i = 0; i < timesteps; i++) {
        float beta = ls_sqrt + amount * ((float)i / (timesteps - 1));
        product *= 1.0f - powf(beta, 2.0f);
        alphas_cumprod[i] = product;
    }
}

void calculate_alphas_cumprod_linear_beta(float* alphas_cumprod,
                                          float beta_start,
                                          float beta_end,
                                          int timesteps = TIMESTEPS) {
    float product = 1.0f;
    for (int i = 0; i < timesteps; i++) {
        float beta = beta_start + (beta_end - beta_start) * ((float)i / (timesteps - 1));
        product *= 1.0f - beta;
        alphas_cumprod[i] = product;
    }
}

template <typename T, typename = void>
struct has_set_runtime_backends : std::false_type {};
template <typename T>
struct has_set_runtime_backends<T,
                                std::void_t<decltype(std::declval<T&>().set_runtime_backends(
                                    std::declval<const std::vector<ggml_backend_t>&>()))>> : std::true_type {};

static_assert(std::atomic<sd_cancel_mode_t>::is_always_lock_free,
              "sd_cancel_mode_t must be lock-free");

StableDiffusionGGML::StableDiffusionGGML()
    : rng(std::make_shared<PhiloxRNG>()),
      conditioning_cache_(std::make_unique<ConditioningCache>()),
      denoiser(std::make_shared<CompVisDenoiser>()) {}

StableDiffusionGGML::~StableDiffusionGGML() = default;

const std::map<StableDiffusionGGML::RunnerGroup, std::set<ModelComponent>>& StableDiffusionGGML::runner_components() {
    static const std::map<RunnerGroup, std::set<ModelComponent>> components{
        {RunnerGroup::Core, {ModelComponent::Conditioner, ModelComponent::Diffusion, ModelComponent::HighNoiseDiffusion, ModelComponent::CLIPVision, ModelComponent::IPAdapter, ModelComponent::AudioEncoder}},
        {RunnerGroup::VAE, {ModelComponent::VAE, ModelComponent::PreviewVAE, ModelComponent::AudioVAE}},
        {RunnerGroup::ControlNet, {ModelComponent::ControlNet}},
        {RunnerGroup::Extensions, {ModelComponent::PhotoMaker, ModelComponent::PuLID}},
    };
    return components;
}

StableDiffusionGGML::RunnerGroups StableDiffusionGGML::all_runner_groups() {
    RunnerGroups groups;
    for (const auto& entry : runner_components()) {
        groups.insert(entry.first);
    }
    return groups;
}

ModelLoader::FileVersions StableDiffusionGGML::runner_source_versions(RunnerGroup group, const ModelLoader& loader) const {
    auto sources = model_manager->source_versions(runner_components().at(group), loader);
    if (group == RunnerGroup::Core) {
        // PhotoMaker's LoRA may already be merged into resident core weights.
        auto extra = loader.file_versions({"alphas_cumprod", "v_pred", "edm_vpred.", "pmid."});
        sources.insert(extra.begin(), extra.end());
    }
    return sources;
}

void StableDiffusionGGML::capture_runner_sources() {
    RunnerState state;
    state.catalog_revision = model_manager->loader().revision();
    for (const auto& entry : runner_components()) {
        state.sources[entry.first] = runner_source_versions(entry.first, model_manager->loader());
    }
    state.ready   = true;
    runner_state_ = std::move(state);
}

void StableDiffusionGGML::end_runners() {
    if (cond_stage_model)
        cond_stage_model->runner_end();
    if (diffusion_model)
        diffusion_model->runner_end();
    if (high_noise_diffusion_model)
        high_noise_diffusion_model->runner_end();
    if (clip_vision)
        clip_vision->runner_end();
    if (ip_adapter)
        ip_adapter->runner_end();
    if (first_stage_model)
        first_stage_model->runner_end();
    if (preview_vae)
        preview_vae->runner_end();
    if (audio_vae_model)
        audio_vae_model->runner_end();
    if (control_net)
        control_net->runner_end();
    for (auto& extension : generation_extensions)
        extension->runner_end();
    for (auto& lora : runtime_lora_models)
        if (lora.model)
            lora.model->runner_end();
}

bool StableDiffusionGGML::reset_runners(const RunnerGroups& groups) {
    conditioning_cache_->clear();
    conditioning_loras_.clear();
    end_runners();
    clear_lora_adapters();
    runtime_lora_models.clear();
    for (auto group : groups) {
        for (auto component : runner_components().at(group)) {
            if (!model_manager->unregister_param_tensors(component)) {
                return false;
            }
        }
    }
    for (auto group : groups) {
        switch (group) {
            case RunnerGroup::Core:
                cond_stage_model.reset();
                diffusion_model.reset();
                high_noise_diffusion_model.reset();
                clip_vision.reset();
                ip_adapter.reset();
                ip_adapter_tokens        = {};
                ip_adapter_uncond_tokens = {};
                runtime_lora_models.clear();
                break;
            case RunnerGroup::VAE:
                first_stage_model.reset();
                preview_vae.reset();
                audio_vae_model.reset();
                break;
            case RunnerGroup::ControlNet:
                control_net.reset();
                break;
            case RunnerGroup::Extensions:
                generation_extensions.clear();
                break;
        }
    }
    return true;
}

bool StableDiffusionGGML::refresh_model_sources() {
    bool changed;
    if (!model_manager->loader().files_changed(changed, false)) {
        return false;
    }
    if (!changed && runner_state_.ready && runner_state_.catalog_revision == model_manager->loader().revision()) {
        return true;
    }
    ModelLoader candidate = model_manager->loader();
    return candidate.refresh_files(false) && apply_model_update(std::move(candidate));
}

bool StableDiffusionGGML::apply_model_update(ModelLoader candidate,
                                             std::unique_ptr<ModelConfig> next_config,
                                             RunnerGroups groups) {
    const SDVersion next_version = candidate.get_sd_version();
    if (next_version == VERSION_COUNT) {
        LOG_ERROR("cannot identify updated diffusion model");
        return false;
    }
    if (!runner_state_.ready || next_version != version) {
        groups = all_runner_groups();
    } else {
        for (const auto& entry : runner_components()) {
            if (runner_state_.sources.at(entry.first) != runner_source_versions(entry.first, candidate)) {
                groups.insert(entry.first);
            }
        }
    }
    if (sd_version_is_z_image(version) && sd_version_is_z_image(next_version)) {
        if (groups.count(RunnerGroup::ControlNet) > 0) {
            groups.insert(RunnerGroup::Core);
        }
        if (groups.count(RunnerGroup::Core) > 0 &&
            model_manager->registered_params_size({ModelComponent::ControlNet}) > 0) {
            groups.insert(RunnerGroup::ControlNet);
        }
    }
    runner_state_.ready = false;
    if (!reset_runners(groups)) {
        return false;
    }
    if (!model_manager->set_loader(std::move(candidate))) {
        reset_runners(all_runner_groups());
        return false;
    }
    if (next_config) {
        config_ = std::move(next_config);
    }
    version = next_version;
    if (!build_runners(groups)) {
        reset_runners(all_runner_groups());
        return false;
    }
    capture_runner_sources();
    return true;
}

ggml_backend_t StableDiffusionGGML::backend_for(SDBackendModule module) {
    ggml_backend_t module_backend = backend_manager.runtime_backend(module);
    if (module_backend == nullptr) {
        LOG_ERROR("failed to initialize %s backend", sd_backend_module_name(module));
    }
    return module_backend;
}

ggml_backend_t StableDiffusionGGML::params_backend_for(SDBackendModule module) {
    ggml_backend_t module_backend = backend_manager.params_backend(module);
    if (module_backend == nullptr) {
        LOG_ERROR("failed to initialize %s params backend", sd_backend_module_name(module));
    }
    return module_backend;
}

void StableDiffusionGGML::set_cancel_flag(enum sd_cancel_mode_t flag) {
    cancellation_flag.store(flag, std::memory_order_release);
}

void StableDiffusionGGML::reset_cancel_flag() {
    set_cancel_flag(SD_CANCEL_RESET);
}

enum sd_cancel_mode_t StableDiffusionGGML::get_cancel_flag() {
    return cancellation_flag.load(std::memory_order_acquire);
}

size_t StableDiffusionGGML::max_graph_vram_bytes_for_module(SDBackendModule module) {
    return max_vram_assignment.bytes_for_backend(backend_for(module));
}

std::vector<size_t> StableDiffusionGGML::layer_split_vram_limits_for_backends(const std::vector<ggml_backend_t>& backends) {
    std::vector<size_t> limits;
    limits.reserve(backends.size());
    for (ggml_backend_t backend : backends) {
        limits.push_back(max_vram_assignment.bytes_for_backend(backend));
    }
    return limits;
}

bool StableDiffusionGGML::ensure_backend_pair(SDBackendModule module) {
    if (backend_for(module) == nullptr) {
        return false;
    }
    return params_backend_for(module) != nullptr;
}

template <typename T>
bool StableDiffusionGGML::register_runner_params(ModelComponent component,
                                                 const std::shared_ptr<T>& model,
                                                 SDBackendModule module,
                                                 size_t* params_mem_size) {
    if (model == nullptr) {
        return true;
    }
    std::map<std::string, ggml_tensor*> group_tensors;
    std::map<ggml_tensor*, enum ggml_op> tensor_ops;
    model->get_param_tensors(group_tensors);
    if constexpr (std::is_base_of_v<Conditioner, T>) {
        model->get_param_tensor_ops(tensor_ops);
    }
    if (model_manager == nullptr) {
        return true;
    }
    ModelManager::ResidencyMode residency_mode =
        backend_manager.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend;

    std::vector<ggml_backend_t> module_backends = backend_manager.runtime_backends(module);
    if (module_backends.size() > 1) {
        if constexpr (has_set_runtime_backends<T>::value) {
            if (module == SDBackendModule::DIFFUSION || module == SDBackendModule::TE) {
                if (backend_manager.split_mode(module) == SDSplitMode::ROW) {
                    return register_row_split_runner_params(component,
                                                            model,
                                                            module,
                                                            module_backends,
                                                            std::move(group_tensors),
                                                            tensor_ops,
                                                            residency_mode,
                                                            params_mem_size);
                }
                return register_layer_split_runner_params(component,
                                                          model,
                                                          module,
                                                          module_backends,
                                                          std::move(group_tensors),
                                                          tensor_ops,
                                                          residency_mode,
                                                          params_mem_size);
            }
        }
        LOG_WARN("%s module does not support multiple runtime backends; using %s",
                 sd_backend_module_name(module),
                 sd::layer_split_backend_device_display_name(module_backends[0]).c_str());
    }
    return model_manager->register_param_tensors(component,
                                                 std::move(group_tensors),
                                                 residency_mode,
                                                 backend_for(module),
                                                 params_backend_for(module),
                                                 params_mem_size,
                                                 false,
                                                 false,
                                                 &tensor_ops);
}

template <typename T>
bool StableDiffusionGGML::register_row_split_runner_params(ModelComponent component,
                                                           const std::shared_ptr<T>& model,
                                                           SDBackendModule module,
                                                           const std::vector<ggml_backend_t>& module_backends,
                                                           std::map<std::string, ggml_tensor*> group_tensors,
                                                           const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
                                                           ModelManager::ResidencyMode residency_mode,
                                                           size_t* params_mem_size) {
    ggml_backend_t main_backend = module_backends[0];

    auto fall_back_to_layer_split = [&](const char* reason) {
        LOG_WARN("%s: row split unavailable (%s); falling back to layer split", model_component_name(component), reason);
        return register_layer_split_runner_params(component,
                                                  model,
                                                  module,
                                                  module_backends,
                                                  std::move(group_tensors),
                                                  tensor_ops,
                                                  residency_mode,
                                                  params_mem_size);
    };

    ggml_backend_dev_t main_dev = ggml_backend_get_device(main_backend);
    ggml_backend_reg_t reg      = main_dev != nullptr ? ggml_backend_dev_backend_reg(main_dev) : nullptr;
    if (reg == nullptr) {
        return fall_back_to_layer_split("no backend registry");
    }
    const size_t reg_dev_count = ggml_backend_reg_dev_count(reg);
    std::vector<float> tensor_split(reg_dev_count, 0.0f);
    constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024;
    for (ggml_backend_t backend : module_backends) {
        ggml_backend_dev_t dev = ggml_backend_get_device(backend);
        int reg_index          = -1;
        for (size_t i = 0; i < reg_dev_count; i++) {
            if (ggml_backend_reg_dev_get(reg, i) == dev) {
                reg_index = (int)i;
                break;
            }
        }
        if (reg_index < 0) {
            return fall_back_to_layer_split("devices span different backend registries");
        }
        size_t free_bytes = 0, total_bytes = 0;
        ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
        int64_t usable_bytes    = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes,
                                                    (int64_t)free_bytes / 8);
        tensor_split[reg_index] = usable_bytes > 0 ? (float)((double)usable_bytes / (1024.0 * 1024.0)) : 1.0f;
    }

    ggml_backend_buffer_type_t split_buft = backend_manager.split_buffer_type(main_backend, tensor_split);
    if (split_buft == nullptr) {
        return fall_back_to_layer_split("backend has no split buffer type");
    }
    std::vector<std::pair<ggml_backend_t, size_t>> split_device_limits;
    for (auto backend : module_backends) {
        split_device_limits.emplace_back(backend, max_vram_assignment.bytes_for_backend(backend));
    }
    model_manager->set_split_buffer_type(main_backend, split_buft, split_device_limits);

    std::map<std::string, ggml_tensor*> split_tensors;
    if constexpr (std::is_base_of_v<Conditioner, T>) {
        model->get_layer_split_param_tensors(split_tensors);
    } else {
        split_tensors = group_tensors;
    }

    std::map<std::string, ggml_tensor*> row_split_map;
    std::map<std::string, ggml_tensor*> regular_map;
    size_t row_split_bytes = 0;
    for (const auto& kv : group_tensors) {
        if (split_tensors.count(kv.first) != 0 &&
            sd::layer_split_tensor_block_index(kv.first) >= 0 &&
            ModelManager::tensor_shape_supports_split_buffer(kv.second)) {
            row_split_map[kv.first] = kv.second;
            row_split_bytes += ggml_nbytes(kv.second);
        } else {
            regular_map[kv.first] = kv.second;
        }
    }
    if (row_split_map.empty()) {
        return fall_back_to_layer_split("no row-splittable transformer block weights found");
    }

    LOG_INFO("%s row split: %zu tensors (%.1f MB) split across %zu devices (main %s)",
             model_component_name(component),
             row_split_map.size(),
             row_split_bytes / (1024.f * 1024.f),
             module_backends.size(),
             sd::layer_split_backend_device_display_name(main_backend).c_str());

    if (!model_manager->register_param_tensors(component,
                                               std::move(row_split_map),
                                               residency_mode,
                                               main_backend,
                                               params_backend_for(module),
                                               params_mem_size,
                                               /*allow_split_buffer=*/true,
                                               false,
                                               &tensor_ops)) {
        return false;
    }
    return model_manager->register_param_tensors(component,
                                                 std::move(regular_map),
                                                 residency_mode,
                                                 main_backend,
                                                 params_backend_for(module),
                                                 params_mem_size,
                                                 false,
                                                 false,
                                                 &tensor_ops);
}

template <typename T>
bool StableDiffusionGGML::register_layer_split_runner_params(ModelComponent component,
                                                             const std::shared_ptr<T>& model,
                                                             SDBackendModule module,
                                                             const std::vector<ggml_backend_t>& module_backends,
                                                             std::map<std::string, ggml_tensor*> group_tensors,
                                                             const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
                                                             ModelManager::ResidencyMode residency_mode,
                                                             size_t* params_mem_size) {
    bool has_cpu_device = false;
    for (ggml_backend_t backend : module_backends) {
        has_cpu_device = has_cpu_device || sd_backend_is_cpu(backend);
    }
    if (has_cpu_device) {
        // The scheduler reserves the CPU slot for its fallback backend, and
        // CPU weight participation is what --params-backend <module>=cpu is
        // for; a CPU device in a split list is almost certainly a mistake.
        LOG_WARN(
            "%s: layer split across a CPU device is not supported; using %s "
            "(use --params-backend %s=cpu to keep weights in RAM)",
            model_component_name(component),
            sd::layer_split_backend_device_display_name(module_backends[0]).c_str(),
            sd_backend_module_name(module));
        return model_manager->register_param_tensors(component,
                                                     std::move(group_tensors),
                                                     residency_mode,
                                                     module_backends[0],
                                                     params_backend_for(module),
                                                     params_mem_size,
                                                     false,
                                                     false,
                                                     &tensor_ops);
    }

    model->set_runtime_backends(module_backends);
    model->set_graph_cut_layer_split_backend_vram_limits(layer_split_vram_limits_for_backends(module_backends));
    model->set_graph_cut_layer_split_enabled(true);
    const bool params_follow_runtime      = backend_manager.params_backend_follows_runtime(module) ||
                                            backend_manager.params_backend_is_disk(module);
    ggml_backend_t initial_params_backend = params_follow_runtime ? module_backends[0] : params_backend_for(module);
    if (initial_params_backend == nullptr) {
        return false;
    }

    LOG_INFO("%s graph-cut layer split: deferring %zu tensors across %zu runtime backends until first graph",
             model_component_name(component),
             group_tensors.size(),
             module_backends.size());

    return model_manager->register_param_tensors(component,
                                                 std::move(group_tensors),
                                                 residency_mode,
                                                 module_backends[0],
                                                 initial_params_backend,
                                                 params_mem_size,
                                                 false,
                                                 params_follow_runtime,
                                                 &tensor_ops);
}

bool StableDiffusionGGML::unload_control_net() {
    ContextOperation operation(*this);
    if (!operation.acquired) {
        return false;
    }
    if (model_manager == nullptr || config_ == nullptr) {
        LOG_ERROR("cannot unload ControlNet: context is not initialized");
        return false;
    }
    ModelLoader candidate = model_manager->loader();
    if (config_->control_net_file != 0 && !candidate.del_file(config_->control_net_file)) {
        return false;
    }
    auto next_config = std::make_unique<ModelConfig>(*config_);
    next_config->set_control_net(0, "");
    return apply_model_update(std::move(candidate), std::move(next_config), {RunnerGroup::ControlNet});
}

bool StableDiffusionGGML::load_control_net_from_file(const std::string& path) {
    ContextOperation operation(*this);
    if (!operation.acquired) {
        return false;
    }
    if (path.empty() || model_manager == nullptr || config_ == nullptr) {
        LOG_ERROR("cannot load ControlNet: invalid path or uninitialized context");
        return false;
    }
    ModelLoader candidate = model_manager->loader();
    ModelLoader::FileId file_id;
    if (!candidate.add_file(path, "", &file_id)) {
        return false;
    }
    if (config_->control_net_file != 0 && config_->control_net_file != file_id && !candidate.del_file(config_->control_net_file)) {
        return false;
    }
    auto next_config = std::make_unique<ModelConfig>(*config_);
    next_config->set_control_net(file_id, path);
    return apply_model_update(std::move(candidate), std::move(next_config), {RunnerGroup::ControlNet});
}

bool StableDiffusionGGML::init_backend() {
    std::string error;
    if (!backend_manager.init(backend_spec.c_str(),
                              params_backend_spec.c_str(),
                              split_mode_spec.c_str(),
                              &error)) {
        LOG_ERROR("backend config failed: %s", error.c_str());
        return false;
    }
    return ensure_backend_pair(SDBackendModule::DIFFUSION);
}

bool StableDiffusionGGML::row_split_active() {
    for (SDBackendModule module : {SDBackendModule::DIFFUSION, SDBackendModule::TE}) {
        if (backend_manager.split_mode(module) == SDSplitMode::ROW &&
            backend_manager.runtime_backends(module).size() > 1) {
            return true;
        }
    }
    return false;
}

bool StableDiffusionGGML::graph_cut_layer_split_active() {
    for (SDBackendModule module : {SDBackendModule::DIFFUSION, SDBackendModule::TE}) {
        if (backend_manager.split_mode(module) == SDSplitMode::LAYER &&
            backend_manager.runtime_backends(module).size() > 1) {
            return true;
        }
    }
    return false;
}

std::shared_ptr<RNG> StableDiffusionGGML::get_rng(rng_type_t rng_type) {
    if (rng_type == STD_DEFAULT_RNG) {
        return std::make_shared<STDDefaultRNG>();
    } else if (rng_type == CPU_RNG) {
        return std::make_shared<MT19937RNG>();
    } else {  // default: CUDA_RNG
        return std::make_shared<PhiloxRNG>();
    }
}

void StableDiffusionGGML::refresh_compvis_denoiser_sigmas() {
    auto comp_vis_denoiser = std::dynamic_pointer_cast<CompVisDenoiser>(denoiser);
    if (!comp_vis_denoiser) {
        return;
    }
    std::vector<float> alphas_cumprod(TIMESTEPS);
    if (file_alphas_cumprod.size() == TIMESTEPS) {
        alphas_cumprod = file_alphas_cumprod;
    } else if (sd_version_is_pixart(version)) {
        // PixArt checkpoints train with a linear beta schedule (0.0001 -> 0.02)
        // instead of the scaled_linear schedule used by SD1.x/SDXL.
        calculate_alphas_cumprod_linear_beta(alphas_cumprod.data(), 0.0001f, 0.02f);
    } else {
        calculate_alphas_cumprod(alphas_cumprod.data());
    }
    for (int i = 0; i < TIMESTEPS; i++) {
        comp_vis_denoiser->sigmas[i]     = std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]);
        comp_vis_denoiser->log_sigmas[i] = std::log(comp_vis_denoiser->sigmas[i]);
    }
}

void StableDiffusionGGML::load_alphas_cumprod() {
    file_alphas_cumprod.clear();

    std::vector<float> loaded_alphas;
    if (!model_manager->load_float_tensor("alphas_cumprod", loaded_alphas)) {
        return;
    }
    if (loaded_alphas.size() != TIMESTEPS) {
        LOG_WARN("ignore alphas_cumprod from model file: expected %d values, got %zu",
                 TIMESTEPS,
                 loaded_alphas.size());
        return;
    }
    for (float alpha : loaded_alphas) {
        if (!std::isfinite(alpha) || alpha <= 0.0f || alpha > 1.0f) {
            LOG_WARN("ignore invalid alphas_cumprod from model file");
            return;
        }
    }

    file_alphas_cumprod = std::move(loaded_alphas);
    LOG_VERBOSE("loaded alphas_cumprod from model file");
}

bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConfig& configuration) {
    const auto* sd_ctx_params = &configuration.params;
    auto& use_tae             = configuration.use_tae;
    auto& use_audio_vae       = configuration.use_audio_vae;
    const auto load_model     = [&](const char* name, const char* path, const std::string& prefix = "") {
        if (!model_loader.init_from_file(path, prefix))
            return false;
        configuration.record_model(name, path);
        return true;
    };
    if (strlen(SAFE_STR(sd_ctx_params->model_path)) > 0) {
        LOG_INFO("loading model from '%s'", sd_ctx_params->model_path);
        if (!load_model("model", sd_ctx_params->model_path)) {
            LOG_ERROR("init model loader from file failed: '%s'", sd_ctx_params->model_path);
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->diffusion_model_path)) > 0) {
        LOG_INFO("loading diffusion model from '%s'", sd_ctx_params->diffusion_model_path);
        if (!load_model("diffusion_model", sd_ctx_params->diffusion_model_path, "model.diffusion_model.")) {
            LOG_WARN("loading diffusion model from '%s' failed", sd_ctx_params->diffusion_model_path);
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) {
        LOG_INFO("loading high noise diffusion model from '%s'", sd_ctx_params->high_noise_diffusion_model_path);
        if (!load_model("high_noise_diffusion_model", sd_ctx_params->high_noise_diffusion_model_path, "model.high_noise_diffusion_model.")) {
            LOG_WARN("loading diffusion model from '%s' failed", sd_ctx_params->high_noise_diffusion_model_path);
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->uncond_diffusion_model_path)) > 0) {
        LOG_INFO("loading unconditional diffusion model from '%s'", sd_ctx_params->uncond_diffusion_model_path);
        if (!load_model("uncond_diffusion_model", sd_ctx_params->uncond_diffusion_model_path, "model.diffusion_model.uncond.")) {
            LOG_WARN("loading unconditional diffusion model from '%s' failed", sd_ctx_params->uncond_diffusion_model_path);
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->clip_l_path)) > 0) {
        LOG_INFO("loading clip_l from '%s'", sd_ctx_params->clip_l_path);
        if (!load_model("clip_l", sd_ctx_params->clip_l_path, "clip_l.")) {
            LOG_WARN("loading clip_l from '%s' failed", sd_ctx_params->clip_l_path);
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->clip_g_path)) > 0) {
        LOG_INFO("loading clip_g from '%s'", sd_ctx_params->clip_g_path);
        if (!load_model("clip_g", sd_ctx_params->clip_g_path, "clip_g.")) {
            LOG_WARN("loading clip_g from '%s' failed", sd_ctx_params->clip_g_path);
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->clip_vision_path)) > 0) {
        LOG_INFO("loading clip_vision from '%s'", sd_ctx_params->clip_vision_path);
        if (!load_model("clip_vision", sd_ctx_params->clip_vision_path, "clip_vision.")) {
            LOG_WARN("loading clip_vision from '%s' failed", sd_ctx_params->clip_vision_path);
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->t5xxl_path)) > 0) {
        LOG_INFO("loading t5xxl from '%s'", sd_ctx_params->t5xxl_path);
        if (!load_model("t5xxl", sd_ctx_params->t5xxl_path, "text_encoders.t5xxl.transformer.")) {
            LOG_WARN("loading t5xxl from '%s' failed", sd_ctx_params->t5xxl_path);
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->pulid_weights_path)) > 0) {
        LOG_INFO("loading PuLID weights from '%s'", sd_ctx_params->pulid_weights_path);
        if (!load_model("pulid_weights", sd_ctx_params->pulid_weights_path,
                        "model.diffusion_model.")) {
            LOG_WARN("loading PuLID weights from '%s' failed", sd_ctx_params->pulid_weights_path);
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->llm_path)) > 0) {
        LOG_INFO("loading llm from '%s'", sd_ctx_params->llm_path);
        if (!load_model("llm", sd_ctx_params->llm_path, "text_encoders.llm.")) {
            LOG_WARN("loading llm from '%s' failed", sd_ctx_params->llm_path);
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->llm_vision_path)) > 0) {
        LOG_INFO("loading llm vision from '%s'", sd_ctx_params->llm_vision_path);
        if (!load_model("llm_vision", sd_ctx_params->llm_vision_path, "text_encoders.llm.visual.")) {
            LOG_WARN("loading llm vision from '%s' failed", sd_ctx_params->llm_vision_path);
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->vae_path)) > 0) {
        LOG_INFO("loading vae from '%s'", sd_ctx_params->vae_path);
        if (!load_model("vae", sd_ctx_params->vae_path, "vae.")) {
            LOG_WARN("loading vae from '%s' failed", sd_ctx_params->vae_path);
            external_vae_is_invalid = true;
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->taesd_path)) > 0) {
        LOG_INFO("loading tae from '%s'", sd_ctx_params->taesd_path);
        if (!load_model("taesd", sd_ctx_params->taesd_path, "tae.")) {
            LOG_WARN("loading tae from '%s' failed", sd_ctx_params->taesd_path);
        } else {
            use_tae = true;
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->embeddings_connectors_path)) > 0) {
        LOG_INFO("loading embeddings connectors from '%s'", sd_ctx_params->embeddings_connectors_path);
        if (!load_model("embeddings_connectors", sd_ctx_params->embeddings_connectors_path, "embeddings_connectors.")) {
            LOG_WARN("loading embeddings connectors from '%s' failed", sd_ctx_params->embeddings_connectors_path);
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->audio_vae_path)) > 0) {
        LOG_INFO("loading audio VAE from '%s'", sd_ctx_params->audio_vae_path);
        if (!load_model("audio_vae", sd_ctx_params->audio_vae_path)) {
            LOG_WARN("loading audio VAE weights from '%s' failed", sd_ctx_params->audio_vae_path);
        } else {
            use_audio_vae = true;
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->audio_encoder_path)) > 0) {
        LOG_INFO("loading audio encoder (wav2vec2) from '%s'", sd_ctx_params->audio_encoder_path);
        if (!load_model("audio_encoder", sd_ctx_params->audio_encoder_path, "wav2vec2.")) {
            LOG_WARN("loading audio encoder weights from '%s' failed", sd_ctx_params->audio_encoder_path);
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->motion_module_path)) > 0) {
        LOG_INFO("loading motion module (AnimateDiff) from '%s'", sd_ctx_params->motion_module_path);
        if (!load_model("motion_module", sd_ctx_params->motion_module_path,
                        "model.diffusion_model.motion_module.")) {
            LOG_WARN("loading motion module from '%s' failed", sd_ctx_params->motion_module_path);
        } else {
            configuration.animatediff_loaded = true;
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->control_net_path)) > 0) {
        if (!model_loader.add_file(sd_ctx_params->control_net_path, "", &configuration.control_net_file)) {
            LOG_ERROR("init control net model loader from file failed: '%s'", sd_ctx_params->control_net_path);
            return false;
        }
    }

    if (configuration.control_net_file)
        configuration.record_model("control_net", sd_ctx_params->control_net_path);

    if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0) {
        if (!load_model("ip_adapter", sd_ctx_params->ip_adapter_path)) {
            LOG_ERROR("init ip-adapter model loader from file failed: '%s'", sd_ctx_params->ip_adapter_path);
            return false;
        }
    }

    if (strlen(SAFE_STR(sd_ctx_params->photo_maker_path)) > 0) {
        configuration.photomaker_source_available = model_loader.add_file(sd_ctx_params->photo_maker_path, "pmid.");
        if (configuration.photomaker_source_available)
            configuration.record_model("photo_maker", sd_ctx_params->photo_maker_path);
        if (!configuration.photomaker_source_available) {
            LOG_WARN("loading stacked ID embedding from '%s' failed", sd_ctx_params->photo_maker_path);
        }
    }

    model_loader.convert_tensors_name();

    ggml_type wtype               = sd_type_to_ggml_type(sd_ctx_params->wtype);
    std::string tensor_type_rules = SAFE_STR(sd_ctx_params->tensor_type_rules);
    if (wtype != GGML_TYPE_COUNT || tensor_type_rules.size() > 0) {
        model_loader.set_wtype_override(wtype, tensor_type_rules);
    }

    return true;
}

bool StableDiffusionGGML::set_sage_attention_enabled(bool enabled) {
    if (!diffusion_model) {
        return false;
    }
    if (enabled) {
#ifndef SD_USE_UPSTREAM_GGML
        auto* ctx = ggml_init({4 * ggml_tensor_overhead(), nullptr, true});
        if (ctx == nullptr) {
            return false;
        }
        auto* q        = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1);
        auto* k        = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1);
        auto* v        = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 128, 128, 1, 1);
        auto* op       = ggml_sage_attn(ctx, q, k, v, 1.f / sqrtf(128.f), GGML_SAGE_ATTN_AUTO);
        bool supported = true;
        for (auto backend : backend_manager.runtime_backends(SDBackendModule::DIFFUSION)) {
            if (!ggml_backend_supports_op(backend, op)) {
                LOG_ERROR("SageAttention is unavailable on %s; it requires patched GGML, CUDA Toolkit 12.0 or newer, and SM80 or newer kernels",
                          ggml_backend_name(backend));
                supported = false;
            }
        }
        ggml_free(ctx);
        if (!supported) {
            return false;
        }
#else
        LOG_ERROR("SageAttention requires -DSD_USE_UPSTREAM_GGML=OFF and a CUDA backend");
        return false;
#endif
    }
    diffusion_model->set_sage_attention_enabled(enabled);
    if (high_noise_diffusion_model) {
        high_noise_diffusion_model->set_sage_attention_enabled(enabled);
    }
    if (enabled) {
        LOG_INFO("Using SageAttention in the diffusion model; CUDA selects the supported kernel, unsupported layers use flash/default attention");
    }
    return true;
}

bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) {
#ifdef SD_USE_UPSTREAM_GGML
    LOG_WARN(
        "Using upstream GGML: INT8 tensorwise/convrot is disabled and FP8 weights are "
        "converted to F16 at load time. Some operators may be unsupported and performance "
        "may be lower than with patched GGML.");
#endif
    if (!validate_tensor_types(sd_ctx_params->wtype, sd_ctx_params->tensor_type_rules)) {
        return false;
    }
    for (float scale : {sd_ctx_params->linear_scale, sd_ctx_params->attn_scale}) {
        if (!std::isfinite(scale) || scale < 0.f || (scale > 0.f && !std::isfinite(1.f / scale))) {
            LOG_ERROR("scale overrides must be finite positive values, or 0 to keep model defaults");
            return false;
        }
    }
    if (sd_ctx_params->conditioning_cache_size < 0) {
        LOG_ERROR("conditioning_cache_size must be non-negative");
        return false;
    }
    const std::string flash_mode = SAFE_STR(sd_ctx_params->flash_attn_mode);
    std::string flash_error;
    if (!sd_parse_flash_attention_assignment(flash_mode, &flash_attn_assignment, &flash_error)) {
        LOG_ERROR("invalid flash_attn_mode '%s': %s", flash_mode.c_str(), flash_error.c_str());
        return false;
    }
    const std::string conv_mode = SAFE_STR(sd_ctx_params->conv_direct_mode);
    std::string conv_error;
    if (!sd_parse_conv_direct_assignment(conv_mode, &conv_direct_assignment, &conv_error)) {
        LOG_ERROR("invalid conv_direct_mode '%s': %s", conv_mode.c_str(), conv_error.c_str());
        return false;
    }
    flash_attn_by_backend.clear();
    conv2d_direct_by_backend.clear();
    conv3d_direct_by_backend.clear();
    conditioning_cache_->set_capacity(static_cast<size_t>(sd_ctx_params->conditioning_cache_size));
    auto configuration = std::make_unique<ModelConfig>(*sd_ctx_params);
    n_threads          = sd_ctx_params->n_threads;
    tensor_executor    = std::make_unique<sd::ParallelExecutor>(n_threads > 0 ? n_threads : sd_get_num_physical_cores());
    sd::ParallelScope tensor_scope(tensor_executor.get());
    enable_mmap               = sd_ctx_params->enable_mmap;
    disable_prefetch          = sd_ctx_params->disable_prefetch;
    disable_segmented_compute = sd_ctx_params->disable_segmented_compute;
    eager_load                = sd_ctx_params->eager_load;
    backend_spec              = SAFE_STR(sd_ctx_params->backend);
    params_backend_spec       = SAFE_STR(sd_ctx_params->params_backend);
    split_mode_spec           = SAFE_STR(sd_ctx_params->split_mode);
    auto_fit_enabled          = sd_ctx_params->auto_fit && params_backend_spec.empty();
    max_vram_assignment.reset(0.f);
    {
        std::string error;
        if (!max_vram_assignment.parse(SAFE_STR(sd_ctx_params->max_vram), &error)) {
            LOG_ERROR("%s", error.c_str());
            return false;
        }
    }

    std::string rpc_servers_spec = SAFE_STR(sd_ctx_params->rpc_servers);
    add_rpc_devices(rpc_servers_spec);

    rng = get_rng(sd_ctx_params->rng_type);
    if (sd_ctx_params->sampler_rng_type != RNG_TYPE_COUNT && sd_ctx_params->sampler_rng_type != sd_ctx_params->rng_type) {
        sampler_rng = get_rng(sd_ctx_params->sampler_rng_type);
    } else {
        sampler_rng = rng;
    }

    ggml_log_set(sd_ggml_log_callback, nullptr);

    model_manager = std::make_shared<ModelManager>();
    model_manager->set_n_threads(n_threads);
    model_manager->set_enable_mmap(enable_mmap);
    model_manager->set_segmented_compute_disabled(disable_segmented_compute);
    model_manager->set_prefetch_disabled(disable_prefetch);
    ModelLoader model_loader;

    if (!init_model_loader(model_loader, *configuration)) {
        return false;
    }

    version = model_loader.get_sd_version();
    if (version == VERSION_COUNT) {
        LOG_ERROR("get sd version from file failed: '%s'", SAFE_STR(sd_ctx_params->model_path));
        return false;
    } else {
        LOG_INFO("Version: %s ", sd_version_name(version));
    }

    if (auto_fit_enabled) {
        if (!sd::backend_fit::derive_backend_specs(model_loader,
                                                   sd_type_to_ggml_type(sd_ctx_params->wtype),
                                                   max_vram_assignment,
                                                   backend_spec,
                                                   params_backend_spec)) {
            return false;
        }
    }

    if (!init_backend()) {
        return false;
    }
    {
        std::string error;
        if (!max_vram_assignment.canonicalize_backend_keys(&error)) {
            LOG_ERROR("%s", error.c_str());
            return false;
        }
    }
    if (eager_load && graph_cut_layer_split_active()) {
        LOG_WARN("--eager-load is not supported with graph-cut layer split; weights will be prepared lazily");
        eager_load = false;
    }

    diffusion_conv_direct = sd_ctx_params->diffusion_conv_direct;
    return apply_model_update(std::move(model_loader), std::move(configuration), all_runner_groups());
}

bool StableDiffusionGGML::uses_tae() const {
    return config_->use_tae || version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09;
}

bool StableDiffusionGGML::tae_preview_only() const {
    return config_->params.tae_preview_only && version != VERSION_SDXS_512_DS && version != VERSION_SDXS_09;
}

void StableDiffusionGGML::configure_weight_loading() {
    const auto* sd_ctx_params = &config_->params;
    const auto& model_loader  = model_manager->loader();
    const auto wtype_stat     = model_loader.get_wtype_stat();
    bool have_int8_tensorwise = false;
    bool have_quant_type_rule = false;
    for (const auto& [_, tensor_storage] : model_loader.get_tensor_storage_map()) {
        if (tensor_storage.is_int8_tensorwise || tensor_storage.w4_convrot_kind) {
            have_int8_tensorwise = true;
        }
        if (tensor_storage.expected_type < GGML_TYPE_COUNT && ggml_is_quantized(tensor_storage.expected_type)) {
            have_quant_type_rule = true;
        }
        if (have_int8_tensorwise && have_quant_type_rule) {
            break;
        }
    }

    if (sd_ctx_params->lora_apply_mode == LORA_APPLY_AUTO) {
        bool have_quantized_weight = have_int8_tensorwise;
        for (const auto& [type, _] : wtype_stat) {
            if (ggml_is_quantized(type)) {
                have_quantized_weight = true;
                break;
            }
        }
        // Avoid full-model LoRA merge buffers on constrained setups.
        const bool params_offloaded      = params_backend_for(SDBackendModule::DIFFUSION) != backend_for(SDBackendModule::DIFFUSION);
        const bool streaming_constrained = params_offloaded ||
                                           backend_manager.params_backend_is_disk(SDBackendModule::DIFFUSION);

        if (have_quant_type_rule || have_quantized_weight || streaming_constrained || row_split_active()) {
            apply_lora_immediately = false;
        } else {
            apply_lora_immediately = true;
        }
    } else if (sd_ctx_params->lora_apply_mode == LORA_APPLY_IMMEDIATELY) {
        if (have_int8_tensorwise) {
            LOG_WARN(
                "ComfyUI INT8/packed INT4 weights do not support the immediately LoRA apply mode; "
                "using at_runtime instead");
            apply_lora_immediately = false;
        } else if (row_split_active()) {
            LOG_WARN(
                "row-split tensors do not support the immediately LoRA apply mode; "
                "LoRAs will not be applied to them (use --lora-apply-mode at_runtime)");
            apply_lora_immediately = false;
        } else {
            apply_lora_immediately = true;
        }
    } else {
        apply_lora_immediately = false;
    }

    bool needs_writable_mmap = enable_mmap && apply_lora_immediately;
    model_manager->set_writable_mmap(needs_writable_mmap);
    if (enable_mmap && apply_lora_immediately) {
        LOG_WARN("in mode 'immediately', LoRAs will cause extra memory usage with mmap");
    }
    model_manager->prepare_file_io();
    load_alphas_cumprod();
}

sd::model_builders::Context StableDiffusionGGML::model_build_context() {
    return {config_->params, version, model_manager->loader().get_tensor_storage_map(), backend_manager, model_manager};
}

bool StableDiffusionGGML::build_core_runners() {
    sd::model_builders::CoreRunners runners;
    if (!sd::model_builders::build_core_runners(model_build_context(), runners)) {
        return false;
    }
    cond_stage_model           = std::move(runners.conditioner);
    diffusion_model            = std::move(runners.diffusion);
    high_noise_diffusion_model = std::move(runners.high_noise_diffusion);
    clip_vision                = std::move(runners.clip_vision);
    ip_adapter                 = std::move(runners.ip_adapter);
    audio_encoder              = std::move(runners.audio_encoder);

    cond_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::TE));
    diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION));
    if (high_noise_diffusion_model) {
        high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION));
    }
    if (clip_vision) {
        clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::CLIP_VISION));
    }
    if (audio_encoder) {
        audio_encoder->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::AUDIO_ENCODER));
    }
    return register_runner_params(ModelComponent::Conditioner, cond_stage_model, SDBackendModule::TE) &&
           register_runner_params(ModelComponent::Diffusion, diffusion_model, SDBackendModule::DIFFUSION) &&
           register_runner_params(ModelComponent::HighNoiseDiffusion, high_noise_diffusion_model, SDBackendModule::DIFFUSION) &&
           register_runner_params(ModelComponent::CLIPVision, clip_vision, SDBackendModule::CLIP_VISION) &&
           register_runner_params(ModelComponent::IPAdapter, ip_adapter, SDBackendModule::DIFFUSION) &&
           register_runner_params(ModelComponent::AudioEncoder, audio_encoder, SDBackendModule::AUDIO_ENCODER);
}

bool StableDiffusionGGML::build_vae_runners() {
    sd::model_builders::VAEOptions options;
    options.use_tae                 = uses_tae();
    options.tae_preview_only        = tae_preview_only();
    options.use_audio_vae           = config_->use_audio_vae;
    options.external_vae_is_invalid = external_vae_is_invalid;
    sd::model_builders::VAERunners runners;
    if (!sd::model_builders::build_vae_runners(model_build_context(), options, runners)) {
        return false;
    }
    first_stage_model = std::move(runners.vae);
    preview_vae       = std::move(runners.preview);
    audio_vae_model   = std::move(runners.audio);

    first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::VAE));
    if (preview_vae) {
        preview_vae->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::VAE));
    }
    return register_runner_params(ModelComponent::VAE, first_stage_model, SDBackendModule::VAE) &&
           register_runner_params(ModelComponent::PreviewVAE, preview_vae, SDBackendModule::VAE) &&
           register_runner_params(ModelComponent::AudioVAE, audio_vae_model, SDBackendModule::VAE);
}

bool StableDiffusionGGML::build_control_net_runner() {
    if (config_->control_net_file == 0) {
        return true;
    }
    if (sd_version_is_z_image(version)) {
        auto z_image = std::dynamic_pointer_cast<ZImage::ZImageRunner>(diffusion_model);
        if (z_image == nullptr || !z_image->control_enabled()) {
            LOG_WARN("z-image control net file present but the diffusion model has no control weights; ignoring");
            return true;
        }
        std::map<std::string, ggml_tensor*> control_tensors;
        z_image->get_control_param_tensors(control_tensors);
        if (control_tensors.empty()) {
            LOG_ERROR("z-image control net has no param tensors to register");
            return false;
        }
        ModelManager::ResidencyMode residency_mode = backend_manager.params_backend_is_disk(SDBackendModule::DIFFUSION)
                                                         ? ModelManager::ResidencyMode::Disk
                                                         : ModelManager::ResidencyMode::ParamBackend;
        return model_manager->register_param_tensors(ModelComponent::ControlNet,
                                                     std::move(control_tensors),
                                                     residency_mode,
                                                     backend_for(SDBackendModule::DIFFUSION),
                                                     params_backend_for(SDBackendModule::DIFFUSION));
    }
    if (!sd::model_builders::build_control_net_runner(model_build_context(), control_net)) {
        return false;
    }
    control_net->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::CONTROL_NET));
    return register_runner_params(ModelComponent::ControlNet, control_net, SDBackendModule::CONTROL_NET);
}

bool StableDiffusionGGML::build_extension_runners() {
    GenerationExtensionInitContext extension_ctx{
        &config_->params,
        version,
        model_manager->loader().get_tensor_storage_map(),
        config_->photomaker_source_available,
        model_manager,
        n_threads,
        [this](SDBackendModule module) { return ensure_backend_pair(module); },
        [this](SDBackendModule module) { return backend_for(module); },
        [this](SDBackendModule module) { return params_backend_for(module); },
    };
    if (!sd::model_builders::build_extension_runners(extension_ctx, generation_extensions)) {
        return false;
    }
    for (auto& extension : generation_extensions) {
        if (!register_runner_params(extension->component(), extension, SDBackendModule::PHOTOMAKER)) {
            return false;
        }
    }
    return true;
}

bool StableDiffusionGGML::flash_attention_for(SDBackendModule module) {
    const auto& params     = config_->params;
    const std::string mode = flash_attn_assignment.get(module);
    if (mode == "on") {
        return true;
    }
    if (mode == "off") {
        return false;
    }
    if (strlen(SAFE_STR(params.flash_attn_mode)) == 0) {
        if (params.flash_attn) {
            return true;
        }
        if (params.diffusion_flash_attn) {
            return module == SDBackendModule::DIFFUSION;
        }
    }
    auto backends    = backend_manager.runtime_backends(module);
    bool enabled     = !backends.empty();
    bool any_enabled = false;
    for (auto backend : backends) {
        auto it = flash_attn_by_backend.find(backend);
        if (it == flash_attn_by_backend.end()) {
            it = flash_attn_by_backend.emplace(backend, sd_probe_flash_attention(backend, n_threads)).first;
        }
        enabled     = enabled && it->second;
        any_enabled = any_enabled || it->second;
    }
    if (any_enabled && !enabled) {
        LOG_INFO("flash attention auto: %s -> off (compute devices disagree)", sd_backend_module_name(module));
    }
    return enabled;
}

bool StableDiffusionGGML::conv_direct_for(SDBackendModule module, bool conv3d) {
    const auto& params     = config_->params;
    const std::string mode = conv_direct_assignment.get(module);
    if (mode == "on" || mode == "off") {
        return mode == "on";
    }
    if (strlen(SAFE_STR(params.conv_direct_mode)) == 0 &&
        (module == SDBackendModule::VAE ? params.vae_conv_direct : params.diffusion_conv_direct)) {
        return true;
    }
    auto& by_backend = conv3d ? conv3d_direct_by_backend : conv2d_direct_by_backend;
    auto backends    = backend_manager.runtime_backends(module);
    bool enabled     = !backends.empty();
    bool any_enabled = false;
    for (auto backend : backends) {
        auto it = by_backend.find(backend);
        if (it == by_backend.end()) {
            it = by_backend.emplace(backend, sd_probe_conv_direct(backend, conv3d)).first;
        }
        enabled     = enabled && it->second;
        any_enabled = any_enabled || it->second;
    }
    if (any_enabled && !enabled) {
        LOG_INFO("conv%dd direct auto: %s -> off (compute devices disagree)", conv3d ? 3 : 2, sd_backend_module_name(module));
    }
    return enabled;
}

bool StableDiffusionGGML::validate_and_load_runners() {
    const auto* sd_ctx_params   = &config_->params;
    const bool use_tae          = uses_tae();
    const bool tae_preview_only = this->tae_preview_only();
    cond_stage_model->set_flash_attention_enabled(flash_attention_for(SDBackendModule::TE));
    if (clip_vision) {
        clip_vision->set_flash_attention_enabled(flash_attention_for(SDBackendModule::CLIP_VISION));
    }
    if (first_stage_model || preview_vae || audio_vae_model) {
        const bool enabled = flash_attention_for(SDBackendModule::VAE);
        if (first_stage_model) {
            first_stage_model->set_flash_attention_enabled(enabled);
        }
        if (preview_vae) {
            preview_vae->set_flash_attention_enabled(enabled);
        }
        if (audio_vae_model) {
            audio_vae_model->set_flash_attention_enabled(enabled);
        }
    }
    const bool diffusion_fa = flash_attention_for(SDBackendModule::DIFFUSION);
    diffusion_model->set_flash_attention_enabled(diffusion_fa);
    if (high_noise_diffusion_model) {
        high_noise_diffusion_model->set_flash_attention_enabled(diffusion_fa);
    }
    if (ip_adapter) {
        ip_adapter->set_flash_attention_enabled(diffusion_fa);
    }
    if (control_net) {
        control_net->set_flash_attention_enabled(flash_attention_for(SDBackendModule::CONTROL_NET));
    }
    if (audio_encoder) {
        audio_encoder->set_flash_attention_enabled(flash_attention_for(SDBackendModule::AUDIO_ENCODER));
    }
    if (sd_ctx_params->sage_attn && !set_sage_attention_enabled(true)) {
        return false;
    }
    const auto set_conv_direct = [this](SDBackendModule module, const std::vector<GGMLRunner*>& runners) {
        const bool conv2d = conv_direct_for(module, false);
        const bool conv3d = conv_direct_for(module, true);
        for (auto runner : runners) {
            if (runner != nullptr) {
                runner->set_conv2d_direct_enabled(conv2d);
                runner->set_conv3d_direct_enabled(conv3d);
            }
        }
    };
    set_conv_direct(SDBackendModule::DIFFUSION, {diffusion_model.get(), high_noise_diffusion_model.get()});
    if (first_stage_model || preview_vae || audio_vae_model) {
        set_conv_direct(SDBackendModule::VAE, {first_stage_model.get(), preview_vae.get(), audio_vae_model.get()});
    }
    if (control_net) {
        set_conv_direct(SDBackendModule::CONTROL_NET, {control_net.get()});
    }
    LOG_VERBOSE("validating model metadata");

    std::set<std::string> ignore_tensors;
    if (version == VERSION_MING_IMAGE) {
        ignore_tensors.insert("text_encoders.llm.vision.");
        ignore_tensors.insert("text_encoders.llm.linear_proj.");
        ignore_tensors.insert("text_encoders.llm.backbone.lm_head.");
    }
    if (use_tae && !tae_preview_only) {
        ignore_tensors.insert("first_stage_model.");
    }
    for (auto& extension : generation_extensions) {
        extension->add_ignore_tensors(ignore_tensors);
    }
    ignore_tensors.insert("model.diffusion_model.__x0__");
    ignore_tensors.insert("model.diffusion_model.__32x32__");
    ignore_tensors.insert("model.diffusion_model.__index_timestep_zero__");

    if (audio_encoder != nullptr) {
        // These wav2vec2 tensors are unused during feature extraction.
        ignore_tensors.insert("wav2vec2.lm_head.");
        ignore_tensors.insert("wav2vec2.masked_spec_embed");
    }

    if (audio_vae_model) {
        if (!sd_version_is_minimax_h3(version)) {
            ignore_tensors.insert("audio_vae.encoder");
        }
    }
    if (version == VERSION_OVIS_IMAGE) {
        ignore_tensors.insert("text_encoders.llm.vision_model.");
        ignore_tensors.insert("text_encoders.llm.visual_tokenizer.");
        ignore_tensors.insert("text_encoders.llm.vte.");
    }
    if (version == VERSION_SVD) {
        ignore_tensors.insert("conditioner.embedders.3");
    }
    if (sd_version_is_ernie_image(version)) {
        ignore_tensors.insert("text_encoders.llm.vision_tower.");
        ignore_tensors.insert("text_encoders.llm.multi_modal_projector.");
    }
    if (sd_version_is_lens(version)) {
        ignore_tensors.insert("text_encoders.llm.tokenizer_json");
        ignore_tensors.insert("text_encoders.llm.model.layers.0.mlp.experts.gate_up_proj.weight_scale_2");
        ignore_tensors.insert("text_encoders.llm.model.layers.0.mlp.experts.down_proj.weight_scale_2");
    }
    if (sd_version_is_ideogram4(version)) {
        ignore_tensors.insert("text_encoders.llm.lm_head.");
        ignore_tensors.insert("text_encoders.llm.visual.");
        ignore_tensors.insert("text_encoders.llm.vision_model.");
        ignore_tensors.insert("text_encoders.llm.tokenizer_json");
    }
    if (version == VERSION_HIDREAM_O1) {
        ignore_tensors.insert("lm_head.");
        ignore_tensors.insert("model.visual.deepstack_merger_list.");
    }

    model_manager->set_common_ignore_tensors(ignore_tensors);
    if (!model_manager->validate_registered_tensors()) {
        LOG_ERROR("model metadata validation failed");
        return false;
    }

    if (eager_load) {
        if (!model_manager->load_all_params_eagerly()) {
            LOG_ERROR("model params eager load failed");
            return false;
        }
        LOG_VERBOSE("model metadata validated; weights pre-loaded to params backend");
    } else {
        LOG_VERBOSE("model metadata validated; weights will be prepared lazily");
    }

    {
        size_t text_encoder_params_mem_size = model_manager->registered_params_size({ModelComponent::Conditioner});
        size_t unet_params_mem_size         = model_manager->registered_params_size({ModelComponent::Diffusion, ModelComponent::HighNoiseDiffusion});
        size_t vae_params_mem_size          = model_manager->registered_params_size(runner_components().at(RunnerGroup::VAE));
        size_t control_net_params_mem_size  = model_manager->registered_params_size({ModelComponent::ControlNet});
        size_t extension_params_mem_size    = model_manager->registered_params_size(runner_components().at(RunnerGroup::Extensions));
        size_t total_params_ram_size        = 0;
        size_t total_params_vram_size       = 0;
        auto add_params_memory              = [&](size_t size, SDBackendModule module) {
            if (size == 0) {
                return true;
            }
            ggml_backend_t module_backend = params_backend_for(module);
            if (module_backend == nullptr) {
                return false;
            }
            if (sd_backend_is_cpu(module_backend)) {
                total_params_ram_size += size;
            } else {
                total_params_vram_size += size;
            }
            return true;
        };
        auto params_memory_location = [&](size_t size, SDBackendModule module) {
            if (size == 0) {
                return "N/A";
            }
            ggml_backend_t module_backend = params_backend_for(module);
            if (module_backend == nullptr) {
                return "N/A";
            }
            return sd_backend_is_cpu(module_backend) ? "RAM" : "VRAM";
        };

        if (!add_params_memory(text_encoder_params_mem_size, SDBackendModule::TE) ||
            !add_params_memory(extension_params_mem_size, SDBackendModule::PHOTOMAKER) ||
            !add_params_memory(unet_params_mem_size, SDBackendModule::DIFFUSION) ||
            !add_params_memory(vae_params_mem_size, SDBackendModule::VAE) ||
            !add_params_memory(control_net_params_mem_size, SDBackendModule::CONTROL_NET)) {
            return false;
        }

        size_t total_params_size = total_params_ram_size + total_params_vram_size;
        LOG_INFO(
            "total params memory size = %.2fMB (VRAM %.2fMB, RAM %.2fMB): "
            "text_encoders %.2fMB(%s), diffusion_model %.2fMB(%s), vae %.2fMB(%s), controlnet %.2fMB(%s), extensions %.2fMB(%s)",
            total_params_size / 1024.0 / 1024.0,
            total_params_vram_size / 1024.0 / 1024.0,
            total_params_ram_size / 1024.0 / 1024.0,
            text_encoder_params_mem_size / 1024.0 / 1024.0,
            params_memory_location(text_encoder_params_mem_size, SDBackendModule::TE),
            unet_params_mem_size / 1024.0 / 1024.0,
            params_memory_location(unet_params_mem_size, SDBackendModule::DIFFUSION),
            vae_params_mem_size / 1024.0 / 1024.0,
            params_memory_location(vae_params_mem_size, SDBackendModule::VAE),
            control_net_params_mem_size / 1024.0 / 1024.0,
            params_memory_location(control_net_params_mem_size, SDBackendModule::CONTROL_NET),
            extension_params_mem_size / 1024.0 / 1024.0,
            params_memory_location(extension_params_mem_size, SDBackendModule::PHOTOMAKER));
    }
    return true;
}

bool StableDiffusionGGML::build_denoiser() {
    const auto* sd_ctx_params      = &config_->params;
    const auto& model_loader       = model_manager->loader();
    const auto& tensor_storage_map = model_loader.get_tensor_storage_map();
    denoiser                       = std::make_shared<CompVisDenoiser>();
    default_flow_shift             = INFINITY;
    prediction_t pred_type         = sd_ctx_params->prediction;

    if (pred_type == PREDICTION_COUNT) {
        const auto* definition = find_model_definition(version);
        if (!definition)
            return false;
        const auto& policy = definition->sampling;
        pred_type          = policy.default_prediction;
        default_flow_shift = policy.flow_shift_from_weights ? policy.flow_shift_from_weights(tensor_storage_map) : policy.default_flow_shift;
        if (policy.probe_v_prediction) {
            pred_type = is_using_v_parameterization_for_sd2(sd_version_is_inpaint(version)) ? V_PRED : EPS_PRED;
        } else if (policy.prediction_from_weights) {
            pred_type = policy.prediction_from_weights(tensor_storage_map);
        }
    }

    switch (pred_type) {
        case EPS_PRED:
            LOG_INFO("running in eps-prediction mode");
            break;
        case V_PRED:
            LOG_INFO("running in v-prediction mode");
            denoiser = std::make_shared<CompVisVDenoiser>();
            break;
        case EDM_V_PRED:
            LOG_INFO("running in v-prediction EDM mode");
            denoiser = std::make_shared<EDMVDenoiser>();
            break;
        case FLOW_PRED: {
            LOG_INFO("running in FLOW mode");
            denoiser = std::make_shared<DiscreteFlowDenoiser>();
            break;
        }
        case FLUX_FLOW_PRED: {
            LOG_INFO("running in Flux FLOW mode");
            denoiser = std::make_shared<FluxFlowDenoiser>();
            break;
        }
        case FLUX3_FLOW_PRED: {
            LOG_INFO("running in Flux.3 FLOW mode");
            denoiser = std::make_shared<Flux3FlowDenoiser>();
            break;
        }
        case H3_AV_FLOW_PRED: {
            LOG_INFO("running in MiniMax H3 AV FLOW mode");
            denoiser = std::make_shared<H3AVFlowDenoiser>(default_flow_shift, version == VERSION_P_H3_IMAGE ? 1.f : 3.f, get_latent_channel());
            break;
        }
        case MING_IMAGE_FLOW_PRED: {
            LOG_INFO("running in Ming-Image FLOW mode");
            denoiser = std::make_shared<MingImageFlowDenoiser>();
            break;
        }
        case SEFI_FLOW_PRED: {
            LOG_INFO("running in SeFi-Image dual-time FLOW mode");
            denoiser = std::make_shared<SefiFlowDenoiser>();
            break;
        }
        case MINIT2I_FLOW_PRED: {
            LOG_INFO("running in MiniT2I FLOW mode");
            denoiser = std::make_shared<MiniT2IFlowDenoiser>();
            break;
        }
        case SENSENOVA_U1_FLOW_PRED: {
            LOG_INFO("running in SenseNova U1.5 FLOW mode");
            denoiser = std::make_shared<SenseNovaU1FlowDenoiser>(default_flow_shift);
            break;
        }
        default: {
            LOG_ERROR("Unknown predition type %i", pred_type);
            return false;
        }
    }

    refresh_compvis_denoiser_sigmas();
    return true;
}

bool StableDiffusionGGML::build_runners(const RunnerGroups& groups) {
    const auto& model_loader                                 = model_manager->loader();
    std::map<ggml_type, uint32_t> wtype_stat                 = model_loader.get_wtype_stat();
    std::map<ggml_type, uint32_t> conditioner_wtype_stat     = model_loader.get_conditioner_wtype_stat();
    std::map<ggml_type, uint32_t> diffusion_model_wtype_stat = model_loader.get_diffusion_model_wtype_stat();
    std::map<ggml_type, uint32_t> vae_wtype_stat             = model_loader.get_vae_wtype_stat();

    auto wtype_stat_to_str = [](const std::map<ggml_type, uint32_t>& m, int key_width = 8, int value_width = 5) -> std::string {
        std::ostringstream oss;
        bool first = true;
        for (const auto& [type, count] : m) {
            if (!first)
                oss << "|";
            first = false;
            oss << std::right << std::setw(key_width) << ggml_type_name(type)
                << ": "
                << std::left << std::setw(value_width) << count;
        }
        return oss.str();
    };

    LOG_INFO("Weight type stat:                 %s", wtype_stat_to_str(wtype_stat).c_str());
    LOG_INFO("Conditioner weight type stat:     %s", wtype_stat_to_str(conditioner_wtype_stat).c_str());
    LOG_INFO("Diffusion model weight type stat: %s", wtype_stat_to_str(diffusion_model_wtype_stat).c_str());
    LOG_INFO("VAE weight type stat:             %s", wtype_stat_to_str(vae_wtype_stat).c_str());

    LOG_DEBUG("ggml tensor size = %d bytes", (int)sizeof(ggml_tensor));

    configure_weight_loading();
    for (auto group : groups) {
        bool success = false;
        switch (group) {
            case RunnerGroup::Core:
                success = build_core_runners();
                break;
            case RunnerGroup::VAE:
                success = build_vae_runners();
                break;
            case RunnerGroup::ControlNet:
                success = build_control_net_runner();
                break;
            case RunnerGroup::Extensions:
                success = build_extension_runners();
                break;
        }
        if (!success) {
            return false;
        }
    }
    if (!validate_and_load_runners()) {
        return false;
    }
    return groups.count(RunnerGroup::Core) == 0 || build_denoiser();
}

bool StableDiffusionGGML::is_using_v_parameterization_for_sd2(bool is_inpaint) {
    struct RunnerEndOnExit {
        GGMLRunner* runner = nullptr;
        ~RunnerEndOnExit() {
            if (runner != nullptr) {
                runner->runner_end();
            }
        }
    };
    RunnerEndOnExit diffusion_runner_end{diffusion_model.get()};

    sd::Tensor<float> x_t   = sd::full<float>({8, 8, 4, 1}, 0.5f);
    sd::Tensor<float> c     = sd::full<float>({1024, 2, 1, 1}, 0.5f);
    sd::Tensor<float> steps = sd::full<float>({1}, 999.0f);
    sd::Tensor<float> concat;
    if (is_inpaint) {
        concat = sd::zeros<float>({8, 8, 5, 1});
    }

    int64_t t0 = ggml_time_ms();
    sd::Tensor<float> out;
    DiffusionParams diffusion_params;
    diffusion_params.x         = &x_t;
    diffusion_params.timesteps = &steps;
    diffusion_params.context   = &c;
    diffusion_params.extra     = UNetDiffusionExtra{};
    if (!concat.empty()) {
        diffusion_params.c_concat = &concat;
    }
    auto out_opt = diffusion_model->compute(n_threads, diffusion_params);
    GGML_ASSERT(!out_opt.empty());
    out = std::move(out_opt);

    double result = static_cast<double>((out - x_t).mean());
    int64_t t1    = ggml_time_ms();
    LOG_VERBOSE("check is_using_v_parameterization_for_sd2, taking %.2fs", (t1 - t0) * 1.0f / 1000);
    return result < -1;
}

std::string StableDiffusionGGML::lora_log_id(const ModelManager::LoraSpec& lora) {
    return lora.is_high_noise ? "|high_noise|" + lora.path : lora.path;
}

std::shared_ptr<LoraModel> StableDiffusionGGML::load_lora_model(const ModelManager::LoraSpec& lora_spec,
                                                                SDBackendModule module,
                                                                LoraModel::filter_t module_filter) {
    if (!ensure_backend_pair(module)) {
        return nullptr;
    }
    if (lora_spec.is_high_noise) {
        LOG_VERBOSE("high noise lora: %s", lora_spec.path.c_str());
    }
    const auto mode                        = backend_manager.params_backend_is_disk(module)
                                                 ? ModelManager::ResidencyMode::Disk
                                                 : ModelManager::ResidencyMode::ParamBackend;
    auto lora                              = std::make_shared<LoraModel>(lora_log_id(lora_spec), backend_for(module), params_backend_for(module),
                                                                         model_manager, lora_spec.file_id, version, mode,
                                                                         backend_manager.params_backend_follows_runtime(module));
    LoraModel::filter_t lora_tensor_filter = module_filter;
    if (!lora_spec.tensor_name_prefix_filter.empty()) {
        lora_tensor_filter = [module_filter, prefix = lora_spec.tensor_name_prefix_filter](const std::string& tensor_name) {
            return starts_with(tensor_name, prefix) && (!module_filter || module_filter(tensor_name));
        };
    }
    if (!lora->init_params(n_threads, lora_tensor_filter)) {
        LOG_WARN("load lora tensors from %s failed", lora_spec.path.c_str());
        return nullptr;
    }

    lora->multiplier = lora_spec.multiplier;
    return lora;
}

void StableDiffusionGGML::clear_lora_adapters() {
    if (cond_stage_model) {
        cond_stage_model->set_weight_adapter(nullptr);
    }
    if (diffusion_model) {
        diffusion_model->set_weight_adapter(nullptr);
    }
    if (high_noise_diffusion_model) {
        high_noise_diffusion_model->set_weight_adapter(nullptr);
    }
    if (first_stage_model) {
        first_stage_model->set_weight_adapter(nullptr);
    }
}

std::vector<std::shared_ptr<LoraModel>> StableDiffusionGGML::load_runtime_loras_for_module(const std::vector<ModelManager::LoraSpec>& loras,
                                                                                           const std::set<std::string>& model_tensor_names,
                                                                                           SDBackendModule module,
                                                                                           LoraModel::filter_t module_filter,
                                                                                           bool& success,
                                                                                           std::vector<RuntimeLora>& next_models) {
    std::vector<std::shared_ptr<LoraModel>> module_lora_models;
    for (const auto& lora_spec : loras) {
        auto cached = std::find_if(runtime_lora_models.begin(), runtime_lora_models.end(), [&](const RuntimeLora& entry) {
            return entry.model != nullptr && entry.module == module && entry.matches(lora_spec);
        });
        auto lora   = cached == runtime_lora_models.end() ? load_lora_model(lora_spec, module, module_filter)
                                                          : std::move(cached->model);
        if (lora == nullptr) {
            if (lora_spec.required) {
                LOG_ERROR("required lora load failed: %s", lora_spec.path.c_str());
                success = false;
            }
            continue;
        }
        if (lora->lora_tensors.empty()) {
            continue;
        }

        lora->preprocess_lora_tensors(model_tensor_names);
        lora->multiplier = lora_spec.multiplier;
        next_models.push_back({lora_spec, module, lora});
        module_lora_models.push_back(std::move(lora));
    }
    return module_lora_models;
}

bool StableDiffusionGGML::apply_loras_immediately(const std::vector<ModelManager::LoraSpec>& loras) {
    if (model_manager == nullptr) {
        if (!loras.empty()) {
            LOG_WARN("model manager is not available for immediate lora");
        }
        return false;
    }

    clear_lora_adapters();
    runtime_lora_models.clear();

    if (!loras.empty()) {
        LOG_INFO("apply lora immediately");
    }
    return model_manager->set_loras(loras, version);
}

bool StableDiffusionGGML::apply_loras_at_runtime(const std::vector<ModelManager::LoraSpec>& loras) {
    if (model_manager != nullptr) {
        if (!model_manager->set_loras({}, version))
            return false;
    }
    clear_lora_adapters();
    if (loras.empty()) {
        runtime_lora_models.clear();
        return true;
    }

    bool success = true;
    std::vector<RuntimeLora> next_models;
    std::set<std::string> model_tensor_names;
    if (model_manager != nullptr) {
        model_tensor_names = model_manager->tensor_names();
    }

    LOG_INFO("apply lora at runtime");
    if (cond_stage_model) {
        auto lora_tensor_filter = [&](const std::string& tensor_name) {
            if (is_cond_stage_model_name(tensor_name)) {
                return true;
            }
            return false;
        };
        auto cond_stage_lora_models =
            load_runtime_loras_for_module(loras,
                                          model_tensor_names,
                                          SDBackendModule::TE,
                                          lora_tensor_filter, success, next_models);
        // Only attach the adapter when there are LoRAs targeting the cond_stage model.
        // An empty MultiLoraAdapter still routes every linear/conv through
        // forward_with_lora() instead of the direct kernel path — slower for no benefit.
        if (!cond_stage_lora_models.empty()) {
            auto multi_lora_adapter = std::make_shared<MultiLoraAdapter>(cond_stage_lora_models);
            cond_stage_model->set_weight_adapter(multi_lora_adapter);
        }
    }
    if (diffusion_model) {
        auto lora_tensor_filter = [&](const std::string& tensor_name) {
            if (is_diffusion_model_name(tensor_name)) {
                return true;
            }
            return false;
        };
        auto diffusion_lora_models =
            load_runtime_loras_for_module(loras,
                                          model_tensor_names,
                                          SDBackendModule::DIFFUSION,
                                          lora_tensor_filter, success, next_models);
        if (!diffusion_lora_models.empty()) {
            auto multi_lora_adapter = std::make_shared<MultiLoraAdapter>(diffusion_lora_models);
            diffusion_model->set_weight_adapter(multi_lora_adapter);
            if (high_noise_diffusion_model) {
                high_noise_diffusion_model->set_weight_adapter(multi_lora_adapter);
            }
        }
    }

    if (first_stage_model) {
        auto lora_tensor_filter = [&](const std::string& tensor_name) {
            if (is_first_stage_model_name(tensor_name)) {
                return true;
            }
            return false;
        };
        auto first_stage_lora_models =
            load_runtime_loras_for_module(loras,
                                          model_tensor_names,
                                          SDBackendModule::VAE,
                                          lora_tensor_filter, success, next_models);
        if (!first_stage_lora_models.empty()) {
            auto multi_lora_adapter = std::make_shared<MultiLoraAdapter>(first_stage_lora_models);
            first_stage_model->set_weight_adapter(multi_lora_adapter);
        }
    }
    runtime_lora_models = std::move(next_models);
    return success;
}

void StableDiffusionGGML::lora_stat() {
    if (!runtime_lora_models.empty()) {
        LOG_INFO("runtime_lora_models:");
        for (auto& lora_model : runtime_lora_models) {
            lora_model.model->stat();
        }
    }
}

bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_count) {
    std::vector<ModelManager::LoraSpec> all_loras;
    all_loras.reserve(lora_count);
    for (uint32_t i = 0; i < lora_count; i++) {
        std::string lora_id = SAFE_STR(loras[i].path);
        ModelManager::LoraSpec lora_spec;
        lora_spec.path          = lora_id;
        lora_spec.multiplier    = loras[i].multiplier;
        lora_spec.is_high_noise = loras[i].is_high_noise;
        all_loras.push_back(std::move(lora_spec));
        if (loras[i].is_high_noise) {
            lora_id = "|high_noise|" + lora_id;
        }
        LOG_VERBOSE("lora %s:%.2f", lora_id.c_str(), loras[i].multiplier);
    }

    for (auto& extension : generation_extensions) {
        extension->collect_loras(all_loras);
    }

    int64_t t0 = ggml_time_ms();
    end_runners();
    clear_lora_adapters();
    if (!model_manager->prepare_lora_sources(all_loras)) {
        conditioning_cache_->clear();
        return false;
    }
    if (!std::equal(all_loras.begin(), all_loras.end(),
                    conditioning_loras_.begin(), conditioning_loras_.end(),
                    [](const ModelManager::LoraSpec& a, const ModelManager::LoraSpec& b) {
                        return a.file_id == b.file_id && a.file_revision == b.file_revision &&
                               a.multiplier == b.multiplier && a.is_high_noise == b.is_high_noise &&
                               a.tensor_name_prefix_filter == b.tensor_name_prefix_filter;
                    })) {
        conditioning_cache_->clear();
    }
    runtime_lora_models.erase(std::remove_if(runtime_lora_models.begin(), runtime_lora_models.end(), [&](const RuntimeLora& entry) {
                                  return std::none_of(all_loras.begin(), all_loras.end(), [&](const ModelManager::LoraSpec& spec) {
                                      return entry.matches(spec);
                                  });
                              }),
                              runtime_lora_models.end());
    const bool success = apply_lora_immediately ? apply_loras_immediately(all_loras)
                                                : apply_loras_at_runtime(all_loras);
    if (!success) {
        conditioning_cache_->clear();
        clear_lora_adapters();
        runtime_lora_models.clear();
        return false;
    }
    runner_state_.catalog_revision = model_manager->loader().revision();
    int64_t t1                     = ggml_time_ms();
    if (!all_loras.empty()) {
        LOG_INFO("apply_loras completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
    }
    conditioning_loras_ = std::move(all_loras);
    return true;
}

SDCondition StableDiffusionGGML::get_learned_condition(const ConditionerParams& params) {
    return conditioning_cache_->get(*cond_stage_model, n_threads, params);
}

void StableDiffusionGGML::reset_generation_extensions() {
    for (auto& extension : generation_extensions) {
        extension->reset_runtime_condition();
    }
}

void StableDiffusionGGML::prepare_generation_extensions(const sd_pm_params_t& pm_params,
                                                        const sd_pulid_params_t& pulid_params,
                                                        ConditionerParams& condition_params,
                                                        int total_steps) {
    reset_generation_extensions();
    GenerationExtensionConditionContext ctx{
        cond_stage_model.get(),
        condition_params,
        pm_params,
        pulid_params,
        n_threads,
        total_steps,
    };

    for (auto& extension : generation_extensions) {
        extension->prepare_condition(ctx);
    }
}

sd::Tensor<float> StableDiffusionGGML::get_clip_vision_output(const sd::Tensor<float>& image,
                                                              bool return_pooled,
                                                              int clip_skip,
                                                              bool zero_out_masked) {
    sd::Tensor<float> output;
    if (zero_out_masked) {
        if (return_pooled) {
            output = sd::zeros<float>({clip_vision->vision_model.projection_dim});
        } else {
            output = sd::zeros<float>({clip_vision->vision_model.hidden_size, 257});
        }
    } else {
        auto pixel_values = clip_preprocess(image, clip_vision->vision_model.image_size, clip_vision->vision_model.image_size);
        auto output_opt   = clip_vision->compute(n_threads, pixel_values, return_pooled, clip_skip);
        if (output_opt.empty()) {
            LOG_ERROR("clip_vision compute failed");
            return {};
        }
        output = std::move(output_opt);
    }
    return output;
}

// Returns 50 Hz wav2vec2 states in sd::Tensor layout: [dim, frames, layers].
sd::Tensor<float> StableDiffusionGGML::get_audio_embedding(const sd_audio_t& audio) {
    if (audio_encoder == nullptr) {
        LOG_ERROR("audio encoder model is not loaded");
        return {};
    }
    if (audio.data == nullptr || audio.sample_count == 0 || audio.channels == 0 || audio.sample_rate == 0) {
        LOG_ERROR("invalid driving audio");
        return {};
    }
    auto mono = sd::audio::downmix_to_mono(audio.data, audio.sample_count, audio.channels);
    if (mono.empty()) {
        LOG_ERROR("audio mono downmix failed");
        return {};
    }
    mono = sd::audio::resample_audio(mono.data(), mono.size(), audio.sample_rate, 16000);
    if (mono.empty()) {
        LOG_ERROR("audio resample to 16 kHz failed");
        return {};
    }
    return audio_encoder->compute(n_threads, mono);
}

void StableDiffusionGGML::compute_ip_adapter_tokens(const sd_image_t& image, float strength) {
    ip_adapter_tokens        = {};
    ip_adapter_uncond_tokens = {};
    ip_adapter_strength      = strength;
    if (ip_adapter == nullptr || clip_vision == nullptr || image.data == nullptr) {
        return;
    }
    auto image_tensor = sd_image_to_tensor(image);
    auto embed        = ip_adapter->is_plus
                            ? get_clip_vision_output(image_tensor, false, 2)
                            : get_clip_vision_output(image_tensor, true, -1);
    if (embed.empty()) {
        return;
    }
    ip_adapter_tokens = ip_adapter->compute(n_threads, embed);
    if (ip_adapter_tokens.empty()) {
        LOG_ERROR("IP-Adapter conditional image projection failed");
        return;
    }
    auto uncond_embed        = sd::Tensor<float>::zeros_like(embed);
    ip_adapter_uncond_tokens = ip_adapter->compute(n_threads, uncond_embed);
    if (ip_adapter_uncond_tokens.empty()) {
        LOG_ERROR("IP-Adapter unconditional image projection failed");
        ip_adapter_tokens = {};
        return;
    }
    LOG_INFO("IP-Adapter: %lld image tokens, strength %.2f",
             (long long)ip_adapter_tokens.shape()[1], strength);
}

std::vector<float> StableDiffusionGGML::process_timesteps(const std::vector<float>& timesteps,
                                                          const sd::Tensor<float>& init_latent,
                                                          const sd::Tensor<float>& denoise_mask,
                                                          int step) {
    if (auto sefi_denoiser = std::dynamic_pointer_cast<SefiFlowDenoiser>(denoiser)) {
        int sched_idx = step > 0 ? step - 1 : 0;
        if (sched_idx >= static_cast<int>(sefi_denoiser->tex_timesteps.size())) {
            sched_idx = static_cast<int>(sefi_denoiser->tex_timesteps.size()) - 1;
        }
        return {sefi_denoiser->sem_timesteps[sched_idx],
                sefi_denoiser->tex_timesteps[sched_idx]};
    }
    if (diffusion_model->get_desc() == "Wan2.2-TI2V-5B") {
        int64_t frame_count = init_latent.shape()[2];
        auto new_timesteps  = std::vector<float>(static_cast<size_t>(frame_count), timesteps[0]);

        if (!denoise_mask.empty() && denoise_mask.dim() >= 4 && denoise_mask.shape()[2] == frame_count) {
            for (int64_t frame = 0; frame < frame_count; ++frame) {
                float value = denoise_mask.dim() == 5 ? denoise_mask.index(0, 0, frame, 0, 0) : denoise_mask.index(0, 0, frame, 0);
                if (value == 0.f) {
                    new_timesteps[static_cast<size_t>(frame)] = 0.f;
                }
            }
        }
        return new_timesteps;
    }
    if (diffusion_model->get_desc() == "Wan2.2-S2V-14B") {
        int64_t frame_count = init_latent.shape()[2];
        return std::vector<float>(static_cast<size_t>(frame_count), timesteps[0]);
    } else {
        return timesteps;
    }
}

std::vector<float> StableDiffusionGGML::process_ltxav_video_timesteps(const std::vector<float>& timesteps,
                                                                      const sd::Tensor<float>& init_latent,
                                                                      const sd::Tensor<float>& denoise_mask) {
    if (timesteps.empty() || denoise_mask.empty() || init_latent.dim() < 4 || denoise_mask.dim() < 4) {
        return timesteps;
    }

    int64_t width  = init_latent.shape()[0];
    int64_t height = init_latent.shape()[1];
    int64_t frames = init_latent.shape()[2];
    if (denoise_mask.shape()[0] != width ||
        denoise_mask.shape()[1] != height ||
        denoise_mask.shape()[2] != frames ||
        denoise_mask.shape()[3] < 1) {
        LOG_WARN("unexpected LTXAV denoise mask shape for timestep processing");
        return timesteps;
    }

    std::vector<float> video_timesteps(static_cast<size_t>(width * height * frames));
    size_t idx = 0;
    for (int64_t t = 0; t < frames; ++t) {
        for (int64_t h = 0; h < height; ++h) {
            for (int64_t w = 0; w < width; ++w) {
                float mask             = denoise_mask.dim() == 5 ? denoise_mask.index(w, h, t, 0, 0)
                                                                 : denoise_mask.index(w, h, t, 0);
                video_timesteps[idx++] = mask * timesteps[0];
            }
        }
    }
    return video_timesteps;
}

void StableDiffusionGGML::preview_image(int step,
                                        const sd::Tensor<float>& latents,
                                        enum SDVersion version,
                                        preview_t preview_mode,
                                        std::function<void(int, int, sd_image_t*, bool, void*)> step_callback,
                                        void* step_callback_data,
                                        bool is_noisy) {
    bool is_video = preview_latent_tensor_is_video(latents);
    uint32_t dim  = is_video ? static_cast<uint32_t>(latents.shape()[3]) : static_cast<uint32_t>(latents.shape()[2]);
    int channels  = get_latent_channel();
    auto _latents = channels != dim ? is_video ? sd::ops::slice(latents, 3, 0, channels)
                                               : sd::ops::slice(latents, 2, 0, channels)
                                    : latents;
    if (preview_mode == PREVIEW_PROJ) {
        int patch_sz                      = 1;
        const float (*latent_rgb_proj)[3] = nullptr;
        float* latent_rgb_bias            = nullptr;
        const float* latent_alpha_proj    = nullptr;
        float latent_alpha_bias           = 1.f;

        if (channels == 128) {
            if (sd_version_uses_flux2_vae(version)) {
                latent_rgb_proj = flux2_latent_rgb_proj;
                latent_rgb_bias = flux2_latent_rgb_bias;
                patch_sz        = 2;
            } else if (version == VERSION_LTXAV) {
                latent_rgb_proj = ltxav_latent_rgb_proj;
                latent_rgb_bias = ltxav_latent_rgb_bias;
            } else {
                LOG_WARN("No latent to RGB projection known for this model");
                return;
            }
        } else if (channels == 64) {
            if (version == VERSION_QWEN_IMAGE_2_1) {
                latent_rgb_proj   = qwen21_latent_rgb_proj;
                latent_rgb_bias   = qwen21_latent_rgb_bias;
                latent_alpha_proj = qwen21_latent_alpha_proj;
                latent_alpha_bias = qwen21_latent_alpha_bias;
            } else {
                LOG_WARN("No latent to RGB projection known for this model");
                return;
            }
        } else if (channels == 48) {
            if (sd_version_is_wan(version)) {
                latent_rgb_proj = wan_22_latent_rgb_proj;
                latent_rgb_bias = wan_22_latent_rgb_bias;
            } else {
                LOG_WARN("No latent to RGB projection known for this model");
                return;
            }
        } else if (channels == 24) {
            if (sd_version_is_minimax_h3(version)) {
                latent_rgb_proj = minimax_latent_rgb_proj;
                latent_rgb_bias = minimax_latent_rgb_bias;
            } else {
                LOG_WARN("No latent to RGB projection known for this model");
                return;
            }
        } else if (channels == 16) {
            if (sd_version_is_sd3(version)) {
                latent_rgb_proj = sd3_latent_rgb_proj;
                latent_rgb_bias = sd3_latent_rgb_bias;
            } else if (sd_version_uses_flux_vae(version)) {
                latent_rgb_proj = flux_latent_rgb_proj;
                latent_rgb_bias = flux_latent_rgb_bias;
            } else if (sd_version_uses_wan_vae(version)) {
                latent_rgb_proj = wan_21_latent_rgb_proj;
                latent_rgb_bias = wan_21_latent_rgb_bias;
            } else {
                LOG_WARN("No latent to RGB projection known for this model");
                return;
            }
        } else if (channels == 4) {
            if (sd_version_is_sdxl(version)) {
                latent_rgb_proj = sdxl_latent_rgb_proj;
                latent_rgb_bias = sdxl_latent_rgb_bias;
            } else if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
                latent_rgb_proj = sd_latent_rgb_proj;
                latent_rgb_bias = sd_latent_rgb_bias;
            } else {
                LOG_WARN("No latent to RGB projection known for this model");
                return;
            }
        } else if (channels != 3) {
            LOG_WARN("No latent to RGB projection known for this model (dim = %d)", dim);
            return;
        }

        uint32_t frames     = is_video ? static_cast<uint32_t>(_latents.shape()[2]) : 1;
        uint32_t img_width  = static_cast<uint32_t>(_latents.shape()[0]) * patch_sz;
        uint32_t img_height = static_cast<uint32_t>(_latents.shape()[1]) * patch_sz;

        uint32_t img_channels = latent_alpha_proj != nullptr ? 4 : 3;
        uint8_t* data         = (uint8_t*)malloc(frames * img_width * img_height * img_channels * sizeof(uint8_t));
        GGML_ASSERT(data != nullptr);
        preview_latent_video(data, _latents, latent_rgb_proj, latent_rgb_bias, patch_sz, latent_alpha_proj, latent_alpha_bias);
        sd_image_t* images = (sd_image_t*)malloc(frames * sizeof(sd_image_t));
        GGML_ASSERT(images != nullptr);
        for (uint32_t i = 0; i < frames; i++) {
            images[i] = {img_width, img_height, img_channels, data + i * img_width * img_height * img_channels};
        }
        step_callback(step, frames, images, is_noisy, step_callback_data);
        free(data);
        free(images);
        return;
    }

    if (preview_mode == PREVIEW_VAE || preview_mode == PREVIEW_TAE) {
        sd::Tensor<float> vae_latents;
        sd::Tensor<float> decoded;
        if (preview_vae) {
            vae_latents = preview_vae->diffusion_to_vae_latents(_latents);
            decoded     = preview_vae->decode(n_threads, vae_latents, vae_tiling_params, is_video, circular_x, circular_y, true);
        } else {
            vae_latents = first_stage_model->diffusion_to_vae_latents(_latents);
            decoded     = first_stage_model->decode(n_threads, vae_latents, vae_tiling_params, is_video, circular_x, circular_y, true);
        }
        if (decoded.empty()) {
            LOG_ERROR("preview decode failed at step %d", step);
            return;
        }

        is_video           = preview_latent_tensor_is_video(decoded);
        uint32_t frames    = is_video ? static_cast<uint32_t>(decoded.shape()[2]) : 1;
        sd_image_t* images = (sd_image_t*)malloc(frames * sizeof(sd_image_t));
        GGML_ASSERT(images != nullptr);
        for (uint32_t i = 0; i < frames; ++i) {
            images[i] = tensor_to_sd_image(decoded, static_cast<int>(i));
        }

        step_callback(step, frames, images, is_noisy, step_callback_data);
        for (uint32_t i = 0; i < frames; ++i) {
            free(images[i].data);
        }
        free(images);
        return;
    }

    if (preview_mode != PREVIEW_NONE) {
        LOG_WARN("Unsupported preview mode: %d", static_cast<int>(preview_mode));
    }
}

std::vector<float> StableDiffusionGGML::prepare_sample_timesteps(float sigma,
                                                                 int shifted_timestep) {
    float t = denoiser->sigma_to_t(sigma);
    if (shifted_timestep > 0) {
        float shifted_t_float = t * (float(shifted_timestep) / float(TIMESTEPS));
        int64_t shifted_t     = static_cast<int64_t>(roundf(shifted_t_float));
        shifted_t             = std::max((int64_t)0, std::min((int64_t)(TIMESTEPS - 1), shifted_t));
        LOG_VERBOSE("shifting timestep from %.2f to %" PRId64 " (sigma: %.4f)", t, shifted_t, sigma);
        return std::vector<float>{(float)shifted_t};
    }
    if (sd_version_is_anima(version)) {
        return std::vector<float>{t / static_cast<float>(TIMESTEPS)};
    }
    if (sd_version_is_boogu_image(version)) {
        return std::vector<float>{t / static_cast<float>(TIMESTEPS)};
    }
    if (version == VERSION_HIDREAM_O1) {
        return std::vector<float>{1.0f - (t / static_cast<float>(TIMESTEPS))};
    }
    if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version) || sd_version_is_ideogram4(version) || version == VERSION_MING_IMAGE) {
        return std::vector<float>{1000.f - t};
    }
    return std::vector<float>{t};
}

void StableDiffusionGGML::adjust_sample_step_scalings(int shifted_timestep,
                                                      const std::vector<float>& timesteps_vec,
                                                      float c_in,
                                                      float* c_skip,
                                                      float* c_out) {
    GGML_ASSERT(c_skip != nullptr);
    GGML_ASSERT(c_out != nullptr);
    if (shifted_timestep <= 0) {
        return;
    }

    int64_t shifted_t_idx              = static_cast<int64_t>(roundf(timesteps_vec[0]));
    float shifted_sigma                = denoiser->t_to_sigma((float)shifted_t_idx);
    std::vector<float> shifted_scaling = denoiser->get_scalings(shifted_sigma);
    float shifted_c_skip               = shifted_scaling[0];
    float shifted_c_out                = shifted_scaling[1];
    float shifted_c_in                 = shifted_scaling[2];

    *c_skip = shifted_c_skip * c_in / shifted_c_in;
    *c_out  = shifted_c_out;
}

StableDiffusionGGML::SamplePreviewContext StableDiffusionGGML::prepare_sample_preview_context() {
    return SamplePreviewContext{sd_get_preview_callback(),
                                sd_get_preview_callback_data(),
                                sd_get_preview_mode()};
}

void StableDiffusionGGML::report_sample_progress(int step,
                                                 size_t total_steps,
                                                 bool terminal_sigma_is_zero,
                                                 int64_t* last_progress_us) {
    if (sd::preview::sample_step_is_complete(step, total_steps, terminal_sigma_is_zero)) {
        int64_t now        = ggml_time_us();
        int showstep       = std::abs(step);
        float step_seconds = last_progress_us != nullptr && *last_progress_us > 0
                                 ? (now - *last_progress_us) / 1000000.f
                                 : 0.f;
        pretty_progress(showstep, (int)total_steps, step_seconds);
        if (last_progress_us != nullptr) {
            *last_progress_us = now;
        }
    }
}

static sd::Tensor<float> batch_two_condition_tensors(const sd::Tensor<float>& a, const sd::Tensor<float>& b) {
    if (a.empty() || b.empty() || a.dim() != b.dim()) {
        return {};
    }
    if (a.dim() == 1) {
        if (a.shape() != b.shape()) {
            return {};
        }
        auto batched = sd::ops::concat(a, b, 0);
        batched.reshape_({a.shape()[0], 2});
        return batched;
    }
    const int64_t batch_dim = a.dim() - 1;
    for (int64_t d = 0; d < batch_dim; d++) {
        if (a.shape()[d] != b.shape()[d]) {
            return {};
        }
    }
    return sd::ops::concat(a, b, static_cast<size_t>(batch_dim));
}

bool StableDiffusionGGML::set_control_type(sd_control_type_t type) {
    return control_net == nullptr || control_net->set_control_type(type);
}

void StableDiffusionGGML::compute_sample_controls(const sd::Tensor<float>& control_image,
                                                  const sd::Tensor<float>& noised_input,
                                                  const sd::Tensor<float>& timesteps_tensor,
                                                  const SDCondition& condition,
                                                  std::vector<sd::Tensor<float>>* controls) {
    GGML_ASSERT(controls != nullptr);
    controls->clear();
    if (control_image.empty() || control_net == nullptr) {
        return;
    }

    auto control_result = control_net->compute(n_threads,
                                               noised_input,
                                               control_image,
                                               timesteps_tensor,
                                               condition.c_crossattn,
                                               condition.c_vector);
    if (!control_result.has_value()) {
        LOG_ERROR("controlnet compute failed");
        return;
    }

    *controls = std::move(*control_result);
}

sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionModelRunner>& work_diffusion_model,
                                              bool inverse_noise_scaling,
                                              const sd::Tensor<float>& init_latent,
                                              sd::Tensor<float> noise,
                                              const SDCondition& cond,
                                              const SDCondition& uncond,
                                              const SDCondition& img_uncond,
                                              const sd::Tensor<float>& control_image,
                                              float control_strength,
                                              const sd_guidance_params_t& guidance,
                                              float eta,
                                              int shifted_timestep,
                                              sample_method_t method,
                                              bool is_flow_denoiser,
                                              const char* extra_sample_args,
                                              const std::vector<float>& sigmas,
                                              const std::vector<sd::Tensor<float>>& ref_latents,
                                              const RefImageParams& ref_image_params,
                                              const sd::Tensor<float>& denoise_mask,
                                              const sd::Tensor<float>& vace_context,
                                              float vace_strength,
                                              int audio_length,
                                              float frame_rate,
                                              const sd_cache_params_t* cache_params,
                                              bool preview_final_step,
                                              const sd::Tensor<float>& video_positions) {
    struct RunnerEndOnExit {
        GGMLRunner* runner = nullptr;
        ~RunnerEndOnExit() {
            if (runner != nullptr) {
                runner->runner_end();
            }
        }
    };
    RunnerEndOnExit sample_diffusion_runner_end{work_diffusion_model.get()};

    // These inputs are immutable for this sampling run. Extensions may replace or
    // modify them per step, so those paths need an explicit stability contract first.
    const bool cache_qwen_prefix = version == VERSION_QWEN_IMAGE_2_1 &&
                                   std::none_of(generation_extensions.begin(), generation_extensions.end(),
                                                [](const auto& extension) { return extension->is_enabled(); });
    using QwenPrefixInputs       = std::tuple<const sd::Tensor<float>*, const sd::Tensor<int32_t>*,
                                              const std::vector<sd::Tensor<float>>*>;
    std::vector<QwenPrefixInputs> qwen_prefix_inputs;

    RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr};

    const bool apply_denoise_mask = !denoise_mask.empty() &&
                                    std::any_of(denoise_mask.values().begin(), denoise_mask.values().end(),
                                                [](float value) { return value != 1.f; });

    std::vector<int> skip_layers(guidance.slg.layers, guidance.slg.layers + guidance.slg.layer_count);
    float cfg_scale     = guidance.txt_cfg;
    float img_cfg_scale = guidance.img_cfg;
    float slg_scale     = guidance.slg.scale;
    bool slg_uncond     = sd::guidance::parse_skip_layer_guidance_uncond_arg(extra_sample_args);

    std::vector<float> guidance_schedule = sd::guidance::parse_guidance_schedule(extra_sample_args);
    if (!guidance_schedule.empty() && guidance_schedule.size() != sigmas.size() - 1) {
        if (guidance_schedule.size() > sigmas.size()) {
            LOG_WARN("guidance_schedule length (%zu) is greater than number of steps (%zu)", guidance_schedule.size(), sigmas.size() - 1);
            LOG_WARN("truncating guidance_schedule to match step count");
            guidance_schedule.resize(sigmas.size() - 1);
        } else {
            LOG_INFO("padding guidance_schedule with cfg_scale");
            while (guidance_schedule.size() < sigmas.size() - 1) {
                guidance_schedule.push_back(cfg_scale);
            }
        }
    }

    if (!guidance_schedule.empty()) {
        std::string schedule_str = "[";
        for (size_t i = 0; i < guidance_schedule.size(); ++i) {
            schedule_str += std::to_string(guidance_schedule[i]);
            if (i < guidance_schedule.size() - 1) {
                schedule_str += ", ";
            }
        }
        schedule_str += "]";
        LOG_VERBOSE("using guidance schedule: %s", schedule_str.c_str());
    }

    sd_sample::SampleCacheRuntime cache_runtime = sd_sample::init_sample_cache_runtime(version,
                                                                                       cache_params,
                                                                                       denoiser.get(),
                                                                                       sigmas);

    bool needs_uncond_denoised = method == EULER_CFG_PP_SAMPLE_METHOD || method == EULER_A_CFG_PP_SAMPLE_METHOD;
    // Spectrum cache is not supported for CFG++ samplers
    if (needs_uncond_denoised) {
        if (cache_runtime.spectrum_enabled) {
            LOG_WARN("Spectrum cache requested but not supported for CFG++ samplers");
            cache_runtime.spectrum_enabled = false;
        }
    }

    size_t steps                = sigmas.size() - 1;
    bool terminal_sigma_is_zero = sigmas.back() == 0.f;
    bool has_skiplayer          = (slg_scale != 0.0f || slg_uncond) && !skip_layers.empty();
    if (has_skiplayer && !sd_version_is_dit(version)) {
        has_skiplayer = false;
        LOG_WARN("SLG is incompatible with this model type");
    }
    sd::guidance::AdaptiveProjectedGuidanceParams apg_params = sd::guidance::parse_adaptive_projected_guidance_args(extra_sample_args);
    bool use_apg_guidance                                    = sd::guidance::is_adaptive_projected_guidance_enabled(apg_params);
    if (use_apg_guidance) {
        LOG_INFO("using Adaptive Projected Guidance (APG)");
    }
    sd::guidance::ClassifierFreeGuidance classifier_free_guidance(cfg_scale, img_cfg_scale);
    sd::guidance::AdaptiveProjectedGuidance adaptive_projected_guidance(cfg_scale, img_cfg_scale, apg_params);
    const sd::guidance::BaseGuidance& primary_guidance = use_apg_guidance
                                                             ? static_cast<const sd::guidance::BaseGuidance&>(adaptive_projected_guidance)
                                                             : static_cast<const sd::guidance::BaseGuidance&>(classifier_free_guidance);
    sd::guidance::SkipLayerGuidance skip_layer_guidance(has_skiplayer ? skip_layers : std::vector<int>(),
                                                        has_skiplayer ? slg_scale : 0.0f,
                                                        guidance.slg.layer_start,
                                                        guidance.slg.layer_end);

    if (version == VERSION_HIDREAM_O1 && !noise.empty()) {
        noise *= eta;
    }

    int64_t last_progress_us = ggml_time_us();
    sd_begin_preview_pass(static_cast<int>(steps));
    SamplePreviewContext preview = prepare_sample_preview_context();

    sd::Tensor<float> processed_init_latent       = denoiser->process_latent_in(init_latent);
    const sd::Tensor<float>& sampling_init_latent = processed_init_latent.empty()
                                                        ? init_latent
                                                        : processed_init_latent;
    sd::Tensor<float> x_t                         = !noise.empty()
                                                        ? denoiser->noise_scaling(sigmas[0], noise, sampling_init_latent)
                                                        : sampling_init_latent;
    sd::Tensor<float> denoised                    = x_t;

    auto denoise = [&](const sd::Tensor<float>& x, float sigma, int step) -> sd::guidance::GuiderOutput {
        if (get_cancel_flag() == SD_CANCEL_ALL) {
            LOG_VERBOSE("cancelling generation");
            return {};
        }

        if (step == 1 || step == -1) {
            pretty_progress(0, (int)steps, 0);
            last_progress_us = ggml_time_us();
        }

        std::vector<float> scaling = denoiser->get_scalings(sigma);
        GGML_ASSERT(scaling.size() == 3);
        float c_skip = scaling[0];
        float c_out  = scaling[1];
        float c_in   = scaling[2];

        bool preview_needed = preview.callback != nullptr &&
                              sd::preview::should_preview_sample_step(step,
                                                                      steps,
                                                                      terminal_sigma_is_zero,
                                                                      sd_get_preview_interval(),
                                                                      preview_final_step);

        std::vector<float> base_timesteps_vec = prepare_sample_timesteps(sigma, shifted_timestep);
        std::vector<float> timesteps_vec      = base_timesteps_vec;
        sd::Tensor<float> audio_timesteps_tensor;
        if (sd_version_is_ltxav(version) && !denoise_mask.empty()) {
            timesteps_vec          = process_ltxav_video_timesteps(base_timesteps_vec, sampling_init_latent, denoise_mask);
            audio_timesteps_tensor = sd::Tensor<float>({static_cast<int64_t>(base_timesteps_vec.size())}, base_timesteps_vec);
        } else {
            timesteps_vec = process_timesteps(timesteps_vec, sampling_init_latent, denoise_mask, step);
        }
        const std::vector<float>& scaling_timesteps_vec = (sd_version_is_ltxav(version) && !denoise_mask.empty())
                                                              ? base_timesteps_vec
                                                              : timesteps_vec;
        adjust_sample_step_scalings(shifted_timestep, scaling_timesteps_vec, c_in, &c_skip, &c_out);

        sd::Tensor<float> timesteps_tensor({static_cast<int64_t>(timesteps_vec.size())}, timesteps_vec);
        sd::Tensor<float> guidance_tensor({1}, std::vector<float>{guidance.distilled_guidance});
        sd::Tensor<float> hunyuan_timestep_r_tensor;
        if (sd_version_is_hunyuan_video(version) && step + 1 < sigmas.size()) {
            hunyuan_timestep_r_tensor = sd::Tensor<float>::from_vector({sigmas[step + 1]});
        }
        sd::Tensor<float> flux3_target_timestep;
        if (version == VERSION_FLUX3) {
            // The sign marks intermediate evaluations; the magnitude is the one-based step.
            const size_t target_index = static_cast<size_t>(std::abs(step));
            GGML_ASSERT(target_index > 0 && target_index < sigmas.size());
            flux3_target_timestep = sd::Tensor<float>::from_vector({sigmas[target_index]});
        }
        sd::Tensor<float> noised_input = x * c_in;
        if (apply_denoise_mask && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version) || sd_version_is_lingbot_video(version))) {
            noised_input = noised_input * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
        }

        if (cache_runtime.spectrum_enabled && cache_runtime.spectrum.should_predict()) {
            cache_runtime.spectrum.predict(&denoised);
            if (apply_denoise_mask) {
                denoised = denoised * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
            }
            if (preview_needed && sd_should_preview_denoised()) {
                preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false);
            }
            report_sample_progress(step, steps, terminal_sigma_is_zero, &last_progress_us);
            sd::guidance::GuiderOutput output;
            output.pred = denoised;
            return output;
        }

        if (preview_needed && sd_should_preview_noisy()) {
            preview_image(step, noised_input, version, preview.mode, preview.callback, preview.data, true);
        }

        sd::Tensor<float> cond_out;
        sd::Tensor<float> uncond_out;
        sd::Tensor<float> img_uncond_out;
        sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma);
        std::vector<sd::Tensor<float>> controls;
        DiffusionParams diffusion_params;
        diffusion_params.x                = &noised_input;
        diffusion_params.timesteps        = &timesteps_tensor;
        diffusion_params.ref_image_params = ref_image_params;
        sd::guidance::GuidanceInput step_guidance_input;
        step_guidance_input.step          = step;
        step_guidance_input.schedule_size = sigmas.size();
        bool is_skiplayer_step            = skip_layer_guidance.is_enabled_for_step(step_guidance_input);

        compute_sample_controls(control_image,
                                noised_input,
                                timesteps_tensor,
                                cond,
                                &controls);
        bool uncond_controls_ready = false;

        static const std::vector<sd::Tensor<float>> empty_ref_latents;
        bool uncond_without_ref_latents = !img_uncond.empty() &&
                                          !ref_latents.empty() &&
                                          sd_version_supports_ref_latent_img_cfg(version);

        auto run_condition = [&](const SDCondition& condition,
                                 const sd::Tensor<float>* c_concat_override                 = nullptr,
                                 const std::vector<int>* local_skip_layers                  = nullptr,
                                 const std::vector<sd::Tensor<float>>* ref_latents_override = nullptr,
                                 bool use_uncond_ip                                         = false) -> sd::Tensor<float> {
            diffusion_params.context     = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn;
            diffusion_params.c_concat    = c_concat_override != nullptr ? c_concat_override : (condition.c_concat.empty() ? nullptr : &condition.c_concat);
            diffusion_params.y           = condition.c_vector.empty() ? nullptr : &condition.c_vector;
            diffusion_params.ref_latents = ref_latents_override != nullptr ? ref_latents_override : (condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images);

            if (sd_version_is_unet(version)) {
                int nvf = -1;
                if (config_->animatediff_loaded && noised_input.dim() >= 4 && noised_input.shape()[3] > 1) {
                    nvf = static_cast<int>(noised_input.shape()[3]);
                }
                UNetDiffusionExtra unet_extra{nvf, &controls, control_strength};
                const auto& ip_tokens = use_uncond_ip ? ip_adapter_uncond_tokens : ip_adapter_tokens;
                if (!ip_tokens.empty()) {
                    unet_extra.ip_context = &ip_tokens;
                    unet_extra.ip_scale   = ip_adapter_strength;
                }
                diffusion_params.extra = unet_extra;
            } else if (sd_version_is_sd3(version) || version == VERSION_MOSS_IMAGE) {
                diffusion_params.extra = SkipLayerDiffusionExtra{local_skip_layers};
            } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) {
                diffusion_params.extra = FluxDiffusionExtra{&guidance_tensor,
                                                            local_skip_layers};
            } else if (sd_version_is_flux3(version)) {
                diffusion_params.extra = Flux3DiffusionExtra{flux3_target_timestep.empty() ? nullptr : &flux3_target_timestep};
            } else if (version == VERSION_QWEN_IMAGE_2_1) {
                diffusion_params.extra = QwenImage21DiffusionExtra{&condition.c_token_types};
            } else if (sd_version_is_anima(version)) {
                diffusion_params.extra = AnimaDiffusionExtra{condition.c_t5_ids.empty() ? nullptr : &condition.c_t5_ids,
                                                             condition.c_t5_weights.empty() ? nullptr : &condition.c_t5_weights};
            } else if (sd_version_is_wan(version)) {
                diffusion_params.extra = WanDiffusionExtra{vace_context.empty() ? nullptr : &vace_context,
                                                           vace_strength,
                                                           condition.c_ref_audios.empty() ? nullptr : &condition.c_ref_audios[0]};
            } else if (sd_version_is_hunyuan_video(version)) {
                diffusion_params.extra = HunyuanVideoDiffusionExtra{
                    &guidance_tensor,
                    condition.extra_c_crossattns.empty() ? nullptr : &condition.extra_c_crossattns[0],
                    condition.c_vector.empty() ? nullptr : &condition.c_vector,
                    hunyuan_timestep_r_tensor.empty() ? nullptr : &hunyuan_timestep_r_tensor};
            } else if (version == VERSION_HIDREAM_O1) {
                diffusion_params.extra = HiDreamO1DiffusionExtra{
                    condition.c_input_ids.empty() ? nullptr : &condition.c_input_ids,
                    condition.c_position_ids.empty() ? nullptr : &condition.c_position_ids,
                    condition.c_token_types.empty() ? nullptr : &condition.c_token_types,
                    condition.c_vinput_mask.empty() ? nullptr : &condition.c_vinput_mask,
                    condition.c_image_embeds.empty() ? nullptr : &condition.c_image_embeds};
            } else if (version == VERSION_MING_IMAGE) {
                diffusion_params.extra = MingImageDiffusionExtra{
                    condition.extra_c_crossattns.empty() ? nullptr : &condition.extra_c_crossattns[0]};
            } else if (sd_version_is_llada_image(version)) {
                diffusion_params.extra = LLaDAImageDiffusionExtra{
                    condition.extra_c_crossattns.empty() ? nullptr : &condition.extra_c_crossattns[0]};
            } else if (sd_version_is_minimax_h3(version)) {
                diffusion_params.extra = MiniMaxH3DiffusionExtra{
                    condition.c_token_types.empty() ? nullptr : &condition.c_token_types,
                    condition.c_position_ids.empty() ? nullptr : &condition.c_position_ids,
                    condition.c_ref_audios.empty() ? nullptr : &condition.c_ref_audios,
                    condition.c_reference_blocks.empty() ? nullptr : &condition.c_reference_blocks,
                    version == VERSION_P_H3_IMAGE ? 2 : audio_length,
                    std::isfinite(active_flow_shift) ? active_flow_shift : 12.f,
                    version == VERSION_P_H3_IMAGE ? 1.f : 3.f};
            } else if (sd_version_is_ltxav(version)) {
                diffusion_params.extra = LTXAVDiffusionExtra{
                    nullptr,
                    audio_timesteps_tensor.empty() ? nullptr : &audio_timesteps_tensor,
                    audio_length,
                    frame_rate,
                    video_positions.empty() ? nullptr : &video_positions};
            } else if (sd_version_is_minit2i(version)) {
                diffusion_params.extra = MiniT2IDiffusionExtra{
                    condition.c_vector.empty() ? nullptr : &condition.c_vector};
            } else if (sd_version_is_sensenova_u1(version)) {
                diffusion_params.extra = SenseNovaU1DiffusionExtra{
                    condition.c_input_ids.empty() ? nullptr : &condition.c_input_ids};
            } else if (sd_version_is_z_image(version)) {
                if (!control_image.empty()) {
                    diffusion_params.extra = ZImageDiffusionExtra{&control_image, control_strength};
                }
            } else {
                diffusion_params.extra = std::monostate{};
            }

            sd::Tensor<float> cached_output;
            if (step_cache.before_condition(&condition, noised_input, &cached_output)) {
                return std::move(cached_output);
            }

            // A re-enabled condition can miss the cache even when the positive pass was reused.
            if (!uncond_controls_ready && !uncond.empty() &&
                (&condition == &uncond || &condition == &img_uncond)) {
                compute_sample_controls(control_image,
                                        noised_input,
                                        timesteps_tensor,
                                        uncond,
                                        &controls);
                uncond_controls_ready = true;
            }

            for (const auto& extension : generation_extensions) {
                extension->before_diffusion(diffusion_params, step);
            }

            if (cache_qwen_prefix) {
                auto* extra = std::get_if<QwenImage21DiffusionExtra>(&diffusion_params.extra);
                if (extra != nullptr) {
                    auto key         = std::make_tuple(diffusion_params.context, extra->image_slots,
                                                       diffusion_params.ref_image_params.pass_to_dit ? diffusion_params.ref_latents : nullptr);
                    auto entry       = std::find(qwen_prefix_inputs.begin(), qwen_prefix_inputs.end(), key);
                    extra->prefix_id = static_cast<uint64_t>(entry - qwen_prefix_inputs.begin()) + 1;
                    if (entry == qwen_prefix_inputs.end()) {
                        qwen_prefix_inputs.push_back(key);
                    }
                }
            }
            auto output_opt = work_diffusion_model->compute(n_threads, diffusion_params);
            if (output_opt.empty()) {
                LOG_ERROR("diffusion model compute failed");
                return sd::Tensor<float>();
            }

            step_cache.after_condition(&condition, noised_input, output_opt);
            return output_opt;
        };

        auto run_batched_condition = [&](const SDCondition& condition,
                                         const sd::Tensor<float>* c_concat_override) -> sd::Tensor<float> {
            const sd::Tensor<float>& condition_concat =
                c_concat_override != nullptr ? *c_concat_override : condition.c_concat;

            sd::Tensor<float> batched_context = batch_two_condition_tensors(condition.c_crossattn, uncond.c_crossattn);
            sd::Tensor<float> batched_y       = batch_two_condition_tensors(condition.c_vector, uncond.c_vector);
            sd::Tensor<float> batched_concat  = batch_two_condition_tensors(condition_concat, uncond.c_concat);
            if (!condition.c_crossattn.empty() && batched_context.empty()) {
                return {};
            }
            if ((!condition.c_vector.empty() || !uncond.c_vector.empty()) && batched_y.empty()) {
                return {};
            }
            if ((!condition_concat.empty() || !uncond.c_concat.empty()) && batched_concat.empty()) {
                return {};
            }

            std::vector<sd::Tensor<float>> uncond_controls;
            compute_sample_controls(control_image, noised_input, timesteps_tensor, uncond, &uncond_controls);
            if (controls.size() != uncond_controls.size()) {
                return {};
            }
            std::vector<sd::Tensor<float>> batched_controls;
            batched_controls.reserve(controls.size());
            for (size_t i = 0; i < controls.size(); i++) {
                sd::Tensor<float> batched_control = batch_two_condition_tensors(controls[i], uncond_controls[i]);
                if (batched_control.empty()) {
                    return {};
                }
                batched_controls.push_back(std::move(batched_control));
            }

            sd::Tensor<float> batched_x =
                sd::ops::concat(noised_input, noised_input, static_cast<size_t>(noised_input.dim() - 1));

            DiffusionParams batched_params = diffusion_params;
            batched_params.x               = &batched_x;
            batched_params.context         = batched_context.empty() ? nullptr : &batched_context;
            batched_params.c_concat        = batched_concat.empty() ? nullptr : &batched_concat;
            batched_params.y               = batched_y.empty() ? nullptr : &batched_y;
            batched_params.ref_latents     = nullptr;
            batched_params.extra           = UNetDiffusionExtra{1, &batched_controls, control_strength};

            sd::Tensor<float> output = work_diffusion_model->compute(n_threads, batched_params);
            if (output.empty()) {
                LOG_ERROR("batched diffusion model compute failed");
            }
            return output;
        };

        const SDCondition* positive_condition      = &cond;
        const sd::Tensor<float>* c_concat_override = nullptr;
        for (const auto& extension : generation_extensions) {
            const SDCondition& next_condition = extension->before_condition(step, *positive_condition);
            if (&next_condition != positive_condition) {
                positive_condition = &next_condition;
                if (positive_condition != &cond) {
                    c_concat_override = cond.c_concat.empty() ? nullptr : &cond.c_concat;
                }
                break;
            }
        }

        float effective_guidance_scale = guidance_schedule.empty()
                                             ? cfg_scale
                                             : guidance_schedule[guidance_schedule.size() - 1 - step];

        float image_guidance_scale = img_cfg_scale;

        constexpr float kEpsilon = 1e-5f;

        bool skip_uncond = false;
        if (!uncond.empty() && !needs_uncond_denoised && !use_apg_guidance) {
            if (!img_uncond.empty()) {
                skip_uncond = std::abs(image_guidance_scale - effective_guidance_scale) < kEpsilon;
            } else {
                skip_uncond = std::abs(effective_guidance_scale - 1.0f) < kEpsilon;
            }
        }

        bool skip_img_uncond = false;
        if (!img_uncond.empty() && !needs_uncond_denoised && !use_apg_guidance) {
            if (!uncond.empty()) {
                skip_img_uncond = std::abs(image_guidance_scale - 1.0f) < kEpsilon;
            } else {
                skip_img_uncond = std::abs(effective_guidance_scale - 1.0f) < kEpsilon;
            }
        }

        const bool batch_cfg_ok = config_->params.batched_cfg &&
                                  sd_version_is_unet(version) &&
                                  !uncond.empty() &&
                                  img_uncond.empty() &&
                                  !skip_uncond &&
                                  !cache_runtime.ucache_enabled() &&
                                  !(is_skiplayer_step && slg_uncond) &&
                                  ip_adapter_tokens.empty() &&
                                  ip_adapter_uncond_tokens.empty() &&
                                  !config_->animatediff_loaded &&
                                  (noised_input.dim() < 4 || noised_input.shape()[3] <= 1) &&
                                  std::none_of(generation_extensions.begin(),
                                               generation_extensions.end(),
                                               [](const std::shared_ptr<GenerationExtension>& extension) {
                                                   return extension->is_enabled();
                                               });

        if (batch_cfg_ok) {
            sd::Tensor<float> batched_out = run_batched_condition(*positive_condition, c_concat_override);
            if (!batched_out.empty() && batched_out.dim() >= 4 && batched_out.shape()[3] == 2) {
                auto parts = sd::ops::chunk(batched_out, 2, 3);
                cond_out   = std::move(parts[0]);
                uncond_out = std::move(parts[1]);
            }
        }

        if (cond_out.empty()) {
            cond_out = run_condition(*positive_condition, c_concat_override);
            if (cond_out.empty()) {
                return {};
            }
        }

        if (uncond_out.empty() && !uncond.empty()) {
            if (!skip_uncond) {
                const std::vector<int>* uncond_skip_layers = nullptr;
                const std::vector<int> moss_skip_layers    = {22};
                if (is_skiplayer_step && slg_uncond) {
                    LOG_VERBOSE("Skipping layers at uncond step %d\n", step);
                    uncond_skip_layers = &skip_layer_guidance.layers();
                } else if (version == VERSION_MOSS_IMAGE) {
                    uncond_skip_layers = &moss_skip_layers;
                }
                uncond_out = run_condition(uncond,
                                           uncond.c_concat.empty() ? nullptr : &uncond.c_concat,
                                           uncond_skip_layers,
                                           nullptr,
                                           true);
                if (uncond_out.empty()) {
                    return {};
                }
            } else {
                step_cache.invalidate_condition(&uncond);
            }
        }

        if (!img_uncond.empty()) {
            if (!skip_img_uncond) {
                img_uncond_out = run_condition(img_uncond,
                                               img_uncond.c_concat.empty() ? nullptr : &img_uncond.c_concat,
                                               nullptr,
                                               uncond_without_ref_latents ? &empty_ref_latents : nullptr,
                                               true);
                if (img_uncond_out.empty()) {
                    return {};
                }
            } else {
                step_cache.invalidate_condition(&img_uncond);
            }
        }
        sd::guidance::GuidanceInput guidance_input;
        guidance_input.step            = step;
        guidance_input.schedule_size   = sigmas.size();
        guidance_input.pred_cond       = &cond_out;
        guidance_input.pred_uncond     = uncond_out.empty() ? nullptr : &uncond_out;
        guidance_input.pred_img_uncond = img_uncond_out.empty() ? nullptr : &img_uncond_out;

        sd::guidance::GuiderOutput guided = primary_guidance.forward(guidance_input, {}, effective_guidance_scale);
        if (guided.pred.empty()) {
            return {};
        }

        if (is_skiplayer_step && slg_scale != 0.0f) {
            LOG_VERBOSE("Skipping layers at step %d\n", step);
            if (!step_cache.is_step_skipped()) {
                guidance_input.predict_skip_layer = [&]() -> sd::Tensor<float> {
                    return run_condition(cond,
                                         cond.c_concat.empty() ? nullptr : &cond.c_concat,
                                         &skip_layer_guidance.layers());
                };
            }
        }

        guided = skip_layer_guidance.forward(guidance_input, std::move(guided));
        if (guided.pred.empty()) {
            return {};
        }

        denoised = guided.pred * c_out + x * c_skip;
        sd::guidance::GuiderOutput output;
        output.pred = denoised;
        if (needs_uncond_denoised) {
            const sd::Tensor<float>& base_uncond = !img_uncond_out.empty()
                                                       ? img_uncond_out
                                                       : (!uncond_out.empty() ? uncond_out : cond_out);
            output.pred_uncond                   = base_uncond * c_out + x * c_skip;
        }
        if (cache_runtime.spectrum_enabled) {
            cache_runtime.spectrum.update(denoised);
        }
        if (apply_denoise_mask) {
            denoised = denoised * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
        }
        if (preview_needed && sd_should_preview_denoised()) {
            preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false);
        }
        report_sample_progress(step, steps, terminal_sigma_is_zero, &last_progress_us);
        output.pred = denoised;
        return output;
    };

    {
        std::ostringstream oss;
        oss << std::setprecision(std::numeric_limits<float>::max_digits10);
        for (size_t i = 0; i < sigmas.size(); ++i) {
            if (i != 0) {
                oss << ",";
            }
            oss << sigmas[i];
        }
        LOG_DEBUG("sigmas: %s", oss.str().c_str());
    }

    auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser, extra_sample_args, denoiser);
    if (x0_opt.empty()) {
        LOG_ERROR("Diffusion model sampling failed");
        if (control_net) {
            control_net->free_control_ctx();
        }
        return {};
    }

    auto x0 = std::move(x0_opt);
    sd_sample::log_sample_cache_summary(cache_runtime, steps);
    if (inverse_noise_scaling) {
        x0 = denoiser->inverse_noise_scaling(sigmas[sigmas.size() - 1], x0);
    }
    x0 = denoiser->process_latent_out(std::move(x0));

    if (control_net) {
        control_net->free_control_ctx();
    }
    return x0;
}

int StableDiffusionGGML::get_vae_scale_factor() {
    if (sd_version_is_pid(version)) {
        return 1;
    }
    return first_stage_model->get_scale_factor();
}

int StableDiffusionGGML::get_diffusion_model_down_factor() {
    return find_model_definition(version)->geometry.diffusion_down_factor;
}

int StableDiffusionGGML::get_latent_channel() {
    return find_model_definition(version)->geometry.latent_channels;
}

int StableDiffusionGGML::get_image_channels() const {
    return find_model_definition(version)->geometry.image_channels;
}

int StableDiffusionGGML::get_image_seq_len(int h, int w) {
    int vae_scale_factor = get_vae_scale_factor();
    return (h / vae_scale_factor) * (w / vae_scale_factor);
}

sd::Tensor<float> StableDiffusionGGML::generate_init_latent(int width,
                                                            int height,
                                                            int frames,
                                                            bool video) {
    int vae_scale_factor = get_vae_scale_factor();
    int W                = width / vae_scale_factor;
    int H                = height / vae_scale_factor;
    int T                = video_frames_to_latent_frames(frames);
    int C                = get_latent_channel();
    if (video) {
        return sd::zeros<float>({W, H, T, C, 1});
    }
    return sd::zeros<float>({W, H, C, 1});
}

int StableDiffusionGGML::video_frames_to_latent_frames(int frames) {
    int latent_frames = frames;
    if (sd_version_is_ltxav(version)) {
        latent_frames = ((frames - 1) / 8) + 1;
    } else if (sd_version_is_minimax_h3(version)) {
        latent_frames = frames <= 5 ? 2 : ((frames - 5) / 17) * 5 + 2;
    } else if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_hunyuan_video(version)) {
        latent_frames = ((frames - 1) / 4) + 1;
    }
    return latent_frames;
}

int StableDiffusionGGML::latent_frames_to_video_frames(int latent_frames) {
    if (latent_frames <= 0) {
        return latent_frames;
    }
    if (sd_version_is_ltxav(version)) {
        return (latent_frames - 1) * 8 + 1;
    }
    if (sd_version_is_minimax_h3(version)) {
        return latent_frames <= 2 ? 5 : ((latent_frames - 2) / 5) * 17 + 5;
    }
    if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_hunyuan_video(version)) {
        return (latent_frames - 1) * 4 + 1;
    }
    return latent_frames;
}

int StableDiffusionGGML::align_video_frames(int frames) {
    if (sd_version_is_minimax_h3(version)) {
        frames = std::max(frames, 5);
        while (frames % 17 != 5) {
            ++frames;
        }
        return frames;
    }
    return latent_frames_to_video_frames(video_frames_to_latent_frames(frames));
}

sd::Tensor<float> StableDiffusionGGML::encode_to_vae_latents(const sd::Tensor<float>& x) {
    auto latents = first_stage_model->encode(n_threads, x, vae_tiling_params, circular_x, circular_y);
    if (latents.empty()) {
        return {};
    }
    latents = first_stage_model->vae_output_to_latents(latents, rng);
    return latents;
}

sd::Tensor<float> StableDiffusionGGML::encode_first_stage(const sd::Tensor<float>& x) {
    auto latents = encode_to_vae_latents(x);
    if (latents.empty()) {
        return {};
    }
    if (version != VERSION_SD1_PIX2PIX) {
        latents = first_stage_model->vae_to_diffusion_latents(latents);
    }
    return latents;
}

sd::Tensor<float> StableDiffusionGGML::decode_first_stage(const sd::Tensor<float>& x, bool decode_video) {
    if (sd_version_is_pid(version) || sd_version_is_minit2i(version) || version == VERSION_IRIS) {
        return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f);
    }
    auto image_latents                = version == VERSION_P_H3_IMAGE && x.dim() == 5
                                            ? sd::ops::slice(x, 3, 0, get_latent_channel())
                                            : x;
    auto latents                      = first_stage_model->diffusion_to_vae_latents(image_latents);
    auto tiling_params                = first_stage_model->resolve_tiling_params(vae_tiling_params);
    const bool prefer_temporal_tiling = decode_video && latents.dim() == 5 && latents.shape()[2] > 1 &&
                                        first_stage_model->can_temporal_tile_decode();
    for (;;) {
        int tile_size_w = static_cast<int>(latents.shape()[0]);
        int tile_size_h = static_cast<int>(latents.shape()[1]);
        float tile_overlap;
        if (tiling_params.enabled &&
            !first_stage_model->get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params,
                                               latents.shape()[0], latents.shape()[1])) {
            return {};
        }
        auto decoded = first_stage_model->decode(n_threads, latents, tiling_params, decode_video, circular_x, circular_y);
        if (version == VERSION_P_H3_IMAGE && !decoded.empty() && decoded.dim() == 5) {
            decoded = decoded.squeeze(2);
        }
        if (!decoded.empty() ||
            !sd::backend_fit::prepare_vae_decode_retry_tiling(tiling_params, prefer_temporal_tiling,
                                                              first_stage_model->last_compute_status(),
                                                              tile_size_w, tile_size_h, first_stage_model->get_scale_factor())) {
            return decoded;
        }
    }
}

sd::Tensor<float> StableDiffusionGGML::normalize_ltx_video_latents(const sd::Tensor<float>& x) {
    auto ltx_vae = std::dynamic_pointer_cast<LTXVideoVAE>(first_stage_model);
    if (!ltx_vae) {
        LOG_ERROR("LTX latent normalization requires LTX video VAE");
        return {};
    }
    return ltx_vae->normalize_latents(n_threads, x);
}

sd::Tensor<float> StableDiffusionGGML::un_normalize_ltx_video_latents(const sd::Tensor<float>& x) {
    auto ltx_vae = std::dynamic_pointer_cast<LTXVideoVAE>(first_stage_model);
    if (!ltx_vae) {
        LOG_ERROR("LTX latent un-normalization requires LTX video VAE");
        return {};
    }
    return ltx_vae->un_normalize_latents(n_threads, x);
}

sd::Tensor<float> StableDiffusionGGML::decode_ltx_audio_latent(const sd::Tensor<float>& audio_latent) {
    if (audio_vae_model == nullptr || audio_latent.empty()) {
        return {};
    }
    auto waveform = audio_vae_model->decode(n_threads, audio_latent);
    return waveform;
}

void StableDiffusionGGML::set_flow_shift(float flow_shift) {
    auto flow_denoiser = std::dynamic_pointer_cast<DiscreteFlowDenoiser>(denoiser);
    if (flow_denoiser) {
        if (flow_shift == INFINITY) {
            flow_shift = default_flow_shift;
        }
        flow_denoiser->set_shift(flow_shift);
        active_flow_shift = flow_shift;
    }
}

bool StableDiffusionGGML::is_flow_denoiser() {
    auto flow_denoiser = std::dynamic_pointer_cast<DiscreteFlowDenoiser>(denoiser);
    return !!flow_denoiser;
}

std::string StableDiffusionGGML::get_default_ref_image_preset(SDVersion version) const {
    return find_model_definition(version)->ref_image_preset;
}

RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_image_args) const {
    RefImageParams params;
    std::string preset_name = get_default_ref_image_preset(version);
    // Keys whose value was rejected. They are already reported, so per-model "no effect"
    // warnings must not repeat them.
    std::set<std::string> rejected_keys;

    for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) {
        if (key == "preset") {
            std::string requested_preset_name = value;
            if (REF_IMAGE_PRESETS.count(requested_preset_name)) {
                preset_name = requested_preset_name;
            } else if (value != "default") {
                std::string valid_list;
                for (auto const& [name, _] : REF_IMAGE_PRESETS) {
                    valid_list += (valid_list.empty() ? "" : ", ") + name;
                }
                LOG_WARN("ignoring invalid reference image preset '%s'. Valid options: [%s]", value.c_str(), valid_list.c_str());
                rejected_keys.insert(key);
            }
            break;
        }
    }
    if (preset_name != "default") {
        LOG_INFO("Using '%s' preset for reference images", preset_name.c_str());
        params = REF_IMAGE_PRESETS.at(preset_name);
    }
    if (version == VERSION_QWEN_IMAGE_2_1) {
        params.vlm_use_vae_input = true;
    }

    for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) {
        if (key == "pass_to_vlm") {
            if (!parse_strict_bool(value, params.pass_to_vlm)) {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
                rejected_keys.insert(key);
            }
        } else if (key == "pass_to_dit") {
            if (!parse_strict_bool(value, params.pass_to_dit)) {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
                rejected_keys.insert(key);
            }
        } else if (key == "ref_index_mode") {
            if (value == "fixed") {
                params.ref_index_mode = Rope::RefIndexMode::FIXED;
            } else if (value == "increase") {
                params.ref_index_mode = Rope::RefIndexMode::INCREASE;
            } else if (value == "decrease") {
                params.ref_index_mode = Rope::RefIndexMode::DECREASE;
            } else {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
                rejected_keys.insert(key);
            }
        } else if (key == "force_ref_timestep_zero") {
            if (!parse_strict_bool(value, params.force_ref_timestep_zero)) {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
                rejected_keys.insert(key);
            }
        } else if (key == "resize_before_vae") {
            if (!parse_strict_bool(value, params.resize_before_vae)) {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
                rejected_keys.insert(key);
            }
        } else if (key == "resize_vae_to_target") {
            if (!parse_strict_bool(value, params.resize_vae_to_target)) {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
            }
        } else if (key == "resize_single_ref_to_target") {
            if (!parse_strict_bool(value, params.resize_single_ref_to_target)) {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
            }
        } else if (key == "vae_size_buckets") {
            if (!parse_strict_bool(value, params.vae_size_buckets)) {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
            }
        } else if (key == "vae_center_crop") {
            if (!parse_strict_bool(value, params.vae_center_crop)) {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
            }
        } else if (key == "vlm_use_vae_input") {
            if (!parse_strict_bool(value, params.vlm_use_vae_input)) {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
            }
        } else if (key == "vae_input_max_pixels") {
            if (!parse_strict_int(value, params.vae_input_max_pixels)) {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
                rejected_keys.insert(key);
            }
        } else if (key == "vlm_resize_mode") {
            if (value == "longest_side") {
                params.vlm_resize_mode = RefImageResizeMode::LONGEST_SIDE;
            } else if (value == "area") {
                params.vlm_resize_mode = RefImageResizeMode::AREA;
            } else if (value == "none") {
                params.vlm_resize_mode = RefImageResizeMode::NONE;
            } else {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
                rejected_keys.insert(key);
            }
        } else if (key == "vlm_max_size") {
            if (!parse_strict_int(value, params.vlm_max_size)) {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
                rejected_keys.insert(key);
            }
        } else if (key == "vlm_min_size") {
            if (!parse_strict_int(value, params.vlm_min_size)) {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
                rejected_keys.insert(key);
            }
        } else if (key == "ref_image_size") {
            bool valid = true;
            if (value == "match") {
                params.ref_image_size = RefImageSizeMode::MATCH;
            } else if (value == "max") {
                params.ref_image_size = RefImageSizeMode::MAX;
            } else {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
                rejected_keys.insert(key);
                valid = false;
            }
            if (valid && !sd_version_is_minimax_h3(version)) {
                LOG_WARN("ref_image_size is only supported by MiniMax-H3 for now");
            }
        } else if (key != "preset" && key != "vlm_size") {
            LOG_WARN("ignoring unknown reference image arg '%s'", key.c_str());
        }
    }
    for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) {
        if (key == "vlm_size") {
            int vlm_size;
            if (!parse_strict_int(value, vlm_size)) {
                LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
                rejected_keys.insert(key);
            } else {
                LOG_INFO("vlm_size override: setting both min and max size to %ld", (long)vlm_size);
                params.vlm_min_size = vlm_size;
                params.vlm_max_size = vlm_size;
            }
            break;
        }
    }
    if (params.force_ref_timestep_zero && !sd_version_is_krea2(version)) {
        LOG_WARN("force_ref_timestep_zero is only supported by Krea2 architecture for now");
    }
    if (version == VERSION_MINIMAX_H3) {
        // H3 consumes references as VAE latents in its own packed sequence rather than
        // through the generic conditioner/DiT reference path, so ref_image_size is the
        // only arg it acts on. Everything else is reported as ineffective rather than
        // enumerated, so newly added RefImageParams members cannot silently do nothing.
        static const std::set<std::string> h3_effective_keys = {"ref_image_size"};
        // Keys that already have their own, more specific warning above.
        static const std::set<std::string> h3_separately_reported_keys = {
            "resize_before_vae",
            "force_ref_timestep_zero",
        };
        for (const auto& arg : parse_key_value_args(ref_image_args, "reference image args")) {
            if (h3_effective_keys.count(arg.first) == 0 &&
                h3_separately_reported_keys.count(arg.first) == 0 &&
                rejected_keys.count(arg.first) == 0) {
                LOG_WARN("reference image arg '%s' has no effect for MiniMax-H3", arg.first.c_str());
            }
        }

        if (!params.resize_before_vae) {
            LOG_WARN("MiniMax-H3 requires 32-pixel-aligned reference images and cannot skip resizing before VAE encoding; ignoring resize_before_vae=false");
        }
    }
    return params;
}

void StableDiffusionGGML::apply_circular_axes(bool circular_x, bool circular_y) {
    this->circular_x = circular_x;
    this->circular_y = circular_y;
    if (this->diffusion_model) {
        this->diffusion_model->set_circular_axes(circular_x, circular_y);
    }
    if (this->high_noise_diffusion_model) {
        this->high_noise_diffusion_model->set_circular_axes(circular_x, circular_y);
    }
    if (this->control_net) {
        this->control_net->set_circular_axes(circular_x, circular_y);
    }
    if (circular_x || circular_y) {
        LOG_INFO("Using circular padding for convolutions (x=%s, y=%s)",
                 circular_x ? "true" : "false",
                 circular_y ? "true" : "false");
    }
}