Add torch_device option to scripts/auto_caption.py

This allows using auto_caption on apple sillicon macs by specifying cpu
as an argument for now, and might allow using mps eventually, once more
operators are implemented.

scripts/auto_caption.py

@@ -66,6 +66,14 @@ def get_parser(**parser_kwargs):
         default=22,
         help="adjusts the likelihood of a word being repeated",
     ),
+    parser.add_argument(
+        "--torch_device",
+        type=str,
+        nargs="?",
+        const=False,
+        default="cuda",
+        help="specify a different torch device, e.g. 'cpu'",
+    ),
 
     return parser
 
@@ -119,9 +127,9 @@ async def main(opt):
     blip_decoder = models.blip.blip_decoder(pretrained=model_cache_path, image_size=SIZE, vit='base', med_config=config_path)
     blip_decoder.eval()
 
-    print("loading model to cuda")
+    print(f"loading model to {opt.torch_device}")
 
-    blip_decoder = blip_decoder.to(torch.device("cuda"))
+    blip_decoder = blip_decoder.to(torch.device(opt.torch_device))
 
     ext = ('.jpg', '.jpeg', '.png', '.webp', '.tif', '.tga', '.tiff', '.bmp', '.gif')
 
@@ -141,7 +149,7 @@ async def main(opt):
                     if not image.mode == "RGB":
                         image = image.convert("RGB")
 
-                    image = load_image(image, device=torch.device("cuda"))
+                    image = load_image(image, device=torch.device(opt.torch_device))
 
                     if opt.nucleus:
                         captions = blip_decoder.generate(image, sample=True, top_p=opt.q_factor)
@@ -207,4 +215,3 @@ if __name__ == "__main__":
     sys.path.append(blip_path)
 
     asyncio.run(main(opt))
-