Upload inference.py with huggingface_hub

inference.py

@@ -1,50 +1,79 @@
 from pathlib import Path
-import torch
+import base64, io, torch
 from PIL import Image
-import base64, io
-
+from huggingface_hub import hf_hub_download
 from enhancer import ESRGANUpscaler, ESRGANUpscalerCheckpoints
 
-# -------- initialise model once at cold-start --------
+# ------------------------------------------------------------------
+# 1️⃣  Download all required checkpoints once (cached by HF DL cache)
+# ------------------------------------------------------------------
+def dl(repo, fname, rev=None):
+    return Path(
+        hf_hub_download(repo_id=repo, filename=fname, revision=rev)
+    )
+
 checkpoints = ESRGANUpscalerCheckpoints(
-    esrgan=Path("checkpoints/4x-UltraSharp.pth")
+    # SD-1.5 Multi-Upscaler core
+    unet              = dl("refiners/juggernaut.reborn.sd1_5.unet",         "model.safetensors",   "347d14c3c782c4959cc4d1bb1e336d19f7dda4d2"),
+    clip_text_encoder = dl("refiners/juggernaut.reborn.sd1_5.text_encoder", "model.safetensors",   "744ad6a5c0437ec02ad826df9f6ede102bb27481"),
+    lda               = dl("refiners/juggernaut.reborn.sd1_5.autoencoder",  "model.safetensors",   "3c1aae3fc3e03e4a2b7e0fa42b62ebb64f1a4c19"),
+    controlnet_tile   = dl("refiners/controlnet.sd1_5.tile",                "model.safetensors",   "48ced6ff8bfa873a8976fa467c3629a240643387"),
+
+    # ESRGAN 4× super-res
+    esrgan            = dl("philz1337x/upscaler",                           "4x-UltraSharp.pth",   "011deacac8270114eb7d2eeff4fe6fa9a837be70"),
+
+    # Negative prompt embedding
+    negative_embedding     = dl("philz1337x/embeddings", "JuggernautNegative-neg.pt", "203caa7e9cc2bc225031a4021f6ab1ded283454a"),
+    negative_embedding_key = "string_to_param.*",
+
+    # LoRAs
+    loras = {
+        "more_details": dl("philz1337x/loras", "more_details.safetensors",     "a3802c0280c0d00c2ab18d37454a8744c44e474e"),
+        "sdxl_render" : dl("philz1337x/loras", "SDXLrender_v2.0.safetensors",  "a3802c0280c0d00c2ab18d37454a8744c44e474e"),
+    },
 )
 
+# ------------------------------------------------------------------
+# 2️⃣  Instantiate the enhancer once (global singleton)
+# ------------------------------------------------------------------
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
+dtype  = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
 
 enhancer = ESRGANUpscaler(checkpoints=checkpoints, device=device, dtype=dtype)
 
-# -------- Hugging Face Hosted Inference API Entry Point --------
+# ------------------------------------------------------------------
+# 3️⃣  Hosted-Inference entry-point
+# ------------------------------------------------------------------
 class EndpointHandler:
+    """
+    Hugging Face Hosted Inference API entrypoint.
+    Expects: {"image": "<BASE64-ENCODED>"}
+    Returns: {enhanced_image, original_size, enhanced_size}
+    """
+
     def __init__(self, path="."):
-        pass  # path is not used since model is already initialized globally
+        pass  # all heavy work done globally above
 
     def __call__(self, inputs: dict) -> dict:
-        """
-        Expected payload:
-          {"image": "<BASE64-STRING>"}
-        Returns:
-          { "enhanced_image": "<BASE64-PNG>", "original_size": [w,h], "enhanced_size": [w,h] }
-        """
         if "image" not in inputs:
             return {"error": "No image provided"}
 
-        # decode base64
-        data = inputs["image"]
-        if data.startswith("data:image"):
-            data = data.split(",")[1]
-        img = Image.open(io.BytesIO(base64.b64decode(data))).convert("RGB")
+        # strip optional data:image/...;base64, prefix
+        payload = inputs["image"].split(",", 1)[-1]
+        try:
+            img = Image.open(io.BytesIO(base64.b64decode(payload))).convert("RGB")
+        except Exception as e:
+            return {"error": f"Bad image payload: {e}"}
 
-        # run ESRGAN ➜ SD-1.5 upscale
-        result = enhancer.upscale(img)
+        # run the two-stage upscaler
+        out = enhancer.upscale(img)
 
         buf = io.BytesIO()
-        result.save(buf, format="PNG")
-        result_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+        out.save(buf, format="PNG")
+        out_b64 = base64.b64encode(buf.getvalue()).decode()
 
         return {
-            "enhanced_image": result_b64,
+            "enhanced_image": out_b64,
             "original_size": img.size,
-            "enhanced_size": result.size,
+            "enhanced_size": out.size,
         }