Update handler.py

handler.py

@@ -2,6 +2,8 @@ from typing import Dict
 import torch
 from diffusers import DiffusionPipeline
 from compel import Compel
+from io import BytesIO
+import base64
 
 class EndpointHandler:
     def __init__(self, path: str = ""):
@@ -12,7 +14,7 @@ class EndpointHandler:
             variant="fp16",
         )
 
-        # Load your LoRA weights hosted in the same repo
+        # Load your LoRA weights from the repo
         self.pipe.load_lora_weights("./c1t3_v1.safetensors")
 
         # Move to GPU if available
@@ -24,20 +26,24 @@ class EndpointHandler:
         # Optional: enable memory optimization
         self.pipe.enable_model_cpu_offload()
 
-        # Initialize Compel (prompt parser)
+        # Initialize Compel (prompt parser for FLUX)
         self.compel = Compel(tokenizer=self.pipe.tokenizer, text_encoder=self.pipe.text_encoder)
 
     def __call__(self, data: Dict[str, str]) -> Dict:
-        # Get the prompt from request
+        # Get prompt from request data
         prompt = data.get("prompt", "")
         if not prompt:
             return {"error": "No prompt provided."}
 
-        # Process the prompt with Compel (recommended for FLUX)
+        # Generate prompt conditioning using Compel
         conditioning = self.compel(prompt)
 
-        # Generate the image
+        # Generate image using FLUX + LoRA
         image = self.pipe(prompt_embeds=conditioning).images[0]
 
-        # Return the result
-        return {"image": image}
+        # Convert image to base64 string for API response
+        buffer = BytesIO()
+        image.save(buffer, format="PNG")
+        base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+        return {"image": base64_image}