Merge pull request 'feat/model_fix' (#17) from feat/model_fix into develop

Reviewed-on: #17

MANIFEST.in

@@ -1,4 +1,4 @@
 include LICENSE
 include README.md
 include requirements.txt
-recursive-include src/aiia *
+recursive-include src/aiunn *

pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools>=45", "wheel"]
 build-backend = "setuptools.build_meta"
 [project]
 name = "aiunn"
-version = "0.1.1"
+version = "0.2.2"
 description = "Finetuner for image upscaling using AIIA"
 readme = "README.md"
 requires-python = ">=3.10"

requirements.txt

@@ -1,6 +1,6 @@
 torch
-aiia
 pillow
+pandas
 torchvision
 scikit-learn
 git+https://gitea.fabelous.app/Machine-Learning/AIIA.git

setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="aiunn",
-    version="0.2.1",
+    version="0.2.2",
     packages=find_packages(where="src"),
     package_dir={"": "src"},
     install_requires=[

src/aiunn/__init__.py

@@ -3,4 +3,4 @@ from .upsampler.aiunn import aiuNN
 from .upsampler.config import aiuNNConfig
 from .inference.inference import aiuNNInference
 
-__version__ = "0.2.1"
+__version__ = "0.2.2"

src/aiunn/upsampler/aiunn.py

@@ -32,7 +32,18 @@ class aiuNN(PreTrainedModel):
     def forward(self, x):
         if self.base_model is None:
             raise ValueError("Base model is not loaded. Call 'load_base_model' before forwarding.")
-        x = self.base_model(x)  # Get base features
+            
+        # Get base features - we need to extract the last hidden state if it's returned as part of a tuple/dict
+        base_output = self.base_model(x)
+        if isinstance(base_output, tuple):
+            x = base_output[0]
+        elif isinstance(base_output, dict):
+            x = base_output.get('last_hidden_state', base_output.get('hidden_states'))
+            if x is None:
+                raise ValueError("Expected 'last_hidden_state' or 'hidden_states' in model output")
+        else:
+            x = base_output
+            
         x = self.pixel_shuffle_conv(x)  # Expand channels for shuffling
         x = self.pixel_shuffle(x)  # Rearrange channels into spatial dimensions
         return x