app.py

import gradio as gr
import requests
import os
import time
import logging
from PIL import Image
import mimetypes

# Cấu hình logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Cấu hình API TDNM
VIDU_API_KEY = os.getenv("VIDU_API_KEY")
TDNM_KEY = os.getenv("TDNM_KEY")
VIDU_API_URL = "https://api.vidu.com"
POLL_INTERVAL = 5  # Giây giữa các lần kiểm tra trạng thái
TIMEOUT = 300  # Thời gian chờ tối đa để tạo video
DEFAULT_MODEL = "vidu2.0"  # Mô hình mặc định

# Hàm kiểm tra TDNM_KEY
def validate_tdn_key(user_key):
    if not TDNM_KEY:
        return False, "Lỗi: Khóa bí mật chưa được cấu hình trên máy chủ."
    if not user_key:
        return False, "Lỗi: Vui lòng nhập khóa bí mật."
    if user_key != TDNM_KEY:
        return False, "Lỗi: Khóa bí mật không hợp lệ."
    return True, "Khóa bí mật được xác thực thành công."

# Hàm kiểm tra yêu cầu hình ảnh
def validate_image(image_path, max_size_mb=10, min_dimensions=(128, 128)):
    if not image_path:
        return False, "Lỗi: Chưa cung cấp hình ảnh."
    
    # Kiểm tra kích thước tệp
    if os.path.getsize(image_path) > max_size_mb * 1024 * 1024:
        return False, f"Lỗi: Kích thước hình ảnh vượt quá {max_size_mb}MB."
    
    # Kiểm tra định dạng
    mime_type, _ = mimetypes.guess_type(image_path)
    if mime_type not in ["image/png", "image/webp", "image/jpeg", "image/jpg"]:
        return False, "Lỗi: Định dạng hình ảnh không được hỗ trợ. Sử dụng PNG, WebP, JPEG hoặc JPG."
    
    # Kiểm tra kích thước và tỷ lệ khung hình
    try:
        img = Image.open(image_path)
        width, height = img.size
        if width < min_dimensions[0] or height < min_dimensions[1]:
            return False, f"Lỗi: Kích thước hình ảnh phải ít nhất {min_dimensions[0]}x{min_dimensions[1]} pixel."
        aspect_ratio = width / height
        if aspect_ratio < 0.25 or aspect_ratio > 4:
            return False, "Lỗi: Tỷ lệ khung hình phải nằm trong khoảng 1:4 đến 4:1."
        return True, None
    except Exception as e:
        return False, f"Lỗi khi kiểm tra hình ảnh: {str(e)}"

# Hàm kiểm tra tỷ lệ mật độ điểm ảnh (cho start-end)
def validate_pixel_density(start_image, end_image):
    try:
        start_img = Image.open(start_image)
        end_img = Image.open(end_image)
        start_pixels = start_img.size[0] * start_img.size[1]
        end_pixels = end_img.size[0] * end_img.size[1]
        ratio = start_pixels / end_pixels
        if not (0.8 <= ratio <= 1.25):
            return False, "Lỗi: Tỷ lệ mật độ điểm ảnh giữa hai hình ảnh phải nằm trong khoảng 0.8 đến 1.25."
        return True, None
    except Exception as e:
        return False, f"Lỗi khi kiểm tra mật độ điểm ảnh: {str(e)}"

# Hàm tải hình ảnh lên TDNM
def upload_image_to_vidu(image_path, max_size_mb=10):
    if not VIDU_API_KEY:
        return None, "Lỗi: Khóa API TDNM chưa được cấu hình."
    
    # Kiểm tra hình ảnh
    valid, error_message = validate_image(image_path, max_size_mb)
    if not valid:
        return None, error_message
    
    # Bước 1: Tạo liên kết tải lên
    url = f"{VIDU_API_URL}/tools/v2/files/uploads"
    headers = {
        "Authorization": f"Token {VIDU_API_KEY}",
        "Content-Type": "application/json"
    }
    payload = {"scene": "vidu"}
    
    try:
        logger.info("Tạo liên kết tải lên")
        response = requests.post(url, json=payload, headers=headers)
        response.raise_for_status()
        result = response.json()
        put_url = result.get("put_url")
        resource_id = result.get("id")
        if not put_url or not resource_id:
            return None, "Lỗi: Không nhận được put_url hoặc resource_id."
    except requests.exceptions.RequestException as e:
        logger.error(f"Lỗi tạo liên kết tải lên: {str(e)}")
        return None, f"Lỗi khi tạo liên kết tải lên: {str(e)}"
    
    # Bước 2: Tải hình ảnh lên put_url
    mime_type, _ = mimetypes.guess_type(image_path)
    with open(image_path, "rb") as f:
        image_data = f.read()
    
    try:
        logger.info("Tải hình ảnh lên put_url")
        response = requests.put(put_url, data=image_data, headers={"Content-Type": mime_type})
        response.raise_for_status()
        etag = response.headers.get("etag")
        if not etag:
            return None, "Lỗi: Không nhận được etag."
    except requests.exceptions.RequestException as e:
        logger.error(f"Lỗi tải hình ảnh: {str(e)}")
        return None, f"Lỗi khi tải hình ảnh: {str(e)}"
    
    # Bước 3: Hoàn tất tải lên
    url = f"{VIDU_API_URL}/tools/v2/files/uploads/{resource_id}/finish"
    payload = {"etag": etag.strip('"')}
    
    try:
        logger.info("Hoàn tất tải lên")
        response = requests.put(url, json=payload, headers=headers)
        response.raise_for_status()
        result = response.json()
        uri = result.get("uri")
        if not uri:
            return None, "Lỗi: Không nhận được URI."
        return uri, "Tải hình ảnh lên thành công."
    except requests.exceptions.RequestException as e:
        logger.error(f"Lỗi hoàn tất tải lên: {str(e)}")
        return None, f"Lỗi khi hoàn tất tải lên: {str(e)}"

# Hàm gọi API TDNM cho Start-End to Video
def start_end_to_video(start_image, end_image, prompt, resolution="720p", duration=4, movement_amplitude="auto", seed=None, user_key=None):
    # Kiểm tra TDNM_KEY
    valid_key, key_message = validate_tdn_key(user_key)
    if not valid_key:
        return None, key_message
    
    if not VIDU_API_KEY:
        return None, "Lỗi: Khóa API TDNM chưa được cấu hình."
    
    if not start_image or not end_image:
        return None, "Lỗi: Cần cung cấp cả hai hình ảnh đầu và cuối."
    
    if prompt and len(prompt) > 1500:
        return None, "Lỗi: Mô tả văn bản không được vượt quá 1500 ký tự."
    
    # Kiểm tra mật độ điểm ảnh
    valid, error_message = validate_pixel_density(start_image, end_image)
    if not valid:
        return None, error_message
    
    # Tải hình ảnh lên TDNM
    start_uri, start_message = upload_image_to_vidu(start_image, max_size_mb=10)
    if not start_uri:
        return None, start_message
    end_uri, end_message = upload_image_to_vidu(end_image, max_size_mb=10)
    if not end_uri:
        return None, end_message
    
    url = f"{VIDU_API_URL}/ent/v2/start-end2video"
    headers = {
        "Authorization": f"Token {VIDU_API_KEY}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "model": DEFAULT_MODEL,
        "images": [start_uri, end_uri],
        "prompt": prompt or "",
        "duration": duration,
        "resolution": resolution,
        "movement_amplitude": movement_amplitude
    }
    if seed is not None:
        payload["seed"] = seed
    
    try:
        logger.info(f"Gửi yêu cầu đến API TDNM Start-End to Video: {payload}")
        response = requests.post(url, json=payload, headers=headers)
        response.raise_for_status()
        result = response.json()
        task_id = result.get("task_id")
        if not task_id:
            return None, "Lỗi: Không nhận được ID tác vụ."
        return task_id, f"Tác vụ được tạo thành công. ID tác vụ: {task_id}"
    except requests.exceptions.RequestException as e:
        logger.error(f"Lỗi API: {str(e)}")
        if response.text:
            logger.error(f"Phản hồi API: {response.text}")
        return None, f"Lỗi: {str(e)} - {response.text}"

# Hàm gọi API TDNM cho Img to Video
def img_to_video(image, prompt, resolution="720p", duration=4, movement_amplitude="auto", seed=None, user_key=None):
    # Kiểm tra TDNM_KEY
    valid_key, key_message = validate_tdn_key(user_key)
    if not valid_key:
        return None, key_message
    
    if not VIDU_API_KEY:
        return None, "Lỗi: Khóa API TDNM chưa được cấu hình."
    
    if not image:
        return None, "Lỗi: Cần cung cấp một hình ảnh."
    
    if prompt and len(prompt) > 1500:
        return None, "Lỗi: Mô tả văn bản không được vượt quá 1500 ký tự."
    
    # Tải hình ảnh lên TDNM
    image_uri, image_message = upload_image_to_vidu(image, max_size_mb=50)
    if not image_uri:
        return None, image_message
    
    url = f"{VIDU_API_URL}/ent/v2/img2video"
    headers = {
        "Authorization": f"Token {VIDU_API_KEY}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "model": DEFAULT_MODEL,
        "images": [image_uri],
        "prompt": prompt or "",
        "duration": duration,
        "resolution": resolution,
        "movement_amplitude": movement_amplitude
    }
    if seed is not None:
        payload["seed"] = seed
    
    try:
        logger.info(f"Gửi yêu cầu đến API TDNM Img to Video: {payload}")
        response = requests.post(url, json=payload, headers=headers)
        response.raise_for_status()
        result = response.json()
        task_id = result.get("task_id")
        if not task_id:
            return None, "Lỗi: Không nhận được ID tác vụ."
        return task_id, f"Tác vụ được tạo thành công. ID tác vụ: {task_id}"
    except requests.exceptions.RequestException as e:
        logger.error(f"Lỗi API: {str(e)}")
        if response.text:
            logger.error(f"Phản hồi API: {response.text}")
        return None, f"Lỗi: {str(e)} - {response.text}"

# Hàm gọi API TDNM cho Reference to Video
def reference_to_video(reference_images, prompt, resolution="720p", duration=4, aspect_ratio="16:9", movement_amplitude="auto", seed=None, user_key=None):
    # Kiểm tra TDNM_KEY
    valid_key, key_message = validate_tdn_key(user_key)
    if not valid_key:
        return None, key_message
    
    if not VIDU_API_KEY:
        return None, "Lỗi: Khóa API TDNM chưa được cấu hình."
    
    if not reference_images:
        return None, "Lỗi: Cần cung cấp ít nhất một hình ảnh tham chiếu."
    
    if len(reference_images) > 3:
        return None, "Lỗi: Chỉ được cung cấp tối đa 3 hình ảnh tham chiếu."
    
    if not prompt:
        return None, "Lỗi: Cần cung cấp mô tả văn bản."
    
    if prompt and len(prompt) > 1500:
        return None, "Lỗi: Mô tả văn bản không được vượt quá 1500 ký tự."
    
    # Tải hình ảnh lên TDNM
    image_uris = []
    for image in reference_images:
        uri, message = upload_image_to_vidu(image, max_size_mb=50)
        if not uri:
            return None, message
        image_uris.append(uri)
    
    url = f"{VIDU_API_URL}/ent/v2/reference2video"
    headers = {
        "Authorization": f"Token {VIDU_API_KEY}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "model": DEFAULT_MODEL,
        "images": image_uris,
        "prompt": prompt,
        "duration": duration,
        "aspect_ratio": aspect_ratio,
        "resolution": resolution,
        "movement_amplitude": movement_amplitude
    }
    if seed is not None:
        payload["seed"] = seed
    
    try:
        logger.info(f"Gửi yêu cầu đến API TDNM Reference to Video: {payload}")
        response = requests.post(url, json=payload, headers=headers)
        response.raise_for_status()
        result = response.json()
        task_id = result.get("task_id")
        if not task_id:
            return None, "Lỗi: Không nhận được ID tác vụ."
        return task_id, f"Tác vụ được tạo thành công. ID tác vụ: {task_id}"
    except requests.exceptions.RequestException as e:
        logger.error(f"Lỗi API: {str(e)}")
        if response.text:
            logger.error(f"Phản hồi API: {response.text}")
        return None, f"Lỗi: {str(e)} - {response.text}"

# Hàm kiểm tra trạng thái tác vụ
def check_task_status(task_id):
    if not task_id:
        return None, "Lỗi: ID tác vụ không hợp lệ."
    
    url = f"{VIDU_API_URL}/ent/v2/tasks/{task_id}/creations"
    headers = {
        "Authorization": f"Token {VIDU_API_KEY}",
        "Content-Type": "application/json"
    }
    
    start_time = time.time()
    while time.time() - start_time < TIMEOUT:
        try:
            logger.info(f"Kiểm tra trạng thái tác vụ cho task_id: {task_id}")
            response = requests.get(url, headers=headers)
            response.raise_for_status()
            result = response.json()
            state = result.get("state")
            creations = result.get("creations", [])
            err_code = result.get("err_code", "")
            
            if state == "success" and creations:
                video_url = creations[0].get("url")
                if video_url:
                    return video_url, "Tạo video thành công!"
                return None, "Lỗi: Không có URL video trong kết quả."
            elif state == "failed":
                return None, f"Lỗi: Tạo video thất bại. Mã lỗi: {err_code or 'Không xác định'}"
            elif state in ["created", "queueing", "processing"]:
                time.sleep(POLL_INTERVAL)
            else:
                return None, f"Lỗi: Trạng thái không xác định {state}."
        except requests.exceptions.RequestException as e:
            logger.error(f"Lỗi kiểm tra trạng thái: {str(e)}")
            if response.text:
                logger.error(f"Phản hồi kiểm tra trạng thái: {response.text}")
            return None, f"Lỗi khi kiểm tra trạng thái tác vụ: {str(e)} - {response.text}"
    
    return None, f"Lỗi: Hết thời gian tạo video. ID tác vụ: {task_id}"

# Hàm giao diện Gradio cho Start-End to Video
def gradio_start_end_to_video(start_image, end_image, prompt, resolution, duration, movement_amplitude, seed, user_key):
    if not start_image or not end_image:
        return None, "Lỗi: Cần cung cấp cả hai hình ảnh đầu và cuối."
    task_id, message = start_end_to_video(start_image, end_image, prompt, resolution, duration, movement_amplitude, seed, user_key)
    if not task_id:
        return None, message
    video_url, status_message = check_task_status(task_id)
    return video_url, status_message

# Hàm giao diện Gradio cho Img to Video
def gradio_img_to_video(image, prompt, resolution, duration, movement_amplitude, seed, user_key):
    if not image:
        return None, "Lỗi: Cần cung cấp một hình ảnh."
    task_id, message = img_to_video(image, prompt, resolution, duration, movement_amplitude, seed, user_key)
    if not task_id:
        return None, message
    video_url, status_message = check_task_status(task_id)
    return video_url, status_message

# Hàm giao diện Gradio cho Reference to Video
def gradio_reference_to_video(image1, image2, image3, prompt, resolution, duration, aspect_ratio, movement_amplitude, seed, user_key):
    reference_images = [img for img in [image1, image2, image3] if img]
    if not reference_images:
        return None, "Lỗi: Cần cung cấp ít nhất một hình ảnh tham chiếu."
    if not prompt:
        return None, "Lỗi: Cần cung cấp mô tả văn bản."
    task_id, message = reference_to_video(reference_images, prompt, resolution, duration, aspect_ratio, movement_amplitude, seed, user_key)
    if not task_id:
        return None, message
    video_url, status_message = check_task_status(task_id)
    return video_url, status_message

# Giao diện Gradio
with gr.Blocks(title="Trình Tạo Video TDNM") as demo:
    gr.Markdown("# Trình Tạo Video TDNM")
    gr.Markdown("Tạo video với TDNM. Vui lòng nhập khóa bí mật (TDNM_KEY) để sử dụng ứng dụng.")

    # Trường nhập khóa bí mật (áp dụng cho tất cả các tab)
    user_key = gr.Textbox(label="Khóa Bí Mật (TDNM_KEY)", type="password", placeholder="Nhập khóa bí mật của bạn")

    # Tab cho Start-End to Video
    with gr.Tab("Video Chuyển Đổi Hình Ảnh"):
        gr.Markdown("Tải lên hai hình ảnh (đầu và cuối) và mô tả văn bản để tạo video chuyển đổi. Hình ảnh phải là PNG, WebP, JPEG hoặc JPG, kích thước dưới 10MB, tỷ lệ khung hình từ 1:4 đến 4:1.")
        start_image = gr.Image(type="filepath", label="Hình Ảnh Đầu")
        end_image = gr.Image(type="filepath", label="Hình Ảnh Cuối")
        prompt_se = gr.Textbox(label="Mô Tả Văn Bản (Tùy Chọn)", placeholder="Ví dụ: 'Chuyển đổi mượt mà từ khung xe thành xe hoàn chỉnh.'")
        resolution_se = gr.Dropdown(choices=["360p", "720p"], label="Độ Phân Giải", value="720p")
        duration_se = gr.Dropdown(choices=[4], label="Thời Lượng (giây)", value=4)
        movement_amplitude_se = gr.Dropdown(choices=["auto", "small", "medium", "large"], label="Biên Độ Chuyển Động", value="auto")
        seed_se = gr.Number(label="Hạt Giống (Tùy Chọn)", value=None, precision=0)
        se_button = gr.Button("Tạo Video")
        se_video_output = gr.Video(label="Video Được Tạo")
        se_message = gr.Textbox(label="Trạng Thái")

        se_button.click(
            fn=gradio_start_end_to_video,
            inputs=[start_image, end_image, prompt_se, resolution_se, duration_se, movement_amplitude_se, seed_se, user_key],
            outputs=[se_video_output, se_message]
        )

    # Tab cho Img to Video
    with gr.Tab("Video Từ Một Ảnh"):
        gr.Markdown("Tải lên một hình ảnh và mô tả văn bản để tạo video. Hình ảnh phải là PNG, WebP, JPEG hoặc JPG, kích thước dưới 50MB, tỷ lệ khung hình từ 1:4 đến 4:1.")
        image_i2v = gr.Image(type="filepath", label="Hình Ảnh")
        prompt_i2v = gr.Textbox(label="Mô Tả Văn Bản (Tùy Chọn)", placeholder="Ví dụ: 'Phi hành gia vẫy tay và camera di chuyển lên.'")
        resolution_i2v = gr.Dropdown(choices=["360p", "720p"], label="Độ Phân Giải", value="720p")
        duration_i2v = gr.Dropdown(choices=[4], label="Thời Lượng (giây)", value=4)
        movement_amplitude_i2v = gr.Dropdown(choices=["auto", "small", "medium", "large"], label="Biên Độ Chuyển Động", value="auto")
        seed_i2v = gr.Number(label="Hạt Giống (Tùy Chọn)", value=None, precision=0)
        i2v_button = gr.Button("Tạo Video")
        i2v_video_output = gr.Video(label="Video Được Tạo")
        i2v_message = gr.Textbox(label="Trạng Thái")

        i2v_button.click(
            fn=gradio_img_to_video,
            inputs=[image_i2v, prompt_i2v, resolution_i2v, duration_i2v, movement_amplitude_i2v, seed_i2v, user_key],
            outputs=[i2v_video_output, i2v_message]
        )

    # Tab cho Reference to Video
    with gr.Tab("Video Từ Hình Ảnh Tham Chiếu"):
        gr.Markdown("Tải lên 1–3 hình ảnh tham chiếu và mô tả văn bản để tạo video với chủ thể nhất quán. Hình ảnh phải là PNG, WebP, JPEG hoặc JPG, kích thước dưới 50MB, độ phân giải tối thiểu 128x128, tỷ lệ khung hình từ 1:4 đến 4:1.")
        image1 = gr.Image(type="filepath", label="Hình Ảnh Tham Chiếu 1")
        image2 = gr.Image(type="filepath", label="Hình Ảnh Tham Chiếu 2 (Tùy Chọn)")
        image3 = gr.Image(type="filepath", label="Hình Ảnh Tham Chiếu 3 (Tùy Chọn)")
        prompt_ref = gr.Textbox(label="Mô Tả Văn Bản", placeholder="Ví dụ: 'Ông già Noel và gấu ôm nhau bên hồ.'")
        resolution_ref = gr.Dropdown(choices=["360p", "720p"], label="Độ Phân Giải", value="720p")
        duration_ref = gr.Dropdown(choices=[4], label="Thời Lượng (giây)", value=4)
        aspect_ratio_ref = gr.Dropdown(choices=["16:9", "9:16", "1:1"], label="Tỷ Lệ Khung Hình", value="16:9")
        movement_amplitude_ref = gr.Dropdown(choices=["auto", "small", "medium", "large"], label="Biên Độ Chuyển Động", value="auto")
        seed_ref = gr.Number(label="Hạt Giống (Tùy Chọn)", value=None, precision=0)
        ref_button = gr.Button("Tạo Video")
        ref_video_output = gr.Video(label="Video Được Tạo")
        ref_message = gr.Textbox(label="Trạng Thái")

        ref_button.click(
            fn=gradio_reference_to_video,
            inputs=[image1, image2, image3, prompt_ref, resolution_ref, duration_ref, aspect_ratio_ref, movement_amplitude_ref, seed_ref, user_key],
            outputs=[ref_video_output, ref_message]
        )

# Khởi chạy ứng dụng (không cần cho HF Spaces, dùng để kiểm tra cục bộ)
if __name__ == "__main__":
    demo.launch()