fix: protect cross_attn and output_proj tokens for Anima LoRAs (#1786)

.github/workflows/close-inactive-issues.yml

@@ -0,0 +1,48 @@
+name: Close inactive issues
+
+on:
+  schedule:
+    # Run daily. GitHub cron schedules use UTC.
+    - cron: "30 1 * * *"
+  workflow_dispatch:
+    inputs:
+      debug_only:
+        description: "Dry run: log intended actions without changing issues"
+        required: false
+        default: false
+        type: boolean
+
+permissions:
+  issues: write
+
+concurrency:
+  group: ${{ github.workflow }}
+  cancel-in-progress: false
+
+jobs:
+  close-inactive-issues:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Comment and close inactive issues
+        uses: actions/stale@v10
+        with:
+          days-before-issue-stale: 365
+          days-before-issue-close: 0
+
+          days-before-pr-stale: -1
+          days-before-pr-close: -1
+
+          stale-issue-label: issue:inactive
+          close-issue-label: issue:auto-closed
+          close-issue-reason: completed
+          stale-issue-message: ""
+          close-issue-message: >
+            This issue has had no activity for one year. The latest version of
+            the code may already have fixed the problem.
+
+            If the issue still exists in the latest version, you can reopen
+            this issue at any time with updated reproduction details.
+
+          remove-issue-stale-when-updated: true
+          operations-per-run: 1000
+          debug-only: ${{ github.event_name == 'workflow_dispatch' && inputs.debug_only || false }}

.gitignore

@@ -13,3 +13,7 @@ output*.png
 models*
 *.log
 preview.png
+.claude/
+CLAUDE.local.md
+.agents/
+.codex/

AGENTS.md

@@ -0,0 +1,183 @@
+# Instructions for stable-diffusion.cpp
+
+This document is for AI coding agents working in this repository. It should
+describe agent-specific workflow, repository routing, editing boundaries, and
+project-specific pitfalls.
+
+For general contribution rules, including PR scope, commit conventions, code
+style, dependency updates, security hygiene, and AI-assisted contribution policy,
+see `CONTRIBUTING.md`.
+
+---
+
+## Agent Operating Rules
+
+Before analyzing or modifying the repository:
+
+1. Read this file.
+2. Use `rg` / `rg --files` or directory listing commands to confirm the current
+   tree before relying on a path.
+3. Start from `src/` and relevant `docs/` for runtime behavior.
+4. Read the relevant code before editing.
+5. Prefer the smallest change that fits the existing architecture.
+6. Report focused verification and mention any tests not run.
+
+Agents must not:
+
+* Run `git push`, create PRs, or submit issue/PR comments on the user's behalf.
+* Create commits unless the user explicitly requests that specific commit.
+* Modify `ggml/`, `thirdparty/`, or `examples/server/frontend/` unless
+  explicitly requested and necessary.
+* Read large local model files or tokenizer vocabulary files.
+* Rewrite unrelated code for style-only reasons.
+* Add secrets, model weights, generated binaries, local absolute paths, or
+  machine-specific output.
+
+When a change is large, architectural, or likely to affect public behavior,
+pause and present a short plan before editing.
+
+---
+
+## Repository Map and Editing Boundaries
+
+This is a routing map for agents, not a full architecture document. The layout
+can change, so verify paths before using them. Do not inspect excluded
+large-data directories while checking the tree.
+
+### Primary Project Code
+
+Core implementation lives under `src/`.
+
+Current source layout includes:
+
+* `src/core/` - shared tensor, ggml integration, backend, graph, RNG, and utility
+  code.
+* `src/model/` - model families and model components.
+* `src/model_io/` - model file loading, GGUF, safetensors, pickle, and related
+  serialization helpers.
+* `src/runtime/` - sampling, denoising, guidance, caching, preprocessing, and
+  runtime execution helpers.
+* `src/tokenizers/` - tokenizer implementations.
+* `src/conditioning/` - conditioning and prompt-related implementation.
+* `src/extensions/` - optional feature extensions.
+* top-level `src/*.cpp` and `src/*.h` files - public implementation entry
+  points, model loading, conversion, versioning, and shared managers.
+
+`src/tokenizers/vocab/` contains large tokenizer vocabulary data. Do not read or
+parse files in this directory; reference the path only when necessary.
+
+### Public API
+
+`include/` contains the C API exposed by the project. Currently the primary
+public header is `include/stable-diffusion.h`.
+
+Treat public headers as stable API. Avoid breaking compatibility unless the user
+explicitly requests it. If public behavior changes, update relevant examples or
+documentation.
+
+### Examples
+
+`examples/` contains programs demonstrating library usage.
+
+* `examples/cli/` - command line program for running models, testing features,
+  and debugging.
+* `examples/common/` - shared example support code.
+* `examples/server/` - server application built on top of the library.
+* `examples/server/frontend/` - git submodule containing independent frontend
+  code. Avoid modifying it unless explicitly requested.
+
+### Documentation and Tooling
+
+* `docs/` - documentation for supported models, build options, behavior, and
+  workflows.
+* `scripts/` - development, model processing, build automation, formatting, and
+  tooling scripts.
+* `cmake/` - CMake support modules.
+* `docker/` - Docker-related project files.
+* `assets/` - documentation assets; not runtime code.
+
+### External, Local, and Generated State
+
+* `ggml/` - git submodule for the ggml dependency.
+* `thirdparty/` - vendored third-party dependencies.
+* `models/` - local model storage. Ignore this directory and do not read model
+  files.
+* `test/` - local testing scripts. Use only when relevant to the task.
+* `build/`, `build_*`, and similar directories - generated build outputs.
+  Inspect them only when debugging a build result.
+
+---
+
+## Agent Workflow for Code Changes
+
+1. Identify the relevant modules under `src/`.
+2. Check whether the change touches the public API in `include/`.
+3. Consult relevant `docs/` and examples before changing user-facing behavior.
+4. Follow existing local patterns before adding new abstractions.
+5. Keep edits scoped to the requested behavior.
+6. Run the narrowest useful build, test, or inspection command available.
+
+Follow `CONTRIBUTING.md` for formatting, naming, PR expectations, dependency
+update policy, and security rules.
+
+---
+
+## Code Comments
+
+Keep comments rare and useful.
+
+Do not add comments that only describe what the code does. Add comments only
+when the code cannot fully express the logic, the logic is unusually complex, or
+there are historical reasons, invariants, constraints, compatibility concerns,
+or known pitfalls that future maintainers need to understand.
+
+Do not add task-specific comments that will be meaningless after review.
+
+Examples from the current codebase:
+
+```cpp
+// GOOD: explains a safety constraint that is not obvious from the assignment.
+// From src/model_io/pickle_io.cpp.
+// Non-tensor checkpoint metadata can use REDUCE for arbitrary
+// Python objects. Do not execute it; keep stack shape only.
+stack.push_back(make_none_value());
+
+// BAD: describes only what the next line does.
+// Set the token count to zero.
+token_count = 0;
+```
+
+---
+
+## Text File Encoding
+
+When reading or editing repository text files:
+
+* Prefer UTF-8 with LF for Markdown, frontend source, JSON, and other text-first
+  project files unless the file already clearly uses a different encoding.
+* Do not assume terminal output encoding matches file encoding on Windows.
+* A file that looks garbled in PowerShell output may still be valid UTF-8.
+* When inspecting UTF-8 files in PowerShell, prefer explicit UTF-8 reads such as:
+  * `[Console]::OutputEncoding = [System.Text.Encoding]::UTF8`
+  * `Get-Content -Encoding utf8 <path>`
+* Avoid rewriting a file purely because console output looked garbled; verify
+  the actual file encoding first.
+
+---
+
+## Tensor and Layout Notes
+
+Additional tensor/layout rules for this codebase:
+
+* `sd::Tensor` shape order is not PyTorch/NumPy-style. `shape()[0]` is the
+  lowest and most contiguous dimension, and higher indices are higher
+  dimensions.
+* Broadcasting for `sd::Tensor` must align dimensions from low to high dimension
+  indices. If one tensor has fewer dimensions, append implicit `1`s at the
+  higher-dimension end.
+* `ggml_n_dims` / `ggml_n_dims(tensor)` can drop trailing singleton high
+  dimensions. Do not assume a logical trailing dimension of `1` will still be
+  counted in ggml metadata.
+* Internal tensor-returning interfaces use an empty `sd::Tensor` to represent
+  null, absent, or failure states. Do not add `std::optional<sd::Tensor<...>>`
+  for internal APIs unless a distinct semantic state is truly required.

CLAUDE.md

@@ -0,0 +1,11 @@
+@AGENTS.md
+
+## Claude Code
+
+Follow `AGENTS.md` as the shared repository instructions.
+
+Do not duplicate contribution, style, PR, dependency, or security policy here;
+use `CONTRIBUTING.md` as the canonical source for those rules.
+
+Keep Claude-specific project notes in this file only when they do not apply to
+other coding agents.

README.md

@@ -57,7 +57,7 @@ API and command-line option may change frequently.***
     - [SeFi-Image](./docs/sefi_image.md)
     - [HiDream-O1-Image](./docs/hidream_o1_image.md)
     - [Ideogram4](./docs/ideogram4.md)
-  - Image Edit Models
+  - [Image Edit Models](./docs/edit.md)
     - [FLUX.1-Kontext-dev](./docs/kontext.md)
     - [Qwen Image Edit series](./docs/qwen_image_edit.md)
     - [LongCat Image Edit](./docs/longcat_image.md)
@@ -65,6 +65,7 @@ API and command-line option may change frequently.***
   - Video Models
     - [Wan2.1/Wan2.2](./docs/wan.md)
     - [LTX-2.3](./docs/ltx2.md)
+    - [LingBot-Video](./docs/lingbot_video.md)
   - [PhotoMaker](./docs/photo_maker.md) support.
   - Control Net support with SD 1.5
   - LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora)

docs/edit.md

@@ -0,0 +1,93 @@
+# Image Editing
+
+Image editing in `stable-diffusion.cpp` allows you to use reference images to guide the generation process, enabling tasks like identity preservation, style transfer, or layout modification.
+
+
+## Supported Models
+
+Depending on the architecture, different models handle reference images differently.
+
+| Model | Default Preset |
+| :--- | :--- |
+| [**FLUX.1-Kontext-dev**](./kontext.md) | `flux_kontext` |
+| [**LongCat Image Edit**](./longcat_image.md) | `longcat` |
+| [**Qwen Image Edit**](./qwen_image_edit.md) | `qwen` |
+| **Qwen Image LAYERED** | `qwen_layered` |
+| [**Flux.2 [Dev] / Flux.2 [Klein]**](./flux2.md) | `flux2` |
+| [**Boogu Image Edit**](./boogu_image.md) | `z_image_omni` |
+| **Krea2 (Community Edit LoRAs)** | `krea2_ostris_edit` |
+| **Anima (Community Edit LoRAs)** | `cosmos_reference` |
+
+Stable-diffusion.spp also supports basic Unet-based editing models like instruct-pix2pix or CosXL-Edit. This document is not about those.
+
+---
+
+## Configuring Reference Modes (`--ref-image-args`)
+
+Different DiT-based editing models require different configurations to process reference images correctly (e.g., whether to use a Vision Language Model (VLM) encoder or pass VAE-encoded images directly to the DiT).
+
+To simplify this, we provide **Presets**. By default, the system automatically selects the best preset based on the model architecture. However, you can override this using the `--ref-image-args` argument.
+
+### Usage
+The `--ref-image-args` argument accepts a comma-separated list of key-value pairs:
+
+**Using a preset:**
+`--ref-image-args "preset=qwen_layered"`
+
+**Using a preset with a specific override:**
+`--ref-image-args "preset=krea2_edit,force_ref_timestep_zero=true"`
+
+### Available Presets
+
+| Preset | Primary Use Case |
+| :--- | :--- |
+| `flux_kontext` | FLUX.1 Kontext |
+| `longcat` | LongCat Image Edit |
+| `flux2` | FLUX.2 models |
+| `qwen` | Qwen Image Edit |
+| `qwen_layered` | Qwen Image Layered |
+| `z_image_omni` | Boogu, Z-Image Omni |
+| `krea2_ostris_edit` | Most Krea2 Community edit LoRAs (trained with Ostris script) |
+| `krea2_edit` | Specifically for [lbouaraba/krea2edit](https://huggingface.co/conradlocke/krea2-identity-edit). (or similar) |
+| `cosmos_reference` | For Anima |
+| `default` | Uses the automatic detection based on model architecture. |
+
+---
+
+## Advanced Parameter Reference
+
+If presets are insufficient, you can manually configure the following parameters via `--ref-image-args`:
+
+| Key | Type | Description | Allowed Values |
+| :--- | :--- | :--- | :--- |
+| `preset` | string | Overrides the automatic preset. | (See the Presets table above) |
+| `pass_to_vlm` | bool | Whether reference images are passed to the VLM encoder. | `true`, `false` |
+| `pass_to_dit` | bool | Whether VAE-encoded references are passed directly to the DiT. | `true`, `false` |
+| `ref_index_mode` | string | Behavior of the RoPE index. | `fixed`, `increase`, `decrease` |
+| `force_ref_timestep_zero` | bool | Forces timestep=0 for reference tokens. | `true`, `false` (Krea2 only) |
+| `resize_before_vae` | bool | Whether reference images are resized before VAE encoding. | `true`, `false` |
+| `vae_input_max_pixels` | int | Maximum pixel area for VAE reference inputs. | Integer |
+| `vlm_resize_mode` | string | How to resize VLM reference inputs. | `longest_side`, `area`, `none` |
+| `vlm_max_size` | int | Maximum VLM input size; interpreted according to `vlm_resize_mode`. | Integer |
+| `vlm_min_size` | int | Minimum VLM input size; interpreted according to `vlm_resize_mode`. | Integer |
+| `vlm_size` | int | Shortcut to set both VLM min and max size to the same value. | Integer |
+
+### Preset Default Values
+
+For a technical overview of how each preset is configured, see the table below.
+
+| Preset | VLM | RoPE Index | Cond Resize | Special Notes |
+| :--- | :---: | :---: | :---: | :--- |
+| `flux_kontext` | No | `fixed` | `none` | |
+| `longcat` | Yes | `fixed` | `area` | |
+| `flux2` | No | `increase` | `none` | |
+| `qwen` | Yes | `increase` | `area` | |
+| `qwen_layered` | Yes | `decrease` | `area` | |
+| `z_image_omni` | Yes | `fixed` | `area` | |
+| `krea2_ostris_edit`| Yes | `increase` | `area` | `force_ref_timestep_zero = true` |
+| `krea2_edit` | Yes | `increase` | `longest` | `vlm_size = 768` |
+| `cosmos_reference` | No | `fixed` | `none` | `resize_before_vae = false` |
+
+**Additional Default Notes:**
+- **VLM Input Sizes:** For most presets, `vlm_max_size` and `vlm_min_size` are set to `-1`, meaning the values are model-dependent and handled automatically. In `area` mode they represent pixel area; in `longest_side` mode they represent a side length in pixels.
+- **VAE Input Size:** `vae_input_max_pixels` defaults to $1024 \times 1024$ pixels (`1048576`).

docs/lingbot_video.md

@@ -0,0 +1,32 @@
+# How to Use
+
+Lingbot Video uses a Lingbot diffusion transformer, the Wan2.1 VAE, and Qwen3-VL 4B as the LLM text encoder.
+
+## Download weights
+
+- Download lingbot-video-dense-1.3b
+    - safetensors: https://huggingface.co/robbyant/lingbot-video-dense-1.3b/tree/main/transformer
+- Download lingbot-video-moe-30b-a3b
+    - safetensors: https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b/tree/main/transformer
+- Download vae
+    - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors
+- Download Qwen3-VL 4B
+    - safetensors: https://huggingface.co/Comfy-Org/Krea-2/tree/main/text_encoders
+    - gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main
+
+## Notes
+
+- Use `-M vid_gen`.
+- T2V uses the text prompt only.
+- I2V uses `-i` as the first video frame. The same image is also passed to the
+  Qwen3-VL prompt enhancer when vision weights are available.
+- Video frames are aligned to Wan-style temporal compression, so use frame counts
+  like 33, 49, or 81.
+
+## Examples
+
+### LingBot-Video T2V
+
+```sh
+.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\lingbot-video-dens-1.3b.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --vae ..\models\vae\wan_2.1_vae.safetensors -p '{"caption":{"comprehensive_description":"A lovely cat sits comfortably on a soft cushion near a sunlit window, looking calm, gentle, and adorable. The cat has soft fluffy fur, bright expressive eyes, small rounded ears, delicate whiskers, and a relaxed posture. Warm daylight falls across the cat from one side, creating soft highlights on the fur and a cozy glow around the scene. The background is softly blurred, showing hints of a peaceful indoor home environment with warm tones and gentle natural light. The overall atmosphere is cute, tender, serene, and photorealistic, emphasizing the cat''s charming appearance, soft texture, and affectionate presence.","camera_info":{"color":"Warm","frame_size":"Close Up","shot_type_angle":"Eye level","lens_size":"Medium Lens","composition":"Centered balanced","lighting":"Soft light","lighting_type":"Daylight"},"world_knowledge":[],"prominent_elements":[{"name":"lovely cat","description":"A cute and gentle domestic cat with soft fluffy fur, expressive eyes, and a calm relaxed presence.","location":"center of the frame","relative_size":"large","shape_and_color":"Small animal body with rounded face, triangular ears, bright eyes, and soft fur in warm natural tones","texture":"soft, fluffy, silky","appearance_details":"The cat has clean well-groomed fur, delicate whiskers, small ears, a cute nose, and bright attentive eyes. Its expression appears calm, affectionate, and slightly curious.","relationship":"Acts as the main subject and emotional focal point of the scene.","orientation":"facing the camera","pose":"sitting comfortably with a relaxed posture","expression":"gentle, adorable, calm, slightly curious","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"cat eyes","description":"Bright expressive eyes that give the cat a sweet and affectionate appearance.","location":"upper center of the cat face","relative_size":"small","shape_and_color":"Round almond-like eyes with glossy reflections","texture":"clear, glossy, reflective","appearance_details":"The eyes catch the soft daylight, creating small natural highlights that make the cat look vivid and alive.","relationship":"Enhance the emotional charm and cuteness of the cat.","orientation":"looking toward the camera","pose":"","expression":"soft and attentive","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"soft cushion","description":"A comfortable cushion or blanket where the cat is resting.","location":"bottom portion of the frame","relative_size":"medium","shape_and_color":"Soft rounded fabric surface in light warm neutral tones","texture":"plush, fabric, cozy","appearance_details":"The cushion gently supports the cat and adds a comfortable home-like feeling to the scene.","relationship":"Provides a cozy resting place for the cat.","orientation":"horizontal","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"sunlit indoor background","description":"A softly blurred indoor background with warm daylight and peaceful home atmosphere.","location":"behind the cat, filling the upper and side areas of the frame","relative_size":"large","shape_and_color":"Soft abstract shapes in warm beige, cream, and pale golden tones","texture":"soft, blurry, bokeh-like","appearance_details":"The background is intentionally out of focus, keeping attention on the cat while creating a cozy and serene mood.","relationship":"Provides a warm and gentle environment that supports the cute domestic scene.","orientation":"upright","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":"","is_cluster":true,"number_of_objects":"numerous"}]}}' -n '{"universal_negative":{"visual_quality":["low quality","worst quality","blurry","pixelated","jpeg artifacts","low resolution","unstable color","color flicker","underexposed","overexposed","invisible subject","subject hidden in darkness"],"artistic_style":["painting","illustration","drawing","cartoon","3d render","cgi","sketch","digital art"],"composition_and_content":["text","watermark","signature","logo","subtitles","pillarboxed","side bars","portrait image in landscape frame"],"temporal_and_motion_stability":["flickering","jittery","motion blur","temporal inconsistency","warping","morphing","incoherent motion","unnatural movement","static object with sudden jump","frame-to-frame inconsistency"],"material_and_structure":["plastic-like glass","unrealistic texture","deformed bottle","liquid freezing improperly","distorted reflections"]}}' --diffusion-fa --offload-to-cpu --cfg-scale 3 --video-frames 33 -v
+```

examples/common/common.cpp

@@ -975,6 +975,11 @@ ArgOptions SDGenerationParams::get_options() {
          "extra VAE tiling args, key=value list. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)",
          (int)',',
          &extra_tiling_args},
+        {"",
+         "--ref-image-args",
+         "Key-value list to set up the way the reference images are processed (empty = auto-detect from model weigths)",
+         (int)',',
+         &ref_image_args},
     };
 
     options.int_options = {
@@ -2418,30 +2423,45 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
         pulid_id_weight,
     };
 
-    params.loras                 = lora_vec.empty() ? nullptr : lora_vec.data();
-    params.lora_count            = static_cast<uint32_t>(lora_vec.size());
-    params.prompt                = prompt.c_str();
-    params.negative_prompt       = negative_prompt.c_str();
-    params.clip_skip             = clip_skip;
-    params.init_image            = init_image.get();
-    params.ref_images            = ref_image_views.empty() ? nullptr : ref_image_views.data();
-    params.ref_images_count      = static_cast<int>(ref_image_views.size());
-    params.auto_resize_ref_image = auto_resize_ref_image;
-    params.increase_ref_index    = increase_ref_index;
-    params.mask_image            = mask_image.get();
-    params.width                 = get_resolved_width();
-    params.height                = get_resolved_height();
-    params.sample_params         = sample_params;
-    params.strength              = strength;
-    params.seed                  = seed;
-    params.batch_count           = batch_count;
-    params.qwen_image_layers     = qwen_image_layers;
-    params.control_image         = control_image.get();
-    params.control_strength      = control_strength;
-    params.pm_params             = pm_params;
-    params.pulid_params          = pulid_params;
-    params.vae_tiling_params     = vae_tiling_params;
-    params.cache                 = cache_params;
+    if (!auto_resize_ref_image) {
+        if (!ref_image_args.empty()) {
+            ref_image_args += ",";
+        }
+        ref_image_args += "resize_before_vae=0";
+        LOG_WARN("Notice: --disable-auto-resize-ref-image is deprecated. Use --ref-image-args \"resize_before_vae=off\" instead.");
+    }
+
+    if (increase_ref_index) {
+        if (!ref_image_args.empty()) {
+            ref_image_args += ",";
+        }
+        ref_image_args += "ref_index_mode=increase";
+        LOG_WARN("Notice: --increase-ref-index is deprecated. Use --ref-image-args \"ref_index_mode=increase\" instead.");
+    }
+
+    params.loras             = lora_vec.empty() ? nullptr : lora_vec.data();
+    params.lora_count        = static_cast<uint32_t>(lora_vec.size());
+    params.prompt            = prompt.c_str();
+    params.negative_prompt   = negative_prompt.c_str();
+    params.clip_skip         = clip_skip;
+    params.init_image        = init_image.get();
+    params.ref_images        = ref_image_views.empty() ? nullptr : ref_image_views.data();
+    params.ref_images_count  = static_cast<int>(ref_image_views.size());
+    params.ref_image_args    = ref_image_args.c_str();
+    params.mask_image        = mask_image.get();
+    params.width             = get_resolved_width();
+    params.height            = get_resolved_height();
+    params.sample_params     = sample_params;
+    params.strength          = strength;
+    params.seed              = seed;
+    params.batch_count       = batch_count;
+    params.qwen_image_layers = qwen_image_layers;
+    params.control_image     = control_image.get();
+    params.control_strength  = control_strength;
+    params.pm_params         = pm_params;
+    params.pulid_params      = pulid_params;
+    params.vae_tiling_params = vae_tiling_params;
+    params.cache             = cache_params;
 
     params.hires.enabled             = hires_enabled;
     params.hires.upscaler            = resolved_hires_upscaler;

examples/common/common.h

@@ -227,6 +227,8 @@ struct SDGenerationParams {
     sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
     std::string extra_tiling_args;
 
+    std::string ref_image_args;
+
     std::string pm_id_images_dir;
     std::string pm_id_embed_path;
     float pm_style_strength = 20.f;

include/stable-diffusion.h

@@ -363,8 +363,7 @@ typedef struct {
     sd_image_t init_image;
     sd_image_t* ref_images;
     int ref_images_count;
-    bool auto_resize_ref_image;
-    bool increase_ref_index;
+    const char* ref_image_args;
     sd_image_t mask_image;
     int width;
     int height;
@@ -428,6 +427,11 @@ 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);
 
+// ControlNet hot-swap APIs are not safe to call while generation is in flight.
+SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path);
+SD_API bool sd_ctx_unload_control_net(sd_ctx_t* sd_ctx);
+SD_API bool sd_ctx_has_control_net(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);
 SD_API const char* sd_rng_type_name(enum rng_type_t rng_type);

src/conditioning/conditioner.hpp

@@ -7,6 +7,7 @@
 
 #include "core/tensor_ggml.hpp"
 #include "core/util.h"
+#include "model/diffusion/model.hpp"
 #include "model/te/clip.hpp"
 #include "model/te/llm.hpp"
 #include "model/te/t5.hpp"
@@ -106,6 +107,7 @@ struct ConditionerParams {
     int height                                       = -1;
     bool zero_out_masked                             = false;
     const std::vector<sd::Tensor<float>>* ref_images = nullptr;  // for qwen image edit
+    RefImageParams ref_image_params;
 };
 
 struct Conditioner {
@@ -1797,7 +1799,7 @@ struct LLMEmbedder : public Conditioner {
             arch = LLM::LLMArch::GPT_OSS_20B;
         } else if (sd_version_is_pid(version)) {
             arch = LLM::LLMArch::GEMMA2_2B;
-        } else if (sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) {
+        } else if (sd_version_is_lingbot_video(version) || sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) {
             arch = LLM::LLMArch::QWEN3_VL;
         } else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
             arch = LLM::LLMArch::QWEN3;
@@ -1993,6 +1995,54 @@ struct LLMEmbedder : public Conditioner {
         return new_hidden_states;
     }
 
+    void resize_image_dims(int height, int width, int& h_bar, int& w_bar, int factor, int min_size, int max_size, RefImageResizeMode mode) {
+        if (min_size > 0 && min_size == max_size) {
+            if (mode == RefImageResizeMode::AREA) {
+                double beta = std::sqrt(static_cast<double>(min_size) / (static_cast<double>(height) * width));
+                h_bar       = std::max(static_cast<int>(factor),
+                                       static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
+                w_bar       = std::max(static_cast<int>(factor),
+                                       static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
+            } else if (mode == RefImageResizeMode::LONGEST_SIDE) {
+                int current_max_side = std::max(height, width);
+                double beta          = static_cast<double>(min_size) / current_max_side;
+                h_bar                = std::max(static_cast<int>(factor),
+                                                static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
+                w_bar                = std::max(static_cast<int>(factor),
+                                                static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
+            }
+            return;
+        }
+
+        if (mode == RefImageResizeMode::AREA) {
+            double current_area = static_cast<double>(h_bar) * w_bar;
+            if (max_size > 0 && current_area > max_size) {
+                double beta = std::sqrt((static_cast<double>(height) * width) / static_cast<double>(max_size));
+                h_bar       = std::max(static_cast<int>(factor),
+                                       static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
+                w_bar       = std::max(static_cast<int>(factor),
+                                       static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
+            } else if (min_size > 0 && current_area < min_size) {
+                double beta = std::sqrt(static_cast<double>(min_size) / (static_cast<double>(height) * width));
+                h_bar       = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
+                w_bar       = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
+            }
+        } else if (mode == RefImageResizeMode::LONGEST_SIDE) {
+            int current_max_side = std::max(height, width);
+            if (max_size > 0 && current_max_side > max_size) {
+                double beta = static_cast<double>(max_size) / current_max_side;
+                h_bar       = std::max(static_cast<int>(factor),
+                                       static_cast<int>(std::floor(height * beta / factor)) * static_cast<int>(factor));
+                w_bar       = std::max(static_cast<int>(factor),
+                                       static_cast<int>(std::floor(width * beta / factor)) * static_cast<int>(factor));
+            } else if (min_size > 0 && current_max_side < min_size) {
+                double beta = static_cast<double>(min_size) / current_max_side;
+                h_bar       = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
+                w_bar       = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
+            }
+        }
+    }
+
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
         std::string prompt;
@@ -2007,38 +2057,126 @@ struct LLMEmbedder : public Conditioner {
         bool spell_quotes                    = false;
         std::set<int> out_layers;
 
-        int64_t t0 = ggml_time_ms();
+        int64_t t0                     = ggml_time_ms();
+        RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode;
 
-        if (sd_version_is_qwen_image(version)) {
+        if (sd_version_is_lingbot_video(version)) {
+            const int pad_token = 151643;
+            const std::string prompt_prefix =
+                "<|im_start|>system\nGiven a user input that may include a text prompt alone, "
+                "a text prompt with an image reference, or a text prompt with a video reference "
+                "or a video reference alone, generate an \"Enhanced prompt\" that provides detailed "
+                "visual descriptions suitable for video generation. Evaluate the level of detail "
+                "in the user's input: if it is simple, enrich it by adding specifics about colors, "
+                "shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal "
+                "progression, and spatial relationships to create vivid, concrete, and temporally "
+                "coherent scenes to create vivid and concrete scenes. Please generate only the "
+                "enhanced description for the prompt below and avoid including any additional "
+                "commentary or evaluations:<|im_end|>\n<|im_start|>user\n";
+
+            auto prefix_tokens               = tokenizer->encode(prompt_prefix, nullptr);
+            prompt_template_encode_start_idx = 0;
+            for (int token : prefix_tokens) {
+                if (token != pad_token) {
+                    prompt_template_encode_start_idx++;
+                }
+            }
+            LOG_DEBUG("prompt_template_encode_start_idx %d", prompt_template_encode_start_idx);
+
+            prompt = prompt_prefix;
+            if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
+                LOG_INFO("LingBotVideoI2VPipeline");
+                const std::string placeholder = "<|image_pad|>";
+                std::string img_prompt;
+
+                for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
+                    const auto& image = (*conditioner_params.ref_images)[i];
+                    const int factor  = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
+                    int height        = static_cast<int>(image.shape()[1]);
+                    int width         = static_cast<int>(image.shape()[0]);
+
+                    int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
+                    if (min_pixels <= 0) {
+                        if (resize_mode == RefImageResizeMode::AREA) {
+                            min_pixels = static_cast<int>(4 * factor * factor);
+                        } else {
+                            min_pixels = static_cast<int>(2 * factor);
+                        }
+                    }
+                    int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
+                    if (max_pixels <= 0) {
+                        if (resize_mode == RefImageResizeMode::AREA) {
+                            max_pixels = static_cast<int>(16384 * factor * factor);
+                        } else {
+                            max_pixels = static_cast<int>(128 * factor);
+                        }
+                    }
+
+                    int h_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(height) / factor) * factor));
+                    int w_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(width) / factor) * factor));
+
+                    if (std::max(height, width) > 200 * std::min(height, width)) {
+                        LOG_WARN("LingBotVideo image aspect ratio is very large: %dx%d", width, height);
+                    }
+
+                    resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
+
+                    LOG_DEBUG("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
+                    auto resized_image = clip_preprocess(image, w_bar, h_bar);
+                    auto image_embed   = llm->encode_image(n_threads, resized_image, false, true, true);
+                    GGML_ASSERT(!image_embed.empty());
+
+                    std::string image_prefix = prompt + img_prompt + "<|vision_start|>";
+                    int image_embed_idx      = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
+                    image_embeds.emplace_back(image_embed_idx, image_embed);
+
+                    img_prompt += "<|vision_start|>";
+                    int64_t num_image_tokens = image_embed.shape()[1];
+                    img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
+                    for (int j = 0; j < num_image_tokens; j++) {
+                        img_prompt += placeholder;
+                    }
+                    img_prompt += "<|vision_end|>";
+                }
+                prompt += img_prompt;
+            }
+
+            prompt += conditioner_params.text;
+            prompt_attn_range = {0, 0};
+            prompt += "<|im_end|>\n<|im_start|>assistant\n";
+        } else if (sd_version_is_qwen_image(version)) {
             if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
                 LOG_INFO("QwenImageEditPlusPipeline");
                 prompt_template_encode_start_idx = 64;
                 int image_embed_idx              = 64 + 6;
 
-                int min_pixels          = 384 * 384;
-                int max_pixels          = 560 * 560;
+                int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
+                if (min_pixels <= 0) {
+                    min_pixels = 384;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        min_pixels *= min_pixels;
+                    }
+                }
+                int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
+                if (max_pixels <= 0) {
+                    max_pixels = 560;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        max_pixels *= max_pixels;
+                    }
+                }
+
                 std::string placeholder = "<|image_pad|>";
                 std::string img_prompt;
 
                 for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
                     const auto& image = (*conditioner_params.ref_images)[i];
-                    double factor     = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
+                    const int factor  = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
                     int height        = static_cast<int>(image.shape()[1]);
                     int width         = static_cast<int>(image.shape()[0]);
-                    int h_bar         = static_cast<int>(std::round(height / factor) * factor);
-                    int w_bar         = static_cast<int>(std::round(width / factor) * factor);
+                    int h_bar         = static_cast<int>(std::round(static_cast<double>(height) / factor) * factor);
+                    int w_bar         = static_cast<int>(std::round(static_cast<double>(width) / factor) * factor);
 
-                    if (static_cast<double>(h_bar) * w_bar > max_pixels) {
-                        double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
-                        h_bar       = std::max(static_cast<int>(factor),
-                                               static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
-                        w_bar       = std::max(static_cast<int>(factor),
-                                               static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
-                    } else if (static_cast<double>(h_bar) * w_bar < min_pixels) {
-                        double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
-                        h_bar       = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
-                        w_bar       = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
-                    }
+                    resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
 
                     LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
 
@@ -2093,16 +2231,33 @@ struct LLMEmbedder : public Conditioner {
                 std::string img_prompt;
                 const std::string placeholder = "<|image_pad|>";
 
+                int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
+                if (min_pixels <= 0) {
+                    min_pixels = 384;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        min_pixels *= min_pixels;
+                    }
+                }
+                int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
+                if (max_pixels <= 0) {
+                    max_pixels = 384;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        max_pixels *= max_pixels;
+                    }
+                }
+
                 for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
                     const auto& image = (*conditioner_params.ref_images)[i];
-                    double factor     = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
+                    const int factor  = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
                     int height        = static_cast<int>(image.shape()[1]);
                     int width         = static_cast<int>(image.shape()[0]);
-                    double beta       = std::sqrt((384.0 * 384.0) / (static_cast<double>(height) * static_cast<double>(width)));
-                    int h_bar         = std::max(static_cast<int>(factor),
-                                                 static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
-                    int w_bar         = std::max(static_cast<int>(factor),
-                                                 static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
+
+                    int h_bar = std::max(factor,
+                                         static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
+                    int w_bar = std::max(factor,
+                                         static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
+
+                    resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
 
                     LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
 
@@ -2141,6 +2296,57 @@ struct LLMEmbedder : public Conditioner {
             out_layers                       = {2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35};
 
             prompt = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n";
+            if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
+                std::string img_prompt        = "";
+                const std::string placeholder = "<|image_pad|>";
+                int min_pixels                = conditioner_params.ref_image_params.vlm_min_size;
+                if (min_pixels <= 0) {
+                    min_pixels = 384;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        min_pixels *= min_pixels;
+                    }
+                }
+                int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
+                if (max_pixels <= 0) {
+                    max_pixels = 1024;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        max_pixels *= max_pixels;
+                    }
+                }
+
+                for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
+                    const auto& image = (*conditioner_params.ref_images)[i];
+                    const int factor  = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
+                    int height        = static_cast<int>(image.shape()[1]);
+                    int width         = static_cast<int>(image.shape()[0]);
+
+                    int h_bar = std::max(factor,
+                                         static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
+                    int w_bar = std::max(factor,
+                                         static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
+
+                    resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
+
+                    LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
+
+                    auto resized_image = clip_preprocess(image, w_bar, h_bar);
+                    auto image_embed   = llm->encode_image(n_threads, resized_image, false, true, true);
+                    GGML_ASSERT(!image_embed.empty());
+
+                    std::string image_prefix = prompt + img_prompt + "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
+                    int image_embed_idx      = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
+                    image_embeds.emplace_back(image_embed_idx, image_embed);
+
+                    img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
+                    int64_t num_image_tokens = image_embed.shape()[1];
+                    img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
+                    for (int j = 0; j < num_image_tokens; j++) {
+                        img_prompt += placeholder;
+                    }
+                    img_prompt += "<|vision_end|>";
+                }
+                prompt += img_prompt;
+            }
 
             prompt_attn_range.first = static_cast<int>(prompt.size());
             prompt += conditioner_params.text;
@@ -2156,30 +2362,33 @@ struct LLMEmbedder : public Conditioner {
                 min_length                       = 512 + prompt_template_encode_start_idx;
                 int image_embed_idx              = 36 + 6;
 
-                int min_pixels          = 384 * 384;
-                int max_pixels          = 560 * 560;
+                int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
+                if (min_pixels <= 0) {
+                    min_pixels = 384;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        min_pixels *= min_pixels;
+                    }
+                }
+                int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
+                if (max_pixels <= 0) {
+                    max_pixels = 560;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        max_pixels *= max_pixels;
+                    }
+                }
+
                 std::string placeholder = "<|image_pad|>";
                 std::string img_prompt;
 
                 for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
                     const auto& image = (*conditioner_params.ref_images)[i];
-                    double factor     = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
+                    const int factor  = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
                     int height        = static_cast<int>(image.shape()[1]);
                     int width         = static_cast<int>(image.shape()[0]);
-                    int h_bar         = static_cast<int>(std::round(height / factor) * factor);
-                    int w_bar         = static_cast<int>(std::round(width / factor) * factor);
+                    int h_bar         = static_cast<int>(std::round(static_cast<double>(height) / factor) * factor);
+                    int w_bar         = static_cast<int>(std::round(static_cast<double>(width) / factor) * factor);
 
-                    if (static_cast<double>(h_bar) * w_bar > max_pixels) {
-                        double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
-                        h_bar       = std::max(static_cast<int>(factor),
-                                               static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
-                        w_bar       = std::max(static_cast<int>(factor),
-                                               static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
-                    } else if (static_cast<double>(h_bar) * w_bar < min_pixels) {
-                        double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
-                        h_bar       = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
-                        w_bar       = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
-                    }
+                    resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
 
                     LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
 

src/model.h

@@ -35,6 +35,7 @@ enum SDVersion {
     VERSION_WAN2,
     VERSION_WAN2_2_I2V,
     VERSION_WAN2_2_TI2V,
+    VERSION_LINGBOT_VIDEO,
     VERSION_QWEN_IMAGE,
     VERSION_QWEN_IMAGE_LAYERED,
     VERSION_ANIMA,
@@ -127,6 +128,13 @@ static inline bool sd_version_is_wan(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_lingbot_video(SDVersion version) {
+    if (version == VERSION_LINGBOT_VIDEO) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_is_qwen_image(SDVersion version) {
     if (version == VERSION_QWEN_IMAGE || version == VERSION_QWEN_IMAGE_LAYERED) {
         return true;
@@ -226,7 +234,7 @@ static inline bool sd_version_uses_flux2_vae(SDVersion version) {
 }
 
 static inline bool sd_version_uses_wan_vae(SDVersion version) {
-    if (sd_version_is_wan(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) {
+    if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) {
         return true;
     }
     return false;
@@ -249,6 +257,7 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_ltxav(version) ||
         sd_version_is_sd3(version) ||
         sd_version_is_wan(version) ||
+        sd_version_is_lingbot_video(version) ||
         sd_version_is_qwen_image(version) ||
         version == VERSION_HIDREAM_O1 ||
         sd_version_is_anima(version) ||

src/model/common/block.hpp

@@ -294,7 +294,7 @@ public:
 
         auto net_0 = std::dynamic_pointer_cast<UnaryBlock>(blocks["net.0"]);
         auto net_2 = std::dynamic_pointer_cast<Linear>(blocks["net.2"]);
-        if (sd_backend_is(ctx->backend, "Vulkan")) {
+        if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
             net_2->set_force_prec_f32(true);
         }
 

src/model/common/rope.hpp

@@ -18,10 +18,6 @@ namespace Rope {
         DECREASE,
     };
 
-    __STATIC_INLINE__ RefIndexMode ref_index_mode_from_bool(bool increase_ref_index) {
-        return increase_ref_index ? RefIndexMode::INCREASE : RefIndexMode::FIXED;
-    }
-
     template <class T>
     __STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
         std::vector<T> result(num);
@@ -759,6 +755,40 @@ namespace Rope {
         return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
     }
 
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_lingbot_video_ids(int t,
+                                                                            int h,
+                                                                            int w,
+                                                                            int pt,
+                                                                            int ph,
+                                                                            int pw,
+                                                                            int bs,
+                                                                            int context_len) {
+        auto vid_ids_repeated = gen_vid_ids(t, h, w, pt, ph, pw, bs, context_len + 1);
+
+        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 + 1);
+            }
+        }
+
+        return concat_ids(vid_ids_repeated, txt_ids, bs);
+    }
+
+    __STATIC_INLINE__ std::vector<float> gen_lingbot_video_pe(int t,
+                                                              int h,
+                                                              int w,
+                                                              int pt,
+                                                              int ph,
+                                                              int pw,
+                                                              int bs,
+                                                              int context_len,
+                                                              int theta,
+                                                              const std::vector<int>& axes_dim) {
+        std::vector<std::vector<float>> ids = gen_lingbot_video_ids(t, h, w, pt, ph, pw, bs, context_len);
+        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
+    }
+
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen2vl_ids(int grid_h,
                                                                       int grid_w,
                                                                       int merge_size,

src/model/diffusion/anima.hpp

@@ -484,10 +484,11 @@ namespace Anima {
                              ggml_tensor* timestep,
                              ggml_tensor* encoder_hidden_states,
                              ggml_tensor* image_pe,
-                             ggml_tensor* t5_ids       = nullptr,
-                             ggml_tensor* t5_weights   = nullptr,
-                             ggml_tensor* adapter_q_pe = nullptr,
-                             ggml_tensor* adapter_k_pe = nullptr) {
+                             ggml_tensor* t5_ids                   = nullptr,
+                             ggml_tensor* t5_weights               = nullptr,
+                             ggml_tensor* adapter_q_pe             = nullptr,
+                             ggml_tensor* adapter_k_pe             = nullptr,
+                             std::vector<ggml_tensor*> ref_latents = {}) {
             GGML_ASSERT(x->ne[3] == 1);
 
             auto x_embedder       = std::dynamic_pointer_cast<XEmbedder>(blocks["x_embedder"]);
@@ -502,8 +503,16 @@ namespace Anima {
             auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, x->ne[0], x->ne[1], 1, x->ne[3]);
             x                 = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2);  // [N, C + 1, H, W]
 
-            x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size);  // [N, h*w, (C+1)*ph*pw]
-
+            x               = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size);  // [N, h*w, (C+1)*ph*pw]
+            int64_t img_len = x->ne[1];
+            if (ref_latents.size() > 0) {
+                for (ggml_tensor* ref : ref_latents) {
+                    auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, ref->ne[0], ref->ne[1], 1, ref->ne[3]);
+                    ref               = ggml_concat(ctx->ggml_ctx, ref, padding_mask, 2);  // [N, C + 1, H, W]
+                    ref               = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
+                    x                 = ggml_concat(ctx->ggml_ctx, x, ref, 1);
+                }
+            }
             x = x_embedder->forward(ctx, x);
 
             auto timestep_proj     = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
@@ -543,6 +552,7 @@ namespace Anima {
                 x          = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe);
                 sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x");
             }
+            x = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, img_len);
 
             x = final_layer->forward(ctx, x, embedded_timestep, temb);  // [N, h*w, ph*pw*C]
 
@@ -602,8 +612,8 @@ namespace Anima {
                                                          const std::vector<int>& axes_dim,
                                                          float h_extrapolation_ratio,
                                                          float w_extrapolation_ratio,
-                                                         float t_extrapolation_ratio) {
-            static const std::vector<ggml_tensor*> empty_ref_latents;
+                                                         float t_extrapolation_ratio,
+                                                         const std::vector<ggml_tensor*>& ref_latents) {
             auto ids = Rope::gen_flux_ids(h,
                                           w,
                                           patch_size,
@@ -611,7 +621,7 @@ namespace Anima {
                                           static_cast<int>(axes_dim.size()),
                                           0,
                                           {},
-                                          empty_ref_latents,
+                                          ref_latents,
                                           Rope::RefIndexMode::FIXED,
                                           1.0f,
                                           false);
@@ -626,14 +636,20 @@ namespace Anima {
 
         ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
                                  const sd::Tensor<float>& timesteps_tensor,
-                                 const sd::Tensor<float>& context_tensor    = {},
-                                 const sd::Tensor<int32_t>& t5_ids_tensor   = {},
-                                 const sd::Tensor<float>& t5_weights_tensor = {}) {
+                                 const sd::Tensor<float>& context_tensor                  = {},
+                                 const sd::Tensor<int32_t>& t5_ids_tensor                 = {},
+                                 const sd::Tensor<float>& t5_weights_tensor               = {},
+                                 const std::vector<sd::Tensor<float>>& ref_latents_tensor = {}) {
             ggml_tensor* x          = make_input(x_tensor);
             ggml_tensor* timesteps  = make_input(timesteps_tensor);
             ggml_tensor* context    = make_optional_input(context_tensor);
             ggml_tensor* t5_ids     = make_optional_input(t5_ids_tensor);
             ggml_tensor* t5_weights = make_optional_input(t5_weights_tensor);
+            std::vector<ggml_tensor*> ref_latents;
+            ref_latents.reserve(ref_latents_tensor.size());
+            for (const auto& ref_latent_tensor : ref_latents_tensor) {
+                ref_latents.push_back(make_input(ref_latent_tensor));
+            }
             GGML_ASSERT(x->ne[3] == 1);
             ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE);
 
@@ -650,7 +666,8 @@ namespace Anima {
                                                            config.axes_dim,
                                                            4.0f,
                                                            4.0f,
-                                                           1.0f);
+                                                           1.0f,
+                                                           ref_latents);
             int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
             auto image_pe         = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len);
             set_backend_tensor_data(image_pe, image_pe_vec.data());
@@ -682,7 +699,8 @@ namespace Anima {
                                           t5_ids,
                                           t5_weights,
                                           adapter_q_pe,
-                                          adapter_k_pe);
+                                          adapter_k_pe,
+                                          ref_latents);
 
             ggml_build_forward_expand(gf, out);
             return gf;
@@ -691,11 +709,13 @@ namespace Anima {
         sd::Tensor<float> compute(int n_threads,
                                   const sd::Tensor<float>& x,
                                   const sd::Tensor<float>& timesteps,
-                                  const sd::Tensor<float>& context    = {},
-                                  const sd::Tensor<int32_t>& t5_ids   = {},
-                                  const sd::Tensor<float>& t5_weights = {}) {
+                                  const sd::Tensor<float>& context                  = {},
+                                  const sd::Tensor<int32_t>& t5_ids                 = {},
+                                  const sd::Tensor<float>& t5_weights               = {},
+                                  const std::vector<sd::Tensor<float>>& ref_latents = {},
+                                  const RefImageParams& ref_image_params            = REF_IMAGE_PRESETS.at("cosmos_reference")) {
             auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(x, timesteps, context, t5_ids, t5_weights);
+                return build_graph(x, timesteps, context, t5_ids, t5_weights, ref_latents);
             };
             return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
@@ -705,12 +725,15 @@ namespace Anima {
             GGML_ASSERT(diffusion_params.x != nullptr);
             GGML_ASSERT(diffusion_params.timesteps != nullptr);
             const auto* extra = diffusion_extra_as<AnimaDiffusionExtra>(diffusion_params);
+            static const std::vector<sd::Tensor<float>> empty_ref_latents;
             return compute(n_threads,
                            *diffusion_params.x,
                            *diffusion_params.timesteps,
                            tensor_or_empty(diffusion_params.context),
                            tensor_or_empty(extra->t5_ids),
-                           tensor_or_empty(extra->t5_weights));
+                           tensor_or_empty(extra->t5_weights),
+                           diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
+                           diffusion_params.ref_image_params);
         }
     };
 }  // namespace Anima

src/model/diffusion/boogu.hpp

@@ -199,7 +199,7 @@ namespace Boogu {
             auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
             auto linear_3 = std::dynamic_pointer_cast<Linear>(blocks["linear_3"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan")) {
+            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
                 linear_2->set_force_prec_f32(true);
             }
 
@@ -259,7 +259,7 @@ namespace Boogu {
             auto norm_k   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
             auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan")) {
+            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
                 to_out_0->set_force_prec_f32(true);
             }
 
@@ -383,7 +383,7 @@ namespace Boogu {
             auto instruct_out  = std::dynamic_pointer_cast<Linear>(blocks["processor.instruct_out"]);
             auto img_out       = std::dynamic_pointer_cast<Linear>(blocks["processor.img_out"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan")) {
+            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
                 to_out_0->set_force_prec_f32(true);
             }
 
@@ -827,7 +827,7 @@ namespace Boogu {
                            *diffusion_params.x,
                            *diffusion_params.timesteps,
                            tensor_or_empty(diffusion_params.context),
-                           diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents);
+                           diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents);
         }
     };
 }  // namespace Boogu

src/model/diffusion/dit.hpp

@@ -104,14 +104,15 @@ namespace DiT {
         return x;
     }
 
-    inline ggml_tensor* patchify(ggml_context* ctx,
-                                 ggml_tensor* x,
-                                 int pt,
-                                 int ph,
-                                 int pw,
-                                 int64_t N = 1) {
+    inline ggml_tensor* patchify_3d(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    int pt,
+                                    int ph,
+                                    int pw,
+                                    int64_t N       = 1,
+                                    bool patch_last = true) {
         // x: [N*C, T, H, W]
-        // return: [N, h*w, C*pt*ph*pw]
+        // return: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C]
         int64_t C     = x->ne[3] / N;
         int64_t T     = x->ne[2];
         int64_t H     = x->ne[1];
@@ -123,15 +124,20 @@ namespace DiT {
         GGML_ASSERT(C * N == x->ne[3]);
         GGML_ASSERT(t_len * pt == T && h_len * ph == H && w_len * pw == W);
 
-        x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N);      // [N*C*t_len, pt, h_len*ph, w_len*pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));          // [N*C*t_len, h_len*ph, pt, w_len*pw]
-        x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N);      // [N*C*t_len*h_len, ph, pt, w_len*pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));          // [N*C*t_len*h_len, pt, ph, w_len*pw]
-        x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N);      // [N*C*t_len*h_len, pt*ph, w_len, pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));          // [N*C*t_len*h_len, w_len, pt*ph, pw]
-        x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N);      // [N, C, t_len*h_len*w_len, pt*ph*pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));          // [N, t_len*h_len*w_len, C, pt*ph*pw]
-        x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1);  // [N, t_len*h_len*w_len, C*pt*ph*pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N);  // [N*C*t_len, pt, h_len*ph, w_len*pw]
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len, h_len*ph, pt, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N);  // [N*C*t_len*h_len, ph, pt, w_len*pw]
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, pt, ph, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N);  // [N*C*t_len*h_len, pt*ph, w_len, pw]
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, w_len, pt*ph, pw]
+        x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N);  // [N, C, t_len*h_len*w_len, pt*ph*pw]
+        if (patch_last) {
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));  // [N, t_len*h_len*w_len, C, pt*ph*pw]
+        } else {
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));  // [N, t_len*h_len*w_len, pt*ph*pw, C]
+        }
+        // [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C]
+        x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1);
         return x;
     }
 
@@ -142,16 +148,23 @@ namespace DiT {
                                       int64_t w_len,
                                       int pt,
                                       int ph,
-                                      int pw) {
-        // x: [N, t_len*h_len*w_len, C*pt*ph*pw]
+                                      int pw,
+                                      bool patch_last = true) {
+        // x: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, pt*ph*pw*C]
         // return: [N*C, t_len*pt, h_len*ph, w_len*pw]
         int64_t N = x->ne[2];
         int64_t C = x->ne[0] / pt / ph / pw;
 
         GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
 
-        x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, C, pt*ph*pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
+        if (patch_last) {
+            x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, C, pt*ph*pw]
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
+        } else {
+            x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, pt*ph*pw, C]
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
+        }
+
         x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);  // [N*C*t_len*h_len, w_len, pt*ph, pw]
         x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, pt*ph, w_len, pw]
         x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);  // [N*C*t_len*h_len, pt, ph, w_len*pw]

src/model/diffusion/flux.hpp

@@ -1642,8 +1642,8 @@ namespace Flux {
                            tensor_or_empty(diffusion_params.c_concat),
                            tensor_or_empty(diffusion_params.y),
                            tensor_or_empty(extra->guidance),
-                           diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.ref_index_mode,
+                           diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
+                           diffusion_params.ref_image_params.ref_index_mode,
                            extra->skip_layers ? *extra->skip_layers : empty_skip_layers,
                            tensor_or_empty(extra->pulid_id),
                            extra->pulid_id_weight);

src/model/diffusion/krea2.hpp

@@ -267,7 +267,7 @@ namespace Krea2 {
             auto knorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["qknorm.knorm"]);
             auto wo    = std::dynamic_pointer_cast<Linear>(blocks["wo"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan")) {
+            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
                 wo->set_force_prec_f32(true);
             }
 
@@ -421,29 +421,88 @@ namespace Krea2 {
         ggml_tensor* forward(GGMLRunnerContext* ctx,
                              ggml_tensor* x,
                              ggml_tensor* vec,
-                             ggml_tensor* pe) {
+                             ggml_tensor* pe,
+                             ggml_tensor* vec_refs = nullptr,
+                             int64_t ref_start     = -1) {
             auto mod      = std::dynamic_pointer_cast<KreaDoubleSharedModulation>(blocks["mod"]);
             auto prenorm  = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["prenorm"]);
             auto postnorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["postnorm"]);
             auto attn     = std::dynamic_pointer_cast<KreaAttention>(blocks["attn"]);
             auto mlp      = std::dynamic_pointer_cast<KreaSwiGLU>(blocks["mlp"]);
 
-            auto mods       = mod->forward(ctx, vec);
-            auto attn_input = Flux::modulate(ctx->ggml_ctx,
-                                             prenorm->forward(ctx, x),
-                                             mods[1],
-                                             mods[0],
-                                             true);
-            auto attn_out   = attn->forward(ctx, attn_input, pe);
-            x               = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2]));
+            if (ref_start >= 0 && vec_refs) {
+                // same as normal, but since vec is different for refs and the rest, needs a lot of views and concats
+                auto mods_main = mod->forward(ctx, vec);
+                auto mods_refs = mod->forward(ctx, vec_refs);
 
-            auto mlp_input = Flux::modulate(ctx->ggml_ctx,
-                                            postnorm->forward(ctx, x),
-                                            mods[4],
-                                            mods[3],
-                                            true);
-            auto mlp_out   = mlp->forward(ctx, mlp_input);
-            x              = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mods[5]));
+                int64_t D  = x->ne[0];
+                int64_t N  = x->ne[1];
+                int64_t B  = x->ne[2];
+                size_t nb1 = x->nb[1];
+                size_t nb2 = x->nb[2];
+
+                int64_t len_main = ref_start;
+                int64_t len_refs = N - ref_start;
+
+                auto pre_x = prenorm->forward(ctx, x);
+
+                auto pre_x_main = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_main, B, nb1, nb2, 0);
+                auto pre_x_refs = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_refs, B, nb1, nb2, len_main * nb1);
+
+                auto attn_in_main = Flux::modulate(ctx->ggml_ctx, pre_x_main, mods_main[1], mods_main[0], true);
+                auto attn_in_refs = Flux::modulate(ctx->ggml_ctx, pre_x_refs, mods_refs[1], mods_refs[0], true);
+
+                auto attn_input = ggml_concat(ctx->ggml_ctx, attn_in_main, attn_in_refs, 1);
+
+                auto attn_out = attn->forward(ctx, attn_input, pe);
+
+                auto attn_out_main = ggml_view_3d(ctx->ggml_ctx, attn_out, D, len_main, B, attn_out->nb[1], attn_out->nb[2], 0);
+                auto attn_out_refs = ggml_view_3d(ctx->ggml_ctx, attn_out, D, len_refs, B, attn_out->nb[1], attn_out->nb[2], len_main * attn_out->nb[1]);
+
+                auto res_main = ggml_mul(ctx->ggml_ctx, attn_out_main, mods_main[2]);
+                auto res_refs = ggml_mul(ctx->ggml_ctx, attn_out_refs, mods_refs[2]);
+
+                auto attn_res = ggml_concat(ctx->ggml_ctx, res_main, res_refs, 1);
+
+                x = ggml_add(ctx->ggml_ctx, x, attn_res);
+
+                auto post_x = postnorm->forward(ctx, x);
+
+                auto post_x_main = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_main, B, post_x->nb[1], post_x->nb[2], 0);
+                auto post_x_refs = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_refs, B, post_x->nb[1], post_x->nb[2], len_main * post_x->nb[1]);
+
+                auto mlp_in_main = Flux::modulate(ctx->ggml_ctx, post_x_main, mods_main[4], mods_main[3], true);
+                auto mlp_in_refs = Flux::modulate(ctx->ggml_ctx, post_x_refs, mods_refs[4], mods_refs[3], true);
+
+                auto mlp_input = ggml_concat(ctx->ggml_ctx, mlp_in_main, mlp_in_refs, 1);
+                auto mlp_out   = mlp->forward(ctx, mlp_input);
+
+                auto mlp_out_main = ggml_view_3d(ctx->ggml_ctx, mlp_out, D, len_main, B, mlp_out->nb[1], mlp_out->nb[2], 0);
+                auto mlp_out_refs = ggml_view_3d(ctx->ggml_ctx, mlp_out, D, len_refs, B, mlp_out->nb[1], mlp_out->nb[2], len_main * mlp_out->nb[1]);
+
+                auto mlp_res_main = ggml_mul(ctx->ggml_ctx, mlp_out_main, mods_main[5]);
+                auto mlp_res_refs = ggml_mul(ctx->ggml_ctx, mlp_out_refs, mods_refs[5]);
+
+                auto mlp_res = ggml_concat(ctx->ggml_ctx, mlp_res_main, mlp_res_refs, 1);
+                x            = ggml_add(ctx->ggml_ctx, x, mlp_res);
+            } else {
+                auto mods       = mod->forward(ctx, vec);
+                auto attn_input = Flux::modulate(ctx->ggml_ctx,
+                                                 prenorm->forward(ctx, x),
+                                                 mods[1],
+                                                 mods[0],
+                                                 true);
+                auto attn_out   = attn->forward(ctx, attn_input, pe);
+                x               = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2]));
+
+                auto mlp_input = Flux::modulate(ctx->ggml_ctx,
+                                                postnorm->forward(ctx, x),
+                                                mods[4],
+                                                mods[3],
+                                                true);
+                auto mlp_out   = mlp->forward(ctx, mlp_input);
+                x              = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mods[5]));
+            }
             return x;
         }
 
@@ -555,7 +614,9 @@ namespace Krea2 {
                              ggml_tensor* x,
                              ggml_tensor* timestep,
                              ggml_tensor* context,
-                             ggml_tensor* pe) {
+                             ggml_tensor* pe,
+                             std::vector<ggml_tensor*> ref_latents = {},
+                             bool zero_timestep_refs               = false) {
             int64_t W = x->ne[0];
             int64_t H = x->ne[1];
             int64_t N = x->ne[3];
@@ -570,6 +631,13 @@ namespace Krea2 {
 
             auto img        = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size, true);
             int64_t img_len = img->ne[1];
+            if (ref_latents.size() > 0) {
+                for (ggml_tensor* ref : ref_latents) {
+                    ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size, true);
+                    img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
+                }
+            }
+            int64_t ref_len = img->ne[1] - img_len;
             img             = first->forward(ctx, img);
 
             auto t    = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.timestep_dim), 10000, 1000.f);
@@ -577,19 +645,30 @@ namespace Krea2 {
             t         = ggml_reshape_3d(ctx->ggml_ctx, t, t->ne[0], 1, t->ne[1]);
             auto tvec = tproj->forward(ctx, t);
 
+            ggml_tensor* tvec_0 = nullptr;
+            if (ref_latents.size() > 0 && zero_timestep_refs) {
+                // "index_timestep_zero" mode: use timestep = 0 for ref latents
+                auto timestep_0 = ggml_scale(ctx->ggml_ctx, timestep, 0.0f);
+                auto t_0        = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_0, static_cast<int>(config.timestep_dim), 10000, 1000.f);
+                t_0             = tmlp->forward(ctx, t_0);
+                t_0             = ggml_reshape_3d(ctx->ggml_ctx, t_0, t_0->ne[0], 1, t_0->ne[1]);
+                tvec_0          = tproj->forward(ctx, t_0);
+            }
+
             auto txt        = txtfusion->forward(ctx, context);
             txt             = txtmlp->forward(ctx, txt);
             int64_t txt_len = txt->ne[1];
 
             auto hidden_states = ggml_concat(ctx->ggml_ctx, txt, img, 1);
+            int64_t ref_start  = hidden_states->ne[1] - ref_len;
             for (int i = 0; i < config.layers; ++i) {
                 auto block    = std::dynamic_pointer_cast<KreaSingleStreamBlock>(blocks["blocks." + std::to_string(i)]);
-                hidden_states = block->forward(ctx, hidden_states, tvec, pe);
+                hidden_states = block->forward(ctx, hidden_states, tvec, pe, tvec_0, ref_start);
                 sd::ggml_graph_cut::mark_graph_cut(hidden_states, "krea2.blocks." + std::to_string(i), "hidden_states");
             }
 
-            hidden_states = last->forward(ctx, hidden_states, t);
             hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, txt_len, txt_len + img_len);
+            hidden_states = last->forward(ctx, hidden_states, t);
             hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, true);
             return hidden_states;
         }
@@ -601,10 +680,16 @@ namespace Krea2 {
                                                       int bs,
                                                       int context_len,
                                                       float theta,
-                                                      const std::vector<int>& axes_dim) {
+                                                      const std::vector<int>& axes_dim,
+                                                      const std::vector<ggml_tensor*>& ref_latents,
+                                                      Rope::RefIndexMode ref_index_mode) {
         auto txt_ids = Rope::gen_flux_txt_ids(bs, context_len, 3, {});
         auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false);
         auto ids     = Rope::concat_ids(txt_ids, img_ids, bs);
+        if (ref_latents.size() > 0) {
+            auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0);
+            ids           = Rope::concat_ids(ids, refs_ids, bs);
+        }
         return Rope::embed_nd(ids, bs, theta, axes_dim);
     }
 
@@ -633,7 +718,9 @@ namespace Krea2 {
 
         ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
                                  const sd::Tensor<float>& timesteps_tensor,
-                                 const sd::Tensor<float>& context_tensor) {
+                                 const sd::Tensor<float>& context_tensor,
+                                 const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
+                                 const RefImageParams& ref_image_params                   = REF_IMAGE_PRESETS.at("krea2_ostris_edit")) {
             ggml_cgraph* gf        = new_graph_custom(KREA2_GRAPH_SIZE);
             ggml_tensor* x         = make_input(x_tensor);
             ggml_tensor* timesteps = make_input(timesteps_tensor);
@@ -641,19 +728,27 @@ namespace Krea2 {
             GGML_ASSERT(!context_tensor.empty());
             ggml_tensor* context = make_input(context_tensor);
 
+            std::vector<ggml_tensor*> ref_latents;
+            ref_latents.reserve(ref_latents_tensor.size());
+            for (const auto& ref_latent_tensor : ref_latents_tensor) {
+                ref_latents.push_back(make_input(ref_latent_tensor));
+            }
+
             pe_vec      = gen_krea2_pe(static_cast<int>(x->ne[1]),
                                        static_cast<int>(x->ne[0]),
                                        config.patch_size,
                                        static_cast<int>(x->ne[3]),
                                        static_cast<int>(context->ne[1]),
                                        config.theta,
-                                       config.axes_dim);
+                                       config.axes_dim,
+                                       ref_latents,
+                                       ref_image_params.ref_index_mode);
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
             set_backend_tensor_data(pe, pe_vec.data());
 
             auto runner_ctx  = get_context();
-            ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe);
+            ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents, ref_image_params.force_ref_timestep_zero);
             ggml_build_forward_expand(gf, out);
             return gf;
         }
@@ -661,9 +756,11 @@ namespace Krea2 {
         sd::Tensor<float> compute(int n_threads,
                                   const sd::Tensor<float>& x,
                                   const sd::Tensor<float>& timesteps,
-                                  const sd::Tensor<float>& context) {
+                                  const sd::Tensor<float>& context,
+                                  const std::vector<sd::Tensor<float>>& ref_latents = {},
+                                  const RefImageParams& ref_image_params            = REF_IMAGE_PRESETS.at("krea2_ostris_edit")) {
             auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(x, timesteps, context);
+                return build_graph(x, timesteps, context, ref_latents, ref_image_params);
             };
             return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
@@ -672,10 +769,13 @@ namespace Krea2 {
                                   const DiffusionParams& diffusion_params) override {
             GGML_ASSERT(diffusion_params.x != nullptr);
             GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            static const std::vector<sd::Tensor<float>> empty_ref_latents;
             return compute(n_threads,
                            *diffusion_params.x,
                            *diffusion_params.timesteps,
-                           tensor_or_empty(diffusion_params.context));
+                           tensor_or_empty(diffusion_params.context),
+                           diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
+                           diffusion_params.ref_image_params);
         }
     };
 }  // namespace Krea2

src/model/diffusion/lingbot_video.hpp

@@ -0,0 +1,692 @@
+#ifndef __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__
+#define __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__
+
+#include <algorithm>
+#include <cinttypes>
+#include <cmath>
+#include <cstdlib>
+#include <memory>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "core/util.h"
+#include "model/common/rope.hpp"
+#include "model/diffusion/dit.hpp"
+#include "model/diffusion/flux.hpp"
+#include "model/diffusion/model.hpp"
+#include "model/diffusion/qwen_image.hpp"
+
+namespace LingBotVideo {
+    constexpr int LINGBOT_VIDEO_GRAPH_SIZE = 65536;
+
+    struct LingBotVideoConfig {
+        int patch_t                   = 1;
+        int patch_h                   = 2;
+        int patch_w                   = 2;
+        int64_t in_channels           = 16;
+        int64_t out_channels          = 16;
+        int64_t hidden_size           = 2048;
+        int64_t num_attention_heads   = 16;
+        int64_t depth                 = 24;
+        int64_t intermediate_size     = 6144;
+        int64_t text_dim              = 2560;
+        int64_t freq_dim              = 256;
+        float norm_eps                = 1e-6f;
+        int rope_theta                = 256;
+        std::vector<int> axes_dim     = {32, 48, 48};
+        int axes_dim_sum              = 128;
+        bool qkv_bias                 = false;
+        bool out_bias                 = true;
+        bool patch_embed_bias         = true;
+        bool timestep_mlp_bias        = true;
+        int64_t num_experts           = 0;
+        int64_t num_experts_per_tok   = 8;
+        int64_t moe_intermediate_size = 512;
+        int64_t decoder_sparse_step   = 1;
+        int64_t n_shared_experts      = 0;
+        bool norm_topk_prob           = true;
+        float routed_scaling_factor   = 1.0f;
+        int64_t n_group               = 0;
+        int64_t topk_group            = 0;
+        std::set<int> sparse_layers;
+
+        static LingBotVideoConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                                      const std::string& prefix) {
+            LingBotVideoConfig config;
+            config.depth = 0;
+
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+
+                if (ends_with(name, "patch_embedder.weight") && tensor_storage.n_dims == 2) {
+                    int64_t patch_dim    = tensor_storage.ne[0];
+                    config.hidden_size   = tensor_storage.ne[1];
+                    int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w;
+                    if (patch_dim % patch_volume == 0) {
+                        config.in_channels = patch_dim / patch_volume;
+                    }
+                } else if (ends_with(name, "text_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
+                    config.text_dim = tensor_storage.ne[0];
+                } else if (ends_with(name, "time_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
+                    config.freq_dim = tensor_storage.ne[0];
+                } else if (ends_with(name, "blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
+                    int64_t head_dim = tensor_storage.ne[0];
+                    if (head_dim > 0) {
+                        config.num_attention_heads = config.hidden_size / head_dim;
+                    }
+                } else if (name.find(".attn.to_q.bias") != std::string::npos) {
+                    config.qkv_bias = true;
+                } else if (name.find(".ffn.gate_proj.weight") != std::string::npos && tensor_storage.n_dims == 2) {
+                    config.intermediate_size = tensor_storage.ne[1];
+                } else if (name.find(".ffn.experts.w1") != std::string::npos && tensor_storage.n_dims == 3) {
+                    config.num_experts           = tensor_storage.ne[2];
+                    config.moe_intermediate_size = tensor_storage.ne[1];
+                } else if (name.find(".ffn.shared_experts.gate_proj.weight") != std::string::npos && tensor_storage.n_dims == 2) {
+                    if (config.moe_intermediate_size > 0) {
+                        config.n_shared_experts = tensor_storage.ne[1] / config.moe_intermediate_size;
+                    }
+                } else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) {
+                    int64_t out_dim      = tensor_storage.ne[1];
+                    int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w;
+                    config.out_channels  = patch_volume > 0 ? out_dim / patch_volume : config.out_channels;
+                }
+
+                size_t block_pos = name.find("blocks.");
+                if (block_pos != std::string::npos) {
+                    auto items = split_string(name.substr(block_pos), '.');
+                    if (items.size() > 1) {
+                        int block_index = atoi(items[1].c_str());
+                        if (block_index + 1 > config.depth) {
+                            config.depth = block_index + 1;
+                        }
+                        if (name.find("blocks." + std::to_string(block_index) + ".ffn.experts.w1") != std::string::npos) {
+                            config.sparse_layers.insert(block_index);
+                        }
+                    }
+                }
+            }
+
+            if (config.depth == 0) {
+                config.depth = 24;
+            }
+            config.axes_dim_sum = 0;
+            for (int axis_dim : config.axes_dim) {
+                config.axes_dim_sum += axis_dim;
+            }
+            if (!config.sparse_layers.empty()) {
+                config.num_experts           = 128;
+                config.num_experts_per_tok   = 8;
+                config.moe_intermediate_size = 768;
+                config.decoder_sparse_step   = 1;
+                config.n_shared_experts      = 1;
+                config.norm_topk_prob        = true;
+                config.n_group               = 4;
+                config.topk_group            = 2;
+                config.routed_scaling_factor = 2.5f;
+            }
+            LOG_DEBUG("lingbot_video: depth = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64 ", text_dim = %" PRId64 ", experts = %" PRId64 ", experts_per_tok = %" PRId64 ", n_group = %" PRId64 ", topk_group = %" PRId64 ", route_scale = %.2f, sparse_layers = %zu",
+                      config.depth,
+                      config.hidden_size,
+                      config.num_attention_heads,
+                      config.text_dim,
+                      config.num_experts,
+                      config.num_experts_per_tok,
+                      config.n_group,
+                      config.topk_group,
+                      config.routed_scaling_factor,
+                      config.sparse_layers.size());
+            return config;
+        }
+    };
+
+    struct LingBotVideoTextEmbedder : public GGMLBlock {
+        LingBotVideoTextEmbedder(int64_t text_dim,
+                                 int64_t hidden_size,
+                                 float eps = 1e-6f) {
+            blocks["norm"]     = std::make_shared<RMSNorm>(text_dim, eps);
+            blocks["linear_1"] = std::make_shared<Linear>(text_dim, hidden_size, true);
+            blocks["linear_2"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto norm     = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
+            auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
+            auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
+
+            x = norm->forward(ctx, x);
+            x = linear_1->forward(ctx, x);
+            x = ggml_silu_inplace(ctx->ggml_ctx, x);
+            x = linear_2->forward(ctx, x);
+            return x;
+        }
+    };
+
+    struct LingBotVideoAttention : public GGMLBlock {
+        int64_t num_heads;
+        int64_t head_dim;
+
+        LingBotVideoAttention(int64_t hidden_size,
+                              int64_t num_heads,
+                              bool qkv_bias = false,
+                              bool out_bias = true,
+                              float eps     = 1e-6f)
+            : num_heads(num_heads),
+              head_dim(hidden_size / num_heads) {
+            int64_t inner_dim = num_heads * head_dim;
+            blocks["to_q"]    = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
+            blocks["to_k"]    = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
+            blocks["to_v"]    = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
+            blocks["norm_q"]  = std::make_shared<RMSNorm>(head_dim, eps);
+            blocks["norm_k"]  = std::make_shared<RMSNorm>(head_dim, eps);
+            blocks["to_out"]  = std::make_shared<Linear>(inner_dim, hidden_size, out_bias);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* pe,
+                             ggml_tensor* attention_mask = nullptr) {
+            // x: [N, video_tokens + text_tokens, hidden_size]
+            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 = std::dynamic_pointer_cast<Linear>(blocks["to_out"]);
+
+            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);
+            k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_heads, S, N);
+            v = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim, num_heads, S, N);
+
+            q = norm_q->forward(ctx, q);
+            k = norm_k->forward(ctx, k);
+
+            x = Rope::attention(ctx, q, k, v, pe, attention_mask);
+            x = to_out->forward(ctx, x);
+            return x;
+        }
+    };
+
+    struct LingBotVideoMLP : public UnaryBlock {
+        LingBotVideoMLP(int64_t hidden_size,
+                        int64_t intermediate_size) {
+            blocks["gate_proj"] = std::make_shared<Linear>(hidden_size, intermediate_size, false);
+            blocks["up_proj"]   = std::make_shared<Linear>(hidden_size, intermediate_size, false);
+            blocks["down_proj"] = std::make_shared<Linear>(intermediate_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 down_proj = std::dynamic_pointer_cast<Linear>(blocks["down_proj"]);
+
+            auto gate = gate_proj->forward(ctx, x);
+            gate      = ggml_silu_inplace(ctx->ggml_ctx, gate);
+            auto up   = up_proj->forward(ctx, x);
+            x         = ggml_mul(ctx->ggml_ctx, gate, up);
+            x         = down_proj->forward(ctx, x);
+            return x;
+        }
+    };
+
+    struct LingBotVideoSparseMoeBlock : public UnaryBlock {
+        int64_t hidden_size;
+        int64_t intermediate_size;
+        int64_t num_experts;
+        int64_t num_experts_per_tok;
+        bool has_shared_experts;
+        bool norm_topk_prob;
+        float routed_scaling_factor;
+        int64_t n_group;
+        int64_t topk_group;
+        std::vector<float> group_expert_mask_vec;
+        bool has_correction_bias = false;
+
+        LingBotVideoSparseMoeBlock(const LingBotVideoConfig& config)
+            : hidden_size(config.hidden_size),
+              intermediate_size(config.moe_intermediate_size),
+              num_experts(config.num_experts),
+              num_experts_per_tok(config.num_experts_per_tok),
+              has_shared_experts(config.n_shared_experts > 0),
+              norm_topk_prob(config.norm_topk_prob),
+              routed_scaling_factor(config.routed_scaling_factor),
+              n_group(config.n_group),
+              topk_group(config.topk_group) {
+            if (n_group > 1) {
+                GGML_ASSERT(num_experts % n_group == 0);
+                int64_t experts_per_group = num_experts / n_group;
+                group_expert_mask_vec.assign(static_cast<size_t>(num_experts * n_group), 0.f);
+                for (int64_t group = 0; group < n_group; ++group) {
+                    int64_t expert_begin = group * experts_per_group;
+                    int64_t expert_end   = expert_begin + experts_per_group;
+                    for (int64_t expert = expert_begin; expert < expert_end; ++expert) {
+                        group_expert_mask_vec[static_cast<size_t>(group * num_experts + expert)] = 1.f;
+                    }
+                }
+            }
+            if (has_shared_experts) {
+                blocks["shared_experts"] = std::make_shared<LingBotVideoMLP>(hidden_size,
+                                                                             intermediate_size * config.n_shared_experts);
+            }
+        }
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            ggml_type router_type   = get_type(prefix + "router.weight", tensor_storage_map, GGML_TYPE_F32);
+            ggml_type w1_type       = get_type(prefix + "experts.w1", tensor_storage_map, GGML_TYPE_F32);
+            ggml_type w2_type       = get_type(prefix + "experts.w2", tensor_storage_map, GGML_TYPE_F32);
+            ggml_type w3_type       = get_type(prefix + "experts.w3", tensor_storage_map, GGML_TYPE_F32);
+            params["router.weight"] = ggml_new_tensor_2d(ctx, router_type, hidden_size, num_experts);
+            if (tensor_storage_map.find(prefix + "router.e_score_correction_bias") != tensor_storage_map.end()) {
+                params["router.e_score_correction_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts);
+                has_correction_bias                      = true;
+            }
+            params["experts.w1"] = ggml_new_tensor_3d(ctx, w1_type, hidden_size, intermediate_size, num_experts);
+            params["experts.w2"] = ggml_new_tensor_3d(ctx, w2_type, intermediate_size, hidden_size, num_experts);
+            params["experts.w3"] = ggml_new_tensor_3d(ctx, w3_type, hidden_size, intermediate_size, num_experts);
+        }
+
+        ggml_tensor* expert_linear(GGMLRunnerContext* ctx,
+                                   const std::string& weight_name,
+                                   ggml_tensor* x,
+                                   ggml_tensor* selected_experts) {
+            return ggml_mul_mat_id(ctx->ggml_ctx, params[weight_name], x, selected_experts);
+        }
+
+        ggml_tensor* select_experts(GGMLRunnerContext* ctx, ggml_tensor* choice_scores) {
+            ggml_context* gctx = ctx->ggml_ctx;
+            if (n_group <= 1 || topk_group <= 0) {
+                return ggml_argsort_top_k(gctx, choice_scores, static_cast<int>(num_experts_per_tok));
+            }
+
+            GGML_ASSERT(choice_scores->ne[0] == num_experts);
+            GGML_ASSERT(num_experts % n_group == 0);
+            GGML_ASSERT(topk_group > 0 && topk_group <= n_group);
+            GGML_ASSERT(!group_expert_mask_vec.empty());
+
+            const int64_t n_token_total     = choice_scores->ne[1];
+            const int64_t experts_per_group = num_experts / n_group;
+            const int group_score_k         = 2;
+            GGML_ASSERT(experts_per_group >= group_score_k);
+
+            ggml_tensor* grouped_scores = ggml_reshape_3d(gctx, choice_scores, experts_per_group, n_group, n_token_total);
+            ggml_tensor* group_top_ids  = ggml_argsort_top_k(gctx, grouped_scores, group_score_k);
+            grouped_scores              = ggml_reshape_3d(gctx, grouped_scores, 1, experts_per_group, n_group * n_token_total);
+            group_top_ids               = ggml_cont(gctx, group_top_ids);
+            group_top_ids               = ggml_reshape_2d(gctx, group_top_ids, group_score_k, n_group * n_token_total);
+
+            ggml_tensor* group_top_values = ggml_get_rows(gctx, grouped_scores, group_top_ids);
+            group_top_values              = ggml_reshape_3d(gctx, group_top_values, group_score_k, n_group, n_token_total);
+
+            ggml_tensor* group_scores = nullptr;
+            for (int rank = 0; rank < group_score_k; ++rank) {
+                ggml_tensor* value = ggml_view_3d(gctx,
+                                                  group_top_values,
+                                                  1,
+                                                  n_group,
+                                                  n_token_total,
+                                                  group_top_values->nb[1],
+                                                  group_top_values->nb[2],
+                                                  rank * group_top_values->nb[0]);
+                group_scores       = group_scores == nullptr ? value : ggml_add(gctx, group_scores, value);
+            }
+            group_scores = ggml_reshape_2d(gctx, group_scores, n_group, n_token_total);
+
+            ggml_tensor* selected_groups = ggml_argsort_top_k(gctx, group_scores, static_cast<int>(topk_group));
+            selected_groups              = ggml_cont(gctx, selected_groups);
+
+            ggml_tensor* group_expert_mask = ggml_new_tensor_3d(gctx, GGML_TYPE_F32, num_experts, n_group, 1);
+            ctx->bind_backend_tensor_data(group_expert_mask, group_expert_mask_vec.data());
+            ggml_tensor* group_expert_mask_template = ggml_new_tensor_3d(gctx, GGML_TYPE_F32, num_experts, n_group, n_token_total);
+            group_expert_mask                       = ggml_repeat(gctx, group_expert_mask, group_expert_mask_template);
+
+            ggml_tensor* selected_group_masks = ggml_get_rows(gctx, group_expert_mask, selected_groups);
+            ggml_tensor* selected_mask        = nullptr;
+            for (int64_t rank = 0; rank < topk_group; ++rank) {
+                ggml_tensor* mask = ggml_view_3d(gctx,
+                                                 selected_group_masks,
+                                                 num_experts,
+                                                 1,
+                                                 n_token_total,
+                                                 selected_group_masks->nb[1],
+                                                 selected_group_masks->nb[2],
+                                                 rank * selected_group_masks->nb[1]);
+                selected_mask     = selected_mask == nullptr ? mask : ggml_add(gctx, selected_mask, mask);
+            }
+            selected_mask = ggml_reshape_2d(gctx, selected_mask, num_experts, n_token_total);
+
+            ggml_tensor* excluded_group_mask = ggml_sub(gctx, selected_mask, ggml_ext_ones_like(gctx, selected_mask));
+            ggml_tensor* masked_scores       = ggml_add(gctx, choice_scores, ggml_scale(gctx, excluded_group_mask, 1.0e9f));
+            return ggml_argsort_top_k(gctx, masked_scores, static_cast<int>(num_experts_per_tok));
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            // x: [N, tokens, hidden_size]
+            GGML_ASSERT(num_experts > 0);
+            GGML_ASSERT(num_experts_per_tok > 0 && num_experts_per_tok <= num_experts);
+
+            ggml_tensor* residual       = x;
+            const int64_t n_token       = x->ne[1];
+            const int64_t N             = x->ne[2];
+            const int64_t n_token_total = n_token * N;
+
+            ggml_tensor* router_logits = ggml_mul_mat(ctx->ggml_ctx, params["router.weight"], x);
+            router_logits              = ggml_reshape_2d(ctx->ggml_ctx, router_logits, num_experts, n_token_total);
+            ggml_tensor* probs         = ggml_sigmoid(ctx->ggml_ctx, router_logits);
+            ggml_tensor* choice_scores = probs;
+            if (has_correction_bias) {
+                choice_scores = ggml_add(ctx->ggml_ctx, choice_scores, params["router.e_score_correction_bias"]);
+            }
+
+            ggml_tensor* selected_experts = select_experts(ctx, choice_scores);
+            ggml_tensor* weights          = ggml_get_rows(ctx->ggml_ctx,
+                                                          ggml_reshape_3d(ctx->ggml_ctx, probs, 1, num_experts, n_token_total),
+                                                          selected_experts);
+            weights                       = ggml_reshape_2d(ctx->ggml_ctx, weights, num_experts_per_tok, n_token_total);
+            if (norm_topk_prob && num_experts_per_tok > 1) {
+                auto weights_sum = ggml_sum_rows(ctx->ggml_ctx, weights);
+                weights_sum      = ggml_clamp(ctx->ggml_ctx, weights_sum, 6.103515625e-5f, INFINITY);
+                weights          = ggml_div(ctx->ggml_ctx, weights, weights_sum);
+            }
+            if (routed_scaling_factor != 1.0f) {
+                weights = ggml_scale(ctx->ggml_ctx, weights, routed_scaling_factor);
+            }
+            weights = ggml_reshape_3d(ctx->ggml_ctx, weights, 1, num_experts_per_tok, n_token_total);
+
+            x         = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size, 1, n_token_total);
+            auto gate = expert_linear(ctx, "experts.w1", x, selected_experts);
+            gate      = ggml_silu_inplace(ctx->ggml_ctx, gate);
+            auto up   = expert_linear(ctx, "experts.w3", x, selected_experts);
+            auto act  = ggml_mul(ctx->ggml_ctx, gate, up);
+            auto out  = expert_linear(ctx, "experts.w2", act, selected_experts);
+            out       = ggml_mul(ctx->ggml_ctx, out, weights);
+
+            ggml_tensor* summed = nullptr;
+            for (int64_t i = 0; i < num_experts_per_tok; ++i) {
+                auto expert_out = ggml_view_2d(ctx->ggml_ctx,
+                                               out,
+                                               hidden_size,
+                                               n_token_total,
+                                               out->nb[2],
+                                               i * out->nb[1]);
+                summed          = summed == nullptr ? expert_out : ggml_add(ctx->ggml_ctx, summed, expert_out);
+            }
+            if (num_experts_per_tok == 1) {
+                summed = ggml_cont(ctx->ggml_ctx, summed);
+            }
+            summed = ggml_reshape_3d(ctx->ggml_ctx, summed, hidden_size, n_token, N);
+
+            if (has_shared_experts) {
+                auto shared_experts = std::dynamic_pointer_cast<LingBotVideoMLP>(blocks["shared_experts"]);
+                summed              = ggml_add(ctx->ggml_ctx, summed, shared_experts->forward(ctx, residual));
+            }
+            return summed;
+        }
+    };
+
+    struct LingBotVideoBlock : public GGMLBlock {
+        int64_t hidden_size;
+
+        LingBotVideoBlock(const LingBotVideoConfig& config,
+                          bool sparse)
+            : hidden_size(config.hidden_size) {
+            blocks["norm1"]          = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
+            blocks["attn"]           = std::make_shared<LingBotVideoAttention>(config.hidden_size,
+                                                                     config.num_attention_heads,
+                                                                     config.qkv_bias,
+                                                                     config.out_bias,
+                                                                     config.norm_eps);
+            blocks["norm_post_attn"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
+            blocks["norm2"]          = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
+            if (sparse) {
+                blocks["ffn"] = std::make_shared<LingBotVideoSparseMoeBlock>(config);
+            } else {
+                blocks["ffn"] = std::make_shared<LingBotVideoMLP>(config.hidden_size, config.intermediate_size);
+            }
+            blocks["norm_post_ffn"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
+        }
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            params["scale_shift_table"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size * 6, 1);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* temb6,
+                             ggml_tensor* pe,
+                             ggml_tensor* attention_mask = nullptr) {
+            // x: [N, tokens, hidden_size], temb6: [N, tokens, 6 * hidden_size]
+            auto norm1          = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
+            auto attn           = std::dynamic_pointer_cast<LingBotVideoAttention>(blocks["attn"]);
+            auto norm_post_attn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_post_attn"]);
+            auto norm2          = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
+            auto ffn            = std::dynamic_pointer_cast<UnaryBlock>(blocks["ffn"]);
+            auto norm_post_ffn  = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_post_ffn"]);
+
+            auto table = ggml_reshape_3d(ctx->ggml_ctx, params["scale_shift_table"], hidden_size * 6, 1, 1);
+            auto mod   = ggml_add(ctx->ggml_ctx, temb6, table);
+            auto mods  = ggml_ext_chunk(ctx->ggml_ctx, mod, 6, 0);
+
+            auto shift_msa = mods[0];
+            auto scale_msa = mods[1];
+            auto gate_msa  = ggml_tanh(ctx->ggml_ctx, mods[2]);
+            auto shift_mlp = mods[3];
+            auto scale_mlp = mods[4];
+            auto gate_mlp  = ggml_tanh(ctx->ggml_ctx, mods[5]);
+
+            auto attn_in  = Flux::modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa, true);
+            auto attn_out = attn->forward(ctx, attn_in, pe, attention_mask);
+            attn_out      = norm_post_attn->forward(ctx, attn_out);
+            x             = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_msa, attn_out));
+
+            auto ffn_in  = Flux::modulate(ctx->ggml_ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp, true);
+            auto ffn_out = ffn->forward(ctx, ffn_in);
+            ffn_out      = norm_post_ffn->forward(ctx, ffn_out);
+            x            = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_mlp, ffn_out));
+            return x;
+        }
+    };
+
+    class LingBotVideoModel : public GGMLBlock {
+    public:
+        LingBotVideoConfig config;
+
+        LingBotVideoModel() = default;
+        LingBotVideoModel(LingBotVideoConfig config)
+            : config(config) {
+            int64_t patch_dim           = config.patch_t * config.patch_h * config.patch_w * config.in_channels;
+            blocks["patch_embedder"]    = std::make_shared<Linear>(patch_dim, config.hidden_size, config.patch_embed_bias);
+            blocks["time_embedder"]     = std::make_shared<Qwen::TimestepEmbedding>(config.freq_dim,
+                                                                                config.hidden_size,
+                                                                                config.hidden_size,
+                                                                                0,
+                                                                                config.timestep_mlp_bias);
+            blocks["time_modulation.1"] = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size, true);
+            blocks["text_embedder"]     = std::make_shared<LingBotVideoTextEmbedder>(config.text_dim,
+                                                                                 config.hidden_size,
+                                                                                 config.norm_eps);
+            for (int i = 0; i < config.depth; i++) {
+                bool sparse                           = config.sparse_layers.find(i) != config.sparse_layers.end();
+                blocks["blocks." + std::to_string(i)] = std::make_shared<LingBotVideoBlock>(config, sparse);
+            }
+            blocks["norm_out"]              = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
+            blocks["norm_out_modulation.1"] = std::make_shared<Linear>(config.hidden_size, 2 * config.hidden_size, true);
+            blocks["proj_out"]              = std::make_shared<Linear>(config.hidden_size,
+                                                          config.patch_t * config.patch_h * config.patch_w * config.out_channels,
+                                                          true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timestep,
+                             ggml_tensor* context,
+                             ggml_tensor* pe) {
+            // x: [N*C, T, H, W], context: [N, text_tokens, text_dim]
+            GGML_ASSERT(context != nullptr);
+            GGML_ASSERT(x->ne[3] == config.in_channels);
+            GGML_ASSERT(x->ne[2] % config.patch_t == 0);
+            GGML_ASSERT(x->ne[1] % config.patch_h == 0);
+            GGML_ASSERT(x->ne[0] % config.patch_w == 0);
+
+            auto patch_embedder      = std::dynamic_pointer_cast<Linear>(blocks["patch_embedder"]);
+            auto time_embedder       = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["time_embedder"]);
+            auto time_modulation     = std::dynamic_pointer_cast<Linear>(blocks["time_modulation.1"]);
+            auto text_embedder       = std::dynamic_pointer_cast<LingBotVideoTextEmbedder>(blocks["text_embedder"]);
+            auto norm_out            = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
+            auto norm_out_modulation = std::dynamic_pointer_cast<Linear>(blocks["norm_out_modulation.1"]);
+            auto proj_out            = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
+
+            int64_t W     = x->ne[0];
+            int64_t H     = x->ne[1];
+            int64_t T     = x->ne[2];
+            int64_t N     = 1;
+            int64_t t_len = T / config.patch_t;
+            int64_t h_len = H / config.patch_h;
+            int64_t w_len = W / config.patch_w;
+            int64_t n_img = t_len * h_len * w_len;
+
+            auto img = DiT::patchify_3d(ctx->ggml_ctx, x, config.patch_t, config.patch_h, config.patch_w, N, false);
+            img      = patch_embedder->forward(ctx, img);
+
+            auto txt           = text_embedder->forward(ctx, context);
+            auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1);
+
+            auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx,
+                                                             timestep,
+                                                             static_cast<int>(config.freq_dim),
+                                                             10000,
+                                                             1.0f);
+            auto t_emb         = time_embedder->forward(ctx, timestep_proj);
+            auto temb_template = ggml_new_tensor_3d(ctx->ggml_ctx, t_emb->type, t_emb->ne[0], hidden_states->ne[1], t_emb->ne[1]);
+            auto temb_tokens   = ggml_repeat(ctx->ggml_ctx,
+                                             ggml_reshape_3d(ctx->ggml_ctx, t_emb, t_emb->ne[0], 1, t_emb->ne[1]),
+                                             temb_template);
+            auto temb6         = time_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens));
+
+            sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.prelude", "hidden_states");
+
+            for (int i = 0; i < config.depth; i++) {
+                auto block    = std::dynamic_pointer_cast<LingBotVideoBlock>(blocks["blocks." + std::to_string(i)]);
+                hidden_states = block->forward(ctx, hidden_states, temb6, pe);
+                sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.blocks." + std::to_string(i), "hidden_states");
+            }
+
+            auto final_mods = ggml_ext_chunk(ctx->ggml_ctx,
+                                             norm_out_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens)),
+                                             2,
+                                             0);
+            hidden_states   = norm_out->forward(ctx, hidden_states);
+            hidden_states   = Flux::modulate(ctx->ggml_ctx, hidden_states, final_mods[0], final_mods[1], true);
+            hidden_states   = proj_out->forward(ctx, hidden_states);
+            hidden_states   = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img);
+
+            auto out = DiT::unpatchify_3d(ctx->ggml_ctx,
+                                          hidden_states,
+                                          t_len,
+                                          h_len,
+                                          w_len,
+                                          config.patch_t,
+                                          config.patch_h,
+                                          config.patch_w,
+                                          false);
+            return out;
+        }
+    };
+
+    struct LingBotVideoRunner : public DiffusionModelRunner {
+        LingBotVideoConfig config;
+        LingBotVideoModel lingbot_video;
+        std::vector<float> pe_vec;
+
+        LingBotVideoRunner(ggml_backend_t backend,
+                           const String2TensorStorage& tensor_storage_map      = {},
+                           const std::string prefix                            = "",
+                           std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                           const char* model_args                              = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(LingBotVideoConfig::detect_from_weights(tensor_storage_map, prefix)) {
+            SD_UNUSED(model_args);
+
+            lingbot_video = LingBotVideoModel(config);
+            lingbot_video.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "lingbot_video";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            lingbot_video.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(LINGBOT_VIDEO_GRAPH_SIZE);
+            ggml_tensor* x         = make_input(x_tensor);
+            ggml_tensor* timesteps = make_input(timesteps_tensor);
+            GGML_ASSERT(x_tensor.dim() == 5);
+            GGML_ASSERT(x->ne[3] == config.in_channels);
+            GGML_ASSERT(!context_tensor.empty());
+            ggml_tensor* context = make_input(context_tensor);
+
+            int64_t W = x->ne[0];
+            int64_t H = x->ne[1];
+            int64_t T = x->ne[2];
+            int64_t N = x_tensor.shape()[4];
+            GGML_ASSERT(N == 1);
+            pe_vec      = Rope::gen_lingbot_video_pe(static_cast<int>(T),
+                                                     static_cast<int>(H),
+                                                     static_cast<int>(W),
+                                                     config.patch_t,
+                                                     config.patch_h,
+                                                     config.patch_w,
+                                                     static_cast<int>(N),
+                                                     static_cast<int>(context->ne[1]),
+                                                     config.rope_theta,
+                                                     config.axes_dim);
+            int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+            auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
+            set_backend_tensor_data(pe, pe_vec.data());
+
+            auto runner_ctx  = get_context();
+            ggml_tensor* out = lingbot_video.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, false, false), x.dim());
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context));
+        }
+    };
+}  // namespace LingBotVideo
+
+#endif  // __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__

src/model/diffusion/model.hpp

@@ -10,6 +10,36 @@
 #include "model/common/rope.hpp"
 #include "model_manager.h"
 
+enum class RefImageResizeMode {
+    NONE,
+    LONGEST_SIDE,
+    AREA,
+};
+
+struct RefImageParams {
+    bool pass_to_vlm                   = false;
+    bool pass_to_dit                   = true;
+    Rope::RefIndexMode ref_index_mode  = Rope::RefIndexMode::FIXED;
+    bool force_ref_timestep_zero       = false;
+    bool resize_before_vae             = true;
+    int vae_input_max_pixels           = -1;
+    RefImageResizeMode vlm_resize_mode = RefImageResizeMode::AREA;
+    int vlm_min_size                   = -1;
+    int vlm_max_size                   = -1;
+};
+
+const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
+    {"flux_kontext", {false, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::NONE, -1, -1}},
+    {"longcat", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
+    {"flux2", {false, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::NONE, -1, -1}},
+    {"qwen", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
+    {"qwen_layered", {true, true, Rope::RefIndexMode::DECREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
+    {"z_image_omni", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
+    {"krea2_ostris_edit", {true, true, Rope::RefIndexMode::INCREASE, true, true, -1, RefImageResizeMode::AREA, -1, -1}},
+    {"krea2_edit", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, 768, 768}},
+    {"cosmos_reference", {false, true, Rope::RefIndexMode::INCREASE, false, false, -1, RefImageResizeMode::NONE, -1, -1}},
+};
+
 struct UNetDiffusionExtra {
     int num_video_frames                           = -1;
     const std::vector<sd::Tensor<float>>* controls = nullptr;
@@ -74,7 +104,7 @@ struct DiffusionParams {
     const sd::Tensor<float>* c_concat                 = nullptr;
     const sd::Tensor<float>* y                        = nullptr;
     const std::vector<sd::Tensor<float>>* ref_latents = nullptr;
-    Rope::RefIndexMode ref_index_mode                 = Rope::RefIndexMode::FIXED;
+    RefImageParams ref_image_params                   = {false, false, Rope::RefIndexMode::FIXED, false};
     DiffusionExtraParams extra                        = std::monostate{};
 };
 

src/model/diffusion/qwen_image.hpp

@@ -183,7 +183,7 @@ namespace Qwen {
             auto to_v     = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
             auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan")) {
+            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
                 to_out_0->set_force_prec_f32(true);
             }
 
@@ -517,7 +517,7 @@ namespace Qwen {
                 if (input->ne[3] == 1) {
                     input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]);
                 }
-                return DiT::patchify(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
+                return DiT::patchify_3d(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
             };
 
             auto img           = patchify_input(x);
@@ -715,8 +715,8 @@ namespace Qwen {
                            *diffusion_params.x,
                            *diffusion_params.timesteps,
                            tensor_or_empty(diffusion_params.context),
-                           diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.ref_index_mode);
+                           diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
+                           diffusion_params.ref_image_params.ref_index_mode);
         }
 
         void test() {

src/model/diffusion/z_image.hpp

@@ -150,6 +150,8 @@ namespace ZImage {
 
             if (sd_backend_is(ctx->backend, "ROCm")) {
                 out_proj->set_scale(1.f / 16.f);
+                out_proj->set_force_prec_f32(true);
+                qkv_proj->set_force_prec_f32(true);
             }
 
             auto qkv = qkv_proj->forward(ctx, x);                                                                            // [N, n_token, (num_heads + num_kv_heads*2)*head_dim]
@@ -227,7 +229,7 @@ namespace ZImage {
             auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
             auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan")) {
+            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
                 w2->set_force_prec_f32(true);
             }
 
@@ -646,8 +648,8 @@ namespace ZImage {
                            *diffusion_params.x,
                            *diffusion_params.timesteps,
                            tensor_or_empty(diffusion_params.context),
-                           diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.ref_index_mode);
+                           diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
+                           diffusion_params.ref_image_params.ref_index_mode);
         }
 
         void test() {

src/model_io/safetensors_io.cpp

@@ -3,15 +3,19 @@
 #include <algorithm>
 #include <cstdint>
 #include <exception>
+#include <filesystem>
 #include <fstream>
 #include <ostream>
 #include <string>
+#include <unordered_set>
 #include <vector>
 
 #include "binary_io.h"
 #include "core/util.h"
 #include "json.hpp"
 
+namespace fs = std::filesystem;
+
 static constexpr size_t ST_HEADER_SIZE_LEN = 8;
 
 static void set_error(std::string* error, const std::string& message) {
@@ -20,6 +24,14 @@ static void set_error(std::string* error, const std::string& message) {
     }
 }
 
+static std::string resolve_index_shard_path(const std::string& index_path, const std::string& shard_path) {
+    fs::path shard_fs_path(shard_path);
+    if (shard_fs_path.is_absolute()) {
+        return shard_fs_path.lexically_normal().string();
+    }
+    return (fs::path(index_path).parent_path() / shard_fs_path).lexically_normal().string();
+}
+
 bool is_safetensors_file(const std::string& file_path) {
     std::ifstream file(file_path, std::ios::binary);
     if (!file.is_open()) {
@@ -230,6 +242,52 @@ bool read_safetensors_file(const std::string& file_path,
     return true;
 }
 
+bool read_safetensors_index_file(const std::string& file_path,
+                                 std::vector<std::string>& shard_paths,
+                                 std::string* error) {
+    shard_paths.clear();
+
+    std::ifstream file(file_path);
+    if (!file.is_open()) {
+        set_error(error, "failed to open '" + file_path + "'");
+        return false;
+    }
+
+    nlohmann::json index;
+    try {
+        index = nlohmann::json::parse(file);
+    } catch (const std::exception&) {
+        set_error(error, "parsing safetensors index failed: '" + file_path + "'");
+        return false;
+    }
+
+    if (!index.is_object() || !index.contains("weight_map") || !index["weight_map"].is_object()) {
+        set_error(error, "invalid safetensors index '" + file_path + "'");
+        return false;
+    }
+
+    std::unordered_set<std::string> seen_shard_paths;
+    for (const auto& item : index["weight_map"].items()) {
+        if (!item.value().is_string()) {
+            set_error(error, "invalid shard path for tensor '" + item.key() + "'");
+            return false;
+        }
+
+        std::string shard_path = resolve_index_shard_path(file_path,
+                                                          item.value().get<std::string>());
+        if (seen_shard_paths.insert(shard_path).second) {
+            shard_paths.push_back(std::move(shard_path));
+        }
+    }
+
+    if (shard_paths.empty()) {
+        set_error(error, "safetensors index has no tensors: '" + file_path + "'");
+        return false;
+    }
+
+    return true;
+}
+
 static bool ggml_type_to_safetensors_dtype(ggml_type type, std::string* dtype) {
     switch (type) {
         case GGML_TYPE_F16:

src/model_io/safetensors_io.h

@@ -11,6 +11,9 @@ 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 read_safetensors_index_file(const std::string& file_path,
+                                 std::vector<std::string>& shard_paths,
+                                 std::string* error = nullptr);
 bool write_safetensors_file(const std::string& file_path,
                             const std::vector<TensorWriteInfo>& tensors,
                             std::string* error = nullptr);

src/model_loader.cpp

@@ -235,6 +235,9 @@ bool ModelLoader::init_from_file(const std::string& file_path, const std::string
     } else if (is_gguf_file(file_path)) {
         LOG_INFO("load %s using gguf format", file_path.c_str());
         return init_from_gguf_file(file_path, prefix);
+    } else if (ends_with(file_path, ".json")) {
+        LOG_INFO("load %s using safetensors index format", file_path.c_str());
+        return init_from_safetensors_index_file(file_path, prefix);
     } else if (is_safetensors_file(file_path)) {
         LOG_INFO("load %s using safetensors format", file_path.c_str());
         return init_from_safetensors_file(file_path, prefix);
@@ -339,6 +342,25 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
     return true;
 }
 
+bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path, const std::string& prefix) {
+    LOG_DEBUG("init from safetensors index '%s', prefix = '%s'", file_path.c_str(), prefix.c_str());
+
+    std::vector<std::string> shard_paths;
+    std::string error;
+    if (!read_safetensors_index_file(file_path, shard_paths, &error)) {
+        LOG_ERROR("%s", error.c_str());
+        return false;
+    }
+
+    for (const std::string& shard_path : shard_paths) {
+        if (!init_from_file(shard_path, prefix)) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
 /*================================================= TorchLegacyModelLoader ==================================================*/
 
 bool ModelLoader::init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix) {
@@ -511,6 +533,9 @@ SDVersion ModelLoader::get_sd_version() {
         if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) {
             is_wan = true;
         }
+        if (tensor_storage.name.find("model.diffusion_model.patch_embedder.weight") != std::string::npos) {
+            return VERSION_LINGBOT_VIDEO;
+        }
         if (tensor_storage.name.find("model.diffusion_model.patch_embedding.weight") != std::string::npos) {
             patch_embedding_channels = tensor_storage.ne[3];
         }

src/model_loader.h

@@ -43,6 +43,7 @@ protected:
 
     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_safetensors_index_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 = "");

src/model_manager.cpp

@@ -179,6 +179,102 @@ bool ModelManager::register_param_tensors(const std::string& desc,
     return true;
 }
 
+bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* registered_tensor_size) {
+    if (desc.empty()) {
+        return true;
+    }
+
+    std::unordered_set<TensorState*> target_states;
+    size_t released_size = 0;
+    for (auto& state : tensor_states_) {
+        if (state == nullptr || state->desc != desc) {
+            continue;
+        }
+        if (state->active_prepare_count > 0) {
+            LOG_ERROR("model manager cannot unregister active %s tensor '%s'",
+                      desc.c_str(),
+                      state->name.c_str());
+            return false;
+        }
+        target_states.insert(state.get());
+        if (state->tensor != nullptr) {
+            released_size += ggml_nbytes(state->tensor);
+        }
+    }
+
+    if (target_states.empty()) {
+        return true;
+    }
+
+    release_compute_staging_blocks(false);
+
+    std::vector<ParamsStorageBlock*> storage_blocks_to_release;
+    std::unordered_set<TensorState*> affected_storage_states;
+    for (const auto& block : params_storage_blocks_) {
+        if (block == nullptr) {
+            continue;
+        }
+        bool has_target_state = false;
+        for (TensorState* state : block->states) {
+            if (state != nullptr && target_states.count(state) > 0) {
+                has_target_state = true;
+                break;
+            }
+        }
+        if (!has_target_state) {
+            continue;
+        }
+        storage_blocks_to_release.push_back(block.get());
+        for (TensorState* state : block->states) {
+            if (state != nullptr) {
+                affected_storage_states.insert(state);
+            }
+        }
+    }
+
+    for (TensorState* state : affected_storage_states) {
+        if (state == nullptr) {
+            continue;
+        }
+        if (state->active_prepare_count > 0 || state->staged_to_compute_backend) {
+            LOG_ERROR("model manager cannot unregister %s while tensor '%s' is active",
+                      desc.c_str(),
+                      state->name.c_str());
+            return false;
+        }
+    }
+
+    for (ParamsStorageBlock* block : storage_blocks_to_release) {
+        if (block != nullptr) {
+            free_params_storage_block(*block);
+            erase_params_storage_block(block);
+        }
+    }
+
+    for (auto it = tensor_states_by_name_.begin(); it != tensor_states_by_name_.end();) {
+        if (target_states.count(it->second) > 0) {
+            it = tensor_states_by_name_.erase(it);
+        } else {
+            ++it;
+        }
+    }
+    tensor_states_.erase(std::remove_if(tensor_states_.begin(),
+                                        tensor_states_.end(),
+                                        [&](const std::unique_ptr<TensorState>& s) {
+                                            return s == nullptr || target_states.count(s.get()) > 0;
+                                        }),
+                         tensor_states_.end());
+
+    if (registered_tensor_size != nullptr) {
+        if (released_size > *registered_tensor_size) {
+            *registered_tensor_size = 0;
+        } else {
+            *registered_tensor_size -= released_size;
+        }
+    }
+    return true;
+}
+
 bool ModelManager::load_all_params_eagerly() {
     std::vector<TensorState*> all_states;
     all_states.reserve(tensor_states_.size());

src/model_manager.h

@@ -134,6 +134,9 @@ public:
                                 bool allow_split_buffer            = false,
                                 bool params_follow_compute_backend = false);
 
+    bool unregister_param_tensors(const std::string& desc,
+                                  size_t* registered_tensor_size = nullptr);
+
     template <typename Runner>
     bool register_runner_params(const std::string& desc,
                                 Runner& runner,

src/name_conversion.cpp

@@ -1102,6 +1102,8 @@ std::string convert_sep_to_dot(std::string name) {
         "norm1_context",
         "ff_context",
         "x_embedder",
+        "cross_attn",
+        "output_proj",
     };
 
     // record the positions of underscores that should NOT be replaced

src/stable-diffusion.cpp

@@ -32,6 +32,7 @@
 #include "model/diffusion/ideogram4.hpp"
 #include "model/diffusion/krea2.hpp"
 #include "model/diffusion/lens.hpp"
+#include "model/diffusion/lingbot_video.hpp"
 #include "model/diffusion/ltxv.hpp"
 #include "model/diffusion/minit2i.hpp"
 #include "model/diffusion/mmdit.hpp"
@@ -87,6 +88,7 @@ const char* model_version_to_str[] = {
     "Wan 2.x",
     "Wan 2.2 I2V",
     "Wan 2.2 TI2V",
+    "LingBot Video",
     "Qwen Image",
     "Qwen Image Layered",
     "Anima",
@@ -222,9 +224,13 @@ public:
     std::string split_mode_spec;
     bool auto_fit_enabled = false;
 
+    bool diffusion_conv_direct = false;
+
     bool is_using_v_parameterization     = false;
     bool is_using_edm_v_parameterization = false;
 
+    size_t control_net_params_mem_size = 0;
+
     std::shared_ptr<ModelManager> model_manager;
 
     std::shared_ptr<Denoiser> denoiser = std::make_shared<CompVisDenoiser>();
@@ -494,6 +500,76 @@ public:
                                                      params_follow_runtime);
     }
 
+    bool unload_control_net() {
+        if (control_net == nullptr) {
+            return true;
+        }
+        if (model_manager != nullptr) {
+            if (!model_manager->unregister_param_tensors("ControlNet", &control_net_params_mem_size)) {
+                return false;
+            }
+        }
+        control_net.reset();
+        control_net_params_mem_size = 0;
+        return true;
+    }
+
+    bool load_control_net_from_file(const std::string& path) {
+        if (path.empty()) {
+            LOG_ERROR("sd_ctx_load_control_net: empty path");
+            return false;
+        }
+        if (model_manager == nullptr) {
+            LOG_ERROR("sd_ctx_load_control_net: model_manager not initialized");
+            return false;
+        }
+
+        if (!unload_control_net()) {
+            return false;
+        }
+
+        ModelLoader& shared_loader = model_manager->loader();
+        if (!shared_loader.init_from_file(path)) {
+            LOG_ERROR("sd_ctx_load_control_net: failed to load '%s'", path.c_str());
+            return false;
+        }
+        shared_loader.convert_tensors_name();
+
+        if (!ensure_backend_pair(SDBackendModule::CONTROL_NET)) {
+            LOG_ERROR("sd_ctx_load_control_net: control_net backend unavailable");
+            return false;
+        }
+
+        control_net = std::make_shared<ControlNet>(backend_for(SDBackendModule::CONTROL_NET),
+                                                   params_backend_for(SDBackendModule::CONTROL_NET),
+                                                   shared_loader.get_tensor_storage_map(),
+                                                   version,
+                                                   "",
+                                                   model_manager);
+        if (diffusion_conv_direct) {
+            LOG_INFO("Using Conv2d direct in the control net");
+            control_net->set_conv2d_direct_enabled(true);
+        }
+        if (!register_runner_params("ControlNet",
+                                    control_net,
+                                    SDBackendModule::CONTROL_NET,
+                                    &control_net_params_mem_size)) {
+            LOG_ERROR("sd_ctx_load_control_net: register_runner_params failed");
+            control_net.reset();
+            control_net_params_mem_size = 0;
+            return false;
+        }
+        if (!model_manager->validate_registered_tensors()) {
+            LOG_ERROR("sd_ctx_load_control_net: registered tensors validation failed");
+            unload_control_net();
+            return false;
+        }
+        LOG_INFO("sd_ctx_load_control_net: loaded '%s' (%.2f MB)",
+                 path.c_str(),
+                 control_net_params_mem_size / 1024.0 / 1024.0);
+        return true;
+    }
+
     bool init_backend() {
         std::string error;
         if (!backend_manager.init(backend_spec.c_str(),
@@ -852,10 +928,12 @@ public:
         model_loader.process_model_files(enable_mmap, needs_writable_mmap);
         load_alphas_cumprod(model_loader);
 
+        diffusion_conv_direct = sd_ctx_params->diffusion_conv_direct;
+
         size_t text_encoder_params_mem_size = 0;
         size_t unet_params_mem_size         = 0;
         size_t vae_params_mem_size          = 0;
-        size_t control_net_params_mem_size  = 0;
+        control_net_params_mem_size         = 0;
         size_t extension_params_mem_size    = 0;
 
         bool tae_preview_only = sd_ctx_params->tae_preview_only;
@@ -905,7 +983,7 @@ public:
                                                                  tensor_storage_map,
                                                                  version,
                                                                  "",
-                                                                 false,
+                                                                 true,
                                                                  model_manager);
                 diffusion_model  = std::make_shared<Krea2::Krea2Runner>(backend_for(SDBackendModule::DIFFUSION),
                                                                        tensor_storage_map,
@@ -1004,6 +1082,25 @@ public:
                         return false;
                     }
                 }
+            } else if (sd_version_is_lingbot_video(version)) {
+                bool enable_vision = false;
+                for (const auto& [name, _] : tensor_storage_map) {
+                    if (starts_with(name, "text_encoders.llm.visual.")) {
+                        enable_vision = true;
+                        break;
+                    }
+                }
+                cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
+                                                                 tensor_storage_map,
+                                                                 version,
+                                                                 "",
+                                                                 enable_vision,
+                                                                 model_manager);
+                diffusion_model  = std::make_shared<LingBotVideo::LingBotVideoRunner>(backend_for(SDBackendModule::DIFFUSION),
+                                                                                     tensor_storage_map,
+                                                                                     "model.diffusion_model",
+                                                                                     model_manager,
+                                                                                     sd_ctx_params->model_args);
             } else if (sd_version_is_qwen_image(version)) {
                 bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED;
                 cond_stage_model   = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
@@ -1503,6 +1600,7 @@ public:
                     }
                 } else if (sd_version_is_sd3(version) ||
                            sd_version_is_wan(version) ||
+                           sd_version_is_lingbot_video(version) ||
                            sd_version_is_qwen_image(version) ||
                            version == VERSION_HIDREAM_O1 ||
                            sd_version_is_anima(version) ||
@@ -2199,7 +2297,7 @@ public:
                              const char* extra_sample_args,
                              const std::vector<float>& sigmas,
                              const std::vector<sd::Tensor<float>>& ref_latents,
-                             bool increase_ref_index,
+                             const RefImageParams& ref_image_params,
                              const sd::Tensor<float>& denoise_mask,
                              const sd::Tensor<float>& vace_context,
                              float vace_strength,
@@ -2332,7 +2430,7 @@ public:
             sd::Tensor<float> timesteps_tensor({static_cast<int64_t>(timesteps_vec.size())}, timesteps_vec);
             sd::Tensor<float> guidance_tensor({1}, std::vector<float>{guidance.distilled_guidance});
             sd::Tensor<float> noised_input = x * c_in;
-            if (!denoise_mask.empty() && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version))) {
+            if (!denoise_mask.empty() && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version) || sd_version_is_lingbot_video(version))) {
                 noised_input = noised_input * denoise_mask + init_latent * (1.0f - denoise_mask);
             }
 
@@ -2360,9 +2458,9 @@ public:
             sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma);
             std::vector<sd::Tensor<float>> controls;
             DiffusionParams diffusion_params;
-            diffusion_params.x              = &noised_input;
-            diffusion_params.timesteps      = &timesteps_tensor;
-            diffusion_params.ref_index_mode = Rope::ref_index_mode_from_bool(increase_ref_index);
+            diffusion_params.x                = &noised_input;
+            diffusion_params.timesteps        = &timesteps_tensor;
+            diffusion_params.ref_image_params = ref_image_params;
             sd::guidance::GuidanceInput step_guidance_input;
             step_guidance_input.step          = step;
             step_guidance_input.schedule_size = sigmas.size();
@@ -2578,7 +2676,7 @@ public:
     int get_diffusion_model_down_factor() {
         int down_factor = 8;  // unet
         if (sd_version_is_dit(version)) {
-            if (sd_version_is_wan(version)) {
+            if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) {
                 down_factor = 2;
             } else {
                 down_factor = 1;
@@ -2641,7 +2739,7 @@ public:
         int latent_frames = frames;
         if (sd_version_is_ltxav(version)) {
             latent_frames = ((frames - 1) / 8) + 1;
-        } else if (sd_version_is_wan(version)) {
+        } else if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) {
             latent_frames = ((frames - 1) / 4) + 1;
         }
         return latent_frames;
@@ -2654,7 +2752,7 @@ public:
         if (sd_version_is_ltxav(version)) {
             return (latent_frames - 1) * 8 + 1;
         }
-        if (sd_version_is_wan(version)) {
+        if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) {
             return (latent_frames - 1) * 4 + 1;
         }
         return latent_frames;
@@ -2742,6 +2840,126 @@ public:
         auto flow_denoiser = std::dynamic_pointer_cast<DiscreteFlowDenoiser>(denoiser);
         return !!flow_denoiser;
     }
+
+    std::string get_default_ref_image_preset(SDVersion version) const {
+        if (sd_version_is_longcat(version)) {
+            return "longcat";
+        } else if (sd_version_is_flux(version)) {
+            return "flux_kontext";
+        } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
+            return "flux2";
+        } else if (version == VERSION_QWEN_IMAGE_LAYERED) {
+            return "qwen_layered";
+        } else if (sd_version_is_qwen_image(version)) {
+            return "qwen";
+        } else if (sd_version_is_z_image(version) || sd_version_is_boogu_image(version)) {
+            return "z_image_omni";
+        } else if (sd_version_is_krea2(version)) {
+            // have to make a choice between "krea2_edit" mode (for lbouaraba/krea2edit)
+            // and "krea2_ostris_edit" (for krea2 ostris edit)
+            // since krea2 ostris edit support predates, it should probably be default
+            return "krea2_ostris_edit";
+        } else if (sd_version_is_anima(version)) {
+            return "cosmos_reference";
+        }
+        return "default";
+    }
+
+    RefImageParams resolve_ref_image_params(const char* ref_image_args) const {
+        RefImageParams params;
+        std::string preset_name = get_default_ref_image_preset(version);
+
+        for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) {
+            if (key == "preset") {
+                std::string requested_preset_name = value;
+                if (REF_IMAGE_PRESETS.count(requested_preset_name)) {
+                    preset_name = requested_preset_name;
+                } else if (value != "default") {
+                    std::string valid_list;
+                    for (auto const& [name, _] : REF_IMAGE_PRESETS) {
+                        valid_list += (valid_list.empty() ? "" : ", ") + name;
+                    }
+                    LOG_WARN("ignoring invalid reference image preset '%s'. Valid options: [%s]", value.c_str(), valid_list.c_str());
+                }
+                break;
+            }
+        }
+        if (preset_name != "default") {
+            LOG_INFO("Using '%s' preset for reference images", preset_name.c_str());
+            params = REF_IMAGE_PRESETS.at(preset_name);
+        }
+
+        for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) {
+            if (key == "pass_to_vlm") {
+                if (!parse_strict_bool(value, params.pass_to_vlm)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "pass_to_dit") {
+                if (!parse_strict_bool(value, params.pass_to_dit)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "ref_index_mode") {
+                if (value == "fixed") {
+                    params.ref_index_mode = Rope::RefIndexMode::FIXED;
+                } else if (value == "increase") {
+                    params.ref_index_mode = Rope::RefIndexMode::INCREASE;
+                } else if (value == "decrease") {
+                    params.ref_index_mode = Rope::RefIndexMode::DECREASE;
+                } else {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "force_ref_timestep_zero") {
+                if (!parse_strict_bool(value, params.force_ref_timestep_zero)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "resize_before_vae") {
+                if (!parse_strict_bool(value, params.resize_before_vae)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "vae_input_max_pixels") {
+                if (!parse_strict_int(value, params.vae_input_max_pixels)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "vlm_resize_mode") {
+                if (value == "longest_side") {
+                    params.vlm_resize_mode = RefImageResizeMode::LONGEST_SIDE;
+                } else if (value == "area") {
+                    params.vlm_resize_mode = RefImageResizeMode::AREA;
+                } else if (value == "none") {
+                    params.vlm_resize_mode = RefImageResizeMode::NONE;
+                } else {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "vlm_max_size") {
+                if (!parse_strict_int(value, params.vlm_max_size)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "vlm_min_size") {
+                if (!parse_strict_int(value, params.vlm_min_size)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key != "preset" && key != "vlm_size") {
+                LOG_WARN("ignoring unknown reference image arg '%s'", key.c_str());
+            }
+        }
+        for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) {
+            if (key == "vlm_size") {
+                int vlm_size;
+                if (!parse_strict_int(value, vlm_size)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                } else {
+                    LOG_INFO("vlm_size override: setting both min and max size to %ld", (long)vlm_size);
+                    params.vlm_min_size = vlm_size;
+                    params.vlm_max_size = vlm_size;
+                }
+                break;
+            }
+        }
+        if (params.force_ref_timestep_zero && !sd_version_is_krea2(version)) {
+            LOG_WARN("force_ref_timestep_zero is only supported by Krea2 architecture for now");
+        }
+        return params;
+    }
 };
 
 /*================================================= SD API ==================================================*/
@@ -3203,6 +3421,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
     sd_sample_params_init(&sd_img_gen_params->sample_params);
     sd_img_gen_params->clip_skip         = -1;
     sd_img_gen_params->ref_images_count  = 0;
+    sd_img_gen_params->ref_image_args    = "";
     sd_img_gen_params->width             = 512;
     sd_img_gen_params->height            = 512;
     sd_img_gen_params->strength          = 0.75f;
@@ -3240,8 +3459,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              "batch_count: %d\n"
              "qwen_image_layers: %d\n"
              "ref_images_count: %d\n"
-             "auto_resize_ref_image: %s\n"
-             "increase_ref_index: %s\n"
+             "ref_image_args: %s\n"
              "control_strength: %.2f\n"
              "photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\n"
              "VAE tiling: %s (temporal=%s, extra_tiling_args=%s)\n"
@@ -3259,8 +3477,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              sd_img_gen_params->batch_count,
              sd_img_gen_params->qwen_image_layers,
              sd_img_gen_params->ref_images_count,
-             BOOL_STR(sd_img_gen_params->auto_resize_ref_image),
-             BOOL_STR(sd_img_gen_params->increase_ref_index),
+             SAFE_STR(sd_img_gen_params->ref_image_args),
              sd_img_gen_params->control_strength,
              sd_img_gen_params->pm_params.style_strength,
              sd_img_gen_params->pm_params.id_images_count,
@@ -3328,7 +3545,7 @@ struct sd_ctx_t {
 };
 
 static bool sd_version_supports_video_generation(SDVersion version) {
-    return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_ltxav(version);
+    return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version);
 }
 
 static bool sd_version_supports_image_generation(SDVersion version) {
@@ -3429,6 +3646,27 @@ SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) {
     return sd_version_supports_video_generation(sd_ctx->sd->version);
 }
 
+SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path) {
+    if (sd_ctx == nullptr || sd_ctx->sd == nullptr || path == nullptr) {
+        return false;
+    }
+    return sd_ctx->sd->load_control_net_from_file(path);
+}
+
+SD_API bool sd_ctx_unload_control_net(sd_ctx_t* sd_ctx) {
+    if (sd_ctx == nullptr || sd_ctx->sd == nullptr) {
+        return false;
+    }
+    return sd_ctx->sd->unload_control_net();
+}
+
+SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx) {
+    if (sd_ctx == nullptr || sd_ctx->sd == nullptr) {
+        return false;
+    }
+    return sd_ctx->sd->control_net != nullptr;
+}
+
 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_pid(sd_ctx->sd->version)) {
@@ -3536,8 +3774,6 @@ struct GenerationRequest {
     float strength                           = 1.f;
     float control_strength                   = 0.f;
     float eta                                = 0.f;
-    bool increase_ref_index                  = false;
-    bool auto_resize_ref_image               = false;
     sd_guidance_params_t guidance            = {};
     sd_guidance_params_t high_noise_guidance = {};
     sd_pm_params_t pm_params                 = {};
@@ -3563,8 +3799,6 @@ struct GenerationRequest {
         strength                    = sd_img_gen_params->strength;
         control_strength            = sd_img_gen_params->control_strength;
         eta                         = sd_img_gen_params->sample_params.eta;
-        increase_ref_index          = sd_img_gen_params->increase_ref_index;
-        auto_resize_ref_image       = sd_img_gen_params->auto_resize_ref_image;
         has_ref_images              = sd_img_gen_params->ref_images_count > 0;
         guidance                    = sd_img_gen_params->sample_params.guidance;
         pm_params                   = sd_img_gen_params->pm_params;
@@ -4312,7 +4546,8 @@ static sd::Tensor<float> ensure_image_tensor_channels(sd::Tensor<float> image, i
 static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd_ctx_t* sd_ctx,
                                                                               const sd_img_gen_params_t* sd_img_gen_params,
                                                                               GenerationRequest* request,
-                                                                              SamplePlan* plan) {
+                                                                              SamplePlan* plan,
+                                                                              const RefImageParams& ref_image_params) {
     int64_t prepare_start_ms = ggml_time_ms();
 
     sd::Tensor<float> init_image_tensor;
@@ -4455,9 +4690,10 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
             continue;
         }
         sd::Tensor<float> ref_latent;
-        if (request->auto_resize_ref_image && !sd_version_is_pid(sd_ctx->sd->version)) {
+        if (ref_image_params.resize_before_vae && !sd_version_is_pid(sd_ctx->sd->version)) {
             LOG_DEBUG("auto resize ref images");
-            int vae_image_size = std::min(1024 * 1024, request->width * request->height);
+            int target_pixels  = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : 1024 * 1024;
+            int vae_image_size = std::min(target_pixels, request->width * request->height);
             double vae_width   = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]);
             double vae_height  = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0];
 
@@ -4584,15 +4820,20 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
                                                                             const sd_img_gen_params_t* sd_img_gen_params,
                                                                             GenerationRequest* request,
                                                                             SamplePlan* plan,
-                                                                            ImageGenerationLatents* latents) {
+                                                                            ImageGenerationLatents* latents,
+                                                                            const RefImageParams& ref_image_params) {
     ConditionerRunnerDoneOnExit conditioner_runner_done{sd_ctx->sd->cond_stage_model.get()};
 
     ConditionerParams condition_params;
-    condition_params.text       = request->prompt;
-    condition_params.clip_skip  = request->clip_skip;
-    condition_params.width      = request->width;
-    condition_params.height     = request->height;
-    condition_params.ref_images = &latents->ref_images;
+    condition_params.text      = request->prompt;
+    condition_params.clip_skip = request->clip_skip;
+    condition_params.width     = request->width;
+    condition_params.height    = request->height;
+    if (ref_image_params.pass_to_vlm) {
+        condition_params.ref_images = &latents->ref_images;
+    }
+
+    condition_params.ref_image_params = ref_image_params;
 
     sd_ctx->sd->prepare_generation_extensions(request->pm_params,
                                               request->pulid_params,
@@ -4602,7 +4843,7 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
     condition_params.zero_out_masked = false;
     auto cond                        = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
                                                                                            condition_params);
-    if (cond.c_concat.empty()) {
+    if (cond.c_concat.empty() && ref_image_params.pass_to_dit) {
         cond.c_concat = latents->concat_latent;  // TODO: optimize
     }
 
@@ -4631,7 +4872,7 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
             uncond                           = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
                                                                                                    condition_params);
         }
-        if (uncond.c_concat.empty()) {
+        if (uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
             uncond.c_concat = latents->concat_latent;  // TODO: optimize
         }
     }
@@ -4655,7 +4896,7 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
             }
             img_uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
                                                                              condition_params);
-            if (img_uncond.c_concat.empty()) {
+            if (img_uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
                 img_uncond.c_concat = latents->img_uncond_concat_latent;  // TODO: optimize
             }
         }
@@ -4972,13 +5213,16 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
     sd_ctx->sd->apply_loras(sd_img_gen_params->loras, sd_img_gen_params->lora_count);
     apply_circular_axes_to_diffusion(sd_ctx, sd_img_gen_params->circular_x, sd_img_gen_params->circular_y);
 
+    const RefImageParams ref_image_params = sd_ctx->sd->resolve_ref_image_params(sd_img_gen_params->ref_image_args);
+
     ImageVaeAxesGuard axes_guard(sd_ctx, sd_img_gen_params, request);
 
     SamplePlan plan(sd_ctx, sd_img_gen_params, request);
     auto latents_opt = prepare_image_generation_latents(sd_ctx,
                                                         sd_img_gen_params,
                                                         &request,
-                                                        &plan);
+                                                        &plan,
+                                                        ref_image_params);
     if (!latents_opt.has_value()) {
         return false;
     }
@@ -4988,7 +5232,8 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
                                                       sd_img_gen_params,
                                                       &request,
                                                       &plan,
-                                                      &latents);
+                                                      &latents,
+                                                      ref_image_params);
     if (!embeds_opt.has_value()) {
         return false;
     }
@@ -5034,7 +5279,7 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
                                                    plan.extra_sample_args,
                                                    plan.sigmas,
                                                    latents.ref_latents,
-                                                   request.increase_ref_index,
+                                                   ref_image_params,
                                                    latents.denoise_mask,
                                                    sd::Tensor<float>(),
                                                    1.f,
@@ -5155,7 +5400,7 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
                                                             plan.extra_sample_args,
                                                             hires_sigma_sched,
                                                             latents.ref_latents,
-                                                            request.increase_ref_index,
+                                                            ref_image_params,
                                                             hires_denoise_mask,
                                                             sd::Tensor<float>(),
                                                             1.f,
@@ -5419,6 +5664,27 @@ static std::optional<ImageGenerationLatents> prepare_video_generation_latents(sd
             sd::ops::fill_slice(&latents.denoise_mask, 2, latents.init_latent.shape()[2] - 1, latents.init_latent.shape()[2], 0.0f);
         }
 
+        int64_t t2 = ggml_time_ms();
+        LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1);
+    } else if (sd_version_is_lingbot_video(sd_ctx->sd->version) && !start_image.empty()) {
+        LOG_INFO("LingBot Video IMG2VID");
+
+        int64_t t1             = ggml_time_ms();
+        auto init_img          = start_image.reshape({start_image.shape()[0], start_image.shape()[1], 1, start_image.shape()[2], 1});
+        auto init_image_latent = sd_ctx->sd->encode_first_stage(init_img);
+        if (init_image_latent.empty()) {
+            LOG_ERROR("failed to encode init video frame");
+            return std::nullopt;
+        }
+
+        latents.init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height, request->frames, true);
+        sd::ops::slice_assign(&latents.init_latent, 2, 0, init_image_latent.shape()[2], init_image_latent);
+
+        latents.denoise_mask = sd::full<float>({latents.init_latent.shape()[0], latents.init_latent.shape()[1], latents.init_latent.shape()[2], 1, 1}, 1.f);
+        sd::ops::fill_slice(&latents.denoise_mask, 2, 0, init_image_latent.shape()[2], 0.0f);
+
+        latents.ref_images.push_back(start_image);
+
         int64_t t2 = ggml_time_ms();
         LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1);
     } else if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-VACE-1.3B" ||
@@ -5524,6 +5790,10 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx,
     condition_params.clip_skip       = request.clip_skip;
     condition_params.text            = request.prompt;
     condition_params.zero_out_masked = true;
+    condition_params.ref_images      = &latents.ref_images;
+    if (sd_version_is_lingbot_video(sd_ctx->sd->version)) {
+        condition_params.ref_image_params.vlm_resize_mode = RefImageResizeMode::AREA;
+    }
 
     int64_t prepare_start_ms = ggml_time_ms();
     embeds.cond              = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
@@ -5804,6 +6074,8 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
 
     sd_ctx->sd->reset_cancel_flag();
 
+    const RefImageParams ref_image_params;
+
     if (num_frames_out != nullptr) {
         *num_frames_out = 0;
     }
@@ -5893,7 +6165,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
                                                            plan.high_noise_extra_sample_args,
                                                            high_noise_sigmas,
                                                            std::vector<sd::Tensor<float>>{},
-                                                           false,
+                                                           ref_image_params,
                                                            latents.denoise_mask,
                                                            latents.vace_context,
                                                            request.vace_strength,
@@ -5935,7 +6207,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
                                                         plan.extra_sample_args,
                                                         plan.sigmas,
                                                         std::vector<sd::Tensor<float>>{},
-                                                        false,
+                                                        ref_image_params,
                                                         latents.denoise_mask,
                                                         latents.vace_context,
                                                         request.vace_strength,
@@ -6073,7 +6345,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
                                             plan.extra_sample_args,
                                             hires_sigma_sched,
                                             std::vector<sd::Tensor<float>>{},
-                                            false,
+                                            ref_image_params,
                                             hires_denoise_mask,
                                             sd::Tensor<float>(),
                                             hires_request.vace_strength,