fix: set default lora_model_dir to . (#1224)

README.md

@@ -15,6 +15,9 @@ API and command-line option may change frequently.***
 
 ## 🔥Important News
 
+* **2026/01/18** 🚀 stable-diffusion.cpp now supports **FLUX.2-klein**  
+  👉 Details: [PR #1193](https://github.com/leejet/stable-diffusion.cpp/pull/1193)
+
 * **2025/12/01** 🚀 stable-diffusion.cpp now supports **Z-Image**  
   👉 Details: [PR #1020](https://github.com/leejet/stable-diffusion.cpp/pull/1020)
 

docs/distilled_sd.md

@@ -1,8 +1,8 @@
-# Running distilled models: SSD1B and SDx.x with tiny U-Nets
+# Running distilled models: SSD1B, Vega and SDx.x with tiny U-Nets
 
 ## Preface 
 
-These models feature a reduced U-Net architecture. Unlike standard SDXL models, the SSD-1B U-Net contains only one middle block and fewer attention layers in its up- and down-blocks, resulting in significantly smaller file sizes. Using these models can reduce inference time by more than 33%. For more details, refer to Segmind's paper: https://arxiv.org/abs/2401.02677v1.
+These models feature a reduced U-Net architecture. Unlike standard SDXL models, the SSD-1B and Vega U-Net contains only one middle block and fewer attention layers in its up- and down-blocks, resulting in significantly smaller file sizes. Using these models can reduce inference time by more than 33%. For more details, refer to Segmind's paper: https://arxiv.org/abs/2401.02677v1.
 Similarly, SD1.x- and SD2.x-style models with a tiny U-Net consist of only 6 U-Net blocks, leading to very small files and time savings of up to 50%. For more information, see the paper: https://arxiv.org/pdf/2305.15798.pdf.
 
 ## SSD1B
@@ -17,7 +17,17 @@ Useful LoRAs are also available:
  * https://huggingface.co/seungminh/lora-swarovski-SSD-1B/resolve/main/pytorch_lora_weights.safetensors
  * https://huggingface.co/kylielee505/mylcmlorassd/resolve/main/pytorch_lora_weights.safetensors
 
-These files can be used out-of-the-box, unlike the models described in the next section.
+## Vega
+
+Segmind's Vega model is available online here:
+
+ * https://huggingface.co/segmind/Segmind-Vega/resolve/main/segmind-vega.safetensors
+ 
+VegaRT is an example for an LCM-LoRA:
+
+ * https://huggingface.co/segmind/Segmind-VegaRT/resolve/main/pytorch_lora_weights.safetensors
+
+Both files can be used out-of-the-box, unlike the models described in next sections.
 
 
 ## SD1.x, SD2.x with tiny U-Nets

docs/esrgan.md

@@ -1,6 +1,6 @@
 ## Using ESRGAN to upscale results
 
-You can use ESRGAN to upscale the generated images. At the moment, only the [RealESRGAN_x4plus_anime_6B.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth) model is supported. Support for more models of this architecture will be added soon.
+You can use ESRGAN—such as the model [RealESRGAN_x4plus_anime_6B.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth)—to upscale the generated images and improve their overall resolution and clarity.
 
 - Specify the model path using the `--upscale-model PATH` parameter. example:
 

examples/cli/main.cpp

@@ -245,7 +245,7 @@ std::string get_image_params(const SDCliParams& cli_params, const SDContextParam
     parameter_string += "Guidance: " + std::to_string(gen_params.sample_params.guidance.distilled_guidance) + ", ";
     parameter_string += "Eta: " + std::to_string(gen_params.sample_params.eta) + ", ";
     parameter_string += "Seed: " + std::to_string(seed) + ", ";
-    parameter_string += "Size: " + std::to_string(gen_params.width) + "x" + std::to_string(gen_params.height) + ", ";
+    parameter_string += "Size: " + std::to_string(gen_params.get_resolved_width()) + "x" + std::to_string(gen_params.get_resolved_height()) + ", ";
     parameter_string += "Model: " + sd_basename(ctx_params.model_path) + ", ";
     parameter_string += "RNG: " + std::string(sd_rng_type_name(ctx_params.rng_type)) + ", ";
     if (ctx_params.sampler_rng_type != RNG_TYPE_COUNT) {
@@ -526,10 +526,10 @@ int main(int argc, const char* argv[]) {
     }
 
     bool vae_decode_only     = true;
-    sd_image_t init_image    = {(uint32_t)gen_params.width, (uint32_t)gen_params.height, 3, nullptr};
-    sd_image_t end_image     = {(uint32_t)gen_params.width, (uint32_t)gen_params.height, 3, nullptr};
-    sd_image_t control_image = {(uint32_t)gen_params.width, (uint32_t)gen_params.height, 3, nullptr};
-    sd_image_t mask_image    = {(uint32_t)gen_params.width, (uint32_t)gen_params.height, 1, nullptr};
+    sd_image_t init_image    = {0, 0, 3, nullptr};
+    sd_image_t end_image     = {0, 0, 3, nullptr};
+    sd_image_t control_image = {0, 0, 3, nullptr};
+    sd_image_t mask_image    = {0, 0, 1, nullptr};
     std::vector<sd_image_t> ref_images;
     std::vector<sd_image_t> pmid_images;
     std::vector<sd_image_t> control_frames;
@@ -556,57 +556,79 @@ int main(int argc, const char* argv[]) {
         control_frames.clear();
     };
 
+    auto load_image_and_update_size = [&](const std::string& path,
+                                          sd_image_t& image,
+                                          bool resize_image    = true,
+                                          int expected_channel = 3) -> bool {
+        int expected_width  = 0;
+        int expected_height = 0;
+        if (resize_image && gen_params.width_and_height_are_set()) {
+            expected_width  = gen_params.width;
+            expected_height = gen_params.height;
+        }
+
+        if (!load_sd_image_from_file(&image, path.c_str(), expected_width, expected_height, expected_channel)) {
+            LOG_ERROR("load image from '%s' failed", path.c_str());
+            release_all_resources();
+            return false;
+        }
+
+        gen_params.set_width_and_height_if_unset(image.width, image.height);
+        return true;
+    };
+
     if (gen_params.init_image_path.size() > 0) {
         vae_decode_only = false;
-
-        int width       = 0;
-        int height      = 0;
-        init_image.data = load_image_from_file(gen_params.init_image_path.c_str(), width, height, gen_params.width, gen_params.height);
-        if (init_image.data == nullptr) {
-            LOG_ERROR("load image from '%s' failed", gen_params.init_image_path.c_str());
-            release_all_resources();
+        if (!load_image_and_update_size(gen_params.init_image_path, init_image)) {
             return 1;
         }
     }
 
     if (gen_params.end_image_path.size() > 0) {
         vae_decode_only = false;
-
-        int width      = 0;
-        int height     = 0;
-        end_image.data = load_image_from_file(gen_params.end_image_path.c_str(), width, height, gen_params.width, gen_params.height);
-        if (end_image.data == nullptr) {
-            LOG_ERROR("load image from '%s' failed", gen_params.end_image_path.c_str());
-            release_all_resources();
+        if (!load_image_and_update_size(gen_params.init_image_path, end_image)) {
             return 1;
         }
     }
 
+    if (gen_params.ref_image_paths.size() > 0) {
+        vae_decode_only = false;
+        for (auto& path : gen_params.ref_image_paths) {
+            sd_image_t ref_image = {0, 0, 3, nullptr};
+            if (!load_image_and_update_size(path, ref_image, false)) {
+                return 1;
+            }
+            ref_images.push_back(ref_image);
+        }
+    }
+
     if (gen_params.mask_image_path.size() > 0) {
-        int c           = 0;
-        int width       = 0;
-        int height      = 0;
-        mask_image.data = load_image_from_file(gen_params.mask_image_path.c_str(), width, height, gen_params.width, gen_params.height, 1);
-        if (mask_image.data == nullptr) {
+        if (load_sd_image_from_file(&mask_image,
+                                    gen_params.mask_image_path.c_str(),
+                                    gen_params.get_resolved_width(),
+                                    gen_params.get_resolved_height(),
+                                    1)) {
             LOG_ERROR("load image from '%s' failed", gen_params.mask_image_path.c_str());
             release_all_resources();
             return 1;
         }
     } else {
-        mask_image.data = (uint8_t*)malloc(gen_params.width * gen_params.height);
+        mask_image.data = (uint8_t*)malloc(gen_params.get_resolved_width() * gen_params.get_resolved_height());
         if (mask_image.data == nullptr) {
             LOG_ERROR("malloc mask image failed");
             release_all_resources();
             return 1;
         }
-        memset(mask_image.data, 255, gen_params.width * gen_params.height);
+        mask_image.width  = gen_params.get_resolved_width();
+        mask_image.height = gen_params.get_resolved_height();
+        memset(mask_image.data, 255, gen_params.get_resolved_width() * gen_params.get_resolved_height());
     }
 
     if (gen_params.control_image_path.size() > 0) {
-        int width          = 0;
-        int height         = 0;
-        control_image.data = load_image_from_file(gen_params.control_image_path.c_str(), width, height, gen_params.width, gen_params.height);
-        if (control_image.data == nullptr) {
+        if (load_sd_image_from_file(&control_image,
+                                    gen_params.control_image_path.c_str(),
+                                    gen_params.get_resolved_width(),
+                                    gen_params.get_resolved_height())) {
             LOG_ERROR("load image from '%s' failed", gen_params.control_image_path.c_str());
             release_all_resources();
             return 1;
@@ -621,29 +643,11 @@ int main(int argc, const char* argv[]) {
         }
     }
 
-    if (gen_params.ref_image_paths.size() > 0) {
-        vae_decode_only = false;
-        for (auto& path : gen_params.ref_image_paths) {
-            int width             = 0;
-            int height            = 0;
-            uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height);
-            if (image_buffer == nullptr) {
-                LOG_ERROR("load image from '%s' failed", path.c_str());
-                release_all_resources();
-                return 1;
-            }
-            ref_images.push_back({(uint32_t)width,
-                                  (uint32_t)height,
-                                  3,
-                                  image_buffer});
-        }
-    }
-
     if (!gen_params.control_video_path.empty()) {
         if (!load_images_from_dir(gen_params.control_video_path,
                                   control_frames,
-                                  gen_params.width,
-                                  gen_params.height,
+                                  gen_params.get_resolved_width(),
+                                  gen_params.get_resolved_height(),
                                   gen_params.video_frames,
                                   cli_params.verbose)) {
             release_all_resources();
@@ -717,8 +721,8 @@ int main(int argc, const char* argv[]) {
                 gen_params.auto_resize_ref_image,
                 gen_params.increase_ref_index,
                 mask_image,
-                gen_params.width,
-                gen_params.height,
+                gen_params.get_resolved_width(),
+                gen_params.get_resolved_height(),
                 gen_params.sample_params,
                 gen_params.strength,
                 gen_params.seed,
@@ -748,8 +752,8 @@ int main(int argc, const char* argv[]) {
                 end_image,
                 control_frames.data(),
                 (int)control_frames.size(),
-                gen_params.width,
-                gen_params.height,
+                gen_params.get_resolved_width(),
+                gen_params.get_resolved_height(),
                 gen_params.sample_params,
                 gen_params.high_noise_sample_params,
                 gen_params.moe_boundary,

examples/common/common.hpp

@@ -445,7 +445,7 @@ struct SDContextParams {
     std::string photo_maker_path;
     sd_type_t wtype = SD_TYPE_COUNT;
     std::string tensor_type_rules;
-    std::string lora_model_dir;
+    std::string lora_model_dir = ".";
 
     std::map<std::string, std::string> embedding_map;
     std::vector<sd_embedding_t> embedding_vec;
@@ -1024,8 +1024,8 @@ struct SDGenerationParams {
     std::string prompt_with_lora;  // for metadata record only
     std::string negative_prompt;
     int clip_skip   = -1;  // <= 0 represents unspecified
-    int width       = 512;
-    int height      = 512;
+    int width       = -1;
+    int height      = -1;
     int batch_count = 1;
     std::string init_image_path;
     std::string end_image_path;
@@ -1705,17 +1705,24 @@ struct SDGenerationParams {
         }
     }
 
+    bool width_and_height_are_set() const {
+        return width > 0 && height > 0;
+    }
+
+    void set_width_and_height_if_unset(int w, int h) {
+        if (!width_and_height_are_set()) {
+            LOG_INFO("set width x height to %d x %d", w, h);
+            width  = w;
+            height = h;
+        }
+    }
+
+    int get_resolved_width() const { return (width > 0) ? width : 512; }
+
+    int get_resolved_height() const { return (height > 0) ? height : 512; }
+
     bool process_and_check(SDMode mode, const std::string& lora_model_dir) {
         prompt_with_lora = prompt;
-        if (width <= 0) {
-            LOG_ERROR("error: the width must be greater than 0\n");
-            return false;
-        }
-
-        if (height <= 0) {
-            LOG_ERROR("error: the height must be greater than 0\n");
-            return false;
-        }
 
         if (sample_params.sample_steps <= 0) {
             LOG_ERROR("error: the sample_steps must be greater than 0\n");
@@ -2083,6 +2090,22 @@ uint8_t* load_image_from_file(const char* image_path,
     return load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel);
 }
 
+bool load_sd_image_from_file(sd_image_t* image,
+                             const char* image_path,
+                             int expected_width   = 0,
+                             int expected_height  = 0,
+                             int expected_channel = 3) {
+    int width;
+    int height;
+    image->data = load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel);
+    if (image->data == nullptr) {
+        return false;
+    }
+    image->width  = width;
+    image->height = height;
+    return true;
+}
+
 uint8_t* load_image_from_memory(const char* image_bytes,
                                 int len,
                                 int& width,

examples/server/main.cpp

@@ -86,21 +86,6 @@ std::vector<uint8_t> base64_decode(const std::string& encoded_string) {
     return ret;
 }
 
-std::string iso_timestamp_now() {
-    using namespace std::chrono;
-    auto now      = system_clock::now();
-    std::time_t t = system_clock::to_time_t(now);
-    std::tm tm{};
-#ifdef _MSC_VER
-    gmtime_s(&tm, &t);
-#else
-    gmtime_r(&t, &tm);
-#endif
-    std::ostringstream oss;
-    oss << std::put_time(&tm, "%Y-%m-%dT%H:%M:%SZ");
-    return oss.str();
-}
-
 struct SDSvrParams {
     std::string listen_ip = "127.0.0.1";
     int listen_port       = 1234;
@@ -404,7 +389,7 @@ int main(int argc, const char** argv) {
             }
 
             json out;
-            out["created"]       = iso_timestamp_now();
+            out["created"]       = static_cast<long long>(std::time(nullptr));
             out["data"]          = json::array();
             out["output_format"] = output_format;
 
@@ -692,7 +677,7 @@ int main(int argc, const char** argv) {
             }
 
             json out;
-            out["created"]       = iso_timestamp_now();
+            out["created"]       = static_cast<long long>(std::time(nullptr));
             out["data"]          = json::array();
             out["output_format"] = output_format;
 

flux.hpp

@@ -748,7 +748,7 @@ namespace Flux {
         int nerf_depth           = 4;
         int nerf_max_freqs       = 8;
         bool use_x0              = false;
-        bool use_patch_size_32   = false;
+        bool fake_patch_size_x2  = false;
     };
 
     struct FluxParams {
@@ -786,8 +786,11 @@ namespace Flux {
         Flux(FluxParams params)
             : params(params) {
             if (params.version == VERSION_CHROMA_RADIANCE) {
-                std::pair<int, int> kernel_size = {16, 16};
-                std::pair<int, int> stride      = kernel_size;
+                std::pair<int, int> kernel_size = {params.patch_size, params.patch_size};
+                if (params.chroma_radiance_params.fake_patch_size_x2) {
+                    kernel_size = {params.patch_size / 2, params.patch_size / 2};
+                }
+                std::pair<int, int> stride = kernel_size;
 
                 blocks["img_in_patch"] = std::make_shared<Conv2d>(params.in_channels,
                                                                   params.hidden_size,
@@ -1082,7 +1085,7 @@ namespace Flux {
             auto img      = pad_to_patch_size(ctx, x);
             auto orig_img = img;
 
-            if (params.chroma_radiance_params.use_patch_size_32) {
+            if (params.chroma_radiance_params.fake_patch_size_x2) {
                 // It's supposed to be using GGML_SCALE_MODE_NEAREST, but this seems more stable
                 // Maybe the implementation of nearest-neighbor interpolation in ggml behaves differently than the one in PyTorch?
                 // img = F.interpolate(img, size=(H//2, W//2), mode="nearest")
@@ -1303,7 +1306,8 @@ namespace Flux {
                 flux_params.ref_index_scale  = 10.f;
                 flux_params.use_mlp_silu_act = true;
             }
-            int64_t head_dim = 0;
+            int64_t head_dim                   = 0;
+            int64_t actual_radiance_patch_size = -1;
             for (auto pair : tensor_storage_map) {
                 std::string tensor_name = pair.first;
                 if (!starts_with(tensor_name, prefix))
@@ -1316,9 +1320,12 @@ namespace Flux {
                     flux_params.chroma_radiance_params.use_x0 = true;
                 }
                 if (tensor_name.find("__32x32__") != std::string::npos) {
-                    LOG_DEBUG("using patch size 32 prediction");
-                    flux_params.chroma_radiance_params.use_patch_size_32 = true;
-                    flux_params.patch_size                               = 32;
+                    LOG_DEBUG("using patch size 32");
+                    flux_params.patch_size = 32;
+                }
+                if (tensor_name.find("img_in_patch.weight") != std::string::npos) {
+                    actual_radiance_patch_size = pair.second.ne[0];
+                    LOG_DEBUG("actual radiance patch size: %d", actual_radiance_patch_size);
                 }
                 if (tensor_name.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) {
                     // Chroma
@@ -1351,6 +1358,11 @@ namespace Flux {
                     head_dim = pair.second.ne[0];
                 }
             }
+            if (actual_radiance_patch_size > 0 && actual_radiance_patch_size != flux_params.patch_size) {
+                GGML_ASSERT(flux_params.patch_size == 2 * actual_radiance_patch_size);
+                LOG_DEBUG("using fake x2 patch size");
+                flux_params.chroma_radiance_params.fake_patch_size_x2 = true;
+            }
 
             flux_params.num_heads = static_cast<int>(flux_params.hidden_size / head_dim);
 

model.cpp

@@ -376,7 +376,11 @@ bool ModelLoader::init_from_file(const std::string& file_path, const std::string
         LOG_INFO("load %s using checkpoint format", file_path.c_str());
         return init_from_ckpt_file(file_path, prefix);
     } else {
-        LOG_WARN("unknown format %s", file_path.c_str());
+        if (file_exists(file_path)) {
+            LOG_WARN("unknown format %s", file_path.c_str());
+        } else {
+            LOG_WARN("file %s not found", file_path.c_str());
+        }
         return false;
     }
 }
@@ -1040,6 +1044,7 @@ SDVersion ModelLoader::get_sd_version() {
     int64_t patch_embedding_channels = 0;
     bool has_img_emb                 = false;
     bool has_middle_block_1          = false;
+    bool has_output_block_311        = false;
     bool has_output_block_71         = false;
 
     for (auto& [name, tensor_storage] : tensor_storage_map) {
@@ -1100,6 +1105,9 @@ SDVersion ModelLoader::get_sd_version() {
             tensor_storage.name.find("unet.mid_block.resnets.1.") != std::string::npos) {
             has_middle_block_1 = true;
         }
+        if (tensor_storage.name.find("model.diffusion_model.output_blocks.3.1.transformer_blocks.1") != std::string::npos) {
+            has_output_block_311 = true;
+        }
         if (tensor_storage.name.find("model.diffusion_model.output_blocks.7.1") != std::string::npos) {
             has_output_block_71 = true;
         }
@@ -1138,6 +1146,9 @@ SDVersion ModelLoader::get_sd_version() {
             return VERSION_SDXL_PIX2PIX;
         }
         if (!has_middle_block_1) {
+            if (!has_output_block_311) {
+                return VERSION_SDXL_VEGA;
+            }
             return VERSION_SDXL_SSD1B;
         }
         return VERSION_SDXL;

model.h

@@ -32,6 +32,7 @@ enum SDVersion {
     VERSION_SDXL,
     VERSION_SDXL_INPAINT,
     VERSION_SDXL_PIX2PIX,
+    VERSION_SDXL_VEGA,
     VERSION_SDXL_SSD1B,
     VERSION_SVD,
     VERSION_SD3,
@@ -66,7 +67,7 @@ static inline bool sd_version_is_sd2(SDVersion version) {
 }
 
 static inline bool sd_version_is_sdxl(SDVersion version) {
-    if (version == VERSION_SDXL || version == VERSION_SDXL_INPAINT || version == VERSION_SDXL_PIX2PIX || version == VERSION_SDXL_SSD1B) {
+    if (version == VERSION_SDXL || version == VERSION_SDXL_INPAINT || version == VERSION_SDXL_PIX2PIX || version == VERSION_SDXL_SSD1B || version == VERSION_SDXL_VEGA) {
         return true;
     }
     return false;

stable-diffusion.cpp

@@ -35,6 +35,7 @@ const char* model_version_to_str[] = {
     "SDXL",
     "SDXL Inpaint",
     "SDXL Instruct-Pix2Pix",
+    "SDXL (Vega)",
     "SDXL (SSD1B)",
     "SVD",
     "SD3.x",
@@ -623,7 +624,7 @@ public:
                         LOG_INFO("Using Conv2d direct in the vae model");
                         first_stage_model->set_conv2d_direct_enabled(true);
                     }
-                    if (version == VERSION_SDXL &&
+                    if (sd_version_is_sdxl(version) &&
                         (strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale)) {
                         float vae_conv_2d_scale = 1.f / 32.f;
                         LOG_WARN(

tae.hpp

@@ -17,22 +17,43 @@ class TAEBlock : public UnaryBlock {
 protected:
     int n_in;
     int n_out;
+    bool use_midblock_gn;
 
 public:
-    TAEBlock(int n_in, int n_out)
-        : n_in(n_in), n_out(n_out) {
+    TAEBlock(int n_in, int n_out, bool use_midblock_gn = false)
+        : n_in(n_in), n_out(n_out), use_midblock_gn(use_midblock_gn) {
         blocks["conv.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_in, n_out, {3, 3}, {1, 1}, {1, 1}));
         blocks["conv.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_out, n_out, {3, 3}, {1, 1}, {1, 1}));
         blocks["conv.4"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_out, n_out, {3, 3}, {1, 1}, {1, 1}));
         if (n_in != n_out) {
             blocks["skip"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_in, n_out, {1, 1}, {1, 1}, {1, 1}, {1, 1}, false));
         }
+        if (use_midblock_gn) {
+            int n_gn         = n_in * 4;
+            blocks["pool.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_in, n_gn, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
+            blocks["pool.1"] = std::shared_ptr<GGMLBlock>(new GroupNorm(4, n_gn));
+            // pool.2 is ReLU, handled in forward
+            blocks["pool.3"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_gn, n_in, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
+        }
     }
 
     struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
         // x: [n, n_in, h, w]
         // return: [n, n_out, h, w]
 
+        if (use_midblock_gn) {
+            auto pool_0 = std::dynamic_pointer_cast<Conv2d>(blocks["pool.0"]);
+            auto pool_1 = std::dynamic_pointer_cast<GroupNorm>(blocks["pool.1"]);
+            auto pool_3 = std::dynamic_pointer_cast<Conv2d>(blocks["pool.3"]);
+
+            auto p = pool_0->forward(ctx, x);
+            p      = pool_1->forward(ctx, p);
+            p      = ggml_relu_inplace(ctx->ggml_ctx, p);
+            p      = pool_3->forward(ctx, p);
+
+            x = ggml_add(ctx->ggml_ctx, x, p);
+        }
+
         auto conv_0 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.0"]);
         auto conv_2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.2"]);
         auto conv_4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.4"]);
@@ -62,7 +83,7 @@ class TinyEncoder : public UnaryBlock {
     int num_blocks  = 3;
 
 public:
-    TinyEncoder(int z_channels = 4)
+    TinyEncoder(int z_channels = 4, bool use_midblock_gn = false)
         : 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}));
@@ -80,7 +101,7 @@ public:
 
         blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
         for (int i = 0; i < num_blocks; i++) {
-            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels, use_midblock_gn));
         }
 
         blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, z_channels, {3, 3}, {1, 1}, {1, 1}));
@@ -107,7 +128,7 @@ class TinyDecoder : public UnaryBlock {
     int num_blocks   = 3;
 
 public:
-    TinyDecoder(int z_channels = 4)
+    TinyDecoder(int z_channels = 4, bool use_midblock_gn = false)
         : z_channels(z_channels) {
         int index = 0;
 
@@ -115,7 +136,7 @@ public:
         index++;  // nn.ReLU()
 
         for (int i = 0; i < num_blocks; i++) {
-            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels, use_midblock_gn));
         }
         index++;  // nn.Upsample()
         blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
@@ -470,29 +491,44 @@ public:
 class TAESD : public GGMLBlock {
 protected:
     bool decode_only;
+    bool taef2 = false;
 
 public:
     TAESD(bool decode_only = true, SDVersion version = VERSION_SD1)
         : decode_only(decode_only) {
-        int z_channels = 4;
+        int z_channels       = 4;
+        bool use_midblock_gn = false;
+        taef2                = sd_version_is_flux2(version);
+
         if (sd_version_is_dit(version)) {
             z_channels = 16;
         }
-        blocks["decoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyDecoder(z_channels));
+        if (taef2) {
+            z_channels      = 32;
+            use_midblock_gn = true;
+        }
+        blocks["decoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyDecoder(z_channels, use_midblock_gn));
 
         if (!decode_only) {
-            blocks["encoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyEncoder(z_channels));
+            blocks["encoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyEncoder(z_channels, use_midblock_gn));
         }
     }
 
     struct ggml_tensor* decode(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
         auto decoder = std::dynamic_pointer_cast<TinyDecoder>(blocks["decoder.layers"]);
+        if (taef2) {
+            z = unpatchify(ctx->ggml_ctx, z, 2);
+        }
         return decoder->forward(ctx, z);
     }
 
     struct ggml_tensor* encode(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         auto encoder = std::dynamic_pointer_cast<TinyEncoder>(blocks["encoder.layers"]);
-        return encoder->forward(ctx, x);
+        auto z       = encoder->forward(ctx, x);
+        if (taef2) {
+            z = patchify(ctx->ggml_ctx, z, 2);
+        }
+        return z;
     }
 };
 

unet.hpp

@@ -201,6 +201,9 @@ public:
             num_head_channels     = 64;
             num_heads             = -1;
             use_linear_projection = true;
+            if (version == VERSION_SDXL_VEGA) {
+                transformer_depth = {1, 1, 2};
+            }
         } else if (version == VERSION_SVD) {
             in_channels           = 8;
             out_channels          = 4;
@@ -319,7 +322,7 @@ public:
         }
         if (!tiny_unet) {
             blocks["middle_block.0"] = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, ch));
-            if (version != VERSION_SDXL_SSD1B) {
+            if (version != VERSION_SDXL_SSD1B && version != VERSION_SDXL_VEGA) {
                 blocks["middle_block.1"] = std::shared_ptr<GGMLBlock>(get_attention_layer(ch,
                                                                                           n_head,
                                                                                           d_head,
@@ -520,7 +523,7 @@ public:
         // middle_block
         if (!tiny_unet) {
             h = resblock_forward("middle_block.0", ctx, h, emb, num_video_frames);  // [N, 4*model_channels, h/8, w/8]
-            if (version != VERSION_SDXL_SSD1B) {
+            if (version != VERSION_SDXL_SSD1B && version != VERSION_SDXL_VEGA) {
                 h = attention_layer_forward("middle_block.1", ctx, h, context, num_video_frames);  // [N, 4*model_channels, h/8, w/8]
                 h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames);             // [N, 4*model_channels, h/8, w/8]
             }