From 8cb475e8fa88bfa7738718152c5db0787554762a Mon Sep 17 00:00:00 2001
From: Nikita Sokovnin <nikitoos67@gmail.com>
Date: Tue, 20 Feb 2024 16:40:45 +0000
Subject: [PATCH] docs: update args description

---
 README.md                                      | 3 ++-
 examples/generate_dataset_and_train_yolo.ipynb | 3 +++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 7ff73df..db3450b 100644
--- a/README.md
+++ b/README.md
@@ -116,6 +116,7 @@ datadreamer --save_dir <directory> --class_names <objects> --prompts_number <num
 - `--image_tester_patience`: Patience level for image tester. Default is 1.
 - `--lm_quantization`: Quantization to use for Mistral language model. Choose between `none` and `4bit`. Default is `none`.
 - `--batch_size_prompt`: Batch size for prompt generation. Default is 64.
+- `--batch_size_image`: Batch size for image generation. Default is 1.
 - `--device`: Choose between `cuda` and `cpu`. Default is cuda.
 - `--seed`: Set a random seed for image and prompt generation. Default is 42.
 
@@ -130,7 +131,7 @@ datadreamer --save_dir <directory> --class_names <objects> --prompts_number <num
 |                   | Simple random generator                                                               | Joins randomly chosen object names      |
 | Image Generation  | [SDXL-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)           | Slow and accurate (1024x1024 images)    |
 |                   | [SDXL-Turbo](https://huggingface.co/stabilityai/sdxl-turbo)                           | Fast and less accurate (512x512 images) |
-| Image Annotation  | [OWLv2](https://huggingface.co/google/owlv2-large-patch14-ensemble)                   | Open-Vocabulary object detector         |
+| Image Annotation  | [OWLv2](https://huggingface.co/google/owlv2-base-patch16-ensemble)                    | Open-Vocabulary object detector         |
 
 <a name="example"></a>
 
diff --git a/examples/generate_dataset_and_train_yolo.ipynb b/examples/generate_dataset_and_train_yolo.ipynb
index 1a18140..ea5ab84 100644
--- a/examples/generate_dataset_and_train_yolo.ipynb
+++ b/examples/generate_dataset_and_train_yolo.ipynb
@@ -84,6 +84,9 @@
         "- `--enhance_class_names`: Enhance class names with synonyms. Default is False.\n",
         "- `--use_image_tester`: Use image tester for image generation. Default is False. Image tester is a model that checks if the generated image contains the desired object.\n",
         "- `--image_tester_patience`: Patience level for image tester. Default is 1. Number of images to run through the image tester before giving up.\n",
+        "- `--lm_quantization`: Quantization to use for Mistral language model. Choose between `none` and `4bit`. Default is `none`.\n",
+        "- `--batch_size_prompt`: Batch size for prompt generation. Default is 64. Increase for faster prompt generation, decrease if there is no enough GPU memory.\n",
+        "- `--batch_size_image`: Batch size for image generation. Default is 1. Increase for faster image generation.\n",
         "- `--seed`: Set a random seed for image and prompt generation. Default is 42.\n"
       ]
     },