feat: add sd-webui style Hires. fix support (#1451)

CMakeLists.txt

@@ -156,10 +156,12 @@ endif()
 
 set(SD_LIB stable-diffusion)
 
-file(GLOB SD_LIB_SOURCES
+file(GLOB SD_LIB_SOURCES CONFIGURE_DEPENDS
     "src/*.h"
     "src/*.cpp"
     "src/*.hpp"
+    "src/model_io/*.h"
+    "src/model_io/*.cpp"
     "src/tokenizers/*.h"
     "src/tokenizers/*.cpp"
     "src/tokenizers/vocab/*.h"

README.md

@@ -77,9 +77,10 @@ API and command-line option may change frequently.***
   - OpenCL
   - SYCL
 - Supported weight formats
-  - Pytorch checkpoint (`.ckpt` or `.pth`)
+  - Pytorch checkpoint (`.ckpt` or `.pth` or `.pt`)
   - Safetensors (`.safetensors`)
   - GGUF (`.gguf`)
+- Convert mode supports converting model weights to `.gguf` or `.safetensors`
 - Supported platforms
     - Linux
     - Mac OS

examples/cli/README.md

@@ -14,6 +14,9 @@ CLI Options:
   --metadata-format <string>  metadata output format, one of [text, json] (default: text)
   --canny                     apply canny preprocessor (edge detection)
   --convert-name              convert tensor name (for convert mode)
+                              convert mode writes `.gguf` or `.safetensors` based on the output extension.
+                              `.safetensors` export currently supports f16, bf16, f32, and i32 tensor types only.
+                              i32 is passthrough only; no f32 <-> i32 conversion is performed
   -v, --verbose               print extra info
   --color                     colors the logging tags according to level
   --taesd-preview-only        prevents usage of taesd for decoding the final image. (for use with --preview tae)

examples/cli/main.cpp

@@ -278,7 +278,9 @@ void parse_args(int argc, const char** argv, SDCliParams& cli_params, SDContextP
     bool valid = cli_params.resolve_and_validate();
     if (valid && cli_params.mode != METADATA) {
         valid = ctx_params.resolve_and_validate(cli_params.mode) &&
-                gen_params.resolve_and_validate(cli_params.mode, ctx_params.lora_model_dir);
+                gen_params.resolve_and_validate(cli_params.mode,
+                                                ctx_params.lora_model_dir,
+                                                ctx_params.hires_upscalers_dir);
     }
 
     if (!valid) {
@@ -688,6 +690,10 @@ int main(int argc, const char* argv[]) {
         vae_decode_only = false;
     }
 
+    if (gen_params.hires_enabled && !gen_params.hires_upscaler_model_path.empty()) {
+        vae_decode_only = false;
+    }
+
     sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(vae_decode_only, true, cli_params.taesd_preview);
 
     SDImageVec results;

examples/common/common.cpp

@@ -351,7 +351,10 @@ ArgOptions SDContextParams::get_options() {
          "--lora-model-dir",
          "lora model directory",
          &lora_model_dir},
-
+        {"",
+         "--hires-upscalers-dir",
+         "highres fix upscaler model directory",
+         &hires_upscalers_dir},
         {"",
          "--tensor-type-rules",
          "weight type per tensor pattern (example: \"^vae\\.=f16,model\\.=q8_0\")",
@@ -649,6 +652,7 @@ std::string SDContextParams::to_string() const {
         << "  wtype: " << sd_type_name(wtype) << ",\n"
         << "  tensor_type_rules: \"" << tensor_type_rules << "\",\n"
         << "  lora_model_dir: \"" << lora_model_dir << "\",\n"
+        << "  hires_upscalers_dir: \"" << hires_upscalers_dir << "\",\n"
         << "  photo_maker_path: \"" << photo_maker_path << "\",\n"
         << "  rng_type: " << sd_rng_type_name(rng_type) << ",\n"
         << "  sampler_rng_type: " << sd_rng_type_name(sampler_rng_type) << ",\n"
@@ -777,6 +781,10 @@ ArgOptions SDGenerationParams::get_options() {
          "--pm-id-embed-path",
          "path to PHOTOMAKER v2 id embed",
          &pm_id_embed_path},
+        {"",
+         "--hires-upscaler",
+         "highres fix upscaler, Latent (nearest) or a model name/path under --hires-upscalers-dir (default: Latent (nearest))",
+         &hires_upscaler},
     };
 
     options.int_options = {
@@ -826,6 +834,22 @@ ArgOptions SDGenerationParams::get_options() {
          "--upscale-tile-size",
          "tile size for ESRGAN upscaling (default: 128)",
          &upscale_tile_size},
+        {"",
+         "--hires-width",
+         "highres fix target width, 0 to use --hires-scale (default: 0)",
+         &hires_width},
+        {"",
+         "--hires-height",
+         "highres fix target height, 0 to use --hires-scale (default: 0)",
+         &hires_height},
+        {"",
+         "--hires-steps",
+         "highres fix second pass sample steps, 0 to reuse --steps (default: 0)",
+         &hires_steps},
+        {"",
+         "--hires-upscale-tile-size",
+         "highres fix upscaler tile size, reserved for model-backed upscalers (default: 128)",
+         &hires_upscale_tile_size},
     };
 
     options.float_options = {
@@ -913,6 +937,14 @@ ArgOptions SDGenerationParams::get_options() {
          "--vae-tile-overlap",
          "tile overlap for vae tiling, in fraction of tile size (default: 0.5)",
          &vae_tiling_params.target_overlap},
+        {"",
+         "--hires-scale",
+         "highres fix scale when target size is not set (default: 2.0)",
+         &hires_scale},
+        {"",
+         "--hires-denoising-strength",
+         "highres fix second pass denoising strength (default: 0.7)",
+         &hires_denoising_strength},
     };
 
     options.bool_options = {
@@ -936,6 +968,11 @@ ArgOptions SDGenerationParams::get_options() {
          "process vae in tiles to reduce memory usage",
          true,
          &vae_tiling_params.enabled},
+        {"",
+         "--hires",
+         "enable highres fix",
+         true,
+         &hires_enabled},
     };
 
     auto on_seed_arg = [&](int argc, const char** argv, int index) {
@@ -1424,6 +1461,37 @@ static bool parse_lora_json_field(const json& parent,
     return true;
 }
 
+static bool resolve_model_file_from_dir(const std::string& model_name,
+                                        const std::string& model_dir,
+                                        const std::vector<std::string>& valid_ext,
+                                        const char* label,
+                                        std::string& resolved_path) {
+    if (model_dir.empty()) {
+        LOG_ERROR("%s directory is empty", label);
+        return false;
+    }
+    if (model_name.empty() ||
+        model_name.find('/') != std::string::npos ||
+        model_name.find('\\') != std::string::npos ||
+        fs::path(model_name).has_root_path() ||
+        fs::path(model_name).has_extension()) {
+        LOG_ERROR("%s must be a model name without path or extension: %s", label, model_name.c_str());
+        return false;
+    }
+
+    fs::path model_dir_path = model_dir;
+    for (const auto& ext : valid_ext) {
+        fs::path try_path = model_dir_path / (model_name + ext);
+        if (fs::exists(try_path) && fs::is_regular_file(try_path)) {
+            resolved_path = try_path.lexically_normal().string();
+            return true;
+        }
+    }
+
+    LOG_ERROR("can not find %s %s in %s", label, model_name.c_str(), model_dir_path.lexically_normal().string().c_str());
+    return false;
+}
+
 bool SDGenerationParams::from_json_str(
     const std::string& json_str,
     const std::function<std::string(const std::string&)>& lora_path_resolver) {
@@ -1487,6 +1555,34 @@ bool SDGenerationParams::from_json_str(
     load_if_exists("increase_ref_index", increase_ref_index);
     load_if_exists("embed_image_metadata", embed_image_metadata);
 
+    if (j.contains("hires") && j["hires"].is_object()) {
+        const json& hires_json = j["hires"];
+        if (hires_json.contains("enabled") && hires_json["enabled"].is_boolean()) {
+            hires_enabled = hires_json["enabled"];
+        }
+        if (hires_json.contains("upscaler") && hires_json["upscaler"].is_string()) {
+            hires_upscaler = hires_json["upscaler"];
+        }
+        if (hires_json.contains("scale") && hires_json["scale"].is_number()) {
+            hires_scale = hires_json["scale"];
+        }
+        if (hires_json.contains("target_width") && hires_json["target_width"].is_number_integer()) {
+            hires_width = hires_json["target_width"];
+        }
+        if (hires_json.contains("target_height") && hires_json["target_height"].is_number_integer()) {
+            hires_height = hires_json["target_height"];
+        }
+        if (hires_json.contains("steps") && hires_json["steps"].is_number_integer()) {
+            hires_steps = hires_json["steps"];
+        }
+        if (hires_json.contains("denoising_strength") && hires_json["denoising_strength"].is_number()) {
+            hires_denoising_strength = hires_json["denoising_strength"];
+        }
+        if (hires_json.contains("upscale_tile_size") && hires_json["upscale_tile_size"].is_number_integer()) {
+            hires_upscale_tile_size = hires_json["upscale_tile_size"];
+        }
+    }
+
     auto parse_sample_params_json = [&](const json& sample_json,
                                         sd_sample_params_t& target_params,
                                         std::vector<int>& target_skip_layers,
@@ -1800,7 +1896,7 @@ bool SDGenerationParams::initialize_cache_params() {
     return true;
 }
 
-bool SDGenerationParams::resolve(const std::string& lora_model_dir, bool strict) {
+bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict) {
     if (high_noise_sample_params.sample_steps <= 0) {
         high_noise_sample_params.sample_steps = -1;
     }
@@ -1819,6 +1915,27 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, bool strict)
         sample_params.sample_steps = std::clamp(sample_params.sample_steps, 1, 100);
     }
 
+    hires_upscaler_model_path.clear();
+    if (hires_enabled) {
+        if (hires_upscaler.empty()) {
+            hires_upscaler = "Latent (nearest)";
+        }
+        resolved_hires_upscaler = str_to_sd_hires_upscaler(hires_upscaler.c_str());
+        if (resolved_hires_upscaler == SD_HIRES_UPSCALER_NONE) {
+            hires_enabled = false;
+        } else if (resolved_hires_upscaler == SD_HIRES_UPSCALER_COUNT) {
+            static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt", ".pth"};
+            if (!resolve_model_file_from_dir(hires_upscaler,
+                                             hires_upscalers_dir,
+                                             valid_ext,
+                                             "hires upscaler",
+                                             hires_upscaler_model_path)) {
+                return false;
+            }
+            resolved_hires_upscaler = SD_HIRES_UPSCALER_MODEL;
+        }
+    }
+
     prompt_with_lora = prompt;
     if (!lora_model_dir.empty()) {
         extract_and_remove_lora(lora_model_dir);
@@ -1883,6 +2000,29 @@ bool SDGenerationParams::validate(SDMode mode) {
         return false;
     }
 
+    if (hires_enabled) {
+        if (hires_width < 0 || hires_height < 0) {
+            LOG_ERROR("error: hires target width and height must be >= 0");
+            return false;
+        }
+        if (hires_scale <= 0.f && hires_width <= 0 && hires_height <= 0) {
+            LOG_ERROR("error: hires scale must be positive when target size is not set");
+            return false;
+        }
+        if (hires_steps < 0) {
+            LOG_ERROR("error: hires steps must be >= 0");
+            return false;
+        }
+        if (hires_denoising_strength <= 0.f || hires_denoising_strength > 1.f) {
+            LOG_ERROR("error: hires denoising strength must be in (0.0, 1.0]");
+            return false;
+        }
+        if (hires_upscale_tile_size < 1) {
+            LOG_ERROR("error: hires upscale tile size must be positive");
+            return false;
+        }
+    }
+
     if (mode == UPSCALE) {
         if (init_image_path.length() == 0) {
             LOG_ERROR("error: upscale mode needs an init image (--init-img)\n");
@@ -1893,8 +2033,11 @@ bool SDGenerationParams::validate(SDMode mode) {
     return true;
 }
 
-bool SDGenerationParams::resolve_and_validate(SDMode mode, const std::string& lora_model_dir, bool strict) {
-    if (!resolve(lora_model_dir, strict)) {
+bool SDGenerationParams::resolve_and_validate(SDMode mode,
+                                              const std::string& lora_model_dir,
+                                              const std::string& hires_upscalers_dir,
+                                              bool strict) {
+    if (!resolve(lora_model_dir, hires_upscalers_dir, strict)) {
         return false;
     }
     if (!validate(mode)) {
@@ -1965,6 +2108,16 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     params.pm_params             = pm_params;
     params.vae_tiling_params     = vae_tiling_params;
     params.cache                 = cache_params;
+
+    params.hires.enabled            = hires_enabled;
+    params.hires.upscaler           = resolved_hires_upscaler;
+    params.hires.model_path         = hires_upscaler_model_path.empty() ? nullptr : hires_upscaler_model_path.c_str();
+    params.hires.scale              = hires_scale;
+    params.hires.target_width       = hires_width;
+    params.hires.target_height      = hires_height;
+    params.hires.steps              = hires_steps;
+    params.hires.denoising_strength = hires_denoising_strength;
+    params.hires.upscale_tile_size  = hires_upscale_tile_size;
     return params;
 }
 
@@ -2089,6 +2242,15 @@ std::string SDGenerationParams::to_string() const {
         << "  seed: " << seed << ",\n"
         << "  upscale_repeats: " << upscale_repeats << ",\n"
         << "  upscale_tile_size: " << upscale_tile_size << ",\n"
+        << "  hires: { enabled: " << (hires_enabled ? "true" : "false")
+        << ", upscaler: \"" << hires_upscaler << "\""
+        << ", model_path: \"" << hires_upscaler_model_path << "\""
+        << ", scale: " << hires_scale
+        << ", target_width: " << hires_width
+        << ", target_height: " << hires_height
+        << ", steps: " << hires_steps
+        << ", denoising_strength: " << hires_denoising_strength
+        << ", upscale_tile_size: " << hires_upscale_tile_size << " },\n"
         << "  vae_tiling_params: { "
         << vae_tiling_params.enabled << ", "
         << vae_tiling_params.tile_size_x << ", "
@@ -2162,6 +2324,13 @@ std::string get_image_params(const SDContextParams& ctx_params, const SDGenerati
     if (gen_params.clip_skip != -1) {
         parameter_string += "Clip skip: " + std::to_string(gen_params.clip_skip) + ", ";
     }
+    if (gen_params.hires_enabled) {
+        parameter_string += "Hires upscale: " + gen_params.hires_upscaler + ", ";
+        parameter_string += "Hires scale: " + std::to_string(gen_params.hires_scale) + ", ";
+        parameter_string += "Hires resize: " + std::to_string(gen_params.hires_width) + "x" + std::to_string(gen_params.hires_height) + ", ";
+        parameter_string += "Hires steps: " + std::to_string(gen_params.hires_steps) + ", ";
+        parameter_string += "Denoising strength: " + std::to_string(gen_params.hires_denoising_strength) + ", ";
+    }
     parameter_string += "Version: stable-diffusion.cpp";
     return parameter_string;
 }

examples/common/common.h

@@ -101,6 +101,7 @@ struct SDContextParams {
     sd_type_t wtype = SD_TYPE_COUNT;
     std::string tensor_type_rules;
     std::string lora_model_dir = ".";
+    std::string hires_upscalers_dir;
 
     std::map<std::string, std::string> embedding_map;
     std::vector<sd_embedding_t> embedding_vec;
@@ -190,12 +191,23 @@ struct SDGenerationParams {
     int upscale_repeats   = 1;
     int upscale_tile_size = 128;
 
+    bool hires_enabled         = false;
+    std::string hires_upscaler = "Latent (nearest)";
+    std::string hires_upscaler_model_path;
+    float hires_scale              = 2.f;
+    int hires_width                = 0;
+    int hires_height               = 0;
+    int hires_steps                = 0;
+    float hires_denoising_strength = 0.7f;
+    int hires_upscale_tile_size    = 128;
+
     std::map<std::string, float> lora_map;
     std::map<std::string, float> high_noise_lora_map;
 
     // Derived and normalized fields.
     std::string prompt_with_lora;  // for metadata record only
     std::vector<sd_lora_t> lora_vec;
+    sd_hires_upscaler_t resolved_hires_upscaler;
 
     // Owned execution payload.
     SDImageOwner init_image;
@@ -225,9 +237,12 @@ struct SDGenerationParams {
     void set_width_and_height_if_unset(int w, int h);
     int get_resolved_width() const;
     int get_resolved_height() const;
-    bool resolve(const std::string& lora_model_dir, bool strict = false);
+    bool resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict = false);
     bool validate(SDMode mode);
-    bool resolve_and_validate(SDMode mode, const std::string& lora_model_dir, bool strict = false);
+    bool resolve_and_validate(SDMode mode,
+                              const std::string& lora_model_dir,
+                              const std::string& hires_upscalers_dir,
+                              bool strict = false);
     sd_img_gen_params_t to_sd_img_gen_params_t();
     sd_vid_gen_params_t to_sd_vid_gen_params_t();
     std::string to_string() const;

examples/server/api.md

@@ -38,6 +38,8 @@ Current generation-related endpoints include:
 - `POST /sdapi/v1/txt2img`
 - `POST /sdapi/v1/img2img`
 - `GET /sdapi/v1/loras`
+- `GET /sdapi/v1/upscalers`
+- `GET /sdapi/v1/latent-upscale-modes`
 - `GET /sdapi/v1/samplers`
 - `GET /sdapi/v1/schedulers`
 - `GET /sdapi/v1/sd-models`
@@ -216,6 +218,13 @@ Currently supported request fields:
 | `scheduler` | `string` | Scheduler name |
 | `lora` | `array<object>` | Structured LoRA list |
 | `extra_images` | `array<string>` | Base64 or data URL images |
+| `enable_hr` | `boolean` | Enable highres fix for `txt2img` |
+| `hr_upscaler` | `string` | `Latent (nearest)` or an upscaler model name from `/sdapi/v1/upscalers` |
+| `hr_scale` | `number` | Highres scale when resize target is not set |
+| `hr_resize_x` | `integer` | Highres target width, `0` to use scale |
+| `hr_resize_y` | `integer` | Highres target height, `0` to use scale |
+| `hr_steps` | `integer` | Highres second-pass sample steps, `0` to reuse `steps` |
+| `denoising_strength` | `number` | Highres denoising strength for `txt2img` |
 
 Native extension fields:
 
@@ -241,6 +250,8 @@ Currently supported request fields:
 | `inpainting_mask_invert` | `integer` or `boolean` | Treated as invert flag |
 | `denoising_strength` | `number` | Clamped to `0.0..1.0` |
 
+Highres fix fields are currently handled for `txt2img`; `img2img` uses `denoising_strength` as image-to-image strength.
+
 Native extension fields:
 
 - any `sdcpp API` fields embedded through `sd_cpp_extra_args` inside `prompt`
@@ -258,6 +269,8 @@ Response fields:
 Currently exposed:
 
 - `GET /sdapi/v1/loras`
+- `GET /sdapi/v1/upscalers`
+- `GET /sdapi/v1/latent-upscale-modes`
 - `GET /sdapi/v1/samplers`
 - `GET /sdapi/v1/schedulers`
 - `GET /sdapi/v1/sd-models`
@@ -272,6 +285,24 @@ Response fields:
 | `[].name` | `string` | Display name derived from file stem |
 | `[].path` | `string` | Relative path under the configured LoRA directory |
 
+`GET /sdapi/v1/upscalers`
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `[].name` | `string` | Built-in name or model stem |
+| `[].model_name` | `string \| null` | Model family label for model-backed upscalers |
+| `[].model_path` | `string \| null` | Absolute model path for model-backed upscalers |
+| `[].model_url` | `string \| null` | Currently always null |
+| `[].scale` | `integer` | Currently `4` |
+
+Built-in entries include `None`, `Lanczos`, and `Nearest`. Model-backed entries are scanned from the top level of `--hires-upscalers-dir`; subdirectories are not scanned.
+
+`GET /sdapi/v1/latent-upscale-modes`
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `[].name` | `string` | WebUI-compatible latent upscale mode name |
+
 `GET /sdapi/v1/samplers`
 
 | Field | Type | Notes |
@@ -388,6 +419,7 @@ Top-level fields:
 | `samplers` | `array<string>` | Available sampling methods |
 | `schedulers` | `array<string>` | Available schedulers |
 | `loras` | `array<object>` | Available LoRA entries |
+| `upscalers` | `array<object>` | Available model-backed highres upscalers |
 | `limits` | `object` | Shared queue and size limits |
 
 `model`
@@ -424,6 +456,14 @@ Shared nested fields:
 | `loras[].name` | `string` |
 | `loras[].path` | `string` |
 
+`upscalers`
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `upscalers[].name` | `string` | Built-in name or model stem; use this value in `hires.upscaler` |
+
+Built-in entries include `None` and `Latent (nearest)`. Model-backed entries are scanned from the top level of `--hires-upscalers-dir`; subdirectories are not scanned.
+
 `limits`
 
 | Field | Type |
@@ -482,6 +522,15 @@ Shared default fields used by both `img_gen` and `vid_gen`:
 | `auto_resize_ref_image` | `boolean` |
 | `increase_ref_index` | `boolean` |
 | `control_strength` | `number` |
+| `hires` | `object` |
+| `hires.enabled` | `boolean` |
+| `hires.upscaler` | `string` |
+| `hires.scale` | `number` |
+| `hires.target_width` | `integer` |
+| `hires.target_height` | `integer` |
+| `hires.steps` | `integer` |
+| `hires.denoising_strength` | `number` |
+| `hires.upscale_tile_size` | `integer` |
 
 `vid_gen`-specific default fields:
 
@@ -514,6 +563,7 @@ Fields returned in `features_by_mode.img_gen`:
 - `ref_images`
 - `lora`
 - `vae_tiling`
+- `hires`
 - `cache`
 - `cancel_queued`
 - `cancel_generating`
@@ -625,6 +675,16 @@ Example:
   },
 
   "lora": [],
+  "hires": {
+    "enabled": false,
+    "upscaler": "Latent (nearest)",
+    "scale": 2.0,
+    "target_width": 0,
+    "target_height": 0,
+    "steps": 0,
+    "denoising_strength": 0.7,
+    "upscale_tile_size": 128
+  },
 
   "vae_tiling_params": {
     "enabled": false,
@@ -729,12 +789,23 @@ Other native fields:
 
 | Field | Type |
 | --- | --- |
+| `hires` | `object` |
+| `hires.enabled` | `boolean` |
+| `hires.upscaler` | `string` |
+| `hires.scale` | `number` |
+| `hires.target_width` | `integer` |
+| `hires.target_height` | `integer` |
+| `hires.steps` | `integer` |
+| `hires.denoising_strength` | `number` |
+| `hires.upscale_tile_size` | `integer` |
 | `vae_tiling_params` | `object` |
 | `cache_mode` | `string` |
 | `cache_option` | `string` |
 | `scm_mask` | `string` |
 | `scm_policy_dynamic` | `boolean` |
 
+For `hires.upscaler`, use `Latent (nearest)` for latent upscale or an `upscalers[].name` value from `GET /sdcpp/v1/capabilities`. Model-backed upscalers are resolved as `--hires-upscalers-dir / (name + ext)` and must live directly in that directory.
+
 HTTP-only output fields:
 
 | Field | Type |

examples/server/main.cpp

@@ -48,7 +48,9 @@ static void parse_args(int argc,
 
     if (!svr_params.resolve_and_validate() ||
         !ctx_params.resolve_and_validate(IMG_GEN) ||
-        !default_gen_params.resolve_and_validate(IMG_GEN, ctx_params.lora_model_dir)) {
+        !default_gen_params.resolve_and_validate(IMG_GEN,
+                                                 ctx_params.lora_model_dir,
+                                                 ctx_params.hires_upscalers_dir)) {
         print_usage(argv[0], options_vec);
         exit(1);
     }
@@ -95,6 +97,8 @@ int main(int argc, const char** argv) {
 
     std::vector<LoraEntry> lora_cache;
     std::mutex lora_mutex;
+    std::vector<UpscalerEntry> upscaler_cache;
+    std::mutex upscaler_mutex;
     AsyncJobManager async_job_manager;
     ServerRuntime runtime = {
         sd_ctx.get(),
@@ -104,6 +108,8 @@ int main(int argc, const char** argv) {
         &default_gen_params,
         &lora_cache,
         &lora_mutex,
+        &upscaler_cache,
+        &upscaler_mutex,
         &async_job_manager,
     };
 

examples/server/routes_openai.cpp

@@ -70,7 +70,7 @@ static bool build_openai_generation_request(const httplib::Request& req,
     }
 
     // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
-    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", true)) {
+    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) {
         error_message = "invalid params";
         return false;
     }
@@ -212,7 +212,7 @@ static bool build_openai_edit_request(const httplib::Request& req,
     }
 
     // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
-    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", true)) {
+    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) {
         error_message = "invalid params";
         return false;
     }

examples/server/routes_sdapi.cpp

@@ -1,6 +1,7 @@
 #include "routes.h"
 
 #include <algorithm>
+#include <cctype>
 #include <cstring>
 #include <regex>
 #include <string_view>
@@ -35,14 +36,20 @@ static fs::path resolve_display_model_path(const ServerRuntime& runtime) {
     return {};
 }
 
+static std::string lower_ascii(std::string value) {
+    std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) {
+        return static_cast<char>(std::tolower(c));
+    });
+    return value;
+}
+
 static enum sample_method_t get_sdapi_sample_method(std::string name) {
     enum sample_method_t result = str_to_sample_method(name.c_str());
     if (result != SAMPLE_METHOD_COUNT) {
         return result;
     }
 
-    std::transform(name.begin(), name.end(), name.begin(),
-                   [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
+    name = lower_ascii(name);
     static const std::unordered_map<std::string_view, sample_method_t> hardcoded{
         {"euler a", EULER_A_SAMPLE_METHOD},
         {"k_euler_a", EULER_A_SAMPLE_METHOD},
@@ -114,6 +121,18 @@ static bool build_sdapi_img_gen_request(const json& j,
     request.gen_params.width                          = j.value("width", -1);
     request.gen_params.height                         = j.value("height", -1);
 
+    if (!img2img && j.value("enable_hr", false)) {
+        request.gen_params.hires_enabled = true;
+        request.gen_params.hires_scale   = j.value("hr_scale", request.gen_params.hires_scale);
+        request.gen_params.hires_width   = j.value("hr_resize_x", request.gen_params.hires_width);
+        request.gen_params.hires_height  = j.value("hr_resize_y", request.gen_params.hires_height);
+        request.gen_params.hires_steps   = j.value("hr_steps", request.gen_params.hires_steps);
+        request.gen_params.hires_denoising_strength =
+            j.value("denoising_strength", request.gen_params.hires_denoising_strength);
+
+        request.gen_params.hires_upscaler = j.value("hr_upscaler", request.gen_params.hires_upscaler);
+    }
+
     std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt);
     if (!sd_cpp_extra_args_str.empty() && !request.gen_params.from_json_str(sd_cpp_extra_args_str)) {
         error_message = "invalid sd_cpp_extra_args";
@@ -228,7 +247,7 @@ static bool build_sdapi_img_gen_request(const json& j,
     }
 
     // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
-    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", true)) {
+    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) {
         error_message = "invalid params";
         return false;
     }
@@ -347,6 +366,45 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
         res.set_content(result.dump(), "application/json");
     });
 
+    svr.Get("/sdapi/v1/upscalers", [runtime](const httplib::Request&, httplib::Response& res) {
+        refresh_upscaler_cache(*runtime);
+
+        auto make_builtin = [](const char* name) {
+            json item;
+            item["name"]       = name;
+            item["model_name"] = nullptr;
+            item["model_path"] = nullptr;
+            item["model_url"]  = nullptr;
+            item["scale"]      = 4;
+            return item;
+        };
+
+        json result = json::array();
+        result.push_back(make_builtin("None"));
+
+        {
+            std::lock_guard<std::mutex> lock(*runtime->upscaler_mutex);
+            for (const auto& e : *runtime->upscaler_cache) {
+                json item;
+                item["name"]       = e.name;
+                item["model_name"] = e.model_name;
+                item["model_path"] = e.fullpath;
+                item["model_url"]  = nullptr;
+                item["scale"]      = e.scale;
+                result.push_back(item);
+            }
+        }
+
+        res.set_content(result.dump(), "application/json");
+    });
+
+    svr.Get("/sdapi/v1/latent-upscale-modes", [](const httplib::Request&, httplib::Response& res) {
+        json result = json::array({
+            {{"name", "Latent (nearest)"}},
+        });
+        res.set_content(result.dump(), "application/json");
+    });
+
     svr.Get("/sdapi/v1/samplers", [runtime](const httplib::Request&, httplib::Response& res) {
         std::vector<std::string> sampler_names;
         sampler_names.push_back("default");

examples/server/routes_sdcpp.cpp

@@ -114,6 +114,17 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
         {"increase_ref_index", defaults.increase_ref_index},
         {"control_strength", defaults.control_strength},
         {"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers)},
+        {"hires",
+         {
+             {"enabled", defaults.hires_enabled},
+             {"upscaler", defaults.hires_upscaler},
+             {"scale", defaults.hires_scale},
+             {"target_width", defaults.hires_width},
+             {"target_height", defaults.hires_height},
+             {"steps", defaults.hires_steps},
+             {"denoising_strength", defaults.hires_denoising_strength},
+             {"upscale_tile_size", defaults.hires_upscale_tile_size},
+         }},
         {"vae_tiling_params", make_vae_tiling_json(defaults.vae_tiling_params)},
         {"cache_mode", defaults.cache_mode},
         {"cache_option", defaults.cache_option},
@@ -157,6 +168,7 @@ static json make_img_gen_features_json() {
         {"ref_images", true},
         {"lora", true},
         {"vae_tiling", true},
+        {"hires", true},
         {"cache", true},
         {"cancel_queued", true},
         {"cancel_generating", false},
@@ -179,6 +191,7 @@ static json make_vid_gen_features_json() {
 
 static json make_capabilities_json(ServerRuntime& runtime) {
     refresh_lora_cache(runtime);
+    refresh_upscaler_cache(runtime);
 
     AsyncJobManager& manager  = *runtime.async_job_manager;
     const auto& defaults      = *runtime.default_gen_params;
@@ -190,6 +203,7 @@ static json make_capabilities_json(ServerRuntime& runtime) {
     json image_output_formats = supported_img_output_formats();
     json video_output_formats = supported_vid_output_formats();
     json available_loras      = json::array();
+    json available_upscalers  = json::array();
     json supported_modes      = json::array();
 
     for (int i = 0; i < SAMPLE_METHOD_COUNT; ++i) {
@@ -210,6 +224,21 @@ static json make_capabilities_json(ServerRuntime& runtime) {
         }
     }
 
+    available_upscalers.push_back({
+        {"name", "None"},
+    });
+    available_upscalers.push_back({
+        {"name", "Latent (nearest)"},
+    });
+    {
+        std::lock_guard<std::mutex> lock(*runtime.upscaler_mutex);
+        for (const auto& entry : *runtime.upscaler_cache) {
+            available_upscalers.push_back({
+                {"name", entry.name},
+            });
+        }
+    }
+
     if (supports_img) {
         supported_modes.push_back("img_gen");
     }
@@ -284,6 +313,7 @@ static json make_capabilities_json(ServerRuntime& runtime) {
     result["features"]               = top_level_features;
     result["features_by_mode"]       = features_by_mode;
     result["loras"]                  = available_loras;
+    result["upscalers"]              = available_upscalers;
     return result;
 }
 
@@ -307,7 +337,7 @@ static bool parse_img_gen_request(const json& body,
         return false;
     }
     // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
-    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", true)) {
+    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) {
         error_message = "invalid generation parameters";
         return false;
     }
@@ -334,7 +364,7 @@ static bool parse_vid_gen_request(const json& body,
         return false;
     }
     // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
-    if (!request.gen_params.resolve_and_validate(VID_GEN, "", true)) {
+    if (!request.gen_params.resolve_and_validate(VID_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) {
         error_message = "invalid generation parameters";
         return false;
     }

examples/server/runtime.cpp

@@ -1,6 +1,7 @@
 #include "runtime.h"
 
 #include <algorithm>
+#include <cctype>
 #include <chrono>
 #include <cstdlib>
 #include <filesystem>
@@ -13,6 +14,18 @@
 
 namespace fs = std::filesystem;
 
+static std::string lower_ascii(std::string value) {
+    std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) {
+        return static_cast<char>(std::tolower(c));
+    });
+    return value;
+}
+
+static bool is_supported_model_ext(const fs::path& p) {
+    auto ext = lower_ascii(p.extension().string());
+    return ext == ".gguf" || ext == ".pt" || ext == ".pth" || ext == ".safetensors";
+}
+
 static const std::string k_base64_chars =
     "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
     "abcdefghijklmnopqrstuvwxyz"
@@ -241,20 +254,12 @@ void refresh_lora_cache(ServerRuntime& rt) {
 
     fs::path lora_dir = rt.ctx_params->lora_model_dir;
     if (fs::exists(lora_dir) && fs::is_directory(lora_dir)) {
-        auto is_lora_ext = [](const fs::path& p) {
-            auto ext = p.extension().string();
-            std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) {
-                return static_cast<char>(std::tolower(c));
-            });
-            return ext == ".gguf" || ext == ".pt" || ext == ".pth" || ext == ".safetensors";
-        };
-
         for (auto& entry : fs::recursive_directory_iterator(lora_dir)) {
             if (!entry.is_regular_file()) {
                 continue;
             }
             const fs::path& p = entry.path();
-            if (!is_lora_ext(p)) {
+            if (!is_supported_model_ext(p)) {
                 continue;
             }
 
@@ -286,6 +291,40 @@ std::string get_lora_full_path(ServerRuntime& rt, const std::string& path) {
     return it != rt.lora_cache->end() ? it->fullpath : "";
 }
 
+void refresh_upscaler_cache(ServerRuntime& rt) {
+    std::vector<UpscalerEntry> new_cache;
+
+    fs::path upscaler_dir = rt.ctx_params->hires_upscalers_dir;
+    if (fs::exists(upscaler_dir) && fs::is_directory(upscaler_dir)) {
+        for (auto& entry : fs::directory_iterator(upscaler_dir)) {
+            if (!entry.is_regular_file()) {
+                continue;
+            }
+            const fs::path& p = entry.path();
+            if (!is_supported_model_ext(p)) {
+                continue;
+            }
+
+            UpscalerEntry upscaler_entry;
+            upscaler_entry.name       = p.stem().u8string();
+            upscaler_entry.fullpath   = fs::absolute(p).lexically_normal().u8string();
+            upscaler_entry.model_name = "ESRGAN_4x";
+            upscaler_entry.path       = p.filename().u8string();
+
+            new_cache.push_back(std::move(upscaler_entry));
+        }
+    }
+
+    std::sort(new_cache.begin(), new_cache.end(), [](const UpscalerEntry& a, const UpscalerEntry& b) {
+        return a.name < b.name;
+    });
+
+    {
+        std::lock_guard<std::mutex> lock(*rt.upscaler_mutex);
+        *rt.upscaler_cache = std::move(new_cache);
+    }
+}
+
 int64_t unix_timestamp_now() {
     return std::chrono::duration_cast<std::chrono::seconds>(
                std::chrono::system_clock::now().time_since_epoch())

examples/server/runtime.h

@@ -37,6 +37,14 @@ struct LoraEntry {
     std::string fullpath;
 };
 
+struct UpscalerEntry {
+    std::string name;
+    std::string path;
+    std::string fullpath;
+    std::string model_name;
+    int scale = 4;
+};
+
 struct ServerRuntime {
     sd_ctx_t* sd_ctx;
     std::mutex* sd_ctx_mutex;
@@ -45,6 +53,8 @@ struct ServerRuntime {
     const SDGenerationParams* default_gen_params;
     std::vector<LoraEntry>* lora_cache;
     std::mutex* lora_mutex;
+    std::vector<UpscalerEntry>* upscaler_cache;
+    std::mutex* upscaler_mutex;
     AsyncJobManager* async_job_manager;
 };
 
@@ -86,4 +96,5 @@ bool runtime_supports_generation_mode(const ServerRuntime& runtime, SDMode mode)
 std::string unsupported_generation_mode_error(SDMode mode);
 void refresh_lora_cache(ServerRuntime& rt);
 std::string get_lora_full_path(ServerRuntime& rt, const std::string& path);
+void refresh_upscaler_cache(ServerRuntime& rt);
 int64_t unix_timestamp_now();

format-code.sh

@@ -1,5 +1,5 @@
 for f in src/*.cpp src/*.h src/*.hpp src/tokenizers/*.h src/tokenizers/*.cpp src/tokenizers/vocab/*.h src/tokenizers/vocab/*.cpp \
-         examples/cli/*.cpp examples/cli/*.h examples/server/*.cpp \
+         src/model_io/*.h src/model_io/*.cpp examples/cli/*.cpp examples/cli/*.h examples/server/*.cpp \
          examples/common/*.hpp examples/common/*.h examples/common/*.cpp; do
   [[ "$f" == vocab* ]] && continue
   echo "formatting '$f'"

include/stable-diffusion.h

@@ -289,6 +289,25 @@ typedef struct {
     const char* path;
 } sd_lora_t;
 
+enum sd_hires_upscaler_t {
+    SD_HIRES_UPSCALER_NONE,
+    SD_HIRES_UPSCALER_LATENT_NEAREST,
+    SD_HIRES_UPSCALER_MODEL,
+    SD_HIRES_UPSCALER_COUNT,
+};
+
+typedef struct {
+    bool enabled;
+    enum sd_hires_upscaler_t upscaler;
+    const char* model_path;
+    float scale;
+    int target_width;
+    int target_height;
+    int steps;
+    float denoising_strength;
+    int upscale_tile_size;
+} sd_hires_params_t;
+
 typedef struct {
     const sd_lora_t* loras;
     uint32_t lora_count;
@@ -312,6 +331,7 @@ typedef struct {
     sd_pm_params_t pm_params;
     sd_tiling_params_t vae_tiling_params;
     sd_cache_params_t cache;
+    sd_hires_params_t hires;
 } sd_img_gen_params_t;
 
 typedef struct {
@@ -365,8 +385,11 @@ SD_API const char* sd_preview_name(enum preview_t preview);
 SD_API enum preview_t str_to_preview(const char* str);
 SD_API const char* sd_lora_apply_mode_name(enum lora_apply_mode_t mode);
 SD_API enum lora_apply_mode_t str_to_lora_apply_mode(const char* str);
+SD_API const char* sd_hires_upscaler_name(enum sd_hires_upscaler_t upscaler);
+SD_API enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str);
 
 SD_API void sd_cache_params_init(sd_cache_params_t* cache_params);
+SD_API void sd_hires_params_init(sd_hires_params_t* hires_params);
 
 SD_API void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params);
 SD_API char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params);

src/convert.cpp

@@ -0,0 +1,138 @@
+#include <cstring>
+#include <mutex>
+#include <regex>
+#include <vector>
+
+#include "model.h"
+#include "model_io/gguf_io.h"
+#include "model_io/safetensors_io.h"
+#include "util.h"
+
+#include "ggml-cpu.h"
+
+static ggml_type get_export_tensor_type(ModelLoader& model_loader,
+                                        const TensorStorage& tensor_storage,
+                                        ggml_type type,
+                                        const TensorTypeRules& tensor_type_rules) {
+    const std::string& name = tensor_storage.name;
+    ggml_type tensor_type   = tensor_storage.type;
+    ggml_type dst_type      = type;
+
+    for (const auto& tensor_type_rule : tensor_type_rules) {
+        std::regex pattern(tensor_type_rule.first);
+        if (std::regex_search(name, pattern)) {
+            dst_type = tensor_type_rule.second;
+            break;
+        }
+    }
+
+    if (model_loader.tensor_should_be_converted(tensor_storage, dst_type)) {
+        tensor_type = dst_type;
+    }
+
+    return tensor_type;
+}
+
+static bool load_tensors_for_export(ModelLoader& model_loader,
+                                    ggml_context* ggml_ctx,
+                                    ggml_type type,
+                                    const TensorTypeRules& tensor_type_rules,
+                                    std::vector<TensorWriteInfo>& tensors) {
+    std::mutex tensor_mutex;
+    auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
+        const std::string& name = tensor_storage.name;
+        ggml_type tensor_type   = get_export_tensor_type(model_loader, tensor_storage, type, tensor_type_rules);
+
+        std::lock_guard<std::mutex> lock(tensor_mutex);
+        ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, tensor_type, tensor_storage.n_dims, tensor_storage.ne);
+        if (tensor == nullptr) {
+            LOG_ERROR("ggml_new_tensor failed");
+            return false;
+        }
+        ggml_set_name(tensor, name.c_str());
+
+        if (!tensor->data) {
+            GGML_ASSERT(ggml_nelements(tensor) == 0);
+            // Avoid crashing writers by setting a dummy pointer for zero-sized tensors.
+            LOG_DEBUG("setting dummy pointer for zero-sized tensor %s", name.c_str());
+            tensor->data = ggml_get_mem_buffer(ggml_ctx);
+        }
+
+        TensorWriteInfo write_info;
+        write_info.tensor = tensor;
+        write_info.n_dims = tensor_storage.n_dims;
+        for (int i = 0; i < tensor_storage.n_dims; ++i) {
+            write_info.ne[i] = tensor_storage.ne[i];
+        }
+
+        *dst_tensor = tensor;
+        tensors.push_back(std::move(write_info));
+
+        return true;
+    };
+
+    bool success = model_loader.load_tensors(on_new_tensor_cb);
+    LOG_INFO("load tensors done");
+    return success;
+}
+
+bool convert(const char* input_path,
+             const char* vae_path,
+             const char* output_path,
+             sd_type_t output_type,
+             const char* tensor_type_rules,
+             bool convert_name) {
+    ModelLoader model_loader;
+
+    if (!model_loader.init_from_file(input_path)) {
+        LOG_ERROR("init model loader from file failed: '%s'", input_path);
+        return false;
+    }
+
+    if (vae_path != nullptr && strlen(vae_path) > 0) {
+        if (!model_loader.init_from_file(vae_path, "vae.")) {
+            LOG_ERROR("init model loader from file failed: '%s'", vae_path);
+            return false;
+        }
+    }
+    if (convert_name) {
+        model_loader.convert_tensors_name();
+    }
+
+    ggml_type type             = (ggml_type)output_type;
+    bool output_is_safetensors = ends_with(output_path, ".safetensors");
+    TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
+
+    auto backend    = ggml_backend_cpu_init();
+    size_t mem_size = 1 * 1024 * 1024;  // for padding
+    mem_size += model_loader.get_tensor_storage_map().size() * ggml_tensor_overhead();
+    mem_size += model_loader.get_params_mem_size(backend, type);
+    LOG_INFO("model tensors mem size: %.2fMB", mem_size / 1024.f / 1024.f);
+    ggml_context* ggml_ctx = ggml_init({mem_size, nullptr, false});
+
+    if (ggml_ctx == nullptr) {
+        LOG_ERROR("ggml_init failed for converter");
+        ggml_backend_free(backend);
+        return false;
+    }
+
+    std::vector<TensorWriteInfo> tensors;
+    bool success = load_tensors_for_export(model_loader, ggml_ctx, type, type_rules, tensors);
+    ggml_backend_free(backend);
+
+    std::string error;
+    if (success) {
+        if (output_is_safetensors) {
+            success = write_safetensors_file(output_path, tensors, &error);
+        } else {
+            success = write_gguf_file(output_path, tensors, &error);
+        }
+    }
+
+    if (!success && !error.empty()) {
+        LOG_ERROR("%s", error.c_str());
+    }
+
+    ggml_free(ggml_ctx);
+    return success;
+}

src/denoiser.hpp

@@ -977,7 +977,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
                                                         float eta = 1.0f) {
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        float sigma = sigmas[i];
+        float sigma    = sigmas[i];
         float sigma_to = sigmas[i + 1];
 
         bool opt_first_step = (1.0 - sigma < 1e-6);
@@ -1040,10 +1040,10 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
                 // and sigma_s = sigma_fn(s) = 1.0f / (exp(s) + 1.0f)
 
                 float exp_s = std::sqrt(((1 - sigma) / sigma) * ((1 - sigma_down) / sigma_down));
-                sigma_s = 1.0f / (exp_s + 1.0f);
+                sigma_s     = 1.0f / (exp_s + 1.0f);
 
                 float sigma_s_i_ratio = sigma_s / sigma;
-                sd::Tensor<float> u = (x * sigma_s_i_ratio) + (denoised * (1.0f - sigma_s_i_ratio));
+                sd::Tensor<float> u   = (x * sigma_s_i_ratio) + (denoised * (1.0f - sigma_s_i_ratio));
 
                 auto denoised2_opt = model(u, sigma_s, i + 1);
                 if (denoised2_opt.empty()) {
@@ -1053,7 +1053,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
             }
 
             float sigma_down_i_ratio = sigma_down / sigma;
-            x = (x * sigma_down_i_ratio) + (D_i * (1.0f - sigma_down_i_ratio));
+            x                        = (x * sigma_down_i_ratio) + (D_i * (1.0f - sigma_down_i_ratio));
 
             if (sigma_to > 0.0f && eta > 0.0f) {
                 x = alpha_scale * x + sd::Tensor<float>::randn_like(x, rng) * sigma_up;
@@ -1064,8 +1064,6 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
     return x;
 }
 
-
-
 static sd::Tensor<float> sample_dpmpp_2m(denoise_cb_t model,
                                          sd::Tensor<float> x,
                                          const std::vector<float>& sigmas) {
@@ -1525,32 +1523,10 @@ static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
                                               const std::vector<float>& sigmas,
                                               std::shared_ptr<RNG> rng,
                                               float eta) {
-    float beta_start = 0.00085f;
-    float beta_end   = 0.0120f;
-    std::vector<double> alphas_cumprod(TIMESTEPS);
-    std::vector<double> compvis_sigmas(TIMESTEPS);
-    for (int i = 0; i < TIMESTEPS; i++) {
-        alphas_cumprod[i] =
-            (i == 0 ? 1.0f : alphas_cumprod[i - 1]) *
-            (1.0f -
-             std::pow(sqrtf(beta_start) +
-                          (sqrtf(beta_end) - sqrtf(beta_start)) *
-                              ((float)i / (TIMESTEPS - 1)),
-                      2));
-        compvis_sigmas[i] =
-            std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]);
-    }
-
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        int timestep      = static_cast<int>(roundf(TIMESTEPS - i * ((float)TIMESTEPS / steps))) - 1;
-        int prev_timestep = timestep - TIMESTEPS / steps;
-        float sigma       = static_cast<float>(compvis_sigmas[timestep]);
-        if (i == 0) {
-            x *= std::sqrt(sigma * sigma + 1) / sigma;
-        } else {
-            x *= std::sqrt(sigma * sigma + 1);
-        }
+        float sigma    = sigmas[i];
+        float sigma_to = sigmas[i + 1];
 
         auto model_output_opt = model(x, sigma, i + 1);
         if (model_output_opt.empty()) {
@@ -1559,8 +1535,8 @@ static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
         sd::Tensor<float> model_output = std::move(model_output_opt);
         model_output                   = (x - model_output) * (1.0f / sigma);
 
-        float alpha_prod_t      = static_cast<float>(alphas_cumprod[timestep]);
-        float alpha_prod_t_prev = static_cast<float>(prev_timestep >= 0 ? alphas_cumprod[prev_timestep] : alphas_cumprod[0]);
+        float alpha_prod_t      = 1.0f / (sigma * sigma + 1.0f);
+        float alpha_prod_t_prev = 1.0f / (sigma_to * sigma_to + 1.0f);
         float beta_prod_t       = 1.0f - alpha_prod_t;
 
         sd::Tensor<float> pred_original_sample = ((x / std::sqrt(sigma * sigma + 1)) -
@@ -1572,11 +1548,11 @@ static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
                          (1.0f - alpha_prod_t / alpha_prod_t_prev);
         float std_dev_t = eta * std::sqrt(variance);
 
-        x = std::sqrt(alpha_prod_t_prev) * pred_original_sample +
-            std::sqrt(1.0f - alpha_prod_t_prev - std::pow(std_dev_t, 2)) * model_output;
+        x = pred_original_sample +
+            std::sqrt((1.0f - alpha_prod_t_prev - std::pow(std_dev_t, 2)) / alpha_prod_t_prev) * model_output;
 
         if (eta > 0) {
-            x += std_dev_t * sd::Tensor<float>::randn_like(x, rng);
+            x += std_dev_t / std::sqrt(alpha_prod_t_prev) * sd::Tensor<float>::randn_like(x, rng);
         }
     }
     return x;
@@ -1603,19 +1579,26 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
             std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]);
     }
 
-    int original_steps = 50;
-    int steps          = static_cast<int>(sigmas.size()) - 1;
-    for (int i = 0; i < steps; i++) {
-        int timestep      = TIMESTEPS - 1 - (TIMESTEPS / original_steps) * (int)floor(i * ((float)original_steps / steps));
-        int prev_timestep = i >= steps - 1 ? 0 : TIMESTEPS - 1 - (TIMESTEPS / original_steps) * (int)floor((i + 1) * ((float)original_steps / steps));
-        int timestep_s    = (int)floor((1 - eta) * prev_timestep);
-        float sigma       = static_cast<float>(compvis_sigmas[timestep]);
-
-        if (i == 0) {
-            x *= std::sqrt(sigma * sigma + 1) / sigma;
-        } else {
-            x *= std::sqrt(sigma * sigma + 1);
+    auto get_timestep_from_sigma = [&](float s) -> int {
+        auto it = std::lower_bound(compvis_sigmas.begin(), compvis_sigmas.end(), s);
+        if (it == compvis_sigmas.begin())
+            return 0;
+        if (it == compvis_sigmas.end())
+            return TIMESTEPS - 1;
+        int idx_high = static_cast<int>(std::distance(compvis_sigmas.begin(), it));
+        int idx_low  = idx_high - 1;
+        if (std::abs(compvis_sigmas[idx_high] - s) < std::abs(compvis_sigmas[idx_low] - s)) {
+            return idx_high;
         }
+        return idx_low;
+    };
+
+    int steps = static_cast<int>(sigmas.size()) - 1;
+    for (int i = 0; i < steps; i++) {
+        float sigma_to    = sigmas[i + 1];
+        int prev_timestep = get_timestep_from_sigma(sigma_to);
+        int timestep_s    = (int)floor((1 - eta) * prev_timestep);
+        float sigma       = sigmas[i];
 
         auto model_output_opt = model(x, sigma, i + 1);
         if (model_output_opt.empty()) {
@@ -1624,9 +1607,9 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
         sd::Tensor<float> model_output = std::move(model_output_opt);
         model_output                   = (x - model_output) * (1.0f / sigma);
 
-        float alpha_prod_t      = static_cast<float>(alphas_cumprod[timestep]);
+        float alpha_prod_t      = 1.0f / (sigma * sigma + 1.0f);
         float beta_prod_t       = 1.0f - alpha_prod_t;
-        float alpha_prod_t_prev = static_cast<float>(prev_timestep >= 0 ? alphas_cumprod[prev_timestep] : alphas_cumprod[0]);
+        float alpha_prod_t_prev = 1.0f / (sigma_to * sigma_to + 1.0f);
         float alpha_prod_s      = static_cast<float>(alphas_cumprod[timestep_s]);
         float beta_prod_s       = 1.0f - alpha_prod_s;
 
@@ -1634,12 +1617,12 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
                                                   std::sqrt(beta_prod_t) * model_output) *
                                                  (1.0f / std::sqrt(alpha_prod_t));
 
-        x = std::sqrt(alpha_prod_s) * pred_original_sample +
-            std::sqrt(beta_prod_s) * model_output;
+        x = std::sqrt(alpha_prod_s / alpha_prod_t_prev) * pred_original_sample +
+            std::sqrt(beta_prod_s / alpha_prod_t_prev) * model_output;
 
-        if (eta > 0 && i != steps - 1) {
+        if (eta > 0 && sigma_to > 0.0f) {
             x = std::sqrt(alpha_prod_t_prev / alpha_prod_s) * x +
-                std::sqrt(1.0f - alpha_prod_t_prev / alpha_prod_s) * sd::Tensor<float>::randn_like(x, rng);
+                std::sqrt(1.0f / alpha_prod_t_prev - 1.0f / alpha_prod_s) * sd::Tensor<float>::randn_like(x, rng);
         }
     }
     return x;

src/model.cpp

@@ -2,6 +2,7 @@
 #include <atomic>
 #include <chrono>
 #include <cstdarg>
+#include <cstdlib>
 #include <fstream>
 #include <functional>
 #include <mutex>
@@ -12,8 +13,11 @@
 #include <unordered_map>
 #include <vector>
 
-#include "gguf_reader.hpp"
 #include "model.h"
+#include "model_io/gguf_io.h"
+#include "model_io/safetensors_io.h"
+#include "model_io/torch_legacy_io.h"
+#include "model_io/torch_zip_io.h"
 #include "stable-diffusion.h"
 #include "util.h"
 
@@ -21,6 +25,7 @@
 #include "ggml-backend.h"
 #include "ggml-cpu.h"
 #include "ggml.h"
+#include "zip.h"
 
 #include "name_conversion.h"
 #include "stable-diffusion.h"
@@ -37,40 +42,6 @@
 #include "ggml-opencl.h"
 #endif
 
-#define ST_HEADER_SIZE_LEN 8
-
-uint64_t read_u64(uint8_t* buffer) {
-    // little endian
-    uint64_t value = 0;
-    value |= static_cast<int64_t>(buffer[7]) << 56;
-    value |= static_cast<int64_t>(buffer[6]) << 48;
-    value |= static_cast<int64_t>(buffer[5]) << 40;
-    value |= static_cast<int64_t>(buffer[4]) << 32;
-    value |= static_cast<int64_t>(buffer[3]) << 24;
-    value |= static_cast<int64_t>(buffer[2]) << 16;
-    value |= static_cast<int64_t>(buffer[1]) << 8;
-    value |= static_cast<int64_t>(buffer[0]);
-    return value;
-}
-
-int32_t read_int(uint8_t* buffer) {
-    // little endian
-    int value = 0;
-    value |= buffer[3] << 24;
-    value |= buffer[2] << 16;
-    value |= buffer[1] << 8;
-    value |= buffer[0];
-    return value;
-}
-
-uint16_t read_short(uint8_t* buffer) {
-    // little endian
-    uint16_t value = 0;
-    value |= buffer[1] << 8;
-    value |= buffer[0];
-    return value;
-}
-
 /*================================================= Preprocess ==================================================*/
 
 const char* unused_tensors[] = {
@@ -110,7 +81,7 @@ const char* unused_tensors[] = {
     "first_stage_model.bn.",
 };
 
-bool is_unused_tensor(std::string name) {
+bool is_unused_tensor(const std::string& name) {
     for (size_t i = 0; i < sizeof(unused_tensors) / sizeof(const char*); i++) {
         if (starts_with(name, unused_tensors[i])) {
             return true;
@@ -250,79 +221,6 @@ void ModelLoader::add_tensor_storage(const TensorStorage& tensor_storage) {
     tensor_storage_map[tensor_storage.name] = tensor_storage;
 }
 
-bool is_zip_file(const std::string& file_path) {
-    zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
-    if (zip == nullptr) {
-        return false;
-    }
-    zip_close(zip);
-    return true;
-}
-
-bool is_gguf_file(const std::string& file_path) {
-    std::ifstream file(file_path, std::ios::binary);
-    if (!file.is_open()) {
-        return false;
-    }
-
-    char magic[4];
-
-    file.read(magic, sizeof(magic));
-    if (!file) {
-        return false;
-    }
-    for (uint32_t i = 0; i < sizeof(magic); i++) {
-        if (magic[i] != GGUF_MAGIC[i]) {
-            return false;
-        }
-    }
-
-    return true;
-}
-
-bool is_safetensors_file(const std::string& file_path) {
-    std::ifstream file(file_path, std::ios::binary);
-    if (!file.is_open()) {
-        return false;
-    }
-
-    // get file size
-    file.seekg(0, file.end);
-    size_t file_size_ = file.tellg();
-    file.seekg(0, file.beg);
-
-    // read header size
-    if (file_size_ <= ST_HEADER_SIZE_LEN) {
-        return false;
-    }
-
-    uint8_t header_size_buf[ST_HEADER_SIZE_LEN];
-    file.read((char*)header_size_buf, ST_HEADER_SIZE_LEN);
-    if (!file) {
-        return false;
-    }
-
-    size_t header_size_ = read_u64(header_size_buf);
-    if (header_size_ >= file_size_ || header_size_ <= 2) {
-        return false;
-    }
-
-    // read header
-    std::vector<char> header_buf;
-    header_buf.resize(header_size_ + 1);
-    header_buf[header_size_] = '\0';
-    file.read(header_buf.data(), header_size_);
-    if (!file) {
-        return false;
-    }
-    try {
-        nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
-    } catch (const std::exception&) {
-        return false;
-    }
-    return true;
-}
-
 bool ModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
     if (is_directory(file_path)) {
         LOG_INFO("load %s using diffusers format", file_path.c_str());
@@ -333,9 +231,12 @@ bool ModelLoader::init_from_file(const std::string& file_path, const std::string
     } else if (is_safetensors_file(file_path)) {
         LOG_INFO("load %s using safetensors format", file_path.c_str());
         return init_from_safetensors_file(file_path, prefix);
-    } else if (is_zip_file(file_path)) {
-        LOG_INFO("load %s using checkpoint format", file_path.c_str());
-        return init_from_ckpt_file(file_path, prefix);
+    } else if (is_torch_zip_file(file_path)) {
+        LOG_INFO("load %s using torch zip format", file_path.c_str());
+        return init_from_torch_zip_file(file_path, prefix);
+    } else if (init_from_torch_legacy_file(file_path, prefix)) {
+        LOG_INFO("load %s using torch legacy format", file_path.c_str());
+        return true;
     } else {
         if (file_exists(file_path)) {
             LOG_WARN("unknown format %s", file_path.c_str());
@@ -375,242 +276,121 @@ bool ModelLoader::init_from_file_and_convert_name(const std::string& file_path,
 
 bool ModelLoader::init_from_gguf_file(const std::string& file_path, const std::string& prefix) {
     LOG_DEBUG("init from '%s'", file_path.c_str());
+
+    std::vector<TensorStorage> tensor_storages;
+    std::string error;
+    if (!read_gguf_file(file_path, tensor_storages, &error)) {
+        LOG_ERROR("%s", error.c_str());
+        return false;
+    }
+
     file_paths_.push_back(file_path);
     size_t file_index = file_paths_.size() - 1;
 
-    gguf_context* ctx_gguf_ = nullptr;
-    ggml_context* ctx_meta_ = nullptr;
+    for (auto& tensor_storage : tensor_storages) {
+        // LOG_DEBUG("%s", tensor_storage.name.c_str());
 
-    ctx_gguf_ = gguf_init_from_file(file_path.c_str(), {true, &ctx_meta_});
-    if (!ctx_gguf_) {
-        LOG_ERROR("failed to open '%s' with gguf_init_from_file. Try to open it with GGUFReader.", file_path.c_str());
-        GGUFReader gguf_reader;
-        if (!gguf_reader.load(file_path)) {
-            LOG_ERROR("failed to open '%s' with GGUFReader.", file_path.c_str());
-            return false;
+        if (!starts_with(tensor_storage.name, prefix)) {
+            tensor_storage.name = prefix + tensor_storage.name;
         }
-
-        size_t data_offset = gguf_reader.data_offset();
-        for (const auto& gguf_tensor_info : gguf_reader.tensors()) {
-            std::string name = gguf_tensor_info.name;
-            if (!starts_with(name, prefix)) {
-                name = prefix + name;
-            }
-
-            TensorStorage tensor_storage(
-                name,
-                gguf_tensor_info.type,
-                gguf_tensor_info.shape.data(),
-                static_cast<int>(gguf_tensor_info.shape.size()),
-                file_index,
-                data_offset + gguf_tensor_info.offset);
-
-            // LOG_DEBUG("%s %s", name.c_str(), tensor_storage.to_string().c_str());
-
-            add_tensor_storage(tensor_storage);
-        }
-
-        return true;
-    }
-
-    int n_tensors = static_cast<int>(gguf_get_n_tensors(ctx_gguf_));
-
-    size_t total_size  = 0;
-    size_t data_offset = gguf_get_data_offset(ctx_gguf_);
-    for (int i = 0; i < n_tensors; i++) {
-        std::string name   = gguf_get_tensor_name(ctx_gguf_, i);
-        ggml_tensor* dummy = ggml_get_tensor(ctx_meta_, name.c_str());
-        size_t offset      = data_offset + gguf_get_tensor_offset(ctx_gguf_, i);
-
-        // LOG_DEBUG("%s", name.c_str());
-
-        if (!starts_with(name, prefix)) {
-            name = prefix + name;
-        }
-
-        TensorStorage tensor_storage(name, dummy->type, dummy->ne, ggml_n_dims(dummy), file_index, offset);
-
-        GGML_ASSERT(ggml_nbytes(dummy) == tensor_storage.nbytes());
+        tensor_storage.file_index = file_index;
 
         add_tensor_storage(tensor_storage);
     }
 
-    gguf_free(ctx_gguf_);
-    ggml_free(ctx_meta_);
-
     return true;
 }
 
 /*================================================= SafeTensorsModelLoader ==================================================*/
 
-ggml_type str_to_ggml_type(const std::string& dtype) {
-    ggml_type ttype = GGML_TYPE_COUNT;
-    if (dtype == "F16") {
-        ttype = GGML_TYPE_F16;
-    } else if (dtype == "BF16") {
-        ttype = GGML_TYPE_BF16;
-    } else if (dtype == "F32") {
-        ttype = GGML_TYPE_F32;
-    } else if (dtype == "F64") {
-        ttype = GGML_TYPE_F32;
-    } else if (dtype == "F8_E4M3") {
-        ttype = GGML_TYPE_F16;
-    } else if (dtype == "F8_E5M2") {
-        ttype = GGML_TYPE_F16;
-    } else if (dtype == "I64") {
-        ttype = GGML_TYPE_I32;
-    }
-    return ttype;
-}
-
-// https://huggingface.co/docs/safetensors/index
 bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const std::string& prefix) {
     LOG_DEBUG("init from '%s', prefix = '%s'", file_path.c_str(), prefix.c_str());
+
+    std::vector<TensorStorage> tensor_storages;
+    std::string error;
+    if (!read_safetensors_file(file_path, tensor_storages, &error)) {
+        LOG_ERROR("%s", error.c_str());
+        return false;
+    }
+
     file_paths_.push_back(file_path);
     size_t file_index = file_paths_.size() - 1;
-    std::ifstream file(file_path, std::ios::binary);
-    if (!file.is_open()) {
-        LOG_ERROR("failed to open '%s'", file_path.c_str());
-        file_paths_.pop_back();
-        return false;
-    }
 
-    // get file size
-    file.seekg(0, file.end);
-    size_t file_size_ = file.tellg();
-    file.seekg(0, file.beg);
-
-    // read header size
-    if (file_size_ <= ST_HEADER_SIZE_LEN) {
-        LOG_ERROR("invalid safetensor file '%s'", file_path.c_str());
-        file_paths_.pop_back();
-        return false;
-    }
-
-    uint8_t header_size_buf[ST_HEADER_SIZE_LEN];
-    file.read((char*)header_size_buf, ST_HEADER_SIZE_LEN);
-    if (!file) {
-        LOG_ERROR("read safetensors header size failed: '%s'", file_path.c_str());
-        return false;
-    }
-
-    size_t header_size_ = read_u64(header_size_buf);
-    if (header_size_ >= file_size_) {
-        LOG_ERROR("invalid safetensor file '%s'", file_path.c_str());
-        file_paths_.pop_back();
-        return false;
-    }
-
-    // read header
-    std::vector<char> header_buf;
-    header_buf.resize(header_size_ + 1);
-    header_buf[header_size_] = '\0';
-    file.read(header_buf.data(), header_size_);
-    if (!file) {
-        LOG_ERROR("read safetensors header failed: '%s'", file_path.c_str());
-        file_paths_.pop_back();
-        return false;
-    }
-
-    nlohmann::json header_;
-    try {
-        header_ = nlohmann::json::parse(header_buf.data());
-    } catch (const std::exception&) {
-        LOG_ERROR("parsing safetensors header failed", file_path.c_str());
-        file_paths_.pop_back();
-        return false;
-    }
-
-    for (auto& item : header_.items()) {
-        std::string name           = item.key();
-        nlohmann::json tensor_info = item.value();
-        // LOG_DEBUG("%s %s\n", name.c_str(), tensor_info.dump().c_str());
-
-        if (name == "__metadata__") {
+    for (auto& tensor_storage : tensor_storages) {
+        if (is_unused_tensor(tensor_storage.name)) {
             continue;
         }
 
-        if (is_unused_tensor(name)) {
-            continue;
-        }
-
-        std::string dtype    = tensor_info["dtype"];
-        nlohmann::json shape = tensor_info["shape"];
-
-        if (dtype == "U8") {
-            continue;
-        }
-
-        size_t begin = tensor_info["data_offsets"][0].get<size_t>();
-        size_t end   = tensor_info["data_offsets"][1].get<size_t>();
-
-        ggml_type type = str_to_ggml_type(dtype);
-        if (type == GGML_TYPE_COUNT) {
-            LOG_ERROR("unsupported dtype '%s' (tensor '%s')", dtype.c_str(), name.c_str());
-            return false;
-        }
-
-        if (shape.size() > SD_MAX_DIMS) {
-            LOG_ERROR("invalid tensor '%s'", name.c_str());
-            return false;
-        }
-
-        int n_dims              = (int)shape.size();
-        int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
-        for (int i = 0; i < n_dims; i++) {
-            ne[i] = shape[i].get<int64_t>();
-        }
-
-        if (n_dims == 5) {
-            n_dims = 4;
-            ne[0]  = ne[0] * ne[1];
-            ne[1]  = ne[2];
-            ne[2]  = ne[3];
-            ne[3]  = ne[4];
-        }
-
-        // ggml_n_dims returns 1 for scalars
-        if (n_dims == 0) {
-            n_dims = 1;
-        }
-
-        if (!starts_with(name, prefix)) {
-            name = prefix + name;
-        }
-
-        TensorStorage tensor_storage(name, type, ne, n_dims, file_index, ST_HEADER_SIZE_LEN + header_size_ + begin);
-        tensor_storage.reverse_ne();
-
-        size_t tensor_data_size = end - begin;
-
-        bool tensor_size_ok;
-        if (dtype == "F8_E4M3") {
-            tensor_storage.is_f8_e4m3 = true;
-            // f8 -> f16
-            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
-        } else if (dtype == "F8_E5M2") {
-            tensor_storage.is_f8_e5m2 = true;
-            // f8 -> f16
-            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
-        } else if (dtype == "F64") {
-            tensor_storage.is_f64 = true;
-            // f64 -> f32
-            tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
-        } else if (dtype == "I64") {
-            tensor_storage.is_i64 = true;
-            // i64 -> i32
-            tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
-        } else {
-            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size);
-        }
-        if (!tensor_size_ok) {
-            LOG_ERROR("size mismatch for tensor '%s' (%s)\n", name.c_str(), dtype.c_str());
-            return false;
+        if (!starts_with(tensor_storage.name, prefix)) {
+            tensor_storage.name = prefix + tensor_storage.name;
         }
+        tensor_storage.file_index = file_index;
 
         add_tensor_storage(tensor_storage);
 
-        // LOG_DEBUG("%s %s", tensor_storage.to_string().c_str(), dtype.c_str());
+        // LOG_DEBUG("%s", tensor_storage.to_string().c_str());
+    }
+
+    return true;
+}
+
+/*================================================= TorchLegacyModelLoader ==================================================*/
+
+bool ModelLoader::init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix) {
+    LOG_DEBUG("init from torch legacy '%s'", file_path.c_str());
+
+    std::vector<TensorStorage> tensor_storages;
+    std::string error;
+    if (!read_torch_legacy_file(file_path, tensor_storages, &error)) {
+        if ((!error.empty()) && (ends_with(file_path, ".pt") || ends_with(file_path, ".pth"))) {
+            LOG_WARN("%s", error.c_str());
+        }
+        return false;
+    }
+
+    file_paths_.push_back(file_path);
+    size_t file_index = file_paths_.size() - 1;
+
+    for (auto& tensor_storage : tensor_storages) {
+        if (is_unused_tensor(tensor_storage.name)) {
+            continue;
+        }
+
+        if (!starts_with(tensor_storage.name, prefix)) {
+            tensor_storage.name = prefix + tensor_storage.name;
+        }
+        tensor_storage.file_index = file_index;
+
+        add_tensor_storage(tensor_storage);
+    }
+
+    return true;
+}
+
+/*================================================= TorchZipModelLoader ==================================================*/
+
+bool ModelLoader::init_from_torch_zip_file(const std::string& file_path, const std::string& prefix) {
+    LOG_DEBUG("init from '%s'", file_path.c_str());
+
+    std::vector<TensorStorage> tensor_storages;
+    std::string error;
+    if (!read_torch_zip_file(file_path, tensor_storages, &error)) {
+        LOG_ERROR("%s", error.c_str());
+        return false;
+    }
+
+    file_paths_.push_back(file_path);
+    size_t file_index = file_paths_.size() - 1;
+
+    for (auto& tensor_storage : tensor_storages) {
+        if (!starts_with(tensor_storage.name, prefix)) {
+            tensor_storage.name = prefix + tensor_storage.name;
+        }
+        tensor_storage.file_index = file_index;
+
+        add_tensor_storage(tensor_storage);
+
+        // LOG_DEBUG("%s", tensor_storage.to_string().c_str());
     }
 
     return true;
@@ -642,367 +422,6 @@ bool ModelLoader::init_from_diffusers_file(const std::string& file_path, const s
     return true;
 }
 
-/*================================================= CkptModelLoader ==================================================*/
-
-// $ python -m pickletools sd-v1-4/archive/data.pkl | head -n 100
-//     0: \x80 PROTO      2
-//     2: }    EMPTY_DICT
-//     3: q    BINPUT     0
-//     5: (    MARK
-//     6: X        BINUNICODE 'epoch'
-//    16: q        BINPUT     1
-//    18: K        BININT1    6
-//    20: X        BINUNICODE 'global_step'
-//    36: q        BINPUT     2
-//    38: J        BININT     470000
-//    43: X        BINUNICODE 'pytorch-lightning_version'
-//    73: q        BINPUT     3
-//    75: X        BINUNICODE '1.4.2'
-//    85: q        BINPUT     4
-//    87: X        BINUNICODE 'state_dict'
-//   102: q        BINPUT     5
-//   104: }        EMPTY_DICT
-//   105: q        BINPUT     6
-//   107: (        MARK
-//   108: X            BINUNICODE 'betas'
-//   118: q            BINPUT     7
-//   120: c            GLOBAL     'torch._utils _rebuild_tensor_v2'
-//   153: q            BINPUT     8
-//   155: (            MARK
-//   156: (                MARK
-//   157: X                    BINUNICODE 'storage'
-//   169: q                    BINPUT     9
-//   171: c                    GLOBAL     'torch FloatStorage'
-//   191: q                    BINPUT     10
-//   193: X                    BINUNICODE '0'
-//   199: q                    BINPUT     11
-//   201: X                    BINUNICODE 'cpu'
-//   209: q                    BINPUT     12
-//   211: M                    BININT2    1000
-//   214: t                    TUPLE      (MARK at 156)
-//   215: q                BINPUT     13
-//   217: Q                BINPERSID
-//   218: K                BININT1    0
-//   220: M                BININT2    1000
-//  ...............................
-//  3201: q            BINPUT     250
-//  3203: R            REDUCE
-//  3204: q            BINPUT     251
-//  3206: X            BINUNICODE 'model.diffusion_model.input_blocks.1.1.proj_in.weight'
-//  3264: q            BINPUT     252
-//  3266: h            BINGET     8
-//  3268: (            MARK
-//  3269: (                MARK
-//  3270: h                    BINGET     9
-//  3272: h                    BINGET     10
-//  3274: X                    BINUNICODE '30'
-//  3281: q                    BINPUT     253
-//  3283: h                    BINGET     12
-//  3285: J                    BININT     102400
-//  3290: t                    TUPLE      (MARK at 3269)
-//  3291: q                BINPUT     254
-//  3293: Q                BINPERSID
-//  3294: K                BININT1    0
-//  3296: (                MARK
-//  3297: M                    BININT2    320
-//  3300: M                    BININT2    320
-//  3303: K                    BININT1    1
-//  3305: K                    BININT1    1
-//  3307: t                    TUPLE      (MARK at 3296)
-//  3308: q                BINPUT     255
-//  3310: (                MARK
-//  3311: M                    BININT2    320
-//  3314: K                    BININT1    1
-//  3316: K                    BININT1    1
-//  3318: K                    BININT1    1
-//  3320: t                    TUPLE      (MARK at 3310)
-//  3321: r                LONG_BINPUT 256
-//  3326: \x89             NEWFALSE
-//  3327: h                BINGET     16
-//  3329: )                EMPTY_TUPLE
-//  3330: R                REDUCE
-//  3331: r                LONG_BINPUT 257
-//  3336: t                TUPLE      (MARK at 3268)
-//  3337: r            LONG_BINPUT 258
-//  3342: R            REDUCE
-//  3343: r            LONG_BINPUT 259
-//  3348: X            BINUNICODE 'model.diffusion_model.input_blocks.1.1.proj_in.bias'
-//  3404: r            LONG_BINPUT 260
-//  3409: h            BINGET     8
-//  3411: (            MARK
-//  3412: (                MARK
-//  3413: h                    BINGET     9
-//  3415: h                    BINGET     10
-//  3417: X                    BINUNICODE '31'
-
-struct PickleTensorReader {
-    enum ReadPhase {
-        READ_NAME,
-        READ_DATA,
-        CHECK_SIZE,
-        READ_DIMENS
-    };
-    ReadPhase phase   = READ_NAME;
-    size_t entry_size = 0;
-    int32_t nelements = 0;
-
-    TensorStorage tensor_storage;
-
-    static ggml_type global_type;  // all pickle_tensors data type
-    static bool read_global_type;
-
-    bool read_int_value(uint32_t value) {
-        if (phase == CHECK_SIZE) {
-            if (entry_size == value * ggml_type_size(tensor_storage.type)) {
-                nelements = value;
-                phase     = READ_DIMENS;
-                return true;
-            } else {
-                phase = READ_NAME;
-            }
-        } else if (phase == READ_DIMENS) {
-            if (tensor_storage.n_dims + 1 > SD_MAX_DIMS) {  // too many dimens
-                phase                 = READ_NAME;
-                tensor_storage.n_dims = 0;
-            }
-            if (nelements % value == 0) {
-                tensor_storage.ne[tensor_storage.n_dims] = value;
-                tensor_storage.n_dims++;
-            }
-        }
-        return false;
-    }
-
-    void read_global(const std::string& str) {
-        if (str == "FloatStorage") {
-            if (read_global_type) {
-                global_type      = GGML_TYPE_F32;
-                read_global_type = false;
-            }
-            tensor_storage.type = GGML_TYPE_F32;
-        } else if (str == "HalfStorage") {
-            if (read_global_type) {
-                global_type      = GGML_TYPE_F16;
-                read_global_type = false;
-            }
-            tensor_storage.type = GGML_TYPE_F16;
-        }
-    }
-
-    void read_string(const std::string& str, zip_t* zip, std::string dir) {
-        if (str == "storage") {
-            read_global_type = true;
-        } else if (str != "state_dict") {
-            if (phase == READ_DATA) {
-                std::string entry_name = dir + "data/" + std::string(str);
-
-                size_t i, n = zip_entries_total(zip);
-                for (i = 0; i < n; ++i) {
-                    zip_entry_openbyindex(zip, i);
-                    {
-                        std::string name = zip_entry_name(zip);
-                        if (name == entry_name) {
-                            tensor_storage.index_in_zip = (int)i;
-                            entry_size                  = zip_entry_size(zip);
-                            zip_entry_close(zip);
-                            break;
-                        }
-                    }
-                    zip_entry_close(zip);
-                }
-
-                phase = entry_size > 0 ? CHECK_SIZE : READ_NAME;
-            }
-            if (!read_global_type && phase == READ_NAME) {
-                tensor_storage.name = str;
-                phase               = READ_DATA;
-                tensor_storage.type = global_type;
-            }
-        }
-    }
-};
-
-ggml_type PickleTensorReader::global_type = GGML_TYPE_F32;  // all pickle_tensors data type
-bool PickleTensorReader::read_global_type = false;
-
-int find_char(uint8_t* buffer, int len, char c) {
-    for (int pos = 0; pos < len; pos++) {
-        if (buffer[pos] == c) {
-            return pos;
-        }
-    }
-    return -1;
-}
-
-#define MAX_STRING_BUFFER 512
-
-bool ModelLoader::parse_data_pkl(uint8_t* buffer,
-                                 size_t buffer_size,
-                                 zip_t* zip,
-                                 std::string dir,
-                                 size_t file_index,
-                                 const std::string prefix) {
-    uint8_t* buffer_end = buffer + buffer_size;
-    if (buffer[0] == 0x80) {  // proto
-        if (buffer[1] != 2) {
-            LOG_ERROR("Unsupported protocol\n");
-            return false;
-        }
-        buffer += 2;  // 0x80 and version
-        char string_buffer[MAX_STRING_BUFFER];
-        bool finish = false;
-        PickleTensorReader reader;
-        // read pickle binary file
-        while (!finish && buffer < buffer_end) {
-            uint8_t opcode = *buffer;
-            buffer++;
-            // https://github.com/python/cpython/blob/3.7/Lib/pickletools.py#L1048
-            // https://github.com/python/cpython/blob/main/Lib/pickle.py#L105
-            switch (opcode) {
-                case '}':  // EMPTY_DICT     = b'}'   # push empty dict
-                    break;
-                case ']':  // EMPTY_LIST     = b']'   # push empty list
-                    break;
-                // skip unused sections
-                case 'h':  // BINGET         = b'h'   #   "    "    "    "   "   "  ;   "    " 1-byte arg
-                case 'q':  // BINPUT         = b'q'   #   "     "    "   "   " ;   "    " 1-byte arg
-                case 'Q':  // BINPERSID      = b'Q'   #  "       "         "  ;  "  "   "     "  stack
-                    buffer++;
-                    break;
-                case 'r':  // LONG_BINPUT    = b'r'   #   "     "    "   "   " ;   "    " 4-byte arg
-                    buffer += 4;
-                    break;
-                case 0x95:  // FRAME            = b'\x95'  # indicate the beginning of a new frame
-                    buffer += 8;
-                    break;
-                case 0x94:  // MEMOIZE          = b'\x94'  # store top of the stack in memo
-                    break;
-                case '(':  // MARK           = b'('   # push special markobject on stack
-                    break;
-                case 'K':  // BININT1        = b'K'   # push 1-byte unsigned int
-                {
-                    uint8_t value = *buffer;
-                    if (reader.read_int_value(value)) {
-                        buffer++;
-                    }
-                    buffer++;
-                } break;
-                case 'M':  // BININT2        = b'M'   # push 2-byte unsigned int
-                {
-                    uint16_t value = read_short(buffer);
-                    if (reader.read_int_value(value)) {
-                        buffer++;
-                    }
-                    buffer += 2;
-                } break;
-                case 'J':  // BININT         = b'J'   # push four-byte signed int
-                {
-                    const int32_t value = read_int(buffer);
-                    if (reader.read_int_value(value)) {
-                        buffer++;  // skip tuple after read num_elements
-                    }
-                    buffer += 4;
-                } break;
-                case 'X':  // BINUNICODE     = b'X'   #   "     "       "  ; counted UTF-8 string argument
-                {
-                    const int32_t len = read_int(buffer);
-                    buffer += 4;
-                    memset(string_buffer, 0, MAX_STRING_BUFFER);
-                    if (len > MAX_STRING_BUFFER) {
-                        LOG_WARN("tensor name very large");
-                    }
-                    memcpy(string_buffer, buffer, len < MAX_STRING_BUFFER ? len : (MAX_STRING_BUFFER - 1));
-                    buffer += len;
-                    reader.read_string(string_buffer, zip, dir);
-                } break;
-                case 0x8C:  // SHORT_BINUNICODE = b'\x8c'  # push short string; UTF-8 length < 256 bytes
-                {
-                    const int8_t len = *buffer;
-                    buffer++;
-                    memset(string_buffer, 0, MAX_STRING_BUFFER);
-                    memcpy(string_buffer, buffer, len);
-                    buffer += len;
-                    // printf("String: '%s'\n", string_buffer);
-                } break;
-                case 'c':  // GLOBAL         = b'c'   # push self.find_class(modname, name); 2 string args
-                {
-                    int len = find_char(buffer, MAX_STRING_BUFFER, '\n');
-
-                    buffer += len + 1;
-                    len = find_char(buffer, MAX_STRING_BUFFER, '\n');
-
-                    memset(string_buffer, 0, MAX_STRING_BUFFER);
-                    memcpy(string_buffer, buffer, len);
-                    buffer += len + 1;
-                    reader.read_global(string_buffer);
-                } break;
-                case 0x86:  // TUPLE2         = b'\x86'  # build 2-tuple from two topmost stack items
-                case 0x85:  // TUPLE1         = b'\x85'  # build 1-tuple from stack top
-                case 't':   // TUPLE          = b't'   # build tuple from topmost stack items
-                    if (reader.phase == PickleTensorReader::READ_DIMENS) {
-                        reader.tensor_storage.reverse_ne();
-                        reader.tensor_storage.file_index = file_index;
-                        // if(strcmp(prefix.c_str(), "scarlett") == 0)
-                        // printf(" ZIP got tensor %s \n ", reader.tensor_storage.name.c_str());
-                        std::string name = reader.tensor_storage.name;
-                        if (!starts_with(name, prefix)) {
-                            name = prefix + name;
-                        }
-                        reader.tensor_storage.name = name;
-                        add_tensor_storage(reader.tensor_storage);
-
-                        // LOG_DEBUG("%s", reader.tensor_storage.name.c_str());
-                        // reset
-                        reader = PickleTensorReader();
-                    }
-                    break;
-                case '.':  // STOP           = b'.'   # every pickle ends with STOP
-                    finish = true;
-                    break;
-                default:
-                    break;
-            }
-        }
-    }
-    return true;
-}
-
-bool ModelLoader::init_from_ckpt_file(const std::string& file_path, const std::string& prefix) {
-    LOG_DEBUG("init from '%s'", file_path.c_str());
-    file_paths_.push_back(file_path);
-    size_t file_index = file_paths_.size() - 1;
-
-    zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
-    if (zip == nullptr) {
-        LOG_ERROR("failed to open '%s'", file_path.c_str());
-        return false;
-    }
-    int n = (int)zip_entries_total(zip);
-    for (int i = 0; i < n; ++i) {
-        zip_entry_openbyindex(zip, i);
-        {
-            std::string name = zip_entry_name(zip);
-            size_t pos       = name.find("data.pkl");
-            if (pos != std::string::npos) {
-                std::string dir = name.substr(0, pos);
-                printf("ZIP %d, name = %s, dir = %s \n", i, name.c_str(), dir.c_str());
-                void* pkl_data = nullptr;
-                size_t pkl_size;
-                zip_entry_read(zip, &pkl_data, &pkl_size);
-
-                // LOG_DEBUG("%lld", pkl_size);
-
-                parse_data_pkl((uint8_t*)pkl_data, pkl_size, zip, dir, file_index, prefix);
-
-                free(pkl_data);
-            }
-        }
-        zip_entry_close(zip);
-    }
-    zip_close(zip);
-    return true;
-}
-
 SDVersion ModelLoader::get_sd_version() {
     TensorStorage token_embedding_weight, input_block_weight;
 
@@ -1268,8 +687,8 @@ std::map<ggml_type, uint32_t> ModelLoader::get_vae_wtype_stat() {
     return wtype_stat;
 }
 
-static std::vector<std::pair<std::string, ggml_type>> parse_tensor_type_rules(const std::string& tensor_type_rules) {
-    std::vector<std::pair<std::string, ggml_type>> result;
+TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules) {
+    TensorTypeRules result;
     for (const auto& item : split_string(tensor_type_rules, ',')) {
         if (item.size() == 0)
             continue;
@@ -1702,76 +1121,6 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage
     return false;
 }
 
-bool ModelLoader::save_to_gguf_file(const std::string& file_path, ggml_type type, const std::string& tensor_type_rules_str) {
-    auto backend    = ggml_backend_cpu_init();
-    size_t mem_size = 1 * 1024 * 1024;  // for padding
-    mem_size += tensor_storage_map.size() * ggml_tensor_overhead();
-    mem_size += get_params_mem_size(backend, type);
-    LOG_INFO("model tensors mem size: %.2fMB", mem_size / 1024.f / 1024.f);
-    ggml_context* ggml_ctx = ggml_init({mem_size, nullptr, false});
-
-    gguf_context* gguf_ctx = gguf_init_empty();
-
-    auto tensor_type_rules = parse_tensor_type_rules(tensor_type_rules_str);
-
-    std::mutex tensor_mutex;
-    auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
-        const std::string& name = tensor_storage.name;
-        ggml_type tensor_type   = tensor_storage.type;
-        ggml_type dst_type      = type;
-
-        for (const auto& tensor_type_rule : tensor_type_rules) {
-            std::regex pattern(tensor_type_rule.first);
-            if (std::regex_search(name, pattern)) {
-                dst_type = tensor_type_rule.second;
-                break;
-            }
-        }
-
-        if (tensor_should_be_converted(tensor_storage, dst_type)) {
-            tensor_type = dst_type;
-        }
-
-        std::lock_guard<std::mutex> lock(tensor_mutex);
-        ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, tensor_type, tensor_storage.n_dims, tensor_storage.ne);
-        if (tensor == nullptr) {
-            LOG_ERROR("ggml_new_tensor failed");
-            return false;
-        }
-        ggml_set_name(tensor, name.c_str());
-
-        // LOG_DEBUG("%s %d %s %d[%d %d %d %d] %d[%d %d %d %d]", name.c_str(),
-        // ggml_nbytes(tensor), ggml_type_name(tensor_type),
-        // tensor_storage.n_dims,
-        // tensor_storage.ne[0], tensor_storage.ne[1], tensor_storage.ne[2], tensor_storage.ne[3],
-        // tensor->n_dims, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);
-
-        if (!tensor->data) {
-            GGML_ASSERT(ggml_nelements(tensor) == 0);
-            // avoid crashing the gguf writer by setting a dummy pointer for zero-sized tensors
-            LOG_DEBUG("setting dummy pointer for zero-sized tensor %s", name.c_str());
-            tensor->data = ggml_get_mem_buffer(ggml_ctx);
-        }
-
-        *dst_tensor = tensor;
-
-        gguf_add_tensor(gguf_ctx, tensor);
-
-        return true;
-    };
-
-    bool success = load_tensors(on_new_tensor_cb);
-    ggml_backend_free(backend);
-    LOG_INFO("load tensors done");
-    LOG_INFO("trying to save tensors to %s", file_path.c_str());
-    if (success) {
-        gguf_write_to_file(gguf_ctx, file_path.c_str(), false);
-    }
-    ggml_free(ggml_ctx);
-    gguf_free(gguf_ctx);
-    return success;
-}
-
 int64_t ModelLoader::get_params_mem_size(ggml_backend_t backend, ggml_type type) {
     size_t alignment = 128;
     if (backend != nullptr) {
@@ -1791,29 +1140,3 @@ int64_t ModelLoader::get_params_mem_size(ggml_backend_t backend, ggml_type type)
 
     return mem_size;
 }
-
-bool convert(const char* input_path,
-             const char* vae_path,
-             const char* output_path,
-             sd_type_t output_type,
-             const char* tensor_type_rules,
-             bool convert_name) {
-    ModelLoader model_loader;
-
-    if (!model_loader.init_from_file(input_path)) {
-        LOG_ERROR("init model loader from file failed: '%s'", input_path);
-        return false;
-    }
-
-    if (vae_path != nullptr && strlen(vae_path) > 0) {
-        if (!model_loader.init_from_file(vae_path, "vae.")) {
-            LOG_ERROR("init model loader from file failed: '%s'", vae_path);
-            return false;
-        }
-    }
-    if (convert_name) {
-        model_loader.convert_tensors_name();
-    }
-    bool success = model_loader.save_to_gguf_file(output_path, (ggml_type)output_type, tensor_type_rules);
-    return success;
-}

src/model.h

@@ -5,20 +5,13 @@
 #include <map>
 #include <memory>
 #include <set>
-#include <sstream>
 #include <string>
-#include <tuple>
-#include <utility>
 #include <vector>
 
 #include "ggml-backend.h"
 #include "ggml.h"
-#include "gguf.h"
-#include "json.hpp"
+#include "model_io/tensor_storage.h"
 #include "ordered_map.hpp"
-#include "zip.h"
-
-#define SD_MAX_DIMS 5
 
 enum SDVersion {
     VERSION_SD1,
@@ -195,116 +188,10 @@ enum PMVersion {
     PM_VERSION_2,
 };
 
-struct TensorStorage {
-    std::string name;
-    ggml_type type          = GGML_TYPE_F32;
-    ggml_type expected_type = GGML_TYPE_COUNT;
-    bool is_f8_e4m3         = false;
-    bool is_f8_e5m2         = false;
-    bool is_f64             = false;
-    bool is_i64             = false;
-    int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
-    int n_dims              = 0;
-
-    size_t file_index = 0;
-    int index_in_zip  = -1;  // >= means stored in a zip file
-    uint64_t offset   = 0;   // offset in file
-
-    TensorStorage() = default;
-
-    TensorStorage(std::string name, ggml_type type, const int64_t* ne, int n_dims, size_t file_index, size_t offset = 0)
-        : name(std::move(name)), type(type), n_dims(n_dims), file_index(file_index), offset(offset) {
-        for (int i = 0; i < n_dims; i++) {
-            this->ne[i] = ne[i];
-        }
-    }
-
-    int64_t nelements() const {
-        int64_t n = 1;
-        for (int i = 0; i < SD_MAX_DIMS; i++) {
-            n *= ne[i];
-        }
-        return n;
-    }
-
-    int64_t nbytes() const {
-        return nelements() * ggml_type_size(type) / ggml_blck_size(type);
-    }
-
-    int64_t nbytes_to_read() const {
-        if (is_f8_e4m3 || is_f8_e5m2) {
-            return nbytes() / 2;
-        } else if (is_f64 || is_i64) {
-            return nbytes() * 2;
-        } else {
-            return nbytes();
-        }
-    }
-
-    void unsqueeze() {
-        if (n_dims == 2) {
-            n_dims = 4;
-            ne[3]  = ne[1];
-            ne[2]  = ne[0];
-            ne[1]  = 1;
-            ne[0]  = 1;
-        }
-    }
-
-    std::vector<TensorStorage> chunk(size_t n) {
-        std::vector<TensorStorage> chunks;
-        uint64_t chunk_size = nbytes_to_read() / n;
-        // printf("%d/%d\n", chunk_size, nbytes_to_read());
-        reverse_ne();
-        for (size_t i = 0; i < n; i++) {
-            TensorStorage chunk_i = *this;
-            chunk_i.ne[0]         = ne[0] / n;
-            chunk_i.offset        = offset + i * chunk_size;
-            chunk_i.reverse_ne();
-            chunks.push_back(chunk_i);
-        }
-        reverse_ne();
-        return chunks;
-    }
-
-    void reverse_ne() {
-        int64_t new_ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
-        for (int i = 0; i < n_dims; i++) {
-            new_ne[i] = ne[n_dims - 1 - i];
-        }
-        for (int i = 0; i < n_dims; i++) {
-            ne[i] = new_ne[i];
-        }
-    }
-
-    std::string to_string() const {
-        std::stringstream ss;
-        const char* type_name = ggml_type_name(type);
-        if (is_f8_e4m3) {
-            type_name = "f8_e4m3";
-        } else if (is_f8_e5m2) {
-            type_name = "f8_e5m2";
-        } else if (is_f64) {
-            type_name = "f64";
-        } else if (is_i64) {
-            type_name = "i64";
-        }
-        ss << name << " | " << type_name << " | ";
-        ss << n_dims << " [";
-        for (int i = 0; i < SD_MAX_DIMS; i++) {
-            ss << ne[i];
-            if (i != SD_MAX_DIMS - 1) {
-                ss << ", ";
-            }
-        }
-        ss << "]";
-        return ss.str();
-    }
-};
-
-typedef std::function<bool(const TensorStorage&, ggml_tensor**)> on_new_tensor_cb_t;
-
 typedef OrderedMap<std::string, TensorStorage> String2TensorStorage;
+using TensorTypeRules = std::vector<std::pair<std::string, ggml_type>>;
+
+TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules);
 
 class ModelLoader {
 protected:
@@ -314,16 +201,10 @@ protected:
 
     void add_tensor_storage(const TensorStorage& tensor_storage);
 
-    bool parse_data_pkl(uint8_t* buffer,
-                        size_t buffer_size,
-                        zip_t* zip,
-                        std::string dir,
-                        size_t file_index,
-                        const std::string prefix);
-
     bool init_from_gguf_file(const std::string& file_path, const std::string& prefix = "");
     bool init_from_safetensors_file(const std::string& file_path, const std::string& prefix = "");
-    bool init_from_ckpt_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_torch_zip_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix = "");
     bool init_from_diffusers_file(const std::string& file_path, const std::string& prefix = "");
 
 public:
@@ -353,7 +234,6 @@ public:
         return names;
     }
 
-    bool save_to_gguf_file(const std::string& file_path, ggml_type type, const std::string& tensor_type_rules);
     bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type);
     int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT);
     ~ModelLoader() = default;

src/model_io/binary_io.h

@@ -0,0 +1,57 @@
+#ifndef __SD_MODEL_IO_BINARY_IO_H__
+#define __SD_MODEL_IO_BINARY_IO_H__
+
+#include <cstdint>
+#include <ostream>
+
+namespace model_io {
+
+    inline int32_t read_int(const uint8_t* buffer) {
+        uint32_t value = 0;
+        value |= static_cast<uint32_t>(buffer[3]) << 24;
+        value |= static_cast<uint32_t>(buffer[2]) << 16;
+        value |= static_cast<uint32_t>(buffer[1]) << 8;
+        value |= static_cast<uint32_t>(buffer[0]);
+        return static_cast<int32_t>(value);
+    }
+
+    inline uint16_t read_short(const uint8_t* buffer) {
+        uint16_t value = 0;
+        value |= static_cast<uint16_t>(buffer[1]) << 8;
+        value |= static_cast<uint16_t>(buffer[0]);
+        return value;
+    }
+
+    inline uint64_t read_u64(const uint8_t* buffer) {
+        uint64_t value = 0;
+        value |= static_cast<uint64_t>(buffer[7]) << 56;
+        value |= static_cast<uint64_t>(buffer[6]) << 48;
+        value |= static_cast<uint64_t>(buffer[5]) << 40;
+        value |= static_cast<uint64_t>(buffer[4]) << 32;
+        value |= static_cast<uint64_t>(buffer[3]) << 24;
+        value |= static_cast<uint64_t>(buffer[2]) << 16;
+        value |= static_cast<uint64_t>(buffer[1]) << 8;
+        value |= static_cast<uint64_t>(buffer[0]);
+        return value;
+    }
+
+    inline void write_u64(std::ostream& stream, uint64_t value) {
+        uint8_t buffer[8];
+        for (int i = 0; i < 8; ++i) {
+            buffer[i] = static_cast<uint8_t>((value >> (8 * i)) & 0xFF);
+        }
+        stream.write((const char*)buffer, sizeof(buffer));
+    }
+
+    inline int find_char(const uint8_t* buffer, int len, char c) {
+        for (int pos = 0; pos < len; pos++) {
+            if (buffer[pos] == (uint8_t)c) {
+                return pos;
+            }
+        }
+        return -1;
+    }
+
+}  // namespace model_io
+
+#endif  // __SD_MODEL_IO_BINARY_IO_H__

src/model_io/gguf_io.cpp

@@ -0,0 +1,123 @@
+#include "gguf_io.h"
+
+#include <cstdint>
+#include <fstream>
+#include <string>
+#include <vector>
+
+#include "gguf.h"
+#include "gguf_reader_ext.h"
+#include "util.h"
+
+static void set_error(std::string* error, const std::string& message) {
+    if (error != nullptr) {
+        *error = message;
+    }
+}
+
+bool is_gguf_file(const std::string& file_path) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        return false;
+    }
+
+    char magic[4];
+
+    file.read(magic, sizeof(magic));
+    if (!file) {
+        return false;
+    }
+    for (uint32_t i = 0; i < sizeof(magic); i++) {
+        if (magic[i] != GGUF_MAGIC[i]) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
+bool read_gguf_file(const std::string& file_path,
+                    std::vector<TensorStorage>& tensor_storages,
+                    std::string* error) {
+    tensor_storages.clear();
+
+    gguf_context* ctx_gguf_ = nullptr;
+    ggml_context* ctx_meta_ = nullptr;
+
+    ctx_gguf_ = gguf_init_from_file(file_path.c_str(), {true, &ctx_meta_});
+    if (!ctx_gguf_) {
+        GGUFReader gguf_reader;
+        if (!gguf_reader.load(file_path)) {
+            set_error(error, "failed to open '" + file_path + "' with GGUFReader");
+            return false;
+        }
+
+        size_t data_offset = gguf_reader.data_offset();
+        for (const auto& gguf_tensor_info : gguf_reader.tensors()) {
+            TensorStorage tensor_storage(
+                gguf_tensor_info.name,
+                gguf_tensor_info.type,
+                gguf_tensor_info.shape.data(),
+                static_cast<int>(gguf_tensor_info.shape.size()),
+                0,
+                data_offset + gguf_tensor_info.offset);
+
+            tensor_storages.push_back(tensor_storage);
+        }
+
+        return true;
+    }
+
+    int n_tensors = static_cast<int>(gguf_get_n_tensors(ctx_gguf_));
+
+    size_t data_offset = gguf_get_data_offset(ctx_gguf_);
+    for (int i = 0; i < n_tensors; i++) {
+        std::string name   = gguf_get_tensor_name(ctx_gguf_, i);
+        ggml_tensor* dummy = ggml_get_tensor(ctx_meta_, name.c_str());
+        size_t offset      = data_offset + gguf_get_tensor_offset(ctx_gguf_, i);
+
+        TensorStorage tensor_storage(name, dummy->type, dummy->ne, ggml_n_dims(dummy), 0, offset);
+
+        if (ggml_nbytes(dummy) != tensor_storage.nbytes()) {
+            gguf_free(ctx_gguf_);
+            ggml_free(ctx_meta_);
+            set_error(error, "size mismatch for tensor '" + name + "'");
+            return false;
+        }
+
+        tensor_storages.push_back(tensor_storage);
+    }
+
+    gguf_free(ctx_gguf_);
+    ggml_free(ctx_meta_);
+
+    return true;
+}
+
+bool write_gguf_file(const std::string& file_path,
+                     const std::vector<TensorWriteInfo>& tensors,
+                     std::string* error) {
+    gguf_context* gguf_ctx = gguf_init_empty();
+    if (gguf_ctx == nullptr) {
+        set_error(error, "gguf_init_empty failed");
+        return false;
+    }
+
+    for (const TensorWriteInfo& write_tensor : tensors) {
+        ggml_tensor* tensor = write_tensor.tensor;
+        if (tensor == nullptr) {
+            set_error(error, "null tensor cannot be written to GGUF");
+            gguf_free(gguf_ctx);
+            return false;
+        }
+        gguf_add_tensor(gguf_ctx, tensor);
+    }
+
+    LOG_INFO("trying to save tensors to %s", file_path.c_str());
+    bool success = gguf_write_to_file(gguf_ctx, file_path.c_str(), false);
+    if (!success) {
+        set_error(error, "failed to write GGUF file '" + file_path + "'");
+    }
+    gguf_free(gguf_ctx);
+    return success;
+}

src/model_io/gguf_io.h

@@ -0,0 +1,17 @@
+#ifndef __SD_MODEL_IO_GGUF_IO_H__
+#define __SD_MODEL_IO_GGUF_IO_H__
+
+#include <string>
+#include <vector>
+
+#include "tensor_storage.h"
+
+bool is_gguf_file(const std::string& file_path);
+bool read_gguf_file(const std::string& file_path,
+                    std::vector<TensorStorage>& tensor_storages,
+                    std::string* error = nullptr);
+bool write_gguf_file(const std::string& file_path,
+                     const std::vector<TensorWriteInfo>& tensors,
+                     std::string* error = nullptr);
+
+#endif  // __SD_MODEL_IO_GGUF_IO_H__

src/model_io/gguf_reader_ext.h

@@ -1,5 +1,5 @@
-#ifndef __GGUF_READER_HPP__
-#define __GGUF_READER_HPP__
+#ifndef __SD_MODEL_IO_GGUF_READER_EXT_H__
+#define __SD_MODEL_IO_GGUF_READER_EXT_H__
 
 #include <cstdint>
 #include <fstream>
@@ -231,4 +231,4 @@ public:
     size_t data_offset() const { return data_offset_; }
 };
 
-#endif  // __GGUF_READER_HPP__
+#endif  // __SD_MODEL_IO_GGUF_READER_EXT_H__

src/model_io/pickle_io.h

@@ -0,0 +1,21 @@
+#ifndef __SD_MODEL_IO_PICKLE_IO_H__
+#define __SD_MODEL_IO_PICKLE_IO_H__
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "tensor_storage.h"
+
+bool skip_pickle_object(const uint8_t* buffer, size_t buffer_size, size_t* object_size);
+bool pickle_object_is_torch_magic_number(const uint8_t* buffer, size_t buffer_size);
+bool parse_pickle_uint32_object(const uint8_t* buffer, size_t buffer_size, uint32_t* value);
+bool parse_torch_state_dict_pickle(const uint8_t* buffer,
+                                   size_t buffer_size,
+                                   std::vector<TensorStorage>& tensor_storages,
+                                   std::unordered_map<std::string, uint64_t>& storage_nbytes,
+                                   std::string* error = nullptr);
+
+#endif  // __SD_MODEL_IO_PICKLE_IO_H__

src/model_io/safetensors_io.cpp

@@ -0,0 +1,316 @@
+#include "safetensors_io.h"
+
+#include <cstdint>
+#include <exception>
+#include <fstream>
+#include <string>
+#include <vector>
+
+#include "binary_io.h"
+#include "json.hpp"
+#include "util.h"
+
+static constexpr size_t ST_HEADER_SIZE_LEN = 8;
+
+static void set_error(std::string* error, const std::string& message) {
+    if (error != nullptr) {
+        *error = message;
+    }
+}
+
+bool is_safetensors_file(const std::string& file_path) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        return false;
+    }
+
+    // get file size
+    file.seekg(0, file.end);
+    size_t file_size_ = file.tellg();
+    file.seekg(0, file.beg);
+
+    // read header size
+    if (file_size_ <= ST_HEADER_SIZE_LEN) {
+        return false;
+    }
+
+    uint8_t header_size_buf[ST_HEADER_SIZE_LEN];
+    file.read((char*)header_size_buf, ST_HEADER_SIZE_LEN);
+    if (!file) {
+        return false;
+    }
+
+    size_t header_size_ = model_io::read_u64(header_size_buf);
+    if (header_size_ >= file_size_ || header_size_ <= 2) {
+        return false;
+    }
+
+    // read header
+    std::vector<char> header_buf;
+    header_buf.resize(header_size_ + 1);
+    header_buf[header_size_] = '\0';
+    file.read(header_buf.data(), header_size_);
+    if (!file) {
+        return false;
+    }
+    try {
+        nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
+    } catch (const std::exception&) {
+        return false;
+    }
+    return true;
+}
+
+static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) {
+    ggml_type ttype = GGML_TYPE_COUNT;
+    if (dtype == "F16") {
+        ttype = GGML_TYPE_F16;
+    } else if (dtype == "BF16") {
+        ttype = GGML_TYPE_BF16;
+    } else if (dtype == "F32") {
+        ttype = GGML_TYPE_F32;
+    } else if (dtype == "F64") {
+        ttype = GGML_TYPE_F32;
+    } else if (dtype == "F8_E4M3") {
+        ttype = GGML_TYPE_F16;
+    } else if (dtype == "F8_E5M2") {
+        ttype = GGML_TYPE_F16;
+    } else if (dtype == "I32") {
+        ttype = GGML_TYPE_I32;
+    } else if (dtype == "I64") {
+        ttype = GGML_TYPE_I32;
+    }
+    return ttype;
+}
+
+// https://huggingface.co/docs/safetensors/index
+bool read_safetensors_file(const std::string& file_path,
+                           std::vector<TensorStorage>& tensor_storages,
+                           std::string* error) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        set_error(error, "failed to open '" + file_path + "'");
+        return false;
+    }
+
+    // get file size
+    file.seekg(0, file.end);
+    size_t file_size_ = file.tellg();
+    file.seekg(0, file.beg);
+
+    // read header size
+    if (file_size_ <= ST_HEADER_SIZE_LEN) {
+        set_error(error, "invalid safetensor file '" + file_path + "'");
+        return false;
+    }
+
+    uint8_t header_size_buf[ST_HEADER_SIZE_LEN];
+    file.read((char*)header_size_buf, ST_HEADER_SIZE_LEN);
+    if (!file) {
+        set_error(error, "read safetensors header size failed: '" + file_path + "'");
+        return false;
+    }
+
+    size_t header_size_ = model_io::read_u64(header_size_buf);
+    if (header_size_ >= file_size_) {
+        set_error(error, "invalid safetensor file '" + file_path + "'");
+        return false;
+    }
+
+    // read header
+    std::vector<char> header_buf;
+    header_buf.resize(header_size_ + 1);
+    header_buf[header_size_] = '\0';
+    file.read(header_buf.data(), header_size_);
+    if (!file) {
+        set_error(error, "read safetensors header failed: '" + file_path + "'");
+        return false;
+    }
+
+    nlohmann::json header_;
+    try {
+        header_ = nlohmann::json::parse(header_buf.data());
+    } catch (const std::exception&) {
+        set_error(error, "parsing safetensors header failed: '" + file_path + "'");
+        return false;
+    }
+
+    tensor_storages.clear();
+    for (auto& item : header_.items()) {
+        std::string name           = item.key();
+        nlohmann::json tensor_info = item.value();
+        // LOG_DEBUG("%s %s\n", name.c_str(), tensor_info.dump().c_str());
+
+        if (name == "__metadata__") {
+            continue;
+        }
+
+        std::string dtype    = tensor_info["dtype"];
+        nlohmann::json shape = tensor_info["shape"];
+
+        if (dtype == "U8") {
+            continue;
+        }
+
+        size_t begin = tensor_info["data_offsets"][0].get<size_t>();
+        size_t end   = tensor_info["data_offsets"][1].get<size_t>();
+
+        ggml_type type = safetensors_dtype_to_ggml_type(dtype);
+        if (type == GGML_TYPE_COUNT) {
+            set_error(error, "unsupported dtype '" + dtype + "' (tensor '" + name + "')");
+            return false;
+        }
+
+        if (shape.size() > SD_MAX_DIMS) {
+            set_error(error, "invalid tensor '" + name + "'");
+            return false;
+        }
+
+        int n_dims              = (int)shape.size();
+        int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
+        for (int i = 0; i < n_dims; i++) {
+            ne[i] = shape[i].get<int64_t>();
+        }
+
+        if (n_dims == 5) {
+            n_dims = 4;
+            ne[0]  = ne[0] * ne[1];
+            ne[1]  = ne[2];
+            ne[2]  = ne[3];
+            ne[3]  = ne[4];
+        }
+
+        // ggml_n_dims returns 1 for scalars
+        if (n_dims == 0) {
+            n_dims = 1;
+        }
+
+        TensorStorage tensor_storage(name, type, ne, n_dims, 0, ST_HEADER_SIZE_LEN + header_size_ + begin);
+        tensor_storage.reverse_ne();
+
+        size_t tensor_data_size = end - begin;
+
+        bool tensor_size_ok;
+        if (dtype == "F8_E4M3") {
+            tensor_storage.is_f8_e4m3 = true;
+            // f8 -> f16
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
+        } else if (dtype == "F8_E5M2") {
+            tensor_storage.is_f8_e5m2 = true;
+            // f8 -> f16
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
+        } else if (dtype == "F64") {
+            tensor_storage.is_f64 = true;
+            // f64 -> f32
+            tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
+        } else if (dtype == "I64") {
+            tensor_storage.is_i64 = true;
+            // i64 -> i32
+            tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
+        } else {
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size);
+        }
+        if (!tensor_size_ok) {
+            set_error(error, "size mismatch for tensor '" + name + "' (" + dtype + ")");
+            return false;
+        }
+
+        tensor_storages.push_back(tensor_storage);
+
+        // LOG_DEBUG("%s %s", tensor_storage.to_string().c_str(), dtype.c_str());
+    }
+
+    return true;
+}
+
+static bool ggml_type_to_safetensors_dtype(ggml_type type, std::string* dtype) {
+    switch (type) {
+        case GGML_TYPE_F16:
+            *dtype = "F16";
+            return true;
+        case GGML_TYPE_BF16:
+            *dtype = "BF16";
+            return true;
+        case GGML_TYPE_F32:
+            *dtype = "F32";
+            return true;
+        case GGML_TYPE_I32:
+            *dtype = "I32";
+            return true;
+        default:
+            return false;
+    }
+}
+
+bool write_safetensors_file(const std::string& file_path,
+                            const std::vector<TensorWriteInfo>& tensors,
+                            std::string* error) {
+    nlohmann::ordered_json header = nlohmann::ordered_json::object();
+
+    uint64_t data_offset = 0;
+    for (const TensorWriteInfo& write_tensor : tensors) {
+        ggml_tensor* tensor = write_tensor.tensor;
+        if (tensor == nullptr) {
+            set_error(error, "null tensor cannot be written to safetensors");
+            return false;
+        }
+
+        const std::string name = ggml_get_name(tensor);
+        std::string dtype;
+        if (!ggml_type_to_safetensors_dtype(tensor->type, &dtype)) {
+            set_error(error,
+                      "unsupported safetensors dtype '" + std::string(ggml_type_name(tensor->type)) +
+                          "' for tensor '" + name + "'");
+            return false;
+        }
+
+        const uint64_t tensor_nbytes = ggml_nbytes(tensor);
+
+        nlohmann::ordered_json json_tensor_info = nlohmann::ordered_json::object();
+        json_tensor_info["dtype"]               = dtype;
+
+        nlohmann::ordered_json shape = nlohmann::ordered_json::array();
+        for (int i = 0; i < write_tensor.n_dims; ++i) {
+            shape.push_back(write_tensor.ne[write_tensor.n_dims - 1 - i]);
+        }
+        json_tensor_info["shape"] = shape;
+
+        nlohmann::ordered_json data_offsets = nlohmann::ordered_json::array();
+        data_offsets.push_back(data_offset);
+        data_offsets.push_back(data_offset + tensor_nbytes);
+        json_tensor_info["data_offsets"] = data_offsets;
+
+        header[name] = json_tensor_info;
+        data_offset += tensor_nbytes;
+    }
+
+    const std::string header_str = header.dump();
+
+    std::ofstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        set_error(error, "failed to open '" + file_path + "' for writing");
+        return false;
+    }
+
+    LOG_INFO("trying to save tensors to %s", file_path.c_str());
+    model_io::write_u64(file, header_str.size());
+    file.write(header_str.data(), header_str.size());
+    if (!file) {
+        set_error(error, "failed to write safetensors header to '" + file_path + "'");
+        return false;
+    }
+
+    for (const TensorWriteInfo& write_tensor : tensors) {
+        ggml_tensor* tensor        = write_tensor.tensor;
+        const std::string name     = ggml_get_name(tensor);
+        const size_t tensor_nbytes = ggml_nbytes(tensor);
+        file.write((const char*)tensor->data, tensor_nbytes);
+        if (!file) {
+            set_error(error,
+                      "failed to write tensor '" + name + "' to '" + file_path + "'");
+            return false;
+        }
+    }
+
+    return true;
+}

src/model_io/safetensors_io.h

@@ -0,0 +1,17 @@
+#ifndef __SD_MODEL_IO_SAFETENSORS_IO_H__
+#define __SD_MODEL_IO_SAFETENSORS_IO_H__
+
+#include <string>
+#include <vector>
+
+#include "tensor_storage.h"
+
+bool is_safetensors_file(const std::string& file_path);
+bool read_safetensors_file(const std::string& file_path,
+                           std::vector<TensorStorage>& tensor_storages,
+                           std::string* error = nullptr);
+bool write_safetensors_file(const std::string& file_path,
+                            const std::vector<TensorWriteInfo>& tensors,
+                            std::string* error = nullptr);
+
+#endif  // __SD_MODEL_IO_SAFETENSORS_IO_H__

src/model_io/tensor_storage.h

@@ -0,0 +1,132 @@
+#ifndef __SD_TENSOR_STORAGE_H__
+#define __SD_TENSOR_STORAGE_H__
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <sstream>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "ggml.h"
+
+#define SD_MAX_DIMS 5
+
+struct TensorStorage {
+    std::string name;
+    ggml_type type          = GGML_TYPE_F32;
+    ggml_type expected_type = GGML_TYPE_COUNT;
+    bool is_f8_e4m3         = false;
+    bool is_f8_e5m2         = false;
+    bool is_f64             = false;
+    bool is_i64             = false;
+    int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
+    int n_dims              = 0;
+
+    std::string storage_key;
+    size_t file_index = 0;
+    int index_in_zip  = -1;  // >= means stored in a zip file
+    uint64_t offset   = 0;   // offset in file
+
+    TensorStorage() = default;
+
+    TensorStorage(std::string name, ggml_type type, const int64_t* ne, int n_dims, size_t file_index, size_t offset = 0)
+        : name(std::move(name)), type(type), n_dims(n_dims), file_index(file_index), offset(offset) {
+        for (int i = 0; i < n_dims; i++) {
+            this->ne[i] = ne[i];
+        }
+    }
+
+    int64_t nelements() const {
+        int64_t n = 1;
+        for (int i = 0; i < SD_MAX_DIMS; i++) {
+            n *= ne[i];
+        }
+        return n;
+    }
+
+    int64_t nbytes() const {
+        return nelements() * ggml_type_size(type) / ggml_blck_size(type);
+    }
+
+    int64_t nbytes_to_read() const {
+        if (is_f8_e4m3 || is_f8_e5m2) {
+            return nbytes() / 2;
+        } else if (is_f64 || is_i64) {
+            return nbytes() * 2;
+        } else {
+            return nbytes();
+        }
+    }
+
+    void unsqueeze() {
+        if (n_dims == 2) {
+            n_dims = 4;
+            ne[3]  = ne[1];
+            ne[2]  = ne[0];
+            ne[1]  = 1;
+            ne[0]  = 1;
+        }
+    }
+
+    std::vector<TensorStorage> chunk(size_t n) {
+        std::vector<TensorStorage> chunks;
+        uint64_t chunk_size = nbytes_to_read() / n;
+        // printf("%d/%d\n", chunk_size, nbytes_to_read());
+        reverse_ne();
+        for (size_t i = 0; i < n; i++) {
+            TensorStorage chunk_i = *this;
+            chunk_i.ne[0]         = ne[0] / n;
+            chunk_i.offset        = offset + i * chunk_size;
+            chunk_i.reverse_ne();
+            chunks.push_back(chunk_i);
+        }
+        reverse_ne();
+        return chunks;
+    }
+
+    void reverse_ne() {
+        int64_t new_ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
+        for (int i = 0; i < n_dims; i++) {
+            new_ne[i] = ne[n_dims - 1 - i];
+        }
+        for (int i = 0; i < n_dims; i++) {
+            ne[i] = new_ne[i];
+        }
+    }
+
+    std::string to_string() const {
+        std::stringstream ss;
+        const char* type_name = ggml_type_name(type);
+        if (is_f8_e4m3) {
+            type_name = "f8_e4m3";
+        } else if (is_f8_e5m2) {
+            type_name = "f8_e5m2";
+        } else if (is_f64) {
+            type_name = "f64";
+        } else if (is_i64) {
+            type_name = "i64";
+        }
+        ss << name << " | " << type_name << " | ";
+        ss << n_dims << " [";
+        for (int i = 0; i < SD_MAX_DIMS; i++) {
+            ss << ne[i];
+            if (i != SD_MAX_DIMS - 1) {
+                ss << ", ";
+            }
+        }
+        ss << "]";
+        return ss.str();
+    }
+};
+
+struct TensorWriteInfo {
+    int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
+    int n_dims              = 0;
+    ggml_tensor* tensor     = nullptr;
+};
+
+typedef std::function<bool(const TensorStorage&, ggml_tensor**)> on_new_tensor_cb_t;
+
+#endif  // __SD_TENSOR_STORAGE_H__

src/model_io/torch_legacy_io.cpp

@@ -0,0 +1,252 @@
+#include "torch_legacy_io.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <fstream>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "pickle_io.h"
+#include "util.h"
+
+// torch.save format background:
+//
+//   - Before PyTorch 1.6.0, torch.save used this legacy non-zip format by
+//     default.
+//   - Since PyTorch 1.6.0, torch.save defaults to an uncompressed ZIP64 archive
+//     containing data.pkl, data/, version, and, since PyTorch 2.1.0, byteorder.
+//   - The old format can still be produced explicitly with:
+//       torch.save(obj, path, _use_new_zipfile_serialization=False)
+//
+// Whether obj is a state_dict or a whole nn.Module does not change the outer
+// container format selected by torch.save. It changes the pickled object inside:
+//
+//   - state_dict: usually an OrderedDict[str, Tensor]. pickle_io.cpp supports a
+//     restricted subset of this layout because tensor metadata and raw storages
+//     can be recovered without executing pickle callables.
+//   - whole module/checkpoint object: arbitrary Python object graph. This may
+//     require importing user classes and executing pickle GLOBAL/REDUCE rebuild
+//     logic, so it is intentionally not supported here.
+//
+// Legacy non-zip PyTorch files are not a single pickle object:
+//
+//   1. pickle object: PyTorch legacy magic number
+//   2. pickle object: legacy protocol version, expected to be 1001
+//   3. pickle object: sys_info metadata, ignored by this reader
+//   4. pickle object: state_dict metadata, parsed by pickle_io.cpp
+//   5. pickle object: serialized storage key list, skipped here
+//   6. raw storage data payloads
+//      - PyTorch writes storages after the pickles, ordered by storage key
+//      - each storage has an 8-byte legacy storage header followed by raw bytes
+static constexpr size_t LEGACY_STORAGE_HEADER_SIZE = 8;
+
+static void set_error(std::string* error, const std::string& message) {
+    if (error != nullptr) {
+        *error = message;
+    }
+}
+
+static std::string bytes_to_hex(const std::vector<uint8_t>& bytes) {
+    static const char* hex = "0123456789ABCDEF";
+    std::string result;
+    result.reserve(bytes.size() * 3);
+    for (size_t i = 0; i < bytes.size(); ++i) {
+        if (i > 0) {
+            result.push_back('-');
+        }
+        result.push_back(hex[(bytes[i] >> 4) & 0x0F]);
+        result.push_back(hex[bytes[i] & 0x0F]);
+    }
+    return result;
+}
+
+static bool is_probably_tar_file(const std::vector<uint8_t>& header) {
+    return header.size() >= 262 &&
+           header[257] == 'u' &&
+           header[258] == 's' &&
+           header[259] == 't' &&
+           header[260] == 'a' &&
+           header[261] == 'r';
+}
+
+static std::string torch_legacy_diagnostics(const std::string& file_path, const std::vector<uint8_t>& buffer) {
+    if (!ends_with(file_path, ".pt") && !ends_with(file_path, ".pth")) {
+        return "";
+    }
+    if (buffer.empty()) {
+        return "unsupported PyTorch file '" + file_path + "': empty file";
+    }
+
+    size_t short_len = std::min<size_t>(buffer.size(), 32);
+    std::vector<uint8_t> short_header(buffer.begin(), buffer.begin() + short_len);
+    const bool raw_pickle = buffer[0] == 0x80;
+    const bool tar_file   = is_probably_tar_file(buffer);
+
+    std::string message = "unsupported PyTorch file '" + file_path + "': first bytes " +
+                          bytes_to_hex(short_header) +
+                          ", raw_pickle=" + (raw_pickle ? "true" : "false") +
+                          ", tar=" + (tar_file ? "true" : "false");
+    if (raw_pickle) {
+        message += "; raw pickle did not match the restricted state_dict layouts currently supported";
+    } else if (tar_file) {
+        message += "; legacy tar PyTorch checkpoints are not supported yet";
+    }
+    return message;
+}
+
+bool read_torch_legacy_file(const std::string& file_path,
+                            std::vector<TensorStorage>& tensor_storages,
+                            std::string* error) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        set_error(error, "failed to open '" + file_path + "'");
+        return false;
+    }
+
+    file.seekg(0, file.end);
+    size_t file_size = (size_t)file.tellg();
+    file.seekg(0, file.beg);
+    if (file_size == 0) {
+        set_error(error, "empty file '" + file_path + "'");
+        return false;
+    }
+
+    std::vector<uint8_t> buffer(file_size);
+    file.read((char*)buffer.data(), file_size);
+    if (!file) {
+        set_error(error, "failed to read '" + file_path + "'");
+        return false;
+    }
+
+    auto finalize_tensor_offsets = [&](size_t storage_data_offset,
+                                       const std::unordered_map<std::string, uint64_t>& legacy_storage_map) -> bool {
+        if (storage_data_offset > file_size) {
+            return false;
+        }
+
+        std::vector<std::string> storage_keys;
+        storage_keys.reserve(legacy_storage_map.size());
+        for (const auto& [storage_key, _] : legacy_storage_map) {
+            storage_keys.push_back(storage_key);
+        }
+        std::sort(storage_keys.begin(), storage_keys.end());
+
+        std::unordered_map<std::string, uint64_t> storage_offsets;
+        uint64_t current_offset = storage_data_offset;
+        for (const auto& storage_key : storage_keys) {
+            auto it = legacy_storage_map.find(storage_key);
+            if (it == legacy_storage_map.end()) {
+                return false;
+            }
+            if (current_offset + LEGACY_STORAGE_HEADER_SIZE + it->second > file_size) {
+                return false;
+            }
+            storage_offsets[storage_key] = current_offset + LEGACY_STORAGE_HEADER_SIZE;
+            current_offset += LEGACY_STORAGE_HEADER_SIZE + it->second;
+        }
+
+        for (auto& tensor_storage : tensor_storages) {
+            if (tensor_storage.storage_key.empty()) {
+                continue;
+            }
+
+            auto it_offset = storage_offsets.find(tensor_storage.storage_key);
+            auto it_size   = legacy_storage_map.find(tensor_storage.storage_key);
+            if (it_offset == storage_offsets.end() || it_size == legacy_storage_map.end()) {
+                return false;
+            }
+
+            uint64_t base_offset    = it_offset->second;
+            uint64_t storage_nbytes = it_size->second;
+            uint64_t tensor_nbytes  = tensor_storage.nbytes_to_read();
+            if (tensor_storage.offset + tensor_nbytes > storage_nbytes) {
+                return false;
+            }
+
+            tensor_storage.offset = base_offset + tensor_storage.offset;
+            tensor_storage.storage_key.clear();
+        }
+
+        return true;
+    };
+
+    auto parse_state_dict_at = [&](size_t state_dict_offset, size_t state_dict_size, size_t* storage_data_offset) -> bool {
+        tensor_storages.clear();
+        std::unordered_map<std::string, uint64_t> legacy_storage_map;
+        if (!parse_torch_state_dict_pickle(buffer.data() + state_dict_offset,
+                                           state_dict_size,
+                                           tensor_storages,
+                                           legacy_storage_map,
+                                           error)) {
+            return false;
+        }
+
+        size_t offset_after_state_dict = state_dict_offset + state_dict_size;
+        size_t storage_keys_size       = 0;
+        if (!skip_pickle_object(buffer.data() + offset_after_state_dict,
+                                buffer.size() - offset_after_state_dict,
+                                &storage_keys_size)) {
+            return false;
+        }
+
+        *storage_data_offset = offset_after_state_dict + storage_keys_size;
+        return finalize_tensor_offsets(*storage_data_offset, legacy_storage_map);
+    };
+
+    size_t object_size_1 = 0;
+    size_t offset        = 0;
+
+    if (skip_pickle_object(buffer.data(), buffer.size(), &object_size_1) &&
+        pickle_object_is_torch_magic_number(buffer.data(), object_size_1)) {
+        offset += object_size_1;
+
+        size_t object_size_2 = 0;
+        if (!skip_pickle_object(buffer.data() + offset, buffer.size() - offset, &object_size_2)) {
+            set_error(error, torch_legacy_diagnostics(file_path, buffer));
+            return false;
+        }
+        uint32_t protocol_version = 0;
+        if (!parse_pickle_uint32_object(buffer.data() + offset, object_size_2, &protocol_version) || protocol_version != 1001) {
+            set_error(error, torch_legacy_diagnostics(file_path, buffer));
+            return false;
+        }
+        offset += object_size_2;
+
+        size_t object_size_3 = 0;
+        if (!skip_pickle_object(buffer.data() + offset, buffer.size() - offset, &object_size_3)) {
+            set_error(error, torch_legacy_diagnostics(file_path, buffer));
+            return false;
+        }
+        offset += object_size_3;
+
+        size_t state_dict_size = 0;
+        if (!skip_pickle_object(buffer.data() + offset, buffer.size() - offset, &state_dict_size)) {
+            set_error(error, torch_legacy_diagnostics(file_path, buffer));
+            return false;
+        }
+
+        size_t storage_data_offset = 0;
+        if (parse_state_dict_at(offset, state_dict_size, &storage_data_offset)) {
+            return true;
+        }
+
+        if (error != nullptr && error->empty()) {
+            set_error(error, torch_legacy_diagnostics(file_path, buffer));
+        }
+        return false;
+    }
+
+    size_t state_dict_size = 0;
+    if (skip_pickle_object(buffer.data(), buffer.size(), &state_dict_size)) {
+        size_t storage_data_offset = 0;
+        if (parse_state_dict_at(0, state_dict_size, &storage_data_offset)) {
+            return true;
+        }
+    }
+
+    if (error != nullptr && error->empty()) {
+        set_error(error, torch_legacy_diagnostics(file_path, buffer));
+    }
+    return false;
+}

src/model_io/torch_legacy_io.h

@@ -0,0 +1,13 @@
+#ifndef __SD_MODEL_IO_TORCH_LEGACY_IO_H__
+#define __SD_MODEL_IO_TORCH_LEGACY_IO_H__
+
+#include <string>
+#include <vector>
+
+#include "tensor_storage.h"
+
+bool read_torch_legacy_file(const std::string& file_path,
+                            std::vector<TensorStorage>& tensor_storages,
+                            std::string* error = nullptr);
+
+#endif  // __SD_MODEL_IO_TORCH_LEGACY_IO_H__

src/model_io/torch_zip_io.cpp

@@ -0,0 +1,140 @@
+#include "torch_zip_io.h"
+
+#include <cstdint>
+#include <cstdlib>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "pickle_io.h"
+
+#include "zip.h"
+
+static void set_error(std::string* error, const std::string& message) {
+    if (error != nullptr) {
+        *error = message;
+    }
+}
+
+bool is_torch_zip_file(const std::string& file_path) {
+    zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
+    if (zip == nullptr) {
+        return false;
+    }
+    zip_close(zip);
+    return true;
+}
+
+static bool find_zip_entry(zip_t* zip, const std::string& entry_name, int* index, uint64_t* size) {
+    size_t n = zip_entries_total(zip);
+    for (size_t i = 0; i < n; ++i) {
+        zip_entry_openbyindex(zip, i);
+        std::string name = zip_entry_name(zip);
+        if (name == entry_name) {
+            *index = (int)i;
+            *size  = zip_entry_size(zip);
+            zip_entry_close(zip);
+            return true;
+        }
+        zip_entry_close(zip);
+    }
+    return false;
+}
+
+static bool parse_zip_data_pkl(const uint8_t* buffer,
+                               size_t buffer_size,
+                               zip_t* zip,
+                               const std::string& dir,
+                               std::vector<TensorStorage>& tensor_storages,
+                               std::string* error) {
+    std::vector<TensorStorage> parsed_tensors;
+    std::unordered_map<std::string, uint64_t> storage_nbytes;
+    if (!parse_torch_state_dict_pickle(buffer, buffer_size, parsed_tensors, storage_nbytes, error)) {
+        if (error != nullptr && error->empty()) {
+            *error = "failed to parse torch zip pickle metadata";
+        }
+        return false;
+    }
+
+    for (auto& tensor_storage : parsed_tensors) {
+        if (tensor_storage.storage_key.empty()) {
+            set_error(error, "tensor '" + tensor_storage.name + "' has no storage key");
+            return false;
+        }
+
+        const std::string entry_name = dir + "data/" + tensor_storage.storage_key;
+        int zip_index                = -1;
+        uint64_t entry_size          = 0;
+        if (!find_zip_entry(zip, entry_name, &zip_index, &entry_size)) {
+            set_error(error, "storage entry '" + entry_name + "' was not found");
+            return false;
+        }
+
+        auto it_storage_size = storage_nbytes.find(tensor_storage.storage_key);
+        if (it_storage_size != storage_nbytes.end() && entry_size < it_storage_size->second) {
+            set_error(error, "storage entry '" + entry_name + "' is smaller than pickle metadata");
+            return false;
+        }
+
+        uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
+        if (tensor_storage.offset + tensor_nbytes > entry_size) {
+            set_error(error, "tensor '" + tensor_storage.name + "' exceeds storage entry '" + entry_name + "'");
+            return false;
+        }
+
+        tensor_storage.index_in_zip = zip_index;
+        tensor_storage.storage_key.clear();
+        tensor_storages.push_back(tensor_storage);
+    }
+
+    return true;
+}
+
+bool read_torch_zip_file(const std::string& file_path,
+                         std::vector<TensorStorage>& tensor_storages,
+                         std::string* error) {
+    zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
+    if (zip == nullptr) {
+        set_error(error, "failed to open '" + file_path + "'");
+        return false;
+    }
+
+    tensor_storages.clear();
+    bool success        = true;
+    bool found_data_pkl = false;
+    int n               = (int)zip_entries_total(zip);
+    for (int i = 0; i < n; ++i) {
+        zip_entry_openbyindex(zip, i);
+        std::string name = zip_entry_name(zip);
+        size_t pos       = name.find("data.pkl");
+        if (pos != std::string::npos) {
+            found_data_pkl  = true;
+            std::string dir = name.substr(0, pos);
+            void* pkl_data  = nullptr;
+            size_t pkl_size = 0;
+            zip_entry_read(zip, &pkl_data, &pkl_size);
+
+            if (pkl_data == nullptr || pkl_size == 0) {
+                set_error(error, "failed to read '" + name + "' from '" + file_path + "'");
+                success = false;
+            } else if (!parse_zip_data_pkl((const uint8_t*)pkl_data, pkl_size, zip, dir, tensor_storages, error)) {
+                success = false;
+            }
+
+            free(pkl_data);
+        }
+        zip_entry_close(zip);
+
+        if (!success) {
+            break;
+        }
+    }
+
+    if (success && !found_data_pkl) {
+        set_error(error, "data.pkl was not found in '" + file_path + "'");
+        success = false;
+    }
+
+    zip_close(zip);
+    return success;
+}

src/model_io/torch_zip_io.h

@@ -0,0 +1,14 @@
+#ifndef __SD_MODEL_IO_TORCH_ZIP_IO_H__
+#define __SD_MODEL_IO_TORCH_ZIP_IO_H__
+
+#include <string>
+#include <vector>
+
+#include "tensor_storage.h"
+
+bool is_torch_zip_file(const std::string& file_path);
+bool read_torch_zip_file(const std::string& file_path,
+                         std::vector<TensorStorage>& tensor_storages,
+                         std::string* error = nullptr);
+
+#endif  // __SD_MODEL_IO_TORCH_ZIP_IO_H__

src/stable-diffusion.cpp

@@ -17,6 +17,7 @@
 #include "pmid.hpp"
 #include "sample-cache.h"
 #include "tae.hpp"
+#include "upscaler.h"
 #include "vae.hpp"
 
 #include "latent-preview.h"
@@ -2113,6 +2114,28 @@ enum lora_apply_mode_t str_to_lora_apply_mode(const char* str) {
     return LORA_APPLY_MODE_COUNT;
 }
 
+const char* hires_upscaler_to_str[] = {
+    "None",
+    "Latent (nearest)",
+    "Model",
+};
+
+const char* sd_hires_upscaler_name(enum sd_hires_upscaler_t upscaler) {
+    if (upscaler < SD_HIRES_UPSCALER_COUNT) {
+        return hires_upscaler_to_str[upscaler];
+    }
+    return NONE_STR;
+}
+
+enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str) {
+    for (int i = 0; i < SD_HIRES_UPSCALER_COUNT; i++) {
+        if (!strcmp(str, hires_upscaler_to_str[i])) {
+            return (enum sd_hires_upscaler_t)i;
+        }
+    }
+    return SD_HIRES_UPSCALER_COUNT;
+}
+
 void sd_cache_params_init(sd_cache_params_t* cache_params) {
     *cache_params                             = {};
     cache_params->mode                        = SD_CACHE_DISABLED;
@@ -2141,6 +2164,19 @@ void sd_cache_params_init(sd_cache_params_t* cache_params) {
     cache_params->spectrum_stop_percent       = 0.9f;
 }
 
+void sd_hires_params_init(sd_hires_params_t* hires_params) {
+    *hires_params                    = {};
+    hires_params->enabled            = false;
+    hires_params->upscaler           = SD_HIRES_UPSCALER_LATENT_NEAREST;
+    hires_params->model_path         = nullptr;
+    hires_params->scale              = 2.0f;
+    hires_params->target_width       = 0;
+    hires_params->target_height      = 0;
+    hires_params->steps              = 0;
+    hires_params->denoising_strength = 0.7f;
+    hires_params->upscale_tile_size  = 128;
+}
+
 void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
     *sd_ctx_params                         = {};
     sd_ctx_params->vae_decode_only         = true;
@@ -2310,6 +2346,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
     sd_img_gen_params->pm_params         = {nullptr, 0, nullptr, 20.f};
     sd_img_gen_params->vae_tiling_params = {false, 0, 0, 0.5f, 0.0f, 0.0f};
     sd_cache_params_init(&sd_img_gen_params->cache);
+    sd_hires_params_init(&sd_img_gen_params->hires);
 }
 
 char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
@@ -2336,7 +2373,8 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              "increase_ref_index: %s\n"
              "control_strength: %.2f\n"
              "photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\n"
-             "VAE tiling: %s\n",
+             "VAE tiling: %s\n"
+             "hires: {enabled=%s, upscaler=%s, model_path=%s, scale=%.2f, target=%dx%d, steps=%d, denoising_strength=%.2f}\n",
              SAFE_STR(sd_img_gen_params->prompt),
              SAFE_STR(sd_img_gen_params->negative_prompt),
              sd_img_gen_params->clip_skip,
@@ -2353,7 +2391,15 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              sd_img_gen_params->pm_params.style_strength,
              sd_img_gen_params->pm_params.id_images_count,
              SAFE_STR(sd_img_gen_params->pm_params.id_embed_path),
-             BOOL_STR(sd_img_gen_params->vae_tiling_params.enabled));
+             BOOL_STR(sd_img_gen_params->vae_tiling_params.enabled),
+             BOOL_STR(sd_img_gen_params->hires.enabled),
+             sd_hires_upscaler_name(sd_img_gen_params->hires.upscaler),
+             SAFE_STR(sd_img_gen_params->hires.model_path),
+             sd_img_gen_params->hires.scale,
+             sd_img_gen_params->hires.target_width,
+             sd_img_gen_params->hires.target_height,
+             sd_img_gen_params->hires.steps,
+             sd_img_gen_params->hires.denoising_strength);
     const char* cache_mode_str = "disabled";
     if (sd_img_gen_params->cache.mode == SD_CACHE_EASYCACHE) {
         cache_mode_str = "easycache";
@@ -2457,8 +2503,10 @@ enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_me
             return EXPONENTIAL_SCHEDULER;
         }
     }
-    if (sample_method == LCM_SAMPLE_METHOD) {
+    if (sample_method == LCM_SAMPLE_METHOD || sample_method == TCD_SAMPLE_METHOD) {
         return LCM_SCHEDULER;
+    } else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) {
+        return SIMPLE_SCHEDULER;
     }
     return DISCRETE_SCHEDULER;
 }
@@ -2532,6 +2580,7 @@ struct GenerationRequest {
     sd_guidance_params_t guidance            = {};
     sd_guidance_params_t high_noise_guidance = {};
     sd_pm_params_t pm_params                 = {};
+    sd_hires_params_t hires                  = {};
     int frames                               = -1;
     float vace_strength                      = 1.f;
 
@@ -2553,6 +2602,7 @@ struct GenerationRequest {
         auto_resize_ref_image       = sd_img_gen_params->auto_resize_ref_image;
         guidance                    = sd_img_gen_params->sample_params.guidance;
         pm_params                   = sd_img_gen_params->pm_params;
+        hires                       = sd_img_gen_params->hires;
         cache_params                = &sd_img_gen_params->cache;
         resolve(sd_ctx);
     }
@@ -2575,26 +2625,76 @@ struct GenerationRequest {
     }
 
     void align_generation_request_size() {
+        align_image_size(&width, &height, "generation request");
+    }
+
+    void align_image_size(int* target_width, int* target_height, const char* label) {
         int spatial_multiple = vae_scale_factor * diffusion_model_down_factor;
-        int width_offset     = align_up_offset(width, spatial_multiple);
-        int height_offset    = align_up_offset(height, spatial_multiple);
+        int width_offset     = align_up_offset(*target_width, spatial_multiple);
+        int height_offset    = align_up_offset(*target_height, spatial_multiple);
         if (width_offset <= 0 && height_offset <= 0) {
             return;
         }
 
-        int original_width  = width;
-        int original_height = height;
+        int original_width  = *target_width;
+        int original_height = *target_height;
 
-        width += width_offset;
-        height += height_offset;
-        LOG_WARN("align up %dx%d to %dx%d (multiple=%d)",
+        *target_width += width_offset;
+        *target_height += height_offset;
+        LOG_WARN("align %s up %dx%d to %dx%d (multiple=%d)",
+                 label,
                  original_width,
                  original_height,
-                 width,
-                 height,
+                 *target_width,
+                 *target_height,
                  spatial_multiple);
     }
 
+    void resolve_hires() {
+        if (!hires.enabled) {
+            return;
+        }
+        if (hires.upscaler == SD_HIRES_UPSCALER_NONE) {
+            hires.enabled = false;
+            return;
+        }
+        if (hires.upscaler < SD_HIRES_UPSCALER_NONE && hires.upscaler >= SD_HIRES_UPSCALER_COUNT) {
+            LOG_WARN("hires upscaler '%d' is invalid, disabling hires", hires.upscaler);
+            hires.enabled = false;
+            return;
+        }
+        if (hires.upscaler == SD_HIRES_UPSCALER_MODEL && strlen(SAFE_STR(hires.model_path)) == 0) {
+            LOG_WARN("hires model upscaler requires a model path, disabling hires");
+            hires.enabled = false;
+            return;
+        }
+        if (hires.scale <= 0.f && hires.target_width <= 0 && hires.target_height <= 0) {
+            LOG_WARN("hires scale must be positive when no target size is set, disabling hires");
+            hires.enabled = false;
+            return;
+        }
+        hires.denoising_strength = std::clamp(hires.denoising_strength, 0.0001f, 1.f);
+        hires.steps              = std::max(0, hires.steps);
+
+        if (hires.target_width > 0 && hires.target_height > 0) {
+            // pass
+        } else if (hires.target_width > 0) {
+            hires.target_height = hires.target_width;
+        } else if (hires.target_height > 0) {
+            hires.target_width = hires.target_height;
+        } else {
+            hires.target_width  = static_cast<int>(std::round(width * hires.scale));
+            hires.target_height = static_cast<int>(std::round(height * hires.scale));
+        }
+
+        if (hires.target_width <= 0 || hires.target_height <= 0) {
+            LOG_WARN("hires target size is not positive, disabling hires");
+            hires.enabled = false;
+            return;
+        }
+        align_image_size(&hires.target_width, &hires.target_height, "hires target");
+    }
+
     static void resolve_guidance(sd_ctx_t* sd_ctx,
                                  sd_guidance_params_t* guidance,
                                  bool* use_uncond,
@@ -2635,6 +2735,7 @@ struct GenerationRequest {
 
     void resolve(sd_ctx_t* sd_ctx) {
         align_generation_request_size();
+        resolve_hires();
         seed = resolve_seed(seed);
 
         resolve_guidance(sd_ctx, &guidance, &use_uncond, &use_img_cond);
@@ -3147,6 +3248,67 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
     return result_images;
 }
 
+static sd::Tensor<float> upscale_hires_latent(sd_ctx_t* sd_ctx,
+                                              const sd::Tensor<float>& latent,
+                                              const GenerationRequest& request,
+                                              UpscalerGGML* upscaler) {
+    if (request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_NEAREST) {
+        std::vector<int64_t> target_shape = latent.shape();
+        if (target_shape.size() < 2) {
+            LOG_ERROR("latent has invalid shape for hires upscale");
+            return {};
+        }
+        target_shape[0] = request.hires.target_width / request.vae_scale_factor;
+        target_shape[1] = request.hires.target_height / request.vae_scale_factor;
+
+        LOG_INFO("hires latent upscale %" PRId64 "x%" PRId64 " -> %" PRId64 "x%" PRId64,
+                 latent.shape()[0],
+                 latent.shape()[1],
+                 target_shape[0],
+                 target_shape[1]);
+        return sd::ops::interpolate(latent, target_shape, sd::ops::InterpolateMode::Nearest);
+    } else if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
+        if (upscaler == nullptr) {
+            LOG_ERROR("hires model upscaler context is null");
+            return {};
+        }
+        if (sd_ctx->sd->vae_decode_only) {
+            LOG_ERROR("hires model upscaler requires VAE encoder weights; create the context with vae_decode_only=false");
+            return {};
+        }
+
+        sd::Tensor<float> decoded = sd_ctx->sd->decode_first_stage(latent);
+        if (decoded.empty()) {
+            LOG_ERROR("decode_first_stage failed before hires model upscale");
+            return {};
+        }
+
+        sd::Tensor<float> upscaled_tensor = upscaler->upscale_tensor(decoded);
+        if (upscaled_tensor.empty()) {
+            LOG_ERROR("hires model upscale failed");
+            return {};
+        }
+
+        if (upscaled_tensor.shape()[0] != request.hires.target_width ||
+            upscaled_tensor.shape()[1] != request.hires.target_height) {
+            upscaled_tensor = sd::ops::interpolate(upscaled_tensor,
+                                                   {request.hires.target_width,
+                                                    request.hires.target_height,
+                                                    upscaled_tensor.shape()[2],
+                                                    upscaled_tensor.shape()[3]});
+        }
+
+        sd::Tensor<float> upscaled_latent = sd_ctx->sd->encode_first_stage(upscaled_tensor);
+        if (upscaled_latent.empty()) {
+            LOG_ERROR("encode_first_stage failed after hires model upscale");
+        }
+        return upscaled_latent;
+    }
+
+    LOG_ERROR("unsupported hires upscaler '%s'", sd_hires_upscaler_name(request.hires.upscaler));
+    return {};
+}
+
 SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params) {
     if (sd_ctx == nullptr || sd_img_gen_params == nullptr) {
         return nullptr;
@@ -3234,7 +3396,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
         }
         return nullptr;
     }
-    if (sd_ctx->sd->free_params_immediately) {
+    if (sd_ctx->sd->free_params_immediately && !request.hires.enabled) {
         sd_ctx->sd->diffusion_model->free_params_buffer();
     }
     int64_t denoise_end = ggml_time_ms();
@@ -3242,6 +3404,131 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
              final_latents.size(),
              (denoise_end - denoise_start) * 1.0f / 1000);
 
+    if (request.hires.enabled && request.hires.target_width > 0) {
+        LOG_INFO("hires fix: upscaling to %dx%d", request.hires.target_width, request.hires.target_height);
+
+        std::unique_ptr<UpscalerGGML> hires_upscaler;
+        if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
+            LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path);
+            hires_upscaler = std::make_unique<UpscalerGGML>(sd_ctx->sd->n_threads,
+                                                            false,
+                                                            request.hires.upscale_tile_size);
+            if (!hires_upscaler->load_from_file(request.hires.model_path,
+                                                sd_ctx->sd->offload_params_to_cpu,
+                                                sd_ctx->sd->n_threads)) {
+                LOG_ERROR("load hires model upscaler failed");
+                if (sd_ctx->sd->free_params_immediately) {
+                    sd_ctx->sd->diffusion_model->free_params_buffer();
+                }
+                return nullptr;
+            }
+        }
+
+        int hires_steps = request.hires.steps > 0 ? request.hires.steps : plan.sample_steps;
+
+        // sd-webui behavior: scale up total steps so trimming by denoising_strength yields exactly hires_steps effective steps,
+        // unlike img2img which trims from a fixed step count
+        hires_steps = static_cast<int>(hires_steps / request.hires.denoising_strength);
+
+        std::vector<float> hires_sigmas = sd_ctx->sd->denoiser->get_sigmas(
+            hires_steps,
+            sd_ctx->sd->get_image_seq_len(request.hires.target_height, request.hires.target_width),
+            sd_img_gen_params->sample_params.scheduler,
+            sd_ctx->sd->version);
+
+        size_t t_enc = static_cast<size_t>(hires_steps * request.hires.denoising_strength);
+        if (t_enc >= static_cast<size_t>(hires_steps)) {
+            t_enc = static_cast<size_t>(hires_steps) - 1;
+        }
+        std::vector<float> hires_sigma_sched(hires_sigmas.begin() + hires_steps - static_cast<int>(t_enc) - 1,
+                                             hires_sigmas.end());
+        LOG_INFO("hires fix: %d steps, denoising_strength=%.2f, sigma_sched_size=%zu",
+                 hires_steps,
+                 request.hires.denoising_strength,
+                 hires_sigma_sched.size());
+
+        std::vector<sd::Tensor<float>> hires_final_latents;
+        int64_t hires_denoise_start = ggml_time_ms();
+        for (int b = 0; b < (int)final_latents.size(); b++) {
+            int64_t cur_seed = request.seed + b;
+            sd_ctx->sd->rng->manual_seed(cur_seed);
+            sd_ctx->sd->sampler_rng->manual_seed(cur_seed);
+
+            sd::Tensor<float> upscaled = upscale_hires_latent(sd_ctx,
+                                                              final_latents[b],
+                                                              request,
+                                                              hires_upscaler.get());
+            if (upscaled.empty()) {
+                if (sd_ctx->sd->free_params_immediately) {
+                    sd_ctx->sd->diffusion_model->free_params_buffer();
+                }
+                return nullptr;
+            }
+
+            sd::Tensor<float> noise = sd::randn_like<float>(upscaled, sd_ctx->sd->rng);
+
+            sd::Tensor<float> hires_denoise_mask;
+            if (!latents.denoise_mask.empty()) {
+                std::vector<int64_t> mask_shape = latents.denoise_mask.shape();
+                mask_shape[0]                   = upscaled.shape()[0];
+                mask_shape[1]                   = upscaled.shape()[1];
+                hires_denoise_mask              = sd::ops::interpolate(latents.denoise_mask,
+                                                                       mask_shape,
+                                                                       sd::ops::InterpolateMode::NearestMax);
+            }
+
+            int64_t hires_sample_start = ggml_time_ms();
+            sd::Tensor<float> x_0      = sd_ctx->sd->sample(sd_ctx->sd->diffusion_model,
+                                                            true,
+                                                            upscaled,
+                                                            std::move(noise),
+                                                            embeds.cond,
+                                                            embeds.uncond,
+                                                            embeds.img_cond,
+                                                            embeds.id_cond,
+                                                            latents.control_image,
+                                                            request.control_strength,
+                                                            request.guidance,
+                                                            plan.eta,
+                                                            request.shifted_timestep,
+                                                            plan.sample_method,
+                                                            sd_ctx->sd->is_flow_denoiser(),
+                                                            hires_sigma_sched,
+                                                            plan.start_merge_step,
+                                                            latents.ref_latents,
+                                                            request.increase_ref_index,
+                                                            hires_denoise_mask,
+                                                            sd::Tensor<float>(),
+                                                            1.f,
+                                                            request.cache_params);
+            int64_t hires_sample_end   = ggml_time_ms();
+            if (!x_0.empty()) {
+                LOG_INFO("hires sampling %d/%d completed, taking %.2fs",
+                         b + 1,
+                         (int)final_latents.size(),
+                         (hires_sample_end - hires_sample_start) * 1.0f / 1000);
+                hires_final_latents.push_back(std::move(x_0));
+                continue;
+            }
+
+            LOG_ERROR("hires sampling for image %d/%d failed after %.2fs",
+                      b + 1,
+                      (int)final_latents.size(),
+                      (hires_sample_end - hires_sample_start) * 1.0f / 1000);
+            if (sd_ctx->sd->free_params_immediately) {
+                sd_ctx->sd->diffusion_model->free_params_buffer();
+            }
+            return nullptr;
+        }
+        if (sd_ctx->sd->free_params_immediately) {
+            sd_ctx->sd->diffusion_model->free_params_buffer();
+        }
+        int64_t hires_denoise_end = ggml_time_ms();
+        LOG_INFO("hires fix completed, taking %.2fs", (hires_denoise_end - hires_denoise_start) * 1.0f / 1000);
+
+        final_latents = std::move(hires_final_latents);
+    }
+
     auto result = decode_image_outputs(sd_ctx, request, final_latents);
     if (result == nullptr) {
         return nullptr;

src/upscaler.cpp

@@ -1,125 +1,115 @@
-#include "esrgan.hpp"
+#include "upscaler.h"
 #include "ggml_extend.hpp"
 #include "model.h"
 #include "stable-diffusion.h"
 #include "util.h"
 
-struct UpscalerGGML {
-    ggml_backend_t backend    = nullptr;  // general backend
-    ggml_type model_data_type = GGML_TYPE_F16;
-    std::shared_ptr<ESRGAN> esrgan_upscaler;
-    std::string esrgan_path;
-    int n_threads;
-    bool direct   = false;
-    int tile_size = 128;
+UpscalerGGML::UpscalerGGML(int n_threads,
+                           bool direct,
+                           int tile_size)
+    : n_threads(n_threads),
+      direct(direct),
+      tile_size(tile_size) {
+}
 
-    UpscalerGGML(int n_threads,
-                 bool direct   = false,
-                 int tile_size = 128)
-        : n_threads(n_threads),
-          direct(direct),
-          tile_size(tile_size) {
-    }
-
-    bool load_from_file(const std::string& esrgan_path,
-                        bool offload_params_to_cpu,
-                        int n_threads) {
-        ggml_log_set(ggml_log_callback_default, nullptr);
+bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
+                                  bool offload_params_to_cpu,
+                                  int n_threads) {
+    ggml_log_set(ggml_log_callback_default, nullptr);
 #ifdef SD_USE_CUDA
-        LOG_DEBUG("Using CUDA backend");
-        backend = ggml_backend_cuda_init(0);
+    LOG_DEBUG("Using CUDA backend");
+    backend = ggml_backend_cuda_init(0);
 #endif
 #ifdef SD_USE_METAL
-        LOG_DEBUG("Using Metal backend");
-        backend = ggml_backend_metal_init();
+    LOG_DEBUG("Using Metal backend");
+    backend = ggml_backend_metal_init();
 #endif
 #ifdef SD_USE_VULKAN
-        LOG_DEBUG("Using Vulkan backend");
-        backend = ggml_backend_vk_init(0);
+    LOG_DEBUG("Using Vulkan backend");
+    backend = ggml_backend_vk_init(0);
 #endif
 #ifdef SD_USE_OPENCL
-        LOG_DEBUG("Using OpenCL backend");
-        backend = ggml_backend_opencl_init();
+    LOG_DEBUG("Using OpenCL backend");
+    backend = ggml_backend_opencl_init();
 #endif
 #ifdef SD_USE_SYCL
-        LOG_DEBUG("Using SYCL backend");
-        backend = ggml_backend_sycl_init(0);
+    LOG_DEBUG("Using SYCL backend");
+    backend = ggml_backend_sycl_init(0);
 #endif
-        ModelLoader model_loader;
-        if (!model_loader.init_from_file_and_convert_name(esrgan_path)) {
-            LOG_ERROR("init model loader from file failed: '%s'", esrgan_path.c_str());
-        }
-        model_loader.set_wtype_override(model_data_type);
-        if (!backend) {
-            LOG_DEBUG("Using CPU backend");
-            backend = ggml_backend_cpu_init();
-        }
-        LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type));
-        esrgan_upscaler = std::make_shared<ESRGAN>(backend, offload_params_to_cpu, tile_size, model_loader.get_tensor_storage_map());
-        if (direct) {
-            esrgan_upscaler->set_conv2d_direct_enabled(true);
-        }
-        if (!esrgan_upscaler->load_from_file(esrgan_path, n_threads)) {
-            return false;
-        }
-        return true;
+    ModelLoader model_loader;
+    if (!model_loader.init_from_file_and_convert_name(esrgan_path)) {
+        LOG_ERROR("init model loader from file failed: '%s'", esrgan_path.c_str());
     }
-
-    sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor) {
-        sd::Tensor<float> upscaled;
-        if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) {
-            upscaled = esrgan_upscaler->compute(n_threads, input_tensor);
-        } else {
-            auto on_processing = [&](const sd::Tensor<float>& input_tile) -> sd::Tensor<float> {
-                auto output_tile = esrgan_upscaler->compute(n_threads, input_tile);
-                if (output_tile.empty()) {
-                    LOG_ERROR("esrgan compute failed while processing a tile");
-                    return {};
-                }
-                return output_tile;
-            };
-
-            upscaled = process_tiles_2d(input_tensor,
-                                        static_cast<int>(input_tensor.shape()[0] * esrgan_upscaler->scale),
-                                        static_cast<int>(input_tensor.shape()[1] * esrgan_upscaler->scale),
-                                        esrgan_upscaler->scale,
-                                        tile_size,
-                                        tile_size,
-                                        0.25f,
-                                        false,
-                                        false,
-                                        on_processing);
-        }
-        esrgan_upscaler->free_compute_buffer();
-        if (upscaled.empty()) {
-            LOG_ERROR("esrgan compute failed");
-            return {};
-        }
-        return upscaled;
+    model_loader.set_wtype_override(model_data_type);
+    if (!backend) {
+        LOG_DEBUG("Using CPU backend");
+        backend = ggml_backend_cpu_init();
     }
+    LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type));
+    esrgan_upscaler = std::make_shared<ESRGAN>(backend, offload_params_to_cpu, tile_size, model_loader.get_tensor_storage_map());
+    if (direct) {
+        esrgan_upscaler->set_conv2d_direct_enabled(true);
+    }
+    if (!esrgan_upscaler->load_from_file(esrgan_path, n_threads)) {
+        return false;
+    }
+    return true;
+}
 
-    sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor) {
-        // upscale_factor, unused for RealESRGAN_x4plus_anime_6B.pth
-        sd_image_t upscaled_image = {0, 0, 0, nullptr};
-        int output_width          = (int)input_image.width * esrgan_upscaler->scale;
-        int output_height         = (int)input_image.height * esrgan_upscaler->scale;
-        LOG_INFO("upscaling from (%i x %i) to (%i x %i)",
-                 input_image.width, input_image.height, output_width, output_height);
+sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_tensor) {
+    sd::Tensor<float> upscaled;
+    if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) {
+        upscaled = esrgan_upscaler->compute(n_threads, input_tensor);
+    } else {
+        auto on_processing = [&](const sd::Tensor<float>& input_tile) -> sd::Tensor<float> {
+            auto output_tile = esrgan_upscaler->compute(n_threads, input_tile);
+            if (output_tile.empty()) {
+                LOG_ERROR("esrgan compute failed while processing a tile");
+                return {};
+            }
+            return output_tile;
+        };
 
-        sd::Tensor<float> input_tensor = sd_image_to_tensor(input_image);
-        sd::Tensor<float> upscaled;
-        int64_t t0 = ggml_time_ms();
-        upscaled   = upscale_tensor(input_tensor);
-        if (upscaled.empty()) {
-            return upscaled_image;
-        }
-        sd_image_t upscaled_data = tensor_to_sd_image(upscaled);
-        int64_t t3               = ggml_time_ms();
-        LOG_INFO("input_image_tensor upscaled, taking %.2fs", (t3 - t0) / 1000.0f);
-        upscaled_image = upscaled_data;
+        upscaled = process_tiles_2d(input_tensor,
+                                    static_cast<int>(input_tensor.shape()[0] * esrgan_upscaler->scale),
+                                    static_cast<int>(input_tensor.shape()[1] * esrgan_upscaler->scale),
+                                    esrgan_upscaler->scale,
+                                    tile_size,
+                                    tile_size,
+                                    0.25f,
+                                    false,
+                                    false,
+                                    on_processing);
+    }
+    esrgan_upscaler->free_compute_buffer();
+    if (upscaled.empty()) {
+        LOG_ERROR("esrgan compute failed");
+        return {};
+    }
+    return upscaled;
+}
+
+sd_image_t UpscalerGGML::upscale(sd_image_t input_image, uint32_t upscale_factor) {
+    // upscale_factor, unused for RealESRGAN_x4plus_anime_6B.pth
+    sd_image_t upscaled_image = {0, 0, 0, nullptr};
+    int output_width          = (int)input_image.width * esrgan_upscaler->scale;
+    int output_height         = (int)input_image.height * esrgan_upscaler->scale;
+    LOG_INFO("upscaling from (%i x %i) to (%i x %i)",
+             input_image.width, input_image.height, output_width, output_height);
+
+    sd::Tensor<float> input_tensor = sd_image_to_tensor(input_image);
+    sd::Tensor<float> upscaled;
+    int64_t t0 = ggml_time_ms();
+    upscaled   = upscale_tensor(input_tensor);
+    if (upscaled.empty()) {
         return upscaled_image;
     }
-};
+    sd_image_t upscaled_data = tensor_to_sd_image(upscaled);
+    int64_t t3               = ggml_time_ms();
+    LOG_INFO("input_image_tensor upscaled, taking %.2fs", (t3 - t0) / 1000.0f);
+    upscaled_image = upscaled_data;
+    return upscaled_image;
+}
 
 struct upscaler_ctx_t {
     UpscalerGGML* upscaler = nullptr;

src/upscaler.h

@@ -0,0 +1,31 @@
+#ifndef __SD_UPSCALER_H__
+#define __SD_UPSCALER_H__
+
+#include "esrgan.hpp"
+#include "stable-diffusion.h"
+#include "tensor.hpp"
+
+#include <memory>
+#include <string>
+
+struct UpscalerGGML {
+    ggml_backend_t backend    = nullptr;  // general backend
+    ggml_type model_data_type = GGML_TYPE_F16;
+    std::shared_ptr<ESRGAN> esrgan_upscaler;
+    std::string esrgan_path;
+    int n_threads;
+    bool direct   = false;
+    int tile_size = 128;
+
+    UpscalerGGML(int n_threads,
+                 bool direct   = false,
+                 int tile_size = 128);
+
+    bool load_from_file(const std::string& esrgan_path,
+                        bool offload_params_to_cpu,
+                        int n_threads);
+    sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor);
+    sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor);
+};
+
+#endif  // __SD_UPSCALER_H__