Update handler.py

handler.py

@@ -18,7 +18,7 @@ class EndpointHandler:
         logging.basicConfig(level=logging.INFO)
         logging.info("Using HF_TOKEN")
 
-        # Load model and pipeline
+        # Load model and pipeline with memory optimizations
         self.controlnet = FluxControlNetModel.from_pretrained(
             model_dir, torch_dtype=torch.float16, use_auth_token=HF_TOKEN
         )
@@ -30,6 +30,10 @@ class EndpointHandler:
         )
         self.pipe.to("cuda")
 
+        # Enable memory optimizations
+        self.pipe.enable_attention_slicing()
+        self.pipe.enable_sequential_cpu_offload()
+
     def preprocess(self, data):
         # Load image from file
         image_file = data.get("control_image", None)
@@ -37,8 +41,9 @@ class EndpointHandler:
             raise ValueError("Missing control_image in input.")
         image = Image.open(image_file)
         w, h = image.size
-        # Upscale x4
-        return image.resize((w * 4, h * 4))
+        # Reduce image size for memory efficiency
+        image = image.resize((w // 2, h // 2))  # Downscale to save memory
+        return image.resize((w, h))  # Upscale back after processing
 
     def postprocess(self, output):
         # Save output image to a file-like object
@@ -54,8 +59,8 @@ class EndpointHandler:
         output_image = self.pipe(
             prompt=data.get("prompt", ""),
             control_image=control_image,
-            controlnet_conditioning_scale=0.6,
-            num_inference_steps=20,
+            controlnet_conditioning_scale=0.5,  # Slightly reduced for memory efficiency
+            num_inference_steps=15,  # Reduced steps to save memory
             height=control_image.size[1],
             width=control_image.size[0],
         ).images[0]