# import gradio as gr

# gr.Interface.load("models/Akbartus/Lora360").launch(show_api=True)

import gradio as gr
import requests
import io
from PIL import Image
import json
import os
import logging
import math 
from tqdm import tqdm
import time

#logging.basicConfig(level=logging.DEBUG)

with open('loras.json', 'r') as f:
    loras = json.load(f)

# Select the default LoRA
default_lora = loras[0]  # Assuming the first LoRA is the default one

def run_lora(prompt, progress=gr.Progress(track_tqdm=True)):
    logging.debug(f"Inside run_lora")
    api_url = f"https://api-inference.huggingface.co/models/{default_lora['repo']}"
    trigger_word = default_lora["trigger_word"]
    payload = {
        "inputs": f"{prompt} {trigger_word}",
        "parameters":{"negative_prompt": "bad art, ugly, watermark, deformed", "num_inference_steps": 30, "scheduler":"DPMSolverMultistepScheduler"},
    }
    
    # Add a print statement to display the API request
    print(f"API Request: {api_url}")
    print(f"API Payload: {payload}")

    error_count = 0
    pbar = tqdm(total=None, desc="Loading model")
    while(True):
        response = requests.post(api_url, json=payload)
        if response.status_code == 200:
            return Image.open(io.BytesIO(response.content))
        elif response.status_code == 503:
            time.sleep(1)
            pbar.update(1)
        elif response.status_code == 500 and error_count < 5:
            print(response.content)
            time.sleep(1)
            error_count += 1
            continue
        else:
            logging.error(f"API Error: {response.status_code}")
            raise gr.Error("API Error: Unable to fetch the image.")  # Raise a Gradio error here

with gr.Blocks(css="custom.css") as app:
    title = gr.Markdown("# LoRA 360 Demonstration")
    description = gr.Markdown(
        "### Lora 360 demonstration and API endpoint."
    )
    with gr.Row():
        prompt_title = gr.Markdown(f"### Type a prompt for {default_lora['title']}")
        with gr.Row():
            prompt = gr.Textbox(label="Prompt", show_label=False, lines=1, max_lines=1, placeholder=f"Type a prompt for {default_lora['title']}")
            button = gr.Button("Run")
        result = gr.Image(interactive=False, label="Generated Image")

    prompt.submit(
        fn=run_lora,
        inputs=[prompt],
        outputs=[result]
    )
    button.click(
        fn=run_lora,
        inputs=[prompt],
        outputs=[result]
    )

app.queue(max_size=20, concurrency_count=5)
app.launch()