From ac3fabd55f070744d0ba65e522a844272a43c614 Mon Sep 17 00:00:00 2001
From: Falko Habel <falko.habel@gmx.de>
Date: Sat, 19 Apr 2025 22:53:13 +0200
Subject: [PATCH] dropped batch processing and dropped fp16 loading

---
 src/aiunn/inference/inference.py | 59 +++-----------------------------
 1 file changed, 5 insertions(+), 54 deletions(-)

diff --git a/src/aiunn/inference/inference.py b/src/aiunn/inference/inference.py
index d288931..1ff5994 100644
--- a/src/aiunn/inference/inference.py
+++ b/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")
@@ -217,10 +174,4 @@ 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"
-    )
\ No newline at end of file
+    
\ No newline at end of file