model_loader.cpp

#include "model_loader.h"

#include <algorithm>
#include <chrono>
#include <cctype>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <sstream>
#include <unordered_map>

#include "pickle_parser.h"
#include "zip_archive.h"

namespace gd {
namespace {

double elapsed_ms(const std::chrono::steady_clock::time_point start,
                  const std::chrono::steady_clock::time_point end) {
    const auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);
    return static_cast<double>(duration.count()) / 1000.0;
}

std::string to_lower(std::string value) {
    std::transform(value.begin(), value.end(), value.begin(), [](unsigned char ch) {
        return static_cast<char>(std::tolower(ch));
    });
    return value;
}

std::string format_bytes(uint64_t value) {
    static const char * units[] = {"B", "KiB", "MiB", "GiB"};
    double size = static_cast<double>(value);
    size_t unit_index = 0;
    while (size >= 1024.0 && unit_index + 1 < std::size(units)) {
        size /= 1024.0;
        ++unit_index;
    }

    std::ostringstream output;
    output << std::fixed << std::setprecision(unit_index == 0 ? 0 : 2) << size << ' ' << units[unit_index];
    return output.str();
}

std::string join_shape(const std::vector<int64_t> & dims) {
    std::ostringstream output;
    output << '[';
    for (size_t index = 0; index < dims.size(); ++index) {
        if (index > 0) {
            output << ", ";
        }
        output << dims[index];
    }
    output << ']';
    return output.str();
}

std::string json_escape(const std::string & value) {
    std::ostringstream out;
    for (const char ch : value) {
        switch (ch) {
            case '\\': out << "\\\\"; break;
            case '"': out << "\\\""; break;
            case '\n': out << "\\n"; break;
            case '\r': out << "\\r"; break;
            case '\t': out << "\\t"; break;
            default:
                out << ch;
                break;
        }
    }
    return out.str();
}

std::string classify_model(const fs::path & model_path, const std::vector<TensorDescriptor> & tensors) {
    const std::string filename = to_lower(model_path.filename().string());
    bool has_guided_diffusion_unet = false;
    bool has_latent_diffusion_prefix = false;

    for (const TensorDescriptor & tensor : tensors) {
        const std::string & name = tensor.name;
        if (name == "input_blocks.0.0.weight" || name == "time_embed.0.weight" || name == "out.2.weight") {
            has_guided_diffusion_unet = true;
        }
        if (name.find("model.diffusion_model.") != std::string::npos ||
            name.find("first_stage_model.") != std::string::npos ||
            name.find("cond_stage_model.") != std::string::npos) {
            has_latent_diffusion_prefix = true;
        }
    }

    if (filename.find("diffusion_uncond") != std::string::npos ||
        (has_guided_diffusion_unet && !has_latent_diffusion_prefix)) {
        if (filename.find("512x512") != std::string::npos) {
            return "OpenAI guided-diffusion UNet checkpoint (unconditional 512x512 model)";
        }
        return "OpenAI guided-diffusion UNet checkpoint";
    }

    if (has_latent_diffusion_prefix) {
        return "Latent diffusion / Stable Diffusion style checkpoint";
    }

    return "Unknown PyTorch checkpoint";
}

bool load_tensor_data(std::ifstream & input,
                      const std::unordered_map<std::string, ZipEntry> & storage_entries,
                      const TensorDescriptor & descriptor,
                      ggml_tensor * tensor) {
    const auto found = storage_entries.find("archive/data/" + descriptor.storage_key);
    if (found == storage_entries.end()) {
        std::cerr << "Missing storage entry for tensor: " << descriptor.name << '\n';
        return false;
    }

    const size_t element_size = static_cast<size_t>(ggml_type_size(descriptor.type));
    const uint64_t tensor_bytes = static_cast<uint64_t>(ggml_nbytes(tensor));
    const uint64_t start = found->second.data_offset + descriptor.storage_offset * element_size;
    const uint64_t end = start + tensor_bytes;
    const uint64_t storage_end = found->second.data_offset + found->second.uncompressed_size;
    if (end > storage_end) {
        std::cerr << "Storage range overflow for tensor: " << descriptor.name << '\n';
        return false;
    }

    input.clear();
    input.seekg(static_cast<std::streamoff>(start), std::ios::beg);
    input.read(reinterpret_cast<char *>(tensor->data), static_cast<std::streamsize>(tensor_bytes));
    if (!input) {
        std::cerr << "Failed to read tensor bytes for: " << descriptor.name << '\n';
        return false;
    }

    return true;
}

}  // namespace

std::optional<LoadedModel> load_model_archive(const fs::path & model_path) {
    const auto t0 = std::chrono::steady_clock::now();

    const auto t_zip_start = std::chrono::steady_clock::now();
    const std::optional<std::vector<ZipEntry>> entries = read_zip_entries(model_path);
    if (!entries.has_value()) {
        return std::nullopt;
    }
    const auto t_zip_end = std::chrono::steady_clock::now();

    LoadedModel model;
    model.path = model_path;
    model.archive_entries = *entries;

    const auto pickle_entry = std::find_if(model.archive_entries.begin(), model.archive_entries.end(), [](const ZipEntry & entry) {
        return entry.name == "archive/data.pkl";
    });
    if (pickle_entry == model.archive_entries.end()) {
        std::cerr << "Missing archive/data.pkl in checkpoint: " << model_path << '\n';
        return std::nullopt;
    }

    const std::optional<std::vector<uint8_t>> payload = read_stored_zip_entry(model_path, *pickle_entry);
    if (!payload.has_value()) {
        return std::nullopt;
    }

    const auto t_pickle_start = std::chrono::steady_clock::now();
    std::vector<TensorDescriptor> descriptors;
    try {
        descriptors = parse_tensor_descriptors(*payload);
    } catch (const std::exception & error) {
        std::cerr << "Failed to parse checkpoint pickle: " << error.what() << '\n';
        return std::nullopt;
    }
    const auto t_pickle_end = std::chrono::steady_clock::now();

    if (descriptors.empty()) {
        std::cerr << "Parsed checkpoint metadata but found zero tensors in: " << model_path << '\n';
        std::cerr << "This likely means the current pickle parser does not support this checkpoint layout yet." << '\n';
        return std::nullopt;
    }

    for (const ZipEntry & entry : model.archive_entries) {
        if (entry.name.rfind("archive/data/", 0) == 0) {
            model.storage_bytes += entry.uncompressed_size;
        }
    }

    model.model_kind = classify_model(model_path, descriptors);

    size_t ctx_size = 16 * 1024;
    for (const TensorDescriptor & descriptor : descriptors) {
        int64_t elements = 1;
        for (int64_t dim : descriptor.shape) {
            elements *= dim;
        }
        ctx_size += ggml_tensor_overhead();
        ctx_size += static_cast<size_t>(elements) * static_cast<size_t>(ggml_type_size(descriptor.type));
    }

    ggml_init_params params = {
        /*.mem_size   =*/ ctx_size,
        /*.mem_buffer =*/ nullptr,
        /*.no_alloc   =*/ false,
    };
    model.ctx = ggml_init(params);
    if (model.ctx == nullptr) {
        std::cerr << "ggml_init failed for model: " << model_path << '\n';
        return std::nullopt;
    }

    std::unordered_map<std::string, ZipEntry> storage_entries;
    for (const ZipEntry & entry : model.archive_entries) {
        storage_entries.emplace(entry.name, entry);
    }

    std::ifstream input(model_path, std::ios::binary);
    if (!input) {
        std::cerr << "Failed to open archive for tensor load: " << model_path << '\n';
        ggml_free(model.ctx);
        return std::nullopt;
    }

    const auto t_load_start = std::chrono::steady_clock::now();
    model.tensors.reserve(descriptors.size());
    for (const TensorDescriptor & descriptor : descriptors) {
        std::vector<int64_t> reversed_shape(descriptor.shape.rbegin(), descriptor.shape.rend());
        ggml_tensor * tensor = ggml_new_tensor(model.ctx,
                                               descriptor.type,
                                               static_cast<int>(reversed_shape.size()),
                                               reversed_shape.data());
        if (tensor == nullptr) {
            std::cerr << "Failed to allocate ggml tensor: " << descriptor.name << '\n';
            ggml_free(model.ctx);
            return std::nullopt;
        }

        ggml_set_name(tensor, descriptor.name.c_str());
        if (!load_tensor_data(input, storage_entries, descriptor, tensor)) {
            ggml_free(model.ctx);
            return std::nullopt;
        }

        model.tensors.push_back(LoadedTensor{descriptor, tensor});
    }
    const auto t_load_end = std::chrono::steady_clock::now();
    const auto t1 = std::chrono::steady_clock::now();

    model.ms_zip_index = elapsed_ms(t_zip_start, t_zip_end);
    model.ms_pickle_parse = elapsed_ms(t_pickle_start, t_pickle_end);
    model.ms_tensor_materialize = elapsed_ms(t_load_start, t_load_end);
    model.ms_total = elapsed_ms(t0, t1);

    return model;
}

std::vector<fs::path> find_model_candidates(const fs::path & models_dir) {
    std::vector<fs::path> models;
    if (!fs::exists(models_dir)) {
        return models;
    }

    for (const fs::directory_entry & entry : fs::directory_iterator(models_dir)) {
        if (!entry.is_regular_file()) {
            continue;
        }

        const std::string extension = to_lower(entry.path().extension().string());
        if (extension == ".pt" || extension == ".ckpt" || extension == ".pth") {
            models.push_back(entry.path());
        }
    }

    std::sort(models.begin(), models.end());
    return models;
}

void print_model_summary(const LoadedModel & model) {
    size_t storage_shards = 0;
    for (const ZipEntry & entry : model.archive_entries) {
        if (entry.name.rfind("archive/data/", 0) == 0) {
            ++storage_shards;
        }
    }

    std::cout << "Model: " << model.path.filename().string() << '\n';
    std::cout << "Path:  " << model.path.string() << '\n';
    std::cout << "Type:  " << model.model_kind << '\n';
    std::cout << "Format: PyTorch zip checkpoint (.pt)" << '\n';
    std::cout << "Archive entries: " << model.archive_entries.size() << '\n';
    std::cout << "Storage shards:  " << storage_shards << '\n';
    std::cout << "Tensor payload:  " << format_bytes(model.storage_bytes) << '\n';
    std::cout << "GGML tensors:    " << model.tensors.size() << '\n';
    std::cout << "GGML context:    " << (model.ctx != nullptr ? "initialized" : "not initialized") << '\n';

    const auto find_tensor = [&](const std::string & name) -> const TensorDescriptor * {
        auto found = std::find_if(model.tensors.begin(), model.tensors.end(), [&](const LoadedTensor & loaded) {
            return loaded.descriptor.name == name;
        });
        return found == model.tensors.end() ? nullptr : &found->descriptor;
    };
    const TensorDescriptor * in_proj = find_tensor("input_blocks.0.0.weight");
    const TensorDescriptor * out_proj = find_tensor("out.2.weight");
    if (in_proj != nullptr && in_proj->shape.size() == 4 && out_proj != nullptr && out_proj->shape.size() == 4) {
        std::cout << "Model shape: in_channels=" << in_proj->shape[1]
                  << ", base_channels=" << in_proj->shape[0]
                  << ", out_channels=" << out_proj->shape[0] << '\n';
        std::cout << "Kernel shape: input=" << join_shape(in_proj->shape)
                  << ", output=" << join_shape(out_proj->shape) << '\n';
    }

    const size_t preview_count = std::min<size_t>(8, model.tensors.size());
    if (preview_count > 0) {
        std::cout << "Sample tensors:" << '\n';
        for (size_t index = 0; index < preview_count; ++index) {
            const TensorDescriptor & tensor = model.tensors[index].descriptor;
            std::cout << "  - " << tensor.name << " | " << ggml_type_name(tensor.type)
                      << " | " << join_shape(tensor.shape) << '\n';
        }
    }

    std::cout << std::fixed << std::setprecision(2);
    std::cout << "Timing (ms): zip_index=" << model.ms_zip_index
              << ", pickle_parse=" << model.ms_pickle_parse
              << ", tensor_load=" << model.ms_tensor_materialize
              << ", total=" << model.ms_total << '\n';

    std::cout << "Instantiation: checkpoint loaded into ggml context" << '\n';
}

bool write_model_manifest_json(const LoadedModel & model, const fs::path & output_path) {
    std::ofstream out(output_path, std::ios::binary);
    if (!out) {
        std::cerr << "Failed to open manifest output path: " << output_path << '\n';
        return false;
    }

    out << "{\n";
    out << "  \"model\": \"" << json_escape(model.path.string()) << "\",\n";
    out << "  \"tensor_count\": " << model.tensors.size() << ",\n";
    out << "  \"tensors\": [\n";
    for (size_t i = 0; i < model.tensors.size(); ++i) {
        const TensorDescriptor & t = model.tensors[i].descriptor;
        out << "    {\"name\": \"" << json_escape(t.name) << "\", \"dtype\": \""
            << ggml_type_name(t.type) << "\", \"shape\": [";
        for (size_t d = 0; d < t.shape.size(); ++d) {
            if (d > 0) {
                out << ", ";
            }
            out << t.shape[d];
        }
        out << "]}";
        if (i + 1 < model.tensors.size()) {
            out << ',';
        }
        out << "\n";
    }
    out << "  ]\n";
    out << "}\n";
    return true;
}

}  // namespace gd