Update handler.py

handler.py

@@ -1,22 +1,43 @@
-from diffusers import DiffusionPipeline
+from typing import Dict
 import torch
+from diffusers import DiffusionPipeline
+from compel import Compel
 
 class EndpointHandler:
-    def __init__(self, path="Texttra/Cityscape_Studio"):
-        print("Loading pipeline...")
-        try:
-            self.pipe = DiffusionPipeline.from_pretrained(
-                path,
-                torch_dtype=torch.float16,
-                revision="fp16",
-                use_safetensors=True
-            ).to("cuda")
-            print("Pipeline loaded successfully.")
-        except Exception as e:
-            print(f"Error loading model: {e}")
-            raise
-
-    def __call__(self, data):
-        inputs = data.pop("inputs", data)
-        prompt = inputs if isinstance(inputs, str) else inputs[0]
-        return self.pipe(prompt).images[0]
+    def __init__(self, path: str = ""):
+        # Load base FLUX pipeline
+        self.pipe = DiffusionPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-dev",
+            torch_dtype=torch.float16,
+            variant="fp16",
+        )
+
+        # Load your LoRA weights hosted in the same repo
+        self.pipe.load_lora_weights("./c1t3_v1.safetensors")
+
+        # Move to GPU if available
+        if torch.cuda.is_available():
+            self.pipe.to("cuda")
+        else:
+            self.pipe.to("cpu")
+
+        # Optional: enable memory optimization
+        self.pipe.enable_model_cpu_offload()
+
+        # Initialize Compel (prompt parser)
+        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
+        prompt = data.get("prompt", "")
+        if not prompt:
+            return {"error": "No prompt provided."}
+
+        # Process the prompt with Compel (recommended for FLUX)
+        conditioning = self.compel(prompt)
+
+        # Generate the image
+        image = self.pipe(prompt_embeds=conditioning).images[0]
+
+        # Return the result
+        return {"image": image}