From 5c668e3c7b569a9cb26a18a98a10fba8962011b5 Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Wed, 26 Mar 2025 22:01:47 +0100 Subject: [PATCH 1/4] 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 From 399a7c0f693cf5e4b82a692cbe31e9759395638a Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Wed, 26 Mar 2025 22:01:59 +0100 Subject: [PATCH 2/4] increamentet release version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index b0e5e10..10e9c2a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "aiunn" -version = "0.1.1" +version = "0.1.2" description = "Finetuner for image upscaling using AIIA" readme = "README.md" requires-python = ">=3.10" From 68a27f00c121f9e84e8c4f6de2fc8305efc5fcc5 Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Wed, 26 Mar 2025 22:02:11 +0100 Subject: [PATCH 3/4] included requirement for aiia --- requirements.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 8e47744..6b388e0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,5 @@ torch aiia pillow torchvision -sklearn \ No newline at end of file +sklearn +https://gitea.fabelous.app/Machine-Learning/AIIA.git \ No newline at end of file From de0da5de82a86854ca7087a811ce74e2d0983e31 Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Thu, 27 Mar 2025 17:18:40 +0100 Subject: [PATCH 4/4] improved and simplified config --- src/aiunn/upsampler/aiunn.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/src/aiunn/upsampler/aiunn.py b/src/aiunn/upsampler/aiunn.py index 5628863..b0c26b2 100644 --- a/src/aiunn/upsampler/aiunn.py +++ b/src/aiunn/upsampler/aiunn.py @@ -2,19 +2,18 @@ import os import torch import torch.nn as nn import warnings -from aiia import AIIAConfig, AIIABase -from aiia.model.Model import AIIA +from aiia.model.Model import AIIA, AIIAConfig, AIIABase from .config import aiuNNConfig import warnings class aiuNN(AIIA): - def __init__(self, base_model: AIIABase): + def __init__(self, base_model: AIIA, config:aiuNNConfig): super().__init__(base_model.config) self.base_model = base_model # Pass the unified base configuration using the new parameter. - self.config = aiuNNConfig(base_config=base_model.config) + self.config = config # Enhanced approach scale_factor = self.config.upsample_scale @@ -29,11 +28,12 @@ class aiuNN(AIIA): def forward(self, x): - x = self.base_model(x) - x = self.upsample(x) - x = self.to_rgb(x) # Ensures output has 3 channels. + x = self.base_model(x) # Get base features + x = self.pixel_shuffle_conv(x) # Expand channels for shuffling + x = self.pixel_shuffle(x) # Rearrange channels into spatial dimensions return x + @classmethod def load(cls, path, precision: str = None, **kwargs): """ @@ -57,7 +57,7 @@ class aiuNN(AIIA): state_dict = torch.load(os.path.join(path, "model.pth"), map_location=device) # Import all model types - from aiia import AIIABase, AIIABaseShared, AIIAExpert, AIIAmoe, AIIAchunked, AIIArecursive + from aiia.model.Model import AIIABase, AIIABaseShared, AIIAExpert, AIIAmoe, AIIAchunked, AIIArecursive # Helper function to detect base class type from key patterns def detect_base_class_type(keys_prefix): @@ -112,7 +112,7 @@ class aiuNN(AIIA): base_model = base_class(config, **kwargs) # Create the aiuNN model with the detected base model - model = cls(base_model) + model = cls(base_model, config=base_model.config) # Handle precision conversion dtype = None @@ -143,9 +143,10 @@ if __name__ == "__main__": from aiia import AIIABase, AIIAConfig # Create a configuration and build a base model. config = AIIAConfig() + ai_config = aiuNNConfig() base_model = AIIABase(config) # Instantiate Upsampler from the base model (works correctly). - upsampler = aiuNN(base_model) + upsampler = aiuNN(base_model, config=ai_config) # Save the model (both configuration and weights). upsampler.save("hehe")