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}