import os
import base64
import requests
from PIL import Image
from io import BytesIO
import gradio as gr

hf_token = os.environ.get('hf_token')
API_URL = os.environ.get('api_url')

headers = {
	"Accept" : "application/json",
	"Authorization": f"Bearer {hf_token}",
	"Content-Type": "application/json" 
}

# helper to decode input image
def decode_base64_image(image_string):
    base64_image = base64.b64decode(image_string)
    buffer = BytesIO(base64_image)
    image = Image.open(buffer)
    return image

def predict(prompt, temperature, guidance_scale, num_inference_steps, seed):
    payload = {
        "inputs": prompt,
        "temperature": temperature,
        "guidance_scale": guidance_scale,
        "num_inference_steps": num_inference_steps,
        "seed": seed
        }
    response = requests.post(API_URL, headers=headers, json=payload)
    output = response.json()
    image = decode_base64_image(image_string=output)
    return image

demo = gr.Interface(
    predict,
    inputs=[
        gr.Textbox(label='Prompt'),
        gr.Slider(0.1, 5.0, label='Temperature', step=0.1, value=1.0),
        gr.Slider(0.1, 15.0, label='Guidance Scale', step=0.1, value=7.5),
        gr.Slider(1, 100, label='Number of inference steps', step=1, value=50),
        gr.Slider(1, 100, label='Seed', step=1, value=42),
    ],
    outputs="image",
    title="Android toy SD demo",
    description="Generate images with android toy!! Just specify an android toy somewhere in the prompt or click on the example provided below. The first run could take up to 2 minutes if the model is sleeping, then approximately 9 seconds per one request",
    examples=[["An android toy near Eiffel Tower", 1.0, 7.5, 50, 42],
              ["A blue android toy in snow near Eiffel Tower in winter", 1.0, 7.5, 50, 42],
              ["An android toy on top of a brick", 1.0, 7.5, 50, 42]],
    cache_examples=True
    )
demo.launch()