Update handler.py

handler.py

@@ -14,11 +14,14 @@ class EndpointHandler:
         if not HF_TOKEN:
             raise ValueError("HF_TOKEN environment variable is not set")
 
-        # Log the token for debugging (remove this in production)
+        # Log the token for debugging
         logging.basicConfig(level=logging.INFO)
         logging.info("Using HF_TOKEN")
 
-        # Load model and pipeline with memory optimizations
+        # Clear GPU memory
+        torch.cuda.empty_cache()
+
+        # Load model and pipeline
         self.controlnet = FluxControlNetModel.from_pretrained(
             model_dir, torch_dtype=torch.float16, use_auth_token=HF_TOKEN
         )
@@ -29,9 +32,7 @@ class EndpointHandler:
             use_auth_token=HF_TOKEN
         )
         self.pipe.to("cuda")
-
-        # Enable memory optimizations
-        self.pipe.enable_attention_slicing()
+        self.pipe.enable_attention_slicing("auto")
         self.pipe.enable_sequential_cpu_offload()
 
     def preprocess(self, data):
@@ -41,9 +42,7 @@ class EndpointHandler:
             raise ValueError("Missing control_image in input.")
         image = Image.open(image_file)
         w, h = image.size
-        # 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
+        return image.resize((w * 2, h * 2))  # Reduce upscale factor to save memory
 
     def postprocess(self, output):
         # Save output image to a file-like object
@@ -55,12 +54,14 @@ class EndpointHandler:
     def inference(self, data):
         # Preprocess input
         control_image = self.preprocess(data)
+        # Clear GPU memory
+        torch.cuda.empty_cache()
         # Generate output
         output_image = self.pipe(
             prompt=data.get("prompt", ""),
             control_image=control_image,
-            controlnet_conditioning_scale=0.5,  # Slightly reduced for memory efficiency
-            num_inference_steps=15,  # Reduced steps to save memory
+            controlnet_conditioning_scale=0.5,  # Reduced to save memory
+            num_inference_steps=15,  # Reduced steps
             height=control_image.size[1],
             width=control_image.size[0],
         ).images[0]