import gradio as gr
import requests
from PIL import Image
from io import BytesIO
import os
from huggingface_hub import InferenceClient
API_TOKEN = os.getenv("HF_API_TOKEN")  # Ensure you've set this environment variable
API_URL = "https://api-inference.huggingface.co/models/enhanceaiteam/Flux-uncensored"

enhancer = InferenceClient(api_key =API_TOKEN)

for message in enhancer.chat_completion(
	model="meta-llama/Llama-3.2-1B-Instruct",
	messages=[{"role": "user", "content": "What is the capital of France?"}],
	max_tokens=500,
	stream=True,
):
    print(message.choices[0].delta.content, end="")
# Load API Token from environment variable

# Hugging Face Inference API URL

# Function to call Hugging Face API and get the generated image
def generate_image(prompt):
    headers = {"Authorization": f"Bearer {API_TOKEN}"}
    data = {"inputs": prompt}
    
    response = requests.post(API_URL, headers=headers, json=data)
    
    if response.status_code == 200:
        image_bytes = BytesIO(response.content)
        image = Image.open(image_bytes)
        return image
    else:
        return f"Error: {response.status_code}, {response.text}"
 
title_html="""
    <center>
        <div id="title-container">
            <h1 id="title-text">FLUX Capacitor</h1>
        </div>
    </center>
"""

css = """
.gradio-container {
    background: url(https://huggingface.co/spaces/K00B404/FLUX.1-Dev-Serverless-darn-enhanced-prompt/resolve/main/edge.png);
    background-size: 900px 880px;
    background-repeat: no-repeat;
    background-position: center;
    background-attachment: fixed;
    color:#000;
}
.dark\:bg-gray-950:is(.dark *) {
  --tw-bg-opacity: 1;
  background-color: rgb(157, 17, 142);
}
.gradio-container-4-41-0 .prose :last-child {
  margin-top: 8px !important;
}
.gradio-container-4-41-0 .prose :last-child {
  margin-bottom: -7px !important;
}
.dark {
    --button-primary-background-fill: #09e60d70;
    --button-primary-background-fill-hover: #00000070;
    --background-fill-primary: #000;
    --background-fill-secondary: #000;
}
.hide-container {
    margin-top;-2px;
}
#app-container3 {
    background-color: rgba(255, 255, 255, 0.001);  /* Corrected to make semi-transparent */
    max-width: 600px;
    margin-left: auto;
    margin-right: auto;
    margin-bottom: 10px;
    border-radius: 125px;
    box-shadow: 0 0 10px rgba(0,0,0,0.1); /* Adjusted shadow opacity */
}
#app-container {
    background-color: rgba(255, 255, 255, 0.001);  /* Semi-transparent background */
    max-width: 600px;
    margin: 0 auto;  /* Center horizontally */
    padding-bottom: 10px;
    border-radius: 25px;
    box-shadow: 0 0 10px rgba(0, 0, 0, 0.1); /* Adjusted shadow opacity */
}
#title-container {
    display: flex;
    align-items: center
    margin-bottom:10px;
    justify-content: center;
}
#title-icon {
    width: 32px;
    height: auto;
    margin-right: 10px;
}
#title-text {
    font-size: 30px;
    font-weight: bold;
    color: #000;
}
"""





# Create Gradio interface
def create_ui():
    with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as ui:
        gr.Markdown("## Flux Uncensored - Text to Image Generator")
        
        with gr.Row():
            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
        generate_button.click(fn=generate_image, inputs=prompt_input, outputs=output_image)
    
    return ui

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