import os
import requests
import gradio as gr
from dotenv import load_dotenv  
import io
from PIL import Image
from mtranslate import translate  

load_dotenv()

api_key = os.getenv("HF_API_KEY")
image_api_url = os.getenv("model") 
headers = {"Authorization": f"Bearer {api_key}"}

censored_words = [
    "sex", "porn", "asshole", "fuck", "cock", "pussy", "dick", "cum","nude", "naked", "fetish"
]

def query(payload):
    response = requests.post(image_api_url, headers=headers, json=payload)
    return response.content

def translate_to_english(text):
    return translate(text, "en", "vi")  

def contains_censored_words(text):
    text = text.lower()
    for word in censored_words:
        if word in text:
            return True
    return False

def generate_image(prompt):
    prompt_en = translate_to_english(prompt)
    
    if contains_censored_words(prompt_en):
        return "Mô tả chứa từ cấm, vui lòng thử lại với mô tả khác."
    
    image_bytes = query({"inputs": prompt_en})
    image = Image.open(io.BytesIO(image_bytes))
    return image

iface = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="image",
    title="Tạo sinh hình ảnh | 0379031662",
    description = """
    # Từ ý tưởng ra hình ảnh  
    **Lưu ý:** Hãy sử dụng từ ngữ đơn giản, có thể gõ bằng Tiếng Việt.  
    **Tác giả:** Như Tuấn  
    **Instagram:** [dk.darkkool](https://instagram.com/dk.darkkool)  
    **Facebook:** [darkkool](https://facebook.com/darkkool)  
    **Website:** [kooldstudio.pixieset.com](https://kooldstudio.pixieset.com)  
    """    
)

iface.launch()