index.html

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Image Caption Prototype</title>
</head>

<body>
    <h1>Image Caption Prototype</h1>

    <input type="file" id="imageInput" accept="image/*">

    <br><br>

    <img
        id="preview"
        style="max-width: 500px; display: none;"
        alt="Selected image"
    >

    <br><br>

    <button id="captionButton" disabled>
        Generate caption
    </button>

    <p id="status"></p>

    <h2>Raw caption</h2>
    <p id="rawCaption"></p>

    <h2>Refined caption</h2>
    <p id="caption"></p>

    <details>
        <summary>Debug</summary>
        <pre id="debug" style="white-space: pre-wrap;"></pre>
    </details>

    <script type="module">
        import {
            Florence2ForConditionalGeneration,
            AutoProcessor,
            AutoTokenizer,
            RawImage,
            pipeline
        } from "https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.8.1";

        const imageInput = document.getElementById("imageInput");
        const preview = document.getElementById("preview");
        const captionButton = document.getElementById("captionButton");
        const status = document.getElementById("status");
        const rawCaptionEl = document.getElementById("rawCaption");
        const caption = document.getElementById("caption");
        const debugEl = document.getElementById("debug");

        const MODEL_ID = "onnx-community/Florence-2-base-ft";
        // Text2text model: much more robust for rewrite tasks than a tiny chat LLM.
        const REFINER_ID = "Xenova/LaMini-Flan-T5-248M";

        let selectedFile = null;
        let model = null;
        let processor = null;
        let tokenizer = null;
        let refiner = null;

        function log(...args) {
            console.log(...args);
            debugEl.textContent += args.map(a =>
                typeof a === "string" ? a : JSON.stringify(a)
            ).join(" ") + "\n";
        }

        // ---------- filename helpers ----------
        function cleanFilename(name) {
            let base = name.replace(/\.[^.]+$/, "");
            base = base.replace(/[_-]+/g, " ");
            base = base.replace(
                /\b(img|dsc|photo|image|picture|pic|file|screenshot|final|copy|edit|v\d+)\b/gi,
                " "
            );
            base = base.replace(/\b\d+\b/g, " ");
            base = base.replace(/\s+/g, " ").trim();
            return base;
        }

        // ---------- model loading ----------
        async function loadCaptionModel() {
            if (model) return;
            status.textContent = "Loading Florence-2 model...";
            model = await Florence2ForConditionalGeneration.from_pretrained(
                MODEL_ID,
                {
                    dtype: {
                        embed_tokens: "fp16",
                        vision_encoder: "fp16",
                        encoder_model: "q4",
                        decoder_model_merged: "q4",
                    },
                    device: "webgpu",
                }
            );
            processor = await AutoProcessor.from_pretrained(MODEL_ID);
            tokenizer = await AutoTokenizer.from_pretrained(MODEL_ID);
        }

        async function loadRefiner() {
            if (refiner) return;
            status.textContent = "Loading caption refiner...";
            refiner = await pipeline("text2text-generation", REFINER_ID, {
                dtype: "q8",
                device: "wasm", // safer than webgpu for this model
            });
        }

        // ---------- sanity checks ----------
        function looksBroken(text, rawCaption) {
            if (!text) return true;
            const t = text.trim();
            if (t.length < 5) return true;
            if (t.length > 220) return true;                       // runaway generation
            if (/\b(\w+)\s+\1\s+\1\b/i.test(t)) return true;       // "the the the"
            // If it contains words like "Raw caption:" or "Filename hints:", it echoed the prompt
            if (/raw caption|filename hints|rewritten caption/i.test(t)) return true;
            // If it is basically the raw caption, prefer the raw caption (no loss)
            const norm = s => s.toLowerCase().replace(/[^a-z0-9 ]/g, "").trim();
            if (norm(t) === norm(rawCaption)) return true;
            return false;
        }

        // ---------- refinement ----------
        async function refine(rawCaption, filename) {
            const hints = cleanFilename(filename);

            // Flan-T5 wants a plain string prompt, and explicit instructions
            // about the output format work best.
            const prompt = hints
                            ? `Condense the caption into one sentence. Then add " — " and a short note ` +
                              `identifying the source or subject using these keywords: ${hints}. ` +
                              `Caption: ${rawCaption} ` +
                              `Output:`
                            : `Condense the caption into one sentence. Caption: ${rawCaption} Output:`;

            log("[refine prompt]", prompt);

            const out = await refiner(prompt, {
                max_new_tokens: 60,
                do_sample: false,
                num_beams: 4,
                repetition_penalty: 1.3,
                no_repeat_ngram_size: 3,
            });

            log("[refine raw output]", JSON.stringify(out));

            const text = (out?.[0]?.generated_text || "").trim();
            return text;
        }

        // ---------- file input ----------
        imageInput.addEventListener("change", (event) => {
            selectedFile = event.target.files[0];
            if (!selectedFile) {
                captionButton.disabled = true;
                preview.style.display = "none";
                return;
            }
            preview.src = URL.createObjectURL(selectedFile);
            preview.style.display = "block";
            captionButton.disabled = false;
            caption.textContent = "";
            rawCaptionEl.textContent = "";
            status.textContent = "";
            debugEl.textContent = "";
        });

        // ---------- main ----------
        captionButton.addEventListener("click", async () => {
            try {
                captionButton.disabled = true;
                caption.textContent = "";
                rawCaptionEl.textContent = "";
                debugEl.textContent = "";

                await loadCaptionModel();
                status.textContent = "Preparing image...";
                const image = await RawImage.fromBlob(selectedFile);
                const task = "<MORE_DETAILED_CAPTION>";
                const visionInputs = await processor(image);
                const prompts = processor.construct_prompts(task);
                const textInputs = tokenizer(prompts);

                status.textContent = "Generating caption...";
                const generatedIds = await model.generate({
                    ...textInputs,
                    ...visionInputs,
                    max_new_tokens: 100,
                });
                const generatedText = tokenizer.batch_decode(
                    generatedIds,
                    { skip_special_tokens: false }
                )[0];
                const result = processor.post_process_generation(
                    generatedText,
                    task,
                    image.size
                );
                const rawCaption = result[task] || "";

                log("[raw caption]", rawCaption);
                log("[filename]", selectedFile.name);
                log("[cleaned filename]", cleanFilename(selectedFile.name));

                rawCaptionEl.textContent = rawCaption || "(no caption)";

                await loadRefiner();
                status.textContent = "Refining caption...";
                const refined = await refine(rawCaption, selectedFile.name);

                log("[refined]", refined);

                if (looksBroken(refined, rawCaption)) {
                    log("[fallback] refined output rejected, using raw caption");
                    caption.textContent = rawCaption || "No caption generated.";
                    status.textContent = "Done (kept raw caption).";
                } else {
                    caption.textContent = refined;
                    status.textContent = "Done.";
                }
            } catch (error) {
                console.error(error);
                status.textContent = "Error: " + (error.message || error);
                log("[error]", String(error));
            } finally {
                captionButton.disabled = false;
            }
        });
    </script>
</body>
</html>