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



api_key = os.environ.get("HF_API_KEY")
API_URL = "https://api-inference.huggingface.co/models/multimodalart/vintage-ads-flux"
headers = {"Authorization": f"Bearer {api_key}"}

# Function to query the model
def query_image(inputs):
    payload = {"inputs": inputs}
    response = requests.post(API_URL, headers=headers, json=payload)
    image_bytes = response.content
    image = Image.open(io.BytesIO(image_bytes))
    return image

# Define Gradio Blocks UI
with gr.Blocks(theme="nevreal/blues") as demo:
    gr.Markdown("# Vintage Ads Image Generator")
    with gr.Row():
        gr.Markdown("model by [multimodalart](https://huggingface.co/multimodalart) this spces by [nevreal](https://huggingface.co/nevreal)")

    with gr.Row():
        prompt = gr.Textbox(label="Enter your prompt", value="Astronaut riding a horse", lines=2)
        submit_button = gr.Button("Generate Image")
        
        with gr.Column():
            output_image = gr.Image(label="Generated Image")

    # Add examples
    examples = gr.Examples(
        examples=[
            "a vintage ad with a woman wearing a VR headset with an apple logo and laughing, vtgads; it reads 'Vision Pro' on the top and 'escape the real world' in the bottom",
            "a vintage ad of Twitter, the twitter bird with a speech baloon that reads 'I will never be X'",
            "a vintage ad of neuralink, a person with a brain computer implant, vtgads; it reads 'neuralink' on the top and 'if you can't defeat the AI, join it!'"
        ],
        inputs=prompt,
    )
    
    # Link button action to function
    submit_button.click(fn=query_image, inputs=prompt, outputs=output_image)

# Launch the web UI
demo.launch()