feat: add FLUX.1 Kontext dev support (#707)

* Kontext support
* add edit mode

---------

Co-authored-by: leejet <leejet714@gmail.com>
This commit is contained in:
stduhpf
2025-06-29 04:08:53 +02:00
committed by GitHub
parent 10c6501bd0
commit c9b5735116
5 changed files with 342 additions and 31 deletions

View File

@@ -618,7 +618,7 @@ public:
int64_t t0 = ggml_time_ms();
struct ggml_tensor* out = ggml_dup_tensor(work_ctx, x_t);
diffusion_model->compute(n_threads, x_t, timesteps, c, concat, NULL, NULL, -1, {}, 0.f, &out);
diffusion_model->compute(n_threads, x_t, timesteps, c, concat, NULL, NULL, {}, -1, {}, 0.f, &out);
diffusion_model->free_compute_buffer();
double result = 0.f;
@@ -800,6 +800,7 @@ public:
const std::vector<float>& sigmas,
int start_merge_step,
SDCondition id_cond,
std::vector<ggml_tensor*> ref_latents = {},
std::vector<int> skip_layers = {},
float slg_scale = 0,
float skip_layer_start = 0.01,
@@ -887,6 +888,7 @@ public:
cond.c_concat,
cond.c_vector,
guidance_tensor,
ref_latents,
-1,
controls,
control_strength,
@@ -899,6 +901,7 @@ public:
cond.c_concat,
id_cond.c_vector,
guidance_tensor,
ref_latents,
-1,
controls,
control_strength,
@@ -919,6 +922,7 @@ public:
uncond.c_concat,
uncond.c_vector,
guidance_tensor,
ref_latents,
-1,
controls,
control_strength,
@@ -939,6 +943,7 @@ public:
cond.c_concat,
cond.c_vector,
guidance_tensor,
ref_latents,
-1,
controls,
control_strength,
@@ -1209,6 +1214,7 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx,
float style_ratio,
bool normalize_input,
std::string input_id_images_path,
std::vector<ggml_tensor*> ref_latents,
std::vector<int> skip_layers = {},
float slg_scale = 0,
float skip_layer_start = 0.01,
@@ -1466,6 +1472,7 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx,
sigmas,
start_merge_step,
id_cond,
ref_latents,
skip_layers,
slg_scale,
skip_layer_start,
@@ -1618,6 +1625,7 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
style_ratio,
normalize_input,
input_id_images_path_c_str,
{},
skip_layers_vec,
slg_scale,
skip_layer_start,
@@ -1798,6 +1806,7 @@ sd_image_t* img2img(sd_ctx_t* sd_ctx,
style_ratio,
normalize_input,
input_id_images_path_c_str,
{},
skip_layers_vec,
slg_scale,
skip_layer_start,
@@ -1943,3 +1952,132 @@ SD_API sd_image_t* img2vid(sd_ctx_t* sd_ctx,
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;
}