Add full diagnostic prints for debugging

handler.py

@@ -10,7 +10,9 @@ from basicsr.archs.rrdbnet_arch import RRDBNet
 
 class EndpointHandler:
     def __init__(self, path="."):
-        print("🔹 Initializing Real-ESRGAN x4 model...")
+        print("🚀 [INIT] Starting EndpointHandler initialization...")
+        print(f"📂 Working directory: {os.getcwd()}")
+        print(f"📁 Model path root: {path}")
 
         self.model_url = (
             "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/"
@@ -19,13 +21,16 @@ class EndpointHandler:
         self.model_path = os.path.join(path, "RealESRGAN_x4plus.pth")
 
         if not os.path.exists(self.model_path):
-            print(f"📥 Downloading RealESRGAN_x4plus.pth ...")
+            print(f"📥 [DOWNLOAD] Fetching model weights from {self.model_url}")
             r = requests.get(self.model_url)
             r.raise_for_status()
             with open(self.model_path, "wb") as f:
                 f.write(r.content)
-            print(f"✅ Downloaded model to {self.model_path}")
+            print(f"✅ [DOWNLOAD] Saved model to {self.model_path}")
+        else:
+            print(f"✅ [CACHE] Model already exists at {self.model_path}")
 
+        print("🧠 [MODEL] Building RRDBNet...")
         model = RRDBNet(
             num_in_ch=3,
             num_out_ch=3,
@@ -35,69 +40,117 @@ class EndpointHandler:
             scale=4,
         )
 
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"💻 [DEVICE] Using device: {device}")
+
         self.upsampler = RealESRGANer(
             scale=4,
             model_path=self.model_path,
             model=model,
             half=False,
-            device="cuda" if torch.cuda.is_available() else "cpu",
+            device=device,
         )
 
-        print("✅ Real-ESRGAN model initialized and ready.")
+        print("✅ [INIT DONE] Real-ESRGAN model initialized and ready.\n\n")
 
     # ==========================================================
-    # Main callable
+    # MAIN CALLABLE
     # ==========================================================
     def __call__(self, data):
+        print("🛰️ [CALL] Endpoint invoked!")
+        print(f"📦 [CALL] Raw data type: {type(data)}")
+        print(f"🔍 [CALL] Data preview: {str(data)[:300]}...")
+
         try:
+            print("➡️ [STEP] Preprocessing input...")
             image = self.preprocess(data)
+            print("✅ [STEP] Preprocessing complete!")
+
+            print("➡️ [STEP] Running inference...")
             output = self.inference(image)
-            return self.postprocess(output)
+            print("✅ [STEP] Inference complete!")
+
+            print("➡️ [STEP] Encoding output image...")
+            result = self.postprocess(output)
+            print("✅ [STEP] Postprocessing complete!")
+
+            return result
         except Exception as e:
-            print("❌ Error:", str(e))
+            print("💥 [ERROR] Exception during inference:", str(e))
             return {"error": str(e)}
 
     # ==========================================================
-    # Steps
+    # PREPROCESS
     # ==========================================================
     def preprocess(self, data):
-        """Accept raw bytes, dicts, or wrapped HF JSON input."""
-        # Case 1: raw bytes
+        print(f"🔧 [PREPROCESS] Type received: {type(data)}")
+
+        # 1️⃣  Hugging Face JSON-wrapped dict {"inputs": ...}
+        if isinstance(data, dict):
+            print("🧩 [PREPROCESS] Detected dict input.")
+            if "inputs" in data:
+                print("📨 [PREPROCESS] Found 'inputs' key in dict.")
+                data = data["inputs"]
+            elif "image" in data:
+                print("📨 [PREPROCESS] Found 'image' key in dict.")
+                data = data["image"]
+            else:
+                print("⚠️ [PREPROCESS] Dict has no 'inputs' or 'image' keys.")
+
+        # 2️⃣  Bytes
         if isinstance(data, (bytes, bytearray)):
+            print("🖼️ [PREPROCESS] Treating input as raw bytes.")
+            print(f"🧮 [BYTES] Length: {len(data)}")
             return Image.open(io.BytesIO(data)).convert("RGB")
 
-        # Case 2: HF JSON payload like {"inputs": ...}
-        if isinstance(data, dict):
-            img_field = data.get("inputs") or data.get("image") or data
-            # If bytes-like
-            if isinstance(img_field, (bytes, bytearray)):
-                return Image.open(io.BytesIO(img_field)).convert("RGB")
-            # If base64 string
-            if isinstance(img_field, str):
-                try:
-                    decoded = base64.b64decode(img_field)
-                    return Image.open(io.BytesIO(decoded)).convert("RGB")
-                except Exception:
-                    raise ValueError("Invalid base64 data in 'inputs' field.")
-
-        # Case 3: Hugging Face wraps even raw bytes under a dict with 'inputs' -> list
-        if isinstance(data, list) and len(data) > 0:
-            item = data[0]
-            if isinstance(item, (bytes, bytearray)):
-                return Image.open(io.BytesIO(item)).convert("RGB")
-            if isinstance(item, str):
-                return Image.open(io.BytesIO(base64.b64decode(item))).convert("RGB")
-
-        raise ValueError("Expected raw image bytes or {'inputs': <bytes/base64>}.")
+        # 3️⃣  Base64 string
+        if isinstance(data, str):
+            print("🧾 [PREPROCESS] Treating input as base64 string.")
+            print(f"🧮 [BASE64] Length: {len(data)}")
+            try:
+                decoded = base64.b64decode(data)
+                print(f"✅ [BASE64] Successfully decoded base64 -> {len(decoded)} bytes.")
+                return Image.open(io.BytesIO(decoded)).convert("RGB")
+            except Exception as e:
+                print("💥 [BASE64] Decode failed:", str(e))
+                raise ValueError("Invalid base64 string in request body.")
+
+        # 4️⃣  List of inputs (some HF payloads wrap inside a list)
+        if isinstance(data, list):
+            print(f"📚 [PREPROCESS] List input detected with length {len(data)}")
+            if len(data) == 0:
+                raise ValueError("Empty input list.")
+            first = data[0]
+            print(f"🔹 [PREPROCESS] First element type: {type(first)}")
+            if isinstance(first, (bytes, bytearray)):
+                return Image.open(io.BytesIO(first)).convert("RGB")
+            if isinstance(first, str):
+                decoded = base64.b64decode(first)
+                return Image.open(io.BytesIO(decoded)).convert("RGB")
 
+        raise ValueError("Unsupported input type. Expected image bytes or base64 data.")
+
+    # ==========================================================
+    # INFERENCE
+    # ==========================================================
     def inference(self, image):
+        print("🎯 [INFERENCE] Running ESRGAN upscaling...")
+        print(f"📐 [INFERENCE] Input image size: {image.size}")
         output, _ = self.upsampler.enhance(image, outscale=4)
+        print(f"✅ [INFERENCE] Output image size: {output.size}")
         return output
 
+    # ==========================================================
+    # POSTPROCESS
+    # ==========================================================
     def postprocess(self, output_image):
+        print("📤 [POSTPROCESS] Encoding image to base64...")
         buf = io.BytesIO()
         output_image.save(buf, format="PNG")
-        encoded = base64.b64encode(buf.getvalue()).decode("utf-8")
+        raw_bytes = buf.getvalue()
+        print(f"📏 [POSTPROCESS] Output image byte size: {len(raw_bytes)}")
+        encoded = base64.b64encode(raw_bytes).decode("utf-8")
+        print(f"✅ [POSTPROCESS] Base64 encoded length: {len(encoded)}")
         buf.close()
         return {"image": encoded}