dropped batch processing and dropped fp16 loading
Gitea Actions For AIIA / Explore-Gitea-Actions (push) Has been cancelled
Details
Gitea Actions For AIIA / Explore-Gitea-Actions (push) Has been cancelled
Details
This commit is contained in:
parent
ced7e8a214
commit
ac3fabd55f
|
@ -12,13 +12,13 @@ class aiuNNInference:
|
||||||
Inference class for aiuNN upsampling model.
|
Inference class for aiuNN upsampling model.
|
||||||
Handles model loading, image upscaling, and output processing.
|
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.
|
Initialize the inference class by loading the aiuNN model.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
model_path: Path to the saved model directory
|
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)
|
device: Optional device specification ('cuda', 'cpu', or None for auto-detection)
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
@ -30,7 +30,7 @@ class aiuNNInference:
|
||||||
self.device = device
|
self.device = device
|
||||||
|
|
||||||
# Load the model with specified precision
|
# 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.to(self.device)
|
||||||
self.model.eval()
|
self.model.eval()
|
||||||
|
|
||||||
|
@ -160,54 +160,11 @@ class aiuNNInference:
|
||||||
|
|
||||||
return binary_data
|
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)
|
# Example usage (can be removed)
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
# Initialize inference with a model path
|
# Initialize inference with a model path
|
||||||
inferencer = aiuNNInference("path/to/model", precision="bf16")
|
inferencer = aiuNNInference("path/to/model")
|
||||||
|
|
||||||
# Upscale a single image
|
# Upscale a single image
|
||||||
upscaled_image = inferencer.upscale("input_image.jpg")
|
upscaled_image = inferencer.upscale("input_image.jpg")
|
||||||
|
@ -218,9 +175,3 @@ if __name__ == "__main__":
|
||||||
# Convert to binary
|
# Convert to binary
|
||||||
binary_data = inferencer.convert_to_binary(upscaled_image)
|
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"
|
|
||||||
)
|
|
Loading…
Reference in New Issue