diff --git a/.coveragerc b/.coveragerc new file mode 100644 index 0000000..7294ebb --- /dev/null +++ b/.coveragerc @@ -0,0 +1,10 @@ +[run] +branch = True +source = src +omit = + */tests/* + */migrations/* + +[report] +show_missing = True +fail_under = 80 \ No newline at end of file diff --git a/.gitea/workflows/embed.yaml b/.gitea/workflows/embed.yaml new file mode 100644 index 0000000..22c5bda --- /dev/null +++ b/.gitea/workflows/embed.yaml @@ -0,0 +1,37 @@ +name: Run VectorLoader Script + +on: + push: + branches: + - main + +jobs: + Explore-Gitea-Actions: + runs-on: ubuntu-latest + container: catthehacker/ubuntu:act-latest + steps: + - name: Check out repository code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v3 + with: + python-version: '3.11.7' + + - name: Clone additional repository + run: | + git config --global credential.helper cache + git clone https://fabel:${{ secrets.CICD }}@gitea.fabelous.app/fabel/VectorLoader.git + + - name: Install dependencies + run: | + cd VectorLoader + python -m pip install --upgrade pip + pip install -r requirements.txt + + - name: Run vectorizing + env: + VECTORDB_TOKEN: ${{ secrets.VECTORDB_TOKEN }} + run: | + cd VectorLoader + python -m src.run --full diff --git a/.gitea/workflows/test.yaml b/.gitea/workflows/test.yaml new file mode 100644 index 0000000..b754440 --- /dev/null +++ b/.gitea/workflows/test.yaml @@ -0,0 +1,37 @@ +name: Gitea Actions For AIIA +run-name: ${{ gitea.actor }} is testing out Gitea Actions 🚀 +on: [push] + +jobs: + Explore-Gitea-Actions: + runs-on: ubuntu-latest + container: catthehacker/ubuntu:act-latest + steps: + - name: Check out repository code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11.7' + + - name: Cache pip and model + uses: actions/cache@v3 + with: + path: | + ~/.cache/pip + ./fabel + key: ${{ runner.os }}-pip-model-${{ hashFiles('requirements.txt', 'requirements-dev.txt') }} + + restore-keys: | + ${{ runner.os }}-pip-model- + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install -r requirements-dev.txt + pip install -e . + - name: Run tests + run: | + pytest tests/ \ No newline at end of file diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..6925c8f --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,4 @@ +include LICENSE +include README.md +include requirements.txt +recursive-include src/aiia * \ No newline at end of file diff --git a/README.md b/README.md index 15bb162..aeb0837 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,130 @@ -# aiunn +# aiuNN - Advanced Image Upscaler using Neural Networks \ No newline at end of file +Adaptive Image Upscaler using Neural Networks + +## Overview + +`aiuNN` is an adaptive image upscaling model built on top of the Adaptive Image Intelligence Architecture (AIIA). This project provides fine-tuned versions of AIIA models specifically designed for high-quality image upscaling. By leveraging neural networks, `aiuNN` can significantly enhance the resolution and detail of images. + +## Features + +- **High-Quality Upscaling**: Achieve superior image quality with detailed and sharp outputs. +- **Fine-Tuned Models**: Pre-trained on a diverse dataset to ensure optimal performance. +- **Easy Integration**: Simple API for integrating upscaling capabilities into your applications. +- **Customizable**: Fine-tune the models further on your own datasets for specific use cases. + +## Installation + +You can install `aiuNN` using pip. Run the following command: + +```sh +pip install git+https://gitea.fabelous.app/Machine-Learning/aiuNN.git +``` + +## Usage + +Here's a basic example of how to use `aiuNN` for image upscaling: + +```python src/main.py +from aiia import AIIABase +from aiunn import aiuNN, aiuNNTrainer +import pandas as pd +from torchvision import transforms + +# Load your base model and upscaler +pretrained_model_path = "path/to/aiia/model" +base_model = AIIABase.load(pretrained_model_path, precision="bf16") +upscaler = aiuNN(base_model) + +# Create trainer with your dataset class +trainer = aiuNNTrainer(upscaler, dataset_class=UpscaleDataset) + +# Load data using parameters for your dataset +dataset_params = { + 'parquet_files': [ + "path/to/dataset1", + "path/to/dataset2" + ], + 'transform': transforms.Compose([transforms.ToTensor()]), + 'samples_per_file': 5000 # Your training samples you want to load per file +} +trainer.load_data(dataset_params=dataset_params, batch_size=1) + +# Fine-tune the model +trainer.finetune(output_path="trained_model") +``` + +## Dataset + +The `UpscaleDataset` class is designed to handle Parquet files containing image data. It loads a subset of images from each file and validates the data types to ensure consistency. + +This is an example dataset that you can use with the AIIUN model: + +```python src/example.py +class UpscaleDataset(Dataset): + def __init__(self, parquet_files: list, transform=None, samples_per_file=10_000): + combined_df = pd.DataFrame() + for parquet_file in parquet_files: + # Load a subset from each parquet file + df = pd.read_parquet(parquet_file, columns=['image_410', 'image_820']).head(samples_per_file) + combined_df = pd.concat([combined_df, df], ignore_index=True) + + # Validate rows (ensuring each value is bytes or str) + self.df = combined_df.apply(self._validate_row, axis=1) + self.transform = transform + self.failed_indices = set() + + def _validate_row(self, row): + for col in ['image_410', 'image_820']: + if not isinstance(row[col], (bytes, str)): + raise ValueError(f"Invalid data type in column {col}: {type(row[col])}") + return row + + def _decode_image(self, data): + try: + if isinstance(data, str): + return base64.b64decode(data) + elif isinstance(data, bytes): + return data + raise ValueError(f"Unsupported data type: {type(data)}") + except Exception as e: + raise RuntimeError(f"Decoding failed: {str(e)}") + + def __len__(self): + return len(self.df) + + def __getitem__(self, idx): + # If previous call failed for this index, use a different index + if idx in self.failed_indices: + return self[(idx + 1) % len(self)] + try: + row = self.df.iloc[idx] + low_res_bytes = self._decode_image(row['image_410']) + high_res_bytes = self._decode_image(row['image_820']) + ImageFile.LOAD_TRUNCATED_IMAGES = True + # Open image bytes with Pillow and convert to RGBA first + low_res_rgba = Image.open(io.BytesIO(low_res_bytes)).convert('RGBA') + high_res_rgba = Image.open(io.BytesIO(high_res_bytes)).convert('RGBA') + + # Create a new RGB image with black background + low_res_rgb = Image.new("RGB", low_res_rgba.size, (0, 0, 0)) + high_res_rgb = Image.new("RGB", high_res_rgba.size, (0, 0, 0)) + + # Composite the original image over the black background + low_res_rgb.paste(low_res_rgba, mask=low_res_rgba.split()[3]) + high_res_rgb.paste(high_res_rgba, mask=high_res_rgba.split()[3]) + + # Now we have true 3-channel RGB images with transparent areas converted to black + low_res = low_res_rgb + high_res = high_res_rgb + + # If a transform is provided (e.g. conversion to Tensor), apply it + if self.transform: + low_res = self.transform(low_res) + high_res = self.transform(high_res) + return low_res, high_res + except Exception as e: + print(f"\nError at index {idx}: {str(e)}") + self.failed_indices.add(idx) + return self[(idx + 1) % len(self)] +``` diff --git a/example.py b/example.py index 58470b0..66764a4 100644 --- a/example.py +++ b/example.py @@ -15,7 +15,7 @@ class UpscaleDataset(Dataset): combined_df = pd.DataFrame() for parquet_file in parquet_files: # Load a subset from each parquet file - df = pd.read_parquet(parquet_file, columns=['image_512', 'image_1024']).head(samples_per_file) + df = pd.read_parquet(parquet_file, columns=['image_410', 'image_820']).head(samples_per_file) combined_df = pd.concat([combined_df, df], ignore_index=True) # Validate rows (ensuring each value is bytes or str) @@ -24,7 +24,7 @@ class UpscaleDataset(Dataset): self.failed_indices = set() def _validate_row(self, row): - for col in ['image_512', 'image_1024']: + for col in ['image_410', 'image_820']: if not isinstance(row[col], (bytes, str)): raise ValueError(f"Invalid data type in column {col}: {type(row[col])}") return row @@ -48,8 +48,8 @@ class UpscaleDataset(Dataset): return self[(idx + 1) % len(self)] try: row = self.df.iloc[idx] - low_res_bytes = self._decode_image(row['image_512']) - high_res_bytes = self._decode_image(row['image_1024']) + low_res_bytes = self._decode_image(row['image_410']) + high_res_bytes = self._decode_image(row['image_820']) ImageFile.LOAD_TRUNCATED_IMAGES = True # Open image bytes with Pillow and convert to RGBA first low_res_rgba = Image.open(io.BytesIO(low_res_bytes)).convert('RGBA') @@ -66,11 +66,7 @@ class UpscaleDataset(Dataset): # Now we have true 3-channel RGB images with transparent areas converted to black low_res = low_res_rgb high_res = high_res_rgb - - # Resize the images to reduce VRAM usage - low_res = low_res.resize((410, 410), Image.LANCZOS) - high_res = high_res.resize((820, 820), Image.LANCZOS) - + # If a transform is provided (e.g. conversion to Tensor), apply it if self.transform: low_res = self.transform(low_res) @@ -98,7 +94,7 @@ if __name__ =="__main__": "/root/training_data/vision-dataset/image_vec_upscaler.parquet" ], 'transform': transforms.Compose([transforms.ToTensor()]), - 'samples_per_file': 5000 + 'samples_per_file': 20_000 } trainer.load_data(dataset_params=dataset_params, batch_size=1) diff --git a/input.jpg b/input.jpg deleted file mode 100644 index 0426a63..0000000 Binary files a/input.jpg and /dev/null differ diff --git a/pyproject.toml b/pyproject.toml index b0e5e10..6be914a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,6 @@ [build-system] requires = ["setuptools>=45", "wheel"] build-backend = "setuptools.build_meta" - [project] name = "aiunn" version = "0.1.1" @@ -10,8 +9,7 @@ readme = "README.md" requires-python = ">=3.10" license = {file = "LICENSE"} authors = [ - {name = "Falko Habel", email = "falko.habel@gmx.de"}, +{name = "Falko Habel", email = "falko.habel@gmx.de"}, ] - [project.urls] -"Homepage" = "https://gitea.fabelous.app/Machine-Learning/aiuNN" \ No newline at end of file +"Homepage"="https://gitea.fabelous.app/Machine-Learning/aiuNN" \ No newline at end of file diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..fa7df11 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +testpaths = tests/ +python_files = test_*.py \ No newline at end of file diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..1d6ed5c --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,2 @@ +pytest +pytest-mock \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 8e47744..0f9fefa 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,5 @@ torch aiia pillow torchvision -sklearn \ No newline at end of file +scikit-learn +git+https://gitea.fabelous.app/Machine-Learning/AIIA.git \ No newline at end of file diff --git a/setup.py b/setup.py index 4a4a835..b6c17c3 100644 --- a/setup.py +++ b/setup.py @@ -2,13 +2,18 @@ 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=[ - line.strip() - for line in open("requirements.txt") - if line.strip() and not line.startswith("#") + "torch", + "aiia", + "pillow", + "torchvision", + "scikit-learn", + ], + dependency_links=[ + "git+https://gitea.fabelous.app/Machine-Learning/AIIA.git#egg=aiia" ], python_requires=">=3.10", ) \ No newline at end of file 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 ca7ba2b..ecb3c21 100644 --- a/src/aiunn/upsampler/aiunn.py +++ b/src/aiunn/upsampler/aiunn.py @@ -2,50 +2,120 @@ import os import torch import torch.nn as nn import warnings -from aiia import AIIA, AIIAConfig, AIIABase +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 - self.upsample = nn.Upsample( - scale_factor=self.config.upsample_scale, - mode=self.config.upsample_mode, - align_corners=self.config.upsample_align_corners - ) - # Conversion layer: change from hidden size channels to 3 channels. - self.to_rgb = nn.Conv2d( + # Enhanced approach + scale_factor = self.config.upsample_scale + out_channels = self.base_model.config.num_channels * (scale_factor ** 2) + self.pixel_shuffle_conv = nn.Conv2d( in_channels=self.base_model.config.hidden_size, - out_channels=3, - kernel_size=1 + out_channels=out_channels, + kernel_size=self.base_model.config.kernel_size, + padding=1 ) + self.pixel_shuffle = nn.PixelShuffle(scale_factor) 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): - # Load the configuration from disk. - config = AIIAConfig.load(path) - # Reconstruct the base model from the loaded configuration. - base_model = AIIABase(config) - # Instantiate the Upsampler using the proper base model. - upsampler = cls(base_model) + def load(cls, path, precision: str = None, **kwargs): + """ + Load a aiuNN model from disk with automatic detection of base model type. - # Load state dict and handle precision conversion if needed. + Args: + path (str): Directory containing the stored configuration and model parameters. + precision (str, optional): Desired precision for the model's parameters. + **kwargs: Additional keyword arguments to override configuration parameters. + + Returns: + An instance of aiuNN with loaded weights. + """ + # Load the configuration + config = aiuNNConfig.load(path) + + # Determine the device device = 'cuda' if torch.cuda.is_available() else 'cpu' - state_dict = torch.load(f"{path}/model.pth", map_location=device) + + # Load the state dictionary + state_dict = torch.load(os.path.join(path, "model.pth"), map_location=device) + + # Import all model types + 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): + if any(f"{keys_prefix}.shared_layer" in key for key in state_dict.keys()): + return AIIABaseShared + else: + return AIIABase + + # Detect base model type + base_model = None + + # Check for AIIAmoe with multiple experts + if any("base_model.experts" in key for key in state_dict.keys()): + # Count the number of experts + max_expert_idx = -1 + for key in state_dict.keys(): + if "base_model.experts." in key: + try: + parts = key.split("base_model.experts.")[1].split(".") + expert_idx = int(parts[0]) + max_expert_idx = max(max_expert_idx, expert_idx) + except (ValueError, IndexError): + pass + + if max_expert_idx >= 0: + # Determine the type of base_cnn each expert is using + base_class_for_experts = detect_base_class_type("base_model.experts.0.base_cnn") + + # Create AIIAmoe with the detected expert count and base class + base_model = AIIAmoe(config, num_experts=max_expert_idx+1, base_class=base_class_for_experts, **kwargs) + + # Check for AIIAchunked or AIIArecursive + elif any("base_model.chunked_cnn" in key for key in state_dict.keys()): + if any("recursion_depth" in key for key in state_dict.keys()): + # This is an AIIArecursive model + base_class = detect_base_class_type("base_model.chunked_cnn.base_cnn") + base_model = AIIArecursive(config, base_class=base_class, **kwargs) + else: + # This is an AIIAchunked model + base_class = detect_base_class_type("base_model.chunked_cnn.base_cnn") + base_model = AIIAchunked(config, base_class=base_class, **kwargs) + + # Check for AIIAExpert + elif any("base_model.base_cnn" in key for key in state_dict.keys()): + # Determine which base class the expert is using + base_class = detect_base_class_type("base_model.base_cnn") + base_model = AIIAExpert(config, base_class=base_class, **kwargs) + + # If none of the above, use AIIABase or AIIABaseShared directly + else: + base_class = detect_base_class_type("base_model") + base_model = base_class(config, **kwargs) + + # Create the aiuNN model with the detected base model + model = cls(base_model, config=base_model.config) + + # Handle precision conversion + dtype = None if precision is not None: if precision.lower() == 'fp16': dtype = torch.float16 @@ -57,12 +127,15 @@ class aiuNN(AIIA): dtype = torch.bfloat16 else: raise ValueError("Unsupported precision. Use 'fp16', 'bf16', or leave as None.") - + + if dtype is not None: for key, param in state_dict.items(): if torch.is_tensor(param): state_dict[key] = param.to(dtype) - upsampler.load_state_dict(state_dict) - return upsampler + + # Load the state dict + model.load_state_dict(state_dict) + return model @@ -70,13 +143,14 @@ 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") + upsampler.save("aiunn") # Now load using the overridden load method; this will load the complete model. - upsampler_loaded = aiuNN.load("hehe", precision="bf16") + upsampler_loaded = aiuNN.load("aiunn", precision="bf16") print("Updated configuration:", upsampler_loaded.config.__dict__) diff --git a/tests/finetune/test_trainer.py b/tests/finetune/test_trainer.py new file mode 100644 index 0000000..773f973 --- /dev/null +++ b/tests/finetune/test_trainer.py @@ -0,0 +1,75 @@ +import pytest +import torch +import torch.nn as nn +from aiunn import aiuNNTrainer + +# Simple mock dataset +class MockDataset(torch.utils.data.Dataset): + def __init__(self, num_samples=10): + self.num_samples = num_samples + + def __len__(self): + return self.num_samples + + def __getitem__(self, idx): + return torch.randn(3, 64, 64), torch.randn(3, 128, 128) + +# Simple mock model +class MockModel(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(3, 3, 3, padding=1) + + def forward(self, x): + return self.conv(x) + + def save(self, path): + pass # Mock save method + +@pytest.fixture +def trainer(): + model = MockModel() + return aiuNNTrainer(model, dataset_class=MockDataset) + +def test_trainer_initialization(trainer): + """Test basic trainer initialization""" + assert trainer.model is not None + assert isinstance(trainer.criterion, nn.MSELoss) + assert trainer.optimizer is None + assert trainer.device in [torch.device('cuda'), torch.device('cpu')] + +def test_load_data_basic(trainer): + """Test basic data loading""" + train_loader, val_loader = trainer.load_data( + dataset_params={'num_samples': 10}, + batch_size=2, + validation_split=0.2 + ) + + assert train_loader is not None + assert val_loader is not None + assert len(train_loader) > 0 + assert len(val_loader) > 0 + +def test_load_custom_datasets(trainer): + """Test loading custom datasets""" + train_dataset = MockDataset(num_samples=10) + val_dataset = MockDataset(num_samples=5) + + train_loader, val_loader = trainer.load_data( + custom_train_dataset=train_dataset, + custom_val_dataset=val_dataset, + batch_size=2 + ) + + assert train_loader is not None + assert val_loader is not None + assert len(train_loader) == 5 # 10 samples with batch size 2 + assert len(val_loader) == 3 # 5 samples with batch size 2 (rounded up) + +def test_error_no_dataset(): + """Test error when no dataset is provided""" + trainer = aiuNNTrainer(MockModel(), dataset_class=None) + + with pytest.raises(ValueError): + trainer.load_data(dataset_params={}) \ No newline at end of file diff --git a/tests/inference/test_inference.py b/tests/inference/test_inference.py new file mode 100644 index 0000000..39dea9a --- /dev/null +++ b/tests/inference/test_inference.py @@ -0,0 +1,99 @@ +# tests/inference/test_inference.py +import pytest +import numpy as np +import torch +from PIL import Image +from aiunn import aiuNNInference +from aiunn.upsampler.aiunn import aiuNN, aiuNNConfig +from aiia import AIIABase, AIIAConfig +import os +import json +from unittest.mock import patch, MagicMock, mock_open + +@pytest.fixture +def real_model(tmp_path): + # Create temporary directory for model + model_dir = tmp_path / "model" + model_dir.mkdir() + + config = AIIAConfig() + ai_config = aiuNNConfig() + base_model = AIIABase(config) + + # Make sure aiuNN is properly configured with all required attributes + upsampler = aiuNN(base_model, config=ai_config) + # Ensure the upsample attribute is properly set if needed + # upsampler.upsample = ... # Add any necessary initialization + + # Save the model and config to temporary directory + save_path = str(model_dir / "save") + os.makedirs(save_path, exist_ok=True) + + # Save config file + config_data = { + "model_type": "test_model", + "scale": 4, + "in_channels": 3, + "out_channels": 3 + } + with open(os.path.join(save_path, "config.json"), "w") as f: + json.dump(config_data, f) + + # Save model + upsampler.save(save_path) + + # Load model in inference mode + inference_model = aiuNNInference(model_path=save_path, precision='fp16', device='cpu') + return inference_model + + +@pytest.fixture +def inference(real_model): + return real_model + +def test_preprocess_image(inference): + # Create a small test image + test_array = np.zeros((100, 100, 3), dtype=np.uint8) + test_image = Image.fromarray(test_array) + + # Test with PIL Image + result = inference.preprocess_image(test_image) + assert isinstance(result, torch.Tensor) + assert result.shape[0] == 1 # batch dimension + assert result.shape[1] == 3 # channels + assert result.shape[2:] == (100, 100) # height, width + +def test_postprocess_tensor(inference): + # Create a test tensor + test_tensor = torch.zeros(1, 3, 100, 100) + + result = inference.postprocess_tensor(test_tensor) + assert isinstance(result, Image.Image) + assert result.size == (100, 100) + assert result.mode == 'RGB' + +def test_save(inference): + # Create a test image + test_image = Image.fromarray(np.zeros((100, 100, 3), dtype=np.uint8)) + output_path = "test_output.png" + + with patch('os.makedirs') as mock_makedirs: + inference.save(test_image, output_path) + mock_makedirs.assert_called_with(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) + +def test_convert_to_binary(inference): + # Create a test image + test_image = Image.fromarray(np.zeros((100, 100, 3), dtype=np.uint8)) + + result = inference.convert_to_binary(test_image) + assert isinstance(result, bytes) + assert len(result) > 0 + +def test_process_batch(inference): + # Create test images + test_array = np.zeros((100, 100, 3), dtype=np.uint8) + test_images = [Image.fromarray(test_array) for _ in range(2)] + + results = inference.process_batch(test_images) + assert len(results) == 2 + assert all(isinstance(img, Image.Image) for img in results) \ No newline at end of file diff --git a/tests/upsampler/test_aiunn.py b/tests/upsampler/test_aiunn.py new file mode 100644 index 0000000..aae0813 --- /dev/null +++ b/tests/upsampler/test_aiunn.py @@ -0,0 +1,48 @@ +import os +import tempfile +from aiia import AIIABase, AIIAConfig +from aiunn import aiuNN, aiuNNConfig + +def test_save_and_load_model(): + # Create a temporary directory to save the model + with tempfile.TemporaryDirectory() as tmpdirname: + # Create configurations and build a base model + config = AIIAConfig() + ai_config = aiuNNConfig() + base_model = AIIABase(config) + upsampler = aiuNN(base_model, config=ai_config) + + # Save the model + save_path = os.path.join(tmpdirname, "model") + upsampler.save(save_path) + + # Load the model + loaded_upsampler = aiuNN.load(save_path) + + # Verify that the loaded model is the same as the original model + assert isinstance(loaded_upsampler, aiuNN) + assert loaded_upsampler.config.__dict__ == upsampler.config.__dict__ + +def test_save_and_load_model_with_precision(): + # Create a temporary directory to save the model + with tempfile.TemporaryDirectory() as tmpdirname: + # Create configurations and build a base model + config = AIIAConfig() + ai_config = aiuNNConfig() + base_model = AIIABase(config) + upsampler = aiuNN(base_model, config=ai_config) + + # Save the model + save_path = os.path.join(tmpdirname, "model") + upsampler.save(save_path) + + # Load the model with precision 'bf16' + loaded_upsampler = aiuNN.load(save_path, precision="bf16") + + # Verify that the loaded model is the same as the original model + assert isinstance(loaded_upsampler, aiuNN) + assert loaded_upsampler.config.__dict__ == upsampler.config.__dict__ + +if __name__ == "__main__": + test_save_and_load_model() + test_save_and_load_model_with_precision() \ No newline at end of file diff --git a/tests/upsampler/test_config.py b/tests/upsampler/test_config.py new file mode 100644 index 0000000..3eb2779 --- /dev/null +++ b/tests/upsampler/test_config.py @@ -0,0 +1,45 @@ +import pytest +from aiunn import aiuNNConfig + +def test_default_initialization(): + config = aiuNNConfig() + assert config.upsample_scale == 2 + assert config.upsample_mode == 'bilinear' + assert not config.upsample_align_corners + assert len(config.layers) == 1 + assert config.layers[0]['name'] == 'Upsample' + +def test_custom_initialization(): + custom_config = { + 'some_key': 'some_value', + } + config = aiuNNConfig(base_config=custom_config, upsample_scale=3, upsample_mode='nearest', upsample_align_corners=True) + assert config.upsample_scale == 3 + assert config.upsample_mode == 'nearest' + assert config.upsample_align_corners + assert len(config.layers) == 1 + assert config.layers[0]['name'] == 'Upsample' + +def test_add_upsample_layer(): + config = aiuNNConfig() + config.add_upsample_layer() + assert len(config.layers) == 1 + +def test_upsample_layer_not_duplicated(): + config = aiuNNConfig() + initial_length = len(config.layers) + # Remove all existing 'Upsample' layers. + for layer in list(config.layers): + if layer.get('name') == 'Upsample': + config.layers.remove(layer) + config.add_upsample_layer() + assert len(config.layers) == 1 + +def test_base_config_with_to_dict(): + class MockBaseConfig: + def to_dict(self): + return {'base_key': 'base_value'} + + base_config = MockBaseConfig() + config = aiuNNConfig(base_config=base_config) + assert config.to_dict()['base_key'] == 'base_value' \ No newline at end of file