Compare commits

..

6 Commits

Author SHA1 Message Date
Cyberhan123
c1790754d3
feat: enhanced third-party integrations (#1632)
* feat: add installation support and configuration files for stable-diffusion

* fix: correct public header setting and update version variable in pkg-config

* fix stable-diffusion install package metadata

---------

Co-authored-by: leejet <leejet714@gmail.com>
2026-06-29 00:48:57 +08:00
leejet
9f855c933b
chore: silence narrowing conversion warnings (#1717) 2026-06-28 23:14:53 +08:00
stduhpf
7b5f34d93e
feat: support Qwen2D VAE (#1714) 2026-06-28 22:50:57 +08:00
stduhpf
d77b8f5ee8
feat: support Qwen-Image/Wan VAE with diffusers naming (#1713) 2026-06-28 22:50:12 +08:00
fszontagh
03e9a22f4d
feat: add SeFi-Image support (#1707) 2026-06-28 22:49:24 +08:00
leejet
f54e45e81c
fix: correct sycl ci (#1716) 2026-06-28 22:45:19 +08:00
22 changed files with 1079 additions and 88 deletions

View File

@ -331,7 +331,8 @@ endif()
add_subdirectory(thirdparty)
target_link_libraries(${SD_LIB} PUBLIC ggml zip)
target_sources(${SD_LIB} PRIVATE $<TARGET_OBJECTS:zip>)
target_link_libraries(${SD_LIB} PUBLIC ggml)
target_include_directories(${SD_LIB} PUBLIC . src include)
target_include_directories(${SD_LIB} PRIVATE src/core)
target_include_directories(${SD_LIB} PUBLIC . thirdparty)
@ -342,7 +343,58 @@ if (SD_BUILD_EXAMPLES)
add_subdirectory(examples)
endif()
set(SD_PUBLIC_HEADERS include/stable-diffusion.h)
set_target_properties(${SD_LIB} PROPERTIES PUBLIC_HEADER "${SD_PUBLIC_HEADERS}")
install(TARGETS ${SD_LIB} LIBRARY PUBLIC_HEADER)
#
# install
#
include(CMakePackageConfigHelpers)
include(GNUInstallDirs)
set(SD_INSTALL_VERSION "${SDCPP_BUILD_VERSION}")
set(SD_INSTALL_COMMIT "${SDCPP_BUILD_COMMIT}")
set(SD_SHARED_LIB ${SD_BUILD_SHARED_LIBS})
set(SD_INCLUDE_INSTALL_DIR ${CMAKE_INSTALL_INCLUDEDIR} CACHE PATH "Location of header files")
set(SD_LIB_INSTALL_DIR ${CMAKE_INSTALL_LIBDIR} CACHE PATH "Location of library files")
set(SD_BIN_INSTALL_DIR ${CMAKE_INSTALL_BINDIR} CACHE PATH "Location of binary files")
set(SD_PUBLIC_HEADERS
${CMAKE_CURRENT_SOURCE_DIR}/include/stable-diffusion.h)
set_target_properties(${SD_LIB}
PROPERTIES
PUBLIC_HEADER "${SD_PUBLIC_HEADERS}")
install(TARGETS ${SD_LIB}
ARCHIVE
LIBRARY
RUNTIME
PUBLIC_HEADER)
configure_package_config_file(
${CMAKE_CURRENT_SOURCE_DIR}/cmake/stable-diffusion-config.cmake.in
${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion-config.cmake
INSTALL_DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/stable-diffusion
PATH_VARS SD_INCLUDE_INSTALL_DIR
SD_LIB_INSTALL_DIR
SD_BIN_INSTALL_DIR )
write_basic_package_version_file(
${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion-version.cmake
VERSION ${SD_INSTALL_VERSION}
COMPATIBILITY SameMajorVersion)
install(FILES ${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion-config.cmake
${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion-version.cmake
DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/stable-diffusion)
configure_file(cmake/stable-diffusion.pc.in
"${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion.pc"
@ONLY)
install(FILES "${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion.pc"
DESTINATION ${CMAKE_INSTALL_LIBDIR}/pkgconfig)

View File

@ -21,23 +21,12 @@ WORKDIR /sd.cpp
COPY . .
RUN mkdir build && cd build && \
cmake .. -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx \
-DSD_SYCL=ON \
-DSD_BUILD_SHARED_LIBS=ON \
-DGGML_NATIVE=OFF \
-DSD_BUILD_SHARED_GGML_LIB=ON \
-DGGML_BACKEND_DL=ON \
-DGGML_CPU_ALL_VARIANTS=ON \
-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
-DCMAKE_INSTALL_RPATH='$ORIGIN' \
-DCMAKE_BUILD_TYPE=Release && \
cmake .. -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DSD_SYCL=ON -DCMAKE_BUILD_TYPE=Release && \
cmake --build . --config Release -j$(nproc)
FROM intel/oneapi-basekit:${SYCL_VERSION}-devel-ubuntu24.04 AS runtime
COPY --from=build /sd.cpp/build/bin /sd.cpp/bin
RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \
printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \
chmod +x /sd-cli /sd-server
COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli
COPY --from=build /sd.cpp/build/bin/sd-server /sd-server
ENTRYPOINT [ "/sd-cli" ]

View File

@ -53,6 +53,7 @@ API and command-line option may change frequently.***
- [ERNIE-Image](./docs/ernie_image.md)
- [Boogu Image](./docs/boogu_image.md)
- [Krea2](./docs/krea2.md)
- [SeFi-Image](./docs/sefi_image.md)
- [HiDream-O1-Image](./docs/hidream_o1_image.md)
- [Ideogram4](./docs/ideogram4.md)
- Image Edit Models

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.7 MiB

View File

@ -0,0 +1,37 @@
set(SD_VERSION "@SD_INSTALL_VERSION@")
set(SD_BUILD_COMMIT "@SD_INSTALL_COMMIT@")
set(SD_SHARED_LIB @SD_SHARED_LIB@)
@PACKAGE_INIT@
set_and_check(SD_INCLUDE_DIR "@PACKAGE_SD_INCLUDE_INSTALL_DIR@")
set_and_check(SD_LIB_DIR "@PACKAGE_SD_LIB_INSTALL_DIR@")
set(SD_BIN_DIR "@PACKAGE_SD_BIN_INSTALL_DIR@")
include(CMakeFindDependencyMacro)
find_dependency(ggml REQUIRED HINTS "${SD_LIB_DIR}/cmake")
if(NOT TARGET stable-diffusion)
find_library(stable-diffusion_LIBRARY stable-diffusion
REQUIRED
HINTS "${SD_LIB_DIR}"
NO_CMAKE_FIND_ROOT_PATH
)
add_library(stable-diffusion UNKNOWN IMPORTED)
set_target_properties(stable-diffusion
PROPERTIES
INTERFACE_INCLUDE_DIRECTORIES "${SD_INCLUDE_DIR}"
INTERFACE_LINK_LIBRARIES "ggml::ggml"
IMPORTED_LINK_INTERFACE_LANGUAGES "CXX"
IMPORTED_LOCATION "${stable-diffusion_LIBRARY}"
INTERFACE_COMPILE_FEATURES "c_std_11;cxx_std_17"
POSITION_INDEPENDENT_CODE ON)
if(SD_SHARED_LIB)
target_compile_definitions(stable-diffusion
INTERFACE SD_BUILD_SHARED_LIB)
endif()
endif()
check_required_components(stable-diffusion)

View File

@ -0,0 +1,11 @@
prefix=@CMAKE_INSTALL_PREFIX@
exec_prefix=${prefix}
libdir=@CMAKE_INSTALL_FULL_LIBDIR@
includedir=@CMAKE_INSTALL_FULL_INCLUDEDIR@
Name: stable-diffusion
Description: Diffusion model(SD,Flux,Wan,Qwen Image,Z-Image,...) inference in pure C/C++
Version: @SDCPP_BUILD_VERSION@
Libs: -L${libdir} -lstable-diffusion
Libs.private: -lggml -lggml-base
Cflags: -I${includedir}

50
docs/sefi_image.md Normal file
View File

@ -0,0 +1,50 @@
# How to Use
SeFi-Image uses a Flux2-style dual-time transformer (semantic + texture streams), the standard Flux2 VAE, and Qwen3-VL as the LLM text encoder. Tech report: [arXiv:2606.22568](https://arxiv.org/abs/2606.22568).
## Download weights
The SeFi-Image family ships in three scales (1B / 2B / 5B) and three families (Base / RL / turbo), all gated on Hugging Face under https://huggingface.co/SeFi-Image.
- 1B and 2B variants pair with Qwen3-VL-2B-Instruct.
- 5B variants pair with Qwen3-VL-4B-Instruct.
- All variants use the standard Flux2 VAE (`flux2_ae.safetensors` from https://huggingface.co/black-forest-labs/FLUX.2-dev).
Convert the transformer and text encoder to sd.cpp safetensors:
```bash
python3 script/convert_sefi.py <hf_repo_dir> <out_dir>/sefi_<scale>_<family>.safetensors
python3 script/convert_qwen3_vl.py <hf_repo_dir>/Qwen3-VL-XB-Instruct <out_dir>/qwen3_vl_<X>b.safetensors
```
## Variant defaults
| Family | timestep_shift_alpha | steps | cfg-scale |
|---|---|---|---|
| Base | 0.3 | 50 | 4.0 |
| RL | 0.3 | 50 | 4.0 |
| turbo | 1.0 | 4 | 1.0 |
The dispatcher picks `alpha` from the filename (`turbo` substring => 1.0, otherwise 0.3). Override via `--extra-sample-args sefi_alpha=<value>` or `sefi_delta_t=<value>`.
## Examples
### 1B / 2B turbo
```
./build/bin/sd-cli --diffusion-model /path/to/sefi_1b_turbo.safetensors --vae /path/to/flux2_ae.safetensors --llm /path/to/qwen3_vl_2b.safetensors -p "a photograph of an orange tabby cat sitting on a couch" --cfg-scale 1.0 --steps 4 -W 1024 -H 1024 -s 42 --diffusion-fa --offload-to-cpu -o out.png
```
### 1B / 2B base
```
./build/bin/sd-cli --diffusion-model /path/to/sefi_1b_base.safetensors --vae /path/to/flux2_ae.safetensors --llm /path/to/qwen3_vl_2b.safetensors -p "a photograph of an orange tabby cat sitting on a couch" --cfg-scale 4.0 --steps 50 -W 1024 -H 1024 -s 42 --diffusion-fa --offload-to-cpu -o out.png
```
### 5B (needs streaming on 12 GiB VRAM)
```
./build/bin/sd-cli --diffusion-model /path/to/sefi_5b_turbo.safetensors --vae /path/to/flux2_ae.safetensors --llm /path/to/qwen3_vl_4b.safetensors -p "a photograph of an orange tabby cat sitting on a couch" --cfg-scale 1.0 --steps 4 -W 1024 -H 1024 -s 42 --diffusion-fa --max-vram 8 --stream-layers --offload-to-cpu -o out.png
```
<img alt="SeFi-Image 5B turbo example" src="../assets/sefi_image/example.png" />

View File

@ -15,7 +15,7 @@ target_include_directories(${TARGET} PRIVATE
"${PROJECT_SOURCE_DIR}/src"
)
install(TARGETS ${TARGET} RUNTIME)
target_link_libraries(${TARGET} PRIVATE stable-diffusion zip ${CMAKE_THREAD_LIBS_INIT})
target_link_libraries(${TARGET} PRIVATE stable-diffusion ${CMAKE_THREAD_LIBS_INIT})
if(SD_WEBP)
target_compile_definitions(${TARGET} PRIVATE SD_USE_WEBP)
target_link_libraries(${TARGET} PRIVATE webp libwebpmux)

View File

@ -81,6 +81,7 @@ enum prediction_t {
FLOW_PRED,
FLUX_FLOW_PRED,
FLUX2_FLOW_PRED,
SEFI_FLOW_PRED,
PREDICTION_COUNT
};

112
script/convert_qwen3_vl.py Normal file
View File

@ -0,0 +1,112 @@
#!/usr/bin/env python3
"""Convert a Qwen3-VL HF safetensors checkpoint into a sd.cpp-loadable form.
The HF dump prefixes text-tower keys with ``model.language_model.`` and
vision-tower keys with ``model.visual.``. sd.cpp expects ``model.<rest>`` for
the text side; the vision side is converted by sd.cpp's own
``convert_qwen3_vl_vision_name`` and is left as-is here.
Operates on raw safetensors bytes so any dtype (BF16/F16/F32) is preserved.
Usage:
python3 script/convert_qwen3_vl.py <hf_qwen3_vl_dir_or_safetensors> <output.safetensors>
"""
import argparse
import json
import os
import struct
import sys
def rewrite_key(key: str) -> str:
if key.startswith("model.language_model."):
return "model." + key[len("model.language_model."):]
return key
def read_safetensors_header(path: str):
with open(path, "rb") as f:
hdr_len = struct.unpack("<Q", f.read(8))[0]
hdr_bytes = f.read(hdr_len)
return json.loads(hdr_bytes), 8 + hdr_len
def collect_shard_paths(path: str):
if os.path.isdir(path):
index_path = os.path.join(path, "model.safetensors.index.json")
if os.path.isfile(index_path):
with open(index_path) as f:
idx = json.load(f)
return sorted({os.path.join(path, n) for n in idx["weight_map"].values()})
single = os.path.join(path, "model.safetensors")
if os.path.isfile(single):
return [single]
raise FileNotFoundError(f"No Qwen3-VL safetensors in {path}")
if os.path.isfile(path):
return [path]
raise FileNotFoundError(path)
def stage_tensors(input_path: str):
entries = []
for shard_path in collect_shard_paths(input_path):
hdr, data_off = read_safetensors_header(shard_path)
for key, info in hdr.items():
if key == "__metadata__":
continue
entries.append((rewrite_key(key), shard_path, data_off, info))
return entries
def write_consolidated(out_path: str, entries):
entries = sorted(entries, key=lambda e: e[0])
new_header = {}
cur_offset = 0
for new_key, shard_path, data_off, info in entries:
start, end = info["data_offsets"]
size = end - start
new_header[new_key] = {
"dtype": info["dtype"],
"shape": info["shape"],
"data_offsets": [cur_offset, cur_offset + size],
}
cur_offset += size
header_json = json.dumps(new_header, separators=(",", ":")).encode("utf-8")
pad = (-len(header_json)) % 8
header_json = header_json + (b" " * pad)
with open(out_path, "wb") as out:
out.write(struct.pack("<Q", len(header_json)))
out.write(header_json)
for new_key, shard_path, data_off, info in entries:
start, end = info["data_offsets"]
with open(shard_path, "rb") as src:
src.seek(data_off + start)
remaining = end - start
while remaining > 0:
chunk = src.read(min(8 * 1024 * 1024, remaining))
if not chunk:
raise IOError(f"Truncated tensor in {shard_path}")
out.write(chunk)
remaining -= len(chunk)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("input", help="HF Qwen3-VL directory or single safetensors file")
parser.add_argument("output", help="Output single safetensors path")
args = parser.parse_args()
entries = stage_tensors(args.input)
print(f"Tensors: {len(entries)}")
print(f"Writing -> {args.output}")
os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
write_consolidated(args.output, entries)
print(f"Done. Output size: {os.path.getsize(args.output) / 1e9:.2f} GB")
if __name__ == "__main__":
main()

279
script/convert_sefi.py Normal file
View File

@ -0,0 +1,279 @@
#!/usr/bin/env python3
"""Convert a SeFi-Image diffusers checkpoint into a single sd.cpp-compatible safetensors.
Operates on raw safetensors bytes so any dtype (BF16, F32, ...) is preserved exactly.
No numpy or torch dependency required.
Usage:
python3 script/convert_sefi.py <sefi_diffusers_dir> <output.safetensors>
"""
import argparse
import json
import os
import re
import struct
import sys
_LINEAR_TO_LIN = re.compile(r"\.linear\.")
_SHARED_MOD_PREFIXES = (
"double_stream_modulation_img",
"double_stream_modulation_txt",
"single_stream_modulation",
)
def rewrite_transformer_key(key: str) -> str:
if key.startswith("backbone."):
key = key[len("backbone."):]
elif key.startswith("dual_time_embed."):
return key
if any(key.startswith(prefix + ".") for prefix in _SHARED_MOD_PREFIXES):
key = _LINEAR_TO_LIN.sub(".lin.", key, count=1)
if key == "context_embedder.weight":
return "txt_in.weight"
if key == "context_embedder.bias":
return "txt_in.bias"
if key == "x_embedder.weight":
return "img_in.weight"
if key == "x_embedder.bias":
return "img_in.bias"
if key == "proj_out.weight":
return "final_layer.linear.weight"
if key == "proj_out.bias":
return "final_layer.linear.bias"
if key == "norm_out.linear.weight":
return "final_layer.adaLN_modulation.1.weight"
if key == "norm_out.linear.bias":
return "final_layer.adaLN_modulation.1.bias"
m = re.match(r"transformer_blocks\.(\d+)\.(.*)$", key)
if m:
return _rewrite_double_stream(m.group(1), m.group(2))
m = re.match(r"single_transformer_blocks\.(\d+)\.(.*)$", key)
if m:
return _rewrite_single_stream(m.group(1), m.group(2))
return key
def _rewrite_double_stream(idx: str, tail: str) -> str:
dst = f"double_blocks.{idx}."
mapping = {
"norm1.linear.weight": "img_mod.lin.weight",
"norm1_context.linear.weight": "txt_mod.lin.weight",
"attn.norm_q.weight": "img_attn.norm.query_norm.scale",
"attn.norm_k.weight": "img_attn.norm.key_norm.scale",
"attn.norm_added_q.weight": "txt_attn.norm.query_norm.scale",
"attn.norm_added_k.weight": "txt_attn.norm.key_norm.scale",
"attn.to_out.0.weight": "img_attn.proj.weight",
"attn.to_add_out.weight": "txt_attn.proj.weight",
"ff.net.0.proj.weight": "img_mlp.0.weight",
"ff.net.2.weight": "img_mlp.2.weight",
"ff_context.net.0.proj.weight": "txt_mlp.0.weight",
"ff_context.net.2.weight": "txt_mlp.2.weight",
"ff.linear_in.weight": "img_mlp.0.weight",
"ff.linear_out.weight": "img_mlp.2.weight",
"ff_context.linear_in.weight": "txt_mlp.0.weight",
"ff_context.linear_out.weight": "txt_mlp.2.weight",
}
return dst + mapping.get(tail, tail)
# QKV triplets to fuse on output: source tails -> target fused tail.
# Each tuple is (q_tail, k_tail, v_tail, fused_target_tail).
QKV_DOUBLE_TRIPLETS = [
("attn.to_q.weight", "attn.to_k.weight", "attn.to_v.weight", "img_attn.qkv.weight"),
("attn.add_q_proj.weight", "attn.add_k_proj.weight", "attn.add_v_proj.weight", "txt_attn.qkv.weight"),
]
def _rewrite_single_stream(idx: str, tail: str) -> str:
dst = f"single_blocks.{idx}."
mapping = {
"norm.linear.weight": "modulation.lin.weight",
"attn.norm_q.weight": "norm.query_norm.scale",
"attn.norm_k.weight": "norm.key_norm.scale",
"attn.to_qkv_mlp_proj.weight": "linear1.weight",
"attn.to_out.weight": "linear2.weight",
}
return dst + mapping.get(tail, tail)
def read_safetensors_header(path: str):
"""Return (header dict, data start byte offset)."""
with open(path, "rb") as f:
hdr_len = struct.unpack("<Q", f.read(8))[0]
hdr_bytes = f.read(hdr_len)
return json.loads(hdr_bytes), 8 + hdr_len
def collect_shard_paths(directory: str, weight_pattern: str):
index_path = os.path.join(directory, f"{weight_pattern}.safetensors.index.json")
if os.path.isfile(index_path):
with open(index_path) as f:
idx = json.load(f)
return sorted({os.path.join(directory, n) for n in idx["weight_map"].values()})
single = os.path.join(directory, f"{weight_pattern}.safetensors")
if not os.path.isfile(single):
raise FileNotFoundError(f"No checkpoint at {directory}: missing {weight_pattern}")
return [single]
def stage_tensors_for_section(section_dir: str, rewrite_fn):
"""Return a list of (new_key, shard_path, data_start_offset, info_dict) entries.
A "qkv_fuse" pseudo-entry with three source descriptors is emitted when a
transformer_blocks.* split q/k/v triplet is found, so the writer can fuse
them into a single output tensor.
"""
entries = []
# First, index all raw keys per shard so we can detect qkv triplets.
raw_by_block = {} # block_idx -> {tail: (key, shard_path, data_off, info)}
raw_others = []
for shard_path in collect_shard_paths(section_dir, "diffusion_pytorch_model"):
hdr, data_off = read_safetensors_header(shard_path)
for key, info in hdr.items():
if key == "__metadata__":
continue
m = re.match(r"backbone\.transformer_blocks\.(\d+)\.(.*)$", key)
if m and any(m.group(2) in trip[:3] for trip in QKV_DOUBLE_TRIPLETS):
idx = m.group(1)
raw_by_block.setdefault(idx, {})[m.group(2)] = (key, shard_path, data_off, info)
else:
raw_others.append((key, shard_path, data_off, info))
for key, shard_path, data_off, info in raw_others:
new_key = rewrite_fn(key)
# Swap the (scale, shift) halves to (shift, scale) at conversion time so
# the on-disk weight matches BFL flux ordering and the runtime stays
# version-agnostic. norm_out.linear weight shape is [2*dim, dim] and bias
# is [2*dim]; both split along axis 0 (outermost == row-major outer).
if new_key in ("final_layer.adaLN_modulation.1.weight",
"final_layer.adaLN_modulation.1.bias"):
info = dict(info)
info["_chunk_swap_halves"] = True
entries.append((new_key, shard_path, data_off, info))
for block_idx, tails in raw_by_block.items():
for q_tail, k_tail, v_tail, fused_tail in QKV_DOUBLE_TRIPLETS:
if q_tail in tails and k_tail in tails and v_tail in tails:
q = tails[q_tail]; k = tails[k_tail]; v = tails[v_tail]
# Validate shapes match.
q_shape = q[3]["shape"]; k_shape = k[3]["shape"]; v_shape = v[3]["shape"]
if q_shape != k_shape or q_shape != v_shape:
raise ValueError(f"qkv shape mismatch at block {block_idx} {q_tail}: q={q_shape} k={k_shape} v={v_shape}")
fused_shape = [q_shape[0] * 3] + list(q_shape[1:])
fused_info = {
"dtype": q[3]["dtype"],
"shape": fused_shape,
"_qkv_sources": [q, k, v], # pseudo field consumed by writer
}
entries.append((f"double_blocks.{block_idx}.{fused_tail}",
None, None, fused_info))
del tails[q_tail]; del tails[k_tail]; del tails[v_tail]
# Anything left in tails was an unmatched single - pass through.
for tail, payload in tails.items():
entries.append((rewrite_fn(payload[0]),) + payload[1:])
return entries
_DTYPE_BYTES = {
"BF16": 2, "F16": 2, "F32": 4, "F64": 8,
"U8": 1, "I8": 1, "I16": 2, "I32": 4, "I64": 8,
"BOOL": 1,
}
def _total_bytes(info: dict) -> int:
if "_qkv_sources" in info:
elems = 1
for d in info["shape"]:
elems *= d
return elems * _DTYPE_BYTES[info["dtype"]]
start, end = info["data_offsets"]
return end - start
def write_consolidated(out_path: str, entries):
"""Write a single safetensors file by streaming raw bytes from each shard.
For qkv-fused entries, q/k/v are concatenated along axis 0 (row-major), so a
simple byte-level concatenation produces the correct fused layout for any
standard dtype.
"""
entries = sorted(entries, key=lambda e: e[0])
new_header = {}
cur_offset = 0
for new_key, shard_path, data_off, info in entries:
size = _total_bytes(info)
new_header[new_key] = {
"dtype": info["dtype"],
"shape": info["shape"],
"data_offsets": [cur_offset, cur_offset + size],
}
cur_offset += size
header_json = json.dumps(new_header, separators=(",", ":")).encode("utf-8")
pad = (-len(header_json)) % 8
header_json = header_json + (b" " * pad)
def copy_range(src_path, src_data_off, src_info, out, byte_range=None):
start, end = src_info["data_offsets"]
if byte_range is not None:
sub_start, sub_end = byte_range
start, end = start + sub_start, start + sub_end
with open(src_path, "rb") as src:
src.seek(src_data_off + start)
remaining = end - start
while remaining > 0:
chunk = src.read(min(8 * 1024 * 1024, remaining))
if not chunk:
raise IOError(f"Truncated tensor in {src_path}")
out.write(chunk)
remaining -= len(chunk)
with open(out_path, "wb") as out:
out.write(struct.pack("<Q", len(header_json)))
out.write(header_json)
for new_key, shard_path, data_off, info in entries:
if "_qkv_sources" in info:
for q_entry in info["_qkv_sources"]:
_, src_path, src_data_off, src_info = q_entry
copy_range(src_path, src_data_off, src_info, out)
elif info.get("_chunk_swap_halves"):
size = _total_bytes(info)
half = size // 2
if size != half * 2:
raise ValueError(f"{new_key}: odd byte size {size} cannot be split into halves")
copy_range(shard_path, data_off, info, out, byte_range=(half, size))
copy_range(shard_path, data_off, info, out, byte_range=(0, half))
else:
copy_range(shard_path, data_off, info, out)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("input_dir", help="SeFi diffusers checkpoint directory")
parser.add_argument("output", help="Output transformer safetensors path (load via --diffusion-model)")
args = parser.parse_args()
transformer_entries = stage_tensors_for_section(
os.path.join(args.input_dir, "transformer"), rewrite_transformer_key)
print(f"Transformer tensors: {len(transformer_entries)}")
print(f"Writing {len(transformer_entries)} tensors -> {args.output}")
os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
write_consolidated(args.output, transformer_entries)
print(f"Done. Output size: {os.path.getsize(args.output) / 1e9:.2f} GB")
if __name__ == "__main__":
main()

View File

@ -1518,7 +1518,7 @@ struct LLMEmbedder : public Conditioner {
arch = LLM::LLMArch::GPT_OSS_20B;
} else if (sd_version_is_pid(version)) {
arch = LLM::LLMArch::GEMMA2_2B;
} else if (sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_krea2(version)) {
} else if (sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) {
arch = LLM::LLMArch::QWEN3_VL;
} else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
arch = LLM::LLMArch::QWEN3;
@ -1997,6 +1997,18 @@ struct LLMEmbedder : public Conditioner {
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
} else if (sd_version_is_sefi_image(version)) {
prompt_template_encode_start_idx = 0;
min_length = 1024;
out_layers = {9, 18, 27};
prompt = "<|im_start|>user\n";
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n<|im_start|>assistant\n";
} else if (version == VERSION_OVIS_IMAGE) {
prompt_template_encode_start_idx = 28;
min_length = prompt_template_encode_start_idx + 256;

View File

@ -49,6 +49,7 @@ enum SDVersion {
VERSION_LONGCAT,
VERSION_PID,
VERSION_IDEOGRAM4,
VERSION_SEFI_IMAGE,
VERSION_KREA2,
VERSION_ESRGAN,
VERSION_COUNT,
@ -187,6 +188,13 @@ static inline bool sd_version_is_ideogram4(SDVersion version) {
return false;
}
static inline bool sd_version_is_sefi_image(SDVersion version) {
if (version == VERSION_SEFI_IMAGE) {
return true;
}
return false;
}
static inline bool sd_version_is_krea2(SDVersion version) {
if (version == VERSION_KREA2) {
return true;
@ -202,7 +210,7 @@ static inline bool sd_version_uses_flux_vae(SDVersion version) {
}
static inline bool sd_version_uses_flux2_vae(SDVersion version) {
if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version)) {
if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version)) {
return true;
}
return false;
@ -242,6 +250,7 @@ static inline bool sd_version_is_dit(SDVersion version) {
sd_version_is_longcat(version) ||
sd_version_is_pid(version) ||
sd_version_is_ideogram4(version) ||
sd_version_is_sefi_image(version) ||
sd_version_is_krea2(version)) {
return true;
}

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, (int)timestep_guidance_in_dim, 10000, 1.f);
auto tex_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_tex, (int)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__

View File

@ -250,7 +250,7 @@ namespace LLM {
config.intermediate_size = tensor_storage.ne[1];
}
}
if (arch == LLMArch::QWEN3 && config.num_layers == 28) {
if ((arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) && config.num_layers == 28) {
config.num_heads = 16;
}
if (detected_vision_layers > 0) {

View File

@ -816,12 +816,13 @@ struct AutoEncoderKL : public VAE {
}
sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
auto latents_ = sd_version_is_sefi_image(version) ? sd::ops::slice(latents, 2, 16, 144) : latents;
if (sd_version_uses_flux2_vae(version)) {
int channel_dim = 2;
auto [mean_tensor, std_tensor] = get_latents_mean_std(latents, channel_dim);
return (latents * std_tensor) / scale_factor + mean_tensor;
auto [mean_tensor, std_tensor] = get_latents_mean_std(latents_, channel_dim);
return (latents_ * std_tensor) / scale_factor + mean_tensor;
}
return (latents / scale_factor) + shift_factor;
return (latents_ / scale_factor) + shift_factor;
}
sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {

View File

@ -113,6 +113,24 @@ namespace WAN {
}
};
class Conv2dBut3d : public Conv2d {
public:
using Conv2d::Conv2d;
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
ggml_tensor* x_swapped = ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2);
x_swapped = ggml_cont(ctx->ggml_ctx, x_swapped);
ggml_tensor* out = Conv2d::forward(ctx, x_swapped);
ggml_tensor* out_swapped = ggml_permute(ctx->ggml_ctx, out, 0, 1, 3, 2);
out_swapped = ggml_cont(ctx->ggml_ctx, out_swapped);
return out_swapped;
}
};
class Resample : public GGMLBlock {
protected:
int64_t dim;
@ -338,21 +356,34 @@ namespace WAN {
protected:
int64_t in_dim;
int64_t out_dim;
bool is_2D;
public:
ResidualBlock(int64_t in_dim, int64_t out_dim)
: in_dim(in_dim), out_dim(out_dim) {
ResidualBlock(int64_t in_dim, int64_t out_dim, bool is_2D = false)
: in_dim(in_dim), out_dim(out_dim), is_2D(is_2D) {
blocks["residual.0"] = std::shared_ptr<GGMLBlock>(new RMS_norm(in_dim));
// residual.1 is nn.SiLU()
if (is_2D) {
blocks["residual.2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(in_dim, out_dim, {3, 3}, {1, 1}, {1, 1}));
} else {
blocks["residual.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(in_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
}
blocks["residual.3"] = std::shared_ptr<GGMLBlock>(new RMS_norm(out_dim));
// residual.4 is nn.SiLU()
// residual.5 is nn.Dropout()
if (is_2D) {
blocks["residual.6"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(out_dim, out_dim, {3, 3}, {1, 1}, {1, 1}));
} else {
blocks["residual.6"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
}
if (in_dim != out_dim) {
if (is_2D) {
blocks["shortcut"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(in_dim, out_dim, {1, 1}));
} else {
blocks["shortcut"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(in_dim, out_dim, {1, 1, 1}));
}
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
@ -363,10 +394,16 @@ namespace WAN {
GGML_ASSERT(b == 1);
ggml_tensor* h = x;
if (in_dim != out_dim) {
if (is_2D) {
auto shortcut = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["shortcut"]);
h = shortcut->forward(ctx, x);
} else {
auto shortcut = std::dynamic_pointer_cast<CausalConv3d>(blocks["shortcut"]);
h = shortcut->forward(ctx, x);
}
}
for (int i = 0; i < 7; i++) {
if (i == 0 || i == 3) { // RMS_norm
@ -385,8 +422,13 @@ namespace WAN {
cache_x,
2);
}
if (is_2D) {
auto layer = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["residual." + std::to_string(i)]);
x = layer->forward(ctx, x);
} else {
x = layer->forward(ctx, x, feat_cache[idx]);
}
feat_cache[idx] = cache_x;
feat_idx += 1;
}
@ -412,13 +454,14 @@ namespace WAN {
int64_t out_dim,
int mult,
bool temperal_downsample = false,
bool down_flag = false)
bool down_flag = false,
bool is_2D = false)
: mult(mult), down_flag(down_flag) {
blocks["avg_shortcut"] = std::shared_ptr<GGMLBlock>(new AvgDown3D(in_dim, out_dim, temperal_downsample ? 2 : 1, down_flag ? 2 : 1));
int i = 0;
for (; i < mult; i++) {
blocks["downsamples." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim));
blocks["downsamples." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim, is_2D));
in_dim = out_dim;
}
if (down_flag) {
@ -472,7 +515,8 @@ namespace WAN {
int64_t out_dim,
int mult,
bool temperal_upsample = false,
bool up_flag = false)
bool up_flag = false,
bool is_2D = false)
: mult(mult), up_flag(up_flag) {
if (up_flag) {
blocks["avg_shortcut"] = std::shared_ptr<GGMLBlock>(new DupUp3D(in_dim, out_dim, temperal_upsample ? 2 : 1, up_flag ? 2 : 1));
@ -480,7 +524,7 @@ namespace WAN {
int i = 0;
for (; i < mult; i++) {
blocks["upsamples." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim));
blocks["upsamples." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim, is_2D));
in_dim = out_dim;
}
if (up_flag) {
@ -592,6 +636,7 @@ namespace WAN {
std::vector<int> dim_mult;
int num_res_blocks;
std::vector<bool> temperal_downsample;
bool is_2D = false;
public:
Encoder3d(int64_t dim = 128,
@ -599,23 +644,26 @@ namespace WAN {
std::vector<int> dim_mult = {1, 2, 4, 4},
int num_res_blocks = 2,
std::vector<bool> temperal_downsample = {false, true, true},
bool wan2_2 = false)
bool wan2_2 = false,
bool is_2D = false)
: dim(dim),
z_dim(z_dim),
dim_mult(dim_mult),
num_res_blocks(num_res_blocks),
temperal_downsample(temperal_downsample),
wan2_2(wan2_2) {
wan2_2(wan2_2),
is_2D(is_2D) {
// attn_scales is always []
std::vector<int64_t> dims = {dim};
for (int u : dim_mult) {
dims.push_back(dim * u);
}
if (wan2_2) {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(12, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
int64_t input_dim = wan2_2 ? 12 : 3;
if (is_2D) {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(input_dim, dims[0], {3, 3}, {1, 1}, {1, 1}));
} else {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(3, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(input_dim, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
}
int index = 0;
@ -630,12 +678,13 @@ namespace WAN {
out_dim,
num_res_blocks,
t_down_flag,
i != dim_mult.size() - 1));
i != dim_mult.size() - 1,
is_2D));
blocks["downsamples." + std::to_string(index++)] = block;
} else {
for (int j = 0; j < num_res_blocks; j++) {
auto block = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim));
auto block = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim, is_2D));
blocks["downsamples." + std::to_string(index++)] = block;
in_dim = out_dim;
}
@ -648,14 +697,18 @@ namespace WAN {
}
}
blocks["middle.0"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(out_dim, out_dim));
blocks["middle.0"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(out_dim, out_dim, is_2D));
blocks["middle.1"] = std::shared_ptr<GGMLBlock>(new AttentionBlock(out_dim));
blocks["middle.2"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(out_dim, out_dim));
blocks["middle.2"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(out_dim, out_dim, is_2D));
blocks["head.0"] = std::shared_ptr<GGMLBlock>(new RMS_norm(out_dim));
// head.1 is nn.SiLU()
if (is_2D) {
blocks["head.2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(out_dim, z_dim, {3, 3}, {1, 1}, {1, 1}));
} else {
blocks["head.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, z_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
@ -673,7 +726,10 @@ namespace WAN {
auto head_2 = std::dynamic_pointer_cast<CausalConv3d>(blocks["head.2"]);
// conv1
if (feat_cache.size() > 0) {
if (is_2D) {
auto conv1 = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["conv1"]);
x = conv1->forward(ctx, x);
} else if (feat_cache.size() > 0) {
int idx = feat_idx;
auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]);
if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) {
@ -728,7 +784,10 @@ namespace WAN {
// head
x = head_0->forward(ctx, x);
x = ggml_silu(ctx->ggml_ctx, x);
if (feat_cache.size() > 0) {
if (is_2D) {
auto head_2 = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["head.2"]);
x = head_2->forward(ctx, x);
} else if (feat_cache.size() > 0) {
int idx = feat_idx;
auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]);
if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) {
@ -758,6 +817,7 @@ namespace WAN {
std::vector<int> dim_mult;
int num_res_blocks;
std::vector<bool> temperal_upsample;
bool is_2D = false;
public:
Decoder3d(int64_t dim = 128,
@ -765,13 +825,15 @@ namespace WAN {
std::vector<int> dim_mult = {1, 2, 4, 4},
int num_res_blocks = 2,
std::vector<bool> temperal_upsample = {true, true, false},
bool wan2_2 = false)
bool wan2_2 = false,
bool is_2D = false)
: dim(dim),
z_dim(z_dim),
dim_mult(dim_mult),
num_res_blocks(num_res_blocks),
temperal_upsample(temperal_upsample),
wan2_2(wan2_2) {
wan2_2(wan2_2),
is_2D(is_2D) {
// attn_scales is always []
std::vector<int64_t> dims = {dim_mult[dim_mult.size() - 1] * dim};
for (int i = static_cast<int>(dim_mult.size()) - 1; i >= 0; i--) {
@ -779,12 +841,16 @@ namespace WAN {
}
// init block
if (is_2D) {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(z_dim, dims[0], {3, 3}, {1, 1}, {1, 1}));
} else {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
}
// middle blocks
blocks["middle.0"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(dims[0], dims[0]));
blocks["middle.0"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(dims[0], dims[0], is_2D));
blocks["middle.1"] = std::shared_ptr<GGMLBlock>(new AttentionBlock(dims[0]));
blocks["middle.2"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(dims[0], dims[0]));
blocks["middle.2"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(dims[0], dims[0], is_2D));
// upsample blocks
int index = 0;
@ -799,7 +865,8 @@ namespace WAN {
out_dim,
num_res_blocks + 1,
t_up_flag,
i != dim_mult.size() - 1));
i != dim_mult.size() - 1,
is_2D));
blocks["upsamples." + std::to_string(index++)] = block;
} else {
@ -807,7 +874,7 @@ namespace WAN {
in_dim = in_dim / 2;
}
for (int j = 0; j < num_res_blocks + 1; j++) {
auto block = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim));
auto block = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim, is_2D));
blocks["upsamples." + std::to_string(index++)] = block;
in_dim = out_dim;
}
@ -822,12 +889,13 @@ namespace WAN {
// output blocks
blocks["head.0"] = std::shared_ptr<GGMLBlock>(new RMS_norm(out_dim));
int64_t final_dim = wan2_2 ? 12 : 3;
// head.1 is nn.SiLU()
if (wan2_2) {
blocks["head.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, 12, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
if (is_2D) {
blocks["head.2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(out_dim, final_dim, {3, 3}, {1, 1}, {1, 1}));
} else {
blocks["head.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, 3, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
blocks["head.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, final_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
}
}
@ -847,7 +915,10 @@ namespace WAN {
auto head_2 = std::dynamic_pointer_cast<CausalConv3d>(blocks["head.2"]);
// conv1
if (feat_cache.size() > 0) {
if (is_2D) {
auto conv1 = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["conv1"]);
x = conv1->forward(ctx, x);
} else if (feat_cache.size() > 0) {
int idx = feat_idx;
auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]);
if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) {
@ -902,7 +973,10 @@ namespace WAN {
// head
x = head_0->forward(ctx, x);
x = ggml_silu(ctx->ggml_ctx, x);
if (feat_cache.size() > 0) {
if (is_2D) {
auto head_2 = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["head.2"]);
x = head_2->forward(ctx, x);
} else if (feat_cache.size() > 0) {
int idx = feat_idx;
auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]);
if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) {
@ -935,6 +1009,7 @@ namespace WAN {
int num_res_blocks = 2;
std::vector<bool> temperal_upsample = {true, true, false};
std::vector<bool> temperal_downsample = {false, true, true};
bool is_2D = false;
int _conv_num = 33;
int _conv_idx = 0;
@ -951,8 +1026,8 @@ namespace WAN {
}
public:
WanVAE(bool decode_only = true, bool wan2_2 = false)
: decode_only(decode_only), wan2_2(wan2_2) {
WanVAE(bool decode_only = true, bool wan2_2 = false, bool is_2D = false)
: decode_only(decode_only), wan2_2(wan2_2), is_2D(is_2D) {
// attn_scales is always []
if (wan2_2) {
dim = 160;
@ -962,13 +1037,27 @@ namespace WAN {
_conv_num = 34;
_enc_conv_num = 26;
}
if (is_2D) {
temperal_upsample = {false, false, false};
temperal_downsample = {false, false, false};
}
if (!decode_only) {
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, temperal_downsample, wan2_2));
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, temperal_downsample, wan2_2, is_2D));
if (is_2D) {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(z_dim * 2, z_dim * 2, {1, 1}));
} else {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim * 2, z_dim * 2, {1, 1, 1}));
}
blocks["decoder"] = std::shared_ptr<GGMLBlock>(new Decoder3d(dec_dim, z_dim, dim_mult, num_res_blocks, temperal_upsample, wan2_2));
}
blocks["decoder"] = std::shared_ptr<GGMLBlock>(new Decoder3d(dec_dim, z_dim, dim_mult, num_res_blocks, temperal_upsample, wan2_2, is_2D));
if (is_2D) {
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(z_dim, z_dim, {1, 1}));
} else {
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, z_dim, {1, 1, 1}));
}
}
static ggml_tensor* patchify(ggml_context* ctx,
ggml_tensor* x,
@ -1054,7 +1143,12 @@ namespace WAN {
out = ggml_concat(ctx->ggml_ctx, out, out_, 2);
}
}
if (is_2D) {
auto conv1 = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["conv1"]);
out = conv1->forward(ctx, out);
} else {
out = conv1->forward(ctx, out);
}
auto mu = ggml_ext_chunk(ctx->ggml_ctx, out, 2, 3)[0];
// sd::ggml_graph_cut::mark_graph_cut(mu, "wan_vae.encode.final", "mu");
clear_cache();
@ -1073,7 +1167,13 @@ namespace WAN {
auto conv2 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv2"]);
int64_t iter_ = z->ne[2];
auto x = conv2->forward(ctx, z);
auto x = z;
if (is_2D) {
auto conv2 = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["conv2"]);
x = conv2->forward(ctx, z);
} else {
x = conv2->forward(ctx, z);
}
// sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decode.prelude", "x");
ggml_tensor* out;
for (int i = 0; i < iter_; i++) {
@ -1129,7 +1229,20 @@ namespace WAN {
bool decode_only = false,
SDVersion version = VERSION_WAN2,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: VAE(version, backend, prefix, weight_manager), decode_only(decode_only), ae(decode_only, version == VERSION_WAN2_2_TI2V) {
: VAE(version, backend, prefix, weight_manager), decode_only(decode_only) {
bool is_2D = false;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (ends_with(name, "decoder.conv1.weight")) {
if (tensor_storage.ne[2] > 3) {
is_2D = true;
}
break;
}
}
if (is_2D) {
LOG_DEBUG("USING 2D VAE");
}
ae = WanVAE(decode_only, version == VERSION_WAN2_2_TI2V, is_2D);
ae.init(params_ctx, tensor_storage_map, prefix);
}

View File

@ -66,7 +66,6 @@ const char* unused_tensors[] = {
// "v_pred", // Used to detect SDXL vpred models
"text_encoders.llm.output.weight",
"text_encoders.llm.lm_head.",
"first_stage_model.bn.",
};
bool is_unused_tensor(const std::string& name) {
@ -480,6 +479,9 @@ SDVersion ModelLoader::get_sd_version() {
if (tensor_storage.name.find("model.diffusion_model.double_stream_modulation_img.lin.weight") != std::string::npos) {
is_flux2 = true;
}
if (tensor_storage.name.find("dual_time_embed.semantic_embedder.linear_1.weight") != std::string::npos) {
return VERSION_SEFI_IMAGE;
}
if (tensor_storage.name.find("single_blocks.47.linear1.weight") != std::string::npos) {
has_single_block_47 = true;
}

View File

@ -743,7 +743,7 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
name = convert_diffusers_unet_to_original_sdxl(name);
} else if (sd_version_is_sd3(version)) {
name = convert_diffusers_dit_to_original_sd3(name);
} else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version)) {
} else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) {
name = convert_diffusers_dit_to_original_flux(name);
} else if (sd_version_is_z_image(version)) {
name = convert_diffusers_dit_to_original_lumina2(name);
@ -850,7 +850,77 @@ std::string convert_diffusers_vae_to_original_sd1(std::string name) {
return result;
}
std::string convert_first_stage_model_name(std::string name, std::string prefix) {
std::string convert_diffusers_to_original_wan_vae(std::string name) {
static const std::vector<std::pair<std::string, std::string>> prefix_map = {
{"quant_conv.", "conv1."},
{"post_quant_conv.", "conv2."},
{"decoder.up_blocks.0.resnets.0.", "decoder.upsamples.0.residual."},
{"decoder.up_blocks.0.resnets.1.", "decoder.upsamples.1.residual."},
{"decoder.up_blocks.0.resnets.2.", "decoder.upsamples.2.residual."},
{"decoder.up_blocks.0.upsamplers.0.", "decoder.upsamples.3."},
{"decoder.up_blocks.1.resnets.0.conv_shortcut.", "decoder.upsamples.4.shortcut."},
{"decoder.up_blocks.1.resnets.0.", "decoder.upsamples.4.residual."},
{"decoder.up_blocks.1.resnets.1.", "decoder.upsamples.5.residual."},
{"decoder.up_blocks.1.resnets.2.", "decoder.upsamples.6.residual."},
{"decoder.up_blocks.1.upsamplers.0.", "decoder.upsamples.7."},
{"decoder.up_blocks.2.resnets.0.", "decoder.upsamples.8.residual."},
{"decoder.up_blocks.2.resnets.1.", "decoder.upsamples.9.residual."},
{"decoder.up_blocks.2.resnets.2.", "decoder.upsamples.10.residual."},
{"decoder.up_blocks.2.upsamplers.0.", "decoder.upsamples.11."},
{"decoder.up_blocks.3.resnets.0.", "decoder.upsamples.12.residual."},
{"decoder.up_blocks.3.resnets.1.", "decoder.upsamples.13.residual."},
{"decoder.up_blocks.3.resnets.2.", "decoder.upsamples.14.residual."},
{"encoder.down_blocks.0.", "encoder.downsamples.0.residual."},
{"encoder.down_blocks.1.", "encoder.downsamples.1.residual."},
{"encoder.down_blocks.2.", "encoder.downsamples.2."},
{"encoder.down_blocks.3.conv_shortcut.", "encoder.downsamples.3.shortcut."},
{"encoder.down_blocks.3.", "encoder.downsamples.3.residual."},
{"encoder.down_blocks.4.", "encoder.downsamples.4.residual."},
{"encoder.down_blocks.5.", "encoder.downsamples.5."},
{"encoder.down_blocks.6.conv_shortcut.", "encoder.downsamples.6.shortcut."},
{"encoder.down_blocks.6.", "encoder.downsamples.6.residual."},
{"encoder.down_blocks.7.", "encoder.downsamples.7.residual."},
{"encoder.down_blocks.8.", "encoder.downsamples.8."},
{"encoder.down_blocks.9.", "encoder.downsamples.9.residual."},
{"encoder.down_blocks.10.", "encoder.downsamples.10.residual."},
};
static const std::vector<std::pair<std::string, std::string>> shared_name_map = {
{".conv_in.", ".conv1."},
{".norm_out.", ".head.0."},
{".conv_out.", ".head.2."},
{".mid_block.attentions.0.", ".middle.1."},
{".mid_block.resnets.0.", ".middle.0.residual."},
{".mid_block.resnets.1.", ".middle.2.residual."},
};
static const std::vector<std::pair<std::string, std::string>> resnet_name_map = {
{".norm1.", ".0."},
{".conv1.", ".2."},
{".norm2.", ".3."},
{".conv2.", ".6."},
};
replace_with_name_map(name, shared_name_map);
replace_with_prefix_map(name, prefix_map);
// Only apply the ResNet-specific renaming if the tensor belongs to a ResNet block.
// This prevents generic ".conv1." or ".conv2." matching on top-level encoder/decoder convolutions.
if (name.find(".residual.") != std::string::npos) {
replace_with_name_map(name, resnet_name_map);
}
return name;
}
std::string convert_first_stage_model_name(std::string name, std::string prefix, SDVersion version) {
if (sd_version_uses_wan_vae(version)) {
return convert_diffusers_to_original_wan_vae(name);
}
static std::unordered_map<std::string, std::string> vae_name_map = {
{"decoder.post_quant_conv.", "post_quant_conv."},
{"encoder.quant_conv.", "quant_conv."},
@ -1239,7 +1309,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
{
for (const auto& prefix : first_stage_model_prefix_vec) {
if (starts_with(name, prefix)) {
name = convert_first_stage_model_name(name.substr(prefix.size()), prefix);
name = convert_first_stage_model_name(name.substr(prefix.size()), prefix, version);
if (version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) {
name = "tae." + name;
} else {

View File

@ -602,7 +602,7 @@ struct LogitNormalScheduler : SigmaScheduler {
}
}
if (image_seq_len > 0 && resolution_aware) {
mean += 0.5 * std::log(static_cast<float>(image_seq_len) / static_cast<float>(known_seq_len));
mean += 0.5f * std::log(static_cast<float>(image_seq_len) / static_cast<float>(known_seq_len));
}
}
@ -735,7 +735,7 @@ struct LogitNormalScheduler : SigmaScheduler {
float t = static_cast<float>(i) / static_cast<float>(n);
// ndtri(1-t) == -ndtri(t)
float z = -ndtri(t);
float z = static_cast<float>(-ndtri(t));
float y = mean + std * z;
@ -1005,6 +1005,8 @@ struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
}
};
struct SefiFlowDenoiser;
struct Flux2FlowDenoiser : public FluxFlowDenoiser {
Flux2FlowDenoiser() = default;
@ -1037,6 +1039,80 @@ struct Flux2FlowDenoiser : public FluxFlowDenoiser {
}
};
struct SefiFlowDenoiser : public Flux2FlowDenoiser {
static constexpr int kNumTrainTimesteps = 1000;
static constexpr int kSemChannels = 16;
static constexpr int kTotalChannels = 144;
float delta_t = 0.1f;
float timestep_shift_alpha = 1.0f;
std::vector<float> sem_sigmas;
std::vector<float> tex_sigmas;
std::vector<float> sem_timesteps;
std::vector<float> tex_timesteps;
SefiFlowDenoiser() = default;
static float apply_alpha_shift(float u_unit, float alpha) {
if (alpha == 1.0f) {
return u_unit;
}
float denom = 1.0f + (alpha - 1.0f) * u_unit;
return (alpha * u_unit) / denom;
}
std::vector<float> get_sigmas(uint32_t n,
int image_seq_len,
scheduler_t scheduler_type,
SDVersion version,
const char* extra_sample_args = nullptr) override {
sem_sigmas.clear();
tex_sigmas.clear();
sem_timesteps.clear();
tex_timesteps.clear();
for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "sefi scheduler arg")) {
if (key == "sefi_alpha") {
if (!parse_strict_float(value, timestep_shift_alpha)) {
LOG_WARN("ignoring invalid sefi scheduler arg '%s=%s'", key.c_str(), value.c_str());
}
} else if (key == "sefi_delta_t") {
if (!parse_strict_float(value, delta_t)) {
LOG_WARN("ignoring invalid sefi scheduler arg '%s=%s'", key.c_str(), value.c_str());
}
}
}
for (uint32_t i = 0; i <= n; ++i) {
float u_base = static_cast<float>(i) / static_cast<float>(n);
float u_shifted = apply_alpha_shift(u_base, timestep_shift_alpha);
float u_sem_raw = u_shifted * (1.0f + delta_t);
float u_sem = std::min(u_sem_raw, 1.0f);
float u_tex = std::max(0.0f, std::min(u_sem_raw - delta_t, 1.0f));
int idx_sem = std::min(kNumTrainTimesteps - 1,
std::max(0, static_cast<int>(u_sem * (kNumTrainTimesteps - 1))));
int idx_tex = std::min(kNumTrainTimesteps - 1,
std::max(0, static_cast<int>(u_tex * (kNumTrainTimesteps - 1))));
float t_sem = static_cast<float>(kNumTrainTimesteps - idx_sem);
float t_tex = static_cast<float>(kNumTrainTimesteps - idx_tex);
float sigma_sem = t_sem / static_cast<float>(kNumTrainTimesteps);
float sigma_tex = t_tex / static_cast<float>(kNumTrainTimesteps);
sem_timesteps.push_back(t_sem);
tex_timesteps.push_back(t_tex);
sem_sigmas.push_back(sigma_sem);
tex_sigmas.push_back(sigma_tex);
}
LOG_DEBUG("SefiFlowDenoiser: built %u-step dual schedule (alpha=%.2f delta_t=%.2f)",
n, timestep_shift_alpha, delta_t);
return tex_sigmas;
}
};
typedef std::function<sd::guidance::GuiderOutput(const sd::Tensor<float>&, float, int)> denoise_cb_t;
static std::pair<float, float> get_ancestral_step(float sigma_from,
@ -1140,6 +1216,40 @@ static sd::Tensor<float> sample_euler_ancestral(denoise_cb_t model,
return x;
}
static sd::Tensor<float> sample_sefi_euler(SefiFlowDenoiser* sefi,
denoise_cb_t model,
sd::Tensor<float> x) {
const std::vector<float>& sigma_tex_vec = sefi->tex_sigmas;
const std::vector<float>& sigma_sem_vec = sefi->sem_sigmas;
int steps = static_cast<int>(sigma_tex_vec.size()) - 1;
for (int i = 0; i < steps; i++) {
float sigma_tex_cur = sigma_tex_vec[i];
float sigma_tex_next = sigma_tex_vec[i + 1];
float sigma_sem_cur = sigma_sem_vec[i];
float sigma_sem_next = sigma_sem_vec[i + 1];
if (sigma_tex_cur <= 1e-9f) {
continue;
}
auto denoised_opt = model(x, sigma_tex_cur, i + 1);
if (denoised_opt.pred.empty()) {
return {};
}
sd::Tensor<float> denoised = std::move(denoised_opt.pred);
sd::Tensor<float> velocity = (x - denoised) / sigma_tex_cur;
auto x_sem = sd::ops::slice(x, 2, 0, SefiFlowDenoiser::kSemChannels);
auto x_tex = sd::ops::slice(x, 2, SefiFlowDenoiser::kSemChannels, SefiFlowDenoiser::kTotalChannels);
auto vel_sem = sd::ops::slice(velocity, 2, 0, SefiFlowDenoiser::kSemChannels);
auto vel_tex = sd::ops::slice(velocity, 2, SefiFlowDenoiser::kSemChannels, SefiFlowDenoiser::kTotalChannels);
auto x_sem_next = x_sem + vel_sem * (sigma_sem_next - sigma_sem_cur);
auto x_tex_next = x_tex + vel_tex * (sigma_tex_next - sigma_tex_cur);
sd::ops::slice_assign(&x, 2, 0, SefiFlowDenoiser::kSemChannels, x_sem_next);
sd::ops::slice_assign(&x, 2, SefiFlowDenoiser::kSemChannels, SefiFlowDenoiser::kTotalChannels, x_tex_next);
}
return x;
}
static sd::Tensor<float> sample_euler(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas) {
@ -2055,7 +2165,13 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
std::shared_ptr<RNG> rng,
float eta,
bool is_flow_denoiser,
const char* extra_sample_args) {
const char* extra_sample_args,
std::shared_ptr<Denoiser> denoiser_for_dispatch = nullptr) {
if (denoiser_for_dispatch) {
if (auto sefi = std::dynamic_pointer_cast<SefiFlowDenoiser>(denoiser_for_dispatch)) {
return sample_sefi_euler(sefi.get(), model, std::move(x));
}
}
SamplerExtraArgs extra_args = parse_key_value_args(extra_sample_args, "extra sample arg");
switch (method) {
case EULER_A_SAMPLE_METHOD:

View File

@ -96,6 +96,7 @@ const char* model_version_to_str[] = {
"Longcat-Image",
"PiD",
"Ideogram 4",
"SeFi-Image",
"Krea2",
"ESRGAN",
};
@ -691,7 +692,7 @@ public:
version,
sd_ctx_params->chroma_use_dit_mask,
model_manager);
} else if (sd_version_is_flux2(version)) {
} else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
bool is_chroma = false;
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
tensor_storage_map,
@ -1295,6 +1296,8 @@ public:
} else if (sd_version_is_krea2(version)) {
default_flow_shift = 1.15f;
}
} else if (sd_version_is_sefi_image(version)) {
pred_type = SEFI_FLOW_PRED;
} else if (sd_version_is_flux2(version)) {
pred_type = FLUX2_FLOW_PRED;
} else {
@ -1334,6 +1337,11 @@ public:
denoiser = std::make_shared<Flux2FlowDenoiser>();
break;
}
case SEFI_FLOW_PRED: {
LOG_INFO("running in SeFi-Image dual-time FLOW mode");
denoiser = std::make_shared<SefiFlowDenoiser>();
break;
}
default: {
LOG_ERROR("Unknown predition type %i", pred_type);
return false;
@ -1639,7 +1647,16 @@ public:
std::vector<float> process_timesteps(const std::vector<float>& timesteps,
const sd::Tensor<float>& init_latent,
const sd::Tensor<float>& denoise_mask) {
const sd::Tensor<float>& denoise_mask,
int step) {
if (auto sefi_denoiser = std::dynamic_pointer_cast<SefiFlowDenoiser>(denoiser)) {
int sched_idx = step > 0 ? step - 1 : 0;
if (sched_idx >= static_cast<int>(sefi_denoiser->tex_timesteps.size())) {
sched_idx = static_cast<int>(sefi_denoiser->tex_timesteps.size()) - 1;
}
return {sefi_denoiser->sem_timesteps[sched_idx],
sefi_denoiser->tex_timesteps[sched_idx]};
}
if (diffusion_model->get_desc() == "Wan2.2-TI2V-5B") {
int64_t frame_count = init_latent.shape()[2];
auto new_timesteps = std::vector<float>(static_cast<size_t>(frame_count), timesteps[0]);
@ -2051,7 +2068,7 @@ public:
timesteps_vec = process_ltxav_video_timesteps(base_timesteps_vec, init_latent, denoise_mask);
audio_timesteps_tensor = sd::Tensor<float>({static_cast<int64_t>(base_timesteps_vec.size())}, base_timesteps_vec);
} else {
timesteps_vec = process_timesteps(timesteps_vec, init_latent, denoise_mask);
timesteps_vec = process_timesteps(timesteps_vec, init_latent, denoise_mask, step);
}
const std::vector<float>& scaling_timesteps_vec = (sd_version_is_ltxav(version) && !denoise_mask.empty())
? base_timesteps_vec
@ -2121,7 +2138,7 @@ public:
diffusion_params.extra = UNetDiffusionExtra{-1, &controls, control_strength};
} else if (sd_version_is_sd3(version)) {
diffusion_params.extra = SkipLayerDiffusionExtra{local_skip_layers};
} else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version)) {
} else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) {
diffusion_params.extra = FluxDiffusionExtra{&guidance_tensor,
local_skip_layers};
} else if (sd_version_is_anima(version)) {
@ -2265,7 +2282,7 @@ public:
return output;
};
auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser, extra_sample_args);
auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser, extra_sample_args, denoiser);
if (x0_opt.empty()) {
LOG_ERROR("Diffusion model sampling failed");
if (control_net) {
@ -2326,6 +2343,8 @@ public:
latent_channel = 3;
} else if (sd_version_is_pid(version)) {
latent_channel = 3;
} else if (sd_version_is_sefi_image(version)) {
latent_channel = 144;
} else if (sd_version_uses_flux2_vae(version)) {
latent_channel = 128;
} else {