Merge branch 'main' into feature/add-mps-support

OrianeN · Sep 17, 2023 · eb85ece · eb85ece
2 parents b2b0d2f + f5d2cd5
commit eb85ece
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 3 deletions.
diff --git a/nougat/_version.py b/nougat/_version.py
@@ -4,4 +4,5 @@
 This source code is licensed under the MIT license found in the
 LICENSE file in the root directory of this source tree.
 """
-__version__ = "0.1.7"
+
+__version__ = "0.1.8"
diff --git a/nougat/model.py b/nougat/model.py
@@ -571,10 +571,12 @@ def inference(
 
         if image_tensors is None:
             image_tensors = self.encoder.prepare_input(image).unsqueeze(0)
+
         if self.device.type != "cpu":
             image_tensors = image_tensors.to(self.device)
         if self.device.type != "mps":
             image_tensors = image_tensors.to(torch.bfloat16)
+
         last_hidden_state = self.encoder(image_tensors)
 
         encoder_outputs = ModelOutput(

diff --git a/nougat/utils/checkpoint.py b/nougat/utils/checkpoint.py
@@ -111,7 +111,7 @@ def get_checkpoint(
         checkpoint = checkpoint.parent
     if download and (not checkpoint.exists() or len(os.listdir(checkpoint)) < 5):
         checkpoint.mkdir(parents=True, exist_ok=True)
-        download_checkpoint(checkpoint, model_tag=model_tag)
+        download_checkpoint(checkpoint, model_tag=model_tag or MODEL_TAG)
     return checkpoint
 
 

diff --git a/predict.py b/predict.py
@@ -58,7 +58,7 @@ def get_args():
         "--model",
         "-m",
         type=str,
-        default=None,
+        default='0.1.0-small',
         help=f"Model tag to use.",
     )
     parser.add_argument("--out", "-o", type=Path, help="Output directory.")