Merge branch 'master' into ltx2.3

.github/workflows/build.yml

@@ -176,6 +176,7 @@ jobs:
 
   build-and-push-docker-images:
     name: Build and push container images
+    if: ${{ github.event_name != 'pull_request' }}
     runs-on: ubuntu-latest
 
     permissions:

CMakeLists.txt

@@ -11,6 +11,10 @@ endif()
 if (MSVC)
     add_compile_definitions(_CRT_SECURE_NO_WARNINGS)
     add_compile_definitions(_SILENCE_CXX17_CODECVT_HEADER_DEPRECATION_WARNING)
+    add_compile_options(
+        $<$<COMPILE_LANGUAGE:C>:/MP>
+        $<$<COMPILE_LANGUAGE:CXX>:/MP>
+    )
 endif()
 
 set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
@@ -152,12 +156,16 @@ endif()
 
 set(SD_LIB stable-diffusion)
 
-file(GLOB SD_LIB_SOURCES
+file(GLOB SD_LIB_SOURCES CONFIGURE_DEPENDS
     "src/*.h"
     "src/*.cpp"
     "src/*.hpp"
-    "src/vocab/*.h"
+    "src/model_io/*.h"
-    "src/vocab/*.cpp"
+    "src/model_io/*.cpp"
+    "src/tokenizers/*.h"
+    "src/tokenizers/*.cpp"
+    "src/tokenizers/vocab/*.h"
+    "src/tokenizers/vocab/*.cpp"
 )
 
 find_program(GIT_EXE NAMES git git.exe NO_CMAKE_FIND_ROOT_PATH)
@@ -250,7 +258,7 @@ endif()
 add_subdirectory(thirdparty)
 
 target_link_libraries(${SD_LIB} PUBLIC ggml zip)
-target_include_directories(${SD_LIB} PUBLIC . include)
+target_include_directories(${SD_LIB} PUBLIC . src include)
 target_include_directories(${SD_LIB} PUBLIC . thirdparty)
 target_compile_features(${SD_LIB} PUBLIC c_std_11 cxx_std_17)
 

README.md

@@ -57,6 +57,7 @@ API and command-line option may change frequently.***
     - [Z-Image](./docs/z_image.md)
     - [Ovis-Image](./docs/ovis_image.md)
     - [Anima](./docs/anima.md)
+    - [ERNIE-Image](./docs/ernie_image.md)
   - Image Edit Models
     - [FLUX.1-Kontext-dev](./docs/kontext.md)
     - [Qwen Image Edit series](./docs/qwen_image_edit.md)
@@ -76,9 +77,10 @@ API and command-line option may change frequently.***
   - OpenCL
   - SYCL
 - Supported weight formats
-  - Pytorch checkpoint (`.ckpt` or `.pth`)
+  - Pytorch checkpoint (`.ckpt` or `.pth` or `.pt`)
   - Safetensors (`.safetensors`)
   - GGUF (`.gguf`)
+- Convert mode supports converting model weights to `.gguf` or `.safetensors`
 - Supported platforms
     - Linux
     - Mac OS
@@ -96,6 +98,7 @@ API and command-line option may change frequently.***
     - `DPM++ 2M`
     - [`DPM++ 2M v2`](https://github.com/AUTOMATIC1111/stable-diffusion-webui/discussions/8457)
     - `DPM++ 2S a`
+    - `ER-SDE`
     - [`LCM`](https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/13952)
 - Cross-platform reproducibility
     - `--rng cuda`, default, consistent with the `stable-diffusion-webui GPU RNG`
@@ -144,6 +147,7 @@ If you want to improve performance or reduce VRAM/RAM usage, please refer to [pe
 - [🔥Z-Image](./docs/z_image.md)
 - [Ovis-Image](./docs/ovis_image.md)
 - [Anima](./docs/anima.md)
+- [ERNIE-Image](./docs/ernie_image.md)
 - [LoRA](./docs/lora.md)
 - [LCM/LCM-LoRA](./docs/lcm.md)
 - [Using PhotoMaker to personalize image generation](./docs/photo_maker.md)

docs/distilled_sd.md

@@ -87,51 +87,32 @@ pipe.save_pretrained("segmindtiny-sd", safe_serialization=True)
 ```bash
 python convert_diffusers_to_original_stable_diffusion.py \
       --model_path  ./segmindtiny-sd \
-      --checkpoint_path ./segmind_tiny-sd.ckpt --half
+      --checkpoint_path ./segmind_tiny-sd.safetensors  --half --use_safetensors
 ```
 
-The file segmind_tiny-sd.ckpt will be generated and is now ready for use with sd.cpp. You can follow a similar process for the other models mentioned above.
+The file segmind_tiny-sd.safetensors will be generated and is now ready for use with sd.cpp. You can follow a similar process for the other models mentioned above.
 
 
-##### Another available .ckpt file:
+### SDXS-512-DreamShaper
-
- * https://huggingface.co/ClashSAN/small-sd/resolve/main/tinySDdistilled.ckpt
-
-To use this file, you must first adjust its non-contiguous tensors:
-
-```python
-import torch
-ckpt = torch.load("tinySDdistilled.ckpt", map_location=torch.device('cpu'))
-for key, value in ckpt['state_dict'].items():
-    if isinstance(value, torch.Tensor):
-        ckpt['state_dict'][key] = value.contiguous()
-torch.save(ckpt, "tinySDdistilled_fixed.ckpt")
-```
-
-
-### SDXS-512
 
 Another very tiny and **incredibly fast**  model is SDXS by IDKiro et al.  The authors refer to it as *"Real-Time One-Step Latent Diffusion Models with Image Conditions"*. For details read the paper: https://arxiv.org/pdf/2403.16627 . Once again the authors removed some more blocks of U-Net part and unlike other SD1 models they use an adjusted _AutoEncoderTiny_ instead of default _AutoEncoderKL_ for the VAE part.
+##### Some ready-to-run SDXS-512 model files are available online, such as:
 
-##### 1. Download the diffusers model from  Hugging Face using Python:
+* https://huggingface.co/akleine/sdxs-512
-
+* https://huggingface.co/concedo/sdxs-512-tinySDdistilled-GGUF
-```python
-from diffusers import StableDiffusionPipeline
-pipe = StableDiffusionPipeline.from_pretrained("IDKiro/sdxs-512-dreamshaper")
-pipe.save_pretrained(save_directory="sdxs")
-```
-##### 2. Create a safetensors file
-
-```bash
-python convert_diffusers_to_original_stable_diffusion.py \
-    --model_path  sdxs  --checkpoint_path sdxs.safetensors --half --use_safetensors
-```
-
-##### 3. Run the model as follows:
 
+##### Run the model as follows:
 ```bash
 ~/stable-diffusion.cpp/build/bin/sd-cli -m sdxs.safetensors -p "portrait of a lovely cat" \
   --cfg-scale 1 --steps 1
 ```
-
 Both options: ``` --cfg-scale 1 ``` and  ``` --steps 1 ``` are mandatory here.
+
+### SDXS-512-0.9
+
+Even though the name "SDXS-512-0.9" is similar to "SDXS-512-DreamShaper", it is *completely different* but also **incredibly fast**. Sometimes it is preferred, so try it yourself.
+##### Download a ready-to-run file from here:
+
+* https://huggingface.co/akleine/sdxs-09
+
+For the use of this model, both options ``` --cfg-scale 1 ``` and ``` --steps 1 ``` are again absolutely necessary.

docs/ernie_image.md

@@ -0,0 +1,35 @@
+# How to Use
+
+You can run ERNIE-Image with stable-diffusion.cpp on GPUs with 4GB of VRAM — or even less.
+
+## Download weights
+
+- Download ERNIE-Image-Turbo
+    - safetensors: https://huggingface.co/Comfy-Org/ERNIE-Image/tree/main/diffusion_models
+    - gguf: https://huggingface.co/unsloth/ERNIE-Image-Turbo-GGUF/tree/main
+- Download ERNIE-Image
+    - safetensors: https://huggingface.co/Comfy-Org/ERNIE-Image/tree/main/diffusion_models
+    - gguf: https://huggingface.co/unsloth/ERNIE-Image-GGUF/tree/main
+- Download vae
+    - safetensors: https://huggingface.co/Comfy-Org/ERNIE-Image/tree/main/vae
+- Download ministral 3b
+    - safetensors: https://huggingface.co/Comfy-Org/ERNIE-Image/tree/main/text_encoders
+    - gguf: https://huggingface.co/unsloth/Ministral-3-3B-Instruct-2512-GGUF/tree/main
+
+## Examples
+
+### ERNIE-Image-Turbo
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\ernie-image-turbo.safetensors --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors  --llm ..\..\ComfyUI\models\text_encoders\ministral-3-3b.safetensors -p "a lovely cat" --cfg-scale 1.0 --steps 8 -v --offload-to-cpu --diffusion-fa
+```
+
+<img width="256" alt="ERNIE-Image Turbo example" src="../assets/ernie_image/turbo_example.png" />
+
+### ERNIE-Image
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\ernie-image-UD-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors  --llm ..\..\ComfyUI\models\text_encoders\ministral-3-3b.safetensors -p "a lovely cat" --cfg-scale 5.0 -v --offload-to-cpu --diffusion-fa
+```
+
+<img width="256" alt="ERNIE-Image example" src="../assets/ernie_image/example.png" />

examples/cli/README.md

@@ -4,14 +4,17 @@
 usage: ./bin/sd-cli  [options]
 
 CLI Options:
-  -o, --output <string>       path to write result image to. you can use printf-style %d format specifiers for image sequences (default:
+  -o, --output <string>         path to write result image to. you can use printf-style %d format specifiers for image
-                              ./output.png) (eg. output_%03d.png). For video generation, single-file outputs support .avi, .webm, and animated .webp
+                                sequences (default: ./output.png) (eg. output_%03d.png). Single-file video outputs
-  --preview-path <string>     path to write preview image to (default: ./preview.png). Multi-frame previews support .avi, .webm, and animated .webp
+                                support .avi, .webm, and animated .webp
-  --preview-interval <int>    interval in denoising steps between consecutive updates of the image preview file (default is 1, meaning updating at
-                              every step)
-  --output-begin-idx <int>    starting index for output image sequence, must be non-negative (default 0 if specified %d in output path, 1 otherwise)
   --image <string>              path to the image to inspect (for metadata mode)
   --metadata-format <string>    metadata output format, one of [text, json] (default: text)
+  --preview-path <string>       path to write preview image to (default: ./preview.png). Multi-frame previews support
+                                .avi, .webm, and animated .webp
+  --preview-interval <int>      interval in denoising steps between consecutive updates of the image preview file
+                                (default is 1, meaning updating at every step)
+  --output-begin-idx <int>      starting index for output image sequence, must be non-negative (default 0 if specified
+                                %d in output path, 1 otherwise)
   --canny                       apply canny preprocessor (edge detection)
   --convert-name                convert tensor name (for convert mode)
   -v, --verbose                 print extra info
@@ -31,7 +34,8 @@ Context Options:
   --clip_g <string>                        path to the clip-g text encoder
   --clip_vision <string>                   path to the clip-vision encoder
   --t5xxl <string>                         path to the t5xxl text encoder
-  --llm <string>                           path to the llm text encoder. For example: (qwenvl2.5 for qwen-image, mistral-small3.2 for flux2, ...)
+  --llm <string>                           path to the llm text encoder. For example: (qwenvl2.5 for qwen-image,
+                                           mistral-small3.2 for flux2, ...)
   --llm_vision <string>                    path to the llm vit
   --qwen2vl <string>                       alias of --llm. Deprecated.
   --qwen2vl_vision <string>                alias of --llm_vision. Deprecated.
@@ -43,16 +47,16 @@ Context Options:
   --control-net <string>                   path to control net model
   --embd-dir <string>                      embeddings directory
   --lora-model-dir <string>                lora model directory
+  --hires-upscalers-dir <string>           highres fix upscaler model directory
   --tensor-type-rules <string>             weight type per tensor pattern (example: "^vae\.=f16,model\.=q8_0")
   --photo-maker <string>                   path to PHOTOMAKER model
   --upscale-model <string>                 path to esrgan model.
-  -t, --threads <int>                      number of threads to use during computation (default: -1). If threads <= 0, then threads will be set to the number of
+  -t, --threads <int>                      number of threads to use during computation (default: -1). If threads <= 0,
-                                           CPU physical cores
+                                           then threads will be set to the number of CPU physical cores
   --chroma-t5-mask-pad <int>               t5 mask pad size of chroma
-  --vae-tile-overlap <float>               tile overlap for vae tiling, in fraction of tile size (default: 0.5)
-  --vae-tiling                             process vae in tiles to reduce memory usage
   --force-sdxl-vae-conv-scale              force use of conv scale on sdxl vae
-  --offload-to-cpu                         place the weights in RAM to save VRAM, and automatically load them into VRAM when needed
+  --offload-to-cpu                         place the weights in RAM to save VRAM, and automatically load them into VRAM
+                                           when needed
   --mmap                                   whether to memory-map model
   --control-net-cpu                        keep controlnet in cpu (for low vram)
   --clip-on-cpu                            keep clip in cpu (for low vram)
@@ -67,20 +71,19 @@ Context Options:
   --chroma-disable-dit-mask                disable dit mask for chroma
   --qwen-image-zero-cond-t                 enable zero_cond_t for qwen image
   --chroma-enable-t5-mask                  enable t5 mask for chroma
-  --type                                   weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, q4_K). If not specified, the default is the
+  --type                                   weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K,
-                                           type of the weight file
+                                           q4_K). If not specified, the default is the type of the weight file
   --rng                                    RNG, one of [std_default, cuda, cpu], default: cuda(sd-webui), cpu(comfyui)
   --sampler-rng                            sampler RNG, one of [std_default, cuda, cpu]. If not specified, use --rng
-  --prediction                             prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, flux2_flow]
+  --prediction                             prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow,
-  --lora-apply-mode                        the way to apply LoRA, one of [auto, immediately, at_runtime], default is auto. In auto mode, if the model weights
+                                           flux2_flow]
-                                           contain any quantized parameters, the at_runtime mode will be used; otherwise,
+  --lora-apply-mode                        the way to apply LoRA, one of [auto, immediately, at_runtime], default is
-                                           immediately will be used.The immediately mode may have precision and
+                                           auto. In auto mode, if the model weights contain any quantized parameters,
-                                           compatibility issues with quantized parameters, but it usually offers faster inference
+                                           the at_runtime mode will be used; otherwise, immediately will be used.The
-                                           speed and, in some cases, lower memory usage. The at_runtime mode, on the
+                                           immediately mode may have precision and compatibility issues with quantized
-                                           other hand, is exactly the opposite.
+                                           parameters, but it usually offers faster inference speed and, in some cases,
-  --vae-tile-size                          tile size for vae tiling, format [X]x[Y] (default: 32x32)
+                                           lower memory usage. The at_runtime mode, on the other hand, is exactly the
-  --vae-relative-tile-size                 relative tile size for vae tiling, format [X]x[Y], in fraction of image size if < 1, in number of tiles per dim if >=1
+                                           opposite.
-                                           (overrides --vae-tile-size)
 
 Generation Options:
   -p, --prompt <string>                    the prompt to render
@@ -89,69 +92,99 @@ Generation Options:
   --end-img <string>                       path to the end image, required by flf2v
   --mask <string>                          path to the mask image
   --control-image <string>                 path to control image, control net
-  --control-video <string>                 path to control video frames, It must be a directory path. The video frames inside should be stored as images in
+  --control-video <string>                 path to control video frames, It must be a directory path. The video frames
-                                           lexicographical (character) order. For example, if the control video path is
+                                           inside should be stored as images in lexicographical (character) order. For
-                                           `frames`, the directory contain images such as 00.png, 01.png, ... etc.
+                                           example, if the control video path is `frames`, the directory contain images
+                                           such as 00.png, 01.png, ... etc.
   --pm-id-images-dir <string>              path to PHOTOMAKER input id images dir
   --pm-id-embed-path <string>              path to PHOTOMAKER v2 id embed
+  --hires-upscaler <string>                highres fix upscaler, Lanczos, Nearest, Latent, Latent (nearest), Latent
+                                           (nearest-exact), Latent (antialiased), Latent (bicubic), Latent (bicubic
+                                           antialiased), or a model name under --hires-upscalers-dir (default: Latent)
   -H, --height <int>                       image height, in pixel space (default: 512)
   -W, --width <int>                        image width, in pixel space (default: 512)
   --steps <int>                            number of sample steps (default: 20)
   --high-noise-steps <int>                 (high noise) number of sample steps (default: -1 = auto)
-  --clip-skip <int>                        ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1). <= 0 represents unspecified,
+  --clip-skip <int>                        ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer
-                                           will be 1 for SD1.x, 2 for SD2.x
+                                           (default: -1). <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x
   -b, --batch-count <int>                  batch count
   --video-frames <int>                     video frames (default: 1)
   --fps <int>                              fps (default: 24)
-  --timestep-shift <int>                   shift timestep for NitroFusion models (default: 0). recommended N for NitroSD-Realism around 250 and 500 for
+  --timestep-shift <int>                   shift timestep for NitroFusion models (default: 0). recommended N for
-                                           NitroSD-Vibrant
+                                           NitroSD-Realism around 250 and 500 for NitroSD-Vibrant
   --upscale-repeats <int>                  Run the ESRGAN upscaler this many times (default: 1)
   --upscale-tile-size <int>                tile size for ESRGAN upscaling (default: 128)
+  --hires-width <int>                      highres fix target width, 0 to use --hires-scale (default: 0)
+  --hires-height <int>                     highres fix target height, 0 to use --hires-scale (default: 0)
+  --hires-steps <int>                      highres fix second pass sample steps, 0 to reuse --steps (default: 0)
+  --hires-upscale-tile-size <int>          highres fix upscaler tile size, reserved for model-backed upscalers (default:
+                                           128)
   --cfg-scale <float>                      unconditional guidance scale: (default: 7.0)
-  --img-cfg-scale <float>                  image guidance scale for inpaint or instruct-pix2pix models: (default: same as --cfg-scale)
+  --img-cfg-scale <float>                  image guidance scale for inpaint or instruct-pix2pix models: (default: same
+                                           as --cfg-scale)
   --guidance <float>                       distilled guidance scale for models with guidance input (default: 3.5)
-  --slg-scale <float>                      skip layer guidance (SLG) scale, only for DiT models: (default: 0). 0 means disabled, a value of 2.5 is nice for sd3.5
+  --slg-scale <float>                      skip layer guidance (SLG) scale, only for DiT models: (default: 0). 0 means
-                                           medium
+                                           disabled, a value of 2.5 is nice for sd3.5 medium
   --skip-layer-start <float>               SLG enabling point (default: 0.01)
   --skip-layer-end <float>                 SLG disabling point (default: 0.2)
-  --eta <float>                            noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a and dpm++2s_a)
+  --eta <float>                            noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and
+                                           res_2s; 1 for euler_a, er_sde and dpm++2s_a)
   --flow-shift <float>                     shift value for Flow models like SD3.x or WAN (default: auto)
   --high-noise-cfg-scale <float>           (high noise) unconditional guidance scale: (default: 7.0)
-  --high-noise-img-cfg-scale <float>       (high noise) image guidance scale for inpaint or instruct-pix2pix models (default: same as --cfg-scale)
+  --high-noise-img-cfg-scale <float>       (high noise) image guidance scale for inpaint or instruct-pix2pix models
-  --high-noise-guidance <float>            (high noise) distilled guidance scale for models with guidance input (default: 3.5)
+                                           (default: same as --cfg-scale)
-  --high-noise-slg-scale <float>           (high noise) skip layer guidance (SLG) scale, only for DiT models: (default: 0)
+  --high-noise-guidance <float>            (high noise) distilled guidance scale for models with guidance input
+                                           (default: 3.5)
+  --high-noise-slg-scale <float>           (high noise) skip layer guidance (SLG) scale, only for DiT models: (default:
+                                           0)
   --high-noise-skip-layer-start <float>    (high noise) SLG enabling point (default: 0.01)
   --high-noise-skip-layer-end <float>      (high noise) SLG disabling point (default: 0.2)
-  --high-noise-eta <float>                 (high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a and dpm++2s_a)
+  --high-noise-eta <float>                 (high noise) noise multiplier (default: 0 for ddim_trailing, tcd,
+                                           res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)
   --strength <float>                       strength for noising/unnoising (default: 0.75)
   --pm-style-strength <float>              
-  --control-strength <float>               strength to apply Control Net (default: 0.9). 1.0 corresponds to full destruction of information in init image
+  --control-strength <float>               strength to apply Control Net (default: 0.9). 1.0 corresponds to full
-  --moe-boundary <float>                   timestep boundary for Wan2.2 MoE model. (default: 0.875). Only enabled if `--high-noise-steps` is set to -1
+                                           destruction of information in init image
+  --moe-boundary <float>                   timestep boundary for Wan2.2 MoE model. (default: 0.875). Only enabled if
+                                           `--high-noise-steps` is set to -1
   --vace-strength <float>                  wan vace strength
-  --increase-ref-index                     automatically increase the indices of references images based on the order they are listed (starting with 1).
+  --vae-tile-overlap <float>               tile overlap for vae tiling, in fraction of tile size (default: 0.5)
+  --hires-scale <float>                    highres fix scale when target size is not set (default: 2.0)
+  --hires-denoising-strength <float>       highres fix second pass denoising strength (default: 0.7)
+  --increase-ref-index                     automatically increase the indices of references images based on the order
+                                           they are listed (starting with 1).
   --disable-auto-resize-ref-image          disable auto resize of ref images
   --disable-image-metadata                 do not embed generation metadata on image files
+  --vae-tiling                             process vae in tiles to reduce memory usage
+  --hires                                  enable highres fix
   -s, --seed                               RNG seed (default: 42, use random seed for < 0)
-  --sampling-method                        sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing,
+  --sampling-method                        sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m,
-                                           tcd, res_multistep, res_2s] (default: euler for Flux/SD3/Wan, euler_a
+                                           dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s,
-                                           otherwise)
+                                           er_sde] (default: euler for Flux/SD3/Wan, euler_a otherwise)
-  --high-noise-sampling-method             (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm,
+  --high-noise-sampling-method             (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a,
-                                           ddim_trailing, tcd, res_multistep, res_2s] default: euler for Flux/SD3/Wan,
+                                           dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep,
-                                           euler_a otherwise
+                                           res_2s, er_sde] default: euler for Flux/SD3/Wan, euler_a otherwise
-  --scheduler                              denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple,
+  --scheduler                              denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits,
-                                           kl_optimal, lcm, bong_tangent], default: discrete
+                                           smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent], default:
-  --sigmas                                 custom sigma values for the sampler, comma-separated (e.g., "14.61,7.8,3.5,0.0").
+                                           discrete
+  --sigmas                                 custom sigma values for the sampler, comma-separated (e.g.,
+                                           "14.61,7.8,3.5,0.0").
   --skip-layers                            layers to skip for SLG steps (default: [7,8,9])
   --high-noise-skip-layers                 (high noise) layers to skip for SLG steps (default: [7,8,9])
   -r, --ref-image                          reference image for Flux Kontext models (can be used multiple times)
-  --cache-mode                             caching method: 'easycache' (DiT), 'ucache' (UNET), 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level),
+  --cache-mode                             caching method: 'easycache' (DiT), 'ucache' (UNET),
-                                           'spectrum' (UNET/DiT Chebyshev+Taylor forecasting)
+                                           'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT
+                                           Chebyshev+Taylor forecasting)
   --cache-option                           named cache params (key=value format, comma-separated). easycache/ucache:
-                                           threshold=,start=,end=,decay=,relative=,reset=; dbcache/taylorseer/cache-dit: Fn=,Bn=,threshold=,warmup=;
+                                           threshold=,start=,end=,decay=,relative=,reset=; dbcache/taylorseer/cache-dit:
-                                           spectrum: w=,m=,lam=,window=,flex=,warmup=,stop=. Examples:
+                                           Fn=,Bn=,threshold=,warmup=; spectrum: w=,m=,lam=,window=,flex=,warmup=,stop=.
-                                           "threshold=0.25" or "threshold=1.5,reset=0" or "w=0.4,window=2"
+                                           Examples: "threshold=0.25" or "threshold=1.5,reset=0"
-  --scm-mask                               SCM steps mask for cache-dit: comma-separated 0/1 (e.g., "1,1,1,0,0,1,0,0,1,0") - 1=compute, 0=can cache
+  --scm-mask                               SCM steps mask for cache-dit: comma-separated 0/1 (e.g.,
+                                           "1,1,1,0,0,1,0,0,1,0") - 1=compute, 0=can cache
   --scm-policy                             SCM policy: 'dynamic' (default) or 'static'
+  --vae-tile-size                          tile size for vae tiling, format [X]x[Y] (default: 32x32)
+  --vae-relative-tile-size                 relative tile size for vae tiling, format [X]x[Y], in fraction of image size
+                                           if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size)
 ```
 
 Metadata mode inspects PNG/JPEG container metadata without loading any model:

examples/cli/main.cpp

@@ -279,7 +279,9 @@ void parse_args(int argc, const char** argv, SDCliParams& cli_params, SDContextP
     bool valid = cli_params.resolve_and_validate();
     if (valid && cli_params.mode != METADATA) {
         valid = ctx_params.resolve_and_validate(cli_params.mode) &&
-                gen_params.resolve_and_validate(cli_params.mode, ctx_params.lora_model_dir);
+                gen_params.resolve_and_validate(cli_params.mode,
+                                                ctx_params.lora_model_dir,
+                                                ctx_params.hires_upscalers_dir);
     }
 
     if (!valid) {
@@ -501,7 +503,7 @@ int main(int argc, const char* argv[]) {
 
     cli_params.verbose = true;
     sd_set_log_callback(sd_log_cb, (void*)&cli_params);
-    LLM::GemmaTokenizer tokenizer;
+    GemmaTokenizer tokenizer;
     auto tokens = tokenizer.tokenize("<html> 一只可爱的小猫");
     for (auto token : tokens) {
         LOG_INFO("%d", token);
@@ -698,6 +700,13 @@ int main(int argc, const char* argv[]) {
         vae_decode_only = false;
     }
 
+    if (gen_params.hires_enabled &&
+        (gen_params.resolved_hires_upscaler == SD_HIRES_UPSCALER_MODEL ||
+         gen_params.resolved_hires_upscaler == SD_HIRES_UPSCALER_LANCZOS ||
+         gen_params.resolved_hires_upscaler == SD_HIRES_UPSCALER_NEAREST)) {
+        vae_decode_only = false;
+    }
+
     sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(vae_decode_only, true, cli_params.taesd_preview);
 
     SDImageVec results;

examples/common/common.cpp

@@ -107,47 +107,60 @@ static bool is_absolute_path(const std::string& p) {
 
 std::string ArgOptions::wrap_text(const std::string& text, size_t width, size_t indent) {
     std::ostringstream oss;
-    size_t line_len = 0;
     size_t pos      = 0;
+    size_t line_len = 0;
 
     while (pos < text.size()) {
-        // Preserve manual newlines
         if (text[pos] == '\n') {
             oss << '\n'
                 << std::string(indent, ' ');
-            line_len = indent;
+            line_len = 0;
             ++pos;
             continue;
         }
 
-        // Add the character
+        if (std::isspace(static_cast<unsigned char>(text[pos]))) {
-        oss << text[pos];
-        ++line_len;
             ++pos;
-
+            continue;
-        // If the current line exceeds width, try to break at the last space
+        }
-        if (line_len >= width) {
+
-            std::string current = oss.str();
+        size_t word_start = pos;
-            size_t back         = current.size();
+        while (pos < text.size() &&
-
+               text[pos] != '\n' &&
-            // Find the last space (for a clean break)
+               !std::isspace(static_cast<unsigned char>(text[pos]))) {
-            while (back > 0 && current[back - 1] != ' ' && current[back - 1] != '\n')
+            ++pos;
-                --back;
+        }
-
+
-            // If found a space to break on
+        std::string word = text.substr(word_start, pos - word_start);
-            if (back > 0 && current[back - 1] != '\n') {
+        while (!word.empty()) {
-                std::string before = current.substr(0, back - 1);
+            size_t separator_len = line_len == 0 ? 0 : 1;
-                std::string after  = current.substr(back);
+            if (line_len + separator_len + word.size() <= width) {
-                oss.str("");
+                if (separator_len > 0) {
-                oss.clear();
+                    oss << ' ';
-                oss << before << "\n"
+                    ++line_len;
-                    << std::string(indent, ' ') << after;
+                }
-            } else {
+                oss << word;
-                // If no space found, just break at width
+                line_len += word.size();
-                oss << "\n"
+                word.clear();
-                    << std::string(indent, ' ');
+                continue;
+            }
+
+            if (line_len > 0) {
+                oss << '\n'
+                    << std::string(indent, ' ');
+                line_len = 0;
+                continue;
+            }
+
+            size_t chunk_len = std::min(width, word.size());
+            oss << word.substr(0, chunk_len);
+            line_len = chunk_len;
+            word.erase(0, chunk_len);
+            if (!word.empty()) {
+                oss << '\n'
+                    << std::string(indent, ' ');
+                line_len = 0;
             }
-            line_len = indent;
         }
     }
 
@@ -351,7 +364,10 @@ ArgOptions SDContextParams::get_options() {
          "--lora-model-dir",
          "lora model directory",
          &lora_model_dir},
-
+        {"",
+         "--hires-upscalers-dir",
+         "highres fix upscaler model directory",
+         &hires_upscalers_dir},
         {"",
          "--tensor-type-rules",
          "weight type per tensor pattern (example: \"^vae\\.=f16,model\\.=q8_0\")",
@@ -649,6 +665,7 @@ std::string SDContextParams::to_string() const {
         << "  wtype: " << sd_type_name(wtype) << ",\n"
         << "  tensor_type_rules: \"" << tensor_type_rules << "\",\n"
         << "  lora_model_dir: \"" << lora_model_dir << "\",\n"
+        << "  hires_upscalers_dir: \"" << hires_upscalers_dir << "\",\n"
         << "  photo_maker_path: \"" << photo_maker_path << "\",\n"
         << "  rng_type: " << sd_rng_type_name(rng_type) << ",\n"
         << "  sampler_rng_type: " << sd_rng_type_name(sampler_rng_type) << ",\n"
@@ -777,6 +794,12 @@ ArgOptions SDGenerationParams::get_options() {
          "--pm-id-embed-path",
          "path to PHOTOMAKER v2 id embed",
          &pm_id_embed_path},
+        {"",
+         "--hires-upscaler",
+         "highres fix upscaler, Lanczos, Nearest, Latent, Latent (nearest), Latent (nearest-exact), "
+         "Latent (antialiased), Latent (bicubic), Latent (bicubic antialiased), or a model name "
+         "under --hires-upscalers-dir (default: Latent)",
+         &hires_upscaler},
     };
 
     options.int_options = {
@@ -826,6 +849,22 @@ ArgOptions SDGenerationParams::get_options() {
          "--upscale-tile-size",
          "tile size for ESRGAN upscaling (default: 128)",
          &upscale_tile_size},
+        {"",
+         "--hires-width",
+         "highres fix target width, 0 to use --hires-scale (default: 0)",
+         &hires_width},
+        {"",
+         "--hires-height",
+         "highres fix target height, 0 to use --hires-scale (default: 0)",
+         &hires_height},
+        {"",
+         "--hires-steps",
+         "highres fix second pass sample steps, 0 to reuse --steps (default: 0)",
+         &hires_steps},
+        {"",
+         "--hires-upscale-tile-size",
+         "highres fix upscaler tile size, reserved for model-backed upscalers (default: 128)",
+         &hires_upscale_tile_size},
     };
 
     options.float_options = {
@@ -855,7 +894,7 @@ ArgOptions SDGenerationParams::get_options() {
          &sample_params.guidance.slg.layer_end},
         {"",
          "--eta",
-         "noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a and dpm++2s_a)",
+         "noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)",
          &sample_params.eta},
         {"",
          "--flow-shift",
@@ -887,7 +926,7 @@ ArgOptions SDGenerationParams::get_options() {
          &high_noise_sample_params.guidance.slg.layer_end},
         {"",
          "--high-noise-eta",
-         "(high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a and dpm++2s_a)",
+         "(high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)",
          &high_noise_sample_params.eta},
         {"",
          "--strength",
@@ -913,6 +952,14 @@ ArgOptions SDGenerationParams::get_options() {
          "--vae-tile-overlap",
          "tile overlap for vae tiling, in fraction of tile size (default: 0.5)",
          &vae_tiling_params.target_overlap},
+        {"",
+         "--hires-scale",
+         "highres fix scale when target size is not set (default: 2.0)",
+         &hires_scale},
+        {"",
+         "--hires-denoising-strength",
+         "highres fix second pass denoising strength (default: 0.7)",
+         &hires_denoising_strength},
     };
 
     options.bool_options = {
@@ -936,6 +983,11 @@ ArgOptions SDGenerationParams::get_options() {
          "process vae in tiles to reduce memory usage",
          true,
          &vae_tiling_params.enabled},
+        {"",
+         "--hires",
+         "enable highres fix",
+         true,
+         &hires_enabled},
     };
 
     auto on_seed_arg = [&](int argc, const char** argv, int index) {
@@ -1185,12 +1237,12 @@ ArgOptions SDGenerationParams::get_options() {
          on_seed_arg},
         {"",
          "--sampling-method",
-         "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s] "
+         "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde] "
          "(default: euler for Flux/SD3/Wan, euler_a otherwise)",
          on_sample_method_arg},
         {"",
          "--high-noise-sampling-method",
-         "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s]"
+         "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde]"
          " default: euler for Flux/SD3/Wan, euler_a otherwise",
          on_high_noise_sample_method_arg},
         {"",
@@ -1424,6 +1476,37 @@ static bool parse_lora_json_field(const json& parent,
     return true;
 }
 
+static bool resolve_model_file_from_dir(const std::string& model_name,
+                                        const std::string& model_dir,
+                                        const std::vector<std::string>& valid_ext,
+                                        const char* label,
+                                        std::string& resolved_path) {
+    if (model_dir.empty()) {
+        LOG_ERROR("%s directory is empty", label);
+        return false;
+    }
+    if (model_name.empty() ||
+        model_name.find('/') != std::string::npos ||
+        model_name.find('\\') != std::string::npos ||
+        fs::path(model_name).has_root_path() ||
+        fs::path(model_name).has_extension()) {
+        LOG_ERROR("%s must be a model name without path or extension: %s", label, model_name.c_str());
+        return false;
+    }
+
+    fs::path model_dir_path = model_dir;
+    for (const auto& ext : valid_ext) {
+        fs::path try_path = model_dir_path / (model_name + ext);
+        if (fs::exists(try_path) && fs::is_regular_file(try_path)) {
+            resolved_path = try_path.lexically_normal().string();
+            return true;
+        }
+    }
+
+    LOG_ERROR("can not find %s %s in %s", label, model_name.c_str(), model_dir_path.lexically_normal().string().c_str());
+    return false;
+}
+
 bool SDGenerationParams::from_json_str(
     const std::string& json_str,
     const std::function<std::string(const std::string&)>& lora_path_resolver) {
@@ -1487,6 +1570,34 @@ bool SDGenerationParams::from_json_str(
     load_if_exists("increase_ref_index", increase_ref_index);
     load_if_exists("embed_image_metadata", embed_image_metadata);
 
+    if (j.contains("hires") && j["hires"].is_object()) {
+        const json& hires_json = j["hires"];
+        if (hires_json.contains("enabled") && hires_json["enabled"].is_boolean()) {
+            hires_enabled = hires_json["enabled"];
+        }
+        if (hires_json.contains("upscaler") && hires_json["upscaler"].is_string()) {
+            hires_upscaler = hires_json["upscaler"];
+        }
+        if (hires_json.contains("scale") && hires_json["scale"].is_number()) {
+            hires_scale = hires_json["scale"];
+        }
+        if (hires_json.contains("target_width") && hires_json["target_width"].is_number_integer()) {
+            hires_width = hires_json["target_width"];
+        }
+        if (hires_json.contains("target_height") && hires_json["target_height"].is_number_integer()) {
+            hires_height = hires_json["target_height"];
+        }
+        if (hires_json.contains("steps") && hires_json["steps"].is_number_integer()) {
+            hires_steps = hires_json["steps"];
+        }
+        if (hires_json.contains("denoising_strength") && hires_json["denoising_strength"].is_number()) {
+            hires_denoising_strength = hires_json["denoising_strength"];
+        }
+        if (hires_json.contains("upscale_tile_size") && hires_json["upscale_tile_size"].is_number_integer()) {
+            hires_upscale_tile_size = hires_json["upscale_tile_size"];
+        }
+    }
+
     auto parse_sample_params_json = [&](const json& sample_json,
                                         sd_sample_params_t& target_params,
                                         std::vector<int>& target_skip_layers,
@@ -1589,10 +1700,18 @@ bool SDGenerationParams::from_json_str(
         LOG_ERROR("invalid init_image");
         return false;
     }
+    if (!parse_image_json_field(j, "end_image", 3, width, height, end_image)) {
+        LOG_ERROR("invalid end_image");
+        return false;
+    }
     if (!parse_image_array_json_field(j, "ref_images", 3, width, height, ref_images)) {
         LOG_ERROR("invalid ref_images");
         return false;
     }
+    if (!parse_image_array_json_field(j, "control_frames", 3, width, height, control_frames)) {
+        LOG_ERROR("invalid control_frames");
+        return false;
+    }
     if (!parse_image_json_field(j, "mask_image", 1, width, height, mask_image)) {
         LOG_ERROR("invalid mask_image");
         return false;
@@ -1792,7 +1911,7 @@ bool SDGenerationParams::initialize_cache_params() {
     return true;
 }
 
-bool SDGenerationParams::resolve(const std::string& lora_model_dir, bool strict) {
+bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict) {
     if (high_noise_sample_params.sample_steps <= 0) {
         high_noise_sample_params.sample_steps = -1;
     }
@@ -1811,6 +1930,27 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, bool strict)
         sample_params.sample_steps = std::clamp(sample_params.sample_steps, 1, 100);
     }
 
+    hires_upscaler_model_path.clear();
+    if (hires_enabled) {
+        if (hires_upscaler.empty()) {
+            hires_upscaler = "Latent";
+        }
+        resolved_hires_upscaler = str_to_sd_hires_upscaler(hires_upscaler.c_str());
+        if (resolved_hires_upscaler == SD_HIRES_UPSCALER_NONE) {
+            hires_enabled = false;
+        } else if (resolved_hires_upscaler == SD_HIRES_UPSCALER_COUNT) {
+            static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt", ".pth"};
+            if (!resolve_model_file_from_dir(hires_upscaler,
+                                             hires_upscalers_dir,
+                                             valid_ext,
+                                             "hires upscaler",
+                                             hires_upscaler_model_path)) {
+                return false;
+            }
+            resolved_hires_upscaler = SD_HIRES_UPSCALER_MODEL;
+        }
+    }
+
     prompt_with_lora = prompt;
     if (!lora_model_dir.empty()) {
         extract_and_remove_lora(lora_model_dir);
@@ -1875,6 +2015,29 @@ bool SDGenerationParams::validate(SDMode mode) {
         return false;
     }
 
+    if (hires_enabled) {
+        if (hires_width < 0 || hires_height < 0) {
+            LOG_ERROR("error: hires target width and height must be >= 0");
+            return false;
+        }
+        if (hires_scale <= 0.f && hires_width <= 0 && hires_height <= 0) {
+            LOG_ERROR("error: hires scale must be positive when target size is not set");
+            return false;
+        }
+        if (hires_steps < 0) {
+            LOG_ERROR("error: hires steps must be >= 0");
+            return false;
+        }
+        if (hires_denoising_strength <= 0.f || hires_denoising_strength > 1.f) {
+            LOG_ERROR("error: hires denoising strength must be in (0.0, 1.0]");
+            return false;
+        }
+        if (hires_upscale_tile_size < 1) {
+            LOG_ERROR("error: hires upscale tile size must be positive");
+            return false;
+        }
+    }
+
     if (mode == UPSCALE) {
         if (init_image_path.length() == 0) {
             LOG_ERROR("error: upscale mode needs an init image (--init-img)\n");
@@ -1885,8 +2048,11 @@ bool SDGenerationParams::validate(SDMode mode) {
     return true;
 }
 
-bool SDGenerationParams::resolve_and_validate(SDMode mode, const std::string& lora_model_dir, bool strict) {
+bool SDGenerationParams::resolve_and_validate(SDMode mode,
-    if (!resolve(lora_model_dir, strict)) {
+                                              const std::string& lora_model_dir,
+                                              const std::string& hires_upscalers_dir,
+                                              bool strict) {
+    if (!resolve(lora_model_dir, hires_upscalers_dir, strict)) {
         return false;
     }
     if (!validate(mode)) {
@@ -1957,6 +2123,16 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     params.pm_params             = pm_params;
     params.vae_tiling_params     = vae_tiling_params;
     params.cache                 = cache_params;
+
+    params.hires.enabled            = hires_enabled;
+    params.hires.upscaler           = resolved_hires_upscaler;
+    params.hires.model_path         = hires_upscaler_model_path.empty() ? nullptr : hires_upscaler_model_path.c_str();
+    params.hires.scale              = hires_scale;
+    params.hires.target_width       = hires_width;
+    params.hires.target_height      = hires_height;
+    params.hires.steps              = hires_steps;
+    params.hires.denoising_strength = hires_denoising_strength;
+    params.hires.upscale_tile_size  = hires_upscale_tile_size;
     return params;
 }
 
@@ -2081,6 +2257,15 @@ std::string SDGenerationParams::to_string() const {
         << "  seed: " << seed << ",\n"
         << "  upscale_repeats: " << upscale_repeats << ",\n"
         << "  upscale_tile_size: " << upscale_tile_size << ",\n"
+        << "  hires: { enabled: " << (hires_enabled ? "true" : "false")
+        << ", upscaler: \"" << hires_upscaler << "\""
+        << ", model_path: \"" << hires_upscaler_model_path << "\""
+        << ", scale: " << hires_scale
+        << ", target_width: " << hires_width
+        << ", target_height: " << hires_height
+        << ", steps: " << hires_steps
+        << ", denoising_strength: " << hires_denoising_strength
+        << ", upscale_tile_size: " << hires_upscale_tile_size << " },\n"
         << "  vae_tiling_params: { "
         << vae_tiling_params.enabled << ", "
         << vae_tiling_params.tile_size_x << ", "
@@ -2154,6 +2339,13 @@ std::string get_image_params(const SDContextParams& ctx_params, const SDGenerati
     if (gen_params.clip_skip != -1) {
         parameter_string += "Clip skip: " + std::to_string(gen_params.clip_skip) + ", ";
     }
+    if (gen_params.hires_enabled) {
+        parameter_string += "Hires upscale: " + gen_params.hires_upscaler + ", ";
+        parameter_string += "Hires scale: " + std::to_string(gen_params.hires_scale) + ", ";
+        parameter_string += "Hires resize: " + std::to_string(gen_params.hires_width) + "x" + std::to_string(gen_params.hires_height) + ", ";
+        parameter_string += "Hires steps: " + std::to_string(gen_params.hires_steps) + ", ";
+        parameter_string += "Denoising strength: " + std::to_string(gen_params.hires_denoising_strength) + ", ";
+    }
     parameter_string += "Version: stable-diffusion.cpp";
     return parameter_string;
 }

examples/common/common.h

@@ -101,6 +101,7 @@ struct SDContextParams {
     sd_type_t wtype = SD_TYPE_COUNT;
     std::string tensor_type_rules;
     std::string lora_model_dir = ".";
+    std::string hires_upscalers_dir;
 
     std::map<std::string, std::string> embedding_map;
     std::vector<sd_embedding_t> embedding_vec;
@@ -190,12 +191,23 @@ struct SDGenerationParams {
     int upscale_repeats   = 1;
     int upscale_tile_size = 128;
 
+    bool hires_enabled         = false;
+    std::string hires_upscaler = "Latent";
+    std::string hires_upscaler_model_path;
+    float hires_scale              = 2.f;
+    int hires_width                = 0;
+    int hires_height               = 0;
+    int hires_steps                = 0;
+    float hires_denoising_strength = 0.7f;
+    int hires_upscale_tile_size    = 128;
+
     std::map<std::string, float> lora_map;
     std::map<std::string, float> high_noise_lora_map;
 
     // Derived and normalized fields.
     std::string prompt_with_lora;  // for metadata record only
     std::vector<sd_lora_t> lora_vec;
+    sd_hires_upscaler_t resolved_hires_upscaler;
 
     // Owned execution payload.
     SDImageOwner init_image;
@@ -225,9 +237,12 @@ struct SDGenerationParams {
     void set_width_and_height_if_unset(int w, int h);
     int get_resolved_width() const;
     int get_resolved_height() const;
-    bool resolve(const std::string& lora_model_dir, bool strict = false);
+    bool resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict = false);
     bool validate(SDMode mode);
-    bool resolve_and_validate(SDMode mode, const std::string& lora_model_dir, bool strict = false);
+    bool resolve_and_validate(SDMode mode,
+                              const std::string& lora_model_dir,
+                              const std::string& hires_upscalers_dir,
+                              bool strict = false);
     sd_img_gen_params_t to_sd_img_gen_params_t();
     sd_vid_gen_params_t to_sd_vid_gen_params_t();
     std::string to_string() const;

examples/common/media_io.cpp

@@ -95,6 +95,57 @@ using WebPMuxPtr         = std::unique_ptr<WebPMux, WebPMuxDeleter>;
 using WebPAnimEncoderPtr = std::unique_ptr<WebPAnimEncoder, WebPAnimEncoderDeleter>;
 #endif
 
+#ifdef SD_USE_WEBM
+class MemoryMkvWriter : public mkvmuxer::IMkvWriter {
+public:
+    mkvmuxer::int32 Write(const void* buf, mkvmuxer::uint32 len) override {
+        if (buf == nullptr && len > 0) {
+            return -1;
+        }
+        const size_t end_pos = position_ + static_cast<size_t>(len);
+        if (end_pos > data_.size()) {
+            data_.resize(end_pos);
+        }
+        if (len > 0) {
+            memcpy(data_.data() + position_, buf, len);
+        }
+        position_ = end_pos;
+        return 0;
+    }
+
+    mkvmuxer::int64 Position() const override {
+        return static_cast<mkvmuxer::int64>(position_);
+    }
+
+    mkvmuxer::int32 Position(mkvmuxer::int64 position) override {
+        if (position < 0) {
+            return -1;
+        }
+        const size_t target = static_cast<size_t>(position);
+        if (target > data_.size()) {
+            data_.resize(target);
+        }
+        position_ = target;
+        return 0;
+    }
+
+    bool Seekable() const override {
+        return true;
+    }
+
+    void ElementStartNotify(mkvmuxer::uint64, mkvmuxer::int64) override {
+    }
+
+    const std::vector<uint8_t>& data() const {
+        return data_;
+    }
+
+private:
+    std::vector<uint8_t> data_;
+    size_t position_ = 0;
+};
+#endif
+
 bool read_binary_file_bytes(const char* path, std::vector<uint8_t>& data) {
     std::ifstream fin(fs::path(path), std::ios::binary);
     if (!fin) {
@@ -570,6 +621,32 @@ void write_u16_le(FILE* f, uint16_t val) {
     fwrite(&val, 2, 1, f);
 }
 
+void write_u32_le(std::vector<uint8_t>& data, uint32_t val) {
+    data.push_back(static_cast<uint8_t>(val & 0xFF));
+    data.push_back(static_cast<uint8_t>((val >> 8) & 0xFF));
+    data.push_back(static_cast<uint8_t>((val >> 16) & 0xFF));
+    data.push_back(static_cast<uint8_t>((val >> 24) & 0xFF));
+}
+
+void write_u16_le(std::vector<uint8_t>& data, uint16_t val) {
+    data.push_back(static_cast<uint8_t>(val & 0xFF));
+    data.push_back(static_cast<uint8_t>((val >> 8) & 0xFF));
+}
+
+void patch_u32_le(std::vector<uint8_t>& data, size_t offset, uint32_t val) {
+    if (offset + 4 > data.size()) {
+        return;
+    }
+    data[offset + 0] = static_cast<uint8_t>(val & 0xFF);
+    data[offset + 1] = static_cast<uint8_t>((val >> 8) & 0xFF);
+    data[offset + 2] = static_cast<uint8_t>((val >> 16) & 0xFF);
+    data[offset + 3] = static_cast<uint8_t>((val >> 24) & 0xFF);
+}
+
+void write_fourcc(std::vector<uint8_t>& data, const char* fourcc) {
+    data.insert(data.end(), fourcc, fourcc + 4);
+}
+
 EncodedImageFormat encoded_image_format_from_path(const std::string& path) {
     std::string ext = fs::path(path).extension().string();
     std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
@@ -699,95 +776,96 @@ uint8_t* load_image_from_memory(const char* image_bytes,
     return load_image_common(true, image_bytes, len, width, height, expected_width, expected_height, expected_channel);
 }
 
-int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
+std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality) {
     if (num_images == 0) {
         fprintf(stderr, "Error: Image array is empty.\n");
-        return -1;
+        return {};
     }
 
-    FilePtr file(fopen(filename, "wb"));
-    if (!file) {
-        perror("Error opening file for writing");
-        return -1;
-    }
-    FILE* f = file.get();
-
     uint32_t width    = images[0].width;
     uint32_t height   = images[0].height;
     uint32_t channels = images[0].channel;
     if (channels != 3 && channels != 4) {
         fprintf(stderr, "Error: Unsupported channel count: %u\n", channels);
-        return -1;
+        return {};
     }
 
-    fwrite("RIFF", 4, 1, f);
+    // stb_image_write changes JPEG sampling behavior above quality 90.
-    long riff_size_pos = ftell(f);
+    // MJPG AVI playback is more compatible when we keep the encoder on the
-    write_u32_le(f, 0);
+    // <= 90 path.
-    fwrite("AVI ", 4, 1, f);
+    const int mjpg_quality = std::clamp(quality, 1, 90);
 
-    fwrite("LIST", 4, 1, f);
+    std::vector<uint8_t> avi_data;
-    write_u32_le(f, 4 + 8 + 56 + 8 + 4 + 8 + 56 + 8 + 40);
+    avi_data.reserve(static_cast<size_t>(num_images) * 1024);
-    fwrite("hdrl", 4, 1, f);
 
-    fwrite("avih", 4, 1, f);
+    write_fourcc(avi_data, "RIFF");
-    write_u32_le(f, 56);
+    const size_t riff_size_pos = avi_data.size();
-    write_u32_le(f, 1000000 / fps);
+    write_u32_le(avi_data, 0);
-    write_u32_le(f, 0);
+    write_fourcc(avi_data, "AVI ");
-    write_u32_le(f, 0);
-    write_u32_le(f, 0x110);
-    write_u32_le(f, num_images);
-    write_u32_le(f, 0);
-    write_u32_le(f, 1);
-    write_u32_le(f, width * height * 3);
-    write_u32_le(f, width);
-    write_u32_le(f, height);
-    write_u32_le(f, 0);
-    write_u32_le(f, 0);
-    write_u32_le(f, 0);
-    write_u32_le(f, 0);
 
-    fwrite("LIST", 4, 1, f);
+    write_fourcc(avi_data, "LIST");
-    write_u32_le(f, 4 + 8 + 56 + 8 + 40);
+    write_u32_le(avi_data, 4 + 8 + 56 + 8 + 4 + 8 + 56 + 8 + 40);
-    fwrite("strl", 4, 1, f);
+    write_fourcc(avi_data, "hdrl");
 
-    fwrite("strh", 4, 1, f);
+    write_fourcc(avi_data, "avih");
-    write_u32_le(f, 56);
+    write_u32_le(avi_data, 56);
-    fwrite("vids", 4, 1, f);
+    write_u32_le(avi_data, 1000000 / fps);
-    fwrite("MJPG", 4, 1, f);
+    write_u32_le(avi_data, 0);
-    write_u32_le(f, 0);
+    write_u32_le(avi_data, 0);
-    write_u16_le(f, 0);
+    write_u32_le(avi_data, 0x110);
-    write_u16_le(f, 0);
+    write_u32_le(avi_data, num_images);
-    write_u32_le(f, 0);
+    write_u32_le(avi_data, 0);
-    write_u32_le(f, 1);
+    write_u32_le(avi_data, 1);
-    write_u32_le(f, fps);
+    write_u32_le(avi_data, width * height * 3);
-    write_u32_le(f, 0);
+    write_u32_le(avi_data, width);
-    write_u32_le(f, num_images);
+    write_u32_le(avi_data, height);
-    write_u32_le(f, width * height * 3);
+    write_u32_le(avi_data, 0);
-    write_u32_le(f, (uint32_t)-1);
+    write_u32_le(avi_data, 0);
-    write_u32_le(f, 0);
+    write_u32_le(avi_data, 0);
-    write_u16_le(f, 0);
+    write_u32_le(avi_data, 0);
-    write_u16_le(f, 0);
-    write_u16_le(f, 0);
-    write_u16_le(f, 0);
 
-    fwrite("strf", 4, 1, f);
+    write_fourcc(avi_data, "LIST");
-    write_u32_le(f, 40);
+    write_u32_le(avi_data, 4 + 8 + 56 + 8 + 40);
-    write_u32_le(f, 40);
+    write_fourcc(avi_data, "strl");
-    write_u32_le(f, width);
-    write_u32_le(f, height);
-    write_u16_le(f, 1);
-    write_u16_le(f, 24);
-    fwrite("MJPG", 4, 1, f);
-    write_u32_le(f, width * height * 3);
-    write_u32_le(f, 0);
-    write_u32_le(f, 0);
-    write_u32_le(f, 0);
-    write_u32_le(f, 0);
 
-    fwrite("LIST", 4, 1, f);
+    write_fourcc(avi_data, "strh");
-    long movi_size_pos = ftell(f);
+    write_u32_le(avi_data, 56);
-    write_u32_le(f, 0);
+    write_fourcc(avi_data, "vids");
-    fwrite("movi", 4, 1, f);
+    write_fourcc(avi_data, "MJPG");
+    write_u32_le(avi_data, 0);
+    write_u16_le(avi_data, 0);
+    write_u16_le(avi_data, 0);
+    write_u32_le(avi_data, 0);
+    write_u32_le(avi_data, 1);
+    write_u32_le(avi_data, fps);
+    write_u32_le(avi_data, 0);
+    write_u32_le(avi_data, num_images);
+    write_u32_le(avi_data, width * height * 3);
+    write_u32_le(avi_data, static_cast<uint32_t>(-1));
+    write_u32_le(avi_data, 0);
+    write_u16_le(avi_data, 0);
+    write_u16_le(avi_data, 0);
+    write_u16_le(avi_data, 0);
+    write_u16_le(avi_data, 0);
+
+    write_fourcc(avi_data, "strf");
+    write_u32_le(avi_data, 40);
+    write_u32_le(avi_data, 40);
+    write_u32_le(avi_data, width);
+    write_u32_le(avi_data, height);
+    write_u16_le(avi_data, 1);
+    write_u16_le(avi_data, 24);
+    write_fourcc(avi_data, "MJPG");
+    write_u32_le(avi_data, width * height * 3);
+    write_u32_le(avi_data, 0);
+    write_u32_le(avi_data, 0);
+    write_u32_le(avi_data, 0);
+    write_u32_le(avi_data, 0);
+
+    write_fourcc(avi_data, "LIST");
+    const size_t movi_size_pos = avi_data.size();
+    write_u32_le(avi_data, 0);
+    write_fourcc(avi_data, "movi");
 
     std::vector<avi_index_entry> index(static_cast<size_t>(num_images));
     std::vector<uint8_t> jpeg_data;
@@ -801,55 +879,61 @@ int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int
             buffer->insert(buffer->end(), src, src + size);
         };
 
-        if (!stbi_write_jpg_to_func(write_to_buf, &jpeg_data, images[i].width, images[i].height, channels, images[i].data, quality)) {
+        if (!stbi_write_jpg_to_func(write_to_buf, &jpeg_data, images[i].width, images[i].height, channels, images[i].data, mjpg_quality)) {
             fprintf(stderr, "Error: Failed to encode JPEG frame.\n");
-            return -1;
+            return {};
         }
 
-        fwrite("00dc", 4, 1, f);
+        index[i].offset = static_cast<uint32_t>(avi_data.size());
-        write_u32_le(f, (uint32_t)jpeg_data.size());
+        write_fourcc(avi_data, "00dc");
-        index[i].offset = ftell(f) - 8;
+        write_u32_le(avi_data, static_cast<uint32_t>(jpeg_data.size()));
         index[i].size = (uint32_t)jpeg_data.size();
-        fwrite(jpeg_data.data(), 1, jpeg_data.size(), f);
+        avi_data.insert(avi_data.end(), jpeg_data.begin(), jpeg_data.end());
 
         if (jpeg_data.size() % 2) {
-            fputc(0, f);
+            avi_data.push_back(0);
         }
     }
 
-    long cur_pos   = ftell(f);
+    const size_t movi_size = avi_data.size() - movi_size_pos - 4;
-    long movi_size = cur_pos - movi_size_pos - 4;
+    patch_u32_le(avi_data, movi_size_pos, static_cast<uint32_t>(movi_size));
-    fseek(f, movi_size_pos, SEEK_SET);
-    write_u32_le(f, movi_size);
-    fseek(f, cur_pos, SEEK_SET);
 
-    fwrite("idx1", 4, 1, f);
+    write_fourcc(avi_data, "idx1");
-    write_u32_le(f, num_images * 16);
+    write_u32_le(avi_data, num_images * 16);
     for (int i = 0; i < num_images; i++) {
-        fwrite("00dc", 4, 1, f);
+        write_fourcc(avi_data, "00dc");
-        write_u32_le(f, 0x10);
+        write_u32_le(avi_data, 0x10);
-        write_u32_le(f, index[i].offset);
+        write_u32_le(avi_data, index[i].offset);
-        write_u32_le(f, index[i].size);
+        write_u32_le(avi_data, index[i].size);
     }
 
-    cur_pos        = ftell(f);
+    const size_t file_size = avi_data.size() - riff_size_pos - 4;
-    long file_size = cur_pos - riff_size_pos - 4;
+    patch_u32_le(avi_data, riff_size_pos, static_cast<uint32_t>(file_size));
-    fseek(f, riff_size_pos, SEEK_SET);
-    write_u32_le(f, file_size);
-    fseek(f, cur_pos, SEEK_SET);
 
+    return avi_data;
+}
+
+int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
+    std::vector<uint8_t> avi_data = create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality);
+    if (avi_data.empty()) {
+        return -1;
+    }
+    if (!write_binary_file_bytes(filename, avi_data)) {
+        perror("Error opening file for writing");
+        return -1;
+    }
     return 0;
 }
 
 #ifdef SD_USE_WEBP
-int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
+std::vector<uint8_t> create_animated_webp_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality) {
     if (num_images == 0) {
         fprintf(stderr, "Error: Image array is empty.\n");
-        return -1;
+        return {};
     }
     if (fps <= 0) {
         fprintf(stderr, "Error: FPS must be positive.\n");
-        return -1;
+        return {};
     }
 
     const int width    = static_cast<int>(images[0].width);
@@ -857,14 +941,14 @@ int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images
     const int channels = static_cast<int>(images[0].channel);
     if (channels != 1 && channels != 3 && channels != 4) {
         fprintf(stderr, "Error: Unsupported channel count: %d\n", channels);
-        return -1;
+        return {};
     }
 
     WebPAnimEncoderOptions anim_options;
     WebPConfig config;
     if (!WebPAnimEncoderOptionsInit(&anim_options) || !WebPConfigInit(&config)) {
         fprintf(stderr, "Error: Failed to initialize WebP animation encoder.\n");
-        return -1;
+        return {};
     }
 
     config.quality      = static_cast<float>(quality);
@@ -875,13 +959,13 @@ int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images
     }
     if (!WebPValidateConfig(&config)) {
         fprintf(stderr, "Error: Invalid WebP encoder configuration.\n");
-        return -1;
+        return {};
     }
 
     WebPAnimEncoderPtr enc(WebPAnimEncoderNew(width, height, &anim_options));
     if (enc == nullptr) {
         fprintf(stderr, "Error: Could not create WebPAnimEncoder object.\n");
-        return -1;
+        return {};
     }
 
     const int frame_duration_ms = std::max(1, static_cast<int>(std::lround(1000.0 / static_cast<double>(fps))));
@@ -891,13 +975,13 @@ int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images
         const sd_image_t& image = images[i];
         if (static_cast<int>(image.width) != width || static_cast<int>(image.height) != height) {
             fprintf(stderr, "Error: Frame dimensions do not match.\n");
-            return -1;
+            return {};
         }
 
         WebPPictureGuard picture;
         if (!picture.initialized) {
             fprintf(stderr, "Error: Failed to initialize WebPPicture.\n");
-            return -1;
+            return {};
         }
         picture.picture.use_argb = 1;
         picture.picture.width    = width;
@@ -921,12 +1005,12 @@ int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images
 
         if (!picture_ok) {
             fprintf(stderr, "Error: Failed to import frame into WebPPicture.\n");
-            return -1;
+            return {};
         }
 
         if (!WebPAnimEncoderAdd(enc.get(), &picture.picture, timestamp_ms, &config)) {
             fprintf(stderr, "Error: Failed to add frame to animated WebP: %s\n", WebPAnimEncoderGetError(enc.get()));
-            return -1;
+            return {};
         }
 
         timestamp_ms += frame_duration_ms;
@@ -934,52 +1018,50 @@ int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images
 
     if (!WebPAnimEncoderAdd(enc.get(), nullptr, timestamp_ms, nullptr)) {
         fprintf(stderr, "Error: Failed to finalize animated WebP frames: %s\n", WebPAnimEncoderGetError(enc.get()));
-        return -1;
+        return {};
     }
 
     WebPDataGuard webp_data;
     if (!WebPAnimEncoderAssemble(enc.get(), &webp_data.data)) {
         fprintf(stderr, "Error: Failed to assemble animated WebP: %s\n", WebPAnimEncoderGetError(enc.get()));
-        return -1;
+        return {};
     }
 
-    FilePtr f(fopen(filename, "wb"));
+    return std::vector<uint8_t>(webp_data.data.bytes, webp_data.data.bytes + webp_data.data.size);
-    if (!f) {
+}
+
+int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
+    std::vector<uint8_t> webp_data = create_animated_webp_from_sd_images_to_vector(images, num_images, fps, quality);
+    if (webp_data.empty()) {
+        return -1;
+    }
+    if (!write_binary_file_bytes(filename, webp_data)) {
         perror("Error opening file for writing");
         return -1;
     }
-    if (webp_data.data.size > 0 && fwrite(webp_data.data.bytes, 1, webp_data.data.size, f.get()) != webp_data.data.size) {
-        fprintf(stderr, "Error: Failed to write animated WebP file.\n");
-        return -1;
-    }
-
     return 0;
 }
 #endif
 
 #ifdef SD_USE_WEBM
-int create_webm_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
+std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality) {
     if (num_images == 0) {
         fprintf(stderr, "Error: Image array is empty.\n");
-        return -1;
+        return {};
     }
     if (fps <= 0) {
         fprintf(stderr, "Error: FPS must be positive.\n");
-        return -1;
+        return {};
     }
 
     const int width  = static_cast<int>(images[0].width);
     const int height = static_cast<int>(images[0].height);
     if (width <= 0 || height <= 0) {
         fprintf(stderr, "Error: Invalid frame dimensions.\n");
-        return -1;
+        return {};
     }
 
-    mkvmuxer::MkvWriter writer;
+    MemoryMkvWriter writer;
-    if (!writer.Open(filename)) {
-        fprintf(stderr, "Error: Could not open WebM file for writing.\n");
-        return -1;
-    }
 
     const int ret = [&]() -> int {
         mkvmuxer::Segment segment;
@@ -1045,30 +1127,63 @@ int create_webm_from_sd_images(const char* filename, sd_image_t* images, int num
         }
         return 0;
     }();
-    writer.Close();
+    if (ret != 0) {
-    return ret;
+        return {};
+    }
+    return writer.data();
+}
+
+int create_webm_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
+    std::vector<uint8_t> webm_data = create_webm_from_sd_images_to_vector(images, num_images, fps, quality);
+    if (webm_data.empty()) {
+        return -1;
+    }
+    if (!write_binary_file_bytes(filename, webm_data)) {
+        perror("Error opening file for writing");
+        return -1;
+    }
+    return 0;
 }
 #endif
 
+std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& output_format,
+                                                           sd_image_t* images,
+                                                           int num_images,
+                                                           int fps,
+                                                           int quality) {
+    std::string format = output_format;
+    std::transform(format.begin(), format.end(), format.begin(),
+                   [](unsigned char c) { return static_cast<char>(tolower(c)); });
+    if (!format.empty() && format[0] == '.') {
+        format.erase(format.begin());
+    }
+
+#ifdef SD_USE_WEBM
+    if (format == "webm") {
+        return create_webm_from_sd_images_to_vector(images, num_images, fps, quality);
+    }
+#endif
+
+#ifdef SD_USE_WEBP
+    if (format == "webp") {
+        return create_animated_webp_from_sd_images_to_vector(images, num_images, fps, quality);
+    }
+#endif
+
+    return create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality);
+}
+
 int create_video_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
     std::string path                = filename ? filename : "";
     auto pos                        = path.find_last_of('.');
     std::string ext                 = pos == std::string::npos ? "" : path.substr(pos);
-    for (char& ch : ext) {
+    std::vector<uint8_t> video_data = create_video_from_sd_images_to_vector(ext, images, num_images, fps, quality);
-        ch = static_cast<char>(tolower(static_cast<unsigned char>(ch)));
+    if (video_data.empty()) {
+        return -1;
     }
-
+    if (!write_binary_file_bytes(filename, video_data)) {
-#ifdef SD_USE_WEBM
+        perror("Error opening file for writing");
-    if (ext == ".webm") {
+        return -1;
-        return create_webm_from_sd_images(filename, images, num_images, fps, quality);
     }
-#endif
+    return 0;
-
-#ifdef SD_USE_WEBP
-    if (ext == ".webp") {
-        return create_animated_webp_from_sd_images(filename, images, num_images, fps, quality);
-    }
-#endif
-
-    return create_mjpg_avi_from_sd_images(filename, images, num_images, fps, quality);
 }

examples/common/media_io.h

@@ -58,6 +58,10 @@ int create_mjpg_avi_from_sd_images(const char* filename,
                                    int num_images,
                                    int fps,
                                    int quality = 90);
+std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images,
+                                                              int num_images,
+                                                              int fps,
+                                                              int quality = 90);
 
 #ifdef SD_USE_WEBP
 int create_animated_webp_from_sd_images(const char* filename,
@@ -65,6 +69,10 @@ int create_animated_webp_from_sd_images(const char* filename,
                                         int num_images,
                                         int fps,
                                         int quality = 90);
+std::vector<uint8_t> create_animated_webp_from_sd_images_to_vector(sd_image_t* images,
+                                                                   int num_images,
+                                                                   int fps,
+                                                                   int quality = 90);
 #endif
 
 #ifdef SD_USE_WEBM
@@ -73,6 +81,10 @@ int create_webm_from_sd_images(const char* filename,
                                int num_images,
                                int fps,
                                int quality = 90);
+std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images,
+                                                          int num_images,
+                                                          int fps,
+                                                          int quality = 90);
 #endif
 
 int create_video_from_sd_images(const char* filename,
@@ -80,5 +92,10 @@ int create_video_from_sd_images(const char* filename,
                                 int num_images,
                                 int fps,
                                 int quality = 90);
+std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& output_format,
+                                                           sd_image_t* images,
+                                                           int num_images,
+                                                           int fps,
+                                                           int quality = 90);
 
 #endif  // __MEDIA_IO_H__

examples/server/CMakeLists.txt

@@ -50,7 +50,13 @@ if(SD_SERVER_BUILD_FRONTEND AND EXISTS "${FRONTEND_DIR}")
 
         set_source_files_properties("${GENERATED_HTML_HEADER}" PROPERTIES GENERATED TRUE)
     else()
-        message(WARNING "pnpm not found, frontend build disabled")
+        if(EXISTS "${GENERATED_HTML_HEADER}")
+            message(STATUS "pnpm not found; using pre-built frontend header detected at ${GENERATED_HTML_HEADER}")
+            set(HAVE_FRONTEND_BUILD ON)
+            add_custom_target(${TARGET}_frontend)
+        else()
+            message(WARNING "pnpm not found; frontend build disabled.")
+        endif()
     endif()
 else()
     message(STATUS "Frontend disabled or directory not found: ${FRONTEND_DIR}")

examples/server/README.md

@@ -136,7 +136,8 @@ Context Options:
   --clip_g <string>                        path to the clip-g text encoder
   --clip_vision <string>                   path to the clip-vision encoder
   --t5xxl <string>                         path to the t5xxl text encoder
-  --llm <string>                           path to the llm text encoder. For example: (qwenvl2.5 for qwen-image, mistral-small3.2 for flux2, ...)
+  --llm <string>                           path to the llm text encoder. For example: (qwenvl2.5 for qwen-image,
+                                           mistral-small3.2 for flux2, ...)
   --llm_vision <string>                    path to the llm vit
   --qwen2vl <string>                       alias of --llm. Deprecated.
   --qwen2vl_vision <string>                alias of --llm_vision. Deprecated.
@@ -148,16 +149,16 @@ Context Options:
   --control-net <string>                   path to control net model
   --embd-dir <string>                      embeddings directory
   --lora-model-dir <string>                lora model directory
+  --hires-upscalers-dir <string>           highres fix upscaler model directory
   --tensor-type-rules <string>             weight type per tensor pattern (example: "^vae\.=f16,model\.=q8_0")
   --photo-maker <string>                   path to PHOTOMAKER model
   --upscale-model <string>                 path to esrgan model.
-  -t, --threads <int>                      number of threads to use during computation (default: -1). If threads <= 0, then threads will be set to the number of
+  -t, --threads <int>                      number of threads to use during computation (default: -1). If threads <= 0,
-                                           CPU physical cores
+                                           then threads will be set to the number of CPU physical cores
   --chroma-t5-mask-pad <int>               t5 mask pad size of chroma
-  --vae-tile-overlap <float>               tile overlap for vae tiling, in fraction of tile size (default: 0.5)
-  --vae-tiling                             process vae in tiles to reduce memory usage
   --force-sdxl-vae-conv-scale              force use of conv scale on sdxl vae
-  --offload-to-cpu                         place the weights in RAM to save VRAM, and automatically load them into VRAM when needed
+  --offload-to-cpu                         place the weights in RAM to save VRAM, and automatically load them into VRAM
+                                           when needed
   --mmap                                   whether to memory-map model
   --control-net-cpu                        keep controlnet in cpu (for low vram)
   --clip-on-cpu                            keep clip in cpu (for low vram)
@@ -172,20 +173,19 @@ Context Options:
   --chroma-disable-dit-mask                disable dit mask for chroma
   --qwen-image-zero-cond-t                 enable zero_cond_t for qwen image
   --chroma-enable-t5-mask                  enable t5 mask for chroma
-  --type                                   weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, q4_K). If not specified, the default is the
+  --type                                   weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K,
-                                           type of the weight file
+                                           q4_K). If not specified, the default is the type of the weight file
   --rng                                    RNG, one of [std_default, cuda, cpu], default: cuda(sd-webui), cpu(comfyui)
   --sampler-rng                            sampler RNG, one of [std_default, cuda, cpu]. If not specified, use --rng
-  --prediction                             prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, flux2_flow]
+  --prediction                             prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow,
-  --lora-apply-mode                        the way to apply LoRA, one of [auto, immediately, at_runtime], default is auto. In auto mode, if the model weights
+                                           flux2_flow]
-                                           contain any quantized parameters, the at_runtime mode will be used; otherwise,
+  --lora-apply-mode                        the way to apply LoRA, one of [auto, immediately, at_runtime], default is
-                                           immediately will be used.The immediately mode may have precision and
+                                           auto. In auto mode, if the model weights contain any quantized parameters,
-                                           compatibility issues with quantized parameters, but it usually offers faster inference
+                                           the at_runtime mode will be used; otherwise, immediately will be used.The
-                                           speed and, in some cases, lower memory usage. The at_runtime mode, on the
+                                           immediately mode may have precision and compatibility issues with quantized
-                                           other hand, is exactly the opposite.
+                                           parameters, but it usually offers faster inference speed and, in some cases,
-  --vae-tile-size                          tile size for vae tiling, format [X]x[Y] (default: 32x32)
+                                           lower memory usage. The at_runtime mode, on the other hand, is exactly the
-  --vae-relative-tile-size                 relative tile size for vae tiling, format [X]x[Y], in fraction of image size if < 1, in number of tiles per dim if >=1
+                                           opposite.
-                                           (overrides --vae-tile-size)
 
 Default Generation Options:
   -p, --prompt <string>                    the prompt to render
@@ -194,65 +194,97 @@ Default Generation Options:
   --end-img <string>                       path to the end image, required by flf2v
   --mask <string>                          path to the mask image
   --control-image <string>                 path to control image, control net
-  --control-video <string>                 path to control video frames, It must be a directory path. The video frames inside should be stored as images in
+  --control-video <string>                 path to control video frames, It must be a directory path. The video frames
-                                           lexicographical (character) order. For example, if the control video path is
+                                           inside should be stored as images in lexicographical (character) order. For
-                                           `frames`, the directory contain images such as 00.png, 01.png, ... etc.
+                                           example, if the control video path is `frames`, the directory contain images
+                                           such as 00.png, 01.png, ... etc.
   --pm-id-images-dir <string>              path to PHOTOMAKER input id images dir
   --pm-id-embed-path <string>              path to PHOTOMAKER v2 id embed
+  --hires-upscaler <string>                highres fix upscaler, Lanczos, Nearest, Latent, Latent (nearest), Latent
+                                           (nearest-exact), Latent (antialiased), Latent (bicubic), Latent (bicubic
+                                           antialiased), or a model name under --hires-upscalers-dir (default: Latent)
   -H, --height <int>                       image height, in pixel space (default: 512)
   -W, --width <int>                        image width, in pixel space (default: 512)
   --steps <int>                            number of sample steps (default: 20)
   --high-noise-steps <int>                 (high noise) number of sample steps (default: -1 = auto)
-  --clip-skip <int>                        ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1). <= 0 represents unspecified,
+  --clip-skip <int>                        ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer
-                                           will be 1 for SD1.x, 2 for SD2.x
+                                           (default: -1). <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x
   -b, --batch-count <int>                  batch count
   --video-frames <int>                     video frames (default: 1)
   --fps <int>                              fps (default: 24)
-  --timestep-shift <int>                   shift timestep for NitroFusion models (default: 0). recommended N for NitroSD-Realism around 250 and 500 for
+  --timestep-shift <int>                   shift timestep for NitroFusion models (default: 0). recommended N for
-                                           NitroSD-Vibrant
+                                           NitroSD-Realism around 250 and 500 for NitroSD-Vibrant
   --upscale-repeats <int>                  Run the ESRGAN upscaler this many times (default: 1)
   --upscale-tile-size <int>                tile size for ESRGAN upscaling (default: 128)
+  --hires-width <int>                      highres fix target width, 0 to use --hires-scale (default: 0)
+  --hires-height <int>                     highres fix target height, 0 to use --hires-scale (default: 0)
+  --hires-steps <int>                      highres fix second pass sample steps, 0 to reuse --steps (default: 0)
+  --hires-upscale-tile-size <int>          highres fix upscaler tile size, reserved for model-backed upscalers (default:
+                                           128)
   --cfg-scale <float>                      unconditional guidance scale: (default: 7.0)
-  --img-cfg-scale <float>                  image guidance scale for inpaint or instruct-pix2pix models: (default: same as --cfg-scale)
+  --img-cfg-scale <float>                  image guidance scale for inpaint or instruct-pix2pix models: (default: same
+                                           as --cfg-scale)
   --guidance <float>                       distilled guidance scale for models with guidance input (default: 3.5)
-  --slg-scale <float>                      skip layer guidance (SLG) scale, only for DiT models: (default: 0). 0 means disabled, a value of 2.5 is nice for sd3.5
+  --slg-scale <float>                      skip layer guidance (SLG) scale, only for DiT models: (default: 0). 0 means
-                                           medium
+                                           disabled, a value of 2.5 is nice for sd3.5 medium
   --skip-layer-start <float>               SLG enabling point (default: 0.01)
   --skip-layer-end <float>                 SLG disabling point (default: 0.2)
-  --eta <float>                            noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a and dpm++2s_a)
+  --eta <float>                            noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and
+                                           res_2s; 1 for euler_a, er_sde and dpm++2s_a)
   --flow-shift <float>                     shift value for Flow models like SD3.x or WAN (default: auto)
   --high-noise-cfg-scale <float>           (high noise) unconditional guidance scale: (default: 7.0)
-  --high-noise-img-cfg-scale <float>       (high noise) image guidance scale for inpaint or instruct-pix2pix models (default: same as --cfg-scale)
+  --high-noise-img-cfg-scale <float>       (high noise) image guidance scale for inpaint or instruct-pix2pix models
-  --high-noise-guidance <float>            (high noise) distilled guidance scale for models with guidance input (default: 3.5)
+                                           (default: same as --cfg-scale)
-  --high-noise-slg-scale <float>           (high noise) skip layer guidance (SLG) scale, only for DiT models: (default: 0)
+  --high-noise-guidance <float>            (high noise) distilled guidance scale for models with guidance input
+                                           (default: 3.5)
+  --high-noise-slg-scale <float>           (high noise) skip layer guidance (SLG) scale, only for DiT models: (default:
+                                           0)
   --high-noise-skip-layer-start <float>    (high noise) SLG enabling point (default: 0.01)
   --high-noise-skip-layer-end <float>      (high noise) SLG disabling point (default: 0.2)
-  --high-noise-eta <float>                 (high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a and dpm++2s_a)
+  --high-noise-eta <float>                 (high noise) noise multiplier (default: 0 for ddim_trailing, tcd,
+                                           res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)
   --strength <float>                       strength for noising/unnoising (default: 0.75)
   --pm-style-strength <float>              
-  --control-strength <float>               strength to apply Control Net (default: 0.9). 1.0 corresponds to full destruction of information in init image
+  --control-strength <float>               strength to apply Control Net (default: 0.9). 1.0 corresponds to full
-  --moe-boundary <float>                   timestep boundary for Wan2.2 MoE model. (default: 0.875). Only enabled if `--high-noise-steps` is set to -1
+                                           destruction of information in init image
+  --moe-boundary <float>                   timestep boundary for Wan2.2 MoE model. (default: 0.875). Only enabled if
+                                           `--high-noise-steps` is set to -1
   --vace-strength <float>                  wan vace strength
-  --increase-ref-index                     automatically increase the indices of references images based on the order they are listed (starting with 1).
+  --vae-tile-overlap <float>               tile overlap for vae tiling, in fraction of tile size (default: 0.5)
+  --hires-scale <float>                    highres fix scale when target size is not set (default: 2.0)
+  --hires-denoising-strength <float>       highres fix second pass denoising strength (default: 0.7)
+  --increase-ref-index                     automatically increase the indices of references images based on the order
+                                           they are listed (starting with 1).
   --disable-auto-resize-ref-image          disable auto resize of ref images
   --disable-image-metadata                 do not embed generation metadata on image files
+  --vae-tiling                             process vae in tiles to reduce memory usage
+  --hires                                  enable highres fix
   -s, --seed                               RNG seed (default: 42, use random seed for < 0)
-  --sampling-method                        sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing,
+  --sampling-method                        sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m,
-                                           tcd, res_multistep, res_2s] (default: euler for Flux/SD3/Wan, euler_a
+                                           dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s,
-                                           otherwise)
+                                           er_sde] (default: euler for Flux/SD3/Wan, euler_a otherwise)
-  --high-noise-sampling-method             (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm,
+  --high-noise-sampling-method             (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a,
-                                           ddim_trailing, tcd, res_multistep, res_2s] default: euler for Flux/SD3/Wan,
+                                           dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep,
-                                           euler_a otherwise
+                                           res_2s, er_sde] default: euler for Flux/SD3/Wan, euler_a otherwise
-  --scheduler                              denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple,
+  --scheduler                              denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits,
-                                           kl_optimal, lcm, bong_tangent], default: discrete
+                                           smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent], default:
-  --sigmas                                 custom sigma values for the sampler, comma-separated (e.g., "14.61,7.8,3.5,0.0").
+                                           discrete
+  --sigmas                                 custom sigma values for the sampler, comma-separated (e.g.,
+                                           "14.61,7.8,3.5,0.0").
   --skip-layers                            layers to skip for SLG steps (default: [7,8,9])
   --high-noise-skip-layers                 (high noise) layers to skip for SLG steps (default: [7,8,9])
   -r, --ref-image                          reference image for Flux Kontext models (can be used multiple times)
-  --cache-mode                             caching method: 'easycache' (DiT), 'ucache' (UNET), 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT Chebyshev+Taylor forecasting)
+  --cache-mode                             caching method: 'easycache' (DiT), 'ucache' (UNET),
+                                           'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT
+                                           Chebyshev+Taylor forecasting)
   --cache-option                           named cache params (key=value format, comma-separated). easycache/ucache:
-                                           threshold=,start=,end=,decay=,relative=,reset=; dbcache/taylorseer/cache-dit: Fn=,Bn=,threshold=,warmup=. Examples:
+                                           threshold=,start=,end=,decay=,relative=,reset=; dbcache/taylorseer/cache-dit:
-                                           "threshold=0.25" or "threshold=1.5,reset=0"
+                                           Fn=,Bn=,threshold=,warmup=; spectrum: w=,m=,lam=,window=,flex=,warmup=,stop=.
-  --scm-mask                               SCM steps mask for cache-dit: comma-separated 0/1 (e.g., "1,1,1,0,0,1,0,0,1,0") - 1=compute, 0=can cache
+                                           Examples: "threshold=0.25" or "threshold=1.5,reset=0"
+  --scm-mask                               SCM steps mask for cache-dit: comma-separated 0/1 (e.g.,
+                                           "1,1,1,0,0,1,0,0,1,0") - 1=compute, 0=can cache
   --scm-policy                             SCM policy: 'dynamic' (default) or 'static'
+  --vae-tile-size                          tile size for vae tiling, format [X]x[Y] (default: 32x32)
+  --vae-relative-tile-size                 relative tile size for vae tiling, format [X]x[Y], in fraction of image size
+                                           if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size)
 ```

examples/server/api.md

@@ -9,7 +9,7 @@ The server currently exposes three API families:
 - `sdcpp API` under `/sdcpp/v1/...`
 
 The `sdcpp API` is the native API surface.
-Its request schema is also the canonical schema for `sd_cpp_extra_args`.
+Its request schema is the same schema used by `sd_cpp_extra_args`.
 
 Global LoRA rule:
 
@@ -38,6 +38,8 @@ Current generation-related endpoints include:
 - `POST /sdapi/v1/txt2img`
 - `POST /sdapi/v1/img2img`
 - `GET /sdapi/v1/loras`
+- `GET /sdapi/v1/upscalers`
+- `GET /sdapi/v1/latent-upscale-modes`
 - `GET /sdapi/v1/samplers`
 - `GET /sdapi/v1/schedulers`
 - `GET /sdapi/v1/sd-models`
@@ -55,8 +57,6 @@ Current endpoints include:
 - `POST /sdcpp/v1/jobs/{id}/cancel`
 - `POST /sdcpp/v1/vid_gen`
 
-`POST /sdcpp/v1/vid_gen` is currently exposed but returns `501 Not Implemented`.
-
 ## `sd_cpp_extra_args`
 
 `sd_cpp_extra_args` is an extension mechanism for the compatibility APIs.
@@ -79,12 +79,12 @@ Behavior:
 - The JSON block is parsed using the same field rules as the `sdcpp API`.
 - The block is removed from the final prompt before generation.
 
-Intended use:
+Supported use:
 
 - extend `OpenAI API` requests with native `stable-diffusion.cpp` controls
 - extend `sdapi` requests with native `stable-diffusion.cpp` controls
 
-Not intended use:
+Unsupported use:
 
 - do not use `sd_cpp_extra_args` with `/sdcpp/v1/*`
 
@@ -218,6 +218,13 @@ Currently supported request fields:
 | `scheduler` | `string` | Scheduler name |
 | `lora` | `array<object>` | Structured LoRA list |
 | `extra_images` | `array<string>` | Base64 or data URL images |
+| `enable_hr` | `boolean` | Enable highres fix for `txt2img` |
+| `hr_upscaler` | `string` | `Lanczos`, `Nearest`, a latent mode such as `Latent (nearest-exact)`, or an upscaler model name from `/sdapi/v1/upscalers` |
+| `hr_scale` | `number` | Highres scale when resize target is not set |
+| `hr_resize_x` | `integer` | Highres target width, `0` to use scale |
+| `hr_resize_y` | `integer` | Highres target height, `0` to use scale |
+| `hr_steps` | `integer` | Highres second-pass sample steps, `0` to reuse `steps` |
+| `denoising_strength` | `number` | Highres denoising strength for `txt2img` |
 
 Native extension fields:
 
@@ -243,6 +250,8 @@ Currently supported request fields:
 | `inpainting_mask_invert` | `integer` or `boolean` | Treated as invert flag |
 | `denoising_strength` | `number` | Clamped to `0.0..1.0` |
 
+Highres fix fields are currently handled for `txt2img`; `img2img` uses `denoising_strength` as image-to-image strength.
+
 Native extension fields:
 
 - any `sdcpp API` fields embedded through `sd_cpp_extra_args` inside `prompt`
@@ -260,6 +269,8 @@ Response fields:
 Currently exposed:
 
 - `GET /sdapi/v1/loras`
+- `GET /sdapi/v1/upscalers`
+- `GET /sdapi/v1/latent-upscale-modes`
 - `GET /sdapi/v1/samplers`
 - `GET /sdapi/v1/schedulers`
 - `GET /sdapi/v1/sd-models`
@@ -274,6 +285,26 @@ Response fields:
 | `[].name` | `string` | Display name derived from file stem |
 | `[].path` | `string` | Relative path under the configured LoRA directory |
 
+`GET /sdapi/v1/upscalers`
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `[].name` | `string` | Built-in name or model stem |
+| `[].model_name` | `string \| null` | Model family label for model-backed upscalers |
+| `[].model_path` | `string \| null` | Absolute model path for model-backed upscalers |
+| `[].model_url` | `string \| null` | Currently always null |
+| `[].scale` | `integer` | Currently `4` |
+
+Built-in entries include `None`, `Lanczos`, and `Nearest`. Model-backed entries are scanned from the top level of `--hires-upscalers-dir`; subdirectories are not scanned.
+
+`GET /sdapi/v1/latent-upscale-modes`
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `[].name` | `string` | WebUI-compatible latent upscale mode name |
+
+Built-in latent modes include `Latent`, `Latent (nearest)`, `Latent (nearest-exact)`, `Latent (antialiased)`, `Latent (bicubic)`, and `Latent (bicubic antialiased)`.
+
 `GET /sdapi/v1/samplers`
 
 | Field | Type | Notes |
@@ -372,20 +403,26 @@ Field types:
 
 Returns frontend-friendly capability metadata.
 
-Typical contents:
+The mode-aware fields are the primary interface. The top-level compatibility fields are deprecated mirrors kept for older clients.
 
-| Field | Type |
+Top-level fields:
-| --- | --- |
-| `model` | `object` |
-| `defaults` | `object` |
-| `loras` | `array<object>` |
-| `samplers` | `array<string>` |
-| `schedulers` | `array<string>` |
-| `output_formats` | `array<string>` |
-| `limits` | `object` |
-| `features` | `object` |
 
-Nested fields currently returned:
+| Field | Type | Notes |
+| --- | --- | --- |
+| `model` | `object` | Loaded model metadata |
+| `current_mode` | `string` | The native generation mode mirrored by top-level compatibility fields |
+| `supported_modes` | `array<string>` | Supported native modes such as `img_gen` or `vid_gen` |
+| `defaults` | `object` | Deprecated compatibility mirror of `defaults_by_mode[current_mode]` |
+| `output_formats` | `array<string>` | Deprecated compatibility mirror of `output_formats_by_mode[current_mode]` |
+| `features` | `object` | Deprecated compatibility mirror of `features_by_mode[current_mode]` |
+| `defaults_by_mode` | `object` | Explicit defaults for each supported mode |
+| `output_formats_by_mode` | `object` | Explicit output formats for each supported mode |
+| `features_by_mode` | `object` | Explicit feature flags for each supported mode |
+| `samplers` | `array<string>` | Available sampling methods |
+| `schedulers` | `array<string>` | Available schedulers |
+| `loras` | `array<object>` | Available LoRA entries |
+| `upscalers` | `array<object>` | Available model-backed highres upscalers |
+| `limits` | `object` | Shared queue and size limits |
 
 `model`
 
@@ -395,50 +432,24 @@ Nested fields currently returned:
 | `model.stem` | `string` |
 | `model.path` | `string` |
 
-`defaults`
+Compatibility rules:
+
+- `defaults`, `output_formats`, and `features` are deprecated compatibility mirrors
+- those three top-level fields always mirror `current_mode`
+- `supported_modes`, `defaults_by_mode`, `output_formats_by_mode`, and `features_by_mode` are the mode-aware fields
+
+Mode-aware objects:
 
 | Field | Type |
 | --- | --- |
-| `defaults.prompt` | `string` |
+| `defaults_by_mode.img_gen` | `object` |
-| `defaults.negative_prompt` | `string` |
+| `defaults_by_mode.vid_gen` | `object` |
-| `defaults.clip_skip` | `integer` |
+| `output_formats_by_mode.img_gen` | `array<string>` |
-| `defaults.width` | `integer` |
+| `output_formats_by_mode.vid_gen` | `array<string>` |
-| `defaults.height` | `integer` |
+| `features_by_mode.img_gen` | `object` |
-| `defaults.strength` | `number` |
+| `features_by_mode.vid_gen` | `object` |
-| `defaults.seed` | `integer` |
+
-| `defaults.batch_count` | `integer` |
+Shared nested fields:
-| `defaults.auto_resize_ref_image` | `boolean` |
-| `defaults.increase_ref_index` | `boolean` |
-| `defaults.control_strength` | `number` |
-| `defaults.sample_params` | `object` |
-| `defaults.sample_params.scheduler` | `string` |
-| `defaults.sample_params.sample_method` | `string` |
-| `defaults.sample_params.sample_steps` | `integer` |
-| `defaults.sample_params.eta` | `number \| null` |
-| `defaults.sample_params.shifted_timestep` | `integer` |
-| `defaults.sample_params.flow_shift` | `number \| null` |
-| `defaults.sample_params.guidance` | `object` |
-| `defaults.sample_params.guidance.txt_cfg` | `number` |
-| `defaults.sample_params.guidance.img_cfg` | `number \| null` |
-| `defaults.sample_params.guidance.distilled_guidance` | `number` |
-| `defaults.sample_params.guidance.slg` | `object` |
-| `defaults.sample_params.guidance.slg.layers` | `array<integer>` |
-| `defaults.sample_params.guidance.slg.layer_start` | `number` |
-| `defaults.sample_params.guidance.slg.layer_end` | `number` |
-| `defaults.sample_params.guidance.slg.scale` | `number` |
-| `defaults.vae_tiling_params` | `object` |
-| `defaults.vae_tiling_params.enabled` | `boolean` |
-| `defaults.vae_tiling_params.tile_size_x` | `integer` |
-| `defaults.vae_tiling_params.tile_size_y` | `integer` |
-| `defaults.vae_tiling_params.target_overlap` | `number` |
-| `defaults.vae_tiling_params.rel_size_x` | `number` |
-| `defaults.vae_tiling_params.rel_size_y` | `number` |
-| `defaults.cache_mode` | `string` |
-| `defaults.cache_option` | `string` |
-| `defaults.scm_mask` | `string` |
-| `defaults.scm_policy_dynamic` | `boolean` |
-| `defaults.output_format` | `string` |
-| `defaults.output_compression` | `integer` |
 
 `loras`
 
@@ -447,6 +458,14 @@ Nested fields currently returned:
 | `loras[].name` | `string` |
 | `loras[].path` | `string` |
 
+`upscalers`
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `upscalers[].name` | `string` | Built-in name or model stem; use this value in `hires.upscaler` |
+
+Built-in entries include `None`, `Lanczos`, `Nearest`, `Latent`, `Latent (nearest)`, `Latent (nearest-exact)`, `Latent (antialiased)`, `Latent (bicubic)`, and `Latent (bicubic antialiased)`. Model-backed entries are scanned from the top level of `--hires-upscalers-dir`; subdirectories are not scanned.
+
 `limits`
 
 | Field | Type |
@@ -458,19 +477,110 @@ Nested fields currently returned:
 | `limits.max_batch_count` | `integer` |
 | `limits.max_queue_size` | `integer` |
 
-`features`
+Shared default fields used by both `img_gen` and `vid_gen`:
 
 | Field | Type |
 | --- | --- |
-| `features.init_image` | `boolean` |
+| `prompt` | `string` |
-| `features.mask_image` | `boolean` |
+| `negative_prompt` | `string` |
-| `features.control_image` | `boolean` |
+| `clip_skip` | `integer` |
-| `features.ref_images` | `boolean` |
+| `width` | `integer` |
-| `features.lora` | `boolean` |
+| `height` | `integer` |
-| `features.vae_tiling` | `boolean` |
+| `strength` | `number` |
-| `features.cache` | `boolean` |
+| `seed` | `integer` |
-| `features.cancel_queued` | `boolean` |
+| `sample_params` | `object` |
-| `features.cancel_generating` | `boolean` |
+| `sample_params.scheduler` | `string` |
+| `sample_params.sample_method` | `string` |
+| `sample_params.sample_steps` | `integer` |
+| `sample_params.eta` | `number \| null` |
+| `sample_params.shifted_timestep` | `integer` |
+| `sample_params.flow_shift` | `number \| null` |
+| `sample_params.guidance.txt_cfg` | `number` |
+| `sample_params.guidance.img_cfg` | `number \| null` |
+| `sample_params.guidance.distilled_guidance` | `number` |
+| `sample_params.guidance.slg.layers` | `array<integer>` |
+| `sample_params.guidance.slg.layer_start` | `number` |
+| `sample_params.guidance.slg.layer_end` | `number` |
+| `sample_params.guidance.slg.scale` | `number` |
+| `vae_tiling_params` | `object` |
+| `vae_tiling_params.enabled` | `boolean` |
+| `vae_tiling_params.tile_size_x` | `integer` |
+| `vae_tiling_params.tile_size_y` | `integer` |
+| `vae_tiling_params.target_overlap` | `number` |
+| `vae_tiling_params.rel_size_x` | `number` |
+| `vae_tiling_params.rel_size_y` | `number` |
+| `cache_mode` | `string` |
+| `cache_option` | `string` |
+| `scm_mask` | `string` |
+| `scm_policy_dynamic` | `boolean` |
+| `output_format` | `string` |
+| `output_compression` | `integer` |
+
+`img_gen`-specific default fields:
+
+| Field | Type |
+| --- | --- |
+| `batch_count` | `integer` |
+| `auto_resize_ref_image` | `boolean` |
+| `increase_ref_index` | `boolean` |
+| `control_strength` | `number` |
+| `hires` | `object` |
+| `hires.enabled` | `boolean` |
+| `hires.upscaler` | `string` |
+| `hires.scale` | `number` |
+| `hires.target_width` | `integer` |
+| `hires.target_height` | `integer` |
+| `hires.steps` | `integer` |
+| `hires.denoising_strength` | `number` |
+| `hires.upscale_tile_size` | `integer` |
+
+`vid_gen`-specific default fields:
+
+| Field | Type |
+| --- | --- |
+| `video_frames` | `integer` |
+| `fps` | `integer` |
+| `moe_boundary` | `number` |
+| `vace_strength` | `number` |
+| `high_noise_sample_params` | `object` |
+| `high_noise_sample_params.scheduler` | `string` |
+| `high_noise_sample_params.sample_method` | `string` |
+| `high_noise_sample_params.sample_steps` | `integer` |
+| `high_noise_sample_params.eta` | `number \| null` |
+| `high_noise_sample_params.shifted_timestep` | `integer` |
+| `high_noise_sample_params.flow_shift` | `number \| null` |
+| `high_noise_sample_params.guidance.txt_cfg` | `number` |
+| `high_noise_sample_params.guidance.img_cfg` | `number \| null` |
+| `high_noise_sample_params.guidance.distilled_guidance` | `number` |
+| `high_noise_sample_params.guidance.slg.layers` | `array<integer>` |
+| `high_noise_sample_params.guidance.slg.layer_start` | `number` |
+| `high_noise_sample_params.guidance.slg.layer_end` | `number` |
+| `high_noise_sample_params.guidance.slg.scale` | `number` |
+
+Fields returned in `features_by_mode.img_gen`:
+
+- `init_image`
+- `mask_image`
+- `control_image`
+- `ref_images`
+- `lora`
+- `vae_tiling`
+- `hires`
+- `cache`
+- `cancel_queued`
+- `cancel_generating`
+
+Fields returned in `features_by_mode.vid_gen`:
+
+- `init_image`
+- `end_image`
+- `control_frames`
+- `high_noise_sample_params`
+- `lora`
+- `vae_tiling`
+- `cache`
+- `cancel_queued`
+- `cancel_generating`
 
 #### `POST /sdcpp/v1/img_gen`
 
@@ -521,9 +631,7 @@ Typical status codes:
 - `409 Conflict`
 - `410 Gone`
 
-### Canonical Request Schema
+### Request Body
-
-The `sdcpp API` request body is the canonical native schema.
 
 Example:
 
@@ -569,6 +677,16 @@ Example:
   },
 
   "lora": [],
+  "hires": {
+    "enabled": false,
+    "upscaler": "Latent",
+    "scale": 2.0,
+    "target_width": 0,
+    "target_height": 0,
+    "steps": 0,
+    "denoising_strength": 0.7,
+    "upscale_tile_size": 128
+  },
 
   "vae_tiling_params": {
     "enabled": false,
@@ -612,7 +730,7 @@ Channel expectations:
 If omitted or null:
 
 - single-image fields map to an empty `sd_image_t`
-- array fields map to `nullptr + count = 0`
+- array fields map to an empty C-style array, represented as `pointer = nullptr` and `count = 0`
 
 ### Field Mapping Summary
 
@@ -673,12 +791,23 @@ Other native fields:
 
 | Field | Type |
 | --- | --- |
+| `hires` | `object` |
+| `hires.enabled` | `boolean` |
+| `hires.upscaler` | `string` |
+| `hires.scale` | `number` |
+| `hires.target_width` | `integer` |
+| `hires.target_height` | `integer` |
+| `hires.steps` | `integer` |
+| `hires.denoising_strength` | `number` |
+| `hires.upscale_tile_size` | `integer` |
 | `vae_tiling_params` | `object` |
 | `cache_mode` | `string` |
 | `cache_option` | `string` |
 | `scm_mask` | `string` |
 | `scm_policy_dynamic` | `boolean` |
 
+For `hires.upscaler`, use `Lanczos`, `Nearest`, `Latent`, `Latent (nearest)`, `Latent (nearest-exact)`, `Latent (antialiased)`, `Latent (bicubic)`, `Latent (bicubic antialiased)`, or an `upscalers[].name` value from `GET /sdcpp/v1/capabilities`. Model-backed upscalers are resolved as `--hires-upscalers-dir / (name + ext)` and must live directly in that directory.
+
 HTTP-only output fields:
 
 | Field | Type |
@@ -686,11 +815,11 @@ HTTP-only output fields:
 | `output_format` | `string` |
 | `output_compression` | `integer` |
 
-### Optional Field Semantics
+### Optional Field Handling
 
-Clients should preserve unset semantics for optional sampling fields.
+Optional sampling fields may be omitted.
 
-If a user has not explicitly provided one of these fields, the client should omit it instead of injecting a guessed fallback:
+When omitted, backend defaults apply to these fields:
 
 - `sample_params.scheduler`
 - `sample_params.sample_method`
@@ -766,29 +895,394 @@ Example cancelled job:
 }
 ```
 
-### Validation and Retention
+### Submission Errors
 
-Recommended behavior:
+`POST /sdcpp/v1/img_gen` may return:
 
-- malformed JSON returns `400`
+- `202 Accepted` when the job is created
-- invalid image payloads return `400`
+- `400 Bad Request` for an empty body, unsupported model mode, invalid JSON, or invalid generation parameters
-- invalid parameter structure returns `400`
+- `429 Too Many Requests` when the job queue is full
-- queue full returns `429` or `503`
+- `500 Internal Server Error` for unexpected server exceptions during submission
-- accepted runtime failures transition the job to `failed`
-- unsupported in-progress cancellation may return `409`
 
-Recommended retention controls:
+### `vid_gen`
 
-- pending job limit
+The following section documents the native async contract for video generation.
-- completed job TTL
-- failed job TTL
 
-### Future `vid_gen`
+#### `POST /sdcpp/v1/vid_gen`
 
-Future `vid_gen` should reuse the same async job model:
+Submits an async video generation job.
 
-- `POST /sdcpp/v1/vid_gen`
+Successful submission returns `202 Accepted`.
-- `GET /sdcpp/v1/jobs/{id}`
-- `POST /sdcpp/v1/jobs/{id}/cancel`
 
-Its request body should mirror `sd_vid_gen_params_t` in the same way that `img_gen` mirrors `sd_img_gen_params_t`.
+Example response:
+
+```json
+{
+  "id": "job_01HTXYZVID",
+  "kind": "vid_gen",
+  "status": "queued",
+  "created": 1775401200,
+  "poll_url": "/sdcpp/v1/jobs/job_01HTXYZVID"
+}
+```
+
+Response fields:
+
+| Field | Type |
+| --- | --- |
+| `id` | `string` |
+| `kind` | `string` |
+| `status` | `string` |
+| `created` | `integer` |
+| `poll_url` | `string` |
+
+### Request Body
+
+Compared with `img_gen`, the `vid_gen` request body:
+
+- `vid_gen` is a single video sequence job, so `batch_count` is not part of the request schema
+- `ref_images`, `mask_image`, `control_image`, `control_strength`, and `embed_image_metadata` are not part of the request schema
+- `vid_gen` adds `end_image`, `control_frames`, `high_noise_sample_params`, `video_frames`, `fps`, `moe_boundary`, and `vace_strength`
+
+Example:
+
+```json
+{
+  "prompt": "a cat walking through a rainy alley",
+  "negative_prompt": "",
+  "clip_skip": -1,
+  "width": 832,
+  "height": 480,
+  "strength": 0.75,
+  "seed": -1,
+  "video_frames": 33,
+  "fps": 16,
+  "moe_boundary": 0.875,
+  "vace_strength": 1.0,
+
+  "init_image": null,
+  "end_image": null,
+  "control_frames": [],
+
+  "sample_params": {
+    "scheduler": "discrete",
+    "sample_method": "euler",
+    "sample_steps": 28,
+    "eta": 1.0,
+    "shifted_timestep": 0,
+    "custom_sigmas": [],
+    "flow_shift": 0.0,
+    "guidance": {
+      "txt_cfg": 7.0,
+      "img_cfg": 7.0,
+      "distilled_guidance": 3.5,
+      "slg": {
+        "layers": [7, 8, 9],
+        "layer_start": 0.01,
+        "layer_end": 0.2,
+        "scale": 0.0
+      }
+    }
+  },
+
+  "high_noise_sample_params": {
+    "scheduler": "discrete",
+    "sample_method": "euler",
+    "sample_steps": -1,
+    "eta": 1.0,
+    "shifted_timestep": 0,
+    "flow_shift": 0.0,
+    "guidance": {
+      "txt_cfg": 7.0,
+      "img_cfg": 7.0,
+      "distilled_guidance": 3.5,
+      "slg": {
+        "layers": [7, 8, 9],
+        "layer_start": 0.01,
+        "layer_end": 0.2,
+        "scale": 0.0
+      }
+    }
+  },
+
+  "lora": [],
+
+  "vae_tiling_params": {
+    "enabled": false,
+    "tile_size_x": 0,
+    "tile_size_y": 0,
+    "target_overlap": 0.5,
+    "rel_size_x": 0.0,
+    "rel_size_y": 0.0
+  },
+
+  "cache_mode": "disabled",
+  "cache_option": "",
+  "scm_mask": "",
+  "scm_policy_dynamic": true,
+
+  "output_format": "webm",
+  "output_compression": 100
+}
+```
+
+### LoRA Rules
+
+- The server only accepts explicit LoRA entries from the `lora` field.
+- Prompt-embedded `<lora:...>` tags are intentionally unsupported.
+- `lora[].is_high_noise` controls whether a LoRA applies only to the high-noise stage.
+
+### Image and Frame Encoding Rules
+
+Any image field accepts:
+
+- a raw base64 string, or
+- a data URL such as `data:image/png;base64,...`
+
+Channel expectations:
+
+- `init_image`: 3 channels
+- `end_image`: 3 channels
+- `control_frames[]`: 3 channels
+
+Frame ordering rules:
+
+- `control_frames[]` order is the conditioning frame order
+- `control_frames[]` is preserved in request order
+
+If omitted or null:
+
+- single-image fields map to an empty `sd_image_t`
+- array fields map to an empty C-style array, represented as `pointer = nullptr` and `count = 0`
+
+### Field Mapping Summary
+
+Top-level scalar fields:
+
+| Field | Type |
+| --- | --- |
+| `prompt` | `string` |
+| `negative_prompt` | `string` |
+| `clip_skip` | `integer` |
+| `width` | `integer` |
+| `height` | `integer` |
+| `strength` | `number` |
+| `seed` | `integer` |
+| `video_frames` | `integer` |
+| `fps` | `integer` |
+| `moe_boundary` | `number` |
+| `vace_strength` | `number` |
+
+Image and frame fields:
+
+| Field | Type |
+| --- | --- |
+| `init_image` | `string \| null` |
+| `end_image` | `string \| null` |
+| `control_frames` | `array<string>` |
+
+LoRA fields:
+
+| Field | Type |
+| --- | --- |
+| `lora[].path` | `string` |
+| `lora[].multiplier` | `number` |
+| `lora[].is_high_noise` | `boolean` |
+
+Sampling fields:
+
+| Field | Type |
+| --- | --- |
+| `sample_params.scheduler` | `string` |
+| `sample_params.sample_method` | `string` |
+| `sample_params.sample_steps` | `integer` |
+| `sample_params.eta` | `number` |
+| `sample_params.shifted_timestep` | `integer` |
+| `sample_params.custom_sigmas` | `array<number>` |
+| `sample_params.flow_shift` | `number` |
+| `sample_params.guidance.txt_cfg` | `number` |
+| `sample_params.guidance.img_cfg` | `number` |
+| `sample_params.guidance.distilled_guidance` | `number` |
+| `sample_params.guidance.slg.layers` | `array<integer>` |
+| `sample_params.guidance.slg.layer_start` | `number` |
+| `sample_params.guidance.slg.layer_end` | `number` |
+| `sample_params.guidance.slg.scale` | `number` |
+
+High-noise sampling fields:
+
+| Field | Type |
+| --- | --- |
+| `high_noise_sample_params.scheduler` | `string` |
+| `high_noise_sample_params.sample_method` | `string` |
+| `high_noise_sample_params.sample_steps` | `integer` |
+| `high_noise_sample_params.eta` | `number` |
+| `high_noise_sample_params.shifted_timestep` | `integer` |
+| `high_noise_sample_params.flow_shift` | `number` |
+| `high_noise_sample_params.guidance.txt_cfg` | `number` |
+| `high_noise_sample_params.guidance.img_cfg` | `number` |
+| `high_noise_sample_params.guidance.distilled_guidance` | `number` |
+| `high_noise_sample_params.guidance.slg.layers` | `array<integer>` |
+| `high_noise_sample_params.guidance.slg.layer_start` | `number` |
+| `high_noise_sample_params.guidance.slg.layer_end` | `number` |
+| `high_noise_sample_params.guidance.slg.scale` | `number` |
+
+Other native fields:
+
+| Field | Type |
+| --- | --- |
+| `vae_tiling_params` | `object` |
+| `cache_mode` | `string` |
+| `cache_option` | `string` |
+| `scm_mask` | `string` |
+| `scm_policy_dynamic` | `boolean` |
+
+HTTP-only output fields:
+
+| Field | Type |
+| --- | --- |
+| `output_format` | `string` |
+| `output_compression` | `integer` |
+
+For `vid_gen`, `output_format` and `output_compression` control container encoding.
+`fps` is request metadata for the generated sequence and is echoed in the completed job result.
+
+Allowed `output_format` values:
+
+- `webm`
+- `webp`
+- `avi`
+
+Output format behavior:
+
+- `output_format` defaults to `webm`
+- `webp` means animated WebP
+- `avi` means MJPG AVI
+- `webm` requires the server to be built with WebM support; otherwise the request returns `400`
+
+### Result Payload
+
+Completed jobs return one encoded container payload, not a list of per-frame images.
+
+Result fields:
+
+- `result.b64_json` contains the whole encoded container file as base64
+- `result.mime_type` identifies the media type
+- `result.output_format` echoes the selected container format
+- `result.fps` echoes the effective playback FPS
+- `result.frame_count` reports the actual decoded frame count used to build the container
+
+Expected MIME types:
+
+| `output_format` | `mime_type` |
+| --- | --- |
+| `webm` | `video/webm` |
+| `webp` | `image/webp` |
+| `avi` | `video/x-msvideo` |
+
+### Optional Field Handling
+
+Optional sampling fields may be omitted.
+
+When omitted, backend defaults apply to these fields:
+
+- `sample_params.scheduler`
+- `sample_params.sample_method`
+- `sample_params.eta`
+- `sample_params.flow_shift`
+- `sample_params.guidance.img_cfg`
+- `high_noise_sample_params.scheduler`
+- `high_noise_sample_params.sample_method`
+- `high_noise_sample_params.eta`
+- `high_noise_sample_params.flow_shift`
+- `high_noise_sample_params.guidance.img_cfg`
+
+`high_noise_sample_params` may also be omitted entirely.
+
+### Frame Count Semantics
+
+`video_frames` is the requested target length, but the current core video path internally normalizes the effective frame count to the largest `4n + 1` value that does not exceed the requested count.
+
+Examples:
+
+- `video_frames = 33` stays `33`
+- `video_frames = 34` becomes `33`
+- `video_frames = 32` becomes `29`
+
+The completed job payload includes the actual decoded `frame_count`.
+
+### Completion Result
+
+Example completed job:
+
+```json
+{
+  "id": "job_01HTXYZVID",
+  "kind": "vid_gen",
+  "status": "completed",
+  "created": 1775401200,
+  "started": 1775401203,
+  "completed": 1775401215,
+  "queue_position": 0,
+  "result": {
+    "output_format": "webm",
+    "mime_type": "video/webm",
+    "fps": 16,
+    "frame_count": 33,
+    "b64_json": "GkXfo59ChoEBQveBAULygQRC84EIQo..."
+  },
+  "error": null
+}
+```
+
+The response returns the encoded `.webm`, animated `.webp`, or `.avi` container payload directly.
+
+### Failure Result
+
+Example failed job:
+
+```json
+{
+  "id": "job_01HTXYZVID",
+  "kind": "vid_gen",
+  "status": "failed",
+  "created": 1775401200,
+  "started": 1775401203,
+  "completed": 1775401204,
+  "queue_position": 0,
+  "result": null,
+  "error": {
+    "code": "generation_failed",
+    "message": "generate_video returned no results"
+  }
+}
+```
+
+### Cancelled Result
+
+Example cancelled job:
+
+```json
+{
+  "id": "job_01HTXYZVID",
+  "kind": "vid_gen",
+  "status": "cancelled",
+  "created": 1775401200,
+  "started": null,
+  "completed": 1775401202,
+  "queue_position": 0,
+  "result": null,
+  "error": {
+    "code": "cancelled",
+    "message": "job cancelled by client"
+  }
+}
+```
+
+### Submission Errors
+
+`POST /sdcpp/v1/vid_gen` may return:
+
+- `202 Accepted` when the job is created
+- `400 Bad Request` for an empty body, unsupported model mode, invalid JSON, invalid generation parameters, or an unsupported output format
+- `429 Too Many Requests` when the job queue is full
+- `500 Internal Server Error` for unexpected server exceptions during submission

examples/server/async_jobs.cpp

@@ -95,6 +95,10 @@ bool cancel_queued_job(AsyncJobManager& manager, AsyncGenerationJob& job) {
     job.status       = AsyncJobStatus::Cancelled;
     job.completed_at = unix_timestamp_now();
     job.result_images_b64.clear();
+    job.result_media_b64.clear();
+    job.result_media_mime_type.clear();
+    job.result_frame_count = 0;
+    job.result_fps         = 0;
     job.error_code         = "cancelled";
     job.error_message      = "job cancelled by client";
     return true;
@@ -122,6 +126,15 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo
     }
 
     if (job.status == AsyncJobStatus::Completed) {
+        if (job.kind == AsyncJobKind::VidGen) {
+            result["result"] = {
+                {"output_format", job.vid_gen.output_format},
+                {"mime_type", job.result_media_mime_type},
+                {"fps", job.result_fps},
+                {"frame_count", job.result_frame_count},
+                {"b64_json", job.result_media_b64},
+            };
+        } else {
             json images = json::array();
             for (size_t i = 0; i < job.result_images_b64.size(); ++i) {
                 images.push_back({{"index", i}, {"b64_json", job.result_images_b64[i]}});
@@ -130,6 +143,7 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo
                 {"output_format", job.img_gen.output_format},
                 {"images", images},
             };
+        }
         result["error"] = nullptr;
     } else if (job.status == AsyncJobStatus::Failed ||
                job.status == AsyncJobStatus::Cancelled) {
@@ -156,16 +170,15 @@ bool execute_img_gen_job(ServerRuntime& runtime,
     sd_img_gen_params_t params = job.img_gen.to_sd_img_gen_params_t();
 
     SDImageVec results;
-    int num_results = 0;
 
     {
         std::lock_guard<std::mutex> lock(*runtime.sd_ctx_mutex);
         sd_image_t* raw_results = generate_image(runtime.sd_ctx, &params);
-        num_results             = params.batch_count;
+        results.adopt(raw_results, params.batch_count);
-        results.adopt(raw_results, num_results);
     }
 
-    if (results.empty() || num_results <= 0) {
+    const int num_results = results.count();
+    if (num_results <= 0) {
         error_message = "generate_image returned no results";
         return false;
     }
@@ -208,6 +221,47 @@ bool execute_img_gen_job(ServerRuntime& runtime,
     return true;
 }
 
+bool execute_vid_gen_job(ServerRuntime& runtime,
+                         AsyncGenerationJob& job,
+                         std::string& output_media_b64,
+                         std::string& output_media_mime_type,
+                         int& output_frame_count,
+                         int& output_fps,
+                         std::string& error_message) {
+    sd_vid_gen_params_t params = job.vid_gen.to_sd_vid_gen_params_t();
+
+    SDImageVec results;
+    int num_results = 0;
+
+    {
+        std::lock_guard<std::mutex> lock(*runtime.sd_ctx_mutex);
+        sd_image_t* raw_results = generate_video(runtime.sd_ctx, &params, &num_results);
+        results.adopt(raw_results, num_results);
+    }
+
+    num_results = results.count();
+    if (num_results <= 0) {
+        error_message = "generate_video returned no results";
+        return false;
+    }
+
+    std::vector<uint8_t> video_bytes = create_video_from_sd_images_to_vector(job.vid_gen.output_format,
+                                                                             results.data(),
+                                                                             num_results,
+                                                                             job.vid_gen.gen_params.fps,
+                                                                             job.vid_gen.output_compression);
+    if (video_bytes.empty()) {
+        error_message = "failed to encode generated video container";
+        return false;
+    }
+
+    output_media_b64       = base64_encode(video_bytes);
+    output_media_mime_type = video_mime_type(job.vid_gen.output_format);
+    output_frame_count     = num_results;
+    output_fps             = job.vid_gen.gen_params.fps;
+    return true;
+}
+
 void async_job_worker(ServerRuntime& runtime) {
     AsyncJobManager& manager = *runtime.async_job_manager;
 
@@ -240,11 +294,23 @@ void async_job_worker(ServerRuntime& runtime) {
         }
 
         std::vector<std::string> output_images;
+        std::string output_media_b64;
+        std::string output_media_mime_type;
+        int output_frame_count = 0;
+        int output_fps         = 0;
         std::string error_message;
         bool ok = false;
 
         if (job->kind == AsyncJobKind::ImgGen) {
             ok = execute_img_gen_job(runtime, *job, output_images, error_message);
+        } else if (job->kind == AsyncJobKind::VidGen) {
+            ok = execute_vid_gen_job(runtime,
+                                     *job,
+                                     output_media_b64,
+                                     output_media_mime_type,
+                                     output_frame_count,
+                                     output_fps,
+                                     error_message);
         } else {
             error_message = "unsupported job kind";
         }
@@ -260,6 +326,10 @@ void async_job_worker(ServerRuntime& runtime) {
             if (ok) {
                 job->status                 = AsyncJobStatus::Completed;
                 job->result_images_b64      = std::move(output_images);
+                job->result_media_b64       = std::move(output_media_b64);
+                job->result_media_mime_type = std::move(output_media_mime_type);
+                job->result_frame_count     = output_frame_count;
+                job->result_fps             = output_fps;
                 job->error_code.clear();
                 job->error_message.clear();
             } else {
@@ -267,6 +337,10 @@ void async_job_worker(ServerRuntime& runtime) {
                 job->error_code    = "generation_failed";
                 job->error_message = error_message.empty() ? "unknown generation error" : error_message;
                 job->result_images_b64.clear();
+                job->result_media_b64.clear();
+                job->result_media_mime_type.clear();
+                job->result_frame_count = 0;
+                job->result_fps         = 0;
             }
 
             purge_expired_jobs(manager);

examples/server/async_jobs.h

@@ -36,7 +36,12 @@ struct AsyncGenerationJob {
     int64_t started_at    = 0;
     int64_t completed_at  = 0;
     ImgGenJobRequest img_gen;
+    VidGenJobRequest vid_gen;
     std::vector<std::string> result_images_b64;
+    std::string result_media_b64;
+    std::string result_media_mime_type;
+    int result_frame_count = 0;
+    int result_fps         = 0;
     std::string error_code;
     std::string error_message;
 };
@@ -63,4 +68,11 @@ bool execute_img_gen_job(ServerRuntime& runtime,
                          AsyncGenerationJob& job,
                          std::vector<std::string>& output_images,
                          std::string& error_message);
+bool execute_vid_gen_job(ServerRuntime& runtime,
+                         AsyncGenerationJob& job,
+                         std::string& output_media_b64,
+                         std::string& output_media_mime_type,
+                         int& output_frame_count,
+                         int& output_fps,
+                         std::string& error_message);
 void async_job_worker(ServerRuntime& runtime);

examples/server/frontend

@@ -1 +1 @@
-Subproject commit 740475a7a6794dc07fb23e8ec5dc56e7e80aa8c1
+Subproject commit 797ccf80825cc035508ba9b599b2a21953e7f835

examples/server/main.cpp

@@ -48,7 +48,9 @@ static void parse_args(int argc,
 
     if (!svr_params.resolve_and_validate() ||
         !ctx_params.resolve_and_validate(IMG_GEN) ||
-        !default_gen_params.resolve_and_validate(IMG_GEN, ctx_params.lora_model_dir)) {
+        !default_gen_params.resolve_and_validate(IMG_GEN,
+                                                 ctx_params.lora_model_dir,
+                                                 ctx_params.hires_upscalers_dir)) {
         print_usage(argv[0], options_vec);
         exit(1);
     }
@@ -95,6 +97,8 @@ int main(int argc, const char** argv) {
 
     std::vector<LoraEntry> lora_cache;
     std::mutex lora_mutex;
+    std::vector<UpscalerEntry> upscaler_cache;
+    std::mutex upscaler_mutex;
     AsyncJobManager async_job_manager;
     ServerRuntime runtime = {
         sd_ctx.get(),
@@ -104,6 +108,8 @@ int main(int argc, const char** argv) {
         &default_gen_params,
         &lora_cache,
         &lora_mutex,
+        &upscaler_cache,
+        &upscaler_mutex,
         &async_job_manager,
     };
 

examples/server/routes_openai.cpp

@@ -70,7 +70,7 @@ static bool build_openai_generation_request(const httplib::Request& req,
     }
 
     // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
-    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", true)) {
+    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) {
         error_message = "invalid params";
         return false;
     }
@@ -212,7 +212,7 @@ static bool build_openai_edit_request(const httplib::Request& req,
     }
 
     // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
-    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", true)) {
+    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) {
         error_message = "invalid params";
         return false;
     }
@@ -253,6 +253,12 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
 
     svr.Post("/v1/images/generations", [runtime](const httplib::Request& req, httplib::Response& res) {
         try {
+            if (!runtime_supports_generation_mode(*runtime, IMG_GEN)) {
+                res.status = 400;
+                res.set_content(json({{"error", unsupported_generation_mode_error(IMG_GEN)}}).dump(), "application/json");
+                return;
+            }
+
             ImgGenJobRequest request;
             std::string error_message;
             if (!build_openai_generation_request(req, *runtime, request, error_message)) {
@@ -319,6 +325,12 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
 
     svr.Post("/v1/images/edits", [runtime](const httplib::Request& req, httplib::Response& res) {
         try {
+            if (!runtime_supports_generation_mode(*runtime, IMG_GEN)) {
+                res.status = 400;
+                res.set_content(json({{"error", unsupported_generation_mode_error(IMG_GEN)}}).dump(), "application/json");
+                return;
+            }
+
             ImgGenJobRequest request;
             std::string error_message;
             if (!build_openai_edit_request(req, *runtime, request, error_message)) {

examples/server/routes_sdapi.cpp

@@ -1,6 +1,7 @@
 #include "routes.h"
 
 #include <algorithm>
+#include <cctype>
 #include <cstring>
 #include <regex>
 #include <string_view>
@@ -35,14 +36,20 @@ static fs::path resolve_display_model_path(const ServerRuntime& runtime) {
     return {};
 }
 
+static std::string lower_ascii(std::string value) {
+    std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) {
+        return static_cast<char>(std::tolower(c));
+    });
+    return value;
+}
+
 static enum sample_method_t get_sdapi_sample_method(std::string name) {
     enum sample_method_t result = str_to_sample_method(name.c_str());
     if (result != SAMPLE_METHOD_COUNT) {
         return result;
     }
 
-    std::transform(name.begin(), name.end(), name.begin(),
+    name = lower_ascii(name);
-                   [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
     static const std::unordered_map<std::string_view, sample_method_t> hardcoded{
         {"euler a", EULER_A_SAMPLE_METHOD},
         {"k_euler_a", EULER_A_SAMPLE_METHOD},
@@ -114,6 +121,18 @@ static bool build_sdapi_img_gen_request(const json& j,
     request.gen_params.width                          = j.value("width", -1);
     request.gen_params.height                         = j.value("height", -1);
 
+    if (!img2img && j.value("enable_hr", false)) {
+        request.gen_params.hires_enabled = true;
+        request.gen_params.hires_scale   = j.value("hr_scale", request.gen_params.hires_scale);
+        request.gen_params.hires_width   = j.value("hr_resize_x", request.gen_params.hires_width);
+        request.gen_params.hires_height  = j.value("hr_resize_y", request.gen_params.hires_height);
+        request.gen_params.hires_steps   = j.value("hr_steps", request.gen_params.hires_steps);
+        request.gen_params.hires_denoising_strength =
+            j.value("denoising_strength", request.gen_params.hires_denoising_strength);
+
+        request.gen_params.hires_upscaler = j.value("hr_upscaler", request.gen_params.hires_upscaler);
+    }
+
     std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt);
     if (!sd_cpp_extra_args_str.empty() && !request.gen_params.from_json_str(sd_cpp_extra_args_str)) {
         error_message = "invalid sd_cpp_extra_args";
@@ -228,7 +247,7 @@ static bool build_sdapi_img_gen_request(const json& j,
     }
 
     // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
-    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", true)) {
+    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) {
         error_message = "invalid params";
         return false;
     }
@@ -246,6 +265,11 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
                 res.set_content(R"({"error":"empty body"})", "application/json");
                 return;
             }
+            if (!runtime_supports_generation_mode(*runtime, IMG_GEN)) {
+                res.status = 400;
+                res.set_content(json({{"error", unsupported_generation_mode_error(IMG_GEN)}}).dump(), "application/json");
+                return;
+            }
 
             json j = json::parse(req.body);
             ImgGenJobRequest request;
@@ -342,6 +366,52 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
         res.set_content(result.dump(), "application/json");
     });
 
+    svr.Get("/sdapi/v1/upscalers", [runtime](const httplib::Request&, httplib::Response& res) {
+        refresh_upscaler_cache(*runtime);
+
+        auto make_builtin = [](const char* name) {
+            json item;
+            item["name"]       = name;
+            item["model_name"] = nullptr;
+            item["model_path"] = nullptr;
+            item["model_url"]  = nullptr;
+            item["scale"]      = 4;
+            return item;
+        };
+
+        json result = json::array();
+        result.push_back(make_builtin("None"));
+        result.push_back(make_builtin("Lanczos"));
+        result.push_back(make_builtin("Nearest"));
+
+        {
+            std::lock_guard<std::mutex> lock(*runtime->upscaler_mutex);
+            for (const auto& e : *runtime->upscaler_cache) {
+                json item;
+                item["name"]       = e.name;
+                item["model_name"] = e.model_name;
+                item["model_path"] = e.fullpath;
+                item["model_url"]  = nullptr;
+                item["scale"]      = e.scale;
+                result.push_back(item);
+            }
+        }
+
+        res.set_content(result.dump(), "application/json");
+    });
+
+    svr.Get("/sdapi/v1/latent-upscale-modes", [](const httplib::Request&, httplib::Response& res) {
+        json result = json::array({
+            {{"name", "Latent"}},
+            {{"name", "Latent (nearest)"}},
+            {{"name", "Latent (nearest-exact)"}},
+            {{"name", "Latent (antialiased)"}},
+            {{"name", "Latent (bicubic)"}},
+            {{"name", "Latent (bicubic antialiased)"}},
+        });
+        res.set_content(result.dump(), "application/json");
+    });
+
     svr.Get("/sdapi/v1/samplers", [runtime](const httplib::Request&, httplib::Response& res) {
         std::vector<std::string> sampler_names;
         sampler_names.push_back("default");

examples/server/routes_sdcpp.cpp

@@ -75,61 +75,9 @@ static fs::path resolve_display_model_path(const ServerRuntime& runtime) {
     return {};
 }
 
-static json make_capabilities_json(ServerRuntime& runtime) {
+static json make_sample_params_json(const sd_sample_params_t& sample_params, const std::vector<int>& skip_layers) {
-    refresh_lora_cache(runtime);
-
-    AsyncJobManager& manager  = *runtime.async_job_manager;
-    const auto& defaults      = *runtime.default_gen_params;
-    const auto& sample_params = defaults.sample_params;
     const auto& guidance = sample_params.guidance;
-    const fs::path model_path = resolve_display_model_path(runtime);
+    return {
-    json samplers             = json::array();
-    json schedulers           = json::array();
-    json output_formats       = json::array({"png", "jpeg"});
-    json available_loras      = json::array();
-
-    for (int i = 0; i < SAMPLE_METHOD_COUNT; ++i) {
-        samplers.push_back(sd_sample_method_name((sample_method_t)i));
-    }
-
-    for (int i = 0; i < SCHEDULER_COUNT; ++i) {
-        schedulers.push_back(sd_scheduler_name((scheduler_t)i));
-    }
-
-#ifdef SD_USE_WEBP
-    output_formats.push_back("webp");
-#endif
-
-    {
-        std::lock_guard<std::mutex> lock(*runtime.lora_mutex);
-        for (const auto& entry : *runtime.lora_cache) {
-            available_loras.push_back({
-                {"name", entry.name},
-                {"path", entry.path},
-            });
-        }
-    }
-
-    json result;
-    result["model"] = {
-        {"name", model_path.filename().u8string()},
-        {"stem", model_path.stem().u8string()},
-        {"path", model_path.u8string()},
-    };
-    result["defaults"] = {
-        {"prompt", defaults.prompt},
-        {"negative_prompt", defaults.negative_prompt},
-        {"clip_skip", defaults.clip_skip},
-        {"width", defaults.width > 0 ? defaults.width : 512},
-        {"height", defaults.height > 0 ? defaults.height : 512},
-        {"strength", defaults.strength},
-        {"seed", defaults.seed},
-        {"batch_count", defaults.batch_count},
-        {"auto_resize_ref_image", defaults.auto_resize_ref_image},
-        {"increase_ref_index", defaults.increase_ref_index},
-        {"control_strength", defaults.control_strength},
-        {"sample_params",
-         {
         {"scheduler", capability_scheduler_name(sample_params.scheduler)},
         {"sample_method", capability_sample_method_name(sample_params.sample_method)},
         {"sample_steps", sample_params.sample_steps},
@@ -143,21 +91,234 @@ static json make_capabilities_json(ServerRuntime& runtime) {
              {"distilled_guidance", guidance.distilled_guidance},
              {"slg",
               {
-                       {"layers", defaults.skip_layers},
+                  {"layers", skip_layers},
                   {"layer_start", guidance.slg.layer_start},
                   {"layer_end", guidance.slg.layer_end},
                   {"scale", guidance.slg.scale},
               }},
          }},
+    };
+}
+
+static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const std::string& output_format) {
+    return {
+        {"prompt", defaults.prompt},
+        {"negative_prompt", defaults.negative_prompt},
+        {"clip_skip", defaults.clip_skip},
+        {"width", defaults.width > 0 ? defaults.width : 512},
+        {"height", defaults.height > 0 ? defaults.height : 512},
+        {"strength", defaults.strength},
+        {"seed", defaults.seed},
+        {"batch_count", defaults.batch_count},
+        {"auto_resize_ref_image", defaults.auto_resize_ref_image},
+        {"increase_ref_index", defaults.increase_ref_index},
+        {"control_strength", defaults.control_strength},
+        {"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers)},
+        {"hires",
+         {
+             {"enabled", defaults.hires_enabled},
+             {"upscaler", defaults.hires_upscaler},
+             {"scale", defaults.hires_scale},
+             {"target_width", defaults.hires_width},
+             {"target_height", defaults.hires_height},
+             {"steps", defaults.hires_steps},
+             {"denoising_strength", defaults.hires_denoising_strength},
+             {"upscale_tile_size", defaults.hires_upscale_tile_size},
          }},
         {"vae_tiling_params", make_vae_tiling_json(defaults.vae_tiling_params)},
         {"cache_mode", defaults.cache_mode},
         {"cache_option", defaults.cache_option},
         {"scm_mask", defaults.scm_mask},
         {"scm_policy_dynamic", defaults.scm_policy_dynamic},
-        {"output_format", "png"},
+        {"output_format", output_format},
         {"output_compression", 100},
     };
+}
+
+static json make_vid_gen_defaults_json(const SDGenerationParams& defaults, const std::string& output_format) {
+    return {
+        {"prompt", defaults.prompt},
+        {"negative_prompt", defaults.negative_prompt},
+        {"clip_skip", defaults.clip_skip},
+        {"width", defaults.width > 0 ? defaults.width : 512},
+        {"height", defaults.height > 0 ? defaults.height : 512},
+        {"strength", defaults.strength},
+        {"seed", defaults.seed},
+        {"video_frames", defaults.video_frames},
+        {"fps", defaults.fps},
+        {"moe_boundary", defaults.moe_boundary},
+        {"vace_strength", defaults.vace_strength},
+        {"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers)},
+        {"high_noise_sample_params", make_sample_params_json(defaults.high_noise_sample_params, defaults.high_noise_skip_layers)},
+        {"vae_tiling_params", make_vae_tiling_json(defaults.vae_tiling_params)},
+        {"cache_mode", defaults.cache_mode},
+        {"cache_option", defaults.cache_option},
+        {"scm_mask", defaults.scm_mask},
+        {"scm_policy_dynamic", defaults.scm_policy_dynamic},
+        {"output_format", output_format},
+        {"output_compression", 100},
+    };
+}
+
+static json make_img_gen_features_json() {
+    return {
+        {"init_image", true},
+        {"mask_image", true},
+        {"control_image", true},
+        {"ref_images", true},
+        {"lora", true},
+        {"vae_tiling", true},
+        {"hires", true},
+        {"cache", true},
+        {"cancel_queued", true},
+        {"cancel_generating", false},
+    };
+}
+
+static json make_vid_gen_features_json() {
+    return {
+        {"init_image", true},
+        {"end_image", true},
+        {"control_frames", true},
+        {"high_noise_sample_params", true},
+        {"lora", true},
+        {"vae_tiling", true},
+        {"cache", true},
+        {"cancel_queued", true},
+        {"cancel_generating", false},
+    };
+}
+
+static json make_capabilities_json(ServerRuntime& runtime) {
+    refresh_lora_cache(runtime);
+    refresh_upscaler_cache(runtime);
+
+    AsyncJobManager& manager  = *runtime.async_job_manager;
+    const auto& defaults      = *runtime.default_gen_params;
+    const fs::path model_path = resolve_display_model_path(runtime);
+    const bool supports_img   = runtime_supports_generation_mode(runtime, IMG_GEN);
+    const bool supports_vid   = runtime_supports_generation_mode(runtime, VID_GEN);
+    json samplers             = json::array();
+    json schedulers           = json::array();
+    json image_output_formats = supported_img_output_formats();
+    json video_output_formats = supported_vid_output_formats();
+    json available_loras      = json::array();
+    json available_upscalers  = json::array();
+    json supported_modes      = json::array();
+
+    for (int i = 0; i < SAMPLE_METHOD_COUNT; ++i) {
+        samplers.push_back(sd_sample_method_name((sample_method_t)i));
+    }
+
+    for (int i = 0; i < SCHEDULER_COUNT; ++i) {
+        schedulers.push_back(sd_scheduler_name((scheduler_t)i));
+    }
+
+    {
+        std::lock_guard<std::mutex> lock(*runtime.lora_mutex);
+        for (const auto& entry : *runtime.lora_cache) {
+            available_loras.push_back({
+                {"name", entry.name},
+                {"path", entry.path},
+            });
+        }
+    }
+
+    available_upscalers.push_back({
+        {"name", "None"},
+    });
+    available_upscalers.push_back({
+        {"name", "Lanczos"},
+    });
+    available_upscalers.push_back({
+        {"name", "Nearest"},
+    });
+    available_upscalers.push_back({
+        {"name", "Latent"},
+    });
+    available_upscalers.push_back({
+        {"name", "Latent (nearest)"},
+    });
+    available_upscalers.push_back({
+        {"name", "Latent (nearest-exact)"},
+    });
+    available_upscalers.push_back({
+        {"name", "Latent (antialiased)"},
+    });
+    available_upscalers.push_back({
+        {"name", "Latent (bicubic)"},
+    });
+    available_upscalers.push_back({
+        {"name", "Latent (bicubic antialiased)"},
+    });
+    {
+        std::lock_guard<std::mutex> lock(*runtime.upscaler_mutex);
+        for (const auto& entry : *runtime.upscaler_cache) {
+            available_upscalers.push_back({
+                {"name", entry.name},
+            });
+        }
+    }
+
+    if (supports_img) {
+        supported_modes.push_back("img_gen");
+    }
+    if (supports_vid) {
+        supported_modes.push_back("vid_gen");
+    }
+
+    std::string default_img_output_format = "png";
+    std::string default_vid_output_format = "avi";
+    if (!image_output_formats.empty()) {
+        default_img_output_format = image_output_formats[0].get<std::string>();
+    }
+    if (!video_output_formats.empty()) {
+        default_vid_output_format = video_output_formats[0].get<std::string>();
+    }
+
+    json defaults_by_mode       = json::object();
+    json output_formats_by_mode = json::object();
+    json features_by_mode       = json::object();
+    if (supports_img) {
+        defaults_by_mode["img_gen"]       = make_img_gen_defaults_json(defaults, default_img_output_format);
+        output_formats_by_mode["img_gen"] = image_output_formats;
+        features_by_mode["img_gen"]       = make_img_gen_features_json();
+    }
+    if (supports_vid) {
+        defaults_by_mode["vid_gen"]       = make_vid_gen_defaults_json(defaults, default_vid_output_format);
+        output_formats_by_mode["vid_gen"] = video_output_formats;
+        features_by_mode["vid_gen"]       = make_vid_gen_features_json();
+    }
+
+    json top_level_defaults       = json::object();
+    json top_level_output_formats = json::array();
+    json top_level_features       = {
+              {"cancel_queued", true},
+              {"cancel_generating", false},
+    };
+    std::string current_mode = "";
+    if (supports_img) {
+        current_mode             = "img_gen";
+        top_level_defaults       = defaults_by_mode["img_gen"];
+        top_level_output_formats = output_formats_by_mode["img_gen"];
+        top_level_features       = features_by_mode["img_gen"];
+    } else if (supports_vid) {
+        current_mode             = "vid_gen";
+        top_level_defaults       = defaults_by_mode["vid_gen"];
+        top_level_output_formats = output_formats_by_mode["vid_gen"];
+        top_level_features       = features_by_mode["vid_gen"];
+    }
+
+    json result;
+    result["model"] = {
+        {"name", model_path.filename().u8string()},
+        {"stem", model_path.stem().u8string()},
+        {"path", model_path.u8string()},
+    };
+    result["current_mode"]     = current_mode;
+    result["supported_modes"]  = supported_modes;
+    result["defaults"]         = top_level_defaults;
+    result["defaults_by_mode"] = defaults_by_mode;
     result["limits"]           = {
                   {"min_width", 64},
                   {"max_width", 4096},
@@ -168,19 +329,12 @@ static json make_capabilities_json(ServerRuntime& runtime) {
     };
     result["samplers"]               = samplers;
     result["schedulers"]             = schedulers;
-    result["output_formats"] = output_formats;
+    result["output_formats"]         = top_level_output_formats;
-    result["features"]       = {
+    result["output_formats_by_mode"] = output_formats_by_mode;
-              {"init_image", true},
+    result["features"]               = top_level_features;
-              {"mask_image", true},
+    result["features_by_mode"]       = features_by_mode;
-              {"control_image", true},
-              {"ref_images", true},
-              {"lora", true},
-              {"vae_tiling", true},
-              {"cache", true},
-              {"cancel_queued", true},
-              {"cancel_generating", false},
-    };
     result["loras"]                  = available_loras;
+    result["upscalers"]              = available_upscalers;
     return result;
 }
 
@@ -204,7 +358,34 @@ static bool parse_img_gen_request(const json& body,
         return false;
     }
     // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
-    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", true)) {
+    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) {
+        error_message = "invalid generation parameters";
+        return false;
+    }
+    return true;
+}
+
+static bool parse_vid_gen_request(const json& body,
+                                  ServerRuntime& runtime,
+                                  VidGenJobRequest& request,
+                                  std::string& error_message) {
+    request.gen_params = *runtime.default_gen_params;
+
+    refresh_lora_cache(runtime);
+    if (!request.gen_params.from_json_str(body.dump(), [&](const std::string& path) {
+            return get_lora_full_path(runtime, path);
+        })) {
+        error_message = "invalid generation parameters";
+        return false;
+    }
+
+    std::string output_format = body.value("output_format", "webm");
+    int output_compression    = body.value("output_compression", 100);
+    if (!assign_output_options(request, output_format, output_compression, error_message)) {
+        return false;
+    }
+    // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
+    if (!request.gen_params.resolve_and_validate(VID_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) {
         error_message = "invalid generation parameters";
         return false;
     }
@@ -226,6 +407,11 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
                 res.set_content(R"({"error":"empty body"})", "application/json");
                 return;
             }
+            if (!runtime_supports_generation_mode(*runtime, IMG_GEN)) {
+                res.status = 400;
+                res.set_content(json({{"error", unsupported_generation_mode_error(IMG_GEN)}}).dump(), "application/json");
+                return;
+            }
 
             json body = json::parse(req.body);
             ImgGenJobRequest request;
@@ -276,9 +462,66 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
         }
     });
 
-    svr.Post("/sdcpp/v1/vid_gen", [](const httplib::Request&, httplib::Response& res) {
+    svr.Post("/sdcpp/v1/vid_gen", [runtime](const httplib::Request& req, httplib::Response& res) {
-        res.status = 501;
+        try {
-        res.set_content(R"({"error":"vid_gen is reserved and not implemented yet"})", "application/json");
+            if (req.body.empty()) {
+                res.status = 400;
+                res.set_content(R"({"error":"empty body"})", "application/json");
+                return;
+            }
+            if (!runtime_supports_generation_mode(*runtime, VID_GEN)) {
+                res.status = 400;
+                res.set_content(json({{"error", unsupported_generation_mode_error(VID_GEN)}}).dump(), "application/json");
+                return;
+            }
+
+            json body = json::parse(req.body);
+            VidGenJobRequest request;
+            std::string error_message;
+            if (!parse_vid_gen_request(body, *runtime, request, error_message)) {
+                res.status = 400;
+                res.set_content(json({{"error", error_message}}).dump(), "application/json");
+                return;
+            }
+
+            AsyncJobManager& manager                = *runtime->async_job_manager;
+            std::shared_ptr<AsyncGenerationJob> job = std::make_shared<AsyncGenerationJob>();
+            job->kind                               = AsyncJobKind::VidGen;
+            job->status                             = AsyncJobStatus::Queued;
+            job->created_at                         = unix_timestamp_now();
+            job->vid_gen                            = std::move(request);
+
+            {
+                std::lock_guard<std::mutex> lock(manager.mutex);
+                purge_expired_jobs(manager);
+                if (count_pending_jobs(manager) >= manager.max_pending_jobs) {
+                    res.status = 429;
+                    res.set_content(R"({"error":"job queue is full"})", "application/json");
+                    return;
+                }
+                job->id               = make_async_job_id(manager);
+                manager.jobs[job->id] = job;
+                manager.queue.push_back(job->id);
+            }
+
+            manager.cv.notify_one();
+
+            json out;
+            out["id"]       = job->id;
+            out["kind"]     = async_job_kind_name(job->kind);
+            out["status"]   = async_job_status_name(job->status);
+            out["created"]  = job->created_at;
+            out["poll_url"] = "/sdcpp/v1/jobs/" + job->id;
+
+            res.status = 202;
+            res.set_content(out.dump(), "application/json");
+        } catch (const json::parse_error& e) {
+            res.status = 400;
+            res.set_content(json({{"error", "invalid json"}, {"message", e.what()}}).dump(), "application/json");
+        } catch (const std::exception& e) {
+            res.status = 500;
+            res.set_content(json({{"error", "server_error"}, {"message", e.what()}}).dump(), "application/json");
+        }
     });
 
     svr.Get(R"(/sdcpp/v1/jobs/([A-Za-z0-9_\-]+))", [runtime](const httplib::Request& req, httplib::Response& res) {

examples/server/runtime.cpp

@@ -1,6 +1,7 @@
 #include "runtime.h"
 
 #include <algorithm>
+#include <cctype>
 #include <chrono>
 #include <cstdlib>
 #include <filesystem>
@@ -13,6 +14,18 @@
 
 namespace fs = std::filesystem;
 
+static std::string lower_ascii(std::string value) {
+    std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) {
+        return static_cast<char>(std::tolower(c));
+    });
+    return value;
+}
+
+static bool is_supported_model_ext(const fs::path& p) {
+    auto ext = lower_ascii(p.extension().string());
+    return ext == ".gguf" || ext == ".pt" || ext == ".pth" || ext == ".safetensors";
+}
+
 static const std::string k_base64_chars =
     "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
     "abcdefghijklmnopqrstuvwxyz"
@@ -45,6 +58,44 @@ std::string normalize_output_format(std::string output_format) {
     return output_format;
 }
 
+std::vector<std::string> supported_img_output_formats(bool allow_webp) {
+    std::vector<std::string> formats = {"png", "jpeg"};
+#ifdef SD_USE_WEBP
+    if (allow_webp) {
+        formats.push_back("webp");
+    }
+#else
+    (void)allow_webp;
+#endif
+    return formats;
+}
+
+std::vector<std::string> supported_vid_output_formats() {
+    std::vector<std::string> formats;
+#ifdef SD_USE_WEBM
+    formats.push_back("webm");
+#endif
+#ifdef SD_USE_WEBP
+    formats.push_back("webp");
+#endif
+    formats.push_back("avi");
+    return formats;
+}
+
+static std::string valid_vid_output_formats_message() {
+    const std::vector<std::string> formats = supported_vid_output_formats();
+
+    std::string message = "invalid output_format, must be one of [";
+    for (size_t i = 0; i < formats.size(); ++i) {
+        if (i > 0) {
+            message += ", ";
+        }
+        message += formats[i];
+    }
+    message += "]";
+    return message;
+}
+
 bool assign_output_options(ImgGenJobRequest& request,
                            std::string output_format,
                            int output_compression,
@@ -53,19 +104,88 @@ bool assign_output_options(ImgGenJobRequest& request,
     request.output_format      = normalize_output_format(std::move(output_format));
     request.output_compression = std::clamp(output_compression, 0, 100);
 
-    const bool valid_format = request.output_format == "png" ||
+    const std::vector<std::string> valid_formats = supported_img_output_formats(allow_webp);
-                              request.output_format == "jpeg" ||
+    const bool valid_format                      = std::find(valid_formats.begin(),
-                              (allow_webp && request.output_format == "webp");
+                                                             valid_formats.end(),
+                                                             request.output_format) != valid_formats.end();
     if (!valid_format) {
-        error_message = allow_webp
+        error_message = "invalid output_format, must be one of [";
-                            ? "invalid output_format, must be one of [png, jpeg, webp]"
+        for (size_t i = 0; i < valid_formats.size(); ++i) {
-                            : "invalid output_format, must be one of [png, jpeg]";
+            if (i > 0) {
+                error_message += ", ";
+            }
+            error_message += valid_formats[i];
+        }
+        error_message += "]";
         return false;
     }
 
     return true;
 }
 
+bool assign_output_options(VidGenJobRequest& request,
+                           std::string output_format,
+                           int output_compression,
+                           std::string& error_message) {
+    request.output_format      = normalize_output_format(std::move(output_format));
+    request.output_compression = std::clamp(output_compression, 0, 100);
+
+    if (request.output_format == "avi") {
+        return true;
+    }
+
+    if (request.output_format == "webm") {
+#ifdef SD_USE_WEBM
+        return true;
+#else
+        error_message = valid_vid_output_formats_message();
+        return false;
+#endif
+    }
+
+    if (request.output_format == "webp") {
+#ifdef SD_USE_WEBP
+        return true;
+#else
+        error_message = valid_vid_output_formats_message();
+        return false;
+#endif
+    }
+
+    error_message = valid_vid_output_formats_message();
+    return false;
+}
+
+std::string video_mime_type(const std::string& output_format) {
+    if (output_format == "webm") {
+        return "video/webm";
+    }
+    if (output_format == "webp") {
+        return "image/webp";
+    }
+    return "video/x-msvideo";
+}
+
+bool runtime_supports_generation_mode(const ServerRuntime& runtime, SDMode mode) {
+    if (mode == VID_GEN) {
+        return sd_ctx_supports_video_generation(runtime.sd_ctx);
+    }
+    if (mode == IMG_GEN) {
+        return sd_ctx_supports_image_generation(runtime.sd_ctx);
+    }
+    return true;
+}
+
+std::string unsupported_generation_mode_error(SDMode mode) {
+    if (mode == VID_GEN) {
+        return "loaded model does not support vid_gen";
+    }
+    if (mode == IMG_GEN) {
+        return "loaded model does not support img_gen";
+    }
+    return "loaded model does not support requested mode";
+}
+
 ArgOptions SDSvrParams::get_options() {
     ArgOptions options;
 
@@ -134,20 +254,12 @@ void refresh_lora_cache(ServerRuntime& rt) {
 
     fs::path lora_dir = rt.ctx_params->lora_model_dir;
     if (fs::exists(lora_dir) && fs::is_directory(lora_dir)) {
-        auto is_lora_ext = [](const fs::path& p) {
-            auto ext = p.extension().string();
-            std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) {
-                return static_cast<char>(std::tolower(c));
-            });
-            return ext == ".gguf" || ext == ".pt" || ext == ".pth" || ext == ".safetensors";
-        };
-
         for (auto& entry : fs::recursive_directory_iterator(lora_dir)) {
             if (!entry.is_regular_file()) {
                 continue;
             }
             const fs::path& p = entry.path();
-            if (!is_lora_ext(p)) {
+            if (!is_supported_model_ext(p)) {
                 continue;
             }
 
@@ -179,6 +291,40 @@ std::string get_lora_full_path(ServerRuntime& rt, const std::string& path) {
     return it != rt.lora_cache->end() ? it->fullpath : "";
 }
 
+void refresh_upscaler_cache(ServerRuntime& rt) {
+    std::vector<UpscalerEntry> new_cache;
+
+    fs::path upscaler_dir = rt.ctx_params->hires_upscalers_dir;
+    if (fs::exists(upscaler_dir) && fs::is_directory(upscaler_dir)) {
+        for (auto& entry : fs::directory_iterator(upscaler_dir)) {
+            if (!entry.is_regular_file()) {
+                continue;
+            }
+            const fs::path& p = entry.path();
+            if (!is_supported_model_ext(p)) {
+                continue;
+            }
+
+            UpscalerEntry upscaler_entry;
+            upscaler_entry.name       = p.stem().u8string();
+            upscaler_entry.fullpath   = fs::absolute(p).lexically_normal().u8string();
+            upscaler_entry.model_name = "ESRGAN_4x";
+            upscaler_entry.path       = p.filename().u8string();
+
+            new_cache.push_back(std::move(upscaler_entry));
+        }
+    }
+
+    std::sort(new_cache.begin(), new_cache.end(), [](const UpscalerEntry& a, const UpscalerEntry& b) {
+        return a.name < b.name;
+    });
+
+    {
+        std::lock_guard<std::mutex> lock(*rt.upscaler_mutex);
+        *rt.upscaler_cache = std::move(new_cache);
+    }
+}
+
 int64_t unix_timestamp_now() {
     return std::chrono::duration_cast<std::chrono::seconds>(
                std::chrono::system_clock::now().time_since_epoch())

examples/server/runtime.h

@@ -37,6 +37,14 @@ struct LoraEntry {
     std::string fullpath;
 };
 
+struct UpscalerEntry {
+    std::string name;
+    std::string path;
+    std::string fullpath;
+    std::string model_name;
+    int scale = 4;
+};
+
 struct ServerRuntime {
     sd_ctx_t* sd_ctx;
     std::mutex* sd_ctx_mutex;
@@ -45,6 +53,8 @@ struct ServerRuntime {
     const SDGenerationParams* default_gen_params;
     std::vector<LoraEntry>* lora_cache;
     std::mutex* lora_mutex;
+    std::vector<UpscalerEntry>* upscaler_cache;
+    std::mutex* upscaler_mutex;
     AsyncJobManager* async_job_manager;
 };
 
@@ -58,13 +68,33 @@ struct ImgGenJobRequest {
     }
 };
 
+struct VidGenJobRequest {
+    SDGenerationParams gen_params;
+    std::string output_format = "webm";
+    int output_compression    = 100;
+
+    sd_vid_gen_params_t to_sd_vid_gen_params_t() {
+        return gen_params.to_sd_vid_gen_params_t();
+    }
+};
+
 std::string base64_encode(const std::vector<uint8_t>& bytes);
 std::string normalize_output_format(std::string output_format);
+std::vector<std::string> supported_img_output_formats(bool allow_webp = true);
+std::vector<std::string> supported_vid_output_formats();
 bool assign_output_options(ImgGenJobRequest& request,
                            std::string output_format,
                            int output_compression,
                            bool allow_webp,
                            std::string& error_message);
+bool assign_output_options(VidGenJobRequest& request,
+                           std::string output_format,
+                           int output_compression,
+                           std::string& error_message);
+std::string video_mime_type(const std::string& output_format);
+bool runtime_supports_generation_mode(const ServerRuntime& runtime, SDMode mode);
+std::string unsupported_generation_mode_error(SDMode mode);
 void refresh_lora_cache(ServerRuntime& rt);
 std::string get_lora_full_path(ServerRuntime& rt, const std::string& path);
+void refresh_upscaler_cache(ServerRuntime& rt);
 int64_t unix_timestamp_now();

format-code.sh

@@ -1,5 +1,5 @@
-for f in src/*.cpp src/*.h src/*.hpp src/vocab/*.h src/vocab/*.cpp \
+for f in src/*.cpp src/*.h src/*.hpp src/tokenizers/*.h src/tokenizers/*.cpp src/tokenizers/vocab/*.h src/tokenizers/vocab/*.cpp \
-         examples/cli/*.cpp examples/cli/*.h examples/server/*.cpp \
+         src/model_io/*.h src/model_io/*.cpp examples/cli/*.cpp examples/cli/*.h examples/server/*.cpp \
          examples/common/*.hpp examples/common/*.h examples/common/*.cpp; do
   [[ "$f" == vocab* ]] && continue
   echo "formatting '$f'"

include/stable-diffusion.h

@@ -50,6 +50,7 @@ enum sample_method_t {
     TCD_SAMPLE_METHOD,
     RES_MULTISTEP_SAMPLE_METHOD,
     RES_2S_SAMPLE_METHOD,
+    ER_SDE_SAMPLE_METHOD,
     SAMPLE_METHOD_COUNT
 };
 
@@ -288,6 +289,32 @@ typedef struct {
     const char* path;
 } sd_lora_t;
 
+enum sd_hires_upscaler_t {
+    SD_HIRES_UPSCALER_NONE,
+    SD_HIRES_UPSCALER_LATENT,
+    SD_HIRES_UPSCALER_LATENT_NEAREST,
+    SD_HIRES_UPSCALER_LATENT_NEAREST_EXACT,
+    SD_HIRES_UPSCALER_LATENT_ANTIALIASED,
+    SD_HIRES_UPSCALER_LATENT_BICUBIC,
+    SD_HIRES_UPSCALER_LATENT_BICUBIC_ANTIALIASED,
+    SD_HIRES_UPSCALER_LANCZOS,
+    SD_HIRES_UPSCALER_NEAREST,
+    SD_HIRES_UPSCALER_MODEL,
+    SD_HIRES_UPSCALER_COUNT,
+};
+
+typedef struct {
+    bool enabled;
+    enum sd_hires_upscaler_t upscaler;
+    const char* model_path;
+    float scale;
+    int target_width;
+    int target_height;
+    int steps;
+    float denoising_strength;
+    int upscale_tile_size;
+} sd_hires_params_t;
+
 typedef struct {
     const sd_lora_t* loras;
     uint32_t lora_count;
@@ -311,6 +338,7 @@ typedef struct {
     sd_pm_params_t pm_params;
     sd_tiling_params_t vae_tiling_params;
     sd_cache_params_t cache;
+    sd_hires_params_t hires;
 } sd_img_gen_params_t;
 
 typedef struct {
@@ -347,6 +375,8 @@ 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 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);
+SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx);
 
 SD_API const char* sd_type_name(enum sd_type_t type);
 SD_API enum sd_type_t str_to_sd_type(const char* str);
@@ -362,8 +392,11 @@ SD_API const char* sd_preview_name(enum preview_t preview);
 SD_API enum preview_t str_to_preview(const char* str);
 SD_API const char* sd_lora_apply_mode_name(enum lora_apply_mode_t mode);
 SD_API enum lora_apply_mode_t str_to_lora_apply_mode(const char* str);
+SD_API const char* sd_hires_upscaler_name(enum sd_hires_upscaler_t upscaler);
+SD_API enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str);
 
 SD_API void sd_cache_params_init(sd_cache_params_t* cache_params);
+SD_API void sd_hires_params_init(sd_hires_params_t* hires_params);
 
 SD_API void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params);
 SD_API char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params);

src/auto_encoder_kl.hpp

@@ -533,7 +533,7 @@ public:
                        const std::string& prefix                      = "")
         : version(version), decode_only(decode_only), use_video_decoder(use_video_decoder) {
         if (sd_version_is_dit(version)) {
-            if (sd_version_is_flux2(version)) {
+            if (sd_version_uses_flux2_vae(version)) {
                 dd_config.z_channels = 32;
                 embed_dim            = 32;
             } else {
@@ -578,7 +578,7 @@ public:
 
     ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
         // z: [N, z_channels, h, w]
-        if (sd_version_is_flux2(version)) {
+        if (sd_version_uses_flux2_vae(version)) {
             // [N, C*p*p, h, w] -> [N, C, h*p, w*p]
             int64_t p = 2;
 
@@ -617,7 +617,7 @@ public:
             auto quant_conv = std::dynamic_pointer_cast<Conv2d>(blocks["quant_conv"]);
             z               = quant_conv->forward(ctx, z);  // [N, 2*embed_dim, h/8, w/8]
         }
-        if (sd_version_is_flux2(version)) {
+        if (sd_version_uses_flux2_vae(version)) {
             z = ggml_ext_chunk(ctx->ggml_ctx, z, 2, 2)[0];
 
             // [N, C, H, W] -> [N, C*p*p, H/p, W/p]
@@ -640,7 +640,7 @@ public:
 
     int get_encoder_output_channels() {
         int factor = dd_config.double_z ? 2 : 1;
-        if (sd_version_is_flux2(version)) {
+        if (sd_version_uses_flux2_vae(version)) {
             return dd_config.z_channels * 4;
         }
         return dd_config.z_channels * factor;
@@ -673,7 +673,7 @@ struct AutoEncoderKL : public VAE {
         } else if (sd_version_is_flux(version) || sd_version_is_z_image(version)) {
             scale_factor = 0.3611f;
             shift_factor = 0.1159f;
-        } else if (sd_version_is_flux2(version)) {
+        } else if (sd_version_uses_flux2_vae(version)) {
             scale_factor = 1.0f;
             shift_factor = 0.f;
         }
@@ -747,7 +747,7 @@ struct AutoEncoderKL : public VAE {
     }
 
     sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
-        if (sd_version_is_flux2(version)) {
+        if (sd_version_uses_flux2_vae(version)) {
             return vae_output;
         } else if (version == VERSION_SD1_PIX2PIX) {
             return sd::ops::chunk(vae_output, 2, 2)[0];
@@ -758,7 +758,7 @@ struct AutoEncoderKL : public VAE {
 
     std::pair<sd::Tensor<float>, sd::Tensor<float>> get_latents_mean_std(const sd::Tensor<float>& latents, int channel_dim) {
         GGML_ASSERT(channel_dim >= 0 && static_cast<size_t>(channel_dim) < static_cast<size_t>(latents.dim()));
-        if (sd_version_is_flux2(version)) {
+        if (sd_version_uses_flux2_vae(version)) {
             GGML_ASSERT(latents.shape()[channel_dim] == 128);
             std::vector<int64_t> stats_shape(static_cast<size_t>(latents.dim()), 1);
             stats_shape[static_cast<size_t>(channel_dim)] = latents.shape()[channel_dim];
@@ -804,7 +804,7 @@ struct AutoEncoderKL : public VAE {
     }
 
     sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
-        if (sd_version_is_flux2(version)) {
+        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;
@@ -813,7 +813,7 @@ struct AutoEncoderKL : public VAE {
     }
 
     sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
-        if (sd_version_is_flux2(version)) {
+        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 - mean_tensor) * scale_factor) / std_tensor;

src/clip.hpp

@@ -3,455 +3,7 @@
 
 #include "ggml_extend.hpp"
 #include "model.h"
-#include "tokenize_util.h"
+#include "tokenizers/clip_tokenizer.h"
-#include "vocab/vocab.h"
-
-/*================================================== CLIPTokenizer ===================================================*/
-
-__STATIC_INLINE__ std::vector<std::pair<int, std::u32string>> bytes_to_unicode() {
-    std::vector<std::pair<int, std::u32string>> byte_unicode_pairs;
-    std::set<int> byte_set;
-    for (int b = static_cast<int>('!'); b <= static_cast<int>('~'); ++b) {
-        byte_set.insert(b);
-        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
-    }
-    for (int b = 161; b <= 172; ++b) {
-        byte_set.insert(b);
-        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
-    }
-    for (int b = 174; b <= 255; ++b) {
-        byte_set.insert(b);
-        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
-    }
-    int n = 0;
-    for (int b = 0; b < 256; ++b) {
-        if (byte_set.find(b) == byte_set.end()) {
-            byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(n + 256)));
-            ++n;
-        }
-    }
-    // LOG_DEBUG("byte_unicode_pairs %d", byte_unicode_pairs.size());
-    return byte_unicode_pairs;
-}
-
-// Ref: https://github.com/openai/CLIP/blob/main/clip/simple_tokenizer.py
-
-typedef std::function<bool(std::string&, std::vector<int32_t>&)> on_new_token_cb_t;
-
-class CLIPTokenizer {
-private:
-    std::map<int, std::u32string> byte_encoder;
-    std::map<std::u32string, int> byte_decoder;
-    std::map<std::u32string, int> encoder;
-    std::map<int, std::u32string> decoder;
-    std::map<std::pair<std::u32string, std::u32string>, int> bpe_ranks;
-    std::regex pat;
-    int encoder_len;
-    int bpe_len;
-
-    std::vector<std::string> special_tokens;
-
-public:
-    const std::string UNK_TOKEN = "<|endoftext|>";
-    const std::string BOS_TOKEN = "<|startoftext|>";
-    const std::string EOS_TOKEN = "<|endoftext|>";
-    const std::string PAD_TOKEN = "<|endoftext|>";
-
-    const int UNK_TOKEN_ID = 49407;
-    const int BOS_TOKEN_ID = 49406;
-    const int EOS_TOKEN_ID = 49407;
-    const int PAD_TOKEN_ID = 49407;
-
-private:
-    static std::string strip(const std::string& str) {
-        std::string::size_type start = str.find_first_not_of(" \t\n\r\v\f");
-        std::string::size_type end   = str.find_last_not_of(" \t\n\r\v\f");
-
-        if (start == std::string::npos) {
-            // String contains only whitespace characters
-            return "";
-        }
-
-        return str.substr(start, end - start + 1);
-    }
-
-    static std::string whitespace_clean(std::string text) {
-        text = std::regex_replace(text, std::regex(R"(\s+)"), " ");
-        text = strip(text);
-        return text;
-    }
-
-    static std::set<std::pair<std::u32string, std::u32string>> get_pairs(const std::vector<std::u32string>& subwords) {
-        std::set<std::pair<std::u32string, std::u32string>> pairs;
-        if (subwords.size() == 0) {
-            return pairs;
-        }
-        std::u32string prev_subword = subwords[0];
-        for (int i = 1; i < subwords.size(); i++) {
-            std::u32string subword = subwords[i];
-            std::pair<std::u32string, std::u32string> pair(prev_subword, subword);
-            pairs.insert(pair);
-            prev_subword = subword;
-        }
-        return pairs;
-    }
-
-    bool is_special_token(const std::string& token) {
-        for (auto& special_token : special_tokens) {
-            if (special_token == token) {
-                return true;
-            }
-        }
-        return false;
-    }
-
-public:
-    CLIPTokenizer(int pad_token_id = 49407, const std::string& merges_utf8_str = "")
-        : PAD_TOKEN_ID(pad_token_id) {
-        if (merges_utf8_str.size() > 0) {
-            load_from_merges(merges_utf8_str);
-        } else {
-            load_from_merges(load_clip_merges());
-        }
-        add_special_token("<|startoftext|>");
-        add_special_token("<|endoftext|>");
-    }
-
-    void load_from_merges(const std::string& merges_utf8_str) {
-        auto byte_unicode_pairs = bytes_to_unicode();
-        // printf("byte_unicode_pairs have %lu pairs \n", byte_unicode_pairs.size());
-        byte_encoder = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
-        for (auto& pair : byte_unicode_pairs) {
-            byte_decoder[pair.second] = pair.first;
-        }
-        // for (auto & pair: byte_unicode_pairs) {
-        //     std::cout << pair.first << ": " << pair.second << std::endl;
-        // }
-        std::vector<std::u32string> merges;
-        size_t start = 0;
-        size_t pos;
-        std::u32string merges_utf32_str = utf8_to_utf32(merges_utf8_str);
-        while ((pos = merges_utf32_str.find('\n', start)) != std::string::npos) {
-            merges.push_back(merges_utf32_str.substr(start, pos - start));
-            start = pos + 1;
-        }
-        // LOG_DEBUG("merges size %llu", merges.size());
-        GGML_ASSERT(merges.size() == 48895);
-        merges = std::vector<std::u32string>(merges.begin() + 1, merges.end());
-        std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
-        for (const auto& merge : merges) {
-            size_t space_pos = merge.find(' ');
-            merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
-            // LOG_DEBUG("%s", utf32_to_utf8(merge.substr(space_pos + 1)).c_str());
-            // printf("%s :: %s | %s \n", utf32_to_utf8(merge).c_str(), utf32_to_utf8(merge.substr(0, space_pos)).c_str(),
-            //                     utf32_to_utf8(merge.substr(space_pos + 1)).c_str());
-        }
-        std::vector<std::u32string> vocab;
-        for (const auto& pair : byte_unicode_pairs) {
-            vocab.push_back(pair.second);
-        }
-        for (const auto& pair : byte_unicode_pairs) {
-            vocab.push_back(pair.second + utf8_to_utf32("</w>"));
-        }
-        for (const auto& merge : merge_pairs) {
-            vocab.push_back(merge.first + merge.second);
-        }
-        vocab.push_back(utf8_to_utf32("<|startoftext|>"));
-        vocab.push_back(utf8_to_utf32("<|endoftext|>"));
-        LOG_DEBUG("vocab size: %llu", vocab.size());
-        int i = 0;
-        for (const auto& token : vocab) {
-            encoder[token] = i;
-            decoder[i]     = token;
-            i++;
-        }
-        encoder_len = i;
-
-        auto it = encoder.find(utf8_to_utf32("img</w>"));
-        if (it != encoder.end()) {
-            LOG_DEBUG("trigger word img already in vocab");
-        } else {
-            LOG_DEBUG("trigger word img not in vocab yet");
-        }
-
-        int rank = 0;
-        for (const auto& merge : merge_pairs) {
-            bpe_ranks[merge] = rank++;
-        }
-        bpe_len = rank;
-    };
-
-    void add_token(const std::string& text) {
-        std::u32string token = utf8_to_utf32(text);
-        auto it              = encoder.find(token);
-        if (it != encoder.end()) {
-            encoder[token]       = encoder_len;
-            decoder[encoder_len] = token;
-            encoder_len++;
-        }
-    }
-
-    void add_special_token(const std::string& token) {
-        special_tokens.push_back(token);
-    }
-
-    std::u32string bpe(const std::u32string& token) {
-        std::vector<std::u32string> word;
-
-        for (int i = 0; i < token.size() - 1; i++) {
-            word.emplace_back(1, token[i]);
-        }
-        word.push_back(token.substr(token.size() - 1) + utf8_to_utf32("</w>"));
-
-        std::set<std::pair<std::u32string, std::u32string>> pairs = get_pairs(word);
-
-        if (pairs.empty()) {
-            return token + utf8_to_utf32("</w>");
-        }
-
-        while (true) {
-            auto min_pair_iter = std::min_element(pairs.begin(),
-                                                  pairs.end(),
-                                                  [&](const std::pair<std::u32string, std::u32string>& a,
-                                                      const std::pair<std::u32string, std::u32string>& b) {
-                                                      if (bpe_ranks.find(a) == bpe_ranks.end()) {
-                                                          return false;
-                                                      } else if (bpe_ranks.find(b) == bpe_ranks.end()) {
-                                                          return true;
-                                                      }
-                                                      return bpe_ranks.at(a) < bpe_ranks.at(b);
-                                                  });
-
-            const std::pair<std::u32string, std::u32string>& bigram = *min_pair_iter;
-
-            if (bpe_ranks.find(bigram) == bpe_ranks.end()) {
-                break;
-            }
-
-            std::u32string first  = bigram.first;
-            std::u32string second = bigram.second;
-            std::vector<std::u32string> new_word;
-            int32_t i = 0;
-
-            while (i < word.size()) {
-                auto it = std::find(word.begin() + i, word.end(), first);
-                if (it == word.end()) {
-                    new_word.insert(new_word.end(), word.begin() + i, word.end());
-                    break;
-                }
-                new_word.insert(new_word.end(), word.begin() + i, it);
-                i = static_cast<int32_t>(std::distance(word.begin(), it));
-
-                if (word[i] == first && i < static_cast<int32_t>(word.size()) - 1 && word[i + 1] == second) {
-                    new_word.push_back(first + second);
-                    i += 2;
-                } else {
-                    new_word.push_back(word[i]);
-                    i += 1;
-                }
-            }
-
-            word = new_word;
-
-            if (word.size() == 1) {
-                break;
-            }
-            pairs = get_pairs(word);
-        }
-
-        std::u32string result;
-        for (int i = 0; i < word.size(); i++) {
-            result += word[i];
-            if (i != word.size() - 1) {
-                result += utf8_to_utf32(" ");
-            }
-        }
-
-        return result;
-    }
-
-    std::vector<int> tokenize(std::string text,
-                              on_new_token_cb_t on_new_token_cb,
-                              size_t max_length = 0,
-                              bool padding      = false) {
-        std::vector<int32_t> tokens = encode(text, on_new_token_cb);
-
-        tokens.insert(tokens.begin(), BOS_TOKEN_ID);
-        if (max_length > 0) {
-            if (tokens.size() > max_length - 1) {
-                tokens.resize(max_length - 1);
-                tokens.push_back(EOS_TOKEN_ID);
-            } else {
-                tokens.push_back(EOS_TOKEN_ID);
-                if (padding) {
-                    tokens.insert(tokens.end(), max_length - tokens.size(), PAD_TOKEN_ID);
-                }
-            }
-        }
-
-        return tokens;
-    }
-
-    void pad_tokens(std::vector<int>& tokens,
-                    std::vector<float>& weights,
-                    size_t max_length = 0,
-                    bool padding      = false) {
-        if (max_length > 0 && padding) {
-            size_t n = static_cast<size_t>(std::ceil(tokens.size() * 1.0 / (max_length - 2)));
-            if (n == 0) {
-                n = 1;
-            }
-            size_t length = max_length * n;
-            LOG_DEBUG("token length: %llu", length);
-            std::vector<int> new_tokens;
-            std::vector<float> new_weights;
-            new_tokens.push_back(BOS_TOKEN_ID);
-            new_weights.push_back(1.0);
-            int token_idx = 0;
-            for (int i = 1; i < length; i++) {
-                if (token_idx >= tokens.size()) {
-                    break;
-                }
-                if (i % max_length == 0) {
-                    new_tokens.push_back(BOS_TOKEN_ID);
-                    new_weights.push_back(1.0);
-                } else if (i % max_length == max_length - 1) {
-                    new_tokens.push_back(EOS_TOKEN_ID);
-                    new_weights.push_back(1.0);
-                } else {
-                    new_tokens.push_back(tokens[token_idx]);
-                    new_weights.push_back(weights[token_idx]);
-                    token_idx++;
-                }
-            }
-
-            new_tokens.push_back(EOS_TOKEN_ID);
-            new_weights.push_back(1.0);
-            tokens  = new_tokens;
-            weights = new_weights;
-
-            if (padding) {
-                tokens.insert(tokens.end(), length - tokens.size(), PAD_TOKEN_ID);
-                weights.insert(weights.end(), length - weights.size(), 1.0);
-            }
-        }
-    }
-
-    std::string clean_up_tokenization(std::string& text) {
-        std::regex pattern(R"( ,)");
-        // Replace " ," with ","
-        std::string result = std::regex_replace(text, pattern, ",");
-        return result;
-    }
-
-    std::string decode(const std::vector<int>& tokens) {
-        std::string text = "";
-        for (int t : tokens) {
-            if (t == 49406 || t == 49407)
-                continue;
-            std::u32string ts = decoder[t];
-            // printf("%d, %s \n", t,  utf32_to_utf8(ts).c_str());
-            std::string s = utf32_to_utf8(ts);
-            if (s.length() >= 4) {
-                if (ends_with(s, "</w>")) {
-                    text += s.replace(s.length() - 4, s.length() - 1, "") + " ";
-                } else {
-                    text += s;
-                }
-            } else {
-                text += " " + s;
-            }
-        }
-        // std::vector<unsigned char> bytes;
-        // for (auto c : text){
-        //     bytes.push_back(byte_decoder[c]);
-        // }
-
-        // std::string s((char *)bytes.data());
-        // std::string s = "";
-        text = clean_up_tokenization(text);
-        return trim(text);
-    }
-
-    std::vector<std::string> token_split(const std::string& text) {
-        std::regex pat(R"('s|'t|'re|'ve|'m|'ll|'d|[[:alpha:]]+|[[:digit:]]|[^[:space:][:alpha:][:digit:]]+)",
-                       std::regex::icase);
-        std::sregex_iterator iter(text.begin(), text.end(), pat);
-        std::sregex_iterator end;
-
-        std::vector<std::string> result;
-        for (; iter != end; ++iter) {
-            result.emplace_back(iter->str());
-        }
-
-        return result;
-    }
-
-    std::vector<int> encode(std::string text, on_new_token_cb_t on_new_token_cb) {
-        std::string original_text = text;
-        std::vector<int32_t> bpe_tokens;
-        text = whitespace_clean(text);
-        std::transform(text.begin(), text.end(), text.begin(), [](unsigned char c) { return std::tolower(c); });
-
-        std::string str = text;
-        std::vector<std::string> token_strs;
-
-        auto splited_texts = split_with_special_tokens(text, special_tokens);
-
-        for (auto& splited_text : splited_texts) {
-            LOG_DEBUG("token %s", splited_text.c_str());
-            if (is_special_token(splited_text)) {
-                LOG_DEBUG("special %s", splited_text.c_str());
-                bool skip = on_new_token_cb(splited_text, bpe_tokens);
-                if (skip) {
-                    token_strs.push_back(splited_text);
-                    continue;
-                }
-                continue;
-            }
-
-            auto tokens = token_split(splited_text);
-            for (auto& token : tokens) {
-                if (on_new_token_cb != nullptr) {
-                    bool skip = on_new_token_cb(token, bpe_tokens);
-                    if (skip) {
-                        token_strs.push_back(token);
-                        continue;
-                    }
-                }
-
-                std::string token_str = token;
-                std::u32string utf32_token;
-                for (int i = 0; i < token_str.length(); i++) {
-                    unsigned char b = token_str[i];
-                    utf32_token += byte_encoder[b];
-                }
-                auto bpe_strs = bpe(utf32_token);
-                size_t start  = 0;
-                size_t pos;
-                while ((pos = bpe_strs.find(' ', start)) != std::u32string::npos) {
-                    auto bpe_str = bpe_strs.substr(start, pos - start);
-                    bpe_tokens.push_back(encoder[bpe_str]);
-                    token_strs.push_back(utf32_to_utf8(bpe_str));
-
-                    start = pos + 1;
-                }
-                auto bpe_str = bpe_strs.substr(start, bpe_strs.size() - start);
-                bpe_tokens.push_back(encoder[bpe_str]);
-                token_strs.push_back(utf32_to_utf8(bpe_str));
-            }
-        }
-        // std::stringstream ss;
-        // ss << "[";
-        // for (auto token : token_strs) {
-        //     ss << "\"" << token << "\", ";
-        // }
-        // ss << "]";
-        // LOG_DEBUG("split prompt \"%s\" to tokens %s", original_text.c_str(), ss.str().c_str());
-        // printf("split prompt \"%s\" to tokens %s \n", original_text.c_str(), ss.str().c_str());
-        return bpe_tokens;
-    }
-};
 
 /*================================================ FrozenCLIPEmbedder ================================================*/
 

src/common_block.hpp

@@ -277,6 +277,7 @@ protected:
     int64_t context_dim;
     int64_t n_head;
     int64_t d_head;
+    bool xtra_dim = false;
 
 public:
     CrossAttention(int64_t query_dim,
@@ -288,7 +289,11 @@ public:
           query_dim(query_dim),
           context_dim(context_dim) {
         int64_t inner_dim = d_head * n_head;
-
+        if (context_dim == 320 && d_head == 320) {
+            // LOG_DEBUG("CrossAttention: temp set dim to 1024 for sdxs_09");
+            xtra_dim    = true;
+            context_dim = 1024;
+        }
         blocks["to_q"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, false));
         blocks["to_k"] = std::shared_ptr<GGMLBlock>(new Linear(context_dim, inner_dim, false));
         blocks["to_v"] = std::shared_ptr<GGMLBlock>(new Linear(context_dim, inner_dim, false));
@@ -314,9 +319,15 @@ public:
         int64_t inner_dim = d_head * n_head;
 
         auto q = to_q->forward(ctx, x);  // [N, n_token, inner_dim]
+        if (xtra_dim) {
+            // LOG_DEBUG("CrossAttention: temp set dim to 1024 for sdxs_09");
+            context->ne[0] = 1024;  // patch dim
+        }
         auto k = to_k->forward(ctx, context);  // [N, n_context, inner_dim]
         auto v = to_v->forward(ctx, context);  // [N, n_context, inner_dim]
-
+        if (xtra_dim) {
+            context->ne[0] = 320;  // reset dim to orig
+        }
         x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, inner_dim]
 
         x = to_out_0->forward(ctx, x);  // [N, n_token, query_dim]

src/conditioner.hpp

@@ -256,15 +256,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         return true;
     }
 
-    std::tuple<std::vector<int>, std::vector<float>, std::vector<bool>>
-    tokenize_with_trigger_token(std::string text,
-                                int num_input_imgs,
-                                int32_t image_token,
-                                bool padding = false) {
-        return tokenize_with_trigger_token(text, num_input_imgs, image_token,
-                                           text_model->model.n_token, padding);
-    }
-
     std::vector<int> convert_token_to_id(std::string text) {
         auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
             auto iter = embedding_map.find(str);
@@ -288,9 +279,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
     std::tuple<std::vector<int>, std::vector<float>, std::vector<bool>>
     tokenize_with_trigger_token(std::string text,
                                 int num_input_imgs,
-                                int32_t image_token,
+                                int32_t image_token) {
-                                size_t max_length = 0,
-                                bool padding      = false) {
         auto parsed_attention = parse_prompt_attention(text);
 
         {
@@ -377,7 +366,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         // tokens.insert(tokens.begin(), tokenizer.BOS_TOKEN_ID);
         // weights.insert(weights.begin(), 1.0);
 
-        tokenizer.pad_tokens(tokens, weights, max_length, padding);
+        tokenizer.pad_tokens(tokens, &weights, nullptr, text_model->model.n_token, text_model->model.n_token, true);
         int offset = pm_version == PM_VERSION_2 ? 2 * num_input_imgs : num_input_imgs;
         for (int i = 0; i < tokens.size(); i++) {
             // if (class_idx + 1 <= i && i < class_idx + 1 + 2*num_input_imgs) // photomaker V2 has num_tokens(=2)*num_input_imgs
@@ -403,13 +392,9 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
     }
 
     std::pair<std::vector<int>, std::vector<float>> tokenize(std::string text,
-                                                             bool padding = false) {
+                                                             size_t min_length          = 0,
-        return tokenize(text, text_model->model.n_token, padding);
-    }
-
-    std::pair<std::vector<int>, std::vector<float>> tokenize(std::string text,
                                                              size_t max_length          = 0,
-                                                             bool padding      = false) {
+                                                             bool allow_overflow_expand = true) {
         auto parsed_attention = parse_prompt_attention(text);
 
         {
@@ -460,7 +445,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
             weights.insert(weights.end(), curr_tokens.size(), curr_weight);
         }
 
-        tokenizer.pad_tokens(tokens, weights, max_length, padding);
+        tokenizer.pad_tokens(tokens, &weights, nullptr, min_length, max_length, allow_overflow_expand);
 
         // for (int i = 0; i < tokens.size(); i++) {
         //     std::cout << tokens[i] << ":" << weights[i] << ", ";
@@ -603,8 +588,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         GGML_ASSERT(image_tokens.size() == 1);
         auto tokens_and_weights     = tokenize_with_trigger_token(conditioner_params.text,
                                                                   conditioner_params.num_input_imgs,
-                                                                  image_tokens[0],
+                                                                  image_tokens[0]);
-                                                                  true);
         std::vector<int>& tokens    = std::get<0>(tokens_and_weights);
         std::vector<float>& weights = std::get<1>(tokens_and_weights);
         std::vector<bool>& clsm     = std::get<2>(tokens_and_weights);
@@ -630,7 +614,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
     std::string remove_trigger_from_prompt(const std::string& prompt) override {
         auto image_tokens = convert_token_to_id(trigger_word);
         GGML_ASSERT(image_tokens.size() == 1);
-        auto tokens_and_weights  = tokenize(prompt, false);
+        auto tokens_and_weights  = tokenize(prompt);
         std::vector<int>& tokens = tokens_and_weights.first;
         auto it                  = std::find(tokens.begin(), tokens.end(), image_tokens[0]);
         GGML_ASSERT(it != tokens.end());  // prompt must have trigger word
@@ -640,7 +624,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
 
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
-        auto tokens_and_weights     = tokenize(conditioner_params.text, true);
+        auto tokens_and_weights     = tokenize(conditioner_params.text, text_model->model.n_token, text_model->model.n_token, true);
         std::vector<int>& tokens    = tokens_and_weights.first;
         std::vector<float>& weights = tokens_and_weights.second;
         return get_learned_condition_common(n_threads,
@@ -822,8 +806,9 @@ struct SD3CLIPEmbedder : public Conditioner {
     }
 
     std::vector<std::pair<std::vector<int>, std::vector<float>>> tokenize(std::string text,
+                                                                          size_t min_length          = 0,
                                                                           size_t max_length          = 0,
-                                                                          bool padding      = false) {
+                                                                          bool allow_overflow_expand = true) {
         auto parsed_attention = parse_prompt_attention(text);
 
         {
@@ -860,20 +845,20 @@ struct SD3CLIPEmbedder : public Conditioner {
                 clip_g_weights.insert(clip_g_weights.end(), curr_tokens.size(), curr_weight);
             }
             if (t5) {
-                std::vector<int> curr_tokens = t5_tokenizer.Encode(curr_text, true);
+                std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
                 t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
                 t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
             }
         }
 
         if (clip_l) {
-            clip_l_tokenizer.pad_tokens(clip_l_tokens, clip_l_weights, max_length, padding);
+            clip_l_tokenizer.pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, min_length, max_length, allow_overflow_expand);
         }
         if (clip_g) {
-            clip_g_tokenizer.pad_tokens(clip_g_tokens, clip_g_weights, max_length, padding);
+            clip_g_tokenizer.pad_tokens(clip_g_tokens, &clip_g_weights, nullptr, min_length, max_length, allow_overflow_expand);
         }
         if (t5) {
-            t5_tokenizer.pad_tokens(t5_tokens, t5_weights, nullptr, max_length, padding);
+            t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, nullptr, min_length, max_length, true);
         }
 
         // for (int i = 0; i < clip_l_tokens.size(); i++) {
@@ -1056,7 +1041,7 @@ struct SD3CLIPEmbedder : public Conditioner {
 
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
-        auto tokens_and_weights = tokenize(conditioner_params.text, 77, true);
+        auto tokens_and_weights = tokenize(conditioner_params.text, 77, 77, true);
         return get_learned_condition_common(n_threads,
                                             tokens_and_weights,
                                             conditioner_params.clip_skip,
@@ -1158,8 +1143,8 @@ struct FluxCLIPEmbedder : public Conditioner {
     }
 
     std::vector<std::pair<std::vector<int>, std::vector<float>>> tokenize(std::string text,
-                                                                          size_t max_length = 0,
+                                                                          size_t min_length = 0,
-                                                                          bool padding      = false) {
+                                                                          size_t max_length = 0) {
         auto parsed_attention = parse_prompt_attention(text);
 
         {
@@ -1189,17 +1174,17 @@ struct FluxCLIPEmbedder : public Conditioner {
                 clip_l_weights.insert(clip_l_weights.end(), curr_tokens.size(), curr_weight);
             }
             if (t5) {
-                std::vector<int> curr_tokens = t5_tokenizer.Encode(curr_text, true);
+                std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
                 t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
                 t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
             }
         }
 
         if (clip_l) {
-            clip_l_tokenizer.pad_tokens(clip_l_tokens, clip_l_weights, 77, padding);
+            clip_l_tokenizer.pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, 77, 77, true);
         }
         if (t5) {
-            t5_tokenizer.pad_tokens(t5_tokens, t5_weights, nullptr, max_length, padding);
+            t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, nullptr, min_length, max_length, true);
         }
 
         // for (int i = 0; i < clip_l_tokens.size(); i++) {
@@ -1300,7 +1285,7 @@ struct FluxCLIPEmbedder : public Conditioner {
 
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
-        auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, true);
+        auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, chunk_len);
         return get_learned_condition_common(n_threads,
                                             tokens_and_weights,
                                             conditioner_params.clip_skip,
@@ -1377,8 +1362,8 @@ struct T5CLIPEmbedder : public Conditioner {
     }
 
     std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
-                                                                                  size_t max_length = 0,
+                                                                                  size_t min_length = 0,
-                                                                                  bool padding      = false) {
+                                                                                  size_t max_length = 0) {
         auto parsed_attention = parse_prompt_attention(text);
 
         {
@@ -1403,12 +1388,15 @@ struct T5CLIPEmbedder : public Conditioner {
                 const std::string& curr_text = item.first;
                 float curr_weight            = item.second;
 
-                std::vector<int> curr_tokens = t5_tokenizer.Encode(curr_text, true);
+                std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
                 t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
                 t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
             }
 
-            t5_tokenizer.pad_tokens(t5_tokens, t5_weights, &t5_mask, max_length, padding);
+            t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, &t5_mask, min_length, max_length, true);
+            for (auto& mask_value : t5_mask) {
+                mask_value = mask_value > 0.0f ? 0.0f : -HUGE_VALF;
+            }
         }
         return {t5_tokens, t5_weights, t5_mask};
     }
@@ -1496,7 +1484,7 @@ struct T5CLIPEmbedder : public Conditioner {
 
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
-        auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, true);
+        auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, chunk_len);
         return get_learned_condition_common(n_threads,
                                             tokens_and_weights,
                                             conditioner_params.clip_skip,
@@ -1505,14 +1493,14 @@ struct T5CLIPEmbedder : public Conditioner {
 };
 
 struct AnimaConditioner : public Conditioner {
-    std::shared_ptr<LLM::BPETokenizer> qwen_tokenizer;
+    std::shared_ptr<BPETokenizer> qwen_tokenizer;
     T5UniGramTokenizer t5_tokenizer;
     std::shared_ptr<LLM::LLMRunner> llm;
 
     AnimaConditioner(ggml_backend_t backend,
                      bool offload_params_to_cpu,
                      const String2TensorStorage& tensor_storage_map = {}) {
-        qwen_tokenizer = std::make_shared<LLM::Qwen2Tokenizer>();
+        qwen_tokenizer = std::make_shared<Qwen2Tokenizer>();
         llm            = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::QWEN3,
                                                backend,
                                                offload_params_to_cpu,
@@ -1578,7 +1566,7 @@ struct AnimaConditioner : public Conditioner {
         for (const auto& item : parsed_attention) {
             const std::string& curr_text = item.first;
             float curr_weight            = item.second;
-            std::vector<int> curr_tokens = t5_tokenizer.Encode(curr_text, true);
+            std::vector<int> curr_tokens = t5_tokenizer.tokenize(curr_text, nullptr, true);
             t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
             t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
         }
@@ -1620,7 +1608,7 @@ struct AnimaConditioner : public Conditioner {
 
 struct LLMEmbedder : public Conditioner {
     SDVersion version;
-    std::shared_ptr<LLM::BPETokenizer> tokenizer;
+    std::shared_ptr<BPETokenizer> tokenizer;
     std::shared_ptr<LLM::LLMRunner> llm;
 
     LLMEmbedder(ggml_backend_t backend,
@@ -1633,13 +1621,15 @@ struct LLMEmbedder : public Conditioner {
         LLM::LLMArch arch = LLM::LLMArch::QWEN2_5_VL;
         if (version == VERSION_FLUX2) {
             arch = LLM::LLMArch::MISTRAL_SMALL_3_2;
+        } else if (sd_version_is_ernie_image(version)) {
+            arch = LLM::LLMArch::MINISTRAL_3_3B;
         } else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
             arch = LLM::LLMArch::QWEN3;
         }
-        if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2) {
+        if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
-            tokenizer = std::make_shared<LLM::MistralTokenizer>();
+            tokenizer = std::make_shared<MistralTokenizer>();
         } else {
-            tokenizer = std::make_shared<LLM::Qwen2Tokenizer>();
+            tokenizer = std::make_shared<Qwen2Tokenizer>();
         }
         llm = std::make_shared<LLM::LLMRunner>(arch,
                                                backend,
@@ -1677,20 +1667,24 @@ struct LLMEmbedder : public Conditioner {
         }
     }
 
-    std::tuple<std::vector<int>, std::vector<float>> tokenize(std::string text,
+    std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
                                                                                   const std::pair<int, int>& attn_range,
-                                                              size_t max_length = 0,
+                                                                                  size_t min_length = 0,
-                                                              bool padding      = false) {
+                                                                                  size_t max_length = 100000000) {
         std::vector<std::pair<std::string, float>> parsed_attention;
         if (attn_range.first >= 0 && attn_range.second > 0) {
+            if (attn_range.first > 0) {
                 parsed_attention.emplace_back(text.substr(0, attn_range.first), 1.f);
+            }
             if (attn_range.second - attn_range.first > 0) {
                 auto new_parsed_attention = parse_prompt_attention(text.substr(attn_range.first, attn_range.second - attn_range.first));
                 parsed_attention.insert(parsed_attention.end(),
                                         new_parsed_attention.begin(),
                                         new_parsed_attention.end());
             }
+            if (attn_range.second < text.size()) {
                 parsed_attention.emplace_back(text.substr(attn_range.second), 1.f);
+            }
         } else {
             parsed_attention.emplace_back(text, 1.f);
         }
@@ -1710,39 +1704,34 @@ struct LLMEmbedder : public Conditioner {
         for (const auto& item : parsed_attention) {
             const std::string& curr_text = item.first;
             float curr_weight            = item.second;
-            std::vector<int> curr_tokens = tokenizer->tokenize(curr_text, nullptr);
+            std::vector<int> curr_tokens = tokenizer->encode(curr_text, nullptr);
             tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
             weights.insert(weights.end(), curr_tokens.size(), curr_weight);
         }
 
-        tokenizer->pad_tokens(tokens, weights, max_length, padding);
+        std::vector<float> mask;
+        tokenizer->pad_tokens(tokens, &weights, &mask, min_length, max_length);
 
         // for (int i = 0; i < tokens.size(); i++) {
         //     std::cout << tokens[i] << ":" << weights[i] << ", " << i << std::endl;
         // }
         // std::cout << std::endl;
 
-        return {tokens, weights};
+        return {tokens, weights, mask};
     }
 
     sd::Tensor<float> encode_prompt(int n_threads,
                                     const std::string prompt,
                                     const std::pair<int, int>& prompt_attn_range,
-                                    int max_length,
                                     int min_length,
+                                    int hidden_states_min_length,
                                     const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
                                     const std::set<int>& out_layers,
                                     int prompt_template_encode_start_idx) {
-        auto tokens_and_weights = tokenize(prompt, prompt_attn_range);
+        auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length);
-        auto& tokens            = std::get<0>(tokens_and_weights);
+        auto& tokens             = std::get<0>(tokens_weights_mask);
-        auto& weights           = std::get<1>(tokens_and_weights);
+        auto& weights            = std::get<1>(tokens_weights_mask);
-        std::vector<float> mask;
+        auto& mask               = std::get<2>(tokens_weights_mask);
-
-        if (max_length > 0 && tokens.size() < max_length) {
-            mask.insert(mask.end(), tokens.size(), 1.f);
-            mask.insert(mask.end(), max_length - tokens.size(), 0.f);
-            tokenizer->pad_tokens(tokens, weights, max_length, true);
-        }
 
         sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
         sd::Tensor<float> attention_mask;
@@ -1769,9 +1758,9 @@ struct LLMEmbedder : public Conditioner {
         GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx);
 
         int64_t zero_pad_len = 0;
-        if (min_length > 0) {
+        if (hidden_states_min_length > 0) {
-            if (hidden_states.shape()[1] - prompt_template_encode_start_idx < min_length) {
+            if (hidden_states.shape()[1] - prompt_template_encode_start_idx < hidden_states_min_length) {
-                zero_pad_len = min_length - hidden_states.shape()[1] + prompt_template_encode_start_idx;
+                zero_pad_len = hidden_states_min_length - hidden_states.shape()[1] + prompt_template_encode_start_idx;
             }
         }
 
@@ -1798,8 +1787,8 @@ struct LLMEmbedder : public Conditioner {
         std::vector<std::pair<int, int>> extra_prompts_attn_range;
         std::vector<std::pair<int, sd::Tensor<float>>> image_embeds;
         int prompt_template_encode_start_idx = 34;
-        int max_length                       = 0;  // pad tokens
+        int min_length                       = 0;  // pad tokens
-        int min_length                       = 0;  // zero pad hidden_states
+        int hidden_states_min_length         = 0;  // zero pad hidden_states
         std::set<int> out_layers;
 
         int64_t t0 = ggml_time_ms();
@@ -1874,7 +1863,7 @@ struct LLMEmbedder : public Conditioner {
             }
         } else if (version == VERSION_FLUX2) {
             prompt_template_encode_start_idx = 0;
-            min_length                       = 512;
+            hidden_states_min_length         = 512;
             out_layers                       = {10, 20, 30};
 
             prompt = "[SYSTEM_PROMPT]You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object\nattribution and actions without speculation.[/SYSTEM_PROMPT][INST]";
@@ -1884,6 +1873,13 @@ struct LLMEmbedder : public Conditioner {
             prompt_attn_range.second = static_cast<int>(prompt.size());
 
             prompt += "[/INST]";
+        } else if (sd_version_is_ernie_image(version)) {
+            prompt_template_encode_start_idx = 0;
+            out_layers                       = {25};  // -2
+
+            prompt_attn_range.first = 0;
+            prompt += conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
         } else if (sd_version_is_z_image(version)) {
             prompt_template_encode_start_idx = 0;
             out_layers                       = {35};  // -2
@@ -1907,7 +1903,7 @@ struct LLMEmbedder : public Conditioner {
             }
         } else if (version == VERSION_FLUX2_KLEIN) {
             prompt_template_encode_start_idx = 0;
-            max_length                       = 512;
+            min_length                       = 512;
             out_layers                       = {9, 18, 27};
 
             prompt = "<|im_start|>user\n";
@@ -1919,7 +1915,7 @@ struct LLMEmbedder : public Conditioner {
             prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
         } else if (version == VERSION_OVIS_IMAGE) {
             prompt_template_encode_start_idx = 28;
-            max_length                       = prompt_template_encode_start_idx + 256;
+            min_length                       = prompt_template_encode_start_idx + 256;
 
             prompt = "<|im_start|>user\nDescribe the image by detailing the color, quantity, text, shape, size, texture, spatial relationships of the objects and background:";
 
@@ -1935,8 +1931,8 @@ struct LLMEmbedder : public Conditioner {
         auto hidden_states = encode_prompt(n_threads,
                                            prompt,
                                            prompt_attn_range,
-                                           max_length,
                                            min_length,
+                                           hidden_states_min_length,
                                            image_embeds,
                                            out_layers,
                                            prompt_template_encode_start_idx);
@@ -1945,8 +1941,8 @@ struct LLMEmbedder : public Conditioner {
             auto extra_hidden_states = encode_prompt(n_threads,
                                                      extra_prompts[i],
                                                      extra_prompts_attn_range[i],
-                                                     max_length,
                                                      min_length,
+                                                     hidden_states_min_length,
                                                      image_embeds,
                                                      out_layers,
                                                      prompt_template_encode_start_idx);

src/convert.cpp

@@ -0,0 +1,138 @@
+#include <cstring>
+#include <mutex>
+#include <regex>
+#include <vector>
+
+#include "model.h"
+#include "model_io/gguf_io.h"
+#include "model_io/safetensors_io.h"
+#include "util.h"
+
+#include "ggml-cpu.h"
+
+static ggml_type get_export_tensor_type(ModelLoader& model_loader,
+                                        const TensorStorage& tensor_storage,
+                                        ggml_type type,
+                                        const TensorTypeRules& tensor_type_rules) {
+    const std::string& name = tensor_storage.name;
+    ggml_type tensor_type   = tensor_storage.type;
+    ggml_type dst_type      = type;
+
+    for (const auto& tensor_type_rule : tensor_type_rules) {
+        std::regex pattern(tensor_type_rule.first);
+        if (std::regex_search(name, pattern)) {
+            dst_type = tensor_type_rule.second;
+            break;
+        }
+    }
+
+    if (model_loader.tensor_should_be_converted(tensor_storage, dst_type)) {
+        tensor_type = dst_type;
+    }
+
+    return tensor_type;
+}
+
+static bool load_tensors_for_export(ModelLoader& model_loader,
+                                    ggml_context* ggml_ctx,
+                                    ggml_type type,
+                                    const TensorTypeRules& tensor_type_rules,
+                                    std::vector<TensorWriteInfo>& tensors) {
+    std::mutex tensor_mutex;
+    auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
+        const std::string& name = tensor_storage.name;
+        ggml_type tensor_type   = get_export_tensor_type(model_loader, tensor_storage, type, tensor_type_rules);
+
+        std::lock_guard<std::mutex> lock(tensor_mutex);
+        ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, tensor_type, tensor_storage.n_dims, tensor_storage.ne);
+        if (tensor == nullptr) {
+            LOG_ERROR("ggml_new_tensor failed");
+            return false;
+        }
+        ggml_set_name(tensor, name.c_str());
+
+        if (!tensor->data) {
+            GGML_ASSERT(ggml_nelements(tensor) == 0);
+            // Avoid crashing writers by setting a dummy pointer for zero-sized tensors.
+            LOG_DEBUG("setting dummy pointer for zero-sized tensor %s", name.c_str());
+            tensor->data = ggml_get_mem_buffer(ggml_ctx);
+        }
+
+        TensorWriteInfo write_info;
+        write_info.tensor = tensor;
+        write_info.n_dims = tensor_storage.n_dims;
+        for (int i = 0; i < tensor_storage.n_dims; ++i) {
+            write_info.ne[i] = tensor_storage.ne[i];
+        }
+
+        *dst_tensor = tensor;
+        tensors.push_back(std::move(write_info));
+
+        return true;
+    };
+
+    bool success = model_loader.load_tensors(on_new_tensor_cb);
+    LOG_INFO("load tensors done");
+    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);
+        return false;
+    }
+
+    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();
+    }
+
+    ggml_type type             = (ggml_type)output_type;
+    bool output_is_safetensors = ends_with(output_path, ".safetensors");
+    TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
+
+    auto backend    = ggml_backend_cpu_init();
+    size_t mem_size = 1 * 1024 * 1024;  // for padding
+    mem_size += model_loader.get_tensor_storage_map().size() * ggml_tensor_overhead();
+    mem_size += model_loader.get_params_mem_size(backend, type);
+    LOG_INFO("model tensors mem size: %.2fMB", mem_size / 1024.f / 1024.f);
+    ggml_context* ggml_ctx = ggml_init({mem_size, nullptr, false});
+
+    if (ggml_ctx == nullptr) {
+        LOG_ERROR("ggml_init failed for converter");
+        ggml_backend_free(backend);
+        return false;
+    }
+
+    std::vector<TensorWriteInfo> tensors;
+    bool success = load_tensors_for_export(model_loader, ggml_ctx, type, type_rules, tensors);
+    ggml_backend_free(backend);
+
+    std::string error;
+    if (success) {
+        if (output_is_safetensors) {
+            success = write_safetensors_file(output_path, tensors, &error);
+        } else {
+            success = write_gguf_file(output_path, tensors, &error);
+        }
+    }
+
+    if (!success && !error.empty()) {
+        LOG_ERROR("%s", error.c_str());
+    }
+
+    ggml_free(ggml_ctx);
+    return success;
+}

src/denoiser.hpp

@@ -953,8 +953,9 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
             float t_next         = t_fn(sigma_down);
             float h              = t_next - t;
             float s              = t + 0.5f * h;
-            sd::Tensor<float> x2 = (sigma_fn(s) / sigma_fn(t)) * x - (exp(-h * 0.5f) - 1) * denoised;
+            float sigma_s        = sigma_fn(s);
-            auto denoised2_opt   = model(x2, sigmas[i + 1], i + 1);
+            sd::Tensor<float> x2 = (sigma_s / sigma_fn(t)) * x - (exp(-h * 0.5f) - 1) * denoised;
+            auto denoised2_opt   = model(x2, sigma_s, i + 1);
             if (denoised2_opt.empty()) {
                 return {};
             }
@@ -969,6 +970,100 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
     return x;
 }
 
+static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
+                                                        sd::Tensor<float> x,
+                                                        const std::vector<float>& sigmas,
+                                                        std::shared_ptr<RNG> rng,
+                                                        float eta = 1.0f) {
+    int steps = static_cast<int>(sigmas.size()) - 1;
+    for (int i = 0; i < steps; i++) {
+        float sigma    = sigmas[i];
+        float sigma_to = sigmas[i + 1];
+
+        bool opt_first_step = (1.0 - sigma < 1e-6);
+
+        auto denoised_opt = model(x, sigma, (opt_first_step ? 1 : -1) * (i + 1));
+        if (denoised_opt.empty()) {
+            return {};
+        }
+        sd::Tensor<float> denoised = std::move(denoised_opt);
+
+        if (sigma_to == 0.0f) {
+            // Euler method (final step, no noise)
+            // sigma_to == 0 --> sigma_down = 0, so:
+            //   x +                       d  * (sigma_down - sigma)
+            // = x + ((x - denoised) / sigma) * (sigma_down - sigma)
+            // = x + ((x - denoised) / sigma) * (         0 - sigma)
+            // = x + ((x - denoised)        ) * -1
+            // = x    -x + denoised
+            x = denoised;
+
+        } else {
+            auto [sigma_down, sigma_up, alpha_scale] = get_ancestral_step_flow(sigma, sigma_to, eta);
+            sd::Tensor<float> D_i;
+
+            if (opt_first_step) {
+                // the reformulated exp_s calc already accounts for this, but we can avoid
+                // a redundant model call for the typical sigma 1 at the first step:
+                // exp_s = sqrt((1-sigma)/sigma * (1-sigma_down)/sigma_down)
+                //       = sqrt((1-    1)/    1 * (1-sigma_down)/sigma_down)
+                //       = 0
+                // so sigma_s = 1 = sigma, and sigma_s_i_ratio = sigma_s / sigma = 1
+                // u = (x*sigma_s_i_ratio)+(denoised*(1.0f-sigma_s_i_ratio))
+                //   = (x*1)+(denoised*0) = x
+                // so D_i = model(u, sigma_s, i + 1)
+                //        = model(x, sigma,   i + 1)
+                //        = denoised
+                D_i = denoised;
+
+            } else {
+                float sigma_s;
+
+                // ref implementation would be:
+                // auto lambda_fn = [](float sigma) -> float {
+                //     return std::log((1.0f - sigma) / sigma); };
+                // auto sigma_fn = [](float lbda) -> float {
+                //     return 1.0f / (std::exp(lbda) + 1.0f); };
+                // t_i     = lambda_fn(sigma);
+                // t_down  = lambda_fn(sigma_down);
+                // float r = 0.5f;
+                // h       = t_down - t_i;
+                // s       = t_i + r * h;
+                // sigma_s = sigma_fn(s);
+
+                // assuming r is constant, we sidestep the singularity at sigma -> 1 by:
+                // s   = 0.5 * (lambda_fn(sigma)     + lambda_fn(sigma_down))
+                //     = 0.5 * (log((1-sigma)/sigma) + log((1-sigma_down)/sigma_down))
+                //     = 0.5 * log(((1-sigma)/sigma) *    ((1-sigma_down)/sigma_down))
+                //     = log(sqrt (((1-sigma)/sigma) *    ((1-sigma_down)/sigma_down)))
+                // so exp(s) = sqrt((1-sigma)/sigma  *     (1-sigma_down)/sigma_down)
+                // and sigma_s = sigma_fn(s) = 1.0f / (exp(s) + 1.0f)
+
+                float exp_s = std::sqrt(((1 - sigma) / sigma) * ((1 - sigma_down) / sigma_down));
+                sigma_s     = 1.0f / (exp_s + 1.0f);
+
+                float sigma_s_i_ratio = sigma_s / sigma;
+                sd::Tensor<float> u   = (x * sigma_s_i_ratio) + (denoised * (1.0f - sigma_s_i_ratio));
+
+                auto denoised2_opt = model(u, sigma_s, i + 1);
+                if (denoised2_opt.empty()) {
+                    return {};
+                }
+                D_i = std::move(denoised2_opt);
+            }
+
+            float sigma_down_i_ratio = sigma_down / sigma;
+            x                        = (x * sigma_down_i_ratio) + (D_i * (1.0f - sigma_down_i_ratio));
+
+            if (sigma_to > 0.0f && eta > 0.0f) {
+                x = alpha_scale * x + sd::Tensor<float>::randn_like(x, rng) * sigma_up;
+            }
+        }
+    }
+
+    return x;
+}
+
 static sd::Tensor<float> sample_dpmpp_2m(denoise_cb_t model,
                                          sd::Tensor<float> x,
                                          const std::vector<float>& sigmas) {
@@ -1040,7 +1135,8 @@ static sd::Tensor<float> sample_dpmpp_2m_v2(denoise_cb_t model,
 static sd::Tensor<float> sample_lcm(denoise_cb_t model,
                                     sd::Tensor<float> x,
                                     const std::vector<float>& sigmas,
-                                    std::shared_ptr<RNG> rng) {
+                                    std::shared_ptr<RNG> rng,
+                                    bool is_flow_denoiser) {
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
         auto denoised_opt = model(x, sigmas[i], i + 1);
@@ -1049,6 +1145,9 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
         }
         x = std::move(denoised_opt);
         if (sigmas[i + 1] > 0) {
+            if (is_flow_denoiser) {
+                x *= (1 - sigmas[i + 1]);
+            }
             x += sd::Tensor<float>::randn_like(x, rng) * sigmas[i + 1];
         }
     }
@@ -1285,37 +1384,149 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
     return x;
 }
 
+static sd::Tensor<float> sample_er_sde(denoise_cb_t model,
+                                       sd::Tensor<float> x,
+                                       std::vector<float> sigmas,
+                                       std::shared_ptr<RNG> rng,
+                                       bool is_flow_denoiser,
+                                       float eta) {
+    constexpr int max_stage                  = 3;
+    constexpr int num_integration_points     = 200;
+    constexpr float num_integration_points_f = 200.0f;
+    float s_noise                            = eta;
+
+    auto er_sde_flow_sigma = [](float sigma) -> float {
+        sigma = std::max(sigma, 1e-6f);
+        sigma = std::min(sigma, 1.0f - 1e-4f);
+        return sigma;
+    };
+
+    auto sigma_to_er_sde_lambda = [&](float sigma, bool is_flow_denoiser) -> float {
+        if (is_flow_denoiser) {
+            sigma = er_sde_flow_sigma(sigma);
+            return sigma / std::max(1.0f - sigma, 1e-6f);
+        }
+        return std::max(sigma, 1e-6f);
+    };
+
+    auto sigma_to_er_sde_alpha = [&](float sigma, bool is_flow_denoiser) -> float {
+        if (is_flow_denoiser) {
+            sigma = er_sde_flow_sigma(sigma);
+            return 1.0f - sigma;
+        }
+        return 1.0f;
+    };
+
+    auto er_sde_noise_scaler = [](float x) -> float {
+        x = std::max(x, 0.0f);
+        return x * (std::exp(std::pow(x, 0.3f)) + 10.0f);
+    };
+
+    if (is_flow_denoiser) {
+        for (size_t i = 0; i + 1 < sigmas.size(); ++i) {
+            if (sigmas[i] > 1.0f) {
+                sigmas[i] = er_sde_flow_sigma(sigmas[i]);
+            }
+        }
+    }
+
+    std::vector<float> er_lambdas(sigmas.size(), 0.0f);
+    for (size_t i = 0; i < sigmas.size(); ++i) {
+        er_lambdas[i] = sigma_to_er_sde_lambda(sigmas[i], is_flow_denoiser);
+    }
+
+    sd::Tensor<float> old_denoised   = x;
+    sd::Tensor<float> old_denoised_d = x;
+    bool have_old_denoised           = false;
+    bool have_old_denoised_d         = false;
+
+    int steps = static_cast<int>(sigmas.size()) - 1;
+    for (int i = 0; i < steps; i++) {
+        sd::Tensor<float> denoised = model(x, sigmas[i], i + 1);
+        if (denoised.empty()) {
+            return {};
+        }
+
+        int stage_used = std::min(max_stage, i + 1);
+
+        if (sigmas[i + 1] == 0.0f) {
+            x = denoised;
+        } else {
+            float er_lambda_s = er_lambdas[i];
+            float er_lambda_t = er_lambdas[i + 1];
+            float alpha_s     = sigma_to_er_sde_alpha(sigmas[i], is_flow_denoiser);
+            float alpha_t     = sigma_to_er_sde_alpha(sigmas[i + 1], is_flow_denoiser);
+            float scaled_s    = er_sde_noise_scaler(er_lambda_s);
+            float scaled_t    = er_sde_noise_scaler(er_lambda_t);
+            float r_alpha     = alpha_s > 0.0f ? alpha_t / alpha_s : 0.0f;
+            float r           = scaled_s > 0.0f ? scaled_t / scaled_s : 0.0f;
+
+            x = r_alpha * r * x + alpha_t * (1.0f - r) * denoised;
+
+            if (stage_used >= 2 && have_old_denoised) {
+                float dt               = er_lambda_t - er_lambda_s;
+                float lambda_step_size = -dt / num_integration_points_f;
+                float s                = 0.0f;
+                float s_u              = 0.0f;
+
+                for (int p = 0; p < num_integration_points; ++p) {
+                    float lambda_pos = er_lambda_t + p * lambda_step_size;
+                    float scaled_pos = er_sde_noise_scaler(lambda_pos);
+                    if (scaled_pos <= 0.0f) {
+                        continue;
+                    }
+
+                    s += 1.0f / scaled_pos;
+                    if (stage_used >= 3 && have_old_denoised_d) {
+                        s_u += (lambda_pos - er_lambda_s) / scaled_pos;
+                    }
+                }
+
+                s *= lambda_step_size;
+
+                float denom_d = er_lambda_s - er_lambdas[i - 1];
+                if (std::fabs(denom_d) > 1e-12f) {
+                    float coeff_d                = alpha_t * (dt + s * scaled_t);
+                    sd::Tensor<float> denoised_d = (denoised - old_denoised) / denom_d;
+                    x += coeff_d * denoised_d;
+
+                    if (stage_used >= 3 && have_old_denoised_d) {
+                        float denom_u = (er_lambda_s - er_lambdas[i - 2]) * 0.5f;
+                        if (std::fabs(denom_u) > 1e-12f) {
+                            s_u *= lambda_step_size;
+                            float coeff_u                = alpha_t * (0.5f * dt * dt + s_u * scaled_t);
+                            sd::Tensor<float> denoised_u = (denoised_d - old_denoised_d) / denom_u;
+                            x += coeff_u * denoised_u;
+                        }
+                    }
+
+                    old_denoised_d      = denoised_d;
+                    have_old_denoised_d = true;
+                }
+            }
+
+            float noise_scale_sq = er_lambda_t * er_lambda_t - er_lambda_s * er_lambda_s * r * r;
+            if (s_noise > 0.0f && noise_scale_sq > 0.0f) {
+                float noise_scale = alpha_t * std::sqrt(std::max(noise_scale_sq, 0.0f));
+                x += sd::Tensor<float>::randn_like(x, rng) * noise_scale;
+            }
+        }
+
+        old_denoised      = denoised;
+        have_old_denoised = true;
+    }
+    return x;
+}
+
 static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
                                               sd::Tensor<float> x,
                                               const std::vector<float>& sigmas,
                                               std::shared_ptr<RNG> rng,
                                               float eta) {
-    float beta_start = 0.00085f;
-    float beta_end   = 0.0120f;
-    std::vector<double> alphas_cumprod(TIMESTEPS);
-    std::vector<double> compvis_sigmas(TIMESTEPS);
-    for (int i = 0; i < TIMESTEPS; i++) {
-        alphas_cumprod[i] =
-            (i == 0 ? 1.0f : alphas_cumprod[i - 1]) *
-            (1.0f -
-             std::pow(sqrtf(beta_start) +
-                          (sqrtf(beta_end) - sqrtf(beta_start)) *
-                              ((float)i / (TIMESTEPS - 1)),
-                      2));
-        compvis_sigmas[i] =
-            std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]);
-    }
-
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        int timestep      = static_cast<int>(roundf(TIMESTEPS - i * ((float)TIMESTEPS / steps))) - 1;
+        float sigma    = sigmas[i];
-        int prev_timestep = timestep - TIMESTEPS / steps;
+        float sigma_to = sigmas[i + 1];
-        float sigma       = static_cast<float>(compvis_sigmas[timestep]);
-        if (i == 0) {
-            x *= std::sqrt(sigma * sigma + 1) / sigma;
-        } else {
-            x *= std::sqrt(sigma * sigma + 1);
-        }
 
         auto model_output_opt = model(x, sigma, i + 1);
         if (model_output_opt.empty()) {
@@ -1324,8 +1535,8 @@ static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
         sd::Tensor<float> model_output = std::move(model_output_opt);
         model_output                   = (x - model_output) * (1.0f / sigma);
 
-        float alpha_prod_t      = static_cast<float>(alphas_cumprod[timestep]);
+        float alpha_prod_t      = 1.0f / (sigma * sigma + 1.0f);
-        float alpha_prod_t_prev = static_cast<float>(prev_timestep >= 0 ? alphas_cumprod[prev_timestep] : alphas_cumprod[0]);
+        float alpha_prod_t_prev = 1.0f / (sigma_to * sigma_to + 1.0f);
         float beta_prod_t       = 1.0f - alpha_prod_t;
 
         sd::Tensor<float> pred_original_sample = ((x / std::sqrt(sigma * sigma + 1)) -
@@ -1337,11 +1548,11 @@ static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
                          (1.0f - alpha_prod_t / alpha_prod_t_prev);
         float std_dev_t = eta * std::sqrt(variance);
 
-        x = std::sqrt(alpha_prod_t_prev) * pred_original_sample +
+        x = pred_original_sample +
-            std::sqrt(1.0f - alpha_prod_t_prev - std::pow(std_dev_t, 2)) * model_output;
+            std::sqrt((1.0f - alpha_prod_t_prev - std::pow(std_dev_t, 2)) / alpha_prod_t_prev) * model_output;
 
         if (eta > 0) {
-            x += std_dev_t * sd::Tensor<float>::randn_like(x, rng);
+            x += std_dev_t / std::sqrt(alpha_prod_t_prev) * sd::Tensor<float>::randn_like(x, rng);
         }
     }
     return x;
@@ -1368,19 +1579,26 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
             std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]);
     }
 
-    int original_steps = 50;
+    auto get_timestep_from_sigma = [&](float s) -> int {
+        auto it = std::lower_bound(compvis_sigmas.begin(), compvis_sigmas.end(), s);
+        if (it == compvis_sigmas.begin())
+            return 0;
+        if (it == compvis_sigmas.end())
+            return TIMESTEPS - 1;
+        int idx_high = static_cast<int>(std::distance(compvis_sigmas.begin(), it));
+        int idx_low  = idx_high - 1;
+        if (std::abs(compvis_sigmas[idx_high] - s) < std::abs(compvis_sigmas[idx_low] - s)) {
+            return idx_high;
+        }
+        return idx_low;
+    };
+
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        int timestep      = TIMESTEPS - 1 - (TIMESTEPS / original_steps) * (int)floor(i * ((float)original_steps / steps));
+        float sigma_to    = sigmas[i + 1];
-        int prev_timestep = i >= steps - 1 ? 0 : TIMESTEPS - 1 - (TIMESTEPS / original_steps) * (int)floor((i + 1) * ((float)original_steps / steps));
+        int prev_timestep = get_timestep_from_sigma(sigma_to);
         int timestep_s    = (int)floor((1 - eta) * prev_timestep);
-        float sigma       = static_cast<float>(compvis_sigmas[timestep]);
+        float sigma       = sigmas[i];
-
-        if (i == 0) {
-            x *= std::sqrt(sigma * sigma + 1) / sigma;
-        } else {
-            x *= std::sqrt(sigma * sigma + 1);
-        }
 
         auto model_output_opt = model(x, sigma, i + 1);
         if (model_output_opt.empty()) {
@@ -1389,9 +1607,9 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
         sd::Tensor<float> model_output = std::move(model_output_opt);
         model_output                   = (x - model_output) * (1.0f / sigma);
 
-        float alpha_prod_t      = static_cast<float>(alphas_cumprod[timestep]);
+        float alpha_prod_t      = 1.0f / (sigma * sigma + 1.0f);
         float beta_prod_t       = 1.0f - alpha_prod_t;
-        float alpha_prod_t_prev = static_cast<float>(prev_timestep >= 0 ? alphas_cumprod[prev_timestep] : alphas_cumprod[0]);
+        float alpha_prod_t_prev = 1.0f / (sigma_to * sigma_to + 1.0f);
         float alpha_prod_s      = static_cast<float>(alphas_cumprod[timestep_s]);
         float beta_prod_s       = 1.0f - alpha_prod_s;
 
@@ -1399,12 +1617,12 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
                                                   std::sqrt(beta_prod_t) * model_output) *
                                                  (1.0f / std::sqrt(alpha_prod_t));
 
-        x = std::sqrt(alpha_prod_s) * pred_original_sample +
+        x = std::sqrt(alpha_prod_s / alpha_prod_t_prev) * pred_original_sample +
-            std::sqrt(beta_prod_s) * model_output;
+            std::sqrt(beta_prod_s / alpha_prod_t_prev) * model_output;
 
-        if (eta > 0 && i != steps - 1) {
+        if (eta > 0 && sigma_to > 0.0f) {
             x = std::sqrt(alpha_prod_t_prev / alpha_prod_s) * x +
-                std::sqrt(1.0f - alpha_prod_t_prev / alpha_prod_s) * sd::Tensor<float>::randn_like(x, rng);
+                std::sqrt(1.0f / alpha_prod_t_prev - 1.0f / alpha_prod_s) * sd::Tensor<float>::randn_like(x, rng);
         }
     }
     return x;
@@ -1431,13 +1649,16 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
         case DPM2_SAMPLE_METHOD:
             return sample_dpm2(model, std::move(x), sigmas);
         case DPMPP2S_A_SAMPLE_METHOD:
+            if (is_flow_denoiser)
+                return sample_dpmpp_2s_ancestral_flow(model, std::move(x), sigmas, rng, eta);
+            else
                 return sample_dpmpp_2s_ancestral(model, std::move(x), sigmas, rng, eta);
         case DPMPP2M_SAMPLE_METHOD:
             return sample_dpmpp_2m(model, std::move(x), sigmas);
         case DPMPP2Mv2_SAMPLE_METHOD:
             return sample_dpmpp_2m_v2(model, std::move(x), sigmas);
         case LCM_SAMPLE_METHOD:
-            return sample_lcm(model, std::move(x), sigmas, rng);
+            return sample_lcm(model, std::move(x), sigmas, rng, is_flow_denoiser);
         case IPNDM_SAMPLE_METHOD:
             return sample_ipndm(model, std::move(x), sigmas);
         case IPNDM_V_SAMPLE_METHOD:
@@ -1446,6 +1667,8 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
             return sample_res_multistep(model, std::move(x), sigmas, rng, eta);
         case RES_2S_SAMPLE_METHOD:
             return sample_res_2s(model, std::move(x), sigmas, rng, eta);
+        case ER_SDE_SAMPLE_METHOD:
+            return sample_er_sde(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
         case DDIM_TRAILING_SAMPLE_METHOD:
             return sample_ddim_trailing(model, std::move(x), sigmas, rng, eta);
         case TCD_SAMPLE_METHOD:

src/diffusion_model.hpp

@@ -3,6 +3,7 @@
 
 #include <optional>
 #include "anima.hpp"
+#include "ernie_image.hpp"
 #include "flux.hpp"
 #include "mmdit.hpp"
 #include "qwen_image.hpp"
@@ -516,4 +517,66 @@ struct ZImageModel : public DiffusionModel {
     }
 };
 
+struct ErnieImageModel : public DiffusionModel {
+    std::string prefix;
+    ErnieImage::ErnieImageRunner ernie_image;
+
+    ErnieImageModel(ggml_backend_t backend,
+                    bool offload_params_to_cpu,
+                    const String2TensorStorage& tensor_storage_map = {},
+                    const std::string prefix                       = "model.diffusion_model")
+        : prefix(prefix), ernie_image(backend, offload_params_to_cpu, tensor_storage_map, prefix) {
+    }
+
+    std::string get_desc() override {
+        return ernie_image.get_desc();
+    }
+
+    void alloc_params_buffer() override {
+        ernie_image.alloc_params_buffer();
+    }
+
+    void free_params_buffer() override {
+        ernie_image.free_params_buffer();
+    }
+
+    void free_compute_buffer() override {
+        ernie_image.free_compute_buffer();
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        ernie_image.get_param_tensors(tensors, prefix);
+    }
+
+    size_t get_params_buffer_size() override {
+        return ernie_image.get_params_buffer_size();
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        ernie_image.set_weight_adapter(adapter);
+    }
+
+    int64_t get_adm_in_channels() override {
+        return 768;
+    }
+
+    void set_flash_attention_enabled(bool enabled) {
+        ernie_image.set_flash_attention_enabled(enabled);
+    }
+
+    void set_circular_axes(bool circular_x, bool circular_y) override {
+        ernie_image.set_circular_axes(circular_x, circular_y);
+    }
+
+    sd::Tensor<float> compute(int n_threads,
+                              const DiffusionParams& diffusion_params) override {
+        GGML_ASSERT(diffusion_params.x != nullptr);
+        GGML_ASSERT(diffusion_params.timesteps != nullptr);
+        return ernie_image.compute(n_threads,
+                                   *diffusion_params.x,
+                                   *diffusion_params.timesteps,
+                                   tensor_or_empty(diffusion_params.context));
+    }
+};
+
 #endif

src/ernie_image.hpp

@@ -0,0 +1,438 @@
+#ifndef __SD_ERNIE_IMAGE_HPP__
+#define __SD_ERNIE_IMAGE_HPP__
+
+#include <memory>
+#include <vector>
+
+#include "common_dit.hpp"
+#include "flux.hpp"
+#include "qwen_image.hpp"
+#include "rope.hpp"
+
+namespace ErnieImage {
+    constexpr int ERNIE_IMAGE_GRAPH_SIZE = 40960;
+
+    __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx,
+                                                              ggml_tensor* timesteps,
+                                                              int dim,
+                                                              int max_period = 10000) {
+        auto emb       = ggml_ext_timestep_embedding(ctx, timesteps, dim, max_period, 1.0f);
+        int64_t half   = dim / 2;
+        auto cos_part  = ggml_view_2d(ctx, emb, half, emb->ne[1], emb->nb[1], 0);
+        auto sin_part  = ggml_view_2d(ctx, emb, half, emb->ne[1], emb->nb[1], half * emb->nb[0]);
+        auto sin_first = ggml_concat(ctx, sin_part, cos_part, 0);
+        return sin_first;
+    }
+
+    __STATIC_INLINE__ ggml_tensor* apply_rotary_emb(ggml_context* ctx, ggml_tensor* x, ggml_tensor* pe) {
+        // x: [N, S, heads, head_dim]
+        // pe: [2, S, 1, head_dim], stored as ggml [head_dim, 1, S, 2].
+        int64_t head_dim = x->ne[0];
+        int64_t heads    = x->ne[1];
+        int64_t S        = x->ne[2];
+        int64_t N        = x->ne[3];
+        int64_t rot_dim  = pe->ne[0];
+        GGML_ASSERT(rot_dim <= head_dim);
+        GGML_ASSERT(rot_dim % 2 == 0);
+        GGML_ASSERT(pe->ne[1] == 1 && pe->ne[2] == S && pe->ne[3] == 2);
+
+        x           = ggml_cont(ctx, x);
+        auto x_rot  = ggml_ext_slice(ctx, x, 0, 0, rot_dim, false);
+        auto x_pass = rot_dim < head_dim ? ggml_ext_slice(ctx, x, 0, rot_dim, head_dim, false) : nullptr;
+
+        int64_t half = rot_dim / 2;
+        auto x1      = ggml_view_4d(ctx, x_rot, half, heads, S, N, x_rot->nb[1], x_rot->nb[2], x_rot->nb[3], 0);
+        auto x2      = ggml_view_4d(ctx, x_rot, half, heads, S, N, x_rot->nb[1], x_rot->nb[2], x_rot->nb[3], half * x_rot->nb[0]);
+        x1           = ggml_cont(ctx, x1);
+        x2           = ggml_cont(ctx, x2);
+        auto rotated = ggml_concat(ctx, ggml_neg(ctx, x2), x1, 0);
+
+        auto cos_emb = ggml_ext_slice(ctx, pe, 3, 0, 1, false);
+        auto sin_emb = ggml_ext_slice(ctx, pe, 3, 1, 2, false);
+
+        auto out = ggml_add(ctx, ggml_mul(ctx, x_rot, cos_emb), ggml_mul(ctx, rotated, sin_emb));
+        if (x_pass != nullptr) {
+            out = ggml_concat(ctx, out, x_pass, 0);
+        }
+        return out;
+    }
+
+    struct ErnieImageAttention : public GGMLBlock {
+        int64_t num_heads;
+        int64_t head_dim;
+
+        ErnieImageAttention(int64_t query_dim,
+                            int64_t heads,
+                            int64_t dim_head,
+                            float eps = 1e-6f)
+            : num_heads(heads), head_dim(dim_head) {
+            int64_t inner_dim  = heads * dim_head;
+            blocks["to_q"]     = std::make_shared<Linear>(query_dim, inner_dim, false);
+            blocks["to_k"]     = std::make_shared<Linear>(query_dim, inner_dim, false);
+            blocks["to_v"]     = std::make_shared<Linear>(query_dim, inner_dim, false);
+            blocks["norm_q"]   = std::make_shared<RMSNorm>(dim_head, eps);
+            blocks["norm_k"]   = std::make_shared<RMSNorm>(dim_head, eps);
+            blocks["to_out.0"] = std::make_shared<Linear>(inner_dim, query_dim, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* pe,
+                             ggml_tensor* attention_mask = nullptr) {
+            // x: [N, S, hidden_size]
+            // pe: [S, head_dim/2, 2, 2], generated in image-token-first order.
+            auto to_q     = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
+            auto to_k     = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
+            auto to_v     = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
+            auto norm_q   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
+            auto norm_k   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
+            auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
+
+            int64_t S = x->ne[1];
+            int64_t N = x->ne[2];
+
+            auto q = to_q->forward(ctx, x);
+            auto k = to_k->forward(ctx, x);
+            auto v = to_v->forward(ctx, x);
+
+            q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, S, N);  // [N, S, heads, head_dim]
+            k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_heads, S, N);  // [N, S, heads, head_dim]
+            v = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim, num_heads, S, N);  // [N, S, heads, head_dim]
+
+            q = norm_q->forward(ctx, q);
+            k = norm_k->forward(ctx, k);
+
+            q = apply_rotary_emb(ctx->ggml_ctx, q, pe);
+            k = apply_rotary_emb(ctx->ggml_ctx, k, pe);
+
+            q = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, q, 0, 2, 1, 3));  // [N, heads, S, head_dim]
+            q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]);
+
+            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);  // [N, S, hidden_size]
+            x = to_out_0->forward(ctx, x);
+            return x;
+        }
+    };
+
+    struct ErnieImageFeedForward : public GGMLBlock {
+    public:
+        ErnieImageFeedForward(int64_t hidden_size, int64_t ffn_hidden_size) {
+            blocks["gate_proj"]  = std::make_shared<Linear>(hidden_size, ffn_hidden_size, false);
+            blocks["up_proj"]    = std::make_shared<Linear>(hidden_size, ffn_hidden_size, false);
+            blocks["linear_fc2"] = std::make_shared<Linear>(ffn_hidden_size, hidden_size, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto gate_proj  = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
+            auto up_proj    = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
+            auto linear_fc2 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc2"]);
+
+            auto gate = gate_proj->forward(ctx, x);
+            gate      = ggml_ext_gelu(ctx->ggml_ctx, gate);
+            x         = up_proj->forward(ctx, x);
+            x         = ggml_mul(ctx->ggml_ctx, x, gate);
+            x         = linear_fc2->forward(ctx, x);
+            return x;
+        }
+    };
+
+    struct ErnieImageSharedAdaLNBlock : public GGMLBlock {
+    public:
+        ErnieImageSharedAdaLNBlock(int64_t hidden_size,
+                                   int64_t num_heads,
+                                   int64_t ffn_hidden_size,
+                                   float eps = 1e-6f) {
+            blocks["adaLN_sa_ln"]    = std::make_shared<RMSNorm>(hidden_size, eps);
+            blocks["self_attention"] = std::make_shared<ErnieImageAttention>(hidden_size,
+                                                                             num_heads,
+                                                                             hidden_size / num_heads,
+                                                                             eps);
+            blocks["adaLN_mlp_ln"]   = std::make_shared<RMSNorm>(hidden_size, eps);
+            blocks["mlp"]            = std::make_shared<ErnieImageFeedForward>(hidden_size, ffn_hidden_size);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* pe,
+                             const std::vector<ggml_tensor*>& temb,
+                             ggml_tensor* attention_mask = nullptr) {
+            // x: [N, image_tokens + text_tokens, hidden_size]
+            auto adaLN_sa_ln    = std::dynamic_pointer_cast<RMSNorm>(blocks["adaLN_sa_ln"]);
+            auto self_attention = std::dynamic_pointer_cast<ErnieImageAttention>(blocks["self_attention"]);
+            auto adaLN_mlp_ln   = std::dynamic_pointer_cast<RMSNorm>(blocks["adaLN_mlp_ln"]);
+            auto mlp            = std::dynamic_pointer_cast<ErnieImageFeedForward>(blocks["mlp"]);
+
+            auto shift_msa = temb[0];
+            auto scale_msa = temb[1];
+            auto gate_msa  = temb[2];
+            auto shift_mlp = temb[3];
+            auto scale_mlp = temb[4];
+            auto gate_mlp  = temb[5];
+
+            auto residual = x;
+            x             = adaLN_sa_ln->forward(ctx, x);
+            x             = Flux::modulate(ctx->ggml_ctx, x, shift_msa, scale_msa, true);
+            auto attn_out = self_attention->forward(ctx, x, pe, attention_mask);
+            x             = ggml_add(ctx->ggml_ctx, residual, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa));
+
+            residual = x;
+            x        = adaLN_mlp_ln->forward(ctx, x);
+            x        = Flux::modulate(ctx->ggml_ctx, x, shift_mlp, scale_mlp, true);
+            x        = ggml_add(ctx->ggml_ctx, residual, ggml_mul(ctx->ggml_ctx, mlp->forward(ctx, x), gate_mlp));
+            return x;
+        }
+    };
+
+    struct ErnieImageAdaLNContinuous : public GGMLBlock {
+    public:
+        ErnieImageAdaLNContinuous(int64_t hidden_size, float eps = 1e-6f) {
+            blocks["norm"]   = std::make_shared<LayerNorm>(hidden_size, eps, false);
+            blocks["linear"] = std::make_shared<Linear>(hidden_size, hidden_size * 2, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* conditioning) {
+            auto norm   = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
+            auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
+
+            auto mods  = ggml_ext_chunk(ctx->ggml_ctx, linear->forward(ctx, conditioning), 2, 0);
+            auto scale = mods[0];
+            auto shift = mods[1];
+
+            x = norm->forward(ctx, x);
+            x = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
+            return x;
+        }
+    };
+
+    struct ErnieImageParams {
+        int64_t hidden_size       = 4096;
+        int64_t num_heads         = 32;
+        int64_t num_layers        = 36;
+        int64_t ffn_hidden_size   = 12288;
+        int64_t in_channels       = 128;
+        int64_t out_channels      = 128;
+        int patch_size            = 1;
+        int64_t text_in_dim       = 3072;
+        int theta                 = 256;
+        std::vector<int> axes_dim = {32, 48, 48};
+        int axes_dim_sum          = 128;
+        float eps                 = 1e-6f;
+    };
+
+    class ErnieImageModel : public GGMLBlock {
+    public:
+        ErnieImageParams params;
+
+        ErnieImageModel() = default;
+        ErnieImageModel(ErnieImageParams params)
+            : params(params) {
+            blocks["x_embedder.proj"] = std::make_shared<Conv2d>(params.in_channels,
+                                                                 params.hidden_size,
+                                                                 std::pair<int, int>{params.patch_size, params.patch_size},
+                                                                 std::pair<int, int>{params.patch_size, params.patch_size},
+                                                                 std::pair<int, int>{0, 0},
+                                                                 std::pair<int, int>{1, 1},
+                                                                 true);
+            if (params.text_in_dim != params.hidden_size) {
+                blocks["text_proj"] = std::make_shared<Linear>(params.text_in_dim, params.hidden_size, false);
+            }
+            blocks["time_embedding"]     = std::make_shared<Qwen::TimestepEmbedding>(params.hidden_size, params.hidden_size);
+            blocks["adaLN_modulation.1"] = std::make_shared<Linear>(params.hidden_size, 6 * params.hidden_size, true);
+
+            for (int i = 0; i < params.num_layers; i++) {
+                blocks["layers." + std::to_string(i)] = std::make_shared<ErnieImageSharedAdaLNBlock>(params.hidden_size,
+                                                                                                     params.num_heads,
+                                                                                                     params.ffn_hidden_size,
+                                                                                                     params.eps);
+            }
+
+            blocks["final_norm"]   = std::make_shared<ErnieImageAdaLNContinuous>(params.hidden_size, params.eps);
+            blocks["final_linear"] = std::make_shared<Linear>(params.hidden_size,
+                                                              params.patch_size * params.patch_size * params.out_channels,
+                                                              true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timestep,
+                             ggml_tensor* context,
+                             ggml_tensor* pe) {
+            // x: [N, C, H, W]
+            // context: [N, text_tokens, 3072]
+            // pe: [image_tokens + text_tokens, head_dim/2, 2, 2]
+            GGML_ASSERT(context != nullptr);
+            GGML_ASSERT(x->ne[1] % params.patch_size == 0 && x->ne[0] % params.patch_size == 0);
+
+            int64_t W     = x->ne[0];
+            int64_t H     = x->ne[1];
+            int64_t Hp    = H / params.patch_size;
+            int64_t Wp    = W / params.patch_size;
+            int64_t n_img = Hp * Wp;
+            int64_t N     = x->ne[3];
+
+            auto x_embedder_proj = std::dynamic_pointer_cast<Conv2d>(blocks["x_embedder.proj"]);
+            auto time_embedding  = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["time_embedding"]);
+            auto adaLN_mod       = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
+            auto final_norm      = std::dynamic_pointer_cast<ErnieImageAdaLNContinuous>(blocks["final_norm"]);
+            auto final_linear    = std::dynamic_pointer_cast<Linear>(blocks["final_linear"]);
+
+            auto img = x_embedder_proj->forward(ctx, x);                                                  // [N, hidden_size, Hp, Wp]
+            img      = ggml_reshape_3d(ctx->ggml_ctx, img, img->ne[0] * img->ne[1], img->ne[2], N);       // [N, hidden_size, image_tokens]
+            img      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));  // [N, image_tokens, hidden_size]
+
+            auto txt       = context;
+            auto text_proj = std::dynamic_pointer_cast<Linear>(blocks["text_proj"]);
+            if (text_proj) {
+                txt = text_proj->forward(ctx, txt);
+            }
+
+            auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1);  // [N, image_tokens + text_tokens, hidden_size]
+
+            auto sample = timestep_embedding_sin_cos(ctx->ggml_ctx, timestep, static_cast<int>(params.hidden_size));
+            auto c      = time_embedding->forward(ctx, sample);  // [N, hidden_size]
+
+            auto mod_params = adaLN_mod->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, 6 * hidden_size]
+            auto chunks     = ggml_ext_chunk(ctx->ggml_ctx, mod_params, 6, 0);
+            std::vector<ggml_tensor*> temb;
+            temb.reserve(6);
+            for (auto chunk : chunks) {
+                temb.push_back(ggml_reshape_3d(ctx->ggml_ctx, chunk, chunk->ne[0], 1, chunk->ne[1]));  // [N, 1, hidden_size]
+            }
+
+            for (int i = 0; i < params.num_layers; i++) {
+                auto layer    = std::dynamic_pointer_cast<ErnieImageSharedAdaLNBlock>(blocks["layers." + std::to_string(i)]);
+                hidden_states = layer->forward(ctx, hidden_states, pe, temb);
+            }
+
+            hidden_states = final_norm->forward(ctx, hidden_states, c);
+            hidden_states = final_linear->forward(ctx, hidden_states);                  // [N, image_tokens, p*p*out_channels]
+            auto patches  = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img);  // [N, image_tokens, hidden_size]
+
+            auto out = DiT::unpatchify(ctx->ggml_ctx,
+                                       patches,
+                                       Hp,
+                                       Wp,
+                                       params.patch_size,
+                                       params.patch_size,
+                                       false);  // [N, out_channels, H, W]
+            return out;
+        }
+    };
+
+    struct ErnieImageRunner : public GGMLRunner {
+        ErnieImageParams ernie_params;
+        ErnieImageModel ernie_image;
+        std::vector<float> pe_vec;
+
+        ErnieImageRunner(ggml_backend_t backend,
+                         bool offload_params_to_cpu,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "")
+            : GGMLRunner(backend, offload_params_to_cpu) {
+            ernie_params.num_layers = 0;
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (ends_with(name, "x_embedder.proj.weight") && tensor_storage.n_dims == 4) {
+                    ernie_params.patch_size  = static_cast<int>(tensor_storage.ne[0]);
+                    ernie_params.in_channels = tensor_storage.ne[2];
+                    ernie_params.hidden_size = tensor_storage.ne[3];
+                } else if (ends_with(name, "text_proj.weight") && tensor_storage.n_dims == 2) {
+                    ernie_params.text_in_dim = tensor_storage.ne[0];
+                } else if (ends_with(name, "layers.0.self_attention.norm_q.weight")) {
+                    int64_t head_dim       = tensor_storage.ne[0];
+                    ernie_params.num_heads = ernie_params.hidden_size / head_dim;
+                } else if (ends_with(name, "layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) {
+                    ernie_params.ffn_hidden_size = tensor_storage.ne[1];
+                } else if (ends_with(name, "final_linear.weight") && tensor_storage.n_dims == 2) {
+                    int64_t out_dim           = tensor_storage.ne[1];
+                    ernie_params.out_channels = out_dim / ernie_params.patch_size / ernie_params.patch_size;
+                }
+
+                size_t pos = name.find("layers.");
+                if (pos != std::string::npos) {
+                    std::string layer_name = name.substr(pos);
+                    auto items             = split_string(layer_name, '.');
+                    if (items.size() > 1) {
+                        int block_index = atoi(items[1].c_str());
+                        if (block_index + 1 > ernie_params.num_layers) {
+                            ernie_params.num_layers = block_index + 1;
+                        }
+                    }
+                }
+            }
+            if (ernie_params.num_layers == 0) {
+                ernie_params.num_layers = 36;
+            }
+            ernie_params.axes_dim_sum = 0;
+            for (int axis_dim : ernie_params.axes_dim) {
+                ernie_params.axes_dim_sum += axis_dim;
+            }
+
+            LOG_INFO("ernie_image: layers = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64
+                     ", ffn_hidden_size = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
+                     ernie_params.num_layers,
+                     ernie_params.hidden_size,
+                     ernie_params.num_heads,
+                     ernie_params.ffn_hidden_size,
+                     ernie_params.in_channels,
+                     ernie_params.out_channels);
+
+            ernie_image = ErnieImageModel(ernie_params);
+            ernie_image.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "ernie_image";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+            ernie_image.get_param_tensors(tensors, prefix);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor) {
+            ggml_cgraph* gf        = new_graph_custom(ERNIE_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(!context_tensor.empty());
+            ggml_tensor* context = make_input(context_tensor);
+
+            pe_vec      = Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
+                                                   static_cast<int>(x->ne[0]),
+                                                   ernie_params.patch_size,
+                                                   static_cast<int>(x->ne[3]),
+                                                   static_cast<int>(context->ne[1]),
+                                                   ernie_params.theta,
+                                                   circular_y_enabled,
+                                                   circular_x_enabled,
+                                                   ernie_params.axes_dim);
+            int pos_len = static_cast<int>(pe_vec.size() / ernie_params.axes_dim_sum / 2);
+            auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, ernie_params.axes_dim_sum, 1, pos_len, 2);
+            set_backend_tensor_data(pe, pe_vec.data());
+
+            auto runner_ctx  = get_context();
+            ggml_tensor* out = ernie_image.forward(&runner_ctx, x, timesteps, context, pe);
+            ggml_build_forward_expand(gf, out);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timesteps,
+                                  const sd::Tensor<float>& context) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(x, timesteps, context);
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        }
+    };
+}  // namespace ErnieImage
+
+#endif  // __SD_ERNIE_IMAGE_HPP__

src/llm.hpp

@@ -14,652 +14,22 @@
 #include <utility>
 #include <vector>
 
-#include "clip.hpp"
 #include "ggml_extend.hpp"
 #include "json.hpp"
 #include "rope.hpp"
-#include "tokenize_util.h"
+#include "tokenizers/bpe_tokenizer.h"
-#include "vocab/vocab.h"
+#include "tokenizers/gemma_tokenizer.h"
+#include "tokenizers/mistral_tokenizer.h"
+#include "tokenizers/qwen2_tokenizer.h"
 
 namespace LLM {
     constexpr int LLM_GRAPH_SIZE = 10240;
 
-    class BPETokenizer {
-    protected:
-        std::map<int, std::u32string> byte_encoder;
-        std::map<std::u32string, int> byte_decoder;
-        std::map<std::u32string, int> encoder;
-        std::map<int, std::u32string> decoder;
-        std::map<std::pair<std::u32string, std::u32string>, int> bpe_ranks;
-        std::regex pat;
-        int encoder_len;
-        int bpe_len;
-
-        std::string UNK_TOKEN;
-        std::string BOS_TOKEN;
-        std::string EOS_TOKEN;
-        std::string PAD_TOKEN;
-
-        int UNK_TOKEN_ID;
-        int BOS_TOKEN_ID;
-        int EOS_TOKEN_ID;
-        int PAD_TOKEN_ID;
-
-        std::vector<std::string> special_tokens;
-
-        bool add_bos_token  = false;
-        bool byte_level_bpe = true;
-        bool byte_fallback  = false;
-
-    protected:
-        virtual std::string preprocess(const std::string& text) const {
-            return text;
-        }
-
-        static std::set<std::pair<std::u32string, std::u32string>> get_pairs(const std::vector<std::u32string>& subwords) {
-            std::set<std::pair<std::u32string, std::u32string>> pairs;
-            if (subwords.size() == 0) {
-                return pairs;
-            }
-            std::u32string prev_subword = subwords[0];
-            for (int i = 1; i < subwords.size(); i++) {
-                std::u32string subword = subwords[i];
-                std::pair<std::u32string, std::u32string> pair(prev_subword, subword);
-                pairs.insert(pair);
-                prev_subword = subword;
-            }
-            return pairs;
-        }
-
-        bool is_special_token(const std::string& token) {
-            for (auto& special_token : special_tokens) {
-                if (special_token == token) {
-                    return true;
-                }
-            }
-            return false;
-        }
-
-        static std::vector<std::u32string> split_utf32(const std::u32string& s, char32_t delim) {
-            std::vector<std::u32string> result;
-            size_t start = 0;
-
-            while (true) {
-                size_t pos = s.find(delim, start);
-                if (pos == std::u32string::npos) {
-                    result.emplace_back(s.substr(start));
-                    break;
-                }
-                result.emplace_back(s.substr(start, pos - start));
-                start = pos + 1;
-            }
-            return result;
-        }
-
-    public:
-        BPETokenizer() = default;
-
-        std::u32string bpe(const std::u32string& token) {
-            std::vector<std::u32string> word;
-
-            for (int i = 0; i < token.size(); i++) {
-                word.emplace_back(1, token[i]);
-            }
-
-            std::set<std::pair<std::u32string, std::u32string>> pairs = get_pairs(word);
-
-            if (pairs.empty()) {
-                return token;
-            }
-
-            while (true) {
-                auto min_pair_iter = std::min_element(pairs.begin(),
-                                                      pairs.end(),
-                                                      [&](const std::pair<std::u32string, std::u32string>& a,
-                                                          const std::pair<std::u32string, std::u32string>& b) {
-                                                          if (bpe_ranks.find(a) == bpe_ranks.end()) {
-                                                              return false;
-                                                          } else if (bpe_ranks.find(b) == bpe_ranks.end()) {
-                                                              return true;
-                                                          }
-                                                          return bpe_ranks.at(a) < bpe_ranks.at(b);
-                                                      });
-
-                const std::pair<std::u32string, std::u32string>& bigram = *min_pair_iter;
-
-                if (bpe_ranks.find(bigram) == bpe_ranks.end()) {
-                    break;
-                }
-
-                std::u32string first  = bigram.first;
-                std::u32string second = bigram.second;
-                std::vector<std::u32string> new_word;
-                int32_t i = 0;
-
-                while (i < word.size()) {
-                    auto it = std::find(word.begin() + i, word.end(), first);
-                    if (it == word.end()) {
-                        new_word.insert(new_word.end(), word.begin() + i, word.end());
-                        break;
-                    }
-                    new_word.insert(new_word.end(), word.begin() + i, it);
-                    i = static_cast<int32_t>(std::distance(word.begin(), it));
-
-                    if (word[i] == first && i < static_cast<int32_t>(word.size()) - 1 && word[i + 1] == second) {
-                        new_word.push_back(first + second);
-                        i += 2;
-                    } else {
-                        new_word.push_back(word[i]);
-                        i += 1;
-                    }
-                }
-
-                word = new_word;
-
-                if (word.size() == 1) {
-                    break;
-                }
-                pairs = get_pairs(word);
-            }
-
-            std::u32string result;
-            for (int i = 0; i < word.size(); i++) {
-                result += word[i];
-                if (i != word.size() - 1) {
-                    result += utf8_to_utf32(" ");
-                }
-            }
-
-            return result;
-        }
-
-        std::vector<int> tokenize(std::string text,
-                                  on_new_token_cb_t on_new_token_cb = nullptr,
-                                  size_t max_length                 = 0,
-                                  bool padding                      = false) {
-            std::vector<int32_t> tokens = encode(text, on_new_token_cb);
-
-            if (max_length > 0) {
-                if (tokens.size() < max_length) {
-                    tokens.resize(max_length);
-                } else {
-                    if (padding) {
-                        tokens.insert(tokens.end(), max_length - tokens.size(), PAD_TOKEN_ID);
-                    }
-                }
-            }
-
-            return tokens;
-        }
-
-        void pad_tokens(std::vector<int>& tokens,
-                        std::vector<float>& weights,
-                        size_t max_length = 0,
-                        bool padding      = false) {
-            if (add_bos_token) {
-                tokens.insert(tokens.begin(), BOS_TOKEN_ID);
-                weights.insert(weights.begin(), 1.f);
-            }
-            if (max_length > 0 && padding) {
-                size_t n = static_cast<size_t>(std::ceil(tokens.size() * 1.f / max_length));
-                if (n == 0) {
-                    n = 1;
-                }
-                size_t length = max_length * n;
-                LOG_DEBUG("token length: %llu", length);
-                tokens.insert(tokens.end(), length - tokens.size(), PAD_TOKEN_ID);
-                weights.insert(weights.end(), length - weights.size(), 1.f);
-            }
-        }
-
-        virtual std::vector<int> encode(std::string text, on_new_token_cb_t on_new_token_cb = nullptr) {
-            std::string original_text = text;
-            std::vector<int32_t> bpe_tokens;
-            std::vector<std::string> token_strs;
-
-            auto splited_texts = split_with_special_tokens(text, special_tokens);
-
-            for (auto& splited_text : splited_texts) {
-                if (is_special_token(splited_text)) {
-                    bpe_tokens.push_back(encoder[utf8_to_utf32(splited_text)]);
-                    token_strs.push_back(splited_text);
-                    continue;
-                }
-                auto tokens = token_split(splited_text);
-                for (auto& token : tokens) {
-                    if (on_new_token_cb != nullptr) {
-                        bool skip = on_new_token_cb(token, bpe_tokens);
-                        if (skip) {
-                            continue;
-                        }
-                    }
-
-                    std::string token_str = preprocess(token);
-                    std::u32string utf32_token;
-                    if (byte_level_bpe) {
-                        for (int i = 0; i < token_str.length(); i++) {
-                            unsigned char b = token_str[i];
-                            utf32_token += byte_encoder[b];
-                        }
-                    } else {
-                        utf32_token = utf8_to_utf32(token_str);
-                    }
-
-                    auto bpe_strs = bpe(utf32_token);
-                    for (const auto& bpe_str : split_utf32(bpe_strs, U' ')) {
-                        int token_id;
-                        auto iter = encoder.find(bpe_str);
-                        if (iter != encoder.end()) {
-                            token_id = iter->second;
-                        } else {
-                            if (byte_fallback) {
-                                auto utf8_token_str = utf32_to_utf8(bpe_str);
-                                for (int i = 0; i < utf8_token_str.length(); i++) {
-                                    unsigned char b = utf8_token_str[i];
-                                    char hex_buf[16];
-                                    snprintf(hex_buf, sizeof(hex_buf), "<0x%02X>", b);
-                                    iter = encoder.find(utf8_to_utf32(hex_buf));
-                                    GGML_ASSERT(iter != encoder.end());
-                                    bpe_tokens.push_back(token_id);
-                                    token_strs.push_back(hex_buf);
-                                }
-                                continue;
-                            } else {
-                                token_id = UNK_TOKEN_ID;
-                            }
-                        }
-                        bpe_tokens.push_back(token_id);
-                        token_strs.push_back(utf32_to_utf8(bpe_str));
-                    }
-                }
-            }
-
-            std::stringstream ss;
-            ss << "[";
-            for (auto token : token_strs) {
-                ss << "\"" << token << "\", ";
-            }
-            ss << "]";
-            LOG_DEBUG("split prompt \"%s\" to tokens %s", original_text.c_str(), ss.str().c_str());
-            return bpe_tokens;
-        }
-    };
-
-    class Qwen2Tokenizer : public BPETokenizer {
-    protected:
-        void load_from_merges(const std::string& merges_utf8_str) {
-            auto byte_unicode_pairs = bytes_to_unicode();
-            byte_encoder            = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
-            for (auto& pair : byte_unicode_pairs) {
-                byte_decoder[pair.second] = pair.first;
-            }
-            std::u32string merges_utf32_str    = utf8_to_utf32(merges_utf8_str);
-            std::vector<std::u32string> merges = split_utf32(merges_utf32_str, U'\n');
-            std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
-            for (const auto& merge : merges) {
-                size_t space_pos = merge.find(' ');
-                merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
-            }
-            LOG_DEBUG("merges size %zu", merge_pairs.size());
-
-            std::vector<std::u32string> tokens;
-            for (const auto& pair : byte_unicode_pairs) {
-                tokens.push_back(pair.second);
-            }
-            for (const auto& merge : merge_pairs) {
-                tokens.push_back(merge.first + merge.second);
-            }
-            for (auto& special_token : special_tokens) {
-                tokens.push_back(utf8_to_utf32(special_token));
-            }
-
-            int i = 0;
-            for (const auto& token : tokens) {
-                encoder[token] = i;
-                decoder[i]     = token;
-                i++;
-            }
-            encoder_len = i;
-            LOG_DEBUG("vocab size: %d", encoder_len);
-
-            int rank = 0;
-            for (const auto& merge : merge_pairs) {
-                bpe_ranks[merge] = rank++;
-            }
-            bpe_len = rank;
-        };
-
-    public:
-        explicit Qwen2Tokenizer(const std::string& merges_utf8_str = "") {
-            UNK_TOKEN = "<|endoftext|>";
-            EOS_TOKEN = "<|endoftext|>";
-            PAD_TOKEN = "<|endoftext|>";
-
-            UNK_TOKEN_ID = 151643;
-            EOS_TOKEN_ID = 151643;
-            PAD_TOKEN_ID = 151643;
-
-            special_tokens = {
-                "<|endoftext|>",
-                "<|im_start|>",
-                "<|im_end|>",
-                "<|object_ref_start|>",
-                "<|object_ref_end|>",
-                "<|box_start|>",
-                "<|box_end|>",
-                "<|quad_start|>",
-                "<|quad_end|>",
-                "<|vision_start|>",
-                "<|vision_end|>",
-                "<|vision_pad|>",
-                "<|image_pad|>",
-                "<|video_pad|>",
-                "<tool_call>",
-                "</tool_call>",
-                "<|fim_prefix|>",
-                "<|fim_middle|>",
-                "<|fim_suffix|>",
-                "<|fim_pad|>",
-                "<|repo_name|>",
-                "<|file_sep|>",
-                "<tool_response>",
-                "</tool_response>",
-                "<think>",
-                "</think>",
-            };
-
-            if (merges_utf8_str.size() > 0) {
-                load_from_merges(merges_utf8_str);
-            } else {
-                load_from_merges(load_qwen2_merges());
-            }
-        }
-    };
-
-    class MistralTokenizer : public BPETokenizer {
-    protected:
-        void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
-            nlohmann::json vocab;
-
-            try {
-                vocab = nlohmann::json::parse(vocab_utf8_str);
-            } catch (const nlohmann::json::parse_error&) {
-                GGML_ABORT("invalid vocab json str");
-            }
-            for (const auto& [key, value] : vocab.items()) {
-                std::u32string token = utf8_to_utf32(key);
-                int i                = value;
-                encoder[token]       = i;
-                decoder[i]           = token;
-            }
-            encoder_len = static_cast<int>(vocab.size());
-            LOG_DEBUG("vocab size: %d", encoder_len);
-
-            auto byte_unicode_pairs = bytes_to_unicode();
-            byte_encoder            = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
-            for (auto& pair : byte_unicode_pairs) {
-                byte_decoder[pair.second] = pair.first;
-            }
-            std::u32string merges_utf32_str    = utf8_to_utf32(merges_utf8_str);
-            std::vector<std::u32string> merges = split_utf32(merges_utf32_str, U'\n');
-            std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
-            for (const auto& merge : merges) {
-                size_t space_pos = merge.find(' ');
-                merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
-            }
-            LOG_DEBUG("merges size %zu", merge_pairs.size());
-
-            int rank = 0;
-            for (const auto& merge : merge_pairs) {
-                bpe_ranks[merge] = rank++;
-            }
-            bpe_len = rank;
-        };
-
-    public:
-        explicit MistralTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "") {
-            add_bos_token = true;
-
-            UNK_TOKEN = "<unk>";
-            BOS_TOKEN = "<s>";
-            EOS_TOKEN = "</s>";
-            PAD_TOKEN = "<pad>";
-
-            UNK_TOKEN_ID = 0;
-            BOS_TOKEN_ID = 1;
-            EOS_TOKEN_ID = 2;
-            PAD_TOKEN_ID = 11;
-
-            special_tokens = {
-                "<unk>",
-                "<s>",
-                "</s>",
-                "[INST]",
-                "[/INST]",
-                "[AVAILABLE_TOOLS]",
-                "[/AVAILABLE_TOOLS]",
-                "[TOOL_RESULTS]",
-                "[/TOOL_RESULTS]",
-                "[TOOL_CALLS]",
-                "[IMG]",
-                "<pad>",
-                "[IMG_BREAK]",
-                "[IMG_END]",
-                "[PREFIX]",
-                "[MIDDLE]",
-                "[SUFFIX]",
-                "[SYSTEM_PROMPT]",
-                "[/SYSTEM_PROMPT]",
-                "[TOOL_CONTENT]",
-            };
-            for (int i = 20; i < 1000; i++) {
-                special_tokens.push_back("<SPECIAL_" + std::to_string(i) + ">");
-            }
-
-            if (merges_utf8_str.size() > 0 && vocab_utf8_str.size() > 0) {
-                load_from_merges(merges_utf8_str, vocab_utf8_str);
-            } else {
-                load_from_merges(load_mistral_merges(), load_mistral_vocab_json());
-            }
-        }
-    };
-
-    class GemmaTokenizer : public BPETokenizer {
-    protected:
-        std::vector<std::string> special_tokens_before_merge;
-        std::vector<std::string> special_tokens_after_merge;
-
-        std::string preprocess(const std::string& text) const override {
-            std::string normalized = text;
-            size_t pos             = 0;
-            while ((pos = normalized.find(' ', pos)) != std::string::npos) {
-                normalized.replace(pos, 1, "\xE2\x96\x81");
-                pos += 3;
-            }
-            return normalized;
-        }
-
-        void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
-            nlohmann::json vocab;
-            try {
-                vocab = nlohmann::json::parse(vocab_utf8_str);
-            } catch (const nlohmann::json::parse_error&) {
-                GGML_ABORT("invalid vocab json str");
-            }
-            for (const auto& [key, value] : vocab.items()) {
-                std::u32string token = utf8_to_utf32(key);
-                int i                = value;
-                encoder[token]       = i;
-                decoder[i]           = token;
-            }
-            encoder_len = static_cast<int>(vocab.size());
-            LOG_DEBUG("vocab size: %d", encoder_len);
-
-            std::u32string merges_utf32_str    = utf8_to_utf32(merges_utf8_str);
-            std::vector<std::u32string> merges = split_utf32(merges_utf32_str, U'\n');
-            std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
-            for (const auto& merge : merges) {
-                size_t space_pos = merge.find(' ');
-                merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
-            }
-            LOG_DEBUG("merges size %zu", merge_pairs.size());
-
-            int rank = 0;
-            for (const auto& merge : merge_pairs) {
-                bpe_ranks[merge] = rank++;
-            }
-            bpe_len = rank;
-        };
-
-    public:
-        explicit GemmaTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_json_utf8_str = "") {
-            byte_level_bpe = false;
-            byte_fallback  = true;
-            add_bos_token = true;
-            PAD_TOKEN      = "<pad>";
-            EOS_TOKEN      = "<eos>";
-            BOS_TOKEN      = "<bos>";
-            UNK_TOKEN      = "<unk>";
-
-            PAD_TOKEN_ID = 0;
-            EOS_TOKEN_ID = 1;
-            BOS_TOKEN_ID = 2;
-            UNK_TOKEN_ID = 3;
-
-            special_tokens_before_merge = {
-                PAD_TOKEN,
-                EOS_TOKEN,
-                BOS_TOKEN,
-                UNK_TOKEN,
-                "<mask>",
-                "[multimodal]",
-            };
-            for (int i = 0; i <= 98; i++) {
-                special_tokens_before_merge.push_back("<unused" + std::to_string(i) + ">");
-            }
-            special_tokens_before_merge.push_back("<start_of_turn>");
-            special_tokens_before_merge.push_back("<end_of_turn>");
-            for (int i = 1; i <= 31; i++) {
-                special_tokens_before_merge.push_back(std::string(i, '\n'));
-            }
-            for (int i = 2; i <= 31; i++) {
-                std::string whitespace_token;
-                for (int j = 0; j < i; j++) {
-                    whitespace_token += "\xE2\x96\x81";
-                }
-                special_tokens_before_merge.push_back(whitespace_token);
-            }
-            std::vector<std::string> html_tokens = {
-                "<table>",
-                "<caption>",
-                "<thead>",
-                "<tbody>",
-                "<tfoot>",
-                "<tr>",
-                "<th>",
-                "<td>",
-                "</table>",
-                "</caption>",
-                "</thead>",
-                "</tbody>",
-                "</tfoot>",
-                "</tr>",
-                "</th>",
-                "</td>",
-                "<h1>",
-                "<h2>",
-                "<h3>",
-                "<h4>",
-                "<h5>",
-                "<h6>",
-                "<blockquote>",
-                "</h1>",
-                "</h2>",
-                "</h3>",
-                "</h4>",
-                "</h5>",
-                "</h6>",
-                "</blockquote>",
-                "<strong>",
-                "<em>",
-                "<b>",
-                "<i>",
-                "<u>",
-                "<s>",
-                "<sub>",
-                "<sup>",
-                "<code>",
-                "</strong>",
-                "</em>",
-                "</b>",
-                "</i>",
-                "</u>",
-                "</s>",
-                "</sub>",
-                "</sup>",
-                "</code>",
-                "<a>",
-                "<html>",
-                "<body>",
-                "<img>",
-                "<span>",
-                "<bbox>",
-                "<ul>",
-                "<li>",
-                "<div>",
-                "<iframe>",
-                "<footer>",
-                "</a>",
-                "</html>",
-                "</body>",
-                "</img>",
-                "</span>",
-                "</bbox>",
-                "</ul>",
-                "</li>",
-                "</div>",
-                "</iframe>",
-                "</footer>",
-            };
-            special_tokens_before_merge.insert(special_tokens_before_merge.end(),
-                                               html_tokens.begin(),
-                                               html_tokens.end());
-            for (int i = 0; i <= 0xFF; i++) {
-                char hex_buf[16];
-                snprintf(hex_buf, sizeof(hex_buf), "<0x%02X>", i);
-                special_tokens_before_merge.push_back(hex_buf);
-            }
-
-            special_tokens_after_merge = {
-                "<start_of_image>",
-                "<end_of_image>",
-            };
-            for (int i = 1; i <= 31; i++) {
-                special_tokens_after_merge.insert(special_tokens_after_merge.begin() + i - 1,
-                                                  std::string(i, '\t'));
-            }
-            for (int i = 99; i <= 6241; i++) {
-                special_tokens_after_merge.push_back("<unused" + std::to_string(i) + ">");
-            }
-            special_tokens_after_merge.push_back("<image_soft_token>");
-
-            special_tokens = special_tokens_before_merge;
-            special_tokens.insert(special_tokens.end(),
-                                  special_tokens_after_merge.begin(),
-                                  special_tokens_after_merge.end());
-
-            if (merges_utf8_str.size() > 0 && vocab_json_utf8_str.size() > 0) {
-                load_from_merges(merges_utf8_str, vocab_json_utf8_str);
-            } else {
-                load_from_merges(load_gemma_merges(), load_gemma_vocab_json());
-            }
-        }
-    };
-
     enum class LLMArch {
         QWEN2_5_VL,
         QWEN3,
         MISTRAL_SMALL_3_2,
+        MINISTRAL_3_3B,
         ARCH_COUNT,
     };
 
@@ -667,6 +37,7 @@ namespace LLM {
         "qwen2.5vl",
         "qwen3",
         "mistral_small3.2",
+        "ministral3.3b",
     };
 
     struct LLMVisionParams {
@@ -1051,6 +422,9 @@ namespace LLM {
             if (arch == LLMArch::MISTRAL_SMALL_3_2) {
                 q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NORMAL, 8192, 1000000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
                 k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NORMAL, 8192, 1000000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
+            } else if (arch == LLMArch::MINISTRAL_3_3B) {
+                q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 262144, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
+                k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 262144, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
             } else if (arch == LLMArch::QWEN3) {
                 q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
                 k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
@@ -1266,7 +640,7 @@ namespace LLM {
                   bool enable_vision_ = false)
             : GGMLRunner(backend, offload_params_to_cpu), enable_vision(enable_vision_) {
             params.arch = arch;
-            if (arch == LLMArch::MISTRAL_SMALL_3_2) {
+            if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
                 params.head_dim     = 128;
                 params.num_heads    = 32;
                 params.num_kv_heads = 8;
@@ -1378,7 +752,7 @@ namespace LLM {
             }
 
             int64_t n_tokens = input_ids->ne[0];
-            if (params.arch == LLMArch::MISTRAL_SMALL_3_2 || params.arch == LLMArch::QWEN3) {
+            if (params.arch == LLMArch::MISTRAL_SMALL_3_2 || params.arch == LLMArch::MINISTRAL_3_3B || params.arch == LLMArch::QWEN3) {
                 input_pos_vec.resize(n_tokens);
                 for (int i = 0; i < n_tokens; ++i) {
                     input_pos_vec[i] = i;
@@ -1614,7 +988,7 @@ namespace LLM {
                     const std::string prefix                       = "",
                     bool enable_vision                             = false)
             : model(arch, backend, offload_params_to_cpu, tensor_storage_map, prefix, enable_vision) {
-            if (arch == LLMArch::MISTRAL_SMALL_3_2) {
+            if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
                 tokenizer = std::make_shared<MistralTokenizer>();
             } else {
                 tokenizer = std::make_shared<Qwen2Tokenizer>();
@@ -1662,7 +1036,7 @@ namespace LLM {
                 weights.insert(weights.end(), curr_tokens.size(), curr_weight);
             }
 
-            tokenizer->pad_tokens(tokens, weights, max_length, padding);
+            tokenizer->pad_tokens(tokens, &weights, nullptr, padding ? max_length : 0, padding ? max_length : 100000000, padding);
 
             // for (int i = 0; i < tokens.size(); i++) {
             //     std::cout << tokens[i] << ":" << weights[i] << ", ";

src/model.h

@@ -5,20 +5,13 @@
 #include <map>
 #include <memory>
 #include <set>
-#include <sstream>
 #include <string>
-#include <tuple>
-#include <utility>
 #include <vector>
 
 #include "ggml-backend.h"
 #include "ggml.h"
-#include "gguf.h"
+#include "model_io/tensor_storage.h"
-#include "json.hpp"
 #include "ordered_map.hpp"
-#include "zip.h"
-
-#define SD_MAX_DIMS 5
 
 enum SDVersion {
     VERSION_SD1,
@@ -28,7 +21,8 @@ enum SDVersion {
     VERSION_SD2,
     VERSION_SD2_INPAINT,
     VERSION_SD2_TINY_UNET,
-    VERSION_SDXS,
+    VERSION_SDXS_512_DS,
+    VERSION_SDXS_09,
     VERSION_SDXL,
     VERSION_SDXL_INPAINT,
     VERSION_SDXL_PIX2PIX,
@@ -50,18 +44,19 @@ enum SDVersion {
     VERSION_FLUX2_KLEIN,
     VERSION_Z_IMAGE,
     VERSION_OVIS_IMAGE,
+    VERSION_ERNIE_IMAGE,
     VERSION_COUNT,
 };
 
 static inline bool sd_version_is_sd1(SDVersion version) {
-    if (version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX || version == VERSION_SD1_TINY_UNET || version == VERSION_SDXS) {
+    if (version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX || version == VERSION_SD1_TINY_UNET || version == VERSION_SDXS_512_DS) {
         return true;
     }
     return false;
 }
 
 static inline bool sd_version_is_sd2(SDVersion version) {
-    if (version == VERSION_SD2 || version == VERSION_SD2_INPAINT || version == VERSION_SD2_TINY_UNET) {
+    if (version == VERSION_SD2 || version == VERSION_SD2_INPAINT || version == VERSION_SD2_TINY_UNET || version == VERSION_SDXS_09) {
         return true;
     }
     return false;
@@ -137,6 +132,20 @@ static inline bool sd_version_is_z_image(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_ernie_image(SDVersion version) {
+    if (version == VERSION_ERNIE_IMAGE) {
+        return true;
+    }
+    return false;
+}
+
+static inline bool sd_version_uses_flux2_vae(SDVersion version) {
+    if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version)) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_is_inpaint(SDVersion version) {
     if (version == VERSION_SD1_INPAINT ||
         version == VERSION_SD2_INPAINT ||
@@ -155,7 +164,8 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_wan(version) ||
         sd_version_is_qwen_image(version) ||
         sd_version_is_anima(version) ||
-        sd_version_is_z_image(version)) {
+        sd_version_is_z_image(version) ||
+        sd_version_is_ernie_image(version)) {
         return true;
     }
     return false;
@@ -178,116 +188,10 @@ enum PMVersion {
     PM_VERSION_2,
 };
 
-struct TensorStorage {
-    std::string name;
-    ggml_type type          = GGML_TYPE_F32;
-    ggml_type expected_type = GGML_TYPE_COUNT;
-    bool is_f8_e4m3         = false;
-    bool is_f8_e5m2         = false;
-    bool is_f64             = false;
-    bool is_i64             = false;
-    int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
-    int n_dims              = 0;
-
-    size_t file_index = 0;
-    int index_in_zip  = -1;  // >= means stored in a zip file
-    uint64_t offset   = 0;   // offset in file
-
-    TensorStorage() = default;
-
-    TensorStorage(std::string name, ggml_type type, const int64_t* ne, int n_dims, size_t file_index, size_t offset = 0)
-        : name(std::move(name)), type(type), n_dims(n_dims), file_index(file_index), offset(offset) {
-        for (int i = 0; i < n_dims; i++) {
-            this->ne[i] = ne[i];
-        }
-    }
-
-    int64_t nelements() const {
-        int64_t n = 1;
-        for (int i = 0; i < SD_MAX_DIMS; i++) {
-            n *= ne[i];
-        }
-        return n;
-    }
-
-    int64_t nbytes() const {
-        return nelements() * ggml_type_size(type) / ggml_blck_size(type);
-    }
-
-    int64_t nbytes_to_read() const {
-        if (is_f8_e4m3 || is_f8_e5m2) {
-            return nbytes() / 2;
-        } else if (is_f64 || is_i64) {
-            return nbytes() * 2;
-        } else {
-            return nbytes();
-        }
-    }
-
-    void unsqueeze() {
-        if (n_dims == 2) {
-            n_dims = 4;
-            ne[3]  = ne[1];
-            ne[2]  = ne[0];
-            ne[1]  = 1;
-            ne[0]  = 1;
-        }
-    }
-
-    std::vector<TensorStorage> chunk(size_t n) {
-        std::vector<TensorStorage> chunks;
-        uint64_t chunk_size = nbytes_to_read() / n;
-        // printf("%d/%d\n", chunk_size, nbytes_to_read());
-        reverse_ne();
-        for (size_t i = 0; i < n; i++) {
-            TensorStorage chunk_i = *this;
-            chunk_i.ne[0]         = ne[0] / n;
-            chunk_i.offset        = offset + i * chunk_size;
-            chunk_i.reverse_ne();
-            chunks.push_back(chunk_i);
-        }
-        reverse_ne();
-        return chunks;
-    }
-
-    void reverse_ne() {
-        int64_t new_ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
-        for (int i = 0; i < n_dims; i++) {
-            new_ne[i] = ne[n_dims - 1 - i];
-        }
-        for (int i = 0; i < n_dims; i++) {
-            ne[i] = new_ne[i];
-        }
-    }
-
-    std::string to_string() const {
-        std::stringstream ss;
-        const char* type_name = ggml_type_name(type);
-        if (is_f8_e4m3) {
-            type_name = "f8_e4m3";
-        } else if (is_f8_e5m2) {
-            type_name = "f8_e5m2";
-        } else if (is_f64) {
-            type_name = "f64";
-        } else if (is_i64) {
-            type_name = "i64";
-        }
-        ss << name << " | " << type_name << " | ";
-        ss << n_dims << " [";
-        for (int i = 0; i < SD_MAX_DIMS; i++) {
-            ss << ne[i];
-            if (i != SD_MAX_DIMS - 1) {
-                ss << ", ";
-            }
-        }
-        ss << "]";
-        return ss.str();
-    }
-};
-
-typedef std::function<bool(const TensorStorage&, ggml_tensor**)> on_new_tensor_cb_t;
-
 typedef OrderedMap<std::string, TensorStorage> String2TensorStorage;
+using TensorTypeRules = std::vector<std::pair<std::string, ggml_type>>;
+
+TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules);
 
 class ModelLoader {
 protected:
@@ -297,16 +201,10 @@ protected:
 
     void add_tensor_storage(const TensorStorage& tensor_storage);
 
-    bool parse_data_pkl(uint8_t* buffer,
-                        size_t buffer_size,
-                        zip_t* zip,
-                        std::string dir,
-                        size_t file_index,
-                        const std::string prefix);
-
     bool init_from_gguf_file(const std::string& file_path, const std::string& prefix = "");
     bool init_from_safetensors_file(const std::string& file_path, const std::string& prefix = "");
-    bool init_from_ckpt_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_torch_zip_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix = "");
     bool init_from_diffusers_file(const std::string& file_path, const std::string& prefix = "");
 
 public:
@@ -336,7 +234,6 @@ public:
         return names;
     }
 
-    bool save_to_gguf_file(const std::string& file_path, ggml_type type, const std::string& tensor_type_rules);
     bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type);
     int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT);
     ~ModelLoader() = default;

src/model_io/binary_io.h

@@ -0,0 +1,57 @@
+#ifndef __SD_MODEL_IO_BINARY_IO_H__
+#define __SD_MODEL_IO_BINARY_IO_H__
+
+#include <cstdint>
+#include <ostream>
+
+namespace model_io {
+
+    inline int32_t read_int(const uint8_t* buffer) {
+        uint32_t value = 0;
+        value |= static_cast<uint32_t>(buffer[3]) << 24;
+        value |= static_cast<uint32_t>(buffer[2]) << 16;
+        value |= static_cast<uint32_t>(buffer[1]) << 8;
+        value |= static_cast<uint32_t>(buffer[0]);
+        return static_cast<int32_t>(value);
+    }
+
+    inline uint16_t read_short(const uint8_t* buffer) {
+        uint16_t value = 0;
+        value |= static_cast<uint16_t>(buffer[1]) << 8;
+        value |= static_cast<uint16_t>(buffer[0]);
+        return value;
+    }
+
+    inline uint64_t read_u64(const uint8_t* buffer) {
+        uint64_t value = 0;
+        value |= static_cast<uint64_t>(buffer[7]) << 56;
+        value |= static_cast<uint64_t>(buffer[6]) << 48;
+        value |= static_cast<uint64_t>(buffer[5]) << 40;
+        value |= static_cast<uint64_t>(buffer[4]) << 32;
+        value |= static_cast<uint64_t>(buffer[3]) << 24;
+        value |= static_cast<uint64_t>(buffer[2]) << 16;
+        value |= static_cast<uint64_t>(buffer[1]) << 8;
+        value |= static_cast<uint64_t>(buffer[0]);
+        return value;
+    }
+
+    inline void write_u64(std::ostream& stream, uint64_t value) {
+        uint8_t buffer[8];
+        for (int i = 0; i < 8; ++i) {
+            buffer[i] = static_cast<uint8_t>((value >> (8 * i)) & 0xFF);
+        }
+        stream.write((const char*)buffer, sizeof(buffer));
+    }
+
+    inline int find_char(const uint8_t* buffer, int len, char c) {
+        for (int pos = 0; pos < len; pos++) {
+            if (buffer[pos] == (uint8_t)c) {
+                return pos;
+            }
+        }
+        return -1;
+    }
+
+}  // namespace model_io
+
+#endif  // __SD_MODEL_IO_BINARY_IO_H__

src/model_io/gguf_io.cpp

@@ -0,0 +1,123 @@
+#include "gguf_io.h"
+
+#include <cstdint>
+#include <fstream>
+#include <string>
+#include <vector>
+
+#include "gguf.h"
+#include "gguf_reader_ext.h"
+#include "util.h"
+
+static void set_error(std::string* error, const std::string& message) {
+    if (error != nullptr) {
+        *error = message;
+    }
+}
+
+bool is_gguf_file(const std::string& file_path) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        return false;
+    }
+
+    char magic[4];
+
+    file.read(magic, sizeof(magic));
+    if (!file) {
+        return false;
+    }
+    for (uint32_t i = 0; i < sizeof(magic); i++) {
+        if (magic[i] != GGUF_MAGIC[i]) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
+bool read_gguf_file(const std::string& file_path,
+                    std::vector<TensorStorage>& tensor_storages,
+                    std::string* error) {
+    tensor_storages.clear();
+
+    gguf_context* ctx_gguf_ = nullptr;
+    ggml_context* ctx_meta_ = nullptr;
+
+    ctx_gguf_ = gguf_init_from_file(file_path.c_str(), {true, &ctx_meta_});
+    if (!ctx_gguf_) {
+        GGUFReader gguf_reader;
+        if (!gguf_reader.load(file_path)) {
+            set_error(error, "failed to open '" + file_path + "' with GGUFReader");
+            return false;
+        }
+
+        size_t data_offset = gguf_reader.data_offset();
+        for (const auto& gguf_tensor_info : gguf_reader.tensors()) {
+            TensorStorage tensor_storage(
+                gguf_tensor_info.name,
+                gguf_tensor_info.type,
+                gguf_tensor_info.shape.data(),
+                static_cast<int>(gguf_tensor_info.shape.size()),
+                0,
+                data_offset + gguf_tensor_info.offset);
+
+            tensor_storages.push_back(tensor_storage);
+        }
+
+        return true;
+    }
+
+    int n_tensors = static_cast<int>(gguf_get_n_tensors(ctx_gguf_));
+
+    size_t data_offset = gguf_get_data_offset(ctx_gguf_);
+    for (int i = 0; i < n_tensors; i++) {
+        std::string name   = gguf_get_tensor_name(ctx_gguf_, i);
+        ggml_tensor* dummy = ggml_get_tensor(ctx_meta_, name.c_str());
+        size_t offset      = data_offset + gguf_get_tensor_offset(ctx_gguf_, i);
+
+        TensorStorage tensor_storage(name, dummy->type, dummy->ne, ggml_n_dims(dummy), 0, offset);
+
+        if (ggml_nbytes(dummy) != tensor_storage.nbytes()) {
+            gguf_free(ctx_gguf_);
+            ggml_free(ctx_meta_);
+            set_error(error, "size mismatch for tensor '" + name + "'");
+            return false;
+        }
+
+        tensor_storages.push_back(tensor_storage);
+    }
+
+    gguf_free(ctx_gguf_);
+    ggml_free(ctx_meta_);
+
+    return true;
+}
+
+bool write_gguf_file(const std::string& file_path,
+                     const std::vector<TensorWriteInfo>& tensors,
+                     std::string* error) {
+    gguf_context* gguf_ctx = gguf_init_empty();
+    if (gguf_ctx == nullptr) {
+        set_error(error, "gguf_init_empty failed");
+        return false;
+    }
+
+    for (const TensorWriteInfo& write_tensor : tensors) {
+        ggml_tensor* tensor = write_tensor.tensor;
+        if (tensor == nullptr) {
+            set_error(error, "null tensor cannot be written to GGUF");
+            gguf_free(gguf_ctx);
+            return false;
+        }
+        gguf_add_tensor(gguf_ctx, tensor);
+    }
+
+    LOG_INFO("trying to save tensors to %s", file_path.c_str());
+    bool success = gguf_write_to_file(gguf_ctx, file_path.c_str(), false);
+    if (!success) {
+        set_error(error, "failed to write GGUF file '" + file_path + "'");
+    }
+    gguf_free(gguf_ctx);
+    return success;
+}

src/model_io/gguf_io.h

@@ -0,0 +1,17 @@
+#ifndef __SD_MODEL_IO_GGUF_IO_H__
+#define __SD_MODEL_IO_GGUF_IO_H__
+
+#include <string>
+#include <vector>
+
+#include "tensor_storage.h"
+
+bool is_gguf_file(const std::string& file_path);
+bool read_gguf_file(const std::string& file_path,
+                    std::vector<TensorStorage>& tensor_storages,
+                    std::string* error = nullptr);
+bool write_gguf_file(const std::string& file_path,
+                     const std::vector<TensorWriteInfo>& tensors,
+                     std::string* error = nullptr);
+
+#endif  // __SD_MODEL_IO_GGUF_IO_H__

src/model_io/gguf_reader_ext.h

@@ -1,5 +1,5 @@
-#ifndef __GGUF_READER_HPP__
+#ifndef __SD_MODEL_IO_GGUF_READER_EXT_H__
-#define __GGUF_READER_HPP__
+#define __SD_MODEL_IO_GGUF_READER_EXT_H__
 
 #include <cstdint>
 #include <fstream>
@@ -231,4 +231,4 @@ public:
     size_t data_offset() const { return data_offset_; }
 };
 
-#endif  // __GGUF_READER_HPP__
+#endif  // __SD_MODEL_IO_GGUF_READER_EXT_H__

src/model_io/pickle_io.h

@@ -0,0 +1,21 @@
+#ifndef __SD_MODEL_IO_PICKLE_IO_H__
+#define __SD_MODEL_IO_PICKLE_IO_H__
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "tensor_storage.h"
+
+bool skip_pickle_object(const uint8_t* buffer, size_t buffer_size, size_t* object_size);
+bool pickle_object_is_torch_magic_number(const uint8_t* buffer, size_t buffer_size);
+bool parse_pickle_uint32_object(const uint8_t* buffer, size_t buffer_size, uint32_t* value);
+bool parse_torch_state_dict_pickle(const uint8_t* buffer,
+                                   size_t buffer_size,
+                                   std::vector<TensorStorage>& tensor_storages,
+                                   std::unordered_map<std::string, uint64_t>& storage_nbytes,
+                                   std::string* error = nullptr);
+
+#endif  // __SD_MODEL_IO_PICKLE_IO_H__

src/model_io/safetensors_io.cpp

@@ -0,0 +1,316 @@
+#include "safetensors_io.h"
+
+#include <cstdint>
+#include <exception>
+#include <fstream>
+#include <string>
+#include <vector>
+
+#include "binary_io.h"
+#include "json.hpp"
+#include "util.h"
+
+static constexpr size_t ST_HEADER_SIZE_LEN = 8;
+
+static void set_error(std::string* error, const std::string& message) {
+    if (error != nullptr) {
+        *error = message;
+    }
+}
+
+bool is_safetensors_file(const std::string& file_path) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        return false;
+    }
+
+    // get file size
+    file.seekg(0, file.end);
+    size_t file_size_ = file.tellg();
+    file.seekg(0, file.beg);
+
+    // read header size
+    if (file_size_ <= ST_HEADER_SIZE_LEN) {
+        return false;
+    }
+
+    uint8_t header_size_buf[ST_HEADER_SIZE_LEN];
+    file.read((char*)header_size_buf, ST_HEADER_SIZE_LEN);
+    if (!file) {
+        return false;
+    }
+
+    size_t header_size_ = model_io::read_u64(header_size_buf);
+    if (header_size_ >= file_size_ || header_size_ <= 2) {
+        return false;
+    }
+
+    // read header
+    std::vector<char> header_buf;
+    header_buf.resize(header_size_ + 1);
+    header_buf[header_size_] = '\0';
+    file.read(header_buf.data(), header_size_);
+    if (!file) {
+        return false;
+    }
+    try {
+        nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
+    } catch (const std::exception&) {
+        return false;
+    }
+    return true;
+}
+
+static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) {
+    ggml_type ttype = GGML_TYPE_COUNT;
+    if (dtype == "F16") {
+        ttype = GGML_TYPE_F16;
+    } else if (dtype == "BF16") {
+        ttype = GGML_TYPE_BF16;
+    } else if (dtype == "F32") {
+        ttype = GGML_TYPE_F32;
+    } else if (dtype == "F64") {
+        ttype = GGML_TYPE_F32;
+    } else if (dtype == "F8_E4M3") {
+        ttype = GGML_TYPE_F16;
+    } else if (dtype == "F8_E5M2") {
+        ttype = GGML_TYPE_F16;
+    } else if (dtype == "I32") {
+        ttype = GGML_TYPE_I32;
+    } else if (dtype == "I64") {
+        ttype = GGML_TYPE_I32;
+    }
+    return ttype;
+}
+
+// https://huggingface.co/docs/safetensors/index
+bool read_safetensors_file(const std::string& file_path,
+                           std::vector<TensorStorage>& tensor_storages,
+                           std::string* error) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        set_error(error, "failed to open '" + file_path + "'");
+        return false;
+    }
+
+    // get file size
+    file.seekg(0, file.end);
+    size_t file_size_ = file.tellg();
+    file.seekg(0, file.beg);
+
+    // read header size
+    if (file_size_ <= ST_HEADER_SIZE_LEN) {
+        set_error(error, "invalid safetensor file '" + file_path + "'");
+        return false;
+    }
+
+    uint8_t header_size_buf[ST_HEADER_SIZE_LEN];
+    file.read((char*)header_size_buf, ST_HEADER_SIZE_LEN);
+    if (!file) {
+        set_error(error, "read safetensors header size failed: '" + file_path + "'");
+        return false;
+    }
+
+    size_t header_size_ = model_io::read_u64(header_size_buf);
+    if (header_size_ >= file_size_) {
+        set_error(error, "invalid safetensor file '" + file_path + "'");
+        return false;
+    }
+
+    // read header
+    std::vector<char> header_buf;
+    header_buf.resize(header_size_ + 1);
+    header_buf[header_size_] = '\0';
+    file.read(header_buf.data(), header_size_);
+    if (!file) {
+        set_error(error, "read safetensors header failed: '" + file_path + "'");
+        return false;
+    }
+
+    nlohmann::json header_;
+    try {
+        header_ = nlohmann::json::parse(header_buf.data());
+    } catch (const std::exception&) {
+        set_error(error, "parsing safetensors header failed: '" + file_path + "'");
+        return false;
+    }
+
+    tensor_storages.clear();
+    for (auto& item : header_.items()) {
+        std::string name           = item.key();
+        nlohmann::json tensor_info = item.value();
+        // LOG_DEBUG("%s %s\n", name.c_str(), tensor_info.dump().c_str());
+
+        if (name == "__metadata__") {
+            continue;
+        }
+
+        std::string dtype    = tensor_info["dtype"];
+        nlohmann::json shape = tensor_info["shape"];
+
+        if (dtype == "U8") {
+            continue;
+        }
+
+        size_t begin = tensor_info["data_offsets"][0].get<size_t>();
+        size_t end   = tensor_info["data_offsets"][1].get<size_t>();
+
+        ggml_type type = safetensors_dtype_to_ggml_type(dtype);
+        if (type == GGML_TYPE_COUNT) {
+            set_error(error, "unsupported dtype '" + dtype + "' (tensor '" + name + "')");
+            return false;
+        }
+
+        if (shape.size() > SD_MAX_DIMS) {
+            set_error(error, "invalid tensor '" + name + "'");
+            return false;
+        }
+
+        int n_dims              = (int)shape.size();
+        int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
+        for (int i = 0; i < n_dims; i++) {
+            ne[i] = shape[i].get<int64_t>();
+        }
+
+        if (n_dims == 5) {
+            n_dims = 4;
+            ne[0]  = ne[0] * ne[1];
+            ne[1]  = ne[2];
+            ne[2]  = ne[3];
+            ne[3]  = ne[4];
+        }
+
+        // ggml_n_dims returns 1 for scalars
+        if (n_dims == 0) {
+            n_dims = 1;
+        }
+
+        TensorStorage tensor_storage(name, type, ne, n_dims, 0, ST_HEADER_SIZE_LEN + header_size_ + begin);
+        tensor_storage.reverse_ne();
+
+        size_t tensor_data_size = end - begin;
+
+        bool tensor_size_ok;
+        if (dtype == "F8_E4M3") {
+            tensor_storage.is_f8_e4m3 = true;
+            // f8 -> f16
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
+        } else if (dtype == "F8_E5M2") {
+            tensor_storage.is_f8_e5m2 = true;
+            // f8 -> f16
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
+        } else if (dtype == "F64") {
+            tensor_storage.is_f64 = true;
+            // f64 -> f32
+            tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
+        } else if (dtype == "I64") {
+            tensor_storage.is_i64 = true;
+            // i64 -> i32
+            tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
+        } else {
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size);
+        }
+        if (!tensor_size_ok) {
+            set_error(error, "size mismatch for tensor '" + name + "' (" + dtype + ")");
+            return false;
+        }
+
+        tensor_storages.push_back(tensor_storage);
+
+        // LOG_DEBUG("%s %s", tensor_storage.to_string().c_str(), dtype.c_str());
+    }
+
+    return true;
+}
+
+static bool ggml_type_to_safetensors_dtype(ggml_type type, std::string* dtype) {
+    switch (type) {
+        case GGML_TYPE_F16:
+            *dtype = "F16";
+            return true;
+        case GGML_TYPE_BF16:
+            *dtype = "BF16";
+            return true;
+        case GGML_TYPE_F32:
+            *dtype = "F32";
+            return true;
+        case GGML_TYPE_I32:
+            *dtype = "I32";
+            return true;
+        default:
+            return false;
+    }
+}
+
+bool write_safetensors_file(const std::string& file_path,
+                            const std::vector<TensorWriteInfo>& tensors,
+                            std::string* error) {
+    nlohmann::ordered_json header = nlohmann::ordered_json::object();
+
+    uint64_t data_offset = 0;
+    for (const TensorWriteInfo& write_tensor : tensors) {
+        ggml_tensor* tensor = write_tensor.tensor;
+        if (tensor == nullptr) {
+            set_error(error, "null tensor cannot be written to safetensors");
+            return false;
+        }
+
+        const std::string name = ggml_get_name(tensor);
+        std::string dtype;
+        if (!ggml_type_to_safetensors_dtype(tensor->type, &dtype)) {
+            set_error(error,
+                      "unsupported safetensors dtype '" + std::string(ggml_type_name(tensor->type)) +
+                          "' for tensor '" + name + "'");
+            return false;
+        }
+
+        const uint64_t tensor_nbytes = ggml_nbytes(tensor);
+
+        nlohmann::ordered_json json_tensor_info = nlohmann::ordered_json::object();
+        json_tensor_info["dtype"]               = dtype;
+
+        nlohmann::ordered_json shape = nlohmann::ordered_json::array();
+        for (int i = 0; i < write_tensor.n_dims; ++i) {
+            shape.push_back(write_tensor.ne[write_tensor.n_dims - 1 - i]);
+        }
+        json_tensor_info["shape"] = shape;
+
+        nlohmann::ordered_json data_offsets = nlohmann::ordered_json::array();
+        data_offsets.push_back(data_offset);
+        data_offsets.push_back(data_offset + tensor_nbytes);
+        json_tensor_info["data_offsets"] = data_offsets;
+
+        header[name] = json_tensor_info;
+        data_offset += tensor_nbytes;
+    }
+
+    const std::string header_str = header.dump();
+
+    std::ofstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        set_error(error, "failed to open '" + file_path + "' for writing");
+        return false;
+    }
+
+    LOG_INFO("trying to save tensors to %s", file_path.c_str());
+    model_io::write_u64(file, header_str.size());
+    file.write(header_str.data(), header_str.size());
+    if (!file) {
+        set_error(error, "failed to write safetensors header to '" + file_path + "'");
+        return false;
+    }
+
+    for (const TensorWriteInfo& write_tensor : tensors) {
+        ggml_tensor* tensor        = write_tensor.tensor;
+        const std::string name     = ggml_get_name(tensor);
+        const size_t tensor_nbytes = ggml_nbytes(tensor);
+        file.write((const char*)tensor->data, tensor_nbytes);
+        if (!file) {
+            set_error(error,
+                      "failed to write tensor '" + name + "' to '" + file_path + "'");
+            return false;
+        }
+    }
+
+    return true;
+}

src/model_io/safetensors_io.h

@@ -0,0 +1,17 @@
+#ifndef __SD_MODEL_IO_SAFETENSORS_IO_H__
+#define __SD_MODEL_IO_SAFETENSORS_IO_H__
+
+#include <string>
+#include <vector>
+
+#include "tensor_storage.h"
+
+bool is_safetensors_file(const std::string& file_path);
+bool read_safetensors_file(const std::string& file_path,
+                           std::vector<TensorStorage>& tensor_storages,
+                           std::string* error = nullptr);
+bool write_safetensors_file(const std::string& file_path,
+                            const std::vector<TensorWriteInfo>& tensors,
+                            std::string* error = nullptr);
+
+#endif  // __SD_MODEL_IO_SAFETENSORS_IO_H__

src/model_io/tensor_storage.h

@@ -0,0 +1,132 @@
+#ifndef __SD_TENSOR_STORAGE_H__
+#define __SD_TENSOR_STORAGE_H__
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <sstream>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "ggml.h"
+
+#define SD_MAX_DIMS 5
+
+struct TensorStorage {
+    std::string name;
+    ggml_type type          = GGML_TYPE_F32;
+    ggml_type expected_type = GGML_TYPE_COUNT;
+    bool is_f8_e4m3         = false;
+    bool is_f8_e5m2         = false;
+    bool is_f64             = false;
+    bool is_i64             = false;
+    int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
+    int n_dims              = 0;
+
+    std::string storage_key;
+    size_t file_index = 0;
+    int index_in_zip  = -1;  // >= means stored in a zip file
+    uint64_t offset   = 0;   // offset in file
+
+    TensorStorage() = default;
+
+    TensorStorage(std::string name, ggml_type type, const int64_t* ne, int n_dims, size_t file_index, size_t offset = 0)
+        : name(std::move(name)), type(type), n_dims(n_dims), file_index(file_index), offset(offset) {
+        for (int i = 0; i < n_dims; i++) {
+            this->ne[i] = ne[i];
+        }
+    }
+
+    int64_t nelements() const {
+        int64_t n = 1;
+        for (int i = 0; i < SD_MAX_DIMS; i++) {
+            n *= ne[i];
+        }
+        return n;
+    }
+
+    int64_t nbytes() const {
+        return nelements() * ggml_type_size(type) / ggml_blck_size(type);
+    }
+
+    int64_t nbytes_to_read() const {
+        if (is_f8_e4m3 || is_f8_e5m2) {
+            return nbytes() / 2;
+        } else if (is_f64 || is_i64) {
+            return nbytes() * 2;
+        } else {
+            return nbytes();
+        }
+    }
+
+    void unsqueeze() {
+        if (n_dims == 2) {
+            n_dims = 4;
+            ne[3]  = ne[1];
+            ne[2]  = ne[0];
+            ne[1]  = 1;
+            ne[0]  = 1;
+        }
+    }
+
+    std::vector<TensorStorage> chunk(size_t n) {
+        std::vector<TensorStorage> chunks;
+        uint64_t chunk_size = nbytes_to_read() / n;
+        // printf("%d/%d\n", chunk_size, nbytes_to_read());
+        reverse_ne();
+        for (size_t i = 0; i < n; i++) {
+            TensorStorage chunk_i = *this;
+            chunk_i.ne[0]         = ne[0] / n;
+            chunk_i.offset        = offset + i * chunk_size;
+            chunk_i.reverse_ne();
+            chunks.push_back(chunk_i);
+        }
+        reverse_ne();
+        return chunks;
+    }
+
+    void reverse_ne() {
+        int64_t new_ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
+        for (int i = 0; i < n_dims; i++) {
+            new_ne[i] = ne[n_dims - 1 - i];
+        }
+        for (int i = 0; i < n_dims; i++) {
+            ne[i] = new_ne[i];
+        }
+    }
+
+    std::string to_string() const {
+        std::stringstream ss;
+        const char* type_name = ggml_type_name(type);
+        if (is_f8_e4m3) {
+            type_name = "f8_e4m3";
+        } else if (is_f8_e5m2) {
+            type_name = "f8_e5m2";
+        } else if (is_f64) {
+            type_name = "f64";
+        } else if (is_i64) {
+            type_name = "i64";
+        }
+        ss << name << " | " << type_name << " | ";
+        ss << n_dims << " [";
+        for (int i = 0; i < SD_MAX_DIMS; i++) {
+            ss << ne[i];
+            if (i != SD_MAX_DIMS - 1) {
+                ss << ", ";
+            }
+        }
+        ss << "]";
+        return ss.str();
+    }
+};
+
+struct TensorWriteInfo {
+    int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
+    int n_dims              = 0;
+    ggml_tensor* tensor     = nullptr;
+};
+
+typedef std::function<bool(const TensorStorage&, ggml_tensor**)> on_new_tensor_cb_t;
+
+#endif  // __SD_TENSOR_STORAGE_H__

src/model_io/torch_legacy_io.cpp

@@ -0,0 +1,252 @@
+#include "torch_legacy_io.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <fstream>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "pickle_io.h"
+#include "util.h"
+
+// torch.save format background:
+//
+//   - Before PyTorch 1.6.0, torch.save used this legacy non-zip format by
+//     default.
+//   - Since PyTorch 1.6.0, torch.save defaults to an uncompressed ZIP64 archive
+//     containing data.pkl, data/, version, and, since PyTorch 2.1.0, byteorder.
+//   - The old format can still be produced explicitly with:
+//       torch.save(obj, path, _use_new_zipfile_serialization=False)
+//
+// Whether obj is a state_dict or a whole nn.Module does not change the outer
+// container format selected by torch.save. It changes the pickled object inside:
+//
+//   - state_dict: usually an OrderedDict[str, Tensor]. pickle_io.cpp supports a
+//     restricted subset of this layout because tensor metadata and raw storages
+//     can be recovered without executing pickle callables.
+//   - whole module/checkpoint object: arbitrary Python object graph. This may
+//     require importing user classes and executing pickle GLOBAL/REDUCE rebuild
+//     logic, so it is intentionally not supported here.
+//
+// Legacy non-zip PyTorch files are not a single pickle object:
+//
+//   1. pickle object: PyTorch legacy magic number
+//   2. pickle object: legacy protocol version, expected to be 1001
+//   3. pickle object: sys_info metadata, ignored by this reader
+//   4. pickle object: state_dict metadata, parsed by pickle_io.cpp
+//   5. pickle object: serialized storage key list, skipped here
+//   6. raw storage data payloads
+//      - PyTorch writes storages after the pickles, ordered by storage key
+//      - each storage has an 8-byte legacy storage header followed by raw bytes
+static constexpr size_t LEGACY_STORAGE_HEADER_SIZE = 8;
+
+static void set_error(std::string* error, const std::string& message) {
+    if (error != nullptr) {
+        *error = message;
+    }
+}
+
+static std::string bytes_to_hex(const std::vector<uint8_t>& bytes) {
+    static const char* hex = "0123456789ABCDEF";
+    std::string result;
+    result.reserve(bytes.size() * 3);
+    for (size_t i = 0; i < bytes.size(); ++i) {
+        if (i > 0) {
+            result.push_back('-');
+        }
+        result.push_back(hex[(bytes[i] >> 4) & 0x0F]);
+        result.push_back(hex[bytes[i] & 0x0F]);
+    }
+    return result;
+}
+
+static bool is_probably_tar_file(const std::vector<uint8_t>& header) {
+    return header.size() >= 262 &&
+           header[257] == 'u' &&
+           header[258] == 's' &&
+           header[259] == 't' &&
+           header[260] == 'a' &&
+           header[261] == 'r';
+}
+
+static std::string torch_legacy_diagnostics(const std::string& file_path, const std::vector<uint8_t>& buffer) {
+    if (!ends_with(file_path, ".pt") && !ends_with(file_path, ".pth")) {
+        return "";
+    }
+    if (buffer.empty()) {
+        return "unsupported PyTorch file '" + file_path + "': empty file";
+    }
+
+    size_t short_len = std::min<size_t>(buffer.size(), 32);
+    std::vector<uint8_t> short_header(buffer.begin(), buffer.begin() + short_len);
+    const bool raw_pickle = buffer[0] == 0x80;
+    const bool tar_file   = is_probably_tar_file(buffer);
+
+    std::string message = "unsupported PyTorch file '" + file_path + "': first bytes " +
+                          bytes_to_hex(short_header) +
+                          ", raw_pickle=" + (raw_pickle ? "true" : "false") +
+                          ", tar=" + (tar_file ? "true" : "false");
+    if (raw_pickle) {
+        message += "; raw pickle did not match the restricted state_dict layouts currently supported";
+    } else if (tar_file) {
+        message += "; legacy tar PyTorch checkpoints are not supported yet";
+    }
+    return message;
+}
+
+bool read_torch_legacy_file(const std::string& file_path,
+                            std::vector<TensorStorage>& tensor_storages,
+                            std::string* error) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        set_error(error, "failed to open '" + file_path + "'");
+        return false;
+    }
+
+    file.seekg(0, file.end);
+    size_t file_size = (size_t)file.tellg();
+    file.seekg(0, file.beg);
+    if (file_size == 0) {
+        set_error(error, "empty file '" + file_path + "'");
+        return false;
+    }
+
+    std::vector<uint8_t> buffer(file_size);
+    file.read((char*)buffer.data(), file_size);
+    if (!file) {
+        set_error(error, "failed to read '" + file_path + "'");
+        return false;
+    }
+
+    auto finalize_tensor_offsets = [&](size_t storage_data_offset,
+                                       const std::unordered_map<std::string, uint64_t>& legacy_storage_map) -> bool {
+        if (storage_data_offset > file_size) {
+            return false;
+        }
+
+        std::vector<std::string> storage_keys;
+        storage_keys.reserve(legacy_storage_map.size());
+        for (const auto& [storage_key, _] : legacy_storage_map) {
+            storage_keys.push_back(storage_key);
+        }
+        std::sort(storage_keys.begin(), storage_keys.end());
+
+        std::unordered_map<std::string, uint64_t> storage_offsets;
+        uint64_t current_offset = storage_data_offset;
+        for (const auto& storage_key : storage_keys) {
+            auto it = legacy_storage_map.find(storage_key);
+            if (it == legacy_storage_map.end()) {
+                return false;
+            }
+            if (current_offset + LEGACY_STORAGE_HEADER_SIZE + it->second > file_size) {
+                return false;
+            }
+            storage_offsets[storage_key] = current_offset + LEGACY_STORAGE_HEADER_SIZE;
+            current_offset += LEGACY_STORAGE_HEADER_SIZE + it->second;
+        }
+
+        for (auto& tensor_storage : tensor_storages) {
+            if (tensor_storage.storage_key.empty()) {
+                continue;
+            }
+
+            auto it_offset = storage_offsets.find(tensor_storage.storage_key);
+            auto it_size   = legacy_storage_map.find(tensor_storage.storage_key);
+            if (it_offset == storage_offsets.end() || it_size == legacy_storage_map.end()) {
+                return false;
+            }
+
+            uint64_t base_offset    = it_offset->second;
+            uint64_t storage_nbytes = it_size->second;
+            uint64_t tensor_nbytes  = tensor_storage.nbytes_to_read();
+            if (tensor_storage.offset + tensor_nbytes > storage_nbytes) {
+                return false;
+            }
+
+            tensor_storage.offset = base_offset + tensor_storage.offset;
+            tensor_storage.storage_key.clear();
+        }
+
+        return true;
+    };
+
+    auto parse_state_dict_at = [&](size_t state_dict_offset, size_t state_dict_size, size_t* storage_data_offset) -> bool {
+        tensor_storages.clear();
+        std::unordered_map<std::string, uint64_t> legacy_storage_map;
+        if (!parse_torch_state_dict_pickle(buffer.data() + state_dict_offset,
+                                           state_dict_size,
+                                           tensor_storages,
+                                           legacy_storage_map,
+                                           error)) {
+            return false;
+        }
+
+        size_t offset_after_state_dict = state_dict_offset + state_dict_size;
+        size_t storage_keys_size       = 0;
+        if (!skip_pickle_object(buffer.data() + offset_after_state_dict,
+                                buffer.size() - offset_after_state_dict,
+                                &storage_keys_size)) {
+            return false;
+        }
+
+        *storage_data_offset = offset_after_state_dict + storage_keys_size;
+        return finalize_tensor_offsets(*storage_data_offset, legacy_storage_map);
+    };
+
+    size_t object_size_1 = 0;
+    size_t offset        = 0;
+
+    if (skip_pickle_object(buffer.data(), buffer.size(), &object_size_1) &&
+        pickle_object_is_torch_magic_number(buffer.data(), object_size_1)) {
+        offset += object_size_1;
+
+        size_t object_size_2 = 0;
+        if (!skip_pickle_object(buffer.data() + offset, buffer.size() - offset, &object_size_2)) {
+            set_error(error, torch_legacy_diagnostics(file_path, buffer));
+            return false;
+        }
+        uint32_t protocol_version = 0;
+        if (!parse_pickle_uint32_object(buffer.data() + offset, object_size_2, &protocol_version) || protocol_version != 1001) {
+            set_error(error, torch_legacy_diagnostics(file_path, buffer));
+            return false;
+        }
+        offset += object_size_2;
+
+        size_t object_size_3 = 0;
+        if (!skip_pickle_object(buffer.data() + offset, buffer.size() - offset, &object_size_3)) {
+            set_error(error, torch_legacy_diagnostics(file_path, buffer));
+            return false;
+        }
+        offset += object_size_3;
+
+        size_t state_dict_size = 0;
+        if (!skip_pickle_object(buffer.data() + offset, buffer.size() - offset, &state_dict_size)) {
+            set_error(error, torch_legacy_diagnostics(file_path, buffer));
+            return false;
+        }
+
+        size_t storage_data_offset = 0;
+        if (parse_state_dict_at(offset, state_dict_size, &storage_data_offset)) {
+            return true;
+        }
+
+        if (error != nullptr && error->empty()) {
+            set_error(error, torch_legacy_diagnostics(file_path, buffer));
+        }
+        return false;
+    }
+
+    size_t state_dict_size = 0;
+    if (skip_pickle_object(buffer.data(), buffer.size(), &state_dict_size)) {
+        size_t storage_data_offset = 0;
+        if (parse_state_dict_at(0, state_dict_size, &storage_data_offset)) {
+            return true;
+        }
+    }
+
+    if (error != nullptr && error->empty()) {
+        set_error(error, torch_legacy_diagnostics(file_path, buffer));
+    }
+    return false;
+}

src/model_io/torch_legacy_io.h

@@ -0,0 +1,13 @@
+#ifndef __SD_MODEL_IO_TORCH_LEGACY_IO_H__
+#define __SD_MODEL_IO_TORCH_LEGACY_IO_H__
+
+#include <string>
+#include <vector>
+
+#include "tensor_storage.h"
+
+bool read_torch_legacy_file(const std::string& file_path,
+                            std::vector<TensorStorage>& tensor_storages,
+                            std::string* error = nullptr);
+
+#endif  // __SD_MODEL_IO_TORCH_LEGACY_IO_H__

src/model_io/torch_zip_io.cpp

@@ -0,0 +1,140 @@
+#include "torch_zip_io.h"
+
+#include <cstdint>
+#include <cstdlib>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "pickle_io.h"
+
+#include "zip.h"
+
+static void set_error(std::string* error, const std::string& message) {
+    if (error != nullptr) {
+        *error = message;
+    }
+}
+
+bool is_torch_zip_file(const std::string& file_path) {
+    zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
+    if (zip == nullptr) {
+        return false;
+    }
+    zip_close(zip);
+    return true;
+}
+
+static bool find_zip_entry(zip_t* zip, const std::string& entry_name, int* index, uint64_t* size) {
+    size_t n = zip_entries_total(zip);
+    for (size_t i = 0; i < n; ++i) {
+        zip_entry_openbyindex(zip, i);
+        std::string name = zip_entry_name(zip);
+        if (name == entry_name) {
+            *index = (int)i;
+            *size  = zip_entry_size(zip);
+            zip_entry_close(zip);
+            return true;
+        }
+        zip_entry_close(zip);
+    }
+    return false;
+}
+
+static bool parse_zip_data_pkl(const uint8_t* buffer,
+                               size_t buffer_size,
+                               zip_t* zip,
+                               const std::string& dir,
+                               std::vector<TensorStorage>& tensor_storages,
+                               std::string* error) {
+    std::vector<TensorStorage> parsed_tensors;
+    std::unordered_map<std::string, uint64_t> storage_nbytes;
+    if (!parse_torch_state_dict_pickle(buffer, buffer_size, parsed_tensors, storage_nbytes, error)) {
+        if (error != nullptr && error->empty()) {
+            *error = "failed to parse torch zip pickle metadata";
+        }
+        return false;
+    }
+
+    for (auto& tensor_storage : parsed_tensors) {
+        if (tensor_storage.storage_key.empty()) {
+            set_error(error, "tensor '" + tensor_storage.name + "' has no storage key");
+            return false;
+        }
+
+        const std::string entry_name = dir + "data/" + tensor_storage.storage_key;
+        int zip_index                = -1;
+        uint64_t entry_size          = 0;
+        if (!find_zip_entry(zip, entry_name, &zip_index, &entry_size)) {
+            set_error(error, "storage entry '" + entry_name + "' was not found");
+            return false;
+        }
+
+        auto it_storage_size = storage_nbytes.find(tensor_storage.storage_key);
+        if (it_storage_size != storage_nbytes.end() && entry_size < it_storage_size->second) {
+            set_error(error, "storage entry '" + entry_name + "' is smaller than pickle metadata");
+            return false;
+        }
+
+        uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
+        if (tensor_storage.offset + tensor_nbytes > entry_size) {
+            set_error(error, "tensor '" + tensor_storage.name + "' exceeds storage entry '" + entry_name + "'");
+            return false;
+        }
+
+        tensor_storage.index_in_zip = zip_index;
+        tensor_storage.storage_key.clear();
+        tensor_storages.push_back(tensor_storage);
+    }
+
+    return true;
+}
+
+bool read_torch_zip_file(const std::string& file_path,
+                         std::vector<TensorStorage>& tensor_storages,
+                         std::string* error) {
+    zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
+    if (zip == nullptr) {
+        set_error(error, "failed to open '" + file_path + "'");
+        return false;
+    }
+
+    tensor_storages.clear();
+    bool success        = true;
+    bool found_data_pkl = false;
+    int n               = (int)zip_entries_total(zip);
+    for (int i = 0; i < n; ++i) {
+        zip_entry_openbyindex(zip, i);
+        std::string name = zip_entry_name(zip);
+        size_t pos       = name.find("data.pkl");
+        if (pos != std::string::npos) {
+            found_data_pkl  = true;
+            std::string dir = name.substr(0, pos);
+            void* pkl_data  = nullptr;
+            size_t pkl_size = 0;
+            zip_entry_read(zip, &pkl_data, &pkl_size);
+
+            if (pkl_data == nullptr || pkl_size == 0) {
+                set_error(error, "failed to read '" + name + "' from '" + file_path + "'");
+                success = false;
+            } else if (!parse_zip_data_pkl((const uint8_t*)pkl_data, pkl_size, zip, dir, tensor_storages, error)) {
+                success = false;
+            }
+
+            free(pkl_data);
+        }
+        zip_entry_close(zip);
+
+        if (!success) {
+            break;
+        }
+    }
+
+    if (success && !found_data_pkl) {
+        set_error(error, "data.pkl was not found in '" + file_path + "'");
+        success = false;
+    }
+
+    zip_close(zip);
+    return success;
+}

src/model_io/torch_zip_io.h

@@ -0,0 +1,14 @@
+#ifndef __SD_MODEL_IO_TORCH_ZIP_IO_H__
+#define __SD_MODEL_IO_TORCH_ZIP_IO_H__
+
+#include <string>
+#include <vector>
+
+#include "tensor_storage.h"
+
+bool is_torch_zip_file(const std::string& file_path);
+bool read_torch_zip_file(const std::string& file_path,
+                         std::vector<TensorStorage>& tensor_storages,
+                         std::string* error = nullptr);
+
+#endif  // __SD_MODEL_IO_TORCH_ZIP_IO_H__

src/name_conversion.cpp

@@ -1120,7 +1120,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);
-                if (version == VERSION_SDXS) {
+                if (version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) {
                     name = "tae." + name;
                 } else {
                     name = prefix + name;

src/rope.hpp

@@ -7,6 +7,11 @@
 #include "ggml_extend.hpp"
 
 namespace Rope {
+    enum class EmbedNDLayout {
+        Matrix,
+        ErnieImage,
+    };
+
     template <class T>
     __STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
         std::vector<T> result(num);
@@ -169,7 +174,8 @@ namespace Rope {
                                                   int bs,
                                                   const std::vector<float>& axis_thetas,
                                                   const std::vector<int>& axes_dim,
-                                                  const std::vector<std::vector<int>>& wrap_dims = {}) {
+                                                  const std::vector<std::vector<int>>& wrap_dims = {},
+                                                  EmbedNDLayout layout                           = EmbedNDLayout::Matrix) {
         std::vector<std::vector<float>> trans_ids = transpose(ids);
         size_t pos_len                            = ids.size() / bs;
         size_t num_axes                           = axes_dim.size();
@@ -204,6 +210,24 @@ namespace Rope {
             offset += rope_emb[0].size();
         }
 
+        if (layout == EmbedNDLayout::ErnieImage) {
+            int head_dim = emb_dim * 2;
+            std::vector<float> ernie_emb(bs * pos_len * head_dim * 2, 0.0f);
+            for (size_t pos_idx = 0; pos_idx < bs * pos_len; ++pos_idx) {
+                for (int i = 0; i < emb_dim; ++i) {
+                    float cos_val             = emb[pos_idx][4 * i];
+                    float sin_val             = emb[pos_idx][4 * i + 2];
+                    size_t cos_offset         = pos_idx * head_dim + 2 * i;
+                    size_t sin_offset         = bs * pos_len * head_dim + cos_offset;
+                    ernie_emb[cos_offset]     = cos_val;
+                    ernie_emb[cos_offset + 1] = cos_val;
+                    ernie_emb[sin_offset]     = sin_val;
+                    ernie_emb[sin_offset + 1] = sin_val;
+                }
+            }
+            return ernie_emb;
+        }
+
         return flatten(emb);
     }
 
@@ -211,9 +235,10 @@ namespace Rope {
                                                   int bs,
                                                   float theta,
                                                   const std::vector<int>& axes_dim,
-                                                  const std::vector<std::vector<int>>& wrap_dims = {}) {
+                                                  const std::vector<std::vector<int>>& wrap_dims = {},
+                                                  EmbedNDLayout layout                           = EmbedNDLayout::Matrix) {
         std::vector<float> axis_thetas(axes_dim.size(), theta);
-        return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims);
+        return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout);
     }
 
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
@@ -437,6 +462,74 @@ namespace Rope {
         return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
     }
 
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_ernie_image_ids(int h,
+                                                                          int w,
+                                                                          int patch_size,
+                                                                          int bs,
+                                                                          int context_len) {
+        int h_len = h / patch_size;
+        int w_len = w / patch_size;
+
+        std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(3, 0.0f));
+        std::vector<float> h_ids = linspace<float>(0.f, static_cast<float>(h_len - 1), h_len);
+        std::vector<float> w_ids = linspace<float>(0.f, static_cast<float>(w_len - 1), w_len);
+        for (int i = 0; i < h_len; ++i) {
+            for (int j = 0; j < w_len; ++j) {
+                img_ids[i * w_len + j][0] = static_cast<float>(context_len);
+                img_ids[i * w_len + j][1] = h_ids[i];
+                img_ids[i * w_len + j][2] = w_ids[j];
+            }
+        }
+
+        std::vector<std::vector<float>> img_ids_repeated(bs * img_ids.size(), std::vector<float>(3, 0.0f));
+        for (int i = 0; i < bs; ++i) {
+            for (int j = 0; j < static_cast<int>(img_ids.size()); ++j) {
+                img_ids_repeated[i * img_ids.size() + j] = img_ids[j];
+            }
+        }
+
+        std::vector<std::vector<float>> txt_ids(bs * context_len, std::vector<float>(3, 0.0f));
+        for (int i = 0; i < bs; ++i) {
+            for (int j = 0; j < context_len; ++j) {
+                txt_ids[i * context_len + j][0] = static_cast<float>(j);
+            }
+        }
+
+        return concat_ids(img_ids_repeated, txt_ids, bs);
+    }
+
+    __STATIC_INLINE__ std::vector<float> gen_ernie_image_pe(int h,
+                                                            int w,
+                                                            int patch_size,
+                                                            int bs,
+                                                            int context_len,
+                                                            int theta,
+                                                            bool circular_h,
+                                                            bool circular_w,
+                                                            const std::vector<int>& axes_dim) {
+        std::vector<std::vector<float>> ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len);
+        std::vector<std::vector<int>> wrap_dims;
+        if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
+            int h_len = h / patch_size;
+            int w_len = w / patch_size;
+            if (h_len > 0 && w_len > 0) {
+                size_t pos_len = ids.size() / bs;
+                wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
+                const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
+                for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
+                    if (circular_h) {
+                        wrap_dims[1][token_i] = h_len;
+                    }
+                    if (circular_w) {
+                        wrap_dims[2][token_i] = w_len;
+                    }
+                }
+            }
+        }
+
+        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage);
+    }
+
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
                                                                   int h,
                                                                   int w,

src/stable-diffusion.cpp

@@ -17,6 +17,7 @@
 #include "pmid.hpp"
 #include "sample-cache.h"
 #include "tae.hpp"
+#include "upscaler.h"
 #include "vae.hpp"
 
 #include "latent-preview.h"
@@ -30,7 +31,8 @@ const char* model_version_to_str[] = {
     "SD 2.x",
     "SD 2.x Inpaint",
     "SD 2.x Tiny UNet",
-    "SDXS",
+    "SDXS (512-DS)",
+    "SDXS (09)",
     "SDXL",
     "SDXL Inpaint",
     "SDXL Instruct-Pix2Pix",
@@ -52,6 +54,7 @@ const char* model_version_to_str[] = {
     "Flux.2 klein",
     "Z-Image",
     "Ovis Image",
+    "Ernie Image",
 };
 
 const char* sampling_methods_str[] = {
@@ -69,6 +72,7 @@ const char* sampling_methods_str[] = {
     "TCD",
     "Res Multistep",
     "Res 2s",
+    "ER-SDE",
 };
 
 /*================================================== Helper Functions ================================================*/
@@ -413,7 +417,7 @@ public:
         }
 
         bool tae_preview_only = sd_ctx_params->tae_preview_only;
-        if (version == VERSION_SDXS) {
+        if (version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) {
             tae_preview_only = false;
             use_tae          = true;
         }
@@ -551,6 +555,15 @@ public:
                                                                 tensor_storage_map,
                                                                 "model.diffusion_model",
                                                                 version);
+            } else if (sd_version_is_ernie_image(version)) {
+                cond_stage_model = std::make_shared<LLMEmbedder>(clip_backend,
+                                                                 offload_params_to_cpu,
+                                                                 tensor_storage_map,
+                                                                 version);
+                diffusion_model  = std::make_shared<ErnieImageModel>(backend,
+                                                                    offload_params_to_cpu,
+                                                                    tensor_storage_map,
+                                                                    "model.diffusion_model");
             } else {  // SD1.x SD2.x SDXL
                 std::map<std::string, std::string> embbeding_map;
                 for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) {
@@ -819,6 +832,10 @@ public:
         if (version == VERSION_SVD) {
             ignore_tensors.insert("conditioner.embedders.3");
         }
+        if (sd_version_is_ernie_image(version)) {
+            ignore_tensors.insert("text_encoders.llm.vision_tower.");
+            ignore_tensors.insert("text_encoders.llm.multi_modal_projector.");
+        }
         bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads, sd_ctx_params->enable_mmap);
         if (!success) {
             LOG_ERROR("load tensors from model loader failed");
@@ -922,10 +939,13 @@ public:
                            sd_version_is_wan(version) ||
                            sd_version_is_qwen_image(version) ||
                            sd_version_is_anima(version) ||
+                           sd_version_is_ernie_image(version) ||
                            sd_version_is_z_image(version)) {
                     pred_type = FLOW_PRED;
                     if (sd_version_is_wan(version)) {
                         default_flow_shift = 5.f;
+                    } else if (sd_version_is_ernie_image(version)) {
+                        default_flow_shift = 4.f;
                     } else {
                         default_flow_shift = 3.f;
                     }
@@ -1395,7 +1415,7 @@ public:
             uint32_t dim                     = is_video ? static_cast<uint32_t>(latents.shape()[3]) : static_cast<uint32_t>(latents.shape()[2]);
 
             if (dim == 128) {
-                if (sd_version_is_flux2(version)) {
+                if (sd_version_uses_flux2_vae(version)) {
                     latent_rgb_proj = flux2_latent_rgb_proj;
                     latent_rgb_bias = flux2_latent_rgb_bias;
                     patch_sz        = 2;
@@ -1844,7 +1864,7 @@ public:
                 latent_channel = 48;
             } else if (version == VERSION_CHROMA_RADIANCE) {
                 latent_channel = 3;
-            } else if (sd_version_is_flux2(version)) {
+            } else if (sd_version_uses_flux2_vae(version)) {
                 latent_channel = 128;
             } else {
                 latent_channel = 16;
@@ -1975,6 +1995,7 @@ const char* sample_method_to_str[] = {
     "tcd",
     "res_multistep",
     "res_2s",
+    "er_sde",
 };
 
 const char* sd_sample_method_name(enum sample_method_t sample_method) {
@@ -2093,6 +2114,35 @@ enum lora_apply_mode_t str_to_lora_apply_mode(const char* str) {
     return LORA_APPLY_MODE_COUNT;
 }
 
+const char* hires_upscaler_to_str[] = {
+    "None",
+    "Latent",
+    "Latent (nearest)",
+    "Latent (nearest-exact)",
+    "Latent (antialiased)",
+    "Latent (bicubic)",
+    "Latent (bicubic antialiased)",
+    "Lanczos",
+    "Nearest",
+    "Model",
+};
+
+const char* sd_hires_upscaler_name(enum sd_hires_upscaler_t upscaler) {
+    if (upscaler >= SD_HIRES_UPSCALER_NONE && upscaler < SD_HIRES_UPSCALER_COUNT) {
+        return hires_upscaler_to_str[upscaler];
+    }
+    return NONE_STR;
+}
+
+enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str) {
+    for (int i = 0; i < SD_HIRES_UPSCALER_COUNT; i++) {
+        if (!strcmp(str, hires_upscaler_to_str[i])) {
+            return (enum sd_hires_upscaler_t)i;
+        }
+    }
+    return SD_HIRES_UPSCALER_COUNT;
+}
+
 void sd_cache_params_init(sd_cache_params_t* cache_params) {
     *cache_params                             = {};
     cache_params->mode                        = SD_CACHE_DISABLED;
@@ -2121,6 +2171,19 @@ void sd_cache_params_init(sd_cache_params_t* cache_params) {
     cache_params->spectrum_stop_percent       = 0.9f;
 }
 
+void sd_hires_params_init(sd_hires_params_t* hires_params) {
+    *hires_params                    = {};
+    hires_params->enabled            = false;
+    hires_params->upscaler           = SD_HIRES_UPSCALER_LATENT;
+    hires_params->model_path         = nullptr;
+    hires_params->scale              = 2.0f;
+    hires_params->target_width       = 0;
+    hires_params->target_height      = 0;
+    hires_params->steps              = 0;
+    hires_params->denoising_strength = 0.7f;
+    hires_params->upscale_tile_size  = 128;
+}
+
 void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
     *sd_ctx_params                         = {};
     sd_ctx_params->vae_decode_only         = true;
@@ -2290,6 +2353,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
     sd_img_gen_params->pm_params         = {nullptr, 0, nullptr, 20.f};
     sd_img_gen_params->vae_tiling_params = {false, 0, 0, 0.5f, 0.0f, 0.0f};
     sd_cache_params_init(&sd_img_gen_params->cache);
+    sd_hires_params_init(&sd_img_gen_params->hires);
 }
 
 char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
@@ -2316,7 +2380,8 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              "increase_ref_index: %s\n"
              "control_strength: %.2f\n"
              "photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\n"
-             "VAE tiling: %s\n",
+             "VAE tiling: %s\n"
+             "hires: {enabled=%s, upscaler=%s, model_path=%s, scale=%.2f, target=%dx%d, steps=%d, denoising_strength=%.2f}\n",
              SAFE_STR(sd_img_gen_params->prompt),
              SAFE_STR(sd_img_gen_params->negative_prompt),
              sd_img_gen_params->clip_skip,
@@ -2333,7 +2398,15 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              sd_img_gen_params->pm_params.style_strength,
              sd_img_gen_params->pm_params.id_images_count,
              SAFE_STR(sd_img_gen_params->pm_params.id_embed_path),
-             BOOL_STR(sd_img_gen_params->vae_tiling_params.enabled));
+             BOOL_STR(sd_img_gen_params->vae_tiling_params.enabled),
+             BOOL_STR(sd_img_gen_params->hires.enabled),
+             sd_hires_upscaler_name(sd_img_gen_params->hires.upscaler),
+             SAFE_STR(sd_img_gen_params->hires.model_path),
+             sd_img_gen_params->hires.scale,
+             sd_img_gen_params->hires.target_width,
+             sd_img_gen_params->hires.target_height,
+             sd_img_gen_params->hires.steps,
+             sd_img_gen_params->hires.denoising_strength);
     const char* cache_mode_str = "disabled";
     if (sd_img_gen_params->cache.mode == SD_CACHE_EASYCACHE) {
         cache_mode_str = "easycache";
@@ -2370,6 +2443,14 @@ struct sd_ctx_t {
     StableDiffusionGGML* sd = nullptr;
 };
 
+static bool sd_version_supports_video_generation(SDVersion version) {
+    return version == VERSION_SVD || sd_version_is_wan(version);
+}
+
+static bool sd_version_supports_image_generation(SDVersion version) {
+    return !sd_version_supports_video_generation(version);
+}
+
 sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) {
     sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t));
     if (sd_ctx == nullptr) {
@@ -2399,6 +2480,20 @@ void free_sd_ctx(sd_ctx_t* sd_ctx) {
     free(sd_ctx);
 }
 
+SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx) {
+    if (sd_ctx == nullptr || sd_ctx->sd == nullptr) {
+        return false;
+    }
+    return sd_version_supports_image_generation(sd_ctx->sd->version);
+}
+
+SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) {
+    if (sd_ctx == nullptr || sd_ctx->sd == nullptr) {
+        return false;
+    }
+    return sd_version_supports_video_generation(sd_ctx->sd->version);
+}
+
 enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) {
     if (sd_ctx != nullptr && sd_ctx->sd != nullptr) {
         if (sd_version_is_dit(sd_ctx->sd->version)) {
@@ -2415,8 +2510,10 @@ enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_me
             return EXPONENTIAL_SCHEDULER;
         }
     }
-    if (sample_method == LCM_SAMPLE_METHOD) {
+    if (sample_method == LCM_SAMPLE_METHOD || sample_method == TCD_SAMPLE_METHOD) {
         return LCM_SCHEDULER;
+    } else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) {
+        return SIMPLE_SCHEDULER;
     }
     return DISCRETE_SCHEDULER;
 }
@@ -2457,6 +2554,7 @@ static float resolve_eta(sd_ctx_t* sd_ctx,
                 return 0.0f;
             case EULER_A_SAMPLE_METHOD:
             case DPMPP2S_A_SAMPLE_METHOD:
+            case ER_SDE_SAMPLE_METHOD:
                 return 1.0f;
             default:;
         }
@@ -2489,6 +2587,7 @@ struct GenerationRequest {
     sd_guidance_params_t guidance            = {};
     sd_guidance_params_t high_noise_guidance = {};
     sd_pm_params_t pm_params                 = {};
+    sd_hires_params_t hires                  = {};
     int frames                               = -1;
     float vace_strength                      = 1.f;
 
@@ -2510,6 +2609,7 @@ struct GenerationRequest {
         auto_resize_ref_image       = sd_img_gen_params->auto_resize_ref_image;
         guidance                    = sd_img_gen_params->sample_params.guidance;
         pm_params                   = sd_img_gen_params->pm_params;
+        hires                       = sd_img_gen_params->hires;
         cache_params                = &sd_img_gen_params->cache;
         resolve(sd_ctx);
     }
@@ -2532,26 +2632,76 @@ struct GenerationRequest {
     }
 
     void align_generation_request_size() {
+        align_image_size(&width, &height, "generation request");
+    }
+
+    void align_image_size(int* target_width, int* target_height, const char* label) {
         int spatial_multiple = vae_scale_factor * diffusion_model_down_factor;
-        int width_offset     = align_up_offset(width, spatial_multiple);
+        int width_offset     = align_up_offset(*target_width, spatial_multiple);
-        int height_offset    = align_up_offset(height, spatial_multiple);
+        int height_offset    = align_up_offset(*target_height, spatial_multiple);
         if (width_offset <= 0 && height_offset <= 0) {
             return;
         }
 
-        int original_width  = width;
+        int original_width  = *target_width;
-        int original_height = height;
+        int original_height = *target_height;
 
-        width += width_offset;
+        *target_width += width_offset;
-        height += height_offset;
+        *target_height += height_offset;
-        LOG_WARN("align up %dx%d to %dx%d (multiple=%d)",
+        LOG_WARN("align %s up %dx%d to %dx%d (multiple=%d)",
+                 label,
                  original_width,
                  original_height,
-                 width,
+                 *target_width,
-                 height,
+                 *target_height,
                  spatial_multiple);
     }
 
+    void resolve_hires() {
+        if (!hires.enabled) {
+            return;
+        }
+        if (hires.upscaler == SD_HIRES_UPSCALER_NONE) {
+            hires.enabled = false;
+            return;
+        }
+        if (hires.upscaler < SD_HIRES_UPSCALER_NONE || hires.upscaler >= SD_HIRES_UPSCALER_COUNT) {
+            LOG_WARN("hires upscaler '%d' is invalid, disabling hires", hires.upscaler);
+            hires.enabled = false;
+            return;
+        }
+        if (hires.upscaler == SD_HIRES_UPSCALER_MODEL && strlen(SAFE_STR(hires.model_path)) == 0) {
+            LOG_WARN("hires model upscaler requires a model path, disabling hires");
+            hires.enabled = false;
+            return;
+        }
+        if (hires.scale <= 0.f && hires.target_width <= 0 && hires.target_height <= 0) {
+            LOG_WARN("hires scale must be positive when no target size is set, disabling hires");
+            hires.enabled = false;
+            return;
+        }
+        hires.denoising_strength = std::clamp(hires.denoising_strength, 0.0001f, 1.f);
+        hires.steps              = std::max(0, hires.steps);
+
+        if (hires.target_width > 0 && hires.target_height > 0) {
+            // pass
+        } else if (hires.target_width > 0) {
+            hires.target_height = hires.target_width;
+        } else if (hires.target_height > 0) {
+            hires.target_width = hires.target_height;
+        } else {
+            hires.target_width  = static_cast<int>(std::round(width * hires.scale));
+            hires.target_height = static_cast<int>(std::round(height * hires.scale));
+        }
+
+        if (hires.target_width <= 0 || hires.target_height <= 0) {
+            LOG_WARN("hires target size is not positive, disabling hires");
+            hires.enabled = false;
+            return;
+        }
+        align_image_size(&hires.target_width, &hires.target_height, "hires target");
+    }
+
     static void resolve_guidance(sd_ctx_t* sd_ctx,
                                  sd_guidance_params_t* guidance,
                                  bool* use_uncond,
@@ -2592,6 +2742,7 @@ struct GenerationRequest {
 
     void resolve(sd_ctx_t* sd_ctx) {
         align_generation_request_size();
+        resolve_hires();
         seed = resolve_seed(seed);
 
         resolve_guidance(sd_ctx, &guidance, &use_uncond, &use_img_cond);
@@ -3104,6 +3255,135 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
     return result_images;
 }
 
+static sd::Tensor<float> upscale_hires_latent(sd_ctx_t* sd_ctx,
+                                              const sd::Tensor<float>& latent,
+                                              const GenerationRequest& request,
+                                              UpscalerGGML* upscaler) {
+    auto get_hires_latent_target_shape = [&]() {
+        std::vector<int64_t> target_shape = latent.shape();
+        if (target_shape.size() < 2) {
+            target_shape.clear();
+            return target_shape;
+        }
+        target_shape[0] = request.hires.target_width / request.vae_scale_factor;
+        target_shape[1] = request.hires.target_height / request.vae_scale_factor;
+        return target_shape;
+    };
+
+    if (request.hires.upscaler == SD_HIRES_UPSCALER_LATENT ||
+        request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_NEAREST ||
+        request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_NEAREST_EXACT ||
+        request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_ANTIALIASED ||
+        request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_BICUBIC ||
+        request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_BICUBIC_ANTIALIASED) {
+        std::vector<int64_t> target_shape = get_hires_latent_target_shape();
+        if (target_shape.empty()) {
+            LOG_ERROR("latent has invalid shape for hires upscale");
+            return {};
+        }
+
+        sd::ops::InterpolateMode mode = sd::ops::InterpolateMode::Nearest;
+        bool antialias                = false;
+        switch (request.hires.upscaler) {
+            case SD_HIRES_UPSCALER_LATENT:
+                mode = sd::ops::InterpolateMode::Bilinear;
+                break;
+            case SD_HIRES_UPSCALER_LATENT_NEAREST:
+                mode = sd::ops::InterpolateMode::Nearest;
+                break;
+            case SD_HIRES_UPSCALER_LATENT_NEAREST_EXACT:
+                mode = sd::ops::InterpolateMode::NearestExact;
+                break;
+            case SD_HIRES_UPSCALER_LATENT_ANTIALIASED:
+                mode      = sd::ops::InterpolateMode::Bilinear;
+                antialias = true;
+                break;
+            case SD_HIRES_UPSCALER_LATENT_BICUBIC:
+                mode = sd::ops::InterpolateMode::Bicubic;
+                break;
+            case SD_HIRES_UPSCALER_LATENT_BICUBIC_ANTIALIASED:
+                mode      = sd::ops::InterpolateMode::Bicubic;
+                antialias = true;
+                break;
+            default:
+                break;
+        }
+
+        LOG_INFO("hires %s upscale %" PRId64 "x%" PRId64 " -> %" PRId64 "x%" PRId64,
+                 sd_hires_upscaler_name(request.hires.upscaler),
+                 latent.shape()[0],
+                 latent.shape()[1],
+                 target_shape[0],
+                 target_shape[1]);
+
+        return sd::ops::interpolate(latent, target_shape, mode, false, antialias);
+    } else if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL ||
+               request.hires.upscaler == SD_HIRES_UPSCALER_LANCZOS ||
+               request.hires.upscaler == SD_HIRES_UPSCALER_NEAREST) {
+        if (sd_ctx->sd->vae_decode_only) {
+            LOG_ERROR("hires %s upscaler requires VAE encoder weights; create the context with vae_decode_only=false",
+                      sd_hires_upscaler_name(request.hires.upscaler));
+            return {};
+        }
+        if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL && upscaler == nullptr) {
+            LOG_ERROR("hires model upscaler context is null");
+            return {};
+        }
+
+        sd::Tensor<float> decoded = sd_ctx->sd->decode_first_stage(latent);
+        if (decoded.empty()) {
+            LOG_ERROR("decode_first_stage failed before hires %s upscale",
+                      sd_hires_upscaler_name(request.hires.upscaler));
+            return {};
+        }
+
+        sd::Tensor<float> upscaled_tensor;
+        if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
+            upscaled_tensor = upscaler->upscale_tensor(decoded);
+            if (upscaled_tensor.empty()) {
+                LOG_ERROR("hires model upscale failed");
+                return {};
+            }
+
+            if (upscaled_tensor.shape()[0] != request.hires.target_width ||
+                upscaled_tensor.shape()[1] != request.hires.target_height) {
+                upscaled_tensor = sd::ops::interpolate(upscaled_tensor,
+                                                       {request.hires.target_width,
+                                                        request.hires.target_height,
+                                                        upscaled_tensor.shape()[2],
+                                                        upscaled_tensor.shape()[3]});
+            }
+        } else {
+            sd::ops::InterpolateMode mode = request.hires.upscaler == SD_HIRES_UPSCALER_LANCZOS
+                                                ? sd::ops::InterpolateMode::Lanczos
+                                                : sd::ops::InterpolateMode::Nearest;
+            LOG_INFO("hires %s image upscale %" PRId64 "x%" PRId64 " -> %dx%d",
+                     sd_hires_upscaler_name(request.hires.upscaler),
+                     decoded.shape()[0],
+                     decoded.shape()[1],
+                     request.hires.target_width,
+                     request.hires.target_height);
+            upscaled_tensor = sd::ops::interpolate(decoded,
+                                                   {request.hires.target_width,
+                                                    request.hires.target_height,
+                                                    decoded.shape()[2],
+                                                    decoded.shape()[3]},
+                                                   mode);
+            upscaled_tensor = sd::ops::clamp(upscaled_tensor, 0.0f, 1.0f);
+        }
+
+        sd::Tensor<float> upscaled_latent = sd_ctx->sd->encode_first_stage(upscaled_tensor);
+        if (upscaled_latent.empty()) {
+            LOG_ERROR("encode_first_stage failed after hires %s upscale",
+                      sd_hires_upscaler_name(request.hires.upscaler));
+        }
+        return upscaled_latent;
+    }
+
+    LOG_ERROR("unsupported hires upscaler '%s'", sd_hires_upscaler_name(request.hires.upscaler));
+    return {};
+}
+
 SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params) {
     if (sd_ctx == nullptr || sd_img_gen_params == nullptr) {
         return nullptr;
@@ -3191,7 +3471,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
         }
         return nullptr;
     }
-    if (sd_ctx->sd->free_params_immediately) {
+    if (sd_ctx->sd->free_params_immediately && !request.hires.enabled) {
         sd_ctx->sd->diffusion_model->free_params_buffer();
     }
     int64_t denoise_end = ggml_time_ms();
@@ -3199,6 +3479,131 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
              final_latents.size(),
              (denoise_end - denoise_start) * 1.0f / 1000);
 
+    if (request.hires.enabled && request.hires.target_width > 0) {
+        LOG_INFO("hires fix: upscaling to %dx%d", request.hires.target_width, request.hires.target_height);
+
+        std::unique_ptr<UpscalerGGML> hires_upscaler;
+        if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
+            LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path);
+            hires_upscaler = std::make_unique<UpscalerGGML>(sd_ctx->sd->n_threads,
+                                                            false,
+                                                            request.hires.upscale_tile_size);
+            if (!hires_upscaler->load_from_file(request.hires.model_path,
+                                                sd_ctx->sd->offload_params_to_cpu,
+                                                sd_ctx->sd->n_threads)) {
+                LOG_ERROR("load hires model upscaler failed");
+                if (sd_ctx->sd->free_params_immediately) {
+                    sd_ctx->sd->diffusion_model->free_params_buffer();
+                }
+                return nullptr;
+            }
+        }
+
+        int hires_steps = request.hires.steps > 0 ? request.hires.steps : plan.sample_steps;
+
+        // sd-webui behavior: scale up total steps so trimming by denoising_strength yields exactly hires_steps effective steps,
+        // unlike img2img which trims from a fixed step count
+        hires_steps = static_cast<int>(hires_steps / request.hires.denoising_strength);
+
+        std::vector<float> hires_sigmas = sd_ctx->sd->denoiser->get_sigmas(
+            hires_steps,
+            sd_ctx->sd->get_image_seq_len(request.hires.target_height, request.hires.target_width),
+            sd_img_gen_params->sample_params.scheduler,
+            sd_ctx->sd->version);
+
+        size_t t_enc = static_cast<size_t>(hires_steps * request.hires.denoising_strength);
+        if (t_enc >= static_cast<size_t>(hires_steps)) {
+            t_enc = static_cast<size_t>(hires_steps) - 1;
+        }
+        std::vector<float> hires_sigma_sched(hires_sigmas.begin() + hires_steps - static_cast<int>(t_enc) - 1,
+                                             hires_sigmas.end());
+        LOG_INFO("hires fix: %d steps, denoising_strength=%.2f, sigma_sched_size=%zu",
+                 hires_steps,
+                 request.hires.denoising_strength,
+                 hires_sigma_sched.size());
+
+        std::vector<sd::Tensor<float>> hires_final_latents;
+        int64_t hires_denoise_start = ggml_time_ms();
+        for (int b = 0; b < (int)final_latents.size(); b++) {
+            int64_t cur_seed = request.seed + b;
+            sd_ctx->sd->rng->manual_seed(cur_seed);
+            sd_ctx->sd->sampler_rng->manual_seed(cur_seed);
+
+            sd::Tensor<float> upscaled = upscale_hires_latent(sd_ctx,
+                                                              final_latents[b],
+                                                              request,
+                                                              hires_upscaler.get());
+            if (upscaled.empty()) {
+                if (sd_ctx->sd->free_params_immediately) {
+                    sd_ctx->sd->diffusion_model->free_params_buffer();
+                }
+                return nullptr;
+            }
+
+            sd::Tensor<float> noise = sd::randn_like<float>(upscaled, sd_ctx->sd->rng);
+
+            sd::Tensor<float> hires_denoise_mask;
+            if (!latents.denoise_mask.empty()) {
+                std::vector<int64_t> mask_shape = latents.denoise_mask.shape();
+                mask_shape[0]                   = upscaled.shape()[0];
+                mask_shape[1]                   = upscaled.shape()[1];
+                hires_denoise_mask              = sd::ops::interpolate(latents.denoise_mask,
+                                                                       mask_shape,
+                                                                       sd::ops::InterpolateMode::NearestMax);
+            }
+
+            int64_t hires_sample_start = ggml_time_ms();
+            sd::Tensor<float> x_0      = sd_ctx->sd->sample(sd_ctx->sd->diffusion_model,
+                                                            true,
+                                                            upscaled,
+                                                            std::move(noise),
+                                                            embeds.cond,
+                                                            embeds.uncond,
+                                                            embeds.img_cond,
+                                                            embeds.id_cond,
+                                                            latents.control_image,
+                                                            request.control_strength,
+                                                            request.guidance,
+                                                            plan.eta,
+                                                            request.shifted_timestep,
+                                                            plan.sample_method,
+                                                            sd_ctx->sd->is_flow_denoiser(),
+                                                            hires_sigma_sched,
+                                                            plan.start_merge_step,
+                                                            latents.ref_latents,
+                                                            request.increase_ref_index,
+                                                            hires_denoise_mask,
+                                                            sd::Tensor<float>(),
+                                                            1.f,
+                                                            request.cache_params);
+            int64_t hires_sample_end   = ggml_time_ms();
+            if (!x_0.empty()) {
+                LOG_INFO("hires sampling %d/%d completed, taking %.2fs",
+                         b + 1,
+                         (int)final_latents.size(),
+                         (hires_sample_end - hires_sample_start) * 1.0f / 1000);
+                hires_final_latents.push_back(std::move(x_0));
+                continue;
+            }
+
+            LOG_ERROR("hires sampling for image %d/%d failed after %.2fs",
+                      b + 1,
+                      (int)final_latents.size(),
+                      (hires_sample_end - hires_sample_start) * 1.0f / 1000);
+            if (sd_ctx->sd->free_params_immediately) {
+                sd_ctx->sd->diffusion_model->free_params_buffer();
+            }
+            return nullptr;
+        }
+        if (sd_ctx->sd->free_params_immediately) {
+            sd_ctx->sd->diffusion_model->free_params_buffer();
+        }
+        int64_t hires_denoise_end = ggml_time_ms();
+        LOG_INFO("hires fix completed, taking %.2fs", (hires_denoise_end - hires_denoise_start) * 1.0f / 1000);
+
+        final_latents = std::move(hires_final_latents);
+    }
+
     auto result = decode_image_outputs(sd_ctx, request, final_latents);
     if (result == nullptr) {
         return nullptr;

src/t5.hpp

@@ -10,452 +10,9 @@
 #include <string>
 #include <unordered_map>
 
-#include "darts.h"
 #include "ggml_extend.hpp"
-#include "json.hpp"
 #include "model.h"
-#include "vocab/vocab.h"
+#include "tokenizers/t5_unigram_tokenizer.h"
-
-// Port from: https://github.com/google/sentencepiece/blob/master/src/unigram_model.h
-// and https://github.com/google/sentencepiece/blob/master/src/unigram_model.h.
-// Original License: https://github.com/google/sentencepiece/blob/master/LICENSE
-//
-// Since tokenization is not the bottleneck in SD, performance was not a major consideration
-// during the migration.
-class MetaspacePreTokenizer {
-private:
-    std::string replacement;
-    bool add_prefix_space;
-
-public:
-    MetaspacePreTokenizer(const std::string replacement = " ", bool add_prefix_space = true)
-        : replacement(replacement), add_prefix_space(add_prefix_space) {}
-
-    std::string tokenize(const std::string& input) const {
-        std::string tokens;
-        std::stringstream ss(input);
-
-        if (add_prefix_space) {
-            tokens += replacement;
-        }
-
-        std::string token;
-        bool firstToken = true;
-        while (std::getline(ss, token, ' ')) {
-            if (!firstToken)
-                tokens += replacement + token;
-            else
-                tokens += token;
-
-            firstToken = false;
-        }
-
-        return tokens;
-    }
-};
-
-using EncodeResult = std::vector<std::pair<std::string, int>>;
-class T5UniGramTokenizer {
-public:
-    enum Status {
-        OK,
-        NO_PIECES_LOADED,
-        NO_ENTRY_FOUND,
-        BUILD_DOUBLE_ARRAY_FAILED,
-        PIECE_ALREADY_DEFINED,
-        INVLIAD_JSON
-    };
-
-protected:
-    MetaspacePreTokenizer pre_tokenizer;
-
-    // all <piece, score> pairs
-    std::vector<std::pair<std::string, float>> piece_score_pairs;
-
-    float min_score_ = 0.0;
-    float max_score_ = 0.0;
-    std::unique_ptr<Darts::DoubleArray> trie_;
-
-    // Maximum size of the return value of Trie, which corresponds
-    // to the maximum size of shared common prefix in the sentence pieces.
-    int trie_results_size_;
-    // unknown id.
-    int unk_id_            = 2;
-    std::string eos_token_ = "</s>";
-    int eos_id_            = 1;
-    int pad_id_            = 0;
-    // status.
-    Status status_ = OK;
-
-    float kUnkPenalty = 10.0;
-
-    std::string replacement;
-    bool add_prefix_space = true;
-
-    void InitializePieces(const std::string& json_str) {
-        nlohmann::json data;
-
-        try {
-            data = nlohmann::json::parse(json_str);
-        } catch (const nlohmann::json::parse_error&) {
-            status_ = INVLIAD_JSON;
-            return;
-        }
-        if (!data.contains("model")) {
-            status_ = INVLIAD_JSON;
-            return;
-        }
-        nlohmann::json model = data["model"];
-        if (!model.contains("vocab")) {
-            status_ = INVLIAD_JSON;
-            return;
-        }
-        if (model.contains("unk_id")) {
-            unk_id_ = model["unk_id"];
-        }
-
-        replacement      = data["pre_tokenizer"]["replacement"];
-        add_prefix_space = data["pre_tokenizer"]["add_prefix_space"];
-
-        pre_tokenizer = MetaspacePreTokenizer(replacement, add_prefix_space);
-
-        for (const auto& item : model["vocab"]) {
-            if (item.size() != 2 || !item[0].is_string() || !item[1].is_number_float()) {
-                status_ = INVLIAD_JSON;
-                return;
-            }
-            std::string piece = item[0];
-            if (piece.empty()) {
-                piece = "<empty_token>";
-            }
-            float score = item[1];
-            piece_score_pairs.emplace_back(piece, score);
-        }
-    }
-
-    // Builds a Trie index.
-    void BuildTrie(std::vector<std::pair<std::string, int>>* pieces) {
-        if (status_ != OK)
-            return;
-
-        if (pieces->empty()) {
-            status_ = NO_PIECES_LOADED;
-            return;
-        }
-
-        // sort by sentencepiece since DoubleArray::build()
-        // only accepts sorted strings.
-        sort(pieces->begin(), pieces->end());
-
-        // Makes key/value set for DoubleArrayTrie.
-        std::vector<const char*> key(pieces->size());
-        std::vector<int> value(pieces->size());
-        for (size_t i = 0; i < pieces->size(); ++i) {
-            // LOG_DEBUG("%s %d", (*pieces)[i].first.c_str(), (*pieces)[i].second);
-            key[i]   = (*pieces)[i].first.data();  // sorted piece.
-            value[i] = (*pieces)[i].second;        // vocab_id
-        }
-
-        trie_ = std::unique_ptr<Darts::DoubleArray>(new Darts::DoubleArray());
-        if (trie_->build(key.size(), const_cast<char**>(&key[0]), nullptr,
-                         &value[0]) != 0) {
-            status_ = BUILD_DOUBLE_ARRAY_FAILED;
-            return;
-        }
-
-        // Computes the maximum number of shared prefixes in the trie.
-        const int kMaxTrieResultsSize = 1024;
-        std::vector<Darts::DoubleArray::result_pair_type> results(
-            kMaxTrieResultsSize);
-        trie_results_size_ = 0;
-        for (const auto& p : *pieces) {
-            const size_t num_nodes = trie_->commonPrefixSearch(
-                p.first.data(), results.data(), results.size(), p.first.size());
-            trie_results_size_ = std::max(trie_results_size_, static_cast<int>(num_nodes));
-        }
-
-        if (trie_results_size_ == 0)
-            status_ = NO_ENTRY_FOUND;
-    }
-
-    // Non-virtual (inlined) implementation for faster execution.
-    inline float GetScoreInlined(int id) const {
-        return piece_score_pairs[id].second;
-    }
-
-    inline bool IsUnusedInlined(int id) const {
-        return false;  // TODO
-    }
-
-    inline bool IsUserDefinedInlined(int id) const {
-        return false;  // TODO
-    }
-
-    inline size_t OneCharLen(const char* src) const {
-        return "\1\1\1\1\1\1\1\1\1\1\1\1\2\2\3\4"[(*src & 0xFF) >> 4];
-    }
-
-    // The optimized Viterbi encode.
-    // Main differences from the original function:
-    // 1. Memorizes the best path at each postion so far,
-    // 2. No need to store the Lattice nodes,
-    // 3. Works in utf-8 directly,
-    // 4. Defines a new struct with fewer fields than Lattice,
-    // 5. Does not depend on `class Lattice` nor call `SetSentence()`,
-    // `PopulateNodes()`, or `Viterbi()`. It does everything in one function.
-    // For detailed explanations please see the comments inside the function body.
-    EncodeResult EncodeOptimized(const std::string& normalized) const {
-        // An optimized Viterbi algorithm for unigram language models. Benchmarking
-        // results show that it generates almost identical outputs and achieves 2.1x
-        // speedup on average for 102 languages compared to the original
-        // implementation. It's based on the following three ideas:
-        //
-        // 1. Because it uses the *unigram* model:
-        //     best_score(x1, x2, ... xt) = best_score(x1, x2, ... x{t-1}) + score(xt)
-        // Deciding the best path (and score) can be decoupled into two isolated
-        // terms: (a) the best path ended before the last token `best_score(x1, x2, ...)`
-        // x{t-1})`, and (b) the last token and its `score(xt)`. The two terms are
-        // not related to each other at all.
-        //
-        // Therefore, we can compute once and store the *best_path ending at
-        // each character position*. In this way, when we know best_path_ends_at[M],
-        // we can reuse it to compute all the best_path_ends_at_[...] where the last
-        // token starts at the same character position M.
-        //
-        // This improves the time complexity from O(n*k*k) to O(n*k) because it
-        // eliminates the extra loop of recomputing the best path ending at the same
-        // position, where n is the input length and k is the maximum number of tokens
-        // that can be recognized starting at each position.
-        //
-        // 2. Again, because it uses the *unigram* model, we don't need to actually
-        // store the lattice nodes. We still recognize all the tokens and lattice
-        // nodes from the input, but along identifying them, we use and discard them
-        // on the fly. There is no need to actually store them for best path Viterbi
-        // decoding. The only thing we need to store is the best_path ending at
-        // each character position.
-        //
-        // This improvement reduces the things needed to store in memory from O(n*k)
-        // to O(n), where n is the input length and k is the maximum number of tokens
-        // that can be recognized starting at each position.
-        //
-        // It also avoids the need of dynamic-size lattice node pool, because the
-        // number of things to store is fixed as n.
-        //
-        // 3. SentencePiece is designed to work with unicode, taking utf-8 encoding
-        // inputs. In the original implementation, the lattice positions are based on
-        // unicode positions. A mapping from unicode position to the utf-8 position is
-        // maintained to recover the utf-8 string piece.
-        //
-        // We found that it is sufficient and beneficial to directly work with utf-8
-        // positions:
-        //
-        // Firstly, it saves the conversion and mapping between unicode positions and
-        // utf-8 positions.
-        //
-        // Secondly, it reduces the number of fields we need to maintain in the
-        // node/path structure. Specifically, there are 8 fields defined in
-        // `Lattice::Node` used by the original encoder, but here in the optimized
-        // encoder we only need to define 3 fields in `BestPathNode`.
-
-        if (status() != OK || normalized.empty()) {
-            return {};
-        }
-        // Represents the last node of the best path.
-        struct BestPathNode {
-            int id = -1;  // The vocab id. (maybe -1 for UNK)
-            float best_path_score =
-                0;  // The total score of the best path ending at this node.
-            int starts_at =
-                -1;  // The starting position (in utf-8) of this node. The entire best
-                     // path can be constructed by backtracking along this link.
-        };
-        const int size        = static_cast<int>(normalized.size());
-        const float unk_score = min_score() - kUnkPenalty;
-        // The ends are exclusive.
-        std::vector<BestPathNode> best_path_ends_at(size + 1);
-        // Generate lattice on-the-fly (not stored) and update best_path_ends_at.
-        int starts_at = 0;
-        while (starts_at < size) {
-            std::size_t node_pos = 0;
-            std::size_t key_pos  = starts_at;
-            const auto best_path_score_till_here =
-                best_path_ends_at[starts_at].best_path_score;
-            bool has_single_node = false;
-            const int mblen =
-                std::min<int>(static_cast<int>(OneCharLen(normalized.data() + starts_at)),
-                              size - starts_at);
-            while (key_pos < size) {
-                const int ret =
-                    trie_->traverse(normalized.data(), node_pos, key_pos, key_pos + 1);
-                if (ret == -2)
-                    break;
-                if (ret >= 0) {
-                    if (IsUnusedInlined(ret))
-                        continue;
-                    // Update the best path node.
-                    auto& target_node = best_path_ends_at[key_pos];
-                    const auto length = (key_pos - starts_at);
-                    // User defined symbol receives extra bonus to always be selected.
-                    const auto score = IsUserDefinedInlined(ret)
-                                           ? (length * max_score_ - 0.1)
-                                           : GetScoreInlined(ret);
-                    const auto candidate_best_path_score =
-                        score + best_path_score_till_here;
-                    if (target_node.starts_at == -1 ||
-                        candidate_best_path_score > target_node.best_path_score) {
-                        target_node.best_path_score = static_cast<float>(candidate_best_path_score);
-                        target_node.starts_at       = starts_at;
-                        target_node.id              = ret;
-                    }
-                    if (!has_single_node && length == mblen) {
-                        has_single_node = true;
-                    }
-                }
-            }
-            if (!has_single_node) {
-                auto& target_node = best_path_ends_at[starts_at + mblen];
-                const auto candidate_best_path_score =
-                    unk_score + best_path_score_till_here;
-                if (target_node.starts_at == -1 ||
-                    candidate_best_path_score > target_node.best_path_score) {
-                    target_node.best_path_score = candidate_best_path_score;
-                    target_node.starts_at       = starts_at;
-                    target_node.id              = unk_id_;
-                }
-            }
-            // Move by one unicode character.
-            starts_at += mblen;
-        }
-        // Backtrack to identify the best path.
-        EncodeResult results;
-        int ends_at = size;
-        while (ends_at > 0) {
-            const auto& node = best_path_ends_at[ends_at];
-            results.emplace_back(
-                normalized.substr(node.starts_at, ends_at - node.starts_at), node.id);
-            ends_at = node.starts_at;
-        }
-        std::reverse(results.begin(), results.end());
-        return results;
-    }
-
-public:
-    explicit T5UniGramTokenizer(bool is_umt5 = false) {
-        if (is_umt5) {
-            InitializePieces(load_umt5_tokenizer_json());
-        } else {
-            InitializePieces(load_t5_tokenizer_json());
-        }
-
-        min_score_ = FLT_MAX;
-        max_score_ = FLT_MIN;
-
-        std::vector<std::pair<std::string, int>> pieces;
-        for (int i = 0; i < piece_score_pairs.size(); i++) {
-            const auto& sp = piece_score_pairs[i];
-
-            min_score_ = std::min(min_score_, sp.second);
-            max_score_ = std::max(max_score_, sp.second);
-
-            pieces.emplace_back(sp.first, i);
-        }
-
-        BuildTrie(&pieces);
-    }
-    ~T5UniGramTokenizer(){};
-
-    std::string Normalize(const std::string& input) const {
-        // Ref: https://github.com/huggingface/tokenizers/blob/1ff56c0c70b045f0cd82da1af9ac08cd4c7a6f9f/bindings/python/py_src/tokenizers/implementations/sentencepiece_unigram.py#L29
-        // TODO: nmt-nfkc
-        std::string normalized = std::regex_replace(input, std::regex(" {2,}"), " ");
-        return normalized;
-    }
-
-    std::vector<int> Encode(const std::string& input, bool append_eos_if_not_present = true) const {
-        std::string normalized = Normalize(input);
-        normalized             = pre_tokenizer.tokenize(normalized);
-        EncodeResult result    = EncodeOptimized(normalized);
-        if (result.size() > 0 && append_eos_if_not_present) {
-            auto item = result[result.size() - 1];
-            if (item.first != eos_token_) {
-                result.emplace_back(eos_token_, eos_id_);
-            }
-        }
-        std::vector<int> tokens;
-        for (auto item : result) {
-            tokens.push_back(item.second);
-        }
-        return tokens;
-    }
-
-    void pad_tokens(std::vector<int>& tokens,
-                    std::vector<float>& weights,
-                    std::vector<float>* attention_mask,
-                    size_t max_length = 0,
-                    bool padding      = false) {
-        if (max_length > 0 && padding) {
-            size_t orig_token_num = tokens.size() - 1;
-            size_t n              = static_cast<size_t>(std::ceil(orig_token_num * 1.0 / (max_length - 1)));
-            if (n == 0) {
-                n = 1;
-            }
-            size_t length = max_length * n;
-            LOG_DEBUG("token length: %llu", length);
-            std::vector<int> new_tokens;
-            std::vector<float> new_weights;
-            std::vector<float> new_attention_mask;
-            int token_idx = 0;
-            for (int i = 0; i < length; i++) {
-                if (token_idx >= orig_token_num) {
-                    break;
-                }
-                if (attention_mask != nullptr) {
-                    new_attention_mask.push_back(0.0);
-                }
-                if (i % max_length == max_length - 1) {
-                    new_tokens.push_back(eos_id_);
-                    new_weights.push_back(1.0);
-                } else {
-                    new_tokens.push_back(tokens[token_idx]);
-                    new_weights.push_back(weights[token_idx]);
-                    token_idx++;
-                }
-            }
-
-            new_tokens.push_back(eos_id_);
-            new_weights.push_back(1.0);
-            if (attention_mask != nullptr) {
-                new_attention_mask.push_back(0.0);
-            }
-
-            tokens  = new_tokens;
-            weights = new_weights;
-            if (attention_mask != nullptr) {
-                *attention_mask = new_attention_mask;
-            }
-
-            if (padding) {
-                int pad_token_id = pad_id_;
-                tokens.insert(tokens.end(), length - tokens.size(), pad_token_id);
-                weights.insert(weights.end(), length - weights.size(), 1.0);
-                if (attention_mask != nullptr) {
-                    // maybe keep some padding tokens unmasked?
-                    attention_mask->insert(attention_mask->end(), length - attention_mask->size(), -HUGE_VALF);
-                }
-            }
-        }
-    }
-
-    // Returns the minimum score in sentence pieces.
-    // min_score() - 10 is used for the cost of unknown sentence.
-    float min_score() const { return min_score_; }
-
-    // Returns the maximum score in sentence pieces.
-    // max_score() is used for the cost of user defined symbols.
-    float max_score() const { return max_score_; }
-
-    Status status() const { return status_; }
-};
 
 class T5LayerNorm : public UnaryBlock {
 protected:
@@ -937,18 +494,17 @@ struct T5Embedder {
         for (const auto& item : parsed_attention) {
             const std::string& curr_text = item.first;
             float curr_weight            = item.second;
-            std::vector<int> curr_tokens = tokenizer.Encode(curr_text, false);
+            std::vector<int> curr_tokens = tokenizer.encode(curr_text);
             tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
             weights.insert(weights.end(), curr_tokens.size(), curr_weight);
         }
 
-        int EOS_TOKEN_ID = 1;
-        tokens.push_back(EOS_TOKEN_ID);
-        weights.push_back(1.0);
-
         std::vector<float> attention_mask;
 
-        tokenizer.pad_tokens(tokens, weights, &attention_mask, max_length, padding);
+        tokenizer.pad_tokens(tokens, &weights, &attention_mask, padding ? max_length : 0, padding ? max_length : 100000000, padding);
+        for (auto& mask_value : attention_mask) {
+            mask_value = mask_value > 0.0f ? 0.0f : -HUGE_VALF;
+        }
 
         // for (int i = 0; i < tokens.size(); i++) {
         //     std::cout << tokens[i] << ":" << weights[i] << ", ";

src/tensor.hpp

@@ -815,11 +815,202 @@ namespace sd {
     namespace ops {
         enum class InterpolateMode {
             Nearest,
+            NearestExact,
             NearestMax,
             NearestMin,
             NearestAvg,
+            Bilinear,
+            Bicubic,
+            Lanczos,
         };
 
+        inline bool is_nearest_like_interpolate_mode(InterpolateMode mode) {
+            return mode == InterpolateMode::Nearest ||
+                   mode == InterpolateMode::NearestExact ||
+                   mode == InterpolateMode::NearestMax ||
+                   mode == InterpolateMode::NearestMin ||
+                   mode == InterpolateMode::NearestAvg;
+        }
+
+        inline bool is_2d_filter_interpolate_mode(InterpolateMode mode) {
+            return mode == InterpolateMode::Bilinear ||
+                   mode == InterpolateMode::Bicubic ||
+                   mode == InterpolateMode::Lanczos;
+        }
+
+        inline int64_t nearest_exact_interpolate_index(int64_t output_index,
+                                                       int64_t input_size,
+                                                       int64_t output_size) {
+            const double scale  = static_cast<double>(input_size) / static_cast<double>(output_size);
+            const double center = (static_cast<double>(output_index) + 0.5) * scale - 0.5;
+            return std::min(std::max<int64_t>(static_cast<int64_t>(std::floor(center + 0.5)), 0), input_size - 1);
+        }
+
+        inline double linear_interpolate_weight(double x) {
+            x = std::abs(x);
+            return x < 1.0 ? 1.0 - x : 0.0;
+        }
+
+        inline double cubic_interpolate_weight(double x) {
+            constexpr double a = -0.75;  // Match PyTorch bicubic interpolation.
+            x                  = std::abs(x);
+            if (x <= 1.0) {
+                return ((a + 2.0) * x - (a + 3.0)) * x * x + 1.0;
+            }
+            if (x < 2.0) {
+                return ((a * x - 5.0 * a) * x + 8.0 * a) * x - 4.0 * a;
+            }
+            return 0.0;
+        }
+
+        inline double sinc(double x) {
+            constexpr double pi = 3.14159265358979323846;
+            if (std::abs(x) < 1e-12) {
+                return 1.0;
+            }
+            const double pix = pi * x;
+            return std::sin(pix) / pix;
+        }
+
+        inline double lanczos_interpolate_weight(double x) {
+            constexpr double radius = 3.0;
+            x                       = std::abs(x);
+            if (x >= radius) {
+                return 0.0;
+            }
+            return sinc(x) * sinc(x / radius);
+        }
+
+        struct InterpolateContributor {
+            int64_t index;
+            double weight;
+        };
+
+        inline std::vector<std::vector<InterpolateContributor>> make_interpolate_contributors(
+            int64_t input_size,
+            int64_t output_size,
+            InterpolateMode mode,
+            bool antialias) {
+            std::vector<std::vector<InterpolateContributor>> contributors(static_cast<size_t>(output_size));
+            const double scale        = static_cast<double>(input_size) / static_cast<double>(output_size);
+            const double filter_scale = antialias ? std::max(1.0, scale) : 1.0;
+
+            for (int64_t out = 0; out < output_size; ++out) {
+                const double center = (static_cast<double>(out) + 0.5) * scale - 0.5;
+                int64_t start       = 0;
+                int64_t end         = 0;
+
+                if (mode == InterpolateMode::Bilinear) {
+                    const double support = filter_scale;
+                    start                = static_cast<int64_t>(std::ceil(center - support));
+                    end                  = static_cast<int64_t>(std::floor(center + support));
+                } else if (mode == InterpolateMode::Bicubic) {
+                    const double support = 2.0 * filter_scale;
+                    start                = static_cast<int64_t>(std::ceil(center - support));
+                    end                  = static_cast<int64_t>(std::floor(center + support));
+                } else if (mode == InterpolateMode::Lanczos) {
+                    const double support = 3.0 * filter_scale;
+                    start                = static_cast<int64_t>(std::ceil(center - support));
+                    end                  = static_cast<int64_t>(std::floor(center + support));
+                } else {
+                    tensor_throw_invalid_argument("Unsupported 2D filter interpolate mode: mode=" +
+                                                  std::to_string(static_cast<int>(mode)));
+                }
+
+                double weight_sum                                      = 0.0;
+                std::vector<InterpolateContributor>& axis_contributors = contributors[static_cast<size_t>(out)];
+                axis_contributors.reserve(static_cast<size_t>(end - start + 1));
+
+                for (int64_t in = start; in <= end; ++in) {
+                    double weight = 0.0;
+                    if (mode == InterpolateMode::Bilinear) {
+                        weight = linear_interpolate_weight((center - static_cast<double>(in)) / filter_scale);
+                    } else if (mode == InterpolateMode::Bicubic) {
+                        weight = cubic_interpolate_weight((center - static_cast<double>(in)) / filter_scale);
+                    } else {
+                        weight = lanczos_interpolate_weight((center - static_cast<double>(in)) / filter_scale);
+                    }
+
+                    if (weight == 0.0) {
+                        continue;
+                    }
+
+                    const int64_t clamped_index = std::min(std::max<int64_t>(in, 0), input_size - 1);
+                    axis_contributors.push_back({clamped_index, weight});
+                    weight_sum += weight;
+                }
+
+                if ((antialias || mode == InterpolateMode::Lanczos) &&
+                    std::abs(weight_sum) > 1e-12) {
+                    for (auto& contributor : axis_contributors) {
+                        contributor.weight /= weight_sum;
+                    }
+                }
+
+                if (axis_contributors.empty()) {
+                    const int64_t nearest = std::min(
+                        std::max<int64_t>(static_cast<int64_t>(std::floor(center + 0.5)), 0),
+                        input_size - 1);
+                    axis_contributors.push_back({nearest, 1.0});
+                }
+            }
+
+            return contributors;
+        }
+
+        template <typename T>
+        inline Tensor<T> interpolate_2d_filter(const Tensor<T>& input,
+                                               const std::vector<int64_t>& output_shape,
+                                               InterpolateMode mode,
+                                               bool antialias) {
+            if (input.dim() < 2) {
+                tensor_throw_invalid_argument("2D filter interpolate requires rank >= 2: input_shape=" +
+                                              tensor_shape_to_string(input.shape()) + ", output_shape=" +
+                                              tensor_shape_to_string(output_shape));
+            }
+            for (size_t i = 2; i < output_shape.size(); ++i) {
+                if (input.shape()[i] != output_shape[i]) {
+                    tensor_throw_invalid_argument("2D filter interpolate only supports resizing dimensions 0 and 1: input_shape=" +
+                                                  tensor_shape_to_string(input.shape()) + ", output_shape=" +
+                                                  tensor_shape_to_string(output_shape));
+                }
+            }
+
+            Tensor<T> output(output_shape);
+            const int64_t input_width   = input.shape()[0];
+            const int64_t input_height  = input.shape()[1];
+            const int64_t output_width  = output_shape[0];
+            const int64_t output_height = output_shape[1];
+            const int64_t input_plane   = input_width * input_height;
+            const int64_t output_plane  = output_width * output_height;
+            const int64_t plane_count   = input.numel() / input_plane;
+
+            auto x_contributors = make_interpolate_contributors(input_width, output_width, mode, antialias);
+            auto y_contributors = make_interpolate_contributors(input_height, output_height, mode, antialias);
+
+            for (int64_t plane = 0; plane < plane_count; ++plane) {
+                const int64_t input_plane_offset  = plane * input_plane;
+                const int64_t output_plane_offset = plane * output_plane;
+                for (int64_t y = 0; y < output_height; ++y) {
+                    const auto& y_axis = y_contributors[static_cast<size_t>(y)];
+                    for (int64_t x = 0; x < output_width; ++x) {
+                        const auto& x_axis = x_contributors[static_cast<size_t>(x)];
+                        double value       = 0.0;
+                        for (const auto& yc : y_axis) {
+                            const int64_t input_row_offset = input_plane_offset + yc.index * input_width;
+                            for (const auto& xc : x_axis) {
+                                value += static_cast<double>(input.data()[input_row_offset + xc.index]) *
+                                         xc.weight * yc.weight;
+                            }
+                        }
+                        output.data()[output_plane_offset + y * output_width + x] = static_cast<T>(value);
+                    }
+                }
+            }
+
+            return output;
+        }
+
         inline int64_t normalize_slice_bound(int64_t index, int64_t dim_size) {
             if (index < 0) {
                 index += dim_size;
@@ -1014,17 +1205,20 @@ namespace sd {
         inline Tensor<T> interpolate(const Tensor<T>& input,
                                      std::vector<int64_t> output_shape,
                                      InterpolateMode mode = InterpolateMode::Nearest,
-                                     bool align_corners   = false) {
+                                     bool align_corners   = false,
-            const bool is_nearest_like_mode = (mode == InterpolateMode::Nearest ||
+                                     bool antialias       = false) {
-                                               mode == InterpolateMode::NearestMax ||
+            const bool is_nearest_like_mode = is_nearest_like_interpolate_mode(mode);
-                                               mode == InterpolateMode::NearestMin ||
+            const bool is_2d_filter_mode    = is_2d_filter_interpolate_mode(mode);
-                                               mode == InterpolateMode::NearestAvg);
+            if (!is_nearest_like_mode && !is_2d_filter_mode) {
-            if (!is_nearest_like_mode) {
+                tensor_throw_invalid_argument("Unsupported interpolate mode: mode=" +
-                tensor_throw_invalid_argument("Only nearest-like interpolate modes are implemented, got mode=" +
+                                              std::to_string(static_cast<int>(mode)));
+            }
+            if (antialias && !is_2d_filter_mode) {
+                tensor_throw_invalid_argument("Tensor interpolate antialias requires a 2D filter mode: mode=" +
                                               std::to_string(static_cast<int>(mode)));
             }
             if (align_corners) {
-                tensor_throw_invalid_argument("align_corners is not supported for nearest-like interpolate: input_shape=" +
+                tensor_throw_invalid_argument("align_corners is not supported for tensor interpolate: input_shape=" +
                                               tensor_shape_to_string(input.shape()) + ", output_shape=" +
                                               tensor_shape_to_string(output_shape));
             }
@@ -1051,6 +1245,10 @@ namespace sd {
                 }
             }
 
+            if (is_2d_filter_mode) {
+                return interpolate_2d_filter(input, output_shape, mode, antialias);
+            }
+
             bool has_downsampling = false;
             for (int64_t i = 0; i < input.dim(); ++i) {
                 if (input.shape()[i] > output_shape[i]) {
@@ -1060,13 +1258,21 @@ namespace sd {
             }
 
             Tensor<T> output(std::move(output_shape));
-            if (mode == InterpolateMode::Nearest || !has_downsampling) {
+            if (mode == InterpolateMode::Nearest ||
+                mode == InterpolateMode::NearestExact ||
+                !has_downsampling) {
                 for (int64_t flat = 0; flat < output.numel(); ++flat) {
                     std::vector<int64_t> output_coord = tensor_unravel_index(flat, output.shape());
                     std::vector<int64_t> input_coord(static_cast<size_t>(input.dim()), 0);
                     for (size_t i = 0; i < static_cast<size_t>(input.dim()); ++i) {
+                        if (mode == InterpolateMode::NearestExact) {
+                            input_coord[i] = nearest_exact_interpolate_index(output_coord[i],
+                                                                             input.shape()[i],
+                                                                             output.shape()[i]);
+                        } else {
                             input_coord[i] = output_coord[i] * input.shape()[i] / output.shape()[i];
                         }
+                    }
                     output[flat] = input.index(input_coord);
                 }
 
@@ -1083,6 +1289,12 @@ namespace sd {
                         return T(0);
                     case InterpolateMode::Nearest:
                         return T(0);
+                    case InterpolateMode::NearestExact:
+                        return T(0);
+                    case InterpolateMode::Bilinear:
+                    case InterpolateMode::Bicubic:
+                    case InterpolateMode::Lanczos:
+                        break;
                 }
 
                 tensor_throw_invalid_argument("Unsupported interpolate mode: mode=" +
@@ -1102,6 +1314,12 @@ namespace sd {
                         break;
                     case InterpolateMode::Nearest:
                         break;
+                    case InterpolateMode::NearestExact:
+                        break;
+                    case InterpolateMode::Bilinear:
+                    case InterpolateMode::Bicubic:
+                    case InterpolateMode::Lanczos:
+                        break;
                 }
             };
 
@@ -1157,17 +1375,20 @@ namespace sd {
                                      const std::optional<std::vector<int64_t>>& size,
                                      const std::optional<std::vector<double>>& scale_factor,
                                      InterpolateMode mode = InterpolateMode::Nearest,
-                                     bool align_corners   = false) {
+                                     bool align_corners   = false,
-            const bool is_nearest_like_mode = (mode == InterpolateMode::Nearest ||
+                                     bool antialias       = false) {
-                                               mode == InterpolateMode::NearestMax ||
+            const bool is_nearest_like_mode = is_nearest_like_interpolate_mode(mode);
-                                               mode == InterpolateMode::NearestMin ||
+            const bool is_2d_filter_mode    = is_2d_filter_interpolate_mode(mode);
-                                               mode == InterpolateMode::NearestAvg);
+            if (!is_nearest_like_mode && !is_2d_filter_mode) {
-            if (!is_nearest_like_mode) {
+                tensor_throw_invalid_argument("Unsupported interpolate mode: mode=" +
-                tensor_throw_invalid_argument("Only nearest-like interpolate modes are implemented, got mode=" +
+                                              std::to_string(static_cast<int>(mode)));
+            }
+            if (antialias && !is_2d_filter_mode) {
+                tensor_throw_invalid_argument("Tensor interpolate antialias requires a 2D filter mode: mode=" +
                                               std::to_string(static_cast<int>(mode)));
             }
             if (align_corners) {
-                tensor_throw_invalid_argument("align_corners is not supported for nearest-like interpolate: input_shape=" +
+                tensor_throw_invalid_argument("align_corners is not supported for tensor interpolate: input_shape=" +
                                               tensor_shape_to_string(input.shape()));
             }
             if (size.has_value() == scale_factor.has_value()) {
@@ -1211,7 +1432,7 @@ namespace sd {
                 }
             }
 
-            return interpolate(input, std::move(output_shape), mode, align_corners);
+            return interpolate(input, std::move(output_shape), mode, align_corners, antialias);
         }
 
         template <typename T>
@@ -1219,12 +1440,14 @@ namespace sd {
                                      const std::optional<std::vector<int64_t>>& size,
                                      double scale_factor,
                                      InterpolateMode mode = InterpolateMode::Nearest,
-                                     bool align_corners   = false) {
+                                     bool align_corners   = false,
+                                     bool antialias       = false) {
             return interpolate(input,
                                size,
                                std::vector<double>(size.has_value() ? size->size() : input.dim(), scale_factor),
                                mode,
-                               align_corners);
+                               align_corners,
+                               antialias);
         }
 
         template <typename T>

src/tokenizers/bpe_tokenizer.cpp

@@ -0,0 +1,213 @@
+#include "bpe_tokenizer.h"
+
+#include <algorithm>
+#include <sstream>
+
+#include "tokenize_util.h"
+#include "util.h"
+
+std::vector<std::pair<int, std::u32string>> BPETokenizer::bytes_to_unicode() {
+    std::vector<std::pair<int, std::u32string>> byte_unicode_pairs;
+    std::set<int> byte_set;
+    for (int b = static_cast<int>('!'); b <= static_cast<int>('~'); ++b) {
+        byte_set.insert(b);
+        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
+    }
+    for (int b = 161; b <= 172; ++b) {
+        byte_set.insert(b);
+        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
+    }
+    for (int b = 174; b <= 255; ++b) {
+        byte_set.insert(b);
+        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
+    }
+    int n = 0;
+    for (int b = 0; b < 256; ++b) {
+        if (byte_set.find(b) == byte_set.end()) {
+            byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(n + 256)));
+            ++n;
+        }
+    }
+    return byte_unicode_pairs;
+}
+
+std::vector<std::string> BPETokenizer::token_split(const std::string& text) const {
+    return ::token_split(text);
+}
+
+std::vector<std::u32string> BPETokenizer::split_utf32(const std::string& text, char32_t delimiter) {
+    std::vector<std::u32string> result;
+    size_t start              = 0;
+    size_t pos                = 0;
+    std::u32string utf32_text = utf8_to_utf32(text);
+    while ((pos = utf32_text.find(delimiter, start)) != std::u32string::npos) {
+        result.push_back(utf32_text.substr(start, pos - start));
+        start = pos + 1;
+    }
+    return result;
+}
+
+static std::set<std::pair<std::u32string, std::u32string>> get_pairs(const std::vector<std::u32string>& subwords) {
+    std::set<std::pair<std::u32string, std::u32string>> pairs;
+    if (subwords.empty()) {
+        return pairs;
+    }
+
+    std::u32string prev_subword = subwords[0];
+    for (int i = 1; i < static_cast<int>(subwords.size()); i++) {
+        std::u32string subword = subwords[i];
+        std::pair<std::u32string, std::u32string> pair(prev_subword, subword);
+        pairs.insert(pair);
+        prev_subword = subword;
+    }
+    return pairs;
+}
+
+std::vector<std::u32string> BPETokenizer::bpe(const std::u32string& token) const {
+    std::vector<std::u32string> word;
+
+    for (int i = 0; i < static_cast<int>(token.size()) - 1; i++) {
+        word.emplace_back(1, token[i]);
+    }
+    word.push_back(token.substr(token.size() - 1) + utf8_to_utf32(end_of_word_suffix));
+
+    std::set<std::pair<std::u32string, std::u32string>> pairs = get_pairs(word);
+
+    if (pairs.empty()) {
+        return {token + utf8_to_utf32(end_of_word_suffix)};
+    }
+
+    while (true) {
+        auto min_pair_iter = std::min_element(pairs.begin(),
+                                              pairs.end(),
+                                              [&](const std::pair<std::u32string, std::u32string>& a,
+                                                  const std::pair<std::u32string, std::u32string>& b) {
+                                                  if (bpe_ranks.find(a) == bpe_ranks.end()) {
+                                                      return false;
+                                                  } else if (bpe_ranks.find(b) == bpe_ranks.end()) {
+                                                      return true;
+                                                  }
+                                                  return bpe_ranks.at(a) < bpe_ranks.at(b);
+                                              });
+
+        const std::pair<std::u32string, std::u32string>& bigram = *min_pair_iter;
+
+        if (bpe_ranks.find(bigram) == bpe_ranks.end()) {
+            break;
+        }
+
+        std::u32string first  = bigram.first;
+        std::u32string second = bigram.second;
+        std::vector<std::u32string> new_word;
+        int32_t i = 0;
+
+        while (i < static_cast<int32_t>(word.size())) {
+            auto it = std::find(word.begin() + i, word.end(), first);
+            if (it == word.end()) {
+                new_word.insert(new_word.end(), word.begin() + i, word.end());
+                break;
+            }
+            new_word.insert(new_word.end(), word.begin() + i, it);
+            i = static_cast<int32_t>(std::distance(word.begin(), it));
+
+            if (word[i] == first && i < static_cast<int32_t>(word.size()) - 1 && word[i + 1] == second) {
+                new_word.push_back(first + second);
+                i += 2;
+            } else {
+                new_word.push_back(word[i]);
+                i += 1;
+            }
+        }
+
+        word = new_word;
+
+        if (word.size() == 1) {
+            break;
+        }
+        pairs = get_pairs(word);
+    }
+
+    return word;
+}
+
+std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t on_new_token_cb) {
+    std::string normalized_text = normalize(text);
+    std::vector<int32_t> bpe_tokens;
+    std::vector<std::string> token_strs;
+
+    auto splited_texts = split_with_special_tokens(normalized_text, special_tokens);
+
+    for (auto& splited_text : splited_texts) {
+        if (is_special_token(splited_text)) {
+            if (on_new_token_cb != nullptr) {
+                bool skip = on_new_token_cb(splited_text, bpe_tokens);
+                if (skip) {
+                    token_strs.push_back(splited_text);
+                    continue;
+                }
+            }
+            bpe_tokens.push_back(encoder[utf8_to_utf32(splited_text)]);
+            token_strs.push_back(splited_text);
+            continue;
+        }
+        auto tokens = token_split(splited_text);
+        for (auto& token : tokens) {
+            if (on_new_token_cb != nullptr) {
+                bool skip = on_new_token_cb(token, bpe_tokens);
+                if (skip) {
+                    token_strs.push_back(splited_text);
+                    continue;
+                }
+            }
+
+            std::string token_str = token;
+            std::u32string utf32_token;
+            if (byte_level_bpe) {
+                for (int i = 0; i < token_str.length(); i++) {
+                    unsigned char b = token_str[i];
+                    utf32_token += byte_encoder[b];
+                }
+            } else {
+                utf32_token = utf8_to_utf32(token_str);
+            }
+            auto bpe_strs = bpe(utf32_token);
+            for (auto bpe_str : bpe_strs) {
+                int token_id;
+                auto iter = encoder.find(bpe_str);
+                if (iter != encoder.end()) {
+                    token_id = iter->second;
+                } else {
+                    if (byte_fallback) {
+                        auto utf8_token_str = utf32_to_utf8(bpe_str);
+                        for (int i = 0; i < utf8_token_str.length(); i++) {
+                            unsigned char b = utf8_token_str[i];
+                            char hex_buf[16];
+                            snprintf(hex_buf, sizeof(hex_buf), "<0x%02X>", b);
+                            iter = encoder.find(utf8_to_utf32(hex_buf));
+                            bpe_tokens.push_back(token_id);
+                            token_strs.push_back(hex_buf);
+                        }
+                        continue;
+                    } else {
+                        token_id = UNK_TOKEN_ID;
+                    }
+                }
+                bpe_tokens.push_back(token_id);
+                token_strs.push_back(utf32_to_utf8(bpe_str));
+            }
+        }
+    }
+
+    std::stringstream ss;
+    ss << "[";
+    for (auto token : token_strs) {
+        ss << "\"" << token << "\", ";
+    }
+    ss << "]";
+    LOG_DEBUG("split prompt \"%s\" to tokens %s", text.c_str(), ss.str().c_str());
+    return bpe_tokens;
+}
+
+std::string BPETokenizer::decode_token(int token_id) const {
+    return utf32_to_utf8(decoder.at(token_id));
+}

src/tokenizers/bpe_tokenizer.h

@@ -0,0 +1,42 @@
+#ifndef __SD_TOKENIZERS_BPE_TOKENIZER_H__
+#define __SD_TOKENIZERS_BPE_TOKENIZER_H__
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <map>
+#include <regex>
+#include <set>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tokenizer.h"
+
+class BPETokenizer : public Tokenizer {
+protected:
+    std::map<int, std::u32string> byte_encoder;
+    std::map<std::u32string, int> byte_decoder;
+    std::map<std::u32string, int> encoder;
+    std::map<int, std::u32string> decoder;
+    std::map<std::pair<std::u32string, std::u32string>, int> bpe_ranks;
+    int encoder_len     = 0;
+    int bpe_len         = 0;
+    bool byte_level_bpe = true;
+    bool byte_fallback  = false;
+
+protected:
+    static std::vector<std::pair<int, std::u32string>> bytes_to_unicode();
+    static std::vector<std::u32string> split_utf32(const std::string& text, char32_t delimiter = U'\n');
+    virtual std::vector<std::string> token_split(const std::string& text) const;
+    std::vector<std::u32string> bpe(const std::u32string& token) const;
+    std::string decode_token(int token_id) const override;
+
+public:
+    BPETokenizer()          = default;
+    virtual ~BPETokenizer() = default;
+
+    std::vector<int> encode(const std::string& text, on_new_token_cb_t on_new_token_cb = nullptr) override;
+};
+
+#endif  // __SD_TOKENIZERS_BPE_TOKENIZER_H__

src/tokenizers/clip_tokenizer.cpp

@@ -0,0 +1,116 @@
+#include "clip_tokenizer.h"
+
+#include <algorithm>
+#include <cctype>
+#include <cmath>
+#include <regex>
+#include <set>
+
+#include "ggml.h"
+#include "tokenize_util.h"
+#include "util.h"
+#include "vocab/vocab.h"
+
+CLIPTokenizer::CLIPTokenizer(int pad_token_id, const std::string& merges_utf8_str) {
+    UNK_TOKEN = "<|endoftext|>";
+    BOS_TOKEN = "<|startoftext|>";
+    EOS_TOKEN = "<|endoftext|>";
+    PAD_TOKEN = "<|endoftext|>";
+
+    UNK_TOKEN_ID = 49407;
+    BOS_TOKEN_ID = 49406;
+    EOS_TOKEN_ID = 49407;
+    PAD_TOKEN_ID = pad_token_id;
+
+    end_of_word_suffix = "</w>";
+    add_bos_token      = true;
+    add_eos_token      = true;
+
+    if (merges_utf8_str.size() > 0) {
+        load_from_merges(merges_utf8_str);
+    } else {
+        load_from_merges(load_clip_merges());
+    }
+    add_special_token("<|startoftext|>");
+    add_special_token("<|endoftext|>");
+}
+
+void CLIPTokenizer::load_from_merges(const std::string& merges_utf8_str) {
+    auto byte_unicode_pairs = bytes_to_unicode();
+    byte_encoder            = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
+    for (auto& pair : byte_unicode_pairs) {
+        byte_decoder[pair.second] = pair.first;
+    }
+
+    std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
+    GGML_ASSERT(merges.size() == 48895);
+    merges = std::vector<std::u32string>(merges.begin() + 1, merges.end());
+    std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
+    for (const auto& merge : merges) {
+        size_t space_pos = merge.find(' ');
+        merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
+    }
+    std::vector<std::u32string> vocab;
+    for (const auto& pair : byte_unicode_pairs) {
+        vocab.push_back(pair.second);
+    }
+    for (const auto& pair : byte_unicode_pairs) {
+        vocab.push_back(pair.second + utf8_to_utf32("</w>"));
+    }
+    for (const auto& merge : merge_pairs) {
+        vocab.push_back(merge.first + merge.second);
+    }
+    vocab.push_back(utf8_to_utf32("<|startoftext|>"));
+    vocab.push_back(utf8_to_utf32("<|endoftext|>"));
+    LOG_DEBUG("vocab size: %llu", vocab.size());
+    int i = 0;
+    for (const auto& token : vocab) {
+        encoder[token] = i;
+        decoder[i]     = token;
+        i++;
+    }
+    encoder_len = i;
+
+    int rank = 0;
+    for (const auto& merge : merge_pairs) {
+        bpe_ranks[merge] = rank++;
+    }
+    bpe_len = rank;
+}
+
+static std::string strip(const std::string& str) {
+    std::string::size_type start = str.find_first_not_of(" \t\n\r\v\f");
+    std::string::size_type end   = str.find_last_not_of(" \t\n\r\v\f");
+
+    if (start == std::string::npos) {
+        return "";
+    }
+
+    return str.substr(start, end - start + 1);
+}
+
+static std::string whitespace_clean(const std::string& text) {
+    auto result = std::regex_replace(text, std::regex(R"(\s+)"), " ");
+    result      = strip(result);
+    return result;
+}
+
+std::string CLIPTokenizer::normalize(const std::string& text) const {
+    auto normalized_text = whitespace_clean(text);
+    std::transform(normalized_text.begin(), normalized_text.end(), normalized_text.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
+    return normalized_text;
+}
+
+std::vector<std::string> CLIPTokenizer::token_split(const std::string& text) const {
+    std::regex clip_pat(R"('s|'t|'re|'ve|'m|'ll|'d|[[:alpha:]]+|[[:digit:]]|[^[:space:][:alpha:][:digit:]]+)",
+                        std::regex::icase);
+    std::sregex_iterator iter(text.begin(), text.end(), clip_pat);
+    std::sregex_iterator end;
+
+    std::vector<std::string> result;
+    for (; iter != end; ++iter) {
+        result.emplace_back(iter->str());
+    }
+
+    return result;
+}

src/tokenizers/clip_tokenizer.h

@@ -0,0 +1,20 @@
+#ifndef __SD_TOKENIZERS_CLIP_TOKENIZER_H__
+#define __SD_TOKENIZERS_CLIP_TOKENIZER_H__
+
+#include <cstddef>
+#include <string>
+#include <vector>
+
+#include "bpe_tokenizer.h"
+
+class CLIPTokenizer : public BPETokenizer {
+protected:
+    void load_from_merges(const std::string& merges_utf8_str);
+    std::string normalize(const std::string& text) const override;
+    std::vector<std::string> token_split(const std::string& text) const override;
+
+public:
+    explicit CLIPTokenizer(int pad_token_id = 49407, const std::string& merges_utf8_str = "");
+};
+
+#endif  // __SD_TOKENIZERS_CLIP_TOKENIZER_H__

src/tokenizers/gemma_tokenizer.cpp

@@ -0,0 +1,190 @@
+#include "gemma_tokenizer.h"
+
+#include "ggml.h"
+#include "json.hpp"
+#include "util.h"
+#include "vocab/vocab.h"
+
+std::string GemmaTokenizer::normalize(const std::string& text) const {
+    std::string normalized = text;
+    size_t pos             = 0;
+    while ((pos = normalized.find(' ', pos)) != std::string::npos) {
+        normalized.replace(pos, 1, "\xE2\x96\x81");
+        pos += 3;
+    }
+    return normalized;
+}
+
+void GemmaTokenizer::load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
+    nlohmann::json vocab;
+    try {
+        vocab = nlohmann::json::parse(vocab_utf8_str);
+    } catch (const nlohmann::json::parse_error&) {
+        GGML_ABORT("invalid vocab json str");
+    }
+    for (const auto& [key, value] : vocab.items()) {
+        std::u32string token = utf8_to_utf32(key);
+        int i                = value;
+        encoder[token]       = i;
+        decoder[i]           = token;
+    }
+    encoder_len = static_cast<int>(vocab.size());
+    LOG_DEBUG("vocab size: %d", encoder_len);
+
+    std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
+    std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
+    for (const auto& merge : merges) {
+        size_t space_pos = merge.find(' ');
+        merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
+    }
+    LOG_DEBUG("merges size %zu", merge_pairs.size());
+
+    int rank = 0;
+    for (const auto& merge : merge_pairs) {
+        bpe_ranks[merge] = rank++;
+    }
+    bpe_len = rank;
+}
+
+GemmaTokenizer::GemmaTokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
+    byte_level_bpe = false;
+    byte_fallback  = true;
+    add_bos_token  = true;
+    PAD_TOKEN      = "<pad>";
+    EOS_TOKEN      = "<eos>";
+    BOS_TOKEN      = "<bos>";
+    UNK_TOKEN      = "<unk>";
+
+    PAD_TOKEN_ID = 0;
+    EOS_TOKEN_ID = 1;
+    BOS_TOKEN_ID = 2;
+    UNK_TOKEN_ID = 3;
+
+    std::vector<std::string> special_tokens_before_merge = {
+        PAD_TOKEN,
+        EOS_TOKEN,
+        BOS_TOKEN,
+        UNK_TOKEN,
+        "<mask>",
+        "[multimodal]",
+    };
+    for (int i = 0; i <= 98; i++) {
+        special_tokens_before_merge.push_back("<unused" + std::to_string(i) + ">");
+    }
+    special_tokens_before_merge.push_back("<start_of_turn>");
+    special_tokens_before_merge.push_back("<end_of_turn>");
+    for (int i = 1; i <= 31; i++) {
+        special_tokens_before_merge.push_back(std::string(i, '\n'));
+    }
+    for (int i = 2; i <= 31; i++) {
+        std::string whitespace_token;
+        for (int j = 0; j < i; j++) {
+            whitespace_token += "\xE2\x96\x81";
+        }
+        special_tokens_before_merge.push_back(whitespace_token);
+    }
+    std::vector<std::string> html_tokens = {
+        "<table>",
+        "<caption>",
+        "<thead>",
+        "<tbody>",
+        "<tfoot>",
+        "<tr>",
+        "<th>",
+        "<td>",
+        "</table>",
+        "</caption>",
+        "</thead>",
+        "</tbody>",
+        "</tfoot>",
+        "</tr>",
+        "</th>",
+        "</td>",
+        "<h1>",
+        "<h2>",
+        "<h3>",
+        "<h4>",
+        "<h5>",
+        "<h6>",
+        "<blockquote>",
+        "</h1>",
+        "</h2>",
+        "</h3>",
+        "</h4>",
+        "</h5>",
+        "</h6>",
+        "</blockquote>",
+        "<strong>",
+        "<em>",
+        "<b>",
+        "<i>",
+        "<u>",
+        "<s>",
+        "<sub>",
+        "<sup>",
+        "<code>",
+        "</strong>",
+        "</em>",
+        "</b>",
+        "</i>",
+        "</u>",
+        "</s>",
+        "</sub>",
+        "</sup>",
+        "</code>",
+        "<a>",
+        "<html>",
+        "<body>",
+        "<img>",
+        "<span>",
+        "<bbox>",
+        "<ul>",
+        "<li>",
+        "<div>",
+        "<iframe>",
+        "<footer>",
+        "</a>",
+        "</html>",
+        "</body>",
+        "</img>",
+        "</span>",
+        "</bbox>",
+        "</ul>",
+        "</li>",
+        "</div>",
+        "</iframe>",
+        "</footer>",
+    };
+    special_tokens_before_merge.insert(special_tokens_before_merge.end(),
+                                       html_tokens.begin(),
+                                       html_tokens.end());
+    for (int i = 0; i <= 0xFF; i++) {
+        char hex_buf[16];
+        snprintf(hex_buf, sizeof(hex_buf), "<0x%02X>", i);
+        special_tokens_before_merge.push_back(hex_buf);
+    }
+
+    std::vector<std::string> special_tokens_after_merge = {
+        "<start_of_image>",
+        "<end_of_image>",
+    };
+    for (int i = 1; i <= 31; i++) {
+        special_tokens_after_merge.insert(special_tokens_after_merge.begin() + i - 1,
+                                          std::string(i, '\t'));
+    }
+    for (int i = 99; i <= 6241; i++) {
+        special_tokens_after_merge.push_back("<unused" + std::to_string(i) + ">");
+    }
+    special_tokens_after_merge.push_back("<image_soft_token>");
+
+    special_tokens = special_tokens_before_merge;
+    special_tokens.insert(special_tokens.end(),
+                          special_tokens_after_merge.begin(),
+                          special_tokens_after_merge.end());
+
+    if (merges_utf8_str.size() > 0 && vocab_utf8_str.size() > 0) {
+        load_from_merges(merges_utf8_str, vocab_utf8_str);
+    } else {
+        load_from_merges(load_gemma_merges(), load_gemma_vocab_json());
+    }
+}

src/tokenizers/gemma_tokenizer.h

@@ -0,0 +1,17 @@
+#ifndef __SD_TOKENIZERS_GEMMA_TOKENIZER_H__
+#define __SD_TOKENIZERS_GEMMA_TOKENIZER_H__
+
+#include <string>
+
+#include "bpe_tokenizer.h"
+
+class GemmaTokenizer : public BPETokenizer {
+protected:
+    void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str);
+    std::string normalize(const std::string& text) const override;
+
+public:
+    explicit GemmaTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
+};
+
+#endif  // __SD_TOKENIZERS_GEMMA_TOKENIZER_H__

src/tokenizers/mistral_tokenizer.cpp

@@ -0,0 +1,89 @@
+#include "mistral_tokenizer.h"
+
+#include "ggml.h"
+#include "json.hpp"
+#include "util.h"
+#include "vocab/vocab.h"
+
+void MistralTokenizer::load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
+    nlohmann::json vocab;
+
+    try {
+        vocab = nlohmann::json::parse(vocab_utf8_str);
+    } catch (const nlohmann::json::parse_error&) {
+        GGML_ABORT("invalid vocab json str");
+    }
+    for (const auto& [key, value] : vocab.items()) {
+        std::u32string token = utf8_to_utf32(key);
+        int i                = value;
+        encoder[token]       = i;
+        decoder[i]           = token;
+    }
+    encoder_len = static_cast<int>(vocab.size());
+    LOG_DEBUG("vocab size: %d", encoder_len);
+
+    auto byte_unicode_pairs = bytes_to_unicode();
+    byte_encoder            = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
+    for (auto& pair : byte_unicode_pairs) {
+        byte_decoder[pair.second] = pair.first;
+    }
+    std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
+    LOG_DEBUG("merges size %llu", merges.size());
+    std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
+    for (const auto& merge : merges) {
+        size_t space_pos = merge.find(' ');
+        merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
+    }
+
+    int rank = 0;
+    for (const auto& merge : merge_pairs) {
+        bpe_ranks[merge] = rank++;
+    }
+    bpe_len = rank;
+}
+
+MistralTokenizer::MistralTokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
+    add_bos_token = true;
+
+    UNK_TOKEN = "<unk>";
+    BOS_TOKEN = "<s>";
+    EOS_TOKEN = "</s>";
+    PAD_TOKEN = "<pad>";
+
+    UNK_TOKEN_ID = 0;
+    BOS_TOKEN_ID = 1;
+    EOS_TOKEN_ID = 2;
+    PAD_TOKEN_ID = 11;
+
+    special_tokens = {
+        "<unk>",
+        "<s>",
+        "</s>",
+        "[INST]",
+        "[/INST]",
+        "[AVAILABLE_TOOLS]",
+        "[/AVAILABLE_TOOLS]",
+        "[TOOL_RESULTS]",
+        "[/TOOL_RESULTS]",
+        "[TOOL_CALLS]",
+        "[IMG]",
+        "<pad>",
+        "[IMG_BREAK]",
+        "[IMG_END]",
+        "[PREFIX]",
+        "[MIDDLE]",
+        "[SUFFIX]",
+        "[SYSTEM_PROMPT]",
+        "[/SYSTEM_PROMPT]",
+        "[TOOL_CONTENT]",
+    };
+    for (int i = 20; i < 1000; i++) {
+        special_tokens.push_back("<SPECIAL_" + std::to_string(i) + ">");
+    }
+
+    if (merges_utf8_str.size() > 0 && vocab_utf8_str.size() > 0) {
+        load_from_merges(merges_utf8_str, vocab_utf8_str);
+    } else {
+        load_from_merges(load_mistral_merges(), load_mistral_vocab_json());
+    }
+}

src/tokenizers/mistral_tokenizer.h

@@ -0,0 +1,16 @@
+#ifndef __SD_TOKENIZERS_MISTRAL_TOKENIZER_H__
+#define __SD_TOKENIZERS_MISTRAL_TOKENIZER_H__
+
+#include <string>
+
+#include "bpe_tokenizer.h"
+
+class MistralTokenizer : public BPETokenizer {
+protected:
+    void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str);
+
+public:
+    explicit MistralTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
+};
+
+#endif  // __SD_TOKENIZERS_MISTRAL_TOKENIZER_H__

src/tokenizers/qwen2_tokenizer.cpp

@@ -0,0 +1,91 @@
+#include "qwen2_tokenizer.h"
+
+#include "util.h"
+#include "vocab/vocab.h"
+
+void Qwen2Tokenizer::load_from_merges(const std::string& merges_utf8_str) {
+    auto byte_unicode_pairs = bytes_to_unicode();
+    byte_encoder            = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
+    for (auto& pair : byte_unicode_pairs) {
+        byte_decoder[pair.second] = pair.first;
+    }
+
+    std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
+    LOG_DEBUG("merges size %llu", merges.size());
+    std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
+    for (const auto& merge : merges) {
+        size_t space_pos = merge.find(' ');
+        merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
+    }
+
+    std::vector<std::u32string> tokens;
+    for (const auto& pair : byte_unicode_pairs) {
+        tokens.push_back(pair.second);
+    }
+    for (const auto& merge : merge_pairs) {
+        tokens.push_back(merge.first + merge.second);
+    }
+    for (auto& special_token : special_tokens) {
+        tokens.push_back(utf8_to_utf32(special_token));
+    }
+
+    int i = 0;
+    for (const auto& token : tokens) {
+        encoder[token] = i;
+        decoder[i]     = token;
+        i++;
+    }
+    encoder_len = i;
+    LOG_DEBUG("vocab size: %d", encoder_len);
+
+    int rank = 0;
+    for (const auto& merge : merge_pairs) {
+        bpe_ranks[merge] = rank++;
+    }
+    bpe_len = rank;
+}
+
+Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str) {
+    UNK_TOKEN = "<|endoftext|>";
+    EOS_TOKEN = "<|endoftext|>";
+    PAD_TOKEN = "<|endoftext|>";
+
+    UNK_TOKEN_ID = 151643;
+    EOS_TOKEN_ID = 151643;
+    PAD_TOKEN_ID = 151643;
+
+    special_tokens = {
+        "<|endoftext|>",
+        "<|im_start|>",
+        "<|im_end|>",
+        "<|object_ref_start|>",
+        "<|object_ref_end|>",
+        "<|box_start|>",
+        "<|box_end|>",
+        "<|quad_start|>",
+        "<|quad_end|>",
+        "<|vision_start|>",
+        "<|vision_end|>",
+        "<|vision_pad|>",
+        "<|image_pad|>",
+        "<|video_pad|>",
+        "<tool_call>",
+        "</tool_call>",
+        "<|fim_prefix|>",
+        "<|fim_middle|>",
+        "<|fim_suffix|>",
+        "<|fim_pad|>",
+        "<|repo_name|>",
+        "<|file_sep|>",
+        "<tool_response>",
+        "</tool_response>",
+        "<think>",
+        "</think>",
+    };
+
+    if (merges_utf8_str.size() > 0) {
+        load_from_merges(merges_utf8_str);
+    } else {
+        load_from_merges(load_qwen2_merges());
+    }
+}

src/tokenizers/qwen2_tokenizer.h

@@ -0,0 +1,16 @@
+#ifndef __SD_TOKENIZERS_QWEN2_TOKENIZER_H__
+#define __SD_TOKENIZERS_QWEN2_TOKENIZER_H__
+
+#include <string>
+
+#include "bpe_tokenizer.h"
+
+class Qwen2Tokenizer : public BPETokenizer {
+protected:
+    void load_from_merges(const std::string& merges_utf8_str);
+
+public:
+    explicit Qwen2Tokenizer(const std::string& merges_utf8_str = "");
+};
+
+#endif  // __SD_TOKENIZERS_QWEN2_TOKENIZER_H__

src/tokenizers/t5_unigram_tokenizer.cpp

@@ -0,0 +1,339 @@
+#include "t5_unigram_tokenizer.h"
+
+#include <algorithm>
+#include <cfloat>
+#include <cmath>
+#include <regex>
+#include <sstream>
+
+#include "json.hpp"
+#include "tokenize_util.h"
+#include "util.h"
+#include "vocab/vocab.h"
+
+// Port from: https://github.com/google/sentencepiece/blob/master/src/unigram_model.h
+// and https://github.com/google/sentencepiece/blob/master/src/unigram_model.h.
+// Original License: https://github.com/google/sentencepiece/blob/master/LICENSE
+//
+// Since tokenization is not the bottleneck in SD, performance was not a major consideration
+// during the migration.
+
+MetaspacePreTokenizer::MetaspacePreTokenizer(const std::string replacement, bool add_prefix_space)
+    : replacement(replacement), add_prefix_space(add_prefix_space) {}
+
+std::string MetaspacePreTokenizer::tokenize(const std::string& input) const {
+    std::string tokens;
+    std::stringstream ss(input);
+
+    if (add_prefix_space) {
+        tokens += replacement;
+    }
+
+    std::string token;
+    bool first_token = true;
+    while (std::getline(ss, token, ' ')) {
+        if (!first_token) {
+            tokens += replacement + token;
+        } else {
+            tokens += token;
+        }
+
+        first_token = false;
+    }
+
+    return tokens;
+}
+
+void T5UniGramTokenizer::InitializePieces(const std::string& json_str) {
+    nlohmann::json data;
+
+    try {
+        data = nlohmann::json::parse(json_str);
+    } catch (const nlohmann::json::parse_error&) {
+        status_ = INVLIAD_JSON;
+        return;
+    }
+    if (!data.contains("model")) {
+        status_ = INVLIAD_JSON;
+        return;
+    }
+    nlohmann::json model = data["model"];
+    if (!model.contains("vocab")) {
+        status_ = INVLIAD_JSON;
+        return;
+    }
+    if (model.contains("unk_id")) {
+        UNK_TOKEN_ID = model["unk_id"];
+    }
+
+    replacement      = data["pre_tokenizer"]["replacement"];
+    add_prefix_space = data["pre_tokenizer"]["add_prefix_space"];
+
+    pre_tokenizer = MetaspacePreTokenizer(replacement, add_prefix_space);
+
+    for (const auto& item : model["vocab"]) {
+        if (item.size() != 2 || !item[0].is_string() || !item[1].is_number_float()) {
+            status_ = INVLIAD_JSON;
+            return;
+        }
+        std::string piece = item[0];
+        if (piece.empty()) {
+            piece = "<empty_token>";
+        }
+        float score = item[1];
+        piece_score_pairs.emplace_back(piece, score);
+    }
+}
+
+void T5UniGramTokenizer::BuildTrie(std::vector<std::pair<std::string, int>>* pieces) {
+    if (status_ != OK) {
+        return;
+    }
+
+    if (pieces->empty()) {
+        status_ = NO_PIECES_LOADED;
+        return;
+    }
+
+    std::sort(pieces->begin(), pieces->end());
+
+    std::vector<const char*> key(pieces->size());
+    std::vector<int> value(pieces->size());
+    for (size_t i = 0; i < pieces->size(); ++i) {
+        key[i]   = (*pieces)[i].first.data();
+        value[i] = (*pieces)[i].second;
+    }
+
+    trie_ = std::unique_ptr<Darts::DoubleArray>(new Darts::DoubleArray());
+    if (trie_->build(key.size(), const_cast<char**>(&key[0]), nullptr, &value[0]) != 0) {
+        status_ = BUILD_DOUBLE_ARRAY_FAILED;
+        return;
+    }
+
+    const int kMaxTrieResultsSize = 1024;
+    std::vector<Darts::DoubleArray::result_pair_type> results(kMaxTrieResultsSize);
+    trie_results_size_ = 0;
+    for (const auto& p : *pieces) {
+        const size_t num_nodes = trie_->commonPrefixSearch(
+            p.first.data(), results.data(), results.size(), p.first.size());
+        trie_results_size_ = std::max(trie_results_size_, static_cast<int>(num_nodes));
+    }
+
+    if (trie_results_size_ == 0) {
+        status_ = NO_ENTRY_FOUND;
+    }
+}
+
+float T5UniGramTokenizer::GetScoreInlined(int id) const {
+    return piece_score_pairs[id].second;
+}
+
+bool T5UniGramTokenizer::IsUnusedInlined(int id) const {
+    (void)id;
+    return false;
+}
+
+bool T5UniGramTokenizer::IsUserDefinedInlined(int id) const {
+    (void)id;
+    return false;
+}
+
+size_t T5UniGramTokenizer::OneCharLen(const char* src) const {
+    return "\1\1\1\1\1\1\1\1\1\1\1\1\2\2\3\4"[(*src & 0xFF) >> 4];
+}
+
+EncodeResult T5UniGramTokenizer::EncodeOptimized(const std::string& normalized) const {
+    if (status() != OK || normalized.empty()) {
+        return {};
+    }
+
+    struct BestPathNode {
+        int id                = -1;
+        float best_path_score = 0;
+        int starts_at         = -1;
+    };
+
+    const int size        = static_cast<int>(normalized.size());
+    const float unk_score = min_score() - kUnkPenalty;
+    std::vector<BestPathNode> best_path_ends_at(size + 1);
+
+    int starts_at = 0;
+    while (starts_at < size) {
+        std::size_t node_pos                 = 0;
+        std::size_t key_pos                  = starts_at;
+        const auto best_path_score_till_here = best_path_ends_at[starts_at].best_path_score;
+        bool has_single_node                 = false;
+        const int mblen                      = std::min<int>(static_cast<int>(OneCharLen(normalized.data() + starts_at)), size - starts_at);
+        while (key_pos < static_cast<size_t>(size)) {
+            const int ret = trie_->traverse(normalized.data(), node_pos, key_pos, key_pos + 1);
+            if (ret == -2) {
+                break;
+            }
+            if (ret >= 0) {
+                if (IsUnusedInlined(ret)) {
+                    continue;
+                }
+                auto& target_node                    = best_path_ends_at[key_pos];
+                const auto length                    = static_cast<int>(key_pos - starts_at);
+                const auto score                     = IsUserDefinedInlined(ret) ? (length * max_score_ - 0.1f) : GetScoreInlined(ret);
+                const auto candidate_best_path_score = score + best_path_score_till_here;
+                if (target_node.starts_at == -1 || candidate_best_path_score > target_node.best_path_score) {
+                    target_node.best_path_score = static_cast<float>(candidate_best_path_score);
+                    target_node.starts_at       = starts_at;
+                    target_node.id              = ret;
+                }
+                if (!has_single_node && length == mblen) {
+                    has_single_node = true;
+                }
+            }
+        }
+        if (!has_single_node) {
+            auto& target_node                    = best_path_ends_at[starts_at + mblen];
+            const auto candidate_best_path_score = unk_score + best_path_score_till_here;
+            if (target_node.starts_at == -1 || candidate_best_path_score > target_node.best_path_score) {
+                target_node.best_path_score = candidate_best_path_score;
+                target_node.starts_at       = starts_at;
+                target_node.id              = UNK_TOKEN_ID;
+            }
+        }
+        starts_at += mblen;
+    }
+
+    EncodeResult results;
+    int ends_at = size;
+    while (ends_at > 0) {
+        const auto& node = best_path_ends_at[ends_at];
+        results.emplace_back(normalized.substr(node.starts_at, ends_at - node.starts_at), node.id);
+        ends_at = node.starts_at;
+    }
+    std::reverse(results.begin(), results.end());
+    return results;
+}
+
+T5UniGramTokenizer::T5UniGramTokenizer(bool is_umt5) {
+    add_bos_token = false;
+    add_eos_token = true;
+
+    if (is_umt5) {
+        PAD_TOKEN_ID = 0;
+        EOS_TOKEN_ID = 1;
+        BOS_TOKEN_ID = 2;
+        UNK_TOKEN_ID = 3;
+
+        PAD_TOKEN = "<pad>";
+        EOS_TOKEN = "</s>";
+        BOS_TOKEN = "<s>";
+        UNK_TOKEN = "<unk>";
+    } else {
+        PAD_TOKEN_ID = 0;
+        EOS_TOKEN_ID = 1;
+        UNK_TOKEN_ID = 2;
+
+        PAD_TOKEN = "<pad>";
+        EOS_TOKEN = "</s>";
+        UNK_TOKEN = "<unk>";
+    }
+
+    special_tokens = {
+        "<pad>",
+        "</s>",
+        "<unk>",
+    };
+
+    if (is_umt5) {
+        special_tokens.push_back("<s>");
+    }
+
+    if (is_umt5) {
+        InitializePieces(load_umt5_tokenizer_json());
+    } else {
+        InitializePieces(load_t5_tokenizer_json());
+    }
+
+    min_score_ = FLT_MAX;
+    max_score_ = FLT_MIN;
+
+    std::vector<std::pair<std::string, int>> pieces;
+    for (int i = 0; i < static_cast<int>(piece_score_pairs.size()); i++) {
+        const auto& sp = piece_score_pairs[i];
+
+        min_score_ = std::min(min_score_, sp.second);
+        max_score_ = std::max(max_score_, sp.second);
+
+        pieces.emplace_back(sp.first, i);
+    }
+
+    BuildTrie(&pieces);
+}
+
+T5UniGramTokenizer::~T5UniGramTokenizer() = default;
+
+std::string T5UniGramTokenizer::decode_token(int token_id) const {
+    if (token_id < 0 || token_id >= static_cast<int>(piece_score_pairs.size())) {
+        return "";
+    }
+
+    const std::string& piece = piece_score_pairs[token_id].first;
+    if (piece == "<empty_token>") {
+        return "";
+    }
+    return piece;
+}
+
+std::string T5UniGramTokenizer::normalize(const std::string& input) const {
+    // Ref: https://github.com/huggingface/tokenizers/blob/1ff56c0c70b045f0cd82da1af9ac08cd4c7a6f9f/bindings/python/py_src/tokenizers/implementations/sentencepiece_unigram.py#L29
+    // TODO: nmt-nfkc
+    std::string normalized = std::regex_replace(input, std::regex(" {2,}"), " ");
+    return normalized;
+}
+
+std::vector<int> T5UniGramTokenizer::encode(const std::string& input, on_new_token_cb_t on_new_token_cb) {
+    std::vector<int32_t> tokens;
+    std::vector<std::string> token_strs;
+    std::string normalized = normalize(input);
+    auto splited_texts     = split_with_special_tokens(normalized, special_tokens);
+    if (splited_texts.empty()) {
+        splited_texts.push_back(normalized);  // for empty string
+    }
+
+    for (auto& splited_text : splited_texts) {
+        if (is_special_token(splited_text)) {
+            if (on_new_token_cb != nullptr) {
+                bool skip = on_new_token_cb(splited_text, tokens);
+                if (skip) {
+                    token_strs.push_back(splited_text);
+                    continue;
+                }
+            }
+
+            if (splited_text == UNK_TOKEN) {
+                tokens.push_back(UNK_TOKEN_ID);
+                token_strs.push_back(UNK_TOKEN);
+            } else if (splited_text == EOS_TOKEN) {
+                tokens.push_back(EOS_TOKEN_ID);
+                token_strs.push_back(EOS_TOKEN);
+            } else if (splited_text == PAD_TOKEN) {
+                tokens.push_back(PAD_TOKEN_ID);
+                token_strs.push_back(PAD_TOKEN);
+            }
+            continue;
+        }
+
+        std::string pretokenized = pre_tokenizer.tokenize(splited_text);
+        EncodeResult result      = EncodeOptimized(pretokenized);
+        for (const auto& item : result) {
+            tokens.push_back(item.second);
+            token_strs.push_back(item.first);
+        }
+    }
+
+    std::stringstream ss;
+    ss << "[";
+    for (const auto& token_str : token_strs) {
+        ss << "\"" << token_str << "\", ";
+    }
+    ss << "]";
+    LOG_DEBUG("split prompt \"%s\" to tokens %s", input.c_str(), ss.str().c_str());
+
+    return tokens;
+}

src/tokenizers/t5_unigram_tokenizer.h

@@ -0,0 +1,70 @@
+#ifndef __SD_TOKENIZERS_T5_UNIGRAM_TOKENIZER_H__
+#define __SD_TOKENIZERS_T5_UNIGRAM_TOKENIZER_H__
+
+#include <cstddef>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "darts.h"
+#include "tokenizer.h"
+
+class MetaspacePreTokenizer {
+private:
+    std::string replacement;
+    bool add_prefix_space;
+
+public:
+    MetaspacePreTokenizer(const std::string replacement = " ", bool add_prefix_space = true);
+
+    std::string tokenize(const std::string& input) const;
+};
+
+using EncodeResult = std::vector<std::pair<std::string, int>>;
+
+class T5UniGramTokenizer : public Tokenizer {
+public:
+    enum Status {
+        OK,
+        NO_PIECES_LOADED,
+        NO_ENTRY_FOUND,
+        BUILD_DOUBLE_ARRAY_FAILED,
+        PIECE_ALREADY_DEFINED,
+        INVLIAD_JSON
+    };
+
+protected:
+    MetaspacePreTokenizer pre_tokenizer;
+    std::vector<std::pair<std::string, float>> piece_score_pairs;
+    float min_score_ = 0.0f;
+    float max_score_ = 0.0f;
+    std::unique_ptr<Darts::DoubleArray> trie_;
+    int trie_results_size_ = 0;
+    Status status_         = OK;
+    float kUnkPenalty      = 10.0f;
+    std::string replacement;
+    bool add_prefix_space = true;
+
+    void InitializePieces(const std::string& json_str);
+    void BuildTrie(std::vector<std::pair<std::string, int>>* pieces);
+    float GetScoreInlined(int id) const;
+    bool IsUnusedInlined(int id) const;
+    bool IsUserDefinedInlined(int id) const;
+    size_t OneCharLen(const char* src) const;
+    EncodeResult EncodeOptimized(const std::string& normalized) const;
+
+    float min_score() const { return min_score_; }
+    float max_score() const { return max_score_; }
+    Status status() const { return status_; }
+    std::string decode_token(int token_id) const override;
+    std::string normalize(const std::string& input) const override;
+
+public:
+    explicit T5UniGramTokenizer(bool is_umt5 = false);
+    ~T5UniGramTokenizer();
+
+    std::vector<int> encode(const std::string& input, on_new_token_cb_t on_new_token_cb = nullptr) override;
+};
+
+#endif  // __SD_TOKENIZERS_T5_UNIGRAM_TOKENIZER_H__

src/tokenizers/tokenize_util.h

@@ -1,5 +1,5 @@
-#ifndef __TOKENIZE_UTIL__
+#ifndef __SD_TOKENIZERS_BPE_TOKENIZE_UTIL_H__
-#define __TOKENIZE_UTIL__
+#define __SD_TOKENIZERS_BPE_TOKENIZE_UTIL_H__
 
 #include <string>
 #include <vector>
@@ -7,4 +7,4 @@
 std::vector<std::string> token_split(const std::string& text);
 std::vector<std::string> split_with_special_tokens(const std::string& text, const std::vector<std::string>& special_tokens);
 
-#endif  // __TOKENIZE_UTIL__
+#endif  // __SD_TOKENIZERS_BPE_TOKENIZE_UTIL_H__

src/tokenizers/tokenizer.cpp

@@ -0,0 +1,222 @@
+#include "tokenizer.h"
+
+#include <algorithm>
+#include <cmath>
+#include <regex>
+
+#include "util.h"
+
+void Tokenizer::add_special_token(const std::string& token) {
+    special_tokens.push_back(token);
+}
+
+bool Tokenizer::is_special_token(const std::string& token) const {
+    for (const auto& special_token : special_tokens) {
+        if (special_token == token) {
+            return true;
+        }
+    }
+    return false;
+}
+
+std::string Tokenizer::normalize(const std::string& text) const {
+    return text;
+}
+
+std::vector<int> Tokenizer::tokenize(const std::string& text,
+                                     on_new_token_cb_t on_new_token_cb,
+                                     bool padding,
+                                     size_t min_length,
+                                     size_t max_length,
+                                     bool allow_overflow_expand) {
+    std::vector<int> tokens = encode(text, on_new_token_cb);
+    if (padding) {
+        pad_tokens(tokens, nullptr, nullptr, min_length, max_length, allow_overflow_expand);
+    }
+    return tokens;
+}
+
+void Tokenizer::pad_tokens(std::vector<int>& tokens,
+                           std::vector<float>* weights,
+                           std::vector<float>* mask,
+                           size_t min_length,
+                           size_t max_length,
+                           bool allow_overflow_expand) {
+    const bool use_weights = weights != nullptr;
+    const bool use_mask    = mask != nullptr;
+
+    if (use_weights && tokens.size() != weights->size()) {
+        LOG_ERROR("tokens size != weights size");
+        return;
+    }
+
+    const size_t bos_count           = add_bos_token ? 1 : 0;
+    const size_t eos_count           = add_eos_token ? 1 : 0;
+    const size_t special_token_count = bos_count + eos_count;
+
+    auto build_sequence = [&](size_t begin,
+                              size_t count,
+                              size_t target_length,
+                              std::vector<int>& out_tokens,
+                              std::vector<float>& out_weights,
+                              std::vector<float>& out_mask) {
+        const size_t base_length  = count + special_token_count;
+        const size_t final_length = std::max(target_length, base_length);
+
+        out_tokens.clear();
+        out_weights.clear();
+        out_mask.clear();
+
+        out_tokens.reserve(final_length);
+        if (use_weights) {
+            out_weights.reserve(final_length);
+        }
+        if (use_mask) {
+            out_mask.reserve(final_length);
+        }
+
+        if (add_bos_token) {
+            out_tokens.push_back(BOS_TOKEN_ID);
+            if (use_weights) {
+                out_weights.push_back(1.0f);
+            }
+            if (use_mask) {
+                out_mask.push_back(1.0f);
+            }
+        }
+
+        for (size_t i = 0; i < count; ++i) {
+            out_tokens.push_back(tokens[begin + i]);
+            if (use_weights) {
+                out_weights.push_back((*weights)[begin + i]);
+            }
+            if (use_mask) {
+                out_mask.push_back(1.0f);
+            }
+        }
+
+        if (add_eos_token) {
+            out_tokens.push_back(EOS_TOKEN_ID);
+            if (use_weights) {
+                out_weights.push_back(1.0f);
+            }
+            if (use_mask) {
+                out_mask.push_back(1.0f);
+            }
+        }
+
+        if (final_length > out_tokens.size()) {
+            const size_t pad_count = final_length - out_tokens.size();
+            if (pad_left) {
+                out_tokens.insert(out_tokens.begin(), pad_count, PAD_TOKEN_ID);
+
+                if (use_weights) {
+                    out_weights.insert(out_weights.begin(), pad_count, 1.0f);
+                }
+                if (use_mask) {
+                    out_mask.insert(out_mask.begin(), pad_count, 0.0f);
+                }
+            } else {
+                out_tokens.insert(out_tokens.end(), pad_count, PAD_TOKEN_ID);
+
+                if (use_weights) {
+                    out_weights.insert(out_weights.end(), pad_count, 1.0f);
+                }
+                if (use_mask) {
+                    out_mask.insert(out_mask.end(), pad_count, 0.0f);
+                }
+            }
+        }
+    };
+
+    const size_t single_length    = std::max(min_length, tokens.size() + special_token_count);
+    const bool exceeds_max_length = max_length > 0 && single_length > max_length;
+
+    std::vector<int> new_tokens;
+    std::vector<float> new_weights;
+    std::vector<float> new_mask;
+
+    if (!exceeds_max_length) {
+        build_sequence(0, tokens.size(), min_length, new_tokens, new_weights, new_mask);
+    } else if (!allow_overflow_expand) {
+        build_sequence(0, tokens.size(), 0, new_tokens, new_weights, new_mask);
+
+        new_tokens.resize(max_length);
+        if (use_weights) {
+            new_weights.resize(max_length);
+        }
+        if (use_mask) {
+            new_mask.resize(max_length);
+        }
+
+        if (add_eos_token && !new_tokens.empty()) {
+            new_tokens.back() = EOS_TOKEN_ID;
+            if (use_weights) {
+                new_weights.back() = 1.0f;
+            }
+            if (use_mask) {
+                new_mask.back() = 1.0f;
+            }
+        }
+    } else if (min_length > special_token_count) {
+        const size_t tokens_per_chunk = min_length - special_token_count;
+        size_t offset                 = 0;
+
+        while (offset < tokens.size()) {
+            const size_t remaining = tokens.size() - offset;
+            const size_t take      = std::min(tokens_per_chunk, remaining);
+
+            std::vector<int> chunk_tokens;
+            std::vector<float> chunk_weights;
+            std::vector<float> chunk_mask;
+
+            build_sequence(offset, take, min_length, chunk_tokens, chunk_weights, chunk_mask);
+
+            new_tokens.insert(new_tokens.end(), chunk_tokens.begin(), chunk_tokens.end());
+            if (use_weights) {
+                new_weights.insert(new_weights.end(), chunk_weights.begin(), chunk_weights.end());
+            }
+            if (use_mask) {
+                new_mask.insert(new_mask.end(), chunk_mask.begin(), chunk_mask.end());
+            }
+
+            offset += take;
+        }
+    } else {
+        build_sequence(0, tokens.size(), min_length, new_tokens, new_weights, new_mask);
+    }
+
+    tokens = std::move(new_tokens);
+    if (use_weights) {
+        *weights = std::move(new_weights);
+    }
+    if (use_mask) {
+        *mask = std::move(new_mask);
+    }
+}
+
+static std::string clean_up_tokenization(std::string& text) {
+    std::regex pattern(R"( ,)");
+    return std::regex_replace(text, pattern, ",");
+}
+
+std::string Tokenizer::decode(const std::vector<int>& tokens) const {
+    std::string text;
+
+    for (int token_id : tokens) {
+        if (token_id == BOS_TOKEN_ID || token_id == EOS_TOKEN_ID || token_id == PAD_TOKEN_ID) {
+            continue;
+        }
+
+        std::string piece = decode_token(token_id);
+        if (!end_of_word_suffix.empty() && ends_with(piece, end_of_word_suffix)) {
+            piece.erase(piece.size() - end_of_word_suffix.size());
+            text += piece + " ";
+        } else {
+            text += piece;
+        }
+    }
+
+    text = clean_up_tokenization(text);
+    return trim(text);
+}

src/tokenizers/tokenizer.h

@@ -0,0 +1,53 @@
+#ifndef __SD_TOKENIZERS_TOKENIZER_H__
+#define __SD_TOKENIZERS_TOKENIZER_H__
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <string>
+#include <vector>
+
+using on_new_token_cb_t = std::function<bool(std::string&, std::vector<int32_t>&)>;
+
+class Tokenizer {
+protected:
+    std::vector<std::string> special_tokens;
+    bool add_bos_token = false;
+    bool add_eos_token = false;
+    bool pad_left      = false;
+    std::string end_of_word_suffix;
+
+    virtual std::string decode_token(int token_id) const = 0;
+    virtual std::string normalize(const std::string& text) const;
+
+public:
+    std::string UNK_TOKEN;
+    std::string BOS_TOKEN;
+    std::string EOS_TOKEN;
+    std::string PAD_TOKEN;
+    int UNK_TOKEN_ID = 0;
+    int BOS_TOKEN_ID = 0;
+    int EOS_TOKEN_ID = 0;
+    int PAD_TOKEN_ID = 0;
+
+    virtual ~Tokenizer() = default;
+
+    void add_special_token(const std::string& token);
+    bool is_special_token(const std::string& token) const;
+    virtual std::vector<int> encode(const std::string& text, on_new_token_cb_t on_new_token_cb = nullptr) = 0;
+    std::vector<int> tokenize(const std::string& text,
+                              on_new_token_cb_t on_new_token_cb = nullptr,
+                              bool padding                      = false,
+                              size_t min_length                 = 0,
+                              size_t max_length                 = 100000000,
+                              bool allow_overflow_expand        = false);
+    void pad_tokens(std::vector<int>& tokens,
+                    std::vector<float>* weights,
+                    std::vector<float>* mask,
+                    size_t min_length          = 0,
+                    size_t max_length          = 100000000,
+                    bool allow_overflow_expand = false);
+    std::string decode(const std::vector<int>& tokens) const;
+};
+
+#endif  // __SD_TOKENIZERS_TOKENIZER_H__

src/tokenizers/vocab/vocab.h

@@ -1,5 +1,5 @@
-#ifndef __VOCAB_H__
+#ifndef __SD_TOKENIZERS_VOCAB_VOCAB_H__
-#define __VOCAB_H__
+#define __SD_TOKENIZERS_VOCAB_VOCAB_H__
 
 #include <string>
 
@@ -12,4 +12,4 @@ std::string load_umt5_tokenizer_json();
 std::string load_gemma_merges();
 std::string load_gemma_vocab_json();
 
-#endif  // __VOCAB_H__
+#endif  // __SD_TOKENIZERS_VOCAB_VOCAB_H__

src/unet.hpp

@@ -217,11 +217,11 @@ public:
         } else if (sd_version_is_unet_edit(version)) {
             in_channels = 8;
         }
-        if (version == VERSION_SD1_TINY_UNET || version == VERSION_SD2_TINY_UNET || version == VERSION_SDXS) {
+        if (version == VERSION_SD1_TINY_UNET || version == VERSION_SD2_TINY_UNET || version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) {
             num_res_blocks = 1;
             channel_mult   = {1, 2, 4};
             tiny_unet      = true;
-            if (version == VERSION_SDXS) {
+            if (version == VERSION_SDXS_512_DS) {
                 attention_resolutions = {4, 2};  // here just like SDXL
             }
         }
@@ -264,6 +264,10 @@ public:
             if (version == VERSION_SVD) {
                 return new SpatialVideoTransformer(in_channels, n_head, d_head, depth, context_dim, use_linear_projection);
             } else {
+                if (version == VERSION_SDXS_09 && n_head == 5) {
+                    n_head = 1;    // to carry a special case of sdxs_09 into CrossAttentionLayer,
+                    d_head = 320;  // works as long the product remains equal (5*64 == 1*320)
+                }
                 return new SpatialTransformer(in_channels, n_head, d_head, depth, context_dim, use_linear_projection);
             }
         };

src/upscaler.cpp

@@ -1,27 +1,18 @@
-#include "esrgan.hpp"
+#include "upscaler.h"
 #include "ggml_extend.hpp"
 #include "model.h"
 #include "stable-diffusion.h"
 #include "util.h"
 
-struct UpscalerGGML {
+UpscalerGGML::UpscalerGGML(int n_threads,
-    ggml_backend_t backend    = nullptr;  // general backend
+                           bool direct,
-    ggml_type model_data_type = GGML_TYPE_F16;
+                           int tile_size)
-    std::shared_ptr<ESRGAN> esrgan_upscaler;
-    std::string esrgan_path;
-    int n_threads;
-    bool direct   = false;
-    int tile_size = 128;
-
-    UpscalerGGML(int n_threads,
-                 bool direct   = false,
-                 int tile_size = 128)
     : n_threads(n_threads),
       direct(direct),
       tile_size(tile_size) {
 }
 
-    bool load_from_file(const std::string& esrgan_path,
+bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
                                   bool offload_params_to_cpu,
                                   int n_threads) {
     ggml_log_set(ggml_log_callback_default, nullptr);
@@ -65,7 +56,7 @@ struct UpscalerGGML {
     return true;
 }
 
-    sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor) {
+sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_tensor) {
     sd::Tensor<float> upscaled;
     if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) {
         upscaled = esrgan_upscaler->compute(n_threads, input_tensor);
@@ -98,7 +89,7 @@ struct UpscalerGGML {
     return upscaled;
 }
 
-    sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor) {
+sd_image_t UpscalerGGML::upscale(sd_image_t input_image, uint32_t upscale_factor) {
     // upscale_factor, unused for RealESRGAN_x4plus_anime_6B.pth
     sd_image_t upscaled_image = {0, 0, 0, nullptr};
     int output_width          = (int)input_image.width * esrgan_upscaler->scale;
@@ -119,7 +110,6 @@ struct UpscalerGGML {
     upscaled_image = upscaled_data;
     return upscaled_image;
 }
-};
 
 struct upscaler_ctx_t {
     UpscalerGGML* upscaler = nullptr;

src/upscaler.h

@@ -0,0 +1,31 @@
+#ifndef __SD_UPSCALER_H__
+#define __SD_UPSCALER_H__
+
+#include "esrgan.hpp"
+#include "stable-diffusion.h"
+#include "tensor.hpp"
+
+#include <memory>
+#include <string>
+
+struct UpscalerGGML {
+    ggml_backend_t backend    = nullptr;  // general backend
+    ggml_type model_data_type = GGML_TYPE_F16;
+    std::shared_ptr<ESRGAN> esrgan_upscaler;
+    std::string esrgan_path;
+    int n_threads;
+    bool direct   = false;
+    int tile_size = 128;
+
+    UpscalerGGML(int n_threads,
+                 bool direct   = false,
+                 int tile_size = 128);
+
+    bool load_from_file(const std::string& esrgan_path,
+                        bool offload_params_to_cpu,
+                        int n_threads);
+    sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor);
+    sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor);
+};
+
+#endif  // __SD_UPSCALER_H__

src/vae.hpp

@@ -69,7 +69,7 @@ public:
         int scale_factor = 8;
         if (version == VERSION_WAN2_2_TI2V) {
             scale_factor = 16;
-        } else if (sd_version_is_flux2(version)) {
+        } else if (sd_version_uses_flux2_vae(version)) {
             scale_factor = 16;
         } else if (version == VERSION_CHROMA_RADIANCE) {
             scale_factor = 1;