Merge pull request #787 from roboflow/florence2-in-app-train

Serve florence2 models created by in app train
roboflow · Nov 8, 2024 · ad704c9 · ad704c9
2 parents 6049ebb + fe8b0ec
commit ad704c9
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/inference/models/transformers/transformers.py b/inference/models/transformers/transformers.py
@@ -1,5 +1,6 @@
 import os
 import re
+import subprocess
 import tarfile
 
 import numpy as np
@@ -179,6 +180,18 @@ def download_model_artifacts_from_roboflow_api(self) -> None:
                 file=filename,
                 model_id=self.endpoint,
             )
+            if filename.endswith("tar.gz"):
+                subprocess.run(
+                    [
+                        "tar",
+                        "-xzf",
+                        os.path.join(self.cache_dir, filename),
+                        "-C",
+                        self.cache_dir,
+                    ],
+                    check=True,
+                )
+
             if perf_counter() - t1 > 120:
                 logger.debug(
                     "Weights download took longer than 120 seconds, refreshing API request"
@@ -234,7 +247,7 @@ def initialize_model(self):
         )
 
         self.processor = self.processor_class.from_pretrained(
-            self.cache_dir, revision=revision
+            model_load_id, revision=revision, cache_dir=cache_dir, token=token
         )
 
     def get_lora_base_from_roboflow(self, repo, revision) -> str:

diff --git a/inference/models/utils.py b/inference/models/utils.py
@@ -267,12 +267,22 @@
             "object-detection",
             "florence-2-base-peft",
         ): LoRAFlorence2,  # TODO: change when we have a new project type
+        (
+            "text-image-pairs",
+            "florence-2-base",
+        ): Florence2,  # TODO: change when we have a new project type
+        ("text-image-pairs", "florence-2-large"): Florence2,
         ("object-detection", "florence-2-large-peft"): LoRAFlorence2,
         (
             "instance-segmentation",
             "florence-2-base-peft",
         ): LoRAFlorence2,  # TODO: change when we have a new project type
         ("instance-segmentation", "florence-2-large-peft"): LoRAFlorence2,
+        (
+            "text-image-pairs",
+            "florence-2-base-peft",
+        ): LoRAFlorence2,
+        ("text-image-pairs", "florence-2-large-peft"): LoRAFlorence2,
     }
     ROBOFLOW_MODEL_TYPES.update(florence2_models)
 except: