diff --git a/docs/model_config.md b/docs/model_config.md index dda93a837f..0d9290c1d5 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/examples/common/common.cpp b/examples/common/common.cpp index df76c517a1..24e1967edb 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -818,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", diff --git a/examples/server/api.md b/examples/server/api.md index fc64fb0296..5e81d05809 100644 --- a/examples/server/api.md +++ b/examples/server/api.md @@ -447,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`: diff --git a/examples/server/openai_edit_policy.h b/examples/server/openai_edit_policy.h deleted file mode 100644 index c963a887f8..0000000000 --- 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_openai.cpp b/examples/server/routes_openai.cpp index 27a7e18f90..ae1fb705ed 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,24 +157,7 @@ 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; @@ -190,17 +172,9 @@ static bool build_openai_edit_request(const httplib::Request& req, 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"; diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 85fe51d463..46c8d2fed6 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 }; @@ -486,8 +489,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); diff --git a/src/model.h b/src/model.h index 73a369301a..3ee9553836 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/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 7a8693ef30..3e8f0fd841 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" diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp index 42599c67a0..cdf65f300f 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/integrations/anima.cpp b/src/model/integrations/anima.cpp new file mode 100644 index 0000000000..0c4600e73c --- /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 0000000000..616cab2fac --- /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 0000000000..65d431fcb4 --- /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 0000000000..15cfaf72bd --- /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 0000000000..4f14237aac --- /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 0000000000..3b11888b08 --- /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 0000000000..a868edd932 --- /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 0000000000..cccb38eab1 --- /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 0000000000..0e283e0a25 --- /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 0000000000..c90e7616f4 --- /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 0000000000..2b54492444 --- /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 0000000000..9414efc75c --- /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 0000000000..f492b7c4ca --- /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 0000000000..d2f502640a --- /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 0000000000..bafa75e8cc --- /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 0000000000..2770ba69d0 --- /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 0000000000..e1141f290d --- /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 0000000000..9be11321e8 --- /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 0000000000..68ec704e04 --- /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 0000000000..31c69f949a --- /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 0000000000..bc023f68a1 --- /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 0000000000..96cb1d4ce9 --- /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 0000000000..63921058ae --- /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 0000000000..93145d15af --- /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 0000000000..199c76ceab --- /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 0000000000..5c37fb1b53 --- /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 0000000000..f85c444953 --- /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 0000000000..519a68040c --- /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 0000000000..75acce029b --- /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 0000000000..f0e5216ef7 --- /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 0000000000..3ba31bfde7 --- /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 0000000000..5ba88dcc73 --- /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_definition.h b/src/model_definition.h new file mode 100644 index 0000000000..2b49994958 --- /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_registry.cpp b/src/model_registry.cpp new file mode 100644 index 0000000000..9a80bb8359 --- /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 a913ae779a..18ef5eb166 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" @@ -55,70 +56,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, @@ -1014,7 +951,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) { @@ -1470,98 +1407,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 +1433,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: { @@ -2999,63 +2855,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,35 +3015,7 @@ 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_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 { diff --git a/src/pipeline/diffusion_engine.h b/src/pipeline/diffusion_engine.h index c5509d9a9e..1e41662df4 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) || diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index c683b53ffc..4b5905f01c 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -804,7 +804,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; } diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index ca4d1998b1..2202314f01 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), + 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); + 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, "", diff --git a/src/pipeline/model_builders.h b/src/pipeline/model_builders.h index 5edc75680f..bb51ccc14a 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 2a46fc9cba..0693532fbf 100644 --- a/src/pipeline/request.cpp +++ b/src/pipeline/request.cpp @@ -6,6 +6,7 @@ #include #include "diffusion_engine.h" +#include "model_definition.h" #include "runtime/denoiser.hpp" namespace sd::pipeline { @@ -43,15 +44,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,16 +59,9 @@ 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; } @@ -184,7 +171,7 @@ namespace sd::pipeline { 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)); } } diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index d476f92d6b..f9d724f0a5 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -171,10 +171,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; @@ -704,13 +707,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; @@ -736,7 +732,7 @@ 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) {