feat: add SeFi-Image support (#1707)

This commit is contained in:
fszontagh
2026-06-28 16:49:24 +02:00
committed by GitHub
parent f54e45e81c
commit 03e9a22f4d
16 changed files with 736 additions and 17 deletions

View File

@@ -8,6 +8,7 @@
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
#include "model/diffusion/model.hpp"
#include "model/diffusion/sefi_image.hpp"
#include "model_loader.h"
#define FLUX_GRAPH_SIZE 10240
@@ -26,6 +27,9 @@ namespace Flux {
struct FluxConfig {
SDVersion version = VERSION_FLUX;
bool is_chroma = false;
bool is_sefi = false;
int64_t semantic_channels = 0;
float sefi_delta_t = 0.1f;
int patch_size = 2;
int64_t in_channels = 64;
int64_t out_channels = 64;
@@ -88,6 +92,21 @@ namespace Flux {
config.share_modulation = true;
config.ref_index_scale = 10.f;
config.use_mlp_silu_act = true;
} else if (sd_version_is_sefi_image(version)) {
config.is_sefi = true;
config.semantic_channels = 16;
config.in_channels = 128 + config.semantic_channels;
config.patch_size = 1;
config.out_channels = 128 + config.semantic_channels;
config.mlp_ratio = 3.f;
config.theta = 2000;
config.axes_dim = {32, 32, 32, 32};
config.vec_in_dim = 0;
config.qkv_bias = false;
config.disable_bias = true;
config.share_modulation = true;
config.ref_index_scale = 10.f;
config.use_mlp_silu_act = true;
} else if (sd_version_is_longcat(version)) {
config.context_in_dim = 3584;
config.vec_in_dim = 0;
@@ -723,8 +742,8 @@ namespace Flux {
auto m = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, 2 * hidden_size]
auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, 2, 0);
shift = m_vec[0]; // [N, hidden_size]
scale = m_vec[1]; // [N, hidden_size]
shift = m_vec[0];
scale = m_vec[1];
}
x = Flux::modulate(ctx->ggml_ctx, norm_final->forward(ctx, x), shift, scale);
@@ -902,6 +921,8 @@ namespace Flux {
}
if (config.is_chroma) {
blocks["distilled_guidance_layer"] = std::make_shared<ChromaApproximator>(config.in_dim, config.hidden_size);
} else if (config.is_sefi) {
blocks["dual_time_embed"] = std::make_shared<SefiImage::SefiDualTimestepEmbeddings>(256, config.hidden_size);
} else {
blocks["time_in"] = std::make_shared<MLPEmbedder>(256, config.hidden_size, !config.disable_bias);
if (config.vec_in_dim > 0) {
@@ -1027,6 +1048,11 @@ namespace Flux {
if (y != nullptr) {
txt_img_mask = ggml_pad(ctx->ggml_ctx, y, static_cast<int>(img->ne[1]), 0, 0, 0);
}
} else if (config.is_sefi) {
auto dual_time_embed = std::dynamic_pointer_cast<SefiImage::SefiDualTimestepEmbeddings>(blocks["dual_time_embed"]);
auto timestep_sem = ggml_view_1d(ctx->ggml_ctx, timesteps, 1, 0);
auto timestep_tex = ggml_view_1d(ctx->ggml_ctx, timesteps, 1, ggml_element_size(timesteps));
vec = dual_time_embed->forward(ctx, timestep_sem, timestep_tex);
} else {
auto time_in = std::dynamic_pointer_cast<MLPEmbedder>(blocks["time_in"]);
vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1000.f));
@@ -1500,7 +1526,7 @@ namespace Flux {
set_backend_tensor_data(mod_index_arange, mod_index_arange_vec.data());
}
std::set<int> txt_arange_dims;
if (sd_version_is_flux2(version)) {
if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
txt_arange_dims = {3};
increase_ref_index = true;
} else if (version == VERSION_OVIS_IMAGE) {

View File

@@ -0,0 +1,91 @@
#ifndef __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__
#define __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__
#include <memory>
#include "model/common/block.hpp"
namespace SefiImage {
struct SefiImageConfig {
int64_t semantic_channels = 16;
int64_t texture_latent_channels = 32;
int64_t timestep_guidance_in_dim = 256;
int64_t hidden_size = 3072;
float timestep_shift_alpha = 0.3f;
float delta_t = 0.1f;
int64_t packed_texture_channels(int patch_size) const {
return texture_latent_channels * patch_size * patch_size;
}
int64_t packed_input_channels(int patch_size) const {
return semantic_channels + packed_texture_channels(patch_size);
}
static SefiImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix) {
SefiImageConfig config;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (ends_with(name, "dual_time_embed.semantic_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
config.timestep_guidance_in_dim = tensor_storage.ne[0];
config.hidden_size = tensor_storage.ne[1] * 2;
}
}
LOG_DEBUG("sefi_image: semantic_channels = %" PRId64 ", texture_latent_channels = %" PRId64 ", hidden_size = %" PRId64,
config.semantic_channels,
config.texture_latent_channels,
config.hidden_size);
return config;
}
};
struct SefiTimestepEmbedding : public GGMLBlock {
public:
SefiTimestepEmbedding(int64_t in_channels, int64_t time_embed_dim) {
blocks["linear_1"] = std::shared_ptr<GGMLBlock>(new Linear(in_channels, time_embed_dim, false));
blocks["linear_2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, time_embed_dim, false));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* sample) {
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
sample = linear_1->forward(ctx, sample);
sample = ggml_silu_inplace(ctx->ggml_ctx, sample);
sample = linear_2->forward(ctx, sample);
return sample;
}
};
struct SefiDualTimestepEmbeddings : public GGMLBlock {
public:
SefiDualTimestepEmbeddings(int64_t in_channels, int64_t embedding_dim) {
GGML_ASSERT(embedding_dim % 2 == 0);
int64_t half_dim = embedding_dim / 2;
blocks["semantic_embedder"] = std::make_shared<SefiTimestepEmbedding>(in_channels, half_dim);
blocks["texture_embedder"] = std::make_shared<SefiTimestepEmbedding>(in_channels, half_dim);
timestep_guidance_in_dim = in_channels;
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* timestep_sem,
ggml_tensor* timestep_tex) {
auto semantic_embedder = std::dynamic_pointer_cast<SefiTimestepEmbedding>(blocks["semantic_embedder"]);
auto texture_embedder = std::dynamic_pointer_cast<SefiTimestepEmbedding>(blocks["texture_embedder"]);
auto sem_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_sem, timestep_guidance_in_dim, 10000, 1.f);
auto tex_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_tex, timestep_guidance_in_dim, 10000, 1.f);
auto sem_emb = semantic_embedder->forward(ctx, sem_proj);
auto tex_emb = texture_embedder->forward(ctx, tex_proj);
return ggml_concat(ctx->ggml_ctx, sem_emb, tex_emb, 0);
}
private:
int64_t timestep_guidance_in_dim = 256;
};
} // namespace SefiImage
#endif // __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__