diff --git a/README.md b/README.md
index e6afdea..6ebd0bb 100644
--- a/README.md
+++ b/README.md
@@ -1,11 +1,36 @@
# TavernAI - Extras
## What is this
A set of unofficial APIs for various [TavernAI](https://github.com/TavernAI/TavernAI) extensions
+
## How to run
* Install Python 3.10
* Run `pip install -r requirements.txt`
* Run `python server.py`
+
## Included functionality
-* BART model for text summarization
-* BERT model for text classification
-* BLIP model for image captioning
\ No newline at end of file
+### BLIP model for image captioning
+`POST /api/caption`
+#### **Input**
+```
+{ "image": "base64 encoded image" }
+```
+#### **Output**
+```
+{ "caption": "caption of the posted image" }
+```
+
+### BART model for text summarization
+Not implemented yet
+
+### BERT model for text classification
+Not implemented yet
+
+## Additional options
+| Flag | Description |
+| -------------- | -------------------------------------------------------------------- |
+| `--port` | Specify the port on which the application is hosted. Default: *5100* |
+| `--listen` | Hosts the app on the local network |
+| `--share` | Shares the app on CloudFlare tunnel |
+| `--bart-model` | Load a custom BART model.
Expects a HuggingFace model ID.
Default: [Qiliang/bart-large-cnn-samsum-ChatGPT_v3](https://huggingface.co/Qiliang/bart-large-cnn-samsum-ChatGPT_v3) |
+| `--bert-model` | Load a custom BERT model.
Expects a HuggingFace model ID.
Default: [bhadresh-savani/distilbert-base-uncased-emotion](https://huggingface.co/bhadresh-savani/distilbert-base-uncased-emotion) |
+| `--blip-model` | Load a custom BLIP model.
Expects a HuggingFace model Id.
Default: [Salesforce/blip-image-captioning-base](https://huggingface.co/Salesforce/blip-image-captioning-base) |
\ No newline at end of file
diff --git a/server.py b/server.py
index 5f168f7..10d99e6 100644
--- a/server.py
+++ b/server.py
@@ -15,18 +15,15 @@ from io import BytesIO
DEFAULT_BART = 'Qiliang/bart-large-cnn-samsum-ChatGPT_v3'
DEFAULT_BERT = 'bhadresh-savani/distilbert-base-uncased-emotion'
DEFAULT_BLIP = 'Salesforce/blip-image-captioning-base'
-UPLOAD_FOLDER = './uploads'
# Script arguments
-parser = argparse.ArgumentParser(
- prog = 'TavernAI Extras',
- description = 'Web API for transformers models')
+parser = argparse.ArgumentParser(prog = 'TavernAI Extras', description = 'Web API for transformers models')
parser.add_argument('--port', type=int, help="Specify the port on which the application is hosted")
parser.add_argument('--listen', action='store_true', help="Hosts the app on the local network")
-parser.add_argument('--share', action='store_true', help="Shares the app on cloudflare tunnel")
-parser.add_argument('--bart-model', help="Customize a BART model to be used by the app")
-parser.add_argument('--bert-model', help="Customize a BERT model to be used by the app")
-parser.add_argument('--blip-model', help="Customize a BLIP model to be used by the app")
+parser.add_argument('--share', action='store_true', help="Shares the app on CloudFlare tunnel")
+parser.add_argument('--bart-model', help="Load a custom BART model")
+parser.add_argument('--bert-model', help="Load a custom BERT model")
+parser.add_argument('--blip-model', help="Load a custom BLIP model")
args = parser.parse_args()
@@ -74,7 +71,6 @@ bert = AutoModelForSequenceClassification.from_pretrained(bert_model)
# Flask init
app = Flask(__name__)
-app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = 100 * 1024 * 1024
# AI stuff