Update handler.py

handler.py

@@ -10,23 +10,23 @@ 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 environment variable is not set")
+        HF_TOKEN = os.getenv('HF_TOKEN')
+        if not HF_TOKEN:
+            raise ValueError("HF_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}")
+        logging.info("Using HF_TOKEN")
 
         # Load model and pipeline
         self.controlnet = FluxControlNetModel.from_pretrained(
-            model_dir, torch_dtype=torch.bfloat16, use_auth_token=HUGGINGFACE_API_TOKEN
+            model_dir, torch_dtype=torch.bfloat16, use_auth_token=HF_TOKEN
         )
         self.pipe = FluxControlNetPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-dev",
             controlnet=self.controlnet,
             torch_dtype=torch.bfloat16,
-            use_auth_token=HUGGINGFACE_API_TOKEN
+            use_auth_token=HF_TOKEN
         )
         self.pipe.to("cuda")
 
@@ -67,4 +67,4 @@ if __name__ == "__main__":
     data = {'control_image': 'path/to/your/image.png', 'prompt': 'Your prompt here'}
     handler = EndpointHandler()
     output = handler.inference(data)
-    print(output)
+    print(output)