Merge pull request 'develop' (#11) from develop into main

Reviewed-on: #11

.gitea/workflows/embed.yaml

@@ -34,4 +34,4 @@ jobs:
           VECTORDB_TOKEN: ${{ secrets.VECTORDB_TOKEN }}
         run: |
           cd VectorLoader
-          python -m src.run                           
+          python -m src.run --full                              

setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="aiunn",
-    version="0.2.0",
+    version="0.1.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.0"
+__version__ = "0.1.2"

src/aiunn/inference/inference.py

@@ -12,13 +12,13 @@ class aiuNNInference:
     Inference class for aiuNN upsampling model.
     Handles model loading, image upscaling, and output processing.
     """
-    def __init__(self, model_path: str, device: Optional[str] = None):
+    def __init__(self, model_path: str, precision: Optional[str] = None, device: Optional[str] = None):
         """
         Initialize the inference class by loading the aiuNN model.
         
         Args:
             model_path: Path to the saved model directory
-
+            precision: Optional precision setting ('fp16', 'bf16', or None for default)
             device: Optional device specification ('cuda', 'cpu', or None for auto-detection)
         """
         
@@ -30,7 +30,7 @@ class aiuNNInference:
             self.device = device
             
         # Load the model with specified precision
-        self.model = aiuNN.from_pretrained(model_path)
+        self.model = aiuNN.load(model_path, precision=precision)
         self.model.to(self.device)
         self.model.eval()
         
@@ -160,11 +160,54 @@ class aiuNNInference:
         
         return binary_data
     
+    def process_batch(self, 
+                      images: List[Union[str, Image.Image]], 
+                      output_dir: Optional[str] = None,
+                      save_format: str = 'PNG',
+                      return_binary: bool = False) -> Union[List[Image.Image], List[bytes], None]:
+        """
+        Process multiple images in batch.
+        
+        Args:
+            images: List of input images (paths or PIL Images)
+            output_dir: Optional directory to save results
+            save_format: Format to use when saving images
+            return_binary: Whether to return binary data instead of PIL Images
+            
+        Returns:
+            List of processed images or binary data, or None if only saving
+        """
+        results = []
+        
+        for i, img in enumerate(images):
+            # Upscale the image
+            upscaled = self.upscale(img)
+            
+            # Save if output directory is provided
+            if output_dir:
+                # Extract filename if input is a path
+                if isinstance(img, str):
+                    filename = os.path.basename(img)
+                    base, _ = os.path.splitext(filename)
+                else:
+                    base = f"upscaled_{i}"
+                
+                output_path = os.path.join(output_dir, f"{base}.{save_format.lower()}")
+                self.save(upscaled, output_path, format=save_format)
+            
+            # Add to results based on return type
+            if return_binary:
+                results.append(self.convert_to_binary(upscaled, format=save_format))
+            else:
+                results.append(upscaled)
+        
+        return results if (not output_dir or return_binary or not save_format) else None
+
 
 # Example usage (can be removed)
 if __name__ == "__main__":
     # Initialize inference with a model path
-    inferencer = aiuNNInference("path/to/model")
+    inferencer = aiuNNInference("path/to/model", precision="bf16")
     
     # Upscale a single image
     upscaled_image = inferencer.upscale("input_image.jpg")
@@ -174,4 +217,10 @@ if __name__ == "__main__":
     
     # Convert to binary
     binary_data = inferencer.convert_to_binary(upscaled_image)
-    
+    
+    # Process a batch of images
+    inferencer.process_batch(
+        ["image1.jpg", "image2.jpg"], 
+        output_dir="output_folder",
+        save_format="PNG"
+    )

src/aiunn/upsampler/aiunn.py

@@ -2,19 +2,19 @@ import os
 import torch
 import torch.nn as nn
 import warnings
-from aiia.model.Model import AIIAConfig, AIIABase
+from aiia.model.Model import AIIA, AIIAConfig, AIIABase
-from transformers import PreTrainedModel
 from .config import aiuNNConfig
 import warnings
 
 
-class aiuNN(PreTrainedModel):
+class aiuNN(AIIA):
-    config_class = aiuNNConfig
+    def __init__(self, base_model: AIIA, config:aiuNNConfig):
-    def __init__(self, config: aiuNNConfig):
+        super().__init__(base_model.config)
-        super().__init__(config)
+        self.base_model = base_model
+        
         # Pass the unified base configuration using the new parameter.
         self.config = config
-
+        
         # Enhanced approach
         scale_factor = self.config.upsample_scale
         out_channels = self.base_model.config.num_channels * (scale_factor ** 2)
@@ -26,18 +26,118 @@ class aiuNN(PreTrainedModel):
         )
         self.pixel_shuffle = nn.PixelShuffle(scale_factor)
 
-    def load_base_model(self, base_model: PreTrainedModel):
+
-        self.base_model = base_model
-        
     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
         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):
+        """
+        Load a aiuNN model from disk with automatic detection of base model type.
+        
+        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'
+        
+        # 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
+            elif precision.lower() == 'bf16':
+                if device == 'cuda' and not torch.cuda.is_bf16_supported():
+                    warnings.warn("BF16 is not supported on this GPU. Falling back to FP16.")
+                    dtype = torch.float16
+                else:
+                    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)
+        
+        # Load the state dict
+        model.load_state_dict(state_dict)
+        return model
+
+
 
 if __name__ == "__main__":
     from aiia import AIIABase, AIIAConfig
@@ -46,11 +146,11 @@ if __name__ == "__main__":
     ai_config = aiuNNConfig()
     base_model = AIIABase(config)
     # Instantiate Upsampler from the base model (works correctly).
-    upsampler = aiuNN(config=ai_config)
+    upsampler = aiuNN(base_model, config=ai_config)
-    upsampler.load_base_model(base_model)
+
     # Save the model (both configuration and weights).
-    upsampler.save_pretrained("aiunn")
+    upsampler.save("aiunn")
 
     # Now load using the overridden load method; this will load the complete model.
-    upsampler_loaded = aiuNN.from_pretrained("aiunn")
+    upsampler_loaded = aiuNN.load("aiunn", precision="bf16")
     print("Updated configuration:", upsampler_loaded.config.__dict__)

tests/inference/test_inference.py

@@ -21,8 +21,9 @@ def real_model(tmp_path):
     base_model = AIIABase(config)
     
     # Make sure aiuNN is properly configured with all required attributes
-    upsampler = aiuNN(config=ai_config)
+    upsampler = aiuNN(base_model, config=ai_config)
-    upsampler.load_base_model(base_model)
+    # 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")
@@ -39,10 +40,10 @@ def real_model(tmp_path):
         json.dump(config_data, f)
     
     # Save model
-    upsampler.save_pretrained(save_path)
+    upsampler.save(save_path)
     
     # Load model in inference mode
-    inference_model = aiuNNInference(model_path=save_path, device='cpu')
+    inference_model = aiuNNInference(model_path=save_path, precision='fp16', device='cpu')
     return inference_model
 
 
@@ -87,3 +88,12 @@ def test_convert_to_binary(inference):
     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)

tests/upsampler/test_aiunn.py

@@ -10,21 +10,39 @@ def test_save_and_load_model():
         config = AIIAConfig()
         ai_config = aiuNNConfig()
         base_model = AIIABase(config)
-        upsampler = aiuNN(config=ai_config)
+        upsampler = aiuNN(base_model, config=ai_config)
-        upsampler.load_base_model(base_model)
+
         # Save the model
         save_path = os.path.join(tmpdirname, "model")
-        upsampler.save_pretrained(save_path)
+        upsampler.save(save_path)
 
         # Load the model
-        loaded_upsampler = aiuNN.from_pretrained(save_path)
+        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.hidden_size == upsampler.config.hidden_size
+        assert loaded_upsampler.config.__dict__ == upsampler.config.__dict__
-        assert loaded_upsampler.config._activation_function == upsampler.config._activation_function
-        assert loaded_upsampler.config.architectures == upsampler.config.architectures
 
+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()