| from typing import Dict |
| import torch |
| from diffusers import FluxPipeline |
| from io import BytesIO |
| import base64 |
|
|
| class EndpointHandler: |
| def __init__(self, path: str = ""): |
| print(f"Initializing model from: {path}") |
| |
| self.pipe = FluxPipeline.from_pretrained( |
| "black-forest-labs/FLUX.1-dev", |
| torch_dtype=torch.float16 |
| ) |
|
|
| print("Loading LoRA weights from: Texttra/Cityscape_Studio") |
| self.pipe.load_lora_weights("Texttra/Cityscape_Studio", weight_name="c1t3_v1.safetensors") |
| self.pipe.fuse_lora(lora_scale=0.9) |
|
|
| self.pipe.to("cuda" if torch.cuda.is_available() else "cpu") |
| print("Model initialized successfully.") |
|
|
| def __call__(self, data: Dict) -> Dict: |
| print("Received data:", data) |
|
|
| inputs = data.get("inputs", {}) |
| prompt = inputs.get("prompt", "") |
| print("Extracted prompt:", prompt) |
|
|
| if not prompt: |
| return {"error": "No prompt provided."} |
|
|
| image = self.pipe( |
| prompt, |
| num_inference_steps=50, |
| guidance_scale=4.5 |
| ).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} |
|
|
|
|
