import gradio as gr
import os
import subprocess
import time
from gradio.components import HTML
from huggingface_hub import hf_hub_download

# Tạo thư mục làm việc
WORK_DIR = "/app/ComfyUI"
if not os.path.exists(WORK_DIR):
    os.makedirs(WORK_DIR)
    subprocess.run(["git", "clone", "--depth", "1", "https://github.com/comfyanonymous/ComfyUI.git", WORK_DIR])

# Tải WAN 2.1 I2V-14B-720P
MODEL_DIR = os.path.join(WORK_DIR, "models")
os.makedirs(MODEL_DIR, exist_ok=True)

def download_model(repo_id, filename, subfolder, target_subfolder):
    target_dir = os.path.join(MODEL_DIR, target_subfolder)
    os.makedirs(target_dir, exist_ok=True)
    file_path = os.path.join(target_dir, filename)
    if not os.path.exists(file_path):
        hf_hub_download(repo_id=repo_id, filename=filename, subfolder=subfolder, local_dir=target_dir)
        print(f"Đã tải {filename} vào {target_dir}")

repo_id = "Comfy-Org/Wan_2.1_ComfyUI_repackaged"
download_model(repo_id, "wan2.1_i2v_720p_14B_fp8_e4m3fn.safetensors", "split_files/diffusion_models", "diffusion_models")
download_model(repo_id, "wan_2.1_vae.safetensors", "split_files/vae", "vae")
download_model(repo_id, "umt5_xxl_fp8_e4m3fn_scaled.safetensors", "split_files/text_encoders", "text_encoders")
download_model(repo_id, "clip_vision_h.safetensors", "split_files/clip_vision", "clip_vision")

# Khởi động ComfyUI
def start_comfyui():
    subprocess.Popen(["python", "main.py", "--port", "7860", "--listen", "0.0.0.0"], cwd=WORK_DIR)
    time.sleep(5)

start_comfyui()

# Nhúng giao diện
comfyui_interface = """
<iframe src="http://localhost:7860" width="100%" height="600px" style="border:none;"></iframe>
"""

demo = gr.Interface(
    fn=lambda x: "Dùng giao diện ComfyUI bên dưới",
    inputs=gr.Textbox(label="Thông tin", value="ComfyUI với WAN 2.1 I2V-14B-720P"),
    outputs=gr.Textbox(label="Trạng thái"),
    title="ComfyUI WAN 2.1 I2V",
    additional_inputs=[HTML(comfyui_interface)]
)

demo.launch(server_name="0.0.0.0", server_port=7860)