Update handler.py

handler.py

@@ -1,55 +1,23 @@
-from typing import Dict
-import torch
-from diffusers import DiffusionPipeline
-from compel import Compel
-from io import BytesIO
-import base64
-import os
+def __call__(self, data: Dict) -> Dict:
+    print("Received data:", data)
 
-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
-        )
-        lora_path = os.path.join(path, "c1t3_v1.safetensors")
-        print(f"Loading LoRA weights from: {lora_path}")
-        self.pipe.load_lora_weights(lora_path)
+    inputs = data.get("inputs", {})
+    prompt = inputs.get("prompt", "")
 
-        if torch.cuda.is_available():
-            self.pipe.to("cuda")
-        else:
-            self.pipe.to("cpu")
+    print("Extracted prompt:", prompt)
 
-        self.pipe.enable_model_cpu_offload()
+    if not prompt:
+        return {"error": "No prompt provided."}
 
-        self.compel = Compel(
-            tokenizer=self.pipe.tokenizer,
-            text_encoder=self.pipe.text_encoder
-        )
-        print("Model initialized.")
+    conditioning = self.compel(prompt)
+    print("Conditioning complete.")
 
-    def __call__(self, data: Dict) -> Dict:
-        print("Received data:", data)
+    image = self.pipe(prompt_embeds=conditioning).images[0]
+    print("Image generated.")
 
-        inputs = data.get("inputs", {})
-        prompt = inputs.get("prompt", "")
-        print("Extracted prompt:", prompt)
+    buffer = BytesIO()
+    image.save(buffer, format="PNG")
+    base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
 
-        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}
+    print("Returning image.")
+    return {"image": base64_image}