app.py

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 và cài đặt custom nodes (từ Dockerfile của bạn)
CUSTOM_NODES_DIR = os.path.join(WORK_DIR, "custom_nodes")
os.makedirs(CUSTOM_NODES_DIR, exist_ok=True)
custom_nodes = [
    "https://github.com/ltdrdata/ComfyUI-Manager.git",
    "https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git",
    "https://github.com/FizzleDorf/ComfyUI_FizzNodes.git",
    "https://github.com/Stability-AI/stability-ComfyUI-nodes.git",
    "https://github.com/EllangoK/ComfyUI-post-processing-nodes.git",
    "https://github.com/TinyTerra/ComfyUI_tinyterraNodes.git",
    "https://github.com/WASasquatch/was-node-suite-comfyui.git",
    "https://github.com/WASasquatch/PowerNoiseSuite.git",
    "https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved.git"
]
for repo in custom_nodes:
    repo_name = repo.split("/")[-1].replace(".git", "")
    target_dir = os.path.join(CUSTOM_NODES_DIR, repo_name)
    if not os.path.exists(target_dir):
        subprocess.run(["git", "clone", "--depth", "1", repo, target_dir])

# Tải mô hình AnimateDiff
animate_diff_dir = os.path.join(CUSTOM_NODES_DIR, "ComfyUI-AnimateDiff-Evolved/models")
os.makedirs(animate_diff_dir, exist_ok=True)
for model in ["mm_sd_v15_v2.ckpt", "mm_sd_v15.ckpt"]:
    model_path = os.path.join(animate_diff_dir, model)
    if not os.path.exists(model_path):
        subprocess.run(["wget", "-c", f"https://huggingface.co/guoyww/animatediff/resolve/main/{model}", "-P", animate_diff_dir])

# Tải mô hình 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_dir = os.path.join(MODEL_DIR, 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, 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", "diffusion_models")
download_model(repo_id, "wan_2.1_vae.safetensors", "vae")
download_model(repo_id, "umt5_xxl_fp8_e4m3fn_scaled.safetensors", "text_encoders")
download_model(repo_id, "clip_vision_h.safetensors", "clip_vision")

# Khởi động ComfyUI
def start_comfyui():
    subprocess.Popen(["python", "main.py", "--port", "7860", "--listen", "0.0.0.0", "--use-split-cross-attention"], 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 Advanced - WAN 2.1 I2V",
    description="Chạy workflow Image-to-Video 720P.",
    additional_inputs=[HTML(comfyui_interface)]
)

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