app.py

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

# Tải biến môi trường từ file .env
load_dotenv()

# Lấy API Key và URL từ biến môi trường
api_key = os.getenv("HF_API_KEY")
image_api_url = os.getenv("IMAGE_API_URL")
headers = {"Authorization": f"Bearer {api_key}"}

# Hàm gọi API Hugging Face để tạo hình ảnh
def query(payload):
    response = requests.post(image_api_url, headers=headers, json=payload)
    return response.content

# Hàm dịch tiếng Việt sang tiếng Anh
def translate_to_english(text):
    return translate(text, "en", "vi")  # Dịch từ tiếng Việt sang tiếng Anh

# Hàm xử lý tạo hình ảnh với các lựa chọn
def generate_image(prompt, view, style, hair_color, eye_color, lip_color):
    # Dịch prompt từ tiếng Việt sang tiếng Anh
    prompt_en = translate_to_english(prompt)
    
    # Thêm các yếu tố lựa chọn vào prompt
    prompt_with_choices = f"{prompt_en}, View: {view}, Style: {style}, Hair color: {hair_color}, Eye color: {eye_color}, Lip color: {lip_color}"

    # Thêm yếu tố ngẫu nhiên vào prompt
    random_factor = random.randint(1, 1000)
    prompt_with_randomness = f"{prompt_with_choices}. Seed: {random_factor}"
    
    # Gửi prompt vào API tạo hình ảnh
    image_bytes = query({"inputs": prompt_with_randomness})
    
    # Mở hình ảnh từ byte dữ liệu
    image = Image.open(io.BytesIO(image_bytes))
    return image

# Tạo giao diện Gradio
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        "text",  # Prompt
        gr.Dropdown(["Cận cảnh", "Nửa người", "Toàn thân"], label="View"),
        gr.Dropdown(["Art", "Anime", "Realistic"], label="Style"),
        gr.ColorPicker(label="Màu tóc"),
        gr.ColorPicker(label="Màu mắt"),
        gr.ColorPicker(label="Màu môi")
    ],
    outputs="image",
    title="Image Generator - Trần Như Tuấn",
    description="Chọn các tùy chọn và nhập mô tả để tạo hình ảnh chân dung"
)

# Khởi chạy ứng dụng
iface.launch()