Text & image generation

/* Ask a vision-language model about an image and stream the answer, then
 * ask a follow-up question in the same conversation.
 *
 * usage: chat MODEL.slab IMAGE LICENSE_FILE
 */

#include <stdio.h>
#include <string.h>

#include "slab.h"

#define S(str) (uint8_t *)(str), strlen(str)

static int read_file(const char *path, char *buf, size_t cap, size_t *len) {
    FILE *f = fopen(path, "rb");
    if (!f) return 0;
    *len = fread(buf, 1, cap, f);
    fclose(f);
    while (*len && (buf[*len - 1] == '\n' || buf[*len - 1] == '\r')) (*len)--;
    return 1;
}

/* start() reads the prompt; next() produces one piece of text per call.
 * There are no callbacks: the app pulls text in a loop. */
static int stream(uint8_t *session, uint8_t *in, uint8_t *out) {
    if (!slab_session_start(session, in, out)) return 0;
    int32_t r;
    while ((r = slab_session_next(session, out)) == 1)
        printf("%.*s", (int)slab_output_text_len(out), (char *)slab_output_text_ptr(out));
    printf("\n");
    return r == 0;
}

int main(int argc, char **argv) {
    if (argc < 4) { fprintf(stderr, "usage: %s MODEL.slab IMAGE LICENSE_FILE\n", argv[0]); return 2; }

    uint8_t *cfg = slab_config_new();
    char token[4096]; size_t token_len;
    if (!read_file(argv[3], token, sizeof token, &token_len) ||
        !slab_config_set_runtime_license(cfg, (uint8_t *)token, token_len)) return 1;

    uint8_t *session = slab_session_new(cfg);
    slab_session_set_context_length(session, 2048);
    if (!slab_session_open(session, S(argv[1]))) {
        fprintf(stderr, "%.*s\n", (int)slab_session_error_message_len(session), (char *)slab_session_error_message_ptr(session));
        return 1;
    }

    uint8_t *photo = slab_frame_new();
    if (!slab_frame_load_file(photo, S(argv[2]))) return 1;

    uint8_t *in = slab_input_new();
    uint8_t *out = slab_output_new();
    slab_input_set_max_tokens(in, 64);
    slab_input_set_temperature(in, 0.0f);
    slab_input_set_image_tiled(in, 1, 1536);   /* large images are split into tiles, not shrunk */

    /* Turn 1: a question about the image. */
    slab_input_add_image_message(in, S("user"), photo, S("What is this document?"));
    if (!stream(session, in, out)) return 1;

    /* Turn 2: the answer and a new question. The shared beginning is reused. */
    char answer[4096];
    snprintf(answer, sizeof answer, "%.*s", (int)slab_output_full_text_len(out), (char *)slab_output_full_text_ptr(out));
    slab_input_add_message(in, S("assistant"), S(answer));
    slab_input_add_message(in, S("user"), S("What is the total amount?"));
    if (!stream(session, in, out)) return 1;
    printf("prompt %d tokens (%d reused), generated %d\n", slab_output_prompt_tokens(out),
           slab_output_cached_prompt_tokens(out), slab_output_generated_tokens(out));

    slab_session_reset(session);   /* forget the conversation before an unrelated one */
    slab_output_free(out);
    slab_input_free(in);
    slab_frame_free(photo);
    slab_session_free(session);
    slab_config_free(cfg);
    return 0;
}