Update handler.py

handler.py

@@ -41,10 +41,10 @@ class EndpointHandler:
         """
         tensor = cast(torch.Tensor, data["inputs"])
         parameters = cast(dict, data.get("parameters", {}))
-        if "height" not in parameters or "width" not in parameters:
+        if tensor.ndim == 3 and ("height" not in parameters or "width" not in parameters):
             raise ValueError("Expected `height` and `width` in parameters.")
-        height = cast(int, parameters.get("height"))
-        width = cast(int, parameters.get("width"))
+        height = cast(int, parameters.get("height", 0))
+        width = cast(int, parameters.get("width", 0))
         do_scaling = cast(bool, parameters.get("do_scaling", True))
         output_type = cast(str, parameters.get("output_type", "pil"))
         partial_postprocess = cast(bool, parameters.get("partial_postprocess", False))
@@ -52,7 +52,8 @@ class EndpointHandler:
             output_type = "pt"
 
         tensor = tensor.to(self.device, self.dtype)
-        tensor = self._unpack_latents(tensor, height, width, self.vae_scale_factor)
+        if tensor.ndim == 3:
+            tensor = self._unpack_latents(tensor, height, width, self.vae_scale_factor)
 
         if do_scaling:
             tensor = (