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

# Load API Token from environment variable
API_TOKEN = os.getenv("HF_API_TOKEN")  # Ensure you've set this environment variable

# Hugging Face Inference API URL
API_URL = "https://api-inference.huggingface.co/models/enhanceaiteam/Flux-uncensored"

# Function to call Hugging Face API and get the generated image
def generate_image(prompt, progress=gr.Progress()):
    headers = {"Authorization": f"Bearer {API_TOKEN}"}
    data = {"inputs": prompt}
    
    # Initialize progress
    progress(0, desc="Starting image generation...")
    
    response = requests.post(API_URL, headers=headers, json=data)
    
    if response.status_code == 200:
        progress(50, desc="Processing image...")
        image_bytes = BytesIO(response.content)
        image = Image.open(image_bytes)
        progress(100, desc="Completed")
        return image
    else:
        progress(100, desc="Error occurred")
        return f"Error: {response.status_code}, {response.text}"

# Create Gradio interface
def create_ui():
    with gr.Blocks(theme="hev832/Applio") as ui:
        gr.Markdown("## Flux Uncensored\nUnofficial Gradio Demo")
        
        prompt_input = gr.Textbox(label="Enter a Prompt", placeholder="Describe the image you want to generate", lines=3)
        generate_button = gr.Button("Generate Image")
        
        with gr.Row():
            output_image = gr.Image(label="Generated Image")
        
        # Link the button to the function with progress
        generate_button.click(fn=generate_image, inputs=prompt_input, outputs=output_image, show_progress=True)
    
    return ui

# Run the interface
if __name__ == "__main__":
    create_ui().launch()