Update model.py

model.py

@@ -1,8 +1,16 @@
 import torch
 import torch.nn as nn
-import json
-from safetensors.torch import load_file
+from transformers import PreTrainedModel, PretrainedConfig
 
+# 1. СНАЧАЛА ОПРЕДЕЛЯЕМ КОНФИГ
+class AlphaDepthConfig(PretrainedConfig):
+    model_type = "alpha-depth"
+    
+    def __init__(self, input_size=[3, 128, 128], **kwargs):
+        self.input_size = input_size
+        super().__init__(**kwargs)
+
+# 2. ВСПОМОГАТЕЛЬНЫЕ БЛОКИ
 class ConvBlock(nn.Module):
     def __init__(self, in_c, out_c):
         super().__init__()
@@ -16,9 +24,12 @@ class ConvBlock(nn.Module):
         )
     def forward(self, x): return self.conv(x)
 
-class AlphaUNet(nn.Module):
-    def __init__(self):
-        super().__init__()
+# 3. САМА МОДЕЛЬ (Наследуемся от PreTrainedModel!)
+class AlphaUNet(PreTrainedModel):
+    config_class = AlphaDepthConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
         # Encoder
         self.enc1 = ConvBlock(3, 32)
         self.pool = nn.MaxPool2d(2)
@@ -47,17 +58,4 @@ class AlphaUNet(nn.Module):
         d1 = torch.cat([d1, e1], dim=1)
         d1 = self.dec1(d1)
         
-        return self.sigmoid(self.final(d1))
-
-    # Метод для удобной загрузки
-    @classmethod
-    def from_pretrained(cls, path="."):
-        model = cls()
-        # Грузим Safetensors
-        try:
-            state_dict = load_file(f"{path}/model.safetensors")
-            model.load_state_dict(state_dict)
-            print(">>> AlphaDepth loaded from .safetensors")
-        except FileNotFoundError:
-            print("Error: model.safetensors not found.")
-        return model
+        return self.sigmoid(self.final(d1))