Update handler.py

handler.py

@@ -1,23 +1,56 @@
-def __call__(self, data: Dict) -> Dict:
-    print("Received data:", data)
+from typing import Dict
+import torch
+from diffusers import DiffusionPipeline
+from compel import Compel
+from io import BytesIO
+import base64
+import os
 
-    inputs = data.get("inputs", {})
-    prompt = inputs.get("prompt", "")
+class EndpointHandler:
+    def __init__(self, path: str = ""):
+        print(f"Initializing model from: {path}")
+        self.pipe = DiffusionPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-dev",
+            torch_dtype=torch.float16,
+            use_auth_token=True
+        )
 
-    print("Extracted prompt:", prompt)
+        lora_path = os.path.join(path, "c1t3_v1.safetensors")
+        print(f"Loading LoRA weights from: {lora_path}")
+        self.pipe.load_lora_weights(lora_path)
 
-    if not prompt:
-        return {"error": "No prompt provided."}
+        if torch.cuda.is_available():
+            self.pipe.to("cuda")
+        else:
+            self.pipe.to("cpu")
 
-    conditioning = self.compel(prompt)
-    print("Conditioning complete.")
+        self.pipe.enable_model_cpu_offload()
 
-    image = self.pipe(prompt_embeds=conditioning).images[0]
-    print("Image generated.")
+        self.compel = Compel(
+            tokenizer=self.pipe.tokenizer,
+            text_encoder=self.pipe.text_encoder
+        )
+        print("Model initialized successfully.")
 
-    buffer = BytesIO()
-    image.save(buffer, format="PNG")
-    base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
+    def __call__(self, data: Dict) -> Dict:
+        print("Received data:", data)
 
-    print("Returning image.")
-    return {"image": base64_image}
+        inputs = data.get("inputs", {})
+        prompt = inputs.get("prompt", "")
+        print("Extracted prompt:", prompt)
+
+        if not prompt:
+            return {"error": "No prompt provided."}
+
+        conditioning = self.compel(prompt)
+        print("Conditioning complete.")
+
+        image = self.pipe(prompt_embeds=conditioning).images[0]
+        print("Image generated.")
+
+        buffer = BytesIO()
+        image.save(buffer, format="PNG")
+        base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        print("Returning image.")
+
+        return {"image": base64_image}