Upload handler.py

handler.py

@@ -2,7 +2,7 @@ import os
 from typing import Any, Dict
 
 from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, TorchAoConfig
-from PIL.Image import Image
+from PIL import Image
 import torch
 
 #import torch._dynamo
@@ -30,7 +30,7 @@ class EndpointHandler:
         self.pipeline.to("cuda")
 
     #@torch.inference_mode()
-    def __call__(self, data: Dict[str, Any]) -> Image:
+    def __call__(self, data: Dict[str, Any]) -> Image.Image:
         #logger.info(f"Received incoming request with {data=}")
 
         if "inputs" in data and isinstance(data["inputs"], str):