import gradio as gr
import random
import requests
import io
from PIL import Image, UnidentifiedImageError
import os

API_URL = "https://api-inference.huggingface.co/models/LucyintheSky/lucy-dream-lora"
headers = {"Authorization": "Bearer " + os.environ.get('TOKEN')}

def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    if response.ok:
        return response.content
    else:
        print("API Error:", response.status_code, response.text)
        return None

def generate(prompt):
    image_bytes = query({
        "inputs": f"lucy dream style {prompt}, perfect body, beautiful face, perfect skin",
        "parameters": {"negative_prompt": "ugly, deformed, deformed face, ugly face, bad quality",
                        "seed": random.randint(0, 9999999)}
    })
    
    if image_bytes is None:
        return "Error: Could not generate image. Please try again."
    
    try:
        image = Image.open(io.BytesIO(image_bytes))
        image_path = "temp_output.png"
        image.save(image_path)
        return image_path
    except UnidentifiedImageError:
        return "Error: Invalid image data received."

theme = gr.themes.Base(
    primary_hue="gray",
    secondary_hue="gray",
    neutral_hue="gray",
    font=['Helvetica', 'ui-sans-serif', 'system-ui', 'sans-serif'],
).set(
    button_large_text_weight='400',
    input_background_fill='#ffffff',
    button_primary_background_fill='#ffd8db',
)

with gr.Blocks(theme=theme) as demo:
    gr.Markdown("""
    <center> <img src="https://cdn-uploads.huggingface.co/production/uploads/650a1281b48ff3906647edd0/rbzPKW_p5XTG9leEO4uFt.png" > </center>
    <br><br><br>
    """)
    img = gr.Image(show_label=False, type="filepath")
    textbox = gr.Textbox(show_label=False, placeholder='Type your prompt here')
    button = gr.Button("Generate", variant="primary")

    gr.Examples(
        [["Emma Stone on prom night, mint green satin maxi dress"], 
         ["Lady Gaga, 50s hair style, prom night"], 
         ["Rihanna wearing a silver prom dress with crystal jewelry"]],
        inputs=textbox,
        outputs=img,
        fn=generate,
        cache_examples=False  # キャッシュ機能をオフにする
    )

    
    gr.Markdown("""
    ![image/png](https://cdn-uploads.huggingface.co/production/uploads/650a1281b48ff3906647edd0/OLPLii-mE6VhIiabvJwOm.png)
    """)
    
    button.click(fn=generate, inputs=textbox, outputs=img)
    textbox.submit(fn=generate, inputs=textbox, outputs=img)

demo.launch()