diff --git a/CMakeLists.txt b/CMakeLists.txt index e64ccb717..04036aded 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -229,6 +229,8 @@ file(GLOB SD_LIB_SOURCES CONFIGURE_DEPENDS "src/model/*/*.h" "src/model/*/*.cpp" "src/model/*/*.hpp" + "src/pipeline/*.h" + "src/pipeline/*.cpp" "src/runtime/*.h" "src/runtime/*.cpp" "src/runtime/*.hpp" diff --git a/scripts/format-code.ps1 b/scripts/format-code.ps1 index c18acfef2..195c5528d 100644 --- a/scripts/format-code.ps1 +++ b/scripts/format-code.ps1 @@ -11,6 +11,8 @@ $patterns = @( "src/extensions/*.cpp" "src/extensions/*.h" "src/extensions/*.hpp" + "src/pipeline/*.cpp" + "src/pipeline/*.h" "src/runtime/*.cpp" "src/runtime/*.h" "src/runtime/*.hpp" diff --git a/scripts/format-code.sh b/scripts/format-code.sh index 733b9e3fc..f5e1d8875 100644 --- a/scripts/format-code.sh +++ b/scripts/format-code.sh @@ -9,6 +9,7 @@ for f in src/*.cpp src/*.h src/*.hpp \ src/conditioning/*.cpp src/conditioning/*.h src/conditioning/*.hpp \ src/core/*.cpp src/core/*.h src/core/*.hpp \ src/extensions/*.cpp src/extensions/*.h src/extensions/*.hpp \ + src/pipeline/*.cpp src/pipeline/*.h \ src/runtime/*.cpp src/runtime/*.h src/runtime/*.hpp \ src/model/*/*.cpp src/model/*/*.h src/model/*/*.hpp \ src/tokenizers/*.h src/tokenizers/*.cpp src/tokenizers/vocab/*.h src/tokenizers/vocab/*.cpp \ diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp new file mode 100644 index 000000000..a8797874a --- /dev/null +++ b/src/pipeline/diffusion_engine.cpp @@ -0,0 +1,2812 @@ +#include "diffusion_engine.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#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 "model.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 "core/backend_fit.h" +#include "extensions/generation_extension.h" +#include "model/adapter/ip_adapter.hpp" +#include "model/adapter/lora.hpp" +#include "model/diffusion/animatediff.hpp" +#include "model/diffusion/control.hpp" +#include "model/diffusion/model.hpp" +#include "model/vae/audio_vae.hpp" +#include "model/vae/ltx_vae.hpp" +#include "model/vae/vae.hpp" +#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 + +const char* model_version_to_str[] = { + "SD 1.x", + "SD 1.x Inpaint", + "Instruct-Pix2Pix", + "SD 1.x Tiny UNet", + "SD 2.x", + "SD 2.x Inpaint", + "SD 2.x Tiny UNet", + "SDXS (512-DS)", + "SDXS (09)", + "SDXL", + "SDXL Inpaint", + "SDXL Instruct-Pix2Pix", + "SDXL (Vega)", + "SDXL (SSD1B)", + "SVD", + "SD3.x", + "Flux", + "Flux Fill", + "Flux Control", + "Flex.2", + "Chroma Radiance", + "Wan 2.x", + "Wan 2.2 I2V", + "Wan 2.2 TI2V", + "LingBot Video", + "Qwen Image", + "Qwen Image Layered", + "Hunyuan Video", + "Anima", + "Flux.2", + "Flux.2 klein", + "LTXAV", + "MiniMax-H3", + "HiDream O1", + "Z-Image", + "Boogu Image", + "Ovis Image", + "Ernie Image", + "Lens", + "MiniT2I", + "Longcat-Image", + "PiD", + "Ideogram 4", + "SeFi-Image", + "Krea2", + "Mage Flow", + "ESRGAN", +}; + +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; + } +} + +template +struct has_set_runtime_backends : std::false_type {}; +template +struct has_set_runtime_backends().set_runtime_backends( + std::declval&>()))>> : std::true_type {}; + +static_assert(std::atomic::is_always_lock_free, + "sd_cancel_mode_t must be lock-free"); + +StableDiffusionGGML::StableDiffusionGGML() + : rng(std::make_shared()), + denoiser(std::make_shared()) {} + +StableDiffusionGGML::~StableDiffusionGGML() = default; + +const std::map>& StableDiffusionGGML::runner_components() { + static const std::map> components{ + {RunnerGroup::Core, {ModelComponent::Conditioner, ModelComponent::Diffusion, ModelComponent::HighNoiseDiffusion, ModelComponent::CLIPVision, ModelComponent::IPAdapter}}, + {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) { + 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 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); + } + } + } + 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 StableDiffusionGGML::layer_split_vram_limits_for_backends(const std::vector& backends) { + std::vector 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 +bool StableDiffusionGGML::register_runner_params(ModelComponent component, + const std::shared_ptr& model, + SDBackendModule module, + size_t* params_mem_size) { + if (model == nullptr) { + return true; + } + std::map group_tensors; + std::map tensor_ops; + model->get_param_tensors(group_tensors); + if constexpr (std::is_base_of_v) { + 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 module_backends = backend_manager.runtime_backends(module); + if (module_backends.size() > 1) { + if constexpr (has_set_runtime_backends::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 +bool StableDiffusionGGML::register_row_split_runner_params(ModelComponent component, + const std::shared_ptr& model, + SDBackendModule module, + const std::vector& module_backends, + std::map group_tensors, + const std::map& 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 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)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> 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 split_tensors; + if constexpr (std::is_base_of_v) { + model->get_layer_split_param_tensors(split_tensors); + } else { + split_tensors = group_tensors; + } + + std::map row_split_map; + std::map 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 +bool StableDiffusionGGML::register_layer_split_runner_params(ModelComponent component, + const std::shared_ptr& model, + SDBackendModule module, + const std::vector& module_backends, + std::map group_tensors, + const std::map& 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 =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(*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(*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 StableDiffusionGGML::get_rng(rng_type_t rng_type) { + if (rng_type == STD_DEFAULT_RNG) { + return std::make_shared(); + } else if (rng_type == CPU_RNG) { + return std::make_shared(); + } else { // default: CUDA_RNG + return std::make_shared(); + } +} + +void StableDiffusionGGML::refresh_compvis_denoiser_sigmas() { + auto comp_vis_denoiser = std::dynamic_pointer_cast(denoiser); + if (!comp_vis_denoiser) { + return; + } + std::vector alphas_cumprod(TIMESTEPS); + if (file_alphas_cumprod.size() == TIMESTEPS) { + alphas_cumprod = file_alphas_cumprod; + } 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 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; + if (strlen(SAFE_STR(sd_ctx_params->model_path)) > 0) { + LOG_INFO("loading model from '%s'", sd_ctx_params->model_path); + if (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(sd_ctx_params->embeddings_connectors_path)) { + 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 (!model_loader.init_from_file(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->motion_module_path)) > 0) { + LOG_INFO("loading motion module (AnimateDiff) from '%s'", sd_ctx_params->motion_module_path); + if (!model_loader.init_from_file(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 (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0) { + if (!model_loader.init_from_file(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) { + 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::init(const sd_ctx_params_t* sd_ctx_params) { + auto configuration = std::make_unique(*sd_ctx_params); + n_threads = sd_ctx_params->n_threads; + 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 && backend_spec.empty() && 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(); + 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 ", model_version_to_str[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; + for (const auto& [_, tensor_storage] : model_loader.get_tensor_storage_map()) { + if (tensor_storage.is_int8_tensorwise) { + have_int8_tensorwise = true; + 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_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( + "INT8 tensorwise 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); + + 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)); + } + 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); +} + +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::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::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(); + if (sd_ctx_params->flash_attn) { + LOG_INFO("Using flash attention"); + cond_stage_model->set_flash_attention_enabled(true); + if (clip_vision) { + clip_vision->set_flash_attention_enabled(true); + } + if (first_stage_model) { + first_stage_model->set_flash_attention_enabled(true); + } + if (preview_vae) { + preview_vae->set_flash_attention_enabled(true); + } + } + + if (sd_ctx_params->flash_attn || sd_ctx_params->diffusion_flash_attn) { + LOG_INFO("Using flash attention in the diffusion model"); + diffusion_model->set_flash_attention_enabled(true); + if (high_noise_diffusion_model) { + high_noise_diffusion_model->set_flash_attention_enabled(true); + } + } + LOG_VERBOSE("validating model metadata"); + + std::set ignore_tensors; + 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_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(); + default_flow_shift = INFINITY; + prediction_t pred_type = sd_ctx_params->prediction; + + if (pred_type == PREDICTION_COUNT) { + if (sd_version_is_sd2(version)) { + pred_type = is_using_v_parameterization_for_sd2(sd_version_is_inpaint(version)) ? V_PRED : EPS_PRED; + } else if (sd_version_is_sdxl(version)) { + if (tensor_storage_map.find("edm_vpred.sigma_max") != tensor_storage_map.end()) { + // CosXL models + // TODO: get sigma_min and sigma_max values from file + pred_type = EDM_V_PRED; + } else if (tensor_storage_map.find("v_pred") != tensor_storage_map.end()) { + pred_type = V_PRED; + } else { + pred_type = EPS_PRED; + } + } else if (sd_version_is_sd3(version) || + sd_version_is_wan(version) || + sd_version_is_hunyuan_video(version) || + sd_version_is_lingbot_video(version) || + sd_version_is_minimax_h3(version) || + sd_version_is_qwen_image(version) || + sd_version_is_mage_flow(version) || + version == VERSION_HIDREAM_O1 || + sd_version_is_anima(version) || + sd_version_is_ernie_image(version) || + sd_version_is_z_image(version) || + sd_version_is_boogu_image(version) || + sd_version_is_pid(version) || + sd_version_is_ideogram4(version)) { + pred_type = FLOW_PRED; + if (sd_version_is_wan(version)) { + default_flow_shift = 5.f; + } else if (sd_version_is_hunyuan_video(version)) { + default_flow_shift = 7.f; + } else if (sd_version_is_minimax_h3(version)) { + default_flow_shift = 12.f; + } else if (sd_version_is_ernie_image(version)) { + default_flow_shift = 4.f; + } else if (sd_version_is_pid(version)) { + default_flow_shift = 1.5f; + } else if (sd_version_is_ideogram4(version)) { + default_flow_shift = 1.0f; + } else if (sd_version_is_boogu_image(version)) { + default_flow_shift = 3.16f; + } else if (sd_version_is_mage_flow(version)) { + default_flow_shift = 6.f; + } else { + default_flow_shift = 3.f; + } + } else if (sd_version_is_flux(version) || + sd_version_is_flux2(version) || + sd_version_is_longcat(version) || + sd_version_is_lens(version) || + sd_version_is_ltxav(version) || + sd_version_is_krea2(version)) { + pred_type = FLUX_FLOW_PRED; + + default_flow_shift = 1.0f; // TODO: validate + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (starts_with(name, "model.diffusion_model.guidance_in.in_layer.weight")) { + default_flow_shift = 1.15f; + break; + } + } + if (sd_version_is_longcat(version)) { + default_flow_shift = 3.0f; + } else if (sd_version_is_lens(version)) { + default_flow_shift = 1.83f; + } else if (sd_version_is_ltxav(version)) { + default_flow_shift = 2.37f; + } else if (sd_version_is_krea2(version)) { + default_flow_shift = 1.15f; + } + } else if (sd_version_is_sefi_image(version)) { + pred_type = SEFI_FLOW_PRED; + } else if (sd_version_is_minit2i(version)) { + pred_type = MINIT2I_FLOW_PRED; + } else { + pred_type = EPS_PRED; + } + } + + 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(); + break; + case EDM_V_PRED: + LOG_INFO("running in v-prediction EDM mode"); + denoiser = std::make_shared(); + break; + case FLOW_PRED: { + if (sd_version_is_ltxav(version)) { + LOG_INFO("running in LTXAV FLOW mode"); + denoiser = std::make_shared(); + } else if (sd_version_is_minimax_h3(version)) { + LOG_INFO("running in MiniMax H3 AV FLOW mode"); + denoiser = std::make_shared(default_flow_shift, 3.f, get_latent_channel()); + } else { + LOG_INFO("running in FLOW mode"); + denoiser = std::make_shared(); + } + break; + } + case FLUX_FLOW_PRED: { + LOG_INFO("running in Flux FLOW mode"); + denoiser = std::make_shared(); + break; + } + case SEFI_FLOW_PRED: { + LOG_INFO("running in SeFi-Image dual-time FLOW mode"); + denoiser = std::make_shared(); + break; + } + case MINIT2I_FLOW_PRED: { + LOG_INFO("running in MiniT2I FLOW mode"); + denoiser = std::make_shared(); + 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 wtype_stat = model_loader.get_wtype_stat(); + std::map conditioner_wtype_stat = model_loader.get_conditioner_wtype_stat(); + std::map diffusion_model_wtype_stat = model_loader.get_diffusion_model_wtype_stat(); + std::map vae_wtype_stat = model_loader.get_vae_wtype_stat(); + + auto wtype_stat_to_str = [](const std::map& 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_VERBOSE("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 x_t = sd::full({8, 8, 4, 1}, 0.5f); + sd::Tensor c = sd::full({1024, 2, 1, 1}, 0.5f); + sd::Tensor steps = sd::full({1}, 999.0f); + sd::Tensor concat; + if (is_inpaint) { + concat = sd::zeros({8, 8, 5, 1}); + } + + int64_t t0 = ggml_time_ms(); + sd::Tensor 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((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 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(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> StableDiffusionGGML::load_runtime_loras_for_module(const std::vector& loras, + const std::set& model_tensor_names, + SDBackendModule module, + LoraModel::filter_t module_filter, + bool& success, + std::vector& next_models) { + std::vector> 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& 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& 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 next_models; + std::set 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(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(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(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 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)) + return false; + 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) { + 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); + } + return true; +} + +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 StableDiffusionGGML::get_clip_vision_output(const sd::Tensor& image, + bool return_pooled, + int clip_skip, + bool zero_out_masked) { + sd::Tensor output; + if (zero_out_masked) { + if (return_pooled) { + output = sd::zeros({clip_vision->vision_model.projection_dim}); + } else { + output = sd::zeros({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; +} + +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::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 StableDiffusionGGML::process_timesteps(const std::vector& timesteps, + const sd::Tensor& init_latent, + const sd::Tensor& denoise_mask, + int step) { + if (auto sefi_denoiser = std::dynamic_pointer_cast(denoiser)) { + int sched_idx = step > 0 ? step - 1 : 0; + if (sched_idx >= static_cast(sefi_denoiser->tex_timesteps.size())) { + sched_idx = static_cast(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(static_cast(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(frame)] = 0.f; + } + } + } + return new_timesteps; + } else { + return timesteps; + } +} + +std::vector StableDiffusionGGML::process_ltxav_video_timesteps(const std::vector& timesteps, + const sd::Tensor& init_latent, + const sd::Tensor& 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 video_timesteps(static_cast(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& latents, + enum SDVersion version, + preview_t preview_mode, + std::function 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(latents.shape()[3]) : static_cast(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; + + 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 == 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(_latents.shape()[2]) : 1; + uint32_t img_width = static_cast(_latents.shape()[0]) * patch_sz; + uint32_t img_height = static_cast(_latents.shape()[1]) * patch_sz; + + uint8_t* data = (uint8_t*)malloc(frames * img_width * img_height * 3 * sizeof(uint8_t)); + GGML_ASSERT(data != nullptr); + preview_latent_video(data, _latents, latent_rgb_proj, latent_rgb_bias, patch_sz); + 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, 3, data + i * img_width * img_height * 3}; + } + 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 vae_latents; + sd::Tensor 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(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(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(preview_mode)); + } +} + +std::vector 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(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)shifted_t}; + } + if (sd_version_is_anima(version)) { + return std::vector{t / static_cast(TIMESTEPS)}; + } + if (sd_version_is_boogu_image(version)) { + return std::vector{t / static_cast(TIMESTEPS)}; + } + if (version == VERSION_HIDREAM_O1) { + return std::vector{1.0f - (t / static_cast(TIMESTEPS))}; + } + if (sd_version_is_z_image(version) || sd_version_is_ideogram4(version)) { + return std::vector{1000.f - t}; + } + return std::vector{t}; +} + +void StableDiffusionGGML::adjust_sample_step_scalings(int shifted_timestep, + const std::vector& 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(roundf(timesteps_vec[0])); + float shifted_sigma = denoiser->t_to_sigma((float)shifted_t_idx); + std::vector 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; + } + } +} + +void StableDiffusionGGML::compute_sample_controls(const sd::Tensor& control_image, + const sd::Tensor& noised_input, + const sd::Tensor& timesteps_tensor, + const SDCondition& condition, + std::vector>* 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 StableDiffusionGGML::sample(const std::shared_ptr& work_diffusion_model, + bool inverse_noise_scaling, + const sd::Tensor& init_latent, + sd::Tensor noise, + const SDCondition& cond, + const SDCondition& uncond, + const SDCondition& img_uncond, + const sd::Tensor& 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& sigmas, + const std::vector>& ref_latents, + const RefImageParams& ref_image_params, + const sd::Tensor& denoise_mask, + const sd::Tensor& 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& video_positions) { + struct RunnerEndOnExit { + GGMLRunner* runner = nullptr; + ~RunnerEndOnExit() { + if (runner != nullptr) { + runner->runner_end(); + } + } + }; + RunnerEndOnExit sample_diffusion_runner_end{work_diffusion_model.get()}; + + RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; + + std::vector 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 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(adaptive_projected_guidance) + : static_cast(classifier_free_guidance); + sd::guidance::SkipLayerGuidance skip_layer_guidance(has_skiplayer ? skip_layers : std::vector(), + 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(); + SamplePreviewContext preview = prepare_sample_preview_context(); + + sd::Tensor processed_init_latent = denoiser->process_latent_in(init_latent); + const sd::Tensor& sampling_init_latent = processed_init_latent.empty() + ? init_latent + : processed_init_latent; + sd::Tensor x_t = !noise.empty() + ? denoiser->noise_scaling(sigmas[0], noise, sampling_init_latent) + : sampling_init_latent; + sd::Tensor denoised = x_t; + + auto denoise = [&](const sd::Tensor& 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 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 base_timesteps_vec = prepare_sample_timesteps(sigma, shifted_timestep); + std::vector timesteps_vec = base_timesteps_vec; + sd::Tensor 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({static_cast(base_timesteps_vec.size())}, base_timesteps_vec); + } else { + timesteps_vec = process_timesteps(timesteps_vec, sampling_init_latent, denoise_mask, step); + } + const std::vector& 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 timesteps_tensor({static_cast(timesteps_vec.size())}, timesteps_vec); + sd::Tensor guidance_tensor({1}, std::vector{guidance.distilled_guidance}); + sd::Tensor hunyuan_timestep_r_tensor; + if (sd_version_is_hunyuan_video(version) && step + 1 < sigmas.size()) { + hunyuan_timestep_r_tensor = sd::Tensor::from_vector({sigmas[step + 1]}); + } + sd::Tensor noised_input = x * c_in; + if (!denoise_mask.empty() && (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 (!denoise_mask.empty()) { + 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 cond_out; + sd::Tensor uncond_out; + sd::Tensor img_uncond_out; + sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma); + std::vector> controls; + DiffusionParams diffusion_params; + diffusion_params.x = &noised_input; + diffusion_params.timesteps = ×teps_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); + + static const std::vector> 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* c_concat_override = nullptr, + const std::vector* local_skip_layers = nullptr, + const std::vector>* ref_latents_override = nullptr, + bool use_uncond_ip = false) -> sd::Tensor { + 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(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)) { + 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_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}; + } 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 (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, + audio_length, + std::isfinite(active_flow_shift) ? active_flow_shift : 12.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 { + diffusion_params.extra = std::monostate{}; + } + + sd::Tensor cached_output; + if (step_cache.before_condition(&condition, noised_input, &cached_output)) { + return std::move(cached_output); + } + + for (const auto& extension : generation_extensions) { + extension->before_diffusion(diffusion_params, step); + } + + auto output_opt = work_diffusion_model->compute(n_threads, diffusion_params); + if (output_opt.empty()) { + LOG_ERROR("diffusion model compute failed"); + return sd::Tensor(); + } + + step_cache.after_condition(&condition, noised_input, output_opt); + return output_opt; + }; + + const SDCondition* positive_condition = &cond; + const sd::Tensor* 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; + } + } + + cond_out = run_condition(*positive_condition, c_concat_override); + if (cond_out.empty()) { + return {}; + } + + if (!uncond.empty()) { + if (!step_cache.is_step_skipped()) { + compute_sample_controls(control_image, + noised_input, + timesteps_tensor, + uncond, + &controls); + } + const std::vector* uncond_skip_layers = nullptr; + if (is_skiplayer_step && slg_uncond) { + LOG_VERBOSE("Skipping layers at uncond step %d\n", step); + uncond_skip_layers = &skip_layer_guidance.layers(); + } + uncond_out = run_condition(uncond, + uncond.c_concat.empty() ? nullptr : &uncond.c_concat, + uncond_skip_layers, + nullptr, + true); + if (uncond_out.empty()) { + return {}; + } + } + if (!img_uncond.empty()) { + 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 {}; + } + } + 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 = guidance_schedule.empty() ? primary_guidance.forward(guidance_input, {}) : primary_guidance.forward(guidance_input, {}, guidance_schedule[guidance_schedule.size() - 1 - step]); + 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 { + 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& 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 (!denoise_mask.empty()) { + 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; + }; + + 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() { + int down_factor = 8; // unet + if (sd_version_is_dit(version)) { + if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version)) { + down_factor = 2; + } else { + down_factor = 1; + } + } + return down_factor; +} + +int StableDiffusionGGML::get_latent_channel() { + int latent_channel = 4; + if (sd_version_is_dit(version)) { + if (sd_version_is_ltxav(version)) { + latent_channel = 128; + } else if (sd_version_is_minimax_h3(version)) { + latent_channel = 24; + } else if (version == VERSION_WAN2_2_TI2V) { + latent_channel = 48; + } else if (sd_version_is_hunyuan_video(version)) { + latent_channel = 32; + } else if (version == VERSION_HIDREAM_O1) { + latent_channel = 3; + } else if (version == VERSION_CHROMA_RADIANCE) { + latent_channel = 3; + } else if (sd_version_is_minit2i(version)) { + latent_channel = 3; + } else if (sd_version_is_pid(version)) { + latent_channel = 3; + } else if (sd_version_is_sefi_image(version)) { + latent_channel = 144; + } else if (sd_version_uses_flux2_vae(version)) { + latent_channel = 128; + } else if (sd_version_is_mage_flow(version)) { + latent_channel = 128; + } else { + latent_channel = 16; + } + } + return latent_channel; +} + +int StableDiffusionGGML::get_image_channels() const { + return version == VERSION_QWEN_IMAGE_LAYERED ? 4 : 3; +} + +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 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({W, H, T, C, 1}); + } + return sd::zeros({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 StableDiffusionGGML::encode_to_vae_latents(const sd::Tensor& 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 StableDiffusionGGML::encode_first_stage(const sd::Tensor& 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 StableDiffusionGGML::decode_first_stage(const sd::Tensor& x, bool decode_video) { + if (sd_version_is_pid(version) || sd_version_is_minit2i(version)) { + return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f); + } + auto latents = first_stage_model->diffusion_to_vae_latents(x); + auto decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); + const bool prefer_temporal_tiling = decode_video && first_stage_model->can_temporal_tile_decode(); + while (decoded.empty() && + auto_fit_enabled && + sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { + decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); + } + return decoded; +} + +sd::Tensor StableDiffusionGGML::normalize_ltx_video_latents(const sd::Tensor& x) { + auto ltx_vae = std::dynamic_pointer_cast(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 StableDiffusionGGML::un_normalize_ltx_video_latents(const sd::Tensor& x) { + auto ltx_vae = std::dynamic_pointer_cast(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 StableDiffusionGGML::decode_ltx_audio_latent(const sd::Tensor& 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(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(denoiser); + return !!flow_denoiser; +} + +std::string StableDiffusionGGML::get_default_ref_image_preset(SDVersion version) const { + if (sd_version_is_longcat(version)) { + return "longcat"; + } else if (sd_version_is_flux(version)) { + return "flux_kontext"; + } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) { + return "flux2"; + } else if (version == VERSION_QWEN_IMAGE_LAYERED) { + return "qwen_layered"; + } else if (sd_version_is_qwen_image(version)) { + return "qwen"; + } else if (sd_version_is_mage_flow(version)) { + return "mage_flow"; + } else if (sd_version_is_z_image(version) || sd_version_is_boogu_image(version)) { + return "z_image_omni"; + } else if (sd_version_is_krea2(version)) { + // have to make a choice between "krea2_edit" mode (for lbouaraba/krea2edit) + // and "krea2_ostris_edit" (for krea2 ostris edit) + // since krea2 ostris edit support predates, it should probably be default + return "krea2_ostris_edit"; + } else if (sd_version_is_anima(version)) { + return "cosmos_reference"; + } + return "default"; +} + +RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_image_args) const { + RefImageParams params; + std::string preset_name = get_default_ref_image_preset(version); + + 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()); + } + break; + } + } + if (preset_name != "default") { + LOG_INFO("Using '%s' preset for reference images", preset_name.c_str()); + params = REF_IMAGE_PRESETS.at(preset_name); + } + + 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()); + } + } 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()); + } + } 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()); + } + } 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()); + } + } 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()); + } + } 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()); + } + } 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()); + } + } 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()); + } + } 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()); + } + } 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()); + } 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"); + } + 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"); + } +} diff --git a/src/pipeline/diffusion_engine.h b/src/pipeline/diffusion_engine.h new file mode 100644 index 000000000..13cb6e8f2 --- /dev/null +++ b/src/pipeline/diffusion_engine.h @@ -0,0 +1,482 @@ +#ifndef __SD_PIPELINE_DIFFUSION_ENGINE_H__ +#define __SD_PIPELINE_DIFFUSION_ENGINE_H__ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/ggml_extend_backend.h" +#include "core/ggml_graph_cut.h" +#include "core/tensor.hpp" +#include "core/util.h" +#include "model/adapter/lora.hpp" +#include "model_builders.h" +#include "model_manager.h" +#include "stable-diffusion.h" + +class RNG; +struct Denoiser; +struct LoraModel; +struct ConditionerParams; +struct SDCondition; +struct RefImageParams; + +extern const char* model_version_to_str[]; + +static inline bool sd_version_supports_ref_latent_img_cfg(SDVersion version) { + return version == VERSION_FLUX || + sd_version_is_flux2(version) || + sd_version_is_qwen_image(version) || + sd_version_is_mage_flow(version) || + sd_version_is_longcat(version) || + sd_version_is_z_image(version) || + sd_version_is_boogu_image(version); +} + +class StableDiffusionGGML { +public: + SDBackendManager backend_manager; + + SDVersion version; + bool external_vae_is_invalid = false; + + bool circular_x = false; + bool circular_y = false; + + std::shared_ptr rng; + std::shared_ptr sampler_rng = nullptr; + int n_threads = -1; + float default_flow_shift = INFINITY; + float active_flow_shift = INFINITY; + + std::shared_ptr cond_stage_model; + std::shared_ptr clip_vision; // for svd or wan2.1 i2v + std::shared_ptr diffusion_model; + std::shared_ptr high_noise_diffusion_model; + std::shared_ptr first_stage_model; + std::shared_ptr preview_vae; + std::shared_ptr audio_vae_model; + std::shared_ptr control_net; + std::shared_ptr ip_adapter; + sd::Tensor ip_adapter_tokens; + sd::Tensor ip_adapter_uncond_tokens; + float ip_adapter_strength = 1.0f; + std::vector> generation_extensions; + struct RuntimeLora { + ModelManager::LoraSpec spec; + SDBackendModule module; + std::shared_ptr model; + + bool matches(const ModelManager::LoraSpec& other) const { + return spec.file_id == other.file_id && spec.file_revision == other.file_revision && + spec.tensor_name_prefix_filter == other.tensor_name_prefix_filter; + } + }; + std::vector runtime_lora_models; + bool apply_lora_immediately = false; + int animatediff_num_frames = 0; + + std::string taesd_path; + sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr}; + bool enable_mmap = false; + sd::ggml_graph_cut::MaxVramAssignment max_vram_assignment; + bool disable_prefetch = false; + bool disable_segmented_compute = false; + bool eager_load = false; + std::string backend_spec; + std::string params_backend_spec; + std::string split_mode_spec; + bool auto_fit_enabled = false; + + bool diffusion_conv_direct = false; + + bool is_using_v_parameterization = false; + bool is_using_edm_v_parameterization = false; + + std::shared_ptr model_manager; + + enum class RunnerGroup { Core, + VAE, + ControlNet, + Extensions }; + using RunnerGroups = std::set; + + struct ModelConfig { + sd_ctx_params_t params{}; + std::list strings; + std::vector embeddings; + ModelLoader::FileId control_net_file = 0; + bool use_tae = false; + bool use_audio_vae = false; + bool photomaker_source_available = false; + bool animatediff_loaded = false; + + explicit ModelConfig(const sd_ctx_params_t& initial) + : params(initial) { + for (auto member : {&sd_ctx_params_t::model_path, &sd_ctx_params_t::clip_l_path, + &sd_ctx_params_t::clip_g_path, &sd_ctx_params_t::clip_vision_path, + &sd_ctx_params_t::t5xxl_path, &sd_ctx_params_t::llm_path, + &sd_ctx_params_t::llm_vision_path, &sd_ctx_params_t::diffusion_model_path, + &sd_ctx_params_t::high_noise_diffusion_model_path, &sd_ctx_params_t::uncond_diffusion_model_path, + &sd_ctx_params_t::embeddings_connectors_path, &sd_ctx_params_t::vae_path, + &sd_ctx_params_t::audio_vae_path, &sd_ctx_params_t::taesd_path, + &sd_ctx_params_t::control_net_path, &sd_ctx_params_t::ip_adapter_path, + &sd_ctx_params_t::motion_module_path, &sd_ctx_params_t::photo_maker_path, + &sd_ctx_params_t::pulid_weights_path, &sd_ctx_params_t::tensor_type_rules, + &sd_ctx_params_t::max_vram, &sd_ctx_params_t::backend, + &sd_ctx_params_t::params_backend, &sd_ctx_params_t::split_mode, + &sd_ctx_params_t::rpc_servers, &sd_ctx_params_t::model_args}) { + strings.emplace_back(SAFE_STR(initial.*member)); + params.*member = strings.back().c_str(); + } + for (uint32_t i = 0; i < initial.embedding_count; ++i) { + strings.emplace_back(SAFE_STR(initial.embeddings[i].name)); + const char* name = strings.back().c_str(); + strings.emplace_back(SAFE_STR(initial.embeddings[i].path)); + embeddings.push_back({name, strings.back().c_str()}); + } + params.embeddings = embeddings.data(); + } + + ModelConfig(const ModelConfig& other) + : ModelConfig(other.params) { + control_net_file = other.control_net_file; + use_tae = other.use_tae; + use_audio_vae = other.use_audio_vae; + photomaker_source_available = other.photomaker_source_available; + animatediff_loaded = other.animatediff_loaded; + } + ModelConfig& operator=(const ModelConfig&) = delete; + + void set_control_net(ModelLoader::FileId id, const std::string& path) { + control_net_file = id; + strings.push_back(path); + params.control_net_path = strings.back().c_str(); + } + }; + + struct RunnerState { + bool ready = false; + uint64_t catalog_revision = 0; + std::map sources; + }; + + std::recursive_mutex execution_mutex; + std::unique_ptr config_; + RunnerState runner_state_; + bool executing_ = false; + + std::shared_ptr denoiser; + std::vector file_alphas_cumprod; + + StableDiffusionGGML(); + ~StableDiffusionGGML(); + + static const std::map>& runner_components(); + + static RunnerGroups all_runner_groups(); + + ModelLoader::FileVersions runner_source_versions(RunnerGroup group, const ModelLoader& loader) const; + + void capture_runner_sources(); + + void end_runners(); + + bool reset_runners(const RunnerGroups& groups); + + bool refresh_model_sources(); + + bool apply_model_update(ModelLoader candidate, + std::unique_ptr next_config = nullptr, + RunnerGroups groups = {}); + + struct ContextOperation { + StableDiffusionGGML& sd; + std::unique_lock lock; + bool acquired = false; + + explicit ContextOperation(StableDiffusionGGML& sd) + : sd(sd), lock(sd.execution_mutex, std::try_to_lock) { + if (!lock.owns_lock() || sd.executing_) { + // The caller may be a log callback, so rejecting it must not log. + return; + } + sd.executing_ = true; + acquired = true; + } + + ~ContextOperation() { + if (acquired) { + sd.executing_ = false; + } + } + }; + + struct ExecutionScope { + ContextOperation operation; + bool ready = false; + + explicit ExecutionScope(StableDiffusionGGML& sd) + : operation(sd) { + ready = operation.acquired && sd.refresh_model_sources(); + } + + ~ExecutionScope() { + if (ready) { + operation.sd.end_runners(); + } + } + }; + + ggml_backend_t backend_for(SDBackendModule module); + + ggml_backend_t params_backend_for(SDBackendModule module); + + std::atomic cancellation_flag = SD_CANCEL_RESET; + + void set_cancel_flag(enum sd_cancel_mode_t flag); + + void reset_cancel_flag(); + + enum sd_cancel_mode_t get_cancel_flag(); + + size_t max_graph_vram_bytes_for_module(SDBackendModule module); + + std::vector layer_split_vram_limits_for_backends(const std::vector& backends); + + bool ensure_backend_pair(SDBackendModule module); + + template + bool register_runner_params(ModelComponent component, + const std::shared_ptr& model, + SDBackendModule module, + size_t* params_mem_size = nullptr); + + template + bool register_row_split_runner_params(ModelComponent component, + const std::shared_ptr& model, + SDBackendModule module, + const std::vector& module_backends, + std::map group_tensors, + const std::map& tensor_ops, + ModelManager::ResidencyMode residency_mode, + size_t* params_mem_size); + + // Register graph-cut layer-split tensors on the primary backend first. + // The first real graph assigns each param tensor to a runtime backend + // before weights are loaded or staged. + template + bool register_layer_split_runner_params(ModelComponent component, + const std::shared_ptr& model, + SDBackendModule module, + const std::vector& module_backends, + std::map group_tensors, + const std::map& tensor_ops, + ModelManager::ResidencyMode residency_mode, + size_t* params_mem_size); + + bool unload_control_net(); + + bool load_control_net_from_file(const std::string& path); + + void apply_circular_axes(bool circular_x, bool circular_y); + + bool init_backend(); + + bool row_split_active(); + + bool graph_cut_layer_split_active(); + + std::shared_ptr get_rng(rng_type_t rng_type); + + void refresh_compvis_denoiser_sigmas(); + + void load_alphas_cumprod(); + + bool init_model_loader(ModelLoader& model_loader, ModelConfig& configuration); + + bool init(const sd_ctx_params_t* sd_ctx_params); + + bool uses_tae() const; + + bool tae_preview_only() const; + + void configure_weight_loading(); + + sd::model_builders::Context model_build_context(); + + bool build_core_runners(); + + bool build_vae_runners(); + + bool build_control_net_runner(); + + bool build_extension_runners(); + + bool validate_and_load_runners(); + + bool build_denoiser(); + + bool build_runners(const RunnerGroups& groups); + + bool is_using_v_parameterization_for_sd2(bool is_inpaint = false); + + static std::string lora_log_id(const ModelManager::LoraSpec& lora); + + std::shared_ptr load_lora_model(const ModelManager::LoraSpec& lora_spec, + SDBackendModule module, + LoraModel::filter_t module_filter = nullptr); + + void clear_lora_adapters(); + + std::vector> load_runtime_loras_for_module(const std::vector& loras, + const std::set& model_tensor_names, + SDBackendModule module, + LoraModel::filter_t module_filter, + bool& success, + std::vector& next_models); + + bool apply_loras_immediately(const std::vector& loras); + + bool apply_loras_at_runtime(const std::vector& loras); + + void lora_stat(); + + bool apply_loras(const sd_lora_t* loras, uint32_t lora_count); + + void reset_generation_extensions(); + + void prepare_generation_extensions(const sd_pm_params_t& pm_params, + const sd_pulid_params_t& pulid_params, + ConditionerParams& condition_params, + int total_steps); + + sd::Tensor get_clip_vision_output(const sd::Tensor& image, + bool return_pooled = true, + int clip_skip = -1, + bool zero_out_masked = false); + + void compute_ip_adapter_tokens(const sd_image_t& image, float strength); + + std::vector process_timesteps(const std::vector& timesteps, + const sd::Tensor& init_latent, + const sd::Tensor& denoise_mask, + int step); + + std::vector process_ltxav_video_timesteps(const std::vector& timesteps, + const sd::Tensor& init_latent, + const sd::Tensor& denoise_mask); + + void preview_image(int step, + const sd::Tensor& latents, + enum SDVersion version, + preview_t preview_mode, + std::function step_callback, + void* step_callback_data, + bool is_noisy); + + std::vector prepare_sample_timesteps(float sigma, + int shifted_timestep); + + void adjust_sample_step_scalings(int shifted_timestep, + const std::vector& timesteps_vec, + float c_in, + float* c_skip, + float* c_out); + + struct SamplePreviewContext { + sd_preview_cb_t callback = nullptr; + void* data = nullptr; + preview_t mode = PREVIEW_NONE; + }; + + SamplePreviewContext prepare_sample_preview_context(); + + void report_sample_progress(int step, + size_t total_steps, + bool terminal_sigma_is_zero, + int64_t* last_progress_us); + + void compute_sample_controls(const sd::Tensor& control_image, + const sd::Tensor& noised_input, + const sd::Tensor& timesteps_tensor, + const SDCondition& condition, + std::vector>* controls); + + sd::Tensor sample(const std::shared_ptr& work_diffusion_model, + bool inverse_noise_scaling, + const sd::Tensor& init_latent, + sd::Tensor noise, + const SDCondition& cond, + const SDCondition& uncond, + const SDCondition& img_uncond, + const sd::Tensor& 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& sigmas, + const std::vector>& ref_latents, + const RefImageParams& ref_image_params, + const sd::Tensor& denoise_mask, + const sd::Tensor& 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& video_positions = {}); + + int get_vae_scale_factor(); + + int get_diffusion_model_down_factor(); + + int get_latent_channel(); + + int get_image_channels() const; + + int get_image_seq_len(int h, int w); + + sd::Tensor generate_init_latent(int width, + int height, + int frames = 1, + bool video = false); + + int video_frames_to_latent_frames(int frames); + + int latent_frames_to_video_frames(int latent_frames); + + int align_video_frames(int frames); + + sd::Tensor encode_to_vae_latents(const sd::Tensor& x); + + sd::Tensor encode_first_stage(const sd::Tensor& x); + + sd::Tensor decode_first_stage(const sd::Tensor& x, bool decode_video = false); + + sd::Tensor normalize_ltx_video_latents(const sd::Tensor& x); + + sd::Tensor un_normalize_ltx_video_latents(const sd::Tensor& x); + + sd::Tensor decode_ltx_audio_latent(const sd::Tensor& audio_latent); + + void set_flow_shift(float flow_shift = INFINITY); + + bool is_flow_denoiser(); + + std::string get_default_ref_image_preset(SDVersion version) const; + + RefImageParams resolve_ref_image_params(const char* ref_image_args) const; +}; + +#endif // __SD_PIPELINE_DIFFUSION_ENGINE_H__ diff --git a/src/pipeline/generation.h b/src/pipeline/generation.h new file mode 100644 index 000000000..1210a8cba --- /dev/null +++ b/src/pipeline/generation.h @@ -0,0 +1,73 @@ +#ifndef __SD_PIPELINE_GENERATION_H__ +#define __SD_PIPELINE_GENERATION_H__ + +#include "conditioning/conditioner.hpp" +#include "stable-diffusion.h" + +class StableDiffusionGGML; + +static inline bool sd_version_supports_animatediff(SDVersion version) { + return version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX; +} + +namespace sd::pipeline { + + struct ImageGenerationLatents { + sd::Tensor init_latent; + sd::Tensor concat_latent; + sd::Tensor img_uncond_concat_latent; + sd::Tensor audio_latent; + sd::Tensor video_positions; + sd::Tensor control_image; + std::vector> ref_images; + std::vector> ref_latents; + std::vector> reference_audio_latents; + std::vector minimax_reference_blocks; + std::vector minimax_presentation_refs; + std::vector keyframe_indices; + sd::Tensor denoise_mask; + sd::Tensor clip_vision_output; + sd::Tensor vace_context; + int64_t ref_image_num = 0; + int64_t video_conditioning_frame_count = 0; + int64_t video_target_frame_count = 0; + int audio_length = 0; + }; + + struct ImageGenerationEmbeds { + SDCondition cond; + SDCondition uncond; + SDCondition img_uncond; + }; + + struct ConditionerRunnerEndOnExit { + Conditioner* conditioner = nullptr; + ~ConditionerRunnerEndOnExit() { + if (conditioner != nullptr) { + conditioner->runner_end(); + } + } + }; + + // Callers hold ExecutionScope; AnimateDiff reuses the image path within the same scope. + bool generate_image(StableDiffusionGGML* sd, + const sd_img_gen_params_t* sd_img_gen_params, + sd_image_t** images_out, + int* num_images_out); + + bool generate_video(StableDiffusionGGML* sd, + const sd_vid_gen_params_t* sd_vid_gen_params, + sd_image_t** frames_out, + int* num_frames_out, + sd_audio_t** audio_out); + + sd::Tensor upscale_ltx_spatial_video_latent(StableDiffusionGGML* sd, + const char* model_path, + const sd::Tensor& packed_latent, + int audio_length); + + sd::Tensor ensure_image_tensor_channels(sd::Tensor image, int channels); + +} // namespace sd::pipeline + +#endif // __SD_PIPELINE_GENERATION_H__ diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp new file mode 100644 index 000000000..9591c5d60 --- /dev/null +++ b/src/pipeline/image.cpp @@ -0,0 +1,1034 @@ +#include "generation.h" + +#include +#include +#include +#include + +#include "core/rng.hpp" +#include "diffusion_engine.h" +#include "model/vae/vae.hpp" +#include "request.h" +#include "runtime/denoiser.hpp" +#include "upscaler.h" + +namespace sd::pipeline { + + struct CircularAxesState { + bool circular_x = false; + bool circular_y = false; + }; + + static CircularAxesState configure_image_vae_axes(StableDiffusionGGML* sd, + const sd_img_gen_params_t* sd_img_gen_params, + const GenerationRequest& request) { + CircularAxesState original_axes = {sd->circular_x, sd->circular_y}; + + if (!sd_img_gen_params->vae_tiling_params.enabled) { + if (sd->first_stage_model) { + sd->first_stage_model->set_circular_axes(sd->circular_x, sd->circular_y); + } + if (sd->preview_vae) { + sd->preview_vae->set_circular_axes(sd->circular_x, sd->circular_y); + } + return original_axes; + } + + int tile_size_x, tile_size_y; + float overlap; + int latent_size_x = request.width / request.vae_scale_factor; + int latent_size_y = request.height / request.vae_scale_factor; + sd->first_stage_model->get_tile_sizes(tile_size_x, + tile_size_y, + overlap, + sd_img_gen_params->vae_tiling_params, + latent_size_x, + latent_size_y); + + sd->circular_x = sd->circular_x && (tile_size_x >= latent_size_x); + sd->circular_y = sd->circular_y && (tile_size_y >= latent_size_y); + + if (sd->first_stage_model) { + sd->first_stage_model->set_circular_axes(sd->circular_x, sd->circular_y); + } + if (sd->preview_vae) { + sd->preview_vae->set_circular_axes(sd->circular_x, sd->circular_y); + } + + sd->circular_x = original_axes.circular_x && (tile_size_x < latent_size_x); + sd->circular_y = original_axes.circular_y && (tile_size_y < latent_size_y); + + return original_axes; + } + + static void restore_image_vae_axes(StableDiffusionGGML* sd, const CircularAxesState& original_axes) { + sd->circular_x = original_axes.circular_x; + sd->circular_y = original_axes.circular_y; + } + + class ImageVaeAxesGuard { + private: + StableDiffusionGGML* sd = nullptr; + CircularAxesState original_axes; + + public: + ImageVaeAxesGuard(StableDiffusionGGML* sd, + const sd_img_gen_params_t* sd_img_gen_params, + const GenerationRequest& request) + : sd(sd), + original_axes(configure_image_vae_axes(sd, sd_img_gen_params, request)) {} + + ~ImageVaeAxesGuard() { + restore_image_vae_axes(sd, original_axes); + } + + ImageVaeAxesGuard(const ImageVaeAxesGuard&) = delete; + ImageVaeAxesGuard& operator=(const ImageVaeAxesGuard&) = delete; + }; + + sd::Tensor ensure_image_tensor_channels(sd::Tensor image, int channels) { + if (image.empty()) { + return image; + } + GGML_ASSERT(image.dim() == 4); + int64_t current_channels = image.shape()[2]; + if (current_channels == channels) { + return image; + } + if (channels == 4) { + sd::Tensor alpha = sd::full({image.shape()[0], image.shape()[1], 1, image.shape()[3]}, 1.f); + if (current_channels == 3) { + return sd::ops::concat(image, alpha, 2); + } + if (current_channels == 1) { + sd::Tensor rgb = sd::ops::concat(image, image, 2); + rgb = sd::ops::concat(rgb, image, 2); + return sd::ops::concat(rgb, alpha, 2); + } + } + if (channels == 3 && current_channels >= 3) { + return sd::ops::slice(image, 2, 0, 3); + } + GGML_ABORT("cannot convert image tensor from %lld to %d channels", + (long long)current_channels, + channels); + } + + static std::optional prepare_image_generation_latents(StableDiffusionGGML* sd, + const sd_img_gen_params_t* sd_img_gen_params, + GenerationRequest* request, + SamplePlan* plan, + const RefImageParams& ref_image_params) { + int64_t prepare_start_ms = ggml_time_ms(); + + sd::Tensor init_image_tensor; + sd::Tensor control_image_tensor; + sd::Tensor mask_image_tensor; + int image_channels = sd->get_image_channels(); + + if (sd_img_gen_params->init_image.data != nullptr) { + LOG_INFO("IMG2IMG"); + + if (request->strength < 1.f) { + bool strength_as_noise_level = false; + bool force_first_sigma = false; + for (const auto& [key, value] : parse_key_value_args(sd_img_gen_params->sample_params.extra_sample_args, "img2img arg")) { + if (key == "strength_as_noise_level") { + if (!parse_strict_bool(value, strength_as_noise_level)) { + LOG_WARN("ignoring invalid img2img sample arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "force_first_sigma") { + if (!parse_strict_bool(value, force_first_sigma)) { + LOG_WARN("ignoring invalid img2img sample arg '%s=%s'", key.c_str(), value.c_str()); + } + } + } + + size_t t_enc; + float target_sigma = -1; + if (!strength_as_noise_level) { + t_enc = static_cast(plan->sample_steps * request->strength); + if (t_enc == static_cast(plan->sample_steps)) { + t_enc--; + } + } else { + LOG_VERBOSE("Interpreting denoise strength as relative noise level"); + // assume x_noised = K * (x * (1-noise_level) + noise * noise_level) = K * lerp(x, noise, noise_level) + // K = 1, noise_level = sigma for flow models + // K = 1+sigma, noise_level=sigma/(1+sigma) for diffusion models + float target_noise_level = request->strength; + target_sigma = sd->denoiser->noise_level_to_sigma(target_noise_level); + size_t start_index = 0; + for (size_t i = 0; i < plan->sigmas.size(); ++i) { + if (plan->sigmas[i] <= target_sigma) { + start_index = i; + break; + } + } + + if (start_index >= plan->sigmas.size() - 1) { + start_index = plan->sigmas.size() - 2; // Leave at least 1 step + } + t_enc = plan->sample_steps - start_index - 1; + } + LOG_INFO("target t_enc is %zu steps", t_enc); + std::vector sigma_sched; + sigma_sched.assign(plan->sigmas.begin() + plan->sample_steps - t_enc - 1, plan->sigmas.end()); + + if (target_sigma > 0 && force_first_sigma && strength_as_noise_level) { + LOG_VERBOSE("force_first_sigma to %.4f (from %.4f)", target_sigma, sigma_sched[0]); + sigma_sched[0] = target_sigma; + } + + plan->sigmas = std::move(sigma_sched); + plan->sample_steps = static_cast(plan->sigmas.size() - 1); + } + + init_image_tensor = ensure_image_tensor_channels(sd_image_to_tensor(sd_img_gen_params->init_image, request->width, request->height), + image_channels); + } + + if (sd_img_gen_params->mask_image.data != nullptr) { + mask_image_tensor = sd_image_to_tensor(sd_img_gen_params->mask_image, request->width, request->height); + mask_image_tensor = sd::ops::round(mask_image_tensor); + } + + if (sd_img_gen_params->control_image.data != nullptr) { + control_image_tensor = sd_image_to_tensor(sd_img_gen_params->control_image, request->width, request->height); + } + + if (init_image_tensor.empty() || mask_image_tensor.empty()) { + if (sd_version_is_inpaint(sd->version)) { + LOG_WARN("inpainting model requires both an init image and a mask image."); + } + } + + if (mask_image_tensor.empty()) { + mask_image_tensor = sd::full({request->width, request->height, 1, 1}, 1.f); + } + + sd::Tensor latent_mask = sd::ops::interpolate(mask_image_tensor, + {request->width / request->vae_scale_factor, + request->height / request->vae_scale_factor, + 1, + 1}, + sd::ops::InterpolateMode::NearestMax); + + sd::Tensor init_latent; + sd::Tensor control_latent; + if (init_image_tensor.empty()) { + if (sd->version == VERSION_QWEN_IMAGE_LAYERED) { + init_latent = sd->generate_init_latent(request->width, request->height, request->qwen_image_layers + 1, true); + } else { + init_latent = sd->generate_init_latent(request->width, request->height); + } + } else { + init_latent = sd->encode_first_stage(init_image_tensor); + if (init_latent.empty()) { + LOG_ERROR("failed to encode init image"); + return std::nullopt; + } + } + + if (sd->animatediff_num_frames > 1 && + init_latent.dim() >= 4 && init_latent.shape()[3] == 1) { + int n_frames = sd->animatediff_num_frames; + std::vector shape(init_latent.shape().begin(), init_latent.shape().end()); + shape[3] = n_frames; + if (!init_image_tensor.empty()) { + sd::Tensor replicated(shape); + for (int f = 0; f < n_frames; ++f) { + sd::ops::slice_assign(&replicated, 3, f, f + 1, init_latent); + } + init_latent = std::move(replicated); + } else { + init_latent = sd::Tensor(std::move(shape)); + } + } + + if (!control_image_tensor.empty()) { + control_latent = sd->encode_first_stage(control_image_tensor); + if (control_latent.empty()) { + LOG_ERROR("failed to encode control image"); + return std::nullopt; + } + } + + std::vector> ref_images; + for (int i = 0; i < sd_img_gen_params->ref_images_count; i++) { + ref_images.push_back(ensure_image_tensor_channels(sd_image_to_tensor(sd_img_gen_params->ref_images[i]), + image_channels)); + } + + if (ref_images.empty() && sd_version_is_unet_edit(sd->version)) { + LOG_WARN("This model needs at least one reference image; using an empty reference"); + ref_images.push_back(sd::zeros({request->width, request->height, image_channels, 1})); + request->guidance.img_cfg = request->guidance.txt_cfg; + request->use_img_uncond = false; + } + + if (!ref_images.empty()) { + LOG_INFO("EDIT mode"); + } + + std::vector> ref_latents; + for (size_t i = 0; i < ref_images.size(); i++) { + if (sd->version == VERSION_HIDREAM_O1) { + continue; + } + sd::Tensor ref_latent; + if (ref_image_params.resize_before_vae && !sd_version_is_pid(sd->version)) { + LOG_VERBOSE("auto resize ref images"); + double vae_width; + double vae_height; + if (ref_image_params.resize_vae_to_target) { + vae_width = request->width; + vae_height = request->height; + } else { + int target_pixels = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : 1024 * 1024; + int vae_image_size = std::min(target_pixels, request->width * request->height); + vae_width = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]); + vae_height = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0]; + } + + int factor = sd_version_is_qwen_image(sd->version) ? 32 : 16; + vae_height = round(vae_height / factor) * factor; + vae_width = round(vae_width / factor) * factor; + + auto resized_ref_img = sd::ops::interpolate(ref_images[i], + {static_cast(vae_width), + static_cast(vae_height), + ref_images[i].shape()[2], + ref_images[i].shape()[3]}); + + LOG_VERBOSE("resize vae ref image %d from %" PRId64 "x%" PRId64 " to %" PRId64 "x%" PRId64, + static_cast(i), + ref_images[i].shape()[1], + ref_images[i].shape()[0], + resized_ref_img.shape()[1], + resized_ref_img.shape()[0]); + + ref_latent = sd->encode_first_stage(resized_ref_img); + } else { + ref_latent = sd->encode_first_stage(ref_images[i]); + } + if (ref_latent.empty()) { + LOG_ERROR("failed to encode reference image %d", static_cast(i)); + return std::nullopt; + } + + ref_latents.push_back(std::move(ref_latent)); + } + + if (sd_version_is_pid(sd->version)) { + if (ref_latents.empty()) { + LOG_ERROR("PiD requires a reference image"); + return std::nullopt; + } + } + + sd::Tensor concat_latent; + sd::Tensor img_uncond_concat_latent; + if (sd_version_is_inpaint(sd->version)) { + sd::Tensor masked_init_latent; + + if (sd->version != VERSION_FLEX_2) { + if (!init_image_tensor.empty()) { + auto masked_image = ((1.0f - mask_image_tensor) * (init_image_tensor - 0.5f)) + 0.5f; + masked_init_latent = sd->encode_first_stage(masked_image); + if (masked_init_latent.empty()) { + LOG_ERROR("failed to encode masked init image"); + return std::nullopt; + } + } else { + masked_init_latent = sd::Tensor::zeros_like(init_latent); + } + } else { + masked_init_latent = ((1.0f - latent_mask) * init_latent); + } + + auto uncond_masked_init_latent = sd::Tensor::zeros_like(masked_init_latent); + + if (sd->version == VERSION_FLUX_FILL) { + auto mask = mask_image_tensor.reshape({request->vae_scale_factor, + request->width / request->vae_scale_factor, + request->vae_scale_factor, + request->height / request->vae_scale_factor}); + mask = mask.permute({1, 3, 0, 2}).reshape({request->width / request->vae_scale_factor, request->height / request->vae_scale_factor, request->vae_scale_factor * request->vae_scale_factor, 1}); + + concat_latent = sd::ops::concat(masked_init_latent, mask, 2); + img_uncond_concat_latent = sd::ops::concat(uncond_masked_init_latent, mask, 2); + } else if (sd->version == VERSION_FLEX_2) { + concat_latent = sd::ops::concat(masked_init_latent, latent_mask, 2); + if (!control_latent.empty()) { + concat_latent = sd::ops::concat(concat_latent, control_latent, 2); + } else { + concat_latent = sd::ops::concat(concat_latent, sd::Tensor::zeros_like(masked_init_latent), 2); + } + + img_uncond_concat_latent = sd::ops::concat(uncond_masked_init_latent, latent_mask, 2); + img_uncond_concat_latent = sd::ops::concat(img_uncond_concat_latent, sd::Tensor::zeros_like(masked_init_latent), 2); + } else { // SD1.x SD2.x SDXL inpaint + concat_latent = sd::ops::concat(latent_mask, masked_init_latent, 2); + img_uncond_concat_latent = sd::ops::concat(latent_mask, uncond_masked_init_latent, 2); + } + } + if (sd_version_is_unet_edit(sd->version)) { + concat_latent = sd::ops::interpolate(ref_latents[0], init_latent.shape()); + img_uncond_concat_latent = sd::Tensor::zeros_like(concat_latent); + } + if (sd->version == VERSION_FLUX_CONTROLS) { + if (!control_latent.empty()) { + concat_latent = control_latent; + } else { + concat_latent = sd::Tensor::zeros_like(init_latent); + } + img_uncond_concat_latent = sd::Tensor::zeros_like(concat_latent); + } + + if (sd_img_gen_params->init_image.data != nullptr || sd_img_gen_params->ref_images_count > 0) { + int64_t t1 = ggml_time_ms(); + LOG_INFO("encode_first_stage completed, taking %.2fs", (t1 - prepare_start_ms) * 1.0f / 1000); + } + + ImageGenerationLatents latents; + latents.init_latent = std::move(init_latent); + latents.concat_latent = std::move(concat_latent); + latents.img_uncond_concat_latent = std::move(img_uncond_concat_latent); + latents.control_image = std::move(control_image_tensor); + latents.ref_images = std::move(ref_images); + latents.ref_latents = std::move(ref_latents); + + if (sd_version_is_inpaint(sd->version)) { + latent_mask = sd::ops::max_pool_2d(latent_mask, + {3, 3}, + {1, 1}, + {1, 1}); + } + latents.denoise_mask = std::move(latent_mask); + + return latents; + } + + static std::optional prepare_image_generation_embeds(StableDiffusionGGML* sd, + const sd_img_gen_params_t* sd_img_gen_params, + GenerationRequest* request, + SamplePlan* plan, + ImageGenerationLatents* latents, + const RefImageParams& ref_image_params) { + ConditionerRunnerEndOnExit conditioner_runner_end{sd->cond_stage_model.get()}; + + ConditionerParams condition_params; + condition_params.text = request->prompt; + condition_params.clip_skip = request->clip_skip; + condition_params.width = request->width; + condition_params.height = request->height; + if (ref_image_params.pass_to_vlm) { + condition_params.ref_images = &latents->ref_images; + } + + condition_params.ref_image_params = ref_image_params; + + sd->prepare_generation_extensions(request->pm_params, + request->pulid_params, + condition_params, + plan->total_steps); + sd->compute_ip_adapter_tokens(sd_img_gen_params->ip_adapter_image, sd_img_gen_params->ip_adapter_strength); + int64_t prepare_start_ms = ggml_time_ms(); + condition_params.zero_out_masked = false; + auto cond = sd->cond_stage_model->get_learned_condition(sd->n_threads, + condition_params); + if (cond.c_concat.empty() && ref_image_params.pass_to_dit) { + cond.c_concat = latents->concat_latent; // TODO: optimize + } + + bool use_ref_latent_img_cfg = request->use_img_uncond && + !latents->ref_images.empty() && + sd_version_supports_ref_latent_img_cfg(sd->version); + + SDCondition uncond; + if (request->use_uncond || request->use_high_noise_uncond) { + if (sd_version_is_ideogram4(sd->version)) { + uncond.c_vector = sd::Tensor::from_vector({1.0f}); + } else if (sd_version_is_minit2i(sd->version)) { + // MiniT2I derives the unconditional signal from the same T5 hidden + // states with a zeroed prompt mask, so no extra text encode is needed. + uncond.c_crossattn = cond.c_crossattn; + uncond.c_vector = sd::Tensor::zeros_like(cond.c_vector); + } else { + bool zero_out_masked = false; + if (sd_version_is_sdxl(sd->version) && + request->negative_prompt.empty() && + !sd->is_using_edm_v_parameterization) { + zero_out_masked = true; + } + condition_params.text = request->negative_prompt; + condition_params.zero_out_masked = zero_out_masked; + uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, + condition_params); + } + if (uncond.c_concat.empty() && ref_image_params.pass_to_dit) { + uncond.c_concat = latents->concat_latent; // TODO: optimize + } + } + + SDCondition img_uncond; + if (request->use_img_uncond) { + if ((request->use_uncond || request->use_high_noise_uncond) && (latents->ref_images.empty() || !use_ref_latent_img_cfg)) { + img_uncond = SDCondition(uncond.c_crossattn, uncond.c_vector, latents->img_uncond_concat_latent); + } else { + bool zero_out_masked = false; + if (sd_version_is_sdxl(sd->version) && + request->negative_prompt.empty() && + !sd->is_using_edm_v_parameterization) { + zero_out_masked = true; + } + condition_params.text = request->negative_prompt; + condition_params.zero_out_masked = zero_out_masked; + std::vector> empty_ref_images; + if (use_ref_latent_img_cfg) { + condition_params.ref_images = &empty_ref_images; + } + img_uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, + condition_params); + if (img_uncond.c_concat.empty() && ref_image_params.pass_to_dit) { + img_uncond.c_concat = latents->img_uncond_concat_latent; // TODO: optimize + } + } + } + + int64_t t1 = ggml_time_ms(); + LOG_INFO("get_learned_condition completed, taking %.2fs", (t1 - prepare_start_ms) * 1.0f / 1000); + + ImageGenerationEmbeds embeds; + embeds.img_uncond = std::move(img_uncond); + embeds.cond = std::move(cond); + embeds.uncond = std::move(uncond); + + return embeds; + } + + static sd_image_t* decode_image_outputs(StableDiffusionGGML* sd, + const GenerationRequest& request, + const std::vector>& final_latents, + int* num_images_out) { + if (final_latents.empty()) { + LOG_ERROR("no latent images to decode"); + return nullptr; + } + if (final_latents.size() > static_cast(request.batch_count)) { + LOG_ERROR("expected at most %d latents, got %zu", request.batch_count, final_latents.size()); + return nullptr; + } + if (final_latents.size() < static_cast(request.batch_count)) { + LOG_INFO("decoding %zu/%d latents", final_latents.size(), request.batch_count); + } else { + LOG_INFO("decoding %zu latents", final_latents.size()); + } + std::vector> decoded_images; + int64_t t0 = ggml_time_ms(); + bool cancelled = false; + + for (size_t i = 0; i < final_latents.size(); i++) { + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling latent decodings"); + cancelled = true; + break; + } + int64_t t1 = ggml_time_ms(); + if (sd->version == VERSION_QWEN_IMAGE_LAYERED) { + int qwen_image_latent_layers = request.qwen_image_layers + 1; + if (final_latents[i].dim() < 5 || final_latents[i].shape()[2] < qwen_image_latent_layers) { + LOG_ERROR("qwen image layered expected at least %d latent layers, got shape dim=%d", + qwen_image_latent_layers, + final_latents[i].dim()); + return nullptr; + } + for (int layer_index = 0; layer_index < qwen_image_latent_layers; layer_index++) { + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling latent decodings"); + cancelled = true; + break; + } + sd::Tensor layer_latent = sd::ops::slice(final_latents[i], 2, layer_index, layer_index + 1); + layer_latent.squeeze_(2); + sd::Tensor image = sd->decode_first_stage(layer_latent); + if (image.empty()) { + LOG_ERROR("decode_first_stage failed for latent %zu layer %d", i + 1, layer_index + 1); + return nullptr; + } + decoded_images.push_back(std::move(image)); + } + if (cancelled) { + break; + } + } else if (sd->animatediff_num_frames > 1 && + final_latents[i].dim() >= 4 && + final_latents[i].shape()[3] == sd->animatediff_num_frames) { + int n_frames = sd->animatediff_num_frames; + for (int f = 0; f < n_frames; ++f) { + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling latent decodings"); + cancelled = true; + break; + } + sd::Tensor frame_latent = sd::ops::slice(final_latents[i], 3, f, f + 1); + sd::Tensor image = sd->decode_first_stage(frame_latent); + if (image.empty()) { + LOG_ERROR("decode_first_stage failed for AnimateDiff frame %d/%d", f + 1, n_frames); + return nullptr; + } + decoded_images.push_back(std::move(image)); + } + } else { + sd::Tensor image = sd->decode_first_stage(final_latents[i]); + if (image.empty()) { + LOG_ERROR("decode_first_stage failed for latent %" PRId64, i + 1); + return nullptr; + } + decoded_images.push_back(std::move(image)); + } + int64_t t2 = ggml_time_ms(); + LOG_INFO("latent %zu decoded, taking %.2fs", i + 1, (t2 - t1) * 1.0f / 1000); + } + + int64_t t4 = ggml_time_ms(); + LOG_INFO("decode_first_stage completed, taking %.2fs", (t4 - t0) * 1.0f / 1000); + if (decoded_images.empty()) { + LOG_ERROR(cancelled ? "cancelled before any latent images were decoded" : "no decoded images"); + return nullptr; + } + + int image_count = static_cast(decoded_images.size()); + sd_image_t* result_images = (sd_image_t*)calloc(image_count, sizeof(sd_image_t)); + if (result_images == nullptr) { + return nullptr; + } + if (num_images_out != nullptr) { + *num_images_out = image_count; + } + + for (size_t i = 0; i < decoded_images.size(); i++) { + result_images[i] = tensor_to_sd_image(decoded_images[i]); + } + + return result_images; + } + + static sd::Tensor upscale_hires_latent(StableDiffusionGGML* sd, + const sd::Tensor& latent, + const GenerationRequest& request, + UpscalerGGML* upscaler) { + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling hires latent upscale"); + return {}; + } + + auto get_hires_latent_target_shape = [&]() { + std::vector target_shape = latent.shape(); + if (target_shape.size() < 2) { + target_shape.clear(); + return target_shape; + } + target_shape[0] = request.hires.target_width / request.vae_scale_factor; + target_shape[1] = request.hires.target_height / request.vae_scale_factor; + return target_shape; + }; + + if (request.hires.upscaler == SD_HIRES_UPSCALER_LATENT || + request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_NEAREST || + request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_NEAREST_EXACT || + request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_ANTIALIASED || + request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_BICUBIC || + request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_BICUBIC_ANTIALIASED) { + std::vector target_shape = get_hires_latent_target_shape(); + if (target_shape.empty()) { + LOG_ERROR("latent has invalid shape for hires upscale"); + return {}; + } + + sd::ops::InterpolateMode mode = sd::ops::InterpolateMode::Nearest; + bool antialias = false; + switch (request.hires.upscaler) { + case SD_HIRES_UPSCALER_LATENT: + mode = sd::ops::InterpolateMode::Bilinear; + break; + case SD_HIRES_UPSCALER_LATENT_NEAREST: + mode = sd::ops::InterpolateMode::Nearest; + break; + case SD_HIRES_UPSCALER_LATENT_NEAREST_EXACT: + mode = sd::ops::InterpolateMode::NearestExact; + break; + case SD_HIRES_UPSCALER_LATENT_ANTIALIASED: + mode = sd::ops::InterpolateMode::Bilinear; + antialias = true; + break; + case SD_HIRES_UPSCALER_LATENT_BICUBIC: + mode = sd::ops::InterpolateMode::Bicubic; + break; + case SD_HIRES_UPSCALER_LATENT_BICUBIC_ANTIALIASED: + mode = sd::ops::InterpolateMode::Bicubic; + antialias = true; + break; + default: + break; + } + + LOG_INFO("hires %s upscale %" PRId64 "x%" PRId64 " -> %" PRId64 "x%" PRId64, + sd_hires_upscaler_name(request.hires.upscaler), + latent.shape()[0], + latent.shape()[1], + target_shape[0], + target_shape[1]); + + return sd::ops::interpolate(latent, target_shape, mode, false, antialias); + } else if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL || + request.hires.upscaler == SD_HIRES_UPSCALER_LANCZOS || + request.hires.upscaler == SD_HIRES_UPSCALER_NEAREST) { + if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL && upscaler == nullptr) { + LOG_ERROR("hires model upscaler context is null"); + return {}; + } + + sd::Tensor decoded = sd->decode_first_stage(latent); + if (decoded.empty()) { + LOG_ERROR("decode_first_stage failed before hires %s upscale", + sd_hires_upscaler_name(request.hires.upscaler)); + return {}; + } + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling hires image upscale"); + return {}; + } + + sd::Tensor upscaled_tensor; + if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) { + upscaled_tensor = upscaler->upscale_tensor(decoded); + if (upscaled_tensor.empty()) { + LOG_ERROR("hires model upscale failed"); + return {}; + } + + if (upscaled_tensor.shape()[0] != request.hires.target_width || + upscaled_tensor.shape()[1] != request.hires.target_height) { + upscaled_tensor = sd::ops::interpolate(upscaled_tensor, + {request.hires.target_width, + request.hires.target_height, + upscaled_tensor.shape()[2], + upscaled_tensor.shape()[3]}); + } + } else { + sd::ops::InterpolateMode mode = request.hires.upscaler == SD_HIRES_UPSCALER_LANCZOS + ? sd::ops::InterpolateMode::Lanczos + : sd::ops::InterpolateMode::Nearest; + LOG_INFO("hires %s image upscale %" PRId64 "x%" PRId64 " -> %dx%d", + sd_hires_upscaler_name(request.hires.upscaler), + decoded.shape()[0], + decoded.shape()[1], + request.hires.target_width, + request.hires.target_height); + upscaled_tensor = sd::ops::interpolate(decoded, + {request.hires.target_width, + request.hires.target_height, + decoded.shape()[2], + decoded.shape()[3]}, + mode); + upscaled_tensor = sd::ops::clamp(upscaled_tensor, 0.0f, 1.0f); + } + + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling hires latent encode"); + return {}; + } + sd::Tensor upscaled_latent = sd->encode_first_stage(upscaled_tensor); + if (upscaled_latent.empty()) { + LOG_ERROR("encode_first_stage failed after hires %s upscale", + sd_hires_upscaler_name(request.hires.upscaler)); + } + return upscaled_latent; + } + + LOG_ERROR("unsupported hires upscaler '%s'", sd_hires_upscaler_name(request.hires.upscaler)); + return {}; + } + + bool generate_image(StableDiffusionGGML* sd, + const sd_img_gen_params_t* sd_img_gen_params, + sd_image_t** images_out, + int* num_images_out) { + if (images_out != nullptr) { + *images_out = nullptr; + } + if (num_images_out != nullptr) { + *num_images_out = 0; + } + if (sd == nullptr || sd_img_gen_params == nullptr) { + return false; + } + + // MiniMax-H3 is video-only. Its denoiser always splits the packed latent into a video and an + // audio half, and only generate_video ever computes the audio length, so reaching this + // function with an H3 checkpoint is guaranteed to die on + // GGML_ASSERT(!audio_input_cache.empty()) with a core dump, after the several minutes it + // takes to load the weights, and with nothing in the output pointing at the missing --mode. + // (The AnimateDiff path below routes vid_gen back through here, but that is SD1.5 plus a + // motion module, never H3.) + if (sd_version_is_minimax_h3(sd->version)) { + LOG_ERROR("MiniMax-H3 is a video model and cannot be run in img_gen mode; use --mode vid_gen"); + return false; + } + + sd->reset_cancel_flag(); + + int64_t t0 = ggml_time_ms(); + sd->vae_tiling_params = sd_img_gen_params->vae_tiling_params; + GenerationRequest request(sd, sd_img_gen_params); + LOG_INFO("generate_image %dx%d", request.width, request.height); + + sd->rng->manual_seed(request.seed); + sd->sampler_rng->manual_seed(request.seed); + sd->set_flow_shift(sd_img_gen_params->sample_params.flow_shift); + if (!sd->apply_loras(sd_img_gen_params->loras, sd_img_gen_params->lora_count)) + return false; + sd->apply_circular_axes(sd_img_gen_params->circular_x, sd_img_gen_params->circular_y); + + const RefImageParams ref_image_params = sd->resolve_ref_image_params(sd_img_gen_params->ref_image_args); + + ImageVaeAxesGuard axes_guard(sd, sd_img_gen_params, request); + + SamplePlan plan(sd, sd_img_gen_params, request); + auto latents_opt = prepare_image_generation_latents(sd, + sd_img_gen_params, + &request, + &plan, + ref_image_params); + if (!latents_opt.has_value()) { + return false; + } + ImageGenerationLatents latents = std::move(*latents_opt); + + auto embeds_opt = prepare_image_generation_embeds(sd, + sd_img_gen_params, + &request, + &plan, + &latents, + ref_image_params); + if (!embeds_opt.has_value()) { + return false; + } + ImageGenerationEmbeds embeds = std::move(*embeds_opt); + + std::vector> final_latents; + int64_t denoise_start = ggml_time_ms(); + for (int b = 0; b < request.batch_count; b++) { + sd_cancel_mode_t cancel = sd->get_cancel_flag(); + if (cancel == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation"); + return false; + } + if (cancel == SD_CANCEL_NEW_LATENTS) { + LOG_INFO("cancelling new latent generation, returning %zu/%d completed latents", + final_latents.size(), + request.batch_count); + break; + } + + int64_t sampling_start = ggml_time_ms(); + int64_t cur_seed = request.seed + b; + LOG_INFO("generating image: %i/%i - seed %" PRId64, b + 1, request.batch_count, cur_seed); + + sd->rng->manual_seed(cur_seed); + sd->sampler_rng->manual_seed(cur_seed); + sd::Tensor noise = sd::randn_like(latents.init_latent, sd->rng); + + sd::Tensor x_0 = sd->sample(sd->diffusion_model, + true, + latents.init_latent, + std::move(noise), + embeds.cond, + embeds.uncond, + embeds.img_uncond, + latents.control_image, + request.control_strength, + request.guidance, + plan.eta, + request.shifted_timestep, + plan.sample_method, + sd->is_flow_denoiser(), + plan.extra_sample_args, + plan.sigmas, + latents.ref_latents, + ref_image_params, + latents.denoise_mask, + sd::Tensor(), + 1.f, + 0, + static_cast(request.fps), + request.cache_params, + true); + int64_t sampling_end = ggml_time_ms(); + if (!x_0.empty()) { + LOG_INFO("sampling completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); + final_latents.push_back(std::move(x_0)); + continue; + } + + LOG_ERROR("sampling for image %d/%d failed after %.2fs", + b + 1, + request.batch_count, + (sampling_end - sampling_start) * 1.0f / 1000); + return false; + } + int64_t denoise_end = ggml_time_ms(); + LOG_INFO("generating %zu latent images completed, taking %.2fs", + final_latents.size(), + (denoise_end - denoise_start) * 1.0f / 1000); + if (final_latents.empty()) { + LOG_ERROR("no latent images generated"); + return false; + } + + if (request.hires.enabled && request.hires.target_width > 0) { + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before hires fix"); + return false; + } + LOG_INFO("hires fix: upscaling to %dx%d", request.hires.target_width, request.hires.target_height); + + std::unique_ptr hires_upscaler; + if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) { + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before hires model load"); + return false; + } + LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path); + hires_upscaler = std::make_unique(sd->n_threads, + false, + request.hires.upscale_tile_size, + sd->backend_spec, + sd->params_backend_spec); + const size_t max_graph_vram_bytes = sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER); + hires_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes); + if (!hires_upscaler->load_from_file(request.hires.model_path, + sd->n_threads)) { + LOG_ERROR("load hires model upscaler failed"); + return false; + } + } + + int hires_scheduler_steps = 0; + std::vector hires_sigma_sched = + make_hires_sigma_schedule(sd, + request.hires, + sd_img_gen_params->sample_params, + plan.sample_method, + plan.sample_steps, + sd->get_image_seq_len(request.hires.target_height, request.hires.target_width), + &hires_scheduler_steps); + LOG_INFO("hires fix: scheduler_steps=%d, denoising_strength=%.2f, sigma_sched_size=%zu%s", + hires_scheduler_steps, + request.hires.denoising_strength, + hires_sigma_sched.size(), + request.hires.custom_sigmas_count > 0 ? ", custom_sigmas=true" : ""); + + std::vector> hires_final_latents; + int64_t hires_denoise_start = ggml_time_ms(); + for (int b = 0; b < (int)final_latents.size(); b++) { + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation during hires fix"); + return false; + } + int64_t cur_seed = request.seed + b; + sd->rng->manual_seed(cur_seed); + sd->sampler_rng->manual_seed(cur_seed); + + sd::Tensor upscaled = upscale_hires_latent(sd, + final_latents[b], + request, + hires_upscaler.get()); + if (upscaled.empty()) { + return false; + } + + sd::Tensor noise = sd::randn_like(upscaled, sd->rng); + + sd::Tensor hires_denoise_mask; + if (!latents.denoise_mask.empty()) { + std::vector mask_shape = latents.denoise_mask.shape(); + mask_shape[0] = upscaled.shape()[0]; + mask_shape[1] = upscaled.shape()[1]; + hires_denoise_mask = sd::ops::interpolate(latents.denoise_mask, + mask_shape, + sd::ops::InterpolateMode::NearestMax); + } + + int64_t hires_sample_start = ggml_time_ms(); + sd::Tensor x_0 = sd->sample(sd->diffusion_model, + true, + upscaled, + std::move(noise), + embeds.cond, + embeds.uncond, + embeds.img_uncond, + latents.control_image, + request.control_strength, + request.guidance, + plan.eta, + request.shifted_timestep, + plan.sample_method, + sd->is_flow_denoiser(), + plan.extra_sample_args, + hires_sigma_sched, + latents.ref_latents, + ref_image_params, + hires_denoise_mask, + sd::Tensor(), + 1.f, + 0, + static_cast(request.fps), + request.cache_params, + false); + int64_t hires_sample_end = ggml_time_ms(); + if (!x_0.empty()) { + LOG_INFO("hires sampling %d/%d completed, taking %.2fs", + b + 1, + (int)final_latents.size(), + (hires_sample_end - hires_sample_start) * 1.0f / 1000); + hires_final_latents.push_back(std::move(x_0)); + continue; + } + + LOG_ERROR("hires sampling for image %d/%d failed after %.2fs", + b + 1, + (int)final_latents.size(), + (hires_sample_end - hires_sample_start) * 1.0f / 1000); + return false; + } + int64_t hires_denoise_end = ggml_time_ms(); + LOG_INFO("hires fix completed, taking %.2fs", (hires_denoise_end - hires_denoise_start) * 1.0f / 1000); + + final_latents = std::move(hires_final_latents); + } + + int num_images = 0; + auto result = decode_image_outputs(sd, request, final_latents, &num_images); + if (result == nullptr) { + return false; + } + + sd->lora_stat(); + + int64_t t1 = ggml_time_ms(); + LOG_INFO("generate_image completed in %.2fs", (t1 - t0) * 1.0f / 1000); + if (num_images_out != nullptr) { + *num_images_out = num_images; + } + if (images_out != nullptr) { + *images_out = result; + } else { + free_sd_images(result, num_images); + } + return true; + } + +} // namespace sd::pipeline diff --git a/src/model_builders.cpp b/src/pipeline/model_builders.cpp similarity index 100% rename from src/model_builders.cpp rename to src/pipeline/model_builders.cpp diff --git a/src/model_builders.h b/src/pipeline/model_builders.h similarity index 93% rename from src/model_builders.h rename to src/pipeline/model_builders.h index ac95a6f65..73d134b3b 100644 --- a/src/model_builders.h +++ b/src/pipeline/model_builders.h @@ -1,5 +1,5 @@ -#ifndef __SD_MODEL_BUILDERS_H__ -#define __SD_MODEL_BUILDERS_H__ +#ifndef __SD_PIPELINE_MODEL_BUILDERS_H__ +#define __SD_PIPELINE_MODEL_BUILDERS_H__ #include #include @@ -60,4 +60,4 @@ namespace sd::model_builders { } // namespace sd::model_builders -#endif // __SD_MODEL_BUILDERS_H__ +#endif // __SD_PIPELINE_MODEL_BUILDERS_H__ diff --git a/src/pipeline/request.cpp b/src/pipeline/request.cpp new file mode 100644 index 000000000..aa78f7e6b --- /dev/null +++ b/src/pipeline/request.cpp @@ -0,0 +1,471 @@ +#include "request.h" + +#include +#include +#include +#include + +#include "diffusion_engine.h" +#include "runtime/denoiser.hpp" + +namespace sd::pipeline { + + const char* sampling_methods_str[] = { + "Euler", + "Euler A", + "Heun", + "DPM2", + "DPM++ (2s)", + "DPM++ (2M)", + "modified DPM++ (2M)", + "iPNDM", + "iPNDM_v", + "LCM", + "DDIM \"trailing\"", + "TCD", + "Res Multistep", + "Res 2s", + "ER-SDE", + "Euler CFG++", + "Euler A CFG++", + "Euler GE", + "DPM++ (2M) SDE", + "DPM++ (2M) SDE BT", + "LMS", + }; + + static_assert(SAMPLE_METHOD_COUNT == sizeof(sampling_methods_str) / sizeof(sampling_methods_str[0]), + "\nnumber of elements in sampling_methods_str[] != SAMPLE_METHOD_COUNT"); + + static bool sd_version_supports_img_cfg(SDVersion version, bool has_ref_images) { + return sd_version_is_inpaint_or_unet_edit(version) || + (has_ref_images && sd_version_supports_ref_latent_img_cfg(version)); + } + + enum sample_method_t default_sample_method(const StableDiffusionGGML* sd) { + if (sd != nullptr) { + if (sd_version_is_pid(sd->version)) { + return LCM_SAMPLE_METHOD; + } + if (sd_version_is_dit(sd->version)) { + return EULER_SAMPLE_METHOD; + } + } + return EULER_A_SAMPLE_METHOD; + } + + enum scheduler_t default_scheduler(const StableDiffusionGGML* sd, enum sample_method_t sample_method) { + if (sd != nullptr) { + auto edm_v_denoiser = std::dynamic_pointer_cast(sd->denoiser); + if (edm_v_denoiser) { + return EXPONENTIAL_SCHEDULER; + } + } + if (sample_method == LCM_SAMPLE_METHOD || sample_method == TCD_SAMPLE_METHOD) { + return LCM_SCHEDULER; + } else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) { + return SIMPLE_SCHEDULER; + } else if (sd != nullptr && sd_version_is_flux(sd->version)) { + return FLUX_SCHEDULER; + } else if (sd != nullptr && sd_version_is_flux2(sd->version)) { + return FLUX2_SCHEDULER; + } else if (sd != nullptr && sd_version_is_ltxav(sd->version)) { + return LTX2_SCHEDULER; + } else if (sd != nullptr && sd_version_is_ideogram4(sd->version)) { + return LOGIT_NORMAL_SCHEDULER; + } + return DISCRETE_SCHEDULER; + } + + static int64_t resolve_seed(int64_t seed) { + if (seed >= 0) { + return seed; + } + srand((int)time(nullptr)); + return rand(); + } + + static enum sample_method_t resolve_sample_method(StableDiffusionGGML* sd, enum sample_method_t sample_method) { + if (sample_method == SAMPLE_METHOD_COUNT) { + return default_sample_method(sd); + } + return sample_method; + } + + static scheduler_t resolve_scheduler(StableDiffusionGGML* sd, + scheduler_t scheduler, + enum sample_method_t sample_method) { + if (scheduler == SCHEDULER_COUNT) { + return default_scheduler(sd, sample_method); + } + return scheduler; + } + + float resolve_eta(StableDiffusionGGML* sd, + float eta, + enum sample_method_t sample_method) { + if (eta == INFINITY) { + if (sd->version == VERSION_HIDREAM_O1) { + return 8.f; + } + switch (sample_method) { + case DDIM_TRAILING_SAMPLE_METHOD: + case TCD_SAMPLE_METHOD: + case RES_MULTISTEP_SAMPLE_METHOD: + case RES_2S_SAMPLE_METHOD: + return 0.0f; + case EULER_A_SAMPLE_METHOD: + case DPMPP2S_A_SAMPLE_METHOD: + case ER_SDE_SAMPLE_METHOD: + case EULER_A_CFG_PP_SAMPLE_METHOD: + case DPMPP2M_SDE_SAMPLE_METHOD: + case DPMPP2M_SDE_BT_SAMPLE_METHOD: + return 1.0f; + default:; + } + return 0.0f; + } + return eta; + } + + GenerationRequest::GenerationRequest(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params) { + prompt = SAFE_STR(sd_img_gen_params->prompt); + negative_prompt = SAFE_STR(sd_img_gen_params->negative_prompt); + width = sd_img_gen_params->width; + height = sd_img_gen_params->height; + vae_scale_factor = sd->get_vae_scale_factor(); + diffusion_model_down_factor = sd->get_diffusion_model_down_factor(); + seed = sd_img_gen_params->seed; + batch_count = sd_img_gen_params->batch_count; + qwen_image_layers = std::max(0, sd_img_gen_params->qwen_image_layers); + clip_skip = sd_img_gen_params->clip_skip; + shifted_timestep = sd_img_gen_params->sample_params.shifted_timestep; + strength = sd_img_gen_params->strength; + control_strength = sd_img_gen_params->control_strength; + eta = sd_img_gen_params->sample_params.eta; + has_ref_images = sd_img_gen_params->ref_images_count > 0; + guidance = sd_img_gen_params->sample_params.guidance; + pm_params = sd_img_gen_params->pm_params; + pulid_params = sd_img_gen_params->pulid_params; + hires = sd_img_gen_params->hires; + cache_params = &sd_img_gen_params->cache; + resolve(sd); + } + + GenerationRequest::GenerationRequest(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params) { + prompt = SAFE_STR(sd_vid_gen_params->prompt); + negative_prompt = SAFE_STR(sd_vid_gen_params->negative_prompt); + width = sd_vid_gen_params->width; + height = sd_vid_gen_params->height; + requested_frames = std::max(1, sd_vid_gen_params->video_frames); + frames = sd->align_video_frames(requested_frames); + clip_skip = sd_vid_gen_params->clip_skip; + fps = std::max(1, sd_vid_gen_params->fps); + if (sd_version_is_minimax_h3(sd->version) && fps != 24) { + LOG_WARN("MiniMax-H3 uses 24 fps; overriding requested fps %d", fps); + fps = 24; + } + vae_scale_factor = sd->get_vae_scale_factor(); + diffusion_model_down_factor = sd->get_diffusion_model_down_factor(); + seed = sd_vid_gen_params->seed; + strength = sd_vid_gen_params->strength; + cache_params = &sd_vid_gen_params->cache; + vace_strength = sd_vid_gen_params->vace_strength; + guidance = sd_vid_gen_params->sample_params.guidance; + high_noise_guidance = sd_vid_gen_params->high_noise_sample_params.guidance; + hires = sd_vid_gen_params->hires; + resolve(sd); + if (frames != requested_frames) { + LOG_WARN("align video frames from %d to %d for %s", + requested_frames, + frames, + model_version_to_str[sd->version]); + } + } + + void GenerationRequest::align_generation_request_size() { + align_image_size(&width, &height, "generation request"); + } + + void GenerationRequest::align_image_size(int* target_width, int* target_height, const char* label) { + int spatial_multiple = vae_scale_factor * diffusion_model_down_factor; + int width_offset = align_up_offset(*target_width, spatial_multiple); + int height_offset = align_up_offset(*target_height, spatial_multiple); + if (width_offset <= 0 && height_offset <= 0) { + return; + } + + int original_width = *target_width; + int original_height = *target_height; + + *target_width += width_offset; + *target_height += height_offset; + LOG_WARN("align %s up %dx%d to %dx%d (multiple=%d)", + label, + original_width, + original_height, + *target_width, + *target_height, + spatial_multiple); + } + + void GenerationRequest::resolve_hires() { + if (!hires.enabled) { + return; + } + if (hires.upscaler == SD_HIRES_UPSCALER_NONE) { + hires.enabled = false; + return; + } + if (hires.upscaler < SD_HIRES_UPSCALER_NONE || hires.upscaler >= SD_HIRES_UPSCALER_COUNT) { + LOG_WARN("hires upscaler '%d' is invalid, disabling hires", hires.upscaler); + hires.enabled = false; + return; + } + if (hires.upscaler == SD_HIRES_UPSCALER_MODEL && strlen(SAFE_STR(hires.model_path)) == 0) { + LOG_WARN("hires model upscaler requires a model path, disabling hires"); + hires.enabled = false; + return; + } + if (hires.scale <= 0.f && hires.target_width <= 0 && hires.target_height <= 0) { + LOG_WARN("hires scale must be positive when no target size is set, disabling hires"); + hires.enabled = false; + return; + } + if (hires.custom_sigmas_count < 0) { + LOG_WARN("hires custom sigmas count is negative, ignoring custom sigmas"); + hires.custom_sigmas = nullptr; + hires.custom_sigmas_count = 0; + } + if (hires.custom_sigmas_count > 0 && hires.custom_sigmas == nullptr) { + LOG_WARN("hires custom sigmas count is positive but custom sigmas are null, ignoring custom sigmas"); + hires.custom_sigmas_count = 0; + } + if (hires.custom_sigmas_count == 1) { + LOG_WARN("hires custom sigmas requires at least two values, ignoring custom sigmas"); + hires.custom_sigmas = nullptr; + hires.custom_sigmas_count = 0; + } + hires.denoising_strength = std::clamp(hires.denoising_strength, 0.0001f, 1.f); + hires.steps = std::max(0, hires.steps); + + if (hires.target_width > 0 && hires.target_height > 0) { + // pass + } else if (hires.target_width > 0) { + hires.target_height = hires.target_width; + } else if (hires.target_height > 0) { + hires.target_width = hires.target_height; + } else { + hires.target_width = static_cast(std::round(width * hires.scale)); + hires.target_height = static_cast(std::round(height * hires.scale)); + } + + if (hires.target_width <= 0 || hires.target_height <= 0) { + LOG_WARN("hires target size is not positive, disabling hires"); + hires.enabled = false; + return; + } + align_image_size(&hires.target_width, &hires.target_height, "hires target"); + } + + void GenerationRequest::resolve_guidance(StableDiffusionGGML* sd, + sd_guidance_params_t* guidance, + bool* use_uncond, + bool* use_img_uncond, + bool has_ref_images, + const char* stage_name) { + GGML_ASSERT(guidance != nullptr); + GGML_ASSERT(use_uncond != nullptr); + GGML_ASSERT(use_img_uncond != nullptr); + // out_img_uncond + text_cfg_scale * (out_cond - out_uncond) + image_cfg_scale * (out_uncond - out_img_uncond) + // -> text_cfg_scale * out_cond + (image_cfg_scale - text_cfg_scale) * out_uncond + (1 - image_cfg_scale) * out_img_uncond + // out_cond : prompt, image latent + // out_uncond : negative prompt, image latent + // out_img_uncond : negative prompt, zero image latent + // image_cfg_scale == 1 reduces 3-cond CFG to 2-cond CFG. + bool img_cfg_was_set = std::isfinite(guidance->img_cfg); + if (!img_cfg_was_set) { + guidance->img_cfg = 1.f; + } + + if (!sd_version_supports_img_cfg(sd->version, has_ref_images)) { + if (img_cfg_was_set && guidance->img_cfg != 1.f) { + LOG_WARN("3-conditioning CFG is not supported with this model, disabling it for better performance"); + } + guidance->img_cfg = 1.f; + } + + if (guidance->img_cfg != guidance->txt_cfg) { + *use_uncond = true; + } + + if (guidance->img_cfg != 1.f) { + *use_img_uncond = true; + } + + if (guidance->txt_cfg < 1.f) { + const char* prefix = stage_name == nullptr ? "" : stage_name; + if (guidance->txt_cfg == 0.f) { + LOG_WARN("%sunconditioned mode, images won't follow the prompt (use cfg-scale=1 for distilled models)", + prefix); + } else { + LOG_WARN("%scfg value out of expected range may produce unexpected results", prefix); + } + } + } + + void GenerationRequest::resolve(StableDiffusionGGML* sd) { + align_generation_request_size(); + resolve_hires(); + seed = resolve_seed(seed); + + resolve_guidance(sd, &guidance, &use_uncond, &use_img_uncond, has_ref_images); + if (sd->high_noise_diffusion_model) { + resolve_guidance(sd, + &high_noise_guidance, + &use_high_noise_uncond, + &use_high_noise_img_uncond, + has_ref_images, + "high noise: "); + } + + if (shifted_timestep > 0 && !sd_version_is_sdxl(sd->version)) { + LOG_WARN("timestep shifting is only supported for SDXL models!"); + shifted_timestep = 0; + } + } + + SamplePlan::SamplePlan(StableDiffusionGGML* sd, + const sd_img_gen_params_t* sd_img_gen_params, + const GenerationRequest& request) { + sample_method = sd_img_gen_params->sample_params.sample_method; + extra_sample_args = sd_img_gen_params->sample_params.extra_sample_args; + eta = sd_img_gen_params->sample_params.eta; + sample_steps = sd_img_gen_params->sample_params.sample_steps; + resolve(sd, &request, &sd_img_gen_params->sample_params); + } + + SamplePlan::SamplePlan(StableDiffusionGGML* sd, + const sd_vid_gen_params_t* sd_vid_gen_params, + const GenerationRequest& request) { + sample_method = sd_vid_gen_params->sample_params.sample_method; + extra_sample_args = sd_vid_gen_params->sample_params.extra_sample_args; + eta = sd_vid_gen_params->sample_params.eta; + sample_steps = sd_vid_gen_params->sample_params.sample_steps; + if (sd->high_noise_diffusion_model) { + high_noise_sample_steps = sd_vid_gen_params->high_noise_sample_params.sample_steps; + high_noise_sample_method = sd_vid_gen_params->high_noise_sample_params.sample_method; + high_noise_extra_sample_args = sd_vid_gen_params->high_noise_sample_params.extra_sample_args; + high_noise_eta = sd_vid_gen_params->high_noise_sample_params.eta; + } + moe_boundary = sd_vid_gen_params->moe_boundary; + resolve(sd, &request, &sd_vid_gen_params->sample_params); + } + + void SamplePlan::resolve(StableDiffusionGGML* sd, + const GenerationRequest* request, + const sd_sample_params_t* sample_params) { + sample_method = resolve_sample_method(sd, sample_method); + + total_steps = sample_steps + std::max(0, high_noise_sample_steps); + + if (sample_params->custom_sigmas_count > 0) { + sigmas = std::vector(sample_params->custom_sigmas, + sample_params->custom_sigmas + sample_params->custom_sigmas_count); + total_steps = static_cast(sigmas.size()) - 1; + LOG_WARN("total_steps != custom_sigmas_count - 1, set total_steps to %d", total_steps); + if (sample_steps >= total_steps) { + sample_steps = total_steps; + LOG_WARN("total_steps != custom_sigmas_count - 1, set sample_steps to %d", sample_steps); + } + if (high_noise_sample_steps > 0) { + high_noise_sample_steps = total_steps - sample_steps; + LOG_WARN("total_steps != custom_sigmas_count - 1, set high_noise_sample_steps to %d", high_noise_sample_steps); + } + } else { + scheduler_t scheduler = resolve_scheduler(sd, + sample_params->scheduler, + sample_method); + int sample_seq_len = sd->get_image_seq_len(request->height, request->width); + if (sd_version_is_ltxav(sd->version) && request->frames > 0) { + int latent_frames = ((request->frames - 1) / 8) + 1; + sample_seq_len *= latent_frames; + } else if (sd_version_is_minimax_h3(sd->version) && request->frames > 0) { + sample_seq_len *= sd->video_frames_to_latent_frames(request->frames); + } + sigmas = sd->denoiser->get_sigmas(total_steps, + sample_seq_len, + scheduler, + sd->version, + sample_params->extra_sample_args); + } + + eta = resolve_eta(sd, eta, sample_method); + + if (high_noise_sample_steps < 0) { + for (size_t i = 0; i < sigmas.size(); ++i) { + if (sigmas[i] < moe_boundary) { + high_noise_sample_steps = static_cast(i); + break; + } + } + LOG_VERBOSE("switching from high noise model at step %d", high_noise_sample_steps); + } + + LOG_INFO("sampling using %s method", sampling_methods_str[sample_method]); + if (high_noise_sample_steps > 0) { + high_noise_sample_method = resolve_sample_method(sd, + high_noise_sample_method); + high_noise_eta = resolve_eta(sd, high_noise_eta, high_noise_sample_method); + LOG_INFO("sampling(high noise) using %s method", sampling_methods_str[high_noise_sample_method]); + } + } + + std::vector make_hires_sigma_schedule(StableDiffusionGGML* sd, + const sd_hires_params_t& hires, + const sd_sample_params_t& sample_params, + sample_method_t sample_method, + int default_steps, + int sample_seq_len, + int* scheduler_steps_out) { + if (scheduler_steps_out != nullptr) { + *scheduler_steps_out = 0; + } + + if (hires.custom_sigmas_count > 0 && hires.custom_sigmas != nullptr) { + std::vector custom_sigmas(hires.custom_sigmas, + hires.custom_sigmas + hires.custom_sigmas_count); + if (scheduler_steps_out != nullptr) { + *scheduler_steps_out = static_cast(custom_sigmas.size()) - 1; + } + return custom_sigmas; + } + + int effective_steps = hires.steps > 0 ? hires.steps : default_steps; + effective_steps = std::max(1, effective_steps); + + // sd-webui behavior: scale up total steps so trimming by denoising_strength yields exactly hires_steps effective steps, + // unlike img2img which trims from a fixed step count. + int scheduler_steps = static_cast(effective_steps / hires.denoising_strength); + scheduler_steps = std::max(1, scheduler_steps); + + scheduler_t scheduler = resolve_scheduler(sd, + sample_params.scheduler, + sample_method); + std::vector sigmas = sd->denoiser->get_sigmas(scheduler_steps, + sample_seq_len, + scheduler, + sd->version, + sample_params.extra_sample_args); + size_t t_enc = static_cast(scheduler_steps * hires.denoising_strength); + if (t_enc >= static_cast(scheduler_steps)) { + t_enc = static_cast(scheduler_steps) - 1; + } + if (scheduler_steps_out != nullptr) { + *scheduler_steps_out = scheduler_steps; + } + return std::vector(sigmas.begin() + scheduler_steps - static_cast(t_enc) - 1, + sigmas.end()); + } + +} // namespace sd::pipeline diff --git a/src/pipeline/request.h b/src/pipeline/request.h new file mode 100644 index 000000000..fac4aba55 --- /dev/null +++ b/src/pipeline/request.h @@ -0,0 +1,110 @@ +#ifndef __SD_PIPELINE_REQUEST_H__ +#define __SD_PIPELINE_REQUEST_H__ + +#include +#include + +#include "stable-diffusion.h" + +class StableDiffusionGGML; + +namespace sd::pipeline { + + extern const char* sampling_methods_str[]; + + enum sample_method_t default_sample_method(const StableDiffusionGGML* sd); + + enum scheduler_t default_scheduler(const StableDiffusionGGML* sd, enum sample_method_t sample_method); + + float resolve_eta(StableDiffusionGGML* sd, + float eta, + enum sample_method_t sample_method); + + struct GenerationRequest { + std::string prompt; + std::string negative_prompt; + int width = -1; + int height = -1; + int clip_skip = -1; + int vae_scale_factor = -1; + int diffusion_model_down_factor = -1; + int64_t seed = -1; + bool use_uncond = false; + bool use_img_uncond = false; + bool use_high_noise_uncond = false; + bool use_high_noise_img_uncond = false; + bool has_ref_images = false; + const sd_cache_params_t* cache_params = nullptr; + int batch_count = 1; + int qwen_image_layers = 3; + int shifted_timestep = 0; + float strength = 1.f; + float control_strength = 0.f; + float eta = 0.f; + sd_guidance_params_t guidance = {}; + sd_guidance_params_t high_noise_guidance = {}; + sd_pm_params_t pm_params = {}; + sd_pulid_params_t pulid_params = {}; + sd_hires_params_t hires = {}; + int frames = -1; + int requested_frames = -1; + int fps = 16; + float vace_strength = 1.f; + + GenerationRequest(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params); + + GenerationRequest(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params); + + void align_generation_request_size(); + + void align_image_size(int* target_width, int* target_height, const char* label); + + void resolve_hires(); + + static void resolve_guidance(StableDiffusionGGML* sd, + sd_guidance_params_t* guidance, + bool* use_uncond, + bool* use_img_uncond, + bool has_ref_images, + const char* stage_name = nullptr); + + void resolve(StableDiffusionGGML* sd); + }; + + struct SamplePlan { + enum sample_method_t sample_method = SAMPLE_METHOD_COUNT; + enum sample_method_t high_noise_sample_method = SAMPLE_METHOD_COUNT; + const char* extra_sample_args = nullptr; + const char* high_noise_extra_sample_args = nullptr; + float eta = 0.f; + float high_noise_eta = 0.f; + int sample_steps = 0; + int high_noise_sample_steps = 0; + int total_steps = 0; + float moe_boundary = 0.f; + std::vector sigmas; + + SamplePlan(StableDiffusionGGML* sd, + const sd_img_gen_params_t* sd_img_gen_params, + const GenerationRequest& request); + + SamplePlan(StableDiffusionGGML* sd, + const sd_vid_gen_params_t* sd_vid_gen_params, + const GenerationRequest& request); + + void resolve(StableDiffusionGGML* sd, + const GenerationRequest* request, + const sd_sample_params_t* sample_params); + }; + + std::vector make_hires_sigma_schedule(StableDiffusionGGML* sd, + const sd_hires_params_t& hires, + const sd_sample_params_t& sample_params, + sample_method_t sample_method, + int default_steps, + int sample_seq_len, + int* scheduler_steps_out); + +} // namespace sd::pipeline + +#endif // __SD_PIPELINE_REQUEST_H__ diff --git a/src/pipeline/video.cpp b/src/pipeline/video.cpp new file mode 100644 index 000000000..cc979c0fa --- /dev/null +++ b/src/pipeline/video.cpp @@ -0,0 +1,1797 @@ +#include "generation.h" + +#include +#include +#include +#include + +#include "core/rng.hpp" +#include "core/rng_philox.hpp" +#include "diffusion_engine.h" +#include "model/diffusion/minimax_h3.hpp" +#include "model/upscaler/ltx_latent_upscaler.hpp" +#include "model/vae/audio_vae.hpp" +#include "model/vae/vae.hpp" +#include "request.h" +#include "runtime/denoiser.hpp" + +namespace sd::pipeline { + + static sd_audio_t* waveform_to_sd_audio(const StableDiffusionGGML* sd, + const sd::Tensor& waveform) { + if (sd == nullptr || waveform.empty()) { + return nullptr; + } + + int64_t sample_count = waveform.shape()[0]; + int64_t channels = waveform.shape().size() > 1 ? waveform.shape()[1] : 1; + if (sample_count <= 0 || channels <= 0) { + return nullptr; + } + + sd_audio_t* audio = (sd_audio_t*)malloc(sizeof(sd_audio_t)); + if (audio == nullptr) { + return nullptr; + } + + audio->sample_rate = static_cast(sd->audio_vae_model != nullptr ? sd->audio_vae_model->output_sample_rate() : 0); + audio->channels = static_cast(channels); + audio->sample_count = static_cast(sample_count); + size_t sample_bytes = waveform.numel() * sizeof(float); + audio->data = (float*)malloc(sample_bytes); + if (audio->data == nullptr) { + free(audio); + return nullptr; + } + + auto wavaform_t = waveform.permute({1, 0, 2, 3}); + std::memcpy(audio->data, wavaform_t.data(), sample_bytes); + + return audio; + } + + static float ltxv_latent_corner_to_pixel_frame(int64_t corner_index, + int temporal_scale, + bool causal_temporal_positioning) { + float pixel_t = static_cast(corner_index * temporal_scale); + if (causal_temporal_positioning) { + pixel_t = std::max(0.f, pixel_t + 1.f - static_cast(temporal_scale)); + } + return pixel_t; + } + + static void set_ltxv_video_position(sd::Tensor* positions, + int64_t token, + float t_start, + float t_end, + float h_start, + float h_end, + float w_start, + float w_end) { + positions->index(0, 0, token, 0) = t_start; + positions->index(1, 0, token, 0) = t_end; + positions->index(0, 1, token, 0) = h_start; + positions->index(1, 1, token, 0) = h_end; + positions->index(0, 2, token, 0) = w_start; + positions->index(1, 2, token, 0) = w_end; + } + + static sd::Tensor build_ltxv_video_positions(int64_t width, + int64_t height, + int64_t target_latent_frames, + int64_t keyframe_latent_frames, + int keyframe_frame_idx, + int keyframe_pixel_frames, + int fps, + int spatial_scale, + int temporal_scale, + bool causal_temporal_positioning) { + GGML_ASSERT(width > 0 && height > 0 && target_latent_frames > 0); + GGML_ASSERT(keyframe_latent_frames > 0); + GGML_ASSERT(fps > 0); + + int64_t total_tokens = width * height * (target_latent_frames + keyframe_latent_frames); + sd::Tensor positions({2, 3, total_tokens, 1}); + int64_t token = 0; + + for (int64_t t = 0; t < target_latent_frames; t++) { + float t_start = ltxv_latent_corner_to_pixel_frame(t, temporal_scale, causal_temporal_positioning) / static_cast(fps); + float t_end = ltxv_latent_corner_to_pixel_frame(t + 1, temporal_scale, causal_temporal_positioning) / static_cast(fps); + for (int64_t h = 0; h < height; h++) { + float h_start = static_cast(h * spatial_scale); + float h_end = static_cast((h + 1) * spatial_scale); + for (int64_t w = 0; w < width; w++) { + float w_start = static_cast(w * spatial_scale); + float w_end = static_cast((w + 1) * spatial_scale); + set_ltxv_video_position(&positions, token++, t_start, t_end, h_start, h_end, w_start, w_end); + } + } + } + + for (int64_t t = 0; t < keyframe_latent_frames; t++) { + float t_start = static_cast(keyframe_frame_idx + t * temporal_scale); + float t_end = static_cast(keyframe_frame_idx + (t + 1) * temporal_scale); + if (keyframe_pixel_frames == 1) { + t_end = t_start + 1.f; + } + t_start /= static_cast(fps); + t_end /= static_cast(fps); + for (int64_t h = 0; h < height; h++) { + float h_start = static_cast(h * spatial_scale); + float h_end = static_cast((h + 1) * spatial_scale); + for (int64_t w = 0; w < width; w++) { + float w_start = static_cast(w * spatial_scale); + float w_end = static_cast((w + 1) * spatial_scale); + set_ltxv_video_position(&positions, token++, t_start, t_end, h_start, h_end, w_start, w_end); + } + } + } + + return positions; + } + + static sd::Tensor pack_ltxav_audio_and_video_latents(const sd::Tensor& video_latent, + const sd::Tensor& audio_latent) { + if (audio_latent.empty()) { + return video_latent; + } + + GGML_ASSERT(video_latent.dim() == 4 || video_latent.dim() == 5); + GGML_ASSERT(audio_latent.dim() == 3 || audio_latent.dim() == 4); + if (video_latent.dim() == 5) { + GGML_ASSERT(video_latent.shape()[4] == 1); + } + if (audio_latent.dim() == 4) { + GGML_ASSERT(audio_latent.shape()[3] == 1); + } + + int64_t width = video_latent.shape()[0]; + int64_t height = video_latent.shape()[1]; + int64_t frames = video_latent.shape()[2]; + int64_t video_ch = video_latent.shape()[3]; + int64_t spatial_size = width * height * frames; + int64_t audio_values = audio_latent.numel(); + int64_t extra_ch = (audio_values + spatial_size - 1) / spatial_size; + + std::vector packed_shape = video_latent.shape(); + packed_shape[3] = video_ch + extra_ch; + sd::Tensor packed = sd::zeros(packed_shape); + + std::copy_n(video_latent.data(), video_latent.numel(), packed.data()); + std::copy_n(audio_latent.data(), audio_latent.numel(), packed.data() + video_latent.numel()); + return packed; + } + + static sd::Tensor pack_ltxav_audio_and_video_denoise_mask(const sd::Tensor& video_mask, + const sd::Tensor& video_latent, + const sd::Tensor& audio_latent) { + if (video_mask.empty() || audio_latent.empty()) { + return video_mask; + } + + GGML_ASSERT(video_latent.dim() == 4 || video_latent.dim() == 5); + GGML_ASSERT(audio_latent.dim() == 3 || audio_latent.dim() == 4); + if (video_latent.dim() == 5) { + GGML_ASSERT(video_latent.shape()[4] == 1); + } + if (audio_latent.dim() == 4) { + GGML_ASSERT(audio_latent.shape()[3] == 1); + } + + int64_t width = video_latent.shape()[0]; + int64_t height = video_latent.shape()[1]; + int64_t frames = video_latent.shape()[2]; + int64_t video_ch = video_latent.shape()[3]; + int64_t spatial_size = width * height * frames; + int64_t audio_values = audio_latent.numel(); + int64_t extra_ch = (audio_values + spatial_size - 1) / spatial_size; + + GGML_ASSERT(video_mask.dim() == video_latent.dim()); + GGML_ASSERT(video_mask.shape()[0] == width); + GGML_ASSERT(video_mask.shape()[1] == height); + GGML_ASSERT(video_mask.shape()[2] == frames); + if (video_mask.dim() == 5) { + GGML_ASSERT(video_mask.shape()[4] == video_latent.shape()[4]); + } + + int64_t mask_ch = video_mask.shape()[3]; + if (mask_ch == video_ch + extra_ch) { + return video_mask; + } + GGML_ASSERT(mask_ch == 1 || mask_ch == video_ch); + + sd::Tensor video_mask_full = video_mask; + if (mask_ch == 1 && video_ch != 1) { + video_mask_full = video_mask * sd::Tensor::ones(video_latent.shape()); + } + + std::vector audio_mask_shape = video_latent.shape(); + audio_mask_shape[3] = extra_ch; + auto audio_mask = sd::Tensor::ones(audio_mask_shape); + return sd::ops::concat(video_mask_full, audio_mask, 3); + } + + static sd::Tensor make_ltxav_video_denoise_mask(const sd::Tensor& video_latent, float value = 1.f) { + if (video_latent.empty()) { + return {}; + } + return sd::full({video_latent.shape()[0], + video_latent.shape()[1], + video_latent.shape()[2], + 1, + 1}, + value); + } + + static sd::Tensor encode_ltxav_condition_image(StableDiffusionGGML* sd, + const sd::Tensor& image, + const char* name) { + if (sd == nullptr || image.empty()) { + return {}; + } + auto condition_image = image.reshape({image.shape()[0], + image.shape()[1], + 1, + image.shape()[2], + image.shape()[3]}); + auto condition_latent = sd->encode_first_stage(condition_image); + if (condition_latent.empty()) { + LOG_ERROR("failed to encode LTXAV %s image", name); + } + return condition_latent; + } + + static bool apply_ltxav_condition_by_latent_index(sd::Tensor* video_latent, + sd::Tensor* video_mask, + const sd::Tensor& condition_latent, + int64_t latent_idx, + const char* name, + float conditioned_mask) { + if (video_latent == nullptr || video_mask == nullptr || video_latent->empty() || video_mask->empty()) { + return false; + } + if (condition_latent.empty() || + condition_latent.shape()[0] != video_latent->shape()[0] || + condition_latent.shape()[1] != video_latent->shape()[1] || + condition_latent.shape()[3] != video_latent->shape()[3]) { + LOG_ERROR("invalid LTXAV %s condition latent shape", name); + return false; + } + int64_t latent_frames = video_latent->shape()[2]; + int64_t condition_frames = condition_latent.shape()[2]; + if (latent_idx < 0 || condition_frames <= 0 || latent_idx + condition_frames > latent_frames) { + LOG_ERROR("invalid LTXAV %s image latent range: start=%" PRId64 ", length=%" PRId64 ", latent_frames=%" PRId64, + name, + latent_idx, + condition_frames, + latent_frames); + return false; + } + + sd::ops::slice_assign(video_latent, 2, latent_idx, latent_idx + condition_frames, condition_latent); + sd::ops::fill_slice(video_mask, 2, latent_idx, latent_idx + condition_frames, conditioned_mask); + return true; + } + + static bool apply_ltxav_condition_image_by_latent_index(StableDiffusionGGML* sd, + const sd::Tensor& image, + sd::Tensor* video_latent, + sd::Tensor* video_mask, + int64_t latent_idx, + const char* name, + float strength) { + auto condition_latent = encode_ltxav_condition_image(sd, image, name); + return !condition_latent.empty() && + apply_ltxav_condition_by_latent_index(video_latent, + video_mask, + condition_latent, + latent_idx, + name, + 1.0f - std::clamp(strength, 0.f, 1.f)); + } + + static sd::Tensor unpack_ltxav_audio_latent(const sd::Tensor& packed_latent, + int audio_length, + int video_channels) { + if (packed_latent.empty() || audio_length <= 0) { + return {}; + } + + GGML_ASSERT(packed_latent.dim() == 4 || packed_latent.dim() == 5); + int64_t width = packed_latent.shape()[0]; + int64_t height = packed_latent.shape()[1]; + int64_t frames = packed_latent.shape()[2]; + int64_t total_channels = packed_latent.shape()[3]; + int64_t spatial_size = width * height * frames; + if (total_channels <= video_channels) { + return {}; + } + + constexpr int kLtxavAudioFrequencyBins = 16; + constexpr int kLtxavAudioChannels = 8; + int64_t required_values = static_cast(audio_length) * kLtxavAudioFrequencyBins * kLtxavAudioChannels; + int64_t packed_values = (total_channels - video_channels) * spatial_size; + if (packed_values < required_values) { + return {}; + } + + sd::Tensor audio_latent({kLtxavAudioFrequencyBins, audio_length, kLtxavAudioChannels, 1}); + const float* audio_src = packed_latent.data() + static_cast(video_channels) * static_cast(spatial_size); + std::copy_n(audio_src, static_cast(required_values), audio_latent.data()); + return audio_latent; + } + + static sd::Tensor make_ltxav_empty_audio_latent(int audio_length) { + if (audio_length <= 0) { + return {}; + } + constexpr int kLtxavAudioFrequencyBins = 16; + constexpr int kLtxavAudioChannels = 8; + return sd::zeros({kLtxavAudioFrequencyBins, audio_length, kLtxavAudioChannels, 1}); + } + + static sd::Tensor resize_ltxav_audio_latent(const sd::Tensor& audio_latent, + int target_audio_length) { + auto resized = make_ltxav_empty_audio_latent(target_audio_length); + if (resized.empty() || audio_latent.empty()) { + return resized; + } + GGML_ASSERT(audio_latent.dim() == 3 || audio_latent.dim() == 4); + int copy_length = std::min(static_cast(audio_latent.shape()[1]), target_audio_length); + if (copy_length > 0) { + auto copied = sd::ops::slice(audio_latent, 1, 0, copy_length); + sd::ops::slice_assign(&resized, 1, 0, copy_length, copied); + } + return resized; + } + + static int get_ltxav_num_audio_latents(int frames, int fps) { + GGML_ASSERT(frames > 0); + GGML_ASSERT(fps > 0); + constexpr float kSampleRate = 16000.0f; + constexpr float kMelHopLength = 160.0f; + constexpr float kAudioLatentDownsample = 4.0f; + constexpr float kLatentsPerSecond = kSampleRate / kMelHopLength / kAudioLatentDownsample; + return static_cast(std::ceil((static_cast(frames) / static_cast(fps)) * kLatentsPerSecond)); + } + + static int get_minimax_h3_num_audio_latents(int frames, int fps) { + GGML_ASSERT(frames > 0 && fps > 0); + return std::max(1, + static_cast(std::lround( + static_cast(frames) * 40.0 / fps))); + } + + static sd::Tensor make_minimax_h3_empty_audio_latent(int audio_length) { + if (audio_length <= 0) { + return {}; + } + return sd::zeros({audio_length, 2, 32, 1}); + } + + static sd::Tensor prepare_minimax_h3_reference_waveform(const sd_audio_t& audio, + int target_sample_rate = 32000) { + if (audio.data == nullptr || audio.sample_count == 0 || audio.channels == 0 || audio.sample_rate == 0) { + return {}; + } + uint64_t output_samples = static_cast(std::llround( + static_cast(audio.sample_count) * target_sample_rate / audio.sample_rate)); + output_samples = std::max(1, output_samples); + uint64_t padded_samples = (output_samples + 799) / 800 * 800; + // Keep stereo streams planar for the mono-per-stream audio encoder: + // [samples, 1, stereo, batch]. This avoids flattening interleaved L/R + // storage into alternating samples when the encoder folds streams into + // its batch dimension. + sd::Tensor waveform({static_cast(padded_samples), 1, 2, 1}); + + for (uint64_t i = 0; i < output_samples; ++i) { + long double source_pos = static_cast(i) * audio.sample_rate / target_sample_rate; + uint64_t source0 = std::min(static_cast(source_pos), audio.sample_count - 1); + uint64_t source1 = std::min(source0 + 1, audio.sample_count - 1); + float fraction = static_cast(source_pos - source0); + for (uint32_t channel = 0; channel < 2; ++channel) { + uint32_t source_channel = audio.channels == 1 ? 0 : std::min(channel, audio.channels - 1); + float a = audio.data[source0 * audio.channels + source_channel]; + float b = audio.data[source1 * audio.channels + source_channel]; + waveform.index(static_cast(i), 0, channel, 0) = + std::clamp(a + (b - a) * fraction, -1.f, 1.f); + } + } + return waveform; + } + + static sd::Tensor unpack_minimax_h3_audio_latent(const sd::Tensor& packed_latent, + int audio_length, + int video_channels) { + if (packed_latent.empty() || audio_length <= 0) { + return {}; + } + GGML_ASSERT(packed_latent.dim() == 4 || packed_latent.dim() == 5); + int64_t spatial_size = packed_latent.shape()[0] * packed_latent.shape()[1] * packed_latent.shape()[2]; + int64_t required = static_cast(audio_length) * 2 * 32; + int64_t available = (packed_latent.shape()[3] - video_channels) * spatial_size; + if (available < required) { + return {}; + } + sd::Tensor audio({audio_length, 2, 32, 1}); + const float* source = packed_latent.data() + + static_cast(video_channels) * static_cast(spatial_size); + std::copy_n(source, static_cast(required), audio.data()); + return audio; + } + + static std::optional prepare_video_generation_latents(StableDiffusionGGML* sd, + const sd_vid_gen_params_t* sd_vid_gen_params, + GenerationRequest* request) { + ImageGenerationLatents latents; + int64_t prepare_start_ms = ggml_time_ms(); + + sd::Tensor start_image; + sd::Tensor end_image; + + if (sd_vid_gen_params->init_image.data) { + start_image = sd_image_to_tensor(sd_vid_gen_params->init_image, request->width, request->height); + } + + if (sd_vid_gen_params->end_image.data) { + end_image = sd_image_to_tensor(sd_vid_gen_params->end_image, request->width, request->height); + } + + if (sd_version_is_minimax_h3(sd->version)) { + if (sd_vid_gen_params->ref_images_count < 0 || sd_vid_gen_params->ref_videos_count < 0 || + sd_vid_gen_params->ref_audios_count < 0 || + (sd_vid_gen_params->ref_images_count > 0 && sd_vid_gen_params->ref_images == nullptr) || + (sd_vid_gen_params->ref_videos_count > 0 && sd_vid_gen_params->ref_videos == nullptr) || + (sd_vid_gen_params->ref_audios_count > 0 && sd_vid_gen_params->ref_audios == nullptr)) { + LOG_ERROR("invalid MiniMax-H3 Ref2VA input arrays"); + return std::nullopt; + } + + latents.audio_length = get_minimax_h3_num_audio_latents(request->frames, + request->fps); + latents.audio_latent = make_minimax_h3_empty_audio_latent(latents.audio_length); + + bool has_references = sd_vid_gen_params->ref_images_count > 0 || + sd_vid_gen_params->ref_videos_count > 0 || + sd_vid_gen_params->ref_audios_count > 0; + if (has_references && (!start_image.empty() || !end_image.empty())) { + LOG_ERROR("MiniMax-H3 keyframes and Ref2VA references cannot be used together"); + return std::nullopt; + } + + if (sd_vid_gen_params->control_frames_size > 0) { + LOG_ERROR("MiniMax-H3 control_frames are not implemented"); + return std::nullopt; + } + + auto add_visual_noise = [&](sd::Tensor latent) { + auto condition_rng = std::make_shared(); + condition_rng->manual_seed(static_cast(request->seed)); + return latent * MiniMaxH3::VISUAL_COND_TIMESTEP + + sd::Tensor::randn_like(latent, condition_rng) * + (1.f - MiniMaxH3::VISUAL_COND_TIMESTEP); + }; + + auto add_keyframe = [&](const sd::Tensor& image, + int32_t frame_index, + const char* name) -> bool { + if (image.empty()) { + return true; + } + auto video_image = image.reshape({image.shape()[0], + image.shape()[1], + 1, + image.shape()[2], + image.shape()[3]}); + auto latent = sd->encode_first_stage(video_image); + if (latent.empty()) { + LOG_ERROR("failed to encode MiniMax-H3 %s keyframe", name); + return false; + } + latents.ref_images.push_back(image); + latents.ref_latents.push_back(add_visual_noise(std::move(latent))); + latents.keyframe_indices.push_back(frame_index); + return true; + }; + + auto resize_reference = [&](const sd::Tensor& image, + int width, + int height) { + return sd::ops::interpolate( + image, + std::vector{width, height, image.shape()[2], image.shape()[3]}); + }; + + auto encode_reference_audio = [&](const sd_audio_t& audio, + int32_t* audio_index) -> bool { + if (sd->audio_vae_model == nullptr) { + LOG_ERROR("MiniMax-H3 Ref2VA audio requires --audio-vae with encoder weights"); + return false; + } + auto waveform = prepare_minimax_h3_reference_waveform( + audio, + sd->audio_vae_model->input_sample_rate()); + if (waveform.empty()) { + LOG_ERROR("invalid MiniMax-H3 reference audio"); + return false; + } + auto encoded = sd->audio_vae_model->encode(sd->n_threads, waveform); + if (encoded.empty()) { + LOG_ERROR("failed to encode MiniMax-H3 reference audio"); + return false; + } + *audio_index = static_cast(latents.reference_audio_latents.size()); + latents.reference_audio_latents.push_back(std::move(encoded)); + return true; + }; + + if (has_references) { + LOG_INFO("MiniMax-H3 Ref2VA: %d image(s), %d video(s), %d audio clip(s)", + sd_vid_gen_params->ref_images_count, + sd_vid_gen_params->ref_videos_count, + sd_vid_gen_params->ref_audios_count); + + for (int i = 0; i < sd_vid_gen_params->ref_images_count; ++i) { + auto image = ensure_image_tensor_channels( + sd_image_to_tensor(sd_vid_gen_params->ref_images[i]), + 3); + if (image.empty()) { + LOG_ERROR("failed to load MiniMax-H3 reference image %d", i + 1); + return std::nullopt; + } + int source_w = static_cast(image.shape()[0]); + int source_h = static_cast(image.shape()[1]); + double source_area = static_cast(source_w) * source_h; + double target_area = static_cast(request->width) * request->height; + double scale = std::min(1.0, std::sqrt(target_area / source_area)); + int width = std::max(32, static_cast(std::round(source_w * scale / 32.f)) * 32); + int height = std::max(32, static_cast(std::round(source_h * scale / 32.f)) * 32); + image = resize_reference(image, width, height); + auto latent = sd->encode_first_stage(image); + if (latent.empty()) { + LOG_ERROR("failed to encode MiniMax-H3 reference image %d", i + 1); + return std::nullopt; + } + int32_t video_index = static_cast(latents.ref_latents.size()); + latents.ref_latents.push_back(add_visual_noise(std::move(latent))); + latents.minimax_reference_blocks.push_back({MiniMaxH3ReferenceKind::IMAGE, + video_index, + -1}); + MiniMaxH3PresentationItem item; + item.kind = MiniMaxH3PresentationKind::IMAGE; + item.frames.push_back(std::move(image)); + latents.minimax_presentation_refs.push_back(std::move(item)); + } + + for (int video_idx = 0; video_idx < sd_vid_gen_params->ref_videos_count; ++video_idx) { + const auto& reference = sd_vid_gen_params->ref_videos[video_idx]; + if (reference.frames == nullptr || reference.frame_count < 1) { + LOG_ERROR("invalid MiniMax-H3 reference video %d", video_idx + 1); + return std::nullopt; + } + int source_fps = reference.fps > 0 ? reference.fps : 24; + int normalized_frames = static_cast(std::lround( + static_cast(reference.frame_count) * 24.0 / source_fps)); + normalized_frames = std::min(normalized_frames, request->frames); + if (normalized_frames < 5) { + LOG_ERROR("MiniMax-H3 reference video %d needs at least 5 frames at 24 fps", + video_idx + 1); + return std::nullopt; + } + while (normalized_frames % 17 != 5) { + --normalized_frames; + } + + auto first = ensure_image_tensor_channels(sd_image_to_tensor(reference.frames[0]), 3); + if (first.empty()) { + LOG_ERROR("invalid first frame in MiniMax-H3 reference video %d", video_idx + 1); + return std::nullopt; + } + int source_w = static_cast(first.shape()[0]); + int source_h = static_cast(first.shape()[1]); + double ratio = static_cast(source_w) / source_h; + double nominal_w = ratio >= 1.0 ? 768.0 * ratio : 768.0; + double nominal_h = ratio >= 1.0 ? 768.0 : 768.0 / ratio; + if (nominal_w * nominal_h > 768.0 * 1344.0) { + double scale = std::sqrt((768.0 * 1344.0) / (nominal_w * nominal_h)); + nominal_w *= scale; + nominal_h *= scale; + } + int width = std::max(32, static_cast(std::round(nominal_w / 32.0)) * 32); + int height = std::max(32, static_cast(std::round(nominal_h / 32.0)) * 32); + if (source_w * source_h < width * height) { + width = std::max(32, static_cast(std::round(source_w / 32.0)) * 32); + height = std::max(32, static_cast(std::round(source_h / 32.0)) * 32); + } + + sd::Tensor video({width, height, normalized_frames, 3, 1}); + for (int frame = 0; frame < normalized_frames; ++frame) { + int source_index = std::min(reference.frame_count - 1, + static_cast(std::floor(frame * source_fps / 24.0))); + auto source = ensure_image_tensor_channels( + sd_image_to_tensor(reference.frames[source_index]), + 3); + if (source.empty()) { + LOG_ERROR("invalid frame %d in MiniMax-H3 reference video %d", + source_index + 1, + video_idx + 1); + return std::nullopt; + } + source = resize_reference(source, width, height); + sd::ops::slice_assign(&video, 2, frame, frame + 1, source.unsqueeze(2)); + } + auto video_latent = sd->encode_first_stage(video); + if (video_latent.empty()) { + LOG_ERROR("failed to encode MiniMax-H3 reference video %d", video_idx + 1); + return std::nullopt; + } + int32_t audio_index = -1; + bool has_audio = reference.audio.data != nullptr && reference.audio.sample_count > 0; + if (has_audio) { + if (!encode_reference_audio(reference.audio, &audio_index)) { + return std::nullopt; + } + MiniMaxH3PresentationItem audio_item; + audio_item.kind = MiniMaxH3PresentationKind::AUDIO; + latents.minimax_presentation_refs.push_back(std::move(audio_item)); + } + + MiniMaxH3PresentationItem video_item; + video_item.kind = MiniMaxH3PresentationKind::VIDEO; + for (int frame = 0; frame < normalized_frames; frame += 12) { + auto sampled = sd::ops::slice(video, 2, frame, frame + 1) + .reshape({width, height, 3, 1}); + video_item.frames.push_back(std::move(sampled)); + video_item.timestamps.push_back(frame / 24.f); + } + latents.minimax_presentation_refs.push_back(std::move(video_item)); + + int32_t video_index = static_cast(latents.ref_latents.size()); + latents.ref_latents.push_back(add_visual_noise(std::move(video_latent))); + latents.minimax_reference_blocks.push_back({has_audio ? MiniMaxH3ReferenceKind::VIDEO_AUDIO + : MiniMaxH3ReferenceKind::VIDEO, + video_index, + audio_index}); + } + + for (int audio_idx = 0; audio_idx < sd_vid_gen_params->ref_audios_count; ++audio_idx) { + int32_t encoded_index = -1; + if (!encode_reference_audio(sd_vid_gen_params->ref_audios[audio_idx], &encoded_index)) { + return std::nullopt; + } + MiniMaxH3PresentationItem item; + item.kind = MiniMaxH3PresentationKind::AUDIO; + latents.minimax_presentation_refs.push_back(std::move(item)); + latents.minimax_reference_blocks.push_back({MiniMaxH3ReferenceKind::AUDIO, + -1, + encoded_index}); + } + } + + if (!has_references && (!start_image.empty() || !end_image.empty())) { + LOG_INFO(!start_image.empty() && !end_image.empty() ? "MiniMax-H3 FL2VA" : !start_image.empty() ? "MiniMax-H3 I2VA" + : "MiniMax-H3 end-frame conditioning"); + } + if (!has_references && + (!add_keyframe(start_image, 0, "start") || + !add_keyframe(end_image, request->frames - 1, "end"))) { + return std::nullopt; + } + } + + if (sd_version_is_ltxav(sd->version)) { + latents.audio_length = get_ltxav_num_audio_latents(request->frames, request->fps); + latents.audio_latent = make_ltxav_empty_audio_latent(latents.audio_length); + } + + if (sd_version_is_ltxav(sd->version)) { + if (sd_vid_gen_params->control_frames_size > 0) { + LOG_ERROR("LTXAV control_frames are not implemented"); + return std::nullopt; + } + + if (!start_image.empty() || !end_image.empty()) { + if (!start_image.empty() && !end_image.empty()) { + LOG_INFO("FLF2V"); + } else if (!start_image.empty()) { + LOG_INFO("IMG2VID"); + } else { + LOG_INFO("END2VID"); + } + + int64_t t1 = ggml_time_ms(); + latents.init_latent = sd->generate_init_latent(request->width, request->height, request->frames, true); + + float conditioning_strength = std::clamp(request->strength, 0.f, 1.f); + float conditioned_mask = 1.0f - conditioning_strength; + latents.denoise_mask = make_ltxav_video_denoise_mask(latents.init_latent, 1.f); + + auto apply_video_condition_by_keyframe_index = [&](const sd::Tensor& keyframes, + int frame_idx, + const char* name) -> bool { + int64_t keyframe_frames = keyframes.shape()[2]; + if (keyframe_frames <= 0 || keyframes.shape()[0] != latents.init_latent.shape()[0] || + keyframes.shape()[1] != latents.init_latent.shape()[1] || + keyframes.shape()[3] != latents.init_latent.shape()[3]) { + LOG_ERROR("invalid LTXAV %s keyframe latent shape", name); + return false; + } + + latents.video_target_frame_count = latents.init_latent.shape()[2]; + latents.video_conditioning_frame_count = keyframe_frames; + latents.init_latent = sd::ops::concat(latents.init_latent, keyframes, 2); + + auto keyframe_mask = sd::full({keyframes.shape()[0], + keyframes.shape()[1], + keyframes.shape()[2], + 1, + 1}, + conditioned_mask); + latents.denoise_mask = sd::ops::concat(latents.denoise_mask, keyframe_mask, 2); + latents.video_positions = build_ltxv_video_positions(latents.init_latent.shape()[0], + latents.init_latent.shape()[1], + latents.video_target_frame_count, + keyframe_frames, + frame_idx, + 1, + request->fps, + request->vae_scale_factor, + 8, + true); + return true; + }; + + if (!start_image.empty()) { + if (!apply_ltxav_condition_image_by_latent_index(sd, + start_image, + &latents.init_latent, + &latents.denoise_mask, + 0, + "init", + conditioning_strength)) { + return std::nullopt; + } + } + + if (!end_image.empty()) { + auto end_image_latent = encode_ltxav_condition_image(sd, end_image, "end"); + if (end_image_latent.empty()) { + return std::nullopt; + } + + int frame_idx = request->frames - 1; + bool ok = frame_idx == 0 ? apply_ltxav_condition_by_latent_index(&latents.init_latent, + &latents.denoise_mask, + end_image_latent, + 0, + "end", + conditioned_mask) + : apply_video_condition_by_keyframe_index(end_image_latent, frame_idx, "end"); + if (!ok) { + return std::nullopt; + } + } + + int64_t t2 = ggml_time_ms(); + LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); + } + } + + if (sd_version_is_hunyuan_video(sd->version) && + (!start_image.empty() || !end_image.empty())) { + LOG_INFO("Hunyuan Video IMG2VID"); + + int64_t t1 = ggml_time_ms(); + auto concat_latent = sd->generate_init_latent(request->width, + request->height, + request->frames, + true); + auto encode_condition_frame = [&](const sd::Tensor& image, + int64_t latent_frame, + const char* name) -> bool { + auto encoded = sd->encode_first_stage(image.unsqueeze(2)); + if (encoded.empty()) { + LOG_ERROR("failed to encode Hunyuan Video %s conditioning frame", name); + return false; + } + if (encoded.dim() == 4) { + encoded.unsqueeze_(2); + } + if (encoded.dim() != 5 || + encoded.shape()[0] != concat_latent.shape()[0] || + encoded.shape()[1] != concat_latent.shape()[1] || + encoded.shape()[3] != concat_latent.shape()[3]) { + LOG_ERROR("invalid Hunyuan Video %s conditioning latent shape", name); + return false; + } + sd::ops::slice_assign(&concat_latent, + 2, + latent_frame, + latent_frame + 1, + sd::ops::slice(encoded, 2, 0, 1)); + return true; + }; + + if (!start_image.empty() && !encode_condition_frame(start_image, 0, "start")) { + return std::nullopt; + } + if (!end_image.empty() && + !encode_condition_frame(end_image, concat_latent.shape()[2] - 1, "end")) { + return std::nullopt; + } + + sd::Tensor concat_mask = sd::zeros({concat_latent.shape()[0], + concat_latent.shape()[1], + concat_latent.shape()[2], + 1, + 1}); + if (!start_image.empty()) { + sd::ops::fill_slice(&concat_mask, 2, 0, 1, 1.0f); + } + if (!end_image.empty()) { + sd::ops::fill_slice(&concat_mask, 2, concat_mask.shape()[2] - 1, concat_mask.shape()[2], 1.0f); + } + latents.concat_latent = sd::ops::concat(concat_latent, concat_mask, 3); + + int64_t t2 = ggml_time_ms(); + LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); + } + + if (sd->diffusion_model->get_desc() == "Wan2.1-I2V-14B" || + sd->diffusion_model->get_desc() == "Wan2.2-I2V-14B" || + sd->diffusion_model->get_desc() == "Wan2.1-I2V-1.3B" || + sd->diffusion_model->get_desc() == "Wan2.1-FLF2V-14B") { + LOG_INFO("IMG2VID"); + + if (sd->diffusion_model->get_desc() == "Wan2.1-I2V-14B" || + sd->diffusion_model->get_desc() == "Wan2.1-I2V-1.3B" || + sd->diffusion_model->get_desc() == "Wan2.1-FLF2V-14B") { + if (!start_image.empty()) { + auto clip_vision_output = sd->get_clip_vision_output(start_image, false, -2); + if (clip_vision_output.empty()) { + LOG_ERROR("failed to compute clip vision output for init image"); + return std::nullopt; + } + latents.clip_vision_output = std::move(clip_vision_output); + } else { + latents.clip_vision_output = sd->get_clip_vision_output(start_image, false, -2, true); + } + + if (sd->diffusion_model->get_desc() == "Wan2.1-FLF2V-14B") { + sd::Tensor end_image_clip_vision_output; + if (!end_image.empty()) { + end_image_clip_vision_output = sd->get_clip_vision_output(end_image, false, -2); + if (end_image_clip_vision_output.empty()) { + LOG_ERROR("failed to compute clip vision output for end image"); + return std::nullopt; + } + } else { + end_image_clip_vision_output = sd->get_clip_vision_output(end_image, false, -2, true); + } + latents.clip_vision_output = sd::ops::concat(latents.clip_vision_output, end_image_clip_vision_output, 1); + } + + int64_t t1 = ggml_time_ms(); + LOG_INFO("get_clip_vision_output completed, taking %" PRId64 " ms", t1 - prepare_start_ms); + } + + int64_t t1 = ggml_time_ms(); + sd::Tensor image = sd::full({request->width, request->height, request->frames, 3, 1}, 0.5f); + if (!start_image.empty()) { + sd::ops::slice_assign(&image, 2, 0, 1, start_image.unsqueeze(2)); + } + if (!end_image.empty()) { + sd::ops::slice_assign(&image, 2, request->frames - 1, request->frames, end_image.unsqueeze(2)); + } + + auto concat_latent = sd->encode_first_stage(image); // [b, c, t, h/vae_scale_factor, w/vae_scale_factor] + if (concat_latent.empty()) { + LOG_ERROR("failed to encode video conditioning frames"); + return std::nullopt; + } + latents.concat_latent = std::move(concat_latent); + + int64_t t2 = ggml_time_ms(); + LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); + + sd::Tensor concat_mask = sd::zeros({latents.concat_latent.shape()[0], + latents.concat_latent.shape()[1], + latents.concat_latent.shape()[2], + 4, + 1}); // [b, 4, t, h/vae_scale_factor, w/vae_scale_factor] + if (!start_image.empty()) { + sd::ops::fill_slice(&concat_mask, 2, 0, 1, 1.0f); + } + if (!end_image.empty()) { + auto last_channel = sd::ops::slice(concat_mask, 3, 3, 4); + sd::ops::fill_slice(&last_channel, 2, last_channel.shape()[2] - 1, last_channel.shape()[2], 1.0f); + sd::ops::slice_assign(&concat_mask, 3, 3, 4, last_channel); + } + latents.concat_latent = sd::ops::concat(concat_mask, latents.concat_latent, 3); // [b, 4+c, t, h/vae_scale_factor, w/vae_scale_factor] + } else if (sd->diffusion_model->get_desc() == "Wan2.2-TI2V-5B" && !start_image.empty()) { + LOG_INFO("IMG2VID"); + + int64_t t1 = ggml_time_ms(); + auto init_img = start_image.reshape({start_image.shape()[0], start_image.shape()[1], 1, start_image.shape()[2], 1}); + auto init_image_latent = sd->encode_first_stage(init_img); // [b, c, 1, h/vae_scale_factor, w/vae_scale_factor] + if (init_image_latent.empty()) { + LOG_ERROR("failed to encode init video frame"); + return std::nullopt; + } + + latents.init_latent = sd->generate_init_latent(request->width, request->height, request->frames, true); // [b, c, t, h/vae_scale_factor, w/vae_scale_factor] + sd::ops::slice_assign(&latents.init_latent, 2, 0, init_image_latent.shape()[2], init_image_latent); + + latents.denoise_mask = sd::full({latents.init_latent.shape()[0], latents.init_latent.shape()[1], latents.init_latent.shape()[2], 1, 1}, 1.f); + sd::ops::fill_slice(&latents.denoise_mask, 2, 0, init_image_latent.shape()[2], 0.0f); + + if (!end_image.empty()) { + auto end_img = end_image.reshape({end_image.shape()[0], end_image.shape()[1], 1, end_image.shape()[2], 1}); + auto end_image_latent = sd->encode_first_stage(end_img); // [b, c, 1, h/vae_scale_factor, w/vae_scale_factor] + if (end_image_latent.empty()) { + LOG_ERROR("failed to encode end video frame"); + return std::nullopt; + } + sd::ops::slice_assign(&latents.init_latent, 2, latents.init_latent.shape()[2] - 1, latents.init_latent.shape()[2], end_image_latent); + sd::ops::fill_slice(&latents.denoise_mask, 2, latents.init_latent.shape()[2] - 1, latents.init_latent.shape()[2], 0.0f); + } + + int64_t t2 = ggml_time_ms(); + LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); + } else if (sd_version_is_lingbot_video(sd->version) && !start_image.empty()) { + LOG_INFO("LingBot Video IMG2VID"); + + int64_t t1 = ggml_time_ms(); + auto init_img = start_image.reshape({start_image.shape()[0], start_image.shape()[1], 1, start_image.shape()[2], 1}); + auto init_image_latent = sd->encode_first_stage(init_img); + if (init_image_latent.empty()) { + LOG_ERROR("failed to encode init video frame"); + return std::nullopt; + } + + latents.init_latent = sd->generate_init_latent(request->width, request->height, request->frames, true); + sd::ops::slice_assign(&latents.init_latent, 2, 0, init_image_latent.shape()[2], init_image_latent); + + latents.denoise_mask = sd::full({latents.init_latent.shape()[0], latents.init_latent.shape()[1], latents.init_latent.shape()[2], 1, 1}, 1.f); + sd::ops::fill_slice(&latents.denoise_mask, 2, 0, init_image_latent.shape()[2], 0.0f); + + latents.ref_images.push_back(start_image); + + int64_t t2 = ggml_time_ms(); + LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); + } else if (sd->diffusion_model->get_desc() == "Wan2.1-VACE-1.3B" || + sd->diffusion_model->get_desc() == "Wan2.x-VACE-14B") { + LOG_INFO("VACE"); + int64_t t1 = ggml_time_ms(); + sd::Tensor ref_image_latent; + if (!start_image.empty()) { + auto ref_img = start_image.reshape({start_image.shape()[0], start_image.shape()[1], 1, start_image.shape()[2], 1}); + auto encoded_ref = sd->encode_first_stage(ref_img); // [b, c, 1, h/vae_scale_factor, w/vae_scale_factor] + if (encoded_ref.empty()) { + LOG_ERROR("failed to encode VACE reference image"); + return std::nullopt; + } + ref_image_latent = sd::ops::concat(encoded_ref, sd::zeros(encoded_ref.shape()), 3); // [b, 2*c, 1, h/vae_scale_factor, w/vae_scale_factor] + } + + sd::Tensor control_video = sd::full({request->width, request->height, request->frames, 3, 1}, 0.5f); + int64_t control_frame_count = std::min(request->frames, sd_vid_gen_params->control_frames_size); + for (int64_t i = 0; i < control_frame_count; ++i) { + auto control_frame = sd_image_to_tensor(sd_vid_gen_params->control_frames[i], request->width, request->height); + sd::ops::slice_assign(&control_video, 2, i, i + 1, control_frame.unsqueeze(2)); + } + + sd::Tensor mask = sd::full({request->width, request->height, request->frames, 1, 1}, 1.0f); + + control_video = control_video - 0.5f; + sd::Tensor inactive = control_video * (1.0f - mask) + 0.5f; + sd::Tensor reactive = control_video * mask + 0.5f; + + inactive = sd->encode_first_stage(inactive); // [b, c, t, h/vae_scale_factor, w/vae_scale_factor] + if (inactive.empty()) { + LOG_ERROR("failed to encode VACE inactive context"); + return std::nullopt; + } + + reactive = sd->encode_first_stage(reactive); // [b, c, t, h/vae_scale_factor, w/vae_scale_factor] + if (reactive.empty()) { + LOG_ERROR("failed to encode VACE reactive context"); + return std::nullopt; + } + + int64_t length = inactive.shape()[2]; + if (!ref_image_latent.empty()) { + length += 1; + request->frames = static_cast((length - 1) * 4 + 1); + latents.ref_image_num = 1; + } + auto vace_context = sd::ops::concat(inactive, reactive, 3); // [b, 2*c, t, h/vae_scale_factor, w/vae_scale_factor] + + mask = sd::full({request->width, request->height, inactive.shape()[2], 1, 1}, 1.0f); + auto mask_context = mask.reshape({request->vae_scale_factor, + inactive.shape()[0], + request->vae_scale_factor, + inactive.shape()[1], + inactive.shape()[2]}); // [t, h/vae_scale_factor, vae_scale_factor, w/vae_scale_factor, vae_scale_factor] + mask_context = mask_context.permute({1, 3, 4, 0, 2}) // [vae_scale_factor, vae_scale_factor, t, h/vae_scale_factor, w/vae_scale_factor] + .reshape({inactive.shape()[0], + inactive.shape()[1], + inactive.shape()[2], + request->vae_scale_factor * request->vae_scale_factor}); // [vae_scale_factor*vae_scale_factor, t, h/vae_scale_factor, w/vae_scale_factor] + + if (!ref_image_latent.empty()) { + vace_context = sd::ops::concat(ref_image_latent, vace_context, 2); // [b, 2*c, t+1, h/vae_scale_factor, w/vae_scale_factor] + auto mask_pad = sd::zeros({mask_context.shape()[0], + mask_context.shape()[1], + 1, + mask_context.shape()[3]}); // [vae_scale_factor*vae_scale_factor, 1, h/vae_scale_factor, w/vae_scale_factor] + mask_context = sd::ops::concat(mask_pad, mask_context, 2); // [vae_scale_factor*vae_scale_factor, t + 1, h/vae_scale_factor, w/vae_scale_factor] + } + + mask_context.unsqueeze_(mask_context.dim()); // [b, vae_scale_factor*vae_scale_factor, t + 1 or t, h/vae_scale_factor, w/vae_scale_factor] + + latents.vace_context = sd::ops::concat(vace_context, mask_context, 3); // [b, 2*c + vae_scale_factor*vae_scale_factor, t + 1 or t, h/vae_scale_factor, w/vae_scale_factor] + int64_t t2 = ggml_time_ms(); + LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); + } + + if (latents.init_latent.empty()) { + latents.init_latent = sd->generate_init_latent(request->width, request->height, request->frames, true); + } + + if ((sd_version_is_ltxav(sd->version) || sd_version_is_minimax_h3(sd->version)) && + !latents.audio_latent.empty()) { + if (!latents.denoise_mask.empty()) { + latents.denoise_mask = pack_ltxav_audio_and_video_denoise_mask(latents.denoise_mask, + latents.init_latent, + latents.audio_latent); + } + latents.init_latent = pack_ltxav_audio_and_video_latents(latents.init_latent, latents.audio_latent); + } + + return latents; + } + + static ImageGenerationEmbeds prepare_video_generation_embeds(StableDiffusionGGML* sd, + const sd_vid_gen_params_t* sd_vid_gen_params, + const GenerationRequest& request, + const ImageGenerationLatents& latents) { + ConditionerRunnerEndOnExit conditioner_runner_end{sd->cond_stage_model.get()}; + + ImageGenerationEmbeds embeds; + ConditionerParams condition_params; + condition_params.clip_skip = request.clip_skip; + condition_params.text = request.prompt; + condition_params.zero_out_masked = true; + condition_params.ref_images = &latents.ref_images; + condition_params.minimax_h3_references = &latents.minimax_presentation_refs; + if (sd_version_is_lingbot_video(sd->version) || sd_version_is_minimax_h3(sd->version)) { + condition_params.ref_image_params.vlm_resize_mode = RefImageResizeMode::AREA; + } + + int64_t prepare_start_ms = ggml_time_ms(); + embeds.cond = sd->cond_stage_model->get_learned_condition(sd->n_threads, + condition_params); + embeds.cond.c_concat = latents.concat_latent; + embeds.cond.c_vector = latents.clip_vision_output; + if (sd_version_is_minimax_h3(sd->version)) { + embeds.cond.c_ref_images = latents.ref_latents; + embeds.cond.c_ref_audios = latents.reference_audio_latents; + embeds.cond.c_reference_blocks = latents.minimax_reference_blocks; + if (!latents.keyframe_indices.empty()) { + embeds.cond.c_position_ids = sd::Tensor( + {static_cast(latents.keyframe_indices.size())}, + latents.keyframe_indices); + } + } + if (request.use_uncond) { + condition_params.text = request.negative_prompt; + embeds.uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, + condition_params); + embeds.uncond.c_concat = latents.concat_latent; + embeds.uncond.c_vector = latents.clip_vision_output; + if (sd_version_is_minimax_h3(sd->version)) { + embeds.uncond.c_ref_images = latents.ref_latents; + embeds.uncond.c_ref_audios = latents.reference_audio_latents; + embeds.uncond.c_reference_blocks = latents.minimax_reference_blocks; + embeds.uncond.c_position_ids = embeds.cond.c_position_ids; + } + } + + int64_t t1 = ggml_time_ms(); + LOG_INFO("get_learned_condition completed, taking %.2fs", (t1 - prepare_start_ms) * 1.0f / 1000); + + return embeds; + } + + static sd_image_t* decode_video_outputs(StableDiffusionGGML* sd, + const GenerationRequest& request, + const sd::Tensor& final_latent, + int* num_frames_out) { + if (final_latent.empty()) { + LOG_ERROR("no latent video to decode"); + return nullptr; + } + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling video decode"); + return nullptr; + } + sd::Tensor video_latent = final_latent; + if ((sd_version_is_ltxav(sd->version) || sd_version_is_minimax_h3(sd->version)) && + video_latent.shape()[3] > sd->get_latent_channel()) { + video_latent = sd::ops::slice(video_latent, 3, 0, sd->get_latent_channel()); + } + LOG_VERBOSE("decode_video_outputs latent %dx%dx%dx%d", + (int)video_latent.shape()[0], + (int)video_latent.shape()[1], + (int)video_latent.shape()[2], + (int)video_latent.shape()[3]); + // auto z = sd::load_tensor_from_file_as_tensor("ltx_vae_z.bin"); + int64_t t4 = ggml_time_ms(); + sd::Tensor vid = sd->decode_first_stage(video_latent, true); + int64_t t5 = ggml_time_ms(); + LOG_INFO("decode_first_stage completed, taking %.2fs", (t5 - t4) * 1.0f / 1000); + if (vid.empty()) { + LOG_ERROR("decode_first_stage failed for video"); + return nullptr; + } + LOG_VERBOSE("decode_video_outputs decoded %dx%dx%dx%d", + (int)vid.shape()[0], + (int)vid.shape()[1], + (int)vid.shape()[2], + (int)vid.shape()[3]); + if (request.frames > 0 && + vid.shape()[2] > request.frames) { + vid = sd::ops::slice(vid, 2, 0, request.frames); + } + + sd_image_t* result_images = (sd_image_t*)calloc(vid.shape()[2], sizeof(sd_image_t)); + if (result_images == nullptr) { + return nullptr; + } + if (num_frames_out != nullptr) { + *num_frames_out = static_cast(vid.shape()[2]); + } + + for (int64_t i = 0; i < vid.shape()[2]; i++) { + result_images[i] = tensor_to_sd_image(vid, static_cast(i)); + } + + return result_images; + } + + sd::Tensor upscale_ltx_spatial_video_latent(StableDiffusionGGML* sd, + const char* model_path, + const sd::Tensor& packed_latent, + int audio_length) { + if (sd == nullptr || sd->model_manager == nullptr || packed_latent.empty()) { + return {}; + } + if (strlen(SAFE_STR(model_path)) == 0) { + LOG_ERROR("LTX latent spatial upscale requires a model path"); + return {}; + } + if (!sd->ensure_backend_pair(SDBackendModule::UPSCALER)) { + return {}; + } + + int latent_channels = sd->get_latent_channel(); + sd::Tensor video_latent = packed_latent; + sd::Tensor audio_latent; + if (packed_latent.shape()[3] > latent_channels) { + video_latent = sd::ops::slice(packed_latent, 3, 0, latent_channels); + audio_latent = unpack_ltxav_audio_latent(packed_latent, audio_length, latent_channels); + } + + LOG_INFO("LTX latent spatial upscale: latent %dx%dx%dx%d -> model output", + (int)video_latent.shape()[0], + (int)video_latent.shape()[1], + (int)video_latent.shape()[2], + (int)video_latent.shape()[3]); + + sd::Tensor unnormalized = sd->un_normalize_ltx_video_latents(video_latent); + if (sd->first_stage_model) { + sd->first_stage_model->runner_end(); + } + if (unnormalized.empty()) { + LOG_ERROR("LTX latent un-normalization failed before spatial upscale"); + return {}; + } + + auto model_manager = sd->model_manager; + struct UpsamplerScope { + ModelManager& manager; + ModelLoader::FileId owned_source = 0; + std::unique_ptr runner; + std::vector params; + + ~UpsamplerScope() { + if (runner) { + runner->runner_end(); + } + GGML_ASSERT(manager.unregister_param_tensors(params)); + if (owned_source != 0) { + GGML_ASSERT(manager.del_file(owned_source)); + } + } + } scope{*model_manager}; + + const std::string prefix = "ltx_latent_upsampler"; + ModelLoader candidate = model_manager->loader(); + ModelLoader::FileId source_file = 0; + if (!candidate.add_file(model_path, prefix + ".", &source_file)) { + LOG_ERROR("init LTX latent upsampler model loader from file failed: '%s'", model_path); + return {}; + } + const bool owns_source = model_manager->loader().file_revision(source_file) == 0; + if (!model_manager->set_loader(std::move(candidate))) { + return {}; + } + scope.owned_source = owns_source ? source_file : 0; + + auto& upsampler = scope.runner; + upsampler = std::make_unique(sd->backend_for(SDBackendModule::UPSCALER), + model_manager->loader().get_tensor_storage_map(), + prefix, + model_manager); + const size_t max_graph_vram_bytes = sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER); + upsampler->set_max_graph_vram_bytes(max_graph_vram_bytes); + if (upsampler->model == nullptr) { + LOG_ERROR("init LTX latent upsampler from metadata failed"); + return {}; + } + + std::map tensors; + upsampler->get_param_tensors(tensors); + for (const auto& entry : tensors) { + scope.params.push_back(entry.second); + } + if (!model_manager->register_param_tensors(ModelComponent::LatentUpsampler, + std::move(tensors), + ModelManager::ResidencyMode::ParamBackend, + sd->backend_for(SDBackendModule::UPSCALER), + sd->params_backend_for(SDBackendModule::UPSCALER)) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("register LTX latent upsampler tensors with model manager failed"); + return {}; + } + + sd::Tensor upscaled = upsampler->compute(sd->n_threads, unnormalized); + upsampler->runner_end(); + if (upscaled.empty()) { + LOG_ERROR("LTX latent spatial upscale failed"); + return {}; + } + + upscaled = sd->normalize_ltx_video_latents(upscaled); + sd->first_stage_model->runner_end(); + if (upscaled.empty()) { + LOG_ERROR("LTX latent normalization failed after spatial upscale"); + return {}; + } + + if (!audio_latent.empty()) { + upscaled = pack_ltxav_audio_and_video_latents(upscaled, audio_latent); + } + return upscaled; + } + + static bool apply_ltxv_refine_image_conditioning(StableDiffusionGGML* sd, + const sd_vid_gen_params_t* sd_vid_gen_params, + const GenerationRequest& request, + const ImageGenerationLatents& latents, + sd::Tensor* latent, + sd::Tensor* denoise_mask, + sd::Tensor* video_positions) { + if (sd == nullptr || sd_vid_gen_params == nullptr || + latent == nullptr || latent->empty() || denoise_mask == nullptr || video_positions == nullptr) { + return true; + } + if (sd_vid_gen_params->init_image.data == nullptr && + sd_vid_gen_params->end_image.data == nullptr) { + return true; + } + constexpr float conditioning_strength = 1.f; + int latent_channels = sd->get_latent_channel(); + sd::Tensor video_latent = *latent; + sd::Tensor audio_latent; + if (latent->shape()[3] > latent_channels) { + video_latent = sd::ops::slice(*latent, 3, 0, latent_channels); + audio_latent = unpack_ltxav_audio_latent(*latent, latents.audio_length, latent_channels); + if (audio_latent.empty()) { + LOG_ERROR("failed to unpack LTXAV audio latent before image-to-video inplace conditioning"); + return false; + } + } + + int image_width = static_cast(video_latent.shape()[0]) * request.vae_scale_factor; + int image_height = static_cast(video_latent.shape()[1]) * request.vae_scale_factor; + sd::Tensor video_mask = make_ltxav_video_denoise_mask(video_latent, 1.f); + + if (sd_vid_gen_params->init_image.data != nullptr) { + sd::Tensor start_image = sd_image_to_tensor(sd_vid_gen_params->init_image, image_width, image_height); + if (!apply_ltxav_condition_image_by_latent_index(sd, + start_image, + &video_latent, + &video_mask, + 0, + "init", + conditioning_strength)) { + return false; + } + } + + if (sd_vid_gen_params->end_image.data != nullptr) { + sd::Tensor end_image = sd_image_to_tensor(sd_vid_gen_params->end_image, image_width, image_height); + sd::Tensor end_image_latent = encode_ltxav_condition_image(sd, end_image, "end"); + if (end_image_latent.empty()) { + return false; + } + + int frame_idx = request.frames - 1; + if (frame_idx == 0) { + if (!apply_ltxav_condition_by_latent_index(&video_latent, + &video_mask, + end_image_latent, + 0, + "end", + 1.f - conditioning_strength)) { + return false; + } + } else { + if (latents.video_conditioning_frame_count <= 0 || latents.video_target_frame_count <= 0) { + LOG_ERROR("LTXV FLF2V refine conditioning requires low-resolution keyframe conditioning metadata"); + return false; + } + int64_t target_latent_frames = latents.video_target_frame_count; + if (!apply_ltxav_condition_by_latent_index(&video_latent, + &video_mask, + end_image_latent, + target_latent_frames, + "end", + 1.f - conditioning_strength)) { + return false; + } + *video_positions = build_ltxv_video_positions(video_latent.shape()[0], + video_latent.shape()[1], + target_latent_frames, + end_image_latent.shape()[2], + frame_idx, + 1, + request.fps, + request.vae_scale_factor, + 8, + true); + } + } + + if (!audio_latent.empty()) { + *latent = pack_ltxav_audio_and_video_latents(video_latent, audio_latent); + *denoise_mask = pack_ltxav_audio_and_video_denoise_mask(video_mask, video_latent, audio_latent); + } else { + *latent = std::move(video_latent); + *denoise_mask = std::move(video_mask); + } + LOG_INFO("LTXV refine image conditioning applied at %dx%d", image_width, image_height); + return true; + } + + static bool generate_animatediff_video(StableDiffusionGGML* sd, + const sd_vid_gen_params_t* sd_vid_gen_params, + sd_image_t** frames_out, + int* num_frames_out) { + int n_frames = sd_vid_gen_params->video_frames; + if (n_frames < 1) { + LOG_ERROR("AnimateDiff: --video-frames must be >= 1"); + return false; + } + if (n_frames > 32) { + LOG_WARN("AnimateDiff motion modules have a 32-frame positional-encoding context; capping to 32"); + n_frames = 32; + } + + sd_img_gen_params_t img_gen_params; + sd_img_gen_params_init(&img_gen_params); + img_gen_params.loras = sd_vid_gen_params->loras; + img_gen_params.lora_count = sd_vid_gen_params->lora_count; + img_gen_params.prompt = sd_vid_gen_params->prompt; + img_gen_params.negative_prompt = sd_vid_gen_params->negative_prompt; + img_gen_params.clip_skip = sd_vid_gen_params->clip_skip; + img_gen_params.width = sd_vid_gen_params->width; + img_gen_params.height = sd_vid_gen_params->height; + img_gen_params.sample_params = sd_vid_gen_params->sample_params; + img_gen_params.strength = sd_vid_gen_params->strength; + img_gen_params.init_image = sd_vid_gen_params->init_image; + img_gen_params.seed = sd_vid_gen_params->seed; + img_gen_params.batch_count = 1; + img_gen_params.control_strength = 1.0f; + img_gen_params.vae_tiling_params = sd_vid_gen_params->vae_tiling_params; + img_gen_params.cache = sd_vid_gen_params->cache; + img_gen_params.hires = sd_vid_gen_params->hires; + img_gen_params.qwen_image_layers = 0; + img_gen_params.circular_x = sd_vid_gen_params->circular_x; + img_gen_params.circular_y = sd_vid_gen_params->circular_y; + + sd->animatediff_num_frames = n_frames; + bool ok = generate_image(sd, &img_gen_params, frames_out, num_frames_out); + sd->animatediff_num_frames = 0; + return ok; + } + + bool generate_video(StableDiffusionGGML* sd, + const sd_vid_gen_params_t* sd_vid_gen_params, + sd_image_t** frames_out, + int* num_frames_out, + sd_audio_t** audio_out) { + if (sd->config_->animatediff_loaded && sd_version_supports_animatediff(sd->version)) { + LOG_INFO("AnimateDiff dispatch: %d frames, %dx%d", + sd_vid_gen_params->video_frames, sd_vid_gen_params->width, sd_vid_gen_params->height); + return generate_animatediff_video(sd, sd_vid_gen_params, frames_out, num_frames_out); + } + + sd->reset_cancel_flag(); + + const RefImageParams ref_image_params; + + int64_t t0 = ggml_time_ms(); + sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params; + sd->apply_circular_axes(sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y); + GenerationRequest request(sd, sd_vid_gen_params); + bool latent_upscale_enabled = request.hires.enabled; + GenerationRequest hires_request = request; + if (latent_upscale_enabled) { + if (!sd_version_is_ltxav(sd->version)) { + LOG_ERROR("LTX latent spatial upscale is only supported for LTX video models"); + return false; + } + if (request.hires.upscaler != SD_HIRES_UPSCALER_MODEL) { + LOG_ERROR("LTX latent spatial upscale currently requires hires upscaler MODEL"); + return false; + } + if (strlen(SAFE_STR(request.hires.model_path)) == 0) { + LOG_ERROR("LTX latent spatial upscale is enabled but hires model path was not provided"); + return false; + } + } + + sd->rng->manual_seed(request.seed); + sd->sampler_rng->manual_seed(request.seed); + sd->set_flow_shift(sd_vid_gen_params->sample_params.flow_shift); + if (!sd->apply_loras(sd_vid_gen_params->loras, sd_vid_gen_params->lora_count)) + return false; + sd->reset_generation_extensions(); + + SamplePlan plan(sd, sd_vid_gen_params, request); + auto latent_inputs_opt = prepare_video_generation_latents(sd, sd_vid_gen_params, &request); + if (!latent_inputs_opt.has_value()) { + return false; + } + ImageGenerationLatents latents = std::move(*latent_inputs_opt); + + ImageGenerationEmbeds embeds = prepare_video_generation_embeds(sd, + sd_vid_gen_params, + request, + latents); + if (latent_upscale_enabled) { + LOG_INFO("generate_video %dx%dx%d -> LTX latent spatial upscale", + request.width, + request.height, + request.frames); + } else { + LOG_INFO("generate_video %dx%dx%d", + request.width, + request.height, + request.frames); + } + + int64_t latent_start = ggml_time_ms(); + int W = request.width / request.vae_scale_factor; + int H = request.height / request.vae_scale_factor; + int T = static_cast(latents.init_latent.shape()[2]); + + sd::Tensor x_t = latents.init_latent; + sd::Tensor noise = sd::Tensor::randn_like(x_t, sd->rng); + + if (plan.high_noise_sample_steps > 0) { + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before high-noise sampling"); + return false; + } + LOG_VERBOSE("sample(high noise) %dx%dx%d", W, H, T); + + int64_t sampling_start = ggml_time_ms(); + std::vector high_noise_sigmas(plan.sigmas.begin(), plan.sigmas.begin() + plan.high_noise_sample_steps + 1); + plan.sigmas = std::vector(plan.sigmas.begin() + plan.high_noise_sample_steps, plan.sigmas.end()); + + sd::Tensor x_t_sampled = sd->sample(sd->high_noise_diffusion_model, + false, + x_t, + std::move(noise), + embeds.cond, + request.use_high_noise_uncond ? embeds.uncond : SDCondition(), + embeds.img_uncond, + sd::Tensor(), + 0.f, + request.high_noise_guidance, + plan.high_noise_eta, + request.shifted_timestep, + plan.high_noise_sample_method, + sd->is_flow_denoiser(), + plan.high_noise_extra_sample_args, + high_noise_sigmas, + std::vector>{}, + ref_image_params, + latents.denoise_mask, + latents.vace_context, + request.vace_strength, + latents.audio_length, + static_cast(request.fps), + request.cache_params, + true, + latents.video_positions); + int64_t sampling_end = ggml_time_ms(); + if (x_t_sampled.empty()) { + LOG_ERROR("sampling(high noise) failed after %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); + return false; + } + + x_t = std::move(x_t_sampled); + noise = {}; + LOG_INFO("sampling(high noise) completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); + } + + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before sampling"); + return false; + } + LOG_VERBOSE("sample %dx%dx%d", W, H, T); + int64_t sampling_start = ggml_time_ms(); + sd::Tensor final_latent = sd->sample(sd->diffusion_model, + true, + x_t, + std::move(noise), + embeds.cond, + request.use_uncond ? embeds.uncond : SDCondition(), + embeds.img_uncond, + sd::Tensor(), + 0.f, + sd_vid_gen_params->sample_params.guidance, + plan.eta, + sd_vid_gen_params->sample_params.shifted_timestep, + plan.sample_method, + sd->is_flow_denoiser(), + plan.extra_sample_args, + plan.sigmas, + std::vector>{}, + ref_image_params, + latents.denoise_mask, + latents.vace_context, + request.vace_strength, + latents.audio_length, + static_cast(request.fps), + request.cache_params, + plan.high_noise_sample_steps <= 0, + latents.video_positions); + + int64_t sampling_end = ggml_time_ms(); + if (final_latent.empty()) { + LOG_ERROR("sampling failed after %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); + return false; + } + LOG_INFO("sampling completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); + + if (latent_upscale_enabled) { + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before latent upscale"); + return false; + } + int64_t upscale_start = ggml_time_ms(); + sd::Tensor upscaled_latent = upscale_ltx_spatial_video_latent(sd, + request.hires.model_path, + final_latent, + latents.audio_length); + int64_t upscale_end = ggml_time_ms(); + if (upscaled_latent.empty()) { + return false; + } + LOG_INFO("LTX latent spatial upscale completed, taking %.2fs", + (upscale_end - upscale_start) * 1.0f / 1000); + + x_t = std::move(upscaled_latent); + hires_request.width = static_cast(x_t.shape()[0]) * hires_request.vae_scale_factor; + hires_request.height = static_cast(x_t.shape()[1]) * hires_request.vae_scale_factor; + int upscaled_latent_frames = static_cast(x_t.shape()[2]); + int upscaled_frames = sd->latent_frames_to_video_frames(upscaled_latent_frames); + if (upscaled_frames != hires_request.frames) { + LOG_INFO("LTX latent upsampler output latent frames %d, frames %d -> %d", + upscaled_latent_frames, + hires_request.frames, + upscaled_frames); + hires_request.frames = upscaled_frames; + } + if (sd_version_is_ltxav(sd->version) && latents.audio_length > 0) { + int target_audio_length = get_ltxav_num_audio_latents(hires_request.frames, hires_request.fps); + if (target_audio_length != latents.audio_length) { + int latent_channels = sd->get_latent_channel(); + sd::Tensor video_latent = x_t; + sd::Tensor audio_latent = latents.audio_latent; + if (x_t.shape()[3] > latent_channels) { + video_latent = sd::ops::slice(x_t, 3, 0, latent_channels); + audio_latent = unpack_ltxav_audio_latent(x_t, latents.audio_length, latent_channels); + } + audio_latent = resize_ltxav_audio_latent(audio_latent, target_audio_length); + if (audio_latent.empty()) { + LOG_ERROR("failed to resize LTX audio latent for latent upscale: %d -> %d", + latents.audio_length, + target_audio_length); + return false; + } + x_t = pack_ltxav_audio_and_video_latents(video_latent, audio_latent); + latents.audio_latent = std::move(audio_latent); + LOG_INFO("LTX audio latent length adjusted for latent upscale: %d -> %d", + latents.audio_length, + target_audio_length); + latents.audio_length = target_audio_length; + } + } + if ((request.hires.target_width > 0 || request.hires.target_height > 0) && + (request.hires.target_width != hires_request.width || request.hires.target_height != hires_request.height)) { + LOG_WARN("LTX latent spatial upsampler output is %dx%d; ignoring hires target %dx%d", + hires_request.width, + hires_request.height, + request.hires.target_width, + request.hires.target_height); + } + sd::Tensor hires_denoise_mask; + sd::Tensor hires_video_positions; + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before latent upscale refine"); + return false; + } + if (!apply_ltxv_refine_image_conditioning(sd, + sd_vid_gen_params, + hires_request, + latents, + &x_t, + &hires_denoise_mask, + &hires_video_positions)) { + return false; + } + noise = sd::Tensor::randn_like(x_t, sd->rng); + + W = hires_request.width / hires_request.vae_scale_factor; + H = hires_request.height / hires_request.vae_scale_factor; + T = static_cast(x_t.shape()[2]); + sample_method_t hires_sample_method = plan.sample_method; + int hires_scheduler_steps = 0; + std::vector hires_sigma_sched = + make_hires_sigma_schedule(sd, + request.hires, + sd_vid_gen_params->sample_params, + hires_sample_method, + plan.sample_steps, + sd->get_image_seq_len(hires_request.height, hires_request.width) * T, + &hires_scheduler_steps); + float hires_eta = resolve_eta(sd, + sd_vid_gen_params->sample_params.eta, + hires_sample_method); + + LOG_VERBOSE("sample(latent upscale) %dx%dx%d", W, H, T); + LOG_INFO("LTX latent spatial upscale refine: scheduler_steps=%d, denoising_strength=%.2f, sampler=%s, sigma_sched_size=%zu%s", + hires_scheduler_steps, + request.hires.denoising_strength, + sampling_methods_str[hires_sample_method], + hires_sigma_sched.size(), + request.hires.custom_sigmas_count > 0 ? ", custom_sigmas=true" : ""); + + sampling_start = ggml_time_ms(); + final_latent = sd->sample(sd->diffusion_model, + true, + x_t, + std::move(noise), + embeds.cond, + hires_request.use_uncond ? embeds.uncond : SDCondition(), + embeds.img_uncond, + sd::Tensor(), + 0.f, + sd_vid_gen_params->sample_params.guidance, + hires_eta, + sd_vid_gen_params->sample_params.shifted_timestep, + hires_sample_method, + sd->is_flow_denoiser(), + plan.extra_sample_args, + hires_sigma_sched, + std::vector>{}, + ref_image_params, + hires_denoise_mask, + sd::Tensor(), + hires_request.vace_strength, + latents.audio_length, + static_cast(hires_request.fps), + hires_request.cache_params, + false, + hires_video_positions); + sampling_end = ggml_time_ms(); + if (final_latent.empty()) { + LOG_ERROR("sampling(latent upscale) failed after %.2fs", + (sampling_end - sampling_start) * 1.0f / 1000); + return false; + } + LOG_INFO("sampling(latent upscale) completed, taking %.2fs", + (sampling_end - sampling_start) * 1.0f / 1000); + } + + int64_t latent_end = ggml_time_ms(); + LOG_INFO("generating latent video completed, taking %.2fs", (latent_end - latent_start) * 1.0f / 1000); + + sd_audio_t* generated_audio = nullptr; + if ((sd_version_is_ltxav(sd->version) || sd_version_is_minimax_h3(sd->version)) && + latents.audio_length > 0 && + sd->audio_vae_model != nullptr) { + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before audio decode"); + return false; + } + int64_t audio_latent_decode_start = ggml_time_ms(); + + auto audio_latent = sd_version_is_minimax_h3(sd->version) + ? unpack_minimax_h3_audio_latent(final_latent, + latents.audio_length, + sd->get_latent_channel()) + : unpack_ltxav_audio_latent(final_latent, + latents.audio_length, + sd->get_latent_channel()); + if (!audio_latent.empty()) { + LOG_VERBOSE("decode audio latent %dx%dx%dx%d", + (int)audio_latent.shape()[0], + (int)audio_latent.shape()[1], + (int)audio_latent.shape()[2], + (int)audio_latent.shape()[3]); + auto waveform = sd->decode_ltx_audio_latent(audio_latent); + if (!waveform.empty()) { + generated_audio = waveform_to_sd_audio(sd, waveform); + } else { + LOG_WARN("audio latent decode failed; continuing with silent video output"); + } + } + int64_t audio_latent_decode_end = ggml_time_ms(); + LOG_INFO("decoding audio latent completed, taking %.2fs", (audio_latent_decode_end - audio_latent_decode_start) * 1.0f / 1000); + } + + if (latents.video_conditioning_frame_count > 0) { + int64_t target_frames = latents.video_target_frame_count > 0 ? latents.video_target_frame_count + : final_latent.shape()[2] - latents.video_conditioning_frame_count; + final_latent = sd::ops::slice(final_latent, 2, 0, target_frames); + } + + if (latents.ref_image_num > 0) { + final_latent = sd::ops::slice(final_latent, 2, latents.ref_image_num, final_latent.shape()[2]); + } + + if (sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before video decode"); + free_sd_audio(generated_audio); + return false; + } + auto result = decode_video_outputs(sd, latent_upscale_enabled ? hires_request : request, final_latent, num_frames_out); + if (result == nullptr) { + free_sd_audio(generated_audio); + return false; + } + + sd->lora_stat(); + + int64_t t1 = ggml_time_ms(); + LOG_INFO("generate_video completed in %.2fs", (t1 - t0) * 1.0f / 1000); + if (frames_out != nullptr) { + *frames_out = result; + } + if (audio_out != nullptr) { + *audio_out = generated_audio; + } else { + free_sd_audio(generated_audio); + } + return true; + } + +} // namespace sd::pipeline diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 442074e80..cda44e92b 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -1,170 +1,17 @@ +#include "stable-diffusion.h" + #include #include #include #include -#include -#include -#include -#include -#include -#include -#include -#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 "model.h" - -#include "core/rng.hpp" -#include "core/rng_mt19937.hpp" -#include "core/rng_philox.hpp" +#include "core/ggml_extend.h" #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 "core/backend_fit.h" -#include "extensions/generation_extension.h" -#include "model/adapter/ip_adapter.hpp" -#include "model/adapter/lora.hpp" -#include "model/diffusion/animatediff.hpp" -#include "model/diffusion/control.hpp" -#include "model/diffusion/minimax_h3.hpp" -#include "model/diffusion/model.hpp" -#include "model/upscaler/esrgan.hpp" -#include "model/upscaler/ltx_latent_upscaler.hpp" -#include "model/vae/audio_vae.hpp" -#include "model/vae/ltx_vae.hpp" -#include "model/vae/vae.hpp" -#include "runtime/denoiser.hpp" -#include "runtime/guidance.h" -#include "runtime/preview_interval.h" -#include "runtime/sample-cache.h" -#include "upscaler.h" - -#include "name_conversion.h" -#include "runtime/latent-preview.h" - -#include - -const char* sd_vae_format_name(enum sd_vae_format_t format); - -static bool sd_version_supports_animatediff(SDVersion version) { - return version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX; -} - -const char* model_version_to_str[] = { - "SD 1.x", - "SD 1.x Inpaint", - "Instruct-Pix2Pix", - "SD 1.x Tiny UNet", - "SD 2.x", - "SD 2.x Inpaint", - "SD 2.x Tiny UNet", - "SDXS (512-DS)", - "SDXS (09)", - "SDXL", - "SDXL Inpaint", - "SDXL Instruct-Pix2Pix", - "SDXL (Vega)", - "SDXL (SSD1B)", - "SVD", - "SD3.x", - "Flux", - "Flux Fill", - "Flux Control", - "Flex.2", - "Chroma Radiance", - "Wan 2.x", - "Wan 2.2 I2V", - "Wan 2.2 TI2V", - "LingBot Video", - "Qwen Image", - "Qwen Image Layered", - "Hunyuan Video", - "Anima", - "Flux.2", - "Flux.2 klein", - "LTXAV", - "MiniMax-H3", - "HiDream O1", - "Z-Image", - "Boogu Image", - "Ovis Image", - "Ernie Image", - "Lens", - "MiniT2I", - "Longcat-Image", - "PiD", - "Ideogram 4", - "SeFi-Image", - "Krea2", - "Mage Flow", - "ESRGAN", -}; - -const char* sampling_methods_str[] = { - "Euler", - "Euler A", - "Heun", - "DPM2", - "DPM++ (2s)", - "DPM++ (2M)", - "modified DPM++ (2M)", - "iPNDM", - "iPNDM_v", - "LCM", - "DDIM \"trailing\"", - "TCD", - "Res Multistep", - "Res 2s", - "ER-SDE", - "Euler CFG++", - "Euler A CFG++", - "Euler GE", - "DPM++ (2M) SDE", - "DPM++ (2M) SDE BT", - "LMS", -}; - -static_assert(SAMPLE_METHOD_COUNT == sizeof(sampling_methods_str) / sizeof(sampling_methods_str[0]), - "\nnumber of elements in sampling_methods_str[] != SAMPLE_METHOD_COUNT"); - -/*================================================== Helper Functions ================================================*/ +#include "pipeline/diffusion_engine.h" +#include "pipeline/generation.h" +#include "pipeline/request.h" -static bool sd_version_supports_ref_latent_img_cfg(SDVersion version) { - return version == VERSION_FLUX || - sd_version_is_flux2(version) || - sd_version_is_qwen_image(version) || - sd_version_is_mage_flow(version) || - sd_version_is_longcat(version) || - sd_version_is_z_image(version) || - sd_version_is_boogu_image(version); -} - -static bool sd_version_supports_img_cfg(SDVersion version, bool has_ref_images) { - return sd_version_is_inpaint_or_unet_edit(version) || - (has_ref_images && sd_version_supports_ref_latent_img_cfg(version)); -} - -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; - } -} +#define NONE_STR "NONE" static float get_cache_reuse_threshold(const sd_cache_params_t& params) { float reuse_threshold = params.reuse_threshold; @@ -178,6507 +25,697 @@ static float get_cache_reuse_threshold(const sd_cache_params_t& params) { return std::max(0.0f, reuse_threshold); } -/*=============================================== StableDiffusionGGML ================================================*/ - -template -struct has_set_runtime_backends : std::false_type {}; -template -struct has_set_runtime_backends().set_runtime_backends( - std::declval&>()))>> : std::true_type {}; - -static_assert(std::atomic::is_always_lock_free, - "sd_cancel_mode_t must be lock-free"); - -class StableDiffusionGGML { -public: - SDBackendManager backend_manager; - - SDVersion version; - bool external_vae_is_invalid = false; - - bool circular_x = false; - bool circular_y = false; - - std::shared_ptr rng = std::make_shared(); - std::shared_ptr sampler_rng = nullptr; - int n_threads = -1; - float default_flow_shift = INFINITY; - float active_flow_shift = INFINITY; - - std::shared_ptr cond_stage_model; - std::shared_ptr clip_vision; // for svd or wan2.1 i2v - std::shared_ptr diffusion_model; - std::shared_ptr high_noise_diffusion_model; - std::shared_ptr first_stage_model; - std::shared_ptr preview_vae; - std::shared_ptr audio_vae_model; - std::shared_ptr control_net; - std::shared_ptr ip_adapter; - sd::Tensor ip_adapter_tokens; - sd::Tensor ip_adapter_uncond_tokens; - float ip_adapter_strength = 1.0f; - std::vector> generation_extensions; - struct RuntimeLora { - ModelManager::LoraSpec spec; - SDBackendModule module; - std::shared_ptr model; - - bool matches(const ModelManager::LoraSpec& other) const { - return spec.file_id == other.file_id && spec.file_revision == other.file_revision && - spec.tensor_name_prefix_filter == other.tensor_name_prefix_filter; - } - }; - std::vector runtime_lora_models; - bool apply_lora_immediately = false; - int animatediff_num_frames = 0; - - std::string taesd_path; - sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr}; - bool enable_mmap = false; - sd::ggml_graph_cut::MaxVramAssignment max_vram_assignment; - bool disable_prefetch = false; - bool disable_segmented_compute = false; - bool eager_load = false; - std::string backend_spec; - std::string params_backend_spec; - std::string split_mode_spec; - bool auto_fit_enabled = false; - - bool diffusion_conv_direct = false; - - bool is_using_v_parameterization = false; - bool is_using_edm_v_parameterization = false; - - std::shared_ptr model_manager; - - enum class RunnerGroup { Core, - VAE, - ControlNet, - Extensions }; - using RunnerGroups = std::set; - - struct ModelConfig { - sd_ctx_params_t params{}; - std::list strings; - std::vector embeddings; - ModelLoader::FileId control_net_file = 0; - bool use_tae = false; - bool use_audio_vae = false; - bool photomaker_source_available = false; - bool animatediff_loaded = false; - - explicit ModelConfig(const sd_ctx_params_t& initial) - : params(initial) { - for (auto member : {&sd_ctx_params_t::model_path, &sd_ctx_params_t::clip_l_path, - &sd_ctx_params_t::clip_g_path, &sd_ctx_params_t::clip_vision_path, - &sd_ctx_params_t::t5xxl_path, &sd_ctx_params_t::llm_path, - &sd_ctx_params_t::llm_vision_path, &sd_ctx_params_t::diffusion_model_path, - &sd_ctx_params_t::high_noise_diffusion_model_path, &sd_ctx_params_t::uncond_diffusion_model_path, - &sd_ctx_params_t::embeddings_connectors_path, &sd_ctx_params_t::vae_path, - &sd_ctx_params_t::audio_vae_path, &sd_ctx_params_t::taesd_path, - &sd_ctx_params_t::control_net_path, &sd_ctx_params_t::ip_adapter_path, - &sd_ctx_params_t::motion_module_path, &sd_ctx_params_t::photo_maker_path, - &sd_ctx_params_t::pulid_weights_path, &sd_ctx_params_t::tensor_type_rules, - &sd_ctx_params_t::max_vram, &sd_ctx_params_t::backend, - &sd_ctx_params_t::params_backend, &sd_ctx_params_t::split_mode, - &sd_ctx_params_t::rpc_servers, &sd_ctx_params_t::model_args}) { - strings.emplace_back(SAFE_STR(initial.*member)); - params.*member = strings.back().c_str(); - } - for (uint32_t i = 0; i < initial.embedding_count; ++i) { - strings.emplace_back(SAFE_STR(initial.embeddings[i].name)); - const char* name = strings.back().c_str(); - strings.emplace_back(SAFE_STR(initial.embeddings[i].path)); - embeddings.push_back({name, strings.back().c_str()}); - } - params.embeddings = embeddings.data(); - } - - ModelConfig(const ModelConfig& other) - : ModelConfig(other.params) { - control_net_file = other.control_net_file; - use_tae = other.use_tae; - use_audio_vae = other.use_audio_vae; - photomaker_source_available = other.photomaker_source_available; - animatediff_loaded = other.animatediff_loaded; - } - ModelConfig& operator=(const ModelConfig&) = delete; - - void set_control_net(ModelLoader::FileId id, const std::string& path) { - control_net_file = id; - strings.push_back(path); - params.control_net_path = strings.back().c_str(); - } - }; - - struct RunnerState { - bool ready = false; - uint64_t catalog_revision = 0; - std::map sources; - }; - - std::recursive_mutex execution_mutex; - std::unique_ptr config_; - RunnerState runner_state_; - bool executing_ = false; - - std::shared_ptr denoiser = std::make_shared(); - std::vector file_alphas_cumprod; - - StableDiffusionGGML() = default; - ~StableDiffusionGGML() = default; - - static const std::map>& runner_components() { - static const std::map> components{ - {RunnerGroup::Core, {ModelComponent::Conditioner, ModelComponent::Diffusion, ModelComponent::HighNoiseDiffusion, ModelComponent::CLIPVision, ModelComponent::IPAdapter}}, - {RunnerGroup::VAE, {ModelComponent::VAE, ModelComponent::PreviewVAE, ModelComponent::AudioVAE}}, - {RunnerGroup::ControlNet, {ModelComponent::ControlNet}}, - {RunnerGroup::Extensions, {ModelComponent::PhotoMaker, ModelComponent::PuLID}}, - }; - return components; +const char* sd_type_name(enum sd_type_t type) { + if ((int)type < std::min(SD_TYPE_COUNT, GGML_TYPE_COUNT)) { + return ggml_type_name((ggml_type)type); } + return NONE_STR; +} - static RunnerGroups all_runner_groups() { - RunnerGroups groups; - for (const auto& entry : runner_components()) { - groups.insert(entry.first); +enum sd_type_t str_to_sd_type(const char* str) { + for (int i = 0; i < std::min(SD_TYPE_COUNT, GGML_TYPE_COUNT); i++) { + auto trait = ggml_get_type_traits((ggml_type)i); + if (!strcmp(str, trait->type_name)) { + return (enum sd_type_t)i; } - return groups; } + return SD_TYPE_COUNT; +} - ModelLoader::FileVersions 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; +const char* rng_type_to_str[] = { + "std_default", + "cuda", + "cpu", +}; + +const char* sd_rng_type_name(enum rng_type_t rng_type) { + if (rng_type < RNG_TYPE_COUNT) { + return rng_type_to_str[rng_type]; } + return NONE_STR; +} - void 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()); +enum rng_type_t str_to_rng_type(const char* str) { + for (int i = 0; i < RNG_TYPE_COUNT; i++) { + if (!strcmp(str, rng_type_to_str[i])) { + return (enum rng_type_t)i; } - state.ready = true; - runner_state_ = std::move(state); } + return RNG_TYPE_COUNT; +} - void 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(); - } +const char* sample_method_to_str[] = { + "euler", + "euler_a", + "heun", + "dpm2", + "dpm++2s_a", + "dpm++2m", + "dpm++2mv2", + "ipndm", + "ipndm_v", + "lcm", + "ddim_trailing", + "tcd", + "res_multistep", + "res_2s", + "er_sde", + "euler_cfg_pp", + "euler_a_cfg_pp", + "euler_ge", + "dpm++2m_sde", + "dpm++2m_sde_bt", + "lms", +}; - bool reset_runners(const RunnerGroups& groups) { - 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; - } +static_assert(SAMPLE_METHOD_COUNT == sizeof(sample_method_to_str) / sizeof(sample_method_to_str[0]), + "\nnumber of elements in sample_method_to_str[] != SAMPLE_METHOD_COUNT"); - bool 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)); +const char* sd_sample_method_name(enum sample_method_t sample_method) { + if (sample_method < SAMPLE_METHOD_COUNT) { + return sample_method_to_str[sample_method]; } + return NONE_STR; +} - bool apply_model_update(ModelLoader candidate, - std::unique_ptr next_config = nullptr, - 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); - } - } - } - 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; +enum sample_method_t str_to_sample_method(const char* str) { + for (int i = 0; i < SAMPLE_METHOD_COUNT; i++) { + if (!strcmp(str, sample_method_to_str[i])) { + return (enum sample_method_t)i; } - capture_runner_sources(); - return true; } + return SAMPLE_METHOD_COUNT; +} - struct ContextOperation { - StableDiffusionGGML& sd; - std::unique_lock lock; - bool acquired = false; - - explicit ContextOperation(StableDiffusionGGML& sd) - : sd(sd), lock(sd.execution_mutex, std::try_to_lock) { - if (!lock.owns_lock() || sd.executing_) { - // The caller may be a log callback, so rejecting it must not log. - return; - } - sd.executing_ = true; - acquired = true; - } +const char* scheduler_to_str[] = { + "discrete", + "karras", + "exponential", + "ays", + "gits", + "sgm_uniform", + "simple", + "smoothstep", + "kl_optimal", + "lcm", + "bong_tangent", + "ltx2", + "logit_normal", + "flux2", + "flux", + "beta", +}; - ~ContextOperation() { - if (acquired) { - sd.executing_ = false; - } - } - }; +static_assert(SCHEDULER_COUNT == sizeof(scheduler_to_str) / sizeof(scheduler_to_str[0]), + "\nnumber of elements in scheduler_to_str[] != SCHEDULER_COUNT"); - struct ExecutionScope { - ContextOperation operation; - bool ready = false; +const char* sd_scheduler_name(enum scheduler_t scheduler) { + if (scheduler < SCHEDULER_COUNT) { + return scheduler_to_str[scheduler]; + } + return NONE_STR; +} - explicit ExecutionScope(StableDiffusionGGML& sd) - : operation(sd) { - ready = operation.acquired && sd.refresh_model_sources(); +enum scheduler_t str_to_scheduler(const char* str) { + if (!strcmp(str, "normal")) { + return DISCRETE_SCHEDULER; + } + for (int i = 0; i < SCHEDULER_COUNT; i++) { + if (!strcmp(str, scheduler_to_str[i])) { + return (enum scheduler_t)i; } + } + return SCHEDULER_COUNT; +} - ~ExecutionScope() { - if (ready) { - operation.sd.end_runners(); - } - } - }; +const char* prediction_to_str[] = { + "eps", + "v", + "edm_v", + "sd3_flow", + "flux_flow", + "sefi_flow", + "minit2i_flow", +}; - ggml_backend_t 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; +const char* sd_prediction_name(enum prediction_t prediction) { + if (prediction < PREDICTION_COUNT) { + return prediction_to_str[prediction]; } + return NONE_STR; +} - ggml_backend_t 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)); +enum prediction_t str_to_prediction(const char* str) { + for (int i = 0; i < PREDICTION_COUNT; i++) { + if (!strcmp(str, prediction_to_str[i])) { + return (enum prediction_t)i; } - return module_backend; } + return PREDICTION_COUNT; +} - std::atomic cancellation_flag = SD_CANCEL_RESET; +const char* preview_to_str[] = { + "none", + "proj", + "tae", + "vae", +}; - void set_cancel_flag(enum sd_cancel_mode_t flag) { - cancellation_flag.store(flag, std::memory_order_release); +const char* sd_preview_name(enum preview_t preview) { + if (preview < PREVIEW_COUNT) { + return preview_to_str[preview]; } + return NONE_STR; +} - void reset_cancel_flag() { - set_cancel_flag(SD_CANCEL_RESET); +enum preview_t str_to_preview(const char* str) { + for (int i = 0; i < PREVIEW_COUNT; i++) { + if (!strcmp(str, preview_to_str[i])) { + return (enum preview_t)i; + } } + return PREVIEW_COUNT; +} - enum sd_cancel_mode_t get_cancel_flag() { - return cancellation_flag.load(std::memory_order_acquire); - } +const char* lora_apply_mode_to_str[] = { + "auto", + "immediately", + "at_runtime", +}; - size_t max_graph_vram_bytes_for_module(SDBackendModule module) { - return max_vram_assignment.bytes_for_backend(backend_for(module)); +const char* sd_lora_apply_mode_name(enum lora_apply_mode_t mode) { + if (mode < LORA_APPLY_MODE_COUNT) { + return lora_apply_mode_to_str[mode]; } + return NONE_STR; +} - std::vector layer_split_vram_limits_for_backends(const std::vector& backends) { - std::vector limits; - limits.reserve(backends.size()); - for (ggml_backend_t backend : backends) { - limits.push_back(max_vram_assignment.bytes_for_backend(backend)); +enum lora_apply_mode_t str_to_lora_apply_mode(const char* str) { + for (int i = 0; i < LORA_APPLY_MODE_COUNT; i++) { + if (!strcmp(str, lora_apply_mode_to_str[i])) { + return (enum lora_apply_mode_t)i; } - return limits; } + return LORA_APPLY_MODE_COUNT; +} - bool ensure_backend_pair(SDBackendModule module) { - if (backend_for(module) == nullptr) { - return false; - } - return params_backend_for(module) != nullptr; - } +const char* hires_upscaler_to_str[] = { + "None", + "Latent", + "Latent (nearest)", + "Latent (nearest-exact)", + "Latent (antialiased)", + "Latent (bicubic)", + "Latent (bicubic antialiased)", + "Lanczos", + "Nearest", + "Model", +}; - template - bool register_runner_params(ModelComponent component, - const std::shared_ptr& model, - SDBackendModule module, - size_t* params_mem_size = nullptr) { - if (model == nullptr) { - return true; - } - std::map group_tensors; - std::map tensor_ops; - model->get_param_tensors(group_tensors); - if constexpr (std::is_base_of_v) { - 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; +const char* sd_hires_upscaler_name(enum sd_hires_upscaler_t upscaler) { + if (upscaler >= SD_HIRES_UPSCALER_NONE && upscaler < SD_HIRES_UPSCALER_COUNT) { + return hires_upscaler_to_str[upscaler]; + } + return NONE_STR; +} - std::vector module_backends = backend_manager.runtime_backends(module); - if (module_backends.size() > 1) { - if constexpr (has_set_runtime_backends::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()); +enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str) { + for (int i = 0; i < SD_HIRES_UPSCALER_COUNT; i++) { + if (!strcmp(str, hires_upscaler_to_str[i])) { + return (enum sd_hires_upscaler_t)i; } - 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); } + return SD_HIRES_UPSCALER_COUNT; +} - template - bool register_row_split_runner_params(ModelComponent component, - const std::shared_ptr& model, - SDBackendModule module, - const std::vector& module_backends, - std::map group_tensors, - const std::map& 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 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)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> 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 split_tensors; - if constexpr (std::is_base_of_v) { - model->get_layer_split_param_tensors(split_tensors); - } else { - split_tensors = group_tensors; - } - - std::map row_split_map; - std::map 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); - } - - // Register graph-cut layer-split tensors on the primary backend first. - // The first real graph assigns each param tensor to a runtime backend - // before weights are loaded or staged. - template - bool register_layer_split_runner_params(ModelComponent component, - const std::shared_ptr& model, - SDBackendModule module, - const std::vector& module_backends, - std::map group_tensors, - const std::map& 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 =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 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(*config_); - next_config->set_control_net(0, ""); - return apply_model_update(std::move(candidate), std::move(next_config), {RunnerGroup::ControlNet}); - } - - bool 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(*config_); - next_config->set_control_net(file_id, path); - return apply_model_update(std::move(candidate), std::move(next_config), {RunnerGroup::ControlNet}); - } - - bool 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 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 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 get_rng(rng_type_t rng_type) { - if (rng_type == STD_DEFAULT_RNG) { - return std::make_shared(); - } else if (rng_type == CPU_RNG) { - return std::make_shared(); - } else { // default: CUDA_RNG - return std::make_shared(); - } - } - - void refresh_compvis_denoiser_sigmas() { - auto comp_vis_denoiser = std::dynamic_pointer_cast(denoiser); - if (!comp_vis_denoiser) { - return; - } - std::vector alphas_cumprod(TIMESTEPS); - if (file_alphas_cumprod.size() == TIMESTEPS) { - alphas_cumprod = file_alphas_cumprod; - } 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 load_alphas_cumprod() { - file_alphas_cumprod.clear(); - - std::vector 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 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; - if (strlen(SAFE_STR(sd_ctx_params->model_path)) > 0) { - LOG_INFO("loading model from '%s'", sd_ctx_params->model_path); - if (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(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 (!model_loader.init_from_file(sd_ctx_params->embeddings_connectors_path)) { - 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 (!model_loader.init_from_file(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->motion_module_path)) > 0) { - LOG_INFO("loading motion module (AnimateDiff) from '%s'", sd_ctx_params->motion_module_path); - if (!model_loader.init_from_file(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 (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0) { - if (!model_loader.init_from_file(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) { - 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 init(const sd_ctx_params_t* sd_ctx_params) { - auto configuration = std::make_unique(*sd_ctx_params); - n_threads = sd_ctx_params->n_threads; - 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 && backend_spec.empty() && 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(); - 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 ", model_version_to_str[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 uses_tae() const { - return config_->use_tae || version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09; - } - - bool tae_preview_only() const { - return config_->params.tae_preview_only && version != VERSION_SDXS_512_DS && version != VERSION_SDXS_09; - } - - void 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; - for (const auto& [_, tensor_storage] : model_loader.get_tensor_storage_map()) { - if (tensor_storage.is_int8_tensorwise) { - have_int8_tensorwise = true; - 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_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( - "INT8 tensorwise 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 model_build_context() { - return {config_->params, version, model_manager->loader().get_tensor_storage_map(), backend_manager, model_manager}; - } - - bool 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); - - 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)); - } - 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); - } - - bool 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 build_control_net_runner() { - if (config_->control_net_file == 0) { - return true; - } - 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 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 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(); - if (sd_ctx_params->flash_attn) { - LOG_INFO("Using flash attention"); - cond_stage_model->set_flash_attention_enabled(true); - if (clip_vision) { - clip_vision->set_flash_attention_enabled(true); - } - if (first_stage_model) { - first_stage_model->set_flash_attention_enabled(true); - } - if (preview_vae) { - preview_vae->set_flash_attention_enabled(true); - } - } - - if (sd_ctx_params->flash_attn || sd_ctx_params->diffusion_flash_attn) { - LOG_INFO("Using flash attention in the diffusion model"); - diffusion_model->set_flash_attention_enabled(true); - if (high_noise_diffusion_model) { - high_noise_diffusion_model->set_flash_attention_enabled(true); - } - } - LOG_VERBOSE("validating model metadata"); - - std::set ignore_tensors; - 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_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 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(); - default_flow_shift = INFINITY; - prediction_t pred_type = sd_ctx_params->prediction; - - if (pred_type == PREDICTION_COUNT) { - if (sd_version_is_sd2(version)) { - pred_type = is_using_v_parameterization_for_sd2(sd_version_is_inpaint(version)) ? V_PRED : EPS_PRED; - } else if (sd_version_is_sdxl(version)) { - if (tensor_storage_map.find("edm_vpred.sigma_max") != tensor_storage_map.end()) { - // CosXL models - // TODO: get sigma_min and sigma_max values from file - pred_type = EDM_V_PRED; - } else if (tensor_storage_map.find("v_pred") != tensor_storage_map.end()) { - pred_type = V_PRED; - } else { - pred_type = EPS_PRED; - } - } else if (sd_version_is_sd3(version) || - sd_version_is_wan(version) || - sd_version_is_hunyuan_video(version) || - sd_version_is_lingbot_video(version) || - sd_version_is_minimax_h3(version) || - sd_version_is_qwen_image(version) || - sd_version_is_mage_flow(version) || - version == VERSION_HIDREAM_O1 || - sd_version_is_anima(version) || - sd_version_is_ernie_image(version) || - sd_version_is_z_image(version) || - sd_version_is_boogu_image(version) || - sd_version_is_pid(version) || - sd_version_is_ideogram4(version)) { - pred_type = FLOW_PRED; - if (sd_version_is_wan(version)) { - default_flow_shift = 5.f; - } else if (sd_version_is_hunyuan_video(version)) { - default_flow_shift = 7.f; - } else if (sd_version_is_minimax_h3(version)) { - default_flow_shift = 12.f; - } else if (sd_version_is_ernie_image(version)) { - default_flow_shift = 4.f; - } else if (sd_version_is_pid(version)) { - default_flow_shift = 1.5f; - } else if (sd_version_is_ideogram4(version)) { - default_flow_shift = 1.0f; - } else if (sd_version_is_boogu_image(version)) { - default_flow_shift = 3.16f; - } else if (sd_version_is_mage_flow(version)) { - default_flow_shift = 6.f; - } else { - default_flow_shift = 3.f; - } - } else if (sd_version_is_flux(version) || - sd_version_is_flux2(version) || - sd_version_is_longcat(version) || - sd_version_is_lens(version) || - sd_version_is_ltxav(version) || - sd_version_is_krea2(version)) { - pred_type = FLUX_FLOW_PRED; - - default_flow_shift = 1.0f; // TODO: validate - for (const auto& [name, tensor_storage] : tensor_storage_map) { - if (starts_with(name, "model.diffusion_model.guidance_in.in_layer.weight")) { - default_flow_shift = 1.15f; - break; - } - } - if (sd_version_is_longcat(version)) { - default_flow_shift = 3.0f; - } else if (sd_version_is_lens(version)) { - default_flow_shift = 1.83f; - } else if (sd_version_is_ltxav(version)) { - default_flow_shift = 2.37f; - } else if (sd_version_is_krea2(version)) { - default_flow_shift = 1.15f; - } - } else if (sd_version_is_sefi_image(version)) { - pred_type = SEFI_FLOW_PRED; - } else if (sd_version_is_minit2i(version)) { - pred_type = MINIT2I_FLOW_PRED; - } else { - pred_type = EPS_PRED; - } - } - - 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(); - break; - case EDM_V_PRED: - LOG_INFO("running in v-prediction EDM mode"); - denoiser = std::make_shared(); - break; - case FLOW_PRED: { - if (sd_version_is_ltxav(version)) { - LOG_INFO("running in LTXAV FLOW mode"); - denoiser = std::make_shared(); - } else if (sd_version_is_minimax_h3(version)) { - LOG_INFO("running in MiniMax H3 AV FLOW mode"); - denoiser = std::make_shared(default_flow_shift, 3.f, get_latent_channel()); - } else { - LOG_INFO("running in FLOW mode"); - denoiser = std::make_shared(); - } - break; - } - case FLUX_FLOW_PRED: { - LOG_INFO("running in Flux FLOW mode"); - denoiser = std::make_shared(); - break; - } - case SEFI_FLOW_PRED: { - LOG_INFO("running in SeFi-Image dual-time FLOW mode"); - denoiser = std::make_shared(); - break; - } - case MINIT2I_FLOW_PRED: { - LOG_INFO("running in MiniT2I FLOW mode"); - denoiser = std::make_shared(); - break; - } - default: { - LOG_ERROR("Unknown predition type %i", pred_type); - return false; - } - } - - refresh_compvis_denoiser_sigmas(); - return true; - } - - bool build_runners(const RunnerGroups& groups) { - const auto& model_loader = model_manager->loader(); - std::map wtype_stat = model_loader.get_wtype_stat(); - std::map conditioner_wtype_stat = model_loader.get_conditioner_wtype_stat(); - std::map diffusion_model_wtype_stat = model_loader.get_diffusion_model_wtype_stat(); - std::map vae_wtype_stat = model_loader.get_vae_wtype_stat(); - - auto wtype_stat_to_str = [](const std::map& 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_VERBOSE("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 is_using_v_parameterization_for_sd2(bool is_inpaint = false) { - struct RunnerEndOnExit { - GGMLRunner* runner = nullptr; - ~RunnerEndOnExit() { - if (runner != nullptr) { - runner->runner_end(); - } - } - }; - RunnerEndOnExit diffusion_runner_end{diffusion_model.get()}; - - sd::Tensor x_t = sd::full({8, 8, 4, 1}, 0.5f); - sd::Tensor c = sd::full({1024, 2, 1, 1}, 0.5f); - sd::Tensor steps = sd::full({1}, 999.0f); - sd::Tensor concat; - if (is_inpaint) { - concat = sd::zeros({8, 8, 5, 1}); - } - - int64_t t0 = ggml_time_ms(); - sd::Tensor 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((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; - } - - static std::string lora_log_id(const ModelManager::LoraSpec& lora) { - return lora.is_high_noise ? "|high_noise|" + lora.path : lora.path; - } - - std::shared_ptr load_lora_model(const ModelManager::LoraSpec& lora_spec, - SDBackendModule module, - LoraModel::filter_t module_filter = nullptr) { - 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(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 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> load_runtime_loras_for_module(const std::vector& loras, - const std::set& model_tensor_names, - SDBackendModule module, - LoraModel::filter_t module_filter, - bool& success, - std::vector& next_models) { - std::vector> 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 apply_loras_immediately(const std::vector& 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 apply_loras_at_runtime(const std::vector& 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 next_models; - std::set 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(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(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(first_stage_lora_models); - first_stage_model->set_weight_adapter(multi_lora_adapter); - } - } - runtime_lora_models = std::move(next_models); - return success; - } - - void lora_stat() { - if (!runtime_lora_models.empty()) { - LOG_INFO("runtime_lora_models:"); - for (auto& lora_model : runtime_lora_models) { - lora_model.model->stat(); - } - } - } - - bool apply_loras(const sd_lora_t* loras, uint32_t lora_count) { - std::vector 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)) - return false; - 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) { - 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); - } - return true; - } - - void reset_generation_extensions() { - for (auto& extension : generation_extensions) { - extension->reset_runtime_condition(); - } - } - - void 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 get_clip_vision_output(const sd::Tensor& image, - bool return_pooled = true, - int clip_skip = -1, - bool zero_out_masked = false) { - sd::Tensor output; - if (zero_out_masked) { - if (return_pooled) { - output = sd::zeros({clip_vision->vision_model.projection_dim}); - } else { - output = sd::zeros({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; - } - - void 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::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 process_timesteps(const std::vector& timesteps, - const sd::Tensor& init_latent, - const sd::Tensor& denoise_mask, - int step) { - if (auto sefi_denoiser = std::dynamic_pointer_cast(denoiser)) { - int sched_idx = step > 0 ? step - 1 : 0; - if (sched_idx >= static_cast(sefi_denoiser->tex_timesteps.size())) { - sched_idx = static_cast(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(static_cast(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(frame)] = 0.f; - } - } - } - return new_timesteps; - } else { - return timesteps; - } - } - - std::vector process_ltxav_video_timesteps(const std::vector& timesteps, - const sd::Tensor& init_latent, - const sd::Tensor& 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 video_timesteps(static_cast(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 preview_image(int step, - const sd::Tensor& latents, - enum SDVersion version, - preview_t preview_mode, - std::function 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(latents.shape()[3]) : static_cast(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; - - 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 == 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(_latents.shape()[2]) : 1; - uint32_t img_width = static_cast(_latents.shape()[0]) * patch_sz; - uint32_t img_height = static_cast(_latents.shape()[1]) * patch_sz; - - uint8_t* data = (uint8_t*)malloc(frames * img_width * img_height * 3 * sizeof(uint8_t)); - GGML_ASSERT(data != nullptr); - preview_latent_video(data, _latents, latent_rgb_proj, latent_rgb_bias, patch_sz); - 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, 3, data + i * img_width * img_height * 3}; - } - 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 vae_latents; - sd::Tensor 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(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(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(preview_mode)); - } - } - - std::vector 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(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)shifted_t}; - } - if (sd_version_is_anima(version)) { - return std::vector{t / static_cast(TIMESTEPS)}; - } - if (sd_version_is_boogu_image(version)) { - return std::vector{t / static_cast(TIMESTEPS)}; - } - if (version == VERSION_HIDREAM_O1) { - return std::vector{1.0f - (t / static_cast(TIMESTEPS))}; - } - if (sd_version_is_z_image(version) || sd_version_is_ideogram4(version)) { - return std::vector{1000.f - t}; - } - return std::vector{t}; - } - - void adjust_sample_step_scalings(int shifted_timestep, - const std::vector& 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(roundf(timesteps_vec[0])); - float shifted_sigma = denoiser->t_to_sigma((float)shifted_t_idx); - std::vector 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; - } - - struct SamplePreviewContext { - sd_preview_cb_t callback = nullptr; - void* data = nullptr; - preview_t mode = PREVIEW_NONE; - }; - - SamplePreviewContext prepare_sample_preview_context() { - return SamplePreviewContext{sd_get_preview_callback(), - sd_get_preview_callback_data(), - sd_get_preview_mode()}; - } - - void 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; - } - } - } - - void compute_sample_controls(const sd::Tensor& control_image, - const sd::Tensor& noised_input, - const sd::Tensor& timesteps_tensor, - const SDCondition& condition, - std::vector>* 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 sample(const std::shared_ptr& work_diffusion_model, - bool inverse_noise_scaling, - const sd::Tensor& init_latent, - sd::Tensor noise, - const SDCondition& cond, - const SDCondition& uncond, - const SDCondition& img_uncond, - const sd::Tensor& 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& sigmas, - const std::vector>& ref_latents, - const RefImageParams& ref_image_params, - const sd::Tensor& denoise_mask, - const sd::Tensor& 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& video_positions = {}) { - struct RunnerEndOnExit { - GGMLRunner* runner = nullptr; - ~RunnerEndOnExit() { - if (runner != nullptr) { - runner->runner_end(); - } - } - }; - RunnerEndOnExit sample_diffusion_runner_end{work_diffusion_model.get()}; - - RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; - - std::vector 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 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(adaptive_projected_guidance) - : static_cast(classifier_free_guidance); - sd::guidance::SkipLayerGuidance skip_layer_guidance(has_skiplayer ? skip_layers : std::vector(), - 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(); - SamplePreviewContext preview = prepare_sample_preview_context(); - - sd::Tensor processed_init_latent = denoiser->process_latent_in(init_latent); - const sd::Tensor& sampling_init_latent = processed_init_latent.empty() - ? init_latent - : processed_init_latent; - sd::Tensor x_t = !noise.empty() - ? denoiser->noise_scaling(sigmas[0], noise, sampling_init_latent) - : sampling_init_latent; - sd::Tensor denoised = x_t; - - auto denoise = [&](const sd::Tensor& 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 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 base_timesteps_vec = prepare_sample_timesteps(sigma, shifted_timestep); - std::vector timesteps_vec = base_timesteps_vec; - sd::Tensor 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({static_cast(base_timesteps_vec.size())}, base_timesteps_vec); - } else { - timesteps_vec = process_timesteps(timesteps_vec, sampling_init_latent, denoise_mask, step); - } - const std::vector& 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 timesteps_tensor({static_cast(timesteps_vec.size())}, timesteps_vec); - sd::Tensor guidance_tensor({1}, std::vector{guidance.distilled_guidance}); - sd::Tensor hunyuan_timestep_r_tensor; - if (sd_version_is_hunyuan_video(version) && step + 1 < sigmas.size()) { - hunyuan_timestep_r_tensor = sd::Tensor::from_vector({sigmas[step + 1]}); - } - sd::Tensor noised_input = x * c_in; - if (!denoise_mask.empty() && (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 (!denoise_mask.empty()) { - 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 cond_out; - sd::Tensor uncond_out; - sd::Tensor img_uncond_out; - sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma); - std::vector> controls; - DiffusionParams diffusion_params; - diffusion_params.x = &noised_input; - diffusion_params.timesteps = ×teps_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); - - static const std::vector> 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* c_concat_override = nullptr, - const std::vector* local_skip_layers = nullptr, - const std::vector>* ref_latents_override = nullptr, - bool use_uncond_ip = false) -> sd::Tensor { - 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(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)) { - 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_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}; - } 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 (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, - audio_length, - std::isfinite(active_flow_shift) ? active_flow_shift : 12.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 { - diffusion_params.extra = std::monostate{}; - } - - sd::Tensor cached_output; - if (step_cache.before_condition(&condition, noised_input, &cached_output)) { - return std::move(cached_output); - } - - for (const auto& extension : generation_extensions) { - extension->before_diffusion(diffusion_params, step); - } - - auto output_opt = work_diffusion_model->compute(n_threads, diffusion_params); - if (output_opt.empty()) { - LOG_ERROR("diffusion model compute failed"); - return sd::Tensor(); - } - - step_cache.after_condition(&condition, noised_input, output_opt); - return output_opt; - }; - - const SDCondition* positive_condition = &cond; - const sd::Tensor* 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; - } - } - - cond_out = run_condition(*positive_condition, c_concat_override); - if (cond_out.empty()) { - return {}; - } - - if (!uncond.empty()) { - if (!step_cache.is_step_skipped()) { - compute_sample_controls(control_image, - noised_input, - timesteps_tensor, - uncond, - &controls); - } - const std::vector* uncond_skip_layers = nullptr; - if (is_skiplayer_step && slg_uncond) { - LOG_VERBOSE("Skipping layers at uncond step %d\n", step); - uncond_skip_layers = &skip_layer_guidance.layers(); - } - uncond_out = run_condition(uncond, - uncond.c_concat.empty() ? nullptr : &uncond.c_concat, - uncond_skip_layers, - nullptr, - true); - if (uncond_out.empty()) { - return {}; - } - } - if (!img_uncond.empty()) { - 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 {}; - } - } - 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 = guidance_schedule.empty() ? primary_guidance.forward(guidance_input, {}) : primary_guidance.forward(guidance_input, {}, guidance_schedule[guidance_schedule.size() - 1 - step]); - 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 { - 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& 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 (!denoise_mask.empty()) { - 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; - }; - - 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 get_vae_scale_factor() { - if (sd_version_is_pid(version)) { - return 1; - } - return first_stage_model->get_scale_factor(); - } - - int get_diffusion_model_down_factor() { - int down_factor = 8; // unet - if (sd_version_is_dit(version)) { - if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version)) { - down_factor = 2; - } else { - down_factor = 1; - } - } - return down_factor; - } - - int get_latent_channel() { - int latent_channel = 4; - if (sd_version_is_dit(version)) { - if (sd_version_is_ltxav(version)) { - latent_channel = 128; - } else if (sd_version_is_minimax_h3(version)) { - latent_channel = 24; - } else if (version == VERSION_WAN2_2_TI2V) { - latent_channel = 48; - } else if (sd_version_is_hunyuan_video(version)) { - latent_channel = 32; - } else if (version == VERSION_HIDREAM_O1) { - latent_channel = 3; - } else if (version == VERSION_CHROMA_RADIANCE) { - latent_channel = 3; - } else if (sd_version_is_minit2i(version)) { - latent_channel = 3; - } else if (sd_version_is_pid(version)) { - latent_channel = 3; - } else if (sd_version_is_sefi_image(version)) { - latent_channel = 144; - } else if (sd_version_uses_flux2_vae(version)) { - latent_channel = 128; - } else if (sd_version_is_mage_flow(version)) { - latent_channel = 128; - } else { - latent_channel = 16; - } - } - return latent_channel; - } - - int get_image_channels() const { - return version == VERSION_QWEN_IMAGE_LAYERED ? 4 : 3; - } - - int 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 generate_init_latent(int width, - int height, - int frames = 1, - bool video = false) { - 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({W, H, T, C, 1}); - } - return sd::zeros({W, H, C, 1}); - } - - int 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 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 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 encode_to_vae_latents(const sd::Tensor& 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 encode_first_stage(const sd::Tensor& 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 decode_first_stage(const sd::Tensor& x, bool decode_video = false) { - if (sd_version_is_pid(version) || sd_version_is_minit2i(version)) { - return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f); - } - auto latents = first_stage_model->diffusion_to_vae_latents(x); - auto decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); - const bool prefer_temporal_tiling = decode_video && first_stage_model->can_temporal_tile_decode(); - while (decoded.empty() && - auto_fit_enabled && - sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { - decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); - } - return decoded; - } - - sd::Tensor normalize_ltx_video_latents(const sd::Tensor& x) { - auto ltx_vae = std::dynamic_pointer_cast(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 un_normalize_ltx_video_latents(const sd::Tensor& x) { - auto ltx_vae = std::dynamic_pointer_cast(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 decode_ltx_audio_latent(const sd::Tensor& audio_latent) { - if (audio_vae_model == nullptr || audio_latent.empty()) { - return {}; - } - auto waveform = audio_vae_model->decode(n_threads, audio_latent); - return waveform; - } - - void set_flow_shift(float flow_shift = INFINITY) { - auto flow_denoiser = std::dynamic_pointer_cast(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 is_flow_denoiser() { - auto flow_denoiser = std::dynamic_pointer_cast(denoiser); - return !!flow_denoiser; - } - - std::string get_default_ref_image_preset(SDVersion version) const { - if (sd_version_is_longcat(version)) { - return "longcat"; - } else if (sd_version_is_flux(version)) { - return "flux_kontext"; - } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) { - return "flux2"; - } else if (version == VERSION_QWEN_IMAGE_LAYERED) { - return "qwen_layered"; - } else if (sd_version_is_qwen_image(version)) { - return "qwen"; - } else if (sd_version_is_mage_flow(version)) { - return "mage_flow"; - } else if (sd_version_is_z_image(version) || sd_version_is_boogu_image(version)) { - return "z_image_omni"; - } else if (sd_version_is_krea2(version)) { - // have to make a choice between "krea2_edit" mode (for lbouaraba/krea2edit) - // and "krea2_ostris_edit" (for krea2 ostris edit) - // since krea2 ostris edit support predates, it should probably be default - return "krea2_ostris_edit"; - } else if (sd_version_is_anima(version)) { - return "cosmos_reference"; - } - return "default"; - } - - RefImageParams resolve_ref_image_params(const char* ref_image_args) const { - RefImageParams params; - std::string preset_name = get_default_ref_image_preset(version); - - 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()); - } - break; - } - } - if (preset_name != "default") { - LOG_INFO("Using '%s' preset for reference images", preset_name.c_str()); - params = REF_IMAGE_PRESETS.at(preset_name); - } - - 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()); - } - } 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()); - } - } 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()); - } - } 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()); - } - } 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()); - } - } 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()); - } - } 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()); - } - } 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()); - } - } 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()); - } - } 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()); - } 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"); - } - return params; - } -}; - -/*================================================= SD API ==================================================*/ - -#define NONE_STR "NONE" - -const char* sd_type_name(enum sd_type_t type) { - if ((int)type < std::min(SD_TYPE_COUNT, GGML_TYPE_COUNT)) { - return ggml_type_name((ggml_type)type); - } - return NONE_STR; -} - -enum sd_type_t str_to_sd_type(const char* str) { - for (int i = 0; i < std::min(SD_TYPE_COUNT, GGML_TYPE_COUNT); i++) { - auto trait = ggml_get_type_traits((ggml_type)i); - if (!strcmp(str, trait->type_name)) { - return (enum sd_type_t)i; - } - } - return SD_TYPE_COUNT; -} - -const char* rng_type_to_str[] = { - "std_default", - "cuda", - "cpu", -}; - -const char* sd_rng_type_name(enum rng_type_t rng_type) { - if (rng_type < RNG_TYPE_COUNT) { - return rng_type_to_str[rng_type]; - } - return NONE_STR; -} - -enum rng_type_t str_to_rng_type(const char* str) { - for (int i = 0; i < RNG_TYPE_COUNT; i++) { - if (!strcmp(str, rng_type_to_str[i])) { - return (enum rng_type_t)i; - } - } - return RNG_TYPE_COUNT; -} - -const char* sample_method_to_str[] = { - "euler", - "euler_a", - "heun", - "dpm2", - "dpm++2s_a", - "dpm++2m", - "dpm++2mv2", - "ipndm", - "ipndm_v", - "lcm", - "ddim_trailing", - "tcd", - "res_multistep", - "res_2s", - "er_sde", - "euler_cfg_pp", - "euler_a_cfg_pp", - "euler_ge", - "dpm++2m_sde", - "dpm++2m_sde_bt", - "lms", -}; - -static_assert(SAMPLE_METHOD_COUNT == sizeof(sample_method_to_str) / sizeof(sample_method_to_str[0]), - "\nnumber of elements in sample_method_to_str[] != SAMPLE_METHOD_COUNT"); - -const char* sd_sample_method_name(enum sample_method_t sample_method) { - if (sample_method < SAMPLE_METHOD_COUNT) { - return sample_method_to_str[sample_method]; - } - return NONE_STR; -} - -enum sample_method_t str_to_sample_method(const char* str) { - for (int i = 0; i < SAMPLE_METHOD_COUNT; i++) { - if (!strcmp(str, sample_method_to_str[i])) { - return (enum sample_method_t)i; - } - } - return SAMPLE_METHOD_COUNT; -} - -const char* scheduler_to_str[] = { - "discrete", - "karras", - "exponential", - "ays", - "gits", - "sgm_uniform", - "simple", - "smoothstep", - "kl_optimal", - "lcm", - "bong_tangent", - "ltx2", - "logit_normal", - "flux2", - "flux", - "beta", -}; - -static_assert(SCHEDULER_COUNT == sizeof(scheduler_to_str) / sizeof(scheduler_to_str[0]), - "\nnumber of elements in scheduler_to_str[] != SCHEDULER_COUNT"); - -const char* sd_scheduler_name(enum scheduler_t scheduler) { - if (scheduler < SCHEDULER_COUNT) { - return scheduler_to_str[scheduler]; - } - return NONE_STR; -} - -enum scheduler_t str_to_scheduler(const char* str) { - if (!strcmp(str, "normal")) { - return DISCRETE_SCHEDULER; - } - for (int i = 0; i < SCHEDULER_COUNT; i++) { - if (!strcmp(str, scheduler_to_str[i])) { - return (enum scheduler_t)i; - } - } - return SCHEDULER_COUNT; -} - -const char* prediction_to_str[] = { - "eps", - "v", - "edm_v", - "sd3_flow", - "flux_flow", - "sefi_flow", - "minit2i_flow", -}; - -const char* sd_prediction_name(enum prediction_t prediction) { - if (prediction < PREDICTION_COUNT) { - return prediction_to_str[prediction]; - } - return NONE_STR; -} - -enum prediction_t str_to_prediction(const char* str) { - for (int i = 0; i < PREDICTION_COUNT; i++) { - if (!strcmp(str, prediction_to_str[i])) { - return (enum prediction_t)i; - } - } - return PREDICTION_COUNT; -} - -const char* preview_to_str[] = { - "none", - "proj", - "tae", - "vae", -}; - -const char* sd_preview_name(enum preview_t preview) { - if (preview < PREVIEW_COUNT) { - return preview_to_str[preview]; - } - return NONE_STR; -} - -enum preview_t str_to_preview(const char* str) { - for (int i = 0; i < PREVIEW_COUNT; i++) { - if (!strcmp(str, preview_to_str[i])) { - return (enum preview_t)i; - } - } - return PREVIEW_COUNT; -} - -const char* lora_apply_mode_to_str[] = { - "auto", - "immediately", - "at_runtime", -}; - -const char* sd_lora_apply_mode_name(enum lora_apply_mode_t mode) { - if (mode < LORA_APPLY_MODE_COUNT) { - return lora_apply_mode_to_str[mode]; - } - return NONE_STR; -} - -enum lora_apply_mode_t str_to_lora_apply_mode(const char* str) { - for (int i = 0; i < LORA_APPLY_MODE_COUNT; i++) { - if (!strcmp(str, lora_apply_mode_to_str[i])) { - return (enum lora_apply_mode_t)i; - } - } - return LORA_APPLY_MODE_COUNT; -} - -const char* hires_upscaler_to_str[] = { - "None", - "Latent", - "Latent (nearest)", - "Latent (nearest-exact)", - "Latent (antialiased)", - "Latent (bicubic)", - "Latent (bicubic antialiased)", - "Lanczos", - "Nearest", - "Model", -}; - -const char* sd_hires_upscaler_name(enum sd_hires_upscaler_t upscaler) { - if (upscaler >= SD_HIRES_UPSCALER_NONE && upscaler < SD_HIRES_UPSCALER_COUNT) { - return hires_upscaler_to_str[upscaler]; - } - return NONE_STR; -} - -enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str) { - for (int i = 0; i < SD_HIRES_UPSCALER_COUNT; i++) { - if (!strcmp(str, hires_upscaler_to_str[i])) { - return (enum sd_hires_upscaler_t)i; - } - } - return SD_HIRES_UPSCALER_COUNT; -} - -const char* sd_vae_format_name(enum sd_vae_format_t format) { - switch (format) { - case SD_VAE_FORMAT_AUTO: - return "auto"; - case SD_VAE_FORMAT_FLUX: - return "flux"; - case SD_VAE_FORMAT_SD3: - return "sd3"; - case SD_VAE_FORMAT_FLUX2: - return "flux2"; - case SD_VAE_FORMAT_WAN: - return "wan"; - default: - return NONE_STR; - } -} - -void sd_cache_params_init(sd_cache_params_t* cache_params) { - *cache_params = {}; - cache_params->mode = SD_CACHE_DISABLED; - cache_params->reuse_threshold = INFINITY; - cache_params->start_percent = 0.15f; - cache_params->end_percent = 0.95f; - cache_params->error_decay_rate = 1.0f; - cache_params->use_relative_threshold = true; - cache_params->reset_error_on_compute = true; - cache_params->Fn_compute_blocks = 8; - cache_params->Bn_compute_blocks = 0; - cache_params->residual_diff_threshold = 0.08f; - cache_params->max_warmup_steps = 8; - cache_params->max_cached_steps = -1; - cache_params->max_continuous_cached_steps = -1; - cache_params->taylorseer_n_derivatives = 1; - cache_params->taylorseer_skip_interval = 1; - cache_params->scm_mask = nullptr; - cache_params->scm_policy_dynamic = true; - cache_params->spectrum_w = 0.40f; - cache_params->spectrum_m = 3; - cache_params->spectrum_lam = 1.0f; - cache_params->spectrum_window_size = 2; - cache_params->spectrum_flex_window = 0.50f; - cache_params->spectrum_warmup_steps = 4; - cache_params->spectrum_stop_percent = 0.9f; -} - -void sd_hires_params_init(sd_hires_params_t* hires_params) { - *hires_params = {}; - hires_params->enabled = false; - hires_params->upscaler = SD_HIRES_UPSCALER_LATENT; - hires_params->model_path = nullptr; - hires_params->scale = 2.0f; - hires_params->target_width = 0; - hires_params->target_height = 0; - hires_params->steps = 0; - hires_params->denoising_strength = 0.7f; - hires_params->upscale_tile_size = 128; - hires_params->custom_sigmas = nullptr; - hires_params->custom_sigmas_count = 0; -} - -void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { - *sd_ctx_params = {}; - sd_ctx_params->n_threads = sd_get_num_physical_cores(); - sd_ctx_params->wtype = SD_TYPE_COUNT; - sd_ctx_params->rng_type = CUDA_RNG; - sd_ctx_params->sampler_rng_type = RNG_TYPE_COUNT; - sd_ctx_params->prediction = PREDICTION_COUNT; - sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO; - sd_ctx_params->max_vram = nullptr; - sd_ctx_params->disable_prefetch = false; - sd_ctx_params->disable_segmented_compute = false; - sd_ctx_params->eager_load = false; - sd_ctx_params->enable_mmap = false; - sd_ctx_params->diffusion_flash_attn = false; - sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO; - sd_ctx_params->backend = nullptr; - sd_ctx_params->params_backend = nullptr; - sd_ctx_params->split_mode = nullptr; - sd_ctx_params->auto_fit = true; - sd_ctx_params->rpc_servers = nullptr; - sd_ctx_params->model_args = nullptr; - sd_ctx_params->pulid_weights_path = nullptr; -} - -char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { - char* buf = (char*)malloc(8192); - if (!buf) - return nullptr; - buf[0] = '\0'; - - snprintf(buf + strlen(buf), 8192 - strlen(buf), - "model_path: %s\n" - "clip_l_path: %s\n" - "clip_g_path: %s\n" - "clip_vision_path: %s\n" - "t5xxl_path: %s\n" - "llm_path: %s\n" - "llm_vision_path: %s\n" - "diffusion_model_path: %s\n" - "high_noise_diffusion_model_path: %s\n" - "uncond_diffusion_model_path: %s\n" - "embeddings_connectors_path: %s\n" - "vae_path: %s\n" - "audio_vae_path: %s\n" - "taesd_path: %s\n" - "control_net_path: %s\n" - "photo_maker_path: %s\n" - "pulid_weights_path: %s\n" - "tensor_type_rules: %s\n" - "n_threads: %d\n" - "wtype: %s\n" - "rng_type: %s\n" - "sampler_rng_type: %s\n" - "prediction: %s\n" - "max_vram: %s\n" - "disable_prefetch: %s\n" - "disable_segmented_compute: %s\n" - "eager_load: %s\n" - "backend: %s\n" - "params_backend: %s\n" - "split_mode: %s\n" - "model_args: %s\n" - "auto_fit: %s\n" - "flash_attn: %s\n" - "diffusion_flash_attn: %s\n" - "vae_format: %s\n", - SAFE_STR(sd_ctx_params->model_path), - SAFE_STR(sd_ctx_params->clip_l_path), - SAFE_STR(sd_ctx_params->clip_g_path), - SAFE_STR(sd_ctx_params->clip_vision_path), - SAFE_STR(sd_ctx_params->t5xxl_path), - SAFE_STR(sd_ctx_params->llm_path), - SAFE_STR(sd_ctx_params->llm_vision_path), - SAFE_STR(sd_ctx_params->diffusion_model_path), - SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path), - SAFE_STR(sd_ctx_params->uncond_diffusion_model_path), - SAFE_STR(sd_ctx_params->embeddings_connectors_path), - SAFE_STR(sd_ctx_params->vae_path), - SAFE_STR(sd_ctx_params->audio_vae_path), - SAFE_STR(sd_ctx_params->taesd_path), - SAFE_STR(sd_ctx_params->control_net_path), - SAFE_STR(sd_ctx_params->photo_maker_path), - SAFE_STR(sd_ctx_params->pulid_weights_path), - SAFE_STR(sd_ctx_params->tensor_type_rules), - sd_ctx_params->n_threads, - sd_type_name(sd_ctx_params->wtype), - sd_rng_type_name(sd_ctx_params->rng_type), - sd_rng_type_name(sd_ctx_params->sampler_rng_type), - sd_prediction_name(sd_ctx_params->prediction), - SAFE_STR(sd_ctx_params->max_vram), - BOOL_STR(sd_ctx_params->disable_prefetch), - BOOL_STR(sd_ctx_params->disable_segmented_compute), - BOOL_STR(sd_ctx_params->eager_load), - SAFE_STR(sd_ctx_params->backend), - SAFE_STR(sd_ctx_params->params_backend), - SAFE_STR(sd_ctx_params->split_mode), - SAFE_STR(sd_ctx_params->model_args), - BOOL_STR(sd_ctx_params->auto_fit), - BOOL_STR(sd_ctx_params->flash_attn), - BOOL_STR(sd_ctx_params->diffusion_flash_attn), - sd_vae_format_name(sd_ctx_params->vae_format)); - - return buf; -} - -void sd_sample_params_init(sd_sample_params_t* sample_params) { - *sample_params = {}; - sample_params->guidance.txt_cfg = 7.0f; - sample_params->guidance.img_cfg = INFINITY; - sample_params->guidance.distilled_guidance = 3.5f; - sample_params->guidance.slg.layer_count = 0; - sample_params->guidance.slg.layer_start = 0.01f; - sample_params->guidance.slg.layer_end = 0.2f; - sample_params->guidance.slg.scale = 0.f; - sample_params->scheduler = SCHEDULER_COUNT; - sample_params->sample_method = SAMPLE_METHOD_COUNT; - sample_params->sample_steps = 20; - sample_params->eta = INFINITY; - sample_params->custom_sigmas = nullptr; - sample_params->custom_sigmas_count = 0; - sample_params->flow_shift = INFINITY; - sample_params->extra_sample_args = nullptr; -} - -char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) { - char* buf = (char*)malloc(4096); - if (!buf) - return nullptr; - buf[0] = '\0'; - - snprintf(buf + strlen(buf), 4096 - strlen(buf), - "(txt_cfg: %.2f, " - "img_cfg: %.2f, " - "distilled_guidance: %.2f, " - "slg.layer_count: %zu, " - "slg.layer_start: %.2f, " - "slg.layer_end: %.2f, " - "slg.scale: %.2f, " - "scheduler: %s, " - "sample_method: %s, " - "sample_steps: %d, " - "eta: %.2f, " - "shifted_timestep: %d, " - "flow_shift: %.2f, " - "extra_sample_args: %s)", - sample_params->guidance.txt_cfg, - std::isfinite(sample_params->guidance.img_cfg) - ? sample_params->guidance.img_cfg - : sample_params->guidance.txt_cfg, - sample_params->guidance.distilled_guidance, - sample_params->guidance.slg.layer_count, - sample_params->guidance.slg.layer_start, - sample_params->guidance.slg.layer_end, - sample_params->guidance.slg.scale, - sd_scheduler_name(sample_params->scheduler), - sd_sample_method_name(sample_params->sample_method), - sample_params->sample_steps, - sample_params->eta, - sample_params->shifted_timestep, - sample_params->flow_shift, - SAFE_STR(sample_params->extra_sample_args)); - - return buf; -} - -void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) { - *sd_img_gen_params = {}; - sd_sample_params_init(&sd_img_gen_params->sample_params); - sd_img_gen_params->clip_skip = -1; - sd_img_gen_params->ref_images_count = 0; - sd_img_gen_params->ref_image_args = ""; - sd_img_gen_params->width = 512; - sd_img_gen_params->height = 512; - sd_img_gen_params->strength = 0.75f; - sd_img_gen_params->seed = -1; - sd_img_gen_params->batch_count = 1; - sd_img_gen_params->control_strength = 0.9f; - sd_img_gen_params->ip_adapter_strength = 1.0f; - sd_img_gen_params->qwen_image_layers = 3; - sd_img_gen_params->circular_x = false; - sd_img_gen_params->circular_y = false; - sd_img_gen_params->pm_params = {nullptr, 0, nullptr, 20.f}; - sd_img_gen_params->pulid_params = {nullptr, 1.0f}; - sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr}; - sd_cache_params_init(&sd_img_gen_params->cache); - sd_hires_params_init(&sd_img_gen_params->hires); -} - -char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) { - char* buf = (char*)malloc(4096); - if (!buf) - return nullptr; - buf[0] = '\0'; - - char* sample_params_str = sd_sample_params_to_str(&sd_img_gen_params->sample_params); - - snprintf(buf + strlen(buf), 4096 - strlen(buf), - "prompt: %s\n" - "negative_prompt: %s\n" - "clip_skip: %d\n" - "width: %d\n" - "height: %d\n" - "sample_params: %s\n" - "strength: %.2f\n" - "seed: %" PRId64 - "\n" - "batch_count: %d\n" - "qwen_image_layers: %d\n" - "ref_images_count: %d\n" - "ref_image_args: %s\n" - "control_strength: %.2f\n" - "photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\n" - "VAE tiling: %s (temporal=%s, extra_tiling_args=%s)\n" - "circular_x: %s\n" - "circular_y: %s\n" - "hires: {enabled=%s, upscaler=%s, model_path=%s, scale=%.2f, target=%dx%d, steps=%d, denoising_strength=%.2f}\n", - SAFE_STR(sd_img_gen_params->prompt), - SAFE_STR(sd_img_gen_params->negative_prompt), - sd_img_gen_params->clip_skip, - sd_img_gen_params->width, - sd_img_gen_params->height, - SAFE_STR(sample_params_str), - sd_img_gen_params->strength, - sd_img_gen_params->seed, - sd_img_gen_params->batch_count, - sd_img_gen_params->qwen_image_layers, - sd_img_gen_params->ref_images_count, - SAFE_STR(sd_img_gen_params->ref_image_args), - sd_img_gen_params->control_strength, - sd_img_gen_params->pm_params.style_strength, - sd_img_gen_params->pm_params.id_images_count, - SAFE_STR(sd_img_gen_params->pm_params.id_embed_path), - BOOL_STR(sd_img_gen_params->vae_tiling_params.enabled), - BOOL_STR(sd_img_gen_params->vae_tiling_params.temporal_tiling), - SAFE_STR(sd_img_gen_params->vae_tiling_params.extra_tiling_args), - BOOL_STR(sd_img_gen_params->circular_x), - BOOL_STR(sd_img_gen_params->circular_y), - BOOL_STR(sd_img_gen_params->hires.enabled), - sd_hires_upscaler_name(sd_img_gen_params->hires.upscaler), - SAFE_STR(sd_img_gen_params->hires.model_path), - sd_img_gen_params->hires.scale, - sd_img_gen_params->hires.target_width, - sd_img_gen_params->hires.target_height, - sd_img_gen_params->hires.steps, - sd_img_gen_params->hires.denoising_strength); - const char* cache_mode_str = "disabled"; - if (sd_img_gen_params->cache.mode == SD_CACHE_EASYCACHE) { - cache_mode_str = "easycache"; - } else if (sd_img_gen_params->cache.mode == SD_CACHE_UCACHE) { - cache_mode_str = "ucache"; - } - snprintf(buf + strlen(buf), 4096 - strlen(buf), - "cache: %s (threshold=%.3f, start=%.2f, end=%.2f)\n", - cache_mode_str, - get_cache_reuse_threshold(sd_img_gen_params->cache), - sd_img_gen_params->cache.start_percent, - sd_img_gen_params->cache.end_percent); - free(sample_params_str); - return buf; -} - -void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params) { - *sd_vid_gen_params = {}; - sd_sample_params_init(&sd_vid_gen_params->sample_params); - sd_sample_params_init(&sd_vid_gen_params->high_noise_sample_params); - sd_vid_gen_params->high_noise_sample_params.sample_steps = -1; - sd_vid_gen_params->width = 512; - sd_vid_gen_params->height = 512; - sd_vid_gen_params->strength = 0.75f; - sd_vid_gen_params->seed = -1; - sd_vid_gen_params->video_frames = 6; - sd_vid_gen_params->fps = 16; - sd_vid_gen_params->moe_boundary = 0.875f; - sd_vid_gen_params->vace_strength = 1.f; - sd_vid_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr}; - sd_vid_gen_params->hires.enabled = false; - sd_vid_gen_params->hires.upscaler = SD_HIRES_UPSCALER_LATENT; - sd_vid_gen_params->hires.scale = 2.f; - sd_vid_gen_params->hires.target_width = 0; - sd_vid_gen_params->hires.target_height = 0; - sd_vid_gen_params->hires.steps = 0; - sd_vid_gen_params->hires.denoising_strength = 0.7f; - sd_vid_gen_params->hires.upscale_tile_size = 128; - sd_vid_gen_params->hires.custom_sigmas = nullptr; - sd_vid_gen_params->hires.custom_sigmas_count = 0; - sd_vid_gen_params->circular_x = false; - sd_vid_gen_params->circular_y = false; - sd_cache_params_init(&sd_vid_gen_params->cache); -} - -struct sd_ctx_t { - StableDiffusionGGML* sd = nullptr; -}; - -static bool sd_version_supports_video_generation(SDVersion version) { - return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version); -} - -static bool sd_version_supports_image_generation(SDVersion version) { - return !sd_version_supports_video_generation(version); -} - -sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) { - sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t)); - if (sd_ctx == nullptr) { - return nullptr; - } - - sd_ctx->sd = new StableDiffusionGGML(); - if (sd_ctx->sd == nullptr) { - free(sd_ctx); - return nullptr; - } - - if (!sd_ctx->sd->init(sd_ctx_params)) { - delete sd_ctx->sd; - sd_ctx->sd = nullptr; - free(sd_ctx); - return nullptr; - } - return sd_ctx; -} - -void free_sd_ctx(sd_ctx_t* sd_ctx) { - if (sd_ctx->sd != nullptr) { - delete sd_ctx->sd; - sd_ctx->sd = nullptr; - } - free(sd_ctx); -} - -SD_API void sd_cancel_generation(sd_ctx_t* sd_ctx, enum sd_cancel_mode_t mode) { - if (sd_ctx && sd_ctx->sd) { - if (mode < SD_CANCEL_ALL || mode > SD_CANCEL_RESET) { - mode = SD_CANCEL_ALL; - } - sd_ctx->sd->set_cancel_flag(mode); - } -} - -static sd_audio_t* waveform_to_sd_audio(const StableDiffusionGGML* sd, - const sd::Tensor& waveform) { - if (sd == nullptr || waveform.empty()) { - return nullptr; - } - - int64_t sample_count = waveform.shape()[0]; - int64_t channels = waveform.shape().size() > 1 ? waveform.shape()[1] : 1; - if (sample_count <= 0 || channels <= 0) { - return nullptr; - } - - sd_audio_t* audio = (sd_audio_t*)malloc(sizeof(sd_audio_t)); - if (audio == nullptr) { - return nullptr; - } - - audio->sample_rate = static_cast(sd->audio_vae_model != nullptr ? sd->audio_vae_model->output_sample_rate() : 0); - audio->channels = static_cast(channels); - audio->sample_count = static_cast(sample_count); - size_t sample_bytes = waveform.numel() * sizeof(float); - audio->data = (float*)malloc(sample_bytes); - if (audio->data == nullptr) { - free(audio); - return nullptr; - } - - auto wavaform_t = waveform.permute({1, 0, 2, 3}); - std::memcpy(audio->data, wavaform_t.data(), sample_bytes); - - return audio; -} - -void free_sd_audio(sd_audio_t* audio) { - if (audio == nullptr) { - return; - } - free(audio->data); - audio->data = nullptr; - free(audio); -} - -SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx) { - if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { - return false; - } - return sd_version_supports_image_generation(sd_ctx->sd->version); -} - -SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) { - if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { - return false; - } - if (sd_ctx->sd->config_->animatediff_loaded && sd_version_supports_animatediff(sd_ctx->sd->version)) { - return true; - } - return sd_version_supports_video_generation(sd_ctx->sd->version); -} - -SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path) { - if (sd_ctx == nullptr || sd_ctx->sd == nullptr || path == nullptr) { - return false; - } - return sd_ctx->sd->load_control_net_from_file(path); -} - -SD_API bool sd_ctx_unload_control_net(sd_ctx_t* sd_ctx) { - if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { - return false; - } - return sd_ctx->sd->unload_control_net(); -} - -SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx) { - if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { - return false; - } - return sd_ctx->sd->control_net != nullptr; -} - -enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) { - if (sd_ctx != nullptr && sd_ctx->sd != nullptr) { - if (sd_version_is_pid(sd_ctx->sd->version)) { - return LCM_SAMPLE_METHOD; - } - if (sd_version_is_dit(sd_ctx->sd->version)) { - return EULER_SAMPLE_METHOD; - } - } - return EULER_A_SAMPLE_METHOD; -} - -enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_method_t sample_method) { - if (sd_ctx != nullptr && sd_ctx->sd != nullptr) { - auto edm_v_denoiser = std::dynamic_pointer_cast(sd_ctx->sd->denoiser); - if (edm_v_denoiser) { - return EXPONENTIAL_SCHEDULER; - } - } - if (sample_method == LCM_SAMPLE_METHOD || sample_method == TCD_SAMPLE_METHOD) { - return LCM_SCHEDULER; - } else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) { - return SIMPLE_SCHEDULER; - } else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_flux(sd_ctx->sd->version)) { - return FLUX_SCHEDULER; - } else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_flux2(sd_ctx->sd->version)) { - return FLUX2_SCHEDULER; - } else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ltxav(sd_ctx->sd->version)) { - return LTX2_SCHEDULER; - } else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ideogram4(sd_ctx->sd->version)) { - return LOGIT_NORMAL_SCHEDULER; - } - return DISCRETE_SCHEDULER; -} - -static int64_t resolve_seed(int64_t seed) { - if (seed >= 0) { - return seed; - } - srand((int)time(nullptr)); - return rand(); -} - -static enum sample_method_t resolve_sample_method(sd_ctx_t* sd_ctx, enum sample_method_t sample_method) { - if (sample_method == SAMPLE_METHOD_COUNT) { - return sd_get_default_sample_method(sd_ctx); - } - return sample_method; -} - -static scheduler_t resolve_scheduler(sd_ctx_t* sd_ctx, - scheduler_t scheduler, - enum sample_method_t sample_method) { - if (scheduler == SCHEDULER_COUNT) { - return sd_get_default_scheduler(sd_ctx, sample_method); - } - return scheduler; -} - -static float resolve_eta(sd_ctx_t* sd_ctx, - float eta, - enum sample_method_t sample_method) { - if (eta == INFINITY) { - if (sd_ctx->sd->version == VERSION_HIDREAM_O1) { - return 8.f; - } - switch (sample_method) { - case DDIM_TRAILING_SAMPLE_METHOD: - case TCD_SAMPLE_METHOD: - case RES_MULTISTEP_SAMPLE_METHOD: - case RES_2S_SAMPLE_METHOD: - return 0.0f; - case EULER_A_SAMPLE_METHOD: - case DPMPP2S_A_SAMPLE_METHOD: - case ER_SDE_SAMPLE_METHOD: - case EULER_A_CFG_PP_SAMPLE_METHOD: - case DPMPP2M_SDE_SAMPLE_METHOD: - case DPMPP2M_SDE_BT_SAMPLE_METHOD: - return 1.0f; - default:; - } - return 0.0f; - } - return eta; -} - -struct GenerationRequest { - std::string prompt; - std::string negative_prompt; - int width = -1; - int height = -1; - int clip_skip = -1; - int vae_scale_factor = -1; - int diffusion_model_down_factor = -1; - int64_t seed = -1; - bool use_uncond = false; - bool use_img_uncond = false; - bool use_high_noise_uncond = false; - bool use_high_noise_img_uncond = false; - bool has_ref_images = false; - const sd_cache_params_t* cache_params = nullptr; - int batch_count = 1; - int qwen_image_layers = 3; - int shifted_timestep = 0; - float strength = 1.f; - float control_strength = 0.f; - float eta = 0.f; - sd_guidance_params_t guidance = {}; - sd_guidance_params_t high_noise_guidance = {}; - sd_pm_params_t pm_params = {}; - sd_pulid_params_t pulid_params = {}; - sd_hires_params_t hires = {}; - int frames = -1; - int requested_frames = -1; - int fps = 16; - float vace_strength = 1.f; - - GenerationRequest(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params) { - prompt = SAFE_STR(sd_img_gen_params->prompt); - negative_prompt = SAFE_STR(sd_img_gen_params->negative_prompt); - width = sd_img_gen_params->width; - height = sd_img_gen_params->height; - vae_scale_factor = sd_ctx->sd->get_vae_scale_factor(); - diffusion_model_down_factor = sd_ctx->sd->get_diffusion_model_down_factor(); - seed = sd_img_gen_params->seed; - batch_count = sd_img_gen_params->batch_count; - qwen_image_layers = std::max(0, sd_img_gen_params->qwen_image_layers); - clip_skip = sd_img_gen_params->clip_skip; - shifted_timestep = sd_img_gen_params->sample_params.shifted_timestep; - strength = sd_img_gen_params->strength; - control_strength = sd_img_gen_params->control_strength; - eta = sd_img_gen_params->sample_params.eta; - has_ref_images = sd_img_gen_params->ref_images_count > 0; - guidance = sd_img_gen_params->sample_params.guidance; - pm_params = sd_img_gen_params->pm_params; - pulid_params = sd_img_gen_params->pulid_params; - hires = sd_img_gen_params->hires; - cache_params = &sd_img_gen_params->cache; - resolve(sd_ctx); - } - - GenerationRequest(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params) { - prompt = SAFE_STR(sd_vid_gen_params->prompt); - negative_prompt = SAFE_STR(sd_vid_gen_params->negative_prompt); - width = sd_vid_gen_params->width; - height = sd_vid_gen_params->height; - requested_frames = std::max(1, sd_vid_gen_params->video_frames); - frames = sd_ctx->sd->align_video_frames(requested_frames); - clip_skip = sd_vid_gen_params->clip_skip; - fps = std::max(1, sd_vid_gen_params->fps); - if (sd_version_is_minimax_h3(sd_ctx->sd->version) && fps != 24) { - LOG_WARN("MiniMax-H3 uses 24 fps; overriding requested fps %d", fps); - fps = 24; - } - vae_scale_factor = sd_ctx->sd->get_vae_scale_factor(); - diffusion_model_down_factor = sd_ctx->sd->get_diffusion_model_down_factor(); - seed = sd_vid_gen_params->seed; - strength = sd_vid_gen_params->strength; - cache_params = &sd_vid_gen_params->cache; - vace_strength = sd_vid_gen_params->vace_strength; - guidance = sd_vid_gen_params->sample_params.guidance; - high_noise_guidance = sd_vid_gen_params->high_noise_sample_params.guidance; - hires = sd_vid_gen_params->hires; - resolve(sd_ctx); - if (frames != requested_frames) { - LOG_WARN("align video frames from %d to %d for %s", - requested_frames, - frames, - model_version_to_str[sd_ctx->sd->version]); - } - } - - void align_generation_request_size() { - align_image_size(&width, &height, "generation request"); - } - - void align_image_size(int* target_width, int* target_height, const char* label) { - int spatial_multiple = vae_scale_factor * diffusion_model_down_factor; - int width_offset = align_up_offset(*target_width, spatial_multiple); - int height_offset = align_up_offset(*target_height, spatial_multiple); - if (width_offset <= 0 && height_offset <= 0) { - return; - } - - int original_width = *target_width; - int original_height = *target_height; - - *target_width += width_offset; - *target_height += height_offset; - LOG_WARN("align %s up %dx%d to %dx%d (multiple=%d)", - label, - original_width, - original_height, - *target_width, - *target_height, - spatial_multiple); - } - - void resolve_hires() { - if (!hires.enabled) { - return; - } - if (hires.upscaler == SD_HIRES_UPSCALER_NONE) { - hires.enabled = false; - return; - } - if (hires.upscaler < SD_HIRES_UPSCALER_NONE || hires.upscaler >= SD_HIRES_UPSCALER_COUNT) { - LOG_WARN("hires upscaler '%d' is invalid, disabling hires", hires.upscaler); - hires.enabled = false; - return; - } - if (hires.upscaler == SD_HIRES_UPSCALER_MODEL && strlen(SAFE_STR(hires.model_path)) == 0) { - LOG_WARN("hires model upscaler requires a model path, disabling hires"); - hires.enabled = false; - return; - } - if (hires.scale <= 0.f && hires.target_width <= 0 && hires.target_height <= 0) { - LOG_WARN("hires scale must be positive when no target size is set, disabling hires"); - hires.enabled = false; - return; - } - if (hires.custom_sigmas_count < 0) { - LOG_WARN("hires custom sigmas count is negative, ignoring custom sigmas"); - hires.custom_sigmas = nullptr; - hires.custom_sigmas_count = 0; - } - if (hires.custom_sigmas_count > 0 && hires.custom_sigmas == nullptr) { - LOG_WARN("hires custom sigmas count is positive but custom sigmas are null, ignoring custom sigmas"); - hires.custom_sigmas_count = 0; - } - if (hires.custom_sigmas_count == 1) { - LOG_WARN("hires custom sigmas requires at least two values, ignoring custom sigmas"); - hires.custom_sigmas = nullptr; - hires.custom_sigmas_count = 0; - } - hires.denoising_strength = std::clamp(hires.denoising_strength, 0.0001f, 1.f); - hires.steps = std::max(0, hires.steps); - - if (hires.target_width > 0 && hires.target_height > 0) { - // pass - } else if (hires.target_width > 0) { - hires.target_height = hires.target_width; - } else if (hires.target_height > 0) { - hires.target_width = hires.target_height; - } else { - hires.target_width = static_cast(std::round(width * hires.scale)); - hires.target_height = static_cast(std::round(height * hires.scale)); - } - - if (hires.target_width <= 0 || hires.target_height <= 0) { - LOG_WARN("hires target size is not positive, disabling hires"); - hires.enabled = false; - return; - } - align_image_size(&hires.target_width, &hires.target_height, "hires target"); - } - - static void resolve_guidance(sd_ctx_t* sd_ctx, - sd_guidance_params_t* guidance, - bool* use_uncond, - bool* use_img_uncond, - bool has_ref_images, - const char* stage_name = nullptr) { - GGML_ASSERT(guidance != nullptr); - GGML_ASSERT(use_uncond != nullptr); - GGML_ASSERT(use_img_uncond != nullptr); - // out_img_uncond + text_cfg_scale * (out_cond - out_uncond) + image_cfg_scale * (out_uncond - out_img_uncond) - // -> text_cfg_scale * out_cond + (image_cfg_scale - text_cfg_scale) * out_uncond + (1 - image_cfg_scale) * out_img_uncond - // out_cond : prompt, image latent - // out_uncond : negative prompt, image latent - // out_img_uncond : negative prompt, zero image latent - // image_cfg_scale == 1 reduces 3-cond CFG to 2-cond CFG. - bool img_cfg_was_set = std::isfinite(guidance->img_cfg); - if (!img_cfg_was_set) { - guidance->img_cfg = 1.f; - } - - if (!sd_version_supports_img_cfg(sd_ctx->sd->version, has_ref_images)) { - if (img_cfg_was_set && guidance->img_cfg != 1.f) { - LOG_WARN("3-conditioning CFG is not supported with this model, disabling it for better performance"); - } - guidance->img_cfg = 1.f; - } - - if (guidance->img_cfg != guidance->txt_cfg) { - *use_uncond = true; - } - - if (guidance->img_cfg != 1.f) { - *use_img_uncond = true; - } - - if (guidance->txt_cfg < 1.f) { - const char* prefix = stage_name == nullptr ? "" : stage_name; - if (guidance->txt_cfg == 0.f) { - LOG_WARN("%sunconditioned mode, images won't follow the prompt (use cfg-scale=1 for distilled models)", - prefix); - } else { - LOG_WARN("%scfg value out of expected range may produce unexpected results", prefix); - } - } - } - - void resolve(sd_ctx_t* sd_ctx) { - align_generation_request_size(); - resolve_hires(); - seed = resolve_seed(seed); - - resolve_guidance(sd_ctx, &guidance, &use_uncond, &use_img_uncond, has_ref_images); - if (sd_ctx->sd->high_noise_diffusion_model) { - resolve_guidance(sd_ctx, - &high_noise_guidance, - &use_high_noise_uncond, - &use_high_noise_img_uncond, - has_ref_images, - "high noise: "); - } - - if (shifted_timestep > 0 && !sd_version_is_sdxl(sd_ctx->sd->version)) { - LOG_WARN("timestep shifting is only supported for SDXL models!"); - shifted_timestep = 0; - } - } -}; - -struct SamplePlan { - enum sample_method_t sample_method = SAMPLE_METHOD_COUNT; - enum sample_method_t high_noise_sample_method = SAMPLE_METHOD_COUNT; - const char* extra_sample_args = nullptr; - const char* high_noise_extra_sample_args = nullptr; - float eta = 0.f; - float high_noise_eta = 0.f; - int sample_steps = 0; - int high_noise_sample_steps = 0; - int total_steps = 0; - float moe_boundary = 0.f; - std::vector sigmas; - - SamplePlan(sd_ctx_t* sd_ctx, - const sd_img_gen_params_t* sd_img_gen_params, - const GenerationRequest& request) { - sample_method = sd_img_gen_params->sample_params.sample_method; - extra_sample_args = sd_img_gen_params->sample_params.extra_sample_args; - eta = sd_img_gen_params->sample_params.eta; - sample_steps = sd_img_gen_params->sample_params.sample_steps; - resolve(sd_ctx, &request, &sd_img_gen_params->sample_params); - } - - SamplePlan(sd_ctx_t* sd_ctx, - const sd_vid_gen_params_t* sd_vid_gen_params, - const GenerationRequest& request) { - sample_method = sd_vid_gen_params->sample_params.sample_method; - extra_sample_args = sd_vid_gen_params->sample_params.extra_sample_args; - eta = sd_vid_gen_params->sample_params.eta; - sample_steps = sd_vid_gen_params->sample_params.sample_steps; - if (sd_ctx->sd->high_noise_diffusion_model) { - high_noise_sample_steps = sd_vid_gen_params->high_noise_sample_params.sample_steps; - high_noise_sample_method = sd_vid_gen_params->high_noise_sample_params.sample_method; - high_noise_extra_sample_args = sd_vid_gen_params->high_noise_sample_params.extra_sample_args; - high_noise_eta = sd_vid_gen_params->high_noise_sample_params.eta; - } - moe_boundary = sd_vid_gen_params->moe_boundary; - resolve(sd_ctx, &request, &sd_vid_gen_params->sample_params); - } - - void resolve(sd_ctx_t* sd_ctx, - const GenerationRequest* request, - const sd_sample_params_t* sample_params) { - sample_method = resolve_sample_method(sd_ctx, sample_method); - - total_steps = sample_steps + std::max(0, high_noise_sample_steps); - - if (sample_params->custom_sigmas_count > 0) { - sigmas = std::vector(sample_params->custom_sigmas, - sample_params->custom_sigmas + sample_params->custom_sigmas_count); - total_steps = static_cast(sigmas.size()) - 1; - LOG_WARN("total_steps != custom_sigmas_count - 1, set total_steps to %d", total_steps); - if (sample_steps >= total_steps) { - sample_steps = total_steps; - LOG_WARN("total_steps != custom_sigmas_count - 1, set sample_steps to %d", sample_steps); - } - if (high_noise_sample_steps > 0) { - high_noise_sample_steps = total_steps - sample_steps; - LOG_WARN("total_steps != custom_sigmas_count - 1, set high_noise_sample_steps to %d", high_noise_sample_steps); - } - } else { - scheduler_t scheduler = resolve_scheduler(sd_ctx, - sample_params->scheduler, - sample_method); - int sample_seq_len = sd_ctx->sd->get_image_seq_len(request->height, request->width); - if (sd_version_is_ltxav(sd_ctx->sd->version) && request->frames > 0) { - int latent_frames = ((request->frames - 1) / 8) + 1; - sample_seq_len *= latent_frames; - } else if (sd_version_is_minimax_h3(sd_ctx->sd->version) && request->frames > 0) { - sample_seq_len *= sd_ctx->sd->video_frames_to_latent_frames(request->frames); - } - sigmas = sd_ctx->sd->denoiser->get_sigmas(total_steps, - sample_seq_len, - scheduler, - sd_ctx->sd->version, - sample_params->extra_sample_args); - } - - eta = resolve_eta(sd_ctx, eta, sample_method); - - if (high_noise_sample_steps < 0) { - for (size_t i = 0; i < sigmas.size(); ++i) { - if (sigmas[i] < moe_boundary) { - high_noise_sample_steps = static_cast(i); - break; - } - } - LOG_VERBOSE("switching from high noise model at step %d", high_noise_sample_steps); - } - - LOG_INFO("sampling using %s method", sampling_methods_str[sample_method]); - if (high_noise_sample_steps > 0) { - high_noise_sample_method = resolve_sample_method(sd_ctx, - high_noise_sample_method); - high_noise_eta = resolve_eta(sd_ctx, high_noise_eta, high_noise_sample_method); - LOG_INFO("sampling(high noise) using %s method", sampling_methods_str[high_noise_sample_method]); - } - } -}; - -struct ImageGenerationLatents { - sd::Tensor init_latent; - sd::Tensor concat_latent; - sd::Tensor img_uncond_concat_latent; - sd::Tensor audio_latent; - sd::Tensor video_positions; - sd::Tensor control_image; - std::vector> ref_images; - std::vector> ref_latents; - std::vector> reference_audio_latents; - std::vector minimax_reference_blocks; - std::vector minimax_presentation_refs; - std::vector keyframe_indices; - sd::Tensor denoise_mask; - sd::Tensor clip_vision_output; - sd::Tensor vace_context; - int64_t ref_image_num = 0; - int64_t video_conditioning_frame_count = 0; - int64_t video_target_frame_count = 0; - int audio_length = 0; -}; - -static float ltxv_latent_corner_to_pixel_frame(int64_t corner_index, - int temporal_scale, - bool causal_temporal_positioning) { - float pixel_t = static_cast(corner_index * temporal_scale); - if (causal_temporal_positioning) { - pixel_t = std::max(0.f, pixel_t + 1.f - static_cast(temporal_scale)); - } - return pixel_t; -} - -static void set_ltxv_video_position(sd::Tensor* positions, - int64_t token, - float t_start, - float t_end, - float h_start, - float h_end, - float w_start, - float w_end) { - positions->index(0, 0, token, 0) = t_start; - positions->index(1, 0, token, 0) = t_end; - positions->index(0, 1, token, 0) = h_start; - positions->index(1, 1, token, 0) = h_end; - positions->index(0, 2, token, 0) = w_start; - positions->index(1, 2, token, 0) = w_end; -} - -static sd::Tensor build_ltxv_video_positions(int64_t width, - int64_t height, - int64_t target_latent_frames, - int64_t keyframe_latent_frames, - int keyframe_frame_idx, - int keyframe_pixel_frames, - int fps, - int spatial_scale, - int temporal_scale, - bool causal_temporal_positioning) { - GGML_ASSERT(width > 0 && height > 0 && target_latent_frames > 0); - GGML_ASSERT(keyframe_latent_frames > 0); - GGML_ASSERT(fps > 0); - - int64_t total_tokens = width * height * (target_latent_frames + keyframe_latent_frames); - sd::Tensor positions({2, 3, total_tokens, 1}); - int64_t token = 0; - - for (int64_t t = 0; t < target_latent_frames; t++) { - float t_start = ltxv_latent_corner_to_pixel_frame(t, temporal_scale, causal_temporal_positioning) / static_cast(fps); - float t_end = ltxv_latent_corner_to_pixel_frame(t + 1, temporal_scale, causal_temporal_positioning) / static_cast(fps); - for (int64_t h = 0; h < height; h++) { - float h_start = static_cast(h * spatial_scale); - float h_end = static_cast((h + 1) * spatial_scale); - for (int64_t w = 0; w < width; w++) { - float w_start = static_cast(w * spatial_scale); - float w_end = static_cast((w + 1) * spatial_scale); - set_ltxv_video_position(&positions, token++, t_start, t_end, h_start, h_end, w_start, w_end); - } - } - } - - for (int64_t t = 0; t < keyframe_latent_frames; t++) { - float t_start = static_cast(keyframe_frame_idx + t * temporal_scale); - float t_end = static_cast(keyframe_frame_idx + (t + 1) * temporal_scale); - if (keyframe_pixel_frames == 1) { - t_end = t_start + 1.f; - } - t_start /= static_cast(fps); - t_end /= static_cast(fps); - for (int64_t h = 0; h < height; h++) { - float h_start = static_cast(h * spatial_scale); - float h_end = static_cast((h + 1) * spatial_scale); - for (int64_t w = 0; w < width; w++) { - float w_start = static_cast(w * spatial_scale); - float w_end = static_cast((w + 1) * spatial_scale); - set_ltxv_video_position(&positions, token++, t_start, t_end, h_start, h_end, w_start, w_end); - } - } - } - - return positions; -} - -static sd::Tensor pack_ltxav_audio_and_video_latents(const sd::Tensor& video_latent, - const sd::Tensor& audio_latent) { - if (audio_latent.empty()) { - return video_latent; - } - - GGML_ASSERT(video_latent.dim() == 4 || video_latent.dim() == 5); - GGML_ASSERT(audio_latent.dim() == 3 || audio_latent.dim() == 4); - if (video_latent.dim() == 5) { - GGML_ASSERT(video_latent.shape()[4] == 1); - } - if (audio_latent.dim() == 4) { - GGML_ASSERT(audio_latent.shape()[3] == 1); - } - - int64_t width = video_latent.shape()[0]; - int64_t height = video_latent.shape()[1]; - int64_t frames = video_latent.shape()[2]; - int64_t video_ch = video_latent.shape()[3]; - int64_t spatial_size = width * height * frames; - int64_t audio_values = audio_latent.numel(); - int64_t extra_ch = (audio_values + spatial_size - 1) / spatial_size; - - std::vector packed_shape = video_latent.shape(); - packed_shape[3] = video_ch + extra_ch; - sd::Tensor packed = sd::zeros(packed_shape); - - std::copy_n(video_latent.data(), video_latent.numel(), packed.data()); - std::copy_n(audio_latent.data(), audio_latent.numel(), packed.data() + video_latent.numel()); - return packed; -} - -static sd::Tensor pack_ltxav_audio_and_video_denoise_mask(const sd::Tensor& video_mask, - const sd::Tensor& video_latent, - const sd::Tensor& audio_latent) { - if (video_mask.empty() || audio_latent.empty()) { - return video_mask; - } - - GGML_ASSERT(video_latent.dim() == 4 || video_latent.dim() == 5); - GGML_ASSERT(audio_latent.dim() == 3 || audio_latent.dim() == 4); - if (video_latent.dim() == 5) { - GGML_ASSERT(video_latent.shape()[4] == 1); - } - if (audio_latent.dim() == 4) { - GGML_ASSERT(audio_latent.shape()[3] == 1); - } - - int64_t width = video_latent.shape()[0]; - int64_t height = video_latent.shape()[1]; - int64_t frames = video_latent.shape()[2]; - int64_t video_ch = video_latent.shape()[3]; - int64_t spatial_size = width * height * frames; - int64_t audio_values = audio_latent.numel(); - int64_t extra_ch = (audio_values + spatial_size - 1) / spatial_size; - - GGML_ASSERT(video_mask.dim() == video_latent.dim()); - GGML_ASSERT(video_mask.shape()[0] == width); - GGML_ASSERT(video_mask.shape()[1] == height); - GGML_ASSERT(video_mask.shape()[2] == frames); - if (video_mask.dim() == 5) { - GGML_ASSERT(video_mask.shape()[4] == video_latent.shape()[4]); - } - - int64_t mask_ch = video_mask.shape()[3]; - if (mask_ch == video_ch + extra_ch) { - return video_mask; - } - GGML_ASSERT(mask_ch == 1 || mask_ch == video_ch); - - sd::Tensor video_mask_full = video_mask; - if (mask_ch == 1 && video_ch != 1) { - video_mask_full = video_mask * sd::Tensor::ones(video_latent.shape()); - } - - std::vector audio_mask_shape = video_latent.shape(); - audio_mask_shape[3] = extra_ch; - auto audio_mask = sd::Tensor::ones(audio_mask_shape); - return sd::ops::concat(video_mask_full, audio_mask, 3); -} - -static sd::Tensor make_ltxav_video_denoise_mask(const sd::Tensor& video_latent, float value = 1.f) { - if (video_latent.empty()) { - return {}; - } - return sd::full({video_latent.shape()[0], - video_latent.shape()[1], - video_latent.shape()[2], - 1, - 1}, - value); -} - -static sd::Tensor encode_ltxav_condition_image(sd_ctx_t* sd_ctx, - const sd::Tensor& image, - const char* name) { - if (sd_ctx == nullptr || sd_ctx->sd == nullptr || image.empty()) { - return {}; - } - auto condition_image = image.reshape({image.shape()[0], - image.shape()[1], - 1, - image.shape()[2], - image.shape()[3]}); - auto condition_latent = sd_ctx->sd->encode_first_stage(condition_image); - if (condition_latent.empty()) { - LOG_ERROR("failed to encode LTXAV %s image", name); - } - return condition_latent; -} - -static bool apply_ltxav_condition_by_latent_index(sd::Tensor* video_latent, - sd::Tensor* video_mask, - const sd::Tensor& condition_latent, - int64_t latent_idx, - const char* name, - float conditioned_mask) { - if (video_latent == nullptr || video_mask == nullptr || video_latent->empty() || video_mask->empty()) { - return false; - } - if (condition_latent.empty() || - condition_latent.shape()[0] != video_latent->shape()[0] || - condition_latent.shape()[1] != video_latent->shape()[1] || - condition_latent.shape()[3] != video_latent->shape()[3]) { - LOG_ERROR("invalid LTXAV %s condition latent shape", name); - return false; - } - int64_t latent_frames = video_latent->shape()[2]; - int64_t condition_frames = condition_latent.shape()[2]; - if (latent_idx < 0 || condition_frames <= 0 || latent_idx + condition_frames > latent_frames) { - LOG_ERROR("invalid LTXAV %s image latent range: start=%" PRId64 ", length=%" PRId64 ", latent_frames=%" PRId64, - name, - latent_idx, - condition_frames, - latent_frames); - return false; - } - - sd::ops::slice_assign(video_latent, 2, latent_idx, latent_idx + condition_frames, condition_latent); - sd::ops::fill_slice(video_mask, 2, latent_idx, latent_idx + condition_frames, conditioned_mask); - return true; -} - -static bool apply_ltxav_condition_image_by_latent_index(sd_ctx_t* sd_ctx, - const sd::Tensor& image, - sd::Tensor* video_latent, - sd::Tensor* video_mask, - int64_t latent_idx, - const char* name, - float strength) { - auto condition_latent = encode_ltxav_condition_image(sd_ctx, image, name); - return !condition_latent.empty() && - apply_ltxav_condition_by_latent_index(video_latent, - video_mask, - condition_latent, - latent_idx, - name, - 1.0f - std::clamp(strength, 0.f, 1.f)); -} - -static sd::Tensor unpack_ltxav_audio_latent(const sd::Tensor& packed_latent, - int audio_length, - int video_channels) { - if (packed_latent.empty() || audio_length <= 0) { - return {}; - } - - GGML_ASSERT(packed_latent.dim() == 4 || packed_latent.dim() == 5); - int64_t width = packed_latent.shape()[0]; - int64_t height = packed_latent.shape()[1]; - int64_t frames = packed_latent.shape()[2]; - int64_t total_channels = packed_latent.shape()[3]; - int64_t spatial_size = width * height * frames; - if (total_channels <= video_channels) { - return {}; - } - - constexpr int kLtxavAudioFrequencyBins = 16; - constexpr int kLtxavAudioChannels = 8; - int64_t required_values = static_cast(audio_length) * kLtxavAudioFrequencyBins * kLtxavAudioChannels; - int64_t packed_values = (total_channels - video_channels) * spatial_size; - if (packed_values < required_values) { - return {}; - } - - sd::Tensor audio_latent({kLtxavAudioFrequencyBins, audio_length, kLtxavAudioChannels, 1}); - const float* audio_src = packed_latent.data() + static_cast(video_channels) * static_cast(spatial_size); - std::copy_n(audio_src, static_cast(required_values), audio_latent.data()); - return audio_latent; -} - -static sd::Tensor make_ltxav_empty_audio_latent(int audio_length) { - if (audio_length <= 0) { - return {}; - } - constexpr int kLtxavAudioFrequencyBins = 16; - constexpr int kLtxavAudioChannels = 8; - return sd::zeros({kLtxavAudioFrequencyBins, audio_length, kLtxavAudioChannels, 1}); -} - -static sd::Tensor resize_ltxav_audio_latent(const sd::Tensor& audio_latent, - int target_audio_length) { - auto resized = make_ltxav_empty_audio_latent(target_audio_length); - if (resized.empty() || audio_latent.empty()) { - return resized; - } - GGML_ASSERT(audio_latent.dim() == 3 || audio_latent.dim() == 4); - int copy_length = std::min(static_cast(audio_latent.shape()[1]), target_audio_length); - if (copy_length > 0) { - auto copied = sd::ops::slice(audio_latent, 1, 0, copy_length); - sd::ops::slice_assign(&resized, 1, 0, copy_length, copied); - } - return resized; -} - -static int get_ltxav_num_audio_latents(int frames, int fps) { - GGML_ASSERT(frames > 0); - GGML_ASSERT(fps > 0); - constexpr float kSampleRate = 16000.0f; - constexpr float kMelHopLength = 160.0f; - constexpr float kAudioLatentDownsample = 4.0f; - constexpr float kLatentsPerSecond = kSampleRate / kMelHopLength / kAudioLatentDownsample; - return static_cast(std::ceil((static_cast(frames) / static_cast(fps)) * kLatentsPerSecond)); -} - -static int get_minimax_h3_num_audio_latents(int frames, int fps) { - GGML_ASSERT(frames > 0 && fps > 0); - return std::max(1, - static_cast(std::lround( - static_cast(frames) * 40.0 / fps))); -} - -static sd::Tensor make_minimax_h3_empty_audio_latent(int audio_length) { - if (audio_length <= 0) { - return {}; - } - return sd::zeros({audio_length, 2, 32, 1}); -} - -static sd::Tensor prepare_minimax_h3_reference_waveform(const sd_audio_t& audio, - int target_sample_rate = 32000) { - if (audio.data == nullptr || audio.sample_count == 0 || audio.channels == 0 || audio.sample_rate == 0) { - return {}; - } - uint64_t output_samples = static_cast(std::llround( - static_cast(audio.sample_count) * target_sample_rate / audio.sample_rate)); - output_samples = std::max(1, output_samples); - uint64_t padded_samples = (output_samples + 799) / 800 * 800; - // Keep stereo streams planar for the mono-per-stream audio encoder: - // [samples, 1, stereo, batch]. This avoids flattening interleaved L/R - // storage into alternating samples when the encoder folds streams into - // its batch dimension. - sd::Tensor waveform({static_cast(padded_samples), 1, 2, 1}); - - for (uint64_t i = 0; i < output_samples; ++i) { - long double source_pos = static_cast(i) * audio.sample_rate / target_sample_rate; - uint64_t source0 = std::min(static_cast(source_pos), audio.sample_count - 1); - uint64_t source1 = std::min(source0 + 1, audio.sample_count - 1); - float fraction = static_cast(source_pos - source0); - for (uint32_t channel = 0; channel < 2; ++channel) { - uint32_t source_channel = audio.channels == 1 ? 0 : std::min(channel, audio.channels - 1); - float a = audio.data[source0 * audio.channels + source_channel]; - float b = audio.data[source1 * audio.channels + source_channel]; - waveform.index(static_cast(i), 0, channel, 0) = - std::clamp(a + (b - a) * fraction, -1.f, 1.f); - } - } - return waveform; -} - -static sd::Tensor unpack_minimax_h3_audio_latent(const sd::Tensor& packed_latent, - int audio_length, - int video_channels) { - if (packed_latent.empty() || audio_length <= 0) { - return {}; - } - GGML_ASSERT(packed_latent.dim() == 4 || packed_latent.dim() == 5); - int64_t spatial_size = packed_latent.shape()[0] * packed_latent.shape()[1] * packed_latent.shape()[2]; - int64_t required = static_cast(audio_length) * 2 * 32; - int64_t available = (packed_latent.shape()[3] - video_channels) * spatial_size; - if (available < required) { - return {}; - } - sd::Tensor audio({audio_length, 2, 32, 1}); - const float* source = packed_latent.data() + - static_cast(video_channels) * static_cast(spatial_size); - std::copy_n(source, static_cast(required), audio.data()); - return audio; -} - -struct ImageGenerationEmbeds { - SDCondition cond; - SDCondition uncond; - SDCondition img_uncond; -}; - -struct ConditionerRunnerEndOnExit { - Conditioner* conditioner = nullptr; - ~ConditionerRunnerEndOnExit() { - if (conditioner != nullptr) { - conditioner->runner_end(); - } - } -}; - -struct CircularAxesState { - bool circular_x = false; - bool circular_y = false; -}; - -static void apply_circular_axes_to_diffusion(sd_ctx_t* sd_ctx, bool circular_x, bool circular_y) { - sd_ctx->sd->circular_x = circular_x; - sd_ctx->sd->circular_y = circular_y; - if (sd_ctx->sd->diffusion_model) { - sd_ctx->sd->diffusion_model->set_circular_axes(circular_x, circular_y); - } - if (sd_ctx->sd->high_noise_diffusion_model) { - sd_ctx->sd->high_noise_diffusion_model->set_circular_axes(circular_x, circular_y); - } - if (sd_ctx->sd->control_net) { - sd_ctx->sd->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"); - } -} - -static CircularAxesState configure_image_vae_axes(sd_ctx_t* sd_ctx, - const sd_img_gen_params_t* sd_img_gen_params, - const GenerationRequest& request) { - CircularAxesState original_axes = {sd_ctx->sd->circular_x, sd_ctx->sd->circular_y}; - - if (!sd_img_gen_params->vae_tiling_params.enabled) { - if (sd_ctx->sd->first_stage_model) { - sd_ctx->sd->first_stage_model->set_circular_axes(sd_ctx->sd->circular_x, sd_ctx->sd->circular_y); - } - if (sd_ctx->sd->preview_vae) { - sd_ctx->sd->preview_vae->set_circular_axes(sd_ctx->sd->circular_x, sd_ctx->sd->circular_y); - } - return original_axes; - } - - int tile_size_x, tile_size_y; - float overlap; - int latent_size_x = request.width / request.vae_scale_factor; - int latent_size_y = request.height / request.vae_scale_factor; - sd_ctx->sd->first_stage_model->get_tile_sizes(tile_size_x, - tile_size_y, - overlap, - sd_img_gen_params->vae_tiling_params, - latent_size_x, - latent_size_y); - - sd_ctx->sd->circular_x = sd_ctx->sd->circular_x && (tile_size_x >= latent_size_x); - sd_ctx->sd->circular_y = sd_ctx->sd->circular_y && (tile_size_y >= latent_size_y); - - if (sd_ctx->sd->first_stage_model) { - sd_ctx->sd->first_stage_model->set_circular_axes(sd_ctx->sd->circular_x, sd_ctx->sd->circular_y); - } - if (sd_ctx->sd->preview_vae) { - sd_ctx->sd->preview_vae->set_circular_axes(sd_ctx->sd->circular_x, sd_ctx->sd->circular_y); - } - - sd_ctx->sd->circular_x = original_axes.circular_x && (tile_size_x < latent_size_x); - sd_ctx->sd->circular_y = original_axes.circular_y && (tile_size_y < latent_size_y); - - return original_axes; -} - -static void restore_image_vae_axes(sd_ctx_t* sd_ctx, const CircularAxesState& original_axes) { - sd_ctx->sd->circular_x = original_axes.circular_x; - sd_ctx->sd->circular_y = original_axes.circular_y; -} - -class ImageVaeAxesGuard { -private: - sd_ctx_t* sd_ctx = nullptr; - CircularAxesState original_axes; - -public: - ImageVaeAxesGuard(sd_ctx_t* sd_ctx, - const sd_img_gen_params_t* sd_img_gen_params, - const GenerationRequest& request) - : sd_ctx(sd_ctx), - original_axes(configure_image_vae_axes(sd_ctx, sd_img_gen_params, request)) {} - - ~ImageVaeAxesGuard() { - restore_image_vae_axes(sd_ctx, original_axes); - } - - ImageVaeAxesGuard(const ImageVaeAxesGuard&) = delete; - ImageVaeAxesGuard& operator=(const ImageVaeAxesGuard&) = delete; -}; - -static sd::Tensor ensure_image_tensor_channels(sd::Tensor image, int channels) { - if (image.empty()) { - return image; - } - GGML_ASSERT(image.dim() == 4); - int64_t current_channels = image.shape()[2]; - if (current_channels == channels) { - return image; - } - if (channels == 4) { - sd::Tensor alpha = sd::full({image.shape()[0], image.shape()[1], 1, image.shape()[3]}, 1.f); - if (current_channels == 3) { - return sd::ops::concat(image, alpha, 2); - } - if (current_channels == 1) { - sd::Tensor rgb = sd::ops::concat(image, image, 2); - rgb = sd::ops::concat(rgb, image, 2); - return sd::ops::concat(rgb, alpha, 2); - } - } - if (channels == 3 && current_channels >= 3) { - return sd::ops::slice(image, 2, 0, 3); - } - GGML_ABORT("cannot convert image tensor from %lld to %d channels", - (long long)current_channels, - channels); -} - -static std::optional prepare_image_generation_latents(sd_ctx_t* sd_ctx, - const sd_img_gen_params_t* sd_img_gen_params, - GenerationRequest* request, - SamplePlan* plan, - const RefImageParams& ref_image_params) { - int64_t prepare_start_ms = ggml_time_ms(); - - sd::Tensor init_image_tensor; - sd::Tensor control_image_tensor; - sd::Tensor mask_image_tensor; - int image_channels = sd_ctx->sd->get_image_channels(); - - if (sd_img_gen_params->init_image.data != nullptr) { - LOG_INFO("IMG2IMG"); - - if (request->strength < 1.f) { - bool strength_as_noise_level = false; - bool force_first_sigma = false; - for (const auto& [key, value] : parse_key_value_args(sd_img_gen_params->sample_params.extra_sample_args, "img2img arg")) { - if (key == "strength_as_noise_level") { - if (!parse_strict_bool(value, strength_as_noise_level)) { - LOG_WARN("ignoring invalid img2img sample arg '%s=%s'", key.c_str(), value.c_str()); - } - } else if (key == "force_first_sigma") { - if (!parse_strict_bool(value, force_first_sigma)) { - LOG_WARN("ignoring invalid img2img sample arg '%s=%s'", key.c_str(), value.c_str()); - } - } - } - - size_t t_enc; - float target_sigma = -1; - if (!strength_as_noise_level) { - t_enc = static_cast(plan->sample_steps * request->strength); - if (t_enc == static_cast(plan->sample_steps)) { - t_enc--; - } - } else { - LOG_VERBOSE("Interpreting denoise strength as relative noise level"); - // assume x_noised = K * (x * (1-noise_level) + noise * noise_level) = K * lerp(x, noise, noise_level) - // K = 1, noise_level = sigma for flow models - // K = 1+sigma, noise_level=sigma/(1+sigma) for diffusion models - float target_noise_level = request->strength; - target_sigma = sd_ctx->sd->denoiser->noise_level_to_sigma(target_noise_level); - size_t start_index = 0; - for (size_t i = 0; i < plan->sigmas.size(); ++i) { - if (plan->sigmas[i] <= target_sigma) { - start_index = i; - break; - } - } - - if (start_index >= plan->sigmas.size() - 1) { - start_index = plan->sigmas.size() - 2; // Leave at least 1 step - } - t_enc = plan->sample_steps - start_index - 1; - } - LOG_INFO("target t_enc is %zu steps", t_enc); - std::vector sigma_sched; - sigma_sched.assign(plan->sigmas.begin() + plan->sample_steps - t_enc - 1, plan->sigmas.end()); - - if (target_sigma > 0 && force_first_sigma && strength_as_noise_level) { - LOG_VERBOSE("force_first_sigma to %.4f (from %.4f)", target_sigma, sigma_sched[0]); - sigma_sched[0] = target_sigma; - } - - plan->sigmas = std::move(sigma_sched); - plan->sample_steps = static_cast(plan->sigmas.size() - 1); - } - - init_image_tensor = ensure_image_tensor_channels(sd_image_to_tensor(sd_img_gen_params->init_image, request->width, request->height), - image_channels); - } - - if (sd_img_gen_params->mask_image.data != nullptr) { - mask_image_tensor = sd_image_to_tensor(sd_img_gen_params->mask_image, request->width, request->height); - mask_image_tensor = sd::ops::round(mask_image_tensor); - } - - if (sd_img_gen_params->control_image.data != nullptr) { - control_image_tensor = sd_image_to_tensor(sd_img_gen_params->control_image, request->width, request->height); - } - - if (init_image_tensor.empty() || mask_image_tensor.empty()) { - if (sd_version_is_inpaint(sd_ctx->sd->version)) { - LOG_WARN("inpainting model requires both an init image and a mask image."); - } - } - - if (mask_image_tensor.empty()) { - mask_image_tensor = sd::full({request->width, request->height, 1, 1}, 1.f); - } - - sd::Tensor latent_mask = sd::ops::interpolate(mask_image_tensor, - {request->width / request->vae_scale_factor, - request->height / request->vae_scale_factor, - 1, - 1}, - sd::ops::InterpolateMode::NearestMax); - - sd::Tensor init_latent; - sd::Tensor control_latent; - if (init_image_tensor.empty()) { - if (sd_ctx->sd->version == VERSION_QWEN_IMAGE_LAYERED) { - init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height, request->qwen_image_layers + 1, true); - } else { - init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height); - } - } else { - init_latent = sd_ctx->sd->encode_first_stage(init_image_tensor); - if (init_latent.empty()) { - LOG_ERROR("failed to encode init image"); - return std::nullopt; - } - } - - if (sd_ctx->sd->animatediff_num_frames > 1 && - init_latent.dim() >= 4 && init_latent.shape()[3] == 1) { - int n_frames = sd_ctx->sd->animatediff_num_frames; - std::vector shape(init_latent.shape().begin(), init_latent.shape().end()); - shape[3] = n_frames; - if (!init_image_tensor.empty()) { - sd::Tensor replicated(shape); - for (int f = 0; f < n_frames; ++f) { - sd::ops::slice_assign(&replicated, 3, f, f + 1, init_latent); - } - init_latent = std::move(replicated); - } else { - init_latent = sd::Tensor(std::move(shape)); - } - } - - if (!control_image_tensor.empty()) { - control_latent = sd_ctx->sd->encode_first_stage(control_image_tensor); - if (control_latent.empty()) { - LOG_ERROR("failed to encode control image"); - return std::nullopt; - } - } - - std::vector> ref_images; - for (int i = 0; i < sd_img_gen_params->ref_images_count; i++) { - ref_images.push_back(ensure_image_tensor_channels(sd_image_to_tensor(sd_img_gen_params->ref_images[i]), - image_channels)); - } - - if (ref_images.empty() && sd_version_is_unet_edit(sd_ctx->sd->version)) { - LOG_WARN("This model needs at least one reference image; using an empty reference"); - ref_images.push_back(sd::zeros({request->width, request->height, image_channels, 1})); - request->guidance.img_cfg = request->guidance.txt_cfg; - request->use_img_uncond = false; - } - - if (!ref_images.empty()) { - LOG_INFO("EDIT mode"); - } - - std::vector> ref_latents; - for (size_t i = 0; i < ref_images.size(); i++) { - if (sd_ctx->sd->version == VERSION_HIDREAM_O1) { - continue; - } - sd::Tensor ref_latent; - if (ref_image_params.resize_before_vae && !sd_version_is_pid(sd_ctx->sd->version)) { - LOG_VERBOSE("auto resize ref images"); - double vae_width; - double vae_height; - if (ref_image_params.resize_vae_to_target) { - vae_width = request->width; - vae_height = request->height; - } else { - int target_pixels = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : 1024 * 1024; - int vae_image_size = std::min(target_pixels, request->width * request->height); - vae_width = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]); - vae_height = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0]; - } - - int factor = sd_version_is_qwen_image(sd_ctx->sd->version) ? 32 : 16; - vae_height = round(vae_height / factor) * factor; - vae_width = round(vae_width / factor) * factor; - - auto resized_ref_img = sd::ops::interpolate(ref_images[i], - {static_cast(vae_width), - static_cast(vae_height), - ref_images[i].shape()[2], - ref_images[i].shape()[3]}); - - LOG_VERBOSE("resize vae ref image %d from %" PRId64 "x%" PRId64 " to %" PRId64 "x%" PRId64, - static_cast(i), - ref_images[i].shape()[1], - ref_images[i].shape()[0], - resized_ref_img.shape()[1], - resized_ref_img.shape()[0]); - - ref_latent = sd_ctx->sd->encode_first_stage(resized_ref_img); - } else { - ref_latent = sd_ctx->sd->encode_first_stage(ref_images[i]); - } - if (ref_latent.empty()) { - LOG_ERROR("failed to encode reference image %d", static_cast(i)); - return std::nullopt; - } - - ref_latents.push_back(std::move(ref_latent)); - } - - if (sd_version_is_pid(sd_ctx->sd->version)) { - if (ref_latents.empty()) { - LOG_ERROR("PiD requires a reference image"); - return std::nullopt; - } - } - - sd::Tensor concat_latent; - sd::Tensor img_uncond_concat_latent; - if (sd_version_is_inpaint(sd_ctx->sd->version)) { - sd::Tensor masked_init_latent; - - if (sd_ctx->sd->version != VERSION_FLEX_2) { - if (!init_image_tensor.empty()) { - auto masked_image = ((1.0f - mask_image_tensor) * (init_image_tensor - 0.5f)) + 0.5f; - masked_init_latent = sd_ctx->sd->encode_first_stage(masked_image); - if (masked_init_latent.empty()) { - LOG_ERROR("failed to encode masked init image"); - return std::nullopt; - } - } else { - masked_init_latent = sd::Tensor::zeros_like(init_latent); - } - } else { - masked_init_latent = ((1.0f - latent_mask) * init_latent); - } - - auto uncond_masked_init_latent = sd::Tensor::zeros_like(masked_init_latent); - - if (sd_ctx->sd->version == VERSION_FLUX_FILL) { - auto mask = mask_image_tensor.reshape({request->vae_scale_factor, - request->width / request->vae_scale_factor, - request->vae_scale_factor, - request->height / request->vae_scale_factor}); - mask = mask.permute({1, 3, 0, 2}).reshape({request->width / request->vae_scale_factor, request->height / request->vae_scale_factor, request->vae_scale_factor * request->vae_scale_factor, 1}); - - concat_latent = sd::ops::concat(masked_init_latent, mask, 2); - img_uncond_concat_latent = sd::ops::concat(uncond_masked_init_latent, mask, 2); - } else if (sd_ctx->sd->version == VERSION_FLEX_2) { - concat_latent = sd::ops::concat(masked_init_latent, latent_mask, 2); - if (!control_latent.empty()) { - concat_latent = sd::ops::concat(concat_latent, control_latent, 2); - } else { - concat_latent = sd::ops::concat(concat_latent, sd::Tensor::zeros_like(masked_init_latent), 2); - } - - img_uncond_concat_latent = sd::ops::concat(uncond_masked_init_latent, latent_mask, 2); - img_uncond_concat_latent = sd::ops::concat(img_uncond_concat_latent, sd::Tensor::zeros_like(masked_init_latent), 2); - } else { // SD1.x SD2.x SDXL inpaint - concat_latent = sd::ops::concat(latent_mask, masked_init_latent, 2); - img_uncond_concat_latent = sd::ops::concat(latent_mask, uncond_masked_init_latent, 2); - } - } - if (sd_version_is_unet_edit(sd_ctx->sd->version)) { - concat_latent = sd::ops::interpolate(ref_latents[0], init_latent.shape()); - img_uncond_concat_latent = sd::Tensor::zeros_like(concat_latent); - } - if (sd_ctx->sd->version == VERSION_FLUX_CONTROLS) { - if (!control_latent.empty()) { - concat_latent = control_latent; - } else { - concat_latent = sd::Tensor::zeros_like(init_latent); - } - img_uncond_concat_latent = sd::Tensor::zeros_like(concat_latent); - } - - if (sd_img_gen_params->init_image.data != nullptr || sd_img_gen_params->ref_images_count > 0) { - int64_t t1 = ggml_time_ms(); - LOG_INFO("encode_first_stage completed, taking %.2fs", (t1 - prepare_start_ms) * 1.0f / 1000); - } - - ImageGenerationLatents latents; - latents.init_latent = std::move(init_latent); - latents.concat_latent = std::move(concat_latent); - latents.img_uncond_concat_latent = std::move(img_uncond_concat_latent); - latents.control_image = std::move(control_image_tensor); - latents.ref_images = std::move(ref_images); - latents.ref_latents = std::move(ref_latents); - - if (sd_version_is_inpaint(sd_ctx->sd->version)) { - latent_mask = sd::ops::max_pool_2d(latent_mask, - {3, 3}, - {1, 1}, - {1, 1}); - } - latents.denoise_mask = std::move(latent_mask); - - return latents; -} - -static std::optional prepare_image_generation_embeds(sd_ctx_t* sd_ctx, - const sd_img_gen_params_t* sd_img_gen_params, - GenerationRequest* request, - SamplePlan* plan, - ImageGenerationLatents* latents, - const RefImageParams& ref_image_params) { - ConditionerRunnerEndOnExit conditioner_runner_end{sd_ctx->sd->cond_stage_model.get()}; - - ConditionerParams condition_params; - condition_params.text = request->prompt; - condition_params.clip_skip = request->clip_skip; - condition_params.width = request->width; - condition_params.height = request->height; - if (ref_image_params.pass_to_vlm) { - condition_params.ref_images = &latents->ref_images; - } - - condition_params.ref_image_params = ref_image_params; - - sd_ctx->sd->prepare_generation_extensions(request->pm_params, - request->pulid_params, - condition_params, - plan->total_steps); - sd_ctx->sd->compute_ip_adapter_tokens(sd_img_gen_params->ip_adapter_image, sd_img_gen_params->ip_adapter_strength); - int64_t prepare_start_ms = ggml_time_ms(); - condition_params.zero_out_masked = false; - auto cond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, - condition_params); - if (cond.c_concat.empty() && ref_image_params.pass_to_dit) { - cond.c_concat = latents->concat_latent; // TODO: optimize - } - - bool use_ref_latent_img_cfg = request->use_img_uncond && - !latents->ref_images.empty() && - sd_version_supports_ref_latent_img_cfg(sd_ctx->sd->version); - - SDCondition uncond; - if (request->use_uncond || request->use_high_noise_uncond) { - if (sd_version_is_ideogram4(sd_ctx->sd->version)) { - uncond.c_vector = sd::Tensor::from_vector({1.0f}); - } else if (sd_version_is_minit2i(sd_ctx->sd->version)) { - // MiniT2I derives the unconditional signal from the same T5 hidden - // states with a zeroed prompt mask, so no extra text encode is needed. - uncond.c_crossattn = cond.c_crossattn; - uncond.c_vector = sd::Tensor::zeros_like(cond.c_vector); - } else { - bool zero_out_masked = false; - if (sd_version_is_sdxl(sd_ctx->sd->version) && - request->negative_prompt.empty() && - !sd_ctx->sd->is_using_edm_v_parameterization) { - zero_out_masked = true; - } - condition_params.text = request->negative_prompt; - condition_params.zero_out_masked = zero_out_masked; - uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, - condition_params); - } - if (uncond.c_concat.empty() && ref_image_params.pass_to_dit) { - uncond.c_concat = latents->concat_latent; // TODO: optimize - } - } - - SDCondition img_uncond; - if (request->use_img_uncond) { - if ((request->use_uncond || request->use_high_noise_uncond) && (latents->ref_images.empty() || !use_ref_latent_img_cfg)) { - img_uncond = SDCondition(uncond.c_crossattn, uncond.c_vector, latents->img_uncond_concat_latent); - } else { - bool zero_out_masked = false; - if (sd_version_is_sdxl(sd_ctx->sd->version) && - request->negative_prompt.empty() && - !sd_ctx->sd->is_using_edm_v_parameterization) { - zero_out_masked = true; - } - condition_params.text = request->negative_prompt; - condition_params.zero_out_masked = zero_out_masked; - std::vector> empty_ref_images; - if (use_ref_latent_img_cfg) { - condition_params.ref_images = &empty_ref_images; - } - img_uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, - condition_params); - if (img_uncond.c_concat.empty() && ref_image_params.pass_to_dit) { - img_uncond.c_concat = latents->img_uncond_concat_latent; // TODO: optimize - } - } - } - - int64_t t1 = ggml_time_ms(); - LOG_INFO("get_learned_condition completed, taking %.2fs", (t1 - prepare_start_ms) * 1.0f / 1000); - - ImageGenerationEmbeds embeds; - embeds.img_uncond = std::move(img_uncond); - embeds.cond = std::move(cond); - embeds.uncond = std::move(uncond); - - return embeds; -} - -static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx, - const GenerationRequest& request, - const std::vector>& final_latents, - int* num_images_out) { - if (final_latents.empty()) { - LOG_ERROR("no latent images to decode"); - return nullptr; - } - if (final_latents.size() > static_cast(request.batch_count)) { - LOG_ERROR("expected at most %d latents, got %zu", request.batch_count, final_latents.size()); - return nullptr; - } - if (final_latents.size() < static_cast(request.batch_count)) { - LOG_INFO("decoding %zu/%d latents", final_latents.size(), request.batch_count); - } else { - LOG_INFO("decoding %zu latents", final_latents.size()); - } - std::vector> decoded_images; - int64_t t0 = ggml_time_ms(); - bool cancelled = false; - - for (size_t i = 0; i < final_latents.size(); i++) { - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling latent decodings"); - cancelled = true; - break; - } - int64_t t1 = ggml_time_ms(); - if (sd_ctx->sd->version == VERSION_QWEN_IMAGE_LAYERED) { - int qwen_image_latent_layers = request.qwen_image_layers + 1; - if (final_latents[i].dim() < 5 || final_latents[i].shape()[2] < qwen_image_latent_layers) { - LOG_ERROR("qwen image layered expected at least %d latent layers, got shape dim=%d", - qwen_image_latent_layers, - final_latents[i].dim()); - return nullptr; - } - for (int layer_index = 0; layer_index < qwen_image_latent_layers; layer_index++) { - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling latent decodings"); - cancelled = true; - break; - } - sd::Tensor layer_latent = sd::ops::slice(final_latents[i], 2, layer_index, layer_index + 1); - layer_latent.squeeze_(2); - sd::Tensor image = sd_ctx->sd->decode_first_stage(layer_latent); - if (image.empty()) { - LOG_ERROR("decode_first_stage failed for latent %zu layer %d", i + 1, layer_index + 1); - return nullptr; - } - decoded_images.push_back(std::move(image)); - } - if (cancelled) { - break; - } - } else if (sd_ctx->sd->animatediff_num_frames > 1 && - final_latents[i].dim() >= 4 && - final_latents[i].shape()[3] == sd_ctx->sd->animatediff_num_frames) { - int n_frames = sd_ctx->sd->animatediff_num_frames; - for (int f = 0; f < n_frames; ++f) { - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling latent decodings"); - cancelled = true; - break; - } - sd::Tensor frame_latent = sd::ops::slice(final_latents[i], 3, f, f + 1); - sd::Tensor image = sd_ctx->sd->decode_first_stage(frame_latent); - if (image.empty()) { - LOG_ERROR("decode_first_stage failed for AnimateDiff frame %d/%d", f + 1, n_frames); - return nullptr; - } - decoded_images.push_back(std::move(image)); - } - } else { - sd::Tensor image = sd_ctx->sd->decode_first_stage(final_latents[i]); - if (image.empty()) { - LOG_ERROR("decode_first_stage failed for latent %" PRId64, i + 1); - return nullptr; - } - decoded_images.push_back(std::move(image)); - } - int64_t t2 = ggml_time_ms(); - LOG_INFO("latent %zu decoded, taking %.2fs", i + 1, (t2 - t1) * 1.0f / 1000); - } - - int64_t t4 = ggml_time_ms(); - LOG_INFO("decode_first_stage completed, taking %.2fs", (t4 - t0) * 1.0f / 1000); - if (decoded_images.empty()) { - LOG_ERROR(cancelled ? "cancelled before any latent images were decoded" : "no decoded images"); - return nullptr; - } - - int image_count = static_cast(decoded_images.size()); - sd_image_t* result_images = (sd_image_t*)calloc(image_count, sizeof(sd_image_t)); - if (result_images == nullptr) { - return nullptr; - } - if (num_images_out != nullptr) { - *num_images_out = image_count; - } - - for (size_t i = 0; i < decoded_images.size(); i++) { - result_images[i] = tensor_to_sd_image(decoded_images[i]); - } - - return result_images; -} - -static sd::Tensor upscale_hires_latent(sd_ctx_t* sd_ctx, - const sd::Tensor& latent, - const GenerationRequest& request, - UpscalerGGML* upscaler) { - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling hires latent upscale"); - return {}; - } - - auto get_hires_latent_target_shape = [&]() { - std::vector target_shape = latent.shape(); - if (target_shape.size() < 2) { - target_shape.clear(); - return target_shape; - } - target_shape[0] = request.hires.target_width / request.vae_scale_factor; - target_shape[1] = request.hires.target_height / request.vae_scale_factor; - return target_shape; - }; - - if (request.hires.upscaler == SD_HIRES_UPSCALER_LATENT || - request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_NEAREST || - request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_NEAREST_EXACT || - request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_ANTIALIASED || - request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_BICUBIC || - request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_BICUBIC_ANTIALIASED) { - std::vector target_shape = get_hires_latent_target_shape(); - if (target_shape.empty()) { - LOG_ERROR("latent has invalid shape for hires upscale"); - return {}; - } - - sd::ops::InterpolateMode mode = sd::ops::InterpolateMode::Nearest; - bool antialias = false; - switch (request.hires.upscaler) { - case SD_HIRES_UPSCALER_LATENT: - mode = sd::ops::InterpolateMode::Bilinear; - break; - case SD_HIRES_UPSCALER_LATENT_NEAREST: - mode = sd::ops::InterpolateMode::Nearest; - break; - case SD_HIRES_UPSCALER_LATENT_NEAREST_EXACT: - mode = sd::ops::InterpolateMode::NearestExact; - break; - case SD_HIRES_UPSCALER_LATENT_ANTIALIASED: - mode = sd::ops::InterpolateMode::Bilinear; - antialias = true; - break; - case SD_HIRES_UPSCALER_LATENT_BICUBIC: - mode = sd::ops::InterpolateMode::Bicubic; - break; - case SD_HIRES_UPSCALER_LATENT_BICUBIC_ANTIALIASED: - mode = sd::ops::InterpolateMode::Bicubic; - antialias = true; - break; - default: - break; - } - - LOG_INFO("hires %s upscale %" PRId64 "x%" PRId64 " -> %" PRId64 "x%" PRId64, - sd_hires_upscaler_name(request.hires.upscaler), - latent.shape()[0], - latent.shape()[1], - target_shape[0], - target_shape[1]); - - return sd::ops::interpolate(latent, target_shape, mode, false, antialias); - } else if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL || - request.hires.upscaler == SD_HIRES_UPSCALER_LANCZOS || - request.hires.upscaler == SD_HIRES_UPSCALER_NEAREST) { - if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL && upscaler == nullptr) { - LOG_ERROR("hires model upscaler context is null"); - return {}; - } - - sd::Tensor decoded = sd_ctx->sd->decode_first_stage(latent); - if (decoded.empty()) { - LOG_ERROR("decode_first_stage failed before hires %s upscale", - sd_hires_upscaler_name(request.hires.upscaler)); - return {}; - } - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling hires image upscale"); - return {}; - } - - sd::Tensor upscaled_tensor; - if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) { - upscaled_tensor = upscaler->upscale_tensor(decoded); - if (upscaled_tensor.empty()) { - LOG_ERROR("hires model upscale failed"); - return {}; - } - - if (upscaled_tensor.shape()[0] != request.hires.target_width || - upscaled_tensor.shape()[1] != request.hires.target_height) { - upscaled_tensor = sd::ops::interpolate(upscaled_tensor, - {request.hires.target_width, - request.hires.target_height, - upscaled_tensor.shape()[2], - upscaled_tensor.shape()[3]}); - } - } else { - sd::ops::InterpolateMode mode = request.hires.upscaler == SD_HIRES_UPSCALER_LANCZOS - ? sd::ops::InterpolateMode::Lanczos - : sd::ops::InterpolateMode::Nearest; - LOG_INFO("hires %s image upscale %" PRId64 "x%" PRId64 " -> %dx%d", - sd_hires_upscaler_name(request.hires.upscaler), - decoded.shape()[0], - decoded.shape()[1], - request.hires.target_width, - request.hires.target_height); - upscaled_tensor = sd::ops::interpolate(decoded, - {request.hires.target_width, - request.hires.target_height, - decoded.shape()[2], - decoded.shape()[3]}, - mode); - upscaled_tensor = sd::ops::clamp(upscaled_tensor, 0.0f, 1.0f); - } - - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling hires latent encode"); - return {}; - } - sd::Tensor upscaled_latent = sd_ctx->sd->encode_first_stage(upscaled_tensor); - if (upscaled_latent.empty()) { - LOG_ERROR("encode_first_stage failed after hires %s upscale", - sd_hires_upscaler_name(request.hires.upscaler)); - } - return upscaled_latent; - } - - LOG_ERROR("unsupported hires upscaler '%s'", sd_hires_upscaler_name(request.hires.upscaler)); - return {}; -} - -static std::vector make_hires_sigma_schedule(sd_ctx_t* sd_ctx, - const sd_hires_params_t& hires, - const sd_sample_params_t& sample_params, - sample_method_t sample_method, - int default_steps, - int sample_seq_len, - int* scheduler_steps_out) { - if (scheduler_steps_out != nullptr) { - *scheduler_steps_out = 0; - } - - if (hires.custom_sigmas_count > 0 && hires.custom_sigmas != nullptr) { - std::vector custom_sigmas(hires.custom_sigmas, - hires.custom_sigmas + hires.custom_sigmas_count); - if (scheduler_steps_out != nullptr) { - *scheduler_steps_out = static_cast(custom_sigmas.size()) - 1; - } - return custom_sigmas; - } - - int effective_steps = hires.steps > 0 ? hires.steps : default_steps; - effective_steps = std::max(1, effective_steps); - - // sd-webui behavior: scale up total steps so trimming by denoising_strength yields exactly hires_steps effective steps, - // unlike img2img which trims from a fixed step count. - int scheduler_steps = static_cast(effective_steps / hires.denoising_strength); - scheduler_steps = std::max(1, scheduler_steps); - - scheduler_t scheduler = resolve_scheduler(sd_ctx, - sample_params.scheduler, - sample_method); - std::vector sigmas = sd_ctx->sd->denoiser->get_sigmas(scheduler_steps, - sample_seq_len, - scheduler, - sd_ctx->sd->version, - sample_params.extra_sample_args); - size_t t_enc = static_cast(scheduler_steps * hires.denoising_strength); - if (t_enc >= static_cast(scheduler_steps)) { - t_enc = static_cast(scheduler_steps) - 1; - } - if (scheduler_steps_out != nullptr) { - *scheduler_steps_out = scheduler_steps; - } - return std::vector(sigmas.begin() + scheduler_steps - static_cast(t_enc) - 1, - sigmas.end()); -} - -static bool generate_image_impl(sd_ctx_t* sd_ctx, - const sd_img_gen_params_t* params, - sd_image_t** images_out, - int* num_images_out); - -SD_API bool generate_image(sd_ctx_t* sd_ctx, - const sd_img_gen_params_t* params, - sd_image_t** images_out, - int* num_images_out) { - if (images_out != nullptr) - *images_out = nullptr; - if (num_images_out != nullptr) - *num_images_out = 0; - if (sd_ctx == nullptr || sd_ctx->sd == nullptr || params == nullptr) { - return false; - } - StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd); - return execution.ready && generate_image_impl(sd_ctx, params, images_out, num_images_out); -} - -static bool generate_image_impl(sd_ctx_t* sd_ctx, - const sd_img_gen_params_t* sd_img_gen_params, - sd_image_t** images_out, - int* num_images_out) { - if (images_out != nullptr) { - *images_out = nullptr; - } - if (num_images_out != nullptr) { - *num_images_out = 0; - } - if (sd_ctx == nullptr || sd_img_gen_params == nullptr) { - return false; - } - - // MiniMax-H3 is video-only. Its denoiser always splits the packed latent into a video and an - // audio half, and only generate_video ever computes the audio length, so reaching this - // function with an H3 checkpoint is guaranteed to die on - // GGML_ASSERT(!audio_input_cache.empty()) with a core dump, after the several minutes it - // takes to load the weights, and with nothing in the output pointing at the missing --mode. - // (The AnimateDiff path below routes vid_gen back through here, but that is SD1.5 plus a - // motion module, never H3.) - if (sd_version_is_minimax_h3(sd_ctx->sd->version)) { - LOG_ERROR("MiniMax-H3 is a video model and cannot be run in img_gen mode; use --mode vid_gen"); - return false; - } - - sd_ctx->sd->reset_cancel_flag(); - - int64_t t0 = ggml_time_ms(); - sd_ctx->sd->vae_tiling_params = sd_img_gen_params->vae_tiling_params; - GenerationRequest request(sd_ctx, sd_img_gen_params); - LOG_INFO("generate_image %dx%d", request.width, request.height); - - sd_ctx->sd->rng->manual_seed(request.seed); - sd_ctx->sd->sampler_rng->manual_seed(request.seed); - sd_ctx->sd->set_flow_shift(sd_img_gen_params->sample_params.flow_shift); - if (!sd_ctx->sd->apply_loras(sd_img_gen_params->loras, sd_img_gen_params->lora_count)) - return false; - apply_circular_axes_to_diffusion(sd_ctx, sd_img_gen_params->circular_x, sd_img_gen_params->circular_y); - - const RefImageParams ref_image_params = sd_ctx->sd->resolve_ref_image_params(sd_img_gen_params->ref_image_args); - - ImageVaeAxesGuard axes_guard(sd_ctx, sd_img_gen_params, request); - - SamplePlan plan(sd_ctx, sd_img_gen_params, request); - auto latents_opt = prepare_image_generation_latents(sd_ctx, - sd_img_gen_params, - &request, - &plan, - ref_image_params); - if (!latents_opt.has_value()) { - return false; - } - ImageGenerationLatents latents = std::move(*latents_opt); - - auto embeds_opt = prepare_image_generation_embeds(sd_ctx, - sd_img_gen_params, - &request, - &plan, - &latents, - ref_image_params); - if (!embeds_opt.has_value()) { - return false; - } - ImageGenerationEmbeds embeds = std::move(*embeds_opt); - - std::vector> final_latents; - int64_t denoise_start = ggml_time_ms(); - for (int b = 0; b < request.batch_count; b++) { - sd_cancel_mode_t cancel = sd_ctx->sd->get_cancel_flag(); - if (cancel == SD_CANCEL_ALL) { - LOG_ERROR("cancelling generation"); - return false; - } - if (cancel == SD_CANCEL_NEW_LATENTS) { - LOG_INFO("cancelling new latent generation, returning %zu/%d completed latents", - final_latents.size(), - request.batch_count); - break; - } - - int64_t sampling_start = ggml_time_ms(); - int64_t cur_seed = request.seed + b; - LOG_INFO("generating image: %i/%i - seed %" PRId64, b + 1, request.batch_count, cur_seed); - - sd_ctx->sd->rng->manual_seed(cur_seed); - sd_ctx->sd->sampler_rng->manual_seed(cur_seed); - sd::Tensor noise = sd::randn_like(latents.init_latent, sd_ctx->sd->rng); - - sd::Tensor x_0 = sd_ctx->sd->sample(sd_ctx->sd->diffusion_model, - true, - latents.init_latent, - std::move(noise), - embeds.cond, - embeds.uncond, - embeds.img_uncond, - latents.control_image, - request.control_strength, - request.guidance, - plan.eta, - request.shifted_timestep, - plan.sample_method, - sd_ctx->sd->is_flow_denoiser(), - plan.extra_sample_args, - plan.sigmas, - latents.ref_latents, - ref_image_params, - latents.denoise_mask, - sd::Tensor(), - 1.f, - 0, - static_cast(request.fps), - request.cache_params, - true); - int64_t sampling_end = ggml_time_ms(); - if (!x_0.empty()) { - LOG_INFO("sampling completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); - final_latents.push_back(std::move(x_0)); - continue; - } - - LOG_ERROR("sampling for image %d/%d failed after %.2fs", - b + 1, - request.batch_count, - (sampling_end - sampling_start) * 1.0f / 1000); - return false; - } - int64_t denoise_end = ggml_time_ms(); - LOG_INFO("generating %zu latent images completed, taking %.2fs", - final_latents.size(), - (denoise_end - denoise_start) * 1.0f / 1000); - if (final_latents.empty()) { - LOG_ERROR("no latent images generated"); - return false; - } - - if (request.hires.enabled && request.hires.target_width > 0) { - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling generation before hires fix"); - return false; - } - LOG_INFO("hires fix: upscaling to %dx%d", request.hires.target_width, request.hires.target_height); - - std::unique_ptr hires_upscaler; - if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) { - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling generation before hires model load"); - return false; - } - LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path); - hires_upscaler = std::make_unique(sd_ctx->sd->n_threads, - false, - request.hires.upscale_tile_size, - sd_ctx->sd->backend_spec, - sd_ctx->sd->params_backend_spec); - const size_t max_graph_vram_bytes = sd_ctx->sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER); - hires_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes); - if (!hires_upscaler->load_from_file(request.hires.model_path, - sd_ctx->sd->n_threads)) { - LOG_ERROR("load hires model upscaler failed"); - return false; - } - } - - int hires_scheduler_steps = 0; - std::vector hires_sigma_sched = - make_hires_sigma_schedule(sd_ctx, - request.hires, - sd_img_gen_params->sample_params, - plan.sample_method, - plan.sample_steps, - sd_ctx->sd->get_image_seq_len(request.hires.target_height, request.hires.target_width), - &hires_scheduler_steps); - LOG_INFO("hires fix: scheduler_steps=%d, denoising_strength=%.2f, sigma_sched_size=%zu%s", - hires_scheduler_steps, - request.hires.denoising_strength, - hires_sigma_sched.size(), - request.hires.custom_sigmas_count > 0 ? ", custom_sigmas=true" : ""); - - std::vector> hires_final_latents; - int64_t hires_denoise_start = ggml_time_ms(); - for (int b = 0; b < (int)final_latents.size(); b++) { - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling generation during hires fix"); - return false; - } - int64_t cur_seed = request.seed + b; - sd_ctx->sd->rng->manual_seed(cur_seed); - sd_ctx->sd->sampler_rng->manual_seed(cur_seed); - - sd::Tensor upscaled = upscale_hires_latent(sd_ctx, - final_latents[b], - request, - hires_upscaler.get()); - if (upscaled.empty()) { - return false; - } - - sd::Tensor noise = sd::randn_like(upscaled, sd_ctx->sd->rng); - - sd::Tensor hires_denoise_mask; - if (!latents.denoise_mask.empty()) { - std::vector mask_shape = latents.denoise_mask.shape(); - mask_shape[0] = upscaled.shape()[0]; - mask_shape[1] = upscaled.shape()[1]; - hires_denoise_mask = sd::ops::interpolate(latents.denoise_mask, - mask_shape, - sd::ops::InterpolateMode::NearestMax); - } - - int64_t hires_sample_start = ggml_time_ms(); - sd::Tensor x_0 = sd_ctx->sd->sample(sd_ctx->sd->diffusion_model, - true, - upscaled, - std::move(noise), - embeds.cond, - embeds.uncond, - embeds.img_uncond, - latents.control_image, - request.control_strength, - request.guidance, - plan.eta, - request.shifted_timestep, - plan.sample_method, - sd_ctx->sd->is_flow_denoiser(), - plan.extra_sample_args, - hires_sigma_sched, - latents.ref_latents, - ref_image_params, - hires_denoise_mask, - sd::Tensor(), - 1.f, - 0, - static_cast(request.fps), - request.cache_params, - false); - int64_t hires_sample_end = ggml_time_ms(); - if (!x_0.empty()) { - LOG_INFO("hires sampling %d/%d completed, taking %.2fs", - b + 1, - (int)final_latents.size(), - (hires_sample_end - hires_sample_start) * 1.0f / 1000); - hires_final_latents.push_back(std::move(x_0)); - continue; - } - - LOG_ERROR("hires sampling for image %d/%d failed after %.2fs", - b + 1, - (int)final_latents.size(), - (hires_sample_end - hires_sample_start) * 1.0f / 1000); - return false; - } - int64_t hires_denoise_end = ggml_time_ms(); - LOG_INFO("hires fix completed, taking %.2fs", (hires_denoise_end - hires_denoise_start) * 1.0f / 1000); - - final_latents = std::move(hires_final_latents); - } - - int num_images = 0; - auto result = decode_image_outputs(sd_ctx, request, final_latents, &num_images); - if (result == nullptr) { - return false; - } - - sd_ctx->sd->lora_stat(); - - int64_t t1 = ggml_time_ms(); - LOG_INFO("generate_image completed in %.2fs", (t1 - t0) * 1.0f / 1000); - if (num_images_out != nullptr) { - *num_images_out = num_images; - } - if (images_out != nullptr) { - *images_out = result; - } else { - free_sd_images(result, num_images); +const char* sd_vae_format_name(enum sd_vae_format_t format) { + switch (format) { + case SD_VAE_FORMAT_AUTO: + return "auto"; + case SD_VAE_FORMAT_FLUX: + return "flux"; + case SD_VAE_FORMAT_SD3: + return "sd3"; + case SD_VAE_FORMAT_FLUX2: + return "flux2"; + case SD_VAE_FORMAT_WAN: + return "wan"; + default: + return NONE_STR; } - return true; } -static std::optional prepare_video_generation_latents(sd_ctx_t* sd_ctx, - const sd_vid_gen_params_t* sd_vid_gen_params, - GenerationRequest* request) { - ImageGenerationLatents latents; - int64_t prepare_start_ms = ggml_time_ms(); - - sd::Tensor start_image; - sd::Tensor end_image; - - if (sd_vid_gen_params->init_image.data) { - start_image = sd_image_to_tensor(sd_vid_gen_params->init_image, request->width, request->height); - } - - if (sd_vid_gen_params->end_image.data) { - end_image = sd_image_to_tensor(sd_vid_gen_params->end_image, request->width, request->height); - } - - if (sd_version_is_minimax_h3(sd_ctx->sd->version)) { - if (sd_vid_gen_params->ref_images_count < 0 || sd_vid_gen_params->ref_videos_count < 0 || - sd_vid_gen_params->ref_audios_count < 0 || - (sd_vid_gen_params->ref_images_count > 0 && sd_vid_gen_params->ref_images == nullptr) || - (sd_vid_gen_params->ref_videos_count > 0 && sd_vid_gen_params->ref_videos == nullptr) || - (sd_vid_gen_params->ref_audios_count > 0 && sd_vid_gen_params->ref_audios == nullptr)) { - LOG_ERROR("invalid MiniMax-H3 Ref2VA input arrays"); - return std::nullopt; - } - - latents.audio_length = get_minimax_h3_num_audio_latents(request->frames, - request->fps); - latents.audio_latent = make_minimax_h3_empty_audio_latent(latents.audio_length); - - bool has_references = sd_vid_gen_params->ref_images_count > 0 || - sd_vid_gen_params->ref_videos_count > 0 || - sd_vid_gen_params->ref_audios_count > 0; - if (has_references && (!start_image.empty() || !end_image.empty())) { - LOG_ERROR("MiniMax-H3 keyframes and Ref2VA references cannot be used together"); - return std::nullopt; - } - - if (sd_vid_gen_params->control_frames_size > 0) { - LOG_ERROR("MiniMax-H3 control_frames are not implemented"); - return std::nullopt; - } - - auto add_visual_noise = [&](sd::Tensor latent) { - auto condition_rng = std::make_shared(); - condition_rng->manual_seed(static_cast(request->seed)); - return latent * MiniMaxH3::VISUAL_COND_TIMESTEP + - sd::Tensor::randn_like(latent, condition_rng) * - (1.f - MiniMaxH3::VISUAL_COND_TIMESTEP); - }; - - auto add_keyframe = [&](const sd::Tensor& image, - int32_t frame_index, - const char* name) -> bool { - if (image.empty()) { - return true; - } - auto video_image = image.reshape({image.shape()[0], - image.shape()[1], - 1, - image.shape()[2], - image.shape()[3]}); - auto latent = sd_ctx->sd->encode_first_stage(video_image); - if (latent.empty()) { - LOG_ERROR("failed to encode MiniMax-H3 %s keyframe", name); - return false; - } - latents.ref_images.push_back(image); - latents.ref_latents.push_back(add_visual_noise(std::move(latent))); - latents.keyframe_indices.push_back(frame_index); - return true; - }; - - auto resize_reference = [&](const sd::Tensor& image, - int width, - int height) { - return sd::ops::interpolate( - image, - std::vector{width, height, image.shape()[2], image.shape()[3]}); - }; - - auto encode_reference_audio = [&](const sd_audio_t& audio, - int32_t* audio_index) -> bool { - if (sd_ctx->sd->audio_vae_model == nullptr) { - LOG_ERROR("MiniMax-H3 Ref2VA audio requires --audio-vae with encoder weights"); - return false; - } - auto waveform = prepare_minimax_h3_reference_waveform( - audio, - sd_ctx->sd->audio_vae_model->input_sample_rate()); - if (waveform.empty()) { - LOG_ERROR("invalid MiniMax-H3 reference audio"); - return false; - } - auto encoded = sd_ctx->sd->audio_vae_model->encode(sd_ctx->sd->n_threads, waveform); - if (encoded.empty()) { - LOG_ERROR("failed to encode MiniMax-H3 reference audio"); - return false; - } - *audio_index = static_cast(latents.reference_audio_latents.size()); - latents.reference_audio_latents.push_back(std::move(encoded)); - return true; - }; - - if (has_references) { - LOG_INFO("MiniMax-H3 Ref2VA: %d image(s), %d video(s), %d audio clip(s)", - sd_vid_gen_params->ref_images_count, - sd_vid_gen_params->ref_videos_count, - sd_vid_gen_params->ref_audios_count); - - for (int i = 0; i < sd_vid_gen_params->ref_images_count; ++i) { - auto image = ensure_image_tensor_channels( - sd_image_to_tensor(sd_vid_gen_params->ref_images[i]), - 3); - if (image.empty()) { - LOG_ERROR("failed to load MiniMax-H3 reference image %d", i + 1); - return std::nullopt; - } - int source_w = static_cast(image.shape()[0]); - int source_h = static_cast(image.shape()[1]); - double source_area = static_cast(source_w) * source_h; - double target_area = static_cast(request->width) * request->height; - double scale = std::min(1.0, std::sqrt(target_area / source_area)); - int width = std::max(32, static_cast(std::round(source_w * scale / 32.f)) * 32); - int height = std::max(32, static_cast(std::round(source_h * scale / 32.f)) * 32); - image = resize_reference(image, width, height); - auto latent = sd_ctx->sd->encode_first_stage(image); - if (latent.empty()) { - LOG_ERROR("failed to encode MiniMax-H3 reference image %d", i + 1); - return std::nullopt; - } - int32_t video_index = static_cast(latents.ref_latents.size()); - latents.ref_latents.push_back(add_visual_noise(std::move(latent))); - latents.minimax_reference_blocks.push_back({MiniMaxH3ReferenceKind::IMAGE, - video_index, - -1}); - MiniMaxH3PresentationItem item; - item.kind = MiniMaxH3PresentationKind::IMAGE; - item.frames.push_back(std::move(image)); - latents.minimax_presentation_refs.push_back(std::move(item)); - } - - for (int video_idx = 0; video_idx < sd_vid_gen_params->ref_videos_count; ++video_idx) { - const auto& reference = sd_vid_gen_params->ref_videos[video_idx]; - if (reference.frames == nullptr || reference.frame_count < 1) { - LOG_ERROR("invalid MiniMax-H3 reference video %d", video_idx + 1); - return std::nullopt; - } - int source_fps = reference.fps > 0 ? reference.fps : 24; - int normalized_frames = static_cast(std::lround( - static_cast(reference.frame_count) * 24.0 / source_fps)); - normalized_frames = std::min(normalized_frames, request->frames); - if (normalized_frames < 5) { - LOG_ERROR("MiniMax-H3 reference video %d needs at least 5 frames at 24 fps", - video_idx + 1); - return std::nullopt; - } - while (normalized_frames % 17 != 5) { - --normalized_frames; - } - - auto first = ensure_image_tensor_channels(sd_image_to_tensor(reference.frames[0]), 3); - if (first.empty()) { - LOG_ERROR("invalid first frame in MiniMax-H3 reference video %d", video_idx + 1); - return std::nullopt; - } - int source_w = static_cast(first.shape()[0]); - int source_h = static_cast(first.shape()[1]); - double ratio = static_cast(source_w) / source_h; - double nominal_w = ratio >= 1.0 ? 768.0 * ratio : 768.0; - double nominal_h = ratio >= 1.0 ? 768.0 : 768.0 / ratio; - if (nominal_w * nominal_h > 768.0 * 1344.0) { - double scale = std::sqrt((768.0 * 1344.0) / (nominal_w * nominal_h)); - nominal_w *= scale; - nominal_h *= scale; - } - int width = std::max(32, static_cast(std::round(nominal_w / 32.0)) * 32); - int height = std::max(32, static_cast(std::round(nominal_h / 32.0)) * 32); - if (source_w * source_h < width * height) { - width = std::max(32, static_cast(std::round(source_w / 32.0)) * 32); - height = std::max(32, static_cast(std::round(source_h / 32.0)) * 32); - } - - sd::Tensor video({width, height, normalized_frames, 3, 1}); - for (int frame = 0; frame < normalized_frames; ++frame) { - int source_index = std::min(reference.frame_count - 1, - static_cast(std::floor(frame * source_fps / 24.0))); - auto source = ensure_image_tensor_channels( - sd_image_to_tensor(reference.frames[source_index]), - 3); - if (source.empty()) { - LOG_ERROR("invalid frame %d in MiniMax-H3 reference video %d", - source_index + 1, - video_idx + 1); - return std::nullopt; - } - source = resize_reference(source, width, height); - sd::ops::slice_assign(&video, 2, frame, frame + 1, source.unsqueeze(2)); - } - auto video_latent = sd_ctx->sd->encode_first_stage(video); - if (video_latent.empty()) { - LOG_ERROR("failed to encode MiniMax-H3 reference video %d", video_idx + 1); - return std::nullopt; - } - int32_t audio_index = -1; - bool has_audio = reference.audio.data != nullptr && reference.audio.sample_count > 0; - if (has_audio) { - if (!encode_reference_audio(reference.audio, &audio_index)) { - return std::nullopt; - } - MiniMaxH3PresentationItem audio_item; - audio_item.kind = MiniMaxH3PresentationKind::AUDIO; - latents.minimax_presentation_refs.push_back(std::move(audio_item)); - } - - MiniMaxH3PresentationItem video_item; - video_item.kind = MiniMaxH3PresentationKind::VIDEO; - for (int frame = 0; frame < normalized_frames; frame += 12) { - auto sampled = sd::ops::slice(video, 2, frame, frame + 1) - .reshape({width, height, 3, 1}); - video_item.frames.push_back(std::move(sampled)); - video_item.timestamps.push_back(frame / 24.f); - } - latents.minimax_presentation_refs.push_back(std::move(video_item)); - - int32_t video_index = static_cast(latents.ref_latents.size()); - latents.ref_latents.push_back(add_visual_noise(std::move(video_latent))); - latents.minimax_reference_blocks.push_back({has_audio ? MiniMaxH3ReferenceKind::VIDEO_AUDIO - : MiniMaxH3ReferenceKind::VIDEO, - video_index, - audio_index}); - } - - for (int audio_idx = 0; audio_idx < sd_vid_gen_params->ref_audios_count; ++audio_idx) { - int32_t encoded_index = -1; - if (!encode_reference_audio(sd_vid_gen_params->ref_audios[audio_idx], &encoded_index)) { - return std::nullopt; - } - MiniMaxH3PresentationItem item; - item.kind = MiniMaxH3PresentationKind::AUDIO; - latents.minimax_presentation_refs.push_back(std::move(item)); - latents.minimax_reference_blocks.push_back({MiniMaxH3ReferenceKind::AUDIO, - -1, - encoded_index}); - } - } - - if (!has_references && (!start_image.empty() || !end_image.empty())) { - LOG_INFO(!start_image.empty() && !end_image.empty() ? "MiniMax-H3 FL2VA" : !start_image.empty() ? "MiniMax-H3 I2VA" - : "MiniMax-H3 end-frame conditioning"); - } - if (!has_references && - (!add_keyframe(start_image, 0, "start") || - !add_keyframe(end_image, request->frames - 1, "end"))) { - return std::nullopt; - } - } - - if (sd_version_is_ltxav(sd_ctx->sd->version)) { - latents.audio_length = get_ltxav_num_audio_latents(request->frames, request->fps); - latents.audio_latent = make_ltxav_empty_audio_latent(latents.audio_length); - } - - if (sd_version_is_ltxav(sd_ctx->sd->version)) { - if (sd_vid_gen_params->control_frames_size > 0) { - LOG_ERROR("LTXAV control_frames are not implemented"); - return std::nullopt; - } - - if (!start_image.empty() || !end_image.empty()) { - if (!start_image.empty() && !end_image.empty()) { - LOG_INFO("FLF2V"); - } else if (!start_image.empty()) { - LOG_INFO("IMG2VID"); - } else { - LOG_INFO("END2VID"); - } - - int64_t t1 = ggml_time_ms(); - latents.init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height, request->frames, true); - - float conditioning_strength = std::clamp(request->strength, 0.f, 1.f); - float conditioned_mask = 1.0f - conditioning_strength; - latents.denoise_mask = make_ltxav_video_denoise_mask(latents.init_latent, 1.f); - - auto apply_video_condition_by_keyframe_index = [&](const sd::Tensor& keyframes, - int frame_idx, - const char* name) -> bool { - int64_t keyframe_frames = keyframes.shape()[2]; - if (keyframe_frames <= 0 || keyframes.shape()[0] != latents.init_latent.shape()[0] || - keyframes.shape()[1] != latents.init_latent.shape()[1] || - keyframes.shape()[3] != latents.init_latent.shape()[3]) { - LOG_ERROR("invalid LTXAV %s keyframe latent shape", name); - return false; - } - - latents.video_target_frame_count = latents.init_latent.shape()[2]; - latents.video_conditioning_frame_count = keyframe_frames; - latents.init_latent = sd::ops::concat(latents.init_latent, keyframes, 2); - - auto keyframe_mask = sd::full({keyframes.shape()[0], - keyframes.shape()[1], - keyframes.shape()[2], - 1, - 1}, - conditioned_mask); - latents.denoise_mask = sd::ops::concat(latents.denoise_mask, keyframe_mask, 2); - latents.video_positions = build_ltxv_video_positions(latents.init_latent.shape()[0], - latents.init_latent.shape()[1], - latents.video_target_frame_count, - keyframe_frames, - frame_idx, - 1, - request->fps, - request->vae_scale_factor, - 8, - true); - return true; - }; - - if (!start_image.empty()) { - if (!apply_ltxav_condition_image_by_latent_index(sd_ctx, - start_image, - &latents.init_latent, - &latents.denoise_mask, - 0, - "init", - conditioning_strength)) { - return std::nullopt; - } - } - - if (!end_image.empty()) { - auto end_image_latent = encode_ltxav_condition_image(sd_ctx, end_image, "end"); - if (end_image_latent.empty()) { - return std::nullopt; - } - - int frame_idx = request->frames - 1; - bool ok = frame_idx == 0 ? apply_ltxav_condition_by_latent_index(&latents.init_latent, - &latents.denoise_mask, - end_image_latent, - 0, - "end", - conditioned_mask) - : apply_video_condition_by_keyframe_index(end_image_latent, frame_idx, "end"); - if (!ok) { - return std::nullopt; - } - } - - int64_t t2 = ggml_time_ms(); - LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); - } - } - - if (sd_version_is_hunyuan_video(sd_ctx->sd->version) && - (!start_image.empty() || !end_image.empty())) { - LOG_INFO("Hunyuan Video IMG2VID"); - - int64_t t1 = ggml_time_ms(); - auto concat_latent = sd_ctx->sd->generate_init_latent(request->width, - request->height, - request->frames, - true); - auto encode_condition_frame = [&](const sd::Tensor& image, - int64_t latent_frame, - const char* name) -> bool { - auto encoded = sd_ctx->sd->encode_first_stage(image.unsqueeze(2)); - if (encoded.empty()) { - LOG_ERROR("failed to encode Hunyuan Video %s conditioning frame", name); - return false; - } - if (encoded.dim() == 4) { - encoded.unsqueeze_(2); - } - if (encoded.dim() != 5 || - encoded.shape()[0] != concat_latent.shape()[0] || - encoded.shape()[1] != concat_latent.shape()[1] || - encoded.shape()[3] != concat_latent.shape()[3]) { - LOG_ERROR("invalid Hunyuan Video %s conditioning latent shape", name); - return false; - } - sd::ops::slice_assign(&concat_latent, - 2, - latent_frame, - latent_frame + 1, - sd::ops::slice(encoded, 2, 0, 1)); - return true; - }; - - if (!start_image.empty() && !encode_condition_frame(start_image, 0, "start")) { - return std::nullopt; - } - if (!end_image.empty() && - !encode_condition_frame(end_image, concat_latent.shape()[2] - 1, "end")) { - return std::nullopt; - } - - sd::Tensor concat_mask = sd::zeros({concat_latent.shape()[0], - concat_latent.shape()[1], - concat_latent.shape()[2], - 1, - 1}); - if (!start_image.empty()) { - sd::ops::fill_slice(&concat_mask, 2, 0, 1, 1.0f); - } - if (!end_image.empty()) { - sd::ops::fill_slice(&concat_mask, 2, concat_mask.shape()[2] - 1, concat_mask.shape()[2], 1.0f); - } - latents.concat_latent = sd::ops::concat(concat_latent, concat_mask, 3); - - int64_t t2 = ggml_time_ms(); - LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); - } - - if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-I2V-14B" || - sd_ctx->sd->diffusion_model->get_desc() == "Wan2.2-I2V-14B" || - sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-I2V-1.3B" || - sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-FLF2V-14B") { - LOG_INFO("IMG2VID"); - - if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-I2V-14B" || - sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-I2V-1.3B" || - sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-FLF2V-14B") { - if (!start_image.empty()) { - auto clip_vision_output = sd_ctx->sd->get_clip_vision_output(start_image, false, -2); - if (clip_vision_output.empty()) { - LOG_ERROR("failed to compute clip vision output for init image"); - return std::nullopt; - } - latents.clip_vision_output = std::move(clip_vision_output); - } else { - latents.clip_vision_output = sd_ctx->sd->get_clip_vision_output(start_image, false, -2, true); - } - - if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-FLF2V-14B") { - sd::Tensor end_image_clip_vision_output; - if (!end_image.empty()) { - end_image_clip_vision_output = sd_ctx->sd->get_clip_vision_output(end_image, false, -2); - if (end_image_clip_vision_output.empty()) { - LOG_ERROR("failed to compute clip vision output for end image"); - return std::nullopt; - } - } else { - end_image_clip_vision_output = sd_ctx->sd->get_clip_vision_output(end_image, false, -2, true); - } - latents.clip_vision_output = sd::ops::concat(latents.clip_vision_output, end_image_clip_vision_output, 1); - } - - int64_t t1 = ggml_time_ms(); - LOG_INFO("get_clip_vision_output completed, taking %" PRId64 " ms", t1 - prepare_start_ms); - } - - int64_t t1 = ggml_time_ms(); - sd::Tensor image = sd::full({request->width, request->height, request->frames, 3, 1}, 0.5f); - if (!start_image.empty()) { - sd::ops::slice_assign(&image, 2, 0, 1, start_image.unsqueeze(2)); - } - if (!end_image.empty()) { - sd::ops::slice_assign(&image, 2, request->frames - 1, request->frames, end_image.unsqueeze(2)); - } - - auto concat_latent = sd_ctx->sd->encode_first_stage(image); // [b, c, t, h/vae_scale_factor, w/vae_scale_factor] - if (concat_latent.empty()) { - LOG_ERROR("failed to encode video conditioning frames"); - return std::nullopt; - } - latents.concat_latent = std::move(concat_latent); - - int64_t t2 = ggml_time_ms(); - LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); - - sd::Tensor concat_mask = sd::zeros({latents.concat_latent.shape()[0], - latents.concat_latent.shape()[1], - latents.concat_latent.shape()[2], - 4, - 1}); // [b, 4, t, h/vae_scale_factor, w/vae_scale_factor] - if (!start_image.empty()) { - sd::ops::fill_slice(&concat_mask, 2, 0, 1, 1.0f); - } - if (!end_image.empty()) { - auto last_channel = sd::ops::slice(concat_mask, 3, 3, 4); - sd::ops::fill_slice(&last_channel, 2, last_channel.shape()[2] - 1, last_channel.shape()[2], 1.0f); - sd::ops::slice_assign(&concat_mask, 3, 3, 4, last_channel); - } - latents.concat_latent = sd::ops::concat(concat_mask, latents.concat_latent, 3); // [b, 4+c, t, h/vae_scale_factor, w/vae_scale_factor] - } else if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.2-TI2V-5B" && !start_image.empty()) { - LOG_INFO("IMG2VID"); - - int64_t t1 = ggml_time_ms(); - auto init_img = start_image.reshape({start_image.shape()[0], start_image.shape()[1], 1, start_image.shape()[2], 1}); - auto init_image_latent = sd_ctx->sd->encode_first_stage(init_img); // [b, c, 1, h/vae_scale_factor, w/vae_scale_factor] - if (init_image_latent.empty()) { - LOG_ERROR("failed to encode init video frame"); - return std::nullopt; - } - - latents.init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height, request->frames, true); // [b, c, t, h/vae_scale_factor, w/vae_scale_factor] - sd::ops::slice_assign(&latents.init_latent, 2, 0, init_image_latent.shape()[2], init_image_latent); - - latents.denoise_mask = sd::full({latents.init_latent.shape()[0], latents.init_latent.shape()[1], latents.init_latent.shape()[2], 1, 1}, 1.f); - sd::ops::fill_slice(&latents.denoise_mask, 2, 0, init_image_latent.shape()[2], 0.0f); - - if (!end_image.empty()) { - auto end_img = end_image.reshape({end_image.shape()[0], end_image.shape()[1], 1, end_image.shape()[2], 1}); - auto end_image_latent = sd_ctx->sd->encode_first_stage(end_img); // [b, c, 1, h/vae_scale_factor, w/vae_scale_factor] - if (end_image_latent.empty()) { - LOG_ERROR("failed to encode end video frame"); - return std::nullopt; - } - sd::ops::slice_assign(&latents.init_latent, 2, latents.init_latent.shape()[2] - 1, latents.init_latent.shape()[2], end_image_latent); - sd::ops::fill_slice(&latents.denoise_mask, 2, latents.init_latent.shape()[2] - 1, latents.init_latent.shape()[2], 0.0f); - } - - int64_t t2 = ggml_time_ms(); - LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); - } else if (sd_version_is_lingbot_video(sd_ctx->sd->version) && !start_image.empty()) { - LOG_INFO("LingBot Video IMG2VID"); - - int64_t t1 = ggml_time_ms(); - auto init_img = start_image.reshape({start_image.shape()[0], start_image.shape()[1], 1, start_image.shape()[2], 1}); - auto init_image_latent = sd_ctx->sd->encode_first_stage(init_img); - if (init_image_latent.empty()) { - LOG_ERROR("failed to encode init video frame"); - return std::nullopt; - } - - latents.init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height, request->frames, true); - sd::ops::slice_assign(&latents.init_latent, 2, 0, init_image_latent.shape()[2], init_image_latent); - - latents.denoise_mask = sd::full({latents.init_latent.shape()[0], latents.init_latent.shape()[1], latents.init_latent.shape()[2], 1, 1}, 1.f); - sd::ops::fill_slice(&latents.denoise_mask, 2, 0, init_image_latent.shape()[2], 0.0f); - - latents.ref_images.push_back(start_image); +void sd_cache_params_init(sd_cache_params_t* cache_params) { + *cache_params = {}; + cache_params->mode = SD_CACHE_DISABLED; + cache_params->reuse_threshold = INFINITY; + cache_params->start_percent = 0.15f; + cache_params->end_percent = 0.95f; + cache_params->error_decay_rate = 1.0f; + cache_params->use_relative_threshold = true; + cache_params->reset_error_on_compute = true; + cache_params->Fn_compute_blocks = 8; + cache_params->Bn_compute_blocks = 0; + cache_params->residual_diff_threshold = 0.08f; + cache_params->max_warmup_steps = 8; + cache_params->max_cached_steps = -1; + cache_params->max_continuous_cached_steps = -1; + cache_params->taylorseer_n_derivatives = 1; + cache_params->taylorseer_skip_interval = 1; + cache_params->scm_mask = nullptr; + cache_params->scm_policy_dynamic = true; + cache_params->spectrum_w = 0.40f; + cache_params->spectrum_m = 3; + cache_params->spectrum_lam = 1.0f; + cache_params->spectrum_window_size = 2; + cache_params->spectrum_flex_window = 0.50f; + cache_params->spectrum_warmup_steps = 4; + cache_params->spectrum_stop_percent = 0.9f; +} - int64_t t2 = ggml_time_ms(); - LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); - } else if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-VACE-1.3B" || - sd_ctx->sd->diffusion_model->get_desc() == "Wan2.x-VACE-14B") { - LOG_INFO("VACE"); - int64_t t1 = ggml_time_ms(); - sd::Tensor ref_image_latent; - if (!start_image.empty()) { - auto ref_img = start_image.reshape({start_image.shape()[0], start_image.shape()[1], 1, start_image.shape()[2], 1}); - auto encoded_ref = sd_ctx->sd->encode_first_stage(ref_img); // [b, c, 1, h/vae_scale_factor, w/vae_scale_factor] - if (encoded_ref.empty()) { - LOG_ERROR("failed to encode VACE reference image"); - return std::nullopt; - } - ref_image_latent = sd::ops::concat(encoded_ref, sd::zeros(encoded_ref.shape()), 3); // [b, 2*c, 1, h/vae_scale_factor, w/vae_scale_factor] - } +void sd_hires_params_init(sd_hires_params_t* hires_params) { + *hires_params = {}; + hires_params->enabled = false; + hires_params->upscaler = SD_HIRES_UPSCALER_LATENT; + hires_params->model_path = nullptr; + hires_params->scale = 2.0f; + hires_params->target_width = 0; + hires_params->target_height = 0; + hires_params->steps = 0; + hires_params->denoising_strength = 0.7f; + hires_params->upscale_tile_size = 128; + hires_params->custom_sigmas = nullptr; + hires_params->custom_sigmas_count = 0; +} - sd::Tensor control_video = sd::full({request->width, request->height, request->frames, 3, 1}, 0.5f); - int64_t control_frame_count = std::min(request->frames, sd_vid_gen_params->control_frames_size); - for (int64_t i = 0; i < control_frame_count; ++i) { - auto control_frame = sd_image_to_tensor(sd_vid_gen_params->control_frames[i], request->width, request->height); - sd::ops::slice_assign(&control_video, 2, i, i + 1, control_frame.unsqueeze(2)); - } +void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { + *sd_ctx_params = {}; + sd_ctx_params->n_threads = sd_get_num_physical_cores(); + sd_ctx_params->wtype = SD_TYPE_COUNT; + sd_ctx_params->rng_type = CUDA_RNG; + sd_ctx_params->sampler_rng_type = RNG_TYPE_COUNT; + sd_ctx_params->prediction = PREDICTION_COUNT; + sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO; + sd_ctx_params->max_vram = nullptr; + sd_ctx_params->disable_prefetch = false; + sd_ctx_params->disable_segmented_compute = false; + sd_ctx_params->eager_load = false; + sd_ctx_params->enable_mmap = false; + sd_ctx_params->diffusion_flash_attn = false; + sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO; + sd_ctx_params->backend = nullptr; + sd_ctx_params->params_backend = nullptr; + sd_ctx_params->split_mode = nullptr; + sd_ctx_params->auto_fit = true; + sd_ctx_params->rpc_servers = nullptr; + sd_ctx_params->model_args = nullptr; + sd_ctx_params->pulid_weights_path = nullptr; +} - sd::Tensor mask = sd::full({request->width, request->height, request->frames, 1, 1}, 1.0f); +char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { + char* buf = (char*)malloc(8192); + if (!buf) + return nullptr; + buf[0] = '\0'; - control_video = control_video - 0.5f; - sd::Tensor inactive = control_video * (1.0f - mask) + 0.5f; - sd::Tensor reactive = control_video * mask + 0.5f; + snprintf(buf + strlen(buf), 8192 - strlen(buf), + "model_path: %s\n" + "clip_l_path: %s\n" + "clip_g_path: %s\n" + "clip_vision_path: %s\n" + "t5xxl_path: %s\n" + "llm_path: %s\n" + "llm_vision_path: %s\n" + "diffusion_model_path: %s\n" + "high_noise_diffusion_model_path: %s\n" + "uncond_diffusion_model_path: %s\n" + "embeddings_connectors_path: %s\n" + "vae_path: %s\n" + "audio_vae_path: %s\n" + "taesd_path: %s\n" + "control_net_path: %s\n" + "photo_maker_path: %s\n" + "pulid_weights_path: %s\n" + "tensor_type_rules: %s\n" + "n_threads: %d\n" + "wtype: %s\n" + "rng_type: %s\n" + "sampler_rng_type: %s\n" + "prediction: %s\n" + "max_vram: %s\n" + "disable_prefetch: %s\n" + "disable_segmented_compute: %s\n" + "eager_load: %s\n" + "backend: %s\n" + "params_backend: %s\n" + "split_mode: %s\n" + "model_args: %s\n" + "auto_fit: %s\n" + "flash_attn: %s\n" + "diffusion_flash_attn: %s\n" + "vae_format: %s\n", + SAFE_STR(sd_ctx_params->model_path), + SAFE_STR(sd_ctx_params->clip_l_path), + SAFE_STR(sd_ctx_params->clip_g_path), + SAFE_STR(sd_ctx_params->clip_vision_path), + SAFE_STR(sd_ctx_params->t5xxl_path), + SAFE_STR(sd_ctx_params->llm_path), + SAFE_STR(sd_ctx_params->llm_vision_path), + SAFE_STR(sd_ctx_params->diffusion_model_path), + SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path), + SAFE_STR(sd_ctx_params->uncond_diffusion_model_path), + SAFE_STR(sd_ctx_params->embeddings_connectors_path), + SAFE_STR(sd_ctx_params->vae_path), + SAFE_STR(sd_ctx_params->audio_vae_path), + SAFE_STR(sd_ctx_params->taesd_path), + SAFE_STR(sd_ctx_params->control_net_path), + SAFE_STR(sd_ctx_params->photo_maker_path), + SAFE_STR(sd_ctx_params->pulid_weights_path), + SAFE_STR(sd_ctx_params->tensor_type_rules), + sd_ctx_params->n_threads, + sd_type_name(sd_ctx_params->wtype), + sd_rng_type_name(sd_ctx_params->rng_type), + sd_rng_type_name(sd_ctx_params->sampler_rng_type), + sd_prediction_name(sd_ctx_params->prediction), + SAFE_STR(sd_ctx_params->max_vram), + BOOL_STR(sd_ctx_params->disable_prefetch), + BOOL_STR(sd_ctx_params->disable_segmented_compute), + BOOL_STR(sd_ctx_params->eager_load), + SAFE_STR(sd_ctx_params->backend), + SAFE_STR(sd_ctx_params->params_backend), + SAFE_STR(sd_ctx_params->split_mode), + SAFE_STR(sd_ctx_params->model_args), + BOOL_STR(sd_ctx_params->auto_fit), + BOOL_STR(sd_ctx_params->flash_attn), + BOOL_STR(sd_ctx_params->diffusion_flash_attn), + sd_vae_format_name(sd_ctx_params->vae_format)); - inactive = sd_ctx->sd->encode_first_stage(inactive); // [b, c, t, h/vae_scale_factor, w/vae_scale_factor] - if (inactive.empty()) { - LOG_ERROR("failed to encode VACE inactive context"); - return std::nullopt; - } + return buf; +} - reactive = sd_ctx->sd->encode_first_stage(reactive); // [b, c, t, h/vae_scale_factor, w/vae_scale_factor] - if (reactive.empty()) { - LOG_ERROR("failed to encode VACE reactive context"); - return std::nullopt; - } +void sd_sample_params_init(sd_sample_params_t* sample_params) { + *sample_params = {}; + sample_params->guidance.txt_cfg = 7.0f; + sample_params->guidance.img_cfg = INFINITY; + sample_params->guidance.distilled_guidance = 3.5f; + sample_params->guidance.slg.layer_count = 0; + sample_params->guidance.slg.layer_start = 0.01f; + sample_params->guidance.slg.layer_end = 0.2f; + sample_params->guidance.slg.scale = 0.f; + sample_params->scheduler = SCHEDULER_COUNT; + sample_params->sample_method = SAMPLE_METHOD_COUNT; + sample_params->sample_steps = 20; + sample_params->eta = INFINITY; + sample_params->custom_sigmas = nullptr; + sample_params->custom_sigmas_count = 0; + sample_params->flow_shift = INFINITY; + sample_params->extra_sample_args = nullptr; +} - int64_t length = inactive.shape()[2]; - if (!ref_image_latent.empty()) { - length += 1; - request->frames = static_cast((length - 1) * 4 + 1); - latents.ref_image_num = 1; - } - auto vace_context = sd::ops::concat(inactive, reactive, 3); // [b, 2*c, t, h/vae_scale_factor, w/vae_scale_factor] +char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) { + char* buf = (char*)malloc(4096); + if (!buf) + return nullptr; + buf[0] = '\0'; - mask = sd::full({request->width, request->height, inactive.shape()[2], 1, 1}, 1.0f); - auto mask_context = mask.reshape({request->vae_scale_factor, - inactive.shape()[0], - request->vae_scale_factor, - inactive.shape()[1], - inactive.shape()[2]}); // [t, h/vae_scale_factor, vae_scale_factor, w/vae_scale_factor, vae_scale_factor] - mask_context = mask_context.permute({1, 3, 4, 0, 2}) // [vae_scale_factor, vae_scale_factor, t, h/vae_scale_factor, w/vae_scale_factor] - .reshape({inactive.shape()[0], - inactive.shape()[1], - inactive.shape()[2], - request->vae_scale_factor * request->vae_scale_factor}); // [vae_scale_factor*vae_scale_factor, t, h/vae_scale_factor, w/vae_scale_factor] + snprintf(buf + strlen(buf), 4096 - strlen(buf), + "(txt_cfg: %.2f, " + "img_cfg: %.2f, " + "distilled_guidance: %.2f, " + "slg.layer_count: %zu, " + "slg.layer_start: %.2f, " + "slg.layer_end: %.2f, " + "slg.scale: %.2f, " + "scheduler: %s, " + "sample_method: %s, " + "sample_steps: %d, " + "eta: %.2f, " + "shifted_timestep: %d, " + "flow_shift: %.2f, " + "extra_sample_args: %s)", + sample_params->guidance.txt_cfg, + std::isfinite(sample_params->guidance.img_cfg) + ? sample_params->guidance.img_cfg + : sample_params->guidance.txt_cfg, + sample_params->guidance.distilled_guidance, + sample_params->guidance.slg.layer_count, + sample_params->guidance.slg.layer_start, + sample_params->guidance.slg.layer_end, + sample_params->guidance.slg.scale, + sd_scheduler_name(sample_params->scheduler), + sd_sample_method_name(sample_params->sample_method), + sample_params->sample_steps, + sample_params->eta, + sample_params->shifted_timestep, + sample_params->flow_shift, + SAFE_STR(sample_params->extra_sample_args)); - if (!ref_image_latent.empty()) { - vace_context = sd::ops::concat(ref_image_latent, vace_context, 2); // [b, 2*c, t+1, h/vae_scale_factor, w/vae_scale_factor] - auto mask_pad = sd::zeros({mask_context.shape()[0], - mask_context.shape()[1], - 1, - mask_context.shape()[3]}); // [vae_scale_factor*vae_scale_factor, 1, h/vae_scale_factor, w/vae_scale_factor] - mask_context = sd::ops::concat(mask_pad, mask_context, 2); // [vae_scale_factor*vae_scale_factor, t + 1, h/vae_scale_factor, w/vae_scale_factor] - } + return buf; +} - mask_context.unsqueeze_(mask_context.dim()); // [b, vae_scale_factor*vae_scale_factor, t + 1 or t, h/vae_scale_factor, w/vae_scale_factor] +void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) { + *sd_img_gen_params = {}; + sd_sample_params_init(&sd_img_gen_params->sample_params); + sd_img_gen_params->clip_skip = -1; + sd_img_gen_params->ref_images_count = 0; + sd_img_gen_params->ref_image_args = ""; + sd_img_gen_params->width = 512; + sd_img_gen_params->height = 512; + sd_img_gen_params->strength = 0.75f; + sd_img_gen_params->seed = -1; + sd_img_gen_params->batch_count = 1; + sd_img_gen_params->control_strength = 0.9f; + sd_img_gen_params->ip_adapter_strength = 1.0f; + sd_img_gen_params->qwen_image_layers = 3; + sd_img_gen_params->circular_x = false; + sd_img_gen_params->circular_y = false; + sd_img_gen_params->pm_params = {nullptr, 0, nullptr, 20.f}; + sd_img_gen_params->pulid_params = {nullptr, 1.0f}; + sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr}; + sd_cache_params_init(&sd_img_gen_params->cache); + sd_hires_params_init(&sd_img_gen_params->hires); +} - latents.vace_context = sd::ops::concat(vace_context, mask_context, 3); // [b, 2*c + vae_scale_factor*vae_scale_factor, t + 1 or t, h/vae_scale_factor, w/vae_scale_factor] - int64_t t2 = ggml_time_ms(); - LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); - } +char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) { + char* buf = (char*)malloc(4096); + if (!buf) + return nullptr; + buf[0] = '\0'; - if (latents.init_latent.empty()) { - latents.init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height, request->frames, true); - } + char* sample_params_str = sd_sample_params_to_str(&sd_img_gen_params->sample_params); - if ((sd_version_is_ltxav(sd_ctx->sd->version) || sd_version_is_minimax_h3(sd_ctx->sd->version)) && - !latents.audio_latent.empty()) { - if (!latents.denoise_mask.empty()) { - latents.denoise_mask = pack_ltxav_audio_and_video_denoise_mask(latents.denoise_mask, - latents.init_latent, - latents.audio_latent); - } - latents.init_latent = pack_ltxav_audio_and_video_latents(latents.init_latent, latents.audio_latent); + snprintf(buf + strlen(buf), 4096 - strlen(buf), + "prompt: %s\n" + "negative_prompt: %s\n" + "clip_skip: %d\n" + "width: %d\n" + "height: %d\n" + "sample_params: %s\n" + "strength: %.2f\n" + "seed: %" PRId64 + "\n" + "batch_count: %d\n" + "qwen_image_layers: %d\n" + "ref_images_count: %d\n" + "ref_image_args: %s\n" + "control_strength: %.2f\n" + "photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\n" + "VAE tiling: %s (temporal=%s, extra_tiling_args=%s)\n" + "circular_x: %s\n" + "circular_y: %s\n" + "hires: {enabled=%s, upscaler=%s, model_path=%s, scale=%.2f, target=%dx%d, steps=%d, denoising_strength=%.2f}\n", + SAFE_STR(sd_img_gen_params->prompt), + SAFE_STR(sd_img_gen_params->negative_prompt), + sd_img_gen_params->clip_skip, + sd_img_gen_params->width, + sd_img_gen_params->height, + SAFE_STR(sample_params_str), + sd_img_gen_params->strength, + sd_img_gen_params->seed, + sd_img_gen_params->batch_count, + sd_img_gen_params->qwen_image_layers, + sd_img_gen_params->ref_images_count, + SAFE_STR(sd_img_gen_params->ref_image_args), + sd_img_gen_params->control_strength, + sd_img_gen_params->pm_params.style_strength, + sd_img_gen_params->pm_params.id_images_count, + SAFE_STR(sd_img_gen_params->pm_params.id_embed_path), + BOOL_STR(sd_img_gen_params->vae_tiling_params.enabled), + BOOL_STR(sd_img_gen_params->vae_tiling_params.temporal_tiling), + SAFE_STR(sd_img_gen_params->vae_tiling_params.extra_tiling_args), + BOOL_STR(sd_img_gen_params->circular_x), + BOOL_STR(sd_img_gen_params->circular_y), + BOOL_STR(sd_img_gen_params->hires.enabled), + sd_hires_upscaler_name(sd_img_gen_params->hires.upscaler), + SAFE_STR(sd_img_gen_params->hires.model_path), + sd_img_gen_params->hires.scale, + sd_img_gen_params->hires.target_width, + sd_img_gen_params->hires.target_height, + sd_img_gen_params->hires.steps, + sd_img_gen_params->hires.denoising_strength); + const char* cache_mode_str = "disabled"; + if (sd_img_gen_params->cache.mode == SD_CACHE_EASYCACHE) { + cache_mode_str = "easycache"; + } else if (sd_img_gen_params->cache.mode == SD_CACHE_UCACHE) { + cache_mode_str = "ucache"; } - - return latents; + snprintf(buf + strlen(buf), 4096 - strlen(buf), + "cache: %s (threshold=%.3f, start=%.2f, end=%.2f)\n", + cache_mode_str, + get_cache_reuse_threshold(sd_img_gen_params->cache), + sd_img_gen_params->cache.start_percent, + sd_img_gen_params->cache.end_percent); + free(sample_params_str); + return buf; } -static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx, - const sd_vid_gen_params_t* sd_vid_gen_params, - const GenerationRequest& request, - const ImageGenerationLatents& latents) { - ConditionerRunnerEndOnExit conditioner_runner_end{sd_ctx->sd->cond_stage_model.get()}; - - ImageGenerationEmbeds embeds; - ConditionerParams condition_params; - condition_params.clip_skip = request.clip_skip; - condition_params.text = request.prompt; - condition_params.zero_out_masked = true; - condition_params.ref_images = &latents.ref_images; - condition_params.minimax_h3_references = &latents.minimax_presentation_refs; - if (sd_version_is_lingbot_video(sd_ctx->sd->version) || sd_version_is_minimax_h3(sd_ctx->sd->version)) { - condition_params.ref_image_params.vlm_resize_mode = RefImageResizeMode::AREA; - } +void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params) { + *sd_vid_gen_params = {}; + sd_sample_params_init(&sd_vid_gen_params->sample_params); + sd_sample_params_init(&sd_vid_gen_params->high_noise_sample_params); + sd_vid_gen_params->high_noise_sample_params.sample_steps = -1; + sd_vid_gen_params->width = 512; + sd_vid_gen_params->height = 512; + sd_vid_gen_params->strength = 0.75f; + sd_vid_gen_params->seed = -1; + sd_vid_gen_params->video_frames = 6; + sd_vid_gen_params->fps = 16; + sd_vid_gen_params->moe_boundary = 0.875f; + sd_vid_gen_params->vace_strength = 1.f; + sd_vid_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr}; + sd_vid_gen_params->hires.enabled = false; + sd_vid_gen_params->hires.upscaler = SD_HIRES_UPSCALER_LATENT; + sd_vid_gen_params->hires.scale = 2.f; + sd_vid_gen_params->hires.target_width = 0; + sd_vid_gen_params->hires.target_height = 0; + sd_vid_gen_params->hires.steps = 0; + sd_vid_gen_params->hires.denoising_strength = 0.7f; + sd_vid_gen_params->hires.upscale_tile_size = 128; + sd_vid_gen_params->hires.custom_sigmas = nullptr; + sd_vid_gen_params->hires.custom_sigmas_count = 0; + sd_vid_gen_params->circular_x = false; + sd_vid_gen_params->circular_y = false; + sd_cache_params_init(&sd_vid_gen_params->cache); +} - int64_t prepare_start_ms = ggml_time_ms(); - embeds.cond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, - condition_params); - embeds.cond.c_concat = latents.concat_latent; - embeds.cond.c_vector = latents.clip_vision_output; - if (sd_version_is_minimax_h3(sd_ctx->sd->version)) { - embeds.cond.c_ref_images = latents.ref_latents; - embeds.cond.c_ref_audios = latents.reference_audio_latents; - embeds.cond.c_reference_blocks = latents.minimax_reference_blocks; - if (!latents.keyframe_indices.empty()) { - embeds.cond.c_position_ids = sd::Tensor( - {static_cast(latents.keyframe_indices.size())}, - latents.keyframe_indices); - } - } - if (request.use_uncond) { - condition_params.text = request.negative_prompt; - embeds.uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, - condition_params); - embeds.uncond.c_concat = latents.concat_latent; - embeds.uncond.c_vector = latents.clip_vision_output; - if (sd_version_is_minimax_h3(sd_ctx->sd->version)) { - embeds.uncond.c_ref_images = latents.ref_latents; - embeds.uncond.c_ref_audios = latents.reference_audio_latents; - embeds.uncond.c_reference_blocks = latents.minimax_reference_blocks; - embeds.uncond.c_position_ids = embeds.cond.c_position_ids; - } - } +struct sd_ctx_t { + StableDiffusionGGML* sd = nullptr; +}; - int64_t t1 = ggml_time_ms(); - LOG_INFO("get_learned_condition completed, taking %.2fs", (t1 - prepare_start_ms) * 1.0f / 1000); +static bool sd_version_supports_video_generation(SDVersion version) { + return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version); +} - return embeds; +static bool sd_version_supports_image_generation(SDVersion version) { + return !sd_version_supports_video_generation(version); } -static sd_image_t* decode_video_outputs(sd_ctx_t* sd_ctx, - const GenerationRequest& request, - const sd::Tensor& final_latent, - int* num_frames_out) { - if (final_latent.empty()) { - LOG_ERROR("no latent video to decode"); - return nullptr; - } - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling video decode"); - return nullptr; - } - sd::Tensor video_latent = final_latent; - if ((sd_version_is_ltxav(sd_ctx->sd->version) || sd_version_is_minimax_h3(sd_ctx->sd->version)) && - video_latent.shape()[3] > sd_ctx->sd->get_latent_channel()) { - video_latent = sd::ops::slice(video_latent, 3, 0, sd_ctx->sd->get_latent_channel()); - } - LOG_VERBOSE("decode_video_outputs latent %dx%dx%dx%d", - (int)video_latent.shape()[0], - (int)video_latent.shape()[1], - (int)video_latent.shape()[2], - (int)video_latent.shape()[3]); - // auto z = sd::load_tensor_from_file_as_tensor("ltx_vae_z.bin"); - int64_t t4 = ggml_time_ms(); - sd::Tensor vid = sd_ctx->sd->decode_first_stage(video_latent, true); - int64_t t5 = ggml_time_ms(); - LOG_INFO("decode_first_stage completed, taking %.2fs", (t5 - t4) * 1.0f / 1000); - if (vid.empty()) { - LOG_ERROR("decode_first_stage failed for video"); +sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) { + sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t)); + if (sd_ctx == nullptr) { return nullptr; } - LOG_VERBOSE("decode_video_outputs decoded %dx%dx%dx%d", - (int)vid.shape()[0], - (int)vid.shape()[1], - (int)vid.shape()[2], - (int)vid.shape()[3]); - if (request.frames > 0 && - vid.shape()[2] > request.frames) { - vid = sd::ops::slice(vid, 2, 0, request.frames); - } - sd_image_t* result_images = (sd_image_t*)calloc(vid.shape()[2], sizeof(sd_image_t)); - if (result_images == nullptr) { + sd_ctx->sd = new StableDiffusionGGML(); + if (sd_ctx->sd == nullptr) { + free(sd_ctx); return nullptr; } - if (num_frames_out != nullptr) { - *num_frames_out = static_cast(vid.shape()[2]); - } - for (int64_t i = 0; i < vid.shape()[2]; i++) { - result_images[i] = tensor_to_sd_image(vid, static_cast(i)); + if (!sd_ctx->sd->init(sd_ctx_params)) { + delete sd_ctx->sd; + sd_ctx->sd = nullptr; + free(sd_ctx); + return nullptr; } - - return result_images; + return sd_ctx; } -static sd::Tensor upscale_ltx_spatial_video_latent(sd_ctx_t* sd_ctx, - const char* model_path, - const sd::Tensor& packed_latent, - int audio_length) { - if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_ctx->sd->model_manager == nullptr || packed_latent.empty()) { - return {}; - } - if (strlen(SAFE_STR(model_path)) == 0) { - LOG_ERROR("LTX latent spatial upscale requires a model path"); - return {}; - } - if (!sd_ctx->sd->ensure_backend_pair(SDBackendModule::UPSCALER)) { - return {}; - } - - int latent_channels = sd_ctx->sd->get_latent_channel(); - sd::Tensor video_latent = packed_latent; - sd::Tensor audio_latent; - if (packed_latent.shape()[3] > latent_channels) { - video_latent = sd::ops::slice(packed_latent, 3, 0, latent_channels); - audio_latent = unpack_ltxav_audio_latent(packed_latent, audio_length, latent_channels); - } - - LOG_INFO("LTX latent spatial upscale: latent %dx%dx%dx%d -> model output", - (int)video_latent.shape()[0], - (int)video_latent.shape()[1], - (int)video_latent.shape()[2], - (int)video_latent.shape()[3]); - - sd::Tensor unnormalized = sd_ctx->sd->un_normalize_ltx_video_latents(video_latent); - if (sd_ctx->sd->first_stage_model) { - sd_ctx->sd->first_stage_model->runner_end(); - } - if (unnormalized.empty()) { - LOG_ERROR("LTX latent un-normalization failed before spatial upscale"); - return {}; +void free_sd_ctx(sd_ctx_t* sd_ctx) { + if (sd_ctx->sd != nullptr) { + delete sd_ctx->sd; + sd_ctx->sd = nullptr; } + free(sd_ctx); +} - auto model_manager = sd_ctx->sd->model_manager; - struct UpsamplerScope { - ModelManager& manager; - ModelLoader::FileId owned_source = 0; - std::unique_ptr runner; - std::vector params; - - ~UpsamplerScope() { - if (runner) { - runner->runner_end(); - } - GGML_ASSERT(manager.unregister_param_tensors(params)); - if (owned_source != 0) { - GGML_ASSERT(manager.del_file(owned_source)); - } +SD_API void sd_cancel_generation(sd_ctx_t* sd_ctx, enum sd_cancel_mode_t mode) { + if (sd_ctx && sd_ctx->sd) { + if (mode < SD_CANCEL_ALL || mode > SD_CANCEL_RESET) { + mode = SD_CANCEL_ALL; } - } scope{*model_manager}; - - const std::string prefix = "ltx_latent_upsampler"; - ModelLoader candidate = model_manager->loader(); - ModelLoader::FileId source_file = 0; - if (!candidate.add_file(model_path, prefix + ".", &source_file)) { - LOG_ERROR("init LTX latent upsampler model loader from file failed: '%s'", model_path); - return {}; - } - const bool owns_source = model_manager->loader().file_revision(source_file) == 0; - if (!model_manager->set_loader(std::move(candidate))) { - return {}; + sd_ctx->sd->set_cancel_flag(mode); } - scope.owned_source = owns_source ? source_file : 0; +} - auto& upsampler = scope.runner; - upsampler = std::make_unique(sd_ctx->sd->backend_for(SDBackendModule::UPSCALER), - model_manager->loader().get_tensor_storage_map(), - prefix, - model_manager); - const size_t max_graph_vram_bytes = sd_ctx->sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER); - upsampler->set_max_graph_vram_bytes(max_graph_vram_bytes); - if (upsampler->model == nullptr) { - LOG_ERROR("init LTX latent upsampler from metadata failed"); - return {}; +void free_sd_audio(sd_audio_t* audio) { + if (audio == nullptr) { + return; } + free(audio->data); + audio->data = nullptr; + free(audio); +} - std::map tensors; - upsampler->get_param_tensors(tensors); - for (const auto& entry : tensors) { - scope.params.push_back(entry.second); - } - if (!model_manager->register_param_tensors(ModelComponent::LatentUpsampler, - std::move(tensors), - ModelManager::ResidencyMode::ParamBackend, - sd_ctx->sd->backend_for(SDBackendModule::UPSCALER), - sd_ctx->sd->params_backend_for(SDBackendModule::UPSCALER)) || - !model_manager->validate_registered_tensors()) { - LOG_ERROR("register LTX latent upsampler tensors with model manager failed"); - return {}; +SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx) { + if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { + return false; } + return sd_version_supports_image_generation(sd_ctx->sd->version); +} - sd::Tensor upscaled = upsampler->compute(sd_ctx->sd->n_threads, unnormalized); - upsampler->runner_end(); - if (upscaled.empty()) { - LOG_ERROR("LTX latent spatial upscale failed"); - return {}; +SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) { + if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { + return false; } - - upscaled = sd_ctx->sd->normalize_ltx_video_latents(upscaled); - sd_ctx->sd->first_stage_model->runner_end(); - if (upscaled.empty()) { - LOG_ERROR("LTX latent normalization failed after spatial upscale"); - return {}; + if (sd_ctx->sd->config_->animatediff_loaded && sd_version_supports_animatediff(sd_ctx->sd->version)) { + return true; } + return sd_version_supports_video_generation(sd_ctx->sd->version); +} - if (!audio_latent.empty()) { - upscaled = pack_ltxav_audio_and_video_latents(upscaled, audio_latent); +SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path) { + if (sd_ctx == nullptr || sd_ctx->sd == nullptr || path == nullptr) { + return false; } - return upscaled; + return sd_ctx->sd->load_control_net_from_file(path); } -static bool apply_ltxv_refine_image_conditioning(sd_ctx_t* sd_ctx, - const sd_vid_gen_params_t* sd_vid_gen_params, - const GenerationRequest& request, - const ImageGenerationLatents& latents, - sd::Tensor* latent, - sd::Tensor* denoise_mask, - sd::Tensor* video_positions) { - if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr || - latent == nullptr || latent->empty() || denoise_mask == nullptr || video_positions == nullptr) { - return true; - } - if (sd_vid_gen_params->init_image.data == nullptr && - sd_vid_gen_params->end_image.data == nullptr) { - return true; - } - constexpr float conditioning_strength = 1.f; - int latent_channels = sd_ctx->sd->get_latent_channel(); - sd::Tensor video_latent = *latent; - sd::Tensor audio_latent; - if (latent->shape()[3] > latent_channels) { - video_latent = sd::ops::slice(*latent, 3, 0, latent_channels); - audio_latent = unpack_ltxav_audio_latent(*latent, latents.audio_length, latent_channels); - if (audio_latent.empty()) { - LOG_ERROR("failed to unpack LTXAV audio latent before image-to-video inplace conditioning"); - return false; - } +SD_API bool sd_ctx_unload_control_net(sd_ctx_t* sd_ctx) { + if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { + return false; } + return sd_ctx->sd->unload_control_net(); +} - int image_width = static_cast(video_latent.shape()[0]) * request.vae_scale_factor; - int image_height = static_cast(video_latent.shape()[1]) * request.vae_scale_factor; - sd::Tensor video_mask = make_ltxav_video_denoise_mask(video_latent, 1.f); - - if (sd_vid_gen_params->init_image.data != nullptr) { - sd::Tensor start_image = sd_image_to_tensor(sd_vid_gen_params->init_image, image_width, image_height); - if (!apply_ltxav_condition_image_by_latent_index(sd_ctx, - start_image, - &video_latent, - &video_mask, - 0, - "init", - conditioning_strength)) { - return false; - } +SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx) { + if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { + return false; } + return sd_ctx->sd->control_net != nullptr; +} - if (sd_vid_gen_params->end_image.data != nullptr) { - sd::Tensor end_image = sd_image_to_tensor(sd_vid_gen_params->end_image, image_width, image_height); - sd::Tensor end_image_latent = encode_ltxav_condition_image(sd_ctx, end_image, "end"); - if (end_image_latent.empty()) { - return false; - } - - int frame_idx = request.frames - 1; - if (frame_idx == 0) { - if (!apply_ltxav_condition_by_latent_index(&video_latent, - &video_mask, - end_image_latent, - 0, - "end", - 1.f - conditioning_strength)) { - return false; - } - } else { - if (latents.video_conditioning_frame_count <= 0 || latents.video_target_frame_count <= 0) { - LOG_ERROR("LTXV FLF2V refine conditioning requires low-resolution keyframe conditioning metadata"); - return false; - } - int64_t target_latent_frames = latents.video_target_frame_count; - if (!apply_ltxav_condition_by_latent_index(&video_latent, - &video_mask, - end_image_latent, - target_latent_frames, - "end", - 1.f - conditioning_strength)) { - return false; - } - *video_positions = build_ltxv_video_positions(video_latent.shape()[0], - video_latent.shape()[1], - target_latent_frames, - end_image_latent.shape()[2], - frame_idx, - 1, - request.fps, - request.vae_scale_factor, - 8, - true); - } - } +enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) { + return sd::pipeline::default_sample_method(sd_ctx != nullptr ? sd_ctx->sd : nullptr); +} - if (!audio_latent.empty()) { - *latent = pack_ltxav_audio_and_video_latents(video_latent, audio_latent); - *denoise_mask = pack_ltxav_audio_and_video_denoise_mask(video_mask, video_latent, audio_latent); - } else { - *latent = std::move(video_latent); - *denoise_mask = std::move(video_mask); - } - LOG_INFO("LTXV refine image conditioning applied at %dx%d", image_width, image_height); - return true; +enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_method_t sample_method) { + return sd::pipeline::default_scheduler(sd_ctx != nullptr ? sd_ctx->sd : nullptr, sample_method); } -static bool generate_animatediff_video(sd_ctx_t* sd_ctx, - const sd_vid_gen_params_t* sd_vid_gen_params, - sd_image_t** frames_out, - int* num_frames_out) { - int n_frames = sd_vid_gen_params->video_frames; - if (n_frames < 1) { - LOG_ERROR("AnimateDiff: --video-frames must be >= 1"); +SD_API bool generate_image(sd_ctx_t* sd_ctx, + const sd_img_gen_params_t* params, + sd_image_t** images_out, + int* num_images_out) { + if (images_out != nullptr) + *images_out = nullptr; + if (num_images_out != nullptr) + *num_images_out = 0; + if (sd_ctx == nullptr || sd_ctx->sd == nullptr || params == nullptr) { return false; } - if (n_frames > 32) { - LOG_WARN("AnimateDiff motion modules have a 32-frame positional-encoding context; capping to 32"); - n_frames = 32; - } - - sd_img_gen_params_t img_gen_params; - sd_img_gen_params_init(&img_gen_params); - img_gen_params.loras = sd_vid_gen_params->loras; - img_gen_params.lora_count = sd_vid_gen_params->lora_count; - img_gen_params.prompt = sd_vid_gen_params->prompt; - img_gen_params.negative_prompt = sd_vid_gen_params->negative_prompt; - img_gen_params.clip_skip = sd_vid_gen_params->clip_skip; - img_gen_params.width = sd_vid_gen_params->width; - img_gen_params.height = sd_vid_gen_params->height; - img_gen_params.sample_params = sd_vid_gen_params->sample_params; - img_gen_params.strength = sd_vid_gen_params->strength; - img_gen_params.init_image = sd_vid_gen_params->init_image; - img_gen_params.seed = sd_vid_gen_params->seed; - img_gen_params.batch_count = 1; - img_gen_params.control_strength = 1.0f; - img_gen_params.vae_tiling_params = sd_vid_gen_params->vae_tiling_params; - img_gen_params.cache = sd_vid_gen_params->cache; - img_gen_params.hires = sd_vid_gen_params->hires; - img_gen_params.qwen_image_layers = 0; - img_gen_params.circular_x = sd_vid_gen_params->circular_x; - img_gen_params.circular_y = sd_vid_gen_params->circular_y; - - sd_ctx->sd->animatediff_num_frames = n_frames; - bool ok = generate_image_impl(sd_ctx, &img_gen_params, frames_out, num_frames_out); - sd_ctx->sd->animatediff_num_frames = 0; - return ok; + StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd); + return execution.ready && sd::pipeline::generate_image(sd_ctx->sd, params, images_out, num_images_out); } SD_API bool generate_video(sd_ctx_t* sd_ctx, @@ -6705,375 +742,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, return false; } - if (sd_ctx->sd->config_->animatediff_loaded && sd_version_supports_animatediff(sd_ctx->sd->version)) { - LOG_INFO("AnimateDiff dispatch: %d frames, %dx%d", - sd_vid_gen_params->video_frames, sd_vid_gen_params->width, sd_vid_gen_params->height); - return generate_animatediff_video(sd_ctx, sd_vid_gen_params, frames_out, num_frames_out); - } - - sd_ctx->sd->reset_cancel_flag(); - - const RefImageParams ref_image_params; - - int64_t t0 = ggml_time_ms(); - sd_ctx->sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params; - apply_circular_axes_to_diffusion(sd_ctx, sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y); - GenerationRequest request(sd_ctx, sd_vid_gen_params); - bool latent_upscale_enabled = request.hires.enabled; - GenerationRequest hires_request = request; - if (latent_upscale_enabled) { - if (!sd_version_is_ltxav(sd_ctx->sd->version)) { - LOG_ERROR("LTX latent spatial upscale is only supported for LTX video models"); - return false; - } - if (request.hires.upscaler != SD_HIRES_UPSCALER_MODEL) { - LOG_ERROR("LTX latent spatial upscale currently requires hires upscaler MODEL"); - return false; - } - if (strlen(SAFE_STR(request.hires.model_path)) == 0) { - LOG_ERROR("LTX latent spatial upscale is enabled but hires model path was not provided"); - return false; - } - } - - sd_ctx->sd->rng->manual_seed(request.seed); - sd_ctx->sd->sampler_rng->manual_seed(request.seed); - sd_ctx->sd->set_flow_shift(sd_vid_gen_params->sample_params.flow_shift); - if (!sd_ctx->sd->apply_loras(sd_vid_gen_params->loras, sd_vid_gen_params->lora_count)) - return false; - sd_ctx->sd->reset_generation_extensions(); - - SamplePlan plan(sd_ctx, sd_vid_gen_params, request); - auto latent_inputs_opt = prepare_video_generation_latents(sd_ctx, sd_vid_gen_params, &request); - if (!latent_inputs_opt.has_value()) { - return false; - } - ImageGenerationLatents latents = std::move(*latent_inputs_opt); - - ImageGenerationEmbeds embeds = prepare_video_generation_embeds(sd_ctx, - sd_vid_gen_params, - request, - latents); - if (latent_upscale_enabled) { - LOG_INFO("generate_video %dx%dx%d -> LTX latent spatial upscale", - request.width, - request.height, - request.frames); - } else { - LOG_INFO("generate_video %dx%dx%d", - request.width, - request.height, - request.frames); - } - - int64_t latent_start = ggml_time_ms(); - int W = request.width / request.vae_scale_factor; - int H = request.height / request.vae_scale_factor; - int T = static_cast(latents.init_latent.shape()[2]); - - sd::Tensor x_t = latents.init_latent; - sd::Tensor noise = sd::Tensor::randn_like(x_t, sd_ctx->sd->rng); - - if (plan.high_noise_sample_steps > 0) { - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling generation before high-noise sampling"); - return false; - } - LOG_VERBOSE("sample(high noise) %dx%dx%d", W, H, T); - - int64_t sampling_start = ggml_time_ms(); - std::vector high_noise_sigmas(plan.sigmas.begin(), plan.sigmas.begin() + plan.high_noise_sample_steps + 1); - plan.sigmas = std::vector(plan.sigmas.begin() + plan.high_noise_sample_steps, plan.sigmas.end()); - - sd::Tensor x_t_sampled = sd_ctx->sd->sample(sd_ctx->sd->high_noise_diffusion_model, - false, - x_t, - std::move(noise), - embeds.cond, - request.use_high_noise_uncond ? embeds.uncond : SDCondition(), - embeds.img_uncond, - sd::Tensor(), - 0.f, - request.high_noise_guidance, - plan.high_noise_eta, - request.shifted_timestep, - plan.high_noise_sample_method, - sd_ctx->sd->is_flow_denoiser(), - plan.high_noise_extra_sample_args, - high_noise_sigmas, - std::vector>{}, - ref_image_params, - latents.denoise_mask, - latents.vace_context, - request.vace_strength, - latents.audio_length, - static_cast(request.fps), - request.cache_params, - true, - latents.video_positions); - int64_t sampling_end = ggml_time_ms(); - if (x_t_sampled.empty()) { - LOG_ERROR("sampling(high noise) failed after %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); - return false; - } - - x_t = std::move(x_t_sampled); - noise = {}; - LOG_INFO("sampling(high noise) completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); - } - - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling generation before sampling"); - return false; - } - LOG_VERBOSE("sample %dx%dx%d", W, H, T); - int64_t sampling_start = ggml_time_ms(); - sd::Tensor final_latent = sd_ctx->sd->sample(sd_ctx->sd->diffusion_model, - true, - x_t, - std::move(noise), - embeds.cond, - request.use_uncond ? embeds.uncond : SDCondition(), - embeds.img_uncond, - sd::Tensor(), - 0.f, - sd_vid_gen_params->sample_params.guidance, - plan.eta, - sd_vid_gen_params->sample_params.shifted_timestep, - plan.sample_method, - sd_ctx->sd->is_flow_denoiser(), - plan.extra_sample_args, - plan.sigmas, - std::vector>{}, - ref_image_params, - latents.denoise_mask, - latents.vace_context, - request.vace_strength, - latents.audio_length, - static_cast(request.fps), - request.cache_params, - plan.high_noise_sample_steps <= 0, - latents.video_positions); - - int64_t sampling_end = ggml_time_ms(); - if (final_latent.empty()) { - LOG_ERROR("sampling failed after %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); - return false; - } - LOG_INFO("sampling completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); - - if (latent_upscale_enabled) { - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling generation before latent upscale"); - return false; - } - int64_t upscale_start = ggml_time_ms(); - sd::Tensor upscaled_latent = upscale_ltx_spatial_video_latent(sd_ctx, - request.hires.model_path, - final_latent, - latents.audio_length); - int64_t upscale_end = ggml_time_ms(); - if (upscaled_latent.empty()) { - return false; - } - LOG_INFO("LTX latent spatial upscale completed, taking %.2fs", - (upscale_end - upscale_start) * 1.0f / 1000); - - x_t = std::move(upscaled_latent); - hires_request.width = static_cast(x_t.shape()[0]) * hires_request.vae_scale_factor; - hires_request.height = static_cast(x_t.shape()[1]) * hires_request.vae_scale_factor; - int upscaled_latent_frames = static_cast(x_t.shape()[2]); - int upscaled_frames = sd_ctx->sd->latent_frames_to_video_frames(upscaled_latent_frames); - if (upscaled_frames != hires_request.frames) { - LOG_INFO("LTX latent upsampler output latent frames %d, frames %d -> %d", - upscaled_latent_frames, - hires_request.frames, - upscaled_frames); - hires_request.frames = upscaled_frames; - } - if (sd_version_is_ltxav(sd_ctx->sd->version) && latents.audio_length > 0) { - int target_audio_length = get_ltxav_num_audio_latents(hires_request.frames, hires_request.fps); - if (target_audio_length != latents.audio_length) { - int latent_channels = sd_ctx->sd->get_latent_channel(); - sd::Tensor video_latent = x_t; - sd::Tensor audio_latent = latents.audio_latent; - if (x_t.shape()[3] > latent_channels) { - video_latent = sd::ops::slice(x_t, 3, 0, latent_channels); - audio_latent = unpack_ltxav_audio_latent(x_t, latents.audio_length, latent_channels); - } - audio_latent = resize_ltxav_audio_latent(audio_latent, target_audio_length); - if (audio_latent.empty()) { - LOG_ERROR("failed to resize LTX audio latent for latent upscale: %d -> %d", - latents.audio_length, - target_audio_length); - return false; - } - x_t = pack_ltxav_audio_and_video_latents(video_latent, audio_latent); - latents.audio_latent = std::move(audio_latent); - LOG_INFO("LTX audio latent length adjusted for latent upscale: %d -> %d", - latents.audio_length, - target_audio_length); - latents.audio_length = target_audio_length; - } - } - if ((request.hires.target_width > 0 || request.hires.target_height > 0) && - (request.hires.target_width != hires_request.width || request.hires.target_height != hires_request.height)) { - LOG_WARN("LTX latent spatial upsampler output is %dx%d; ignoring hires target %dx%d", - hires_request.width, - hires_request.height, - request.hires.target_width, - request.hires.target_height); - } - sd::Tensor hires_denoise_mask; - sd::Tensor hires_video_positions; - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling generation before latent upscale refine"); - return false; - } - if (!apply_ltxv_refine_image_conditioning(sd_ctx, - sd_vid_gen_params, - hires_request, - latents, - &x_t, - &hires_denoise_mask, - &hires_video_positions)) { - return false; - } - noise = sd::Tensor::randn_like(x_t, sd_ctx->sd->rng); - - W = hires_request.width / hires_request.vae_scale_factor; - H = hires_request.height / hires_request.vae_scale_factor; - T = static_cast(x_t.shape()[2]); - sample_method_t hires_sample_method = plan.sample_method; - int hires_scheduler_steps = 0; - std::vector hires_sigma_sched = - make_hires_sigma_schedule(sd_ctx, - request.hires, - sd_vid_gen_params->sample_params, - hires_sample_method, - plan.sample_steps, - sd_ctx->sd->get_image_seq_len(hires_request.height, hires_request.width) * T, - &hires_scheduler_steps); - float hires_eta = resolve_eta(sd_ctx, - sd_vid_gen_params->sample_params.eta, - hires_sample_method); - - LOG_VERBOSE("sample(latent upscale) %dx%dx%d", W, H, T); - LOG_INFO("LTX latent spatial upscale refine: scheduler_steps=%d, denoising_strength=%.2f, sampler=%s, sigma_sched_size=%zu%s", - hires_scheduler_steps, - request.hires.denoising_strength, - sampling_methods_str[hires_sample_method], - hires_sigma_sched.size(), - request.hires.custom_sigmas_count > 0 ? ", custom_sigmas=true" : ""); - - sampling_start = ggml_time_ms(); - final_latent = sd_ctx->sd->sample(sd_ctx->sd->diffusion_model, - true, - x_t, - std::move(noise), - embeds.cond, - hires_request.use_uncond ? embeds.uncond : SDCondition(), - embeds.img_uncond, - sd::Tensor(), - 0.f, - sd_vid_gen_params->sample_params.guidance, - hires_eta, - sd_vid_gen_params->sample_params.shifted_timestep, - hires_sample_method, - sd_ctx->sd->is_flow_denoiser(), - plan.extra_sample_args, - hires_sigma_sched, - std::vector>{}, - ref_image_params, - hires_denoise_mask, - sd::Tensor(), - hires_request.vace_strength, - latents.audio_length, - static_cast(hires_request.fps), - hires_request.cache_params, - false, - hires_video_positions); - sampling_end = ggml_time_ms(); - if (final_latent.empty()) { - LOG_ERROR("sampling(latent upscale) failed after %.2fs", - (sampling_end - sampling_start) * 1.0f / 1000); - return false; - } - LOG_INFO("sampling(latent upscale) completed, taking %.2fs", - (sampling_end - sampling_start) * 1.0f / 1000); - } - - int64_t latent_end = ggml_time_ms(); - LOG_INFO("generating latent video completed, taking %.2fs", (latent_end - latent_start) * 1.0f / 1000); - - sd_audio_t* generated_audio = nullptr; - if ((sd_version_is_ltxav(sd_ctx->sd->version) || sd_version_is_minimax_h3(sd_ctx->sd->version)) && - latents.audio_length > 0 && - sd_ctx->sd->audio_vae_model != nullptr) { - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling generation before audio decode"); - return false; - } - int64_t audio_latent_decode_start = ggml_time_ms(); - - auto audio_latent = sd_version_is_minimax_h3(sd_ctx->sd->version) - ? unpack_minimax_h3_audio_latent(final_latent, - latents.audio_length, - sd_ctx->sd->get_latent_channel()) - : unpack_ltxav_audio_latent(final_latent, - latents.audio_length, - sd_ctx->sd->get_latent_channel()); - if (!audio_latent.empty()) { - LOG_VERBOSE("decode audio latent %dx%dx%dx%d", - (int)audio_latent.shape()[0], - (int)audio_latent.shape()[1], - (int)audio_latent.shape()[2], - (int)audio_latent.shape()[3]); - auto waveform = sd_ctx->sd->decode_ltx_audio_latent(audio_latent); - if (!waveform.empty()) { - generated_audio = waveform_to_sd_audio(sd_ctx->sd, waveform); - } else { - LOG_WARN("audio latent decode failed; continuing with silent video output"); - } - } - int64_t audio_latent_decode_end = ggml_time_ms(); - LOG_INFO("decoding audio latent completed, taking %.2fs", (audio_latent_decode_end - audio_latent_decode_start) * 1.0f / 1000); - } - - if (latents.video_conditioning_frame_count > 0) { - int64_t target_frames = latents.video_target_frame_count > 0 ? latents.video_target_frame_count - : final_latent.shape()[2] - latents.video_conditioning_frame_count; - final_latent = sd::ops::slice(final_latent, 2, 0, target_frames); - } - - if (latents.ref_image_num > 0) { - final_latent = sd::ops::slice(final_latent, 2, latents.ref_image_num, final_latent.shape()[2]); - } - - if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { - LOG_ERROR("cancelling generation before video decode"); - free_sd_audio(generated_audio); - return false; - } - auto result = decode_video_outputs(sd_ctx, latent_upscale_enabled ? hires_request : request, final_latent, num_frames_out); - if (result == nullptr) { - free_sd_audio(generated_audio); - return false; - } - - sd_ctx->sd->lora_stat(); - - int64_t t1 = ggml_time_ms(); - LOG_INFO("generate_video completed in %.2fs", (t1 - t0) * 1.0f / 1000); - if (frames_out != nullptr) { - *frames_out = result; - } - if (audio_out != nullptr) { - *audio_out = generated_audio; - } else { - free_sd_audio(generated_audio); - } - return true; + return sd::pipeline::generate_video(sd_ctx->sd, sd_vid_gen_params, frames_out, num_frames_out, audio_out); } SD_API void free_sd_images(sd_image_t* result_images, int num_images) {