diff --git a/xllm/core/framework/model/model_args.h b/xllm/core/framework/model/model_args.h index bfdf1334f5..4f6b7cc570 100644 --- a/xllm/core/framework/model/model_args.h +++ b/xllm/core/framework/model/model_args.h @@ -460,12 +460,18 @@ struct ModelArgs { PROPERTY(int64_t, vae_scale_factor_spatial) = 0; PROPERTY(bool, vae_is_residual) = false; + PROPERTY(float, batch_norm_eps) = 1e-04f; + PROPERTY(float, batch_norm_momentum) = 0.1f; + PROPERTY(std::vector, ae_patch_size) = {}; + // dit related args PROPERTY(int64_t, joint_attention_dim) = 0; PROPERTY(int64_t, pooled_projection_dim) = 0; PROPERTY(bool, guidance_embeds) = true; PROPERTY(std::vector, axes_dims_rope) = {}; PROPERTY(int64_t, num_single_layers) = 0; + + PROPERTY(float, mlp_ratio) = 3.0f; PROPERTY(int, timestep_guidance_channels) = 256; PROPERTY(int64_t, patch_size) = 1; PROPERTY(std::vector, wan_patch_size) = { 1, 2, 2 }; diff --git a/xllm/core/layers/common/add_matmul.cpp b/xllm/core/layers/common/add_matmul.cpp index ff9a157a0a..a11dbc9712 100644 --- a/xllm/core/layers/common/add_matmul.cpp +++ b/xllm/core/layers/common/add_matmul.cpp @@ -194,10 +194,9 @@ void AddMatmulWeightTransposedImpl::load_state_dict( } else { if (state_dict.has("weight")) { weight::load_weight(state_dict, "weight", weight_, weight_is_loaded_); - if (with_bias_) { - torch::Tensor transposed = weight_.data().transpose(0, 1).contiguous(); - weight_.set_data(transposed); - } + + torch::Tensor transposed = weight_.data().transpose(0, 1).contiguous(); + weight_.set_data(transposed); } } if (with_bias_) { diff --git a/xllm/models/dit/autoencoders/autoencoder_kl_flux2.h b/xllm/models/dit/autoencoders/autoencoder_kl_flux2.h new file mode 100644 index 0000000000..76d153a61e --- /dev/null +++ b/xllm/models/dit/autoencoders/autoencoder_kl_flux2.h @@ -0,0 +1,213 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once +#include "models/dit/autoencoders/autoencoder_kl.h" + +namespace xllm { + +// VAE implementation for Flux2, including encoder and decoder with BatchNorm2d +class AutoencoderKLFlux2Impl final : public torch::nn::Module { + public: + explicit AutoencoderKLFlux2Impl(const ModelContext& context) + : args_(context.get_model_args()) { + encoder_ = register_module("encoder", VAEEncoder(context)); + decoder_ = register_module("decoder", VAEDecoder(context)); + if (args_.use_quant_conv()) { + quant_conv_ = register_module( + "quant_conv", + torch::nn::Conv2d(torch::nn::Conv2dOptions( + 2 * args_.latent_channels(), 2 * args_.latent_channels(), 1))); + } + if (args_.use_post_quant_conv()) { + post_quant_conv_ = register_module( + "post_quant_conv", + torch::nn::Conv2d(torch::nn::Conv2dOptions( + args_.latent_channels(), args_.latent_channels(), 1))); + } + + auto dtype = context.get_tensor_options().dtype().toScalarType(); + encoder_->to(dtype); + decoder_->to(dtype); + if (args_.use_quant_conv()) { + quant_conv_->to(dtype); + } + if (args_.use_post_quant_conv()) { + post_quant_conv_->to(dtype); + } + + int64_t patch_size_prod = 1; + for (int64_t ps : args_.ae_patch_size()) { + patch_size_prod *= ps; + } + int64_t bn_num_features = patch_size_prod * args_.latent_channels(); + + bn_ = register_module( + "bn", + torch::nn::BatchNorm2d(torch::nn::BatchNorm2dOptions(bn_num_features) + .eps(args_.batch_norm_eps()) + .momentum(args_.batch_norm_momentum()) + .affine(false) + .track_running_stats(true))); + bn_->to(dtype); + } + + torch::Tensor encode(const torch::Tensor& images, int64_t seed) { + auto enc = encoder_(images); + if (args_.use_quant_conv()) { + enc = quant_conv_(enc); + } + auto posterior = DiagonalGaussianDistribution(enc); + return posterior.sample(seed); + } + + torch::Tensor decode(const torch::Tensor& latents) { + torch::Tensor processed_latents = latents; + + if (args_.use_post_quant_conv()) { + processed_latents = post_quant_conv_(processed_latents); + } + + auto dec = decoder_(processed_latents); + return dec; + } + + void load_model(std::unique_ptr loader) { + for (const auto& state_dict : loader->get_state_dicts()) { + encoder_->load_state_dict(state_dict->get_dict_with_prefix("encoder.")); + decoder_->load_state_dict(state_dict->get_dict_with_prefix("decoder.")); + if (args_.use_quant_conv()) { + weight::load_weight(state_dict->get_dict_with_prefix("quant_conv."), + "weight", + quant_conv_->weight, + is_quant_conv_weight_); + weight::load_weight(state_dict->get_dict_with_prefix("quant_conv."), + "bias", + quant_conv_->bias, + is_quant_conv_bias_); + } + if (args_.use_post_quant_conv()) { + weight::load_weight( + state_dict->get_dict_with_prefix("post_quant_conv."), + "weight", + post_quant_conv_->weight, + is_post_quant_conv_weight_); + weight::load_weight( + state_dict->get_dict_with_prefix("post_quant_conv."), + "bias", + post_quant_conv_->bias, + is_post_quant_conv_bias_); + } + + weight::load_weight(state_dict->get_dict_with_prefix("bn."), + "running_mean", + bn_->running_mean, + is_bn_running_mean_); + weight::load_weight(state_dict->get_dict_with_prefix("bn."), + "running_var", + bn_->running_var, + is_bn_running_var_); + /*weight::load_weight(state_dict->get_dict_with_prefix("bn."), + "num_batches_tracked", + bn_->num_batches_tracked, + is_bn_num_batches_tracked_);*/ + } + verify_loaded_weights(""); + LOG(INFO) << "VAE model loaded successfully."; + } + + void verify_loaded_weights(const std::string& prefix) { + encoder_->verify_loaded_weights(prefix + "encoder."); + decoder_->verify_loaded_weights(prefix + "decoder."); + if (args_.use_quant_conv()) { + CHECK(is_quant_conv_weight_) + << "weight is not loaded for " << prefix + "quant_conv.weight"; + CHECK(is_quant_conv_bias_) + << "bias is not loaded for " << prefix + "quant_conv.bias"; + } + if (args_.use_post_quant_conv()) { + CHECK(is_post_quant_conv_weight_) + << "weight is not loaded for " << prefix + "post_quant_conv.weight"; + CHECK(is_post_quant_conv_bias_) + << "bias is not loaded for " << prefix + "post_quant_conv.bias"; + } + CHECK(is_bn_running_mean_) + << "running_mean is not loaded for " << prefix + "bn.running_mean"; + CHECK(is_bn_running_var_) + << "running_var is not loaded for " << prefix + "bn.running_var"; + /* CHECK(is_bn_num_batches_tracked_) + << "num_batches_tracked is not loaded for " + << prefix + "bn.num_batches_tracked";*/ + } + + torch::Tensor get_bn_running_mean() const { return bn_->running_mean; } + + torch::Tensor get_bn_running_var() const { return bn_->running_var; } + + float get_batch_norm_eps() const { return args_.batch_norm_eps(); } + + private: + VAEEncoder encoder_ = nullptr; + VAEDecoder decoder_ = nullptr; + torch::nn::Conv2d quant_conv_ = nullptr; + torch::nn::Conv2d post_quant_conv_ = nullptr; + torch::nn::BatchNorm2d bn_ = nullptr; + bool use_post_quant_conv_ = false; + + bool is_quant_conv_weight_ = false; + bool is_quant_conv_bias_ = false; + bool is_post_quant_conv_weight_ = false; + bool is_post_quant_conv_bias_ = false; + bool is_bn_running_mean_ = false; + bool is_bn_running_var_ = false; + /*bool is_bn_num_batches_tracked_ = false;*/ + ModelArgs args_; +}; +TORCH_MODULE(AutoencoderKLFlux2); + +// register VAE model with the model registry +REGISTER_MODEL_ARGS(AutoencoderKLFlux2, [&] { + LOAD_ARG_OR(dtype, "dtype", "bfloat16"); + LOAD_ARG_OR(in_channels, "in_channels", 3); + LOAD_ARG_OR(out_channels, "out_channels", 3); + LOAD_ARG_OR(down_block_types, + "down_block_types", + (std::vector{"DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D"})); + LOAD_ARG_OR(up_block_types, + "up_block_types", + (std::vector{"UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D"})); + LOAD_ARG_OR(block_out_channels, + "block_out_channels", + (std::vector{128, 256, 512, 512})); + LOAD_ARG_OR(layers_per_block, "layers_per_block", 2); + LOAD_ARG_OR(latent_channels, "latent_channels", 32); + LOAD_ARG_OR(norm_num_groups, "norm_num_groups", 32); + LOAD_ARG_OR(sample_size, "sample_size", 1024); + LOAD_ARG_OR(mid_block_add_attention, "mid_block_add_attention", true); + LOAD_ARG_OR(force_upcast, "force_upcast", true); + LOAD_ARG_OR(use_quant_conv, "use_quant_conv", false); + LOAD_ARG_OR(use_post_quant_conv, "use_post_quant_conv", false); + LOAD_ARG_OR(batch_norm_eps, "batch_norm_eps", 1e-04f); + LOAD_ARG_OR(act_fn, "act_fn", "silu"); + LOAD_ARG_OR(batch_norm_momentum, "batch_norm_momentum", 0.1f); + LOAD_ARG_OR(ae_patch_size, "patch_size", (std::vector{2, 2})); +}); +} // namespace xllm diff --git a/xllm/models/dit/pipelines/pipeline_flux2.h b/xllm/models/dit/pipelines/pipeline_flux2.h new file mode 100644 index 0000000000..7d4bcf1476 --- /dev/null +++ b/xllm/models/dit/pipelines/pipeline_flux2.h @@ -0,0 +1,380 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once +#include +#include + +#include "core/framework/dit_cache/dit_cache.h" +#include "models/dit/autoencoders/autoencoder_kl_flux2.h" +#include "models/dit/pipelines/pipeline_flux2_base.h" +#include "models/dit/processors/flux2_image_processor.h" +#include "models/dit/schedulers/flowmatch_euler_discrete_scheduler.h" +#include "models/dit/transformers/transformer_flux2.h" + +namespace xllm { + +class Flux2PipelineImpl final : public Flux2PipelineBaseImpl { + public: + explicit Flux2PipelineImpl(const DiTModelContext& context) { + const auto& model_args = context.get_model_args("vae"); + options_ = context.get_tensor_options(); + vae_scale_factor_ = 1 << (model_args.block_out_channels().size() - 1); + + vae_shift_factor_ = model_args.shift_factor(); + vae_scaling_factor_ = model_args.scale_factor(); + tokenizer_max_length_ = 512; + default_sample_size_ = 128; + + flux2_image_processor_ = Flux2ImageProcessor( + context.get_model_context("vae"), vae_scale_factor_ * 2); + vae_ = AutoencoderKLFlux2(context.get_model_context("vae")); + pos_embed_ = register_module( + "pos_embed", + Flux2PosEmbed(context.get_model_args("transformer").rope_theta(), + context.get_model_args("transformer").axes_dims_rope())); + + transformer_ = Flux2DiTModel(context.get_model_context("transformer"), + context.get_parallel_args()); + num_single_layers_ = + context.get_model_args("transformer").num_single_layers(); + scheduler_ = + FlowMatchEulerDiscreteScheduler(context.get_model_context("scheduler")); + register_module("vae", vae_); + register_module("flux2_image_processor", flux2_image_processor_); + register_module("transformer", transformer_); + register_module("scheduler", scheduler_); + } + + DiTForwardOutput forward(const DiTForwardInput& input) { + const auto& generation_params = input.generation_params; + + int64_t seed = generation_params.seed > 0 ? generation_params.seed : 42; + auto prompts = std::make_optional(input.prompts); + auto latents = input.latents.defined() ? std::make_optional(input.latents) + : std::nullopt; + auto prompt_embeds = input.prompt_embeds.defined() + ? std::make_optional(input.prompt_embeds) + : std::nullopt; + auto images = input.images.defined() ? std::make_optional(input.images) + : std::nullopt; + + auto output = forward_impl( + prompts, // prompt + generation_params.height, // height + generation_params.width, // width + generation_params.num_inference_steps, // num_inference_steps + generation_params.guidance_scale, // guidance_scale + generation_params.num_images_per_prompt, // num_images_per_prompt + seed, // seed + latents, // latents + prompt_embeds, // prompt_embeds + images, // images + generation_params.max_sequence_length // max_sequence_length + ); + + DiTForwardOutput out; + out.tensors = torch::chunk(output, input.batch_size); + return out; + } + + void load_model(std::unique_ptr loader) { + std::string model_path = loader->model_root_path(); + auto transformer_loader = loader->take_component_loader("transformer"); + auto vae_loader = loader->take_component_loader("vae"); + transformer_->load_model(std::move(transformer_loader)); + transformer_->to(options_.device()); + vae_->load_model(std::move(vae_loader)); + vae_->to(options_.device()); + } + + private: + std::pair prepare_latents( + int64_t batch_size, + int64_t num_channels_latents, + int64_t height, + int64_t width, + int64_t seed, + std::optional latents = std::nullopt) { + int64_t adjusted_height = 2 * (height / (vae_scale_factor_ * 2)); + int64_t adjusted_width = 2 * (width / (vae_scale_factor_ * 2)); + std::vector shape = {batch_size, + num_channels_latents * 4, + adjusted_height / 2, + adjusted_width / 2}; + if (latents.has_value()) { + torch::Tensor latent_image_ids = + prepare_latent_image_ids(latents.value()); + return {latents.value(), latent_image_ids}; + } + torch::Tensor latents_tensor = + xllm::dit::randn_tensor(shape, seed, options_); + torch::Tensor packed_latents = pack_latents(latents_tensor); + torch::Tensor latent_image_ids = prepare_latent_image_ids(latents_tensor); + return {packed_latents, latent_image_ids}; + } + + std::pair prepare_image_latents( + const std::vector& images, + int64_t batch_size, + int64_t seed) { + std::vector image_latents; + for (const auto& image : images) { + auto image_latent = vae_->encode(image, seed); + auto patched_latent = patchify_latents(image_latent); + auto latents_bn_mean = + vae_->get_bn_running_mean() + .view({1, -1, 1, 1}) + .to(image_latent.device(), image_latent.dtype()); + auto latents_bn_std = + torch::sqrt(vae_->get_bn_running_var().view({1, -1, 1, 1}) + + vae_->get_batch_norm_eps()); + image_latent = (patched_latent - latents_bn_mean) / latents_bn_std; + image_latents.push_back(image_latent); + } + auto concatenated_latents = pack_latents_for_images(image_latents); + auto image_latent_ids = _prepare_image_ids(image_latents); + + concatenated_latents = concatenated_latents.unsqueeze(0); + auto repeated_latents = concatenated_latents.repeat({batch_size, 1, 1}); + auto repeated_ids = image_latent_ids.repeat({batch_size, 1, 1}); + return {repeated_latents, repeated_ids}; + } + + torch::Tensor pack_latents_for_images( + const std::vector& image_latents) { + std::vector packed_latents; + packed_latents.reserve(image_latents.size()); + for (const auto& latent : image_latents) { + auto packed = pack_latents(latent); + packed = packed.squeeze(0); + packed_latents.emplace_back(packed); + } + return torch::cat(packed_latents, 0); + } + + torch::Tensor forward_impl( + std::optional> prompt, + int64_t height = 512, + int64_t width = 512, + int64_t num_inference_steps = 50, + float guidance_scale = 4.0f, + int64_t num_images_per_prompt = 1, + std::optional seed = std::nullopt, + std::optional latents = std::nullopt, + std::optional prompt_embeds = std::nullopt, + std::optional images = std::nullopt, + int64_t max_sequence_length = 512, + const std::vector& hidden_states_layers = {10, 20, 30}, + const std::string& system_message = SYSTEM_MESSAGE) { + torch::NoGradGuard no_grad; + + int64_t batch_size = prompt_embeds.value().size(0); + int64_t total_batch_size = batch_size * num_images_per_prompt; + + int64_t num_channels = prompt_embeds.value().size(1); + int64_t seq_len = prompt_embeds.value().size(2); + int64_t hidden_dim = prompt_embeds.value().size(3); + + auto prompt_embeds_value = + prompt_embeds.value() + .permute({0, 2, 1, 3}) + .reshape({batch_size, seq_len, num_channels * hidden_dim}); + + prompt_embeds_value = + prompt_embeds_value.repeat({1, num_images_per_prompt, 1}); + prompt_embeds_value = prompt_embeds_value.view( + {batch_size * num_images_per_prompt, seq_len, -1}); + + device_ = options_.device(); + + torch::Tensor encoded_prompt_embeds = prompt_embeds_value; + torch::Tensor text_ids = prepare_text_ids(prompt_embeds_value); + encoded_prompt_embeds = + encoded_prompt_embeds.to(device_).to(torch::kBFloat16); + text_ids = text_ids.to(device_).to(torch::kLong); + + // process images + std::vector condition_images_list; + if (images.has_value()) { + auto input_images = images.value(); + if (input_images.dim() == 3) { + input_images = input_images.unsqueeze(0); + } + + for (int64_t i = 0; i < input_images.size(0); ++i) { + auto img = input_images[i]; + flux2_image_processor_->check_image_input(img); + int64_t image_width = img.size(-1); + int64_t image_height = img.size(-2); + + if (image_width * image_height > 1024 * 1024) { + img = flux2_image_processor_->resize_to_target_area(img, 1024 * 1024); + image_width = img.size(-1); + image_height = img.size(-2); + } + int64_t multiple_of = vae_scale_factor_ * 2; + image_width = (image_width / multiple_of) * multiple_of; + image_height = (image_height / multiple_of) * multiple_of; + img = + flux2_image_processor_->preprocess(img, image_height, image_width); + condition_images_list.push_back(img); + } + } + + // prepare latent + int64_t num_channels_latents = transformer_->in_channels() / 4; + auto [prepared_latents, latent_image_ids] = + prepare_latents(total_batch_size, + num_channels_latents, + height, + width, + seed.has_value() ? seed.value() : 42, + latents); + prepared_latents = prepared_latents.to(torch::kBFloat16); + latent_image_ids = latent_image_ids.to(torch::kLong); + + torch::Tensor image_latents; + torch::Tensor image_latent_ids; + if (!condition_images_list.empty()) { + std::tie(image_latents, image_latent_ids) = + prepare_image_latents(condition_images_list, + total_batch_size, + seed.has_value() ? seed.value() : 42); + } + + // prepare timestep + std::vector new_sigmas; + new_sigmas.reserve(num_inference_steps); + for (int64_t i = 0; i < num_inference_steps; ++i) { + new_sigmas.emplace_back(1.0f - static_cast(i) / + (num_inference_steps - 1) * + (1.0f - 1.0f / num_inference_steps)); + } + + int64_t image_seq_len = prepared_latents.size(1); + float mu = compute_empirical_mu(image_seq_len, num_inference_steps); + auto [timesteps, num_inference_steps_actual] = flux2_retrieve_timesteps( + scheduler_, num_inference_steps, options_.device(), new_sigmas, mu); + + torch::Tensor guidance; + torch::TensorOptions options = + torch::dtype(torch::kFloat32).device(options_.device()); + + guidance = torch::full(at::IntArrayRef({1}), guidance_scale, options); + guidance = guidance.expand({prepared_latents.size(0)}); + + auto [rot_emb1, rot_emb2] = + pos_embed_->forward_cache(text_ids, + latent_image_ids, + height / (vae_scale_factor_ * 2), + width / (vae_scale_factor_ * 2)); + + torch::Tensor image_rotary_emb = + torch::stack({rot_emb1, rot_emb2}, 0).to(options_.dtype()); + + // denosing loop + DiTCache::get_instance().set_infer_steps(num_inference_steps); + DiTCache::get_instance().set_num_blocks(num_single_layers_); + scheduler_->set_begin_index(0); + torch::Tensor timestep = + torch::empty({prepared_latents.size(0)}, prepared_latents.options()); + + for (int64_t i = 0; i < timesteps.numel(); ++i) { + torch::Tensor t = timesteps[i].unsqueeze(0); + timestep.fill_(t.item()) + .to(prepared_latents.dtype()) + .div_(1000.0f); + + torch::Tensor latent_model_input = prepared_latents.to(options_.dtype()); + torch::Tensor latent_image_ids_input = latent_image_ids; + if (image_latents.defined()) { + latent_model_input = torch::cat({prepared_latents, image_latents}, 1) + .to(options_.dtype()); + latent_image_ids_input = + torch::cat({latent_image_ids, image_latent_ids}, 1); + auto [rot_emb1_img, rot_emb2_img] = + pos_embed_->forward_cache(text_ids, + latent_image_ids_input, + height / (vae_scale_factor_ * 2), + width / (vae_scale_factor_ * 2)); + image_rotary_emb = + torch::stack({rot_emb1_img, rot_emb2_img}, 0).to(options_.dtype()); + } + torch::Tensor noise_pred = transformer_->forward(latent_model_input, + encoded_prompt_embeds, + timestep, + guidance, + image_rotary_emb, + /*step_idx=*/i); + + if (image_latents.defined()) { + noise_pred = noise_pred.narrow(1, 0, prepared_latents.size(1)); + } + + auto prev_latents = scheduler_->step(noise_pred, t, prepared_latents); + + prepared_latents = prev_latents.detach(); + std::vector tensors = {prepared_latents, noise_pred}; + noise_pred.reset(); + prev_latents = torch::Tensor(); + + if (latents.has_value() && + prepared_latents.dtype() != latents.value().dtype()) { + prepared_latents = prepared_latents.to(latents.value().dtype()); + } + } + + torch::Tensor image; + + torch::Tensor unpacked_latents = + unpack_latents_with_ids(prepared_latents, latent_image_ids); + auto latents_bn_mean = + vae_->get_bn_running_mean() + .view({1, -1, 1, 1}) + .to(unpacked_latents.device(), unpacked_latents.dtype()); + auto latents_bn_std = + torch::sqrt(vae_->get_bn_running_var().view({1, -1, 1, 1}) + + vae_->get_batch_norm_eps()) + .to(unpacked_latents.device(), unpacked_latents.dtype()); + unpacked_latents = unpacked_latents * latents_bn_std + latents_bn_mean; + unpacked_latents = unpatchify_latents(unpacked_latents); + + image = vae_->decode(unpacked_latents); + + image = flux2_image_processor_->postprocess(image); + + return image; + } + + private: + FlowMatchEulerDiscreteScheduler scheduler_{nullptr}; + AutoencoderKLFlux2 vae_{nullptr}; + Flux2ImageProcessor flux2_image_processor_{nullptr}; + Flux2DiTModel transformer_{nullptr}; + // Mistral3EncoderModel mistral3_{nullptr}; + float vae_scaling_factor_; + float vae_shift_factor_; + int32_t tokenizer_max_length_; + int32_t default_sample_size_; + int32_t vae_scale_factor_; + int64_t num_single_layers_; + Flux2PosEmbed pos_embed_{nullptr}; + // std::unique_ptr chat_template_; +}; +TORCH_MODULE(Flux2Pipeline); + +REGISTER_DIT_MODEL(flux2, Flux2Pipeline); +}; // namespace xllm diff --git a/xllm/models/dit/pipelines/pipeline_flux2_base.h b/xllm/models/dit/pipelines/pipeline_flux2_base.h new file mode 100644 index 0000000000..89c6c2429b --- /dev/null +++ b/xllm/models/dit/pipelines/pipeline_flux2_base.h @@ -0,0 +1,467 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once +#include +#include + +#include +#include +#include +#include +#include + +#include "core/framework/dit_model_loader.h" +#include "core/framework/model_context.h" +#include "core/framework/request/dit_request_state.h" +#include "core/framework/state_dict/state_dict.h" +#include "core/framework/state_dict/utils.h" +#include "models/dit/autoencoders/autoencoder_kl_flux2.h" +#include "models/dit/schedulers/flowmatch_euler_discrete_scheduler.h" +#include "models/dit/transformers/transformer_flux2.h" +#include "models/model_registry.h" + +namespace xllm { + +inline std::string SYSTEM_MESSAGE = + "You are an AI that reasons about image descriptions. You give structured " + "responses focusing on object relationships, object attribution and " + "actions " + "without speculation."; + +inline std::string SYSTEM_MESSAGE_UPSAMPLING_T2I = + "You are an expert prompt engineer for FLUX.2 by Black Forest Labs. " + "Rewrite user prompts to be more descriptive while strictly preserving " + "their " + "core subject and intent.\n" + "Guidelines:\n" + "1. Structure: Keep structured inputs structured (enhance within fields). " + "Convert natural language to detailed paragraphs.\n" + "2. Details: Add concrete visual specifics - form, scale, textures, " + "materials, lighting (quality, direction, color), shadows, spatial " + "relationships, and environmental context.\n" + "3. Text in Images: Put ALL text in quotation marks, matching the " + "prompt's language. Always provide explicit quoted text for objects that " + "would " + "contain text in reality (signs, labels, screens, etc.) - without it, " + "the model generates gibberish.\n" + "Output only the revised prompt and nothing else."; + +inline std::string SYSTEM_MESSAGE_UPSAMPLING_I2I = + "You are FLUX.2 by Black Forest Labs, an image-editing expert. You " + "convert editing requests into one concise instruction (50-80 words, ~30 " + "for " + "brief requests).\n" + "Rules:\n" + "- Single instruction only, no commentary\n" + "- Use clear, analytical language (avoid \"whimsical,\" \"cascading,\" " + "etc.)\n" + "- Specify what changes AND what stays the same (face, lighting, " + "composition)\n" + "- Reference actual image elements\n" + "- Turn negatives into positives (\"don't change X\" → \"keep X\")\n" + "- Make abstractions concrete (\"futuristic\" → \"glowing cyan neon, " + "metallic " + "panels\")\n" + "- Keep content PG-13\n" + "Output only the final instruction in plain text and nothing else."; + +inline std::vector>> +format_input(const std::vector& prompts, + const std::string& system_message = SYSTEM_MESSAGE, + const std::vector>& images = + std::vector>()) { + std::vector>> + messages_batch; + messages_batch.reserve(prompts.size()); + + for (const auto& prompt : prompts) { + std::vector> messages; + + messages.push_back( + {{"role", "system"}, + {"content", + "[{\"type\": \"text\", \"text\": \"" + system_message + "\"}]"}}); + + messages.push_back( + {{"role", "user"}, + {"content", "[{\"type\": \"text\", \"text\": \"" + prompt + "\"}]"}}); + + messages_batch.push_back(messages); + } + + return messages_batch; +} + +inline float compute_empirical_mu(int64_t image_seq_len, int64_t num_steps) { + double a1 = 8.73809524e-05, b1 = 1.89833333; + double a2 = 0.00016927, b2 = 0.45666666; + + double mu; + if (image_seq_len > 4300) { + mu = a2 * image_seq_len + b2; + return static_cast(mu); + } + + double m_200 = a2 * image_seq_len + b2; + double m_10 = a1 * image_seq_len + b1; + + double a = (m_200 - m_10) / 190.0; + double b = m_200 - 200.0 * a; + mu = a * num_steps + b; + + return static_cast(mu); +} + +inline std::pair flux2_retrieve_timesteps( + FlowMatchEulerDiscreteScheduler scheduler, + int64_t num_inference_steps = 0, + torch::Device device = torch::kCPU, + std::optional> sigmas = std::nullopt, + std::optional mu = std::nullopt) { + torch::Tensor scheduler_timesteps; + int64_t steps; + if (sigmas.has_value()) { + steps = sigmas->size(); + scheduler->set_timesteps( + static_cast(steps), device, *sigmas, mu, std::nullopt); + + scheduler_timesteps = scheduler->timesteps(); + } else { + steps = num_inference_steps; + scheduler->set_timesteps( + static_cast(steps), device, std::nullopt, mu, std::nullopt); + scheduler_timesteps = scheduler->timesteps(); + } + if (scheduler_timesteps.device() != device) { + scheduler_timesteps = scheduler_timesteps.to(device); + } + return {scheduler_timesteps, steps}; +} + +class Flux2PosEmbedImpl final : public torch::nn::Module { + public: + Flux2PosEmbedImpl(int64_t theta, std::vector axes_dim) { + theta_ = theta; + axes_dim_ = axes_dim; + } + + std::pair forward_cache( + const torch::Tensor& txt_ids, + const torch::Tensor& img_ids, + int64_t height = -1, + int64_t width = -1) { + int64_t seq_len = txt_ids.size(0); + + if (height != cached_image_height_ || width != cached_image_width_ || + seq_len != max_seq_len_) { + torch::Tensor ids = torch::cat({txt_ids, img_ids}, 1); + cached_image_height_ = height; + cached_image_width_ = width; + max_seq_len_ = seq_len; + auto [cos, sin] = forward(ids); + freqs_cos_cache_ = std::move(cos); + freqs_sin_cache_ = std::move(sin); + } + return {freqs_cos_cache_, freqs_sin_cache_}; + } + std::pair forward(const torch::Tensor& ids) { + int64_t n_axes = axes_dim_.size(); + std::vector cos_out, sin_out; + auto pos = ids.to(torch::kFloat32); + torch::Dtype freqs_dtype = torch::kFloat64; + for (int64_t i = 0; i < n_axes; ++i) { + auto pos_slice = pos.select(-1, i).squeeze(0); + auto result = get_1d_rotary_pos_embed( + axes_dim_[i], pos_slice, theta_, true, 1, 1, true, freqs_dtype); + auto cos = result[0]; + auto sin = result[1]; + cos_out.push_back(cos); + sin_out.push_back(sin); + } + + auto freqs_cos = torch::cat(cos_out, -1); + auto freqs_sin = torch::cat(sin_out, -1); + return {freqs_cos, freqs_sin}; + } + + private: + int64_t theta_; + std::vector axes_dim_; + torch::Tensor freqs_cos_cache_; + torch::Tensor freqs_sin_cache_; + int64_t max_seq_len_ = -1; + int64_t cached_image_height_ = -1; + int64_t cached_image_width_ = -1; +}; +TORCH_MODULE(Flux2PosEmbed); + +class Flux2PipelineBaseImpl : public torch::nn::Module { + protected: + torch::Tensor get_mistral3_prompt_embeds( + std::vector& prompt, + int64_t num_images_per_prompt = 1, + int64_t max_sequence_length = 512, + const std::vector& hidden_states_layers = {10, 20, 30}, + const std::string& system_message = SYSTEM_MESSAGE) { + int64_t batch_size = prompt.size(); + + auto messages_batch = format_input(prompt, system_message); + + std::vector formatted_prompts; + formatted_prompts.reserve(batch_size); + + std::vector> text_input_ids; + text_input_ids.reserve(batch_size); + CHECK(tokenizer_->batch_encode(formatted_prompts, &text_input_ids)); + for (auto& ids : text_input_ids) { + ids.resize(max_sequence_length, 0); + } + + std::vector text_input_ids_flat; + text_input_ids_flat.reserve(batch_size * max_sequence_length); + for (const auto& ids : text_input_ids) { + text_input_ids_flat.insert( + text_input_ids_flat.end(), ids.begin(), ids.end()); + } + auto input_ids = + torch::tensor(text_input_ids_flat, torch::dtype(torch::kLong)) + .view({batch_size, max_sequence_length}) + .to(options_.device()); + + // TODO: implement actual text encoder forward call, e.g.: + // auto prompt_embeds = hidden_states_output... + LOG(FATAL) << "Text encoder is not yet implemented. " + "Please provide pre-computed prompt_embeds directly."; + + return {}; + } + + torch::Tensor prepare_text_ids(const torch::Tensor& prompt_embeds) { + int64_t batch_size = prompt_embeds.size(0); + int64_t seq_len = prompt_embeds.size(1); + + std::vector out_ids; + out_ids.reserve(batch_size); + auto int_opts = torch::TensorOptions().dtype(torch::kInt64); + for (int64_t i = 0; i < batch_size; ++i) { + auto t = torch::arange(1, int_opts); + auto h = torch::arange(1, int_opts); + auto w = torch::arange(1, int_opts); + auto l = torch::arange(seq_len, int_opts); + auto grid = torch::meshgrid({t, h, w, l}, "ij"); + auto coords = torch::stack({grid[0].flatten(), + grid[1].flatten(), + grid[2].flatten(), + grid[3].flatten()}, + -1); + out_ids.push_back(coords); + } + + auto text_ids = torch::stack(out_ids, 0); + return text_ids; + } + + std::tuple encode_prompt( + std::optional> prompt, + std::optional prompt_embeds, + int64_t num_images_per_prompt = 1, + int64_t max_sequence_length = 512, + const std::vector& hidden_states_layers = {10, 20, 30}, + const std::string& system_message = SYSTEM_MESSAGE) { + std::vector prompt_list; + if (prompt.has_value()) { + prompt_list = prompt.value(); + } + if (prompt_list.empty()) { + prompt_list = {""}; + } + if (!prompt_embeds.has_value()) { + prompt_embeds = get_mistral3_prompt_embeds(prompt_list, + num_images_per_prompt, + max_sequence_length, + hidden_states_layers); + } + torch::Tensor text_ids = prepare_text_ids(prompt_embeds.value()); + + return std::make_tuple(prompt_embeds.value(), text_ids); + } + + torch::Tensor prepare_latent_image_ids(const torch::Tensor& latents) { + int64_t batch_size = latents.size(0); + int64_t num_channels = latents.size(1); + int64_t height = latents.size(2); + int64_t width = latents.size(3); + + auto t = torch::arange(1, options_); + auto h = torch::arange(height, options_); + auto w = torch::arange(width, options_); + auto l = torch::arange(1, options_); + + auto grid = torch::meshgrid({t, h, w, l}, "ij"); + auto coords = torch::stack({grid[0].flatten(), + grid[1].flatten(), + grid[2].flatten(), + grid[3].flatten()}, + -1); + + auto latent_image_ids = coords.unsqueeze(0).expand({batch_size, -1, -1}); + return latent_image_ids; + } + + torch::Tensor pack_latents(const torch::Tensor& latents) { + int64_t batch_size = latents.size(0); + int64_t num_channels = latents.size(1); + int64_t height = latents.size(2); + int64_t width = latents.size(3); + + torch::Tensor latents_packed = + latents.reshape({batch_size, num_channels, height * width}); + latents_packed = latents_packed.permute({0, 2, 1}); + + return latents_packed; + } + + torch::Tensor patchify_latents(const torch::Tensor& latents) { + int64_t batch_size = latents.size(0); + int64_t num_channels_latents = latents.size(1); + int64_t height = latents.size(2); + int64_t width = latents.size(3); + + torch::Tensor latents_patched = latents.view( + {batch_size, num_channels_latents, height / 2, 2, width / 2, 2}); + latents_patched = latents_patched.permute({0, 1, 3, 5, 2, 4}); + latents_patched = latents_patched.reshape( + {batch_size, num_channels_latents * 4, height / 2, width / 2}); + + return latents_patched; + } + + torch::Tensor unpatchify_latents(const torch::Tensor& latents) { + int64_t batch_size = latents.size(0); + int64_t num_channels_latents = latents.size(1); + int64_t height = latents.size(2); + int64_t width = latents.size(3); + + torch::Tensor latents_unpatched = latents.reshape( + {batch_size, num_channels_latents / (2 * 2), 2, 2, height, width}); + latents_unpatched = latents_unpatched.permute({0, 1, 4, 2, 5, 3}); + latents_unpatched = latents_unpatched.reshape( + {batch_size, num_channels_latents / (2 * 2), height * 2, width * 2}); + + return latents_unpatched; + } + + torch::Tensor unpack_latents(const torch::Tensor& latents, + int64_t height, + int64_t width, + int64_t vae_scale_factor) { + int64_t batch_size = latents.size(0); + int64_t num_patches = latents.size(1); + int64_t channels = latents.size(2); + height = 2 * (height / (vae_scale_factor_ * 2)); + width = 2 * (width / (vae_scale_factor_ * 2)); + + torch::Tensor latents_unpacked = + latents.view({batch_size, height / 2, width / 2, channels / 4, 2, 2}); + latents_unpacked = latents_unpacked.permute({0, 3, 1, 4, 2, 5}); + latents_unpacked = latents_unpacked.reshape( + {batch_size, channels / (2 * 2), height, width}); + + return latents_unpacked; + } + + torch::Tensor unpack_latents_with_ids(const torch::Tensor& latents, + const torch::Tensor& latent_ids) { + int64_t batch_size = latents.size(0); + int64_t seq_len = latents.size(1); + int64_t channels = latents.size(2); + + std::vector x_list; + for (int64_t i = 0; i < batch_size; ++i) { + torch::Tensor data = latents[i]; + torch::Tensor pos = latent_ids[i]; + + torch::Tensor h_ids = pos.select(1, 1).to(torch::kInt64); + torch::Tensor w_ids = pos.select(1, 2).to(torch::kInt64); + + int64_t h = h_ids.max().item() + 1; + int64_t w = w_ids.max().item() + 1; + + torch::Tensor flat_ids = h_ids * w + w_ids; + + torch::Tensor out = torch::zeros({h * w, channels}, data.options()); + out.scatter_(0, flat_ids.unsqueeze(1).expand({-1, channels}), data); + + out = out.view({h, w, channels}).permute({2, 0, 1}); + x_list.push_back(out); + } + + return torch::stack(x_list, 0); + } + + torch::Tensor _prepare_image_ids( + const std::vector& image_latents, + int64_t scale = 10) { + if (image_latents.empty()) { + LOG(FATAL) << "image_latents list cannot be empty!"; + } + + int64_t num_latents = image_latents.size(); + torch::Tensor t_indices = torch::arange(0, num_latents, torch::kInt64); + torch::Tensor t_coords = scale + scale * t_indices; + t_coords = t_coords.unsqueeze(1); + + std::vector image_latent_ids; + for (int64_t i = 0; i < num_latents; ++i) { + torch::Tensor x = image_latents[i]; + x = x.squeeze(0); + if (x.dim() != 3) { + throw std::invalid_argument( + "Each image latent must be 3D (C, H, W) or 4D (1, C, H, W), got " + + std::to_string(x.dim()) + "D tensor!"); + } + + int64_t height = x.size(1); + int64_t width = x.size(2); + + torch::Tensor h_coords = torch::arange(0, height, torch::kInt64); + torch::Tensor w_coords = torch::arange(0, width, torch::kInt64); + torch::Tensor l_coords = torch::arange(0, 1, torch::kInt64); + + torch::Tensor t = t_coords[i].expand({1}); + auto t_exp = t.repeat({height * width * 1}); + auto h_exp = h_coords.repeat_interleave(width * 1).repeat({1}); + auto w_exp = w_coords.repeat({height}).repeat_interleave(1); + auto l_exp = l_coords.repeat({height * width}).repeat({1}); + + torch::Tensor coords = torch::stack({t_exp, h_exp, w_exp, l_exp}, 1); + image_latent_ids.push_back(coords); + } + + torch::Tensor combined_coords = torch::cat(image_latent_ids, 0); + combined_coords = combined_coords.unsqueeze(0); + return combined_coords; + } + + protected: + torch::Device device_ = torch::kCPU; + torch::ScalarType dtype_; + std::unique_ptr tokenizer_; + torch::TensorOptions options_; + int32_t tokenizer_max_length_; + int32_t vae_scale_factor_; +}; +} // namespace xllm diff --git a/xllm/models/dit/processors/flux2_image_processor.h b/xllm/models/dit/processors/flux2_image_processor.h new file mode 100644 index 0000000000..8b48b13bf5 --- /dev/null +++ b/xllm/models/dit/processors/flux2_image_processor.h @@ -0,0 +1,153 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once +#include +#include +#include + +#include "models/dit/autoencoders/autoencoder_kl.h" +#include "torch/torch.h" + +namespace xllm { + +class Flux2ImageProcessorImpl final : public VAEImageProcessorImpl { + public: + explicit Flux2ImageProcessorImpl(const ModelContext& context, + int64_t vae_scale_factor = 16) + : VAEImageProcessorImpl(context, true, true, false, true, false, 32), + vae_scale_factor_(vae_scale_factor), + vae_latent_channels_(32) {} + + torch::Tensor check_image_input(const torch::Tensor& image, + int64_t max_aspect_ratio = 8, + int64_t min_side_length = 64, + int64_t max_area = 1024 * 1024) { + if (image.dim() != 3 && image.dim() != 4) { + LOG(FATAL) << "Image must be 3D (C, H, W) or 4D (B, C, H, W), got dim: " + << image.dim(); + } + + int64_t height = image.size(-2); + int64_t width = image.size(-1); + + if (width < min_side_length || height < min_side_length) { + LOG(FATAL) << "Image too small: " << width << "x" << height + << ". Both dimensions must be at least " << min_side_length + << "px"; + } + + float aspect_ratio = std::max(static_cast(width) / height, + static_cast(height) / width); + if (aspect_ratio > max_aspect_ratio) { + LOG(FATAL) << "Aspect ratio too extreme: " << width << "x" << height + << " (ratio: " << aspect_ratio + << ":1). Maximum allowed ratio is " << max_aspect_ratio + << ":1"; + } + + return image; + } + + torch::Tensor resize_to_target_area(const torch::Tensor& image, + int64_t target_area = 1024 * 1024) { + int64_t image_width = image.size(-1); + int64_t image_height = image.size(-2); + + float scale = std::sqrt(static_cast(target_area) / + static_cast(image_width * image_height)); + int64_t width = static_cast(image_width * scale); + int64_t height = static_cast(image_height * scale); + return torch::nn::functional::interpolate( + image.unsqueeze(0), // [1, 3, H, W] + torch::nn::functional::InterpolateFuncOptions() + .mode(torch::kBilinear) + .align_corners(false) + .size(std::vector{height, width})) + .squeeze(0); + } + + torch::Tensor resize_if_exceeds_area(const torch::Tensor& image, + int64_t target_area = 1024 * 1024) { + int64_t image_width = image.size(-1); + int64_t image_height = image.size(-2); + int64_t pixel_count = image_width * image_height; + + if (pixel_count <= target_area) { + return image; + } + + return resize_to_target_area(image, target_area); + } + + torch::Tensor resize_and_crop(const torch::Tensor& image, + int64_t width, + int64_t height) { + int64_t image_width = image.size(-1); + int64_t image_height = image.size(-2); + + int64_t left = (image_width - width) / 2; + int64_t top = (image_height - height) / 2; + int64_t right = left + width; + int64_t bottom = top + height; + + return image.slice(-2, top, bottom).slice(-1, left, right); + } + + torch::Tensor concatenate_images(const std::vector& images) { + if (images.empty()) { + LOG(FATAL) << "Cannot concatenate empty image list"; + } + + if (images.size() == 1) { + return images[0].clone(); + } + + int64_t total_width = 0; + int64_t max_height = 0; + + for (const auto& img : images) { + total_width += img.size(-1); + max_height = std::max(max_height, img.size(-2)); + } + + torch::TensorOptions options = images[0].options(); + auto background_color = torch::full({1, 3, max_height, 1}, 1.0f, options); + + auto new_img = torch::full({1, 3, max_height, total_width}, 1.0f, options); + + int64_t x_offset = 0; + for (const auto& img : images) { + int64_t img_height = img.size(-2); + int64_t img_width = img.size(-1); + + int64_t y_offset = (max_height - img_height) / 2; + + new_img.slice(-1, x_offset, x_offset + img_width) + .slice(-2, y_offset, y_offset + img_height) + .copy_(img); + + x_offset += img_width; + } + + return new_img.squeeze(0); + } + + private: + int64_t vae_scale_factor_; + int64_t vae_latent_channels_; +}; +TORCH_MODULE(Flux2ImageProcessor); +} // namespace xllm diff --git a/xllm/models/dit/transformers/transformer_flux.h b/xllm/models/dit/transformers/transformer_flux.h index 9a16c96835..4f4adf945c 100644 --- a/xllm/models/dit/transformers/transformer_flux.h +++ b/xllm/models/dit/transformers/transformer_flux.h @@ -902,7 +902,8 @@ TORCH_MODULE(AdaLayerNormZeroSingle); class AdaLayerNormContinuousImpl : public torch::nn::Module { public: - explicit AdaLayerNormContinuousImpl(ModelContext context) + explicit AdaLayerNormContinuousImpl(ModelContext context, + bool with_bias = true) : options_(context.get_tensor_options()) { ModelArgs model_args = context.get_model_args(); auto num_attention_heads = model_args.n_heads(); @@ -913,7 +914,7 @@ class AdaLayerNormContinuousImpl : public torch::nn::Module { linear_ = register_module("linear", layer::AddMatmul(conditioning_embedding_dim, 2 * embedding_dim, - /*with_bias=*/true, + /*with_bias=*/with_bias, options_)); norm_ = register_module( "norm", diff --git a/xllm/models/dit/transformers/transformer_flux2.h b/xllm/models/dit/transformers/transformer_flux2.h new file mode 100644 index 0000000000..a9f2d60946 --- /dev/null +++ b/xllm/models/dit/transformers/transformer_flux2.h @@ -0,0 +1,1513 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +// sequence parallel pad manager +#include "core/framework/config/dit_config.h" +#include "core/framework/config/parallel_config.h" +#include "core/framework/dit_cache/dit_cache.h" +#include "core/framework/dit_model_loader.h" +#include "core/framework/model/model_input_params.h" +#include "core/framework/parallel_state/parallel_args.h" +#include "core/framework/state_dict/state_dict.h" +#include "core/framework/state_dict/utils.h" +#include "core/layers/common/add_matmul.h" +#include "core/layers/common/rms_norm.h" +#include "framework/model_context.h" +#include "framework/parallel_state/parallel_state.h" +#include "models/dit/transformers/transformer_flux.h" +#include "models/dit/utils/dit_parallel_linear.h" +#include "models/model_registry.h" +#if defined(USE_NPU) +#include "torch_npu/csrc/aten/CustomFunctions.h" +#endif + +namespace xllm { +using dit::DiTParallelLinear; +using dit::TpOptions; + +class Flux2SwiGLUImpl final : public torch::nn::Module { + public: + Flux2SwiGLUImpl() { gate_fn_ = torch::nn::SiLU(); } + + torch::Tensor forward(torch::Tensor x) { + auto chunks = torch::chunk(x, 2, /*dim=*/-1); + torch::Tensor x1 = chunks[0]; + torch::Tensor x2 = chunks[1]; + + torch::Tensor x_out = gate_fn_(x1) * x2; + return x_out; + } + + private: + torch::nn::SiLU gate_fn_; +}; +TORCH_MODULE(Flux2SwiGLU); + +class Flux2FeedForwardImpl final : public torch::nn::Module { + public: + explicit Flux2FeedForwardImpl(const ModelContext& context, + const ParallelArgs& parallel_args) + : options_(context.get_tensor_options()), parallel_args_(parallel_args) { + auto model_args = context.get_model_args(); + float eps = model_args.mlp_ratio(); + int64_t num_attention_heads = model_args.n_heads(); + int64_t attention_head_dim = model_args.head_dim(); + int64_t inner_dim = num_attention_heads * attention_head_dim; + std::optional tp_options = std::nullopt; + + if (::xllm::ParallelConfig::get_instance().tp_size() > 1) { + tp_options = TpOptions( + /*column_parallel=*/true, + /*gather_output=*/false, + /*need_scatter=*/false, + /*process_group=*/parallel_args_.dit_tp_group_); + + if (!parallel_args_.dit_tp_group_ || !tp_options.has_value() || + !tp_options->process_group) { + } + } + + auto linear_in = DiTParallelLinear( + inner_dim, inner_dim * 6, false, options_, std::nullopt, tp_options); + linear_in_ = register_module("linear_in", linear_in); + act_fn_ = register_module("act_fn", Flux2SwiGLU()); + + std::optional tp_out_options = std::nullopt; + + if (::xllm::ParallelConfig::get_instance().tp_size() > 1) { + tp_out_options = TpOptions( + /*column_parallel=*/false, + /*gather_output=*/false, + /*need_scatter=*/false, + /*process_group=*/parallel_args_.dit_tp_group_); + } + + auto linear_out = DiTParallelLinear(inner_dim * 3, + inner_dim, + false, + options_, + std::nullopt, + tp_out_options); + linear_out_ = register_module("linear_out", linear_out); + } + + torch::Tensor forward(const torch::Tensor& hidden_states, + const bool is_save) { + auto out = linear_in_->forward(hidden_states); + out = act_fn_->forward(out); + out = linear_out_->forward(out); + return out; + } + + void load_state_dict(const StateDict& state_dict) { + int64_t tp_size = ::xllm::ParallelConfig::get_instance().tp_size() > 1 + ? ::xllm::ParallelConfig::get_instance().tp_size() + : 1; + if (tp_size > 1) { + int64_t tp_rank = parallel_args_.dit_tp_group_->rank(); + + // Manually reorder linear_in weight: split gate/up separately, + // take rank-th slice of each, then recombine so each rank has + // matching gate-up pairs for SwiGLU. + auto linear_in_weight = state_dict.get_tensor("linear_in.weight"); + if (linear_in_weight.defined()) { + // Full weight shape: [inner_dim * 6, inner_dim] + // Rows [0, inner_dim*3) = gate projection + // Rows [inner_dim*3, inner_dim*6) = up projection + int64_t gate_up_dim = linear_in_weight.size(0) / 2; + auto gate_weight = linear_in_weight.slice(0, 0, gate_up_dim); + auto up_weight = linear_in_weight.slice(0, gate_up_dim); + + auto gate_chunks = gate_weight.chunk(tp_size, 0); + auto up_chunks = up_weight.chunk(tp_size, 0); + + auto gate_rank = gate_chunks[tp_rank]; + auto up_rank = up_chunks[tp_rank]; + + auto rank_weight = torch::cat({gate_rank, up_rank}, 0).to(options_); + + { + torch::NoGradGuard no_grad; + for (auto& p : linear_in_->named_parameters()) { + if (p.key() == "weight") { + p.value().copy_(rank_weight); + break; + } + } + } + } + } else { + linear_in_->as()->load_state_dict( + state_dict.get_dict_with_prefix("linear_in.")); + } + + // linear_out uses default load_state_dict which handles + // row split (dim 1) correctly for each rank. + linear_out_->as()->load_state_dict( + state_dict.get_dict_with_prefix("linear_out.")); + } + + void verify_loaded_weights(const std::string& prefix) const { + int64_t tp_size = ::xllm::ParallelConfig::get_instance().tp_size() > 1 + ? ::xllm::ParallelConfig::get_instance().tp_size() + : 1; + if (tp_size <= 1) { + linear_in_->as()->verify_loaded_weights(prefix + + "linear_in."); + } + linear_out_->as()->verify_loaded_weights(prefix + + "linear_out."); + } + + private: + int64_t dim_; + int64_t inner_dim_; + float mult_; + DiTParallelLinear linear_in_{nullptr}; + Flux2SwiGLU act_fn_{nullptr}; + DiTParallelLinear linear_out_{nullptr}; + torch::TensorOptions options_; + ParallelArgs parallel_args_; +}; +TORCH_MODULE(Flux2FeedForward); + +class Flux2AttentionImpl : public torch::nn::Module { + public: + explicit Flux2AttentionImpl(const ModelContext& context, + const ParallelArgs& parallel_args) + : options_(context.get_tensor_options()), parallel_args_(parallel_args) { + auto model_args = context.get_model_args(); + heads_ = model_args.n_heads(); + head_dim_ = model_args.head_dim(); + query_dim_ = heads_ * head_dim_; + out_dim_ = query_dim_; + added_kv_proj_dim_ = query_dim_; + + // Determine parallelism strategy + std::optional tp_options = std::nullopt; + + if (::xllm::ParallelConfig::get_instance().tp_size() > 1) { + tp_options = TpOptions( + /*column_parallel=*/true, + /*gather_output=*/false, + /*need_scatter=*/false, + /*process_group=*/parallel_args_.dit_tp_group_); + } + + auto to_q = DiTParallelLinear( + query_dim_, out_dim_, false, options_, std::nullopt, tp_options); + to_q_ = register_module("to_q", to_q); + + auto to_k = DiTParallelLinear( + query_dim_, out_dim_, false, options_, std::nullopt, tp_options); + to_k_ = register_module("to_k", to_k); + + auto to_v = DiTParallelLinear( + query_dim_, out_dim_, false, options_, std::nullopt, tp_options); + to_v_ = register_module("to_v", to_v); + + norm_q_ = + register_module("norm_q", layer::RMSNorm(head_dim_, 1e-6f, options_)); + norm_k_ = + register_module("norm_k", layer::RMSNorm(head_dim_, 1e-6f, options_)); + + std::optional tp_out_options = std::nullopt; + + if (::xllm::ParallelConfig::get_instance().tp_size() > 1) { + tp_out_options = TpOptions( + /*column_parallel=*/false, + /*gather_output=*/false, + /*need_scatter=*/false, + /*process_group=*/parallel_args_.dit_tp_group_); + } + + auto to_out = DiTParallelLinear( + out_dim_, query_dim_, false, options_, std::nullopt, tp_out_options); + to_out_ = register_module("to_out", to_out); + + // ── Added KV projections (cross-attention for text stream) ── + if (added_kv_proj_dim_ > 0) { + std::optional tp_add_out_options = std::nullopt; + + if (::xllm::ParallelConfig::get_instance().tp_size() > 1) { + tp_add_out_options = TpOptions( + /*column_parallel=*/false, + /*gather_output=*/false, + /*need_scatter=*/false, + /*process_group=*/parallel_args_.dit_tp_group_); + } + + auto to_add_out = DiTParallelLinear(out_dim_, + added_kv_proj_dim_, + false, + options_, + std::nullopt, + tp_add_out_options); + to_add_out_ = register_module("to_add_out", to_add_out); + norm_added_q_ = register_module( + "norm_added_q", layer::RMSNorm(head_dim_, 1e-6f, options_)); + norm_added_k_ = register_module( + "norm_added_k", layer::RMSNorm(head_dim_, 1e-6f, options_)); + + auto to_add_q = DiTParallelLinear(added_kv_proj_dim_, + out_dim_, + false, + options_, + std::nullopt, + tp_options); + to_add_q_ = register_module("to_add_q", to_add_q); + + auto to_add_k = DiTParallelLinear(added_kv_proj_dim_, + out_dim_, + false, + options_, + std::nullopt, + tp_options); + to_add_k_ = register_module("to_add_k", to_add_k); + + auto to_add_v = DiTParallelLinear(added_kv_proj_dim_, + out_dim_, + false, + options_, + std::nullopt, + tp_options); + to_add_v_ = register_module("to_add_v", to_add_v); + } + } + + std::tuple forward( + const torch::Tensor& hidden_states, + const torch::Tensor& encoder_hidden_states, + const torch::Tensor& image_rotary_emb) { + int64_t input_ndim = hidden_states.dim(); + torch::Tensor hidden_states_reshaped = hidden_states; + + if (input_ndim == 4) { + auto shape = hidden_states.sizes(); + int64_t batch_size = shape[0]; + int64_t channel = shape[1]; + int64_t height = shape[2]; + int64_t width = shape[3]; + hidden_states_reshaped = + hidden_states.view({batch_size, channel, height * width}) + .transpose(1, 2); + } + + int64_t context_input_ndim = encoder_hidden_states.dim(); + torch::Tensor encoder_hidden_states_reshaped = encoder_hidden_states; + if (context_input_ndim == 4) { + auto shape = encoder_hidden_states.sizes(); + int64_t batch_size = shape[0]; + int64_t channel = shape[1]; + int64_t height = shape[2]; + int64_t width = shape[3]; + encoder_hidden_states_reshaped = + encoder_hidden_states.view({batch_size, channel, height * width}) + .transpose(1, 2); + } + + int64_t batch_size = encoder_hidden_states_reshaped.size(0); + + hidden_states_reshaped = hidden_states_reshaped.squeeze(1).squeeze(1); + encoder_hidden_states_reshaped = + encoder_hidden_states_reshaped.squeeze(1).squeeze(1); + + auto query = to_q_->forward(hidden_states_reshaped); + auto key = to_k_->forward(hidden_states_reshaped); + auto value = to_v_->forward(hidden_states_reshaped); + int64_t inner_dim = key.size(-1); + int64_t attn_heads = inner_dim / head_dim_; + + int64_t head_dim = head_dim_; + query = query.view({batch_size, -1, attn_heads, head_dim}); + key = key.view({batch_size, -1, attn_heads, head_dim}); + value = value.view({batch_size, -1, attn_heads, head_dim}); + if (norm_q_) { + query = std::get<0>(norm_q_(query)); + } + if (norm_k_) { + key = std::get<0>(norm_k_(key)); + } + auto encoder_hidden_states_query_proj = + to_add_q_->forward(encoder_hidden_states_reshaped); + auto encoder_hidden_states_key_proj = + to_add_k_->forward(encoder_hidden_states_reshaped); + auto encoder_hidden_states_value_proj = + to_add_v_->forward(encoder_hidden_states_reshaped); + encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( + {batch_size, -1, attn_heads, head_dim}); + encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view( + {batch_size, -1, attn_heads, head_dim}); + encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( + {batch_size, -1, attn_heads, head_dim}); + if (norm_added_q_) { + encoder_hidden_states_query_proj = + std::get<0>(norm_added_q_(encoder_hidden_states_query_proj)); + } + + if (norm_added_k_) { + encoder_hidden_states_key_proj = + std::get<0>(norm_added_k_(encoder_hidden_states_key_proj)); + } + + // Concatenate for joint attention: order [text, image] + auto query1 = torch::cat({encoder_hidden_states_query_proj, query}, 1); + auto key1 = torch::cat({encoder_hidden_states_key_proj, key}, 1); + auto value1 = torch::cat({encoder_hidden_states_value_proj, value}, 1); + if (image_rotary_emb.defined()) { + query1 = apply_rotary_emb(query1, image_rotary_emb, false); + key1 = apply_rotary_emb(key1, image_rotary_emb, false); + } + + // After SP all2all (before_attention=true), attn_heads already equals + // heads/sp_size. No further division needed. + int64_t local_heads = attn_heads; + +#if defined(USE_NPU) + int64_t head_num_ = query1.size(2); + int64_t head_dim_ = query1.size(-1); + auto results = + at_npu::native::custom_ops::npu_fusion_attention(query1, + key1, + value1, + local_heads, + "BSND", + torch::nullopt, + torch::nullopt, + torch::nullopt, + pow(head_dim_, -0.5), + 1.0, + 65535, + 65535); + auto attn_output = std::get<0>(results); + + attn_output = attn_output.reshape({batch_size, -1, local_heads * head_dim}); + +#elif defined(USE_CUDA) + query1 = query1.transpose(1, 2); + key1 = key1.transpose(1, 2); + value1 = value1.transpose(1, 2); + torch::Tensor attn_output = torch::scaled_dot_product_attention( + query1, key1, value1, torch::nullopt, 0.0, false); + attn_output = attn_output.transpose(1, 2).reshape( + {batch_size, -1, local_heads * head_dim}); +#else + NOT_IMPLEMENTED(); +#endif + + attn_output = attn_output.to(query.dtype()); + int64_t encoder_length = encoder_hidden_states_query_proj.size(1); + torch::Tensor encoder_output = attn_output.slice(1, 0, encoder_length); + torch::Tensor hidden_output = attn_output.slice(1, encoder_length); + encoder_output = encoder_output.flatten(2); + hidden_output = hidden_output.flatten(2); + + hidden_output = to_out_->forward(hidden_output); + encoder_output = to_add_out_->forward(encoder_output); + + return std::make_tuple(hidden_output, encoder_output); + } + + void load_state_dict(const StateDict& state_dict) { + to_q_->as()->load_state_dict( + state_dict.get_dict_with_prefix("to_q.")); + to_k_->as()->load_state_dict( + state_dict.get_dict_with_prefix("to_k.")); + to_v_->as()->load_state_dict( + state_dict.get_dict_with_prefix("to_v.")); + norm_q_->load_state_dict(state_dict.get_dict_with_prefix("norm_q.")); + norm_k_->load_state_dict(state_dict.get_dict_with_prefix("norm_k.")); + + to_out_->as()->load_state_dict( + state_dict.get_dict_with_prefix("to_out.0.")); + + if (added_kv_proj_dim_ > 0) { + norm_added_q_->load_state_dict( + state_dict.get_dict_with_prefix("norm_added_q.")); + norm_added_k_->load_state_dict( + state_dict.get_dict_with_prefix("norm_added_k.")); + to_add_q_->as()->load_state_dict( + state_dict.get_dict_with_prefix("add_q_proj.")); + to_add_k_->as()->load_state_dict( + state_dict.get_dict_with_prefix("add_k_proj.")); + to_add_v_->as()->load_state_dict( + state_dict.get_dict_with_prefix("add_v_proj.")); + to_add_out_->as()->load_state_dict( + state_dict.get_dict_with_prefix("to_add_out.")); + } + } + + void verify_loaded_weights(const std::string& prefix) const { + to_q_->as()->verify_loaded_weights(prefix + "to_q."); + to_k_->as()->verify_loaded_weights(prefix + "to_k."); + to_v_->as()->verify_loaded_weights(prefix + "to_v."); + + to_out_->as()->verify_loaded_weights(prefix + + "to_out.0."); + if (added_kv_proj_dim_ > 0) { + to_add_q_->as()->verify_loaded_weights(prefix + + "add_q_proj."); + to_add_k_->as()->verify_loaded_weights(prefix + + "add_k_proj."); + to_add_v_->as()->verify_loaded_weights(prefix + + "add_v_proj."); + to_add_out_->as()->verify_loaded_weights( + prefix + "to_add_out."); + } + } + + private: + int64_t heads_; + int64_t head_dim_; + int64_t query_dim_; + int64_t out_dim_; + int64_t added_kv_proj_dim_; + DiTParallelLinear to_q_{nullptr}; + DiTParallelLinear to_k_{nullptr}; + DiTParallelLinear to_v_{nullptr}; + layer::RMSNorm norm_q_{nullptr}; + layer::RMSNorm norm_k_{nullptr}; + DiTParallelLinear to_out_{nullptr}; + layer::RMSNorm norm_added_q_{nullptr}; + layer::RMSNorm norm_added_k_{nullptr}; + DiTParallelLinear to_add_out_{nullptr}; + DiTParallelLinear to_add_q_{nullptr}; + DiTParallelLinear to_add_k_{nullptr}; + DiTParallelLinear to_add_v_{nullptr}; + torch::TensorOptions options_; + ParallelArgs parallel_args_; +}; +TORCH_MODULE(Flux2Attention); + +class Flux2ModulationImpl : public torch::nn::Module { + public: + explicit Flux2ModulationImpl(const ModelContext& context, + int64_t dim, + int64_t mod_param_sets, + bool bias = false) + : options_(context.get_tensor_options()) { + auto model_args = context.get_model_args(); + + linear_ = register_module( + "linear", + layer::AddMatmul(dim, dim * 3 * mod_param_sets, bias, options_)); + act_fn_ = torch::nn::SiLU(); + } + + torch::Tensor forward(const torch::Tensor& temb) { + auto mod = act_fn_->forward(temb); + mod = linear_->forward(mod); + return mod; + } + + void load_state_dict(const StateDict& state_dict) { + linear_->load_state_dict(state_dict.get_dict_with_prefix("linear.")); + } + + void verify_loaded_weights(const std::string& prefix) const { + linear_->verify_loaded_weights(prefix + "linear."); + } + + static std::vector> + split(const torch::Tensor& mod, int64_t mod_param_sets) { + torch::Tensor mod_reshaped; + if (mod.dim() == 2) { + mod_reshaped = mod.unsqueeze(1); + } else { + mod_reshaped = mod; + } + + auto mod_params = torch::chunk(mod_reshaped, 3 * mod_param_sets, -1); + + std::vector> result; + for (int64_t i = 0; i < mod_param_sets; ++i) { + int64_t start_idx = 3 * i; + auto param_tuple = std::make_tuple(mod_params[start_idx], + mod_params[start_idx + 1], + mod_params[start_idx + 2]); + result.push_back(param_tuple); + } + + return result; + } + + private: + layer::AddMatmul linear_{nullptr}; + torch::nn::SiLU act_fn_{nullptr}; + torch::TensorOptions options_; +}; +TORCH_MODULE(Flux2Modulation); + +class Flux2TimestepEmbeddingImpl : public torch::nn::Module { + public: + Flux2TimestepEmbeddingImpl(int64_t in_channels, + int64_t time_embed_dim, + int64_t out_dim = -1, + bool sample_proj_bias = true) + : options_(torch::dtype(torch::kFloat32)) { + linear_1_ = register_module( + "linear_1", + layer::AddMatmul( + in_channels, time_embed_dim, sample_proj_bias, options_)); + + act_ = register_module("act", torch::nn::SiLU()); + + int64_t time_embed_dim_out = (out_dim > 0) ? out_dim : time_embed_dim; + linear_2_ = register_module( + "linear_2", + layer::AddMatmul( + time_embed_dim, time_embed_dim_out, sample_proj_bias, options_)); + } + + torch::Tensor forward(const torch::Tensor& sample) { + torch::Tensor result = sample; + + result = linear_1_->forward(result); + + if (act_) { + result = act_->forward(result); + } + + result = linear_2_->forward(result); + return result; + } + + void load_state_dict(const StateDict& state_dict) { + linear_1_->load_state_dict(state_dict.get_dict_with_prefix("linear_1.")); + linear_2_->load_state_dict(state_dict.get_dict_with_prefix("linear_2.")); + } + + void verify_loaded_weights(const std::string& prefix) const { + linear_1_->verify_loaded_weights(prefix + "linear_1."); + linear_2_->verify_loaded_weights(prefix + "linear_2."); + } + + private: + torch::TensorOptions options_; + layer::AddMatmul linear_1_{nullptr}; + torch::nn::SiLU act_{nullptr}; + layer::AddMatmul linear_2_{nullptr}; +}; +TORCH_MODULE(Flux2TimestepEmbedding); + +class Flux2TimestepsImpl : public torch::nn::Module { + public: + explicit Flux2TimestepsImpl(int64_t num_channels, + bool flip_sin_to_cos = true, + float downscale_freq_shift = 0.0, + int64_t scale = 1) + : num_channels_(num_channels), + flip_sin_to_cos_(flip_sin_to_cos), + downscale_freq_shift_(downscale_freq_shift), + scale_(scale) {} + + torch::Tensor forward(const torch::Tensor& timesteps) { + return get_timestep_embedding(timesteps, + num_channels_, + flip_sin_to_cos_, + downscale_freq_shift_, + scale_); + } + + private: + int64_t num_channels_; + bool flip_sin_to_cos_; + float downscale_freq_shift_; + int64_t scale_; + + torch::Tensor get_timestep_embedding(const torch::Tensor& timesteps, + int embedding_dim, + bool flip_sin_to_cos = false, + float downscale_freq_shift = 1.0f, + float scale = 1.0f, + int max_period = 10000) { + int half_dim = embedding_dim / 2; + auto exponent = -std::log(static_cast(max_period)) * + torch::arange(0, + half_dim, + torch::TensorOptions() + .dtype(torch::kFloat32) + .device(timesteps.device())); + exponent = exponent / (half_dim - downscale_freq_shift); + + auto emb = torch::exp(exponent); + emb = timesteps.unsqueeze(1).to(torch::kFloat32) * emb.unsqueeze(0); + emb = scale * emb; + emb = torch::cat({torch::sin(emb), torch::cos(emb)}, /*dim=*/-1); + + if (flip_sin_to_cos) { + emb = torch::cat({emb.slice(/*dim=*/-1, /*start=*/half_dim), + emb.slice(/*dim=*/-1, /*start=*/0, /*end=*/half_dim)}, + /*dim=*/-1); + } + + if (embedding_dim % 2 == 1) { + emb = torch::nn::functional::pad( + emb, torch::nn::functional::PadFuncOptions({0, 1, 0, 0})); + } + + return emb; + } +}; +TORCH_MODULE(Flux2Timesteps); + +class Flux2TimestepGuidanceEmbeddingsImpl : public torch::nn::Module { + public: + explicit Flux2TimestepGuidanceEmbeddingsImpl(const ModelContext& context, + int64_t embedding_dim, + bool bias = false, + bool guidance_embeds = true) + : options_(context.get_tensor_options()) { + auto model_args = context.get_model_args(); + in_channels_ = model_args.timestep_guidance_channels(); + embedding_dim_ = embedding_dim; + guidance_embeds_ = guidance_embeds; + + time_proj_ = + register_module("time_proj", Flux2Timesteps(in_channels_, true, 0.0)); + timestep_embedder_ = register_module( + "timestep_embedder", + Flux2TimestepEmbedding(in_channels_, embedding_dim_, -1, bias)); + if (guidance_embeds_) { + guidance_embedder_ = register_module( + "guidance_embedder", + Flux2TimestepEmbedding(in_channels_, embedding_dim_, -1, bias)); + } + } + + torch::Tensor forward(const torch::Tensor& timestep, + const torch::Tensor& guidance) { + auto timesteps_proj = time_proj_->forward(timestep); + auto timesteps_emb = + timestep_embedder_->forward(timesteps_proj.to(timestep.dtype())); + + if (guidance_embeds_ && guidance.defined()) { + auto guidance_proj = time_proj_->forward(guidance); + auto guidance_emb = + guidance_embedder_->forward(guidance_proj.to(guidance.dtype())); + return timesteps_emb + guidance_emb; + } else { + return timesteps_emb; + } + } + + void load_state_dict(const StateDict& state_dict) { + timestep_embedder_->load_state_dict( + state_dict.get_dict_with_prefix("timestep_embedder.")); + if (guidance_embeds_) { + guidance_embedder_->load_state_dict( + state_dict.get_dict_with_prefix("guidance_embedder.")); + } + } + + void verify_loaded_weights(const std::string& prefix) const { + timestep_embedder_->verify_loaded_weights(prefix + "timestep_embedder."); + if (guidance_embeds_) { + guidance_embedder_->verify_loaded_weights(prefix + "guidance_embedder."); + } + } + + private: + int64_t in_channels_; + int64_t embedding_dim_; + bool guidance_embeds_; + Flux2Timesteps time_proj_{nullptr}; + Flux2TimestepEmbedding timestep_embedder_{nullptr}; + Flux2TimestepEmbedding guidance_embedder_{nullptr}; + torch::TensorOptions options_; +}; +TORCH_MODULE(Flux2TimestepGuidanceEmbeddings); + +class Flux2TransformerBlockImpl : public torch::nn::Module { + public: + explicit Flux2TransformerBlockImpl(const ModelContext& context, + int64_t dim, + const ParallelArgs& parallel_args) + : options_(context.get_tensor_options()), parallel_args_(parallel_args) { + auto model_args = context.get_model_args(); + double eps = model_args.eps(); + + norm1_ = register_module( + "norm1", + torch::nn::LayerNorm( + torch::nn::LayerNormOptions({dim}).elementwise_affine(false).eps( + eps))); + + norm1_context_ = register_module( + "norm1_context", + torch::nn::LayerNorm( + torch::nn::LayerNormOptions({dim}).elementwise_affine(false).eps( + eps))); + + attn_ = register_module("attn", Flux2Attention(context, parallel_args)); + + norm2_ = register_module( + "norm2", + torch::nn::LayerNorm( + torch::nn::LayerNormOptions({dim}).elementwise_affine(false).eps( + eps))); + + ff_ = register_module("ff", Flux2FeedForward(context, parallel_args)); + + norm2_context_ = register_module( + "norm2_context", + torch::nn::LayerNorm( + torch::nn::LayerNormOptions({dim}).elementwise_affine(false).eps( + eps))); + + ff_context_ = + register_module("ff_context", Flux2FeedForward(context, parallel_args)); + } + + std::tuple forward( + const torch::Tensor& hidden_states, + const torch::Tensor& encoder_hidden_states, + const torch::Tensor& temb_img, + const torch::Tensor& temb_txt, + const torch::Tensor& image_rotary_emb) { + auto img_mod_params = Flux2ModulationImpl::split(temb_img, 2); + auto txt_mod_params = Flux2ModulationImpl::split(temb_txt, 2); + + auto [shift_msa_img, scale_msa_img, gate_msa_img] = img_mod_params[0]; + auto [shift_mlp_img, scale_mlp_img, gate_mlp_img] = img_mod_params[1]; + + auto [shift_msa_txt, scale_msa_txt, gate_msa_txt] = txt_mod_params[0]; + auto [shift_mlp_txt, scale_mlp_txt, gate_mlp_txt] = txt_mod_params[1]; + + auto norm_hidden_states = norm1_->forward(hidden_states); + norm_hidden_states = + (1 + scale_msa_img) * norm_hidden_states + shift_msa_img; + + auto norm_encoder_hidden_states = + norm1_context_->forward(encoder_hidden_states); + norm_encoder_hidden_states = + (1 + scale_msa_txt) * norm_encoder_hidden_states + shift_msa_txt; + + auto [attn_output, context_attn_output] = attn_->forward( + norm_hidden_states, norm_encoder_hidden_states, image_rotary_emb); + + attn_output = gate_msa_img * attn_output; + torch::Tensor new_hidden_states = hidden_states + attn_output; + + auto norm_hs = norm2_->forward(new_hidden_states); + norm_hs = norm_hs * (1 + scale_mlp_img) + shift_mlp_img; + + auto ff_output = ff_->forward(norm_hs, /*is_save*/ true); + + new_hidden_states = new_hidden_states + gate_mlp_img * ff_output; + + context_attn_output = gate_msa_txt * context_attn_output; + torch::Tensor new_encoder_hidden_states = + encoder_hidden_states + context_attn_output; + + auto norm_enc_hs = norm2_context_->forward(new_encoder_hidden_states); + norm_enc_hs = norm_enc_hs * (1 + scale_mlp_txt) + shift_mlp_txt; + + auto ff_context_out = ff_context_->forward(norm_enc_hs, /*is_save*/ false); + new_encoder_hidden_states = + new_encoder_hidden_states + gate_mlp_txt * ff_context_out; + + if (new_encoder_hidden_states.scalar_type() == torch::kFloat16) { + new_encoder_hidden_states = + torch::clamp(new_encoder_hidden_states, -65504.0f, 65504.0f); + } + + return std::make_tuple(new_encoder_hidden_states, new_hidden_states); + } + + void load_state_dict(const StateDict& state_dict) { + attn_->load_state_dict(state_dict.get_dict_with_prefix("attn.")); + ff_->load_state_dict(state_dict.get_dict_with_prefix("ff.")); + ff_context_->load_state_dict( + state_dict.get_dict_with_prefix("ff_context.")); + } + + void verify_loaded_weights(const std::string& prefix) const { + attn_->verify_loaded_weights(prefix + "attn."); + ff_->verify_loaded_weights(prefix + "ff."); + ff_context_->verify_loaded_weights(prefix + "ff_context."); + } + + private: + torch::nn::LayerNorm norm1_{nullptr}; + torch::nn::LayerNorm norm1_context_{nullptr}; + Flux2Attention attn_{nullptr}; + torch::nn::LayerNorm norm2_{nullptr}; + Flux2FeedForward ff_{nullptr}; + torch::nn::LayerNorm norm2_context_{nullptr}; + Flux2FeedForward ff_context_{nullptr}; + torch::TensorOptions options_; + ParallelArgs parallel_args_; +}; +TORCH_MODULE(Flux2TransformerBlock); + +class Flux2ParallelSelfAttentionImpl : public torch::nn::Module { + public: + explicit Flux2ParallelSelfAttentionImpl(const ModelContext& context, + const ParallelArgs& parallel_args) + : options_(context.get_tensor_options()), parallel_args_(parallel_args) { + auto model_args = context.get_model_args(); + heads_ = model_args.n_heads(); + head_dim_ = model_args.head_dim(); + query_dim_ = heads_ * head_dim_; + out_dim_ = query_dim_; + mlp_ratio_ = model_args.mlp_ratio(); + mlp_hidden_dim_ = static_cast(query_dim_ * mlp_ratio_); + mlp_mult_factor_ = 2; + + // ── Determine parallelism strategy for QKV ── + std::optional tp_options = std::nullopt; + + if (::xllm::ParallelConfig::get_instance().tp_size() > 1) { + tp_options = TpOptions( + /*column_parallel=*/true, + /*gather_output=*/false, + /*need_scatter=*/false, + /*process_group=*/parallel_args_.dit_tp_group_); + } + + // Individual Q, K, V projections (instead of fused) + auto to_q = DiTParallelLinear( + query_dim_, query_dim_, false, options_, std::nullopt, tp_options); + to_q_ = register_module("to_q", to_q); + + auto to_k = DiTParallelLinear( + query_dim_, query_dim_, false, options_, std::nullopt, tp_options); + to_k_ = register_module("to_k", to_k); + + auto to_v = DiTParallelLinear( + query_dim_, query_dim_, false, options_, std::nullopt, tp_options); + to_v_ = register_module("to_v", to_v); + + // ── MLP projection (not SP-aware, local computation) ── + std::optional mlp_tp_options = std::nullopt; + + if (::xllm::ParallelConfig::get_instance().tp_size() > 1) { + mlp_tp_options = TpOptions( + /*column_parallel=*/true, + /*gather_output=*/false, + /*need_scatter=*/false, + /*process_group=*/parallel_args_.dit_tp_group_); + } + + int64_t mlp_proj_dim = mlp_hidden_dim_ * mlp_mult_factor_; + auto mlp_proj = DiTParallelLinear(query_dim_, + mlp_proj_dim, + false, + options_, + std::nullopt, + mlp_tp_options); + mlp_proj_ = register_module("mlp_proj", mlp_proj); + + mlp_act_fn_ = register_module("mlp_act_fn", Flux2SwiGLU()); + norm_q_ = + register_module("norm_q", layer::RMSNorm(head_dim_, 1e-6f, options_)); + norm_k_ = + register_module("norm_k", layer::RMSNorm(head_dim_, 1e-6f, options_)); + + // ── Output projection ── + // In SP mode, we handle the all2all manually in forward() before this + // layer, so to_out_ is always Default or TP (not SP). + int64_t out_input_dim = query_dim_ + mlp_hidden_dim_; + std::optional tp_out_options = std::nullopt; + + if (::xllm::ParallelConfig::get_instance().tp_size() > 1) { + tp_out_options = TpOptions( + /*column_parallel=*/false, + /*gather_output=*/true, + /*need_scatter=*/true, + /*process_group=*/parallel_args_.dit_tp_group_); + } + + auto to_out = DiTParallelLinear( + out_input_dim, out_dim_, false, options_, std::nullopt, tp_out_options); + to_out_ = register_module("to_out", to_out); + } + + torch::Tensor forward(const torch::Tensor& hidden_states, + const torch::Tensor& image_rotary_emb) { + int64_t batch_size = hidden_states.size(0); + + // ── Separate Q/K/V and MLP projections ── + auto q = to_q_->forward(hidden_states); + auto k = to_k_->forward(hidden_states); + auto v = to_v_->forward(hidden_states); + auto mlp_output = mlp_proj_->forward(hidden_states); + + int64_t attn_heads = q.size(-1) / head_dim_; + int64_t head_dim = head_dim_; + + q = q.view({batch_size, -1, attn_heads, head_dim}); + k = k.view({batch_size, -1, attn_heads, head_dim}); + v = v.view({batch_size, -1, attn_heads, head_dim}); + + if (norm_q_) { + q = std::get<0>(norm_q_->forward(q)); + } + if (norm_k_) { + k = std::get<0>(norm_k_->forward(k)); + } + + if (image_rotary_emb.defined()) { + q = apply_rotary_emb(q, image_rotary_emb, false); + k = apply_rotary_emb(k, image_rotary_emb, false); + } + + // After SP all2all (before_attention=true), attn_heads already equals + // heads/sp_size. No further division needed. + int64_t local_heads = attn_heads; + +#if defined(USE_NPU) + int64_t head_num_ = q.size(2); + int64_t head_dim_ = q.size(-1); + auto results = + at_npu::native::custom_ops::npu_fusion_attention(q, + k, + v, + local_heads, + "BSND", + torch::nullopt, + torch::nullopt, + torch::nullopt, + pow(head_dim_, -0.5), + 1.0, + 65535, + 65535); + auto attn_output = std::get<0>(results); + + attn_output = attn_output.reshape({batch_size, -1, local_heads * head_dim}); + +#elif defined(USE_CUDA) + q = q.transpose(1, 2); + k = k.transpose(1, 2); + v = v.transpose(1, 2); + + torch::Tensor attn_output = torch::scaled_dot_product_attention( + q, k, v, torch::nullopt, 0.0, false); + attn_output = attn_output.transpose(1, 2).reshape( + {batch_size, -1, local_heads * head_dim}); +#else + NOT_IMPLEMENTED(); +#endif + + attn_output = attn_output.to(q.dtype()); + + if (::xllm::ParallelConfig::get_instance().tp_size() > 1) { + mlp_output = mlp_output.contiguous(); + mlp_output = + parallel_state::gather(mlp_output, parallel_args_.dit_tp_group_, -1); + + attn_output = attn_output.contiguous(); + attn_output = + parallel_state::gather(attn_output, parallel_args_.dit_tp_group_, -1); + } + + mlp_output = mlp_act_fn_(mlp_output); + auto output = + torch::cat(std::vector{attn_output, mlp_output}, -1); + + output = to_out_->forward(output); + + return output; + } + + void load_state_dict(const StateDict& state_dict) { + // The original checkpoint has a fused "to_qkv_mlp_proj.weight" tensor. + // We split it into separate Q/K/V/MLP weights and load individually. + auto fused_weight = state_dict.get_tensor("to_qkv_mlp_proj.weight"); + if (fused_weight.defined()) { + int64_t tp_size = ::xllm::ParallelConfig::get_instance().tp_size() > 1 + ? ::xllm::ParallelConfig::get_instance().tp_size() + : 1; + // Original fused layout: rows = [Q, K, V, MLP_gate, MLP_up] + // each of Q/K/V has query_dim_ rows, MLP has mlp_hidden_dim_*2 rows + auto q_weight = fused_weight.slice(0, 0, query_dim_); + auto k_weight = fused_weight.slice(0, query_dim_, query_dim_ * 2); + auto v_weight = fused_weight.slice(0, query_dim_ * 2, query_dim_ * 3); + auto mlp_weight = fused_weight.slice( + 0, + query_dim_ * 3, + query_dim_ * 3 + mlp_hidden_dim_ * mlp_mult_factor_); + + if (tp_size > 1) { + // TP or Combined TP+SP mode: shard each weight along dim 0 by rank. + // SP doesn't shard weights, so combined mode uses the same TP sharding. + int64_t tp_rank = parallel_args_.dit_tp_group_->rank(); + auto q_rank = q_weight.chunk(tp_size, 0)[tp_rank]; + auto k_rank = k_weight.chunk(tp_size, 0)[tp_rank]; + auto v_rank = v_weight.chunk(tp_size, 0)[tp_rank]; + auto mlp_rank = mlp_weight.chunk(tp_size, 0)[tp_rank]; + + // Directly copy already-sharded weights to parameters. + // Same pattern as transformer_flux2.h: bypass + // DiTParallelLinear::load_state_dict which would double-shard. + torch::NoGradGuard no_grad; + for (auto& p : to_q_->named_parameters()) { + if (p.key() == "weight") { + p.value().copy_(q_rank.to(options_)); + break; + } + } + for (auto& p : to_k_->named_parameters()) { + if (p.key() == "weight") { + p.value().copy_(k_rank.to(options_)); + break; + } + } + for (auto& p : to_v_->named_parameters()) { + if (p.key() == "weight") { + p.value().copy_(v_rank.to(options_)); + break; + } + } + for (auto& p : mlp_proj_->named_parameters()) { + if (p.key() == "weight") { + p.value().copy_(mlp_rank.to(options_)); + break; + } + } + } else { + // Default mode: full weights + { + std::unordered_map temp_dict; + temp_dict["weight"] = q_weight; + StateDict temp_state_dict(temp_dict); + to_q_->as()->load_state_dict(temp_state_dict); + } + { + std::unordered_map temp_dict; + temp_dict["weight"] = k_weight; + StateDict temp_state_dict(temp_dict); + to_k_->as()->load_state_dict(temp_state_dict); + } + { + std::unordered_map temp_dict; + temp_dict["weight"] = v_weight; + StateDict temp_state_dict(temp_dict); + to_v_->as()->load_state_dict(temp_state_dict); + } + { + std::unordered_map temp_dict; + temp_dict["weight"] = mlp_weight; + StateDict temp_state_dict(temp_dict); + mlp_proj_->as()->load_state_dict(temp_state_dict); + } + } + } + norm_q_->load_state_dict(state_dict.get_dict_with_prefix("norm_q.")); + norm_k_->load_state_dict(state_dict.get_dict_with_prefix("norm_k.")); + to_out_->as()->load_state_dict( + state_dict.get_dict_with_prefix("to_out.")); + } + + void verify_loaded_weights(const std::string& prefix) const { + to_out_->as()->verify_loaded_weights(prefix + "to_out."); + } + + private: + int64_t heads_; + int64_t head_dim_; + int64_t query_dim_; + int64_t out_dim_; + float mlp_ratio_; + int64_t mlp_hidden_dim_; + int64_t mlp_mult_factor_; + DiTParallelLinear to_q_{nullptr}; + DiTParallelLinear to_k_{nullptr}; + DiTParallelLinear to_v_{nullptr}; + DiTParallelLinear mlp_proj_{nullptr}; + Flux2SwiGLU mlp_act_fn_{nullptr}; + layer::RMSNorm norm_q_{nullptr}; + layer::RMSNorm norm_k_{nullptr}; + DiTParallelLinear to_out_{nullptr}; + torch::TensorOptions options_; + ParallelArgs parallel_args_; +}; +TORCH_MODULE(Flux2ParallelSelfAttention); + +class Flux2SingleTransformerBlockImpl : public torch::nn::Module { + public: + explicit Flux2SingleTransformerBlockImpl(const ModelContext& context, + int64_t inner_dim, + const ParallelArgs& parallel_args) + : options_(context.get_tensor_options()), parallel_args_(parallel_args) { + auto model_args = context.get_model_args(); + + norm_ = register_module( + "norm", + torch::nn::LayerNorm(torch::nn::LayerNormOptions({inner_dim}) + .elementwise_affine(false) + .eps(1e-6))); + + attn_ = register_module("attn", + Flux2ParallelSelfAttention(context, parallel_args)); + } + + torch::Tensor forward(const torch::Tensor& hidden_states, + const torch::Tensor& temb_mod, + const torch::Tensor& image_rotary_emb, + bool split_hidden_states, + int64_t text_seq_len) { + auto mod_params = Flux2ModulationImpl::split(temb_mod, 1); + auto [shift, scale, gate] = mod_params[0]; + + auto norm_hidden_states = norm_->forward(hidden_states); + norm_hidden_states = (1 + scale) * norm_hidden_states + shift; + auto attn_output = attn_->forward(norm_hidden_states, image_rotary_emb); + + auto gate_attn = gate * attn_output; + auto output = hidden_states + gate_attn; + if (output.dtype() == torch::kFloat16) { + output = torch::clamp(output, -65504.0f, 65504.0f); + } + + return output; + } + + void load_state_dict(const StateDict& state_dict) { + attn_->load_state_dict(state_dict.get_dict_with_prefix("attn.")); + } + + void verify_loaded_weights(const std::string& prefix) const { + attn_->verify_loaded_weights(prefix + "attn."); + } + + private: + torch::nn::LayerNorm norm_{nullptr}; + Flux2ParallelSelfAttention attn_{nullptr}; + torch::TensorOptions options_; + ParallelArgs parallel_args_; +}; +TORCH_MODULE(Flux2SingleTransformerBlock); + +class Flux2Transformer2DModelImpl : public torch::nn::Module { + public: + explicit Flux2Transformer2DModelImpl(const ModelContext& context, + const ParallelArgs& parallel_args) + : options_(context.get_tensor_options()), parallel_args_(parallel_args) { + auto model_args = context.get_model_args(); + int64_t num_attention_heads = model_args.n_heads(); + int64_t attention_head_dim = model_args.head_dim(); + int64_t joint_attention_dim = model_args.joint_attention_dim(); + int64_t num_layers = model_args.num_layers(); + int64_t num_single_layers = model_args.num_single_layers(); + int64_t patch_size = model_args.patch_size(); + float rope_theta = model_args.rope_theta(); + auto axes_dims_rope = model_args.axes_dims_rope(); + in_channels_ = model_args.in_channels(); + out_channels_ = model_args.out_channels(); + + int64_t inner_dim = num_attention_heads * attention_head_dim; + + time_guidance_embed_ = + Flux2TimestepGuidanceEmbeddings(context, inner_dim, false, true); + register_module("time_guidance_embed", time_guidance_embed_); + + double_stream_modulation_img_ = + register_module("double_stream_modulation_img", + Flux2Modulation(context, inner_dim, 2, false)); + double_stream_modulation_txt_ = + register_module("double_stream_modulation_txt", + Flux2Modulation(context, inner_dim, 2, false)); + single_stream_modulation_ = + register_module("single_stream_modulation", + Flux2Modulation(context, inner_dim, 1, false)); + + x_embedder_ = register_module( + "x_embedder", + layer::AddMatmul(in_channels_, inner_dim, false, options_)); + context_embedder_ = register_module( + "context_embedder", + layer::AddMatmul(joint_attention_dim, inner_dim, false, options_)); + + transformer_blocks_ = + register_module("transformer_blocks", torch::nn::ModuleList()); + single_transformer_blocks_ = + register_module("single_transformer_blocks", torch::nn::ModuleList()); + + transformer_block_layers_.reserve(num_layers); + for (int64_t i = 0; i < num_layers; ++i) { + auto block = Flux2TransformerBlock(context, inner_dim, parallel_args); + transformer_blocks_->push_back(block); + transformer_block_layers_.push_back(block); + } + + single_transformer_block_layers_.reserve(num_single_layers); + for (int64_t i = 0; i < num_single_layers; ++i) { + auto block = + Flux2SingleTransformerBlock(context, inner_dim, parallel_args); + single_transformer_blocks_->push_back(block); + single_transformer_block_layers_.push_back(block); + } + + norm_out_ = + register_module("norm_out", AdaLayerNormContinuous(context, false)); + proj_out_ = register_module( + "proj_out", + layer::AddMatmul(inner_dim, + patch_size * patch_size * out_channels_, + false, + options_)); + } + + torch::Tensor forward(const torch::Tensor& hidden_states_input, + const torch::Tensor& encoder_hidden_states_input, + const torch::Tensor& timestep, + const torch::Tensor& guidance, + const torch::Tensor& image_rotary_emb, + int64_t step_idx = 0) { + torch::Tensor hidden_states = x_embedder_->forward(hidden_states_input); + torch::Tensor encoder_hidden_states = + context_embedder_->forward(encoder_hidden_states_input); + + auto timestep_scaled = timestep.to(hidden_states.dtype()) * 1000.0f; + auto guidance_scaled = guidance.defined() + ? guidance.to(hidden_states.dtype()) * 1000.0f + : torch::Tensor(); + auto temb = time_guidance_embed_->forward(timestep_scaled, guidance_scaled); + + auto double_stream_mod_img = double_stream_modulation_img_->forward(temb); + auto double_stream_mod_txt = double_stream_modulation_txt_->forward(temb); + auto single_stream_mod = single_stream_modulation_->forward(temb); + + // ── Double-stream transformer blocks ── + for (int64_t i = 0; i < transformer_block_layers_.size(); ++i) { + auto block = transformer_block_layers_[i]; + auto [new_encoder_hidden, new_hidden] = + block->forward(hidden_states, + encoder_hidden_states, + double_stream_mod_img, + double_stream_mod_txt, + image_rotary_emb); + + hidden_states = new_hidden; + encoder_hidden_states = new_encoder_hidden; + } + + // ── Merge into single stream: [txt_seq, img_seq] + hidden_states = torch::cat({encoder_hidden_states, hidden_states}, 1); + + // ── Single-stream transformer blocks (DiTCache: use_cfg=true) ── + // NOTE: Flux2 does NOT support CFG. The "use_cfg" parameter routes to + // the second cache instance (active_cond_cache_) to isolate single-stream. + torch::Tensor ss_original_hidden_states = hidden_states; + auto dummy_encoder = + torch::zeros({hidden_states.size(0), 1, hidden_states.size(2)}, + hidden_states.options()); + + // Step start for single-stream phase + TensorMap ss_step_in_map = { + {"hidden_states", hidden_states}, + {"original_hidden_states", ss_original_hidden_states}}; + CacheStepIn ss_stepin_before(step_idx, ss_step_in_map); + bool ss_use_step_cache = + DiTCache::get_instance().on_before_step(ss_stepin_before, + /*use_cfg=*/true); + + if (!ss_use_step_cache) { + for (int64_t i = 0; i < single_transformer_block_layers_.size(); ++i) { + TensorMap ss_block_in_before_map = {}; + CacheBlockIn ss_blockin_before(i, ss_block_in_before_map); + bool ss_use_block_cache = + DiTCache::get_instance().on_before_block(ss_blockin_before, + /*use_cfg=*/true); + + if (!ss_use_block_cache) { + auto block = single_transformer_block_layers_[i]; + hidden_states = block->forward( + hidden_states, single_stream_mod, image_rotary_emb, false, 0); + } + + TensorMap ss_block_in_after_map = { + {"hidden_states", hidden_states}, + {"encoder_hidden_states", dummy_encoder}, + {"original_hidden_states", ss_original_hidden_states}}; + CacheBlockIn ss_blockin_after(i, ss_block_in_after_map); + CacheBlockOut ss_blockout_after = + DiTCache::get_instance().on_after_block(ss_blockin_after, + /*use_cfg=*/true); + + hidden_states = ss_blockout_after.tensors.at("hidden_states"); + } + } + + // Step end for single-stream phase + TensorMap ss_step_after_map = { + {"hidden_states", hidden_states}, + {"original_hidden_states", ss_original_hidden_states}}; + CacheStepIn ss_stepin_after(step_idx, ss_step_after_map); + CacheStepOut ss_stepout_after = + DiTCache::get_instance().on_after_step(ss_stepin_after, + /*use_cfg=*/true); + hidden_states = ss_stepout_after.tensors.at("hidden_states"); + + int64_t start = encoder_hidden_states.size(1); + int64_t length = hidden_states.size(1) - start; + auto output_hidden = + hidden_states.narrow(1, start, std::max(length, int64_t(0))); + + auto output_hidden_final = norm_out_->forward(output_hidden, temb); + + auto final_output = proj_out_->forward(output_hidden_final); + + return final_output; + } + + void load_model(std::unique_ptr loader) { + for (const auto& state_dict : loader->get_state_dicts()) { + context_embedder_->load_state_dict( + state_dict->get_dict_with_prefix("context_embedder.")); + x_embedder_->load_state_dict( + state_dict->get_dict_with_prefix("x_embedder.")); + time_guidance_embed_->load_state_dict( + state_dict->get_dict_with_prefix("time_guidance_embed.")); + double_stream_modulation_img_->load_state_dict( + state_dict->get_dict_with_prefix("double_stream_modulation_img.")); + double_stream_modulation_txt_->load_state_dict( + state_dict->get_dict_with_prefix("double_stream_modulation_txt.")); + single_stream_modulation_->load_state_dict( + state_dict->get_dict_with_prefix("single_stream_modulation.")); + for (int64_t i = 0; i < transformer_block_layers_.size(); ++i) { + auto block = transformer_block_layers_[i]; + block->load_state_dict(state_dict->get_dict_with_prefix( + "transformer_blocks." + std::to_string(i) + ".")); + } + for (int64_t i = 0; i < single_transformer_block_layers_.size(); ++i) { + auto block = single_transformer_block_layers_[i]; + block->load_state_dict(state_dict->get_dict_with_prefix( + "single_transformer_blocks." + std::to_string(i) + ".")); + } + norm_out_->load_state_dict(state_dict->get_dict_with_prefix("norm_out.")); + proj_out_->load_state_dict(state_dict->get_dict_with_prefix("proj_out.")); + } + } + + void verify_loaded_weights(const std::string& prefix) { + context_embedder_->verify_loaded_weights(prefix + "context_embedder."); + x_embedder_->verify_loaded_weights(prefix + "x_embedder."); + time_guidance_embed_->verify_loaded_weights(prefix + + "time_guidance_embed."); + double_stream_modulation_img_->verify_loaded_weights( + prefix + "double_stream_modulation_img."); + double_stream_modulation_txt_->verify_loaded_weights( + prefix + "double_stream_modulation_txt."); + single_stream_modulation_->verify_loaded_weights( + prefix + "single_stream_modulation."); + for (int64_t i = 0; i < transformer_block_layers_.size(); ++i) { + auto block = transformer_block_layers_[i]; + block->verify_loaded_weights(prefix + "transformer_blocks." + + std::to_string(i) + "."); + } + for (int64_t i = 0; i < single_transformer_block_layers_.size(); ++i) { + auto block = single_transformer_block_layers_[i]; + block->verify_loaded_weights(prefix + "single_transformer_blocks." + + std::to_string(i) + "."); + } + norm_out_->verify_loaded_weights(prefix + "norm_out."); + proj_out_->verify_loaded_weights(prefix + "proj_out."); + } + + int64_t in_channels() { return in_channels_; } + + private: + int64_t in_channels_; + int64_t out_channels_; + layer::AddMatmul context_embedder_{nullptr}; + layer::AddMatmul x_embedder_{nullptr}; + Flux2TimestepGuidanceEmbeddings time_guidance_embed_{nullptr}; + Flux2Modulation double_stream_modulation_img_{nullptr}; + Flux2Modulation double_stream_modulation_txt_{nullptr}; + Flux2Modulation single_stream_modulation_{nullptr}; + torch::nn::ModuleList transformer_blocks_{nullptr}; + std::vector transformer_block_layers_; + torch::nn::ModuleList single_transformer_blocks_{nullptr}; + std::vector single_transformer_block_layers_; + AdaLayerNormContinuous norm_out_{nullptr}; + layer::AddMatmul proj_out_{nullptr}; + torch::TensorOptions options_; + ParallelArgs parallel_args_; +}; +TORCH_MODULE(Flux2Transformer2DModel); + +class Flux2DiTModelImpl : public torch::nn::Module { + public: + explicit Flux2DiTModelImpl(const ModelContext& context, + const ParallelArgs& parallel_args) + : options_(context.get_tensor_options()), parallel_args_(parallel_args) { + flux2_transformer_2d_model_ = + register_module("flux2_transformer_2d_model_", + Flux2Transformer2DModel(context, parallel_args)); + } + + torch::Tensor forward(const torch::Tensor& hidden_states_input, + const torch::Tensor& encoder_hidden_states_input, + const torch::Tensor& timestep, + const torch::Tensor& guidance, + const torch::Tensor& image_rotary_emb, + int64_t step_idx = 0) { + torch::Tensor output = + flux2_transformer_2d_model_->forward(hidden_states_input, + encoder_hidden_states_input, + timestep, + guidance, + image_rotary_emb, + step_idx); + return output; + } + int64_t in_channels() { return flux2_transformer_2d_model_->in_channels(); } + + void load_model(std::unique_ptr loader) { + flux2_transformer_2d_model_->load_model(std::move(loader)); + flux2_transformer_2d_model_->verify_loaded_weights(""); + } + + private: + Flux2Transformer2DModel flux2_transformer_2d_model_{nullptr}; + torch::TensorOptions options_; + ParallelArgs parallel_args_; +}; +TORCH_MODULE(Flux2DiTModel); + +REGISTER_MODEL_ARGS(Flux2Transformer2DModel, [&] { + LOAD_ARG_OR(head_dim, "attention_head_dim", 128); + LOAD_ARG_OR(n_heads, "num_attention_heads", 48); + LOAD_ARG_OR( + axes_dims_rope, "axes_dims_rope", (std::vector{32, 32, 32, 32})); + LOAD_ARG_OR(eps, "eps", 1e-6); + LOAD_ARG_OR(in_channels, "in_channels", 128); + LOAD_ARG_OR(joint_attention_dim, "joint_attention_dim", 15360); + LOAD_ARG_OR(mlp_ratio, "mlp_ratio", 3.0f); + LOAD_ARG_OR(num_layers, "num_layers", 8); + LOAD_ARG_OR(num_single_layers, "num_single_layers", 48); + LOAD_ARG_OR(out_channels, "out_channels", 128); + LOAD_ARG_OR(patch_size, "patch_size", 1); + LOAD_ARG_OR(rope_theta, "rope_theta", 2000.0f); + LOAD_ARG_OR(timestep_guidance_channels, "timestep_guidance_channels", 256); +}); + +} // namespace xllm diff --git a/xllm/models/models.h b/xllm/models/models.h index a94510eb08..9bc4711c3c 100644 --- a/xllm/models/models.h +++ b/xllm/models/models.h @@ -17,6 +17,7 @@ limitations under the License. #if defined(USE_NPU) #include "dit/pipelines/pipeline_flux.h" // IWYU pragma: keep +#include "dit/pipelines/pipeline_flux2.h" // IWYU pragma: keep #include "dit/pipelines/pipeline_flux_control.h" // IWYU pragma: keep #include "dit/pipelines/pipeline_flux_fill.h" // IWYU pragma: keep #include "dit/pipelines/pipeline_qwenimage_edit_plus.h" // IWYU pragma: keep