update reqs

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
+einops==0.8.0
+timm==1.0.8

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: