1 files changed, 76 insertions, 13 deletions

diff --git a/main.c b/main.c
index 3564425..873afdb 100644
--- a/main.c
+++ b/main.c
@@ -8,8 +8,10 @@
 #include "mtmd-helper.h"
 #include "llama.h"
 
-#define N_CTX        1024
+#define N_CTX        4096
+#define N_BATCH      2048
 #define N_THREADS       4
+#define N_PREDICT     256
 
 #define MODEL_PATH   "model.gguf"
 #define MMPROJ_PATH  "mmproj.gguf"
@@ -55,6 +57,7 @@ static char * build_prompt(struct llama_model *model, int n_images)
 	if (n < 0) {
 		fprintf(stderr, "Chat template size check failed\n");
 		free(content);
+		return NULL;
 	}
 	
 	buf_len = n + 1;
@@ -63,6 +66,7 @@ static char * build_prompt(struct llama_model *model, int n_images)
 	if (n < 0) {
 		fprintf(stderr, "Failed to apply chat template\n");
 		free(content);
+		return NULL;
 	}
 
 	buf[n] = '\0';
@@ -71,7 +75,7 @@ static char * build_prompt(struct llama_model *model, int n_images)
 
 int main(int argc, char *argv[])
 {
-	int n_images;
+	int i, n_images, rc;
 	const char *prompt;
 	const struct llama_vocab *vocab;
 
@@ -82,18 +86,22 @@ int main(int argc, char *argv[])
 	struct llama_sampler *smpl = NULL;
 	struct llama_sampler_chain_params sparams;
 
-	int i;
-	mtmd_input_chunks *chunks;
+	int n_bitmaps = 0;
+	mtmd_input_chunks *chunks = NULL;
 	mtmd_bitmap **bitmaps = NULL;
 	struct mtmd_context *mctx = NULL;
 	struct mtmd_context_params mtparams;
 	struct mtmd_helper_init_opt init_opt;
 	struct mtmd_helper_bitmap_wrapper wrapper;
 
+	llama_pos n_past;
+
+	rc = 1;
+
 	n_images = argc - 1;
 	if (n_images < 1) {
 		fprintf(stderr, "Usage: %s image1 [image2 ...]\n", argv[0]);
-		return 1;
+		return rc;
 	}
 
 	llama_backend_init();
@@ -157,16 +165,16 @@ int main(int argc, char *argv[])
 
 	init_opt = mtmd_helper_init_opt_default();
 	bitmaps = CALLOC(n_images, sizeof(*bitmaps));
-	for (i = 0; i < n_images; i++) {
+	for (n_bitmaps = 0; n_bitmaps < n_images; n_bitmaps++) {
 		wrapper = mtmd_helper_bitmap_init_from_file(
 		                 mctx, 
-		                 argv[i + 1], 
+		                 argv[n_bitmaps + 1], 
 		                 false, 
 		                 init_opt);
 
-		bitmaps[i] = wrapper.bitmap;
-		if (!bitmaps[i]) {
-			fprintf(stderr, "Failed to load image: %s\n", argv[i + 1]);
+		bitmaps[n_bitmaps] = wrapper.bitmap;
+		if (!bitmaps[n_bitmaps]) {
+			fprintf(stderr, "Failed to load image: %s\n", argv[n_bitmaps + 1]);
 			goto cleanup;
 		}
 	}
@@ -177,21 +185,76 @@ int main(int argc, char *argv[])
 		goto cleanup;
 	}
 
-	printf("Init success\n");
+	{
+		mtmd_input_text text;
+
+		memset(&text, 0, sizeof(text));
+		text.text = prompt;
+		text.text_len = strlen(prompt);
+		text.add_special = true;   /* use model's BOS/EOS tokens */
+		text.parse_special = true; /* treat BOS/EOS tokens as special (Qwen/Llama) */
+		
+		chunks = mtmd_input_chunks_init();
+		if (mtmd_tokenize(mctx, chunks, &text, 
+			(const mtmd_bitmap **)bitmaps, n_images) != 0) {
+			fprintf(stderr, "mtmd_tokenize error\n");
+			goto cleanup;
+		}
+	}
+
+	/* Encode images + decode prompt; logits for the last one */
+	if (mtmd_helper_eval_chunks(mctx, ctx, chunks, 0, 0, N_BATCH, 
+		true, &n_past) != 0) {
+		fprintf(stderr, "Failed to evaluate prompt/images\n");
+		goto cleanup;
+	}
+
+	/* Generation loop */
+	{
+		int t, n;
+		llama_token token;
+
+		rc = 0;
+		for (t = 0; t < N_PREDICT && n_past < N_CTX; t++) {
+			token = llama_sampler_sample(smpl, ctx, -1);
+			if (llama_vocab_is_eog(vocab, token))
+				break;
+
+			n = 0;
+			char buf[256];
+			n = llama_token_to_piece(vocab, token, buf, sizeof(buf), 0, true);
+			if (n > 0)
+				fwrite(buf, 1, n, stdout);
+			fflush(stdout);
+
+			if (llama_decode(ctx, llama_batch_get_one(&token, 1)) != 0) {
+				rc = 1;
+				fprintf(stderr, "llama_decode error\n");
+				break;
+			}
+			
+			n_past++;
+		}
+		printf("\n");
+	}
 
 cleanup:
 	if (bitmaps) {
-		for (i = 0; bitmaps[i]; i++)
+		for (i = 0; i < n_bitmaps && bitmaps[i]; i++)
 			mtmd_bitmap_free(bitmaps[i]);
 		free(bitmaps);
 	}
+
+	if (chunks)
+		mtmd_input_chunks_free(chunks);
 	if (smpl)
 		llama_sampler_free(smpl);
 	if (ctx)
 		llama_free(ctx);
 	if (model)
 		llama_model_free(model);
+
 	llama_backend_free();
 
-	return 0;
+	return rc;
 }