import os
import gradio as gr
import requests
import base64

# Get API key from environment variable
api_key = os.environ.get("NVCF_API_KEY")

if not api_key:
    raise ValueError("Please set the NVCF_API_KEY environment variable.")

# API details
invoke_url = "https://api.nvcf.nvidia.com/v2/nvcf/pexec/functions/89848fb8-549f-41bb-88cb-95d6597044a4"
fetch_url_format = "https://api.nvcf.nvidia.com/v2/nvcf/pexec/status/"
headers = {
    "Authorization": f"Bearer {api_key}",
    "Accept": "application/json",
}

# Function to generate image using the API
def generate_image(prompt, negative_prompt, sampler, seed, guidance_scale, inference_steps):
    payload = {
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "sampler": sampler,
        "seed": seed,
        "guidance_scale": guidance_scale,
        "inference_steps": inference_steps
    }

    session = requests.Session()
    response = session.post(invoke_url, headers=headers, json=payload)

    while response.status_code == 202:
        request_id = response.headers.get("NVCF-REQID")
        fetch_url = fetch_url_format + request_id
        response = session.get(fetch_url, headers=headers)


    response.raise_for_status()
    response_body = response.json()

    # Print the API response for debugging
    print("API Response:", response_body)

    # Decode the base64-encoded image data
    b64_image_data = response_body.get("b64_json")
    if b64_image_data is None:
        return "Error: API response does not contain 'b64_json' key."
    
    image_data = base64.b64decode(b64_image_data)
    
    # Since Gradio's Image component can accept a PIL Image, convert the binary data to a PIL Image
    image = Image.open(BytesIO(image_data))
    
    return image


# Create Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Describe the image you want to generate"),
        gr.Textbox(label="Negative Prompt", placeholder="What should not be in the image"),
        gr.Dropdown(label="Sampler", choices=["DPM", "DDPM", "PLMS"], value="DPM"),
        gr.Number(label="Seed", value=0),
        gr.Slider(label="Guidance Scale", minimum=0, maximum=20, value=5),
        gr.Slider(label="Inference Steps", minimum=1, maximum=50, value=25)
    ],
    outputs=gr.Image(label="Generated Image"),
)

# Launch the Gradio app
iface.launch()