sync with sd.cpp

This commit is contained in:
Concedo
2025-06-30 00:10:51 +08:00
parent e5af9b5ea9
commit 186227fc26
8 changed files with 234 additions and 82 deletions
+142 -5
View File
@@ -159,7 +159,10 @@ public:
bool clip_on_cpu,
bool control_net_cpu,
bool vae_on_cpu,
bool diffusion_flash_attn) {
bool diffusion_flash_attn,
bool chroma_use_dit_mask,
bool chroma_use_t5_mask,
int chroma_t5_mask_pad) {
use_tiny_autoencoder = taesd_path.size() > 0;
std::string taesd_path_fixed = taesd_path;
is_loaded_chroma = false;
@@ -391,11 +394,11 @@ public:
}
}
if (is_chroma) {
cond_stage_model = std::make_shared<PixArtCLIPEmbedder>(clip_backend, model_loader.tensor_storages_types);
cond_stage_model = std::make_shared<PixArtCLIPEmbedder>(clip_backend, model_loader.tensor_storages_types, -1, chroma_use_t5_mask, chroma_t5_mask_pad);
} else {
cond_stage_model = std::make_shared<FluxCLIPEmbedder>(clip_backend, model_loader.tensor_storages_types);
}
diffusion_model = std::make_shared<FluxModel>(backend, model_loader.tensor_storages_types, version, diffusion_flash_attn);
diffusion_model = std::make_shared<FluxModel>(backend, model_loader.tensor_storages_types, version, diffusion_flash_attn, chroma_use_dit_mask);
} else {
if (id_embeddings_path.find("v2") != std::string::npos) {
cond_stage_model = std::make_shared<FrozenCLIPEmbedderWithCustomWords>(clip_backend, model_loader.tensor_storages_types, embeddings_path, version, PM_VERSION_2);
@@ -1337,7 +1340,10 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
bool keep_clip_on_cpu,
bool keep_control_net_cpu,
bool keep_vae_on_cpu,
bool diffusion_flash_attn) {
bool diffusion_flash_attn,
bool chroma_use_dit_mask,
bool chroma_use_t5_mask,
int chroma_t5_mask_pad) {
sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t));
if (sd_ctx == NULL) {
return NULL;
@@ -1379,7 +1385,10 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
keep_clip_on_cpu,
keep_control_net_cpu,
keep_vae_on_cpu,
diffusion_flash_attn)) {
diffusion_flash_attn,
chroma_use_dit_mask,
chroma_use_t5_mask,
chroma_t5_mask_pad)) {
delete sd_ctx->sd;
sd_ctx->sd = NULL;
free(sd_ctx);
@@ -2231,5 +2240,133 @@ SD_API sd_image_t* img2vid(sd_ctx_t* sd_ctx,
LOG_INFO("img2vid completed in %.2fs", (t3 - t0) * 1.0f / 1000);
return result_images;
}
sd_image_t* edit(sd_ctx_t* sd_ctx,
sd_image_t* ref_images,
int ref_images_count,
const char* prompt_c_str,
const char* negative_prompt_c_str,
int clip_skip,
float cfg_scale,
float guidance,
float eta,
int width,
int height,
sample_method_t sample_method,
int sample_steps,
float strength,
int64_t seed,
int batch_count,
const sd_image_t* control_cond,
float control_strength,
float style_ratio,
bool normalize_input,
int* skip_layers = NULL,
size_t skip_layers_count = 0,
float slg_scale = 0,
float skip_layer_start = 0.01,
float skip_layer_end = 0.2) {
std::vector<int> skip_layers_vec(skip_layers, skip_layers + skip_layers_count);
LOG_DEBUG("edit %dx%d", width, height);
if (sd_ctx == NULL) {
return NULL;
}
if (ref_images_count <= 0) {
LOG_ERROR("ref images count should > 0");
return NULL;
}
struct ggml_init_params params;
params.mem_size = static_cast<size_t>(30 * 1024 * 1024); // 10 MB
params.mem_size += width * height * 3 * sizeof(float) * 3 * ref_images_count;
params.mem_size *= batch_count;
params.mem_buffer = NULL;
params.no_alloc = false;
// LOG_DEBUG("mem_size %u ", params.mem_size);
struct ggml_context* work_ctx = ggml_init(params);
if (!work_ctx) {
LOG_ERROR("ggml_init() failed");
return NULL;
}
if (seed < 0) {
srand((int)time(NULL));
seed = rand();
}
sd_ctx->sd->rng->manual_seed(seed);
int C = 4;
if (sd_version_is_sd3(sd_ctx->sd->version)) {
C = 16;
} else if (sd_version_is_flux(sd_ctx->sd->version)) {
C = 16;
}
int W = width / 8;
int H = height / 8;
ggml_tensor* init_latent = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, W, H, C, 1);
if (sd_version_is_sd3(sd_ctx->sd->version)) {
ggml_set_f32(init_latent, 0.0609f);
} else if (sd_version_is_flux(sd_ctx->sd->version)) {
ggml_set_f32(init_latent, 0.1159f);
} else {
ggml_set_f32(init_latent, 0.f);
}
size_t t0 = ggml_time_ms();
std::vector<struct ggml_tensor*> ref_latents;
for (int i = 0; i < ref_images_count; i++) {
ggml_tensor* img = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, ref_images[i].width, ref_images[i].height, 3, 1);
sd_image_to_tensor(ref_images[i].data, img);
ggml_tensor* latent = NULL;
if (!sd_ctx->sd->use_tiny_autoencoder) {
ggml_tensor* moments = sd_ctx->sd->encode_first_stage(work_ctx, img);
latent = sd_ctx->sd->get_first_stage_encoding(work_ctx, moments);
} else {
latent = sd_ctx->sd->encode_first_stage(work_ctx, img);
}
ref_latents.push_back(latent);
}
size_t t1 = ggml_time_ms();
LOG_INFO("encode_first_stage completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
std::vector<float> sigmas = sd_ctx->sd->denoiser->get_sigmas(sample_steps);
sd_image_t* result_images = generate_image(sd_ctx,
work_ctx,
init_latent,
prompt_c_str,
negative_prompt_c_str,
clip_skip,
cfg_scale,
guidance,
eta,
width,
height,
sample_method,
sigmas,
seed,
batch_count,
control_cond,
control_strength,
style_ratio,
normalize_input,
"",
ref_latents,
skip_layers_vec,
slg_scale,
skip_layer_start,
skip_layer_end,
NULL);
size_t t2 = ggml_time_ms();
LOG_INFO("edit completed in %.2fs", (t2 - t0) * 1.0f / 1000);
return result_images;
}