Upload model

modeling_basnet.py

@@ -14,7 +14,7 @@ logger = logging.getLogger(__name__)
 
 
 @dataclass
-class BASNetModelOutput(ModelOutput):
+class BasNetSideOutput(ModelOutput):
  dout: torch.Tensor
  d1: Optional[torch.Tensor] = None
  d2: Optional[torch.Tensor] = None
@@ -25,6 +25,11 @@ class BASNetModelOutput(ModelOutput):
  db: Optional[torch.Tensor] = None
 
 
+@dataclass
+class BASNetModelOutput(ModelOutput):
+ activated: BasNetSideOutput
+
+
 class RefUnet(nn.Module):
  def __init__(self, in_ch: int, inc_ch: int) -> None:
  super().__init__()
@@ -466,27 +471,21 @@ class BASNetModel(PreTrainedModel):
  d6_act = torch.sigmoid(d6)
  db_act = torch.sigmoid(db)
 
+ side_outputs = (
+ dout_act,
+ d1_act,
+ d2_act,
+ d3_act,
+ d4_act,
+ d5_act,
+ d6_act,
+ db_act,
+ )
  if not return_dict:
- return (
- dout_act,
- d1_act,
- d2_act,
- d3_act,
- d4_act,
- d5_act,
- d6_act,
- db_act,
- )
+ return (side_outputs,)
 
  return BASNetModelOutput(
- dout=dout_act,
- d1=d1_act,
- d2=d2_act,
- d3=d3_act,
- d4=d4_act,
- d5=d5_act,
- d6=d6_act,
- db=db_act,
+ activated=BasNetSideOutput(*side_outputs),
  )