Update handler.py

handler.py

@@ -2,10 +2,9 @@ import torch
 from diffusers.utils import load_image
 from diffusers import FluxControlNetModel
 from diffusers.pipelines import FluxControlNetPipeline
-from PIL import Image
-import io
+from io import BytesIO
 
-class CustomHandler:
+class EndpointHandler:
     def __init__(self, model_dir="huyai123/Flux.1-dev-Image-Upscaler"):
         # Load model and pipeline
         self.controlnet = FluxControlNetModel.from_pretrained(
@@ -30,7 +29,7 @@ class CustomHandler:
 
     def postprocess(self, output):
         # Save output image to a file-like object
-        buffer = io.BytesIO()
+        buffer = BytesIO()
         output.save(buffer, format="PNG")
         buffer.seek(0)  # Reset buffer pointer
         return buffer
@@ -44,9 +43,8 @@ class CustomHandler:
             control_image=control_image,
             controlnet_conditioning_scale=0.6,
             num_inference_steps=28,
-            guidance_scale=3.5,
             height=control_image.size[1],
             width=control_image.size[0],
         ).images[0]
         # Postprocess output
-        return self.postprocess(output_image)
+        return self.postprocess(output_image)