Upload model

modeling.py

@@ -19,19 +19,13 @@ from transformers.utils.backbone_utils import load_backbone
 from .configuration import LSPDetrConfig, STAConfig
 
 
-def _meta_safe_is_orthogonal(Q, eps=None):
-    return True if Q.device == torch.device("meta") else _is_orthogonal(Q, eps=eps)
-
-
+flex_attention = torch.compile(flex_attention, dynamic=True)
 patch(
     "torch.nn.utils.parametrizations._is_orthogonal",
-    _meta_safe_is_orthogonal,
+    lambda Q, eps=None: Q.device == torch.device("meta") or _is_orthogonal(Q, eps=eps),
 ).start()
 
 
-flex_attention = torch.compile(flex_attention, dynamic=True)
-
-
 class CayleySTRING(nn.Module):
     """Implements the Cayley-STRING positional encoding.
 
@@ -464,6 +458,7 @@ class LSPTransformer(nn.Module):
             "absolute_points": relative_to_absolute_pos(
                 ref_points, self.query_block_size, self.query_block_size
             ).flatten(1, 2),
+            "embeddings": tgt.flatten(1, 2),
             "aux_outputs": [
                 {
                     "logits": a,
@@ -525,6 +520,3 @@ class LSPDetrModel(PreTrainedModel):
         )
 
         return self.decode_head(tgt, ref_points, features, h, w)
-
-
-LSPDetrModel.from_pretrained("RationAI/LSP-DETR", trust_remote_code=True)