sd: sync with master-746-2574f59 (#2291)

* sd: clean up SD_USE_ defines

* sd: sync with master-721-8caa3f9+5 (9956436)

* sd: simplify taehv selection

* sd: sync with master-731-9f855c9

* sd: sync with master-737-3b6c9ca

* sd: sync with master-741-484baa4

* sd: sync with master-743-3590aa8

* sd: sync with master-746-2574f59

* sd: fix ggml_ext_pad_ext call
This commit is contained in:
Wagner Bruna 2026-07-03 04:28:51 -03:00 committed by GitHub
parent 2ad093af75
commit 6482a596e1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
43 changed files with 3009 additions and 487 deletions

View file

@ -101,9 +101,9 @@ NONECFLAGS =
LLAMA_USE_BUNDLED_GLSLC := 1
FAILSAFE_FLAGS = -DUSE_FAILSAFE
VULKAN_FLAGS = -DGGML_USE_VULKAN -DSD_USE_VULKAN
VULKAN_FLAGS = -DGGML_USE_VULKAN
ifdef LLAMA_CUBLAS
CUBLAS_FLAGS = -DGGML_USE_CUDA -DSD_USE_CUDA
CUBLAS_FLAGS = -DGGML_USE_CUDA
else
CUBLAS_FLAGS =
endif
@ -215,7 +215,7 @@ OBJS_CUDA_TEMP_INST += \
ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.o
ifdef LLAMA_CUBLAS
CUBLAS_FLAGS = -DGGML_USE_CUDA -DSD_USE_CUDA -I/usr/local/cuda/include -I/opt/cuda/include -I$(CUDA_PATH)/targets/x86_64-linux/include
CUBLAS_FLAGS = -DGGML_USE_CUDA -I/usr/local/cuda/include -I/opt/cuda/include -I$(CUDA_PATH)/targets/x86_64-linux/include
CUBLASLD_FLAGS = -lcuda -lcublas -lcudart -lcublasLt -lpthread -ldl -lrt -L/usr/local/cuda/lib64 -L/opt/cuda/lib64 -L$(CUDA_PATH)/targets/x86_64-linux/lib -L$(CUDA_PATH)/lib64/stubs -L/usr/local/cuda/targets/aarch64-linux/lib -L/usr/local/cuda/targets/sbsa-linux/lib -L/usr/lib/wsl/lib
CUBLAS_OBJS = ggml-cuda.o ggml_v3-cuda.o ggml_v2-cuda.o ggml_v2-cuda-legacy.o
CUBLAS_OBJS += $(patsubst %.cu,%.o,$(filter-out ggml/src/ggml-cuda/ggml-cuda.cu, $(wildcard ggml/src/ggml-cuda/*.cu)))
@ -315,7 +315,7 @@ HIPFLAGS += -DGGML_HIP_NO_ROCWMMA_FATTN
endif
endif
HIPFLAGS += -DGGML_USE_HIP -DGGML_HIP_NO_VMM -DGGML_USE_CUDA -DSD_USE_CUDA $(shell $(ROCM_PATH)/bin/hipconfig -C)
HIPFLAGS += -DGGML_USE_HIP -DGGML_HIP_NO_VMM -DGGML_USE_CUDA $(shell $(ROCM_PATH)/bin/hipconfig -C)
HIPLDFLAGS += -L$(ROCM_PATH)/lib -Wl,-rpath=$(ROCM_PATH)/lib
HIPLDFLAGS += -L$(ROCM_PATH)/lib64 -Wl,-rpath=$(ROCM_PATH)/lib64
HIPLDFLAGS += -lhipblas -lamdhip64 -lrocblas
@ -339,8 +339,8 @@ endif # LLAMA_HIPBLAS
ifdef LLAMA_METAL
CFLAGS += -DGGML_USE_METAL -DGGML_METAL_NDEBUG -DSD_USE_METAL
CXXFLAGS += -DGGML_USE_METAL -DSD_USE_METAL
CFLAGS += -DGGML_USE_METAL -DGGML_METAL_NDEBUG
CXXFLAGS += -DGGML_USE_METAL
LDFLAGS += -framework Foundation -framework Metal -framework MetalKit -framework MetalPerformanceShaders
OBJS += ggml-metal.o ggml-metal-device.o ggml-metal-device-m.o ggml-metal-context-m.o ggml-metal-common.o ggml-metal-ops.o
@ -703,7 +703,7 @@ budget.o: common/reasoning-budget.cpp common/reasoning-budget.h
chat.o: common/chat.cpp common/chat.h
$(CXX) $(CXXFLAGS) -c $< -o $@
SDCPP_COMMON_BASENAMES := include/stable-diffusion.h src/conditioning/conditioner.hpp src/core/ggml_extend_backend.cpp src/core/ggml_extend_backend.h src/core/ggml_extend.hpp src/core/ggml_graph_cut.cpp src/core/ggml_graph_cut.h src/core/ordered_map.hpp src/core/rng.hpp src/core/rng_mt19937.hpp src/core/rng_philox.hpp src/core/tensor_ggml.hpp src/core/tensor.hpp src/core/util.cpp src/core/util.h src/extensions/generation_extension.h src/extensions/photomaker_extension.cpp src/extensions/pulid_extension.cpp src/kcpp_sd_extensions.h src/model/adapter/lora.hpp src/model/adapter/pmid.hpp src/model/adapter/pulid.hpp src/model/common/block.hpp src/model/common/rope.hpp src/model/diffusion/anima.hpp src/model/diffusion/boogu.hpp src/model/diffusion/control.hpp src/model/diffusion/dit.hpp src/model/diffusion/ernie_image.hpp src/model/diffusion/flux.hpp src/model/diffusion/hidream_o1.hpp src/model/diffusion/ideogram4.hpp src/model/diffusion/krea2.hpp src/model/diffusion/lens.hpp src/model/diffusion/ltxv.hpp src/model/diffusion/mmdit.hpp src/model/diffusion/model.hpp src/model/diffusion/pid.hpp src/model/diffusion/qwen_image.hpp src/model/diffusion/unet.hpp src/model/diffusion/wan.hpp src/model/diffusion/z_image.hpp src/model.h src/model_io/binary_io.h src/model_io/gguf_io.cpp src/model_io/gguf_io.h src/model_io/gguf_reader_ext.h src/model_io/pickle_io.cpp src/model_io/pickle_io.h src/model_io/safetensors_io.cpp src/model_io/safetensors_io.h src/model_io/tensor_storage.h src/model_io/torch_legacy_io.cpp src/model_io/torch_legacy_io.h src/model_io/torch_zip_io.cpp src/model_io/torch_zip_io.h src/model_loader.cpp src/model_loader.h src/model_manager.cpp src/model_manager.h src/model/te/clip.hpp src/model/te/llm.hpp src/model/te/t5.hpp src/model/upscaler/esrgan.hpp src/model/upscaler/ltx_latent_upscaler.hpp src/model/vae/auto_encoder_kl.hpp src/model/vae/ltx_audio_vae.hpp src/model/vae/ltx_vae.hpp src/model/vae/tae.hpp src/model/vae/vae.hpp src/model/vae/wan_vae.hpp src/name_conversion.cpp src/name_conversion.h src/runtime/cache_dit.hpp src/runtime/condition_cache_utils.hpp src/runtime/denoiser.hpp src/runtime/easycache.hpp src/runtime/gits_noise.h src/runtime/guidance.cpp src/runtime/guidance.h src/runtime/latent-preview.h src/runtime/preprocessing.hpp src/runtime/sample-cache.cpp src/runtime/sample-cache.h src/runtime/spectrum.hpp src/runtime/ucache.hpp src/stable-diffusion.cpp src/tokenizers/bpe_tokenizer.cpp src/tokenizers/bpe_tokenizer.h src/tokenizers/clip_tokenizer.cpp src/tokenizers/clip_tokenizer.h src/tokenizers/gemma_tokenizer.cpp src/tokenizers/gemma_tokenizer.h src/tokenizers/gpt_oss_tokenizer.cpp src/tokenizers/gpt_oss_tokenizer.h src/tokenizers/mistral_tokenizer.cpp src/tokenizers/mistral_tokenizer.h src/tokenizers/qwen2_tokenizer.cpp src/tokenizers/qwen2_tokenizer.h src/tokenizers/t5_unigram_tokenizer.cpp src/tokenizers/t5_unigram_tokenizer.h src/tokenizers/tokenizer.cpp src/tokenizers/tokenizer.h src/tokenizers/tokenize_util.cpp src/tokenizers/tokenize_util.h src/tokenizers/vocab/vocab.h src/upscaler.cpp src/upscaler.h src/weight_manager.h
SDCPP_COMMON_BASENAMES := include/stable-diffusion.h src/conditioning/conditioner.hpp src/core/ggml_extend_backend.cpp src/core/ggml_extend_backend.h src/core/ggml_extend.hpp src/core/ggml_graph_cut.cpp src/core/ggml_graph_cut.h src/core/ordered_map.hpp src/core/rng.hpp src/core/rng_mt19937.hpp src/core/rng_philox.hpp src/core/tensor_ggml.hpp src/core/tensor.hpp src/core/util.cpp src/core/util.h src/extensions/generation_extension.h src/extensions/photomaker_extension.cpp src/extensions/pulid_extension.cpp src/kcpp_sd_extensions.h src/model/adapter/lora.hpp src/model/adapter/pmid.hpp src/model/adapter/pulid.hpp src/model/common/block.hpp src/model/common/rope.hpp src/model/diffusion/anima.hpp src/model/diffusion/boogu.hpp src/model/diffusion/control.hpp src/model/diffusion/dit.hpp src/model/diffusion/ernie_image.hpp src/model/diffusion/flux.hpp src/model/diffusion/hidream_o1.hpp src/model/diffusion/ideogram4.hpp src/model/diffusion/krea2.hpp src/model/diffusion/lens.hpp src/model/diffusion/ltxv.hpp src/model/diffusion/minit2i.hpp src/model/diffusion/mmdit.hpp src/model/diffusion/model.hpp src/model/diffusion/pid.hpp src/model/diffusion/qwen_image.hpp src/model/diffusion/sefi_image.hpp src/model/diffusion/unet.hpp src/model/diffusion/wan.hpp src/model/diffusion/z_image.hpp src/model.h src/model_io/binary_io.h src/model_io/gguf_io.cpp src/model_io/gguf_io.h src/model_io/gguf_reader_ext.h src/model_io/pickle_io.cpp src/model_io/pickle_io.h src/model_io/safetensors_io.cpp src/model_io/safetensors_io.h src/model_io/tensor_storage.h src/model_io/torch_legacy_io.cpp src/model_io/torch_legacy_io.h src/model_io/torch_zip_io.cpp src/model_io/torch_zip_io.h src/model_loader.cpp src/model_loader.h src/model_manager.cpp src/model_manager.h src/model/te/clip.hpp src/model/te/llm.hpp src/model/te/t5.hpp src/model/upscaler/esrgan.hpp src/model/upscaler/ltx_latent_upscaler.hpp src/model/vae/auto_encoder_kl.hpp src/model/vae/ltx_audio_vae.hpp src/model/vae/ltx_vae.hpp src/model/vae/tae.hpp src/model/vae/vae.hpp src/model/vae/wan_vae.hpp src/name_conversion.cpp src/name_conversion.h src/runtime/cache_dit.hpp src/runtime/condition_cache_utils.hpp src/runtime/denoiser.hpp src/runtime/easycache.hpp src/runtime/gits_noise.h src/runtime/guidance.cpp src/runtime/guidance.h src/runtime/imatrix.cpp src/runtime/imatrix.h src/runtime/latent-preview.h src/runtime/preprocessing.hpp src/runtime/sample-cache.cpp src/runtime/sample-cache.h src/runtime/spectrum.hpp src/runtime/ucache.hpp src/stable-diffusion.cpp src/tokenizers/bpe_tokenizer.cpp src/tokenizers/bpe_tokenizer.h src/tokenizers/clip_tokenizer.cpp src/tokenizers/clip_tokenizer.h src/tokenizers/gemma_tokenizer.cpp src/tokenizers/gemma_tokenizer.h src/tokenizers/gpt_oss_tokenizer.cpp src/tokenizers/gpt_oss_tokenizer.h src/tokenizers/mistral_tokenizer.cpp src/tokenizers/mistral_tokenizer.h src/tokenizers/qwen2_tokenizer.cpp src/tokenizers/qwen2_tokenizer.h src/tokenizers/t5_unigram_tokenizer.cpp src/tokenizers/t5_unigram_tokenizer.h src/tokenizers/tokenizer.cpp src/tokenizers/tokenizer.h src/tokenizers/tokenize_util.cpp src/tokenizers/tokenize_util.h src/tokenizers/vocab/vocab.h src/upscaler.cpp src/upscaler.h src/weight_manager.h
SDCPP_MAIN_BASENAMES := examples/cli/image_metadata.cpp examples/cli/image_metadata.h examples/cli/main.cpp examples/cli/msf_gif.h examples/common/common.cpp examples/common/common.h examples/common/log.cpp examples/common/log.h examples/common/media_io.cpp examples/common/media_io.h examples/common/resource_owners.hpp src/tokenizers/vocab/clip_merges.hpp src/tokenizers/vocab/gemma2_merges.hpp src/tokenizers/vocab/gemma2_vocab.hpp src/tokenizers/vocab/gemma_merges.hpp src/tokenizers/vocab/gemma_vocab.hpp src/tokenizers/vocab/gpt_oss_merges.hpp src/tokenizers/vocab/gpt_oss_vocab.hpp src/tokenizers/vocab/mistral_merges.hpp src/tokenizers/vocab/mistral_vocab.hpp src/tokenizers/vocab/qwen_merges.hpp src/tokenizers/vocab/t5.hpp src/tokenizers/vocab/umt5.hpp src/tokenizers/vocab/vocab.cpp src/convert.cpp src/version.cpp
@ -782,9 +782,9 @@ clean:
main: tools/completion/main.cpp tools/completion/completion.cpp common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
mainvk: tools/completion/main.cpp tools/completion/completion.cpp common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
fitparams: tools/fit-params/main.cpp tools/fit-params/fit-params.cpp common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
sdmain: $(OBJS_SDCOMMON) $(OBJS_SDMAIN) build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
whispermain: otherarch/whispercpp/main.cpp otherarch/whispercpp/whisper.cpp build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
@ -798,17 +798,17 @@ mtmd-cli: tools/mtmd/mtmd-cli.cpp tools/mtmd/clip.cpp common/debug.cpp common/ar
embedding: examples/embedding/embedding.cpp common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) src/llama-cparams.cpp build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
embeddingvk: examples/embedding/embedding.cpp common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) src/llama-cparams.cpp build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
ttscppmain: otherarch/ttscpp/cli/cli.cpp otherarch/ttscpp/cli/playback.cpp otherarch/ttscpp/cli/playback.h otherarch/ttscpp/cli/write_file.cpp otherarch/ttscpp/cli/write_file.h otherarch/ttscpp/cli/vad.cpp otherarch/ttscpp/cli/vad.h otherarch/ttscpp/src/ttscpp.cpp otherarch/ttscpp/src/ttstokenizer.cpp otherarch/ttscpp/src/ttssampler.cpp otherarch/ttscpp/src/parler_model.cpp otherarch/ttscpp/src/dac_model.cpp otherarch/ttscpp/src/ttsutil.cpp otherarch/ttscpp/src/ttsargs.cpp otherarch/ttscpp/src/ttst5_encoder_model.cpp otherarch/ttscpp/src/phonemizer.cpp otherarch/ttscpp/src/tts_model.cpp otherarch/ttscpp/src/kokoro_model.cpp otherarch/ttscpp/src/dia_model.cpp otherarch/ttscpp/src/orpheus_model.cpp otherarch/ttscpp/src/snac_model.cpp otherarch/ttscpp/src/general_neural_audio_codec.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
qwen3tts: otherarch/qwen3tts/q3ttsmain.cpp otherarch/qwen3tts/qwen3_tts.cpp otherarch/qwen3tts/text_tokenizer.cpp otherarch/qwen3tts/gguf_loader.cpp otherarch/qwen3tts/tts_transformer.cpp otherarch/qwen3tts/audio_tokenizer_decoder.cpp otherarch/qwen3tts/audio_tokenizer_encoder.cpp otherarch/qwen3tts/coreml_code_predictor_stub.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
rpcserver: tools/rpc/rpc-server.cpp common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
llamaserver: $(LLAMASERVER_SRCS) $(LLAMASERVER_COMMON_SRCS) build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
$(CXX) $(CXXFLAGS) $(LLAMASERVER_CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
llamaservervk: $(LLAMASERVER_SRCS) $(LLAMASERVER_COMMON_SRCS) build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
$(CXX) $(CXXFLAGS) $(LLAMASERVER_CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
$(CXX) $(CXXFLAGS) $(LLAMASERVER_CXXFLAGS) -DGGML_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
ggml/src/ggml-vulkan-shaders.cpp: ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
ifdef VULKAN_BUILD

View file

@ -1,6 +1,7 @@
#include <stdio.h>
#include <string.h>
#include <time.h>
#include <algorithm>
#include <cctype>
#include <filesystem>
#include <functional>
@ -53,6 +54,9 @@ struct SDCliParams {
bool metadata_brief = false;
bool metadata_all = false;
std::string imatrix_out;
std::vector<std::string> imatrix_in;
bool normal_exit = false;
ArgOptions get_options() {
@ -79,6 +83,11 @@ struct SDCliParams {
"path to write preview image to (default: ./preview.png). Multi-frame previews support .avi, .webm, and animated .webp",
0,
&preview_path},
{"",
"--imat-out",
"compute the imatrix for this run and save it to the provided path",
0,
&imatrix_out},
};
options.int_options = {
@ -179,6 +188,14 @@ struct SDCliParams {
return -1;
};
auto on_imatrix_in_arg = [&](int argc, const char** argv, int index) {
if (++index >= argc) {
return -1;
}
imatrix_in.push_back(argv[index]);
return 1;
};
options.manual_options = {
{"-M",
"--mode",
@ -192,6 +209,10 @@ struct SDCliParams {
"--help",
"show this help message and exit",
on_help_arg},
{"",
"--imat-in",
"load an imatrix file for quantization or continued collection; can be specified multiple times",
on_imatrix_in_arg},
};
return options;
@ -253,6 +274,7 @@ struct SDCliParams {
<< " preview_fps: " << preview_fps << ",\n"
<< " taesd_preview: " << (taesd_preview ? "true" : "false") << ",\n"
<< " preview_noisy: " << (preview_noisy ? "true" : "false") << ",\n"
<< " imatrix_out: \"" << imatrix_out << "\",\n"
<< " metadata_raw: " << (metadata_raw ? "true" : "false") << ",\n"
<< " metadata_brief: " << (metadata_brief ? "true" : "false") << ",\n"
<< " metadata_all: " << (metadata_all ? "true" : "false") << "\n"
@ -459,7 +481,8 @@ bool save_results(const SDCliParams& cli_params,
if (!img.data)
return false;
const int64_t metadata_seed = cli_params.mode == VID_GEN ? gen_params.seed : gen_params.seed + idx;
int images_per_batch = gen_params.batch_count > 0 ? std::max(1, num_results / gen_params.batch_count) : 1;
const int64_t metadata_seed = cli_params.mode == VID_GEN ? gen_params.seed : gen_params.seed + idx / images_per_batch;
std::string params = gen_params.embed_image_metadata
? get_image_params(ctx_params, gen_params, metadata_seed, cli_params.mode)
: "";
@ -605,13 +628,32 @@ int main(int argc, const char* argv[]) {
LOG_DEBUG("%s", ctx_params.to_string().c_str());
LOG_DEBUG("%s", gen_params.to_string().c_str());
if (!cli_params.imatrix_out.empty()) {
if (fs::exists(cli_params.imatrix_out) &&
std::find(cli_params.imatrix_in.begin(), cli_params.imatrix_in.end(), cli_params.imatrix_out) == cli_params.imatrix_in.end()) {
LOG_WARN("imatrix file '%s' already exists and will be overwritten", cli_params.imatrix_out.c_str());
}
enable_imatrix_collection();
}
for (const auto& in_file : cli_params.imatrix_in) {
LOG_INFO("loading imatrix from '%s'", in_file.c_str());
if (!load_imatrix(in_file.c_str())) {
LOG_WARN("failed to load imatrix from '%s'", in_file.c_str());
}
}
if (cli_params.mode == CONVERT) {
bool success = convert(ctx_params.model_path.c_str(),
ctx_params.vae_path.c_str(),
cli_params.output_path.c_str(),
ctx_params.wtype,
ctx_params.tensor_type_rules.c_str(),
cli_params.convert_name);
bool success = convert_with_components(ctx_params.model_path.c_str(),
ctx_params.clip_l_path.c_str(),
ctx_params.clip_g_path.c_str(),
ctx_params.t5xxl_path.c_str(),
ctx_params.diffusion_model_path.c_str(),
ctx_params.vae_path.c_str(),
cli_params.output_path.c_str(),
ctx_params.wtype,
ctx_params.tensor_type_rules.c_str(),
cli_params.convert_name);
if (!success) {
LOG_ERROR("convert '%s'/'%s' to '%s' failed",
ctx_params.model_path.c_str(),
@ -766,8 +808,12 @@ int main(int argc, const char* argv[]) {
if (cli_params.mode == IMG_GEN) {
sd_img_gen_params_t img_gen_params = gen_params.to_sd_img_gen_params_t();
num_results = gen_params.batch_count;
results.adopt(generate_image(sd_ctx.get(), &img_gen_params), num_results);
sd_image_t* generated_images = nullptr;
if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) {
generated_images = nullptr;
num_results = 0;
}
results.adopt(generated_images, num_results);
} else if (cli_params.mode == VID_GEN) {
sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t();
sd_image_t* generated_video = nullptr;
@ -802,12 +848,22 @@ int main(int argc, const char* argv[]) {
SDImageOwner current_image(results[i]);
results[i] = {0, 0, 0, nullptr};
for (int u = 0; u < gen_params.upscale_repeats; ++u) {
SDImageOwner upscaled_image(upscale(upscaler_ctx.get(), current_image.get(), upscale_factor));
if (upscaled_image.get().data == nullptr) {
sd_image_t* upscaled_images = nullptr;
int upscaled_count = 0;
bool upscale_ok = upscale(upscaler_ctx.get(),
current_image.get(),
upscale_factor,
&upscaled_images,
&upscaled_count);
if (!upscale_ok || upscaled_count <= 0 || upscaled_images[0].data == nullptr) {
free_sd_images(upscaled_images, upscaled_count);
LOG_ERROR("upscale failed");
break;
}
current_image = std::move(upscaled_image);
sd_image_t upscaled_image = upscaled_images[0];
upscaled_images[0] = {0, 0, 0, nullptr};
free_sd_images(upscaled_images, upscaled_count);
current_image.reset(upscaled_image);
}
results[i] = current_image.release(); // Set the final upscaled image as the result
}
@ -819,6 +875,11 @@ int main(int argc, const char* argv[]) {
return 1;
}
if (!cli_params.imatrix_out.empty()) {
LOG_INFO("saving imatrix to '%s'", cli_params.imatrix_out.c_str());
save_imatrix(cli_params.imatrix_out.c_str());
}
free_sd_audio(generated_audio);
return 0;

View file

@ -653,7 +653,7 @@ ArgOptions SDContextParams::get_options() {
on_sampler_rng_arg},
{"",
"--prediction",
"prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, flux2_flow]",
"prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, sefi_flow]",
on_prediction_arg},
{"",
"--lora-apply-mode",
@ -710,7 +710,18 @@ bool SDContextParams::resolve(SDMode mode) {
}
bool SDContextParams::validate(SDMode mode) {
if (mode != UPSCALE && mode != METADATA && model_path.length() == 0 && diffusion_model_path.length() == 0) {
if (mode == CONVERT) {
const bool has_convert_input = model_path.length() != 0 ||
clip_l_path.length() != 0 ||
clip_g_path.length() != 0 ||
t5xxl_path.length() != 0 ||
diffusion_model_path.length() != 0 ||
vae_path.length() != 0;
if (!has_convert_input) {
LOG_ERROR("error: convert mode needs at least one model input path\n");
return false;
}
} else if (mode != UPSCALE && mode != METADATA && model_path.length() == 0 && diffusion_model_path.length() == 0) {
LOG_ERROR("error: the following arguments are required: model_path/diffusion_model\n");
return false;
}
@ -960,7 +971,7 @@ ArgOptions SDGenerationParams::get_options() {
&hires_upscaler},
{"",
"--extra-sample-args",
"extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma;; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware",
"extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma;; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware",
(int)',',
&extra_sample_args},
{"",
@ -996,6 +1007,10 @@ ArgOptions SDGenerationParams::get_options() {
"--batch-count",
"batch count",
&batch_count},
{"",
"--qwen-image-layers",
"number of Qwen Image Layered layers; latent/output count is layers + 1 (default: 3)",
&qwen_image_layers},
{"",
"--video-frames",
"video frames (default: 1)",
@ -1475,7 +1490,7 @@ ArgOptions SDGenerationParams::get_options() {
on_high_noise_sample_method_arg},
{"",
"--scheduler",
"denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2, logit_normal], default: model-specific",
"denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2, logit_normal, flux2, flux, beta], alias: normal=discrete, default: model-specific",
on_scheduler_arg},
{"",
"--sigmas",
@ -1816,6 +1831,7 @@ bool SDGenerationParams::from_json_str(
load_if_exists("width", width);
load_if_exists("height", height);
load_if_exists("batch_count", batch_count);
load_if_exists("qwen_image_layers", qwen_image_layers);
load_if_exists("video_frames", video_frames);
load_if_exists("fps", fps);
load_if_exists("upscale_repeats", upscale_repeats);
@ -2240,6 +2256,11 @@ bool SDGenerationParams::validate(SDMode mode) {
return false;
}
if (qwen_image_layers < 0) {
LOG_ERROR("error: qwen_image_layers must be non-negative");
return false;
}
if (sample_params.sample_steps <= 0) {
LOG_ERROR("error: the sample_steps must be greater than 0\n");
return false;
@ -2406,6 +2427,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
params.strength = strength;
params.seed = seed;
params.batch_count = batch_count;
params.qwen_image_layers = qwen_image_layers;
params.control_image = control_image.get();
params.control_strength = control_strength;
params.pm_params = pm_params;
@ -2531,6 +2553,7 @@ std::string SDGenerationParams::to_string() const {
<< " width: " << width << ",\n"
<< " height: " << height << ",\n"
<< " batch_count: " << batch_count << ",\n"
<< " qwen_image_layers: " << qwen_image_layers << ",\n"
<< " init_image_path: \"" << init_image_path << "\",\n"
<< " end_image_path: \"" << end_image_path << "\",\n"
<< " mask_image_path: \"" << mask_image_path << "\",\n"

View file

@ -197,6 +197,7 @@ struct SDGenerationParams {
int width = -1;
int height = -1;
int batch_count = 1;
int qwen_image_layers = 3;
int64_t seed = 42;
float strength = 0.75f;
float control_strength = 0.9f;

View file

@ -71,6 +71,9 @@ enum scheduler_t {
BONG_TANGENT_SCHEDULER,
LTX2_SCHEDULER,
LOGIT_NORMAL_SCHEDULER,
FLUX2_SCHEDULER,
FLUX_SCHEDULER,
BETA_SCHEDULER,
SCHEDULER_COUNT
};
@ -80,7 +83,8 @@ enum prediction_t {
EDM_V_PRED,
FLOW_PRED,
FLUX_FLOW_PRED,
FLUX2_FLOW_PRED,
SEFI_FLOW_PRED,
MINIT2I_FLOW_PRED,
PREDICTION_COUNT
};
@ -376,6 +380,7 @@ typedef struct {
sd_tiling_params_t vae_tiling_params;
sd_cache_params_t cache;
sd_hires_params_t hires;
int qwen_image_layers;
} sd_img_gen_params_t;
typedef struct {
@ -405,14 +410,17 @@ typedef struct {
} sd_vid_gen_params_t;
typedef struct sd_ctx_t sd_ctx_t;
struct ggml_tensor;
typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data);
typedef void (*sd_progress_cb_t)(int step, int steps, float time, void* data);
typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data);
typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void* user_data);
SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data);
SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data);
SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data);
SD_API void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data);
SD_API int32_t sd_get_num_physical_cores();
SD_API const char* sd_get_system_info();
SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx);
@ -453,7 +461,10 @@ SD_API enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sa
SD_API void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params);
SD_API char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params);
SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params);
SD_API bool generate_image(sd_ctx_t* sd_ctx,
const sd_img_gen_params_t* sd_img_gen_params,
sd_image_t** images_out,
int* num_images_out);
enum sd_cancel_mode_t {
// Stop the current generation as soon as possible.
@ -483,9 +494,11 @@ SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path,
const char* params_backend);
SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx);
SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx,
sd_image_t input_image,
uint32_t upscale_factor);
SD_API bool upscale(upscaler_ctx_t* upscaler_ctx,
sd_image_t input_image,
uint32_t upscale_factor,
sd_image_t** images_out,
int* num_images_out);
SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx);
@ -496,6 +509,17 @@ SD_API bool convert(const char* input_path,
const char* tensor_type_rules,
bool convert_name);
SD_API bool convert_with_components(const char* model_path,
const char* clip_l_path,
const char* clip_g_path,
const char* t5xxl_path,
const char* diffusion_model_path,
const char* vae_path,
const char* output_path,
enum sd_type_t output_type,
const char* tensor_type_rules,
bool convert_name);
SD_API bool preprocess_canny(sd_image_t image,
float high_threshold,
float low_threshold,
@ -503,6 +527,11 @@ SD_API bool preprocess_canny(sd_image_t image,
float strong,
bool inverse);
SD_API bool load_imatrix(const char* imatrix_path);
SD_API void save_imatrix(const char* imatrix_path);
SD_API void enable_imatrix_collection(void);
SD_API void disable_imatrix_collection(void);
SD_API const char* sd_commit(void);
SD_API const char* sd_version(void);

View file

@ -977,6 +977,21 @@ bool supports_reference_images(kcpp_sd::model_info info)
return supported;
}
static std::string upscale_image_to_png_base64(upscaler_ctx_t* upscaler_ctx, const sd_image_t& input_image, int upscale_factor = 2, const std::string& meta_image_info = "")
{
std::string gen_data;
sd_image_t* upscaled = nullptr;
int upscaled_count = 0;
if (upscale(upscaler_ctx, input_image, upscale_factor, &upscaled, &upscaled_count)) {
gen_data = raw_image_to_png_base64(*upscaled, meta_image_info);
free_sd_images(upscaled, upscaled_count);
} else {
printf("Upscaling failed!\n");
gen_data = raw_image_to_png_base64(input_image, meta_image_info);
}
return gen_data;
}
sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
{
if(sd_ctx == nullptr || sd_params == nullptr)
@ -984,6 +999,7 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
return sd_generation.error("Warning: KCPP image generation not initialized!");
}
sd_image_t * results = nullptr;
int generated_num_results = 0;
std::string img2img_data = std::string(inputs.init_images);
std::string img2img_mask = std::string(inputs.mask);
@ -1309,7 +1325,6 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
//the below params are only used in video models. May move into standalone object in future
int vid_req_frames = inputs.vid_req_frames;
int video_output_type = inputs.video_output_type;
int generated_num_results = 1;
int vid_fps = inputs.vid_fps;
remove_limits = inputs.remove_limits;
@ -1399,7 +1414,10 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
fflush(stdout);
results = generate_image(sd_ctx, &params);
if (!generate_image(sd_ctx, &params, &results, &generated_num_results)) {
results = nullptr;
generated_num_results = 0;
}
} else {
@ -1466,7 +1484,10 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
if (is_passthrough) {
printf("No generation triggered, passthrough mode.\n");
} else {
results = generate_image(sd_ctx, &params);
if (!generate_image(sd_ctx, &params, &results, &generated_num_results)) {
results = nullptr;
generated_num_results = 0;
}
}
}
@ -1511,111 +1532,102 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
jsoninfo["all_seeds"] = nlohmann::json::array();
jsoninfo["version"] = "KoboldCpp";
}
sd_image_t upscaled_image;
upscaled_image.data = nullptr;
sd_image_t* upscaled_image = nullptr;
std::string gen_data;
std::string gen_data2;
std::string final_frame_data;
if (is_passthrough)
{
//either return original image or upscale if needed
sd_image_t *result_image = &input_image;
if(inputs.upscale && upscaler_ctx != nullptr)
{
printf("Upscaling original image (passthrough)...\n");
upscaled_image = upscale(upscaler_ctx, input_image, 2);
result_image = &upscaled_image;
gen_data = upscale_image_to_png_base64(upscaler_ctx, input_image, 2);
}
gen_data = raw_image_to_png_base64(*result_image);
else {
gen_data = raw_image_to_png_base64(input_image);
}
}
else if (isanim)
{
//if multiframe, make a video
if (generated_num_results > 0 && results && results->data)
{
if(!sd_is_quiet && sddebugmode==1)
{
printf("\nSaving video buffer, VIDEO_OUTPUT_TYPE=%d...",video_output_type);
}
uint8_t * out_data = nullptr;
uint8_t * out_data2 = nullptr;
size_t out_len = 0;
size_t out_len2 = 0;
int status = 0;
int status2 = 0;
if(video_output_type==0 || video_output_type==2)
{
status = create_gif_buf_from_sd_images_msf(results, generated_num_results, vid_fps, &out_data,&out_len);
}
if(video_output_type==1 || video_output_type==2)
{
status2 = create_mjpg_avi_membuf_from_sd_images(results, generated_num_results, vid_fps, 40, &out_data2,&out_len2, generated_audio);
}
if(generated_num_results>1)
{
sd_image_t *final_frame_image = &results[generated_num_results-1];
final_frame_data = raw_image_to_png_base64(*final_frame_image);
}
if(!sd_is_quiet && sddebugmode==1)
{
printf("Video Output Sizes: GIF=%zu AVI=%zu\n",out_len,out_len2);
if(status==0 && status2==0)
{
printf("Video(s) Saved (Len %zu)!\n",out_len);
} else {
printf("Save Failed!\n");
}
}
if(status==0 && out_len>0)
{
gen_data = kcpp_base64_encode(out_data, out_len);
free(out_data);
}
if (status2 == 0 && out_len2 > 0) {
if (gen_data == "") {
gen_data = kcpp_base64_encode(out_data2, out_len2);
} else {
gen_data2 = kcpp_base64_encode(out_data2, out_len2);
}
free(out_data2);
}
}
free_sd_images(results, generated_num_results);
}
else
{
for (int i = 0; i < params.batch_count; i++)
for (int i = 0; i < generated_num_results; i++)
{
if (results[i].data == NULL) {
sd_image_t& result_image = results[i];
if (result_image.data == NULL) {
continue;
}
//if multiframe, make a video
if(isanim)
std::string meta_image_info = get_image_params(params, lora_meta, i);
if(inputs.upscale && upscaler_ctx != nullptr)
{
if(!sd_is_quiet && sddebugmode==1)
{
printf("\nSaving video buffer, VIDEO_OUTPUT_TYPE=%d...",video_output_type);
}
uint8_t * out_data = nullptr;
uint8_t * out_data2 = nullptr;
size_t out_len = 0;
size_t out_len2 = 0;
int status = 0;
int status2 = 0;
if(video_output_type==0 || video_output_type==2)
{
status = create_gif_buf_from_sd_images_msf(results, generated_num_results, vid_fps, &out_data,&out_len);
}
if(video_output_type==1 || video_output_type==2)
{
status2 = create_mjpg_avi_membuf_from_sd_images(results, generated_num_results, vid_fps, 40, &out_data2,&out_len2, generated_audio);
}
if(generated_num_results>1)
{
sd_image_t *final_frame_image = &results[generated_num_results-1];
final_frame_data = raw_image_to_png_base64(*final_frame_image);
}
if(!sd_is_quiet && sddebugmode==1)
{
printf("Video Output Sizes: GIF=%zu AVI=%zu\n",out_len,out_len2);
if(status==0 && status2==0)
{
printf("Video(s) Saved (Len %zu)!\n",out_len);
} else {
printf("Save Failed!\n");
}
}
if(status==0 && out_len>0)
{
gen_data = kcpp_base64_encode(out_data, out_len);
free(out_data);
}
if (status2 == 0 && out_len2 > 0) {
if (gen_data == "") {
gen_data = kcpp_base64_encode(out_data2, out_len2);
} else {
gen_data2 = kcpp_base64_encode(out_data2, out_len2);
}
free(out_data2);
}
printf("Upscaling output image...\n");
gen_data = upscale_image_to_png_base64(upscaler_ctx, result_image, 2, meta_image_info);
} else {
gen_data = raw_image_to_png_base64(result_image, meta_image_info);
}
else
{
sd_image_t *result_image = &results[i];
if(inputs.upscale && upscaler_ctx != nullptr)
{
printf("Upscaling output image...\n");
upscaled_image = upscale(upscaler_ctx, results[i], 2);
result_image = &upscaled_image;
}
std::string meta_image_info = get_image_params(params, lora_meta, i);
gen_data = raw_image_to_png_base64(*result_image, meta_image_info);
jsoninfo["infotexts"][i] = meta_image_info;
jsoninfo["all_seeds"][i] = params.seed + i;
jsoninfo["all_prompts"][i] = params.prompt;
jsoninfo["all_negative_prompts"][i] = params.negative_prompt;
}
free(results[i].data);
results[i].data = NULL;
jsoninfo["infotexts"][i] = meta_image_info;
jsoninfo["all_seeds"][i] = params.seed + i;
jsoninfo["all_prompts"][i] = params.prompt;
jsoninfo["all_negative_prompts"][i] = params.negative_prompt;
}
}
if(upscaled_image.data)
{
free(upscaled_image.data);
upscaled_image.data = nullptr;
free_sd_images(results, generated_num_results);
results = nullptr;
}
if (generated_audio) {
@ -1627,8 +1639,6 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
input_audio.data = nullptr;
}
free(results);
total_img_gens += 1;
if(!sd_is_quiet)
{
@ -1662,9 +1672,7 @@ sd_generation_outputs sdtype_upscale(const sd_upscale_inputs inputs)
}
upscale_src_buffer = load_image_from_b64(rawb64,nx,ny);
sd_image_t source_img;
sd_image_t upscaled_image;
source_img.data = nullptr;
upscaled_image.data = nullptr;
std::string result;
if(upscale_src_buffer)
{
@ -1673,10 +1681,7 @@ sd_generation_outputs sdtype_upscale(const sd_upscale_inputs inputs)
source_img.channel = 3;
source_img.data = upscale_src_buffer;
upscaled_image = upscale(upscaler_ctx, source_img, inputs.upscaling_resize);
result = raw_image_to_png_base64(upscaled_image);
free(upscaled_image.data);
result = upscale_image_to_png_base64(upscaler_ctx, source_img, inputs.upscaling_resize);
}
if (result == "") {

View file

@ -1,4 +1,4 @@
#ifndef __SD_CONDITIONING_CONDITIONER_HPP__
#ifndef __SD_CONDITIONING_CONDITIONER_HPP__
#define __SD_CONDITIONING_CONDITIONER_HPP__
#include <cmath>
@ -1378,6 +1378,101 @@ struct T5CLIPEmbedder : public Conditioner {
}
};
struct MiniT2IConditioner : public Conditioner {
T5UniGramTokenizer tokenizer;
std::shared_ptr<T5Runner> t5;
size_t prompt_length = 256;
MiniT2IConditioner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr) {
bool use_t5 = false;
for (const auto& pair : tensor_storage_map) {
if (pair.first.find("text_encoders.t5xxl") != std::string::npos) {
use_t5 = true;
break;
}
}
if (!use_t5) {
LOG_WARN("IMPORTANT NOTICE: No MiniT2I T5 text encoder provided, cannot process prompts!");
return;
}
t5 = std::make_shared<T5Runner>(backend, tensor_storage_map, "text_encoders.t5xxl.transformer", false, weight_manager);
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
if (t5) {
t5->set_max_graph_vram_bytes(max_vram_bytes);
}
}
void set_stream_layers_enabled(bool enabled) override {
if (t5) {
t5->set_stream_layers_enabled(enabled);
}
}
void set_flash_attention_enabled(bool enabled) override {
if (t5) {
t5->set_flash_attention_enabled(enabled);
}
}
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
if (t5) {
t5->set_weight_adapter(adapter);
}
}
void runner_done() override {
if (t5) {
t5->runner_done();
}
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
SDCondition result;
if (!t5) {
result.c_crossattn = sd::Tensor<float>::zeros({1024, static_cast<int64_t>(prompt_length)});
result.c_vector = sd::Tensor<float>::zeros({static_cast<int64_t>(prompt_length)});
return result;
}
std::vector<int> tokens = tokenizer.encode(conditioner_params.text);
if (tokens.size() > prompt_length) {
tokens.resize(prompt_length);
}
std::vector<float> mask(tokens.size(), 1.0f);
while (tokens.size() < prompt_length) {
tokens.push_back(tokenizer.PAD_TOKEN_ID);
mask.push_back(0.0f);
}
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
std::vector<float> t5_mask(mask.size(), 0.0f);
for (size_t i = 0; i < mask.size(); ++i) {
t5_mask[i] = mask[i] > 0.0f ? 0.0f : -HUGE_VALF;
}
sd::Tensor<float> hidden_states = t5->compute(n_threads,
input_ids,
sd::Tensor<float>::from_vector(t5_mask),
false,
true,
true);
GGML_ASSERT(!hidden_states.empty());
result.c_crossattn = std::move(hidden_states);
result.c_vector = sd::Tensor<float>::from_vector(mask);
return result;
}
};
struct AnimaConditioner : public Conditioner {
std::shared_ptr<BPETokenizer> qwen_tokenizer;
T5UniGramTokenizer t5_tokenizer;
@ -1518,7 +1613,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 +2092,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

@ -76,29 +76,22 @@ static bool load_tensors_for_export(ModelLoader& model_loader,
return success;
}
bool convert(const char* input_path,
const char* vae_path,
const char* output_path,
sd_type_t output_type,
const char* tensor_type_rules,
bool convert_name) {
ModelLoader model_loader;
if (!model_loader.init_from_file(input_path)) {
LOG_ERROR("init model loader from file failed: '%s'", input_path);
static bool init_convert_path(ModelLoader& model_loader, const char* path, const char* prefix, bool& loaded_any) {
if (path == nullptr || strlen(path) == 0) {
return true;
}
if (!model_loader.init_from_file(path, prefix)) {
LOG_ERROR("init model loader from file failed: '%s'", path);
return false;
}
loaded_any = true;
return true;
}
if (vae_path != nullptr && strlen(vae_path) > 0) {
if (!model_loader.init_from_file(vae_path, "vae.")) {
LOG_ERROR("init model loader from file failed: '%s'", vae_path);
return false;
}
}
if (convert_name) {
model_loader.convert_tensors_name();
}
static bool export_loaded_model(ModelLoader& model_loader,
const char* output_path,
sd_type_t output_type,
const char* tensor_type_rules) {
ggml_type type = sd_type_to_ggml_type(output_type);
bool output_is_safetensors = ends_with(output_path, ".safetensors");
TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
@ -136,3 +129,55 @@ bool convert(const char* input_path,
ggml_free(ggml_ctx);
return success;
}
bool convert_with_components(const char* model_path,
const char* clip_l_path,
const char* clip_g_path,
const char* t5xxl_path,
const char* diffusion_model_path,
const char* vae_path,
const char* output_path,
sd_type_t output_type,
const char* tensor_type_rules,
bool convert_name) {
ModelLoader model_loader;
bool loaded_any = false;
if (!init_convert_path(model_loader, model_path, "", loaded_any) ||
!init_convert_path(model_loader, clip_l_path, "text_encoders.clip_l.transformer.", loaded_any) ||
!init_convert_path(model_loader, clip_g_path, "text_encoders.clip_g.transformer.", loaded_any) ||
!init_convert_path(model_loader, t5xxl_path, "text_encoders.t5xxl.transformer.", loaded_any) ||
!init_convert_path(model_loader, diffusion_model_path, "model.diffusion_model.", loaded_any) ||
!init_convert_path(model_loader, vae_path, "vae.", loaded_any)) {
return false;
}
if (!loaded_any) {
LOG_ERROR("no input model path provided for convert");
return false;
}
if (convert_name) {
model_loader.convert_tensors_name();
}
return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules);
}
bool convert(const char* input_path,
const char* vae_path,
const char* output_path,
sd_type_t output_type,
const char* tensor_type_rules,
bool convert_name) {
return convert_with_components(input_path,
nullptr,
nullptr,
nullptr,
nullptr,
vae_path,
output_path,
output_type,
tensor_type_rules,
convert_name);
}

View file

@ -391,7 +391,7 @@ __STATIC_INLINE__ uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, uint8_t*
int64_t width = input->ne[0];
int64_t height = input->ne[1];
int64_t channels = input->ne[2];
GGML_ASSERT(channels == 3 && input->type == GGML_TYPE_F32);
GGML_ASSERT(input->type == GGML_TYPE_F32);
if (image_data == nullptr) {
image_data = (uint8_t*)malloc(width * height * channels);
}
@ -1038,6 +1038,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_linear(ggml_context* ctx,
}
__STATIC_INLINE__ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
ggml_backend_t backend,
ggml_tensor* x,
int lp0,
int rp0,
@ -1063,7 +1064,17 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
}
if (lp0 != 0 || rp0 != 0 || lp1 != 0 || rp1 != 0 || lp2 != 0 || rp2 != 0 || lp3 != 0 || rp3 != 0) {
x = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
ggml_tensor* padded = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
if (backend == nullptr || ggml_backend_supports_op(backend, padded)) {
x = padded;
} else {
// Some backends (e.g. Metal) only implement right-padding for
// GGML_OP_PAD (see #850): pad right by lp+rp instead, then roll
// the padding around to the left. shift < ne always holds because
// ne grew by lp+rp.
x = ggml_pad_ext(ctx, x, 0, lp0 + rp0, 0, lp1 + rp1, 0, lp2 + rp2, 0, lp3 + rp3);
x = ggml_roll(ctx, x, lp0, lp1, lp2, lp3);
}
}
return x;
}
@ -1076,7 +1087,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_pad(ggml_context* ctx,
int p3 = 0,
bool circular_x = false,
bool circular_y = false) {
return ggml_ext_pad_ext(ctx, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y);
return ggml_ext_pad_ext(ctx, nullptr, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y);
}
// w: [OCIC, KH, KW]
@ -1105,7 +1116,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
}
if ((p0 != 0 || p1 != 0) && (circular_x || circular_y)) {
x = ggml_ext_pad_ext(ctx, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y);
x = ggml_ext_pad_ext(ctx, nullptr, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y);
p0 = 0;
p1 = 0;
}
@ -1130,6 +1141,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
// b: [OC,]
// result: [N*OC, OD, OH, OW]
__STATIC_INLINE__ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
ggml_backend_t backend,
ggml_tensor* x,
ggml_tensor* w,
ggml_tensor* b,
@ -1159,7 +1171,21 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 1, 3, 2));
x = ggml_reshape_4d(ctx, x, im2col->ne[1], im2col->ne[2], OD, OC * N);
} else {
x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2);
// ggml_conv_3d decomposes into GGML_OP_IM2COL_3D, which some backends
// (e.g. Metal, see #850) do not implement. Fall back to
// GGML_OP_CONV_3D on those backends.
bool im2col_3d_supported = true;
if (backend != nullptr) {
ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type);
im2col_3d_supported = ggml_backend_supports_op(backend, im2col);
}
if (im2col_3d_supported) {
x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2);
} else {
int64_t OC = w->ne[3] / IC;
int64_t N = x->ne[3] / IC;
x = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC);
}
}
if (b != nullptr) {
@ -1362,6 +1388,9 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
}
auto out = ggml_flash_attn_ext(ctx, q_in, k_in, v_in, mask_in, scale / kv_scale, 0, 0);
if (!ggml_backend_supports_op(backend, out)) {
return nullptr;
}
ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32);
if (kv_scale != 1.0f) {
out = ggml_ext_scale(ctx, out, 1.0f / kv_scale);
@ -1379,9 +1408,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
if (can_use_flash_attn) {
kqv = build_kqv(q, k, v, mask);
if (!ggml_backend_supports_op(backend, kqv)) {
kqv = nullptr;
} else {
if (kqv != nullptr) {
kqv = ggml_view_4d(ctx,
kqv,
d_head,
@ -2473,7 +2500,10 @@ protected:
sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
}
ggml_status status = ggml_backend_graph_compute(runtime_backend, gf);
ggml_status status = sd_backend_graph_compute_with_eval_callback(runtime_backend,
gf,
sd_get_backend_eval_callback(),
sd_get_backend_eval_callback_data());
if (status != GGML_STATUS_SUCCESS) {
LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
return std::nullopt;
@ -3365,7 +3395,7 @@ public:
b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
}
}
return ggml_ext_conv_3d(ctx->ggml_ctx, x, w, b, in_channels,
return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
std::get<2>(padding), std::get<1>(padding), std::get<0>(padding),
std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),

View file

@ -9,6 +9,7 @@
#include <vector>
#include "core/util.h"
#include "ggml/src/ggml-impl.h"
#include "stable-diffusion.h"
static std::string trim_copy(const std::string& value) {
@ -110,7 +111,67 @@ static std::string resolve_first_device_by_type(enum ggml_backend_dev_type type)
if (dev == nullptr) {
return "";
}
return ggml_backend_dev_name(dev);
const char* dev_name = ggml_backend_dev_name(dev);
if (dev_name != nullptr && dev_name[0] != '\0') {
return dev_name;
}
ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
const char* reg_name = reg != nullptr ? ggml_backend_reg_name(reg) : nullptr;
return reg_name != nullptr ? reg_name : "";
}
static ggml_backend_dev_t resolve_first_device_by_registry_name(const std::string& name) {
std::string lower = lower_copy(trim_copy(name));
if (lower == "metal") {
lower = "mtl";
}
if (lower.empty()) {
return nullptr;
}
const size_t device_count = ggml_backend_dev_count();
for (size_t i = 0; i < device_count; ++i) {
ggml_backend_dev_t dev = ggml_backend_dev_get(i);
ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
if (reg == nullptr) {
continue;
}
const char* reg_name = ggml_backend_reg_name(reg);
if (reg_name != nullptr && lower_copy(reg_name) == lower) {
return dev;
}
}
return nullptr;
}
static ggml_backend_dev_t resolve_device_by_name(const std::string& name) {
const std::string lower = lower_copy(trim_copy(name));
if (lower.empty()) {
return nullptr;
}
const size_t device_count = ggml_backend_dev_count();
for (size_t i = 0; i < device_count; ++i) {
ggml_backend_dev_t dev = ggml_backend_dev_get(i);
const char* dev_name = ggml_backend_dev_name(dev);
if (dev_name != nullptr && lower_copy(dev_name) == lower) {
return dev;
}
}
return nullptr;
}
static std::string backend_device_name(ggml_backend_dev_t dev) {
if (dev == nullptr) {
return "";
}
const char* name = ggml_backend_dev_name(dev);
if (name != nullptr && name[0] != '\0') {
return name;
}
ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
const char* reg_name = reg != nullptr ? ggml_backend_reg_name(reg) : nullptr;
return reg_name != nullptr ? reg_name : "";
}
static ggml_backend_buffer_t ggml_backend_tensor_buffer(const struct ggml_tensor* tensor) {
@ -296,6 +357,10 @@ std::string sd_backend_resolve_name(const std::string& name) {
return resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_IGPU);
}
if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(requested)) {
return backend_device_name(dev);
}
const size_t device_count = ggml_backend_dev_count();
for (size_t i = 0; i < device_count; ++i) {
ggml_backend_dev_t dev = ggml_backend_dev_get(i);
@ -328,7 +393,20 @@ static ggml_backend_t init_named_backend(const std::string& name) {
return ggml_backend_init_best();
}
if (ggml_backend_dev_t dev = resolve_device_by_name(name)) {
return ggml_backend_dev_init(dev, nullptr);
}
if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(name)) {
return ggml_backend_dev_init(dev, nullptr);
}
std::string resolved = sd_backend_resolve_name(name);
if (ggml_backend_dev_t dev = resolve_device_by_name(resolved)) {
return ggml_backend_dev_init(dev, nullptr);
}
if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(resolved)) {
return ggml_backend_dev_init(dev, nullptr);
}
if (resolved.empty()) {
return nullptr;
}
@ -364,6 +442,68 @@ bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) {
return false;
}
static ggml_cgraph sd_ggml_graph_view(ggml_cgraph* cgraph0, int i0, int i1) {
ggml_cgraph cgraph = {
/*.size =*/0,
/*.n_nodes =*/i1 - i0,
/*.n_leafs =*/0,
/*.nodes =*/cgraph0->nodes + i0,
/*.grads =*/nullptr,
/*.grad_accs =*/nullptr,
/*.leafs =*/nullptr,
/*.use_counts =*/cgraph0->use_counts,
/*.visited_hash_set =*/cgraph0->visited_hash_set,
/*.order =*/cgraph0->order,
/*.uid =*/0,
};
return cgraph;
}
ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
ggml_cgraph* gf,
sd_graph_eval_callback_t callback_eval,
void* callback_eval_user_data) {
if (callback_eval == nullptr) {
return ggml_backend_graph_compute(backend, gf);
}
ggml_status status = GGML_STATUS_SUCCESS;
const int n_nodes = ggml_graph_n_nodes(gf);
bool stopped = false;
for (int j0 = 0; j0 < n_nodes; ++j0) {
ggml_tensor* t = ggml_graph_node(gf, j0);
bool need = callback_eval(t, true, callback_eval_user_data);
int j1 = j0;
while (!need && j1 < n_nodes - 1) {
t = ggml_graph_node(gf, ++j1);
need = callback_eval(t, true, callback_eval_user_data);
}
ggml_cgraph gv = sd_ggml_graph_view(gf, j0, j1 + 1);
status = ggml_backend_graph_compute_async(backend, &gv);
if (status != GGML_STATUS_SUCCESS) {
break;
}
ggml_backend_synchronize(backend);
if (need && !callback_eval(t, false, callback_eval_user_data)) {
stopped = true;
break;
}
j0 = j1;
}
ggml_backend_synchronize(backend);
if (stopped && status == GGML_STATUS_SUCCESS) {
status = GGML_STATUS_ABORTED;
}
return status;
}
const char* sd_get_system_info() {
static std::string cache_info = []() -> std::string {
ggml_backend_load_all_once();
@ -599,7 +739,7 @@ bool SDBackendManager::validate(std::string* error) const {
}
return false;
}
if (!sd_backend_resolve_name(name).empty()) {
if (!sd_backend_resolve_name(name).empty() || resolve_first_device_by_registry_name(name) != nullptr) {
return true;
}
if (error != nullptr) {

View file

@ -9,6 +9,7 @@
#include "ggml-backend.h"
#include "ggml.h"
#include "stable-diffusion.h"
enum class SDBackendModule {
DIFFUSION,
@ -71,6 +72,10 @@ bool sd_backend_is(ggml_backend_t backend, const std::string& name);
bool sd_backend_is_cpu(ggml_backend_t backend);
ggml_backend_t sd_backend_cpu_init();
bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
ggml_cgraph* gf,
sd_graph_eval_callback_t callback_eval,
void* callback_eval_user_data);
std::string sd_backend_resolve_name(const std::string& name);
const char* sd_backend_module_name(SDBackendModule module);
void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);

View file

@ -360,6 +360,9 @@ int sd_preview_interval = 1;
bool sd_preview_denoised = true;
bool sd_preview_noisy = false;
static sd_graph_eval_callback_t sd_backend_eval_cb = nullptr;
static void* sd_backend_eval_cb_data = nullptr;
std::u32string utf8_to_utf32(const std::string& utf8_str) {
std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
return converter.from_bytes(utf8_str);
@ -662,6 +665,11 @@ void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode, int interval, b
sd_preview_noisy = noisy;
}
void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data) {
sd_backend_eval_cb = cb;
sd_backend_eval_cb_data = data;
}
sd_preview_cb_t sd_get_preview_callback() {
return sd_preview_cb;
}
@ -682,6 +690,14 @@ bool sd_should_preview_noisy() {
return sd_preview_noisy;
}
sd_graph_eval_callback_t sd_get_backend_eval_callback() {
return sd_backend_eval_cb;
}
void* sd_get_backend_eval_callback_data() {
return sd_backend_eval_cb_data;
}
sd_progress_cb_t sd_get_progress_callback() {
return sd_progress_cb;
}

View file

@ -101,6 +101,9 @@ int sd_get_preview_interval();
bool sd_should_preview_denoised();
bool sd_should_preview_noisy();
sd_graph_eval_callback_t sd_get_backend_eval_callback();
void* sd_get_backend_eval_callback_data();
// test if the backend is a specific one, e.g. "CUDA", "ROCm", "Vulkan" etc.
bool sd_backend_is(ggml_backend_t backend, const std::string& name);

View file

@ -36,6 +36,7 @@ enum SDVersion {
VERSION_WAN2_2_I2V,
VERSION_WAN2_2_TI2V,
VERSION_QWEN_IMAGE,
VERSION_QWEN_IMAGE_LAYERED,
VERSION_ANIMA,
VERSION_FLUX2,
VERSION_FLUX2_KLEIN,
@ -46,9 +47,11 @@ enum SDVersion {
VERSION_OVIS_IMAGE,
VERSION_ERNIE_IMAGE,
VERSION_LENS,
VERSION_MINIT2I,
VERSION_LONGCAT,
VERSION_PID,
VERSION_IDEOGRAM4,
VERSION_SEFI_IMAGE,
VERSION_KREA2,
VERSION_ESRGAN,
VERSION_COUNT,
@ -125,7 +128,7 @@ static inline bool sd_version_is_wan(SDVersion version) {
}
static inline bool sd_version_is_qwen_image(SDVersion version) {
if (version == VERSION_QWEN_IMAGE) {
if (version == VERSION_QWEN_IMAGE || version == VERSION_QWEN_IMAGE_LAYERED) {
return true;
}
return false;
@ -173,6 +176,13 @@ static inline bool sd_version_is_lens(SDVersion version) {
return false;
}
static inline bool sd_version_is_minit2i(SDVersion version) {
if (version == VERSION_MINIT2I) {
return true;
}
return false;
}
static inline bool sd_version_is_pid(SDVersion version) {
if (version == VERSION_PID) {
return true;
@ -187,6 +197,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 +219,14 @@ 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;
}
static inline bool sd_version_uses_wan_vae(SDVersion version) {
if (sd_version_is_wan(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) {
return true;
}
return false;
@ -232,9 +256,11 @@ static inline bool sd_version_is_dit(SDVersion version) {
sd_version_is_boogu_image(version) ||
sd_version_is_ernie_image(version) ||
sd_version_is_lens(version) ||
sd_version_is_minit2i(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

@ -12,6 +12,16 @@ namespace Rope {
ErnieImage,
};
enum class RefIndexMode {
FIXED,
INCREASE,
DECREASE,
};
__STATIC_INLINE__ RefIndexMode ref_index_mode_from_bool(bool increase_ref_index) {
return increase_ref_index ? RefIndexMode::INCREASE : RefIndexMode::FIXED;
}
template <class T>
__STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
std::vector<T> result(num);
@ -346,7 +356,7 @@ namespace Rope {
int axes_dim_num,
int start_index,
const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index,
RefIndexMode ref_index_mode,
float ref_index_scale,
bool scale_rope,
int base_offset = 0) {
@ -357,13 +367,15 @@ namespace Rope {
for (ggml_tensor* ref : ref_latents) {
int h_offset = 0;
int w_offset = 0;
if (!increase_ref_index) {
if (ref_index_mode == RefIndexMode::FIXED) {
if (ref->ne[1] + curr_h_offset > ref->ne[0] + curr_w_offset) {
w_offset = curr_w_offset;
} else {
h_offset = curr_h_offset;
}
scale_rope = false;
} else if (ref_index_mode == RefIndexMode::DECREASE) {
index--;
}
auto ref_ids = gen_flux_img_ids(static_cast<int>(ref->ne[1]),
@ -377,7 +389,7 @@ namespace Rope {
scale_rope);
ids = concat_ids(ids, ref_ids, bs);
if (increase_ref_index) {
if (ref_index_mode == RefIndexMode::INCREASE) {
index++;
}
@ -395,7 +407,7 @@ namespace Rope {
int context_len,
std::set<int> txt_arange_dims,
const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index,
RefIndexMode ref_index_mode,
float ref_index_scale,
bool is_longcat) {
int x_index = is_longcat ? 1 : 0;
@ -406,7 +418,7 @@ namespace Rope {
auto ids = concat_ids(txt_ids, img_ids, bs);
if (ref_latents.size() > 0) {
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, increase_ref_index, ref_index_scale, false, offset);
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset);
ids = concat_ids(ids, refs_ids, bs);
}
return ids;
@ -420,7 +432,7 @@ namespace Rope {
int context_len,
std::set<int> txt_arange_dims,
const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index,
RefIndexMode ref_index_mode,
float ref_index_scale,
int theta,
bool circular_h,
@ -435,7 +447,7 @@ namespace Rope {
context_len,
txt_arange_dims,
ref_latents,
increase_ref_index,
ref_index_mode,
ref_index_scale,
is_longcat);
std::vector<std::vector<int>> wrap_dims;
@ -481,17 +493,64 @@ namespace Rope {
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int h,
__STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
int h,
int w,
int pt,
int ph,
int pw,
int bs,
int t_offset = 0,
int h_offset = 0,
int w_offset = 0,
bool scale_rope = false) {
int t_len = (t + (pt / 2)) / pt;
int h_len = (h + (ph / 2)) / ph;
int w_len = (w + (pw / 2)) / pw;
std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
if (scale_rope) {
h_offset -= h_len / 2;
w_offset -= w_len / 2;
}
std::vector<float> t_ids = linspace<float>(1.f * t_offset, 1.f * t_len - 1 + t_offset, t_len);
std::vector<float> h_ids = linspace<float>(1.f * h_offset, 1.f * h_len - 1 + h_offset, h_len);
std::vector<float> w_ids = linspace<float>(1.f * w_offset, 1.f * w_len - 1 + w_offset, w_len);
for (int i = 0; i < t_len; ++i) {
for (int j = 0; j < h_len; ++j) {
for (int k = 0; k < w_len; ++k) {
int idx = i * h_len * w_len + j * w_len + k;
vid_ids[idx][0] = t_ids[i];
vid_ids[idx][1] = h_ids[j];
vid_ids[idx][2] = w_ids[k];
}
}
}
std::vector<std::vector<float>> vid_ids_repeated(bs * vid_ids.size(), std::vector<float>(3));
for (int i = 0; i < bs; ++i) {
for (int j = 0; j < vid_ids.size(); ++j) {
vid_ids_repeated[i * vid_ids.size() + j] = vid_ids[j];
}
}
return vid_ids_repeated;
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int t,
int h,
int w,
int patch_size,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index) {
RefIndexMode ref_index_mode) {
int h_len = (h + (patch_size / 2)) / patch_size;
int w_len = (w + (patch_size / 2)) / patch_size;
int txt_id_start = std::max(h_len, w_len);
auto txt_ids = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
int txt_id_start = std::max(h_len, w_len) / 2;
auto txt_ids = linspace<float>(1.f * txt_id_start, 1.f * txt_id_start + context_len - 1, context_len);
std::vector<std::vector<float>> txt_ids_repeated(bs * context_len, std::vector<float>(3));
for (int i = 0; i < bs; ++i) {
for (int j = 0; j < txt_ids.size(); ++j) {
@ -499,28 +558,30 @@ namespace Rope {
}
}
int axes_dim_num = 3;
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, 0, 0, 0, true);
auto img_ids = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true);
auto ids = concat_ids(txt_ids_repeated, img_ids, bs);
if (ref_latents.size() > 0) {
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, 1, ref_latents, increase_ref_index, 1.f, true);
ids = concat_ids(ids, refs_ids, bs);
int ref_start_index = ref_index_mode == RefIndexMode::DECREASE ? 0 : 1;
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true);
ids = concat_ids(ids, refs_ids, bs);
}
return ids;
}
// Generate qwen_image positional embeddings
__STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int h,
__STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int t,
int h,
int w,
int patch_size,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index,
RefIndexMode ref_index_mode,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_qwen_image_ids(h, w, patch_size, bs, context_len, ref_latents, increase_ref_index);
std::vector<std::vector<float>> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode);
std::vector<std::vector<int>> wrap_dims;
// This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
@ -533,7 +594,7 @@ namespace Rope {
// Track per-token wrap lengths for the row/column axes so only spatial tokens become periodic.
wrap_dims.assign(axes_dim.size(), std::vector<int>(total_tokens / bs, 0));
size_t cursor = context_len; // ignore text tokens
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
const size_t img_tokens = static_cast<size_t>(t) * static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len;
@ -684,46 +745,6 @@ namespace Rope {
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage);
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
int h,
int w,
int pt,
int ph,
int pw,
int bs,
int t_offset = 0,
int h_offset = 0,
int w_offset = 0) {
int t_len = (t + (pt / 2)) / pt;
int h_len = (h + (ph / 2)) / ph;
int w_len = (w + (pw / 2)) / pw;
std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
std::vector<float> t_ids = linspace<float>(1.f * t_offset, 1.f * t_len - 1 + t_offset, t_len);
std::vector<float> h_ids = linspace<float>(1.f * h_offset, 1.f * h_len - 1 + h_offset, h_len);
std::vector<float> w_ids = linspace<float>(1.f * w_offset, 1.f * w_len - 1 + w_offset, w_len);
for (int i = 0; i < t_len; ++i) {
for (int j = 0; j < h_len; ++j) {
for (int k = 0; k < w_len; ++k) {
int idx = i * h_len * w_len + j * w_len + k;
vid_ids[idx][0] = t_ids[i];
vid_ids[idx][1] = h_ids[j];
vid_ids[idx][2] = w_ids[k];
}
}
}
std::vector<std::vector<float>> vid_ids_repeated(bs * vid_ids.size(), std::vector<float>(3));
for (int i = 0; i < bs; ++i) {
for (int j = 0; j < vid_ids.size(); ++j) {
vid_ids_repeated[i * vid_ids.size() + j] = vid_ids[j];
}
}
return vid_ids_repeated;
}
// Generate wan positional embeddings
__STATIC_INLINE__ std::vector<float> gen_wan_pe(int t,
int h,
@ -785,7 +806,8 @@ namespace Rope {
int context_len,
int seq_multi_of,
const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index) {
RefIndexMode ref_index_mode) {
SD_UNUSED(ref_index_mode);
int padded_context_len = context_len + bound_mod(context_len, seq_multi_of);
auto txt_ids = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
for (int i = 0; i < bs * padded_context_len; i++) {
@ -816,12 +838,12 @@ namespace Rope {
int context_len,
int seq_multi_of,
const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index,
RefIndexMode ref_index_mode,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, increase_ref_index);
std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode);
std::vector<std::vector<int>> wrap_dims;
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
int pad_h = (patch_size - (h % patch_size)) % patch_size;

View file

@ -227,7 +227,6 @@ namespace Anima {
k4 = k_norm->forward(ctx, k4);
ggml_tensor* attn_out = nullptr;
float scale = (sd_backend_is(ctx->backend, "Vulkan") && ctx->flash_attn_enabled) ? 1.0f / 32.0f : 1.0f;
if (pe_q != nullptr || pe_k != nullptr) {
if (pe_q == nullptr) {
pe_q = pe_k;
@ -245,8 +244,7 @@ namespace Anima {
num_heads,
nullptr,
true,
ctx->flash_attn_enabled,
scale);
ctx->flash_attn_enabled);
} else {
auto q_flat = ggml_reshape_3d(ctx->ggml_ctx, q4, head_dim * num_heads, L_q, N);
auto k_flat = ggml_reshape_3d(ctx->ggml_ctx, k4, head_dim * num_heads, L_k, N);
@ -258,8 +256,7 @@ namespace Anima {
num_heads,
nullptr,
false,
ctx->flash_attn_enabled,
scale);
ctx->flash_attn_enabled);
}
return out_proj->forward(ctx, attn_out);
@ -615,7 +612,7 @@ namespace Anima {
0,
{},
empty_ref_latents,
false,
Rope::RefIndexMode::FIXED,
1.0f,
false);

View file

@ -1,4 +1,4 @@
#ifndef __SD_MODEL_DIFFUSION_CONTROL_HPP__
#ifndef __SD_MODEL_DIFFUSION_CONTROL_HPP__
#define __SD_MODEL_DIFFUSION_CONTROL_HPP__
#include "model/common/block.hpp"

View file

@ -135,23 +135,23 @@ namespace DiT {
return x;
}
inline ggml_tensor* unpatchify(ggml_context* ctx,
ggml_tensor* x,
int64_t t_len,
int64_t h_len,
int64_t w_len,
int pt,
int ph,
int pw) {
// x: [N, t_len*h_len*w_len, pt*ph*pw*C]
inline ggml_tensor* unpatchify_3d(ggml_context* ctx,
ggml_tensor* x,
int64_t t_len,
int64_t h_len,
int64_t w_len,
int pt,
int ph,
int pw) {
// x: [N, t_len*h_len*w_len, C*pt*ph*pw]
// return: [N*C, t_len*pt, h_len*ph, w_len*pw]
int64_t N = x->ne[3];
int64_t N = x->ne[2];
int64_t C = x->ne[0] / pt / ph / pw;
GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, C, pt*ph*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw]
x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw]

View file

@ -162,8 +162,6 @@ namespace ErnieImage {
int64_t S = x->ne[1];
int64_t N = x->ne[2];
float scale = (sd_backend_is(ctx->backend, "Vulkan") && ctx->flash_attn_enabled) ? 1.0f / 32.0f : 1.0f;
auto q = to_q->forward(ctx, x);
auto k = to_k->forward(ctx, x);
auto v = to_v->forward(ctx, x);
@ -184,7 +182,7 @@ namespace ErnieImage {
k = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, k, 0, 2, 1, 3)); // [N, heads, S, head_dim]
k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled, scale); // [N, S, hidden_size]
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled); // [N, S, hidden_size]
x = to_out_0->forward(ctx, x);
return x;
}

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));
@ -1459,7 +1485,7 @@ namespace Flux {
const sd::Tensor<float>& y_tensor = {},
const sd::Tensor<float>& guidance_tensor = {},
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
bool increase_ref_index = false,
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED,
std::vector<int> skip_layers = {},
const sd::Tensor<float>& pulid_id_tensor = {},
float pulid_id_weight = 1.0f) {
@ -1500,9 +1526,9 @@ 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)) {
txt_arange_dims = {3};
increase_ref_index = true;
if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
txt_arange_dims = {3};
ref_index_mode = Rope::RefIndexMode::INCREASE;
} else if (version == VERSION_OVIS_IMAGE) {
txt_arange_dims = {1, 2};
}
@ -1513,7 +1539,7 @@ namespace Flux {
static_cast<int>(context->ne[1]),
txt_arange_dims,
ref_latents,
increase_ref_index,
ref_index_mode,
config.ref_index_scale,
config.theta,
circular_y_enabled,
@ -1573,7 +1599,7 @@ namespace Flux {
const sd::Tensor<float>& y = {},
const sd::Tensor<float>& guidance = {},
const std::vector<sd::Tensor<float>>& ref_latents = {},
bool increase_ref_index = false,
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED,
std::vector<int> skip_layers = std::vector<int>(),
const sd::Tensor<float>& pulid_id = {},
float pulid_id_weight = 1.0f) {
@ -1584,7 +1610,7 @@ namespace Flux {
// guidance: [N, ]
// pulid_id: empty (no injection) or [N, num_id_tokens=32, kv_dim=2048]
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, skip_layers, pulid_id, pulid_id_weight);
return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, ref_index_mode, skip_layers, pulid_id, pulid_id_weight);
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@ -1606,7 +1632,7 @@ namespace Flux {
tensor_or_empty(diffusion_params.y),
tensor_or_empty(extra->guidance),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.increase_ref_index,
diffusion_params.ref_index_mode,
extra->skip_layers ? *extra->skip_layers : empty_skip_layers,
tensor_or_empty(extra->pulid_id),
extra->pulid_id_weight);
@ -1657,7 +1683,7 @@ namespace Flux {
{},
guidance,
{},
false);
Rope::RefIndexMode::FIXED);
int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty());

View file

@ -1,4 +1,4 @@
#ifndef __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
#ifndef __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
#define __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
#include <algorithm>

View file

@ -0,0 +1,611 @@
#ifndef __SD_MODEL_DIFFUSION_MINIT2I_HPP__
#define __SD_MODEL_DIFFUSION_MINIT2I_HPP__
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <memory>
#include <string>
#include <vector>
#include "core/ggml_extend.hpp"
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
#include "model/diffusion/model.hpp"
#include "model_loader.h"
namespace MiniT2I {
constexpr int MINIT2I_GRAPH_SIZE = 196608;
struct MiniT2IConfig {
int64_t image_size = 512;
int64_t patch_size = 16;
int64_t in_channels = 3;
int64_t txt_input_size = 1024;
int64_t hidden_size = 768;
int64_t txt_hidden_size = 768;
int64_t cond_vec_size = 768;
int64_t depth_double = 17;
int64_t txt_preamble_depth = 2;
int64_t num_heads = 12;
int64_t head_dim = 64;
float mlp_ratio = 2.6667f;
int64_t pca_channels = 128;
int64_t prompt_length = 256;
int64_t n_T = 100;
float cfg_interval_start = 0.0f;
float cfg_interval_end = 1.0f;
static MiniT2IConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
MiniT2IConfig config;
config.depth_double = 0;
config.txt_preamble_depth = 0;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (ends_with(name, "img_embedder.proj1.weight") && tensor_storage.n_dims == 4) {
config.patch_size = tensor_storage.ne[0];
config.in_channels = tensor_storage.ne[2];
config.pca_channels = tensor_storage.ne[3];
} else if (ends_with(name, "img_embedder.proj2.weight") && tensor_storage.n_dims == 4) {
config.pca_channels = tensor_storage.ne[2];
config.hidden_size = tensor_storage.ne[3];
} else if (ends_with(name, "txt_embedder.weight") && tensor_storage.n_dims == 2) {
config.txt_input_size = tensor_storage.ne[0];
config.txt_hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "pooled_embedder.weight") && tensor_storage.n_dims == 2) {
config.cond_vec_size = tensor_storage.ne[1];
} else if (ends_with(name, "double_blocks.0.img_qkv.weight") && tensor_storage.n_dims == 2) {
int64_t inner3 = tensor_storage.ne[1];
int64_t inner = inner3 / 3;
config.hidden_size = tensor_storage.ne[0];
if (config.hidden_size == 768) {
config.num_heads = 12;
config.head_dim = 64;
} else if (config.hidden_size == 1248) {
config.num_heads = 24;
config.head_dim = 52;
} else if (inner > 0) {
config.head_dim = 64;
config.num_heads = std::max<int64_t>(1, inner / config.head_dim);
}
} else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) {
int64_t patch_area = config.patch_size * config.patch_size;
config.hidden_size = tensor_storage.ne[0];
config.in_channels = patch_area > 0 ? tensor_storage.ne[1] / patch_area : config.in_channels;
} else if (ends_with(name, "mask_token") && tensor_storage.n_dims >= 2) {
config.prompt_length = tensor_storage.ne[1];
}
size_t pos = name.find("double_blocks.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int64_t idx = atoi(items[1].c_str());
config.depth_double = std::max<int64_t>(config.depth_double, idx + 1);
}
}
pos = name.find("txt_preamble_blocks.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int64_t idx = atoi(items[1].c_str());
config.txt_preamble_depth = std::max<int64_t>(config.txt_preamble_depth, idx + 1);
}
}
}
if (config.depth_double <= 0) {
config.depth_double = config.hidden_size == 1248 ? 23 : 17;
}
if (config.txt_preamble_depth <= 0) {
config.txt_preamble_depth = 2;
}
if (config.head_dim <= 0 || config.num_heads <= 0) {
config.head_dim = config.hidden_size == 1248 ? 52 : 64;
config.num_heads = config.hidden_size / config.head_dim;
}
LOG_DEBUG("minit2i: hidden_size=%" PRId64 ", txt_hidden_size=%" PRId64 ", heads=%" PRId64 ", head_dim=%" PRId64 ", double_blocks=%" PRId64 ", txt_blocks=%" PRId64 ", patch=%" PRId64 ", in_channels=%" PRId64,
config.hidden_size,
config.txt_hidden_size,
config.num_heads,
config.head_dim,
config.depth_double,
config.txt_preamble_depth,
config.patch_size,
config.in_channels);
return config;
}
};
inline std::vector<float> make_2d_sincos_pos_embed(int grid_size, int dim) {
GGML_ASSERT(dim % 4 == 0);
int half_dim = dim / 2;
int quarter = half_dim / 2;
std::vector<float> out(static_cast<size_t>(grid_size) * grid_size * dim);
std::vector<float> omega(quarter);
for (int i = 0; i < quarter; ++i) {
omega[i] = 1.0f / std::pow(10000.0f, static_cast<float>(i) / static_cast<float>(quarter));
}
for (int y = 0; y < grid_size; ++y) {
for (int x = 0; x < grid_size; ++x) {
size_t base = static_cast<size_t>(y * grid_size + x) * dim;
for (int i = 0; i < quarter; ++i) {
float ay = y * omega[i];
float ax = x * omega[i];
out[base + i] = std::sin(ax);
out[base + quarter + i] = std::cos(ax);
out[base + half_dim + i] = std::sin(ay);
out[base + half_dim + quarter + i] = std::cos(ay);
}
}
}
return out;
}
inline std::vector<float> make_text_rope(int length, int head_dim) {
return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), head_dim, 10000.f));
}
inline std::vector<float> make_vision_rope(int side, int head_dim) {
GGML_ASSERT(head_dim % 4 == 0);
int dim = head_dim / 2;
int quarter = dim / 2;
int length = side * side;
std::vector<float> out(static_cast<size_t>(length) * (head_dim / 2) * 4);
std::vector<float> freqs(quarter);
for (int i = 0; i < quarter; ++i) {
freqs[i] = 1.0f / std::pow(10000.0f, static_cast<float>(2 * i) / static_cast<float>(dim));
}
for (int y = 0; y < side; ++y) {
for (int x = 0; x < side; ++x) {
int pos = y * side + x;
size_t base = static_cast<size_t>(pos) * (head_dim / 2) * 4;
for (int i = 0; i < quarter; ++i) {
float ay = y * freqs[i];
float ax = x * freqs[i];
float angles[2] = {ay, ax};
for (int axis = 0; axis < 2; ++axis) {
int j = axis * quarter + i;
out[base + 4 * j] = std::cos(angles[axis]);
out[base + 4 * j + 1] = -std::sin(angles[axis]);
out[base + 4 * j + 2] = std::sin(angles[axis]);
out[base + 4 * j + 3] = std::cos(angles[axis]);
}
}
}
}
return out;
}
struct SwiGLUMlp : public GGMLBlock {
SwiGLUMlp(int64_t in_features, int64_t hidden_features) {
int64_t hidden_dim = ((hidden_features + 7) / 8) * 8;
blocks["w1"] = std::make_shared<Linear>(in_features, hidden_dim, false);
blocks["w3"] = std::make_shared<Linear>(in_features, hidden_dim, false);
blocks["w2"] = std::make_shared<Linear>(hidden_dim, in_features, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
auto gate = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x));
auto up = w3->forward(ctx, x);
return w2->forward(ctx, ggml_mul(ctx->ggml_ctx, gate, up));
}
};
struct BottleneckPatchEmbed : public GGMLBlock {
int64_t patch_size;
BottleneckPatchEmbed(int64_t patch_size, int64_t in_channels, int64_t pca_channels, int64_t hidden_size)
: patch_size(patch_size) {
blocks["proj1"] = std::make_shared<Conv2d>(in_channels,
pca_channels,
std::pair<int, int>{static_cast<int>(patch_size), static_cast<int>(patch_size)},
std::pair<int, int>{static_cast<int>(patch_size), static_cast<int>(patch_size)},
std::pair<int, int>{0, 0},
std::pair<int, int>{1, 1},
false);
blocks["proj2"] = std::make_shared<Conv2d>(pca_channels,
hidden_size,
std::pair<int, int>{1, 1},
std::pair<int, int>{1, 1},
std::pair<int, int>{0, 0},
std::pair<int, int>{1, 1},
true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto proj1 = std::dynamic_pointer_cast<Conv2d>(blocks["proj1"]);
auto proj2 = std::dynamic_pointer_cast<Conv2d>(blocks["proj2"]);
x = proj1->forward(ctx, x);
x = proj2->forward(ctx, x);
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
return x;
}
};
struct TimestepEmbedder : public GGMLBlock {
int frequency_embedding_size;
TimestepEmbedder(int64_t hidden_size, int frequency_embedding_size = 256)
: frequency_embedding_size(frequency_embedding_size) {
blocks["mlp.0"] = std::make_shared<Linear>(frequency_embedding_size, hidden_size, true, true);
blocks["mlp.2"] = std::make_shared<Linear>(hidden_size, hidden_size, true, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10000, 1.0f);
t_emb = mlp_0->forward(ctx, t_emb);
t_emb = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
return mlp_2->forward(ctx, t_emb);
}
};
inline std::vector<ggml_tensor*> split_qkv(ggml_context* ctx, ggml_tensor* qkv, int64_t num_heads, int64_t head_dim) {
int64_t N = qkv->ne[2];
int64_t L = qkv->ne[1];
auto q = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N,
qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], 0);
auto k = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N,
qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], qkv->nb[0] * head_dim * num_heads);
auto v = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N,
qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], qkv->nb[0] * head_dim * num_heads * 2);
return {q, k, v};
}
struct PlainTextTransformerBlock : public GGMLBlock {
int64_t num_heads;
int64_t head_dim;
PlainTextTransformerBlock(int64_t hidden_size, int64_t num_heads, int64_t head_dim, float mlp_ratio)
: num_heads(num_heads), head_dim(head_dim) {
int64_t inner_dim = num_heads * head_dim;
blocks["norm1"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["norm2"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["qkv"] = std::make_shared<Linear>(hidden_size, inner_dim * 3, true);
blocks["attn_proj"] = std::make_shared<Linear>(inner_dim, hidden_size, true);
blocks["mlp"] = std::make_shared<SwiGLUMlp>(hidden_size, static_cast<int64_t>(hidden_size * mlp_ratio));
blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* pe) {
auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto attn_proj = std::dynamic_pointer_cast<Linear>(blocks["attn_proj"]);
auto mlp = std::dynamic_pointer_cast<SwiGLUMlp>(blocks["mlp"]);
auto q_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
auto k_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
auto qkv = split_qkv(ctx->ggml_ctx, qkv_proj->forward(ctx, norm1->forward(ctx, txt)), num_heads, head_dim);
auto q = q_norm->forward(ctx, qkv[0]);
auto k = k_norm->forward(ctx, qkv[1]);
auto v = qkv[2];
auto out = Rope::attention(ctx, q, k, v, pe, nullptr, 1.0f, false);
txt = ggml_add(ctx->ggml_ctx, txt, attn_proj->forward(ctx, out));
txt = ggml_add(ctx->ggml_ctx, txt, mlp->forward(ctx, norm2->forward(ctx, txt)));
return txt;
}
};
struct DoubleStreamDiTBlock : public GGMLBlock {
int64_t num_heads;
int64_t head_dim;
DoubleStreamDiTBlock(int64_t hidden_size, int64_t txt_hidden_size, int64_t num_heads, int64_t head_dim, float mlp_ratio)
: num_heads(num_heads), head_dim(head_dim) {
int64_t inner_dim = num_heads * head_dim;
blocks["img_norm1"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["img_norm2"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["txt_norm1"] = std::make_shared<RMSNorm>(txt_hidden_size, 1e-6f);
blocks["txt_norm2"] = std::make_shared<RMSNorm>(txt_hidden_size, 1e-6f);
blocks["img_qkv"] = std::make_shared<Linear>(hidden_size, inner_dim * 3, true);
blocks["txt_qkv"] = std::make_shared<Linear>(txt_hidden_size, inner_dim * 3, true);
blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["img_attn_proj"] = std::make_shared<Linear>(inner_dim, hidden_size, true);
blocks["txt_attn_proj"] = std::make_shared<Linear>(inner_dim, txt_hidden_size, true);
blocks["img_mlp"] = std::make_shared<SwiGLUMlp>(hidden_size, static_cast<int64_t>(hidden_size * mlp_ratio));
blocks["txt_mlp"] = std::make_shared<SwiGLUMlp>(txt_hidden_size, static_cast<int64_t>(txt_hidden_size * mlp_ratio));
}
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* img,
ggml_tensor* txt,
ggml_tensor* pe) {
auto img_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_norm1"]);
auto img_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_norm2"]);
auto txt_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm1"]);
auto txt_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm2"]);
auto img_qkv_p = std::dynamic_pointer_cast<Linear>(blocks["img_qkv"]);
auto txt_qkv_p = std::dynamic_pointer_cast<Linear>(blocks["txt_qkv"]);
auto q_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
auto k_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
auto img_proj = std::dynamic_pointer_cast<Linear>(blocks["img_attn_proj"]);
auto txt_proj = std::dynamic_pointer_cast<Linear>(blocks["txt_attn_proj"]);
auto img_mlp = std::dynamic_pointer_cast<SwiGLUMlp>(blocks["img_mlp"]);
auto txt_mlp = std::dynamic_pointer_cast<SwiGLUMlp>(blocks["txt_mlp"]);
int64_t li = img->ne[1];
int64_t lt = txt->ne[1];
auto img_qkv = split_qkv(ctx->ggml_ctx, img_qkv_p->forward(ctx, img_norm1->forward(ctx, img)), num_heads, head_dim);
auto txt_qkv = split_qkv(ctx->ggml_ctx, txt_qkv_p->forward(ctx, txt_norm1->forward(ctx, txt)), num_heads, head_dim);
auto q = ggml_concat(ctx->ggml_ctx, q_norm->forward(ctx, txt_qkv[0]), q_norm->forward(ctx, img_qkv[0]), 2);
auto k = ggml_concat(ctx->ggml_ctx, k_norm->forward(ctx, txt_qkv[1]), k_norm->forward(ctx, img_qkv[1]), 2);
auto v = ggml_concat(ctx->ggml_ctx, txt_qkv[2], img_qkv[2], 2);
auto out = Rope::attention(ctx, q, k, v, pe, nullptr, 1.0f, false);
auto out_txt = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, lt);
auto out_img = ggml_ext_slice(ctx->ggml_ctx, out, 1, lt, lt + li);
img = ggml_add(ctx->ggml_ctx, img, img_proj->forward(ctx, out_img));
txt = ggml_add(ctx->ggml_ctx, txt, txt_proj->forward(ctx, out_txt));
img = ggml_add(ctx->ggml_ctx, img, img_mlp->forward(ctx, img_norm2->forward(ctx, img)));
txt = ggml_add(ctx->ggml_ctx, txt, txt_mlp->forward(ctx, txt_norm2->forward(ctx, txt)));
return {img, txt};
}
};
struct FinalLayer : public GGMLBlock {
FinalLayer(int64_t hidden_size, int64_t patch_size, int64_t out_channels) {
blocks["norm_final"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["linear"] = std::make_shared<Linear>(hidden_size, patch_size * patch_size * out_channels, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm_final = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_final"]);
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
return linear->forward(ctx, norm_final->forward(ctx, x));
}
};
struct MMJiT : public GGMLBlock {
MiniT2IConfig config;
MMJiT(const MiniT2IConfig& config)
: config(config) {
blocks["img_embedder"] = std::make_shared<BottleneckPatchEmbed>(config.patch_size, config.in_channels, config.pca_channels, config.hidden_size);
blocks["txt_embedder"] = std::make_shared<Linear>(config.txt_input_size, config.txt_hidden_size, false);
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>(config.cond_vec_size);
blocks["pooled_embedder"] = std::make_shared<Linear>(config.txt_input_size, config.cond_vec_size, false);
for (int64_t i = 0; i < config.txt_preamble_depth; ++i) {
blocks["txt_preamble_blocks." + std::to_string(i)] = std::make_shared<PlainTextTransformerBlock>(config.txt_hidden_size, config.num_heads, config.head_dim, config.mlp_ratio);
}
for (int64_t i = 0; i < config.depth_double; ++i) {
blocks["double_blocks." + std::to_string(i)] = std::make_shared<DoubleStreamDiTBlock>(config.hidden_size, config.txt_hidden_size, config.num_heads, config.head_dim, config.mlp_ratio);
}
blocks["final_layer"] = std::make_shared<FinalLayer>(config.hidden_size, config.patch_size, config.in_channels);
}
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
enum ggml_type wtype = get_type(prefix + "mask_token", tensor_storage_map, GGML_TYPE_F32);
params["mask_token"] = ggml_new_tensor_3d(ctx, wtype, config.txt_input_size, 1, 1);
}
ggml_tensor* apply_text_mask(GGMLRunnerContext* ctx, ggml_tensor* context, ggml_tensor* mask) {
if (mask == nullptr) {
return context;
}
mask = ggml_reshape_3d(ctx->ggml_ctx, mask, 1, mask->ne[0], mask->ne[1]);
mask = ggml_repeat(ctx->ggml_ctx, mask, context);
auto keep = ggml_mul(ctx->ggml_ctx, context, mask);
auto inv = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx->ggml_ctx, mask), mask);
auto mask_token = ggml_repeat(ctx->ggml_ctx, params["mask_token"], context);
return ggml_add(ctx->ggml_ctx, keep, ggml_mul(ctx->ggml_ctx, mask_token, inv));
}
ggml_tensor* pool_context(GGMLRunnerContext* ctx, ggml_tensor* context) {
int64_t dim = context->ne[0];
int64_t len = context->ne[1];
int64_t N = context->ne[2];
auto x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, context, 1, 0, 2, 3));
x = ggml_reshape_3d(ctx->ggml_ctx, x, len, dim, N);
x = ggml_mean(ctx->ggml_ctx, x);
x = ggml_reshape_2d(ctx->ggml_ctx, x, dim, N);
return x;
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* img,
ggml_tensor* context,
ggml_tensor* mask,
ggml_tensor* pos_embed,
ggml_tensor* txt_pe,
ggml_tensor* joint_pe) {
auto img_embedder = std::dynamic_pointer_cast<BottleneckPatchEmbed>(blocks["img_embedder"]);
auto txt_embedder = std::dynamic_pointer_cast<Linear>(blocks["txt_embedder"]);
auto final_layer = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
int64_t W = img->ne[0];
int64_t H = img->ne[1];
int64_t hp = H / config.patch_size;
int64_t wp = W / config.patch_size;
context = apply_text_mask(ctx, context, mask);
auto x = img_embedder->forward(ctx, img);
x = ggml_add(ctx->ggml_ctx, x, pos_embed);
auto txt = txt_embedder->forward(ctx, context);
for (int64_t i = 0; i < config.txt_preamble_depth; ++i) {
auto block = std::dynamic_pointer_cast<PlainTextTransformerBlock>(blocks["txt_preamble_blocks." + std::to_string(i)]);
txt = block->forward(ctx, txt, txt_pe);
sd::ggml_graph_cut::mark_graph_cut(txt, "minit2i.txt_preamble_blocks." + std::to_string(i), "txt");
}
for (int64_t i = 0; i < config.depth_double; ++i) {
auto block = std::dynamic_pointer_cast<DoubleStreamDiTBlock>(blocks["double_blocks." + std::to_string(i)]);
auto out = block->forward(ctx, x, txt, joint_pe);
x = out.first;
txt = out.second;
sd::ggml_graph_cut::mark_graph_cut(x, "minit2i.double_blocks." + std::to_string(i), "x");
sd::ggml_graph_cut::mark_graph_cut(txt, "minit2i.double_blocks." + std::to_string(i), "txt");
}
auto combined = ggml_concat(ctx->ggml_ctx, txt, x, 1);
auto out = final_layer->forward(ctx, combined);
auto img_out = ggml_ext_slice(ctx->ggml_ctx, out, 1, txt->ne[1], txt->ne[1] + x->ne[1]);
return DiT::unpatchify(ctx->ggml_ctx, img_out, hp, wp, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), false);
}
};
struct MiniT2IRunner : public DiffusionModelRunner {
MiniT2IConfig config;
MMJiT model;
ggml_context* position_cache_ctx = nullptr;
ggml_backend_buffer_t position_cache_buffer = nullptr;
ggml_tensor* cached_pos_embed = nullptr;
ggml_tensor* cached_txt_pe = nullptr;
ggml_tensor* cached_joint_pe = nullptr;
int64_t cached_img_side = -1;
int64_t cached_txt_len = -1;
int64_t cached_hidden_size = -1;
int64_t cached_head_dim = -1;
MiniT2IRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
config(MiniT2IConfig::detect_from_weights(tensor_storage_map, this->prefix)),
model(config) {
model.init(params_ctx, tensor_storage_map, this->prefix);
}
~MiniT2IRunner() override {
free_position_cache();
}
std::string get_desc() override {
return "MiniT2I";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
model.get_param_tensors(tensors, prefix);
}
void free_position_cache() {
if (position_cache_buffer != nullptr) {
ggml_backend_buffer_free(position_cache_buffer);
position_cache_buffer = nullptr;
}
if (position_cache_ctx != nullptr) {
ggml_free(position_cache_ctx);
position_cache_ctx = nullptr;
}
cached_pos_embed = nullptr;
cached_txt_pe = nullptr;
cached_joint_pe = nullptr;
cached_img_side = -1;
cached_txt_len = -1;
cached_hidden_size = -1;
cached_head_dim = -1;
}
void ensure_position_cache(int64_t img_side, int64_t txt_len) {
if (cached_img_side == img_side &&
cached_txt_len == txt_len &&
cached_hidden_size == config.hidden_size &&
cached_head_dim == config.head_dim &&
cached_pos_embed != nullptr &&
cached_txt_pe != nullptr &&
cached_joint_pe != nullptr) {
return;
}
free_position_cache();
auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast<int>(img_side), static_cast<int>(config.hidden_size));
auto txt_pe_vec = make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim));
auto img_pe_vec = make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim));
auto joint_pe_vec = txt_pe_vec;
joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end());
ggml_init_params params;
params.mem_size = static_cast<size_t>(3 * ggml_tensor_overhead());
params.mem_buffer = nullptr;
params.no_alloc = true;
position_cache_ctx = ggml_init(params);
GGML_ASSERT(position_cache_ctx != nullptr);
cached_pos_embed = ggml_new_tensor_3d(position_cache_ctx, GGML_TYPE_F32, config.hidden_size, img_side * img_side, 1);
ggml_set_name(cached_pos_embed, "minit2i.pos_embed");
cached_txt_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len);
ggml_set_name(cached_txt_pe, "minit2i.txt_pe");
cached_joint_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len + img_side * img_side);
ggml_set_name(cached_joint_pe, "minit2i.joint_pe");
position_cache_buffer = ggml_backend_alloc_ctx_tensors(position_cache_ctx, runtime_backend);
GGML_ASSERT(position_cache_buffer != nullptr);
ggml_backend_buffer_set_usage(position_cache_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
ggml_backend_tensor_set(cached_pos_embed, pos_embed_vec.data(), 0, ggml_nbytes(cached_pos_embed));
ggml_backend_tensor_set(cached_txt_pe, txt_pe_vec.data(), 0, ggml_nbytes(cached_txt_pe));
ggml_backend_tensor_set(cached_joint_pe, joint_pe_vec.data(), 0, ggml_nbytes(cached_joint_pe));
ggml_backend_synchronize(runtime_backend);
cached_img_side = img_side;
cached_txt_len = txt_len;
cached_hidden_size = config.hidden_size;
cached_head_dim = config.head_dim;
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
const sd::Tensor<float>& mask_tensor) {
ggml_cgraph* gf = new_graph_custom(MINIT2I_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* context = make_input(context_tensor);
ggml_tensor* mask = make_input(mask_tensor);
SD_UNUSED(timesteps_tensor);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t img_side = H / config.patch_size;
int64_t txt_len = context->ne[1];
ensure_position_cache(img_side, txt_len);
auto runner_ctx = get_context();
auto out = model.forward(&runner_ctx, x, context, mask, cached_pos_embed, cached_txt_pe, cached_joint_pe);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context,
const sd::Tensor<float>& mask) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, mask);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
GGML_ASSERT(diffusion_params.context != nullptr);
const auto* extra = diffusion_extra_as<MiniT2IDiffusionExtra>(diffusion_params);
GGML_ASSERT(extra->mask != nullptr);
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
*diffusion_params.context,
*extra->mask);
}
};
} // namespace MiniT2I
#endif // __SD_MODEL_DIFFUSION_MINIT2I_HPP__

View file

@ -1,4 +1,4 @@
#ifndef __SD_MODEL_DIFFUSION_MODEL_HPP__
#ifndef __SD_MODEL_DIFFUSION_MODEL_HPP__
#define __SD_MODEL_DIFFUSION_MODEL_HPP__
#include <string>
@ -7,6 +7,7 @@
#include "core/ggml_extend.hpp"
#include "core/tensor_ggml.hpp"
#include "model/common/rope.hpp"
#include "model_manager.h"
struct UNetDiffusionExtra {
@ -52,6 +53,10 @@ struct LTXAVDiffusionExtra {
const sd::Tensor<float>* video_positions = nullptr;
};
struct MiniT2IDiffusionExtra {
const sd::Tensor<float>* mask = nullptr;
};
using DiffusionExtraParams = std::variant<std::monostate,
UNetDiffusionExtra,
SkipLayerDiffusionExtra,
@ -59,7 +64,8 @@ using DiffusionExtraParams = std::variant<std::monostate,
AnimaDiffusionExtra,
WanDiffusionExtra,
HiDreamO1DiffusionExtra,
LTXAVDiffusionExtra>;
LTXAVDiffusionExtra,
MiniT2IDiffusionExtra>;
struct DiffusionParams {
const sd::Tensor<float>* x = nullptr;
@ -68,7 +74,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;
bool increase_ref_index = false;
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED;
DiffusionExtraParams extra = std::monostate{};
};

View file

@ -4,6 +4,7 @@
#include <memory>
#include "model/common/block.hpp"
#include "model/diffusion/dit.hpp"
#include "model/diffusion/flux.hpp"
#include "model/diffusion/model.hpp"
#include "model_loader.h"
@ -23,6 +24,7 @@ namespace Qwen {
std::vector<int> axes_dim = {16, 56, 56};
int axes_dim_sum = 128;
bool zero_cond_t = false;
bool use_additional_t_cond = false;
static QwenImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
QwenImageConfig config;
@ -88,19 +90,33 @@ namespace Qwen {
};
struct QwenTimestepProjEmbeddings : public GGMLBlock {
protected:
bool use_additional_t_cond = false;
public:
QwenTimestepProjEmbeddings(int64_t embedding_dim) {
QwenTimestepProjEmbeddings(int64_t embedding_dim, bool use_additional_t_cond = false)
: use_additional_t_cond(use_additional_t_cond) {
blocks["timestep_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedding(256, embedding_dim));
if (use_additional_t_cond) {
blocks["addition_t_embedding"] = std::shared_ptr<GGMLBlock>(new Embedding(2, embedding_dim));
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* timesteps) {
ggml_tensor* timesteps,
ggml_tensor* addition_t_cond = nullptr) {
// timesteps: [N,]
// return: [N, embedding_dim]
auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
auto timesteps_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1.f);
auto timesteps_emb = timestep_embedder->forward(ctx, timesteps_proj);
if (use_additional_t_cond) {
GGML_ASSERT(addition_t_cond != nullptr);
auto addition_t_embedding = std::dynamic_pointer_cast<Embedding>(blocks["addition_t_embedding"]);
auto addition_t_emb = addition_t_embedding->forward(ctx, addition_t_cond);
timesteps_emb = ggml_add(ctx->ggml_ctx, timesteps_emb, addition_t_emb);
}
return timesteps_emb;
}
};
@ -402,7 +418,7 @@ namespace Qwen {
QwenImageModel(QwenImageConfig config)
: config(config) {
int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim));
blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim, config.use_additional_t_cond));
blocks["txt_norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(config.joint_attention_dim, 1e-6f));
blocks["img_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.in_channels, inner_dim));
blocks["txt_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.joint_attention_dim, inner_dim));
@ -424,6 +440,7 @@ namespace Qwen {
ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* addition_t_cond,
ggml_tensor* context,
ggml_tensor* pe,
ggml_tensor* modulate_index = nullptr) {
@ -434,9 +451,9 @@ namespace Qwen {
auto norm_out = std::dynamic_pointer_cast<AdaLayerNormContinuous>(blocks["norm_out"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
auto t_emb = time_text_embed->forward(ctx, timestep);
auto t_emb = time_text_embed->forward(ctx, timestep, addition_t_cond);
if (config.zero_cond_t) {
auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep));
auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), addition_t_cond);
t_emb = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1);
}
auto img = img_in->forward(ctx, x);
@ -469,33 +486,50 @@ namespace Qwen {
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* addition_t_cond,
ggml_tensor* context,
ggml_tensor* pe,
std::vector<ggml_tensor*> ref_latents = {},
ggml_tensor* modulate_index = nullptr) {
// Forward pass of DiT.
// x: [N, C, H, W]
// x: [N, C, H, W] or [N*C, T, H, W]
// timestep: [N,]
// context: [N, L, D]
// pe: [L, d_head/2, 2, 2]
// return: [N, C, H, W]
// return: [N, C, H, W] or [N*C, T, H, W]
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t C = x->ne[2];
int64_t N = x->ne[3];
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t T = 1;
int64_t N = addition_t_cond != nullptr ? addition_t_cond->ne[0] : x->ne[3];
bool has_time_axis = false;
if (x->ne[3] != 1) {
T = x->ne[2];
has_time_axis = true;
}
auto img = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size);
auto patchify_input = [&](ggml_tensor* input) -> ggml_tensor* {
input = DiT::pad_to_patch_size(ctx, input, config.patch_size, config.patch_size);
if (!has_time_axis) {
return DiT::patchify(ctx->ggml_ctx, input, config.patch_size, config.patch_size);
}
if (input->ne[3] == 1) {
input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]);
}
return DiT::patchify(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
};
auto img = patchify_input(x);
int64_t img_tokens = img->ne[1];
if (ref_latents.size() > 0) {
for (ggml_tensor* ref : ref_latents) {
ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
ref = patchify_input(ref);
img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
}
}
auto out = forward_orig(ctx, img, timestep, context, pe, modulate_index); // [N, h_len*w_len, ph*pw*C]
auto out = forward_orig(ctx, img, timestep, addition_t_cond, context, pe, modulate_index); // [N, h_len*w_len, ph*pw*C]
if (out->ne[1] > img_tokens) {
out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3)); // [num_tokens, N, C * patch_size * patch_size]
@ -503,7 +537,17 @@ namespace Qwen {
out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3)); // [N, h*w, C * patch_size * patch_size]
}
out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size); // [N, C, H, W]
if (has_time_axis) {
int pad_h = (config.patch_size - H % config.patch_size) % config.patch_size;
int pad_w = (config.patch_size - W % config.patch_size) % config.patch_size;
int h_len = static_cast<int>((H + pad_h) / config.patch_size);
int w_len = static_cast<int>((W + pad_w) / config.patch_size);
out = DiT::unpatchify_3d(ctx->ggml_ctx, out, T, h_len, w_len, 1, config.patch_size, config.patch_size);
out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H); // [N*C, T, H, W + pad_w]
out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W); // [N*C, T, H, W]
} else {
out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size); // [N, C, H, W]
}
return out;
}
@ -515,6 +559,7 @@ namespace Qwen {
QwenImageModel qwen_image;
std::vector<float> pe_vec;
std::vector<float> modulate_index_vec;
std::vector<int32_t> additional_t_cond_vec;
SDVersion version;
QwenImageRunner(ggml_backend_t backend,
@ -524,9 +569,13 @@ namespace Qwen {
bool zero_cond_t = false,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)),
version(version) {
config.zero_cond_t = config.zero_cond_t || zero_cond_t;
qwen_image = QwenImageModel(config);
if (version == VERSION_QWEN_IMAGE_LAYERED) {
config.use_additional_t_cond = true;
}
qwen_image = QwenImageModel(config);
qwen_image.init(params_ctx, tensor_storage_map, prefix);
}
@ -542,11 +591,11 @@ namespace Qwen {
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
bool increase_ref_index = false) {
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::INCREASE) {
ggml_cgraph* gf = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
GGML_ASSERT(x->ne[3] == 1);
GGML_ASSERT(x->ne[3] == 1 || x_tensor.dim() == 5);
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
std::vector<ggml_tensor*> ref_latents;
@ -555,13 +604,29 @@ namespace Qwen {
ref_latents.push_back(make_input(ref_latent_tensor));
}
pe_vec = Rope::gen_qwen_image_pe(static_cast<int>(x->ne[1]),
int batch_size = static_cast<int>(x->ne[3]);
int time_len = 1;
if (x_tensor.dim() == 5) {
time_len = static_cast<int>(x_tensor.shape()[2]);
batch_size = static_cast<int>(x_tensor.shape()[4]);
}
ggml_tensor* addition_t_cond = nullptr;
if (version == VERSION_QWEN_IMAGE_LAYERED) {
additional_t_cond_vec.assign(static_cast<size_t>(batch_size), 0);
addition_t_cond = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, batch_size);
set_backend_tensor_data(addition_t_cond, additional_t_cond_vec.data());
ref_index_mode = Rope::RefIndexMode::DECREASE;
}
pe_vec = Rope::gen_qwen_image_pe(time_len,
static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
batch_size,
static_cast<int>(context->ne[1]),
ref_latents,
increase_ref_index,
ref_index_mode,
config.theta,
circular_y_enabled,
circular_x_enabled,
@ -604,6 +669,7 @@ namespace Qwen {
ggml_tensor* out = qwen_image.forward(&runner_ctx,
x,
timesteps,
addition_t_cond,
context,
pe,
ref_latents,
@ -619,12 +685,12 @@ namespace Qwen {
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context,
const std::vector<sd::Tensor<float>>& ref_latents = {},
bool increase_ref_index = false) {
// x: [N, in_channels, h, w]
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::INCREASE) {
// x: [N, C, H, W] or [N*C, T, H, W]
// timesteps: [N, ]
// context: [N, max_position, hidden_size]
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@ -640,7 +706,7 @@ namespace Qwen {
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.increase_ref_index);
diffusion_params.ref_index_mode);
}
void test() {
@ -674,7 +740,7 @@ namespace Qwen {
timesteps,
context,
{},
false);
Rope::RefIndexMode::FIXED);
int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty());

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

@ -575,7 +575,7 @@ namespace ZImage {
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
bool increase_ref_index = false) {
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) {
ggml_cgraph* gf = new_graph_custom(Z_IMAGE_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
@ -595,7 +595,7 @@ namespace ZImage {
static_cast<int>(context->ne[1]),
SEQ_MULTI_OF,
ref_latents,
increase_ref_index,
ref_index_mode,
config.theta,
circular_y_enabled,
circular_x_enabled,
@ -626,12 +626,12 @@ namespace ZImage {
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context,
const std::vector<sd::Tensor<float>>& ref_latents = {},
bool increase_ref_index = false) {
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) {
// x: [N, in_channels, h, w]
// timesteps: [N, ]
// context: [N, max_position, hidden_size]
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@ -647,7 +647,7 @@ namespace ZImage {
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.increase_ref_index);
diffusion_params.ref_index_mode);
}
void test() {
@ -681,7 +681,7 @@ namespace ZImage {
timesteps,
context,
{},
false);
Rope::RefIndexMode::FIXED);
int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty());

View file

@ -1,4 +1,4 @@
#ifndef __SD_MODEL_TE_CLIP_HPP__
#ifndef __SD_MODEL_TE_CLIP_HPP__
#define __SD_MODEL_TE_CLIP_HPP__
#include "core/ggml_extend.hpp"

View file

@ -1,4 +1,4 @@
#ifndef __SD_MODEL_TE_LLM_HPP__
#ifndef __SD_MODEL_TE_LLM_HPP__
#define __SD_MODEL_TE_LLM_HPP__
#include <algorithm>
@ -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

@ -1,4 +1,4 @@
#ifndef __SD_MODEL_TE_T5_HPP__
#ifndef __SD_MODEL_TE_T5_HPP__
#define __SD_MODEL_TE_T5_HPP__
#include <cfloat>
@ -26,13 +26,66 @@ struct T5Config {
static T5Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
bool is_umt5 = false) {
(void)tensor_storage_map;
(void)prefix;
T5Config config;
if (is_umt5) {
config.vocab_size = 256384;
config.relative_attention = false;
}
auto find_tensor = [&](const std::string& suffix) -> const TensorStorage* {
auto it = tensor_storage_map.find(prefix + "." + suffix);
if (it != tensor_storage_map.end()) {
return &it->second;
}
it = tensor_storage_map.find(prefix + suffix);
if (it != tensor_storage_map.end()) {
return &it->second;
}
return nullptr;
};
if (const TensorStorage* shared = find_tensor("shared.weight")) {
if (shared->n_dims == 2) {
config.vocab_size = shared->ne[1];
config.model_dim = shared->ne[0];
}
}
if (const TensorStorage* q = find_tensor("encoder.block.0.layer.0.SelfAttention.q.weight")) {
if (q->n_dims == 2) {
config.model_dim = q->ne[0];
int64_t inner_dim = q->ne[1];
// Flan-T5/T5 uses d_kv=64 for common sizes.
if (inner_dim % 64 == 0) {
config.num_heads = inner_dim / 64;
}
}
}
if (const TensorStorage* wi = find_tensor("encoder.block.0.layer.1.DenseReluDense.wi_0.weight")) {
if (wi->n_dims == 2) {
config.model_dim = wi->ne[0];
config.ff_dim = wi->ne[1];
}
}
int64_t detected_layers = 0;
for (const auto& [name, _] : tensor_storage_map) {
std::string base = prefix;
if (!base.empty() && base.back() != '.') {
base += ".";
}
std::string layer_prefix = base + "encoder.block.";
if (!starts_with(name, layer_prefix)) {
continue;
}
size_t pos = layer_prefix.size();
size_t dot = name.find('.', pos);
if (dot == std::string::npos) {
continue;
}
int64_t layer = atoi(name.substr(pos, dot - pos).c_str());
detected_layers = std::max(detected_layers, layer + 1);
}
if (detected_layers > 0) {
config.num_layers = detected_layers;
}
return config;
}
};

View file

@ -1,4 +1,4 @@
#ifndef __SD_MODEL_UPSCALER_ESRGAN_HPP__
#ifndef __SD_MODEL_UPSCALER_ESRGAN_HPP__
#define __SD_MODEL_UPSCALER_ESRGAN_HPP__
#include <algorithm>

View file

@ -1,4 +1,4 @@
#ifndef __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
#ifndef __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
#define __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
#include <algorithm>

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

@ -1,4 +1,4 @@
#ifndef __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__
#ifndef __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__
#define __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__
#include <algorithm>
@ -327,7 +327,7 @@ namespace LTXV {
auto x = ggml_reshape_3d(ctx, waveform, time, 1, channels * batch);
if (left_pad > 0) {
x = ggml_pad_ext(ctx, x, static_cast<int>(left_pad), 0, 0, 0, 0, 0, 0, 0);
x = ggml_ext_pad_ext(ctx, runner_ctx->backend, x, static_cast<int>(left_pad), 0, 0, 0, 0, 0, 0, 0);
}
auto frames = ggml_conv_1d(ctx, forward_basis, x, hop_length, 0, 1);
@ -564,6 +564,7 @@ namespace LTXV {
int pad_h = kernel_size.first - 1;
int pad_w = kernel_size.second - 1;
x = ggml_ext_pad_ext(ctx->ggml_ctx,
ctx->backend,
x,
pad_w / 2,
pad_w - pad_w / 2,
@ -603,6 +604,7 @@ namespace LTXV {
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
x = ggml_ext_pad_ext(ctx->ggml_ctx,
ctx->backend,
x,
0,
1,

View file

@ -408,7 +408,7 @@ public:
h = conv->forward(ctx, h);
for (int j = 0; j < num_blocks; j++) {
auto block = std::dynamic_pointer_cast<MemBlock>(blocks[std::to_string(index++)]);
auto mem = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
auto mem = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, h, 0, 0, 0, 0, 0, 0, 1, 0);
mem = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
h = block->forward(ctx, h, mem);
}
@ -479,7 +479,7 @@ public:
int index = 3;
for (int i = 0; i < num_layers; i++) {
for (int j = 0; j < num_blocks; j++) {
auto mem = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
auto mem = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, h, 0, 0, 0, 0, 0, 0, 1, 0);
mem = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
if (is_wide) {
auto block = std::dynamic_pointer_cast<WideMemBlock>(blocks[std::to_string(index++)]);
@ -548,7 +548,7 @@ public:
}
auto result = decoder->forward(ctx, z);
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
// (W, H, C, T) -> (W, H, T, C)
// (W, H, T, C) -> (W, H, C, T)
result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2));
}
return result;
@ -556,8 +556,10 @@ public:
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);
// (W, H, T, C) -> (W, H, C, T)
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
// (W, H, T, C) -> (W, H, C, T)
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
}
int64_t num_frames = x->ne[3];
if (num_frames % encoder->t_downscale) {
// pad to multiple of encoder->t_downscale at the end
@ -567,7 +569,10 @@ public:
}
}
x = encoder->forward(ctx, x);
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
// (W, H, C, T) -> (W, H, T, C)
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
}
return x;
}
};

View file

@ -1,4 +1,4 @@
#ifndef __SD_MODEL_VAE_VAE_HPP__
#ifndef __SD_MODEL_VAE_VAE_HPP__
#define __SD_MODEL_VAE_VAE_HPP__
#include "core/tensor_ggml.hpp"
@ -78,7 +78,7 @@ public:
scale_factor = 16;
} else if (sd_version_uses_flux2_vae(version)) {
scale_factor = 16;
} else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) {
} else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) {
scale_factor = 1;
}
return scale_factor;

View file

@ -72,8 +72,8 @@ namespace WAN {
lp2 -= (int)cache_x->ne[2];
}
x = ggml_ext_pad_ext(ctx->ggml_ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
return ggml_ext_conv_3d(ctx->ggml_ctx, x, w, b, in_channels,
x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
0, 0, 0,
std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation));
@ -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;
@ -177,7 +195,7 @@ namespace WAN {
2);
}
if (chunk_idx == 1 && cache_x->ne[2] < 2) { // Rep
cache_x = ggml_pad_ext(ctx->ggml_ctx, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0);
cache_x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0);
// aka cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device),cache_x],dim=2)
}
if (chunk_idx == 1) {
@ -265,7 +283,7 @@ namespace WAN {
int pad_t = (factor_t - T % factor_t) % factor_t;
x = ggml_pad_ext(ctx->ggml_ctx, x, 0, 0, 0, 0, pad_t, 0, 0, 0);
x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, 0, 0, 0, 0, pad_t, 0, 0, 0);
T = x->ne[2];
x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, factor_t, T / factor_t, C); // [C, T/factor_t, factor_t, H*W]
@ -338,19 +356,32 @@ 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()
blocks["residual.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(in_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
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()
blocks["residual.6"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
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) {
blocks["shortcut"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(in_dim, out_dim, {1, 1, 1}));
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}));
}
}
}
@ -363,9 +394,15 @@ namespace WAN {
GGML_ASSERT(b == 1);
ggml_tensor* h = x;
if (in_dim != out_dim) {
auto shortcut = std::dynamic_pointer_cast<CausalConv3d>(blocks["shortcut"]);
if (is_2D) {
auto shortcut = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["shortcut"]);
h = shortcut->forward(ctx, x);
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++) {
@ -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, feat_cache[idx]);
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) {
@ -587,35 +631,41 @@ namespace WAN {
class Encoder3d : public GGMLBlock {
protected:
bool wan2_2;
int64_t in_channels;
int64_t dim;
int64_t z_dim;
std::vector<int> dim_mult;
int num_res_blocks;
std::vector<bool> temperal_downsample;
bool is_2D = false;
public:
Encoder3d(int64_t dim = 128,
int64_t z_dim = 4,
int64_t in_channels = 3,
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)
: dim(dim),
bool wan2_2 = false,
bool is_2D = false)
: in_channels(in_channels),
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}));
if (is_2D) {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(in_channels, 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(in_channels, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
}
int index = 0;
@ -630,12 +680,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,13 +699,17 @@ 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()
blocks["head.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, z_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
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,
@ -673,7 +728,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 +786,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) {
@ -753,25 +814,31 @@ namespace WAN {
class Decoder3d : public GGMLBlock {
protected:
bool wan2_2;
int64_t out_channels;
int64_t dim;
int64_t z_dim;
std::vector<int> dim_mult;
int num_res_blocks;
std::vector<bool> temperal_upsample;
bool is_2D = false;
public:
Decoder3d(int64_t dim = 128,
int64_t z_dim = 4,
int64_t out_channels = 3,
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)
: dim(dim),
bool wan2_2 = false,
bool is_2D = false)
: out_channels(out_channels),
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 +846,16 @@ namespace WAN {
}
// init block
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
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 +870,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 +879,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;
}
@ -821,13 +893,14 @@ namespace WAN {
}
// output blocks
blocks["head.0"] = std::shared_ptr<GGMLBlock>(new RMS_norm(out_dim));
blocks["head.0"] = std::shared_ptr<GGMLBlock>(new RMS_norm(out_dim));
int64_t final_dim = out_channels;
// 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 +920,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 +978,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) {
@ -928,6 +1007,8 @@ namespace WAN {
public:
bool wan2_2 = false;
bool decode_only = true;
int64_t input_channels = 3;
int patch_size = 1;
int64_t dim = 96;
int64_t dec_dim = 96;
int64_t z_dim = 16;
@ -935,6 +1016,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,23 +1033,43 @@ 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, SDVersion version = VERSION_WAN2, bool is_2D = false)
: decode_only(decode_only),
wan2_2(version == VERSION_WAN2_2_TI2V),
is_2D(is_2D) {
// attn_scales is always []
if (wan2_2) {
dim = 160;
dec_dim = 256;
z_dim = 48;
dim = 160;
dec_dim = 256;
z_dim = 48;
input_channels = 12;
patch_size = 2;
_conv_num = 34;
_enc_conv_num = 26;
} else if (version == VERSION_QWEN_IMAGE_LAYERED) {
input_channels = 4;
}
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["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim * 2, z_dim * 2, {1, 1, 1}));
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new Encoder3d(dim, z_dim * 2, input_channels, 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, input_channels, 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}));
}
blocks["decoder"] = std::shared_ptr<GGMLBlock>(new Decoder3d(dec_dim, z_dim, dim_mult, num_res_blocks, temperal_upsample, wan2_2));
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, z_dim, {1, 1, 1}));
}
static ggml_tensor* patchify(ggml_context* ctx,
@ -1030,11 +1132,13 @@ namespace WAN {
GGML_ASSERT(b == 1);
GGML_ASSERT(decode_only == false);
if (x->ne[2] > 1 && is_2D) {
LOG_WARN("Using 2D VAE to encode video, expect poor results");
}
clear_cache();
if (wan2_2) {
x = patchify(ctx->ggml_ctx, x, 2, b);
}
x = patchify(ctx->ggml_ctx, x, patch_size, b);
// sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encode.prelude", "x");
auto encoder = std::dynamic_pointer_cast<Encoder3d>(blocks["encoder"]);
@ -1049,12 +1153,18 @@ namespace WAN {
auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1); // [b*c, 1, h, w]
out = encoder->forward(ctx, in, b, _enc_feat_map, _enc_conv_idx, i);
} else {
auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1 + 4 * (i - 1), 1 + 4 * i); // [b*c, 4, h, w]
// if is_2D, drop 3 out of 4 frames
auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1 + 4 * (i - 1), (is_2D ? 1 - 3 : 1) + 4 * i); // [b*c, 4, h, w]
auto out_ = encoder->forward(ctx, in, b, _enc_feat_map, _enc_conv_idx, i);
out = ggml_concat(ctx->ggml_ctx, out, out_, 2);
}
}
out = conv1->forward(ctx, out);
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();
@ -1067,13 +1177,23 @@ namespace WAN {
// z: [b*c, t, h, w]
GGML_ASSERT(b == 1);
if (z->ne[2] > 1 && is_2D) {
LOG_WARN("Using 2D VAE to decode video, expect poor results");
}
clear_cache();
auto decoder = std::dynamic_pointer_cast<Decoder3d>(blocks["decoder"]);
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++) {
@ -1085,11 +1205,15 @@ namespace WAN {
auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1); // [b*c, 1, h, w]
auto out_ = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i);
out = ggml_concat(ctx->ggml_ctx, out, out_, 2);
if (is_2D) {
// repeat frames to avoid mismatch
for (int j = 0; j < 4 - 1; j++) {
out = ggml_concat(ctx->ggml_ctx, out, out_, 2);
}
}
}
}
if (wan2_2) {
out = unpatchify(ctx->ggml_ctx, out, 2, b);
}
out = unpatchify(ctx->ggml_ctx, out, patch_size, b);
// sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode.final", "out");
clear_cache();
return out;
@ -1110,9 +1234,7 @@ namespace WAN {
auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1); // [b*c, 1, h, w]
_conv_idx = 0;
auto out = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i);
if (wan2_2) {
out = unpatchify(ctx->ggml_ctx, out, 2, b);
}
out = unpatchify(ctx->ggml_ctx, out, patch_size, b);
// sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode_partial.final", "out");
return out;
}
@ -1129,7 +1251,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, is_2D);
ae.init(params_ctx, tensor_storage_map, prefix);
}

View file

@ -20,6 +20,7 @@
#include "model_io/torch_legacy_io.h"
#include "model_io/torch_zip_io.h"
#include "model_loader.h"
#include "runtime/imatrix.h"
#include "stable-diffusion.h"
#include "core/ggml_extend_backend.h"
@ -66,7 +67,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) {
@ -171,7 +171,8 @@ void convert_tensor(void* src,
void* dst,
ggml_type dst_type,
int nrows,
int n_per_row) {
int n_per_row,
std::vector<float> imatrix = {}) {
int n = nrows * n_per_row;
if (src_type == dst_type) {
size_t nbytes = n * ggml_type_size(src_type) / ggml_blck_size(src_type);
@ -180,7 +181,7 @@ void convert_tensor(void* src,
if (dst_type == GGML_TYPE_F16) {
ggml_fp32_to_fp16_row((float*)src, (ggml_fp16_t*)dst, n);
} else {
std::vector<float> imatrix(n_per_row, 1.0f); // dummy importance matrix
imatrix.resize(n_per_row, 1.0f);
const float* im = imatrix.data();
ggml_quantize_chunk(dst_type, (float*)src, dst, 0, nrows, n_per_row, im);
}
@ -210,7 +211,7 @@ void convert_tensor(void* src,
if (dst_type == GGML_TYPE_F16) {
ggml_fp32_to_fp16_row((float*)src_data_f32, (ggml_fp16_t*)dst, n);
} else {
std::vector<float> imatrix(n_per_row, 1.0f); // dummy importance matrix
imatrix.resize(n_per_row, 1.0f);
const float* im = imatrix.data();
ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, nrows, n_per_row, im);
}
@ -499,7 +500,13 @@ SDVersion ModelLoader::get_sd_version() {
tensor_storage_map.find("model.diffusion_model.transformer_blocks.0.img_mlp.w1.weight") != tensor_storage_map.end()) {
return VERSION_LENS;
}
if (tensor_storage.name.find("net.img_embedder.proj1.weight") != std::string::npos) {
return VERSION_MINIT2I;
}
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
if (tensor_storage_map.find("model.diffusion_model.time_text_embed.addition_t_embedding.weight") != tensor_storage_map.end()) {
return VERSION_QWEN_IMAGE_LAYERED;
}
return VERSION_QWEN_IMAGE;
}
if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) {
@ -508,6 +515,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;
}
@ -996,6 +1006,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
size_t total_tensors_processed = 0;
const int64_t t_start = start_time;
int last_n_threads = 1;
SDVersion imatrix_version = (version_ == VERSION_COUNT) ? get_sd_version() : version_;
for (size_t file_index = 0; file_index < file_data.size(); ++file_index) {
auto& fdata = file_data[file_index];
@ -1180,12 +1191,15 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
failed = true;
return;
}
std::string processed_name = convert_tensor_name(tensor_storage.name, imatrix_version);
std::vector<float> imatrix = get_imatrix_collector().get_values(processed_name);
convert_tensor((void*)target_buf,
tensor_storage.type,
convert_buf,
dst_tensor->type,
(int)tensor_storage.nelements() / (int)tensor_storage.ne[0],
(int)tensor_storage.ne[0]);
(int)tensor_storage.ne[0],
std::move(imatrix));
} else {
convert_buf = read_buf;
}

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

@ -302,6 +302,137 @@ struct KarrasScheduler : SigmaScheduler {
}
};
struct BetaScheduler : SigmaScheduler {
static constexpr double alpha = 0.6;
static constexpr double beta = 0.6;
static double log_beta(double a, double b) {
return std::lgamma(a) + std::lgamma(b) - std::lgamma(a + b);
}
static double incbeta(double x, double a, double b) {
if (x <= 0.0) {
return 0.0;
}
if (x >= 1.0) {
return 1.0;
}
// Continued fraction approximation using Lentz's method.
const int max_iter = 200;
const double epsilon = 3.0e-7;
const double tiny = 1e-30;
const double qab = a + b;
const double qap = a + 1.0;
const double qam = a - 1.0;
double c = 1.0;
double d = 1.0 - qab * x / qap;
if (std::abs(d) < tiny) {
d = tiny;
}
d = 1.0 / d;
double h = d;
for (int m = 1; m <= max_iter; m++) {
const int m2 = 2 * m;
double aa = m * (b - m) * x / ((qam + m2) * (a + m2));
d = 1.0 + aa * d;
if (std::abs(d) < tiny) {
d = tiny;
}
c = 1.0 + aa / c;
if (std::abs(c) < tiny) {
c = tiny;
}
d = 1.0 / d;
h *= d * c;
aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2));
d = 1.0 + aa * d;
if (std::abs(d) < tiny) {
d = tiny;
}
c = 1.0 + aa / c;
if (std::abs(c) < tiny) {
c = tiny;
}
d = 1.0 / d;
const double del = d * c;
h *= del;
if (std::abs(del - 1.0) < epsilon) {
break;
}
}
return std::exp(a * std::log(x) + b * std::log(1.0 - x) - log_beta(a, b)) / a * h;
}
static double beta_cdf(double x, double a, double b) {
if (x == 0.0) {
return 0.0;
}
if (x == 1.0) {
return 1.0;
}
if (x < (a + 1.0) / (a + b + 2.0)) {
return incbeta(x, a, b);
}
return 1.0 - incbeta(1.0 - x, b, a);
}
static double beta_ppf(double u, double a, double b, int max_iter = 30) {
double x = 0.5;
for (int i = 0; i < max_iter; i++) {
const double f = beta_cdf(x, a, b) - u;
if (std::abs(f) < 1e-10) {
break;
}
const double df = std::exp((a - 1.0) * std::log(x) + (b - 1.0) * std::log(1.0 - x) - log_beta(a, b));
x -= f / df;
if (x <= 0.0) {
x = 1e-10;
}
if (x >= 1.0) {
x = 1.0 - 1e-10;
}
}
return x;
}
std::vector<float> get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t t_to_sigma) override {
std::vector<float> result;
result.reserve(n + 1);
const int t_max = TIMESTEPS - 1;
if (n == 0) {
return result;
} else if (n == 1) {
result.push_back(t_to_sigma(static_cast<float>(t_max)));
result.push_back(0.f);
return result;
}
int last_t = -1;
for (uint32_t i = 0; i < n; i++) {
const double u = 1.0 - static_cast<double>(i) / static_cast<double>(n);
const double t_cont = beta_ppf(u, alpha, beta) * t_max;
const int t = static_cast<int>(std::lround(t_cont));
if (t != last_t) {
result.push_back(t_to_sigma(static_cast<float>(t)));
last_t = t;
}
}
result.push_back(0.f);
return result;
}
};
struct SimpleScheduler : SigmaScheduler {
std::vector<float> get_sigmas(uint32_t n, float sigma_min, float sigma_max, t_to_sigma_t t_to_sigma) override {
std::vector<float> result_sigmas;
@ -559,6 +690,122 @@ struct LTX2Scheduler : SigmaScheduler {
}
};
inline float flux_time_shift(float mu, float sigma, float t) {
return ::expf(mu) / (::expf(mu) + ::powf((1.0f / t - 1.0f), sigma));
}
// https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py#L289
struct FluxScheduler : SigmaScheduler {
int image_seq_len = 0;
float base_shift = 0.5f;
float max_shift = 1.15f;
explicit FluxScheduler(int image_seq_len, const char* extra_sample_args = nullptr)
: image_seq_len(image_seq_len) {
parse_extra_sample_args(extra_sample_args);
}
void parse_extra_sample_args(const char* extra_sample_args) {
for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "flux scheduler arg")) {
if (key == "base_shift") {
if (!parse_strict_float(value, base_shift)) {
LOG_WARN("ignoring invalid flux scheduler arg '%s=%s'", key.c_str(), value.c_str());
}
} else if (key == "max_shift") {
if (!parse_strict_float(value, max_shift)) {
LOG_WARN("ignoring invalid flux scheduler arg '%s=%s'", key.c_str(), value.c_str());
}
}
}
}
float compute_mu() const {
constexpr float base_shift_anchor = 256.0f;
constexpr float max_shift_anchor = 4096.0f;
float m = (max_shift - base_shift) / (max_shift_anchor - base_shift_anchor);
float b = base_shift - m * base_shift_anchor;
return static_cast<float>(image_seq_len) * m + b;
}
std::vector<float> get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override {
std::vector<float> sigmas;
sigmas.reserve(n + 1);
float mu = compute_mu();
LOG_DEBUG("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu);
if (n == 0) {
sigmas.push_back(1.0f);
return sigmas;
}
for (uint32_t i = 0; i <= n; ++i) {
float t = 1.0f - static_cast<float>(i) / static_cast<float>(n);
if (t <= 0.0f) {
sigmas.push_back(0.0f);
} else {
sigmas.push_back(flux_time_shift(mu, 1.0f, t));
}
}
sigmas[n] = 0.0f;
return sigmas;
}
};
// https://github.com/black-forest-labs/flux2/blob/main/src/flux2/sampling.py#L244
struct Flux2Scheduler : SigmaScheduler {
int image_seq_len = 0;
explicit Flux2Scheduler(int image_seq_len)
: image_seq_len(image_seq_len) {}
static float compute_empirical_mu(int image_seq_len, uint32_t num_steps) {
const float a1 = 8.73809524e-05f;
const float b1 = 1.89833333f;
const float a2 = 0.00016927f;
const float b2 = 0.45666666f;
if (image_seq_len > 4300) {
return a2 * image_seq_len + b2;
}
float m_200 = a2 * image_seq_len + b2;
float m_10 = a1 * image_seq_len + b1;
float a = (m_200 - m_10) / 190.0f;
float b = m_200 - 200.0f * a;
return a * num_steps + b;
}
std::vector<float> get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override {
std::vector<float> sigmas;
sigmas.reserve(n + 1);
float mu = compute_empirical_mu(image_seq_len, n);
LOG_DEBUG("Flux2 scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu);
if (n == 0) {
sigmas.push_back(1.0f);
return sigmas;
}
for (uint32_t i = 0; i <= n; ++i) {
float t = 1.0f - static_cast<float>(i) / static_cast<float>(n);
if (t <= 0.0f) {
sigmas.push_back(0.0f);
} else if (t >= 1.0f) {
sigmas.push_back(1.0f);
} else {
sigmas.push_back(flux_time_shift(mu, 1.0f, t));
}
}
sigmas[n] = 0.0f;
return sigmas;
}
};
/*
* Logit-Normal Scheduler
* Based on: https://github.com/ideogram-oss/ideogram4/blob/main/src/ideogram4/scheduler.py
@ -602,7 +849,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));
}
}
@ -616,7 +863,6 @@ struct LogitNormalScheduler : SigmaScheduler {
one_minus_t_min = sigmoid(0.5f * logsnr_max);
// t_max = 1.0f / (1.0f + std::exp(0.5f * logsnr_min));
one_minus_t_max = sigmoid(0.5f * logsnr_min);
}
LogitNormalScheduler(int image_seq_len = 0, const char* extra_sample_args = nullptr) {
@ -736,7 +982,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;
@ -780,6 +1026,10 @@ struct Denoiser {
LOG_INFO("get_sigmas with Karras scheduler");
scheduler = std::make_shared<KarrasScheduler>();
break;
case BETA_SCHEDULER:
LOG_INFO("get_sigmas with Beta scheduler");
scheduler = std::make_shared<BetaScheduler>();
break;
case EXPONENTIAL_SCHEDULER:
LOG_INFO("get_sigmas exponential scheduler");
scheduler = std::make_shared<ExponentialScheduler>();
@ -825,6 +1075,16 @@ struct Denoiser {
scheduler = std::make_shared<LogitNormalScheduler>(image_seq_len, extra_sample_args);
break;
}
case FLUX2_SCHEDULER: {
LOG_INFO("get_sigmas with Flux2 scheduler");
scheduler = std::make_shared<Flux2Scheduler>(image_seq_len);
break;
}
case FLUX_SCHEDULER: {
LOG_INFO("get_sigmas with Flux scheduler");
scheduler = std::make_shared<FluxScheduler>(image_seq_len, extra_sample_args);
break;
}
default:
LOG_INFO("get_sigmas with discrete scheduler (default)");
scheduler = std::make_shared<DiscreteScheduler>();
@ -989,10 +1249,6 @@ struct DiscreteFlowDenoiser : public Denoiser {
}
};
inline float flux_time_shift(float mu, float sigma, float t) {
return ::expf(mu) / (::expf(mu) + ::powf((1.0f / t - 1.0f), sigma));
}
struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
FluxFlowDenoiser() = default;
@ -1006,35 +1262,141 @@ struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
}
};
struct Flux2FlowDenoiser : public FluxFlowDenoiser {
Flux2FlowDenoiser() = default;
struct SefiFlowDenoiser;
float compute_empirical_mu(uint32_t n, int image_seq_len) {
const float a1 = 8.73809524e-05f;
const float b1 = 1.89833333f;
const float a2 = 0.00016927f;
const float b2 = 0.45666666f;
struct SefiFlowDenoiser : public FluxFlowDenoiser {
static constexpr int kNumTrainTimesteps = 1000;
static constexpr int kSemChannels = 16;
static constexpr int kTotalChannels = 144;
if (image_seq_len > 4300) {
float mu = a2 * image_seq_len + b2;
return mu;
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());
}
}
}
float m_200 = a2 * image_seq_len + b2;
float m_10 = a1 * image_seq_len + b1;
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 a = (m_200 - m_10) / 190.0f;
float b = m_200 - 200.0f * a;
float mu = a * n + b;
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));
return mu;
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;
}
};
// MiniT2I predicts x0 directly and integrates a linear flow ODE:
// x_{t+dt} = x_t + (x0 - x_t)/(1 - t) * dt, t in [0, 1), x0 = start = noise * 2.
// Mapping sigma = 1 - t makes the generic Euler update
// x += (x - denoised)/sigma * (sigma_next - sigma)
// exactly reproduce that step when denoised == x0. To make the generic
// `denoised = pred * c_out + x * c_skip` yield x0 from the model's raw x0
// prediction we use c_skip = 0, c_out = 1, c_in = 1. Sigmas run linearly 1 -> 0.
struct MiniT2IFlowDenoiser : public Denoiser {
float sigma_min() override {
return 0.0f;
}
float sigma_max() override {
return 1.0f;
}
float sigma_to_t(float sigma) override {
return 1.0f - sigma;
}
float t_to_sigma(float t) override {
return 1.0f - t;
}
std::vector<float> get_scalings(float sigma) override {
SD_UNUSED(sigma);
float c_skip = 0.0f;
float c_out = 1.0f;
float c_in = 1.0f;
return {c_skip, c_out, c_in};
}
sd::Tensor<float> noise_scaling(float sigma,
const sd::Tensor<float>& noise,
const sd::Tensor<float>& latent) override {
SD_UNUSED(sigma);
SD_UNUSED(latent);
// Sampling starts from x0_init = noise * 2 (see MiniT2I reference).
return noise * 2.0f;
}
sd::Tensor<float> inverse_noise_scaling(float sigma, const sd::Tensor<float>& latent) override {
SD_UNUSED(sigma);
return latent;
}
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 {
float mu = compute_empirical_mu(n, image_seq_len);
LOG_DEBUG("Flux2FlowDenoiser: set shift to %.3f", mu);
set_shift(mu);
return Denoiser::get_sigmas(n, image_seq_len, scheduler_type, version, extra_sample_args);
SD_UNUSED(image_seq_len);
SD_UNUSED(scheduler_type);
SD_UNUSED(version);
SD_UNUSED(extra_sample_args);
// Uniform t schedule 0 -> 1 => sigma 1 -> 0, matching the reference loop.
std::vector<float> sigmas;
sigmas.reserve(n + 1);
for (uint32_t i = 0; i < n; ++i) {
sigmas.push_back(1.0f - static_cast<float>(i) / static_cast<float>(n));
}
sigmas.push_back(0.0f);
return sigmas;
}
};
@ -1141,6 +1503,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) {
@ -2056,7 +2452,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

@ -0,0 +1,308 @@
#include "runtime/imatrix.h"
/* Adapted from llama.cpp (credits: Kawrakow). */
#include "core/util.h"
#include "ggml-backend.h"
#include "ggml.h"
#include "stable-diffusion.h"
#include <cmath>
#include <cstdlib>
#include <cstring>
static IMatrixCollector imatrix_collector;
IMatrixCollector& get_imatrix_collector() {
return imatrix_collector;
}
// remove any prefix and suffixes from the name
// CUDA0#blk.0.attn_k.weight#0 => blk.0.attn_k.weight
static std::string filter_tensor_name(const char* name) {
std::string wname;
const char* p = strchr(name, '#');
if (p != NULL) {
p = p + 1;
const char* q = strchr(p, '#');
if (q != NULL) {
wname = std::string(p, q - p);
} else {
wname = p;
}
} else {
wname = name;
}
return wname;
}
bool IMatrixCollector::collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data) {
GGML_UNUSED(user_data);
if (t == nullptr) {
return false;
}
if (t->op != GGML_OP_MUL_MAT && t->op != GGML_OP_MUL_MAT_ID) {
return false;
}
const struct ggml_tensor* src0 = t->src[0];
const struct ggml_tensor* src1 = t->src[1];
if (src0 == nullptr || src1 == nullptr) {
return false;
}
std::string wname = filter_tensor_name(src0->name);
// when ask is true, the scheduler wants to know if we are interested in data from this tensor
// if we return true, a follow-up call will be made with ask=false in which we can do the actual collection
if (ask) {
if (t->op == GGML_OP_MUL_MAT_ID) {
return true; // collect all indirect matrix multiplications
}
// why are small batches ignored (<16 tokens)?
// if (src1->ne[1] < 16 || src1->type != GGML_TYPE_F32) return false;
if (!(wname.substr(0, 6) == "model." || wname.substr(0, 17) == "cond_stage_model." || wname.substr(0, 14) == "text_encoders.")) {
return false;
}
return true;
}
std::lock_guard<std::mutex> lock(mutex_);
// copy the data from the GPU memory if needed
const bool is_host = src1->buffer == NULL || ggml_backend_buffer_is_host(src1->buffer);
if (!is_host) {
src1_data_.resize(ggml_nelements(src1));
ggml_backend_tensor_get(src1, src1_data_.data(), 0, ggml_nbytes(src1));
}
const float* data = is_host ? (const float*)src1->data : src1_data_.data();
// this has been adapted to the new format of storing merged experts in a single 3d tensor
// ref: https://github.com/ggml-org/llama.cpp/pull/6387
if (t->op == GGML_OP_MUL_MAT_ID) {
// ids -> [n_experts_used, n_tokens]
// src1 -> [cols, n_expert_used, n_tokens]
const ggml_tensor* ids = t->src[2];
const int n_as = static_cast<int>(src0->ne[2]);
const int n_ids = static_cast<int>(ids->ne[0]);
// the top-k selected expert ids are stored in the ids tensor
// for simplicity, always copy ids to host, because it is small
// take into account that ids is not contiguous!
GGML_ASSERT(ids->ne[1] == src1->ne[2]);
ids_.resize(ggml_nbytes(ids));
ggml_backend_tensor_get(ids, ids_.data(), 0, ggml_nbytes(ids));
auto& e = stats_[wname];
++e.ncall;
if (e.values.empty()) {
e.values.resize(src1->ne[0] * n_as, 0);
e.counts.resize(src1->ne[0] * n_as, 0);
} else if (e.values.size() != (size_t)src1->ne[0] * n_as) {
LOG_ERROR("inconsistent size for %s (%d vs %d)\n", wname.c_str(), (int)e.values.size(), (int)src1->ne[0] * n_as);
exit(1); // GGML_ABORT("fatal error");
}
// loop over all possible experts, regardless if they are used or not in the batch
for (int ex = 0; ex < n_as; ++ex) {
size_t e_start = ex * src1->ne[0];
for (int idx = 0; idx < n_ids; ++idx) {
for (int row = 0; row < (int)src1->ne[2]; ++row) {
const int excur = *(const int32_t*)(ids_.data() + row * ids->nb[1] + idx * ids->nb[0]);
GGML_ASSERT(excur >= 0 && excur < n_as); // sanity check
if (excur != ex)
continue;
const int64_t i11 = idx % src1->ne[1];
const int64_t i12 = row;
const float* x = (const float*)((const char*)data + i11 * src1->nb[1] + i12 * src1->nb[2]);
for (int j = 0; j < (int)src1->ne[0]; ++j) {
e.values[e_start + j] += x[j] * x[j];
e.counts[e_start + j]++;
if (!std::isfinite(e.values[e_start + j])) {
LOG_ERROR("%f detected in %s\n", e.values[e_start + j], wname.c_str());
exit(1);
}
}
}
}
}
} else {
auto& e = stats_[wname];
if (e.values.empty()) {
e.values.resize(src1->ne[0], 0);
e.counts.resize(src1->ne[0], 0);
} else if (e.values.size() != (size_t)src1->ne[0]) {
LOG_WARN("inconsistent size for %s (%d vs %d)\n", wname.c_str(), (int)e.values.size(), (int)src1->ne[0]);
exit(1); // GGML_ABORT("fatal error");
}
++e.ncall;
for (int row = 0; row < (int)src1->ne[1]; ++row) {
const float* x = data + row * src1->ne[0];
for (int j = 0; j < (int)src1->ne[0]; ++j) {
if (std::isfinite(x[j])) {
e.values[j] += x[j] * x[j];
e.counts[j]++;
if (!std::isfinite(e.values[j])) {
LOG_WARN("%f detected in %s\n", e.values[j], wname.c_str());
exit(1);
}
} else {
// Likely something from an attention mask?
}
}
}
}
return true;
}
bool load_imatrix(const char* imatrix_path) {
return imatrix_collector.load_imatrix(imatrix_path);
}
void save_imatrix(const char* imatrix_path) {
imatrix_collector.save_imatrix(imatrix_path);
}
static bool collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data) {
return imatrix_collector.collect_imatrix(t, ask, user_data);
}
void enable_imatrix_collection() {
sd_set_backend_eval_callback(collect_imatrix, nullptr);
}
void disable_imatrix_collection() {
sd_set_backend_eval_callback(nullptr, nullptr);
}
void IMatrixCollector::save_imatrix(std::string fname, int ncall) const {
if (ncall > 0) {
fname += ".at_";
fname += std::to_string(ncall);
}
// avoid writing imatrix entries that do not have full data
// this can happen with MoE models where some of the experts end up not being exercised by the provided training data
int n_entries = 0;
std::vector<std::string> to_store;
for (const auto& kv : stats_) {
const int n_all = static_cast<int>(kv.second.counts.size());
if (n_all == 0) {
continue;
}
int n_zeros = 0;
for (const int c : kv.second.counts) {
if (c == 0) {
n_zeros++;
}
}
if (n_zeros == n_all) {
LOG_WARN("entry '%40s' has no data - skipping\n", kv.first.c_str());
continue;
}
if (n_zeros > 0) {
LOG_WARN("entry '%40s' has partial data (%.2f%%) - skipping\n", kv.first.c_str(), 100.0f * (n_all - n_zeros) / n_all);
continue;
}
n_entries++;
to_store.push_back(kv.first);
}
if (to_store.size() < stats_.size()) {
LOG_WARN("storing only %zu out of %zu entries\n", to_store.size(), stats_.size());
}
std::ofstream out(fname, std::ios::binary);
out.write((const char*)&n_entries, sizeof(n_entries));
for (const auto& name : to_store) {
const auto& stat = stats_.at(name);
int len = static_cast<int>(name.size());
out.write((const char*)&len, sizeof(len));
out.write(name.c_str(), len);
out.write((const char*)&stat.ncall, sizeof(stat.ncall));
int nval = static_cast<int>(stat.values.size());
out.write((const char*)&nval, sizeof(nval));
if (nval > 0) {
std::vector<float> tmp(nval);
for (int i = 0; i < nval; i++) {
tmp[i] = (stat.values[i] / static_cast<float>(stat.counts[i])) * static_cast<float>(stat.ncall);
}
out.write((const char*)tmp.data(), nval * sizeof(float));
}
}
// Write the number of call the matrix was computed with
out.write((const char*)&last_call_, sizeof(last_call_));
}
bool IMatrixCollector::load_imatrix(const char* fname) {
std::ifstream in(fname, std::ios::binary);
if (!in) {
LOG_ERROR("failed to open %s\n", fname);
return false;
}
int n_entries;
in.read((char*)&n_entries, sizeof(n_entries));
if (in.fail() || n_entries < 1) {
LOG_ERROR("no data in file %s\n", fname);
return false;
}
for (int i = 0; i < n_entries; ++i) {
int len;
in.read((char*)&len, sizeof(len));
std::vector<char> name_as_vec(len + 1);
in.read((char*)name_as_vec.data(), len);
if (in.fail()) {
LOG_ERROR("failed reading name for entry %d from %s\n", i + 1, fname);
return false;
}
name_as_vec[len] = 0;
std::string name{name_as_vec.data()};
auto& e = stats_[std::move(name)];
int ncall;
in.read((char*)&ncall, sizeof(ncall));
int nval;
in.read((char*)&nval, sizeof(nval));
if (in.fail() || nval < 1) {
LOG_ERROR("failed reading number of values for entry %d\n", i);
stats_ = {};
return false;
}
if (e.values.empty()) {
e.values.resize(nval, 0);
e.counts.resize(nval, 0);
}
std::vector<float> tmp(nval);
in.read((char*)tmp.data(), nval * sizeof(float));
if (in.fail()) {
LOG_ERROR("failed reading data for entry %d\n", i);
stats_ = {};
return false;
}
// Recreate the state as expected by save_imatrix(), and correct for weighted sum.
for (int i = 0; i < nval; i++) {
e.values[i] += tmp[i];
e.counts[i] += ncall;
}
e.ncall += ncall;
}
return true;
}

View file

@ -0,0 +1,45 @@
#ifndef __SD_RUNTIME_IMATRIX_H__
#define __SD_RUNTIME_IMATRIX_H__
#include <fstream>
#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>
/* Adapted from llama.cpp (credits: Kawrakow). */
struct ggml_tensor;
struct IMatrixStats {
std::vector<float> values{};
std::vector<int> counts{};
int ncall = 0;
};
class IMatrixCollector {
private:
std::unordered_map<std::string, IMatrixStats> stats_ = {};
std::mutex mutex_;
int last_call_ = 0;
std::vector<float> src1_data_;
std::vector<char> ids_; // the expert ids from ggml_mul_mat_id
public:
IMatrixCollector() = default;
bool collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data);
void save_imatrix(std::string fname, int ncall = -1) const;
bool load_imatrix(const char* fname);
std::vector<float> get_values(const std::string& key) const {
auto it = stats_.find(key);
if (it != stats_.end()) {
return it->second.values;
} else {
return {};
}
}
};
IMatrixCollector& get_imatrix_collector();
#endif // __SD_RUNTIME_IMATRIX_H__

View file

@ -31,6 +31,7 @@
#include "model/diffusion/krea2.hpp"
#include "model/diffusion/lens.hpp"
#include "model/diffusion/ltxv.hpp"
#include "model/diffusion/minit2i.hpp"
#include "model/diffusion/mmdit.hpp"
#include "model/diffusion/model.hpp"
#include "model/diffusion/pid.hpp"
@ -85,6 +86,7 @@ const char* model_version_to_str[] = {
"Wan 2.2 I2V",
"Wan 2.2 TI2V",
"Qwen Image",
"Qwen Image Layered",
"Anima",
"Flux.2",
"Flux.2 klein",
@ -95,9 +97,11 @@ const char* model_version_to_str[] = {
"Ovis Image",
"Ernie Image",
"Lens",
"MiniT2I",
"Longcat-Image",
"PiD",
"Ideogram 4",
"SeFi-Image",
"Krea2",
"ESRGAN",
};
@ -497,7 +501,8 @@ public:
bool is_ideogram = sd_version_is_ideogram4(tempver);
bool is_boogu = sd_version_is_boogu_image(tempver);
bool is_krea2 = sd_version_is_krea2(tempver);
bool conditioner_is_llm = (is_qwenimg || iszimg || isflux2 || is_ovis || is_anima || is_ernie || is_longcat || is_lens || is_ltx || is_ideogram || is_boogu || is_krea2);
bool is_sefi = sd_version_is_sefi_image(tempver);
bool conditioner_is_llm = (is_qwenimg || iszimg || isflux2 || is_ovis || is_anima || is_ernie || is_longcat || is_lens || is_ltx || is_ideogram || is_boogu || is_krea2 || is_sefi);
bool has_llm_vision = (is_qwenimg || is_longcat || is_boogu);
//kcpp qol fallback: if a llm was loaded as t5 by mistake
@ -603,7 +608,7 @@ public:
{
to_replace = "taesd_f2.embd";
}
else if(is_wan21||is_qwenimg||is_anima||is_krea2)
else if(sd_version_uses_wan_vae(tempver))
{
to_replace = "taesd_w21.embd";
}
@ -939,7 +944,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,
@ -999,13 +1004,14 @@ public:
}
}
} else if (sd_version_is_qwen_image(version)) {
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED;
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
tensor_storage_map,
version,
"",
true,
enable_vision,
model_manager);
diffusion_model = std::make_shared<Qwen::QwenImageRunner>(backend_for(SDBackendModule::DIFFUSION),
diffusion_model = std::make_shared<Qwen::QwenImageRunner>(backend_for(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
version,
@ -1032,6 +1038,14 @@ public:
tensor_storage_map,
"model",
model_manager);
} else if (sd_version_is_minit2i(version)) {
cond_stage_model = std::make_shared<MiniT2IConditioner>(backend_for(SDBackendModule::TE),
tensor_storage_map,
model_manager);
diffusion_model = std::make_shared<MiniT2I::MiniT2IRunner>(backend_for(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model.model.net",
model_manager);
} else if (sd_version_is_anima(version)) {
cond_stage_model = std::make_shared<AnimaConditioner>(backend_for(SDBackendModule::TE),
tensor_storage_map,
@ -1140,11 +1154,7 @@ public:
}
auto create_tae = [&](bool decode_only) -> std::shared_ptr<VAE> {
if (sd_version_is_wan(version) ||
sd_version_is_qwen_image(version) ||
sd_version_is_krea2(version) ||
sd_version_is_anima(version) ||
sd_version_is_ltxav(version)) {
if (sd_version_uses_wan_vae(version) || sd_version_is_ltxav(version)) {
return std::make_shared<TinyVideoAutoEncoder>(backend_for(SDBackendModule::VAE),
tensor_storage_map,
"decoder",
@ -1181,10 +1191,7 @@ public:
false,
version,
model_manager);
} else if (sd_version_is_wan(version) ||
sd_version_is_qwen_image(version) ||
sd_version_is_krea2(version) ||
sd_version_is_anima(version)) {
} else if (sd_version_uses_wan_vae(version)) {
return std::make_shared<WAN::WanVAERunner>(backend_for(SDBackendModule::VAE),
tensor_storage_map,
"first_stage_model",
@ -1212,7 +1219,7 @@ public:
}
};
if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) {
if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) {
LOG_INFO("using FakeVAE");
first_stage_model = std::make_shared<FakeVAE>(version,
backend_for(SDBackendModule::VAE),
@ -1528,6 +1535,7 @@ public:
default_flow_shift = 3.f;
}
} else if (sd_version_is_flux(version) ||
sd_version_is_flux2(version) ||
sd_version_is_longcat(version) ||
sd_version_is_lens(version) ||
sd_version_is_ltxav(version) ||
@ -1550,8 +1558,10 @@ public:
} else if (sd_version_is_krea2(version)) {
default_flow_shift = 1.15f;
}
} else if (sd_version_is_flux2(version)) {
pred_type = FLUX2_FLOW_PRED;
} else if (sd_version_is_sefi_image(version)) {
pred_type = SEFI_FLOW_PRED;
} else if (sd_version_is_minit2i(version)) {
pred_type = MINIT2I_FLOW_PRED;
} else {
pred_type = EPS_PRED;
}
@ -1584,9 +1594,14 @@ public:
denoiser = std::make_shared<FluxFlowDenoiser>();
break;
}
case FLUX2_FLOW_PRED: {
LOG_INFO("running in Flux2 FLOW mode");
denoiser = std::make_shared<Flux2FlowDenoiser>();
case SEFI_FLOW_PRED: {
LOG_INFO("running in SeFi-Image dual-time FLOW mode");
denoiser = std::make_shared<SefiFlowDenoiser>();
break;
}
case MINIT2I_FLOW_PRED: {
LOG_INFO("running in MiniT2I FLOW mode");
denoiser = std::make_shared<MiniT2IFlowDenoiser>();
break;
}
default: {
@ -1913,7 +1928,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]);
@ -2009,7 +2033,7 @@ public:
} else if (sd_version_uses_flux_vae(version)) {
latent_rgb_proj = flux_latent_rgb_proj;
latent_rgb_bias = flux_latent_rgb_bias;
} else if (sd_version_is_wan(version) || sd_version_is_qwen_image(version) || sd_version_is_anima(version) || sd_version_is_krea2(version)) {
} else if (sd_version_uses_wan_vae(version)) {
latent_rgb_proj = wan_21_latent_rgb_proj;
latent_rgb_bias = wan_21_latent_rgb_bias;
} else {
@ -2295,12 +2319,13 @@ public:
}
int64_t last_progress_us = ggml_time_us();
sd::Tensor<float> x_t = !noise.empty()
? denoiser->noise_scaling(sigmas[0], noise, init_latent)
: init_latent;
sd::Tensor<float> denoised = x_t;
SamplePreviewContext preview = prepare_sample_preview_context();
sd::Tensor<float> x_t = !noise.empty()
? denoiser->noise_scaling(sigmas[0], noise, init_latent)
: init_latent;
sd::Tensor<float> denoised = x_t;
auto denoise = [&](const sd::Tensor<float>& x, float sigma, int step) -> sd::guidance::GuiderOutput {
if (get_cancel_flag() == SD_CANCEL_ALL) {
LOG_DEBUG("cancelling generation");
@ -2325,7 +2350,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
@ -2363,9 +2388,9 @@ public:
sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma);
std::vector<sd::Tensor<float>> controls;
DiffusionParams diffusion_params;
diffusion_params.x = &noised_input;
diffusion_params.timesteps = &timesteps_tensor;
diffusion_params.increase_ref_index = increase_ref_index;
diffusion_params.x = &noised_input;
diffusion_params.timesteps = &timesteps_tensor;
diffusion_params.ref_index_mode = Rope::ref_index_mode_from_bool(increase_ref_index);
sd::guidance::GuidanceInput step_guidance_input;
step_guidance_input.step = step;
step_guidance_input.schedule_size = sigmas.size();
@ -2395,7 +2420,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)) {
@ -2418,6 +2443,9 @@ public:
audio_length,
frame_rate,
video_positions.empty() ? nullptr : &video_positions};
} else if (sd_version_is_minit2i(version)) {
diffusion_params.extra = MiniT2IDiffusionExtra{
condition.c_vector.empty() ? nullptr : &condition.c_vector};
} else {
diffusion_params.extra = std::monostate{};
}
@ -2539,7 +2567,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) {
@ -2598,8 +2626,12 @@ public:
latent_channel = 3;
} else if (version == VERSION_CHROMA_RADIANCE) {
latent_channel = 3;
} else if (sd_version_is_minit2i(version)) {
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 {
@ -2609,6 +2641,10 @@ public:
return latent_channel;
}
int get_image_channels() const {
return version == VERSION_QWEN_IMAGE_LAYERED ? 4 : 3;
}
int get_image_seq_len(int h, int w) {
int vae_scale_factor = get_vae_scale_factor();
return (h / vae_scale_factor) * (w / vae_scale_factor);
@ -2677,7 +2713,7 @@ public:
}
sd::Tensor<float> decode_first_stage(const sd::Tensor<float>& x, bool decode_video = false) {
if (sd_version_is_pid(version)) {
if (sd_version_is_pid(version) || sd_version_is_minit2i(version)) {
return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f);
}
auto latents = first_stage_model->diffusion_to_vae_latents(x);
@ -2838,6 +2874,9 @@ const char* scheduler_to_str[] = {
"bong_tangent",
"ltx2",
"logit_normal",
"flux2",
"flux",
"beta",
};
const char* sd_scheduler_name(enum scheduler_t scheduler) {
@ -2848,6 +2887,9 @@ const char* sd_scheduler_name(enum scheduler_t scheduler) {
}
enum scheduler_t str_to_scheduler(const char* str) {
if (!strcmp(str, "normal")) {
return DISCRETE_SCHEDULER;
}
for (int i = 0; i < SCHEDULER_COUNT; i++) {
if (!strcmp(str, scheduler_to_str[i])) {
return (enum scheduler_t)i;
@ -2862,7 +2904,8 @@ const char* prediction_to_str[] = {
"edm_v",
"sd3_flow",
"flux_flow",
"flux2_flow",
"sefi_flow",
"minit2i_flow",
};
const char* sd_prediction_name(enum prediction_t prediction) {
@ -3206,6 +3249,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
sd_img_gen_params->seed = -1;
sd_img_gen_params->batch_count = 1;
sd_img_gen_params->control_strength = 0.9f;
sd_img_gen_params->qwen_image_layers = 3;
sd_img_gen_params->pm_params = {nullptr, 0, nullptr, 20.f};
sd_img_gen_params->pulid_params = {nullptr, 1.0f};
sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
@ -3232,6 +3276,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
"seed: %" PRId64
"\n"
"batch_count: %d\n"
"qwen_image_layers: %d\n"
"ref_images_count: %d\n"
"auto_resize_ref_image: %s\n"
"increase_ref_index: %s\n"
@ -3248,6 +3293,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
sd_img_gen_params->strength,
sd_img_gen_params->seed,
sd_img_gen_params->batch_count,
sd_img_gen_params->qwen_image_layers,
sd_img_gen_params->ref_images_count,
BOOL_STR(sd_img_gen_params->auto_resize_ref_image),
BOOL_STR(sd_img_gen_params->increase_ref_index),
@ -3510,9 +3556,13 @@ enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_me
return LCM_SCHEDULER;
} else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) {
return SIMPLE_SCHEDULER;
} else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_flux(sd_ctx->sd->version)) {
return FLUX_SCHEDULER;
} else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_flux2(sd_ctx->sd->version)) {
return FLUX2_SCHEDULER;
} else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ltxav(sd_ctx->sd->version)) {
return LTX2_SCHEDULER;
} else if(sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ideogram4(sd_ctx->sd->version)) {
} else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ideogram4(sd_ctx->sd->version)) {
return LOGIT_NORMAL_SCHEDULER;
}
return DISCRETE_SCHEDULER;
@ -3583,6 +3633,7 @@ struct GenerationRequest {
bool has_ref_images = false;
const sd_cache_params_t* cache_params = nullptr;
int batch_count = 1;
int qwen_image_layers = 3;
int shifted_timestep = 0;
float strength = 1.f;
float control_strength = 0.f;
@ -3608,6 +3659,7 @@ struct GenerationRequest {
diffusion_model_down_factor = sd_ctx->sd->get_diffusion_model_down_factor();
seed = sd_img_gen_params->seed;
batch_count = sd_img_gen_params->batch_count;
qwen_image_layers = std::max(0, sd_img_gen_params->qwen_image_layers);
clip_skip = sd_img_gen_params->clip_skip;
shifted_timestep = sd_img_gen_params->sample_params.shifted_timestep;
strength = sd_img_gen_params->strength;
@ -4313,6 +4365,34 @@ public:
ImageVaeAxesGuard& operator=(const ImageVaeAxesGuard&) = delete;
};
static sd::Tensor<float> ensure_image_tensor_channels(sd::Tensor<float> image, int channels) {
if (image.empty()) {
return image;
}
GGML_ASSERT(image.dim() == 4);
int64_t current_channels = image.shape()[2];
if (current_channels == channels) {
return image;
}
if (channels == 4) {
sd::Tensor<float> alpha = sd::full<float>({image.shape()[0], image.shape()[1], 1, image.shape()[3]}, 1.f);
if (current_channels == 3) {
return sd::ops::concat(image, alpha, 2);
}
if (current_channels == 1) {
sd::Tensor<float> rgb = sd::ops::concat(image, image, 2);
rgb = sd::ops::concat(rgb, image, 2);
return sd::ops::concat(rgb, alpha, 2);
}
}
if (channels == 3 && current_channels >= 3) {
return sd::ops::slice(image, 2, 0, 3);
}
GGML_ABORT("cannot convert image tensor from %lld to %d channels",
(long long)current_channels,
channels);
}
static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd_ctx_t* sd_ctx,
const sd_img_gen_params_t* sd_img_gen_params,
GenerationRequest* request,
@ -4322,6 +4402,7 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
sd::Tensor<float> init_image_tensor;
sd::Tensor<float> control_image_tensor;
sd::Tensor<float> mask_image_tensor;
int image_channels = sd_ctx->sd->get_image_channels();
if (sd_img_gen_params->init_image.data != nullptr) {
LOG_INFO("IMG2IMG");
@ -4338,7 +4419,8 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
plan->sample_steps = static_cast<int>(plan->sigmas.size() - 1);
}
init_image_tensor = sd_image_to_tensor(sd_img_gen_params->init_image, request->width, request->height);
init_image_tensor = ensure_image_tensor_channels(sd_image_to_tensor(sd_img_gen_params->init_image, request->width, request->height),
image_channels);
}
if (sd_img_gen_params->mask_image.data != nullptr) {
@ -4370,7 +4452,11 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
sd::Tensor<float> init_latent;
sd::Tensor<float> control_latent;
if (init_image_tensor.empty()) {
init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height);
if (sd_ctx->sd->version == VERSION_QWEN_IMAGE_LAYERED) {
init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height, request->qwen_image_layers + 1, true);
} else {
init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height);
}
} else {
init_latent = sd_ctx->sd->encode_first_stage(init_image_tensor);
if (init_latent.empty()) {
@ -4389,12 +4475,13 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
std::vector<sd::Tensor<float>> ref_images;
for (int i = 0; i < sd_img_gen_params->ref_images_count; i++) {
ref_images.push_back(sd_image_to_tensor(sd_img_gen_params->ref_images[i]));
ref_images.push_back(ensure_image_tensor_channels(sd_image_to_tensor(sd_img_gen_params->ref_images[i]),
image_channels));
}
if (ref_images.empty() && sd_version_is_unet_edit(sd_ctx->sd->version)) {
LOG_WARN("This model needs at least one reference image; using an empty reference");
ref_images.push_back(sd::zeros<float>({request->width, request->height, 3, 1}));
ref_images.push_back(sd::zeros<float>({request->width, request->height, image_channels, 1}));
request->guidance.img_cfg = request->guidance.txt_cfg;
request->use_img_uncond = false;
}
@ -4420,7 +4507,10 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
vae_width = round(vae_width / factor) * factor;
auto resized_ref_img = sd::ops::interpolate(ref_images[i],
{static_cast<int>(vae_width), static_cast<int>(vae_height), 3, 1});
{static_cast<int>(vae_width),
static_cast<int>(vae_height),
ref_images[i].shape()[2],
ref_images[i].shape()[3]});
LOG_DEBUG("resize vae ref image %d from %" PRId64 "x%" PRId64 " to %" PRId64 "x%" PRId64,
static_cast<int>(i),
@ -4565,6 +4655,11 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
if (request->use_uncond || request->use_high_noise_uncond) {
if (sd_version_is_ideogram4(sd_ctx->sd->version)) {
uncond.c_vector = sd::Tensor<float>::from_vector({1.0f});
} else if (sd_version_is_minit2i(sd_ctx->sd->version)) {
// MiniT2I derives the unconditional signal from the same T5 hidden
// states with a zeroed prompt mask, so no extra text encode is needed.
uncond.c_crossattn = cond.c_crossattn;
uncond.c_vector = sd::Tensor<float>::zeros_like(cond.c_vector);
} else {
bool zero_out_masked = false;
if (sd_version_is_sdxl(sd_ctx->sd->version) &&
@ -4620,7 +4715,8 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
const GenerationRequest& request,
const std::vector<sd::Tensor<float>>& final_latents) {
const std::vector<sd::Tensor<float>>& final_latents,
int* num_images_out) {
if (final_latents.empty()) {
LOG_ERROR("no latent images to decode");
return nullptr;
@ -4644,13 +4740,41 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
cancelled = true;
break;
}
int64_t t1 = ggml_time_ms();
sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(final_latents[i]);
if (image.empty()) {
LOG_ERROR("decode_first_stage failed for latent %" PRId64, i + 1);
return nullptr;
int64_t t1 = ggml_time_ms();
if (sd_ctx->sd->version == VERSION_QWEN_IMAGE_LAYERED) {
int qwen_image_latent_layers = request.qwen_image_layers + 1;
if (final_latents[i].dim() < 5 || final_latents[i].shape()[2] < qwen_image_latent_layers) {
LOG_ERROR("qwen image layered expected at least %d latent layers, got shape dim=%d",
qwen_image_latent_layers,
final_latents[i].dim());
return nullptr;
}
for (int layer_index = 0; layer_index < qwen_image_latent_layers; layer_index++) {
if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
LOG_ERROR("cancelling latent decodings");
cancelled = true;
break;
}
sd::Tensor<float> layer_latent = sd::ops::slice(final_latents[i], 2, layer_index, layer_index + 1);
layer_latent.squeeze_(2);
sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(layer_latent);
if (image.empty()) {
LOG_ERROR("decode_first_stage failed for latent %zu layer %d", i + 1, layer_index + 1);
return nullptr;
}
decoded_images.push_back(std::move(image));
}
if (cancelled) {
break;
}
} else {
sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(final_latents[i]);
if (image.empty()) {
LOG_ERROR("decode_first_stage failed for latent %" PRId64, i + 1);
return nullptr;
}
decoded_images.push_back(std::move(image));
}
decoded_images.push_back(std::move(image));
int64_t t2 = ggml_time_ms();
LOG_INFO("latent %zu decoded, taking %.2fs", i + 1, (t2 - t1) * 1.0f / 1000);
}
@ -4662,11 +4786,14 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
return nullptr;
}
sd_image_t* result_images = (sd_image_t*)calloc(request.batch_count, sizeof(sd_image_t));
int image_count = static_cast<int>(decoded_images.size());
sd_image_t* result_images = (sd_image_t*)calloc(image_count, sizeof(sd_image_t));
if (result_images == nullptr) {
return nullptr;
}
memset(result_images, 0, request.batch_count * sizeof(sd_image_t));
if (num_images_out != nullptr) {
*num_images_out = image_count;
}
for (size_t i = 0; i < decoded_images.size(); i++) {
result_images[i] = tensor_to_sd_image(decoded_images[i]);
@ -4859,9 +4986,18 @@ static std::vector<float> make_hires_sigma_schedule(sd_ctx_t* sd_ctx,
sigmas.end());
}
SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params) {
SD_API bool generate_image(sd_ctx_t* sd_ctx,
const sd_img_gen_params_t* sd_img_gen_params,
sd_image_t** images_out,
int* num_images_out) {
if (images_out != nullptr) {
*images_out = nullptr;
}
if (num_images_out != nullptr) {
*num_images_out = 0;
}
if (sd_ctx == nullptr || sd_img_gen_params == nullptr) {
return nullptr;
return false;
}
sd_ctx->sd->reset_cancel_flag();
@ -4884,7 +5020,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
&request,
&plan);
if (!latents_opt.has_value()) {
return nullptr;
return false;
}
ImageGenerationLatents latents = std::move(*latents_opt);
@ -4894,7 +5030,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
&plan,
&latents);
if (!embeds_opt.has_value()) {
return nullptr;
return false;
}
ImageGenerationEmbeds embeds = std::move(*embeds_opt);
@ -4904,7 +5040,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
sd_cancel_mode_t cancel = sd_ctx->sd->get_cancel_flag();
if (cancel == SD_CANCEL_ALL) {
LOG_ERROR("cancelling generation");
return nullptr;
return false;
}
if (cancel == SD_CANCEL_NEW_LATENTS) {
LOG_INFO("cancelling new latent generation, returning %zu/%d completed latents",
@ -4956,7 +5092,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
b + 1,
request.batch_count,
(sampling_end - sampling_start) * 1.0f / 1000);
return nullptr;
return false;
}
int64_t denoise_end = ggml_time_ms();
LOG_INFO("generating %zu latent images completed, taking %.2fs",
@ -4964,13 +5100,13 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
(denoise_end - denoise_start) * 1.0f / 1000);
if (final_latents.empty()) {
LOG_ERROR("no latent images generated");
return nullptr;
return false;
}
if (request.hires.enabled && request.hires.target_width > 0) {
if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
LOG_ERROR("cancelling generation before hires fix");
return nullptr;
return false;
}
LOG_INFO("hires fix: upscaling to %dx%d", request.hires.target_width, request.hires.target_height);
@ -4978,7 +5114,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
LOG_ERROR("cancelling generation before hires model load");
return nullptr;
return false;
}
LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path);
hires_upscaler = std::make_unique<UpscalerGGML>(sd_ctx->sd->n_threads,
@ -4991,7 +5127,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
if (!hires_upscaler->load_from_file(request.hires.model_path,
sd_ctx->sd->n_threads)) {
LOG_ERROR("load hires model upscaler failed");
return nullptr;
return false;
}
}
@ -5015,7 +5151,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
for (int b = 0; b < (int)final_latents.size(); b++) {
if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
LOG_ERROR("cancelling generation during hires fix");
return nullptr;
return false;
}
int64_t cur_seed = request.seed + b;
sd_ctx->sd->rng->manual_seed(cur_seed);
@ -5026,7 +5162,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
request,
hires_upscaler.get());
if (upscaled.empty()) {
return nullptr;
return false;
}
sd::Tensor<float> noise = sd::randn_like<float>(upscaled, sd_ctx->sd->rng);
@ -5080,7 +5216,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
b + 1,
(int)final_latents.size(),
(hires_sample_end - hires_sample_start) * 1.0f / 1000);
return nullptr;
return false;
}
int64_t hires_denoise_end = ggml_time_ms();
LOG_INFO("hires fix completed, taking %.2fs", (hires_denoise_end - hires_denoise_start) * 1.0f / 1000);
@ -5088,16 +5224,25 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
final_latents = std::move(hires_final_latents);
}
auto result = decode_image_outputs(sd_ctx, request, final_latents);
int num_images = 0;
auto result = decode_image_outputs(sd_ctx, request, final_latents, &num_images);
if (result == nullptr) {
return nullptr;
return false;
}
sd_ctx->sd->lora_stat();
int64_t t1 = ggml_time_ms();
LOG_INFO("generate_image completed in %.2fs", (t1 - t0) * 1.0f / 1000);
return result;
if (num_images_out != nullptr) {
*num_images_out = num_images;
}
if (images_out != nullptr) {
*images_out = result;
} else {
free_sd_images(result, num_images);
}
return true;
}
static std::optional<ImageGenerationLatents> prepare_video_generation_latents(sd_ctx_t* sd_ctx,

View file

@ -4,6 +4,7 @@
#include "model_loader.h"
#include "stable-diffusion.h"
#include <cstdlib>
#include <utility>
UpscalerGGML::UpscalerGGML(int n_threads,
@ -198,8 +199,41 @@ upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str,
return upscaler_ctx;
}
sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_t upscale_factor) {
return upscaler_ctx->upscaler->upscale(input_image, upscale_factor);
bool upscale(upscaler_ctx_t* upscaler_ctx,
sd_image_t input_image,
uint32_t upscale_factor,
sd_image_t** images_out,
int* num_images_out) {
if (images_out != nullptr) {
*images_out = nullptr;
}
if (num_images_out != nullptr) {
*num_images_out = 0;
}
if (upscaler_ctx == nullptr || upscaler_ctx->upscaler == nullptr) {
return false;
}
sd_image_t* result_images = (sd_image_t*)calloc(1, sizeof(sd_image_t));
if (result_images == nullptr) {
return false;
}
result_images[0] = upscaler_ctx->upscaler->upscale(input_image, upscale_factor);
if (result_images[0].data == nullptr) {
free(result_images);
return false;
}
if (num_images_out != nullptr) {
*num_images_out = 1;
}
if (images_out != nullptr) {
*images_out = result_images;
} else {
free_sd_images(result_images, 1);
}
return true;
}
int get_upscale_factor(upscaler_ctx_t* upscaler_ctx) {