From 6482a596e1d391539b08b8c40cd555e4fbd372b2 Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Fri, 3 Jul 2026 04:28:51 -0300 Subject: [PATCH] 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 --- Makefile | 24 +- otherarch/sdcpp/examples/cli/main.cpp | 85 ++- otherarch/sdcpp/examples/common/common.cpp | 31 +- otherarch/sdcpp/examples/common/common.h | 1 + otherarch/sdcpp/include/stable-diffusion.h | 39 +- otherarch/sdcpp/sdtype_adapter.cpp | 199 +++--- .../sdcpp/src/conditioning/conditioner.hpp | 111 +++- otherarch/sdcpp/src/convert.cpp | 85 ++- otherarch/sdcpp/src/core/ggml_extend.hpp | 50 +- .../sdcpp/src/core/ggml_extend_backend.cpp | 144 ++++- .../sdcpp/src/core/ggml_extend_backend.h | 5 + otherarch/sdcpp/src/core/util.cpp | 16 + otherarch/sdcpp/src/core/util.h | 3 + otherarch/sdcpp/src/model.h | 30 +- otherarch/sdcpp/src/model/common/rope.hpp | 144 +++-- otherarch/sdcpp/src/model/diffusion/anima.hpp | 9 +- .../sdcpp/src/model/diffusion/control.hpp | 2 +- otherarch/sdcpp/src/model/diffusion/dit.hpp | 24 +- .../sdcpp/src/model/diffusion/ernie_image.hpp | 4 +- otherarch/sdcpp/src/model/diffusion/flux.hpp | 48 +- .../sdcpp/src/model/diffusion/hidream_o1.hpp | 2 +- .../sdcpp/src/model/diffusion/minit2i.hpp | 611 ++++++++++++++++++ otherarch/sdcpp/src/model/diffusion/model.hpp | 12 +- .../sdcpp/src/model/diffusion/qwen_image.hpp | 120 +++- .../sdcpp/src/model/diffusion/sefi_image.hpp | 91 +++ .../sdcpp/src/model/diffusion/z_image.hpp | 12 +- otherarch/sdcpp/src/model/te/clip.hpp | 2 +- otherarch/sdcpp/src/model/te/llm.hpp | 4 +- otherarch/sdcpp/src/model/te/t5.hpp | 59 +- otherarch/sdcpp/src/model/upscaler/esrgan.hpp | 2 +- .../model/upscaler/ltx_latent_upscaler.hpp | 2 +- .../sdcpp/src/model/vae/auto_encoder_kl.hpp | 7 +- .../sdcpp/src/model/vae/ltx_audio_vae.hpp | 6 +- otherarch/sdcpp/src/model/vae/tae.hpp | 17 +- otherarch/sdcpp/src/model/vae/vae.hpp | 4 +- otherarch/sdcpp/src/model/vae/wan_vae.hpp | 265 ++++++-- otherarch/sdcpp/src/model_loader.cpp | 24 +- otherarch/sdcpp/src/name_conversion.cpp | 76 ++- otherarch/sdcpp/src/runtime/denoiser.hpp | 458 ++++++++++++- otherarch/sdcpp/src/runtime/imatrix.cpp | 308 +++++++++ otherarch/sdcpp/src/runtime/imatrix.h | 45 ++ otherarch/sdcpp/src/stable-diffusion.cpp | 277 ++++++-- otherarch/sdcpp/src/upscaler.cpp | 38 +- 43 files changed, 3009 insertions(+), 487 deletions(-) create mode 100644 otherarch/sdcpp/src/model/diffusion/minit2i.hpp create mode 100644 otherarch/sdcpp/src/model/diffusion/sefi_image.hpp create mode 100644 otherarch/sdcpp/src/runtime/imatrix.cpp create mode 100644 otherarch/sdcpp/src/runtime/imatrix.h diff --git a/Makefile b/Makefile index d09f5d2a4..ae2351e5b 100644 --- a/Makefile +++ b/Makefile @@ -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 diff --git a/otherarch/sdcpp/examples/cli/main.cpp b/otherarch/sdcpp/examples/cli/main.cpp index 84e9e7853..7892d5213 100644 --- a/otherarch/sdcpp/examples/cli/main.cpp +++ b/otherarch/sdcpp/examples/cli/main.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include #include @@ -53,6 +54,9 @@ struct SDCliParams { bool metadata_brief = false; bool metadata_all = false; + std::string imatrix_out; + std::vector 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; diff --git a/otherarch/sdcpp/examples/common/common.cpp b/otherarch/sdcpp/examples/common/common.cpp index 744005af9..ac1cf32f6 100644 --- a/otherarch/sdcpp/examples/common/common.cpp +++ b/otherarch/sdcpp/examples/common/common.cpp @@ -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" diff --git a/otherarch/sdcpp/examples/common/common.h b/otherarch/sdcpp/examples/common/common.h index e7c25015b..941fa3317 100644 --- a/otherarch/sdcpp/examples/common/common.h +++ b/otherarch/sdcpp/examples/common/common.h @@ -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; diff --git a/otherarch/sdcpp/include/stable-diffusion.h b/otherarch/sdcpp/include/stable-diffusion.h index 7058852cc..d5bace723 100644 --- a/otherarch/sdcpp/include/stable-diffusion.h +++ b/otherarch/sdcpp/include/stable-diffusion.h @@ -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); diff --git a/otherarch/sdcpp/sdtype_adapter.cpp b/otherarch/sdcpp/sdtype_adapter.cpp index 5dff78031..73c0b73f7 100644 --- a/otherarch/sdcpp/sdtype_adapter.cpp +++ b/otherarch/sdcpp/sdtype_adapter.cpp @@ -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, ¶ms); + if (!generate_image(sd_ctx, ¶ms, &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, ¶ms); + if (!generate_image(sd_ctx, ¶ms, &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 == "") { diff --git a/otherarch/sdcpp/src/conditioning/conditioner.hpp b/otherarch/sdcpp/src/conditioning/conditioner.hpp index e037fe76b..d63303a82 100644 --- a/otherarch/sdcpp/src/conditioning/conditioner.hpp +++ b/otherarch/sdcpp/src/conditioning/conditioner.hpp @@ -1,4 +1,4 @@ -#ifndef __SD_CONDITIONING_CONDITIONER_HPP__ +#ifndef __SD_CONDITIONING_CONDITIONER_HPP__ #define __SD_CONDITIONING_CONDITIONER_HPP__ #include @@ -1378,6 +1378,101 @@ struct T5CLIPEmbedder : public Conditioner { } }; +struct MiniT2IConditioner : public Conditioner { + T5UniGramTokenizer tokenizer; + std::shared_ptr t5; + size_t prompt_length = 256; + + MiniT2IConditioner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + std::shared_ptr 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(backend, tensor_storage_map, "text_encoders.t5xxl.transformer", false, weight_manager); + } + + void get_param_tensors(std::map& 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& 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::zeros({1024, static_cast(prompt_length)}); + result.c_vector = sd::Tensor::zeros({static_cast(prompt_length)}); + return result; + } + + std::vector tokens = tokenizer.encode(conditioner_params.text); + if (tokens.size() > prompt_length) { + tokens.resize(prompt_length); + } + std::vector mask(tokens.size(), 1.0f); + while (tokens.size() < prompt_length) { + tokens.push_back(tokenizer.PAD_TOKEN_ID); + mask.push_back(0.0f); + } + + sd::Tensor input_ids({static_cast(tokens.size())}, tokens); + std::vector 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 hidden_states = t5->compute(n_threads, + input_ids, + sd::Tensor::from_vector(t5_mask), + false, + true, + true); + GGML_ASSERT(!hidden_states.empty()); + result.c_crossattn = std::move(hidden_states); + result.c_vector = sd::Tensor::from_vector(mask); + return result; + } +}; + struct AnimaConditioner : public Conditioner { std::shared_ptr 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(prompt.size()); prompt += "<|im_end|>\n<|im_start|>assistant\n\n\n\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(prompt.size()); + prompt += conditioner_params.text; + prompt_attn_range.second = static_cast(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; diff --git a/otherarch/sdcpp/src/convert.cpp b/otherarch/sdcpp/src/convert.cpp index 27d377ec0..0b7fe2cfb 100644 --- a/otherarch/sdcpp/src/convert.cpp +++ b/otherarch/sdcpp/src/convert.cpp @@ -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); +} diff --git a/otherarch/sdcpp/src/core/ggml_extend.hpp b/otherarch/sdcpp/src/core/ggml_extend.hpp index 9883103e2..48a7ad740 100644 --- a/otherarch/sdcpp/src/core/ggml_extend.hpp +++ b/otherarch/sdcpp/src/core/ggml_extend.hpp @@ -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: [OC，IC, 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), diff --git a/otherarch/sdcpp/src/core/ggml_extend_backend.cpp b/otherarch/sdcpp/src/core/ggml_extend_backend.cpp index f3e2cceba..f29bdb696 100644 --- a/otherarch/sdcpp/src/core/ggml_extend_backend.cpp +++ b/otherarch/sdcpp/src/core/ggml_extend_backend.cpp @@ -9,6 +9,7 @@ #include #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) { diff --git a/otherarch/sdcpp/src/core/ggml_extend_backend.h b/otherarch/sdcpp/src/core/ggml_extend_backend.h index 9aecf97c0..19b71d432 100644 --- a/otherarch/sdcpp/src/core/ggml_extend_backend.h +++ b/otherarch/sdcpp/src/core/ggml_extend_backend.h @@ -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); diff --git a/otherarch/sdcpp/src/core/util.cpp b/otherarch/sdcpp/src/core/util.cpp index 05c308d9b..cf1460cb9 100644 --- a/otherarch/sdcpp/src/core/util.cpp +++ b/otherarch/sdcpp/src/core/util.cpp @@ -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, 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; } diff --git a/otherarch/sdcpp/src/core/util.h b/otherarch/sdcpp/src/core/util.h index ec4c45595..07968cbcd 100644 --- a/otherarch/sdcpp/src/core/util.h +++ b/otherarch/sdcpp/src/core/util.h @@ -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); diff --git a/otherarch/sdcpp/src/model.h b/otherarch/sdcpp/src/model.h index cce309138..2de82006a 100644 --- a/otherarch/sdcpp/src/model.h +++ b/otherarch/sdcpp/src/model.h @@ -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; } diff --git a/otherarch/sdcpp/src/model/common/rope.hpp b/otherarch/sdcpp/src/model/common/rope.hpp index 2e21ef7c2..d6cbb2e6a 100644 --- a/otherarch/sdcpp/src/model/common/rope.hpp +++ b/otherarch/sdcpp/src/model/common/rope.hpp @@ -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 __STATIC_INLINE__ std::vector linspace(T start, T end, int num) { std::vector result(num); @@ -346,7 +356,7 @@ namespace Rope { int axes_dim_num, int start_index, const std::vector& 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(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 txt_arange_dims, const std::vector& 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 txt_arange_dims, const std::vector& 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> wrap_dims; @@ -481,17 +493,64 @@ namespace Rope { return embed_nd(ids, bs, static_cast(theta), axes_dim, wrap_dims); } - __STATIC_INLINE__ std::vector> gen_qwen_image_ids(int h, + __STATIC_INLINE__ std::vector> 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> vid_ids(t_len * h_len * w_len, std::vector(3, 0.0)); + + if (scale_rope) { + h_offset -= h_len / 2; + w_offset -= w_len / 2; + } + + std::vector t_ids = linspace(1.f * t_offset, 1.f * t_len - 1 + t_offset, t_len); + std::vector h_ids = linspace(1.f * h_offset, 1.f * h_len - 1 + h_offset, h_len); + std::vector w_ids = linspace(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> vid_ids_repeated(bs * vid_ids.size(), std::vector(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> gen_qwen_image_ids(int t, + int h, int w, int patch_size, int bs, int context_len, const std::vector& 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(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(1.f * txt_id_start, 1.f * txt_id_start + context_len - 1, context_len); std::vector> txt_ids_repeated(bs * context_len, std::vector(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 gen_qwen_image_pe(int h, + __STATIC_INLINE__ std::vector gen_qwen_image_pe(int t, + int h, int w, int patch_size, int bs, int context_len, const std::vector& ref_latents, - bool increase_ref_index, + RefIndexMode ref_index_mode, int theta, bool circular_h, bool circular_w, const std::vector& axes_dim) { - std::vector> ids = gen_qwen_image_ids(h, w, patch_size, bs, context_len, ref_latents, increase_ref_index); + std::vector> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode); std::vector> 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(total_tokens / bs, 0)); size_t cursor = context_len; // ignore text tokens - const size_t img_tokens = static_cast(h_len) * static_cast(w_len); + const size_t img_tokens = static_cast(t) * static_cast(h_len) * static_cast(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(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage); } - __STATIC_INLINE__ std::vector> 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> vid_ids(t_len * h_len * w_len, std::vector(3, 0.0)); - - std::vector t_ids = linspace(1.f * t_offset, 1.f * t_len - 1 + t_offset, t_len); - std::vector h_ids = linspace(1.f * h_offset, 1.f * h_len - 1 + h_offset, h_len); - std::vector w_ids = linspace(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> vid_ids_repeated(bs * vid_ids.size(), std::vector(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 gen_wan_pe(int t, int h, @@ -785,7 +806,8 @@ namespace Rope { int context_len, int seq_multi_of, const std::vector& 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>(bs * padded_context_len, std::vector(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& ref_latents, - bool increase_ref_index, + RefIndexMode ref_index_mode, int theta, bool circular_h, bool circular_w, const std::vector& axes_dim) { - std::vector> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, increase_ref_index); + std::vector> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode); std::vector> wrap_dims; if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { int pad_h = (patch_size - (h % patch_size)) % patch_size; diff --git a/otherarch/sdcpp/src/model/diffusion/anima.hpp b/otherarch/sdcpp/src/model/diffusion/anima.hpp index 504904d41..1c1ec2ed2 100644 --- a/otherarch/sdcpp/src/model/diffusion/anima.hpp +++ b/otherarch/sdcpp/src/model/diffusion/anima.hpp @@ -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); diff --git a/otherarch/sdcpp/src/model/diffusion/control.hpp b/otherarch/sdcpp/src/model/diffusion/control.hpp index 57e3616f2..eeb8f5109 100644 --- a/otherarch/sdcpp/src/model/diffusion/control.hpp +++ b/otherarch/sdcpp/src/model/diffusion/control.hpp @@ -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" diff --git a/otherarch/sdcpp/src/model/diffusion/dit.hpp b/otherarch/sdcpp/src/model/diffusion/dit.hpp index c18518564..a88e7546d 100644 --- a/otherarch/sdcpp/src/model/diffusion/dit.hpp +++ b/otherarch/sdcpp/src/model/diffusion/dit.hpp @@ -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] diff --git a/otherarch/sdcpp/src/model/diffusion/ernie_image.hpp b/otherarch/sdcpp/src/model/diffusion/ernie_image.hpp index 0427b3b38..12fcada59 100644 --- a/otherarch/sdcpp/src/model/diffusion/ernie_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/ernie_image.hpp @@ -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; } diff --git a/otherarch/sdcpp/src/model/diffusion/flux.hpp b/otherarch/sdcpp/src/model/diffusion/flux.hpp index b5e6c63bf..e3cd6ba7c 100644 --- a/otherarch/sdcpp/src/model/diffusion/flux.hpp +++ b/otherarch/sdcpp/src/model/diffusion/flux.hpp @@ -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(config.in_dim, config.hidden_size); + } else if (config.is_sefi) { + blocks["dual_time_embed"] = std::make_shared(256, config.hidden_size); } else { blocks["time_in"] = std::make_shared(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(img->ne[1]), 0, 0, 0); } + } else if (config.is_sefi) { + auto dual_time_embed = std::dynamic_pointer_cast(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(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& y_tensor = {}, const sd::Tensor& guidance_tensor = {}, const std::vector>& ref_latents_tensor = {}, - bool increase_ref_index = false, + Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED, std::vector skip_layers = {}, const sd::Tensor& 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 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(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& y = {}, const sd::Tensor& guidance = {}, const std::vector>& ref_latents = {}, - bool increase_ref_index = false, + Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED, std::vector skip_layers = std::vector(), const sd::Tensor& 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(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()); diff --git a/otherarch/sdcpp/src/model/diffusion/hidream_o1.hpp b/otherarch/sdcpp/src/model/diffusion/hidream_o1.hpp index 559f61bcf..9d3df0391 100644 --- a/otherarch/sdcpp/src/model/diffusion/hidream_o1.hpp +++ b/otherarch/sdcpp/src/model/diffusion/hidream_o1.hpp @@ -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 diff --git a/otherarch/sdcpp/src/model/diffusion/minit2i.hpp b/otherarch/sdcpp/src/model/diffusion/minit2i.hpp new file mode 100644 index 000000000..284661054 --- /dev/null +++ b/otherarch/sdcpp/src/model/diffusion/minit2i.hpp @@ -0,0 +1,611 @@ +#ifndef __SD_MODEL_DIFFUSION_MINIT2I_HPP__ +#define __SD_MODEL_DIFFUSION_MINIT2I_HPP__ + +#include +#include +#include +#include +#include +#include +#include + +#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(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(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(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 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 out(static_cast(grid_size) * grid_size * dim); + std::vector omega(quarter); + for (int i = 0; i < quarter; ++i) { + omega[i] = 1.0f / std::pow(10000.0f, static_cast(i) / static_cast(quarter)); + } + for (int y = 0; y < grid_size; ++y) { + for (int x = 0; x < grid_size; ++x) { + size_t base = static_cast(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 make_text_rope(int length, int head_dim) { + return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast(length - 1), length), head_dim, 10000.f)); + } + + inline std::vector 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 out(static_cast(length) * (head_dim / 2) * 4); + std::vector freqs(quarter); + for (int i = 0; i < quarter; ++i) { + freqs[i] = 1.0f / std::pow(10000.0f, static_cast(2 * i) / static_cast(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(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(in_features, hidden_dim, false); + blocks["w3"] = std::make_shared(in_features, hidden_dim, false); + blocks["w2"] = std::make_shared(hidden_dim, in_features, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto w1 = std::dynamic_pointer_cast(blocks["w1"]); + auto w3 = std::dynamic_pointer_cast(blocks["w3"]); + auto w2 = std::dynamic_pointer_cast(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(in_channels, + pca_channels, + std::pair{static_cast(patch_size), static_cast(patch_size)}, + std::pair{static_cast(patch_size), static_cast(patch_size)}, + std::pair{0, 0}, + std::pair{1, 1}, + false); + blocks["proj2"] = std::make_shared(pca_channels, + hidden_size, + std::pair{1, 1}, + std::pair{1, 1}, + std::pair{0, 0}, + std::pair{1, 1}, + true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto proj1 = std::dynamic_pointer_cast(blocks["proj1"]); + auto proj2 = std::dynamic_pointer_cast(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(frequency_embedding_size, hidden_size, true, true); + blocks["mlp.2"] = std::make_shared(hidden_size, hidden_size, true, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) { + auto mlp_0 = std::dynamic_pointer_cast(blocks["mlp.0"]); + auto mlp_2 = std::dynamic_pointer_cast(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 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(hidden_size, 1e-6f); + blocks["norm2"] = std::make_shared(hidden_size, 1e-6f); + blocks["qkv"] = std::make_shared(hidden_size, inner_dim * 3, true); + blocks["attn_proj"] = std::make_shared(inner_dim, hidden_size, true); + blocks["mlp"] = std::make_shared(hidden_size, static_cast(hidden_size * mlp_ratio)); + blocks["q_norm"] = std::make_shared(head_dim, 1e-6f); + blocks["k_norm"] = std::make_shared(head_dim, 1e-6f); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* pe) { + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); + auto attn_proj = std::dynamic_pointer_cast(blocks["attn_proj"]); + auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); + auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); + auto k_norm = std::dynamic_pointer_cast(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(hidden_size, 1e-6f); + blocks["img_norm2"] = std::make_shared(hidden_size, 1e-6f); + blocks["txt_norm1"] = std::make_shared(txt_hidden_size, 1e-6f); + blocks["txt_norm2"] = std::make_shared(txt_hidden_size, 1e-6f); + blocks["img_qkv"] = std::make_shared(hidden_size, inner_dim * 3, true); + blocks["txt_qkv"] = std::make_shared(txt_hidden_size, inner_dim * 3, true); + blocks["q_norm"] = std::make_shared(head_dim, 1e-6f); + blocks["k_norm"] = std::make_shared(head_dim, 1e-6f); + blocks["img_attn_proj"] = std::make_shared(inner_dim, hidden_size, true); + blocks["txt_attn_proj"] = std::make_shared(inner_dim, txt_hidden_size, true); + blocks["img_mlp"] = std::make_shared(hidden_size, static_cast(hidden_size * mlp_ratio)); + blocks["txt_mlp"] = std::make_shared(txt_hidden_size, static_cast(txt_hidden_size * mlp_ratio)); + } + + std::pair forward(GGMLRunnerContext* ctx, + ggml_tensor* img, + ggml_tensor* txt, + ggml_tensor* pe) { + auto img_norm1 = std::dynamic_pointer_cast(blocks["img_norm1"]); + auto img_norm2 = std::dynamic_pointer_cast(blocks["img_norm2"]); + auto txt_norm1 = std::dynamic_pointer_cast(blocks["txt_norm1"]); + auto txt_norm2 = std::dynamic_pointer_cast(blocks["txt_norm2"]); + auto img_qkv_p = std::dynamic_pointer_cast(blocks["img_qkv"]); + auto txt_qkv_p = std::dynamic_pointer_cast(blocks["txt_qkv"]); + auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); + auto k_norm = std::dynamic_pointer_cast(blocks["k_norm"]); + auto img_proj = std::dynamic_pointer_cast(blocks["img_attn_proj"]); + auto txt_proj = std::dynamic_pointer_cast(blocks["txt_attn_proj"]); + auto img_mlp = std::dynamic_pointer_cast(blocks["img_mlp"]); + auto txt_mlp = std::dynamic_pointer_cast(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(hidden_size, 1e-6f); + blocks["linear"] = std::make_shared(hidden_size, patch_size * patch_size * out_channels, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto norm_final = std::dynamic_pointer_cast(blocks["norm_final"]); + auto linear = std::dynamic_pointer_cast(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(config.patch_size, config.in_channels, config.pca_channels, config.hidden_size); + blocks["txt_embedder"] = std::make_shared(config.txt_input_size, config.txt_hidden_size, false); + blocks["t_embedder"] = std::make_shared(config.cond_vec_size); + blocks["pooled_embedder"] = std::make_shared(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(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(config.hidden_size, config.txt_hidden_size, config.num_heads, config.head_dim, config.mlp_ratio); + } + blocks["final_layer"] = std::make_shared(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(blocks["img_embedder"]); + auto txt_embedder = std::dynamic_pointer_cast(blocks["txt_embedder"]); + auto final_layer = std::dynamic_pointer_cast(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(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(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(config.patch_size), static_cast(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 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& 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(img_side), static_cast(config.hidden_size)); + auto txt_pe_vec = make_text_rope(static_cast(txt_len), static_cast(config.head_dim)); + auto img_pe_vec = make_vision_rope(static_cast(img_side), static_cast(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(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& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor, + const sd::Tensor& 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 compute(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timesteps, + const sd::Tensor& context, + const sd::Tensor& mask) { + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(x, timesteps, context, mask); + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + } + + sd::Tensor 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(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__ diff --git a/otherarch/sdcpp/src/model/diffusion/model.hpp b/otherarch/sdcpp/src/model/diffusion/model.hpp index 67f0fee02..8ef00023e 100644 --- a/otherarch/sdcpp/src/model/diffusion/model.hpp +++ b/otherarch/sdcpp/src/model/diffusion/model.hpp @@ -1,4 +1,4 @@ -#ifndef __SD_MODEL_DIFFUSION_MODEL_HPP__ +#ifndef __SD_MODEL_DIFFUSION_MODEL_HPP__ #define __SD_MODEL_DIFFUSION_MODEL_HPP__ #include @@ -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* video_positions = nullptr; }; +struct MiniT2IDiffusionExtra { + const sd::Tensor* mask = nullptr; +}; + using DiffusionExtraParams = std::variant; + LTXAVDiffusionExtra, + MiniT2IDiffusionExtra>; struct DiffusionParams { const sd::Tensor* x = nullptr; @@ -68,7 +74,7 @@ struct DiffusionParams { const sd::Tensor* c_concat = nullptr; const sd::Tensor* y = nullptr; const std::vector>* ref_latents = nullptr; - bool increase_ref_index = false; + Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED; DiffusionExtraParams extra = std::monostate{}; }; diff --git a/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp b/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp index aecd8bcef..5dabbb176 100644 --- a/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp @@ -4,6 +4,7 @@ #include #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 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(new TimestepEmbedding(256, embedding_dim)); + if (use_additional_t_cond) { + blocks["addition_t_embedding"] = std::shared_ptr(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(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(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(new QwenTimestepProjEmbeddings(inner_dim)); + blocks["time_text_embed"] = std::shared_ptr(new QwenTimestepProjEmbeddings(inner_dim, config.use_additional_t_cond)); blocks["txt_norm"] = std::shared_ptr(new RMSNorm(config.joint_attention_dim, 1e-6f)); blocks["img_in"] = std::shared_ptr(new Linear(config.in_channels, inner_dim)); blocks["txt_in"] = std::shared_ptr(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(blocks["norm_out"]); auto proj_out = std::dynamic_pointer_cast(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 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((H + pad_h) / config.patch_size); + int w_len = static_cast((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 pe_vec; std::vector modulate_index_vec; + std::vector 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 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& timesteps_tensor, const sd::Tensor& context_tensor, const std::vector>& 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 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(x->ne[1]), + int batch_size = static_cast(x->ne[3]); + int time_len = 1; + if (x_tensor.dim() == 5) { + time_len = static_cast(x_tensor.shape()[2]); + batch_size = static_cast(x_tensor.shape()[4]); + } + + ggml_tensor* addition_t_cond = nullptr; + if (version == VERSION_QWEN_IMAGE_LAYERED) { + additional_t_cond_vec.assign(static_cast(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(x->ne[1]), static_cast(x->ne[0]), config.patch_size, - static_cast(x->ne[3]), + batch_size, static_cast(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& timesteps, const sd::Tensor& context, const std::vector>& 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(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()); diff --git a/otherarch/sdcpp/src/model/diffusion/sefi_image.hpp b/otherarch/sdcpp/src/model/diffusion/sefi_image.hpp new file mode 100644 index 000000000..271919882 --- /dev/null +++ b/otherarch/sdcpp/src/model/diffusion/sefi_image.hpp @@ -0,0 +1,91 @@ +#ifndef __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__ +#define __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__ + +#include + +#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(new Linear(in_channels, time_embed_dim, false)); + blocks["linear_2"] = std::shared_ptr(new Linear(time_embed_dim, time_embed_dim, false)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* sample) { + auto linear_1 = std::dynamic_pointer_cast(blocks["linear_1"]); + auto linear_2 = std::dynamic_pointer_cast(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(in_channels, half_dim); + blocks["texture_embedder"] = std::make_shared(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(blocks["semantic_embedder"]); + auto texture_embedder = std::dynamic_pointer_cast(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__ diff --git a/otherarch/sdcpp/src/model/diffusion/z_image.hpp b/otherarch/sdcpp/src/model/diffusion/z_image.hpp index d23c28563..362192f19 100644 --- a/otherarch/sdcpp/src/model/diffusion/z_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/z_image.hpp @@ -575,7 +575,7 @@ namespace ZImage { const sd::Tensor& timesteps_tensor, const sd::Tensor& context_tensor, const std::vector>& 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(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& timesteps, const sd::Tensor& context, const std::vector>& 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(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()); diff --git a/otherarch/sdcpp/src/model/te/clip.hpp b/otherarch/sdcpp/src/model/te/clip.hpp index 6dc8a947b..2fde3de79 100644 --- a/otherarch/sdcpp/src/model/te/clip.hpp +++ b/otherarch/sdcpp/src/model/te/clip.hpp @@ -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" diff --git a/otherarch/sdcpp/src/model/te/llm.hpp b/otherarch/sdcpp/src/model/te/llm.hpp index 12daf5637..3c5e9bb54 100644 --- a/otherarch/sdcpp/src/model/te/llm.hpp +++ b/otherarch/sdcpp/src/model/te/llm.hpp @@ -1,4 +1,4 @@ -#ifndef __SD_MODEL_TE_LLM_HPP__ +#ifndef __SD_MODEL_TE_LLM_HPP__ #define __SD_MODEL_TE_LLM_HPP__ #include @@ -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) { diff --git a/otherarch/sdcpp/src/model/te/t5.hpp b/otherarch/sdcpp/src/model/te/t5.hpp index 23da08222..6d2326f94 100644 --- a/otherarch/sdcpp/src/model/te/t5.hpp +++ b/otherarch/sdcpp/src/model/te/t5.hpp @@ -1,4 +1,4 @@ -#ifndef __SD_MODEL_TE_T5_HPP__ +#ifndef __SD_MODEL_TE_T5_HPP__ #define __SD_MODEL_TE_T5_HPP__ #include @@ -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; } }; diff --git a/otherarch/sdcpp/src/model/upscaler/esrgan.hpp b/otherarch/sdcpp/src/model/upscaler/esrgan.hpp index 4afbab07d..21c977124 100644 --- a/otherarch/sdcpp/src/model/upscaler/esrgan.hpp +++ b/otherarch/sdcpp/src/model/upscaler/esrgan.hpp @@ -1,4 +1,4 @@ -#ifndef __SD_MODEL_UPSCALER_ESRGAN_HPP__ +#ifndef __SD_MODEL_UPSCALER_ESRGAN_HPP__ #define __SD_MODEL_UPSCALER_ESRGAN_HPP__ #include diff --git a/otherarch/sdcpp/src/model/upscaler/ltx_latent_upscaler.hpp b/otherarch/sdcpp/src/model/upscaler/ltx_latent_upscaler.hpp index 5343ad036..b70e16136 100644 --- a/otherarch/sdcpp/src/model/upscaler/ltx_latent_upscaler.hpp +++ b/otherarch/sdcpp/src/model/upscaler/ltx_latent_upscaler.hpp @@ -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 diff --git a/otherarch/sdcpp/src/model/vae/auto_encoder_kl.hpp b/otherarch/sdcpp/src/model/vae/auto_encoder_kl.hpp index e41f5fd46..604347d77 100644 --- a/otherarch/sdcpp/src/model/vae/auto_encoder_kl.hpp +++ b/otherarch/sdcpp/src/model/vae/auto_encoder_kl.hpp @@ -816,12 +816,13 @@ struct AutoEncoderKL : public VAE { } sd::Tensor diffusion_to_vae_latents(const sd::Tensor& 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 vae_to_diffusion_latents(const sd::Tensor& latents) override { diff --git a/otherarch/sdcpp/src/model/vae/ltx_audio_vae.hpp b/otherarch/sdcpp/src/model/vae/ltx_audio_vae.hpp index 49847445d..9912a6136 100644 --- a/otherarch/sdcpp/src/model/vae/ltx_audio_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/ltx_audio_vae.hpp @@ -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 @@ -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(left_pad), 0, 0, 0, 0, 0, 0, 0); + x = ggml_ext_pad_ext(ctx, runner_ctx->backend, x, static_cast(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(blocks["conv"]); x = ggml_ext_pad_ext(ctx->ggml_ctx, + ctx->backend, x, 0, 1, diff --git a/otherarch/sdcpp/src/model/vae/tae.hpp b/otherarch/sdcpp/src/model/vae/tae.hpp index 7c6e1d35c..558c0544c 100644 --- a/otherarch/sdcpp/src/model/vae/tae.hpp +++ b/otherarch/sdcpp/src/model/vae/tae.hpp @@ -408,7 +408,7 @@ public: h = conv->forward(ctx, h); for (int j = 0; j < num_blocks; j++) { auto block = std::dynamic_pointer_cast(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(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(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; } }; diff --git a/otherarch/sdcpp/src/model/vae/vae.hpp b/otherarch/sdcpp/src/model/vae/vae.hpp index af091bb57..8b8c46ded 100644 --- a/otherarch/sdcpp/src/model/vae/vae.hpp +++ b/otherarch/sdcpp/src/model/vae/vae.hpp @@ -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; diff --git a/otherarch/sdcpp/src/model/vae/wan_vae.hpp b/otherarch/sdcpp/src/model/vae/wan_vae.hpp index 8a845c7ca..949562096 100644 --- a/otherarch/sdcpp/src/model/vae/wan_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/wan_vae.hpp @@ -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(new RMS_norm(in_dim)); // residual.1 is nn.SiLU() - blocks["residual.2"] = std::shared_ptr(new CausalConv3d(in_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + if (is_2D) { + blocks["residual.2"] = std::shared_ptr(new Conv2dBut3d(in_dim, out_dim, {3, 3}, {1, 1}, {1, 1})); + } else { + blocks["residual.2"] = std::shared_ptr(new CausalConv3d(in_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + } blocks["residual.3"] = std::shared_ptr(new RMS_norm(out_dim)); // residual.4 is nn.SiLU() // residual.5 is nn.Dropout() - blocks["residual.6"] = std::shared_ptr(new CausalConv3d(out_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + if (is_2D) { + blocks["residual.6"] = std::shared_ptr(new Conv2dBut3d(out_dim, out_dim, {3, 3}, {1, 1}, {1, 1})); + } else { + blocks["residual.6"] = std::shared_ptr(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(new CausalConv3d(in_dim, out_dim, {1, 1, 1})); + if (is_2D) { + blocks["shortcut"] = std::shared_ptr(new Conv2dBut3d(in_dim, out_dim, {1, 1})); + } else { + blocks["shortcut"] = std::shared_ptr(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(blocks["shortcut"]); + if (is_2D) { + auto shortcut = std::dynamic_pointer_cast(blocks["shortcut"]); - h = shortcut->forward(ctx, x); + h = shortcut->forward(ctx, x); + } else { + auto shortcut = std::dynamic_pointer_cast(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(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(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(new ResidualBlock(in_dim, out_dim)); + blocks["downsamples." + std::to_string(i)] = std::shared_ptr(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(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(new ResidualBlock(in_dim, out_dim)); + blocks["upsamples." + std::to_string(i)] = std::shared_ptr(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 dim_mult; int num_res_blocks; std::vector temperal_downsample; + bool is_2D = false; public: Encoder3d(int64_t dim = 128, int64_t z_dim = 4, + int64_t in_channels = 3, std::vector dim_mult = {1, 2, 4, 4}, int num_res_blocks = 2, std::vector 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 dims = {dim}; for (int u : dim_mult) { dims.push_back(dim * u); } - if (wan2_2) { - blocks["conv1"] = std::shared_ptr(new CausalConv3d(12, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + if (is_2D) { + blocks["conv1"] = std::shared_ptr(new Conv2dBut3d(in_channels, dims[0], {3, 3}, {1, 1}, {1, 1})); } else { - blocks["conv1"] = std::shared_ptr(new CausalConv3d(3, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + blocks["conv1"] = std::shared_ptr(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(new ResidualBlock(in_dim, out_dim)); + auto block = std::shared_ptr(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(new ResidualBlock(out_dim, out_dim)); + blocks["middle.0"] = std::shared_ptr(new ResidualBlock(out_dim, out_dim, is_2D)); blocks["middle.1"] = std::shared_ptr(new AttentionBlock(out_dim)); - blocks["middle.2"] = std::shared_ptr(new ResidualBlock(out_dim, out_dim)); + blocks["middle.2"] = std::shared_ptr(new ResidualBlock(out_dim, out_dim, is_2D)); blocks["head.0"] = std::shared_ptr(new RMS_norm(out_dim)); // head.1 is nn.SiLU() - blocks["head.2"] = std::shared_ptr(new CausalConv3d(out_dim, z_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + if (is_2D) { + blocks["head.2"] = std::shared_ptr(new Conv2dBut3d(out_dim, z_dim, {3, 3}, {1, 1}, {1, 1})); + } else { + blocks["head.2"] = std::shared_ptr(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(blocks["head.2"]); // conv1 - if (feat_cache.size() > 0) { + if (is_2D) { + auto conv1 = std::dynamic_pointer_cast(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(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 dim_mult; int num_res_blocks; std::vector temperal_upsample; + bool is_2D = false; public: Decoder3d(int64_t dim = 128, int64_t z_dim = 4, + int64_t out_channels = 3, std::vector dim_mult = {1, 2, 4, 4}, int num_res_blocks = 2, std::vector 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 dims = {dim_mult[dim_mult.size() - 1] * dim}; for (int i = static_cast(dim_mult.size()) - 1; i >= 0; i--) { @@ -779,12 +846,16 @@ namespace WAN { } // init block - blocks["conv1"] = std::shared_ptr(new CausalConv3d(z_dim, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + if (is_2D) { + blocks["conv1"] = std::shared_ptr(new Conv2dBut3d(z_dim, dims[0], {3, 3}, {1, 1}, {1, 1})); + } else { + blocks["conv1"] = std::shared_ptr(new CausalConv3d(z_dim, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + } // middle blocks - blocks["middle.0"] = std::shared_ptr(new ResidualBlock(dims[0], dims[0])); + blocks["middle.0"] = std::shared_ptr(new ResidualBlock(dims[0], dims[0], is_2D)); blocks["middle.1"] = std::shared_ptr(new AttentionBlock(dims[0])); - blocks["middle.2"] = std::shared_ptr(new ResidualBlock(dims[0], dims[0])); + blocks["middle.2"] = std::shared_ptr(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(new ResidualBlock(in_dim, out_dim)); + auto block = std::shared_ptr(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(new RMS_norm(out_dim)); + blocks["head.0"] = std::shared_ptr(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(new CausalConv3d(out_dim, 12, {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + if (is_2D) { + blocks["head.2"] = std::shared_ptr(new Conv2dBut3d(out_dim, final_dim, {3, 3}, {1, 1}, {1, 1})); } else { - blocks["head.2"] = std::shared_ptr(new CausalConv3d(out_dim, 3, {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + blocks["head.2"] = std::shared_ptr(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(blocks["head.2"]); // conv1 - if (feat_cache.size() > 0) { + if (is_2D) { + auto conv1 = std::dynamic_pointer_cast(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(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 temperal_upsample = {true, true, false}; std::vector 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(new Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, temperal_downsample, wan2_2)); - blocks["conv1"] = std::shared_ptr(new CausalConv3d(z_dim * 2, z_dim * 2, {1, 1, 1})); + blocks["encoder"] = std::shared_ptr(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(new Conv2dBut3d(z_dim * 2, z_dim * 2, {1, 1})); + } else { + blocks["conv1"] = std::shared_ptr(new CausalConv3d(z_dim * 2, z_dim * 2, {1, 1, 1})); + } + } + blocks["decoder"] = std::shared_ptr(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(new Conv2dBut3d(z_dim, z_dim, {1, 1})); + } else { + blocks["conv2"] = std::shared_ptr(new CausalConv3d(z_dim, z_dim, {1, 1, 1})); } - blocks["decoder"] = std::shared_ptr(new Decoder3d(dec_dim, z_dim, dim_mult, num_res_blocks, temperal_upsample, wan2_2)); - blocks["conv2"] = std::shared_ptr(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(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(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(blocks["decoder"]); auto conv2 = std::dynamic_pointer_cast(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(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 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); } diff --git a/otherarch/sdcpp/src/model_loader.cpp b/otherarch/sdcpp/src/model_loader.cpp index 788663103..fda1f3e82 100644 --- a/otherarch/sdcpp/src/model_loader.cpp +++ b/otherarch/sdcpp/src/model_loader.cpp @@ -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 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 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 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 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; } diff --git a/otherarch/sdcpp/src/name_conversion.cpp b/otherarch/sdcpp/src/name_conversion.cpp index ccc8347b7..64b7c681d 100644 --- a/otherarch/sdcpp/src/name_conversion.cpp +++ b/otherarch/sdcpp/src/name_conversion.cpp @@ -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> 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> 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> 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 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 { diff --git a/otherarch/sdcpp/src/runtime/denoiser.hpp b/otherarch/sdcpp/src/runtime/denoiser.hpp index 28b29ef27..812eebe61 100644 --- a/otherarch/sdcpp/src/runtime/denoiser.hpp +++ b/otherarch/sdcpp/src/runtime/denoiser.hpp @@ -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 get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t t_to_sigma) override { + std::vector 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(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(i) / static_cast(n); + const double t_cont = beta_ppf(u, alpha, beta) * t_max; + const int t = static_cast(std::lround(t_cont)); + + if (t != last_t) { + result.push_back(t_to_sigma(static_cast(t))); + last_t = t; + } + } + + result.push_back(0.f); + return result; + } +}; + struct SimpleScheduler : SigmaScheduler { std::vector get_sigmas(uint32_t n, float sigma_min, float sigma_max, t_to_sigma_t t_to_sigma) override { std::vector 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(image_seq_len) * m + b; + } + + std::vector get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override { + std::vector 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(i) / static_cast(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 get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override { + std::vector 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(i) / static_cast(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(image_seq_len) / static_cast(known_seq_len)); + mean += 0.5f * std::log(static_cast(image_seq_len) / static_cast(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(i) / static_cast(n); // ndtri(1-t) == -ndtri(t) - float z = -ndtri(t); + float z = static_cast(-ndtri(t)); float y = mean + std * z; @@ -780,6 +1026,10 @@ struct Denoiser { LOG_INFO("get_sigmas with Karras scheduler"); scheduler = std::make_shared(); break; + case BETA_SCHEDULER: + LOG_INFO("get_sigmas with Beta scheduler"); + scheduler = std::make_shared(); + break; case EXPONENTIAL_SCHEDULER: LOG_INFO("get_sigmas exponential scheduler"); scheduler = std::make_shared(); @@ -825,6 +1075,16 @@ struct Denoiser { scheduler = std::make_shared(image_seq_len, extra_sample_args); break; } + case FLUX2_SCHEDULER: { + LOG_INFO("get_sigmas with Flux2 scheduler"); + scheduler = std::make_shared(image_seq_len); + break; + } + case FLUX_SCHEDULER: { + LOG_INFO("get_sigmas with Flux scheduler"); + scheduler = std::make_shared(image_seq_len, extra_sample_args); + break; + } default: LOG_INFO("get_sigmas with discrete scheduler (default)"); scheduler = std::make_shared(); @@ -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 sem_sigmas; + std::vector tex_sigmas; + std::vector sem_timesteps; + std::vector 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 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(i) / static_cast(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(u_sem * (kNumTrainTimesteps - 1)))); + int idx_tex = std::min(kNumTrainTimesteps - 1, + std::max(0, static_cast(u_tex * (kNumTrainTimesteps - 1)))); + + float t_sem = static_cast(kNumTrainTimesteps - idx_sem); + float t_tex = static_cast(kNumTrainTimesteps - idx_tex); + float sigma_sem = t_sem / static_cast(kNumTrainTimesteps); + float sigma_tex = t_tex / static_cast(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 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 noise_scaling(float sigma, + const sd::Tensor& noise, + const sd::Tensor& latent) override { + SD_UNUSED(sigma); + SD_UNUSED(latent); + // Sampling starts from x0_init = noise * 2 (see MiniT2I reference). + return noise * 2.0f; + } + + sd::Tensor inverse_noise_scaling(float sigma, const sd::Tensor& latent) override { + SD_UNUSED(sigma); + return latent; } std::vector 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 sigmas; + sigmas.reserve(n + 1); + for (uint32_t i = 0; i < n; ++i) { + sigmas.push_back(1.0f - static_cast(i) / static_cast(n)); + } + sigmas.push_back(0.0f); + return sigmas; } }; @@ -1141,6 +1503,40 @@ static sd::Tensor sample_euler_ancestral(denoise_cb_t model, return x; } +static sd::Tensor sample_sefi_euler(SefiFlowDenoiser* sefi, + denoise_cb_t model, + sd::Tensor x) { + const std::vector& sigma_tex_vec = sefi->tex_sigmas; + const std::vector& sigma_sem_vec = sefi->sem_sigmas; + int steps = static_cast(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 denoised = std::move(denoised_opt.pred); + sd::Tensor 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 sample_euler(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas) { @@ -2056,7 +2452,13 @@ static sd::Tensor sample_k_diffusion(sample_method_t method, std::shared_ptr rng, float eta, bool is_flow_denoiser, - const char* extra_sample_args) { + const char* extra_sample_args, + std::shared_ptr denoiser_for_dispatch = nullptr) { + if (denoiser_for_dispatch) { + if (auto sefi = std::dynamic_pointer_cast(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: diff --git a/otherarch/sdcpp/src/runtime/imatrix.cpp b/otherarch/sdcpp/src/runtime/imatrix.cpp new file mode 100644 index 000000000..313eadc68 --- /dev/null +++ b/otherarch/sdcpp/src/runtime/imatrix.cpp @@ -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 +#include +#include + +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 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(src0->ne[2]); + const int n_ids = static_cast(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 to_store; + + for (const auto& kv : stats_) { + const int n_all = static_cast(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(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(stat.values.size()); + out.write((const char*)&nval, sizeof(nval)); + if (nval > 0) { + std::vector tmp(nval); + for (int i = 0; i < nval; i++) { + tmp[i] = (stat.values[i] / static_cast(stat.counts[i])) * static_cast(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 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 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; +} diff --git a/otherarch/sdcpp/src/runtime/imatrix.h b/otherarch/sdcpp/src/runtime/imatrix.h new file mode 100644 index 000000000..341c18c8b --- /dev/null +++ b/otherarch/sdcpp/src/runtime/imatrix.h @@ -0,0 +1,45 @@ +#ifndef __SD_RUNTIME_IMATRIX_H__ +#define __SD_RUNTIME_IMATRIX_H__ + +#include +#include +#include +#include +#include + +/* Adapted from llama.cpp (credits: Kawrakow). */ + +struct ggml_tensor; + +struct IMatrixStats { + std::vector values{}; + std::vector counts{}; + int ncall = 0; +}; + +class IMatrixCollector { +private: + std::unordered_map stats_ = {}; + std::mutex mutex_; + int last_call_ = 0; + std::vector src1_data_; + std::vector 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 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__ diff --git a/otherarch/sdcpp/src/stable-diffusion.cpp b/otherarch/sdcpp/src/stable-diffusion.cpp index e951e4143..a0ee643b3 100644 --- a/otherarch/sdcpp/src/stable-diffusion.cpp +++ b/otherarch/sdcpp/src/stable-diffusion.cpp @@ -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(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(backend_for(SDBackendModule::TE), + bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED; + cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), tensor_storage_map, version, "", - true, + enable_vision, model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), + diffusion_model = std::make_shared(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(backend_for(SDBackendModule::TE), + tensor_storage_map, + model_manager); + diffusion_model = std::make_shared(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(backend_for(SDBackendModule::TE), tensor_storage_map, @@ -1140,11 +1154,7 @@ public: } auto create_tae = [&](bool decode_only) -> std::shared_ptr { - 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(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(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(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(); break; } - case FLUX2_FLOW_PRED: { - LOG_INFO("running in Flux2 FLOW mode"); - denoiser = std::make_shared(); + case SEFI_FLOW_PRED: { + LOG_INFO("running in SeFi-Image dual-time FLOW mode"); + denoiser = std::make_shared(); + break; + } + case MINIT2I_FLOW_PRED: { + LOG_INFO("running in MiniT2I FLOW mode"); + denoiser = std::make_shared(); break; } default: { @@ -1913,7 +1928,16 @@ public: std::vector process_timesteps(const std::vector& timesteps, const sd::Tensor& init_latent, - const sd::Tensor& denoise_mask) { + const sd::Tensor& denoise_mask, + int step) { + if (auto sefi_denoiser = std::dynamic_pointer_cast(denoiser)) { + int sched_idx = step > 0 ? step - 1 : 0; + if (sched_idx >= static_cast(sefi_denoiser->tex_timesteps.size())) { + sched_idx = static_cast(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(static_cast(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 x_t = !noise.empty() - ? denoiser->noise_scaling(sigmas[0], noise, init_latent) - : init_latent; - sd::Tensor denoised = x_t; SamplePreviewContext preview = prepare_sample_preview_context(); + sd::Tensor x_t = !noise.empty() + ? denoiser->noise_scaling(sigmas[0], noise, init_latent) + : init_latent; + sd::Tensor denoised = x_t; + auto denoise = [&](const sd::Tensor& 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({static_cast(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& 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> controls; DiffusionParams diffusion_params; - diffusion_params.x = &noised_input; - diffusion_params.timesteps = ×teps_tensor; - diffusion_params.increase_ref_index = increase_ref_index; + diffusion_params.x = &noised_input; + diffusion_params.timesteps = ×teps_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 decode_first_stage(const sd::Tensor& 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 ensure_image_tensor_channels(sd::Tensor 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 alpha = sd::full({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 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 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 prepare_image_generation_latents(sd sd::Tensor init_image_tensor; sd::Tensor control_image_tensor; sd::Tensor 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 prepare_image_generation_latents(sd plan->sample_steps = static_cast(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 prepare_image_generation_latents(sd sd::Tensor init_latent; sd::Tensor 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 prepare_image_generation_latents(sd std::vector> 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({request->width, request->height, 3, 1})); + ref_images.push_back(sd::zeros({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 prepare_image_generation_latents(sd vae_width = round(vae_width / factor) * factor; auto resized_ref_img = sd::ops::interpolate(ref_images[i], - {static_cast(vae_width), static_cast(vae_height), 3, 1}); + {static_cast(vae_width), + static_cast(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(i), @@ -4565,6 +4655,11 @@ static std::optional 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::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::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 prepare_image_generation_embeds(sd_c static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx, const GenerationRequest& request, - const std::vector>& final_latents) { + const std::vector>& 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 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 layer_latent = sd::ops::slice(final_latents[i], 2, layer_index, layer_index + 1); + layer_latent.squeeze_(2); + sd::Tensor 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 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(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 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(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 noise = sd::randn_like(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 prepare_video_generation_latents(sd_ctx_t* sd_ctx, diff --git a/otherarch/sdcpp/src/upscaler.cpp b/otherarch/sdcpp/src/upscaler.cpp index d02366ecb..88a8a6336 100644 --- a/otherarch/sdcpp/src/upscaler.cpp +++ b/otherarch/sdcpp/src/upscaler.cpp @@ -4,6 +4,7 @@ #include "model_loader.h" #include "stable-diffusion.h" +#include #include 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) {