Fix color hue and add RGB output conversion

handler.py

@@ -21,7 +21,7 @@ class EndpointHandler:
         self.half = self.device == "cuda"
         self.path = path
 
-        # URLs
+        # Model URLs (GFPGAN + RealESRGAN)
         self.gfpgan_model_url = (
             "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
         )
@@ -29,18 +29,21 @@ class EndpointHandler:
             "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"
         )
 
-        # Local model paths
+        # Local cache paths
         self.gfpgan_model_path = os.path.join(path, "GFPGANv1.4.pth")
         self.realesr_model_path = os.path.join(path, "realesr-general-x4v3.pth")
 
         self.bg_upsampler = None
         self.restorer = None
 
+        # Ensure model weights exist
         self._ensure_model(self.gfpgan_model_url, self.gfpgan_model_path)
         self._ensure_model(self.realesr_model_url, self.realesr_model_path)
+
         logger.info(f"🧠 Device: {self.device}, half precision: {self.half}")
 
     def _ensure_model(self, url, path):
+        """Download model if missing."""
         if not os.path.exists(path):
             logger.info(f"⬇️ Downloading model from {url}")
             r = requests.get(url, timeout=60)
@@ -52,7 +55,7 @@ class EndpointHandler:
             logger.info(f"📁 Found cached model: {path}")
 
     def _init_models(self):
-        """Lazy-load models"""
+        """Lazy-load ESRGAN + GFPGAN models."""
         if self.bg_upsampler is None:
             logger.info("🧩 Initializing Real-ESRGAN upsampler...")
             model = SRVGGNetCompact(
@@ -82,7 +85,7 @@ class EndpointHandler:
             logger.info("✅ Models ready!")
 
     def _load_image(self, data):
-        """Handle different input formats."""
+        """Accept base64, raw bytes, or URL and return PIL image."""
         if isinstance(data, dict) and "inputs" in data:
             data = data["inputs"]
 
@@ -96,7 +99,7 @@ class EndpointHandler:
                 resp = requests.get(data)
                 return Image.open(io.BytesIO(resp.content)).convert("RGB")
             else:
-                # assume base64
+                # Base64
                 logger.info("🧬 Decoding base64 image input")
                 try:
                     decoded = base64.b64decode(data)
@@ -108,30 +111,34 @@ class EndpointHandler:
         raise ValueError("Unsupported input type")
 
     def __call__(self, data):
+        logger.info("⚙️ Starting GFPGAN inference pipeline...")
         self._init_models()
-        logger.info("⚙️ Starting inference...")
 
         # Load input
         image = self._load_image(data)
         input_img = np.array(image, dtype=np.uint8)
-
         logger.info(f"📏 Input image shape: {input_img.shape}")
 
+        # Restore face(s)
         cropped_faces, restored_faces, restored_img = self.restorer.enhance(
             input_img, has_aligned=False, only_center_face=False, paste_back=True
         )
 
-        logger.info("🖼️ Restoration complete, encoding output...")
+        logger.info("🖼️ Restoration complete, preparing output...")
 
-        # Encode result as base64
+        # ✅ Convert color from BGR → RGB (fix hue issue)
         restored_img_rgb = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
-        _, buffer = cv2.imencode(".jpg", restored_img)
+        restored_img_rgb = np.clip(restored_img_rgb, 0, 255).astype(np.uint8)
+
+        # ✅ Encode output as base64 string for JSON
+        _, buffer = cv2.imencode(".jpg", restored_img_rgb)
         b64_output = base64.b64encode(buffer).decode("utf-8")
 
-        logger.info("✅ Returning base64 image JSON")
+        logger.info("✅ Returning base64-encoded image JSON response")
+
         return {
             "image": b64_output,
             "status": "success",
-            "info": "Restored with GFPGAN v1.4 + Real-ESRGAN x4v3"
+            "info": "Restored with GFPGAN v1.4 + Real-ESRGAN x4v3 (RGB fixed)"
         }