Update handler.py

handler.py

@@ -27,6 +27,9 @@ class EndpointHandler:
         tensor = cast(torch.Tensor, data["inputs"])
         parameters = cast(dict, data.get("parameters", {}))
         do_scaling = cast(bool, parameters.get("do_scaling", True))
+        scaling_factor = cast(float, parameters.get("scaling_factor", None))
+        if do_scaling and scaling_factor is None:
+            scaling_factor = self.vae.config.scaling_factor
         output_type = cast(str, parameters.get("output_type", "pil"))
         partial_postprocess = cast(bool, parameters.get("partial_postprocess", False))
         if partial_postprocess and output_type != "pt":
@@ -34,8 +37,8 @@ class EndpointHandler:
 
         tensor = tensor.to(self.device, self.dtype)
 
-        if do_scaling:
-            tensor = tensor / self.vae.config.scaling_factor
+        if scaling_factor is not None:
+            tensor = tensor / scaling_factor
 
         with torch.no_grad():
             frames = cast(torch.Tensor, self.vae.decode(tensor, return_dict=False)[0])
@@ -55,4 +58,4 @@ class EndpointHandler:
         elif output_type == "pt":
             frames = frames
 
-        return frames
+        return frames