From 87aa2e68728d33ecf0a887cd7341c349ca8a1a62 Mon Sep 17 00:00:00 2001
From: aiXander <xandersteenbruggex@gmail.com>
Date: Wed, 14 Aug 2024 15:27:32 -0700
Subject: [PATCH] update reqs

---
 requirements.txt      | 10 +++++-----
 trainer/preprocess.py |  4 +---
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 2959012..81d1b40 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,6 @@
-torch==2.1.0
-torchaudio==2.1.0
-torchvision==0.16.0
+torch==2.2.1
+torchaudio==2.2.1
+torchvision==0.17.1
 transformers==4.38.0
 diffusers==0.26.0
 tokenizers==0.15.2
@@ -21,5 +21,5 @@ ujson==5.10.0
 bitsandbytes==0.43.1
 setuptools==70.3.0
 torchtyping==0.1.5
-einops
-timm
\ No newline at end of file
+einops==0.8.0
+timm==1.0.8
\ No newline at end of file
diff --git a/trainer/preprocess.py b/trainer/preprocess.py
index 18b83e5..b471a7a 100755
--- a/trainer/preprocess.py
+++ b/trainer/preprocess.py
@@ -597,8 +597,6 @@ def caption_dataset(
         caption_model: Literal["blip", "gpt4-v", "florence"] = "blip"
     ) -> List[str]:
 
-    print(f"Captioning images using {caption_model}...")
-
     if "blip" in caption_model:
         captions = blip_caption_dataset(images, captions)
     elif "gpt4-v" in caption_model:
@@ -820,7 +818,7 @@ def load_and_save_masks_and_captions(
         captions = captions + captions
 
 
-    print(f"Generating {len(images)} captions using mode: {concept_mode}...")
+    print(f"Generating {len(images)} captions using {caption_model} in {concept_mode} mode...")
     captions = caption_dataset(images, captions, caption_model = caption_model)
 
     # It's nice if we can achieve the gpt pass, so if we're not losing too much, cut-off the n_images to just match what we're allowed to give to gpt: