app.py

import gradio as gr  # Để tạo giao diện web
import requests  # Để thực hiện các yêu cầu HTTP
import json  # Để làm việc với dữ liệu JSON
import hashlib  # Để tạo hash MD5
import time  # Để xử lý thời gian
import os  # Để làm việc với biến môi trường
from PIL import Image  # Để mở và xử lý hình ảnh
from pathlib import Path  # Để làm việc với đường dẫn
from io import BytesIO  # Để xử lý dữ liệu binary như hình ảnh


images = ["bxh.jpg", "bxh1.jpg", "bxh2.jpg", "bxh3.jpg"]  # Các file ảnh của bạn

# URL và thông tin cá nhân từ API
url_pre = "https://ap-east-1.tensorart.cloud/v1"

# Thư mục để lưu ảnh đã tạo
SAVE_DIR = "generated_images"
Path(SAVE_DIR).mkdir(exist_ok=True)

# Lấy API key từ biến môi trường
api_key_token = os.getenv("api_key_token")
if not api_key_token:
    raise ValueError("API key token không tìm thấy trong biến môi trường.")

# Hàm yêu cầu API để tạo ảnh
def txt2img(prompt, width, height):
    model_id = "770694094415489962"  # Model ID cố định
    vae_id = "sdxl-vae-fp16-fix.safetensors"  # VAE cố định
    lora_items = [
        {"loraModel": "808188171067237481", "weight": 1.0}
    ]

    txt2img_data = {
        "request_id": hashlib.md5(str(int(time.time())).encode()).hexdigest(),
        "stages": [
            {
                "type": "INPUT_INITIALIZE",
                "inputInitialize": {
                    "seed": -1,
                    "count": 1
                }
            },
            {
                "type": "DIFFUSION",
                "diffusion": {
                    "width": width,
                    "height": height,
                    "prompts": [
                        {
                            "text": prompt
                        }
                    ],
                    "negativePrompts": [
                        {
                            "text": "nsfw"
                        }
                    ],
                    "sdModel": model_id,
                    "sdVae": vae_id,
                    "sampler": "Euler a",
                    "steps": 20,
                    "cfgScale": 3,
                    "clipSkip": 1,
                    "etaNoiseSeedDelta": 31337,
                    "lora": {
                        "items": lora_items
                    }
                }
            }
        ]
    }

    body = json.dumps(txt2img_data)
    
    # Sử dụng Bearer token để xác thực
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json',
        'Authorization': f'Bearer {api_key_token}'
    }

    # Tạo một job mới
    response = requests.post(f"{url_pre}/jobs", json=txt2img_data, headers=headers)

    if response.status_code != 200:
        return f"Error: {response.status_code} - {response.text}"

    response_data = response.json()
    job_id = response_data['job']['id']
    print(f"Job created. ID: {job_id}")

    # Kiểm tra kết quả của job
    start_time = time.time()
    timeout = 300  # Giới hạn thời gian chờ là 300 giây (5 phút)

    while True:
        time.sleep(10)  # Chờ 10 giây giữa các lần kiểm tra

        # Kiểm tra nếu quá thời gian chờ
        elapsed_time = time.time() - start_time
        if elapsed_time > timeout:
            return f"Error: Job timed out after {timeout} seconds."

        # Gửi yêu cầu để lấy trạng thái của job
        response = requests.get(f"{url_pre}/jobs/{job_id}", headers=headers)

        if response.status_code != 200:
            return f"Error: {response.status_code} - {response.text}"

        get_job_response_data = response.json()
        print("Job response data:", get_job_response_data)
        job_status = get_job_response_data['job']['status']
        print(f"Job status: {job_status}")

        if job_status == 'SUCCESS':
            if 'successInfo' in get_job_response_data['job']:  # Kiểm tra nếu 'successInfo' tồn tại
                image_url = get_job_response_data['job']['successInfo']['images'][0]['url']
                print(f"Job succeeded. Image URL: {image_url}")
                
                # Tải ảnh từ URL và trả về dưới dạng ảnh PIL để hiển thị trong Gradio
                response_image = requests.get(image_url)
                img = Image.open(BytesIO(response_image.content))
                return img  # Trả về ảnh dưới dạng đối tượng PIL
            else:
                return "Error: Output is missing in the job response."
        elif job_status == 'FAILED':
            return "Error: Job failed. Please try again with different settings."
        else:
            print("Job is still in progress...")

# Cài đặt giao diện Gradio
with gr.Blocks() as demo:
    gr.Markdown("""
    # Flux ENHANCED ALL - Realistic
    <p style='font-size: 14px;'>Feel free to explore my work and contact me for any inquiries via email: dung.ngt1988@gmail.com</p>
    """)    
    # Sử dụng hai cột để sắp xếp giao diện
    with gr.Row():
        with gr.Column():
            # Cột thứ nhất: mô tả ảnh và chọn kích thước
            prompt_input = gr.Textbox(label="Mô tả ảnh cần tạo", 
                                      placeholder="A cover image for podcast at Youtube channel about buddhism with a girl stands on the left of image. She points a big perfect typography '4 things Buddha can not help' on the center of image")
            size_options = gr.Dropdown(choices=["1152x768","1024x1024", "768x1152"], label="Chọn cỡ ảnh")
            generate_button = gr.Button("Generate")  # Nút tạo ảnh

        with gr.Column():
            # Cột thứ hai: hiển thị ảnh đã tạo
            output_image = gr.Image(label="Ảnh đã tạo")  # Để hiển thị ảnh sau khi tạo

    # Hàm để sinh ảnh khi nhấn nút
    def generate(prompt, size_choice):
        if size_choice == "1152x768":
            width, height = 1152, 768
        elif size_choice == "768x1152":
            width, height = 768, 1152
        elif size_choice == "1024x1024":
            width, height = 1024, 1024
        else:
            raise ValueError("Invalid size choice")

        return txt2img(prompt, width, height)
    
    # Kết nối sự kiện của nút sinh ảnh với hàm generate
    generate_button.click(fn=generate, inputs=[prompt_input, size_options], outputs=output_image)

# Chạy ứng dụng
demo.launch()