From 5c668e3c7b569a9cb26a18a98a10fba8962011b5 Mon Sep 17 00:00:00 2001
From: Falko Habel <falko.habel@gmx.de>
Date: Wed, 26 Mar 2025 22:01:47 +0100
Subject: [PATCH] updated to support new SMoe Version of AIIA

---
 setup.py                     | 2 +-
 src/aiunn/__init__.py        | 2 +-
 src/aiunn/upsampler/aiunn.py | 3 ++-
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/setup.py b/setup.py
index 4a4a835..baca931 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="aiunn",
-    version="0.1.1",
+    version="0.1.2",
     packages=find_packages(where="src"),
     package_dir={"": "src"},
     install_requires=[
diff --git a/src/aiunn/__init__.py b/src/aiunn/__init__.py
index a097c72..16f468f 100644
--- a/src/aiunn/__init__.py
+++ b/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.1.1"
\ No newline at end of file
+__version__ = "0.1.2"
\ No newline at end of file
diff --git a/src/aiunn/upsampler/aiunn.py b/src/aiunn/upsampler/aiunn.py
index a224e11..5628863 100644
--- a/src/aiunn/upsampler/aiunn.py
+++ b/src/aiunn/upsampler/aiunn.py
@@ -2,7 +2,8 @@ import os
 import torch
 import torch.nn as nn
 import warnings
-from aiia import AIIA, AIIAConfig, AIIABase
+from aiia import AIIAConfig, AIIABase
+from aiia.model.Model import AIIA
 from .config import aiuNNConfig
 import warnings