vrr/koboldcpp
2
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2025-09-11 01:24:36 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

taesd now supports flux and sd3

Browse source
This commit is contained in:
Concedo 2025-01-08 22:35:50 +08:00
parent c73d99ccac
commit 3732bb2686
15 changed files with 50 additions and 52 deletions
Download patch file Download diff file Expand all files Collapse all files

23
otherarch/sdcpp/tae.hpp
View file

@ -62,7 +62,8 @@ class TinyEncoder : public UnaryBlock {
int num_blocks = 3;
public:
TinyEncoder() {
TinyEncoder(int z_channels = 4)
: z_channels(z_channels) {
int index = 0;
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, channels, {3, 3}, {1, 1}, {1, 1}));
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
@ -106,7 +107,10 @@ class TinyDecoder : public UnaryBlock {
int num_blocks = 3;
public:
TinyDecoder(int index = 0) {
TinyDecoder(int z_channels = 4)
: z_channels(z_channels) {
int index = 0;
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(z_channels, channels, {3, 3}, {1, 1}, {1, 1}));
index++; // nn.ReLU()
@ -163,12 +167,16 @@ protected:
bool decode_only;
public:
TAESD(bool decode_only = true)
TAESD(bool decode_only = true, SDVersion version = VERSION_SD1)
: decode_only(decode_only) {
blocks["decoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyDecoder());
int z_channels = 4;
if (sd_version_is_dit(version)) {
z_channels = 16;
}
blocks["decoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyDecoder(z_channels));
if (!decode_only) {
blocks["encoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyEncoder());
blocks["encoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyEncoder(z_channels));
}
}
@ -190,9 +198,10 @@ struct TinyAutoEncoder : public GGMLRunner {
TinyAutoEncoder(ggml_backend_t backend,
std::map<std::string, enum ggml_type>& tensor_types,
const std::string prefix,
bool decoder_only = true)
bool decoder_only = true,
SDVersion version = VERSION_SD1)
: decode_only(decoder_only),
taesd(decode_only),
taesd(decode_only, version),
GGMLRunner(backend) {
taesd.init(params_ctx, tensor_types, prefix);
}