Update modeling_upscaler.py

modeling_upscaler.py

@@ -96,38 +96,6 @@ class ESPCNExactUpsampler(nn.Module):
         return x
 
 
-class TwoStageSR(nn.Module):
-    def __init__(
-        self,
-        in_channels=3,
-        scale=4,
-        width=32,
-        num_blocks=3,
-        espcn_channels=64,
-    ):
-        super().__init__()
-        self.scale = scale
-        self.restoration = RestorationNet(
-            in_channels=in_channels,
-            width=width,
-            num_blocks=num_blocks,
-        )
-        self.upsampler = ESPCNExactUpsampler(
-            scale_factor=scale,
-            num_channels=in_channels,
-            channels=espcn_channels,
-        )
-
-    def forward(self, lr):
-        lr_clean = self.restoration(lr)
-        hr_pred = self.upsampler(lr_clean)
-        return hr_pred, lr_clean
-
-
-# -------------------------
-# Transformers output
-# -------------------------
-
 @dataclass
 class UpscalerOutput(ModelOutput):
     sr: torch.FloatTensor
@@ -141,16 +109,25 @@ class UpscalerModel(PreTrainedModel):
     def __init__(self, config: UpscalerConfig):
         super().__init__(config)
 
-        self.model = TwoStageSR(
+        # IMPORTANT :
+        # on expose directement les sous-modules pour que les clés
+        # correspondent au checkpoint Lightning :
+        # restoration.* / upsampler.*
+        self.scale = config.scale
+        self.restoration = RestorationNet(
             in_channels=config.in_channels,
-            scale=config.scale,
             width=config.width,
             num_blocks=config.num_blocks,
-            espcn_channels=config.espcn_channels,
+        )
+        self.upsampler = ESPCNExactUpsampler(
+            scale_factor=config.scale,
+            num_channels=config.in_channels,
+            channels=config.espcn_channels,
         )
 
         self.post_init()
 
     def forward(self, pixel_values: torch.FloatTensor, **kwargs) -> UpscalerOutput:
-        sr, restored = self.model(pixel_values)
+        restored = self.restoration(pixel_values)
+        sr = self.upsampler(restored)
         return UpscalerOutput(sr=sr, restored=restored)