Update handler.py

handler.py

@@ -2,28 +2,31 @@ import torch
 from diffusers import StableDiffusionXLPipeline
 import base64
 from io import BytesIO
+import os
 
 class InferenceHandler:
     def __init__(self):
+        # Determine the device to run on
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         model_name = "colt12/maxcushion"
 
-        # If your model is private, include the use_auth_token parameter
+        # Load the pipeline with authentication from environment variable
         self.pipe = StableDiffusionXLPipeline.from_pretrained(
             model_name,
             torch_dtype=torch.float16,
             use_safetensors=True,
-            # Uncomment the line below and replace with your token if needed
-            # use_auth_token="your_huggingface_token"
+            use_auth_token=os.getenv("HUGGINGFACE_TOKEN")  # Securely get the token
         ).to(self.device)
 
     def __call__(self, inputs):
+        # Extract the prompt from inputs
         prompt = inputs.get("prompt", "")
         if not prompt:
             raise ValueError("A prompt must be provided")
         
         negative_prompt = inputs.get("negative_prompt", "")
-        
+
+        # Generate the image using the pipeline
         image = self.pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -31,10 +34,13 @@ class InferenceHandler:
             guidance_scale=7.5
         ).images[0]
         
+        # Convert the image to base64 encoding
         buffered = BytesIO()
         image.save(buffered, format="PNG")
         image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
         
+        # Return the base64 image
         return {"image_base64": image_base64}
 
+# Instantiate the handler
 handler = InferenceHandler()