import gradio as gr
import requests
import io
from PIL import Image
import json
import os

# Load LoRAs from JSON
with open('loras.json', 'r') as f:
    loras = json.load(f)

# API call function
def query(payload, api_url, token):
    headers = {"Authorization": f"Bearer {token}"}
    print(f"Sending API request with payload: {payload}")
    response = requests.post(api_url, headers=headers, json=payload)
    if response.status_code == 200:
        return io.BytesIO(response.content)
    else:
        print(f"API Error: {response.text}")
        return None

# Define the function to run when the button is clicked
def run_lora(prompt, selected_lora_index):
    if selected_lora_index is not None:
        selected_lora_index = selected_lora_index[0]  # Unpack the tuple
        selected_lora = loras[selected_lora_index]
        api_url = f"https://api-inference.huggingface.co/models/{selected_lora['repo']}"
        trigger_word = selected_lora["trigger_word"]
        token = os.getenv("API_TOKEN")
        payload = {"inputs": f"{prompt} {trigger_word}"}
        
        # API call
        image_bytes = query(payload, api_url, token)
        
        if image_bytes:
            return Image.open(image_bytes)
    return "API Error or No LoRA selected"

# Placeholder for gallery.select function
def update_selection(selected=None):
    if selected is None:
        return None,
    return selected[0],

# Gradio UI
with gr.Blocks(css="custom.css") as app:
    title = gr.HTML("<h1>LoRA the Explorer</h1>")
    selected_state = gr.State(0)  # Initialize with the index of the first LoRA
    with gr.Row():
        gallery = gr.Gallery(
            [(item["image"], item["title"]) for item in loras],
            label="LoRA Gallery",
            allow_preview=False,
            columns=3
        )
        with gr.Column():
            prompt_title = gr.Markdown("### Click on a LoRA in the gallery to select it")
            with gr.Row():
                prompt = gr.Textbox(label="Prompt", show_label=False, lines=1, max_lines=1, placeholder="Type a prompt after selecting a LoRA")
                button = gr.Button("Run")
            result = gr.Image(interactive=False, label="Generated Image")

    gallery.select(
        update_selection,
        outputs=[selected_state]
    )
    button.click(
        fn=run_lora,
        inputs=[prompt, selected_state],
        outputs=[result]
    )

app.queue(max_size=20)
app.launch()