Update handler.py

handler.py

@@ -20,12 +20,12 @@ class EndpointHandler:
 
         # Load model and pipeline
         self.controlnet = FluxControlNetModel.from_pretrained(
-            model_dir, torch_dtype=torch.bfloat16, use_auth_token=HF_TOKEN
+            model_dir, torch_dtype=torch.float16, use_auth_token=HF_TOKEN
         )
         self.pipe = FluxControlNetPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-dev",
             controlnet=self.controlnet,
-            torch_dtype=torch.bfloat16,
+            torch_dtype=torch.float16,
             use_auth_token=HF_TOKEN
         )
         self.pipe.to("cuda")
@@ -55,7 +55,7 @@ class EndpointHandler:
             prompt=data.get("prompt", ""),
             control_image=control_image,
             controlnet_conditioning_scale=0.6,
-            num_inference_steps=28,
+            num_inference_steps=20,
             height=control_image.size[1],
             width=control_image.size[0],
         ).images[0]