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"

# Função para enviar o prompt para a API e receber a imagem gerada
def generate_image(prompt):
    api_url = f"https://api-inference.huggingface.co/models/{repo}"
    #token = os.getenv("API_TOKEN")  # Substitua pelo seu token da API do Hugging Face
    headers = {
        #"Authorization": f"Bearer {token}"
    }
    full_prompt = f"{prompt} {trigger_word}"  # Combina o prompt do usuário com a 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:
        try:
            response = requests.post(api_url, headers=headers, json=payload)
            if response.status_code == 200:
                image_data = response.json().get("data", [])[0]  # Obtém os dados da imagem
                return Image.open(BytesIO(image_data))
            elif response.status_code == 503:
                # 503 é quando o modelo está em inicialização
                time.sleep(1)
                pbar.update(1)
            elif response.status_code == 500 and error_count < 5:
                # Erro interno do servidor, tenta novamente até 5 vezes
                time.sleep(1)
                error_count += 1
            else:
                raise Exception(f"API Error: {response.status_code}")
        except requests.exceptions.HTTPError as http_err:
            print(f"HTTP error occurred: {http_err}")
        except requests.exceptions.RequestException as err:
            print(f"An error occurred: {err}")
        except requests.exceptions.JSONDecodeError as json_err:
            print(f"JSON decode error: {json_err}")
            # Se a resposta não for JSON, tenta carregar a imagem diretamente
            return Image.open(BytesIO(response.content))
        except Exception as e:
            print(f"An unexpected error occurred: {e}")
            break  # Sai do loop se um erro inesperado ocorrer

# Criando a interface do Gradio
iface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(lines=2, placeholder="Digite seu prompt aqui..."),
    outputs="image",
    title="Gerador de Imagens Studio Ghibli V2",
    description="Insira um prompt de texto para gerar uma imagem com o estilo do Studio Ghibli."
)

# Iniciando a interface
iface.launch()