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

# Definindo as informações do repositório e a trigger word
repo = "artificialguybr/StudioGhibli.Redmond-V2"
trigger_word = "Studio Ghibli, StdGBRedmAF"

def generate_image(prompt):
    api_url = f"https://api-inference.huggingface.co/models/{repo}"
    # Descomente e use seu token da API do Hugging Face
    # token = os.getenv("API_TOKEN")
    headers = {
        # "Authorization": f"Bearer {token}"
    }
    full_prompt = f"{prompt} {trigger_word}"
    payload = {
        "inputs": full_prompt,
        "parameters": {
            "negative_prompt": "bad art, ugly, watermark, deformed",
            "num_inference_steps": 30,
            "scheduler": "DPMSolverMultistepScheduler"
        },
    }

    error_count = 0
    pbar = tqdm(total=None, desc="Loading model")
    while True:
        response = requests.post(api_url, headers=headers, json=payload)
        if response.status_code == 200:
            return Image.open(BytesIO(response.content))
        elif response.status_code == 503:
            time.sleep(1)
            pbar.update(1)
        elif response.status_code == 500 and error_count < 5:
            time.sleep(1)
            error_count += 1
        else:
            raise Exception(f"API Error: {response.status_code}")

css = """
body { font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif; background-color: #f0f0f0; }
h1 { color: #333; text-align: center; }
h2 { color: #555; }
.gradio-container { max-width: 800px; margin: auto; padding: 20px; border-radius: 10px; background-color: #fff; box-shadow: 0 0 10px rgba(0,0,0,0.1); }
.gradio-input, .gradio-output { margin-bottom: 20px; }
.gradio-button { background-color: #4CAF50; color: white; padding: 14px 20px; margin: 10px 0; border: none; border-radius: 4px; cursor: pointer; }
.gradio-button:hover { background-color: #45a049; }
"""

iface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(lines=2, placeholder="Enter a Studio Ghibli style prompt..."),
    outputs="image",
    title="Studio Ghibli Style Image Generator",
    description="Type a prompt and generate an image with the enchanting Studio Ghibli aesthetic.",
    css=css
)

iface.launch()