sd: sync to master-391-5865b5e (#1878)

This commit is contained in:
Wagner Bruna
2025-12-08 08:53:52 -03:00
committed by GitHub
parent 242ae8b8f3
commit 801840d3bd
23 changed files with 250 additions and 216 deletions
+126 -125
View File
@@ -393,13 +393,6 @@ public:
}
auto& tensor_storage_map = model_loader.get_tensor_storage_map();
for (auto& [name, tensor_storage] : tensor_storage_map) {
if (contains(name, "llm") &&
ends_with(name, "weight") &&
(tensor_storage.type == GGML_TYPE_F32 || tensor_storage.type == GGML_TYPE_BF16)) {
tensor_storage.expected_type = GGML_TYPE_F16;
}
}
LOG_INFO("Version: %s ", model_version_to_str[version]);
@@ -834,7 +827,7 @@ public:
return false;
}
// LOG_DEBUG("model size = %.2fMB", total_size / 1024.0 / 1024.0);
LOG_DEBUG("finished loaded file");
{
size_t clip_params_mem_size = cond_stage_model->get_params_buffer_size();
@@ -909,8 +902,59 @@ public:
ggml_backend_is_cpu(clip_backend) ? "RAM" : "VRAM");
}
if (sd_ctx_params->prediction != DEFAULT_PRED) {
switch (sd_ctx_params->prediction) {
// init denoiser
{
prediction_t pred_type = sd_ctx_params->prediction;
float flow_shift = sd_ctx_params->flow_shift;
if (pred_type == PREDICTION_COUNT) {
if (sd_version_is_sd2(version)) {
// check is_using_v_parameterization_for_sd2
if (is_using_v_parameterization_for_sd2(ctx, sd_version_is_inpaint(version))) {
pred_type = V_PRED;
} else {
pred_type = 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_qwen_image(version) ||
sd_version_is_z_image(version)) {
pred_type = FLOW_PRED;
if (flow_shift == INFINITY) {
if (sd_version_is_wan(version)) {
flow_shift = 5.f;
} else {
flow_shift = 3.f;
}
}
} else if (sd_version_is_flux(version)) {
pred_type = FLUX_FLOW_PRED;
if (flow_shift == INFINITY) {
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")) {
flow_shift = 1.15f;
}
}
}
} else if (sd_version_is_flux2(version)) {
pred_type = FLUX2_FLOW_PRED;
} else {
pred_type = EPS_PRED;
}
}
switch (pred_type) {
case EPS_PRED:
LOG_INFO("running in eps-prediction mode");
break;
@@ -922,22 +966,14 @@ public:
LOG_INFO("running in v-prediction EDM mode");
denoiser = std::make_shared<EDMVDenoiser>();
break;
case SD3_FLOW_PRED: {
case FLOW_PRED: {
LOG_INFO("running in FLOW mode");
float shift = sd_ctx_params->flow_shift;
if (shift == INFINITY) {
shift = 3.0;
}
denoiser = std::make_shared<DiscreteFlowDenoiser>(shift);
denoiser = std::make_shared<DiscreteFlowDenoiser>(flow_shift);
break;
}
case FLUX_FLOW_PRED: {
LOG_INFO("running in Flux FLOW mode");
float shift = sd_ctx_params->flow_shift;
if (shift == INFINITY) {
shift = 3.0;
}
denoiser = std::make_shared<FluxFlowDenoiser>(shift);
denoiser = std::make_shared<FluxFlowDenoiser>(flow_shift);
break;
}
case FLUX2_FLOW_PRED: {
@@ -946,93 +982,21 @@ public:
break;
}
default: {
LOG_ERROR("Unknown parametrization %i", sd_ctx_params->prediction);
LOG_ERROR("Unknown predition type %i", pred_type);
ggml_free(ctx);
return false;
}
}
} else {
if (sd_version_is_sd2(version)) {
// check is_using_v_parameterization_for_sd2
if (is_using_v_parameterization_for_sd2(ctx, sd_version_is_inpaint(version))) {
is_using_v_parameterization = true;
}
} 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
is_using_edm_v_parameterization = true;
}
if (tensor_storage_map.find("v_pred") != tensor_storage_map.end()) {
is_using_v_parameterization = true;
}
} else if (version == VERSION_SVD) {
// TODO: V_PREDICTION_EDM
is_using_v_parameterization = true;
}
if (sd_version_is_sd3(version)) {
LOG_INFO("running in FLOW mode");
float shift = sd_ctx_params->flow_shift;
if (shift == INFINITY) {
shift = 3.0;
auto comp_vis_denoiser = std::dynamic_pointer_cast<CompVisDenoiser>(denoiser);
if (comp_vis_denoiser) {
for (int i = 0; i < TIMESTEPS; i++) {
comp_vis_denoiser->sigmas[i] = std::sqrt((1 - ((float*)alphas_cumprod_tensor->data)[i]) / ((float*)alphas_cumprod_tensor->data)[i]);
comp_vis_denoiser->log_sigmas[i] = std::log(comp_vis_denoiser->sigmas[i]);
}
denoiser = std::make_shared<DiscreteFlowDenoiser>(shift);
} else if (sd_version_is_flux(version)) {
LOG_INFO("running in Flux FLOW mode");
float shift = sd_ctx_params->flow_shift;
if (shift == INFINITY) {
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")) {
shift = 1.15f;
}
}
}
denoiser = std::make_shared<FluxFlowDenoiser>(shift);
} else if (sd_version_is_flux2(version)) {
LOG_INFO("running in Flux2 FLOW mode");
denoiser = std::make_shared<Flux2FlowDenoiser>();
} else if (sd_version_is_wan(version)) {
LOG_INFO("running in FLOW mode");
float shift = sd_ctx_params->flow_shift;
if (shift == INFINITY) {
shift = 5.0;
}
denoiser = std::make_shared<DiscreteFlowDenoiser>(shift);
} else if (sd_version_is_qwen_image(version)) {
LOG_INFO("running in FLOW mode");
float shift = sd_ctx_params->flow_shift;
if (shift == INFINITY) {
shift = 3.0;
}
denoiser = std::make_shared<DiscreteFlowDenoiser>(shift);
} else if (sd_version_is_z_image(version)) {
LOG_INFO("running in FLOW mode");
float shift = sd_ctx_params->flow_shift;
if (shift == INFINITY) {
shift = 3.0f;
}
denoiser = std::make_shared<DiscreteFlowDenoiser>(shift);
} else if (is_using_v_parameterization) {
LOG_INFO("running in v-prediction mode");
denoiser = std::make_shared<CompVisVDenoiser>();
} else if (is_using_edm_v_parameterization) {
LOG_INFO("running in v-prediction EDM mode");
denoiser = std::make_shared<EDMVDenoiser>();
} else {
LOG_INFO("running in eps-prediction mode");
}
}
auto comp_vis_denoiser = std::dynamic_pointer_cast<CompVisDenoiser>(denoiser);
if (comp_vis_denoiser) {
for (int i = 0; i < TIMESTEPS; i++) {
comp_vis_denoiser->sigmas[i] = std::sqrt((1 - ((float*)alphas_cumprod_tensor->data)[i]) / ((float*)alphas_cumprod_tensor->data)[i]);
comp_vis_denoiser->log_sigmas[i] = std::log(comp_vis_denoiser->sigmas[i]);
}
}
LOG_DEBUG("finished loaded file");
ggml_free(ctx);
use_tiny_autoencoder = use_tiny_autoencoder && !sd_ctx_params->tae_preview_only;
return true;
@@ -1164,6 +1128,12 @@ public:
lora_state_diff[lora_name] -= curr_multiplier;
}
if (lora_state_diff.empty()) {
return;
}
LOG_INFO("apply lora immediately");
size_t rm = lora_state_diff.size() - lora_state.size();
if (rm != 0) {
LOG_INFO("attempting to apply %lu LoRAs (removing %lu applied LoRAs)", lora_state.size(), rm);
@@ -1193,6 +1163,10 @@ public:
cond_stage_lora_models.clear();
diffusion_lora_models.clear();
first_stage_lora_models.clear();
if (lora_state.empty()) {
return;
}
LOG_INFO("apply lora at runtime");
if (cond_stage_model) {
std::vector<std::shared_ptr<LoraModel>> lora_models;
auto lora_state_diff = lora_state;
@@ -1333,10 +1307,8 @@ public:
#else
int64_t t0 = ggml_time_ms();
if (apply_lora_immediately) {
LOG_INFO("apply lora immediately");
apply_loras_immediately(lora_f2m);
} else {
LOG_INFO("apply at runtime");
apply_loras_at_runtime(lora_f2m);
}
int64_t t1 = ggml_time_ms();
@@ -1856,8 +1828,11 @@ public:
std::vector<struct ggml_tensor*> controls;
if (control_hint != nullptr && control_net != nullptr) {
control_net->compute(n_threads, noised_input, control_hint, timesteps, cond.c_crossattn, cond.c_vector);
controls = control_net->controls;
if (control_net->compute(n_threads, noised_input, control_hint, timesteps, cond.c_crossattn, cond.c_vector)) {
controls = control_net->controls;
} else {
LOG_ERROR("controlnet compute failed");
}
// print_ggml_tensor(controls[12]);
// GGML_ASSERT(0);
}
@@ -1889,9 +1864,12 @@ public:
bool skip_model = easycache_before_condition(active_condition, *active_output);
if (!skip_model) {
work_diffusion_model->compute(n_threads,
diffusion_params,
active_output);
if (!work_diffusion_model->compute(n_threads,
diffusion_params,
active_output)) {
LOG_ERROR("diffusion model compute failed");
return nullptr;
}
easycache_after_condition(active_condition, *active_output);
}
@@ -1901,8 +1879,11 @@ public:
if (has_unconditioned) {
// uncond
if (!current_step_skipped && control_hint != nullptr && control_net != nullptr) {
control_net->compute(n_threads, noised_input, control_hint, timesteps, uncond.c_crossattn, uncond.c_vector);
controls = control_net->controls;
if (control_net->compute(n_threads, noised_input, control_hint, timesteps, uncond.c_crossattn, uncond.c_vector)) {
controls = control_net->controls;
} else {
LOG_ERROR("controlnet compute failed");
}
}
current_step_skipped = easycache_step_is_skipped();
diffusion_params.controls = controls;
@@ -1911,9 +1892,12 @@ public:
diffusion_params.y = uncond.c_vector;
bool skip_uncond = easycache_before_condition(&uncond, out_uncond);
if (!skip_uncond) {
work_diffusion_model->compute(n_threads,
diffusion_params,
&out_uncond);
if (!work_diffusion_model->compute(n_threads,
diffusion_params,
&out_uncond)) {
LOG_ERROR("diffusion model compute failed");
return nullptr;
}
easycache_after_condition(&uncond, out_uncond);
}
negative_data = (float*)out_uncond->data;
@@ -1926,9 +1910,12 @@ public:
diffusion_params.y = img_cond.c_vector;
bool skip_img_cond = easycache_before_condition(&img_cond, out_img_cond);
if (!skip_img_cond) {
work_diffusion_model->compute(n_threads,
diffusion_params,
&out_img_cond);
if (!work_diffusion_model->compute(n_threads,
diffusion_params,
&out_img_cond)) {
LOG_ERROR("diffusion model compute failed");
return nullptr;
}
easycache_after_condition(&img_cond, out_img_cond);
}
img_cond_data = (float*)out_img_cond->data;
@@ -1945,9 +1932,12 @@ public:
diffusion_params.c_concat = cond.c_concat;
diffusion_params.y = cond.c_vector;
diffusion_params.skip_layers = skip_layers;
work_diffusion_model->compute(n_threads,
diffusion_params,
&out_skip);
if (!work_diffusion_model->compute(n_threads,
diffusion_params,
&out_skip)) {
LOG_ERROR("diffusion model compute failed");
return nullptr;
}
}
skip_layer_data = (float*)out_skip->data;
}
@@ -2010,7 +2000,15 @@ public:
return denoised;
};
sample_k_diffusion(method, denoise, work_ctx, x, sigmas, sampler_rng, eta);
if (!sample_k_diffusion(method, denoise, work_ctx, x, sigmas, sampler_rng, eta)) {
LOG_ERROR("Diffusion model sampling failed");
if (control_net) {
control_net->free_control_ctx();
control_net->free_compute_buffer();
}
diffusion_model->free_compute_buffer();
return NULL;
}
if (easycache_enabled) {
size_t total_steps = sigmas.size() > 0 ? sigmas.size() - 1 : 0;
@@ -2565,7 +2563,6 @@ enum scheduler_t str_to_scheduler(const char* str) {
}
const char* prediction_to_str[] = {
"default",
"eps",
"v",
"edm_v",
@@ -2651,7 +2648,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
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 = DEFAULT_PRED;
sd_ctx_params->prediction = PREDICTION_COUNT;
sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO;
sd_ctx_params->offload_params_to_cpu = false;
sd_ctx_params->keep_clip_on_cpu = false;
@@ -3237,10 +3234,14 @@ sd_image_t* generate_image_internal(sd_ctx_t* sd_ctx,
nullptr,
1.0f,
easycache_params);
// print_ggml_tensor(x_0);
int64_t sampling_end = ggml_time_ms();
LOG_INFO("sampling completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000);
final_latents.push_back(x_0);
int64_t sampling_end = ggml_time_ms();
if (x_0 != nullptr) {
// print_ggml_tensor(x_0);
LOG_INFO("sampling completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000);
final_latents.push_back(x_0);
} else {
LOG_ERROR("sampling for image %d/%d failed after %.2fs", b + 1, batch_count, (sampling_end - sampling_start) * 1.0f / 1000);
}
}
if (sd_ctx->sd->free_params_immediately) {