Update handler.py

handler.py

@@ -1,6 +1,9 @@
 from diffusers import AutoencoderKL
+from PIL import Image
+import io
 from transformers import CLIPProcessor, CLIPModel
 from model import Model
+from pathlib import Path
 from noise_scheduler import NoiseSchedule
 import torch
 import base64
@@ -35,13 +38,16 @@ class CLIP:
 
 class Inference:
     def __init__(self):
+        here = Path(__file__).resolve().parent
+        ckpt_path = here / "unet.pt"
+
         self.clip = CLIP()
         self.ae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae").to('cuda' if torch.cuda.is_available() else "cpu")
         self.ae.eval()
         for name, param in self.ae.named_parameters():
             param.requires_grad = False
         self.unet = Model(T=T, filters=[64,128,256,512], t_dim=t_dim, depth=depth, LDM=LDM)
-        self.unet.load_state_dict(torch.load("repository/unet.pt", weights_only=False, map_location=torch.device('cpu')))
+        self.unet.load_state_dict(torch.load(ckpt_path, weights_only=False, map_location=torch.device('cpu')))
         self.unet.eval()
         for name, param in self.unet.named_parameters():
             param.requires_grad = False
@@ -65,11 +71,17 @@ class Inference:
 
 class EndpointHandler:
     def __init__(self, path: str = ""):
-        # path -> repo directory on the endpoint container
-        # you can read files via Path(path)/"unet.pt" if needed
         self.engine = Inference()
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        png_bytes, score = self.engine(num_images=1)
+        img_tensor, score, _ = self.engine(num_images=1)  # (H,W,C) in [0,1]
+        img_uint8 = (img_tensor.clamp(0,1).numpy() * 255).astype("uint8")
+        pil_img = Image.fromarray(img_uint8)
+        
+        buf = io.BytesIO()
+        pil_img.save(buf, format="PNG")
+        png_bytes = buf.getvalue()
+
         b64 = base64.b64encode(png_bytes).decode("utf-8")
         return {"image": b64, "score": float(score)}
+