optimum-rbln 0.9.2a8__tar.gz → 0.9.2a10__tar.gz

{optimum_rbln-0.9.2a8 → optimum_rbln-0.9.2a10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.9.2a8
+Version: 0.9.2a10
 Summary: Optimum RBLN is the interface between the HuggingFace Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai

{optimum_rbln-0.9.2a8 → optimum_rbln-0.9.2a10}/src/optimum/rbln/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.9.2a8'
-__version_tuple__ = version_tuple = (0, 9, 2, 'a8')
+__version__ = version = '0.9.2a10'
+__version_tuple__ = version_tuple = (0, 9, 2, 'a10')
 
 __commit_id__ = commit_id = None

{optimum_rbln-0.9.2a8 → optimum_rbln-0.9.2a10}/src/optimum/rbln/transformers/modeling_generic.py

@@ -267,6 +267,20 @@ class RBLNModelForImageClassification(RBLNImageModel):
 class RBLNModelForDepthEstimation(RBLNImageModel):
     auto_model_class = AutoModelForDepthEstimation
 
+    @classmethod
+    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNImageModelConfig):
+        class ImageModelWrapper(nn.Module):
+            def __init__(self, model: "PreTrainedModel", rbln_config: RBLNImageModelConfig):
+                super().__init__()
+                self.model = model
+                self.rbln_config = rbln_config
+
+            def forward(self, *args, **kwargs):
+                output = self.model(*args, return_dict=True, **kwargs)
+                return output.predicted_depth
+
+        return ImageModelWrapper(model, rbln_config).eval()
+
 
 class RBLNModelForAudioClassification(RBLNModel):
     """