Compare commits

...

4 Commits

44 changed files with 2949 additions and 182 deletions

View File

@ -312,7 +312,7 @@ set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
if (NOT SD_USE_SYSTEM_GGML) if (NOT SD_USE_SYSTEM_GGML)
# see https://github.com/ggerganov/ggml/pull/682 # see https://github.com/ggerganov/ggml/pull/682
add_definitions(-DGGML_MAX_NAME=128) add_definitions(-DGGML_MAX_NAME=160)
endif() endif()
# deps # deps

View File

@ -57,7 +57,7 @@ API and command-line option may change frequently.***
- [SeFi-Image](./docs/sefi_image.md) - [SeFi-Image](./docs/sefi_image.md)
- [HiDream-O1-Image](./docs/hidream_o1_image.md) - [HiDream-O1-Image](./docs/hidream_o1_image.md)
- [Ideogram4](./docs/ideogram4.md) - [Ideogram4](./docs/ideogram4.md)
- Image Edit Models - [Image Edit Models](./docs/edit.md)
- [FLUX.1-Kontext-dev](./docs/kontext.md) - [FLUX.1-Kontext-dev](./docs/kontext.md)
- [Qwen Image Edit series](./docs/qwen_image_edit.md) - [Qwen Image Edit series](./docs/qwen_image_edit.md)
- [LongCat Image Edit](./docs/longcat_image.md) - [LongCat Image Edit](./docs/longcat_image.md)
@ -68,6 +68,7 @@ API and command-line option may change frequently.***
- [LingBot-Video](./docs/lingbot_video.md) - [LingBot-Video](./docs/lingbot_video.md)
- [PhotoMaker](./docs/photo_maker.md) support. - [PhotoMaker](./docs/photo_maker.md) support.
- Control Net support with SD 1.5 - Control Net support with SD 1.5
- [ADetailer](./docs/adetailer.md)
- LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora) - LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora)
- Latent Consistency Models support (LCM/LCM-LoRA) - Latent Consistency Models support (LCM/LCM-LoRA)
- Faster and memory efficient latent decoding with [TAESD](./docs/taesd.md) - Faster and memory efficient latent decoding with [TAESD](./docs/taesd.md)

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.6 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1002 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.1 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.2 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 MiB

110
docs/adetailer.md Normal file
View File

@ -0,0 +1,110 @@
# ADetailer
`sd-cli` can run a YOLOv8 object detector on an existing or newly generated
image and perform a cropped inpaint pass for every detected object. The first
implementation supports YOLOv8 detection checkpoints. YOLOv8 segmentation and
MediaPipe models are not supported yet.
## Convert a detector
Ultralytics checkpoints must be converted before use. The converter fuses
BatchNorm into convolution layers and writes a safetensors file with the weight
names expected by the native GGML implementation.
```bash
python scripts/convert_yolov8_to_safetensors.py face_yolov8n.pt face_yolov8n.safetensors
```
The converter requires Python packages `ultralytics`, `torch`, and
`safetensors`.
Only YOLOv8 detection checkpoints are accepted.
PyTorch checkpoints use pickle internally, so only convert `.pt` files from a
trusted source.
## Repair an existing image
Use the dedicated `adetailer` mode to detect and repair objects in an existing
image:
```bash
./bin/sd-cli \
-M adetailer \
-m model.safetensors \
-i input.png \
-o repaired.png \
-p "detailed portrait photo" \
--negative-prompt "deformed face" \
--steps 24 \
--cfg-scale 6 \
--strength 0.4 \
--sampling-method dpm++2m \
--scheduler karras \
--ad-model face_yolov8n.safetensors \
--extra-ad-args "confidence=0.3,inpaint_padding=32,mask_blur=4"
```
This mode reuses the normal image-generation options for the detail pass:
- `--init-img`, `--output`, `--prompt`, and `--negative-prompt`
- `--steps`, `--cfg-scale`, `--sampling-method`, and `--scheduler`
- `--strength`, `--seed`, LoRA settings, VAE tiling, and backend assignments
- `--width` and `--height`, which also resize the input when specified
`--ad-prompt` and `--ad-negative-prompt` optionally override the normal prompts.
Values provided in `--extra-ad-args`, such as `steps`, `cfg_scale`,
`denoising_strength`, or `inpaint_width`, take precedence over inherited values.
## Repair generated images
ADetailer can also run automatically after normal image generation:
```bash
./bin/sd-cli \
-m model.safetensors \
-p "portrait photo" \
--ad-model face_yolov8n.safetensors \
--ad-prompt "[PROMPT], detailed face" \
--ad-negative-prompt "" \
--extra-ad-args "confidence=0.3,denoising_strength=0.4,inpaint_width=512,inpaint_height=512"
```
An empty ADetailer prompt inherits the main prompt. `[PROMPT]` inserts the main
prompt, `[SEP]` assigns different prompts to consecutive masks, and `[SKIP]`
skips the corresponding mask.
All settings other than the detector path and prompts are passed through
`--extra-ad-args` as a comma-separated `key=value` list:
| Key | Default | Description |
| --- | ---: | --- |
| `input_size` | `640` | Square YOLO input size; must be a multiple of 32 |
| `confidence` | `0.3` | Detection confidence threshold |
| `nms` | `0.45` | NMS IoU threshold |
| `max_detections` | `100` | Maximum detections retained after NMS |
| `mask_k_largest` | `0` | Keep only the largest K masks; zero keeps all |
| `mask_min_ratio` | `0` | Minimum bbox area relative to the image |
| `mask_max_ratio` | `1` | Maximum bbox area relative to the image |
| `dilate_erode` | `4` | Positive values dilate; negative values erode |
| `x_offset`, `y_offset` | `0` | Mask offset in pixels; positive Y moves upward |
| `mask_mode` | `none` | `none`, `merge`, or `merge_invert` |
| `merge_masks`, `invert_mask` | `false` | Boolean alternatives to `mask_mode` |
| `mask_blur` | `4` | Final composite feather radius |
| `inpaint_padding` | `32` | Padding around the detected region |
| `inpaint_width`, `inpaint_height` | mode-specific | `512x512` after generation; input/output size in `adetailer` mode |
| `denoising_strength` | mode-specific | `0.4` after generation; inherits `--strength` in `adetailer` mode |
| `steps` | `0` | Detail steps; zero inherits the main generation |
| `cfg_scale` | `-1` | Detail CFG; a negative value inherits the main generation |
| `sample_method` | inherited | Detail sampler name |
| `scheduler` | inherited | Detail scheduler name |
| `sort_by` | `none` | `none`, `left_to_right`, `center_to_edge`, or `area` |
Multiple masks are processed serially. Each completed inpaint becomes the input
for the next mask, and the seed is incremented by the mask index. Use
`mask_mode=merge` to process all detections in one inpaint pass.
The detector uses the `detector` backend module. For example, keep detection on
the CPU while diffusion runs on CUDA:
```bash
--backend "diffusion=cuda0,detector=cpu"
```

152
docs/animatediff.md Normal file
View File

@ -0,0 +1,152 @@
# AnimateDiff (SD 1.5)
AnimateDiff adds motion to a frozen Stable Diffusion 1.5 checkpoint by
injecting a temporal-attention module at 20 UNet slots. The base SD 1.5
model, VAE, and text encoder are unchanged; only the motion module produces
the temporal residual that turns a batch of independent frames into a
coherent animation. Reference: Guo et al., "AnimateDiff: Animate Your
Personalized Text-to-Image Diffusion Models without Specific Tuning"
(https://arxiv.org/abs/2307.04725).
## Download weights
- Motion module (v3, recommended)
- fp16 safetensors: https://huggingface.co/conrevo/AnimateDiff-A1111/resolve/main/motion_module/mm_sd15_v3.safetensors
- original checkpoint: https://huggingface.co/guoyww/animatediff/resolve/main/v3_sd15_mm.ckpt
- SD 1.5 base model
- any SD 1.5 checkpoint works. `realisticVisionV60B1` and `toonyou_beta3`
are the ones used in guoyww's reference configs.
- Domain Adapter LoRA (optional, v3 only, sharpens the base UNet's output
toward the motion module's trained distribution)
- ckpt: https://huggingface.co/guoyww/animatediff/resolve/main/v3_sd15_adapter.ckpt
- place under your `--lora-model-dir` and reference in the prompt as
`<lora:v3_sd15_adapter:1.0>`.
The motion module is `~836 MB` and loads alongside the SD 1.5 UNet via
`--motion-module`.
## Motion module versions
Per [animatediff.net/models](https://animatediff.net/models):
| Module | Base | Native res | Character |
|---------------------|------|------------|-----------|
| `mm_sd_v14.ckpt` | 1.5 | 256x256 | earliest, more jittery |
| `mm_sd_v15.ckpt` | 1.5 | 256x256 | improved stability over v1.4 |
| `mm_sd_v15_v2.ckpt` | 1.5 | 384x384 | significantly better motion dynamics |
| `v3_sd15_mm.ckpt` | 1.5 | 512x512 | smoothest, highest quality; pairs with a Domain Adapter LoRA |
| `mm_sdxl_v10_beta` | SDXL | 512x512 | experimental, not yet supported here |
Match your `-H -W` to the module's native resolution for best results. v3 is
trained at 512x512 - going smaller (e.g. 384x384) still works but the motion
character is closer to v2.
## Examples
Generate an 8-frame animation at 512x512, seed 42, 20 steps. The sampler /
scheduler / CFG values below match what mm_sd15_v3 was trained with; using
SD 1.5 defaults (euler_a, low CFG) produces noise-like output.
```
.\bin\Release\sd-cli.exe -M vid_gen \
--model ..\models\checkpoints\realisticVisionV60B1.safetensors \
--motion-module ..\models\animatediff\mm_sd15_v3.safetensors \
--offload-to-cpu --diffusion-fa \
-p "a red apple on a wooden table" \
--cfg-scale 8.0 --sampling-method euler --scheduler discrete \
-H 512 -W 512 --video-frames 8 --fps 8 --steps 20 -s 42 \
-o out.avi
```
Generate at the motion module's native 16-frame context (recommended for
best temporal quality). Needs more VRAM at 512x512, so drop to 384x384 or
use layer streaming:
```
.\bin\Release\sd-cli.exe -M vid_gen \
--model ..\models\checkpoints\realisticVisionV60B1.safetensors \
--motion-module ..\models\animatediff\mm_sd15_v3.safetensors \
--offload-to-cpu --diffusion-fa \
-p "photo of coastline, rocks, storm weather, wind, waves, lightning" \
--cfg-scale 8.0 --sampling-method euler --scheduler discrete \
-H 384 -W 384 --video-frames 16 --fps 8 --steps 20 -s 42 \
-o out.avi
```
Low-VRAM streaming (verified with a 2 GiB cap on RTX 3060):
```
.\bin\Release\sd-cli.exe -M vid_gen \
--model ..\models\checkpoints\realisticVisionV60B1.safetensors \
--motion-module ..\models\animatediff\mm_sd15_v3.safetensors \
--max-vram 2.0 --stream-layers --diffusion-fa \
-p "photo of coastline, rocks, storm weather, wind, waves, lightning" \
--cfg-scale 8.0 --sampling-method euler --scheduler discrete \
-H 384 -W 384 --video-frames 8 --fps 8 --steps 20 -s 42 \
-o out.avi
```
## Reference-quality reproduction
Using guoyww's official reference configs on this impl (RealisticVision v6.0
base + `mm_sd15_v3` or `mm_sd_v15_v2` + native resolution + 16 frames + euler
+ 25 steps + CFG 8 + linear beta schedule) reproduces the reference
AnimateDiff output style.
### v3 (512x512, `mm_sd15_v3`)
| Prompt | Sample |
|---------------------------------------|--------|
| B&W man on stormy coastline | <img src="../assets/animatediff/v3_man.gif" width="256"/> |
| Close-up rabbit macro shot | <img src="../assets/animatediff/v3_rabbit.gif" width="256"/> |
| Coastline, storm, waves, lightning | <img src="../assets/animatediff/v3_coast.gif" width="256"/> |
| Old house, storm, forest, night | <img src="../assets/animatediff/v3_house.gif" width="256"/> |
### v2 (384x384, `mm_sd_v15_v2.ckpt`)
| Prompt | Sample |
|---------------------------------------|--------|
| B&W man on stormy coastline | <img src="../assets/animatediff/v2_man.gif" width="256"/> |
| Close-up rabbit macro shot | <img src="../assets/animatediff/v2_rabbit.gif" width="256"/> |
| Coastline, storm, waves, lightning | <img src="../assets/animatediff/v2_coast.gif" width="256"/> |
| Old house, storm, forest, night | <img src="../assets/animatediff/v2_house.gif" width="256"/> |
Motion is strong for scenes with motion cues in the prompt (storm/waves/wind)
and subtle for static subjects (close-up macro), matching reference behavior.
v2 has an additional motion module at the UNet middle block that v3 dropped;
this impl auto-detects the topology from the checkpoint.
### v3 + Domain Adapter LoRA
Attaching the `v3_sd15_adapter` LoRA sharpens the base UNet output toward
the training distribution the motion module was fine-tuned against. Same
prompt, seed, config as above:
<img src="../assets/animatediff/v3_rabbit_domain_lora.gif" width="256"/>
Individual fur strands, glowing inner-ear, and richer forest detail become
visible compared to the no-LoRA rendering.
```
sd-cli -M vid_gen --model realisticVisionV60B1.safetensors \
--motion-module mm_sd15_v3.safetensors \
--lora-model-dir ./loras \
-p "close up photo of a rabbit ...<lora:v3_sd15_adapter:1.0>" ...
```
## Notes
- The motion module was trained at `video_length=16`. Running with
`--video-frames 16` gives the best coherence; F=8 works but shows a shorter
motion arc. Frame counts up to 32 are supported by the positional encoding
but exceed the trained regime and produce more static output.
- At `--video-frames 1` the motion module is skipped entirely and the output
is bit-identical to `-M img_gen`. This avoids the single-token
temporal-attention degeneracy that would otherwise emit an untrained-magnitude
residual on a single-frame sample.
- The base UNet is frozen, so character identity and style follow the SD 1.5
checkpoint you pass to `--model`. LoRAs and prompt weighting attach to the
base model in the usual way.
- No mid_block motion module in v3. `mm_sdxl_v10_beta` (SDXL variant) is not
supported yet.
- Output is written as MJPEG AVI. Use `--fps` to set playback speed.

View File

@ -153,6 +153,7 @@ still runs out of memory, tiling is enabled and the decode retried once.
| `controlnet` | ControlNet | `controlnet`, `control` | | `controlnet` | ControlNet | `controlnet`, `control` |
| `photomaker` | PhotoMaker ID encoder and PhotoMaker LoRA | `photomaker`, `photomakerid`, `pmid`, `photo` | | `photomaker` | PhotoMaker ID encoder and PhotoMaker LoRA | `photomaker`, `photomakerid`, `pmid`, `photo` |
| `upscaler` | ESRGAN upscaler | `upscaler`, `esrgan`, `hires` | | `upscaler` | ESRGAN upscaler | `upscaler`, `esrgan`, `hires` |
| `detector` | ADetailer YOLOv8 detector | `detector`, `adetailer`, `yolo` |
`te` is the preferred module name for text encoders. `clip` is kept as an accepted alias because many existing commands and model names use CLIP terminology. `te` is the preferred module name for text encoders. `clip` is kept as an accepted alias because many existing commands and model names use CLIP terminology.

93
docs/edit.md Normal file
View File

@ -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`).

View File

@ -6,6 +6,9 @@ For detailed command-line arguments, run:
./bin/sd-cli -h ./bin/sd-cli -h
``` ```
For direct image repair or automatic post-generation YOLOv8 detection followed by cropped inpainting, see
[ADetailer](../../docs/adetailer.md).
Metadata mode inspects PNG/JPEG container metadata without loading any model: Metadata mode inspects PNG/JPEG container metadata without loading any model:
```bash ```bash

View File

@ -199,7 +199,7 @@ struct SDCliParams {
options.manual_options = { options.manual_options = {
{"-M", {"-M",
"--mode", "--mode",
"run mode, one of [img_gen, vid_gen, upscale, convert, metadata], default: img_gen", "run mode, one of [img_gen, adetailer, vid_gen, upscale, convert, metadata], default: img_gen",
on_mode_arg}, on_mode_arg},
{"", {"",
"--preview", "--preview",
@ -566,6 +566,65 @@ bool save_results(const SDCliParams& cli_params,
return sucessful_reults != 0; return sucessful_reults != 0;
} }
static bool apply_adetailer(sd_ctx_t* sd_ctx,
const sd_ctx_params_t& sd_ctx_params,
const SDContextParams& ctx_params,
const SDGenerationParams& gen_params,
const sd_img_gen_params_t& img_gen_params,
SDMode mode,
SDImageVec& results,
int num_results) {
if (gen_params.ad_model_path.empty()) {
return true;
}
sd_adetailer_params_t ad_params{};
ad_params.prompt = gen_params.ad_prompt.empty() ? nullptr : gen_params.ad_prompt.c_str();
ad_params.negative_prompt = gen_params.ad_negative_prompt.empty() ? nullptr : gen_params.ad_negative_prompt.c_str();
ad_params.extra_ad_args = gen_params.extra_ad_args.c_str();
ADetailerCtxPtr ad_ctx(new_adetailer_ctx(gen_params.ad_model_path.c_str(),
ctx_params.n_threads,
sd_ctx_params.backend,
sd_ctx_params.params_backend));
if (ad_ctx == nullptr) {
LOG_ERROR("new_adetailer_ctx failed");
return false;
}
for (int i = 0; i < num_results; ++i) {
if (results[i].data == nullptr) {
continue;
}
sd_img_gen_params_t ad_generation_params = img_gen_params;
ad_generation_params.seed = img_gen_params.seed + i;
if (mode == IMG_GEN) {
ad_generation_params.width = 512;
ad_generation_params.height = 512;
ad_generation_params.strength = 0.4f;
}
sd_image_t* detailed_images = nullptr;
int detailed_count = 0;
if (!adetail_image(ad_ctx.get(),
sd_ctx,
results[i],
&ad_params,
&ad_generation_params,
&detailed_images,
&detailed_count) ||
detailed_count <= 0 || detailed_images == nullptr || detailed_images[0].data == nullptr) {
free_sd_images(detailed_images, detailed_count);
LOG_ERROR("ADetailer failed for image %d", i + 1);
return false;
}
free(results[i].data);
results[i] = detailed_images[0];
detailed_images[0] = {0, 0, 0, nullptr};
free_sd_images(detailed_images, detailed_count);
}
return true;
}
int main(int argc, const char* argv[]) { int main(int argc, const char* argv[]) {
if (argc > 1 && std::string(argv[1]) == "--version") { if (argc > 1 && std::string(argv[1]) == "--version") {
std::cout << version_string() << "\n"; std::cout << version_string() << "\n";
@ -598,6 +657,11 @@ int main(int argc, const char* argv[]) {
return 0; return 0;
} }
if (!gen_params.ad_model_path.empty() && cli_params.mode != IMG_GEN && cli_params.mode != ADETAILER) {
LOG_ERROR("--ad-model is only supported in image generation and adetailer modes");
return 1;
}
if (gen_params.video_frames > 4) { if (gen_params.video_frames > 4) {
size_t last_dot_pos = cli_params.preview_path.find_last_of("."); size_t last_dot_pos = cli_params.preview_path.find_last_of(".");
std::string base_path = cli_params.preview_path; std::string base_path = cli_params.preview_path;
@ -806,15 +870,22 @@ int main(int argc, const char* argv[]) {
gen_params.sample_params.scheduler = sd_get_default_scheduler(sd_ctx.get(), gen_params.sample_params.sample_method); gen_params.sample_params.scheduler = sd_get_default_scheduler(sd_ctx.get(), gen_params.sample_params.sample_method);
} }
if (cli_params.mode == IMG_GEN) { sd_img_gen_params_t img_gen_params{};
sd_img_gen_params_t img_gen_params = gen_params.to_sd_img_gen_params_t(); const bool use_img_gen_params = cli_params.mode == IMG_GEN || cli_params.mode == ADETAILER;
if (use_img_gen_params) {
img_gen_params = gen_params.to_sd_img_gen_params_t();
}
if (cli_params.mode == IMG_GEN) {
sd_image_t* generated_images = nullptr; sd_image_t* generated_images = nullptr;
if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) { if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) {
generated_images = nullptr; generated_images = nullptr;
num_results = 0; num_results = 0;
} }
results.adopt(generated_images, num_results); results.adopt(generated_images, num_results);
} else if (cli_params.mode == ADETAILER) {
num_results = 1;
results.push_back(gen_params.init_image.release());
} else if (cli_params.mode == VID_GEN) { } else if (cli_params.mode == VID_GEN) {
sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t(); sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t();
sd_image_t* generated_video = nullptr; sd_image_t* generated_video = nullptr;
@ -828,6 +899,18 @@ int main(int argc, const char* argv[]) {
LOG_ERROR("generate failed"); LOG_ERROR("generate failed");
return 1; return 1;
} }
if (use_img_gen_params &&
!apply_adetailer(sd_ctx.get(),
sd_ctx_params,
ctx_params,
gen_params,
img_gen_params,
cli_params.mode,
results,
num_results)) {
return 1;
}
} }
int upscale_factor = 4; // unused for RealESRGAN_x4plus_anime_6B.pth int upscale_factor = 4; // unused for RealESRGAN_x4plus_anime_6B.pth

View File

@ -30,6 +30,7 @@ namespace fs = std::filesystem;
const char* const modes_str[] = { const char* const modes_str[] = {
"img_gen", "img_gen",
"adetailer",
"vid_gen", "vid_gen",
"convert", "convert",
"upscale", "upscale",
@ -423,6 +424,11 @@ ArgOptions SDContextParams::get_options() {
"path to control net model", "path to control net model",
0, 0,
&control_net_path}, &control_net_path},
{"",
"--motion-module",
"path to AnimateDiff motion module (SD 1.5); enables video generation on --video-frames > 1",
0,
&motion_module_path},
{"", {"",
"--embd-dir", "--embd-dir",
"embeddings directory", "embeddings directory",
@ -672,7 +678,7 @@ ArgOptions SDContextParams::get_options() {
} }
void SDContextParams::build_embedding_map() { void SDContextParams::build_embedding_map() {
static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt"}; static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt", ".ckpt"};
if (!fs::exists(embedding_dir) || !fs::is_directory(embedding_dir)) { if (!fs::exists(embedding_dir) || !fs::is_directory(embedding_dir)) {
return; return;
@ -867,6 +873,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
sd_ctx_params.audio_vae_path = audio_vae_path.c_str(); sd_ctx_params.audio_vae_path = audio_vae_path.c_str();
sd_ctx_params.taesd_path = taesd_path.c_str(); sd_ctx_params.taesd_path = taesd_path.c_str();
sd_ctx_params.control_net_path = control_net_path.c_str(); sd_ctx_params.control_net_path = control_net_path.c_str();
sd_ctx_params.motion_module_path = motion_module_path.c_str();
sd_ctx_params.embeddings = embedding_vec.data(); sd_ctx_params.embeddings = embedding_vec.data();
sd_ctx_params.embedding_count = static_cast<uint32_t>(embedding_vec.size()); sd_ctx_params.embedding_count = static_cast<uint32_t>(embedding_vec.size());
sd_ctx_params.photo_maker_path = photo_maker_path.c_str(); sd_ctx_params.photo_maker_path = photo_maker_path.c_str();
@ -916,6 +923,26 @@ ArgOptions SDGenerationParams::get_options() {
"the negative prompt (default: \"\")", "the negative prompt (default: \"\")",
0, 0,
&negative_prompt}, &negative_prompt},
{"",
"--ad-model",
"path to a converted YOLOv8 detection model for ADetailer",
0,
&ad_model_path},
{"",
"--ad-prompt",
"ADetailer prompt; empty inherits the main prompt, supports [PROMPT], [SEP], and [SKIP]",
0,
&ad_prompt},
{"",
"--ad-negative-prompt",
"ADetailer negative prompt; empty inherits the main negative prompt, supports [PROMPT] and [SEP]",
0,
&ad_negative_prompt},
{"",
"--extra-ad-args",
"extra ADetailer args, key=value list. Supports input_size, confidence, nms, max_detections, mask_k_largest, mask_min_ratio, mask_max_ratio, dilate_erode, x_offset, y_offset, mask_mode, merge_masks, invert_mask, mask_blur, inpaint_padding, inpaint_width, inpaint_height, denoising_strength, steps, cfg_scale, sample_method, scheduler, sort_by",
(int)',',
&extra_ad_args},
{"-i", {"-i",
"--init-img", "--init-img",
"path to the init image", "path to the init image",
@ -975,6 +1002,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)", "extra VAE tiling args, key=value list. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)",
(int)',', (int)',',
&extra_tiling_args}, &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 = { options.int_options = {
@ -1831,6 +1863,10 @@ bool SDGenerationParams::from_json_str(
load_if_exists("prompt", prompt); load_if_exists("prompt", prompt);
load_if_exists("negative_prompt", negative_prompt); load_if_exists("negative_prompt", negative_prompt);
load_if_exists("ad_model", ad_model_path);
load_if_exists("ad_prompt", ad_prompt);
load_if_exists("ad_negative_prompt", ad_negative_prompt);
load_if_exists("extra_ad_args", extra_ad_args);
load_if_exists("cache_mode", cache_mode); load_if_exists("cache_mode", cache_mode);
load_if_exists("cache_option", cache_option); load_if_exists("cache_option", cache_option);
load_if_exists("scm_mask", scm_mask); load_if_exists("scm_mask", scm_mask);
@ -2027,7 +2063,7 @@ void SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_d
return; return;
} }
static const std::regex re(R"(<lora:([^:>]+):([^>]+)>)"); static const std::regex re(R"(<lora:([^:>]+):([^>]+)>)");
static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt"}; static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt", ".ckpt"};
std::smatch m; std::smatch m;
std::string tmp = prompt; std::string tmp = prompt;
@ -2347,13 +2383,19 @@ bool SDGenerationParams::validate(SDMode mode) {
} }
} }
if (mode == UPSCALE) { if (mode == UPSCALE || mode == ADETAILER) {
if (init_image_path.length() == 0) { if (init_image_path.length() == 0) {
LOG_ERROR("error: upscale mode needs an init image (--init-img)\n"); LOG_ERROR("error: %s mode needs an init image (--init-img)\n",
mode == UPSCALE ? "upscale" : "adetailer");
return false; return false;
} }
} }
if (mode == ADETAILER && ad_model_path.empty()) {
LOG_ERROR("error: adetailer mode needs a detector model (--ad-model)\n");
return false;
}
return true; return true;
} }
@ -2418,30 +2460,45 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
pulid_id_weight, pulid_id_weight,
}; };
params.loras = lora_vec.empty() ? nullptr : lora_vec.data(); if (!auto_resize_ref_image) {
params.lora_count = static_cast<uint32_t>(lora_vec.size()); if (!ref_image_args.empty()) {
params.prompt = prompt.c_str(); ref_image_args += ",";
params.negative_prompt = negative_prompt.c_str(); }
params.clip_skip = clip_skip; ref_image_args += "resize_before_vae=0";
params.init_image = init_image.get(); LOG_WARN("Notice: --disable-auto-resize-ref-image is deprecated. Use --ref-image-args \"resize_before_vae=off\" instead.");
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; if (increase_ref_index) {
params.increase_ref_index = increase_ref_index; if (!ref_image_args.empty()) {
params.mask_image = mask_image.get(); ref_image_args += ",";
params.width = get_resolved_width(); }
params.height = get_resolved_height(); ref_image_args += "ref_index_mode=increase";
params.sample_params = sample_params; LOG_WARN("Notice: --increase-ref-index is deprecated. Use --ref-image-args \"ref_index_mode=increase\" instead.");
params.strength = strength; }
params.seed = seed;
params.batch_count = batch_count; params.loras = lora_vec.empty() ? nullptr : lora_vec.data();
params.qwen_image_layers = qwen_image_layers; params.lora_count = static_cast<uint32_t>(lora_vec.size());
params.control_image = control_image.get(); params.prompt = prompt.c_str();
params.control_strength = control_strength; params.negative_prompt = negative_prompt.c_str();
params.pm_params = pm_params; params.clip_skip = clip_skip;
params.pulid_params = pulid_params; params.init_image = init_image.get();
params.vae_tiling_params = vae_tiling_params; params.ref_images = ref_image_views.empty() ? nullptr : ref_image_views.data();
params.cache = cache_params; 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.enabled = hires_enabled;
params.hires.upscaler = resolved_hires_upscaler; params.hires.upscaler = resolved_hires_upscaler;
@ -2561,6 +2618,10 @@ std::string SDGenerationParams::to_string() const {
<< " high_noise_loras: \"" << high_noise_loras_str << "\",\n" << " high_noise_loras: \"" << high_noise_loras_str << "\",\n"
<< " prompt: \"" << prompt << "\",\n" << " prompt: \"" << prompt << "\",\n"
<< " negative_prompt: \"" << negative_prompt << "\",\n" << " negative_prompt: \"" << negative_prompt << "\",\n"
<< " ad_model_path: \"" << ad_model_path << "\",\n"
<< " ad_prompt: \"" << ad_prompt << "\",\n"
<< " ad_negative_prompt: \"" << ad_negative_prompt << "\",\n"
<< " extra_ad_args: \"" << extra_ad_args << "\",\n"
<< " clip_skip: " << clip_skip << ",\n" << " clip_skip: " << clip_skip << ",\n"
<< " width: " << width << ",\n" << " width: " << width << ",\n"
<< " height: " << height << ",\n" << " height: " << height << ",\n"
@ -2675,8 +2736,13 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
int64_t seed, int64_t seed,
SDMode mode) { SDMode mode) {
json root; json root;
root["schema"] = "sdcpp.image.params/v1"; root["schema"] = "sdcpp.image.params/v1";
root["mode"] = mode == VID_GEN ? "vid_gen" : "img_gen"; root["mode"] = "img_gen";
if (mode == VID_GEN) {
root["mode"] = "vid_gen";
} else if (mode == ADETAILER) {
root["mode"] = "adetailer";
}
root["generator"] = { root["generator"] = {
{"name", "stable-diffusion.cpp"}, {"name", "stable-diffusion.cpp"},
{"version", safe_json_string(sd_version())}, {"version", safe_json_string(sd_version())},
@ -2690,6 +2756,14 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
{"positive", gen_params.prompt}, {"positive", gen_params.prompt},
{"negative", gen_params.negative_prompt}, {"negative", gen_params.negative_prompt},
}; };
if (!gen_params.ad_model_path.empty()) {
root["adetailer"] = {
{"model", sd_basename(gen_params.ad_model_path)},
{"prompt", gen_params.ad_prompt},
{"negative_prompt", gen_params.ad_negative_prompt},
{"extra_args", gen_params.extra_ad_args},
};
}
root["sampling"] = build_sampling_metadata_json(gen_params.sample_params, root["sampling"] = build_sampling_metadata_json(gen_params.sample_params,
gen_params.skip_layers, gen_params.skip_layers,
&gen_params.custom_sigmas); &gen_params.custom_sigmas);
@ -2844,6 +2918,18 @@ std::string get_image_params(const SDContextParams& ctx_params,
if (!gen_params.extra_sample_args.empty()) { if (!gen_params.extra_sample_args.empty()) {
parameter_string += "Extra sample args: " + gen_params.extra_sample_args + ", "; parameter_string += "Extra sample args: " + gen_params.extra_sample_args + ", ";
} }
if (!gen_params.ad_model_path.empty()) {
parameter_string += "ADetailer model: " + sd_basename(gen_params.ad_model_path) + ", ";
if (!gen_params.ad_prompt.empty()) {
parameter_string += "ADetailer prompt: " + gen_params.ad_prompt + ", ";
}
if (!gen_params.ad_negative_prompt.empty()) {
parameter_string += "ADetailer negative prompt: " + gen_params.ad_negative_prompt + ", ";
}
if (!gen_params.extra_ad_args.empty()) {
parameter_string += "ADetailer args: " + gen_params.extra_ad_args + ", ";
}
}
parameter_string += "Seed: " + std::to_string(seed) + ", "; parameter_string += "Seed: " + std::to_string(seed) + ", ";
parameter_string += "Size: " + std::to_string(gen_params.get_resolved_width()) + "x" + std::to_string(gen_params.get_resolved_height()) + ", "; parameter_string += "Size: " + std::to_string(gen_params.get_resolved_width()) + "x" + std::to_string(gen_params.get_resolved_height()) + ", ";
parameter_string += "Model: " + sd_basename(ctx_params.model_path) + ", "; parameter_string += "Model: " + sd_basename(ctx_params.model_path) + ", ";

View File

@ -16,10 +16,11 @@
#define BOOL_STR(b) ((b) ? "true" : "false") #define BOOL_STR(b) ((b) ? "true" : "false")
extern const char* const modes_str[]; extern const char* const modes_str[];
#define SD_ALL_MODES_STR "img_gen, vid_gen, convert, upscale, metadata" #define SD_ALL_MODES_STR "img_gen, adetailer, vid_gen, convert, upscale, metadata"
enum SDMode { enum SDMode {
IMG_GEN, IMG_GEN,
ADETAILER,
VID_GEN, VID_GEN,
CONVERT, CONVERT,
UPSCALE, UPSCALE,
@ -132,6 +133,7 @@ struct SDContextParams {
std::string taesd_path; std::string taesd_path;
std::string esrgan_path; std::string esrgan_path;
std::string control_net_path; std::string control_net_path;
std::string motion_module_path;
std::string embedding_dir; std::string embedding_dir;
std::string photo_maker_path; std::string photo_maker_path;
std::string pulid_weights_path; std::string pulid_weights_path;
@ -186,6 +188,10 @@ struct SDGenerationParams {
// User-facing input fields. // User-facing input fields.
std::string prompt; std::string prompt;
std::string negative_prompt; std::string negative_prompt;
std::string ad_model_path;
std::string ad_prompt;
std::string ad_negative_prompt;
std::string extra_ad_args;
int clip_skip = -1; // <= 0 represents unspecified int clip_skip = -1; // <= 0 represents unspecified
int width = -1; int width = -1;
int height = -1; int height = -1;
@ -227,6 +233,8 @@ struct SDGenerationParams {
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr}; 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 extra_tiling_args;
std::string ref_image_args;
std::string pm_id_images_dir; std::string pm_id_images_dir;
std::string pm_id_embed_path; std::string pm_id_embed_path;
float pm_style_strength = 20.f; float pm_style_strength = 20.f;

View File

@ -40,12 +40,21 @@ struct UpscalerCtxDeleter {
} }
}; };
struct ADetailerCtxDeleter {
void operator()(adetailer_ctx_t* ctx) const {
if (ctx != nullptr) {
free_adetailer_ctx(ctx);
}
}
};
template <typename T> template <typename T>
using FreeUniquePtr = std::unique_ptr<T, FreeDeleter>; using FreeUniquePtr = std::unique_ptr<T, FreeDeleter>;
using FilePtr = std::unique_ptr<FILE, FileCloser>; using FilePtr = std::unique_ptr<FILE, FileCloser>;
using SDCtxPtr = std::unique_ptr<sd_ctx_t, SDCtxDeleter>; using SDCtxPtr = std::unique_ptr<sd_ctx_t, SDCtxDeleter>;
using UpscalerCtxPtr = std::unique_ptr<upscaler_ctx_t, UpscalerCtxDeleter>; using UpscalerCtxPtr = std::unique_ptr<upscaler_ctx_t, UpscalerCtxDeleter>;
using ADetailerCtxPtr = std::unique_ptr<adetailer_ctx_t, ADetailerCtxDeleter>;
class SDImageOwner { class SDImageOwner {
private: private:

View File

@ -199,6 +199,7 @@ typedef struct {
const char* audio_vae_path; const char* audio_vae_path;
const char* taesd_path; const char* taesd_path;
const char* control_net_path; const char* control_net_path;
const char* motion_module_path;
const sd_embedding_t* embeddings; const sd_embedding_t* embeddings;
uint32_t embedding_count; uint32_t embedding_count;
const char* photo_maker_path; const char* photo_maker_path;
@ -363,8 +364,7 @@ typedef struct {
sd_image_t init_image; sd_image_t init_image;
sd_image_t* ref_images; sd_image_t* ref_images;
int ref_images_count; int ref_images_count;
bool auto_resize_ref_image; const char* ref_image_args;
bool increase_ref_index;
sd_image_t mask_image; sd_image_t mask_image;
int width; int width;
int height; int height;
@ -509,6 +509,27 @@ SD_API bool upscale(upscaler_ctx_t* upscaler_ctx,
SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx); SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx);
typedef struct adetailer_ctx_t adetailer_ctx_t;
typedef struct {
const char* prompt;
const char* negative_prompt;
const char* extra_ad_args;
} sd_adetailer_params_t;
SD_API adetailer_ctx_t* new_adetailer_ctx(const char* detector_path,
int n_threads,
const char* backend,
const char* params_backend);
SD_API void free_adetailer_ctx(adetailer_ctx_t* adetailer_ctx);
SD_API bool adetail_image(adetailer_ctx_t* adetailer_ctx,
sd_ctx_t* sd_ctx,
sd_image_t input_image,
const sd_adetailer_params_t* adetailer_params,
const sd_img_gen_params_t* inpaint_params,
sd_image_t** images_out,
int* num_images_out);
SD_API bool convert(const char* input_path, SD_API bool convert(const char* input_path,
const char* vae_path, const char* vae_path,
const char* output_path, const char* output_path,

View File

@ -0,0 +1,86 @@
#!/usr/bin/env python3
"""Convert an Ultralytics YOLOv8 detection checkpoint for sd.cpp ADetailer."""
import argparse
import json
from pathlib import Path
def parse_args():
parser = argparse.ArgumentParser(
description="Convert an Ultralytics YOLOv8 detection .pt checkpoint to safetensors."
)
parser.add_argument("input", type=Path, help="input YOLOv8 detection checkpoint")
parser.add_argument("output", type=Path, help="output safetensors path")
parser.add_argument(
"--input-size", type=int, default=640, help="detector input size metadata (default: 640)"
)
return parser.parse_args()
def main():
args = parse_args()
if args.input_size < 32 or args.input_size % 32 != 0:
raise ValueError("--input-size must be a positive multiple of 32")
if args.output.suffix.lower() != ".safetensors":
raise ValueError("output path must use the .safetensors extension")
try:
import torch
from safetensors.torch import save_file
from ultralytics import YOLO
from ultralytics.nn.modules.head import Detect
except ImportError as exc:
raise SystemExit("conversion requires ultralytics, torch, and safetensors") from exc
torch_load = torch.load
def load_trusted_checkpoint(*load_args, **load_kwargs):
load_kwargs.setdefault("weights_only", False)
return torch_load(*load_args, **load_kwargs)
torch.load = load_trusted_checkpoint
try:
yolo = YOLO(str(args.input))
finally:
torch.load = torch_load
network = yolo.model
if not isinstance(network.model[-1], Detect) or network.model[-1].__class__.__name__ != "Detect":
raise ValueError("only YOLOv8 detection checkpoints are supported; segmentation is not yet supported")
network.eval()
network.fuse()
state_dict = network.state_dict()
required = {
"model.0.conv.weight",
"model.22.cv2.0.2.weight",
"model.22.cv3.0.2.weight",
}
missing = sorted(required.difference(state_dict))
if missing:
raise ValueError(f"checkpoint does not match the supported YOLOv8 layout; missing {missing}")
tensors = {}
for name, tensor in state_dict.items():
if not name.startswith("model.") or ".bn." in name or name.endswith("dfl.conv.weight"):
continue
if not (name.endswith(".weight") or name.endswith(".bias")):
continue
dtype = torch.float16 if name.endswith(".weight") else torch.float32
tensors[name] = tensor.detach().to(device="cpu", dtype=dtype).contiguous()
metadata = {
"format": "pt",
"yolov8.variant": "detect",
"yolov8.input_size": str(args.input_size),
"yolov8.num_classes": str(int(network.model[-1].nc)),
"yolov8.reg_max": str(int(network.model[-1].reg_max)),
"yolov8.names": json.dumps(yolo.names, ensure_ascii=False),
}
args.output.parent.mkdir(parents=True, exist_ok=True)
save_file(tensors, str(args.output), metadata=metadata)
print(f"wrote {args.output}: {len(tensors)} tensors")
if __name__ == "__main__":
main()

View File

@ -7,6 +7,7 @@
#include "core/tensor_ggml.hpp" #include "core/tensor_ggml.hpp"
#include "core/util.h" #include "core/util.h"
#include "model/diffusion/model.hpp"
#include "model/te/clip.hpp" #include "model/te/clip.hpp"
#include "model/te/llm.hpp" #include "model/te/llm.hpp"
#include "model/te/t5.hpp" #include "model/te/t5.hpp"
@ -106,6 +107,7 @@ struct ConditionerParams {
int height = -1; int height = -1;
bool zero_out_masked = false; bool zero_out_masked = false;
const std::vector<sd::Tensor<float>>* ref_images = nullptr; // for qwen image edit const std::vector<sd::Tensor<float>>* ref_images = nullptr; // for qwen image edit
RefImageParams ref_image_params;
}; };
struct Conditioner { struct Conditioner {
@ -1993,6 +1995,54 @@ struct LLMEmbedder : public Conditioner {
return new_hidden_states; 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, SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override { const ConditionerParams& conditioner_params) override {
std::string prompt; std::string prompt;
@ -2007,7 +2057,8 @@ struct LLMEmbedder : public Conditioner {
bool spell_quotes = false; bool spell_quotes = false;
std::set<int> out_layers; 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_lingbot_video(version)) { if (sd_version_is_lingbot_video(version)) {
const int pad_token = 151643; const int pad_token = 151643;
@ -2040,28 +2091,35 @@ struct LLMEmbedder : public Conditioner {
for (int i = 0; i < conditioner_params.ref_images->size(); i++) { for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[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 height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]); int width = static_cast<int>(image.shape()[0]);
int min_pixels = static_cast<int>(4 * factor * factor);
int max_pixels = static_cast<int>(16384 * factor * factor); int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
int h_bar = std::max(static_cast<int>(factor), static_cast<int>(std::round(height / factor) * factor)); if (min_pixels <= 0) {
int w_bar = std::max(static_cast<int>(factor), static_cast<int>(std::round(width / factor) * factor)); 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)) { if (std::max(height, width) > 200 * std::min(height, width)) {
LOG_WARN("LingBotVideo image aspect ratio is very large: %dx%d", width, height); LOG_WARN("LingBotVideo image aspect ratio is very large: %dx%d", width, height);
} }
if (h_bar * w_bar > max_pixels) {
double beta = std::sqrt((height * width) / static_cast<double>(max_pixels)); resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
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 (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);
}
LOG_DEBUG("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); 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 resized_image = clip_preprocess(image, w_bar, h_bar);
@ -2092,30 +2150,33 @@ struct LLMEmbedder : public Conditioner {
prompt_template_encode_start_idx = 64; prompt_template_encode_start_idx = 64;
int image_embed_idx = 64 + 6; int image_embed_idx = 64 + 6;
int min_pixels = 384 * 384; int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
int max_pixels = 560 * 560; 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 placeholder = "<|image_pad|>";
std::string img_prompt; std::string img_prompt;
for (int i = 0; i < conditioner_params.ref_images->size(); i++) { for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[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 height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]); int width = static_cast<int>(image.shape()[0]);
int h_bar = static_cast<int>(std::round(height / factor) * factor); int h_bar = static_cast<int>(std::round(static_cast<double>(height) / factor) * factor);
int w_bar = static_cast<int>(std::round(width / 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) { resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
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);
}
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
@ -2170,16 +2231,33 @@ struct LLMEmbedder : public Conditioner {
std::string img_prompt; std::string img_prompt;
const std::string placeholder = "<|image_pad|>"; 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++) { for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[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 height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]); 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), int h_bar = std::max(factor,
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor)); static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
int w_bar = std::max(static_cast<int>(factor), int w_bar = std::max(factor,
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(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); LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
@ -2221,17 +2299,33 @@ struct LLMEmbedder : public Conditioner {
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
std::string img_prompt = ""; std::string img_prompt = "";
const std::string placeholder = "<|image_pad|>"; 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++) { for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[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 height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]); 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), int h_bar = std::max(factor,
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor)); static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
int w_bar = std::max(static_cast<int>(factor), int w_bar = std::max(factor,
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(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); LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
@ -2268,30 +2362,33 @@ struct LLMEmbedder : public Conditioner {
min_length = 512 + prompt_template_encode_start_idx; min_length = 512 + prompt_template_encode_start_idx;
int image_embed_idx = 36 + 6; int image_embed_idx = 36 + 6;
int min_pixels = 384 * 384; int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
int max_pixels = 560 * 560; 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 placeholder = "<|image_pad|>";
std::string img_prompt; std::string img_prompt;
for (int i = 0; i < conditioner_params.ref_images->size(); i++) { for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[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 height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]); int width = static_cast<int>(image.shape()[0]);
int h_bar = static_cast<int>(std::round(height / factor) * factor); int h_bar = static_cast<int>(std::round(static_cast<double>(height) / factor) * factor);
int w_bar = static_cast<int>(std::round(width / 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) { resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
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);
}
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);

View File

@ -83,6 +83,10 @@ static bool parse_backend_module(const std::string& raw_name, SDBackendModule* m
*module = SDBackendModule::UPSCALER; *module = SDBackendModule::UPSCALER;
return true; return true;
} }
if (name == "detector" || name == "adetailer" || name == "yolo") {
*module = SDBackendModule::DETECTOR;
return true;
}
return false; return false;
} }
@ -956,6 +960,8 @@ const char* sd_backend_module_name(SDBackendModule module) {
return "photomaker"; return "photomaker";
case SDBackendModule::UPSCALER: case SDBackendModule::UPSCALER:
return "upscaler"; return "upscaler";
case SDBackendModule::DETECTOR:
return "detector";
} }
return "unknown"; return "unknown";
} }

View File

@ -20,6 +20,7 @@ enum class SDBackendModule {
CONTROL_NET, CONTROL_NET,
PHOTOMAKER, PHOTOMAKER,
UPSCALER, UPSCALER,
DETECTOR,
}; };
struct SDBackendAssignment { struct SDBackendAssignment {

1020
src/detailer.cpp Normal file

File diff suppressed because it is too large Load Diff

75
src/detailer.h Normal file
View File

@ -0,0 +1,75 @@
#ifndef __SD_DETAILER_H__
#define __SD_DETAILER_H__
#include <memory>
#include <string>
#include <vector>
#include "core/ggml_extend_backend.h"
#include "model/detector/yolov8.h"
#include "model_manager.h"
#include "stable-diffusion.h"
struct ADetailerDetection {
float x1 = 0.f;
float y1 = 0.f;
float x2 = 0.f;
float y2 = 0.f;
float confidence = 0.f;
int class_id = 0;
};
enum ADetailerSort {
ADETAILER_SORT_NONE,
ADETAILER_SORT_LEFT_TO_RIGHT,
ADETAILER_SORT_CENTER_TO_EDGE,
ADETAILER_SORT_AREA,
};
struct ADetailerParams {
const char* prompt = nullptr;
const char* negative_prompt = nullptr;
int input_size = 640;
float confidence = 0.3f;
float nms_threshold = 0.45f;
int max_detections = 100;
int mask_k_largest = 0;
float mask_min_ratio = 0.f;
float mask_max_ratio = 1.f;
int dilate_erode = 4;
int x_offset = 0;
int y_offset = 0;
bool merge_masks = false;
bool invert_mask = false;
int mask_blur = 4;
int inpaint_padding = 32;
int inpaint_width = 512;
int inpaint_height = 512;
float denoising_strength = 0.4f;
int steps = 0;
float cfg_scale = -1.f;
sample_method_t sample_method = SAMPLE_METHOD_COUNT;
scheduler_t scheduler = SCHEDULER_COUNT;
ADetailerSort sort_by = ADETAILER_SORT_NONE;
};
struct ADetailerGGML {
SDBackendManager backend_manager;
std::shared_ptr<ModelManager> model_manager;
std::shared_ptr<YOLOv8Runner> detector;
std::vector<std::string> class_names;
int n_threads = 1;
std::string backend_spec;
std::string params_backend_spec;
ADetailerGGML(int n_threads,
std::string backend_spec,
std::string params_backend_spec);
~ADetailerGGML();
bool load_from_file(const std::string& detector_path);
std::vector<ADetailerDetection> predict(sd_image_t image,
const ADetailerParams& params);
};
#endif // __SD_DETAILER_H__

View File

@ -18,10 +18,6 @@ namespace Rope {
DECREASE, DECREASE,
}; };
__STATIC_INLINE__ RefIndexMode ref_index_mode_from_bool(bool increase_ref_index) {
return increase_ref_index ? RefIndexMode::INCREASE : RefIndexMode::FIXED;
}
template <class T> template <class T>
__STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) { __STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
std::vector<T> result(num); std::vector<T> result(num);

362
src/model/detector/yolov8.h Normal file
View File

@ -0,0 +1,362 @@
#ifndef __SD_MODEL_DETECTOR_YOLOV8_H__
#define __SD_MODEL_DETECTOR_YOLOV8_H__
#include <algorithm>
#include <array>
#include <map>
#include <memory>
#include <string>
#include <vector>
#include "core/ggml_extend.hpp"
#include "core/util.h"
struct YOLOv8Config {
std::array<int, 23> out_channels{};
std::map<int, int> hidden_channels;
std::map<int, int> repeats;
int detect_box_channels = 0;
int detect_cls_channels = 0;
int reg_max = 0;
int num_classes = 0;
bool valid = false;
static YOLOv8Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "") {
YOLOv8Config config;
auto full_name = [&](const std::string& name) {
return prefix.empty() ? name : prefix + "." + name;
};
auto find_weight = [&](const std::string& name) -> const TensorStorage* {
auto iter = tensor_storage_map.find(full_name(name));
return iter == tensor_storage_map.end() ? nullptr : &iter->second;
};
auto conv_out = [&](const std::string& name) -> int {
const TensorStorage* weight = find_weight(name);
return weight != nullptr && weight->n_dims == 4 ? static_cast<int>(weight->ne[3]) : 0;
};
for (int layer : {0, 1, 3, 5, 7, 16, 19}) {
config.out_channels[layer] = conv_out("model." + std::to_string(layer) + ".conv.weight");
}
for (int layer : {2, 4, 6, 8, 12, 15, 18, 21}) {
const std::string base = "model." + std::to_string(layer);
config.out_channels[layer] = conv_out(base + ".cv2.conv.weight");
config.hidden_channels[layer] = conv_out(base + ".cv1.conv.weight") / 2;
int repeat_count = 0;
while (find_weight(base + ".m." + std::to_string(repeat_count) + ".cv1.conv.weight") != nullptr) {
++repeat_count;
}
config.repeats[layer] = repeat_count;
}
config.out_channels[9] = conv_out("model.9.cv2.conv.weight");
config.detect_box_channels = conv_out("model.22.cv2.0.0.conv.weight");
config.detect_cls_channels = conv_out("model.22.cv3.0.0.conv.weight");
const int box_outputs = conv_out("model.22.cv2.0.2.weight");
config.num_classes = conv_out("model.22.cv3.0.2.weight");
config.reg_max = box_outputs / 4;
config.valid = config.out_channels[0] > 0 && config.out_channels[9] > 0 &&
config.out_channels[15] > 0 && config.out_channels[18] > 0 &&
config.out_channels[21] > 0 && config.detect_box_channels > 0 &&
config.detect_cls_channels > 0 && box_outputs > 0 && box_outputs % 4 == 0 &&
config.num_classes > 0;
for (int layer : {2, 4, 6, 8, 12, 15, 18, 21}) {
config.valid = config.valid && config.hidden_channels[layer] > 0 && config.repeats[layer] > 0;
}
if (config.valid) {
LOG_DEBUG("yolov8: classes=%d, reg_max=%d, p3=%d, p4=%d, p5=%d",
config.num_classes,
config.reg_max,
config.out_channels[15],
config.out_channels[18],
config.out_channels[21]);
}
return config;
}
};
class YOLOConv : public UnaryBlock {
int out_channels_ = 0;
public:
YOLOConv(int in_channels, int out_channels, int kernel, int stride = 1)
: out_channels_(out_channels) {
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels,
out_channels,
{kernel, kernel},
{stride, stride},
{kernel / 2, kernel / 2},
{1, 1},
true));
}
int out_channels() const {
return out_channels_;
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
return ggml_silu_inplace(ctx->ggml_ctx, conv->forward(ctx, x));
}
};
class YOLOBottleneck : public UnaryBlock {
bool shortcut_ = false;
public:
YOLOBottleneck(int channels, bool shortcut)
: shortcut_(shortcut) {
blocks["cv1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(channels, channels, 3));
blocks["cv2"] = std::shared_ptr<GGMLBlock>(new YOLOConv(channels, channels, 3));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto cv1 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv1"]);
auto cv2 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv2"]);
auto out = cv2->forward(ctx, cv1->forward(ctx, x));
return shortcut_ ? ggml_add(ctx->ggml_ctx, x, out) : out;
}
};
class YOLOC2f : public UnaryBlock {
int hidden_channels_ = 0;
int repeats_ = 0;
public:
YOLOC2f(int in_channels,
int out_channels,
int hidden_channels,
int repeats,
bool shortcut)
: hidden_channels_(hidden_channels), repeats_(repeats) {
blocks["cv1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels, hidden_channels * 2, 1));
blocks["cv2"] = std::shared_ptr<GGMLBlock>(new YOLOConv(hidden_channels * (2 + repeats), out_channels, 1));
for (int i = 0; i < repeats; ++i) {
blocks["m." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new YOLOBottleneck(hidden_channels, shortcut));
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto cv1 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv1"]);
auto cv2 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv2"]);
auto split = cv1->forward(ctx, x);
// split: [N, 2*C, H, W], ggml layout [W, H, 2*C, N].
auto y0 = ggml_view_4d(ctx->ggml_ctx,
split,
split->ne[0],
split->ne[1],
hidden_channels_,
split->ne[3],
split->nb[1],
split->nb[2],
split->nb[3],
0);
auto y1 = ggml_view_4d(ctx->ggml_ctx,
split,
split->ne[0],
split->ne[1],
hidden_channels_,
split->ne[3],
split->nb[1],
split->nb[2],
split->nb[3],
static_cast<size_t>(hidden_channels_) * split->nb[2]);
auto joined = ggml_concat(ctx->ggml_ctx, y0, y1, 2);
auto last = y1;
for (int i = 0; i < repeats_; ++i) {
auto block = std::dynamic_pointer_cast<YOLOBottleneck>(blocks["m." + std::to_string(i)]);
last = block->forward(ctx, last);
joined = ggml_concat(ctx->ggml_ctx, joined, last, 2);
}
return cv2->forward(ctx, joined);
}
};
class YOLOSPPF : public UnaryBlock {
public:
YOLOSPPF(int in_channels, int out_channels) {
blocks["cv1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels, in_channels / 2, 1));
blocks["cv2"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels * 2, out_channels, 1));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto cv1 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv1"]);
auto cv2 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv2"]);
x = cv1->forward(ctx, x);
auto y1 = ggml_pool_2d(ctx->ggml_ctx, x, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2);
auto y2 = ggml_pool_2d(ctx->ggml_ctx, y1, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2);
auto y3 = ggml_pool_2d(ctx->ggml_ctx, y2, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2);
auto out = ggml_concat(ctx->ggml_ctx, x, y1, 2);
out = ggml_concat(ctx->ggml_ctx, out, y2, 2);
out = ggml_concat(ctx->ggml_ctx, out, y3, 2);
return cv2->forward(ctx, out);
}
};
class YOLODetect : public GGMLBlock {
int num_classes_ = 0;
int reg_max_ = 0;
public:
YOLODetect(const std::array<int, 3>& in_channels,
int box_channels,
int cls_channels,
int reg_max,
int num_classes)
: num_classes_(num_classes), reg_max_(reg_max) {
for (int i = 0; i < 3; ++i) {
const std::string box = "cv2." + std::to_string(i);
blocks[box + ".0"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels[i], box_channels, 3));
blocks[box + ".1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(box_channels, box_channels, 3));
blocks[box + ".2"] = std::shared_ptr<GGMLBlock>(new Conv2d(box_channels, reg_max * 4, {1, 1}, {1, 1}, {0, 0}, {1, 1}, true));
const std::string cls = "cv3." + std::to_string(i);
blocks[cls + ".0"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels[i], cls_channels, 3));
blocks[cls + ".1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(cls_channels, cls_channels, 3));
blocks[cls + ".2"] = std::shared_ptr<GGMLBlock>(new Conv2d(cls_channels, num_classes, {1, 1}, {1, 1}, {0, 0}, {1, 1}, true));
}
}
ggml_tensor* forward_scale(GGMLRunnerContext* ctx, ggml_tensor* x, int index) {
const std::string box = "cv2." + std::to_string(index);
auto box0 = std::dynamic_pointer_cast<YOLOConv>(blocks[box + ".0"]);
auto box1 = std::dynamic_pointer_cast<YOLOConv>(blocks[box + ".1"]);
auto box2 = std::dynamic_pointer_cast<Conv2d>(blocks[box + ".2"]);
const std::string cls = "cv3." + std::to_string(index);
auto cls0 = std::dynamic_pointer_cast<YOLOConv>(blocks[cls + ".0"]);
auto cls1 = std::dynamic_pointer_cast<YOLOConv>(blocks[cls + ".1"]);
auto cls2 = std::dynamic_pointer_cast<Conv2d>(blocks[cls + ".2"]);
auto boxes = box2->forward(ctx, box1->forward(ctx, box0->forward(ctx, x)));
auto classes = cls2->forward(ctx, cls1->forward(ctx, cls0->forward(ctx, x)));
return ggml_concat(ctx->ggml_ctx, boxes, classes, 2);
}
int output_channels() const {
return reg_max_ * 4 + num_classes_;
}
};
class YOLOv8Model : public GGMLBlock {
YOLOv8Config config_;
std::shared_ptr<YOLOC2f> make_c2f(int layer, int in_channels, bool shortcut) {
return std::make_shared<YOLOC2f>(in_channels,
config_.out_channels[layer],
config_.hidden_channels.at(layer),
config_.repeats.at(layer),
shortcut);
}
public:
explicit YOLOv8Model(YOLOv8Config config)
: config_(std::move(config)) {
blocks["model.0"] = std::make_shared<YOLOConv>(3, config_.out_channels[0], 3, 2);
blocks["model.1"] = std::make_shared<YOLOConv>(config_.out_channels[0], config_.out_channels[1], 3, 2);
blocks["model.2"] = make_c2f(2, config_.out_channels[1], true);
blocks["model.3"] = std::make_shared<YOLOConv>(config_.out_channels[2], config_.out_channels[3], 3, 2);
blocks["model.4"] = make_c2f(4, config_.out_channels[3], true);
blocks["model.5"] = std::make_shared<YOLOConv>(config_.out_channels[4], config_.out_channels[5], 3, 2);
blocks["model.6"] = make_c2f(6, config_.out_channels[5], true);
blocks["model.7"] = std::make_shared<YOLOConv>(config_.out_channels[6], config_.out_channels[7], 3, 2);
blocks["model.8"] = make_c2f(8, config_.out_channels[7], true);
blocks["model.9"] = std::make_shared<YOLOSPPF>(config_.out_channels[8], config_.out_channels[9]);
blocks["model.12"] = make_c2f(12, config_.out_channels[9] + config_.out_channels[6], false);
blocks["model.15"] = make_c2f(15, config_.out_channels[12] + config_.out_channels[4], false);
blocks["model.16"] = std::make_shared<YOLOConv>(config_.out_channels[15], config_.out_channels[16], 3, 2);
blocks["model.18"] = make_c2f(18, config_.out_channels[16] + config_.out_channels[12], false);
blocks["model.19"] = std::make_shared<YOLOConv>(config_.out_channels[18], config_.out_channels[19], 3, 2);
blocks["model.21"] = make_c2f(21, config_.out_channels[19] + config_.out_channels[9], false);
blocks["model.22"] = std::make_shared<YOLODetect>(
std::array<int, 3>{config_.out_channels[15], config_.out_channels[18], config_.out_channels[21]},
config_.detect_box_channels,
config_.detect_cls_channels,
config_.reg_max,
config_.num_classes);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto run = [&](int layer, ggml_tensor* input) {
return std::dynamic_pointer_cast<UnaryBlock>(blocks["model." + std::to_string(layer)])->forward(ctx, input);
};
auto x0 = run(0, x);
auto x1 = run(1, x0);
auto x2 = run(2, x1);
auto x3 = run(3, x2);
auto x4 = run(4, x3);
auto x5 = run(5, x4);
auto x6 = run(6, x5);
auto x7 = run(7, x6);
auto x8 = run(8, x7);
auto x9 = run(9, x8);
auto x12 = run(12, ggml_concat(ctx->ggml_ctx, ggml_upscale(ctx->ggml_ctx, x9, 2, GGML_SCALE_MODE_NEAREST), x6, 2));
auto x15 = run(15, ggml_concat(ctx->ggml_ctx, ggml_upscale(ctx->ggml_ctx, x12, 2, GGML_SCALE_MODE_NEAREST), x4, 2));
auto x16 = run(16, x15);
auto x18 = run(18, ggml_concat(ctx->ggml_ctx, x16, x12, 2));
auto x19 = run(19, x18);
auto x21 = run(21, ggml_concat(ctx->ggml_ctx, x19, x9, 2));
auto detect = std::dynamic_pointer_cast<YOLODetect>(blocks["model.22"]);
auto p3 = detect->forward_scale(ctx, x15, 0);
auto p4 = detect->forward_scale(ctx, x18, 1);
auto p5 = detect->forward_scale(ctx, x21, 2);
p3 = ggml_reshape_2d(ctx->ggml_ctx, p3, p3->ne[0] * p3->ne[1], detect->output_channels());
p4 = ggml_reshape_2d(ctx->ggml_ctx, p4, p4->ne[0] * p4->ne[1], detect->output_channels());
p5 = ggml_reshape_2d(ctx->ggml_ctx, p5, p5->ne[0] * p5->ne[1], detect->output_channels());
return ggml_concat(ctx->ggml_ctx, ggml_concat(ctx->ggml_ctx, p3, p4, 0), p5, 0);
}
};
struct YOLOv8Runner : public GGMLRunner {
YOLOv8Config config;
std::unique_ptr<YOLOv8Model> model;
YOLOv8Runner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager),
config(YOLOv8Config::detect_from_weights(tensor_storage_map)) {
if (config.valid) {
model = std::make_unique<YOLOv8Model>(config);
model->init(params_ctx, tensor_storage_map, "");
}
}
std::string get_desc() override {
return "yolov8";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
if (model) {
model->get_param_tensors(tensors);
}
}
ggml_cgraph* build_graph(const sd::Tensor<float>& input) {
if (!model) {
return nullptr;
}
ggml_cgraph* graph = new_graph_custom(1 << 16);
auto x = make_input(input);
auto runner_ctx = get_context();
auto output = model->forward(&runner_ctx, x);
ggml_build_forward_expand(graph, output);
return graph;
}
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& input) {
auto get_graph = [&]() { return build_graph(input); };
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false));
}
};
#endif // __SD_MODEL_DETECTOR_YOLOV8_H__

View File

@ -484,10 +484,11 @@ namespace Anima {
ggml_tensor* timestep, ggml_tensor* timestep,
ggml_tensor* encoder_hidden_states, ggml_tensor* encoder_hidden_states,
ggml_tensor* image_pe, ggml_tensor* image_pe,
ggml_tensor* t5_ids = nullptr, ggml_tensor* t5_ids = nullptr,
ggml_tensor* t5_weights = nullptr, ggml_tensor* t5_weights = nullptr,
ggml_tensor* adapter_q_pe = nullptr, ggml_tensor* adapter_q_pe = nullptr,
ggml_tensor* adapter_k_pe = nullptr) { ggml_tensor* adapter_k_pe = nullptr,
std::vector<ggml_tensor*> ref_latents = {}) {
GGML_ASSERT(x->ne[3] == 1); GGML_ASSERT(x->ne[3] == 1);
auto x_embedder = std::dynamic_pointer_cast<XEmbedder>(blocks["x_embedder"]); 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]); 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 = 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); x = x_embedder->forward(ctx, x);
auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size)); 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); 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"); 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] 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, const std::vector<int>& axes_dim,
float h_extrapolation_ratio, float h_extrapolation_ratio,
float w_extrapolation_ratio, float w_extrapolation_ratio,
float t_extrapolation_ratio) { float t_extrapolation_ratio,
static const std::vector<ggml_tensor*> empty_ref_latents; const std::vector<ggml_tensor*>& ref_latents) {
auto ids = Rope::gen_flux_ids(h, auto ids = Rope::gen_flux_ids(h,
w, w,
patch_size, patch_size,
@ -611,7 +621,7 @@ namespace Anima {
static_cast<int>(axes_dim.size()), static_cast<int>(axes_dim.size()),
0, 0,
{}, {},
empty_ref_latents, ref_latents,
Rope::RefIndexMode::FIXED, Rope::RefIndexMode::FIXED,
1.0f, 1.0f,
false); false);
@ -626,14 +636,20 @@ namespace Anima {
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor, ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor, const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor = {}, const sd::Tensor<float>& context_tensor = {},
const sd::Tensor<int32_t>& t5_ids_tensor = {}, const sd::Tensor<int32_t>& t5_ids_tensor = {},
const sd::Tensor<float>& t5_weights_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* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor); ggml_tensor* timesteps = make_input(timesteps_tensor);
ggml_tensor* context = make_optional_input(context_tensor); ggml_tensor* context = make_optional_input(context_tensor);
ggml_tensor* t5_ids = make_optional_input(t5_ids_tensor); ggml_tensor* t5_ids = make_optional_input(t5_ids_tensor);
ggml_tensor* t5_weights = make_optional_input(t5_weights_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_ASSERT(x->ne[3] == 1);
ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE); ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE);
@ -650,7 +666,8 @@ namespace Anima {
config.axes_dim, config.axes_dim,
4.0f, 4.0f,
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)); 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); 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()); set_backend_tensor_data(image_pe, image_pe_vec.data());
@ -682,7 +699,8 @@ namespace Anima {
t5_ids, t5_ids,
t5_weights, t5_weights,
adapter_q_pe, adapter_q_pe,
adapter_k_pe); adapter_k_pe,
ref_latents);
ggml_build_forward_expand(gf, out); ggml_build_forward_expand(gf, out);
return gf; return gf;
@ -691,11 +709,13 @@ namespace Anima {
sd::Tensor<float> compute(int n_threads, sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x, const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps, const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context = {}, const sd::Tensor<float>& context = {},
const sd::Tensor<int32_t>& t5_ids = {}, const sd::Tensor<int32_t>& t5_ids = {},
const sd::Tensor<float>& t5_weights = {}) { 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* { 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()); 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.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr); GGML_ASSERT(diffusion_params.timesteps != nullptr);
const auto* extra = diffusion_extra_as<AnimaDiffusionExtra>(diffusion_params); const auto* extra = diffusion_extra_as<AnimaDiffusionExtra>(diffusion_params);
static const std::vector<sd::Tensor<float>> empty_ref_latents;
return compute(n_threads, return compute(n_threads,
*diffusion_params.x, *diffusion_params.x,
*diffusion_params.timesteps, *diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context), tensor_or_empty(diffusion_params.context),
tensor_or_empty(extra->t5_ids), 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 } // namespace Anima

View File

@ -0,0 +1,182 @@
#ifndef __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__
#define __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__
#include "core/ggml_extend.hpp"
#include "model/common/block.hpp"
// AnimateDiff (https://arxiv.org/abs/2307.04725) SD 1.5 motion modules.
namespace AnimateDiff {
struct MotionModuleConfig {
int max_frames = 32;
int64_t num_heads = 8;
int norm_num_groups = 32;
std::vector<int64_t> down_channels = {320, 640, 1280, 1280};
std::vector<int64_t> up_channels = {1280, 1280, 640, 320};
int num_down_motion_per_block = 2;
int num_up_motion_per_block = 3;
bool enable_mid_block = false;
int64_t mid_channels = 1280;
};
class TemporalAttention : public GGMLBlock {
protected:
int64_t channels;
int64_t num_heads;
int max_frames;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
params["pos_encoder.pe"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, channels, max_frames, 1);
}
public:
TemporalAttention(int64_t channels, int64_t num_heads, int max_frames)
: channels(channels), num_heads(num_heads), max_frames(max_frames) {
blocks["to_q"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, false));
blocks["to_k"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, false));
blocks["to_v"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, false));
blocks["to_out.0"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, true));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
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 to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
int64_t C = x->ne[0];
int64_t F = x->ne[1];
auto pe = params["pos_encoder.pe"];
auto pe_f = (F == pe->ne[1])
? pe
: ggml_view_3d(ctx->ggml_ctx, pe, C, F, 1, pe->nb[1], pe->nb[2], 0);
auto x_pe = ggml_add(ctx->ggml_ctx, x, ggml_repeat(ctx->ggml_ctx, pe_f, x));
auto q = to_q->forward(ctx, x_pe);
auto k = to_k->forward(ctx, x_pe);
auto v = to_v->forward(ctx, x_pe);
auto a = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, (int)num_heads, nullptr, false);
return to_out->forward(ctx, a);
}
};
class TemporalTransformerBlock : public GGMLBlock {
public:
TemporalTransformerBlock(int64_t channels, int64_t num_heads, int max_frames) {
blocks["attention_blocks.0"] = std::make_shared<TemporalAttention>(channels, num_heads, max_frames);
blocks["attention_blocks.1"] = std::make_shared<TemporalAttention>(channels, num_heads, max_frames);
blocks["norms.0"] = std::shared_ptr<GGMLBlock>(new LayerNorm(channels));
blocks["norms.1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(channels));
blocks["ff"] = std::make_shared<FeedForward>(channels, channels, 4, FeedForward::Activation::GEGLU);
blocks["ff_norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(channels));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto attn0 = std::dynamic_pointer_cast<TemporalAttention>(blocks["attention_blocks.0"]);
auto attn1 = std::dynamic_pointer_cast<TemporalAttention>(blocks["attention_blocks.1"]);
auto norm0 = std::dynamic_pointer_cast<LayerNorm>(blocks["norms.0"]);
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norms.1"]);
auto ff = std::dynamic_pointer_cast<FeedForward>(blocks["ff"]);
auto ff_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["ff_norm"]);
auto r = x;
x = ggml_add(ctx->ggml_ctx, attn0->forward(ctx, norm0->forward(ctx, x)), r);
r = x;
x = ggml_add(ctx->ggml_ctx, attn1->forward(ctx, norm1->forward(ctx, x)), r);
r = x;
x = ggml_add(ctx->ggml_ctx, ff->forward(ctx, ff_norm->forward(ctx, x)), r);
return x;
}
};
class TemporalTransformer : public GGMLBlock {
public:
TemporalTransformer(int64_t channels, int64_t num_heads, int norm_num_groups, int max_frames) {
blocks["norm"] = std::shared_ptr<GGMLBlock>(new GroupNorm(norm_num_groups, channels));
blocks["proj_in"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, true));
blocks["transformer_blocks.0"] = std::make_shared<TemporalTransformerBlock>(channels, num_heads, max_frames);
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, true));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t num_frames) {
auto norm = std::dynamic_pointer_cast<GroupNorm>(blocks["norm"]);
auto proj_in = std::dynamic_pointer_cast<Linear>(blocks["proj_in"]);
auto tb0 = std::dynamic_pointer_cast<TemporalTransformerBlock>(blocks["transformer_blocks.0"]);
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 C = x->ne[2];
GGML_ASSERT(x->ne[3] == num_frames);
auto residual = x;
auto h = norm->forward(ctx, x);
h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1));
h = ggml_reshape_3d(ctx->ggml_ctx, h, C, num_frames, W * H);
h = proj_in->forward(ctx, h);
h = tb0->forward(ctx, h);
h = proj_out->forward(ctx, h);
h = ggml_reshape_4d(ctx->ggml_ctx, h, C, num_frames, W, H);
h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1));
return ggml_add(ctx->ggml_ctx, h, residual);
}
};
class MotionModule : public GGMLBlock {
public:
MotionModule(int64_t channels, int64_t num_heads, int norm_num_groups, int max_frames) {
blocks["temporal_transformer"] = std::make_shared<TemporalTransformer>(channels, num_heads, norm_num_groups, max_frames);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t num_frames) {
auto tt = std::dynamic_pointer_cast<TemporalTransformer>(blocks["temporal_transformer"]);
return tt->forward(ctx, x, num_frames);
}
};
class AnimateDiffModel : public GGMLBlock {
public:
MotionModuleConfig config;
AnimateDiffModel(const MotionModuleConfig& cfg)
: config(cfg) {
for (int i = 0; i < static_cast<int>(cfg.down_channels.size()); ++i) {
int64_t ch = cfg.down_channels[i];
for (int j = 0; j < cfg.num_down_motion_per_block; ++j) {
blocks["down_blocks." + std::to_string(i) + ".motion_modules." + std::to_string(j)] =
std::make_shared<MotionModule>(ch, cfg.num_heads, cfg.norm_num_groups, cfg.max_frames);
}
}
for (int i = 0; i < static_cast<int>(cfg.up_channels.size()); ++i) {
int64_t ch = cfg.up_channels[i];
for (int j = 0; j < cfg.num_up_motion_per_block; ++j) {
blocks["up_blocks." + std::to_string(i) + ".motion_modules." + std::to_string(j)] =
std::make_shared<MotionModule>(ch, cfg.num_heads, cfg.norm_num_groups, cfg.max_frames);
}
}
if (cfg.enable_mid_block) {
blocks["mid_block.motion_modules.0"] =
std::make_shared<MotionModule>(cfg.mid_channels, cfg.num_heads, cfg.norm_num_groups, cfg.max_frames);
}
}
std::shared_ptr<MotionModule> motion(const std::string& key) {
auto it = blocks.find(key);
if (it == blocks.end())
return nullptr;
return std::dynamic_pointer_cast<MotionModule>(it->second);
}
};
} // namespace AnimateDiff
#endif // __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__

View File

@ -827,7 +827,7 @@ namespace Boogu {
*diffusion_params.x, *diffusion_params.x,
*diffusion_params.timesteps, *diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context), 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 } // namespace Boogu

View File

@ -1642,8 +1642,8 @@ namespace Flux {
tensor_or_empty(diffusion_params.c_concat), tensor_or_empty(diffusion_params.c_concat),
tensor_or_empty(diffusion_params.y), tensor_or_empty(diffusion_params.y),
tensor_or_empty(extra->guidance), tensor_or_empty(extra->guidance),
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,
diffusion_params.ref_index_mode, diffusion_params.ref_image_params.ref_index_mode,
extra->skip_layers ? *extra->skip_layers : empty_skip_layers, extra->skip_layers ? *extra->skip_layers : empty_skip_layers,
tensor_or_empty(extra->pulid_id), tensor_or_empty(extra->pulid_id),
extra->pulid_id_weight); extra->pulid_id_weight);

View File

@ -615,7 +615,8 @@ namespace Krea2 {
ggml_tensor* timestep, ggml_tensor* timestep,
ggml_tensor* context, ggml_tensor* context,
ggml_tensor* pe, ggml_tensor* pe,
std::vector<ggml_tensor*> ref_latents = {}) { std::vector<ggml_tensor*> ref_latents = {},
bool zero_timestep_refs = false) {
int64_t W = x->ne[0]; int64_t W = x->ne[0];
int64_t H = x->ne[1]; int64_t H = x->ne[1];
int64_t N = x->ne[3]; int64_t N = x->ne[3];
@ -645,7 +646,7 @@ namespace Krea2 {
auto tvec = tproj->forward(ctx, t); auto tvec = tproj->forward(ctx, t);
ggml_tensor* tvec_0 = nullptr; ggml_tensor* tvec_0 = nullptr;
if (ref_latents.size() > 0) { if (ref_latents.size() > 0 && zero_timestep_refs) {
// "index_timestep_zero" mode: use timestep = 0 for ref latents // "index_timestep_zero" mode: use timestep = 0 for ref latents
auto timestep_0 = ggml_scale(ctx->ggml_ctx, timestep, 0.0f); 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); auto t_0 = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_0, static_cast<int>(config.timestep_dim), 10000, 1000.f);
@ -719,7 +720,7 @@ namespace Krea2 {
const sd::Tensor<float>& timesteps_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 std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) { const RefImageParams& ref_image_params = REF_IMAGE_PRESETS.at("krea2_ostris_edit")) {
ggml_cgraph* gf = new_graph_custom(KREA2_GRAPH_SIZE); ggml_cgraph* gf = new_graph_custom(KREA2_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor); ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor); ggml_tensor* timesteps = make_input(timesteps_tensor);
@ -741,13 +742,13 @@ namespace Krea2 {
config.theta, config.theta,
config.axes_dim, config.axes_dim,
ref_latents, ref_latents,
ref_index_mode); ref_image_params.ref_index_mode);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); 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); 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()); set_backend_tensor_data(pe, pe_vec.data());
auto runner_ctx = get_context(); auto runner_ctx = get_context();
ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents); 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); ggml_build_forward_expand(gf, out);
return gf; return gf;
} }
@ -757,9 +758,9 @@ namespace Krea2 {
const sd::Tensor<float>& timesteps, const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context, const sd::Tensor<float>& context,
const std::vector<sd::Tensor<float>>& ref_latents = {}, const std::vector<sd::Tensor<float>>& ref_latents = {},
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) { const RefImageParams& ref_image_params = REF_IMAGE_PRESETS.at("krea2_ostris_edit")) {
auto get_graph = [&]() -> ggml_cgraph* { auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, ref_latents, ref_index_mode); 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()); return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
} }
@ -773,8 +774,8 @@ namespace Krea2 {
*diffusion_params.x, *diffusion_params.x,
*diffusion_params.timesteps, *diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context), 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,
diffusion_params.ref_index_mode); diffusion_params.ref_image_params);
} }
}; };
} // namespace Krea2 } // namespace Krea2

View File

@ -10,6 +10,36 @@
#include "model/common/rope.hpp" #include "model/common/rope.hpp"
#include "model_manager.h" #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 { struct UNetDiffusionExtra {
int num_video_frames = -1; int num_video_frames = -1;
const std::vector<sd::Tensor<float>>* controls = nullptr; 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>* c_concat = nullptr;
const sd::Tensor<float>* y = nullptr; const sd::Tensor<float>* y = nullptr;
const std::vector<sd::Tensor<float>>* ref_latents = 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{}; DiffusionExtraParams extra = std::monostate{};
}; };

View File

@ -715,8 +715,8 @@ namespace Qwen {
*diffusion_params.x, *diffusion_params.x,
*diffusion_params.timesteps, *diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context), 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,
diffusion_params.ref_index_mode); diffusion_params.ref_image_params.ref_index_mode);
} }
void test() { void test() {

View File

@ -6,6 +6,7 @@
#include "model.h" #include "model.h"
#include "model/common/block.hpp" #include "model/common/block.hpp"
#include "model/diffusion/animatediff.hpp"
#include "model/diffusion/model.hpp" #include "model/diffusion/model.hpp"
/*==================================================== UnetModel =====================================================*/ /*==================================================== UnetModel =====================================================*/
@ -29,6 +30,8 @@ struct UNetConfig {
bool tiny_unet = false; bool tiny_unet = false;
int model_channels = 320; int model_channels = 320;
int adm_in_channels = 2816; // only for VERSION_SDXL/SVD int adm_in_channels = 2816; // only for VERSION_SDXL/SVD
bool enable_animatediff = false;
bool animatediff_has_mid_block = false;
static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix, const std::string& prefix,
@ -84,6 +87,13 @@ struct UNetConfig {
return &it->second; return &it->second;
}; };
if (find_weight("motion_module.down_blocks.0.motion_modules.0.temporal_transformer.proj_in.weight") != nullptr) {
config.enable_animatediff = true;
if (find_weight("motion_module.mid_block.motion_modules.0.temporal_transformer.proj_in.weight") != nullptr) {
config.animatediff_has_mid_block = true;
}
}
if (const TensorStorage* input = find_weight("input_blocks.0.0.weight")) { if (const TensorStorage* input = find_weight("input_blocks.0.0.weight")) {
if (input->n_dims == 4) { if (input->n_dims == 4) {
config.in_channels = static_cast<int>(input->ne[2]); config.in_channels = static_cast<int>(input->ne[2]);
@ -473,6 +483,12 @@ public:
blocks["out.0"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(ch)); // ch == model_channels blocks["out.0"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(ch)); // ch == model_channels
// out_1 is nn.SiLU() // out_1 is nn.SiLU()
blocks["out.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(model_channels, out_channels, {3, 3}, {1, 1}, {1, 1})); blocks["out.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(model_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
if (this->config.enable_animatediff) {
AnimateDiff::MotionModuleConfig mm_cfg;
mm_cfg.enable_mid_block = this->config.animatediff_has_mid_block;
blocks["motion_module"] = std::make_shared<AnimateDiff::AnimateDiffModel>(mm_cfg);
}
} }
ggml_tensor* resblock_forward(std::string name, ggml_tensor* resblock_forward(std::string name,
@ -583,6 +599,42 @@ public:
ggml_set_name(h, "bench-start"); ggml_set_name(h, "bench-start");
hs.push_back(h); hs.push_back(h);
auto motion_root = config.enable_animatediff && num_video_frames > 1
? std::dynamic_pointer_cast<AnimateDiff::AnimateDiffModel>(blocks["motion_module"])
: nullptr;
auto apply_motion_input = [&](int input_block_idx, ggml_tensor* h_in) -> ggml_tensor* {
if (!motion_root)
return h_in;
int di = (input_block_idx - 1) / 3;
int mj = (input_block_idx - 1) % 3;
if (di < 0 || di >= (int)channel_mult.size() || mj < 0 || mj >= num_res_blocks)
return h_in;
auto mm = motion_root->motion("down_blocks." + std::to_string(di) + ".motion_modules." + std::to_string(mj));
if (!mm)
return h_in;
return mm->forward(ctx, h_in, num_video_frames);
};
auto apply_motion_output = [&](int output_block_idx, ggml_tensor* h_in) -> ggml_tensor* {
if (!motion_root)
return h_in;
int ui = output_block_idx / 3;
int mj = output_block_idx % 3;
if (ui < 0 || ui >= (int)channel_mult.size() || mj < 0 || mj > num_res_blocks)
return h_in;
auto mm = motion_root->motion("up_blocks." + std::to_string(ui) + ".motion_modules." + std::to_string(mj));
if (!mm)
return h_in;
return mm->forward(ctx, h_in, num_video_frames);
};
auto apply_motion_mid = [&](ggml_tensor* h_in) -> ggml_tensor* {
if (!motion_root)
return h_in;
auto mm = motion_root->motion("mid_block.motion_modules.0");
if (!mm)
return h_in;
return mm->forward(ctx, h_in, num_video_frames);
};
// input block 1-11 // input block 1-11
size_t len_mults = channel_mult.size(); size_t len_mults = channel_mult.size();
int input_block_idx = 0; int input_block_idx = 0;
@ -597,6 +649,7 @@ public:
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1"; std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
h = attention_layer_forward(name, ctx, h, context, num_video_frames); // [N, mult*model_channels, h, w] h = attention_layer_forward(name, ctx, h, context, num_video_frames); // [N, mult*model_channels, h, w]
} }
h = apply_motion_input(input_block_idx, h);
sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h"); sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h");
hs.push_back(h); hs.push_back(h);
} }
@ -624,6 +677,7 @@ public:
h = attention_layer_forward("middle_block.1", ctx, h, context, num_video_frames); // [N, 4*model_channels, h/8, w/8] h = attention_layer_forward("middle_block.1", ctx, h, context, num_video_frames); // [N, 4*model_channels, h/8, w/8]
h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames); // [N, 4*model_channels, h/8, w/8] h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames); // [N, 4*model_channels, h/8, w/8]
} }
h = apply_motion_mid(h);
} }
sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h"); sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h");
if (controls.size() > 0) { if (controls.size() > 0) {
@ -660,6 +714,8 @@ public:
up_sample_idx++; up_sample_idx++;
} }
h = apply_motion_output(output_block_idx, h);
if (i > 0 && j == num_res_blocks) { if (i > 0 && j == num_res_blocks) {
if (tiny_unet) { if (tiny_unet) {
output_block_idx++; output_block_idx++;

View File

@ -648,8 +648,8 @@ namespace ZImage {
*diffusion_params.x, *diffusion_params.x,
*diffusion_params.timesteps, *diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context), 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,
diffusion_params.ref_index_mode); diffusion_params.ref_image_params.ref_index_mode);
} }
void test() { void test() {

View File

@ -100,7 +100,8 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) {
// https://huggingface.co/docs/safetensors/index // https://huggingface.co/docs/safetensors/index
bool read_safetensors_file(const std::string& file_path, bool read_safetensors_file(const std::string& file_path,
std::vector<TensorStorage>& tensor_storages, std::vector<TensorStorage>& tensor_storages,
std::string* error) { std::string* error,
std::map<std::string, std::string>* metadata) {
std::ifstream file(file_path, std::ios::binary); std::ifstream file(file_path, std::ios::binary);
if (!file.is_open()) { if (!file.is_open()) {
set_error(error, "failed to open '" + file_path + "'"); set_error(error, "failed to open '" + file_path + "'");
@ -150,6 +151,18 @@ bool read_safetensors_file(const std::string& file_path,
return false; return false;
} }
if (metadata != nullptr) {
metadata->clear();
auto metadata_item = header_.find("__metadata__");
if (metadata_item != header_.end() && metadata_item->is_object()) {
for (const auto& item : metadata_item->items()) {
if (item.value().is_string()) {
metadata->emplace(item.key(), item.value().get<std::string>());
}
}
}
}
tensor_storages.clear(); tensor_storages.clear();
for (auto& item : header_.items()) { for (auto& item : header_.items()) {
std::string name = item.key(); std::string name = item.key();

View File

@ -1,6 +1,7 @@
#ifndef __SD_MODEL_IO_SAFETENSORS_IO_H__ #ifndef __SD_MODEL_IO_SAFETENSORS_IO_H__
#define __SD_MODEL_IO_SAFETENSORS_IO_H__ #define __SD_MODEL_IO_SAFETENSORS_IO_H__
#include <map>
#include <string> #include <string>
#include <vector> #include <vector>
@ -10,7 +11,8 @@
bool is_safetensors_file(const std::string& file_path); bool is_safetensors_file(const std::string& file_path);
bool read_safetensors_file(const std::string& file_path, bool read_safetensors_file(const std::string& file_path,
std::vector<TensorStorage>& tensor_storages, std::vector<TensorStorage>& tensor_storages,
std::string* error = nullptr); std::string* error = nullptr,
std::map<std::string, std::string>* metadata = nullptr);
bool read_safetensors_index_file(const std::string& file_path, bool read_safetensors_index_file(const std::string& file_path,
std::vector<std::string>& shard_paths, std::vector<std::string>& shard_paths,
std::string* error = nullptr); std::string* error = nullptr);

View File

@ -317,7 +317,7 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
std::vector<TensorStorage> tensor_storages; std::vector<TensorStorage> tensor_storages;
std::string error; std::string error;
if (!read_safetensors_file(file_path, tensor_storages, &error)) { if (!read_safetensors_file(file_path, tensor_storages, &error, &metadata_)) {
LOG_ERROR("%s", error.c_str()); LOG_ERROR("%s", error.c_str());
return false; return false;
} }

View File

@ -36,6 +36,7 @@ protected:
std::vector<ModelFileData> file_data; std::vector<ModelFileData> file_data;
bool model_files_processed = false; bool model_files_processed = false;
String2TensorStorage tensor_storage_map; String2TensorStorage tensor_storage_map;
std::map<std::string, std::string> metadata_;
int n_threads_; int n_threads_;
size_t add_file_path(const std::string& file_path); size_t add_file_path(const std::string& file_path);
@ -63,6 +64,7 @@ public:
std::map<ggml_type, uint32_t> get_vae_wtype_stat(); std::map<ggml_type, uint32_t> get_vae_wtype_stat();
String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; } String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
const String2TensorStorage& get_tensor_storage_map() const { return tensor_storage_map; } const String2TensorStorage& get_tensor_storage_map() const { return tensor_storage_map; }
const std::map<std::string, std::string>& get_metadata() const { return metadata_; }
void set_n_threads(int n_threads); void set_n_threads(int n_threads);
void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = ""); void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = "");
void process_model_files(bool enable_mmap = false, bool writable_mmap = true); void process_model_files(bool enable_mmap = false, bool writable_mmap = true);

View File

@ -304,6 +304,12 @@ std::string convert_diffusers_unet_to_original_sd1(std::string name) {
} }
} }
static const std::vector<std::pair<std::string, std::string>> name_map{
{"to_out.weight", "to_out.0.weight"},
{"to_out.bias", "to_out.0.bias"},
};
replace_with_name_map(result, name_map);
return result; return result;
} }
@ -1102,6 +1108,8 @@ std::string convert_sep_to_dot(std::string name) {
"norm1_context", "norm1_context",
"ff_context", "ff_context",
"x_embedder", "x_embedder",
"cross_attn",
"output_proj",
}; };
// record the positions of underscores that should NOT be replaced // record the positions of underscores that should NOT be replaced
@ -1341,13 +1349,25 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
// diffusion model // diffusion model
{ {
bool matched = false;
for (const auto& prefix : diffuison_model_prefix_vec) { for (const auto& prefix : diffuison_model_prefix_vec) {
if (starts_with(name, prefix)) { if (starts_with(name, prefix)) {
name = convert_diffusion_model_name(name.substr(prefix.size()), prefix, version); name = convert_diffusion_model_name(name.substr(prefix.size()), prefix, version);
name = prefix + name; name = prefix + name;
matched = true;
break; break;
} }
} }
if (is_lora && !matched && !diffuison_model_prefix_vec.empty()) {
if (starts_with(name, "down_blocks.") || starts_with(name, "up_blocks.") ||
starts_with(name, "mid_block.") || starts_with(name, "conv_in.") ||
starts_with(name, "conv_out.") || starts_with(name, "time_embedding.") ||
starts_with(name, "conv_norm_out.")) {
const std::string& canonical_prefix = diffuison_model_prefix_vec.front();
name = convert_diffusion_model_name(name, canonical_prefix, version);
name = canonical_prefix + name;
}
}
} }
// cond_stage_model // cond_stage_model

View File

@ -24,6 +24,7 @@
#include "extensions/generation_extension.h" #include "extensions/generation_extension.h"
#include "model/adapter/lora.hpp" #include "model/adapter/lora.hpp"
#include "model/diffusion/anima.hpp" #include "model/diffusion/anima.hpp"
#include "model/diffusion/animatediff.hpp"
#include "model/diffusion/boogu.hpp" #include "model/diffusion/boogu.hpp"
#include "model/diffusion/control.hpp" #include "model/diffusion/control.hpp"
#include "model/diffusion/ernie_image.hpp" #include "model/diffusion/ernie_image.hpp"
@ -63,6 +64,10 @@
const char* sd_vae_format_name(enum sd_vae_format_t format); const char* sd_vae_format_name(enum sd_vae_format_t format);
static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback); static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback);
static bool sd_version_supports_animatediff(SDVersion version) {
return version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX;
}
const char* model_version_to_str[] = { const char* model_version_to_str[] = {
"SD 1.x", "SD 1.x",
"SD 1.x Inpaint", "SD 1.x Inpaint",
@ -212,6 +217,8 @@ public:
std::vector<std::shared_ptr<GenerationExtension>> generation_extensions; std::vector<std::shared_ptr<GenerationExtension>> generation_extensions;
std::vector<std::shared_ptr<LoraModel>> runtime_lora_models; std::vector<std::shared_ptr<LoraModel>> runtime_lora_models;
bool apply_lora_immediately = false; bool apply_lora_immediately = false;
bool animatediff_loaded = false;
int animatediff_num_frames = 0;
std::string taesd_path; std::string taesd_path;
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr}; sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr};
@ -808,6 +815,16 @@ public:
} }
} }
if (strlen(SAFE_STR(sd_ctx_params->motion_module_path)) > 0) {
LOG_INFO("loading motion module (AnimateDiff) from '%s'", sd_ctx_params->motion_module_path);
if (!model_loader.init_from_file(sd_ctx_params->motion_module_path,
"model.diffusion_model.motion_module.")) {
LOG_WARN("loading motion module from '%s' failed", sd_ctx_params->motion_module_path);
} else {
animatediff_loaded = true;
}
}
if (strlen(SAFE_STR(sd_ctx_params->control_net_path)) > 0) { if (strlen(SAFE_STR(sd_ctx_params->control_net_path)) > 0) {
if (!model_loader.init_from_file(sd_ctx_params->control_net_path)) { if (!model_loader.init_from_file(sd_ctx_params->control_net_path)) {
LOG_ERROR("init control net model loader from file failed: '%s'", sd_ctx_params->control_net_path); LOG_ERROR("init control net model loader from file failed: '%s'", sd_ctx_params->control_net_path);
@ -2297,7 +2314,7 @@ public:
const char* extra_sample_args, const char* extra_sample_args,
const std::vector<float>& sigmas, const std::vector<float>& sigmas,
const std::vector<sd::Tensor<float>>& ref_latents, 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>& denoise_mask,
const sd::Tensor<float>& vace_context, const sd::Tensor<float>& vace_context,
float vace_strength, float vace_strength,
@ -2458,9 +2475,9 @@ public:
sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma); sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma);
std::vector<sd::Tensor<float>> controls; std::vector<sd::Tensor<float>> controls;
DiffusionParams diffusion_params; DiffusionParams diffusion_params;
diffusion_params.x = &noised_input; diffusion_params.x = &noised_input;
diffusion_params.timesteps = &timesteps_tensor; diffusion_params.timesteps = &timesteps_tensor;
diffusion_params.ref_index_mode = Rope::ref_index_mode_from_bool(increase_ref_index); diffusion_params.ref_image_params = ref_image_params;
sd::guidance::GuidanceInput step_guidance_input; sd::guidance::GuidanceInput step_guidance_input;
step_guidance_input.step = step; step_guidance_input.step = step;
step_guidance_input.schedule_size = sigmas.size(); step_guidance_input.schedule_size = sigmas.size();
@ -2487,7 +2504,11 @@ public:
diffusion_params.ref_latents = ref_latents_override != nullptr ? ref_latents_override : (condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images); diffusion_params.ref_latents = ref_latents_override != nullptr ? ref_latents_override : (condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images);
if (sd_version_is_unet(version)) { if (sd_version_is_unet(version)) {
diffusion_params.extra = UNetDiffusionExtra{-1, &controls, control_strength}; int nvf = -1;
if (animatediff_loaded && noised_input.dim() >= 4 && noised_input.shape()[3] > 1) {
nvf = static_cast<int>(noised_input.shape()[3]);
}
diffusion_params.extra = UNetDiffusionExtra{nvf, &controls, control_strength};
} else if (sd_version_is_sd3(version)) { } else if (sd_version_is_sd3(version)) {
diffusion_params.extra = SkipLayerDiffusionExtra{local_skip_layers}; diffusion_params.extra = SkipLayerDiffusionExtra{local_skip_layers};
} else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) { } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) {
@ -2840,6 +2861,126 @@ public:
auto flow_denoiser = std::dynamic_pointer_cast<DiscreteFlowDenoiser>(denoiser); auto flow_denoiser = std::dynamic_pointer_cast<DiscreteFlowDenoiser>(denoiser);
return !!flow_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 ==================================================*/ /*================================================= SD API ==================================================*/
@ -3301,6 +3442,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_sample_params_init(&sd_img_gen_params->sample_params);
sd_img_gen_params->clip_skip = -1; sd_img_gen_params->clip_skip = -1;
sd_img_gen_params->ref_images_count = 0; 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->width = 512;
sd_img_gen_params->height = 512; sd_img_gen_params->height = 512;
sd_img_gen_params->strength = 0.75f; sd_img_gen_params->strength = 0.75f;
@ -3338,8 +3480,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
"batch_count: %d\n" "batch_count: %d\n"
"qwen_image_layers: %d\n" "qwen_image_layers: %d\n"
"ref_images_count: %d\n" "ref_images_count: %d\n"
"auto_resize_ref_image: %s\n" "ref_image_args: %s\n"
"increase_ref_index: %s\n"
"control_strength: %.2f\n" "control_strength: %.2f\n"
"photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\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" "VAE tiling: %s (temporal=%s, extra_tiling_args=%s)\n"
@ -3357,8 +3498,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->batch_count,
sd_img_gen_params->qwen_image_layers, sd_img_gen_params->qwen_image_layers,
sd_img_gen_params->ref_images_count, sd_img_gen_params->ref_images_count,
BOOL_STR(sd_img_gen_params->auto_resize_ref_image), SAFE_STR(sd_img_gen_params->ref_image_args),
BOOL_STR(sd_img_gen_params->increase_ref_index),
sd_img_gen_params->control_strength, sd_img_gen_params->control_strength,
sd_img_gen_params->pm_params.style_strength, sd_img_gen_params->pm_params.style_strength,
sd_img_gen_params->pm_params.id_images_count, sd_img_gen_params->pm_params.id_images_count,
@ -3524,6 +3664,9 @@ SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) {
if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { if (sd_ctx == nullptr || sd_ctx->sd == nullptr) {
return false; return false;
} }
if (sd_ctx->sd->animatediff_loaded && sd_version_supports_animatediff(sd_ctx->sd->version)) {
return true;
}
return sd_version_supports_video_generation(sd_ctx->sd->version); return sd_version_supports_video_generation(sd_ctx->sd->version);
} }
@ -3655,8 +3798,6 @@ struct GenerationRequest {
float strength = 1.f; float strength = 1.f;
float control_strength = 0.f; float control_strength = 0.f;
float eta = 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 guidance = {};
sd_guidance_params_t high_noise_guidance = {}; sd_guidance_params_t high_noise_guidance = {};
sd_pm_params_t pm_params = {}; sd_pm_params_t pm_params = {};
@ -3682,8 +3823,6 @@ struct GenerationRequest {
strength = sd_img_gen_params->strength; strength = sd_img_gen_params->strength;
control_strength = sd_img_gen_params->control_strength; control_strength = sd_img_gen_params->control_strength;
eta = sd_img_gen_params->sample_params.eta; 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; has_ref_images = sd_img_gen_params->ref_images_count > 0;
guidance = sd_img_gen_params->sample_params.guidance; guidance = sd_img_gen_params->sample_params.guidance;
pm_params = sd_img_gen_params->pm_params; pm_params = sd_img_gen_params->pm_params;
@ -4431,7 +4570,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, static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd_ctx_t* sd_ctx,
const sd_img_gen_params_t* sd_img_gen_params, const sd_img_gen_params_t* sd_img_gen_params,
GenerationRequest* request, GenerationRequest* request,
SamplePlan* plan) { SamplePlan* plan,
const RefImageParams& ref_image_params) {
int64_t prepare_start_ms = ggml_time_ms(); int64_t prepare_start_ms = ggml_time_ms();
sd::Tensor<float> init_image_tensor; sd::Tensor<float> init_image_tensor;
@ -4543,6 +4683,14 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
} }
} }
if (sd_ctx->sd->animatediff_num_frames > 1 &&
init_latent.dim() >= 4 && init_latent.shape()[3] == 1) {
int n_frames = sd_ctx->sd->animatediff_num_frames;
std::vector<int64_t> shape(init_latent.shape().begin(), init_latent.shape().end());
shape[3] = n_frames;
init_latent = sd::Tensor<float>(std::move(shape)); // zero-filled batch of N frames; per-frame noise is generated later via randn_like.
}
if (!control_image_tensor.empty()) { if (!control_image_tensor.empty()) {
control_latent = sd_ctx->sd->encode_first_stage(control_image_tensor); control_latent = sd_ctx->sd->encode_first_stage(control_image_tensor);
if (control_latent.empty()) { if (control_latent.empty()) {
@ -4574,9 +4722,10 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
continue; continue;
} }
sd::Tensor<float> ref_latent; 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"); 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_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]; double vae_height = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0];
@ -4703,15 +4852,20 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
const sd_img_gen_params_t* sd_img_gen_params, const sd_img_gen_params_t* sd_img_gen_params,
GenerationRequest* request, GenerationRequest* request,
SamplePlan* plan, SamplePlan* plan,
ImageGenerationLatents* latents) { ImageGenerationLatents* latents,
const RefImageParams& ref_image_params) {
ConditionerRunnerDoneOnExit conditioner_runner_done{sd_ctx->sd->cond_stage_model.get()}; ConditionerRunnerDoneOnExit conditioner_runner_done{sd_ctx->sd->cond_stage_model.get()};
ConditionerParams condition_params; ConditionerParams condition_params;
condition_params.text = request->prompt; condition_params.text = request->prompt;
condition_params.clip_skip = request->clip_skip; condition_params.clip_skip = request->clip_skip;
condition_params.width = request->width; condition_params.width = request->width;
condition_params.height = request->height; condition_params.height = request->height;
condition_params.ref_images = &latents->ref_images; 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, sd_ctx->sd->prepare_generation_extensions(request->pm_params,
request->pulid_params, request->pulid_params,
@ -4721,7 +4875,7 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
condition_params.zero_out_masked = false; condition_params.zero_out_masked = false;
auto cond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, auto cond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
condition_params); 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 cond.c_concat = latents->concat_latent; // TODO: optimize
} }
@ -4750,7 +4904,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, uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
condition_params); 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 uncond.c_concat = latents->concat_latent; // TODO: optimize
} }
} }
@ -4774,7 +4928,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, img_uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
condition_params); 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 img_uncond.c_concat = latents->img_uncond_concat_latent; // TODO: optimize
} }
} }
@ -4845,6 +4999,24 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
if (cancelled) { if (cancelled) {
break; break;
} }
} else if (sd_ctx->sd->animatediff_num_frames > 1 &&
final_latents[i].dim() >= 4 &&
final_latents[i].shape()[3] == sd_ctx->sd->animatediff_num_frames) {
int n_frames = sd_ctx->sd->animatediff_num_frames;
for (int f = 0; f < n_frames; ++f) {
if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
LOG_ERROR("cancelling latent decodings");
cancelled = true;
break;
}
sd::Tensor<float> frame_latent = sd::ops::slice(final_latents[i], 3, f, f + 1);
sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(frame_latent);
if (image.empty()) {
LOG_ERROR("decode_first_stage failed for AnimateDiff frame %d/%d", f + 1, n_frames);
return nullptr;
}
decoded_images.push_back(std::move(image));
}
} else { } else {
sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(final_latents[i]); sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(final_latents[i]);
if (image.empty()) { if (image.empty()) {
@ -5091,13 +5263,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); 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); 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); ImageVaeAxesGuard axes_guard(sd_ctx, sd_img_gen_params, request);
SamplePlan plan(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, auto latents_opt = prepare_image_generation_latents(sd_ctx,
sd_img_gen_params, sd_img_gen_params,
&request, &request,
&plan); &plan,
ref_image_params);
if (!latents_opt.has_value()) { if (!latents_opt.has_value()) {
return false; return false;
} }
@ -5107,7 +5282,8 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
sd_img_gen_params, sd_img_gen_params,
&request, &request,
&plan, &plan,
&latents); &latents,
ref_image_params);
if (!embeds_opt.has_value()) { if (!embeds_opt.has_value()) {
return false; return false;
} }
@ -5153,7 +5329,7 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
plan.extra_sample_args, plan.extra_sample_args,
plan.sigmas, plan.sigmas,
latents.ref_latents, latents.ref_latents,
request.increase_ref_index, ref_image_params,
latents.denoise_mask, latents.denoise_mask,
sd::Tensor<float>(), sd::Tensor<float>(),
1.f, 1.f,
@ -5274,7 +5450,7 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
plan.extra_sample_args, plan.extra_sample_args,
hires_sigma_sched, hires_sigma_sched,
latents.ref_latents, latents.ref_latents,
request.increase_ref_index, ref_image_params,
hires_denoise_mask, hires_denoise_mask,
sd::Tensor<float>(), sd::Tensor<float>(),
1.f, 1.f,
@ -5665,6 +5841,9 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx,
condition_params.text = request.prompt; condition_params.text = request.prompt;
condition_params.zero_out_masked = true; condition_params.zero_out_masked = true;
condition_params.ref_images = &latents.ref_images; 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(); int64_t prepare_start_ms = ggml_time_ms();
embeds.cond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, embeds.cond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
@ -5928,6 +6107,47 @@ static bool apply_ltxv_refine_image_conditioning(sd_ctx_t* sd_ctx,
return true; return true;
} }
static bool generate_animatediff_video(sd_ctx_t* sd_ctx,
const sd_vid_gen_params_t* sd_vid_gen_params,
sd_image_t** frames_out,
int* num_frames_out) {
int n_frames = sd_vid_gen_params->video_frames;
if (n_frames < 1) {
LOG_ERROR("AnimateDiff: --video-frames must be >= 1");
return false;
}
if (n_frames > 32) {
LOG_WARN("AnimateDiff motion modules have a 32-frame positional-encoding context; capping to 32");
n_frames = 32;
}
sd_img_gen_params_t img_gen_params;
sd_img_gen_params_init(&img_gen_params);
img_gen_params.loras = sd_vid_gen_params->loras;
img_gen_params.lora_count = sd_vid_gen_params->lora_count;
img_gen_params.prompt = sd_vid_gen_params->prompt;
img_gen_params.negative_prompt = sd_vid_gen_params->negative_prompt;
img_gen_params.clip_skip = sd_vid_gen_params->clip_skip;
img_gen_params.width = sd_vid_gen_params->width;
img_gen_params.height = sd_vid_gen_params->height;
img_gen_params.sample_params = sd_vid_gen_params->sample_params;
img_gen_params.strength = sd_vid_gen_params->strength;
img_gen_params.seed = sd_vid_gen_params->seed;
img_gen_params.batch_count = 1;
img_gen_params.control_strength = 1.0f;
img_gen_params.vae_tiling_params = sd_vid_gen_params->vae_tiling_params;
img_gen_params.cache = sd_vid_gen_params->cache;
img_gen_params.hires = sd_vid_gen_params->hires;
img_gen_params.qwen_image_layers = 0;
img_gen_params.circular_x = sd_vid_gen_params->circular_x;
img_gen_params.circular_y = sd_vid_gen_params->circular_y;
sd_ctx->sd->animatediff_num_frames = n_frames;
bool ok = generate_image(sd_ctx, &img_gen_params, frames_out, num_frames_out);
sd_ctx->sd->animatediff_num_frames = 0;
return ok;
}
SD_API bool generate_video(sd_ctx_t* sd_ctx, SD_API bool generate_video(sd_ctx_t* sd_ctx,
const sd_vid_gen_params_t* sd_vid_gen_params, const sd_vid_gen_params_t* sd_vid_gen_params,
sd_image_t** frames_out, sd_image_t** frames_out,
@ -5942,12 +6162,20 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
if (audio_out != nullptr) { if (audio_out != nullptr) {
*audio_out = nullptr; *audio_out = nullptr;
} }
sd_ctx->sd->reset_cancel_flag();
if (num_frames_out != nullptr) { if (num_frames_out != nullptr) {
*num_frames_out = 0; *num_frames_out = 0;
} }
if (sd_ctx->sd->animatediff_loaded && sd_version_supports_animatediff(sd_ctx->sd->version)) {
LOG_INFO("AnimateDiff dispatch: %d frames, %dx%d",
sd_vid_gen_params->video_frames, sd_vid_gen_params->width, sd_vid_gen_params->height);
return generate_animatediff_video(sd_ctx, sd_vid_gen_params, frames_out, num_frames_out);
}
sd_ctx->sd->reset_cancel_flag();
const RefImageParams ref_image_params;
int64_t t0 = ggml_time_ms(); int64_t t0 = ggml_time_ms();
sd_ctx->sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params; sd_ctx->sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params;
apply_circular_axes_to_diffusion(sd_ctx, sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y); apply_circular_axes_to_diffusion(sd_ctx, sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y);
@ -6034,7 +6262,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
plan.high_noise_extra_sample_args, plan.high_noise_extra_sample_args,
high_noise_sigmas, high_noise_sigmas,
std::vector<sd::Tensor<float>>{}, std::vector<sd::Tensor<float>>{},
false, ref_image_params,
latents.denoise_mask, latents.denoise_mask,
latents.vace_context, latents.vace_context,
request.vace_strength, request.vace_strength,
@ -6076,7 +6304,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
plan.extra_sample_args, plan.extra_sample_args,
plan.sigmas, plan.sigmas,
std::vector<sd::Tensor<float>>{}, std::vector<sd::Tensor<float>>{},
false, ref_image_params,
latents.denoise_mask, latents.denoise_mask,
latents.vace_context, latents.vace_context,
request.vace_strength, request.vace_strength,
@ -6214,7 +6442,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
plan.extra_sample_args, plan.extra_sample_args,
hires_sigma_sched, hires_sigma_sched,
std::vector<sd::Tensor<float>>{}, std::vector<sd::Tensor<float>>{},
false, ref_image_params,
hires_denoise_mask, hires_denoise_mask,
sd::Tensor<float>(), sd::Tensor<float>(),
hires_request.vace_strength, hires_request.vace_strength,