diff --git a/docs/edit.md b/docs/edit.md index 9f33862d2..1312854b7 100644 --- a/docs/edit.md +++ b/docs/edit.md @@ -85,6 +85,9 @@ If presets are insufficient, you can manually configure the following parameters | `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 | +| `ref_image_size` | string | MiniMax-H3 reference image sizing before VAE encoding. `match` scales each reference down to the generation's pixel area; `max` uses a 2048 px short edge. Only references above the limit are downscaled; each axis is then snapped to the nearest 32 px multiple (an axis of at least 16 px grows by at most 16 px, and references are never narrower than 32 px). | `match`, `max` | + +`ref_image_size` applies to `vid_gen` with MiniMax-H3 Ref2VA. Reference latents are attended to at every sampling step, so `max` can be several times slower. ### Preset Default Values diff --git a/docs/generation_metadata.md b/docs/generation_metadata.md new file mode 100644 index 000000000..1f7fc3338 --- /dev/null +++ b/docs/generation_metadata.md @@ -0,0 +1,165 @@ +# Generation results and metadata + +`generate_image()`, `generate_video()`, `adetail_image()` and `upscale()` return +owned results containing the media and its generation record. This is a breaking +C API change: applications and bindings must update their calls and rebuild. + +```c +sd_img_gen_params_t params; +sd_img_gen_params_init(¶ms); +params.prompt = "a small house"; + +sd_image_generation_result_t* result = NULL; +if (generate_image(ctx, ¶ms, &result)) { + for (int i = 0; i < result->count; ++i) { + const sd_image_result_t* item = &result->items[i]; + /* Consume item->image and item->generation_info together. */ + } + free_sd_image_generation_result(result); +} +``` + +`result_out` is required and is set to `NULL` on failure. A successful result +owns its pixels and UTF-8 JSON strings. The video result additionally owns its +audio and contains the effective `fps` and `frame_count`. Results remain valid +after freeing the generation context. Use the matching library result destructor; +do not separately free fields unless transferring ownership and clearing them. +Both result destructors accept `NULL`. + +## Reproduction parameters + +The `sdcpp.generation/v1` JSON is a snapshot of generation parameters for the +same loaded context and the same input images, masks, reference media and other +assets. It is created after `GenerationRequest` resolves the request, before +preprocessing or sampling. A failed call returns no result or metadata. + +Fields follow the C generation parameter structs: `prompt`, `negative_prompt`, +`width`, `height`, `seed`, `sample_params`, `hires`, `cache`, +`vae_tiling_params`, `image_preprocess`, and the image/video-specific options. +`kind` distinguishes image and video parameters. There is no inferred `task`. +Enums use their API string names. Seeds are decimal strings to preserve their +full integer range. Omitted optional fields retain the parameter initializer's +defaults, including automatic values represented by infinity in the C API. + +Disabled optional features are omitted entirely: `cache`, `hires`, VAE tiling, +zero/nonpositive `shifted_timestep`, inactive SLG, false circular axes, and empty +preprocessing/reference arguments. SLG layers remain present when `slg_uncond` +is enabled even with zero SLG scale. Enabled caches store their mode's parameters; +TaylorSeer retains the shared DBCache thresholds and limits used by the runtime. +Image CFG at its neutral value of 1 is omitted and resolves back to 1 by default. + +Image strength, ControlNet/IP-Adapter strength and PhotoMaker settings are stored +when their corresponding input images are supplied. PuLID requires an embedding +path. Layer count is stored for layered generation; video VACE and high-noise +settings require those models. High-noise settings omit scheduler, custom sigmas, +flow shift and timestep shifting because that pass shares the base configuration. +Explicit nonempty argument strings are retained without rewriting their contents. +Core sampling values such as seed, sampler, scheduler, steps and guidance remain +explicit even when equal to their defaults. ADetailer always retains its inpaint +strength because its input crops are created later. +Model-required VAE tiling is resolved before serialization, so forced tiling +(such as MiniMax-H3) remains present even if the caller did not enable it. + +The request fixes the seed, dimensions, clip skip, guidance, sampler, scheduler +and eta before taking the snapshot. `custom_sigmas` is stored only when supplied +as an input. Automatically computed schedules, preprocessing transforms, +conditioning/cache decisions, detector regions and output statistics are not +stored: they are recomputed from the same parameters and inputs. + +Each image records the seed used to sample its latent. `batch_count` is omitted; +replay image parameters with `batch_count = 1`. Layers or frames decoded from +the same latent share its seed, so layered generation reproduces that group of +layers. Video frames share the video request. Partial batch cancellation returns +metadata only for completed outputs. Full cancellation returns no result. + +The current image pipeline prepares input latents once using the starting seed +and shares them across a batch. With a stochastic VAE encoder, replaying a later +image alone using its sampling seed can change those input latents. Per-image +seeds therefore do not guarantee independent reproduction for these paths. + +To replay a record, initialize the matching C parameter struct, populate its +fields from the JSON, restore the same input buffers, and call generation with +the same context. Keep custom-sigma/SLG arrays and string storage alive for the +call. Paths supplied as generation parameters (LoRAs, hires models and identity +embeddings) are retained. `models` identifies successfully loaded context model +files by component (`model`, `diffusion_model`, `clip_l`, `vae`, etc.), using +filenames with extensions and omitting directories. Optional components disabled +for the request and preview-only TAE are omitted. WebUI text includes `Model` +(the diffusion model when supplied separately, otherwise the checkpoint, without +the extension) and `VAE` for a separate active VAE. These names are identifiers, +not hashes or sufficient configuration to rebuild a context; replay still requires +the same model files and context settings. The JSON does not embed media buffers +or model weights. + +When `--lora-model-dir` is configured, CLI/server output records express LoRA +paths relative to that directory, preserving subdirectories and file extensions +and using `/` separators. This applies to file metadata, API responses and +ADetailer records. Resolve these paths against the same LoRA directory before +passing them back to the C API. Core results retain the supplied loading paths. +Files outside the directory retain `..` components; paths on another volume, +or records without a configured LoRA directory, retain their original paths. + +## Postprocessing + +Upscale records `upscale.upscale_factor`. ADetailer records `adetailer` and +`inpaint_params`, with a concrete starting seed chosen before detection. These +operations require the same upscaler/detector context and source image on replay. +They do not store detection results or per-region generation traces. + +The CLI/server append operation parameter snapshots to `postprocess` in call +order. The source generation parameters remain unchanged. Frame postprocessing +uses the same operation parameters for each video frame. Standalone +postprocessing does not invent source diffusion parameters for an input image. + +## File metadata + +Image encoders accept the native record and derive WebUI infotext from it. +Images use the existing metadata locations and combined text layout: +`, SDCPP: `. Standalone postprocessing without +source diffusion parameters stores only native JSON. Video metadata stores +only native JSON. + +| Container | Metadata location | +| --- | --- | +| PNG | `tEXt`, keyword `parameters` | +| JPEG | COM segment, `parameters` label followed by a NUL byte | +| WebP image / animated WebP | XMP property `sdcpp:parameters` | +| AVI | RIFF INFO `ICMT` | +| WebM | `COMMENT` tag | + +Inspect images or videos with `sd-cli -M metadata --image `; the same +command supports all containers above, with `--metadata-format json` for a +structured report. Video metadata inspection does not require a model or video +decoder. +Matroska containers are also supported, including `.webm` outputs whose audio +codec causes the muxer to select the Matroska document type. + +The XMP namespace is `https://github.com/leejet/stable-diffusion.cpp/ns/1.0/`. +PNG and JPEG retain the existing UTF-8 text payload convention. Memory encoding +and file saving use the same implementation. The core result and API +`generation_info` fields remain pure JSON; the combined layout applies only +to embedded image metadata. + +WebUI conversion is lossy: known samplers and schedules have explicit name +mappings; unsupported ones use `SDCPP Sampler`/`SDCPP Schedule` fields. Hires +keeps the base request size and separate configured hires size/steps. The native +record retains the full generation parameters. Standalone postprocessing has no WebUI +infotext because its original diffusion settings are unknown. + +WebUI prompt text appends ordinary LoRAs as ``. +Directory components are omitted to match WebUI's filename lookup; the native +`loras` array retains relative paths and weights. High-noise-only LoRAs and names +that cannot be represented by this tag syntax remain in the native record only. +The native prompt is not modified, and postprocessing LoRAs are not added to the +source generation prompt. + +JPEG requires the COM payload, including the label and terminating NUL bytes, +to fit within 65,533 bytes. Larger records fail encoding with an error rather +than being truncated. Use PNG or WebP for larger records. + +Disabling file metadata omits both image representations, but does not disable +the core generation record or its API response. Native async image results and +OpenAI image response items contain `generation_info`; video results contain a +sequence-level `generation_info`. The SD API `info` object contains per-output +`generation_info` alongside derived `infotexts` and each image's seed. Its outer +`parameters` field continues to echo the request. diff --git a/docs/image_preprocessing.md b/docs/image_preprocessing.md index 719781750..8ec236e71 100644 --- a/docs/image_preprocessing.md +++ b/docs/image_preprocessing.md @@ -128,7 +128,8 @@ To disable reference resizing before VAE encoding, use `--ref-image-args "resize_before_vae=false"` or the server field `"ref_image_args": "resize_before_vae=false"`. This is separate from `target=ref,mode=none`, which only skips input geometry. Model constraints -still apply. +still apply; MiniMax-H3 requires 32-pixel-aligned references and warns instead +of honoring the switch. ## Server requests @@ -148,12 +149,12 @@ an array of rule strings: OpenAI-compatible requests accept it through `{...}` in the prompt. Request rules replace server-default rules. Generation metadata records the -user rules; image encodings and channel conventions are unchanged. +user rules for replay with the same input images; image encodings and channel conventions are unchanged. ## C API Set `image_preprocess` on the existing image/video generation parameters. -The `generate_image()` and `generate_video()` signatures are unchanged: +Generation returns media and actual parameters together; see [generation metadata](generation_metadata.md): ```c sd_img_gen_params_t params; @@ -161,7 +162,10 @@ sd_img_gen_params_init(¶ms); /* Set prompt, original-resolution input images, and generation options. */ params.image_preprocess.rules = "target=init,mode=crop-resize,filter=lanczos;" "target=mask,filter=nearest-exact"; -bool ok = generate_image(ctx, ¶ms, &images, &count); +sd_image_generation_result_t* result = NULL; +bool ok = generate_image(ctx, ¶ms, &result); +/* Consume result->items and their generation_info on success. */ +free_sd_image_generation_result(result); ``` Both generation parameter initializers set `image_preprocess.rules` to `NULL`, diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index 64eb78f53..0c17ba0e0 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -80,8 +80,20 @@ latents condition the diffusion transformer. Paired video and audio references share the same timeline. Ref2VA cannot be combined with `--init-img` or `--end-img` in one request. -Reference images keep their aspect ratio and are only downscaled when their -pixel area exceeds the requested generation canvas. +Reference images keep their aspect ratio and are only downscaled when they +exceed the configured limit; neither mode scales a reference up. The limit is +selected by `ref_image_size`. The default, `match`, uses the pixel area of the +requested generation canvas. Pass `--ref-image-args "ref_image_size=max"` (or +`ref_image_size` in `sd_vid_gen_params_t.ref_image_args`) to use a 2048 px short +edge instead, which preserves more identity detail at a higher sampling cost: +reference latents are attended at every sampling step. + +Whichever limit applies, each axis is then snapped to the nearest 32-pixel +multiple. An axis of at least 16 pixels can therefore grow by up to 16 pixels +past the scaled size (a 48-pixel axis becomes 64), and references are never +narrower than 32 pixels. +`resize_before_vae=false` is not supported here; MiniMax-H3 keeps references on +the 32-pixel grid and reports a warning. The C API exposes the same inputs through `ref_images`, `ref_videos`, and `ref_audios` in `sd_vid_gen_params_t`. Each `sd_ref_video_t` supplies its own diff --git a/docs/model_config.md b/docs/model_config.md index dda93a837..0d9290c1d 100644 --- a/docs/model_config.md +++ b/docs/model_config.md @@ -1,7 +1,7 @@ # Model Configuration Conventions -This document describes the conventions for model configuration structs and -weight-based configuration detection. +This document describes model configuration structs, weight-based configuration +detection, and the definitions used to integrate models into the runtime. ## Config Types @@ -116,3 +116,113 @@ struct. If a model has no weight-derived config today, it may still provide `detect_from_weights` for API consistency, but it should not print a config detection log unless it actually derives values from weights. + +## Model Integration Definitions + +Every `SDVersion` has an explicit `ModelDefinition` in +`src/model/integrations/`. Related variants share an integration file and runner +factory. Network blocks and their weight-derived configuration remain under +`src/model/diffusion/`; reusable text encoder implementations remain under +`src/model/te/`. Weight-name conversion remains in `src/name_conversion.cpp`. + +To add a model: + +1. Add its identity to `SDVersion` without renumbering existing values. +2. Implement the network and `*Config::detect_from_weights` following the + conventions above, or reuse an existing runner. +3. Add a definition in its integration file, including its component + factory, traits, geometry, and sampling policy. +4. Add the family's explicit registration in `src/model_registry.cpp`. For a new + variant of a registered family, add its definition to that family's list. +5. Add version detection signatures and any required weight-name mappings in + `src/model_loader.cpp` and `src/name_conversion.cpp` respectively. Update shared + conditioning or runtime code only when existing behavior cannot serve the model. + +Use the existing family integration as the starting point for a variant. For a +new family, `src/model/integrations/moss_image.cpp` is a compact example of a +conditioner/diffusion factory and its definition. Definitions collect model +construction and defaults; they do not replace version detection, name conversion, +or model-specific inference algorithms. + +The registry checks for missing, duplicate, and invalid version registrations. +Explicit references keep integrations reachable in static-library builds. The +existing CMake glob includes integration sources automatically. No legacy +builder, display-name table, or prediction-default fallback needs updating. + +### Detection And Weight Conversion + +Version detection remains in `ModelLoader::get_sd_version()` in +`src/model_loader.cpp`. Add model signatures there and preserve the order of +overlapping checks. Detection uses tensor names and metadata, and can run before +and after name conversion, so signatures must work at both stages. Do not load +tensor data to identify an architecture. The registry looks up definitions by +the detected version; it does not perform detection. + +Weight-name mappings and their version-based dispatch remain centralized in +`src/name_conversion.cpp`. Add new mappings there when needed. Inference, GGUF +export, quantization, and LoRA conversion share the `convert_tensor_name` entry +point; preserve its prefix normalization and component conversion order. + +### Components And Conditioning + +The core factory receives `model_builders::Context`, the tokenizer configuration, +and the output runners. Backend checks, optional IP-Adapter setup, scale +overrides, and tokenizer usage validation remain in the common builder. A +factory returns failure if any required additional backend cannot initialize. + +LLM architecture selection, tokenizer requirements, prompt preparation, and +encoding remain centralized in `src/conditioning/conditioner.hpp`. Integration +factories instantiate the conditioner with the model version; models requiring +new LLM conditioning behavior still need changes in that shared implementation. + +`vae_kind` is the single source of truth for the VAE family; compatibility +helpers derive their results from it. `KL`, `FLUX`, and `FLUX2` share the +`AutoEncoderKL` implementation while retaining their distinct latent formats. +The related traits have separate responsibilities: + +- `supports_video_generation` describes the model's video generation capability. +- `video_tae` selects `TinyVideoAutoEncoder` instead of `TinyImageAutoEncoder`. + Image models can use a video TAE, so this is independent of video generation. +- `fake_vae` selects FakeVAE for pixel-space models. It is separate from + `vae_kind`, which also controls latent-format compatibility helpers. +- `vae_format_override` allows PiD's VAE format override. + +`configure_vae` handles model-specific KL adjustments such as +PixArt's scale and SDXL's convolution scale. VAE compatibility traits also drive +latent transformations and previews through the existing helper functions. + +### Sampling And Geometry + +`ModelSamplingPolicy` groups the defaults first: `default_prediction`, +`default_flow_shift`, `default_sample_method`, and `default_scheduler`. Set only +the values that differ from the struct defaults. `prediction_from_weights` and +`flow_shift_from_weights` preserve checkpoint-dependent defaults during automatic +prediction selection; `probe_v_prediction` enables SD2's runtime v-prediction +probe. Explicit prediction settings take precedence over automatic selection. +MOSS's default unconditional skip layer stays +in the diffusion engine, with active unconditional SLG taking precedence. + +The resolved `prediction_t` selects the denoiser directly. Flux3, MiniMax H3 AV, +and Ming-Image use `FLUX3_FLOW_PRED`, `H3_AV_FLOW_PRED`, and +`MING_IMAGE_FLOW_PRED`; their CLI overrides are `flux3_flow`, `h3_av_flow`, and +`ming_image_flow`. Explicit `sd3_flow` always selects `DiscreteFlowDenoiser`, +and `flux_flow` always selects `FluxFlowDenoiser`, regardless of model version. +Automatic selection retains each model's default implementation. New prediction +values are appended before `PREDICTION_COUNT`, preserving existing valid values; +clients using the automatic-selection sentinel must rebuild with the new header. + +`ModelGeometry` declares `diffusion_down_factor`, `latent_channels`, and +`image_channels`. `ref_image_preset` defaults to `"default"`; set it only for models +that need another reference-image preset. + +The common pipeline executes these policies; specialized video, multimodal, +cache, and extension algorithms remain in their existing runtime components. +Model-specific diffusion extra parameters are assembled in the diffusion engine. +Video frame conversion and alignment also remain in the diffusion engine. +`get_vae_scale_factor()` retains PiD's pixel-space special case and otherwise +uses the VAE's scale factor. + +Validate migrations against the previous implementation using synthetic metadata, +name conversions, prompt preparation, default/override selection, and focused +builds. Real generation checks are still needed to validate numerical behavior +when changing network or encoding algorithms. diff --git a/docs/pe_cache.md b/docs/pe_cache.md new file mode 100644 index 000000000..e18195999 --- /dev/null +++ b/docs/pe_cache.md @@ -0,0 +1,60 @@ +# Position embedding cache + +`sd::PECache` stores CPU-side `sd::Tensor` data. Each model runner owns +its cache and calls the existing `make_input()` on every graph build. Cache +hits skip PE generation; allocation and upload remain ordinary graph-input +operations. `GGMLRunner` has no PE-specific code or resource hooks. + +The cache keeps the most recent key and tensor for each named PE input. A +matching key returns the existing tensor by reference. Changing its parameters +replaces that input's cached result. Different named inputs have independent +storage. Alternating two layouts under the same name recomputes their PE; +there is no LRU, device-buffer cache, memory budget, or graph lifecycle callback. +CPU data lives with the model runner, including across `runner_end()` calls, +and is released when its named entry is replaced or the cache is destroyed. + +## Adding a model + +Include `core/pe_cache.h` and declare a member: + +```cpp +sd::PECache pe_cache_; +``` + +Construct the key before generating PE and place the generation code in the +lazy callback: + +```cpp +sd::PECacheKey key("model.image_pe"); +key.add(height).add(width).add(head_dim).add(theta) + .add(circular_x_enabled).add(circular_y_enabled); +auto pe = make_input(pe_cache_.get_or_create(key, [&]() { + auto values = generate_position_embedding(/* existing arguments */); + return sd::Tensor({2, 2, head_dim / 2, token_count}, std::move(values)); +})); +``` + +Keys must cover every variable affecting PE values or shape: geometry, text +length, frequency settings, reference-image ordering and offsets, circular +settings, and video/audio positions or frame rates as applicable. +`add_shape()` records all four ggml dimensions, including trailing singleton +dimensions. `add_shapes()` records an ordered list of reference shapes. +`add(sd::Tensor)` records both shape and values for explicit position inputs. +Keys compare their complete encoded fields rather than a data pointer or hash. + +Use distinct names for PE inputs that coexist in a graph, such as Anima's +adapter query and key PE. `make_input()` defers reading the CPU data until graph +execution: do not replace a named entry while an unfinished graph references +it. Adding or updating other named entries does not invalidate its storage. +Return tensor shapes with the contiguous dimension first, and preserve each +model's existing F32 values, layout, and final transformations. + +## Current integrations + +The cache is used by MiniT2I, the Flux runner and its variants, Flux3, Qwen Image, +Qwen Image 2.1, Z-Image and L2P, Anima, Boogu, ERNIE Image, Fibo2, Ideogram4, +Iris, Krea2, Lens, LLaDA Image, MAGE Flow, Ming Image, MOSS Image, PiD, PixArt, +Wan, Hunyuan Video, LingBot Video, and LTX audio/video. + +Position embeddings stored as model weights and graph-side RoPE operations +remain on their existing paths, including MiniMax H3 and SenseNova U1. diff --git a/docs/taesd.md b/docs/taesd.md index 954086d18..8331f441b 100644 --- a/docs/taesd.md +++ b/docs/taesd.md @@ -49,3 +49,5 @@ curl -L -o taeqi2_1.safetensors https://huggingface.co/madebyollin/taeqi2_1/reso ``` Replace `--vae PATH` with `--taesd taeqi2_1.safetensors` in the [Qwen Image 2.1 examples](qwen_image_2.1.md) to use it for encoding and decoding. For previews only, keep `--vae PATH` and add `--taesd taeqi2_1.safetensors --taesd-preview-only --preview tae`. + +`sd-server` accepts the same `--taesd` and `--taesd-preview-only` options; the preview method is selected per request instead of via `--preview`. diff --git a/docs/wan.md b/docs/wan.md index 87a960f30..2031de7df 100644 --- a/docs/wan.md +++ b/docs/wan.md @@ -184,6 +184,8 @@ Notes: +TI2V 5B can decode multiple latent frames together when temporal tiling is enabled. For example, `--temporal-tiling --vae-tiling --vae-tile-size 256x256 --extra-tiling-args "temporal_tile_frames=4,temporal_decode_frames=2"` keeps up to four latent frames in each temporal tile and decodes up to two at once, with spatial tiles of 256×256 pixels. `temporal_decode_frames` defaults to `1`, preserving frame-by-frame decoding; other model versions ignore it. Increasing it requires more memory. Adjust the spatial tile size to fit the available memory. + ### Wan2.1 FLF2V 14B ``` diff --git a/examples/cli/CMakeLists.txt b/examples/cli/CMakeLists.txt index 6b1266cf5..90164cfc0 100644 --- a/examples/cli/CMakeLists.txt +++ b/examples/cli/CMakeLists.txt @@ -2,8 +2,10 @@ set(TARGET sd-cli) add_executable(${TARGET} ../common/common.cpp + ../common/generation_metadata.cpp ../common/log.cpp ../common/media_io.cpp + ../common/media_metadata.cpp image_metadata.cpp main.cpp ) diff --git a/examples/cli/README.md b/examples/cli/README.md index bf9de1862..578945549 100644 --- a/examples/cli/README.md +++ b/examples/cli/README.md @@ -31,14 +31,21 @@ rule for edge detection. See [Image preprocessing](../../docs/image_preprocessing.md) for input selectors, input defaults, downstream model processing, mask alignment, and examples. -Metadata mode inspects PNG/JPEG container metadata without loading any model: +Metadata mode inspects PNG, JPEG, WebP (including animated WebP), AVI, WebM, and Matroska +container metadata without loading any model. Use `--image` for both images and videos: ```bash ./bin/sd-cli -M metadata --image ./output.png ./bin/sd-cli -M metadata --image ./output.jpg --metadata-format json +./bin/sd-cli -M metadata --image ./output.webm +./bin/sd-cli -M metadata --image ./output.avi --metadata-format json ./bin/sd-cli -M metadata --image ./output.png --metadata-raw ./bin/sd-cli -M metadata --image ./output.png --metadata-all ``` For completely black or white images or videos, NaNs, and the `--linear-scale` / `--attn-scale` workaround, see [Troubleshooting](../../docs/troubleshooting.md). + +Generation outputs embed the core [generation record](../../docs/generation_metadata.md). +Images retain the native JSON and an additional WebUI-compatible infotext. +`--disable-image-metadata` omits both from files. diff --git a/examples/cli/image_metadata.cpp b/examples/cli/image_metadata.cpp index ba0d6da23..bb5ab7cf8 100644 --- a/examples/cli/image_metadata.cpp +++ b/examples/cli/image_metadata.cpp @@ -2,12 +2,14 @@ #include #include +#include #include #include #include #include #include #include +#include #include #include #include @@ -391,8 +393,39 @@ namespace { json& report, std::string& error); + json parse_exif_user_comment(const uint8_t* data, size_t size) { + if (size >= 8 && memcmp(data, "ASCII\0\0\0", 8) == 0) { + return trim_trailing_nuls(bytes_to_string(data + 8, data + size)); + } + if (size < 8 || memcmp(data, "UNICODE\0", 8) != 0) { + return json{{"raw_hex_preview", hex_preview(data, size)}}; + } + size_t offset = 8; + // WebUI's piexif writer uses big-endian UTF-16 without a BOM. + bool little_endian = false; + if (size >= 10 && ((data[8] == 0xfe && data[9] == 0xff) || (data[8] == 0xff && data[9] == 0xfe))) { + little_endian = data[8] == 0xff; + offset = 10; + } + if ((size - offset) % 2 != 0) { + return json{{"error", "truncated UTF-16 UserComment"}, {"raw_hex_preview", hex_preview(data, size)}}; + } + std::u16string text; + for (; offset < size; offset += 2) { + text.push_back(static_cast(little_endian ? data[offset] | (uint16_t(data[offset + 1]) << 8) + : (uint16_t(data[offset]) << 8) | data[offset + 1])); + } + try { + std::wstring_convert, char16_t> converter; + return trim_trailing_nuls(converter.to_bytes(text)); + } catch (const std::range_error&) { + return json{{"error", "invalid UTF-16 UserComment"}, {"raw_hex_preview", hex_preview(data, size)}}; + } + } + json parse_exif_value(const std::vector& data, size_t value_offset, + uint16_t tag, uint16_t type, uint32_t count, bool little_endian, @@ -466,8 +499,9 @@ namespace { break; } case 7: - value = bytes_to_string(data.data() + value_offset, - data.data() + value_offset + count); + value = tag == 0x9286 ? parse_exif_user_comment(data.data() + value_offset, count) + : json(bytes_to_string(data.data() + value_offset, + data.data() + value_offset + count)); break; case 9: { if (count == 1) { @@ -591,6 +625,7 @@ namespace { tag_json["value"] = parse_exif_value(data, value_offset, + tag, type, value_count, little_endian, @@ -1098,6 +1133,178 @@ namespace { return true; } + static bool parse_avi_chunks(const std::vector& data, size_t offset, size_t end, bool info, bool include_raw, json& result, std::string& error) { + while (offset < end) { + if (end - offset < 8) { + error = "truncated AVI chunk header"; + return false; + } + const auto name = bytes_to_string(data.data() + offset, data.data() + offset + 4); + const size_t size = read_u32_le(data, offset + 4); + offset += 8; + if (size > end - offset || (size & 1u) > end - offset - size) { + error = "AVI chunk exceeds its container"; + return false; + } + json entry = {{"entry_type", "chunk"}, {"name", name}, {"length", size}, {"metadata_like", info}}; + if (name == "LIST") { + if (size < 4) { + error = "truncated AVI LIST type"; + return false; + } + const auto type = bytes_to_string(data.data() + offset, data.data() + offset + 4); + entry["data"]["list_type"] = type; + if (type == "INFO" && !info && + !parse_avi_chunks(data, offset + 4, offset + size, true, include_raw, result, error)) + return false; + } else if (info) { + entry["data"]["text"] = trim_trailing_nuls(bytes_to_string(data.data() + offset, data.data() + offset + size)); + } else { + append_raw_preview(entry, data.data() + offset, size, include_raw); + } + result["entries"].push_back(std::move(entry)); + offset += size + (size & 1u); + } + return true; + } + + static bool parse_avi(const std::vector& data, bool include_raw, json& result, std::string& error) { + result = {{"format", "AVI"}, {"entries", json::array()}}; + size_t offset = 0; + while (offset < data.size()) { + if (data.size() - offset < 12 || memcmp(data.data() + offset, "RIFF", 4) != 0 || + memcmp(data.data() + offset + 8, offset == 0 ? "AVI " : "AVIX", 4) != 0) { + error = "invalid AVI RIFF header"; + return false; + } + const size_t size = read_u32_le(data, offset + 4); + if (size < 4 || size > data.size() - offset - 8 || + (size & 1u) > data.size() - offset - 8 - size) { + error = "AVI RIFF size exceeds file size"; + return false; + } + if (!parse_avi_chunks(data, offset + 12, offset + 8 + size, false, include_raw, result, error)) + return false; + offset += 8 + size + (size & 1u); + } + return true; + } + + static bool read_ebml_uint(const std::vector& data, size_t& offset, size_t end, bool id, uint64_t& value, bool& unknown) { + if (offset == end || data[offset] == 0) + return false; + uint8_t mask = 0x80; + size_t length = 1; + while (!(data[offset] & mask)) { + mask >>= 1; + ++length; + } + if (length > (id ? 4u : 8u) || length > end - offset) + return false; + value = data[offset++] & (id ? 0xff : mask - 1); + for (size_t i = 1; i < length; ++i) + value = (value << 8) | data[offset++]; + unknown = !id && value == (uint64_t(1) << (7 * length)) - 1; + return true; + } + + static bool is_webm_segment_child(uint64_t id) { + return id == 0x114D9B74 || id == 0x1549A966 || id == 0x1654AE6B || + id == 0x1F43B675 || id == 0x1C53BB6B || id == 0x1941A469 || + id == 0x1043A770 || id == 0x1254C367; + } + + static bool parse_webm_elements(const std::vector& data, size_t& offset, size_t end, uint64_t parent, bool unknown_cluster, int depth, bool include_raw, json& result, std::string& error) { + if (depth > 16) { + error = "WebM tag nesting exceeds limit"; + return false; + } + std::string tag_name; + std::string tag_text; + bool has_text = false; + while (offset < end) { + const auto start = offset; + uint64_t id, size; + bool unknown; + if (!read_ebml_uint(data, offset, end, true, id, unknown)) { + error = "invalid or truncated WebM element ID"; + return false; + } + // An unknown-size Cluster ends at the next Segment-level element. + if (unknown_cluster && is_webm_segment_child(id)) { + offset = start; + return true; + } + if (!read_ebml_uint(data, offset, end, false, size, unknown) || + (!unknown && size > end - offset)) { + error = "invalid or truncated WebM element size"; + return false; + } + if (unknown && id != 0x18538067 && id != 0x1F43B675) { + error = "unsupported unknown-size WebM element"; + return false; + } + const size_t stop = unknown ? end : offset + static_cast(size); + const bool container = (parent == 0 && (id == 0x1A45DFA3 || id == 0x18538067)) || + (parent == 0x18538067 && (id == 0x1254C367 || (id == 0x1F43B675 && unknown))) || + (parent == 0x1254C367 && id == 0x7373) || + ((parent == 0x7373 || parent == 0x67C8) && id == 0x67C8); + if (parent == 0 && id == 0x18538067) + result["has_segment"] = true; + if (parent == 0x1A45DFA3 && id == 0x4282) { + const auto type = bytes_to_string(data.data() + offset, data.data() + stop); + if (type != "webm" && type != "matroska") { + error = "EBML document is not WebM or Matroska"; + return false; + } + result["doc_type"] = type; + result["format"] = type == "webm" ? "WEBM" : "MATROSKA"; + } + if (parent == 0x67C8 && (id == 0x45A3 || id == 0x4487)) { + auto value = bytes_to_string(data.data() + offset, data.data() + stop); + if (id == 0x45A3) + tag_name = std::move(value); + else { + tag_text = std::move(value); + has_text = true; + } + } + if (parent == 0 || parent == 0x18538067) { + std::ostringstream name; + name << "0x" << std::hex << std::uppercase << id; + json entry = {{"entry_type", "element"}, {"name", name.str()}, {"length", stop - offset}, {"metadata_like", false}}; + if (unknown) + entry["unknown_size"] = true; + append_raw_preview(entry, data.data() + offset, stop - offset, include_raw); + result["entries"].push_back(std::move(entry)); + } + if (container) { + if (!parse_webm_elements(data, offset, stop, id, id == 0x1F43B675 && unknown, + depth + 1, include_raw, result, error)) + return false; + } else { + offset = stop; + } + } + if (parent == 0x67C8 && has_text) { + result["entries"].push_back({{"entry_type", "tag"}, {"name", tag_name}, {"metadata_like", true}, {"data", {{"text", tag_text}}}}); + } + return true; + } + + static bool parse_webm(const std::vector& data, bool include_raw, json& result, std::string& error) { + result = {{"format", "WEBM"}, {"entries", json::array()}}; + size_t offset = 0; + if (!parse_webm_elements(data, offset, data.size(), 0, false, 0, include_raw, result, error)) + return false; + if (!result.contains("doc_type") || !result.value("has_segment", false)) { + error = "missing WebM document type or Segment"; + return false; + } + result.erase("has_segment"); + return true; + } + std::string abbreviate(const std::string& value, bool brief) { if (!brief || value.size() <= 240) { return value; @@ -1116,7 +1323,7 @@ namespace { return; } if (value.is_primitive()) { - out << prefix << key << ": " << value.dump() << "\n"; + out << prefix << key << ": " << value.dump(-1, ' ', false, json::error_handler_t::replace) << "\n"; return; } if (value.is_array()) { @@ -1125,7 +1332,7 @@ namespace { if (item.is_string()) { out << prefix << " - " << abbreviate(item.get(), brief) << "\n"; } else if (item.is_primitive()) { - out << prefix << " - " << item.dump() << "\n"; + out << prefix << " - " << item.dump(-1, ' ', false, json::error_handler_t::replace) << "\n"; } else { out << prefix << " -\n"; for (auto it = item.begin(); it != item.end(); ++it) { @@ -1158,10 +1365,12 @@ namespace { } for (const auto& entry : entries) { - const bool is_chunk = entry.value("entry_type", "") == "chunk"; - const std::string name = is_chunk ? entry.value("name", "unknown") - : entry.value("marker", "unknown"); - out << (is_chunk ? "Chunk: " : "Segment: ") << name << "\n"; + const auto type = entry.value("entry_type", ""); + const auto name = entry.value("name", entry.value("marker", "unknown")); + const char* label = type == "chunk" ? "Chunk" : type == "tag" ? "Tag" + : type == "element" ? "Element" + : "Segment"; + out << label << ": " << name << "\n"; for (auto it = entry.begin(); it != entry.end(); ++it) { if (it.key() == "entry_type" || it.key() == "name" || it.key() == "marker" || it.key() == "metadata_like") { @@ -1211,7 +1420,13 @@ namespace { return parse_webp(data, include_raw, report, error); } - error = "unsupported image format; only PNG, JPEG, and WebP are supported"; + if (data.size() >= 12 && memcmp(data.data(), "RIFF", 4) == 0 && + memcmp(data.data() + 8, "AVI ", 4) == 0) + return parse_avi(data, include_raw, report, error); + if (data.size() >= 4 && read_u32_be(data, 0) == 0x1A45DFA3) + return parse_webm(data, include_raw, report, error); + + error = "unsupported media format; only PNG, JPEG, WebP, AVI, WebM, and Matroska are supported"; return false; } @@ -1230,7 +1445,7 @@ bool print_image_metadata(const std::string& image_path, if (options.output_format == MetadataOutputFormat::JSON) { visible_report["file"] = image_path; - out << visible_report.dump(2) << "\n"; + out << visible_report.dump(2, ' ', false, json::error_handler_t::replace) << "\n"; } else { print_text_report(image_path, visible_report, options.include_structural, options.brief, out); } diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 8ee51865f..591a15112 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -47,7 +47,6 @@ struct SDCliParams { int preview_interval = 1; std::string preview_path = "preview.png"; int preview_fps = 16; - bool taesd_preview = false; bool preview_noisy = false; bool color = false; bool metadata_raw = false; @@ -70,7 +69,7 @@ struct SDCliParams { &output_path}, {"", "--image", - "path to the image to inspect (for metadata mode)", + "path to the image or video to inspect (for metadata mode)", 0, &image_path}, {"", @@ -114,10 +113,6 @@ struct SDCliParams { "--color", "colors the logging tags according to level", true, &color}, - {"", - "--taesd-preview-only", - std::string("prevents usage of taesd for decoding the final image. (for use with --preview ") + previews_str[PREVIEW_TAE] + ")", - true, &taesd_preview}, {"", "--preview-noisy", "enables previewing noisy inputs of the models rather than the denoised outputs", @@ -232,7 +227,7 @@ struct SDCliParams { } } else { if (image_path.empty()) { - LOG_ERROR("error: metadata mode needs an image path (--image)"); + LOG_ERROR("error: metadata mode needs an image or video path (--image)"); return false; } if (metadata_format != "text" && metadata_format != "json") { @@ -268,7 +263,6 @@ struct SDCliParams { << " preview_interval: " << preview_interval << ",\n" << " preview_path: \"" << preview_path << "\",\n" << " preview_fps: " << preview_fps << ",\n" - << " taesd_preview: " << (taesd_preview ? "true" : "false") << ",\n" << " preview_noisy: " << (preview_noisy ? "true" : "false") << ",\n" << " imatrix_out: \"" << imatrix_out << "\",\n" << " metadata_raw: " << (metadata_raw ? "true" : "false") << ",\n" @@ -447,10 +441,11 @@ static fs::path get_video_audio_sidecar_path(const SDCliParams& cli_params) { bool save_results(const SDCliParams& cli_params, const SDContextParams& ctx_params, const SDGenerationParams& gen_params, - sd_image_t* results, + SDImageVec& results, int num_results, - const sd_audio_t* generated_audio = nullptr) { - if (results == nullptr || num_results <= 0) { + const sd_audio_t* generated_audio, + const std::string& video_info) { + if (results.empty() || num_results <= 0) { return false; } @@ -493,12 +488,10 @@ bool save_results(const SDCliParams& cli_params, if (!img.data) return false; - int images_per_batch = gen_params.batch_count > 0 ? std::max(1, num_results / gen_params.batch_count) : 1; - const int64_t metadata_seed = cli_params.mode == VID_GEN ? gen_params.seed : gen_params.seed + idx / images_per_batch; - std::string params = gen_params.embed_image_metadata - ? get_image_params(ctx_params, gen_params, metadata_seed, cli_params.mode) - : ""; - const bool ok = write_image_to_file(path.u8string(), img.data, img.width, img.height, img.channel, params, cli_params.compression_quality); + std::string params = gen_params.embed_image_metadata + ? results.generation_info(idx) + : ""; + const bool ok = write_image_to_file(path.u8string(), img.data, img.width, img.height, img.channel, params, cli_params.compression_quality); LOG_INFO("save result image %d to '%s' (%s)", idx, path.u8string().c_str(), ok ? "success" : "failure"); return ok; }; @@ -540,7 +533,7 @@ bool save_results(const SDCliParams& cli_params, if (ext_lower != ".avi" && ext_lower != ".webp" && ext_lower != ".webm") ext = ".avi"; std::string params = gen_params.embed_image_metadata - ? get_image_params(ctx_params, gen_params, gen_params.seed, cli_params.mode) + ? video_info : ""; fs::path video_path = base_path; @@ -548,7 +541,7 @@ bool save_results(const SDCliParams& cli_params, std::string final_ext_lower = ext.u8string(); std::transform(final_ext_lower.begin(), final_ext_lower.end(), final_ext_lower.begin(), ::tolower); const bool mux_audio = generated_audio != nullptr && (final_ext_lower == ".avi" || final_ext_lower == ".webm"); - if (create_video_from_sd_images(video_path.u8string().c_str(), results, num_results, gen_params.fps, cli_params.compression_quality, mux_audio ? generated_audio : nullptr, params) == 0) { + if (create_video_from_sd_images(video_path.u8string().c_str(), results.data(), num_results, gen_params.fps, cli_params.compression_quality, mux_audio ? generated_audio : nullptr, params) == 0) { LOG_INFO("save result video to '%s'", video_path.u8string().c_str()); if (generated_audio != nullptr && !mux_audio) { fs::path wav_path = video_path; @@ -613,30 +606,29 @@ static bool apply_adetailer(sd_ctx_t* sd_ctx, continue; } sd_img_gen_params_t ad_generation_params = img_gen_params; - ad_generation_params.seed = img_gen_params.seed + i; + ad_generation_params.seed = generation_seed(results.generation_info(i), img_gen_params.seed); 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; + sd_image_generation_result_t* detailed = nullptr; 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); + &detailed) || + detailed == nullptr || detailed->count <= 0 || detailed->items[0].image.data == nullptr) { + free_sd_image_generation_result(detailed); LOG_ERROR("ADetailer failed for image %d", i + 1); return false; } + results.append_generation_info(i, detailed->items[0].generation_info, ctx_params.lora_model_dir); free(results[i].data); - results[i] = detailed_images[0]; - detailed_images[0] = {0, 0, 0, nullptr}; - free_sd_images(detailed_images, detailed_count); + results[i] = detailed->items[0].image; + detailed->items[0].image = {}; + free_sd_image_generation_result(detailed); } return true; } @@ -885,11 +877,12 @@ int main(int argc, const char* argv[]) { } } - sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(cli_params.taesd_preview); + sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(); SDImageVec results; int num_results = 0; sd_audio_t* generated_audio = nullptr; + std::string video_info; if (cli_params.mode == UPSCALE) { num_results = 1; @@ -902,18 +895,6 @@ int main(int argc, const char* argv[]) { return 1; } - if (gen_params.sample_params.sample_method == SAMPLE_METHOD_COUNT) { - gen_params.sample_params.sample_method = sd_get_default_sample_method(sd_ctx.get()); - } - - if (gen_params.high_noise_sample_params.sample_method == SAMPLE_METHOD_COUNT) { - gen_params.high_noise_sample_params.sample_method = sd_get_default_sample_method(sd_ctx.get()); - } - - if (gen_params.sample_params.scheduler == SCHEDULER_COUNT) { - gen_params.sample_params.scheduler = sd_get_default_scheduler(sd_ctx.get(), gen_params.sample_params.sample_method); - } - sd_img_gen_params_t img_gen_params{}; const bool use_img_gen_params = cli_params.mode == IMG_GEN || cli_params.mode == ADETAILER; if (use_img_gen_params) { @@ -921,23 +902,24 @@ int main(int argc, const char* argv[]) { } if (cli_params.mode == IMG_GEN) { - sd_image_t* generated_images = nullptr; - if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) { - generated_images = nullptr; - num_results = 0; - } - results.adopt(generated_images, num_results); + sd_image_generation_result_t* generated = nullptr; + generate_image(sd_ctx.get(), &img_gen_params, &generated); + results.adopt(generated, ctx_params.lora_model_dir); + num_results = results.count(); } else if (cli_params.mode == ADETAILER) { num_results = 1; results.push_back(gen_params.init_image.release()); } else if (cli_params.mode == VID_GEN) { - sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t(); - sd_image_t* generated_video = nullptr; - if (!generate_video(sd_ctx.get(), &vid_gen_params, &generated_video, &num_results, &generated_audio, &cli_params.preview_fps)) { - generated_video = nullptr; + sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t(); + sd_video_generation_result_t* generated = nullptr; + if (generate_video(sd_ctx.get(), &vid_gen_params, &generated)) { + gen_params.fps = cli_params.preview_fps = generated->fps; + generated_audio = generated->audio; + generated->audio = nullptr; + num_results = generated->frame_count; + results.adopt_video(generated, ctx_params.lora_model_dir); + video_info = results.generation_info(0); } - gen_params.fps = cli_params.preview_fps; - results.adopt(generated_video, num_results); } if (!results) { @@ -977,21 +959,20 @@ int main(int argc, const char* argv[]) { SDImageOwner current_image(results[i]); results[i] = {0, 0, 0, nullptr}; for (int u = 0; u < gen_params.upscale_repeats; ++u) { - sd_image_t* upscaled_images = nullptr; - int upscaled_count = 0; - bool upscale_ok = upscale(upscaler_ctx.get(), - current_image.get(), - upscale_factor, - &upscaled_images, - &upscaled_count); - if (!upscale_ok || upscaled_count <= 0 || upscaled_images[0].data == nullptr) { - free_sd_images(upscaled_images, upscaled_count); + sd_image_generation_result_t* upscaled = nullptr; + bool upscale_ok = upscale(upscaler_ctx.get(), + current_image.get(), + upscale_factor, + &upscaled); + if (!upscale_ok || !upscaled || upscaled->count <= 0 || upscaled->items[0].image.data == nullptr) { + free_sd_image_generation_result(upscaled); LOG_ERROR("upscale failed"); break; } - sd_image_t upscaled_image = upscaled_images[0]; - upscaled_images[0] = {0, 0, 0, nullptr}; - free_sd_images(upscaled_images, upscaled_count); + results.append_generation_info(i, upscaled->items[0].generation_info); + sd_image_t upscaled_image = upscaled->items[0].image; + upscaled->items[0].image = {}; + free_sd_image_generation_result(upscaled); current_image.reset(upscaled_image); } results[i] = current_image.release(); // Set the final upscaled image as the result @@ -999,7 +980,9 @@ int main(int argc, const char* argv[]) { } } - if (!save_results(cli_params, ctx_params, gen_params, results.data(), num_results, generated_audio)) { + if (!video_info.empty() && results.count() > 0) + video_info = results.generation_info(0); + if (!save_results(cli_params, ctx_params, gen_params, results, num_results, generated_audio, video_info)) { free_sd_audio(generated_audio); return 1; } diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 7f65edac2..00301dc9c 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -641,6 +641,10 @@ ArgOptions SDContextParams::get_options() { "--sage-attn", "use native CUDA SageAttention in the diffusion model, with flash/default attention fallback", true, &sage_attn}, + {"", + "--taesd-preview-only", + "prevents usage of taesd for decoding the final image; the VAE keeps encoding / decoding and taesd is only used for step previews", + true, &taesd_preview_only}, }; auto on_type_arg = [&](int argc, const char** argv, int index) { @@ -814,7 +818,7 @@ ArgOptions SDContextParams::get_options() { on_sampler_rng_arg}, {"", "--prediction", - "prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, sefi_flow]", + "prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, sefi_flow, minit2i_flow, sensenova_u1_flow, flux3_flow, h3_av_flow, ming_image_flow]", on_prediction_arg}, {"", "--lora-apply-mode", @@ -1014,6 +1018,7 @@ std::string SDContextParams::to_string() const { << " attn_scale: " << attn_scale << ",\n" << " diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n" << " vae_conv_direct: " << (vae_conv_direct ? "true" : "false") << ",\n" + << " taesd_preview_only: " << (taesd_preview_only ? "true" : "false") << ",\n" << " prediction: " << sd_prediction_name(prediction) << ",\n" << " lora_apply_mode: " << sd_lora_apply_mode_name(lora_apply_mode) << ",\n" << " force_sdxl_vae_conv_scale: " << (force_sdxl_vae_conv_scale ? "true" : "false") << "\n" @@ -1021,7 +1026,7 @@ std::string SDContextParams::to_string() const { return oss.str(); } -sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { +sd_ctx_params_t SDContextParams::to_sd_ctx_params_t() { prepare_backend_assignments(); embedding_vec.clear(); embedding_vec.reserve(embedding_map.size()); @@ -1071,7 +1076,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.sage_attn = sage_attn; sd_ctx_params.linear_scale = linear_scale; sd_ctx_params.attn_scale = attn_scale; - sd_ctx_params.tae_preview_only = taesd_preview; + sd_ctx_params.tae_preview_only = taesd_preview_only; sd_ctx_params.diffusion_conv_direct = diffusion_conv_direct; sd_ctx_params.vae_conv_direct = vae_conv_direct; sd_ctx_params.force_sdxl_vae_conv_scale = force_sdxl_vae_conv_scale; @@ -1189,7 +1194,7 @@ ArgOptions SDGenerationParams::get_options() { &extra_sample_args}, {"", "--extra-tiling-args", - "extra VAE tiling args, key=value list. Supported video VAEs accept temporal_tile_frames/temporal_tile_size (default: 4), temporal_tile_overlap (default: 1)", + "extra VAE tiling args, key=value list. Supported video VAEs accept temporal_tile_frames/temporal_tile_size (default: 4), temporal_tile_overlap (default: 1); Wan2.2 TI2V 5B accepts temporal_decode_frames (default: 1)", (int)',', &extra_tiling_args}, {"", @@ -2543,11 +2548,6 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::s return false; } - if (seed < 0) { - srand((int)time(nullptr)); - seed = rand(); - } - if (strict) { batch_count = std::clamp(batch_count, 1, 8); sample_params.sample_steps = std::clamp(sample_params.sample_steps, 1, 100); @@ -2914,6 +2914,7 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() { params.circular_x = circular || circular_x; params.circular_y = circular || circular_y; params.image_preprocess = {image_preprocess.c_str()}; + params.ref_image_args = ref_image_args.c_str(); return params; } @@ -3023,301 +3024,3 @@ std::string SDGenerationParams::to_string() const { std::string version_string() { return std::string("stable-diffusion.cpp version ") + sd_version() + ", commit " + sd_commit(); } - -static std::string safe_json_string(const char* value) { - return value ? value : ""; -} - -static void set_json_basename_if_not_empty(json& target, const char* key, const std::string& path) { - if (!path.empty()) { - target[key] = sd_basename(path); - } -} - -static json build_sampling_metadata_json(const sd_sample_params_t& sample_params, - const std::vector& skip_layers, - const std::vector* custom_sigmas = nullptr) { - json sampling = { - {"steps", sample_params.sample_steps}, - {"eta", sample_params.eta}, - {"shifted_timestep", sample_params.shifted_timestep}, - {"flow_shift", sample_params.flow_shift}, - {"extra_sample_args", safe_json_string(sample_params.extra_sample_args)}, - {"guidance", - { - {"txt_cfg", sample_params.guidance.txt_cfg}, - {"img_cfg", sample_params.guidance.img_cfg}, - {"distilled_guidance", sample_params.guidance.distilled_guidance}, - {"slg", - { - {"scale", sample_params.guidance.slg.scale}, - {"layers", skip_layers}, - {"start", sample_params.guidance.slg.layer_start}, - {"end", sample_params.guidance.slg.layer_end}, - }}, - }}, - }; - if (sample_params.sample_method != SAMPLE_METHOD_COUNT) { - sampling["method"] = safe_json_string(sd_sample_method_name(sample_params.sample_method)); - } - if (sample_params.scheduler != SCHEDULER_COUNT) { - sampling["scheduler"] = safe_json_string(sd_scheduler_name(sample_params.scheduler)); - } - if (custom_sigmas != nullptr) { - sampling["custom_sigmas"] = *custom_sigmas; - } - return sampling; -} - -std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params, - const SDGenerationParams& gen_params, - int64_t seed, - SDMode mode) { - json root; - root["schema"] = "sdcpp.image.params/v1"; - root["mode"] = "img_gen"; - if (mode == VID_GEN) { - root["mode"] = "vid_gen"; - } else if (mode == ADETAILER) { - root["mode"] = "adetailer"; - } - root["generator"] = { - {"name", "stable-diffusion.cpp"}, - {"version", safe_json_string(sd_version())}, - {"commit", safe_json_string(sd_commit())}, - }; - root["seed"] = seed; - root["width"] = gen_params.get_resolved_width(); - root["height"] = gen_params.get_resolved_height(); - - root["prompt"] = { - {"positive", gen_params.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, - gen_params.skip_layers, - &gen_params.custom_sigmas); - - json models; - set_json_basename_if_not_empty(models, "model", ctx_params.model_path); - set_json_basename_if_not_empty(models, "clip_l", ctx_params.clip_l_path); - set_json_basename_if_not_empty(models, "clip_g", ctx_params.clip_g_path); - set_json_basename_if_not_empty(models, "clip_vision", ctx_params.clip_vision_path); - set_json_basename_if_not_empty(models, "t5xxl", ctx_params.t5xxl_path); - set_json_basename_if_not_empty(models, "llm", ctx_params.llm_path); - set_json_basename_if_not_empty(models, "llm_vision", ctx_params.llm_vision_path); - set_json_basename_if_not_empty(models, "diffusion_model", ctx_params.diffusion_model_path); - set_json_basename_if_not_empty(models, "high_noise_diffusion_model", ctx_params.high_noise_diffusion_model_path); - set_json_basename_if_not_empty(models, "uncond_diffusion_model", ctx_params.uncond_diffusion_model_path); - set_json_basename_if_not_empty(models, "vae", ctx_params.vae_path); - set_json_basename_if_not_empty(models, "taesd", ctx_params.taesd_path); - set_json_basename_if_not_empty(models, "control_net", ctx_params.control_net_path); - root["models"] = std::move(models); - - root["clip_skip"] = gen_params.clip_skip; - root["strength"] = gen_params.strength; - root["control_strength"] = gen_params.control_strength; - root["ip_adapter_strength"] = gen_params.ip_adapter_strength; - root["ref_image_args"] = gen_params.ref_image_args; - root["image_preprocess"] = gen_params.image_preprocess; - root["increase_ref_index"] = gen_params.increase_ref_index; - if (mode == VID_GEN) { - root["video"] = { - {"frame_count", gen_params.video_frames}, - {"fps", gen_params.fps}, - }; - root["moe_boundary"] = gen_params.moe_boundary; - root["vace_strength"] = gen_params.vace_strength; - root["high_noise_sampling"] = build_sampling_metadata_json(gen_params.high_noise_sample_params, - gen_params.high_noise_skip_layers); - } - - root["rng"] = safe_json_string(sd_rng_type_name(ctx_params.rng_type)); - if (ctx_params.sampler_rng_type != RNG_TYPE_COUNT) { - root["sampler_rng"] = safe_json_string(sd_rng_type_name(ctx_params.sampler_rng_type)); - } - - json loras = json::array(); - for (const auto& entry : gen_params.lora_map) { - loras.push_back({ - {"name", sd_basename(entry.first)}, - {"multiplier", entry.second}, - {"is_high_noise", false}, - }); - } - for (const auto& entry : gen_params.high_noise_lora_map) { - loras.push_back({ - {"name", sd_basename(entry.first)}, - {"multiplier", entry.second}, - {"is_high_noise", true}, - }); - } - if (!loras.empty()) { - root["loras"] = std::move(loras); - } - - if (gen_params.hires_enabled) { - root["hires"] = { - {"enabled", gen_params.hires_enabled}, - {"upscaler", gen_params.hires_upscaler}, - {"model", gen_params.hires_upscaler_model_path.empty() ? "" : sd_basename(gen_params.hires_upscaler_model_path)}, - {"scale", gen_params.hires_scale}, - {"target_width", gen_params.hires_width}, - {"target_height", gen_params.hires_height}, - {"steps", gen_params.hires_steps}, - {"denoising_strength", gen_params.hires_denoising_strength}, - {"custom_sigmas", gen_params.hires_custom_sigmas}, - {"upscale_tile_size", gen_params.hires_upscale_tile_size}, - }; - } - - if (gen_params.cache_params.mode != SD_CACHE_DISABLED) { - root["cache"] = { - {"requested_mode", gen_params.cache_mode}, - {"requested_option", gen_params.cache_option}, - {"mode", gen_params.cache_params.mode}, - {"scm_mask", gen_params.scm_mask}, - {"scm_policy_dynamic", gen_params.scm_policy_dynamic}, - {"reuse_threshold", gen_params.cache_params.reuse_threshold}, - {"start_percent", gen_params.cache_params.start_percent}, - {"end_percent", gen_params.cache_params.end_percent}, - {"error_decay_rate", gen_params.cache_params.error_decay_rate}, - {"use_relative_threshold", gen_params.cache_params.use_relative_threshold}, - {"reset_error_on_compute", gen_params.cache_params.reset_error_on_compute}, - {"Fn_compute_blocks", gen_params.cache_params.Fn_compute_blocks}, - {"Bn_compute_blocks", gen_params.cache_params.Bn_compute_blocks}, - {"residual_diff_threshold", gen_params.cache_params.residual_diff_threshold}, - {"max_warmup_steps", gen_params.cache_params.max_warmup_steps}, - {"max_cached_steps", gen_params.cache_params.max_cached_steps}, - {"max_continuous_cached_steps", gen_params.cache_params.max_continuous_cached_steps}, - {"taylorseer_n_derivatives", gen_params.cache_params.taylorseer_n_derivatives}, - {"taylorseer_skip_interval", gen_params.cache_params.taylorseer_skip_interval}, - {"spectrum_w", gen_params.cache_params.spectrum_w}, - {"spectrum_m", gen_params.cache_params.spectrum_m}, - {"spectrum_lam", gen_params.cache_params.spectrum_lam}, - {"spectrum_window_size", gen_params.cache_params.spectrum_window_size}, - {"spectrum_flex_window", gen_params.cache_params.spectrum_flex_window}, - {"spectrum_warmup_steps", gen_params.cache_params.spectrum_warmup_steps}, - {"spectrum_stop_percent", gen_params.cache_params.spectrum_stop_percent}, - }; - } - - if (gen_params.vae_tiling_params.enabled || - gen_params.vae_tiling_params.temporal_tiling || - !gen_params.extra_tiling_args.empty()) { - root["vae_tiling"] = { - {"enabled", gen_params.vae_tiling_params.enabled}, - {"temporal_tiling", gen_params.vae_tiling_params.temporal_tiling}, - {"tile_size_w", gen_params.vae_tiling_params.tile_size_w}, - {"tile_size_h", gen_params.vae_tiling_params.tile_size_h}, - {"target_overlap", gen_params.vae_tiling_params.target_overlap}, - {"rel_size_w", gen_params.vae_tiling_params.rel_size_w}, - {"rel_size_h", gen_params.vae_tiling_params.rel_size_h}, - {"extra_tiling_args", gen_params.extra_tiling_args}, - }; - } - - return root.dump(); -} - -std::string get_image_params(const SDContextParams& ctx_params, - const SDGenerationParams& gen_params, - int64_t seed, - SDMode mode) { - std::string parameter_string; - if (gen_params.prompt_with_lora.size() != 0) { - parameter_string += gen_params.prompt_with_lora + "\n"; - } else { - parameter_string += gen_params.prompt + "\n"; - } - if (gen_params.negative_prompt.size() != 0) { - parameter_string += "Negative prompt: " + gen_params.negative_prompt + "\n"; - } - parameter_string += "Steps: " + std::to_string(gen_params.sample_params.sample_steps) + ", "; - parameter_string += "CFG scale: " + std::to_string(gen_params.sample_params.guidance.txt_cfg) + ", "; - if (gen_params.sample_params.guidance.slg.scale != 0 && gen_params.skip_layers.size() != 0) { - parameter_string += "SLG scale: " + std::to_string(gen_params.sample_params.guidance.slg.scale) + ", "; - parameter_string += "Skip layers: ["; - for (const auto& layer : gen_params.skip_layers) { - parameter_string += std::to_string(layer) + ", "; - } - parameter_string += "], "; - parameter_string += "Skip layer start: " + std::to_string(gen_params.sample_params.guidance.slg.layer_start) + ", "; - parameter_string += "Skip layer end: " + std::to_string(gen_params.sample_params.guidance.slg.layer_end) + ", "; - } - parameter_string += "Guidance: " + std::to_string(gen_params.sample_params.guidance.distilled_guidance) + ", "; - parameter_string += "Eta: " + std::to_string(gen_params.sample_params.eta) + ", "; - if (!gen_params.extra_sample_args.empty()) { - 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 += "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 += "RNG: " + std::string(sd_rng_type_name(ctx_params.rng_type)) + ", "; - if (ctx_params.sampler_rng_type != RNG_TYPE_COUNT) { - parameter_string += "Sampler RNG: " + std::string(sd_rng_type_name(ctx_params.sampler_rng_type)) + ", "; - } - parameter_string += "Sampler: " + std::string(sd_sample_method_name(gen_params.sample_params.sample_method)); - if (!gen_params.custom_sigmas.empty()) { - parameter_string += ", Custom Sigmas: ["; - for (size_t i = 0; i < gen_params.custom_sigmas.size(); ++i) { - std::ostringstream oss; - oss << std::fixed << std::setprecision(4) << gen_params.custom_sigmas[i]; - parameter_string += oss.str() + (i == gen_params.custom_sigmas.size() - 1 ? "" : ", "); - } - parameter_string += "]"; - } else if (gen_params.sample_params.scheduler != SCHEDULER_COUNT) { // Only show schedule if not using custom sigmas - parameter_string += " " + std::string(sd_scheduler_name(gen_params.sample_params.scheduler)); - } - parameter_string += ", "; - for (const auto& te : {ctx_params.clip_l_path, ctx_params.clip_g_path, ctx_params.t5xxl_path, ctx_params.llm_path, ctx_params.llm_vision_path}) { - if (!te.empty()) { - parameter_string += "TE: " + sd_basename(te) + ", "; - } - } - if (!ctx_params.diffusion_model_path.empty()) { - parameter_string += "Unet: " + sd_basename(ctx_params.diffusion_model_path) + ", "; - } - if (!ctx_params.uncond_diffusion_model_path.empty()) { - parameter_string += "Uncond Unet: " + sd_basename(ctx_params.uncond_diffusion_model_path) + ", "; - } - if (!ctx_params.vae_path.empty()) { - parameter_string += "VAE: " + sd_basename(ctx_params.vae_path) + ", "; - } - if (gen_params.clip_skip != -1) { - parameter_string += "Clip skip: " + std::to_string(gen_params.clip_skip) + ", "; - } - if (gen_params.hires_enabled) { - parameter_string += "Hires upscale: " + gen_params.hires_upscaler + ", "; - parameter_string += "Hires scale: " + std::to_string(gen_params.hires_scale) + ", "; - parameter_string += "Hires resize: " + std::to_string(gen_params.hires_width) + "x" + std::to_string(gen_params.hires_height) + ", "; - parameter_string += "Hires steps: " + std::to_string(gen_params.hires_steps) + ", "; - parameter_string += "Denoising strength: " + std::to_string(gen_params.hires_denoising_strength) + ", "; - if (!gen_params.hires_custom_sigmas.empty()) { - parameter_string += "Hires custom sigmas: " + vec_to_string(gen_params.hires_custom_sigmas) + ", "; - } - } - parameter_string += "Version: stable-diffusion.cpp"; - parameter_string += ", SDCPP: " + build_sdcpp_image_metadata_json(ctx_params, gen_params, seed, mode); - return parameter_string; -} diff --git a/examples/common/common.h b/examples/common/common.h index 992f74c17..39d6011af 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -1,6 +1,8 @@ #ifndef __EXAMPLES_COMMON_COMMON_H__ #define __EXAMPLES_COMMON_COMMON_H__ +#include "generation_metadata.h" + #include #include #include @@ -180,6 +182,7 @@ struct SDContextParams { bool sage_attn = false; bool diffusion_conv_direct = false; bool vae_conv_direct = false; + bool taesd_preview_only = false; prediction_t prediction = PREDICTION_COUNT; lora_apply_mode_t lora_apply_mode = LORA_APPLY_AUTO; @@ -197,7 +200,7 @@ struct SDContextParams { bool validate(SDMode mode); bool resolve_and_validate(SDMode mode); std::string to_string() const; - sd_ctx_params_t to_sd_ctx_params_t(bool taesd_preview); + sd_ctx_params_t to_sd_ctx_params_t(); }; struct SDGenerationParams { @@ -337,13 +340,4 @@ struct SDGenerationParams { }; std::string version_string(); -std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params, - const SDGenerationParams& gen_params, - int64_t seed, - SDMode mode = IMG_GEN); -std::string get_image_params(const SDContextParams& ctx_params, - const SDGenerationParams& gen_params, - int64_t seed, - SDMode mode = IMG_GEN); - #endif // __EXAMPLES_COMMON_COMMON_H__ diff --git a/examples/common/generation_metadata.cpp b/examples/common/generation_metadata.cpp new file mode 100644 index 000000000..989e91446 --- /dev/null +++ b/examples/common/generation_metadata.cpp @@ -0,0 +1,155 @@ +#include "generation_metadata.h" + +#include +#include + +#include "json.hpp" + +static void relativize_lora_paths(nlohmann::json& info, const std::filesystem::path& root) { + if (info.contains("loras")) { + for (auto& lora : info["loras"]) { + const auto path = lora.value("path", ""); + if (path.empty()) + continue; + std::error_code error; + const auto absolute = std::filesystem::absolute(std::filesystem::u8path(path), error); + if (error) + continue; + const auto relative = absolute.lexically_normal().lexically_relative(root); + if (!relative.empty()) + lora["path"] = relative.generic_u8string(); + } + } + if (info.contains("inpaint_params")) + relativize_lora_paths(info["inpaint_params"], root); + if (info.contains("postprocess")) { + for (auto& operation : info["postprocess"]) + relativize_lora_paths(operation, root); + } +} + +std::string generation_info_with_relative_lora_paths(const std::string& info, const std::string& lora_model_dir) { + if (info.empty() || lora_model_dir.empty()) + return info; + std::error_code error; + const auto root = std::filesystem::absolute(std::filesystem::u8path(lora_model_dir), error); + if (error) + return info; + auto result = nlohmann::json::parse(info); + relativize_lora_paths(result, root.lexically_normal()); + return result.dump(); +} + +std::string merge_generation_info(const std::string& source, const std::string& operation) { + if (source.empty()) + return operation; + auto result = nlohmann::json::parse(source); + const auto next = nlohmann::json::parse(operation); + result["postprocess"].push_back(next); + return result.dump(); +} + +int64_t generation_seed(const std::string& info, int64_t fallback) { + if (info.empty()) + return fallback; + const auto value = nlohmann::json::parse(info); + const auto& params = value.contains("inpaint_params") ? value.at("inpaint_params") : value; + return params.contains("seed") ? std::stoll(params.at("seed").get()) : fallback; +} + +std::string generation_info_to_webui(const std::string& text) { + using nlohmann::json; + if (text.empty()) + return ""; + const auto info = json::parse(text); + if (info.value("schema", "") != "sdcpp.generation/v1") + return ""; + // Standalone postprocessing has no source diffusion settings to import. + if (!info.contains("sample_params")) + return ""; + const auto& sampling = info.at("sample_params"); + std::string result = info.value("prompt", ""); + if (info.contains("loras")) { + for (const auto& lora : info.at("loras")) { + // WebUI tags cannot restrict a LoRA to the high-noise model. + if (lora.value("is_high_noise", false)) + continue; + auto name = lora.at("path").get(); + const auto separator = name.find_last_of("/\\"); + if (separator != std::string::npos) + name.erase(0, separator + 1); + const auto extension = name.find_last_of('.'); + if (extension != std::string::npos && extension > 0) + name.erase(extension); + if (name.empty() || name.find_first_of(":<>\r\n") != std::string::npos) + continue; + if (!result.empty()) + result += " "; + result += ""; + } + } + const auto negative = info.value("negative_prompt", ""); + if (!negative.empty()) + result += "\nNegative prompt: " + negative; + result += "\n"; + bool first = true; + const auto field = [&](const char* key, const json& value) { + std::string v = value.is_string() ? value.get() : value.dump(); + if (v.find_first_of(",\n\r:\"") != std::string::npos) + v = json(v).dump(-1, ' ', false); + if (!first) + result += ", "; + first = false; + result += std::string(key) + ": " + v; + }; + field("Steps", sampling.contains("custom_sigmas") ? json(sampling.at("custom_sigmas").size() - 1) : sampling.at("sample_steps")); + static const std::map samplers = { + {"euler", "Euler"}, {"euler_a", "Euler a"}, {"heun", "Heun"}, {"dpm2", "DPM2"}, {"dpm++2s_a", "DPM++ 2S a"}, {"dpm++2m", "DPM++ 2M"}, {"dpm++2m_sde", "DPM++ 2M SDE"}, {"lcm", "LCM"}}; + const auto method = sampling.at("sample_method").get(); + auto mapped = samplers.find(method); + field(mapped != samplers.end() ? "Sampler" : "SDCPP Sampler", mapped != samplers.end() ? mapped->second : method); + static const std::map schedules = { + {"discrete", "Normal"}, {"karras", "Karras"}, {"exponential", "Exponential"}, {"sgm_uniform", "SGM Uniform"}, {"simple", "Simple"}, {"ddim_uniform", "DDIM"}, {"beta", "Beta"}}; + if (sampling.contains("scheduler")) { + const auto scheduler = sampling.at("scheduler").get(); + auto mapped_schedule = schedules.find(scheduler); + field(mapped_schedule != schedules.end() ? "Schedule type" : "SDCPP Schedule", mapped_schedule != schedules.end() ? mapped_schedule->second : scheduler); + } + field("CFG scale", sampling.at("guidance").at("txt_cfg")); + field("Seed", info.at("seed")); + field("Size", std::to_string(info.at("width").get()) + "x" + std::to_string(info.at("height").get())); + if (info.contains("models")) { + const auto& models = info.at("models"); + auto model = models.value("diffusion_model", models.value("model", "")); + const auto extension = model.find_last_of('.'); + if (extension != std::string::npos && extension > 0) + model.erase(extension); + if (!model.empty()) + field("Model", model); + if (models.contains("vae")) + field("VAE", models.at("vae")); + } + if (sampling.contains("eta")) + field("Eta", sampling.at("eta")); + if (info.value("clip_skip", 0) > 0) + field("Clip skip", info.at("clip_skip")); + if (info.contains("strength")) + field("Denoising strength", info.at("strength")); + if (info.contains("hires") && info.at("hires").value("enabled", false)) { + const auto& hires = info.at("hires"); + field("Hires resize", std::to_string(hires.at("target_width").get()) + "x" + std::to_string(hires.at("target_height").get())); + if (hires.value("steps", 0) > 0) + field("Hires steps", hires.at("steps")); + field("Denoising strength", hires.at("denoising_strength")); + field("Hires upscaler", hires.at("upscaler")); + } + field("Version", "stable-diffusion.cpp"); + return result; +} + +std::string generation_info_to_image_metadata(const std::string& info) { + if (info.empty()) + return ""; + const auto parameters = generation_info_to_webui(info); + return parameters.empty() ? info : parameters + ", SDCPP: " + info; +} diff --git a/examples/common/generation_metadata.h b/examples/common/generation_metadata.h new file mode 100644 index 000000000..7c9e9ffd4 --- /dev/null +++ b/examples/common/generation_metadata.h @@ -0,0 +1,13 @@ +#ifndef __EXAMPLES_COMMON_GENERATION_METADATA_H__ +#define __EXAMPLES_COMMON_GENERATION_METADATA_H__ + +#include +#include + +std::string merge_generation_info(const std::string& source, const std::string& operation); +std::string generation_info_with_relative_lora_paths(const std::string& info, const std::string& lora_model_dir); +int64_t generation_seed(const std::string& info, int64_t fallback = -1); +std::string generation_info_to_webui(const std::string& text); +std::string generation_info_to_image_metadata(const std::string& info); + +#endif // __EXAMPLES_COMMON_GENERATION_METADATA_H__ diff --git a/examples/common/media_io.cpp b/examples/common/media_io.cpp index baac6d200..f09c2166e 100644 --- a/examples/common/media_io.cpp +++ b/examples/common/media_io.cpp @@ -1,5 +1,7 @@ #include "media_io.h" +#include "generation_metadata.h" #include "log.h" +#include "media_metadata.h" #include "resource_owners.hpp" #include @@ -316,23 +318,22 @@ uint8_t* decode_webp_image_to_buffer(const uint8_t* data, return output.release(); } -std::string build_webp_xmp_packet(const std::string& parameters) { - if (parameters.empty()) { - return ""; - } - - const std::string escaped_parameters = xml_escape(parameters); - return "\n" - "\n" - " \n" - " \n" - " " + - escaped_parameters + - "\n" - " \n" - " \n" - "\n" - ""; +static bool embed_webp_metadata(std::vector& image, const std::string& parameters) { + if (parameters.empty()) + return true; + WebPData data{image.data(), image.size()}; + WebPMuxPtr mux(WebPMuxCreate(&data, 1)); + if (!mux) + return false; + const auto xmp = generation_xmp(parameters); + WebPData xmp_data{reinterpret_cast(xmp.data()), xmp.size()}; + if (WebPMuxSetChunk(mux.get(), "XMP ", &xmp_data, 1) != WEBP_MUX_OK) + return false; + WebPDataGuard assembled; + if (WebPMuxAssemble(mux.get(), &assembled.data) != WEBP_MUX_OK) + return false; + image.assign(assembled.data.bytes, assembled.data.bytes + assembled.data.size); + return true; } bool encode_webp_image_to_vector(const uint8_t* image, @@ -376,37 +377,7 @@ bool encode_webp_image_to_vector(const uint8_t* image, out.assign(encoded.get(), encoded.get() + encoded_size); - if (parameters.empty()) { - return true; - } - - WebPData image_data; - WebPDataInit(&image_data); - WebPDataGuard assembled_data; - - image_data.bytes = out.data(); - image_data.size = out.size(); - - WebPMuxPtr mux(WebPMuxNew()); - if (mux == nullptr) { - return false; - } - - const std::string xmp_packet = build_webp_xmp_packet(parameters); - WebPData xmp_data; - WebPDataInit(&xmp_data); - xmp_data.bytes = reinterpret_cast(xmp_packet.data()); - xmp_data.size = xmp_packet.size(); - - const bool ok = WebPMuxSetImage(mux.get(), &image_data, 1) == WEBP_MUX_OK && - WebPMuxSetChunk(mux.get(), "XMP ", &xmp_data, 1) == WEBP_MUX_OK && - WebPMuxAssemble(mux.get(), &assembled_data.data) == WEBP_MUX_OK; - - if (ok) { - out.assign(assembled_data.data.bytes, assembled_data.data.bytes + assembled_data.data.size); - } - - return ok; + return embed_webp_metadata(out, generation_info_to_image_metadata(parameters)); } #ifdef SD_USE_WEBM @@ -739,6 +710,7 @@ std::vector encode_image_to_vector(EncodedImageFormat format, const std::string& parameters, int quality) { std::vector buffer; + const auto metadata = format == EncodedImageFormat::WEBP ? std::string() : generation_info_to_image_metadata(parameters); auto write_func = [&buffer](void* context, void* data, int size) { (void)context; @@ -761,7 +733,7 @@ std::vector encode_image_to_vector(EncodedImageFormat format, result = stbi_write_jpg_to_func(c_func, &ctx, width, height, channels, image, quality); break; case EncodedImageFormat::PNG: - result = stbi_ext_write_png_to_func(c_func, &ctx, width, height, channels, image, width * channels, parameters.empty() ? nullptr : parameters.c_str()); + result = stbi_ext_write_png_to_func(c_func, &ctx, width, height, channels, image, width * channels, metadata.empty() ? nullptr : metadata.c_str()); break; case EncodedImageFormat::WEBP: #ifdef SD_USE_WEBP @@ -779,9 +751,12 @@ std::vector encode_image_to_vector(EncodedImageFormat format, break; } - if (!result) { - buffer.clear(); + if (result && format == EncodedImageFormat::JPEG && !embed_jpeg_metadata(buffer, metadata)) { + LOG_ERROR("generation metadata exceeds the container capacity or the encoded image is invalid"); + result = 0; } + if (!result) + buffer.clear(); return buffer; } @@ -794,18 +769,8 @@ bool write_image_to_file(const std::string& path, int quality) { const EncodedImageFormat format = encoded_image_format_from_path(path); - switch (format) { - case EncodedImageFormat::JPEG: - return stbi_write_jpg(path.c_str(), width, height, channels, image, quality, parameters.empty() ? nullptr : parameters.c_str()) != 0; - case EncodedImageFormat::PNG: - return stbi_write_png(path.c_str(), width, height, channels, image, 0, parameters.empty() ? nullptr : parameters.c_str()) != 0; - case EncodedImageFormat::WEBP: { - const std::vector encoded = encode_image_to_vector(format, image, width, height, channels, parameters, quality); - return !encoded.empty() && write_binary_file_bytes(path, encoded); - } - default: - return false; - } + const auto encoded = encode_image_to_vector(format, image, width, height, channels, parameters, quality); + return !encoded.empty() && write_binary_file_bytes(path, encoded); } uint8_t* load_image_from_file(const char* image_path, @@ -1384,7 +1349,10 @@ std::vector create_video_from_sd_images_to_vector(const std::string& ou #ifdef SD_USE_WEBP if (format == "webp") { - return create_animated_webp_from_sd_images_to_vector(images, num_images, fps, quality); + auto data = create_animated_webp_from_sd_images_to_vector(images, num_images, fps, quality); + if (!data.empty() && !embed_webp_metadata(data, parameters)) + return {}; + return data; } #endif diff --git a/examples/common/media_io.h b/examples/common/media_io.h index b2b56eccd..18042a8ba 100644 --- a/examples/common/media_io.h +++ b/examples/common/media_io.h @@ -16,6 +16,7 @@ enum class EncodedImageFormat { EncodedImageFormat encoded_image_format_from_path(const std::string& path); +// parameters is the core sdcpp.generation/v1 JSON; empty disables both metadata formats. std::vector encode_image_to_vector(EncodedImageFormat format, const uint8_t* image, int width, diff --git a/examples/common/media_metadata.cpp b/examples/common/media_metadata.cpp new file mode 100644 index 000000000..d5e9af3fa --- /dev/null +++ b/examples/common/media_metadata.cpp @@ -0,0 +1,47 @@ +#include "media_metadata.h" + +std::string generation_xmp(const std::string& info) { + std::string escaped; + for (char c : info) { + switch (c) { + case '&': + escaped += "&"; + break; + case '<': + escaped += "<"; + break; + case '>': + escaped += ">"; + break; + case '"': + escaped += """; + break; + case '\'': + escaped += "'"; + break; + default: + escaped += c; + break; + } + } + return "" + "" + "" + + escaped + ""; +} + +bool embed_jpeg_metadata(std::vector& image, const std::string& parameters) { + if (parameters.empty()) + return true; + const char label[] = "parameters"; + const size_t overhead = sizeof(label) + 1; + if (image.size() < 2 || parameters.size() > 65533 - overhead) + return false; + const size_t length = overhead + parameters.size() + 2; + std::vector segment = {0xff, 0xfe, static_cast(length >> 8), static_cast(length)}; + segment.insert(segment.end(), label, label + sizeof(label)); + segment.insert(segment.end(), parameters.begin(), parameters.end()); + segment.push_back(0); + image.insert(image.begin() + 2, segment.begin(), segment.end()); + return true; +} diff --git a/examples/common/media_metadata.h b/examples/common/media_metadata.h new file mode 100644 index 000000000..1843dee19 --- /dev/null +++ b/examples/common/media_metadata.h @@ -0,0 +1,11 @@ +#ifndef __EXAMPLES_COMMON_MEDIA_METADATA_H__ +#define __EXAMPLES_COMMON_MEDIA_METADATA_H__ + +#include +#include +#include + +std::string generation_xmp(const std::string& parameters); +bool embed_jpeg_metadata(std::vector& image, const std::string& parameters); + +#endif // __EXAMPLES_COMMON_MEDIA_METADATA_H__ diff --git a/examples/common/resource_owners.hpp b/examples/common/resource_owners.hpp index 73f8bdd81..8a8f84314 100644 --- a/examples/common/resource_owners.hpp +++ b/examples/common/resource_owners.hpp @@ -8,6 +8,7 @@ #include #include +#include "generation_metadata.h" #include "stable-diffusion.h" struct FreeDeleter { @@ -175,6 +176,7 @@ class SDAudioOwner { class SDImageVec { private: std::vector images_; + std::vector generation_infos_; public: SDImageVec() = default; @@ -183,13 +185,14 @@ class SDImageVec { SDImageVec& operator=(const SDImageVec&) = delete; SDImageVec(SDImageVec&& other) noexcept - : images_(std::move(other.images_)) { + : images_(std::move(other.images_)), generation_infos_(std::move(other.generation_infos_)) { } SDImageVec& operator=(SDImageVec&& other) noexcept { if (this != &other) { clear(); - images_ = std::move(other.images_); + images_ = std::move(other.images_); + generation_infos_ = std::move(other.generation_infos_); } return *this; } @@ -210,6 +213,41 @@ class SDImageVec { images_.reserve(count); } + const std::string& generation_info(size_t index) const { + static const std::string empty; + return index < generation_infos_.size() ? generation_infos_[index] : empty; + } + + void append_generation_info(size_t index, const std::string& info, const std::string& lora_model_dir = {}) { + generation_infos_.resize(images_.size()); + generation_infos_.at(index) = merge_generation_info(generation_infos_.at(index), generation_info_with_relative_lora_paths(info, lora_model_dir)); + } + + void adopt(sd_image_generation_result_t* result, const std::string& lora_model_dir = {}) { + clear(); + if (!result) + return; + for (int i = 0; i < result->count; ++i) { + images_.push_back(result->items[i].image); + result->items[i].image = {}; + generation_infos_.push_back(generation_info_with_relative_lora_paths(result->items[i].generation_info, lora_model_dir)); + } + free_sd_image_generation_result(result); + } + + void adopt_video(sd_video_generation_result_t* result, const std::string& lora_model_dir = {}) { + clear(); + if (!result) + return; + const auto info = generation_info_with_relative_lora_paths(result->generation_info, lora_model_dir); + for (int i = 0; i < result->frame_count; ++i) { + images_.push_back(result->frames[i]); + result->frames[i] = {}; + generation_infos_.push_back(info); + } + free_sd_video_generation_result(result); + } + void adopt(sd_image_t* images, int count) { clear(); if (images == nullptr || count <= 0) { @@ -270,6 +308,7 @@ class SDImageVec { image.data = nullptr; } images_.clear(); + generation_infos_.clear(); } }; diff --git a/examples/server/CMakeLists.txt b/examples/server/CMakeLists.txt index 58260be0d..d01f7a7e2 100644 --- a/examples/server/CMakeLists.txt +++ b/examples/server/CMakeLists.txt @@ -64,8 +64,10 @@ endif() add_executable(${TARGET} ../common/common.cpp + ../common/generation_metadata.cpp ../common/log.cpp ../common/media_io.cpp + ../common/media_metadata.cpp main.cpp runtime.cpp model_catalog.cpp diff --git a/examples/server/api.md b/examples/server/api.md index bfa98f6f9..1258b1af5 100644 --- a/examples/server/api.md +++ b/examples/server/api.md @@ -303,11 +303,15 @@ new model architectures. Optional `context` settings apply to each model load and override the matching startup setting. Supported fields are `vae_format` (`auto`, `flux`, `sd3`, -`flux2`, or `wan`) and `model_args` (a string in the same format as -`--model-args`; replaces the startup string). Other startup settings remain -unchanged. Switching to a recipe that omits these fields restores the startup -values. For example, PiD uses `{"vae_format": "flux"}` and Qwen Image Edit -2511 uses `{"model_args": "qwen_image_zero_cond_t=true"}`. +`flux2`, or `wan`), `model_args` (a string in the same format as +`--model-args`; replaces the startup string), and `backend` / `params_backend` +(strings in the same format as `--backend` / `--params-backend`; replace the +startup assignments for that load, so a recipe can place components on +specific devices, e.g. keep a large text encoder on CPU while another recipe +uses the GPU). Other startup settings remain unchanged. Switching to a +recipe that omits these fields restores the startup values. For example, PiD +uses `{"vae_format": "flux"}` and Qwen Image Edit 2511 uses +`{"model_args": "qwen_image_zero_cond_t=true"}`. Supported defaults: @@ -443,10 +447,6 @@ image establishes the generation dimensions if `size` is omitted. Input geometry follows `image_preprocess`: references preserve their dimensions by default, while init and mask use the generation canvas preset. -When the request starts without preconfigured reference images, the first uploaded image becomes the img2img init image if the loaded model does not support reference-image conditioning, the final reference list is empty, a mask is present, or an explicit numeric `strength` override is provided. Otherwise, only the final reference list is used for image conditioning, without initializing the img2img latent. Requests with preconfigured references retain their existing init-image selection. - -Routing uses the reference list and mask after applying `sd_cpp_extra_args`. Both `ref_images: []` and `ref_images: null` clear the reference list; `mask_image: null` clears an uploaded mask. An explicit `init_image` in `sd_cpp_extra_args`, including `null`, takes precedence over the implicit init-image selection. - Reference encoding then follows model presets and `ref_image_args`. To skip input geometry for references and disable resizing before VAE encoding, include this in `prompt`: @@ -873,7 +873,7 @@ Shared default fields used by both `img_gen` and `vid_gen`: `vae_tiling_params.tile_size_w` and `tile_size_h` are in **image pixels**, with `0` selecting the 256-pixel default. Both encode and decode use these sizes without an encoding multiplier. Positive `rel_size_w`/`rel_size_h` values override the corresponding absolute size: values up to 1 are dimension fractions, and values greater than 1 are target tile counts. Set `enabled` to use spatial tiling. Sizes are aligned down to the VAE scale factor and capped at the input dimensions; explicit sizes below the minimum supported tile size are rejected. These fields previously used latent units; see [VAE tiling](../../docs/performance.md#use-vae-tiling-to-reduce-encode-and-decode-memory-usage) for migration and OOM retry behavior. -`vae_tiling_params.extra_tiling_args` accepts a key=value list. Supported video VAEs accept `temporal_tile_frames` (alias `temporal_tile_size`, default `4`) and `temporal_tile_overlap` (default `1`). +`vae_tiling_params.extra_tiling_args` accepts a key=value list. Supported video VAEs accept `temporal_tile_frames` (alias `temporal_tile_size`, default `4`) and `temporal_tile_overlap` (default `1`). Wan2.2 TI2V 5B uses `temporal_decode_frames` (default `1`) to decode multiple latent frames per forward pass, trading additional memory for throughput; other model versions ignore this option. LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEHV use overlap blending. MiniMax H3 keeps its model-specific fixed temporal windows because its latent-to-frame mapping is non-linear. `img_gen`-specific default fields: @@ -902,6 +902,7 @@ LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEH | --- | --- | | `video_frames` | `integer` | | `fps` | `integer` | +| `ref_image_args` | `string` | | `moe_boundary` | `number` | | `vace_strength` | `number` | | `high_noise_sample_params` | `object` | @@ -1033,7 +1034,11 @@ Notes: #### `POST /sdcpp/v1/jobs/{id}/cancel` -Attempts to cancel an accepted job. +Cancels a queued job immediately, or requests interruption of a generating image or +video job. A generating job returns `200 OK` with `cancel_requested: true` and +remains `generating` until execution stops at a cancellation checkpoint. Continue +polling until its status becomes `cancelled`. Cancelled jobs have no result. +Repeated cancellation requests are safe; terminal jobs are returned unchanged. Typical status codes: @@ -1395,7 +1400,7 @@ Compared with `img_gen`, the `vid_gen` request body: - `vid_gen` is a single video sequence job, so `batch_count` is not part of the request schema - `mask_image`, `control_image`, `control_strength`, `ip_adapter_image`, `ip_adapter_strength`, and `embed_image_metadata` are not part of the request schema - `ref_images` is accepted for MiniMax-H3 Ref2VA conditioning; other video model families currently ignore it -- `vid_gen` adds `end_image`, `control_frames`, `high_noise_sample_params`, `video_frames`, `fps`, `moe_boundary`, and `vace_strength` +- `vid_gen` adds `end_image`, `control_frames`, `high_noise_sample_params`, `video_frames`, `fps`, `ref_image_args`, `moe_boundary`, and `vace_strength` Example: @@ -1410,6 +1415,7 @@ Example: "seed": -1, "video_frames": 33, "fps": 16, + "ref_image_args": "", "moe_boundary": 0.875, "vace_strength": 1.0, @@ -1491,6 +1497,7 @@ Example: - Images retain request order and correspond to ``, ``, and so on in the prompt. - For MiniMax-H3, non-empty `ref_images` cannot be combined with `init_image` or `end_image`. A conflicting API request fails during generation; the WebUI checks this before submission. - MiniMax-H3 does not support `control_frames`; leave that array empty. +- `ref_image_args` applies to `vid_gen`; MiniMax-H3 honors `ref_image_size` there (`match` default, `max` for a 2048 px short edge). Only references above the limit are downscaled; each axis is then snapped to the nearest 32 px multiple. - Other video model families currently ignore `ref_images`. - `features_by_mode.vid_gen` does not currently advertise `ref_images`; its absence is not an indication that MiniMax-H3 reference images are unsupported. @@ -1541,6 +1548,7 @@ Top-level scalar fields: | `seed` | `integer` | | `video_frames` | `integer` | | `fps` | `integer` | +| `ref_image_args` | `string` | | `moe_boundary` | `number` | | `vace_strength` | `number` | | `preview` | `string` | @@ -1768,3 +1776,14 @@ Example cancelled job: - `400 Bad Request` for an empty body, unsupported model mode, invalid JSON, invalid generation parameters, or an unsupported output format - `429 Too Many Requests` when the job queue is full - `500 Internal Server Error` for unexpected server exceptions during submission + +## Actual generation parameters + +Generation responses expose the core `sdcpp.generation/v1` record as +`generation_info`, independently of `embed_image_metadata`. Native async image +result items and OpenAI image items contain the parameter snapshot with each +image's actual seed and without `batch_count`. Native video +results contain the video request snapshot. SD API `info.generation_info` is an array +aligned with `info.infotexts` and `info.all_seeds`; the outer `parameters` field +remains the request echo. See [generation metadata](../../docs/generation_metadata.md) +for field semantics, postprocessing and container storage. diff --git a/examples/server/async_jobs.cpp b/examples/server/async_jobs.cpp index edd93593c..7a8010789 100644 --- a/examples/server/async_jobs.cpp +++ b/examples/server/async_jobs.cpp @@ -3,13 +3,46 @@ #include "async_jobs.h" #include +#include #include #include +#include #include "common/log.h" #include "common/media_io.h" #include "common/resource_owners.hpp" +class JobCancellationMonitor { + std::mutex mutex_; + std::condition_variable cv_; + bool done_ = false; + std::thread thread_; + +public: + // The caller holds sd_ctx_mutex until this monitor has joined, keeping the + // context alive and preventing cancellation from reaching another request. + JobCancellationMonitor(sd_ctx_t* ctx, const AsyncGenerationJob& job) + : thread_([this, ctx, &job]() { + std::unique_lock lock(mutex_); + while (!done_) { + if (job.cancel_requested.load()) { + // Each generation pass resets the library flag, including ADetailer. + sd_cancel_generation(ctx, SD_CANCEL_ALL); + } + cv_.wait_for(lock, std::chrono::milliseconds(10), [this]() { return done_; }); + } + }) {} + + ~JobCancellationMonitor() { + { + std::lock_guard lock(mutex_); + done_ = true; + } + cv_.notify_one(); + thread_.join(); + } +}; + static void preview_callback(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data) { (void)is_noisy; if (frame_count <= 0 || frames == nullptr || frames[0].data == nullptr) { @@ -142,9 +175,11 @@ bool cancel_queued_job(AsyncJobManager& manager, AsyncGenerationJob& job) { } manager.queue.erase(new_end, manager.queue.end()); + job.cancel_requested.store(true); job.status = AsyncJobStatus::Cancelled; job.completed_at = unix_timestamp_now_ms(); job.result_images_b64.clear(); + job.result_generation_info = nullptr; job.result_media_b64.clear(); job.result_media_mime_type.clear(); job.result_frame_count = 0; @@ -156,13 +191,14 @@ bool cancel_queued_job(AsyncJobManager& manager, AsyncGenerationJob& job) { json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJob& job) { json result; - result["id"] = job.id; - result["kind"] = async_job_kind_name(job.kind); - result["status"] = async_job_status_name(job.status); - result["created"] = job.created_at; - result["started"] = job.started_at == 0 ? json(nullptr) : json(job.started_at); - result["completed"] = job.completed_at == 0 ? json(nullptr) : json(job.completed_at); - result["queue_position"] = 0; + result["id"] = job.id; + result["kind"] = async_job_kind_name(job.kind); + result["status"] = async_job_status_name(job.status); + result["cancel_requested"] = job.cancel_requested.load(); + result["created"] = job.created_at; + result["started"] = job.started_at == 0 ? json(nullptr) : json(job.started_at); + result["completed"] = job.completed_at == 0 ? json(nullptr) : json(job.completed_at); + result["queue_position"] = 0; if (job.status == AsyncJobStatus::Queued) { size_t position = 1; @@ -183,11 +219,12 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo {"fps", job.result_fps}, {"frame_count", job.result_frame_count}, {"b64_json", job.result_media_b64}, + {"generation_info", job.result_generation_info}, }; } else { json images = json::array(); for (size_t i = 0; i < job.result_images_b64.size(); ++i) { - images.push_back({{"index", i}, {"b64_json", job.result_images_b64[i]}}); + images.push_back({{"index", i}, {"b64_json", job.result_images_b64[i]}, {"generation_info", job.result_generation_info.at(i)}}); } result["result"] = { {"output_format", job.img_gen.output_format}, @@ -199,11 +236,11 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo job.status == AsyncJobStatus::Cancelled) { result["result"] = nullptr; result["error"] = { - {"code", + {"code", job.error_code.empty() - ? (job.status == AsyncJobStatus::Cancelled ? "cancelled" : "generation_failed") - : job.error_code}, - {"message", job.error_message}, + ? (job.status == AsyncJobStatus::Cancelled ? "cancelled" : "generation_failed") + : job.error_code}, + {"message", job.error_message}, }; } else { result["result"] = nullptr; @@ -228,22 +265,27 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo bool execute_img_gen_job(ServerRuntime& runtime, AsyncGenerationJob& job, std::vector& output_images, + json& output_info, std::string& error_message) { + output_info = json::array(); sd_img_gen_params_t params = job.img_gen.to_sd_img_gen_params_t(); SDImageVec results; { std::lock_guard lock(*runtime.sd_ctx_mutex); - set_preview_callback(job); - sd_image_t* raw_results = nullptr; - int num_results = 0; - if (!generate_image(runtime.sd_ctx, ¶ms, &raw_results, &num_results)) { - raw_results = nullptr; - num_results = 0; + if (job.cancel_requested.load()) { + return false; } + JobCancellationMonitor cancellation(runtime.sd_ctx, job); + set_preview_callback(job); + sd_image_generation_result_t* generated = nullptr; + generate_image(runtime.sd_ctx, ¶ms, &generated); + results.adopt(generated, runtime.ctx_params->lora_model_dir); clear_preview_callback(); - results.adopt(raw_results, num_results); + if (job.cancel_requested.load()) { + return false; + } if (!results.empty() && !apply_server_adetailer(runtime, job.img_gen, params, results)) { error_message = "ADetailer failed"; return false; @@ -251,7 +293,7 @@ bool execute_img_gen_job(ServerRuntime& runtime, } const int num_results = results.count(); - if (num_results <= 0) { + if (job.cancel_requested.load() || num_results <= 0) { error_message = "generate_image returned no results"; return false; } @@ -263,17 +305,13 @@ bool execute_img_gen_job(ServerRuntime& runtime, encoded_format = EncodedImageFormat::WEBP; } - int batch_count = job.img_gen.gen_params.batch_count; - int images_per_batch = batch_count > 0 ? std::max(1, num_results / batch_count) : 1; for (int i = 0; i < num_results; ++i) { if (results[i].data == nullptr) { continue; } const std::string metadata = job.img_gen.gen_params.embed_image_metadata - ? get_image_params(*runtime.ctx_params, - job.img_gen.gen_params, - job.img_gen.gen_params.seed + i / images_per_batch) + ? results.generation_info(i) : ""; auto image_bytes = encode_image_to_vector(encoded_format, results[i].data, @@ -286,6 +324,7 @@ bool execute_img_gen_job(ServerRuntime& runtime, continue; } output_images.push_back(base64_encode(image_bytes)); + output_info.push_back(json::parse(results.generation_info(i))); } if (output_images.empty()) { @@ -302,11 +341,10 @@ bool execute_vid_gen_job(ServerRuntime& runtime, std::string& output_media_mime_type, int& output_frame_count, int& output_fps, + json& output_info, std::string& error_message) { sd_vid_gen_params_t params = job.vid_gen.to_sd_vid_gen_params_t(); - std::string str_params = job.vid_gen.gen_params.embed_image_metadata - ? get_image_params(*runtime.ctx_params, job.vid_gen.gen_params, job.vid_gen.gen_params.seed, VID_GEN) - : ""; + std::string str_params; SDImageVec results; int num_results = 0; @@ -314,17 +352,26 @@ bool execute_vid_gen_job(ServerRuntime& runtime, { std::lock_guard lock(*runtime.sd_ctx_mutex); + if (job.cancel_requested.load()) { + return false; + } + JobCancellationMonitor cancellation(runtime.sd_ctx, job); set_preview_callback(job); - sd_image_t* raw_results = nullptr; - if (!generate_video(runtime.sd_ctx, ¶ms, &raw_results, &num_results, &generated_audio, &output_fps)) { - raw_results = nullptr; + sd_video_generation_result_t* generated = nullptr; + if (generate_video(runtime.sd_ctx, ¶ms, &generated)) { + output_fps = generated->fps; + generated_audio = generated->audio; + generated->audio = nullptr; + results.adopt_video(generated, runtime.ctx_params->lora_model_dir); + output_info = json::parse(results.generation_info(0)); + if (job.vid_gen.gen_params.embed_image_metadata) + str_params = results.generation_info(0); } clear_preview_callback(); - results.adopt(raw_results, num_results); } num_results = results.count(); - if (num_results <= 0) { + if (job.cancel_requested.load() || num_results <= 0) { free_sd_audio(generated_audio); error_message = "generate_video returned no results"; return false; @@ -381,6 +428,7 @@ void async_job_worker(ServerRuntime& runtime) { } std::vector output_images; + json output_info; std::string output_media_b64; std::string output_media_mime_type; int output_frame_count = 0; @@ -389,7 +437,7 @@ void async_job_worker(ServerRuntime& runtime) { bool ok = false; if (job->kind == AsyncJobKind::ImgGen) { - ok = execute_img_gen_job(runtime, *job, output_images, error_message); + ok = execute_img_gen_job(runtime, *job, output_images, output_info, error_message); } else if (job->kind == AsyncJobKind::VidGen) { ok = execute_vid_gen_job(runtime, *job, @@ -397,6 +445,7 @@ void async_job_worker(ServerRuntime& runtime) { output_media_mime_type, output_frame_count, output_fps, + output_info, error_message); } else { error_message = "unsupported job kind"; @@ -410,9 +459,14 @@ void async_job_worker(ServerRuntime& runtime) { } job->completed_at = unix_timestamp_now_ms(); - if (ok) { + if (job->cancel_requested.load()) { + job->status = AsyncJobStatus::Cancelled; + job->error_code = "cancelled"; + job->error_message = "job cancelled by client"; + } else if (ok) { job->status = AsyncJobStatus::Completed; job->result_images_b64 = std::move(output_images); + job->result_generation_info = std::move(output_info); job->result_media_b64 = std::move(output_media_b64); job->result_media_mime_type = std::move(output_media_mime_type); job->result_frame_count = output_frame_count; @@ -424,6 +478,7 @@ void async_job_worker(ServerRuntime& runtime) { job->error_code = "generation_failed"; job->error_message = error_message.empty() ? "unknown generation error" : error_message; job->result_images_b64.clear(); + job->result_generation_info = nullptr; job->result_media_b64.clear(); job->result_media_mime_type.clear(); job->result_frame_count = 0; diff --git a/examples/server/async_jobs.h b/examples/server/async_jobs.h index 049a9d73b..6fa3cbd13 100644 --- a/examples/server/async_jobs.h +++ b/examples/server/async_jobs.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -32,12 +33,14 @@ struct AsyncGenerationJob { std::string id; AsyncJobKind kind = AsyncJobKind::ImgGen; AsyncJobStatus status = AsyncJobStatus::Queued; - int64_t created_at = unix_timestamp_now_ms(); - int64_t started_at = 0; - int64_t completed_at = 0; + std::atomic cancel_requested{false}; + int64_t created_at = unix_timestamp_now_ms(); + int64_t started_at = 0; + int64_t completed_at = 0; ImgGenJobRequest img_gen; VidGenJobRequest vid_gen; std::vector result_images_b64; + json result_generation_info; std::string result_media_b64; std::string result_media_mime_type; int result_frame_count = 0; @@ -72,6 +75,7 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo bool execute_img_gen_job(ServerRuntime& runtime, AsyncGenerationJob& job, std::vector& output_images, + json& output_info, std::string& error_message); bool execute_vid_gen_job(ServerRuntime& runtime, AsyncGenerationJob& job, @@ -79,5 +83,6 @@ bool execute_vid_gen_job(ServerRuntime& runtime, std::string& output_media_mime_type, int& output_frame_count, int& output_fps, + json& output_info, std::string& error_message); void async_job_worker(ServerRuntime& runtime); diff --git a/examples/server/frontend b/examples/server/frontend index e36e34a2e..34e5ad6ed 160000 --- a/examples/server/frontend +++ b/examples/server/frontend @@ -1 +1 @@ -Subproject commit e36e34a2e2204e3df654b16cadc112d8ce8f3174 +Subproject commit 34e5ad6edcaf827b7fe1e9ffa55ee86cd2fc43d2 diff --git a/examples/server/main.cpp b/examples/server/main.cpp index d79c6e2f3..a10cb66f1 100644 --- a/examples/server/main.cpp +++ b/examples/server/main.cpp @@ -93,7 +93,7 @@ int main(int argc, const char** argv) { LOG_DEBUG("%s", ctx_params.to_string().c_str()); LOG_DEBUG("%s", default_gen_params.to_string().c_str()); - sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(false); + sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(); ServerModels models; if (!ctx_params.model_path.empty() || !ctx_params.diffusion_model_path.empty()) { const auto started = std::chrono::steady_clock::now(); diff --git a/examples/server/model_catalog.cpp b/examples/server/model_catalog.cpp index e8a6414c6..eef685614 100644 --- a/examples/server/model_catalog.cpp +++ b/examples/server/model_catalog.cpp @@ -19,7 +19,7 @@ struct ModelSlot { const char* id; const char* label; std::vector categories; - std::string SDContextParams::*field; + std::string SDContextParams::* field; }; static const ModelSlot k_model_slots[] = { @@ -273,7 +273,7 @@ static void apply_recipe_defaults(const json& defaults, SDGenerationParams& para } static void apply_recipe_context(const json& settings, SDContextParams& params) { - check_keys(settings, {"vae_format", "model_args"}); + check_keys(settings, {"vae_format", "model_args", "backend", "params_backend"}); if (settings.contains("vae_format")) { const auto format = settings.at("vae_format").get(); if (format != "auto" && format != "flux" && format != "sd3" && format != "flux2" && format != "wan") @@ -283,6 +283,12 @@ static void apply_recipe_context(const json& settings, SDContextParams& params) if (settings.contains("model_args")) { params.model_args = settings.at("model_args").get(); } + if (settings.contains("backend")) { + params.backend = settings.at("backend").get(); + } + if (settings.contains("params_backend")) { + params.params_backend = settings.at("params_backend").get(); + } } static const std::map k_recipe_features = { diff --git a/examples/server/openai_edit_policy.h b/examples/server/openai_edit_policy.h deleted file mode 100644 index c963a887f..000000000 --- a/examples/server/openai_edit_policy.h +++ /dev/null @@ -1,44 +0,0 @@ -#ifndef __SD_EXAMPLES_SERVER_OPENAI_EDIT_POLICY_H__ -#define __SD_EXAMPLES_SERVER_OPENAI_EDIT_POLICY_H__ - -#include "common/resource_owners.hpp" - -enum class OpenAIEditInitImageAction { - KeepExisting, - UseUploaded, - Clear, -}; - -inline bool openai_edit_uses_uploaded_image_as_init(bool supports_ref_image_conditioning, - bool has_ref_images, - bool has_mask, - bool has_explicit_strength_override) { - return !supports_ref_image_conditioning || !has_ref_images || has_mask || has_explicit_strength_override; -} - -inline OpenAIEditInitImageAction apply_openai_edit_init_image_policy( - SDImageOwner& init_image, - SDImageOwner& uploaded_image, - bool uploaded_image_was_first_reference, - bool supports_ref_image_conditioning, - bool has_ref_images, - bool has_mask, - bool has_extra_init_image_override, - bool has_explicit_strength_override) { - if (!uploaded_image_was_first_reference || has_extra_init_image_override) { - return OpenAIEditInitImageAction::KeepExisting; - } - - if (openai_edit_uses_uploaded_image_as_init(supports_ref_image_conditioning, - has_ref_images, - has_mask, - has_explicit_strength_override)) { - init_image = std::move(uploaded_image); - return OpenAIEditInitImageAction::UseUploaded; - } - - init_image.reset(); - return OpenAIEditInitImageAction::Clear; -} - -#endif // __SD_EXAMPLES_SERVER_OPENAI_EDIT_POLICY_H__ diff --git a/examples/server/routes_models.cpp b/examples/server/routes_models.cpp index ea14eb821..9e4e3963c 100644 --- a/examples/server/routes_models.cpp +++ b/examples/server/routes_models.cpp @@ -94,7 +94,7 @@ static void change_model(ServerRuntime& runtime, const httplib::Request& req, ht *runtime.ctx_params = unload ? models.base_context : std::move(context); *runtime.default_gen_params = unload ? models.base_defaults : std::move(defaults); if (!unload) { - auto params = runtime.ctx_params->to_sd_ctx_params_t(false); + auto params = runtime.ctx_params->to_sd_ctx_params_t(); models.context.reset(new_sd_ctx(¶ms)); if (!models.context) { throw std::runtime_error("Model initialization failed; check the server log for component compatibility and memory errors"); diff --git a/examples/server/routes_openai.cpp b/examples/server/routes_openai.cpp index 27a7e18f9..f721e1495 100644 --- a/examples/server/routes_openai.cpp +++ b/examples/server/routes_openai.cpp @@ -7,7 +7,6 @@ #include "common/common.h" #include "common/media_io.h" #include "common/resource_owners.hpp" -#include "openai_edit_policy.h" static std::string extract_and_remove_sd_cpp_extra_args(std::string& text) { std::regex re("(.*?)"); @@ -158,49 +157,24 @@ static bool build_openai_edit_request(const httplib::Request& req, request.gen_params.height = height; request.gen_params.batch_count = n; - std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt); - const json extra_args = sd_cpp_extra_args_str.empty() - ? json::object() - : json::parse(sd_cpp_extra_args_str, nullptr, false); - const bool has_extra_init_image_override = extra_args.is_object() && extra_args.contains("init_image"); - const bool has_extra_ref_images_override = extra_args.is_object() && extra_args.contains("ref_images"); - const bool has_extra_strength_override = extra_args.is_object() && - extra_args.contains("strength") && - extra_args["strength"].is_number(); - const bool has_extra_mask_image = extra_args.is_object() && - extra_args.contains("mask_image") && - !extra_args["mask_image"].is_null(); - const bool supports_ref_image_conditioning = sd_ctx_supports_ref_image_conditioning(runtime.sd_ctx); - const bool may_need_uploaded_init_image = !has_extra_init_image_override && - (!supports_ref_image_conditioning || !mask_bytes.empty() || - has_extra_mask_image || has_extra_strength_override || has_extra_ref_images_override); - SDImageOwner uploaded_image_for_init; - bool uploaded_image_was_first_reference = false; + std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt); for (auto& bytes : images_bytes) { int img_w = 0; int img_h = 0; int resolved_channel = 0; uint8_t* raw_pixels = load_image_from_memory( - reinterpret_cast(bytes.data()), - static_cast(bytes.size()), - img_w, img_h, resolved_channel, - 0, 0, - 0); + reinterpret_cast(bytes.data()), + static_cast(bytes.size()), + img_w, img_h, resolved_channel, + 0, 0, + 0); if (raw_pixels == nullptr) { continue; } - const bool is_first_ref_image = request.gen_params.ref_images.empty(); SDImageOwner image_owner({(uint32_t)img_w, (uint32_t)img_h, (uint32_t)resolved_channel, raw_pixels}); request.gen_params.set_width_and_height_if_unset(image_owner.get().width, image_owner.get().height); - if (is_first_ref_image) { - uploaded_image_was_first_reference = true; - if (may_need_uploaded_init_image) { - uploaded_image_for_init = image_owner; - } - } - request.gen_params.ref_images.push_back(std::move(image_owner)); } @@ -231,21 +205,6 @@ static bool build_openai_edit_request(const httplib::Request& req, return false; } - const auto init_image_action = apply_openai_edit_init_image_policy( - request.gen_params.init_image, - uploaded_image_for_init, - uploaded_image_was_first_reference, - supports_ref_image_conditioning, - !request.gen_params.ref_images.empty(), - request.gen_params.mask_image.get().data != nullptr, - has_extra_init_image_override, - has_extra_strength_override); - if (init_image_action == OpenAIEditInitImageAction::UseUploaded && - request.gen_params.init_image.get().data == nullptr) { - error_message = "could not allocate init image"; - return false; - } - // Intentionally disable prompt-embedded LoRA tag parsing for server APIs. if (!request.gen_params.resolve_and_validate(IMG_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) { error_message = "invalid params"; @@ -264,12 +223,10 @@ static bool execute_sync_img_gen_request(ServerRuntime& runtime, { std::lock_guard lock(*runtime.sd_ctx_mutex); - sd_image_t* raw_results = nullptr; - if (!generate_image(runtime.sd_ctx, &img_gen_params, &raw_results, &num_results)) { - raw_results = nullptr; - num_results = 0; - } - results.adopt(raw_results, num_results); + sd_image_generation_result_t* generated = nullptr; + generate_image(runtime.sd_ctx, &img_gen_params, &generated); + results.adopt(generated, runtime.ctx_params->lora_model_dir); + num_results = results.count(); if (!results.empty() && !apply_server_adetailer(runtime, request, img_gen_params, results)) { error_message = "ADetailer failed"; return false; @@ -326,35 +283,34 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { out["data"] = json::array(); out["output_format"] = request.output_format; - int result_count = results.count(); - int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, result_count / request.gen_params.batch_count) : 1; + int result_count = results.count(); + for (int i = 0; i < result_count; ++i) { if (results[i].data == nullptr) { continue; } std::string params = request.gen_params.embed_image_metadata - ? get_image_params(*runtime->ctx_params, - request.gen_params, - request.gen_params.seed + i / images_per_batch) + ? results.generation_info(i) : ""; auto image_bytes = encode_image_to_vector(request.output_format == "jpeg" ? EncodedImageFormat::JPEG : request.output_format == "webp" ? EncodedImageFormat::WEBP : EncodedImageFormat::PNG, - results[i].data, - results[i].width, - results[i].height, - results[i].channel, - params, - request.output_compression); + results[i].data, + results[i].width, + results[i].height, + results[i].channel, + params, + request.output_compression); if (image_bytes.empty()) { LOG_ERROR("write image to mem failed"); continue; } json item; - item["b64_json"] = base64_encode(image_bytes); + item["b64_json"] = base64_encode(image_bytes); + item["generation_info"] = json::parse(results.generation_info(i)); out["data"].push_back(item); } @@ -400,26 +356,25 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { out["data"] = json::array(); out["output_format"] = request.output_format; - int result_count = results.count(); - int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, result_count / request.gen_params.batch_count) : 1; + int result_count = results.count(); + for (int i = 0; i < result_count; ++i) { if (results[i].data == nullptr) { continue; } std::string params = request.gen_params.embed_image_metadata - ? get_image_params(*runtime->ctx_params, - request.gen_params, - request.gen_params.seed + i / images_per_batch) + ? results.generation_info(i) : ""; auto image_bytes = encode_image_to_vector(request.output_format == "jpeg" ? EncodedImageFormat::JPEG : EncodedImageFormat::PNG, - results[i].data, - results[i].width, - results[i].height, - results[i].channel, - params, - request.output_compression); + results[i].data, + results[i].width, + results[i].height, + results[i].channel, + params, + request.output_compression); json item; - item["b64_json"] = base64_encode(image_bytes); + item["b64_json"] = base64_encode(image_bytes); + item["generation_info"] = json::parse(results.generation_info(i)); out["data"].push_back(item); } diff --git a/examples/server/routes_sdapi.cpp b/examples/server/routes_sdapi.cpp index 014b0b5f8..0d37898a8 100644 --- a/examples/server/routes_sdapi.cpp +++ b/examples/server/routes_sdapi.cpp @@ -462,46 +462,27 @@ static bool build_sdapi_img_gen_request(const json& j, validate_server_img_gen_request(runtime, request, error_message); } -static nlohmann::json prepare_info_field(const SDContextParams& ctx_params, - const SDGenerationParams& gen_params, - bool img2img) { - nlohmann::json jsoninfo = nlohmann::json::object(); - jsoninfo["prompt"] = gen_params.prompt; - if (!gen_params.negative_prompt.empty()) { - jsoninfo["negative_prompt"] = gen_params.negative_prompt; - } - jsoninfo["seed"] = gen_params.seed; - jsoninfo["cfg_scale"] = gen_params.sample_params.guidance.txt_cfg; - jsoninfo["width"] = gen_params.get_resolved_width(); - jsoninfo["height"] = gen_params.get_resolved_height(); - jsoninfo["steps"] = gen_params.sample_params.sample_steps; - jsoninfo["sampler_name"] = sd_sample_method_name(gen_params.sample_params.sample_method); - if (gen_params.clip_skip != -1) { - jsoninfo["clip_skip"] = gen_params.clip_skip; - } - if (gen_params.sample_params.scheduler != scheduler_t::SCHEDULER_COUNT) { - jsoninfo["extra_generation_params"] = nlohmann::json::object(); - jsoninfo["extra_generation_params"]["Schedule type"] = sd_scheduler_name(gen_params.sample_params.scheduler); - } - if (img2img) { - jsoninfo["denoising_strength"] = gen_params.strength; - } - // not clear what should happen if we have both model and diffusion_model - if (!ctx_params.diffusion_model_path.empty()) { - jsoninfo["sd_model_name"] = sd_basename(ctx_params.diffusion_model_path); - } else if (!ctx_params.model_path.empty()) { - jsoninfo["sd_model_name"] = sd_basename(ctx_params.model_path); - } - if (!ctx_params.vae_path.empty()) { - jsoninfo["sd_vae_name"] = sd_basename(ctx_params.vae_path); - } - jsoninfo["version"] = "stable-diffusion.cpp"; - - jsoninfo["infotexts"] = nlohmann::json::array(); - jsoninfo["all_prompts"] = nlohmann::json::array(); - jsoninfo["all_negative_prompts"] = nlohmann::json::array(); - jsoninfo["all_seeds"] = nlohmann::json::array(); - return jsoninfo; +static nlohmann::json prepare_info_field(const std::string& generation_info) { + const auto info = nlohmann::json::parse(generation_info); + const auto& sampling = info.at("sample_params"); + nlohmann::json result = {{"prompt", info.at("prompt")}, + {"negative_prompt", info.at("negative_prompt")}, + {"seed", generation_seed(generation_info)}, + {"cfg_scale", sampling.at("guidance").at("txt_cfg")}, + {"width", info.at("width")}, + {"height", info.at("height")}, + {"steps", sampling.contains("custom_sigmas") ? json(sampling.at("custom_sigmas").size() - 1) : sampling.at("sample_steps")}, + {"sampler_name", sampling.at("sample_method")}, + {"version", "stable-diffusion.cpp"}}; + if (info.contains("clip_skip")) + result["clip_skip"] = info.at("clip_skip"); + if (info.contains("strength")) + result["denoising_strength"] = info.at("strength"); + if (sampling.contains("scheduler")) + result["extra_generation_params"]["Schedule type"] = sampling.at("scheduler"); + for (const auto* key : {"infotexts", "all_prompts", "all_negative_prompts", "all_seeds", "generation_info"}) + result[key] = nlohmann::json::array(); + return result; } struct SDAPIPreviewState { @@ -543,7 +524,7 @@ static void preview_callback(int step, int frame_count, sd_image_t* frames, bool std::lock_guard lock(g_preview.mutex); g_preview.b64 = std::move(b64); g_preview.step = static_cast(std::min(step < 0 ? -static_cast(step) : step, - std::max(0, info.total_steps))); + std::max(0, info.total_steps))); g_preview.total = std::max(0, info.total_steps); g_preview.job_no = std::max(0, info.sample_pass - 1); } @@ -621,12 +602,10 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { std::lock_guard lock(*runtime->sd_ctx_mutex); const bool detailing = request.adetailer_enabled && !runtime->default_gen_params->ad_model_path.empty(); SDAPIPreviewGuard preview_guard(id_task, img_gen_params, img2img ? "img2img" : "txt2img", detailing); - sd_image_t* raw_results = nullptr; - if (!generate_image(runtime->sd_ctx, &img_gen_params, &raw_results, &num_results)) { - raw_results = nullptr; - num_results = 0; - } - results.adopt(raw_results, num_results); + sd_image_generation_result_t* generated = nullptr; + generate_image(runtime->sd_ctx, &img_gen_params, &generated); + results.adopt(generated, runtime->ctx_params->lora_model_dir); + num_results = results.count(); if (detailing && !results.empty()) preview_guard.begin_detailing(); if (!results.empty() && !apply_server_adetailer(*runtime, request, img_gen_params, results)) { @@ -645,9 +624,8 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { json out; out["images"] = json::array(); out["parameters"] = j; - json jsoninfo = prepare_info_field(*runtime->ctx_params, request.gen_params, img2img); + json jsoninfo = prepare_info_field(results.generation_info(0)); - int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, num_results / request.gen_params.batch_count) : 1; for (int i = 0; i < num_results; ++i) { if (results[i].data == nullptr) { continue; @@ -655,16 +633,14 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { bool embed_meta = request.gen_params.embed_image_metadata; - std::string params = get_image_params(*runtime->ctx_params, - request.gen_params, - request.gen_params.seed + i / images_per_batch); + std::string params = generation_info_to_webui(results.generation_info(i)); auto image_bytes = encode_image_to_vector(EncodedImageFormat::PNG, results[i].data, results[i].width, results[i].height, results[i].channel, - embed_meta ? params : ""); + embed_meta ? results.generation_info(i) : ""); if (image_bytes.empty()) { LOG_ERROR("write image to mem failed"); @@ -673,10 +649,12 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { out["images"].push_back(base64_encode(image_bytes)); - jsoninfo["infotexts"][i] = params; - jsoninfo["all_seeds"][i] = request.gen_params.seed + i; - jsoninfo["all_prompts"][i] = request.gen_params.prompt; - jsoninfo["all_negative_prompts"][i] = request.gen_params.negative_prompt; + const auto actual = json::parse(results.generation_info(i)); + jsoninfo["infotexts"].push_back(params); + jsoninfo["all_seeds"].push_back(generation_seed(results.generation_info(i))); + jsoninfo["generation_info"].push_back(actual); + jsoninfo["all_prompts"].push_back(actual["prompt"]); + jsoninfo["all_negative_prompts"].push_back(actual["negative_prompt"]); } // not a mistake: it is supposed to be a string in json format diff --git a/examples/server/routes_sdcpp.cpp b/examples/server/routes_sdcpp.cpp index 8eb9c9ab3..b972768b9 100644 --- a/examples/server/routes_sdcpp.cpp +++ b/examples/server/routes_sdcpp.cpp @@ -178,6 +178,7 @@ static json make_vid_gen_defaults_json(const SDGenerationParams& defaults, const {"seed", defaults.seed}, {"video_frames", defaults.video_frames}, {"image_preprocess", defaults.image_preprocess}, + {"ref_image_args", defaults.ref_image_args}, {"fps", defaults.fps}, {"moe_boundary", defaults.moe_boundary}, {"vace_strength", defaults.vace_strength}, @@ -206,7 +207,7 @@ static json make_img_gen_features_json() { {"hires", true}, {"cache", true}, {"cancel_queued", true}, - {"cancel_generating", false}, + {"cancel_generating", true}, {"preview", true}, }; } @@ -221,7 +222,7 @@ static json make_vid_gen_features_json() { {"vae_tiling", true}, {"cache", true}, {"cancel_queued", true}, - {"cancel_generating", false}, + {"cancel_generating", true}, {"preview", true}, }; } @@ -380,8 +381,8 @@ static json make_capabilities_json(ServerRuntime& runtime) { json top_level_defaults = json::object(); json top_level_output_formats = json::array(); json top_level_features = { - {"cancel_queued", true}, - {"cancel_generating", false}, + {"cancel_queued", true}, + {"cancel_generating", true}, }; std::string current_mode = ""; if (supports_img) { @@ -409,14 +410,14 @@ static json make_capabilities_json(ServerRuntime& runtime) { result["defaults"] = top_level_defaults; result["defaults_by_mode"] = defaults_by_mode; result["limits"] = { - {"min_width", 64}, - {"max_width", k_max_generation_dimension}, - {"min_height", 64}, - {"max_height", k_max_generation_dimension}, - {"max_batch_count", 8}, - {"max_queue_size", manager.max_pending_jobs}, - {"max_upscale_width", k_max_upscale_dimension}, - {"max_upscale_height", k_max_upscale_dimension}, + {"min_width", 64}, + {"max_width", k_max_generation_dimension}, + {"min_height", 64}, + {"max_height", k_max_generation_dimension}, + {"max_batch_count", 8}, + {"max_queue_size", manager.max_pending_jobs}, + {"max_upscale_width", k_max_upscale_dimension}, + {"max_upscale_height", k_max_upscale_dimension}, }; result["samplers"] = samplers; result["schedulers"] = schedulers; @@ -519,7 +520,7 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { const bool valid = value.is_number_unsigned() ? value.get() <= static_cast(std::numeric_limits::max()) : value.is_number_integer() && value.get() >= std::numeric_limits::min() && - value.get() <= std::numeric_limits::max(); + value.get() <= std::numeric_limits::max(); if (!valid) { res.status = 400; res.set_content(json({{"error", std::string(key) + " must be a 32-bit integer"}}).dump(), "application/json"); @@ -608,19 +609,20 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { } SDImageOwner current(input.release()); + std::string generation_info; for (int i = 0; i < repeats; ++i) { - sd_image_t* out_images = nullptr; - int out_count = 0; - if (!upscale(upscaler_ctx.get(), current.get(), (uint32_t)factor, &out_images, &out_count) || - out_count <= 0 || out_images[0].data == nullptr) { - free_sd_images(out_images, out_count); + sd_image_generation_result_t* generated = nullptr; + if (!upscale(upscaler_ctx.get(), current.get(), (uint32_t)factor, &generated) || + !generated || generated->count <= 0 || generated->items[0].image.data == nullptr) { + free_sd_image_generation_result(generated); res.status = 500; res.set_content(R"({"error":"upscale failed"})", "application/json"); return; } - sd_image_t produced = out_images[0]; - out_images[0] = {0, 0, 0, nullptr}; - free_sd_images(out_images, out_count); + generation_info = merge_generation_info(generation_info, generated->items[0].generation_info); + sd_image_t produced = generated->items[0].image; + generated->items[0].image = {}; + free_sd_image_generation_result(generated); current.reset(produced); } @@ -630,12 +632,12 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { auto image_bytes = encode_image_to_vector(format == "jpeg" ? EncodedImageFormat::JPEG : format == "webp" ? EncodedImageFormat::WEBP : EncodedImageFormat::PNG, - result.data, - result.width, - result.height, - result.channel, - "", - compression); + result.data, + result.width, + result.height, + result.channel, + generation_info, + compression); if (image_bytes.empty()) { res.status = 500; res.set_content(R"({"error":"the result could not be encoded"})", "application/json"); @@ -643,13 +645,14 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { } json out; - out["upscaler"] = used_name; - out["scale"] = factor; - out["repeats"] = repeats; - out["width"] = result.width; - out["height"] = result.height; - out["output_format"] = format; - json images = json::array(); + out["generation_info"] = json::parse(generation_info); + out["upscaler"] = used_name; + out["scale"] = factor; + out["repeats"] = repeats; + out["width"] = result.width; + out["height"] = result.height; + out["output_format"] = format; + json images = json::array(); images.push_back({{"index", 0}, {"b64_json", base64_encode(image_bytes)}}); out["images"] = std::move(images); res.set_content(out.dump(), "application/json"); @@ -840,8 +843,9 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { } if (job.status == AsyncJobStatus::Generating) { - res.status = 409; - res.set_content(R"({"error":"job is currently generating and cannot be interrupted yet"})", "application/json"); + job.cancel_requested.store(true); + res.status = 200; + res.set_content(make_async_job_json(manager, job).dump(), "application/json"); return; } diff --git a/examples/server/runtime.cpp b/examples/server/runtime.cpp index 5b33b338a..b187ff85a 100644 --- a/examples/server/runtime.cpp +++ b/examples/server/runtime.cpp @@ -136,7 +136,7 @@ bool apply_server_adetailer(ServerRuntime& runtime, continue; } sd_img_gen_params_t ad_generation_params = params; - ad_generation_params.seed = params.seed + i; + ad_generation_params.seed = generation_seed(results.generation_info(i)); ad_generation_params.width = pass.width; ad_generation_params.height = pass.height; ad_generation_params.strength = pass.strength; @@ -147,24 +147,23 @@ bool apply_server_adetailer(ServerRuntime& runtime, ad_generation_params.sample_params.scheduler = SCHEDULER_COUNT; } - sd_image_t* detailed_images = nullptr; - int detailed_count = 0; + sd_image_generation_result_t* detailed = nullptr; if (!adetail_image(ad_ctx.get(), runtime.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); + &detailed) || + detailed == nullptr || detailed->count <= 0 || detailed->items[0].image.data == nullptr) { + free_sd_image_generation_result(detailed); LOG_ERROR("ADetailer failed for image %d", i + 1); return false; } + results.append_generation_info(i, detailed->items[0].generation_info, runtime.ctx_params->lora_model_dir); free(results[i].data); - results[i] = detailed_images[0]; - detailed_images[0] = {0, 0, 0, nullptr}; - free_sd_images(detailed_images, detailed_count); + results[i] = detailed->items[0].image; + detailed->items[0].image = {}; + free_sd_image_generation_result(detailed); } } return true; diff --git a/ggml b/ggml index dc998771e..a5915eb0e 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit dc998771e700d49bbdf9a39cede22bb5b3311767 +Subproject commit a5915eb0e47adef21c3f25b889a008925ea838c2 diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 85fe51d46..f4ea0eaf1 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -94,6 +94,9 @@ enum prediction_t { SEFI_FLOW_PRED, MINIT2I_FLOW_PRED, SENSENOVA_U1_FLOW_PRED, + FLUX3_FLOW_PRED, + H3_AV_FLOW_PRED, + MING_IMAGE_FLOW_PRED, PREDICTION_COUNT }; @@ -269,6 +272,25 @@ typedef struct { uint8_t* data; } sd_image_t; +typedef struct { + sd_image_t image; + // Owned UTF-8 JSON using the sdcpp.generation/v1 schema. + char* generation_info; +} sd_image_result_t; + +typedef struct { + sd_image_result_t* items; + int count; +} sd_image_generation_result_t; + +typedef struct { + sd_image_t* frames; + int frame_count; + sd_audio_t* audio; + int fps; + char* generation_info; +} sd_video_generation_result_t; + typedef struct { // Semicolon-separated target=...,key=value rules. NULL preserves defaults. const char* rules; @@ -456,6 +478,7 @@ typedef struct { bool circular_x; bool circular_y; sd_image_preprocess_params_t image_preprocess; + const char* ref_image_args; } sd_vid_gen_params_t; typedef struct sd_ctx_t sd_ctx_t; @@ -486,8 +509,6 @@ SD_API int32_t sd_get_num_physical_cores(); SD_API const char* sd_get_system_info(); SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx); SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx); -/** Returns whether the loaded model consumes ref_images as conditioning inputs. */ -SD_API bool sd_ctx_supports_ref_image_conditioning(const sd_ctx_t* sd_ctx); // ControlNet hot-swap APIs are not safe to call while generation is in flight. SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path); @@ -534,8 +555,12 @@ SD_API void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params); SD_API char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params); SD_API bool generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params, - sd_image_t** images_out, - int* num_images_out); + sd_image_generation_result_t** result_out); + +// Results own their pixels, audio and generation_info, independently of the context. +// Generation calls require result_out and set it to NULL on failure. +SD_API void free_sd_image_generation_result(sd_image_generation_result_t* result); +SD_API void free_sd_video_generation_result(sd_video_generation_result_t* result); enum sd_cancel_mode_t { // Stop the current generation as soon as possible. @@ -549,13 +574,9 @@ enum sd_cancel_mode_t { SD_API void sd_cancel_generation(sd_ctx_t* sd_ctx, enum sd_cancel_mode_t mode); SD_API void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params); -// If non-NULL, fps_out receives the effective encoding frame rate before preview callbacks. SD_API bool generate_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, - sd_audio_t** audio_out, - int* fps_out); + sd_video_generation_result_t** result_out); typedef struct upscaler_ctx_t upscaler_ctx_t; @@ -570,8 +591,7 @@ SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx); SD_API bool upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_t upscale_factor, - sd_image_t** images_out, - int* num_images_out); + sd_image_generation_result_t** result_out); SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx); // Reads model metadata only; returns 0 if the file is not a recognized RGB ESRGAN model. @@ -595,8 +615,7 @@ SD_API bool adetail_image(adetailer_ctx_t* adetailer_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_image_generation_result_t** result_out); SD_API bool convert(const char* input_path, const char* vae_path, diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index c0e0c03b1..e3832a1d8 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -141,6 +141,7 @@ struct ConditionerParams { }; struct Conditioner { + virtual int resolve_clip_skip(int) const { return 0; } virtual ~Conditioner() = default; public: @@ -173,6 +174,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { std::vector token_embed_custom2; std::map> embedding_pos_map; + int resolve_clip_skip(int value) const override { + return value > 0 ? value : ((sd_version_is_sd2(version) || sd_version_is_sdxl(version)) ? 2 : 1); + } + FrozenCLIPEmbedderWithCustomWords(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map, const std::map& orig_embedding_map, @@ -456,10 +461,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { sd::Tensor hidden_states; // [n_token, hidden_size] or [n_token, hidden_size + hidden_size2] sd::Tensor pooled; - if (clip_skip <= 0) { - clip_skip = (sd_version_is_sd2(version) || sd_version_is_sdxl(version)) ? 2 : 1; - } - size_t chunk_len = 77; size_t chunk_count = tokens.size() / chunk_len; for (int chunk_idx = 0; chunk_idx < chunk_count; chunk_idx++) { @@ -651,6 +652,8 @@ struct SD3CLIPEmbedder : public Conditioner { std::shared_ptr clip_g; std::shared_ptr t5; + int resolve_clip_skip(int value) const override { return value > 0 ? value : 2; } + SD3CLIPEmbedder(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, std::shared_ptr weight_manager = nullptr, @@ -898,10 +901,6 @@ struct SD3CLIPEmbedder : public Conditioner { auto& t5_tokens = token_and_weights[2].first; auto& t5_weights = token_and_weights[2].second; - if (clip_skip <= 0) { - clip_skip = 2; - } - size_t chunk_len = 77; int64_t t0 = ggml_time_ms(); sd::Tensor hidden_states; @@ -1097,6 +1096,8 @@ struct FluxCLIPEmbedder : public Conditioner { std::shared_ptr t5; size_t chunk_len = 256; + int resolve_clip_skip(int value) const override { return value > 0 ? value : 2; } + FluxCLIPEmbedder(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, std::shared_ptr weight_manager = nullptr, @@ -1293,10 +1294,6 @@ struct FluxCLIPEmbedder : public Conditioner { auto& t5_tokens = token_and_weights[1].first; auto& t5_weights = token_and_weights[1].second; - if (clip_skip <= 0) { - clip_skip = 2; - } - int64_t t0 = ggml_time_ms(); sd::Tensor hidden_states; // [N, n_token, 4096] sd::Tensor pooled; // [768,] @@ -1833,11 +1830,11 @@ struct AnimaConditioner : public Conditioner { std::shared_ptr weight_manager = nullptr, const TokenizerConfig& tokenizers = {}) { llm = std::make_shared(LLM::LLMArch::QWEN3, - backend, - tensor_storage_map, - "text_encoders.llm", - false, - weight_manager); + backend, + tensor_storage_map, + "text_encoders.llm", + false, + weight_manager); qwen_tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 151643); if (!qwen_tokenizer) { qwen_tokenizer = std::make_shared(); @@ -2026,11 +2023,11 @@ struct LLMEmbedder : public Conditioner { arch = LLM::LLMArch::QWEN3_5; } llm = std::make_shared(arch, - backend, - tensor_storage_map, - "text_encoders.llm", - enable_vision, - weight_manager); + backend, + tensor_storage_map, + "text_encoders.llm", + enable_vision, + weight_manager); int pad_id = 151643; if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) { pad_id = 11; @@ -2148,10 +2145,11 @@ struct LLMEmbedder : public Conditioner { std::tuple, std::vector, std::vector> tokenize(std::string text, const std::pair& attn_range, - size_t min_length = 0, - size_t max_length = 100000000, - bool spell_quotes = false, - const std::string& suffix = "") { + size_t min_length = 0, + size_t max_length = 100000000, + bool spell_quotes = false, + const std::string& suffix = "", + const std::vector>>& image_embeds = {}) { std::vector> parsed_attention; if (attn_range.first >= 0 && attn_range.second > 0) { if (attn_range.first > 0) { @@ -2196,6 +2194,19 @@ struct LLMEmbedder : public Conditioner { weights.insert(weights.end(), curr_tokens.size(), curr_weight); } + for (const auto& image_embed : image_embeds) { + const int index = image_embed.first; + const int64_t count = image_embed.second.shape()[1]; + if (index < 0 || static_cast(index) >= tokens.size() || count < 1) { + LOG_ERROR("Invalid reference image token range"); + return {}; + } + const int token = tokens[index]; + const float weight = weights[index]; + tokens.insert(tokens.begin() + index, count - 1, token); + weights.insert(weights.begin() + index, count - 1, weight); + } + if (!suffix.empty()) { std::vector suffix_tokens; if (!tokenizer->encode(suffix, suffix_tokens) || suffix_tokens.size() >= max_length) { @@ -2216,6 +2227,17 @@ struct LLMEmbedder : public Conditioner { std::vector mask; tokenizer->pad_tokens(tokens, &weights, &mask, min_length, max_length); + // Keep masks and weights in expanded coordinates, but only look up one token per image. + for (auto it = image_embeds.rbegin(); it != image_embeds.rend(); ++it) { + const int64_t begin = it->first; + const int64_t end = begin + it->second.shape()[1]; + if (begin < 0 || end > static_cast(tokens.size())) { + LOG_ERROR("Reference image exceeds the tokenized context"); + return {}; + } + tokens.erase(tokens.begin() + begin + 1, tokens.begin() + end); + } + // for (int i = 0; i < tokens.size(); i++) { // std::cout << tokens[i] << ":" << weights[i] << ", " << i << std::endl; // } @@ -2238,7 +2260,7 @@ struct LLMEmbedder : public Conditioner { const std::vector& image_grids = {}, const std::string& prompt_suffix = "", sd::Tensor* output_mask = nullptr) { - auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes, prompt_suffix); + auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes, prompt_suffix, image_embeds); auto& tokens = std::get<0>(tokens_weights_mask); auto& weights = std::get<1>(tokens_weights_mask); auto& mask = std::get<2>(tokens_weights_mask); @@ -2268,10 +2290,10 @@ struct LLMEmbedder : public Conditioner { sd::Tensor position_ids; if (sd_version_is_flux3(version)) { - position_ids = sd::Tensor::zeros({static_cast(tokens.size() * 4)}); + position_ids = sd::Tensor::zeros({static_cast(weights.size() * 4)}); for (size_t axis = 0; axis < 4; ++axis) { - for (size_t i = 0; i < tokens.size(); ++i) { - position_ids[axis * tokens.size() + i] = mask[i] == 0.f ? 1 : static_cast(i); + for (size_t i = 0; i < weights.size(); ++i) { + position_ids[axis * weights.size() + i] = mask[i] == 0.f ? 1 : static_cast(i); } } } @@ -2368,7 +2390,7 @@ struct LLMEmbedder : public Conditioner { } } - bool expand_vision_marker(std::string& prompt, + bool expand_vision_marker(const std::string& prompt, size_t& marker_start, std::vector> outputs, int grid_h, @@ -2391,6 +2413,9 @@ struct LLMEmbedder : public Conditioner { return false; } int index = static_cast(prefix_tokens.size()); + for (const auto& image_embed : image_embeds) { + index += static_cast(image_embed.second.shape()[1]) - 1; + } int count = static_cast(outputs[0].shape()[1]); if (index > max_length || count > max_length - index) { LOG_ERROR("Reference images exceed the %d-token context", max_length); @@ -2405,13 +2430,7 @@ struct LLMEmbedder : public Conditioner { deepstack_image_embeds[layer - 1].emplace_back(index, std::move(outputs[layer])); } image_grids.push_back({index, count, grid_h, grid_w}); - std::string markers; - markers.reserve(static_cast(count) * placeholder.size()); - for (int i = 0; i < count; ++i) { - markers += placeholder; - } - prompt.replace(position, placeholder.size(), markers); - marker_start = position + markers.size(); + marker_start = position + placeholder.size(); return true; } @@ -2638,14 +2657,13 @@ struct LLMEmbedder : public Conditioner { return {}; } int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); img_prompt += "<|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(img_prompt.size() + static_cast(num_image_tokens) * placeholder.size() + 32); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } prompt += img_prompt; @@ -2707,7 +2725,6 @@ struct LLMEmbedder : public Conditioner { if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { LOG_INFO("%s", sd_version_is_mage_flow(version) ? "MageFlowEditPipeline" : "QwenImageEditPlusPipeline"); prompt_template_encode_start_idx = 64; - int image_embed_idx = 64 + 6; int min_pixels = conditioner_params.ref_image_params.vlm_min_size; if (min_pixels <= 0) { @@ -2724,6 +2741,7 @@ struct LLMEmbedder : public Conditioner { } } + prompt = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n"; std::string placeholder = "<|image_pad|>"; std::string img_prompt; @@ -2743,19 +2761,22 @@ struct LLMEmbedder : public Conditioner { auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); + std::string image_prefix = prompt + img_prompt + (sd_version_is_mage_flow(version) ? "Image " : "Picture ") + std::to_string(i + 1) + ": <|vision_start|>"; + std::vector prefix_tokens; + if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) { + return {}; + } + int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); - image_embed_idx += 1 + static_cast(image_embed.shape()[1]) + 6; img_prompt += (sd_version_is_mage_flow(version) ? "Image " : "Picture ") + std::to_string(i + 1) + ": <|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(num_image_tokens * placeholder.size()); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } - prompt = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n"; prompt += img_prompt; prompt_attn_range.first = static_cast(prompt.size()); @@ -2833,14 +2854,13 @@ struct LLMEmbedder : public Conditioner { return {}; } int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); img_prompt += "<|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(img_prompt.size() + static_cast(num_image_tokens) * placeholder.size() + 32); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } @@ -2905,14 +2925,13 @@ struct LLMEmbedder : public Conditioner { return {}; } int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(img_prompt.size() + static_cast(num_image_tokens) * placeholder.size() + 32); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } prompt += img_prompt; @@ -2930,7 +2949,6 @@ struct LLMEmbedder : public Conditioner { LOG_INFO("LongCatEditPipeline"); prompt_template_encode_start_idx = 67; min_length = 512 + prompt_template_encode_start_idx; - int image_embed_idx = 36 + 6; int min_pixels = conditioner_params.ref_image_params.vlm_min_size; if (min_pixels <= 0) { @@ -2947,6 +2965,7 @@ struct LLMEmbedder : public Conditioner { } } + prompt = "<|im_start|>system\nAs an image editing expert, first analyze the content and attributes of the input image(s). Then, based on the user's editing instructions, clearly and precisely determine how to modify the given image(s), ensuring that only the specified parts are altered and all other aspects remain consistent with the original(s).<|im_end|>\n<|im_start|>user\n"; std::string placeholder = "<|image_pad|>"; std::string img_prompt; @@ -2965,19 +2984,22 @@ struct LLMEmbedder : public Conditioner { auto resized_image = clip_preprocess(image, w_bar, h_bar); auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); + std::string image_prefix = prompt + img_prompt + "<|vision_start|>"; + std::vector prefix_tokens; + if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) { + return {}; + } + int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); - image_embed_idx += 1 + static_cast(image_embed.shape()[1]) + 6; img_prompt += "<|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(num_image_tokens * placeholder.size()); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } - prompt = "<|im_start|>system\nAs an image editing expert, first analyze the content and attributes of the input image(s). Then, based on the user's editing instructions, clearly and precisely determine how to modify the given image(s), ensuring that only the specified parts are altered and all other aspects remain consistent with the original(s).<|im_end|>\n<|im_start|>user\n"; prompt += img_prompt; } else { prompt_template_encode_start_idx = 36; @@ -3045,7 +3067,10 @@ struct LLMEmbedder : public Conditioner { } else if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version)) { prompt_template_encode_start_idx = 0; out_layers = {35}; // -2 - + if (conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { + LOG_ERROR("Z-Image currently does not support reference-image conditioning"); + return {}; + } if (conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { LOG_INFO("ZImageOmniPipeline"); prompt = "<|im_start|>user\n<|vision_start|>"; @@ -3284,7 +3309,7 @@ struct LLMEmbedder : public Conditioner { for (const auto& [index, image_embed] : image_embeds) { int64_t begin = std::max(0, index - 1); int64_t end = std::min(static_cast(tags.size()), - index + image_embed.shape()[1] + 1); + index + image_embed.shape()[1] + 1); std::fill(tags.begin() + begin, tags.begin() + end, 0); } int64_t tag_count = static_cast(tags.size()); @@ -3464,9 +3489,9 @@ struct LLaDAImageEmbedder : public Conditioner { // <|endoftext|> doubles as the pad token in LLaDA2's tokenizer.json. tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 156892); query_former = std::make_shared(backend, - tensor_storage_map, - query_former_prefix, - weight_manager); + tensor_storage_map, + query_former_prefix, + weight_manager); text_projection = std::make_shared(backend, tensor_storage_map, text_projection_prefix, @@ -3580,12 +3605,10 @@ struct LLaDAImageEmbedder : public Conditioner { auto inputs_embeds = llm->compute_input_embeds(n_threads, text_ids); auto query_embeds = query_former->compute(n_threads, inputs_embeds); - // splice_image_embeds() replaces tokens in place, so the query slots have to exist in - // input_ids; their ids are irrelevant because the embeddings are overwritten. - std::vector padded(tokens.begin(), tokens.end()); - padded.resize(static_cast(n_text + num_queries), tokenizer->PAD_TOKEN_ID); - int64_t n_total = static_cast(padded.size()); - sd::Tensor input_ids({n_total}, padded); + std::vector input_tokens(tokens.begin(), tokens.end()); + input_tokens.push_back(tokenizer->PAD_TOKEN_ID); + sd::Tensor input_ids({n_text + 1}, std::move(input_tokens)); + int64_t n_total = n_text + num_queries; // Bidirectional everywhere except that the text tokens must not see the appended // queries, matching backbone_attention_mask[:, :, :text_length, text_length:] = min. @@ -3670,20 +3693,20 @@ struct LTXAVEmbedder : public Conditioner { LLM::LLMArch arch = detect_gemma_arch(tensor_storage_map, llm_prefix); LOG_INFO("ltxav text encoder: %s", arch == LLM::LLMArch::GEMMA4_12B ? "gemma 4" : "gemma 3"); llm = std::make_shared(arch, - backend, - tensor_storage_map, - llm_prefix, - false, - weight_manager); + backend, + tensor_storage_map, + llm_prefix, + false, + weight_manager); tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 0, true); if (!tokenizer) { tokenizer = std::make_shared(); } dual_projection = tensor_storage_map.find(projector_prefix + ".video_aggregate_embed.weight") != tensor_storage_map.end(); projector = std::make_shared(backend, - tensor_storage_map, - projector_prefix, - weight_manager); + tensor_storage_map, + projector_prefix, + weight_manager); } void get_param_tensors(std::map& tensors) override { diff --git a/src/conditioning/conditioning_cache.h b/src/conditioning/conditioning_cache.h index 6e6cbcaf7..34648c9ad 100644 --- a/src/conditioning/conditioning_cache.h +++ b/src/conditioning/conditioning_cache.h @@ -48,7 +48,8 @@ class ConditioningCache { return std::tie(p.text, p.clip_skip, p.width, p.height, p.zero_out_masked, r.pass_to_vlm, r.pass_to_dit, r.ref_index_mode, r.force_ref_timestep_zero, r.resize_before_vae, r.vae_input_max_pixels, - r.vlm_resize_mode, r.vlm_min_size, r.vlm_max_size, r.resize_vae_to_target); + r.vlm_resize_mode, r.vlm_min_size, r.vlm_max_size, r.resize_vae_to_target, + r.ref_image_size); }; if (fields(a) != fields(b) || (a.ref_images == nullptr) != (b.ref_images == nullptr) || diff --git a/src/core/backend_fit.cpp b/src/core/backend_fit.cpp index 84667ef36..91dfedb48 100644 --- a/src/core/backend_fit.cpp +++ b/src/core/backend_fit.cpp @@ -183,7 +183,7 @@ namespace sd::backend_fit { vm_statistics64_data_t stats{}; mach_msg_type_number_t count = HOST_VM_INFO64_COUNT; const bool ok = host_page_size(host, &page_size) == KERN_SUCCESS && - host_statistics64(host, HOST_VM_INFO64, (host_info64_t)&stats, &count) == KERN_SUCCESS; + host_statistics64(host, HOST_VM_INFO64, (host_info64_t)&stats, &count) == KERN_SUCCESS; mach_port_deallocate(mach_task_self(), host); if (ok) { return ((int64_t)stats.free_count + stats.inactive_count) * page_size; @@ -262,10 +262,10 @@ namespace sd::backend_fit { const auto& runtime_devices = plan.runtimes[ci].devices; const bool fits_runtime = !runtime_devices.empty() && - std::all_of(runtime_devices.begin(), runtime_devices.end(), [&](size_t di) { + std::all_of(runtime_devices.begin(), runtime_devices.end(), [&](size_t di) { const int64_t headroom = headroom_for(ci, di); return headroom <= remaining[di] && comp.params_bytes <= remaining[di] - headroom; - }); + }); if (fits_runtime) { decision.params_location = ParamsLocation::MAIN_GPU; decision.params_device = runtime_devices.front(); diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 64c153329..1099f4b03 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -554,9 +554,9 @@ std::vector split_image_qkv(ggml_context* ctx, qkv = ggml_cont(ctx, ggml_ext_torch_permute(ctx, qkv, 0, 1, 3, 2)); // [3, N, C, H*W] int64_t offset = qkv->nb[2] * qkv->ne[2]; - auto q = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 0); // [N, C, H, W] - auto k = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 1); // [N, C, H, W] - auto v = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 2); // [N, C, H, W] + auto q = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 0); // [N, C, H, W] + auto k = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 1); // [N, C, H, W] + auto v = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 2); // [N, C, H, W] return {q, k, v}; } diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 4c156ed43..1efcb9fa2 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -896,7 +896,7 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n const bool last = index + 1 == plan.segments.size(); auto fail_segment = [&](const char* phase) { LOG_ERROR("%s segment %zu/%zu (%s) failed during %s", get_desc().c_str(), - index + 1, plan.segments.size(), segment.group_name.c_str(), phase); + index + 1, plan.segments.size(), segment.group_name.c_str(), phase); return std::nullopt; }; cut_cache_.prune(segment.live_cut_names); diff --git a/src/core/pe_cache.cpp b/src/core/pe_cache.cpp new file mode 100644 index 000000000..df4aa62cb --- /dev/null +++ b/src/core/pe_cache.cpp @@ -0,0 +1,30 @@ +#include "core/pe_cache.h" + +namespace sd { + PECacheKey& PECacheKey::add_shape(const ggml_tensor* tensor) { + add(tensor != nullptr); + if (tensor != nullptr) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + add(tensor->ne[i]); + } + } + return *this; + } + + PECacheKey& PECacheKey::add_shapes(const std::vector& tensors) { + add(tensors.size()); + for (const auto* tensor : tensors) { + add_shape(tensor); + } + return *this; + } + + const Tensor& PECache::get_or_create(const PECacheKey& key, const std::function()>& build) { + auto entry = entries_.find(key.name()); + if (entry != entries_.end() && entry->second.key == key) { + return entry->second.value; + } + auto value = build(); + return entries_.insert_or_assign(key.name(), Entry{key, std::move(value)}).first->second.value; + } +} diff --git a/src/core/pe_cache.h b/src/core/pe_cache.h new file mode 100644 index 000000000..3539707bd --- /dev/null +++ b/src/core/pe_cache.h @@ -0,0 +1,82 @@ +#ifndef __SD_CORE_PE_CACHE_H__ +#define __SD_CORE_PE_CACHE_H__ + +#include +#include +#include +#include + +#include "core/tensor.hpp" +#include "ggml.h" + +namespace sd { + class PECacheKey { + std::string name_; + std::string fields_; + + public: + explicit PECacheKey(std::string name) + : name_(std::move(name)) {} + const std::string& name() const { return name_; } + + template || std::is_enum_v, int> = 0> + PECacheKey& add(T value) { + if constexpr (std::is_floating_point_v) { + const double field = static_cast(value); + fields_.push_back('f'); + fields_.append(reinterpret_cast(&field), sizeof(field)); + } else { + const uint64_t field = static_cast(value); + fields_.push_back('i'); + fields_.append(reinterpret_cast(&field), sizeof(field)); + } + return *this; + } + + template + PECacheKey& add(const std::vector& values) { + fields_.push_back('v'); + add(values.size()); + for (const auto& value : values) { + add(value); + } + return *this; + } + + template + PECacheKey& add(const std::set& values) { + return add(std::vector(values.begin(), values.end())); + } + + template + PECacheKey& add(const Tensor& values) { + add(values.shape()); + for (int64_t i = 0; i < values.numel(); ++i) { + add(values.data()[i]); + } + return *this; + } + + PECacheKey& add_shape(const ggml_tensor* tensor); + PECacheKey& add_shapes(const std::vector& tensors); + + bool operator==(const PECacheKey& other) const { + return name_ == other.name_ && fields_ == other.fields_; + } + }; + + class PECache { + struct Entry { + PECacheKey key; + Tensor value; + }; + + std::map entries_; + + public: + const Tensor& get_or_create(const PECacheKey& key, const std::function()>& build); + }; + +} + +#endif // __SD_CORE_PE_CACHE_H__ diff --git a/src/core/segment_graph_bindings.cpp b/src/core/segment_graph_bindings.cpp index 3c40a3b03..ba72116f3 100644 --- a/src/core/segment_graph_bindings.cpp +++ b/src/core/segment_graph_bindings.cpp @@ -112,7 +112,7 @@ namespace sd { input_tensor->data = cached_tensor->data == nullptr ? nullptr : static_cast(static_cast(cached_tensor->data) + - input_tensor->view_offs); + input_tensor->view_offs); input_tensor->extra = cached_tensor->extra; } else { input_tensor->buffer = cached_tensor->buffer; diff --git a/src/core/util.cpp b/src/core/util.cpp index 7bd5c0bd1..a691b423c 100644 --- a/src/core/util.cpp +++ b/src/core/util.cpp @@ -343,7 +343,7 @@ int32_t sd_get_num_physical_cores() { // enumerate the set of thread siblings, num entries is num cores std::unordered_set siblings; for (uint32_t cpu = 0; cpu < UINT32_MAX; ++cpu) { - std::ifstream thread_siblings("/sys/devices/system/cpu" + std::to_string(cpu) + "/topology/thread_siblings"); + std::ifstream thread_siblings("/sys/devices/system/cpu/cpu" + std::to_string(cpu) + "/topology/thread_siblings"); if (!thread_siblings.is_open()) { break; // no more cpus } diff --git a/src/detailer.cpp b/src/detailer.cpp index ec4d6065a..9c48ca478 100644 --- a/src/detailer.cpp +++ b/src/detailer.cpp @@ -1,4 +1,6 @@ #include "detailer.h" +#include "generation_info.h" +#include "pipeline/request.h" #include #include @@ -915,14 +917,11 @@ bool adetail_image(adetailer_ctx_t* context, sd_image_t input_image, const sd_adetailer_params_t* public_params, const sd_img_gen_params_t* inpaint_params, - sd_image_t** images_out, - int* num_images_out) { - if (images_out != nullptr) { - *images_out = nullptr; - } - if (num_images_out != nullptr) { - *num_images_out = 0; - } + sd_image_generation_result_t** result_out) try { + if (result_out) + *result_out = nullptr; + if (!result_out) + return false; if (context == nullptr || context->detailer == nullptr || sd_ctx == nullptr || public_params == nullptr || inpaint_params == nullptr || input_image.data == nullptr) { return false; @@ -933,7 +932,10 @@ bool adetail_image(adetailer_ctx_t* context, return false; } - auto detections = context->detailer->predict(input_image, params); + auto resolved_inpaint = *inpaint_params; + resolved_inpaint.seed = sd::pipeline::resolve_seed(resolved_inpaint.seed); + const auto info = sd::adetailer_generation_info(*public_params, resolved_inpaint); + auto detections = context->detailer->predict(input_image, params); filter_and_sort_detections(&detections, input_image.width, input_image.height, params); auto masks = make_masks(detections, input_image.width, input_image.height, params); @@ -942,6 +944,7 @@ bool adetail_image(adetailer_ctx_t* context, return false; } + std::unique_ptr current_owner(current.data, &free); int applied = 0; for (size_t i = 0; i < masks.size(); ++i) { std::string prompt = resolve_prompt(params.prompt, inpaint_params->prompt, static_cast(i)); @@ -968,11 +971,10 @@ bool adetail_image(adetailer_ctx_t* context, if (local_image.data == nullptr || local_mask.data == nullptr) { free(local_image.data); free(local_mask.data); - free(current.data); return false; } - sd_img_gen_params_t generation = *inpaint_params; + sd_img_gen_params_t generation = resolved_inpaint; generation.prompt = prompt.c_str(); generation.negative_prompt = negative_prompt.c_str(); generation.init_image = local_image; @@ -980,7 +982,7 @@ bool adetail_image(adetailer_ctx_t* context, generation.width = params.inpaint_width; generation.height = params.inpaint_height; generation.strength = params.denoising_strength; - generation.seed = inpaint_params->seed + static_cast(i); + generation.seed = static_cast((static_cast(resolved_inpaint.seed) + i) & INT64_MAX); generation.batch_count = 1; generation.ref_images = nullptr; generation.ref_images_count = 0; @@ -1004,37 +1006,32 @@ bool adetail_image(adetailer_ctx_t* context, generation.sample_params.scheduler = params.scheduler; } - sd_image_t* generated = nullptr; - int generated_count = 0; - bool success = generate_image(sd_ctx, &generation, &generated, &generated_count); + sd_image_generation_result_t* generated = nullptr; + bool success = generate_image(sd_ctx, &generation, &generated); free(local_image.data); free(local_mask.data); - if (!success || generated_count <= 0 || generated == nullptr || generated[0].data == nullptr) { - free_sd_images(generated, generated_count); - free(current.data); + if (!success || generated == nullptr || generated->count <= 0 || generated->items[0].image.data == nullptr) { + free_sd_image_generation_result(generated); return false; } + std::unique_ptr generated_owner(generated, &free_sd_image_generation_result); Mask feather = gaussian_blur_mask(masks[i], params.mask_blur); - composite_crop(¤t, generated[0], feather, crop); - free_sd_images(generated, generated_count); + composite_crop(¤t, generated->items[0].image, feather, crop); ++applied; } + const std::vector infos{info}; auto* outputs = static_cast(calloc(1, sizeof(sd_image_t))); - if (outputs == nullptr) { - free(current.data); + if (!outputs) return false; - } outputs[0] = current; - if (images_out != nullptr) { - *images_out = outputs; - } else { - free_sd_images(outputs, 1); - } - if (num_images_out != nullptr) { - *num_images_out = 1; - } + current_owner.release(); + if (!sd::make_image_generation_result(outputs, 1, infos, result_out)) + return false; LOG_INFO("ADetailer applied %d mask(s)", applied); return true; +} catch (const std::exception& e) { + LOG_ERROR("ADetailer failed: %s", e.what()); + return false; } diff --git a/src/extensions/photomaker_extension.cpp b/src/extensions/photomaker_extension.cpp index f904c8f73..59b9a2c2e 100644 --- a/src/extensions/photomaker_extension.cpp +++ b/src/extensions/photomaker_extension.cpp @@ -154,12 +154,12 @@ struct PhotoMakerExtension : public GenerationExtension { PMVersion pm_version = std::strstr(model_path.c_str(), "v2") != nullptr ? PM_VERSION_2 : PM_VERSION_1; pmid_model = std::make_shared(ctx.backend_for(SDBackendModule::PHOTOMAKER), - ctx.tensor_storage_map, - "pmid", - ctx.version, - pm_version, - 20.f, - ctx.model_manager); + ctx.tensor_storage_map, + "pmid", + ctx.version, + pm_version, + 20.f, + ctx.model_manager); pmid_model->set_scale_overrides(ctx.params->linear_scale, ctx.params->attn_scale); if (pm_version == PM_VERSION_2) { LOG_INFO("using PhotoMaker Version 2"); diff --git a/src/generation_info.cpp b/src/generation_info.cpp new file mode 100644 index 000000000..2f49054ea --- /dev/null +++ b/src/generation_info.cpp @@ -0,0 +1,276 @@ +#include "generation_info.h" + +#include +#include +#include + +#include "pipeline/request.h" +#include "runtime/guidance.h" + +namespace sd { + using nlohmann::json; + + static const char* text(const char* value) { + return value ? value : ""; + } + + static json sampling_info(const sd_sample_params_t& p, bool high_noise = false) { + const auto& g = p.guidance; + json info = {{"sample_steps", p.sample_steps}, {"guidance", {{"txt_cfg", g.txt_cfg}, {"distilled_guidance", g.distilled_guidance}}}}; + if (!high_noise && p.shifted_timestep > 0) + info["shifted_timestep"] = p.shifted_timestep; + if (std::isfinite(g.img_cfg) && g.img_cfg != 1.f) + info["guidance"]["img_cfg"] = g.img_cfg; + if (p.sample_method != SAMPLE_METHOD_COUNT) + info["sample_method"] = sd_sample_method_name(p.sample_method); + if (!high_noise && p.scheduler != SCHEDULER_COUNT) + info["scheduler"] = sd_scheduler_name(p.scheduler); + if (std::isfinite(p.eta)) + info["eta"] = p.eta; + if (!high_noise && std::isfinite(p.flow_shift)) + info["flow_shift"] = p.flow_shift; + if (!high_noise && p.custom_sigmas && p.custom_sigmas_count > 0) + info["custom_sigmas"] = std::vector(p.custom_sigmas, p.custom_sigmas + p.custom_sigmas_count); + if (p.extra_sample_args && *p.extra_sample_args) + info["extra_sample_args"] = p.extra_sample_args; + if (g.slg.layers && g.slg.layer_count > 0 && + (g.slg.scale != 0.f || guidance::parse_skip_layer_guidance_uncond_arg(p.extra_sample_args))) + info["guidance"]["slg"] = {{"scale", g.slg.scale}, {"layer_start", g.slg.layer_start}, {"layer_end", g.slg.layer_end}, {"layers", std::vector(g.slg.layers, g.slg.layers + g.slg.layer_count)}}; + return info; + } + + static json hires_info(const sd_hires_params_t& p) { + json info = {{"enabled", p.enabled}}; + if (!p.enabled) + return info; + info.update({{"upscaler", sd_hires_upscaler_name(p.upscaler)}, {"scale", p.scale}, {"target_width", p.target_width}, {"target_height", p.target_height}, {"steps", p.steps}, {"denoising_strength", p.denoising_strength}}); + if (p.upscaler == SD_HIRES_UPSCALER_MODEL) { + info["model_path"] = text(p.model_path); + info["upscale_tile_size"] = p.upscale_tile_size; + } + if (p.custom_sigmas && p.custom_sigmas_count > 0) + info["custom_sigmas"] = std::vector(p.custom_sigmas, p.custom_sigmas + p.custom_sigmas_count); + return info; + } + + static json cache_info(const sd_cache_params_t& p) { + static const char* modes[] = {"disabled", "easycache", "ucache", "dbcache", "taylorseer", "cache_dit", "spectrum"}; + if (p.mode < SD_CACHE_DISABLED || p.mode > SD_CACHE_SPECTRUM) + return {{"mode", static_cast(p.mode)}}; + json info = {{"mode", modes[p.mode]}}; + if (p.mode == SD_CACHE_DISABLED) + return info; + if (p.mode == SD_CACHE_EASYCACHE || p.mode == SD_CACHE_UCACHE) { + if (std::isfinite(p.reuse_threshold)) + info["reuse_threshold"] = p.reuse_threshold; + info.update({{"start_percent", p.start_percent}, {"end_percent", p.end_percent}}); + if (p.mode == SD_CACHE_UCACHE) + info.update({{"error_decay_rate", p.error_decay_rate}, {"use_relative_threshold", p.use_relative_threshold}, {"reset_error_on_compute", p.reset_error_on_compute}}); + } else if (p.mode == SD_CACHE_DBCACHE || p.mode == SD_CACHE_TAYLORSEER || p.mode == SD_CACHE_CACHE_DIT) { + // TaylorSeer also uses the shared DBCache thresholds and limits. + info.update({{"Fn_compute_blocks", p.Fn_compute_blocks}, + {"Bn_compute_blocks", p.Bn_compute_blocks}, + {"residual_diff_threshold", p.residual_diff_threshold}, + {"max_warmup_steps", p.max_warmup_steps}, + {"max_cached_steps", p.max_cached_steps}, + {"max_continuous_cached_steps", p.max_continuous_cached_steps}}); + if (p.mode == SD_CACHE_TAYLORSEER || p.mode == SD_CACHE_CACHE_DIT) + info.update({{"taylorseer_n_derivatives", p.taylorseer_n_derivatives}, {"taylorseer_skip_interval", p.taylorseer_skip_interval}}); + if (p.scm_mask && *p.scm_mask) + info.update({{"scm_mask", p.scm_mask}, {"scm_policy_dynamic", p.scm_policy_dynamic}}); + } else if (p.mode == SD_CACHE_SPECTRUM) { + info.update({{"spectrum_w", p.spectrum_w}, + {"spectrum_m", p.spectrum_m}, + {"spectrum_lam", p.spectrum_lam}, + {"spectrum_window_size", p.spectrum_window_size}, + {"spectrum_flex_window", p.spectrum_flex_window}, + {"spectrum_warmup_steps", p.spectrum_warmup_steps}, + {"spectrum_stop_percent", p.spectrum_stop_percent}}); + } + return info; + } + + template + static json common_info(const Params& p, const char* kind) { + json info = {{"schema", "sdcpp.generation/v1"}, {"kind", kind}, {"prompt", text(p.prompt)}, {"negative_prompt", text(p.negative_prompt)}, {"width", p.width}, {"height", p.height}, {"clip_skip", p.clip_skip}, {"seed", std::to_string(p.seed)}, {"sample_params", sampling_info(p.sample_params)}}; + if (p.hires.enabled) + info["hires"] = hires_info(p.hires); + if (p.cache.mode != SD_CACHE_DISABLED) + info["cache"] = cache_info(p.cache); + if (p.circular_x) + info["circular_x"] = true; + if (p.circular_y) + info["circular_y"] = true; + if (p.ref_image_args && *p.ref_image_args) + info["ref_image_args"] = p.ref_image_args; + if (p.image_preprocess.rules && *p.image_preprocess.rules) + info["image_preprocess"] = {{"rules", p.image_preprocess.rules}}; + const auto& t = p.vae_tiling_params; + if (t.enabled || t.temporal_tiling) { + auto& tiling = info["vae_tiling_params"]; + if (t.enabled) + tiling.update({{"enabled", true}, {"tile_size_w", t.tile_size_w}, {"tile_size_h", t.tile_size_h}, {"target_overlap", t.target_overlap}, {"rel_size_w", t.rel_size_w}, {"rel_size_h", t.rel_size_h}}); + if (t.temporal_tiling) + tiling["temporal_tiling"] = true; + if (t.extra_tiling_args && *t.extra_tiling_args) + tiling["extra_tiling_args"] = t.extra_tiling_args; + } + if (p.lora_count > 0) { + info["loras"] = json::array(); + for (uint32_t i = 0; i < p.lora_count; ++i) + info["loras"].push_back({{"path", text(p.loras[i].path)}, {"multiplier", p.loras[i].multiplier}, {"is_high_noise", p.loras[i].is_high_noise}}); + } + return info; + } + + static json image_info(const sd_img_gen_params_t& p) { + auto info = common_info(p, "image"); + if (p.init_image.data) + info["strength"] = p.strength; + if (p.control_image.data) + info["control_strength"] = p.control_strength; + if (p.ip_adapter_image.data) + info["ip_adapter_strength"] = p.ip_adapter_strength; + if (p.pm_params.id_images && p.pm_params.id_images_count > 0) { + info["pm_params"] = {{"style_strength", p.pm_params.style_strength}}; + if (p.pm_params.id_embed_path && *p.pm_params.id_embed_path) + info["pm_params"]["id_embed_path"] = p.pm_params.id_embed_path; + } + if (p.pulid_params.id_embedding_path && *p.pulid_params.id_embedding_path) + info["pulid_params"] = {{"id_embedding_path", p.pulid_params.id_embedding_path}, {"id_weight", p.pulid_params.id_weight}}; + return info; + } + + template + static Params resolved_params(const pipeline::GenerationRequest& request, Params p) { + p.width = request.width; + p.height = request.height; + p.seed = request.seed; + p.clip_skip = request.clip_skip; + p.sample_params = request.sample_params; + p.hires = request.hires; + p.vae_tiling_params = request.vae_tiling_params; + return p; + } + + json generation_info(const pipeline::GenerationRequest& request, const sd_img_gen_params_t& params) { + auto p = resolved_params(request, params); + p.qwen_image_layers = request.qwen_image_layers; + auto info = image_info(p); + if (!request.models.empty()) + info["models"] = request.models; + if (request.is_layered_image) + info["qwen_image_layers"] = p.qwen_image_layers; + return info; + } + + json generation_info(const pipeline::GenerationRequest& request, const sd_vid_gen_params_t& params) { + const auto p = resolved_params(request, params); + auto info = common_info(p, "video"); + if (!request.models.empty()) + info["models"] = request.models; + info.update({{"video_frames", request.frames}, {"fps", request.fps}}); + if (p.init_image.data || p.end_image.data) + info["strength"] = p.strength; + if (request.use_vace) + info["vace_strength"] = p.vace_strength; + if (request.use_high_noise) + info.update({{"moe_boundary", p.moe_boundary}, {"high_noise_sample_params", sampling_info(request.high_noise_sample_params, true)}}); + if (p.ref_videos_count > 0) { + info["ref_videos"] = json::array(); + for (int i = 0; i < p.ref_videos_count; ++i) + info["ref_videos"].push_back({{"fps", p.ref_videos[i].fps}}); + } + return info; + } + + json adetailer_generation_info(const sd_adetailer_params_t& params, const sd_img_gen_params_t& inpaint_params) { + auto inpaint = image_info(inpaint_params); + inpaint["strength"] = inpaint_params.strength; + // ADetailer clears these inputs before generating each crop. + for (const auto* key : {"control_strength", "pm_params", "pulid_params", "hires", "image_preprocess"}) + inpaint.erase(key); + sd_img_gen_params_t defaults; + sd_img_gen_params_init(&defaults); + if (inpaint_params.qwen_image_layers != defaults.qwen_image_layers) + inpaint["qwen_image_layers"] = inpaint_params.qwen_image_layers; + json adetailer = json::object(); + if (params.prompt && *params.prompt) + adetailer["prompt"] = params.prompt; + if (params.negative_prompt && *params.negative_prompt) + adetailer["negative_prompt"] = params.negative_prompt; + if (params.extra_ad_args && *params.extra_ad_args) + adetailer["extra_ad_args"] = params.extra_ad_args; + return {{"schema", "sdcpp.generation/v1"}, {"kind", "image"}, {"adetailer", adetailer}, {"inpaint_params", inpaint}}; + } + + json upscale_generation_info(uint32_t upscale_factor) { + return {{"schema", "sdcpp.generation/v1"}, {"kind", "image"}, {"upscale", {{"upscale_factor", upscale_factor}}}}; + } + + char* copy_generation_info(const json& info) { + const std::string text = info.dump(-1, ' ', false, json::error_handler_t::replace); + char* result = static_cast(malloc(text.size() + 1)); + if (result) + memcpy(result, text.c_str(), text.size() + 1); + return result; + } + + bool make_image_generation_result(sd_image_t* images, int count, const std::vector& infos, sd_image_generation_result_t** result_out) { + if (!result_out || !images || count <= 0 || infos.size() != static_cast(count)) { + free_sd_images(images, count); + return false; + } + auto* result = static_cast(calloc(1, sizeof(sd_image_generation_result_t))); + if (!result) { + free_sd_images(images, count); + return false; + } + result->items = static_cast(calloc(count, sizeof(sd_image_result_t))); + if (!result->items) { + free(result); + free_sd_images(images, count); + return false; + } + result->count = count; + try { + for (int i = 0; i < count; ++i) { + result->items[i].image = images[i]; + images[i] = {}; + result->items[i].generation_info = copy_generation_info(infos[i]); + if (!result->items[i].generation_info) { + free_sd_images(images, count); + free_sd_image_generation_result(result); + return false; + } + } + } catch (...) { + free_sd_images(images, count); + free_sd_image_generation_result(result); + return false; + } + free(images); + *result_out = result; + return true; + } +} + +void free_sd_image_generation_result(sd_image_generation_result_t* result) { + if (!result) + return; + for (int i = 0; i < result->count; ++i) { + free(result->items[i].image.data); + free(result->items[i].generation_info); + } + free(result->items); + free(result); +} + +void free_sd_video_generation_result(sd_video_generation_result_t* result) { + if (!result) + return; + free_sd_images(result->frames, result->frame_count); + free_sd_audio(result->audio); + free(result->generation_info); + free(result); +} diff --git a/src/generation_info.h b/src/generation_info.h new file mode 100644 index 000000000..56688ee81 --- /dev/null +++ b/src/generation_info.h @@ -0,0 +1,23 @@ +#ifndef __SD_GENERATION_INFO_H__ +#define __SD_GENERATION_INFO_H__ + +#include + +#include "json.hpp" +#include "stable-diffusion.h" + +namespace sd::pipeline { + struct GenerationRequest; +} + +namespace sd { + nlohmann::json generation_info(const pipeline::GenerationRequest& request, const sd_img_gen_params_t& params); + nlohmann::json generation_info(const pipeline::GenerationRequest& request, const sd_vid_gen_params_t& params); + nlohmann::json adetailer_generation_info(const sd_adetailer_params_t& params, const sd_img_gen_params_t& inpaint_params); + nlohmann::json upscale_generation_info(uint32_t upscale_factor); + char* copy_generation_info(const nlohmann::json& info); + // Consumes images even on failure. + bool make_image_generation_result(sd_image_t* images, int count, const std::vector& infos, sd_image_generation_result_t** result_out); +} + +#endif // __SD_GENERATION_INFO_H__ diff --git a/src/model.h b/src/model.h index 73a369301..3ee955383 100644 --- a/src/model.h +++ b/src/model.h @@ -72,6 +72,30 @@ enum SDVersion { VERSION_COUNT, }; +enum class ModelVAEKind { KL, + FLUX, + FLUX2, + WAN, + HUNYUAN, + LTXAV, + MINIMAX_H3, + MAGE }; + +struct ModelTraits { + bool is_dit = false; + bool supports_video_generation = false; + bool is_inpaint = false; + bool is_unet_edit = false; + bool is_control = false; + ModelVAEKind vae_kind = ModelVAEKind::KL; + bool video_tae = false; + bool fake_vae = false; + bool vae_format_override = false; +}; + +const ModelTraits* find_model_traits(SDVersion version); +const char* sd_version_name(SDVersion version); + static inline bool sd_version_is_sd1(SDVersion version) { if (version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX || version == VERSION_SD1_TINY_UNET || version == VERSION_SDXS_512_DS) { return true; @@ -272,7 +296,8 @@ static inline bool sd_version_is_pixart(SDVersion version) { } static inline bool sd_version_supports_video_generation(SDVersion version) { - return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version); + const auto* traits = find_model_traits(version); + return traits && traits->supports_video_generation; } static inline bool sd_version_supports_image_generation(SDVersion version) { @@ -280,116 +305,43 @@ static inline bool sd_version_supports_image_generation(SDVersion version) { } static inline bool sd_version_uses_flux_vae(SDVersion version) { - if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) { - return true; - } - return false; + const auto* traits = find_model_traits(version); + return traits && traits->vae_kind == ModelVAEKind::FLUX; } static inline bool sd_version_uses_flux2_vae(SDVersion version) { - if (sd_version_is_flux2(version) || - sd_version_is_flux3(version) || - sd_version_is_ernie_image(version) || - sd_version_is_lens(version) || - sd_version_is_ideogram4(version) || - sd_version_is_sefi_image(version) || - sd_version_is_llada_image(version) || - version == VERSION_MOSS_IMAGE || - version == VERSION_FIBO2) { - return true; - } - return false; + const auto* traits = find_model_traits(version); + return traits && traits->vae_kind == ModelVAEKind::FLUX2; } static inline bool sd_version_uses_wan_vae(SDVersion version) { - if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version) || version == VERSION_MING_IMAGE) { - return true; - } - return false; + const auto* traits = find_model_traits(version); + return traits && traits->vae_kind == ModelVAEKind::WAN; } static inline bool sd_version_uses_hunyuan_video_vae(SDVersion version) { - return sd_version_is_hunyuan_video(version); + const auto* traits = find_model_traits(version); + return traits && traits->vae_kind == ModelVAEKind::HUNYUAN; } static inline bool sd_version_is_inpaint(SDVersion version) { - if (version == VERSION_SD1_INPAINT || - version == VERSION_SD2_INPAINT || - version == VERSION_SDXL_INPAINT || - version == VERSION_FLUX_FILL || - version == VERSION_FLEX_2) { - return true; - } - return false; + const auto* traits = find_model_traits(version); + return traits && traits->is_inpaint; } static inline bool sd_version_is_dit(SDVersion version) { - if (sd_version_is_flux(version) || - sd_version_is_flux2(version) || - sd_version_is_flux3(version) || - sd_version_is_ltxav(version) || - sd_version_is_minimax_h3(version) || - sd_version_is_sd3(version) || - sd_version_is_wan(version) || - sd_version_is_lingbot_video(version) || - sd_version_is_qwen_image(version) || - sd_version_is_hunyuan_video(version) || - version == VERSION_HIDREAM_O1 || - sd_version_is_anima(version) || - sd_version_is_z_image(version) || - sd_version_is_z_image_l2p(version) || - version == VERSION_MING_IMAGE || - sd_version_is_llada_image(version) || - sd_version_is_boogu_image(version) || - sd_version_is_ernie_image(version) || - sd_version_is_lens(version) || - sd_version_is_minit2i(version) || - sd_version_is_longcat(version) || - sd_version_is_pid(version) || - sd_version_is_ideogram4(version) || - sd_version_is_sefi_image(version) || - sd_version_is_krea2(version) || - sd_version_is_mage_flow(version) || - sd_version_is_sensenova_u1(version) || - sd_version_is_pixart(version) || - version == VERSION_IRIS || - version == VERSION_MOSS_IMAGE || - version == VERSION_FIBO2) { - return true; - } - return false; + const auto* traits = find_model_traits(version); + return traits && traits->is_dit; } static inline bool sd_version_is_unet_edit(SDVersion version) { - return version == VERSION_SD1_PIX2PIX || version == VERSION_SDXL_PIX2PIX; + const auto* traits = find_model_traits(version); + return traits && traits->is_unet_edit; } static inline bool sd_version_is_control(SDVersion version) { - return version == VERSION_FLUX_CONTROLS || version == VERSION_FLEX_2; -} - -// These model families consume uploaded reference images as generation -// conditioning. Inpaint and control variants keep their init-image semantics. -static inline bool sd_version_supports_ref_image_conditioning(SDVersion version) { - if (sd_version_is_inpaint(version) || sd_version_is_control(version) || version == VERSION_CHROMA_RADIANCE) { - return false; - } - - return sd_version_is_unet_edit(version) || - sd_version_is_flux(version) || - sd_version_is_flux3(version) || - sd_version_is_flux2(version) || - sd_version_is_qwen_image(version) || - sd_version_is_longcat(version) || - sd_version_is_z_image(version) || - sd_version_is_boogu_image(version) || - sd_version_is_llada_image(version) || - sd_version_is_krea2(version) || - sd_version_is_mage_flow(version) || - sd_version_is_anima(version) || - sd_version_is_sefi_image(version) || - sd_version_is_pid(version) || - version == VERSION_FIBO2; + const auto* traits = find_model_traits(version); + return traits && traits->is_control; } static bool sd_version_is_inpaint_or_unet_edit(SDVersion version) { diff --git a/src/model/adapter/lora.hpp b/src/model/adapter/lora.hpp index 1e5377ccc..b2657a369 100644 --- a/src/model/adapter/lora.hpp +++ b/src/model/adapter/lora.hpp @@ -69,7 +69,8 @@ struct LoraModel : public GGMLRunner { metadata_alpha = -1.f; const auto metadata = model_manager->loader().file_metadata(source_file); const auto alpha = metadata.find("alpha"); - if (alpha != metadata.end()) { + // Converters that write "baked_scale" have already folded alpha / rank into lora_up. + if (alpha != metadata.end() && metadata.find("baked_scale") == metadata.end()) { try { size_t parsed = 0; metadata_alpha = std::stof(alpha->second, &parsed); diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index 93e457d4b..63584bc5c 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -660,7 +660,7 @@ namespace Rope { const int axes_dim_num = static_cast(axes_dim.size()); auto make_image_ids = [=, &result](int image_h, int image_w, int image_index) { std::vector> image_ids(static_cast(bs) * image_h * image_w, - std::vector(axes_dim_num, 0.f)); + std::vector(axes_dim_num, 0.f)); result.positions.append_image(image_h, image_w); int h_start = -(image_h - image_h / 2); int w_start = -(image_w - image_w / 2); diff --git a/src/model/diffusion/anima.hpp b/src/model/diffusion/anima.hpp index a414d96e6..66e1a0ba3 100644 --- a/src/model/diffusion/anima.hpp +++ b/src/model/diffusion/anima.hpp @@ -7,6 +7,7 @@ #include #include #include +#include "core/pe_cache.h" #include "model/common/block.hpp" #include "model/common/rope.hpp" @@ -562,10 +563,9 @@ namespace Anima { }; struct AnimaRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: - std::vector image_pe_vec; - std::vector adapter_q_pe_vec; - std::vector adapter_k_pe_vec; AnimaConfig config; AnimaNet net; @@ -660,19 +660,23 @@ namespace Anima { int64_t h_pad = x->ne[1] + pad_h; int64_t w_pad = x->ne[0] + pad_w; - image_pe_vec = finish_rope_pe(gen_anima_image_pe_vec(1, - static_cast(h_pad), - static_cast(w_pad), - static_cast(config.patch_size), - config.theta, - config.axes_dim, - 4.0f, - 4.0f, - 1.0f, - ref_latents)); - int64_t image_pos_len = static_cast(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2)); - auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len); - set_backend_tensor_data(image_pe, image_pe_vec.data()); + sd::PECacheKey image_pe_key("anima.image_pe"); + image_pe_key.add(h_pad).add(w_pad).add(config.patch_size).add(config.theta).add(config.axes_dim).add(config.head_dim); + image_pe_key.add_shapes(ref_latents).add(circular_x_enabled).add(circular_y_enabled); + auto image_pe = make_input(pe_cache_.get_or_create(image_pe_key, [&]() { + auto image_pe_vec = finish_rope_pe(gen_anima_image_pe_vec(1, + static_cast(h_pad), + static_cast(w_pad), + static_cast(config.patch_size), + config.theta, + config.axes_dim, + 4.0f, + 4.0f, + 1.0f, + ref_latents)); + int64_t image_pos_len = static_cast(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2)); + return sd::Tensor({2, 2, config.head_dim / 2, image_pos_len}, std::move(image_pe_vec)); + })); ggml_tensor* adapter_q_pe = nullptr; ggml_tensor* adapter_k_pe = nullptr; @@ -680,16 +684,16 @@ namespace Anima { int64_t target_len = t5_ids->ne[0]; int64_t source_len = context->ne[1]; - adapter_q_pe_vec = gen_1d_rope_pe_vec(target_len, 64, 10000.f); - adapter_k_pe_vec = gen_1d_rope_pe_vec(source_len, 64, 10000.f); - - int64_t target_pos_len = static_cast(adapter_q_pe_vec.size()) / (2 * 2 * 32); - int64_t source_pos_len = static_cast(adapter_k_pe_vec.size()) / (2 * 2 * 32); - - adapter_q_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 32, target_pos_len); - adapter_k_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 32, source_pos_len); - set_backend_tensor_data(adapter_q_pe, adapter_q_pe_vec.data()); - set_backend_tensor_data(adapter_k_pe, adapter_k_pe_vec.data()); + sd::PECacheKey q_key("anima.adapter_q_pe"); + q_key.add(target_len); + adapter_q_pe = make_input(pe_cache_.get_or_create(q_key, [&]() { + return sd::Tensor({2, 2, 32, target_len}, gen_1d_rope_pe_vec(target_len, 64, 10000.f)); + })); + sd::PECacheKey k_key("anima.adapter_k_pe"); + k_key.add(source_len); + adapter_k_pe = make_input(pe_cache_.get_or_create(k_key, [&]() { + return sd::Tensor({2, 2, 32, source_len}, gen_1d_rope_pe_vec(source_len, 64, 10000.f)); + })); } auto runner_ctx = get_context(); diff --git a/src/model/diffusion/boogu.hpp b/src/model/diffusion/boogu.hpp index 1ca137098..ac67ff9cd 100644 --- a/src/model/diffusion/boogu.hpp +++ b/src/model/diffusion/boogu.hpp @@ -6,6 +6,7 @@ #include #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_runner.h" @@ -543,18 +544,18 @@ namespace Boogu { blocks["x_embedder"] = std::make_shared(this->config.patch_size * this->config.patch_size * this->config.in_channels, this->config.hidden_size, true); blocks["ref_image_patch_embedder"] = std::make_shared(this->config.patch_size * this->config.patch_size * this->config.in_channels, this->config.hidden_size, true); blocks["time_caption_embed"] = std::make_shared(this->config.hidden_size, - this->config.instruction_feat_dim, - 256, - this->config.norm_eps, - this->config.timestep_scale); + this->config.instruction_feat_dim, + 256, + this->config.norm_eps, + this->config.timestep_scale); for (int i = 0; i < this->config.num_refiner_layers; i++) { blocks["noise_refiner." + std::to_string(i)] = std::make_shared(this->config.hidden_size, - this->config.num_attention_heads, - this->config.num_kv_heads, - this->config.multiple_of, - this->config.norm_eps, - true); + this->config.num_attention_heads, + this->config.num_kv_heads, + this->config.multiple_of, + this->config.norm_eps, + true); blocks["ref_image_refiner." + std::to_string(i)] = std::make_shared(this->config.hidden_size, this->config.num_attention_heads, this->config.num_kv_heads, @@ -562,11 +563,11 @@ namespace Boogu { this->config.norm_eps, true); blocks["context_refiner." + std::to_string(i)] = std::make_shared(this->config.hidden_size, - this->config.num_attention_heads, - this->config.num_kv_heads, - this->config.multiple_of, - this->config.norm_eps, - false); + this->config.num_attention_heads, + this->config.num_kv_heads, + this->config.multiple_of, + this->config.norm_eps, + false); } for (int i = 0; i < this->config.num_double_stream_layers; i++) { @@ -759,9 +760,9 @@ namespace Boogu { } struct BooguImageRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; BooguConfig config; BooguImageModel boogu; - std::vector pe_vec; BooguImageRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -799,17 +800,21 @@ namespace Boogu { ref_latents.push_back(make_input(ref_latent_tensor)); } - pe_vec = finish_rope_pe(gen_boogu_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - ref_latents, - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("boogu.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shapes(ref_latents); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(gen_boogu_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ref_latents, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = boogu.forward(&runner_ctx, x, timesteps, context, pe, ref_latents); diff --git a/src/model/diffusion/ernie_image.hpp b/src/model/diffusion/ernie_image.hpp index b284cb83f..d4f774fc6 100644 --- a/src/model/diffusion/ernie_image.hpp +++ b/src/model/diffusion/ernie_image.hpp @@ -4,6 +4,7 @@ #include #include #include +#include "core/pe_cache.h" #include "model/common/rope.hpp" #include "model/diffusion/dit.hpp" @@ -383,9 +384,9 @@ namespace ErnieImage { }; struct ErnieImageRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; ErnieImageConfig config; ErnieImageModel ernie_image; - std::vector pe_vec; ErnieImageRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -415,16 +416,20 @@ namespace ErnieImage { GGML_ASSERT(!context_tensor.empty()); ggml_tensor* context = make_input(context_tensor); - pe_vec = finish_rope_pe(Rope::gen_ernie_image_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, config.axes_dim_sum, 1, pos_len, 2); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("ernie_image.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_ernie_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({config.axes_dim_sum, 1, pos_len, 2}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = ernie_image.forward(&runner_ctx, x, timesteps, context, pe); diff --git a/src/model/diffusion/fibo2.hpp b/src/model/diffusion/fibo2.hpp index f53b6ecd1..a642cb4df 100644 --- a/src/model/diffusion/fibo2.hpp +++ b/src/model/diffusion/fibo2.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_DIFFUSION_FIBO2_HPP__ #define __SD_MODEL_DIFFUSION_FIBO2_HPP__ +#include "core/pe_cache.h" #include "model/diffusion/z_image.hpp" // https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_bria_fibo2.py @@ -333,9 +334,9 @@ namespace Fibo2 { }; struct Fibo2Runner : public DiffusionModelRunner { + sd::PECache pe_cache_; Fibo2Config config; Fibo2Model model; - std::vector pe_values; Fibo2Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr weight_manager = nullptr) : DiffusionModelRunner(backend, prefix, weight_manager), @@ -360,28 +361,35 @@ namespace Fibo2 { auto x = make_input(*p.x); GGML_ASSERT(x->ne[3] == 1 && x->ne[2] == config.channels); int count = config.gist_count(static_cast(p.context->shape()[1])); - Rope::Embedding embedding; - for (int i = 0; i < count; ++i) { - embedding.ids.push_back({static_cast(i + 1), 0.f, 0.f}); - } - embedding.positions.append_tokens(count); - auto append_image = [&](ggml_tensor* image, int plane) { - auto ids = Rope::gen_flux_img_ids(static_cast(image->ne[1]), static_cast(image->ne[0]), - 1, 1, 3, plane, 0, 0, false, &embedding.positions); - embedding.ids.insert(embedding.ids.end(), ids.begin(), ids.end()); - }; - append_image(x, static_cast(config.max_gist + 1)); std::vector refs; if (p.ref_latents && p.ref_image_params.pass_to_dit) { for (const auto& ref : *p.ref_latents) { refs.push_back(make_input(ref)); - append_image(refs.back(), static_cast(config.max_gist + 1 + refs.size())); } } - embedding.values = Rope::embed_nd(embedding.ids, 1, 256.f, {32, 48, 48}, embedding.layout, &embedding.frequencies); - pe_values = finish_rope_pe(std::move(embedding)); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 64, static_cast(pe_values.size() / 256)); - set_backend_tensor_data(pe, pe_values.data()); + sd::PECacheKey pe_key("fibo2.pe"); + pe_key.add_shape(x).add(count).add(config.max_gist).add_shapes(refs).add(circular_x_enabled); + pe_key.add(circular_y_enabled); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + Rope::Embedding embedding; + for (int i = 0; i < count; ++i) { + embedding.ids.push_back({static_cast(i + 1), 0.f, 0.f}); + } + embedding.positions.append_tokens(count); + auto append_image = [&](ggml_tensor* image, int plane) { + auto ids = Rope::gen_flux_img_ids(static_cast(image->ne[1]), static_cast(image->ne[0]), + 1, 1, 3, plane, 0, 0, false, &embedding.positions); + embedding.ids.insert(embedding.ids.end(), ids.begin(), ids.end()); + }; + append_image(x, static_cast(config.max_gist + 1)); + for (size_t i = 0; i < refs.size(); ++i) { + append_image(refs[i], static_cast(config.max_gist + 2 + i)); + } + embedding.values = Rope::embed_nd(embedding.ids, 1, 256.f, {32, 48, 48}, embedding.layout, &embedding.frequencies); + auto values = finish_rope_pe(std::move(embedding)); + const int64_t tokens = static_cast(values.size() / 256); + return sd::Tensor({2, 2, 64, tokens}, std::move(values)); + })); auto ctx = get_context(); auto out = model.forward(&ctx, x, make_input(*p.timesteps), make_input(*p.context), pe, count, refs); ggml_build_forward_expand(graph, out); diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp index 825dabe13..2f4272eb1 100644 --- a/src/model/diffusion/flux.hpp +++ b/src/model/diffusion/flux.hpp @@ -6,6 +6,7 @@ #include #include "core/ggml_extend_backend.h" #include "core/ggml_tensor_utils.h" +#include "core/pe_cache.h" #include "core/util.h" #include "model/adapter/pulid.hpp" @@ -1398,10 +1399,11 @@ namespace Flux { }; struct FluxRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: FluxConfig config; Flux flux; - std::vector pe_vec; std::vector mod_index_arange_vec; std::vector dct_vec; sd::Tensor guidance_tensor; @@ -1548,25 +1550,26 @@ namespace Flux { } else if (version == VERSION_OVIS_IMAGE) { txt_arange_dims = {1, 2}; } - pe_vec = finish_rope_pe(Rope::gen_flux_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - txt_arange_dims, - ref_latents, - ref_index_mode, - config.ref_index_scale, - config.theta, - config.axes_dim, - sd_version_is_longcat(version))); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_VERBOSE("pos_len %d", pos_len); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - // pe->data = pe_vec.data(); - // print_ggml_tensor(pe); - // pe->data = nullptr; - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("flux.pe"); + pe_key.add(version).add_shape(x).add_shape(context).add(config.patch_size).add(txt_arange_dims); + pe_key.add_shapes(ref_latents).add(ref_index_mode).add(config.ref_index_scale).add(config.theta).add(config.axes_dim); + pe_key.add(circular_x_enabled).add(circular_y_enabled); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_flux_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + txt_arange_dims, + ref_latents, + ref_index_mode, + config.ref_index_scale, + config.theta, + config.axes_dim, + sd_version_is_longcat(version))); + int64_t pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); if (version == VERSION_CHROMA_RADIANCE) { int patch_size = config.patch_size; diff --git a/src/model/diffusion/flux3.hpp b/src/model/diffusion/flux3.hpp index 952587b01..223fc62db 100644 --- a/src/model/diffusion/flux3.hpp +++ b/src/model/diffusion/flux3.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_DIFFUSION_FLUX3_HPP__ #define __SD_MODEL_DIFFUSION_FLUX3_HPP__ +#include "core/pe_cache.h" #include "model/diffusion/flux.hpp" namespace Flux3 { @@ -183,9 +184,9 @@ namespace Flux3 { }; struct Flux3Runner : public DiffusionModelRunner { + sd::PECache pe_cache_; Flux3Config config; Model model; - std::vector pe_values; sd::Tensor zero = sd::Tensor::zeros({1}); Flux3Runner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr weight_manager = nullptr) @@ -213,37 +214,44 @@ namespace Flux3 { auto x = make_input(*params.x); auto text = make_input(*params.context); GGML_ASSERT(x->ne[3] == 1 && text->ne[2] == 1); - Rope::Embedding embedding; - auto append_image = [&](int64_t w, int64_t h, float index) { - embedding.positions.append_image(static_cast(h), static_cast(w), 1, 0.5f, 0.5f); - for (int64_t y = 0; y < h; ++y) { - for (int64_t x = 0; x < w; ++x) { - embedding.ids.push_back({0.f, y * 0.5f, x * 0.5f, index}); - } - } - }; - append_image(x->ne[0], x->ne[1], 0.f); + std::vector refs; ggml_tensor* references = nullptr; if (params.ref_latents && params.ref_image_params.pass_to_dit) { - int index = 0; for (const auto& ref : *params.ref_latents) { auto r = make_input(ref); - append_image(r->ne[0], r->ne[1], -float(++index)); + refs.push_back(r); r = DiT::patchify(compute_ctx, r, 1, 1); references = references ? ggml_concat(compute_ctx, references, r, 1) : r; } } - for (int64_t i = 0; i < text->ne[1]; ++i) { - embedding.ids.push_back({0.f, 0.f, 0.f, static_cast(i)}); - } - embedding.positions.append_tokens(text->ne[1]); - embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, {32, 32, 32, 32}, embedding.layout, &embedding.frequencies); - pe_values = finish_rope_pe(std::move(embedding)); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 64, pe_values.size() / 256); - set_backend_tensor_data(pe, pe_values.data()); + sd::PECacheKey pe_key("flux3.pe"); + pe_key.add_shape(x).add_shape(text).add_shapes(refs).add(circular_x_enabled).add(circular_y_enabled); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + Rope::Embedding embedding; + auto append_image = [&](int64_t w, int64_t h, float index) { + embedding.positions.append_image(static_cast(h), static_cast(w), 1, 0.5f, 0.5f); + for (int64_t y = 0; y < h; ++y) { + for (int64_t x = 0; x < w; ++x) { + embedding.ids.push_back({0.f, y * 0.5f, x * 0.5f, index}); + } + } + }; + append_image(x->ne[0], x->ne[1], 0.f); + for (size_t i = 0; i < refs.size(); ++i) { + append_image(refs[i]->ne[0], refs[i]->ne[1], -float(i + 1)); + } + for (int64_t i = 0; i < text->ne[1]; ++i) { + embedding.ids.push_back({0.f, 0.f, 0.f, static_cast(i)}); + } + embedding.positions.append_tokens(text->ne[1]); + embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, {32, 32, 32, 32}, embedding.layout, &embedding.frequencies); + auto values = finish_rope_pe(std::move(embedding)); + const int64_t tokens = static_cast(values.size() / 256); + return sd::Tensor({2, 2, 64, tokens}, std::move(values)); + })); auto runner_ctx = get_context(); auto out = model.forward(&runner_ctx, x, text, references, make_input(*params.timesteps), make_input(zero), - config.flow_map ? make_input(*extra->target_timestep) : nullptr, pe); + config.flow_map ? make_input(*extra->target_timestep) : nullptr, pe); ggml_build_forward_expand(gf, out); return gf; }; diff --git a/src/model/diffusion/hidream_o1.hpp b/src/model/diffusion/hidream_o1.hpp index 2befb723a..90bfb51b3 100644 --- a/src/model/diffusion/hidream_o1.hpp +++ b/src/model/diffusion/hidream_o1.hpp @@ -56,15 +56,6 @@ namespace HiDreamO1 { } }; - static inline std::string repeat_special_token(const std::string& token, int64_t count) { - std::string out; - out.reserve(static_cast(count) * token.size()); - for (int64_t i = 0; i < count; ++i) { - out += token; - } - return out; - } - static inline std::pair calculate_dimensions(int max_size, double ratio) { int width = static_cast(std::sqrt(max_size * max_size * ratio)); int height = static_cast(width / ratio); @@ -247,10 +238,10 @@ namespace HiDreamO1 { blocks["language_model"] = std::make_shared(this->config.llm); blocks["t_embedder1"] = std::make_shared(this->config.llm.hidden_size); blocks["x_embedder"] = std::make_shared(this->config.patch_size * this->config.patch_size * 3, - this->config.llm.hidden_size / 4, - this->config.llm.hidden_size); + this->config.llm.hidden_size / 4, + this->config.llm.hidden_size); blocks["final_layer2"] = std::make_shared(this->config.llm.hidden_size, - this->config.patch_size * this->config.patch_size * 3); + this->config.patch_size * this->config.patch_size * 3); } std::shared_ptr text_model() { @@ -405,7 +396,7 @@ namespace HiDreamO1 { txt = LLM::splice_image_embeds(&runner_ctx, txt, image_embeds); auto t_emb = t_embedder1->forward(&runner_ctx, timestep); - int64_t txt_seq_len = input_ids->ne[0]; + int64_t txt_seq_len = txt->ne[1]; if (txt_seq_len > 1) { auto prefix = ggml_ext_slice(compute_ctx, txt, 1, 0, txt_seq_len - 1); txt = ggml_concat(compute_ctx, prefix, ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1), 1); @@ -616,8 +607,7 @@ namespace HiDreamO1 { sd::Tensor vlm_image = sd::ops::interpolate( resized_ref, {dims.first, dims.second, resized_ref.shape()[2], resized_ref.shape()[3]}); - vlm_image = vlm_image * 2.0f - 1.0f; - int64_t image_tokens = static_cast(dims.first / PATCH_SIZE) * static_cast(dims.second / PATCH_SIZE); + vlm_image = vlm_image * 2.0f - 1.0f; auto patch_img = resized_ref * 2.0f - 1.0f; result.c_ref_images.push_back(std::move(patch_img)); @@ -626,8 +616,11 @@ namespace HiDreamO1 { return {}; } int64_t prompt_start = static_cast(prefix_tokens.size()); + for (const auto& grid : image_grids) { + prompt_start += static_cast(grid[1]) * grid[2] - 1; + } prompt += "<|vision_start|>"; - prompt += repeat_special_token("<|image_pad|>", image_tokens); + prompt += "<|image_pad|>"; prompt += "<|vision_end|>"; vlm_images.emplace_back(static_cast(prompt_start), std::move(vlm_image)); image_grids.push_back({1, dims.second / PATCH_SIZE, dims.first / PATCH_SIZE}); @@ -642,6 +635,13 @@ namespace HiDreamO1 { } std::vector input_ids_pad = input_ids; + for (size_t i = 0; i < vlm_images.size(); ++i) { + const auto& grid = image_grids[i]; + const int64_t count = static_cast(grid[1]) * grid[2]; + input_ids_pad.insert(input_ids_pad.begin() + vlm_images[i].first, count - 1, IMAGE_TOKEN_ID); + } + + int txt_seq_len = static_cast(input_ids_pad.size()); input_ids_pad.push_back(VISION_START_TOKEN_ID); input_ids_pad.insert(input_ids_pad.end(), target_image_len - 1, IMAGE_TOKEN_ID); image_grids.push_back({1, static_cast(height / PATCH_SIZE), static_cast(width / PATCH_SIZE)}); @@ -656,8 +656,7 @@ namespace HiDreamO1 { } std::vector token_types(input_ids_pad.size(), 0); - int txt_seq_len = static_cast(input_ids.size()); - int bgn = txt_seq_len - TIMESTEP_TOKEN_NUM; + int bgn = txt_seq_len - TIMESTEP_TOKEN_NUM; for (int i = bgn; i < static_cast(token_types.size()); ++i) { token_types[i] = 1; } diff --git a/src/model/diffusion/hunyuan.hpp b/src/model/diffusion/hunyuan.hpp index c106b9cd5..5fc29b6b0 100644 --- a/src/model/diffusion/hunyuan.hpp +++ b/src/model/diffusion/hunyuan.hpp @@ -3,6 +3,7 @@ #include #include +#include "core/pe_cache.h" #include "model/common/block.hpp" #include "model/diffusion/flux.hpp" @@ -302,9 +303,9 @@ namespace Hunyuan { int64_t head_dim = this->config.hidden_size / this->config.num_heads; blocks["txt_in"] = std::make_shared(this->config.context_in_dim, this->config.num_heads, head_dim, 2); blocks["img_in"] = std::make_shared(static_cast(224) /*Not used*/, - this->config.patch_size, - this->config.in_channels, - this->config.hidden_size); + this->config.patch_size, + this->config.in_channels, + this->config.hidden_size); blocks["time_in"] = std::make_shared(256, this->config.hidden_size); if (this->config.vec_in_dim > 0) { blocks["vector_in"] = std::make_shared(this->config.vec_in_dim, this->config.hidden_size); @@ -540,10 +541,11 @@ namespace Hunyuan { }; struct HunyuanVideoRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: HunyuanVideoConfig config; HunyuanVideoModel hunyuan_video; - std::vector pe_vec; SDVersion version; HunyuanVideoRunner(ggml_backend_t backend, @@ -605,23 +607,24 @@ namespace Hunyuan { if (vision != nullptr) { text_len += static_cast(vision->ne[1]); } - pe_vec = Rope::gen_hunyuan_video_pe(static_cast(x->ne[2]), - static_cast(x->ne[1]), - static_cast(x->ne[0]), - std::get<0>(config.patch_size), - std::get<1>(config.patch_size), - std::get<2>(config.patch_size), - 1, - text_len, - config.theta, - config.axes_dim); - int64_t pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_VERBOSE("pos_len %d", pos_len); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - // pe->data = pe_vec.data(); - // print_ggml_tensor(pe, true, "pe"); - // pe->data = nullptr; - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("hunyuan.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(std::get<0>(config.patch_size)); + pe_key.add(std::get<1>(config.patch_size)).add(std::get<2>(config.patch_size)).add(text_len); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = Rope::gen_hunyuan_video_pe(static_cast(x->ne[2]), + static_cast(x->ne[1]), + static_cast(x->ne[0]), + std::get<0>(config.patch_size), + std::get<1>(config.patch_size), + std::get<2>(config.patch_size), + 1, + text_len, + config.theta, + config.axes_dim); + int64_t pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); diff --git a/src/model/diffusion/ideogram4.hpp b/src/model/diffusion/ideogram4.hpp index 3254571db..664c827d5 100644 --- a/src/model/diffusion/ideogram4.hpp +++ b/src/model/diffusion/ideogram4.hpp @@ -8,6 +8,7 @@ #include #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_graph_cut.h" @@ -422,6 +423,8 @@ namespace Ideogram4 { }; class Ideogram4Runner : public DiffusionModelRunner { + sd::PECache pe_cache_; + protected: bool should_use_uncond_model(const DiffusionParams& diffusion_params) const { return has_uncond_model && @@ -436,7 +439,6 @@ namespace Ideogram4 { Ideogram4Transformer uncond_model; bool has_uncond_model = false; std::string uncond_prefix; - std::vector pe_vec; std::vector image_indicator_vec; Ideogram4Runner(ggml_backend_t backend, @@ -496,15 +498,19 @@ namespace Ideogram4 { int64_t head_dim = config.emb_dim / config.num_heads; auto runner_ctx = get_context(); - pe_vec = finish_rope_pe(gen_ideogram4_pe(static_cast(grid_h), + sd::PECacheKey pe_key("ideogram4.pe"); + pe_key.add(grid_h).add(grid_w).add(x->ne[3]).add(context_len).add(head_dim).add(pos_len).add(config.rope_theta); + pe_key.add(config.mrope_section).add(circular_x_enabled).add(circular_y_enabled); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(gen_ideogram4_pe(static_cast(grid_h), static_cast(grid_w), static_cast(x->ne[3]), static_cast(context_len), static_cast(head_dim), static_cast(config.rope_theta), config.mrope_section)); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + return sd::Tensor({2, 2, head_dim / 2, pos_len}, std::move(pe_vec)); + })); image_indicator_vec.assign(static_cast(pos_len), 1); for (int64_t i = 0; i < context_len; ++i) { diff --git a/src/model/diffusion/iris.hpp b/src/model/diffusion/iris.hpp index b90fbe9d4..3681c65ab 100644 --- a/src/model/diffusion/iris.hpp +++ b/src/model/diffusion/iris.hpp @@ -3,6 +3,7 @@ #include #include +#include "core/pe_cache.h" #include "model/diffusion/pid.hpp" namespace Iris { @@ -382,11 +383,9 @@ namespace Iris { } struct IrisRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; IrisConfig config; IrisModel model; - std::vector joint_pe_data; - std::vector pixel_pe_data; - std::vector pixel_pos_data; std::vector text_mask_data; IrisRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr weight_manager = nullptr) @@ -418,24 +417,35 @@ namespace Iris { GGML_ASSERT(h % config.patch_size == 0 && w % config.patch_size == 0); GGML_ASSERT(text->ne[0] == config.text_dim * config.text_layers && text->ne[1] == nt); GGML_ASSERT(inputs.y->numel() == nt && x->ne[3] == 1); - joint_pe_data = Pid::make_rope_1d(nt, static_cast(config.head_dim), 10000.f); - auto img = image_rope(h / config.patch_size, w / config.patch_size, config.head_dim); - joint_pe_data.insert(joint_pe_data.end(), img.begin(), img.end()); - pixel_pe_data = image_rope(h / config.patch_size, w / config.patch_size, config.pixel_attn_dim / config.pixel_heads); - pixel_pos_data = Pid::make_pixel_abs_pos(h, w, static_cast(config.pixel_dim)); + sd::PECacheKey joint_key("iris.joint_pe"); + joint_key.add(h).add(w).add(nt).add(config.patch_size).add(config.head_dim); + auto pe = make_input(pe_cache_.get_or_create(joint_key, [&]() { + auto values = Pid::make_rope_1d(nt, static_cast(config.head_dim), 10000.f); + auto img = image_rope(h / config.patch_size, w / config.patch_size, config.head_dim); + values.insert(values.end(), img.begin(), img.end()); + const int64_t tokens = nt + static_cast(img.size() / (2 * config.head_dim)); + return sd::Tensor({2, 2, config.head_dim / 2, tokens}, std::move(values)); + })); + sd::PECacheKey pixel_key("iris.pixel_pe"); + pixel_key.add(h).add(w).add(config.patch_size).add(config.pixel_attn_dim).add(config.pixel_heads); + auto pp = make_input(pe_cache_.get_or_create(pixel_key, [&]() { + return sd::Tensor({2, 2, config.pixel_attn_dim / config.pixel_heads / 2, + (h / config.patch_size) * (w / config.patch_size)}, + image_rope(h / config.patch_size, w / config.patch_size, config.pixel_attn_dim / config.pixel_heads)); + })); + sd::PECacheKey pos_key("iris.pixel_pos"); + pos_key.add(h).add(w).add(config.pixel_dim); + auto pos = make_input(pe_cache_.get_or_create(pos_key, [&]() { + return sd::Tensor({config.pixel_dim, h * w, 1}, + Pid::make_pixel_abs_pos(h, w, static_cast(config.pixel_dim))); + })); text_mask_data.resize(nt * nt); for (int q = 0; q < nt; ++q) { for (int k = 0; k < nt; ++k) { text_mask_data[q * nt + k] = inputs.y->values()[k] != 0.f || q == k ? 0.f : -INFINITY; } } - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, nt + img.size() / (2 * config.head_dim)); - auto pp = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.pixel_attn_dim / config.pixel_heads / 2, (h / config.patch_size) * (w / config.patch_size)); - auto pos = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.pixel_dim, h * w, 1); auto mask = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, nt, nt); - set_backend_tensor_data(pe, joint_pe_data.data()); - set_backend_tensor_data(pp, pixel_pe_data.data()); - set_backend_tensor_data(pos, pixel_pos_data.data()); set_backend_tensor_data(mask, text_mask_data.data()); auto ctx = get_context(); auto out = model.forward(&ctx, x, t, text, mask, pe, pp, pos); diff --git a/src/model/diffusion/krea2.hpp b/src/model/diffusion/krea2.hpp index e58df4108..1558e4e9b 100644 --- a/src/model/diffusion/krea2.hpp +++ b/src/model/diffusion/krea2.hpp @@ -11,6 +11,7 @@ #include #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_graph_cut.h" @@ -73,9 +74,9 @@ namespace Krea2 { int64_t dim_head = head_dim(); int64_t unit = dim_head / 16; axes_dim = { - static_cast(dim_head - 12 * unit), - static_cast(6 * unit), - static_cast(6 * unit), + static_cast(dim_head - 12 * unit), + static_cast(6 * unit), + static_cast(6 * unit), }; axes_dim_sum = axes_dim[0] + axes_dim[1] + axes_dim[2]; } @@ -387,10 +388,10 @@ namespace Krea2 { this->config.mlp_multiplier, this->config.norm_eps); blocks["refiner_blocks." + std::to_string(i)] = std::make_shared(this->config.text_dim, - this->config.text_heads, - this->config.text_kv_heads, - this->config.mlp_multiplier, - this->config.norm_eps); + this->config.text_heads, + this->config.text_kv_heads, + this->config.mlp_multiplier, + this->config.norm_eps); } blocks["projector"] = std::make_shared(this->config.text_layers, 1, false); } @@ -613,8 +614,8 @@ namespace Krea2 { explicit Krea2Model(Krea2Config config) : config(std::move(config)) { blocks["first"] = std::make_shared(this->config.patch_size * this->config.patch_size * this->config.in_channels, - this->config.features, - true); + this->config.features, + true); blocks["tmlp"] = std::make_shared(this->config); blocks["txtfusion"] = std::make_shared(this->config); blocks["txtmlp"] = std::make_shared(this->config); @@ -714,9 +715,9 @@ namespace Krea2 { } struct Krea2Runner : public DiffusionModelRunner { + sd::PECache pe_cache_; Krea2Config config; Krea2Model model; - std::vector pe_vec; Krea2Runner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -754,18 +755,23 @@ namespace Krea2 { ref_latents.push_back(make_input(ref_latent_tensor)); } - pe_vec = finish_rope_pe(gen_krea2_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - config.axes_dim, - ref_latents, - ref_image_params.ref_index_mode)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("krea2.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shapes(ref_latents); + pe_key.add(ref_image_params.ref_index_mode); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(gen_krea2_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim, + ref_latents, + ref_image_params.ref_index_mode)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents, ref_image_params.force_ref_timestep_zero); diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp index 8fb1512b4..76f7bee1f 100644 --- a/src/model/diffusion/lens.hpp +++ b/src/model/diffusion/lens.hpp @@ -4,6 +4,7 @@ #include #include #include +#include "core/pe_cache.h" #include "model/common/block.hpp" #include "model/common/rope.hpp" @@ -352,9 +353,9 @@ namespace Lens { }; struct LensRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; LensConfig config; LensModel lens; - std::vector pe_vec; LensRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -384,15 +385,19 @@ namespace Lens { GGML_ASSERT(!context_tensor.empty()); ggml_tensor* context = make_input(context_tensor); - pe_vec = finish_rope_pe(Rope::gen_lens_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("lens.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_lens_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = lens.forward(&runner_ctx, x, timesteps, context, pe); diff --git a/src/model/diffusion/lingbot_video.hpp b/src/model/diffusion/lingbot_video.hpp index 1a93b458e..194bf7912 100644 --- a/src/model/diffusion/lingbot_video.hpp +++ b/src/model/diffusion/lingbot_video.hpp @@ -9,6 +9,7 @@ #include #include #include +#include "core/pe_cache.h" #include "core/util.h" #include "model/common/rope.hpp" @@ -442,10 +443,10 @@ namespace LingBotVideo { : hidden_size(config.hidden_size) { blocks["norm1"] = std::make_shared(config.hidden_size, config.norm_eps); blocks["attn"] = std::make_shared(config.hidden_size, - config.num_attention_heads, - config.qkv_bias, - config.out_bias, - config.norm_eps); + config.num_attention_heads, + config.qkv_bias, + config.out_bias, + config.norm_eps); blocks["norm_post_attn"] = std::make_shared(config.hidden_size, config.norm_eps); blocks["norm2"] = std::make_shared(config.hidden_size, config.norm_eps); if (sparse) { @@ -509,14 +510,14 @@ namespace LingBotVideo { int64_t patch_dim = config.patch_t * config.patch_h * config.patch_w * config.in_channels; blocks["patch_embedder"] = std::make_shared(patch_dim, config.hidden_size, config.patch_embed_bias); blocks["time_embedder"] = std::make_shared(config.freq_dim, - config.hidden_size, - config.hidden_size, - 0, - config.timestep_mlp_bias); + config.hidden_size, + config.hidden_size, + 0, + config.timestep_mlp_bias); blocks["time_modulation.1"] = std::make_shared(config.hidden_size, 6 * config.hidden_size, true); blocks["text_embedder"] = std::make_shared(config.text_dim, - config.hidden_size, - config.norm_eps); + config.hidden_size, + config.norm_eps); for (int i = 0; i < config.depth; i++) { bool sparse = config.sparse_layers.find(i) != config.sparse_layers.end(); blocks["blocks." + std::to_string(i)] = std::make_shared(config, sparse); @@ -524,8 +525,8 @@ namespace LingBotVideo { blocks["norm_out"] = std::make_shared(config.hidden_size, config.norm_eps, false); blocks["norm_out_modulation.1"] = std::make_shared(config.hidden_size, 2 * config.hidden_size, true); blocks["proj_out"] = std::make_shared(config.hidden_size, - config.patch_t * config.patch_h * config.patch_w * config.out_channels, - true); + config.patch_t * config.patch_h * config.patch_w * config.out_channels, + true); } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -606,9 +607,9 @@ namespace LingBotVideo { }; struct LingBotVideoRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; LingBotVideoConfig config; LingBotVideoModel lingbot_video; - std::vector pe_vec; LingBotVideoRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -647,19 +648,23 @@ namespace LingBotVideo { int64_t T = x->ne[2]; int64_t N = x_tensor.shape()[4]; GGML_ASSERT(N == 1); - pe_vec = Rope::gen_lingbot_video_pe(static_cast(T), - static_cast(H), - static_cast(W), - config.patch_t, - config.patch_h, - config.patch_w, - static_cast(N), - static_cast(context->ne[1]), - config.rope_theta, - config.axes_dim); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("lingbot_video.pe"); + pe_key.add(T).add(H).add(W).add(N).add_shape(context).add(config.patch_t).add(config.patch_h).add(config.patch_w); + pe_key.add(config.rope_theta).add(config.axes_dim).add(config.axes_dim_sum); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = Rope::gen_lingbot_video_pe(static_cast(T), + static_cast(H), + static_cast(W), + config.patch_t, + config.patch_h, + config.patch_w, + static_cast(N), + static_cast(context->ne[1]), + config.rope_theta, + config.axes_dim); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = lingbot_video.forward(&runner_ctx, x, timesteps, context, pe); diff --git a/src/model/diffusion/llada_image.hpp b/src/model/diffusion/llada_image.hpp index e69c3ecd3..4a9ea1e5a 100644 --- a/src/model/diffusion/llada_image.hpp +++ b/src/model/diffusion/llada_image.hpp @@ -3,6 +3,7 @@ #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_runner.h" @@ -379,10 +380,11 @@ namespace LLaDAImage { }; struct LLaDAImageRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: LLaDAImageConfig config; LLaDAImageModel llada_image; - std::vector pe_vec; LLaDAImageRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -412,17 +414,21 @@ namespace LLaDAImage { GGML_ASSERT(!context_tensor.empty()); ggml_tensor* context = make_input(context_tensor); - pe_vec = finish_rope_pe(Rope::gen_llada_image_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - ZImage::SEQ_MULTI_OF, - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("llada_image.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_llada_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ZImage::SEQ_MULTI_OF, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = llada_image.forward(&runner_ctx, x, timesteps, context, pe); @@ -459,17 +465,21 @@ namespace LLaDAImage { ggml_tensor* source = make_input(source_tensor); GGML_ASSERT(x->ne[3] == 1); - pe_vec = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(context->ne[1]), - semantic != nullptr ? static_cast(semantic->ne[1]) : 0, - ZImage::SEQ_MULTI_OF, - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("llada_image.pe1"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shape(semantic); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(context->ne[1]), + semantic != nullptr ? static_cast(semantic->ne[1]) : 0, + ZImage::SEQ_MULTI_OF, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); int64_t W = x->ne[0]; diff --git a/src/model/diffusion/ltxv.hpp b/src/model/diffusion/ltxv.hpp index e10d5e829..23c5f5f29 100644 --- a/src/model/diffusion/ltxv.hpp +++ b/src/model/diffusion/ltxv.hpp @@ -11,6 +11,7 @@ #include #include "core/ggml_extend_backend.h" #include "core/ggml_tensor_utils.h" +#include "core/pe_cache.h" #include "model/common/block.hpp" #include "model/common/rope.hpp" @@ -597,9 +598,9 @@ namespace LTXV { : embedding_dim(embedding_dim), embedding_coefficient(embedding_coefficient) { blocks["emb.timestep_embedder"] = std::make_shared(embedding_dim); blocks["linear"] = std::make_shared(embedding_dim, - embedding_coefficient * embedding_dim, - true, - true); + embedding_coefficient * embedding_dim, + true, + true); } std::pair forward(GGMLRunnerContext* ctx, @@ -1257,7 +1258,7 @@ namespace LTXV { v_context, attn2.get(), v_table, - cross_attention_adaln ? params["prompt_scale_shift_table"] : nullptr, + cross_attention_adaln ? params["prompt_scale_shift_table"] : nullptr, v_timestep, v_prompt_timestep, v_dim, @@ -1275,7 +1276,7 @@ namespace LTXV { a_context, audio_attn2.get(), a_table, - cross_attention_adaln ? params["audio_prompt_scale_shift_table"] : nullptr, + cross_attention_adaln ? params["audio_prompt_scale_shift_table"] : nullptr, a_timestep, a_prompt_timestep, a_dim, @@ -1736,14 +1737,9 @@ namespace LTXV { }; struct LTXAVRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; LTXAVConfig config; LTXAVModelBlock model; - std::vector video_pe_vec; - std::vector audio_pe_vec; - std::vector video_cross_pe_vec; - std::vector audio_cross_pe_vec; - std::vector connector_pe_vec; - std::vector audio_connector_pe_vec; sd::Tensor vx_input_cache; sd::Tensor ax_input_cache; @@ -1853,79 +1849,101 @@ namespace LTXV { float video_frame_rate = frame_rate > 0.f ? frame_rate : 24.f; int64_t video_token_count = vx->ne[0] * vx->ne[1] * vx->ne[2]; bool has_video_positions = !video_positions_tensor.empty(); - if (has_video_positions) { - GGML_ASSERT(video_positions_tensor.shape()[2] == video_token_count); - video_pe_vec = build_video_rope_matrix_from_positions(video_positions_tensor, - static_cast(config.hidden_size), - static_cast(config.num_attention_heads), - config.positional_embedding_theta, - config.positional_embedding_max_pos, - config.use_middle_indices_grid); - } else { - video_pe_vec = build_video_rope_matrix(vx->ne[0], - vx->ne[1], - vx->ne[2], - static_cast(config.hidden_size), - static_cast(config.num_attention_heads), - video_frame_rate, - config.positional_embedding_theta, - config.positional_embedding_max_pos, - config.vae_scale_factors, - config.causal_temporal_positioning, - config.use_middle_indices_grid); - } - auto video_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.attention_head_dim / 2, video_token_count * config.num_attention_heads); - ggml_set_name(video_pe, "ltxav_video_pe"); - set_backend_tensor_data(video_pe, video_pe_vec.data()); + sd::PECacheKey video_pe_key("ltxav_video_pe"); + video_pe_key.add_shape(vx).add(video_frame_rate).add(video_positions_tensor).add(config.positional_embedding_theta); + video_pe_key.add(config.causal_temporal_positioning).add(config.use_middle_indices_grid).add(config.hidden_size); + video_pe_key.add(config.num_attention_heads).add(config.attention_head_dim).add(config.positional_embedding_max_pos); + video_pe_key.add(std::get<0>(config.vae_scale_factors)).add(std::get<1>(config.vae_scale_factors)); + video_pe_key.add(std::get<2>(config.vae_scale_factors)); + auto video_pe = make_input(pe_cache_.get_or_create(video_pe_key, [&]() { + std::vector video_pe_vec; + if (has_video_positions) { + GGML_ASSERT(video_positions_tensor.shape()[2] == video_token_count); + video_pe_vec = build_video_rope_matrix_from_positions(video_positions_tensor, + static_cast(config.hidden_size), + static_cast(config.num_attention_heads), + config.positional_embedding_theta, + config.positional_embedding_max_pos, + config.use_middle_indices_grid); + } else { + video_pe_vec = build_video_rope_matrix(vx->ne[0], + vx->ne[1], + vx->ne[2], + static_cast(config.hidden_size), + static_cast(config.num_attention_heads), + video_frame_rate, + config.positional_embedding_theta, + config.positional_embedding_max_pos, + config.vae_scale_factors, + config.causal_temporal_positioning, + config.use_middle_indices_grid); + } + return sd::Tensor({2, 2, config.attention_head_dim / 2, video_token_count * config.num_attention_heads}, std::move(video_pe_vec)); + })); ggml_tensor* audio_pe = nullptr; ggml_tensor* video_cross_pe = nullptr; ggml_tensor* audio_cross_pe = nullptr; if (ax != nullptr && ggml_nelements(ax) > 0 && ax->ne[1] > 0) { - audio_pe_vec = build_audio_rope_matrix(ax->ne[1], - static_cast(config.audio_hidden_size), - static_cast(config.audio_num_attention_heads), - config.positional_embedding_theta, - config.audio_positional_embedding_max_pos[0], - config.use_middle_indices_grid); - audio_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads); - ggml_set_name(audio_pe, "ltxav_audio_pe"); - set_backend_tensor_data(audio_pe, audio_pe_vec.data()); + sd::PECacheKey audio_pe_key("ltxav_audio_pe"); + audio_pe_key.add(ax->ne[1]).add(config.audio_hidden_size).add(config.audio_num_attention_heads); + audio_pe_key.add(config.audio_attention_head_dim).add(config.positional_embedding_theta); + audio_pe_key.add(config.audio_positional_embedding_max_pos[0]).add(config.use_middle_indices_grid); + audio_pe = make_input(pe_cache_.get_or_create(audio_pe_key, [&]() { + auto audio_pe_vec = build_audio_rope_matrix(ax->ne[1], + static_cast(config.audio_hidden_size), + static_cast(config.audio_num_attention_heads), + config.positional_embedding_theta, + config.audio_positional_embedding_max_pos[0], + config.use_middle_indices_grid); + return sd::Tensor({2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads}, std::move(audio_pe_vec)); + })); int temporal_max_pos = std::max(config.positional_embedding_max_pos[0], config.audio_positional_embedding_max_pos[0]); - if (has_video_positions) { - video_cross_pe_vec = build_video_temporal_rope_matrix_from_positions(video_positions_tensor, - static_cast(config.audio_cross_attention_dim), - static_cast(config.audio_num_attention_heads), - config.positional_embedding_theta, - temporal_max_pos, - true); - } else { - video_cross_pe_vec = build_video_temporal_rope_matrix(vx->ne[0], - vx->ne[1], - vx->ne[2], - static_cast(config.audio_cross_attention_dim), - static_cast(config.audio_num_attention_heads), - video_frame_rate, - config.positional_embedding_theta, - temporal_max_pos, - std::get<0>(config.vae_scale_factors), - config.causal_temporal_positioning, - true); - } - video_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, video_token_count * config.audio_num_attention_heads); - ggml_set_name(video_cross_pe, "ltxav_video_cross_pe"); - set_backend_tensor_data(video_cross_pe, video_cross_pe_vec.data()); - - audio_cross_pe_vec = build_audio_rope_matrix(ax->ne[1], - static_cast(config.audio_cross_attention_dim), - static_cast(config.audio_num_attention_heads), - config.positional_embedding_theta, - temporal_max_pos, - true); - audio_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads); - ggml_set_name(audio_cross_pe, "ltxav_audio_cross_pe"); - set_backend_tensor_data(audio_cross_pe, audio_cross_pe_vec.data()); + sd::PECacheKey video_cross_pe_key("ltxav_video_cross_pe"); + video_cross_pe_key.add_shape(vx).add(video_frame_rate).add(video_positions_tensor); + video_cross_pe_key.add(config.positional_embedding_theta).add(config.causal_temporal_positioning); + video_cross_pe_key.add(config.use_middle_indices_grid).add(config.audio_cross_attention_dim); + video_cross_pe_key.add(config.audio_num_attention_heads).add(config.audio_attention_head_dim); + video_cross_pe_key.add(temporal_max_pos).add(std::get<0>(config.vae_scale_factors)); + video_cross_pe = make_input(pe_cache_.get_or_create(video_cross_pe_key, [&]() { + std::vector video_cross_pe_vec; + if (has_video_positions) { + video_cross_pe_vec = build_video_temporal_rope_matrix_from_positions(video_positions_tensor, + static_cast(config.audio_cross_attention_dim), + static_cast(config.audio_num_attention_heads), + config.positional_embedding_theta, + temporal_max_pos, + true); + } else { + video_cross_pe_vec = build_video_temporal_rope_matrix(vx->ne[0], + vx->ne[1], + vx->ne[2], + static_cast(config.audio_cross_attention_dim), + static_cast(config.audio_num_attention_heads), + video_frame_rate, + config.positional_embedding_theta, + temporal_max_pos, + std::get<0>(config.vae_scale_factors), + config.causal_temporal_positioning, + true); + } + return sd::Tensor({2, 2, config.audio_attention_head_dim / 2, video_token_count * config.audio_num_attention_heads}, std::move(video_cross_pe_vec)); + })); + + sd::PECacheKey audio_cross_pe_key("ltxav_audio_cross_pe"); + audio_cross_pe_key.add(ax->ne[1]).add(config.audio_cross_attention_dim).add(config.audio_num_attention_heads); + audio_cross_pe_key.add(config.audio_attention_head_dim).add(config.positional_embedding_theta); + audio_cross_pe_key.add(temporal_max_pos); + audio_cross_pe = make_input(pe_cache_.get_or_create(audio_cross_pe_key, [&]() { + auto audio_cross_pe_vec = build_audio_rope_matrix(ax->ne[1], + static_cast(config.audio_cross_attention_dim), + static_cast(config.audio_num_attention_heads), + config.positional_embedding_theta, + temporal_max_pos, + true); + return sd::Tensor({2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads}, std::move(audio_cross_pe_vec)); + })); } bool needs_video_connector_pe = @@ -1941,10 +1959,13 @@ namespace LTXV { int64_t target_len = std::max(1024, seq_len); int64_t duplications = (target_len + config.connector_num_registers - 1) / config.connector_num_registers; int64_t full_len = seq_len + duplications * config.connector_num_registers - seq_len; - connector_pe_vec = build_1d_rope_matrix(full_len, static_cast(config.connector_hidden_size), static_cast(config.connector_num_heads), 10000.f, 4096.f, true); - video_connector_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.connector_head_dim / 2, full_len * config.connector_num_heads); - ggml_set_name(video_connector_pe, "ltxav_video_connector_pe"); - set_backend_tensor_data(video_connector_pe, connector_pe_vec.data()); + sd::PECacheKey video_connector_pe_key("ltxav_video_connector_pe"); + video_connector_pe_key.add(full_len).add(config.connector_hidden_size).add(config.connector_num_heads); + video_connector_pe_key.add(config.connector_head_dim); + video_connector_pe = make_input(pe_cache_.get_or_create(video_connector_pe_key, [&]() { + auto connector_pe_vec = build_1d_rope_matrix(full_len, static_cast(config.connector_hidden_size), static_cast(config.connector_num_heads), 10000.f, 4096.f, true); + return sd::Tensor({2, 2, config.connector_head_dim / 2, full_len * config.connector_num_heads}, std::move(connector_pe_vec)); + })); } bool run_audio_context = @@ -1961,14 +1982,17 @@ namespace LTXV { context->ne[1] < 1024)); ggml_tensor* audio_connector_pe = nullptr; if (needs_audio_connector_pe) { - int64_t seq_len = context->ne[1]; - int64_t target_len = std::max(1024, seq_len); - int64_t duplications = (target_len + config.audio_connector_num_registers - 1) / config.audio_connector_num_registers; - int64_t full_len = seq_len + duplications * config.audio_connector_num_registers - seq_len; - audio_connector_pe_vec = build_1d_rope_matrix(full_len, static_cast(config.audio_connector_hidden_size), static_cast(config.audio_connector_num_heads), 10000.f, 4096.f, true); - audio_connector_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_connector_head_dim / 2, full_len * config.audio_connector_num_heads); - ggml_set_name(audio_connector_pe, "ltxav_audio_connector_pe"); - set_backend_tensor_data(audio_connector_pe, audio_connector_pe_vec.data()); + int64_t seq_len = context->ne[1]; + int64_t target_len = std::max(1024, seq_len); + int64_t duplications = (target_len + config.audio_connector_num_registers - 1) / config.audio_connector_num_registers; + int64_t full_len = seq_len + duplications * config.audio_connector_num_registers - seq_len; + sd::PECacheKey audio_connector_pe_key("ltxav_audio_connector_pe"); + audio_connector_pe_key.add(full_len).add(config.audio_connector_hidden_size).add(config.audio_connector_num_heads); + audio_connector_pe_key.add(config.audio_connector_head_dim); + audio_connector_pe = make_input(pe_cache_.get_or_create(audio_connector_pe_key, [&]() { + auto audio_connector_pe_vec = build_1d_rope_matrix(full_len, static_cast(config.audio_connector_hidden_size), static_cast(config.audio_connector_num_heads), 10000.f, 4096.f, true); + return sd::Tensor({2, 2, config.audio_connector_head_dim / 2, full_len * config.audio_connector_num_heads}, std::move(audio_connector_pe_vec)); + })); } auto runner_ctx = get_context(); diff --git a/src/model/diffusion/mage_flow.hpp b/src/model/diffusion/mage_flow.hpp index 1b9b9f9be..415c0f722 100644 --- a/src/model/diffusion/mage_flow.hpp +++ b/src/model/diffusion/mage_flow.hpp @@ -3,6 +3,7 @@ #include #include +#include "core/pe_cache.h" #include "model/diffusion/qwen_image.hpp" @@ -57,11 +58,12 @@ namespace MageFlow { }; struct MageFlowRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: Qwen::QwenImageConfig config; Qwen::QwenImageModel mage_flow; std::shared_ptr time_text_embed; - std::vector pe_vec; MageFlowRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -110,16 +112,20 @@ namespace MageFlow { } int batch_size = static_cast(x->ne[3]); - pe_vec = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - batch_size, - static_cast(context->ne[1]), - ref_latents, - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("mage_flow.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(batch_size).add_shapes(ref_latents); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + batch_size, + static_cast(context->ne[1]), + ref_latents, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); time_text_embed->prepare(timesteps_tensor); auto runner_ctx = get_context(); diff --git a/src/model/diffusion/ming_image.hpp b/src/model/diffusion/ming_image.hpp index f19be9903..1de9ac157 100644 --- a/src/model/diffusion/ming_image.hpp +++ b/src/model/diffusion/ming_image.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__ #define __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__ +#include "core/pe_cache.h" #include "z_image.hpp" namespace MingImage { @@ -72,9 +73,9 @@ namespace MingImage { }; struct MingImageRunner : DiffusionModelRunner { + sd::PECache pe_cache_; MingImageConfig config; MingImageModel model; - std::vector pe_values; MingImageRunner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr weight_manager = nullptr) : DiffusionModelRunner(backend, prefix, weight_manager), @@ -109,19 +110,23 @@ namespace MingImage { const int64_t n_txt = context->ne[1] + direct->ne[1]; const int64_t n_img = ((x->ne[0] + config.patch_size - 1) / config.patch_size) * ((x->ne[1] + config.patch_size - 1) / config.patch_size); - auto padded = finish_rope_pe(Rope::gen_z_image_pe( - static_cast(x->ne[1]), static_cast(x->ne[0]), config.patch_size, 1, - static_cast(n_txt), ZImage::SEQ_MULTI_OF, {}, Rope::RefIndexMode::FIXED, - config.theta, config.axes_dim)); - // Zero-masked alignment tokens cannot affect valid queries. Omit them while - // retaining the padded caption length used to position image tokens. - const size_t stride = config.axes_dim_sum * 2; - const int64_t padded_txt = n_txt + Rope::bound_mod(static_cast(n_txt), ZImage::SEQ_MULTI_OF); - pe_values.assign(padded.begin(), padded.begin() + n_txt * stride); - pe_values.insert(pe_values.end(), padded.begin() + padded_txt * stride, - padded.begin() + (padded_txt + n_img) * stride); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, n_txt + n_img); - set_backend_tensor_data(pe, pe_values.data()); + sd::PECacheKey pe_key("ming_image.pe"); + pe_key.add_shape(x).add(n_txt).add(config.patch_size).add(config.theta).add(config.axes_dim); + pe_key.add(config.axes_dim_sum).add(circular_x_enabled).add(circular_y_enabled); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto padded = finish_rope_pe(Rope::gen_z_image_pe( + static_cast(x->ne[1]), static_cast(x->ne[0]), config.patch_size, 1, + static_cast(n_txt), ZImage::SEQ_MULTI_OF, {}, Rope::RefIndexMode::FIXED, + config.theta, config.axes_dim)); + // Zero-masked alignment tokens cannot affect valid queries. Omit them while + // retaining the padded caption length used to position image tokens. + const size_t stride = config.axes_dim_sum * 2; + const int64_t padded_txt = n_txt + Rope::bound_mod(static_cast(n_txt), ZImage::SEQ_MULTI_OF); + std::vector pe_values(padded.begin(), padded.begin() + n_txt * stride); + pe_values.insert(pe_values.end(), padded.begin() + padded_txt * stride, + padded.begin() + (padded_txt + n_img) * stride); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, n_txt + n_img}, std::move(pe_values)); + })); auto ctx = get_context(); auto out = model.forward(&ctx, x, t, context, direct, pe); ggml_build_forward_expand(gf, out); diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 7a8693ef3..e93c70f33 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -12,6 +12,7 @@ #include "core/ggml_tensor_utils.h" #include "core/ggml_graph_cut.h" +#include "model/common/block.hpp" #include "model/diffusion/dit.hpp" #include "model/diffusion/model.hpp" @@ -234,9 +235,9 @@ namespace MiniMaxH3 { blocks["norm1"] = std::make_shared(config.hidden_size, config.norm_eps); blocks["norm2"] = std::make_shared(config.hidden_size, config.norm_eps); blocks["attn"] = std::make_shared(config.hidden_size, - config.num_attention_heads, - config.attention_head_dim, - config.qk_norm_eps); + config.num_attention_heads, + config.attention_head_dim, + config.qk_norm_eps); blocks["mlp"] = std::make_shared(config.hidden_size, config.ffn_hidden_size); } @@ -391,11 +392,11 @@ namespace MiniMaxH3 { blocks["norm1"] = std::make_shared(config.hidden_size, config.norm_eps); blocks["norm2"] = std::make_shared(config.hidden_size, config.norm_eps); blocks["attn"] = std::make_shared(config.hidden_size, - config.num_attention_heads, - config.attention_head_dim, - config.qk_norm_eps); + config.num_attention_heads, + config.attention_head_dim, + config.qk_norm_eps); blocks["mlp"] = std::make_shared(config.hidden_size, - config.ffn_hidden_size); + config.ffn_hidden_size); blocks["adaln_proj"] = std::make_shared(config.time_embed_dim, config.hidden_size, 6, @@ -482,8 +483,8 @@ namespace MiniMaxH3 { auto row = modulation_row(ctx->ggml_ctx, mods, config.hidden_size, 2, 1, segment.modulation_row); auto value = norm->forward(ctx, ggml_ext_slice(ctx->ggml_ctx, x, 1, segment.start, segment.end)); return ggml_add(ctx->ggml_ctx, - ggml_add(ctx->ggml_ctx, value, ggml_mul(ctx->ggml_ctx, value, row[1])), - row[0]); + ggml_add(ctx->ggml_ctx, value, ggml_mul(ctx->ggml_ctx, value, row[1])), + row[0]); }; return {video_out->forward(ctx, apply(video)), audio_out->forward(ctx, apply(audio))}; @@ -846,10 +847,10 @@ namespace MiniMaxH3 { const auto& condition = condition_videos[index]; auto axes = condition_spatial_axes(condition); int64_t count = condition.shape()[2] * - static_cast(axes.first.size() * axes.second.size()); - bool is_first = keyframe_indices.empty() || keyframe_indices[static_cast(index)] == 0; - float keyframe_t = is_first ? static_cast(text_len) - : static_cast(text_len) + video_duration - FRAME_RESCALE; + static_cast(axes.first.size() * axes.second.size()); + bool is_first = keyframe_indices.empty() || keyframe_indices[static_cast(index)] == 0; + float keyframe_t = is_first ? static_cast(text_len) + : static_cast(text_len) + video_duration - FRAME_RESCALE; for (int64_t t = 0; t < condition.shape()[2]; ++t) { for (float h : axes.first) { for (float w : axes.second) { @@ -903,9 +904,9 @@ namespace MiniMaxH3 { if (block.kind != MiniMaxH3ReferenceKind::AUDIO) { GGML_ASSERT(ref_video != nullptr); - auto axes = condition_spatial_axes(*ref_video); - int64_t count = ref_video->shape()[2] * - static_cast(axes.first.size() * axes.second.size()); + auto axes = condition_spatial_axes(*ref_video); + int64_t count = ref_video->shape()[2] * + static_cast(axes.first.size() * axes.second.size()); float video_end = append_video_positions(*ref_video, cursor); layout.sequence_segments.push_back({row, row + count, diff --git a/src/model/diffusion/minit2i.hpp b/src/model/diffusion/minit2i.hpp index 0ddacf02e..6dc9c1798 100644 --- a/src/model/diffusion/minit2i.hpp +++ b/src/model/diffusion/minit2i.hpp @@ -9,6 +9,7 @@ #include #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_runner.h" @@ -473,20 +474,9 @@ namespace MiniT2I { }; struct MiniT2IRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; MiniT2IConfig config; MMJiT model; - ggml_context* position_cache_ctx = nullptr; - ggml_backend_buffer_t position_cache_buffer = nullptr; - ggml_tensor* cached_pos_embed = nullptr; - ggml_tensor* cached_txt_pe = nullptr; - ggml_tensor* cached_joint_pe = nullptr; - int64_t cached_img_side = -1; - int64_t cached_txt_len = -1; - int64_t cached_hidden_size = -1; - int64_t cached_head_dim = -1; - bool cached_circular_x = false; - bool cached_circular_y = false; - MiniT2IRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "", @@ -497,10 +487,6 @@ namespace MiniT2I { model.init(params_ctx, tensor_storage_map, this->prefix); } - ~MiniT2IRunner() override { - free_position_cache(); - } - std::string get_desc() override { return "MiniT2I"; } @@ -509,75 +495,6 @@ namespace MiniT2I { model.get_param_tensors(tensors, prefix); } - void free_position_cache() { - if (position_cache_buffer != nullptr) { - ggml_backend_buffer_free(position_cache_buffer); - position_cache_buffer = nullptr; - } - if (position_cache_ctx != nullptr) { - ggml_free(position_cache_ctx); - position_cache_ctx = nullptr; - } - cached_pos_embed = nullptr; - cached_txt_pe = nullptr; - cached_joint_pe = nullptr; - cached_img_side = -1; - cached_txt_len = -1; - cached_hidden_size = -1; - cached_head_dim = -1; - } - - void ensure_position_cache(int64_t img_side, int64_t txt_len) { - if (cached_img_side == img_side && - cached_txt_len == txt_len && - cached_hidden_size == config.hidden_size && - cached_head_dim == config.head_dim && - cached_circular_x == circular_x_enabled && - cached_circular_y == circular_y_enabled && - cached_pos_embed != nullptr && - cached_txt_pe != nullptr && - cached_joint_pe != nullptr) { - return; - } - - free_position_cache(); - - auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast(img_side), static_cast(config.hidden_size)); - auto txt_pe_vec = make_text_rope(static_cast(txt_len), static_cast(config.head_dim)); - auto img_pe_vec = finish_rope_pe(make_vision_rope(static_cast(img_side), static_cast(config.head_dim))); - auto joint_pe_vec = txt_pe_vec; - joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end()); - - ggml_init_params params; - params.mem_size = static_cast(3 * ggml_tensor_overhead()); - params.mem_buffer = nullptr; - params.no_alloc = true; - position_cache_ctx = ggml_init(params); - GGML_ASSERT(position_cache_ctx != nullptr); - - cached_pos_embed = ggml_new_tensor_3d(position_cache_ctx, GGML_TYPE_F32, config.hidden_size, img_side * img_side, 1); - ggml_set_name(cached_pos_embed, "minit2i.pos_embed"); - cached_txt_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len); - ggml_set_name(cached_txt_pe, "minit2i.txt_pe"); - cached_joint_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len + img_side * img_side); - ggml_set_name(cached_joint_pe, "minit2i.joint_pe"); - - position_cache_buffer = ggml_backend_alloc_ctx_tensors(position_cache_ctx, runtime_backend); - GGML_ASSERT(position_cache_buffer != nullptr); - ggml_backend_buffer_set_usage(position_cache_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - ggml_backend_tensor_set(cached_pos_embed, pos_embed_vec.data(), 0, ggml_nbytes(cached_pos_embed)); - ggml_backend_tensor_set(cached_txt_pe, txt_pe_vec.data(), 0, ggml_nbytes(cached_txt_pe)); - ggml_backend_tensor_set(cached_joint_pe, joint_pe_vec.data(), 0, ggml_nbytes(cached_joint_pe)); - ggml_backend_synchronize(runtime_backend); - - cached_img_side = img_side; - cached_txt_len = txt_len; - cached_hidden_size = config.hidden_size; - cached_head_dim = config.head_dim; - cached_circular_x = circular_x_enabled; - cached_circular_y = circular_y_enabled; - } - ggml_cgraph* build_graph(const sd::Tensor& x_tensor, const sd::Tensor& timesteps_tensor, const sd::Tensor& context_tensor, @@ -592,10 +509,29 @@ namespace MiniT2I { int64_t H = x->ne[1]; int64_t img_side = H / config.patch_size; int64_t txt_len = context->ne[1]; - ensure_position_cache(img_side, txt_len); + sd::PECacheKey pos_key("minit2i.pos_embed"); + pos_key.add(img_side).add(config.hidden_size); + auto pos_embed = make_input(pe_cache_.get_or_create(pos_key, [&]() { + return sd::Tensor({config.hidden_size, img_side * img_side, 1}, + make_2d_sincos_pos_embed(static_cast(img_side), static_cast(config.hidden_size))); + })); + sd::PECacheKey txt_key("minit2i.txt_pe"); + txt_key.add(txt_len).add(config.head_dim); + auto txt_pe = make_input(pe_cache_.get_or_create(txt_key, [&]() { + return sd::Tensor({2, 2, config.head_dim / 2, txt_len}, + make_text_rope(static_cast(txt_len), static_cast(config.head_dim))); + })); + sd::PECacheKey joint_key("minit2i.joint_pe"); + joint_key.add(img_side).add(txt_len).add(config.head_dim).add(circular_x_enabled).add(circular_y_enabled); + auto joint_pe = make_input(pe_cache_.get_or_create(joint_key, [&]() { + auto values = make_text_rope(static_cast(txt_len), static_cast(config.head_dim)); + auto image = finish_rope_pe(make_vision_rope(static_cast(img_side), static_cast(config.head_dim))); + values.insert(values.end(), image.begin(), image.end()); + return sd::Tensor({2, 2, config.head_dim / 2, txt_len + img_side * img_side}, std::move(values)); + })); auto runner_ctx = get_context(); - auto out = model.forward(&runner_ctx, x, context, mask, cached_pos_embed, cached_txt_pe, cached_joint_pe); + auto out = model.forward(&runner_ctx, x, context, mask, pos_embed, txt_pe, joint_pe); ggml_build_forward_expand(gf, out); return gf; } diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp index 42599c67a..cdf65f300 100644 --- a/src/model/diffusion/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -14,6 +14,7 @@ #include "core/util.h" #include "model/common/block.hpp" #include "model/common/ggml_block.hpp" +#include "model/diffusion/dit.hpp" #include "model/diffusion/model.hpp" #include "model_loader.h" diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index 7c14d9d1f..4ff91c416 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -16,6 +16,11 @@ enum class RefImageResizeMode { AREA, }; +enum class RefImageSizeMode { + MATCH, + MAX, +}; + struct RefImageParams { bool pass_to_vlm = false; bool pass_to_dit = true; @@ -31,6 +36,7 @@ struct RefImageParams { bool vae_size_buckets = false; bool vae_center_crop = false; bool vlm_use_vae_input = false; + RefImageSizeMode ref_image_size = RefImageSizeMode::MATCH; }; const std::unordered_map REF_IMAGE_PRESETS = { @@ -165,6 +171,11 @@ struct MingImageDiffusionExtra { const sd::Tensor* direct_context = nullptr; }; +struct ZImageDiffusionExtra { + const sd::Tensor* control_latent = nullptr; + float strength = 1.f; +}; + using DiffusionExtraParams = std::variant; + MingImageDiffusionExtra, + ZImageDiffusionExtra>; struct DiffusionParams { const sd::Tensor* x = nullptr; diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp index f1cbbc7f3..0eace334e 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -3,6 +3,7 @@ #include #include +#include "core/pe_cache.h" #include "model/common/ggml_block.hpp" #include "model/common/rope.hpp" @@ -193,9 +194,10 @@ namespace MOSSImage { }; struct MOSSImageRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; MOSSImageConfig config; Model model; - std::vector pe_data, time_data; + std::vector time_data; MOSSImageRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr weight_manager = nullptr) : DiffusionModelRunner(backend, prefix, weight_manager), config(MOSSImageConfig::detect_from_weights(weights, prefix)), model(config) { @@ -224,31 +226,34 @@ namespace MOSSImage { auto text = make_input(*inputs.context); GGML_ASSERT(x->ne[3] == 1 && x->ne[2] == config.channels && text->ne[0] == config.text_dim); int64_t h = x->ne[1], w = x->ne[0], nt = text->ne[1]; - Rope::Embedding embedding; - embedding.positions.append_image(static_cast(h), static_cast(w)); - embedding.positions.images.back().height_axis = 2; - embedding.positions.images.back().width_axis = 3; - embedding.positions.append_tokens(nt); - for (int64_t token = 0; token < h * w + nt; ++token) { - float pos[4] = {0.f, float(token - h * w), 0.f, 0.f}; - if (token < h * w) { - pos[1] = float(nt); - pos[2] = float(token / w) - float(h - 1) / 2.f; - pos[3] = float(token % w) - float(w - 1) / 2.f; + sd::PECacheKey pe_key("moss_image.pe"); + pe_key.add(h).add(w).add(nt).add(config.head_dim).add(circular_x_enabled).add(circular_y_enabled); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + Rope::Embedding embedding; + embedding.positions.append_image(static_cast(h), static_cast(w)); + embedding.positions.images.back().height_axis = 2; + embedding.positions.images.back().width_axis = 3; + embedding.positions.append_tokens(nt); + for (int64_t token = 0; token < h * w + nt; ++token) { + float pos[4] = {0.f, float(token - h * w), 0.f, 0.f}; + if (token < h * w) { + pos[1] = float(nt); + pos[2] = float(token / w) - float(h - 1) / 2.f; + pos[3] = float(token % w) - float(w - 1) / 2.f; + } + embedding.ids.push_back({pos[0], pos[1], pos[2], pos[3]}); } - embedding.ids.push_back({pos[0], pos[1], pos[2], pos[3]}); - } - embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, std::vector(4, static_cast(config.head_dim / 4)), embedding.layout, &embedding.frequencies); - pe_data = finish_rope_pe(std::move(embedding)); + embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, std::vector(4, static_cast(config.head_dim / 4)), embedding.layout, &embedding.frequencies); + auto values = finish_rope_pe(std::move(embedding)); + return sd::Tensor({2, 2, config.head_dim / 2, h * w + nt}, std::move(values)); + })); time_data.resize(config.freq_dim); for (int64_t i = 0; i < config.freq_dim / 2; ++i) { double angle = inputs.timesteps->values()[0] * std::pow(10000.0, -double(i) / double(config.freq_dim / 2)); time_data[i] = float(std::cos(angle)); time_data[i + config.freq_dim / 2] = float(std::sin(angle)); } - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, h * w + nt); auto time = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.freq_dim, 1, 1); - set_backend_tensor_data(pe, pe_data.data()); set_backend_tensor_data(time, time_data.data()); auto ctx = get_context(); ggml_build_forward_expand(gf, model.forward(&ctx, x, time, text, pe, skip)); diff --git a/src/model/diffusion/pid.hpp b/src/model/diffusion/pid.hpp index 6edc83422..6c5a59be6 100644 --- a/src/model/diffusion/pid.hpp +++ b/src/model/diffusion/pid.hpp @@ -7,6 +7,7 @@ #include #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_runner.h" @@ -644,7 +645,7 @@ namespace Pid { int64_t effective_down_factor = config.lq_latent_down_factor / unpatchify_factor; float z_to_patch_ratio = static_cast(config.lq_sr_scale * effective_down_factor) / - static_cast(config.patch_size); + static_cast(config.patch_size); GGML_ASSERT(z_to_patch_ratio >= 1.0f); if (lq_latent->ne[0] != target_pW || lq_latent->ne[1] != target_pH) { lq_latent = ggml_interpolate(ctx->ggml_ctx, @@ -822,12 +823,9 @@ namespace Pid { }; struct PiDRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; PixelDiTConfig config; PixelDiT model; - std::vector pos_img_vec; - std::vector pos_txt_vec; - std::vector pixel_pos_vec; - std::vector pixel_pos_comp_vec; PiDRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map, @@ -867,57 +865,53 @@ namespace Pid { int64_t Hs = Hp / config.patch_size; int64_t Ws = Wp / config.patch_size; - pos_img_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), - static_cast(Ws), - static_cast(config.hidden_size / config.num_groups), - 10000.f, - 16.f, - static_cast(config.rope_ref_grid_h), - static_cast(config.rope_ref_grid_w))); - auto pos_img = ggml_new_tensor_4d(compute_ctx, - GGML_TYPE_F32, - 2, - 2, - config.hidden_size / config.num_groups / 2, - Hs * Ws); - set_backend_tensor_data(pos_img, pos_img_vec.data()); + sd::PECacheKey pos_img_key("pid.pos_img"); + pos_img_key.add(Hs).add(Ws).add(config.hidden_size).add(config.num_groups).add(config.rope_ref_grid_h); + pos_img_key.add(config.rope_ref_grid_w).add(circular_x_enabled).add(circular_y_enabled); + auto pos_img = make_input(pe_cache_.get_or_create(pos_img_key, [&]() { + auto pos_img_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), + static_cast(Ws), + static_cast(config.hidden_size / config.num_groups), + 10000.f, + 16.f, + static_cast(config.rope_ref_grid_h), + static_cast(config.rope_ref_grid_w))); + return sd::Tensor({2, 2, config.hidden_size / config.num_groups / 2, Hs * Ws}, std::move(pos_img_vec)); + })); int64_t Ltxt = std::min(context->ne[1], config.txt_max_length); - pos_txt_vec = make_rope_1d(static_cast(Ltxt), - static_cast(config.hidden_size / config.num_groups), - config.text_rope_theta); - auto pos_txt = ggml_new_tensor_4d(compute_ctx, - GGML_TYPE_F32, - 2, - 2, - config.hidden_size / config.num_groups / 2, - Ltxt); - set_backend_tensor_data(pos_txt, pos_txt_vec.data()); - - pixel_pos_vec = make_pixel_abs_pos(static_cast(Hp), - static_cast(Wp), - static_cast(config.pixel_hidden_size)); - auto pixel_pos = ggml_new_tensor_3d(compute_ctx, - GGML_TYPE_F32, - config.pixel_hidden_size, - Wp * Hp, - 1); - set_backend_tensor_data(pixel_pos, pixel_pos_vec.data()); - - pixel_pos_comp_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), - static_cast(Ws), - static_cast(config.pixel_attn_hidden_size / config.pixel_num_groups), - 10000.f, - 16.f, - static_cast(config.rope_ref_grid_h), - static_cast(config.rope_ref_grid_w))); - auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx, - GGML_TYPE_F32, - 2, - 2, - config.pixel_attn_hidden_size / config.pixel_num_groups / 2, - Hs * Ws); - set_backend_tensor_data(pixel_pos_comp, pixel_pos_comp_vec.data()); + sd::PECacheKey pos_txt_key("pid.pos_txt"); + pos_txt_key.add(Ltxt).add(config.hidden_size).add(config.num_groups).add(config.text_rope_theta); + auto pos_txt = make_input(pe_cache_.get_or_create(pos_txt_key, [&]() { + auto pos_txt_vec = make_rope_1d(static_cast(Ltxt), + static_cast(config.hidden_size / config.num_groups), + config.text_rope_theta); + return sd::Tensor({2, 2, config.hidden_size / config.num_groups / 2, Ltxt}, std::move(pos_txt_vec)); + })); + + sd::PECacheKey pixel_pos_key("pid.pixel_pos"); + pixel_pos_key.add(Hp).add(Wp).add(config.pixel_hidden_size); + auto pixel_pos = make_input(pe_cache_.get_or_create(pixel_pos_key, [&]() { + auto pixel_pos_vec = make_pixel_abs_pos(static_cast(Hp), + static_cast(Wp), + static_cast(config.pixel_hidden_size)); + return sd::Tensor({config.pixel_hidden_size, Wp * Hp, 1}, std::move(pixel_pos_vec)); + })); + + sd::PECacheKey pixel_pos_comp_key("pid.pixel_pos_comp"); + pixel_pos_comp_key.add(Hs).add(Ws).add(config.pixel_attn_hidden_size).add(config.pixel_num_groups); + pixel_pos_comp_key.add(config.rope_ref_grid_h).add(config.rope_ref_grid_w).add(circular_x_enabled); + pixel_pos_comp_key.add(circular_y_enabled); + auto pixel_pos_comp = make_input(pe_cache_.get_or_create(pixel_pos_comp_key, [&]() { + auto pixel_pos_comp_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), + static_cast(Ws), + static_cast(config.pixel_attn_hidden_size / config.pixel_num_groups), + 10000.f, + 16.f, + static_cast(config.rope_ref_grid_h), + static_cast(config.rope_ref_grid_w))); + return sd::Tensor({2, 2, config.pixel_attn_hidden_size / config.pixel_num_groups / 2, Hs * Ws}, std::move(pixel_pos_comp_vec)); + })); auto runner_ctx = get_context(); auto out = model.forward(&runner_ctx, diff --git a/src/model/diffusion/pixart.hpp b/src/model/diffusion/pixart.hpp index cc0629192..45f9e4317 100644 --- a/src/model/diffusion/pixart.hpp +++ b/src/model/diffusion/pixart.hpp @@ -5,6 +5,7 @@ #include #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_runner.h" @@ -224,9 +225,9 @@ namespace PixArt { PixArtModel(const PixArtConfig& config) : config(config) { blocks["x_embedder.proj"] = std::make_shared(config.in_channels, - config.hidden_size, - std::pair{static_cast(config.patch_size), static_cast(config.patch_size)}, - std::pair{static_cast(config.patch_size), static_cast(config.patch_size)}); + config.hidden_size, + std::pair{static_cast(config.patch_size), static_cast(config.patch_size)}, + std::pair{static_cast(config.patch_size), static_cast(config.patch_size)}); blocks["t_embedder"] = std::make_shared(ADALN_EMBED_DIM, config.hidden_size); blocks["t_block.1"] = std::make_shared(config.hidden_size, 6 * config.hidden_size); blocks["y_embedder.y_proj.fc1"] = std::make_shared(config.caption_channels, config.hidden_size); @@ -237,7 +238,7 @@ namespace PixArt { } blocks["final_layer.norm_final"] = std::make_shared(config.hidden_size, 1e-6f, false); blocks["final_layer.linear"] = std::make_shared(config.hidden_size, - config.patch_size * config.patch_size * config.out_channels); + config.patch_size * config.patch_size * config.out_channels); } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -295,9 +296,9 @@ namespace PixArt { }; struct PixArtRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; PixArtConfig config; PixArtModel model; - std::vector pos_vec; PixArtRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -357,10 +358,13 @@ namespace PixArt { int64_t wp = W / config.patch_size; int64_t hp = H / config.patch_size; - pos_vec = gen_2d_sincos_pos_embed(config.hidden_size, hp, wp, - config.pos_embed_base_size, config.interpolation_scale); - auto pos = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.hidden_size, wp * hp, 1); - set_backend_tensor_data(pos, pos_vec.data()); + sd::PECacheKey pos_key("pixart.pos"); + pos_key.add(hp).add(wp).add(config.hidden_size).add(config.pos_embed_base_size).add(config.interpolation_scale); + auto pos = make_input(pe_cache_.get_or_create(pos_key, [&]() { + auto pos_vec = gen_2d_sincos_pos_embed(config.hidden_size, hp, wp, + config.pos_embed_base_size, config.interpolation_scale); + return sd::Tensor({config.hidden_size, wp * hp, 1}, std::move(pos_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pos, context_mask); diff --git a/src/model/diffusion/qwen_image.hpp b/src/model/diffusion/qwen_image.hpp index 82d173005..4c473da35 100644 --- a/src/model/diffusion/qwen_image.hpp +++ b/src/model/diffusion/qwen_image.hpp @@ -4,6 +4,7 @@ #include #include "core/ggml_extend_backend.h" #include "core/ggml_tensor_utils.h" +#include "core/pe_cache.h" #include "core/util.h" #include "model/common/block.hpp" @@ -561,10 +562,11 @@ namespace Qwen { }; struct QwenImageRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: QwenImageConfig config; QwenImageModel qwen_image; - std::vector pe_vec; std::vector modulate_index_vec; std::vector additional_t_cond_vec; SDVersion version; @@ -635,23 +637,24 @@ namespace Qwen { ref_index_mode = Rope::RefIndexMode::DECREASE; } - pe_vec = finish_rope_pe(Rope::gen_qwen_image_pe(time_len, - static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - batch_size, - static_cast(context->ne[1]), - ref_latents, - ref_index_mode, - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_VERBOSE("pos_len %d", pos_len); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - // pe->data = pe_vec.data(); - // print_ggml_tensor(pe, true, "pe"); - // pe->data = nullptr; - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("qwen_image.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(time_len).add(batch_size).add(config.patch_size); + pe_key.add_shapes(ref_latents).add(ref_index_mode); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_qwen_image_pe(time_len, + static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + batch_size, + static_cast(context->ne[1]), + ref_latents, + ref_index_mode, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); ggml_tensor* modulate_index = nullptr; if (config.zero_cond_t) { diff --git a/src/model/diffusion/qwen_image_2_1.hpp b/src/model/diffusion/qwen_image_2_1.hpp index 1d37d6afe..042544acd 100644 --- a/src/model/diffusion/qwen_image_2_1.hpp +++ b/src/model/diffusion/qwen_image_2_1.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_DIFFUSION_QWEN_IMAGE_2_1_H__ #define __SD_MODEL_DIFFUSION_QWEN_IMAGE_2_1_H__ +#include "core/pe_cache.h" #include "model/diffusion/qwen_image.hpp" namespace Qwen { @@ -370,9 +371,9 @@ namespace Qwen { }; struct QwenImage21Runner : public DiffusionModelRunner { + sd::PECache pe_cache_; QwenImage21Config config; QwenImage21Model model; - std::vector pe_data; std::vector> mask_data; ggml_type prefix_cache_type = GGML_TYPE_COUNT; bool prefix_cache_enabled = true; @@ -486,8 +487,8 @@ namespace Qwen { } QwenImage21PrefixCache cache; if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) { - cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id) + - ".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled); + cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id) + + ".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled); cache.prefix_length = layout.prefix_length; if (prefix_cache_type != GGML_TYPE_COUNT) { cache.type = prefix_cache_type; @@ -504,20 +505,7 @@ namespace Qwen { if (prefix_cache_type == GGML_TYPE_COUNT && active_cache.type == GGML_TYPE_F16) { checked_cache.flash_attn_used = &flash_attn_used; } - const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE; - const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0); - Rope::Embedding embedding; - embedding.ids.assign(first_position, layout.positions.end()); - const size_t offset = cached ? static_cast(layout.prefix_length) : 0; - embedding.positions.token_count = embedding.ids.size(); - for (auto region : layout.rope_layout.images) { - if (region.begin >= offset) { - region.begin -= offset; - embedding.positions.images.push_back(region); - } - } - embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, config.axes_dim, embedding.layout, &embedding.frequencies); - pe_data = finish_rope_pe(std::move(embedding)); + const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE; mask_data.clear(); if (!cached) { for (const auto& segment : layout.segments) { @@ -535,9 +523,30 @@ namespace Qwen { } auto build = [&]() { auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, - layout.positions.size() - (cached ? layout.prefix_length : 0)); - set_backend_tensor_data(pe, pe_data.data()); + sd::PECacheKey pe_key("qwen_image_2_1.pe"); + pe_key.add(layout.positions).add(cached).add(layout.prefix_length).add(config.axes_dim).add(config.head_dim); + pe_key.add(circular_x_enabled).add(circular_y_enabled); + for (const auto& region : layout.rope_layout.images) { + pe_key.add(region.begin).add(region.count).add(region.height_period).add(region.width_period); + pe_key.add(region.height_axis).add(region.width_axis); + } + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0); + Rope::Embedding embedding; + embedding.ids.assign(first_position, layout.positions.end()); + const size_t offset = cached ? static_cast(layout.prefix_length) : 0; + embedding.positions.token_count = embedding.ids.size(); + for (auto region : layout.rope_layout.images) { + if (region.begin >= offset) { + region.begin -= offset; + embedding.positions.images.push_back(region); + } + } + embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, config.axes_dim, embedding.layout, &embedding.frequencies); + auto values = finish_rope_pe(std::move(embedding)); + const int64_t tokens = static_cast(layout.positions.size()) - (cached ? layout.prefix_length : 0); + return sd::Tensor({2, 2, config.head_dim / 2, tokens}, std::move(values)); + })); std::vector masks, ref_inputs; for (const auto& mask : mask_data) { masks.push_back(mask.empty() ? nullptr : make_input(mask)); @@ -549,7 +558,7 @@ namespace Qwen { } auto ctx = get_context(graph); auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), cached ? nullptr : make_input(context), - ref_inputs, pe, layout, masks, checked_cache); + ref_inputs, pe, layout, masks, checked_cache); if (!flash_attn_used) { return static_cast(nullptr); } diff --git a/src/model/diffusion/sensenova_u1.h b/src/model/diffusion/sensenova_u1.h index 1702ed854..302df6b26 100644 --- a/src/model/diffusion/sensenova_u1.h +++ b/src/model/diffusion/sensenova_u1.h @@ -563,7 +563,7 @@ namespace SenseNovaU1 { blocks["language_model.model"] = std::make_shared(config); blocks["fm_modules.vision_model_mot_gen.embeddings"] = std::make_shared(config); blocks["fm_modules.timestep_embedder"] = std::make_shared(config.hidden_size, - config.timestep_embedding_size); + config.timestep_embedding_size); if (config.add_noise_scale_embedding) { blocks["fm_modules.noise_scale_embedder"] = std::make_shared(config.hidden_size, config.timestep_embedding_size); diff --git a/src/model/diffusion/wan.hpp b/src/model/diffusion/wan.hpp index 4b75203d3..3c76c4ea1 100644 --- a/src/model/diffusion/wan.hpp +++ b/src/model/diffusion/wan.hpp @@ -8,6 +8,7 @@ #include #include "core/ggml_extend_backend.h" #include "core/ggml_tensor_utils.h" +#include "core/pe_cache.h" #include "model/common/block.hpp" #include "model/common/rope.hpp" @@ -885,11 +886,12 @@ namespace WAN { }; struct WanRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: std::string desc = "wan"; WanConfig config; Wan wan; - std::vector pe_vec; SDVersion version; WanRunner(ggml_backend_t backend, @@ -1000,37 +1002,36 @@ namespace WAN { ggml_tensor* audio_embed = make_optional_input(audio_embed_tensor); ggml_tensor* ref_latent = make_optional_input(ref_latent_tensor); - pe_vec = Rope::gen_wan_pe(static_cast(x->ne[2]), - static_cast(x->ne[1]), - static_cast(x->ne[0]), - std::get<0>(config.patch_size), - std::get<1>(config.patch_size), - std::get<2>(config.patch_size), - 1, - config.theta, - config.axes_dim); - if (ref_latent != nullptr) { - // Match S2V's reference-frame temporal offset. - int t_start = std::max(30, static_cast(x->ne[2]) + 9); - auto ref_pe = Rope::gen_wan_pe(static_cast(ref_latent->ne[2]), - static_cast(ref_latent->ne[1]), - static_cast(ref_latent->ne[0]), + sd::PECacheKey pe_key("wan.pe"); + pe_key.add_shape(x).add_shape(ref_latent).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(std::get<0>(config.patch_size)).add(std::get<1>(config.patch_size)).add(std::get<2>(config.patch_size)); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = Rope::gen_wan_pe(static_cast(x->ne[2]), + static_cast(x->ne[1]), + static_cast(x->ne[0]), std::get<0>(config.patch_size), std::get<1>(config.patch_size), std::get<2>(config.patch_size), 1, config.theta, - config.axes_dim, - t_start); - pe_vec.insert(pe_vec.end(), ref_pe.begin(), ref_pe.end()); - } - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_VERBOSE("pos_len %d", pos_len); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - // pe->data = pe_vec.data(); - // print_ggml_tensor(pe); - // pe->data = nullptr; - set_backend_tensor_data(pe, pe_vec.data()); + config.axes_dim); + if (ref_latent != nullptr) { + int t_start = std::max(30, static_cast(x->ne[2]) + 9); + auto ref_pe = Rope::gen_wan_pe(static_cast(ref_latent->ne[2]), + static_cast(ref_latent->ne[1]), + static_cast(ref_latent->ne[0]), + std::get<0>(config.patch_size), + std::get<1>(config.patch_size), + std::get<2>(config.patch_size), + 1, + config.theta, + config.axes_dim, + t_start); + pe_vec.insert(pe_vec.end(), ref_pe.begin(), ref_pe.end()); + } + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); if (c_concat != nullptr) { x = ggml_concat(compute_ctx, x, c_concat, 3); diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp index ed2644047..229c6aaee 100644 --- a/src/model/diffusion/z_image.hpp +++ b/src/model/diffusion/z_image.hpp @@ -3,6 +3,7 @@ #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_extend_backend.h" @@ -23,9 +24,10 @@ #endif namespace ZImage { - constexpr int Z_IMAGE_GRAPH_SIZE = 20480; - constexpr int ADALN_EMBED_DIM = 256; - constexpr int SEQ_MULTI_OF = 32; + constexpr int Z_IMAGE_GRAPH_SIZE = 24576; + constexpr int ADALN_EMBED_DIM = 256; + constexpr int SEQ_MULTI_OF = 32; + constexpr int CONTROL_IN_CHANNELS = 33; struct ZImageConfig { int patch_size = 2; @@ -391,6 +393,135 @@ namespace ZImage { } }; + struct ZImageControlTransformerBlock : public GGMLBlock { + protected: + int block_id; + bool has_before_proj; + int64_t hidden_size; + int64_t head_dim; + int64_t num_heads; + int64_t num_kv_heads; + bool qk_norm; + + public: + ZImageControlTransformerBlock(int block_id, + int64_t hidden_size, + int64_t head_dim, + int64_t num_heads, + int64_t num_kv_heads, + int64_t multiple_of, + float ffn_dim_multiplier, + float norm_eps, + bool qk_norm, + bool with_before_proj = false, + bool norm_elementwise_affine = true, + float qk_norm_eps = 1e-6f) + : block_id(block_id), + has_before_proj(with_before_proj), + hidden_size(hidden_size), + head_dim(head_dim), + num_heads(num_heads), + num_kv_heads(num_kv_heads), + qk_norm(qk_norm) { + blocks["attention.to_q"] = std::make_shared(hidden_size, num_heads * head_dim, false); + blocks["attention.to_k"] = std::make_shared(hidden_size, num_kv_heads * head_dim, false); + blocks["attention.to_v"] = std::make_shared(hidden_size, num_kv_heads * head_dim, false); + blocks["attention.to_out.0"] = std::make_shared(num_heads * head_dim, hidden_size, false); + if (qk_norm) { + blocks["attention.norm_q"] = std::make_shared(head_dim, qk_norm_eps, norm_elementwise_affine); + blocks["attention.norm_k"] = std::make_shared(head_dim, qk_norm_eps, norm_elementwise_affine); + } + blocks["attention_norm1"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["ffn_norm1"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["attention_norm2"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["ffn_norm2"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["feed_forward"] = std::make_shared(hidden_size, hidden_size, multiple_of, ffn_dim_multiplier); + blocks["adaLN_modulation.0"] = std::make_shared(MIN(hidden_size, ADALN_EMBED_DIM), 4 * hidden_size); + if (with_before_proj) { + blocks["before_proj"] = std::make_shared(hidden_size, hidden_size); + } + blocks["after_proj"] = std::make_shared(hidden_size, hidden_size); + } + + ggml_tensor* forward_attention(GGMLRunnerContext* ctx, + ggml_tensor* c, + ggml_tensor* pe) { + auto to_q = std::dynamic_pointer_cast(blocks["attention.to_q"]); + auto to_k = std::dynamic_pointer_cast(blocks["attention.to_k"]); + auto to_v = std::dynamic_pointer_cast(blocks["attention.to_v"]); + auto out_proj = std::dynamic_pointer_cast(blocks["attention.to_out.0"]); + + if (out_proj->has_quantized_weight()) { + out_proj->set_scale(1.f / 16.f); + } + + int64_t n_token = c->ne[1]; + int64_t N = c->ne[2]; + + if (sd_backend_is(ctx->backend, "ROCm")) { + out_proj->set_force_prec_f32(true); + to_q->set_force_prec_f32(true); + to_k->set_force_prec_f32(true); + to_v->set_force_prec_f32(true); + } + + auto q = ggml_reshape_4d(ctx->ggml_ctx, to_q->forward(ctx, c), head_dim, num_heads, n_token, N); + auto k = ggml_reshape_4d(ctx->ggml_ctx, to_k->forward(ctx, c), head_dim, num_kv_heads, n_token, N); + auto v = ggml_reshape_4d(ctx->ggml_ctx, to_v->forward(ctx, c), head_dim, num_kv_heads, n_token, N); + + if (qk_norm) { + q = std::dynamic_pointer_cast(blocks["attention.norm_q"])->forward(ctx, q); + k = std::dynamic_pointer_cast(blocks["attention.norm_k"])->forward(ctx, k); + } + + auto out = Rope::attention(ctx, q, k, v, pe, nullptr, 1.f / 128.f); + return out_proj->forward(ctx, out); + } + + std::pair forward(GGMLRunnerContext* ctx, + ggml_tensor* c, + ggml_tensor* pe, + ggml_tensor* skip, + ggml_tensor* t_emb) { + auto attention_norm1 = std::dynamic_pointer_cast(blocks["attention_norm1"]); + auto ffn_norm1 = std::dynamic_pointer_cast(blocks["ffn_norm1"]); + auto attention_norm2 = std::dynamic_pointer_cast(blocks["attention_norm2"]); + auto ffn_norm2 = std::dynamic_pointer_cast(blocks["ffn_norm2"]); + auto feed_forward = std::dynamic_pointer_cast(blocks["feed_forward"]); + auto adaLN_modulation_0 = std::dynamic_pointer_cast(blocks["adaLN_modulation.0"]); + auto after_proj = std::dynamic_pointer_cast(blocks["after_proj"]); + + if (block_id == 0 && has_before_proj) { + GGML_ASSERT(skip != nullptr); + auto before_proj = std::dynamic_pointer_cast(blocks["before_proj"]); + c = ggml_add(ctx->ggml_ctx, before_proj->forward(ctx, c), skip); + } + + auto m = adaLN_modulation_0->forward(ctx, t_emb); + auto mods = ggml_ext_chunk(ctx->ggml_ctx, m, 4, 0); + auto scale_msa = mods[0]; + auto gate_msa = mods[1]; + auto scale_mlp = mods[2]; + auto gate_mlp = mods[3]; + + auto residual = c; + c = modulate(ctx->ggml_ctx, attention_norm1->forward(ctx, c), scale_msa); + c = forward_attention(ctx, c, pe); + c = attention_norm2->forward(ctx, c); + c = ggml_mul(ctx->ggml_ctx, c, ggml_tanh(ctx->ggml_ctx, gate_msa)); + c = ggml_add(ctx->ggml_ctx, c, residual); + + residual = c; + c = modulate(ctx->ggml_ctx, ffn_norm1->forward(ctx, c), scale_mlp); + c = feed_forward->forward(ctx, c); + c = ffn_norm2->forward(ctx, c); + c = ggml_mul(ctx->ggml_ctx, c, ggml_tanh(ctx->ggml_ctx, gate_mlp)); + c = ggml_add(ctx->ggml_ctx, c, residual); + + return {after_proj->forward(ctx, c), c}; + } + }; + struct FinalLayer : public GGMLBlock { public: FinalLayer(int64_t hidden_size, @@ -423,15 +554,24 @@ namespace ZImage { class ZImageModel : public GGMLBlock { protected: ZImageConfig config; + GGMLBlockMap control_blocks; + int control_num_layers = 0; + int control_num_refiner_layers = 0; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size); params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size); + + for (auto& [name, block] : control_blocks) { + block->init(ctx, tensor_storage_map, name); + } } public: + bool control_enabled = false; + ZImageModel() = default; - ZImageModel(ZImageConfig config) + ZImageModel(ZImageConfig config, const String2TensorStorage& tensor_storage_map = {}) : config(config) { blocks["x_embedder"] = std::make_shared(config.patch_size * config.patch_size * config.in_channels, config.hidden_size); blocks["t_embedder"] = std::make_shared(MIN(config.hidden_size, 1024), 256, 256); @@ -484,13 +624,154 @@ namespace ZImage { } blocks["final_layer"] = std::make_shared(config.hidden_size, config.patch_size, config.out_channels); + + if (tensor_storage_map.find("control_all_x_embedder.2-1.weight") != tensor_storage_map.end()) { + control_enabled = true; + + for (const auto& pair : tensor_storage_map) { + if (starts_with(pair.first, "control_layers.")) { + auto items = split_string(pair.first, '.'); + if (items.size() > 1) { + control_num_layers = std::max(control_num_layers, atoi(items[1].c_str()) + 1); + } + } else if (starts_with(pair.first, "control_noise_refiner.")) { + auto items = split_string(pair.first, '.'); + if (items.size() > 1) { + control_num_refiner_layers = std::max(control_num_refiner_layers, atoi(items[1].c_str()) + 1); + } + } + } + + GGML_ASSERT(control_num_layers > 0); + GGML_ASSERT(control_num_refiner_layers == config.num_refiner_layers); + + control_blocks["control_all_x_embedder.2-1"] = std::make_shared(config.patch_size * config.patch_size * CONTROL_IN_CHANNELS, config.hidden_size); + + for (int i = 0; i < control_num_refiner_layers; i++) { + std::string name = "control_noise_refiner." + std::to_string(i); + auto block = std::make_shared(i, + config.hidden_size, + config.head_dim, + config.num_heads, + config.num_kv_heads, + config.multiple_of, + config.ffn_dim_multiplier, + config.norm_eps, + config.qk_norm, + tensor_storage_map.find(name + ".before_proj.weight") != tensor_storage_map.end()); + + control_blocks[name] = block; + } + + for (int i = 0; i < control_num_layers; i++) { + std::string name = "control_layers." + std::to_string(i); + auto block = std::make_shared(i, + config.hidden_size, + config.head_dim, + config.num_heads, + config.num_kv_heads, + config.multiple_of, + config.ffn_dim_multiplier, + config.norm_eps, + config.qk_norm, + tensor_storage_map.find(name + ".before_proj.weight") != tensor_storage_map.end()); + + control_blocks[name] = block; + } + } + } + + void get_control_param_tensors(std::map& tensors) { + for (auto& [name, block] : control_blocks) { + block->get_param_tensors(tensors, name); + } + } + + ggml_tensor* forward_control_refiner(GGMLRunnerContext* ctx, + ggml_tensor* control, + ggml_tensor* img, + ggml_tensor* img_pe, + ggml_tensor* t_emb, + float strength, + std::vector& ref_hints) { + auto all_x_embedder = std::dynamic_pointer_cast(control_blocks["control_all_x_embedder.2-1"]); + + auto c = all_x_embedder->forward(ctx, control); + sd::ggml_graph_cut::mark_graph_cut(c, "z_image.control.prelude", "control"); + + int64_t N = c->ne[2]; + int64_t n_token = c->ne[1]; + + int64_t n_pad_token = Rope::bound_mod(static_cast(n_token), SEQ_MULTI_OF); + if (n_pad_token > 0) { + auto img_pad_token = params["x_pad_token"]; + auto ctrl_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, img_pad_token, img_pad_token->ne[0], n_pad_token, N, 1); + c = ggml_concat(ctx->ggml_ctx, c, ctrl_pad_tokens, 1); + } + + for (int i = 0; i < control_num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(control_blocks["control_noise_refiner." + std::to_string(i)]); + + auto out = block->forward(ctx, c, img_pe, img, t_emb); + c = out.second; + ref_hints.push_back(ggml_ext_scale(ctx->ggml_ctx, out.first, strength)); + sd::ggml_graph_cut::mark_graph_cut(c, "z_image.control.noise_refiner." + std::to_string(i), "control"); + } + + return c; + } + + ggml_tensor* forward_noise_refiner_hinted(GGMLRunnerContext* ctx, + ggml_tensor* img, + ggml_tensor* img_pe, + ggml_tensor* t_emb, + const std::vector& ref_hints) { + GGML_ASSERT(ref_hints.size() == static_cast(config.num_refiner_layers)); + + auto c = img; + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["noise_refiner." + std::to_string(i)]); + + c = block->forward(ctx, c, img_pe, nullptr, t_emb); + c = ggml_add(ctx->ggml_ctx, c, ref_hints[i]); + sd::ggml_graph_cut::mark_graph_cut(c, "z_image.noise_refiner_hinted." + std::to_string(i), "img"); + } + + return c; + } + + std::vector forward_control_main(GGMLRunnerContext* ctx, + ggml_tensor* control_ctx, + ggml_tensor* img_c, + ggml_tensor* txt, + ggml_tensor* pe, + ggml_tensor* t_emb, + float strength) { + auto unified_c = ggml_concat(ctx->ggml_ctx, txt, img_c, 1); + auto ctrl_unified = ggml_concat(ctx->ggml_ctx, txt, control_ctx, 1); + + std::vector main_hints; + main_hints.reserve(control_num_layers); + + for (int i = 0; i < control_num_layers; i++) { + auto block = std::dynamic_pointer_cast(control_blocks["control_layers." + std::to_string(i)]); + + auto out = block->forward(ctx, ctrl_unified, pe, unified_c, t_emb); + ctrl_unified = out.second; + main_hints.push_back(ggml_ext_scale(ctx->ggml_ctx, out.first, strength)); + sd::ggml_graph_cut::mark_graph_cut(ctrl_unified, "z_image.control.layers." + std::to_string(i), "control"); + } + + return main_hints; } ggml_tensor* forward_core(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, - ggml_tensor* pe) { + ggml_tensor* pe, + ggml_tensor* control = nullptr, + float strength = 0.f) { auto x_embedder = std::dynamic_pointer_cast(blocks["x_embedder"]); auto t_embedder = std::dynamic_pointer_cast(blocks["t_embedder"]); auto cap_embedder_0 = std::dynamic_pointer_cast(blocks["cap_embedder.0"]); @@ -530,6 +811,14 @@ namespace ZImage { auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]); auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]); + std::vector ref_hints; + ggml_tensor* control_ctx = nullptr; + if (control != nullptr) { + control_ctx = forward_control_refiner(ctx, control, img, img_pe, t_emb, strength, ref_hints); + } + + ggml_tensor* img_pre_refiner = img; + for (int i = 0; i < config.num_refiner_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)]); @@ -544,13 +833,27 @@ namespace ZImage { sd::ggml_graph_cut::mark_graph_cut(img, "z_image.noise_refiner." + std::to_string(i), "img"); } + ggml_tensor* img_c = nullptr; + if (control_ctx != nullptr) { + img_c = forward_noise_refiner_hinted(ctx, img_pre_refiner, img_pe, t_emb, ref_hints); + } + auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, hidden_size] sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.prelude", "txt_img"); + std::vector main_hints; + if (control_ctx != nullptr) { + main_hints = forward_control_main(ctx, control_ctx, img_c, txt, pe, t_emb, strength); + } + GGML_ASSERT(main_hints.empty() || main_hints.size() == static_cast(control_num_layers)); + for (int i = 0; i < config.num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb); + if (i % 2 == 0 && (i / 2) < static_cast(main_hints.size())) { + txt_img = ggml_add(ctx->ggml_ctx, txt_img, main_hints[i / 2]); + } sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.layers." + std::to_string(i), "txt_img"); } @@ -566,7 +869,9 @@ namespace ZImage { ggml_tensor* timestep, ggml_tensor* context, ggml_tensor* pe, - std::vector ref_latents = {}) { + std::vector ref_latents = {}, + ggml_tensor* control_latent = nullptr, + float control_strength = 0.f) { // Forward pass of DiT. // x: [N, C, H, W] // timestep: [N,] @@ -591,7 +896,16 @@ namespace ZImage { } } - auto out = forward_core(ctx, img, timestep, context, pe); + ggml_tensor* control = nullptr; + if (control_latent != nullptr) { + GGML_ASSERT(control_enabled); + GGML_ASSERT(control_latent->ne[2] == CONTROL_IN_CHANNELS); + GGML_ASSERT(control_latent->ne[0] == W && control_latent->ne[1] == H && control_latent->ne[3] == N); + GGML_ASSERT(ref_latents.empty()); + control = DiT::pad_and_patchify(ctx, control_latent, patch_size, patch_size, false); + } + + auto out = forward_core(ctx, img, timestep, context, pe, control, control_strength); out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, n_img_token); // [N, n_img_token, ph*pw*C] out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false); // [N, C, H, W] @@ -603,10 +917,11 @@ namespace ZImage { }; struct ZImageRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: ZImageConfig config; ZImageModel z_image; - std::vector pe_vec; std::vector timestep_vec; SDVersion version; @@ -617,7 +932,7 @@ namespace ZImage { std::shared_ptr weight_manager = nullptr) : DiffusionModelRunner(backend, prefix, weight_manager), config(ZImageConfig::detect_from_weights(tensor_storage_map, prefix)) { - z_image = ZImageModel(config); + z_image = ZImageModel(config, tensor_storage_map); z_image.init(params_ctx, tensor_storage_map, prefix); } @@ -629,11 +944,21 @@ namespace ZImage { z_image.get_param_tensors(tensors, prefix); } + bool control_enabled() const { + return z_image.control_enabled; + } + + void get_control_param_tensors(std::map& tensors) { + z_image.get_control_param_tensors(tensors); + } + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, const sd::Tensor& timesteps_tensor, const sd::Tensor& context_tensor, const std::vector>& ref_latents_tensor = {}, - Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) { + Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED, + const sd::Tensor& control_latent_tensor = {}, + float control_strength = 0.f) { ggml_cgraph* gf = new_graph_custom(Z_IMAGE_GRAPH_SIZE); ggml_tensor* x = make_input(x_tensor); ggml_tensor* timesteps = make_input(timesteps_tensor); @@ -646,23 +971,26 @@ namespace ZImage { ref_latents.push_back(make_input(ref_latent_tensor)); } - pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - SEQ_MULTI_OF, - ref_latents, - ref_index_mode, - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_VERBOSE("pos_len %d", pos_len); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - // pe->data = pe_vec.data(); - // print_ggml_tensor(pe, true, "pe"); - // pe->data = nullptr; - set_backend_tensor_data(pe, pe_vec.data()); + ggml_tensor* control_latent = make_optional_input(control_latent_tensor); + + sd::PECacheKey pe_key("z_image.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shapes(ref_latents); + pe_key.add(ref_index_mode); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + SEQ_MULTI_OF, + ref_latents, + ref_index_mode, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = z_image.forward(&runner_ctx, @@ -670,7 +998,9 @@ namespace ZImage { timesteps, context, pe, - ref_latents); + ref_latents, + control_latent, + control_strength); ggml_build_forward_expand(gf, out); @@ -682,12 +1012,14 @@ namespace ZImage { const sd::Tensor& timesteps, const sd::Tensor& context, const std::vector>& ref_latents = {}, - Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) { + Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED, + const sd::Tensor& control_latent = {}, + float control_strength = 0.f) { // x: [N, in_channels, h, w] // timesteps: [N, ] // context: [N, max_position, hidden_size] auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, ref_latents, ref_index_mode); + return build_graph(x, timesteps, context, ref_latents, ref_index_mode, control_latent, control_strength); }; return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); @@ -698,12 +1030,21 @@ namespace ZImage { GGML_ASSERT(diffusion_params.x != nullptr); GGML_ASSERT(diffusion_params.timesteps != nullptr); static const std::vector> empty_ref_latents; + static const sd::Tensor empty_control_latent; + + const auto* z_image_extra = std::get_if(&diffusion_params.extra); + const sd::Tensor& control_latent = (z_image_extra != nullptr && z_image_extra->control_latent != nullptr) ? *z_image_extra->control_latent + : empty_control_latent; + float control_strength = z_image_extra != nullptr ? z_image_extra->strength : 0.f; + return compute(n_threads, *diffusion_params.x, *diffusion_params.timesteps, tensor_or_empty(diffusion_params.context), diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents, - diffusion_params.ref_image_params.ref_index_mode); + diffusion_params.ref_image_params.ref_index_mode, + control_latent, + control_strength); } void test() { diff --git a/src/model/diffusion/z_image_l2p.hpp b/src/model/diffusion/z_image_l2p.hpp index 0eb0b14d3..f44f2ff9f 100644 --- a/src/model/diffusion/z_image_l2p.hpp +++ b/src/model/diffusion/z_image_l2p.hpp @@ -2,6 +2,7 @@ #define __SD_MODEL_DIFFUSION_Z_IMAGE_L2P_HPP__ #include +#include "core/pe_cache.h" #include "z_image.hpp" @@ -183,9 +184,9 @@ namespace ZImageL2P { }; struct ZImageL2PRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; ZImageL2PConfig config; ZImageL2PModel model; - std::vector pe_vec; ZImageL2PRunner(ggml_backend_t backend, const String2TensorStorage& tensors, @@ -221,19 +222,23 @@ namespace ZImageL2P { auto context = make_input(*inputs.context); GGML_ASSERT(x->ne[3] == 1); GGML_ASSERT(x->ne[0] % config.patch_size == 0 && x->ne[1] % config.patch_size == 0); - pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - ZImage::SEQ_MULTI_OF, - {}, - Rope::RefIndexMode::FIXED, - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("z_image_l2p.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ZImage::SEQ_MULTI_OF, + {}, + Rope::RefIndexMode::FIXED, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto ctx = get_context(); auto out = model.forward(&ctx, x, t, context, pe); ggml_build_forward_expand(gf, out); diff --git a/src/model/integrations/anima.cpp b/src/model/integrations/anima.cpp new file mode 100644 index 000000000..2f9b2babc --- /dev/null +++ b/src/model/integrations/anima.cpp @@ -0,0 +1,43 @@ +#include "model/diffusion/anima.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_anima(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& anima_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_ANIMA; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "cosmos_reference"; + definition.name = "Anima"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_anima; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/boogu_image.cpp b/src/model/integrations/boogu_image.cpp new file mode 100644 index 000000000..3d6e8c06c --- /dev/null +++ b/src/model/integrations/boogu_image.cpp @@ -0,0 +1,46 @@ +#include "conditioning/conditioner.hpp" +#include "model/diffusion/boogu.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_boogu_image(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager); + return true; +} + +const std::vector& boogu_image_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_BOOGU_IMAGE; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "z_image_omni"; + definition.name = "Boogu Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.sampling = {FLOW_PRED, 3.16f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_boogu_image; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/ernie_image.cpp b/src/model/integrations/ernie_image.cpp new file mode 100644 index 000000000..f1d185cdd --- /dev/null +++ b/src/model/integrations/ernie_image.cpp @@ -0,0 +1,44 @@ +#include "model/diffusion/ernie_image.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_ernie_image(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& ernie_image_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_ERNIE_IMAGE; + definition.geometry = {1, 128, 3}; + definition.name = "Ernie Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLOW_PRED, 4.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_ernie_image; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/fibo2.cpp b/src/model/integrations/fibo2.cpp new file mode 100644 index 000000000..d16a0cf85 --- /dev/null +++ b/src/model/integrations/fibo2.cpp @@ -0,0 +1,39 @@ +#include "model/diffusion/fibo2.hpp" +#include "conditioning/conditioner.hpp" +#include "conditioning/fibo2.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_fibo2(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, weight_manager, tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, "model.diffusion_model", weight_manager); + return true; +} + +const std::vector& fibo2_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_FIBO2; + definition.geometry = {1, 128, 3}; + definition.ref_image_preset = "fibo2"; + definition.name = "Fibo-2.0"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_fibo2; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/flux.cpp b/src/model/integrations/flux.cpp new file mode 100644 index 000000000..4f14237aa --- /dev/null +++ b/src/model/integrations/flux.cpp @@ -0,0 +1,159 @@ +#include "model/diffusion/flux.hpp" +#include "conditioning/conditioner.hpp" +#include "model/integrations/flux.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_flux(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + bool is_chroma = false; + for (auto pair : tensor_storage_map) { + if (pair.first.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) { + is_chroma = true; + break; + } + } + if (is_chroma) { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + false, + 1, + false, + weight_manager, + sd_ctx_params->model_args); + } else if (version == VERSION_OVIS_IMAGE) { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + } else { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager, + tokenizers); + } + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager, + sd_ctx_params->model_args); + return true; +} + +const std::vector& flux_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_FLUX; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "flux_kontext"; + definition.name = "Flux"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_flux; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_FLUX_FILL; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "flux_kontext"; + definition.name = "Flux Fill"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.traits.is_inpaint = true; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_flux; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_FLUX_CONTROLS; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "flux_kontext"; + definition.name = "Flux Control"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.traits.is_control = true; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_flux; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_FLEX_2; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "flux_kontext"; + definition.name = "Flex.2"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.traits.is_inpaint = true; + definition.traits.is_control = true; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_flux; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_CHROMA_RADIANCE; + definition.geometry = {1, 3, 3}; + definition.ref_image_preset = "flux_kontext"; + definition.name = "Chroma Radiance"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.traits.fake_vae = true; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_flux; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_OVIS_IMAGE; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "flux_kontext"; + definition.name = "Ovis Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_flux; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} + +float flux_guidance_flow_shift(const String2TensorStorage& weights) { + for (const auto& [name, tensor] : weights) { + if (starts_with(name, "model.diffusion_model.guidance_in.in_layer.weight")) + return 1.15f; + } + return 1.f; +} diff --git a/src/model/integrations/flux.hpp b/src/model/integrations/flux.hpp new file mode 100644 index 000000000..3b11888b0 --- /dev/null +++ b/src/model/integrations/flux.hpp @@ -0,0 +1,5 @@ +#ifndef __SD_MODEL_INTEGRATIONS_FLUX_HPP__ +#define __SD_MODEL_INTEGRATIONS_FLUX_HPP__ +#include "model.h" +float flux_guidance_flow_shift(const String2TensorStorage& weights); +#endif diff --git a/src/model/integrations/flux2.cpp b/src/model/integrations/flux2.cpp new file mode 100644 index 000000000..d895df2e3 --- /dev/null +++ b/src/model/integrations/flux2.cpp @@ -0,0 +1,79 @@ +#include "conditioning/conditioner.hpp" +#include "model/diffusion/flux.hpp" +#include "model/integrations/flux.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_flux2(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + bool is_chroma = false; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager, + sd_ctx_params->model_args); + return true; +} + +const std::vector& flux2_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_FLUX2; + definition.geometry = {1, 128, 3}; + definition.ref_image_preset = "flux2"; + definition.name = "Flux.2"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX2_SCHEDULER; + definition.build = build_flux2; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_FLUX2_KLEIN; + definition.geometry = {1, 128, 3}; + definition.ref_image_preset = "flux2"; + definition.name = "Flux.2 klein"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX2_SCHEDULER; + definition.build = build_flux2; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SEFI_IMAGE; + definition.geometry = {1, 144, 3}; + definition.ref_image_preset = "flux2"; + definition.name = "SeFi-Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {SEFI_FLOW_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_flux2; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/flux3.cpp b/src/model/integrations/flux3.cpp new file mode 100644 index 000000000..9ee8d3db4 --- /dev/null +++ b/src/model/integrations/flux3.cpp @@ -0,0 +1,38 @@ +#include "model/diffusion/flux3.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_flux3(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, version, "", false, weight_manager, tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, "model.diffusion_model", weight_manager); + return true; +} + +const std::vector& flux3_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_FLUX3; + definition.geometry = {1, 128, 3}; + definition.ref_image_preset = "flux2"; + definition.name = "Flux.3 Image Klein"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLUX3_FLOW_PRED, 6.93f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_flux3; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/hidream_o1.cpp b/src/model/integrations/hidream_o1.cpp new file mode 100644 index 000000000..66f6a9fc3 --- /dev/null +++ b/src/model/integrations/hidream_o1.cpp @@ -0,0 +1,41 @@ +#include "model/diffusion/hidream_o1.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_hidream_o1(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model", + weight_manager); + return true; +} + +const std::vector& hidream_o1_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_HIDREAM_O1; + definition.geometry = {1, 3, 3}; + definition.name = "HiDream O1"; + definition.traits.is_dit = true; + definition.traits.fake_vae = true; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_hidream_o1; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/hunyuan_video.cpp b/src/model/integrations/hunyuan_video.cpp new file mode 100644 index 000000000..85bb741f8 --- /dev/null +++ b/src/model/integrations/hunyuan_video.cpp @@ -0,0 +1,47 @@ +#include "conditioning/conditioner.hpp" +#include "model/diffusion/hunyuan.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_hunyuan_video(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager); + return true; +} + +const std::vector& hunyuan_video_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_HUNYUAN_VIDEO; + definition.geometry = {1, 32, 3}; + definition.name = "Hunyuan Video"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::HUNYUAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 7.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_hunyuan_video; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/ideogram4.cpp b/src/model/integrations/ideogram4.cpp new file mode 100644 index 000000000..77baef272 --- /dev/null +++ b/src/model/integrations/ideogram4.cpp @@ -0,0 +1,44 @@ +#include "model/diffusion/ideogram4.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_ideogram4(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& ideogram4_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_IDEOGRAM4; + definition.geometry = {1, 128, 3}; + definition.name = "Ideogram 4"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = LOGIT_NORMAL_SCHEDULER; + definition.build = build_ideogram4; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/iris.cpp b/src/model/integrations/iris.cpp new file mode 100644 index 000000000..a75c88ca4 --- /dev/null +++ b/src/model/integrations/iris.cpp @@ -0,0 +1,37 @@ +#include "model/diffusion/iris.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_iris(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, version, "", false, weight_manager, tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, "model.diffusion_model", weight_manager); + return true; +} + +const std::vector& iris_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_IRIS; + definition.geometry = {16, 3, 3}; + definition.name = "Iris-3B"; + definition.traits.is_dit = true; + definition.traits.fake_vae = true; + definition.sampling = {FLOW_PRED, 4.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_iris; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/krea2.cpp b/src/model/integrations/krea2.cpp new file mode 100644 index 000000000..367044537 --- /dev/null +++ b/src/model/integrations/krea2.cpp @@ -0,0 +1,46 @@ +#include "model/diffusion/krea2.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_krea2(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& krea2_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_KREA2; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "krea2_ostris_edit"; + definition.name = "Krea2"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLUX_FLOW_PRED, 1.15f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_krea2; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/lens.cpp b/src/model/integrations/lens.cpp new file mode 100644 index 000000000..da66d2cd8 --- /dev/null +++ b/src/model/integrations/lens.cpp @@ -0,0 +1,44 @@ +#include "model/diffusion/lens.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_lens(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& lens_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_LENS; + definition.geometry = {1, 128, 3}; + definition.name = "Lens"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLUX_FLOW_PRED, 1.83f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_lens; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/lingbot_video.cpp b/src/model/integrations/lingbot_video.cpp new file mode 100644 index 000000000..d4cc6cc55 --- /dev/null +++ b/src/model/integrations/lingbot_video.cpp @@ -0,0 +1,54 @@ +#include "model/diffusion/lingbot_video.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_lingbot_video(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + bool enable_vision = false; + for (const auto& [name, _] : tensor_storage_map) { + if (starts_with(name, "text_encoders.llm.visual.")) { + enable_vision = true; + break; + } + } + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + enable_vision, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager, + sd_ctx_params->model_args); + return true; +} + +const std::vector& lingbot_video_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_LINGBOT_VIDEO; + definition.geometry = {2, 16, 3}; + definition.name = "LingBot Video"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_lingbot_video; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/llada_image.cpp b/src/model/integrations/llada_image.cpp new file mode 100644 index 000000000..58c47e8f8 --- /dev/null +++ b/src/model/integrations/llada_image.cpp @@ -0,0 +1,50 @@ +#include "model/diffusion/llada_image.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_llada_image(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + "text_encoders.llm", + "queryformer", + "text_projection", + "sigvq", + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& llada_image_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_LLADA_IMAGE; + definition.geometry = {1, 128, 3}; + definition.ref_image_preset = "llada_image"; + definition.name = "LLaDA-Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = LLADA_IMAGE_SCHEDULER; + definition.build = build_llada_image; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} + +// The attention projections keep their diffusers names (JointAttention's split_qkv mode), so +// only the patch-size-keyed dicts need flattening. Latents arrive already patchified from the +// Flux2 VAE, so the only patch key is 1-1. diff --git a/src/model/integrations/longcat.cpp b/src/model/integrations/longcat.cpp new file mode 100644 index 000000000..d371710c4 --- /dev/null +++ b/src/model/integrations/longcat.cpp @@ -0,0 +1,47 @@ +#include "conditioning/conditioner.hpp" +#include "model/diffusion/flux.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_longcat(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager, + sd_ctx_params->model_args); + return true; +} + +const std::vector& longcat_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_LONGCAT; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "longcat"; + definition.name = "Longcat-Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.sampling = {FLUX_FLOW_PRED, 3.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_longcat; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/ltxav.cpp b/src/model/integrations/ltxav.cpp new file mode 100644 index 000000000..999523924 --- /dev/null +++ b/src/model/integrations/ltxav.cpp @@ -0,0 +1,45 @@ +#include "conditioning/conditioner.hpp" +#include "model/diffusion/ltxv.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_ltxav(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + "text_encoders.llm", + "text_embedding_projection", + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& ltxav_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_LTXAV; + definition.geometry = {1, 128, 3}; + definition.name = "LTXAV"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::LTXAV; + definition.traits.video_tae = true; + definition.sampling = {FLUX_FLOW_PRED, 2.37f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = LTX2_SCHEDULER; + definition.build = build_ltxav; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/mage_flow.cpp b/src/model/integrations/mage_flow.cpp new file mode 100644 index 000000000..000fa1571 --- /dev/null +++ b/src/model/integrations/mage_flow.cpp @@ -0,0 +1,45 @@ +#include "model/diffusion/mage_flow.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_mage_flow(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& mage_flow_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_MAGE_FLOW; + definition.geometry = {1, 128, 3}; + definition.ref_image_preset = "mage_flow"; + definition.name = "Mage Flow"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::MAGE; + definition.sampling = {FLOW_PRED, 6.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_mage_flow; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/ming_image.cpp b/src/model/integrations/ming_image.cpp new file mode 100644 index 000000000..8587df8c8 --- /dev/null +++ b/src/model/integrations/ming_image.cpp @@ -0,0 +1,38 @@ +#include "model/diffusion/ming_image.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_ming_image(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, weight_manager, tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, "model.diffusion_model", weight_manager); + return true; +} + +const std::vector& ming_image_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_MING_IMAGE; + definition.geometry = {2, 16, 4}; + definition.name = "Ming-Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {MING_IMAGE_FLOW_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_ming_image; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/minimax_h3.cpp b/src/model/integrations/minimax_h3.cpp new file mode 100644 index 000000000..2569b626e --- /dev/null +++ b/src/model/integrations/minimax_h3.cpp @@ -0,0 +1,46 @@ +#include "model/diffusion/minimax_h3.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_minimax_h3(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& minimax_h3_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_MINIMAX_H3; + definition.geometry = {2, 24, 3}; + definition.name = "MiniMax-H3"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::MINIMAX_H3; + definition.traits.video_tae = true; + definition.sampling = {H3_AV_FLOW_PRED, 12.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_minimax_h3; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/minit2i.cpp b/src/model/integrations/minit2i.cpp new file mode 100644 index 000000000..fb305943d --- /dev/null +++ b/src/model/integrations/minit2i.cpp @@ -0,0 +1,40 @@ +#include "model/diffusion/minit2i.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_minit2i(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model.model.net", + weight_manager); + return true; +} + +const std::vector& minit2i_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_MINIT2I; + definition.geometry = {1, 3, 3}; + definition.name = "MiniT2I"; + definition.traits.is_dit = true; + definition.traits.fake_vae = true; + definition.sampling = {MINIT2I_FLOW_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_minit2i; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/moss_image.cpp b/src/model/integrations/moss_image.cpp new file mode 100644 index 000000000..4a85d169a --- /dev/null +++ b/src/model/integrations/moss_image.cpp @@ -0,0 +1,33 @@ +#include "model_definition.h" + +#include "conditioning/conditioner.hpp" +#include "model/diffusion/moss_image.hpp" +#include "pipeline/model_builders.h" + +static bool build_moss_image(const sd::model_builders::Context& ctx, + const TokenizerConfig& tokenizers, + sd::model_builders::CoreRunners& runners) { + runners.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + ctx.tensor_storage_map, ctx.version, "", false, + ctx.weight_manager, tokenizers); + runners.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + ctx.tensor_storage_map, "model.diffusion_model", ctx.weight_manager); + return true; +} + +const std::vector& moss_image_definitions() { + static const std::vector definitions = [] { + ModelDefinition definition; + definition.version = VERSION_MOSS_IMAGE; + definition.geometry = {1, 128, 3}; + definition.name = "MOSS-Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLOW_PRED, 5.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_moss_image; + return std::vector{std::move(definition)}; + }(); + return definitions; +} diff --git a/src/model/integrations/pid.cpp b/src/model/integrations/pid.cpp new file mode 100644 index 000000000..d307fb6be --- /dev/null +++ b/src/model/integrations/pid.cpp @@ -0,0 +1,44 @@ +#include "model/diffusion/pid.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_pid(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model.net", + weight_manager); + return true; +} + +const std::vector& pid_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_PID; + definition.geometry = {1, 3, 3}; + definition.name = "PiD"; + definition.traits.is_dit = true; + definition.traits.vae_format_override = true; + definition.sampling = {FLOW_PRED, 1.5f}; + definition.sampling.default_sample_method = LCM_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_pid; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/pixart.cpp b/src/model/integrations/pixart.cpp new file mode 100644 index 000000000..d0a4aa709 --- /dev/null +++ b/src/model/integrations/pixart.cpp @@ -0,0 +1,67 @@ +#include "model/diffusion/pixart.hpp" +#include "conditioning/conditioner.hpp" +#include "model/vae/auto_encoder_kl.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static void configure_pixart_vae(const sd::model_builders::Context& ctx, + const sd::model_builders::VAEOptions&, + VAE& vae) { + auto& model = static_cast(vae); + // Alpha-512 and Sigma share tensor layouts; Alpha-512 needs an explicit scale override. + if (ctx.tensor_storage_map.count("model.diffusion_model.csize_embedder.mlp.0.weight") != 0) { + model.scale_factor = 0.18215f; + } + for (const auto& [key, value] : parse_key_value_args(ctx.params.model_args, "model arg")) { + if (key == "pixart_vae_scale_factor") { + float parsed = 0.f; + if (parse_strict_float(value, parsed) && std::isfinite(parsed) && parsed > 0.f) { + model.scale_factor = parsed; + } else { + LOG_WARN("ignoring invalid PixArt model arg '%s=%s'", key.c_str(), value.c_str()); + } + } + } + LOG_VERBOSE("pixart: VAE scale factor = %.5f", model.scale_factor); +} + +static bool build_pixart(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + true, + 0, + false, + weight_manager, + sd_ctx_params->model_args); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager, + sd_ctx_params->model_args); + return true; +} + +const std::vector& pixart_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_PIXART; + definition.geometry = {2, 4, 3}; + definition.name = "PixArt"; + definition.traits.is_dit = true; + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_pixart; + definition.configure_vae = configure_pixart_vae; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/qwen_image.cpp b/src/model/integrations/qwen_image.cpp new file mode 100644 index 000000000..49505bb77 --- /dev/null +++ b/src/model/integrations/qwen_image.cpp @@ -0,0 +1,87 @@ +#include "model/diffusion/qwen_image.hpp" +#include "conditioning/conditioner.hpp" +#include "model/diffusion/qwen_image_2_1.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_qwen_image(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + enable_vision, + weight_manager, + tokenizers); + if (version == VERSION_QWEN_IMAGE_2_1) { + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager, + sd_ctx_params->model_args); + } else { + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager, + sd_ctx_params->model_args); + } + return true; +} + +const std::vector& qwen_image_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_QWEN_IMAGE; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "qwen"; + definition.name = "Qwen Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_qwen_image; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_QWEN_IMAGE_LAYERED; + definition.geometry = {1, 16, 4}; + definition.ref_image_preset = "qwen_layered"; + definition.name = "Qwen Image Layered"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_qwen_image; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_QWEN_IMAGE_2_1; + definition.geometry = {2, 64, 4}; + definition.ref_image_preset = "qwen"; + definition.name = "Qwen Image 2.1"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_qwen_image; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/sd.cpp b/src/model/integrations/sd.cpp new file mode 100644 index 000000000..a36ff8e36 --- /dev/null +++ b/src/model/integrations/sd.cpp @@ -0,0 +1,261 @@ +#include "conditioning/conditioner.hpp" +#include "model/diffusion/unet.hpp" +#include "model/vae/auto_encoder_kl.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static void configure_sdxl_vae(const sd::model_builders::Context& ctx, + const sd::model_builders::VAEOptions& options, + VAE& vae) { + if (strlen(SAFE_STR(ctx.params.vae_path)) == 0 || ctx.params.force_sdxl_vae_conv_scale || options.external_vae_is_invalid) { + float vae_conv_2d_scale = 1.f / 32.f; + LOG_WARN( + "No valid VAE specified with --vae or --force-sdxl-vae-conv-scale flag set, " + "using Conv2D scale %.3f", + vae_conv_2d_scale); + static_cast(vae).set_conv2d_scale(vae_conv_2d_scale); + } +} + +static bool build_sd(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + // SD1.x SD2.x SDXL + std::map embbeding_map; + for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) { + embbeding_map.emplace(SAFE_STR(sd_ctx_params->embeddings[i].name), SAFE_STR(sd_ctx_params->embeddings[i].path)); + } + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + embbeding_map, + version, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager); + return true; +} +static prediction_t detect_sdxl_prediction(const String2TensorStorage& weights) { + if (weights.count("edm_vpred.sigma_max")) + return EDM_V_PRED; + return weights.count("v_pred") ? V_PRED : EPS_PRED; +} + +const std::vector& sd_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_SD1; + definition.geometry = {8, 4, 3}; + definition.name = "SD 1.x"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SD1_INPAINT; + definition.geometry = {8, 4, 3}; + definition.name = "SD 1.x Inpaint"; + definition.traits.is_inpaint = true; + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SD1_PIX2PIX; + definition.geometry = {8, 4, 3}; + definition.name = "Instruct-Pix2Pix"; + definition.traits.is_unet_edit = true; + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SD1_TINY_UNET; + definition.geometry = {8, 4, 3}; + definition.name = "SD 1.x Tiny UNet"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SD2; + definition.geometry = {8, 4, 3}; + definition.name = "SD 2.x"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.probe_v_prediction = true; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SD2_INPAINT; + definition.geometry = {8, 4, 3}; + definition.name = "SD 2.x Inpaint"; + definition.traits.is_inpaint = true; + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.probe_v_prediction = true; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SD2_TINY_UNET; + definition.geometry = {8, 4, 3}; + definition.name = "SD 2.x Tiny UNet"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.probe_v_prediction = true; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SDXS_512_DS; + definition.geometry = {8, 4, 3}; + definition.name = "SDXS (512-DS)"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SDXS_09; + definition.geometry = {8, 4, 3}; + definition.name = "SDXS (09)"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.probe_v_prediction = true; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SDXL; + definition.geometry = {8, 4, 3}; + definition.name = "SDXL"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.prediction_from_weights = detect_sdxl_prediction; + definition.configure_vae = configure_sdxl_vae; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SDXL_INPAINT; + definition.geometry = {8, 4, 3}; + definition.name = "SDXL Inpaint"; + definition.traits.is_inpaint = true; + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.prediction_from_weights = detect_sdxl_prediction; + definition.configure_vae = configure_sdxl_vae; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SDXL_PIX2PIX; + definition.geometry = {8, 4, 3}; + definition.name = "SDXL Instruct-Pix2Pix"; + definition.traits.is_unet_edit = true; + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.prediction_from_weights = detect_sdxl_prediction; + definition.configure_vae = configure_sdxl_vae; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SDXL_VEGA; + definition.geometry = {8, 4, 3}; + definition.name = "SDXL (Vega)"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.prediction_from_weights = detect_sdxl_prediction; + definition.configure_vae = configure_sdxl_vae; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SDXL_SSD1B; + definition.geometry = {8, 4, 3}; + definition.name = "SDXL (SSD1B)"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.prediction_from_weights = detect_sdxl_prediction; + definition.configure_vae = configure_sdxl_vae; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SVD; + definition.geometry = {8, 4, 3}; + definition.name = "SVD"; + definition.traits.supports_video_generation = true; + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_ESRGAN; + definition.geometry = {8, 4, 3}; + definition.name = "ESRGAN"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/sd3.cpp b/src/model/integrations/sd3.cpp new file mode 100644 index 000000000..f9c2ad780 --- /dev/null +++ b/src/model/integrations/sd3.cpp @@ -0,0 +1,40 @@ +#include "conditioning/conditioner.hpp" +#include "model/diffusion/mmdit.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_sd3(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& sd3_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_SD3; + definition.geometry = {1, 16, 3}; + definition.name = "SD3.x"; + definition.traits.is_dit = true; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd3; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/sensenova_u1.cpp b/src/model/integrations/sensenova_u1.cpp new file mode 100644 index 000000000..900f222e1 --- /dev/null +++ b/src/model/integrations/sensenova_u1.cpp @@ -0,0 +1,38 @@ +#include "model/diffusion/sensenova_u1.h" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_sensenova_u1(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "", + weight_manager); + return true; +} + +const std::vector& sensenova_u1_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_SENSENOVA_U1_5; + definition.geometry = {32, 3, 3}; + definition.name = "SenseNova U1.5"; + definition.traits.is_dit = true; + definition.traits.fake_vae = true; + definition.sampling = {SENSENOVA_U1_FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sensenova_u1; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/wan.cpp b/src/model/integrations/wan.cpp new file mode 100644 index 000000000..2daada294 --- /dev/null +++ b/src/model/integrations/wan.cpp @@ -0,0 +1,119 @@ +#include "model/diffusion/wan.hpp" +#include "conditioning/conditioner.hpp" +#include "model/audio/wav2vec2.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_wan(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + true, + 0, + true, + weight_manager); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager); + if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) { + result.high_noise_diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.high_noise_diffusion_model", + version, + weight_manager); + } + if (result.diffusion->get_desc() == "Wan2.1-I2V-14B" || + result.diffusion->get_desc() == "Wan2.1-FLF2V-14B" || + result.diffusion->get_desc() == "Wan2.1-I2V-1.3B") { + if (!sd::model_builders::ensure_backend_pair(ctx.backends, SDBackendModule::CLIP_VISION)) { + return false; + } + result.clip_vision = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::CLIP_VISION), + tensor_storage_map, + weight_manager); + } + if (version == VERSION_WAN2_2_S2V && + tensor_storage_map.count("wav2vec2.encoder.layer_norm.bias") > 0) { + if (!sd::model_builders::ensure_backend_pair(ctx.backends, SDBackendModule::AUDIO_ENCODER)) { + return false; + } + result.audio_encoder = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::AUDIO_ENCODER), + tensor_storage_map, + "wav2vec2.", + weight_manager); + } + return true; +} + +const std::vector& wan_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_WAN2; + definition.geometry = {2, 16, 3}; + definition.name = "Wan 2.x"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 5.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_wan; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_WAN2_2_I2V; + definition.geometry = {2, 16, 3}; + definition.name = "Wan 2.2 I2V"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 5.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_wan; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_WAN2_2_TI2V; + definition.geometry = {2, 48, 3}; + definition.name = "Wan 2.2 TI2V"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 5.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_wan; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_WAN2_2_S2V; + definition.geometry = {2, 16, 3}; + definition.name = "Wan 2.2 S2V"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 5.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_wan; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/z_image.cpp b/src/model/integrations/z_image.cpp new file mode 100644 index 000000000..29054d473 --- /dev/null +++ b/src/model/integrations/z_image.cpp @@ -0,0 +1,45 @@ +#include "model/diffusion/z_image.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_z_image(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager); + return true; +} + +const std::vector& z_image_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_Z_IMAGE; + definition.geometry = {1, 16, 3}; + definition.name = "Z-Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_z_image; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/z_image_l2p.cpp b/src/model/integrations/z_image_l2p.cpp new file mode 100644 index 000000000..d8104f853 --- /dev/null +++ b/src/model/integrations/z_image_l2p.cpp @@ -0,0 +1,44 @@ +#include "model/diffusion/z_image_l2p.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_z_image_l2p(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& z_image_l2p_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_Z_IMAGE_L2P; + definition.geometry = {16, 3, 3}; + definition.name = "Z-Image L2P"; + definition.traits.is_dit = true; + definition.traits.fake_vae = true; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_z_image_l2p; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/te/llada_image_te.hpp b/src/model/te/llada_image_te.hpp index 1c97a4529..2da78413b 100644 --- a/src/model/te/llada_image_te.hpp +++ b/src/model/te/llada_image_te.hpp @@ -347,8 +347,8 @@ namespace LLaDAImageTE { blocks["visual.blocks." + std::to_string(i)] = std::make_shared(config); } blocks["vqmodel.quant_conv"] = std::make_shared(config.hidden_size, - config.codebook_embed_dim, - std::make_pair(1, 1)); + config.codebook_embed_dim, + std::make_pair(1, 1)); blocks["prior_projector.net.0.proj"] = std::make_shared(config.semantic_embed_dim, config.semantic_embed_dim, true); blocks["prior_projector.net.2"] = std::make_shared(config.semantic_embed_dim, config.semantic_embed_dim, true); } diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index 8ea918025..85f751169 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -924,6 +924,8 @@ namespace LLM { int64_t txt_token_start = 0; int64_t txt_token_end = 0; ggml_tensor* input_embed = nullptr; + // Image and DeepStack indices remain in the expanded sequence coordinates. + int64_t removed_tokens = 0; for (int i = 0; i < image_embeds.size(); i++) { if (i == 0) { @@ -931,7 +933,10 @@ namespace LLM { } else { txt_token_start = image_embeds[i - 1].first + image_embeds[i - 1].second->ne[1]; } - txt_token_end = image_embeds[i].first; + txt_token_start -= removed_tokens; + txt_token_end = image_embeds[i].first - removed_tokens; + removed_tokens += image_embeds[i].second->ne[1] - 1; + GGML_ASSERT(txt_token_start >= 0 && txt_token_end >= txt_token_start && txt_token_end < raw_x->ne[1]); // An embed can sit flush against the previous one or at the very start/end of the // sequence, leaving no text tokens to splice around it. @@ -951,14 +956,14 @@ namespace LLM { } } - txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1]; + txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1] - removed_tokens; txt_token_end = raw_x->ne[1]; if (txt_token_end > txt_token_start) { auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); input_embed = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1); } - GGML_ASSERT(raw_x->ne[1] == input_embed->ne[1]); + GGML_ASSERT(raw_x->ne[1] + removed_tokens == input_embed->ne[1]); return input_embed; } @@ -2506,6 +2511,9 @@ namespace LLM { } int64_t n_tokens = input_ids->ne[0]; + for (const auto& image_embed : image_embeds) { + n_tokens += image_embed.second->ne[1] - 1; + } if (config.arch == LLMArch::MISTRAL_SMALL_3_2 || config.arch == LLMArch::MINISTRAL_3_3B || config.arch == LLMArch::QWEN3 || @@ -3014,20 +3022,17 @@ namespace LLM { LOG_VERBOSE("llm encode_image test done in %lldms", t1 - t0); } - std::string placeholder = "<|image_pad|>"; - std::string img_prompt = "Picture 1: <|vision_start|>"; // [24669, 220, 16, 25, 220, 151652] - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(num_image_tokens * placeholder.size()); - for (int i = 0; i < num_image_tokens; i++) { - img_prompt += placeholder; - } - img_prompt += "<|vision_end|>"; - - std::vector>> image_embeds; - image_embeds.emplace_back(64, image_embed); + const std::string image_prefix = "Picture 1: <|vision_start|>"; + const std::string img_prompt = image_prefix + "<|image_pad|><|vision_end|>"; std::pair prompt_attn_range; std::string text = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n"; + std::vector prefix_tokens; + if (!tokenizer->encode(text + image_prefix, prefix_tokens, nullptr)) { + return; + } + std::vector>> image_embeds; + image_embeds.emplace_back(static_cast(prefix_tokens.size()), image_embed); text += img_prompt; prompt_attn_range.first = static_cast(text.size()); text += "change 'flux.cpp' to 'edit.cpp'"; diff --git a/src/model/vae/ltx_audio_vae.hpp b/src/model/vae/ltx_audio_vae.hpp index 97c434146..a05ce86a4 100644 --- a/src/model/vae/ltx_audio_vae.hpp +++ b/src/model/vae/ltx_audio_vae.hpp @@ -1087,9 +1087,9 @@ namespace LTXV { auto& tensor_storage_map = model_loader.get_tensor_storage_map(); auto ltx_audio_vae = std::make_shared(backend, - tensor_storage_map, - prefix, - model_manager); + tensor_storage_map, + prefix, + model_manager); if (!model_manager->set_loader(std::move(model_loader)) || !model_manager->register_runner_params(ModelComponent::AudioVAE, diff --git a/src/model/vae/ltx_vae.hpp b/src/model/vae/ltx_vae.hpp index f680f2008..3fa42a166 100644 --- a/src/model/vae/ltx_vae.hpp +++ b/src/model/vae/ltx_vae.hpp @@ -497,12 +497,12 @@ namespace LTXVAE { GGML_ASSERT(out_channels % factor == 0); blocks["conv"] = std::make_shared(in_channels, - out_channels / factor, - 3, - std::tuple{1, 1, 1}, - 1, - true, - force_conv_prec_f32); + out_channels / factor, + 3, + std::tuple{1, 1, 1}, + 1, + true, + force_conv_prec_f32); blocks["skip_downsample"] = std::make_shared(in_channels, out_channels, factor_t, factor_s); blocks["conv_downsample"] = std::make_shared(out_channels / factor, out_channels, factor_t, factor_s); } @@ -1077,11 +1077,11 @@ namespace LTXVAE { patch_size); } blocks["decoder"] = std::make_shared(version, - tensor_storage_map, - prefix, - patch_size, - false, - timestep_conditioning); + tensor_storage_map, + prefix, + patch_size, + false, + timestep_conditioning); blocks["per_channel_statistics"] = std::make_shared(); } diff --git a/src/model/vae/minimax_h3_audio_vae.hpp b/src/model/vae/minimax_h3_audio_vae.hpp index 59c30a27d..02829a4d0 100644 --- a/src/model/vae/minimax_h3_audio_vae.hpp +++ b/src/model/vae/minimax_h3_audio_vae.hpp @@ -76,10 +76,10 @@ namespace MiniMaxH3 { blocks["block.2"] = std::make_shared(in_channels, 9); blocks["block.3"] = std::make_shared(in_channels); blocks["block.4"] = std::make_shared(in_channels, - out_channels, - 2 * stride, - stride, - static_cast(std::ceil(stride / 2.f))); + out_channels, + 2 * stride, + stride, + static_cast(std::ceil(stride / 2.f))); } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { @@ -300,12 +300,12 @@ namespace MiniMaxH3 { } blocks["activation_post"] = std::make_shared(channels); blocks["conv_post"] = std::make_shared(channels, - 1, - 7, - 1, - 3, - 1, - false); + 1, + 7, + 1, + 3, + 1, + false); } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index 98cf19c18..d397a40df 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -449,8 +449,8 @@ namespace MiniMaxH3VAE { MiniMaxH3VideoVAE() { blocks["encoder"] = std::make_shared(); blocks["quant_conv"] = std::make_shared(48, - 48, - std::tuple{1, 1, 1}); + 48, + std::tuple{1, 1, 1}); blocks["post_quant_conv"] = std::make_shared(24, 24, std::tuple{1, 1, 1}); diff --git a/src/model/vae/vae_tiling.hpp b/src/model/vae/vae_tiling.hpp index 50cf08355..acad04623 100644 --- a/src/model/vae/vae_tiling.hpp +++ b/src/model/vae/vae_tiling.hpp @@ -15,8 +15,9 @@ enum class VAETemporalDirection { }; struct VAETemporalTilingConfig { - int tile_frames = 1; - int overlap = 0; + int tile_frames = 1; + int overlap = 0; + int decode_frames = 1; }; struct VAETemporalTile { @@ -43,7 +44,8 @@ inline VAETemporalTilingConfig resolve_vae_temporal_tiling_config(const sd_tilin config.overlap = std::max(0, default_overlap); for (const auto& [key, value] : parse_key_value_args(params.extra_tiling_args, "VAE extra tiling arg")) { - if (key != "temporal_tile_frames" && key != "temporal_tile_size" && key != "temporal_tile_overlap") { + if (key != "temporal_tile_frames" && key != "temporal_tile_size" && key != "temporal_tile_overlap" && + key != "temporal_decode_frames") { continue; } @@ -52,6 +54,8 @@ inline VAETemporalTilingConfig resolve_vae_temporal_tiling_config(const sd_tilin LOG_WARN("ignoring invalid VAE extra tiling arg '%s=%s'", key.c_str(), value.c_str()); } else if (key == "temporal_tile_overlap") { config.overlap = std::max(0, parsed); + } else if (key == "temporal_decode_frames") { + config.decode_frames = std::max(1, parsed); } else { config.tile_frames = std::max(1, parsed); } diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index 9dd30fdff..6b94bcb0d 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -1287,9 +1287,11 @@ namespace WAN { ggml_tensor* decode_tiled_chunk(GGMLRunnerContext* ctx, ggml_tensor* z, int chunk_idx, - int64_t b = 1) { + int64_t b = 1, + int decode_frames = 1) { // z: [b*c, t, h, w] GGML_ASSERT(b == 1); + GGML_ASSERT(decode_frames >= 1); auto decoder = std::dynamic_pointer_cast(blocks["decoder"]); auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); @@ -1303,11 +1305,13 @@ namespace WAN { } ggml_tensor* out = nullptr; - for (int64_t frame = 0; frame < x->ne[2]; ++frame) { + for (int64_t frame = 0; frame < x->ne[2];) { const int global_frame = chunk_idx + static_cast(frame); - auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + 1); - _conv_idx = 0; - auto out_frame = decoder->forward(ctx, in, b, _feat_map, _conv_idx, global_frame); + // The first latent frame bypasses temporal upsampling and primes the causal caches. + const int64_t count = !is_2D && global_frame > 0 ? std::min(decode_frames, x->ne[2] - frame) : 1; + auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + count); + _conv_idx = 0; + auto out_frame = decoder->forward(ctx, in, b, _feat_map, _conv_idx, global_frame); if (is_2D && global_frame > 0) { auto repeated = out_frame; for (int repeat = 1; repeat < 4; ++repeat) { @@ -1316,6 +1320,7 @@ namespace WAN { out_frame = repeated; } out = out == nullptr ? out_frame : ggml_concat(ctx->ggml_ctx, out, out_frame, 2); + frame += count; } return unpatchify(ctx->ggml_ctx, out, patch_size, b); } @@ -1455,7 +1460,7 @@ namespace WAN { return gf; } - ggml_cgraph* build_temporal_tile_graph(const sd::Tensor& z_tensor, int chunk_idx) { + ggml_cgraph* build_temporal_tile_graph(const sd::Tensor& z_tensor, int chunk_idx, int decode_frames) { ggml_cgraph* gf = new_graph_custom(std::max(20480, 10240 * z_tensor.shape()[2])); ae.clear_cache(); @@ -1469,7 +1474,8 @@ namespace WAN { auto runner_ctx = get_context(gf); - ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx); + ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx, 1, + version == VERSION_WAN2_2_TI2V ? decode_frames : 1); for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) { ggml_tensor* feat_cache = ae._feat_map[feat_idx]; @@ -1492,8 +1498,19 @@ namespace WAN { stateful_config.overlap = 0; auto plan = make_vae_temporal_tile_plan(input.shape()[2], stateful_config); - LOG_VERBOSE("Wan VAE stateful temporal tiling: tile_frames=%d, total latent frames=%lld, tiles=%d", + if (sd_version_is_wan(version) && !ae.is_2D && plan.tiles.size() > 1 && + plan.tiles.size() <= 3 && plan.tile_frames > 2 && + plan.tiles.back().end - plan.tiles.back().start == 1) { + // Longer plans retain full middle tiles, so balancing their tail cannot shrink the largest decoded output. + auto& previous = plan.tiles[plan.tiles.size() - 2]; + auto& last = plan.tiles.back(); + --previous.end; + --last.start; + } + + LOG_VERBOSE("Wan VAE stateful temporal tiling: tile_frames=%d, decode_frames=%d, total latent frames=%lld, tiles=%d", plan.tile_frames, + version == VERSION_WAN2_2_TI2V && !ae.is_2D ? config.decode_frames : 1, (long long)input.shape()[2], (int)plan.tiles.size()); @@ -1507,7 +1524,7 @@ namespace WAN { (long long)tile.start, (long long)tile.end); auto get_graph = [&]() -> ggml_cgraph* { - return build_temporal_tile_graph(input_tile, static_cast(tile.start)); + return build_temporal_tile_graph(input_tile, static_cast(tile.start), config.decode_frames); }; return restore_trailing_singleton_dims( GGMLRunner::compute(get_graph, n_threads, false), diff --git a/src/model_definition.h b/src/model_definition.h new file mode 100644 index 000000000..2b4999495 --- /dev/null +++ b/src/model_definition.h @@ -0,0 +1,47 @@ +#ifndef __SD_MODEL_DEFINITION_H__ +#define __SD_MODEL_DEFINITION_H__ + +#include +#include "model.h" +#include "stable-diffusion.h" + +class TokenizerConfig; +struct VAE; +namespace sd::model_builders { + struct Context; + struct CoreRunners; + struct VAEOptions; +} + +struct ModelSamplingPolicy { + prediction_t default_prediction = EPS_PRED; + float default_flow_shift = std::numeric_limits::infinity(); + sample_method_t default_sample_method = EULER_A_SAMPLE_METHOD; + scheduler_t default_scheduler = DISCRETE_SCHEDULER; + + bool probe_v_prediction = false; + + prediction_t (*prediction_from_weights)(const String2TensorStorage&) = nullptr; + float (*flow_shift_from_weights)(const String2TensorStorage&) = nullptr; +}; + +struct ModelGeometry { + int diffusion_down_factor = 8; + int latent_channels = 4; + int image_channels = 3; +}; + +struct ModelDefinition { + SDVersion version = VERSION_COUNT; + const char* name = nullptr; + ModelTraits traits; + ModelSamplingPolicy sampling; + ModelGeometry geometry; + const char* ref_image_preset = "default"; + void (*configure_vae)(const sd::model_builders::Context&, const sd::model_builders::VAEOptions&, VAE&) = nullptr; + bool (*build)(const sd::model_builders::Context&, const TokenizerConfig&, sd::model_builders::CoreRunners&) = nullptr; +}; + +const ModelDefinition* find_model_definition(SDVersion version); + +#endif // __SD_MODEL_DEFINITION_H__ diff --git a/src/model_loader.cpp b/src/model_loader.cpp index a0f500819..63d681de6 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -340,7 +340,7 @@ bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path, const auto expected = (std::filesystem::u8path(file_path).parent_path() / std::filesystem::u8path(entry.value().get())) .lexically_normal(); - bool found = false; + bool found = false; for (const auto& shard : parsed_tensor_names_) { if (std::filesystem::u8path(shard.first).lexically_normal() == expected) { found = shard.second.count(entry.key()) != 0; diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 921989df7..243811d9e 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -568,7 +568,7 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector& states // The temporary runner is destroyed before this manager call returns. auto borrowed_manager = std::shared_ptr(this, [](ModelManager*) {}); auto lora = std::make_shared(id, compute_backend, target->params_backend, - borrowed_manager, lora_spec.file_id, lora_version_, - target->residency_mode); + borrowed_manager, lora_spec.file_id, lora_version_, + target->residency_mode); LoraModel::filter_t lora_tensor_filter = nullptr; if (!lora_spec.tensor_name_prefix_filter.empty()) { @@ -1409,8 +1409,8 @@ bool ModelManager::assign_compute_backend(const std::vector& tenso const bool params_follow_compute = state->params_follow_compute_backend || state->residency_mode == ResidencyMode::Disk; - const bool compute_changes = state->compute_backend != compute_backend; - const bool params_changes = params_follow_compute && state->params_backend != compute_backend; + const bool compute_changes = state->compute_backend != compute_backend; + const bool params_changes = params_follow_compute && state->params_backend != compute_backend; if (!compute_changes && !params_changes) { continue; } @@ -1501,7 +1501,7 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe ? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device : std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) { return ggml_backend_get_device(entry.first) == compute_device; - }); + }); if (!on_device) { return; } @@ -1782,10 +1782,10 @@ ModelManager::CapacityCheck ModelManager::check_capacity( const size_t resident = add(weights_resident, add(other_runtime, request.runtime_resident_bytes)); if (log_details) { LOG_WARN("model manager memory on %s: reported free %.2f MB / total %.2f MB, tracked weights %.2f MB / other runtime %.2f MB / current runtime %.2f MB", - ggml_backend_name(backend), - free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0), - weights_resident / (1024.0 * 1024.0), other_runtime / (1024.0 * 1024.0), - request.runtime_resident_bytes / (1024.0 * 1024.0)); + ggml_backend_name(backend), + free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0), + weights_resident / (1024.0 * 1024.0), other_runtime / (1024.0 * 1024.0), + request.runtime_resident_bytes / (1024.0 * 1024.0)); } if (free_bytes == 0 && total_bytes == 0) { return SIZE_MAX; @@ -1797,7 +1797,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity( } // A wrapped report may reflect real pressure; reclaim before accepting the tracked estimate. LOG_VERBOSE("ignoring impossible free memory report on %s (%.2f MB free / %.2f MB total), using tracked residency", - ggml_backend_name(backend), free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0)); + ggml_backend_name(backend), free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0)); free_bytes = total_bytes; result.requires_reclamation = true; } @@ -1907,9 +1907,9 @@ bool ModelManager::ensure_compute_backend_capacity( } const bool reloadable = state->residency_mode == ResidencyMode::Disk || state->compute_backend != state->params_backend; - const bool resident = state->compute_backend == state->params_backend - ? state->loaded_to_params_backend - : state->staged_to_compute_backend; + const bool resident = state->compute_backend == state->params_backend + ? state->loaded_to_params_backend + : state->staged_to_compute_backend; if (reloadable && resident) { eviction_states.insert(state); } diff --git a/src/model_registry.cpp b/src/model_registry.cpp new file mode 100644 index 000000000..9a80bb835 --- /dev/null +++ b/src/model_registry.cpp @@ -0,0 +1,102 @@ +#include +#include +#include "model_definition.h" + +const std::vector& fibo2_definitions(); +const std::vector& sd3_definitions(); +const std::vector& pid_definitions(); +const std::vector& ideogram4_definitions(); +const std::vector& krea2_definitions(); +const std::vector& flux_definitions(); +const std::vector& flux3_definitions(); +const std::vector& flux2_definitions(); +const std::vector& ltxav_definitions(); +const std::vector& minimax_h3_definitions(); +const std::vector& hunyuan_video_definitions(); +const std::vector& wan_definitions(); +const std::vector& lingbot_video_definitions(); +const std::vector& qwen_image_definitions(); +const std::vector& pixart_definitions(); +const std::vector& mage_flow_definitions(); +const std::vector& longcat_definitions(); +const std::vector& hidream_o1_definitions(); +const std::vector& minit2i_definitions(); +const std::vector& sensenova_u1_definitions(); +const std::vector& anima_definitions(); +const std::vector& ming_image_definitions(); +const std::vector& z_image_definitions(); +const std::vector& z_image_l2p_definitions(); +const std::vector& llada_image_definitions(); +const std::vector& boogu_image_definitions(); +const std::vector& ernie_image_definitions(); +const std::vector& lens_definitions(); +const std::vector& iris_definitions(); +const std::vector& sd_definitions(); +const std::vector& moss_image_definitions(); + +static const std::array& model_definitions() { + static const auto definitions = [] { + std::array result{}; + auto append = [&](const std::vector& family) { + for (const auto& definition : family) { + if (definition.version < 0 || definition.version >= VERSION_COUNT || result[definition.version]) { + throw std::logic_error("duplicate or invalid model registration"); + } + result[definition.version] = &definition; + } + }; + append(fibo2_definitions()); + append(sd3_definitions()); + append(pid_definitions()); + append(ideogram4_definitions()); + append(krea2_definitions()); + append(flux_definitions()); + append(flux3_definitions()); + append(flux2_definitions()); + append(ltxav_definitions()); + append(minimax_h3_definitions()); + append(hunyuan_video_definitions()); + append(wan_definitions()); + append(lingbot_video_definitions()); + append(qwen_image_definitions()); + append(pixart_definitions()); + append(mage_flow_definitions()); + append(longcat_definitions()); + append(hidream_o1_definitions()); + append(minit2i_definitions()); + append(sensenova_u1_definitions()); + append(anima_definitions()); + append(ming_image_definitions()); + append(z_image_definitions()); + append(z_image_l2p_definitions()); + append(llada_image_definitions()); + append(boogu_image_definitions()); + append(ernie_image_definitions()); + append(lens_definitions()); + append(iris_definitions()); + append(sd_definitions()); + append(moss_image_definitions()); + for (const auto* definition : result) { + if (!definition) + throw std::logic_error("missing model registration"); + } + return result; + }(); + return definitions; +} + +const ModelDefinition* find_model_definition(SDVersion version) { + if (version < 0 || version >= VERSION_COUNT) + return nullptr; + return model_definitions()[version]; +} + +const ModelTraits* find_model_traits(SDVersion version) { + const auto* definition = find_model_definition(version); + return definition ? &definition->traits : nullptr; +} + +const char* sd_version_name(SDVersion version) { + const auto* definition = find_model_definition(version); + return definition ? definition->name : "Unknown"; +} diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 1e3c2e24b..d7a0f005c 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -21,6 +21,7 @@ #include "core/probe/conv_direct.h" #include "core/probe/flash_attn.h" #include "model.h" +#include "model_definition.h" #include "core/rng.hpp" #include "core/rng_mt19937.hpp" @@ -41,6 +42,7 @@ #include "model/diffusion/animatediff.hpp" #include "model/diffusion/control.hpp" #include "model/diffusion/model.hpp" +#include "model/diffusion/z_image.hpp" #include "model/vae/audio_vae.hpp" #include "model/vae/ltx_vae.hpp" #include "model/vae/vae.hpp" @@ -55,70 +57,6 @@ #include -const char* model_version_to_str[] = { - "SD 1.x", - "SD 1.x Inpaint", - "Instruct-Pix2Pix", - "SD 1.x Tiny UNet", - "SD 2.x", - "SD 2.x Inpaint", - "SD 2.x Tiny UNet", - "SDXS (512-DS)", - "SDXS (09)", - "SDXL", - "SDXL Inpaint", - "SDXL Instruct-Pix2Pix", - "SDXL (Vega)", - "SDXL (SSD1B)", - "SVD", - "SD3.x", - "Flux", - "Flux Fill", - "Flux Control", - "Flex.2", - "Chroma Radiance", - "Wan 2.x", - "Wan 2.2 I2V", - "Wan 2.2 TI2V", - "Wan 2.2 S2V", - "LingBot Video", - "Qwen Image", - "Qwen Image Layered", - "Qwen Image 2.1", - "Hunyuan Video", - "Anima", - "Flux.2", - "Flux.2 klein", - "LTXAV", - "MiniMax-H3", - "HiDream O1", - "Z-Image", - "Boogu Image", - "Ovis Image", - "Ernie Image", - "Lens", - "MiniT2I", - "Longcat-Image", - "PiD", - "Ideogram 4", - "SeFi-Image", - "Krea2", - "Mage Flow", - "SenseNova U1.5", - "LLaDA-Image", - "ESRGAN", - "PixArt", - "Ming-Image", - "Z-Image L2P", - "Iris-3B", - "Flux.3 Image Klein", - "MOSS-Image", - "Fibo-2.0", -}; - -static_assert(VERSION_COUNT == sizeof(model_version_to_str) / sizeof(model_version_to_str[0]), - "\nnumber of elements in model_version_to_str[] != VERSION_COUNT"); - void calculate_alphas_cumprod(float* alphas_cumprod, float linear_start = 0.00085f, float linear_end = 0.0120f, @@ -297,6 +235,15 @@ bool StableDiffusionGGML::apply_model_update(ModelLoader candidate, } } } + if (sd_version_is_z_image(version) && sd_version_is_z_image(next_version)) { + if (groups.count(RunnerGroup::ControlNet) > 0) { + groups.insert(RunnerGroup::Core); + } + if (groups.count(RunnerGroup::Core) > 0 && + model_manager->registered_params_size({ModelComponent::ControlNet}) > 0) { + groups.insert(RunnerGroup::ControlNet); + } + } runner_state_.ready = false; if (!reset_runners(groups)) { return false; @@ -470,7 +417,7 @@ bool StableDiffusionGGML::register_row_split_runner_params(ModelComponent compon size_t free_bytes = 0, total_bytes = 0; ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); int64_t usable_bytes = std::max((int64_t)free_bytes - compute_headroom_bytes, - (int64_t)free_bytes / 8); + (int64_t)free_bytes / 8); tensor_split[reg_index] = usable_bytes > 0 ? (float)((double)usable_bytes / (1024.0 * 1024.0)) : 1.0f; } @@ -574,8 +521,8 @@ bool StableDiffusionGGML::register_layer_split_runner_params(ModelComponent comp model->set_runtime_backends(module_backends); model->set_graph_cut_layer_split_backend_vram_limits(layer_split_vram_limits_for_backends(module_backends)); model->set_graph_cut_layer_split_enabled(true); - const bool params_follow_runtime = backend_manager.params_backend_follows_runtime(module) || - backend_manager.params_backend_is_disk(module); + const bool params_follow_runtime = backend_manager.params_backend_follows_runtime(module) || + backend_manager.params_backend_is_disk(module); ggml_backend_t initial_params_backend = params_follow_runtime ? module_backends[0] : params_backend_for(module); if (initial_params_backend == nullptr) { return false; @@ -728,87 +675,93 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf const auto* sd_ctx_params = &configuration.params; auto& use_tae = configuration.use_tae; auto& use_audio_vae = configuration.use_audio_vae; + const auto load_model = [&](const char* name, const char* path, const std::string& prefix = "") { + if (!model_loader.init_from_file(path, prefix)) + return false; + configuration.record_model(name, path); + return true; + }; if (strlen(SAFE_STR(sd_ctx_params->model_path)) > 0) { LOG_INFO("loading model from '%s'", sd_ctx_params->model_path); - if (!model_loader.init_from_file(sd_ctx_params->model_path)) { + if (!load_model("model", sd_ctx_params->model_path)) { LOG_ERROR("init model loader from file failed: '%s'", sd_ctx_params->model_path); } } if (strlen(SAFE_STR(sd_ctx_params->diffusion_model_path)) > 0) { LOG_INFO("loading diffusion model from '%s'", sd_ctx_params->diffusion_model_path); - if (!model_loader.init_from_file(sd_ctx_params->diffusion_model_path, "model.diffusion_model.")) { + if (!load_model("diffusion_model", sd_ctx_params->diffusion_model_path, "model.diffusion_model.")) { LOG_WARN("loading diffusion model from '%s' failed", sd_ctx_params->diffusion_model_path); } } if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) { LOG_INFO("loading high noise diffusion model from '%s'", sd_ctx_params->high_noise_diffusion_model_path); - if (!model_loader.init_from_file(sd_ctx_params->high_noise_diffusion_model_path, "model.high_noise_diffusion_model.")) { + if (!load_model("high_noise_diffusion_model", sd_ctx_params->high_noise_diffusion_model_path, "model.high_noise_diffusion_model.")) { LOG_WARN("loading diffusion model from '%s' failed", sd_ctx_params->high_noise_diffusion_model_path); } } if (strlen(SAFE_STR(sd_ctx_params->uncond_diffusion_model_path)) > 0) { LOG_INFO("loading unconditional diffusion model from '%s'", sd_ctx_params->uncond_diffusion_model_path); - if (!model_loader.init_from_file(sd_ctx_params->uncond_diffusion_model_path, "model.diffusion_model.uncond.")) { + if (!load_model("uncond_diffusion_model", sd_ctx_params->uncond_diffusion_model_path, "model.diffusion_model.uncond.")) { LOG_WARN("loading unconditional diffusion model from '%s' failed", sd_ctx_params->uncond_diffusion_model_path); } } if (strlen(SAFE_STR(sd_ctx_params->clip_l_path)) > 0) { LOG_INFO("loading clip_l from '%s'", sd_ctx_params->clip_l_path); - if (!model_loader.init_from_file(sd_ctx_params->clip_l_path, "clip_l.")) { + if (!load_model("clip_l", sd_ctx_params->clip_l_path, "clip_l.")) { LOG_WARN("loading clip_l from '%s' failed", sd_ctx_params->clip_l_path); } } if (strlen(SAFE_STR(sd_ctx_params->clip_g_path)) > 0) { LOG_INFO("loading clip_g from '%s'", sd_ctx_params->clip_g_path); - if (!model_loader.init_from_file(sd_ctx_params->clip_g_path, "clip_g.")) { + if (!load_model("clip_g", sd_ctx_params->clip_g_path, "clip_g.")) { LOG_WARN("loading clip_g from '%s' failed", sd_ctx_params->clip_g_path); } } if (strlen(SAFE_STR(sd_ctx_params->clip_vision_path)) > 0) { LOG_INFO("loading clip_vision from '%s'", sd_ctx_params->clip_vision_path); - if (!model_loader.init_from_file(sd_ctx_params->clip_vision_path, "clip_vision.")) { + if (!load_model("clip_vision", sd_ctx_params->clip_vision_path, "clip_vision.")) { LOG_WARN("loading clip_vision from '%s' failed", sd_ctx_params->clip_vision_path); } } if (strlen(SAFE_STR(sd_ctx_params->t5xxl_path)) > 0) { LOG_INFO("loading t5xxl from '%s'", sd_ctx_params->t5xxl_path); - if (!model_loader.init_from_file(sd_ctx_params->t5xxl_path, "text_encoders.t5xxl.transformer.")) { + if (!load_model("t5xxl", sd_ctx_params->t5xxl_path, "text_encoders.t5xxl.transformer.")) { LOG_WARN("loading t5xxl from '%s' failed", sd_ctx_params->t5xxl_path); } } if (strlen(SAFE_STR(sd_ctx_params->pulid_weights_path)) > 0) { LOG_INFO("loading PuLID weights from '%s'", sd_ctx_params->pulid_weights_path); - if (!model_loader.init_from_file(sd_ctx_params->pulid_weights_path, - "model.diffusion_model.")) { + if (!load_model("pulid_weights", sd_ctx_params->pulid_weights_path, + "model.diffusion_model.")) { LOG_WARN("loading PuLID weights from '%s' failed", sd_ctx_params->pulid_weights_path); } } if (strlen(SAFE_STR(sd_ctx_params->llm_path)) > 0) { LOG_INFO("loading llm from '%s'", sd_ctx_params->llm_path); - if (!model_loader.init_from_file(sd_ctx_params->llm_path, "text_encoders.llm.")) { + if (!load_model("llm", sd_ctx_params->llm_path, "text_encoders.llm.")) { LOG_WARN("loading llm from '%s' failed", sd_ctx_params->llm_path); } } if (strlen(SAFE_STR(sd_ctx_params->llm_vision_path)) > 0) { LOG_INFO("loading llm vision from '%s'", sd_ctx_params->llm_vision_path); - if (!model_loader.init_from_file(sd_ctx_params->llm_vision_path, "text_encoders.llm.visual.")) { + if (!load_model("llm_vision", sd_ctx_params->llm_vision_path, "text_encoders.llm.visual.")) { LOG_WARN("loading llm vision from '%s' failed", sd_ctx_params->llm_vision_path); } } if (strlen(SAFE_STR(sd_ctx_params->vae_path)) > 0) { LOG_INFO("loading vae from '%s'", sd_ctx_params->vae_path); - if (!model_loader.init_from_file(sd_ctx_params->vae_path, "vae.")) { + if (!load_model("vae", sd_ctx_params->vae_path, "vae.")) { LOG_WARN("loading vae from '%s' failed", sd_ctx_params->vae_path); external_vae_is_invalid = true; } @@ -816,7 +769,7 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf if (strlen(SAFE_STR(sd_ctx_params->taesd_path)) > 0) { LOG_INFO("loading tae from '%s'", sd_ctx_params->taesd_path); - if (!model_loader.init_from_file(sd_ctx_params->taesd_path, "tae.")) { + if (!load_model("taesd", sd_ctx_params->taesd_path, "tae.")) { LOG_WARN("loading tae from '%s' failed", sd_ctx_params->taesd_path); } else { use_tae = true; @@ -825,14 +778,14 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf if (strlen(SAFE_STR(sd_ctx_params->embeddings_connectors_path)) > 0) { LOG_INFO("loading embeddings connectors from '%s'", sd_ctx_params->embeddings_connectors_path); - if (!model_loader.init_from_file(sd_ctx_params->embeddings_connectors_path)) { + if (!load_model("embeddings_connectors", sd_ctx_params->embeddings_connectors_path)) { LOG_WARN("loading embeddings connectors from '%s' failed", sd_ctx_params->embeddings_connectors_path); } } if (strlen(SAFE_STR(sd_ctx_params->audio_vae_path)) > 0) { LOG_INFO("loading audio VAE from '%s'", sd_ctx_params->audio_vae_path); - if (!model_loader.init_from_file(sd_ctx_params->audio_vae_path)) { + if (!load_model("audio_vae", sd_ctx_params->audio_vae_path)) { LOG_WARN("loading audio VAE weights from '%s' failed", sd_ctx_params->audio_vae_path); } else { use_audio_vae = true; @@ -841,15 +794,15 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf if (strlen(SAFE_STR(sd_ctx_params->audio_encoder_path)) > 0) { LOG_INFO("loading audio encoder (wav2vec2) from '%s'", sd_ctx_params->audio_encoder_path); - if (!model_loader.init_from_file(sd_ctx_params->audio_encoder_path, "wav2vec2.")) { + if (!load_model("audio_encoder", sd_ctx_params->audio_encoder_path, "wav2vec2.")) { LOG_WARN("loading audio encoder weights from '%s' failed", sd_ctx_params->audio_encoder_path); } } 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.")) { + if (!load_model("motion_module", 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 { configuration.animatediff_loaded = true; @@ -863,8 +816,11 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf } } + if (configuration.control_net_file) + configuration.record_model("control_net", sd_ctx_params->control_net_path); + if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0) { - if (!model_loader.init_from_file(sd_ctx_params->ip_adapter_path)) { + if (!load_model("ip_adapter", sd_ctx_params->ip_adapter_path)) { LOG_ERROR("init ip-adapter model loader from file failed: '%s'", sd_ctx_params->ip_adapter_path); return false; } @@ -872,6 +828,8 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf if (strlen(SAFE_STR(sd_ctx_params->photo_maker_path)) > 0) { configuration.photomaker_source_available = model_loader.add_file(sd_ctx_params->photo_maker_path, "pmid."); + if (configuration.photomaker_source_available) + configuration.record_model("photo_maker", sd_ctx_params->photo_maker_path); if (!configuration.photomaker_source_available) { LOG_WARN("loading stacked ID embedding from '%s' failed", sd_ctx_params->photo_maker_path); } @@ -1014,7 +972,7 @@ bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) { LOG_ERROR("get sd version from file failed: '%s'", SAFE_STR(sd_ctx_params->model_path)); return false; } else { - LOG_INFO("Version: %s ", model_version_to_str[version]); + LOG_INFO("Version: %s ", sd_version_name(version)); } if (auto_fit_enabled) { @@ -1179,6 +1137,27 @@ bool StableDiffusionGGML::build_control_net_runner() { if (config_->control_net_file == 0) { return true; } + if (sd_version_is_z_image(version)) { + auto z_image = std::dynamic_pointer_cast(diffusion_model); + if (z_image == nullptr || !z_image->control_enabled()) { + LOG_WARN("z-image control net file present but the diffusion model has no control weights; ignoring"); + return true; + } + std::map control_tensors; + z_image->get_control_param_tensors(control_tensors); + if (control_tensors.empty()) { + LOG_ERROR("z-image control net has no param tensors to register"); + return false; + } + ModelManager::ResidencyMode residency_mode = backend_manager.params_backend_is_disk(SDBackendModule::DIFFUSION) + ? ModelManager::ResidencyMode::Disk + : ModelManager::ResidencyMode::ParamBackend; + return model_manager->register_param_tensors(ModelComponent::ControlNet, + std::move(control_tensors), + residency_mode, + backend_for(SDBackendModule::DIFFUSION), + params_backend_for(SDBackendModule::DIFFUSION)); + } if (!sd::model_builders::build_control_net_runner(model_build_context(), control_net)) { return false; } @@ -1470,98 +1449,16 @@ bool StableDiffusionGGML::build_denoiser() { prediction_t pred_type = sd_ctx_params->prediction; if (pred_type == PREDICTION_COUNT) { - if (sd_version_is_sd2(version)) { + const auto* definition = find_model_definition(version); + if (!definition) + return false; + const auto& policy = definition->sampling; + pred_type = policy.default_prediction; + default_flow_shift = policy.flow_shift_from_weights ? policy.flow_shift_from_weights(tensor_storage_map) : policy.default_flow_shift; + if (policy.probe_v_prediction) { pred_type = is_using_v_parameterization_for_sd2(sd_version_is_inpaint(version)) ? V_PRED : EPS_PRED; - } else if (sd_version_is_sdxl(version)) { - if (tensor_storage_map.find("edm_vpred.sigma_max") != tensor_storage_map.end()) { - // CosXL models - // TODO: get sigma_min and sigma_max values from file - pred_type = EDM_V_PRED; - } else if (tensor_storage_map.find("v_pred") != tensor_storage_map.end()) { - pred_type = V_PRED; - } else { - pred_type = EPS_PRED; - } - } else if (sd_version_is_sd3(version) || - sd_version_is_wan(version) || - sd_version_is_hunyuan_video(version) || - sd_version_is_lingbot_video(version) || - sd_version_is_minimax_h3(version) || - sd_version_is_qwen_image(version) || - sd_version_is_mage_flow(version) || - version == VERSION_HIDREAM_O1 || - sd_version_is_anima(version) || - sd_version_is_ernie_image(version) || - sd_version_is_z_image(version) || - sd_version_is_z_image_l2p(version) || - version == VERSION_MING_IMAGE || - sd_version_is_llada_image(version) || - sd_version_is_boogu_image(version) || - sd_version_is_pid(version) || - sd_version_is_ideogram4(version) || - version == VERSION_IRIS || - version == VERSION_MOSS_IMAGE || - version == VERSION_FIBO2) { - pred_type = FLOW_PRED; - if (sd_version_is_wan(version) || version == VERSION_MOSS_IMAGE) { - default_flow_shift = 5.f; - } else if (sd_version_is_hunyuan_video(version)) { - default_flow_shift = 7.f; - } else if (sd_version_is_minimax_h3(version)) { - default_flow_shift = 12.f; - } else if (sd_version_is_ernie_image(version) || version == VERSION_IRIS) { - default_flow_shift = 4.f; - } else if (sd_version_is_pid(version)) { - default_flow_shift = 1.5f; - } else if (sd_version_is_ideogram4(version)) { - default_flow_shift = 1.0f; - } else if (sd_version_is_boogu_image(version)) { - default_flow_shift = 3.16f; - } else if (sd_version_is_mage_flow(version)) { - default_flow_shift = 6.f; - } else if (version == VERSION_MING_IMAGE) { - default_flow_shift = INFINITY; - } else if (sd_version_is_llada_image(version)) { - default_flow_shift = 1.0f; // unused: LLADA_IMAGE_SCHEDULER builds a fixed grid - } else { - default_flow_shift = 3.f; - } - } else if (sd_version_is_flux3(version)) { - pred_type = FLUX_FLOW_PRED; - default_flow_shift = 6.93f; - } else if (sd_version_is_flux(version) || - sd_version_is_flux2(version) || - sd_version_is_longcat(version) || - sd_version_is_lens(version) || - sd_version_is_ltxav(version) || - sd_version_is_krea2(version)) { - pred_type = FLUX_FLOW_PRED; - - default_flow_shift = 1.0f; // TODO: validate - for (const auto& [name, tensor_storage] : tensor_storage_map) { - if (starts_with(name, "model.diffusion_model.guidance_in.in_layer.weight")) { - default_flow_shift = 1.15f; - break; - } - } - if (sd_version_is_longcat(version)) { - default_flow_shift = 3.0f; - } else if (sd_version_is_lens(version)) { - default_flow_shift = 1.83f; - } else if (sd_version_is_ltxav(version)) { - default_flow_shift = 2.37f; - } else if (sd_version_is_krea2(version)) { - default_flow_shift = 1.15f; - } - } else if (sd_version_is_sefi_image(version)) { - pred_type = SEFI_FLOW_PRED; - } else if (sd_version_is_minit2i(version)) { - pred_type = MINIT2I_FLOW_PRED; - } else if (sd_version_is_sensenova_u1(version)) { - pred_type = SENSENOVA_U1_FLOW_PRED; - default_flow_shift = 3.f; - } else { - pred_type = EPS_PRED; + } else if (policy.prediction_from_weights) { + pred_type = policy.prediction_from_weights(tensor_storage_map); } } @@ -1578,27 +1475,28 @@ bool StableDiffusionGGML::build_denoiser() { denoiser = std::make_shared(); break; case FLOW_PRED: { - if (sd_version_is_ltxav(version)) { - LOG_INFO("running in LTXAV FLOW mode"); - denoiser = std::make_shared(); - } else if (sd_version_is_minimax_h3(version)) { - LOG_INFO("running in MiniMax H3 AV FLOW mode"); - denoiser = std::make_shared(default_flow_shift, 3.f, get_latent_channel()); - } else if (version == VERSION_MING_IMAGE) { - denoiser = std::make_shared(); - } else { - LOG_INFO("running in FLOW mode"); - denoiser = std::make_shared(); - } + LOG_INFO("running in FLOW mode"); + denoiser = std::make_shared(); break; } case FLUX_FLOW_PRED: { LOG_INFO("running in Flux FLOW mode"); - if (sd_version_is_flux3(version)) { - denoiser = std::make_shared(); - } else { - denoiser = std::make_shared(); - } + denoiser = std::make_shared(); + break; + } + case FLUX3_FLOW_PRED: { + LOG_INFO("running in Flux.3 FLOW mode"); + denoiser = std::make_shared(); + break; + } + case H3_AV_FLOW_PRED: { + LOG_INFO("running in MiniMax H3 AV FLOW mode"); + denoiser = std::make_shared(default_flow_shift, 3.f, get_latent_channel()); + break; + } + case MING_IMAGE_FLOW_PRED: { + LOG_INFO("running in Ming-Image FLOW mode"); + denoiser = std::make_shared(); break; } case SEFI_FLOW_PRED: { @@ -1737,8 +1635,8 @@ std::shared_ptr StableDiffusionGGML::load_lora_model(const ModelManag ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend; auto lora = std::make_shared(lora_log_id(lora_spec), backend_for(module), params_backend_for(module), - model_manager, lora_spec.file_id, version, mode, - backend_manager.params_backend_follows_runtime(module)); + model_manager, lora_spec.file_id, version, mode, + backend_manager.params_backend_follows_runtime(module)); LoraModel::filter_t lora_tensor_filter = module_filter; if (!lora_spec.tensor_name_prefix_filter.empty()) { lora_tensor_filter = [module_filter, prefix = lora_spec.tensor_name_prefix_filter](const std::string& tensor_name) { @@ -2152,11 +2050,11 @@ void StableDiffusionGGML::preview_image(int step, : sd::ops::slice(latents, 2, 0, channels) : latents; if (preview_mode == PREVIEW_PROJ) { - int patch_sz = 1; - const float(*latent_rgb_proj)[3] = nullptr; - float* latent_rgb_bias = nullptr; - const float* latent_alpha_proj = nullptr; - float latent_alpha_bias = 1.f; + int patch_sz = 1; + const float (*latent_rgb_proj)[3] = nullptr; + float* latent_rgb_bias = nullptr; + const float* latent_alpha_proj = nullptr; + float latent_alpha_bias = 1.f; if (channels == 128) { if (sd_version_uses_flux2_vae(version)) { @@ -2438,8 +2336,8 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptris_enabled(); }); - using QwenPrefixInputs = std::tuple*, const sd::Tensor*, - const std::vector>*>; + using QwenPrefixInputs = std::tuple*, const sd::Tensor*, + const std::vector>*>; std::vector qwen_prefix_inputs; RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; @@ -2712,6 +2610,10 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr StableDiffusionGGML::sample(const std::shared_ptr(&diffusion_params.extra); if (extra != nullptr) { auto key = std::make_tuple(diffusion_params.context, extra->image_slots, - diffusion_params.ref_image_params.pass_to_dit ? diffusion_params.ref_latents : nullptr); + diffusion_params.ref_image_params.pass_to_dit ? diffusion_params.ref_latents : nullptr); auto entry = std::find(qwen_prefix_inputs.begin(), qwen_prefix_inputs.end(), key); extra->prefix_id = static_cast(entry - qwen_prefix_inputs.begin()) + 1; if (entry == qwen_prefix_inputs.end()) { @@ -2999,63 +2901,15 @@ int StableDiffusionGGML::get_vae_scale_factor() { } int StableDiffusionGGML::get_diffusion_model_down_factor() { - int down_factor = 8; // unet - if (sd_version_is_dit(version)) { - if (sd_version_is_sensenova_u1(version)) { - down_factor = 32; - } else if (sd_version_is_z_image_l2p(version) || version == VERSION_IRIS) { - down_factor = 16; - } else if (version == VERSION_QWEN_IMAGE_2_1 || version == VERSION_MING_IMAGE || sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version) || sd_version_is_pixart(version)) { - down_factor = 2; - } else { - down_factor = 1; - } - } - return down_factor; + return find_model_definition(version)->geometry.diffusion_down_factor; } int StableDiffusionGGML::get_latent_channel() { - int latent_channel = 4; - if (sd_version_is_dit(version)) { - if (sd_version_is_ltxav(version)) { - latent_channel = 128; - } else if (sd_version_is_minimax_h3(version)) { - latent_channel = 24; - } else if (version == VERSION_QWEN_IMAGE_2_1) { - latent_channel = 64; - } else if (version == VERSION_WAN2_2_TI2V) { - latent_channel = 48; - } else if (sd_version_is_hunyuan_video(version)) { - latent_channel = 32; - } else if (version == VERSION_HIDREAM_O1) { - latent_channel = 3; - } else if (version == VERSION_CHROMA_RADIANCE) { - latent_channel = 3; - } else if (sd_version_is_minit2i(version)) { - latent_channel = 3; - } else if (sd_version_is_sensenova_u1(version)) { - latent_channel = 3; - } else if (sd_version_is_z_image_l2p(version)) { - latent_channel = 3; - } else if (sd_version_is_pid(version) || version == VERSION_IRIS) { - latent_channel = 3; - } else if (sd_version_is_sefi_image(version)) { - latent_channel = 144; - } else if (sd_version_uses_flux2_vae(version)) { - latent_channel = 128; - } else if (sd_version_is_mage_flow(version)) { - latent_channel = 128; - } else if (sd_version_is_pixart(version)) { - latent_channel = 4; - } else { - latent_channel = 16; - } - } - return latent_channel; + return find_model_definition(version)->geometry.latent_channels; } int StableDiffusionGGML::get_image_channels() const { - return version == VERSION_QWEN_IMAGE_LAYERED || version == VERSION_QWEN_IMAGE_2_1 || version == VERSION_MING_IMAGE ? 4 : 3; + return find_model_definition(version)->geometry.image_channels; } int StableDiffusionGGML::get_image_seq_len(int h, int w) { @@ -3207,40 +3061,15 @@ bool StableDiffusionGGML::is_flow_denoiser() { } std::string StableDiffusionGGML::get_default_ref_image_preset(SDVersion version) const { - if (version == VERSION_FIBO2) { - return "fibo2"; - } else if (sd_version_is_longcat(version)) { - return "longcat"; - } else if (sd_version_is_flux(version)) { - return "flux_kontext"; - } else if (sd_version_is_flux3(version)) { - return "flux2"; - } 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_mage_flow(version)) { - return "mage_flow"; - } else if (sd_version_is_z_image(version) || sd_version_is_boogu_image(version)) { - return "z_image_omni"; - } else if (sd_version_is_llada_image(version)) { - return "llada_image"; - } 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"; + return find_model_definition(version)->ref_image_preset; } RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_image_args) const { RefImageParams params; std::string preset_name = get_default_ref_image_preset(version); + // Keys whose value was rejected. They are already reported, so per-model "no effect" + // warnings must not repeat them. + std::set rejected_keys; for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) { if (key == "preset") { @@ -3253,6 +3082,7 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima valid_list += (valid_list.empty() ? "" : ", ") + name; } LOG_WARN("ignoring invalid reference image preset '%s'. Valid options: [%s]", value.c_str(), valid_list.c_str()); + rejected_keys.insert(key); } break; } @@ -3269,10 +3099,12 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima 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()); + rejected_keys.insert(key); } } 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()); + rejected_keys.insert(key); } } else if (key == "ref_index_mode") { if (value == "fixed") { @@ -3283,14 +3115,17 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima params.ref_index_mode = Rope::RefIndexMode::DECREASE; } else { LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + rejected_keys.insert(key); } } 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()); + rejected_keys.insert(key); } } 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()); + rejected_keys.insert(key); } } else if (key == "resize_vae_to_target") { if (!parse_strict_bool(value, params.resize_vae_to_target)) { @@ -3315,6 +3150,7 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima } 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()); + rejected_keys.insert(key); } } else if (key == "vlm_resize_mode") { if (value == "longest_side") { @@ -3325,14 +3161,31 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima params.vlm_resize_mode = RefImageResizeMode::NONE; } else { LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + rejected_keys.insert(key); } } 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()); + rejected_keys.insert(key); } } 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()); + rejected_keys.insert(key); + } + } else if (key == "ref_image_size") { + bool valid = true; + if (value == "match") { + params.ref_image_size = RefImageSizeMode::MATCH; + } else if (value == "max") { + params.ref_image_size = RefImageSizeMode::MAX; + } else { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + rejected_keys.insert(key); + valid = false; + } + if (valid && !sd_version_is_minimax_h3(version)) { + LOG_WARN("ref_image_size is only supported by MiniMax-H3 for now"); } } else if (key != "preset" && key != "vlm_size") { LOG_WARN("ignoring unknown reference image arg '%s'", key.c_str()); @@ -3343,6 +3196,7 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima 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()); + rejected_keys.insert(key); } else { LOG_INFO("vlm_size override: setting both min and max size to %ld", (long)vlm_size); params.vlm_min_size = vlm_size; @@ -3354,6 +3208,29 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima 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"); } + if (sd_version_is_minimax_h3(version)) { + // H3 consumes references as VAE latents in its own packed sequence rather than + // through the generic conditioner/DiT reference path, so ref_image_size is the + // only arg it acts on. Everything else is reported as ineffective rather than + // enumerated, so newly added RefImageParams members cannot silently do nothing. + static const std::set h3_effective_keys = {"ref_image_size"}; + // Keys that already have their own, more specific warning above. + static const std::set h3_separately_reported_keys = { + "resize_before_vae", + "force_ref_timestep_zero", + }; + for (const auto& arg : parse_key_value_args(ref_image_args, "reference image args")) { + if (h3_effective_keys.count(arg.first) == 0 && + h3_separately_reported_keys.count(arg.first) == 0 && + rejected_keys.count(arg.first) == 0) { + LOG_WARN("reference image arg '%s' has no effect for MiniMax-H3", arg.first.c_str()); + } + } + + if (!params.resize_before_vae) { + LOG_WARN("MiniMax-H3 requires 32-pixel-aligned reference images and cannot skip resizing before VAE encoding; ignoring resize_before_vae=false"); + } + } return params; } diff --git a/src/pipeline/diffusion_engine.h b/src/pipeline/diffusion_engine.h index c5509d9a9..c7295d8f7 100644 --- a/src/pipeline/diffusion_engine.h +++ b/src/pipeline/diffusion_engine.h @@ -34,8 +34,6 @@ namespace Wav2Vec2 { class Wav2Vec2ModelRunner; } -extern const char* model_version_to_str[]; - static inline bool sd_version_supports_ref_latent_img_cfg(SDVersion version) { return version == VERSION_FLUX || sd_version_is_flux2(version) || @@ -125,6 +123,7 @@ class StableDiffusionGGML { sd_ctx_params_t params{}; std::list strings; std::vector embeddings; + std::map models; ModelLoader::FileId control_net_file = 0; bool use_tae = false; bool use_audio_vae = false; @@ -166,13 +165,24 @@ class StableDiffusionGGML { use_audio_vae = other.use_audio_vae; photomaker_source_available = other.photomaker_source_available; animatediff_loaded = other.animatediff_loaded; + models = other.models; } ModelConfig& operator=(const ModelConfig&) = delete; + void record_model(const std::string& name, const std::string& path) { + if (path.empty()) { + models.erase(name); + return; + } + const auto separator = path.find_last_of("/\\"); + models[name] = separator == std::string::npos ? path : path.substr(separator + 1); + } + void set_control_net(ModelLoader::FileId id, const std::string& path) { control_net_file = id; strings.push_back(path); params.control_net_path = strings.back().c_str(); + record_model("control_net", path); } }; diff --git a/src/pipeline/generation.h b/src/pipeline/generation.h index 34aa000aa..0bc391315 100644 --- a/src/pipeline/generation.h +++ b/src/pipeline/generation.h @@ -2,6 +2,7 @@ #define __SD_PIPELINE_GENERATION_H__ #include "conditioning/conditioner.hpp" +#include "json.hpp" #include "stable-diffusion.h" class StableDiffusionGGML; @@ -54,14 +55,16 @@ namespace sd::pipeline { bool generate_image(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params, sd_image_t** images_out, - int* num_images_out); + int* num_images_out, + std::vector* infos_out = nullptr); bool generate_video(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params, sd_image_t** frames_out, int* num_frames_out, sd_audio_t** audio_out, - int* fps_out); + int* fps_out, + nlohmann::json* info_out = nullptr); sd::Tensor upscale_ltx_spatial_video_latent(StableDiffusionGGML* sd, const char* model_path, diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index c683b53ff..7d264e22c 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -1,4 +1,5 @@ #include "generation.h" +#include "generation_info.h" #include #include @@ -117,6 +118,32 @@ namespace sd::pipeline { channels); } + static sd::Tensor encode_z_image_control_latent(StableDiffusionGGML* sd, const sd::Tensor& control_rgb) { + sd::Tensor vae_output = sd->first_stage_model->encode(sd->n_threads, + control_rgb, + sd->vae_tiling_params, + sd->circular_x, + sd->circular_y); + if (vae_output.empty()) { + LOG_ERROR("failed to encode control image"); + return {}; + } + // TAE encode already returns the latent; KL/Flux VAE returns a mean/logvar pair. + sd::Tensor ctrl_latent = sd->first_stage_model->get_desc() == "taesd" + ? vae_output + : sd::ops::chunk(vae_output, 2, 2)[0]; + ctrl_latent = sd->first_stage_model->vae_to_diffusion_latents(ctrl_latent); + if (ctrl_latent.empty() || ctrl_latent.dim() != 4 || ctrl_latent.shape()[2] > 33) { + LOG_ERROR("failed to encode control image"); + return {}; + } + std::vector padded_shape = ctrl_latent.shape(); + padded_shape[2] = 33; + sd::Tensor padded(padded_shape); + sd::ops::slice_assign(&padded, 2, 0, ctrl_latent.shape()[2], ctrl_latent); + return padded; + } + static std::optional prepare_image_generation_latents(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params, GenerationRequest* request, @@ -250,10 +277,22 @@ namespace sd::pipeline { } if (!control_image_tensor.empty()) { - control_latent = sd->encode_first_stage(control_image_tensor); - if (control_latent.empty()) { - LOG_ERROR("failed to encode control image"); - return std::nullopt; + const bool z_image_control = sd_version_is_z_image(sd->version) && + sd->model_manager->registered_params_size({ModelComponent::ControlNet}) > 0; + if (z_image_control) { + control_image_tensor = encode_z_image_control_latent(sd, control_image_tensor); + if (control_image_tensor.empty()) { + return std::nullopt; + } + } else if (sd_version_is_z_image(sd->version)) { + LOG_WARN("control image provided but the z-image model has no control net; ignoring"); + control_image_tensor = {}; + } else { + control_latent = sd->encode_first_stage(control_image_tensor); + if (control_latent.empty()) { + LOG_ERROR("failed to encode control image"); + return std::nullopt; + } } } @@ -274,9 +313,13 @@ namespace sd::pipeline { LOG_INFO("EDIT mode"); } + bool skip_vae_encoding = (sd->version == VERSION_HIDREAM_O1) || !( + ref_image_params.pass_to_dit || ref_image_params.vlm_use_vae_input || + sd_version_is_unet_edit(sd->version) || sd_version_is_pid(sd->version)); + std::vector> ref_latents; for (size_t i = 0; i < ref_images.size(); i++) { - if (sd->version == VERSION_HIDREAM_O1) { + if (skip_vae_encoding) { continue; } sd::Tensor ref_latent; @@ -630,7 +673,7 @@ namespace sd::pipeline { int64_t t4 = ggml_time_ms(); LOG_INFO("decode_first_stage completed, taking %.2fs", (t4 - t0) * 1.0f / 1000); - if (decoded_images.empty()) { + if (decoded_images.empty() || cancelled) { LOG_ERROR(cancelled ? "cancelled before any latent images were decoded" : "no decoded images"); return nullptr; } @@ -791,7 +834,10 @@ namespace sd::pipeline { bool generate_image(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params, sd_image_t** images_out, - int* num_images_out) { + int* num_images_out, + std::vector* infos_out) { + if (infos_out) + infos_out->clear(); if (images_out != nullptr) { *images_out = nullptr; } @@ -804,7 +850,7 @@ namespace sd::pipeline { if (!sd_version_supports_image_generation(sd->version)) { LOG_ERROR("%s cannot be run with generate_image(); use generate_video() or --mode vid_gen in the CLI", - model_version_to_str[sd->version]); + sd_version_name(sd->version)); return false; } @@ -813,6 +859,7 @@ namespace sd::pipeline { int64_t t0 = ggml_time_ms(); sd->vae_tiling_params = sd_img_gen_params->vae_tiling_params; GenerationRequest request(sd, sd_img_gen_params); + const auto info = sd::generation_info(request, *sd_img_gen_params); sd::ImagePreprocessor preprocessing(sd_img_gen_params->image_preprocess.rules); sd_img_gen_params_t processed_params = *sd_img_gen_params; if (!preprocessing.prepare_inputs(processed_params, request.width, request.height)) @@ -853,7 +900,6 @@ namespace sd::pipeline { return false; } ImageGenerationEmbeds embeds = std::move(*embeds_opt); - std::vector> final_latents; int64_t denoise_start = ggml_time_ms(); for (int b = 0; b < request.batch_count; b++) { @@ -870,7 +916,7 @@ namespace sd::pipeline { } int64_t sampling_start = ggml_time_ms(); - int64_t cur_seed = request.seed + b; + int64_t cur_seed = request.seed_for_batch(b); LOG_INFO("generating image: %i/%i - seed %" PRId64, b + 1, request.batch_count, cur_seed); sd->rng->manual_seed(cur_seed); @@ -939,10 +985,10 @@ namespace sd::pipeline { } LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path); hires_upscaler = std::make_unique(sd->n_threads, - false, - request.hires.upscale_tile_size, - sd->backend_spec, - sd->params_backend_spec); + false, + request.hires.upscale_tile_size, + sd->backend_spec, + sd->params_backend_spec); const size_t max_graph_vram_bytes = sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER); hires_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes); if (!hires_upscaler->load_from_file(request.hires.model_path, @@ -956,7 +1002,7 @@ namespace sd::pipeline { std::vector hires_sigma_sched = make_hires_sigma_schedule(sd, request.hires, - sd_img_gen_params->sample_params, + request.sample_params, plan.sample_method, plan.sample_steps, sd->get_image_seq_len(request.hires.target_height, request.hires.target_width), @@ -969,12 +1015,35 @@ namespace sd::pipeline { std::vector> hires_final_latents; int64_t hires_denoise_start = ggml_time_ms(); + sd::Tensor hires_control_image; + if (!latents.control_image.empty() && + sd_version_is_z_image(sd->version) && + sd->model_manager->registered_params_size({ModelComponent::ControlNet}) > 0) { + const int64_t hires_latent_w = request.hires.target_width / request.vae_scale_factor; + const int64_t hires_latent_h = request.hires.target_height / request.vae_scale_factor; + if (latents.control_image.shape()[0] != hires_latent_w || + latents.control_image.shape()[1] != hires_latent_h) { + LOG_INFO("hires fix: re-encoding control latent for %dx%d", request.hires.target_width, request.hires.target_height); + hires_control_image = encode_z_image_control_latent(sd, + sd_image_to_tensor(sd_img_gen_params->control_image, + request.hires.target_width, + request.hires.target_height)); + if (hires_control_image.empty()) { + std::vector ctrl_shape = latents.control_image.shape(); + ctrl_shape[0] = hires_latent_w; + ctrl_shape[1] = hires_latent_h; + hires_control_image = sd::ops::interpolate(latents.control_image, ctrl_shape, sd::ops::InterpolateMode::Bicubic); + } + } else { + hires_control_image = latents.control_image; + } + } for (int b = 0; b < (int)final_latents.size(); b++) { if (sd->get_cancel_flag() == SD_CANCEL_ALL) { LOG_ERROR("cancelling generation during hires fix"); return false; } - int64_t cur_seed = request.seed + b; + int64_t cur_seed = request.seed_for_batch(b); sd->rng->manual_seed(cur_seed); sd->sampler_rng->manual_seed(cur_seed); @@ -1006,7 +1075,7 @@ namespace sd::pipeline { embeds.cond, embeds.uncond, embeds.img_uncond, - latents.control_image, + !hires_control_image.empty() ? hires_control_image : latents.control_image, request.control_strength, request.guidance, plan.eta, @@ -1052,6 +1121,19 @@ namespace sd::pipeline { return false; } + if (infos_out) { + try { + infos_out->assign(num_images, info); + // Layers and frames decoded from the same latent share its seed. + const int images_per_latent = num_images / static_cast(final_latents.size()); + for (int i = 0; i < num_images; ++i) { + (*infos_out)[i]["seed"] = std::to_string(request.seed_for_batch(i / images_per_latent)); + } + } catch (...) { + free_sd_images(result, num_images); + throw; + } + } sd->lora_stat(); int64_t t1 = ggml_time_ms(); diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index ca4d1998b..8b85f9b46 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -5,44 +5,13 @@ #include #include "conditioning/conditioner.hpp" -#include "conditioning/fibo2.hpp" #include "core/ggml_extend_backend.h" #include "core/util.h" #include "extensions/generation_extension.h" #include "model/adapter/ip_adapter.hpp" #include "model/audio/wav2vec2.hpp" -#include "model/diffusion/anima.hpp" -#include "model/diffusion/boogu.hpp" #include "model/diffusion/control.hpp" -#include "model/diffusion/ernie_image.hpp" -#include "model/diffusion/fibo2.hpp" -#include "model/diffusion/flux.hpp" -#include "model/diffusion/flux3.hpp" -#include "model/diffusion/hidream_o1.hpp" -#include "model/diffusion/hunyuan.hpp" -#include "model/diffusion/ideogram4.hpp" -#include "model/diffusion/iris.hpp" -#include "model/diffusion/krea2.hpp" -#include "model/diffusion/lens.hpp" -#include "model/diffusion/lingbot_video.hpp" -#include "model/diffusion/llada_image.hpp" -#include "model/diffusion/ltxv.hpp" -#include "model/diffusion/mage_flow.hpp" -#include "model/diffusion/ming_image.hpp" -#include "model/diffusion/minimax_h3.hpp" -#include "model/diffusion/minit2i.hpp" -#include "model/diffusion/mmdit.hpp" #include "model/diffusion/model.hpp" -#include "model/diffusion/moss_image.hpp" -#include "model/diffusion/pid.hpp" -#include "model/diffusion/pixart.hpp" -#include "model/diffusion/qwen_image.hpp" -#include "model/diffusion/qwen_image_2_1.hpp" -#include "model/diffusion/sensenova_u1.h" -#include "model/diffusion/unet.hpp" -#include "model/diffusion/wan.hpp" -#include "model/diffusion/z_image.hpp" -#include "model/diffusion/z_image_l2p.hpp" #include "model/vae/auto_encoder_kl.hpp" #include "model/vae/hunyuan_vae.hpp" #include "model/vae/ltx_audio_vae.hpp" @@ -53,10 +22,11 @@ #include "model/vae/tae.hpp" #include "model/vae/vae.hpp" #include "model/vae/wan_vae.hpp" +#include "model_definition.h" namespace sd::model_builders { - static bool ensure_backend_pair(SDBackendManager& backends, SDBackendModule module) { + bool ensure_backend_pair(SDBackendManager& backends, SDBackendModule module) { if (backends.runtime_backend(module) == nullptr) { LOG_ERROR("failed to initialize %s backend", sd_backend_module_name(module)); return false; @@ -96,407 +66,10 @@ namespace sd::model_builders { return false; } - if (version == VERSION_FIBO2) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, weight_manager, tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, "model.diffusion_model", weight_manager); - } else if (sd_version_is_sd3(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_pid(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model.net", - weight_manager); - } else if (sd_version_is_ideogram4(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_krea2(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_flux(version)) { - bool is_chroma = false; - for (auto pair : tensor_storage_map) { - if (pair.first.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) { - is_chroma = true; - break; - } - } - if (is_chroma) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - false, - 1, - false, - weight_manager, - sd_ctx_params->model_args); - } else if (version == VERSION_OVIS_IMAGE) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - } else { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - weight_manager, - tokenizers); - } - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager, - sd_ctx_params->model_args); - } else if (sd_version_is_flux3(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, version, "", false, weight_manager, tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, "model.diffusion_model", weight_manager); - } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) { - bool is_chroma = false; - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager, - sd_ctx_params->model_args); - } else if (sd_version_is_ltxav(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - "text_encoders.llm", - "text_embedding_projection", - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_minimax_h3(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_hunyuan_video(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager); - } else if (sd_version_is_wan(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - true, - 0, - true, - weight_manager); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager); - if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) { - result.high_noise_diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.high_noise_diffusion_model", - version, - weight_manager); - } - if (result.diffusion->get_desc() == "Wan2.1-I2V-14B" || - result.diffusion->get_desc() == "Wan2.1-FLF2V-14B" || - result.diffusion->get_desc() == "Wan2.1-I2V-1.3B") { - if (!ensure_backend_pair(ctx.backends, SDBackendModule::CLIP_VISION)) { - return false; - } - result.clip_vision = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::CLIP_VISION), - tensor_storage_map, - weight_manager); - } - if (version == VERSION_WAN2_2_S2V && - tensor_storage_map.count("wav2vec2.encoder.layer_norm.bias") > 0) { - if (!ensure_backend_pair(ctx.backends, SDBackendModule::AUDIO_ENCODER)) { - return false; - } - result.audio_encoder = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::AUDIO_ENCODER), - tensor_storage_map, - "wav2vec2.", - weight_manager); - } - } else if (sd_version_is_lingbot_video(version)) { - bool enable_vision = false; - for (const auto& [name, _] : tensor_storage_map) { - if (starts_with(name, "text_encoders.llm.visual.")) { - enable_vision = true; - break; - } - } - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - enable_vision, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager, - sd_ctx_params->model_args); - } else if (sd_version_is_qwen_image(version)) { - bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED; - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - enable_vision, - weight_manager, - tokenizers); - if (version == VERSION_QWEN_IMAGE_2_1) { - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager, - sd_ctx_params->model_args); - } else { - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager, - sd_ctx_params->model_args); - } - } else if (version == VERSION_PIXART) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - true, - 0, - false, - weight_manager, - sd_ctx_params->model_args); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager, - sd_ctx_params->model_args); - } else if (sd_version_is_mage_flow(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_longcat(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager, - sd_ctx_params->model_args); - } else if (version == VERSION_HIDREAM_O1) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model", - weight_manager); - } else if (sd_version_is_minit2i(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - weight_manager); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model.model.net", - weight_manager); - } else if (sd_version_is_sensenova_u1(version)) { - result.conditioner = std::make_shared(); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "", - weight_manager); - } else if (sd_version_is_anima(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (version == VERSION_MING_IMAGE) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, weight_manager, tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, "model.diffusion_model", weight_manager); - } else if (sd_version_is_z_image(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager); - } else if (sd_version_is_z_image_l2p(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_llada_image(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - "text_encoders.llm", - "queryformer", - "text_projection", - "sigvq", - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_boogu_image(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager); - } else if (sd_version_is_ernie_image(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_lens(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (version == VERSION_IRIS) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, version, "", false, weight_manager, tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, "model.diffusion_model", weight_manager); - } else if (version == VERSION_MOSS_IMAGE) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, version, "", false, weight_manager, tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, "model.diffusion_model", weight_manager); - } else { // SD1.x SD2.x SDXL - std::map embbeding_map; - for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) { - embbeding_map.emplace(SAFE_STR(sd_ctx_params->embeddings[i].name), SAFE_STR(sd_ctx_params->embeddings[i].path)); - } - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - embbeding_map, - version, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager); + const auto* definition = find_model_definition(version); + if (!definition || !definition->build || !definition->build(ctx, tokenizers, result)) { + LOG_ERROR("failed to build model runners for %s", sd_version_name(version)); + return false; } if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0 && result.clip_vision == nullptr) { @@ -548,7 +121,7 @@ namespace sd::model_builders { } auto create_tae = [&](bool decode_only) -> std::shared_ptr { - if ((sd_version_uses_wan_vae(version) && version != VERSION_QWEN_IMAGE_2_1) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) { + if (find_model_traits(version)->video_tae) { return std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), tensor_storage_map, "decoder", @@ -573,36 +146,36 @@ namespace sd::model_builders { vae_format = SD_VAE_FORMAT_AUTO; } SDVersion vae_version = version; - if (sd_version_is_pid(version) && vae_format != SD_VAE_FORMAT_AUTO) { + if (find_model_traits(version)->vae_format_override && vae_format != SD_VAE_FORMAT_AUTO) { vae_version = sd_vae_format_to_version(vae_format, vae_version); } auto create_vae = [&]() -> std::shared_ptr { - if (sd_version_is_ltxav(version)) { + if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::LTXAV) { return std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), tensor_storage_map, "first_stage_model", false, version, weight_manager); - } else if (sd_version_is_minimax_h3(version)) { + } else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::MINIMAX_H3) { return std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), tensor_storage_map, "first_stage_model", weight_manager); - } else if (sd_version_is_mage_flow(vae_version)) { + } else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::MAGE) { return std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), tensor_storage_map, "first_stage_model", weight_manager); - } else if (sd_version_uses_hunyuan_video_vae(vae_version)) { + } else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::HUNYUAN) { return std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), tensor_storage_map, "first_stage_model", false, vae_version, weight_manager); - } else if (sd_version_uses_wan_vae(vae_version)) { + } else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::WAN) { return std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), tensor_storage_map, "first_stage_model", @@ -610,49 +183,22 @@ namespace sd::model_builders { vae_version, weight_manager); } else { - auto model = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), - tensor_storage_map, - "first_stage_model", - false, - false, - vae_version, - weight_manager); - if (sd_version_is_pixart(version)) { - // Alpha-512 and Sigma share tensor layouts; Alpha-512 needs an explicit scale override. - if (tensor_storage_map.count("model.diffusion_model.csize_embedder.mlp.0.weight") != 0) { - model->scale_factor = 0.18215f; - } - for (const auto& [key, value] : parse_key_value_args(sd_ctx_params->model_args, "model arg")) { - if (key == "pixart_vae_scale_factor") { - float parsed = 0.f; - if (parse_strict_float(value, parsed) && std::isfinite(parsed) && parsed > 0.f) { - model->scale_factor = parsed; - } else { - LOG_WARN("ignoring invalid PixArt model arg '%s=%s'", key.c_str(), value.c_str()); - } - } - } - LOG_VERBOSE("pixart: VAE scale factor = %.5f", model->scale_factor); - } - if (sd_version_is_sdxl(version) && - (strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale || options.external_vae_is_invalid)) { - float vae_conv_2d_scale = 1.f / 32.f; - LOG_WARN( - "No valid VAE specified with --vae or --force-sdxl-vae-conv-scale flag set, " - "using Conv2D scale %.3f", - vae_conv_2d_scale); - model->set_conv2d_scale(vae_conv_2d_scale); + auto model = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), + tensor_storage_map, + "first_stage_model", + false, + false, + vae_version, + weight_manager); + const auto* definition = find_model_definition(version); + if (definition && definition->configure_vae) { + definition->configure_vae(ctx, options, *model); } return model; } }; - if (version == VERSION_CHROMA_RADIANCE || - version == VERSION_HIDREAM_O1 || - sd_version_is_minit2i(version) || - sd_version_is_sensenova_u1(version) || - sd_version_is_z_image_l2p(version) || - version == VERSION_IRIS) { + if (find_model_traits(version)->fake_vae) { LOG_INFO("using FakeVAE"); result.vae = std::make_shared(version, ctx.backends.runtime_backend(SDBackendModule::VAE), @@ -670,7 +216,7 @@ namespace sd::model_builders { } if (options.use_audio_vae) { - if (sd_version_is_minimax_h3(version)) { + if (find_model_traits(version)->vae_kind == ModelVAEKind::MINIMAX_H3) { result.audio = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), tensor_storage_map, "", @@ -701,6 +247,9 @@ namespace sd::model_builders { const auto& tensor_storage_map = ctx.tensor_storage_map; const auto version = ctx.version; const auto& weight_manager = ctx.weight_manager; + if (sd_version_is_z_image(version)) { + return true; + } if (!ensure_backend_pair(ctx.backends, SDBackendModule::CONTROL_NET)) { return false; } diff --git a/src/pipeline/model_builders.h b/src/pipeline/model_builders.h index 5edc75680..bb51ccc14 100644 --- a/src/pipeline/model_builders.h +++ b/src/pipeline/model_builders.h @@ -4,6 +4,7 @@ #include #include +#include "core/ggml_extend_backend.h" #include "model.h" #include "stable-diffusion.h" @@ -56,6 +57,8 @@ namespace sd::model_builders { std::shared_ptr audio; }; + bool ensure_backend_pair(SDBackendManager& backends, SDBackendModule module); + bool build_core_runners(const Context& ctx, CoreRunners& runners); bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners); bool build_control_net_runner(const Context& ctx, std::shared_ptr& runner); diff --git a/src/pipeline/request.cpp b/src/pipeline/request.cpp index 1df3e9f19..2647efb8e 100644 --- a/src/pipeline/request.cpp +++ b/src/pipeline/request.cpp @@ -1,4 +1,5 @@ #include "request.h" +#include "conditioning/conditioner.hpp" #include #include @@ -6,6 +7,9 @@ #include #include "diffusion_engine.h" +#include "generation.h" +#include "model/vae/vae.hpp" +#include "model_definition.h" #include "runtime/denoiser.hpp" namespace sd::pipeline { @@ -43,15 +47,8 @@ namespace sd::pipeline { } enum sample_method_t default_sample_method(const StableDiffusionGGML* sd) { - if (sd != nullptr) { - if (sd_version_is_pid(sd->version)) { - return LCM_SAMPLE_METHOD; - } - if (sd_version_is_dit(sd->version)) { - return EULER_SAMPLE_METHOD; - } - } - return EULER_A_SAMPLE_METHOD; + const auto* definition = sd ? find_model_definition(sd->version) : nullptr; + return definition ? definition->sampling.default_sample_method : EULER_A_SAMPLE_METHOD; } enum scheduler_t default_scheduler(const StableDiffusionGGML* sd, enum sample_method_t sample_method) { @@ -65,21 +62,14 @@ namespace sd::pipeline { return LCM_SCHEDULER; } else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) { return SIMPLE_SCHEDULER; - } else if (sd != nullptr && (sd_version_is_flux(sd->version) || sd->version == VERSION_QWEN_IMAGE_2_1 || sd->version == VERSION_FIBO2)) { - return FLUX_SCHEDULER; - } else if (sd != nullptr && sd_version_is_flux2(sd->version)) { - return FLUX2_SCHEDULER; - } else if (sd != nullptr && sd_version_is_ltxav(sd->version)) { - return LTX2_SCHEDULER; - } else if (sd != nullptr && sd_version_is_ideogram4(sd->version)) { - return LOGIT_NORMAL_SCHEDULER; - } else if (sd != nullptr && sd_version_is_llada_image(sd->version)) { - return LLADA_IMAGE_SCHEDULER; + } + if (const auto* definition = sd ? find_model_definition(sd->version) : nullptr) { + return definition->sampling.default_scheduler; } return DISCRETE_SCHEDULER; } - static int64_t resolve_seed(int64_t seed) { + int64_t resolve_seed(int64_t seed) { if (seed >= 0) { return seed; } @@ -140,6 +130,7 @@ namespace sd::pipeline { seed = sd_img_gen_params->seed; batch_count = sd_img_gen_params->batch_count; qwen_image_layers = std::max(0, sd_img_gen_params->qwen_image_layers); + is_layered_image = sd->version == VERSION_QWEN_IMAGE_LAYERED; clip_skip = sd_img_gen_params->clip_skip; shifted_timestep = sd_img_gen_params->sample_params.shifted_timestep; strength = sd_img_gen_params->strength; @@ -150,8 +141,18 @@ namespace sd::pipeline { pm_params = sd_img_gen_params->pm_params; pulid_params = sd_img_gen_params->pulid_params; hires = sd_img_gen_params->hires; + vae_tiling_params = sd_img_gen_params->vae_tiling_params; cache_params = &sd_img_gen_params->cache; + sample_params = sd_img_gen_params->sample_params; resolve(sd); + if (!sd_img_gen_params->control_image.data) + models.erase("control_net"); + if (!sd_img_gen_params->ip_adapter_image.data) + models.erase("ip_adapter"); + if (!pm_params.id_images || pm_params.id_images_count <= 0) + models.erase("photo_maker"); + if (!pulid_params.id_embedding_path || !*pulid_params.id_embedding_path) + models.erase("pulid_weights"); } GenerationRequest::GenerationRequest(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params) { @@ -179,12 +180,25 @@ namespace sd::pipeline { guidance = sd_vid_gen_params->sample_params.guidance; high_noise_guidance = sd_vid_gen_params->high_noise_sample_params.guidance; hires = sd_vid_gen_params->hires; + vae_tiling_params = sd_vid_gen_params->vae_tiling_params; + sample_params = sd_vid_gen_params->sample_params; + high_noise_sample_params = sd_vid_gen_params->high_noise_sample_params; + use_high_noise = sd->high_noise_diffusion_model != nullptr; + use_vace = sd->diffusion_model && (sd->diffusion_model->get_desc() == "Wan2.1-VACE-1.3B" || + sd->diffusion_model->get_desc() == "Wan2.x-VACE-14B"); + shifted_timestep = sample_params.shifted_timestep; + if (sd->config_ && sd->config_->animatediff_loaded && sd_version_supports_animatediff(sd->version)) + frames = std::min(requested_frames, 32); resolve(sd); + models.erase("control_net"); + models.erase("ip_adapter"); + models.erase("photo_maker"); + models.erase("pulid_weights"); if (frames != requested_frames) { LOG_WARN("align video frames from %d to %d for %s", requested_frames, frames, - model_version_to_str[sd->version]); + sd_version_name(sd->version)); } } @@ -214,6 +228,10 @@ namespace sd::pipeline { spatial_multiple); } + int64_t GenerationRequest::seed_for_batch(int index) const { + return static_cast((static_cast(seed) + index) & INT64_MAX); + } + void GenerationRequest::resolve_hires() { if (!hires.enabled) { return; @@ -320,9 +338,35 @@ namespace sd::pipeline { } void GenerationRequest::resolve(StableDiffusionGGML* sd) { + if (sd->config_) { + models = sd->config_->models; + if (!use_high_noise) + models.erase("high_noise_diffusion_model"); + if (!sd->clip_vision) + models.erase("clip_vision"); + if (!sd->audio_vae_model || frames <= 0) + models.erase("audio_vae"); + if (!sd->audio_encoder || frames <= 0) + models.erase("audio_encoder"); + if (!sd->config_->animatediff_loaded || !sd_version_supports_animatediff(sd->version) || frames <= 0) + models.erase("motion_module"); + if (sd->uses_tae() && !sd->tae_preview_only()) + models.erase("vae"); + else + models.erase("taesd"); + } + if (sd_version_is_minit2i(sd->version) && (width != 512 || height != 512)) { + width = 512; + height = 512; + LOG_WARN("align to fixed size 512x512 for MiniT2I"); + } align_generation_request_size(); resolve_hires(); + if (sd->first_stage_model) + vae_tiling_params = sd->first_stage_model->resolve_tiling_params(vae_tiling_params); seed = resolve_seed(seed); + if (sd->cond_stage_model) + clip_skip = sd->cond_stage_model->resolve_clip_skip(clip_skip); resolve_guidance(sd, &guidance, &use_uncond, &use_img_uncond, has_ref_images); if (sd->high_noise_diffusion_model) { @@ -338,33 +382,44 @@ namespace sd::pipeline { LOG_WARN("timestep shifting is only supported for SDXL models!"); shifted_timestep = 0; } + sample_params.guidance = guidance; + sample_params.shifted_timestep = shifted_timestep; + sample_params.sample_method = resolve_sample_method(sd, sample_params.sample_method); + sample_params.scheduler = resolve_scheduler(sd, sample_params.scheduler, sample_params.sample_method); + sample_params.eta = resolve_eta(sd, sample_params.eta, sample_params.sample_method); + if (frames > 0) { + high_noise_sample_params.guidance = high_noise_guidance; + high_noise_sample_params.sample_method = resolve_sample_method(sd, high_noise_sample_params.sample_method); + high_noise_sample_params.scheduler = resolve_scheduler(sd, high_noise_sample_params.scheduler, high_noise_sample_params.sample_method); + high_noise_sample_params.eta = resolve_eta(sd, high_noise_sample_params.eta, high_noise_sample_params.sample_method); + } } SamplePlan::SamplePlan(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params, const GenerationRequest& request) { - sample_method = sd_img_gen_params->sample_params.sample_method; - extra_sample_args = sd_img_gen_params->sample_params.extra_sample_args; - eta = sd_img_gen_params->sample_params.eta; - sample_steps = sd_img_gen_params->sample_params.sample_steps; - resolve(sd, &request, &sd_img_gen_params->sample_params); + sample_method = request.sample_params.sample_method; + extra_sample_args = request.sample_params.extra_sample_args; + eta = request.sample_params.eta; + sample_steps = request.sample_params.sample_steps; + resolve(sd, &request, &request.sample_params); } SamplePlan::SamplePlan(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params, const GenerationRequest& request) { - sample_method = sd_vid_gen_params->sample_params.sample_method; - extra_sample_args = sd_vid_gen_params->sample_params.extra_sample_args; - eta = sd_vid_gen_params->sample_params.eta; - sample_steps = sd_vid_gen_params->sample_params.sample_steps; + sample_method = request.sample_params.sample_method; + extra_sample_args = request.sample_params.extra_sample_args; + eta = request.sample_params.eta; + sample_steps = request.sample_params.sample_steps; if (sd->high_noise_diffusion_model) { - high_noise_sample_steps = sd_vid_gen_params->high_noise_sample_params.sample_steps; - high_noise_sample_method = sd_vid_gen_params->high_noise_sample_params.sample_method; - high_noise_extra_sample_args = sd_vid_gen_params->high_noise_sample_params.extra_sample_args; - high_noise_eta = sd_vid_gen_params->high_noise_sample_params.eta; + high_noise_sample_steps = request.high_noise_sample_params.sample_steps; + high_noise_sample_method = request.high_noise_sample_params.sample_method; + high_noise_extra_sample_args = request.high_noise_sample_params.extra_sample_args; + high_noise_eta = request.high_noise_sample_params.eta; } moe_boundary = sd_vid_gen_params->moe_boundary; - resolve(sd, &request, &sd_vid_gen_params->sample_params); + resolve(sd, &request, &request.sample_params); } void SamplePlan::resolve(StableDiffusionGGML* sd, @@ -390,10 +445,10 @@ namespace sd::pipeline { } } } else { - scheduler_t scheduler = resolve_scheduler(sd, - sample_params->scheduler, - sample_method); - int sample_seq_len = sd->get_image_seq_len(request->height, request->width); + scheduler = resolve_scheduler(sd, + sample_params->scheduler, + sample_method); + int sample_seq_len = sd->get_image_seq_len(request->height, request->width); if (sd_version_is_ltxav(sd->version) && request->frames > 0) { int latent_frames = ((request->frames - 1) / 8) + 1; sample_seq_len *= latent_frames; diff --git a/src/pipeline/request.h b/src/pipeline/request.h index fac4aba55..1cf1cee6a 100644 --- a/src/pipeline/request.h +++ b/src/pipeline/request.h @@ -1,6 +1,7 @@ #ifndef __SD_PIPELINE_REQUEST_H__ #define __SD_PIPELINE_REQUEST_H__ +#include #include #include @@ -16,6 +17,8 @@ namespace sd::pipeline { enum scheduler_t default_scheduler(const StableDiffusionGGML* sd, enum sample_method_t sample_method); + int64_t resolve_seed(int64_t seed); + float resolve_eta(StableDiffusionGGML* sd, float eta, enum sample_method_t sample_method); @@ -23,38 +26,47 @@ namespace sd::pipeline { struct GenerationRequest { std::string prompt; std::string negative_prompt; - int width = -1; - int height = -1; - int clip_skip = -1; - int vae_scale_factor = -1; - int diffusion_model_down_factor = -1; - int64_t seed = -1; - bool use_uncond = false; - bool use_img_uncond = false; - bool use_high_noise_uncond = false; - bool use_high_noise_img_uncond = false; - bool has_ref_images = false; - const sd_cache_params_t* cache_params = nullptr; - int batch_count = 1; - int qwen_image_layers = 3; - int shifted_timestep = 0; - float strength = 1.f; - float control_strength = 0.f; - float eta = 0.f; - sd_guidance_params_t guidance = {}; - sd_guidance_params_t high_noise_guidance = {}; - sd_pm_params_t pm_params = {}; - sd_pulid_params_t pulid_params = {}; - sd_hires_params_t hires = {}; - int frames = -1; - int requested_frames = -1; - int fps = 16; - float vace_strength = 1.f; + std::map models; + int width = -1; + int height = -1; + int clip_skip = -1; + int vae_scale_factor = -1; + int diffusion_model_down_factor = -1; + int64_t seed = -1; + bool use_uncond = false; + bool use_img_uncond = false; + bool use_high_noise_uncond = false; + bool use_high_noise_img_uncond = false; + bool has_ref_images = false; + bool is_layered_image = false; + bool use_high_noise = false; + bool use_vace = false; + const sd_cache_params_t* cache_params = nullptr; + int batch_count = 1; + int qwen_image_layers = 3; + int shifted_timestep = 0; + float strength = 1.f; + float control_strength = 0.f; + float eta = 0.f; + sd_sample_params_t sample_params = {}; + sd_sample_params_t high_noise_sample_params = {}; + sd_guidance_params_t guidance = {}; + sd_guidance_params_t high_noise_guidance = {}; + sd_pm_params_t pm_params = {}; + sd_pulid_params_t pulid_params = {}; + sd_hires_params_t hires = {}; + sd_tiling_params_t vae_tiling_params = {}; + int frames = -1; + int requested_frames = -1; + int fps = 16; + float vace_strength = 1.f; GenerationRequest(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params); GenerationRequest(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params); + int64_t seed_for_batch(int index) const; + void align_generation_request_size(); void align_image_size(int* target_width, int* target_height, const char* label); @@ -81,6 +93,7 @@ namespace sd::pipeline { int sample_steps = 0; int high_noise_sample_steps = 0; int total_steps = 0; + scheduler_t scheduler = SCHEDULER_COUNT; float moe_boundary = 0.f; std::vector sigmas; diff --git a/src/pipeline/video.cpp b/src/pipeline/video.cpp index 2fcff4c9d..a07b93175 100644 --- a/src/pipeline/video.cpp +++ b/src/pipeline/video.cpp @@ -1,9 +1,11 @@ #include "generation.h" +#include "generation_info.h" #include #include #include #include +#include #include "conditioning/wan_audio.h" #include "core/rng.hpp" @@ -461,9 +463,35 @@ namespace sd::pipeline { return window; } + // ComfyUI-compatible reference sizing: per-axis rounding to 32 px keeps the reference + // divisible by VAE 16 x patch 2 without padding in the H3 graph. Both modes round rather + // than floor: flooring would under-size in-limit references and diverge from the + // upstream node, which rounds in both branches. + static std::pair minimax_h3_reference_size(int source_w, + int source_h, + int target_w, + int target_h, + RefImageSizeMode mode) { + constexpr double kCanvasMultiple = 32.0; + constexpr double kRefImageShortEdge = 2048.0; + double scale = 1.0; + if (mode == RefImageSizeMode::MAX) { + scale = std::min(1.0, kRefImageShortEdge / std::min(source_w, source_h)); + } else { + double source_area = static_cast(source_w) * source_h; + double target_area = static_cast(target_w) * target_h; + scale = std::min(1.0, std::sqrt(target_area / source_area)); + } + auto align32 = [kCanvasMultiple](double value) { + return std::max(32, static_cast(std::round(value / kCanvasMultiple)) * 32); + }; + return {align32(source_w * scale), align32(source_h * scale)}; + } + static std::optional prepare_video_generation_latents(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params, - GenerationRequest* request) { + GenerationRequest* request, + const RefImageParams& ref_image_params) { ImageGenerationLatents latents; int64_t prepare_start_ms = ggml_time_ms(); @@ -576,6 +604,7 @@ namespace sd::pipeline { sd_vid_gen_params->ref_videos_count, sd_vid_gen_params->ref_audios_count); + int64_t ref_image_tokens = 0; for (int i = 0; i < sd_vid_gen_params->ref_images_count; ++i) { auto image = ensure_image_tensor_channels( sd_image_to_tensor(sd_vid_gen_params->ref_images[i]), @@ -584,15 +613,18 @@ namespace sd::pipeline { LOG_ERROR("failed to load MiniMax-H3 reference image %d", i + 1); return std::nullopt; } - int source_w = static_cast(image.shape()[0]); - int source_h = static_cast(image.shape()[1]); - double source_area = static_cast(source_w) * source_h; - double target_area = static_cast(request->width) * request->height; - double scale = std::min(1.0, std::sqrt(target_area / source_area)); - int width = std::max(32, static_cast(std::round(source_w * scale / 32.f)) * 32); - int height = std::max(32, static_cast(std::round(source_h * scale / 32.f)) * 32); - image = resize_reference(image, width, height); - auto latent = sd->encode_first_stage(image); + int source_w = static_cast(image.shape()[0]); + int source_h = static_cast(image.shape()[1]); + auto [width, height] = + minimax_h3_reference_size(source_w, source_h, request->width, request->height, ref_image_params.ref_image_size); + LOG_INFO("MiniMax-H3 reference image %d: %dx%d -> %dx%d", + i + 1, + source_w, + source_h, + width, + height); + image = resize_reference(image, width, height); + auto latent = sd->encode_first_stage(image); if (latent.empty()) { LOG_ERROR("failed to encode MiniMax-H3 reference image %d", i + 1); return std::nullopt; @@ -602,6 +634,10 @@ namespace sd::pipeline { latents.minimax_reference_blocks.push_back({MiniMaxH3ReferenceKind::IMAGE, video_index, -1}); + // The DiT patchifies latents 2 x 2, so the token count is a quarter of the + // latent grid. + ref_image_tokens += static_cast(latents.ref_latents.back().shape()[0] / 2) * + (latents.ref_latents.back().shape()[1] / 2); MiniMaxH3PresentationItem item; item.kind = MiniMaxH3PresentationKind::IMAGE; item.frames.push_back(std::move(image)); @@ -654,8 +690,8 @@ namespace sd::pipeline { int source_index = std::min(reference.frame_count - 1, static_cast(std::floor(frame * source_fps / 24.0))); auto source = ensure_image_tensor_channels( - sd_image_to_tensor(reference.frames[source_index]), - 3); + sd_image_to_tensor(reference.frames[source_index]), + 3); if (source.empty()) { LOG_ERROR("invalid frame %d in MiniMax-H3 reference video %d", source_index + 1, @@ -711,6 +747,18 @@ namespace sd::pipeline { -1, encoded_index}); } + + if (ref_image_tokens > 0) { + int64_t target_tokens = sd->get_image_seq_len(request->width, request->height) / 4 * + sd->video_frames_to_latent_frames(request->frames); + if (target_tokens > 0 && ref_image_tokens > target_tokens * 4) { + LOG_WARN("MiniMax-H3 reference images add %" PRId64 + " tokens versus %" PRId64 + " for the generated video; sampling cost scales with every attended token", + ref_image_tokens, + target_tokens); + } + } } if (!has_references && (!start_image.empty() || !end_image.empty())) { @@ -771,7 +819,7 @@ namespace sd::pipeline { keyframes.shape()[2], 1, 1}, - conditioned_mask); + conditioned_mask); latents.denoise_mask = sd::ops::concat(latents.denoise_mask, keyframe_mask, 2); latents.video_positions = build_ltxv_video_positions(latents.init_latent.shape()[0], latents.init_latent.shape()[1], @@ -1059,10 +1107,10 @@ namespace sd::pipeline { inactive.shape()[1], inactive.shape()[2]}); // [t, h/vae_scale_factor, vae_scale_factor, w/vae_scale_factor, vae_scale_factor] mask_context = mask_context.permute({1, 3, 4, 0, 2}) // [vae_scale_factor, vae_scale_factor, t, h/vae_scale_factor, w/vae_scale_factor] - .reshape({inactive.shape()[0], - inactive.shape()[1], - inactive.shape()[2], - request->vae_scale_factor * request->vae_scale_factor}); // [vae_scale_factor*vae_scale_factor, t, h/vae_scale_factor, w/vae_scale_factor] + .reshape({inactive.shape()[0], + inactive.shape()[1], + inactive.shape()[2], + request->vae_scale_factor * request->vae_scale_factor}); // [vae_scale_factor*vae_scale_factor, t, h/vae_scale_factor, w/vae_scale_factor] if (!ref_image_latent.empty()) { vace_context = sd::ops::concat(ref_image_latent, vace_context, 2); // [b, 2*c, t+1, h/vae_scale_factor, w/vae_scale_factor] @@ -1147,7 +1195,8 @@ namespace sd::pipeline { static std::optional prepare_video_generation_embeds(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params, const GenerationRequest& request, - const ImageGenerationLatents& latents) { + const ImageGenerationLatents& latents, + const RefImageParams& ref_image_params) { ConditionerRunnerEndOnExit conditioner_runner_end{sd->cond_stage_model.get()}; ImageGenerationEmbeds embeds; @@ -1157,9 +1206,7 @@ namespace sd::pipeline { condition_params.zero_out_masked = true; condition_params.ref_images = &latents.ref_images; condition_params.minimax_h3_references = &latents.minimax_presentation_refs; - if (sd_version_is_lingbot_video(sd->version) || sd_version_is_minimax_h3(sd->version)) { - condition_params.ref_image_params.vlm_resize_mode = RefImageResizeMode::AREA; - } + condition_params.ref_image_params = ref_image_params; int64_t prepare_start_ms = ggml_time_ms(); embeds.cond = sd->get_learned_condition(condition_params); @@ -1341,9 +1388,9 @@ namespace sd::pipeline { auto& upsampler = scope.runner; upsampler = std::make_unique(sd->backend_for(SDBackendModule::UPSCALER), - model_manager->loader().get_tensor_storage_map(), - prefix, - model_manager); + model_manager->loader().get_tensor_storage_map(), + prefix, + model_manager); const size_t max_graph_vram_bytes = sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER); upsampler->set_max_graph_vram_bytes(max_graph_vram_bytes); if (upsampler->model == nullptr) { @@ -1526,9 +1573,16 @@ namespace sd::pipeline { img_gen_params.circular_x = sd_vid_gen_params->circular_x; img_gen_params.circular_y = sd_vid_gen_params->circular_y; img_gen_params.image_preprocess = sd_vid_gen_params->image_preprocess; + img_gen_params.ref_image_args = sd_vid_gen_params->ref_image_args; sd->animatediff_num_frames = n_frames; - bool ok = generate_image(sd, &img_gen_params, frames_out, num_frames_out); + bool ok = false; + try { + ok = generate_image(sd, &img_gen_params, frames_out, num_frames_out); + } catch (...) { + sd->animatediff_num_frames = 0; + throw; + } sd->animatediff_num_frames = 0; return ok; } @@ -1538,24 +1592,33 @@ namespace sd::pipeline { sd_image_t** frames_out, int* num_frames_out, sd_audio_t** audio_out, - int* fps_out) { + int* fps_out, + nlohmann::json* info_out) { + if (info_out) + *info_out = nlohmann::json(); + GenerationRequest request(sd, sd_vid_gen_params); + const auto info = sd::generation_info(request, *sd_vid_gen_params); if (sd->config_->animatediff_loaded && sd_version_supports_animatediff(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); if (fps_out != nullptr) { *fps_out = std::max(1, sd_vid_gen_params->fps); } - return generate_animatediff_video(sd, sd_vid_gen_params, frames_out, num_frames_out); + auto params = *sd_vid_gen_params; + params.seed = request.seed; + const bool ok = generate_animatediff_video(sd, ¶ms, frames_out, num_frames_out); + if (ok && info_out) + *info_out = info; + return ok; } sd->reset_cancel_flag(); - const RefImageParams ref_image_params; + const RefImageParams ref_image_params = sd->resolve_ref_image_params(sd_vid_gen_params->ref_image_args); int64_t t0 = ggml_time_ms(); sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params; sd->apply_circular_axes(sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y); - GenerationRequest request(sd, sd_vid_gen_params); sd::ImagePreprocessor preprocessing(sd_vid_gen_params->image_preprocess.rules); sd_vid_gen_params_t processed_params = *sd_vid_gen_params; if (!preprocessing.prepare_inputs(processed_params, request.width, request.height)) @@ -1589,7 +1652,7 @@ namespace sd::pipeline { sd->reset_generation_extensions(); SamplePlan plan(sd, sd_vid_gen_params, request); - auto latent_inputs_opt = prepare_video_generation_latents(sd, sd_vid_gen_params, &request); + auto latent_inputs_opt = prepare_video_generation_latents(sd, sd_vid_gen_params, &request, ref_image_params); if (!latent_inputs_opt.has_value()) { return false; } @@ -1598,7 +1661,8 @@ namespace sd::pipeline { auto embeds_opt = prepare_video_generation_embeds(sd, sd_vid_gen_params, request, - latents); + latents, + ref_image_params); if (!embeds_opt) { return false; } @@ -1686,9 +1750,9 @@ namespace sd::pipeline { embeds.img_uncond, sd::Tensor(), 0.f, - sd_vid_gen_params->sample_params.guidance, + request.guidance, plan.eta, - sd_vid_gen_params->sample_params.shifted_timestep, + request.shifted_timestep, plan.sample_method, sd->is_flow_denoiser(), plan.extra_sample_args, @@ -1798,7 +1862,7 @@ namespace sd::pipeline { std::vector hires_sigma_sched = make_hires_sigma_schedule(sd, request.hires, - sd_vid_gen_params->sample_params, + request.sample_params, hires_sample_method, plan.sample_steps, sd->get_image_seq_len(hires_request.height, hires_request.width) * T, @@ -1821,13 +1885,13 @@ namespace sd::pipeline { x_t, std::move(noise), embeds.cond, - hires_request.use_uncond ? embeds.uncond : SDCondition(), + hires_request.use_uncond ? embeds.uncond : SDCondition(), embeds.img_uncond, sd::Tensor(), 0.f, - sd_vid_gen_params->sample_params.guidance, + request.guidance, hires_eta, - sd_vid_gen_params->sample_params.shifted_timestep, + request.shifted_timestep, hires_sample_method, sd->is_flow_denoiser(), plan.extra_sample_args, @@ -1959,6 +2023,14 @@ namespace sd::pipeline { generated_audio->sample_count = want_samples; } } + if (info_out) { + try { + *info_out = info; + } catch (...) { + free_sd_audio(generated_audio); + throw; + } + } if (audio_out != nullptr) { *audio_out = generated_audio; } else { diff --git a/src/runtime/cache_dit.hpp b/src/runtime/cache_dit.hpp index 273aa0ef3..666e61cad 100644 --- a/src/runtime/cache_dit.hpp +++ b/src/runtime/cache_dit.hpp @@ -842,7 +842,7 @@ struct CacheDitConditionState { const float* input_data = input.data(); float diff = CacheDitState::calculate_residual_diff( - it->second.prev_input.data(), input_data, ne); + it->second.prev_input.data(), input_data, ne); float effective_threshold = config.residual_diff_threshold; if (config.Fn_compute_blocks > 0) { diff --git a/src/runtime/guidance.cpp b/src/runtime/guidance.cpp index bfb773b0a..854b55661 100644 --- a/src/runtime/guidance.cpp +++ b/src/runtime/guidance.cpp @@ -164,8 +164,8 @@ namespace sd::guidance { if (has_tensor(input.pred_img_uncond)) { const sd::Tensor& pred_img_uncond = *input.pred_img_uncond; output.pred = pred_img_uncond + - image_guidance_scale_ * (pred_uncond - pred_img_uncond) + - guidance_scale * (pred_cond - pred_uncond); + image_guidance_scale_ * (pred_uncond - pred_img_uncond) + + guidance_scale * (pred_cond - pred_uncond); } else { output.pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond); @@ -224,8 +224,8 @@ namespace sd::guidance { if (has_tensor(input.pred_img_uncond)) { const sd::Tensor& pred_img_uncond = *input.pred_img_uncond; output.pred = pred_img_uncond + - image_guidance_scale_ * (pred_uncond - pred_img_uncond) + - guidance_scale * (pred_cond - pred_uncond); + image_guidance_scale_ * (pred_uncond - pred_img_uncond) + + guidance_scale * (pred_cond - pred_uncond); } else { output.pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond); } diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index d476f92d6..bca4b79ba 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -1,4 +1,5 @@ #include "stable-diffusion.h" +#include "generation_info.h" #include #include @@ -171,10 +172,13 @@ const char* prediction_to_str[] = { "sefi_flow", "minit2i_flow", "sensenova_u1_flow", + "flux3_flow", + "h3_av_flow", + "ming_image_flow", }; const char* sd_prediction_name(enum prediction_t prediction) { - if (prediction < PREDICTION_COUNT) { + if (prediction >= EPS_PRED && prediction < PREDICTION_COUNT) { return prediction_to_str[prediction]; } return NONE_STR; @@ -633,6 +637,7 @@ void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params) { sd_vid_gen_params->hires.custom_sigmas_count = 0; sd_vid_gen_params->circular_x = false; sd_vid_gen_params->circular_y = false; + sd_vid_gen_params->ref_image_args = ""; sd_cache_params_init(&sd_vid_gen_params->cache); } @@ -704,13 +709,6 @@ SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) { return sd_version_supports_video_generation(sd_ctx->sd->version); } -SD_API bool sd_ctx_supports_ref_image_conditioning(const sd_ctx_t* sd_ctx) { - if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { - return false; - } - return sd_version_supports_ref_image_conditioning(sd_ctx->sd->version); -} - SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path) { if (sd_ctx == nullptr || sd_ctx->sd == nullptr || path == nullptr) { return false; @@ -729,14 +727,20 @@ SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx) { if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { return false; } - return sd_ctx->sd->control_net != nullptr; + if (sd_ctx->sd->control_net != nullptr) { + return true; + } + // z-image keeps control net weights inside the diffusion runner, registered + // under ModelComponent::ControlNet instead of a separate control_net runner. + return sd_ctx->sd->model_manager != nullptr && + sd_ctx->sd->model_manager->registered_params_size({ModelComponent::ControlNet}) > 0; } const char* sd_get_model_version_name(const sd_ctx_t* sd_ctx) { if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_ctx->sd->version >= VERSION_COUNT) { return "Unknown"; } - return model_version_to_str[sd_ctx->sd->version]; + return sd_version_name(sd_ctx->sd->version); } enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) { @@ -749,43 +753,62 @@ enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_me SD_API bool generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* params, - sd_image_t** images_out, - int* num_images_out) { - if (images_out != nullptr) - *images_out = nullptr; - if (num_images_out != nullptr) - *num_images_out = 0; - if (sd_ctx == nullptr || sd_ctx->sd == nullptr || params == nullptr) { + sd_image_generation_result_t** result_out) { + if (result_out) + *result_out = nullptr; + if (!result_out || !sd_ctx || !sd_ctx->sd || !params) return false; - } StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd); - return execution.ready && sd::pipeline::generate_image(sd_ctx->sd, params, images_out, num_images_out); + if (!execution.ready) + return false; + sd_image_t* images = nullptr; + int count = 0; + std::vector infos; + try { + if (!sd::pipeline::generate_image(sd_ctx->sd, params, &images, &count, &infos)) { + free_sd_images(images, count); + return false; + } + } catch (const std::exception& e) { + LOG_ERROR("image generation failed: %s", e.what()); + free_sd_images(images, count); + return false; + } + return sd::make_image_generation_result(images, count, infos, result_out); } SD_API bool generate_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, - sd_audio_t** audio_out, - int* fps_out) { - if (frames_out != nullptr) - *frames_out = nullptr; - if (audio_out != nullptr) - *audio_out = nullptr; - if (num_frames_out != nullptr) - *num_frames_out = 0; - if (fps_out != nullptr) - *fps_out = 0; - if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) { + const sd_vid_gen_params_t* params, + sd_video_generation_result_t** result_out) { + if (result_out) + *result_out = nullptr; + if (!result_out || !sd_ctx || !sd_ctx->sd || !params) return false; - } - StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd); - if (!execution.ready) { + if (!execution.ready) + return false; + auto* result = static_cast(calloc(1, sizeof(sd_video_generation_result_t))); + if (!result) + return false; + try { + nlohmann::json info; + if (!sd::pipeline::generate_video(sd_ctx->sd, params, &result->frames, &result->frame_count, + &result->audio, &result->fps, &info)) { + free_sd_video_generation_result(result); + return false; + } + result->generation_info = sd::copy_generation_info(info); + if (!result->generation_info) { + free_sd_video_generation_result(result); + return false; + } + } catch (const std::exception& e) { + LOG_ERROR("video generation failed: %s", e.what()); + free_sd_video_generation_result(result); return false; } - - return sd::pipeline::generate_video(sd_ctx->sd, sd_vid_gen_params, frames_out, num_frames_out, audio_out, fps_out); + *result_out = result; + return true; } SD_API void free_sd_images(sd_image_t* result_images, int num_images) { diff --git a/src/upscaler.cpp b/src/upscaler.cpp index dfe5e4b22..c9bbbcb6b 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -1,6 +1,7 @@ #include "upscaler.h" #include "core/ggml_extend_backend.h" #include "core/util.h" +#include "generation_info.h" #include "model_loader.h" #include "runtime/tiling.h" #include "stable-diffusion.h" @@ -259,14 +260,11 @@ upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str, bool upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_t upscale_factor, - sd_image_t** images_out, - int* num_images_out) { - if (images_out != nullptr) { - *images_out = nullptr; - } - if (num_images_out != nullptr) { - *num_images_out = 0; - } + sd_image_generation_result_t** result_out) try { + if (result_out) + *result_out = nullptr; + if (!result_out) + return false; if (upscaler_ctx == nullptr || upscaler_ctx->upscaler == nullptr) { return false; } @@ -276,21 +274,17 @@ bool upscale(upscaler_ctx_t* upscaler_ctx, return false; } + auto release_images = [](sd_image_t* images) { free_sd_images(images, 1); }; + std::unique_ptr image_owner(result_images, release_images); + const std::vector infos{sd::upscale_generation_info(upscale_factor)}; result_images[0] = upscaler_ctx->upscaler->upscale(input_image, upscale_factor); - if (result_images[0].data == nullptr) { - free(result_images); + if (!result_images[0].data) return false; - } - if (num_images_out != nullptr) { - *num_images_out = 1; - } - if (images_out != nullptr) { - *images_out = result_images; - } else { - free_sd_images(result_images, 1); - } - return true; + return sd::make_image_generation_result(image_owner.release(), 1, infos, result_out); +} catch (const std::exception& e) { + LOG_ERROR("upscale failed: %s", e.what()); + return false; } int get_upscale_factor(upscaler_ctx_t* upscaler_ctx) { @@ -329,6 +323,8 @@ int get_upscaler_model_scale(const char* model_path) { } void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx) { + if (!upscaler_ctx) + return; if (upscaler_ctx->upscaler != nullptr) { delete upscaler_ctx->upscaler; upscaler_ctx->upscaler = nullptr;