Skip to content

sd: cherry-pick changes up to master-827-97d2990 #2408

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wbruna wants to merge 1 commit into
base: concedo_experimental
Choose a base branch
from
Open

sd: cherry-pick changes up to master-827-97d2990 #2408

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 17 additions & 6 deletions otherarch/sdcpp/examples/common/common.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1008,7 +1008,7 @@ ArgOptions SDGenerationParams::get_options() {
&hires_upscaler},
{"",
"--extra-sample-args",
"extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_divisions",
"extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_max_order, lms_shift, lms_divisions",
(int)',',
&extra_sample_args},
{"",
Expand Down Expand Up @@ -1555,24 +1555,35 @@ ArgOptions SDGenerationParams::get_options() {
return 1;
};

std::string sample_methods = sample_method_to_str[0];
for (int i = 1; i < SAMPLE_METHOD_COUNT; i++) {
sample_methods += ", " + std::string(sample_method_to_str[i]);
}

std::string schedulers = scheduler_to_str[0];
for (int i = 1; i < SCHEDULER_COUNT; i++) {
schedulers += ", " + std::string(scheduler_to_str[i]);
}

options.manual_options = {
{"-s",
"--seed",
"RNG seed (default: 42, use random seed for < 0)",
on_seed_arg},
{"",
"--sampling-method",
"sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp, lms]"
"(default: euler for Flux/SD3/Wan, euler_a otherwise)",
"sampling method, one of [" + sample_methods + "], "
"default: euler for Flux/SD3/Wan, euler_a otherwise",
on_sample_method_arg},
{"",
"--high-noise-sampling-method",
"(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp, lms]"
" default: euler for Flux/SD3/Wan, euler_a otherwise",
"(high noise) sampling method, one of [" + sample_methods + "], "
"default: euler for Flux/SD3/Wan, euler_a otherwise",
on_high_noise_sample_method_arg},
{"",
"--scheduler",
"denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2, logit_normal, flux2, flux, beta], alias: normal=discrete, default: model-specific",
"denoiser sigma scheduler, one of [" + schedulers + "], "
"alias: normal=discrete, default: model-specific",
on_scheduler_arg},
{"",
"--sigmas",
Expand Down
4 changes: 4 additions & 0 deletions otherarch/sdcpp/include/stable-diffusion.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,8 @@ enum sample_method_t {
SAMPLE_METHOD_COUNT
};

extern SD_API const char* sample_method_to_str[];

enum scheduler_t {
DISCRETE_SCHEDULER,
KARRAS_SCHEDULER,
Expand All @@ -80,6 +82,8 @@ enum scheduler_t {
SCHEDULER_COUNT
};

extern SD_API const char* scheduler_to_str[];

enum prediction_t {
EPS_PRED,
V_PRED,
Expand Down
7 changes: 5 additions & 2 deletions otherarch/sdcpp/src/core/ggml_extend.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2408,13 +2408,14 @@ struct GGMLRunner {
GGML_ASSERT(gf != nullptr);

size_t effective_budget = max_graph_vram_bytes;
size_t free_clamp = SIZE_MAX;
if (stream_layers_enabled && max_graph_vram_bytes > 0 && runtime_backend != nullptr) {
ggml_backend_dev_t dev = ggml_backend_get_device(runtime_backend);
if (dev != nullptr && ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) {
size_t free_vram = 0, total_vram = 0;
ggml_backend_dev_memory(dev, &free_vram, &total_vram);
constexpr size_t safety_margin = 512ull * 1024 * 1024;
size_t free_clamp = (free_vram > safety_margin) ? (free_vram - safety_margin) : 0;
free_clamp = (free_vram > safety_margin) ? (free_vram - safety_margin) : 0;
if (free_clamp < effective_budget) {
LOG_DEBUG("%s clamping streaming budget: actual free VRAM %.2f MB < user cap %.2f MB",
get_desc().c_str(),
Expand All @@ -2431,7 +2432,9 @@ struct GGMLRunner {
observed_max_effective_budget_ = effective_budget;
budget_increased = true;
} else {
effective_budget = observed_max_effective_budget_;
// Keep the plan cache stable, but never plan above what is free now:
// another model or process can take VRAM after the first measurement.
effective_budget = std::min(observed_max_effective_budget_, free_clamp);
}
}

Expand Down
55 changes: 45 additions & 10 deletions otherarch/sdcpp/src/core/ggml_graph_cut.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -453,11 +453,12 @@ namespace sd::ggml_graph_cut {
if (tensor == nullptr || tensor->name[0] == '\0') {
return false;
}
return std::strncmp(tensor->name, GGML_RUNNER_CUT_PREFIX, std::strlen(GGML_RUNNER_CUT_PREFIX)) == 0;
return starts_with(tensor->name, GGML_RUNNER_CUT_PREFIX) &&
ends_with(tensor->name, GGML_RUNNER_CUT_SUFFIX);
}

std::string make_graph_cut_name(const std::string& group, const std::string& output) {
return std::string(GGML_RUNNER_CUT_PREFIX) + group + "|" + output;
return std::string(GGML_RUNNER_CUT_PREFIX) + group + "|" + output + GGML_RUNNER_CUT_SUFFIX;
}

void mark_graph_cut(ggml_tensor* tensor, const std::string& group, const std::string& output) {
Expand Down Expand Up @@ -603,7 +604,43 @@ namespace sd::ggml_graph_cut {
GGML_ASSERT(gf != nullptr);
GGML_ASSERT(graph_ctx_out != nullptr);

const size_t graph_size = segment.internal_node_indices.size() + segment.input_refs.size() + 8;
// Collect leaf inputs and internal nodes, then any tensor they
// reference that is not already represented, notably the view_src of a
// view-typed input leaf. ggml_gallocr sizes its hash set from
// n_nodes + n_leafs (plus a 25% margin that rounds down to zero for a
// one-node segment), so every distinct tensor it will hash must be
// counted here or a tiny segment overflows the hash set and aborts.
std::vector<ggml_tensor*> leaves;
std::unordered_set<ggml_tensor*> represented;
for (const auto& input : segment.input_refs) {
ggml_tensor* current_input = input_tensor(gf, input);
if (current_input == nullptr) {
continue;
}
if (represented.insert(current_input).second) {
leaves.push_back(current_input);
}
}
for (int node_idx : segment.internal_node_indices) {
represented.insert(ggml_graph_node(gf, node_idx));
}
auto add_reference = [&](ggml_tensor* tensor) {
if (tensor != nullptr && represented.insert(tensor).second) {
leaves.push_back(tensor);
}
};
for (int node_idx : segment.internal_node_indices) {
ggml_tensor* node = ggml_graph_node(gf, node_idx);
for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
add_reference(node->src[src_idx]);
}
add_reference(node->view_src);
}
for (size_t i = 0; i < leaves.size(); ++i) {
add_reference(leaves[i]->view_src);
}

const size_t graph_size = segment.internal_node_indices.size() + leaves.size() + 8;
ggml_init_params params = {
/*.mem_size =*/ggml_graph_overhead_custom(graph_size, false) + 1024,
/*.mem_buffer =*/nullptr,
Expand All @@ -614,13 +651,9 @@ namespace sd::ggml_graph_cut {
ggml_cgraph* segment_graph = ggml_new_graph_custom(graph_ctx, graph_size, false);
GGML_ASSERT(segment_graph != nullptr);

for (const auto& input : segment.input_refs) {
ggml_tensor* current_input = input_tensor(gf, input);
if (current_input == nullptr) {
continue;
}
for (ggml_tensor* leaf : leaves) {
GGML_ASSERT(segment_graph->n_leafs < segment_graph->size);
segment_graph->leafs[segment_graph->n_leafs++] = current_input;
segment_graph->leafs[segment_graph->n_leafs++] = leaf;
}

for (int output_node_index : segment.output_node_indices) {
Expand Down Expand Up @@ -751,7 +784,9 @@ namespace sd::ggml_graph_cut {

plan.has_cuts = true;
std::string full_name(node->name);
std::string payload = full_name.substr(std::strlen(GGML_RUNNER_CUT_PREFIX));
size_t prefix_len = std::strlen(GGML_RUNNER_CUT_PREFIX);
size_t suffix_len = std::strlen(GGML_RUNNER_CUT_SUFFIX);
std::string payload = full_name.substr(prefix_len, full_name.size() - prefix_len - suffix_len);
size_t sep = payload.find('|');
std::string group = sep == std::string::npos ? payload : payload.substr(0, sep);

Expand Down
1 change: 1 addition & 0 deletions otherarch/sdcpp/src/core/ggml_graph_cut.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ namespace sd::ggml_graph_cut {
};

static constexpr const char* GGML_RUNNER_CUT_PREFIX = "ggml_runner_cut:";
static constexpr const char* GGML_RUNNER_CUT_SUFFIX = "|";

struct MaxVramAssignment {
float default_gib = 0.f;
Expand Down
113 changes: 108 additions & 5 deletions otherarch/sdcpp/src/model/vae/tae.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -426,10 +426,10 @@ class TinyVideoDecoder : public UnaryBlock {
static const int num_layers = 3;
int channels[num_layers + 1] = {256, 128, 64, 64};
int patch_size = 1;
int t_upscale = 1;
bool is_wide = false;

public:
int t_upscale = 1;
TinyVideoDecoder(int z_channels = 4, int patch_size = 1, std::vector<bool> time_upscale = {false, true, true}, bool is_wide = false)
: z_channels(z_channels), patch_size(patch_size), is_wide(is_wide) {
t_upscale = 1;
Expand Down Expand Up @@ -536,6 +536,10 @@ class TAEHV : public GGMLBlock {
patch = 4;
time_downscale = {true, true, true};
time_upscale = {true, true, true};
} else if (sd_version_is_minimax_h3(version)) {
z_channels = 24;
patch = 2;
time_downscale = {true, true, false};
}
blocks["decoder"] = std::shared_ptr<GGMLBlock>(new TinyVideoDecoder(z_channels, patch, time_upscale, is_wide));
if (!decode_only) {
Expand All @@ -545,24 +549,123 @@ class TAEHV : public GGMLBlock {

ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
auto decoder = std::dynamic_pointer_cast<TinyVideoDecoder>(blocks["decoder"]);
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) {
// (W, H, C, T) -> (W, H, T, C)
z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 1, 3, 2));
}
auto result = decoder->forward(ctx, z);
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {

if (sd_version_is_minimax_h3(version)) {
int64_t num_frames = result->ne[3];
int64_t chunk_frames = 5 * decoder->t_upscale;
int64_t pad = (chunk_frames - (num_frames % chunk_frames)) % chunk_frames;

result = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, result, 0, 0, 0, 0, 0, 0, 0, pad, false, false);

int64_t num_chunks = (num_frames + pad) / chunk_frames;
auto to_trim = decoder->t_upscale - 1;
std::vector<ggml_tensor*> to_concat = {};
for (int i = 0; i < num_chunks; i++) {
auto chunk = ggml_view_4d(ctx->ggml_ctx, result,
result->ne[0], result->ne[1], result->ne[2], chunk_frames - to_trim,
result->nb[1], result->nb[2], result->nb[3],
i * chunk_frames * result->nb[3]);
to_concat.push_back(chunk);
}
result = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3);
result = ggml_view_4d(ctx->ggml_ctx, result,
result->ne[0], result->ne[1], result->ne[2],
result->ne[3] - decoder->t_upscale * 3,
result->nb[1], result->nb[2], result->nb[3], 0);
}

if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) {
// (W, H, T, C) -> (W, H, C, T)
result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2));
}
return result;
}

ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
ggml_tensor* encode_h3(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {

int64_t num_frames = x->ne[3];
int64_t pad = (17 - (num_frames % 17)) % 17;

if (pad > 0) {
auto last_frame = ggml_view_4d(ctx->ggml_ctx, x,
x->ne[0], x->ne[1], x->ne[2], 1,
x->nb[1], x->nb[2], x->nb[3],
(num_frames - 1) * x->nb[3]);
for (int i = 0; i < pad; i++) {
x = ggml_concat(ctx->ggml_ctx, x, last_frame, 3);
}
}

int64_t T_padded = x->ne[3];
int64_t num_chunks = T_padded / 17;

auto zero_frame = ggml_view_4d(ctx->ggml_ctx, x,
x->ne[0], x->ne[1], x->ne[2], 1,
x->nb[1], x->nb[2], x->nb[3], 0);
auto zeros_1 = ggml_scale(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, zero_frame), 0.0f);
auto zeros_3 = zeros_1;
for (int i = 1; i < 3; i++) {
zeros_3 = ggml_concat(ctx->ggml_ctx, zeros_3, zeros_1, 3);
}
ggml_tensor* out = nullptr;
if (false) {
std::vector<ggml_tensor*> to_concat = {};
for (int i = 0; i < num_chunks; i++) {
auto chunk = ggml_view_4d(ctx->ggml_ctx, x,
x->ne[0], x->ne[1], x->ne[2], 17,
x->nb[1], x->nb[2], x->nb[3],
i * 17 * x->nb[3]);

auto chunk_padded = ggml_concat(ctx->ggml_ctx, zeros_3, chunk, 3);

to_concat.push_back(chunk_padded);
}
ggml_tensor* x_in = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3);
out = encoder->forward(ctx, x_in);
} else {
std::vector<ggml_tensor*> to_concat = {};
for (int i = 0; i < num_chunks; i++) {
auto chunk = ggml_view_4d(ctx->ggml_ctx, x,
x->ne[0], x->ne[1], x->ne[2], 17,
x->nb[1], x->nb[2], x->nb[3],
i * 17 * x->nb[3]);

auto chunk_padded = ggml_concat(ctx->ggml_ctx, zeros_3, chunk, 3);

auto chunk_out = encoder->forward(ctx, chunk_padded);
// auto chunk_out = encoder->forward_seq(ctx, chunk_padded); // ~same vram usage, and straight-up slower. it's already sequential enough

to_concat.push_back(chunk_out);
}
out = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3);
}

// Return x[:, :-3] - drop the last 3 elements in the T dimension
int64_t out_T = out->ne[3];
out = ggml_view_4d(ctx->ggml_ctx, out,
out->ne[0], out->ne[1], out->ne[2], out_T - 3,
out->nb[1], out->nb[2], out->nb[3], 0);

return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 1, 3, 2));
}

ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || (sd_version_is_minimax_h3(version) && x->ne[3] > 1)) {
// (W, H, T, C) -> (W, H, C, T)
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
}
if (sd_version_is_minimax_h3(version)) {
return encode_h3(ctx, x);
}

auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);

int64_t num_frames = x->ne[3];
if (num_frames % encoder->t_downscale) {
// pad to multiple of encoder->t_downscale at the end
Expand Down
Loading