import gradio as gr
import requests
import io
import os
import logging
from PIL import Image, ImageEnhance
from gradio_client import Client

# API URLs and headers
API_URL_1 = "https://api-inference.huggingface.co/models/sd-community/sdxl-flash"

API_KEY = os.getenv("HF_API_KEY")
headers = {"Authorization": f"Bearer {API_KEY}"}

client = Client("AP123/SDXL-Lightning")

def query(api_url, payload):
    try:
        response = requests.post(api_url, headers=headers, json=payload)
        response.raise_for_status()
        return response.content
    except requests.exceptions.HTTPError as http_err:
        logging.error(f"HTTP error occurred: {http_err}")
        print(f"HTTP error occurred: {http_err}")
    except Exception as err:
        logging.error(f"An error occurred: {err}")
        print(f"An error occurred: {err}")
    return None

def map_inference_steps(steps):
    if steps <= 1:
        return "1-Step"
    elif steps <= 2:
        return "2-Step"
    elif steps <= 4:
        return "4-Step"
    else:
        return "8-Step"

def gen_img(prompt, num_inference_steps, guidance_scale, contrast):
    payload = {
        "inputs": prompt,
        "parameters": {
            "num_inference_steps": num_inference_steps,
            "guidance_scale": guidance_scale
        }
    }

    # Try generating image with the first model
    image_bytes = query(API_URL_1, payload)
    if image_bytes is None:
        # print("First model failed, trying second model...")
        try:
            mapped_steps = map_inference_steps(num_inference_steps)
            result = client.predict(prompt, mapped_steps, api_name="/generate_image")
            with open(result, "rb") as img_file:
                image = Image.open(img_file)
                enhancer = ImageEnhance.Contrast(image)
                image = enhancer.enhance(contrast)
                return image, "Image generated by the second model"
        except Exception as e:
            # logging.error(f"Error with the second model: {e}")
            # print(f"Error with the second model: {e}")
            return None, "Error with the second model"
    
    try:
        image = Image.open(io.BytesIO(image_bytes))
        enhancer = ImageEnhance.Contrast(image)
        image = enhancer.enhance(contrast)
        return image, "Image generated by the first model"
    except Exception as e:
        # logging.error(f"Error processing image: {e}")
        return None, "Error processing image"

with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.lime, spacing_size="md", radius_size="lg", 
                                    font=[gr.themes.GoogleFont("Copperplate"), "Fantasy", "sans-serif"])) as demo:
    
    gr.Markdown("""
    <div style='text-align: center; 
                font-size: 36px; 
                font-family: Copperplate, Fantasy; 
                color: orange;'>
        Pencil Vision (early beta 1.3)
    </div>
    """)
    with gr.Row():
        prompt = gr.Textbox(label="Enter your prompt to generate the image", placeholder="A fantasy landscape with mountains and a river")

    with gr.Row():
        num_inference_steps = gr.Slider(minimum=1, maximum=15, value=6, step=1, label="Number of Inference Steps")
        guidance_scale = gr.Slider(minimum=1.0, maximum=5.0, value=2.5, step=0.1, label="Guidance Scale")
        contrast_scale = gr.Slider(minimum=1.0, maximum=3.0, value=1.0, step=0.1, label="Contrast Scale")  

    with gr.Row():
        generate_button = gr.Button("Generate Image")

    output_image = gr.Image(label="Generated Image")
    # model_info = gr.Textbox(label="Model Info", interactive=False)
    error_message = gr.Markdown(visible=False)
    
    def handle_click(prompt, num_inference_steps, guidance_scale, contrast_scale):
        image, info = gen_img(prompt, num_inference_steps, guidance_scale, contrast_scale)
        if image is None:
            return None, info, gr.update(visible=True, value="Error generating image. Please try again.")
        return image, info, gr.update(visible=True)
    
    generate_button.click(handle_click, inputs=[prompt, num_inference_steps, guidance_scale, contrast_scale], outputs=[output_image, error_message])

demo.launch(share=True)