litserve

configs/infer.yaml

@@ -40,3 +40,13 @@ seed: 42
 
 # name of the experiment
 name: "catdog_experiment"
+
+server:
+  port: 8080
+  max_batch_size: 8
+  batch_timeout: 0.01
+  accelerator: "auto"
+  devices: "auto"
+  workers_per_device: 2
+
+labels: ["cat", "dog"]

src/client.py

@@ -0,0 +1,85 @@
+import requests
+from urllib.request import urlopen
+import base64
+import os
+
+
+def fetch_image(url):
+    """
+    Fetch image data from a URL.
+    """
+    return urlopen(url).read()
+
+
+def encode_image_to_base64(img_data):
+    """
+    Encode image bytes to a base64 string.
+    """
+    return base64.b64encode(img_data).decode("utf-8")
+
+
+def send_prediction_request(base64_image, server_url):
+    """
+    Send a single base64 image to the prediction API and retrieve predictions.
+    """
+    try:
+        response = requests.post(f"{server_url}/predict", json={"image": base64_image})
+        return response
+    except requests.exceptions.RequestException as e:
+        print(f"Error connecting to the server: {e}")
+        return None
+
+
+def send_batch_prediction_request(base64_images, server_url):
+    """
+    Send a batch of base64 images to the prediction API and retrieve predictions.
+    """
+    try:
+        response = requests.post(
+            f"{server_url}/predict", json=[{"image": img} for img in base64_images]
+        )
+        return response
+    except requests.exceptions.RequestException as e:
+        print(f"Error connecting to the server: {e}")
+        return None
+
+
+def main():
+    # Server URL (default or from environment)
+    server_url = os.getenv("SERVER_URL", "http://localhost:8080")
+
+    # Example URLs for testing
+    image_urls = [
+        "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png"
+    ]
+
+    # Fetch and encode images
+    try:
+        print("Fetching and encoding images...")
+        base64_images = [encode_image_to_base64(fetch_image(url)) for url in image_urls]
+        print("Images fetched and encoded successfully.")
+    except Exception as e:
+        print(f"Error fetching or encoding images: {e}")
+        return
+
+    # Test single image prediction
+    try:
+        print("\n--- Single Image Prediction ---")
+        single_response = send_prediction_request(base64_images[0], server_url)
+        if single_response and single_response.status_code == 200:
+            predictions = single_response.json().get("predictions", [])
+            if predictions:
+                print("Top 5 Predictions:")
+                for pred in predictions:
+                    print(f"{pred['label']}: {pred['probability']:.2%}")
+            else:
+                print("No predictions returned.")
+        elif single_response:
+            print(f"Error: {single_response.status_code}")
+            print(single_response.text)
+    except Exception as e:
+        print(f"Error sending single prediction request: {e}")
+
+
+if __name__ == "__main__":
+    main()

src/infer.py

@@ -82,7 +82,7 @@ def download_image(cfg: DictConfig):
         logger.error(f"Failed to download image. Status code: {response.status_code}")
 
 
-@hydra.main(config_path="../configs", config_name="infer", version_base="1.1")
+@hydra.main(config_path="../configs", config_name="infer", version_base="1.3")
 def main_infer(cfg: DictConfig):
     # Print the configuration
     logger.info(OmegaConf.to_yaml(cfg))

src/litserve_test_client.py

@@ -50,8 +50,7 @@ def main():
 
     # Example URLs for testing
     image_urls = [
-        "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png",
-        "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png",
+        "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png"
     ]
 
     # Fetch and encode images

src/server.py

@@ -0,0 +1,185 @@
+import torch
+from PIL import Image
+import io
+import litserve as lit
+import base64
+from torchvision import transforms
+from src.models.catdog_model import ViTTinyClassifier
+import hydra
+from omegaconf import DictConfig, OmegaConf
+from dotenv import load_dotenv, find_dotenv
+import rootutils
+from loguru import logger
+from src.utils.logging_utils import setup_logger
+from pathlib import Path
+
+# Load environment variables
+load_dotenv(find_dotenv(".env"))
+
+# Setup root directory
+root = rootutils.setup_root(__file__, indicator=".project-root")
+logger.info(f"Root directory set to: {root}")
+
+
+class ImageClassifierAPI(lit.LitAPI):
+    def __init__(self, cfg: DictConfig):
+        """
+        Initialize the API with Hydra configuration.
+        """
+        super().__init__()
+        self.cfg = cfg
+
+        # Validate required config keys
+        required_keys = ["ckpt_path", "data.image_size", "labels"]
+        missing_keys = [key for key in required_keys if not OmegaConf.select(cfg, key)]
+        if missing_keys:
+            logger.error(f"Missing required config keys: {missing_keys}")
+            raise ValueError(f"Missing required config keys: {missing_keys}")
+        logger.info(f"Configuration validated: {OmegaConf.to_yaml(cfg)}")
+
+    def setup(self, device):
+        """Initialize the model and necessary components."""
+        self.device = device
+        logger.info("Setting up the model and components.")
+
+        # Log the configuration for debugging
+        logger.debug(f"Configuration passed to setup: {OmegaConf.to_yaml(self.cfg)}")
+
+        # Load the model from checkpoint
+        try:
+            self.model = ViTTinyClassifier.load_from_checkpoint(
+                checkpoint_path=self.cfg.ckpt_path
+            )
+            self.model = self.model.to(device).eval()
+            logger.info("Model loaded and moved to device.")
+        except FileNotFoundError:
+            logger.error(f"Checkpoint file not found: {self.cfg.ckpt_path}")
+            raise
+        except Exception as e:
+            logger.error(f"Error loading model: {e}")
+            raise
+
+        # Define transforms
+        self.transforms = transforms.Compose(
+            [
+                transforms.Resize((self.cfg.data.image_size, self.cfg.data.image_size)),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406],  # Hard-coded mean
+                    std=[0.229, 0.224, 0.225],  # Hard-coded std
+                ),
+            ]
+        )
+        logger.info("Transforms initialized.")
+
+        # Load labels
+        try:
+            self.labels = self.cfg.labels
+            logger.info(f"Labels loaded: {self.labels}")
+        except Exception as e:
+            logger.error(f"Error loading labels: {e}")
+            raise ValueError("Failed to load labels from the configuration.")
+
+    def decode_request(self, request):
+        """Handle both single and batch inputs."""
+        logger.info(f"decode_request received: {request}")
+        if not isinstance(request, dict) or "image" not in request:
+            logger.error(
+                "Invalid request format. Expected a dictionary with key 'image'."
+            )
+            raise ValueError(
+                "Invalid request format. Expected a dictionary with key 'image'."
+            )
+        return request["image"]
+
+    def batch(self, inputs):
+        """Batch process images."""
+        logger.info(f"batch received inputs: {inputs}")
+        if not isinstance(inputs, list):
+            raise ValueError("Input to batch must be a list.")
+
+        batch_tensors = []
+        try:
+            for image_bytes in inputs:
+                if not isinstance(image_bytes, str):  # Ensure input is a base64 string
+                    raise ValueError(
+                        f"Input must be a base64-encoded string, got: {type(image_bytes)}"
+                    )
+
+                # Decode base64 string to bytes
+                img_bytes = base64.b64decode(image_bytes)
+
+                # Convert bytes to PIL Image
+                try:
+                    image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+                except Exception as img_error:
+                    logger.error(f"Failed to process image: {img_error}")
+                    raise
+
+                # Apply transforms and add to batch
+                tensor = self.transforms(image)
+                batch_tensors.append(tensor)
+
+            return torch.stack(batch_tensors).to(self.device)
+        except Exception as e:
+            logger.error(f"Error decoding image: {e}")
+            raise ValueError("Failed to decode and process the images.")
+
+    def predict(self, x):
+        """Make predictions on the input batch."""
+        with torch.inference_mode():
+            outputs = self.model(x)
+            probabilities = torch.nn.functional.softmax(outputs, dim=1)
+        logger.info("Prediction completed.")
+        return probabilities
+
+    def unbatch(self, output):
+        """Unbatch the output."""
+        return [output[i] for i in range(output.size(0))]
+
+    def encode_response(self, output):
+        """Convert model output to API response for batches."""
+        try:
+            probs, indices = torch.topk(output, k=1)
+            responses = {
+                "predictions": [
+                    {
+                        "label": self.labels[idx.item()],
+                        "probability": prob.item(),
+                    }
+                    for prob, idx in zip(probs, indices)
+                ]
+            }
+            logger.info("Batch response successfully encoded.")
+            return responses
+        except Exception as e:
+            logger.error(f"Error encoding batch response: {e}")
+            raise ValueError("Failed to encode the batch response.")
+
+
+@hydra.main(config_path="../configs", config_name="infer", version_base="1.3")
+def main(cfg: DictConfig):
+    # Initialize loguru
+    setup_logger(Path(cfg.paths.log_dir) / "infer.log")
+    logger.info("Starting the Image Classifier API server.")
+
+    # Log configuration
+    logger.info(f"Configuration: {OmegaConf.to_yaml(cfg)}")
+
+    # Create the API instance with the Hydra config
+    api = ImageClassifierAPI(cfg)
+
+    # Configure the server
+    server = lit.LitServer(
+        api,
+        accelerator=cfg.server.accelerator,
+        max_batch_size=cfg.server.max_batch_size,
+        batch_timeout=cfg.server.batch_timeout,
+        devices=cfg.server.devices,
+        workers_per_device=cfg.server.workers_per_device,
+    )
+    server.run(port=cfg.server.port)
+
+
+if __name__ == "__main__":
+    main()