Spaces:
Configuration error
Configuration error
added initial files
Browse files- .gitattributes +35 -35
- README.md +70 -13
- app.py +81 -0
- paligemma_tokenizer.model +3 -0
- requirements.txt +13 -0
- utils.py +204 -0
.gitattributes
CHANGED
|
@@ -1,35 +1,35 @@
|
|
| 1 |
-
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
-
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
-
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
-
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
-
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
-
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
-
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
-
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
-
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
-
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
-
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
-
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
-
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
-
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
-
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
-
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
-
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
-
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
-
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
-
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
-
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
-
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
-
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
-
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
-
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
-
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
-
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
-
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
-
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
-
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
-
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
-
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
-
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
-
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
-
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
| 1 |
+
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
+
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
+
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
+
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
+
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
+
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
+
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
+
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
+
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
+
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
+
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
+
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
+
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
+
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
+
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
+
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
+
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
+
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
+
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
+
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
+
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
+
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
+
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
+
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
+
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
+
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
+
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
+
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
+
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
+
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
+
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
+
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
README.md
CHANGED
|
@@ -1,13 +1,70 @@
|
|
| 1 |
-
|
| 2 |
-
|
| 3 |
-
|
| 4 |
-
|
| 5 |
-
|
| 6 |
-
|
| 7 |
-
|
| 8 |
-
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
|
| 13 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
## Inhabitr Design Engine 1.0
|
| 2 |
+
|
| 3 |
+
#### PaliGemma Vision Model(unifying attributes model)
|
| 4 |
+
Leveraging the PaliGemma vision model to unify attribute extraction, this custom trained Model identifies and outputs furniture attributes—such as primary color and secondary color, primary material, secondary material, and design style—in JSON format.
|
| 5 |
+
|
| 6 |
+
# Features
|
| 7 |
+
Paligemma Model: A 3-billion parameter multi-model capable of generating captions from images, performing object detection, object segmentation ,OCR and VQA.
|
| 8 |
+
Custom Training: The model is custom trained on approximately 500 pairs of cropped images (sofa, accent chair) along with their captions.
|
| 9 |
+
|
| 10 |
+
# Requirements
|
| 11 |
+
Python 3.9+
|
| 12 |
+
|
| 13 |
+
Minimum 12GB of GPU for local inference
|
| 14 |
+
|
| 15 |
+
Pip (Python package installer)
|
| 16 |
+
|
| 17 |
+
# Installation
|
| 18 |
+
First, clone the repository:
|
| 19 |
+
|
| 20 |
+
git clone git clone https://Apoorva_inhabitr@bitbucket.org/Inhabitr/vision_models.git
|
| 21 |
+
|
| 22 |
+
# Create and Activate Python Virtual Environment
|
| 23 |
+
|
| 24 |
+
For Unix/macOS
|
| 25 |
+
|
| 26 |
+
python3 -m venv env
|
| 27 |
+
source env/bin/activate
|
| 28 |
+
|
| 29 |
+
For Windows
|
| 30 |
+
|
| 31 |
+
python -m venv env
|
| 32 |
+
.\env\Scripts\activate
|
| 33 |
+
|
| 34 |
+
# Install the required Python libraries
|
| 35 |
+
|
| 36 |
+
pip install -r requirements.txt
|
| 37 |
+
|
| 38 |
+
# Download Model Weights and Tokenizer
|
| 39 |
+
|
| 40 |
+
Download the model weights and tokenizer from Google Drive
|
| 41 |
+
|
| 42 |
+
https://drive.google.com/drive/folders/1S-z374V-yd3izeBatAMQlitmZ0mbMv6s?usp=drive_link.
|
| 43 |
+
|
| 44 |
+
Place the downloaded files in the models/ folder located in the root directory of the project.
|
| 45 |
+
|
| 46 |
+
# Start the API
|
| 47 |
+
python main.py
|
| 48 |
+
If running locally, the API will be available at http://localhost:5000.
|
| 49 |
+
|
| 50 |
+
# API Endpoints for Captions
|
| 51 |
+
predict attributes
|
| 52 |
+
|
| 53 |
+
POST /predict
|
| 54 |
+
|
| 55 |
+
Request Parameters
|
| 56 |
+
|
| 57 |
+
Send an image file
|
| 58 |
+
|
| 59 |
+
image: The image file used as the basis for the search.
|
| 60 |
+
|
| 61 |
+
Example Request with Image File
|
| 62 |
+
|
| 63 |
+
curl -X POST http://127.0.0.1:5000/predict \
|
| 64 |
+
-H "Content-Type: multipart/form-data" \
|
| 65 |
+
-F "file=@/path/to/your/image.jpg"
|
| 66 |
+
|
| 67 |
+
# Performance Considerations
|
| 68 |
+
GPU: Model parameters are loaded into memory for fast inference within 1 sec after initial loading.
|
| 69 |
+
|
| 70 |
+
CPU: Initial model loading may take 4-6 minutes based on your system configuration.
|
app.py
ADDED
|
@@ -0,0 +1,81 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import streamlit as st
|
| 2 |
+
from PIL import Image
|
| 3 |
+
import os
|
| 4 |
+
from fastapi import FastAPI, File, UploadFile
|
| 5 |
+
from pydantic import BaseModel
|
| 6 |
+
import uvicorn
|
| 7 |
+
import threading
|
| 8 |
+
from werkzeug.utils import secure_filename
|
| 9 |
+
from utils import setup_and_predict
|
| 10 |
+
from io import BytesIO
|
| 11 |
+
|
| 12 |
+
# Configure upload folder and allowed extensions
|
| 13 |
+
UPLOAD_FOLDER = 'uploads'
|
| 14 |
+
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
|
| 15 |
+
|
| 16 |
+
if not os.path.exists(UPLOAD_FOLDER):
|
| 17 |
+
os.makedirs(UPLOAD_FOLDER)
|
| 18 |
+
|
| 19 |
+
def allowed_file(filename):
|
| 20 |
+
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
|
| 21 |
+
|
| 22 |
+
def save_uploaded_file(uploaded_file):
|
| 23 |
+
filename = secure_filename(uploaded_file.name)
|
| 24 |
+
filepath = os.path.join(UPLOAD_FOLDER, filename)
|
| 25 |
+
with open(filepath, "wb") as f:
|
| 26 |
+
f.write(uploaded_file.getbuffer())
|
| 27 |
+
return filepath
|
| 28 |
+
|
| 29 |
+
# FastAPI app for Postman requests
|
| 30 |
+
api = FastAPI()
|
| 31 |
+
|
| 32 |
+
class PredictionResponse(BaseModel):
|
| 33 |
+
filename: str
|
| 34 |
+
result: str
|
| 35 |
+
|
| 36 |
+
@api.post("/predict", response_model=PredictionResponse)
|
| 37 |
+
async def predict_api(file: UploadFile = File(...)):
|
| 38 |
+
if allowed_file(file.filename):
|
| 39 |
+
filename = secure_filename(file.filename)
|
| 40 |
+
filepath = os.path.join(UPLOAD_FOLDER, filename)
|
| 41 |
+
|
| 42 |
+
with open(filepath, "wb") as f:
|
| 43 |
+
f.write(await file.read())
|
| 44 |
+
|
| 45 |
+
try:
|
| 46 |
+
result = setup_and_predict(filepath)
|
| 47 |
+
except Exception as e:
|
| 48 |
+
result = f"Unable to process the request: {e}"
|
| 49 |
+
|
| 50 |
+
return PredictionResponse(filename=filename, result=result)
|
| 51 |
+
|
| 52 |
+
return {"error": "Invalid file type"}
|
| 53 |
+
|
| 54 |
+
def run_api():
|
| 55 |
+
uvicorn.run(api, host="0.0.0.0", port=8000)
|
| 56 |
+
|
| 57 |
+
# Streamlit app
|
| 58 |
+
st.title("Upload an Image for Prediction")
|
| 59 |
+
|
| 60 |
+
uploaded_file = st.file_uploader("Choose an image...", type=['png', 'jpg', 'jpeg', 'gif'])
|
| 61 |
+
|
| 62 |
+
if uploaded_file is not None and allowed_file(uploaded_file.name):
|
| 63 |
+
filepath = save_uploaded_file(uploaded_file)
|
| 64 |
+
|
| 65 |
+
try:
|
| 66 |
+
result = setup_and_predict(filepath)
|
| 67 |
+
st.image(filepath, caption='Uploaded Image.', use_column_width=True)
|
| 68 |
+
st.write("Prediction Result:")
|
| 69 |
+
st.write(result)
|
| 70 |
+
except Exception as e:
|
| 71 |
+
st.error(f"Unable to process the request: {e}")
|
| 72 |
+
|
| 73 |
+
# Option to download the image
|
| 74 |
+
with open(filepath, "rb") as file:
|
| 75 |
+
st.download_button(label="Download Image", data=file, file_name=uploaded_file.name)
|
| 76 |
+
|
| 77 |
+
else:
|
| 78 |
+
st.warning("Please upload a valid image file.")
|
| 79 |
+
|
| 80 |
+
# Run FastAPI in a separate thread
|
| 81 |
+
threading.Thread(target=run_api).start()
|
paligemma_tokenizer.model
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:8986bb4f423f07f8c7f70d0dbe3526fb2316056c17bae71b1ea975e77a168fc6
|
| 3 |
+
size 4264023
|
requirements.txt
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Flask
|
| 2 |
+
jax
|
| 3 |
+
jaxlib
|
| 4 |
+
numpy
|
| 5 |
+
tensorflow
|
| 6 |
+
sentencepiece
|
| 7 |
+
Pillow
|
| 8 |
+
ml_collections
|
| 9 |
+
overrides
|
| 10 |
+
einops~=0.7
|
| 11 |
+
numpy>=1.26
|
| 12 |
+
flax
|
| 13 |
+
pillow
|
utils.py
ADDED
|
@@ -0,0 +1,204 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import sys
|
| 3 |
+
import base64
|
| 4 |
+
import functools
|
| 5 |
+
import html
|
| 6 |
+
import io
|
| 7 |
+
import warnings
|
| 8 |
+
import jax
|
| 9 |
+
import jax.numpy as jnp
|
| 10 |
+
import numpy as np
|
| 11 |
+
import ml_collections
|
| 12 |
+
import tensorflow as tf
|
| 13 |
+
import sentencepiece
|
| 14 |
+
from PIL import Image
|
| 15 |
+
|
| 16 |
+
# TPUs with
|
| 17 |
+
if "COLAB_TPU_ADDR" in os.environ:
|
| 18 |
+
raise "It seems you are using Colab with remote TPUs which is not supported."
|
| 19 |
+
# Append big_vision code to python import path
|
| 20 |
+
if "big_vision_repo" not in sys.path:
|
| 21 |
+
sys.path.append("big_vision_repo")
|
| 22 |
+
|
| 23 |
+
# Import model definition from big_vision
|
| 24 |
+
from big_vision.models.proj.paligemma import paligemma
|
| 25 |
+
from big_vision.trainers.proj.paligemma import predict_fns
|
| 26 |
+
# Import big vision utilities
|
| 27 |
+
import big_vision.datasets.jsonl
|
| 28 |
+
import big_vision.utils
|
| 29 |
+
import big_vision.sharding
|
| 30 |
+
# Don't let TF use the GPU or TPUs
|
| 31 |
+
tf.config.set_visible_devices([], "GPU")
|
| 32 |
+
tf.config.set_visible_devices([], "TPU")
|
| 33 |
+
backend = jax.lib.xla_bridge.get_backend()
|
| 34 |
+
|
| 35 |
+
model_path = './Sofa-attributes-paligemma-ckpt.npz'
|
| 36 |
+
tokenizer_path = './paligemma_tokenizer.model'
|
| 37 |
+
|
| 38 |
+
# Define model
|
| 39 |
+
model_config = ml_collections.FrozenConfigDict({
|
| 40 |
+
"llm": {"vocab_size": 257_152},
|
| 41 |
+
"img": {"variant": "So400m/14", "pool_type": "none", "scan": True, "dtype_mm": "float16"}
|
| 42 |
+
})
|
| 43 |
+
model = paligemma.Model(**model_config)
|
| 44 |
+
tokenizer = sentencepiece.SentencePieceProcessor(tokenizer_path)
|
| 45 |
+
# Load params - this can take up to 1 minute in T4 colabs.
|
| 46 |
+
params = paligemma.load(None, model_path, model_config)
|
| 47 |
+
# Define `decode` function to sample outputs from the model.
|
| 48 |
+
decode_fn = predict_fns.get_all(model)['decode']
|
| 49 |
+
decode = functools.partial(decode_fn, devices=jax.devices(), eos_token=tokenizer.eos_id())
|
| 50 |
+
|
| 51 |
+
# Create a pytree mask of the trainable params.
|
| 52 |
+
def is_trainable_param(name, param): # pylint: disable=unused-argument
|
| 53 |
+
if name.startswith("llm/layers/attn/"): return True
|
| 54 |
+
if name.startswith("llm/"): return False
|
| 55 |
+
if name.startswith("img/"): return False
|
| 56 |
+
raise ValueError(f"Unexpected param name {name}")
|
| 57 |
+
trainable_mask = big_vision.utils.tree_map_with_names(is_trainable_param, params)
|
| 58 |
+
# If more than one device is available (e.g. multiple GPUs) the parameters can
|
| 59 |
+
# be sharded across them to reduce HBM usage per device.
|
| 60 |
+
mesh = jax.sharding.Mesh(jax.devices(), ("data"))
|
| 61 |
+
|
| 62 |
+
data_sharding = jax.sharding.NamedSharding(
|
| 63 |
+
mesh, jax.sharding.PartitionSpec("data"))
|
| 64 |
+
|
| 65 |
+
params_sharding = big_vision.sharding.infer_sharding(
|
| 66 |
+
params, strategy=[('.*', 'fsdp(axis="data")')], mesh=mesh)
|
| 67 |
+
|
| 68 |
+
# Yes: Some donated buffers are not usable.
|
| 69 |
+
warnings.filterwarnings(
|
| 70 |
+
"ignore", message="Some donated buffers were not usable")
|
| 71 |
+
|
| 72 |
+
@functools.partial(jax.jit, donate_argnums=(0,), static_argnums=(1,))
|
| 73 |
+
def maybe_cast_to_f32(params, trainable):
|
| 74 |
+
return jax.tree.map(lambda p, m: p.astype(jnp.float32) if m else p,
|
| 75 |
+
params, trainable)
|
| 76 |
+
|
| 77 |
+
# Loading all params in simultaneous - albeit much faster and more succinct -
|
| 78 |
+
# requires more RAM than the T4 colab runtimes have by default (12GB RAM).
|
| 79 |
+
# Instead we do it param by param.
|
| 80 |
+
params, treedef = jax.tree.flatten(params)
|
| 81 |
+
sharding_leaves = jax.tree.leaves(params_sharding)
|
| 82 |
+
trainable_leaves = jax.tree.leaves(trainable_mask)
|
| 83 |
+
for idx, (sharding, trainable) in enumerate(zip(sharding_leaves, trainable_leaves)):
|
| 84 |
+
params[idx] = big_vision.utils.reshard(params[idx], sharding)
|
| 85 |
+
params[idx] = maybe_cast_to_f32(params[idx], trainable)
|
| 86 |
+
params[idx].block_until_ready()
|
| 87 |
+
params = jax.tree.unflatten(treedef, params)
|
| 88 |
+
|
| 89 |
+
# Print params to show what the model is made of.
|
| 90 |
+
def parameter_overview(params):
|
| 91 |
+
for path, arr in big_vision.utils.tree_flatten_with_names(params)[0]:
|
| 92 |
+
print(f"{path:80s} {str(arr.shape):22s} {arr.dtype}")
|
| 93 |
+
|
| 94 |
+
print(" == Model params == ")
|
| 95 |
+
parameter_overview(params)
|
| 96 |
+
def setup_and_predict(image_path):
|
| 97 |
+
# Preprocess image and tokens
|
| 98 |
+
def preprocess_image(image, size=224):
|
| 99 |
+
# Model has been trained to handle images of different aspects ratios
|
| 100 |
+
# resized to 224x224 in the range [-1, 1]. Bilinear and antialias resize
|
| 101 |
+
# options are helpful to improve quality in some tasks.
|
| 102 |
+
image = np.asarray(image)
|
| 103 |
+
if image.ndim == 2: # Convert image without last channel into greyscale.
|
| 104 |
+
image = np.stack((image,)*3, axis=-1)
|
| 105 |
+
image = image[..., :3] # Remove alpha layer.
|
| 106 |
+
assert image.shape[-1] == 3
|
| 107 |
+
|
| 108 |
+
image = tf.constant(image)
|
| 109 |
+
image = tf.image.resize(image, (size, size), method='bilinear', antialias=True)
|
| 110 |
+
return image.numpy() / 127.5 - 1.0 # [0, 255]->[-1,1]
|
| 111 |
+
|
| 112 |
+
def preprocess_tokens(prefix, suffix=None, seqlen=None):
|
| 113 |
+
# Model has been trained to handle tokenized text composed of a prefix with
|
| 114 |
+
# full attention and a suffix with causal attention.
|
| 115 |
+
separator = "\n"
|
| 116 |
+
tokens = tokenizer.encode(prefix, add_bos=True) + tokenizer.encode(separator)
|
| 117 |
+
mask_ar = [0] * len(tokens) # 0 to use full attention for prefix.
|
| 118 |
+
mask_loss = [0] * len(tokens) # 0 to not use prefix tokens in the loss.
|
| 119 |
+
|
| 120 |
+
if suffix:
|
| 121 |
+
suffix = tokenizer.encode(suffix, add_eos=True)
|
| 122 |
+
tokens += suffix
|
| 123 |
+
mask_ar += [1] * len(suffix) # 1 to use causal attention for suffix.
|
| 124 |
+
mask_loss += [1] * len(suffix) # 1 to use suffix tokens in the loss.
|
| 125 |
+
|
| 126 |
+
mask_input = [1] * len(tokens) # 1 if its a token, 0 if padding.
|
| 127 |
+
if seqlen:
|
| 128 |
+
padding = [0] * max(0, seqlen - len(tokens))
|
| 129 |
+
tokens = tokens[:seqlen] + padding
|
| 130 |
+
mask_ar = mask_ar[:seqlen] + padding
|
| 131 |
+
mask_loss = mask_loss[:seqlen] + padding
|
| 132 |
+
mask_input = mask_input[:seqlen] + padding
|
| 133 |
+
|
| 134 |
+
return jax.tree.map(np.array, (tokens, mask_ar, mask_loss, mask_input))
|
| 135 |
+
|
| 136 |
+
def postprocess_tokens(tokens):
|
| 137 |
+
tokens = tokens.tolist() # np.array to list[int]
|
| 138 |
+
try: # Remove tokens at and after EOS if any.
|
| 139 |
+
eos_pos = tokens.index(tokenizer.eos_id())
|
| 140 |
+
tokens = tokens[:eos_pos]
|
| 141 |
+
except ValueError:
|
| 142 |
+
pass
|
| 143 |
+
return tokenizer.decode(tokens)
|
| 144 |
+
|
| 145 |
+
# Make predictions
|
| 146 |
+
# Evaluation/inference loop.
|
| 147 |
+
SEQLEN = 128
|
| 148 |
+
def make_predictions(data_iterator, *, num_examples=None,
|
| 149 |
+
batch_size=4, seqlen=SEQLEN, sampler="greedy"):
|
| 150 |
+
outputs = []
|
| 151 |
+
while True:
|
| 152 |
+
# Construct a list of examples in the batch.
|
| 153 |
+
examples = []
|
| 154 |
+
try:
|
| 155 |
+
for _ in range(batch_size):
|
| 156 |
+
examples.append(next(data_iterator))
|
| 157 |
+
examples[-1]["_mask"] = np.array(True) # Indicates true example.
|
| 158 |
+
except StopIteration:
|
| 159 |
+
if len(examples) == 0:
|
| 160 |
+
return outputs
|
| 161 |
+
|
| 162 |
+
# Not enough examples to complete a batch. Pad by repeating last example.
|
| 163 |
+
while len(examples) % batch_size:
|
| 164 |
+
examples.append(dict(examples[-1]))
|
| 165 |
+
examples[-1]["_mask"] = np.array(False) # Indicates padding example.
|
| 166 |
+
|
| 167 |
+
# Convert list of examples into a dict of np.arrays and load onto devices.
|
| 168 |
+
batch = jax.tree.map(lambda *x: np.stack(x), *examples)
|
| 169 |
+
batch = big_vision.utils.reshard(batch, data_sharding)
|
| 170 |
+
|
| 171 |
+
# Make model predictions
|
| 172 |
+
tokens = decode({"params": params}, batch=batch,
|
| 173 |
+
max_decode_len=seqlen, sampler=sampler)
|
| 174 |
+
|
| 175 |
+
# Fetch model predictions to device and detokenize.
|
| 176 |
+
tokens, mask = jax.device_get((tokens, batch["_mask"]))
|
| 177 |
+
tokens = tokens[mask] # remove padding examples.
|
| 178 |
+
responses = [postprocess_tokens(t) for t in tokens]
|
| 179 |
+
|
| 180 |
+
# Append to html output.
|
| 181 |
+
for example, response in zip(examples, responses):
|
| 182 |
+
outputs.append((example["image"], response))
|
| 183 |
+
if num_examples and len(outputs) >= num_examples:
|
| 184 |
+
return outputs
|
| 185 |
+
|
| 186 |
+
def test_data_iterator(file_name):
|
| 187 |
+
image = Image.open(file_name)
|
| 188 |
+
image = preprocess_image(image)
|
| 189 |
+
|
| 190 |
+
prefix = "caption en"
|
| 191 |
+
tokens, mask_ar, _, mask_input = preprocess_tokens(prefix, seqlen=SEQLEN)
|
| 192 |
+
|
| 193 |
+
yield {
|
| 194 |
+
"image": np.asarray(image),
|
| 195 |
+
"text": np.asarray(tokens),
|
| 196 |
+
"mask_ar": np.asarray(mask_ar),
|
| 197 |
+
"mask_input": np.asarray(mask_input)
|
| 198 |
+
}
|
| 199 |
+
|
| 200 |
+
# Call the prediction function and print the result
|
| 201 |
+
image, caption = make_predictions(test_data_iterator(file_name=image_path), batch_size=1)[0]
|
| 202 |
+
return caption
|
| 203 |
+
|
| 204 |
+
|