fast api app

Dockerfile

@@ -0,0 +1,21 @@
+# Use a small CUDA image if you plan to request GPU; CPU also works.
+FROM python:3.11-slim
+
+# Create non-root user (required by Spaces Docker) and prepare workdir
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user PATH=/home/user/.local/bin:$PATH
+WORKDIR $HOME/app
+
+# Copy files
+COPY --chown=user requirements.txt ./
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY --chown=user app.py ./
+
+# Spaces expects your app to listen on port 7860 unless overridden in README yaml
+ENV PORT=7860
+EXPOSE 7860
+
+# Run FastAPI
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,53 @@
+import os, io, time, contextlib, torch
+from fastapi import FastAPI, UploadFile, File
+from PIL import Image
+from transformers import (VisionEncoderDecoderModel, AutoTokenizer, AutoImageProcessor,
+                          BitsAndBytesConfig)
+
+MODEL_ID = os.getenv("MODEL_ID", "Parsa2025AI/r2gen-swin-cerebras-ft")
+GEN_MAX_LEN = int(os.getenv("GEN_MAX_LEN", "192"))
+NUM_BEAMS = int(os.getenv("NUM_BEAMS", "1"))
+
+app = FastAPI(title="R2Gen API (FastAPI on Spaces)")
+
+# Quantization + auto device map works on CPU or GPU Space
+bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_use_double_quant=True,
+                         bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16)
+
+image_processor = AutoImageProcessor.from_pretrained(MODEL_ID)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = VisionEncoderDecoderModel.from_pretrained(
+    MODEL_ID, quantization_config=bnb, device_map="auto", offload_folder="/data/offload"
+)
+model.eval()
+
+# IDs for generation
+if model.config.pad_token_id is None and tokenizer.pad_token_id is not None:
+    model.config.pad_token_id = tokenizer.pad_token_id
+if model.config.eos_token_id is None and tokenizer.eos_token_id is not None:
+    model.config.eos_token_id = tokenizer.eos_token_id
+
+@app.get("/health")
+def health():
+    return {"ok": True, "model": MODEL_ID}
+
+@app.post("/generate")
+def generate(file: UploadFile = File(...)):
+    img = Image.open(io.BytesIO(file.file.read())).convert("RGB")
+    inputs = image_processor(img, return_tensors="pt")
+
+    # Match encoder dtype/device (important when quantized/offloaded)
+    enc_param = next(model.encoder.parameters())
+    pixel_values = inputs.pixel_values.to(device=enc_param.device, dtype=enc_param.dtype)
+
+    gen_kwargs = dict(max_length=GEN_MAX_LEN, num_beams=NUM_BEAMS,
+                      pad_token_id=model.config.pad_token_id, eos_token_id=model.config.eos_token_id)
+
+    t0 = time.time()
+    with torch.inference_mode():
+        use_amp = (enc_param.device.type == "cuda" and enc_param.dtype in (torch.float16, torch.bfloat16))
+        ctx = torch.autocast("cuda", dtype=enc_param.dtype) if use_amp else contextlib.nullcontext()
+        with ctx:
+            out = model.generate(pixel_values=pixel_values, **gen_kwargs)
+    text = tokenizer.decode(out[0], skip_special_tokens=True).strip()
+    return {"text": text, "ms": int((time.time() - t0) * 1000)}

requirements.txt

@@ -0,0 +1,7 @@
+fastapi
+uvicorn
+transformers>=4.42
+accelerate
+bitsandbytes
+torch
+pillow