feat: add configurable reference image processing for edit models (#1780)

This commit is contained in:
stduhpf
2026-07-14 16:58:22 +02:00
committed by GitHub
parent b5d812008e
commit 74bce049d0
15 changed files with 551 additions and 159 deletions

View File

@@ -10,6 +10,36 @@
#include "model/common/rope.hpp"
#include "model_manager.h"
enum class RefImageResizeMode {
NONE,
LONGEST_SIDE,
AREA,
};
struct RefImageParams {
bool pass_to_vlm = false;
bool pass_to_dit = true;
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED;
bool force_ref_timestep_zero = false;
bool resize_before_vae = true;
int vae_input_max_pixels = -1;
RefImageResizeMode vlm_resize_mode = RefImageResizeMode::AREA;
int vlm_min_size = -1;
int vlm_max_size = -1;
};
const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
{"flux_kontext", {false, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::NONE, -1, -1}},
{"longcat", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"flux2", {false, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::NONE, -1, -1}},
{"qwen", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"qwen_layered", {true, true, Rope::RefIndexMode::DECREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"z_image_omni", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"krea2_ostris_edit", {true, true, Rope::RefIndexMode::INCREASE, true, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"krea2_edit", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, 768, 768}},
{"cosmos_reference", {false, true, Rope::RefIndexMode::INCREASE, false, false, -1, RefImageResizeMode::NONE, -1, -1}},
};
struct UNetDiffusionExtra {
int num_video_frames = -1;
const std::vector<sd::Tensor<float>>* controls = nullptr;
@@ -74,7 +104,7 @@ struct DiffusionParams {
const sd::Tensor<float>* c_concat = nullptr;
const sd::Tensor<float>* y = nullptr;
const std::vector<sd::Tensor<float>>* ref_latents = nullptr;
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED;
RefImageParams ref_image_params = {false, false, Rope::RefIndexMode::FIXED, false};
DiffusionExtraParams extra = std::monostate{};
};