Update handler.py

handler.py

@@ -5,13 +5,18 @@ from diffusers.utils import load_image
 from diffusers import FluxControlNetModel
 from diffusers.pipelines import FluxControlNetPipeline
 from io import BytesIO
+import logging
 
 class EndpointHandler:
     def __init__(self, model_dir="huyai123/Flux.1-dev-Image-Upscaler"):
         # Access the environment variable
         HUGGINGFACE_API_TOKEN = os.getenv('HUGGINGFACE_API_TOKEN')
         if not HUGGINGFACE_API_TOKEN:
-            raise ValueError("HUGGINGFACE_API_TOKEN")
+            raise ValueError("HUGGINGFACE_API_TOKEN environment variable is not set")
+
+        # Log the token for debugging (remove this in production)
+        logging.basicConfig(level=logging.INFO)
+        logging.info(f"Using HUGGINGFACE_API_TOKEN: {HUGGINGFACE_API_TOKEN}")
 
         # Load model and pipeline
         self.controlnet = FluxControlNetModel.from_pretrained(
@@ -46,13 +51,20 @@ class EndpointHandler:
         # Preprocess input
         control_image = self.preprocess(data)
         # Generate output
-        output_image = self.pipe(
-            prompt=data.get("prompt", ""),
-            control_image=control_image,
-            controlnet_conditioning_scale=0.6,
+        outputImage = self.pipe(
+            prompt=data.get("prompt","),
+            control_image=controlImage,
+            controlnet_condition_scale=0.6,
             num_inference_steps=28,
-            height=control_image.size[1],
-            width=control_image.size[0],
+            height=controlImage.size[1],
+            width=controlImage.size[0],
         ).images[0]
         # Postprocess output
-        return self.postprocess(output_image)
+        return self.postprocess(outputImage)
+
+if __name__ == "__main__":
+    # Example usage
+    data = {'control_image': 'path/to/your/image.png', 'prompt': 'Your prompt here'}
+    handler = EndpointHandler()
+    output = handler.inference(data)
+    print(output)