Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
139 lines
4.4 KiB
C++
139 lines
4.4 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/assets/tensor_source.h"
|
|
#include "engine/framework/core/backend.h"
|
|
#include "engine/framework/core/module.h"
|
|
#include "engine/framework/modules/linear_module.h"
|
|
#include "engine/models/dramabox/assets.h"
|
|
#include "engine/models/dramabox/prompt_connector.h"
|
|
|
|
#include <cstdint>
|
|
#include <memory>
|
|
#include <optional>
|
|
#include <vector>
|
|
|
|
namespace engine::core {
|
|
class BackendWeightStore;
|
|
class ExecutionContext;
|
|
}
|
|
|
|
namespace engine::models::dramabox {
|
|
|
|
struct DramaBoxAdaLayerNormWeights {
|
|
modules::LinearWeights timestep_linear_1;
|
|
modules::LinearWeights timestep_linear_2;
|
|
modules::LinearWeights output_linear;
|
|
int64_t output_coefficient = 0;
|
|
};
|
|
|
|
struct DramaBoxDiTSelfAttentionWeights {
|
|
modules::LinearWeights qkv_gate;
|
|
std::optional<modules::LinearWeights> q;
|
|
std::optional<modules::LinearWeights> k;
|
|
std::optional<modules::LinearWeights> v;
|
|
std::optional<modules::LinearWeights> gate;
|
|
modules::LinearWeights out;
|
|
core::TensorValue q_norm;
|
|
core::TensorValue k_norm;
|
|
};
|
|
|
|
struct DramaBoxDiTCrossAttentionWeights {
|
|
modules::LinearWeights q_gate;
|
|
modules::LinearWeights kv;
|
|
modules::LinearWeights out;
|
|
core::TensorValue q_norm;
|
|
core::TensorValue k_norm;
|
|
};
|
|
|
|
struct DramaBoxDiTBlockWeights {
|
|
core::TensorValue audio_scale_shift_table;
|
|
core::TensorValue audio_prompt_scale_shift_table;
|
|
DramaBoxDiTSelfAttentionWeights self_attention;
|
|
DramaBoxDiTCrossAttentionWeights cross_attention;
|
|
modules::LinearWeights ff_in;
|
|
modules::LinearWeights ff_out;
|
|
};
|
|
|
|
struct DramaBoxDiTWeights {
|
|
std::shared_ptr<core::BackendWeightStore> store;
|
|
modules::LinearWeights patchify_proj;
|
|
DramaBoxAdaLayerNormWeights adaln;
|
|
DramaBoxAdaLayerNormWeights prompt_adaln;
|
|
std::vector<DramaBoxDiTBlockWeights> blocks;
|
|
core::TensorValue output_scale_shift_table;
|
|
modules::LinearWeights output_proj;
|
|
};
|
|
|
|
struct DramaBoxDiTInputs {
|
|
int64_t batch = 0;
|
|
int64_t tokens = 0;
|
|
bool stg_enabled = false;
|
|
int64_t ref_tokens = 0;
|
|
const std::vector<float> * latent = nullptr;
|
|
const std::vector<float> * sigma_features = nullptr;
|
|
};
|
|
|
|
DramaBoxDiTWeights load_dramabox_dit_weights(
|
|
const DramaBoxAssets & assets,
|
|
ggml_backend_t backend,
|
|
core::BackendType backend_type,
|
|
size_t weight_context_bytes,
|
|
assets::TensorStorageType weight_storage_type,
|
|
DramaBoxPerfMode perf_mode = DramaBoxPerfMode::Exact);
|
|
|
|
class DramaBoxDiTRuntime {
|
|
public:
|
|
DramaBoxDiTRuntime(
|
|
core::ExecutionContext & execution,
|
|
std::shared_ptr<const DramaBoxAssets> assets,
|
|
assets::TensorStorageType weight_storage_type,
|
|
DramaBoxPerfMode perf_mode);
|
|
~DramaBoxDiTRuntime();
|
|
|
|
DramaBoxDiTRuntime(const DramaBoxDiTRuntime &) = delete;
|
|
DramaBoxDiTRuntime & operator=(const DramaBoxDiTRuntime &) = delete;
|
|
|
|
void prepare(int64_t batch, int64_t tokens, int64_t context_tokens, bool stg_enabled, int64_t ref_tokens) const;
|
|
void prepare_static_inputs(
|
|
int64_t batch,
|
|
int64_t tokens,
|
|
bool stg_enabled,
|
|
int64_t ref_tokens,
|
|
const DramaBoxConditioningEncoding & conditioning,
|
|
const std::vector<float> & rope_cos,
|
|
const std::vector<float> & rope_sin,
|
|
const std::vector<float> & timestep_mask) const;
|
|
std::vector<float> forward(const DramaBoxDiTInputs & inputs) const;
|
|
void release_runtime_state() const;
|
|
|
|
private:
|
|
class Graph;
|
|
|
|
core::ExecutionContext * execution_ = nullptr;
|
|
std::shared_ptr<const DramaBoxAssets> assets_;
|
|
assets::TensorStorageType weight_storage_type_ = assets::TensorStorageType::Native;
|
|
DramaBoxPerfMode perf_mode_ = DramaBoxPerfMode::Exact;
|
|
mutable std::unique_ptr<DramaBoxDiTWeights> weights_;
|
|
mutable std::unique_ptr<Graph> graph_;
|
|
};
|
|
|
|
std::vector<float> make_dramabox_timestep_features(const std::vector<float> & timesteps);
|
|
void fill_dramabox_timestep_features(
|
|
const std::vector<float> & timesteps,
|
|
std::vector<float> & out);
|
|
void make_dramabox_audio_rope(
|
|
const std::vector<float> & positions,
|
|
int64_t batch,
|
|
int64_t tokens,
|
|
const DramaBoxConfig & config,
|
|
std::vector<float> & cos,
|
|
std::vector<float> & sin);
|
|
void make_dramabox_audio_rope_repeated(
|
|
const std::vector<float> & positions,
|
|
int64_t repeat_count,
|
|
int64_t tokens,
|
|
const DramaBoxConfig & config,
|
|
std::vector<float> & cos,
|
|
std::vector<float> & sin);
|
|
|
|
} // namespace engine::models::dramabox
|