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
100 lines
3.9 KiB
C++
100 lines
3.9 KiB
C++
#pragma once
|
|
|
|
#include "engine/models/universr/assets.h"
|
|
#include "engine/framework/core/backend_weight_store.h"
|
|
#include "engine/framework/modules/conv_modules.h"
|
|
#include "engine/framework/modules/linear_module.h"
|
|
#include "engine/framework/modules/norm_modules.h"
|
|
|
|
namespace engine::models::universr {
|
|
|
|
struct ConvNeXtWeights {
|
|
modules::Conv2dWeights depthwise;
|
|
modules::NormWeights norm;
|
|
modules::LinearWeights expansion;
|
|
core::TensorValue grn_gamma;
|
|
modules::LinearWeights projection;
|
|
};
|
|
|
|
ConvNeXtWeights load_convnext_weights(core::BackendWeightStore & store,
|
|
const assets::TensorSource & source, const std::string & prefix,
|
|
int64_t channels, assets::TensorStorageType storage);
|
|
|
|
// Matches upstream Block: NCHW input/output, reflection padding, and spatial GRN.
|
|
core::TensorValue build_convnext_block(core::ModuleBuildContext & ctx,
|
|
const core::TensorValue & input, const ConvNeXtWeights & weights);
|
|
struct ConditioningWeights {
|
|
core::TensorValue frequency_pe;
|
|
core::TensorValue sample_rate_embedding;
|
|
core::TensorValue unconditional_embedding;
|
|
core::TensorValue time_frequencies;
|
|
modules::LinearWeights sample_rate_projector;
|
|
modules::LinearWeights low_film;
|
|
modules::LinearWeights high_film;
|
|
modules::LinearWeights head;
|
|
modules::LinearWeights sample_rate_hidden;
|
|
modules::LinearWeights sample_rate_film;
|
|
std::vector<ConvNeXtWeights> blocks;
|
|
};
|
|
|
|
ConditioningWeights load_conditioning_weights(core::BackendWeightStore & store,
|
|
const assets::TensorSource & source, const UniverSRConfig & config,
|
|
assets::TensorStorageType storage);
|
|
|
|
core::TensorValue build_conditioning_encoder(core::ModuleBuildContext & ctx,
|
|
const core::TensorValue & low_spectrum, const core::TensorValue & sample_rate_embedding,
|
|
const ConditioningWeights & weights);
|
|
|
|
core::TensorValue build_spatial_conditioning(core::ModuleBuildContext & ctx,
|
|
const core::TensorValue & condition, const UniverSRConfig & config,
|
|
const ConditioningWeights & weights);
|
|
|
|
core::TensorValue build_time_embedding(core::ModuleBuildContext & ctx,
|
|
const core::TensorValue & time, const core::TensorValue & sample_rate_embedding,
|
|
const ConditioningWeights & weights);
|
|
|
|
core::TensorValue build_projected_conditioning(core::ModuleBuildContext & ctx,
|
|
const core::TensorValue & condition, const UniverSRConfig & config,
|
|
const ConditioningWeights & weights, const core::TensorValue & projection);
|
|
struct TimeBlockWeights {
|
|
modules::LinearWeights time_hidden;
|
|
modules::LinearWeights time_output;
|
|
ConvNeXtWeights block;
|
|
};
|
|
|
|
struct EncoderWeights {
|
|
std::vector<TimeBlockWeights> blocks;
|
|
modules::NormWeights norm;
|
|
modules::Conv2dWeights downsample;
|
|
};
|
|
|
|
struct DecoderWeights {
|
|
// [4 * output_channels, input_channels], with each output channel's 2x2 patch in row-major order.
|
|
core::TensorValue upsample_weight;
|
|
core::TensorValue upsample_bias;
|
|
std::vector<TimeBlockWeights> blocks;
|
|
};
|
|
|
|
struct UNetBackboneWeights {
|
|
modules::LinearWeights input;
|
|
modules::NormWeights input_norm;
|
|
std::vector<EncoderWeights> encoders;
|
|
std::vector<TimeBlockWeights> middle;
|
|
std::vector<DecoderWeights> decoders;
|
|
modules::LinearWeights output;
|
|
};
|
|
|
|
UNetBackboneWeights load_unet_backbone_weights(core::BackendWeightStore & store,
|
|
const assets::TensorSource & source, const UniverSRConfig & config,
|
|
assets::TensorStorageType storage);
|
|
|
|
core::TensorValue build_time_block(core::ModuleBuildContext & ctx,
|
|
const core::TensorValue & input, const core::TensorValue & time,
|
|
const TimeBlockWeights & weights);
|
|
|
|
// Input is the NCHW concatenation of noisy spectrum and spatial conditioning.
|
|
core::TensorValue build_unet_backbone(core::ModuleBuildContext & ctx,
|
|
const core::TensorValue & input, const core::TensorValue & time,
|
|
const UNetBackboneWeights & weights, bool input_projected = false);
|
|
|
|
} // namespace engine::models::universr
|