Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
57 lines
2.1 KiB
C++
57 lines
2.1 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/runtime/errors.h"
|
|
#include "flow_sampler_runtime.h"
|
|
|
|
#include <algorithm>
|
|
#include <cctype>
|
|
#include <cmath>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace engine::modules {
|
|
|
|
// Create a linear or sway timestep schedule for RF sampling.
|
|
inline std::vector<FlowSamplerScheduleStep>
|
|
make_rf_schedule(int64_t steps, std::string model,
|
|
std::string mode = "linear", float sway_coeff = -1.0F,
|
|
float init_scale = 0.999F) {
|
|
using runtime::InvalidRequestError;
|
|
if (steps <= 0) {
|
|
throw InvalidRequestError(model + " RF Schedule requires positive steps");
|
|
}
|
|
const auto first = mode.find_first_not_of(" \t\r\n");
|
|
mode = first == std::string::npos ? "" :
|
|
mode.substr(first, mode.find_last_not_of(" \t\r\n") - first + 1);
|
|
std::transform(mode.begin(), mode.end(), mode.begin(),
|
|
[](unsigned char c) { return static_cast<char>(std::tolower(c)); });
|
|
if (mode != "linear" && mode != "sway") {
|
|
throw InvalidRequestError(model + " t_schedule_mode must be linear or sway");
|
|
}
|
|
if (!std::isfinite(sway_coeff)) {
|
|
throw InvalidRequestError(model + " sway_coeff must be finite");
|
|
}
|
|
auto time_at = [&](int64_t index) {
|
|
float u = static_cast<float>(index) / static_cast<float>(steps);
|
|
if (mode == "sway") {
|
|
u = std::clamp(u + sway_coeff * (std::cos(1.5707963267948966F * u) + u - 1.0F),
|
|
0.0F, 1.0F);
|
|
}
|
|
return (1.0F - u) * init_scale;
|
|
};
|
|
std::vector<FlowSamplerScheduleStep> schedule;
|
|
schedule.reserve(static_cast<size_t>(steps));
|
|
float t = time_at(0);
|
|
for (int64_t step = 0; step < steps; ++step) {
|
|
const float next = time_at(step + 1);
|
|
if (!(t > next)) {
|
|
throw runtime::InvalidRequestError(
|
|
model + " t_schedule must be strictly decreasing; adjust num_inference_steps or sway_coeff");
|
|
}
|
|
schedule.push_back({step, t, next, 0.0F, 0.0F});
|
|
t = next;
|
|
}
|
|
return schedule;
|
|
}
|
|
|
|
} // namespace engine::modules
|