fix: added

.dockerignore

@@ -1,3 +1,5 @@
-node_modules
 .git
 .env
+__pycache__
+*.pyc
+.venv

.gitignore

@@ -1 +1,4 @@
 token.txt
+__pycache__/
+*.pyc
+.venv/

Dockerfile

@@ -1,25 +1,23 @@
-FROM node:22-slim
+FROM python:3.11-slim
 
-# sharp needs libvips
-RUN apt-get update && \
-    apt-get install -y --no-install-recommends libvips-dev && \
-    rm -rf /var/lib/apt/lists/*
+ENV PYTHONDONTWRITEBYTECODE=1
+ENV PYTHONUNBUFFERED=1
 
 WORKDIR /app
 
-COPY package.json ./
-RUN npm install --omit=dev
-
-COPY src/ src/
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends libgl1 libglib2.0-0 && \
+    rm -rf /var/lib/apt/lists/*
 
-# Give node user ownership of everything (including node_modules/.cache)
-RUN chown -R node:node /app
+COPY requirements.txt ./
+RUN pip install --no-cache-dir -r requirements.txt
 
-USER node
+COPY src/ src/
+COPY README.md ./
 
-# Download model at build time so container starts fast
-RUN node -e "import('./src/model.js').then(m => m.loadModel()).then(() => process.exit(0))"
+# Download model weights at build time so cold start is faster.
+RUN python -c "from src.model import load_model; load_model()"
 
 EXPOSE 7860
 
-CMD ["node", "src/server.js"]
+CMD ["uvicorn", "src.server:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -8,6 +8,6 @@ app_port: 7860
 pinned: false
 ---
 
-Image captioning API using FastVLM (ONNX). Open `/docs` for Swagger UI.
+Image captioning API using `microsoft/Florence-2-base` with a Python FastAPI backend. Open `/docs` for Swagger UI.
 
-Speed tuning env vars: `DEFAULT_MAX_TOKENS` (default `64`), `MAX_IMAGE_SIDE` (default `896`), `MAX_MAX_TOKENS` (default `256`).
+Speed tuning env vars: `DEFAULT_MAX_TOKENS` (default `64`), `MAX_IMAGE_SIDE` (default `896`), `MAX_MAX_TOKENS` (default `256`), `MODEL_ID` (default `microsoft/Florence-2-base`).

docker-compose.yml

@@ -8,4 +8,5 @@ services:
       - DEFAULT_MAX_TOKENS=64
       - MAX_IMAGE_SIDE=896
       - MAX_MAX_TOKENS=256
+      - MODEL_ID=microsoft/Florence-2-base
     restart: unless-stopped

package.json

@@ -1,19 +0,0 @@
-{
-  "name": "img3txt",
-  "version": "1.0.0",
-  "description": "Image captioning API using FastVLM ONNX model",
-  "type": "module",
-  "scripts": {
-    "start": "node src/server.js",
-    "dev": "node --watch src/server.js"
-  },
-  "dependencies": {
-    "@huggingface/transformers": "^3.4.1",
-    "fastify": "^5.2.1",
-    "@fastify/multipart": "^9.0.3",
-    "@fastify/swagger": "^9.4.2",
-    "@fastify/swagger-ui": "^5.2.1",
-    "@fastify/cors": "^10.0.2",
-    "sharp": "^0.33.5"
-  }
-}

requirements.txt

@@ -0,0 +1,6 @@
+fastapi==0.116.1
+uvicorn[standard]==0.35.0
+transformers==4.55.4
+torch==2.8.0
+pillow==11.3.0
+python-multipart==0.0.20

src/model.js

@@ -1,106 +0,0 @@
-import {
-  AutoModelForImageTextToText,
-  AutoProcessor,
-  RawImage,
-} from "@huggingface/transformers";
-import sharp from "sharp";
-
-const MODEL_ID = "onnx-community/FastVLM-0.5B-ONNX";
-const DEFAULT_MAX_TOKENS = parseInt(process.env.DEFAULT_MAX_TOKENS || "64", 10);
-const MAX_MAX_TOKENS = parseInt(process.env.MAX_MAX_TOKENS || "256", 10);
-const MAX_IMAGE_SIDE = parseInt(process.env.MAX_IMAGE_SIDE || "896", 10);
-
-let model = null;
-let processor = null;
-
-/** Supported task instructions for FastVLM */
-export const TASKS = {
-  caption: "Describe this image.",
-  detailed_caption: "Describe this image in detail.",
-  more_detailed_caption:
-    "Provide a very detailed description of this image.",
-  ocr: "Extract all readable text from this image.",
-  ocr_with_region:
-    "Extract all readable text and include where it appears in the image.",
-  object_detection: "List the visible objects in this image.",
-  dense_region_caption:
-    "Describe this image region by region with detailed observations.",
-  region_proposal:
-    "Propose important regions in this image and explain what each region contains.",
-};
-
-export async function loadModel() {
-  if (!model) {
-    console.log("Loading FastVLM model...");
-    model = await AutoModelForImageTextToText.from_pretrained(MODEL_ID, {
-      dtype: {
-        embed_tokens: "fp16",
-        vision_encoder: "q4",
-        decoder_model_merged: "q4",
-      },
-    });
-    processor = await AutoProcessor.from_pretrained(MODEL_ID);
-    console.log("Model loaded.");
-  }
-  return { model, processor };
-}
-
-/**
- * Generate text from an image buffer.
- * @param {Buffer} imageBuffer - Raw image bytes
- * @param {string} task - One of the TASKS keys (default: "caption")
- * @param {string|null} textInput - Optional extra text input for the task
- * @param {number} maxTokens - Max new tokens to generate
- * @returns {Promise<object>} Generated result from FastVLM
- */
-export async function generateCaption(
-  imageBuffer,
-  task = "caption",
-  textInput = null,
-  maxTokens = DEFAULT_MAX_TOKENS
-) {
-  const { model: m, processor: p } = await loadModel();
-
-  const safeMaxTokens = Number.isFinite(maxTokens)
-    ? Math.min(Math.max(maxTokens, 8), MAX_MAX_TOKENS)
-    : DEFAULT_MAX_TOKENS;
-
-  // Downscale large uploads to reduce encoder latency.
-  const metadata = await sharp(imageBuffer).metadata();
-  let preparedBuffer = imageBuffer;
-  if (
-    metadata.width &&
-    metadata.height &&
-    (metadata.width > MAX_IMAGE_SIDE || metadata.height > MAX_IMAGE_SIDE)
-  ) {
-    preparedBuffer = await sharp(imageBuffer)
-      .resize({
-        width: MAX_IMAGE_SIDE,
-        height: MAX_IMAGE_SIDE,
-        fit: "inside",
-        withoutEnlargement: true,
-      })
-      .toBuffer();
-  }
-
-  const image = await RawImage.fromBlob(new Blob([preparedBuffer]));
-
-  const baseInstruction = TASKS[task] || TASKS.caption;
-  const instruction = textInput
-    ? `${baseInstruction}\nAdditional instruction: ${textInput}`
-    : baseInstruction;
-  const messages = [{ role: "user", content: `<image>${instruction}` }];
-  const prompt = p.apply_chat_template(messages, { add_generation_prompt: true });
-  const inputs = await p(image, prompt, { add_special_tokens: false });
-
-  const generatedIds = await m.generate({
-    ...inputs,
-    do_sample: false,
-    max_new_tokens: safeMaxTokens,
-  });
-
-  const generatedText = p.batch_decode(generatedIds, {
-    skip_special_tokens: true,
-  })[0];
-  return { text: generatedText.trim() };
-}

src/model.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import os
+from io import BytesIO
+from typing import Any
+
+import torch
+from PIL import Image
+from transformers import AutoModelForCausalLM, AutoProcessor
+
+MODEL_ID = os.getenv("MODEL_ID", "microsoft/Florence-2-base")
+DEFAULT_MAX_TOKENS = int(os.getenv("DEFAULT_MAX_TOKENS", "64"))
+MAX_MAX_TOKENS = int(os.getenv("MAX_MAX_TOKENS", "256"))
+MAX_IMAGE_SIDE = int(os.getenv("MAX_IMAGE_SIDE", "896"))
+
+TASKS = {
+    "caption": "<CAPTION>",
+    "detailed_caption": "<DETAILED_CAPTION>",
+    "more_detailed_caption": "<MORE_DETAILED_CAPTION>",
+    "ocr": "<OCR>",
+    "ocr_with_region": "<OCR_WITH_REGION>",
+    "object_detection": "<OD>",
+    "dense_region_caption": "<DENSE_REGION_CAPTION>",
+    "region_proposal": "<REGION_PROPOSAL>",
+}
+
+_model = None
+_processor = None
+_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+_dtype = torch.float16 if _device.type == "cuda" else torch.float32
+
+
+def _prepare_image(image_bytes: bytes) -> Image.Image:
+    image = Image.open(BytesIO(image_bytes)).convert("RGB")
+    width, height = image.size
+    if width <= MAX_IMAGE_SIDE and height <= MAX_IMAGE_SIDE:
+        return image
+
+    ratio = min(MAX_IMAGE_SIDE / width, MAX_IMAGE_SIDE / height)
+    new_size = (max(1, int(width * ratio)), max(1, int(height * ratio)))
+    return image.resize(new_size, Image.Resampling.LANCZOS)
+
+
+def load_model() -> tuple[Any, Any]:
+    global _model, _processor
+    if _model is None or _processor is None:
+        _processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+        _model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            trust_remote_code=True,
+            torch_dtype=_dtype,
+        ).to(_device)
+        _model.eval()
+    return _model, _processor
+
+
+def generate_caption(
+    image_bytes: bytes,
+    task: str = "caption",
+    text_input: str | None = None,
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+) -> dict[str, Any]:
+    model, processor = load_model()
+    prompt_task = TASKS.get(task, TASKS["caption"])
+    prompt = f"{prompt_task} {text_input.strip()}" if text_input else prompt_task
+
+    safe_max_tokens = min(max(int(max_tokens), 8), MAX_MAX_TOKENS)
+    image = _prepare_image(image_bytes)
+
+    inputs = processor(text=prompt, images=image, return_tensors="pt")
+    input_ids = inputs["input_ids"].to(_device)
+    pixel_values = inputs["pixel_values"].to(_device, _dtype)
+
+    with torch.inference_mode():
+        generated_ids = model.generate(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            do_sample=False,
+            max_new_tokens=safe_max_tokens,
+            num_beams=1,
+        )
+
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+
+    parsed = None
+    post_process = getattr(processor, "post_process_generation", None)
+    if callable(post_process):
+        try:
+            parsed = post_process(generated_text, task=prompt_task, image_size=image.size)
+        except Exception:
+            parsed = None
+
+    return {"text": generated_text, "parsed": parsed} if parsed else {"text": generated_text}

src/server.js

@@ -1,238 +0,0 @@
-import Fastify from "fastify";
-import multipart from "@fastify/multipart";
-import swagger from "@fastify/swagger";
-import swaggerUi from "@fastify/swagger-ui";
-import cors from "@fastify/cors";
-import { generateCaption, loadModel, TASKS } from "./model.js";
-
-const app = Fastify({ logger: true });
-const DEFAULT_MAX_TOKENS = parseInt(process.env.DEFAULT_MAX_TOKENS || "64", 10);
-
-// --- Plugins ---
-await app.register(cors);
-await app.register(multipart, { limits: { fileSize: 20 * 1024 * 1024 } });
-
-await app.register(swagger, {
-  openapi: {
-    info: {
-      title: "img3txt — Image Captioning API",
-      description:
-        "Generate captions, OCR, object detection and more from images using FastVLM (ONNX).",
-      version: "1.0.0",
-    },
-    servers: [{ url: "/" }],
-    tags: [
-      { name: "caption", description: "Image captioning endpoints" },
-      { name: "health", description: "Health check" },
-    ],
-  },
-});
-
-await app.register(swaggerUi, {
-  routePrefix: "/docs",
-  uiConfig: { docExpansion: "list", deepLinking: true },
-});
-
-// --- Schemas ---
-const taskEnum = Object.keys(TASKS);
-
-const captionResponseSchema = {
-  type: "object",
-  properties: {
-    task: { type: "string", example: "caption" },
-    result: { type: "object", additionalProperties: true },
-  },
-};
-
-const batchResponseSchema = {
-  type: "object",
-  properties: {
-    results: {
-      type: "array",
-      items: {
-        type: "object",
-        properties: {
-          filename: { type: "string" },
-          task: { type: "string" },
-          result: { type: "object", additionalProperties: true },
-        },
-      },
-    },
-  },
-};
-
-const errorSchema = {
-  type: "object",
-  properties: {
-    error: { type: "string" },
-  },
-};
-
-// --- Routes ---
-
-// Landing page — HF Spaces iframe shows this
-app.get(
-  "/",
-  { schema: { hide: true } },
-  async (req, reply) => {
-    reply.type("text/html").send(`<!DOCTYPE html>
-<html lang="en"><head><meta charset="utf-8">
-<meta name="viewport" content="width=device-width,initial-scale=1">
-<title>img3txt — FastVLM Image Captioning API</title>
-<style>
-*{margin:0;padding:0;box-sizing:border-box}
-body{font-family:system-ui,sans-serif;background:#0f172a;color:#e2e8f0;display:flex;align-items:center;justify-content:center;min-height:100vh}
-.card{background:#1e293b;border-radius:16px;padding:2.5rem;max-width:520px;width:90%;text-align:center;box-shadow:0 25px 50px rgba(0,0,0,.4)}
-h1{font-size:1.8rem;margin-bottom:.5rem}
-.sub{color:#94a3b8;margin-bottom:1.5rem}
-.btn{display:inline-block;padding:.75rem 1.5rem;background:#3b82f6;color:#fff;border-radius:8px;text-decoration:none;font-weight:600;margin:.25rem}
-.btn:hover{background:#2563eb}
-.tasks{margin-top:1.5rem;text-align:left;background:#0f172a;border-radius:8px;padding:1rem}
-.tasks code{color:#38bdf8}
-</style></head><body>
-<div class="card">
-<h1>img3txt</h1>
-<p class="sub">Image captioning, OCR &amp; object detection powered by FastVLM (ONNX)</p>
-<a class="btn" href="/docs">Swagger UI</a>
-<a class="btn" href="/health">Health Check</a>
-<div class="tasks">
-<p><strong>POST /caption</strong> with form fields:</p>
-<ul style="margin:.5rem 0 0 1.2rem;color:#94a3b8">
-<li><code>file</code> — image (required)</li>
-<li><code>task</code> — caption, detailed_caption, more_detailed_caption, ocr, ocr_with_region, object_detection, dense_region_caption, region_proposal</li>
-<li><code>max_tokens</code> — default 64 (smaller = faster)</li>
-</ul>
-</div>
-</div></body></html>`);
-  }
-);
-
-app.get(
-  "/health",
-  {
-    schema: {
-      tags: ["health"],
-      summary: "Health check",
-      response: {
-        200: {
-          type: "object",
-          properties: {
-            status: { type: "string", example: "ok" },
-            model: { type: "string" },
-            tasks: { type: "array", items: { type: "string" } },
-          },
-        },
-      },
-    },
-  },
-  async () => ({
-    status: "ok",
-    model: "onnx-community/FastVLM-0.5B-ONNX",
-    tasks: taskEnum,
-  })
-);
-
-app.post(
-  "/caption",
-  {
-    schema: {
-      tags: ["caption"],
-      summary: "Generate caption / OCR / detection for a single image",
-      description: `Upload an image as multipart form data. Supported tasks: ${taskEnum.join(", ")}`,
-      consumes: ["multipart/form-data"],
-      response: {
-        200: captionResponseSchema,
-        400: errorSchema,
-      },
-    },
-  },
-  async (req, reply) => {
-    const data = await req.file();
-    if (!data) {
-      return reply.code(400).send({ error: "No file uploaded" });
-    }
-
-    const task = data.fields.task?.value || "caption";
-    const textInput = data.fields.text?.value || null;
-    const maxTokens = parseInt(
-      data.fields.max_tokens?.value || String(DEFAULT_MAX_TOKENS),
-      10
-    );
-
-    if (!TASKS[task]) {
-      return reply
-        .code(400)
-        .send({ error: `Invalid task. Choose from: ${taskEnum.join(", ")}` });
-    }
-
-    const buffer = await data.toBuffer();
-    const result = await generateCaption(buffer, task, textInput, maxTokens);
-
-    return { task, result };
-  }
-);
-
-app.post(
-  "/caption/batch",
-  {
-    schema: {
-      tags: ["caption"],
-      summary: "Generate captions for multiple images",
-      description:
-        "Upload multiple images as multipart form data. All images share the same task and settings.",
-      consumes: ["multipart/form-data"],
-      response: {
-        200: batchResponseSchema,
-        400: errorSchema,
-      },
-    },
-  },
-  async (req, reply) => {
-    const parts = await req.parts();
-    const files = [];
-    let task = "caption";
-    let textInput = null;
-    let maxTokens = DEFAULT_MAX_TOKENS;
-
-    for await (const part of parts) {
-      if (part.type === "file") {
-        files.push({ filename: part.filename, buffer: await part.toBuffer() });
-      } else if (part.fieldname === "task") {
-        task = part.value;
-      } else if (part.fieldname === "text") {
-        textInput = part.value;
-      } else if (part.fieldname === "max_tokens") {
-        maxTokens = parseInt(part.value, 10);
-      }
-    }
-
-    if (files.length === 0) {
-      return reply.code(400).send({ error: "No files uploaded" });
-    }
-    if (!TASKS[task]) {
-      return reply
-        .code(400)
-        .send({ error: `Invalid task. Choose from: ${taskEnum.join(", ")}` });
-    }
-
-    const results = [];
-    for (const f of files) {
-      const result = await generateCaption(f.buffer, task, textInput, maxTokens);
-      results.push({ filename: f.filename, task, result });
-    }
-
-    return { results };
-  }
-);
-
-// --- Start ---
-const PORT = process.env.PORT || 7860;
-
-// Pre-load model then start server
-await loadModel();
-app.listen({ host: "0.0.0.0", port: PORT }, (err) => {
-  if (err) {
-    app.log.error(err);
-    process.exit(1);
-  }
-});

src/server.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import os
+from typing import Any
+
+from fastapi import FastAPI, File, Form, HTTPException, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import HTMLResponse
+
+from .model import MODEL_ID, TASKS, DEFAULT_MAX_TOKENS, generate_caption, load_model
+
+app = FastAPI(
+    title="img3txt - Florence-2 API",
+    description="Generate captions, OCR, object detection and more from images using Florence-2.",
+    version="1.0.0",
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@app.on_event("startup")
+def warmup_model() -> None:
+    load_model()
+
+
+@app.get("/", response_class=HTMLResponse, include_in_schema=False)
+def root() -> str:
+    return """<!DOCTYPE html>
+<html lang=\"en\"><head><meta charset=\"utf-8\">
+<meta name=\"viewport\" content=\"width=device-width,initial-scale=1\">
+<title>img3txt - Florence-2 Image Captioning API</title>
+<style>
+*{margin:0;padding:0;box-sizing:border-box}
+body{font-family:system-ui,sans-serif;background:#0f172a;color:#e2e8f0;display:flex;align-items:center;justify-content:center;min-height:100vh}
+.card{background:#1e293b;border-radius:16px;padding:2.5rem;max-width:520px;width:90%;text-align:center;box-shadow:0 25px 50px rgba(0,0,0,.4)}
+h1{font-size:1.8rem;margin-bottom:.5rem}
+.sub{color:#94a3b8;margin-bottom:1.5rem}
+.btn{display:inline-block;padding:.75rem 1.5rem;background:#3b82f6;color:#fff;border-radius:8px;text-decoration:none;font-weight:600;margin:.25rem}
+.btn:hover{background:#2563eb}
+.tasks{margin-top:1.5rem;text-align:left;background:#0f172a;border-radius:8px;padding:1rem}
+.tasks code{color:#38bdf8}
+</style></head><body>
+<div class=\"card\">
+<h1>img3txt</h1>
+<p class=\"sub\">Image captioning, OCR &amp; object detection powered by Florence-2</p>
+<a class=\"btn\" href=\"/docs\">Swagger UI</a>
+<a class=\"btn\" href=\"/health\">Health Check</a>
+<div class=\"tasks\">
+<p><strong>POST /caption</strong> with form fields:</p>
+<ul style=\"margin:.5rem 0 0 1.2rem;color:#94a3b8\">
+<li><code>file</code> - image (required)</li>
+<li><code>task</code> - caption, detailed_caption, more_detailed_caption, ocr, ocr_with_region, object_detection, dense_region_caption, region_proposal</li>
+<li><code>max_tokens</code> - default 64 (smaller = faster)</li>
+</ul>
+</div>
+</div></body></html>"""
+
+
+@app.get("/health")
+def health() -> dict[str, Any]:
+    return {"status": "ok", "model": MODEL_ID, "tasks": list(TASKS.keys())}
+
+
+@app.post("/caption")
+async def caption(
+    file: UploadFile = File(...),
+    task: str = Form("caption"),
+    text: str | None = Form(None),
+    max_tokens: int = Form(DEFAULT_MAX_TOKENS),
+) -> dict[str, Any]:
+    if task not in TASKS:
+        raise HTTPException(status_code=400, detail=f"Invalid task. Choose from: {', '.join(TASKS.keys())}")
+
+    image_bytes = await file.read()
+    if not image_bytes:
+        raise HTTPException(status_code=400, detail="Empty file uploaded")
+
+    result = generate_caption(image_bytes, task, text, max_tokens)
+    return {"task": task, "result": result}
+
+
+@app.post("/caption/batch")
+async def caption_batch(
+    files: list[UploadFile] = File(...),
+    task: str = Form("caption"),
+    text: str | None = Form(None),
+    max_tokens: int = Form(DEFAULT_MAX_TOKENS),
+) -> dict[str, Any]:
+    if task not in TASKS:
+        raise HTTPException(status_code=400, detail=f"Invalid task. Choose from: {', '.join(TASKS.keys())}")
+
+    results: list[dict[str, Any]] = []
+    for upload in files:
+        image_bytes = await upload.read()
+        if not image_bytes:
+            continue
+        result = generate_caption(image_bytes, task, text, max_tokens)
+        results.append({"filename": upload.filename, "task": task, "result": result})
+
+    if not results:
+        raise HTTPException(status_code=400, detail="No files uploaded")
+
+    return {"results": results}
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    port = int(os.getenv("PORT", "7860"))
+    uvicorn.run("src.server:app", host="0.0.0.0", port=port)