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

# 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):
    if seed is None or seed == 0:
        seed = random.randint(11111111, 999999999999999)
    payload = {
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "sampler": sampler,
        "seed": seed,
        "guidance_scale": guidance_scale,
        "inference_steps": inference_steps
    }
    print(payload)
    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)
    
    # 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", "EulerA", "LMS", "DDIM"], value="DPM"),
        gr.Number(label="Seed", value=0),
        gr.Slider(label="Guidance Scale", minimum=1, maximum=9, value=5),
        gr.Slider(label="Inference Steps", minimum=5, maximum=100, value=35)
    ],
    outputs=gr.Image(label="Generated Image"),
    description = """
        <div style="text-align: center; font-size: 1.5em; margin-bottom: 20px;">
            <strong>Generate Stunning Images with Stable Diffusion XL</strong>
        </div>
    
        <p>
            This Gradio app harnesses the power of Stable Diffusion XL image generation capabilities to bring your creative visions to life. Using NVIDIA NGC. 
            Simply provide a text prompt describing the image you desire, and let the AI do its magic!
        </p>
    
        <p>
            <strong>How to Use:</strong>
        </p>
    
        <ol>
            <li>Enter a detailed <strong>prompt</strong> describing the image you want to generate.</li>
            <li>Optionally, add a <strong>negative prompt</strong> to specify elements you want to avoid.</li>
            <li>Choose a <strong>sampler</strong> (algorithm) for image generation.</li>
            <li>Set a <strong>seed</strong> for reproducibility (or leave it blank for random results).</li>
            <li>Adjust the <strong>guidance scale</strong> to influence how closely the image follows the prompt.</li>
            <li>Set the <strong>inference steps</strong> to control the level of detail and processing time.</li>
            <li>Click <strong>Generate</strong> and marvel at your creation!
        </ol>
    
        <p>
            <strong>This service is powered by NVIDIA NGC and is completely free to use.</strong>
        </p>
    
        <p>
            <strong>Created by:</strong> @artificialguybr (<a href="https://twitter.com/artificialguybr">Twitter</a>)
        </p>
    
        <p>
            <strong>Explore more:</strong> <a href="https://artificialguy.com">artificialguy.com</a>
        </p>
    """
)

# Launch the Gradio app
iface.launch()