increaed software version

.gitea/workflows/embed.yaml

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

setup.py

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

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, precision: Optional[str] = None, device: Optional[str] = None):
+    def __init__(self, model_path: str, 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.load(model_path, precision=precision)
+        self.model = aiuNN.from_pretrained(model_path)
         self.model.to(self.device)
         self.model.eval()
         
@@ -160,54 +160,11 @@ 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", precision="bf16")
+    inferencer = aiuNNInference("path/to/model")
     
     # Upscale a single image
     upscaled_image = inferencer.upscale("input_image.jpg")
@@ -218,9 +175,3 @@ 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,16 +2,16 @@ import os
 import torch
 import torch.nn as nn
 import warnings
-from aiia.model.Model import AIIA, AIIAConfig, AIIABase
+from aiia.model.Model import AIIAConfig, AIIABase
+from transformers import PreTrainedModel
 from .config import aiuNNConfig
 import warnings
 
 
-class aiuNN(AIIA):
+class aiuNN(PreTrainedModel):
-    def __init__(self, base_model: AIIA, config:aiuNNConfig):
+    config_class = aiuNNConfig
-        super().__init__(base_model.config)
+    def __init__(self, config: aiuNNConfig):
-        self.base_model = base_model
+        super().__init__(config)
-        
         # Pass the unified base configuration using the new parameter.
         self.config = config
 
@@ -26,118 +26,18 @@ class aiuNN(AIIA):
         )
         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
@@ -146,11 +46,11 @@ if __name__ == "__main__":
     ai_config = aiuNNConfig()
     base_model = AIIABase(config)
     # Instantiate Upsampler from the base model (works correctly).
-    upsampler = aiuNN(base_model, config=ai_config)
+    upsampler = aiuNN(config=ai_config)
-
+    upsampler.load_base_model(base_model)
     # Save the model (both configuration and weights).
-    upsampler.save("aiunn")
+    upsampler.save_pretrained("aiunn")
 
     # Now load using the overridden load method; this will load the complete model.
-    upsampler_loaded = aiuNN.load("aiunn", precision="bf16")
+    upsampler_loaded = aiuNN.from_pretrained("aiunn")
     print("Updated configuration:", upsampler_loaded.config.__dict__)

tests/inference/test_inference.py

@@ -21,9 +21,8 @@ def real_model(tmp_path):
     base_model = AIIABase(config)
     
     # Make sure aiuNN is properly configured with all required attributes
-    upsampler = aiuNN(base_model, config=ai_config)
+    upsampler = aiuNN(config=ai_config)
-    # Ensure the upsample attribute is properly set if needed
+    upsampler.load_base_model(base_model)
-    # upsampler.upsample = ... # Add any necessary initialization
     
     # Save the model and config to temporary directory
     save_path = str(model_dir / "save")
@@ -40,10 +39,10 @@ def real_model(tmp_path):
         json.dump(config_data, f)
     
     # Save model
-    upsampler.save(save_path)
+    upsampler.save_pretrained(save_path)
     
     # Load model in inference mode
-    inference_model = aiuNNInference(model_path=save_path, precision='fp16', device='cpu')
+    inference_model = aiuNNInference(model_path=save_path, device='cpu')
     return inference_model
 
 
@@ -88,12 +87,3 @@ 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,39 +10,21 @@ def test_save_and_load_model():
         config = AIIAConfig()
         ai_config = aiuNNConfig()
         base_model = AIIABase(config)
-        upsampler = aiuNN(base_model, config=ai_config)
+        upsampler = aiuNN(config=ai_config)
-
+        upsampler.load_base_model(base_model)
         # Save the model
         save_path = os.path.join(tmpdirname, "model")
-        upsampler.save(save_path)
+        upsampler.save_pretrained(save_path)
 
         # Load the model
-        loaded_upsampler = aiuNN.load(save_path)
+        loaded_upsampler = aiuNN.from_pretrained(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__
+        assert loaded_upsampler.config.hidden_size == upsampler.config.hidden_size
+        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()